From 3b53be001b274b6f40ffde4650f3fd2ad6a64d73 Mon Sep 17 00:00:00 2001
From: Dwayne Robinson <dwayner@microsoft.com>
Date: Wed, 22 Jul 2020 20:51:03 +0000
Subject: [PATCH] Merged PR 4954300: Enable preliminary 8D support in DirectML
 EP

- Update TensorDesc maximum dimension count to 8.
- Update OneHot kernel dimension range.

Related work items: #27678579
---
 .../core/providers/dml/DmlExecutionProvider/src/DmlCommon.h | 2 +-
 .../src/Operators/DmlOperatorOneHot.cpp                     | 6 ++++--
 .../providers/dml/DmlExecutionProvider/src/TensorDesc.h     | 2 --
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommon.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommon.h
index e1fd0695cd..536fa9beeb 100644
--- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommon.h
+++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommon.h
@@ -10,7 +10,7 @@ namespace Dml
 {
     using namespace OperatorHelper;
 
-    static const int MaximumDimensionCount = DML_TENSOR_DIMENSION_COUNT_MAX;
+    static const int MaximumDimensionCount = DML_TENSOR_DIMENSION_COUNT_MAX1;
 
     DML_TENSOR_DATA_TYPE GetDmlDataTypeFromMlDataType(MLOperatorTensorDataType tensorDataType);
     DML_TENSOR_DATA_TYPE GetDmlDataTypeFromMlDataTypeNoThrow(MLOperatorTensorDataType tensorDataType) noexcept;
diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorOneHot.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorOneHot.cpp
index ec01a54acc..f1ba681d99 100644
--- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorOneHot.cpp
+++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorOneHot.cpp
@@ -37,7 +37,7 @@ public:
                 TensorAxis::DoNotCoerce,
                 TensorAxis::W,
                 TensorAxis::RightAligned,
-                NchwDimensionCount, // minDimensionCount
+                1, // minDimensionCount
                 0
             );
 
@@ -49,10 +49,12 @@ public:
                 TensorAxis::DoNotCoerce,
                 TensorAxis::W,
                 TensorAxis::RightAligned,
-                NchwDimensionCount, // minDimensionCount
+                1, // minDimensionCount
                 0
             );
 
+        DmlOperator::Remap64bitDmlDataTypesTo32bitIfNeeded();
+
         // Adjust the axis so it's in DML's terms rather than the original ONNX indexing.
         uint32_t dmlAxis = GetDmlAdjustedAxis(
             m_absoluteAxis,
diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/TensorDesc.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/TensorDesc.h
index ea1af04ad6..39f4b0e56c 100644
--- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/TensorDesc.h
+++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/TensorDesc.h
@@ -59,8 +59,6 @@ namespace Dml
         DML_BUFFER_TENSOR_DESC m_bufferTensorDesc = {};
     };
 
-
-
     class TensorDescBuilder
     {
     public: