[LLVMCPU][RISCV] Add RVV bf16 mmt4d ukernel using the Zvfbfwma widening MAC (bf16*bf16->f32) (#24695)
## Summary
On RISC-V, a data-tiled bf16 matmul was computed by promoting LHS/RHS to
f32, ignoring the `Zvfbfwma` bf16 widening multiply-accumulate extension
(`vfwmaccbf16`) that LLVM already lowers to and that shipping silicon
implements. This PR adds a native `bf16 x bf16 -> f32` `mmt4d`
microkernel for RVV+`Zvfbfwma`. Naturally, this change is related to
(and was implementing looking at) the existing f32, `int8` (#23734), and
f16/`Zvfh` (#22231) ukernels.
## How the change was tested
1. Compiler e2e
A `128x256 bf16 * 256x128 bf16 -> 128x128 f32` matmul compiled for
`riscv64 +v,+zvfbfwma` with `--iree-opt-data-tiling
--iree-llvmcpu-enable-ukernels=all`:
Before: dispatch contains `arith.extf` f32-promotion, no ukernel.
After: `Mmt4dTilingExpert` dispatch `matmul_128x128x256_bf16xbf16xf32`
with `iree_uk_mmt4d` + `iree_uk_pack`; **zero `arith.extf`**; all four
`iree_uk_mmt4d_tile_bf16bf16f32_*_riscv_64_zvfbfwma` symbols present;
generated RISC-V asm inner loop is `vfwmaccbf16.vf`.
2. Numerics under `qemu-riscv64`
A standalone harness calling the real tile functions matches an f32
reference exactly. A control run without the extension crashes,
confirming `vfwmaccbf16` is genuinely executed. I used qemu 8.2 which
exposes the extension as the experimental `-cpu max,x-zvfbfwma=true`
property.
3. Tile selection under `qemu-riscv64`.
Verified `iree_uk_mmt4d_select_tile_func_arch` (the
`entry_point.c`/`tiles.inl` wiring) returns the correct bf16 tile for
each M0 value when `cpu_data` has the `Zvfbfwma` bit, and does not
select it when the bit is clear — i.e. the feature is only enabled for
`Zvfbfwma`.
4. In-tree `mmt4d_test`.
The added case exercises the same path in CI's RISC-V cross-build (falls
back to the generic tile if the CI toolchain can't assemble `Zvfbfwma`,
i think it's the same limitation as `Zvfh` in #22303).
### A note on CI coverage
If I understand correctly, the RISC-V CI job runs under an older
`qemu-riscv64` on a stock `rv64gcv` profile that doesn't expose
`Zvfbfwma`, and `mmt4d_test` *skips* (does not fail) any case whose
feature the runtime CPU lacks. Zvfbfwma is only assembled/executed on a
toolchain + qemu new enough to support it (qemu ≥ 8.2 exposes it as the
experimental `x-zvfbfwma` property). My point is: a green CI does not
mean the bf16 kernel ran in CI — the functional validation of record is
the local `qemu-riscv64` 8.2 numerics + tile-selection runs described
above.
Disclosure: this change was assisted by Claude Code, but the change was
reviewed and thoroughly tested before being submitted.
---------
Signed-off-by: Zmicier Prybysh <zprybysh@baylibre.com>
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_riscv.mlir b/compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_riscv.mlir
index a7ed1af..327a069 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_riscv.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_riscv.mlir
@@ -361,3 +361,82 @@
// SCALABLE-SAME: inner_tiles = [7, %[[C8_VSCALE]]]
// CHECK: return %[[UNPACK]]
+
+// -----
+
+// RISC-V64 with V extension and zvfbfwma - full matmul lowering (bf16 -> f32)
+
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#encoding_lhs = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [bf16, bf16, f32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]>
+#encoding_rhs = #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [bf16, bf16, f32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]>
+#encoding_result = #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [bf16, bf16, f32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]>
+func.func @full_matmul_lowering_bf16bf16f32_riscv64(
+ %lhs: tensor<?x?xbf16>,
+ %rhs: tensor<?x?xbf16>,
+ %acc: tensor<?x?xf32>,
+ %m : index, %n : index, %k : index
+) -> tensor<?x?xf32> attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {target_triple="riscv64-xyz-xyz", cpu_features="+v,+zvfbfwma", iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %0 = iree_encoding.set_encoding %lhs encoding_dims{%m, %n, %k} : tensor<?x?xbf16> -> tensor<?x?xbf16, #encoding_lhs>
+ %1 = iree_encoding.set_encoding %rhs encoding_dims{%m, %n, %k} : tensor<?x?xbf16> -> tensor<?x?xbf16, #encoding_rhs>
+ %2 = iree_encoding.set_encoding %acc encoding_dims{%m, %n, %k} : tensor<?x?xf32> -> tensor<?x?xf32, #encoding_result>
+ %3 = linalg.matmul
+ ins(%0, %1 : tensor<?x?xbf16, #encoding_lhs>,
+ tensor<?x?xbf16, #encoding_rhs>)
+ outs(%2 : tensor<?x?xf32, #encoding_result>)
+ -> tensor<?x?xf32, #encoding_result>
+ %4 = iree_encoding.unset_encoding %3 encoding_dims{%m, %n, %k} : tensor<?x?xf32, #encoding_result> -> tensor<?x?xf32>{%m, %n}
+ return %4 : tensor<?x?xf32>
+}
+
+/// BF16 -> F32 on RISC-V64 with V+zvfbfwma extension:
+/// For the static case: M=7, N=16 (vlen=128), K=1
+/// For the scalable case: M=7, N=8*vscale (vlen=64*vscale), K=1
+/// The bf16 operands are LMUL=2 and the f32 accumulators EMUL=4.
+
+// CHECK-LABEL: func @full_matmul_lowering_bf16bf16f32_riscv64(
+// CHECK-SAME: %[[LHS:[a-zA-Z0-9]+]]: tensor<?x?xbf16>
+// CHECK-SAME: %[[RHS:[a-zA-Z0-9]+]]: tensor<?x?xbf16>
+// CHECK-SAME: %[[ACC:[a-zA-Z0-9]+]]: tensor<?x?xf32>
+
+// SCALABLE: %[[C8:.+]] = arith.constant 8
+
+/// Pack LHS: [M, K] -> [?, ?, 7, 1] (this happens before vscale computation)
+// CHECK: %[[PACK_LHS:.+]] = linalg.pack %[[LHS]]
+// CHECK-SAME: outer_dims_perm = [0, 1]
+// CHECK-SAME: inner_dims_pos = [0, 1]
+// CHECK-SAME: inner_tiles = [7, 1]
+
+// SCALABLE: %[[VSCALE:.+]] = vector.vscale
+// SCALABLE: %[[C8_VSCALE:.+]] = arith.muli %[[VSCALE]], %[[C8]]
+
+/// Pack RHS: [K, N] -> [?, ?, N_tile, 1] with outer_dims_perm = [1, 0]
+// CHECK: %[[PACK_RHS:.+]] = linalg.pack %[[RHS]]
+// CHECK-SAME: outer_dims_perm = [1, 0]
+// CHECK-SAME: inner_dims_pos = [1, 0]
+// STATIC-SAME: inner_tiles = [16, 1]
+// SCALABLE-SAME: inner_tiles = [%[[C8_VSCALE]], 1]
+
+/// Pack ACC: [M, N] -> [?, ?, 7, N_tile]
+// CHECK: %[[PACK_ACC:.+]] = linalg.pack %[[ACC]]
+// CHECK-SAME: outer_dims_perm = [0, 1]
+// CHECK-SAME: inner_dims_pos = [0, 1]
+// STATIC-SAME: inner_tiles = [7, 16]
+// SCALABLE-SAME: inner_tiles = [7, %[[C8_VSCALE]]]
+
+/// The mmt4d operation
+// CHECK: %[[MMT4D:.+]] = linalg.mmt4d
+// CHECK-SAME: ins(%[[PACK_LHS]], %[[PACK_RHS]] :
+// CHECK-SAME: outs(%[[PACK_ACC]] :
+
+/// Unpack result: [?, ?, 7, N_tile] -> [M, N]
+// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[MMT4D]]
+// CHECK-SAME: outer_dims_perm = [0, 1]
+// CHECK-SAME: inner_dims_pos = [0, 1]
+// STATIC-SAME: inner_tiles = [7, 16]
+// SCALABLE-SAME: inner_tiles = [7, %[[C8_VSCALE]]]
+
+// CHECK: return %[[UNPACK]]
diff --git a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp
index cd14497..e3d2324 100644
--- a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp
+++ b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp
@@ -1090,6 +1090,23 @@
};
}
}
+ // BF16 path: RVV bf16 widening multiply-accumulate via Zvfbfwma.
+ // Same formula as the f32 path: N0 = VLEN/8, so the bf16 operands are LMUL=2
+ // (m2) and the f32 accumulators EMUL=4 (m4). vfwmaccbf16 fuses the bf16 ->
+ // f32 widening into the MAC, so the accumulator is f32 and there is no
+ // bf16-accumulator variant.
+ // M0=7 with EMUL=4 accumulators: 7*4 + 2(rhs) = 30 regs.
+ if (lhs.isBF16() && rhs.isBF16() && out.isF32()) {
+ int N0 = vlen / 8;
+ if (hasFeature(config, "+zvfbfwma")) {
+ return {
+ TileMxNxK{7, N0, 1}, // Aim to use VFWMACCBF16.VF.
+ TileMxNxK{4, N0, 1}, // Truncation of the above.
+ TileMxNxK{2, N0, 1}, // Truncation of the above.
+ TileMxNxK{1, N0, 1}, // Truncation of the above.
+ };
+ }
+ }
// Integer 8 path: standard RVV widening multiply-accumulate.
// Same formula as the f32 path: N0 = VLEN/8, targeting LMUL=4 for i32
// accumulators. The widening chain is i8(m1) -> i16(m2) -> i32(m4),
diff --git a/runtime/src/iree/builtins/ukernel/arch/riscv_64/BUILD.bazel b/runtime/src/iree/builtins/ukernel/arch/riscv_64/BUILD.bazel
index 108ca5b..b41e2f4 100644
--- a/runtime/src/iree/builtins/ukernel/arch/riscv_64/BUILD.bazel
+++ b/runtime/src/iree/builtins/ukernel/arch/riscv_64/BUILD.bazel
@@ -78,6 +78,16 @@
internal_hdrs = UKERNEL_RISCV_64_INTERNAL_HEADERS,
)
+iree_bitcode_library(
+ name = "ukernel_bitcode_arch_riscv_64_zvfbfwma",
+ srcs = [
+ "mmt4d_riscv_64_zvfbfwma.c",
+ ],
+ arch = "riscv_64",
+ copts = ["-march=rv64gcv_zvfbfwma"],
+ internal_hdrs = UKERNEL_RISCV_64_INTERNAL_HEADERS,
+)
+
iree_link_bitcode(
name = "ukernel_bitcode_arch_riscv_64",
bitcode_files = [
@@ -85,6 +95,7 @@
"ukernel_bitcode_arch_riscv_64_v.bc",
"ukernel_bitcode_arch_riscv_64_zvfhmin.bc",
"ukernel_bitcode_arch_riscv_64_zvfh.bc",
+ "ukernel_bitcode_arch_riscv_64_zvfbfwma.bc",
],
)
diff --git a/runtime/src/iree/builtins/ukernel/arch/riscv_64/CMakeLists.txt b/runtime/src/iree/builtins/ukernel/arch/riscv_64/CMakeLists.txt
index b02f6c7..7cd59ef 100644
--- a/runtime/src/iree/builtins/ukernel/arch/riscv_64/CMakeLists.txt
+++ b/runtime/src/iree/builtins/ukernel/arch/riscv_64/CMakeLists.txt
@@ -90,12 +90,32 @@
"-march=rv64gc_zvfh"
)
+iree_bitcode_library(
+ NAME
+ ukernel_bitcode_arch_riscv_64_zvfbfwma
+ ARCH
+ riscv_64
+ INTERNAL_HDRS
+ "${PROJECT_BINARY_DIR}/runtime/src/iree/builtins/ukernel/internal_headers_filegroup.stamp"
+ "${PROJECT_BINARY_DIR}/runtime/src/iree/schemas/cpu_data_headers_filegroup.stamp"
+ "common_riscv_64.h"
+ "mmt4d_riscv_64_internal.h"
+ "mmt4d_riscv_64_tiles.inl"
+ "pack_riscv_64_internal.h"
+ "unpack_riscv_64_internal.h"
+ SRCS
+ "mmt4d_riscv_64_zvfbfwma.c"
+ COPTS
+ "-march=rv64gcv_zvfbfwma"
+)
+
iree_link_bitcode(
NAME
ukernel_bitcode_arch_riscv_64
SRCS
"ukernel_bitcode_arch_riscv_64_entry_points.bc"
"ukernel_bitcode_arch_riscv_64_v.bc"
+ "ukernel_bitcode_arch_riscv_64_zvfbfwma.bc"
"ukernel_bitcode_arch_riscv_64_zvfh.bc"
"ukernel_bitcode_arch_riscv_64_zvfhmin.bc"
@@ -126,8 +146,14 @@
"-march=rv64gc_zvfh"
)
+iree_select_compiler_opts(IREE_UK_COPTS_RISCV_64_ZVFBFWMA
+ CLANG_OR_GCC
+ "-march=rv64gcv_zvfbfwma"
+)
+
check_cxx_compiler_flag("${IREE_UK_COPTS_RISCV_64_V}" IREE_UK_BUILD_RISCV_64_V)
check_cxx_compiler_flag("${IREE_UK_COPTS_RISCV_64_ZVFH}" IREE_UK_BUILD_RISCV_64_ZVFH)
+check_cxx_compiler_flag("${IREE_UK_COPTS_RISCV_64_ZVFBFWMA}" IREE_UK_BUILD_RISCV_64_ZVFBFWMA)
# CPU features that we will try checking compiler support for, unless
# we set them to OFF below.
@@ -215,6 +241,20 @@
list(APPEND IREE_UK_RISCV_64_DEPS "::riscv_64_zvfh")
endif() # IREE_UK_BUILD_RISCV_64_ZVFH
+if(IREE_UK_BUILD_RISCV_64_ZVFBFWMA)
+iree_cc_library(
+ NAME
+ riscv_64_zvfbfwma
+ SRCS
+ "mmt4d_riscv_64_zvfbfwma.c"
+ COPTS
+ "${IREE_UK_COPTS_RISCV_64_ZVFBFWMA}"
+ DEPS
+ iree::builtins::ukernel::internal_headers
+)
+list(APPEND IREE_UK_RISCV_64_DEPS "::riscv_64_zvfbfwma")
+endif() # IREE_UK_BUILD_RISCV_64_ZVFBFWMA
+
iree_cc_library(
NAME
riscv_64
diff --git a/runtime/src/iree/builtins/ukernel/arch/riscv_64/common_riscv_64.h b/runtime/src/iree/builtins/ukernel/arch/riscv_64/common_riscv_64.h
index 62942fd..b68d961 100644
--- a/runtime/src/iree/builtins/ukernel/arch/riscv_64/common_riscv_64.h
+++ b/runtime/src/iree/builtins/ukernel/arch/riscv_64/common_riscv_64.h
@@ -15,6 +15,7 @@
#define IREE_UK_BUILD_RISCV_64_V
#define IREE_UK_BUILD_RISCV_64_ZVFHMIN
#define IREE_UK_BUILD_RISCV_64_ZVFH
+#define IREE_UK_BUILD_RISCV_64_ZVFBFWMA
#else
// Compiling with the system toolchain. Include the configured header.
#include "iree/builtins/ukernel/arch/riscv_64/config_riscv_64.h"
@@ -38,4 +39,9 @@
return iree_uk_all_bits_set(cpu_data[0], IREE_CPU_DATA0_RISCV_64_ZVFH);
}
+static inline bool iree_uk_cpu_riscv_64_zvfbfwma(
+ const iree_uk_uint64_t* cpu_data) {
+ return iree_uk_all_bits_set(cpu_data[0], IREE_CPU_DATA0_RISCV_64_ZVFBFWMA);
+}
+
#endif // IREE_BUILTINS_UKERNEL_ARCH_RISCV_64_COMMON_RISCV_64_H_
diff --git a/runtime/src/iree/builtins/ukernel/arch/riscv_64/config_riscv_64.h.in b/runtime/src/iree/builtins/ukernel/arch/riscv_64/config_riscv_64.h.in
index 42858b8..bee185d 100644
--- a/runtime/src/iree/builtins/ukernel/arch/riscv_64/config_riscv_64.h.in
+++ b/runtime/src/iree/builtins/ukernel/arch/riscv_64/config_riscv_64.h.in
@@ -14,5 +14,6 @@
#cmakedefine IREE_UK_BUILD_RISCV_64_V
#cmakedefine IREE_UK_BUILD_RISCV_64_ZVFHMIN
#cmakedefine IREE_UK_BUILD_RISCV_64_ZVFH
+#cmakedefine IREE_UK_BUILD_RISCV_64_ZVFBFWMA
#endif // IREE_BUILTINS_UKERNEL_ARCH_RISCV_64_CONFIG_RISCV_64_H_
diff --git a/runtime/src/iree/builtins/ukernel/arch/riscv_64/mmt4d_riscv_64_entry_point.c b/runtime/src/iree/builtins/ukernel/arch/riscv_64/mmt4d_riscv_64_entry_point.c
index 1c26680..bdde1c5 100644
--- a/runtime/src/iree/builtins/ukernel/arch/riscv_64/mmt4d_riscv_64_entry_point.c
+++ b/runtime/src/iree/builtins/ukernel/arch/riscv_64/mmt4d_riscv_64_entry_point.c
@@ -44,6 +44,13 @@
#define IREE_UK_MMT4D_TILE_riscv_64_zvfh(lhs, rhs, out, m0, k0)
#endif
+#ifdef IREE_UK_BUILD_RISCV_64_ZVFBFWMA
+#define IREE_UK_MMT4D_TILE_riscv_64_zvfbfwma(lhs, rhs, out, m0, k0) \
+ IREE_UK_MMT4D_TILE_IMPL_riscv_64(lhs, rhs, out, m0, k0, _zvfbfwma)
+#else
+#define IREE_UK_MMT4D_TILE_riscv_64_zvfbfwma(lhs, rhs, out, m0, k0)
+#endif
+
#define IREE_UK_MMT4D_TILE(arch, lhs, rhs, out, m0, k0, suffix) \
IREE_UK_MMT4D_TILE_riscv_64##suffix(lhs, rhs, out, m0, k0)
diff --git a/runtime/src/iree/builtins/ukernel/arch/riscv_64/mmt4d_riscv_64_tiles.inl b/runtime/src/iree/builtins/ukernel/arch/riscv_64/mmt4d_riscv_64_tiles.inl
index 7e8713f..c67baa5 100644
--- a/runtime/src/iree/builtins/ukernel/arch/riscv_64/mmt4d_riscv_64_tiles.inl
+++ b/runtime/src/iree/builtins/ukernel/arch/riscv_64/mmt4d_riscv_64_tiles.inl
@@ -37,3 +37,8 @@
IREE_UK_MMT4D_TILE(riscv_64, f16, f16, f16, 2, 1, _zvfh)
IREE_UK_MMT4D_TILE(riscv_64, f16, f16, f16, 4, 1, _zvfh)
IREE_UK_MMT4D_TILE(riscv_64, f16, f16, f16, 7, 1, _zvfh)
+
+IREE_UK_MMT4D_TILE(riscv_64, bf16, bf16, f32, 1, 1, _zvfbfwma)
+IREE_UK_MMT4D_TILE(riscv_64, bf16, bf16, f32, 2, 1, _zvfbfwma)
+IREE_UK_MMT4D_TILE(riscv_64, bf16, bf16, f32, 4, 1, _zvfbfwma)
+IREE_UK_MMT4D_TILE(riscv_64, bf16, bf16, f32, 7, 1, _zvfbfwma)
diff --git a/runtime/src/iree/builtins/ukernel/arch/riscv_64/mmt4d_riscv_64_zvfbfwma.c b/runtime/src/iree/builtins/ukernel/arch/riscv_64/mmt4d_riscv_64_zvfbfwma.c
new file mode 100644
index 0000000..7a7a3e9
--- /dev/null
+++ b/runtime/src/iree/builtins/ukernel/arch/riscv_64/mmt4d_riscv_64_zvfbfwma.c
@@ -0,0 +1,153 @@
+// Copyright 2026 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include <riscv_vector.h>
+
+#include "iree/builtins/ukernel/arch/riscv_64/common_riscv_64.h"
+#include "iree/builtins/ukernel/arch/riscv_64/mmt4d_riscv_64_internal.h"
+
+// bf16bf16f32 tile using the Zvfbfwma bf16 widening multiply-accumulate
+// (`vfwmaccbf16`). Structurally identical to the f16f16f32 path in
+// mmt4d_riscv_64_zvfh.c, with the f16 vector/scalar types swapped for bf16 and
+// the widening MAC swapped for its bf16 counterpart. The accumulator is f32 in
+// both, so there is no accumulator-type branching and no narrowing store:
+// unlike f16 (native f16 arithmetic via Zvfh), bf16 has no non-widening vector
+// FMA, so only the bf16bf16f32 case is provided here.
+IREE_UK_ATTRIBUTE_ALWAYS_INLINE static inline void
+iree_uk_mmt4d_tile_bf16bf16f32_1xXXx1_to_7xXXx1_riscv_64_zvfbfwma(
+ 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_ASSERT(M0 >= 1 && M0 <= 7);
+ const __bf16* IREE_UK_RESTRICT lhs_ptr = lhs_panel;
+ const __bf16* IREE_UK_RESTRICT rhs_ptr = rhs_panel;
+ float* IREE_UK_RESTRICT out_ptr = out_tile;
+
+ vfloat32m4_t acc0, acc1, acc2, acc3, acc4, acc5, acc6;
+ int N0 = params->N0;
+ size_t vl = N0;
+
+ if (M0 == 1) {
+ if (params->flags & IREE_UK_FLAG_MMT4D_ACCUMULATE) {
+ acc0 = __riscv_vle32_v_f32m4(out_ptr, vl);
+ } else {
+ acc0 = __riscv_vfmv_v_f_f32m4(0.0f, vl);
+ }
+ for (int k = 0; k < params->K; ++k) {
+ vbfloat16m2_t rhs = __riscv_vle16_v_bf16m2(rhs_ptr, vl);
+ rhs_ptr += N0;
+ __bf16 lhs0 = *lhs_ptr++;
+
+ acc0 = __riscv_vfwmaccbf16_vf_f32m4(acc0, lhs0, rhs, vl);
+ }
+ __riscv_vse32_v_f32m4(out_ptr, acc0, vl);
+ } else if (M0 == 2) {
+ if (params->flags & IREE_UK_FLAG_MMT4D_ACCUMULATE) {
+ acc0 = __riscv_vle32_v_f32m4(out_ptr, vl);
+ acc1 = __riscv_vle32_v_f32m4(out_ptr + N0, vl);
+ } else {
+ acc0 = __riscv_vfmv_v_f_f32m4(0.0f, vl);
+ acc1 = __riscv_vfmv_v_f_f32m4(0.0f, vl);
+ }
+ for (int k = 0; k < params->K; ++k) {
+ vbfloat16m2_t rhs = __riscv_vle16_v_bf16m2(rhs_ptr, vl);
+ rhs_ptr += N0;
+ __bf16 lhs0 = *lhs_ptr++;
+ __bf16 lhs1 = *lhs_ptr++;
+
+ acc0 = __riscv_vfwmaccbf16_vf_f32m4(acc0, lhs0, rhs, vl);
+ acc1 = __riscv_vfwmaccbf16_vf_f32m4(acc1, lhs1, rhs, vl);
+ }
+ __riscv_vse32_v_f32m4(out_ptr, acc0, vl);
+ __riscv_vse32_v_f32m4(out_ptr + N0, acc1, vl);
+ } else if (M0 == 4) {
+ if (params->flags & IREE_UK_FLAG_MMT4D_ACCUMULATE) {
+ acc0 = __riscv_vle32_v_f32m4(out_ptr, vl);
+ acc1 = __riscv_vle32_v_f32m4(out_ptr + N0, vl);
+ acc2 = __riscv_vle32_v_f32m4(out_ptr + N0 * 2, vl);
+ acc3 = __riscv_vle32_v_f32m4(out_ptr + N0 * 3, vl);
+ } else {
+ acc0 = __riscv_vfmv_v_f_f32m4(0.0f, vl);
+ acc1 = __riscv_vfmv_v_f_f32m4(0.0f, vl);
+ acc2 = __riscv_vfmv_v_f_f32m4(0.0f, vl);
+ acc3 = __riscv_vfmv_v_f_f32m4(0.0f, vl);
+ }
+ for (int k = 0; k < params->K; ++k) {
+ vbfloat16m2_t rhs = __riscv_vle16_v_bf16m2(rhs_ptr, vl);
+ rhs_ptr += N0;
+ __bf16 lhs0 = *lhs_ptr++;
+ __bf16 lhs1 = *lhs_ptr++;
+ __bf16 lhs2 = *lhs_ptr++;
+ __bf16 lhs3 = *lhs_ptr++;
+
+ acc0 = __riscv_vfwmaccbf16_vf_f32m4(acc0, lhs0, rhs, vl);
+ acc1 = __riscv_vfwmaccbf16_vf_f32m4(acc1, lhs1, rhs, vl);
+ acc2 = __riscv_vfwmaccbf16_vf_f32m4(acc2, lhs2, rhs, vl);
+ acc3 = __riscv_vfwmaccbf16_vf_f32m4(acc3, lhs3, rhs, vl);
+ }
+ __riscv_vse32_v_f32m4(out_ptr, acc0, vl);
+ __riscv_vse32_v_f32m4(out_ptr + N0, acc1, vl);
+ __riscv_vse32_v_f32m4(out_ptr + N0 * 2, acc2, vl);
+ __riscv_vse32_v_f32m4(out_ptr + N0 * 3, acc3, vl);
+ } else if (M0 == 7) {
+ if (params->flags & IREE_UK_FLAG_MMT4D_ACCUMULATE) {
+ acc0 = __riscv_vle32_v_f32m4(out_ptr, vl);
+ acc1 = __riscv_vle32_v_f32m4(out_ptr + N0, vl);
+ acc2 = __riscv_vle32_v_f32m4(out_ptr + N0 * 2, vl);
+ acc3 = __riscv_vle32_v_f32m4(out_ptr + N0 * 3, vl);
+ acc4 = __riscv_vle32_v_f32m4(out_ptr + N0 * 4, vl);
+ acc5 = __riscv_vle32_v_f32m4(out_ptr + N0 * 5, vl);
+ acc6 = __riscv_vle32_v_f32m4(out_ptr + N0 * 6, vl);
+ } else {
+ acc0 = __riscv_vfmv_v_f_f32m4(0.0f, vl);
+ acc1 = __riscv_vfmv_v_f_f32m4(0.0f, vl);
+ acc2 = __riscv_vfmv_v_f_f32m4(0.0f, vl);
+ acc3 = __riscv_vfmv_v_f_f32m4(0.0f, vl);
+ acc4 = __riscv_vfmv_v_f_f32m4(0.0f, vl);
+ acc5 = __riscv_vfmv_v_f_f32m4(0.0f, vl);
+ acc6 = __riscv_vfmv_v_f_f32m4(0.0f, vl);
+ }
+ for (int k = 0; k < params->K; ++k) {
+ vbfloat16m2_t rhs = __riscv_vle16_v_bf16m2(rhs_ptr, vl);
+ rhs_ptr += N0;
+ __bf16 lhs0 = *lhs_ptr++;
+ __bf16 lhs1 = *lhs_ptr++;
+ __bf16 lhs2 = *lhs_ptr++;
+ __bf16 lhs3 = *lhs_ptr++;
+ __bf16 lhs4 = *lhs_ptr++;
+ __bf16 lhs5 = *lhs_ptr++;
+ __bf16 lhs6 = *lhs_ptr++;
+
+ acc0 = __riscv_vfwmaccbf16_vf_f32m4(acc0, lhs0, rhs, vl);
+ acc1 = __riscv_vfwmaccbf16_vf_f32m4(acc1, lhs1, rhs, vl);
+ acc2 = __riscv_vfwmaccbf16_vf_f32m4(acc2, lhs2, rhs, vl);
+ acc3 = __riscv_vfwmaccbf16_vf_f32m4(acc3, lhs3, rhs, vl);
+ acc4 = __riscv_vfwmaccbf16_vf_f32m4(acc4, lhs4, rhs, vl);
+ acc5 = __riscv_vfwmaccbf16_vf_f32m4(acc5, lhs5, rhs, vl);
+ acc6 = __riscv_vfwmaccbf16_vf_f32m4(acc6, lhs6, rhs, vl);
+ }
+ __riscv_vse32_v_f32m4(out_ptr, acc0, vl);
+ __riscv_vse32_v_f32m4(out_ptr + N0, acc1, vl);
+ __riscv_vse32_v_f32m4(out_ptr + N0 * 2, acc2, vl);
+ __riscv_vse32_v_f32m4(out_ptr + N0 * 3, acc3, vl);
+ __riscv_vse32_v_f32m4(out_ptr + N0 * 4, acc4, vl);
+ __riscv_vse32_v_f32m4(out_ptr + N0 * 5, acc5, vl);
+ __riscv_vse32_v_f32m4(out_ptr + N0 * 6, acc6, vl);
+ }
+}
+
+IREE_UK_MMT4D_TILE_FUNC_IMPL_FOR_M0(
+ iree_uk_mmt4d_tile_bf16bf16f32_1xXXx1_to_7xXXx1_riscv_64_zvfbfwma,
+ iree_uk_mmt4d_tile_bf16bf16f32_1xXXx1_riscv_64_zvfbfwma, 1)
+IREE_UK_MMT4D_TILE_FUNC_IMPL_FOR_M0(
+ iree_uk_mmt4d_tile_bf16bf16f32_1xXXx1_to_7xXXx1_riscv_64_zvfbfwma,
+ iree_uk_mmt4d_tile_bf16bf16f32_2xXXx1_riscv_64_zvfbfwma, 2)
+IREE_UK_MMT4D_TILE_FUNC_IMPL_FOR_M0(
+ iree_uk_mmt4d_tile_bf16bf16f32_1xXXx1_to_7xXXx1_riscv_64_zvfbfwma,
+ iree_uk_mmt4d_tile_bf16bf16f32_4xXXx1_riscv_64_zvfbfwma, 4)
+IREE_UK_MMT4D_TILE_FUNC_IMPL_FOR_M0(
+ iree_uk_mmt4d_tile_bf16bf16f32_1xXXx1_to_7xXXx1_riscv_64_zvfbfwma,
+ iree_uk_mmt4d_tile_bf16bf16f32_7xXXx1_riscv_64_zvfbfwma, 7)
diff --git a/runtime/src/iree/builtins/ukernel/tools/mmt4d_test.c b/runtime/src/iree/builtins/ukernel/tools/mmt4d_test.c
index d931467..8b6e17a 100644
--- a/runtime/src/iree/builtins/ukernel/tools/mmt4d_test.c
+++ b/runtime/src/iree/builtins/ukernel/tools/mmt4d_test.c
@@ -615,6 +615,7 @@
iree_uk_test_mmt4d(IREE_UK_FLAG_MMT4D_SKIP_INTERMEDIATE_ROUNDINGS |
IREE_UK_FLAG_MMT4D_TYPE_F16F16F16,
7, 16, 1, "zvfh");
+ iree_uk_test_mmt4d(IREE_UK_FLAG_MMT4D_TYPE_BF16BF16F32, 7, 16, 1, "zvfbfwma");
#endif // defined(IREE_ARCH_ARM_64)
diff --git a/runtime/src/iree/schemas/cpu_feature_bits.inl b/runtime/src/iree/schemas/cpu_feature_bits.inl
index 95e92a3..b634143 100644
--- a/runtime/src/iree/schemas/cpu_feature_bits.inl
+++ b/runtime/src/iree/schemas/cpu_feature_bits.inl
@@ -125,3 +125,4 @@
IREE_CPU_FEATURE_BIT(RISCV_64, 0, 0, V, "v")
IREE_CPU_FEATURE_BIT(RISCV_64, 0, 1, ZVFHMIN, "zvfhmin")
IREE_CPU_FEATURE_BIT(RISCV_64, 0, 2, ZVFH, "zvfh")
+IREE_CPU_FEATURE_BIT(RISCV_64, 0, 3, ZVFBFWMA, "zvfbfwma")
