Skip to content

Commit e93d226

Browse files
clementvalclementval
authored
[flang][cuda] Update CompilerGeneratedNames pass to work on gpu module (#120660)
- Update `CompilerGeneratedNames` so it can perform renaming in gpu.module - Update Codegen so it look in the correct module for the type descriptor.
1 parent 4e3c0bb commit e93d226

File tree

3 files changed

+68
-30
lines changed

3 files changed

+68
-30
lines changed

Diff for: ‎flang/lib/Optimizer/CodeGen/CodeGen.cpp

+27-14
Original file line numberDiff line numberDiff line change
@@ -1341,10 +1341,10 @@ struct EmboxCommonConversion : public fir::FIROpConversion<OP> {
13411341

13421342
/// Get the address of the type descriptor global variable that was created by
13431343
/// lowering for derived type \p recType.
1344-
mlir::Value getTypeDescriptor(mlir::ModuleOp mod,
1345-
mlir::ConversionPatternRewriter &rewriter,
1346-
mlir::Location loc,
1347-
fir::RecordType recType) const {
1344+
template <typename ModOpTy>
1345+
mlir::Value
1346+
getTypeDescriptor(ModOpTy mod, mlir::ConversionPatternRewriter &rewriter,
1347+
mlir::Location loc, fir::RecordType recType) const {
13481348
std::string name =
13491349
this->options.typeDescriptorsRenamedForAssembly
13501350
? fir::NameUniquer::getTypeDescriptorAssemblyName(recType.getName())
@@ -1369,7 +1369,8 @@ struct EmboxCommonConversion : public fir::FIROpConversion<OP> {
13691369
return rewriter.create<mlir::LLVM::ZeroOp>(loc, llvmPtrTy);
13701370
}
13711371

1372-
mlir::Value populateDescriptor(mlir::Location loc, mlir::ModuleOp mod,
1372+
template <typename ModOpTy>
1373+
mlir::Value populateDescriptor(mlir::Location loc, ModOpTy mod,
13731374
fir::BaseBoxType boxTy, mlir::Type inputType,
13741375
mlir::ConversionPatternRewriter &rewriter,
13751376
unsigned rank, mlir::Value eleSize,
@@ -1508,10 +1509,16 @@ struct EmboxCommonConversion : public fir::FIROpConversion<OP> {
15081509
extraField =
15091510
this->getExtraFromBox(loc, sourceBoxTyPair, sourceBox, rewriter);
15101511
}
1511-
auto mod = box->template getParentOfType<mlir::ModuleOp>();
1512-
mlir::Value descriptor =
1513-
populateDescriptor(loc, mod, boxTy, inputType, rewriter, rank, eleSize,
1514-
cfiTy, typeDesc, allocatorIdx, extraField);
1512+
1513+
mlir::Value descriptor;
1514+
if (auto gpuMod = box->template getParentOfType<mlir::gpu::GPUModuleOp>())
1515+
descriptor = populateDescriptor(loc, gpuMod, boxTy, inputType, rewriter,
1516+
rank, eleSize, cfiTy, typeDesc,
1517+
allocatorIdx, extraField);
1518+
else if (auto mod = box->template getParentOfType<mlir::ModuleOp>())
1519+
descriptor = populateDescriptor(loc, mod, boxTy, inputType, rewriter,
1520+
rank, eleSize, cfiTy, typeDesc,
1521+
allocatorIdx, extraField);
15151522

15161523
return {boxTy, descriptor, eleSize};
15171524
}
@@ -1554,11 +1561,17 @@ struct EmboxCommonConversion : public fir::FIROpConversion<OP> {
15541561
mlir::Value extraField =
15551562
this->getExtraFromBox(loc, inputBoxTyPair, loweredBox, rewriter);
15561563

1557-
auto mod = box->template getParentOfType<mlir::ModuleOp>();
1558-
mlir::Value descriptor =
1559-
populateDescriptor(loc, mod, boxTy, box.getBox().getType(), rewriter,
1560-
rank, eleSize, cfiTy, typeDesc,
1561-
/*allocatorIdx=*/kDefaultAllocator, extraField);
1564+
mlir::Value descriptor;
1565+
if (auto gpuMod = box->template getParentOfType<mlir::gpu::GPUModuleOp>())
1566+
descriptor =
1567+
populateDescriptor(loc, gpuMod, boxTy, box.getBox().getType(),
1568+
rewriter, rank, eleSize, cfiTy, typeDesc,
1569+
/*allocatorIdx=*/kDefaultAllocator, extraField);
1570+
else if (auto mod = box->template getParentOfType<mlir::ModuleOp>())
1571+
descriptor =
1572+
populateDescriptor(loc, mod, boxTy, box.getBox().getType(), rewriter,
1573+
rank, eleSize, cfiTy, typeDesc,
1574+
/*allocatorIdx=*/kDefaultAllocator, extraField);
15621575

15631576
return {boxTy, descriptor, eleSize};
15641577
}

Diff for: ‎flang/lib/Optimizer/Transforms/CompilerGeneratedNames.cpp

+24-16
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
#include "flang/Optimizer/Dialect/FIROpsSupport.h"
1212
#include "flang/Optimizer/Support/InternalNames.h"
1313
#include "flang/Optimizer/Transforms/Passes.h"
14+
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
1415
#include "mlir/IR/Attributes.h"
1516
#include "mlir/IR/SymbolTable.h"
1617
#include "mlir/Pass/Pass.h"
@@ -42,24 +43,31 @@ void CompilerGeneratedNamesConversionPass::runOnOperation() {
4243
auto *context = &getContext();
4344

4445
llvm::DenseMap<mlir::StringAttr, mlir::FlatSymbolRefAttr> remappings;
45-
for (auto &funcOrGlobal : op->getRegion(0).front()) {
46-
if (llvm::isa<mlir::func::FuncOp>(funcOrGlobal) ||
47-
llvm::isa<fir::GlobalOp>(funcOrGlobal)) {
48-
auto symName = funcOrGlobal.getAttrOfType<mlir::StringAttr>(
49-
mlir::SymbolTable::getSymbolAttrName());
50-
auto deconstructedName = fir::NameUniquer::deconstruct(symName);
51-
if (deconstructedName.first != fir::NameUniquer::NameKind::NOT_UNIQUED &&
52-
!fir::NameUniquer::isExternalFacingUniquedName(deconstructedName)) {
53-
std::string newName =
54-
fir::NameUniquer::replaceSpecialSymbols(symName.getValue().str());
55-
if (newName != symName) {
56-
auto newAttr = mlir::StringAttr::get(context, newName);
57-
mlir::SymbolTable::setSymbolName(&funcOrGlobal, newAttr);
58-
auto newSymRef = mlir::FlatSymbolRefAttr::get(newAttr);
59-
remappings.try_emplace(symName, newSymRef);
60-
}
46+
47+
auto processOp = [&](mlir::Operation &op) {
48+
auto symName = op.getAttrOfType<mlir::StringAttr>(
49+
mlir::SymbolTable::getSymbolAttrName());
50+
auto deconstructedName = fir::NameUniquer::deconstruct(symName);
51+
if (deconstructedName.first != fir::NameUniquer::NameKind::NOT_UNIQUED &&
52+
!fir::NameUniquer::isExternalFacingUniquedName(deconstructedName)) {
53+
std::string newName =
54+
fir::NameUniquer::replaceSpecialSymbols(symName.getValue().str());
55+
if (newName != symName) {
56+
auto newAttr = mlir::StringAttr::get(context, newName);
57+
mlir::SymbolTable::setSymbolName(&op, newAttr);
58+
auto newSymRef = mlir::FlatSymbolRefAttr::get(newAttr);
59+
remappings.try_emplace(symName, newSymRef);
6160
}
6261
}
62+
};
63+
for (auto &op : op->getRegion(0).front()) {
64+
if (llvm::isa<mlir::func::FuncOp>(op) || llvm::isa<fir::GlobalOp>(op))
65+
processOp(op);
66+
else if (auto gpuMod = mlir::dyn_cast<mlir::gpu::GPUModuleOp>(&op))
67+
for (auto &op : gpuMod->getRegion(0).front())
68+
if (llvm::isa<mlir::func::FuncOp>(op) || llvm::isa<fir::GlobalOp>(op) ||
69+
llvm::isa<mlir::gpu::GPUFuncOp>(op))
70+
processOp(op);
6371
}
6472

6573
if (remappings.empty())

Diff for: ‎flang/test/Fir/CUDA/cuda-compiler-generated-names.mlir

+17
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
// RUN: fir-opt --split-input-file --compiler-generated-names --fir-to-llvm-ir="target=x86_64-unknown-linux-gnu type-descriptors-renamed-for-assembly=true" %s | FileCheck %s
2+
3+
module @mod1 attributes {gpu.container} {
4+
gpu.module @gpu1 {
5+
fir.global linkonce @_QMtest_dinitE.dt.tseq constant : i8
6+
7+
func.func @embox1(%arg0: !fir.ref<!fir.type<_QMtest_dinitTtseq{i:i32}>>) {
8+
%0 = fir.embox %arg0() : (!fir.ref<!fir.type<_QMtest_dinitTtseq{i:i32}>>) -> !fir.box<!fir.type<_QMtest_dinitTtseq{i:i32}>>
9+
return
10+
}
11+
}
12+
}
13+
14+
// CHECK-LABEL: gpu.module @gpu1
15+
// CHECK: llvm.mlir.global linkonce constant @_QMtest_dinitEXdtXtseq
16+
// CHECK: llvm.mlir.addressof @_QMtest_dinitEXdtXtseq : !llvm.ptr
17+

0 commit comments

Comments
 (0)