[mlir] [sroa] Add support for MemRef.

[Moxinilian]

Summary

This patch implements SROA interfaces for MemRef, up to a given fixed
size.

[gysit]

That looks good to me. I left some nit comments and suggestions.

mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp

61

nit: I think the call to the ArrayRef constructor can be dropped.

115

nit: can you use MemRefType memrefType = getType(); here? Also I would prefer memrefType or here since it is a type and not a memref value.

140

nit: I think get type can be use directly as well?

auto destructurable = getType().cast<DestructurableTypeInterface>();

189

I would consider renaming to getAttributeIndexFromIndexOperands?

190

nit: Can you use OperandRange or better ValueRange since we are only interested in the values.

194

nit: I believe the namespace is not needed here?

284

nit: I would prefer memrefType.

301

nit: Consider using memrefType as well.

310

Assuming the coordinate is a signed integer I would add a check >= 0 as well.

mlir/test/Dialect/MemRef/sroa.mlir

10

Does CHECK-DAG make sense here? I would assume we can use CHECK: here since it is one dimensional? Otherwise I would go for the same matching solution as in the high dimensional test case?

29

This tests if there is no memref of a different type before right? If we want to check this we should also do this in the other tests e.g. resolve_alias? Also it may be worth considering to use CHECK-NOT: = memref.alloca() : memref<2x2xi32> to make a bit clearer what we are checking for?

30–32

75

nit: CHECK-DAG -> CHECK?

97

nit: CHECK-DAG -> CHECK

137

nit: CHECK-DAG -> CHECK

[Moxinilian]

Address comments.

[Moxinilian]

mlir/test/Dialect/MemRef/sroa.mlir

29

I'm not sure because I am also checking for absence of any other kind of alloca... But I think it should definitely be everywhere.

[Dinistro]

LGTM!

[gysit]

LGTM!

mlir/test/Dialect/MemRef/sroa.mlir

29

Sure!

mlir/include/mlir/Dialect/MemRef/IR/MemRefMemorySlot.h

+//===- MemRefMemorySlot.h - Implementation of Memory Slot Interfaces ------===//
+//
+// 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_DIALECT_MEMREF_IR_MEMREFMEMORYSLOT_H
+#define MLIR_DIALECT_MEMREF_IR_MEMREFMEMORYSLOT_H
+
+namespace mlir {
+class DialectRegistry;
+
+namespace memref {
+void registerMemorySlotExternalModels(DialectRegistry &registry);
+} // namespace memref
+} // namespace mlir
+
+#endif // MLIR_DIALECT_MEMREF_IR_MEMREFMEMORYSLOT_H

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

 }
 
 //===----------------------------------------------------------------------===//
 // AllocaOp
 //===----------------------------------------------------------------------===//
 
 def MemRef_AllocaOp : AllocLikeOp<"alloca", AutomaticAllocationScopeResource,[
     DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>,
-    DeclareOpInterfaceMethods<PromotableAllocationOpInterface>]> {
+    DeclareOpInterfaceMethods<PromotableAllocationOpInterface>,
+    DeclareOpInterfaceMethods<DestructurableAllocationOpInterface>]> {
   let summary = "stack memory allocation operation";
   let description = [{
     The `alloca` operation allocates memory on the stack, to be automatically
     released when control transfers back from the region of its closest
     surrounding operation with an
     [`AutomaticAllocationScope`](../Traits.md/#automaticallocationscope) trait.
     The amount of memory allocated is specified by its memref and additional
     operands. For example:

 // LoadOp
 //===----------------------------------------------------------------------===//
 
 def LoadOp : MemRef_Op<"load",
      [TypesMatchWith<"result type matches element type of 'memref'",
                      "memref", "result",
                      "::llvm::cast<MemRefType>($_self).getElementType()">,
       MemRefsNormalizable,
-      DeclareOpInterfaceMethods<PromotableMemOpInterface>]> {
+      DeclareOpInterfaceMethods<PromotableMemOpInterface>,
+      DeclareOpInterfaceMethods<DestructurableAccessorOpInterface>]> {
   let summary = "load operation";
   let description = [{
     The `load` op reads an element from a memref specified by an index list. The
     output of load is a new value with the same type as the elements of the
     memref. The arity of indices is the rank of the memref (i.e., if the memref
     loaded from is of rank 3, then 3 indices are required for the load following
     the memref identifier).
 

 // StoreOp
 //===----------------------------------------------------------------------===//
 
 def MemRef_StoreOp : MemRef_Op<"store",
      [TypesMatchWith<"type of 'value' matches element type of 'memref'",
                      "memref", "value",
                      "::llvm::cast<MemRefType>($_self).getElementType()">,
       MemRefsNormalizable,
-      DeclareOpInterfaceMethods<PromotableMemOpInterface>]> {
+      DeclareOpInterfaceMethods<PromotableMemOpInterface>,
+      DeclareOpInterfaceMethods<DestructurableAccessorOpInterface>]> {
   let summary = "store operation";
   let description = [{
     Store a value to a memref location given by indices. The value stored should
     have the same type as the elemental type of the memref. The number of
     arguments provided within brackets need to match the rank of the memref.
 
     In an affine context, the indices of a store are restricted to SSA values
     bound to surrounding loop induction variables,

mlir/include/mlir/InitAllDialects.h

 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/TransformOps/DialectExtension.h"
 #include "mlir/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/Transforms/TilingInterfaceImpl.h"
 #include "mlir/Dialect/MLProgram/IR/MLProgram.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/MemRef/IR/MemRefMemorySlot.h"
 #include "mlir/Dialect/MemRef/IR/ValueBoundsOpInterfaceImpl.h"
 #include "mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.h"
 #include "mlir/Dialect/MemRef/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/MemRef/Transforms/RuntimeOpVerification.h"
 #include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
 #include "mlir/Dialect/OpenACC/OpenACC.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/Dialect/PDL/IR/PDL.h"
inline void registerAllDialects(DialectRegistry &registry) {
   bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(
       registry);
   linalg::registerBufferizableOpInterfaceExternalModels(registry);
   linalg::registerTilingInterfaceExternalModels(registry);
   linalg::registerValueBoundsOpInterfaceExternalModels(registry);
   memref::registerBufferizableOpInterfaceExternalModels(registry);
   memref::registerRuntimeVerifiableOpInterfaceExternalModels(registry);
   memref::registerValueBoundsOpInterfaceExternalModels(registry);
+  memref::registerMemorySlotExternalModels(registry);
   scf::registerBufferizableOpInterfaceExternalModels(registry);
   scf::registerValueBoundsOpInterfaceExternalModels(registry);
   shape::registerBufferizableOpInterfaceExternalModels(registry);
   sparse_tensor::registerBufferizableOpInterfaceExternalModels(registry);
   tensor::registerBufferizableOpInterfaceExternalModels(registry);
   tensor::registerInferTypeOpInterfaceExternalModels(registry);
   tensor::registerTilingInterfaceExternalModels(registry);
   tensor::registerValueBoundsOpInterfaceExternalModels(registry);

mlir/lib/Dialect/MemRef/IR/CMakeLists.txt

 add_mlir_dialect_library(MLIRMemRefDialect
   MemRefDialect.cpp
-  MemRefMem2Reg.cpp
+  MemRefMemorySlot.cpp
   MemRefOps.cpp
   ValueBoundsOpInterfaceImpl.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${PROJECT_SOURCE_DIR}/inlude/mlir/Dialect/MemRefDialect
 
   DEPENDS
   MLIRMemRefOpsIncGen
 
   LINK_COMPONENTS
   Core
 
   LINK_LIBS PUBLIC
   MLIRArithDialect
   MLIRArithUtils
   MLIRControlFlowInterfaces
   MLIRDialect
   MLIRDialectUtils
   MLIRInferTypeOpInterface
   MLIRIR
+  MLIRMemorySlotInterfaces
   MLIRShapedOpInterfaces
   MLIRSideEffectInterfaces
   MLIRValueBoundsOpInterface
   MLIRViewLikeInterface
 )

mlir/lib/Dialect/MemRef/IR/MemRefMem2Reg.cpp

-//===- MemRefMem2Reg.cpp - Mem2Reg 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 Mem2Reg-related interfaces for MemRef dialect
-// operations.
-//
-//===----------------------------------------------------------------------===//
-
-#include "mlir/Dialect/Arith/IR/Arith.h"
-#include "mlir/Dialect/MemRef/IR/MemRef.h"
-#include "llvm/ADT/TypeSwitch.h"
-
-using namespace mlir;
-
-//===----------------------------------------------------------------------===//
-//  AllocaOp interfaces
-//===----------------------------------------------------------------------===//
-
-static bool isSupportedElementType(Type type) {
-  return llvm::isa<MemRefType>(type) ||
-         OpBuilder(type.getContext()).getZeroAttr(type);
-}
-
-SmallVector<MemorySlot> memref::AllocaOp::getPromotableSlots() {
-  MemRefType type = getType();
-  if (!isSupportedElementType(type.getElementType()))
-    return {};
-  if (!type.hasStaticShape())
-    return {};
-  // Make sure the memref contains only a single element.
-  if (any_of(type.getShape(), [](uint64_t dim) { return dim != 1; }))
-    return {};
-
-  return {MemorySlot{getResult(), type.getElementType()}};
-}
-
-Value memref::AllocaOp::getDefaultValue(const MemorySlot &slot,
-                                        RewriterBase &rewriter) {
-  assert(isSupportedElementType(slot.elemType));
-  // TODO: support more types.
-  return TypeSwitch<Type, Value>(slot.elemType)
-      .Case([&](MemRefType t) {
-        return rewriter.create<memref::AllocaOp>(getLoc(), t);
-      })
-      .Default([&](Type t) {
-        return rewriter.create<arith::ConstantOp>(getLoc(), t,
-                                                  rewriter.getZeroAttr(t));
-      });
-}
-
-void memref::AllocaOp::handlePromotionComplete(const MemorySlot &slot,
-                                               Value defaultValue,
-                                               RewriterBase &rewriter) {
-  if (defaultValue.use_empty())
-    rewriter.eraseOp(defaultValue.getDefiningOp());
-  rewriter.eraseOp(*this);
-}
-
-void memref::AllocaOp::handleBlockArgument(const MemorySlot &slot,
-                                           BlockArgument argument,
-                                           RewriterBase &rewriter) {}
-
-//===----------------------------------------------------------------------===//
-//  LoadOp/StoreOp interfaces
-//===----------------------------------------------------------------------===//
-
-bool memref::LoadOp::loadsFrom(const MemorySlot &slot) {
-  return getMemRef() == slot.ptr;
-}
-
-Value memref::LoadOp::getStored(const MemorySlot &slot) { return {}; }
-
-bool memref::LoadOp::canUsesBeRemoved(
-    const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
-    SmallVectorImpl<OpOperand *> &newBlockingUses) {
-  if (blockingUses.size() != 1)
-    return false;
-  Value blockingUse = (*blockingUses.begin())->get();
-  return blockingUse == slot.ptr && getMemRef() == slot.ptr &&
-         getResult().getType() == slot.elemType;
-}
-
-DeletionKind memref::LoadOp::removeBlockingUses(
-    const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
-    RewriterBase &rewriter, Value reachingDefinition) {
-  // `canUsesBeRemoved` checked this blocking use must be the loaded slot
-  // pointer.
-  rewriter.replaceAllUsesWith(getResult(), reachingDefinition);
-  return DeletionKind::Delete;
-}
-
-bool memref::StoreOp::loadsFrom(const MemorySlot &slot) { return false; }
-
-Value memref::StoreOp::getStored(const MemorySlot &slot) {
-  if (getMemRef() != slot.ptr)
-    return {};
-  return getValue();
-}
-
-bool memref::StoreOp::canUsesBeRemoved(
-    const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
-    SmallVectorImpl<OpOperand *> &newBlockingUses) {
-  if (blockingUses.size() != 1)
-    return false;
-  Value blockingUse = (*blockingUses.begin())->get();
-  return blockingUse == slot.ptr && getMemRef() == slot.ptr &&
-         getValue() != slot.ptr && getValue().getType() == slot.elemType;
-}
-
-DeletionKind memref::StoreOp::removeBlockingUses(
-    const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
-    RewriterBase &rewriter, Value reachingDefinition) {
-  return DeletionKind::Delete;
-}

mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp

+//===- MemRefMemorySlot.cpp - Memory Slot 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 Mem2Reg-related interfaces for MemRef dialect
+// operations.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/MemRef/IR/MemRefMemorySlot.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Complex/IR/Complex.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/BuiltinDialect.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Matchers.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/Value.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
+#include "mlir/Interfaces/MemorySlotInterfaces.h"
+#include "mlir/Support/LogicalResult.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/TypeSwitch.h"
+
+using namespace mlir;
+
+//===----------------------------------------------------------------------===//
+//  Utilities
+//===----------------------------------------------------------------------===//
+
+/// Walks over the indices of the elements of a tensor of a given `shape` by
+/// updating `index` in place to the next index. This returns failure if the
+/// provided index was the last index.
+static LogicalResult nextIndex(ArrayRef<int64_t> shape,
+                               MutableArrayRef<int64_t> index) {
+  for (size_t i = 0; i < shape.size(); ++i) {
+    index[i]++;
+    if (index[i] < shape[i])
+      return success();
+    index[i] = 0;
+  }
+  return failure();
+}
+
+/// Calls `walker` for each index within a tensor of a given `shape`, providing
+/// the index as an array attribute of the coordinates.
+template <typename CallableT>
+static void walkIndicesAsAttr(MLIRContext *ctx, ArrayRef<int64_t> shape,
+                              CallableT &&walker) {
+  Type indexType = IndexType::get(ctx);
+  SmallVector<int64_t> shapeIter(shape.size(), 0);
+  do {
+    SmallVector<Attribute> indexAsAttr;
+    for (int64_t dim : shapeIter)
+      indexAsAttr.push_back(IntegerAttr::get(indexType, dim));
+    walker(ArrayAttr::get(ctx, indexAsAttr));
+  } while (succeeded(nextIndex(shape, shapeIter)));
+}
+
+//===----------------------------------------------------------------------===//
+//  Interfaces for AllocaOp
+//===----------------------------------------------------------------------===//
+
+static bool isSupportedElementType(Type type) {
+  return type.isa<MemRefType>() ||
+         OpBuilder(type.getContext()).getZeroAttr(type);
+}
+
+SmallVector<MemorySlot> memref::AllocaOp::getPromotableSlots() {
+  MemRefType type = getType();
+  if (!isSupportedElementType(type.getElementType()))
+    return {};
+  if (!type.hasStaticShape())
+    return {};
+  // Make sure the memref contains only a single element.
+  if (type.getNumElements() != 1)
+    return {};
+
+  return {MemorySlot{getResult(), type.getElementType()}};
+}
+
+Value memref::AllocaOp::getDefaultValue(const MemorySlot &slot,
+                                        RewriterBase &rewriter) {
+  assert(isSupportedElementType(slot.elemType));
+  // TODO: support more types.
+  return TypeSwitch<Type, Value>(slot.elemType)
+      .Case([&](MemRefType t) {
+        return rewriter.create<memref::AllocaOp>(getLoc(), t);
+      })
+      .Default([&](Type t) {
+        return rewriter.create<arith::ConstantOp>(getLoc(), t,
+                                                  rewriter.getZeroAttr(t));
+      });
+}
+
+void memref::AllocaOp::handlePromotionComplete(const MemorySlot &slot,
+                                               Value defaultValue,
+                                               RewriterBase &rewriter) {
+  if (defaultValue.use_empty())
+    rewriter.eraseOp(defaultValue.getDefiningOp());
+  rewriter.eraseOp(*this);
+}
+
+void memref::AllocaOp::handleBlockArgument(const MemorySlot &slot,
+                                           BlockArgument argument,
+                                           RewriterBase &rewriter) {}
+
+SmallVector<DestructurableMemorySlot>
+memref::AllocaOp::getDestructurableSlots() {
+  MemRefType memrefType = getType();
+  auto destructurable = memrefType.dyn_cast<DestructurableTypeInterface>();
+  if (!destructurable)
+    return {};
+
+  Optional<DenseMap<Attribute, Type>> destructuredType =
+      destructurable.getSubelementIndexMap();
+  if (!destructuredType)
+    return {};
+
+  DenseMap<Attribute, Type> indexMap;
+  for (auto const &[index, type] : *destructuredType)
+    indexMap.insert({index, MemRefType::get({}, type)});
+
+  return {DestructurableMemorySlot{{getMemref(), memrefType}, indexMap}};
+}
+
+DenseMap<Attribute, MemorySlot>
+memref::AllocaOp::destructure(const DestructurableMemorySlot &slot,
+                              const SmallPtrSetImpl<Attribute> &usedIndices,
+                              RewriterBase &rewriter) {
+  rewriter.setInsertionPointAfter(*this);
+
+  DenseMap<Attribute, MemorySlot> slotMap;
+
+  auto memrefType = getType().cast<DestructurableTypeInterface>();
+  for (Attribute usedIndex : usedIndices) {
+    Type elemType = memrefType.getTypeAtIndex(usedIndex);
+    MemRefType elemPtr = MemRefType::get({}, elemType);
+    auto subAlloca = rewriter.create<memref::AllocaOp>(getLoc(), elemPtr);
+    slotMap.try_emplace<MemorySlot>(usedIndex,
+                                    {subAlloca.getResult(), elemType});
+  }
+
+  return slotMap;
+}
+
+void memref::AllocaOp::handleDestructuringComplete(
+    const DestructurableMemorySlot &slot, RewriterBase &rewriter) {
+  assert(slot.ptr == getResult());
+  rewriter.eraseOp(*this);
+}
+
+//===----------------------------------------------------------------------===//
+//  Interfaces for LoadOp/StoreOp
+//===----------------------------------------------------------------------===//
+
+bool memref::LoadOp::loadsFrom(const MemorySlot &slot) {
+  return getMemRef() == slot.ptr;
+}
+
+Value memref::LoadOp::getStored(const MemorySlot &slot) { return {}; }
+
+bool memref::LoadOp::canUsesBeRemoved(
+    const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
+    SmallVectorImpl<OpOperand *> &newBlockingUses) {
+  if (blockingUses.size() != 1)
+    return false;
+  Value blockingUse = (*blockingUses.begin())->get();
+  return blockingUse == slot.ptr && getMemRef() == slot.ptr &&
+         getResult().getType() == slot.elemType;
+}
+
+DeletionKind memref::LoadOp::removeBlockingUses(
+    const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
+    RewriterBase &rewriter, Value reachingDefinition) {
+  // `canUsesBeRemoved` checked this blocking use must be the loaded slot
+  // pointer.
+  rewriter.replaceAllUsesWith(getResult(), reachingDefinition);
+  return DeletionKind::Delete;
+}
+
+/// Returns the index of a memref in attribute form, given its indices.
+static Attribute getAttributeIndexFromIndexOperands(MLIRContext *ctx,
+                                                    ValueRange indices) {
+  SmallVector<Attribute> index;
+  for (Value coord : indices) {
+    IntegerAttr coordAttr;
+    if (!matchPattern(coord, m_Constant<IntegerAttr>(&coordAttr)))
+      return {};
+    index.push_back(coordAttr);
+  }
+  return ArrayAttr::get(ctx, index);
+}
+
+bool memref::LoadOp::canRewire(const DestructurableMemorySlot &slot,
+                               SmallPtrSetImpl<Attribute> &usedIndices,
+                               SmallVectorImpl<MemorySlot> &mustBeSafelyUsed) {
+  if (slot.ptr != getMemRef())
+    return false;
+  Attribute index =
+      getAttributeIndexFromIndexOperands(getContext(), getIndices());
+  if (!index)
+    return false;
+  usedIndices.insert(index);
+  return true;
+}
+
+DeletionKind memref::LoadOp::rewire(const DestructurableMemorySlot &slot,
+                                    DenseMap<Attribute, MemorySlot> &subslots,
+                                    RewriterBase &rewriter) {
+  Attribute index =
+      getAttributeIndexFromIndexOperands(getContext(), getIndices());
+  const MemorySlot &memorySlot = subslots.at(index);
+  rewriter.updateRootInPlace(*this, [&]() {
+    setMemRef(memorySlot.ptr);
+    getIndicesMutable().clear();
+  });
+  return DeletionKind::Keep;
+}
+
+bool memref::StoreOp::loadsFrom(const MemorySlot &slot) { return false; }
+
+Value memref::StoreOp::getStored(const MemorySlot &slot) {
+  if (getMemRef() != slot.ptr)
+    return {};
+  return getValue();
+}
+
+bool memref::StoreOp::canUsesBeRemoved(
+    const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
+    SmallVectorImpl<OpOperand *> &newBlockingUses) {
+  if (blockingUses.size() != 1)
+    return false;
+  Value blockingUse = (*blockingUses.begin())->get();
+  return blockingUse == slot.ptr && getMemRef() == slot.ptr &&
+         getValue() != slot.ptr && getValue().getType() == slot.elemType;
+}
+
+DeletionKind memref::StoreOp::removeBlockingUses(
+    const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
+    RewriterBase &rewriter, Value reachingDefinition) {
+  return DeletionKind::Delete;
+}
+
+bool memref::StoreOp::canRewire(const DestructurableMemorySlot &slot,
+                                SmallPtrSetImpl<Attribute> &usedIndices,
+                                SmallVectorImpl<MemorySlot> &mustBeSafelyUsed) {
+  if (slot.ptr != getMemRef() || getValue() == slot.ptr)
+    return false;
+  Attribute index =
+      getAttributeIndexFromIndexOperands(getContext(), getIndices());
+  if (!index || !slot.elementPtrs.contains(index))
+    return false;
+  usedIndices.insert(index);
+  return true;
+}
+
+DeletionKind memref::StoreOp::rewire(const DestructurableMemorySlot &slot,
+                                     DenseMap<Attribute, MemorySlot> &subslots,
+                                     RewriterBase &rewriter) {
+  Attribute index =
+      getAttributeIndexFromIndexOperands(getContext(), getIndices());
+  const MemorySlot &memorySlot = subslots.at(index);
+  rewriter.updateRootInPlace(*this, [&]() {
+    setMemRef(memorySlot.ptr);
+    getIndicesMutable().clear();
+  });
+  return DeletionKind::Keep;
+}
+
+//===----------------------------------------------------------------------===//
+//  Interfaces for destructurable types
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+struct MemRefDestructurableTypeExternalModel
+    : public DestructurableTypeInterface::ExternalModel<
+          MemRefDestructurableTypeExternalModel, MemRefType> {
+  std::optional<DenseMap<Attribute, Type>>
+  getSubelementIndexMap(Type type) const {
+    auto memrefType = type.cast<MemRefType>();
+    constexpr int64_t maxMemrefSizeForDestructuring = 16;
+    if (!memrefType.hasStaticShape() ||
+        memrefType.getNumElements() > maxMemrefSizeForDestructuring ||
+        memrefType.getNumElements() == 1)
+      return {};
+
+    DenseMap<Attribute, Type> destructured;
+    walkIndicesAsAttr(
+        memrefType.getContext(), memrefType.getShape(), [&](Attribute index) {
+          destructured.insert({index, memrefType.getElementType()});
+        });
+
+    return destructured;
+  }
+
+  Type getTypeAtIndex(Type type, Attribute index) const {
+    auto memrefType = type.cast<MemRefType>();
+    auto coordArrAttr = index.dyn_cast<ArrayAttr>();
+    if (!coordArrAttr || coordArrAttr.size() != memrefType.getShape().size())
+      return {};
+
+    Type indexType = IndexType::get(memrefType.getContext());
+    for (const auto &[coordAttr, dimSize] :
+         llvm::zip(coordArrAttr, memrefType.getShape())) {
+      auto coord = coordAttr.dyn_cast<IntegerAttr>();
+      if (!coord || coord.getType() != indexType || coord.getInt() < 0 ||
+          coord.getInt() >= dimSize)
+        return {};
+    }
+
+    return memrefType.getElementType();
+  }
+};
+
+} // namespace
+
+//===----------------------------------------------------------------------===//
+//  Register external models
+//===----------------------------------------------------------------------===//
+
+void mlir::memref::registerMemorySlotExternalModels(DialectRegistry &registry) {
+  registry.addExtension(+[](MLIRContext *ctx, BuiltinDialect *dialect) {
+    MemRefType::attachInterface<MemRefDestructurableTypeExternalModel>(*ctx);
+  });
+}

mlir/test/Dialect/MemRef/sroa.mlir

+// RUN: mlir-opt %s --pass-pipeline="builtin.module(func.func(sroa))" --split-input-file | FileCheck %s
+
+// CHECK-LABEL: func.func @basic
+// CHECK-SAME: (%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32)
+func.func @basic(%arg0: i32, %arg1: i32) -> i32 {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  // CHECK-NOT: = memref.alloca()
+  // CHECK-COUNT-2: = memref.alloca() : memref<i32>
+  // CHECK-NOT: = memref.alloca()
+  %alloca = memref.alloca() : memref<2xi32>
+  // CHECK: memref.store %[[ARG0]], %[[ALLOCA0:.*]][]
+  memref.store %arg0, %alloca[%c0] : memref<2xi32>
+  // CHECK: memref.store %[[ARG1]], %[[ALLOCA1:.*]][]
+  memref.store %arg1, %alloca[%c1] : memref<2xi32>
+  // CHECK: %[[RES:.*]] = memref.load %[[ALLOCA0]][]
+  %res = memref.load %alloca[%c0] : memref<2xi32>
+  // CHECK: return %[[RES]] : i32
+  return %res : i32
+}
+
+// -----
+
+// CHECK-LABEL: func.func @basic_high_dimensions
+// CHECK-SAME: (%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32, %[[ARG2:.*]]: i32)
+func.func @basic_high_dimensions(%arg0: i32, %arg1: i32, %arg2: i32) -> i32 {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  // CHECK-NOT: = memref.alloca()
+  // CHECK-COUNT-3: = memref.alloca() : memref<i32>
+  // CHECK-NOT: = memref.alloca()
+  %alloca = memref.alloca() : memref<2x2xi32>
+  // CHECK: memref.store %[[ARG0]], %[[ALLOCA0:.*]][]
+  memref.store %arg0, %alloca[%c0, %c0] : memref<2x2xi32>
+  // CHECK: memref.store %[[ARG1]], %[[ALLOCA1:.*]][]
+  memref.store %arg1, %alloca[%c0, %c1] : memref<2x2xi32>
+  // CHECK: memref.store %[[ARG2]], %[[ALLOCA2:.*]][]
+  memref.store %arg2, %alloca[%c1, %c0] : memref<2x2xi32>
+  // CHECK: %[[RES:.*]] = memref.load %[[ALLOCA1]][]
+  %res = memref.load %alloca[%c0, %c1] : memref<2x2xi32>
+  // CHECK: return %[[RES]] : i32
+  return %res : i32
+}
+
+// -----
+
+// CHECK-LABEL: func.func @resolve_alias
+// CHECK-SAME: (%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32)
+func.func @resolve_alias(%arg0: i32, %arg1: i32) -> i32 {
+  %c0 = arith.constant 0 : index
+  // CHECK-NOT: = memref.alloca()
+  // CHECK: %[[ALLOCA:.*]] = memref.alloca() : memref<i32>
+  // CHECK-NOT: = memref.alloca()
+  %alloca = memref.alloca() : memref<2xi32>
+  // CHECK: memref.store %[[ARG0]], %[[ALLOCA]][]
+  memref.store %arg0, %alloca[%c0] : memref<2xi32>
+  // CHECK: memref.store %[[ARG1]], %[[ALLOCA]][]
+  memref.store %arg1, %alloca[%c0] : memref<2xi32>
+  // CHECK: %[[RES:.*]] = memref.load %[[ALLOCA]][]
+  %res = memref.load %alloca[%c0] : memref<2xi32>
+  // CHECK: return %[[RES]] : i32
+  return %res : i32
+}
+
+// -----
+
+// CHECK-LABEL: func.func @no_direct_use
+// CHECK-SAME: (%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32)
+func.func @no_direct_use(%arg0: i32, %arg1: i32) -> i32 {
+  // CHECK: %[[C0:.*]] = arith.constant 0 : index
+  %c0 = arith.constant 0 : index
+  // CHECK: %[[C1:.*]] = arith.constant 1 : index
+  %c1 = arith.constant 1 : index
+  // CHECK-NOT: = memref.alloca()
+  // CHECK: %[[ALLOCA:.*]] = memref.alloca() : memref<2xi32>
+  // CHECK-NOT: = memref.alloca()
+  %alloca = memref.alloca() : memref<2xi32>
+  // CHECK: memref.store %[[ARG0]], %[[ALLOCA]][%[[C0]]]
+  memref.store %arg0, %alloca[%c0] : memref<2xi32>
+  // CHECK: memref.store %[[ARG1]], %[[ALLOCA]][%[[C1]]]
+  memref.store %arg1, %alloca[%c1] : memref<2xi32>
+  // CHECK: %[[RES:.*]] = memref.load %[[ALLOCA]][%[[C0]]]
+  %res = memref.load %alloca[%c0] : memref<2xi32>
+  call @use(%alloca) : (memref<2xi32>) -> ()
+  // CHECK: return %[[RES]] : i32
+  return %res : i32
+}
+
+func.func @use(%foo: memref<2xi32>) { return }
+
+// -----
+
+// CHECK-LABEL: func.func @no_dynamic_indexing
+// CHECK-SAME: (%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32, %[[INDEX:.*]]: index)
+func.func @no_dynamic_indexing(%arg0: i32, %arg1: i32, %index: index) -> i32 {
+  // CHECK: %[[C0:.*]] = arith.constant 0 : index
+  %c0 = arith.constant 0 : index
+  // CHECK-NOT: = memref.alloca()
+  // CHECK: %[[ALLOCA:.*]] = memref.alloca() : memref<2xi32>
+  // CHECK-NOT: = memref.alloca()
+  %alloca = memref.alloca() : memref<2xi32>
+  // CHECK: memref.store %[[ARG0]], %[[ALLOCA]][%[[C0]]]
+  memref.store %arg0, %alloca[%c0] : memref<2xi32>
+  // CHECK: memref.store %[[ARG1]], %[[ALLOCA]][%[[INDEX]]]
+  memref.store %arg1, %alloca[%index] : memref<2xi32>
+  // CHECK: %[[RES:.*]] = memref.load %[[ALLOCA]][%[[C0]]]
+  %res = memref.load %alloca[%c0] : memref<2xi32>
+  // CHECK: return %[[RES]] : i32
+  return %res : i32
+}
+
+// -----
+
+// CHECK-LABEL: func.func @no_dynamic_shape
+// CHECK-SAME: (%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32)
+func.func @no_dynamic_shape(%arg0: i32, %arg1: i32) -> i32 {
+  // CHECK: %[[C0:.*]] = arith.constant 0 : index
+  %c0 = arith.constant 0 : index
+  // CHECK: %[[C1:.*]] = arith.constant 1 : index
+  %c1 = arith.constant 1 : index
+  // CHECK-NOT: = memref.alloca()
+  // CHECK: %[[ALLOCA:.*]] = memref.alloca(%[[C1]]) : memref<?x2xi32>
+  // CHECK-NOT: = memref.alloca()
+  %alloca = memref.alloca(%c1) : memref<?x2xi32>
+  // CHECK: memref.store %[[ARG0]], %[[ALLOCA]][%[[C0]], %[[C0]]]
+  memref.store %arg0, %alloca[%c0, %c0] : memref<?x2xi32>
+  // CHECK: %[[RES:.*]] = memref.load %[[ALLOCA]][%[[C0]], %[[C0]]]
+  %res = memref.load %alloca[%c0, %c0] : memref<?x2xi32>
+  // CHECK: return %[[RES]] : i32
+  return %res : i32
+}
+
+// -----
+
+// CHECK-LABEL: func.func @no_out_of_bounds
+// CHECK-SAME: (%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32)
+func.func @no_out_of_bounds(%arg0: i32, %arg1: i32) -> i32 {
+  // CHECK: %[[C0:.*]] = arith.constant 0 : index
+  %c0 = arith.constant 0 : index
+  // CHECK: %[[C100:.*]] = arith.constant 100 : index
+  %c100 = arith.constant 100 : index
+  // CHECK-NOT: = memref.alloca()
+  // CHECK: %[[ALLOCA:.*]] = memref.alloca() : memref<2xi32>
+  // CHECK-NOT: = memref.alloca()
+  %alloca = memref.alloca() : memref<2xi32>
+  // CHECK: memref.store %[[ARG0]], %[[ALLOCA]][%[[C0]]]
+  memref.store %arg0, %alloca[%c0] : memref<2xi32>
+  // CHECK: memref.store %[[ARG1]], %[[ALLOCA]][%[[C100]]]
+  memref.store %arg1, %alloca[%c100] : memref<2xi32>
+  // CHECK: %[[RES:.*]] = memref.load %[[ALLOCA]][%[[C0]]]
+  %res = memref.load %alloca[%c0] : memref<2xi32>
+  // CHECK: return %[[RES]] : i32
+  return %res : i32
+}