Skip to content

Optimize Arm bf16 qb4w fully connected conversion path - #11112

Open
JakeStevens wants to merge 2 commits into
google:masterfrom
JakeStevens:arm-bf16-qb4w-fully-connected
Open

Optimize Arm bf16 qb4w fully connected conversion path#11112
JakeStevens wants to merge 2 commits into
google:masterfrom
JakeStevens:arm-bf16-qb4w-fully-connected

Conversation

@JakeStevens

Copy link
Copy Markdown
Contributor

Two related changes:

Keep qd8_bf16_qb4w fully connected native on Arm only when an optimized i8mm or NEON dot-product kernel is available and the input, output, and weight scales use the supported BF16 types. Otherwise, retain the fp32 rewrite fallback.

Add NEON v8 BF16-to-QS8 conversion microkernels and select the u32 kernel for Arm packed-LHS conversion. This replaces the scalar-imagic conversion path on Arm while preserving scalar and x86 fallback behavior.

Test plan:

      cmake --build build-test --target bf16-qs8-vcvt-test -j8
      ./build-test/test/bf16-qs8-vcvt-test  # 64/64 pass
      cmake --build build-android-arm64 --target bf16-qs8-vcvt-test -j8

Benchmark (Pixel 10, Android 17/API 37, arm64-v8a):
Configured a Release build for Android API 24 with NDK r29/Clang 21.

      cmake --build build-android-arm64 --target vunary-bench -j8
      adb push build-android-arm64/bench/vunary-bench /data/local/tmp/
      adb shell /data/local/tmp/vunary-bench --benchmark_filter='bf16_qs8_vcvt.*(neonv8_u32|scalar_imagic_u4)'

Results (real time, CPU scaling and ASLR enabled):

        N=8192:  neonv8_u32 1110 ns, 22.1348 GB/s; scalar_imagic_u4 1182 ns, 20.7856 GB/s
        N=65536: neonv8_u32 8886 ns, 22.1246 GB/s; scalar_imagic_u4 9455 ns, 20.7941 GB/s
        neonv8_u32 was 1.064-1.065x faster than scalar_imagic_u4.

JakeStevens and others added 2 commits September 1, 2026 12:19
Keep qd8_bf16_qb4w fully connected native on Arm only when an optimized i8mm or NEON dot-product kernel is available and the input, output, and weight scales use the supported BF16 types. Otherwise, retain the fp32 rewrite fallback.

Add NEON v8 BF16-to-QS8 conversion microkernels and select the u32 kernel for Arm packed-LHS conversion. This replaces the scalar-imagic conversion path on Arm while preserving scalar and x86 fallback behavior.

Test plan:
  cmake --build build-test --target bf16-qs8-vcvt-test -j8
  ./build-test/test/bf16-qs8-vcvt-test  # 64/64 pass
  cmake --build build-android-arm64 --target bf16-qs8-vcvt-test -j8

Benchmark (Pixel 10, Android 17/API 37, arm64-v8a):
  Configured a Release build for Android API 24 with NDK r29/Clang 21.
  cmake --build build-android-arm64 --target vunary-bench -j8
  adb push build-android-arm64/bench/vunary-bench /data/local/tmp/
  adb shell /data/local/tmp/vunary-bench --benchmark_filter='bf16_qs8_vcvt.*(neonv8_u32|scalar_imagic_u4)'

  Results (real time, CPU scaling and ASLR enabled):
    N=8192:  neonv8_u32 1110 ns, 22.1348 GB/s; scalar_imagic_u4 1182 ns, 20.7856 GB/s
    N=65536: neonv8_u32 8886 ns, 22.1246 GB/s; scalar_imagic_u4 9455 ns, 20.7941 GB/s
    neonv8_u32 was 1.064-1.065x faster than scalar_imagic_u4.
std::vector<uint16_t> weights_scale = {0x3F80, 0x3F80};
std::vector<uint8_t> weights_data = std::vector<uint8_t>(32, 0x88);

void Build(enum xnn_datatype input_datatype,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm wondering how difficult it would be to use the Tensor API to build this test graph like was done in

std::unique_ptr<XnnpackGraph> graph;
{
XnnTensor input({.type = Type::kFP16, .shape = {3, 4}});
XnnTensor weights({.type = Type::kFP16,
.shape = {2, 4},
.buffer = OwningCpuBuffer::Copy<Type::kFP16>(
{1, 2, 3, 4, 5, 6, 7, 8})});
XnnTensor output = FullyConnected(input, weights);
LRT_TENSOR_ASSERT_OK_AND_ASSIGN(graph, BuildXnnpackGraph({output}));
}
std::unique_ptr<XnnpackGraph> expected_graph;
{
XnnTensor input({.type = Type::kFP16, .shape = {3, 4}});
input = Cast(input, Type::kFP32);
XnnTensor weights({.type = Type::kFP32,
.shape = {2, 4},
.buffer = OwningCpuBuffer::Copy<Type::kFP32>(
{1, 2, 3, 4, 5, 6, 7, 8})});
XnnTensor output = FullyConnected(input, weights);
output = Cast(output, Type::kFP16);
LRT_TENSOR_ASSERT_OK_AND_ASSIGN(expected_graph,
BuildXnnpackGraph({output}));
}
xnn_subgraph_t subgraph = graph->subgraph();
ASSERT_THAT(xnn_subgraph_fallback_from_fp16_to_fp32(subgraph,
/*optimization_flags=*/0),
Eq(xnn_status_success));
EXPECT_THAT(graph, IsIsomorphicTo(expected_graph));

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants