[Mlir-commits] [mlir] [mlir][linalg] Fix UBSan division-by-zero in PackOp folding (PR #186271)

[Andrzej Warzyński]

banach-space wrote:

> It is not OK for any transformation to generate invalid IR from valid IR: that's always a bug.

I know what you mean, but this particular case raises a wider question for  me. The following Op is already invalid, but we don't let the verify to catch that (it could and does ATM, but by design MLIR verifiers should only check "locally"):
```mlir
  %tile_size = arith.constant 0 : index
  %empty = tensor.empty(%tile_size) : tensor<1x16x?x1xi32>
  %pack = linalg.pack %A
    padding_value(%pad_val : i32)
    inner_dims_pos = [0, 1]
    inner_tiles = [%tile_size, 1]
    into %empty : tensor<7x16xi32> -> tensor<1x16x?x1xi32>
 ```
 
 The folder turns that into something that the verifier can and indeed does reject:
 ```mlir
     %2 = "tensor.empty"() : () -> tensor<1x16x0x1xi32>
     %3 = "tensor.cast"(%2) : (tensor<1x16x0x1xi32>) -> tensor<1x16x?x1xi32>
     %4 = "linalg.pack"(%arg0, %2, %0) <{inner_dims_pos = array<i64: 0, 1>, operandSegmentSizes = array<i32: 1, 1, 1, 0>, static_inner_tiles = array<i64: 0, 1>}> : (tensor<7x16xi32>, tensor<1x16x0x1xi32>, i32) -> tensor<1x16x0x1xi32>
 ```
So, in the context of the Op verifier, `FoldTensorCastPackOp` turns something "valid" into something "invalid", agreed. However, I feel that we are missing something at the framework level that would prevent/capture this sort of issues (e.g. **runtime  Op verification**).

The approach that you are taking here is consistent with how MLIR has been approaching this and my comment is tangential.

> The crash happens in the builder of the op: we never get to this point.

Hm, the bulider seems to run fine. It's `PackOp::canonicalize() -> paddingIsNotNeeded() -> requirePaddingValue()` that leads to UB, right?  In fact, how about an assert in `requirePaddingValue`? That method effectively assumes that tile sizes are non-zero.

--- 
All in all LGTM % comments
* Remove references to `expand_shape` - the relation is not clear to me.
* [optional] Comments + assert re assumption on tile sizes in `paddingIsNotNeeded()`.

Btw, IMO we are missing a wider discussion on 0-dim in Linalg. That's outside the scope here and I am only mentioning so that we refrain from special-casing for dim-0 if there are more issues like this.

Thanks for digging into this and proposing a fix!

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