Skip to content

Commit 5bf9f4b

Browse files
agozillonagozillon
committed
[Flang][MLIR] Extend DataLAyout utilities to have basic GPU Module support
As there is now certain areas where we now have the possibility of having either a ModuleOp or GPUModuleOp and both of these modules can have DataLayout's and we may require utilising the DataLayout utilities in these areas I've taken the liberty of trying to extend them for use with both. Those with more knowledge of how they wish the GPUModuleOp's to interact with their parent ModuleOp's DataLayout may have further alterations they wish to make in the future, but for the moment, it'll simply utilise the basic data layout construction which I believe combines parent and child datalayouts from the ModuleOp and GPUModuleOp. If there is no GPUModuleOp DataLayout it should default to the parent ModuleOp. It's worth noting there is some weirdness if you have two module operations defining builtin dialect DataLayout Entries, it appears the combinatorial functionality for DataLayouts doesn't support the merging of these. This behaviour is useful for areas like: https://github.com/llvm/llvm-project/pull/119585/files#diff-19fc4bcb38829d085e25d601d344bbd85bf7ef749ca359e348f4a7c750eae89dR1412 where we have a crossroads between the two different module operations.
1 parent c805df6 commit 5bf9f4b

File tree

10 files changed

+79
-25
lines changed

10 files changed

+79
-25
lines changed

‎flang/include/flang/Optimizer/Support/DataLayout.h‎

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,14 @@
1818

1919
namespace mlir {
2020
class ModuleOp;
21-
}
21+
namespace gpu {
22+
class GPUModuleOp;
23+
} // namespace gpu
24+
} // namespace mlir
25+
2226
namespace llvm {
2327
class DataLayout;
24-
}
28+
} // namespace llvm
2529

2630
namespace fir::support {
2731
/// Create an mlir::DataLayoutSpecInterface attribute from an llvm::DataLayout
@@ -30,21 +34,30 @@ namespace fir::support {
3034
/// the llvm::DataLayout on the module.
3135
/// These attributes are replaced if they were already set.
3236
void setMLIRDataLayout(mlir::ModuleOp mlirModule, const llvm::DataLayout &dl);
37+
void setMLIRDataLayout(mlir::gpu::GPUModuleOp mlirModule,
38+
const llvm::DataLayout &dl);
3339

3440
/// Create an mlir::DataLayoutSpecInterface from the llvm.data_layout attribute
3541
/// if one is provided. If such attribute is not available, create a default
3642
/// target independent layout when allowDefaultLayout is true. Otherwise do
3743
/// nothing.
3844
void setMLIRDataLayoutFromAttributes(mlir::ModuleOp mlirModule,
3945
bool allowDefaultLayout);
46+
void setMLIRDataLayoutFromAttributes(mlir::gpu::GPUModuleOp mlirModule,
47+
bool allowDefaultLayout);
4048

4149
/// Create mlir::DataLayout from the data layout information on the
4250
/// mlir::Module. Creates the data layout information attributes with
4351
/// setMLIRDataLayoutFromAttributes if the DLTI attribute is not yet set. If no
4452
/// information is present at all and \p allowDefaultLayout is false, returns
4553
/// std::nullopt.
4654
std::optional<mlir::DataLayout>
47-
getOrSetDataLayout(mlir::ModuleOp mlirModule, bool allowDefaultLayout = false);
55+
getOrSetMLIRDataLayout(mlir::ModuleOp mlirModule,
56+
bool allowDefaultLayout = false);
57+
std::optional<mlir::DataLayout>
58+
getOrSetMLIRDataLayout(mlir::gpu::GPUModuleOp mlirModule,
59+
bool allowDefaultLayout = false);
60+
4861
} // namespace fir::support
4962

5063
#endif // FORTRAN_OPTIMIZER_SUPPORT_DATALAYOUT_H

‎flang/lib/Optimizer/CodeGen/CodeGen.cpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1190,7 +1190,7 @@ genCUFAllocDescriptor(mlir::Location loc,
11901190
mlir::ModuleOp mod, fir::BaseBoxType boxTy,
11911191
const fir::LLVMTypeConverter &typeConverter) {
11921192
std::optional<mlir::DataLayout> dl =
1193-
fir::support::getOrSetDataLayout(mod, /*allowDefaultLayout=*/true);
1193+
fir::support::getOrSetMLIRDataLayout(mod, /*allowDefaultLayout=*/true);
11941194
if (!dl)
11951195
mlir::emitError(mod.getLoc(),
11961196
"module operation must carry a data layout attribute "
@@ -3908,7 +3908,7 @@ class FIRToLLVMLowering
39083908
return signalPassFailure();
39093909

39103910
std::optional<mlir::DataLayout> dl =
3911-
fir::support::getOrSetDataLayout(mod, /*allowDefaultLayout=*/true);
3911+
fir::support::getOrSetMLIRDataLayout(mod, /*allowDefaultLayout=*/true);
39123912
if (!dl) {
39133913
mlir::emitError(mod.getLoc(),
39143914
"module operation must carry a data layout attribute "

‎flang/lib/Optimizer/CodeGen/TargetRewrite.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -107,7 +107,7 @@ class TargetRewrite : public fir::impl::TargetRewritePassBase<TargetRewrite> {
107107
// TargetRewrite will require querying the type storage sizes, if it was
108108
// not set already, create a DataLayoutSpec for the ModuleOp now.
109109
std::optional<mlir::DataLayout> dl =
110-
fir::support::getOrSetDataLayout(mod, /*allowDefaultLayout=*/true);
110+
fir::support::getOrSetMLIRDataLayout(mod, /*allowDefaultLayout=*/true);
111111
if (!dl) {
112112
mlir::emitError(mod.getLoc(),
113113
"module operation must carry a data layout attribute "

‎flang/lib/Optimizer/Support/DataLayout.cpp‎

Lines changed: 51 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
#include "flang/Optimizer/Dialect/Support/FIRContext.h"
1111
#include "flang/Optimizer/Support/FatalError.h"
1212
#include "mlir/Dialect/DLTI/DLTI.h"
13+
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
1314
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
1415
#include "mlir/IR/BuiltinOps.h"
1516
#include "mlir/Interfaces/DataLayoutInterfaces.h"
@@ -20,8 +21,9 @@
2021
#include "llvm/Support/TargetSelect.h"
2122
#include "llvm/Target/TargetMachine.h"
2223

23-
void fir::support::setMLIRDataLayout(mlir::ModuleOp mlirModule,
24-
const llvm::DataLayout &dl) {
24+
namespace {
25+
template <typename ModOpTy>
26+
void setDataLayout(ModOpTy mlirModule, const llvm::DataLayout &dl) {
2527
mlir::MLIRContext *context = mlirModule.getContext();
2628
mlirModule->setAttr(
2729
mlir::LLVM::LLVMDialect::getDataLayoutAttrName(),
@@ -30,12 +32,13 @@ void fir::support::setMLIRDataLayout(mlir::ModuleOp mlirModule,
3032
mlirModule->setAttr(mlir::DLTIDialect::kDataLayoutAttrName, dlSpec);
3133
}
3234

33-
void fir::support::setMLIRDataLayoutFromAttributes(mlir::ModuleOp mlirModule,
34-
bool allowDefaultLayout) {
35+
template <typename ModOpTy>
36+
void setDataLayoutFromAttributes(ModOpTy mlirModule, bool allowDefaultLayout) {
3537
if (mlirModule.getDataLayoutSpec())
3638
return; // Already set.
37-
if (auto dataLayoutString = mlirModule->getAttrOfType<mlir::StringAttr>(
38-
mlir::LLVM::LLVMDialect::getDataLayoutAttrName())) {
39+
if (auto dataLayoutString =
40+
mlirModule->template getAttrOfType<mlir::StringAttr>(
41+
mlir::LLVM::LLVMDialect::getDataLayoutAttrName())) {
3942
llvm::DataLayout llvmDataLayout(dataLayoutString);
4043
fir::support::setMLIRDataLayout(mlirModule, llvmDataLayout);
4144
return;
@@ -46,15 +49,53 @@ void fir::support::setMLIRDataLayoutFromAttributes(mlir::ModuleOp mlirModule,
4649
fir::support::setMLIRDataLayout(mlirModule, llvmDataLayout);
4750
}
4851

49-
std::optional<mlir::DataLayout>
50-
fir::support::getOrSetDataLayout(mlir::ModuleOp mlirModule,
51-
bool allowDefaultLayout) {
52+
template <typename ModOpTy>
53+
std::optional<mlir::DataLayout> getOrSetDataLayout(ModOpTy mlirModule,
54+
bool allowDefaultLayout) {
5255
if (!mlirModule.getDataLayoutSpec()) {
5356
fir::support::setMLIRDataLayoutFromAttributes(mlirModule,
5457
allowDefaultLayout);
55-
if (!mlirModule.getDataLayoutSpec()) {
58+
// if it is a GPU module, we let it proceed, as it's contained within
59+
// a module, its parent may have a DataLayout that can take its
60+
// place.
61+
if (!mlirModule.getDataLayoutSpec() &&
62+
!mlir::isa<mlir::gpu::GPUModuleOp>(mlirModule)) {
5663
return std::nullopt;
5764
}
5865
}
5966
return mlir::DataLayout(mlirModule);
6067
}
68+
69+
} // namespace
70+
71+
void fir::support::setMLIRDataLayout(mlir::ModuleOp mlirModule,
72+
const llvm::DataLayout &dl) {
73+
setDataLayout(mlirModule, dl);
74+
}
75+
76+
void fir::support::setMLIRDataLayout(mlir::gpu::GPUModuleOp mlirModule,
77+
const llvm::DataLayout &dl) {
78+
setDataLayout(mlirModule, dl);
79+
}
80+
81+
void fir::support::setMLIRDataLayoutFromAttributes(mlir::ModuleOp mlirModule,
82+
bool allowDefaultLayout) {
83+
setDataLayoutFromAttributes(mlirModule, allowDefaultLayout);
84+
}
85+
86+
void fir::support::setMLIRDataLayoutFromAttributes(
87+
mlir::gpu::GPUModuleOp mlirModule, bool allowDefaultLayout) {
88+
setDataLayoutFromAttributes(mlirModule, allowDefaultLayout);
89+
}
90+
91+
std::optional<mlir::DataLayout>
92+
fir::support::getOrSetMLIRDataLayout(mlir::ModuleOp mlirModule,
93+
bool allowDefaultLayout) {
94+
return getOrSetDataLayout(mlirModule, allowDefaultLayout);
95+
}
96+
97+
std::optional<mlir::DataLayout>
98+
fir::support::getOrSetMLIRDataLayout(mlir::gpu::GPUModuleOp mlirModule,
99+
bool allowDefaultLayout) {
100+
return getOrSetDataLayout(mlirModule, allowDefaultLayout);
101+
}

‎flang/lib/Optimizer/Transforms/AddDebugInfo.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -418,7 +418,7 @@ void AddDebugInfoPass::runOnOperation() {
418418
llvm::StringRef fileName;
419419
std::string filePath;
420420
std::optional<mlir::DataLayout> dl =
421-
fir::support::getOrSetDataLayout(module, /*allowDefaultLayout=*/true);
421+
fir::support::getOrSetMLIRDataLayout(module, /*allowDefaultLayout=*/true);
422422
if (!dl) {
423423
mlir::emitError(module.getLoc(), "Missing data layout attribute in module");
424424
signalPassFailure();

‎flang/lib/Optimizer/Transforms/CUFAddConstructor.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ struct CUFAddConstructor
5757
auto funcTy =
5858
mlir::LLVM::LLVMFunctionType::get(voidTy, {}, /*isVarArg=*/false);
5959
std::optional<mlir::DataLayout> dl =
60-
fir::support::getOrSetDataLayout(mod, /*allowDefaultLayout=*/false);
60+
fir::support::getOrSetMLIRDataLayout(mod, /*allowDefaultLayout=*/false);
6161
if (!dl) {
6262
mlir::emitError(mod.getLoc(),
6363
"data layout attribute is required to perform " +

‎flang/lib/Optimizer/Transforms/CUFGPUToLLVMConversion.cpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -182,8 +182,8 @@ class CUFGPUToLLVMConversion
182182
if (!module)
183183
return signalPassFailure();
184184

185-
std::optional<mlir::DataLayout> dl =
186-
fir::support::getOrSetDataLayout(module, /*allowDefaultLayout=*/false);
185+
std::optional<mlir::DataLayout> dl = fir::support::getOrSetMLIRDataLayout(
186+
module, /*allowDefaultLayout=*/false);
187187
fir::LLVMTypeConverter typeConverter(module, /*applyTBAA=*/false,
188188
/*forceUnifiedTBAATree=*/false, *dl);
189189
cuf::populateCUFGPUToLLVMConversionPatterns(typeConverter, patterns);

‎flang/lib/Optimizer/Transforms/CUFOpConversion.cpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -884,8 +884,8 @@ class CUFOpConversion : public fir::impl::CUFOpConversionBase<CUFOpConversion> {
884884
return signalPassFailure();
885885
mlir::SymbolTable symtab(module);
886886

887-
std::optional<mlir::DataLayout> dl =
888-
fir::support::getOrSetDataLayout(module, /*allowDefaultLayout=*/false);
887+
std::optional<mlir::DataLayout> dl = fir::support::getOrSetMLIRDataLayout(
888+
module, /*allowDefaultLayout=*/false);
889889
fir::LLVMTypeConverter typeConverter(module, /*applyTBAA=*/false,
890890
/*forceUnifiedTBAATree=*/false, *dl);
891891
target.addLegalDialect<fir::FIROpsDialect, mlir::arith::ArithDialect,

‎flang/lib/Optimizer/Transforms/LoopVersioning.cpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -312,8 +312,8 @@ void LoopVersioningPass::runOnOperation() {
312312
mlir::ModuleOp module = func->getParentOfType<mlir::ModuleOp>();
313313
fir::KindMapping kindMap = fir::getKindMapping(module);
314314
mlir::SmallVector<ArgInfo, 4> argsOfInterest;
315-
std::optional<mlir::DataLayout> dl =
316-
fir::support::getOrSetDataLayout(module, /*allowDefaultLayout=*/false);
315+
std::optional<mlir::DataLayout> dl = fir::support::getOrSetMLIRDataLayout(
316+
module, /*allowDefaultLayout=*/false);
317317
if (!dl)
318318
mlir::emitError(module.getLoc(),
319319
"data layout attribute is required to perform " DEBUG_TYPE

‎flang/test/lib/OpenACC/TestOpenACCInterfaces.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ struct TestFIROpenACCInterfaces
2828
void runOnOperation() override {
2929
mlir::ModuleOp mod = getOperation();
3030
auto datalayout =
31-
fir::support::getOrSetDataLayout(mod, /*allowDefaultLayout=*/true);
31+
fir::support::getOrSetMLIRDataLayout(mod, /*allowDefaultLayout=*/true);
3232
mlir::OpBuilder builder(mod);
3333
getOperation().walk([&](Operation *op) {
3434
if (isa<ACC_DATA_ENTRY_OPS>(op)) {

0 commit comments

Comments
 (0)