[Mlir-commits] [mlir] [MLIR][Vector] Generalize broadcast lowering for single-element vector to nD (PR #206501)

[Andrzej Warzyński]

banach-space wrote:

Thanks!

> This PR generalizes the broadcast optimization from `Before: 0D/1D to 1D` to `After: single-element vector to nD vector`. Example:

Please, instead, include IR. That tends to be much easier to parse. Also, you seem to be focusing on a relatively short example.

> 
> ```mlir
> vector.broadcast %v : vector<1xf8E8M0FNU> to vector<16x1xf8E8M0FNU>
> ```
> 
> becomes
> 
> ```mlir
>     %790 = llvm.shufflevector %789, %28 [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] : vector<16xi8>
> ```
> 
> instead of falling into the branch of:
> 
> ```c++
>     if (srcRank < dstRank) {
>       // Duplication.
>       VectorType resType = VectorType::Builder(dstType).dropDim(0);
>       Value bcst =
>           vector::BroadcastOp::create(rewriter, loc, resType, op.getSource());
>       Value result = ub::PoisonOp::create(rewriter, loc, dstType);
>       for (int64_t d = 0, dim = dstType.getDimSize(0); d < dim; ++d)
>         result = vector::InsertOp::create(rewriter, loc, bcst, result, d);
>       rewriter.replaceOp(op, result);
>       return success();
>     }
> ```

Please consistently use IR for _BEFORE_ and _AFTER_ (as opposed to mixing IR and source C++). 

> that would generate 16 inserts for one broadcast op, leading to excessive IR that may survive up to LLVM dialect, further expanding to 32 ops (extract + insert).

In such a case, could you add a test for the conversion so that we can see the impact? 

> 
> Application: MX block scale for `arith.scaling_extf` elementwise requirement.

[nit] Those MX types are tricky to parse (the cognitive load is high). Could you instead use something shorter in the summary? Ultimately, the element type is irrelevant for the transformation that you are updating.

All in all LG.

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