diff --git a/onnxruntime/core/mlas/lib/convolve.cpp b/onnxruntime/core/mlas/lib/convolve.cpp index 696314f20267e..4378ec1948fdb 100644 --- a/onnxruntime/core/mlas/lib/convolve.cpp +++ b/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(MLAS_CONV_WORKING_BUFFER_SIZE_PER_THREAD)}); TargetThreadCount = MaximumThreadCount; if (static_cast(TargetThreadCount) >= Parameters->BatchCount * Parameters->GroupCount) { TargetThreadCount = static_cast(Parameters->BatchCount * Parameters->GroupCount); } - *WorkingBufferSize = TargetThreadCount * WorkingBufferSizePerThread; + *WorkingBufferSize = TargetThreadCount * MLAS_CONV_WORKING_BUFFER_SIZE_PER_THREAD; } } } diff --git a/onnxruntime/test/contrib_ops/fused_conv_test.cc b/onnxruntime/test/contrib_ops/fused_conv_test.cc index 7bfacb996526f..608ccadff8f1d 100644 --- a/onnxruntime/test/contrib_ops/fused_conv_test.cc +++ b/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(1)); + test.AddAttribute("kernel_shape", vector{7, 7, 7}); + test.AddAttribute("pads", vector{3, 3, 3, 3, 3, 3}); + test.AddAttribute("strides", vector{1, 1, 1}); + test.AddAttribute("dilations", vector{1, 1, 1}); + test.AddAttribute("activation", string("Relu")); + + const vector X_shape = {static_cast(batch_count), + static_cast(input_channels), + static_cast(input_depth), + static_cast(input_height), + static_cast(input_width)}; + const vector W_shape = {static_cast(filter_count), + static_cast(input_channels), + static_cast(kernel_depth), + static_cast(kernel_height), + static_cast(kernel_width)}; + const vector Y_shape = {static_cast(batch_count), + static_cast(filter_count), + static_cast(input_depth), + static_cast(input_height), + static_cast(input_width)}; + + vector X(batch_count * input_channels * input_depth * input_height * input_width, 1.0f); + vector 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(0, pad - pos); + const int64_t hi = std::min(kernel - 1, dim - 1 - pos + pad); + return hi - lo + 1; + }; + + vector 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(d), + static_cast(input_depth), + static_cast(kernel_depth)); + for (size_t h = 0; h < input_height; ++h) { + const int64_t ch = valid_count(static_cast(h), + static_cast(input_height), + static_cast(kernel_height)); + for (size_t w = 0; w < input_width; ++w) { + const int64_t cw = valid_count(static_cast(w), + static_cast(input_width), + static_cast(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(cd * ch * cw); + } + } + } + } + } + + test.AddInput("X", X_shape, X); + test.AddInput("W", W_shape, W, true); + test.AddOutput("Y", Y_shape, Y); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + #endif } // namespace test diff --git a/onnxruntime/test/mlas/unittest/test_conv2d.h b/onnxruntime/test/mlas/unittest/test_conv2d.h index 6ac47c69ae0b8..ef65da4adb031 100644 --- a/onnxruntime/test/mlas/unittest/test_conv2d.h +++ b/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(); } };