[Mlir-commits] [mlir] [mlir][x86vector] Sink Vector.transfer_reads and vector.load before the consumer (PR #169333)

[Renato Golin]

================
@@ -0,0 +1,93 @@
+//===- SinkVectorProducerOps.cpp ------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Dialect/Vector/Utils/VectorUtils.h"
+#include "mlir/Dialect/X86Vector/Transforms.h"
+#include "mlir/Dialect/X86Vector/X86VectorDialect.h"
+
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/Dominance.h"
+#include "mlir/IR/PatternMatch.h"
+
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+using namespace mlir::vector;
+using namespace mlir::x86vector;
+
+/// Sink vector producers forward to reduce live ranges.
+/// This pattern applies to ops such as vector.load and vector.transfer_read.
+template <typename producerOp>
+struct SinkVectorProducerOps final : public OpRewritePattern<producerOp> {
+  using OpRewritePattern<producerOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(producerOp op,
+                                PatternRewriter &rewriter) const override {
+
+    // Collect all users of the producer op.
+    llvm::SmallVector<Operation *> users;
+    for (OpResult result : op->getResults())
+      for (Operation *user : result.getUsers())
+        users.push_back(user);
+
+    // If there are no users, nothing to sink.
+    if (users.empty())
+      return failure();
+
+    // If the next op is already a user, do not move.
+    Operation *nextOp = op->getNextNode();
+    if (llvm::is_contained(users, nextOp))
+      return failure();
+
+    // Prevent pathological looping:
+    // If the next op produces values used by any of op's users, don't move.
----------------
rengolin wrote:

I see the problem:
```
  %0 = my_op
  %1 = next_op
  ...
  %m = first_user %0, %1
```
Both %0 and %1 can be in any order, so if %0 moves now, it will become %1 and if you run the pass again, it'll match and go back being %0, in a cycle.

However, this is not only true for the _next_ operation, but any operation between %0 and %m. This is a high polynomial time algorithm, especially with intersection checks on two lists.

What if we just intersect users?
```
  users = all users of all results of "op";
  other = op;
  while (other =  other->getNextNode()) {
    // If ops share users: O(num_ops * num_users)
    if (is_contained(users, other->getUsers()) {
      // Move "op" just _before_ "nextOp" and stop
      op->moveBefore(other);
      break;
    }
  }
```

Now, "op" will never _pass_ "other", so if you run again, it will do nothing.

Running the pass multiple times should _compact_ the IR:
```
  // Original
  %a = my_op
  ...
  %b = my_other_op
  ...
  %c = my_last_op
  ...
  %val = first_user (%a, %b, %c)

  // First pass
  ...
  %a = my_op
  ...
  %b = my_other_op
  ...
  %c = my_last_op
  %val = first_user (%a, %b, %c)

  // Second pass
  ...
  ...
  %a = my_op
  ...
  %b = my_other_op
  %c = my_last_op
  %val = first_user (%a, %b, %c)

  // Third pass
  ...
  ...
  ...
  %a = my_op
  %b = my_other_op
  %c = my_last_op
  %val = first_user (%a, %b, %c)
```

Alternatively, you can take the first op (`%a`), get it's first user (`%val`), and then iterate _backwards_, moving `%c` first, then %b`, then %a`, and you only need to pass once to _compact_ the IR.

https://github.com/llvm/llvm-project/pull/169333