[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:

Yes. If we do `intersect` between the `current op users` and `next op users`, few valid patterns can buy-pass this transformation. For example, in the below code: 

Example:
`%a = prod1`
`%b = prod2`
`%c = prod3`
`%d = prod4`
`%c1 = fma %a %c %arg`
`%c2 = fma %a %d %c1`
`%c3 = fma %b %c %c2`
`%c4 = fma %b %d %c3`

prod2 users: {`c3, c4`} and prod3 users: {`c1, c3`}. If the logic is based on intersect, the rewrite will be:
`%b = prod2`
`%c = prod3`
`%a = prod1`
`%c1 = fma %a %c %arg`
`%d = prod4`
`%c2 = fma %a %d %c1`
`%c3 = fma %b %c %c2`
`%c4 = fma %b %d %c3`

`prod2` can be moved before `%c3`, but didn't as `%c3` is a user for both of them. So, we check for the `first` users of both the `current` and `next` producer to do the shift.


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