[flang-commits] [flang] df55ebf - [flang][cuda] Create predefined variable reads at the location of use (#222058)
via flang-commits
flang-commits at lists.llvm.org
Tue Sep 8 15:56:46 PDT 2026
Author: jiel-nv
Date: 2026-09-08T15:56:41-07:00
New Revision: df55ebfbe36ec8c6e19fd5b06f218a56cc6f5e23
URL: https://github.com/llvm/llvm-project/commit/df55ebfbe36ec8c6e19fd5b06f218a56cc6f5e23
DIFF: https://github.com/llvm/llvm-project/commit/df55ebfbe36ec8c6e19fd5b06f218a56cc6f5e23.diff
LOG: [flang][cuda] Create predefined variable reads at the location of use (#222058)
Added:
Modified:
flang/lib/Optimizer/Transforms/CUDA/CUFPredefinedVarToGPU.cpp
flang/test/Fir/CUDA/predefined-variables.mlir
Removed:
################################################################################
diff --git a/flang/lib/Optimizer/Transforms/CUDA/CUFPredefinedVarToGPU.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFPredefinedVarToGPU.cpp
index 0bf18ed69050c..4f669461d92d8 100644
--- a/flang/lib/Optimizer/Transforms/CUDA/CUFPredefinedVarToGPU.cpp
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFPredefinedVarToGPU.cpp
@@ -22,25 +22,6 @@ using namespace mlir;
namespace {
-template <typename OpTyX, typename OpTyY, typename OpTyZ>
-static void createForAllDimensions(mlir::OpBuilder &builder, mlir::Location loc,
- mlir::Value c1,
- SmallVectorImpl<mlir::Value> &values,
- bool incrementByOne = false) {
- if (incrementByOne) {
- auto baseX = OpTyX::create(builder, loc, builder.getI32Type());
- values.push_back(mlir::arith::AddIOp::create(builder, loc, baseX, c1));
- auto baseY = OpTyY::create(builder, loc, builder.getI32Type());
- values.push_back(mlir::arith::AddIOp::create(builder, loc, baseY, c1));
- auto baseZ = OpTyZ::create(builder, loc, builder.getI32Type());
- values.push_back(mlir::arith::AddIOp::create(builder, loc, baseZ, c1));
- } else {
- values.push_back(OpTyX::create(builder, loc, builder.getI32Type()));
- values.push_back(OpTyY::create(builder, loc, builder.getI32Type()));
- values.push_back(OpTyZ::create(builder, loc, builder.getI32Type()));
- }
-}
-
static constexpr llvm::StringRef builtinsModuleName = "__fortran_builtins";
static constexpr llvm::StringRef builtinVarPrefix = "__builtin_";
static constexpr llvm::StringRef threadidx = "threadidx";
@@ -57,10 +38,10 @@ std::string mangleBuiltin(llvm::StringRef varName) {
varName.str();
}
+template <typename OpTy>
static void
-processCoordinateOp(mlir::OpBuilder &builder, mlir::Location loc,
- fir::CoordinateOp coordOp, unsigned fieldIdx,
- mlir::Value &gpuValue,
+processCoordinateOp(mlir::OpBuilder &builder, fir::CoordinateOp coordOp,
+ unsigned fieldIdx, bool incrementByOne,
llvm::SmallVectorImpl<mlir::Operation *> &opsToDelete) {
std::optional<llvm::ArrayRef<int32_t>> fieldIndices =
coordOp.getFieldIndices();
@@ -71,28 +52,37 @@ processCoordinateOp(mlir::OpBuilder &builder, mlir::Location loc,
assert(mlir::isa<fir::LoadOp>(coordUse.getOwner()) &&
"only expect load op");
auto loadOp = mlir::dyn_cast<fir::LoadOp>(coordUse.getOwner());
+ // Use the loadOp's loc as the location info for the register read op.
+ mlir::Location loc = loadOp.getLoc();
+ builder.setInsertionPoint(loadOp);
+ mlir::Value gpuValue = OpTy::create(builder, loc, builder.getI32Type());
+ if (incrementByOne) {
+ auto c1 = mlir::arith::ConstantOp::create(
+ builder, loc, builder.getI32Type(), builder.getI32IntegerAttr(1));
+ gpuValue = mlir::arith::AddIOp::create(builder, loc, gpuValue, c1);
+ }
loadOp.getResult().replaceAllUsesWith(gpuValue);
opsToDelete.push_back(loadOp);
}
}
}
+template <typename OpTyX, typename OpTyY, typename OpTyZ>
static void
-processDeclareOp(mlir::OpBuilder &builder, mlir::Location loc,
- fir::DeclareOp declareOp, llvm::StringRef builtinVar,
- llvm::SmallVectorImpl<mlir::Value> &gpuValues,
+processDeclareOp(mlir::OpBuilder &builder, fir::DeclareOp declareOp,
+ llvm::StringRef builtinVar, bool incrementByOne,
llvm::SmallVectorImpl<mlir::Operation *> &opsToDelete,
llvm::SmallPtrSetImpl<mlir::Operation *> &memrefDefiningOps) {
if (declareOp.getUniqName().str().compare(builtinVar) == 0) {
for (mlir::OpOperand &use : declareOp.getResult().getUses()) {
fir::CoordinateOp coordOp =
mlir::dyn_cast<fir::CoordinateOp>(use.getOwner());
- processCoordinateOp(builder, loc, coordOp, field_x, gpuValues[0],
- opsToDelete);
- processCoordinateOp(builder, loc, coordOp, field_y, gpuValues[1],
- opsToDelete);
- processCoordinateOp(builder, loc, coordOp, field_z, gpuValues[2],
- opsToDelete);
+ processCoordinateOp<OpTyX>(builder, coordOp, field_x, incrementByOne,
+ opsToDelete);
+ processCoordinateOp<OpTyY>(builder, coordOp, field_y, incrementByOne,
+ opsToDelete);
+ processCoordinateOp<OpTyZ>(builder, coordOp, field_z, incrementByOne,
+ opsToDelete);
opsToDelete.push_back(coordOp);
}
opsToDelete.push_back(declareOp.getOperation());
@@ -107,7 +97,7 @@ processDeclareOp(mlir::OpBuilder &builder, mlir::Location loc,
struct CUFPredefinedVarToGPU
: public fir::impl::CUFPredefinedVarToGPUBase<CUFPredefinedVarToGPU> {
- void rewritePredefinedVars(mlir::Region ®ion, mlir::Location loc) {
+ void rewritePredefinedVars(mlir::Region ®ion) {
if (region.empty())
return;
@@ -123,33 +113,25 @@ struct CUFPredefinedVarToGPU
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;
llvm::SmallPtrSet<mlir::Operation *, 4> memrefDefiningOps;
region.walk([&](fir::DeclareOp declareOp) {
- processDeclareOp(builder, loc, declareOp, mangleBuiltin(threadidx),
- threadids, opsToDelete, memrefDefiningOps);
- processDeclareOp(builder, loc, declareOp, mangleBuiltin(blockidx),
- blockids, opsToDelete, memrefDefiningOps);
- processDeclareOp(builder, loc, declareOp, mangleBuiltin(blockdim),
- blockdims, opsToDelete, memrefDefiningOps);
- processDeclareOp(builder, loc, declareOp, mangleBuiltin(griddim),
- griddims, opsToDelete, memrefDefiningOps);
+ processDeclareOp<mlir::NVVM::ThreadIdXOp, mlir::NVVM::ThreadIdYOp,
+ mlir::NVVM::ThreadIdZOp>(
+ builder, declareOp, mangleBuiltin(threadidx),
+ /*incrementByOne=*/true, opsToDelete, memrefDefiningOps);
+ processDeclareOp<mlir::NVVM::BlockIdXOp, mlir::NVVM::BlockIdYOp,
+ mlir::NVVM::BlockIdZOp>(
+ builder, declareOp, mangleBuiltin(blockidx),
+ /*incrementByOne=*/true, opsToDelete, memrefDefiningOps);
+ processDeclareOp<mlir::NVVM::BlockDimXOp, mlir::NVVM::BlockDimYOp,
+ mlir::NVVM::BlockDimZOp>(
+ builder, declareOp, mangleBuiltin(blockdim),
+ /*incrementByOne=*/false, opsToDelete, memrefDefiningOps);
+ processDeclareOp<mlir::NVVM::GridDimXOp, mlir::NVVM::GridDimYOp,
+ mlir::NVVM::GridDimZOp>(
+ builder, declareOp, mangleBuiltin(griddim),
+ /*incrementByOne=*/false, opsToDelete, memrefDefiningOps);
});
for (auto *op : opsToDelete)
@@ -174,7 +156,7 @@ struct CUFPredefinedVarToGPU
cudaProcAttr.getValue() == cuf::ProcAttribute::Global ||
cudaProcAttr.getValue() == cuf::ProcAttribute::GridGlobal ||
cudaProcAttr.getValue() == cuf::ProcAttribute::HostDevice) {
- rewritePredefinedVars(funcOp.getRegion(), funcOp.getLoc());
+ rewritePredefinedVars(funcOp.getRegion());
rewrittenWholeFunction = true;
}
}
@@ -185,10 +167,10 @@ struct CUFPredefinedVarToGPU
// Host functions containing cuf.kernel or OpenACC compute regions can
// still carry predefined vars in the kernel body. Rewrite them in-place.
funcOp.walk([&](cuf::KernelOp kernelOp) {
- rewritePredefinedVars(kernelOp.getRegion(), kernelOp.getLoc());
+ rewritePredefinedVars(kernelOp.getRegion());
});
funcOp.walk([&](mlir::acc::ComputeRegionOpInterface computeOp) {
- rewritePredefinedVars(computeOp->getRegion(0), computeOp->getLoc());
+ rewritePredefinedVars(computeOp->getRegion(0));
});
}
};
diff --git a/flang/test/Fir/CUDA/predefined-variables.mlir b/flang/test/Fir/CUDA/predefined-variables.mlir
index 0df0c43c5ea76..3b63128a260d0 100644
--- a/flang/test/Fir/CUDA/predefined-variables.mlir
+++ b/flang/test/Fir/CUDA/predefined-variables.mlir
@@ -1,5 +1,6 @@
// RUN: fir-opt --split-input-file --cuf-predefined-var-to-gpu --canonicalize %s | FileCheck %s
// RUN: fir-opt --split-input-file --cuf-predefined-var-to-gpu --canonicalize %s | fir-opt --cuf-predefined-var-to-gpu --canonicalize | FileCheck %s
+// RUN: fir-opt --split-input-file --cuf-predefined-var-to-gpu --mlir-print-debuginfo --mlir-print-local-scope %s | FileCheck %s --check-prefix=LOC
// attributes(device) subroutine sub1(i)
// integer :: i
@@ -50,26 +51,31 @@ func.func @_QPsub1(%arg0: !fir.ref<i32> {fir.bindc_name = "i", cuf.data_attr = #
return
}
+// Each read is created where the variable is read, so the reads appear in use
+// order and the two reads of threadidx%x are distinct values.
+
// CHECK-LABEL: func.func @_QPsub1
// CHECK: %[[WARPSIZE:.*]] = arith.constant 32 : i32
+// CHECK: %[[I:.*]] = fir.declare %{{.*}} {uniq_name = "_QFsub1Ei"} : (!fir.ref<i32>) -> !fir.ref<i32>
+
// CHECK: %[[BASE_THREAD_ID_X:.*]] = nvvm.read.ptx.sreg.tid.x : i32
// CHECK: %[[THREAD_ID_X:.*]] = arith.addi %[[BASE_THREAD_ID_X]], %c1{{.*}} : i32
-// CHECK: %[[BASE_BLOCK_ID_X:.*]] = nvvm.read.ptx.sreg.ctaid.x : i32
-// CHECK: %[[BLOCK_ID_X:.*]] = arith.addi %[[BASE_BLOCK_ID_X]], %c1{{.*}} : i32
-// CHECK: %[[GRID_DIM_Y:.*]] = nvvm.read.ptx.sreg.nctaid.y : i32
-// CHECK: %[[BLOCK_DIM_X:.*]] = nvvm.read.ptx.sreg.ntid.x : i32
-
-// CHECK: %[[I:.*]] = fir.declare %{{.*}} {uniq_name = "_QFsub1Ei"} : (!fir.ref<i32>) -> !fir.ref<i32>
// CHECK: fir.store %[[THREAD_ID_X]] to %[[I]] : !fir.ref<i32>
+// CHECK: %[[BLOCK_DIM_X:.*]] = nvvm.read.ptx.sreg.ntid.x : i32
// CHECK: fir.store %[[BLOCK_DIM_X]] to %[[I]] : !fir.ref<i32>
+// CHECK: %[[BASE_BLOCK_ID_X:.*]] = nvvm.read.ptx.sreg.ctaid.x : i32
+// CHECK: %[[BLOCK_ID_X:.*]] = arith.addi %[[BASE_BLOCK_ID_X]], %c1{{.*}} : i32
// CHECK: fir.store %[[BLOCK_ID_X]] to %[[I]] : !fir.ref<i32>
+// CHECK: %[[GRID_DIM_Y:.*]] = nvvm.read.ptx.sreg.nctaid.y : i32
// CHECK: fir.store %[[GRID_DIM_Y]] to %[[I]] : !fir.ref<i32>
// CHECK: fir.store %[[WARPSIZE]] to %[[I]] : !fir.ref<i32>
-// CHECK: %[[CMP:.*]] = arith.cmpi eq, %[[THREAD_ID_X]], %c0{{.*}} : i32
+// CHECK: %[[BASE_THREAD_ID_X2:.*]] = nvvm.read.ptx.sreg.tid.x : i32
+// CHECK: %[[THREAD_ID_X2:.*]] = arith.addi %[[BASE_THREAD_ID_X2]], %c1{{.*}} : i32
+// CHECK: %[[CMP:.*]] = arith.cmpi eq, %[[THREAD_ID_X2]], %c0{{.*}} : i32
// CHECK: fir.if %[[CMP]] {
// CHECK: fir.store %c0{{.*}} to %[[I]] : !fir.ref<i32>
// CHECK: }
@@ -129,10 +135,10 @@ func.func @_QPsub1(%arg0: !fir.ref<i32> {fir.bindc_name = "i", cuf.data_attr = #
// CHECK: %[[BASE_THREAD_ID_X:.*]] = nvvm.read.ptx.sreg.tid.x : i32
// CHECK: %{{.*}} = arith.addi %[[BASE_THREAD_ID_X]], %c1{{.*}} : i32
+// CHECK: %{{.*}} = nvvm.read.ptx.sreg.ntid.x : i32
// CHECK: %[[BASE_BLOCK_ID_X:.*]] = nvvm.read.ptx.sreg.ctaid.x : i32
// CHECK: %{{.*}} = arith.addi %[[BASE_BLOCK_ID_X]], %c1{{.*}} : i32
// CHECK: %{{.*}} = nvvm.read.ptx.sreg.nctaid.y : i32
-// CHECK: %{{.*}} = nvvm.read.ptx.sreg.ntid.x : i32
// -----
@@ -180,12 +186,15 @@ func.func @_QPsub1(%arg0: !fir.ref<i32> {fir.bindc_name = "i", cuf.data_attr = #
// CHECK: %{{.*}} = arith.constant 32 : i32
+// The read returned by the function is not touched by the pass.
+// CHECK: %{{.*}} = nvvm.read.ptx.sreg.tid.x : i32
+
// CHECK: %[[BASE_THREAD_ID_X:.*]] = nvvm.read.ptx.sreg.tid.x : i32
// CHECK: %{{.*}} = arith.addi %[[BASE_THREAD_ID_X]], %c1{{.*}} : i32
+// CHECK: %{{.*}} = nvvm.read.ptx.sreg.ntid.x : i32
// CHECK: %[[BASE_BLOCK_ID_X:.*]] = nvvm.read.ptx.sreg.ctaid.x : i32
// CHECK: %{{.*}} = arith.addi %[[BASE_BLOCK_ID_X]], %c1{{.*}} : i32
// CHECK: %{{.*}} = nvvm.read.ptx.sreg.nctaid.y : i32
-// CHECK: %{{.*}} = nvvm.read.ptx.sreg.ntid.x : i32
// -----
@@ -202,11 +211,12 @@ func.func @_QMbarPgfoo(%arg0: !fir.ref<i32> {cuf.data_attr = #cuf.cuda<device>,
return
}
+// The only use is inside the fir.if, so the read is created there.
// CHECK-LABEL: func.func @_QMbarPgfoo
-// CHECK: %[[THREAD_ID_X:.*]] = nvvm.read.ptx.sreg.tid.x : i32
-// CHECK: %[[ADD:.*]] = arith.addi %[[THREAD_ID_X]], %c1_i32 : i32
// CHECK: fir.if
-// CHECK: fir.store %[[ADD]] to %{{.*}} : !fir.ref<i32>
+// CHECK: %[[THREAD_ID_X:.*]] = nvvm.read.ptx.sreg.tid.x : i32
+// CHECK: %[[ADD:.*]] = arith.addi %[[THREAD_ID_X]], %c1_i32 : i32
+// CHECK: fir.store %[[ADD]] to %{{.*}} : !fir.ref<i32>
// -----
@@ -233,7 +243,9 @@ func.func @_QMbarPgfoo2(%arg0: !fir.ref<i32> {cuf.data_attr = #cuf.cuda<device>,
// CHECK: %[[ADD:.*]] = arith.addi %[[THREAD_ID_X]], %c1_i32 : i32
// CHECK: fir.store %[[ADD]] to %{{.*}} : !fir.ref<i32>
// CHECK: fir.if
-// CHECK: fir.store %[[ADD]] to %{{.*}} : !fir.ref<i32>
+// CHECK: %[[THREAD_ID_X2:.*]] = nvvm.read.ptx.sreg.tid.x : i32
+// CHECK: %[[ADD2:.*]] = arith.addi %[[THREAD_ID_X2]], %c1_i32 : i32
+// CHECK: fir.store %[[ADD2]] to %{{.*}} : !fir.ref<i32>
// -----
@@ -443,4 +455,39 @@ func.func @_QMdevmodPkernel() attributes {cuf.proc_attr = #cuf.cuda_proc<global>
// CHECK: %[[TID:.*]] = nvvm.read.ptx.sreg.tid.x : i32
// CHECK: %[[ADD:.*]] = arith.addi %[[TID]], %c1{{.*}} : i32
// CHECK: fir.store %[[ADD]] to %{{.*}} : !fir.ref<i32>
-// CHECK: fir.store %[[ADD]] to %{{.*}} : !fir.ref<i32>
+// CHECK: %[[TID2:.*]] = nvvm.read.ptx.sreg.tid.x : i32
+// CHECK: %[[ADD2:.*]] = arith.addi %[[TID2]], %c1{{.*}} : i32
+// CHECK: fir.store %[[ADD2]] to %{{.*}} : !fir.ref<i32>
+
+// -----
+
+// attributes(global) subroutine sub4(i)
+// integer :: i
+// i = blockidx%x
+// i = blockidx%y
+// end subroutine
+
+// Each read is attributed to the line of the use it comes from.
+func.func @_QPsub4(%arg0: !fir.ref<i32> {fir.bindc_name = "i", cuf.data_attr = #cuf.cuda<device>}) attributes {cuf.proc_attr = #cuf.cuda_proc<global>} {
+ %0 = fir.address_of(@_QM__fortran_builtinsE__builtin_blockidx) : !fir.ref<!fir.type<_QM__fortran_builtinsT__builtin_dim3{x:i32,y:i32,z:i32}>> loc(#loc1)
+ %1 = fir.declare %0 {uniq_name = "_QM__fortran_builtinsE__builtin_blockidx"} : (!fir.ref<!fir.type<_QM__fortran_builtinsT__builtin_dim3{x:i32,y:i32,z:i32}>>) -> !fir.ref<!fir.type<_QM__fortran_builtinsT__builtin_dim3{x:i32,y:i32,z:i32}>> loc(#loc1)
+ %2 = fir.declare %arg0 {uniq_name = "_QFsub4Ei"} : (!fir.ref<i32>) -> !fir.ref<i32> loc(#loc1)
+ %3 = fir.coordinate_of %1, x : (!fir.ref<!fir.type<_QM__fortran_builtinsT__builtin_dim3{x:i32,y:i32,z:i32}>>) -> !fir.ref<i32> loc(#loc2)
+ %4 = fir.load %3 : !fir.ref<i32> loc(#loc2)
+ fir.store %4 to %2 : !fir.ref<i32> loc(#loc2)
+ %5 = fir.coordinate_of %1, y : (!fir.ref<!fir.type<_QM__fortran_builtinsT__builtin_dim3{x:i32,y:i32,z:i32}>>) -> !fir.ref<i32> loc(#loc3)
+ %6 = fir.load %5 : !fir.ref<i32> loc(#loc3)
+ fir.store %6 to %2 : !fir.ref<i32> loc(#loc3)
+ return loc(#loc4)
+} loc(#loc1)
+
+#loc1 = loc("sub4.cuf":1:1)
+#loc2 = loc("sub4.cuf":3:3)
+#loc3 = loc("sub4.cuf":4:3)
+#loc4 = loc("sub4.cuf":5:1)
+
+// LOC-LABEL: func.func @_QPsub4
+// LOC: nvvm.read.ptx.sreg.ctaid.x : i32 loc("sub4.cuf":3:3)
+// LOC: arith.addi {{.*}} : i32 loc("sub4.cuf":3:3)
+// LOC: nvvm.read.ptx.sreg.ctaid.y : i32 loc("sub4.cuf":4:3)
+// LOC: arith.addi {{.*}} : i32 loc("sub4.cuf":4:3)
More information about the flang-commits
mailing list