[Mlir-commits] [mlir] 9d69bca - [mlir] Add stack alignment to the data layout dialect.
Tobias Gysi
llvmlistbot at llvm.org
Mon Apr 3 00:02:12 PDT 2023
Author: Tobias Gysi
Date: 2023-04-03T06:58:37Z
New Revision: 9d69bca18787e0861c1bc6db8e82d4f21b77ae75
URL: https://github.com/llvm/llvm-project/commit/9d69bca18787e0861c1bc6db8e82d4f21b77ae75
DIFF: https://github.com/llvm/llvm-project/commit/9d69bca18787e0861c1bc6db8e82d4f21b77ae75.diff
LOG: [mlir] Add stack alignment to the data layout dialect.
The revision adds the stack alignment to the
data layout dialect and it extends the LLVM dialect
import and export to support the new data layout
entry.
One possible use case for the flag is the LLVM dialect
inliner. The LLVM inliner queries the flag to
determine if it is safe to update the alignment of an
existing alloca. We may want to perform the same
optimization inside of MLIR.
Reviewed By: Dinistro
Differential Revision: https://reviews.llvm.org/D147332
Added:
Modified:
mlir/include/mlir/Dialect/DLTI/DLTI.h
mlir/include/mlir/Dialect/DLTI/DLTIBase.td
mlir/include/mlir/Interfaces/DataLayoutInterfaces.h
mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
mlir/lib/Dialect/DLTI/DLTI.cpp
mlir/lib/Interfaces/DataLayoutInterfaces.cpp
mlir/lib/Target/LLVMIR/DataLayoutImporter.cpp
mlir/lib/Target/LLVMIR/DataLayoutImporter.h
mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
mlir/test/Dialect/LLVMIR/layout.mlir
mlir/test/Target/LLVMIR/Import/data-layout.ll
mlir/test/Target/LLVMIR/data-layout.mlir
mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp
mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/DLTI/DLTI.h b/mlir/include/mlir/Dialect/DLTI/DLTI.h
index 77a036754f966..cf78b2312c246 100644
--- a/mlir/include/mlir/Dialect/DLTI/DLTI.h
+++ b/mlir/include/mlir/Dialect/DLTI/DLTI.h
@@ -101,6 +101,9 @@ class DataLayoutSpecAttr
/// Returns the alloca memory space identifier.
StringAttr getAllocaMemorySpaceIdentifier(MLIRContext *context) const;
+ /// Returns the stack alignment identifier.
+ StringAttr getStackAlignmentIdentifier(MLIRContext *context) const;
+
/// Parses an instance of this attribute.
static DataLayoutSpecAttr parse(AsmParser &parser);
diff --git a/mlir/include/mlir/Dialect/DLTI/DLTIBase.td b/mlir/include/mlir/Dialect/DLTI/DLTIBase.td
index a1f9eb361d717..37187bf9e2055 100644
--- a/mlir/include/mlir/Dialect/DLTI/DLTIBase.td
+++ b/mlir/include/mlir/Dialect/DLTI/DLTIBase.td
@@ -39,6 +39,9 @@ def DLTI_Dialect : Dialect {
constexpr const static ::llvm::StringLiteral
kDataLayoutAllocaMemorySpaceKey = "dlti.alloca_memory_space";
+
+ constexpr const static ::llvm::StringLiteral
+ kDataLayoutStackAlignmentKey = "dlti.stack_alignment";
}];
let useDefaultAttributePrinterParser = 1;
diff --git a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h
index 3d4d3796ecca0..e6e31d171c6a9 100644
--- a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h
+++ b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h
@@ -60,6 +60,10 @@ getDefaultPreferredAlignment(Type type, const DataLayout &dataLayout,
/// DataLayoutInterface if specified, otherwise returns the default.
Attribute getDefaultAllocaMemorySpace(DataLayoutEntryInterface entry);
+/// Default handler for the stack alignment request. Dispatches to the
+/// DataLayoutInterface if specified, otherwise returns the default.
+unsigned getDefaultStackAlignment(DataLayoutEntryInterface entry);
+
/// Given a list of data layout entries, returns a new list containing the
/// entries with keys having the given type ID, i.e. belonging to the same type
/// class.
@@ -166,6 +170,12 @@ class DataLayout {
/// Returns the memory space used for AllocaOps.
Attribute getAllocaMemorySpace() const;
+ /// Returns the natural alignment of the stack in bits. Alignment promotion of
+ /// stack variables should be limited to the natural stack alignment to
+ /// prevent dynamic stack alignment. Returns zero if the stack alignment is
+ /// unspecified.
+ unsigned getStackAlignment() const;
+
private:
/// Combined layout spec at the given scope.
const DataLayoutSpecInterface originalLayout;
@@ -190,6 +200,9 @@ class DataLayout {
/// Cache for alloca memory space.
mutable std::optional<Attribute> allocaMemorySpace;
+
+ /// Cache for stack alignment.
+ mutable std::optional<unsigned> stackAlignment;
};
} // namespace mlir
diff --git a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
index 2d9d6acbedae5..bac8c23ceec01 100644
--- a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
+++ b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
@@ -112,6 +112,12 @@ def DataLayoutSpecInterface : AttrInterface<"DataLayoutSpecInterface"> {
/*methodName=*/"getAllocaMemorySpaceIdentifier",
/*args=*/(ins "::mlir::MLIRContext *":$context)
>,
+ InterfaceMethod<
+ /*description=*/"Returns the stack alignment identifier.",
+ /*retTy=*/"::mlir::StringAttr",
+ /*methodName=*/"getStackAlignmentIdentifier",
+ /*args=*/(ins "::mlir::MLIRContext *":$context)
+ >,
// Implementations may override this if they have an efficient lookup
// mechanism.
InterfaceMethod<
@@ -274,6 +280,18 @@ def DataLayoutOpInterface : OpInterface<"DataLayoutOpInterface"> {
return ::mlir::detail::getDefaultAllocaMemorySpace(entry);
}]
>,
+ StaticInterfaceMethod<
+ /*description=*/"Returns the natural stack alignment in bits computed "
+ "using the relevant entries. The data layout object "
+ "can be used for recursive queries.",
+ /*retTy=*/"unsigned",
+ /*methodName=*/"getStackAlignment",
+ /*args=*/(ins "::mlir::DataLayoutEntryInterface":$entry),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return ::mlir::detail::getDefaultStackAlignment(entry);
+ }]
+ >,
];
let verify = [{ return ::mlir::detail::verifyDataLayoutOp($_op); }];
diff --git a/mlir/lib/Dialect/DLTI/DLTI.cpp b/mlir/lib/Dialect/DLTI/DLTI.cpp
index 86cc3ea766031..889778a1177cb 100644
--- a/mlir/lib/Dialect/DLTI/DLTI.cpp
+++ b/mlir/lib/Dialect/DLTI/DLTI.cpp
@@ -108,6 +108,7 @@ void DataLayoutEntryAttr::print(AsmPrinter &os) const {
constexpr const StringLiteral mlir::DataLayoutSpecAttr::kAttrKeyword;
constexpr const StringLiteral
mlir::DLTIDialect::kDataLayoutAllocaMemorySpaceKey;
+constexpr const StringLiteral mlir::DLTIDialect::kDataLayoutStackAlignmentKey;
namespace mlir {
namespace impl {
@@ -281,6 +282,12 @@ DataLayoutSpecAttr::getAllocaMemorySpaceIdentifier(MLIRContext *context) const {
DLTIDialect::kDataLayoutAllocaMemorySpaceKey);
}
+StringAttr
+DataLayoutSpecAttr::getStackAlignmentIdentifier(MLIRContext *context) const {
+ return Builder(context).getStringAttr(
+ DLTIDialect::kDataLayoutStackAlignmentKey);
+}
+
/// Parses an attribute with syntax
/// attr ::= `#target.` `dl_spec` `<` attr-list? `>`
/// attr-list ::= attr
@@ -337,7 +344,8 @@ class TargetDataLayoutInterface : public DataLayoutDialectInterface {
<< DLTIDialect::kDataLayoutEndiannessBig << "' or '"
<< DLTIDialect::kDataLayoutEndiannessLittle << "'";
}
- if (entryName == DLTIDialect::kDataLayoutAllocaMemorySpaceKey)
+ if (entryName == DLTIDialect::kDataLayoutAllocaMemorySpaceKey ||
+ entryName == DLTIDialect::kDataLayoutStackAlignmentKey)
return success();
return emitError(loc) << "unknown data layout entry name: " << entryName;
}
diff --git a/mlir/lib/Interfaces/DataLayoutInterfaces.cpp b/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
index ba6a321360fa5..ae8e680f0a068 100644
--- a/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
+++ b/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
@@ -225,6 +225,17 @@ mlir::detail::getDefaultAllocaMemorySpace(DataLayoutEntryInterface entry) {
return entry.getValue();
}
+// Returns the stack alignment if specified in the given entry. If the entry is
+// empty the default alignment zero is returned.
+unsigned
+mlir::detail::getDefaultStackAlignment(DataLayoutEntryInterface entry) {
+ if (entry == DataLayoutEntryInterface())
+ return 0;
+
+ auto value = entry.getValue().cast<IntegerAttr>();
+ return value.getValue().getZExtValue();
+}
+
DataLayoutEntryList
mlir::detail::filterEntriesForType(DataLayoutEntryListRef entries,
TypeID typeID) {
@@ -359,7 +370,7 @@ mlir::DataLayout::DataLayout() : DataLayout(ModuleOp()) {}
mlir::DataLayout::DataLayout(DataLayoutOpInterface op)
: originalLayout(getCombinedDataLayout(op)), scope(op),
- allocaMemorySpace(std::nullopt) {
+ allocaMemorySpace(std::nullopt), stackAlignment(std::nullopt) {
#if LLVM_ENABLE_ABI_BREAKING_CHECKS
checkMissingLayout(originalLayout, op);
collectParentLayouts(op, layoutStack);
@@ -368,7 +379,7 @@ mlir::DataLayout::DataLayout(DataLayoutOpInterface op)
mlir::DataLayout::DataLayout(ModuleOp op)
: originalLayout(getCombinedDataLayout(op)), scope(op),
- allocaMemorySpace(std::nullopt) {
+ allocaMemorySpace(std::nullopt), stackAlignment(std::nullopt) {
#if LLVM_ENABLE_ABI_BREAKING_CHECKS
checkMissingLayout(originalLayout, op);
collectParentLayouts(op, layoutStack);
@@ -486,6 +497,22 @@ mlir::Attribute mlir::DataLayout::getAllocaMemorySpace() const {
return *allocaMemorySpace;
}
+unsigned mlir::DataLayout::getStackAlignment() const {
+ checkValid();
+ MLIRContext *context = scope->getContext();
+ if (stackAlignment)
+ return *stackAlignment;
+ DataLayoutEntryInterface entry;
+ if (originalLayout)
+ entry = originalLayout.getSpecForIdentifier(
+ originalLayout.getStackAlignmentIdentifier(context));
+ if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
+ stackAlignment = iface.getStackAlignment(entry);
+ else
+ stackAlignment = detail::getDefaultStackAlignment(entry);
+ return *stackAlignment;
+}
+
//===----------------------------------------------------------------------===//
// DataLayoutSpecInterface
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Target/LLVMIR/DataLayoutImporter.cpp b/mlir/lib/Target/LLVMIR/DataLayoutImporter.cpp
index af480c597a97f..d19590600bb9a 100644
--- a/mlir/lib/Target/LLVMIR/DataLayoutImporter.cpp
+++ b/mlir/lib/Target/LLVMIR/DataLayoutImporter.cpp
@@ -183,6 +183,26 @@ DataLayoutImporter::tryToEmplaceAllocaAddrSpaceEntry(StringRef token) {
return success();
}
+LogicalResult
+DataLayoutImporter::tryToEmplaceStackAlignmentEntry(StringRef token) {
+ auto key =
+ StringAttr::get(context, DLTIDialect::kDataLayoutStackAlignmentKey);
+ if (keyEntries.count(key))
+ return success();
+
+ FailureOr<unsigned> alignment = tryToParseInt(token);
+ if (failed(alignment))
+ return failure();
+
+ // Only store the stack alignment if it has a non-default value.
+ if (*alignment == 0)
+ return success();
+ OpBuilder builder(context);
+ keyEntries.try_emplace(key, DataLayoutEntryAttr::get(
+ key, builder.getI32IntegerAttr(*alignment)));
+ return success();
+}
+
void DataLayoutImporter::translateDataLayout(
const llvm::DataLayout &llvmDataLayout) {
dataLayout = {};
@@ -230,6 +250,12 @@ void DataLayoutImporter::translateDataLayout(
return;
continue;
}
+ // Parse the stack alignment.
+ if (*prefix == "S") {
+ if (failed(tryToEmplaceStackAlignmentEntry(token)))
+ return;
+ continue;
+ }
// Parse integer alignment specifications.
if (*prefix == "i") {
FailureOr<unsigned> width = tryToParseInt(token);
diff --git a/mlir/lib/Target/LLVMIR/DataLayoutImporter.h b/mlir/lib/Target/LLVMIR/DataLayoutImporter.h
index 2b2fa663a63fb..5c1606216037e 100644
--- a/mlir/lib/Target/LLVMIR/DataLayoutImporter.h
+++ b/mlir/lib/Target/LLVMIR/DataLayoutImporter.h
@@ -38,11 +38,11 @@ namespace detail {
FloatType getFloatType(MLIRContext *context, unsigned width);
/// Helper class that translates an LLVM data layout to an MLIR data layout
-/// specification. Only integer, float, pointer, alloca memory space, and
-/// endianness entries are translated. The class also returns all entries from
-/// the default data layout specification found in the language reference
-/// (https://llvm.org/docs/LangRef.html#data-layout) if they are not overwritten
-/// by the provided data layout.
+/// specification. Only integer, float, pointer, alloca memory space, stack
+/// alignment, and endianness entries are translated. The class also returns all
+/// entries from the default data layout specification found in the language
+/// reference (https://llvm.org/docs/LangRef.html#data-layout) if they are not
+/// overwritten by the provided data layout.
class DataLayoutImporter {
public:
DataLayoutImporter(MLIRContext *context,
@@ -99,6 +99,9 @@ class DataLayoutImporter {
/// Adds an alloca address space entry if there is none yet.
LogicalResult tryToEmplaceAllocaAddrSpaceEntry(StringRef token);
+ /// Adds a stack alignment entry if there is none yet.
+ LogicalResult tryToEmplaceStackAlignmentEntry(StringRef token);
+
std::string layoutStr = {};
StringRef lastToken = {};
SmallVector<StringRef> unhandledTokens;
diff --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
index 7159aa49540ef..a08bb372fc738 100644
--- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
@@ -75,17 +75,28 @@ translateDataLayout(DataLayoutSpecInterface attribute,
auto value = entry.getValue().cast<StringAttr>();
bool isLittleEndian =
value.getValue() == DLTIDialect::kDataLayoutEndiannessLittle;
- layoutStream << (isLittleEndian ? "e" : "E");
+ layoutStream << "-" << (isLittleEndian ? "e" : "E");
layoutStream.flush();
continue;
}
if (key.getValue() == DLTIDialect::kDataLayoutAllocaMemorySpaceKey) {
auto value = entry.getValue().cast<IntegerAttr>();
- if (value != 0) {
- // Only emit non-default address space.
- layoutStream << "A" << value;
- layoutStream.flush();
- }
+ uint64_t space = value.getValue().getZExtValue();
+ // Skip the default address space.
+ if (space == 0)
+ continue;
+ layoutStream << "-A" << space;
+ layoutStream.flush();
+ continue;
+ }
+ if (key.getValue() == DLTIDialect::kDataLayoutStackAlignmentKey) {
+ auto value = entry.getValue().cast<IntegerAttr>();
+ uint64_t alignment = value.getValue().getZExtValue();
+ // Skip the default stack alignment.
+ if (alignment == 0)
+ continue;
+ layoutStream << "-S" << alignment;
+ layoutStream.flush();
continue;
}
emitError(*loc) << "unsupported data layout key " << key;
diff --git a/mlir/test/Dialect/LLVMIR/layout.mlir b/mlir/test/Dialect/LLVMIR/layout.mlir
index ba388c134682b..363222b5ad912 100644
--- a/mlir/test/Dialect/LLVMIR/layout.mlir
+++ b/mlir/test/Dialect/LLVMIR/layout.mlir
@@ -8,18 +8,21 @@ module {
// CHECK: bitsize = 64
// CHECK: preferred = 8
// CHECK: size = 8
+ // CHECK: stack_alignment = 0
"test.data_layout_query"() : () -> !llvm.ptr
// CHECK: alignment = 8
// CHECK: alloca_memory_space = 0
// CHECK: bitsize = 64
// CHECK: preferred = 8
// CHECK: size = 8
+ // CHECK: stack_alignment = 0
"test.data_layout_query"() : () -> !llvm.ptr<3>
// CHECK: alignment = 8
// CHECK: alloca_memory_space = 0
// CHECK: bitsize = 64
// CHECK: preferred = 8
// CHECK: size = 8
+ // CHECK: stack_alignment = 0
"test.data_layout_query"() : () -> !llvm.ptr<5>
return
}
@@ -31,7 +34,8 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
#dlti.dl_entry<!llvm.ptr, dense<[32, 32, 64]> : vector<3xi32>>,
#dlti.dl_entry<!llvm.ptr<5>, dense<[64, 64, 64]> : vector<3xi32>>,
#dlti.dl_entry<!llvm.ptr<4>, dense<[32, 64, 64]> : vector<3xi32>>,
- #dlti.dl_entry<"dlti.alloca_memory_space", 5 : ui32>
+ #dlti.dl_entry<"dlti.alloca_memory_space", 5 : ui32>,
+ #dlti.dl_entry<"dlti.stack_alignment", 128 : i32>
>} {
// CHECK: @spec
func.func @spec() {
@@ -40,31 +44,36 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
// CHECK: bitsize = 32
// CHECK: preferred = 8
// CHECK: size = 4
+ // CHECK: stack_alignment = 128
"test.data_layout_query"() : () -> !llvm.ptr
// CHECK: alignment = 4
// CHECK: alloca_memory_space = 5
// CHECK: bitsize = 32
// CHECK: preferred = 8
// CHECK: size = 4
+ // CHECK: stack_alignment = 128
"test.data_layout_query"() : () -> !llvm.ptr<3>
// CHECK: alignment = 8
// CHECK: alloca_memory_space = 5
// CHECK: bitsize = 64
// CHECK: preferred = 8
// CHECK: size = 8
+ // CHECK: stack_alignment = 128
"test.data_layout_query"() : () -> !llvm.ptr<5>
// CHECK: alignment = 4
// CHECK: alloca_memory_space = 5
// CHECK: bitsize = 32
// CHECK: preferred = 8
// CHECK: size = 4
+ // CHECK: stack_alignment = 128
"test.data_layout_query"() : () -> !llvm.ptr<3>
// CHECK: alignment = 8
// CHECK: alloca_memory_space = 5
// CHECK: bitsize = 32
// CHECK: preferred = 8
// CHECK: size = 4
- "test.data_layout_query"() : () -> !llvm.ptr<4>
+ // CHECK: stack_alignment = 128
+ "test.data_layout_query"() : () -> !llvm.ptr<4>
return
}
}
diff --git a/mlir/test/Target/LLVMIR/Import/data-layout.ll b/mlir/test/Target/LLVMIR/Import/data-layout.ll
index e29e1744c562a..b46f5fc4f2ae2 100644
--- a/mlir/test/Target/LLVMIR/Import/data-layout.ll
+++ b/mlir/test/Target/LLVMIR/Import/data-layout.ll
@@ -28,6 +28,7 @@ target datalayout = ""
; CHECK-DAG: #dlti.dl_entry<!llvm.ptr<270>, dense<[32, 64, 64, 32]> : vector<4xi32>>
; CHECK-DAG: #dlti.dl_entry<!llvm.ptr<271>, dense<32> : vector<4xi32>>
; CHECK-DAG: #dlti.dl_entry<!llvm.ptr<272>, dense<64> : vector<4xi32>>
+; CHECK-DAG: #dlti.dl_entry<"dlti.stack_alignment", 128 : i32>
target datalayout = "e-m:e-p270:32:64-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
; // -----
@@ -36,13 +37,15 @@ target datalayout = "e-m:e-p270:32:64-p271:32:32-p272:64:64-i64:64-f80:128-n8:16
; CHECK: #dlti.dl_spec<
; CHECK-DAG: #dlti.dl_entry<"dlti.endianness", "big">
; CHECK-DAG: #dlti.dl_entry<!llvm.ptr<270>, dense<[16, 32, 64, 128]> : vector<4xi32>>
+; CHECK-DAG: #dlti.dl_entry<!llvm.ptr<271>, dense<[16, 32, 64, 16]> : vector<4xi32>>
; CHECK-DAG: #dlti.dl_entry<"dlti.alloca_memory_space", 1 : ui32>
; CHECK-DAG: #dlti.dl_entry<i64, dense<[64, 128]> : vector<2xi32>>
-target datalayout = "E-p270:16:32:64:128-A1-i64:64:128"
+target datalayout = "A1-E-p270:16:32:64:128-p271:16:32:64-i64:64:128"
; // -----
; CHECK: dlti.dl_spec =
; CHECK: #dlti.dl_spec<
; CHECK-NOT: #dlti.dl_entry<"dlti.alloca_memory_space"
-target datalayout = "E-A0"
+; CHECK-NOT: #dlti.dl_entry<"dlti.stack_alignment"
+target datalayout = "A0-S0"
diff --git a/mlir/test/Target/LLVMIR/data-layout.mlir b/mlir/test/Target/LLVMIR/data-layout.mlir
index 5840bcf700dd2..74941d84d7d21 100644
--- a/mlir/test/Target/LLVMIR/data-layout.mlir
+++ b/mlir/test/Target/LLVMIR/data-layout.mlir
@@ -2,12 +2,16 @@
// CHECK: target datalayout
// CHECK: E-
+// CHECK: A4-
+// CHECK: S128-
// CHECK: i64:64:128
// CHECK: f80:128:256
// CHECK: p0:32:64:128
// CHECK: p1:32:32:32:64
module attributes {dlti.dl_spec = #dlti.dl_spec<
#dlti.dl_entry<"dlti.endianness", "big">,
+#dlti.dl_entry<"dlti.alloca_memory_space", 4 : ui32>,
+#dlti.dl_entry<"dlti.stack_alignment", 128 : i32>,
#dlti.dl_entry<index, 64>,
#dlti.dl_entry<i64, dense<[64,128]> : vector<2xi32>>,
#dlti.dl_entry<f80, dense<[128,256]> : vector<2xi32>>,
@@ -21,6 +25,22 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<
// -----
+// CHECK: target datalayout
+// CHECK: e
+// CHECK-NOT: A0
+// CHECK-NOT: S0
+module attributes {dlti.dl_spec = #dlti.dl_spec<
+#dlti.dl_entry<"dlti.endianness", "little">,
+#dlti.dl_entry<"dlti.alloca_memory_space", 0 : ui32>,
+#dlti.dl_entry<"dlti.stack_alignment", 0 : i32>
+>} {
+ llvm.func @bar() {
+ llvm.return
+ }
+}
+
+// -----
+
// expected-error at below {{unsupported data layout for non-signless integer 'ui64'}}
module attributes {dlti.dl_spec = #dlti.dl_spec<
#dlti.dl_entry<ui64, dense<[64,128]> : vector<2xi32>>>
diff --git a/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp b/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp
index 9e71bd4e999f7..c326a746601b2 100644
--- a/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp
+++ b/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp
@@ -41,6 +41,7 @@ struct TestDataLayoutQuery
unsigned alignment = layout.getTypeABIAlignment(op.getType());
unsigned preferred = layout.getTypePreferredAlignment(op.getType());
Attribute allocaMemorySpace = layout.getAllocaMemorySpace();
+ unsigned stackAlignment = layout.getStackAlignment();
op->setAttrs(
{builder.getNamedAttr("size", builder.getIndexAttr(size)),
builder.getNamedAttr("bitsize", builder.getIndexAttr(bitsize)),
@@ -49,9 +50,9 @@ struct TestDataLayoutQuery
builder.getNamedAttr("alloca_memory_space",
allocaMemorySpace == Attribute()
? builder.getUI32IntegerAttr(0)
- : allocaMemorySpace)
-
- });
+ : allocaMemorySpace),
+ builder.getNamedAttr("stack_alignment",
+ builder.getIndexAttr(stackAlignment))});
});
}
};
diff --git a/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp b/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp
index c35e176035227..7e0a8f55a5526 100644
--- a/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp
+++ b/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp
@@ -24,6 +24,8 @@ namespace {
constexpr static llvm::StringLiteral kAttrName = "dltest.layout";
constexpr static llvm::StringLiteral kAllocaKeyName =
"dltest.alloca_memory_space";
+constexpr static llvm::StringLiteral kStackAlignmentKeyName =
+ "dltest.stack_alignment";
/// Trivial array storage for the custom data layout spec attribute, just a list
/// of entries.
@@ -67,6 +69,9 @@ struct CustomDataLayoutSpec
StringAttr getAllocaMemorySpaceIdentifier(MLIRContext *context) const {
return Builder(context).getStringAttr(kAllocaKeyName);
}
+ StringAttr getStackAlignmentIdentifier(MLIRContext *context) const {
+ return Builder(context).getStringAttr(kStackAlignmentKeyName);
+ }
};
/// A type subject to data layout that exits the program if it is queried more
@@ -276,6 +281,7 @@ module {}
EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 2u);
EXPECT_EQ(layout.getAllocaMemorySpace(), Attribute());
+ EXPECT_EQ(layout.getStackAlignment(), 0u);
}
TEST(DataLayout, NullSpec) {
@@ -302,6 +308,7 @@ TEST(DataLayout, NullSpec) {
EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 32u);
EXPECT_EQ(layout.getAllocaMemorySpace(), Attribute());
+ EXPECT_EQ(layout.getStackAlignment(), 0u);
}
TEST(DataLayout, EmptySpec) {
@@ -327,6 +334,7 @@ TEST(DataLayout, EmptySpec) {
EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 32u);
EXPECT_EQ(layout.getAllocaMemorySpace(), Attribute());
+ EXPECT_EQ(layout.getStackAlignment(), 0u);
}
TEST(DataLayout, SpecWithEntries) {
@@ -334,7 +342,8 @@ TEST(DataLayout, SpecWithEntries) {
"dltest.op_with_layout"() { dltest.layout = #dltest.spec<
#dlti.dl_entry<i42, 5>,
#dlti.dl_entry<i16, 6>,
- #dlti.dl_entry<"dltest.alloca_memory_space", 5 : i32>
+ #dlti.dl_entry<"dltest.alloca_memory_space", 5 : i32>,
+ #dlti.dl_entry<"dltest.stack_alignment", 128 : i32>
> } : () -> ()
)MLIR";
@@ -365,6 +374,7 @@ TEST(DataLayout, SpecWithEntries) {
EXPECT_EQ(layout.getTypePreferredAlignment(Float32Type::get(&ctx)), 64u);
EXPECT_EQ(layout.getAllocaMemorySpace(), Builder(&ctx).getI32IntegerAttr(5));
+ EXPECT_EQ(layout.getStackAlignment(), 128u);
}
TEST(DataLayout, Caching) {
More information about the Mlir-commits
mailing list