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

[springerm]

Summary

Do not erase those ops during the traversal. Also improve error handling. Fixes a crash introduced by D111924.

mlir/include/mlir/Dialect/Linalg/Transforms/ComprehensiveBufferize.h

public:
   /// operands. Return true if it is possible to find an inplace write W that
   /// creates a conflict.
   bool
   wouldCreateReadAfterWriteInterference(OpOperand &operand, OpResult result,
                                         const DominanceInfo &domInfo) const;
 
   /// 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);
   }
 
   /// Return true if the source of an `insertSliceOp` bufferizes to an
   /// equivalent ExtractSliceOp.
   bool isSourceEquivalentToAMatchingInplaceExtractSliceOp(
       tensor::InsertSliceOp insertSliceOp) const;

mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp

for (OpOperand &opOperand : callOp->getOpOperands()) {
     }
     newOperands.push_back(buffer);
   }
 
   // 4. Create the new CallOp.
   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();
 }
 
 /// tensor::CastOp bufferizes to memref::CastOp.
 static LogicalResult bufferize(OpBuilder &b, tensor::CastOp castOp,
                                BlockAndValueMapping &bvm,
                                BufferizationAliasInfo &aliasInfo) {
   // Take a guard before anything else.
static LogicalResult bufferizeFuncOpInternals(
   LLVM_DEBUG(llvm::dbgs() << "\n\n");
   LDBG("Begin BufferizeFuncOpInternals:\n" << funcOp << '\n');
   OpBuilder b(funcOp->getContext());
 
   // Start by bufferizing `funcOp` arguments.
   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;
   auto walkFunc = [&](Operation *op) -> WalkResult {
     // Collect ops that need to be bufferized before `op`.
     SmallVector<Operation *> preorderBufferize;
     Operation *parentOp = op->getParentOp();
     // scf::ForOp and TiledLoopOp must be bufferized before their blocks
for (Operation *op : llvm::reverse(preorderBufferize))
       if (failed(bufferizeOp(op, bvm, aliasInfo, &bufferizedFunctionTypes,
                              &globalCreator)))
         return failure();
 
     if (!bufferizedOps.contains(op) &&
         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())
     return failure();
 
   LDBG("End BufferizeFuncOpInternals:\n" << funcOp << '\n');
 
+  for (Operation *op : toErase)
+    op->erase();
+
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization entry-point for modules.
 //===----------------------------------------------------------------------===//
 
 /// Return the op with Allocate MemoryEffect if `v` is equivalent to such an