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

[Tobias Gysi]

================
@@ -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;
----------------
gysit wrote:

Would it possible to just return null if there are multiple operands that point to the same root slot? That way the separate `isUsingAtMostOneSlotAlias` check may not be necessary. Or would that break some uses of this helper?

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