Stream and Flow canonicalization handle workgroup_count_from_slice workload deduplication inconsistently

[boydjohnson]

Intro

I am working on a project to write a compiler and runtime for RK3588 NPUs using IREE (https://github.com/boydjohnson/iree-rocket). My general approach for the compiler is using MLIR Transform Dialect with IREE.

Issue

I noticed that Stream --iree-stream-fold-uniform-operands leaves workgroup_count_from_slice(%H, %H, %C) when H and W can be shown to be the same value. I would have expected workgroup_count_from_slice(%H, %C) because the ordinals get remapped in the the same pass.

Before Pass

dim	ordinal
H	0
W	1
C	2

After Pass

dim	ordinal
H	0
W	0
C	1

when workgroup_count_from_slice(%H, %C) you get 0 -> %H, 0-> %H, 1 -> %C
when workgroup_count_from_slice(%H, %H, %C) you get 0 -> %H 0-> %H, 1 -> %H

This causes

• H=128, C=64: C is interpreted as 128 → over-dispatch [8,1,1]
• H=32, C=128: C is interpreted as 32 → under-dispatch [1,1,1]

Flow --iree-flow-canonicalize gives workgroup_count_from_slice(%H, %C)

Stream --iree-stream-fold-uniform-operands give workgroup_count_from_slice(%H, %H, %C)

Flow

Stream unit test

workgroup_count_from_slice lowering