[mlir][bufferization] Add `BufferizableOpInterface::hasTensorSemantics` #75273

matthias-springer · 2023-12-13T02:18:54Z

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.

Depends on #77574.

llvmbot · 2023-12-13T02:19:23Z

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-bufferization

Author: Matthias Springer (matthias-springer)

Changes

Add a new interface method to BufferizableOpInterface: hasTensorSemanticsForBufferization. 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.

Full diff: https://github.com/llvm/llvm-project/pull/75273.diff

6 Files Affected:

(modified) mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h (+10)
(modified) mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td (+22)
(modified) mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp (+25)
(modified) mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp (+4-29)
(modified) mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp (+22)
(modified) mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp (+9)

diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
index 7c09a43f96397b..d3dc5683772e27 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
+++ b/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 hasTensorSemanticsForBufferization(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::hasTensorSemanticsForBufferization.
+bool defaultHasTensorSemanticsForBufferization(Operation *op);
 } // namespace detail
 
 } // namespace bufferization
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td
index fd1ceb68af5dd9..e7445cb4e63da2 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td
@@ -575,6 +575,28 @@ 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.
+
+          TODO: This interface method should be called `hasTensorSemantics`, but
+          that name is already in use in `DestinationStyleOpInterface`.
+        }],
+        /*retType=*/"bool",
+        /*methodName=*/"hasTensorSemanticsForBufferization",
+        /*args=*/(ins),
+        /*methodBody=*/"",
+        /*defaultImplementation=*/[{
+          return ::mlir::bufferization::detail
+              ::defaultHasTensorSemanticsForBufferization($_op.getOperation());
+        }]
+      >,
       StaticInterfaceMethod<
         /*desc=*/[{
           Return `true` if the op and this interface implementation supports
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index 1e8dc4387ed4f0..8f577947f8ddbf 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
@@ -689,6 +689,12 @@ bufferization::getBufferType(Value value, const BufferizationOptions &options,
                        *options.defaultMemorySpace);
 }
 
+bool bufferization::hasTensorSemanticsForBufferization(Operation *op) {
+  if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op))
+    return bufferizableOp.hasTensorSemanticsForBufferization();
+  return detail::defaultHasTensorSemanticsForBufferization(op);
+}
+
 void bufferization::replaceOpWithBufferizedValues(RewriterBase &rewriter,
                                                   Operation *op,
                                                   ValueRange values) {
@@ -989,3 +995,22 @@ bufferization::detail::unknownGetAliasingValues(OpOperand &opOperand) {
           r.addAlias({bbArg, BufferRelation::Unknown, /*isDefinite=*/false});
   return r;
 }
+
+static bool isaTensor(Type t) { return isa<TensorType>(t); }
+
+bool bufferization::detail::defaultHasTensorSemanticsForBufferization(
+    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;
+
+  bool hasTensorResult = any_of(op->getResultTypes(), isaTensor);
+  bool hasTensorOperand = any_of(op->getOperandTypes(), isaTensor);
+  return hasTensorResult || hasTensorOperand;
+}
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
index f2125feeda5415..79be7bed79db38 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
+++ b/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 {
@@ -417,7 +392,7 @@ class BufferizationRewriter : public IRRewriter, public RewriterBase::Listener {
       return;
 
     // Skip non-tensor ops.
-    if (!hasTensorSemantics(op))
+    if (!hasTensorSemanticsForBufferization(op))
       return;
 
     // Skip ops that are not allowed to be bufferized.
@@ -470,7 +445,7 @@ LogicalResult bufferization::bufferizeOp(Operation *op,
   // canonicalize away (or canonicalize to more precise layouts).
   SmallVector<Operation *> worklist;
   op->walk<WalkOrder::PostOrder>([&](Operation *op) {
-    if (hasTensorSemantics(op))
+    if (hasTensorSemanticsForBufferization(op))
       worklist.push_back(op);
   });
 
@@ -492,7 +467,7 @@ LogicalResult bufferization::bufferizeOp(Operation *op,
     if (!options.isOpAllowed(nextOp))
       continue;
     // Skip ops that no longer have tensor semantics.
-    if (!hasTensorSemantics(nextOp))
+    if (!hasTensorSemanticsForBufferization(nextOp))
       continue;
     // Check for unsupported unstructured control flow.
     if (!bufferizableOp.supportsUnstructuredControlFlow())
@@ -546,7 +521,7 @@ LogicalResult bufferization::bufferizeOp(Operation *op,
       continue;
     // Ops that no longer have tensor semantics (because they were updated
     // in-place) are allowed.
-    if (!hasTensorSemantics(op))
+    if (!hasTensorSemanticsForBufferization(op))
       continue;
     // Continue ops that are not allowed.
     if (!options.isOpAllowed(op))
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
index 3a8c397c02a809..5aab035f9f30fd 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
@@ -325,6 +325,28 @@ struct FuncOpInterface
 
   static bool supportsUnstructuredControlFlow() { return true; }
 
+  bool hasTensorSemanticsForBufferization(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 {
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp
index 1404ed8f43f964..aeda995fd585aa 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp
+++ b/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);

jpienaar · 2024-01-04T20:51:24Z

hasTensorSemanticsForBufferization is a bit of a weird name ... This represents ops that have value semantics and need to be bufferized? Or is there more to it?

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

jpienaar · 2024-01-04T20:52:55Z

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

+    return true;
+
+  bool hasTensorResult = any_of(op->getResultTypes(), isaTensor);
+  bool hasTensorOperand = any_of(op->getOperandTypes(), isaTensor);


I was wondering: is outputs only sufficient?

In practice, results are usually sufficient. But we could have a hypothetical op such as test.dump_tensor, which takes a tensor operand and has no result.

jpienaar · 2024-01-04T20:53:34Z

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

+  // Bufferize all other ops.
+  for (Operation &op : moduleOp.getOps()) {
+    // Functions were already bufferized.
+    if (isa<func::FuncOp>(&op))


Should this be on function like interface?

For consistency reasons, func::FuncOp would be better here. Other function-like ops are not supported at the moment (they will fail to bufferize) and we exclusively use func::FuncOp, func::CallOp, func::ReturnOp in this file at the moment.

matthias-springer · 2024-01-05T15:03:46Z

hasTensorSemanticsForBufferization is a bit of a weird name ... This represents ops that have value semantics and need to be bufferized? Or is there more to it?

It should really be called hasTensorSemantics. I just can't use that name because of a name collision.

Side note: This is not the first time that I'm running into such naming issues. There is an inconsistency in MLIR when (a) implementing an interface directly on an op and (b) implementing an interface with an external model. In the latter case, the interface methods are properly encapsulated: you have to explicitly cast<MyInterface>(my_op.getOperation()).interface_method() before calling an interface method. In the former case, you can directly call my_op.interface_method(). Name collisions could be avoided if we would treat all interface implementations like (b). (That would be a pretty large API change API. There may also be performance implications as (b) is no longer based on C++ inheritance and the compiled code may be less efficient.)

matthias-springer · 2024-01-10T11:15:48Z

This PR now depends on #77574, which aligns DestinationStyleOpInterface::hasTensorSemantics with the default implementation of BufferizableOpInterface::hasTensorSemantics.

For all ops that implement both DestinationStyleOpInterface and BufferizableOpInterface directly on the op without an external model (currently, the only such op is MaterializeInDestinationOp), those two implementations now "agree".

See comment in BufferizationOps.td:

    // Both `DestinationStyleOpInterface` and `BufferizableOpInterface` define
    // `hasTensorSemantics`. Both return the same result, but we have to choose
    // one to disambiguate the method lookup.
    using DestinationStyleOpInterface::Trait<MaterializeInDestinationOp>
        ::hasTensorSemantics;

It is possible to have different implementations as long as at least one of the two interfaces is implemented as an external model. In such a case, the functions of the external model can only be accessed by explicitly casting to the interface, and at that point it is clear which implementation is meant. In practice, I don't think that we will see cases where an op implements the two interfaces (external model or not) but has different implementations of hasTensorSemantics.

…fferSemantics` Rename interface functions as follows: * `hasTensorSemantics` -> `hasPureTensorSemantics` * `hasBufferSemantics` -> `hasPureBufferSemantics` These two functions return "true" if the op has tensor/buffer operands but not buffer/tensor operands. Add two new interface functions: * `hasTensorSemantics`: Return "true" if the op has tensor operands. Whether the op has buffer operands or not does not matter. * `hasBufferSemantics`: Return "true" if the op has buffer operands. Whether the op has tensor operands or not does not matter. Also drop the "ranked" part from the interface, i.e., do not distinguish between ranked/unranked types. This change aligns the meaning of "tensor semantics" with the bufferization framework. (An op is supposed to be bufferized if it has tensor operands, and we don't care if it also has memref operands.) This change is in preparation of llvm#75273, which adds `BufferizableOpInterface::hasTensorSemantics`.

…ufferSemantics` (#77574) Rename interface functions as follows: * `hasTensorSemantics` -> `hasPureTensorSemantics` * `hasBufferSemantics` -> `hasPureBufferSemantics` These two functions return "true" if the op has tensor/buffer operands but not buffer/tensor operands. Also drop the "ranked" part from the interface, i.e., do not distinguish between ranked/unranked types. The new function names describe the functions more accurately. They also align their semantics with the notion of "tensor semantics" with the bufferization framework. (An op is supposed to be bufferized if it has tensor operands, and we don't care if it also has memref operands.) This change is in preparation of #75273, which adds `BufferizableOpInterface::hasTensorSemantics`. By renaming the functions in the `DestinationStyleOpInterface`, we can avoid name clashes between the two interfaces.

matthias-springer · 2024-01-12T09:18:41Z

DestinationStyleOpInterface::hasTensorSemantics was renamed to hasPureTensorSemantics in #77574. There are no name clashes anymore.

jpienaar

Looks good, thanks!

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

Add a new interface method to `BufferizableOpInterface`: `hasTensorSemanticsForBufferization`. 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 (llvm#75103). The new interface method can return "true" for such ops. This change also decouples `bufferization::bufferizeOp` a bit from the func dialect.

…ufferSemantics` (llvm#77574) Rename interface functions as follows: * `hasTensorSemantics` -> `hasPureTensorSemantics` * `hasBufferSemantics` -> `hasPureBufferSemantics` These two functions return "true" if the op has tensor/buffer operands but not buffer/tensor operands. Also drop the "ranked" part from the interface, i.e., do not distinguish between ranked/unranked types. The new function names describe the functions more accurately. They also align their semantics with the notion of "tensor semantics" with the bufferization framework. (An op is supposed to be bufferized if it has tensor operands, and we don't care if it also has memref operands.) This change is in preparation of llvm#75273, which adds `BufferizableOpInterface::hasTensorSemantics`. By renaming the functions in the `DestinationStyleOpInterface`, we can avoid name clashes between the two interfaces.

…s` (llvm#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 (llvm#75103). The new interface method can return "true" for such ops. This change also decouples `bufferization::bufferizeOp` a bit from the func dialect.

matthias-springer requested review from jpienaar and nicolasvasilache December 13, 2023 02:18

llvmbot added mlir mlir:bufferization Bufferization infrastructure labels Dec 13, 2023

matthias-springer mentioned this pull request Dec 13, 2023

[mlir][MLProgram] Add MLProgram to MemRef bufferization pass #75103

Merged

matthias-springer requested a review from stellaraccident December 20, 2023 02:11

jpienaar reviewed Jan 4, 2024

View reviewed changes

matthias-springer requested a review from jpienaar January 5, 2024 15:04

matthias-springer mentioned this pull request Jan 10, 2024

[mlir][Interfaces] DestinationStyleOpInterface: Rename hasTensor/BufferSemantics #77574

Merged

matthias-springer force-pushed the bufferizable_op_interface_tensor_semantics branch from 7ed95b1 to 81e929c Compare January 10, 2024 11:07

matthias-springer force-pushed the bufferizable_op_interface_tensor_semantics branch from 81e929c to f4cd395 Compare January 12, 2024 09:17

jpienaar approved these changes Jan 15, 2024

View reviewed changes

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp Outdated Show resolved Hide resolved

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp Outdated Show resolved Hide resolved

matthias-springer force-pushed the bufferizable_op_interface_tensor_semantics branch from f4cd395 to 9cb1e5d Compare January 16, 2024 09:02

matthias-springer merged commit 8f2d83d into llvm:main Jan 16, 2024

tzunghanjuang mentioned this pull request Oct 10, 2024

Migrating to One-Shot Bufferization PennyLaneAI/catalyst#1027

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[mlir][bufferization] Add `BufferizableOpInterface::hasTensorSemantics` #75273

[mlir][bufferization] Add `BufferizableOpInterface::hasTensorSemantics` #75273

Uh oh!

matthias-springer commented Dec 13, 2023 •

edited

Loading

Uh oh!

llvmbot commented Dec 13, 2023 •

edited

Loading

Uh oh!

jpienaar commented Jan 4, 2024

Uh oh!

Uh oh!

jpienaar Jan 4, 2024

Uh oh!

matthias-springer Jan 5, 2024

Uh oh!

jpienaar Jan 4, 2024

Uh oh!

matthias-springer Jan 5, 2024

Uh oh!

matthias-springer commented Jan 5, 2024

Uh oh!

matthias-springer commented Jan 10, 2024 •

edited

Loading

Uh oh!

matthias-springer commented Jan 12, 2024

Uh oh!

jpienaar left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

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

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

Uh oh!

Conversation

matthias-springer commented Dec 13, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Dec 13, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jpienaar commented Jan 4, 2024

Uh oh!

Uh oh!

jpienaar Jan 4, 2024

Choose a reason for hiding this comment

Uh oh!

matthias-springer Jan 5, 2024

Choose a reason for hiding this comment

Uh oh!

jpienaar Jan 4, 2024

Choose a reason for hiding this comment

Uh oh!

matthias-springer Jan 5, 2024

Choose a reason for hiding this comment

Uh oh!

matthias-springer commented Jan 5, 2024

Uh oh!

matthias-springer commented Jan 10, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

matthias-springer commented Jan 12, 2024

Uh oh!

jpienaar left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

[mlir][bufferization] Add `BufferizableOpInterface::hasTensorSemantics` #75273

[mlir][bufferization] Add `BufferizableOpInterface::hasTensorSemantics` #75273

matthias-springer commented Dec 13, 2023 •

edited

Loading

llvmbot commented Dec 13, 2023 •

edited

Loading

matthias-springer commented Jan 10, 2024 •

edited

Loading