[ARM64] Add base bf16 mmt4d tiles for cores without the BF16 extension (#24900)
On Arm64 cores without the Armv8.6 BF16 extension, bf16 `mmt4d`
currently falls back to the scalar generic kernel. All existing Arm64
bf16 tiles require the `_bf16` feature:
```text
IREE_UK_MMT4D_TILE(arm_64, bf16, bf16, f32, 1..8, 8, 4, _bf16)
IREE_UK_MMT4D_TILE(arm_64, bf16, bf16, bf16, 1..8, 8, 4, _bf16)
```
This PR adds base Arm64 bf16 tiles for cores without that extension.
The new tiles widen bf16 operands to f32 using NEON, accumulate with
FMLA, and narrow back to bf16 on store when required. The same
implementation serves both bf16 and f32 accumulator/output variants.
No compiler changes are required: `enumerateMatmulTileArm64` already
selects `8x8x1` tiles for bf16 when the BF16 extension is unavailable.
## Performance
i.MX95 (Cortex-A55, six cores), 512x512x512, median of three runs:
| case | before | after | speedup |
|---|---:|---:|---:|
| bf16 -> bf16 | 216 ms | 7.64 ms | 28x |
| bf16 -> f32 | 216 ms | 7.53 ms | 29x |
| f32 -> f32 | — | 8.22 ms | — |
For narrow M, using Mx512 by 512x512:
| M | before | after | speedup |
|---|---:|---:|---:|
| 1 | 1.52 ms | 0.879 ms | 1.7x |
| 2 | 2.31 ms | 0.874 ms | 2.6x |
| 4 | 3.37 ms | 0.887 ms | 3.8x |
## Correctness
Tested on device with `iree-run-module` for M=1,2,4,8 for bf16 -> bf16
and M=8 for bf16 -> f32.
The registered tiles match the generic implementation for normal inputs.
The widening behavior for bf16 subnormals and NaN payloads follows the
native NEON conversion behavior and can differ from the scalar helper,
similarly to the existing Arm64 f16 tiles.
The bf16 -> bf16 path accumulates in f32 and rounds once when storing,
matching the behavior of the existing Arm64 f16 tile.
Assisted by: Claude
---------
Signed-off-by: Thomas Ziereis <ziereis@roofline.ai>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
diff --git a/build_tools/linters/typos.toml b/build_tools/linters/typos.toml
index 56ee027..c018595 100644
--- a/build_tools/linters/typos.toml
+++ b/build_tools/linters/typos.toml
@@ -25,6 +25,8 @@
valu = "valu"
# ARM Scalable Matrix Extension.
sme = "sme"
+# AArch64 Shift Left Long instruction.
+shll = "shll"
# MLIR OpFoldResult abbreviation.
ofr = "ofr"
# GPU instruction conversion suffix (dp4xi8toi32).
diff --git a/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_base.c b/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_base.c
index 5789b9c..19d5ecf 100644
--- a/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_base.c
+++ b/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_base.c
@@ -216,6 +216,173 @@
iree_uk_mmt4d_tile_f16f16f16_1x8x1_to_8x8x1_arm_64,
iree_uk_mmt4d_tile_f16f16f16_8x8x1_arm_64, 8)
+// Widens 4 bf16 values to f32. A bf16 value is the top 16 bits of the f32 with
+// the same value, so widening is a shift by the element width, which SHLL does
+// for 4 lanes at once.
+IREE_UK_ATTRIBUTE_ALWAYS_INLINE static inline float32x4_t
+iree_uk_bf16x4_to_f32x4_arm_64(const iree_uk_uint16_t* IREE_UK_RESTRICT ptr) {
+ return vreinterpretq_f32_u32(vshll_n_u16(vld1_u16(ptr), 16));
+}
+
+// Narrows 4 f32 values to bf16. Both formats share an exponent field, so one
+// round-to-nearest-even bias over the whole word covers normals and subnormals
+// alike. Runs once per output tile.
+IREE_UK_ATTRIBUTE_ALWAYS_INLINE static inline uint16x4_t
+iree_uk_f32x4_to_bf16x4_arm_64(float32x4_t value) {
+ uint32x4_t bits = vreinterpretq_u32_f32(value);
+ uint32x4_t exp = vandq_u32(bits, vdupq_n_u32(0x7F800000u));
+ uint32x4_t mantissa = vandq_u32(bits, vdupq_n_u32(0x007FFFFFu));
+ // Round to nearest even: bias by half an interval plus the low bit of the
+ // result that is about to be kept, then truncate.
+ uint32x4_t keep_lsb = vandq_u32(vshrq_n_u32(bits, 16), vdupq_n_u32(1));
+ uint32x4_t rounded =
+ vaddq_u32(bits, vaddq_u32(vdupq_n_u32(0x7FFFu), keep_lsb));
+ // Inf needs no case of its own: a zero mantissa cannot carry into the
+ // exponent. NaN does, or a payload held entirely in the truncated bits would
+ // round up into an infinity. Setting the quiet bit keeps the mantissa
+ // nonzero while preserving whatever payload survives, as compiler-rt does.
+ uint32x4_t is_nan = vandq_u32(vceqq_u32(exp, vdupq_n_u32(0x7F800000u)),
+ vmvnq_u32(vceqzq_u32(mantissa)));
+ uint32x4_t quiet_nan = vorrq_u32(bits, vdupq_n_u32(0x00400000u));
+ rounded = vbslq_u32(is_nan, quiet_nan, rounded);
+ return vshrn_n_u32(rounded, 16);
+}
+
+// Shared implementation for bf16bf16bf16 and bf16bf16f32 on arm_64 cores
+// without the BF16 extension: widen both operands to f32 and use FMLA. The
+// widening is one SHLL per 4 elements, so the inner loop still runs at the FMLA
+// throughput the f32 tile achieves while reading half as many bytes.
+// In the bf16bf16bf16 case, intermediate roundings are skipped, as in the
+// f16f16f16 tile above: the accumulator stays f32 and is rounded once on store.
+IREE_UK_ATTRIBUTE_ALWAYS_INLINE static inline void
+iree_uk_mmt4d_tile_bf16bf16fXX_1x8x1_to_8x8x1_arm_64(
+ void* IREE_UK_RESTRICT out_tile, const void* IREE_UK_RESTRICT lhs_panel,
+ const void* IREE_UK_RESTRICT rhs_panel,
+ const iree_uk_mmt4d_params_t* params, iree_uk_type_t acc_type, int M0) {
+ IREE_UK_ASSERT(M0 >= 1 && M0 <= 8 && iree_uk_is_po2_u32(M0));
+ const iree_uk_uint16_t* IREE_UK_RESTRICT lhs_ptr = lhs_panel;
+ const iree_uk_uint16_t* IREE_UK_RESTRICT rhs_ptr = rhs_panel;
+ float32x4_t acc[16];
+ if (params->flags & IREE_UK_FLAG_MMT4D_ACCUMULATE) {
+ if (acc_type == IREE_UK_TYPE_FLOAT_32) {
+ const float* IREE_UK_RESTRICT out_ptr = out_tile;
+ IREE_UK_UNROLL for (int i = 0; i < 2 * M0; ++i) {
+ acc[i] = vld1q_f32(out_ptr + 4 * i);
+ }
+ } else {
+ const iree_uk_uint16_t* IREE_UK_RESTRICT out_ptr = out_tile;
+ IREE_UK_UNROLL for (int i = 0; i < 2 * M0; ++i) {
+ acc[i] = iree_uk_bf16x4_to_f32x4_arm_64(out_ptr + 4 * i);
+ }
+ }
+ } else {
+ IREE_UK_UNROLL for (int i = 0; i < 2 * M0; ++i) { acc[i] = vdupq_n_f32(0); }
+ }
+ for (int k = 0; k < params->K; ++k) {
+ float32x4_t rhs[2];
+ IREE_UK_UNROLL for (int i = 0; i < 2; ++i) {
+ rhs[i] = iree_uk_bf16x4_to_f32x4_arm_64(rhs_ptr + 4 * i);
+ }
+ rhs_ptr += 8;
+
+ if (M0 == 1) {
+ float32x4_t lhs =
+ vreinterpretq_f32_u32(vshll_n_u16(vld1_dup_u16(lhs_ptr), 16));
+ ++lhs_ptr;
+ acc[0] = vfmaq_lane_f32(acc[0], rhs[0], vget_low_f32(lhs), 0);
+ acc[1] = vfmaq_lane_f32(acc[1], rhs[1], vget_low_f32(lhs), 0);
+ } else if (M0 == 2) {
+ uint16x4_t lhs_bf16 = vld1_dup_u16(lhs_ptr);
+ lhs_bf16 = vld1_lane_u16(lhs_ptr + 1, lhs_bf16, 1);
+ lhs_ptr += 2;
+ float32x4_t lhs = vreinterpretq_f32_u32(vshll_n_u16(lhs_bf16, 16));
+ acc[0] = vfmaq_lane_f32(acc[0], rhs[0], vget_low_f32(lhs), 0);
+ acc[1] = vfmaq_lane_f32(acc[1], rhs[1], vget_low_f32(lhs), 0);
+ acc[2] = vfmaq_lane_f32(acc[2], rhs[0], vget_low_f32(lhs), 1);
+ acc[3] = vfmaq_lane_f32(acc[3], rhs[1], vget_low_f32(lhs), 1);
+ } else {
+ float32x4_t lhs[2];
+ IREE_UK_UNROLL for (int i = 0; i < M0 / 4; ++i) {
+ lhs[i] = iree_uk_bf16x4_to_f32x4_arm_64(lhs_ptr + 4 * i);
+ }
+ lhs_ptr += M0;
+ acc[0] = vfmaq_lane_f32(acc[0], rhs[0], vget_low_f32(lhs[0]), 0);
+ acc[1] = vfmaq_lane_f32(acc[1], rhs[1], vget_low_f32(lhs[0]), 0);
+ acc[2] = vfmaq_lane_f32(acc[2], rhs[0], vget_low_f32(lhs[0]), 1);
+ acc[3] = vfmaq_lane_f32(acc[3], rhs[1], vget_low_f32(lhs[0]), 1);
+ acc[4] = vfmaq_lane_f32(acc[4], rhs[0], vget_high_f32(lhs[0]), 0);
+ acc[5] = vfmaq_lane_f32(acc[5], rhs[1], vget_high_f32(lhs[0]), 0);
+ acc[6] = vfmaq_lane_f32(acc[6], rhs[0], vget_high_f32(lhs[0]), 1);
+ acc[7] = vfmaq_lane_f32(acc[7], rhs[1], vget_high_f32(lhs[0]), 1);
+ if (M0 == 8) {
+ acc[8] = vfmaq_lane_f32(acc[8], rhs[0], vget_low_f32(lhs[1]), 0);
+ acc[9] = vfmaq_lane_f32(acc[9], rhs[1], vget_low_f32(lhs[1]), 0);
+ acc[10] = vfmaq_lane_f32(acc[10], rhs[0], vget_low_f32(lhs[1]), 1);
+ acc[11] = vfmaq_lane_f32(acc[11], rhs[1], vget_low_f32(lhs[1]), 1);
+ acc[12] = vfmaq_lane_f32(acc[12], rhs[0], vget_high_f32(lhs[1]), 0);
+ acc[13] = vfmaq_lane_f32(acc[13], rhs[1], vget_high_f32(lhs[1]), 0);
+ acc[14] = vfmaq_lane_f32(acc[14], rhs[0], vget_high_f32(lhs[1]), 1);
+ acc[15] = vfmaq_lane_f32(acc[15], rhs[1], vget_high_f32(lhs[1]), 1);
+ }
+ }
+ }
+ if (acc_type == IREE_UK_TYPE_FLOAT_32) {
+ float* IREE_UK_RESTRICT out_ptr = out_tile;
+ IREE_UK_UNROLL for (int i = 0; i < 2 * M0; ++i) {
+ vst1q_f32(out_ptr + 4 * i, acc[i]);
+ }
+ } else {
+ iree_uk_uint16_t* IREE_UK_RESTRICT out_ptr = out_tile;
+ IREE_UK_UNROLL for (int i = 0; i < 2 * M0; ++i) {
+ vst1_u16(out_ptr + 4 * i, iree_uk_f32x4_to_bf16x4_arm_64(acc[i]));
+ }
+ }
+}
+
+IREE_UK_ATTRIBUTE_ALWAYS_INLINE static inline void
+iree_uk_mmt4d_tile_bf16bf16bf16_1x8x1_to_8x8x1_arm_64(
+ void* IREE_UK_RESTRICT out_tile, const void* IREE_UK_RESTRICT lhs_panel,
+ const void* IREE_UK_RESTRICT rhs_panel,
+ const iree_uk_mmt4d_params_t* params, int M0) {
+ iree_uk_mmt4d_tile_bf16bf16fXX_1x8x1_to_8x8x1_arm_64(
+ out_tile, lhs_panel, rhs_panel, params, IREE_UK_TYPE_BFLOAT_16, M0);
+}
+
+IREE_UK_ATTRIBUTE_ALWAYS_INLINE static inline void
+iree_uk_mmt4d_tile_bf16bf16f32_1x8x1_to_8x8x1_arm_64(
+ void* IREE_UK_RESTRICT out_tile, const void* IREE_UK_RESTRICT lhs_panel,
+ const void* IREE_UK_RESTRICT rhs_panel,
+ const iree_uk_mmt4d_params_t* params, int M0) {
+ iree_uk_mmt4d_tile_bf16bf16fXX_1x8x1_to_8x8x1_arm_64(
+ out_tile, lhs_panel, rhs_panel, params, IREE_UK_TYPE_FLOAT_32, M0);
+}
+
+IREE_UK_MMT4D_TILE_FUNC_IMPL_FOR_M0(
+ iree_uk_mmt4d_tile_bf16bf16f32_1x8x1_to_8x8x1_arm_64,
+ iree_uk_mmt4d_tile_bf16bf16f32_1x8x1_arm_64, 1)
+IREE_UK_MMT4D_TILE_FUNC_IMPL_FOR_M0(
+ iree_uk_mmt4d_tile_bf16bf16f32_1x8x1_to_8x8x1_arm_64,
+ iree_uk_mmt4d_tile_bf16bf16f32_2x8x1_arm_64, 2)
+IREE_UK_MMT4D_TILE_FUNC_IMPL_FOR_M0(
+ iree_uk_mmt4d_tile_bf16bf16f32_1x8x1_to_8x8x1_arm_64,
+ iree_uk_mmt4d_tile_bf16bf16f32_4x8x1_arm_64, 4)
+IREE_UK_MMT4D_TILE_FUNC_IMPL_FOR_M0(
+ iree_uk_mmt4d_tile_bf16bf16f32_1x8x1_to_8x8x1_arm_64,
+ iree_uk_mmt4d_tile_bf16bf16f32_8x8x1_arm_64, 8)
+
+IREE_UK_MMT4D_TILE_FUNC_IMPL_FOR_M0(
+ iree_uk_mmt4d_tile_bf16bf16bf16_1x8x1_to_8x8x1_arm_64,
+ iree_uk_mmt4d_tile_bf16bf16bf16_1x8x1_arm_64, 1)
+IREE_UK_MMT4D_TILE_FUNC_IMPL_FOR_M0(
+ iree_uk_mmt4d_tile_bf16bf16bf16_1x8x1_to_8x8x1_arm_64,
+ iree_uk_mmt4d_tile_bf16bf16bf16_2x8x1_arm_64, 2)
+IREE_UK_MMT4D_TILE_FUNC_IMPL_FOR_M0(
+ iree_uk_mmt4d_tile_bf16bf16bf16_1x8x1_to_8x8x1_arm_64,
+ iree_uk_mmt4d_tile_bf16bf16bf16_4x8x1_arm_64, 4)
+IREE_UK_MMT4D_TILE_FUNC_IMPL_FOR_M0(
+ iree_uk_mmt4d_tile_bf16bf16bf16_1x8x1_to_8x8x1_arm_64,
+ iree_uk_mmt4d_tile_bf16bf16bf16_8x8x1_arm_64, 8)
+
IREE_UK_ATTRIBUTE_ALWAYS_INLINE static inline void
iree_uk_mmt4d_tile_s8s8s32_1x8x1_to_8x8x1_arm_64(
void* IREE_UK_RESTRICT out_tile, const void* IREE_UK_RESTRICT lhs_panel,
diff --git a/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_tiles.inl b/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_tiles.inl
index 6bc173a..ad45057 100644
--- a/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_tiles.inl
+++ b/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_tiles.inl
@@ -27,6 +27,16 @@
IREE_UK_MMT4D_TILE(arm_64, f16, f16, f16, 2, 8, 1, _fullfp16)
IREE_UK_MMT4D_TILE(arm_64, f16, f16, f16, 4, 8, 1, _fullfp16)
IREE_UK_MMT4D_TILE(arm_64, f16, f16, f16, 8, 8, 1, _fullfp16)
+// Cores without the BF16 extension widen to f32 and use FMLA, matching the
+// f32/f16 tile shape the compiler enumerates for them.
+IREE_UK_MMT4D_TILE(arm_64, bf16, bf16, f32, 1, 8, 1, )
+IREE_UK_MMT4D_TILE(arm_64, bf16, bf16, f32, 2, 8, 1, )
+IREE_UK_MMT4D_TILE(arm_64, bf16, bf16, f32, 4, 8, 1, )
+IREE_UK_MMT4D_TILE(arm_64, bf16, bf16, f32, 8, 8, 1, )
+IREE_UK_MMT4D_TILE(arm_64, bf16, bf16, bf16, 1, 8, 1, )
+IREE_UK_MMT4D_TILE(arm_64, bf16, bf16, bf16, 2, 8, 1, )
+IREE_UK_MMT4D_TILE(arm_64, bf16, bf16, bf16, 4, 8, 1, )
+IREE_UK_MMT4D_TILE(arm_64, bf16, bf16, bf16, 8, 8, 1, )
IREE_UK_MMT4D_TILE(arm_64, bf16, bf16, f32, 1, 8, 4, _bf16)
IREE_UK_MMT4D_TILE(arm_64, bf16, bf16, f32, 2, 8, 4, _bf16)
IREE_UK_MMT4D_TILE(arm_64, bf16, bf16, f32, 4, 8, 4, _bf16)
diff --git a/runtime/src/iree/builtins/ukernel/common.h b/runtime/src/iree/builtins/ukernel/common.h
index ee7f3cf..7159049 100644
--- a/runtime/src/iree/builtins/ukernel/common.h
+++ b/runtime/src/iree/builtins/ukernel/common.h
@@ -662,6 +662,15 @@
const int prefix##exp_mask IREE_UK_ATTRIBUTE_UNUSED = \
(1u << prefix##sign_shift) - (1u << prefix##exp_shift);
+// Rounds to nearest even by adding a bias. Does not shift; the result is ready
+// to be right-shifted by shift_amount.
+static inline iree_uk_uint32_t iree_uk_bias_to_nearest_even(
+ iree_uk_uint32_t input, int shift_amount) {
+ iree_uk_uint32_t even_bit = 1u << shift_amount;
+ iree_uk_uint32_t odd_bit = even_bit >> 1;
+ return input + ((input & even_bit) ? odd_bit : (odd_bit - 1));
+}
+
static inline float iree_uk_generic_fp16_to_f32(iree_uk_uint16_t f16_value,
int exp_bits) {
IREE_UK_FP_FORMAT_CONSTANTS(f16_, 16, exp_bits)
@@ -683,7 +692,12 @@
// Inf. Leave zero mantissa.
}
} else if (f16_exp == 0) {
- // Zero or subnormal. Generate zero. Leave zero mantissa.
+ // Zero or subnormal. bf16 shares f32's exponent field, so its subnormals
+ // are f32 subnormals with the mantissa shifted into place. Narrower
+ // exponents would need renormalizing; leave those zero.
+ if (f16_exp_bits == f32_exp_bits) {
+ f32_mantissa = f16_mantissa << (f32_mantissa_bits - f16_mantissa_bits);
+ }
} else {
// Normal finite value.
int arithmetic_f16_exp = f16_exp >> f16_exp_shift;
@@ -721,7 +735,14 @@
// Inf. Leave zero mantissa.
}
} else if (f32_exp == 0) {
- // Zero or subnormal. Generate zero. Leave zero mantissa.
+ // Zero or subnormal. bf16 shares f32's exponent field, so f32 subnormals
+ // stay representable: round the mantissa and let a carry out of the field
+ // land in the exponent. Narrower exponents cannot represent them.
+ if (f16_exp_bits == f32_exp_bits) {
+ int shift_amount = f32_mantissa_bits - f16_mantissa_bits;
+ f16_mantissa = iree_uk_bias_to_nearest_even(f32_mantissa, shift_amount) >>
+ shift_amount;
+ }
} else {
// Normal finite value.
int arithmetic_exp = (f32_exp >> f32_exp_shift) - (1 << (f32_exp_bits - 1));
@@ -733,13 +754,8 @@
f16_exp = 0;
} else {
// Normal case.
- // Implement round-to-nearest-even, by adding a bias before truncating.
- // truncating.
- int even_bit = 1u << (f32_mantissa_bits - f16_mantissa_bits);
- int odd_bit = even_bit >> 1;
- iree_uk_uint32_t biased_f32_mantissa =
- f32_mantissa +
- ((f32_mantissa & even_bit) ? (odd_bit) : (odd_bit - 1));
+ iree_uk_uint32_t biased_f32_mantissa = iree_uk_bias_to_nearest_even(
+ f32_mantissa, f32_mantissa_bits - f16_mantissa_bits);
// Adding the bias may cause an exponent increment.
if (biased_f32_mantissa > f32_mantissa_mask) {
// Note: software implementations that try to be fast tend to get this
diff --git a/runtime/src/iree/builtins/ukernel/tools/mmt4d_benchmark.c b/runtime/src/iree/builtins/ukernel/tools/mmt4d_benchmark.c
index ea537ac..d56ac3e 100644
--- a/runtime/src/iree/builtins/ukernel/tools/mmt4d_benchmark.c
+++ b/runtime/src/iree/builtins/ukernel/tools/mmt4d_benchmark.c
@@ -143,9 +143,13 @@
"");
iree_uk_benchmark_register_mmt4d(IREE_UK_FLAG_MMT4D_TYPE_F16F16F16, 8, 8, 1,
"fullfp16");
+ iree_uk_benchmark_register_mmt4d(IREE_UK_FLAG_MMT4D_TYPE_BF16BF16F32, 8, 8, 1,
+ "");
iree_uk_benchmark_register_mmt4d(IREE_UK_FLAG_MMT4D_TYPE_BF16BF16F32, 8, 8, 4,
"bf16");
iree_uk_benchmark_register_mmt4d(IREE_UK_FLAG_MMT4D_TYPE_BF16BF16BF16, 8, 8,
+ 1, "");
+ iree_uk_benchmark_register_mmt4d(IREE_UK_FLAG_MMT4D_TYPE_BF16BF16BF16, 8, 8,
4, "bf16");
iree_uk_benchmark_register_mmt4d(IREE_UK_FLAG_MMT4D_TYPE_S8S8S32, 8, 8, 1,
"");
diff --git a/runtime/src/iree/builtins/ukernel/tools/mmt4d_test.c b/runtime/src/iree/builtins/ukernel/tools/mmt4d_test.c
index a09f227..fd5f6b5 100644
--- a/runtime/src/iree/builtins/ukernel/tools/mmt4d_test.c
+++ b/runtime/src/iree/builtins/ukernel/tools/mmt4d_test.c
@@ -568,6 +568,10 @@
iree_uk_test_mmt4d(IREE_UK_FLAG_MMT4D_SKIP_INTERMEDIATE_ROUNDINGS |
IREE_UK_FLAG_MMT4D_TYPE_F16F16F16,
8, 8, 1, "");
+ iree_uk_test_mmt4d(IREE_UK_FLAG_MMT4D_TYPE_BF16BF16F32, 8, 8, 1, "");
+ iree_uk_test_mmt4d(IREE_UK_FLAG_MMT4D_SKIP_INTERMEDIATE_ROUNDINGS |
+ IREE_UK_FLAG_MMT4D_TYPE_BF16BF16BF16,
+ 8, 8, 1, "");
iree_uk_test_mmt4d(IREE_UK_FLAG_MMT4D_TYPE_S8S8S32, 8, 8, 1, "");
iree_uk_test_mmt4d(IREE_UK_FLAG_MMT4D_TYPE_S8S4S32, 4, 16, 2, "");
iree_uk_test_mmt4d(IREE_UK_FLAG_MMT4D_TYPE_F16F16F32, 8, 8, 1, "fp16fml");
