[flang-commits] [flang] Revert "[flang][cuda] Update stream type for cuf kernel op" (PR #136644)

[Valentin Clement バレンタイン クレメン via flang-commits]

https://github.com/clementval created https://github.com/llvm/llvm-project/pull/136644

Reverts llvm/llvm-project#136627

>From f84b83e552e411409870ed67d264f7bef6aebfb1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Valentin=20Clement=20=28=E3=83=90=E3=83=AC=E3=83=B3?=
 =?UTF-8?q?=E3=82=BF=E3=82=A4=E3=83=B3=20=E3=82=AF=E3=83=AC=E3=83=A1?=
 =?UTF-8?q?=E3=83=B3=29?= <clementval at gmail.com>
Date: Mon, 21 Apr 2025 19:37:09 -0700
Subject: [PATCH] Revert "[flang][cuda] Update stream type for cuf kernel op
 (#136627)"

This reverts commit 46e734746db7176f6e32b3c98beacf1e94fced37.
---
 .../flang/Optimizer/Dialect/CUF/CUFOps.td     | 19 ++++++++++++-------
 flang/lib/Lower/Bridge.cpp                    | 10 ++++++----
 flang/lib/Optimizer/Dialect/CUF/CUFOps.cpp    |  2 +-
 .../Lower/CUDA/cuda-kernel-loop-directive.cuf |  4 +++-
 4 files changed, 22 insertions(+), 13 deletions(-)

diff --git a/flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td b/flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td
index 46cc59cda1612..926983d364ed1 100644
--- a/flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td
+++ b/flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td
@@ -254,19 +254,24 @@ def cuf_KernelOp : cuf_Op<"kernel", [AttrSizedOperandSegments,
     represented by a 0 constant value.
   }];
 
-  let arguments = (ins Variadic<I32>:$grid, // empty means `*`
-      Variadic<I32>:$block,                 // empty means `*`
-      Optional<fir_ReferenceType>:$stream, Variadic<Index>:$lowerbound,
-      Variadic<Index>:$upperbound, Variadic<Index>:$step,
-      OptionalAttr<I64Attr>:$n, Variadic<AnyType>:$reduceOperands,
-      OptionalAttr<ArrayAttr>:$reduceAttrs);
+  let arguments = (ins
+    Variadic<I32>:$grid, // empty means `*`
+    Variadic<I32>:$block, // empty means `*`
+    Optional<I32>:$stream,
+    Variadic<Index>:$lowerbound,
+    Variadic<Index>:$upperbound,
+    Variadic<Index>:$step,
+    OptionalAttr<I64Attr>:$n,
+    Variadic<AnyType>:$reduceOperands,
+    OptionalAttr<ArrayAttr>:$reduceAttrs
+  );
 
   let regions = (region AnyRegion:$region);
 
   let assemblyFormat = [{
     `<` `<` `<` custom<CUFKernelValues>($grid, type($grid)) `,` 
                 custom<CUFKernelValues>($block, type($block))
-        ( `,` `stream` `=` $stream^ `:` qualified(type($stream)))? `>` `>` `>`
+        ( `,` `stream` `=` $stream^ )? `>` `>` `>`
         ( `reduce` `(` $reduceOperands^ `:` type($reduceOperands) `:` $reduceAttrs `)` )?
         custom<CUFKernelLoopControl>($region, $lowerbound, type($lowerbound),
             $upperbound, type($upperbound), $step, type($step))
diff --git a/flang/lib/Lower/Bridge.cpp b/flang/lib/Lower/Bridge.cpp
index 7b76845b5af05..1652a86ed7e63 100644
--- a/flang/lib/Lower/Bridge.cpp
+++ b/flang/lib/Lower/Bridge.cpp
@@ -3097,7 +3097,7 @@ class FirConverter : public Fortran::lower::AbstractConverter {
 
     llvm::SmallVector<mlir::Value> gridValues;
     llvm::SmallVector<mlir::Value> blockValues;
-    mlir::Value streamAddr;
+    mlir::Value streamValue;
 
     if (launchConfig) {
       const std::list<Fortran::parser::CUFKernelDoConstruct::StarOrExpr> &grid =
@@ -3130,8 +3130,10 @@ class FirConverter : public Fortran::lower::AbstractConverter {
       }
 
       if (stream)
-        streamAddr = fir::getBase(
-            genExprAddr(*Fortran::semantics::GetExpr(*stream), stmtCtx));
+        streamValue = builder->createConvert(
+            loc, builder->getI32Type(),
+            fir::getBase(
+                genExprValue(*Fortran::semantics::GetExpr(*stream), stmtCtx)));
     }
 
     const auto &outerDoConstruct =
@@ -3265,7 +3267,7 @@ class FirConverter : public Fortran::lower::AbstractConverter {
     }
 
     auto op = builder->create<cuf::KernelOp>(
-        loc, gridValues, blockValues, streamAddr, lbs, ubs, steps, n,
+        loc, gridValues, blockValues, streamValue, lbs, ubs, steps, n,
         mlir::ValueRange(reduceOperands), builder->getArrayAttr(reduceAttrs));
     builder->createBlock(&op.getRegion(), op.getRegion().end(), ivTypes,
                          ivLocs);
diff --git a/flang/lib/Optimizer/Dialect/CUF/CUFOps.cpp b/flang/lib/Optimizer/Dialect/CUF/CUFOps.cpp
index 24033bc15b8eb..a86f12c2c4a55 100644
--- a/flang/lib/Optimizer/Dialect/CUF/CUFOps.cpp
+++ b/flang/lib/Optimizer/Dialect/CUF/CUFOps.cpp
@@ -271,7 +271,7 @@ llvm::LogicalResult cuf::KernelOp::verify() {
         return emitOpError("expect reduce attributes to be ReduceAttr");
     }
   }
-  return checkStreamType(*this);
+  return mlir::success();
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/flang/test/Lower/CUDA/cuda-kernel-loop-directive.cuf b/flang/test/Lower/CUDA/cuda-kernel-loop-directive.cuf
index 10f0b9e3d1215..0fceb292f10d2 100644
--- a/flang/test/Lower/CUDA/cuda-kernel-loop-directive.cuf
+++ b/flang/test/Lower/CUDA/cuda-kernel-loop-directive.cuf
@@ -75,7 +75,9 @@ subroutine sub1()
   end do
 end
 
-! CHECK: cuf.kernel<<<*, *, stream = %[[STREAM]]#0 : !fir.ref<i64>>>>
+! CHECK: %[[STREAM_LOAD:.*]] = fir.load %[[STREAM]]#0 : !fir.ref<i64>
+! CHECK: %[[STREAM_I32:.*]] = fir.convert %[[STREAM_LOAD]] : (i64) -> i32
+! CHECK: cuf.kernel<<<*, *, stream = %[[STREAM_I32]]>>>
 
 
 ! Test lowering with unstructured construct inside.