[Codegen][CPU] Fill in the bf16 and i8 ukernel bodies + e2e tests.#24572
Open
bjacob wants to merge 1 commit into
Open
[Codegen][CPU] Fill in the bf16 and i8 ukernel bodies + e2e tests.#24572bjacob wants to merge 1 commit into
bjacob wants to merge 1 commit into
Conversation
0907e85 to
121f292
Compare
8af2850 to
cb00012
Compare
121f292 to
d829f18
Compare
cb00012 to
6cd0524
Compare
d829f18 to
6dbb0eb
Compare
6cd0524 to
fc522fd
Compare
6dbb0eb to
a5ed591
Compare
fc522fd to
5fb2905
Compare
a5ed591 to
e72c8e2
Compare
5fb2905 to
6ec426f
Compare
e72c8e2 to
bdc6f02
Compare
6ec426f to
6602e64
Compare
bdc6f02 to
e576a0c
Compare
18678c3 to
a9f88d9
Compare
e576a0c to
2209e59
Compare
2209e59 to
189618e
Compare
a9f88d9 to
b60f897
Compare
189618e to
7422d39
Compare
b60f897 to
6c0ae40
Compare
Replaces the bf16 and i8-VNNI seeds' stub bodies with real SIMD
implementations, both generic over the `intrinsics_{m,n,k}` unrolling
factors and structured like the AMDGPU C ukernel they were adapted from
(`iree_uk_amdgpu_multi_mma_mfma_i32_16x16x32_i8`): accumulators in
registers, an outer loop over the K tiles (`k_outer`), and inside it the
`(intrinsics_m, intrinsics_n, intrinsics_k)` unroll. The `intrinsics_*`
arrive as constants at the inlined call site, so the loops fully unroll
and the `acc_regs` arrays become fixed register files -- the bitcode-LTO
equivalent of a C++ template, as the README describes.
- bf16 (`MMA_X86_AVX512BF16_1x16x2_F32_BF16`): one `_mm512_dpbf16_ps`
per (m, n, k), with the LHS K-pair broadcast via `set1_ps`.
- i8 (`MMA_X86_AVX512VNNI_16x16x2_I32_I8_CASTI16`): the 16x16x2 tile is
bit-compatible with the codegen path `lowerX86Avx512Vnni16x16x2I8` --
one `vpmovsxbw` widen of each i8 panel to i16, the `vpshufd` /
`vbroadcasti32x4` fan-out, and 16 `vpdpwssd` over the block-interleaved
(rlo, chi, rhi, clo) ACC layout. The i8 ukernel needs `-mavx512bw` for
the widen, so it is added to the VNNI copts.
`LLVMCPUSelectUKernels` now only selects a ukernel when its bitcode
actually exists (via `attachUKernelBitcodeOnOp`'s bool return), so an
`MMAIntrinsic` the cost model picks but for which no seed exists -- e.g.
the M<->N-swapped `MMA_X86_AVX512BF16_16x1x2_F32_BF16` -- falls back to
codegen instead of dangling an undefined symbol.
Adds two execution/numerical tests, the first of the new C-bitcode
ukernel path: `e2e_matmul_cpu_dt_inner_tiled_llvm_ukernel_bf16_f32`
(avx512bf16) and `..._i8_i32` (avx512vnni). Each compiles a data-tiled
matmul with `--iree-llvmcpu-enable-llvm-ukernels=inner_tiled`, links the
ukernel bitcode, runs on host and checks results against a reference --
exercising the operand threading and generic `intrinsics_{m,n,k}`
unrolling that the IR-level lit tests cannot. Both were confirmed to
actually select their ukernel (not silently fall back to codegen).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Benoit Jacob <jacob.benoit.1@gmail.com>
7422d39 to
5cd7454
Compare
6c0ae40 to
62ad0cd
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Replaces the bf16 and i8-VNNI seeds' stub bodies with real SIMD implementations, both generic over the
intrinsics_{m,n,k}unrolling factors and structured like the AMDGPU C ukernel they were adapted from (iree_uk_amdgpu_multi_mma_mfma_i32_16x16x32_i8): accumulators in registers, an outer loop over the K tiles (k_outer), and inside it the(intrinsics_m, intrinsics_n, intrinsics_k)unroll. Theintrinsics_*arrive as constants at the inlined call site, so the loops fully unroll and theacc_regsarrays become fixed register files -- the bitcode-LTO equivalent of a C++ template, as the README describes.MMA_X86_AVX512BF16_1x16x2_F32_BF16): one_mm512_dpbf16_psper (m, n, k), with the LHS K-pair broadcast viaset1_ps.MMA_X86_AVX512VNNI_16x16x2_I32_I8_CASTI16): the 16x16x2 tile is bit-compatible with the codegen pathlowerX86Avx512Vnni16x16x2I8-- onevpmovsxbwwiden of each i8 panel to i16, thevpshufd/vbroadcasti32x4fan-out, and 16vpdpwssdover the block-interleaved (rlo, chi, rhi, clo) ACC layout. The i8 ukernel needs-mavx512bwfor the widen, so it is added to the VNNI copts.LLVMCPUSelectUKernelsnow only selects a ukernel when its bitcode actually exists (viaattachUKernelBitcodeOnOp's bool return), so anMMAIntrinsicthe cost model picks but for which no seed exists -- e.g. the M<->N-swappedMMA_X86_AVX512BF16_16x1x2_F32_BF16-- falls back to codegen instead of dangling an undefined symbol.Adds two execution/numerical tests, the first of the new C-bitcode ukernel path:
e2e_matmul_cpu_dt_inner_tiled_llvm_ukernel_bf16_f32(avx512bf16) and..._i8_i32(avx512vnni). Each compiles a data-tiled matmul with--iree-llvmcpu-enable-llvm-ukernels=inner_tiled, links the ukernel bitcode, runs on host and checks results against a reference -- exercising the operand threading and genericintrinsics_{m,n,k}unrolling that the IR-level lit tests cannot. Both were confirmed to actually select their ukernel (not silently fall back to codegen).Progress towards #24574.