[spirv] Fix legalize standard to spir-v for transfer ops

[ThomasRaoux]

Summary

Forward missing attributes when creating the new transfer op otherwise the builder would use default values.

[antiagainst]

LGTM.

BTW, it seems we put these vector ops in std to spv patterns.. Should we move them to vector to spv given we have that directory now?

[mravishankar]

I think this is more related to legalizing the subview rather than the vector operation.

[ThomasRaoux]

In D89907#2345595, @antiagainst wrote:

LGTM.

BTW, it seems we put these vector ops in std to spv patterns.. Should we move them to vector to spv given we have that directory now?

The downside is that we would not be able to share code as easily, there is very little code specific to transfer ops. As Mahesh mentioned this pass seems to deal more with subview ops than load/transfer ops

mlir/lib/Conversion/StandardToSPIRV/LegalizeStandardForSPIRV.cpp

rewriter.replaceOpWithNewOp<LoadOp>(loadOp, subViewOp.source(),
                                       sourceIndices);
 }
 
 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);
+      loadOp, loadOp.getVectorType(), subViewOp.source(), sourceIndices,
+      loadOp.permutation_map(), loadOp.padding(), loadOp.maskedAttr());
 }
 
 template <>
 void StoreOpOfSubViewFolder<StoreOp>::replaceOp(
     StoreOp storeOp, SubViewOp subViewOp, ArrayRef<Value> sourceIndices,
     PatternRewriter &rewriter) const {
   rewriter.replaceOpWithNewOp<StoreOp>(storeOp, storeOp.value(),
                                        subViewOp.source(), sourceIndices);
 }
 
 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);
+      sourceIndices, tranferWriteOp.permutation_map(),
+      tranferWriteOp.maskedAttr());
 }
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Utility functions for op legalization.
 //===----------------------------------------------------------------------===//
 
 /// Given the 'indices' of an load/store operation where the memref is a result

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) {
   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: [[F1:%.*]] = constant 1.000000e+00 : f32
   // 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
+  // CHECK: vector.transfer_read [[ARG0]]{{\[}}[[INDEX1]], [[INDEX2]]{{\]}}, [[F1]] {masked = [false]}
+  %f1 = constant 1.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>
+  %1 = vector.transfer_read %0[%arg3, %arg4], %f1 {masked = [false]} : 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]]{{\]}}
+  // CHECK: vector.transfer_write [[ARG5]], [[ARG0]]{{\[}}[[INDEX1]], [[INDEX2]]{{\]}} {masked = [false]}
   %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]>
+  vector.transfer_write %arg5, %0[%arg3, %arg4] {masked = [false]} : vector<4xf32>, memref<4x4xf32, offset:?, strides: [64, 3]>
   return
 }