[mlir][linalg][bufferize] Remove remaining linalg dependencies

[springerm]

Summary

• Move "linalg.inplaceable" attr name literals to BufferizableOpInterface.
• Use memref.copy by default. Override to linalg.copy in ComprehensiveBufferizePass.

These are the last remaining code dependencies on Linalg in Comprehensive Bufferize. The next commit will make ComprehensiveBufferize independent of the Linalg dialect.

Depends On D113388

[aartbik]

Good for the sparse change.

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

let extraClassDeclaration = [{
     bool bufferizesToAliasOnly(OpOperand &opOperand) {
       auto bufferizableOp =
           cast<BufferizableOpInterface>(getOperation());
       return !bufferizableOp.bufferizesToMemoryRead(opOperand)
           && !bufferizableOp.bufferizesToMemoryWrite(opOperand)
           && static_cast<bool>(
               bufferizableOp.getAliasingOpResult(opOperand));
     }
+
+    // TODO: The following two attributes should belong to the tensor dialect.
+    // The corresponding verifier should also be in the tensor dialect.
+    /// Attribute name used to mark region arguments that can be bufferized
+    /// in-place during linalg comprehensive bufferization.
+    constexpr const static ::llvm::StringLiteral
+      kInplaceableAttrName = "linalg.inplaceable";
+
+    /// Attribute name used to mark the bufferization layout for region
+    /// arguments during linalg comprehensive bufferization.
+    constexpr const static ::llvm::StringLiteral
+      kBufferLayoutAttrName = "linalg.buffer_layout";
   }];
 }
 
 #endif  // BUFFERIZABLE_OP_INTERFACE

mlir/include/mlir/Dialect/Linalg/IR/LinalgBase.td

def Linalg_Dialect : Dialect {
   let hasCanonicalizer = 1;
   let hasOperationAttrVerify = 1;
   let hasConstantMaterializer = 1;
   let extraClassDeclaration = [{
     /// Attribute name used to to memoize indexing maps for named ops.
     constexpr const static ::llvm::StringLiteral
         kMemoizedIndexingMapsAttrName = "linalg.memoized_indexing_maps";
 
-    /// Attribute name used to mark region arguments that can be bufferized
-    /// in-place during linalg comprehensive bufferization.
-    constexpr const static ::llvm::StringLiteral
-      kInplaceableAttrName = "linalg.inplaceable";
-
-    /// Attribute name used to mark the bufferization layout for region
-    /// arguments during linalg comprehensive bufferization.
-    constexpr const static ::llvm::StringLiteral
-      kBufferLayoutAttrName = "linalg.buffer_layout";
-
     using RegionBuilderFunType =
       llvm::function_ref<void(ImplicitLocOpBuilder &b, Block &)>;
     RegionBuilderFunType getRegionBuilder(StringRef name) {
       return namedStructuredOpRegionBuilders.lookup(name);
     }
     private:
       llvm::StringMap<RegionBuilderFunType> namedStructuredOpRegionBuilders;
   }];
 }
 
 // Whether a type is a RangeType.
 def LinalgIsRangeTypePred : CPred<"$_self.isa<RangeType>()">;
 def Range : DialectType<Linalg_Dialect, LinalgIsRangeTypePred, "range">;
 
 #endif // LINALG_BASE

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

static void setInPlaceOpResult(OpResult opResult, bool inPlace) {
   op->setAttr(kInPlaceResultsAttrName,
               OpBuilder(op).getStrArrayAttr(inPlaceVector));
 }
 
 /// Set the attribute that triggers inplace bufferization on a FuncOp argument
 /// `bbArg`.
 static void setInPlaceFuncArgument(BlockArgument bbArg, bool inPlace) {
   auto funcOp = cast<FuncOp>(bbArg.getOwner()->getParentOp());
-  funcOp.setArgAttr(bbArg.getArgNumber(), LinalgDialect::kInplaceableAttrName,
+  funcOp.setArgAttr(bbArg.getArgNumber(),
+                    BufferizableOpInterface::kInplaceableAttrName,
                     BoolAttr::get(bbArg.getContext(), inPlace));
 }
 
 /// Remove the attribute that triggers inplace bufferization on a FuncOp
 /// argument `bbArg`.
 static void removeBufferizationFuncArguments(BlockArgument bbArg) {
   auto funcOp = cast<FuncOp>(bbArg.getOwner()->getParentOp());
   funcOp.removeArgAttr(bbArg.getArgNumber(),
-                       LinalgDialect::kBufferLayoutAttrName);
+                       BufferizableOpInterface::kBufferLayoutAttrName);
   funcOp.removeArgAttr(bbArg.getArgNumber(),
-                       LinalgDialect::kInplaceableAttrName);
+                       BufferizableOpInterface::kInplaceableAttrName);
 }
 
 //===----------------------------------------------------------------------===//
 // Printing helpers.
 //===----------------------------------------------------------------------===//
 
 #ifndef NDEBUG
 /// Helper method printing the bufferization information of a buffer / tensor.
if (none_of(op->getOperandTypes(), isaTensor) &&
       return;
     ops.push_back(op);
   });
 
   // Set the function arguments marked with inplaceable to be known as
   // bufferizing to a writeable memory.
   for (BlockArgument bbArg : funcOp.getArguments()) {
     BoolAttr inplaceAttr = funcOp.getArgAttrOfType<BoolAttr>(
-        bbArg.getArgNumber(), LinalgDialect::kInplaceableAttrName);
+        bbArg.getArgNumber(), BufferizableOpInterface::kInplaceableAttrName);
     if (inplaceAttr && inplaceAttr.getValue())
       aliasInfo.setBufferizesToWritableMemory(bbArg);
   }
 
   LogicalResult res =
       inPlaceAnalysis(ops, aliasInfo, domInfo, analysisFuzzerSeed);
   LDBG("End InPlaceAnalysisFuncOpInternals:\n" << funcOp << '\n');
 
void mlir::linalg::comprehensive_bufferize::defaultDeallocationFn(
     OpBuilder &b, Location loc, Value allocatedBuffer) {
   b.create<memref::DeallocOp>(loc, allocatedBuffer);
 }
 
 void mlir::linalg::comprehensive_bufferize::defaultMemCpyFn(OpBuilder &b,
                                                             Location loc,
                                                             Value from,
                                                             Value to) {
-  b.create<CopyOp>(loc, from, to);
+  b.create<memref::CopyOp>(loc, from, to);
 }
 
 LogicalResult mlir::linalg::comprehensive_bufferize::bufferizeOp(
     Operation *op, BufferizationState &state,
     DenseMap<FuncOp, FunctionType> *bufferizedFunctionTypes) {
   OpBuilder b(op->getContext());
 
   // CallOps are handled separately.
for (FuncOp funcOp : orderedFuncOps) {
     SmallVector<Type> argumentTypes;
     // Iterate on each function argument and check it it was marked with a
     // desired layout.
     for (auto it : llvm::enumerate(funcOp.getType().getInputs())) {
       int argNumber = it.index();
       Type inputType = it.value();
       auto memrefType = inputType.dyn_cast<MemRefType>();
       auto layoutAttr = funcOp.getArgAttrOfType<AffineMapAttr>(
-          argNumber, LinalgDialect::kBufferLayoutAttrName);
+          argNumber, BufferizableOpInterface::kBufferLayoutAttrName);
       AffineMap desiredLayoutMap =
           layoutAttr ? layoutAttr.getValue() : AffineMap();
       AffineMap currentLayoutMap =
           memrefType ? getStridedLinearLayoutMap(memrefType) : AffineMap();
       if (!memrefType || !layoutAttr || desiredLayoutMap == currentLayoutMap) {
         argumentTypes.push_back(inputType);
         foreachCaller(callerMap, funcOp, [&](Operation *caller) {
           operandsPerCaller.find(caller)->getSecond().push_back(

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

 add_mlir_dialect_library(MLIRLinalg
   LinalgInterfaces.cpp
   LinalgOps.cpp
   LinalgTypes.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Linalg
 
   DEPENDS
   MLIRLinalgInterfacesIncGen
   MLIRLinalgOpsIncGen
   MLIRLinalgStructuredOpsIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffine
   MLIRArithmetic
+  MLIRBufferizableOpInterface
   MLIRDialectUtils
   MLIRInferTypeOpInterface
   MLIRIR
   MLIRParser
   MLIRSideEffectInterfaces
   MLIRSCF
   MLIRStandard
   MLIRMath
   MLIRMemRef
   MLIRTensor
   MLIRTilingInterface
   MLIRViewLikeInterface
   )

mlir/lib/Dialect/Linalg/IR/LinalgTypes.cpp

 //===- Dialect.cpp - Implementation of the linalg dialect and types -------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the Linalg dialect types and dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
+#include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/FunctionSupport.h"
 #include "mlir/Parser.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/InliningUtils.h"

 //===----------------------------------------------------------------------===//
 
 /// Attribute name used to to memoize indexing maps for named ops.
 constexpr const ::llvm::StringLiteral
     LinalgDialect::kMemoizedIndexingMapsAttrName;
 
 /// Attribute name used to mark the bufferization layout for region
 /// arguments during linalg comprehensive bufferization.
-constexpr const ::llvm::StringLiteral LinalgDialect::kBufferLayoutAttrName;
+constexpr const ::llvm::StringLiteral
+    comprehensive_bufferize::BufferizableOpInterface::kBufferLayoutAttrName;
 
 /// Attribute name used to mark region arguments that can be bufferized
 /// in-place during linalg comprehensive bufferization.
-constexpr const ::llvm::StringLiteral LinalgDialect::kInplaceableAttrName;
+constexpr const ::llvm::StringLiteral
+    comprehensive_bufferize::BufferizableOpInterface::kInplaceableAttrName;
 
 /// Trait to check if T provides a `regionBuilder` method.
 template <typename T, typename... Args>
 using has_region_builder = decltype(T::regionBuilder);
 template <typename T>
 using detect_has_region_builder = llvm::is_detected<has_region_builder, T>;
 
 /// SFINAE helper for single C++ class without a `regionBuilder` method (e.g.

 
 void mlir::linalg::LinalgDialect::printType(Type type,
                                             DialectAsmPrinter &os) const {
   print(type.cast<RangeType>(), os);
 }
 
 LogicalResult LinalgDialect::verifyOperationAttribute(Operation *op,
                                                       NamedAttribute attr) {
-  if (attr.first == LinalgDialect::kInplaceableAttrName) {
+  using comprehensive_bufferize::BufferizableOpInterface;
+
+  if (attr.first == BufferizableOpInterface::kInplaceableAttrName) {
     if (!attr.second.isa<BoolAttr>()) {
-      return op->emitError() << "'" << LinalgDialect::kInplaceableAttrName
-                             << "' is expected to be a boolean attribute";
+      return op->emitError()
+             << "'" << BufferizableOpInterface::kInplaceableAttrName
     }
     if (!op->hasTrait<OpTrait::FunctionLike>())
       return op->emitError() << "expected " << attr.first
                              << " to be used on function-like operations";
     return success();
   }
-  if (attr.first == LinalgDialect::kBufferLayoutAttrName) {
+  if (attr.first == BufferizableOpInterface::kBufferLayoutAttrName) {
     if (!attr.second.isa<AffineMapAttr>()) {
-      return op->emitError() << "'" << LinalgDialect::kBufferLayoutAttrName
-                             << "' is expected to be a affine map attribute";
+      return op->emitError()
+             << "'" << BufferizableOpInterface::kBufferLayoutAttrName
     }
     if (!op->hasTrait<OpTrait::FunctionLike>())
       return op->emitError() << "expected " << attr.first
                              << " to be used on function-like operations";
     return success();
   }
   if (attr.first == LinalgDialect::kMemoizedIndexingMapsAttrName)
     return success();
   return op->emitError() << "attribute '" << attr.first
                          << "' not supported by the linalg dialect";
 }

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

 
 void LinalgComprehensiveModuleBufferize::runOnOperation() {
   BufferizationOptions options;
   if (useAlloca) {
     options.allocationFns->allocationFn = allocationFnUsingAlloca;
     options.allocationFns->deallocationFn = [](OpBuilder &b, Location loc,
                                                Value v) {};
   }
+  // TODO: Change to memref::CopyOp (default memCpyFn).
+  options.allocationFns->memCpyFn = [](OpBuilder &b, Location loc, Value from,
+                                       Value to) {
+    b.create<linalg::CopyOp>(loc, from, to);
+  };
+
   options.allowReturnMemref = allowReturnMemref;
   options.analysisFuzzerSeed = analysisFuzzerSeed;
   options.testAnalysisOnly = testAnalysisOnly;
 
   ModuleOp moduleOp = getOperation();
   applyEnablingTransformations(moduleOp);
 
   if (failed(runComprehensiveBufferize(moduleOp, options))) {

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

 add_mlir_dialect_library(MLIRSparseTensorTransforms
   Sparsification.cpp
   SparseTensorConversion.cpp
   SparseTensorPasses.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SparseTensor
 
   DEPENDS
   MLIRSparseTensorPassIncGen
 
   LINK_LIBS PUBLIC
   MLIRArithmetic
+  MLIRBufferizableOpInterface
   MLIRIR
   MLIRLLVMIR
   MLIRLinalg
   MLIRLinalgTransforms
   MLIRMemRef
   MLIRPass
   MLIRSCF
   MLIRStandard
   MLIRSparseTensor
   MLIRSparseTensorUtils
   MLIRTensor
   MLIRTransforms
   MLIRVector
 )

mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp

 //===- Sparsification.cpp - Implementation of sparsification --------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements converting sparse tensor types to actual sparse code.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/SCF/Transforms.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Dialect/SparseTensor/Utils/Merger.h"
static bool computeIterationGraph(Merger &merger, linalg::GenericOp op,
   return true;
 }
 
 /// Returns true if tensor has an in-place annotation.
 static bool isInPlace(Value val) {
   if (auto arg = val.dyn_cast<BlockArgument>())
     if (auto funcOp = dyn_cast<FuncOp>(arg.getOwner()->getParentOp()))
       if (auto attr = funcOp.getArgAttrOfType<BoolAttr>(
-              arg.getArgNumber(), linalg::LinalgDialect::kInplaceableAttrName))
+              arg.getArgNumber(),
+              linalg::comprehensive_bufferize::BufferizableOpInterface::
+                  kInplaceableAttrName))
         return attr.getValue();
   return false;
 }
 
 /// Returns true if tensor materializes into the computation.
 static bool isMaterializing(Value val) {
   return val.getDefiningOp<linalg::InitTensorOp>() ||
          val.getDefiningOp<InitOp>();

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

cc_library(
     hdrs = [
         "include/mlir/Dialect/SparseTensor/Transforms/Passes.h",
         "include/mlir/ExecutionEngine/SparseTensorUtils.h",
     ],
     includes = ["include"],
     deps = [
         ":Affine",
         ":ArithmeticDialect",
+        ":BufferizableOpInterface",
         ":IR",
         ":LLVMDialect",
         ":LinalgOps",
         ":LinalgTransforms",
         ":MemRefDialect",
         ":Pass",
         ":SCFDialect",
         ":SparseTensor",
cc_library(
     hdrs = [
         "include/mlir/Dialect/Linalg/IR/LinalgOps.h",
         "include/mlir/Dialect/Linalg/IR/LinalgTypes.h",
     ],
     includes = ["include"],
     deps = [
         ":Affine",
         ":ArithmeticDialect",
+        ":BufferizableOpInterface",
         ":CopyOpInterface",
         ":DialectUtils",
         ":IR",
         ":InferTypeOpInterface",
         ":LinalgInterfaces",
         ":LinalgInterfacesIncGen",
         ":LinalgNamedStructuredOpsYamlIncGen",
         ":LinalgOpsIncGen",