[Mlir-commits] [mlir] [OpenMP][MLIR] Fix initTargetRuntimeAttrs, LoopTripCount calculation when host_eval vars is empty (PR #157717)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-llvm

@llvm/pr-subscribers-mlir

Author: Chaitanya (skc7)

<details>
<summary>Changes</summary>

In OpenMPToLLVMIRTranslation, If `host_eval` vars are empty, LoopTripCount cannot be calculated.

This PR fixes this issue by adding a check to presence of host_eval vars before calculating the loopTripCount.

---
Full diff: https://github.com/llvm/llvm-project/pull/157717.diff


2 Files Affected:

- (modified) mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp (+7-2) 
- (added) mlir/test/Target/LLVMIR/omptarget-loopnest-no-host-eval.mlir (+28) 


``````````diff
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 6694de8383534..019ce626903b2 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -5393,8 +5393,11 @@ initTargetRuntimeAttrs(llvm::IRBuilderBase &builder,
   Value numThreads, numTeamsLower, numTeamsUpper, teamsThreadLimit;
   llvm::SmallVector<Value> lowerBounds(numLoops), upperBounds(numLoops),
       steps(numLoops);
-  extractHostEvalClauses(targetOp, numThreads, numTeamsLower, numTeamsUpper,
-                         teamsThreadLimit, &lowerBounds, &upperBounds, &steps);
+  if (!targetOp.getHostEvalVars().empty()) {
+    extractHostEvalClauses(targetOp, numThreads, numTeamsLower, numTeamsUpper,
+                           teamsThreadLimit, &lowerBounds, &upperBounds,
+                           &steps);
+  }
 
   // TODO: Handle constant 'if' clauses.
   if (Value targetThreadLimit = targetOp.getThreadLimit())
@@ -5424,6 +5427,8 @@ initTargetRuntimeAttrs(llvm::IRBuilderBase &builder,
     // here, since we're only interested in the trip count.
     for (auto [loopLower, loopUpper, loopStep] :
          llvm::zip_equal(lowerBounds, upperBounds, steps)) {
+      if (!loopLower || !loopUpper || !loopStep)
+        break;
       llvm::Value *lowerBound = moduleTranslation.lookupValue(loopLower);
       llvm::Value *upperBound = moduleTranslation.lookupValue(loopUpper);
       llvm::Value *step = moduleTranslation.lookupValue(loopStep);
diff --git a/mlir/test/Target/LLVMIR/omptarget-loopnest-no-host-eval.mlir b/mlir/test/Target/LLVMIR/omptarget-loopnest-no-host-eval.mlir
new file mode 100644
index 0000000000000..8628c115364f9
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/omptarget-loopnest-no-host-eval.mlir
@@ -0,0 +1,28 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+// Ensure that the -mlir-to-llmvir pass doesn't crash.
+
+// CHECK: define void @_QQmain()
+
+module attributes {llvm.target_triple = "x86_64-unknown-linux-gnu", omp.is_gpu = false, omp.is_target_device = false} {
+  llvm.func @_QQmain()  {
+    omp.target {
+      %0 = llvm.mlir.constant(1000 : i32) : i32
+      %1 = llvm.mlir.constant(1 : i32) : i32
+      omp.teams {
+        omp.parallel {
+          omp.distribute {
+            omp.wsloop {
+              omp.loop_nest (%arg0) : i32 = (%1) to (%0) inclusive step (%1) {
+                omp.yield
+              }
+            } {omp.composite}
+          } {omp.composite}
+          omp.terminator
+        } {omp.composite}
+        omp.terminator
+      }
+      omp.terminator
+    }
+    llvm.return
+  }
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/157717