[all-commits] [llvm/llvm-project] 97043e: [mlir][Vector][GPU] Distribute expanding `shape_ca...

[Artem Gindinson via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: 97043e50ad41c9da94a5cb48417f5139d6a84c8d
      https://github.com/llvm/llvm-project/commit/97043e50ad41c9da94a5cb48417f5139d6a84c8d
  Author: Artem Gindinson <gindinson at roofline.ai>
  Date:   2026-03-03 (Tue, 03 Mar 2026)

  Changed paths:
    M mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
    M mlir/test/Dialect/Vector/vector-warp-distribute.mlir

  Log Message:
  -----------
  [mlir][Vector][GPU] Distribute expanding `shape_cast` ops (#183830)

The initial implementation of `shape_cast` distribution only focused on
scenarios with collapsing shape casts. Within downstream pipelines such
as IREE, commit 962a9a3 exposes an issue with this implementation, where
the rank-expanding cast ops (stemming from the new `vector.broadcast`
canonicalization) silently fall through to the "collapsing-or-no-op"
logic. This brings about bugs with rank mismatches and firing validation
assertions when distributing rather common reshaping sequences
encountered after CSE/ canonicalization, such as below:
```
  // Example 1: gather op
  %weight = arith.constant dense_resource<__elided__> : tensor<256xi8>
  %c0 = arith.constant 0 : index
  ...
  %expand = vector.shape_cast <...> : vector<1xindex> to vector<1x1xindex>
  %gather = vector.gather %weight[%c0] [%expand], <...>, <...> : memref<256xi8>, vector<1x1xindex>, vector<1x1xi1>, vector<1x1xi8> into vector<1x1xi8>
  %collapse_back = vector.shape_cast %gather : vector<1x1xi8> to vector<1xi8>
  // Example 2: multi-reduction
  %expand = vector.shape_cast <...>: vector<1x96xi32> to vector<1x2x48xi32>
  %reduce = vector.multi_reduction <add>, %expand, <...> [1, 2]: vector<1x2x48xi32> to vector<1x1xi32>
  %collapse = vector.shape_cast %reduce : vector<1x1xi32> to vector<1xi32>
```

This commit adds initial handling of expanding `shape_cast`s, going
through the three scenarios:
- if all "excess" dimensions in the front of the destination shape are
unit, it's clear that the work is not distributed across those, so we
strip the same number of unit dimensions from the per-lane yielded type;
- if the source type within the warp code is of rank 1, we still
determine the corresponding output type cleanly by multiplying the
dimensions of the per-lane yield type;
- if neither of the above are true, explicitly fail the pattern for such
expanding `shape_cast`'s. Dimension-specific distribution parameters are
deemed ambiguous, at least from within this pattern's context.

---------

Signed-off-by: Artem Gindinson <gindinson at roofline.ai>



To unsubscribe from these emails, change your notification settings at https://github.com/llvm/llvm-project/settings/notifications