[mlir][linalg] Improve aliasing approximation for hoisting transfer read/write.

[ThomasRaoux]

Summary

Improve the logic deciding if it is safe to hoist vector transfer read/write out of the loop. Change the logic to prevent hoisting operations if there are any unknown access to the memref in the loop no matter where the operation is. For other transfer read/write in the loop check if we can prove that they access disjoint memory and ignore them in this case.

[nicolasvasilache]

Made a quick pass but I do not think the algorithm is complete atm: in order to check lack of aliasing you need to look at whether the static intervals between read/write pairs overlap.
Here the positive example only has trivial intervals (i.e. that reduce to a single point).

The testing itself would need at least these cases:
fail:

vector.transfer_write %u00, %memref1[%c0, %c0] : vector<2xf32>, memref<?x?xf32>
vector.transfer_write %u01, %memref1[%c0 %c1] : vector<2xf32>, memref<?x?xf32>

success:

vector.transfer_write %u00, %memref1[%c0, %c0] : vector<2xf32>, memref<?x?xf32>
vector.transfer_write %u01, %memref1[%c0 %c2] : vector<2xf32>, memref<?x?xf32>

You also want to distinguish between the dimensions that require interval analysis (i.e. the ones that become new vectors) and the ones that are just indexed (and that only need to be provably different across writes).

[ThomasRaoux]

Thanks Nicolas. I addressed the review, please take another look to confirm that I got it right. Based on what you mentioned my understanding is that if any indices are different in the first ranks (from 0 to memrefRank-vectorRank) then the slices are disjoint for the rest of the ranks I check that the distance between the offset is smaller that the vector size (since I check that the vector types must match)

[nicolasvasilache]

This looks good to me, there are more relaxations we can do in the future but we should be good for now.
Thanks, Thomas!

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

94

Note memref may itself have a vector element type and we need to account for that.
I added xferOp.getLeadingMemRefRank() recently to more easily capture the dims that are just indexed.

112

Great, thanks!

171

typos: transfer_read/transfer_write ?

173

Let's also check transferRead.getVectorType() == transferWrite.getVectorType(), I missed that in the original impl.

mlir/test/Dialect/Linalg/hoisting.mlir

179

// CHECK-SAME: (vector<3xf32>, vector<3xf32>, vector<4xf32>, vector<4xf32>) {

here and below?

[ThomasRaoux]

Address review comments

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

   }
 }
 
+/// Return true if we can prove that the transfer operations access dijoint
+/// memory.
+template <typename TransferTypeA, typename TransferTypeB>
+static bool isDisjoint(TransferTypeA transferA, TransferTypeB transferB) {
+  if (transferA.memref() != transferB.memref())
+    return false;
+  for (unsigned i = 0, e = std::min(transferA.indices().size(),
+                                    transferB.indices().size());
+       i < e; i++) {
+    auto indexA = transferA.indices()[i].template getDefiningOp<ConstantOp>();
+    auto indexB = transferB.indices()[i].template getDefiningOp<ConstantOp>();
+    // If any of the indices are dynamic we cannot prove anything.
+    if (!indexA || !indexB)
+      return false;
+    // If we find a different static index then different slices are being
+    // accessed.
+    if (indexA.getValue().template cast<IntegerAttr>().getInt() !=
+        indexB.getValue().template cast<IntegerAttr>().getInt())
+      return true;
+  }
+  return false;
+}
+
 void mlir::linalg::hoistRedundantVectorTransfers(FuncOp func) {
   bool changed = true;
   while (changed) {

 
       // Approximate aliasing by checking that:
       //   1. indices are the same,
-      //   2. no other use either dominates the transfer_read or is dominated
-      //   by the transfer_write (i.e. aliasing between the write and the read
-      //   across the loop).
+      //   2. no other operations in the loop access the same memref except
+      //      for transferRead/TramsferWrite accessing statically disjoint
+      //      slices.
       if (transferRead.indices() != transferWrite.indices())
         return WalkResult::advance();
 

       DominanceInfo dom(loop);
       if (!dom.properlyDominates(transferRead.getOperation(), transferWrite))
         return WalkResult::advance();
-      for (auto &use : transferRead.memref().getUses())
-        if (dom.properlyDominates(use.getOwner(),
-                                  transferRead.getOperation()) ||
-            dom.properlyDominates(transferWrite, use.getOwner()))
+      for (auto &use : transferRead.memref().getUses()) {
+        if (!dom.properlyDominates(loop, use.getOwner()))
+          continue;
+        if (use.getOwner() == transferRead.getOperation() ||
+            use.getOwner() == transferWrite.getOperation())
+          continue;
+        if (auto transferWriteUse =
+                dyn_cast<vector::TransferWriteOp>(use.getOwner())) {
+          if (!isDisjoint(transferWrite, transferWriteUse))
+            return WalkResult::advance();
+        } else if (auto transferReadUse =
+                       dyn_cast<vector::TransferReadOp>(use.getOwner())) {
+          if (!isDisjoint(transferWrite, transferReadUse))
+            return WalkResult::advance();
+        } else {
+          // Unknown use, we cannot prove that it doesn't alias with the
+          // transferRead/transferWrite operations.
           return WalkResult::advance();
+        }
+      }
 
       // Hoist read before.
       if (failed(loop.moveOutOfLoop({transferRead})))

mlir/test/Dialect/Linalg/hoisting.mlir

       %r3 = vector.transfer_read %memref3[%c0, %c0], %cst: memref<?x?xf32>, vector<4xf32>
       "some_crippling_use"(%memref4) : (memref<?x?xf32>) -> ()
       %r4 = vector.transfer_read %memref4[%c0, %c0], %cst: memref<?x?xf32>, vector<5xf32>
+      %r5 = vector.transfer_read %memref5[%c0, %c0], %cst: memref<?x?xf32>, vector<6xf32>
+      "some_crippling_use"(%memref5) : (memref<?x?xf32>) -> ()
       %u0 = "some_use"(%r0) : (vector<1xf32>) -> vector<1xf32>
       %u1 = "some_use"(%r1) : (vector<2xf32>) -> vector<2xf32>
       %u2 = "some_use"(%memref2) : (memref<?x?xf32>) -> vector<3xf32>
       %u3 = "some_use"(%r3) : (vector<4xf32>) -> vector<4xf32>
       %u4 = "some_use"(%r4) : (vector<5xf32>) -> vector<5xf32>
+      %u5 = "some_use"(%r5) : (vector<6xf32>) -> vector<6xf32>
       vector.transfer_write %u0, %memref1[%c0, %c0] : vector<1xf32>, memref<?x?xf32>
       vector.transfer_write %u1, %memref0[%i, %i] : vector<2xf32>, memref<?x?xf32>
       vector.transfer_write %u2, %memref2[%c0, %c0] : vector<3xf32>, memref<?x?xf32>
       vector.transfer_write %u3, %memref3[%c0, %c0] : vector<4xf32>, memref<?x?xf32>
       vector.transfer_write %u4, %memref4[%c0, %c0] : vector<5xf32>, memref<?x?xf32>
+      vector.transfer_write %u5, %memref5[%c0, %c0] : vector<6xf32>, memref<?x?xf32>
       "some_crippling_use"(%memref3) : (memref<?x?xf32>) -> ()
     }
   }
   return
 }
+
+// VECTOR_TRANSFERS-LABEL: func @hoist_vector_transfer_pairs_disjoint(
+//  VECTOR_TRANSFERS-SAME:   %[[MEMREF0:[a-zA-Z0-9]*]]: memref<?x?xf32>,
+//  VECTOR_TRANSFERS-SAME:   %[[MEMREF1:[a-zA-Z0-9]*]]: memref<?x?xf32>,
+//  VECTOR_TRANSFERS-SAME:   %[[VAL:[a-zA-Z0-9]*]]: index,
+//  VECTOR_TRANSFERS-SAME:   %[[LB:[a-zA-Z0-9]*]]: index,
+//  VECTOR_TRANSFERS-SAME:   %[[UB:[a-zA-Z0-9]*]]: index,
+//  VECTOR_TRANSFERS-SAME:   %[[STEP:[a-zA-Z0-9]*]]: index,
+//  VECTOR_TRANSFERS-SAME:   %[[RANDOM:[a-zA-Z0-9]*]]: index,
+//  VECTOR_TRANSFERS-SAME:   %[[CMP:[a-zA-Z0-9]*]]: i1
+func @hoist_vector_transfer_pairs_disjoint(
+    %memref0: memref<?x?xf32>, %memref1: memref<?x?xf32>, %val: index,
+    %lb : index, %ub : index, %step: index, %random_index : index,
+    %cmp: i1) {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %cst = constant 0.0 : f32
+
+// VECTOR_TRANSFERS: vector.transfer_read %{{.*}} : memref<?x?xf32>, vector<1xf32>
+// VECTOR_TRANSFERS: vector.transfer_read %{{.*}} : memref<?x?xf32>, vector<1xf32>
+// VECTOR_TRANSFERS: scf.for %[[I:.*]] = %[[LB]] to %[[UB]] step %[[STEP]] iter_args({{.*}}) -> (vector<1xf32>, vector<1xf32>) {
+// VECTOR_TRANSFERS:   scf.for %[[J:.*]] = %[[LB]] to %[[UB]] step %[[STEP]] iter_args({{.*}}) -> (vector<1xf32>, vector<1xf32>) {
+// VECTOR_TRANSFERS:     vector.transfer_read %{{.*}} : memref<?x?xf32>, vector<2xf32>
+// VECTOR_TRANSFERS:     vector.transfer_read %{{.*}} : memref<?x?xf32>, vector<2xf32>
+// VECTOR_TRANSFERS:     "some_use"(%{{.*}}) : (vector<1xf32>) -> vector<1xf32>
+// VECTOR_TRANSFERS:     "some_use"(%{{.*}}) : (vector<1xf32>) -> vector<1xf32>
+// VECTOR_TRANSFERS:     "some_use"(%{{.*}}) : (vector<2xf32>) -> vector<2xf32>
+// VECTOR_TRANSFERS:     "some_use"(%{{.*}}) : (vector<2xf32>) -> vector<2xf32>
+// VECTOR_TRANSFERS:     vector.transfer_write %{{.*}} : vector<2xf32>, memref<?x?xf32>
+// VECTOR_TRANSFERS:     vector.transfer_write %{{.*}} : vector<2xf32>, memref<?x?xf32>
+// VECTOR_TRANSFERS:     scf.yield {{.*}} : vector<1xf32>, vector<1xf32>
+// VECTOR_TRANSFERS:   }
+// VECTOR_TRANSFERS:   scf.yield {{.*}} : vector<1xf32>, vector<1xf32>
+// VECTOR_TRANSFERS: }
+// VECTOR_TRANSFERS: vector.transfer_write %{{.*}} : vector<1xf32>, memref<?x?xf32>
+// VECTOR_TRANSFERS: vector.transfer_write %{{.*}} : vector<1xf32>, memref<?x?xf32>
+  scf.for %i = %lb to %ub step %step {
+    scf.for %j = %lb to %ub step %step {
+      %r00 = vector.transfer_read %memref1[%c0, %c0], %cst: memref<?x?xf32>, vector<1xf32>
+      %r01 = vector.transfer_read %memref1[%c1, %c1], %cst: memref<?x?xf32>, vector<1xf32>
+      %r10 = vector.transfer_read %memref0[%i, %i], %cst: memref<?x?xf32>, vector<2xf32>
+      %r11 = vector.transfer_read %memref0[%random_index, %random_index], %cst: memref<?x?xf32>, vector<2xf32>
+      %u00 = "some_use"(%r00) : (vector<1xf32>) -> vector<1xf32>
+      %u01 = "some_use"(%r01) : (vector<1xf32>) -> vector<1xf32>
+      %u10 = "some_use"(%r10) : (vector<2xf32>) -> vector<2xf32>
+      %u11 = "some_use"(%r11) : (vector<2xf32>) -> vector<2xf32>
+      vector.transfer_write %u00, %memref1[%c0, %c0] : vector<1xf32>, memref<?x?xf32>
+      vector.transfer_write %u01, %memref1[%c1, %c1] : vector<1xf32>, memref<?x?xf32>
+      vector.transfer_write %u10, %memref0[%i, %i] : vector<2xf32>, memref<?x?xf32>
+      vector.transfer_write %u11, %memref0[%random_index, %random_index] : vector<2xf32>, memref<?x?xf32>
+    }
+  }
+  return
+}