[flang-commits] [flang] [flang][CodeGen] Run PreCGRewrite on omp reduction declare ops (PR #84954)

[Tom Eccles via flang-commits]

https://github.com/tblah updated https://github.com/llvm/llvm-project/pull/84954

>From 104a57ab6ec548ce3c202eccf7dddd1544d93aec Mon Sep 17 00:00:00 2001
From: Tom Eccles <tom.eccles at arm.com>
Date: Fri, 1 Mar 2024 16:59:09 +0000
Subject: [PATCH 1/5] [flang][CodeGen] Run PreCGRewrite on omp reduction
 declare ops

OpenMP reduction declare operations can contain FIR code which needs to
be lowered to LLVM. With array reductions, these regions can contain
more complicated operations which need PreCGRewriting. A similar extra
case was already needed for fir::GlobalOp.
---
 flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp b/flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp
index 9d48a2f08aba03..04a3a693658474 100644
--- a/flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp
+++ b/flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp
@@ -22,6 +22,7 @@
 #include "mlir/Transforms/RegionUtils.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/Debug.h"
+#include <mlir/Dialect/OpenMP/OpenMPDialect.h>
 
 namespace fir {
 #define GEN_PASS_DEF_CODEGENREWRITE
@@ -318,6 +319,10 @@ class CodeGenRewrite : public fir::impl::CodeGenRewriteBase<CodeGenRewrite> {
       runOn(func, func.getBody());
     for (auto global : mod.getOps<fir::GlobalOp>())
       runOn(global, global.getRegion());
+    for (auto omp : mod.getOps<mlir::omp::ReductionDeclareOp>()) {
+      runOn(omp, omp.getInitializerRegion());
+      runOn(omp, omp.getReductionRegion());
+    }
   }
 };
 

>From 07fc05338353d8288dcf609709c5fb4a045a859a Mon Sep 17 00:00:00 2001
From: Tom Eccles <tom.eccles at arm.com>
Date: Thu, 14 Mar 2024 07:54:12 +0000
Subject: [PATCH 2/5] Fix include

---
 flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp b/flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp
index 04a3a693658474..433e4b87aa4a55 100644
--- a/flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp
+++ b/flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp
@@ -18,11 +18,11 @@
 #include "flang/Optimizer/Dialect/FIROps.h"
 #include "flang/Optimizer/Dialect/FIRType.h"
 #include "flang/Optimizer/Dialect/Support/FIRContext.h"
+#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/RegionUtils.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/Debug.h"
-#include <mlir/Dialect/OpenMP/OpenMPDialect.h>
 
 namespace fir {
 #define GEN_PASS_DEF_CODEGENREWRITE

>From 88448175f0e60c321f1293fd284138dfd4d9a80d Mon Sep 17 00:00:00 2001
From: Tom Eccles <tom.eccles at arm.com>
Date: Thu, 14 Mar 2024 08:02:40 +0000
Subject: [PATCH 3/5] Run PreCGRewrite on all regions in the module

---
 flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp | 13 +++----------
 1 file changed, 3 insertions(+), 10 deletions(-)

diff --git a/flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp b/flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp
index 433e4b87aa4a55..074f165a87257b 100644
--- a/flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp
+++ b/flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp
@@ -284,7 +284,7 @@ class DeclareOpConversion : public mlir::OpRewritePattern<fir::DeclareOp> {
 
 class CodeGenRewrite : public fir::impl::CodeGenRewriteBase<CodeGenRewrite> {
 public:
-  void runOn(mlir::Operation *op, mlir::Region &region) {
+  void runOn(mlir::Operation *op) {
     auto &context = getContext();
     mlir::ConversionTarget target(context);
     target.addLegalDialect<mlir::arith::ArithDialect, fir::FIROpsDialect,
@@ -314,15 +314,8 @@ class CodeGenRewrite : public fir::impl::CodeGenRewriteBase<CodeGenRewrite> {
 
   void runOnOperation() override final {
     // Call runOn on all top level regions that may contain emboxOp/arrayCoorOp.
-    auto mod = getOperation();
-    for (auto func : mod.getOps<mlir::func::FuncOp>())
-      runOn(func, func.getBody());
-    for (auto global : mod.getOps<fir::GlobalOp>())
-      runOn(global, global.getRegion());
-    for (auto omp : mod.getOps<mlir::omp::ReductionDeclareOp>()) {
-      runOn(omp, omp.getInitializerRegion());
-      runOn(omp, omp.getReductionRegion());
-    }
+    mlir::ModuleOp mod = getOperation();
+    runOn(mod);
   }
 };
 

>From 55614ce2141812cb5a8c6150ba6864bc2dc5b6ec Mon Sep 17 00:00:00 2001
From: Tom Eccles <tom.eccles at arm.com>
Date: Thu, 14 Mar 2024 14:03:13 +0000
Subject: [PATCH 4/5] Remove unessecary include

---
 flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp | 1 -
 1 file changed, 1 deletion(-)

diff --git a/flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp b/flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp
index 074f165a87257b..6e43cdcd8d1fdf 100644
--- a/flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp
+++ b/flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp
@@ -18,7 +18,6 @@
 #include "flang/Optimizer/Dialect/FIROps.h"
 #include "flang/Optimizer/Dialect/FIRType.h"
 #include "flang/Optimizer/Dialect/Support/FIRContext.h"
-#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/RegionUtils.h"
 #include "llvm/ADT/STLExtras.h"

>From 4eef5d2232b9385e33f13c0f527c7aee8fa4d98e Mon Sep 17 00:00:00 2001
From: Tom Eccles <tom.eccles at arm.com>
Date: Fri, 15 Mar 2024 11:28:30 +0000
Subject: [PATCH 5/5] Remove runOn

---
 flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp | 14 +++++---------
 1 file changed, 5 insertions(+), 9 deletions(-)

diff --git a/flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp b/flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp
index 6e43cdcd8d1fdf..665bf09b8fc33b 100644
--- a/flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp
+++ b/flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp
@@ -283,7 +283,9 @@ class DeclareOpConversion : public mlir::OpRewritePattern<fir::DeclareOp> {
 
 class CodeGenRewrite : public fir::impl::CodeGenRewriteBase<CodeGenRewrite> {
 public:
-  void runOn(mlir::Operation *op) {
+  void runOnOperation() override final {
+    mlir::ModuleOp mod = getOperation();
+
     auto &context = getContext();
     mlir::ConversionTarget target(context);
     target.addLegalDialect<mlir::arith::ArithDialect, fir::FIROpsDialect,
@@ -300,7 +302,7 @@ class CodeGenRewrite : public fir::impl::CodeGenRewriteBase<CodeGenRewrite> {
     mlir::RewritePatternSet patterns(&context);
     fir::populatePreCGRewritePatterns(patterns);
     if (mlir::failed(
-            mlir::applyPartialConversion(op, target, std::move(patterns)))) {
+            mlir::applyPartialConversion(mod, target, std::move(patterns)))) {
       mlir::emitError(mlir::UnknownLoc::get(&context),
                       "error in running the pre-codegen conversions");
       signalPassFailure();
@@ -308,13 +310,7 @@ class CodeGenRewrite : public fir::impl::CodeGenRewriteBase<CodeGenRewrite> {
     }
     // Erase any residual (fir.shape, fir.slice...).
     mlir::IRRewriter rewriter(&context);
-    (void)mlir::runRegionDCE(rewriter, op->getRegions());
-  }
-
-  void runOnOperation() override final {
-    // Call runOn on all top level regions that may contain emboxOp/arrayCoorOp.
-    mlir::ModuleOp mod = getOperation();
-    runOn(mod);
+    (void)mlir::runRegionDCE(rewriter, mod->getRegions());
   }
 };