[Mlir-commits] [mlir] 63d22f7 - [MLIR][LLVM][SROA] Make GEP handling type agnostic (#86950)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Apr 2 05:35:58 PDT 2024
Author: Christian Ulmann
Date: 2024-04-02T14:35:55+02:00
New Revision: 63d22f7a5b6afc515799f67c388bf5a8864274e4
URL: https://github.com/llvm/llvm-project/commit/63d22f7a5b6afc515799f67c388bf5a8864274e4
DIFF: https://github.com/llvm/llvm-project/commit/63d22f7a5b6afc515799f67c388bf5a8864274e4.diff
LOG: [MLIR][LLVM][SROA] Make GEP handling type agnostic (#86950)
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.
Added:
Modified:
mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
mlir/test/Dialect/LLVMIR/sroa.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
index f171bf7cc4bec3..06c1fdd2eb2d95 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,147 @@ 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) {
+ 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 {
+/// A struct that stores both the index into the aggregate type of the slot as
+/// well as the corresponding byte offset in memory.
+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 is in the bounds of the GEP index
+ // representation. LLVM dialects's GEP arguments have a limited bitwidth, thus
+ // this additional check is necessary.
+ 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 +581,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>(
+ {getRes(), getByteArrayType(getContext(), slotSize - *gepOffset)});
return success();
}
@@ -464,23 +602,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 +628,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..fe1531d988a4f5 100644
--- a/mlir/test/Dialect/LLVMIR/sroa.mlir
+++ b/mlir/test/Dialect/LLVMIR/sroa.mlir
@@ -82,6 +82,27 @@ llvm.func @multi_level_indirect() -> i32 {
// -----
+// This verifies that a nested GEP's users are checked properly. In this case
+// the load goes over the bounds of the memory slot and thus should block the
+// splitting of the alloca.
+
+// CHECK-LABEL: llvm.func @nested_access_over_slot_bound
+llvm.func @nested_access_over_slot_bound() -> i64 {
+ %0 = llvm.mlir.constant(1 : i32) : i32
+ // CHECK: %[[ALLOCA:.*]] = llvm.alloca %{{.*}} x !llvm.struct<(i32, struct<(
+ %1 = llvm.alloca %0 x !llvm.struct<(i32, struct<(array<10 x i32>)>, i32)> {alignment = 8 : i64} : (i32) -> !llvm.ptr
+ // CHECK: %[[GEP0:.*]] = llvm.getelementptr inbounds %[[ALLOCA]]
+ %2 = llvm.getelementptr inbounds %1[0, 1, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<(i32, struct<(array<10 x i32>)>, i32)>
+ // CHECK: %[[GEP1:.*]] = llvm.getelementptr inbounds %[[GEP0]]
+ %3 = llvm.getelementptr inbounds %2[0, 9] : (!llvm.ptr) -> !llvm.ptr, !llvm.array<10 x i32>
+ // CHECK: %[[RES:.*]] = llvm.load %[[GEP1]]
+ %4 = llvm.load %3 : !llvm.ptr -> i64
+ // CHECK: llvm.return %[[RES]] : i64
+ llvm.return %4 : i64
+}
+
+// -----
+
// CHECK-LABEL: llvm.func @resolve_alias
// CHECK-SAME: (%[[ARG:.*]]: i32)
llvm.func @resolve_alias(%arg: i32) -> i32 {
@@ -318,3 +339,112 @@ 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
+ %2 = llvm.getelementptr %1[8] : (!llvm.ptr) -> !llvm.ptr, i8
+ // CHECK-NEXT: llvm.store %[[ARG]], %[[ALLOCA]]
+ llvm.store %arg, %2 : 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
+ %2 = llvm.getelementptr %1[0, 1] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<(i32, i32)>
+ // CHECK-NEXT: llvm.store %[[ARG]], %[[ALLOCA]]
+ llvm.store %arg, %2 : 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
+ %2 = llvm.getelementptr %1[7] : (!llvm.ptr) -> !llvm.ptr, i8
+ llvm.store %arg, %2 : 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
+ %2 = llvm.getelementptr %1[9] : (!llvm.ptr) -> !llvm.ptr, i8
+ llvm.store %arg, %2 : 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
+ %2 = llvm.getelementptr %1[2] : (!llvm.ptr) -> !llvm.ptr, i8
+ llvm.store %arg, %2 : 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
+ %2 = llvm.getelementptr %1[2] : (!llvm.ptr) -> !llvm.ptr, i8
+ // CHECK-NEXT: llvm.store %[[ARG]], %[[ALLOCA]]
+ llvm.store %arg, %2 : 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
+}
+
+// -----
+
+// CHECK-LABEL: llvm.func @out_of_bound_gep_array_access
+// CHECK-SAME: %[[ARG:.*]]: i32
+llvm.func @out_of_bound_gep_array_access(%arg: i32) {
+ %0 = llvm.mlir.constant(1 : i32) : i32
+ // CHECK: %[[ALLOCA:.*]] = llvm.alloca %{{.*}} x i32
+ %1 = llvm.alloca %0 x !llvm.struct<"foo", (i32, i32)> : (i32) -> !llvm.ptr
+ %2 = llvm.getelementptr %1[0, 4] : (!llvm.ptr) -> !llvm.ptr, !llvm.array<4 x i8>
+ // CHECK-NEXT: llvm.store %[[ARG]], %[[ALLOCA]]
+ llvm.store %arg, %2 : i32, !llvm.ptr
+ llvm.return
+}
More information about the Mlir-commits
mailing list