[flang-commits] [flang] [flang][cuda] Update FIROps.td to add $grid_z to CudaKernelLaunch (PR #85318)

[Iman Hosseini via flang-commits]

https://github.com/ImanHosseini updated https://github.com/llvm/llvm-project/pull/85318

>From a10d6a393bf0c107db1bc5d70d4c137bc98f5fcb Mon Sep 17 00:00:00 2001
From: Iman Hosseini <hosseini.iman at yahoo.com>
Date: Thu, 14 Mar 2024 21:46:38 +0000
Subject: [PATCH 1/2] Update FIROps.td to add $grid_z to CudaKernelLaunch

grid can be 3 dimensional.
---
 flang/include/flang/Optimizer/Dialect/FIROps.td | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/flang/include/flang/Optimizer/Dialect/FIROps.td b/flang/include/flang/Optimizer/Dialect/FIROps.td
index 65a86d25333b5d..f4792637f481c0 100644
--- a/flang/include/flang/Optimizer/Dialect/FIROps.td
+++ b/flang/include/flang/Optimizer/Dialect/FIROps.td
@@ -2454,6 +2454,7 @@ def fir_CUDAKernelLaunch : fir_Op<"cuda_kernel_launch", [CallOpInterface,
     SymbolRefAttr:$callee,
     I32:$grid_x,
     I32:$grid_y,
+    I32:$grid_z,
     I32:$block_x,
     I32:$block_y,
     I32:$block_z,
@@ -2463,8 +2464,8 @@ def fir_CUDAKernelLaunch : fir_Op<"cuda_kernel_launch", [CallOpInterface,
   );
 
   let assemblyFormat = [{
-    $callee `<` `<` `<` $grid_x `,` $grid_y `,` $block_x `,` $block_y `,`
-        $block_z ( `,` $bytes^ ( `,` $stream^ )? )? `>` `>` `>`
+    $callee `<` `<` `<` $grid_x `,` $grid_y `,` $grid_z `,`$block_x `,`
+        $block_y `,` $block_z ( `,` $bytes^ ( `,` $stream^ )? )? `>` `>` `>`
         `` `(` ( $args^ `:` type($args) )? `)` attr-dict
   }];
 

>From 4a469090b910ff8f201d87831a507f04eea7f8f7 Mon Sep 17 00:00:00 2001
From: Iman Hosseini <hosseini.iman at yahoo.com>
Date: Thu, 14 Mar 2024 22:14:17 +0000
Subject: [PATCH 2/2] Update ConvertCall.cpp for 3d grid

---
 flang/lib/Lower/ConvertCall.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/flang/lib/Lower/ConvertCall.cpp b/flang/lib/Lower/ConvertCall.cpp
index 6e3ce101ef1af9..990912195d1445 100644
--- a/flang/lib/Lower/ConvertCall.cpp
+++ b/flang/lib/Lower/ConvertCall.cpp
@@ -466,8 +466,8 @@ std::pair<fir::ExtendedValue, bool> Fortran::lower::genCallOpAndResult(
               caller.getCallDescription().chevrons()[3], stmtCtx)));
 
     builder.create<fir::CUDAKernelLaunch>(
-        loc, funcType.getResults(), funcSymbolAttr, grid_x, grid_y, block_x,
-        block_y, block_z, bytes, stream, operands);
+        loc, funcType.getResults(), funcSymbolAttr, grid_x, grid_y, one,
+        block_x, block_y, block_z, bytes, stream, operands);
     callNumResults = 0;
   } else if (caller.requireDispatchCall()) {
     // Procedure call requiring a dynamic dispatch. Call is created with