unify memory opimization passes

better reports
cleanups

src/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.cpp

@@ -2,70 +2,204 @@
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Interfaces/DestinationStyleOpInterface.h"
 
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 
 using namespace mlir;
 
-namespace onnx_mlir {
-namespace pim {
+namespace onnx_mlir::pim {
+namespace {
+
+struct OperationOrdering {
+  DenseMap<Operation*, uint64_t> position;
+  DenseMap<Operation*, uint64_t> subtreeEnd;
+  uint64_t next = 0;
+};
+
+static void orderOperation(Operation* op, OperationOrdering& ordering) {
+  uint64_t end = ordering.next;
+  ordering.position[op] = ordering.next++;
+  for (Region& region : op->getRegions())
+    for (Block& block : region)
+      for (Operation& nested : block) {
+        orderOperation(&nested, ordering);
+        end = std::max(end, ordering.subtreeEnd.lookup(&nested));
+      }
+  ordering.subtreeEnd[op] = end;
+}
+
+static OperationOrdering buildOrdering(Operation* coreLikeOp) {
+  OperationOrdering ordering;
+  if (!coreLikeOp || coreLikeOp->getNumRegions() != 1 || coreLikeOp->getRegion(0).empty())
+    return ordering;
+  for (Operation& op : coreLikeOp->getRegion(0).front())
+    orderOperation(&op, ordering);
+  return ordering;
+}
+
+static bool isRuntimeMemoryTouch(Operation* op) {
+  return isa<PimMemCopyHostToDevOp,
+             PimMemCopyDevToHostOp,
+             PimMemCopyOp,
+             PimReceiveOp,
+             PimSendOp,
+             PimConcatOp,
+             PimVMMOp,
+             PimTransposeOp,
+             PimVVAddOp,
+             PimVVSubOp,
+             PimVVMulOp,
+             PimVVMaxOp,
+             PimVVDMulOp,
+             PimVAvgOp,
+             PimVReluOp,
+             PimVTanhOp,
+             PimVSigmOp,
+             PimVSoftmaxOp>(op);
+}
+
+static bool isIgnoredUser(Operation* op) {
+  return isLocalMemoryAliasOp(op) || isa<scf::ForOp, scf::YieldOp, memref::DeallocOp>(op)
+      || isCoreStaticAddressOp(op);
+}
+
+static bool isWithin(Value value, Region* region) {
+  if (auto argument = dyn_cast<BlockArgument>(value))
+    return argument.getOwner()->getParent() == region;
+  if (Operation* definingOp = value.getDefiningOp())
+    return definingOp->getParentRegion() == region || region->isAncestor(definingOp->getParentRegion());
+  return false;
+}
+
+static std::pair<uint64_t, uint64_t> getTouchRange(Value definingValue, Operation* user, const OperationOrdering& ordering) {
+  uint64_t start = ordering.position.lookup(user);
+  uint64_t end = start;
+  for (Operation* current = user; current; current = current->getParentOp())
+    if (auto forOp = dyn_cast<scf::ForOp>(current); forOp && !isWithin(definingValue, &forOp.getRegion())) {
+      start = std::min(start, ordering.position.lookup(forOp));
+      end = std::max(end, ordering.subtreeEnd.lookup(forOp));
+    }
+  return {start, end};
+}
+
+static FailureOr<size_t> getAllocationSize(memref::AllocOp allocation) {
+  auto type = dyn_cast<ShapedType>(allocation.getType());
+  if (!type)
+    return failure();
+  auto bytes = getCheckedShapedTypeSizeInBytes(type, allocation, "memory allocation byte size");
+  if (failed(bytes))
+    return failure();
+  return checkedSize(*bytes, allocation, "memory allocation byte size");
+}
+
+} // namespace
 
 bool isLocalMemoryAliasOp(Operation* op) {
   return isa<memref::SubViewOp, memref::CastOp, memref::CollapseShapeOp, memref::ExpandShapeOp>(op);
 }
 
-LogicalResult walkLocalMemoryUses(Value root,
-                                  llvm::function_ref<LogicalResult(Value, Operation*)> visitUser) {
+LogicalResult walkLocalMemoryUses(Value root, LocalMemoryUseCallback callback) {
   llvm::SmallPtrSet<Value, 16> visitedValues;
-  llvm::SmallPtrSet<Operation*, 32> visitedUsers;
-  llvm::SmallVector<Value> pendingValues {root};
-  auto addAlias = [&](Value value) { pendingValues.push_back(value); };
-
-  while (!pendingValues.empty()) {
-    Value value = pendingValues.pop_back_val();
+  llvm::SmallPtrSet<OpOperand*, 32> visitedUses;
+  SmallVector<Value> pending = {root};
+  while (!pending.empty()) {
+    Value value = pending.pop_back_val();
     if (!visitedValues.insert(value).second)
       continue;
-    for (Operation* user : value.getUsers()) {
-      if (!visitedUsers.insert(user).second)
-        continue;
-      if (failed(visitUser(value, user)))
+    for (OpOperand& use : value.getUses()) {
+      if (!visitedUses.insert(&use).second || failed(callback(use)))
         return failure();
+      Operation* user = use.getOwner();
 
-      if (isLocalMemoryAliasOp(user))
-        for (Value result : user->getResults())
-          addAlias(result);
+      if (isLocalMemoryAliasOp(user) && use.getOperandNumber() == 0) {
+        if (user->getNumResults() != 1)
+          return failure();
+        pending.push_back(user->getResult(0));
+      }
 
-      if (auto dpsOp = dyn_cast<DestinationStyleOpInterface>(user))
-        for (OpResult result : user->getResults())
-          if (OpOperand* tied = dpsOp.getTiedOpOperand(result); tied && tied->get() == value)
-            addAlias(result);
+      if (auto dpsOp = dyn_cast<DestinationStyleOpInterface>(user); dpsOp && dpsOp.isDpsInit(&use))
+        pending.push_back(dpsOp.getTiedOpResult(&use));
 
       if (auto forOp = dyn_cast<scf::ForOp>(user))
-        for (auto [index, initArg] : llvm::enumerate(forOp.getInitArgs()))
-          if (initArg == value) {
-            addAlias(forOp.getRegionIterArgs()[index]);
-            addAlias(forOp.getResult(index));
+        for (auto [index, init] : llvm::enumerate(forOp.getInitArgsMutable()))
+          if (&init == &use) {
+            pending.push_back(forOp.getRegionIterArgs()[index]);
+            pending.push_back(forOp.getResult(index));
           }
 
-      auto yieldOp = dyn_cast<scf::YieldOp>(user);
-      if (!yieldOp)
-        continue;
-      for (auto [index, operand] : llvm::enumerate(yieldOp.getOperands())) {
-        if (operand != value)
-          continue;
-        if (auto forOp = dyn_cast<scf::ForOp>(yieldOp->getParentOp()))
-          addAlias(forOp.getResult(index));
-        else if (auto ifOp = dyn_cast<scf::IfOp>(yieldOp->getParentOp()))
-          addAlias(ifOp.getResult(index));
-        else if (auto switchOp = dyn_cast<scf::IndexSwitchOp>(yieldOp->getParentOp()))
-          addAlias(switchOp.getResult(index));
+      if (auto yieldOp = dyn_cast<scf::YieldOp>(user)) {
+        unsigned index = use.getOperandNumber();
+        Operation* parent = yieldOp->getParentOp();
+        if (auto forOp = dyn_cast<scf::ForOp>(parent))
+          pending.push_back(forOp.getResult(index));
+        else if (auto ifOp = dyn_cast<scf::IfOp>(parent))
+          pending.push_back(ifOp.getResult(index));
+        else if (auto switchOp = dyn_cast<scf::IndexSwitchOp>(parent))
+          pending.push_back(switchOp.getResult(index));
+        else
+          return failure();
       }
     }
   }
   return success();
 }
 
-} // namespace pim
-} // namespace onnx_mlir
+FailureOr<SmallVector<LocalMemoryInterval, 0>> analyzeLocalMemoryLifetimes(Operation* coreLikeOp) {
+  SmallVector<LocalMemoryInterval, 0> intervals;
+  OperationOrdering ordering = buildOrdering(coreLikeOp);
+  if (ordering.position.empty())
+    return intervals;
+  uint64_t fallbackEnd = ordering.next - 1;
+  bool failedSize = false;
+  coreLikeOp->walk([&](memref::AllocOp allocation) {
+    auto size = getAllocationSize(allocation);
+    if (failed(size)) {
+      failedSize = true;
+      return;
+    }
+
+    LocalMemoryInterval interval {allocation, ordering.position.lookup(allocation), ordering.position.lookup(allocation), *size};
+    bool hasRuntimeTouch = false;
+    bool needsFallback = false;
+    if (failed(walkLocalMemoryUses(allocation.getResult(), [&](OpOperand& use) {
+      Operation* user = use.getOwner();
+      if (isRuntimeMemoryTouch(user)) {
+        auto [start, end] = getTouchRange(allocation.getResult(), user, ordering);
+        if (!hasRuntimeTouch) {
+          interval.start = start;
+          interval.end = end;
+          hasRuntimeTouch = true;
+        }
+        else {
+          interval.start = std::min(interval.start, start);
+          interval.end = std::max(interval.end, end);
+        }
+      }
+      else if (!isIgnoredUser(user)) {
+        needsFallback = true;
+      }
+      return success();
+    })))
+      needsFallback = true;
+    if (!hasRuntimeTouch || needsFallback) {
+      interval.start = std::min(interval.start, ordering.position.lookup(allocation));
+      interval.end = fallbackEnd;
+    }
+    intervals.push_back(interval);
+  });
+  if (failedSize)
+    return failure();
+  return intervals;
+}
+
+bool localMemoryLifetimesOverlap(const LocalMemoryInterval& lhs, const LocalMemoryInterval& rhs) {
+  return !(lhs.end < rhs.start || rhs.end < lhs.start);
+}
+
+} // namespace onnx_mlir::pim