[Mlir-commits] [mlir] [mlir][vector] Refine vector.transfer_read hoisting (PR #65770)

[Andrzej Warzyński]

https://github.com/banach-space created https://github.com/llvm/llvm-project/pull/65770:

Make sure that when analysing a `vector.transfer_read` that's a
candidate for hoisting, the users of relevant `memref.collapse_shape`
Op are included in the alias analysis.


>From 06a0d866c30e134e2fa3edf2dacc5e3f35669a77 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Fri, 8 Sep 2023 15:24:02 +0000
Subject: [PATCH] [mlir][vector] Refine vector.transfer_read hoisting

Make sure that when analysing a `vector.transfer_read` that's a
candidate for hoisting, the users of relevant `memref.collapse_shape`
Op are included in the alias analysis.
---
 .../Dialect/Linalg/Transforms/Hoisting.cpp    |  4 +++
 mlir/test/Dialect/Linalg/hoisting.mlir        | 36 +++++++++++++++++++
 2 files changed, 40 insertions(+)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
index 5f20ea42e499241..7dae1caa48b9077 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
@@ -69,6 +69,10 @@ static bool noAliasingUseInLoop(vector::TransferReadOp transferRead,
       users.append(subView->getUsers().begin(), subView->getUsers().end());
       continue;
     }
+    if (auto collapsed = dyn_cast<memref::CollapseShapeOp>(user)) {
+      users.append(collapsed->getUsers().begin(), collapsed->getUsers().end());
+      continue;
+    }
     if (isMemoryEffectFree(user) || isa<vector::TransferReadOp>(user))
       continue;
     if (!loop->isAncestor(user))
diff --git a/mlir/test/Dialect/Linalg/hoisting.mlir b/mlir/test/Dialect/Linalg/hoisting.mlir
index 5b1bb3fc15e09ea..6052775c7e8dea7 100644
--- a/mlir/test/Dialect/Linalg/hoisting.mlir
+++ b/mlir/test/Dialect/Linalg/hoisting.mlir
@@ -756,3 +756,39 @@ transform.sequence failures(propagate) {
   transform.structured.hoist_redundant_vector_transfers %0
     : (!transform.any_op) -> !transform.any_op
 }
+
+// -----
+
+// The users of `memref.collapse_shape %alloca` below write to `%alloca`, which
+// means that hoisting `vector.transfer_read %alloca` further down is not safe.
+
+// CHECK-LABEL:  func.func @collapse_shape
+//       CHECK:    scf.for {{.*}} {
+//       CHECK:      vector.transfer_write
+//       CHECK:      vector.transfer_read
+//       CHECK:      vector.transfer_write
+//       CHECK:    }
+
+func.func @collapse_shape(%2: memref<1x20x1xi32>, %0:  memref<1x28x1xi32>, %1: memref<9x1xi32>, %idx1: index, %idx2: index, %vec: vector<4xi32>) {
+  %c0_i32 = arith.constant 0 : i32
+  %c0 = arith.constant 0 : index
+  %c4 = arith.constant 4 : index
+  %c20 = arith.constant 20 : index
+  %alloca = memref.alloca() {alignment = 64 : i64} : memref<1x4x1xi32>
+  scf.for %arg0 = %c0 to %c20 step %c4 {
+    %subview = memref.subview %2[0, %arg0, 0] [1, 4, 1] [1, 1, 1] : memref<1x20x1xi32> to memref<1x4x1xi32, strided<[20, 1, 1], offset: ?>>
+    %collapse_shape = memref.collapse_shape %alloca [[0, 1, 2]] : memref<1x4x1xi32> into memref<4xi32>
+    vector.transfer_write %vec, %collapse_shape[%c0] {in_bounds = [true]} : vector<4xi32>, memref<4xi32>
+    %37 = vector.transfer_read %alloca[%c0, %c0, %c0], %c0_i32 {in_bounds = [true, true, true]} : memref<1x4x1xi32>, vector<1x4x1xi32>
+    vector.transfer_write %37, %subview[%c0, %c0, %c0] {in_bounds = [true, true, true]} : vector<1x4x1xi32>, memref<1x4x1xi32, strided<[20, 1, 1], offset: ?>>
+  }
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["func.func"]} in %arg1
+    : (!transform.any_op) -> !transform.any_op
+  transform.structured.hoist_redundant_vector_transfers %0
+    : (!transform.any_op) -> !transform.any_op
+}