[MLAS] Fix scratch buffer over-allocation in batched/grouped NCHW Conv (microsoft#28482)

### Description

microsoft#26103 which was shipped
with ORT 1.23.2 introduced logic to parallelize on the batch and/or
group count parameters for NCHW Conv if either were > 1 to improve
performance of batched grouped convolutions - prior to that the
parallelization was only on the output spatial map. The per-thread
scratch buffer allocated by the B/G parallelization path introduced in
microsoft#26103 is excessive and is
not needed for the algorithm being used. Under the hood,
`MlasConvOperation()` uses a double tiled approach which means that the
per-thread scratch buffer needs just enough memory to fit that tile of
data and nothing more than that. This change fixes the conservative
over-allocation which shows up as a memory bloat in OP's model in issue
microsoft#28466

Note: the work partitioning across (batch × group) introduced in microsoft#26103
is unchanged — only the per-thread scratch buffer sizing is corrected.
The batched/grouped Conv speedup from microsoft#26103 is retained.

### Motivation and Context
Fix microsoft#28466

Author: hariharans29

onnxruntime/core/mlas/lib/convolve.cpp

@@ -849,7 +849,7 @@ Return Value:
         if (bias != nullptr) {
             bias += group * FilterCount;
         }
-        float* ColumnBuffer = WorkBlock->WorkingBuffer + Index * OutputSize * K;
+        float* ColumnBuffer = WorkBlock->WorkingBuffer + Index * MLAS_CONV_WORKING_BUFFER_SIZE_PER_THREAD;
 
         MlasConvOperation(Parameters, input, filter, bias, ColumnBuffer, output, 0, OutputSize);
     }
@@ -1712,14 +1712,11 @@ Return Value:
 
         if (Parameters->BatchCount > 1 || Parameters->GroupCount > 1) {
 
-            size_t WorkingBufferSizePerThread = std::max({Parameters->OutputSize * Parameters->K,
-                                                          Parameters->FilterCount * Parameters->OutputSize,
-                                                          static_cast<size_t>(MLAS_CONV_WORKING_BUFFER_SIZE_PER_THREAD)});
             TargetThreadCount = MaximumThreadCount;
             if (static_cast<size_t>(TargetThreadCount) >= Parameters->BatchCount * Parameters->GroupCount) {
                 TargetThreadCount = static_cast<ptrdiff_t>(Parameters->BatchCount * Parameters->GroupCount);
             }
-            *WorkingBufferSize = TargetThreadCount * WorkingBufferSizePerThread;
+            *WorkingBufferSize = TargetThreadCount * MLAS_CONV_WORKING_BUFFER_SIZE_PER_THREAD;
         }
     }
 }

onnxruntime/test/contrib_ops/fused_conv_test.cc

@@ -374,6 +374,90 @@ TEST(FusedConvTest, Cpu_NhwcConv2D_AutoPadSameUpper) {
 }
 #endif
 
+TEST(FusedConvTest, Cpu_Conv3D_Batched_Relu) {
+  constexpr size_t batch_count = 4;
+  constexpr size_t input_channels = 1;
+  constexpr size_t input_depth = 8;
+  constexpr size_t input_height = 8;
+  constexpr size_t input_width = 8;
+  constexpr size_t filter_count = 6;
+  constexpr size_t kernel_depth = 7;
+  constexpr size_t kernel_height = 7;
+  constexpr size_t kernel_width = 7;
+
+  OpTester test("FusedConv", 1, onnxruntime::kMSDomain);
+  test.AddAttribute("group", static_cast<int64_t>(1));
+  test.AddAttribute("kernel_shape", vector<int64_t>{7, 7, 7});
+  test.AddAttribute("pads", vector<int64_t>{3, 3, 3, 3, 3, 3});
+  test.AddAttribute("strides", vector<int64_t>{1, 1, 1});
+  test.AddAttribute("dilations", vector<int64_t>{1, 1, 1});
+  test.AddAttribute("activation", string("Relu"));
+
+  const vector<int64_t> X_shape = {static_cast<int64_t>(batch_count),
+                                   static_cast<int64_t>(input_channels),
+                                   static_cast<int64_t>(input_depth),
+                                   static_cast<int64_t>(input_height),
+                                   static_cast<int64_t>(input_width)};
+  const vector<int64_t> W_shape = {static_cast<int64_t>(filter_count),
+                                   static_cast<int64_t>(input_channels),
+                                   static_cast<int64_t>(kernel_depth),
+                                   static_cast<int64_t>(kernel_height),
+                                   static_cast<int64_t>(kernel_width)};
+  const vector<int64_t> Y_shape = {static_cast<int64_t>(batch_count),
+                                   static_cast<int64_t>(filter_count),
+                                   static_cast<int64_t>(input_depth),
+                                   static_cast<int64_t>(input_height),
+                                   static_cast<int64_t>(input_width)};
+
+  vector<float> X(batch_count * input_channels * input_depth * input_height * input_width, 1.0f);
+  vector<float> W(filter_count * input_channels * kernel_depth * kernel_height * kernel_width, 1.0f);
+
+  // With X = 1, W = 1, no bias, and a single input channel, the pre-activation output at
+  // [b][f][d][h][w] equals the number of valid kernel positions that fall inside the input
+  // volume at (d, h, w). For a kernel of size K with stride 1 and pad = K/2, the per-axis
+  // valid count at position p in a dimension of length L is:
+  //     count(p, L, K) = min(K - 1, L - 1 - p + K/2) - max(0, K/2 - p) + 1.
+  // The post-Relu output is the product of the per-axis counts (all values are positive).
+  auto valid_count = [](int64_t pos, int64_t dim, int64_t kernel) -> int64_t {
+    const int64_t pad = kernel / 2;
+    const int64_t lo = std::max<int64_t>(0, pad - pos);
+    const int64_t hi = std::min<int64_t>(kernel - 1, dim - 1 - pos + pad);
+    return hi - lo + 1;
+  };
+
+  vector<float> Y(batch_count * filter_count * input_depth * input_height * input_width);
+  for (size_t b = 0; b < batch_count; ++b) {
+    for (size_t f = 0; f < filter_count; ++f) {
+      for (size_t d = 0; d < input_depth; ++d) {
+        const int64_t cd = valid_count(static_cast<int64_t>(d),
+                                       static_cast<int64_t>(input_depth),
+                                       static_cast<int64_t>(kernel_depth));
+        for (size_t h = 0; h < input_height; ++h) {
+          const int64_t ch = valid_count(static_cast<int64_t>(h),
+                                         static_cast<int64_t>(input_height),
+                                         static_cast<int64_t>(kernel_height));
+          for (size_t w = 0; w < input_width; ++w) {
+            const int64_t cw = valid_count(static_cast<int64_t>(w),
+                                           static_cast<int64_t>(input_width),
+                                           static_cast<int64_t>(kernel_width));
+            const size_t idx = ((b * filter_count + f) * input_depth + d) * input_height * input_width +
+                               h * input_width + w;
+            Y[idx] = static_cast<float>(cd * ch * cw);
+          }
+        }
+      }
+    }
+  }
+
+  test.AddInput<float>("X", X_shape, X);
+  test.AddInput<float>("W", W_shape, W, true);
+  test.AddOutput<float>("Y", Y_shape, Y);
+
+  std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
+  execution_providers.push_back(DefaultCpuExecutionProvider());
+  test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
+}
+
 #endif
 
 }  // namespace test

onnxruntime/test/mlas/unittest/test_conv2d.h

@@ -492,6 +492,52 @@ class MlasConv2DTest : public MlasTestBase {
     }
   }
 
+  void TestBatchedConv3DWorkingBufferUsesThreadTileSize() {
+    constexpr size_t Dimensions = 3;
+    constexpr size_t BatchCount = 12;
+    constexpr size_t GroupCount = 1;
+    constexpr size_t InputChannels = 1;
+    constexpr size_t FilterCount = 6;
+    int64_t InputShape[] = {32, 64, 64};
+    int64_t KernelShape[] = {7, 7, 7};
+    int64_t DilationShape[] = {1, 1, 1};
+    int64_t Padding[] = {3, 3, 3, 3, 3, 3};
+    int64_t StrideShape[] = {1, 1, 1};
+    int64_t OutputShape[] = {32, 64, 64};
+
+    MLAS_ACTIVATION Activation;
+    Activation.ActivationKind = MlasIdentityActivation;
+
+    MLAS_CONV_PARAMETERS Parameters;
+    size_t WorkingBufferSize = 0;
+
+    MlasConvPrepare(&Parameters,
+                    Dimensions,
+                    BatchCount,
+                    GroupCount,
+                    InputChannels,
+                    InputShape,
+                    KernelShape,
+                    DilationShape,
+                    Padding,
+                    StrideShape,
+                    OutputShape,
+                    FilterCount,
+                    &Activation,
+                    &WorkingBufferSize,
+                    false,
+                    0.0f,
+                    threadpool_);
+
+    if (Parameters.Algorithm != MlasConvAlgorithmExpandThenGemmSegmented) {
+      GTEST_SKIP() << "This platform uses a different Conv3D algorithm.";
+    }
+
+    const size_t full_column_buffer_size = Parameters.OutputSize * Parameters.K;
+    ASSERT_LT(WorkingBufferSize, full_column_buffer_size)
+        << "Batched Conv3D should not allocate a full im2col buffer per worker.";
+  }
+
   void Test(
       size_t BatchCount,
       size_t GroupCount,
@@ -688,5 +734,6 @@ class MlasConv2DTest : public MlasTestBase {
     TestMobileClipBetaActivationRegression(64, 64, 64);
     TestMobileClipBetaActivationRegression(128, 32, 32);
     TestMobileClipBetaActivationRegression(256, 16, 16);
+    TestBatchedConv3DWorkingBufferUsesThreadTileSize();
   }
 };