[Mlir-commits] [mlir] [mlir][vector] Prevent incorrect vector.transfer_{read|write} hoisting (PR #66930)

[Andrzej Warzyński]

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

Extracted from https://github.com/openxla/iree/issues/14994. This is a
draft to facilitate a discussion on the right approach for this.


>From b5452fef01939c564b98f572b6b24e641904f5c6 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Wed, 20 Sep 2023 17:15:09 +0000
Subject: [PATCH] [mlir][vector] Prevent incorrect vector.transfer_{read|write}
 hoisting

Extracted from https://github.com/openxla/iree/issues/14994. This is a
draft to facilitate a discussion on the right approach for this.
---
 .../Dialect/Linalg/Transforms/Hoisting.cpp    |  4 +++
 mlir/test/Dialect/Linalg/hoisting.mlir        | 34 +++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
index 7c6639304d97c58..68afc3f0649588f 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
@@ -152,6 +152,10 @@ void mlir::linalg::hoistRedundantVectorTransfers(func::FuncOp func) {
           transferRead.getPermutationMap() != transferWrite.getPermutationMap())
         return WalkResult::advance();
 
+      // TODO: Is this the right fix?
+      if (!transferRead.use_empty())
+        return WalkResult::advance();
+
       // TODO: may want to memoize this information for performance but it
       // likely gets invalidated often.
       DominanceInfo dom(loop);
diff --git a/mlir/test/Dialect/Linalg/hoisting.mlir b/mlir/test/Dialect/Linalg/hoisting.mlir
index e25914620726b9b..ff1ff33ea7d1cd0 100644
--- a/mlir/test/Dialect/Linalg/hoisting.mlir
+++ b/mlir/test/Dialect/Linalg/hoisting.mlir
@@ -827,3 +827,37 @@ transform.sequence failures(propagate) {
   transform.structured.hoist_redundant_vector_transfers %0
     : (!transform.any_op) -> !transform.any_op
 }
+
+// -----
+
+// Regression test - the `vector.transfer_read`/`vector.transfer_write` pair
+// reads/writes to a memref, so it re-uesed in every iteration. Hoisting this
+// pair would change the semantics of this loop.
+
+// CHECK-LABEL:  func.func @no_hoisting_collapse_shape_2
+//       CHECK:    scf.for {{.*}} {
+//       CHECK:      vector.transfer_read
+//       CHECK:      vector.transfer_write
+
+func.func @no_hoisting_write_to_memref(%arg0: i32, %arg1: vector<1xi32>) {
+  %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<1x1x1xi32>
+  scf.for %_ = %c0 to %c20 step %c4 {
+    %collapsed = memref.collapse_shape %alloca [[0, 1, 2]] : memref<1x1x1xi32> into memref<1xi32>
+    %10 = vector.transfer_read %collapsed[%c0], %c0_i32 {in_bounds = [true]} : memref<1xi32>, vector<1xi32>
+    %17 = vector.outerproduct %arg1, %arg0, %10 {kind = #vector.kind<add>} : vector<1xi32>, i32
+    vector.transfer_write %17, %collapsed[%c0] {in_bounds = [true]} : vector<1xi32>, memref<1xi32>
+  }
+  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
+}