[Mlir-commits] [mlir] [mlir][mem2reg] Promote memory slots through transparent view operations (PR #196924)

[llvmlistbot at llvm.org]

================
@@ -8,5 +8,153 @@
 
 #include "mlir/Interfaces/MemorySlotInterfaces.h"
 
+#include "llvm/ADT/SmallVector.h"
+
 #include "mlir/Interfaces/MemorySlotOpInterfaces.cpp.inc"
 #include "mlir/Interfaces/MemorySlotTypeInterfaces.cpp.inc"
+
+using namespace mlir;
+
+/// Returns the slot describing `aliasPtr`: `rootSlot` if it is the root,
+/// the entry in `aliasMap` if it's a known alias, or `nullopt` otherwise.
+static std::optional<MemorySlot>
+getParentSlot(Value aliasPtr, const MemorySlot &rootSlot,
+              const PromotableAliasMap &aliasMap) {
+  if (aliasPtr == rootSlot.ptr)
+    return rootSlot;
+  auto it = aliasMap.find(aliasPtr);
+  if (it == aliasMap.end())
+    return std::nullopt;
+  return it->second.slot;
+}
+
+void mlir::populatePromotableAliasMap(PromotableAliaserInterface aliaser,
+                                      const MemorySlot &rootSlot,
+                                      PromotableAliasMap &aliasMap) {
+  for (OpOperand &operand : aliaser->getOpOperands()) {
+    std::optional<MemorySlot> parentSlot =
+        getParentSlot(operand.get(), rootSlot, aliasMap);
+    if (!parentSlot)
+      continue;
+    SmallVector<MemorySlot, 2> newSlots;
+    aliaser.getPromotableSlotAliases(operand, *parentSlot, newSlots);
+    for (const MemorySlot &alias : newSlots)
+      aliasMap.try_emplace(alias.ptr, PromotableSlotAliasInfo{alias, &operand});
+  }
+}
+
+std::optional<MemorySlot>
+mlir::getOpAliasSlot(Operation *op, const MemorySlot &rootSlot,
+                     const PromotableAliasMap &aliasMap) {
+  for (Value operand : op->getOperands())
+    if (std::optional<MemorySlot> slot =
+            getParentSlot(operand, rootSlot, aliasMap))
+      return slot;
+  return std::nullopt;
+}
+
+bool mlir::isUsingAtMostOneSlotAlias(Operation *op, const MemorySlot &rootSlot,
+                                     const PromotableAliasMap &aliasMap) {
+  Value uniqueAliasPtr;
+  for (Value operand : op->getOperands()) {
+    std::optional<MemorySlot> slot = getParentSlot(operand, rootSlot, aliasMap);
+    if (!slot)
+      continue;
+    if (uniqueAliasPtr && uniqueAliasPtr != slot->ptr)
+      return false;
+    uniqueAliasPtr = slot->ptr;
+  }
+  return true;
+}
+
+namespace {
+/// A step in an alias chain, from leaf to root. `parentSlot` is one step
+/// closer to the root; `aliasSlot` is the slot exposed at this step.
+struct ChainStep {
+  PromotableAliaserInterface aliaser;
+  OpOperand *aliasedSlotPointerOperand;
+  MemorySlot parentSlot;
+  MemorySlot aliasSlot;
+};
+} // namespace
+
+/// Walks back from `aliasPtr` to `rootSlot.ptr` through `aliasMap`,
+/// populating `chainOut` from leaf to root. Returns false if `aliasPtr`
+/// is not a known alias of `rootSlot`.
+static bool buildAliasChain(Value aliasPtr, const MemorySlot &rootSlot,
----------------
jeanPerier wrote:

Agree this is cleaner. Updated.

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