[Mlir-commits] [mlir] [mlir][vector] fix contraction lowering crash on scalar results after folding a broadcast-like shape_cast (PR #229824)

[llvmlistbot at llvm.org]

Cerulime wrote:

> Wasn't the problem discovered and solution proposed by @Cerulime ? We should make sure that their contribution is attributed accordingly.
> 
> It's not entirely clear to me what is happening here - the reproduce is too complex. Below's a reduced version (I removed one TD Op). It generates invalid IR and that needs to be investigated _first_. Otherwise, we are tweaking other patterns so that they stay correct even when _invalid_ input is fed. We don't do that in MLIR.
> 
> ```mlir
> func.func @contract_scalar_result_unit_dims_lowering(%x: vector<2xbf16>, %scale: bf16, %acc: bf16) -> bf16 {
>   %lhs = vector.shape_cast %x : vector<2xbf16> to vector<1x1x2xbf16>
>   %rhs = vector.broadcast %scale : bf16 to vector<1x1x2xbf16>
>   %out = vector.contract {
>     indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
>                      affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
>                      affine_map<(d0, d1, d2) -> ()>],
>     iterator_types = ["reduction", "reduction", "reduction"],
>     kind = #vector.kind<add>
>   } %lhs, %rhs, %acc : vector<1x1x2xbf16>, vector<1x1x2xbf16> into bf16
>   return %out : bf16
> }
> 
> module attributes {transform.with_named_sequence} {
>   transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {
>     %func = transform.structured.match ops{["func.func"]} in %module_op : (!transform.any_op) -> !transform.any_op
>     transform.apply_patterns to %func {
>       transform.apply_patterns.vector.reduction_to_contract
>     } : !transform.any_op
>     transform.yield
>   }
> }
> ```
> 
> Thanks!

Hi @banach-space ,

Thank you first for looking into this and for looking out for the attribution!

Yes, you point out another problem in this file. I also submit a issue #230084 to describe it. I am not sure should we merge these two issues.

The reduced test case makes the root cause crystal clear:

In CombineContractBroadcastMask (combineContractAndBroadcast), the LHS vector.shape_cast is folded, dropping the leading unit reduction dimensions (d0, d1) from the LHS map ((d0, d1, d2) -> (d2)).
However, the RHS is a scalar broadcast (vector.broadcast %scale : bf16 to vector<1x1x2xbf16>), which combineContractAndBroadcast skips folding because the source is scalar rather than a vector. The RHS map remains (d0, d1, d2) -> (d0, d1, d2).
This creates an asymmetric vector.contract where d0 and d1 are still marked as reduction iterators but only appear on the RHS. In ContractionOp::verify() / verifyOutputShape(), these orphaned reduction dimensions fail to form contracting pairs and are treated as free dimensions, leading to a verification failure because the result is expected to be a vector rather than a scalar.

The crash in ContractionOpLowering::lowerParallel was purely a downstream symptom of this invalid intermediate IR. Other details is in that issue.

Thanks for all contributors. I am glad to help.

https://github.com/llvm/llvm-project/pull/229824