1 files changed, 49 insertions, 9 deletions

diff --git a/flang/lib/Optimizer/HLFIR/Transforms/LowerHLFIROrderedAssignments.cpp b/flang/lib/Optimizer/HLFIR/Transforms/LowerHLFIROrderedAssignments.cpp
index 85e32ffa0ec2..a0dbd46975cd 100644
--- a/flang/lib/Optimizer/HLFIR/Transforms/LowerHLFIROrderedAssignments.cpp
+++ b/flang/lib/Optimizer/HLFIR/Transforms/LowerHLFIROrderedAssignments.cpp
@@ -12,37 +12,76 @@
 // where.
 // The pass lowers these operations to regular hlfir.assign, loops and, if
 // needed, introduces temporary storage to fulfill Fortran semantics.
+//
+// For each rewrite, an analysis builds an evaluation schedule, and then the
+// new code is generated by following the evaluation schedule.
 //===----------------------------------------------------------------------===//
 
+#include "ScheduleOrderedAssignments.h"
 #include "flang/Optimizer/Builder/Todo.h"
-#include "flang/Optimizer/HLFIR/HLFIROps.h"
 #include "flang/Optimizer/HLFIR/Passes.h"
 #include "mlir/Transforms/DialectConversion.h"
+#include "llvm/Support/Debug.h"
 
 namespace hlfir {
 #define GEN_PASS_DEF_LOWERHLFIRORDEREDASSIGNMENTS
 #include "flang/Optimizer/HLFIR/Passes.h.inc"
 } // namespace hlfir
 
-using namespace mlir;
+#define DEBUG_TYPE "flang-ordered-assignment"
+
+// Test option only to test the scheduling part only (operations are erased
+// without codegen). The only goal is to allow printing and testing the debug
+// info.
+static llvm::cl::opt<bool> dbgScheduleOnly(
+    "flang-dbg-order-assignment-schedule-only",
+    llvm::cl::desc("Only run ordered assignment scheduling with no codegen"),
+    llvm::cl::init(false));
+
+/// Shared rewrite entry point for all the ordered assignment tree root
+/// operations. It calls the scheduler and then apply the schedule.
+static mlir::LogicalResult
+rewrite(hlfir::OrderedAssignmentTreeOpInterface &root,
+        bool tryFusingAssignments, mlir::PatternRewriter &rewriter) {
+  (void)hlfir::buildEvaluationSchedule(root, tryFusingAssignments);
+
+  LLVM_DEBUG(
+      /// Debug option to print the scheduling debug info without doing
+      /// any code generation. The operations are simply erased to avoid
+      /// failing and calling the rewrite patterns on nested operations.
+      /// The only purpose of this is to help testing scheduling without
+      /// having to test generated code.
+      if (dbgScheduleOnly) {
+        rewriter.eraseOp(root);
+        return mlir::success();
+      });
+  // TODO: lower to loops according to schedule.
+  return mlir::failure();
+}
 
 namespace {
 
 class ForallOpConversion : public mlir::OpRewritePattern<hlfir::ForallOp> {
 public:
-  explicit ForallOpConversion(mlir::MLIRContext *ctx) : OpRewritePattern{ctx} {}
+  explicit ForallOpConversion(mlir::MLIRContext *ctx, bool tryFusingAssignments)
+      : OpRewritePattern{ctx}, tryFusingAssignments{tryFusingAssignments} {}
 
   mlir::LogicalResult
   matchAndRewrite(hlfir::ForallOp forallOp,
                   mlir::PatternRewriter &rewriter) const override {
-    TODO(forallOp.getLoc(), "FORALL construct or statement in HLFIR");
-    return mlir::failure();
+    auto root = mlir::cast<hlfir::OrderedAssignmentTreeOpInterface>(
+        forallOp.getOperation());
+    if (mlir::failed(::rewrite(root, tryFusingAssignments, rewriter)))
+      TODO(forallOp.getLoc(), "FORALL construct or statement in HLFIR");
+    return mlir::success();
   }
+  const bool tryFusingAssignments;
 };
 
 class WhereOpConversion : public mlir::OpRewritePattern<hlfir::WhereOp> {
 public:
-  explicit WhereOpConversion(mlir::MLIRContext *ctx) : OpRewritePattern{ctx} {}
+  explicit WhereOpConversion(mlir::MLIRContext *ctx, bool tryFusingAssignments)
+      : OpRewritePattern{ctx}, tryFusingAssignments{tryFusingAssignments} {}
 
   mlir::LogicalResult
   matchAndRewrite(hlfir::WhereOp whereOp,
@@ -50,6 +89,7 @@ public:
     TODO(whereOp.getLoc(), "WHERE construct or statement in HLFIR");
     return mlir::failure();
   }
+  const bool tryFusingAssignments;
 };
 
 class RegionAssignConversion
@@ -84,9 +124,9 @@ public:
     // operations that can be the root of ordered assignments. The other
     // operations will be taken care of while rewriting these trees (they
     // cannot exist outside of these operations given their verifiers/traits).
-    patterns
-        .insert<ForallOpConversion, WhereOpConversion, RegionAssignConversion>(
-            context);
+    patterns.insert<ForallOpConversion, WhereOpConversion>(
+        context, this->tryFusingAssignments.getValue());
+    patterns.insert<RegionAssignConversion>(context);
     mlir::ConversionTarget target(*context);
     target.markUnknownOpDynamicallyLegal([](mlir::Operation *op) {
       return !mlir::isa<hlfir::OrderedAssignmentTreeOpInterface>(op);