[Mlir-commits] [mlir] [mlir][DataLayout] Add getTypeAllocSize (PR #231262)
Adam Smith
llvmlistbot at llvm.org
Tue Oct 13 08:08:34 PDT 2026
https://github.com/adams381 updated https://github.com/llvm/llvm-project/pull/231262
>From eca10f019ad9cf8468b3225acc976177ec2753a4 Mon Sep 17 00:00:00 2001
From: Adam Smith <adams at nvidia.com>
Date: Mon, 12 Oct 2026 11:52:58 -0700
Subject: [PATCH 1/2] [mlir][DataLayout] Add getTypeAllocSize
Add DataLayout::getTypeAllocSize, the type's size rounded up to its ABI
alignment, as in llvm::DataLayout. Use it for the three places that
computed this by hand: LLVM dialect arrays and structs, and the OpenMP
offload privatization pass. ClangIR's type layout hooks also need it.
Assisted-by: Cursor / Claude Opus 5.5
---
mlir/docs/DataLayout.md | 1 +
.../mlir/Interfaces/DataLayoutInterfaces.h | 4 +++
mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp | 7 ++---
.../OpenMPOffloadPrivatizationPrepare.cpp | 4 +--
mlir/lib/Interfaces/DataLayoutInterfaces.cpp | 4 +++
mlir/test/Dialect/LLVMIR/layout.mlir | 27 +++++++++++++++++
.../DataLayoutInterfaces/query.mlir | 1 +
.../lib/Dialect/DLTI/TestDataLayoutQuery.cpp | 2 ++
.../Interfaces/DataLayoutInterfacesTest.cpp | 29 +++++++++++++++++++
9 files changed, 71 insertions(+), 8 deletions(-)
diff --git a/mlir/docs/DataLayout.md b/mlir/docs/DataLayout.md
index 21430c44ca95fb..a6ee6120e8b921 100644
--- a/mlir/docs/DataLayout.md
+++ b/mlir/docs/DataLayout.md
@@ -75,6 +75,7 @@ public:
llvm::TypeSize getTypeSize(Type type) const;
llvm::TypeSize getTypeSizeInBits(Type type) const;
+ llvm::TypeSize getTypeAllocSize(Type type) const;
uint64_t getTypeABIAlignment(Type type) const;
uint64_t getTypePreferredAlignment(Type type) const;
std::optional<uint64_t> getTypeIndexBitwidth(Type type) const;
diff --git a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h
index 20de0c6ba7e43f..02a9eaefad8444 100644
--- a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h
+++ b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h
@@ -187,6 +187,10 @@ class DataLayout {
/// Returns the size in bits of the given type in the current scope.
llvm::TypeSize getTypeSizeInBits(Type t) const;
+ /// Returns the offset in bytes between successive objects of the given type
+ /// in the current scope, that is, its size rounded up to its ABI alignment.
+ llvm::TypeSize getTypeAllocSize(Type t) const;
+
/// Returns the required alignment of the given type in the current scope.
uint64_t getTypeABIAlignment(Type t) const;
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
index 92df593860c01a..7ee1ae5afd062b 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
@@ -187,9 +187,7 @@ LLVMArrayType::getTypeSizeInBits(const DataLayout &dataLayout,
llvm::TypeSize LLVMArrayType::getTypeSize(const DataLayout &dataLayout,
DataLayoutEntryListRef params) const {
- return llvm::alignTo(dataLayout.getTypeSize(getElementType()),
- dataLayout.getTypeABIAlignment(getElementType())) *
- getNumElements();
+ return dataLayout.getTypeAllocSize(getElementType()) * getNumElements();
}
uint64_t LLVMArrayType::getABIAlignment(const DataLayout &dataLayout,
@@ -563,8 +561,7 @@ LLVMStructType::getTypeSizeInBits(const DataLayout &dataLayout,
// Add padding to align the element unless the struct is packed.
structSize = llvm::alignTo(structSize, elementAlignment);
// Elements occupy their allocation size, even in packed structs.
- structSize += llvm::alignTo(dataLayout.getTypeSize(element),
- dataLayout.getTypeABIAlignment(element));
+ structSize += dataLayout.getTypeAllocSize(element);
// The alignment requirement of a struct is equal to the strictest alignment
// requirement of its elements.
diff --git a/mlir/lib/Dialect/OpenMP/Transforms/OpenMPOffloadPrivatizationPrepare.cpp b/mlir/lib/Dialect/OpenMP/Transforms/OpenMPOffloadPrivatizationPrepare.cpp
index d3182a51b3a3f1..144d6b605e624e 100644
--- a/mlir/lib/Dialect/OpenMP/Transforms/OpenMPOffloadPrivatizationPrepare.cpp
+++ b/mlir/lib/Dialect/OpenMP/Transforms/OpenMPOffloadPrivatizationPrepare.cpp
@@ -361,9 +361,7 @@ class PrepareForOMPOffloadPrivatizationPass
// Get the (compile-time constant) size of varType as per the
// given DataLayout dl.
std::int64_t getSizeInBytes(const DataLayout &dl, Type varType) const {
- llvm::TypeSize size = dl.getTypeSize(varType);
- unsigned short alignment = dl.getTypeABIAlignment(varType);
- return llvm::alignTo(size, alignment);
+ return dl.getTypeAllocSize(varType);
}
LLVM::LLVMFuncOp getMalloc(ModuleOp mod, IRRewriter &rewriter) const {
diff --git a/mlir/lib/Interfaces/DataLayoutInterfaces.cpp b/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
index 4aec320ddeff50..c00a2989ad35cc 100644
--- a/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
+++ b/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
@@ -576,6 +576,10 @@ llvm::TypeSize mlir::DataLayout::getTypeSizeInBits(Type t) const {
});
}
+llvm::TypeSize mlir::DataLayout::getTypeAllocSize(Type t) const {
+ return llvm::alignTo(getTypeSize(t), getTypeABIAlignment(t));
+}
+
uint64_t mlir::DataLayout::getTypeABIAlignment(Type t) const {
checkValid();
return cachedLookup<uint64_t>(t, abiAlignments, [&](Type ty) {
diff --git a/mlir/test/Dialect/LLVMIR/layout.mlir b/mlir/test/Dialect/LLVMIR/layout.mlir
index 2bdad3539a4f05..b838096c7614cd 100644
--- a/mlir/test/Dialect/LLVMIR/layout.mlir
+++ b/mlir/test/Dialect/LLVMIR/layout.mlir
@@ -228,6 +228,33 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
>} {
// CHECK-LABEL: @struct_element_allocation_size
func.func @struct_element_allocation_size() {
+ // The allocation size is the size rounded up to the ABI alignment.
+ // CHECK: alignment = 16
+ // CHECK-SAME: alloc_size = 16
+ // CHECK-SAME: bitsize = 80
+ // CHECK-SAME: size = 10
+ "test.data_layout_query"() : () -> f80
+
+ // CHECK: alignment = 4
+ // CHECK-SAME: alloc_size = 4
+ // CHECK-SAME: bitsize = 24
+ // CHECK-SAME: preferred = 8
+ // CHECK-SAME: size = 3
+ "test.data_layout_query"() : () -> i24
+
+ // Array elements are spaced by their allocation size.
+ // CHECK: alignment = 16
+ // CHECK-SAME: alloc_size = 32
+ // CHECK-SAME: bitsize = 256
+ // CHECK-SAME: size = 32
+ "test.data_layout_query"() : () -> !llvm.array<2 x f80>
+
+ // CHECK: alignment = 4
+ // CHECK-SAME: alloc_size = 12
+ // CHECK-SAME: bitsize = 96
+ // CHECK-SAME: size = 12
+ "test.data_layout_query"() : () -> !llvm.array<3 x i24>
+
// A packed element still occupies its allocation size, including padding.
// CHECK: alignment = 1
// CHECK-SAME: bitsize = 128
diff --git a/mlir/test/Interfaces/DataLayoutInterfaces/query.mlir b/mlir/test/Interfaces/DataLayoutInterfaces/query.mlir
index 97ef8b2a8ae1cc..4bc17afd246f23 100644
--- a/mlir/test/Interfaces/DataLayoutInterfaces/query.mlir
+++ b/mlir/test/Interfaces/DataLayoutInterfaces/query.mlir
@@ -39,6 +39,7 @@ func.func @no_layout_builtin() {
// CHECK: size = 16
"test.data_layout_query"() : () -> vector<4xi32>
// CHECK: alignment = 16
+ // CHECK: alloc_size = {minimal_size = 16 : index, scalable}
// CHECK: bitsize = {minimal_size = 128 : index, scalable}
// CHECK: index = 0
// CHECK: preferred = 16
diff --git a/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp b/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp
index 58c77cb94b47f8..bf53a0875571f6 100644
--- a/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp
+++ b/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp
@@ -38,6 +38,7 @@ struct TestDataLayoutQuery
const DataLayout &layout = layouts.getAbove(op);
llvm::TypeSize size = layout.getTypeSize(op.getType());
llvm::TypeSize bitsize = layout.getTypeSizeInBits(op.getType());
+ llvm::TypeSize allocSize = layout.getTypeAllocSize(op.getType());
uint64_t alignment = layout.getTypeABIAlignment(op.getType());
uint64_t preferred = layout.getTypePreferredAlignment(op.getType());
uint64_t index = layout.getTypeIndexBitwidth(op.getType()).value_or(0);
@@ -66,6 +67,7 @@ struct TestDataLayoutQuery
op->setAttrs(
{builder.getNamedAttr("size", convertTypeSizeToAttr(size)),
builder.getNamedAttr("bitsize", convertTypeSizeToAttr(bitsize)),
+ builder.getNamedAttr("alloc_size", convertTypeSizeToAttr(allocSize)),
builder.getNamedAttr("alignment", builder.getIndexAttr(alignment)),
builder.getNamedAttr("preferred", builder.getIndexAttr(preferred)),
builder.getNamedAttr("index", builder.getIndexAttr(index)),
diff --git a/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp b/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp
index f72c00070bd01a..fa61dc0cf114ff 100644
--- a/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp
+++ b/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp
@@ -511,6 +511,33 @@ module {}
EXPECT_EQ(layout.getManglingMode(), Attribute());
}
+TEST(DataLayout, AllocSize) {
+ const char *ir = R"MLIR(
+module {}
+ )MLIR";
+
+ DialectRegistry registry;
+ registry.insert<DLTIDialect, DLTestDialect>();
+ MLIRContext ctx(registry);
+
+ OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(ir, &ctx);
+ DataLayout layout(module.get());
+ Type i42 = IntegerType::get(&ctx, 42);
+ EXPECT_EQ(layout.getTypeSize(i42), 6u);
+ EXPECT_EQ(layout.getTypeAllocSize(i42), 8u);
+ EXPECT_EQ(layout.getTypeAllocSize(Float16Type::get(&ctx)), 2u);
+
+ Type i72 = IntegerType::get(&ctx, 72);
+ EXPECT_EQ(layout.getTypeSize(i72), 9u);
+ EXPECT_EQ(layout.getTypeABIAlignment(i72), 4u);
+ EXPECT_EQ(layout.getTypeAllocSize(i72), 12u);
+
+ auto scalableVec = VectorType::get({1}, i42, /*scalableDims=*/{true});
+ EXPECT_EQ(layout.getTypeSize(scalableVec), llvm::TypeSize::getScalable(6));
+ EXPECT_EQ(layout.getTypeAllocSize(scalableVec),
+ llvm::TypeSize::getScalable(8));
+}
+
TEST(DataLayout, NullSpec) {
const char *ir = R"MLIR(
"dltest.op_with_layout"() : () -> ()
@@ -631,6 +658,8 @@ TEST(DataLayout, SpecWithEntries) {
EXPECT_EQ(layout.getTypeSizeInBits(Float16Type::get(&ctx)), 48u);
EXPECT_EQ(layout.getTypeABIAlignment(IntegerType::get(&ctx, 42)), 8u);
EXPECT_EQ(layout.getTypeABIAlignment(Float16Type::get(&ctx)), 8u);
+ EXPECT_EQ(layout.getTypeAllocSize(IntegerType::get(&ctx, 42)), 8u);
+ EXPECT_EQ(layout.getTypeAllocSize(Float16Type::get(&ctx)), 8u);
EXPECT_EQ(layout.getTypePreferredAlignment(IntegerType::get(&ctx, 42)), 16u);
EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 16u);
EXPECT_EQ(layout.getTypeIndexBitwidth(Float16Type::get(&ctx)), std::nullopt);
>From 8707b62299c5c6422ab7ac27c61bd9359218e127 Mon Sep 17 00:00:00 2001
From: Adam Smith <adams at nvidia.com>
Date: Tue, 13 Oct 2026 08:07:08 -0700
Subject: [PATCH 2/2] [mlir] Check alloc_size for every type in query.mlir
Assisted-by: Cursor / Claude Opus 5.5
---
.../DataLayoutInterfaces/query.mlir | 30 +++++++++++++++++++
1 file changed, 30 insertions(+)
diff --git a/mlir/test/Interfaces/DataLayoutInterfaces/query.mlir b/mlir/test/Interfaces/DataLayoutInterfaces/query.mlir
index 4bc17afd246f23..24a76881da23cd 100644
--- a/mlir/test/Interfaces/DataLayoutInterfaces/query.mlir
+++ b/mlir/test/Interfaces/DataLayoutInterfaces/query.mlir
@@ -3,36 +3,42 @@
// CHECK-LABEL: @no_layout_builtin
func.func @no_layout_builtin() {
// CHECK: alignment = 4
+ // CHECK: alloc_size = 4
// CHECK: bitsize = 32
// CHECK: index = 0
// CHECK: preferred = 4
// CHECK: size = 4
"test.data_layout_query"() : () -> i32
// CHECK: alignment = 8
+ // CHECK: alloc_size = 8
// CHECK: bitsize = 64
// CHECK: index = 0
// CHECK: preferred = 8
// CHECK: size = 8
"test.data_layout_query"() : () -> f64
// CHECK: alignment = 4
+ // CHECK: alloc_size = 8
// CHECK: bitsize = 64
// CHECK: index = 0
// CHECK: preferred = 4
// CHECK: size = 8
"test.data_layout_query"() : () -> complex<f32>
// CHECK: alignment = 1
+ // CHECK: alloc_size = 2
// CHECK: bitsize = 14
// CHECK: index = 0
// CHECK: preferred = 1
// CHECK: size = 2
"test.data_layout_query"() : () -> complex<i6>
// CHECK: alignment = 4
+ // CHECK: alloc_size = 8
// CHECK: bitsize = 64
// CHECK: index = 64
// CHECK: preferred = 8
// CHECK: size = 8
"test.data_layout_query"() : () -> index
// CHECK: alignment = 16
+ // CHECK: alloc_size = 16
// CHECK: bitsize = 128
// CHECK: index = 0
// CHECK: preferred = 16
@@ -46,6 +52,7 @@ func.func @no_layout_builtin() {
// CHECK: size = {minimal_size = 16 : index, scalable}
"test.data_layout_query"() : () -> vector<[4]xi32>
// CHECK: alignment = 2
+ // CHECK: alloc_size = {minimal_size = 2 : index, scalable}
// CHECK: bitsize = {minimal_size = 16 : index, scalable}
// CHECK: index = 0
// CHECK: preferred = 2
@@ -58,6 +65,7 @@ func.func @no_layout_builtin() {
// CHECK-LABEL: @no_layout_custom
func.func @no_layout_custom() {
// CHECK: alignment = 1
+ // CHECK: alloc_size = 1
// CHECK: bitsize = 1
// CHECK: index = 1
// CHECK: preferred = 1
@@ -70,6 +78,7 @@ func.func @no_layout_custom() {
func.func @layout_op_no_layout() {
"test.op_with_data_layout"() ({
// CHECK: alignment = 1
+ // CHECK: alloc_size = 1
// CHECK: bitsize = 1
// CHECK: index = 1
// CHECK: preferred = 1
@@ -84,6 +93,7 @@ func.func @layout_op_no_layout() {
func.func @layout_op() {
"test.op_with_data_layout"() ({
// CHECK: alignment = 20
+ // CHECK: alloc_size = 20
// CHECK: bitsize = 10
// CHECK: index = 30
// CHECK: preferred = 1
@@ -104,6 +114,7 @@ func.func @nested_inner_only() {
"test.op_with_data_layout"() ({
"test.op_with_data_layout"() ({
// CHECK: alignment = 20
+ // CHECK: alloc_size = 20
// CHECK: bitsize = 10
// CHECK: index = 30
// CHECK: preferred = 1
@@ -126,6 +137,7 @@ func.func @nested_outer_only() {
"test.op_with_data_layout"() ({
"test.op_with_data_layout"() ({
// CHECK: alignment = 20
+ // CHECK: alloc_size = 20
// CHECK: bitsize = 10
// CHECK: index = 30
// CHECK: preferred = 1
@@ -148,6 +160,7 @@ func.func @nested_middle_only() {
"test.op_with_data_layout"() ({
"test.op_with_data_layout"() ({
// CHECK: alignment = 20
+ // CHECK: alloc_size = 20
// CHECK: bitsize = 10
// CHECK: index = 30
// CHECK: preferred = 1
@@ -172,6 +185,7 @@ func.func @nested_combine_with_missing() {
"test.op_with_data_layout"() ({
"test.op_with_data_layout"() ({
// CHECK: alignment = 20
+ // CHECK: alloc_size = 20
// CHECK: bitsize = 10
// CHECK: index = 21
// CHECK: preferred = 30
@@ -185,6 +199,7 @@ func.func @nested_combine_with_missing() {
#dlti.dl_entry<!test.test_type_with_layout<20>, ["alignment", 20]>
>} : () -> ()
// CHECK: alignment = 1
+ // CHECK: alloc_size = 6
// CHECK: bitsize = 42
// CHECK: index = 21
// CHECK: preferred = 30
@@ -205,6 +220,7 @@ func.func @nested_combine_all() {
"test.op_with_data_layout"() ({
"test.op_with_data_layout"() ({
// CHECK: alignment = 20
+ // CHECK: alloc_size = 20
// CHECK: bitsize = 3
// CHECK: index = 40
// CHECK: preferred = 30
@@ -216,6 +232,7 @@ func.func @nested_combine_all() {
#dlti.dl_entry<!test.test_type_with_layout<30>, ["preferred", 30]>
>} : () -> ()
// CHECK: alignment = 20
+ // CHECK: alloc_size = 20
// CHECK: bitsize = 10
// CHECK: index = 40
// CHECK: preferred = 30
@@ -228,6 +245,7 @@ func.func @nested_combine_all() {
#dlti.dl_entry<!test.test_type_with_layout<40>, ["index", 40]>
>} : () -> ()
// CHECK: alignment = 1
+ // CHECK: alloc_size = 6
// CHECK: bitsize = 42
// CHECK: index = 1
// CHECK: preferred = 30
@@ -245,21 +263,25 @@ func.func @nested_combine_all() {
func.func @integers() {
"test.op_with_data_layout"() ({
// CHECK: alignment = 8
+ // CHECK: alloc_size = 8
// CHECK: bitsize = 32
// CHECK: index = 0
// CHECK: preferred = 8
"test.data_layout_query"() : () -> i32
// CHECK: alignment = 16
+ // CHECK: alloc_size = 16
// CHECK: bitsize = 56
// CHECK: index = 0
// CHECK: preferred = 16
"test.data_layout_query"() : () -> i56
// CHECK: alignment = 16
+ // CHECK: alloc_size = 16
// CHECK: bitsize = 64
// CHECK: index = 0
// CHECK: preferred = 16
"test.data_layout_query"() : () -> i64
// CHECK: alignment = 16
+ // CHECK: alloc_size = 16
// CHECK: bitsize = 128
// CHECK: index = 0
// CHECK: preferred = 16
@@ -271,21 +293,25 @@ func.func @integers() {
>} : () -> ()
"test.op_with_data_layout"() ({
// CHECK: alignment = 8
+ // CHECK: alloc_size = 8
// CHECK: bitsize = 32
// CHECK: index = 0
// CHECK: preferred = 16
"test.data_layout_query"() : () -> i32
// CHECK: alignment = 16
+ // CHECK: alloc_size = 16
// CHECK: bitsize = 56
// CHECK: index = 0
// CHECK: preferred = 32
"test.data_layout_query"() : () -> i56
// CHECK: alignment = 16
+ // CHECK: alloc_size = 16
// CHECK: bitsize = 64
// CHECK: index = 0
// CHECK: preferred = 32
"test.data_layout_query"() : () -> i64
// CHECK: alignment = 16
+ // CHECK: alloc_size = 16
// CHECK: bitsize = 128
// CHECK: index = 0
// CHECK: preferred = 32
@@ -301,11 +327,13 @@ func.func @integers() {
func.func @floats() {
"test.op_with_data_layout"() ({
// CHECK: alignment = 8
+ // CHECK: alloc_size = 8
// CHECK: bitsize = 32
// CHECK: index = 0
// CHECK: preferred = 8
"test.data_layout_query"() : () -> f32
// CHECK: alignment = 16
+ // CHECK: alloc_size = 16
// CHECK: bitsize = 80
// CHECK: index = 0
// CHECK: preferred = 16
@@ -317,11 +345,13 @@ func.func @floats() {
>} : () -> ()
"test.op_with_data_layout"() ({
// CHECK: alignment = 8
+ // CHECK: alloc_size = 8
// CHECK: bitsize = 32
// CHECK: index = 0
// CHECK: preferred = 16
"test.data_layout_query"() : () -> f32
// CHECK: alignment = 16
+ // CHECK: alloc_size = 16
// CHECK: bitsize = 80
// CHECK: index = 0
// CHECK: preferred = 32
More information about the Mlir-commits
mailing list