[Mlir-commits] [mlir] [MLIR][LLVM][SROA] Make GEP handling type agnostic (PR #86950)
Christian Ulmann
llvmlistbot at llvm.org
Thu Mar 28 06:31:37 PDT 2024
https://github.com/Dinistro created https://github.com/llvm/llvm-project/pull/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.
>From 3d7f1bfd46778b4eeb42615fa6df9eb86a9f16b6 Mon Sep 17 00:00:00 2001
From: Christian Ulmann <christian.ulmann at nextsilicon.com>
Date: Thu, 28 Mar 2024 13:21:30 +0000
Subject: [PATCH] [MLIR][LLVM][SROA] Make GEP handling type agnostic
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.
---
mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp | 242 +++++++++++++-----
mlir/test/Dialect/LLVMIR/sroa.mlir | 95 +++++++
2 files changed, 278 insertions(+), 59 deletions(-)
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
+}
More information about the Mlir-commits
mailing list