[Mlir-commits] [mlir] [mlir][LLVM] Add LLVMStructType::getElementOffset() (PR #234171)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Oct 21 04:57:23 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Ivan Kosarev (kosarev)
<details>
<summary>Changes</summary>
Since bb6f7fcb46cd LLVMStructType::getTypeSizeInBits() uses allocation sizes, but there are a few places where we still compute field offsets based on store sizes.
This change factors out the computations into a single method, a counterpart of LLVM's StructLayout::getElementOffset(), so that all now agree in how they do that.
Also simplify the surrounding code a bit while there.
---
Full diff: https://github.com/llvm/llvm-project/pull/234171.diff
4 Files Affected:
- (modified) mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.td (+4)
- (modified) mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp (+15-43)
- (modified) mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp (+39-15)
- (modified) mlir/unittests/Dialect/LLVMIR/LLVMTypeTest.cpp (+33)
``````````diff
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.td
index 080f78e080d63..e31323651ceda 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.td
@@ -258,6 +258,10 @@ def LLVMStructType : LLVMType<"LLVMStruct", "struct", [
/// Checks if a struct is packed.
bool isPacked() const;
+ /// Returns the byte offset of the element at the given index.
+ uint64_t getElementOffset(const DataLayout &dataLayout,
+ unsigned index) const;
+
/// Checks if a struct is identified.
bool isIdentified() const;
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
index b76dc8647b91b..a67dea264c3a8 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
@@ -726,20 +726,8 @@ static std::optional<uint64_t> gepToByteOffset(const DataLayout &dataLayout,
return false;
})
.Case([&](LLVM::LLVMStructType structType) {
- ArrayRef<Type> body = structType.getBody();
- assert(index < body.size() && "expected valid struct indexing");
- for (uint32_t i : llvm::seq(index)) {
- if (!structType.isPacked())
- offset = llvm::alignTo(
- offset, dataLayout.getTypeABIAlignment(body[i]));
- offset += dataLayout.getTypeSize(body[i]);
- }
-
- // Align for the current type as well.
- if (!structType.isPacked())
- offset = llvm::alignTo(
- offset, dataLayout.getTypeABIAlignment(body[index]));
- currentType = body[index];
+ offset += structType.getElementOffset(dataLayout, index);
+ currentType = structType.getBody()[index];
return false;
})
.Default([&](Type type) {
@@ -800,31 +788,22 @@ getSubslotAccessInfo(const DestructurableMemorySlot &slot,
})
.Case([&](LLVM::LLVMStructType structType)
-> std::optional<SubslotAccessInfo> {
- uint64_t distanceToStart = 0;
// Walk over the elements of the struct to find in which of
// them the offset is.
for (auto [index, elem] : llvm::enumerate(structType.getBody())) {
- uint64_t elemSize = dataLayout.getTypeSize(elem);
- if (!structType.isPacked()) {
- distanceToStart = llvm::alignTo(
- distanceToStart, dataLayout.getTypeABIAlignment(elem));
- // If the offset is in padding, cancel the rewrite.
- if (offset < distanceToStart)
- return {};
- }
+ uint64_t elemOffset = structType.getElementOffset(dataLayout, index);
+ // If the offset is in padding, cancel the rewrite.
+ if (*offset < elemOffset)
+ return {};
- if (offset < distanceToStart + elemSize) {
+ if (*offset < elemOffset + dataLayout.getTypeSize(elem)) {
if (isOutOfBoundsGEPIndex(index))
return {};
// The offset is within this element, stop iterating the
// struct and return the index.
return SubslotAccessInfo{static_cast<uint32_t>(index),
- *offset - distanceToStart};
+ *offset - elemOffset};
}
-
- // The offset is not within this element, continue walking
- // over the struct.
- distanceToStart += elemSize;
}
return {};
@@ -1145,35 +1124,28 @@ memsetRewire(MemsetIntr op, const DestructurableMemorySlot &slot,
IntegerAttr memsetLenAttr = createMemsetLenAttr(op);
- bool packed = false;
- if (auto structType = dyn_cast<LLVM::LLVMStructType>(slot.valueType))
- packed = structType.isPacked();
+ auto structType = dyn_cast<LLVM::LLVMStructType>(slot.valueType);
Type i32 = IntegerType::get(op.getContext(), 32);
uint64_t memsetLen = memsetLenAttr.getValue().getZExtValue();
- uint64_t covered = 0;
for (size_t i = 0; i < types->size(); i++) {
// Create indices on the fly to get elements in the right order.
Attribute index = IntegerAttr::get(i32, i);
Type elemType = types->at(index);
uint64_t typeSize = dataLayout.getTypeSize(elemType);
-
- if (!packed)
- covered =
- llvm::alignTo(covered, dataLayout.getTypeABIAlignment(elemType));
-
- if (covered >= memsetLen)
+ uint64_t offset =
+ structType
+ ? structType.getElementOffset(dataLayout, i)
+ : i * dataLayout.getTypeAllocSize(elemType);
+ if (offset >= memsetLen)
break;
// If this subslot is used, apply a new memset to it.
- // Otherwise, only compute its offset within the original memset.
if (subslots.contains(index)) {
- uint64_t newMemsetSize = std::min(memsetLen - covered, typeSize);
+ uint64_t newMemsetSize = std::min(memsetLen - offset, typeSize);
createMemsetIntr(builder, op, memsetLenAttr, newMemsetSize, subslots,
index);
}
-
- covered += typeSize;
}
return DeletionKind::Delete;
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
index 7ee1ae5afd062..2cc2db897b7f4 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
@@ -550,27 +550,51 @@ LLVMStructType::verifyInvariants(function_ref<InFlightDiagnostic()> emitError,
return success();
}
+/// Returns the byte offset of the element of the specified structure at the
+/// given index and the strictest alignment requirement of the elements up to
+/// the index. For the index past the last element it thus returns the size and
+/// alignment of the whole structure.
+static uint64_t getOffsetAndAlignment(const DataLayout &dataLayout,
+ LLVMStructType structType, unsigned index,
+ uint64_t &structAlignment) {
+ uint64_t offset = 0;
+ structAlignment = 1;
+ for (auto [i, element] : llvm::enumerate(structType.getBody())) {
+ // Align the element unless the structure is packed.
+ uint64_t elementAlignment = dataLayout.getTypeABIAlignment(element);
+ if (!structType.isPacked()) {
+ offset = llvm::alignTo(offset, elementAlignment);
+
+ // The structure alignment is the strictest alignment of its
+ // elements.
+ structAlignment = std::max(structAlignment, elementAlignment);
+ }
+ if (i == index)
+ return offset;
+
+ // Each element gets its allocation size, even in packed structures.
+ offset += dataLayout.getTypeAllocSize(element);
+ }
+ return offset;
+}
+
+uint64_t LLVMStructType::getElementOffset(const DataLayout &dataLayout,
+ unsigned index) const {
+ assert(index < getBody().size() && "expected a valid struct element index");
+ uint64_t structAlignment;
+ return getOffsetAndAlignment(dataLayout, *this, index, structAlignment);
+}
+
llvm::TypeSize
LLVMStructType::getTypeSizeInBits(const DataLayout &dataLayout,
DataLayoutEntryListRef params) const {
- auto structSize = llvm::TypeSize::getFixed(0);
- uint64_t structAlignment = 1;
- for (Type element : getBody()) {
- uint64_t elementAlignment =
- isPacked() ? 1 : dataLayout.getTypeABIAlignment(element);
- // 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 += dataLayout.getTypeAllocSize(element);
-
- // The alignment requirement of a struct is equal to the strictest alignment
- // requirement of its elements.
- structAlignment = std::max(elementAlignment, structAlignment);
- }
+ uint64_t structAlignment;
+ uint64_t structSize = getOffsetAndAlignment(
+ dataLayout, *this, getBody().size(), structAlignment);
// At the end, add padding to the struct to satisfy its own alignment
// requirement. Otherwise structs inside of arrays would be misaligned.
structSize = llvm::alignTo(structSize, structAlignment);
- return structSize * kBitsInByte;
+ return llvm::TypeSize::getFixed(structSize * kBitsInByte);
}
namespace {
diff --git a/mlir/unittests/Dialect/LLVMIR/LLVMTypeTest.cpp b/mlir/unittests/Dialect/LLVMIR/LLVMTypeTest.cpp
index df19caec69397..72d8867d232aa 100644
--- a/mlir/unittests/Dialect/LLVMIR/LLVMTypeTest.cpp
+++ b/mlir/unittests/Dialect/LLVMIR/LLVMTypeTest.cpp
@@ -27,3 +27,36 @@ TEST_F(LLVMIRDeathTest, StructTypeGetVerifiesElementTypes) {
"invalid LLVM structure element type");
}
#endif
+
+TEST_F(LLVMIRTest, StructElementOffsets) {
+ DataLayout dataLayout;
+ Type i8 = IntegerType::get(&context, 8);
+ Type i32 = IntegerType::get(&context, 32);
+ Type f64 = Float64Type::get(&context);
+ Type f80 = Float80Type::get(&context);
+
+ // Check ABI alignment for unpacked structures.
+ auto unpacked = LLVMStructType::getLiteral(&context, {i8, f64, i32});
+ EXPECT_EQ(unpacked.getElementOffset(dataLayout, 0), 0u);
+ EXPECT_EQ(unpacked.getElementOffset(dataLayout, 1), 8u);
+ EXPECT_EQ(unpacked.getElementOffset(dataLayout, 2), 16u);
+ EXPECT_EQ(dataLayout.getTypeSize(unpacked), 24u);
+
+ // Check no alignment for packed structures.
+ auto packed = LLVMStructType::getLiteral(&context, {i8, f64, i32},
+ /*isPacked=*/true);
+ EXPECT_EQ(packed.getElementOffset(dataLayout, 0), 0u);
+ EXPECT_EQ(packed.getElementOffset(dataLayout, 1), 1u);
+ EXPECT_EQ(packed.getElementOffset(dataLayout, 2), 9u);
+ EXPECT_EQ(dataLayout.getTypeSize(packed), 13u);
+
+ // Check that fields get their allocation size, even in packed
+ // structures.
+ auto alloc = LLVMStructType::getLiteral(&context, {f80, i8});
+ EXPECT_EQ(alloc.getElementOffset(dataLayout, 1), 16u);
+ EXPECT_EQ(dataLayout.getTypeSize(alloc), 32u);
+ auto packedAlloc =
+ LLVMStructType::getLiteral(&context, {f80, i8}, /*isPacked=*/true);
+ EXPECT_EQ(packedAlloc.getElementOffset(dataLayout, 1), 16u);
+ EXPECT_EQ(dataLayout.getTypeSize(packedAlloc), 17u);
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/234171
More information about the Mlir-commits
mailing list