[flang][cuda] Handle lowering of stars in cuf kernel launch parameters

[clementval]

Parsing of the cuf kernel loop directive has been updated to handle variants with the * syntax. This patch updates the lowering to make use of them.

• If the grid or block syntax uses only stars then the operation variadic operand remains empty.
• If there is values and stars, then stars are represented as a zero constant value.

[llvmbot]

@llvm/pr-subscribers-flang-fir-hlfir

Author: Valentin Clement (バレンタイン クレメン) (clementval)

Changes

Parsing of the cuf kernel loop directive has been updated to handle variants with the * syntax. This patch updates the lowering to make use of them.

• If the grid or block syntax uses only stars then the operation variadic operand remains empty.
• If there is values and stars, then stars are represented as a zero constant value.

---

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

2 Files Affected:

• (modified) flang/lib/Lower/Bridge.cpp (+32-13)
• (modified) flang/test/Lower/CUDA/cuda-kernel-loop-directive.cuf (+16-3)

diff --git a/flang/lib/Lower/Bridge.cpp b/flang/lib/Lower/Bridge.cpp
index 650ec5db2d0ccb..1b9a8a867b0804 100644
--- a/flang/lib/Lower/Bridge.cpp
+++ b/flang/lib/Lower/Bridge.cpp
@@ -2529,23 +2529,42 @@ class FirConverter : public Fortran::lower::AbstractConverter {
     const std::optional<Fortran::parser::ScalarIntExpr> &stream =
         std::get<3>(dir.t);
 
+    auto isOnlyStars =
+        [&](const std::list<Fortran::parser::CUFKernelDoConstruct::StarOrExpr>
+                &list) -> bool {
+      for (const Fortran::parser::CUFKernelDoConstruct::StarOrExpr &expr :
+           list) {
+        if (expr.v)
+          return false;
+      }
+      return true;
+    };
+
+    mlir::Value zero =
+        builder->createIntegerConstant(loc, builder->getI32Type(), 0);
+
     llvm::SmallVector<mlir::Value> gridValues;
-    for (const Fortran::parser::CUFKernelDoConstruct::StarOrExpr &expr : grid) {
-      if (expr.v) {
-        gridValues.push_back(fir::getBase(
-            genExprValue(*Fortran::semantics::GetExpr(*expr.v), stmtCtx)));
-      } else {
-        // TODO: '*'
+    if (!isOnlyStars(grid)) {
+      for (const Fortran::parser::CUFKernelDoConstruct::StarOrExpr &expr :
+           grid) {
+        if (expr.v) {
+          gridValues.push_back(fir::getBase(
+              genExprValue(*Fortran::semantics::GetExpr(*expr.v), stmtCtx)));
+        } else {
+          gridValues.push_back(zero);
+        }
       }
     }
     llvm::SmallVector<mlir::Value> blockValues;
-    for (const Fortran::parser::CUFKernelDoConstruct::StarOrExpr &expr :
-         block) {
-      if (expr.v) {
-        blockValues.push_back(fir::getBase(
-            genExprValue(*Fortran::semantics::GetExpr(*expr.v), stmtCtx)));
-      } else {
-        // TODO: '*'
+    if (!isOnlyStars(block)) {
+      for (const Fortran::parser::CUFKernelDoConstruct::StarOrExpr &expr :
+           block) {
+        if (expr.v) {
+          blockValues.push_back(fir::getBase(
+              genExprValue(*Fortran::semantics::GetExpr(*expr.v), stmtCtx)));
+        } else {
+          blockValues.push_back(zero);
+        }
       }
     }
     mlir::Value streamValue;
diff --git a/flang/test/Lower/CUDA/cuda-kernel-loop-directive.cuf b/flang/test/Lower/CUDA/cuda-kernel-loop-directive.cuf
index c017561447f85d..6179e609db383c 100644
--- a/flang/test/Lower/CUDA/cuda-kernel-loop-directive.cuf
+++ b/flang/test/Lower/CUDA/cuda-kernel-loop-directive.cuf
@@ -42,7 +42,20 @@ subroutine sub1()
 ! CHECK: fir.cuda_kernel<<<%c1{{.*}}, (%c256{{.*}}, %c1{{.*}})>>> (%{{.*}} : index, %{{.*}} : index) = (%{{.*}}, %{{.*}} : index, index) to (%{{.*}}, %{{.*}} : index, index) step (%{{.*}}, %{{.*}} : index, index)
 ! CHECK: {n = 2 : i64}
 
-! TODO: lowering for these cases
-! !$cuf kernel do(2) <<< (1,*), (256,1) >>>
-! !$cuf kernel do(2) <<< (*,*), (32,4) >>>
+  !$cuf kernel do(2) <<< (1,*), (256,1) >>>
+  do i = 1, n
+    do j = 1, n
+      c(i,j) = c(i,j) * d(i,j)
+    end do
+  end do
+! CHECK: fir.cuda_kernel<<<(%c1{{.*}}, %c0{{.*}}), (%c256{{.*}}, %c1{{.*}})>>> (%{{.*}} : index, %{{.*}} : index) = (%{{.*}}, %{{.*}} : index, index) to (%{{.*}}, %{{.*}} : index, index)  step (%{{.*}}, %{{.*}} : index, index)
+
+!$cuf kernel do(2) <<< (*,*), (32,4) >>>
+  do i = 1, n
+    do j = 1, n
+      c(i,j) = c(i,j) * d(i,j)
+    end do
+  end do
+
+! CHECK: fir.cuda_kernel<<<*, (%c32{{.*}}, %c4{{.*}})>>> (%{{.*}} : index, %{{.*}} : index) = (%{{.*}}, %{{.*}} : index, index) to (%{{.*}}, %{{.*}} : index, index)  step (%{{.*}}, %{{.*}} : index, index)
 end

[vzakhari]

LGTM!

Adding a comment about special meaning of 0 into the operation description might be useful.

[vzakhari]

Thanks!