Skip to content

Expose DirectML provider to Python #3359

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 4 commits into from
Closed

Expose DirectML provider to Python #3359

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
1d9b586
expose directml provider to python
RobinKa Mar 29, 2020
da6a0b7
add -dml to package name when use_dml is specified
RobinKa Mar 29, 2020
e9d2ae5
Merge remote-tracking branch 'upstream/master'
RobinKa May 5, 2020
dae9a6f
fixed typo (wrong argument order in build_python_wheel for dml)
RobinKa May 5, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 21 additions & 1 deletion onnxruntime/python/onnxruntime_pybind_state.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,7 +88,13 @@
#define BACKEND_OPENBLAS ""
#endif

#define BACKEND_DEVICE BACKEND_PROC BACKEND_DNNL BACKEND_MKLML BACKEND_NGRAPH BACKEND_NUPHAR BACKEND_OPENBLAS BACKEND_OPENVINO
#if USE_DML
#define BACKEND_DML "-DML"
#else
#define BACKEND_DML ""
#endif

#define BACKEND_DEVICE BACKEND_PROC BACKEND_DNNL BACKEND_MKLML BACKEND_NGRAPH BACKEND_OPENVINO BACKEND_NUPHAR BACKEND_OPENBLAS BACKEND_DML
#include "core/session/onnxruntime_cxx_api.h"
#include "core/providers/providers.h"
#include "core/providers/cpu/cpu_execution_provider.h"
Expand Down Expand Up @@ -117,6 +123,9 @@ std::string openvino_device;
#include "core/providers/nuphar/nuphar_provider_factory.h"
std::string nuphar_settings;
#endif
#ifdef USE_DML
#include "core/providers/dml/dml_provider_factory.h"
#endif

namespace onnxruntime {
std::shared_ptr<IExecutionProviderFactory> CreateExecutionProviderFactory_CPU(int use_arena);
Expand All @@ -128,6 +137,7 @@ std::shared_ptr<IExecutionProviderFactory> CreateExecutionProviderFactory_Dnnl(i
std::shared_ptr<IExecutionProviderFactory> CreateExecutionProviderFactory_NGraph(const char* ng_backend_type);
std::shared_ptr<IExecutionProviderFactory> CreateExecutionProviderFactory_OpenVINO(const char* device);
std::shared_ptr<IExecutionProviderFactory> CreateExecutionProviderFactory_Nuphar(bool, const char*);
std::shared_ptr<IExecutionProviderFactory> CreateExecutionProviderFactory_DML(int device_id);
} // namespace onnxruntime

#if defined(_MSC_VER)
Expand Down Expand Up @@ -293,6 +303,9 @@ const std::vector<std::string>& GetAvailableProviders() {
#endif
#ifdef USE_NUPHAR
available_providers.push_back(kNupharExecutionProvider);
#endif
#ifdef USE_DML
available_providers.push_back(kDmlExecutionProvider);
#endif
return available_providers;
};
Expand Down Expand Up @@ -332,6 +345,10 @@ void RegisterExecutionProviders(InferenceSession* sess, const std::vector<std::s
#if USE_NUPHAR
RegisterExecutionProvider(sess, *onnxruntime::CreateExecutionProviderFactory_Nuphar(true, nuphar_settings.c_str()));
nuphar_settings.clear(); // clear nuphar_settings after use to avoid it being accidentally passed on to next session
#endif
} else if (type == kDmlExecutionProvider) {
#ifdef USE_DML
RegisterExecutionProvider(sess, *onnxruntime::CreateExecutionProviderFactory_DML(0));
#endif
} else {
// unknown provider
Expand Down Expand Up @@ -420,6 +437,9 @@ void addGlobalMethods(py::module& m, const Environment& env) {
#endif
#ifdef USE_TENSORRT
onnxruntime::CreateExecutionProviderFactory_Tensorrt(0)
#endif
#ifdef USE_DML
onnxruntime::CreateExecutionProviderFactory_DML(0)
#endif
};

Expand Down
3 changes: 3 additions & 0 deletions setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,9 @@
elif '--use_nuphar' in sys.argv:
package_name = 'onnxruntime-nuphar'
sys.argv.remove('--use_nuphar')
elif '--use_dml' in sys.argv:
package_name = 'onnxruntime-dml'
sys.argv.remove('--use_dml')
# --use_acl is specified in build.py, but not parsed here


Expand Down
6 changes: 5 additions & 1 deletion tools/ci_build/build.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1343,7 +1343,8 @@ def nuphar_run_python_tests(build_dir, configs):

def build_python_wheel(
source_dir, build_dir, configs, use_cuda, use_ngraph, use_dnnl,
use_tensorrt, use_openvino, use_nuphar, wheel_name_suffix, use_acl,
use_tensorrt, use_openvino, use_nuphar, use_dml, wheel_name_suffix,
use_acl,
nightly_build=False,
featurizers_build=False):
for config in configs:
Expand Down Expand Up @@ -1375,6 +1376,8 @@ def build_python_wheel(
args.append('--use_dnnl')
elif use_nuphar:
args.append('--use_nuphar')
elif use_dml:
args.append('--use_dml')
elif use_acl:
args.append('--use_acl')

Expand Down Expand Up @@ -1711,6 +1714,7 @@ def main():
args.use_tensorrt,
args.use_openvino,
args.use_nuphar,
args.use_dml,
args.wheel_name_suffix,
args.use_acl,
nightly_build=nightly_build,
Expand Down