microsoft · JonathanC-ARM · Jun 30, 2026
diff --git 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
@@ -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
@@ -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
@@ -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 <typename UkernelFn>
-struct KaiMatmulKernel {
+struct KaiKernel {
     const char* name;
     UkernelFn ukernel;
 };
 
+template <typename UkernelFn>
+using KaiMatmulKernel = KaiKernel<UkernelFn>;
+
 // Wrapper for FP32 GEMM kernels where both LHS and RHS are pre-packed (common SGEMM path).
 using KaiF32SgemmKernel = KaiMatmulKernel<kai_matmul_clamp_f32_f32p_f32p_ukernel>;
 
@@ -43,6 +48,9 @@ using KaiF32IMatmulKernel = KaiMatmulKernel<kai_imatmul_clamp_f32_f32p_f32p_uker
 
 using KaiBF16SBgemmKernel = KaiMatmulKernel<kai_matmul_clamp_f32_bf16p_bf16p_ukernel>;
 
+// Wrapper for FP32 planar depthwise convolution kernels.
+using KaiF32DepthwiseConvKernel = KaiKernel<kai_dwconv_clamp_f32_f32_f32p_planar_ukernel>;
+
 // 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();
+const KaiBF16SBgemmKernel& GetKleidiAISBGemmUKernel();
+
+// Returns the selected FP32 depthwise convolution ukernel based on runtime CPU capabilities.
+const KaiF32DepthwiseConvKernel& GetKleidiAIDepthwiseConvUKernel();
diff --git 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<float>::max(),
+                                  std::numeric_limits<float>::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<const std::byte*>(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<const std::byte*>(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;
 }