[MLIR][LLVM][SROA] Make GEP handling type agnostic

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,29 +433,165 @@ 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(
     const MemorySlot &slot, SmallVectorImpl<MemorySlot> &mustBeSafelyUsed,
     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,60 +602,45 @@ 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;
 }
 
 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;
 }
 
 //===----------------------------------------------------------------------===//