[mlir][bufferization] Add BufferizableOpInterface::hasTensorSemantics (#75273)

Add a new interface method to `BufferizableOpInterface`:
`hasTensorSemantics`. This method returns "true" if the op has tensor
semantics and should be bufferized.

Until now, we assumed that an op has tensor semantics if it has tensor
operands and/or tensor op results. However, there are ops like
`ml_program.global` that do not have any results/operands but must still
be bufferized (#75103). The new interface method can return "true" for
such ops.

This change also decouples `bufferization::bufferizeOp` a bit from the
func dialect.

Author: matthias-springer

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

@@ -601,6 +601,12 @@ FailureOr<BaseMemRefType> getBufferType(Value value,
                                         const BufferizationOptions &options,
                                         SmallVector<Value> &invocationStack);
 
+/// Return "true" if the given op has tensor semantics and should be bufferized.
+/// If the op is bufferizable, the BufferizableOpInterface is queried.
+/// Otherwise, an op has tensor semantics if it has tensor operands, tensor
+/// op results and/or tensor block arguments.
+bool hasTensorSemantics(Operation *op);
+
 /// Replace an op with replacement values. The op is deleted. Tensor OpResults
 /// must be replaced with memref values.
 void replaceOpWithBufferizedValues(RewriterBase &rewriter, Operation *op,
@@ -694,6 +700,10 @@ AliasingOpOperandList unknownGetAliasingOpOperands(Value value);
 /// This is the default implementation of getAliasingValues in case the owner
 /// op does not implement the BufferizableOpInterface.
 AliasingValueList unknownGetAliasingValues(OpOperand &opOperand);
+
+/// This is the default implementation of
+/// BufferizableOpInterface::hasTensorSemantics
+bool defaultHasTensorSemantics(Operation *op);
 } // namespace detail
 
 } // namespace bufferization

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td

@@ -575,6 +575,25 @@ def BufferizableOpInterface : OpInterface<"BufferizableOpInterface"> {
           return false;
         }]
       >,
+      InterfaceMethod<
+        /*desc=*/[{
+          Return "true" if the this op has tensor semantics and should be
+          bufferized. By default, ops with tensor operands, tensor op results
+          and/or tensor block arguments have tensor semantics.
+
+          This interface methods can be implemented by ops that should be
+          bufferized but do not have tensor semantics according to the above
+          definition. E.g., this function can return "true" for symbols.
+        }],
+        /*retType=*/"bool",
+        /*methodName=*/"hasTensorSemantics",
+        /*args=*/(ins),
+        /*methodBody=*/"",
+        /*defaultImplementation=*/[{
+          return ::mlir::bufferization::detail
+              ::defaultHasTensorSemantics($_op.getOperation());
+        }]
+      >,
       StaticInterfaceMethod<
         /*desc=*/[{
           Return `true` if the op and this interface implementation supports

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

@@ -689,6 +689,12 @@ bufferization::getBufferType(Value value, const BufferizationOptions &options,
                        *options.defaultMemorySpace);
 }
 
+bool bufferization::hasTensorSemantics(Operation *op) {
+  if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op))
+    return bufferizableOp.hasTensorSemantics();
+  return detail::defaultHasTensorSemantics(op);
+}
+
 void bufferization::replaceOpWithBufferizedValues(RewriterBase &rewriter,
                                                   Operation *op,
                                                   ValueRange values) {
@@ -989,3 +995,20 @@ bufferization::detail::unknownGetAliasingValues(OpOperand &opOperand) {
           r.addAlias({bbArg, BufferRelation::Unknown, /*isDefinite=*/false});
   return r;
 }
+
+bool bufferization::detail::defaultHasTensorSemantics(Operation *op) {
+  auto isaTensor = [](Type t) { return isa<TensorType>(t); };
+  bool hasTensorBlockArgument = any_of(op->getRegions(), [&](Region &r) {
+    return any_of(r.getBlocks(), [&](Block &b) {
+      return any_of(b.getArguments(), [&](BlockArgument bbArg) {
+        return isaTensor(bbArg.getType());
+      });
+    });
+  });
+  if (hasTensorBlockArgument)
+    return true;
+
+  if (any_of(op->getResultTypes(), isaTensor))
+    return true;
+  return any_of(op->getOperandTypes(), isaTensor);
+}

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

@@ -350,31 +350,6 @@ mlir::bufferization::createFinalizingBufferizePass() {
 // BufferizableOpInterface-based Bufferization
 //===----------------------------------------------------------------------===//
 
-static bool isaTensor(Type t) { return isa<TensorType>(t); }
-
-/// Return true if the given op has a tensor result or a tensor operand.
-static bool hasTensorSemantics(Operation *op) {
-  bool hasTensorBlockArgument = any_of(op->getRegions(), [](Region &r) {
-    return any_of(r.getBlocks(), [](Block &b) {
-      return any_of(b.getArguments(), [](BlockArgument bbArg) {
-        return isaTensor(bbArg.getType());
-      });
-    });
-  });
-  if (hasTensorBlockArgument)
-    return true;
-
-  if (auto funcOp = dyn_cast<FunctionOpInterface>(op)) {
-    bool hasTensorArg = any_of(funcOp.getArgumentTypes(), isaTensor);
-    bool hasTensorResult = any_of(funcOp.getResultTypes(), isaTensor);
-    return hasTensorArg || hasTensorResult;
-  }
-
-  bool hasTensorResult = any_of(op->getResultTypes(), isaTensor);
-  bool hasTensorOperand = any_of(op->getOperandTypes(), isaTensor);
-  return hasTensorResult || hasTensorOperand;
-}
-
 namespace {
 /// A rewriter that keeps track of extra information during bufferization.
 class BufferizationRewriter : public IRRewriter, public RewriterBase::Listener {

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

@@ -325,6 +325,28 @@ struct FuncOpInterface
 
   static bool supportsUnstructuredControlFlow() { return true; }
 
+  bool hasTensorSemantics(Operation *op) const {
+    auto isaTensor = [](Type type) { return isa<TensorType>(type); };
+
+    // A function has tensor semantics if it has tensor arguments/results.
+    auto funcOp = cast<FuncOp>(op);
+    bool hasTensorArg = any_of(funcOp.getArgumentTypes(), isaTensor);
+    bool hasTensorResult = any_of(funcOp.getResultTypes(), isaTensor);
+    if (hasTensorArg || hasTensorResult)
+      return true;
+
+    // It also has tensor semantics if it has tensor block arguments.
+    // TODO: Decouple bufferization of unstructured control flow from
+    // BufferizableOpInterface implementations. We should only care about
+    // region entry block arguments here (which are already covered by the
+    // argument types of the function).
+    for (Block &block : funcOp.getBody())
+      if (any_of(block.getArgumentTypes(), isaTensor))
+        return true;
+
+    return false;
+  }
+
   AliasingOpOperandList
   getAliasingOpOperands(Operation *op, Value value,
                         const AnalysisState &state) const {

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

@@ -1030,13 +1030,6 @@ OneShotAnalysisState::analyzeSingleOp(Operation *op,
   return success();
 }
 
-/// Return true if the given op has a tensor result or a tensor operand.
-static bool hasTensorSemantics(Operation *op) {
-  bool hasTensorResult = any_of(op->getResultTypes(), isaTensor);
-  bool hasTensorOperand = any_of(op->getOperandTypes(), isaTensor);
-  return hasTensorResult || hasTensorOperand;
-}
-
 /// Analyze equivalence of tied OpResult/OpOperand pairs of the given ops.
 static void equivalenceAnalysis(SmallVector<Operation *> &ops,
                                 OneShotAnalysisState &state) {

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

@@ -458,6 +458,15 @@ LogicalResult mlir::bufferization::bufferizeModuleOp(
       foldMemRefCasts(funcOp);
   }
 
+  // Bufferize all other ops.
+  for (Operation &op : moduleOp.getOps()) {
+    // Functions were already bufferized.
+    if (isa<func::FuncOp>(&op))
+      continue;
+    if (failed(bufferizeOp(&op, options, statistics)))
+      return failure();
+  }
+
   // Post-pass cleanup of function argument attributes.
   removeBufferizationAttributesInModule(moduleOp);