[mlir][linalg][bufferize] Add BufferizableOpInterface

[springerm]

Summary

This commit adds a new op interface: BufferizableOpInterface. In the future, ops that implement this interface can be bufferized using Comprehensive Bufferize.

Note: The interface methods of this interface correspond to the "op interface" in ComprehensiveBufferize.cpp.

[springerm]

Just fyi, we talked about removing methods from the interface. That is done in D112902. This commit just translates the existing "op interface" from ComprehensiveBufferize.cpp into a .td file. The interface is exactly the one that's described in ComprehensiveBufferize.cpp (NFC).

[nicolasvasilache]

mlir/include/mlir/Dialect/Linalg/Transforms/BufferizableOpInterface.td

40

This would crash.
The doc says it's optional, should you remove it ?

61

Same here.

79

Is this the one we discussed as "property derived from other interface methods"?

From https://mlir.llvm.org/docs/Interfaces/, there is a variant you are interested in:

StaticInterfaceMethod<[{
  This method corresponds to a static interface method that has an explicit
  implementation of the method body. Given that the method body has been
  explicitly implemented, this method should not be defined by the operation
  implementing this method.

[nicolasvasilache]

mlir/include/mlir/Dialect/Linalg/Transforms/BufferizableOpInterface.td

79

This is fine as a followup at the place where you propose removing it.

[springerm]

mlir/include/mlir/Dialect/Linalg/Transforms/BufferizableOpInterface.td

40

Yes, I want it it to crash.

If an op does not have tensor OpOperands, this method should not be called in the first place. If it is called, there's a problem with our code.

It will also crash if a programmer forgot to implement this interface method.

79

Nice, will make it static in D112902 with a default impl that does the right thing.

mlir/include/mlir/Dialect/Linalg/CMakeLists.txt

 add_subdirectory(IR)
+add_subdirectory(Transforms)
 
 set(LLVM_TARGET_DEFINITIONS Passes.td)
 mlir_tablegen(Passes.h.inc -gen-pass-decls -name Linalg)
 mlir_tablegen(Passes.capi.h.inc -gen-pass-capi-header --prefix Linalg)
 mlir_tablegen(Passes.capi.cpp.inc -gen-pass-capi-impl --prefix Linalg)
 add_public_tablegen_target(MLIRLinalgPassIncGen)
 
 add_mlir_doc(Passes LinalgPasses ./ -gen-pass-doc)

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

+//===- BufferizableOpInterface.h - Comprehensive Bufferize ------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_LINALG_TRANSFORMS_BUFFERIZABLEOPINTERFACE_H_
+#define MLIR_DIALECT_LINALG_TRANSFORMS_BUFFERIZABLEOPINTERFACE_H_
+
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/Support/LLVM.h"
+
+namespace mlir {
+class BlockAndValueMapping;
+
+namespace linalg {
+class AllocationCallbacks;
+class BufferizationAliasInfo;
+
+/// Specify fine-grain relationship between buffers to enable more analysis.
+enum class BufferRelation {
+  None,
+  // TODO: ResultContainsOperand,
+  // TODO: OperandContainsResult,
+  Equivalent
+};
+} // namespace linalg
+} // namespace mlir
+
+#include "mlir/Dialect/Linalg/Transforms/BufferizableOpInterface.h.inc"
+
+#endif // MLIR_DIALECT_LINALG_TRANSFORMS_BUFFERIZABLEOPINTERFACE_H_

mlir/include/mlir/Dialect/Linalg/Transforms/BufferizableOpInterface.td

+//===-- BufferizableOpInterface.td - Compreh. Bufferize ----*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_LINALG_TRANSFORMS_BUFFERIZABLEOPINTERFACE
+#define MLIR_DIALECT_LINALG_TRANSFORMS_BUFFERIZABLEOPINTERFACE
+
+include "mlir/IR/OpBase.td"
+
+def BufferizableOpInterface : OpInterface<"BufferizableOpInterface"> {
+  let description = [{
+    An op interface for Comprehensive Bufferization. Ops that implement this
+    interface can be bufferized using Comprehensive Bufferization.
+  }];
+  let cppNamespace = "::mlir::linalg";
+  let methods = [
+      InterfaceMethod<
+        /*desc=*/[{
+          Return `true` if the given OpOperand bufferizes to a memory read. This
+          method will never be called on OpOperands that do not have a tensor
+          type.
+
+          Note: It is always safe to consider an OpOperand as a memory read,
+          even if it does actually not read; however, this can introduce
+          unnecessary out-of-place bufferization decisions. The analysis of
+          Comprehensive Bufferize considers OpOperands of unknown ops (that do
+          not implement this interface) as reading OpOperands.
+        }],
+        /*retType=*/"bool",
+        /*methodName=*/"bufferizesToMemoryRead",
+        /*args=*/(ins "OpOperand &":$opOperand),
+        /*methodBody=*/"",
+        /*defaultImplementation=*/[{
+          // Does not have to be implemented for ops without tensor OpOperands.
+          llvm_unreachable("bufferizesToMemoryRead not implemented");
+         }]
+      >,
+      InterfaceMethod<
+        /*desc=*/[{
+          Return `true` if the given OpOperand bufferizes to a memory write.
+          This method will never be called on OpOperands that do not have a
+          tensor type.
+
+          Note: It is always safe to consider an OpOperand as a memory write,
+          even if it does actually not write; however, this can introduce
+          unnecessary out-of-place bufferization decisions. The analysis of
+          Comprehensive Bufferize considers OpOperands of unknown ops (that do
+          not implement this interface) as writing OpOperands.
+        }],
+        /*retType=*/"bool",
+        /*methodName=*/"bufferizesToMemoryWrite",
+        /*args=*/(ins "OpOperand &":$opOperand),
+        /*methodBody=*/"",
+        /*defaultImplementation=*/[{
+          // Does not have to be implemented for ops without tensor OpOperands.
+          llvm_unreachable("bufferizesToMemoryWrite not implemented");
+         }]
+      >,
+      // TODO: Simplify this interface by removing `bufferizesToAliasOnly` and
+      // `getInplaceableOpResult`. Instead, always use `getAliasingOpResult`. If
+      // `getAliasingOpResult` returns a non-null value, we know that an alias
+      // is created. If `bufferizesToMemoryRead` and `bufferizesToMemoryWrite`
+      // return `false`, we know that the operands "bufferizes to alias only".
+      InterfaceMethod<
+        /*desc=*/[{
+          Return `true` if the given OpOperand creates an alias but does neither
+          read nor write. This implies that `bufferizesToMemoryRead` and
+          `bufferizesToMemoryWrite` must return `false`. This method will never
+          be called on OpOperands that do not have a tensor type.
+
+          Examples of such ops are `tensor.extract_slice` and `tensor.cast`.
+        }],
+        /*retType=*/"bool",
+        /*methodName=*/"bufferizesToAliasOnly",
+        /*args=*/(ins "OpOperand &":$opOperand),
+        /*methodBody=*/"",
+        /*defaultImplementation=*/[{
+          // For better debugging, run `bufferizesToMemoryWrite`, which fires an
+          // assertion when called on an op that should not have tensor
+          // OpOperands.
+          (void) cast<BufferizableOpInterface>($_op.getOperation())
+              .bufferizesToMemoryWrite(opOperand);
+          // Return `false` by default, as most ops are not "alias only".
+          return false;
+         }]
+      >,
+      InterfaceMethod<
+        /*desc=*/[{
+          Return the OpResult that can bufferize in-place with a given
+          OpOperand. Return a null value if the OpOperand cannot bufferize
+          in-place. This method will never be called on OpOperands that do not
+          have a tensor type.
+        }],
+        /*retType=*/"OpResult",
+        /*methodName=*/"getInplaceableOpResult",
+        /*args=*/(ins "OpOperand &":$opOperand),
+        /*methodBody=*/"",
+        /*defaultImplementation=*/[{
+          // Does not have to be implemented for ops without tensor OpOperands.
+          llvm_unreachable("getInplaceableOpResult not implemented");
+        }]
+      >,
+      InterfaceMethod<
+        /*desc=*/[{
+          Return the OpResult that aliases with a given OpOperand when
+          bufferized in-place. This is a superset of `getInplaceableOpResult`.
+          This method will never be called on OpOperands that do not have a
+          tensor type.
+        }],
+        /*retType=*/"OpResult",
+        /*methodName=*/"getAliasingOpResult",
+        /*args=*/(ins "OpOperand &":$opOperand),
+        /*methodBody=*/"",
+        /*defaultImplementation=*/[{
+          return cast<BufferizableOpInterface>($_op.getOperation())
+              .getInplaceableOpResult(opOperand);
+        }]
+      >,
+      InterfaceMethod<
+        /*desc=*/[{
+          Return the OpOperands that alias with a given OpResult when
+          bufferized in-place. This method will never be called on OpResults
+          that do not have a tensor type.
+
+          Note: This method can return multiple OpOperands, indicating that the
+          given OpResult may at runtime alias with any of the OpOperands. This
+          is useful for branches and for ops such as `std.select`.
+        }],
+        /*retType=*/"SmallVector<OpOperand *>",
+        /*methodName=*/"getAliasingOpOperand",
+        /*args=*/(ins "OpResult":$opResult),
+        /*methodBody=*/"",
+        /*defaultImplementation=*/[{
+          // Does not have to be implemented for ops without tensor OpResults.
+          llvm_unreachable("getInplaceableOpResult not implemented");
+        }]
+      >,
+      InterfaceMethod<
+        /*desc=*/[{
+          Return the buffer relation between the given OpOperand and its
+          aliasing OpResult when bufferized in-place. Most OpOperands have an
+          "equivalence" relation.
+
+          TODO: Support other relations such as "OpOperand is included in
+          OpResult".
+        }],
+        /*retType=*/"BufferRelation",
+        /*methodName=*/"bufferRelation",
+        /*args=*/(ins "OpOperand &":$opOperand),
+        /*methodBody=*/"",
+        /*defaultImplementation=*/[{
+          // Does not have to be implemented for ops without tensor OpOperands.
+          llvm_unreachable("bufferRelation not implemented");
+        }]
+      >,
+      // TODO: Simplify method signature: Pass an OpBuilder and a
+      // BufferizationState object.
+      InterfaceMethod<
+        /*desc=*/[{
+          Bufferize this op, i.e., rewrite it into a memref-based equivalent.
+          `bvm` maps tensor values to memref values and this method should map
+          tensor results to memref results after creating/modifying ops.
+        }],
+        /*retType=*/"LogicalResult",
+        /*methodName=*/"bufferize",
+        /*args=*/(ins "OpBuilder &":$b,
+                      "BlockAndValueMapping &":$bvm,
+                      "BufferizationAliasInfo &":$aliasInfo,
+                      "AllocationCallbacks &":$allocationFn),
+        /*methodBody=*/"",
+        /*defaultImplementation=*/[{
+          llvm_unreachable("bufferize not implemented");
+          return failure();
+        }]
+      >,
+  ];
+}
+
+#endif  // MLIR_DIALECT_LINALG_TRANSFORMS_BUFFERIZABLEOPINTERFACE

mlir/include/mlir/Dialect/Linalg/Transforms/CMakeLists.txt

+set(LLVM_TARGET_DEFINITIONS BufferizableOpInterface.td)
+mlir_tablegen(BufferizableOpInterface.h.inc -gen-op-interface-decls)
+mlir_tablegen(BufferizableOpInterface.cpp.inc -gen-op-interface-defs)
+add_public_tablegen_target(MLIRBufferizableOpInterfaceIncGen)
+add_dependencies(mlir-headers MLIRBufferizableOpInterfaceIncGen)

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

 /// needs to look through their uses.
 /// ExtractSliceOp + InsertSliceOp have special joint behavior: they may
 /// bufferize to the same buffer (i.e. subview), which is what introduces the
 /// need for bufferization classes.
 /// Some of these functionalities could be refactored in a Bufferizer class that
 /// uses BufferizationAliasInfo.
 class BufferizationAliasInfo {
 public:
-  /// Specify fine-grain relationship between buffers to enable more analysis.
-  enum class BufferRelation {
-    None,
-    // TODO: ResultContainsOperand,
-    // TODO: OperandContainsResult,
-    Equivalent
-  };
-
   explicit BufferizationAliasInfo(Operation *rootOp);
 
   /// Add a new entry for `v` in the `aliasInfo` and `equivalentInfo`. In the
   /// beginning the alias and equivalence sets only contain `v` itself.
   void createAliasInfoEntry(Value v);
 
   /// Insert an info entry for `newValue` and merge its alias set with that of
   /// `alias`.

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

+//===- BufferizableOpInterface.cpp - Comprehensive Bufferize --------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Linalg/Transforms/BufferizableOpInterface.h"
+
+namespace mlir {
+namespace linalg {
+
+#include "mlir/Dialect/Linalg/Transforms/BufferizableOpInterface.cpp.inc"
+
+} // namespace linalg
+} // namespace mlir

mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt

+set(LLVM_OPTIONAL_SOURCES
+  Bufferize.cpp
+  BufferizableOpInterface.cpp
+  CodegenStrategy.cpp
+  ComprehensiveBufferize.cpp
+  Detensorize.cpp
+  Distribution.cpp
+  DropUnitDims.cpp
+  ElementwiseOpFusion.cpp
+  ElementwiseToLinalg.cpp
+  Fusion.cpp
+  FusionOnTensors.cpp
+  Generalization.cpp
+  Hoisting.cpp
+  HoistPadding.cpp
+  InlineScalarOperands.cpp
+  Interchange.cpp
+  Loops.cpp
+  LinalgStrategyPasses.cpp
+  Promotion.cpp
+  Tiling.cpp
+  Transforms.cpp
+  Vectorization.cpp
+)
+
+add_mlir_dialect_library(MLIRBufferizableOpInterface
+  BufferizableOpInterface.cpp
+
+  DEPENDS
+  MLIRBufferizableOpInterfaceIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+)
+
 add_mlir_dialect_library(MLIRLinalgTransforms
   Bufferize.cpp
   CodegenStrategy.cpp
   ComprehensiveBufferize.cpp
   Detensorize.cpp
   Distribution.cpp
   DropUnitDims.cpp
   ElementwiseOpFusion.cpp
add_mlir_dialect_library(MLIRLinalgTransforms
   DEPENDS
   MLIRLinalgPassIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffine
   MLIRAffineUtils
   MLIRAnalysis
   MLIRArithmetic
+  MLIRBufferizableOpInterface
   MLIRComplex
   MLIRInferTypeOpInterface
   MLIRIR
   MLIRMemRef
   MLIRLinalgAnalysis
   MLIRLinalg
   MLIRLinalgUtils
   MLIRSCF

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

 
 #include "mlir/Dialect/Linalg/Transforms/ComprehensiveBufferize.h"
 
 #include <random>
 
 #include "PassDetail.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/Linalg/Passes.h"
+#include "mlir/Dialect/Linalg/Transforms/BufferizableOpInterface.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/Operation.h"

 #include "llvm/Support/FormatVariadic.h"
 
 #define DEBUG_TYPE "comprehensive-module-bufferize"
 
 using namespace mlir;
 using namespace linalg;
 using namespace tensor;
 
-using BufferRelation = BufferizationAliasInfo::BufferRelation;
-
 #define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
 #define LDBG(X) LLVM_DEBUG(DBGS() << X)
 
 // TODO: from some HW description.
 static constexpr int64_t kBufferAlignments = 128;
 
 // Forward declarations.
 static std::string printOperationInfo(Operation *, bool prefix = true);

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

tbl_outs = [
         ),
     ],
     tblgen = ":mlir-tblgen",
     td_file = "include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td",
     deps = [":LinalgStructuredOpsTdFiles"],
 )
 
 td_library(
+    name = "BufferizableOpInterfaceTdFiles",
+    srcs = [
+        "include/mlir/Dialect/Linalg/Transforms/BufferizableOpInterface.td",
+    ],
+    includes = ["include"],
+    deps = [
+        ":OpBaseTdFiles",
+    ],
+)
+
+gentbl_cc_library(
+    name = "BufferizableOpInterfaceIncGen",
+    strip_include_prefix = "include",
+    tbl_outs = [
+        (
+            ["-gen-op-interface-decls"],
+            "include/mlir/Dialect/Linalg/Transforms/BufferizableOpInterface.h.inc",
+        ),
+        (
+            ["-gen-op-interface-defs"],
+            "include/mlir/Dialect/Linalg/Transforms/BufferizableOpInterface.cpp.inc",
+        ),
+    ],
+    tblgen = ":mlir-tblgen",
+    td_file = "include/mlir/Dialect/Linalg/Transforms/BufferizableOpInterface.td",
+    deps = [
+        ":BufferizableOpInterfaceTdFiles",
+    ],
+)
+
+cc_library(
+    name = "BufferizableOpInterface",
+    srcs = [
+        "lib/Dialect/Linalg/Transforms/BufferizableOpInterface.cpp",
+    ],
+    hdrs = [
+        "include/mlir/Dialect/Linalg/Transforms/BufferizableOpInterface.h",
+    ],
+    includes = ["include"],
+    deps = [
+        ":BufferizableOpInterfaceIncGen",
+        ":IR",
+        ":Support",
+        "//llvm:Support",
+    ],
+)
+
+td_library(
     name = "LinalgDocTdFiles",
     srcs = ["include/mlir/Dialect/Linalg/IR/LinalgDoc.td"],
     includes = ["include"],
     deps = [
         ":LinalgOpsTdFiles",
         ":LinalgStructuredOpsTdFiles",
     ],
 )
gentbl_cc_library(
     ],
     tblgen = ":mlir-tblgen",
     td_file = "include/mlir/Dialect/Linalg/Passes.td",
     deps = [":PassBaseTdFiles"],
 )
 
 cc_library(
     name = "LinalgTransforms",
-    srcs = glob([
+    srcs = glob(
+        [
             "lib/Dialect/Linalg/Transforms/*.cpp",
             "lib/Dialect/Linalg/Transforms/*.h",
-    ]) + [
+        ],
+        exclude = ["lib/Dialect/Linalg/Transforms/BufferizableOpInterface.cpp"],
+    ) + [
         "lib/Dialect/Linalg/Analysis/DependenceAnalysis.cpp",
         "lib/Dialect/Linalg/Utils/Utils.cpp",
     ],
     hdrs = [
         "include/mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h",
         "include/mlir/Dialect/Linalg/Passes.h",
         "include/mlir/Dialect/Linalg/Transforms/CodegenStrategy.h",
         "include/mlir/Dialect/Linalg/Transforms/ComprehensiveBufferize.h",
         "include/mlir/Dialect/Linalg/Transforms/HoistPadding.h",
         "include/mlir/Dialect/Linalg/Transforms/Hoisting.h",
         "include/mlir/Dialect/Linalg/Transforms/Transforms.h",
         "include/mlir/Dialect/Linalg/Utils/Utils.h",
     ],
     includes = ["include"],
     deps = [
         ":Affine",
         ":AffineUtils",
         ":Analysis",
         ":ArithmeticDialect",
+        ":BufferizableOpInterface",
         ":ComplexDialect",
         ":DialectUtils",
         ":IR",
         ":InferTypeOpInterface",
         ":LinalgOps",
         ":LinalgPassIncGen",
         ":LinalgStructuredOpsIncGen",
         ":MathDialect",