[MLIR][LLVM][SROA] Make GEP handling type agnostic #86950
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This commit removes SROA's type consistency constraints from LLVM dialect's GEPOp. The checks for valid indexing are now purely done by computing the GEP's offset with the aid of the data layout. To simplify handling of nested "subslots", we are tricking the SROA by handing in memory slots that hold byte array types. This ensures that subsequent accesses only need to check if their access will be in-bounds. This lifts the requirement of detemining the sub-types for all but the first level of subslots.
|
@llvm/pr-subscribers-mlir @llvm/pr-subscribers-mlir-llvm Author: Christian Ulmann (Dinistro) ChangesThis commit removes SROA's type consistency constraints from LLVM dialect's GEPOp. The checks for valid indexing are now purely done by computing the GEP's offset with the aid of the data layout. To simplify handling of "nested subslots", we are tricking the SROA by handing in memory slots that hold byte array types. This ensures that subsequent accesses only need to check if their access will be in-bounds. This lifts the requirement of determining the sub-types for all but the first level of subslots. Full diff: https://github.com/llvm/llvm-project/pull/86950.diff 2 Files Affected:
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
index f171bf7cc4bec3..1128166255764d 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
@@ -20,6 +20,8 @@
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/TypeSwitch.h"
+#define DEBUG_TYPE "sroa"
+
using namespace mlir;
//===----------------------------------------------------------------------===//
@@ -431,10 +433,148 @@ DeletionKind LLVM::GEPOp::removeBlockingUses(
return DeletionKind::Delete;
}
-static bool isFirstIndexZero(LLVM::GEPOp gep) {
- IntegerAttr index =
- llvm::dyn_cast_if_present<IntegerAttr>(gep.getIndices()[0]);
- return index && index.getInt() == 0;
+/// Returns the amount of bytes the provided GEP elements will offset the
+/// pointer by. Returns nullopt if no constant offset could be computed.
+static std::optional<uint64_t> gepToByteOffset(const DataLayout &dataLayout,
+ LLVM::GEPOp gep) {
+ // Collects all indices.
+ SmallVector<uint64_t> indices;
+ for (auto index : gep.getIndices()) {
+ auto constIndex = dyn_cast<IntegerAttr>(index);
+ if (!constIndex)
+ return {};
+ int64_t gepIndex = constIndex.getInt();
+ // Negative indices are not supported.
+ if (gepIndex < 0)
+ return {};
+ indices.push_back(gepIndex);
+ }
+
+ Type currentType = gep.getElemType();
+ uint64_t offset = indices[0] * dataLayout.getTypeSize(currentType);
+
+ for (uint64_t index : llvm::drop_begin(indices)) {
+ bool shouldCancel =
+ TypeSwitch<Type, bool>(currentType)
+ .Case([&](LLVM::LLVMArrayType arrayType) {
+ // TODO: Support out-of-bounds accesses.
+ if (arrayType.getNumElements() <= index)
+ return true;
+ offset +=
+ index * dataLayout.getTypeSize(arrayType.getElementType());
+ currentType = arrayType.getElementType();
+ 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];
+ return false;
+ })
+ .Default([&](Type type) {
+ LLVM_DEBUG(llvm::dbgs()
+ << "[sroa] Unsupported type for offset computations"
+ << type << "\n");
+ return true;
+ });
+
+ if (shouldCancel)
+ return std::nullopt;
+ }
+
+ return offset;
+}
+
+namespace {
+/// Helper that contains information about accesses into a subslot.
+struct SubslotAccessInfo {
+ /// The parent slot's index that the access falls into.
+ uint32_t index;
+ /// The offset into the subslot of the access.
+ uint64_t subslotOffset;
+};
+} // namespace
+
+/// Computes subslot access information for an access into `slot` with the given
+/// offset.
+/// Returns nullopt when the offset is out-of-bounds or when the access is into
+/// the padding of `slot`.
+static std::optional<SubslotAccessInfo>
+getSubslotAccessInfo(const DestructurableMemorySlot &slot,
+ const DataLayout &dataLayout, LLVM::GEPOp gep) {
+ std::optional<uint64_t> offset = gepToByteOffset(dataLayout, gep);
+ if (!offset)
+ return {};
+
+ // Helper to check that a constant index in the bounds of the GEP index
+ // representation.
+ auto isOutOfBoundsGEPIndex = [](uint64_t index) {
+ return index > (1 << LLVM::kGEPConstantBitWidth);
+ };
+
+ Type type = slot.elemType;
+ if (*offset >= dataLayout.getTypeSize(type))
+ return {};
+ return TypeSwitch<Type, std::optional<SubslotAccessInfo>>(type)
+ .Case([&](LLVM::LLVMArrayType arrayType)
+ -> std::optional<SubslotAccessInfo> {
+ // Find which element of the array contains the offset.
+ uint64_t elemSize = dataLayout.getTypeSize(arrayType.getElementType());
+ uint64_t index = *offset / elemSize;
+ if (isOutOfBoundsGEPIndex(index))
+ return {};
+ return SubslotAccessInfo{static_cast<uint32_t>(index),
+ *offset - (index * elemSize)};
+ })
+ .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 {};
+ }
+
+ if (offset < distanceToStart + elemSize) {
+ 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};
+ }
+
+ // The offset is not within this element, continue walking
+ // over the struct.
+ distanceToStart += elemSize;
+ }
+
+ return {};
+ });
+}
+
+/// Constructs a byte array type of the given size.
+static LLVM::LLVMArrayType getByteArrayType(MLIRContext *context,
+ unsigned size) {
+ auto byteType = IntegerType::get(context, 8);
+ return LLVM::LLVMArrayType::get(context, byteType, size);
}
LogicalResult LLVM::GEPOp::ensureOnlySafeAccesses(
@@ -442,18 +582,17 @@ LogicalResult LLVM::GEPOp::ensureOnlySafeAccesses(
const DataLayout &dataLayout) {
if (getBase() != slot.ptr)
return success();
- if (slot.elemType != getElemType())
- return failure();
- if (!isFirstIndexZero(*this))
+ std::optional<uint64_t> gepOffset = gepToByteOffset(dataLayout, *this);
+ if (!gepOffset)
return failure();
- // Dynamic indices can be out-of-bounds (even negative), so an access with
- // dynamic indices can never be considered safe.
- if (!getDynamicIndices().empty())
+ uint64_t slotSize = dataLayout.getTypeSize(slot.elemType);
+ // Check that the access is strictly inside the slot.
+ if (*gepOffset >= slotSize)
return failure();
- Type reachedType = getResultPtrElementType();
- if (!reachedType)
- return failure();
- mustBeSafelyUsed.emplace_back<MemorySlot>({getResult(), reachedType});
+ // Every access that remains in bounds of the remaining slot is considered
+ // legal.
+ mustBeSafelyUsed.emplace_back<MemorySlot>(
+ {getBase(), getByteArrayType(getContext(), slotSize - *gepOffset)});
return success();
}
@@ -464,23 +603,25 @@ bool LLVM::GEPOp::canRewire(const DestructurableMemorySlot &slot,
if (!isa<LLVM::LLVMPointerType>(getBase().getType()))
return false;
- if (getBase() != slot.ptr || slot.elemType != getElemType())
- return false;
- if (!isFirstIndexZero(*this))
- return false;
- // Dynamic indices can be out-of-bounds (even negative), so an access with
- // dynamic indices can never be properly rewired.
- if (!getDynamicIndices().empty())
- return false;
- Type reachedType = getResultPtrElementType();
- if (!reachedType || getIndices().size() < 2)
+ if (getBase() != slot.ptr)
return false;
- auto firstLevelIndex = dyn_cast<IntegerAttr>(getIndices()[1]);
- if (!firstLevelIndex)
+ std::optional<SubslotAccessInfo> accessInfo =
+ getSubslotAccessInfo(slot, dataLayout, *this);
+ if (!accessInfo)
return false;
- mustBeSafelyUsed.emplace_back<MemorySlot>({getResult(), reachedType});
- assert(slot.elementPtrs.contains(firstLevelIndex));
- usedIndices.insert(firstLevelIndex);
+ auto indexAttr =
+ IntegerAttr::get(IntegerType::get(getContext(), 32), accessInfo->index);
+ assert(slot.elementPtrs.contains(indexAttr));
+ usedIndices.insert(indexAttr);
+
+ // The remainder of the subslot should be accesses in-bounds. Thus, we create
+ // a dummy slot with the size of the remainder.
+ Type subslotType = slot.elementPtrs.lookup(indexAttr);
+ uint64_t slotSize = dataLayout.getTypeSize(subslotType);
+ LLVM::LLVMArrayType remainingSlotType =
+ getByteArrayType(getContext(), slotSize - accessInfo->subslotOffset);
+ mustBeSafelyUsed.emplace_back<MemorySlot>({getRes(), remainingSlotType});
+
return true;
}
@@ -488,36 +629,19 @@ DeletionKind LLVM::GEPOp::rewire(const DestructurableMemorySlot &slot,
DenseMap<Attribute, MemorySlot> &subslots,
RewriterBase &rewriter,
const DataLayout &dataLayout) {
- IntegerAttr firstLevelIndex =
- llvm::dyn_cast_if_present<IntegerAttr>(getIndices()[1]);
- const MemorySlot &newSlot = subslots.at(firstLevelIndex);
-
- ArrayRef<int32_t> remainingIndices = getRawConstantIndices().slice(2);
-
- // If the GEP would become trivial after this transformation, eliminate it.
- // A GEP should only be eliminated if it has no indices (except the first
- // pointer index), as simplifying GEPs with all-zero indices would eliminate
- // structure information useful for further destruction.
- if (remainingIndices.empty()) {
- rewriter.replaceAllUsesWith(getResult(), newSlot.ptr);
- return DeletionKind::Delete;
- }
-
- rewriter.modifyOpInPlace(*this, [&]() {
- // Rewire the indices by popping off the second index.
- // Start with a single zero, then add the indices beyond the second.
- SmallVector<int32_t> newIndices(1);
- newIndices.append(remainingIndices.begin(), remainingIndices.end());
- setRawConstantIndices(newIndices);
-
- // Rewire the pointed type.
- setElemType(newSlot.elemType);
-
- // Rewire the pointer.
- getBaseMutable().assign(newSlot.ptr);
- });
-
- return DeletionKind::Keep;
+ std::optional<SubslotAccessInfo> accessInfo =
+ getSubslotAccessInfo(slot, dataLayout, *this);
+ assert(accessInfo && "expected access info to be checked before");
+ auto indexAttr =
+ IntegerAttr::get(IntegerType::get(getContext(), 32), accessInfo->index);
+ const MemorySlot &newSlot = subslots.at(indexAttr);
+
+ auto byteType = IntegerType::get(rewriter.getContext(), 8);
+ auto newPtr = rewriter.createOrFold<LLVM::GEPOp>(
+ getLoc(), getResult().getType(), byteType, newSlot.ptr,
+ ArrayRef<GEPArg>(accessInfo->subslotOffset), getInbounds());
+ rewriter.replaceAllUsesWith(getResult(), newPtr);
+ return DeletionKind::Delete;
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/LLVMIR/sroa.mlir b/mlir/test/Dialect/LLVMIR/sroa.mlir
index 3f4d17c6a43f97..1b82514500a0d7 100644
--- a/mlir/test/Dialect/LLVMIR/sroa.mlir
+++ b/mlir/test/Dialect/LLVMIR/sroa.mlir
@@ -318,3 +318,98 @@ llvm.func @store_to_memory(%arg: !llvm.ptr) {
llvm.store %1, %arg : !llvm.ptr, !llvm.ptr
llvm.return
}
+
+// -----
+
+// CHECK-LABEL: llvm.func @type_mismatch_array_access
+// CHECK-SAME: %[[ARG:.*]]: i32
+llvm.func @type_mismatch_array_access(%arg: i32) {
+ %0 = llvm.mlir.constant(1 : i32) : i32
+ // CHECK: %[[ALLOCA:.*]] = llvm.alloca %{{.*}} x i32
+ %1 = llvm.alloca %0 x !llvm.struct<(i32, i32, i32)> : (i32) -> !llvm.ptr
+ %7 = llvm.getelementptr %1[8] : (!llvm.ptr) -> !llvm.ptr, i8
+ // CHECK-NEXT: llvm.store %[[ARG]], %[[ALLOCA]]
+ llvm.store %arg, %7 : i32, !llvm.ptr
+ llvm.return
+}
+
+// -----
+
+// CHECK-LABEL: llvm.func @type_mismatch_struct_access
+// CHECK-SAME: %[[ARG:.*]]: i32
+llvm.func @type_mismatch_struct_access(%arg: i32) {
+ %0 = llvm.mlir.constant(1 : i32) : i32
+ // CHECK: %[[ALLOCA:.*]] = llvm.alloca %{{.*}} x i32
+ %1 = llvm.alloca %0 x !llvm.struct<(i32, i32, i32)> : (i32) -> !llvm.ptr
+ %7 = llvm.getelementptr %1[0, 1] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<(i32, i32)>
+ // CHECK-NEXT: llvm.store %[[ARG]], %[[ALLOCA]]
+ llvm.store %arg, %7 : i32, !llvm.ptr
+ llvm.return
+}
+
+// -----
+
+// CHECK-LABEL: llvm.func @index_in_final_padding
+llvm.func @index_in_final_padding(%arg: i32) {
+ %0 = llvm.mlir.constant(1 : i32) : i32
+ // CHECK: %[[ALLOCA:.*]] = llvm.alloca %{{.*}} x !llvm.struct<"foo", (i32, i8)>
+ %1 = llvm.alloca %0 x !llvm.struct<"foo", (i32, i8)> : (i32) -> !llvm.ptr
+ // CHECK: = llvm.getelementptr %[[ALLOCA]][7] : (!llvm.ptr) -> !llvm.ptr, i8
+ %7 = llvm.getelementptr %1[7] : (!llvm.ptr) -> !llvm.ptr, i8
+ llvm.store %arg, %7 : i32, !llvm.ptr
+ llvm.return
+}
+
+// -----
+
+// CHECK-LABEL: llvm.func @index_out_of_bounds
+llvm.func @index_out_of_bounds(%arg: i32) {
+ %0 = llvm.mlir.constant(1 : i32) : i32
+ // CHECK: %[[ALLOCA:.*]] = llvm.alloca %{{.*}} x !llvm.struct<"foo", (i32, i32)>
+ %1 = llvm.alloca %0 x !llvm.struct<"foo", (i32, i32)> : (i32) -> !llvm.ptr
+ // CHECK: = llvm.getelementptr %[[ALLOCA]][9] : (!llvm.ptr) -> !llvm.ptr, i8
+ %7 = llvm.getelementptr %1[9] : (!llvm.ptr) -> !llvm.ptr, i8
+ llvm.store %arg, %7 : i32, !llvm.ptr
+ llvm.return
+}
+
+// -----
+
+// CHECK-LABEL: llvm.func @index_in_padding
+llvm.func @index_in_padding(%arg: i16) {
+ %0 = llvm.mlir.constant(1 : i32) : i32
+ // CHECK: %[[ALLOCA:.*]] = llvm.alloca %{{.*}} x !llvm.struct<"foo", (i16, i32)>
+ %1 = llvm.alloca %0 x !llvm.struct<"foo", (i16, i32)> : (i32) -> !llvm.ptr
+ // CHECK: = llvm.getelementptr %[[ALLOCA]][2] : (!llvm.ptr) -> !llvm.ptr, i8
+ %7 = llvm.getelementptr %1[2] : (!llvm.ptr) -> !llvm.ptr, i8
+ llvm.store %arg, %7 : i16, !llvm.ptr
+ llvm.return
+}
+
+// -----
+
+// CHECK-LABEL: llvm.func @index_not_in_padding_because_packed
+// CHECK-SAME: %[[ARG:.*]]: i16
+llvm.func @index_not_in_padding_because_packed(%arg: i16) {
+ %0 = llvm.mlir.constant(1 : i32) : i32
+ // CHECK: %[[ALLOCA:.*]] = llvm.alloca %{{.*}} x i32
+ %1 = llvm.alloca %0 x !llvm.struct<"foo", packed (i16, i32)> : (i32) -> !llvm.ptr
+ %7 = llvm.getelementptr %1[2] : (!llvm.ptr) -> !llvm.ptr, i8
+ // CHECK-NEXT: llvm.store %[[ARG]], %[[ALLOCA]]
+ llvm.store %arg, %7 : i16, !llvm.ptr
+ llvm.return
+}
+
+// -----
+
+// CHECK-LABEL: llvm.func @no_crash_on_negative_gep_index
+// CHECK-SAME: %[[ARG:.*]]: f16
+llvm.func @no_crash_on_negative_gep_index(%arg: f16) {
+ %0 = llvm.mlir.constant(1 : i32) : i32
+ // CHECK: %[[ALLOCA:.*]] = llvm.alloca %{{.*}} x !llvm.struct<"foo", (i32, i32, i32)>
+ %1 = llvm.alloca %0 x !llvm.struct<"foo", (i32, i32, i32)> : (i32) -> !llvm.ptr
+ // CHECK: llvm.getelementptr %[[ALLOCA]][-1] : (!llvm.ptr) -> !llvm.ptr, f32
+ %2 = llvm.getelementptr %1[-1] : (!llvm.ptr) -> !llvm.ptr, f32
+ llvm.store %arg, %2 : f16, !llvm.ptr
+ llvm.return
+}
|
gysit
reviewed
Mar 28, 2024
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This commit removes SROA's type consistency constraints from LLVM dialect's GEPOp. The checks for valid indexing are now purely done by computing the GEP's offset with the aid of the data layout.
To simplify handling of "nested subslots", we are tricking the SROA by handing in memory slots that hold byte array types. This ensures that subsequent accesses only need to check if their access will be in-bounds. This lifts the requirement of determining the sub-types for all but the first level of subslots.