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

[Arun Thangamani]

================
@@ -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.
----------------
arun-thmn wrote:

Thanks. Now, in the `first` pass we iterate through all the instructions in the IR and capture the first user of all `vector.load` or `vector.transfer_read`. Then, move the `prod` right before its first `use`. It should take `O(n)` for best and average case. The worst case is `O(n2)`.

For the second and next `pass`, the function `checkLooping` returns `failure()` if the immediate nextOp of the `producer` is it's `user`. 

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