From f55c1d6287be8b44314b34da57aa9f7f99db9910 Mon Sep 17 00:00:00 2001
From: Harsh Chauhan <73985490+harz05@users.noreply.github.com>
Date: Fri, 5 Jun 2026 08:17:12 +0530
Subject: [PATCH] [tmva][sofie] Fix double-counted dilation in Conv im2col

Conv generated invalid code and crashed (segfault) for dilation > 1. The operator builds a dilated `_f` weight buffer and expands `fAttrKernelShape` to the effective kernel size, but the im2col call was still passed the real dilation, so dilation was applied twice (expanded kernel + dilation). This produced a negative output dimension and an out-of-bounds write.

Fix: after the dilated `_f` is built, the dense im2col must use dilation 1, since the dilation is already folded into the expanded kernel shape and the `_f` layout. This is a no-op for dilation 1, so existing Conv tests are unaffected.

Also adds a `ConvWithDilation` test (3x3 kernel, dilation 2) with reference output. The suite had no dilation > 1 case before, which is why this was never caught. The full SOFIE test suite passes locally on master.

(cherry picked from commit d42a7a7780cbf4ac673b3a8344e3e1c27e2032ab)
---
 tmva/sofie/inc/TMVA/ROperator_Conv.hxx        |   4 ++++
 tmva/sofie/test/TestCustomModelsFromONNX.cxx  |  22 ++++++++++++++++++
 .../test/input_models/ConvWithDilation.onnx   | Bin 0 -> 269 bytes
 .../references/ConvWithDilation.ref.hxx       |   3 +++
 4 files changed, 29 insertions(+)
 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..4e2e90cbc130c 100644
--- a/tmva/sofie/inc/TMVA/ROperator_Conv.hxx
+++ b/tmva/sofie/inc/TMVA/ROperator_Conv.hxx
@@ -429,6 +429,10 @@ public:
       out << SP << SP << "}\n";
       out << SP << "}\n";
 
+      // Dilation is already folded into the expanded kernel shape and the dilated tensor_<X>_f
+      // layout above, so the dense im2col below must use dilation 1 to avoid double-counting it.
+      fAttrDilations = std::vector<size_t>(3, 1);
+
       //out << SP << "char " << OpName << "_transA = 'T';\n";
       out << SP << "char " << OpName << "_transA = 'N';\n";
       out << SP << "char " << OpName << "_transB = 'N';\n";
diff --git a/tmva/sofie/test/TestCustomModelsFromONNX.cxx b/tmva/sofie/test/TestCustomModelsFromONNX.cxx
index b9857830820ac..33eb135bcdbd7 100644
--- a/tmva/sofie/test/TestCustomModelsFromONNX.cxx
+++ b/tmva/sofie/test/TestCustomModelsFromONNX.cxx
@@ -22,6 +22,7 @@ constexpr auto modelDataSuffix = "_FromONNX.dat";
 #include "input_models/references/Erf.ref.hxx"
 #include "input_models/references/LinearWithSigmoid.ref.hxx"
 #include "input_models/references/ConvWithPadding.ref.hxx"
+#include "input_models/references/ConvWithDilation.ref.hxx"
 #include "input_models/references/ConvWithoutPadding.ref.hxx"
 #include "input_models/references/ConvWithAutopadSameLower.ref.hxx"
 #include "input_models/references/ConvWithAutopadSameUpper.ref.hxx"
@@ -726,6 +727,27 @@ TEST(ONNX, ConvWithStridesPadding)
 }
 
 
+TEST(ONNX, ConvWithDilation)
+{
+   constexpr float TOLERANCE = DEFAULT_TOLERANCE;
+
+   // Preparing the standard all-ones input
+   std::vector<float> input(49);
+   std::iota(input.begin(), input.end(), 0.0f);
+   ASSERT_INCLUDE_AND_RUN(std::vector<float>, "ConvWithDilation", input);
+
+   // Checking output size
+   EXPECT_EQ(output.size(), std::size(ConvWithDilation_ExpectedOutput::all_ones));
+
+   float *correct = ConvWithDilation_ExpectedOutput::all_ones;
+
+   // Checking every output value, one by one
+   for (size_t i = 0; i < output.size(); ++i) {
+      EXPECT_LE(std::abs(output[i] - correct[i]), TOLERANCE);
+   }
+}
+
+
 TEST(ONNX, ConvWithStridesNoPadding)
 {
    constexpr float TOLERANCE = DEFAULT_TOLERANCE;
diff --git a/tmva/sofie/test/input_models/ConvWithDilation.onnx b/tmva/sofie/test/input_models/ConvWithDilation.onnx
new file mode 100644
index 0000000000000000000000000000000000000000..c49707afa5be0b363b34f128e695f65d18bbfb2b
GIT binary patch
literal 269
zcmZus!AiqG6x_F)y6dCDx>V7VBA#;1O)$IEp2T~=i$s#PizNx!O+UcDNKgHs^q;!H
z9A$XCVVK9v6XNBU;vQ$VYC4aQsf)Vvfj&iVb+wNbUKVzjds|n|P(zQv<2lYYZQV6@
za0D9P;QXg-t8(|{ws}(;X4q)D#;C~)XUHE$Ng{W?wMFR+m~hLL95XWOx2q{5P6Pxn
z!6Ympdw)I+!|Q4INDqCV2A&c^^r64kNs{QG>&<$tKOd(0>J}FQlte0E8Xh+Pe#Ggk
Ne;h10*fK&Qug>tKKQ{mX

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..76559a18954bc
--- /dev/null
+++ b/tmva/sofie/test/input_models/references/ConvWithDilation.ref.hxx
@@ -0,0 +1,3 @@
+namespace ConvWithDilation_ExpectedOutput {
+float all_ones[] = {98.400002f, 102.900009f, 107.400002f, 129.899994f, 134.399994f, 138.899994f, 161.399994f, 165.899994f, 170.399994f};
+} // namespace ConvWithDilation_ExpectedOutput