From 385768aff2412c1b8e1e0e10f6911c5cd15b50e3 Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy@amd.com>
Date: Mon, 11 Mar 2024 06:41:50 -0500
Subject: [PATCH] Handle review comments.

---
 .../flang/Optimizer/Transforms/Passes.td      |  2 +-
 .../Transforms/DoConcurrentConversion.cpp     | 40 +++++++++++++++----
 2 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/flang/include/flang/Optimizer/Transforms/Passes.td b/flang/include/flang/Optimizer/Transforms/Passes.td
index 06de4a1d28a929..95ffd199bc0fe4 100644
--- a/flang/include/flang/Optimizer/Transforms/Passes.td
+++ b/flang/include/flang/Optimizer/Transforms/Passes.td
@@ -400,7 +400,7 @@ def FunctionAttr : Pass<"function-attr", "mlir::func::FuncOp"> {
 def DoConcurrentConversionPass : Pass<"fopenmp-do-concurrent-conversion", "mlir::func::FuncOp"> {
   let summary = "Map `DO CONCURRENT` loops to OpenMP worksharing loops.";
 
-  let description = [{ This is an experimental pass to map `DO CONCURRENR` loops
+  let description = [{ This is an experimental pass to map `DO CONCURRENT` loops
      to their correspnding equivalent OpenMP worksharing constructs.
 
      For now the following is supported:
diff --git a/flang/lib/Optimizer/Transforms/DoConcurrentConversion.cpp b/flang/lib/Optimizer/Transforms/DoConcurrentConversion.cpp
index 4534d514b86d76..f5a3d925ab5d9a 100644
--- a/flang/lib/Optimizer/Transforms/DoConcurrentConversion.cpp
+++ b/flang/lib/Optimizer/Transforms/DoConcurrentConversion.cpp
@@ -70,13 +70,39 @@ class DoConcurrentConversion : public mlir::OpConversionPattern<fir::DoLoopOp> {
     // loop header preparation/allocation operations.
 
     // Clone the LB, UB, step defining ops inside the parallel region.
+    mlir::Operation* lbOp = doLoop.getLowerBound().getDefiningOp();
+    mlir::Operation* ubOp = doLoop.getUpperBound().getDefiningOp();
+    mlir::Operation* stepOp = doLoop.getStep().getDefiningOp();
+
+    if (lbOp == nullptr || ubOp == nullptr || stepOp == nullptr) {
+      return rewriter.notifyMatchFailure(
+          doLoop, "At least one of the loop's LB, UB, or step doesn't have a "
+                  "defining operation.");
+    }
+
+    std::function<bool(mlir::Operation *)> isOpUltimatelyConstant =
+        [&](mlir::Operation *operation) {
+          if (mlir::isa_and_present<mlir::arith::ConstantOp>(operation))
+            return true;
+
+          if (fir::ConvertOp convertOp =
+                  mlir::dyn_cast_if_present<fir::ConvertOp>(operation))
+            return isOpUltimatelyConstant(convertOp.getValue().getDefiningOp());
+
+          return false;
+        };
+
+    if (!isOpUltimatelyConstant(lbOp) || !isOpUltimatelyConstant(ubOp) ||
+        !isOpUltimatelyConstant(stepOp)) {
+      return rewriter.notifyMatchFailure(
+          doLoop, "`do concurrent` conversion is currently only supported for "
+                  "constant LB, UB, and step values.");
+    }
+
     llvm::SmallVector<mlir::Value> lowerBound, upperBound, step;
-    lowerBound.push_back(
-        rewriter.clone(*doLoop.getLowerBound().getDefiningOp())->getResult(0));
-    upperBound.push_back(
-        rewriter.clone(*doLoop.getUpperBound().getDefiningOp())->getResult(0));
-    step.push_back(
-        rewriter.clone(*doLoop.getStep().getDefiningOp())->getResult(0));
+    lowerBound.push_back(rewriter.clone(*lbOp)->getResult(0));
+    upperBound.push_back(rewriter.clone(*ubOp)->getResult(0));
+    step.push_back(rewriter.clone(*stepOp)->getResult(0));
     // ==== TODO (1) End ====
 
     auto wsLoopOp = rewriter.create<mlir::omp::WsLoopOp>(
@@ -127,7 +153,7 @@ class DoConcurrentConversion : public mlir::OpConversionPattern<fir::DoLoopOp> {
       workList.remove(item);
     }
 
-    // For each collected `fir.sotre`, find the target memref's alloca's and
+    // For each collected `fir.store`, find the target memref's alloca's and
     // declare ops.
     llvm::SmallSetVector<mlir::Operation *, 4> declareAndAllocasToClone;
     for (auto storeOp : inductionVarTargetStores) {