Enhance config for DML EP (#1286)

mengniwang95 · pre-commit-ci[bot] · web-flow · commit 6cad5008cc4e · 2023-10-08T13:10:17.000+08:00
Signed-off-by: Mengni Wang &lt;mengni.wang@intel.com&gt;
Co-authored-by: pre-commit-ci[bot] &lt;66853113+pre-commit-ci[bot]@users.noreply.github.com&gt;
diff --git a/docs/source/quantization.md b/docs/source/quantization.md
@@ -497,7 +497,7 @@ Intel(R) Neural Compressor support multi-framework: PyTorch, Tensorflow, ONNX Ru
             <td align="left">DmlExecutionProvider*</td>
             <td align="left">OneDNN</td>
             <td align="left">"onnxrt_dml_ep"</td>
-            <td align="left">NA</td>
+            <td align="left">npu</td>
         </tr>
         <tr>
             <td rowspan="2" align="left">Tensorflow</td>
diff --git a/neural_compressor/adaptor/onnxrt.py b/neural_compressor/adaptor/onnxrt.py
@@ -350,6 +350,10 @@ def _check_backend_available(self, backend):
             logger.warning("Backend `{}` requires a GPU device. Reset device to 'gpu'.".format(backend))
             self.device = "gpu"
 
+        if backend in ["onnxrt_dml_ep"] and self.device != "npu":
+            logger.warning("Backend `{}` requires a NPU device. Reset device to 'npu'.".format(backend))
+            self.device = "npu"
+
         ep = PROVIDERS[backend]
         if ep not in ort.get_available_providers():
             logger.warning(
diff --git a/neural_compressor/config.py b/neural_compressor/config.py
@@ -364,7 +364,7 @@ def device(self):
 
     @device.setter
     def device(self, device):
-        if _check_value("device", device, str, ["cpu", "gpu"]):
+        if _check_value("device", device, str, ["cpu", "gpu", "npu"]):
             self._device = device
 
     @property
@@ -756,7 +756,7 @@ class _BaseQuantizationConfig:
                  'optypes_to_exclude_output_quant': don't quantize output of specified optypes
                  'dedicated_qdq_pair': whether dedicate QDQ pair, only valid for onnxrt_trt_ep
         quant_format: Support 'default', 'QDQ' and 'QOperator', only required in ONNXRuntime.
-        device: Support 'cpu' and 'gpu'.
+        device: Support 'cpu', 'gpu' and 'npu'.
         calibration_sampling_size: Number of calibration sample.
         op_type_dict: Tuning constraints on optype-wise  for advance user to reduce tuning space.
                       User can specify the quantization config by op type:
@@ -1129,7 +1129,7 @@ def device(self):
 
     @device.setter
     def device(self, device):
-        if _check_value("device", device, str, ["cpu", "gpu"]):
+        if _check_value("device", device, str, ["cpu", "gpu", "npu"]):
             self._device = device
 
     @property
@@ -1196,7 +1196,7 @@ class PostTrainingQuantConfig(_BaseQuantizationConfig):
     """Config Class for Post Training Quantization.
 
     Args:
-        device: Support 'cpu' and 'gpu'.
+        device: Support 'cpu', 'gpu' and 'npu'.
         backend: Backend for model execution.
                  Support 'default', 'itex', 'ipex', 'onnxrt_trt_ep', 'onnxrt_cuda_ep', 'onnxrt_dnnl_ep',
                  'onnxrt_dml_ep'
@@ -1381,7 +1381,7 @@ class QuantizationAwareTrainingConfig(_BaseQuantizationConfig):
     """Config Class for Quantization Aware Training.
 
     Args:
-        device: Support 'cpu' and 'gpu'.
+        device: Support 'cpu', 'gpu' and 'npu'.
         backend: Backend for model execution.
                  Support 'default', 'itex', 'ipex', 'onnxrt_trt_ep', 'onnxrt_cuda_ep', 'onnxrt_dnnl_ep',
                  'onnxrt_dml_ep'
@@ -1869,7 +1869,7 @@ class MixedPrecisionConfig(object):
 
     Args:
         device (str, optional): Device for execution.
-                                Support 'cpu' and 'gpu', default is 'cpu'.
+                                Support 'cpu', 'gpu' and 'npu', default is 'cpu'.
         backend (str, optional): Backend for model execution.
                                  Support 'default', 'itex', 'ipex', 'onnxrt_trt_ep', 'onnxrt_cuda_ep', 'onnxrt_dnnl_ep',
                                  'onnxrt_dml_ep'. Default is 'default'.
@@ -2023,7 +2023,7 @@ def device(self):
     @device.setter
     def device(self, device):
         """Set device."""
-        if _check_value("device", device, str, ["cpu", "gpu"]):
+        if _check_value("device", device, str, ["cpu", "gpu", "npu"]):
             self._device = device
 
     @property
diff --git a/test/adaptor/onnxrt_adaptor/test_adaptor_onnxrt.py b/test/adaptor/onnxrt_adaptor/test_adaptor_onnxrt.py
@@ -1712,6 +1712,21 @@ def test_backend(self, mock_warning):
 
         self.assertEqual(mock_warning.call_count, 2)
 
+        framework_specific_info = {
+            "device": "cpu",
+            "backend": "onnxrt_dml_ep",
+            "approach": "post_training_static_quant",
+            "workspace_path": "./nc_workspace",
+        }
+        framework = "onnxrt_qlinearops"
+        adaptor = FRAMEWORKS[framework](framework_specific_info)
+
+        call_args_list = mock_warning.call_args_list
+        first_warning_args = call_args_list[2][0]
+        self.assertEqual(first_warning_args[0], "Backend `onnxrt_dml_ep` requires a NPU device. Reset device to 'npu'.")
+        second_warning_args = call_args_list[3][0]
+        self.assertIn("not in available provider names. Fallback to available providers", second_warning_args[0])
+
     def test_cuda_ep_env_set(self):
         config = PostTrainingQuantConfig(approach="static", backend="onnxrt_cuda_ep", device="gpu", quant_level=1)
         quantization.fit(
