From 61535b9606a400cc6d9d07aece16d8fa3e4324ed Mon Sep 17 00:00:00 2001 From: Orlaith Monahan Date: Tue, 19 May 2026 12:48:41 +0100 Subject: [PATCH 1/7] Update the MLAS NHWC Sans transposes feature to also support Depthwise * Allow for NHWC Depthwise convolutions when groups are values other than 1 * Added verification tests * Changed the fallback / skip tests to now check for asymettric padding, non-depthwise grouped conv, and multiplier > 1 Signed-off-by: Orlaith Monahan --- onnxruntime/core/mlas/lib/convolve.cpp | 9 +- .../mlas/lib/kleidiai/convolve_kleidiai.cpp | 42 +++++- .../core/optimizer/nhwc_transformer.cc | 14 +- .../test/contrib_ops/fused_conv_test.cc | 129 ++++++++++++++++++ .../test/optimizer/nhwc_transformer_test.cc | 98 ++++++++++++- 5 files changed, 279 insertions(+), 13 deletions(-) diff --git a/onnxruntime/core/mlas/lib/convolve.cpp b/onnxruntime/core/mlas/lib/convolve.cpp index 4378ec1948fdb..5dc5eb1281b12 100644 --- a/onnxruntime/core/mlas/lib/convolve.cpp +++ b/onnxruntime/core/mlas/lib/convolve.cpp @@ -1377,7 +1377,7 @@ MlasConvSupportsSymmetricChannelsLast2DFloatKernel( return false; } - if (Dimensions != 2 || BatchCount != 1 || GroupCount != 1 || Beta != 0.0f) { + if (Dimensions != 2 || BatchCount != 1 || Beta != 0.0f) { return false; } @@ -1395,7 +1395,12 @@ MlasConvSupportsSymmetricChannelsLast2DFloatKernel( return false; } - if (FilterCount <= 1 || KernelShape[0] < 3 || KernelShape[1] < 3) { + const bool is_depthwise = GroupCount > 1 && FilterCount == 1; + if (GroupCount > 1 && !is_depthwise) { + return false; + } + + if (!is_depthwise && (FilterCount <= 1 || KernelShape[0] < 3 || KernelShape[1] < 3)) { return false; } diff --git a/onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp b/onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp index cca4f5a19c417..a5bcaf18e94ae 100644 --- a/onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp +++ b/onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp @@ -550,15 +550,38 @@ static void ConvolveSme(const size_t co, //channels out dim[1] = MlasDivRoundup(m, m_step); dim[2] = MlasDivRoundup(co, n_step); + const bool grouped_channels_last = input_is_channels_last && groups > 1; + const size_t input_channels_total = ci * groups; + const size_t output_channels_total = co * groups; + const float* input_base = in; + float* output_base = out; + for (size_t g = 0; g < groups; ++g) { + const float* input_group = in; + std::vector input_group_buffer; + if (grouped_channels_last) { + input_group_buffer.resize(ih * iw * ci); + for (size_t pixel = 0; pixel < ih * iw; ++pixel) { + const float* src = input_base + pixel * input_channels_total + g * ci; + std::copy_n(src, ci, input_group_buffer.data() + pixel * ci); + } + input_group = input_group_buffer.data(); + } auto result = out; const bool need_transpose = (!input_is_channels_last) && (co > 1); + const bool use_temp_output = grouped_channels_last || need_transpose; if (need_transpose) { result = tmp_mlas_aligned; } + if (grouped_channels_last) { + result = tmp_mlas_aligned; + } - auto lhs = LhsPackImageDataSme(ci, ih, iw, d_kh, d_kw, sh, sw, padding, in, input_is_channels_last, ThreadPool); + auto lhs = LhsPackImageDataSme(ci, ih, iw, d_kh, d_kw, sh, sw, padding, + input_group, + input_is_channels_last, + ThreadPool); const std::byte* rhs_data = packed_rhs ? packed_rhs + g * packed_rhs_group_stride : nullptr; std::unique_ptr rhs_storage; if (rhs_data == nullptr) { @@ -613,13 +636,26 @@ static void ConvolveSme(const size_t co, //channels out ); }); + if (grouped_channels_last) { + for (size_t pixel = 0; pixel < m; ++pixel) { + float* dst = output_base + pixel * output_channels_total + g * co; + const float* src = result + pixel * co; + std::copy_n(src, co, dst); + } + } + if (need_transpose) { //Note: this could be absorbed into post conv activation MlasTranspose(tmp_mlas_aligned, out, m, co, ThreadPool); } - in += ci * ih * iw; - out += m * co; + if (!grouped_channels_last) { + in += ci * ih * iw; + out += use_temp_output ? 0 : m * co; + if (need_transpose) { + out += m * co; + } + } weights += co * ci * kh * kw; if(bias){ bias += co; diff --git a/onnxruntime/core/optimizer/nhwc_transformer.cc b/onnxruntime/core/optimizer/nhwc_transformer.cc index 6c0717865b135..004a756ef08a5 100644 --- a/onnxruntime/core/optimizer/nhwc_transformer.cc +++ b/onnxruntime/core/optimizer/nhwc_transformer.cc @@ -191,9 +191,10 @@ bool FloatNhwcWrapperFilter(const onnx_transpose_optimization::api::GraphRef& gr } const auto group = node.GetAttributeInt("group").value_or(1); - if (group != 1) { + if (group <= 0) { return false; } + const auto group_count = narrow(group); std::array input_spatial_shape{}; std::array kernel_spatial_shape{}; @@ -201,17 +202,22 @@ bool FloatNhwcWrapperFilter(const onnx_transpose_optimization::api::GraphRef& gr std::array strides{1, 1}; std::array pads{}; size_t batch_count = 0; - size_t filter_count = 0; + size_t total_filter_count = 0; if (!TryGetDimValueAsSizeT(*input_shape, 0, batch_count) || !TryGetDimValueAsSizeT(*input_shape, 2, input_spatial_shape[0]) || !TryGetDimValueAsSizeT(*input_shape, 3, input_spatial_shape[1]) || - !TryGetDimValueAsSizeT(*weight_shape, 0, filter_count) || + !TryGetDimValueAsSizeT(*weight_shape, 0, total_filter_count) || !TryGetDimValueAsSizeT(*weight_shape, 2, kernel_spatial_shape[0]) || !TryGetDimValueAsSizeT(*weight_shape, 3, kernel_spatial_shape[1])) { return false; } + if (total_filter_count == 0 || total_filter_count % group_count != 0) { + return false; + } + const size_t filter_count = total_filter_count / group_count; + const auto dilations_opt = node.GetAttributeInts("dilations"); if (dilations_opt.has_value() && !TryReadPositiveInts(*dilations_opt, dilations)) { return false; @@ -229,7 +235,7 @@ bool FloatNhwcWrapperFilter(const onnx_transpose_optimization::api::GraphRef& gr return MlasConvSupportsSymmetricChannelsLast2DFloatKernel( /*Dimensions*/ 2, batch_count, - /*GroupCount*/ 1, + group_count, input_spatial_shape.data(), kernel_spatial_shape.data(), dilations.data(), diff --git a/onnxruntime/test/contrib_ops/fused_conv_test.cc b/onnxruntime/test/contrib_ops/fused_conv_test.cc index 608ccadff8f1d..f453cb3bca746 100644 --- a/onnxruntime/test/contrib_ops/fused_conv_test.cc +++ b/onnxruntime/test/contrib_ops/fused_conv_test.cc @@ -3,11 +3,19 @@ #include "gtest/gtest.h" +#include + +#include "core/common/narrow.h" +#include "core/framework/kernel_registry.h" #include "test/common/cuda_op_test_utils.h" #include "test/common/tensor_op_test_utils.h" #include "test/providers/provider_test_utils.h" #include "test/util/include/default_providers.h" +#if defined(USE_KLEIDIAI) && defined(MLAS_TARGET_ARM64) +#include "core/mlas/lib/mlasi.h" +#endif + namespace onnxruntime { namespace test { @@ -121,6 +129,88 @@ void RunConvOp(const ConvOpAndTestAttributes& attributes, } #ifdef USE_KLEIDIAI +namespace { + +#if defined(MLAS_TARGET_ARM64) +bool HasFloatNhwcFusedConvKernel() { + auto cpu_ep = DefaultCpuExecutionProvider(); + if (cpu_ep == nullptr) { + return false; + } + + auto kernel_registry = cpu_ep->GetKernelRegistry(); + if (!kernel_registry) { + return false; + } + + KernelRegistry::TypeConstraintMap type_constraints{ + {"T", DataTypeImpl::GetTensorType()}, + }; + + const KernelCreateInfo* kernel_create_info{}; + const auto status = kernel_registry->TryFindKernel( + kCpuExecutionProvider, + "NhwcFusedConv", + kMSDomain, + 1, + type_constraints, + DefaultLoggingManager().DefaultLogger(), + &kernel_create_info); + + return status.IsOK() && kernel_create_info != nullptr; +} + +bool HasFloatNhwcNoTransposeSupport(const vector& input_shape, + const vector& weight_shape, + const vector& pads, + const vector& strides, + int64_t group) { + if (!HasFloatNhwcFusedConvKernel() || !MLAS_CPUIDINFO::GetCPUIDInfo().HasArm_SME()) { + return false; + } + + if (group <= 0 || input_shape.size() != 4 || weight_shape.size() != 4 || + pads.size() != 4 || strides.size() != 2 || + weight_shape[0] <= 0 || weight_shape[0] % group != 0) { + return false; + } + + std::array input_spatial_shape{ + narrow(input_shape[1]), + narrow(input_shape[2]), + }; + std::array kernel_spatial_shape{ + narrow(weight_shape[2]), + narrow(weight_shape[3]), + }; + std::array dilations{1, 1}; + std::array strides_size_t{ + narrow(strides[0]), + narrow(strides[1]), + }; + std::array pads_size_t{ + narrow(pads[0]), + narrow(pads[1]), + narrow(pads[2]), + narrow(pads[3]), + }; + + return MlasConvSupportsSymmetricChannelsLast2DFloatKernel( + /*Dimensions*/ 2, + narrow(input_shape[0]), + narrow(group), + input_spatial_shape.data(), + kernel_spatial_shape.data(), + dilations.data(), + pads_size_t.data(), + strides_size_t.data(), + narrow(weight_shape[0] / group), + /*Beta*/ 0.0f); +} +#endif + +} // namespace + void TestNhwcFusedConvFloatOp(const ConvOpAndTestAttributes& attributes, const vector>& inputs, const vector>& input_shapes, @@ -372,6 +462,45 @@ TEST(FusedConvTest, Cpu_NhwcConv2D_AutoPadSameUpper) { TestNhwcFusedConvFloatOp(attrs, {X, W}, {X_shape, W_shape}, expected_vals, Y_shape); TestNhwcFusedConvFloatOp(attrs, {X, W}, {X_shape, W_shape}, expected_vals, Y_shape, true); } + +TEST(FusedConvTest, Cpu_NhwcDepthwiseConv2D_SymmetricPadding) { +#if !defined(MLAS_TARGET_ARM64) + GTEST_SKIP() << "Float NHWC depthwise fast-path requires Arm64."; +#else + ConvOpAndTestAttributes attrs = { + "", // auto_pad + vector{1, 1}, // dilations + 2, // group + vector{3, 3}, // kernel_shape + vector{1, 1, 1, 1}, // pads + vector{1, 1}, // strides + "Relu" // activation + }; + + vector X_shape = {1, 3, 3, 2}; + vector X = {1.0f, 10.0f, 2.0f, 20.0f, 3.0f, 30.0f, + 4.0f, 40.0f, 5.0f, 50.0f, 6.0f, 60.0f, + 7.0f, 70.0f, 8.0f, 80.0f, 9.0f, 90.0f}; + vector W_shape = {2, 1, 3, 3}; + vector W = {1.0f, 1.0f, 1.0f, + 1.0f, 1.0f, 1.0f, + 1.0f, 1.0f, 1.0f, + 1.0f, 0.0f, 1.0f, + 0.0f, 2.0f, 0.0f, + 1.0f, 0.0f, 1.0f}; + vector Y_shape = {1, 3, 3, 2}; + auto expected_vals = {12.0f, 70.0f, 21.0f, 140.0f, 16.0f, 110.0f, + 27.0f, 180.0f, 45.0f, 300.0f, 33.0f, 220.0f, + 24.0f, 190.0f, 39.0f, 260.0f, 28.0f, 230.0f}; + + if (!HasFloatNhwcNoTransposeSupport(X_shape, W_shape, attrs.pads, attrs.strides, attrs.group)) { + GTEST_SKIP() << "Float NHWC depthwise fast-path is not available on this configuration."; + } + + TestNhwcFusedConvFloatOp(attrs, {X, W}, {X_shape, W_shape}, expected_vals, Y_shape); + TestNhwcFusedConvFloatOp(attrs, {X, W}, {X_shape, W_shape}, expected_vals, Y_shape, true); +#endif +} #endif TEST(FusedConvTest, Cpu_Conv3D_Batched_Relu) { diff --git a/onnxruntime/test/optimizer/nhwc_transformer_test.cc b/onnxruntime/test/optimizer/nhwc_transformer_test.cc index b73929efab8a6..ff5a93360a8a9 100644 --- a/onnxruntime/test/optimizer/nhwc_transformer_test.cc +++ b/onnxruntime/test/optimizer/nhwc_transformer_test.cc @@ -81,9 +81,15 @@ static bool HasFloatNhwcNoTransposeSupport(const std::vector& input_sha return false; } - if (has_sum_input || group != 1 || input_shape.size() != 4 || weight_shape.size() != 4) { + if (has_sum_input || group <= 0 || input_shape.size() != 4 || weight_shape.size() != 4) { return false; } + const auto group_count = narrow(group); + + if (weight_shape[0] <= 0 || weight_shape[0] % group != 0) { + return false; + } + const auto filter_count = narrow(weight_shape[0] / group); std::array input_spatial_shape{ narrow(input_shape[2]), @@ -169,13 +175,13 @@ static bool HasFloatNhwcNoTransposeSupport(const std::vector& input_sha return MlasConvSupportsSymmetricChannelsLast2DFloatKernel( /*Dimensions*/ 2, narrow(input_shape[0]), - /*GroupCount*/ 1, + group_count, input_spatial_shape.data(), kernel_spatial_shape.data(), dilations_size_t.data(), pads_size_t.data(), strides_size_t.data(), - narrow(weight_shape[0]), + filter_count, /*Beta*/ 0.0f); #else ORT_UNUSED_PARAMETER(input_shape); @@ -407,7 +413,7 @@ TEST(NhwcTransformerTests, ConvGlobalAveragePool) { TransformerLevel::Level3); } -TEST(NhwcTransformerTests, ConvDepthwiseFloat_SkipNhwc) { +TEST(NhwcTransformerTests, ConvDepthwiseFloat_UsesHelperCapability) { auto build_test_case = [&](ModelTestBuilder& builder) { auto* input_arg = builder.MakeInput({1, 8, 7, 7}, -1.0f, 1.0f); auto* weight_arg = builder.MakeInitializer({8, 1, 3, 3}, -1.0f, 1.0f); @@ -435,6 +441,90 @@ TEST(NhwcTransformerTests, ConvDepthwiseFloat_SkipNhwc) { /*relative_per_sample_tolerance*/ 1e-6); } +TEST(NhwcTransformerTests, ConvDepthwiseFloat_AsymmetricPaddingSkipsNhwc) { + auto build_test_case = [&](ModelTestBuilder& builder) { + auto* input_arg = builder.MakeInput({1, 8, 7, 7}, -1.0f, 1.0f); + auto* weight_arg = builder.MakeInitializer({8, 1, 3, 3}, -1.0f, 1.0f); + auto* output_arg = builder.MakeOutput(); + + Node& conv_node = builder.AddConvNode(input_arg, weight_arg, output_arg); + conv_node.AddAttribute("group", static_cast(8)); + conv_node.AddAttribute("pads", std::vector{0, 1, 1, 1}); + }; + + auto check_nhwc_graph = [&](InferenceSessionWrapper& session) { + auto op_to_count = CountOpsInGraph(session.GetGraph()); + const bool expect_nhwc = HasFloatNhwcNoTransposeSupport({1, 8, 7, 7}, {8, 1, 3, 3}, {0, 1, 1, 1}, {}, {}, 8); + EXPECT_FALSE(expect_nhwc); + EXPECT_EQ(op_to_count["com.microsoft.NhwcFusedConv"], 0); + EXPECT_EQ(op_to_count["Transpose"], 0); + }; + + TransformerTester(build_test_case, + check_nhwc_graph, + TransformerLevel::Level2, + TransformerLevel::Level3, + /*opset_version*/ 12, + /*per_sample_tolerance*/ 1e-6, + /*relative_per_sample_tolerance*/ 1e-6); +} + +TEST(NhwcTransformerTests, ConvGroupedFloat_SkipNhwc) { + auto build_test_case = [&](ModelTestBuilder& builder) { + auto* input_arg = builder.MakeInput({1, 8, 7, 7}, -1.0f, 1.0f); + auto* weight_arg = builder.MakeInitializer({8, 2, 3, 3}, -1.0f, 1.0f); + auto* output_arg = builder.MakeOutput(); + + Node& conv_node = builder.AddConvNode(input_arg, weight_arg, output_arg); + conv_node.AddAttribute("group", static_cast(4)); + conv_node.AddAttribute("pads", std::vector{1, 1, 1, 1}); + }; + + auto check_nhwc_graph = [&](InferenceSessionWrapper& session) { + auto op_to_count = CountOpsInGraph(session.GetGraph()); + const bool expect_nhwc = HasFloatNhwcNoTransposeSupport({1, 8, 7, 7}, {8, 2, 3, 3}, {1, 1, 1, 1}, {}, {}, 4); + EXPECT_FALSE(expect_nhwc); + EXPECT_EQ(op_to_count["com.microsoft.NhwcFusedConv"], 0); + EXPECT_EQ(op_to_count["Transpose"], 0); + }; + + TransformerTester(build_test_case, + check_nhwc_graph, + TransformerLevel::Level2, + TransformerLevel::Level3, + /*opset_version*/ 12, + /*per_sample_tolerance*/ 1e-6, + /*relative_per_sample_tolerance*/ 1e-6); +} + +TEST(NhwcTransformerTests, ConvDepthwiseMultiplier2Float_SkipNhwc) { + auto build_test_case = [&](ModelTestBuilder& builder) { + auto* input_arg = builder.MakeInput({1, 8, 7, 7}, -1.0f, 1.0f); + auto* weight_arg = builder.MakeInitializer({16, 1, 3, 3}, -1.0f, 1.0f); + auto* output_arg = builder.MakeOutput(); + + Node& conv_node = builder.AddConvNode(input_arg, weight_arg, output_arg); + conv_node.AddAttribute("group", static_cast(8)); + conv_node.AddAttribute("pads", std::vector{1, 1, 1, 1}); + }; + + auto check_nhwc_graph = [&](InferenceSessionWrapper& session) { + auto op_to_count = CountOpsInGraph(session.GetGraph()); + const bool expect_nhwc = HasFloatNhwcNoTransposeSupport({1, 8, 7, 7}, {16, 1, 3, 3}, {1, 1, 1, 1}, {}, {}, 8); + EXPECT_FALSE(expect_nhwc); + EXPECT_EQ(op_to_count["com.microsoft.NhwcFusedConv"], 0); + EXPECT_EQ(op_to_count["Transpose"], 0); + }; + + TransformerTester(build_test_case, + check_nhwc_graph, + TransformerLevel::Level2, + TransformerLevel::Level3, + /*opset_version*/ 12, + /*per_sample_tolerance*/ 1e-6, + /*relative_per_sample_tolerance*/ 1e-6); +} + TEST(NhwcTransformerTests, ConvFloat_UsesNhwcOnlyWithKleidi) { auto build_test_case = [&](ModelTestBuilder& builder) { auto* input_arg = builder.MakeInput({1, 8, 7, 7}, -1.0f, 1.0f); From 6a897220e67c0186813a84f5579c0be6b5b22110 Mon Sep 17 00:00:00 2001 From: Orlaith Monahan Date: Tue, 19 May 2026 14:06:55 +0100 Subject: [PATCH 2/7] Adding benchmark tests removed unnecessary linux ifdefs Signed-off-by: Orlaith Monahan --- onnxruntime/core/providers/cpu/nn/conv.cc | 8 +- onnxruntime/core/providers/cpu/nn/conv.h | 4 +- onnxruntime/test/mlas/bench/bench_sconv.cpp | 154 ++++++++++++++++++++ 3 files changed, 160 insertions(+), 6 deletions(-) diff --git a/onnxruntime/core/providers/cpu/nn/conv.cc b/onnxruntime/core/providers/cpu/nn/conv.cc index 87ce1b05caae2..eff72e4498b3f 100644 --- a/onnxruntime/core/providers/cpu/nn/conv.cc +++ b/onnxruntime/core/providers/cpu/nn/conv.cc @@ -23,7 +23,7 @@ #include "core/common/safeint.h" #include "core/util/math_cpuonly.h" -#if defined(USE_KLEIDIAI) && defined(__aarch64__) && defined(__linux__) +#if defined(USE_KLEIDIAI) && defined(MLAS_TARGET_ARM64) #include "core/mlas/lib/kleidiai/mlasi_kleidiai.h" #endif @@ -191,7 +191,7 @@ Status Conv::Compute(OpKernelContext* context) const { return Status::OK(); } -#if defined(USE_KLEIDIAI) && defined(__aarch64__) && defined(__linux__) +#if defined(USE_KLEIDIAI) && defined(MLAS_TARGET_ARM64) Status Conv::EnsurePackedChannelsLastFilter(concurrency::ThreadPool* thread_pool, size_t filter_count_per_group, size_t input_channels_per_group, @@ -329,7 +329,7 @@ Status Conv::Compute(OpKernelContext* context) const { narrow(M / conv_attrs_.group), /*Beta*/ 0.0f); -#if defined(USE_KLEIDIAI) && defined(__aarch64__) && defined(__linux__) +#if defined(USE_KLEIDIAI) && defined(MLAS_TARGET_ARM64) if (nhwc_fastpath && can_cache_packed_filter_) { ORT_RETURN_IF_ERROR(EnsurePackedChannelsLastFilter(thread_pool, narrow(M / conv_attrs_.group), @@ -385,7 +385,7 @@ Status Conv::Compute(OpKernelContext* context) const { nhwc_fastpath ? 0.0f : Beta, thread_pool); -#if defined(USE_KLEIDIAI) && defined(__aarch64__) && defined(__linux__) +#if defined(USE_KLEIDIAI) && defined(MLAS_TARGET_ARM64) if (nhwc_fastpath && packed_filter_ != nullptr) { Parameters.FilterIsPacked = true; Parameters.PackedFilter = packed_filter_.get(); diff --git a/onnxruntime/core/providers/cpu/nn/conv.h b/onnxruntime/core/providers/cpu/nn/conv.h index 1cbe417cdbd96..9e073df545328 100644 --- a/onnxruntime/core/providers/cpu/nn/conv.h +++ b/onnxruntime/core/providers/cpu/nn/conv.h @@ -31,7 +31,7 @@ class Conv : public OpKernel { activation_.ActivationKind = MlasIdentityActivation; SetupMlasBackendKernelSelectorFromConfigOptions(mlas_backend_kernel_selector_config_, info.GetConfigOptions()); -#if defined(USE_KLEIDIAI) && defined(__aarch64__) && defined(__linux__) +#if defined(USE_KLEIDIAI) && defined(MLAS_TARGET_ARM64) if (channels_last_) { const auto& input_defs = info.node().InputDefs(); const bool has_bias_input = input_defs.size() >= 3 && input_defs[2] != nullptr; @@ -56,7 +56,7 @@ class Conv : public OpKernel { ConvAttributes conv_attrs_; bool channels_last_{false}; -#if defined(USE_KLEIDIAI) && defined(__aarch64__) && defined(__linux__) +#if defined(USE_KLEIDIAI) && defined(MLAS_TARGET_ARM64) private: Status EnsurePackedChannelsLastFilter(concurrency::ThreadPool* thread_pool, size_t filter_count_per_group, diff --git a/onnxruntime/test/mlas/bench/bench_sconv.cpp b/onnxruntime/test/mlas/bench/bench_sconv.cpp index 9df09728ffa17..59fca6d152699 100644 --- a/onnxruntime/test/mlas/bench/bench_sconv.cpp +++ b/onnxruntime/test/mlas/bench/bench_sconv.cpp @@ -146,6 +146,145 @@ static MLAS_THREADPOOL* GetMlasThreadPoolForConvBenchmark(void) { return threadpool.get(); } +void SCONV_NHWC_KLEIDIAI(benchmark::State& state, const char* /*dummy*/) { + const int64_t rank = state.range(0); // Rank + const int64_t batch_size = state.range(1); // N + const int64_t groups = state.range(2); // G + const int64_t input_channels_per_group = state.range(3); // Cpg + const int64_t output_channels_per_group = state.range(4); // Fpg + + if (rank <= 0) throw std::invalid_argument("Kernel rank must greater than 0!"); + if (batch_size <= 0) throw std::invalid_argument("Batch size must greater than 0!"); + if (groups <= 0) throw std::invalid_argument("Group count must greater than 0!"); + if (input_channels_per_group <= 0) throw std::invalid_argument("input_channels_per_group must greater than 0!"); + if (output_channels_per_group <= 0) throw std::invalid_argument("output_channels_per_group must greater than 0!"); + + size_t arg_position = 5; + const auto input_shape = BenchArgsVector(state, arg_position, rank); + const auto kernel_shape = BenchArgsVector(state, arg_position, rank); + const auto paddings = BenchArgsVector(state, arg_position, rank * 2); + const auto strides = BenchArgsVector(state, arg_position, rank); + const auto dilations = BenchArgsVector(state, arg_position, rank); + + if (std::any_of(input_shape.begin(), input_shape.end(), [](const int64_t& dim) { return dim <= 0; })) { + throw std::invalid_argument("all input image dim must > 0"); + } + + if (std::any_of(kernel_shape.begin(), kernel_shape.end(), [](const int64_t& dim) { return dim <= 0; })) { + throw std::invalid_argument("all kernel dim must > 0"); + } + + if (std::any_of(strides.begin(), strides.end(), [](const int64_t& dim) { return dim <= 0; })) { + throw std::invalid_argument("all strides dim must > 0"); + } + + if (std::any_of(dilations.begin(), dilations.end(), [](const int64_t& dim) { return dim <= 0; })) { + throw std::invalid_argument("all dilations dim must > 0"); + } + + if (rank != 2 || batch_size != 1) { + state.SkipWithError("KleidiAI NHWC benchmark requires 2D convolution with batch size 1."); + return; + } + + std::vector input_shape_size_t(static_cast(rank)); + std::vector kernel_shape_size_t(static_cast(rank)); + std::vector paddings_size_t(static_cast(rank * 2)); + std::vector strides_size_t(static_cast(rank)); + std::vector dilations_size_t(static_cast(rank)); + for (int64_t i = 0; i < rank; ++i) { + input_shape_size_t[static_cast(i)] = static_cast(input_shape[static_cast(i)]); + kernel_shape_size_t[static_cast(i)] = static_cast(kernel_shape[static_cast(i)]); + strides_size_t[static_cast(i)] = static_cast(strides[static_cast(i)]); + dilations_size_t[static_cast(i)] = static_cast(dilations[static_cast(i)]); + paddings_size_t[static_cast(i)] = static_cast(paddings[static_cast(i)]); + paddings_size_t[static_cast(i + rank)] = static_cast(paddings[static_cast(i + rank)]); + } + + if (!MlasConvSupportsSymmetricChannelsLast2DFloatKernel( + static_cast(rank), + static_cast(batch_size), + static_cast(groups), + input_shape_size_t.data(), + kernel_shape_size_t.data(), + dilations_size_t.data(), + paddings_size_t.data(), + strides_size_t.data(), + static_cast(output_channels_per_group), + 0.0f)) { + state.SkipWithError("KleidiAI NHWC kernel is not supported for this benchmark shape on the current platform."); + return; + } + + const int64_t GC = groups * input_channels_per_group; + const int64_t GF = groups * output_channels_per_group; + std::vector x_shape = {batch_size}; + x_shape.insert(x_shape.end(), input_shape.begin(), input_shape.end()); + x_shape.push_back(GC); + + std::vector f_shape = {GF, input_channels_per_group}; + f_shape.insert(f_shape.end(), kernel_shape.begin(), kernel_shape.end()); + + std::vector output_shape(static_cast(rank)); + for (int64_t i = 0; i < rank; ++i) { + auto km = 1 + dilations[static_cast(i)] * (kernel_shape[static_cast(i)] - 1); + output_shape[static_cast(i)] = + (paddings[static_cast(i)] + paddings[static_cast(i + rank)] + input_shape[static_cast(i)] - km) / + strides[static_cast(i)] + + 1; + } + + std::vector y_shape = {batch_size}; + y_shape.insert(y_shape.end(), output_shape.begin(), output_shape.end()); + y_shape.push_back(GF); + + MLAS_ACTIVATION activation; + activation.ActivationKind = MlasIdentityActivation; + MLAS_CONV_PARAMETERS Parameters; + size_t WorkingBufferSize = 0; + MlasConvPrepare(&Parameters, + static_cast(rank), + static_cast(batch_size), + static_cast(groups), + static_cast(input_channels_per_group), + input_shape.data(), + kernel_shape.data(), + dilations.data(), + paddings.data(), + strides.data(), + output_shape.data(), + static_cast(output_channels_per_group), + &activation, + &WorkingBufferSize, + true, + 0.0f, + nullptr); + + auto X = RandomVectorUniform(x_shape, -2.0, 2.0); + auto F = RandomVectorUniform(f_shape, -1.0, 1.0); + int64_t y_size = std::accumulate(y_shape.begin(), y_shape.end(), 1LL, std::multiplies()); + std::vector Y(static_cast(y_size)); + std::vector working_buffer(WorkingBufferSize); + + MlasConv(&Parameters, + X.data(), + F.data(), + nullptr, + working_buffer.data(), + Y.data(), + nullptr); + + for (auto _ : state) { + MlasConv(&Parameters, + X.data(), + F.data(), + nullptr, + working_buffer.data(), + Y.data(), + nullptr); + } +} + void SCONV_NCHW_THREADED(benchmark::State& state, const char* /*dummy*/) { MLAS_THREADPOOL* tp = GetMlasThreadPoolForConvBenchmark(); @@ -354,6 +493,21 @@ static void MobileClip(benchmark::internal::Benchmark* b) { BENCHMARK_CAPTURE(SCONV_NCHW, MobileClip, "")->Apply(MobileClip)->UseRealTime(); BENCHMARK_CAPTURE(SCONV_NCHW_THREADED, MobileClip, "")->Apply(MobileClip)->UseRealTime(); +static void KleidiAiNhwcComparison(benchmark::internal::Benchmark* b) { + b->ArgNames(ArgNamesForConv(2)); + + // Dense 3x3 conv shapes that fit the Arm SME / KleidiAI NHWC fast-path envelope. + b->Args({2, 1, 1, 64, 64, 56, 56, 3, 3, 1, 1, 1, 1, 1, 1, 1, 1}); + b->Args({2, 1, 1, 128, 128, 28, 28, 3, 3, 1, 1, 1, 1, 1, 1, 1, 1}); + + // Classic depthwise shapes now supported by the NHWC helper gate. + b->Args({2, 1, 64, 1, 1, 56, 56, 3, 3, 1, 1, 1, 1, 1, 1, 1, 1}); + b->Args({2, 1, 72, 1, 1, 48, 80, 3, 3, 1, 1, 1, 1, 2, 2, 1, 1}); +} + +BENCHMARK_CAPTURE(SCONV_NCHW, KleidiAiNhwcComparison_NchwBaseline, "")->Apply(KleidiAiNhwcComparison)->UseRealTime(); +BENCHMARK_CAPTURE(SCONV_NHWC_KLEIDIAI, KleidiAiNhwcComparison_NhwcFastPath, "")->Apply(KleidiAiNhwcComparison)->UseRealTime(); + static void General_Conv2d(benchmark::internal::Benchmark* b) { b->ArgNames(ArgNamesForConv(2)); b->ArgsProduct( From ef85de289fdb112a7c0e31c2f6f0829a1440ee6c Mon Sep 17 00:00:00 2001 From: Orlaith Monahan Date: Thu, 21 May 2026 15:13:16 +0100 Subject: [PATCH 3/7] Add a limit to FloatNhwcWrapperFilter to avoid narrow oversize errors Signed-off-by: Orlaith Monahan --- onnxruntime/core/optimizer/nhwc_transformer.cc | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/onnxruntime/core/optimizer/nhwc_transformer.cc b/onnxruntime/core/optimizer/nhwc_transformer.cc index 004a756ef08a5..2737dc9a02c0b 100644 --- a/onnxruntime/core/optimizer/nhwc_transformer.cc +++ b/onnxruntime/core/optimizer/nhwc_transformer.cc @@ -5,6 +5,7 @@ #include #include #include +#include #include #include "core/common/cpuid_info.h" #include "core/graph/constants.h" @@ -194,6 +195,10 @@ bool FloatNhwcWrapperFilter(const onnx_transpose_optimization::api::GraphRef& gr if (group <= 0) { return false; } + constexpr uint64_t kSizeTMax = static_cast(std::numeric_limits::max()); + if (static_cast(group) > kSizeTMax) { + return false; + } const auto group_count = narrow(group); std::array input_spatial_shape{}; From 2c43663d6c1dcf824536114c51d7a9b42b043296 Mon Sep 17 00:00:00 2001 From: Orlaith Monahan Date: Thu, 21 May 2026 15:18:50 +0100 Subject: [PATCH 4/7] Move the input_group_buffer allocation outside of the loop in ConvolveSme and only size it once * Fixed some grammer in throw statements Signed-off-by: Orlaith Monahan --- .../core/mlas/lib/kleidiai/convolve_kleidiai.cpp | 6 ++++-- onnxruntime/test/mlas/bench/bench_sconv.cpp | 10 +++++----- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp b/onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp index a5bcaf18e94ae..1da8530350d75 100644 --- a/onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp +++ b/onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp @@ -555,12 +555,14 @@ static void ConvolveSme(const size_t co, //channels out const size_t output_channels_total = co * groups; const float* input_base = in; float* output_base = out; + std::vector input_group_buffer; + if (grouped_channels_last) { + input_group_buffer.resize(ih * iw * ci); + } for (size_t g = 0; g < groups; ++g) { const float* input_group = in; - std::vector input_group_buffer; if (grouped_channels_last) { - input_group_buffer.resize(ih * iw * ci); for (size_t pixel = 0; pixel < ih * iw; ++pixel) { const float* src = input_base + pixel * input_channels_total + g * ci; std::copy_n(src, ci, input_group_buffer.data() + pixel * ci); diff --git a/onnxruntime/test/mlas/bench/bench_sconv.cpp b/onnxruntime/test/mlas/bench/bench_sconv.cpp index 59fca6d152699..85cb68076b638 100644 --- a/onnxruntime/test/mlas/bench/bench_sconv.cpp +++ b/onnxruntime/test/mlas/bench/bench_sconv.cpp @@ -153,11 +153,11 @@ void SCONV_NHWC_KLEIDIAI(benchmark::State& state, const char* /*dummy*/) { const int64_t input_channels_per_group = state.range(3); // Cpg const int64_t output_channels_per_group = state.range(4); // Fpg - if (rank <= 0) throw std::invalid_argument("Kernel rank must greater than 0!"); - if (batch_size <= 0) throw std::invalid_argument("Batch size must greater than 0!"); - if (groups <= 0) throw std::invalid_argument("Group count must greater than 0!"); - if (input_channels_per_group <= 0) throw std::invalid_argument("input_channels_per_group must greater than 0!"); - if (output_channels_per_group <= 0) throw std::invalid_argument("output_channels_per_group must greater than 0!"); + if (rank <= 0) throw std::invalid_argument("Kernel rank must be greater than 0"); + if (batch_size <= 0) throw std::invalid_argument("Batch size must be greater than 0"); + if (groups <= 0) throw std::invalid_argument("Group count must be greater than 0"); + if (input_channels_per_group <= 0) throw std::invalid_argument("input_channels_per_group must be greater than 0"); + if (output_channels_per_group <= 0) throw std::invalid_argument("output_channels_per_group must be greater than 0"); size_t arg_position = 5; const auto input_shape = BenchArgsVector(state, arg_position, rank); From 745c64ad0d9089558f825f7b2380cae7f12047bd Mon Sep 17 00:00:00 2001 From: Orlaith Monahan Date: Tue, 2 Jun 2026 16:37:09 +0100 Subject: [PATCH 5/7] Fix for Activation not scaling to the entire group when greater than 1 * Update the capabilites for when grouped is supported and adapt the execution path correctly * Added filter validation for the Depthwise branch * Added a regression test for PreActivation Signed-off-by: Orlaith Monahan --- onnxruntime/core/mlas/lib/convolve.cpp | 10 +++-- .../mlas/lib/kleidiai/convolve_kleidiai.cpp | 13 ++++++- .../test/contrib_ops/fused_conv_test.cc | 39 +++++++++++++++++++ 3 files changed, 57 insertions(+), 5 deletions(-) diff --git a/onnxruntime/core/mlas/lib/convolve.cpp b/onnxruntime/core/mlas/lib/convolve.cpp index 5dc5eb1281b12..53a275ea70c84 100644 --- a/onnxruntime/core/mlas/lib/convolve.cpp +++ b/onnxruntime/core/mlas/lib/convolve.cpp @@ -1395,12 +1395,16 @@ MlasConvSupportsSymmetricChannelsLast2DFloatKernel( return false; } - const bool is_depthwise = GroupCount > 1 && FilterCount == 1; - if (GroupCount > 1 && !is_depthwise) { + const bool is_depthwise = GroupCount > 1; + if (is_depthwise) { + if (FilterCount != 1) { + return false; + } + } else if (FilterCount <= 1) { return false; } - if (!is_depthwise && (FilterCount <= 1 || KernelShape[0] < 3 || KernelShape[1] < 3)) { + if (KernelShape[0] < 3 || KernelShape[1] < 3) { return false; } diff --git a/onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp b/onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp index 1da8530350d75..0f6bc63f0a640 100644 --- a/onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp +++ b/onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp @@ -109,6 +109,12 @@ static size_t ComputeMlasWorkingBufferSize(const size_t co, } static bool CheckCapabilitiesSme(const MLAS_CONV_PARAMETERS* Parameters) { + // Grouped support in this override is only implemented for channels-last + // layout. The generic grouped path still assumes contiguous per-group CHW. + if (Parameters->GroupCount > 1 && !Parameters->ChannelsLast) { + return false; + } + if (!MlasConvSupportsSymmetricChannelsLast2DFloatKernel( Parameters->Dimensions, Parameters->BatchCount, @@ -777,7 +783,10 @@ ArmKleidiAI::MlasConv( Parameters->PackedFilterGroupStride, Input, Output, WorkingBuffer, Parameters->ChannelsLast, ThreadPool); - MlasActivation(Parameters->Activation, Output, nullptr, Parameters->FilterCount, Parameters->OutputSize, - Parameters->OutputSize); + const bool grouped_channels_last = Parameters->ChannelsLast && Parameters->GroupCount > 1; + const size_t activation_rows = grouped_channels_last ? Parameters->OutputSize : Parameters->FilterCount; + const size_t activation_cols = + grouped_channels_last ? Parameters->GroupCount * Parameters->FilterCount : Parameters->OutputSize; + MlasActivation(Parameters->Activation, Output, nullptr, activation_rows, activation_cols, activation_cols); return true; } diff --git a/onnxruntime/test/contrib_ops/fused_conv_test.cc b/onnxruntime/test/contrib_ops/fused_conv_test.cc index f453cb3bca746..daec14ef8bd0e 100644 --- a/onnxruntime/test/contrib_ops/fused_conv_test.cc +++ b/onnxruntime/test/contrib_ops/fused_conv_test.cc @@ -501,6 +501,45 @@ TEST(FusedConvTest, Cpu_NhwcDepthwiseConv2D_SymmetricPadding) { TestNhwcFusedConvFloatOp(attrs, {X, W}, {X_shape, W_shape}, expected_vals, Y_shape, true); #endif } + +TEST(FusedConvTest, Cpu_NhwcDepthwiseConv2D_Relu_NegativePreActivation) { +#if !defined(MLAS_TARGET_ARM64) + GTEST_SKIP() << "Float NHWC depthwise fast-path requires Arm64."; +#else + ConvOpAndTestAttributes attrs = { + "", // auto_pad + vector{1, 1}, // dilations + 2, // group + vector{3, 3}, // kernel_shape + vector{1, 1, 1, 1}, // pads + vector{1, 1}, // strides + "Relu" // activation + }; + + vector X_shape = {1, 3, 3, 2}; + vector X = {1.0f, 10.0f, 2.0f, 20.0f, 3.0f, 30.0f, + 4.0f, 40.0f, 5.0f, 50.0f, 6.0f, 60.0f, + 7.0f, 70.0f, 8.0f, 80.0f, 9.0f, 90.0f}; + vector W_shape = {2, 1, 3, 3}; + vector W = {-1.0f, -1.0f, -1.0f, + -1.0f, -1.0f, -1.0f, + -1.0f, -1.0f, -1.0f, + 1.0f, 1.0f, 1.0f, + 1.0f, 1.0f, 1.0f, + 1.0f, 1.0f, 1.0f}; + vector Y_shape = {1, 3, 3, 2}; + auto expected_vals = {0.0f, 120.0f, 0.0f, 210.0f, 0.0f, 160.0f, + 0.0f, 270.0f, 0.0f, 450.0f, 0.0f, 330.0f, + 0.0f, 240.0f, 0.0f, 390.0f, 0.0f, 280.0f}; + + if (!HasFloatNhwcNoTransposeSupport(X_shape, W_shape, attrs.pads, attrs.strides, attrs.group)) { + GTEST_SKIP() << "Float NHWC depthwise fast-path is not available on this configuration."; + } + + TestNhwcFusedConvFloatOp(attrs, {X, W}, {X_shape, W_shape}, expected_vals, Y_shape); + TestNhwcFusedConvFloatOp(attrs, {X, W}, {X_shape, W_shape}, expected_vals, Y_shape, true); +#endif +} #endif TEST(FusedConvTest, Cpu_Conv3D_Batched_Relu) { From 671a09c2fc740979d921c50876d70dcb7a3c3a28 Mon Sep 17 00:00:00 2001 From: Orlaith Monahan Date: Thu, 4 Jun 2026 14:25:34 +0100 Subject: [PATCH 6/7] Refactors to reduce code duplication Signed-off-by: Orlaith Monahan --- .../mlas/lib/kleidiai/convolve_kleidiai.cpp | 25 +- onnxruntime/test/mlas/bench/bench_sconv.cpp | 370 ++++++++---------- 2 files changed, 167 insertions(+), 228 deletions(-) diff --git a/onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp b/onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp index 0f6bc63f0a640..0c7b50032bad1 100644 --- a/onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp +++ b/onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp @@ -7,6 +7,7 @@ #include #include #include +#include #include #include #include @@ -92,6 +93,15 @@ static constexpr size_t ComputeConvOutSize(const size_t L, const size_t K, const return 0; } +static inline void CopyChannelBlock(float* dst, const float* src, size_t channels) { + if (channels == 1) { + *dst = *src; + return; + } + + std::memcpy(dst, src, channels * sizeof(float)); +} + static size_t ComputeMlasWorkingBufferSize(const size_t co, const size_t ih, const size_t iw, const size_t kh, const size_t kw, @@ -571,18 +581,14 @@ static void ConvolveSme(const size_t co, //channels out if (grouped_channels_last) { for (size_t pixel = 0; pixel < ih * iw; ++pixel) { const float* src = input_base + pixel * input_channels_total + g * ci; - std::copy_n(src, ci, input_group_buffer.data() + pixel * ci); + CopyChannelBlock(input_group_buffer.data() + pixel * ci, src, ci); } input_group = input_group_buffer.data(); } auto result = out; const bool need_transpose = (!input_is_channels_last) && (co > 1); - const bool use_temp_output = grouped_channels_last || need_transpose; - if (need_transpose) { - result = tmp_mlas_aligned; - } - if (grouped_channels_last) { + if (need_transpose || grouped_channels_last) { result = tmp_mlas_aligned; } @@ -648,7 +654,7 @@ static void ConvolveSme(const size_t co, //channels out for (size_t pixel = 0; pixel < m; ++pixel) { float* dst = output_base + pixel * output_channels_total + g * co; const float* src = result + pixel * co; - std::copy_n(src, co, dst); + CopyChannelBlock(dst, src, co); } } @@ -659,10 +665,7 @@ static void ConvolveSme(const size_t co, //channels out if (!grouped_channels_last) { in += ci * ih * iw; - out += use_temp_output ? 0 : m * co; - if (need_transpose) { - out += m * co; - } + out += m * co; } weights += co * ci * kh * kw; if(bias){ diff --git a/onnxruntime/test/mlas/bench/bench_sconv.cpp b/onnxruntime/test/mlas/bench/bench_sconv.cpp index 85cb68076b638..c62bdf906c522 100644 --- a/onnxruntime/test/mlas/bench/bench_sconv.cpp +++ b/onnxruntime/test/mlas/bench/bench_sconv.cpp @@ -5,6 +5,7 @@ #include "bench_util.h" #include "core/util/thread_utils.h" +#include #include #include @@ -39,80 +40,158 @@ static const std::vector& ArgNamesForConv(size_t rank) { return rank_to_args_name[rank]; } -// dummy for some strange build error when using Bench capture -void SCONV_NCHW(benchmark::State& state, const char* /*dummy*/) { - const int64_t rank = state.range(0); // Rank - const int64_t batch_size = state.range(1); // N - const int64_t groups = state.range(2); // G - const int64_t input_channels_per_group = state.range(3); // Cpg - const int64_t output_channels_per_group = state.range(4); // Fpg - - if (rank <= 0) throw std::invalid_argument("Kernel rank must greater than 0!"); - if (batch_size <= 0) throw std::invalid_argument("Batch size must greater than 0!"); - if (groups <= 0) throw std::invalid_argument("Group count must greater than 0!"); - if (input_channels_per_group <= 0) throw std::invalid_argument("input_channels_per_group must greater than 0!"); - if (output_channels_per_group <= 0) throw std::invalid_argument("output_channels_per_group must greater than 0!"); +struct ConvBenchmarkArgs { + int64_t rank; + int64_t batch_size; + int64_t groups; + int64_t input_channels_per_group; + int64_t output_channels_per_group; + int64_t total_input_channels; + int64_t total_output_channels; + std::vector input_shape; + std::vector kernel_shape; + std::vector paddings; + std::vector strides; + std::vector dilations; + std::vector output_shape; +}; + +static ConvBenchmarkArgs ParseConvBenchmarkArgs(benchmark::State& state) { + ConvBenchmarkArgs args{ + state.range(0), // rank + state.range(1), // batch_size + state.range(2), // groups + state.range(3), // input_channels_per_group + state.range(4), // output_channels_per_group + 0, + 0, + {}, + {}, + {}, + {}, + {}, + {}}; + + if (args.rank <= 0) throw std::invalid_argument("Kernel rank must be greater than 0"); + if (args.batch_size <= 0) throw std::invalid_argument("Batch size must be greater than 0"); + if (args.groups <= 0) throw std::invalid_argument("Group count must be greater than 0"); + if (args.input_channels_per_group <= 0) throw std::invalid_argument("input_channels_per_group must be greater than 0"); + if (args.output_channels_per_group <= 0) throw std::invalid_argument("output_channels_per_group must be greater than 0"); size_t arg_position = 5; - const auto input_shape = BenchArgsVector(state, arg_position, rank); - const auto kernel_shape = BenchArgsVector(state, arg_position, rank); - const auto paddings = BenchArgsVector(state, arg_position, rank * 2); - const auto strides = BenchArgsVector(state, arg_position, rank); - const auto dilations = BenchArgsVector(state, arg_position, rank); - - // do not check the size of each vector as they are forced from args. - if (std::any_of(input_shape.begin(), input_shape.end(), [](const int64_t& dim) { return dim <= 0; })) { + args.input_shape = BenchArgsVector(state, arg_position, args.rank); + args.kernel_shape = BenchArgsVector(state, arg_position, args.rank); + args.paddings = BenchArgsVector(state, arg_position, args.rank * 2); + args.strides = BenchArgsVector(state, arg_position, args.rank); + args.dilations = BenchArgsVector(state, arg_position, args.rank); + + if (std::any_of(args.input_shape.begin(), args.input_shape.end(), [](const int64_t& dim) { return dim <= 0; })) { throw std::invalid_argument("all input image dim must > 0"); } - if (std::any_of(kernel_shape.begin(), kernel_shape.end(), [](const int64_t& dim) { return dim <= 0; })) { + if (std::any_of(args.kernel_shape.begin(), args.kernel_shape.end(), [](const int64_t& dim) { return dim <= 0; })) { throw std::invalid_argument("all kernel dim must > 0"); } - if (std::any_of(strides.begin(), strides.end(), [](const int64_t& dim) { return dim <= 0; })) { + if (std::any_of(args.strides.begin(), args.strides.end(), [](const int64_t& dim) { return dim <= 0; })) { throw std::invalid_argument("all strides dim must > 0"); } - if (std::any_of(dilations.begin(), dilations.end(), [](const int64_t& dim) { return dim <= 0; })) { + if (std::any_of(args.dilations.begin(), args.dilations.end(), [](const int64_t& dim) { return dim <= 0; })) { throw std::invalid_argument("all dilations dim must > 0"); } - const int64_t GC = groups * input_channels_per_group; - const int64_t GF = groups * output_channels_per_group; - std::vector x_shape = {batch_size, GC}; - x_shape.insert(x_shape.end(), input_shape.begin(), input_shape.end()); - std::vector f_shape = {GF, input_channels_per_group}; - f_shape.insert(f_shape.end(), kernel_shape.begin(), kernel_shape.end()); - - std::vector output_shape((size_t)rank); - for (int64_t i = 0; i < rank; ++i) { - auto km = 1 + dilations[i] * (kernel_shape[i] - 1); - output_shape[i] = (paddings[i] + paddings[i + rank] + input_shape[i] - km) / strides[i] + 1; + args.total_input_channels = args.groups * args.input_channels_per_group; + args.total_output_channels = args.groups * args.output_channels_per_group; + args.output_shape.resize(static_cast(args.rank)); + for (int64_t i = 0; i < args.rank; ++i) { + const auto index = static_cast(i); + const auto km = 1 + args.dilations[index] * (args.kernel_shape[index] - 1); + args.output_shape[index] = + (args.paddings[index] + args.paddings[index + static_cast(args.rank)] + args.input_shape[index] - km) / + args.strides[index] + + 1; + } + + return args; +} + +static std::vector MakeInputShape(const ConvBenchmarkArgs& args, bool channels_last) { + std::vector shape = {args.batch_size}; + if (channels_last) { + shape.insert(shape.end(), args.input_shape.begin(), args.input_shape.end()); + shape.push_back(args.total_input_channels); + } else { + shape.push_back(args.total_input_channels); + shape.insert(shape.end(), args.input_shape.begin(), args.input_shape.end()); + } + + return shape; +} + +static std::vector MakeFilterShape(const ConvBenchmarkArgs& args) { + std::vector shape = {args.total_output_channels, args.input_channels_per_group}; + shape.insert(shape.end(), args.kernel_shape.begin(), args.kernel_shape.end()); + return shape; +} + +static std::vector MakeOutputShape(const ConvBenchmarkArgs& args, bool channels_last) { + std::vector shape = {args.batch_size}; + if (channels_last) { + shape.insert(shape.end(), args.output_shape.begin(), args.output_shape.end()); + shape.push_back(args.total_output_channels); + } else { + shape.push_back(args.total_output_channels); + shape.insert(shape.end(), args.output_shape.begin(), args.output_shape.end()); } - std::vector y_shape = {batch_size, GF}; - y_shape.insert(y_shape.end(), output_shape.begin(), output_shape.end()); + return shape; +} + +static std::vector ToSizeT(const std::vector& values) { + std::vector result(values.size()); + std::transform(values.begin(), values.end(), result.begin(), [](int64_t value) { + return static_cast(value); + }); + return result; +} + +static void PrepareConvParameters(const ConvBenchmarkArgs& args, + bool channels_last, + MLAS_THREADPOOL* thread_pool, + MLAS_CONV_PARAMETERS* parameters, + size_t* working_buffer_size) { MLAS_ACTIVATION activation; activation.ActivationKind = MlasIdentityActivation; - MLAS_CONV_PARAMETERS Parameters; - size_t WorkingBufferSize = 0; - MlasConvPrepare(&Parameters, - static_cast(rank), - static_cast(batch_size), - static_cast(groups), - static_cast(input_channels_per_group), - input_shape.data(), - kernel_shape.data(), - dilations.data(), - paddings.data(), - strides.data(), - output_shape.data(), - static_cast(output_channels_per_group), + MlasConvPrepare(parameters, + static_cast(args.rank), + static_cast(args.batch_size), + static_cast(args.groups), + static_cast(args.input_channels_per_group), + args.input_shape.data(), + args.kernel_shape.data(), + args.dilations.data(), + args.paddings.data(), + args.strides.data(), + args.output_shape.data(), + static_cast(args.output_channels_per_group), &activation, - &WorkingBufferSize, - false, + working_buffer_size, + channels_last, 0.0f, - nullptr); + thread_pool); +} + +// dummy for some strange build error when using Bench capture +void SCONV_NCHW(benchmark::State& state, const char* /*dummy*/) { + const auto args = ParseConvBenchmarkArgs(state); + const auto x_shape = MakeInputShape(args, false); + const auto f_shape = MakeFilterShape(args); + const auto y_shape = MakeOutputShape(args, false); + + MLAS_CONV_PARAMETERS Parameters; + size_t WorkingBufferSize = 0; + PrepareConvParameters(args, false, nullptr, &Parameters, &WorkingBufferSize); auto X = RandomVectorUniform(x_shape, -2.0, 2.0); auto F = RandomVectorUniform(f_shape, -1.0, 1.0); @@ -147,118 +226,40 @@ static MLAS_THREADPOOL* GetMlasThreadPoolForConvBenchmark(void) { } void SCONV_NHWC_KLEIDIAI(benchmark::State& state, const char* /*dummy*/) { - const int64_t rank = state.range(0); // Rank - const int64_t batch_size = state.range(1); // N - const int64_t groups = state.range(2); // G - const int64_t input_channels_per_group = state.range(3); // Cpg - const int64_t output_channels_per_group = state.range(4); // Fpg - - if (rank <= 0) throw std::invalid_argument("Kernel rank must be greater than 0"); - if (batch_size <= 0) throw std::invalid_argument("Batch size must be greater than 0"); - if (groups <= 0) throw std::invalid_argument("Group count must be greater than 0"); - if (input_channels_per_group <= 0) throw std::invalid_argument("input_channels_per_group must be greater than 0"); - if (output_channels_per_group <= 0) throw std::invalid_argument("output_channels_per_group must be greater than 0"); - - size_t arg_position = 5; - const auto input_shape = BenchArgsVector(state, arg_position, rank); - const auto kernel_shape = BenchArgsVector(state, arg_position, rank); - const auto paddings = BenchArgsVector(state, arg_position, rank * 2); - const auto strides = BenchArgsVector(state, arg_position, rank); - const auto dilations = BenchArgsVector(state, arg_position, rank); - - if (std::any_of(input_shape.begin(), input_shape.end(), [](const int64_t& dim) { return dim <= 0; })) { - throw std::invalid_argument("all input image dim must > 0"); - } - - if (std::any_of(kernel_shape.begin(), kernel_shape.end(), [](const int64_t& dim) { return dim <= 0; })) { - throw std::invalid_argument("all kernel dim must > 0"); - } - - if (std::any_of(strides.begin(), strides.end(), [](const int64_t& dim) { return dim <= 0; })) { - throw std::invalid_argument("all strides dim must > 0"); - } - - if (std::any_of(dilations.begin(), dilations.end(), [](const int64_t& dim) { return dim <= 0; })) { - throw std::invalid_argument("all dilations dim must > 0"); - } - - if (rank != 2 || batch_size != 1) { + const auto args = ParseConvBenchmarkArgs(state); + if (args.rank != 2 || args.batch_size != 1) { state.SkipWithError("KleidiAI NHWC benchmark requires 2D convolution with batch size 1."); return; } - std::vector input_shape_size_t(static_cast(rank)); - std::vector kernel_shape_size_t(static_cast(rank)); - std::vector paddings_size_t(static_cast(rank * 2)); - std::vector strides_size_t(static_cast(rank)); - std::vector dilations_size_t(static_cast(rank)); - for (int64_t i = 0; i < rank; ++i) { - input_shape_size_t[static_cast(i)] = static_cast(input_shape[static_cast(i)]); - kernel_shape_size_t[static_cast(i)] = static_cast(kernel_shape[static_cast(i)]); - strides_size_t[static_cast(i)] = static_cast(strides[static_cast(i)]); - dilations_size_t[static_cast(i)] = static_cast(dilations[static_cast(i)]); - paddings_size_t[static_cast(i)] = static_cast(paddings[static_cast(i)]); - paddings_size_t[static_cast(i + rank)] = static_cast(paddings[static_cast(i + rank)]); - } + const auto input_shape_size_t = ToSizeT(args.input_shape); + const auto kernel_shape_size_t = ToSizeT(args.kernel_shape); + const auto paddings_size_t = ToSizeT(args.paddings); + const auto strides_size_t = ToSizeT(args.strides); + const auto dilations_size_t = ToSizeT(args.dilations); if (!MlasConvSupportsSymmetricChannelsLast2DFloatKernel( - static_cast(rank), - static_cast(batch_size), - static_cast(groups), + static_cast(args.rank), + static_cast(args.batch_size), + static_cast(args.groups), input_shape_size_t.data(), kernel_shape_size_t.data(), dilations_size_t.data(), paddings_size_t.data(), strides_size_t.data(), - static_cast(output_channels_per_group), + static_cast(args.output_channels_per_group), 0.0f)) { state.SkipWithError("KleidiAI NHWC kernel is not supported for this benchmark shape on the current platform."); return; } - const int64_t GC = groups * input_channels_per_group; - const int64_t GF = groups * output_channels_per_group; - std::vector x_shape = {batch_size}; - x_shape.insert(x_shape.end(), input_shape.begin(), input_shape.end()); - x_shape.push_back(GC); - - std::vector f_shape = {GF, input_channels_per_group}; - f_shape.insert(f_shape.end(), kernel_shape.begin(), kernel_shape.end()); - - std::vector output_shape(static_cast(rank)); - for (int64_t i = 0; i < rank; ++i) { - auto km = 1 + dilations[static_cast(i)] * (kernel_shape[static_cast(i)] - 1); - output_shape[static_cast(i)] = - (paddings[static_cast(i)] + paddings[static_cast(i + rank)] + input_shape[static_cast(i)] - km) / - strides[static_cast(i)] + - 1; - } - - std::vector y_shape = {batch_size}; - y_shape.insert(y_shape.end(), output_shape.begin(), output_shape.end()); - y_shape.push_back(GF); + const auto x_shape = MakeInputShape(args, true); + const auto f_shape = MakeFilterShape(args); + const auto y_shape = MakeOutputShape(args, true); - MLAS_ACTIVATION activation; - activation.ActivationKind = MlasIdentityActivation; MLAS_CONV_PARAMETERS Parameters; size_t WorkingBufferSize = 0; - MlasConvPrepare(&Parameters, - static_cast(rank), - static_cast(batch_size), - static_cast(groups), - static_cast(input_channels_per_group), - input_shape.data(), - kernel_shape.data(), - dilations.data(), - paddings.data(), - strides.data(), - output_shape.data(), - static_cast(output_channels_per_group), - &activation, - &WorkingBufferSize, - true, - 0.0f, - nullptr); + PrepareConvParameters(args, true, nullptr, &Parameters, &WorkingBufferSize); auto X = RandomVectorUniform(x_shape, -2.0, 2.0); auto F = RandomVectorUniform(f_shape, -1.0, 1.0); @@ -287,79 +288,14 @@ void SCONV_NHWC_KLEIDIAI(benchmark::State& state, const char* /*dummy*/) { void SCONV_NCHW_THREADED(benchmark::State& state, const char* /*dummy*/) { MLAS_THREADPOOL* tp = GetMlasThreadPoolForConvBenchmark(); + const auto args = ParseConvBenchmarkArgs(state); + const auto x_shape = MakeInputShape(args, false); + const auto f_shape = MakeFilterShape(args); + const auto y_shape = MakeOutputShape(args, false); - const int64_t rank = state.range(0); // Rank - const int64_t batch_size = state.range(1); // N - const int64_t groups = state.range(2); // G - const int64_t input_channels_per_group = state.range(3); // Cpg - const int64_t output_channels_per_group = state.range(4); // Fpg - - if (rank <= 0) throw std::invalid_argument("Kernel rank must greater than 0!"); - if (batch_size <= 0) throw std::invalid_argument("Batch size must greater than 0!"); - if (groups <= 0) throw std::invalid_argument("Group count must greater than 0!"); - if (input_channels_per_group <= 0) throw std::invalid_argument("input_channels_per_group must greater than 0!"); - if (output_channels_per_group <= 0) throw std::invalid_argument("output_channels_per_group must greater than 0!"); - - size_t arg_position = 5; - const auto input_shape = BenchArgsVector(state, arg_position, rank); - const auto kernel_shape = BenchArgsVector(state, arg_position, rank); - const auto paddings = BenchArgsVector(state, arg_position, rank * 2); - const auto strides = BenchArgsVector(state, arg_position, rank); - const auto dilations = BenchArgsVector(state, arg_position, rank); - - // do not check the size of each vector as they are forced from args. - if (std::any_of(input_shape.begin(), input_shape.end(), [](const int64_t& dim) { return dim <= 0; })) { - throw std::invalid_argument("all input image dim must > 0"); - } - - if (std::any_of(kernel_shape.begin(), kernel_shape.end(), [](const int64_t& dim) { return dim <= 0; })) { - throw std::invalid_argument("all kernel dim must > 0"); - } - - if (std::any_of(strides.begin(), strides.end(), [](const int64_t& dim) { return dim <= 0; })) { - throw std::invalid_argument("all strides dim must > 0"); - } - - if (std::any_of(dilations.begin(), dilations.end(), [](const int64_t& dim) { return dim <= 0; })) { - throw std::invalid_argument("all dilations dim must > 0"); - } - - const int64_t GC = groups * input_channels_per_group; - const int64_t GF = groups * output_channels_per_group; - std::vector x_shape = {batch_size, GC}; - x_shape.insert(x_shape.end(), input_shape.begin(), input_shape.end()); - std::vector f_shape = {GF, input_channels_per_group}; - f_shape.insert(f_shape.end(), kernel_shape.begin(), kernel_shape.end()); - - std::vector output_shape((size_t)rank); - for (int64_t i = 0; i < rank; ++i) { - auto km = 1 + dilations[i] * (kernel_shape[i] - 1); - output_shape[i] = (paddings[i] + paddings[i + rank] + input_shape[i] - km) / strides[i] + 1; - } - std::vector y_shape = {batch_size, GF}; - y_shape.insert(y_shape.end(), output_shape.begin(), output_shape.end()); - - MLAS_ACTIVATION activation; - activation.ActivationKind = MlasIdentityActivation; MLAS_CONV_PARAMETERS Parameters; size_t WorkingBufferSize = 0; - MlasConvPrepare(&Parameters, - static_cast(rank), - static_cast(batch_size), - static_cast(groups), - static_cast(input_channels_per_group), - input_shape.data(), - kernel_shape.data(), - dilations.data(), - paddings.data(), - strides.data(), - output_shape.data(), - static_cast(output_channels_per_group), - &activation, - &WorkingBufferSize, - false, - 0.0f, - tp); + PrepareConvParameters(args, false, tp, &Parameters, &WorkingBufferSize); auto X = RandomVectorUniform(x_shape, -2.0, 2.0); auto F = RandomVectorUniform(f_shape, -1.0, 1.0); From 3439f0343dbd7f88f404f6e084bc98f268b785fb Mon Sep 17 00:00:00 2001 From: Orlaith Monahan Date: Fri, 12 Jun 2026 12:39:05 +0100 Subject: [PATCH 7/7] Lintrunner fixes Signed-off-by: Orlaith Monahan --- onnxruntime/test/contrib_ops/fused_conv_test.cc | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/onnxruntime/test/contrib_ops/fused_conv_test.cc b/onnxruntime/test/contrib_ops/fused_conv_test.cc index daec14ef8bd0e..ec67a8e3328a5 100644 --- a/onnxruntime/test/contrib_ops/fused_conv_test.cc +++ b/onnxruntime/test/contrib_ops/fused_conv_test.cc @@ -524,9 +524,9 @@ TEST(FusedConvTest, Cpu_NhwcDepthwiseConv2D_Relu_NegativePreActivation) { vector W = {-1.0f, -1.0f, -1.0f, -1.0f, -1.0f, -1.0f, -1.0f, -1.0f, -1.0f, - 1.0f, 1.0f, 1.0f, - 1.0f, 1.0f, 1.0f, - 1.0f, 1.0f, 1.0f}; + 1.0f, 1.0f, 1.0f, + 1.0f, 1.0f, 1.0f, + 1.0f, 1.0f, 1.0f}; vector Y_shape = {1, 3, 3, 2}; auto expected_vals = {0.0f, 120.0f, 0.0f, 210.0f, 0.0f, 160.0f, 0.0f, 270.0f, 0.0f, 450.0f, 0.0f, 330.0f,