[mlir][bufferize][NFC] Add vector-bufferize pass and remove obsolete patterns

[springerm]

Summary

This is towards removing the remaining dialect conversion patterns from Linalg/Transforms/Bufferize.cpp. Some of these patterns are for vector ops, which should have been in the vector dialect in the first place.

This commit removes the bufferization patterns for transfer_read and transfer_write. In existing code, vector-bufferize is used, which bufferizes those ops using BufferizableOpInterface.

[aartbik]

mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco_utils.py

160

apologies, another revision just crossed, changing this to the sparse-compiler pipeline
(so less changes required ;-), but you will have to rebase

[springerm]

rebase

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

-# This dialect does currently not have any passes.
+set(LLVM_TARGET_DEFINITIONS Passes.td)
+mlir_tablegen(Passes.h.inc -gen-pass-decls -name Vector)
+add_public_tablegen_target(MLIRVectorTransformsIncGen)
+
+add_mlir_doc(Passes VectorPasses ./ -gen-pass-doc)

mlir/include/mlir/Dialect/Vector/Transforms/Passes.h

+//===- Passes.h - Pass Entrypoints ------------------------------*- 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_VECTOR_TRANSFORMS_PASSES_H_
+#define MLIR_DIALECT_VECTOR_TRANSFORMS_PASSES_H_
+
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+namespace vector {
+/// Creates an instance of the `vector` dialect bufferization pass.
+std::unique_ptr<Pass> createVectorBufferizePass();
+
+//===----------------------------------------------------------------------===//
+// Registration
+//===----------------------------------------------------------------------===//
+
+/// Generate the code for registering passes.
+#define GEN_PASS_REGISTRATION
+#include "mlir/Dialect/Vector/Transforms/Passes.h.inc"
+} // namespace vector
+
+} // namespace mlir
+
+#endif // MLIR_DIALECT_VECTOR_TRANSFORMS_PASSES_H_

mlir/include/mlir/Dialect/Vector/Transforms/Passes.td

+//===-- Passes.td - Vector pass definition file ------------*- 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_VECTOR_TRANSFORMS_PASSES
+#define MLIR_DIALECT_VECTOR_TRANSFORMS_PASSES
+
+include "mlir/Pass/PassBase.td"
+
+def VectorBufferize : Pass<"vector-bufferize", "FuncOp"> {
+  let summary = "Bufferize Vector dialect ops";
+  let constructor = "mlir::vector::createVectorBufferizePass()";
+}
+
+#endif // MLIR_DIALECT_VECTOR_TRANSFORMS_PASSES

mlir/include/mlir/InitAllPasses.h

 #include "mlir/Dialect/SCF/Passes.h"
 #include "mlir/Dialect/SPIRV/Transforms/Passes.h"
 #include "mlir/Dialect/Shape/Transforms/Passes.h"
 #include "mlir/Dialect/SparseTensor/Pipelines/Passes.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Dialect/StandardOps/Transforms/Passes.h"
 #include "mlir/Dialect/Tensor/Transforms/Passes.h"
 #include "mlir/Dialect/Tosa/Transforms/Passes.h"
+#include "mlir/Dialect/Vector/Transforms/Passes.h"
 #include "mlir/Transforms/Passes.h"
 
 #include <cstdlib>
 
 namespace mlir {
 
 // This function may be called to register the MLIR passes with the
 // global registry.
inline void registerAllPasses() {
   memref::registerMemRefPasses();
   quant::registerQuantPasses();
   registerSCFPasses();
   registerShapePasses();
   spirv::registerSPIRVPasses();
   registerStandardPasses();
   tensor::registerTensorPasses();
   tosa::registerTosaOptPasses();
+  vector::registerVectorPasses();
 
   // Dialect pipelines
   sparse_tensor::registerSparseTensorPipelines();
 }
 
 } // namespace mlir
 
 #endif // MLIR_INITALLPASSES_H_

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

Value subview = rewriter.create<memref::SubViewOp>(
         op.getLoc(), destMemRef, op.getMixedOffsets(), op.getMixedSizes(),
         op.getMixedStrides());
     // Copy the small memref.
     rewriter.create<memref::CopyOp>(op.getLoc(), sourceMemRef, subview);
     rewriter.replaceOp(op, destMemRef);
     return success();
   }
 };
-
-class VectorTransferReadOpConverter
-    : public OpConversionPattern<vector::TransferReadOp> {
-public:
-  using OpConversionPattern<vector::TransferReadOp>::OpConversionPattern;
-
-  LogicalResult
-  matchAndRewrite(vector::TransferReadOp readOp, OpAdaptor adaptor,
-                  ConversionPatternRewriter &rewriter) const final {
-    if (readOp.getShapedType().isa<MemRefType>())
-      return failure();
-    rewriter.replaceOpWithNewOp<vector::TransferReadOp>(
-        readOp, readOp.getType(), adaptor.source(), adaptor.indices(),
-        adaptor.permutation_mapAttr(), adaptor.padding(), adaptor.mask(),
-        adaptor.in_boundsAttr());
-    return success();
-  }
-};
-
-class VectorTransferWriteOpConverter
-    : public OpConversionPattern<vector::TransferWriteOp> {
-public:
-  using OpConversionPattern<vector::TransferWriteOp>::OpConversionPattern;
-
-  LogicalResult
-  matchAndRewrite(vector::TransferWriteOp writeOp, OpAdaptor adaptor,
-                  ConversionPatternRewriter &rewriter) const final {
-    if (writeOp.getShapedType().isa<MemRefType>())
-      return failure();
-    rewriter.create<vector::TransferWriteOp>(
-        writeOp.getLoc(), adaptor.vector(), adaptor.source(), adaptor.indices(),
-        adaptor.permutation_mapAttr(),
-        adaptor.in_bounds() ? adaptor.in_boundsAttr() : ArrayAttr());
-    rewriter.replaceOp(writeOp, adaptor.source());
-    return success();
-  }
-};
 } // namespace
 
 namespace {
 /// Converts Linalg operations that work on tensor-type operands or results to
 /// work on buffers.
 struct LinalgBufferizePass : public LinalgBufferizeBase<LinalgBufferizePass> {
   void runOnOperation() override {
     MLIRContext &context = getContext();
     ConversionTarget target(context);
     bufferization::BufferizeTypeConverter typeConverter;
 
     // Mark all Standard operations legal.
     target.addLegalDialect<arith::ArithmeticDialect, AffineDialect,
                            memref::MemRefDialect, StandardOpsDialect,
                            tensor::TensorDialect>();
     target.addIllegalOp<InitTensorOp, tensor::PadOp, tensor::CollapseShapeOp,
                         tensor::ExpandShapeOp, tensor::ExtractSliceOp,
                         tensor::InsertSliceOp>();
 
     // Mark all Linalg operations illegal as long as they work on tensors.
     auto isLegalOperation = [&](Operation *op) {
       return typeConverter.isLegal(op);
     };
     target.addDynamicallyLegalDialect<linalg::LinalgDialect>(isLegalOperation);
-    target
-        .addDynamicallyLegalOp<vector::TransferReadOp, vector::TransferWriteOp>(
-            isLegalOperation);
 
     RewritePatternSet patterns(&context);
     populateLinalgBufferizePatterns(typeConverter, patterns);
     if (failed(applyPartialConversion(getOperation(), target,
                                       std::move(patterns))))
       signalPassFailure();
   }
 };
void mlir::linalg::populateLinalgBufferizePatterns(
   // clang-format off
   patterns.add<
       BufferizeAnyLinalgOp,
       BufferizeFillOp,
       BufferizeInitTensorOp,
       BufferizeTensorReshapeOp<tensor::ExpandShapeOp>,
       BufferizeTensorReshapeOp<tensor::CollapseShapeOp>,
       ExtractSliceOpConverter,
-      InsertSliceOpConverter,
-      VectorTransferReadOpConverter,
-      VectorTransferWriteOpConverter
+      InsertSliceOpConverter
     >(typeConverter, patterns.getContext());
   // clang-format on
   patterns.add<GeneralizePadOpPattern>(patterns.getContext());
 }

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

add_mlir_dialect_library(MLIRSparseTensorPipelines
   MLIRPass
   MLIRReconcileUnrealizedCasts
   MLIRSCFToControlFlow
   MLIRSparseTensor
   MLIRSparseTensorTransforms
   MLIRStandardOpsTransforms
   MLIRTensorTransforms
   MLIRVectorToLLVM
+  MLIRVectorTransforms
 )

mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp

 #include "mlir/Conversion/Passes.h"
 #include "mlir/Dialect/Arithmetic/Transforms/Passes.h"
 #include "mlir/Dialect/Bufferization/Transforms/Passes.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Dialect/StandardOps/Transforms/Passes.h"
 #include "mlir/Dialect/Tensor/Transforms/Passes.h"
+#include "mlir/Dialect/Vector/Transforms/Passes.h"
 #include "mlir/Pass/PassManager.h"
 
 using namespace mlir;
 using namespace mlir::sparse_tensor;
 
 //===----------------------------------------------------------------------===//
 // Pipeline implementation.
 //===----------------------------------------------------------------------===//
 
 void mlir::sparse_tensor::buildSparseCompiler(
     OpPassManager &pm, const SparseCompilerOptions &options) {
   // TODO(wrengr): ensure the original `pm` is for ModuleOp
   pm.addPass(createSparsificationPass(options.sparsificationOptions()));
   pm.addPass(createSparseTensorConversionPass());
   pm.addNestedPass<FuncOp>(createLinalgBufferizePass());
+  pm.addNestedPass<FuncOp>(vector::createVectorBufferizePass());
   pm.addNestedPass<FuncOp>(createConvertLinalgToLoopsPass());
   pm.addNestedPass<FuncOp>(createConvertVectorToSCFPass());
   pm.addNestedPass<FuncOp>(createConvertSCFToCFPass());
   pm.addPass(createFuncBufferizePass());
   pm.addPass(arith::createConstantBufferizePass());
   pm.addNestedPass<FuncOp>(createTensorBufferizePass());
   pm.addNestedPass<FuncOp>(
       mlir::bufferization::createFinalizingBufferizePass());

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

+//===- Bufferize.cpp - Bufferization for `vector` dialect 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 bufferization of `vector` dialect ops
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"
+#include "PassDetail.h"
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.h"
+#include "mlir/Dialect/Vector/Transforms/Passes.h"
+
+using namespace mlir;
+using namespace bufferization;
+
+namespace {
+struct VectorBufferizePass : public VectorBufferizeBase<VectorBufferizePass> {
+  void runOnOperation() override {
+    BufferizationOptions options = getPartialBufferizationOptions();
+    options.addToDialectFilter<vector::VectorDialect>();
+
+    if (failed(bufferizeOp(getOperation(), options)))
+      signalPassFailure();
+  }
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<bufferization::BufferizationDialect, memref::MemRefDialect,
+                    tensor::TensorDialect, vector::VectorDialect>();
+    vector::registerBufferizableOpInterfaceExternalModels(registry);
+  }
+};
+} // namespace
+
+std::unique_ptr<Pass> mlir::vector::createVectorBufferizePass() {
+  return std::make_unique<VectorBufferizePass>();
+}

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

 add_mlir_dialect_library(MLIRVectorTransforms
   BufferizableOpInterfaceImpl.cpp
+  Bufferize.cpp
   VectorDropLeadUnitDim.cpp
   VectorInsertExtractStridedSliceRewritePatterns.cpp
   VectorMultiDimReductionTransforms.cpp
   VectorTransferOpTransforms.cpp
   VectorTransferSplitRewritePatterns.cpp
   VectorTransferPermutationMapRewritePatterns.cpp
   VectorTransforms.cpp
   VectorUnrollDistribute.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Vector/Transforms
 
   LINK_LIBS PUBLIC
   MLIRAffine
   MLIRAffineAnalysis
   MLIRAffineUtils
   MLIRArithmetic
   MLIRBufferization
+  MLIRBufferizationTransforms
   MLIRDialectUtils
   MLIRIR
   MLIRLinalg
   MLIRMemRef
   MLIRSCF
+  MLIRTransforms
   MLIRVector
   MLIRVectorInterfaces
   MLIRVectorUtils
   )

mlir/lib/Dialect/Vector/Transforms/PassDetail.h

+//===- PassDetail.h - Vector Pass class details -----------------*- 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 DIALECT_VECTOR_TRANSFORMS_PASSDETAIL_H_
+#define DIALECT_VECTOR_TRANSFORMS_PASSDETAIL_H_
+
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+
+namespace bufferization {
+class BufferizationDialect;
+} // namespace bufferization
+
+namespace memref {
+class MemRefDialect;
+} // namespace memref
+
+#define GEN_PASS_CLASSES
+#include "mlir/Dialect/Vector/Transforms/Passes.h.inc"
+
+} // namespace mlir
+
+#endif // DIALECT_VECTOR_TRANSFORMS_PASSDETAIL_H_

mlir/test/Dialect/Linalg/bufferize.mlir

 // CHECK:           %[[OUT:.*]] = memref.alloc(%[[DIM1]], %[[OUT_DIM2]], %[[OUT_DIM3]]) : memref<4x?x?x?xf32>
 // CHECK:           memref.copy %[[FILLED]], %[[OUT]] : memref<4x?x?x?xf32> to memref<4x?x?x?xf32>
 // CHECK:           %[[INTERIOR:.*]] = memref.subview %[[OUT]][0, 0, %[[OFFSET]], 0] [4, %[[DIM1]], 2, %[[DIM3]]] [1, 1, 1, 1] : memref<4x?x?x?xf32> to memref<4x?x2x?xf32, #map>
 // CHECK:           memref.copy %[[IN_MEMREF]], %[[INTERIOR]] : memref<4x?x2x?xf32> to memref<4x?x2x?xf32, #map>
 // CHECK:           %[[OUT_TENSOR:.*]] = bufferization.to_tensor %[[OUT]] : memref<4x?x?x?xf32>
 // CHECK:           return %[[OUT_TENSOR]] : tensor<4x?x?x?xf32>
 // CHECK:         }
 
-
-// -----
-
-// CHECK-LABEL:   func @vector_transfer
-func @vector_transfer(%in: tensor<4xf32>, %out: tensor<4xf32>) {
-  %c0 = arith.constant 0 : index
-  %cst = arith.constant 0.000000e+00 : f32
-  %read = vector.transfer_read %in[%c0], %cst {in_bounds = [true]}
-      : tensor<4xf32>, vector<4xf32>
-  %tanh = math.tanh %read : vector<4xf32>
-  %write = vector.transfer_write %tanh, %out[%c0] {in_bounds = [true]}
-      : vector<4xf32>, tensor<4xf32>
-  return
-  // CHECK: vector.transfer_read {{.*}} : memref<4xf32>, vector<4xf32>
-  // CHECK: vector.transfer_write {{.*}} : vector<4xf32>, memref<4xf32>
-}
-
 // -----
 
 // CHECK-LABEL:   func @bufferize_dot
 func @bufferize_dot(%in: tensor<4xf32>, %out: tensor<f32>) -> tensor<f32> {
   %dot = linalg.dot ins(%in, %in : tensor<4xf32>, tensor<4xf32>)
                           outs(%out : tensor<f32>) -> tensor<f32>
   return %dot : tensor<f32>
   // CHECK: linalg.dot ins(%{{.*}}, %{{.*}} : memref<4xf32>, memref<4xf32>)
   // CHECK-SAME:       outs(%[[OUT:.*]] : memref<f32>)
   // CHECK: %[[OUT_TENSOR:.*]] = bufferization.to_tensor %[[OUT]] : memref<f32>
   // CHECK: return %[[OUT_TENSOR]]
 }

mlir/test/Dialect/Vector/bufferize.mlir

+// RUN: mlir-opt %s -vector-bufferize -split-input-file | FileCheck %s
+
+// CHECK-LABEL: func @transfer_read(
+//  CHECK-SAME:     %[[t:.*]]: tensor<?x?xf32>, %[[o1:.*]]: index, %[[o2:.*]]: index, %[[pad:.*]]: f32)
+//       CHECK:   %[[m:.*]] = bufferization.to_memref %[[t]] : memref<?x?xf32>
+//       CHECK:   %[[r:.*]] = vector.transfer_read %[[m]][%[[o1]], %[[o2]]], %[[pad]] {in_bounds = [true, false]} : memref<?x?xf32>, vector<5x6xf32>
+//       CHECK:   return %[[r]]
+func @transfer_read(%t: tensor<?x?xf32>, %o1: index,
+                    %o2: index, %pad: f32) -> vector<5x6xf32> {
+  %0 = vector.transfer_read %t[%o1, %o2], %pad {in_bounds = [true, false]}
+      : tensor<?x?xf32>, vector<5x6xf32>
+  return %0 : vector<5x6xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @transfer_write(
+//  CHECK-SAME:     %[[t:.*]]: tensor<?x?xf32>, %[[o1:.*]]: index, %[[o2:.*]]: index, %[[vec:.*]]: vector<5x6xf32>)
+//       CHECK:   %[[m:.*]] = bufferization.to_memref %[[t]] : memref<?x?xf32>
+//       CHECK:   %[[alloc:.*]] = memref.alloc(%{{.*}}, %{{.*}}) {{.*}} : memref<?x?xf32>
+//       CHECK:   memref.copy %[[m]], %[[alloc]]
+//       CHECK:   vector.transfer_write %[[vec]], %[[alloc]][%[[o1]], %[[o2]]] {in_bounds = [true, false]} : vector<5x6xf32>, memref<?x?xf32>
+//       CHECK:   %[[r:.*]] = bufferization.to_tensor %[[alloc]] : memref<?x?xf32>
+//       CHECK:   return %[[r]]
+func @transfer_write(%t: tensor<?x?xf32>, %o1: index,
+                     %o2: index, %vec: vector<5x6xf32>) -> tensor<?x?xf32> {
+  %0 = vector.transfer_write %vec, %t[%o1, %o2] {in_bounds = [true, false]}
+      : vector<5x6xf32>, tensor<?x?xf32>
+  return %0 : tensor<?x?xf32>
+}

mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco_utils.py

def compile_and_build_engine(
 
   Returns:
     A JIT execution engine for the MLIR module.
 
   """
   pipeline = (
       f"sparsification,"
       f"sparse-tensor-conversion,"
-      f"builtin.func(linalg-bufferize,convert-linalg-to-loops,convert-vector-to-scf),"
+      f"builtin.func(vector-bufferize,linalg-bufferize,convert-linalg-to-loops,convert-vector-to-scf),"
       f"convert-scf-to-cf,"
       f"func-bufferize,"
       f"arith-bufferize,"
       f"builtin.func(tensor-bufferize,finalizing-bufferize),"
       f"convert-vector-to-llvm{{reassociate-fp-reductions=1 enable-index-optimizations=1}},"
       f"lower-affine,"
       f"convert-memref-to-llvm,"
       f"convert-std-to-llvm,"

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

deps = [
         ":ConversionPasses",
         ":LinalgTransforms",
         ":Pass",
         ":SparseTensor",
         ":SparseTensorTransforms",
         ":StandardOpsTransforms",
         ":TensorTransforms",
         ":VectorToLLVM",
+        ":VectorTransforms",
     ],
 )
 
 td_library(
     name = "StdOpsTdFiles",
     srcs = [
         "include/mlir/Dialect/StandardOps/IR/Ops.td",
     ],
deps = [
         ":TensorDialect",
         ":VectorInterfaces",
         ":VectorOpsIncGen",
         ":ViewLikeInterface",
         "//llvm:Support",
     ],
 )
 
+gentbl_cc_library(
+    name = "VectorPassIncGen",
+    strip_include_prefix = "include",
+    tbl_outs = [
+        (
+            [
+                "-gen-pass-decls",
+                "-name=Vector",
+            ],
+            "include/mlir/Dialect/Vector/Transforms/Passes.h.inc",
+        ),
+    ],
+    tblgen = ":mlir-tblgen",
+    td_file = "include/mlir/Dialect/Vector/Transforms/Passes.td",
+    deps = [":PassBaseTdFiles"],
+)
+
 cc_library(
     name = "VectorTransforms",
     srcs = glob(
         [
             "lib/Dialect/Vector/Transforms/*.cpp",
+            "lib/Dialect/Vector/Transforms/*.h",
         ],
     ),
     hdrs = glob([
         "include/mlir/Dialect/Vector/Transforms/*.h",
     ]),
     includes = ["include"],
     deps = [
         ":Affine",
         ":AffineAnalysis",
         ":Analysis",
         ":ArithmeticDialect",
         ":BufferizationDialect",
+        ":BufferizationTransforms",
         ":DialectUtils",
         ":IR",
         ":LinalgOps",
         ":MemRefDialect",
+        ":Pass",
         ":SCFDialect",
         ":StandardOps",
         ":Support",
         ":TensorDialect",
+        ":Transforms",
         ":VectorInterfaces",
         ":VectorOps",
+        ":VectorPassIncGen",
         ":VectorUtils",
         "//llvm:Support",
     ],
 )
 
 cc_library(
     name = "VectorUtils",
     srcs = glob(
deps = [
         ":TosaToLinalg",
         ":Transforms",
         ":TransformsPassIncGen",
         ":VectorOps",
         ":VectorToLLVM",
         ":VectorToROCDL",
         ":VectorToSCF",
         ":VectorToSPIRV",
+        ":VectorTransforms",
         ":X86Vector",
         ":X86VectorTransforms",
     ],
 )
 
 cc_binary(
     name = "mlir-lsp-server",
     srcs = ["tools/mlir-lsp-server/mlir-lsp-server.cpp"],