From 589612106aea3033d4f574cd194a532ec421bd46 Mon Sep 17 00:00:00 2001
From: Patrice Vignola <vignola.patrice@gmail.com>
Date: Tue, 3 Jan 2023 12:46:33 -0800
Subject: [PATCH] [DML EP] Force layer norm inputs to be 4D to better target
 metacommands (#14022)

### Description
Force layer norm inputs to be 4D to better target metacommands

### Motivation and Context
This may improve performance on some hardware by allowing the driver to
return valid layouts to DML when querying for metacommand support.
---
 .../DmlOperatorLayerNormalization.cpp         | 29 ++++++++++++++++---
 1 file changed, 25 insertions(+), 4 deletions(-)

diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorLayerNormalization.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorLayerNormalization.cpp
index 9b366649de..5c64059f7c 100644
--- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorLayerNormalization.cpp
+++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorLayerNormalization.cpp
@@ -25,12 +25,33 @@ public:
             std::nullopt,
             kernelCreationContext.GetTensorShapeDescription().GetInputTensorDimensionCount(0));
 
-        const float epsilon = kernelCreationContext.GetOptionalAttribute<float>(AttrName::Epsilon, DefaultEpsilon);
+        constexpr static uint32_t minimumDimensionCount = 4;
 
+        // Pad the input and the output with trailing 1's until they are at least 4D
+        for (uint32_t i = 0; i < kernelCreationContext.GetInputCount(); ++i)
+        {
+            auto sizes = m_inputTensorDescs[i].GetSizes();
+            std::vector<uint32_t> tensorShape(sizes.begin(), sizes.end());
+            tensorShape.resize(std::max<size_t>(tensorShape.size(), minimumDimensionCount), 1);
+
+            if (m_inputTensorDescs[i].GetDmlDataType() != DML_TENSOR_TYPE_INVALID)
+            {
+                m_inputTensorDescs[i] = TensorDesc(
+                    m_inputTensorDescs[i].GetDmlDataType(),
+                    tensorShape);
+            }
+        }
+
+        m_outputTensorDescs[0] = TensorDesc(
+            m_outputTensorDescs[0].GetDmlDataType(),
+            m_inputTensorDescs[0].GetSizes());
+
+        const float epsilon = kernelCreationContext.GetOptionalAttribute<float>(AttrName::Epsilon, DefaultEpsilon);
         int32_t onnxAxis = kernelCreationContext.GetOptionalAttribute<int32_t>(AttrName::Axis, -1);
-        uint32_t inputDimCount = kernelCreationContext.GetTensorShapeDescription().GetInputTensorDimensionCount(0);
-        onnxAxis = OperatorHelper::HandleNegativeAxis(onnxAxis, inputDimCount);
-        std::vector<uint32_t> onnxAxes(static_cast<size_t>(inputDimCount) - static_cast<size_t>(onnxAxis));
+        uint32_t onnxDimCount = kernelCreationContext.GetTensorShapeDescription().GetInputTensorDimensionCount(0);
+        uint32_t dmlDimCount = m_inputTensorDescs[0].GetDimensionCount();
+        onnxAxis = OperatorHelper::HandleNegativeAxis(onnxAxis, onnxDimCount);
+        std::vector<uint32_t> onnxAxes(static_cast<size_t>(dmlDimCount) - static_cast<size_t>(onnxAxis));
         std::iota(onnxAxes.begin(), onnxAxes.end(), onnxAxis);
 
         assert(m_inputTensorDescs.size() == 3);