[mlir][tensor] Move BufferizableOpInterface impl to tensor dialect

[springerm]

Summary

This is in preparation of unifying the existing bufferization with One-Shot bufferization.

A subsequent commit will replace tensor-bufferize's implementation with the BufferizableOpInterface-based implementation and move over missing test cases.

This commit also introduces a new AlwaysCopyBufferizationState that will be used for partial bufferization and always copies tensor OpOperands that bufferize to memory writes. We will start using this in a subsequent commit.

Depends On D117775

[nicolasvasilache]

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

220

This seems like this should be separated and tested ?

[springerm]

rebase

mlir/include/mlir/Dialect/Bufferization/Transforms/Bufferize.h

 #ifndef MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_BUFFERIZE_H
 #define MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_BUFFERIZE_H
 
 #include "mlir/Transforms/DialectConversion.h"
 
 namespace mlir {
 namespace bufferization {
 
+class BufferizationState;
+struct BufferizationOptions;
+
 /// A helper type converter class that automatically populates the relevant
 /// materializations and type conversions for bufferization.
 class BufferizeTypeConverter : public TypeConverter {
 public:
   BufferizeTypeConverter();
 };
 
 /// Marks ops used by bufferization for type conversion materializations as
 /// "legal" in the given ConversionTarget.
 ///
 /// This function should be called by all bufferization passes using
 /// BufferizeTypeConverter so that materializations work properly. One exception
 /// is bufferization passes doing "full" conversions, where it can be desirable
 /// for even the materializations to remain illegal so that they are eliminated,
 /// such as via the patterns in
 /// populateEliminateBufferizeMaterializationsPatterns.
 void populateBufferizeMaterializationLegality(ConversionTarget &target);
 
 /// Populate patterns to eliminate bufferize materializations.
 ///
 /// In particular, these are the tensor_load/buffer_cast ops.
 void populateEliminateBufferizeMaterializationsPatterns(
     BufferizeTypeConverter &typeConverter, RewritePatternSet &patterns);
 
-class BufferizationState;
-
 /// Bufferize `op` and its nested ops that implement `BufferizableOpInterface`.
 /// Whether buffer copies are needed or not is queried from `state`.
 ///
 /// Note: If `allowUnknownOps` is set to false, bufferization fails when an
 /// unknown op (that does not implement `BufferizableOpInterface`) is found. No
 /// to_tensor/to_memref ops are inserted in that case.
 ///
-/// Note: Tje layout map chosen to bufferize is the most dynamic canonical
+/// Note: The layout map chosen to bufferize is the most dynamic canonical
 /// strided layout of the proper rank. This ensures compatibility with expected
 /// layouts after transformations. Combinations of memref.cast +
 /// canonicalization are responsible for clean ups.
 // TODO: Extract `options` from `state` and pass as separate argument.
 LogicalResult bufferizeOp(Operation *op, const BufferizationState &state);
 
 } // namespace bufferization
 } // namespace mlir
 
 #endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_BUFFERIZE_H

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

-//===- LinalgInterfaceImpl.h - Linalg Impl. of BufferizableOpInterface ----===//
-//
-// 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_COMPREHENSIVEBUFFERIZE_TENSORINTERFACEIMPL_H
-#define MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_TENSORINTERFACEIMPL_H
-
-namespace mlir {
-
-class DialectRegistry;
-
-namespace linalg {
-namespace comprehensive_bufferize {
-namespace tensor_ext {
-
-void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
-
-} // namespace tensor_ext
-} // namespace comprehensive_bufferize
-} // namespace linalg
-} // namespace mlir
-
-#endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_TENSORINTERFACEIMPL_H

mlir/include/mlir/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.h

+//===- BufferizableOpInterfaceImpl.h - Impl. of BufferizableOpInterface ---===//
+//
+// 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_TENSOR_BUFFERIZABLEOPINTERFACEIMPL_H
+#define MLIR_DIALECT_TENSOR_BUFFERIZABLEOPINTERFACEIMPL_H
+
+namespace mlir {
+class DialectRegistry;
+
+namespace tensor {
+void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
+} // namespace tensor
+} // namespace mlir
+
+#endif // MLIR_DIALECT_TENSOR_BUFFERIZABLEOPINTERFACEIMPL_H

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

 };
 } // namespace
 
 std::unique_ptr<OperationPass<FuncOp>>
 mlir::bufferization::createFinalizingBufferizePass() {
   return std::make_unique<FinalizingBufferizePass>();
 }
 
+//===----------------------------------------------------------------------===//
+// BufferizableOpInterface-based Bufferization
+//===----------------------------------------------------------------------===//
+
 static bool isaTensor(Type t) { return t.isa<TensorType>(); }
 
 /// 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;
 }
LogicalResult bufferization::bufferizeOp(Operation *op,
                                          const BufferizationState &state) {
   // Bufferize the op and its nested ops.
   OwningRewritePatternList patterns(op->getContext());
   patterns.add<BufferizationPattern>(op->getContext(), state);
   if (failed(applyPatternsAndFoldGreedily(op, std::move(patterns))))
     return failure();
 
   return checkBufferizationResult(op, state.getOptions());
 }

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

 set(LLVM_OPTIONAL_SOURCES
   AffineInterfaceImpl.cpp
   ArithInterfaceImpl.cpp
   LinalgInterfaceImpl.cpp
   ModuleBufferization.cpp
   SCFInterfaceImpl.cpp
   StdInterfaceImpl.cpp
-  TensorInterfaceImpl.cpp
   VectorInterfaceImpl.cpp
 )
 
 add_mlir_dialect_library(MLIRAffineBufferizableOpInterfaceImpl
   AffineInterfaceImpl.cpp
 
   LINK_LIBS PUBLIC
   MLIRAffine

 add_mlir_dialect_library(MLIRStdBufferizableOpInterfaceImpl
   StdInterfaceImpl.cpp
 
   LINK_LIBS PUBLIC
   MLIRBufferizableOpInterface
   MLIRStandard
 )
 
-add_mlir_dialect_library(MLIRTensorBufferizableOpInterfaceImpl
-  TensorInterfaceImpl.cpp
-
-  LINK_LIBS PUBLIC
-  MLIRBufferizableOpInterface
-  MLIRIR
-  MLIRMemRef
-  MLIRTensor
-)
-
 add_mlir_dialect_library(MLIRVectorBufferizableOpInterfaceImpl
   VectorInterfaceImpl.cpp
 
   LINK_LIBS PUBLIC
   MLIRBufferizableOpInterface
   MLIRIR
   MLIRVector
 )

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

add_mlir_dialect_library(MLIRLinalgTransforms
   MLIRSCFBufferizableOpInterfaceImpl
   MLIRSCFTransforms
   MLIRStdBufferizableOpInterfaceImpl
   MLIRPass
   MLIRStandard
   MLIRStandardOpsTransforms
   MLIRStandardToLLVM
   MLIRTensor
-  MLIRTensorBufferizableOpInterfaceImpl
+  MLIRTensorTransforms
   MLIRTransforms
   MLIRTransformUtils
   MLIRVector
   MLIRVectorBufferizableOpInterfaceImpl
   MLIRX86VectorTransforms
   MLIRVectorToSCF
 )

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

 #include "mlir/Dialect/Bufferization/IR/BufferizationInterfaceImpl.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/AffineInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/ArithInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/LinalgInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/SCFInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/StdInterfaceImpl.h"
-#include "mlir/Dialect/Linalg/ComprehensiveBufferize/TensorInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/VectorInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/Passes.h"
+#include "mlir/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/Passes.h"
 
 using namespace mlir;
 using namespace mlir::bufferization;
 using namespace mlir::linalg;
registry
                 arith::ArithmeticDialect, StandardOpsDialect, AffineDialect>();
     affine_ext::registerBufferizableOpInterfaceExternalModels(registry);
     arith_ext::registerBufferizableOpInterfaceExternalModels(registry);
     bufferization_ext::registerBufferizableOpInterfaceExternalModels(registry);
     linalg_ext::registerBufferizableOpInterfaceExternalModels(registry);
     scf_ext::registerBufferizableOpInterfaceExternalModels(registry);
     std_ext::registerModuleBufferizationExternalModels(registry);
     std_ext::registerBufferizableOpInterfaceExternalModels(registry);
-    tensor_ext::registerBufferizableOpInterfaceExternalModels(registry);
+    tensor::registerBufferizableOpInterfaceExternalModels(registry);
     vector_ext::registerBufferizableOpInterfaceExternalModels(registry);
   }
 };
 } // namespace
 
 static void applyEnablingTransformations(ModuleOp moduleOp) {
   RewritePatternSet patterns(moduleOp.getContext());
   patterns.add<GeneralizePadOpPattern>(moduleOp.getContext());

mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp

-//===- TensorInterfaceImpl.cpp - Tensor Impl. of BufferizableOpInterface --===//
+//===- BufferizableOpInterfaceImpl.cpp - Impl. of BufferizableOpInterface -===//
 //
 // 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/ComprehensiveBufferize/TensorInterfaceImpl.h"
+#include "mlir/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Operation.h"
 
 using namespace mlir;
 using namespace mlir::bufferization;
+using namespace mlir::tensor;
 
 namespace mlir {
-namespace linalg {
-namespace comprehensive_bufferize {
-namespace tensor_ext {
-
-using tensor::ExtractSliceOp;
-using tensor::InsertSliceOp;
+namespace tensor {
+namespace {
 
 struct CastOpInterface
     : public BufferizableOpInterface::ExternalModel<CastOpInterface,
                                                     tensor::CastOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                               const BufferizationState &state) const {
     return false;
   }
if (failed(createMemCpy(rewriter, loc, srcMemref, subView,
                             state.getOptions())))
       return failure();
 
     replaceOpWithBufferizedValues(rewriter, op, *dstMemref);
     return success();
   }
 };
 
-} // namespace tensor_ext
-} // namespace comprehensive_bufferize
-} // namespace linalg
+} // namespace
+} // namespace tensor
 } // namespace mlir
 
-void mlir::linalg::comprehensive_bufferize::tensor_ext::
-    registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry) {
-  registry.addOpInterface<tensor::CastOp, tensor_ext::CastOpInterface>();
-  registry.addOpInterface<tensor::DimOp, tensor_ext::DimOpInterface>();
-  registry.addOpInterface<tensor::ExtractSliceOp,
-                          tensor_ext::ExtractSliceOpInterface>();
-  registry.addOpInterface<tensor::ExtractOp, tensor_ext::ExtractOpInterface>();
-  registry.addOpInterface<tensor::InsertOp, tensor_ext::InsertOpInterface>();
-  registry.addOpInterface<tensor::InsertSliceOp,
-                          tensor_ext::InsertSliceOpInterface>();
+void mlir::tensor::registerBufferizableOpInterfaceExternalModels(
+    DialectRegistry &registry) {
+  registry.addOpInterface<CastOp, CastOpInterface>();
+  registry.addOpInterface<DimOp, DimOpInterface>();
+  registry.addOpInterface<ExtractSliceOp, ExtractSliceOpInterface>();
+  registry.addOpInterface<ExtractOp, ExtractOpInterface>();
+  registry.addOpInterface<InsertOp, InsertOpInterface>();
+  registry.addOpInterface<InsertSliceOp, InsertSliceOpInterface>();
 }

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

 add_mlir_dialect_library(MLIRTensorTransforms
+  BufferizableOpInterfaceImpl.cpp
   Bufferize.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Tensor/Transforms
 
   DEPENDS
   MLIRTensorTransformsIncGen
 
   LINK_LIBS PUBLIC
   MLIRArithmetic
+  MLIRBufferizableOpInterface
   MLIRBufferizationTransforms
   MLIRIR
   MLIRMemRef
   MLIRPass
   MLIRSCF
   MLIRTensor
   MLIRTransforms
   )

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

add_mlir_library(MLIRLinalgTestPasses
   MLIRLLVMToLLVMIRTranslation
   MLIRMemRef
   MLIRPass
   MLIRSCF
   MLIRSCFBufferizableOpInterfaceImpl
   MLIRStdBufferizableOpInterfaceImpl
   MLIRStandard
   MLIRTensor
-  MLIRTensorBufferizableOpInterfaceImpl
+  MLIRTensorTransforms
   MLIRTransformUtils
   MLIRVector
   MLIRVectorBufferizableOpInterfaceImpl
   MLIRVectorToSCF
   )

mlir/test/lib/Dialect/Linalg/TestComprehensiveBufferize.cpp

 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Bufferization/IR/BufferizationInterfaceImpl.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/AffineInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/ArithInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/LinalgInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/SCFInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/StdInterfaceImpl.h"
-#include "mlir/Dialect/Linalg/ComprehensiveBufferize/TensorInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/VectorInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/Passes.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 using namespace mlir::linalg::comprehensive_bufferize;
 using namespace mlir::bufferization;
registry.insert<bufferization::BufferizationDialect, linalg::LinalgDialect,
                     vector::VectorDialect, scf::SCFDialect, StandardOpsDialect,
                     arith::ArithmeticDialect, AffineDialect>();
     affine_ext::registerBufferizableOpInterfaceExternalModels(registry);
     arith_ext::registerBufferizableOpInterfaceExternalModels(registry);
     bufferization_ext::registerBufferizableOpInterfaceExternalModels(registry);
     linalg_ext::registerBufferizableOpInterfaceExternalModels(registry);
     scf_ext::registerBufferizableOpInterfaceExternalModels(registry);
     std_ext::registerBufferizableOpInterfaceExternalModels(registry);
-    tensor_ext::registerBufferizableOpInterfaceExternalModels(registry);
+    tensor::registerBufferizableOpInterfaceExternalModels(registry);
     vector_ext::registerBufferizableOpInterfaceExternalModels(registry);
   }
 
   void runOnOperation() override;
 
   Option<bool> allowReturnMemref{
       *this, "allow-return-memref",
       llvm::cl::desc("Allow returning/yielding memrefs from functions/blocks"),

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

 cc_library(
     name = "TensorTransforms",
     srcs = glob(
         [
             "lib/Dialect/Tensor/Transforms/*.cpp",
             "lib/Dialect/Tensor/Transforms/*.h",
         ],
     ),
-    hdrs = ["include/mlir/Dialect/Tensor/Transforms/Passes.h"],
+    hdrs = [
+        "include/mlir/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.h",
+        "include/mlir/Dialect/Tensor/Transforms/Passes.h",
+    ],
     includes = ["include"],
     deps = [
         ":ArithmeticDialect",
         ":Async",
+        ":BufferizableOpInterface",
         ":BufferizationDialect",
         ":BufferizationTransforms",
         ":IR",
         ":MemRefDialect",
         ":ParallelLoopMapperAttrGen",
         ":Pass",
         ":SCFDialect",
         ":StandardOps",
deps = [
         ":IR",
         ":StandardOps",
         ":Support",
         "//llvm:Support",
     ],
 )
 
 cc_library(
-    name = "TensorBufferizableOpInterfaceImpl",
-    srcs = [
-        "lib/Dialect/Linalg/ComprehensiveBufferize/TensorInterfaceImpl.cpp",
-    ],
-    hdrs = [
-        "include/mlir/Dialect/Linalg/ComprehensiveBufferize/TensorInterfaceImpl.h",
-    ],
-    includes = ["include"],
-    deps = [
-        ":BufferizableOpInterface",
-        ":IR",
-        ":MemRefDialect",
-        ":Support",
-        ":TensorDialect",
-        "//llvm:Support",
-    ],
-)
-
-cc_library(
     name = "VectorBufferizableOpInterfaceImpl",
     srcs = [
         "lib/Dialect/Linalg/ComprehensiveBufferize/VectorInterfaceImpl.cpp",
     ],
     hdrs = [
         "include/mlir/Dialect/Linalg/ComprehensiveBufferize/VectorInterfaceImpl.h",
     ],
     includes = ["include"],
deps = [
         ":Pass",
         ":SCFBufferizableOpInterfaceImpl",
         ":SCFDialect",
         ":SCFTransforms",
         ":StandardOps",
         ":StandardOpsTransforms",
         ":StdBufferizableOpInterfaceImpl",
         ":Support",
-        ":TensorBufferizableOpInterfaceImpl",
         ":TensorDialect",
+        ":TensorTransforms",
         ":TensorUtils",
         ":TransformUtils",
         ":VectorBufferizableOpInterfaceImpl",
         ":VectorOps",
         ":VectorToSCF",
         ":X86VectorTransforms",
         "//llvm:Support",
     ],

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

deps = [
         "//mlir:LinalgTransforms",
         "//mlir:MemRefDialect",
         "//mlir:Pass",
         "//mlir:SCFBufferizableOpInterfaceImpl",
         "//mlir:SCFDialect",
         "//mlir:SCFTransforms",
         "//mlir:StandardOps",
         "//mlir:StdBufferizableOpInterfaceImpl",
-        "//mlir:TensorBufferizableOpInterfaceImpl",
         "//mlir:TensorDialect",
+        "//mlir:TensorTransforms",
         "//mlir:TransformUtils",
         "//mlir:VectorBufferizableOpInterfaceImpl",
         "//mlir:VectorOps",
         "//mlir:VectorToSCF",
     ],
 )
 
 cc_library(