[mlir] Linalg refactor for using "bufferize" terminology.

[silvas]

Summary

Part of the refactor discussed in:
https://llvm.discourse.group/t/what-is-the-strategy-for-tensor-memref-conversion-bufferization/1938/17

[silvas]

update

[nicolasvasilache]

Thanks!

[silvas]

update

[aartbik]

bufferize it is!

mlir/include/mlir/Dialect/Linalg/Passes.h

 
 /// Create a pass to convert Linalg operations to affine.for loops and
 /// affine_load/affine_store accesses.
 /// Placeholder for now, this is NYI.
 std::unique_ptr<OperationPass<FuncOp>> createConvertLinalgToAffineLoopsPass();
 
 /// Create a pass to convert Linalg operations which work on tensors to use
 /// buffers instead.
-std::unique_ptr<OperationPass<ModuleOp>>
-createConvertLinalgOnTensorsToBuffersPass();
+std::unique_ptr<OperationPass<ModuleOp>> createLinalgBufferizePass();
 
 /// Patterns for fusing linalg operation on tensors.
 void populateLinalgTensorOpsFusionPatterns(MLIRContext *context,
                                            OwningRewritePatternList &patterns);
 
 /// Patterns to fold unit-extent dimensions in operands/results of linalg ops on
 /// tensors.
 void populateLinalgFoldUnitExtentDimsPatterns(

mlir/include/mlir/Dialect/Linalg/Passes.td

 }
 
 def LinalgLowerToLoops : FunctionPass<"convert-linalg-to-loops"> {
   let summary = "Lower the operations from the linalg dialect into loops";
   let constructor = "mlir::createConvertLinalgToLoopsPass()";
   let dependentDialects = ["linalg::LinalgDialect", "scf::SCFDialect", "AffineDialect"];
 }
 
-def LinalgOnTensorsToBuffers : Pass<"convert-linalg-on-tensors-to-buffers", "ModuleOp"> {
-  let summary = "Convert the Linalg operations which work on tensor-type "
-                "operands or results to use buffers instead";
-  let constructor = "mlir::createConvertLinalgOnTensorsToBuffersPass()";
+def LinalgBufferize : Pass<"linalg-bufferize", "ModuleOp"> {
+  let summary = "Bufferize the linalg dialect";
+  let constructor = "mlir::createLinalgBufferizePass()";
   let dependentDialects = ["linalg::LinalgDialect", "vector::VectorDialect"];
 }
 
 def LinalgLowerToParallelLoops
     : FunctionPass<"convert-linalg-to-parallel-loops"> {
   let summary = "Lower the operations from the linalg dialect into parallel "
                 "loops";
   let constructor = "mlir::createConvertLinalgToParallelLoopsPass()";

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

struct TiledAndFusedLinalgOps {
   SmallVector<Operation *, 4> unfusedLoops;
 };
 
 /// Populates patterns for vectorization of all ConvN-D ops.
 void populateConvVectorizationPatterns(
     MLIRContext *context, SmallVectorImpl<OwningRewritePatternList> &patterns,
     ArrayRef<int64_t> tileSizes);
 
-/// Populates the given list with patterns to convert Linalg operations on
-/// tensors to buffers.
-void populateConvertLinalgOnTensorsToBuffersPatterns(
-    MLIRContext *context, BufferAssignmentTypeConverter &converter,
+/// Populates the given list with patterns to bufferize linalg ops.
+void populateLinalgBufferizePatterns(MLIRContext *context,
+                                     BufferAssignmentTypeConverter &converter,
                                      OwningRewritePatternList &patterns);
 
 /// Performs standalone tiling of a single LinalgOp by `tileSizes`.
 /// and permute the loop nest according to `interchangeVector`
 /// The permutation is expressed as a list of integers that specify
 /// the new ordering of the loop nest. The length of `interchangeVector`
 /// must be equal to the length of `tileSizes`.
 /// An empty vector is interpreted as the identity permutation and the
 /// transformation returns early.

mlir/integration_test/Dialect/Linalg/CPU/test-tensor-e2e.mlir

-// RUN: mlir-opt %s -convert-linalg-on-tensors-to-buffers -convert-linalg-to-loops -convert-linalg-to-llvm -convert-std-to-llvm | \
+// RUN: mlir-opt %s -linalg-bufferize -convert-linalg-to-loops -convert-linalg-to-llvm -convert-std-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 func @foo() -> tensor<4xf32> {
   %0 = constant dense<[1.0, 2.0, 3.0, 4.0]> : tensor<4xf32>
   return %0 : tensor<4xf32>
 }

mlir/integration_test/Dialect/Linalg/CPU/test-tensor-matmul.mlir

-// RUN: mlir-opt %s -convert-linalg-on-tensors-to-buffers -convert-linalg-to-loops -convert-linalg-to-llvm -convert-std-to-llvm | \
+// RUN: mlir-opt %s -linalg-bufferize -convert-linalg-to-loops -convert-linalg-to-llvm -convert-std-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 func @main() {
   %A = constant dense<[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]> : tensor<2x3xf32>
   %B = constant dense<[[1.0, 2.0, 3.0, 4.0],
                        [5.0, 6.0, 7.0, 8.0],

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

-//===- TensorsToBuffers.cpp - Transformation from tensors to buffers ------===//
+//===- Bufferize.cpp - Bufferization of linalg ops ------------------===//
 //
 // 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 conversion from tensors to buffers on Linalg
-// operations.
-//
-//===----------------------------------------------------------------------===//
 
+#include "mlir/Transforms/Bufferize.h"
 #include "PassDetail.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Pass/Pass.h"
-#include "mlir/Transforms/Bufferize.h"
 
 using namespace ::mlir;
 using namespace ::mlir::linalg;
 
 static SmallVector<Range, 4> computeLoopRanges(Location loc, LinalgOp linalgOp,
                                                OpBuilder &b) {
   auto indexingMaps = llvm::to_vector<4>(
       linalgOp.indexing_maps().getAsValueRange<AffineMapAttr>());
LogicalResult mlir::linalg::TensorCastOpConverter::matchAndRewrite(
   rewriter.replaceOpWithNewOp<MemRefCastOp>(op, t, operands.front());
   return success();
 }
 
 namespace {
 
 /// Converts Linalg operations that work on tensor-type operands or results to
 /// work on buffers.
-struct ConvertLinalgOnTensorsToBuffers
-    : public LinalgOnTensorsToBuffersBase<ConvertLinalgOnTensorsToBuffers> {
+struct LinalgBufferizePass : public LinalgBufferizeBase<LinalgBufferizePass> {
   void runOnOperation() override {
     MLIRContext &context = getContext();
     ConversionTarget target(context);
     BufferAssignmentTypeConverter converter;
 
     // Mark all Standard operations legal.
     target.addLegalDialect<StandardOpsDialect, vector::VectorDialect>();
     target.addLegalOp<ModuleOp>();
target.addDynamicallyLegalOp<FuncOp>([&](FuncOp funcOp) {
                            [&](Type type) { return type.isa<MemRefType>(); }) &&
              converter.isLegal(&funcOp.getBody());
     });
 
     converter.setResultConversionKind<RankedTensorType, MemRefType>(
         BufferAssignmentTypeConverter::AppendToArgumentsList);
 
     OwningRewritePatternList patterns;
-    populateConvertLinalgOnTensorsToBuffersPatterns(&context, converter,
-                                                    patterns);
+    populateLinalgBufferizePatterns(&context, converter, patterns);
     populateWithBufferAssignmentOpConversionPatterns<
         mlir::ReturnOp, mlir::ReturnOp, linalg::CopyOp>(&context, converter,
                                                         patterns);
     if (failed(applyFullConversion(this->getOperation(), target, patterns)))
       this->signalPassFailure();
   }
 };
 } // end anonymous namespace
 
-std::unique_ptr<OperationPass<ModuleOp>>
-mlir::createConvertLinalgOnTensorsToBuffersPass() {
-  return std::make_unique<ConvertLinalgOnTensorsToBuffers>();
+std::unique_ptr<OperationPass<ModuleOp>> mlir::createLinalgBufferizePass() {
+  return std::make_unique<LinalgBufferizePass>();
 }
-void mlir::linalg::populateConvertLinalgOnTensorsToBuffersPatterns(
-
+void mlir::linalg::populateLinalgBufferizePatterns(
     MLIRContext *context, BufferAssignmentTypeConverter &converter,
     OwningRewritePatternList &patterns) {
   patterns.insert<
       // clang-format off
       LinalgOpConverter,
       TensorCastOpConverter,
       TensorConstantOpConverter
       // clang-format on
       >(context, converter);
 }

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

 add_mlir_dialect_library(MLIRLinalgTransforms
+  Bufferize.cpp
   DropUnitDims.cpp
   Fusion.cpp
   FusionOnTensors.cpp
   Hoisting.cpp
   Interchange.cpp
   Loops.cpp
   Promotion.cpp
-  TensorsToBuffers.cpp
   Tiling.cpp
   Transforms.cpp
   Vectorization.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Linalg
 
   DEPENDS

mlir/test/Dialect/Linalg/bufferize.mlir

-// RUN: mlir-opt -convert-linalg-on-tensors-to-buffers -buffer-placement -split-input-file %s | FileCheck %s
+// RUN: mlir-opt -linalg-bufferize -buffer-placement -split-input-file %s | FileCheck %s
 
 #map0 = affine_map<(d0) -> (d0)>
 
 // CHECK-LABEL: func @multiple_results
 func @multiple_results(%arg0: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>) {
     %0, %1 = linalg.generic {
       indexing_maps = [#map0, #map0, #map0],
       iterator_types = ["parallel"]