[Mlir-commits] [mlir] c5e67b8 - [mlir] [tensor] Crash in getPackOpResultTypeShape for tensor.pack/unpack ops. (#90641)

[llvmlistbot at llvm.org]

Author: Sayan Saha
Date: 2024-05-13T18:18:20-04:00
New Revision: c5e67b86ef6718585120e3cabb04c1fc2d292cfb

URL: https://github.com/llvm/llvm-project/commit/c5e67b86ef6718585120e3cabb04c1fc2d292cfb
DIFF: https://github.com/llvm/llvm-project/commit/c5e67b86ef6718585120e3cabb04c1fc2d292cfb.diff

LOG: [mlir] [tensor] Crash in getPackOpResultTypeShape for tensor.pack/unpack ops. (#90641)

Windows build of `mlir` with Visual Studio (19.36.32538 for x64) using
with the following command:

`cmake.exe -GNinja -DCMAKE_BUILD_TYPE=Release
-DLLVM_ENABLE_PROJECTS=mlir -DLLVM_ENABLE_EH=ON -DLLVM_ENABLE_RTTI=1
-DLLVM_TARGETS_TO_BUILD=host ../llvm`

is leading to a crash when calling canonicalization on
`tensor.pack`/`tensor.unpack` ops `mlir-opt --canonicalize input.mlir`
where the `input.mlir` is as follows (this is taken from one of the
filecheck tests for `tensor.pack`):

```
func.func @pack_unpack(%arg0: tensor<128x256xf32>) -> tensor<128x256xf32> {
          %pack_dest = tensor.empty() : tensor<8x16x8x32xf32>
          %unpack_dest = tensor.empty() : tensor<128x256xf32>
          %tp = tensor.pack %arg0 outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [8, 32] into %pack_dest : tensor<128x256xf32> -> tensor<8x16x8x32xf32>
          %tup = tensor.unpack %tp outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [8, 32] into %unpack_dest : tensor<8x16x8x32xf32> -> tensor<128x256xf32>
          return %tup : tensor<128x256xf32>
        }
```

The crash is seemingly coming from invalid memory access during
iterating over `innerDimsPos` within `getPackOpResultTypeShape`.

This crash is also causing the following tests to fail:


```
MLIR :: Dialect/Linalg/canonicalize.mlir                                                                                                                    
MLIR :: Dialect/Linalg/data-layout-propagation.mlir                                                                                                         
MLIR :: Dialect/Linalg/generalize-tensor-pack-tile.mlir                                                                                                     
MLIR :: Dialect/Linalg/generalize-tensor-pack.mlir                                                                                                          
MLIR :: Dialect/Linalg/generalize-tensor-unpack-tile.mlir                                                                                                   
MLIR :: Dialect/Linalg/generalize-tensor-unpack.mlir                                                                                                        
MLIR :: Dialect/Linalg/transform-lower-pack.mlir                                                                                                            
MLIR :: Dialect/Linalg/transform-op-fuse.mlir                                                                                                               
MLIR :: Dialect/Linalg/transform-op-pack.mlir                                                                                                               
MLIR :: Dialect/Linalg/transform-pack-greedily.mlir                                                                                                         
MLIR :: Dialect/Tensor/canonicalize.mlir                                                                                                                   
MLIR :: Dialect/Tensor/fold-into-pack-and-unpack.mlir                                                                                                       
MLIR :: Dialect/Tensor/invalid.mlir                                                                                                                         
MLIR :: Dialect/Tensor/ops.mlir                                                                                                                             
MLIR :: Dialect/Tensor/simplify-pack-unpack.mlir                                                                                                           
MLIR :: Dialect/Tensor/tiling.mlir
```

Added: 
    

Modified: 
    mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index e41d59a0e0b94..414bd7459af8f 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -3963,7 +3963,7 @@ static SmallVector<int64_t> getPackOpResultTypeShape(
     ArrayRef<int64_t> sourceShape, ArrayRef<int64_t> innerTileSizes,
     ArrayRef<int64_t> innerDimsPos, ArrayRef<int64_t> outerDimsPerm) {
   SmallVector<int64_t> resultShape = llvm::to_vector(sourceShape);
-  for (auto tiledDim : llvm::enumerate(innerDimsPos)) {
+  for (auto tiledDim : llvm::enumerate(llvm::to_vector(innerDimsPos))) {
     if (ShapedType::isDynamic(resultShape[tiledDim.value()]))
       continue;
     if (ShapedType::isDynamic(innerTileSizes[tiledDim.index()])) {
@@ -3992,7 +3992,7 @@ SmallVector<OpFoldResult> PackOp::getResultShape(
   AffineExpr s0, s1;
   bindSymbols(builder.getContext(), s0, s1);
   AffineExpr ceilDivExpr = s0.ceilDiv(s1);
-  for (auto tiledDim : llvm::enumerate(innerDimsPos)) {
+  for (auto tiledDim : llvm::enumerate(llvm::to_vector(innerDimsPos))) {
     resultDims[tiledDim.value()] = affine::makeComposedFoldedAffineApply(
         builder, loc, ceilDivExpr,
         {resultDims[tiledDim.value()], innerTileSizes[tiledDim.index()]});