[NFC][Flang][OpenMP] Refactor OpenMP.cpp::genOpenMPReduction

[DylanFleming-arm]

Summary

This patch serves two main purposes:
Firstly, to split some of the logic into a seperate method
to try and improve readability

On top of this, it aims to make creating the reductions more generic.
That way, subsequent patches adding reductions shouldn't need
to add a significant amount of extra logic checks, such as checking
for specific operators.

[awarzynski]

Makes sense to me, minor suggestions inline.

flang/include/flang/Lower/OpenMP.h

56

Could you add a comment to describe this method? Might be worth including names for the arguments, too. Perhaps this would make sense:

/// Checks whether /p loadVal is is used in an operation the result of which is then stored into /p reductionVal. If yes, then the operation corresponding to the reduction is returned. /p loadVal is assumed to be the value of a load operation and /p reductionVal is the results of an OpenMP reduction operation.
mlir::Operation *getReductionInChain(mlir::Value reductionVal, mlir::Value loadVal);

flang/lib/Lower/OpenMP.cpp

1662

Although the type of loadOperands is an mlir::OpOperand, it's actually something like loadUse (the actual type and what it represents don't match that well). Also, this is definitely singular (rather than plural) as it's just one item at a time :) Same for reductionOperands.

[DylanFleming-arm]

Added comment explaining getReductionInChain, renamed obscure/incorrect variables

[awarzynski]

Thanks for the updates - my remaining comments are nits, so feel free to ignore!

LGTM, though please wait till tomorrow before merging (in case there are more comments).

Also, the failure in the pre-merge CI looks bogus to me, so I would just ignore it. Apparently it's clang-format that failed, but it doesn't say why :) Unless you can see the actual reason for the failure?

flang/lib/Lower/OpenMP.cpp

1660–1664

[nit] Thanks for adding the comment! It should be moved into the header file and written in doxygen format, i.e.:

/// Checks whether /p loadVal is used in an operation, the result of which is then stored into /p reductionVal.
/// If yes, then the operation corresponding to the reduction is returned. 
/// /p loadVal is assumed to be the value of a load operation and /p reductionVal is the results
/// of an OpenMP reduction operation.

I see that you've just followed the style already present in the file. That's also fine and perhaps better for consistency.

1669

[nit] reductionUse would be more consistent with loadUse. And, from what I can tell, equally accurate.

flang/include/flang/Lower/OpenMP.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef FORTRAN_LOWER_OPENMP_H
 #define FORTRAN_LOWER_OPENMP_H
 
 #include <cinttypes>
 
+namespace mlir {
+class Value;
+class Operation;
+} // namespace mlir
+
+namespace fir {
+class FirOpBuilder;
+} // namespace fir
+
 namespace Fortran {
 namespace parser {
 struct OpenMPConstruct;
 struct OpenMPDeclarativeConstruct;
 struct OmpEndLoopDirective;
 struct OmpClauseList;
 } // namespace parser
 
void genOpenMPConstruct(AbstractConverter &, pft::Evaluation &,
                         const parser::OpenMPConstruct &);
 void genOpenMPDeclarativeConstruct(AbstractConverter &, pft::Evaluation &,
                                    const parser::OpenMPDeclarativeConstruct &);
 int64_t getCollapseValue(const Fortran::parser::OmpClauseList &clauseList);
 void genThreadprivateOp(AbstractConverter &, const pft::Variable &);
 void genOpenMPReduction(AbstractConverter &,
                         const Fortran::parser::OmpClauseList &clauseList);
 
+void updateReduction(mlir::Operation *, fir::FirOpBuilder &, mlir::Value,
+                     mlir::Value);
+
+mlir::Operation *getReductionInChain(mlir::Value, mlir::Value);
+
 } // namespace lower
 } // namespace Fortran
 
 #endif // FORTRAN_LOWER_OPENMP_H

flang/lib/Lower/OpenMP.cpp

if (const auto &reductionClause =
                 reductionOp->u)};
         if (intrinsicOp !=
             Fortran::parser::DefinedOperator::IntrinsicOperator::Add)
           continue;
         for (const auto &ompObject : objectList.v) {
           if (const auto *name{
                   Fortran::parser::Unwrap<Fortran::parser::Name>(ompObject)}) {
             if (const auto *symbol{name->symbol}) {
-              mlir::Value symVal = converter.getSymbolAddress(*symbol);
-              mlir::Type redType =
-                  symVal.getType().cast<fir::ReferenceType>().getEleTy();
-              if (!redType.isIntOrIndex())
+              mlir::Value reductionVal = converter.getSymbolAddress(*symbol);
+              mlir::Type reductionType =
+                  reductionVal.getType().cast<fir::ReferenceType>().getEleTy();
+              if (!reductionType.isIntOrIndex())
                 continue;
-              for (mlir::OpOperand &use1 : symVal.getUses()) {
-                if (auto load = mlir::dyn_cast<fir::LoadOp>(use1.getOwner())) {
-                  mlir::Value loadVal = load.getRes();
-                  for (mlir::OpOperand &use2 : loadVal.getUses()) {
-                    if (auto add = mlir::dyn_cast<mlir::arith::AddIOp>(
-                            use2.getOwner())) {
-                      mlir::Value addRes = add.getResult();
-                      for (mlir::OpOperand &use3 : addRes.getUses()) {
-                        if (auto store =
-                                mlir::dyn_cast<fir::StoreOp>(use3.getOwner())) {
-                          if (store.getMemref() == symVal) {
-                            // Chain found! Now replace load->reduction->store
-                            // with the OpenMP reduction operation.
-                            mlir::OpBuilder::InsertPoint insertPtDel =
-                                firOpBuilder.saveInsertionPoint();
-                            firOpBuilder.setInsertionPoint(add);
-                            if (add.getLhs() == loadVal) {
-                              firOpBuilder.create<mlir::omp::ReductionOp>(
-                                  add.getLoc(), add.getRhs(), symVal);
-                            } else {
-                              firOpBuilder.create<mlir::omp::ReductionOp>(
-                                  add.getLoc(), add.getLhs(), symVal);
-                            }
-                            store.erase();
-                            add.erase();
-                            load.erase();
-                            firOpBuilder.restoreInsertionPoint(insertPtDel);
-                          }
-                        }
-                      }
-                    }
-                  }
-                }
-              }
-            }
-          }
+
+              for (mlir::OpOperand &reductionValUse : reductionVal.getUses()) {
+
+                if (auto loadOp =
+                        mlir::dyn_cast<fir::LoadOp>(reductionValUse.getOwner())) {
+                  mlir::Value loadVal = loadOp.getRes();
+                  if (auto reductionOp = getReductionInChain(reductionVal, loadVal)) {
+                    updateReduction(reductionOp, firOpBuilder, loadVal, reductionVal);
+
+// Checks whether loadVal is used in an operation,
+// the result of which is then stored into reductionVal. 
+// If yes, then the operation corresponding to the reduction is returned. 
+// loadVal is assumed to be the value of a load operation
+// reductionVal is the results of an OpenMP reduction operation.
+mlir::Operation *Fortran::lower::getReductionInChain(mlir::Value reductionVal,
+                                                     mlir::Value loadVal) {
+  for (mlir::OpOperand &loadUse : loadVal.getUses()) {
+    if (auto reductionOp = loadUse.getOwner()) {
+      for (mlir::OpOperand &reductionOperand : reductionOp->getUses()) {
+        if (auto store =
+                mlir::dyn_cast<fir::StoreOp>(reductionOperand.getOwner())) {
+          if (store.getMemref() == reductionVal) {
+            store.erase();
+            return reductionOp;
+          }
         }
       }
     }
   }
+  return nullptr;
+}
+
+void Fortran::lower::updateReduction(mlir::Operation *op,
+                                     fir::FirOpBuilder &firOpBuilder,
+                                     mlir::Value loadVal, mlir::Value reductionVal) {
+  mlir::OpBuilder::InsertPoint insertPtDel = firOpBuilder.saveInsertionPoint();
+  firOpBuilder.setInsertionPoint(op);
+
+  if (op->getOperand(0) == loadVal)
+    firOpBuilder.create<mlir::omp::ReductionOp>(op->getLoc(), op->getOperand(1),
+                                                reductionVal);
+  else
+    firOpBuilder.create<mlir::omp::ReductionOp>(op->getLoc(), op->getOperand(0),
+                                                reductionVal);
+
+  firOpBuilder.restoreInsertionPoint(insertPtDel);
 }