Memory Liveness

2026-06-03 18:15:30 +02:00
parent 2a8faf9c6b
commit 20cf40c9ba
15 changed files with 1263 additions and 112 deletions
@@ -2,7 +2,6 @@
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
-#include "mlir/IR/AsmState.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
@@ -23,11 +22,11 @@
 #include <absl/types/compare.h>
 #include <algorithm>
 #include <cassert>
-#include <climits>
 #include <cstdint>
 #include <fstream>
 #include <limits>
 #include <memory>
+#include <numeric>
 #include <string>
 #include <utility>

@@ -38,10 +37,12 @@
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/WeightUtils.hpp"
+#include "src/Accelerators/PIM/Common/Support/FileSystemUtils.hpp"
 #include "src/Accelerators/PIM/Compiler/PimArtifactWriter.hpp"
 #include "src/Accelerators/PIM/Compiler/PimBinaryFormat.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCodeGen.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
+#include "src/Accelerators/PIM/Compiler/PimMemoryLiveness.hpp"
 #include "src/Accelerators/PIM/Compiler/PimWeightEmitter.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"

@@ -97,6 +98,51 @@ static int32_t getVectorByteSizeOrCrash(ShapedType type) {
  return pim::checkedI32OrCrash(*byteSize, "vector byte size");
 }

+static Operation *getDiagnosticAnchor(mlir::Value value) {
+  if (Operation *definingOp = value.getDefiningOp())
+    return definingOp;
+  if (auto blockArg = dyn_cast<BlockArgument>(value))
+    return blockArg.getOwner()->getParentOp();
+  return nullptr;
+}
+
+// PIM instruction immediates are serialized as signed int32_t fields today
+// (`sldi` goes through checkedI32OrCrash), so local addresses must stay within
+// the non-negative int32_t range.
+static constexpr size_t kPimAddressLimit = static_cast<size_t>(std::numeric_limits<int32_t>::max());
+
+static FailureOr<size_t> checkedAlignTo(size_t value, size_t alignment, Operation *anchor, StringRef fieldName) {
+  if (alignment == 0)
+    return value;
+  size_t remainder = value % alignment;
+  if (remainder == 0)
+    return value;
+  return pim::checkedAdd(value, alignment - remainder, anchor, fieldName);
+}
+
+static void printMemoryOverflowDiagnostic(mlir::Value value,
+                                          const MemoryValueKey &key,
+                                          size_t requestedSize,
+                                          size_t currentFirstAvailableAddress,
+                                          size_t alignedEndAddress) {
+  llvm::errs() << "PIM local memory allocation overflow\n";
+  llvm::errs() << "Requested allocation size: " << requestedSize << " bytes\n";
+  llvm::errs() << "Current firstAvailableAddress: " << currentFirstAvailableAddress << "\n";
+  llvm::errs() << "Aligned end address: " << alignedEndAddress << "\n";
+  llvm::errs() << "Address limit: " << kPimAddressLimit << " (signed int32_t immediate range)\n";
+  if (key.lane)
+    llvm::errs() << "Lane: " << *key.lane << "\n";
+  llvm::errs() << "Value: ";
+  value.print(llvm::errs());
+  llvm::errs() << "\n";
+  llvm::errs() << "Value type: " << value.getType() << "\n";
+  if (Operation *definingOp = value.getDefiningOp()) {
+    llvm::errs() << "Defining op:\n";
+    definingOp->print(llvm::errs());
+    llvm::errs() << "\n";
+  }
+}
+
 } // namespace

 MemEntry* PimMemory::gatherMemEntry(mlir::Value value, std::optional<unsigned> lane) {
@@ -124,20 +170,30 @@ void PimMemory::allocateGatheredMemory() {

 void PimMemory::allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memEntry, MemoryReportKind reportKind) {
  memEntry.address = firstAvailableAddress;
-  assert(memEntry.address < (size_t) INT_MAX && "Address allocated bigger than 32bit");
-  firstAvailableAddress += memEntry.size;
-  // Alignment
-  if (size_t remainder = firstAvailableAddress % minAlignment)
-    firstAvailableAddress += minAlignment - remainder;
+  Operation *anchor = getDiagnosticAnchor(key.value);
+  auto checkedEnd = pim::checkedAdd(memEntry.address, memEntry.size, anchor, "local memory end");
+  FailureOr<size_t> checkedAlignedEnd = failure();
+  if (succeeded(checkedEnd))
+    checkedAlignedEnd = checkedAlignTo(*checkedEnd, minAlignment, anchor, "local memory alignment");
+  bool startFits = memEntry.address <= kPimAddressLimit;
+  bool endFits = succeeded(checkedEnd) && *checkedEnd <= kPimAddressLimit;
+  bool alignedEndFits = succeeded(checkedAlignedEnd) && *checkedAlignedEnd <= kPimAddressLimit;
+  if (!startFits || !endFits || !alignedEndFits) {
+    printMemoryOverflowDiagnostic(
+      key.value,
+      key,
+      memEntry.size,
+      firstAvailableAddress,
+      succeeded(checkedAlignedEnd) ? *checkedAlignedEnd : kPimAddressLimit);
+    llvm_unreachable("PIM local memory allocation overflow");
+  }
+  firstAvailableAddress = *checkedAlignedEnd;

  ownedMemEntriesMap[key] = memEntry;
  globalMemEntriesMap[key] = memEntry;

  switch (reportKind) {
-  case MemoryReportKind::Alloca:
-    ++reportRow.numAlloca;
-    reportRow.sizeAlloca += memEntry.size;
-    break;
+  case MemoryReportKind::Alloca: break;
  case MemoryReportKind::Global:
    ++reportRow.numGlobal;
    reportRow.sizeGlobal += memEntry.size;
@@ -147,6 +203,31 @@ void PimMemory::allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memE
  }
 }

+PhysicalSlotInfo PimMemory::allocatePhysicalSlot(size_t slotSize, const MemoryValueKey& key) {
+  PhysicalSlotInfo slot;
+  slot.id = nextPhysicalSlotId++;
+  slot.address = firstAvailableAddress;
+  slot.size = slotSize;
+
+  Operation *anchor = getDiagnosticAnchor(key.value);
+  auto checkedEnd = pim::checkedAdd(slot.address, slot.size, anchor, "local memory end");
+  FailureOr<size_t> checkedAlignedEnd = failure();
+  if (succeeded(checkedEnd))
+    checkedAlignedEnd = checkedAlignTo(*checkedEnd, minAlignment, anchor, "local memory alignment");
+  bool startFits = slot.address <= kPimAddressLimit;
+  bool endFits = succeeded(checkedEnd) && *checkedEnd <= kPimAddressLimit;
+  bool alignedEndFits = succeeded(checkedAlignedEnd) && *checkedAlignedEnd <= kPimAddressLimit;
+  if (!startFits || !endFits || !alignedEndFits) {
+    printMemoryOverflowDiagnostic(
+      key.value, key, slot.size, firstAvailableAddress, succeeded(checkedAlignedEnd) ? *checkedAlignedEnd : kPimAddressLimit);
+    llvm_unreachable("PIM local memory allocation overflow");
+  }
+
+  firstAvailableAddress = *checkedAlignedEnd;
+  localPhysicalSlots.push_back(slot);
+  return slot;
+}
+
 void PimMemory::allocateHost(ModuleOp moduleOp, func::FuncOp funcOp) {
  SmallDenseMap<memref::GlobalOp, mlir::Value, 8> globalConstants;
  SmallVector<std::pair<mlir::Value, mlir::Value>, 16> globalAliases;
@@ -186,9 +267,71 @@ void PimMemory::allocateHost(ModuleOp moduleOp, func::FuncOp funcOp) {
 }

 void PimMemory::allocateCore(Operation* op, std::optional<unsigned> lane) {
-  op->walk([&](memref::AllocOp allocOp) { gatherMemEntry(allocOp, lane); });
+  auto intervals = buildLocalAllocIntervals(op, lane);
+  SmallVector<PlannedPhysicalSlot> plannedSlots = planPhysicalSlots(intervals);

-  allocateGatheredMemory();
+  SmallVector<size_t> slotOrder(plannedSlots.size());
+  std::iota(slotOrder.begin(), slotOrder.end(), 0);
+  llvm::stable_sort(slotOrder, [&](size_t lhsIndex, size_t rhsIndex) {
+    const PlannedPhysicalSlot &lhs = plannedSlots[lhsIndex];
+    const PlannedPhysicalSlot &rhs = plannedSlots[rhsIndex];
+    if (lhs.requiredSize != rhs.requiredSize)
+      return lhs.requiredSize > rhs.requiredSize;
+    return lhs.id < rhs.id;
+  });
+
+  SmallVector<bool, 16> usedExistingSlots(localPhysicalSlots.size(), false);
+  for (size_t slotIndex : slotOrder) {
+    PlannedPhysicalSlot &slot = plannedSlots[slotIndex];
+    size_t bestExistingIndex = std::numeric_limits<size_t>::max();
+    auto bestKey = std::tuple<size_t, size_t, size_t>(
+      std::numeric_limits<size_t>::max(), std::numeric_limits<size_t>::max(), std::numeric_limits<size_t>::max());
+
+    for (size_t existingIndex = 0; existingIndex < localPhysicalSlots.size(); ++existingIndex) {
+      if (usedExistingSlots[existingIndex])
+        continue;
+      const PhysicalSlotInfo &existingSlot = localPhysicalSlots[existingIndex];
+      if (existingSlot.size < slot.requiredSize)
+        continue;
+      auto candidateKey = std::tuple<size_t, size_t, size_t>(
+        existingSlot.size - slot.requiredSize, existingSlot.size, existingSlot.id);
+      if (candidateKey < bestKey) {
+        bestKey = candidateKey;
+        bestExistingIndex = existingIndex;
+      }
+    }
+
+    if (bestExistingIndex != std::numeric_limits<size_t>::max()) {
+      const PhysicalSlotInfo &existingSlot = localPhysicalSlots[bestExistingIndex];
+      slot.id = existingSlot.id;
+      slot.address = existingSlot.address;
+      slot.size = existingSlot.size;
+      usedExistingSlots[bestExistingIndex] = true;
+    }
+    else {
+      PhysicalSlotInfo newSlot = allocatePhysicalSlot(slot.requiredSize, intervals[slot.intervalIndices.front()].key);
+      slot.id = newSlot.id;
+      slot.address = newSlot.address;
+      slot.size = newSlot.size;
+      usedExistingSlots.push_back(true);
+    }
+
+    for (size_t intervalIndex : slot.intervalIndices) {
+      LocalAllocInterval &interval = intervals[intervalIndex];
+      interval.physicalSlotId = slot.id;
+      interval.assignedAddress = slot.address;
+      interval.physicalSlotSize = slot.size;
+      MemEntry memEntry {slot.address, interval.size};
+      ownedMemEntriesMap[interval.key] = memEntry;
+      globalMemEntriesMap[interval.key] = memEntry;
+    }
+  }
+
+  if (pimMemoryReport != PimMemoryReportNone) {
+    MemoryPlanArtifacts artifacts =
+      buildMemoryPlanArtifacts(op, lane, intervals, plannedSlots, kPimAddressLimit, pimMemoryReport);
+    livenessArtifacts.textReport += artifacts.textReport;
+  }
 }

 static void printHostMemoryReportRow(raw_ostream& os, const MemoryReportRow& row) {
@@ -228,7 +371,14 @@ static MemoryReportRow addMemoryReportRows(const MemoryReportRow& lhs, const Mem
  return result;
 }

-MemoryReportRow PimMemory::getReportRow() const { return reportRow; }
+MemoryReportRow PimMemory::getReportRow() const {
+  MemoryReportRow row = reportRow;
+  row.numAlloca = localPhysicalSlots.size();
+  row.sizeAlloca = 0;
+  for (const PhysicalSlotInfo &slot : localPhysicalSlots)
+    row.sizeAlloca += slot.size;
+  return row;
+}

 void PimMemory::remove(mlir::Value val) {
  for (auto it = ownedMemEntriesMap.begin(); it != ownedMemEntriesMap.end();)
@@ -847,6 +997,7 @@ struct CoreEmissionResult {
  OnnxMlirCompilerErrorCodes status = CompilerSuccess;
  MemoryReportRow reportRow;
  llvm::SmallVector<ResolvedWeightView, 8> usedWeights;
+  MemoryPlanArtifacts livenessArtifacts;
 };

 template <typename MapTy>
@@ -1319,6 +1470,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
      assert(processedOperations > 0);
      result.reportRow = deviceMemory.getReportRow();
      result.usedWeights = std::move(usedWeights);
+      result.livenessArtifacts = deviceMemory.getLivenessArtifacts();
    }
    else {
      auto coreBatchOp = cast<pim::PimCoreBatchOp>(job.coreLikeOp);
@@ -1349,6 +1501,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::

      result.reportRow = deviceMemory.getReportRow();
      result.usedWeights = std::move(usedWeights);
+      result.livenessArtifacts = deviceMemory.getLivenessArtifacts();
    }

    pim_binary::patchInstructionCount(coreBinaryStream, coreCodeGen.getEmittedInstructionCount());
@@ -1382,6 +1535,18 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
  auto weightEmission = createAndPopulateWeightFolder(weightRequests, outputDirPath);
  memory.setTotalWeightBytes(weightEmission.totalWeightBytes);
  auto& mapCoreWeightToFileName = weightEmission.mapCoreWeightToFileName;
+  if (std::string reportsRoot = getOutputDir(); !reportsRoot.empty()) {
+    std::string reportsDir = reportsRoot + "/reports";
+    sys::fs::remove(reportsDir + "/pim_memory_liveness_report.txt");
+    sys::fs::remove(reportsDir + "/pim_memory_liveness_report.json");
+    sys::fs::remove(reportsDir + "/pim_memory_liveness_timeline.dot");
+  }
+  std::fstream livenessReportFile;
+  std::unique_ptr<llvm::raw_os_ostream> livenessReportOs;
+  if (pimMemoryReport != PimMemoryReportNone) {
+    livenessReportFile = openReportFileWithExtension("pim_memory_liveness_report", "txt");
+    livenessReportOs = std::make_unique<llvm::raw_os_ostream>(livenessReportFile);
+  }

  for (size_t jobIndex = 0; jobIndex < jobs.size(); ++jobIndex) {
    const CoreEmissionJob& job = jobs[jobIndex];
@@ -1393,6 +1558,8 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
        return err;
      xbarsPerArrayGroup["core" + std::to_string(job.emittedCoreId)] = std::move(xbarsPerGroup);
      memory.recordCoreReport(job.emittedCoreId, result.reportRow);
+      if (livenessReportFile.is_open())
+        *livenessReportOs << "Core " << job.emittedCoreId << ":\n" << result.livenessArtifacts.textReport;
      continue;
    }
  }
@@ -1421,10 +1588,18 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
                             batchPerCoreRow.value_or(MemoryReportRow {}),
                             batchRow.numAlloca,
                             batchRow.sizeAlloca);
+    if (livenessReportFile.is_open())
+      for (size_t jobIndex : group)
+        *livenessReportOs << "Batch " << batchReportId << " core " << jobs[jobIndex].emittedCoreId << ":\n"
+                          << jobResults[jobIndex].livenessArtifacts.textReport;
  }

  maxCoreId = nextEmittedCoreId == 0 ? 0 : nextEmittedCoreId - 1;

+  if (livenessReportFile.is_open()) {
+    livenessReportOs->flush();
+    livenessReportFile.close();
+  }
  memory.flushReport();
  return writeConfigJson(funcOp, memory, maxCoreId, std::move(xbarsPerArrayGroup), outputDirPath);
 }