[flang-commits] [flang] 6ca9a30 - [flang][cuda] Update stream operand type for cuf.kernel_launch op (#135222)

[via flang-commits]

Author: Valentin Clement (バレンタイン クレメン)
Date: 2025-04-10T11:13:27-07:00
New Revision: 6ca9a30c26641ac5f836cf794bed01ff61972f37

URL: https://github.com/llvm/llvm-project/commit/6ca9a30c26641ac5f836cf794bed01ff61972f37
DIFF: https://github.com/llvm/llvm-project/commit/6ca9a30c26641ac5f836cf794bed01ff61972f37.diff

LOG: [flang][cuda] Update stream operand type for cuf.kernel_launch op (#135222)

Added: 
    

Modified: 
    flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td
    flang/lib/Lower/ConvertCall.cpp
    flang/test/Lower/CUDA/cuda-kernel-calls.cuf

Removed: 
    


################################################################################
diff  --git a/flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td b/flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td
index e95e27bfa5ad3..feef5485194f8 100644
--- a/flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td
+++ b/flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td
@@ -197,24 +197,16 @@ def cuf_KernelLaunchOp : cuf_Op<"kernel_launch", [CallOpInterface,
     ```
   }];
 
-  let arguments = (ins
-    SymbolRefAttr:$callee,
-    I32:$grid_x,
-    I32:$grid_y,
-    I32:$grid_z,
-    I32:$block_x,
-    I32:$block_y,
-    I32:$block_z,
-    Optional<I32>:$bytes,
-    Optional<I32>:$stream,
-    Variadic<AnyType>:$args,
-    OptionalAttr<DictArrayAttr>:$arg_attrs,
-    OptionalAttr<DictArrayAttr>:$res_attrs
-  );
+  let arguments = (ins SymbolRefAttr:$callee, I32:$grid_x, I32:$grid_y,
+      I32:$grid_z, I32:$block_x, I32:$block_y, I32:$block_z,
+      Optional<I32>:$bytes, Optional<AnyIntegerType>:$stream,
+      Variadic<AnyType>:$args, OptionalAttr<DictArrayAttr>:$arg_attrs,
+      OptionalAttr<DictArrayAttr>:$res_attrs);
 
   let assemblyFormat = [{
     $callee `<` `<` `<` $grid_x `,` $grid_y `,` $grid_z `,`$block_x `,`
-        $block_y `,` $block_z ( `,` $bytes^ ( `,` $stream^ )? )? `>` `>` `>`
+        $block_y `,` $block_z
+        ( `,` $bytes^ ( `,` $stream^ `:` type($stream) )? )? `>` `>` `>`
         `` `(` $args `)` ( `:` `(` type($args)^ `)` )? attr-dict
   }];
 

diff  --git a/flang/lib/Lower/ConvertCall.cpp b/flang/lib/Lower/ConvertCall.cpp
index 6a0f4d1090adc..d674775ffb522 100644
--- a/flang/lib/Lower/ConvertCall.cpp
+++ b/flang/lib/Lower/ConvertCall.cpp
@@ -588,10 +588,8 @@ Fortran::lower::genCallOpAndResult(
 
     mlir::Value stream; // stream is optional.
     if (caller.getCallDescription().chevrons().size() > 3)
-      stream = builder.createConvert(
-          loc, i32Ty,
-          fir::getBase(converter.genExprValue(
-              caller.getCallDescription().chevrons()[3], stmtCtx)));
+      stream = fir::getBase(converter.genExprValue(
+          caller.getCallDescription().chevrons()[3], stmtCtx));
 
     builder.create<cuf::KernelLaunchOp>(
         loc, funcType.getResults(), funcSymbolAttr, grid_x, grid_y, grid_z,

diff  --git a/flang/test/Lower/CUDA/cuda-kernel-calls.cuf b/flang/test/Lower/CUDA/cuda-kernel-calls.cuf
index 7db8ef4fcb680..d66d2811f7a8b 100644
--- a/flang/test/Lower/CUDA/cuda-kernel-calls.cuf
+++ b/flang/test/Lower/CUDA/cuda-kernel-calls.cuf
@@ -15,6 +15,8 @@ contains
 
   subroutine host()
     real, device :: a
+    integer(8) :: stream
+
 ! CHECK-LABEL: func.func @_QMtest_callPhost()
 ! CHECK: %[[A:.*]]:2 = hlfir.declare %{{.*}} {data_attr = #cuf.cuda<device>, uniq_name = "_QMtest_callFhostEa"} : (!fir.ref<f32>) -> (!fir.ref<f32>, !fir.ref<f32>)
 
@@ -51,6 +53,10 @@ contains
 
     call dev_kernel1<<<*, 32>>>(a)
 ! CHECK: cuf.kernel_launch @_QMtest_callPdev_kernel1<<<%c-1{{.*}}, %c1{{.*}}, %c1{{.*}}, %c32{{.*}}, %c1{{.*}}, %c1{{.*}}>>>(%{{.*}})
+
+    call dev_kernel1<<<*,32,0,stream>>>(a)
+! CHECK: cuf.kernel_launch @_QMtest_callPdev_kernel1<<<%c-1{{.*}}, %c1{{.*}}, %c1{{.*}}, %c32{{.*}}, %c1{{.*}}, %c1{{.*}}, %c0{{.*}}, %{{.*}} : i64>>>(%{{.*}}) : (!fir.ref<f32>)
+
   end
 
 end