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 2 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 @@ -79,7 +79,13 @@
#define BACKEND_OPENBLAS ""
#endif

#define BACKEND_DEVICE BACKEND_PROC BACKEND_DNNL BACKEND_MKLML BACKEND_NGRAPH BACKEND_OPENVINO BACKEND_NUPHAR BACKEND_OPENBLAS
#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 @@ -107,6 +113,9 @@ std::string nuphar_settings;
#ifdef USE_BRAINSLICE
#include "core/providers/brainslice/brainslice_provider_factory.h"
#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 @@ -117,6 +126,7 @@ std::shared_ptr<IExecutionProviderFactory> CreateExecutionProviderFactory_NGraph
std::shared_ptr<IExecutionProviderFactory> CreateExecutionProviderFactory_OpenVINO(const char* device);
std::shared_ptr<IExecutionProviderFactory> CreateExecutionProviderFactory_Nuphar(bool, const char*);
std::shared_ptr<IExecutionProviderFactory> CreateExecutionProviderFactory_BrainSlice(uint32_t ip, int, int, bool, const char*, const char*, 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_BRAINSLICE
available_providers.push_back(kBrainSliceExecutionProvider);
#endif
#ifdef USE_DML
available_providers.push_back(kDmlExecutionProvider);
#endif
return available_providers;
};
Expand Down Expand Up @@ -333,6 +346,10 @@ void RegisterExecutionProviders(InferenceSession* sess, const std::vector<std::s
} else if (type == kBrainSliceExecutionProvider) {
#ifdef USE_BRAINSLICE
RegisterExecutionProvider(sess, *onnxruntime::CreateExecutionProviderFactory_BrainSlice(0, -1, -1, false, "", "", ""));
#endif
} else if (type == kDmlExecutionProvider) {
#ifdef USE_DML
RegisterExecutionProvider(sess, *onnxruntime::CreateExecutionProviderFactory_DML(0));
#endif
} else {
// unknown provider
Expand Down Expand Up @@ -418,6 +435,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
4 changes: 4 additions & 0 deletions setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,10 @@
package_name = 'onnxruntime-nuphar'
sys.argv.remove('--use_nuphar')

elif '--use_dml' in sys.argv:
package_name = 'onnxruntime-dml'
sys.argv.remove('--use_dml')

if '--nightly_build' in sys.argv:
package_name = 'ort-nightly'
nightly_build = True
Expand Down
6 changes: 4 additions & 2 deletions tools/ci_build/build.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -782,7 +782,7 @@ def nuphar_run_python_tests(build_dir, configs):
run_subprocess([sys.executable, 'onnxruntime_test_python_nuphar.py'], cwd=cwd, dll_path=dll_path)


def build_python_wheel(source_dir, build_dir, configs, use_cuda, use_ngraph, use_dnnl, use_tensorrt, use_openvino, use_nuphar, nightly_build = False):
def build_python_wheel(source_dir, build_dir, configs, use_cuda, use_ngraph, use_dnnl, use_tensorrt, use_openvino, use_nuphar, use_dml, nightly_build = False):
for config in configs:
cwd = get_config_build_dir(build_dir, config)
if is_windows():
Expand All @@ -802,6 +802,8 @@ def build_python_wheel(source_dir, build_dir, configs, use_cuda, use_ngraph, use
args.append('--use_openvino')
elif use_nuphar:
args.append('--use_nuphar')
elif use_dml:
args.append('--use_dml')
run_subprocess(args, cwd=cwd)

def build_protoc_for_host(cmake_path, source_dir, build_dir, args):
Expand Down Expand Up @@ -1051,7 +1053,7 @@ def main():
if args.build:
if args.build_wheel:
nightly_build = bool(os.getenv('NIGHTLY_BUILD') == '1')
build_python_wheel(source_dir, build_dir, configs, args.use_cuda, args.use_ngraph, args.use_dnnl, args.use_tensorrt, args.use_openvino, args.use_nuphar, nightly_build)
build_python_wheel(source_dir, build_dir, configs, args.use_cuda, args.use_ngraph, args.use_dnnl, args.use_tensorrt, args.use_openvino, args.use_nuphar, args.use_dml, nightly_build)

if args.gen_doc and (args.build or args.test):
generate_documentation(source_dir, build_dir, configs)
Expand Down