Hacks for NCHWC Conv

Rohanjames1997 · Rohanjames1997 · commit b63aac123247 · 2025-12-20T00:11:48.000Z
diff --git a/onnxruntime/core/mlas/lib/mlasi.h b/onnxruntime/core/mlas/lib/mlasi.h
@@ -964,7 +964,9 @@ extern "C" {
 #endif
 #if defined(MLAS_TARGET_ARM64) && defined(MLAS_USE_ARM_NEON_NCHWC)
     MLAS_CONV_FLOAT_KERNEL MlasConvNchwFloatKernelNeon;
+    MLAS_CONV_FLOAT_KERNEL MlasConvNchwBf16KernelNeon;
     MLAS_CONV_FLOAT_KERNEL MlasConvNchwcFloatKernelNeon;
+    MLAS_CONV_FLOAT_KERNEL MlasConvNchwcBf16KernelNeon;
     MLAS_CONV_DEPTHWISE_FLOAT_KERNEL MlasConvDepthwiseFloatKernelNeon;
     MLAS_CONV_DEPTHWISE_FLOAT_KERNEL MlasConvDepthwiseBf16KernelNeon;
     MLAS_CONV_POINTWISE_FLOAT_KERNEL MlasConvPointwiseFloatKernelNeon;
diff --git a/onnxruntime/core/mlas/lib/platform.cpp b/onnxruntime/core/mlas/lib/platform.cpp
@@ -567,7 +567,9 @@ Return Value:
 
 #if defined(MLAS_USE_ARM_NEON_NCHWC)
     this->ConvNchwFloatKernel = MlasConvNchwFloatKernelNeon;
+    this->ConvNchwFloatKernel = MlasConvNchwBf16KernelNeon;
     this->ConvNchwcFloatKernel = MlasConvNchwcFloatKernelNeon;
+    this->ConvNchwcFloatKernel = MlasConvNchwcBf16KernelNeon;
     this->ConvDepthwiseFloatKernel = MlasConvDepthwiseFloatKernelNeon;
     // this->ConvDepthwiseFloatKernel = MlasConvDepthwiseBf16KernelNeon;
     this->ConvPointwiseFloatKernel = MlasConvPointwiseFloatKernelNeon;
diff --git a/onnxruntime/core/mlas/lib/sbconv_kernel_neon.cpp b/onnxruntime/core/mlas/lib/sbconv_kernel_neon.cpp
@@ -24,52 +24,6 @@ Module Name:
 
 constexpr size_t BlockSize = MLAS_PLATFORM::MLAS_NEON_NCHWC_BLOCK_SIZE;
 
-inline void MLASCALL
-MlasRowDot(const float* Arow, const float* Brow, float* out, int index, size_t len)
-{
-    float32x4_t acc4 = MlasBroadcastFloat32x4(0.f);
-
-    size_t i = 0;
-    for (; i + 8 <= len; i += 8) {
-        float32x4_t a0 = MlasLoadFloat32x4(Arow + i);
-        float32x4_t a1 = MlasLoadFloat32x4(Arow + i + 4);
-        float32x4_t b0 = MlasLoadFloat32x4(Brow + i);
-        float32x4_t b1 = MlasLoadFloat32x4(Brow + i + 4);
-
-        bfloat16x8_t a_bf16 = vcvtq_low_bf16_f32(a0);
-        a_bf16 = vcvtq_high_bf16_f32(a_bf16, a1);
-
-        bfloat16x8_t b_bf16 = vcvtq_low_bf16_f32(b0);
-        b_bf16 = vcvtq_high_bf16_f32(b_bf16, b1);
-
-        acc4 = vbfdotq_f32(acc4, a_bf16, b_bf16);
-    }
-
-    float sum = vaddvq_f32(acc4);
-
-    for (; i < len; i++)
-        sum += Arow[i] * Brow[i];
-
-    out[index] = sum;
-}
-
-inline void MLASCALL
-MlasSBDotRowWise(const float* A, const float* B, size_t len, float* out)
-{
-    float tmpA[len];
-    float tmpB[len];
-
-    for (size_t r = 0; r < 16; r++) {
-        for (size_t j = 0; j < len; j++)
-            tmpA[j] = A[j * 16 + r];
-
-        for (size_t j = 0; j < len; j++)
-            tmpB[j] = B[j * 16 + r];
-
-        MlasRowDot(tmpA, tmpB, out, r, len);
-    }
-}
-
 void
     MLASCALL
     MlasConvDepthwiseBf16KernelNeon(
@@ -91,97 +45,251 @@ void
         unsigned KernelFlags
     )
 {
-    const float32x4_t ZeroVector = MlasBroadcastFloat32x4(0.0f);
-    const float32x4_t AccumulateMask = vreinterpretq_f32_s32(MlasBroadcastInt32x4(-(KernelFlags & MLAS_CONV_KERNEL_FLAG_ACCUMULATE_OUTPUT)));
-    const float32x4_t BiasMask = vreinterpretq_f32_s32(MlasBroadcastInt32x4(-(KernelFlags & MLAS_CONV_KERNEL_FLAG_BIAS_ADDITION)));
-    const float32x4_t ReluMask = vreinterpretq_f32_s32(MlasBroadcastInt32x4(-(KernelFlags & MLAS_CONV_KERNEL_FLAG_RELU_ACTIVATION)));
+    const bool AccumulateOutput = KernelFlags & MLAS_CONV_KERNEL_FLAG_ACCUMULATE_OUTPUT;
+    const bool BiasAddition = KernelFlags & MLAS_CONV_KERNEL_FLAG_BIAS_ADDITION;
+    const bool ReluActivation = KernelFlags & MLAS_CONV_KERNEL_FLAG_RELU_ACTIVATION;
 
     const size_t StrideWidthElements = StrideWidth / sizeof(float);
     const size_t DilationWidthElements = DilationWidth / sizeof(float);
     const size_t DilatedInputWidthElements = DilatedInputWidth / sizeof(float);
+    const size_t InputWidthElements = InputWidth / sizeof(float);
+    const size_t TotalOutputCount = OutputCountLeftPad + OutputCount + OutputCountRightPad;
+    const size_t KernelSize = KernelHeight * KernelWidth;
 
     MLAS_UNREFERENCED_PARAMETER(InputStride);
 
+    // Depthwise: 16 independent channels, each doing [TotalOutputCount][KernelSize] x [KernelSize][1]
+    // Batch all 16 channels into one MlasSBGemmBatch call
+    
+    std::vector<float> im2col_buffer(BlockSize * TotalOutputCount * KernelSize);
+    std::vector<float> filter_cols(BlockSize * KernelSize);
+    std::vector<float> output_buffer(BlockSize * TotalOutputCount);
+    
+    // Prepare filter columns: transpose [KernelSize][16] -> 16 separate [KernelSize] vectors
+    for (size_t c = 0; c < BlockSize; c++) {
+        for (size_t k = 0; k < KernelSize; k++) {
+            filter_cols[c * KernelSize + k] = Filter[k * BlockSize + c];
+        }
+    }
+    
+    // im2col for all channels: [c][out_idx][kpos]
+    for (size_t c = 0; c < BlockSize; c++) {
+        for (size_t out_idx = 0; out_idx < TotalOutputCount; out_idx++) {
+            for (size_t kpos = 0; kpos < KernelSize; kpos++) {
+                size_t kh = kpos / KernelWidth;
+                size_t kw = kpos % KernelWidth;
+                const float* input_ptr = Input + out_idx * StrideWidthElements +
+                                          kh * DilatedInputWidthElements + kw * DilationWidthElements + c;
+                const float* row_start = InputBase + kh * DilatedInputWidthElements;
+                const float* row_end = row_start + InputWidthElements;
+                im2col_buffer[c * TotalOutputCount * KernelSize + out_idx * KernelSize + kpos] = 
+                    (input_ptr >= row_start && input_ptr < row_end) ? *input_ptr : 0.0f;
+            }
+        }
+    }
+    
+    // Batched SBGEMM: 16 independent GEMMs, each M=TotalOutputCount, N=1, K=KernelSize
+    MLAS_SBGEMM_DATA_PARAMS params[16];
+    for (size_t c = 0; c < BlockSize; c++) {
+        params[c].A = &im2col_buffer[c * TotalOutputCount * KernelSize];
+        params[c].B = &filter_cols[c * KernelSize];
+        params[c].C = &output_buffer[c * TotalOutputCount];
+        params[c].lda = KernelSize;
+        params[c].ldb = 1;
+        params[c].ldc = 1;
+        params[c].Bias = nullptr;
+        params[c].AIsfp32 = true;
+        params[c].BIsfp32 = true;
+        params[c].ZeroMode = true;
+        params[c].OutputProcessor = nullptr;
+    }
+    MlasSBGemmBatch(TotalOutputCount, 1, KernelSize, BlockSize, params, nullptr);
+    
+    // Scatter results back to output and apply post-processing
+    for (size_t out_idx = 0; out_idx < TotalOutputCount; out_idx++) {
+        float* output_ptr = &Output[out_idx * BlockSize];
+        for (size_t c = 0; c < BlockSize; c++) {
+            float val = output_buffer[c * TotalOutputCount + out_idx];
+            if (AccumulateOutput) val += output_ptr[c];
+            if (BiasAddition) val += Bias[c];
+            if (ReluActivation && val < 0) val = 0;
+            output_ptr[c] = val;
+        }
+    }
+}
+
+//
+// BF16 NCHW/NCHWc Convolution Kernel using im2col + SBGEMM.
+//   NCHW: 1 input channel per kernel position, single GEMM with K=KernelSize
+//   NCHWc: BlockSize input channels per kernel position, loop over kpos with K=BlockSize
+//
+// BF16 NCHW/NCHWc Convolution Kernel using im2col + SBGEMM.
+//   NCHW: 1 input channel per kernel position, single GEMM with K=KernelSize
+//   NCHWc: BlockSize input channels per kernel position, loop over kpos with K=BlockSize
+//
+template <bool IsNchwcFormat>
+void MLASCALL
+MlasConvBf16KernelNeonImpl(
+    const float* Input,
+    const float* Filter,
+    float* Output,
+    size_t StrideWidth,
+    size_t DilationWidth,
+    size_t FilterCount,
+    size_t InputStride,
+    size_t FilterStride,
+    size_t OutputStride,
+    size_t KernelHeight,
+    size_t KernelWidth,
+    const float* InputBase,
+    size_t InputWidth,
+    size_t DilatedInputWidth,
+    size_t OutputCountLeftPad,
+    size_t OutputCount,
+    size_t OutputCountRightPad,
+    const float* Bias,
+    unsigned KernelFlags
+)
+{
+    const bool AccumulateOutput = (KernelFlags & MLAS_CONV_KERNEL_FLAG_ACCUMULATE_OUTPUT) != 0;
+    const bool BiasAddition = (KernelFlags & MLAS_CONV_KERNEL_FLAG_BIAS_ADDITION) != 0;
+    const bool ReluActivation = (KernelFlags & MLAS_CONV_KERNEL_FLAG_RELU_ACTIVATION) != 0;
+
+    const size_t StrideWidthElements = StrideWidth / sizeof(float);
+    const size_t DilationWidthElements = DilationWidth / sizeof(float);
+    const size_t FilterStrideElements = FilterStride / sizeof(float);
+    const size_t OutputStrideElements = OutputStride / sizeof(float);
     const size_t InputWidthElements = InputWidth / sizeof(float);
+    const size_t DilatedInputWidthElements = DilatedInputWidth / sizeof(float);
+
+    MLAS_UNREFERENCED_PARAMETER(InputStride);
 
     const size_t TotalOutputCount = OutputCountLeftPad + OutputCount + OutputCountRightPad;
-    
-    const size_t MaxKernelPositions = KernelHeight * KernelWidth;
-    float tmpInput[MaxKernelPositions * BlockSize];
-
-    for (size_t output_idx = 0; output_idx < TotalOutputCount; output_idx++) {
-        float* output_ptr = &Output[output_idx * BlockSize];
-
-        float32x4_t OldOutput0 = MlasLoadFloat32x4(output_ptr);
-        float32x4_t OldOutput1 = MlasLoadFloat32x4(output_ptr + 4);
-        float32x4_t OldOutput2 = MlasLoadFloat32x4(output_ptr + 8);
-        float32x4_t OldOutput3 = MlasLoadFloat32x4(output_ptr + 12);
-        
-        for (size_t kernel_pos = 0; kernel_pos < MaxKernelPositions; kernel_pos++) {
-            size_t kh = kernel_pos / KernelWidth;
-            size_t kw = kernel_pos % KernelWidth;
-            const float* input_base = Input + output_idx * StrideWidthElements +
+    const size_t KernelSize = KernelHeight * KernelWidth;
+
+    std::vector<float> im2col_buffer(TotalOutputCount * (IsNchwcFormat ? BlockSize : KernelSize));
+
+    if (BiasAddition && AccumulateOutput) {
+        for (size_t f = 0; f < FilterCount; f++) {
+            float* output = Output + f * OutputStrideElements;
+            const float32x4_t b0 = MlasLoadFloat32x4(&Bias[f * BlockSize]);
+            const float32x4_t b1 = MlasLoadFloat32x4(&Bias[f * BlockSize + 4]);
+            const float32x4_t b2 = MlasLoadFloat32x4(&Bias[f * BlockSize + 8]);
+            const float32x4_t b3 = MlasLoadFloat32x4(&Bias[f * BlockSize + 12]);
+            for (size_t i = 0; i < TotalOutputCount; i++) {
+                MlasStoreFloat32x4(&output[i * BlockSize], MlasAddFloat32x4(b0, MlasLoadFloat32x4(&output[i * BlockSize])));
+                MlasStoreFloat32x4(&output[i * BlockSize + 4], MlasAddFloat32x4(b1, MlasLoadFloat32x4(&output[i * BlockSize + 4])));
+                MlasStoreFloat32x4(&output[i * BlockSize + 8], MlasAddFloat32x4(b2, MlasLoadFloat32x4(&output[i * BlockSize + 8])));
+                MlasStoreFloat32x4(&output[i * BlockSize + 12], MlasAddFloat32x4(b3, MlasLoadFloat32x4(&output[i * BlockSize + 12])));
+            }
+        }
+    }
+
+    MLAS_SBGEMM_DATA_PARAMS gemm_params[16];
+    const size_t K = IsNchwcFormat ? BlockSize : KernelSize;
+
+    // Helper lambda for im2col extraction at a kernel position
+    auto extractIm2Col = [&](size_t kpos, float* col_base, size_t col_stride) {
+        size_t kh = kpos / KernelWidth;
+        size_t kw = kpos % KernelWidth;
+        const float* row_start = InputBase + kh * DilatedInputWidthElements;
+        const float* row_end = row_start + InputWidthElements;
+
+        for (size_t out_idx = 0; out_idx < TotalOutputCount; out_idx++) {
+            const float* input_base = Input + out_idx * StrideWidthElements +
                                       kh * DilatedInputWidthElements + kw * DilationWidthElements;
-            const float* row_start = InputBase + kh * DilatedInputWidthElements;
-            const float* row_end = row_start + InputWidthElements;
-
-            bool valid = (input_base >= row_start) && (input_base + 15 < row_end);
-            const float* safe_ptr = input_base;
-            safe_ptr = (input_base < row_start) ? row_start : safe_ptr;
-            safe_ptr = (input_base + 15 >= row_end) ? row_start : safe_ptr;
-            
-            float32x4_t validMask = vreinterpretq_f32_s32(MlasBroadcastInt32x4(valid ? -1 : 0));
-            
-            float32x4_t loaded0 = MlasLoadFloat32x4(safe_ptr);
-            float32x4_t loaded1 = MlasLoadFloat32x4(safe_ptr + 4);
-            float32x4_t loaded2 = MlasLoadFloat32x4(safe_ptr + 8);
-            float32x4_t loaded3 = MlasLoadFloat32x4(safe_ptr + 12);
-            
-            float32x4_t data0 = MlasBlendFloat32x4(ZeroVector, loaded0, validMask);
-            float32x4_t data1 = MlasBlendFloat32x4(ZeroVector, loaded1, validMask);
-            float32x4_t data2 = MlasBlendFloat32x4(ZeroVector, loaded2, validMask);
-            float32x4_t data3 = MlasBlendFloat32x4(ZeroVector, loaded3, validMask);
-            
-            MlasStoreFloat32x4(&tmpInput[kernel_pos * BlockSize], data0);
-            MlasStoreFloat32x4(&tmpInput[kernel_pos * BlockSize + 4], data1);
-            MlasStoreFloat32x4(&tmpInput[kernel_pos * BlockSize + 8], data2);
-            MlasStoreFloat32x4(&tmpInput[kernel_pos * BlockSize + 12], data3);
+            float* col_ptr = col_base + out_idx * col_stride;
+
+            if constexpr (IsNchwcFormat) {
+                for (size_t ic = 0; ic < BlockSize; ic++) {
+                    const float* ie = input_base + ic;
+                    col_ptr[ic] = (ie >= row_start && ie < row_end) ? *ie : 0.0f;
+                }
+            } else {
+                col_ptr[kpos] = (input_base >= row_start && input_base < row_end) ? *input_base : 0.0f;
+            }
         }
+    };
 
-        MlasSBDotRowWise(tmpInput, Filter, MaxKernelPositions, output_ptr);
-
-        float32x4_t Accumulator0 = MlasLoadFloat32x4(output_ptr);
-        float32x4_t Accumulator1 = MlasLoadFloat32x4(output_ptr + 4);
-        float32x4_t Accumulator2 = MlasLoadFloat32x4(output_ptr + 8);
-        float32x4_t Accumulator3 = MlasLoadFloat32x4(output_ptr + 12);
-
-        Accumulator0 = MlasAddFloat32x4(Accumulator0, MlasAndFloat32x4(OldOutput0, AccumulateMask));
-        Accumulator1 = MlasAddFloat32x4(Accumulator1, MlasAndFloat32x4(OldOutput1, AccumulateMask));
-        Accumulator2 = MlasAddFloat32x4(Accumulator2, MlasAndFloat32x4(OldOutput2, AccumulateMask));
-        Accumulator3 = MlasAddFloat32x4(Accumulator3, MlasAndFloat32x4(OldOutput3, AccumulateMask));
-
-        Accumulator0 = MlasAddFloat32x4(Accumulator0, MlasAndFloat32x4(MlasLoadFloat32x4(Bias), BiasMask));
-        Accumulator1 = MlasAddFloat32x4(Accumulator1, MlasAndFloat32x4(MlasLoadFloat32x4(Bias + 4), BiasMask));
-        Accumulator2 = MlasAddFloat32x4(Accumulator2, MlasAndFloat32x4(MlasLoadFloat32x4(Bias + 8), BiasMask));
-        Accumulator3 = MlasAddFloat32x4(Accumulator3, MlasAndFloat32x4(MlasLoadFloat32x4(Bias + 12), BiasMask));
-
-        float32x4_t Relu0 = MlasMaximumFloat32x4(Accumulator0, ZeroVector);
-        float32x4_t Relu1 = MlasMaximumFloat32x4(Accumulator1, ZeroVector);
-        float32x4_t Relu2 = MlasMaximumFloat32x4(Accumulator2, ZeroVector);
-        float32x4_t Relu3 = MlasMaximumFloat32x4(Accumulator3, ZeroVector);
-
-        Accumulator0 = MlasBlendFloat32x4(Accumulator0, Relu0, ReluMask);
-        Accumulator1 = MlasBlendFloat32x4(Accumulator1, Relu1, ReluMask);
-        Accumulator2 = MlasBlendFloat32x4(Accumulator2, Relu2, ReluMask);
-        Accumulator3 = MlasBlendFloat32x4(Accumulator3, Relu3, ReluMask);
-
-        MlasStoreFloat32x4(output_ptr, Accumulator0);
-        MlasStoreFloat32x4(output_ptr + 4, Accumulator1);
-        MlasStoreFloat32x4(output_ptr + 8, Accumulator2);
-        MlasStoreFloat32x4(output_ptr + 12, Accumulator3);
+    // Helper lambda to setup GEMM params
+    auto setupGemmParams = [&](size_t filter_offset, bool zeroMode) {
+        size_t idx = 0;
+        for (size_t f = 0; f < FilterCount; f++) {
+            gemm_params[idx].A = im2col_buffer.data();
+            gemm_params[idx].B = Filter + f * FilterStrideElements + filter_offset;
+            gemm_params[idx].C = Output + f * OutputStrideElements;
+            gemm_params[idx].lda = K;
+            gemm_params[idx].ldb = BlockSize;
+            gemm_params[idx].ldc = BlockSize;
+            gemm_params[idx].Bias = BiasAddition ? (Bias + f * BlockSize) : nullptr;
+            gemm_params[idx].AIsfp32 = true;
+            gemm_params[idx].BIsfp32 = true;
+            gemm_params[idx].ZeroMode = zeroMode;
+            gemm_params[idx].OutputProcessor = nullptr;
+            idx++;
+        }
+        return idx;
+    };
+
+    const size_t numGemmCalls = IsNchwcFormat ? KernelSize : 1;
+    for (size_t g = 0; g < numGemmCalls; g++) {
+        if constexpr (IsNchwcFormat) {
+            extractIm2Col(g, im2col_buffer.data(), BlockSize);
+        } else {
+            for (size_t kpos = 0; kpos < KernelSize; kpos++) {
+                extractIm2Col(kpos, im2col_buffer.data(), KernelSize);
+            }
+        }
+        size_t kh = g / KernelWidth, kw = g % KernelWidth;
+        size_t filter_offset = IsNchwcFormat ? kh * (KernelWidth * BlockSize * BlockSize) + kw * (BlockSize * BlockSize) : 0;
+        size_t idx = setupGemmParams(filter_offset, (g == 0) && !AccumulateOutput);
+        MlasSBGemmBatch(TotalOutputCount, BlockSize, K, idx, gemm_params, nullptr);
+    }
+
+    if (ReluActivation) {
+        const float32x4_t ZeroVector = MlasBroadcastFloat32x4(0.0f);
+        for (size_t f = 0; f < FilterCount; f++) {
+            float* output = Output + f * OutputStrideElements;
+            for (size_t i = 0; i < TotalOutputCount; i++) {
+                MlasStoreFloat32x4(&output[i * BlockSize], MlasMaximumFloat32x4(MlasLoadFloat32x4(&output[i * BlockSize]), ZeroVector));
+                MlasStoreFloat32x4(&output[i * BlockSize + 4], MlasMaximumFloat32x4(MlasLoadFloat32x4(&output[i * BlockSize + 4]), ZeroVector));
+                MlasStoreFloat32x4(&output[i * BlockSize + 8], MlasMaximumFloat32x4(MlasLoadFloat32x4(&output[i * BlockSize + 8]), ZeroVector));
+                MlasStoreFloat32x4(&output[i * BlockSize + 12], MlasMaximumFloat32x4(MlasLoadFloat32x4(&output[i * BlockSize + 12]), ZeroVector));
+            }
+        }
     }
 }
 
+void MLASCALL MlasConvNchwcBf16KernelNeon(
+    const float* Input, const float* Filter, float* Output,
+    size_t StrideWidth, size_t DilationWidth, size_t FilterCount,
+    size_t InputStride, size_t FilterStride, size_t OutputStride,
+    size_t KernelHeight, size_t KernelWidth, const float* InputBase,
+    size_t InputWidth, size_t DilatedInputWidth,
+    size_t OutputCountLeftPad, size_t OutputCount, size_t OutputCountRightPad,
+    const float* Bias, unsigned KernelFlags)
+{
+    MlasConvBf16KernelNeonImpl<true>(Input, Filter, Output, StrideWidth, DilationWidth,
+        FilterCount, InputStride, FilterStride, OutputStride, KernelHeight, KernelWidth,
+        InputBase, InputWidth, DilatedInputWidth, OutputCountLeftPad, OutputCount,
+        OutputCountRightPad, Bias, KernelFlags);
+}
+
+void MLASCALL MlasConvNchwBf16KernelNeon(
+    const float* Input, const float* Filter, float* Output,
+    size_t StrideWidth, size_t DilationWidth, size_t FilterCount,
+    size_t InputStride, size_t FilterStride, size_t OutputStride,
+    size_t KernelHeight, size_t KernelWidth, const float* InputBase,
+    size_t InputWidth, size_t DilatedInputWidth,
+    size_t OutputCountLeftPad, size_t OutputCount, size_t OutputCountRightPad,
+    const float* Bias, unsigned KernelFlags)
+{
+    MlasConvBf16KernelNeonImpl<false>(Input, Filter, Output, StrideWidth, DilationWidth,
+        FilterCount, InputStride, FilterStride, OutputStride, KernelHeight, KernelWidth,
+        InputBase, InputWidth, DilatedInputWidth, OutputCountLeftPad, OutputCount,
+        OutputCountRightPad, Bias, KernelFlags);
+}
+
 //
 // BF16 Pointwise (1x1) Convolution Kernel using SBGEMM.
 //
