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

[Andrzej Warzyński]

banach-space wrote:

> Right now the problem is that the folder produced invalid IR, leading to the crash.
> 
> ```
> %2 = "linalg.pack"(%arg0, %1, %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>
> ```
> 
> This op is rejected by the verifier, because of the static dimension being 0. Basically this is the check I'm modifying to only check the static case and not the dynamic one. 

+1

> We were folding the dynamic into static, which we need to stop doing to align with the verifier.

This is a bit unclear. The proposed fix will only prevent folding when the size is effectively `0` (`isZeroInteger`), right? As opposed to folding any dynamic size into a static one? That feels arbitrary. 

The current fold seems correct to me. It's replacing IR that's invalid (dynamic shapes mean that this can only be verified at runtime) with other invalid IR (with static shapes, so verification can kick in at compile time). Why not let `PackOp::verify` do its job post-folding? Wouldn't the verifier fail before hitting UB and wouldn't that be sufficient to fix the original issue.

Btw,
> Add a regression test that ensures a pack with a zero tile size is not folded into tensor.expand_shape.

I don't quite see how `tensor.expand_shape` is involved here. I run your example and didn't see any relevant patterns being run.

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