[MLIR][LLVM][DLTI] Handle data layout token 'n32:64' (#141299)

This commit is contained in:
Bruno Cardoso Lopes
2025-05-28 11:07:03 -07:00
committed by GitHub
parent 7ed185a8f2
commit 86685b95bf
14 changed files with 139 additions and 5 deletions

View File

@@ -95,6 +95,9 @@ def DLTI_DataLayoutSpecAttr :
/// Returns the function pointer alignment identifier.
StringAttr getFunctionPointerAlignmentIdentifier(MLIRContext *context) const;
/// Returns the legal int widths identifier.
StringAttr getLegalIntWidthsIdentifier(MLIRContext *context) const;
/// Returns the attribute associated with the key.
FailureOr<Attribute> query(DataLayoutEntryKey key) {
return ::llvm::cast<mlir::DataLayoutSpecInterface>(*this).queryHelper(key);

View File

@@ -70,6 +70,9 @@ def DLTI_Dialect : Dialect {
constexpr const static ::llvm::StringLiteral
kDataLayoutFunctionPointerAlignmentKey = "dlti.function_pointer_alignment";
constexpr const static ::llvm::StringLiteral
kDataLayoutLegalIntWidthsKey = "dlti.legal_int_widths";
}];
let useDefaultAttributePrinterParser = 1;

View File

@@ -105,6 +105,10 @@ uint64_t getDefaultStackAlignment(DataLayoutEntryInterface entry);
/// the DataLayoutInterface if specified, otherwise returns the default.
Attribute getDefaultFunctionPointerAlignment(DataLayoutEntryInterface entry);
/// Default handler for the legal int widths request. Dispatches to the
/// DataLayoutInterface if specified, otherwise returns the default.
Attribute getDefaultLegalIntWidths(DataLayoutEntryInterface entry);
/// Returns the value of the property from the specified DataLayoutEntry. If the
/// property is missing from the entry, returns std::nullopt.
std::optional<Attribute> getDevicePropertyValue(DataLayoutEntryInterface entry);
@@ -266,6 +270,9 @@ public:
/// Returns function pointer alignment.
Attribute getFunctionPointerAlignment() const;
/// Returns the legal int widths.
Attribute getLegalIntWidths() const;
/// Returns the value of the specified property if the property is defined for
/// the given device ID, otherwise returns std::nullopt.
std::optional<Attribute>
@@ -312,6 +319,8 @@ private:
mutable std::optional<uint64_t> stackAlignment;
/// Cache for function pointer alignment.
mutable std::optional<Attribute> functionPointerAlignment;
/// Cache for legal int widths.
mutable std::optional<Attribute> legalIntWidths;
};
} // namespace mlir

View File

@@ -177,6 +177,12 @@ def DataLayoutSpecInterface : AttrInterface<"DataLayoutSpecInterface", [DLTIQuer
/*methodName=*/"getFunctionPointerAlignmentIdentifier",
/*args=*/(ins "::mlir::MLIRContext *":$context)
>,
InterfaceMethod<
/*description=*/"Returns the legal int widths identifier.",
/*retTy=*/"::mlir::StringAttr",
/*methodName=*/"getLegalIntWidthsIdentifier",
/*args=*/(ins "::mlir::MLIRContext *":$context)
>,
// Implementations may override this if they have an efficient lookup
// mechanism.
InterfaceMethod<
@@ -571,6 +577,18 @@ def DataLayoutOpInterface : OpInterface<"DataLayoutOpInterface"> {
return ::mlir::detail::getDefaultFunctionPointerAlignment(entry);
}]
>,
StaticInterfaceMethod<
/*description=*/"Returns the legal int widths, each width in bits computed "
"using the relevant entries. The data layout object "
"can be used for recursive queries.",
/*retTy=*/"Attribute",
/*methodName=*/"getLegalIntWidths",
/*args=*/(ins "::mlir::DataLayoutEntryInterface":$entry),
/*methodBody=*/"",
/*defaultImplementation=*/[{
return ::mlir::detail::getDefaultLegalIntWidths(entry);
}]
>,
StaticInterfaceMethod<
/*description=*/"Returns the value of the property, if the property is "
"defined. Otherwise, it returns std::nullopt.",

View File

@@ -426,6 +426,12 @@ StringAttr DataLayoutSpecAttr::getFunctionPointerAlignmentIdentifier(
DLTIDialect::kDataLayoutFunctionPointerAlignmentKey);
}
StringAttr
DataLayoutSpecAttr::getLegalIntWidthsIdentifier(MLIRContext *context) const {
return Builder(context).getStringAttr(
DLTIDialect::kDataLayoutLegalIntWidthsKey);
}
/// Parses an attribute with syntax:
/// dl-spec-attr ::= `#dlti.` `dl_spec` `<` entry-list `>`
/// entry-list ::= | entry | entry `,` entry-list
@@ -632,6 +638,7 @@ public:
entryName == DLTIDialect::kDataLayoutGlobalMemorySpaceKey ||
entryName == DLTIDialect::kDataLayoutStackAlignmentKey ||
entryName == DLTIDialect::kDataLayoutFunctionPointerAlignmentKey ||
entryName == DLTIDialect::kDataLayoutLegalIntWidthsKey ||
entryName == DLTIDialect::kDataLayoutManglingModeKey)
return success();
return emitError(loc) << "unknown data layout entry name: " << entryName;

View File

@@ -321,6 +321,16 @@ Attribute mlir::detail::getDefaultFunctionPointerAlignment(
return entry.getValue();
}
// Returns the legal int widths if specified in the given entry. If the entry is
// empty the default legal int widths represented by an empty attribute is
// returned.
Attribute
mlir::detail::getDefaultLegalIntWidths(DataLayoutEntryInterface entry) {
if (entry == DataLayoutEntryInterface())
return Attribute();
return entry.getValue();
}
std::optional<Attribute>
mlir::detail::getDevicePropertyValue(DataLayoutEntryInterface entry) {
if (entry == DataLayoutEntryInterface())
@@ -736,6 +746,22 @@ Attribute mlir::DataLayout::getFunctionPointerAlignment() const {
return *functionPointerAlignment;
}
Attribute mlir::DataLayout::getLegalIntWidths() const {
checkValid();
if (legalIntWidths)
return *legalIntWidths;
DataLayoutEntryInterface entry;
if (originalLayout)
entry = originalLayout.getSpecForIdentifier(
originalLayout.getLegalIntWidthsIdentifier(
originalLayout.getContext()));
if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
legalIntWidths = iface.getLegalIntWidths(entry);
else
legalIntWidths = detail::getDefaultLegalIntWidths(entry);
return *legalIntWidths;
}
std::optional<Attribute> mlir::DataLayout::getDevicePropertyValue(
TargetSystemSpecInterface::DeviceID deviceID,
StringAttr propertyName) const {

View File

@@ -63,20 +63,25 @@ FailureOr<uint64_t> DataLayoutImporter::tryToParseInt(StringRef &token) const {
return parameter;
}
FailureOr<SmallVector<uint64_t>>
DataLayoutImporter::tryToParseIntList(StringRef token) const {
template <class T>
static FailureOr<SmallVector<T>> tryToParseIntListImpl(StringRef token) {
SmallVector<StringRef> tokens;
token.consume_front(":");
token.split(tokens, ':');
// Parse an integer list.
SmallVector<uint64_t> results(tokens.size());
SmallVector<T> results(tokens.size());
for (auto [result, token] : llvm::zip(results, tokens))
if (token.getAsInteger(/*Radix=*/10, result))
return failure();
return results;
}
FailureOr<SmallVector<uint64_t>>
DataLayoutImporter::tryToParseIntList(StringRef token) const {
return tryToParseIntListImpl<uint64_t>(token);
}
FailureOr<DenseIntElementsAttr>
DataLayoutImporter::tryToParseAlignment(StringRef token) const {
FailureOr<SmallVector<uint64_t>> alignment = tryToParseIntList(token);
@@ -251,6 +256,25 @@ LogicalResult DataLayoutImporter::tryToEmplaceFunctionPointerAlignmentEntry(
return success();
}
LogicalResult
DataLayoutImporter::tryToEmplaceLegalIntWidthsEntry(StringRef token) {
auto key =
StringAttr::get(context, DLTIDialect::kDataLayoutLegalIntWidthsKey);
if (keyEntries.count(key))
return success();
FailureOr<SmallVector<int32_t>> intWidths =
tryToParseIntListImpl<int32_t>(token);
if (failed(intWidths) || intWidths->empty())
return failure();
OpBuilder builder(context);
keyEntries.try_emplace(
key,
DataLayoutEntryAttr::get(key, builder.getDenseI32ArrayAttr(*intWidths)));
return success();
}
void DataLayoutImporter::translateDataLayout(
const llvm::DataLayout &llvmDataLayout) {
dataLayout = {};
@@ -360,6 +384,12 @@ void DataLayoutImporter::translateDataLayout(
return;
continue;
}
// Parse native integer widths specifications.
if (*prefix == "n") {
if (failed(tryToEmplaceLegalIntWidthsEntry(token)))
return;
continue;
}
// Parse function pointer alignment specifications.
// Note that prefix here is "Fn" or "Fi", not a single character.
if (prefix->starts_with("F")) {

View File

@@ -113,6 +113,9 @@ private:
tryToEmplaceFunctionPointerAlignmentEntry(StringRef fnPtrAlignEntry,
StringRef token);
/// Adds legal int widths entry if there is none yet.
LogicalResult tryToEmplaceLegalIntWidthsEntry(StringRef token);
std::string layoutStr = {};
StringRef lastToken = {};
SmallVector<StringRef> unhandledTokens;

View File

@@ -35,6 +35,7 @@
#include "mlir/Target/LLVMIR/TypeToLLVM.h"
#include "llvm/ADT/PostOrderIterator.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/StringExtras.h"
#include "llvm/ADT/TypeSwitch.h"
@@ -52,6 +53,7 @@
#include "llvm/IR/Module.h"
#include "llvm/IR/Verifier.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/ErrorHandling.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Transforms/Utils/BasicBlockUtils.h"
#include "llvm/Transforms/Utils/Cloning.h"
@@ -250,6 +252,13 @@ translateDataLayout(DataLayoutSpecInterface attribute,
<< alignment;
continue;
}
if (key.getValue() == DLTIDialect::kDataLayoutLegalIntWidthsKey) {
layoutStream << "-n";
llvm::interleave(
cast<DenseI32ArrayAttr>(entry.getValue()).asArrayRef(), layoutStream,
[&](int32_t val) { layoutStream << val; }, ":");
continue;
}
emitError(*loc) << "unsupported data layout key " << key;
return failure();
}

View File

@@ -12,6 +12,7 @@ module {
// CHECK-SAME: #dlti.function_pointer_alignment<0, function_dependent = false>,
// CHECK: global_memory_space = 0
// CHECK: index = 64
// CHECK: legal_int_widths = array<i32>
// CHECK: mangling_mode = ""
// CHECK: preferred = 8
// CHECK: program_memory_space = 0
@@ -27,6 +28,7 @@ module {
// CHECK-SAME: #dlti.function_pointer_alignment<0, function_dependent = false>,
// CHECK: global_memory_space = 0
// CHECK: index = 64
// CHECK: legal_int_widths = array<i32>
// CHECK: mangling_mode = ""
// CHECK: preferred = 8
// CHECK: program_memory_space = 0
@@ -42,6 +44,7 @@ module {
// CHECK-SAME: #dlti.function_pointer_alignment<0, function_dependent = false>,
// CHECK: global_memory_space = 0
// CHECK: index = 64
// CHECK: legal_int_widths = array<i32>
// CHECK: mangling_mode = ""
// CHECK: preferred = 8
// CHECK: program_memory_space = 0
@@ -65,6 +68,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
#dlti.dl_entry<"dlti.program_memory_space", 3 : ui64>,
#dlti.dl_entry<"dlti.stack_alignment", 128 : i64>,
#dlti.dl_entry<"dlti.mangling_mode", "e">,
#dlti.dl_entry<"dlti.legal_int_widths", array<i32: 32, 64>>,
#dlti.dl_entry<"dlti.function_pointer_alignment",
"#dlti.function_pointer_alignment<32, function_dependent = true>">
>} {
@@ -79,6 +83,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
// CHECK-SAME: "#dlti.function_pointer_alignment<32, function_dependent = true>",
// CHECK: global_memory_space = 2
// CHECK: index = 32
// CHECK: legal_int_widths = array<i32: 32, 64>
// CHECK: mangling_mode = "e"
// CHECK: preferred = 8
// CHECK: program_memory_space = 3
@@ -94,6 +99,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
// CHECK-SAME: "#dlti.function_pointer_alignment<32, function_dependent = true>",
// CHECK: global_memory_space = 2
// CHECK: index = 32
// CHECK: legal_int_widths = array<i32: 32, 64>
// CHECK: preferred = 8
// CHECK: program_memory_space = 3
// CHECK: size = 4
@@ -108,6 +114,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
// CHECK-SAME: "#dlti.function_pointer_alignment<32, function_dependent = true>",
// CHECK: global_memory_space = 2
// CHECK: index = 64
// CHECK: legal_int_widths = array<i32: 32, 64>
// CHECK: mangling_mode = "e"
// CHECK: preferred = 8
// CHECK: program_memory_space = 3
@@ -123,6 +130,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
// CHECK-SAME: "#dlti.function_pointer_alignment<32, function_dependent = true>",
// CHECK: global_memory_space = 2
// CHECK: index = 24
// CHECK: legal_int_widths = array<i32: 32, 64>
// CHECK: mangling_mode = "e"
// CHECK: preferred = 8
// CHECK: program_memory_space = 3

View File

@@ -29,6 +29,7 @@ target datalayout = ""
; CHECK-SAME: i8 = dense<8> : vector<2xi64>
; CHECK-SAME: "dlti.endianness" = "little"
; CHECK-SAME: "dlti.mangling_mode" = "e"
; CHECK-SAME: "dlti.legal_int_widths" = array<i32: 8, 16, 32, 64>
; CHECK-SAME: "dlti.stack_alignment" = 128 : i64
; CHECK-SAME: "dlti.function_pointer_alignment" = #dlti.function_pointer_alignment<32, function_dependent = true>
target datalayout = "e-m:e-p270:32:64-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128-Fn32"

View File

@@ -5,6 +5,7 @@
// CHECK: A4-
// CHECK: S128-
// CHECK: m:e-
// CHECK: n8:16:32:64-
// CHECK: Fn32
// CHECK: i64:64:128
// CHECK: f80:128:256
@@ -15,6 +16,7 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<
#dlti.dl_entry<"dlti.alloca_memory_space", 4 : ui32>,
#dlti.dl_entry<"dlti.stack_alignment", 128 : i32>,
#dlti.dl_entry<"dlti.mangling_mode", "e">,
#dlti.dl_entry<"dlti.legal_int_widths", array<i32: 8, 16, 32, 64>>,
#dlti.dl_entry<"dlti.function_pointer_alignment",
#dlti.function_pointer_alignment<32, function_dependent = true>>,
#dlti.dl_entry<index, 64>,

View File

@@ -49,6 +49,7 @@ struct TestDataLayoutQuery
Attribute globalMemorySpace = layout.getGlobalMemorySpace();
uint64_t stackAlignment = layout.getStackAlignment();
Attribute functionPointerAlignment = layout.getFunctionPointerAlignment();
Attribute legalIntWidths = layout.getLegalIntWidths();
auto convertTypeSizeToAttr = [&](llvm::TypeSize typeSize) -> Attribute {
if (!typeSize.isScalable())
@@ -97,7 +98,11 @@ struct TestDataLayoutQuery
? FunctionPointerAlignmentAttr::get(
builder.getContext(), 0,
/*function_dependent=*/false)
: functionPointerAlignment)
: functionPointerAlignment),
builder.getNamedAttr("legal_int_widths",
legalIntWidths == Attribute()
? builder.getDenseI32ArrayAttr({})
: legalIntWidths)
});

View File

@@ -37,6 +37,8 @@ constexpr static llvm::StringLiteral kStackAlignmentKeyName =
"dltest.stack_alignment";
constexpr static llvm::StringLiteral kFunctionPointerAlignmentKeyName =
"dltest.function_pointer_alignment";
constexpr static llvm::StringLiteral kLegalIntWidthsKeyName =
"dltest.legal_int_widths";
constexpr static llvm::StringLiteral kTargetSystemDescAttrName =
"dl_target_sys_desc_test.target_system_spec";
@@ -107,6 +109,9 @@ struct CustomDataLayoutSpec
StringAttr getFunctionPointerAlignmentIdentifier(MLIRContext *context) const {
return Builder(context).getStringAttr(kFunctionPointerAlignmentKeyName);
}
StringAttr getLegalIntWidthsIdentifier(MLIRContext *context) const {
return Builder(context).getStringAttr(kLegalIntWidthsKeyName);
}
FailureOr<Attribute> query(DataLayoutEntryKey key) const {
return llvm::cast<mlir::DataLayoutSpecInterface>(*this).queryHelper(key);
}
@@ -500,6 +505,7 @@ module {}
EXPECT_EQ(layout.getGlobalMemorySpace(), Attribute());
EXPECT_EQ(layout.getStackAlignment(), 0u);
EXPECT_EQ(layout.getFunctionPointerAlignment(), Attribute());
EXPECT_EQ(layout.getLegalIntWidths(), Attribute());
EXPECT_EQ(layout.getManglingMode(), Attribute());
}
@@ -578,6 +584,7 @@ TEST(DataLayout, EmptySpec) {
EXPECT_EQ(layout.getStackAlignment(), 0u);
EXPECT_EQ(layout.getManglingMode(), Attribute());
EXPECT_EQ(layout.getFunctionPointerAlignment(), Attribute());
EXPECT_EQ(layout.getLegalIntWidths(), Attribute());
EXPECT_EQ(layout.getDevicePropertyValue(
Builder(&ctx).getStringAttr("CPU" /* device ID*/),
@@ -603,7 +610,8 @@ TEST(DataLayout, SpecWithEntries) {
#dlti.dl_entry<"dltest.stack_alignment", 128 : i32>,
#dlti.dl_entry<"dltest.mangling_mode", "o">,
#dlti.dl_entry<"dltest.function_pointer_alignment",
#dlti.function_pointer_alignment<64, function_dependent = true>>
#dlti.function_pointer_alignment<64, function_dependent = true>>,
#dlti.dl_entry<"dltest.legal_int_widths", array<i32: 64>>
> } : () -> ()
)MLIR";
@@ -645,6 +653,8 @@ TEST(DataLayout, SpecWithEntries) {
EXPECT_EQ(
layout.getFunctionPointerAlignment(),
FunctionPointerAlignmentAttr::get(&ctx, 64, /*function_dependent=*/true));
EXPECT_EQ(layout.getLegalIntWidths(),
Builder(&ctx).getDenseI32ArrayAttr({64}));
}
TEST(DataLayout, SpecWithTargetSystemDescEntries) {