shylie/llvm-project
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[flang][cuda] cuf.alloc in device context should be converted to fir.alloc (#116110)

Browse Source 
Update `inDeviceContext` to account for the gpu.func operation.
This commit is contained in:
Valentin Clement (バレンタイン クレメン) 2024-11-13 14:57:42 -08:00 committed by GitHub
parent fd2e4004cd
commit ec066d30e2
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 12 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
flang/lib/Optimizer/Transforms/CUFOpConversion.cpp
View File

@ -251,6 +251,8 @@ struct CUFDeallocateOpConversion
static bool inDeviceContext(mlir::Operation *op) {
if (op->getParentOfType<cuf::KernelOp>())
return true;
if (auto funcOp = op->getParentOfType<mlir::gpu::GPUFuncOp>())
return true;
if (auto funcOp = op->getParentOfType<mlir::func::FuncOp>()) {
if (auto cudaProcAttr =
funcOp.getOperation()->getAttrOfType<cuf::ProcAttributeAttr>(

10
flang/test/Fir/CUDA/cuda-alloc-free.fir
View File

@ -73,4 +73,14 @@ func.func @_QPtest_type() {
// CHECK: %[[CONV_BYTES:.*]] = fir.convert %[[BYTES]] : (index) -> i64
// CHECK: fir.call @_FortranACUFMemAlloc(%[[CONV_BYTES]], %c0{{.*}}, %{{.*}}, %{{.*}}) : (i64, i32, !fir.ref<i8>, i32) -> !fir.llvm_ptr<i8>
gpu.module @cuda_device_mod [#nvvm.target] {
gpu.func @_QMalloc() kernel {
%0 = cuf.alloc !fir.box<!fir.heap<!fir.array<?xf32>>> {bindc_name = "a", data_attr = #cuf.cuda<device>, uniq_name = "_QMallocEa"} -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
gpu.return
}
}
// CHECK-LABEL: gpu.func @_QMalloc() kernel
// CHECK: fir.alloca !fir.box<!fir.heap<!fir.array<?xf32>>> {bindc_name = "a", uniq_name = "_QMallocEa"}
} // end module