[flang][hlfir] Add ordered assignment pass with TODOs

[jeanPerier]

Summary

I plan to implement lowering from parse tree to HLFIR first for forall
and where to ease testing of the rewrite pass while writing it.

To avoid cryptic errors in ConvertToFir pass about unhandled operations,
this patch already defines the pass that will further lower these
operations and make it throw clear TODO messages.

[clementval]

LGTM

flang/include/flang/Optimizer/HLFIR/Passes.h

 namespace hlfir {
 #define GEN_PASS_DECL
 #include "flang/Optimizer/HLFIR/Passes.h.inc"
 
 std::unique_ptr<mlir::Pass> createConvertHLFIRtoFIRPass();
 std::unique_ptr<mlir::Pass> createBufferizeHLFIRPass();
 std::unique_ptr<mlir::Pass> createLowerHLFIRIntrinsicsPass();
 std::unique_ptr<mlir::Pass> createSimplifyHLFIRIntrinsicsPass();
+std::unique_ptr<mlir::Pass> createLowerHLFIROrderedAssignmentsPass();
 
 #define GEN_PASS_REGISTRATION
 #include "flang/Optimizer/HLFIR/Passes.h.inc"
 } // namespace hlfir
 
 #endif // FORTRAN_OPTIMIZER_HLFIR_PASSES_H

flang/include/flang/Optimizer/HLFIR/Passes.td

def BufferizeHLFIR : Pass<"bufferize-hlfir", "::mlir::ModuleOp"> {
   let constructor = "hlfir::createBufferizeHLFIRPass()";
 }
 
 def LowerHLFIRIntrinsics : Pass<"lower-hlfir-intrinsics", "::mlir::ModuleOp"> {
   let summary = "Lower HLFIR transformational intrinsic operations";
   let constructor = "hlfir::createLowerHLFIRIntrinsicsPass()";
 }
 
+def LowerHLFIROrderedAssignments : Pass<"lower-hlfir-ordered-assignments", "::mlir::ModuleOp"> {
+  let summary = "Lower HLFIR ordered assignments like forall and where operations";
+  let constructor = "hlfir::createLowerHLFIROrderedAssignmentsPass()";
+}
+
 def SimplifyHLFIRIntrinsics : Pass<"simplify-hlfir-intrinsics", "::mlir::func::FuncOp"> {
   let summary = "Simplify HLFIR intrinsic operations that don't need to result in runtime calls";
   let constructor = "hlfir::createSimplifyHLFIRIntrinsicsPass()";
 }
 
 #endif //FORTRAN_DIALECT_HLFIR_PASSES

flang/include/flang/Tools/CLOptions.inc

 /// \param optLevel - optimization level used for creating FIR optimization
 ///   passes pipeline
 inline void createHLFIRToFIRPassPipeline(
     mlir::PassManager &pm, llvm::OptimizationLevel optLevel = defaultOptLevel) {
   if (optLevel.isOptimizingForSpeed()) {
     addCanonicalizerPassWithoutRegionSimplification(pm);
     pm.addPass(hlfir::createSimplifyHLFIRIntrinsicsPass());
   }
+  pm.addPass(hlfir::createLowerHLFIROrderedAssignmentsPass());
   pm.addPass(hlfir::createLowerHLFIRIntrinsicsPass());
   pm.addPass(hlfir::createBufferizeHLFIRPass());
   pm.addPass(hlfir::createConvertHLFIRtoFIRPass());
 }
 
 #if !defined(FLANG_EXCLUDE_CODEGEN)
 inline void createDebugPasses(
     mlir::PassManager &pm, llvm::codegenoptions::DebugInfoKind debugLevel) {

flang/lib/Optimizer/HLFIR/Transforms/CMakeLists.txt

 get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
 
 add_flang_library(HLFIRTransforms
   BufferizeHLFIR.cpp
   ConvertToFIR.cpp
   LowerHLFIRIntrinsics.cpp
+  LowerHLFIROrderedAssignments.cpp
   SimplifyHLFIRIntrinsics.cpp
 
   DEPENDS
   FIRDialect
   HLFIROpsIncGen
   ${dialect_libs}
 
   LINK_LIBS

flang/lib/Optimizer/HLFIR/Transforms/LowerHLFIROrderedAssignments.cpp

+//===- LowerHLFIROrderedAssignments.cpp - Lower HLFIR ordered assignments -===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+// This file defines a pass to lower HLFIR ordered assignments.
+// Ordered assignments are all the operations with the
+// OrderedAssignmentTreeOpInterface that implements user defined assignments,
+// assignment to vector subscripted entities, and assignments inside forall and
+// where.
+// The pass lowers these operations to regular hlfir.assign, loops and, if
+// needed, introduces temporary storage to fulfill Fortran semantics.
+//===----------------------------------------------------------------------===//
+
+#include "flang/Optimizer/Builder/Todo.h"
+#include "flang/Optimizer/HLFIR/HLFIROps.h"
+#include "flang/Optimizer/HLFIR/Passes.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace hlfir {
+#define GEN_PASS_DEF_LOWERHLFIRORDEREDASSIGNMENTS
+#include "flang/Optimizer/HLFIR/Passes.h.inc"
+} // namespace hlfir
+
+using namespace mlir;
+
+namespace {
+
+class ForallOpConversion : public mlir::OpRewritePattern<hlfir::ForallOp> {
+public:
+  explicit ForallOpConversion(mlir::MLIRContext *ctx) : OpRewritePattern{ctx} {}
+
+  mlir::LogicalResult
+  matchAndRewrite(hlfir::ForallOp forallOp,
+                  mlir::PatternRewriter &rewriter) const override {
+    TODO(forallOp.getLoc(), "FORALL construct or statement in HLFIR");
+    return mlir::failure();
+  }
+};
+
+class WhereOpConversion : public mlir::OpRewritePattern<hlfir::WhereOp> {
+public:
+  explicit WhereOpConversion(mlir::MLIRContext *ctx) : OpRewritePattern{ctx} {}
+
+  mlir::LogicalResult
+  matchAndRewrite(hlfir::WhereOp whereOp,
+                  mlir::PatternRewriter &rewriter) const override {
+    TODO(whereOp.getLoc(), "WHERE construct or statement in HLFIR");
+    return mlir::failure();
+  }
+};
+
+class RegionAssignConversion
+    : public mlir::OpRewritePattern<hlfir::RegionAssignOp> {
+public:
+  explicit RegionAssignConversion(mlir::MLIRContext *ctx)
+      : OpRewritePattern{ctx} {}
+
+  mlir::LogicalResult
+  matchAndRewrite(hlfir::RegionAssignOp regionAssignOp,
+                  mlir::PatternRewriter &rewriter) const override {
+    if (!regionAssignOp.getUserDefinedAssignment().empty())
+      TODO(regionAssignOp.getLoc(), "user defined assignment in HLFIR");
+    else
+      TODO(regionAssignOp.getLoc(),
+           "assignments to vector subscripted entity in HLFIR");
+    return mlir::failure();
+  }
+};
+
+class LowerHLFIROrderedAssignments
+    : public hlfir::impl::LowerHLFIROrderedAssignmentsBase<
+          LowerHLFIROrderedAssignments> {
+public:
+  void runOnOperation() override {
+    // Running on a ModuleOp because this pass may generate FuncOp declaration
+    // for runtime calls. This could be a FuncOp pass otherwise.
+    auto module = this->getOperation();
+    auto *context = &getContext();
+    mlir::RewritePatternSet patterns(context);
+    // Patterns are only defined for the OrderedAssignmentTreeOpInterface
+    // 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);
+    mlir::ConversionTarget target(*context);
+    target.markUnknownOpDynamicallyLegal([](mlir::Operation *op) {
+      return !mlir::isa<hlfir::OrderedAssignmentTreeOpInterface>(op);
+    });
+    if (mlir::failed(mlir::applyPartialConversion(module, target,
+                                                  std::move(patterns)))) {
+      mlir::emitError(mlir::UnknownLoc::get(context),
+                      "failure in HLFIR ordered assignments lowering pass");
+      signalPassFailure();
+    }
+  }
+};
+} // namespace
+
+std::unique_ptr<mlir::Pass> hlfir::createLowerHLFIROrderedAssignmentsPass() {
+  return std::make_unique<LowerHLFIROrderedAssignments>();
+}

flang/test/Driver/mlir-debug-pass-pipeline.f90

 ! DEBUG-DIRECTIVES: warning: Unsupported debug option: line-directives-only
 !
 ! DEBUG-ERR: error: invalid value 'invalid' in '-debug-info-kind=invalid'
 ! DEBUG-ERR-NOT: Pass statistics report
 
 ! ALL: Pass statistics report
 
 ! ALL: Fortran::lower::VerifierPass
+! ALL-NEXT: LowerHLFIROrderedAssignments
 ! ALL-NEXT: LowerHLFIRIntrinsics
 ! ALL-NEXT: BufferizeHLFIR
 ! ALL-NEXT: ConvertHLFIRtoFIR
 ! ALL-NEXT: CSE
 ! Ideally, we need an output with only the pass names, but
 ! there is currently no way to get that, so in order to
 ! guarantee that the passes are in the expected order
 ! (i.e. use -NEXT) we have to check the statistics output as well.

flang/test/Driver/mlir-pass-pipeline.f90

 end program
 
 ! ALL: Pass statistics report
 
 ! ALL: Fortran::lower::VerifierPass
 ! O2-NEXT: Canonicalizer
 ! O2-NEXT: 'func.func' Pipeline
 ! O2-NEXT:   SimplifyHLFIRIntrinsics
+! ALL-NEXT: LowerHLFIROrderedAssignments
 ! ALL-NEXT: LowerHLFIRIntrinsics
 ! ALL-NEXT: BufferizeHLFIR
 ! ALL-NEXT: ConvertHLFIRtoFIR
 ! ALL-NEXT: CSE
 ! Ideally, we need an output with only the pass names, but
 ! there is currently no way to get that, so in order to
 ! guarantee that the passes are in the expected order
 ! (i.e. use -NEXT) we have to check the statistics output as well.

flang/test/Fir/basic-program.fir

 // CHECK-LABEL: define void @_QQmain()
 // CHECK:       ret void
 
 // PASSES: Pass statistics report
 
 // PASSES:        Canonicalizer
 // PASSES-NEXT: 'func.func' Pipeline
 // PASSES-NEXT:   SimplifyHLFIRIntrinsics
+// PASSES-NEXT:   LowerHLFIROrderedAssignments
 // PASSES-NEXT:   LowerHLFIRIntrinsics
 // PASSES-NEXT:   BufferizeHLFIR
 // PASSES-NEXT:   ConvertHLFIRtoFIR
 // PASSES-NEXT:   CSE
 // PASSES-NEXT:   (S) 0 num-cse'd - Number of operations CSE'd
 // PASSES-NEXT:   (S) 0 num-dce'd - Number of operations DCE'd
 
 // PASSES-NEXT: 'func.func' Pipeline

flang/test/HLFIR/ordered-assignments-codegen-todo.fir

+// Just test that Ordered assignment pass TODOs are properly reported.
+// RUN: %not_todo_cmd fir-opt --lower-hlfir-ordered-assignments %s 2>&1 | FileCheck %s
+
+
+// CHECK: not yet implemented: FORALL construct or statement in HLFIR
+
+func.func @forall_todo(%arg0: !fir.ref<!fir.array<10xf32>>, %arg1: !fir.ref<!fir.array<10xf32>>) {
+  %c1 = arith.constant 1 : index
+  %c10 = arith.constant 10 : index
+  hlfir.forall lb {
+    hlfir.yield %c1 : index
+  } ub {
+    hlfir.yield %c10 : index
+  }  (%arg2: i64) {
+    hlfir.region_assign {
+      %1 = hlfir.designate %arg1 (%arg2)  : (!fir.ref<!fir.array<10xf32>>, i64) -> !fir.ref<f32>
+      hlfir.yield %1 : !fir.ref<f32>
+    } to {
+      %1 = hlfir.designate %arg0 (%arg2)  : (!fir.ref<!fir.array<10xf32>>, i64) -> !fir.ref<f32>
+      hlfir.yield %1 : !fir.ref<f32>
+    }
+  }
+  return
+}