[mlir][spirv] Legalize subviewop when used with vector transfer

[ThomasRaoux]

Summary

Subview operations are not natively supported downstream in the spirv path. This change allows removing subview when used by vector transfer the same way we already do it when they are used by LoadOp/StoreOp

[rriddle]

mlir/lib/Conversion/StandardToSPIRV/LegalizeStandardForSPIRV.cpp

37

Same here.

52

Use ArrayRef<Value> instead.

mlir/lib/Conversion/StandardToSPIRV/LegalizeStandardForSPIRV.cpp

 // dialect to handle ops that cannot be lowered directly.
 //
 //===----------------------------------------------------------------------===//
 
 #include "../PassDetail.h"
 #include "mlir/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.h"
 #include "mlir/Conversion/StandardToSPIRV/ConvertStandardToSPIRVPass.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/StandardTypes.h"
 
 using namespace mlir;
 
 namespace {
-/// Merges subview operation with load operation.
-class LoadOpOfSubViewFolder final : public OpRewritePattern<LoadOp> {
+/// Merges subview operation with load/transferRead operation.
+template <typename OpTy>
+class LoadOpOfSubViewFolder final : public OpRewritePattern<OpTy> {
 public:
-  using OpRewritePattern<LoadOp>::OpRewritePattern;
+  using OpRewritePattern<OpTy>::OpRewritePattern;
 
-  LogicalResult matchAndRewrite(LoadOp loadOp,
+  LogicalResult matchAndRewrite(OpTy loadOp,
                                 PatternRewriter &rewriter) const override;
+
+private:
+  void replaceOp(OpTy loadOp, SubViewOp subViewOp,
+                 ArrayRef<Value> sourceIndices,
+                 PatternRewriter &rewriter) const;
 };
 
-/// Merges subview operation with store operation.
-class StoreOpOfSubViewFolder final : public OpRewritePattern<StoreOp> {
+/// Merges subview operation with store/transferWriteOp operation.
+template <typename OpTy>
+class StoreOpOfSubViewFolder final : public OpRewritePattern<OpTy> {
 public:
-  using OpRewritePattern<StoreOp>::OpRewritePattern;
+  using OpRewritePattern<OpTy>::OpRewritePattern;
 
-  LogicalResult matchAndRewrite(StoreOp storeOp,
+  LogicalResult matchAndRewrite(OpTy storeOp,
                                 PatternRewriter &rewriter) const override;
+
+private:
+  void replaceOp(OpTy StoreOp, SubViewOp subViewOp,
+                 ArrayRef<Value> sourceIndices,
+                 PatternRewriter &rewriter) const;
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Utility functions for op legalization.
 //===----------------------------------------------------------------------===//
 
 /// Given the 'indices' of an load/store operation where the memref is a result
for (auto index : llvm::enumerate(indices)) {
     auto mul = rewriter.create<MulIOp>(loc, index.value(), stride);
     sourceIndices[index.index()] =
         rewriter.create<AddIOp>(loc, offset, mul).getResult();
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
-// Folding SubViewOp and LoadOp.
+// Folding SubViewOp and LoadOp/TransferReadOp.
 //===----------------------------------------------------------------------===//
 
+template <typename OpTy>
 LogicalResult
-LoadOpOfSubViewFolder::matchAndRewrite(LoadOp loadOp,
+LoadOpOfSubViewFolder<OpTy>::matchAndRewrite(OpTy loadOp,
                                              PatternRewriter &rewriter) const {
-  auto subViewOp = loadOp.memref().getDefiningOp<SubViewOp>();
+  auto subViewOp = loadOp.memref().template getDefiningOp<SubViewOp>();
   if (!subViewOp) {
     return failure();
   }
   SmallVector<Value, 4> sourceIndices;
   if (failed(resolveSourceIndices(loadOp.getLoc(), rewriter, subViewOp,
                                   loadOp.indices(), sourceIndices)))
     return failure();
 
+  replaceOp(loadOp, subViewOp, sourceIndices, rewriter);
+  return success();
+}
+
+template <>
+void LoadOpOfSubViewFolder<LoadOp>::replaceOp(LoadOp loadOp,
+                                              SubViewOp subViewOp,
+                                              ArrayRef<Value> sourceIndices,
+                                              PatternRewriter &rewriter) const {
   rewriter.replaceOpWithNewOp<LoadOp>(loadOp, subViewOp.source(),
                                       sourceIndices);
-  return success();
+}
+
+template <>
+void LoadOpOfSubViewFolder<vector::TransferReadOp>::replaceOp(
+    vector::TransferReadOp loadOp, SubViewOp subViewOp,
+    ArrayRef<Value> sourceIndices, PatternRewriter &rewriter) const {
+  rewriter.replaceOpWithNewOp<vector::TransferReadOp>(
+      loadOp, loadOp.getVectorType(), subViewOp.source(), sourceIndices);
 }
 
 //===----------------------------------------------------------------------===//
-// Folding SubViewOp and StoreOp.
+// Folding SubViewOp and StoreOp/TransferWriteOp.
 //===----------------------------------------------------------------------===//
 
+template <typename OpTy>
 LogicalResult
-StoreOpOfSubViewFolder::matchAndRewrite(StoreOp storeOp,
+StoreOpOfSubViewFolder<OpTy>::matchAndRewrite(OpTy storeOp,
                                               PatternRewriter &rewriter) const {
-  auto subViewOp = storeOp.memref().getDefiningOp<SubViewOp>();
+  auto subViewOp = storeOp.memref().template getDefiningOp<SubViewOp>();
   if (!subViewOp) {
     return failure();
   }
   SmallVector<Value, 4> sourceIndices;
   if (failed(resolveSourceIndices(storeOp.getLoc(), rewriter, subViewOp,
                                   storeOp.indices(), sourceIndices)))
     return failure();
 
+  replaceOp(storeOp, subViewOp, sourceIndices, rewriter);
+  return success();
+}
+
+template <>
+void StoreOpOfSubViewFolder<StoreOp>::replaceOp(
+    StoreOp storeOp, SubViewOp subViewOp, ArrayRef<Value> sourceIndices,
+    PatternRewriter &rewriter) const {
   rewriter.replaceOpWithNewOp<StoreOp>(storeOp, storeOp.value(),
                                        subViewOp.source(), sourceIndices);
-  return success();
+}
+
+template <>
+void StoreOpOfSubViewFolder<vector::TransferWriteOp>::replaceOp(
+    vector::TransferWriteOp tranferWriteOp, SubViewOp subViewOp,
+    ArrayRef<Value> sourceIndices, PatternRewriter &rewriter) const {
+  rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
+      tranferWriteOp, tranferWriteOp.vector(), subViewOp.source(),
+      sourceIndices);
 }
 
 //===----------------------------------------------------------------------===//
 // Hook for adding patterns.
 //===----------------------------------------------------------------------===//
 
 void mlir::populateStdLegalizationPatternsForSPIRVLowering(
     MLIRContext *context, OwningRewritePatternList &patterns) {
-  patterns.insert<LoadOpOfSubViewFolder, StoreOpOfSubViewFolder>(context);
+  patterns.insert<LoadOpOfSubViewFolder<LoadOp>,
+                  LoadOpOfSubViewFolder<vector::TransferReadOp>,
+                  StoreOpOfSubViewFolder<StoreOp>,
+                  StoreOpOfSubViewFolder<vector::TransferWriteOp>>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // Pass for testing just the legalization patterns.
 //===----------------------------------------------------------------------===//
 
 namespace {
 struct SPIRVLegalization final

mlir/test/Conversion/StandardToSPIRV/legalization.mlir

func @fold_dynamic_stride_subview_with_store(%arg0 : memref<12x32xf32>, %arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index, %arg5 : index, %arg6 : index, %arg7 : f32) {
   // CHECK: [[STRIDE2:%.*]] = muli [[ARG4]], [[ARG6]] : index
   // CHECK: [[INDEX2:%.*]] = addi [[ARG2]], [[STRIDE2]] : index
   // CHECK: store [[ARG7]], [[ARG0]]{{\[}}[[INDEX1]], [[INDEX2]]{{\]}}
   %0 = subview %arg0[%arg1, %arg2][4, 4][%arg5, %arg6] :
     memref<12x32xf32> to memref<4x4xf32, offset:?, strides: [?, ?]>
   store %arg7, %0[%arg3, %arg4] : memref<4x4xf32, offset:?, strides: [?, ?]>
   return
 }
+
+// CHECK-LABEL: @fold_static_stride_subview_with_transfer_read
+// CHECK-SAME: [[ARG0:%.*]]: memref<12x32xf32>, [[ARG1:%.*]]: index, [[ARG2:%.*]]: index, [[ARG3:%.*]]: index, [[ARG4:%.*]]: index
+func @fold_static_stride_subview_with_transfer_read(%arg0 : memref<12x32xf32>, %arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index) -> vector<4xf32> {
+  // CHECK-NOT: subview
+  // CHECK: [[C2:%.*]] = constant 2 : index
+  // CHECK: [[C3:%.*]] = constant 3 : index
+  // CHECK: [[STRIDE1:%.*]] = muli [[ARG3]], [[C2]] : index
+  // CHECK: [[INDEX1:%.*]] = addi [[ARG1]], [[STRIDE1]] : index
+  // CHECK: [[STRIDE2:%.*]] = muli [[ARG4]], [[C3]] : index
+  // CHECK: [[INDEX2:%.*]] = addi [[ARG2]], [[STRIDE2]] : index
+  // CHECK: vector.transfer_read [[ARG0]]{{\[}}[[INDEX1]], [[INDEX2]]{{\]}}
+  %f0 = constant 0.0 : f32
+  %0 = subview %arg0[%arg1, %arg2][4, 4][2, 3] : memref<12x32xf32> to memref<4x4xf32, offset:?, strides: [64, 3]>
+  %1 = vector.transfer_read %0[%arg3, %arg4], %f0 : memref<4x4xf32, offset:?, strides: [64, 3]>, vector<4xf32>
+  return %1 : vector<4xf32>
+}
+
+// CHECK-LABEL: @fold_static_stride_subview_with_transfer_write
+// CHECK-SAME: [[ARG0:%.*]]: memref<12x32xf32>, [[ARG1:%.*]]: index, [[ARG2:%.*]]: index, [[ARG3:%.*]]: index, [[ARG4:%.*]]: index, [[ARG5:%.*]]: vector<4xf32>
+func @fold_static_stride_subview_with_transfer_write(%arg0 : memref<12x32xf32>, %arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index, %arg5 : vector<4xf32>) {
+  // CHECK-NOT: subview
+  // CHECK: [[C2:%.*]] = constant 2 : index
+  // CHECK: [[C3:%.*]] = constant 3 : index
+  // CHECK: [[STRIDE1:%.*]] = muli [[ARG3]], [[C2]] : index
+  // CHECK: [[INDEX1:%.*]] = addi [[ARG1]], [[STRIDE1]] : index
+  // CHECK: [[STRIDE2:%.*]] = muli [[ARG4]], [[C3]] : index
+  // CHECK: [[INDEX2:%.*]] = addi [[ARG2]], [[STRIDE2]] : index
+  // CHECK: vector.transfer_write [[ARG5]], [[ARG0]]{{\[}}[[INDEX1]], [[INDEX2]]{{\]}}
+  %0 = subview %arg0[%arg1, %arg2][4, 4][2, 3] :
+    memref<12x32xf32> to memref<4x4xf32, offset:?, strides: [64, 3]>
+  vector.transfer_write %arg5, %0[%arg3, %arg4] : vector<4xf32>, memref<4x4xf32, offset:?, strides: [64, 3]>
+  return
+}