From 6a356d53313a8d63223e6f8281d76d6652bd1cd0 Mon Sep 17 00:00:00 2001
From: uwumarogie <uwumarogiephillips@gmail.com>
Date: Thu, 4 Jun 2026 03:23:23 +0200
Subject: [PATCH] Fix SOFIE convolution dilation codegen

---
 tmva/sofie/inc/TMVA/ROperator_Conv.hxx         |  16 +++++++---------
 tmva/sofie/test/TestCustomModelsFromONNX.cxx   |  17 +++++++++++++++++
 .../test/input_models/ConvWithDilation.onnx    | Bin 0 -> 262 bytes
 .../references/ConvWithDilation.ref.hxx        |   3 +++
 4 files changed, 27 insertions(+), 9 deletions(-)
 create mode 100644 tmva/sofie/test/input_models/ConvWithDilation.onnx
 create mode 100644 tmva/sofie/test/input_models/references/ConvWithDilation.ref.hxx

diff --git a/tmva/sofie/inc/TMVA/ROperator_Conv.hxx b/tmva/sofie/inc/TMVA/ROperator_Conv.hxx
index 78c1dcebc8030..be18559fddc69 100644
--- a/tmva/sofie/inc/TMVA/ROperator_Conv.hxx
+++ b/tmva/sofie/inc/TMVA/ROperator_Conv.hxx
@@ -454,6 +454,7 @@ public:
       // trick for speed is using caffe im2col and output a matrix which contains filtered values as rows.
       // By doing this one has consecutive memory reads and writes
       // Resulting matrix op_xcol is (input channels * filter_h * filter_w , output_h * output_w)
+      // The generated filter tensor is already expanded for fAttrDilations, so Im2col scans it contiguously.
       if (fDim ==1) {
          if (fAttrPads[0] != fAttrPads[1] ) {
             std::cout << "TMVA SOFIE Operator Conv:  asymmetric padding not supported. Assume an average padding "
@@ -493,11 +494,10 @@ public:
                 << fShapeW[1] << "," << iHeight << "," << iWidth << ",";
             if (fDim == 1)
                out << "1, " << fAttrKernelShape[0] << ",0," << fAttrPads[0] << ",1," << fAttrStrides[0] << ",1,"
-                   << fAttrDilations[0];
+                   << 1;
             else // dim ==2
                out << fAttrKernelShape[0] << "," << fAttrKernelShape[1] << "," << fAttrPads[0] << "," << fAttrPads[1]
-                   << "," << fAttrStrides[0] << "," << fAttrStrides[1] << "," << fAttrDilations[0] << ","
-                   << fAttrDilations[1];
+                   << "," << fAttrStrides[0] << "," << fAttrStrides[1] << ",1,1";
             out << "," << "tensor_" <<fNX << "_xcol);\n\n ";
          } else {
             // 3d im2col
@@ -509,8 +509,7 @@ public:
                 << fShapeW[1] << "," << iDepth << "," << iHeight << "," << iWidth << ","
                 << fAttrKernelShape[0] << "," << fAttrKernelShape[1] << "," << fAttrKernelShape[2] << ","
                 << fAttrPads[0] << "," << fAttrPads[1] << "," << fAttrPads[2] << ","
-                << fAttrStrides[0] << "," << fAttrStrides[1] << "," << fAttrStrides[2] << ","
-                << fAttrDilations[0] << "," << fAttrDilations[1] << "," << fAttrDilations[2] << ","
+                << fAttrStrides[0] << "," << fAttrStrides[1] << "," << fAttrStrides[2] << ",1,1,1,"
                 << "tensor_" << fNX << "_xcol);\n\n ";
          }
          // BLAS
@@ -551,11 +550,10 @@ public:
                 << fShapeW[1] << "," << iHeight << "," << iWidth << ",";
             if (fDim == 1)
                out << "1, " << fAttrKernelShape[0] << ",0," << fAttrPads[0] << ",1," << fAttrStrides[0] << ",1,"
-                   << fAttrDilations[0];
+                   << 1;
             else // dim ==2
                out << fAttrKernelShape[0] << "," << fAttrKernelShape[1] << "," << fAttrPads[0] << "," << fAttrPads[1]
-                   << "," << fAttrStrides[0] << "," << fAttrStrides[1] << "," << fAttrDilations[0] << ","
-                   << fAttrDilations[1];
+                   << "," << fAttrStrides[0] << "," << fAttrStrides[1] << ",1,1";
             out << ", tensor_" << fNX << "_xcol);\n\n ";
          } else {
             // 3d im2col
@@ -567,7 +565,7 @@ public:
                 << fShapeW[1] << "," << iDepth << "," << iHeight << "," << iWidth << "," << fAttrKernelShape[0] << ","
                 << fAttrKernelShape[1] << "," << fAttrKernelShape[2] << "," << fAttrPads[0] << "," << fAttrPads[1]
                 << "," << fAttrPads[2] << "," << fAttrStrides[0] << "," << fAttrStrides[1] << "," << fAttrStrides[2]
-                << "," << fAttrDilations[0] << "," << fAttrDilations[1] << "," << fAttrDilations[2] << ",tensor_" << fNX
+                << ",1,1,1,tensor_" << fNX
                 << "_xcol);\n\n ";
          }
 
diff --git a/tmva/sofie/test/TestCustomModelsFromONNX.cxx b/tmva/sofie/test/TestCustomModelsFromONNX.cxx
index b9857830820ac..7c6be791086c7 100644
--- a/tmva/sofie/test/TestCustomModelsFromONNX.cxx
+++ b/tmva/sofie/test/TestCustomModelsFromONNX.cxx
@@ -28,6 +28,7 @@ constexpr auto modelDataSuffix = "_FromONNX.dat";
 #include "input_models/references/ConvWithStridesPadding.ref.hxx"
 #include "input_models/references/ConvWithStridesNoPadding.ref.hxx"
 #include "input_models/references/ConvWithAsymmetricPadding.ref.hxx"
+#include "input_models/references/ConvWithDilation.ref.hxx"
 #include "input_models/references/MaxPool1d.ref.hxx"
 #include "input_models/references/MaxPool2d.ref.hxx"
 #include "input_models/references/MaxPool2d_CeilMode.ref.hxx"
@@ -663,6 +664,22 @@ TEST(ONNX, ConvWithoutPadding)
    }
 }
 
+TEST(ONNX, ConvWithDilation)
+{
+   constexpr float TOLERANCE = DEFAULT_TOLERANCE;
+
+   std::vector<float> input(49);
+   std::iota(input.begin(), input.end(), 0.0f);
+   ASSERT_INCLUDE_AND_RUN(std::vector<float>, "ConvWithDilation", input);
+
+   EXPECT_EQ(output.size(), std::size(ConvWithDilation_ExpectedOutput::outputs));
+
+   float *correct = ConvWithDilation_ExpectedOutput::outputs;
+
+   for (size_t i = 0; i < output.size(); ++i) {
+      EXPECT_LE(std::abs(output[i] - correct[i]), TOLERANCE);
+   }
+}
 
 TEST(ONNX, ConvWithAutopadSameLower)
 {
diff --git a/tmva/sofie/test/input_models/ConvWithDilation.onnx b/tmva/sofie/test/input_models/ConvWithDilation.onnx
new file mode 100644
index 0000000000000000000000000000000000000000..e32e0b7aedff8d5f35cab813891ccc95445c0ced
GIT binary patch
literal 262
zcmd;J7ve0+&o9v}Ni8n1`ozeU&&3$Q#TYKc7%3&-oS#<~o>`LNl9`iOl9`{U!~zo0
z66NB_PA$qy&51A0NGwQoVq#`qz{sv8z{Or%Qk0pJTI|Hc$OsbQ=VB>HOeuC^VE_V<
zfDjjF3e+s1A|@uFn2=C1(9HO9pf&L@IW0pj7Dh&9W&uVem9uBg*q%M(V>fH&OgkXk
z4rr47teKMb#>U3>KymxDv^0As#_%X<pu2^{xI{P@g#@^mI2fUr9Z0h$0VN}`NHPN@
JnVpyfcmXclJ-7e>

literal 0
HcmV?d00001

diff --git a/tmva/sofie/test/input_models/references/ConvWithDilation.ref.hxx b/tmva/sofie/test/input_models/references/ConvWithDilation.ref.hxx
new file mode 100644
index 0000000000000..c343b99b7c7eb
--- /dev/null
+++ b/tmva/sofie/test/input_models/references/ConvWithDilation.ref.hxx
@@ -0,0 +1,3 @@
+namespace ConvWithDilation_ExpectedOutput {
+float outputs[] = {98.4f, 102.9f, 107.4f, 129.9f, 134.4f, 138.9f, 161.4f, 165.9f, 170.4f};
+} // namespace ConvWithDilation_ExpectedOutput