finally fast googlenet with correct latency artifacts for fair comparison

src/PIM/Compiler/PimCodeGen.cpp

@@ -92,11 +92,12 @@ static MemoryReportKind classifyMemoryReportKind(mlir::Value value) {
   return MemoryReportKind::None;
 }
 
-static int32_t getVectorByteSizeOrCrash(ShapedType type) {
-  auto byteSize = pim::getCheckedShapedTypeSizeInBytes(type, UnknownLoc::get(type.getContext()), "vector byte size");
-  if (failed(byteSize))
-    llvm_unreachable("Failed to compute checked vector byte size");
-  return pim::checkedI32OrCrash(*byteSize, "vector byte size");
+static int32_t getVectorElementCountOrCrash(ShapedType type) {
+  return pim::checkedI32OrCrash(type.getNumElements(), "vector element count");
+}
+
+static int32_t getVectorElementBitwidthOrCrash(ShapedType type) {
+  return pim::checkedI32OrCrash(static_cast<uint64_t>(type.getElementTypeBitWidth()), "vector element bitwidth");
 }
 
 static Operation* getDiagnosticAnchor(mlir::Value value) {
@@ -165,8 +166,7 @@ size_t PimMemory::allocateAddress(size_t size, const MemoryValueKey& key) {
     checkedAlignedEnd = checkedAlignTo(*checkedEnd, minAlignment, anchor, "local memory alignment");
   if (address > kPimLocalMemoryAddressLimit || failed(checkedEnd) || *checkedEnd > kPimLocalMemoryAddressLimit
       || failed(checkedAlignedEnd) || *checkedAlignedEnd > kPimLocalMemoryAddressLimit) {
-    printMemoryOverflowDiagnostic(
-      key,
+    printMemoryOverflowDiagnostic(key,
       size,
       firstAvailableAddress,
       succeeded(checkedAlignedEnd) ? *checkedAlignedEnd : kPimLocalMemoryAddressLimit);
@@ -208,9 +208,7 @@ void PimMemory::allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memE
   switch (reportKind) {
   case MemoryReportKind::Alloca:
   case MemoryReportKind::Global:
-  case MemoryReportKind::Input:
-    ++reportRow.hostObjectCount;
-    break;
+  case MemoryReportKind::Input:  ++reportRow.hostObjectCount; break;
   case MemoryReportKind::None:  break;
   }
 }
@@ -259,8 +257,7 @@ void PimMemory::allocateCore(const CompiledCoreMemoryPlan& plan, std::optional<u
     reportRow.logicalLocalAllocationCount = plan.logicalAllocationCount;
     reportRow.logicalLocalBytes = plan.logicalBytes;
   }
-  else if (*localArenaSize != plan.arenaSize
-           || reportRow.logicalLocalAllocationCount != plan.logicalAllocationCount
+  else if (*localArenaSize != plan.arenaSize || reportRow.logicalLocalAllocationCount != plan.logicalAllocationCount
            || reportRow.logicalLocalBytes != plan.logicalBytes)
     llvm_unreachable("inconsistent PIM local-memory plan across core-batch lanes");
   for (const CompiledLocalMemoryEntry& entry : plan.entries) {
@@ -358,8 +355,8 @@ llvm::FailureOr<int64_t> PimAcceleratorMemory::getIndexValue(mlir::Value value,
 PimAcceleratorMemory::PimAcceleratorMemory()
 : hostMem(memEntriesMap), fileReport(openMemoryReport(pimMemoryReport == PimMemoryReportSummary)) {}
 
-PimAcceleratorMemory::PimAcceleratorMemory(
-  const llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32>& initialMemEntries, bool enableReport)
+PimAcceleratorMemory::PimAcceleratorMemory(const llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32>& initialMemEntries,
+                                           bool enableReport)
 : memEntriesMap(initialMemEntries),
   hostMem(memEntriesMap),
   fileReport(enableReport ? openMemoryReport(true) : std::fstream()) {}
@@ -367,10 +364,8 @@ PimAcceleratorMemory::PimAcceleratorMemory(
 void PimAcceleratorMemory::reportHost() { hostReportRow = hostMem.getReportRow(); }
 
 void PimAcceleratorMemory::recordCoreReport(size_t coreId, const MemoryReportRow& row) {
-  reportEntries.push_back({MemoryReportEntry::Kind::Core,
-                           coreId,
-                           {pim::checkedI32OrCrash(coreId, "memory report core id")},
-                           row});
+  reportEntries.push_back(
+    {MemoryReportEntry::Kind::Core, coreId, {pim::checkedI32OrCrash(coreId, "memory report core id")}, row});
 }
 
 void PimAcceleratorMemory::recordBatchReport(uint64_t batchId,
@@ -441,8 +436,7 @@ void PimAcceleratorMemory::flushReport() {
   os << "  Weights memory: " << formatReportMemory(totalWeightBytes) << "\n";
   os << "  Local memory before reuse: " << formatReportMemory(logicalBytes) << "\n";
   os << "  Local memory after reuse: " << formatReportMemory(physicalBytes) << "\n";
-  os << "  Saved local memory: " << formatReportMemory(savedBytes) << " ("
-     << formatv("{0:F1}%", savedPercent) << ")\n";
+  os << "  Saved local memory: " << formatReportMemory(savedBytes) << " (" << formatv("{0:F1}%", savedPercent) << ")\n";
   os << "  Largest core local memory: " << formatReportMemory(largest) << "\n";
   if (!groups.empty()) {
     os << "    ";
@@ -456,20 +450,15 @@ void PimAcceleratorMemory::flushReport() {
     printLabel(group);
     os << "\n";
     uint64_t groupSaved = group.row.logicalLocalBytes - group.row.physicalLocalBytes;
-    double groupPercent = group.row.logicalLocalBytes == 0
-                            ? 0.0
-                            : 100.0 * groupSaved / group.row.logicalLocalBytes;
+    double groupPercent = group.row.logicalLocalBytes == 0 ? 0.0 : 100.0 * groupSaved / group.row.logicalLocalBytes;
     if (group.coreIds.size() == 1) {
       os << "  Local memory: " << formatReportMemory(group.row.logicalLocalBytes) << " → "
-         << formatReportMemory(group.row.physicalLocalBytes) << " (" << formatv("{0:F1}% saved", groupPercent)
-         << ")\n";
+         << formatReportMemory(group.row.physicalLocalBytes) << " (" << formatv("{0:F1}% saved", groupPercent) << ")\n";
     }
     else {
       os << "  Per core: " << formatReportMemory(group.row.logicalLocalBytes) << " → "
-         << formatReportMemory(group.row.physicalLocalBytes) << " (" << formatv("{0:F1}% saved", groupPercent)
-         << ")\n";
-      os << "  Total after reuse: " << formatReportMemory(group.row.physicalLocalBytes * group.coreIds.size())
-         << "\n";
+         << formatReportMemory(group.row.physicalLocalBytes) << " (" << formatv("{0:F1}% saved", groupPercent) << ")\n";
+      os << "  Total after reuse: " << formatReportMemory(group.row.physicalLocalBytes * group.coreIds.size()) << "\n";
     }
   }
   if (groups.size() > kGroupLimit)
@@ -479,12 +468,6 @@ void PimAcceleratorMemory::flushReport() {
   fileReport.close();
 }
 
-size_t PimCodeGen::remapCoreId(size_t coreId) const {
-  auto it = emittedCoreIds.find(coreId);
-  assert(it != emittedCoreIds.end() && "Missing emitted core id remapping");
-  return it->second;
-}
-
 void PimCodeGen::emitInstruction(const pim_binary::InstructionRecord& instruction) const {
   if (failed(instructionWriter.append(instruction)))
     return;
@@ -493,6 +476,19 @@ void PimCodeGen::emitInstruction(const pim_binary::InstructionRecord& instructio
   updateScalarRegisterCache(instruction);
 }
 
+void PimCodeGen::ensureVectorBitwidth(int32_t inputBitwidth, int32_t outputBitwidth) const {
+  std::array<int32_t, 2> requested = {inputBitwidth, outputBitwidth};
+  if (vectorBitwidths == requested)
+    return;
+
+  pim_binary::InstructionRecord instruction;
+  instruction.opcode = pim_binary::Opcode::setbw;
+  instruction.generic1 = inputBitwidth;
+  instruction.generic2 = outputBitwidth;
+  emitInstruction(instruction);
+  vectorBitwidths = requested;
+}
+
 void PimCodeGen::updateScalarRegisterCache(const pim_binary::InstructionRecord& instruction) const {
   switch (instruction.opcode) {
   case pim_binary::Opcode::sldi:  scalarRegisterValues[instruction.rd] = instruction.r2OrImm; break;
@@ -563,7 +559,7 @@ void PimCodeGen::emitCommunicationOp(pim_binary::Opcode opcode, size_t bufferAdd
   pim_binary::InstructionRecord instruction;
   instruction.opcode = opcode;
   instruction.rd = 0;
-  instruction.r2OrImm = pim::checkedI32OrCrash(remapCoreId(coreId), "communication core id");
+  instruction.r2OrImm = pim::checkedI32OrCrash(coreId, "physical communication core id");
   instruction.generic1 = 0;
   instruction.generic2 = 0;
   instruction.generic3 = pim::checkedI32OrCrash(size, "communication byte size");
@@ -679,6 +675,8 @@ void PimCodeGen::codeGenMVMLikeOp(size_t mvmId,
                                   MVMTy mvmLikeOp,
                                   bool transposeMatrix,
                                   const StaticValueKnowledge& knowledge) {
+  ensureVectorBitwidth(getVectorElementBitwidthOrCrash(cast<ShapedType>(mvmLikeOp.getInput().getType())),
+                       getVectorElementBitwidthOrCrash(cast<ShapedType>(mvmLikeOp.getOutputBuffer().getType())));
   emitMvmOp(mvmId, addressOf(mvmLikeOp.getOutputBuffer(), knowledge), 0, addressOf(mvmLikeOp.getInput(), knowledge), 0);
 
   // TODO: save weights somewhere (if transposeMatrix=true, transpose the weight matrix)
@@ -688,25 +686,29 @@ void PimCodeGen::emitBinaryVectorOp(pim_binary::Opcode opcode,
                                     mlir::Value output,
                                     mlir::Value lhs,
                                     mlir::Value rhs,
-                                    size_t byteSize,
                                     const StaticValueKnowledge& knowledge) const {
+  auto inputType = cast<ShapedType>(lhs.getType());
+  ensureVectorBitwidth(getVectorElementBitwidthOrCrash(inputType),
+                       getVectorElementBitwidthOrCrash(cast<ShapedType>(output.getType())));
   setupRdRs1Rs2(addressOf(output, knowledge), 0, addressOf(lhs, knowledge), 0, addressOf(rhs, knowledge), 0);
   pim_binary::InstructionRecord instruction;
   instruction.opcode = opcode;
   instruction.rd = 0;
   instruction.r1 = 1;
   instruction.r2OrImm = 2;
-  instruction.generic3 = pim::checkedI32OrCrash(byteSize, "vector byte size");
+  instruction.generic3 = getVectorElementCountOrCrash(inputType);
   emitInstruction(instruction);
 }
 
 void PimCodeGen::emitUnaryVectorOp(pim_binary::Opcode opcode,
                                    mlir::Value output,
                                    mlir::Value input,
-                                   size_t byteSize,
                                    const StaticValueKnowledge& knowledge,
                                    int32_t r2OrImm,
                                    int32_t generic1) const {
+  auto inputType = cast<ShapedType>(input.getType());
+  ensureVectorBitwidth(getVectorElementBitwidthOrCrash(inputType),
+                       getVectorElementBitwidthOrCrash(cast<ShapedType>(output.getType())));
   setupRdRs1(addressOf(output, knowledge), 0, addressOf(input, knowledge), 0);
   pim_binary::InstructionRecord instruction;
   instruction.opcode = opcode;
@@ -714,7 +716,7 @@ void PimCodeGen::emitUnaryVectorOp(pim_binary::Opcode opcode,
   instruction.r1 = 1;
   instruction.r2OrImm = r2OrImm;
   instruction.generic1 = generic1;
-  instruction.generic3 = pim::checkedI32OrCrash(byteSize, "vector byte size");
+  instruction.generic3 = getVectorElementCountOrCrash(inputType);
   emitInstruction(instruction);
 }
 
@@ -812,8 +814,7 @@ static SmallVector<Operation*> collectTopLevelCoreLikeOps(func::FuncOp funcOp) {
 static FailureOr<CompiledCoreMemoryPlan> compileCoreMemoryPlan(Operation* coreLikeOp) {
   CompiledCoreMemoryPlan plan;
   auto arenaAttr = coreLikeOp->getAttrOfType<IntegerAttr>(kLocalMemorySizeAttrName);
-  if (!arenaAttr || arenaAttr.getInt() < 0
-      || static_cast<uint64_t>(arenaAttr.getInt()) > kPimLocalMemoryAddressLimit) {
+  if (!arenaAttr || arenaAttr.getInt() < 0 || static_cast<uint64_t>(arenaAttr.getInt()) > kPimLocalMemoryAddressLimit) {
     coreLikeOp->emitError("requires a valid pim.local_memory_size attribute before codegen");
     return failure();
   }
@@ -841,7 +842,9 @@ static FailureOr<CompiledCoreMemoryPlan> compileCoreMemoryPlan(Operation* coreLi
       hasFailure = true;
       return;
     }
-    plan.entries.push_back({allocOp.getResult(), {address, static_cast<size_t>(*checkedSize)}});
+    plan.entries.push_back({
+      allocOp.getResult(), {address, static_cast<size_t>(*checkedSize)}
+    });
     auto logicalBytes = pim::checkedAdd(
       static_cast<size_t>(plan.logicalBytes), static_cast<size_t>(*checkedSize), allocOp, "logical local bytes");
     if (failed(logicalBytes)) {
@@ -995,21 +998,10 @@ 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(), knowledge);
     };
     auto emitUnary = [&](auto op, pim_binary::Opcode opcode, int32_t r2OrImm, int32_t generic1) {
-      coreCodeGen.emitUnaryVectorOp(opcode,
-                                    op.getOutputBuffer(),
-                                    op.getInput(),
-                                    getVectorByteSizeOrCrash(cast<ShapedType>(op.getInput().getType())),
-                                    knowledge,
-                                    r2OrImm,
-                                    generic1);
+      coreCodeGen.emitUnaryVectorOp(opcode, op.getOutputBuffer(), op.getInput(), knowledge, r2OrImm, generic1);
     };
 
     switch (node.opKind) {
@@ -1092,8 +1084,8 @@ static void aliasMaterializedHostGlobals(CoreLikeOpTy coreLikeOp,
   });
 }
 
-static OnnxMlirCompilerErrorCodes emitEmptyCoreArtifacts(StringRef outputDirPath, size_t emittedCoreId) {
-  std::string outputCorePath = (outputDirPath + "/core_" + std::to_string(emittedCoreId) + ".pim").str();
+static OnnxMlirCompilerErrorCodes emitEmptyCoreArtifacts(StringRef outputDirPath, size_t physicalCoreId) {
+  std::string outputCorePath = (outputDirPath + "/core_" + std::to_string(physicalCoreId) + ".pim").str();
   std::error_code errorCode;
   raw_fd_ostream coreBinaryStream(outputCorePath, errorCode, sys::fs::OF_None);
   if (errorCode) {
@@ -1117,7 +1109,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(physicalCoreId) + ".json").str();
   errorCode = std::error_code();
   raw_fd_ostream coreJsonStream(outputCoreJsonPath, errorCode);
   if (errorCode) {
@@ -1179,36 +1171,15 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
     return it->second.get();
   };
 
-  llvm::DenseMap<size_t, size_t> emittedCoreIds;
-  size_t nextEmittedCoreId = 0;
-
-  for (Operation* op : coreLikeOps) {
-    if (auto coreOp = dyn_cast<pim::PimCoreOp>(op)) {
-      size_t originalCoreId = static_cast<size_t>(coreOp.getCoreId());
-      if (!emittedCoreIds.contains(originalCoreId))
-        emittedCoreIds[originalCoreId] = nextEmittedCoreId++;
-      continue;
-    }
-
-    auto coreBatchOp = cast<pim::PimCoreBatchOp>(op);
-    auto batchCoreIds = getBatchCoreIds(coreBatchOp);
-    for (unsigned lane = 0; lane < static_cast<unsigned>(coreBatchOp.getLaneCount()); ++lane) {
-      size_t originalCoreId = static_cast<size_t>(batchCoreIds[lane]);
-      if (!emittedCoreIds.contains(originalCoreId))
-        emittedCoreIds[originalCoreId] = nextEmittedCoreId++;
-    }
-  }
-
   SmallVector<CoreEmissionJob> jobs;
   SmallVector<SmallVector<size_t>> batchJobIndices;
   for (Operation* op : coreLikeOps) {
     if (auto coreOp = dyn_cast<pim::PimCoreOp>(op)) {
-      size_t originalCoreId = static_cast<size_t>(coreOp.getCoreId());
       CoreEmissionJob job;
       job.coreLikeOp = coreOp;
       job.program = getCompiledProgram(op);
       job.memoryPlan = getMemoryPlan(op);
-      job.emittedCoreId = emittedCoreIds.lookup(originalCoreId);
+      job.physicalCoreId = static_cast<size_t>(coreOp.getCoreId());
       jobs.push_back(std::move(job));
       continue;
     }
@@ -1220,16 +1191,15 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
       lanesByCoreId[static_cast<size_t>(batchCoreIds[lane])].push_back(lane);
 
     SmallVector<size_t> jobIndices;
-    SmallVector<size_t> orderedOriginalCoreIds = llvm::to_vector(lanesByCoreId.keys());
-    llvm::sort(orderedOriginalCoreIds,
-               [&](size_t lhs, size_t rhs) { return emittedCoreIds.lookup(lhs) < emittedCoreIds.lookup(rhs); });
-    for (size_t originalCoreId : orderedOriginalCoreIds) {
+    SmallVector<size_t> physicalCoreIds = llvm::to_vector(lanesByCoreId.keys());
+    llvm::sort(physicalCoreIds);
+    for (size_t physicalCoreId : physicalCoreIds) {
       CoreEmissionJob job;
       job.coreLikeOp = coreBatchOp;
       job.program = getCompiledProgram(op);
       job.memoryPlan = getMemoryPlan(op);
-      job.emittedCoreId = emittedCoreIds.lookup(originalCoreId);
-      job.lanes = lanesByCoreId.lookup(originalCoreId);
+      job.physicalCoreId = physicalCoreId;
+      job.lanes = lanesByCoreId.lookup(physicalCoreId);
       job.batchReportId = nextBatchReportId;
       jobIndices.push_back(jobs.size());
       jobs.push_back(std::move(job));
@@ -1238,8 +1208,11 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
     ++nextBatchReportId;
   }
 
-  auto linkCoreWeights =
-    [&](size_t coreId, ArrayRef<std::string> weightFiles, json::Array& xbarsPerGroup) -> OnnxMlirCompilerErrorCodes {
+  auto linkCoreWeights = [&](size_t coreId,
+                             ArrayRef<std::string> weightFiles,
+                             ArrayRef<ResolvedWeightView> weights,
+                             json::Array& xbarsPerGroup) -> OnnxMlirCompilerErrorCodes {
+    assert(weightFiles.size() == weights.size() && "weight files must match resolved weight views");
     auto coreWeightsDirPath = outputDirPath + "/core_" + std::to_string(coreId);
     if (auto error = sys::fs::create_directory(coreWeightsDirPath); error && error != std::errc::file_exists) {
       errs() << "Error creating core directory: " << coreWeightsDirPath << ": " << error.message() << '\n';
@@ -1247,7 +1220,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
     }
 
     for (auto [slot, fileName] : llvm::enumerate(weightFiles)) {
-      xbarsPerGroup.push_back(1);
+      xbarsPerGroup.push_back(weights[slot].shape[1] / static_cast<int64_t>(crossbarSize));
       std::string sourcePath = outputDirPath + "/weights/" + fileName;
       std::string targetPath = coreWeightsDirPath + "/crossbar_" + std::to_string(slot) + ".bin";
       sys::fs::remove(targetPath);
@@ -1295,7 +1268,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
     };
 
     std::error_code errorCode;
-    auto outputCorePath = outputDirPath + "/core_" + std::to_string(job.emittedCoreId) + ".pim";
+    auto outputCorePath = outputDirPath + "/core_" + std::to_string(job.physicalCoreId) + ".pim";
     raw_fd_ostream coreBinaryStream(outputCorePath, errorCode, sys::fs::OF_None);
     if (errorCode) {
       errs() << "Error while opening core file `" << outputCorePath << "`: " << errorCode.message() << '\n';
@@ -1305,7 +1278,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
 
     std::unique_ptr<raw_fd_ostream> coreJsonStream;
     if (pimEmitJson.getValue()) {
-      std::string outputCoreJsonPath = outputDirPath + "/core_" + std::to_string(job.emittedCoreId) + ".json";
+      std::string outputCoreJsonPath = outputDirPath + "/core_" + std::to_string(job.physicalCoreId) + ".json";
       errorCode = std::error_code();
       coreJsonStream = std::make_unique<raw_fd_ostream>(outputCoreJsonPath, errorCode);
       if (errorCode) {
@@ -1317,7 +1290,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
     }
 
     PimInstructionWriter instructionWriter(coreBinaryStream);
-    PimCodeGen coreCodeGen(jobMemory, instructionWriter, coreJsonStream.get(), emittedCoreIds);
+    PimCodeGen coreCodeGen(jobMemory, instructionWriter, coreJsonStream.get());
     auto finalizeInstructions = [&]() {
       bool succeeded = mlir::succeeded(instructionWriter.finalize());
       coreBinaryStream.close();
@@ -1330,7 +1303,7 @@ 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);
+      auto& deviceMemory = jobMemory.getOrCreateDeviceMem(job.physicalCoreId);
       deviceMemory.allocateCore(*job.memoryPlan);
 
       StaticValueKnowledge knowledge = seedCoreCodegenKnowledge(coreOp);
@@ -1345,7 +1318,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
     else {
       auto coreBatchOp = cast<pim::PimCoreBatchOp>(job.coreLikeOp);
       aliasMaterializedHostGlobals(coreBatchOp, moduleOp, materializedHostGlobals, jobMemory);
-      auto& deviceMemory = jobMemory.getOrCreateDeviceMem(job.emittedCoreId);
+      auto& deviceMemory = jobMemory.getOrCreateDeviceMem(job.physicalCoreId);
 
       for (unsigned lane : job.lanes) {
         StaticValueKnowledge knowledge = seedCoreBatchCodegenKnowledge(coreBatchOp, lane);
@@ -1398,18 +1371,29 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
     if (jobResults[jobIndex].status != CompilerSuccess)
       return jobResults[jobIndex].status;
 
-  if (jobs.empty()) {
-    if (auto err = emitEmptyCoreArtifacts(outputDirPath, 0))
+  size_t maxPhysicalCoreId = 0;
+  for (const CoreEmissionJob& job : jobs)
+    maxPhysicalCoreId = std::max(maxPhysicalCoreId, job.physicalCoreId);
+  std::vector<bool> activePhysicalCores(maxPhysicalCoreId + 1);
+  for (const CoreEmissionJob& job : jobs)
+    activePhysicalCores[job.physicalCoreId] = true;
+  for (size_t physicalCoreId = 0;
+       physicalCoreId < activePhysicalCores.size(); ++physicalCoreId) {
+    if (activePhysicalCores[physicalCoreId])
+      continue;
+    if (auto err =
+          emitEmptyCoreArtifacts(outputDirPath, physicalCoreId))
       return err;
-    xbarsPerArrayGroup["core0"] = json::Array {};
-    memory.recordCoreReport(0, MemoryReportRow {});
+    xbarsPerArrayGroup["core" + std::to_string(physicalCoreId)] =
+      json::Array {};
+    memory.recordCoreReport(physicalCoreId, MemoryReportRow {});
   }
 
   llvm::SmallVector<WeightFileRequest, 8> weightRequests;
   weightRequests.reserve(jobs.size());
   for (size_t jobIndex = 0; jobIndex < jobs.size(); ++jobIndex) {
     WeightFileRequest request;
-    request.coreId = jobs[jobIndex].emittedCoreId;
+    request.coreId = jobs[jobIndex].physicalCoreId;
     request.weights = jobResults[jobIndex].usedWeights;
     weightRequests.push_back(std::move(request));
   }
@@ -1422,10 +1406,11 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
     json::Array xbarsPerGroup;
 
     if (auto coreOp = dyn_cast<pim::PimCoreOp>(job.coreLikeOp)) {
-      if (auto err = linkCoreWeights(job.emittedCoreId, mapCoreWeightToFileName[job.emittedCoreId], xbarsPerGroup))
+      if (auto err = linkCoreWeights(
+            job.physicalCoreId, mapCoreWeightToFileName[job.physicalCoreId], result.usedWeights, xbarsPerGroup))
         return err;
-      xbarsPerArrayGroup["core" + std::to_string(job.emittedCoreId)] = std::move(xbarsPerGroup);
-      memory.recordCoreReport(job.emittedCoreId, result.reportRow);
+      xbarsPerArrayGroup["core" + std::to_string(job.physicalCoreId)] = std::move(xbarsPerGroup);
+      memory.recordCoreReport(job.physicalCoreId, result.reportRow);
       continue;
     }
   }
@@ -1438,10 +1423,11 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
       const CoreEmissionJob& job = jobs[jobIndex];
       const CoreEmissionResult& result = jobResults[jobIndex];
       json::Array xbarsPerGroup;
-      if (auto err = linkCoreWeights(job.emittedCoreId, mapCoreWeightToFileName[job.emittedCoreId], xbarsPerGroup))
+      if (auto err = linkCoreWeights(
+            job.physicalCoreId, mapCoreWeightToFileName[job.physicalCoreId], result.usedWeights, xbarsPerGroup))
         return err;
-      xbarsPerArrayGroup["core" + std::to_string(job.emittedCoreId)] = std::move(xbarsPerGroup);
-      reportedCoreIds.push_back(pim::checkedI32OrCrash(job.emittedCoreId, "batch report core id"));
+      xbarsPerArrayGroup["core" + std::to_string(job.physicalCoreId)] = std::move(xbarsPerGroup);
+      reportedCoreIds.push_back(pim::checkedI32OrCrash(job.physicalCoreId, "batch report physical core id"));
       if (!batchPerCoreRow)
         batchPerCoreRow = result.reportRow;
       else if (!(*batchPerCoreRow == result.reportRow))
@@ -1449,11 +1435,10 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
     }
 
     uint64_t batchReportId = jobs[group.front()].batchReportId.value_or(0);
-    memory.recordBatchReport(
-      batchReportId, reportedCoreIds, batchPerCoreRow.value_or(MemoryReportRow {}));
+    memory.recordBatchReport(batchReportId, reportedCoreIds, batchPerCoreRow.value_or(MemoryReportRow {}));
   }
 
-  maxCoreId = nextEmittedCoreId == 0 ? 0 : nextEmittedCoreId - 1;
+  maxCoreId = maxPhysicalCoreId;
 
   memory.flushReport();
   return writeConfigJson(funcOp, memory, maxCoreId, std::move(xbarsPerArrayGroup), outputDirPath);