[RVV] Enable SpaceMiT vendor IME (#24706)
## Summary
This PR adds an s8s8s32 mmt4d ukernel for SpaceMiT IME (`+xsmtvdot` /
`smt.vmadot`)
and wires the compiler to use it. The tile is fixed at 12x16x8 (3x4 grid
of 4x4x8 atoms). Under `+xsmtvdot` we
select that shape and auto-enable mmt4d; plain `+v` targets are
unchanged.
Also, runtime detection can't use the usual ISA hwprobe bits for this
vendor extension, so we set the feature bit from SpaceMiT X60
`mvendorid`/`marchid`. And adds lit coverage, mmt4d unit/benchmark
entries, and an
`riscv_64:ime:+v,+zvl256b,+xsmtvdot` e2e matmul variant. Validated on
Banana Pi (BPI-F3 / X60): cpuinfo, mmt4d_test, and the ukernel e2e path.
Fixes #24576
Co-authored-by: Coden <github.com/coden-s>
Co-authored-by: Andrew <github.com/akuhlken>
Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: Terence <terrencemahlatini@gmail.com>
---------
Signed-off-by: Terence <terrencemahlatini@gmail.com>
Co-authored-by: Cursor <cursoragent@cursor.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 327a069..651353c 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
@@ -242,6 +242,150 @@
// -----
+// RISC-V 64 + xsmtvdot + zvl256b: IME 12x16x8 tile (3x4 of 4x4x8).
+#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 = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]>
+#encoding_rhs = #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]>
+#encoding_result = #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]>
+func.func @matmul_lowering_i8i8i32_riscv64_zvl256b(
+ %lhs: tensor<?x?xi8, #encoding_lhs>,
+ %rhs: tensor<?x?xi8, #encoding_rhs>,
+ %result: tensor<?x?xi32, #encoding_result>
+) -> tensor<?x?xi32, #encoding_result> attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {target_triple="riscv64-xyz-xyz", cpu_features="+v,+zvl256b,+xsmtvdot", iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %out = linalg.matmul
+ ins(%lhs, %rhs : tensor<?x?xi8, #encoding_lhs>,
+ tensor<?x?xi8, #encoding_rhs>)
+ outs(%result : tensor<?x?xi32, #encoding_result>)
+ -> tensor<?x?xi32, #encoding_result>
+ return %out : tensor<?x?xi32, #encoding_result>
+}
+// CHECK-LABEL: func @matmul_lowering_i8i8i32_riscv64_zvl256b(
+// STATIC-SAME: %[[LHS:[a-zA-Z0-9]+]]: tensor<?x?x12x8xi8>
+// STATIC-SAME: %[[RHS:[a-zA-Z0-9]+]]: tensor<?x?x16x8xi8>
+// STATIC-SAME: %[[ACC:[a-zA-Z0-9]+]]: tensor<?x?x12x16xi32>
+// SCALABLE-SAME: %[[LHS:[a-zA-Z0-9]+]]: tensor<?x?x7x1xi8>
+// SCALABLE-SAME: %[[RHS:[a-zA-Z0-9]+]]: tensor<?x?x?x1xi8>
+// SCALABLE-SAME: %[[ACC:[a-zA-Z0-9]+]]: tensor<?x?x7x?xi32>
+// CHECK: %[[MMT4D:.+]] = linalg.mmt4d
+// CHECK-SAME: ins(%[[LHS]], %[[RHS]]
+// CHECK-SAME: outs(%[[ACC]]
+// CHECK: return %[[MMT4D]]
+
+// -----
+
+// RISC-V 64 + xsmtvdot + zvl1024b: IME 24x32x16 tile (3x4 of 8x8x16).
+#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 = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]>
+#encoding_rhs = #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]>
+#encoding_result = #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]>
+func.func @matmul_lowering_i8i8i32_riscv64_zvl1024b(
+ %lhs: tensor<?x?xi8, #encoding_lhs>,
+ %rhs: tensor<?x?xi8, #encoding_rhs>,
+ %result: tensor<?x?xi32, #encoding_result>
+) -> tensor<?x?xi32, #encoding_result> attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {target_triple="riscv64-xyz-xyz", cpu_features="+v,+zvl1024b,+xsmtvdot", iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %out = linalg.matmul
+ ins(%lhs, %rhs : tensor<?x?xi8, #encoding_lhs>,
+ tensor<?x?xi8, #encoding_rhs>)
+ outs(%result : tensor<?x?xi32, #encoding_result>)
+ -> tensor<?x?xi32, #encoding_result>
+ return %out : tensor<?x?xi32, #encoding_result>
+}
+// CHECK-LABEL: func @matmul_lowering_i8i8i32_riscv64_zvl1024b(
+// STATIC-SAME: %[[LHS:[a-zA-Z0-9]+]]: tensor<?x?x24x16xi8>
+// STATIC-SAME: %[[RHS:[a-zA-Z0-9]+]]: tensor<?x?x32x16xi8>
+// STATIC-SAME: %[[ACC:[a-zA-Z0-9]+]]: tensor<?x?x24x32xi32>
+// SCALABLE-SAME: %[[LHS:[a-zA-Z0-9]+]]: tensor<?x?x7x1xi8>
+// SCALABLE-SAME: %[[RHS:[a-zA-Z0-9]+]]: tensor<?x?x?x1xi8>
+// SCALABLE-SAME: %[[ACC:[a-zA-Z0-9]+]]: tensor<?x?x7x?xi32>
+// CHECK: %[[MMT4D:.+]] = linalg.mmt4d
+// CHECK-SAME: ins(%[[LHS]], %[[RHS]]
+// CHECK-SAME: outs(%[[ACC]]
+// CHECK: return %[[MMT4D]]
+
+// -----
+
+// RISC-V 64 + xsmtvdot + zvl4096b: IME 48x64x32 tile (3x4 of 16x16x32).
+#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 = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]>
+#encoding_rhs = #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]>
+#encoding_result = #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]>
+func.func @matmul_lowering_i8i8i32_riscv64_zvl4096b(
+ %lhs: tensor<?x?xi8, #encoding_lhs>,
+ %rhs: tensor<?x?xi8, #encoding_rhs>,
+ %result: tensor<?x?xi32, #encoding_result>
+) -> tensor<?x?xi32, #encoding_result> attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {target_triple="riscv64-xyz-xyz", cpu_features="+v,+zvl4096b,+xsmtvdot", iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %out = linalg.matmul
+ ins(%lhs, %rhs : tensor<?x?xi8, #encoding_lhs>,
+ tensor<?x?xi8, #encoding_rhs>)
+ outs(%result : tensor<?x?xi32, #encoding_result>)
+ -> tensor<?x?xi32, #encoding_result>
+ return %out : tensor<?x?xi32, #encoding_result>
+}
+// CHECK-LABEL: func @matmul_lowering_i8i8i32_riscv64_zvl4096b(
+// STATIC-SAME: %[[LHS:[a-zA-Z0-9]+]]: tensor<?x?x48x32xi8>
+// STATIC-SAME: %[[RHS:[a-zA-Z0-9]+]]: tensor<?x?x64x32xi8>
+// STATIC-SAME: %[[ACC:[a-zA-Z0-9]+]]: tensor<?x?x48x64xi32>
+// SCALABLE-SAME: %[[LHS:[a-zA-Z0-9]+]]: tensor<?x?x7x1xi8>
+// SCALABLE-SAME: %[[RHS:[a-zA-Z0-9]+]]: tensor<?x?x?x1xi8>
+// SCALABLE-SAME: %[[ACC:[a-zA-Z0-9]+]]: tensor<?x?x7x?xi32>
+// CHECK: %[[MMT4D:.+]] = linalg.mmt4d
+// CHECK-SAME: ins(%[[LHS]], %[[RHS]]
+// CHECK-SAME: outs(%[[ACC]]
+// CHECK: return %[[MMT4D]]
+
+// -----
+
+// RISC-V 64 + xsmtvdot + zvl512b: VLEN not in {256, 1024, 4096}, so fall
+// through to standard RVV i8 tiles. IME vmadot is not selected.
+// Static: N0=vlen/8=64. Scalable: N is vscale-based (base VLEN=64).
+#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 = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]>
+#encoding_rhs = #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]>
+#encoding_result = #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]>
+func.func @matmul_lowering_i8i8i32_riscv64_zvl512b_fallback(
+ %lhs: tensor<?x?xi8, #encoding_lhs>,
+ %rhs: tensor<?x?xi8, #encoding_rhs>,
+ %result: tensor<?x?xi32, #encoding_result>
+) -> tensor<?x?xi32, #encoding_result> attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {target_triple="riscv64-xyz-xyz", cpu_features="+v,+zvl512b,+xsmtvdot", iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %out = linalg.matmul
+ ins(%lhs, %rhs : tensor<?x?xi8, #encoding_lhs>,
+ tensor<?x?xi8, #encoding_rhs>)
+ outs(%result : tensor<?x?xi32, #encoding_result>)
+ -> tensor<?x?xi32, #encoding_result>
+ return %out : tensor<?x?xi32, #encoding_result>
+}
+// CHECK-LABEL: func @matmul_lowering_i8i8i32_riscv64_zvl512b_fallback(
+// STATIC-SAME: %[[LHS:[a-zA-Z0-9]+]]: tensor<?x?x7x1xi8>
+// STATIC-SAME: %[[RHS:[a-zA-Z0-9]+]]: tensor<?x?x64x1xi8>
+// STATIC-SAME: %[[ACC:[a-zA-Z0-9]+]]: tensor<?x?x7x64xi32>
+
+// SCALABLE-SAME: %[[LHS:[a-zA-Z0-9]+]]: tensor<?x?x7x1xi8>
+// SCALABLE-SAME: %[[RHS:[a-zA-Z0-9]+]]: tensor<?x?x?x1xi8>
+// SCALABLE-SAME: %[[ACC:[a-zA-Z0-9]+]]: tensor<?x?x7x?xi32>
+
+// CHECK: %[[MMT4D:.+]] = linalg.mmt4d
+// CHECK-SAME: ins(%[[LHS]], %[[RHS]]
+// CHECK-SAME: outs(%[[ACC]]
+// CHECK: return %[[MMT4D]]
+
+// -----
+
// Checks that we don't transpose M/N for the narrow-N case when scalable inner tiles are present.
// For scalable vectors, we want to keep the N dimension scalable.
diff --git a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp
index 7260db7..4ea7b07 100644
--- a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp
+++ b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp
@@ -1240,14 +1240,30 @@
};
}
}
- // 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),
- // fully utilizing the register capacity at any VLEN.
- // M0=7 with LMUL=4 accumulators: 7*4 + 2(prod) + 1(rhs) = 31 regs.
+ // Integer 8 path.
if (lhs.isSignlessInteger(8) && rhs.isSignlessInteger(8) &&
out.isSignlessInteger(32)) {
int N0 = vlen / 8;
+ // SpaceMiT IME (vmadot): fixed hardware tiles for VLEN ∈ {256, 1024,
+ // 4096}. Tile is a 3×4 grid of the SEW=8 MAC atom (4×4×8 / 8×8×16 /
+ // 16×16×32). Other VLENs / scalable mode fall through to standard RVV
+ // tiles.
+ if (hasFeature(config, "+xsmtvdot") && !isScalableVectorizationEnabled()) {
+ int atom = 0;
+ if (vlen == 256) {
+ atom = 4;
+ } else if (vlen == 1024) {
+ atom = 8;
+ } else if (vlen == 4096) {
+ atom = 16;
+ }
+ if (atom != 0) {
+ return {TileMxNxK{3 * atom, 4 * atom, 2 * atom}};
+ }
+ }
+ // Standard RVV widening multiply-accumulate.
+ // Widening chain: i8(m1) -> i16(m2) -> i32(m4).
+ // M0=7 with LMUL=4 accumulators: 7*4 + 2(prod) + 1(rhs) = 31 regs.
return {
TileMxNxK{7, N0, 1}, // Primary shape for RVV widening int8 kernels.
TileMxNxK{4, N0, 1}, // Truncation of the above.
diff --git a/compiler/src/iree/compiler/Codegen/Utils/Utils.cpp b/compiler/src/iree/compiler/Codegen/Utils/Utils.cpp
index 1e30292..34d30e4 100644
--- a/compiler/src/iree/compiler/Codegen/Utils/Utils.cpp
+++ b/compiler/src/iree/compiler/Codegen/Utils/Utils.cpp
@@ -209,6 +209,12 @@
if (isAArch64(targetConfig)) {
return "mmt4d";
}
+ // Auto-enable mmt4d for IME-capable RISC-V targets so callers do not need
+ // --iree-llvmcpu-enable-ukernels=mmt4d explicitly. Plain +v targets are
+ // unaffected; they still require the explicit flag.
+ if (isRISCV64(targetConfig) && hasFeature(targetConfig, "+xsmtvdot")) {
+ return "mmt4d";
+ }
return kNone;
}
diff --git a/runtime/src/iree/base/internal/cpu.c b/runtime/src/iree/base/internal/cpu.c
index 1862851..3ba85de 100644
--- a/runtime/src/iree/base/internal/cpu.c
+++ b/runtime/src/iree/base/internal/cpu.c
@@ -402,14 +402,32 @@
// NOTE: not all kernel versions have all of the constants we need defined so as
// a practice we always define the feature bits we need locally.
+#define IREE_RISCV_HWPROBE_KEY_MVENDORID 0
+#define IREE_RISCV_HWPROBE_KEY_MARCHID 1
#define IREE_RISCV_HWPROBE_KEY_IMA_EXT_0 4
#define IREE_RISCV_HWPROBE_IMA_V (1 << 2)
#define IREE_RISCV_HWPROBE_EXT_ZVFH (1 << 30)
#define IREE_RISCV_HWPROBE_EXT_ZVFHMIN (1ULL << 31)
#define IREE_RISCV_HWPROBE_EXT_ZVFBFWMA (1ULL << 54)
+// SpaceMiT K1 machine identity, used to enable the vendor IME (XSMTVDot /
+// `vmadot`) matrix extension.
+#define IREE_RISCV_MVENDORID_SPACEMIT 0x710ULL
+#define IREE_RISCV_MARCHID_SPACEMIT_X60 0x8000000058000001ULL
+
+// Indices into `kv_pairs` below. hwprobe fills each pair's `value` in place, so
+// these must match the array element order.
+enum {
+ IREE_RISCV_HWPROBE_PAIR_IMA_EXT_0 = 0,
+ IREE_RISCV_HWPROBE_PAIR_MVENDORID,
+ IREE_RISCV_HWPROBE_PAIR_MARCHID,
+};
+
static struct riscv_hwprobe kv_pairs[] = {
- {IREE_RISCV_HWPROBE_KEY_IMA_EXT_0, 0}};
+ {IREE_RISCV_HWPROBE_KEY_IMA_EXT_0, 0},
+ {IREE_RISCV_HWPROBE_KEY_MVENDORID, 0},
+ {IREE_RISCV_HWPROBE_KEY_MARCHID, 0},
+};
long riscv_hwprobe(struct riscv_hwprobe* pairs, size_t pair_count,
size_t cpusetsize, cpu_set_t* cpus, unsigned int flags) {
@@ -418,12 +436,13 @@
}
static void iree_cpu_initialize_from_platform_riscv_64(uint64_t* out_fields) {
- long rc = riscv_hwprobe(&kv_pairs[0], 1, 0, NULL, 0);
+ long rc = riscv_hwprobe(&kv_pairs[0], IREE_ARRAYSIZE(kv_pairs), 0, NULL, 0);
if (rc != 0) {
fprintf(stderr, "riscv_hwprobe syscall failed");
exit(1);
}
- unsigned long long hwprobe = kv_pairs[0].value;
+ unsigned long long hwprobe =
+ kv_pairs[IREE_RISCV_HWPROBE_PAIR_IMA_EXT_0].value;
IREE_COPY_BITS(out_fields[0], IREE_CPU_DATA0_RISCV_64_V, hwprobe,
IREE_RISCV_HWPROBE_IMA_V);
IREE_COPY_BITS(out_fields[0], IREE_CPU_DATA0_RISCV_64_ZVFHMIN, hwprobe,
@@ -432,6 +451,28 @@
IREE_RISCV_HWPROBE_EXT_ZVFH);
IREE_COPY_BITS(out_fields[0], IREE_CPU_DATA0_RISCV_64_ZVFBFWMA, hwprobe,
IREE_RISCV_HWPROBE_EXT_ZVFBFWMA);
+
+ // XSMTVDot (SpaceMiT IME): absent from hwprobe / cpuinfo. Approximating via
+ // X60 mvendorid/marchid until a vendor hwprobe key exists. this host bit
+ // mainly unblocks ukernel tests/benchmarks.
+ unsigned long long mvendorid =
+ kv_pairs[IREE_RISCV_HWPROBE_PAIR_MVENDORID].value;
+ unsigned long long marchid = kv_pairs[IREE_RISCV_HWPROBE_PAIR_MARCHID].value;
+ if (mvendorid == IREE_RISCV_MVENDORID_SPACEMIT &&
+ marchid == IREE_RISCV_MARCHID_SPACEMIT_X60) {
+ out_fields[0] |= IREE_CPU_DATA0_RISCV_64_XSMTVDOT;
+ }
+
+ // No stable hwprobe VLEN key yet; reading vlenb (CSR 0xc22) when V is
+ // present. Sets only the three IME-relevant zvl*b bits (inclusive).
+ if (out_fields[0] & IREE_CPU_DATA0_RISCV_64_V) {
+ uint64_t vlenb = 0;
+ __asm__ volatile("csrr %0, 0xc22" : "=r"(vlenb));
+ uint64_t vlen_bits = vlenb * 8;
+ if (vlen_bits >= 256) out_fields[0] |= IREE_CPU_DATA0_RISCV_64_ZVL256B;
+ if (vlen_bits >= 1024) out_fields[0] |= IREE_CPU_DATA0_RISCV_64_ZVL1024B;
+ if (vlen_bits >= 4096) out_fields[0] |= IREE_CPU_DATA0_RISCV_64_ZVL4096B;
+ }
}
#else
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 02dae9e..79fee62 100644
--- a/runtime/src/iree/builtins/ukernel/arch/riscv_64/BUILD.bazel
+++ b/runtime/src/iree/builtins/ukernel/arch/riscv_64/BUILD.bazel
@@ -80,6 +80,16 @@
)
iree_bitcode_library(
+ name = "ukernel_bitcode_arch_riscv_64_xsmtvdot",
+ srcs = [
+ "mmt4d_riscv_64_xsmtvdot.c",
+ ],
+ arch = "riscv_64",
+ copts = ["-march=rv64gc_xsmtvdot"],
+ internal_hdrs = UKERNEL_RISCV_64_INTERNAL_HEADERS,
+)
+
+iree_bitcode_library(
name = "ukernel_bitcode_arch_riscv_64_zvfbfwma",
srcs = [
"mmt4d_riscv_64_zvfbfwma.c",
@@ -96,6 +106,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_xsmtvdot.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 f155c06..3ba7d6f 100644
--- a/runtime/src/iree/builtins/ukernel/arch/riscv_64/CMakeLists.txt
+++ b/runtime/src/iree/builtins/ukernel/arch/riscv_64/CMakeLists.txt
@@ -93,6 +93,25 @@
iree_bitcode_library(
NAME
+ ukernel_bitcode_arch_riscv_64_xsmtvdot
+ 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_xsmtvdot.c"
+ COPTS
+ "-march=rv64gc_xsmtvdot"
+)
+
+iree_bitcode_library(
+ NAME
ukernel_bitcode_arch_riscv_64_zvfbfwma
ARCH
riscv_64
@@ -116,6 +135,7 @@
SRCS
"ukernel_bitcode_arch_riscv_64_entry_points.bc"
"ukernel_bitcode_arch_riscv_64_v.bc"
+ "ukernel_bitcode_arch_riscv_64_xsmtvdot.bc"
"ukernel_bitcode_arch_riscv_64_zvfbfwma.bc"
"ukernel_bitcode_arch_riscv_64_zvfh.bc"
"ukernel_bitcode_arch_riscv_64_zvfhmin.bc"
@@ -147,6 +167,11 @@
"-march=rv64gc_zvfh"
)
+iree_select_compiler_opts(IREE_UK_COPTS_RISCV_64_XSMTVDOT
+ CLANG_OR_GCC
+ "-march=rv64gc_xsmtvdot"
+)
+
iree_select_compiler_opts(IREE_UK_COPTS_RISCV_64_ZVFBFWMA
CLANG_OR_GCC
"-march=rv64gcv_zvfbfwma"
@@ -154,6 +179,7 @@
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_XSMTVDOT}" IREE_UK_BUILD_RISCV_64_XSMTVDOT)
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
@@ -242,6 +268,20 @@
list(APPEND IREE_UK_RISCV_64_DEPS "::riscv_64_zvfh")
endif() # IREE_UK_BUILD_RISCV_64_ZVFH
+if(IREE_UK_BUILD_RISCV_64_XSMTVDOT)
+iree_cc_library(
+ NAME
+ riscv_64_xsmtvdot
+ SRCS
+ "mmt4d_riscv_64_xsmtvdot.c"
+ COPTS
+ "${IREE_UK_COPTS_RISCV_64_XSMTVDOT}"
+ DEPS
+ iree::builtins::ukernel::internal_headers
+)
+list(APPEND IREE_UK_RISCV_64_DEPS "::riscv_64_xsmtvdot")
+endif() # IREE_UK_BUILD_RISCV_64_XSMTVDOT
+
if(IREE_UK_BUILD_RISCV_64_ZVFBFWMA)
iree_cc_library(
NAME
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 b68d961..e864e59 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_XSMTVDOT
#define IREE_UK_BUILD_RISCV_64_ZVFBFWMA
#else
// Compiling with the system toolchain. Include the configured header.
@@ -39,6 +40,32 @@
return iree_uk_all_bits_set(cpu_data[0], IREE_CPU_DATA0_RISCV_64_ZVFH);
}
+static inline bool iree_uk_cpu_riscv_64_xsmtvdot(
+ const iree_uk_uint64_t* cpu_data) {
+ return iree_uk_all_bits_set(cpu_data[0], IREE_CPU_DATA0_RISCV_64_XSMTVDOT);
+}
+
+// IME tile shapes require a full vmadot operand panel in one vector register:
+// 12x16x8 needs VLEN>=256, 24x32x16 needs VLEN>=1024, 48x64x32 needs
+// VLEN>=4096.
+static inline bool iree_uk_cpu_riscv_64_xsmtvdot_zvl256b(
+ const iree_uk_uint64_t* cpu_data) {
+ return iree_uk_cpu_riscv_64_xsmtvdot(cpu_data) &&
+ iree_uk_all_bits_set(cpu_data[0], IREE_CPU_DATA0_RISCV_64_ZVL256B);
+}
+
+static inline bool iree_uk_cpu_riscv_64_xsmtvdot_zvl1024b(
+ const iree_uk_uint64_t* cpu_data) {
+ return iree_uk_cpu_riscv_64_xsmtvdot(cpu_data) &&
+ iree_uk_all_bits_set(cpu_data[0], IREE_CPU_DATA0_RISCV_64_ZVL1024B);
+}
+
+static inline bool iree_uk_cpu_riscv_64_xsmtvdot_zvl4096b(
+ const iree_uk_uint64_t* cpu_data) {
+ return iree_uk_cpu_riscv_64_xsmtvdot(cpu_data) &&
+ iree_uk_all_bits_set(cpu_data[0], IREE_CPU_DATA0_RISCV_64_ZVL4096B);
+}
+
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);
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 bee185d..a1dbaab 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
@@ -15,5 +15,6 @@
#cmakedefine IREE_UK_BUILD_RISCV_64_ZVFHMIN
#cmakedefine IREE_UK_BUILD_RISCV_64_ZVFH
#cmakedefine IREE_UK_BUILD_RISCV_64_ZVFBFWMA
+#cmakedefine IREE_UK_BUILD_RISCV_64_XSMTVDOT
#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 bdde1c5..4f9026d 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,19 @@
#define IREE_UK_MMT4D_TILE_riscv_64_zvfh(lhs, rhs, out, m0, k0)
#endif
+#ifdef IREE_UK_BUILD_RISCV_64_XSMTVDOT
+#define IREE_UK_MMT4D_TILE_riscv_64_xsmtvdot_zvl256b(lhs, rhs, out, m0, k0) \
+ IREE_UK_MMT4D_TILE_IMPL_riscv_64(lhs, rhs, out, m0, k0, _xsmtvdot_zvl256b)
+#define IREE_UK_MMT4D_TILE_riscv_64_xsmtvdot_zvl1024b(lhs, rhs, out, m0, k0) \
+ IREE_UK_MMT4D_TILE_IMPL_riscv_64(lhs, rhs, out, m0, k0, _xsmtvdot_zvl1024b)
+#define IREE_UK_MMT4D_TILE_riscv_64_xsmtvdot_zvl4096b(lhs, rhs, out, m0, k0) \
+ IREE_UK_MMT4D_TILE_IMPL_riscv_64(lhs, rhs, out, m0, k0, _xsmtvdot_zvl4096b)
+#else
+#define IREE_UK_MMT4D_TILE_riscv_64_xsmtvdot_zvl256b(lhs, rhs, out, m0, k0)
+#define IREE_UK_MMT4D_TILE_riscv_64_xsmtvdot_zvl1024b(lhs, rhs, out, m0, k0)
+#define IREE_UK_MMT4D_TILE_riscv_64_xsmtvdot_zvl4096b(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)
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 c67baa5..af48484 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
@@ -38,6 +38,10 @@
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, s8, s8, s32, 12, 8, _xsmtvdot_zvl256b)
+IREE_UK_MMT4D_TILE(riscv_64, s8, s8, s32, 24, 16, _xsmtvdot_zvl1024b)
+IREE_UK_MMT4D_TILE(riscv_64, s8, s8, s32, 48, 32, _xsmtvdot_zvl4096b)
+
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)
diff --git a/runtime/src/iree/builtins/ukernel/arch/riscv_64/mmt4d_riscv_64_xsmtvdot.c b/runtime/src/iree/builtins/ukernel/arch/riscv_64/mmt4d_riscv_64_xsmtvdot.c
new file mode 100644
index 0000000..08c84b9
--- /dev/null
+++ b/runtime/src/iree/builtins/ukernel/arch/riscv_64/mmt4d_riscv_64_xsmtvdot.c
@@ -0,0 +1,154 @@
+// 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
+
+// SpaceMiT IME (XSMTVDot) s8s8s32 mmt4d.
+//
+// One `smt.vmadot` at SEW=8 is a square MAC atom selected by vl (not VLMAX):
+// atom=4 (vl*SEW=256) -> 4x4x8 -> tile 12x16x8
+// atom=8 (vl*SEW=1024) -> 8x8x16 -> tile 24x32x16
+// atom=16 (vl*SEW=4096) -> 16x16x32 -> tile 48x64x32
+// The macro-tile is a fixed 3x4 atom grid (12 vmadot per K panel). Each
+// {vd,vd+1} int32 accumulator group holds one atom x atom output block,
+// row-major and contiguous within the group, and stays resident across K.
+//
+// Panels:
+// lhs : int8 [K1][M0][K0] K0-contiguous
+// rhs : int8 [K1][N0][K0] transposed (N0,K0), K0-contiguous
+// out : int32 [M0][N0]
+
+#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"
+
+enum { IME_MT = 3, IME_NT = 4 };
+
+// Byte-offset index vector mapping accumulator-group element i (row-major
+// atom x atom, i = r*atom + c) to its position in the strided (M0,N0) output
+// block: out offset (in elements) = r*N0 + c = i + (N0 - atom) * (i / atom).
+static inline vuint32m2_t ime_out_index(int N0, int atom, int log2_atom,
+ size_t vl32) {
+ vuint32m2_t vi = __riscv_vid_v_u32m2(vl32);
+ vuint32m2_t row = __riscv_vsrl_vx_u32m2(vi, (size_t)log2_atom, vl32);
+ vuint32m2_t off =
+ __riscv_vmacc_vx_u32m2(vi, (uint32_t)(N0 - atom), row, vl32);
+ return __riscv_vsll_vx_u32m2(off, 2, vl32);
+}
+
+// Shared 3x4-atom xsmtvdot body. The macro-tile is a fixed 3x4 grid of the
+// SEW=8 MAC atom, so M0 alone determines the atom edge (4 / 8 / 16) and hence
+// N0 = 4*atom and K0 = 2*atom.
+IREE_UK_ATTRIBUTE_ALWAYS_INLINE static inline void
+iree_uk_mmt4d_tile_s8s8s32_12xXXx8_to_48xXXx32_riscv_64_xsmtvdot(
+ 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) {
+ const int atom = M0 / IME_MT;
+ const int N0 = IME_NT * atom;
+ const int K0 = 2 * atom;
+ IREE_UK_ASSERT(atom == 4 || atom == 8 || atom == 16);
+ IREE_UK_ASSERT(params->M0 == M0);
+ IREE_UK_ASSERT(params->N0 == N0);
+ IREE_UK_ASSERT(params->K0 == K0);
+
+ const iree_uk_int8_t* IREE_UK_RESTRICT lhs = lhs_panel;
+ const iree_uk_int8_t* IREE_UK_RESTRICT rhs = rhs_panel;
+ iree_uk_int32_t* IREE_UK_RESTRICT out = out_tile;
+
+ const int K1 = (int)(params->K);
+ const int accumulate = (params->flags & IREE_UK_FLAG_MMT4D_ACCUMULATE) != 0;
+ const int atom_elems = atom * atom;
+ const int panel_stride = atom * K0;
+ const int log2_atom = __builtin_ctz((unsigned int)atom);
+ const size_t vl8 = (size_t)panel_stride;
+ const size_t vl32 = (size_t)atom_elems;
+
+ // out_tile element address of accumulator group (mt,nt), where
+ // acc<mt*IME_NT+nt> holds A[mt] . B[nt].
+#define IME_OUT(mt, nt) (out + ((mt) * atom) * N0 + (nt) * atom)
+
+ vint32m2_t acc0, acc1, acc2, acc3, acc4, acc5, acc6, acc7, acc8, acc9, acc10,
+ acc11;
+ if (accumulate) {
+ vuint32m2_t idx = ime_out_index(N0, atom, log2_atom, vl32);
+ acc0 = __riscv_vluxei32_v_i32m2(IME_OUT(0, 0), idx, vl32);
+ acc1 = __riscv_vluxei32_v_i32m2(IME_OUT(0, 1), idx, vl32);
+ acc2 = __riscv_vluxei32_v_i32m2(IME_OUT(0, 2), idx, vl32);
+ acc3 = __riscv_vluxei32_v_i32m2(IME_OUT(0, 3), idx, vl32);
+ acc4 = __riscv_vluxei32_v_i32m2(IME_OUT(1, 0), idx, vl32);
+ acc5 = __riscv_vluxei32_v_i32m2(IME_OUT(1, 1), idx, vl32);
+ acc6 = __riscv_vluxei32_v_i32m2(IME_OUT(1, 2), idx, vl32);
+ acc7 = __riscv_vluxei32_v_i32m2(IME_OUT(1, 3), idx, vl32);
+ acc8 = __riscv_vluxei32_v_i32m2(IME_OUT(2, 0), idx, vl32);
+ acc9 = __riscv_vluxei32_v_i32m2(IME_OUT(2, 1), idx, vl32);
+ acc10 = __riscv_vluxei32_v_i32m2(IME_OUT(2, 2), idx, vl32);
+ acc11 = __riscv_vluxei32_v_i32m2(IME_OUT(2, 3), idx, vl32);
+ } else {
+ acc0 = __riscv_vmv_v_x_i32m2(0, vl32);
+ acc1 = acc2 = acc3 = acc4 = acc5 = acc6 = acc7 = acc8 = acc9 = acc10 =
+ acc11 = acc0;
+ }
+
+ for (int k = 0; k < K1; ++k) {
+ vint8m1_t a0 = __riscv_vle8_v_i8m1(lhs + 0 * panel_stride, vl8);
+ vint8m1_t a1 = __riscv_vle8_v_i8m1(lhs + 1 * panel_stride, vl8);
+ vint8m1_t a2 = __riscv_vle8_v_i8m1(lhs + 2 * panel_stride, vl8);
+ vint8m1_t b0 = __riscv_vle8_v_i8m1(rhs + 0 * panel_stride, vl8);
+ vint8m1_t b1 = __riscv_vle8_v_i8m1(rhs + 1 * panel_stride, vl8);
+ vint8m1_t b2 = __riscv_vle8_v_i8m1(rhs + 2 * panel_stride, vl8);
+ vint8m1_t b3 = __riscv_vle8_v_i8m1(rhs + 3 * panel_stride, vl8);
+ lhs += IME_MT * panel_stride;
+ rhs += IME_NT * panel_stride;
+
+ __asm__ volatile(
+ " vsetvli x0, %[vl], e8, m1, ta, ma \n\t"
+ " smt.vmadot %[c0], %[a0], %[b0] \n\t"
+ " smt.vmadot %[c1], %[a0], %[b1] \n\t"
+ " smt.vmadot %[c2], %[a0], %[b2] \n\t"
+ " smt.vmadot %[c3], %[a0], %[b3] \n\t"
+ " smt.vmadot %[c4], %[a1], %[b0] \n\t"
+ " smt.vmadot %[c5], %[a1], %[b1] \n\t"
+ " smt.vmadot %[c6], %[a1], %[b2] \n\t"
+ " smt.vmadot %[c7], %[a1], %[b3] \n\t"
+ " smt.vmadot %[c8], %[a2], %[b0] \n\t"
+ " smt.vmadot %[c9], %[a2], %[b1] \n\t"
+ " smt.vmadot %[c10], %[a2], %[b2] \n\t"
+ " smt.vmadot %[c11], %[a2], %[b3] \n\t"
+ : [c0] "+vr"(acc0), [c1] "+vr"(acc1), [c2] "+vr"(acc2),
+ [c3] "+vr"(acc3), [c4] "+vr"(acc4), [c5] "+vr"(acc5),
+ [c6] "+vr"(acc6), [c7] "+vr"(acc7), [c8] "+vr"(acc8),
+ [c9] "+vr"(acc9), [c10] "+vr"(acc10), [c11] "+vr"(acc11)
+ : [a0] "vr"(a0), [a1] "vr"(a1), [a2] "vr"(a2), [b0] "vr"(b0),
+ [b1] "vr"(b1), [b2] "vr"(b2), [b3] "vr"(b3), [vl] "r"(vl8)
+ :);
+ }
+
+ vuint32m2_t idx = ime_out_index(N0, atom, log2_atom, vl32);
+ __riscv_vsuxei32_v_i32m2(IME_OUT(0, 0), idx, acc0, vl32);
+ __riscv_vsuxei32_v_i32m2(IME_OUT(0, 1), idx, acc1, vl32);
+ __riscv_vsuxei32_v_i32m2(IME_OUT(0, 2), idx, acc2, vl32);
+ __riscv_vsuxei32_v_i32m2(IME_OUT(0, 3), idx, acc3, vl32);
+ __riscv_vsuxei32_v_i32m2(IME_OUT(1, 0), idx, acc4, vl32);
+ __riscv_vsuxei32_v_i32m2(IME_OUT(1, 1), idx, acc5, vl32);
+ __riscv_vsuxei32_v_i32m2(IME_OUT(1, 2), idx, acc6, vl32);
+ __riscv_vsuxei32_v_i32m2(IME_OUT(1, 3), idx, acc7, vl32);
+ __riscv_vsuxei32_v_i32m2(IME_OUT(2, 0), idx, acc8, vl32);
+ __riscv_vsuxei32_v_i32m2(IME_OUT(2, 1), idx, acc9, vl32);
+ __riscv_vsuxei32_v_i32m2(IME_OUT(2, 2), idx, acc10, vl32);
+ __riscv_vsuxei32_v_i32m2(IME_OUT(2, 3), idx, acc11, vl32);
+
+#undef IME_OUT
+}
+
+IREE_UK_MMT4D_TILE_FUNC_IMPL_FOR_M0(
+ iree_uk_mmt4d_tile_s8s8s32_12xXXx8_to_48xXXx32_riscv_64_xsmtvdot,
+ iree_uk_mmt4d_tile_s8s8s32_12xXXx8_riscv_64_xsmtvdot_zvl256b, 12)
+IREE_UK_MMT4D_TILE_FUNC_IMPL_FOR_M0(
+ iree_uk_mmt4d_tile_s8s8s32_12xXXx8_to_48xXXx32_riscv_64_xsmtvdot,
+ iree_uk_mmt4d_tile_s8s8s32_24xXXx16_riscv_64_xsmtvdot_zvl1024b, 24)
+IREE_UK_MMT4D_TILE_FUNC_IMPL_FOR_M0(
+ iree_uk_mmt4d_tile_s8s8s32_12xXXx8_to_48xXXx32_riscv_64_xsmtvdot,
+ iree_uk_mmt4d_tile_s8s8s32_48xXXx32_riscv_64_xsmtvdot_zvl4096b, 48)
diff --git a/runtime/src/iree/builtins/ukernel/tools/mmt4d_benchmark.c b/runtime/src/iree/builtins/ukernel/tools/mmt4d_benchmark.c
index 7474bb5..c0f631c 100644
--- a/runtime/src/iree/builtins/ukernel/tools/mmt4d_benchmark.c
+++ b/runtime/src/iree/builtins/ukernel/tools/mmt4d_benchmark.c
@@ -5,6 +5,7 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#include <stdio.h>
+#include <string.h>
#include "iree/base/api.h"
#include "iree/base/tooling/flags.h"
@@ -110,10 +111,12 @@
int N0, int K0,
const char* cpu_features) {
// Test narrowed, power-of-two values of M0, as mmt4d kernels tend to have
- // narrow variants for handling these cases.
- for (int narrowM0 = 1; narrowM0 < M0; narrowM0 *= 2) {
- iree_uk_benchmark_register_mmt4d_impl(flags, narrowM0, N0, K0, cpu_features,
- "");
+ // narrow variants for handling these cases. xsmtvdot has no narrow variants.
+ if (!strstr(cpu_features, "xsmtvdot")) {
+ for (int narrowM0 = 1; narrowM0 < M0; narrowM0 *= 2) {
+ iree_uk_benchmark_register_mmt4d_impl(flags, narrowM0, N0, K0,
+ cpu_features, "");
+ }
}
iree_uk_benchmark_register_mmt4d_impl(flags, M0, N0, K0, cpu_features, "");
}
@@ -199,6 +202,12 @@
IREE_UK_FLAG_MMT4D_SKIP_INTERMEDIATE_ROUNDINGS |
IREE_UK_FLAG_MMT4D_TYPE_F16F16F16,
7, 16, 1, "zvfh");
+ iree_uk_benchmark_register_mmt4d(IREE_UK_FLAG_MMT4D_TYPE_S8S8S32, 12, 16, 8,
+ "xsmtvdot,zvl256b");
+ iree_uk_benchmark_register_mmt4d(IREE_UK_FLAG_MMT4D_TYPE_S8S8S32, 24, 32, 16,
+ "xsmtvdot,zvl1024b");
+ iree_uk_benchmark_register_mmt4d(IREE_UK_FLAG_MMT4D_TYPE_S8S8S32, 48, 64, 32,
+ "xsmtvdot,zvl4096b");
#else // defined(IREE_ARCH_ARM_64)
// Architectures on which we do not have any optimized ukernel code.
// Benchmark some arbitrary tile shape.
diff --git a/runtime/src/iree/builtins/ukernel/tools/mmt4d_test.c b/runtime/src/iree/builtins/ukernel/tools/mmt4d_test.c
index 8b6e17a..68bcc1b 100644
--- a/runtime/src/iree/builtins/ukernel/tools/mmt4d_test.c
+++ b/runtime/src/iree/builtins/ukernel/tools/mmt4d_test.c
@@ -4,6 +4,8 @@
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+#include <string.h>
+
#include "iree/base/api.h"
#include "iree/base/internal/math.h"
#include "iree/builtins/ukernel/api.h"
@@ -532,9 +534,11 @@
// currently an implementation detail within each arch/ subdirectory.
flags |= IREE_UK_FLAG_MMT4D_ALLOW_GENERIC_FALLBACK_TILE_FUNCTION;
// Test narrowed, power-of-two values of M0, as mmt4d kernels tend to have
- // narrow variants for handling these cases.
- for (int narrowM0 = 1; narrowM0 < M0; narrowM0 *= 2) {
- iree_uk_test_mmt4d_impl(flags, narrowM0, N0, K0, cpu_features);
+ // narrow variants for handling these cases. xsmtvdot has no narrow variants.
+ if (!strstr(cpu_features, "xsmtvdot")) {
+ for (int narrowM0 = 1; narrowM0 < M0; narrowM0 *= 2) {
+ iree_uk_test_mmt4d_impl(flags, narrowM0, N0, K0, cpu_features);
+ }
}
iree_uk_test_mmt4d_impl(flags, M0, N0, K0, cpu_features);
}
@@ -615,6 +619,12 @@
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_S8S8S32, 12, 16, 8,
+ "xsmtvdot,zvl256b");
+ iree_uk_test_mmt4d(IREE_UK_FLAG_MMT4D_TYPE_S8S8S32, 24, 32, 16,
+ "xsmtvdot,zvl1024b");
+ iree_uk_test_mmt4d(IREE_UK_FLAG_MMT4D_TYPE_S8S8S32, 48, 64, 32,
+ "xsmtvdot,zvl4096b");
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 b634143..b0aa262 100644
--- a/runtime/src/iree/schemas/cpu_feature_bits.inl
+++ b/runtime/src/iree/schemas/cpu_feature_bits.inl
@@ -126,3 +126,9 @@
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")
+IREE_CPU_FEATURE_BIT(RISCV_64, 0, 4, XSMTVDOT, "xsmtvdot")
+// Minimum VLEN guarantees (LLVM zvl*b) used to gate SpaceMiT IME
+// tile shapes that require a full vmadot operand panel in one vector register.
+IREE_CPU_FEATURE_BIT(RISCV_64, 0, 5, ZVL256B, "zvl256b")
+IREE_CPU_FEATURE_BIT(RISCV_64, 0, 6, ZVL1024B, "zvl1024b")
+IREE_CPU_FEATURE_BIT(RISCV_64, 0, 7, ZVL4096B, "zvl4096b")
diff --git a/tests/e2e/matmul/BUILD.bazel b/tests/e2e/matmul/BUILD.bazel
index cf42715..ef83b21 100644
--- a/tests/e2e/matmul/BUILD.bazel
+++ b/tests/e2e/matmul/BUILD.bazel
@@ -176,6 +176,7 @@
"arm_64:dotprod:+dotprod",
"arm_64:i8mm:+i8mm",
"x86_64:avx512vnni:" + ",".join(X86_64_AVX512_VNNI),
+ "riscv_64:ime:+v,+zvl256b,+xsmtvdot",
] if lhs_rhs_type == "i8" and acc_type == "i32" else [
"x86_64:avx2:" + ",".join(X86_64_AVX2),
"x86_64:avx512:" + ",".join(X86_64_AVX512),
diff --git a/tests/e2e/matmul/CMakeLists.txt b/tests/e2e/matmul/CMakeLists.txt
index 40f516d..5e93079 100644
--- a/tests/e2e/matmul/CMakeLists.txt
+++ b/tests/e2e/matmul/CMakeLists.txt
@@ -144,6 +144,7 @@
"arm_64:dotprod:+dotprod"
"arm_64:i8mm:+i8mm"
"x86_64:avx512vnni:+avx,+avx2,+fma,+f16c,+avx512f,+avx512vl,+avx512cd,+avx512bw,+avx512dq,+avx512vnni"
+ "riscv_64:ime:+v,+zvl256b,+xsmtvdot"
)
iree_generated_e2e_runner_test(
@@ -344,6 +345,7 @@
"arm_64:dotprod:+dotprod"
"arm_64:i8mm:+i8mm"
"x86_64:avx512vnni:+avx,+avx2,+fma,+f16c,+avx512f,+avx512vl,+avx512cd,+avx512bw,+avx512dq,+avx512vnni"
+ "riscv_64:ime:+v,+zvl256b,+xsmtvdot"
)
iree_generated_e2e_runner_test(
