Skip to content

mlas/arm64: add BF16 fast-math conv kernels for NCHW/NCHWc paths #27878

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
hariharans29 merged 11 commits into from
Apr 9, 2026
Merged

mlas/arm64: add BF16 fast-math conv kernels for NCHW/NCHWc paths #27878

Show file tree
Hide file tree
Changes from 8 commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
b1a32c9
mlas/arm64: add BF16 fast-math conv kernels for NCHW/NCHWc paths
milpuz01 Mar 26, 2026
5138c56
Refactor BF16 code to be solely in sbconv_kernel_neon.cpp file
milpuz01 Mar 28, 2026
286f6be
Fix NEON conv flag masks for bias and ReLU post-processing
milpuz01 Mar 28, 2026
5e0315a
Guard BF16 conv helpers to avoid unused-function warnings
milpuz01 Mar 28, 2026
cfb7109
Add BF16 direct/depthwise conv coverage and guard direct 3x3 edge rows
milpuz01 Mar 31, 2026
68108bd
Only build BF16 conv asm when the toolchain supports BF16
milpuz01 Mar 31, 2026
8a69461
Centralize BF16 conv runtime guards and wrapper gating
milpuz01 Mar 31, 2026
27f4467
Inline Arm Linux NCHWC guards for BF16 conv dispatch
milpuz01 Apr 1, 2026
7abe8ba
Probe BF16 compiler support and relax BF16 conv guards
milpuz01 Apr 7, 2026
da3044a
Drop redundant BF16 feature guard from pointwise asm
milpuz01 Apr 7, 2026
6bd0bfb
Simplify BF16 conv source wiring and Linux wrapper guards
milpuz01 Apr 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 37 additions & 4 deletions cmake/onnxruntime_mlas.cmake
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -349,7 +349,14 @@ function (setup_arm_neon_nchwc)
${MLAS_SRC_DIR}/aarch64/SconvNchwcKernelNeon.S
${MLAS_SRC_DIR}/aarch64/SconvDepthwiseKernelNeon.S
${MLAS_SRC_DIR}/aarch64/SconvPointwiseKernelNeon.S
)
)
if(MLAS_ARM64_BFLOAT16_COMPILE_FLAGS)
target_sources(onnxruntime_mlas PRIVATE
${MLAS_SRC_DIR}/aarch64/SconvDepthwiseKernelNeonBf16.S
${MLAS_SRC_DIR}/aarch64/SconvKernelNeonBf16.S
${MLAS_SRC_DIR}/aarch64/SconvPointwiseKernelNeonBf16.S
)
Comment thread
milpuz01 marked this conversation as resolved.
Outdated
endif()
endif()
list(APPEND mlas_private_compile_definitions MLAS_USE_ARM_NEON_NCHWC)
set(mlas_private_compile_definitions ${mlas_private_compile_definitions} PARENT_SCOPE)
Expand Down Expand Up @@ -468,6 +475,24 @@ else()
endif()
if(ARM64 AND MLAS_SOURCE_IS_NOT_SET )
enable_language(ASM)
set(MLAS_ARM64_BFLOAT16_COMPILE_FLAGS "")
if(APPLE)
set(MLAS_OLD_REQUIRED_FLAGS "${CMAKE_REQUIRED_FLAGS}")
set(CMAKE_REQUIRED_FLAGS "${CMAKE_REQUIRED_FLAGS} -arch arm64")
check_cxx_compiler_flag("-march=armv8.2-a+bf16" HAS_ARM64_BFLOAT16_MARCH_APPLE)
if(HAS_ARM64_BFLOAT16_MARCH_APPLE)
set(MLAS_ARM64_BFLOAT16_COMPILE_FLAGS " -march=armv8.2-a+bf16 ")
else()
check_cxx_compiler_flag("-Xclang -target-feature -Xclang +bf16" HAS_ARM64_BFLOAT16_TARGET_FEATURE_APPLE)
if(HAS_ARM64_BFLOAT16_TARGET_FEATURE_APPLE)
set(MLAS_ARM64_BFLOAT16_COMPILE_FLAGS " -Xclang -target-feature -Xclang +bf16 ")
endif()
endif()
set(CMAKE_REQUIRED_FLAGS "${MLAS_OLD_REQUIRED_FLAGS}")
unset(MLAS_OLD_REQUIRED_FLAGS)
elseif(HAS_ARM64_BFLOAT16)
set(MLAS_ARM64_BFLOAT16_COMPILE_FLAGS " -march=armv8.2-a+bf16 ")
Comment thread
milpuz01 marked this conversation as resolved.
Outdated
endif()
set(mlas_platform_srcs
${MLAS_SRC_DIR}/aarch64/ConvSymS8KernelDot.S
${MLAS_SRC_DIR}/aarch64/ConvSymS8KernelDotLd64.S
Expand Down Expand Up @@ -550,12 +575,9 @@ else()
set_source_files_properties(${MLAS_SRC_DIR}/aarch64/HalfGemmKernelNeon.S PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+fp16 ")
set_source_files_properties(${MLAS_SRC_DIR}/aarch64/QgemmS8S8KernelSmmla.S PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+i8mm ")
set_source_files_properties(${MLAS_SRC_DIR}/aarch64/QgemmU8X8KernelUmmla.S PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+i8mm ")
set_source_files_properties(${MLAS_SRC_DIR}/aarch64/SbgemmKernelNeon.S PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+bf16 ")
set_source_files_properties(${MLAS_SRC_DIR}/activate_fp16.cpp PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+fp16 ")
set_source_files_properties(${MLAS_SRC_DIR}/dwconv.cpp PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+fp16 ")
set_source_files_properties(${MLAS_SRC_DIR}/pooling_fp16.cpp PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+fp16 ")
set_source_files_properties(${MLAS_SRC_DIR}/sbgemm_kernel_neon.cpp PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+bf16 ")
set_source_files_properties(${MLAS_SRC_DIR}/sbconv_kernel_neon.cpp PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+bf16 ")
set_source_files_properties(${MLAS_SRC_DIR}/cast_kernel_neon.cpp PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+fp16 ")
set_source_files_properties(${MLAS_SRC_DIR}/hqnbitgemm_kernel_neon_fp16.cpp PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+fp16 ")
set_source_files_properties(${MLAS_SRC_DIR}/hqnbitgemm_kernel_neon_fp16_8bit.cpp PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+fp16 ")
Expand All @@ -565,6 +587,17 @@ else()
set_source_files_properties(${MLAS_SRC_DIR}/eltwise_kernel_neon_fp16.cpp PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+fp16 ")
endif()

if(MLAS_ARM64_BFLOAT16_COMPILE_FLAGS)
set_source_files_properties(${MLAS_SRC_DIR}/aarch64/SconvDepthwiseKernelNeonBf16.S PROPERTIES COMPILE_FLAGS "${MLAS_ARM64_BFLOAT16_COMPILE_FLAGS}")
set_source_files_properties(${MLAS_SRC_DIR}/aarch64/SconvKernelNeonBf16.S PROPERTIES COMPILE_FLAGS "${MLAS_ARM64_BFLOAT16_COMPILE_FLAGS}")
set_source_files_properties(${MLAS_SRC_DIR}/aarch64/SconvPointwiseKernelNeonBf16.S PROPERTIES COMPILE_FLAGS "${MLAS_ARM64_BFLOAT16_COMPILE_FLAGS}")
if(NOT APPLE)
set_source_files_properties(${MLAS_SRC_DIR}/aarch64/SbgemmKernelNeon.S PROPERTIES COMPILE_FLAGS "${MLAS_ARM64_BFLOAT16_COMPILE_FLAGS}")
set_source_files_properties(${MLAS_SRC_DIR}/sbgemm_kernel_neon.cpp PROPERTIES COMPILE_FLAGS "${MLAS_ARM64_BFLOAT16_COMPILE_FLAGS}")
set_source_files_properties(${MLAS_SRC_DIR}/sbconv_kernel_neon.cpp PROPERTIES COMPILE_FLAGS "${MLAS_ARM64_BFLOAT16_COMPILE_FLAGS}")
endif()
Comment thread
milpuz01 marked this conversation as resolved.
endif()

if(ONNXRUNTIME_MLAS_MULTI_ARCH)
onnxruntime_add_static_library(onnxruntime_mlas_arm64 ${mlas_platform_srcs})
set_target_properties(onnxruntime_mlas_arm64 PROPERTIES OSX_ARCHITECTURES "arm64")
Expand Down
Loading
Loading