[Mlir-commits] [mlir] [MLIR][LLVM][SROA] Make GEP handling type agnostic (PR #86950)
Christian Ulmann
llvmlistbot at llvm.org
Tue Apr 2 01:36:48 PDT 2024
https://github.com/Dinistro updated https://github.com/llvm/llvm-project/pull/86950
>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 1/3] [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
+}
>From a20d3db5aa3634f04981956a04fee792a519341e Mon Sep 17 00:00:00 2001
From: Christian Ulmann <christian.ulmann at nextsilicon.com>
Date: Thu, 28 Mar 2024 13:39:21 +0000
Subject: [PATCH 2/3] add GEP overflow support + test cleanup
---
mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp | 3 --
mlir/test/Dialect/LLVMIR/sroa.mlir | 38 +++++++++++++------
2 files changed, 26 insertions(+), 15 deletions(-)
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
index 1128166255764d..6420d9607b5448 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
@@ -457,9 +457,6 @@ static std::optional<uint64_t> gepToByteOffset(const DataLayout &dataLayout,
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();
diff --git a/mlir/test/Dialect/LLVMIR/sroa.mlir b/mlir/test/Dialect/LLVMIR/sroa.mlir
index 1b82514500a0d7..6ba693feb584d6 100644
--- a/mlir/test/Dialect/LLVMIR/sroa.mlir
+++ b/mlir/test/Dialect/LLVMIR/sroa.mlir
@@ -327,9 +327,9 @@ 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
+ %2 = llvm.getelementptr %1[8] : (!llvm.ptr) -> !llvm.ptr, i8
// CHECK-NEXT: llvm.store %[[ARG]], %[[ALLOCA]]
- llvm.store %arg, %7 : i32, !llvm.ptr
+ llvm.store %arg, %2 : i32, !llvm.ptr
llvm.return
}
@@ -341,9 +341,9 @@ 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)>
+ %2 = 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.store %arg, %2 : i32, !llvm.ptr
llvm.return
}
@@ -355,8 +355,8 @@ llvm.func @index_in_final_padding(%arg: 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
+ %2 = llvm.getelementptr %1[7] : (!llvm.ptr) -> !llvm.ptr, i8
+ llvm.store %arg, %2 : i32, !llvm.ptr
llvm.return
}
@@ -368,8 +368,8 @@ llvm.func @index_out_of_bounds(%arg: 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
+ %2 = llvm.getelementptr %1[9] : (!llvm.ptr) -> !llvm.ptr, i8
+ llvm.store %arg, %2 : i32, !llvm.ptr
llvm.return
}
@@ -381,8 +381,8 @@ llvm.func @index_in_padding(%arg: i16) {
// 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
+ %2 = llvm.getelementptr %1[2] : (!llvm.ptr) -> !llvm.ptr, i8
+ llvm.store %arg, %2 : i16, !llvm.ptr
llvm.return
}
@@ -394,9 +394,9 @@ 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
+ %2 = llvm.getelementptr %1[2] : (!llvm.ptr) -> !llvm.ptr, i8
// CHECK-NEXT: llvm.store %[[ARG]], %[[ALLOCA]]
- llvm.store %arg, %7 : i16, !llvm.ptr
+ llvm.store %arg, %2 : i16, !llvm.ptr
llvm.return
}
@@ -413,3 +413,17 @@ llvm.func @no_crash_on_negative_gep_index(%arg: f16) {
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
+}
>From 31835b321df0912f1a66903ab925b06368adaadf Mon Sep 17 00:00:00 2001
From: Christian Ulmann <christian.ulmann at nextsilicon.com>
Date: Tue, 2 Apr 2024 08:36:35 +0000
Subject: [PATCH 3/3] address review comments
---
mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp | 12 ++++++-----
mlir/test/Dialect/LLVMIR/sroa.mlir | 21 +++++++++++++++++++
2 files changed, 28 insertions(+), 5 deletions(-)
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
index 6420d9607b5448..06c1fdd2eb2d95 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
@@ -494,7 +494,8 @@ static std::optional<uint64_t> gepToByteOffset(const DataLayout &dataLayout,
}
namespace {
-/// Helper that contains information about accesses into a subslot.
+/// 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;
@@ -514,10 +515,11 @@ getSubslotAccessInfo(const DestructurableMemorySlot &slot,
if (!offset)
return {};
- // Helper to check that a constant index in the bounds of the GEP index
- // representation.
+ // 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);
+ return index >= (1 << LLVM::kGEPConstantBitWidth);
};
Type type = slot.elemType;
@@ -589,7 +591,7 @@ LogicalResult LLVM::GEPOp::ensureOnlySafeAccesses(
// Every access that remains in bounds of the remaining slot is considered
// legal.
mustBeSafelyUsed.emplace_back<MemorySlot>(
- {getBase(), getByteArrayType(getContext(), slotSize - *gepOffset)});
+ {getRes(), getByteArrayType(getContext(), slotSize - *gepOffset)});
return success();
}
diff --git a/mlir/test/Dialect/LLVMIR/sroa.mlir b/mlir/test/Dialect/LLVMIR/sroa.mlir
index 6ba693feb584d6..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 {
More information about the Mlir-commits
mailing list