[mlir] Add patterns for vector.transfer_read/write to Linalg bufferization.

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<linalg::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, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const final {
+    if (readOp.getShapedType().isa<MemRefType>())
+      return failure();
+    vector::TransferReadOp::Adaptor adaptor(operands,
+                                            readOp->getAttrDictionary());
+    rewriter.replaceOpWithNewOp<vector::TransferReadOp>(
+        readOp, readOp.getType(), adaptor.source(), adaptor.indices(),
+        adaptor.permutation_map(), adaptor.padding(), adaptor.mask(),
+        adaptor.in_bounds());
+    return success();
+  }
+};
+
+class VectorTransferWriteOpConverter
+    : public OpConversionPattern<vector::TransferWriteOp> {
+public:
+  using OpConversionPattern<vector::TransferWriteOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(vector::TransferWriteOp writeOp, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const final {
+    if (writeOp.getShapedType().isa<MemRefType>())
+      return failure();
+    vector::TransferWriteOp::Adaptor adaptor(operands,
+                                             writeOp->getAttrDictionary());
+    rewriter.create<vector::TransferWriteOp>(
+        writeOp.getLoc(), adaptor.vector(), adaptor.source(), adaptor.indices(),
+        adaptor.permutation_map(),
+        adaptor.in_bounds() ? adaptor.in_bounds() : 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);
     BufferizeTypeConverter typeConverter;
 
     // Mark all Standard operations legal.
     target.addLegalDialect<AffineDialect, math::MathDialect,
                            memref::MemRefDialect, StandardOpsDialect,
                            tensor::TensorDialect>();
     target.addIllegalOp<InitTensorOp, tensor::ExtractSliceOp,
                         tensor::InsertSliceOp, PadTensorOp>();
 
     // 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<ConstantOp>(isLegalOperation);
+    target.addDynamicallyLegalOp<ConstantOp, vector::TransferReadOp,
+                                 vector::TransferWriteOp>(isLegalOperation);
 
     RewritePatternSet patterns(&context);
-    patterns.add<GeneralizePadTensorOpPattern>(patterns.getContext());
     populateLinalgBufferizePatterns(typeConverter, patterns);
     if (failed(applyPartialConversion(getOperation(), target,
                                       std::move(patterns))))
       signalPassFailure();
   }
 };
 } // end anonymous namespace
 
 std::unique_ptr<OperationPass<FuncOp>> mlir::createLinalgBufferizePass() {
   return std::make_unique<LinalgBufferizePass>();
 }
 
 void mlir::linalg::populateLinalgBufferizePatterns(
     BufferizeTypeConverter &typeConverter, RewritePatternSet &patterns) {
   // TODO: Drop this once tensor constants work in standard.
   // clang-format off
   patterns.add<
       BufferizeAnyLinalgOp,
       BufferizeFillOp,
       BufferizeInitTensorOp,
       BufferizeTensorReshapeOp<TensorExpandShapeOp>,
       BufferizeTensorReshapeOp<TensorCollapseShapeOp>,
       ExtractSliceOpConverter,
-      InsertSliceOpConverter
+      InsertSliceOpConverter,
+      VectorTransferReadOpConverter,
+      VectorTransferWriteOpConverter
     >(typeConverter, patterns.getContext());
   // clang-format on
+  patterns.add<GeneralizePadTensorOpPattern>(patterns.getContext());
 }

mlir/test/Dialect/Linalg/bufferize.mlir

 // CHECK:           %[[IN_MEMREF:.*]] = memref.buffer_cast %[[IN]] : memref<4x?x2x?xf32>
 // CHECK:           %[[OUT:.*]] = memref.alloc(%[[DIM1]], %[[OUT_DIM2]], %[[OUT_DIM3]]) : memref<4x?x?x?xf32>
 // CHECK:           linalg.copy(%[[FILLED]], %[[OUT]]) : memref<4x?x?x?xf32>, 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:           linalg.copy(%[[IN_MEMREF]], %[[INTERIOR]]) : memref<4x?x2x?xf32>, memref<4x?x2x?xf32, #map>
 // CHECK:           %[[OUT_TENSOR:.*]] = memref.tensor_load %[[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 = constant 0 : index
+  %cst = 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>
+}