Skip to content

cpu: aarch64: add f16 and bf16 support to jit_io_helper and jit_uni_binary#5092

Open
Sqvid wants to merge 2 commits into
mainfrom
sqvid/jit-io-halfword
Open

cpu: aarch64: add f16 and bf16 support to jit_io_helper and jit_uni_binary#5092
Sqvid wants to merge 2 commits into
mainfrom
sqvid/jit-io-halfword

Conversation

@Sqvid
Copy link
Copy Markdown
Contributor

@Sqvid Sqvid commented Apr 29, 2026

Description

This PR adds f16 and bf16 support to jit_io_helper_t. I have also demonstrated the impact of these changes by adding support for these data-types to jit_uni_binary_t.

# Dispatch before:
$ ./build/tests/benchdnn/benchdnn --binary --batch=test_binary_{float16,bfloat16} --batch=harness_binary_{f16,bf16}
...
============================================================
= Implementation statistics (--summary=no-impl to disable) =
============================================================
| ref:any : 8900 (100%)                                    |
============================================================
tests:20000 passed:8872 skipped:11100 mistrusted:28 unimplemented:0 invalid_arguments:0 failed:0 listed:0

# Dispatch after:
$ ./build/tests/benchdnn/benchdnn --binary --batch=test_binary_{float16,bfloat16} --batch=harness_binary_{f16,bf16}
...
============================================================
= Implementation statistics (--summary=no-impl to disable) =
============================================================
| ref:any : 6240 (54%)                                     |
| jit:uni : 5320 (46%)                                     |
============================================================
tests:20000 passed:11520 skipped:8440 mistrusted:40 unimplemented:0 invalid_arguments:0 failed:0 listed:0

Checklist

General

  • Do all unit and benchdnn tests (make test and make test_benchdnn_*) pass locally for each commit?
  • Have you formatted the code using clang-format?

@Sqvid Sqvid requested a review from a team as a code owner April 29, 2026 10:29
@github-actions github-actions Bot added the platform:cpu-aarch64 Codeowner: @oneapi-src/onednn-cpu-aarch64 label Apr 29, 2026
@Sqvid Sqvid force-pushed the sqvid/jit-io-halfword branch from 1ee6669 to b77feda Compare April 29, 2026 10:43
@Sqvid Sqvid marked this pull request as draft April 29, 2026 10:59
@Sqvid Sqvid marked this pull request as ready for review April 29, 2026 12:23
Sqvid added 2 commits May 5, 2026 09:07
Signed-off-by: Siddhartha Menon <siddhartha.menon@arm.com>
Signed-off-by: Siddhartha Menon <siddhartha.menon@arm.com>
@Sqvid Sqvid force-pushed the sqvid/jit-io-halfword branch from b77feda to a9b1bf0 Compare May 5, 2026 08:08
Comment thread src/cpu/aarch64/utils/jit_io_helper.cpp
@Sqvid Sqvid requested a review from michalowski-arm May 7, 2026 14:29
Copy link
Copy Markdown
Contributor

@michalowski-arm michalowski-arm left a comment

Choose a reason for hiding this comment

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

LGTM

return;
}

host_->str(host_->z31,
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

How do we know we can always use z31? Also, I think a comment explaining that this is a fallback would be useful.

}
case data_type::bf16: {
if (dst_vmm.getIdx() != src_vmm.getIdx())
host_->mov(dst_vmm.s, host_->P_ALL_ONE / Xbyak_aarch64::T_m,
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Can we get rid of this mov with

host_->lsl(dst_vmm.s, src_vmm.s, 16);

If not, a comment explaining why would be useful

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

Labels

platform:cpu-aarch64 Codeowner: @oneapi-src/onednn-cpu-aarch64

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants