diff --git a/cmake/onnxruntime_mlas.cmake b/cmake/onnxruntime_mlas.cmake index d55a4d49fb455..103d748d4f169 100644 --- a/cmake/onnxruntime_mlas.cmake +++ b/cmake/onnxruntime_mlas.cmake @@ -320,6 +320,7 @@ function(setup_kleidiai) ${MLAS_SRC_DIR}/kleidiai/sbgemm_kleidiai.cpp ${MLAS_SRC_DIR}/kleidiai/convolve_kleidiai.cpp ${MLAS_SRC_DIR}/kleidiai/qgemm_kleidiai.cpp + ${MLAS_SRC_DIR}/kleidiai/dw_conv_kleidiai.cpp ) target_link_libraries(onnxruntime_mlas PRIVATE kleidiai) list(APPEND onnxruntime_EXTERNAL_LIBRARIES kleidiai) diff --git a/onnxruntime/core/mlas/lib/convolve.cpp b/onnxruntime/core/mlas/lib/convolve.cpp index 9bff72b29d8fb..e0770e1d505df 100644 --- a/onnxruntime/core/mlas/lib/convolve.cpp +++ b/onnxruntime/core/mlas/lib/convolve.cpp @@ -1348,6 +1348,9 @@ static constexpr size_t ComputeChannelsLastConvOutSize(size_t input, size_t kern return 0; } + +constexpr size_t kKleidiAIDepthwiseRowsPerTile = 4; +constexpr size_t kKleidiAIDepthwiseColsPerTile = 4; #endif } // namespace @@ -1450,22 +1453,46 @@ MlasConvSupportsDepthwiseChannelsLast2DFloatKernel( MLAS_UNREFERENCED_PARAMETER(Beta); return false; #else - MLAS_UNREFERENCED_PARAMETER(Dimensions); - MLAS_UNREFERENCED_PARAMETER(BatchCount); - MLAS_UNREFERENCED_PARAMETER(GroupCount); - MLAS_UNREFERENCED_PARAMETER(InputChannelsPerGroup); - MLAS_UNREFERENCED_PARAMETER(InputShape); - MLAS_UNREFERENCED_PARAMETER(KernelShape); - MLAS_UNREFERENCED_PARAMETER(DilationShape); - MLAS_UNREFERENCED_PARAMETER(Padding); - MLAS_UNREFERENCED_PARAMETER(StrideShape); - MLAS_UNREFERENCED_PARAMETER(FilterCount); - MLAS_UNREFERENCED_PARAMETER(Beta); + if (GetMlasPlatform().MlasConvPrepareOverride == nullptr || + GetMlasPlatform().MlasConvOverride == nullptr) { + return false; + } - // TODO: enable only for shapes supported by the dedicated - // depthwise kernel. Until then, keep depthwise/grouped convolutions out of - // the Arm® KleidiAI™ NHWC path. - return false; + if (!MLAS_CPUIDINFO::GetCPUIDInfo().HasArm_SME2()) { + return false; + } + + if (Dimensions != 2 || BatchCount != 1 || Beta != 0.0f || GroupCount == 0) { + return false; + } + + if (InputChannelsPerGroup != 1 || FilterCount != 1) { + return false; + } + + if (KernelShape[0] != 3 || KernelShape[1] != 3) { + return false; + } + + if (StrideShape[0] != 1 || StrideShape[1] != 1) { + return false; + } + + if (DilationShape[0] != 1 || DilationShape[1] != 1) { + return false; + } + + const bool zero_padding = Padding[0] == 0 && Padding[1] == 0 && Padding[2] == 0 && Padding[3] == 0; + const bool unit_padding = Padding[0] == 1 && Padding[1] == 1 && Padding[2] == 1 && Padding[3] == 1; + if (!zero_padding && !unit_padding) { + return false; + } + + const size_t output_h = + ComputeChannelsLastConvOutSize(InputShape[0], KernelShape[0], Padding[0], StrideShape[0]); + const size_t output_w = + ComputeChannelsLastConvOutSize(InputShape[1], KernelShape[1], Padding[1], StrideShape[1]); + return output_h >= kKleidiAIDepthwiseRowsPerTile && output_w >= kKleidiAIDepthwiseColsPerTile; #endif } diff --git a/onnxruntime/core/mlas/lib/kai_ukernel_interface.cpp b/onnxruntime/core/mlas/lib/kai_ukernel_interface.cpp index 6ee80594c6b49..a808cbe301d64 100644 --- a/onnxruntime/core/mlas/lib/kai_ukernel_interface.cpp +++ b/onnxruntime/core/mlas/lib/kai_ukernel_interface.cpp @@ -40,6 +40,8 @@ #include "kai/ukernels/matmul/matmul_clamp_f32_f32_f32p/kai_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla.h" // IMATMUL #include "kai/ukernels/matmul/imatmul_clamp_f32_f32p_f32p/kai_imatmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme2_mopa.h" +// DWCONV +#include "kai/ukernels/dwconv/dwconv_f32_f32_f32p/kai_dwconv_clamp_f32_f32_f32p1vlx1b_3x3_s1_4xc_sme2_mla.h" #if defined(ENABLE_QMX_KERNELS) // QMX kernels (optional) @@ -232,6 +234,9 @@ const KaiF32IMatmulKernel imatmul_conv_sme2 = const KaiBF16SBgemmKernel sbgemm_gemm_sme2 = KAI_WRAP_UKERNEL_RUN_MATMUL_11(matmul_clamp_f32_bf16p2vlx2_bf16p2vlx2_2vlx2vl_sme2_mopa); +const KaiF32DepthwiseConvKernel dwconv_sme2 = + KAI_WRAP_UKERNEL_RUN_DWCONV_PLANAR_4(dwconv_clamp_f32_f32_f32p1vlx1b_3x3_s1_4xc_sme2_mla); + #if defined(ENABLE_QMX_KERNELS) const KaiF32IMatmulKernel imatmul_conv_qmx = KAI_WRAP_UKERNEL_RUN_IMATMUL_PACKED_7(imatmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_qmx_mopa); @@ -372,3 +377,8 @@ const KaiBF16SBgemmKernel& GetKleidiAISBGemmUKernel() { // Currently only SME2 variant exists for bfloat16/SBGEMM kernel return sbgemm_gemm_sme2; } + +const KaiF32DepthwiseConvKernel& GetKleidiAIDepthwiseConvUKernel() { + // Currently only an SME2 variant exists for FP32 depthwise convolution. + return dwconv_sme2; +} diff --git a/onnxruntime/core/mlas/lib/kai_ukernel_interface.h b/onnxruntime/core/mlas/lib/kai_ukernel_interface.h index 155ecf1762b3b..0e3365cf9d445 100644 --- a/onnxruntime/core/mlas/lib/kai_ukernel_interface.h +++ b/onnxruntime/core/mlas/lib/kai_ukernel_interface.h @@ -18,14 +18,19 @@ #include "kai/ukernels/matmul/imatmul_clamp_f32_f32p_f32p/kai_imatmul_clamp_f32_f32p_f32p_interface.h" +#include "kai/ukernels/dwconv/dwconv_f32_f32_f32p/kai_dwconv_clamp_f32_f32_f32p_interface.h" + // Wrapper type that carries a stable "name" alongside the KAI ukernel interface. // This avoids needing to infer which underlying microkernel was selected from a function pointer. template -struct KaiMatmulKernel { +struct KaiKernel { const char* name; UkernelFn ukernel; }; +template +using KaiMatmulKernel = KaiKernel; + // Wrapper for FP32 GEMM kernels where both LHS and RHS are pre-packed (common SGEMM path). using KaiF32SgemmKernel = KaiMatmulKernel; @@ -43,6 +48,9 @@ using KaiF32IMatmulKernel = KaiMatmulKernel; +// Wrapper for FP32 planar depthwise convolution kernels. +using KaiF32DepthwiseConvKernel = KaiKernel; + // Returns the selected Qnbit GEMM ukernel based on runtime CPU capabilities. const KaiQnbitGemmKernel& GetKleidiAIGemmUKernel(); @@ -62,4 +70,7 @@ const KaiF32SgemvKernel& GetKleidiAISGemvUKernel(); const KaiF32IMatmulKernel& GetKleidiAIF32IMatmulUKernel(); // Returns the selected BF16 SBGEMM ukernel used by the KleidiAI based on runtime CPU capabilities. -const KaiBF16SBgemmKernel& GetKleidiAISBGemmUKernel(); \ No newline at end of file +const KaiBF16SBgemmKernel& GetKleidiAISBGemmUKernel(); + +// Returns the selected FP32 depthwise convolution ukernel based on runtime CPU capabilities. +const KaiF32DepthwiseConvKernel& GetKleidiAIDepthwiseConvUKernel(); \ No newline at end of file diff --git a/onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp b/onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp index 4f62ab24648ef..466916e8b9aa6 100644 --- a/onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp +++ b/onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp @@ -123,7 +123,17 @@ static size_t ComputeMlasWorkingBufferSize(const size_t co, return m * co; } -static bool CheckCapabilitiesSme(const MLAS_CONV_PARAMETERS* Parameters) { +static bool CheckIgemmRouteCapabilities(const MLAS_CONV_PARAMETERS* Parameters, + const ArmKleidiAI::ConvRouteSelection& route_selection) { + if (route_selection.route != ArmKleidiAI::ConvRoute::IGemm) { + if (route_selection.route == ArmKleidiAI::ConvRoute::SGemmFallback) { + KLEIDIAI_DEBUG_LOG("CheckIgemmRouteCapabilities returning false to prefer SGEMM-backed conv path."); + } else { + KLEIDIAI_DEBUG_LOG("CheckIgemmRouteCapabilities returning false on non-IGEMM route."); + } + return false; + } + // 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) { @@ -141,35 +151,24 @@ static bool CheckCapabilitiesSme(const MLAS_CONV_PARAMETERS* Parameters) { Parameters->StrideShape, Parameters->FilterCount, Parameters->Beta)) { - KLEIDIAI_DEBUG_LOG("CheckCapabilitiesSme returning false on shared capability checks."); + KLEIDIAI_DEBUG_LOG("CheckIgemmRouteCapabilities returning false on shared capability checks."); return false; } - const auto route_selection = ArmKleidiAI::SelectConvRoute(Parameters); - const auto route = route_selection.route; - - if (route == ArmKleidiAI::ConvRoute::IGemm) { - // ensure LHS packed buffer size is non-zero - const size_t d_kh = route_selection.effective_kernel_h; - const size_t d_kw = route_selection.effective_kernel_w; - const size_t m_step = imatmul_conv.ukernel.get_m_step(); + // ensure LHS packed buffer size is non-zero + const size_t d_kh = route_selection.effective_kernel_h; + const size_t d_kw = route_selection.effective_kernel_w; + const size_t m_step = imatmul_conv.ukernel.get_m_step(); - const size_t bytes_per_m_step = kai_get_lhs_packed_size_lhs_imatmul_pack_x32p2vlx1_x32p_sme( - m_step, d_kh * d_kw, Parameters->InputChannels); + const size_t bytes_per_m_step = kai_get_lhs_packed_size_lhs_imatmul_pack_x32p2vlx1_x32p_sme( + m_step, d_kh * d_kw, Parameters->InputChannels); - if (bytes_per_m_step == 0) { - KLEIDIAI_DEBUG_LOG("CheckCapabilitiesSME returning false on zero LHS packed size"); - return false; - } - return true; + if (bytes_per_m_step == 0) { + KLEIDIAI_DEBUG_LOG("CheckIgemmRouteCapabilities returning false on zero LHS packed size"); + return false; } - if (route == ArmKleidiAI::ConvRoute::SGemmFallback) { - KLEIDIAI_DEBUG_LOG("CheckCapabilitiesSme returning false to prefer SGEMM-backed conv path."); - } else { - KLEIDIAI_DEBUG_LOG("CheckCapabilitiesSme returning false on functional or optimization checks."); - } - return false; + return true; } //General purpose axis swapping @@ -758,7 +757,13 @@ ArmKleidiAI::MlasConvPrepare(MLAS_CONV_PARAMETERS* Parameters, Parameters->ThreadCount = MlasGetMaximumThreadCount(ThreadPool); - if(!CheckCapabilitiesSme(Parameters)){ + const auto route_selection = ArmKleidiAI::SelectConvRoute(Parameters); + if (route_selection.route == ArmKleidiAI::ConvRoute::Depthwise) { + *WorkingBufferSize = 0; + return true; + } + + if (!CheckIgemmRouteCapabilities(Parameters, route_selection)) { return false; } @@ -790,26 +795,59 @@ ArmKleidiAI::MlasConv( return false; } - if(!CheckCapabilitiesSme(Parameters)){ - // Fallback to Default Mlas + const auto route_selection = ArmKleidiAI::SelectConvRoute(Parameters); + + if (route_selection.route == ArmKleidiAI::ConvRoute::Depthwise) { + if (DepthwiseConvKleidiAI(Parameters->BatchCount, + Parameters->InputShape[0], + Parameters->InputShape[1], + Parameters->GroupCount * Parameters->InputChannels, + Parameters->KernelShape[0], + Parameters->KernelShape[1], + Parameters->Padding[0], + Parameters->Padding[1], + Parameters->Padding[2], + Parameters->Padding[3], + Parameters->ChannelsLast, + Input, + Filter, + Bias, + Output, + -std::numeric_limits::max(), + std::numeric_limits::max())) { + const size_t activation_rows = Parameters->ChannelsLast + ? Parameters->OutputSize + : Parameters->GroupCount * Parameters->FilterCount; + const size_t activation_cols = Parameters->ChannelsLast + ? Parameters->GroupCount * Parameters->FilterCount + : Parameters->OutputSize; + MlasActivation(Parameters->Activation, Output, nullptr, activation_rows, activation_cols, activation_cols); + return true; + } + return false; - }; - ConvolveSme(Parameters->FilterCount, Parameters->InputChannels, // channel out, in - Parameters->InputShape[0], Parameters->InputShape[1], // image dimensions - Parameters->KernelShape[0], Parameters->KernelShape[1], // kernel dimensions - Parameters->StrideShape[0], Parameters->StrideShape[1], // kernel stride dimensions - Parameters->DilationShape[0], Parameters->DilationShape[1], // kernel dilation - Parameters->Padding[0], // image padding - Parameters->GroupCount, // filter groups - Filter, Bias, - reinterpret_cast(Parameters->PackedFilter), - Parameters->PackedFilterGroupStride, - Input, Output, WorkingBuffer, Parameters->ChannelsLast, ThreadPool); - - 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; + } + + if (CheckIgemmRouteCapabilities(Parameters, route_selection)) { + ConvolveSme(Parameters->FilterCount, Parameters->InputChannels, // channel out, in + Parameters->InputShape[0], Parameters->InputShape[1], // image dimensions + Parameters->KernelShape[0], Parameters->KernelShape[1], // kernel dimensions + Parameters->StrideShape[0], Parameters->StrideShape[1], // kernel stride dimensions + Parameters->DilationShape[0], Parameters->DilationShape[1], // kernel dilation + Parameters->Padding[0], // image padding + Parameters->GroupCount, // filter groups + Filter, Bias, + reinterpret_cast(Parameters->PackedFilter), + Parameters->PackedFilterGroupStride, + Input, Output, WorkingBuffer, Parameters->ChannelsLast, ThreadPool); + + 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; + } + + return false; } diff --git a/onnxruntime/core/mlas/lib/kleidiai/dw_conv_kleidiai.cpp b/onnxruntime/core/mlas/lib/kleidiai/dw_conv_kleidiai.cpp new file mode 100644 index 0000000000000..972363a610af5 --- /dev/null +++ b/onnxruntime/core/mlas/lib/kleidiai/dw_conv_kleidiai.cpp @@ -0,0 +1,333 @@ +// +// SPDX-FileCopyrightText: Copyright 2025-2026 Arm Limited and/or its affiliates +// +// SPDX-License-Identifier: MIT +// + +#include "mlasi_kleidiai.h" + +#include +#include +#include + +#include "kai_ukernel_interface.h" + +#include "kai/ukernels/dwconv/pack/kai_rhs_dwconv_pack_x32p1vlx1b_x32_x32_sme.h" + +const KaiF32DepthwiseConvKernel& dwconv = GetKleidiAIDepthwiseConvUKernel(); + +namespace ArmKleidiAI { +namespace { + +struct DwconvTlsBuffers { + std::vector feature_map_nhwc; + std::vector weights_hwcn; + std::vector weights_packed; + std::vector nhwc_out; + std::vector bias_fallback; +}; + +thread_local DwconvTlsBuffers g_dwconv_tls; + +constexpr size_t kDwconvColsPerTile = 4; + +static void ConvertNchwToNhwc(const float* src, + float* dst, + size_t batches, + size_t channels, + size_t height, + size_t width) { + const size_t src_stride_n = channels * height * width; + const size_t src_stride_c = height * width; + const size_t src_stride_h = width; + const size_t dst_stride_n = height * width * channels; + const size_t dst_stride_h = width * channels; + const size_t dst_stride_w = channels; + + for (size_t n = 0; n < batches; ++n) { + for (size_t h = 0; h < height; ++h) { + for (size_t w = 0; w < width; ++w) { + for (size_t c = 0; c < channels; ++c) { + const size_t src_index = n * src_stride_n + c * src_stride_c + h * src_stride_h + w; + const size_t dst_index = n * dst_stride_n + h * dst_stride_h + w * dst_stride_w + c; + dst[dst_index] = src[src_index]; + } + } + } + } +} + +static void ConvertNhwcToNchw(const float* src, + float* dst, + size_t batches, + size_t channels, + size_t height, + size_t width) { + const size_t dst_stride_n = channels * height * width; + const size_t dst_stride_c = height * width; + const size_t dst_stride_h = width; + const size_t src_stride_n = height * width * channels; + const size_t src_stride_h = width * channels; + const size_t src_stride_w = channels; + + for (size_t n = 0; n < batches; ++n) { + for (size_t c = 0; c < channels; ++c) { + for (size_t h = 0; h < height; ++h) { + for (size_t w = 0; w < width; ++w) { + const size_t src_index = n * src_stride_n + h * src_stride_h + w * src_stride_w + c; + const size_t dst_index = n * dst_stride_n + c * dst_stride_c + h * dst_stride_h + w; + dst[dst_index] = src[src_index]; + } + } + } + } +} + +static void ConvertDepthwiseWeightsToHwcn(const float* src, + float* dst, + size_t channels, + size_t filter_height, + size_t filter_width) { + const size_t kernel_size = filter_height * filter_width; + for (size_t c = 0; c < channels; ++c) { + const float* channel_weights = src + c * kernel_size; + for (size_t kh = 0; kh < filter_height; ++kh) { + for (size_t kw = 0; kw < filter_width; ++kw) { + const size_t dst_index = (kh * filter_width + kw) * channels + c; + dst[dst_index] = channel_weights[kh * filter_width + kw]; + } + } + } +} + +static bool TryComputeDepthwiseOutputShape(size_t in_height, + size_t in_width, + size_t filter_height, + size_t filter_width, + size_t pad_top, + size_t pad_left, + size_t pad_bottom, + size_t pad_right, + size_t& out_height, + size_t& out_width) { + const size_t padded_height = in_height + pad_top + pad_bottom; + const size_t padded_width = in_width + pad_left + pad_right; + + if (padded_height < filter_height || padded_width < filter_width) { + return false; + } + + out_height = padded_height + 1 - filter_height; + out_width = padded_width + 1 - filter_width; + return true; +} + +} // namespace + +bool +MLASCALL +DepthwiseConvKleidiAISupported(const MLAS_CONV_PARAMETERS* Parameters) { + if (Parameters == nullptr) { + return false; + } + + if (!UseSME2) { + return false; + } + + if (Parameters->BackendKernelSelectorConfig && !Parameters->BackendKernelSelectorConfig->use_kleidiai) { + return false; + } + + if (Parameters->Dimensions != 2) { + return false; + } + + // The current direct kernel path processes a single batch at a time. + if (Parameters->BatchCount != 1) { + return false; + } + + if (Parameters->Beta != 0.0f) { + return false; + } + + // Depthwise convolution with multiplier 1: one input channel and one filter per group. + if (Parameters->InputChannels != 1 || Parameters->FilterCount != 1) { + return false; + } + + // Kernel specialization is 3x3 with unit stride and dilation. + if (Parameters->KernelShape[0] != 3 || Parameters->KernelShape[1] != 3) { + return false; + } + + if (Parameters->StrideShape[0] != 1 || Parameters->StrideShape[1] != 1) { + return false; + } + + if (Parameters->DilationShape[0] != 1 || Parameters->DilationShape[1] != 1) { + return false; + } + + const bool zero_padding = Parameters->Padding[0] == 0 && Parameters->Padding[1] == 0 && + Parameters->Padding[2] == 0 && Parameters->Padding[3] == 0; + const bool unit_padding = Parameters->Padding[0] == 1 && Parameters->Padding[1] == 1 && + Parameters->Padding[2] == 1 && Parameters->Padding[3] == 1; + if (!zero_padding && !unit_padding) { + return false; + } + + size_t out_height = 0; + size_t out_width = 0; + if (!TryComputeDepthwiseOutputShape(Parameters->InputShape[0], + Parameters->InputShape[1], + Parameters->KernelShape[0], + Parameters->KernelShape[1], + Parameters->Padding[0], + Parameters->Padding[1], + Parameters->Padding[2], + Parameters->Padding[3], + out_height, + out_width)) { + return false; + } + + return out_height >= dwconv.ukernel.get_m_step() && out_width >= kDwconvColsPerTile; +} + +bool +MLASCALL +DepthwiseConvKleidiAI(size_t batches, + size_t in_height, + size_t in_width, + size_t channels, + size_t filter_height, + size_t filter_width, + size_t pad_top, + size_t pad_left, + size_t pad_bottom, + size_t pad_right, + bool channels_last, + const float* feature_map, + const float* weights, + const float* bias, + float* out, + float clamp_min, + float clamp_max) { + if (!UseSME2 || feature_map == nullptr || weights == nullptr || out == nullptr) { + return false; + } + + if (batches != 1 || channels == 0) { + return false; + } + + size_t out_height = 0; + size_t out_width = 0; + if (!TryComputeDepthwiseOutputShape(in_height, + in_width, + filter_height, + filter_width, + pad_top, + pad_left, + pad_bottom, + pad_right, + out_height, + out_width)) { + return false; + } + + const size_t rows_handled = dwconv.ukernel.get_m_step(); + if (out_height < rows_handled || out_width < kDwconvColsPerTile) { + return false; + } + + auto& tls = g_dwconv_tls; + + const float* feature_map_nhwc = feature_map; + if (!channels_last) { + const size_t input_size = batches * in_height * in_width * channels; + tls.feature_map_nhwc.resize(input_size); + ConvertNchwToNhwc(feature_map, tls.feature_map_nhwc.data(), batches, channels, in_height, in_width); + feature_map_nhwc = tls.feature_map_nhwc.data(); + } + + const size_t weights_size = filter_height * filter_width * channels; + tls.weights_hwcn.resize(weights_size); + ConvertDepthwiseWeightsToHwcn(weights, tls.weights_hwcn.data(), channels, filter_height, filter_width); + + const float* bias_data = bias; + if (bias_data == nullptr) { + tls.bias_fallback.assign(channels, 0.0f); + bias_data = tls.bias_fallback.data(); + } + + const size_t packed_size_bytes = + kai_rhs_get_dst_size_dwconv_pack_x32p1vlx1b_x32_x32_sme(filter_height, filter_width, channels); + tls.weights_packed.resize(packed_size_bytes); + KLEIDIAI_KERNEL_LOG("kai_run_rhs_dwconv_pack_x32p1vlx1b_x32_x32_sme" + << " filter_height=" << filter_height << " filter_width=" << filter_width + << " channels=" << channels); + kai_run_rhs_dwconv_pack_x32p1vlx1b_x32_x32_sme(filter_height, + filter_width, + filter_height, + filter_width, + channels, + tls.weights_hwcn.data(), + bias_data, + tls.weights_packed.data()); + + float* nhwc_out = out; + if (!channels_last) { + const size_t output_size = batches * out_height * out_width * channels; + tls.nhwc_out.assign(output_size, 0.0f); + nhwc_out = tls.nhwc_out.data(); + } + + const size_t in_row_stride_elements = in_width * channels; + const size_t out_row_stride_elements = out_width * channels; + for (size_t out_row = 0; out_row < out_height; out_row += rows_handled) { + const ptrdiff_t start_in_row = static_cast(out_row) - static_cast(pad_top); + const size_t kernel_pad_top = start_in_row < 0 ? static_cast(-start_in_row) : 0; + const size_t in_row = start_in_row < 0 ? 0 : static_cast(start_in_row); + + const size_t rows_to_process = std::min(rows_handled, out_height - out_row); + size_t valid_input_rows = 0; + if (in_row < in_height) { + const size_t max_rows_available = in_height - in_row; + const size_t needed_rows = filter_height + rows_to_process - 1; + valid_input_rows = std::min(max_rows_available, needed_rows); + } + + const float* inptr = feature_map_nhwc + in_row * in_row_stride_elements; + float* outptr = nhwc_out + out_row * out_row_stride_elements; + + KLEIDIAI_KERNEL_LOG(dwconv.name + << " valid_input_rows=" << valid_input_rows + << " valid_dst_rows=" << rows_to_process + << " pad_left=" << pad_left << " pad_top=" << kernel_pad_top); + dwconv.ukernel.run_dwconv(inptr, + tls.weights_packed.data(), + outptr, + in_row_stride_elements * sizeof(float), + channels * sizeof(float), + out_row_stride_elements * sizeof(float), + channels * sizeof(float), + valid_input_rows, + rows_to_process, + pad_left, + kernel_pad_top, + 0.0f, + clamp_min, + clamp_max); + } + + if (!channels_last) { + ConvertNhwcToNchw(nhwc_out, out, batches, channels, out_height, out_width); + } + + return true; +} + +} // namespace ArmKleidiAI diff --git a/onnxruntime/core/mlas/lib/kleidiai/mlasi_kleidiai.h b/onnxruntime/core/mlas/lib/kleidiai/mlasi_kleidiai.h index ae000168f7004..b71971a5c1c75 100644 --- a/onnxruntime/core/mlas/lib/kleidiai/mlasi_kleidiai.h +++ b/onnxruntime/core/mlas/lib/kleidiai/mlasi_kleidiai.h @@ -56,9 +56,16 @@ inline const bool UseSME2 = MLAS_CPUIDINFO::GetCPUIDInfo().HasArm_SME2(); inline const bool UseSME = MLAS_CPUIDINFO::GetCPUIDInfo().HasArm_SME(); inline const std::string_view vendor_name = MLAS_CPUIDINFO::GetCPUIDInfo().GetCPUVendor(); +bool +MLASCALL +DepthwiseConvKleidiAISupported( + const MLAS_CONV_PARAMETERS* Parameters + ); + // Selects the convolution route for Arm® KleidiAI™ enum class ConvRoute { NoKleidiAi, // decline the conv, caller runs unchanged + Depthwise, // handle the whole conv via SME2 depthwise kernel IGemm, // handle the whole conv via SME IGEMM kernel SGemmFallback, // decline IGEMM, but still route the per-segment SGEMM slices through MlasGemm // so that the Arm® KleidiAI™ SGEMM backend override can pick them up @@ -117,6 +124,12 @@ inline ConvRouteSelection SelectConvRoute(const MLAS_CONV_PARAMETERS* Parameters return ConvRouteSelection{}; } + if (Parameters->InputChannels == 1 && Parameters->FilterCount == 1) { + return DepthwiseConvKleidiAISupported(Parameters) + ? ConvRouteSelection{ConvRoute::Depthwise, Parameters->KernelShape[0], Parameters->KernelShape[1]} + : ConvRouteSelection{}; + } + size_t effective_kernel_h; size_t effective_kernel_w; if (!TryComputeDilatedKernelSize(Parameters->DilationShape[0], Parameters->KernelShape[0], &effective_kernel_h) || @@ -338,6 +351,28 @@ MlasConv( MLAS_THREADPOOL* ThreadPool ); +bool +MLASCALL +DepthwiseConvKleidiAI( + size_t batches, + size_t in_height, + size_t in_width, + size_t channels, + size_t filter_height, + size_t filter_width, + size_t pad_top, + size_t pad_left, + size_t pad_bottom, + size_t pad_right, + bool channels_last, + const float* feature_map, + const float* weights, + const float* bias, + float* out, + float clamp_min, + float clamp_max + ); + size_t MLASCALL MlasConvSymmetricChannelsLast2DFloatPackWSize( diff --git a/onnxruntime/core/optimizer/nhwc_transformer.cc b/onnxruntime/core/optimizer/nhwc_transformer.cc index b324f6ec6da5f..10429de34df95 100644 --- a/onnxruntime/core/optimizer/nhwc_transformer.cc +++ b/onnxruntime/core/optimizer/nhwc_transformer.cc @@ -207,12 +207,14 @@ bool FloatNhwcWrapperFilter(const onnx_transpose_optimization::api::GraphRef& gr std::array strides{1, 1}; std::array pads{}; size_t batch_count = 0; + size_t input_channels_per_group = 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, total_filter_count) || + !TryGetDimValueAsSizeT(*weight_shape, 1, input_channels_per_group) || !TryGetDimValueAsSizeT(*weight_shape, 2, kernel_spatial_shape[0]) || !TryGetDimValueAsSizeT(*weight_shape, 3, kernel_spatial_shape[1])) { return false; @@ -251,11 +253,6 @@ bool FloatNhwcWrapperFilter(const onnx_transpose_optimization::api::GraphRef& gr return true; } - size_t input_channels_per_group = 0; - if (!TryGetDimValueAsSizeT(*weight_shape, 1, input_channels_per_group)) { - return false; - } - return MlasConvSupportsDepthwiseChannelsLast2DFloatKernel( /*Dimensions*/ 2, batch_count, diff --git a/onnxruntime/test/mlas/unittest/test_dwconv_kleidiai.cpp b/onnxruntime/test/mlas/unittest/test_dwconv_kleidiai.cpp new file mode 100644 index 0000000000000..a2a0cc3fe712b --- /dev/null +++ b/onnxruntime/test/mlas/unittest/test_dwconv_kleidiai.cpp @@ -0,0 +1,305 @@ +// +// SPDX-FileCopyrightText: Copyright 2025-2026 Arm Limited and/or its affiliates +// +// SPDX-License-Identifier: MIT +// + +#if defined(USE_KLEIDIAI) && !defined(_MSC_VER) + +#include "test_util.h" +#include "core/mlas/lib/kleidiai/mlasi_kleidiai.h" + +#include +#include +#include +#include +#include + +namespace { + +void ConvertNchwToNhwc(const float* src, + float* dst, + size_t batches, + size_t channels, + size_t height, + size_t width) { + const size_t src_stride_n = channels * height * width; + const size_t src_stride_c = height * width; + const size_t src_stride_h = width; + const size_t dst_stride_n = height * width * channels; + const size_t dst_stride_h = width * channels; + const size_t dst_stride_w = channels; + + for (size_t n = 0; n < batches; ++n) { + for (size_t h = 0; h < height; ++h) { + for (size_t w = 0; w < width; ++w) { + for (size_t c = 0; c < channels; ++c) { + const size_t src_index = n * src_stride_n + c * src_stride_c + h * src_stride_h + w; + const size_t dst_index = n * dst_stride_n + h * dst_stride_h + w * dst_stride_w + c; + dst[dst_index] = src[src_index]; + } + } + } + } +} + +void DepthwiseReferenceNchw(const float* input, + const float* weights, + const float* bias, + size_t batches, + size_t channels, + size_t in_height, + size_t in_width, + size_t filter_height, + size_t filter_width, + size_t pad_top, + size_t pad_left, + size_t pad_bottom, + size_t pad_right, + float clamp_min, + float clamp_max, + float* output) { + const size_t out_height = in_height + pad_top + pad_bottom + 1 - filter_height; + const size_t out_width = in_width + pad_left + pad_right + 1 - filter_width; + + for (size_t b = 0; b < batches; ++b) { + for (size_t c = 0; c < channels; ++c) { + for (size_t oh = 0; oh < out_height; ++oh) { + for (size_t ow = 0; ow < out_width; ++ow) { + float acc = bias != nullptr ? bias[c] : 0.0f; + for (size_t kh = 0; kh < filter_height; ++kh) { + const int in_y = static_cast(oh) + static_cast(kh) - static_cast(pad_top); + if (in_y < 0 || in_y >= static_cast(in_height)) { + continue; + } + for (size_t kw = 0; kw < filter_width; ++kw) { + const int in_x = static_cast(ow) + static_cast(kw) - static_cast(pad_left); + if (in_x < 0 || in_x >= static_cast(in_width)) { + continue; + } + const size_t input_idx = + (((b * channels) + c) * in_height + static_cast(in_y)) * in_width + + static_cast(in_x); + const size_t weight_idx = (c * filter_height + kh) * filter_width + kw; + acc += input[input_idx] * weights[weight_idx]; + } + } + const size_t output_idx = (((b * channels) + c) * out_height + oh) * out_width + ow; + output[output_idx] = std::clamp(acc, clamp_min, clamp_max); + } + } + } + } +} + +void RunDepthwiseConvCase(size_t channels, size_t in_height, size_t in_width, size_t padding, bool channels_last) { + if (!MLAS_CPUIDINFO::GetCPUIDInfo().HasArm_SME2()) { + GTEST_SKIP() << "DepthwiseConvKleidiAI requires ARM64 SME2. Skipping test."; + } + + constexpr size_t batches = 1; + constexpr size_t filter_height = 3; + constexpr size_t filter_width = 3; + + const size_t pad_top = padding; + const size_t pad_left = padding; + const size_t pad_bottom = padding; + const size_t pad_right = padding; + + const size_t input_size = batches * channels * in_height * in_width; + const size_t weights_size = channels * filter_height * filter_width; + const size_t out_height = in_height + pad_top + pad_bottom + 1 - filter_height; + const size_t out_width = in_width + pad_left + pad_right + 1 - filter_width; + const size_t output_size = batches * channels * out_height * out_width; + + std::vector input(input_size); + std::vector weights(weights_size); + std::vector bias(channels); + std::vector input_nhwc(input_size); + std::vector expected(output_size); + std::vector expected_nhwc(output_size); + std::vector output(output_size, std::numeric_limits::quiet_NaN()); + + std::mt19937 rng(static_cast(channels * 131 + in_height * 17 + padding)); + std::uniform_real_distribution dist(-1.0f, 1.0f); + + auto fill_buffer = [&](std::vector& buffer) { + for (float& v : buffer) { + v = dist(rng); + } + }; + + fill_buffer(input); + fill_buffer(weights); + fill_buffer(bias); + + const float clamp_min = -std::numeric_limits::max(); + const float clamp_max = std::numeric_limits::max(); + + DepthwiseReferenceNchw(input.data(), + weights.data(), + bias.data(), + batches, + channels, + in_height, + in_width, + filter_height, + filter_width, + pad_top, + pad_left, + pad_bottom, + pad_right, + clamp_min, + clamp_max, + expected.data()); + + ConvertNchwToNhwc(input.data(), input_nhwc.data(), batches, channels, in_height, in_width); + ConvertNchwToNhwc(expected.data(), expected_nhwc.data(), batches, channels, out_height, out_width); + + const float* input_data = channels_last ? input_nhwc.data() : input.data(); + const float* expected_data = channels_last ? expected_nhwc.data() : expected.data(); + + const bool status = ArmKleidiAI::DepthwiseConvKleidiAI(batches, + in_height, + in_width, + channels, + filter_height, + filter_width, + pad_top, + pad_left, + pad_bottom, + pad_right, + channels_last, + input_data, + weights.data(), + bias.data(), + output.data(), + clamp_min, + clamp_max); + ASSERT_TRUE(status); + + for (size_t i = 0; i < output_size; ++i) { + EXPECT_NEAR(expected_data[i], output[i], 1e-4f) << "Mismatch at element " << i; + } +} + +void RunMlasConvChannelsLastCase(size_t channels, size_t in_height, size_t in_width, size_t padding) { + if (!MLAS_CPUIDINFO::GetCPUIDInfo().HasArm_SME2()) { + GTEST_SKIP() << "DepthwiseConvKleidiAI requires ARM64 SME2. Skipping test."; + } + + constexpr size_t batches = 1; + constexpr size_t filter_height = 3; + constexpr size_t filter_width = 3; + constexpr size_t filter_count = 1; + + const size_t out_height = in_height + padding + padding + 1 - filter_height; + const size_t out_width = in_width + padding + padding + 1 - filter_width; + const size_t input_size = batches * channels * in_height * in_width; + const size_t output_size = batches * channels * out_height * out_width; + + std::vector input(input_size); + std::vector input_nhwc(input_size); + std::vector weights(channels * filter_height * filter_width); + std::vector bias(channels); + std::vector expected(output_size); + std::vector expected_nhwc(output_size); + std::vector output(output_size, std::numeric_limits::quiet_NaN()); + + std::mt19937 rng(static_cast(channels * 197 + in_width * 23 + padding)); + std::uniform_real_distribution dist(-1.0f, 1.0f); + for (float& v : input) { + v = dist(rng); + } + for (float& v : weights) { + v = dist(rng); + } + for (float& v : bias) { + v = dist(rng); + } + + DepthwiseReferenceNchw(input.data(), + weights.data(), + bias.data(), + batches, + channels, + in_height, + in_width, + filter_height, + filter_width, + padding, + padding, + padding, + padding, + -std::numeric_limits::max(), + std::numeric_limits::max(), + expected.data()); + ConvertNchwToNhwc(input.data(), input_nhwc.data(), batches, channels, in_height, in_width); + ConvertNchwToNhwc(expected.data(), expected_nhwc.data(), batches, channels, out_height, out_width); + + const int64_t input_shape[] = {static_cast(in_height), static_cast(in_width)}; + const int64_t kernel_shape[] = {static_cast(filter_height), static_cast(filter_width)}; + const int64_t dilation_shape[] = {1, 1}; + const int64_t pads[] = {static_cast(padding), static_cast(padding), + static_cast(padding), static_cast(padding)}; + const int64_t strides[] = {1, 1}; + const int64_t output_shape[] = {static_cast(out_height), static_cast(out_width)}; + MLAS_ACTIVATION activation; + activation.ActivationKind = MlasIdentityActivation; + + MLAS_CONV_PARAMETERS parameters{}; + size_t working_buffer_size = 0; + MlasConvPrepare(¶meters, + 2, + batches, + channels, + 1, + input_shape, + kernel_shape, + dilation_shape, + pads, + strides, + output_shape, + filter_count, + &activation, + &working_buffer_size, + true, + 0.0f, + nullptr); + std::vector working_buffer(working_buffer_size); + MlasConv(¶meters, + input_nhwc.data(), + weights.data(), + bias.data(), + working_buffer.data(), + output.data(), + nullptr); + + for (size_t i = 0; i < output_size; ++i) { + EXPECT_NEAR(expected_nhwc[i], output[i], 1e-4f) << "Mismatch at element " << i; + } +} + +} // namespace + +TEST(MlasKleidiDepthwiseTest, ZeroPaddingNchw) { + RunDepthwiseConvCase(/*channels=*/32, /*in_height=*/8, /*in_width=*/8, /*padding=*/0, /*channels_last=*/false); +} + +TEST(MlasKleidiDepthwiseTest, UnitPaddingNchw) { + RunDepthwiseConvCase(/*channels=*/32, /*in_height=*/8, /*in_width=*/8, /*padding=*/1, /*channels_last=*/false); +} + +TEST(MlasKleidiDepthwiseTest, ZeroPaddingNhwc) { + RunDepthwiseConvCase(/*channels=*/32, /*in_height=*/8, /*in_width=*/8, /*padding=*/0, /*channels_last=*/true); +} + +TEST(MlasKleidiDepthwiseTest, UnitPaddingNhwc) { + RunDepthwiseConvCase(/*channels=*/32, /*in_height=*/8, /*in_width=*/8, /*padding=*/1, /*channels_last=*/true); +} + +TEST(MlasKleidiDepthwiseTest, MlasConvChannelsLast) { + RunMlasConvChannelsLastCase(/*channels=*/32, /*in_height=*/8, /*in_width=*/8, /*padding=*/1); +} + +#endif // defined(USE_KLEIDIAI) && !defined(_MSC_VER)