diff --git a/flang/lib/Optimizer/Builder/IntrinsicCall.cpp b/flang/lib/Optimizer/Builder/IntrinsicCall.cpp index ec0c802fb209c..29eedfb0ce9cd 100644 --- a/flang/lib/Optimizer/Builder/IntrinsicCall.cpp +++ b/flang/lib/Optimizer/Builder/IntrinsicCall.cpp @@ -3516,11 +3516,23 @@ static mlir::Value getAddrFromBox(fir::FirOpBuilder &builder, return addr; } +static void clocDeviceArgRewrite(fir::ExtendedValue arg) { + // Special case for device address in c_loc. + if (auto emboxOp = mlir::dyn_cast_or_null( + fir::getBase(arg).getDefiningOp())) + if (auto declareOp = mlir::dyn_cast_or_null( + emboxOp.getMemref().getDefiningOp())) + if (declareOp.getDataAttr() && + declareOp.getDataAttr() == cuf::DataAttribute::Device) + emboxOp.getMemrefMutable().assign(declareOp.getMemref()); +} + static fir::ExtendedValue genCLocOrCFunLoc(fir::FirOpBuilder &builder, mlir::Location loc, mlir::Type resultType, llvm::ArrayRef args, bool isFunc = false, bool isDevLoc = false) { assert(args.size() == 1); + clocDeviceArgRewrite(args[0]); mlir::Value res = fir::AllocaOp::create(builder, loc, resultType); mlir::Value resAddr; if (isDevLoc) diff --git a/flang/test/Lower/CUDA/cuda-cloc.cuf b/flang/test/Lower/CUDA/cuda-cloc.cuf new file mode 100644 index 0000000000000..87a98d47f4bfe --- /dev/null +++ b/flang/test/Lower/CUDA/cuda-cloc.cuf @@ -0,0 +1,19 @@ +! RUN: bbc -emit-hlfir -fcuda %s -o - | FileCheck %s + +module symbols + integer(4), device, target :: sdev(100) +end module + +subroutine sub1 + use iso_c_binding + use symbols + print*, c_loc(sdev) +end subroutine + +! CHECK-LABEL: func.func @_QPsub1() +! CHECK: %[[ADDR:.*]] = fir.address_of(@_QMsymbolsEsdev) : !fir.ref> +! CHECK: %[[EMBOX:.*]] = fir.embox %[[ADDR]](%{{.*}}) : (!fir.ref>, !fir.shape<1>) -> !fir.box> +! CHECK: %[[__ADDRESS:.*]] = fir.coordinate_of %{{.*}}, __address : (!fir.ref>) -> !fir.ref +! CHECK: %[[BOX_ADDR:.*]] = fir.box_addr %[[EMBOX]] : (!fir.box>) -> !fir.ref> +! CHECK: %[[CONV:.*]] = fir.convert %[[BOX_ADDR]] : (!fir.ref>) -> i64 +! CHECK: fir.store %[[CONV]] to %[[__ADDRESS]] : !fir.ref