[flang-commits] [flang] [flang][cuda] Rewrite predefined var in cuf kernel inside host function (PR #205974)
Valentin Clement バレンタイン クレメン via flang-commits
flang-commits at lists.llvm.org
Thu Jun 25 22:26:19 PDT 2026
https://github.com/clementval created https://github.com/llvm/llvm-project/pull/205974
Predefined variables are normally set inside device function but can appear in cuf kernel when a device function is inlined inside the cuf kernel.
Perviously the pass was only rewriting the predefined variables inside device function and not host function with cuf kernels. Update the pass to support this case.
>From e4c5b685a5191303d9d2787f4074df5df27135a1 Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Thu, 25 Jun 2026 22:24:17 -0700
Subject: [PATCH] [flang][cuda] Rewrite predefined var in cuf kernel inside
host function
---
.../Transforms/CUDA/CUFPredefinedVarToGPU.cpp | 95 ++++++++++++-------
flang/test/Fir/CUDA/predefined-variables.mlir | 5 +-
2 files changed, 64 insertions(+), 36 deletions(-)
diff --git a/flang/lib/Optimizer/Transforms/CUDA/CUFPredefinedVarToGPU.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFPredefinedVarToGPU.cpp
index 901f15d792173..2dd332488269d 100644
--- a/flang/lib/Optimizer/Transforms/CUDA/CUFPredefinedVarToGPU.cpp
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFPredefinedVarToGPU.cpp
@@ -6,6 +6,7 @@
//
//===----------------------------------------------------------------------===//
+#include "flang/Optimizer/Dialect/CUF/CUFOps.h"
#include "flang/Optimizer/Dialect/FIROps.h"
#include "flang/Optimizer/Dialect/FIROpsSupport.h"
#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
@@ -101,11 +102,60 @@ processDeclareOp(mlir::OpBuilder &builder, mlir::Location loc,
struct CUFPredefinedVarToGPU
: public fir::impl::CUFPredefinedVarToGPUBase<CUFPredefinedVarToGPU> {
+ void rewritePredefinedVars(mlir::Region ®ion, mlir::Location loc) {
+ if (region.empty())
+ return;
+
+ bool hasPredefinedDeclares = false;
+ region.walk([&](fir::DeclareOp declareOp) {
+ llvm::StringRef uniqName = declareOp.getUniqName();
+ hasPredefinedDeclares |= uniqName == mangleBuiltin(threadidx) ||
+ uniqName == mangleBuiltin(blockidx) ||
+ uniqName == mangleBuiltin(blockdim) ||
+ uniqName == mangleBuiltin(griddim);
+ });
+ if (!hasPredefinedDeclares)
+ return;
+
+ mlir::OpBuilder builder(region.getContext());
+ builder.setInsertionPointToStart(®ion.front());
+ auto c1 = mlir::arith::ConstantOp::create(
+ builder, loc, builder.getI32Type(), builder.getI32IntegerAttr(1));
+ llvm::SmallVector<mlir::Value, 3> threadids, blockids, blockdims, griddims;
+ createForAllDimensions<mlir::NVVM::ThreadIdXOp, mlir::NVVM::ThreadIdYOp,
+ mlir::NVVM::ThreadIdZOp>(builder, loc, c1, threadids,
+ /*incrementByOne=*/true);
+ createForAllDimensions<mlir::NVVM::BlockIdXOp, mlir::NVVM::BlockIdYOp,
+ mlir::NVVM::BlockIdZOp>(builder, loc, c1, blockids,
+ /*incrementByOne=*/true);
+ createForAllDimensions<mlir::NVVM::GridDimXOp, mlir::NVVM::GridDimYOp,
+ mlir::NVVM::GridDimZOp>(builder, loc, c1, griddims);
+ createForAllDimensions<mlir::NVVM::BlockDimXOp, mlir::NVVM::BlockDimYOp,
+ mlir::NVVM::BlockDimZOp>(builder, loc, c1,
+ blockdims);
+
+ llvm::SmallVector<mlir::Operation *> opsToDelete;
+ region.walk([&](fir::DeclareOp declareOp) {
+ processDeclareOp(builder, loc, declareOp, mangleBuiltin(threadidx),
+ threadids, opsToDelete);
+ processDeclareOp(builder, loc, declareOp, mangleBuiltin(blockidx),
+ blockids, opsToDelete);
+ processDeclareOp(builder, loc, declareOp, mangleBuiltin(blockdim),
+ blockdims, opsToDelete);
+ processDeclareOp(builder, loc, declareOp, mangleBuiltin(griddim),
+ griddims, opsToDelete);
+ });
+
+ for (auto *op : opsToDelete)
+ op->erase();
+ }
+
void runOnOperation() override {
func::FuncOp funcOp = getOperation();
if (funcOp.getBody().empty())
return;
+ bool rewrittenWholeFunction = false;
if (auto cudaProcAttr =
funcOp.getOperation()->getAttrOfType<cuf::ProcAttributeAttr>(
cuf::getProcAttrName())) {
@@ -113,42 +163,19 @@ struct CUFPredefinedVarToGPU
cudaProcAttr.getValue() == cuf::ProcAttribute::Global ||
cudaProcAttr.getValue() == cuf::ProcAttribute::GridGlobal ||
cudaProcAttr.getValue() == cuf::ProcAttribute::HostDevice) {
- mlir::Location loc = funcOp.getLoc();
- mlir::OpBuilder builder(funcOp.getContext());
- builder.setInsertionPointToStart(&funcOp.getBody().front());
- auto c1 = mlir::arith::ConstantOp::create(
- builder, loc, builder.getI32Type(), builder.getI32IntegerAttr(1));
- llvm::SmallVector<mlir::Value, 3> threadids, blockids, blockdims,
- griddims;
- createForAllDimensions<mlir::NVVM::ThreadIdXOp, mlir::NVVM::ThreadIdYOp,
- mlir::NVVM::ThreadIdZOp>(
- builder, loc, c1, threadids, /*incrementByOne=*/true);
- createForAllDimensions<mlir::NVVM::BlockIdXOp, mlir::NVVM::BlockIdYOp,
- mlir::NVVM::BlockIdZOp>(
- builder, loc, c1, blockids, /*incrementByOne=*/true);
- createForAllDimensions<mlir::NVVM::GridDimXOp, mlir::NVVM::GridDimYOp,
- mlir::NVVM::GridDimZOp>(builder, loc, c1,
- griddims);
- createForAllDimensions<mlir::NVVM::BlockDimXOp, mlir::NVVM::BlockDimYOp,
- mlir::NVVM::BlockDimZOp>(builder, loc, c1,
- blockdims);
-
- llvm::SmallVector<mlir::Operation *> opsToDelete;
- funcOp.walk([&](fir::DeclareOp declareOp) {
- processDeclareOp(builder, loc, declareOp, mangleBuiltin(threadidx),
- threadids, opsToDelete);
- processDeclareOp(builder, loc, declareOp, mangleBuiltin(blockidx),
- blockids, opsToDelete);
- processDeclareOp(builder, loc, declareOp, mangleBuiltin(blockdim),
- blockdims, opsToDelete);
- processDeclareOp(builder, loc, declareOp, mangleBuiltin(griddim),
- griddims, opsToDelete);
- });
-
- for (auto *op : opsToDelete)
- op->erase();
+ rewritePredefinedVars(funcOp.getRegion(), funcOp.getLoc());
+ rewrittenWholeFunction = true;
}
}
+
+ if (rewrittenWholeFunction)
+ return;
+
+ // Host functions containing cuf.kernel regions can still carry predefined
+ // vars in the kernel body. Rewrite them in-place as well.
+ funcOp.walk([&](cuf::KernelOp kernelOp) {
+ rewritePredefinedVars(kernelOp.getRegion(), kernelOp.getLoc());
+ });
}
};
diff --git a/flang/test/Fir/CUDA/predefined-variables.mlir b/flang/test/Fir/CUDA/predefined-variables.mlir
index 90b72f3bc1bcb..50f269e4c65db 100644
--- a/flang/test/Fir/CUDA/predefined-variables.mlir
+++ b/flang/test/Fir/CUDA/predefined-variables.mlir
@@ -237,7 +237,7 @@ func.func @_QMbarPgfoo2(%arg0: !fir.ref<i32> {cuf.data_attr = #cuf.cuda<device>,
// -----
-func.func @surviving_predefined_vars(%arg0: i32, %arg1: i32, %arg2: i32) attributes {cuf.proc_attr = #cuf.cuda_proc<global>} {
+func.func @surviving_predefined_vars(%arg0: i32, %arg1: i32, %arg2: i32) {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%c4 = arith.constant 4 : index
@@ -279,7 +279,7 @@ func.func @surviving_predefined_vars(%arg0: i32, %arg1: i32, %arg2: i32) attribu
// -----
-func.func @surviving_predefined_vars(%arg0: i32, %arg1: i32, %arg2: i32) attributes {cuf.proc_attr = #cuf.cuda_proc<global>} {
+func.func @surviving_predefined_vars(%arg0: i32, %arg1: i32, %arg2: i32) {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%c4 = arith.constant 4 : index
@@ -317,3 +317,4 @@ func.func @surviving_predefined_vars(%arg0: i32, %arg1: i32, %arg2: i32) attribu
// CHECK-NOT: _QM__fortran_builtinsE__builtin_blockidx
// CHECK-NOT: _QM__fortran_builtinsE__builtin_griddim
// CHECK-NOT: _QM__fortran_builtinsE__builtin_threadidx
+// CHECK: nvvm.read.ptx.sreg.tid.x
More information about the flang-commits
mailing list