Skip to content

[Flang][MLIR] Extend DataLayout utilities to have basic GPU Module support #123149

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jan 30, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 16 additions & 3 deletions flang/include/flang/Optimizer/Support/DataLayout.h
Original file line number Diff line number Diff line change
Expand Up @@ -18,10 +18,14 @@

namespace mlir {
class ModuleOp;
}
namespace gpu {
class GPUModuleOp;
} // namespace gpu
} // namespace mlir

namespace llvm {
class DataLayout;
}
} // namespace llvm

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

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

/// Create mlir::DataLayout from the data layout information on the
/// mlir::Module. Creates the data layout information attributes with
/// setMLIRDataLayoutFromAttributes if the DLTI attribute is not yet set. If no
/// information is present at all and \p allowDefaultLayout is false, returns
/// std::nullopt.
std::optional<mlir::DataLayout>
getOrSetDataLayout(mlir::ModuleOp mlirModule, bool allowDefaultLayout = false);
getOrSetMLIRDataLayout(mlir::ModuleOp mlirModule,
bool allowDefaultLayout = false);
std::optional<mlir::DataLayout>
getOrSetMLIRDataLayout(mlir::gpu::GPUModuleOp mlirModule,
bool allowDefaultLayout = false);

} // namespace fir::support

#endif // FORTRAN_OPTIMIZER_SUPPORT_DATALAYOUT_H
4 changes: 2 additions & 2 deletions flang/lib/Optimizer/CodeGen/CodeGen.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1190,7 +1190,7 @@ genCUFAllocDescriptor(mlir::Location loc,
mlir::ModuleOp mod, fir::BaseBoxType boxTy,
const fir::LLVMTypeConverter &typeConverter) {
std::optional<mlir::DataLayout> dl =
fir::support::getOrSetDataLayout(mod, /*allowDefaultLayout=*/true);
fir::support::getOrSetMLIRDataLayout(mod, /*allowDefaultLayout=*/true);
if (!dl)
mlir::emitError(mod.getLoc(),
"module operation must carry a data layout attribute "
Expand Down Expand Up @@ -3942,7 +3942,7 @@ class FIRToLLVMLowering
return signalPassFailure();

std::optional<mlir::DataLayout> dl =
fir::support::getOrSetDataLayout(mod, /*allowDefaultLayout=*/true);
fir::support::getOrSetMLIRDataLayout(mod, /*allowDefaultLayout=*/true);
if (!dl) {
mlir::emitError(mod.getLoc(),
"module operation must carry a data layout attribute "
Expand Down
2 changes: 1 addition & 1 deletion flang/lib/Optimizer/CodeGen/TargetRewrite.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -107,7 +107,7 @@ class TargetRewrite : public fir::impl::TargetRewritePassBase<TargetRewrite> {
// TargetRewrite will require querying the type storage sizes, if it was
// not set already, create a DataLayoutSpec for the ModuleOp now.
std::optional<mlir::DataLayout> dl =
fir::support::getOrSetDataLayout(mod, /*allowDefaultLayout=*/true);
fir::support::getOrSetMLIRDataLayout(mod, /*allowDefaultLayout=*/true);
if (!dl) {
mlir::emitError(mod.getLoc(),
"module operation must carry a data layout attribute "
Expand Down
65 changes: 51 additions & 14 deletions flang/lib/Optimizer/Support/DataLayout.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
#include "flang/Optimizer/Dialect/Support/FIRContext.h"
#include "flang/Optimizer/Support/FatalError.h"
#include "mlir/Dialect/DLTI/DLTI.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/Interfaces/DataLayoutInterfaces.h"
Expand All @@ -20,8 +21,9 @@
#include "llvm/Support/TargetSelect.h"
#include "llvm/Target/TargetMachine.h"

void fir::support::setMLIRDataLayout(mlir::ModuleOp mlirModule,
const llvm::DataLayout &dl) {
namespace {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: For local functions, use static instead: https://llvm.org/docs/CodingStandards.html#anonymous-namespaces.

template <typename ModOpTy>
static void setDataLayout(ModOpTy mlirModule, const llvm::DataLayout &dl) {
mlir::MLIRContext *context = mlirModule.getContext();
mlirModule->setAttr(
mlir::LLVM::LLVMDialect::getDataLayoutAttrName(),
Expand All @@ -30,12 +32,14 @@ void fir::support::setMLIRDataLayout(mlir::ModuleOp mlirModule,
mlirModule->setAttr(mlir::DLTIDialect::kDataLayoutAttrName, dlSpec);
}

void fir::support::setMLIRDataLayoutFromAttributes(mlir::ModuleOp mlirModule,
bool allowDefaultLayout) {
template <typename ModOpTy>
static void setDataLayoutFromAttributes(ModOpTy mlirModule,
bool allowDefaultLayout) {
if (mlirModule.getDataLayoutSpec())
return; // Already set.
if (auto dataLayoutString = mlirModule->getAttrOfType<mlir::StringAttr>(
mlir::LLVM::LLVMDialect::getDataLayoutAttrName())) {
if (auto dataLayoutString =
mlirModule->template getAttrOfType<mlir::StringAttr>(
mlir::LLVM::LLVMDialect::getDataLayoutAttrName())) {
llvm::DataLayout llvmDataLayout(dataLayoutString);
fir::support::setMLIRDataLayout(mlirModule, llvmDataLayout);
return;
Expand All @@ -46,15 +50,48 @@ void fir::support::setMLIRDataLayoutFromAttributes(mlir::ModuleOp mlirModule,
fir::support::setMLIRDataLayout(mlirModule, llvmDataLayout);
}

std::optional<mlir::DataLayout>
fir::support::getOrSetDataLayout(mlir::ModuleOp mlirModule,
bool allowDefaultLayout) {
if (!mlirModule.getDataLayoutSpec()) {
template <typename ModOpTy>
static std::optional<mlir::DataLayout>
getOrSetDataLayout(ModOpTy mlirModule, bool allowDefaultLayout) {
if (!mlirModule.getDataLayoutSpec())
fir::support::setMLIRDataLayoutFromAttributes(mlirModule,
allowDefaultLayout);
if (!mlirModule.getDataLayoutSpec()) {
return std::nullopt;
}
}
if (!mlirModule.getDataLayoutSpec() &&
!mlir::isa<mlir::gpu::GPUModuleOp>(mlirModule))
return std::nullopt;
return mlir::DataLayout(mlirModule);
}

} // namespace

void fir::support::setMLIRDataLayout(mlir::ModuleOp mlirModule,
const llvm::DataLayout &dl) {
setDataLayout(mlirModule, dl);
}

void fir::support::setMLIRDataLayout(mlir::gpu::GPUModuleOp mlirModule,
const llvm::DataLayout &dl) {
setDataLayout(mlirModule, dl);
}

void fir::support::setMLIRDataLayoutFromAttributes(mlir::ModuleOp mlirModule,
bool allowDefaultLayout) {
setDataLayoutFromAttributes(mlirModule, allowDefaultLayout);
}

void fir::support::setMLIRDataLayoutFromAttributes(
mlir::gpu::GPUModuleOp mlirModule, bool allowDefaultLayout) {
setDataLayoutFromAttributes(mlirModule, allowDefaultLayout);
}

std::optional<mlir::DataLayout>
fir::support::getOrSetMLIRDataLayout(mlir::ModuleOp mlirModule,
bool allowDefaultLayout) {
return getOrSetDataLayout(mlirModule, allowDefaultLayout);
}

std::optional<mlir::DataLayout>
fir::support::getOrSetMLIRDataLayout(mlir::gpu::GPUModuleOp mlirModule,
bool allowDefaultLayout) {
return getOrSetDataLayout(mlirModule, allowDefaultLayout);
}
2 changes: 1 addition & 1 deletion flang/lib/Optimizer/Transforms/AddDebugInfo.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -523,7 +523,7 @@ void AddDebugInfoPass::runOnOperation() {
llvm::StringRef fileName;
std::string filePath;
std::optional<mlir::DataLayout> dl =
fir::support::getOrSetDataLayout(module, /*allowDefaultLayout=*/true);
fir::support::getOrSetMLIRDataLayout(module, /*allowDefaultLayout=*/true);
if (!dl) {
mlir::emitError(module.getLoc(), "Missing data layout attribute in module");
signalPassFailure();
Expand Down
2 changes: 1 addition & 1 deletion flang/lib/Optimizer/Transforms/CUFAddConstructor.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@ struct CUFAddConstructor
auto funcTy =
mlir::LLVM::LLVMFunctionType::get(voidTy, {}, /*isVarArg=*/false);
std::optional<mlir::DataLayout> dl =
fir::support::getOrSetDataLayout(mod, /*allowDefaultLayout=*/false);
fir::support::getOrSetMLIRDataLayout(mod, /*allowDefaultLayout=*/false);
if (!dl) {
mlir::emitError(mod.getLoc(),
"data layout attribute is required to perform " +
Expand Down
4 changes: 2 additions & 2 deletions flang/lib/Optimizer/Transforms/CUFGPUToLLVMConversion.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -188,8 +188,8 @@ class CUFGPUToLLVMConversion
if (!module)
return signalPassFailure();

std::optional<mlir::DataLayout> dl =
fir::support::getOrSetDataLayout(module, /*allowDefaultLayout=*/false);
std::optional<mlir::DataLayout> dl = fir::support::getOrSetMLIRDataLayout(
module, /*allowDefaultLayout=*/false);
fir::LLVMTypeConverter typeConverter(module, /*applyTBAA=*/false,
/*forceUnifiedTBAATree=*/false, *dl);
cuf::populateCUFGPUToLLVMConversionPatterns(typeConverter, patterns);
Expand Down
4 changes: 2 additions & 2 deletions flang/lib/Optimizer/Transforms/CUFOpConversion.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -913,8 +913,8 @@ class CUFOpConversion : public fir::impl::CUFOpConversionBase<CUFOpConversion> {
return signalPassFailure();
mlir::SymbolTable symtab(module);

std::optional<mlir::DataLayout> dl =
fir::support::getOrSetDataLayout(module, /*allowDefaultLayout=*/false);
std::optional<mlir::DataLayout> dl = fir::support::getOrSetMLIRDataLayout(
module, /*allowDefaultLayout=*/false);
fir::LLVMTypeConverter typeConverter(module, /*applyTBAA=*/false,
/*forceUnifiedTBAATree=*/false, *dl);
target.addLegalDialect<fir::FIROpsDialect, mlir::arith::ArithDialect,
Expand Down
4 changes: 2 additions & 2 deletions flang/lib/Optimizer/Transforms/LoopVersioning.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -312,8 +312,8 @@ void LoopVersioningPass::runOnOperation() {
mlir::ModuleOp module = func->getParentOfType<mlir::ModuleOp>();
fir::KindMapping kindMap = fir::getKindMapping(module);
mlir::SmallVector<ArgInfo, 4> argsOfInterest;
std::optional<mlir::DataLayout> dl =
fir::support::getOrSetDataLayout(module, /*allowDefaultLayout=*/false);
std::optional<mlir::DataLayout> dl = fir::support::getOrSetMLIRDataLayout(
module, /*allowDefaultLayout=*/false);
if (!dl)
mlir::emitError(module.getLoc(),
"data layout attribute is required to perform " DEBUG_TYPE
Expand Down
2 changes: 1 addition & 1 deletion flang/test/lib/OpenACC/TestOpenACCInterfaces.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ struct TestFIROpenACCInterfaces
void runOnOperation() override {
mlir::ModuleOp mod = getOperation();
auto datalayout =
fir::support::getOrSetDataLayout(mod, /*allowDefaultLayout=*/true);
fir::support::getOrSetMLIRDataLayout(mod, /*allowDefaultLayout=*/true);
mlir::OpBuilder builder(mod);
getOperation().walk([&](Operation *op) {
if (isa<ACC_DATA_ENTRY_OPS>(op)) {
Expand Down