blazingly fasterer

2026-07-21 15:43:35 +02:00
parent a893d23a74
commit 3e468b58c8
37 changed files with 1119 additions and 933 deletions
@@ -0,0 +1,71 @@
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Interfaces/DestinationStyleOpInterface.h"
+
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallVector.h"
+
+#include "src/Accelerators/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace pim {
+
+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) {
+  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();
+    if (!visitedValues.insert(value).second)
+      continue;
+    for (Operation* user : value.getUsers()) {
+      if (!visitedUsers.insert(user).second)
+        continue;
+      if (failed(visitUser(value, user)))
+        return failure();
+
+      if (isLocalMemoryAliasOp(user))
+        for (Value result : user->getResults())
+          addAlias(result);
+
+      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 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));
+          }
+
+      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));
+      }
+    }
+  }
+  return success();
+}
+
+} // namespace pim
+} // namespace onnx_mlir