Skip to content

Commit efc89b5

Browse files
authored
mlas/arm64: add BF16 fast-math conv kernels for NCHW/NCHWc paths (#27878)
### Description Add Arm64 BF16 fast-math convolution support in MLAS: - direct NCHW conv - depthwise 3x3 NCHWc conv - pointwise 1x1 NCHWc conv This change adds new AArch64 BF16 asm kernels, wires them into MLAS platform dispatch, keeps accumulated pointwise batches on the custom BF16 path instead of falling back to generic SBGEMM, and adds the required BF16 build flags. The new paths are only used when Arm64 BF16 fast-math is enabled via the existing session option. Baseline FP32 behavior is unchanged. ### Performance Individual convolution improvements when running on `c8g` AWS instance where in columns base is FP32 execution, fast-math when enabled without this PR and PR is fast-math with this change: | Type | Shape | fast-math vs base | PR w/ fast-math vs base | PR w/ fast-math vs fast-math | |---|---|---:|---:|---:| | depthwise | N1 IC32 OC32 H112xW112->112x112 K3x3 S1x1 D1x1 P1/1/1/1 G32 | 0.991x | 1.047x | 1.057x | | depthwise | N1 IC96 OC96 H112xW112->56x56 K3x3 S2x2 D1x1 P1/1/1/1 G96 | 1.015x | 1.015x | 1.000x | | depthwise | N1 IC144 OC144 H56xW56->28x28 K3x3 S2x2 D1x1 P1/1/1/1 G144 | 1.020x | 1.004x | 0.984x | | depthwise | N1 IC144 OC144 H56xW56->56x56 K3x3 S1x1 D1x1 P1/1/1/1 G144 | 1.034x | 1.138x | 1.101x | | depthwise | N1 IC192 OC192 H28xW28->28x28 K3x3 S1x1 D1x1 P1/1/1/1 G192 | 0.997x | 1.033x | 1.037x | | depthwise | N1 IC384 OC384 H28xW28->14x14 K3x3 S2x2 D1x1 P1/1/1/1 G384 | 1.016x | 1.021x | 1.005x | | depthwise | N1 IC384 OC384 H28xW28->28x28 K3x3 S1x1 D1x1 P1/1/1/1 G384 | 1.011x | 1.090x | 1.077x | | depthwise | N1 IC576 OC576 H14xW14->7x7 K3x3 S2x2 D1x1 P1/1/1/1 G576 | 1.029x | 0.995x | 0.967x | | depthwise | N1 IC576 OC576 H14xW14->14x14 K3x3 S1x1 D1x1 P1/1/1/1 G576 | 1.025x | 1.006x | 0.982x | | depthwise | N1 IC960 OC960 H7xW7->7x7 K3x3 S1x1 D1x1 P1/1/1/1 G960 | 1.002x | 0.941x | 0.939x | | nchw | N1 IC3 OC32 H224xW224->112x112 K3x3 S2x2 D1x1 P1/1/1/1 G1 | 1.001x | 1.058x | 1.058x | | pointwise | N1 IC16 OC96 H112xW112->112x112 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.213x | 1.328x | 1.095x | | pointwise | N1 IC32 OC16 H112xW112->112x112 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.020x | 1.019x | 0.998x | | pointwise | N1 IC32 OC32 H112xW112->112x112 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.118x | 1.196x | 1.069x | | pointwise | N1 IC32 OC144 H56xW56->56x56 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.220x | 1.528x | 1.252x | | pointwise | N1 IC32 OC192 H28xW28->28x28 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.199x | 1.418x | 1.183x | | pointwise | N1 IC64 OC384 H28xW28->28x28 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.294x | 1.938x | 1.497x | | pointwise | N1 IC96 OC32 H56xW56->56x56 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.080x | 1.426x | 1.320x | | pointwise | N1 IC96 OC576 H14xW14->14x14 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.280x | 1.961x | 1.532x | | pointwise | N1 IC144 OC32 H28xW28->28x28 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.132x | 1.351x | 1.193x | | pointwise | N1 IC144 OC32 H56xW56->56x56 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.073x | 1.374x | 1.281x | | pointwise | N1 IC160 OC960 H7xW7->7x7 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.133x | 1.744x | 1.539x | | pointwise | N1 IC192 OC32 H28xW28->28x28 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.166x | 1.411x | 1.210x | | pointwise | N1 IC192 OC64 H28xW28->28x28 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.212x | 1.763x | 1.454x | | pointwise | N1 IC320 OC1280 H7xW7->7x7 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.136x | 2.059x | 1.812x | | pointwise | N1 IC384 OC64 H28xW28->28x28 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.256x | 1.904x | 1.516x | | pointwise | N1 IC384 OC96 H14xW14->14x14 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.206x | 1.929x | 1.600x | | pointwise | N1 IC576 OC96 H14xW14->14x14 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.250x | 2.055x | 1.644x | | pointwise | N1 IC576 OC160 H7xW7->7x7 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 0.902x | 1.423x | 1.577x | | pointwise | N1 IC960 OC160 H7xW7->7x7 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 0.915x | 1.527x | 1.668x | | pointwise | N1 IC960 OC320 H7xW7->7x7 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.020x | 1.756x | 1.723x | | pointwise | N1 IC1280 OC1008 H1xW1->1x1 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 0.747x | 1.149x | 1.538x | When running the full models the performance improvements are on `c8g` (AWS Graviton 4) and `Standard_D32plds_v6` (Azure Cobalt-100) when running [MobileNet v2.7](https://github.com/onnx/models/blob/main/validated/vision/classification/mobilenet/model/mobilenetv2-7.onnx) with 8 threads are: | Instance | PR w/ fast-math vs base | PR w/ fast-math vs fast-mat | |---|---|---| `c8g` | 1.892x | 1.647x | `Standard_D32plds_v6` | 2.884x | 1.692x | (cc: @Rohanjames1997 @snadampal) --------- Signed-off-by: Milos Puzovic <milos.puzovic@arm.com>
1 parent 3ac6040 commit efc89b5

10 files changed

Lines changed: 4550 additions & 33 deletions

File tree

cmake/onnxruntime_mlas.cmake

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -351,7 +351,7 @@ function (setup_arm_neon_nchwc)
351351
${MLAS_SRC_DIR}/aarch64/SconvNchwcKernelNeon.S
352352
${MLAS_SRC_DIR}/aarch64/SconvDepthwiseKernelNeon.S
353353
${MLAS_SRC_DIR}/aarch64/SconvPointwiseKernelNeon.S
354-
)
354+
)
355355
endif()
356356
list(APPEND mlas_private_compile_definitions MLAS_USE_ARM_NEON_NCHWC)
357357
set(mlas_private_compile_definitions ${mlas_private_compile_definitions} PARENT_SCOPE)
@@ -549,6 +549,13 @@ else()
549549
${MLAS_SRC_DIR}/softmax_kernel_neon_fp16.cpp
550550
${MLAS_SRC_DIR}/eltwise_kernel_neon_fp16.cpp
551551
)
552+
if (onnxruntime_USE_ARM_NEON_NCHWC)
553+
list(APPEND mlas_platform_srcs
554+
${MLAS_SRC_DIR}/aarch64/SconvDepthwiseKernelNeonBf16.S
555+
${MLAS_SRC_DIR}/aarch64/SconvKernelNeonBf16.S
556+
${MLAS_SRC_DIR}/aarch64/SconvPointwiseKernelNeonBf16.S
557+
)
558+
endif()
552559
set_source_files_properties(${MLAS_SRC_DIR}/aarch64/HalfGemmKernelNeon.S PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+fp16 ")
553560
set_source_files_properties(${MLAS_SRC_DIR}/aarch64/QgemmS8S8KernelSmmla.S PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+i8mm ")
554561
set_source_files_properties(${MLAS_SRC_DIR}/aarch64/QgemmU8X8KernelUmmla.S PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+i8mm ")
@@ -565,6 +572,11 @@ else()
565572
set_source_files_properties(${MLAS_SRC_DIR}/halfgemm_kernel_neon_fp16.cpp PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+fp16 ")
566573
set_source_files_properties(${MLAS_SRC_DIR}/softmax_kernel_neon_fp16.cpp PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+fp16 ")
567574
set_source_files_properties(${MLAS_SRC_DIR}/eltwise_kernel_neon_fp16.cpp PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+fp16 ")
575+
if (onnxruntime_USE_ARM_NEON_NCHWC)
576+
set_source_files_properties(${MLAS_SRC_DIR}/aarch64/SconvDepthwiseKernelNeonBf16.S PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+bf16 ")
577+
set_source_files_properties(${MLAS_SRC_DIR}/aarch64/SconvKernelNeonBf16.S PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+bf16 ")
578+
set_source_files_properties(${MLAS_SRC_DIR}/aarch64/SconvPointwiseKernelNeonBf16.S PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+bf16 ")
579+
endif()
568580
endif()
569581

570582
if(ONNXRUNTIME_MLAS_MULTI_ARCH)

0 commit comments

Comments
 (0)