[Mlir-commits] [mlir] [MLIR][LLVM][SROA] Make GEP handling type agnostic (PR #86950)

[Tobias Gysi]

================
@@ -431,29 +433,163 @@ 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 {
+/// 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
----------------
gysit wrote:

```suggestion
  // Helper to check that a constant index is in the bounds of the GEP index
```
? 

I think it may make sense to add a brief comment that LLVM dialect GEP indices have a limited bitwidth and we need to verify we are within these bounds when creating new GEP instructions.

https://github.com/llvm/llvm-project/pull/86950