[mlir][bufferize] Disallow adding new bufferizable ops during bufferization

[springerm]

Summary

Ops that are created during the bufferization were not analyzed (when run with One-Shot Bufferize), and users should instead create memref ops directly.

Futhermore, this fixes an issue where an op was erased (and put on the erasedOps list), but subsequently a new tensor op was created at the same memory location. This op was then not bufferized. Disallowing the creation of new tensor ops simplifies the bufferization and fixes such issues.

[springerm]

update

[springerm]

update

[pifon2a]

Thank you a lot!

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

case OpFilterEntry::DENY:
         return false;
     };
   }
   return isAllowed;
 }
 
 BufferizableOpInterface
 BufferizationOptions::dynCastBufferizableOp(Operation *op) const {
-  if (isOpAllowed(op))
-    return dyn_cast<BufferizableOpInterface>(op);
+  auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op);
+  if (!bufferizableOp)
+    return nullptr;
+  if (!isOpAllowed(op))
     return nullptr;
+  return bufferizableOp;
 }
 
 BufferizableOpInterface
 BufferizationOptions::dynCastBufferizableOp(Value value) const {
   if (auto bufferizableOp = value.getDefiningOp<BufferizableOpInterface>())
     if (isOpAllowed(bufferizableOp.getOperation()))
       return bufferizableOp;
   return nullptr;

mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp

 }
 
 namespace {
 /// A rewriter that keeps track of extra information during bufferization.
 class BufferizationRewriter : public IRRewriter {
 public:
   BufferizationRewriter(MLIRContext *ctx, DenseSet<Operation *> &erasedOps,
                         DenseSet<Operation *> &toMemrefOps,
-                        SmallVector<Operation *> &worklist)
+                        const BufferizationOptions &options)
       : IRRewriter(ctx), erasedOps(erasedOps), toMemrefOps(toMemrefOps),
-        worklist(worklist) {}
+        options(options) {}
 
 protected:
   void notifyOperationRemoved(Operation *op) override {
     IRRewriter::notifyOperationRemoved(op);
     erasedOps.insert(op);
+    // Erase if present.
+    toMemrefOps.erase(op);
   }
 
   void notifyOperationInserted(Operation *op) override {
     IRRewriter::notifyOperationInserted(op);
 
     // Keep track of to_memref ops.
     if (isa<ToMemrefOp>(op)) {
       toMemrefOps.insert(op);
       return;
     }
 
     // Skip to_tensor ops.
     if (isa<ToTensorOp>(op))
       return;
 
-    // A new bufferizable op was inserted. Add it to the worklist.
-    if (hasTensorSemantics(op))
-      worklist.push_back(op);
+    // Adding new bufferizable ops is not allowed during bufferization. Such ops
+    // would not be analyzed and can lead to surprising behavior.
+    assert((!hasTensorSemantics(op) || !options.isOpAllowed(op)) &&
+           "creating new tensor ops is not allowed during bufferization");
   }
 
 private:
   /// A set of all erased ops.
   DenseSet<Operation *> &erasedOps;
 
   /// A set of all to_memref ops.
   DenseSet<Operation *> &toMemrefOps;
 
-  /// The list of bufferizable ops.
-  SmallVector<Operation *> &worklist;
+  /// The bufferization options.
+  const BufferizationOptions &options;
 };
 } // namespace
 
 LogicalResult
 bufferization::bufferizeOp(Operation *op,
                            BufferizationState &bufferizationState) {
   const auto &options = bufferizationState.getOptions();
 
if (hasTensorSemantics(op))
       worklist.push_back(op);
   });
 
   // Keep track of all erased ops.
   DenseSet<Operation *> erasedOps;
 
   // Bufferize all ops.
   BufferizationRewriter rewriter(op->getContext(), erasedOps, toMemrefOps,
-                                 worklist);
+                                 bufferizationState.getOptions());
   for (unsigned i = 0; i < worklist.size(); ++i) {
     Operation *op = worklist[i];
     // Skip ops that were erased.
     if (erasedOps.contains(op))
       continue;
-    // Skip ops that are not bufferizable.
-    auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op);
+    // Skip ops that are not bufferizable or not allowed.
+    auto bufferizableOp = options.dynCastBufferizableOp(op);
     if (!bufferizableOp)
       continue;
-    // Continue ops that are not allowed.
-    if (!options.isOpAllowed(op))
+    // Skip ops that no longer have tensor semantics.
+    if (!hasTensorSemantics(op))
       continue;
     // Bufferize the op.
     rewriter.setInsertionPoint(op);
     (void)bufferizableOp.bufferize(rewriter, bufferizationState);
   }
 
   // Fold all to_memref(to_tensor(x)) pairs.
   for (Operation *op : toMemrefOps) {
-    if (erasedOps.contains(op))
-      continue;
     rewriter.setInsertionPoint(op);
     (void)bufferization::foldToMemrefToTensorPair(rewriter,
                                                   cast<ToMemrefOp>(op));
   }
 
   /// Check the result of bufferization. Return an error if an op was not
   /// bufferized, unless partial bufferization is allowed.
   if (bufferizationState.getOptions().allowUnknownOps)