Update DirectML python bindings to latest SDK

Python/.gitignore

@@ -0,0 +1,4 @@
+*.py[co]
+.eggs
+__pycache__/
+*.egg-info/

Python/CMakeLists.txt

@@ -36,5 +36,18 @@ pybind11_add_module(
 
 target_compile_options(pydirectml PRIVATE /W4 /WX)
 target_include_directories(pydirectml PRIVATE pybind11/include gpgmm/src/include gpgmm/ ${DML_PATH}/include ${DMLX_PATH})
-target_link_directories(pydirectml PRIVATE ${DML_PATH}/bin/x64-win)
+
+# if not windows, raise error
+if(NOT WIN32)
+    message(FATAL_ERROR "This project is only supported on Windows.")
+endif()
+
+if(CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64" OR CMAKE_SYSTEM_PROCESSOR MATCHES "AMD64" OR CMAKE_SYSTEM_PROCESSOR MATCHES "x64")
+    target_link_directories(pydirectml PRIVATE ${DML_PATH}/bin/x64-win)
+elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "ARM64")
+    target_link_directories(pydirectml PRIVATE ${DML_PATH}/bin/arm64-win)
+else()
+    message(FATAL_ERROR "Unsupported architecture: ${CMAKE_SYSTEM_PROCESSOR}")
+endif()
+
 target_link_libraries(pydirectml PRIVATE gpgmm dxgi.lib d3d12.lib directml.lib)

Python/samples/candy.py

@@ -56,7 +56,7 @@
 # instance_norm5
 instance_norm5_scale = dml.input_tensor(builder, 3, dml.TensorDesc(data_type, [1,16,1,1]))
 instance_norm5_bias = dml.input_tensor(builder, 4, dml.TensorDesc(data_type, flags, [1,16,1,1]))
-instance_norm5 = dml.mean_variance_normalization(conv4, instance_norm5_scale, instance_norm5_bias, [0,2,3], 1, 0.000009999999747378752, dml.FusedActivation(dml.OperatorType.ACTIVATION_RELU))
+instance_norm5 = dml.mean_variance_normalization(conv4, instance_norm5_scale, instance_norm5_bias, [0,2,3], True, True, 0.000009999999747378752, dml.FusedActivation(dml.OperatorType.ACTIVATION_RELU))
 
 # conv7
 conv7_filter = dml.input_tensor(builder, 5, dml.TensorDesc(data_type, flags, [32,16,3,3]))
@@ -66,7 +66,7 @@
 # instance_norm8
 instance_norm8_scale = dml.input_tensor(builder, 7, dml.TensorDesc(data_type, [1,32,1,1]))
 instance_norm8_bias = dml.input_tensor(builder, 8, dml.TensorDesc(data_type, flags, [1,32,1,1]))
-instance_norm8 = dml.mean_variance_normalization(conv7, instance_norm8_scale, instance_norm8_bias, [0,2,3], 1, 0.000009999999747378752, dml.FusedActivation(dml.OperatorType.ACTIVATION_RELU))
+instance_norm8 = dml.mean_variance_normalization(conv7, instance_norm8_scale, instance_norm8_bias, [0,2,3], True, True, 0.000009999999747378752, dml.FusedActivation(dml.OperatorType.ACTIVATION_RELU))
 
 # conv10
 conv10_filter = dml.input_tensor(builder, 9, dml.TensorDesc(data_type, flags, [64,32,3,3]))
@@ -76,7 +76,7 @@
 # instance_norm11
 instance_norm11_scale = dml.input_tensor(builder, 11, dml.TensorDesc(data_type, [1,64,1,1]))
 instance_norm11_bias = dml.input_tensor(builder, 12, dml.TensorDesc(data_type, flags, [1,64,1,1]))
-instance_norm11 = dml.mean_variance_normalization(conv10, instance_norm11_scale, instance_norm11_bias, [0,2,3], 1, 0.000009999999747378752, dml.FusedActivation(dml.OperatorType.ACTIVATION_RELU))
+instance_norm11 = dml.mean_variance_normalization(conv10, instance_norm11_scale, instance_norm11_bias, [0,2,3], True, True, 0.000009999999747378752, dml.FusedActivation(dml.OperatorType.ACTIVATION_RELU))
 
 # conv16
 conv16_filter = dml.input_tensor(builder, 13, dml.TensorDesc(data_type, flags, [64,64,3,3]))
@@ -86,7 +86,7 @@
 # instance_norm17
 instance_norm17_scale = dml.input_tensor(builder, 15, dml.TensorDesc(data_type, [1,64,1,1]))
 instance_norm17_bias = dml.input_tensor(builder, 16, dml.TensorDesc(data_type, flags, [1,64,1,1]))
-instance_norm17 = dml.mean_variance_normalization(conv16, instance_norm17_scale, instance_norm17_bias, [0,2,3], 1, 0.000009999999747378752, dml.FusedActivation(dml.OperatorType.ACTIVATION_RELU))
+instance_norm17 = dml.mean_variance_normalization(conv16, instance_norm17_scale, instance_norm17_bias, [0,2,3], True, True, 0.000009999999747378752, dml.FusedActivation(dml.OperatorType.ACTIVATION_RELU))
 
 # conv19
 conv19_filter = dml.input_tensor(builder, 17, dml.TensorDesc(data_type, flags, [64,64,3,3]))
@@ -96,7 +96,7 @@
 # instance_norm15
 instance_norm15_scale = dml.input_tensor(builder, 19, dml.TensorDesc(data_type, [1,64,1,1]))
 instance_norm15_bias = dml.input_tensor(builder, 20, dml.TensorDesc(data_type, flags, [1,64,1,1]))
-instance_norm15 = dml.mean_variance_normalization(conv19, instance_norm15_scale, instance_norm15_bias, [0,2,3], 1, 0.000009999999747378752)
+instance_norm15 = dml.mean_variance_normalization(conv19, instance_norm15_scale, instance_norm15_bias, [0,2,3], True, True, 0.000009999999747378752)
 
 # crop21
 crop21 = dml.slice(instance_norm11,[0,0,2,2],[1,64,196,196],[1,1,1,1])
@@ -112,7 +112,7 @@
 # instance_norm27
 instance_norm27_scale = dml.input_tensor(builder, 23, dml.TensorDesc(data_type, [1,64,1,1]))
 instance_norm27_bias = dml.input_tensor(builder, 24, dml.TensorDesc(data_type, flags, [1,64,1,1]))
-instance_norm27 = dml.mean_variance_normalization(conv26, instance_norm27_scale, instance_norm27_bias, [0,2,3], 1, 0.000009999999747378752, dml.FusedActivation(dml.OperatorType.ACTIVATION_RELU))
+instance_norm27 = dml.mean_variance_normalization(conv26, instance_norm27_scale, instance_norm27_bias, [0,2,3], True, True, 0.000009999999747378752, dml.FusedActivation(dml.OperatorType.ACTIVATION_RELU))
 
 # conv29
 conv29_filter = dml.input_tensor(builder, 25, dml.TensorDesc(data_type, flags, [64,64,3,3]))
@@ -122,7 +122,7 @@
 # instance_norm25
 instance_norm25_scale = dml.input_tensor(builder, 27, dml.TensorDesc(data_type, [1,64,1,1]))
 instance_norm25_bias = dml.input_tensor(builder, 28, dml.TensorDesc(data_type, flags, [1,64,1,1]))
-instance_norm25 = dml.mean_variance_normalization(conv29, instance_norm25_scale, instance_norm25_bias, [0,2,3], 1, 0.000009999999747378752)
+instance_norm25 = dml.mean_variance_normalization(conv29, instance_norm25_scale, instance_norm25_bias, [0,2,3], True, True, 0.000009999999747378752)
 
 # crop21
 crop31 = dml.slice(add,[0,0,2,2],[1,64,196-2-2,196-2-2],[1,1,1,1])
@@ -138,7 +138,7 @@
 # instance_norm37
 instance_norm37_scale = dml.input_tensor(builder, 31, dml.TensorDesc(data_type, [1,64,1,1]))
 instance_norm37_bias = dml.input_tensor(builder, 32, dml.TensorDesc(data_type, flags, [1,64,1,1]))
-instance_norm37 = dml.mean_variance_normalization(conv36, instance_norm37_scale, instance_norm37_bias, [0,2,3], 1, 0.000009999999747378752, dml.FusedActivation(dml.OperatorType.ACTIVATION_RELU))
+instance_norm37 = dml.mean_variance_normalization(conv36, instance_norm37_scale, instance_norm37_bias, [0,2,3], True, True, 0.000009999999747378752, dml.FusedActivation(dml.OperatorType.ACTIVATION_RELU))
 
 # conv39
 conv39_filter = dml.input_tensor(builder, 33, dml.TensorDesc(data_type, flags, [64,64,3,3]))
@@ -148,7 +148,7 @@
 # instance_norm35
 instance_norm35_scale = dml.input_tensor(builder, 35, dml.TensorDesc(data_type, [1,64,1,1]))
 instance_norm35_bias = dml.input_tensor(builder, 36, dml.TensorDesc(data_type, flags, [1,64,1,1]))
-instance_norm35 = dml.mean_variance_normalization(conv39, instance_norm35_scale, instance_norm35_bias, [0,2,3], 1, 0.000009999999747378752)
+instance_norm35 = dml.mean_variance_normalization(conv39, instance_norm35_scale, instance_norm35_bias, [0,2,3], True, True, 0.000009999999747378752)
 
 # crop41
 crop41 = dml.slice(add1,[0,0,2,2],[1, 64, 192-2-2, 192-2-2],[1,1,1,1])
@@ -164,7 +164,7 @@
 # instance_norm47
 instance_norm47_scale = dml.input_tensor(builder, 39, dml.TensorDesc(data_type, [1,64,1,1]))
 instance_norm47_bias = dml.input_tensor(builder, 40, dml.TensorDesc(data_type, flags, [1,64,1,1]))
-instance_norm47 = dml.mean_variance_normalization(conv46, instance_norm47_scale, instance_norm47_bias, [0,2,3], 1, 0.000009999999747378752, dml.FusedActivation(dml.OperatorType.ACTIVATION_RELU))
+instance_norm47 = dml.mean_variance_normalization(conv46, instance_norm47_scale, instance_norm47_bias, [0,2,3], True, True, 0.000009999999747378752, dml.FusedActivation(dml.OperatorType.ACTIVATION_RELU))
 
 # conv49
 conv49_filter = dml.input_tensor(builder, 41, dml.TensorDesc(data_type, flags, [64,64,3,3]))
@@ -174,7 +174,7 @@
 # instance_norm45
 instance_norm45_scale = dml.input_tensor(builder, 43, dml.TensorDesc(data_type, [1,64,1,1]))
 instance_norm45_bias = dml.input_tensor(builder, 44, dml.TensorDesc(data_type, flags, [1,64,1,1]))
-instance_norm45 = dml.mean_variance_normalization(conv49, instance_norm45_scale, instance_norm45_bias, [0,2,3], 1, 0.000009999999747378752)
+instance_norm45 = dml.mean_variance_normalization(conv49, instance_norm45_scale, instance_norm45_bias, [0,2,3], True, True, 0.000009999999747378752)
 
 # crop51
 crop51 = dml.slice(instance_norm35,[0,0,2,2],[1, 64, 188-2-2, 188-2-2],[1,1,1,1])
@@ -190,7 +190,7 @@
 # instance_norm57
 instance_norm57_scale = dml.input_tensor(builder, 47, dml.TensorDesc(data_type, [1,64,1,1]))
 instance_norm57_bias = dml.input_tensor(builder, 48, dml.TensorDesc(data_type, flags, [1,64,1,1]))
-instance_norm57 = dml.mean_variance_normalization(conv56, instance_norm57_scale, instance_norm57_bias, [0,2,3], 1, 0.000009999999747378752, dml.FusedActivation(dml.OperatorType.ACTIVATION_RELU))
+instance_norm57 = dml.mean_variance_normalization(conv56, instance_norm57_scale, instance_norm57_bias, [0,2,3], True, True, 0.000009999999747378752, dml.FusedActivation(dml.OperatorType.ACTIVATION_RELU))
 
 # conv59
 conv59_filter = dml.input_tensor(builder, 49, dml.TensorDesc(data_type, flags, [64,64,3,3]))
@@ -200,7 +200,7 @@
 # instance_norm55
 instance_norm55_scale = dml.input_tensor(builder, 51, dml.TensorDesc(data_type, [1,64,1,1]))
 instance_norm55_bias = dml.input_tensor(builder, 52, dml.TensorDesc(data_type, flags, [1,64,1,1]))
-instance_norm55 = dml.mean_variance_normalization(conv59, instance_norm55_scale, instance_norm55_bias, [0,2,3], 1, 0.000009999999747378752)
+instance_norm55 = dml.mean_variance_normalization(conv59, instance_norm55_scale, instance_norm55_bias, [0,2,3], True, True, 0.000009999999747378752)
 
 # crop61
 crop61 = dml.slice(add3,[0,0,2,2],[1, 64, 184-2-2, 184-2-2],[1,1,1,1])
@@ -219,7 +219,7 @@
 # instance_norm65
 instance_norm65_scale = dml.input_tensor(builder, 55, dml.TensorDesc(data_type, [1,32,1,1]))
 instance_norm65_bias = dml.input_tensor(builder, 56, dml.TensorDesc(data_type, flags, [1,32,1,1]))
-instance_norm65 = dml.mean_variance_normalization(crop63, instance_norm65_scale, instance_norm65_bias, [0,2,3], 1, 0.000009999999747378752, dml.FusedActivation(dml.OperatorType.ACTIVATION_RELU))
+instance_norm65 = dml.mean_variance_normalization(crop63, instance_norm65_scale, instance_norm65_bias, [0,2,3], True, True, 0.000009999999747378752, dml.FusedActivation(dml.OperatorType.ACTIVATION_RELU))
 
 # conv67
 conv67_filter = dml.input_tensor(builder, 57, dml.TensorDesc(data_type, flags, [32,16,3,3]))
@@ -232,7 +232,7 @@
 # instance_norm69
 instance_norm69_scale = dml.input_tensor(builder, 59, dml.TensorDesc(data_type, [1,16,1,1]))
 instance_norm69_bias = dml.input_tensor(builder, 60, dml.TensorDesc(data_type, flags, [1,16,1,1]))
-instance_norm69 = dml.mean_variance_normalization(crop67, instance_norm69_scale, instance_norm69_bias, [0,2,3], 1, 0.000009999999747378752, dml.FusedActivation(dml.OperatorType.ACTIVATION_RELU))
+instance_norm69 = dml.mean_variance_normalization(crop67, instance_norm69_scale, instance_norm69_bias, [0,2,3], True, True, 0.000009999999747378752, dml.FusedActivation(dml.OperatorType.ACTIVATION_RELU))
 
 # conv71
 conv71_filter = dml.input_tensor(builder, 61, dml.TensorDesc(data_type, flags, [3,16,9,9]))

Python/samples/mobilenet.py

@@ -720,7 +720,7 @@ def append_input_tensor(builder: dml.GraphBuilder, input_bindings: list, input_t
 batch_norm53 = dml.batch_normalization(conv53, batch_norm53_mean, batch_norm53_variance, batch_norm53_scale, batch_norm53_bias, 1, 0.000009999999747378752, dml.FusedActivation(dml.OperatorType.ACTIVATION_RELU))
 
 # avg_pool1
-avg_pool1 = dml.average_pooling(batch_norm53, [1,1], [7,7], [0,0], [0,0], 0)
+avg_pool1 = dml.average_pooling(batch_norm53, [1,1], [7,7], [0,0], [0,0], [0,0], 0)
 
 # conv54
 conv54_filter = append_input_tensor(builder, input_bindings, dml.TensorDesc(data_type, flags, [1000,1280,1,1]), "mobilenetv20_output_pred_weight.npy")

Python/samples/squeezenet.py

@@ -223,7 +223,7 @@ def append_input_tensor(builder: dml.GraphBuilder, input_bindings: list, input_t
 squeezenet0_conv25_fwd = dml.convolution(squeezenet0_dropout0_fwd, squeezenet0_conv25_weight, squeezenet0_conv25_bias, fused_activation = dml.FusedActivation(dml.OperatorType.ACTIVATION_RELU))
 
 # squeezenet0_pool3_fwd
-squeezenet0_pool3_fwd = dml.average_pooling(squeezenet0_conv25_fwd, [13,13], [13,13], [0,0], [0,0], 0)
+squeezenet0_pool3_fwd = dml.average_pooling(squeezenet0_conv25_fwd, [13,13], [13,13], [0,0], [0,0], [0,0], 0)
 
 # squeezenet0_flatten0_reshape0
 squeezenet0_flatten0_reshape0 =dml.reinterpret(squeezenet0_pool3_fwd, dml.TensorDataType.FLOAT32, [1,1,1,1000], [1000,1000,1000,1])

Python/setup.py

@@ -12,22 +12,28 @@
 from setuptools.command.build_ext import build_ext
 from distutils.version import LooseVersion
 
+target_arch = platform.machine().lower()
+if target_arch == 'amd64' or target_arch == 'x86_64':
+    target_arch = 'x64'
+elif target_arch == 'aarch64' or target_arch == 'arm64':
+    target_arch = 'arm64'
+
 dml_feed_url = 'https://api.nuget.org/v3/index.json'
 dml_resource_id = 'microsoft.ai.directml'
-dml_resource_version = '1.9.1'
+dml_resource_version = '1.15.4'
 
 dependency_dir = 'dependencies'
-dml_bin_path = f'{dependency_dir}/{dml_resource_id}.{dml_resource_version}/bin/x64-win/'
-lib_dir = '..\libraries'
+dml_bin_path = f'{dependency_dir}/{dml_resource_id}.{dml_resource_version}/bin/{target_arch}-win/'
+lib_dir = '..\\libraries'
 base_path = os.path.dirname(os.path.realpath(__file__))
 dependency_path = os.path.join(base_path, dependency_dir)
 
 dml_resource_name = '.'.join([dml_resource_id, dml_resource_version])
-dml_path = '%s\%s' % (dependency_path, dml_resource_name)
+dml_path = os.path.join(dependency_path, dml_resource_name)
 
 dmlx = 'DirectMLX'
 dmlx_file = dmlx + '.h'
-dmlx_source_path = '%s\%s' % (os.path.join(base_path, lib_dir), dmlx_file)
+dmlx_source_path = os.path.join(base_path, lib_dir, dmlx_file)
 dmlx_path = os.path.join(dependency_path, dmlx)
 
 class CMakeExtension(Extension):
@@ -105,7 +111,7 @@ def build_extension(self, ext):
         if platform.system() == "Windows":
             cmake_args += ['-DCMAKE_LIBRARY_OUTPUT_DIRECTORY_{}={}'.format(cfg.upper(), extdir)]
             if sys.maxsize > 2**32:
-                cmake_args += ['-A', 'x64']
+                cmake_args += ['-A', target_arch]
             cmake_args += ['-DDML_PATH={}'.format(dml_path)]
             cmake_args += ['-DDMLX_PATH={}'.format(dmlx_path)]
             build_args += ['--', '/m']

Python/src/module.cpp

@@ -425,15 +425,17 @@ PYBIND11_MODULE(pydirectml, module)
         dml::Optional<dml::Expression> bias,
         std::vector<uint32_t> axes,
         bool normalizeVariance,
+        bool normalizeMean,
         float epsilon,
         dml::FusedActivation fusedActivation) {
-            return dml::MeanVarianceNormalization(input, scale, bias, axes, normalizeVariance, epsilon, fusedActivation);
+            return dml::MeanVarianceNormalization(input, scale, bias, axes, normalizeVariance, normalizeMean, epsilon, fusedActivation);
         }, "Normalize inputs using output = scale * (input - mean) / sqrt(variance + epsilon) + bias, where mean and variance are computed per instance per channel.",
         py::arg("input"),
         py::arg("scale") = dml::NullOpt,
         py::arg("bias") = dml::NullOpt,
         py::arg("axes") = std::vector<uint32_t>{},
         py::arg("normalize_variance"),
+        py::arg("normalize_mean"),
         py::arg("epsilon"),
         py::arg("fused_activation") = dml::FusedActivation::None());
 
@@ -511,15 +513,17 @@ PYBIND11_MODULE(pydirectml, module)
         std::vector<uint32_t> windowSizes,
         std::vector<uint32_t> startPadding,
         std::vector<uint32_t> endPadding,
+        std::vector<uint32_t> dilations,
         bool includePadding) {
-            return dml::AveragePooling(input, strides, windowSizes, startPadding, endPadding, includePadding);
+            return dml::AveragePooling(input, strides, windowSizes, startPadding, endPadding, dilations, includePadding);
         },
         "Average all elements in each pool.",
         py::arg("input"),
         py::arg("strides"),
         py::arg("window_sizes"),
         py::arg("start_padding"),
         py::arg("end_padding"),
+        py::arg("dilations"),
         py::arg("include_padding"));
 
     module.def("max_pooling", [](
@@ -554,7 +558,7 @@ PYBIND11_MODULE(pydirectml, module)
         py::arg("new_size"),
         py::arg("new_strides"));
 
-    module.def("activation_soft_max", &dml::ActivationSoftmax, "Raise all elements to e, and divide all the elements in each batch by that batch's sum.",
+    module.def("activation_soft_max", py::overload_cast<dml::Expression>(&dml::ActivationSoftmax), "Raise all elements to e, and divide all the elements in each batch by that batch's sum.",
         py::arg("input"));
 
     module.def("join", [](