[RVV] Enable SpaceMiT vendor IME by TMahlatini · Pull Request #24706 · iree-org/iree

TMahlatini · Jul 15, 2026

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

github-actions · Jul 15, 2026

Hello @TMahlatini 👋

Thank you for submitting a Pull Request to IREE! It looks like this is your first one. Below are some useful links and pointers.

General guidance

Our general Contributing guide contains information and links to detailed guides on code quality, testing, commit summaries and our CI system.

A common point for new PRs: if a DCO signing check fails for you, check out the section on Developer Certificate of Origin.
In these cases, it should suffice to amend your commit signature(s) per the guide and force-push the PR branch.

Action required: acknowledge IREE project policies

IREE is a Linux Foundation project. All participants are expected to follow the LF Projects Code of Conduct.

Please also note that all contributions to IREE must follow our IREE AI Tool Use Policy. In particular:

Contributors must fully understand, and vouch for, all submitted changes and the intent behind them.
Contributors should write PR descriptions themselves.
Substantial use of LLM/generative AI tools must be noted in the PR description, e.g. via Assisted-by: tool-name or Co-authored-by: tool-name tool@email trailers.
GitHub issues labeled as "Good first issue" are explicitly designated as learning opportunities for newcomers to the project. We discourage AI tool usage for resolutions to those issues, and substantial use (e.g. fully prompting the fix and/or the tests) is forbidden.

We kindly ask you to reply to this message and confirm that you understand and accept the cited policies, particularly the AI Tool Use Policy.

If you have any questions, feel free to leave a comment here, or ask away on IREE Discord.

Thank you,
The IREE Community

Add xsmtvdot microkernel: 12x16x8. Updated contracts, build, tests and benchmarks to validate the kernel. Fixes: iree-org#24576 Co-authored-by: Coden <github.com/coden-s> Co-authored-by: Andrew <github.com/akuhlken> Signed-off-by: Terence <terrencemahlatini@gmail.com> Co-authored-by: Cursor <cursoragent@cursor.com>

TMahlatini · Jul 15, 2026

I understand and accept the cited policies, including the AI Tool Use Policy.

- Added definitions for SpaceMiT vendor and architecture IDs to support the XSMTVDot matrix extension. - Modified the `iree_cpu_initialize_from_platform_riscv_64` function to detect the presence of the XSMTVDot extension based on the machine identity. Fixes: iree-org#24576 Signed-off-by: Terence <terrencemahlatini@gmail.com> Co-authored-by: Cursor <cursoragent@cursor.com>

egebeysel

Initial comments from my side, haven't checked the plumbing nor the ukernel itself completely yet, though I have an intiial comment regarding support for the VLEN that SpaceMiT IME allows. I think it should work quite easily without changing your ukernel much, but let me know if that's not the case.

egebeysel · Jul 17, 2026

+    // SpaceMiT IME (vmadot): fixed tile shape derived from the 4×4×8 atom.
+    // M0=12 (3×4 atom grid) matches the registered runtime ukernel tile.
+    if (hasFeature(config, "+xsmtvdot")) {
+      return {TileMxNxK{12, 16, 8}};


could you remind me why we didn't have the 8, 4 fallback sizes for M0 here?

we prototyped M0=8 and 4 earlier, but weren’t confident in the performance. Our early structure put the K loop in C around separate asm blocks, so vector accumulators weren’t held live across K and the results were hard to trust. For this PR we intentionally only register the M0=12 primary. Happy to follow up with M0=8/4 using the same single-asm structure (2×4 and 1×4 atom grids).

egebeysel · Jul 17, 2026

+    // SpaceMiT IME (vmadot): fixed tile shape derived from the 4×4×8 atom.
+    // M0=12 (3×4 atom grid) matches the registered runtime ukernel tile.
+    if (hasFeature(config, "+xsmtvdot")) {
+      return {TileMxNxK{12, 16, 8}};


Also, the SpaceMiT IME extension technically allows for vector lengths from 128 to 4096 as powers of two. I know the board you has VLEN=256bits, but we could easily support other powers of two as well.

Though there is a caveat: depending on the VLEN, the layout that the instruction operates on changes, i.e. there is a new variable called copy or Copies that can be 1 or 2:

=1 for VLEN=256, 1024, 4096

=2 for VLEN=128, 512, 2048

Your microkernel and your layout should be able to cover the =1 case, so you can maybe take the vector length from the +zvl*b features and set the M and N dims of the layout accordingly, i.e. instead of 3*4 x 4*4 x 8 we can set 3*(n0/8) x 4*(n0/8) x 8 AND bail out (don't pick vmadot) if we have the =2 case. We can add support for that later 👍 The microkernel should not change regardless, only the input layout would be different and more complex, but that is currently not representable cleanly with what we have with mmt4d at the moment. Reach out to me if you want more details from the RISC-V Discord channel of IREE.

The SpaceMiT IME documentation is here: https://github.com/spacemit-com/riscv-ime-extension-spec/releases/tag/v1.0, check section 3.1.1.

(I think the sentence that starts with "Input data is stored in VS1 and VS1+1 (the index of VS1 must be even)" for the non-sliding vmadot is wrong, you don't have VS1 + 1 involved here, the layout figure is correct though)

Thanks. I will update this PR so that it covers Copies=1 and fall through when Copies=2. Agreed on the non-sliding vmadot comment.

Also, the SpaceMiT IME extension technically allows for vector lengths from 128 to 4096 as powers of two. I know the board you has VLEN=256bits, but we could easily support other powers of two as well.

Though there is a caveat: depending on the VLEN, the layout that the instruction operates on changes, i.e. there is a new variable called copy or Copies that can be 1 or 2:

* `=1` for `VLEN=256, 1024, 4096` * `=2` for `VLEN=128, 512, 2048`

Your microkernel and your layout should be able to cover the =1 case, so you can maybe take the vector length from the +zvl*b features and set the M and N dims of the layout accordingly, i.e. instead of 3*4 x 4*4 x 8 we can set 3*(n0/8) x 4*(n0/8) x 8 AND bail out (don't pick vmadot) if we have the =2 case. We can add support for that later 👍 The microkernel should not change regardless, only the input layout would be different and more complex, but that is currently not representable cleanly with what we have with mmt4d at the moment. Reach out to me if you want more details from the RISC-V Discord channel of IREE.

Latest commit extends support to 1024, 4096 however only 256 has been validated on hardware.

egebeysel · Jul 17, 2026

+  // 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.


jfyi, there's #24620 that wants to enable mmt4d ukernels for RVV regardless, and I think it's fine to do so. No action item needed from your side.

egebeysel · Jul 17, 2026

 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)


todo for myself (and you can also check this 😄), what happens when M0 can also take multiple values? (i.e. dependent on VLEN)

…gurations that cover vlen=1024 and 4096. Signed-off-by: Terence <terrencemahlatini@gmail.com>

Re-add IME tile cases dropped when resolving the main merge conflict. Signed-off-by: Terence <terrencemahlatini@gmail.com>

egebeysel · Jul 21, 2026

Reflecting the offline feedback here:

hey, I was going to try and give it another go tomorrow, but now that you're at it: I would actually try and contain the inline_asm to the instructions for which you do not have intrinsics for: if I'm not mistaken, that is just the vmadot's.

If you do that, you can also use IREE_UK_UNROLL and declutter the code a bit:

iree/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_i8mm.c

Lines 56 to 57 in e53a335

    
           IREE_UK_UNROLL for (int i = 0; i < mtiles; ++i) { 
        
             IREE_UK_UNROLL for (int j = 0; j < 4; ++j) { acc[i][j] = vdupq_n_s32(0); }

Can also wrap them around like this:

iree/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_fp16fml.c

Lines 19 to 26 in e53a335

    
           #define iree_vfmlalq_laneq_x_f16(INSTR, A, B, C, L) \ 
        
             ({                                                \ 
        
               asm(INSTR " %[a].4s, %[b].4h, %[c].h[" #L "]"   \ 
        
                   : [a] "+w"(A)                               \ 
        
                   : [b] "w"(B), [c] "x"(C)                    \ 
        
                   :);                                         \ 
        
               A;                                              \ 
        
             })

As far as I can see, you are not doing much hand-scheduling that would require a full inline assembly kernel with pipelining and stuff (not saying you should). In that case, it's fine to keep the inline asm to the minimum and wrap it around.

Finally: for the scatter/gather of the accumulator, AArch64 has these zip/unzip-like instructions. I think RVV also has similar instructions or shuffles, maybe it's worth checking that out :slight_smile: You can do that directly without storing and reading after the K loop:

iree/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_i8mm.c

Lines 83 to 96 in e53a335

    
             // Swizzle accumulator 2x2 register tiles back to row-major and store. 
        
             IREE_UK_UNROLL for (int i = 0; i < mtiles; ++i) { 
        
               IREE_UK_UNROLL for (int j = 0; j < 2; ++j) { 
        
                 int32x4_t acc_1x4_0 = 
        
                     iree_uk_neon_uzp1_s32_as_s64(acc[i][2 * j + 0], acc[i][2 * j + 1]); 
        
                 vst1q_s32(out_ptr + 8 * (2 * i + 0) + 4 * j, acc_1x4_0); 
        
                 if (M0 > 1) { 
        
                   int32x4_t acc_1x4_1 = 
        
                       iree_uk_neon_uzp2_s32_as_s64(acc[i][2 * j + 0], acc[i][2 * j + 1]); 
        
                   vst1q_s32(out_ptr + 8 * (2 * i + 1) + 4 * j, acc_1x4_1); 
        
                 } 
        
               } 
        
             } 
        
           }

TMahlatini · Jul 21, 2026

Reflecting the offline feedback here:

hey, I was going to try and give it another go tomorrow, but now that you're at it: I would actually try and contain the inline_asm to the instructions for which you do not have intrinsics for: if I'm not mistaken, that is just the vmadot's.

If you do that, you can also use IREE_UK_UNROLL and declutter the code a bit:

iree/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_i8mm.c

Lines 56 to 57 in e53a335

IREE_UK_UNROLL for (int i = 0; i < mtiles; ++i) {

IREE_UK_UNROLL for (int j = 0; j < 4; ++j) { acc[i][j] = vdupq_n_s32(0); }

Can also wrap them around like this:

iree/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_fp16fml.c

Lines 19 to 26 in e53a335

#define iree_vfmlalq_laneq_x_f16(INSTR, A, B, C, L) \

({ \

asm(INSTR " %[a].4s, %[b].4h, %[c].h[" #L "]" \

: [a] "+w"(A) \

: [b] "w"(B), [c] "x"(C) \

:); \

A; \

})

As far as I can see, you are not doing much hand-scheduling that would require a full inline assembly kernel with pipelining and stuff (not saying you should). In that case, it's fine to keep the inline asm to the minimum and wrap it around.

Finally: for the scatter/gather of the accumulator, AArch64 has these zip/unzip-like instructions. I think RVV also has similar instructions or shuffles, maybe it's worth checking that out :slight_smile: You can do that directly without storing and reading after the K loop:

iree/runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_i8mm.c

Lines 83 to 96 in e53a335

// Swizzle accumulator 2x2 register tiles back to row-major and store.

IREE_UK_UNROLL for (int i = 0; i < mtiles; ++i) {

IREE_UK_UNROLL for (int j = 0; j < 2; ++j) {

int32x4_t acc_1x4_0 =

iree_uk_neon_uzp1_s32_as_s64(acc[i][2 * j + 0], acc[i][2 * j + 1]);

vst1q_s32(out_ptr + 8 * (2 * i + 0) + 4 * j, acc_1x4_0);

if (M0 > 1) {

int32x4_t acc_1x4_1 =

iree_uk_neon_uzp2_s32_as_s64(acc[i][2 * j + 0], acc[i][2 * j + 1]);

vst1q_s32(out_ptr + 8 * (2 * i + 1) + 4 * j, acc_1x4_1);

}

}

}

}

Will work on it and update the PR.

TMahlatini changed the title ~~YOUR TITLE~~ [RVV] Enable SpaceMiT vendor IME Jul 15, 2026

TMahlatini force-pushed the users/tmahlatini/xsmtvdot-ime-mmt4d branch 4 times, most recently from eafb0c9 to 60332b4 Compare July 15, 2026 20:05

TMahlatini force-pushed the users/tmahlatini/xsmtvdot-ime-mmt4d branch from 60332b4 to 4354747 Compare July 15, 2026 20:37

TMahlatini marked this pull request as ready for review July 15, 2026 20:46

TMahlatini requested review from MaheshRavishankar, Max191, benvanik, hanhanW, jtuyls and qedawkins as code owners July 15, 2026 20:46

AGindinson requested review from egebeysel and phemashekar July 16, 2026 06:48

TMahlatini force-pushed the users/tmahlatini/xsmtvdot-ime-mmt4d branch from 4354747 to e33a309 Compare July 16, 2026 14:20

egebeysel mentioned this pull request Jul 17, 2026

Upgrade RISC-V toolchain on CI #22303

Open

egebeysel reviewed Jul 17, 2026

View reviewed changes

TMahlatini and others added 3 commits July 20, 2026 16:51

[LLVMCPU][RISCV] Extended xsmtvdot support with additional tile confi…

bf21633

…gurations that cover vlen=1024 and 4096. Signed-off-by: Terence <terrencemahlatini@gmail.com>

Merge branch 'main' into users/tmahlatini/xsmtvdot-ime-mmt4d

23c8714

[LLVMCPU][RISCV] Restore xsmtvdot materialize_encoding lit tests

e2f44bb

Re-add IME tile cases dropped when resolving the main merge conflict. Signed-off-by: Terence <terrencemahlatini@gmail.com>

Merge branch 'main' into users/tmahlatini/xsmtvdot-ime-mmt4d

d96aea7

Search code, repositories, users, issues, pull requests...

Uh oh!

Conversation

TMahlatini commented Jul 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Uh oh!

github-actions Bot commented Jul 15, 2026

General guidance

Action required: acknowledge IREE project policies

Uh oh!

TMahlatini commented Jul 15, 2026

Uh oh!

egebeysel left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

egebeysel commented Jul 21, 2026

Uh oh!

TMahlatini commented Jul 21, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

TMahlatini commented Jul 15, 2026 •

edited

Loading