Skip to content

VMI lacks reduce_mini / reduce_maxi for i32 — only f32 variants exist (reduce_minf / reduce_maxf) #502

Description

@peanutchan

Problem

VMI currently provides reduce_minf and reduce_maxf for f32 only. There are no i32 variants (reduce_mini / reduce_maxi).

This means any kernel that needs to reduce an i32 vector (e.g., finding the minimum/maximum index across a vector of expert indices) must bitcast i32 → f32, call reduce_minf/reduce_maxf, then bitcast the result back to i32.

What exists today

Op Element type Source
reduce_minf f32/f16/bf16 only VMIOps.td:371-390
reduce_maxf f32/f16/bf16 only VMIOps.td:392-410
group_reduce_maxi i32 (grouped only) VMIOps.td:443
reduce_mini does not exist
reduce_maxi does not exist

The grouped variant group_reduce_maxi exists, but there is no non-grouped reduce_mini / reduce_maxi for i32.

This is also documented in VMILowerUnifiedToLegacy.cpp:62-66:

//   vcmax → reduce_maxf / group_reduce_maxf/group_reduce_maxi
//           (full-int → skipped, no legacy reduce_maxi)
//   vcmin → reduce_minf only (full-float)
//           (group / full-int → skipped, no legacy group_reduce_min*
//            nor reduce_mini)

What CCE VF does natively

CCE VF uses i32 vmin + vcmin directly — no type conversion needed:

// Elementwise min across index vectors (i32)
vmin(acc_min_idx, acc_min_idx, tmp_idx_vec[i], preg_all, MODE_ZEROING);
// Cross-lane reduce (i32)
vcmin(idx, acc_min_idx, preg_all, MODE_ZEROING);

What VMI must do (workaround)

// Step 1: bitcast entire i32 index vector to f32
%idx_f32 = pto.vmi.bitcast %indices
    : !pto.vmi.vreg<64xi32> -> !pto.vmi.vreg<64xf32>

// Step 2: reduce_minf on the f32 view (lowers to vcmin + vmin)
%min_idx_f32 = pto.vmi.reduce_minf %idx_f32, %init, %mask
    : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<1xf32>, !pto.vmi.mask<64xpred>
    -> !pto.vmi.vreg<1xf32>

// Step 3: bitcast result back to i32
%min_idx_i32 = pto.vmi.bitcast %min_idx_f32
    : !pto.vmi.vreg<1xf32> -> !pto.vmi.vreg<1xi32>

Performance impact

vbitcast has zero ISA cost — it is a pure register alias and does not appear in any PMU/PERF/simd dispatch log. Confirmed by simulation: the SIMD dispatch dump shows RV_VCMIN and RV_VMIN execute correctly, but no cast instruction is ever dispatched.

So the workaround does not affect runtime performance. However, it adds code verbosity and makes VMI kernels harder to read compared to CCE VF.

Why not just use 1xf32 to represent the result?

After reduce_minf, the result IS already !pto.vmi.vreg<1xf32>. The problem is the input: the index vector is 64xi32, and reduce_minf requires f32 input. So the entire index vector must be bitcast to f32 before the reduce, then the 1xf32 result bitcast back to 1xi32 after.

If reduce_mini existed for i32, we could skip both bitcasts and work directly with i32 throughout — matching CCE VF 1:1 at the source level.

Safety of the workaround

For non-negative indices (0..N), the i32 bit patterns are positive f32 subnormals, which preserve the same ordering as unsigned i32. So reduce_minf on the bitcast values gives the correct minimum index. This is safe for all practical use cases (expert indices, token indices, etc.).

Requested feature

Add reduce_mini and reduce_maxi for i32:

%result = pto.vmi.reduce_mini %source, %init, %mask
    : !pto.vmi.vreg<64xi32>, !pto.vmi.vreg<1xi32>, !pto.vmi.mask<64xpred>
    -> !pto.vmi.vreg<1xi32>

These would lower to the same vcmin/vmax/vcmax/vmin ISA instructions — no new hardware support needed, just a new VMI op that accepts i32 input directly.

Microbenchmark

A standalone test case is attached at:
pto-skills/testing-pto-kernels/generated/moe_a5/topk_gate/test_reduce_mini_gap.pto

It demonstrates the 3-step workaround (bitcast → reduce_minf → bitcast) vs the desired single-step reduce_mini.

Related

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions