[mlir][vector] Add memory effects to transfer_read transfer_write ops

[ThomasRaoux]

Summary

This allow more accurate modeling of the side effects and allow dead code elimination to remove dead transfer ops.

[aartbik]

do we need similar extensions for the other load/store flavors (masked load/store, gather/scatter, compress/expand)?
[if so, please let's wait until my CL is in to avoid conflicts, since I have been making some mods to those ops]

mlir/test/Conversion/VectorToSCF/vector-to-loops.mlir

25

Add a slightly more clarifying comment on why that is here, below too

mlir/test/Dialect/Vector/canonicalize.mlir

671

can you indent this line up with the other parameters?

[ThomasRaoux]

In D94318#2487343, @aartbik wrote:

do we need similar extensions for the other load/store flavors (masked load/store, gather/scatter, compress/expand)?
[if so, please let's wait until my CL is in to avoid conflicts, since I have been making some mods to those ops]

We do need to mark those as well, it looks like right now DCE wouldn't be able to remove a dead masked_load for instance. If those ops are not going to be expanded to tensor it can done in a slightly simpler way, we can just add MemRead<$base> to the properties directly in tablegen. I can definitely wait and update all of them once your patch lands.

[ThomasRaoux]

Address review comments.

mlir/include/mlir/Dialect/Vector/VectorOps.td

def Vector_ExtractStridedSliceOp :
   let hasCanonicalizer = 1;
   let hasFolder = 1;
   let assemblyFormat = "$vector attr-dict `:` type($vector) `to` type(results)";
 }
 
 def Vector_TransferReadOp :
   Vector_Op<"transfer_read", [
       DeclareOpInterfaceMethods<VectorTransferOpInterface>,
-      DeclareOpInterfaceMethods<VectorUnrollOpInterface, ["getShapeForUnroll"]>
+      DeclareOpInterfaceMethods<VectorUnrollOpInterface, ["getShapeForUnroll"]>,
+      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>
     ]>,
     Arguments<(ins AnyShaped:$source, Variadic<Index>:$indices,
                AffineMapAttr:$permutation_map, AnyType:$padding,
                OptionalAttr<BoolArrayAttr>:$masked)>,
     Results<(outs AnyVector:$vector)> {
 
   let summary = "Reads a supervector from memory into an SSA vector value.";
 
def Vector_TransferReadOp :
   ];
 
   let hasFolder = 1;
 }
 
 def Vector_TransferWriteOp :
   Vector_Op<"transfer_write", [
       DeclareOpInterfaceMethods<VectorTransferOpInterface>,
-      DeclareOpInterfaceMethods<VectorUnrollOpInterface, ["getShapeForUnroll"]>
+      DeclareOpInterfaceMethods<VectorUnrollOpInterface, ["getShapeForUnroll"]>,
+      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>
   ]>,
     Arguments<(ins AnyVector:$vector, AnyShaped:$source,
                Variadic<Index>:$indices,
                AffineMapAttr:$permutation_map,
                OptionalAttr<BoolArrayAttr>:$masked)>,
     Results<(outs Optional<AnyRankedTensor>:$result)> {
 
   let summary = "The vector.transfer_write op writes a supervector to memory.";

mlir/lib/Dialect/Vector/VectorOps.cpp

OpFoldResult TransferReadOp::fold(ArrayRef<Attribute>) {
   return OpFoldResult();
 }
 
 Optional<SmallVector<int64_t, 4>> TransferReadOp::getShapeForUnroll() {
   auto s = getVectorType().getShape();
   return SmallVector<int64_t, 4>{s.begin(), s.end()};
 }
 
+void TransferReadOp::getEffects(
+    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
+        &effects) {
+  if (getShapedType().isa<MemRefType>())
+    effects.emplace_back(MemoryEffects::Read::get(), source(),
+                         SideEffects::DefaultResource::get());
+}
+
 //===----------------------------------------------------------------------===//
 // TransferWriteOp
 //===----------------------------------------------------------------------===//
 
 /// Builder that sets permutation map to 'getMinorIdentityMap'.
 void TransferWriteOp::build(OpBuilder &builder, OperationState &result,
                             Value vector, Value source, ValueRange indices,
                             ArrayRef<bool> maybeMasked) {
if (succeeded(foldTransferMaskAttribute(*this)))
     return success();
   return foldMemRefCast(*this);
 }
 
 Optional<SmallVector<int64_t, 4>> TransferWriteOp::getShapeForUnroll() {
   return llvm::to_vector<4>(getVectorType().getShape());
 }
 
+void TransferWriteOp::getEffects(
+    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
+        &effects) {
+  if (getShapedType().isa<MemRefType>())
+    effects.emplace_back(MemoryEffects::Write::get(), source(),
+                         SideEffects::DefaultResource::get());
+}
+
 //===----------------------------------------------------------------------===//
 // MaskedLoadOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(MaskedLoadOp op) {
   VectorType maskVType = op.getMaskVectorType();
   VectorType passVType = op.getPassThruVectorType();
   VectorType resVType = op.getResultVectorType();

mlir/test/Conversion/VectorToSCF/vector-to-loops.mlir

-// RUN: mlir-opt %s -convert-vector-to-scf -split-input-file | FileCheck %s
-// RUN: mlir-opt %s -convert-vector-to-scf=full-unroll=true -split-input-file | FileCheck %s --check-prefix=FULL-UNROLL
+// RUN: mlir-opt %s -convert-vector-to-scf -split-input-file -allow-unregistered-dialect | FileCheck %s
+// RUN: mlir-opt %s -convert-vector-to-scf=full-unroll=true -split-input-file -allow-unregistered-dialect | FileCheck %s --check-prefix=FULL-UNROLL
 
 // CHECK-LABEL: func @materialize_read_1d() {
 func @materialize_read_1d() {
   %f0 = constant 0.0: f32
   %A = alloc () : memref<7x42xf32>
   affine.for %i0 = 0 to 7 step 4 {
     affine.for %i1 = 0 to 42 step 4 {
       %f1 = vector.transfer_read %A[%i0, %i1], %f0 {permutation_map = affine_map<(d0, d1) -> (d0)>} : memref<7x42xf32>, vector<4xf32>
       %ip1 = affine.apply affine_map<(d0) -> (d0 + 1)> (%i1)
       %f2 = vector.transfer_read %A[%i0, %ip1], %f0 {permutation_map = affine_map<(d0, d1) -> (d0)>} : memref<7x42xf32>, vector<4xf32>
       %ip2 = affine.apply affine_map<(d0) -> (d0 + 2)> (%i1)
       %f3 = vector.transfer_read %A[%i0, %ip2], %f0 {permutation_map = affine_map<(d0, d1) -> (d0)>} : memref<7x42xf32>, vector<4xf32>
       %ip3 = affine.apply affine_map<(d0) -> (d0 + 3)> (%i1)
       %f4 = vector.transfer_read %A[%i0, %ip3], %f0 {permutation_map = affine_map<(d0, d1) -> (d0)>} : memref<7x42xf32>, vector<4xf32>
       // Both accesses in the load must be clipped otherwise %i1 + 2 and %i1 + 3 will go out of bounds.
       // CHECK: scf.if
       // CHECK-NEXT: load
       // CHECK-NEXT: vector.insertelement
       // CHECK-NEXT: store
       // CHECK-NEXT: else
       // CHECK-NEXT: vector.insertelement
       // CHECK-NEXT: store
+      // Add a dummy use to prevent dead code elimination from removing transfer
+      // read ops.
+      "dummy_use"(%f1, %f2, %f3, %f4) : (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> ()
     }
   }
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @materialize_read_1d_partially_specialized
 func @materialize_read_1d_partially_specialized(%dyn1 : index, %dyn2 : index, %dyn4 : index) {
   %f0 = constant 0.0: f32
   %A = alloc (%dyn1, %dyn2, %dyn4) : memref<7x?x?x42x?xf32>
   affine.for %i0 = 0 to 7 {
     affine.for %i1 = 0 to %dyn1 {
       affine.for %i2 = 0 to %dyn2 {
         affine.for %i3 = 0 to 42 step 2 {
           affine.for %i4 = 0 to %dyn4 {
             %f1 = vector.transfer_read %A[%i0, %i1, %i2, %i3, %i4], %f0 {permutation_map = affine_map<(d0, d1, d2, d3, d4) -> (d3)>} : memref<7x?x?x42x?xf32>, vector<4xf32>
             %i3p1 = affine.apply affine_map<(d0) -> (d0 + 1)> (%i3)
             %f2 = vector.transfer_read %A[%i0, %i1, %i2, %i3p1, %i4], %f0 {permutation_map = affine_map<(d0, d1, d2, d3, d4) -> (d3)>} : memref<7x?x?x42x?xf32>, vector<4xf32>
+            // Add a dummy use to prevent dead code elimination from removing
+            // transfer read ops.
+            "dummy_use"(%f1, %f2) : (vector<4xf32>, vector<4xf32>) -> ()
           }
         }
       }
     }
   }
   // CHECK: %[[tensor:[0-9]+]] = alloc
   // CHECK-NOT: {{.*}} dim %[[tensor]], %c0
   // CHECK-NOT: {{.*}} dim %[[tensor]], %c3
func @materialize_read(%M: index, %N: index, %O: index, %P: index) {
   // CHECK-NEXT:                  store %[[RVEC]], %[[ALLOC]][%[[I6]], %[[I5]]] : memref<5x4xvector<3xf32>>
   // CHECK-NEXT:                } else {
   // CHECK-NEXT:                  %[[CVEC:.*]] = vector.insertelement
   // CHECK-NEXT:                  store %[[CVEC]], %[[ALLOC]][%[[I6]], %[[I5]]] : memref<5x4xvector<3xf32>>
   // CHECK-NEXT:                }
   // CHECK-NEXT:              }
   // CHECK-NEXT:            }
   // CHECK-NEXT:          }
+  // CHECK-NEXT:          %[[ALLOC_CAST:.*]] = vector.type_cast %[[ALLOC]] : memref<5x4xvector<3xf32>> to memref<vector<5x4x3xf32>>
+  // CHECK-NEXT:          %[[LD:.*]] = load %[[ALLOC_CAST]][] : memref<vector<5x4x3xf32>>
+  // CHECK-NEXT:          "dummy_use"(%[[LD]]) : (vector<5x4x3xf32>) -> ()
   // CHECK-NEXT:        }
   // CHECK-NEXT:      }
   // CHECK-NEXT:    }
   // CHECK-NEXT:  }
   // CHECK-NEXT:  return
   // CHECK-NEXT:}
 
   // Check that I0 + I4 (of size 3) read from first index load(L0, ...) and write into last index store(..., I4)
   // Check that I3 + I6 (of size 5) read from last index load(..., L3) and write into first index store(I6, ...)
   // Other dimensions are just accessed with I1, I2 resp.
   %A = alloc (%M, %N, %O, %P) : memref<?x?x?x?xf32, 0>
   affine.for %i0 = 0 to %M step 3 {
     affine.for %i1 = 0 to %N {
       affine.for %i2 = 0 to %O {
         affine.for %i3 = 0 to %P step 5 {
           %f = vector.transfer_read %A[%i0, %i1, %i2, %i3], %f0 {permutation_map = affine_map<(d0, d1, d2, d3) -> (d3, 0, d0)>} : memref<?x?x?x?xf32>, vector<5x4x3xf32>
+          // Add a dummy use to prevent dead code elimination from removing
+          // transfer read ops.
+          "dummy_use"(%f) : (vector<5x4x3xf32>) -> ()
         }
       }
     }
   }
   return
 }
 
 // -----

mlir/test/Dialect/Vector/canonicalize.mlir

 // CHECK-LABEL: broadcast_to_shapecast
 //       CHECK:   %[[C:.*]] = vector.shape_cast %{{.*}} : vector<4x4xf16> to vector<1x4x4xf16>
 //  CHECK-NEXT:   return %[[C]] : vector<1x4x4xf16>
 func @broadcast_to_shapecast(%arg0: vector<4x4xf16>) -> vector<1x4x4xf16> {
   %0 = vector.broadcast %arg0 : vector<4x4xf16> to vector<1x4x4xf16>
   return %0 : vector<1x4x4xf16>
 }
 
+// -----
+
+// CHECK-LABEL: func @dead_transfer_op
+//   CHECK-NOT:   vector.transfer_read
+//   CHECK-NOT:   vector.transfer_write
+//       CHECK:   return
+func @dead_transfer_op(%arg0 : tensor<4x4xf32>, %arg1 : memref<4x4xf32>,
+                       %v0 : vector<1x4xf32>) {
+  %c0 = constant 0 : index
+  %cf0 = constant 0.0 : f32
+  %r = vector.transfer_read %arg1[%c0, %c0], %cf0 :
+    memref<4x4xf32>, vector<1x4xf32>
+  %w = vector.transfer_write %v0, %arg0[%c0, %c0] :
+    vector<1x4xf32>, tensor<4x4xf32>
+  return
+}