[Mlir-commits] [mlir] 7bce6bb - [mlir][linalg][bufferize] Fix crash when bufferizing CallOpInterface

[Matthias Springer]

Author: Matthias Springer
Date: 2021-10-25T21:07:49+09:00
New Revision: 7bce6bb34b70fb94871fc00f164b2b8f151cf309

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

LOG: [mlir][linalg][bufferize] Fix crash when bufferizing CallOpInterface

Do not erase those ops during the traversal. Also improve error handling.

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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Linalg/Transforms/ComprehensiveBufferize.h
    mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Linalg/Transforms/ComprehensiveBufferize.h b/mlir/include/mlir/Dialect/Linalg/Transforms/ComprehensiveBufferize.h
index 9b1c9a18acda..5226ab394cd7 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/ComprehensiveBufferize.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/ComprehensiveBufferize.h
@@ -90,6 +90,11 @@ class BufferizationAliasInfo {
 
   /// Return true if `v1` and `v2` bufferize to equivalent buffers.
   bool areEquivalentBufferizedValues(Value v1, Value v2) const {
+    // Return `false` if we have no information about `v1` or `v2`.
+    if (equivalentInfo.findValue(v1) == equivalentInfo.end() ||
+        equivalentInfo.findValue(v2) == equivalentInfo.end())
+      return false;
+
     return equivalentInfo.getLeaderValue(v1) ==
            equivalentInfo.getLeaderValue(v2);
   }

diff  --git a/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp b/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp
index b4af0fd82903..9e970f3e5e14 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp
@@ -1748,7 +1748,7 @@ bufferize(OpBuilder &b, CallOpInterface callOp, BlockAndValueMapping &bvm,
   Operation *newCallOp = b.create<CallOp>(callOp.getLoc(), funcOp.sym_name(),
                                           resultTypes, newOperands);
   newCallOp->setAttrs(callOp->getAttrs());
-  callOp->erase();
+  // Delete the op at the end of bufferization.
   return success();
 }
 
@@ -2496,6 +2496,8 @@ static LogicalResult bufferizeFuncOpInternals(
   if (failed(bufferize(b, funcOp, bvm, aliasInfo)))
     return failure();
 
+  // Cannot erase ops during the traversal. Do that afterwards.
+  SmallVector<Operation *> toErase;
   // Bufferize the function body. `bufferizedOps` keeps track ops that were
   // already bufferized with pre-order traversal.
   DenseSet<Operation *> bufferizedOps;
@@ -2522,6 +2524,13 @@ static LogicalResult bufferizeFuncOpInternals(
         failed(bufferizeOp(op, bvm, aliasInfo, &bufferizedFunctionTypes,
                            &globalCreator)))
       return failure();
+
+    // Register post-walk erasure, if necessary.
+    if (isa<CallOpInterface>(op))
+      if (llvm::any_of(op->getOperandTypes(), isaTensor) ||
+          llvm::any_of(op->getResultTypes(), isaTensor))
+        toErase.push_back(op);
+
     return success();
   };
   if (funcOp.walk(walkFunc).wasInterrupted())
@@ -2529,6 +2538,9 @@ static LogicalResult bufferizeFuncOpInternals(
 
   LDBG("End BufferizeFuncOpInternals:\n" << funcOp << '\n');
 
+  for (Operation *op : toErase)
+    op->erase();
+
   return success();
 }