blazingly fasterer

2026-07-21 15:43:35 +02:00
parent a893d23a74
commit 3e468b58c8
37 changed files with 1119 additions and 933 deletions
@@ -15,6 +15,7 @@
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/FileSystem.h"
+#include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/JSON.h"
 #include "llvm/Support/raw_ostream.h"

@@ -31,8 +32,8 @@

 #include "Common/IR/CompactAsmUtils.hpp"
 #include "Common/PimCommon.hpp"
-#include "Common/Support/Diagnostics.hpp"
 #include "Common/Support/CheckedArithmetic.hpp"
+#include "Common/Support/Diagnostics.hpp"
 #include "Common/Support/ReportUtils.hpp"
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
@@ -43,7 +44,6 @@
 #include "src/Accelerators/PIM/Compiler/PimCodeGen.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCoreProgram.hpp"
-#include "src/Accelerators/PIM/Compiler/PimMemoryLiveness.hpp"
 #include "src/Accelerators/PIM/Compiler/PimWeightEmitter.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"

@@ -110,8 +110,6 @@ static Operation* getDiagnosticAnchor(mlir::Value value) {
 // 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;
@@ -129,7 +127,7 @@ static void printMemoryOverflowDiagnostic(const MemoryValueKey& key,
  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";
+  llvm::errs() << "Address limit: " << kPimLocalMemoryAddressLimit << " (signed int32_t immediate range)\n";
  if (key.lane)
    llvm::errs() << "Lane: " << *key.lane << "\n";
  llvm::errs() << "Value: ";
@@ -152,10 +150,13 @@ size_t PimMemory::allocateAddress(size_t size, const MemoryValueKey& key) {
  FailureOr<size_t> checkedAlignedEnd = failure();
  if (succeeded(checkedEnd))
    checkedAlignedEnd = checkedAlignTo(*checkedEnd, minAlignment, anchor, "local memory alignment");
-  if (address > kPimAddressLimit || failed(checkedEnd) || *checkedEnd > kPimAddressLimit
-      || failed(checkedAlignedEnd) || *checkedAlignedEnd > kPimAddressLimit) {
+  if (address > kPimLocalMemoryAddressLimit || failed(checkedEnd) || *checkedEnd > kPimLocalMemoryAddressLimit
+      || failed(checkedAlignedEnd) || *checkedAlignedEnd > kPimLocalMemoryAddressLimit) {
    printMemoryOverflowDiagnostic(
-      key, size, firstAvailableAddress, succeeded(checkedAlignedEnd) ? *checkedAlignedEnd : kPimAddressLimit);
+      key,
+      size,
+      firstAvailableAddress,
+      succeeded(checkedAlignedEnd) ? *checkedAlignedEnd : kPimLocalMemoryAddressLimit);
    llvm_unreachable("PIM local memory allocation overflow");
  }
  firstAvailableAddress = *checkedAlignedEnd;
@@ -202,15 +203,6 @@ void PimMemory::allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memE
  }
 }

-PhysicalSlotInfo PimMemory::allocatePhysicalSlot(size_t slotSize, const MemoryValueKey& key) {
-  PhysicalSlotInfo slot;
-  slot.id = nextPhysicalSlotId++;
-  slot.address = allocateAddress(slotSize, key);
-  slot.size = slotSize;
-  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;
@@ -249,71 +241,17 @@ void PimMemory::allocateHost(ModuleOp moduleOp, func::FuncOp funcOp) {
    globalMemEntriesMap[getMemoryValueKey(alias)] = getMemEntry(getMemoryValueKey(original));
 }

-void PimMemory::allocateCore(Operation* op, std::optional<unsigned> lane) {
-  auto intervals = buildLocalAllocIntervals(op, lane, pimMemoryReport == PimMemoryReportFull);
-  SmallVector<PlannedPhysicalSlot> plannedSlots = planPhysicalSlots(intervals);
-
-  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;
-    }
+void PimMemory::allocateCore(const CompiledCoreMemoryPlan& plan, std::optional<unsigned> lane) {
+  if (localPhysicalSlots.empty()) {
+    llvm::append_range(localPhysicalSlots, plan.slots);
+    reportRow.logicalAllocaBytes = plan.logicalBytes;
+    reportRow.fallbackIntervals = plan.fallbackIntervals;
+    reportRow.nestedSingleUseIntervals = plan.nestedSingleUseIntervals;
  }
-
-  if (pimMemoryReport != PimMemoryReportNone) {
-    MemoryPlanArtifacts artifacts =
-      buildMemoryPlanArtifacts(op, lane, intervals, plannedSlots, kPimAddressLimit, pimMemoryReport);
-    livenessArtifacts += artifacts;
+  for (const CompiledLocalMemoryEntry& entry : plan.entries) {
+    MemoryValueKey key = getMemoryValueKey(entry.value, lane);
+    ownedMemEntriesMap[key] = entry.memory;
+    globalMemEntriesMap[key] = entry.memory;
  }
 }

@@ -475,13 +413,41 @@ void PimAcceleratorMemory::flushReport() {
  uint64_t totalGlobalMemory = hostReportRow.has_value() ? hostReportRow->sizeGlobal : 0;
  uint64_t totalWeightsMemory = totalWeightBytes;
  uint64_t totalCoresMemory = 0;
-  for (const MemoryReportEntry& entry : reportEntries)
+  uint64_t totalLogicalCoreMemory = 0;
+  uint64_t totalFallbackIntervals = 0;
+  uint64_t totalNestedSingleUseIntervals = 0;
+  uint64_t maximumCorePeak = 0;
+  std::string worstEntry = "<none>";
+  for (const MemoryReportEntry& entry : reportEntries) {
    totalCoresMemory += entry.totalAllocaBytes;
+    uint64_t factor = entry.kind == MemoryReportEntry::Kind::Batch ? entry.coreIds.size() : 1;
+    totalLogicalCoreMemory += entry.row.logicalAllocaBytes * factor;
+    totalFallbackIntervals += entry.row.fallbackIntervals * factor;
+    totalNestedSingleUseIntervals += entry.row.nestedSingleUseIntervals * factor;
+    if (entry.row.sizeAlloca <= maximumCorePeak)
+      continue;
+    maximumCorePeak = entry.row.sizeAlloca;
+    worstEntry = entry.kind == MemoryReportEntry::Kind::Batch ? "Batch " + std::to_string(entry.id)
+                                                              : "Core " + std::to_string(entry.coreIds.front());
+  }
+  uint64_t savedCoreMemory =
+    totalLogicalCoreMemory >= totalCoresMemory ? totalLogicalCoreMemory - totalCoresMemory : 0;
+  double savedPercent = totalLogicalCoreMemory == 0
+                        ? 0.0
+                        : 100.0 * static_cast<double>(savedCoreMemory) / static_cast<double>(totalLogicalCoreMemory);

-  llvm::SmallVector<ReportField, 3> totalFields = {
-    {"Global memory",  formatReportMemory(totalGlobalMemory) },
-    {"Weights memory", formatReportMemory(totalWeightsMemory)},
-    {"Cores memory",   formatReportMemory(totalCoresMemory)  }
+  llvm::SmallVector<ReportField, 10> totalFields = {
+    {"Global memory",               formatReportMemory(totalGlobalMemory)               },
+    {"Weights memory",              formatReportMemory(totalWeightsMemory)              },
+    {"Logical core allocations",    formatReportMemory(totalLogicalCoreMemory)           },
+    {"Physical core memory",        formatReportMemory(totalCoresMemory)                 },
+    {"Cores memory",                formatReportMemory(totalCoresMemory)                 },
+    {"Saved core memory",           formatReportMemory(savedCoreMemory)                  },
+    {"Saved core memory percent",   formatv("{0:F2}%", savedPercent).str()               },
+    {"Maximum core peak",           formatReportMemory(maximumCorePeak)                  },
+    {"Worst core/batch",            worstEntry                                           },
+    {"Fallback intervals",          std::to_string(totalFallbackIntervals)               },
+    {"Nested single-use intervals", std::to_string(totalNestedSingleUseIntervals)        }
  };
  printReportTotalsBlock(os, totalFields);

@@ -552,19 +518,14 @@ void PimCodeGen::emitInstruction(const pim_binary::InstructionRecord& instructio

 void PimCodeGen::updateScalarRegisterCache(const pim_binary::InstructionRecord& instruction) const {
  switch (instruction.opcode) {
-  case pim_binary::Opcode::sldi:
-    scalarRegisterValues[instruction.rd] = instruction.r2OrImm;
-    break;
+  case pim_binary::Opcode::sldi:  scalarRegisterValues[instruction.rd] = instruction.r2OrImm; break;
  case pim_binary::Opcode::sld:
  case pim_binary::Opcode::sadd:
  case pim_binary::Opcode::ssub:
  case pim_binary::Opcode::smul:
  case pim_binary::Opcode::saddi:
-  case pim_binary::Opcode::smuli:
-    scalarRegisterValues[instruction.rd].reset();
-    break;
-  default:
-    break;
+  case pim_binary::Opcode::smuli: scalarRegisterValues[instruction.rd].reset(); break;
+  default:                        break;
  }
 }

@@ -580,8 +541,8 @@ void PimCodeGen::genSetRegisterImmediate(uint8_t registerNumber, int32_t immedia
 }

 void PimCodeGen::genSetRegisterImmediateUnsigned(size_t registerNumber, size_t immediate) const {
-  genSetRegisterImmediate(
-    pim::checkedU8OrCrash(registerNumber, "register number"), pim::checkedI32OrCrash(immediate, "register immediate"));
+  genSetRegisterImmediate(pim::checkedU8OrCrash(registerNumber, "register number"),
+                          pim::checkedI32OrCrash(immediate, "register immediate"));
 }

 void PimCodeGen::setupRd(size_t rdAddress, size_t rdOffset) const {
@@ -729,8 +690,7 @@ void PimCodeGen::codeGenConcatOp(pim::PimConcatOp concatOp, const StaticValueKno
    for (size_t outerIndex = 0; outerIndex < outerCount; ++outerIndex) {
      size_t dstOffset = (outerIndex * outputConcatDim + concatOffset) * innerCount * elementSize;
      size_t srcOffset = outerIndex * inputConcatDim * innerCount * elementSize;
-      emitMemCopyOp(
-        pim_binary::Opcode::lmv, outputAddr, dstOffset, inputAddr, srcOffset, blockSizeInBytes, "len");
+      emitMemCopyOp(pim_binary::Opcode::lmv, outputAddr, dstOffset, inputAddr, srcOffset, blockSizeInBytes, "len");
    }

    concatOffset += inputConcatDim;
@@ -788,10 +748,11 @@ void PimCodeGen::codeGenTransposeOp(const CompiledTransposePlan& plan, const Sta
  size_t maxElementOffset = plan.totalBytes == 0 ? 0 : plan.totalBytes - plan.elementBytes;
  int32_t maxSourceAddress = pim::checkedI32OrCrash(
    pim::checkedAddOrCrash(srcAddr, maxElementOffset, "transpose source address"), "transpose source address");
-  int32_t maxDestinationAddress = pim::checkedI32OrCrash(
-    pim::checkedAddOrCrash(dstAddr, maxElementOffset, "transpose destination address"), "transpose destination address");
-  (void)maxSourceAddress;
-  (void)maxDestinationAddress;
+  int32_t maxDestinationAddress =
+    pim::checkedI32OrCrash(pim::checkedAddOrCrash(dstAddr, maxElementOffset, "transpose destination address"),
+                           "transpose destination address");
+  (void) maxSourceAddress;
+  (void) maxDestinationAddress;

  pim_binary::InstructionRecord copyInstruction;
  copyInstruction.opcode = pim_binary::Opcode::lmv;
@@ -871,6 +832,60 @@ static SmallVector<Operation*> collectTopLevelCoreLikeOps(func::FuncOp funcOp) {
  return coreLikeOps;
 }

+static FailureOr<CompiledCoreMemoryPlan> compileCoreMemoryPlan(Operation* coreLikeOp) {
+  CompiledCoreMemoryPlan plan;
+  auto fallbackAttr = coreLikeOp->getAttrOfType<IntegerAttr>(kLocalMemoryFallbackCountAttrName);
+  auto nestedSingleUseAttr = coreLikeOp->getAttrOfType<IntegerAttr>(kLocalMemoryNestedSingleUseCountAttrName);
+  if (!fallbackAttr || !nestedSingleUseAttr || fallbackAttr.getInt() < 0 || nestedSingleUseAttr.getInt() < 0) {
+    coreLikeOp->emitError("requires complete PIM local-memory planning summary attributes before codegen");
+    return failure();
+  }
+  plan.fallbackIntervals = static_cast<uint64_t>(fallbackAttr.getInt());
+  plan.nestedSingleUseIntervals = static_cast<uint64_t>(nestedSingleUseAttr.getInt());
+  SmallDenseMap<size_t, size_t, 32> slotIndices;
+  bool hasFailure = false;
+  coreLikeOp->walk([&](memref::AllocOp allocOp) {
+    if (hasFailure)
+      return;
+    auto addressAttr = allocOp->getAttrOfType<IntegerAttr>(kLocalMemoryAddressAttrName);
+    auto slotAttr = allocOp->getAttrOfType<IntegerAttr>(kLocalMemorySlotAttrName);
+    auto slotSizeAttr = allocOp->getAttrOfType<IntegerAttr>(kLocalMemorySlotSizeAttrName);
+    if (!addressAttr || !slotAttr || !slotSizeAttr || addressAttr.getInt() < 0 || slotAttr.getInt() < 0
+        || slotSizeAttr.getInt() < 0) {
+      allocOp.emitOpError("requires a complete non-negative PIM local-memory plan before codegen");
+      hasFailure = true;
+      return;
+    }
+
+    auto checkedSize = pim::getCheckedShapedTypeSizeInBytes(
+      cast<ShapedType>(allocOp.getType()), allocOp, "planned local allocation byte size");
+    if (failed(checkedSize)) {
+      hasFailure = true;
+      return;
+    }
+    size_t address = static_cast<size_t>(addressAttr.getInt());
+    size_t slotId = static_cast<size_t>(slotAttr.getInt());
+    size_t slotSize = static_cast<size_t>(slotSizeAttr.getInt());
+    plan.entries.push_back({allocOp.getResult(), {address, static_cast<size_t>(*checkedSize)}});
+    plan.logicalBytes += *checkedSize;
+
+    auto [slotIt, inserted] = slotIndices.try_emplace(slotId, plan.slots.size());
+    if (inserted) {
+      plan.slots.push_back({slotId, 0, slotSize});
+      return;
+    }
+    const PhysicalSlotInfo& existing = plan.slots[slotIt->second];
+    if (existing.size != slotSize) {
+      allocOp.emitOpError("has a PIM local-memory arena inconsistent with another allocation");
+      hasFailure = true;
+    }
+  });
+  if (hasFailure)
+    return failure();
+  llvm::sort(plan.slots, [](const PhysicalSlotInfo& lhs, const PhysicalSlotInfo& rhs) { return lhs.id < rhs.id; });
+  return plan;
+}
+
 struct CoreEmissionResult {
  static constexpr size_t kMaxStoredCodegenDiagnostics = 8;

@@ -882,10 +897,8 @@ struct CoreEmissionResult {
  OnnxMlirCompilerErrorCodes status = CompilerSuccess;
  MemoryReportRow reportRow;
  llvm::SmallVector<ResolvedWeightView, 8> usedWeights;
-  MemoryPlanArtifacts livenessArtifacts;
  llvm::SmallVector<DiagnosticRecord, kMaxStoredCodegenDiagnostics> diagnostics;
  size_t diagnosticCount = 0;
-
  void recordDiagnostic(Operation* op, StringRef message) {
    ++diagnosticCount;
    if (diagnostics.size() < kMaxStoredCodegenDiagnostics)
@@ -1012,13 +1025,21 @@ static LogicalResult executeCompiledCorePlan(
    }

    auto emitBinary = [&](auto op, pim_binary::Opcode opcode) {
-      coreCodeGen.emitBinaryVectorOp(opcode, op.getOutputBuffer(), op.getLhs(), op.getRhs(),
-                                     getVectorByteSizeOrCrash(cast<ShapedType>(op.getLhs().getType())), knowledge);
+      coreCodeGen.emitBinaryVectorOp(opcode,
+                                     op.getOutputBuffer(),
+                                     op.getLhs(),
+                                     op.getRhs(),
+                                     getVectorByteSizeOrCrash(cast<ShapedType>(op.getLhs().getType())),
+                                     knowledge);
    };
    auto emitUnary = [&](auto op, pim_binary::Opcode opcode, int32_t r2OrImm, int32_t generic1) {
-      coreCodeGen.emitUnaryVectorOp(opcode, op.getOutputBuffer(), op.getInput(),
+      coreCodeGen.emitUnaryVectorOp(opcode,
+                                    op.getOutputBuffer(),
+                                    op.getInput(),
                                    getVectorByteSizeOrCrash(cast<ShapedType>(op.getInput().getType())),
-                                    knowledge, r2OrImm, generic1);
+                                    knowledge,
+                                    r2OrImm,
+                                    generic1);
    };

    switch (node.opKind) {
@@ -1038,20 +1059,19 @@ static LogicalResult executeCompiledCorePlan(
      else
        return failure();
      break;
-    case CompiledCoreOpKind::Transpose:
-      coreCodeGen.codeGenTransposeOp(*node.transposePlan, knowledge);
-      break;
-    case CompiledCoreOpKind::VVAdd:  emitBinary(cast<pim::PimVVAddOp>(node.op), pim_binary::Opcode::vvadd); break;
-    case CompiledCoreOpKind::VVSub:  emitBinary(cast<pim::PimVVSubOp>(node.op), pim_binary::Opcode::vvsub); break;
-    case CompiledCoreOpKind::VVMul:  emitBinary(cast<pim::PimVVMulOp>(node.op), pim_binary::Opcode::vvmul); break;
-    case CompiledCoreOpKind::VVMax:  emitBinary(cast<pim::PimVVMaxOp>(node.op), pim_binary::Opcode::vvmax); break;
-    case CompiledCoreOpKind::VVDMul: emitBinary(cast<pim::PimVVDMulOp>(node.op), pim_binary::Opcode::vvdmul); break;
-    case CompiledCoreOpKind::VAvg:   emitUnary(cast<pim::PimVAvgOp>(node.op), pim_binary::Opcode::vavg, 1, 1); break;
-    case CompiledCoreOpKind::VRelu:  emitUnary(cast<pim::PimVReluOp>(node.op), pim_binary::Opcode::vrelu, 0, 0); break;
-    case CompiledCoreOpKind::VTanh:  emitUnary(cast<pim::PimVTanhOp>(node.op), pim_binary::Opcode::vtanh, 0, 0); break;
-    case CompiledCoreOpKind::VSigm:  emitUnary(cast<pim::PimVSigmOp>(node.op), pim_binary::Opcode::vsigm, 0, 0); break;
+    case CompiledCoreOpKind::Transpose: coreCodeGen.codeGenTransposeOp(*node.transposePlan, knowledge); break;
+    case CompiledCoreOpKind::VVAdd:     emitBinary(cast<pim::PimVVAddOp>(node.op), pim_binary::Opcode::vvadd); break;
+    case CompiledCoreOpKind::VVSub:     emitBinary(cast<pim::PimVVSubOp>(node.op), pim_binary::Opcode::vvsub); break;
+    case CompiledCoreOpKind::VVMul:     emitBinary(cast<pim::PimVVMulOp>(node.op), pim_binary::Opcode::vvmul); break;
+    case CompiledCoreOpKind::VVMax:     emitBinary(cast<pim::PimVVMaxOp>(node.op), pim_binary::Opcode::vvmax); break;
+    case CompiledCoreOpKind::VVDMul:    emitBinary(cast<pim::PimVVDMulOp>(node.op), pim_binary::Opcode::vvdmul); break;
+    case CompiledCoreOpKind::VAvg:      emitUnary(cast<pim::PimVAvgOp>(node.op), pim_binary::Opcode::vavg, 1, 1); break;
+    case CompiledCoreOpKind::VRelu:     emitUnary(cast<pim::PimVReluOp>(node.op), pim_binary::Opcode::vrelu, 0, 0); break;
+    case CompiledCoreOpKind::VTanh:     emitUnary(cast<pim::PimVTanhOp>(node.op), pim_binary::Opcode::vtanh, 0, 0); break;
+    case CompiledCoreOpKind::VSigm:     emitUnary(cast<pim::PimVSigmOp>(node.op), pim_binary::Opcode::vsigm, 0, 0); break;
    case CompiledCoreOpKind::VSoftmax:
-      emitUnary(cast<pim::PimVSoftmaxOp>(node.op), pim_binary::Opcode::vsoftmax, 0, 0); break;
+      emitUnary(cast<pim::PimVSoftmaxOp>(node.op), pim_binary::Opcode::vsoftmax, 0, 0);
+      break;
    }
  }
  return success();
@@ -1103,8 +1123,7 @@ static void aliasMaterializedHostGlobals(CoreLikeOpTy coreLikeOp,
 }

 static OnnxMlirCompilerErrorCodes emitEmptyCoreArtifacts(StringRef outputDirPath, size_t emittedCoreId) {
-  std::string outputCorePath =
-    (outputDirPath + "/core_" + std::to_string(emittedCoreId) + ".pim").str();
+  std::string outputCorePath = (outputDirPath + "/core_" + std::to_string(emittedCoreId) + ".pim").str();
  std::error_code errorCode;
  raw_fd_ostream coreBinaryStream(outputCorePath, errorCode, sys::fs::OF_None);
  if (errorCode) {
@@ -1128,8 +1147,7 @@ static OnnxMlirCompilerErrorCodes emitEmptyCoreArtifacts(StringRef outputDirPath
  if (!pimEmitJson.getValue())
    return CompilerSuccess;

-  std::string outputCoreJsonPath =
-    (outputDirPath + "/core_" + std::to_string(emittedCoreId) + ".json").str();
+  std::string outputCoreJsonPath = (outputDirPath + "/core_" + std::to_string(emittedCoreId) + ".json").str();
  errorCode = std::error_code();
  raw_fd_ostream coreJsonStream(outputCoreJsonPath, errorCode);
  if (errorCode) {
@@ -1169,18 +1187,27 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
  SmallDenseMap<memref::GlobalOp, MemEntry, 16> materializedHostGlobals =
    collectMaterializedHostGlobals(moduleOp, funcOp, memory);
  llvm::DenseMap<Operation*, std::unique_ptr<CompiledCoreProgram>> compiledPrograms;
+  llvm::DenseMap<Operation*, std::unique_ptr<CompiledCoreMemoryPlan>> memoryPlans;
  for (Operation* op : coreLikeOps) {
    auto program = std::make_unique<CompiledCoreProgram>();
    if (failed(compileCoreProgram(op, *program)))
      return CompilerFailure;
    compiledPrograms.try_emplace(op, std::move(program));
+    auto memoryPlan = compileCoreMemoryPlan(op);
+    if (failed(memoryPlan))
+      return CompilerFailure;
+    memoryPlans.try_emplace(op, std::make_unique<CompiledCoreMemoryPlan>(std::move(*memoryPlan)));
  }
-
  auto getCompiledProgram = [&](Operation* op) {
    auto it = compiledPrograms.find(op);
    assert(it != compiledPrograms.end() && "missing compiled PIM core program");
    return it->second.get();
  };
+  auto getMemoryPlan = [&](Operation* op) {
+    auto it = memoryPlans.find(op);
+    assert(it != memoryPlans.end() && "missing PIM core memory plan");
+    return it->second.get();
+  };

  llvm::DenseMap<size_t, size_t> emittedCoreIds;
  size_t nextEmittedCoreId = 0;
@@ -1210,6 +1237,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
      CoreEmissionJob job;
      job.coreLikeOp = coreOp;
      job.program = getCompiledProgram(op);
+      job.memoryPlan = getMemoryPlan(op);
      job.emittedCoreId = emittedCoreIds.lookup(originalCoreId);
      jobs.push_back(std::move(job));
      continue;
@@ -1229,6 +1257,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
      CoreEmissionJob job;
      job.coreLikeOp = coreBatchOp;
      job.program = getCompiledProgram(op);
+      job.memoryPlan = getMemoryPlan(op);
      job.emittedCoreId = emittedCoreIds.lookup(originalCoreId);
      job.lanes = lanesByCoreId.lookup(originalCoreId);
      job.batchReportId = nextBatchReportId;
@@ -1332,17 +1361,16 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
    if (auto coreOp = dyn_cast<pim::PimCoreOp>(job.coreLikeOp)) {
      aliasMaterializedHostGlobals(coreOp, moduleOp, materializedHostGlobals, jobMemory);
      auto& deviceMemory = jobMemory.getOrCreateDeviceMem(job.emittedCoreId);
-      deviceMemory.allocateCore(coreOp);
+      deviceMemory.allocateCore(*job.memoryPlan);

      StaticValueKnowledge knowledge = seedCoreCodegenKnowledge(coreOp);
      if (failed(executeCompiledCoreProgram(*job.program, coreCodeGen, knowledge, resolveWeightSlot))) {
-        (void)finalizeInstructions();
+        (void) finalizeInstructions();
        result.status = CompilerFailure;
        return result;
      }
      result.reportRow = deviceMemory.getReportRow();
      result.usedWeights = std::move(usedWeights);
-      result.livenessArtifacts = deviceMemory.getLivenessArtifacts();
    }
    else {
      auto coreBatchOp = cast<pim::PimCoreBatchOp>(job.coreLikeOp);
@@ -1352,10 +1380,10 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
      for (unsigned lane : job.lanes) {
        StaticValueKnowledge knowledge = seedCoreBatchCodegenKnowledge(coreBatchOp, lane);

-        deviceMemory.allocateCore(coreBatchOp, lane);
+        deviceMemory.allocateCore(*job.memoryPlan, lane);
        coreCodeGen.setBatchLane(lane);
        if (failed(executeCompiledCoreProgram(*job.program, coreCodeGen, knowledge, resolveWeightSlot))) {
-          (void)finalizeInstructions();
+          (void) finalizeInstructions();
          result.status = CompilerFailure;
          return result;
        }
@@ -1363,7 +1391,6 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::

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

    if (!finalizeInstructions()) {
@@ -1388,9 +1415,8 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
  for (const CoreEmissionResult& result : jobResults) {
    if (!summaryAnchor && !result.diagnostics.empty())
      summaryAnchor = result.diagnostics.front().op;
-    for (const CoreEmissionResult::DiagnosticRecord& diagnostic : result.diagnostics) {
+    for (const CoreEmissionResult::DiagnosticRecord& diagnostic : result.diagnostics)
      diagnostics.report(diagnostic.op, [&](Operation* op) { op->emitError() << diagnostic.message; });
-    }
    size_t unreportedCount = result.diagnosticCount - result.diagnostics.size();
    diagnostics.noteFailures(static_cast<int64_t>(unreportedCount));
  }
@@ -1420,19 +1446,6 @@ 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];
    const CoreEmissionResult& result = jobResults[jobIndex];
@@ -1443,8 +1456,6 @@ 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;
      continue;
    }
  }
@@ -1473,18 +1484,10 @@ 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;
  }

  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);
 }