Add support for 3d convolution

microsoft · Sep 21, 2023 · 42bfcc9 · 42bfcc9
1 parent 02fe4b5
commit 42bfcc9
Show file tree

Hide file tree

Showing 4 changed files with 72 additions and 7 deletions.
diff --git a/DxDispatch/CMakeLists.txt b/DxDispatch/CMakeLists.txt
@@ -251,7 +251,8 @@ if(DXD_TESTS)
         set_tests_properties(test_${model_name} PROPERTIES PASS_REGULAR_EXPRESSION ${expected_output})
     endfunction()
 
-    model_test(dml_convolution "Resource 'output': 6, 8, 12, 14")
+    model_test(dml_convolution_2d "Resource 'output': 6, 8, 12, 14")
+    model_test(dml_convolution_3d "Resource 'output': 4, 4, 4, 4, 4, 4, 4, 4")
     model_test(dml_cumulative_product "Resource 'Out': 2, 8, 64, 192")
     model_test(dml_element_wise_add "Resource 'Out': 6, 10, -2")
     model_test(dml_element_wise_add_npy "Resource 'Out': 2, 4, 6, 8, 10, 12")

diff --git a/DxDispatch/models/dml_convolution.json → DxDispatch/models/dml_convolution_2d.json b/DxDispatch/models/dml_convolution.json → DxDispatch/models/dml_convolution_2d.json
diff --git a/DxDispatch/models/dml_convolution_3d.json b/DxDispatch/models/dml_convolution_3d.json
@@ -0,0 +1,60 @@
+{
+    "$schema": "./_schema.json",
+
+    "resources": 
+    {
+        "input": 
+        {
+            "initialValuesDataType": "FLOAT32",
+            "initialValues": { "valueCount": 27, "valueStart": 1, "valueDelta": 0 }
+        },
+        "filter": 
+        {
+            "initialValuesDataType": "FLOAT32",
+            "initialValues": [0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5]
+        },
+        "output": 
+        {
+            "initialValuesDataType": "FLOAT32",
+            "initialValues": { "valueCount": 4, "value": 0 }
+        }
+    },
+
+    "dispatchables": 
+    {
+        "conv3d": 
+        {
+            "type": "DML_OPERATOR_CONVOLUTION",
+            "desc": 
+            {
+                "InputTensor": { "DataType": "FLOAT32", "Sizes": [1,1,3,3,3] },
+                "FilterTensor": { "DataType": "FLOAT32", "Sizes": [1,1,2,2,2] },
+                "OutputTensor": { "DataType": "FLOAT32", "Sizes": [1,1,2,2,2] },
+                "Mode": "DML_CONVOLUTION_MODE_CROSS_CORRELATION",
+                "Direction": "DML_CONVOLUTION_DIRECTION_FORWARD",
+                "DimensionCount": 3,
+                "Strides": [1,1,1],
+                "Dilations": [1,1,1],
+                "StartPadding": [0,0,0],
+                "EndPadding": [0,0,0],
+                "OutputPadding": [0,0,0],
+                "GroupCount": 1
+            }
+        }
+    },
+
+    "commands": 
+    [
+        {
+            "type": "dispatch",
+            "dispatchable": "conv3d",
+            "bindings":
+            {
+                "InputTensor": "input",
+                "FilterTensor": "filter",
+                "OutputTensor": "output"
+            }
+        },
+        { "type": "print", "resource": "output" }
+    ]
+}
diff --git a/Libraries/DirectMLX.h b/Libraries/DirectMLX.h
@@ -3547,11 +3547,13 @@ namespace dml
         uint32_t dimensionCount = static_cast<uint32_t>(inputTensor.sizes.size());
 
         // todo: support 1d convolution?
-        assert(dimensionCount == 4);
+        assert(dimensionCount == 4 || dimensionCount == 5);
         uint32_t spatialDimensionCount = dimensionCount - 2;
 
-        const uint32_t defaultStridesAndDilations[2] = { 1, 1 };
-        const uint32_t defaultPadding[2] = { 0, 0 };
+        // If the spatial dimension count is 2, we'll just use the first two elements by setting
+        // DimensionCount = 2 in the desc
+        const uint32_t defaultStridesAndDilations[3] = { 1, 1, 1 };
+        const uint32_t defaultPadding[3] = { 0, 0, 0 };
 
         assert(strides.empty() || strides.size() == spatialDimensionCount);
         assert(dilations.empty() || dilations.size() == spatialDimensionCount);
@@ -3659,11 +3661,13 @@ namespace dml
         uint32_t dimensionCount = static_cast<uint32_t>(inputTensor.sizes.size());
 
         // todo: suppord 1d convolution?
-        assert(dimensionCount == 4);
+        assert(dimensionCount == 4 || dimensionCount == 5);
         const uint32_t spatialDimensionCount = dimensionCount - 2;
 
-        const uint32_t defaultStridesAndDilations[2] = { 1, 1 };
-        const uint32_t defaultPadding[2] = { 1, 1 };
+        // If the spatial dimension count is 2, we'll just use the first two elements by setting
+        // DimensionCount = 2 in the desc
+        const uint32_t defaultStridesAndDilations[3] = { 1, 1, 1 };
+        const uint32_t defaultPadding[3] = { 0, 0, 0 };
 
         assert(strides.empty() || strides.size() == spatialDimensionCount);
         assert(dilations.empty() || dilations.size() == spatialDimensionCount);