[mlir] Add a generic SROA implementation.

[Moxinilian]

Summary

This revision introduces a generic implementation of Scalar Replacement
Of Aggregates. In contrast to the implementation in LLVM, this focuses
on the core of SROA: destructuring aggregates. By implementing
interfaces on allocators and accessors, memory allocators can be
destructured into smaller allocators, through the MemorySlot
abstraction.

This pass only works on aggregates that are accessed in a "type-safe"
way, that is within the bounds and respecting the type of a given memory
slot. The destructuring pattern and functions only peel off the first
layer of aggregates and can safely be applied repeatedly. For
convenience, the transformation is also available as a pass that will
apply the pattern repeatedly.

Depends on D149958

[tschuett]

Thanks. I am interested in the pass and of course mem2reg as a pass. My only annoyance is that I have to run the pass several times.

[Dinistro]

I added a bunch of comments. The pass already looks great.

Make sure to update the commit message, as the pass applies the pattern until convergence, not just once.

mlir/include/mlir/Interfaces/MemorySlotInterfaces.h

28

Ultra NIT: I prefer explicit public inheritance

mlir/include/mlir/Interfaces/MemorySlotInterfaces.td

263

This is a very general name for an interface. We should consider to rename this to something that ties it to the MemorySlot abstraction.

351

Only allowing to use attributes to index into the list of types seems a bit cumbersome for other usages.
Adding a second InterfaceMethod that allows you to access this with an unsigned might be nice to have.

The function with an attribute parameter could then even have a default implementation that dispatches to the other one.

mlir/include/mlir/Transforms/SROA.h

19

NIT: Top-level comment.

28

NIT: llvm::Optional is deprecated.

mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp

72

NIT: destructuredTypes or destructuredTypeMap

283

NIT: Top level comment

295

std::optional

301

NIT: Drop auto, as the type is not directly clear.

344

According to the GEPOp definition, the type of base can also be a vector. In that case, the cast would assert.
Maybe adding a test for this case would also be nice.

396

std::optional

400

Ultra NIT: Usually, const is used in front of the type of a reference, i.e., const auto &. There is no semantic difference, but it is what is mainly used in the MLIR codebase.

mlir/lib/Transforms/SROA.cpp

26

std::optional

35

NIT: Can you use the SmallPtrSetImpl type here?

157

NIT: use notifyMatchFailure() instead.

177

IIRC, the PatternRewriter class inherits from OpBuilder. If there is a reason why you cannot use the rewriter, comment it.

[Moxinilian]

Thank you for the review!

@tschuett Sorry, that was not clear. While the pattern only peels off one layer of aggregate (as it would typically be used recursively in pattern application), the pass applies it repeatedly. I updated the summary accordingly.

mlir/include/mlir/Interfaces/MemorySlotInterfaces.td

263

I went with SafeMemorySlotAccessOpInterface, does that sound good?

351

I'm worried about this. LLVM is so far the only use case I need to implement that uses integers directly. In MemRef you want to use multi-dimensional indices, not just integers. And then in an hypothetical hashmap dialect you may want to use strings or any attribute really. Blessing the "integer-only" use case feels very ad hoc to me, especially considering implementers will have to handle Attribute anyway because they may accidentally receive or want to process Attribute indices no matter what. Could you be more specific of the approach you want for this?

mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp

72

So the intent was that this represents the type in its destructured form. I went for subelementIndexMap then.

344

Great catch. For the test, I am not sure, considering this only gets invoked currently when the base is the result of an alloca, which must be a pointer. Any idea?

[Moxinilian]

Address comments.

[tschuett]

Thanks. Then I will take the pass.

[Moxinilian]

Improve destructuring algorithm in AllocaOp.

[gysit]

Very nice work!

mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h

205

There are still at least two uses of Optional. Please search&replace with std::optional.

mlir/include/mlir/Interfaces/MemorySlotInterfaces.td

231

249

nit: I assume usedIndices an input only? Could it be const and or documented?

mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp

425

Should size limit be a pass parameter and also apply to structs? This maybe part of a followup revision.

mlir/lib/Transforms/SROA.cpp

17

nit: Can you check the includes? I don't think there TypeSwitch is used in this file?

125

nit: Drop llvm:: here and in llvm::reverse below (I believe).

146

I would define a SROA DEBUG_TYPE as in the mem2reg revision.

172

nit: Let's use *info to access optional values.

[Dinistro]

mlir/include/mlir/Interfaces/MemorySlotInterfaces.td

263

Sounds good to me.

351

Oh, I didn't realize that the index can be multidimensional. In that case it can probably stay an attribute. Just be aware that attributes are always uniqued and aren't cleaned up, even if they are no longer used.

mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp

344

In that case you cannot add a test, I guess.

[Moxinilian]

Rebase, add statistics and improve pattern.

[Moxinilian]

Address comments.

mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp

425

Yes probably. The ideal maximum size would be +inf, the limit being only for practical reasons. I think it is fine that the limit does not apply to structs for example, as you cannot create a struct with a trillion fields as easily as an array with a trillion elements. Let's save that for a follow up if it proves useful, though.

mlir/lib/Transforms/SROA.cpp

125

So I kind of prefer having llvm:: in front of this sort of utilities, like topological sort, to signify that this comes from the infra and is not a local concept (so the reader does not assume enumerate has some form of contextual local logic). Does that sound good to you? Otherwise I can drop them, it's fine.

[gysit]

Nice we are almost there! I have some nit comment but am otherwise happy!

The build bot seems to be unable to apply the patch. I assume a rebase should fix this.

mlir/include/mlir/Interfaces/MemorySlotInterfaces.td

243–244

nit: I would try to make the last sentence a bit more direct, e.g., "Only generate subslots for the indices found in usedIndices since all other subslots are unused"

mlir/include/mlir/Transforms/Passes.td

251

nit: Total number of memory slots (since they seem countable?)

Note this appears in several places!

256

And if possible try to make the lines shorter to stay within the 80 columns. I expect [{}] should work.

261

In think I would also use number / Maximal number?

mlir/lib/Transforms/SROA.cpp

192–193

nit: Would it make sense to move this check into computeDestructuringInfo to make it more self contained.

197

nit: I would suggest to inverse the condition to save on indent in the body of the current if, i.e.:

if (!info)
  continue;
destructureSlot(slot, allocator, rewriter, *info, statistics);
destructuredAny = true;

mlir/test/Dialect/LLVMIR/sroa.mlir

184

nit: I would stick to no_dynamic_indexing. The double negation is a bit hard to parse.

[Moxinilian]

Address comments.

mlir/include/mlir/Transforms/Passes.td

256

I tried! Unfortunately it does not escape carriage returns and breaks the generated code entirely...

[Moxinilian]

Fix tests + improve code layout.

[Moxinilian]

Rebase.

[Moxinilian]

Fix forgotten s. Sorry!

[Moxinilian]

Fix GEPIndidicesAdaptor bug in release mode.

[gysit]

LGTM!

Do you want to try another rebase. Maybe the CI then turns green assuming the underlying problem has been fixed in meantime?

[Moxinilian]

Refresh commit for CI.

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

 def LLVM_FSubOp : LLVM_FloatArithmeticOp<"fsub", "FSub">;
 def LLVM_FMulOp : LLVM_FloatArithmeticOp<"fmul", "FMul">;
 def LLVM_FDivOp : LLVM_FloatArithmeticOp<"fdiv", "FDiv">;
 def LLVM_FRemOp : LLVM_FloatArithmeticOp<"frem", "FRem">;
 def LLVM_FNegOp : LLVM_UnaryFloatArithmeticOp<
   LLVM_ScalarOrVectorOf<LLVM_AnyFloat>, "fneg", "FNeg">;
 
 // Memory-related operations.
-def LLVM_AllocaOp : LLVM_Op<"alloca",
-    [DeclareOpInterfaceMethods<PromotableAllocationOpInterface>]>,
+def LLVM_AllocaOp : LLVM_Op<"alloca", 
+    [DeclareOpInterfaceMethods<PromotableAllocationOpInterface>,
+     DeclareOpInterfaceMethods<DestructurableAllocationOpInterface>]>,
   LLVM_MemOpPatterns {
   let arguments = (ins AnyInteger:$arraySize,
                    OptionalAttr<I64Attr>:$alignment,
                    OptionalAttr<TypeAttr>:$elem_type,
                    UnitAttr:$inalloca);
   let results = (outs Res<LLVM_AnyPointer, "",
                           [MemAlloc<AutomaticAllocationScopeResource>]>:$res);
   string llvmInstName = "Alloca";
let builders = [
 
     }]>
     ];
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
 }
 
 def LLVM_GEPOp : LLVM_Op<"getelementptr", [Pure,
-    DeclareOpInterfaceMethods<PromotableOpInterface>]> {
+    DeclareOpInterfaceMethods<PromotableOpInterface>,
+    DeclareOpInterfaceMethods<SafeMemorySlotAccessOpInterface>,
+    DeclareOpInterfaceMethods<DestructurableAccessorOpInterface>]> {
   let arguments = (ins LLVM_ScalarOrVectorOf<LLVM_AnyPointer>:$base,
                    Variadic<LLVM_ScalarOrVectorOf<AnyInteger>>:$dynamicIndices,
                    DenseI32ArrayAttr:$rawConstantIndices,
                    OptionalAttr<TypeAttr>:$elem_type,
                    UnitAttr:$inbounds);
   let results = (outs LLVM_ScalarOrVectorOf<LLVM_AnyPointer>:$res);
   let skipDefaultBuilders = 1;
 
let extraClassDeclaration = [{
 
     GEPIndicesAdaptor<ValueRange> getIndices();
   }];
   let hasFolder = 1;
   let hasVerifier = 1;
 }
 
 def LLVM_LoadOp : LLVM_MemAccessOpBase<"load",
-    [DeclareOpInterfaceMethods<PromotableMemOpInterface>]> {
+    [DeclareOpInterfaceMethods<PromotableMemOpInterface>,
+     DeclareOpInterfaceMethods<SafeMemorySlotAccessOpInterface>]> {
   dag args = (ins Arg<LLVM_PointerTo<LLVM_LoadableType>, "", [MemRead]>:$addr,
               OptionalAttr<I64Attr>:$alignment,
               UnitAttr:$volatile_,
               UnitAttr:$nontemporal,
               DefaultValuedAttr<
                 AtomicOrdering, "AtomicOrdering::not_atomic">:$ordering,
               OptionalAttr<StrAttr>:$syncscope);
   // Append the aliasing related attributes define in LLVM_MemAccessOpBase.
OpBuilder<(ins "Type":$type, "Value":$addr,
       CArg<"bool", "false">:$isNonTemporal,
       CArg<"AtomicOrdering", "AtomicOrdering::not_atomic">:$ordering,
       CArg<"StringRef", "StringRef()">:$syncscope)>
   ];
   let hasVerifier = 1;
 }
 
 def LLVM_StoreOp : LLVM_MemAccessOpBase<"store",
-    [DeclareOpInterfaceMethods<PromotableMemOpInterface>]> {
+    [DeclareOpInterfaceMethods<PromotableMemOpInterface>,
+     DeclareOpInterfaceMethods<SafeMemorySlotAccessOpInterface>]> {
   dag args = (ins LLVM_LoadableType:$value,
               Arg<LLVM_PointerTo<LLVM_LoadableType>,"",[MemWrite]>:$addr,
               OptionalAttr<I64Attr>:$alignment,
               UnitAttr:$volatile_,
               UnitAttr:$nontemporal,
               DefaultValuedAttr<
                 AtomicOrdering, "AtomicOrdering::not_atomic">:$ordering,
               OptionalAttr<StrAttr>:$syncscope);

mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_LLVMIR_LLVMTYPES_H_
 #define MLIR_DIALECT_LLVMIR_LLVMTYPES_H_
 
 #include "mlir/IR/Types.h"
 #include "mlir/Interfaces/DataLayoutInterfaces.h"
+#include "mlir/Interfaces/MemorySlotInterfaces.h"
 #include <optional>
 
 namespace llvm {
 class ElementCount;
 class TypeSize;
 } // namespace llvm
 
 namespace mlir {

 /// body, one is not allowed to set another body (however, one can set exactly
 /// the same body).
 ///
 /// Note that the packedness of the struct takes place in uniquing of literal
 /// structs, but does not in uniquing of identified structs.
 class LLVMStructType
     : public Type::TypeBase<LLVMStructType, Type, detail::LLVMStructTypeStorage,
                             DataLayoutTypeInterface::Trait,
+                            DestructurableTypeInterface::Trait,
                             TypeTrait::IsMutable> {
 public:
   /// Inherit base constructors.
   using Base::Base;
 
   /// Checks if the given type can be contained in a structure type.
   static bool isValidElementType(Type type);
 
public:
   unsigned getPreferredAlignment(const DataLayout &dataLayout,
                                  DataLayoutEntryListRef params) const;
 
   bool areCompatible(DataLayoutEntryListRef oldLayout,
                      DataLayoutEntryListRef newLayout) const;
 
   LogicalResult verifyEntries(DataLayoutEntryListRef entries,
                               Location loc) const;
+
+  /// Destructs the struct into its indexed field types.
+  std::optional<DenseMap<Attribute, Type>> getSubelementIndexMap();
+
+  /// Returns which type is stored at a given integer index within the struct.
+  Type getTypeAtIndex(Attribute index);
 };
 
 //===----------------------------------------------------------------------===//
 // Printing and parsing.
 //===----------------------------------------------------------------------===//
 
 namespace detail {
 /// Parses an LLVM dialect type.

mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.td

 //===-- LLVMOps.td - LLVM IR dialect op definition file ----*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVMTYPES_TD
 #define LLVMTYPES_TD
 
 include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
 include "mlir/IR/AttrTypeBase.td"
 include "mlir/Interfaces/DataLayoutInterfaces.td"
+include "mlir/Interfaces/MemorySlotInterfaces.td"
 
 /// Base class for all LLVM dialect types.
 class LLVMType<string typeName, string typeMnemonic, list<Trait> traits = []>
     : TypeDef<LLVM_Dialect, typeName, traits> {
   let mnemonic = typeMnemonic;
 }
 
 //===----------------------------------------------------------------------===//
 // LLVMArrayType
 //===----------------------------------------------------------------------===//
 
 def LLVMArrayType : LLVMType<"LLVMArray", "array", [
-    DeclareTypeInterfaceMethods<DataLayoutTypeInterface, ["getTypeSize"]>]> {
+    DeclareTypeInterfaceMethods<DataLayoutTypeInterface, ["getTypeSize"]>,
+    DeclareTypeInterfaceMethods<DestructurableTypeInterface>]> {
   let summary = "LLVM array type";
   let description = [{
     The `!llvm.array` type represents a fixed-size array of element types.
     It is an aggregate type representing consecutive elements in memory,
     parameterized by the number of elements and the element type.
 
     Example:
 

mlir/include/mlir/Interfaces/CMakeLists.txt

 add_mlir_interface(TilingInterface)
 add_mlir_interface(ValueBoundsOpInterface)
 add_mlir_interface(VectorInterfaces)
 add_mlir_interface(ViewLikeInterface)
 
 set(LLVM_TARGET_DEFINITIONS MemorySlotInterfaces.td)
 mlir_tablegen(MemorySlotOpInterfaces.h.inc -gen-op-interface-decls)
 mlir_tablegen(MemorySlotOpInterfaces.cpp.inc -gen-op-interface-defs)
+mlir_tablegen(MemorySlotTypeInterfaces.h.inc -gen-type-interface-decls)
+mlir_tablegen(MemorySlotTypeInterfaces.cpp.inc -gen-type-interface-defs)
 add_public_tablegen_target(MLIRMemorySlotInterfacesIncGen)
 add_dependencies(mlir-generic-headers MLIRMemorySlotInterfacesIncGen)
 
 set(LLVM_TARGET_DEFINITIONS DataLayoutInterfaces.td)
 mlir_tablegen(DataLayoutAttrInterface.h.inc -gen-attr-interface-decls)
 mlir_tablegen(DataLayoutAttrInterface.cpp.inc -gen-attr-interface-defs)
 mlir_tablegen(DataLayoutOpInterface.h.inc -gen-op-interface-decls)
 mlir_tablegen(DataLayoutOpInterface.cpp.inc -gen-op-interface-defs)

mlir/include/mlir/Interfaces/MemorySlotInterfaces.h

 /// (for example alloca).
 struct MemorySlot {
   /// Pointer to the memory slot, used by operations to refer to it.
   Value ptr;
   /// Type of the value contained in the slot.
   Type elemType;
 };
 
+/// Memory slot attached with information about its destructuring procedure.
+struct DestructurableMemorySlot : public MemorySlot {
+  /// Maps an index within the memory slot to the type of the pointer that
+  /// will be generated to access the element directly.
+  DenseMap<Attribute, Type> elementPtrs;
+};
+
 /// Returned by operation promotion logic requesting the deletion of an
 /// operation.
 enum class DeletionKind {
   /// Keep the operation after promotion.
   Keep,
   /// Delete the operation after promotion.
   Delete,
 };
 
 } // namespace mlir
 
 #include "mlir/Interfaces/MemorySlotOpInterfaces.h.inc"
+#include "mlir/Interfaces/MemorySlotTypeInterfaces.h.inc"
 
 #endif // MLIR_INTERFACES_MEMORYSLOTINTERFACES_H

mlir/include/mlir/Interfaces/MemorySlotInterfaces.td

InterfaceMethod<[{
       "::mlir::DeletionKind",
       "removeBlockingUses",
       (ins "const ::llvm::SmallPtrSetImpl<mlir::OpOperand *> &":$blockingUses,
            "::mlir::RewriterBase &":$rewriter)
     >,
   ];
 }
 
+def DestructurableAllocationOpInterface
+  : OpInterface<"DestructurableAllocationOpInterface"> {
+  let description = [{
+    Describes operations allocating memory slots of aggregates that can be
+    destructured into multiple smaller allocations.
+  }];
+  let cppNamespace = "::mlir";
+
+  let methods = [
+    InterfaceMethod<[{
+        Returns the list of slots for which destructuring should be attempted,
+        specifying in which way the slot should be destructured into subslots.
+        The subslots are indexed by attributes. This computes the type of the
+        pointer for each subslot to be generated. The type of the memory slot
+        must implement `DestructurableTypeInterface`.
+
+        No IR mutation is allowed in this method.
+      }],
+      "::llvm::SmallVector<::mlir::DestructurableMemorySlot>",
+      "getDestructurableSlots",
+      (ins)
+    >,
+    InterfaceMethod<[{
+        Destructures this slot into multiple subslots. The newly generated slots
+        may belong to a different allocator. The original slot must still exist
+        at the end of this call. Only generates subslots for the indices found in
+        `usedIndices` since all other subslots are unused.
+
+        The rewriter is located at the beginning of the block where the slot
+        pointer is defined. All IR mutations must happen through the rewriter.
+      }],
+      "::llvm::DenseMap<::mlir::Attribute, ::mlir::MemorySlot>",
+      "destructure",
+      (ins "const ::mlir::DestructurableMemorySlot &":$slot,
+           "const ::llvm::SmallPtrSetImpl<::mlir::Attribute> &":$usedIndices,
+           "::mlir::RewriterBase &":$rewriter)
+    >,
+    InterfaceMethod<[{
+        Hook triggered once the destructuring of a slot is complete, meaning the
+        original slot is no longer being refered to and could be deleted.
+        This will only be called for slots declared by this operation.
+
+        All IR mutations must happen through the rewriter.
+      }],
+      "void", "handleDestructuringComplete",
+      (ins "const ::mlir::DestructurableMemorySlot &":$slot,
+           "::mlir::RewriterBase &":$rewriter)
+    >,
+  ];
+}
+
+def SafeMemorySlotAccessOpInterface
+  : OpInterface<"SafeMemorySlotAccessOpInterface"> {
+  let description = [{
+    Describes operations using memory slots in a type-safe manner.
+  }];
+  let cppNamespace = "::mlir";
+
+  let methods = [
+    InterfaceMethod<[{
+        Returns whether all accesses in this operation to the provided slot are
+        done in a type-safe manner. To be type-safe, the access must only load
+        the value in this type as the type of the slot, and without assuming any
+        context around the slot. For example, a type-safe load must not load
+        outside the bounds of the slot.
+
+        If the type-safety of the accesses depends on the type-safety of the
+        accesses to further memory slots, the result of this method will be
+        conditioned to the type-safety of the accesses to the slots added by
+        this method to `mustBeSafelyUsed`.
+
+        No IR mutation is allowed in this method.
+      }],
+      "::mlir::LogicalResult",
+      "ensureOnlySafeAccesses",
+      (ins "const ::mlir::MemorySlot &":$slot,
+           "::mlir::SmallVectorImpl<::mlir::MemorySlot> &":$mustBeSafelyUsed)
+    >
+  ];
+}
+
+def DestructurableAccessorOpInterface
+  : OpInterface<"DestructurableAccessorOpInterface"> {
+  let description = [{
+    Describes operations that can access a sub-element of a destructurable slot.
+  }];
+  let cppNamespace = "::mlir";
+
+  let methods = [
+    InterfaceMethod<[{
+        For a given destructurable memory slot, returns whether this operation can
+        rewire its uses of the slot to use the slots generated after
+        destructuring. This may involve creating new operations, and usually
+        amounts to checking if the pointer types match.
+
+        This method must also register the indices it will access within the
+        `usedIndices` set. If the accessor generates new slots mapping to
+        subelements, they must be registered in `mustBeSafelyUsed` to ensure
+        they are used in a locally type-safe manner.
+
+        No IR mutation is allowed in this method.
+      }],
+      "bool",
+      "canRewire",
+      (ins "const ::mlir::DestructurableMemorySlot &":$slot,
+           "::llvm::SmallPtrSetImpl<::mlir::Attribute> &":$usedIndices,
+           "::mlir::SmallVectorImpl<::mlir::MemorySlot> &":$mustBeSafelyUsed)
+    >,
+    InterfaceMethod<[{
+        Rewires the use of a slot to the generated subslots, without deleting
+        any operation. Returns whether the accessor should be deleted.
+
+        All IR mutations must happen through the rewriter. Deletion of
+        operations is not allowed, only the accessor can be scheduled for
+        deletion by returning the appropriate value.
+      }],
+      "::mlir::DeletionKind",
+      "rewire",
+      (ins "const ::mlir::DestructurableMemorySlot &":$slot,
+           "::llvm::DenseMap<::mlir::Attribute, ::mlir::MemorySlot> &":$subslots,
+           "::mlir::RewriterBase &":$rewriter)
+    >
+  ];
+}
+
+def DestructurableTypeInterface
+  : TypeInterface<"DestructurableTypeInterface"> {
+  let description = [{
+    Describes a type that can be broken down into indexable sub-element types.
+  }];
+  let cppNamespace = "::mlir";
+
+  let methods = [
+    InterfaceMethod<[{
+        Destructures the type into subelements into a map of index attributes to
+        types of subelements. Returns nothing if the type cannot be destructured.
+      }],
+      "::std::optional<::llvm::DenseMap<::mlir::Attribute, ::mlir::Type>>",
+      "getSubelementIndexMap",
+      (ins)
+    >,
+    InterfaceMethod<[{
+        Indicates which type is held at the provided index, returning a null
+        Type if no type could be computed. While this can return information
+        even when the type cannot be completely destructured, it must be coherent
+        with the types returned by `getSubelementIndexMap` when they exist.
+      }],
+      "::mlir::Type",
+      "getTypeAtIndex",
+      (ins "::mlir::Attribute":$index)
+    >
+  ];
+}
+
 #endif // MLIR_INTERFACES_MEMORYSLOTINTERFACES

mlir/include/mlir/Transforms/Mem2Reg.h

 #include "mlir/IR/Dominance.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/MemorySlotInterfaces.h"
 #include "llvm/ADT/Statistic.h"
 
 namespace mlir {
 
+/// Statistics collected while applying mem2reg.
 struct Mem2RegStatistics {
+  /// Total amount of memory slot promoted.
   llvm::Statistic *promotedAmount = nullptr;
+  /// Total amount of new block argument inserted in blocks.
   llvm::Statistic *newBlockArgumentAmount = nullptr;
 };
 
 /// Pattern applying mem2reg to the regions of the operations on which it
 /// matches.
 class Mem2RegPattern
     : public OpInterfaceRewritePattern<PromotableAllocationOpInterface> {
 public:
   using OpInterfaceRewritePattern::OpInterfaceRewritePattern;
 
   Mem2RegPattern(MLIRContext *context, Mem2RegStatistics statistics = {},
                  PatternBenefit benefit = 1)
       : OpInterfaceRewritePattern(context, benefit), statistics(statistics) {}
 
   LogicalResult matchAndRewrite(PromotableAllocationOpInterface allocator,
                                 PatternRewriter &rewriter) const override;
 
 private:
   Mem2RegStatistics statistics;
 };
 
-/// Pattern applying mem2reg to the regions of the operations on which it
-/// matches.
-class Mem2RegPattern : public RewritePattern {
-public:
-  using RewritePattern::RewritePattern;
-
-  Mem2RegPattern(MLIRContext *ctx, PatternBenefit benefit = 1)
-      : RewritePattern(MatchAnyOpTypeTag(), benefit, ctx) {}
-
-  LogicalResult matchAndRewrite(Operation *op,
-                                PatternRewriter &rewriter) const override;
-};
-
 /// Attempts to promote the memory slots of the provided allocators. Succeeds if
 /// at least one memory slot was promoted.
 LogicalResult
 tryToPromoteMemorySlots(ArrayRef<PromotableAllocationOpInterface> allocators,
                         RewriterBase &rewriter,
                         Mem2RegStatistics statistics = {});
 
 } // namespace mlir
 
 #endif // MLIR_TRANSFORMS_MEM2REG_H

mlir/include/mlir/Transforms/Passes.h

 #define GEN_PASS_DECL_CANONICALIZER
 #define GEN_PASS_DECL_CONTROLFLOWSINK
 #define GEN_PASS_DECL_CSEPASS
 #define GEN_PASS_DECL_INLINER
 #define GEN_PASS_DECL_LOOPINVARIANTCODEMOTION
 #define GEN_PASS_DECL_MEM2REG
 #define GEN_PASS_DECL_PRINTIRPASS
 #define GEN_PASS_DECL_PRINTOPSTATS
+#define GEN_PASS_DECL_SROA
 #define GEN_PASS_DECL_STRIPDEBUGINFO
 #define GEN_PASS_DECL_SCCP
 #define GEN_PASS_DECL_SYMBOLDCE
 #define GEN_PASS_DECL_SYMBOLPRIVATIZE
 #define GEN_PASS_DECL_TOPOLOGICALSORT
 #include "mlir/Transforms/Passes.h.inc"
 
 /// Creates an instance of the Canonicalizer pass, configured with default

mlir/include/mlir/Transforms/Passes.td

let options = [
     Option<"enableRegionSimplification", "region-simplify", "bool",
        /*default=*/"true",
        "Perform control flow optimizations to the region tree">,
   ];
 
   let statistics = [
     Statistic<"promotedAmount",
               "promoted slots",
-              "Number of promoted memory slot">,
+              "Total amount of memory slot promoted">,
     Statistic<"newBlockArgumentAmount",
               "new block args",
-              "Total number of block arguments added">,
+              "Total amount of new block argument inserted in blocks">,
   ];
 }
 
 def PrintOpStats : Pass<"print-op-stats"> {
   let summary = "Print statistics of operations";
   let constructor = "mlir::createPrintOpStatsPass()";
   let options = [
     Option<"printAsJSON", "json", "bool", /*default=*/"false",
let description = [{
     constant are replaced, and removed if possible.
 
     This implementation is based on the algorithm described by Wegman and Zadeck
     in [“Constant Propagation with Conditional Branches”](https://dl.acm.org/doi/10.1145/103135.103136) (1991).
   }];
   let constructor = "mlir::createSCCPPass()";
 }
 
+def SROA : Pass<"sroa"> {
+  let summary = "Scalar Replacement of Aggregates";
+  let description = [{
+    Scalar Replacement of Aggregates. Replaces allocations of aggregates into
+    independant allocations of its elements.
+
+    Allocators must implement `DestructurableAllocationOpInterface` to provide
+    the list of memory slots for which destructuring should be attempted.
+
+    This pass will only be applied if all accessors of the aggregate implement
+    the `DestructurableAccessorOpInterface`. If the accessors provide a view
+    into the struct, users of the view must ensure it is used in a type-safe
+    manner and within bounds by implementing `TypeSafeOpInterface`.
+  }];
+
+  let statistics = [
+    Statistic<
+      "destructuredAmount",
+      "destructured slots",
+      "Total amount of memory slots destructured"
+    >,
+    Statistic<
+      "slotsWithMemoryBenefit",
+      "slots with memory benefit",
+      "Total amount of memory slots in which the destructured size was smaller "
+      "than the total size after eliminating unused fields"
+    >,
+    Statistic<
+      "maxSubelementAmount",
+      "max subelement number",
+      "Maximal number of sub-elements a successfully destructured slot "
+      "initially had"
+    >,
+  ];
+}
+
 def StripDebugInfo : Pass<"strip-debuginfo"> {
   let summary = "Strip debug info from all operations";
   let description = [{
     This pass strips the IR of any location information, by replacing all
     operation locations with [`unknown`](Dialects/Builtin.md/#unknownloc).
   }];
   let constructor = "mlir::createStripDebugInfoPass()";
 }

mlir/include/mlir/Transforms/SROA.h

+//===-- SROA.h - Scalar Replacement Of Aggregates ---------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TRANSFORMS_SROA_H
+#define MLIR_TRANSFORMS_SROA_H
+
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Interfaces/MemorySlotInterfaces.h"
+#include "mlir/Support/LogicalResult.h"
+#include "llvm/ADT/Statistic.h"
+
+namespace mlir {
+
+/// Statistics collected while applying SROA.
+struct SROAStatistics {
+  /// Total amount of memory slots destructured.
+  llvm::Statistic *destructuredAmount = nullptr;
+  /// Total amount of memory slots in which the destructured size was smaller
+  /// than the total size after eliminating unused fields.
+  llvm::Statistic *slotsWithMemoryBenefit = nullptr;
+  /// Maximal number of sub-elements a successfully destructured slot initially
+  /// had.
+  llvm::Statistic *maxSubelementAmount = nullptr;
+};
+
+/// Pattern applying SROA to the regions of the operations on which it
+/// matches.
+class SROAPattern
+    : public OpInterfaceRewritePattern<DestructurableAllocationOpInterface> {
+public:
+  using OpInterfaceRewritePattern::OpInterfaceRewritePattern;
+
+  SROAPattern(MLIRContext *context, SROAStatistics statistics = {},
+              PatternBenefit benefit = 1)
+      : OpInterfaceRewritePattern(context, benefit), statistics(statistics) {}
+
+  LogicalResult matchAndRewrite(DestructurableAllocationOpInterface allocator,
+                                PatternRewriter &rewriter) const override;
+
+private:
+  SROAStatistics statistics;
+};
+
+/// Attempts to destructure the slots of destructurable allocators. Returns
+/// failure if no slot was destructured.
+LogicalResult tryToDestructureMemorySlots(
+    ArrayRef<DestructurableAllocationOpInterface> allocators,
+    RewriterBase &rewriter, SROAStatistics statistics = {});
+
+} // namespace mlir
+
+#endif // MLIR_TRANSFORMS_SROA_H

mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp

 //===- LLVMMemorySlot.cpp - MemorySlot interfaces ---------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements MemorySlot-related interfaces for LLVM dialect
 // operations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/Matchers.h"
+#include "mlir/IR/ValueRange.h"
 #include "mlir/Interfaces/MemorySlotInterfaces.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // Interfaces for AllocaOp
 //===----------------------------------------------------------------------===//
 

 void LLVM::AllocaOp::handlePromotionComplete(const MemorySlot &slot,
                                              Value defaultValue,
                                              RewriterBase &rewriter) {
   if (defaultValue && defaultValue.use_empty())
     rewriter.eraseOp(defaultValue.getDefiningOp());
   rewriter.eraseOp(*this);
 }
 
+SmallVector<DestructurableMemorySlot> LLVM::AllocaOp::getDestructurableSlots() {
+  if (!mlir::matchPattern(getArraySize(), m_One()))
+    return {};
+
+  Type elemType = getAllocaElementType(*this);
+  auto destructurable = dyn_cast<DestructurableTypeInterface>(elemType);
+  if (!destructurable)
+    return {};
+
+  std::optional<DenseMap<Attribute, Type>> destructuredType =
+      destructurable.getSubelementIndexMap();
+  if (!destructuredType)
+    return {};
+
+  DenseMap<Attribute, Type> allocaTypeMap;
+  for (Attribute index : llvm::make_first_range(destructuredType.value()))
+    allocaTypeMap.insert({index, LLVM::LLVMPointerType::get(getContext())});
+
+  return {DestructurableMemorySlot{{getResult(), elemType}, {allocaTypeMap}}};
+}
+
+DenseMap<Attribute, MemorySlot>
+LLVM::AllocaOp::destructure(const DestructurableMemorySlot &slot,
+                            const SmallPtrSetImpl<Attribute> &usedIndices,
+                            RewriterBase &rewriter) {
+  assert(slot.ptr == getResult());
+  Type elemType =
+      getElemType() ? *getElemType() : getResult().getType().getElementType();
+
+  rewriter.setInsertionPointAfter(*this);
+
+  auto destructurableType = cast<DestructurableTypeInterface>(elemType);
+  DenseMap<Attribute, MemorySlot> slotMap;
+  for (Attribute index : usedIndices) {
+    Type elemType = destructurableType.getTypeAtIndex(index);
+    assert(elemType && "used index must exist");
+    auto subAlloca = rewriter.create<LLVM::AllocaOp>(
+        getLoc(), LLVM::LLVMPointerType::get(getContext()), elemType,
+        getArraySize());
+    slotMap.try_emplace<MemorySlot>(index, {subAlloca.getResult(), elemType});
+  }
+
+  return slotMap;
+}
+
+void LLVM::AllocaOp::handleDestructuringComplete(
+    const DestructurableMemorySlot &slot, RewriterBase &rewriter) {
+  assert(slot.ptr == getResult());
+  rewriter.eraseOp(*this);
+}
+
 //===----------------------------------------------------------------------===//
 // Interfaces for LoadOp/StoreOp
 //===----------------------------------------------------------------------===//
 
 bool LLVM::LoadOp::loadsFrom(const MemorySlot &slot) {
   return getAddr() == slot.ptr;
 }
 

 }
 
 DeletionKind LLVM::StoreOp::removeBlockingUses(
     const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
     RewriterBase &rewriter, Value reachingDefinition) {
   // `canUsesBeRemoved` checked this blocking use must be the stored slot
   // pointer.
   for (Operation *user : slot.ptr.getUsers())
-    if (auto declareOp = llvm::dyn_cast<LLVM::DbgDeclareOp>(user))
+    if (auto declareOp = dyn_cast<LLVM::DbgDeclareOp>(user))
       rewriter.create<LLVM::DbgValueOp>(declareOp->getLoc(), getValue(),
                                         declareOp.getVarInfo());
   return DeletionKind::Delete;
 }
 
+LogicalResult LLVM::LoadOp::ensureOnlySafeAccesses(
+    const MemorySlot &slot, SmallVectorImpl<MemorySlot> &mustBeSafelyUsed) {
+  return success(getAddr() != slot.ptr || getType() == slot.elemType);
+}
+
+LogicalResult LLVM::StoreOp::ensureOnlySafeAccesses(
+    const MemorySlot &slot, SmallVectorImpl<MemorySlot> &mustBeSafelyUsed) {
+  return success(getAddr() != slot.ptr ||
+                 getValue().getType() == slot.elemType);
+}
+
 //===----------------------------------------------------------------------===//
 // Interfaces for discardable OPs
 //===----------------------------------------------------------------------===//
 
 /// Conditions the deletion of the operation to the removal of all its uses.
 static bool forwardToUsers(Operation *op,
                            SmallVectorImpl<OpOperand *> &newBlockingUses) {
   for (Value result : op->getResults())
bool LLVM::DbgDeclareOp::canUsesBeRemoved(
   return true;
 }
 
 DeletionKind LLVM::DbgDeclareOp::removeBlockingUses(
     const SmallPtrSetImpl<OpOperand *> &blockingUses, RewriterBase &rewriter) {
   return DeletionKind::Delete;
 }
 
+//===----------------------------------------------------------------------===//
+// Interfaces for GEPOp
+//===----------------------------------------------------------------------===//
+
 static bool hasAllZeroIndices(LLVM::GEPOp gepOp) {
   return llvm::all_of(gepOp.getIndices(), [](auto index) {
     auto indexAttr = index.template dyn_cast<IntegerAttr>();
     return indexAttr && indexAttr.getValue() == 0;
   });
 }
 
-//===----------------------------------------------------------------------===//
-// Interfaces for GEPOp
-//===----------------------------------------------------------------------===//
-
 bool LLVM::GEPOp::canUsesBeRemoved(
     const SmallPtrSetImpl<OpOperand *> &blockingUses,
     SmallVectorImpl<OpOperand *> &newBlockingUses) {
   // GEP can be removed as long as it is a no-op and its users can be removed.
   if (!hasAllZeroIndices(*this))
     return false;
   return forwardToUsers(*this, newBlockingUses);
 }
 
 DeletionKind LLVM::GEPOp::removeBlockingUses(
     const SmallPtrSetImpl<OpOperand *> &blockingUses, RewriterBase &rewriter) {
   return DeletionKind::Delete;
 }
 
+/// Returns the type the resulting pointer of the GEP points to. If such a type
+/// is not clear, returns null type.
+static Type computeReachedGEPType(LLVM::GEPOp gep) {
+  if (gep.getIndices().empty())
+    return {};
+
+  // Check the pointer indexing only targets the first element.
+  auto firstIndex = gep.getIndices()[0];
+  IntegerAttr indexInt = firstIndex.dyn_cast<IntegerAttr>();
+  if (!indexInt || indexInt.getInt() != 0)
+    return {};
+
+  // Set the initial type currently being used for indexing. This will be
+  // updated as the indices get walked over.
+  std::optional<Type> maybeSelectedType = gep.getElemType();
+  if (!maybeSelectedType)
+    return {};
+  Type selectedType = *maybeSelectedType;
+
+  // Follow the indexed elements in the gep.
+  for (const LLVM::GEPIndicesAdaptor<ValueRange>::value_type &index :
+       llvm::drop_begin(gep.getIndices())) {
+    // Ensure the index is static and obtain it.
+    IntegerAttr indexInt = index.dyn_cast<IntegerAttr>();
+    if (!indexInt)
+      return {};
+
+    // Ensure the structure of the type being indexed can be reasoned about.
+    // This includes rejecting any potential typed pointer.
+    auto destructurable = selectedType.dyn_cast<DestructurableTypeInterface>();
+    if (!destructurable)
+      return {};
+
+    // Follow the type at the index the gep is accessing, making it the new type
+    // used for indexing.
+    Type field = destructurable.getTypeAtIndex(indexInt);
+    if (!field)
+      return {};
+    selectedType = field;
+  }
+
+  // When there are no more indices, the type currently being used for indexing
+  // is the type of the value pointed at by the returned indexed pointer.
+  return selectedType;
+}
+
+LogicalResult LLVM::GEPOp::ensureOnlySafeAccesses(
+    const MemorySlot &slot, SmallVectorImpl<MemorySlot> &mustBeSafelyUsed) {
+  if (getBase() != slot.ptr)
+    return success();
+  if (slot.elemType != getElemType())
+    return failure();
+  Type reachedType = computeReachedGEPType(*this);
+  if (!reachedType)
+    return failure();
+  mustBeSafelyUsed.emplace_back<MemorySlot>({getResult(), reachedType});
+  return success();
+}
+
+bool LLVM::GEPOp::canRewire(const DestructurableMemorySlot &slot,
+                            SmallPtrSetImpl<Attribute> &usedIndices,
+                            SmallVectorImpl<MemorySlot> &mustBeSafelyUsed) {
+  auto basePtrType = getBase().getType().dyn_cast<LLVM::LLVMPointerType>();
+  if (!basePtrType)
+    return false;
+
+  // Typed pointers are not supported. This should be removed once typed
+  // pointers are removed from the LLVM dialect.
+  if (!basePtrType.isOpaque())
+    return false;
+
+  if (getBase() != slot.ptr || slot.elemType != getElemType())
+    return false;
+  Type reachedType = computeReachedGEPType(*this);
+  if (!reachedType || getIndices().size() < 2)
+    return false;
+  auto firstLevelIndex = cast<IntegerAttr>(getIndices()[1]);
+  assert(slot.elementPtrs.contains(firstLevelIndex));
+  if (!slot.elementPtrs.at(firstLevelIndex).isa<LLVM::LLVMPointerType>())
+    return false;
+  mustBeSafelyUsed.emplace_back<MemorySlot>({getResult(), reachedType});
+  usedIndices.insert(firstLevelIndex);
+  return true;
+}
+
+DeletionKind LLVM::GEPOp::rewire(const DestructurableMemorySlot &slot,
+                                 DenseMap<Attribute, MemorySlot> &subslots,
+                                 RewriterBase &rewriter) {
+  IntegerAttr firstLevelIndex = getIndices()[1].dyn_cast<IntegerAttr>();
+  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.updateRootInPlace(*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;
+}
+
+//===----------------------------------------------------------------------===//
+// Interfaces for destructurable types
+//===----------------------------------------------------------------------===//
+
+std::optional<DenseMap<Attribute, Type>>
+LLVM::LLVMStructType::getSubelementIndexMap() {
+  Type i32 = IntegerType::get(getContext(), 32);
+  DenseMap<Attribute, Type> destructured;
+  for (const auto &[index, elemType] : llvm::enumerate(getBody()))
+    destructured.insert({IntegerAttr::get(i32, index), elemType});
+  return destructured;
+}
+
+Type LLVM::LLVMStructType::getTypeAtIndex(Attribute index) {
+  auto indexAttr = index.dyn_cast<IntegerAttr>();
+  if (!indexAttr || !indexAttr.getType().isInteger(32))
+    return {};
+  int32_t indexInt = indexAttr.getInt();
+  ArrayRef<Type> body = getBody();
+  if (indexInt < 0 || body.size() <= static_cast<uint32_t>(indexInt))
+    return {};
+  return body[indexInt];
+}
+
+std::optional<DenseMap<Attribute, Type>>
+LLVM::LLVMArrayType::getSubelementIndexMap() const {
+  constexpr size_t maxArraySizeForDestructuring = 16;
+  if (getNumElements() > maxArraySizeForDestructuring)
+    return {};
+  int32_t numElements = getNumElements();
+
+  Type i32 = IntegerType::get(getContext(), 32);
+  DenseMap<Attribute, Type> destructured;
+  for (int32_t index = 0; index < numElements; ++index)
+    destructured.insert({IntegerAttr::get(i32, index), getElementType()});
+  return destructured;
+}
+
+Type LLVM::LLVMArrayType::getTypeAtIndex(Attribute index) const {
+  auto indexAttr = index.dyn_cast<IntegerAttr>();
+  if (!indexAttr || !indexAttr.getType().isInteger(32))
+    return {};
+  int32_t indexInt = indexAttr.getInt();
+  if (indexInt < 0 || getNumElements() <= static_cast<uint32_t>(indexInt))
+    return {};
+  return getElementType();
+}

mlir/lib/Interfaces/MemorySlotInterfaces.cpp

 //===-- MemorySlotInterfaces.cpp - MemorySlot interfaces --------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Interfaces/MemorySlotInterfaces.h"
 
 #include "mlir/Interfaces/MemorySlotOpInterfaces.cpp.inc"
+#include "mlir/Interfaces/MemorySlotTypeInterfaces.cpp.inc"

mlir/lib/Transforms/CMakeLists.txt

 add_subdirectory(Utils)
 
 add_mlir_library(MLIRTransforms
   Canonicalizer.cpp
   ControlFlowSink.cpp
   CSE.cpp
   GenerateRuntimeVerification.cpp
   Inliner.cpp
   LocationSnapshot.cpp
   LoopInvariantCodeMotion.cpp
   Mem2Reg.cpp
   OpStats.cpp
   PrintIR.cpp
   SCCP.cpp
+  SROA.cpp
   StripDebugInfo.cpp
   SymbolDCE.cpp
   SymbolPrivatize.cpp
   TopologicalSort.cpp
   ViewOpGraph.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Transforms

mlir/lib/Transforms/Mem2Reg.cpp

if (statistics.promotedAmount)
     (*statistics.promotedAmount)++;
 
   allocator.handlePromotionComplete(slot, defaultValue, rewriter);
 }
 
 LogicalResult mlir::tryToPromoteMemorySlots(
     ArrayRef<PromotableAllocationOpInterface> allocators,
     RewriterBase &rewriter, Mem2RegStatistics statistics) {
-  DominanceInfo dominance;
-
   bool promotedAny = false;
 
   for (PromotableAllocationOpInterface allocator : allocators) {
     for (MemorySlot slot : allocator.getPromotableSlots()) {
       if (slot.ptr.use_empty())
         continue;
 
       DominanceInfo dominance;

 
 LogicalResult
 Mem2RegPattern::matchAndRewrite(PromotableAllocationOpInterface allocator,
                                 PatternRewriter &rewriter) const {
   hasBoundedRewriteRecursion();
   return tryToPromoteMemorySlots({allocator}, rewriter, statistics);
 }
 
-LogicalResult Mem2RegPattern::matchAndRewrite(Operation *op,
-                                              PatternRewriter &rewriter) const {
-  hasBoundedRewriteRecursion();
+namespace {
 
 struct Mem2Reg : impl::Mem2RegBase<Mem2Reg> {
   using impl::Mem2RegBase<Mem2Reg>::Mem2RegBase;
 
   void runOnOperation() override {
     Operation *scopeOp = getOperation();
 
     Mem2RegStatistics statictics{&promotedAmount, &newBlockArgumentAmount};

mlir/lib/Transforms/SROA.cpp

+//===-- SROA.cpp - Scalar Replacement Of Aggregates -------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Transforms/SROA.h"
+#include "mlir/Analysis/SliceAnalysis.h"
+#include "mlir/Interfaces/MemorySlotInterfaces.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "mlir/Transforms/Passes.h"
+
+namespace mlir {
+#define GEN_PASS_DEF_SROA
+#include "mlir/Transforms/Passes.h.inc"
+} // namespace mlir
+
+#define DEBUG_TYPE "sroa"
+
+using namespace mlir;
+
+namespace {
+
+/// Information computed by destructurable memory slot analysis used to perform
+/// actual destructuring of the slot. This struct is only constructed if
+/// destructuring is possible, and contains the necessary data to perform it.
+struct MemorySlotDestructuringInfo {
+  /// Set of the indices that are actually used when accessing the subelements.
+  SmallPtrSet<Attribute, 8> usedIndices;
+  /// Blocking uses of a given user of the memory slot that must be eliminated.
+  DenseMap<Operation *, SmallPtrSet<OpOperand *, 4>> userToBlockingUses;
+  /// List of potentially indirect accessors of the memory slot that need
+  /// rewiring.
+  SmallVector<DestructurableAccessorOpInterface> accessors;
+};
+
+} // namespace
+
+/// Computes information for slot destructuring. This will compute whether this
+/// slot can be destructured and data to perform the destructuring. Returns
+/// nothing if the slot cannot be destructured or if there is no useful work to
+/// be done.
+static std::optional<MemorySlotDestructuringInfo>
+computeDestructuringInfo(DestructurableMemorySlot &slot) {
+  assert(isa<DestructurableTypeInterface>(slot.elemType));
+
+  if (slot.ptr.use_empty())
+    return {};
+
+  MemorySlotDestructuringInfo info;
+
+  SmallVector<MemorySlot> usedSafelyWorklist;
+
+  auto scheduleAsBlockingUse = [&](OpOperand &use) {
+    SmallPtrSetImpl<OpOperand *> &blockingUses =
+        info.userToBlockingUses.getOrInsertDefault(use.getOwner());
+    blockingUses.insert(&use);
+  };
+
+  // Initialize the analysis with the immediate users of the slot.
+  for (OpOperand &use : slot.ptr.getUses()) {
+    if (auto accessor =
+            dyn_cast<DestructurableAccessorOpInterface>(use.getOwner())) {
+      if (accessor.canRewire(slot, info.usedIndices, usedSafelyWorklist)) {
+        info.accessors.push_back(accessor);
+        continue;
+      }
+    }
+
+    // If it cannot be shown that the operation uses the slot safely, maybe it
+    // can be promoted out of using the slot?
+    scheduleAsBlockingUse(use);
+  }
+
+  SmallPtrSet<OpOperand *, 16> visited;
+  while (!usedSafelyWorklist.empty()) {
+    MemorySlot mustBeUsedSafely = usedSafelyWorklist.pop_back_val();
+    for (OpOperand &subslotUse : mustBeUsedSafely.ptr.getUses()) {
+      if (!visited.insert(&subslotUse).second)
+        continue;
+      Operation *subslotUser = subslotUse.getOwner();
+
+      if (auto memOp = dyn_cast<SafeMemorySlotAccessOpInterface>(subslotUser))
+        if (succeeded(memOp.ensureOnlySafeAccesses(mustBeUsedSafely,
+                                                   usedSafelyWorklist)))
+          continue;
+
+      // If it cannot be shown that the operation uses the slot safely, maybe it
+      // can be promoted out of using the slot?
+      scheduleAsBlockingUse(subslotUse);
+    }
+  }
+
+  SetVector<Operation *> forwardSlice;
+  mlir::getForwardSlice(slot.ptr, &forwardSlice);
+  for (Operation *user : forwardSlice) {
+    // If the next operation has no blocking uses, everything is fine.
+    if (!info.userToBlockingUses.contains(user))
+      continue;
+
+    SmallPtrSet<OpOperand *, 4> &blockingUses = info.userToBlockingUses[user];
+    auto promotable = dyn_cast<PromotableOpInterface>(user);
+
+    // An operation that has blocking uses must be promoted. If it is not
+    // promotable, destructuring must fail.
+    if (!promotable)
+      return {};
+
+    SmallVector<OpOperand *> newBlockingUses;
+    // If the operation decides it cannot deal with removing the blocking uses,
+    // destructuring must fail.
+    if (!promotable.canUsesBeRemoved(blockingUses, newBlockingUses))
+      return {};
+
+    // Then, register any new blocking uses for coming operations.
+    for (OpOperand *blockingUse : newBlockingUses) {
+      assert(llvm::is_contained(user->getResults(), blockingUse->get()));
+
+      SmallPtrSetImpl<OpOperand *> &newUserBlockingUseSet =
+          info.userToBlockingUses.getOrInsertDefault(blockingUse->getOwner());
+      newUserBlockingUseSet.insert(blockingUse);
+    }
+  }
+
+  return info;
+}
+
+/// Performs the destructuring of a destructible slot given associated
+/// destructuring information. The provided slot will be destructured in
+/// subslots as specified by its allocator.
+static void destructureSlot(DestructurableMemorySlot &slot,
+                            DestructurableAllocationOpInterface allocator,
+                            RewriterBase &rewriter,
+                            MemorySlotDestructuringInfo &info,
+                            const SROAStatistics &statistics) {
+  RewriterBase::InsertionGuard guard(rewriter);
+
+  rewriter.setInsertionPointToStart(slot.ptr.getParentBlock());
+  DenseMap<Attribute, MemorySlot> subslots =
+      allocator.destructure(slot, info.usedIndices, rewriter);
+
+  if (statistics.slotsWithMemoryBenefit &&
+      slot.elementPtrs.size() != info.usedIndices.size())
+    (*statistics.slotsWithMemoryBenefit)++;
+
+  if (statistics.maxSubelementAmount)
+    statistics.maxSubelementAmount->updateMax(slot.elementPtrs.size());
+
+  SetVector<Operation *> usersToRewire;
+  for (Operation *user : llvm::make_first_range(info.userToBlockingUses))
+    usersToRewire.insert(user);
+  for (DestructurableAccessorOpInterface accessor : info.accessors)
+    usersToRewire.insert(accessor);
+  usersToRewire = mlir::topologicalSort(usersToRewire);
+
+  llvm::SmallVector<Operation *> toErase;
+  for (Operation *toRewire : llvm::reverse(usersToRewire)) {
+    rewriter.setInsertionPointAfter(toRewire);
+    if (auto accessor = dyn_cast<DestructurableAccessorOpInterface>(toRewire)) {
+      if (accessor.rewire(slot, subslots, rewriter) == DeletionKind::Delete)
+        toErase.push_back(accessor);
+      continue;
+    }
+
+    auto promotable = cast<PromotableOpInterface>(toRewire);
+    if (promotable.removeBlockingUses(info.userToBlockingUses[promotable],
+                                      rewriter) == DeletionKind::Delete)
+      toErase.push_back(promotable);
+  }
+
+  for (Operation *toEraseOp : toErase)
+    rewriter.eraseOp(toEraseOp);
+
+  assert(slot.ptr.use_empty() && "after destructuring, the original slot "
+                                 "pointer should no longer be used");
+
+  LLVM_DEBUG(llvm::dbgs() << "[sroa] Destructured memory slot: " << slot.ptr
+                          << "\n");
+
+  if (statistics.destructuredAmount)
+    (*statistics.destructuredAmount)++;
+
+  allocator.handleDestructuringComplete(slot, rewriter);
+}
+
+LogicalResult mlir::tryToDestructureMemorySlots(
+    ArrayRef<DestructurableAllocationOpInterface> allocators,
+    RewriterBase &rewriter, SROAStatistics statistics) {
+  bool destructuredAny = false;
+
+  for (DestructurableAllocationOpInterface allocator : allocators) {
+    for (DestructurableMemorySlot slot : allocator.getDestructurableSlots()) {
+      std::optional<MemorySlotDestructuringInfo> info =
+          computeDestructuringInfo(slot);
+      if (!info)
+        continue;
+
+      destructureSlot(slot, allocator, rewriter, *info, statistics);
+      destructuredAny = true;
+    }
+  }
+
+  return success(destructuredAny);
+}
+
+LogicalResult
+SROAPattern::matchAndRewrite(DestructurableAllocationOpInterface allocator,
+                             PatternRewriter &rewriter) const {
+  hasBoundedRewriteRecursion();
+  return tryToDestructureMemorySlots({allocator}, rewriter, statistics);
+}
+
+namespace {
+
+struct SROA : public impl::SROABase<SROA> {
+  using impl::SROABase<SROA>::SROABase;
+
+  void runOnOperation() override {
+    Operation *scopeOp = getOperation();
+
+    SROAStatistics statistics{&destructuredAmount, &slotsWithMemoryBenefit,
+                              &maxSubelementAmount};
+
+    RewritePatternSet rewritePatterns(&getContext());
+    rewritePatterns.add<SROAPattern>(&getContext(), statistics);
+    FrozenRewritePatternSet frozen(std::move(rewritePatterns));
+
+    if (failed(applyPatternsAndFoldGreedily(scopeOp, frozen)))
+      signalPassFailure();
+  }
+};
+
+} // namespace

mlir/test/Dialect/LLVMIR/sroa-statistics.mlir

+// RUN: mlir-opt %s --pass-pipeline="builtin.module(llvm.func(sroa))" --split-input-file --mlir-pass-statistics 2>&1 >/dev/null | FileCheck %s
+
+// CHECK: SROA
+// CHECK-NEXT: (S) 1 destructured slots
+// CHECK-NEXT: (S) 2 max subelement number
+// CHECK-NEXT: (S) 1 slots with memory benefit
+llvm.func @basic() -> i32 {
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", (i32, i32)> {alignment = 8 : i64} : (i32) -> !llvm.ptr
+  %2 = llvm.getelementptr inbounds %1[0, 1] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i32, i32)>
+  %3 = llvm.load %2 : !llvm.ptr -> i32
+  llvm.return %3 : i32
+}
+
+// -----
+
+// CHECK: SROA
+// CHECK-NEXT: (S) 1 destructured slots
+// CHECK-NEXT: (S) 2 max subelement number
+// CHECK-NEXT: (S) 0 slots with memory benefit
+llvm.func @basic_no_memory_benefit() -> i32 {
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", (i32, i32)> {alignment = 8 : i64} : (i32) -> !llvm.ptr
+  %2 = llvm.getelementptr inbounds %1[0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i32, i32)>
+  %3 = llvm.getelementptr inbounds %1[0, 1] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i32, i32)>
+  %4 = llvm.load %2 : !llvm.ptr -> i32
+  %5 = llvm.load %3 : !llvm.ptr -> i32
+  %6 = llvm.add %4, %5 : i32
+  llvm.return %6 : i32
+}
+
+// -----
+
+// CHECK: SROA
+// CHECK-NEXT: (S)  1 destructured slots
+// CHECK-NEXT: (S) 10 max subelement number
+// CHECK-NEXT: (S)  1 slots with memory benefit
+llvm.func @basic_array() -> i32 {
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  %1 = llvm.alloca %0 x !llvm.array<10 x i32> {alignment = 8 : i64} : (i32) -> !llvm.ptr
+  %2 = llvm.getelementptr inbounds %1[0, 2] : (!llvm.ptr) -> !llvm.ptr, !llvm.array<10 x i32>
+  %3 = llvm.load %2 : !llvm.ptr -> i32
+  llvm.return %3 : i32
+}
+
+// -----
+
+// SROA is applied repeatedly here, peeling off layers of aggregates one after
+// the other, four times.
+
+// CHECK: SROA
+// CHECK-NEXT: (S)  4 destructured slots
+// CHECK-NEXT: (S) 10 max subelement number
+// CHECK-NEXT: (S)  4 slots with memory benefit
+llvm.func @multi_level_direct() -> i32 {
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", (i32, f64, struct<"bar", (i8, array<10 x array<10 x i32>>, i8)>)> {alignment = 8 : i64} : (i32) -> !llvm.ptr
+  %2 = llvm.getelementptr inbounds %1[0, 2, 1, 5, 8] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i32, f64, struct<"bar", (i8, array<10 x array<10 x i32>>, i8)>)>
+  %3 = llvm.load %2 : !llvm.ptr -> i32
+  llvm.return %3 : i32
+}

mlir/test/Dialect/LLVMIR/sroa.mlir

+// RUN: mlir-opt %s --pass-pipeline="builtin.module(llvm.func(sroa))" --split-input-file | FileCheck %s
+
+// CHECK-LABEL: llvm.func @basic_struct
+llvm.func @basic_struct() -> i32 {
+  // CHECK: %[[SIZE:.*]] = llvm.mlir.constant(1 : i32)
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %[[SIZE]] x i32
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", (i32, f64, i32)> {alignment = 8 : i64} : (i32) -> !llvm.ptr
+  %2 = llvm.getelementptr inbounds %1[0, 2] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i32, f64, i32)>
+  // CHECK: %[[RES:.*]] = llvm.load %[[ALLOCA]]
+  %3 = llvm.load %2 : !llvm.ptr -> i32
+  // CHECK: llvm.return %[[RES]] : i32
+  llvm.return %3 : i32
+}
+
+// -----
+
+// CHECK-LABEL: llvm.func @basic_array
+llvm.func @basic_array() -> i32 {
+  // CHECK: %[[SIZE:.*]] = llvm.mlir.constant(1 : i32)
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %[[SIZE]] x i32
+  %1 = llvm.alloca %0 x !llvm.array<10 x i32> {alignment = 8 : i64} : (i32) -> !llvm.ptr
+  %2 = llvm.getelementptr inbounds %1[0, 2] : (!llvm.ptr) -> !llvm.ptr, !llvm.array<10 x i32>
+  // CHECK: %[[RES:.*]] = llvm.load %[[ALLOCA]]
+  %3 = llvm.load %2 : !llvm.ptr -> i32
+  // CHECK: llvm.return %[[RES]] : i32
+  llvm.return %3 : i32
+}
+
+// -----
+
+// CHECK-LABEL: llvm.func @multi_level_direct
+llvm.func @multi_level_direct() -> i32 {
+  // CHECK: %[[SIZE:.*]] = llvm.mlir.constant(1 : i32)
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %[[SIZE]] x i32
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", (i32, f64, struct<"bar", (i8, array<10 x array<10 x i32>>, i8)>)> {alignment = 8 : i64} : (i32) -> !llvm.ptr
+  %2 = llvm.getelementptr inbounds %1[0, 2, 1, 5, 8] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i32, f64, struct<"bar", (i8, array<10 x array<10 x i32>>, i8)>)>
+  // CHECK: %[[RES:.*]] = llvm.load %[[ALLOCA]]
+  %3 = llvm.load %2 : !llvm.ptr -> i32
+  // CHECK: llvm.return %[[RES]] : i32
+  llvm.return %3 : i32
+}
+
+// -----
+
+// The first application of SROA would generate a GEP with indices [0, 0]. This
+// test ensures this GEP is not eliminated during the first application. Even
+// though doing it would be correct, it would prevent the second application
+// of SROA to eliminate the array. GEPs should be eliminated only when they are
+// truly trivial (with indices [0]).
+
+// CHECK-LABEL: llvm.func @multi_level_direct_two_applications
+llvm.func @multi_level_direct_two_applications() -> i32 {
+  // CHECK: %[[SIZE:.*]] = llvm.mlir.constant(1 : i32)
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %[[SIZE]] x i32
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", (i32, f64, array<10 x i32>, i8)> {alignment = 8 : i64} : (i32) -> !llvm.ptr
+  %2 = llvm.getelementptr inbounds %1[0, 2, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i32, f64, array<10 x i32>, i8)>
+  // CHECK: %[[RES:.*]] = llvm.load %[[ALLOCA]]
+  %3 = llvm.load %2 : !llvm.ptr -> i32
+  // CHECK: llvm.return %[[RES]] : i32
+  llvm.return %3 : i32
+}
+
+// -----
+
+// CHECK-LABEL: llvm.func @multi_level_indirect
+llvm.func @multi_level_indirect() -> i32 {
+  // CHECK: %[[SIZE:.*]] = llvm.mlir.constant(1 : i32)
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %[[SIZE]] x i32
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", (i32, f64, struct<"bar", (i8, array<10 x array<10 x i32>>, i8)>)> {alignment = 8 : i64} : (i32) -> !llvm.ptr
+  %2 = llvm.getelementptr inbounds %1[0, 2, 1, 5] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i32, f64, struct<"bar", (i8, array<10 x array<10 x i32>>, i8)>)>
+  %3 = llvm.getelementptr inbounds %2[0, 8] : (!llvm.ptr) -> !llvm.ptr, !llvm.array<10 x i32>
+  // CHECK: %[[RES:.*]] = llvm.load %[[ALLOCA]]
+  %4 = llvm.load %3 : !llvm.ptr -> i32
+  // CHECK: llvm.return %[[RES]] : i32
+  llvm.return %4 : i32
+}
+
+// -----
+
+// CHECK-LABEL: llvm.func @resolve_alias
+// CHECK-SAME: (%[[ARG:.*]]: i32)
+llvm.func @resolve_alias(%arg: i32) -> i32 {
+  // CHECK: %[[SIZE:.*]] = llvm.mlir.constant(1 : i32)
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %[[SIZE]] x i32
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", (i32, f64, i32)> {alignment = 8 : i64} : (i32) -> !llvm.ptr
+  %2 = llvm.getelementptr %1[0, 2] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i32, f64, i32)>
+  %3 = llvm.getelementptr inbounds %1[0, 2] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i32, f64, i32)>
+  // CHECK: llvm.store %[[ARG]], %[[ALLOCA]]
+  llvm.store %arg, %2 : i32, !llvm.ptr
+  // CHECK: %[[RES:.*]] = llvm.load %[[ALLOCA]]
+  %4 = llvm.load %3 : !llvm.ptr -> i32
+  // CHECK: llvm.return %[[RES]] : i32
+  llvm.return %4 : i32
+}
+
+// -----
+
+// CHECK-LABEL: llvm.func @no_non_single_support
+llvm.func @no_non_single_support() -> i32 {
+  // CHECK: %[[SIZE:.*]] = llvm.mlir.constant
+  %0 = llvm.mlir.constant(2 : i32) : i32
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %[[SIZE]] x !llvm.struct<"foo", (i32, f64, i32)> {alignment = 8 : i64} : (i32) -> !llvm.ptr
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", (i32, f64, i32)> {alignment = 8 : i64} : (i32) -> !llvm.ptr
+  // CHECK-NOT: = llvm.alloca
+  %2 = llvm.getelementptr inbounds %1[0, 2] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i32, f64, i32)>
+  %3 = llvm.load %2 : !llvm.ptr -> i32
+  llvm.return %3 : i32
+}
+
+// -----
+
+// CHECK-LABEL: llvm.func @no_pointer_indexing
+llvm.func @no_pointer_indexing() -> i32 {
+  // CHECK: %[[SIZE:.*]] = llvm.mlir.constant(1 : i32)
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %[[SIZE]] x !llvm.struct<"foo", (i32, f64, i32)> {alignment = 8 : i64} : (i32) -> !llvm.ptr
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", (i32, f64, i32)> {alignment = 8 : i64} : (i32) -> !llvm.ptr
+  // CHECK-NOT: = llvm.alloca
+  %2 = llvm.getelementptr %1[1, 2] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i32, f64, i32)>
+  %3 = llvm.load %2 : !llvm.ptr -> i32
+  llvm.return %3 : i32
+}
+
+// -----
+
+// CHECK-LABEL: llvm.func @no_direct_use
+llvm.func @no_direct_use() -> i32 {
+  // CHECK: %[[SIZE:.*]] = llvm.mlir.constant(1 : i32)
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %[[SIZE]] x !llvm.struct<"foo", (i32, f64, i32)> {alignment = 8 : i64} : (i32) -> !llvm.ptr
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", (i32, f64, i32)> {alignment = 8 : i64} : (i32) -> !llvm.ptr
+  // CHECK-NOT: = llvm.alloca
+  %2 = llvm.getelementptr %1[0, 2] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i32, f64, i32)>
+  %3 = llvm.load %2 : !llvm.ptr -> i32
+  llvm.call @use(%1) : (!llvm.ptr) -> ()
+  llvm.return %3 : i32
+}
+
+llvm.func @use(!llvm.ptr)
+
+// -----
+
+// CHECK-LABEL: llvm.func @direct_promotable_use_is_fine
+llvm.func @direct_promotable_use_is_fine() -> i32 {
+  // CHECK: %[[SIZE:.*]] = llvm.mlir.constant(1 : i32)
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %[[SIZE]] x i32
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", (i32, f64, i32)> {alignment = 8 : i64} : (i32) -> !llvm.ptr
+  %2 = llvm.getelementptr %1[0, 2] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i32, f64, i32)>
+  // CHECK: %[[RES:.*]] = llvm.load %[[ALLOCA]]
+  %3 = llvm.load %2 : !llvm.ptr -> i32
+  // This is a direct use of the slot but it can be removed because it implements PromotableOpInterface.
+  llvm.intr.lifetime.start 2, %1 : !llvm.ptr
+  // CHECK: llvm.return %[[RES]] : i32
+  llvm.return %3 : i32
+}
+
+// -----
+
+// CHECK-LABEL: llvm.func @direct_promotable_use_is_fine_on_accessor
+llvm.func @direct_promotable_use_is_fine_on_accessor() -> i32 {
+  // CHECK: %[[SIZE:.*]] = llvm.mlir.constant(1 : i32)
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %[[SIZE]] x i32
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", (i32, f64, i32)> {alignment = 8 : i64} : (i32) -> !llvm.ptr
+  %2 = llvm.getelementptr %1[0, 2] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i32, f64, i32)>
+  // CHECK: %[[RES:.*]] = llvm.load %[[ALLOCA]]
+  %3 = llvm.load %2 : !llvm.ptr -> i32
+  // This does not provide side-effect info but it can be removed because it implements PromotableOpInterface.
+  llvm.intr.lifetime.start 2, %2 : !llvm.ptr
+  // CHECK: llvm.return %[[RES]] : i32
+  llvm.return %3 : i32
+}
+
+// -----
+
+// CHECK-LABEL: llvm.func @no_dynamic_indexing
+// CHECK-SAME: (%[[ARG:.*]]: i32)
+llvm.func @no_dynamic_indexing(%arg: i32) -> i32 {
+  // CHECK: %[[SIZE:.*]] = llvm.mlir.constant(1 : i32)
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %[[SIZE]] x !llvm.array<10 x i32> {alignment = 8 : i64} : (i32) -> !llvm.ptr
+  %1 = llvm.alloca %0 x !llvm.array<10 x i32> {alignment = 8 : i64} : (i32) -> !llvm.ptr
+  // CHECK-NOT: = llvm.alloca
+  // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, %[[ARG]]]
+  %2 = llvm.getelementptr %1[0, %arg] : (!llvm.ptr, i32) -> !llvm.ptr, !llvm.array<10 x i32>
+  // CHECK: %[[RES:.*]] = llvm.load %[[GEP]]
+  %3 = llvm.load %2 : !llvm.ptr -> i32
+  // CHECK: llvm.return %[[RES]] : i32
+  llvm.return %3 : i32
+}
+
+// -----
+
+// CHECK-LABEL: llvm.func @no_typed_pointers
+llvm.func @no_typed_pointers() -> i32 {
+  // CHECK: %[[SIZE:.*]] = llvm.mlir.constant(1 : i32)
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %[[SIZE]] x !llvm.array<10 x i32> {alignment = 8 : i64} : (i32) -> !llvm.ptr<array<10 x i32>>
+  %1 = llvm.alloca %0 x !llvm.array<10 x i32> {alignment = 8 : i64} : (i32) -> !llvm.ptr<array<10 x i32>>
+  // CHECK-NOT: = llvm.alloca
+  %2 = llvm.getelementptr %1[0, 1] : (!llvm.ptr<array<10 x i32>>) -> !llvm.ptr<i32>
+  %3 = llvm.load %2 : !llvm.ptr<i32>
+  llvm.return %3 : i32
+}