[Mlir-commits] [mlir] 82dd977 - [mlir][Linalg] Tighten canonicalization of InsertSliceOp that triggers infinite loop

[Nicolas Vasilache]

Author: Nicolas Vasilache
Date: 2021-10-14T15:26:03Z
New Revision: 82dd977bafa8ff6e0e0efa16cafcfd7d1d9249ba

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

LOG: [mlir][Linalg] Tighten canonicalization of InsertSliceOp that triggers infinite loop

I am unclear this is reproducible with correct IR but atm the verifier for InsertSliceOp
is not powerful enough and this triggers an infinite loop that is worth fixing independently.

Differential Revision: https://reviews.llvm.org/D111812

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 51c9c5b01e355..1a278a54910a9 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -1279,12 +1279,19 @@ struct InsertSliceOpSourceCastInserter final
               getConstantIntValue(insertSliceOp.getMixedSizes()[i]))
         newSrcShape[i] = *constInt;
     }
+
     RankedTensorType newSrcType =
         RankedTensorType::get(newSrcShape, srcType.getElementType());
-    if (srcType == newSrcType)
+    if (srcType == newSrcType ||
+        !preservesStaticInformation(srcType, newSrcType) ||
+        !tensor::CastOp::areCastCompatible(srcType, newSrcType))
       return failure();
 
-    // srcType and newSrcType are 
diff erent. Insert a cast.
+    // newSrcType is:
+    //   1) Different from srcType.
+    //   2) "More static" than srcType.
+    //   3) Cast-compatible with srcType.
+    // Insert the cast.
     Value cast = rewriter.create<tensor::CastOp>(
         insertSliceOp.getLoc(), newSrcType, insertSliceOp.source());
     rewriter.replaceOpWithNewOp<InsertSliceOp>(

diff  --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir
index 2887b6ea87fb1..9d9da02c0220f 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -559,3 +559,13 @@ func @fold_overlapping_insert(%input : tensor<?x?x?xf32>, %slice1: tensor<4x?x8x
   // CHECK: return %[[INSERT]]
   return %1 : tensor<?x?x?xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func @folding_incorrect_ir_triggers_infinite_loop
+func @folding_incorrect_ir_triggers_infinite_loop(
+  %A : tensor<4x4xf32>, %C : tensor<?x?xf32>) -> tensor<?x?xf32> {
+  %rC = tensor.insert_slice %A into %C[0, 0][12345, 67890][1, 1] :
+    tensor<4x4xf32> into tensor<?x?xf32>
+  return %rC: tensor<?x?xf32>
+}