#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/Threading.h"
#include "mlir/IR/Value.h"
#include "mlir/IR/Verifier.h"

#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
#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"

#include <absl/types/compare.h>
#include <algorithm>
#include <cassert>
#include <cstdint>
#include <fstream>
#include <limits>
#include <memory>
#include <numeric>
#include <string>
#include <utility>

#include "Common/IR/CompactAsmUtils.hpp"
#include "Common/PimCommon.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"
#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/PimCoreProgram.hpp"
#include "src/Accelerators/PIM/Compiler/PimWeightEmitter.hpp"
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"

using namespace llvm;
using namespace mlir;
using namespace onnx_mlir;
using namespace onnx_mlir::compact_asm;

namespace {

static std::optional<unsigned> getLaneForMemoryValue(mlir::Value value, std::optional<unsigned> lane) {
  if (!lane)
    return std::nullopt;
  auto allocOp = value.getDefiningOp<memref::AllocOp>();
  if (!allocOp || !allocOp->getParentOfType<pim::PimCoreBatchOp>())
    return std::nullopt;
  return lane;
}

static mlir::Value resolveCachedAlias(mlir::Value value, const StaticValueKnowledge& knowledge) {
  auto iter = knowledge.aliases.find(value);
  while (iter != knowledge.aliases.end()) {
    value = iter->second;
    iter = knowledge.aliases.find(value);
  }
  return value;
}

static MemoryValueKey getMemoryValueKey(mlir::Value value, std::optional<unsigned> lane = std::nullopt) {
  return {value, getLaneForMemoryValue(value, lane)};
}

static bool isInsidePimCoreLikeOp(memref::AllocOp allocOp) {
  return allocOp->getParentOfType<pim::PimCoreOp>() || allocOp->getParentOfType<pim::PimCoreBatchOp>();
}

static MemoryReportKind classifyMemoryReportKind(mlir::Value value) {
  if (isa<mlir::BlockArgument>(value))
    return MemoryReportKind::Input;
  if (auto* op = value.getDefiningOp()) {
    if (isa<memref::AllocOp>(op))
      return MemoryReportKind::Alloca;
    if (isa<memref::GetGlobalOp>(op))
      return MemoryReportKind::Global;
  }
  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 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 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(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: " << kPimLocalMemoryAddressLimit << " (signed int32_t immediate range)\n";
  if (key.lane)
    llvm::errs() << "Lane: " << *key.lane << "\n";
  llvm::errs() << "Value: ";
  key.value.print(llvm::errs());
  llvm::errs() << "\n";
  llvm::errs() << "Value type: " << key.value.getType() << "\n";
  if (Operation* definingOp = key.value.getDefiningOp()) {
    llvm::errs() << "Defining op:\n";
    definingOp->print(llvm::errs());
    llvm::errs() << "\n";
  }
}

static std::fstream openMemoryReport(bool enabled) {
  if (std::string outputRoot = getOutputDir(); !outputRoot.empty()) {
    for (std::string name : {std::string("memory_") + "coalescing_report.txt",
                             std::string("pim_memory_") + "liveness_report.txt",
                             std::string("pim_memory_") + "liveness_report.json",
                             std::string("pim_memory_") + "liveness_timeline.dot"})
      sys::fs::remove(outputRoot + "/reports/" + name);
    if (!enabled)
      sys::fs::remove(outputRoot + "/reports/memory_report.txt");
  }
  return enabled ? openReportFile("memory_report") : std::fstream();
}

} // namespace

size_t PimMemory::allocateAddress(size_t size, const MemoryValueKey& key) {
  size_t address = firstAvailableAddress;
  Operation* anchor = getDiagnosticAnchor(key.value);
  auto checkedEnd = pim::checkedAdd(address, size, anchor, "local memory end");
  FailureOr<size_t> checkedAlignedEnd = failure();
  if (succeeded(checkedEnd))
    checkedAlignedEnd = checkedAlignTo(*checkedEnd, minAlignment, anchor, "local memory alignment");
  if (address > kPimLocalMemoryAddressLimit || failed(checkedEnd) || *checkedEnd > kPimLocalMemoryAddressLimit
      || failed(checkedAlignedEnd) || *checkedAlignedEnd > kPimLocalMemoryAddressLimit) {
    printMemoryOverflowDiagnostic(
      key,
      size,
      firstAvailableAddress,
      succeeded(checkedAlignedEnd) ? *checkedAlignedEnd : kPimLocalMemoryAddressLimit);
    llvm_unreachable("PIM local memory allocation overflow");
  }
  firstAvailableAddress = *checkedAlignedEnd;
  return address;
}

MemEntry* PimMemory::gatherMemEntry(mlir::Value value, std::optional<unsigned> lane) {
  auto type = cast<ShapedType>(value.getType());
  assert("Only static shape is supported" && type.hasStaticShape());
  auto checkedAllocSize =
    pim::getCheckedShapedTypeSizeInBytes(type, UnknownLoc::get(type.getContext()), "memory allocation byte size");
  if (failed(checkedAllocSize))
    llvm_unreachable("Failed to compute checked allocation byte size");
  PendingMemEntry pending;
  pending.memEntry = {0, *checkedAllocSize};
  pending.key = getMemoryValueKey(value, lane);
  pending.reportKind = classifyMemoryReportKind(value);
  return &memEntries.emplace_back(std::move(pending)).memEntry;
}

void PimMemory::allocateGatheredMemory() {
  llvm::sort(memEntries, [](const PendingMemEntry& lhs, const PendingMemEntry& rhs) {
    return lhs.memEntry.size > rhs.memEntry.size;
  });
  for (PendingMemEntry& pending : memEntries)
    allocateMemoryForValue(pending.key, pending.memEntry, pending.reportKind);
  memEntries.clear();
}

void PimMemory::allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memEntry, MemoryReportKind reportKind) {
  memEntry.address = allocateAddress(memEntry.size, key);

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

  switch (reportKind) {
  case MemoryReportKind::Alloca:
  case MemoryReportKind::Global:
  case MemoryReportKind::Input:
    ++reportRow.hostObjectCount;
    break;
  case MemoryReportKind::None:  break;
  }
}

void PimMemory::allocateHost(ModuleOp moduleOp, func::FuncOp funcOp) {
  SmallDenseMap<memref::GlobalOp, mlir::Value, 8> globalConstants;
  SmallVector<std::pair<mlir::Value, mlir::Value>, 16> globalAliases;
  SmallVector<mlir::Value> args;

  for (mlir::Value arg : funcOp.getArguments()) {
    gatherMemEntry(arg);
    args.push_back(arg);
  }

  funcOp.walk([&](memref::GetGlobalOp getGlobalOp) {
    if (!hasWeightAlways(getGlobalOp)) {
      auto globalMemrefOp = lookupGlobalForGetGlobal(moduleOp, getGlobalOp);
      if (globalMemrefOp.getName().starts_with("arg")) {
        StringRef indexStr = globalMemrefOp.getName().substr(4);
        int index = 0;
        llvm::to_integer(indexStr, index, 10);
        globalAliases.push_back({getGlobalOp.getResult(), args[index]});
      }
      auto [iter, inserted] = globalConstants.try_emplace(globalMemrefOp, getGlobalOp.getResult());
      if (inserted)
        gatherMemEntry(getGlobalOp.getResult());
      else
        globalAliases.push_back({getGlobalOp.getResult(), iter->second});
    }
  });

  funcOp.walk([&](memref::AllocOp allocOp) {
    if (!isInsidePimCoreLikeOp(allocOp))
      gatherMemEntry(allocOp.getResult());
  });

  allocateGatheredMemory();

  for (auto [alias, original] : globalAliases)
    globalMemEntriesMap[getMemoryValueKey(alias)] = getMemEntry(getMemoryValueKey(original));
}

void PimMemory::allocateCore(const CompiledCoreMemoryPlan& plan, std::optional<unsigned> lane) {
  if (!localArenaSize) {
    localArenaSize = plan.arenaSize;
    reportRow.logicalLocalAllocationCount = plan.logicalAllocationCount;
    reportRow.logicalLocalBytes = plan.logicalBytes;
  }
  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) {
    MemoryValueKey key = getMemoryValueKey(entry.value, lane);
    ownedMemEntriesMap[key] = entry.memory;
    globalMemEntriesMap[key] = entry.memory;
  }
}

MemoryReportRow PimMemory::getReportRow() const {
  MemoryReportRow row = reportRow;
  row.hostBytes = firstAvailableAddress;
  row.physicalLocalBytes = localArenaSize.value_or(0);
  return row;
}

MemEntry PimMemory::getMemEntry(const MemoryValueKey& key) const {
  auto iter = globalMemEntriesMap.find(key);
  assert("Missing memEntry for value" && iter != globalMemEntriesMap.end());
  return iter->second;
}

PimMemory& PimAcceleratorMemory::getOrCreateDeviceMem(size_t id) {
  return deviceMem.try_emplace(id, memEntriesMap).first->second;
}

size_t PimAcceleratorMemory::getValueAddress(mlir::Value value,
                                             const StaticValueKnowledge& knowledge,
                                             std::optional<unsigned> lane) const {
  value = resolveCachedAlias(value, knowledge);

  FailureOr<ResolvedContiguousAddress> resolvedAddress = resolveContiguousAddress(value, knowledge);
  if (failed(resolvedAddress)) {
    auto compiledIt = compiledAddressExprs.find(value);
    if (compiledIt == compiledAddressExprs.end()) {
      auto compiledExpr = compileContiguousAddressExpr(value);
      if (failed(compiledExpr)) {
        errs() << "Failed to compile contiguous address for value: ";
        value.print(errs());
        errs() << " : " << value.getType();
        errs() << "\n";
        llvm_unreachable("Failed to compile contiguous address");
      }
      compiledIt = compiledAddressExprs.try_emplace(value, *compiledExpr).first;
    }

    resolvedAddress = compiledIt->second.evaluate(knowledge, lane);
    if (failed(resolvedAddress)) {
      errs() << "Failed to evaluate contiguous address for value: ";
      value.print(errs());
      errs() << " : " << value.getType();
      errs() << "\n";
      if (auto* definingOp = value.getDefiningOp()) {
        errs() << "Defining op:\n";
        definingOp->print(errs());
        errs() << "\n";
      }
      llvm_unreachable("Failed to resolve contiguous address");
    }
  }

  MemoryValueKey key = getMemoryValueKey(resolvedAddress->base, lane);
  auto iter = memEntriesMap.find(key);
  if (iter == memEntriesMap.end()) {
    errs() << "Missing mem entry for value: ";
    resolvedAddress->base.print(errs());
    errs() << "\n";
    if (key.lane)
      errs() << "Lane: " << *key.lane << "\n";
    if (auto* definingOp = resolvedAddress->base.getDefiningOp()) {
      errs() << "Defining op:\n";
      definingOp->print(errs());
      errs() << "\n";
    }
    llvm_unreachable("Missing mem entry");
  }

  size_t byteOffset = pim::checkedSizeOrCrash(resolvedAddress->byteOffset, "resolved PIM byte offset");
  return pim::checkedAddOrCrash(iter->second.address, byteOffset, "resolved PIM address");
}

llvm::FailureOr<int64_t> PimAcceleratorMemory::getIndexValue(mlir::Value value,
                                                             const StaticValueKnowledge& knowledge) const {
  value = resolveCachedAlias(value, knowledge);
  auto compiledIt = compiledIndexExprs.find(value);
  if (compiledIt == compiledIndexExprs.end()) {
    auto compiledExpr = compileIndexExpr(value);
    if (failed(compiledExpr))
      return mlir::failure();
    compiledIt = compiledIndexExprs.try_emplace(value, *compiledExpr).first;
  }
  return compiledIt->second.evaluate(knowledge);
}

PimAcceleratorMemory::PimAcceleratorMemory()
: hostMem(memEntriesMap), fileReport(openMemoryReport(pimMemoryReport == PimMemoryReportSummary)) {}

PimAcceleratorMemory::PimAcceleratorMemory(
  const llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32>& initialMemEntries, bool enableReport)
: memEntriesMap(initialMemEntries),
  hostMem(memEntriesMap),
  fileReport(enableReport ? openMemoryReport(true) : std::fstream()) {}

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

void PimAcceleratorMemory::recordBatchReport(uint64_t batchId,
                                             ArrayRef<int32_t> coreIds,
                                             const MemoryReportRow& perCoreRow) {
  MemoryReportEntry entry;
  entry.kind = MemoryReportEntry::Kind::Batch;
  entry.id = batchId;
  llvm::append_range(entry.coreIds, coreIds);
  entry.row = perCoreRow;
  reportEntries.push_back(std::move(entry));
}

void PimAcceleratorMemory::flushReport() {
  if (!fileReport.is_open())
    return;

  llvm::raw_os_ostream os(fileReport);
  struct ReportGroup {
    MemoryReportRow row;
    SmallVector<int32_t, 8> coreIds;
    std::optional<uint64_t> batchId;
  };
  SmallVector<ReportGroup, 32> groups;
  for (const MemoryReportEntry& entry : reportEntries) {
    auto group = llvm::find_if(groups, [&](const ReportGroup& candidate) {
      return candidate.row.logicalLocalBytes == entry.row.logicalLocalBytes
          && candidate.row.physicalLocalBytes == entry.row.physicalLocalBytes;
    });
    if (group == groups.end()) {
      ReportGroup newGroup;
      newGroup.row = entry.row;
      llvm::append_range(newGroup.coreIds, entry.coreIds);
      if (entry.kind == MemoryReportEntry::Kind::Batch)
        newGroup.batchId = entry.id;
      groups.push_back(std::move(newGroup));
      continue;
    }
    llvm::append_range(group->coreIds, entry.coreIds);
    group->batchId.reset();
  }
  llvm::sort(groups, [](const ReportGroup& lhs, const ReportGroup& rhs) {
    if (lhs.row.physicalLocalBytes != rhs.row.physicalLocalBytes)
      return lhs.row.physicalLocalBytes > rhs.row.physicalLocalBytes;
    return lhs.row.logicalLocalBytes > rhs.row.logicalLocalBytes;
  });
  uint64_t logicalBytes = 0;
  uint64_t physicalBytes = 0;
  for (const ReportGroup& group : groups) {
    logicalBytes += group.row.logicalLocalBytes * group.coreIds.size();
    physicalBytes += group.row.physicalLocalBytes * group.coreIds.size();
  }
  uint64_t savedBytes = logicalBytes - physicalBytes;
  double savedPercent = logicalBytes == 0 ? 0.0 : 100.0 * savedBytes / logicalBytes;
  uint64_t largest = groups.empty() ? 0 : groups.front().row.physicalLocalBytes;
  auto printLabel = [&](const ReportGroup& group) {
    if (group.batchId)
      os << "Batch " << *group.batchId << " — cores ";
    else if (group.coreIds.size() == 1)
      os << "Core ";
    else
      os << "Cores ";
    printCompressedIntegerEntries(os, ArrayRef<int32_t>(group.coreIds));
  };

  os << "Summary\n";
  os << "  Host memory: " << formatReportMemory(hostReportRow ? hostReportRow->hostBytes : 0) << "\n";
  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 << "  Largest core local memory: " << formatReportMemory(largest) << "\n";
  if (!groups.empty()) {
    os << "    ";
    printLabel(groups.front());
    os << "\n";
  }

  os << "\nLargest core groups\n";
  constexpr size_t kGroupLimit = 8;
  for (const ReportGroup& group : ArrayRef(groups).take_front(kGroupLimit)) {
    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;
    if (group.coreIds.size() == 1) {
      os << "  Local memory: " << formatReportMemory(group.row.logicalLocalBytes) << " → "
         << 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";
    }
  }
  if (groups.size() > kGroupLimit)
    os << "  " << groups.size() - kGroupLimit << " smaller groups omitted\n";

  os.flush();
  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;
  if (coreJsonStream)
    *coreJsonStream << json::Value(pim_binary::makeInstructionJson(instruction)) << ',';
  updateScalarRegisterCache(instruction);
}

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

void PimCodeGen::genSetRegisterImmediate(uint8_t registerNumber, int32_t immediate) const {
  if (scalarRegisterValues[registerNumber] == immediate)
    return;

  pim_binary::InstructionRecord instruction;
  instruction.opcode = pim_binary::Opcode::sldi;
  instruction.rd = registerNumber;
  instruction.r2OrImm = immediate;
  emitInstruction(instruction);
}

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

void PimCodeGen::setupRd(size_t rdAddress, size_t rdOffset) const {
  genSetRegisterImmediateUnsigned(0, pim::checkedAddOrCrash(rdAddress, rdOffset, "rd address"));
}

void PimCodeGen::setupRdRs1(size_t rdAddress, size_t rdOffset, size_t rs1Address, size_t rs1Offset) const {
  genSetRegisterImmediateUnsigned(0, pim::checkedAddOrCrash(rdAddress, rdOffset, "rd address"));
  genSetRegisterImmediateUnsigned(1, pim::checkedAddOrCrash(rs1Address, rs1Offset, "rs1 address"));
}

void PimCodeGen::setupRdRs1Rs2(
  size_t rdAddress, size_t rdOffset, size_t rs1Address, size_t rs1Offset, size_t rs2Address, size_t rs2Offset) const {
  genSetRegisterImmediateUnsigned(0, pim::checkedAddOrCrash(rdAddress, rdOffset, "rd address"));
  genSetRegisterImmediateUnsigned(1, pim::checkedAddOrCrash(rs1Address, rs1Offset, "rs1 address"));
  genSetRegisterImmediateUnsigned(2, pim::checkedAddOrCrash(rs2Address, rs2Offset, "rs2 address"));
}

void PimCodeGen::emitMemCopyOp(pim_binary::Opcode opcode,
                               size_t rdAddr,
                               size_t rdOffset,
                               size_t rs1Addr,
                               size_t rs1Offset,
                               size_t size,
                               StringRef sizeFieldName) const {
  setupRdRs1(rdAddr, rdOffset, rs1Addr, rs1Offset);

  pim_binary::InstructionRecord instruction;
  instruction.opcode = opcode;
  instruction.rd = 0;
  instruction.r1 = 1;
  instruction.generic1 = 0;
  instruction.generic2 = 0;
  instruction.generic3 = pim::checkedI32OrCrash(size, sizeFieldName);
  emitInstruction(instruction);
}

void PimCodeGen::emitCommunicationOp(pim_binary::Opcode opcode, size_t bufferAddr, size_t coreId, size_t size) const {
  setupRd(bufferAddr, 0);

  pim_binary::InstructionRecord instruction;
  instruction.opcode = opcode;
  instruction.rd = 0;
  instruction.r2OrImm = pim::checkedI32OrCrash(remapCoreId(coreId), "communication core id");
  instruction.generic1 = 0;
  instruction.generic2 = 0;
  instruction.generic3 = pim::checkedI32OrCrash(size, "communication byte size");
  emitInstruction(instruction);
}

void PimCodeGen::emitMvmOp(size_t groupId, size_t rdAddr, size_t rdOffset, size_t rs1Addr, size_t rs1Offset) const {
  setupRdRs1(rdAddr, rdOffset, rs1Addr, rs1Offset);

  pim_binary::InstructionRecord instruction;
  instruction.opcode = pim_binary::Opcode::mvmul;
  instruction.rd = 0;
  instruction.r1 = 1;
  instruction.r2OrImm = 8;
  instruction.generic1 = 0;
  instruction.generic2 = pim::checkedI32OrCrash(groupId, "mvm group id");
  emitInstruction(instruction);
}

void PimCodeGen::codeGenLoadOp(pim::PimMemCopyHostToDevOp loadOp, const StaticValueKnowledge& knowledge) const {
  auto deviceTargetOffset = indexOf(loadOp.getDeviceTargetOffset(), knowledge);
  auto hostSourceOffset = indexOf(loadOp.getHostSourceOffset(), knowledge);
  assert(succeeded(deviceTargetOffset) && succeeded(hostSourceOffset)
         && "pim.memcp_hd offsets must be statically resolvable during codegen");
  emitMemCopyOp(pim_binary::Opcode::ld,
                addressOf(loadOp.getDeviceTarget(), knowledge),
                *deviceTargetOffset,
                addressOf(loadOp.getHostSource(), knowledge),
                *hostSourceOffset,
                loadOp.getSize());
}

void PimCodeGen::codeGenStoreOp(pim::PimMemCopyDevToHostOp storeOp, const StaticValueKnowledge& knowledge) const {
  auto hostTargetOffset = indexOf(storeOp.getHostTargetOffset(), knowledge);
  auto deviceSourceOffset = indexOf(storeOp.getDeviceSourceOffset(), knowledge);
  assert(succeeded(hostTargetOffset) && succeeded(deviceSourceOffset)
         && "pim.memcp_dh offsets must be statically resolvable during codegen");
  emitMemCopyOp(pim_binary::Opcode::st,
                addressOf(storeOp.getHostTarget(), knowledge),
                *hostTargetOffset,
                addressOf(storeOp.getDeviceSource(), knowledge),
                *deviceSourceOffset,
                storeOp.getSize());
}

void PimCodeGen::codeGenLmvOp(pim::PimMemCopyOp lmvOp, const StaticValueKnowledge& knowledge) const {
  auto targetOffset = indexOf(lmvOp.getTargetOffset(), knowledge);
  auto sourceOffset = indexOf(lmvOp.getSourceOffset(), knowledge);
  assert(succeeded(targetOffset) && succeeded(sourceOffset)
         && "pim.memcp offsets must be statically resolvable during codegen");
  emitMemCopyOp(pim_binary::Opcode::lmv,
                addressOf(lmvOp.getTarget(), knowledge),
                *targetOffset,
                addressOf(lmvOp.getSource(), knowledge),
                *sourceOffset,
                lmvOp.getSize(),
                "len");
}

void PimCodeGen::codeGenReceiveOp(pim::PimReceiveOp receiveOp, const StaticValueKnowledge& knowledge) const {
  auto sourceCoreId = indexOf(receiveOp.getSourceCoreId(), knowledge);
  assert(succeeded(sourceCoreId) && "pim.receive source core id must be statically resolvable during codegen");
  emitCommunicationOp(
    pim_binary::Opcode::recv, addressOf(receiveOp.getOutputBuffer(), knowledge), *sourceCoreId, receiveOp.getSize());
}

void PimCodeGen::codeGenSendOp(pim::PimSendOp sendOp, const StaticValueKnowledge& knowledge) const {
  auto targetCoreId = indexOf(sendOp.getTargetCoreId(), knowledge);
  assert(succeeded(targetCoreId) && "pim.send target core id must be statically resolvable during codegen");
  emitCommunicationOp(
    pim_binary::Opcode::send, addressOf(sendOp.getInput(), knowledge), *targetCoreId, sendOp.getSize());
}

void PimCodeGen::codeGenConcatOp(pim::PimConcatOp concatOp, const StaticValueKnowledge& knowledge) const {
  auto outputType = cast<ShapedType>(concatOp.getOutputBuffer().getType());
  assert(outputType.hasStaticShape() && "concat codegen requires static output shape");

  int64_t axis = concatOp.getAxis();
  ArrayRef<int64_t> outputShape = outputType.getShape();
  size_t elementSize = getElementTypeSizeInBytes(outputType.getElementType());
  size_t outputAddr = addressOf(concatOp.getOutputBuffer(), knowledge);

  size_t outerCount = 1;
  for (int64_t dim = 0; dim < axis; ++dim)
    outerCount *= static_cast<size_t>(outputShape[dim]);

  size_t innerCount = 1;
  for (size_t dim = static_cast<size_t>(axis) + 1; dim < outputShape.size(); ++dim)
    innerCount *= static_cast<size_t>(outputShape[dim]);

  size_t outputConcatDim = static_cast<size_t>(outputShape[axis]);
  size_t concatOffset = 0;
  for (mlir::Value input : concatOp.getInputs()) {
    auto inputType = cast<ShapedType>(input.getType());
    assert(inputType.hasStaticShape() && "concat codegen requires static input shapes");

    size_t inputConcatDim = static_cast<size_t>(inputType.getDimSize(axis));
    size_t blockSizeInBytes = inputConcatDim * innerCount * elementSize;
    size_t inputAddr = addressOf(input, knowledge);

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

    concatOffset += inputConcatDim;
  }
}

template <typename MVMTy>
void PimCodeGen::codeGenMVMLikeOp(size_t mvmId,
                                  MVMTy mvmLikeOp,
                                  bool transposeMatrix,
                                  const StaticValueKnowledge& knowledge) {
  emitMvmOp(mvmId, addressOf(mvmLikeOp.getOutputBuffer(), knowledge), 0, addressOf(mvmLikeOp.getInput(), knowledge), 0);

  // TODO: save weights somewhere (if transposeMatrix=true, transpose the weight matrix)
}

void PimCodeGen::emitBinaryVectorOp(pim_binary::Opcode opcode,
                                    mlir::Value output,
                                    mlir::Value lhs,
                                    mlir::Value rhs,
                                    size_t byteSize,
                                    const StaticValueKnowledge& knowledge) const {
  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");
  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 {
  setupRdRs1(addressOf(output, knowledge), 0, addressOf(input, knowledge), 0);
  pim_binary::InstructionRecord instruction;
  instruction.opcode = opcode;
  instruction.rd = 0;
  instruction.r1 = 1;
  instruction.r2OrImm = r2OrImm;
  instruction.generic1 = generic1;
  instruction.generic3 = pim::checkedI32OrCrash(byteSize, "vector byte size");
  emitInstruction(instruction);
}

void PimCodeGen::codeGenTransposeOp(const CompiledTransposePlan& plan, const StaticValueKnowledge& knowledge) const {
  auto srcAddr = addressOf(plan.source, knowledge);
  auto dstAddr = addressOf(plan.destination, knowledge);

  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;

  pim_binary::InstructionRecord copyInstruction;
  copyInstruction.opcode = pim_binary::Opcode::lmv;
  copyInstruction.rd = 0;
  copyInstruction.r1 = 1;

  size_t maxRunElements = static_cast<size_t>(std::numeric_limits<int32_t>::max()) / plan.elementBytes;
  auto emitRun = [&](size_t sourceStart, size_t destinationStart, size_t runLength) {
    while (runLength != 0) {
      size_t chunkElements = std::min(runLength, maxRunElements);
      // totalBytes was checked when the plan was compiled, so these bounded
      // products cannot overflow.
      size_t sourceOffset = sourceStart * plan.elementBytes;
      size_t destinationOffset = destinationStart * plan.elementBytes;
      size_t byteSize = chunkElements * plan.elementBytes;
      assert(sourceOffset <= plan.totalBytes - byteSize && destinationOffset <= plan.totalBytes - byteSize);
      genSetRegisterImmediate(0, static_cast<int32_t>(dstAddr + destinationOffset));
      genSetRegisterImmediate(1, static_cast<int32_t>(srcAddr + sourceOffset));
      copyInstruction.generic3 = static_cast<int32_t>(byteSize);
      emitInstruction(copyInstruction);
      sourceStart += chunkElements;
      destinationStart += chunkElements;
      runLength -= chunkElements;
    }
  };

  if (plan.storagePreserving) {
    emitRun(0, 0, plan.totalElements);
    return;
  }

  size_t rank = plan.sourceShape.size();
  SmallVector<size_t> sourceIndices(rank, 0);
  size_t destinationFlat = 0;
  size_t runSourceStart = 0;
  size_t runDestinationStart = 0;
  size_t runLength = 0;

  for (size_t sourceFlat = 0; sourceFlat < plan.totalElements; ++sourceFlat) {
    if (runLength != 0 && destinationFlat != runDestinationStart + runLength) {
      emitRun(runSourceStart, runDestinationStart, runLength);
      runSourceStart = sourceFlat;
      runDestinationStart = destinationFlat;
      runLength = 0;
    }
    if (runLength == 0) {
      runSourceStart = sourceFlat;
      runDestinationStart = destinationFlat;
    }
    ++runLength;

    if (runLength == maxRunElements) {
      emitRun(runSourceStart, runDestinationStart, runLength);
      runLength = 0;
    }

    if (sourceFlat + 1 == plan.totalElements)
      break;
    for (size_t sourceDim = rank; sourceDim-- > 0;) {
      destinationFlat += plan.destinationStrides[plan.destinationDimensionForSource[sourceDim]];
      if (++sourceIndices[sourceDim] < plan.sourceShape[sourceDim])
        break;
      sourceIndices[sourceDim] = 0;
      destinationFlat -= plan.destinationRewinds[sourceDim];
    }
  }

  if (runLength != 0)
    emitRun(runSourceStart, runDestinationStart, runLength);
}

static SmallVector<Operation*> collectTopLevelCoreLikeOps(func::FuncOp funcOp) {
  SmallVector<Operation*> coreLikeOps;
  for (Operation& op : funcOp.getBody().front())
    if (dyn_cast<pim::PimCoreOp>(&op) || dyn_cast<pim::PimCoreBatchOp>(&op))
      coreLikeOps.push_back(&op);
  return coreLikeOps;
}

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) {
    coreLikeOp->emitError("requires a valid pim.local_memory_size attribute before codegen");
    return failure();
  }
  plan.arenaSize = static_cast<size_t>(arenaAttr.getInt());
  bool hasFailure = false;
  coreLikeOp->walk([&](memref::AllocOp allocOp) {
    if (hasFailure)
      return;
    auto addressAttr = allocOp->getAttrOfType<IntegerAttr>(kLocalMemoryAddressAttrName);
    if (!addressAttr || addressAttr.getInt() < 0) {
      allocOp.emitOpError("requires a non-negative pim.local_memory_address attribute 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());
    if (address > plan.arenaSize || *checkedSize > plan.arenaSize - address) {
      allocOp.emitOpError("has a planned local-memory range outside its core arena");
      hasFailure = true;
      return;
    }
    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)) {
      hasFailure = true;
      return;
    }
    plan.logicalBytes = *logicalBytes;
    ++plan.logicalAllocationCount;
  });
  if (hasFailure)
    return failure();
  return plan;
}

struct CoreEmissionResult {
  static constexpr size_t kMaxStoredCodegenDiagnostics = 8;

  struct DiagnosticRecord {
    Operation* op = nullptr;
    std::string message;
  };

  OnnxMlirCompilerErrorCodes status = CompilerSuccess;
  MemoryReportRow reportRow;
  llvm::SmallVector<ResolvedWeightView, 8> usedWeights;
  llvm::SmallVector<DiagnosticRecord, kMaxStoredCodegenDiagnostics> diagnostics;
  size_t diagnosticCount = 0;
  void recordDiagnostic(Operation* op, StringRef message) {
    ++diagnosticCount;
    if (diagnostics.size() < kMaxStoredCodegenDiagnostics)
      diagnostics.push_back({op, message.str()});
  }
};

static StaticValueKnowledge seedCoreCodegenKnowledge(pim::PimCoreOp coreOp) {
  StaticValueKnowledge knowledge;
  for (auto [index, weight] : llvm::enumerate(coreOp.getWeights()))
    knowledge.aliases[coreOp.getWeightArgument(index)] = weight;
  return knowledge;
}

static StaticValueKnowledge seedCoreBatchCodegenKnowledge(pim::PimCoreBatchOp coreBatchOp, unsigned lane) {
  StaticValueKnowledge knowledge;
  knowledge.indexValues[coreBatchOp.getLaneArgument()] = lane;
  for (auto [index, weight] : llvm::enumerate(coreBatchOp.getWeights()))
    knowledge.aliases[coreBatchOp.getWeightArgument(index)] = weight;
  for (auto [index, input] : llvm::enumerate(coreBatchOp.getInputs()))
    knowledge.aliases[coreBatchOp.getInputArgument(index)] = input;
  return knowledge;
}

template <typename MapTy>
class ScopedMapBindings {
  using KeyTy = typename MapTy::key_type;
  using ValueTy = typename MapTy::mapped_type;

  MapTy& map;
  llvm::SmallVector<std::pair<KeyTy, std::optional<ValueTy>>, 8> savedEntries;

public:
  explicit ScopedMapBindings(MapTy& map)
  : map(map) {}

  void bind(const KeyTy& key, const ValueTy& value) {
    auto it = map.find(key);
    if (it == map.end())
      savedEntries.emplace_back(key, std::nullopt);
    else
      savedEntries.emplace_back(key, it->second);
    map[key] = value;
  }

  ~ScopedMapBindings() {
    for (auto it = savedEntries.rbegin(); it != savedEntries.rend(); ++it)
      if (it->second)
        map[it->first] = *it->second;
      else
        map.erase(it->first);
  }
};

static LogicalResult executeCompiledCorePlan(
  const llvm::SmallVectorImpl<CompiledCoreNode>& plan,
  PimCodeGen& coreCodeGen,
  StaticValueKnowledge& knowledge,
  llvm::function_ref<llvm::FailureOr<unsigned>(pim::PimVMMOp, const StaticValueKnowledge&)> resolveWeightSlot) {
  for (const CompiledCoreNode& node : plan) {
    if (node.kind == CompiledCoreNode::Kind::Loop) {
      auto lowerBound = node.lowerBound.evaluate(knowledge);
      auto upperBound = node.upperBound.evaluate(knowledge);
      auto step = node.step.evaluate(knowledge);
      auto forOp = cast<mlir::scf::ForOp>(node.op);
      if (failed(lowerBound) || failed(upperBound) || failed(step) || *step <= 0) {
        forOp.emitOpError("requires statically evaluable scf.for bounds for PIM codegen");
        return failure();
      }

      llvm::SmallVector<mlir::Value> iterValues(forOp.getInitArgs().begin(), forOp.getInitArgs().end());
      for (int64_t inductionValue = *lowerBound; inductionValue < *upperBound; inductionValue += *step) {
        ScopedMapBindings<decltype(knowledge.indexValues)> indexBindings(knowledge.indexValues);
        ScopedMapBindings<decltype(knowledge.aliases)> aliasBindings(knowledge.aliases);
        indexBindings.bind(forOp.getInductionVar(), inductionValue);
        for (auto [iterArg, iterValue] : llvm::zip_equal(forOp.getRegionIterArgs(), iterValues))
          aliasBindings.bind(iterArg, iterValue);

        if (failed(executeCompiledCorePlan(*node.loopBody, coreCodeGen, knowledge, resolveWeightSlot)))
          return failure();

        auto yieldOp = cast<mlir::scf::YieldOp>(forOp.getRegion().front().getTerminator());
        for (auto [index, yieldedValue] : llvm::enumerate(yieldOp.getOperands()))
          iterValues[index] = resolveLoopCarriedAlias(yieldedValue, knowledge);
      }
      continue;
    }

    if (node.kind == CompiledCoreNode::Kind::If) {
      auto condition = node.condition.evaluate(knowledge);
      auto ifOp = cast<mlir::scf::IfOp>(node.op);
      if (failed(condition)) {
        ifOp.emitOpError("requires statically evaluable scf.if condition for PIM codegen");
        return failure();
      }

      const auto& selectedBody = *condition != 0 ? node.thenBody : node.elseBody;
      if (selectedBody && failed(executeCompiledCorePlan(*selectedBody, coreCodeGen, knowledge, resolveWeightSlot)))
        return failure();
      continue;
    }

    if (node.kind == CompiledCoreNode::Kind::IndexSwitch) {
      auto selector = node.condition.evaluate(knowledge);
      auto switchOp = cast<mlir::scf::IndexSwitchOp>(node.op);
      if (failed(selector)) {
        switchOp.emitOpError("requires a statically evaluable scf.index_switch selector for PIM codegen");
        return failure();
      }
      const llvm::SmallVectorImpl<CompiledCoreNode>* selectedBody = node.defaultBody.get();
      mlir::Region* selectedRegion = &switchOp.getDefaultRegion();
      for (auto [index, caseValue] : llvm::enumerate(node.caseValues))
        if (caseValue == *selector) {
          selectedBody = node.caseBodies[index].get();
          selectedRegion = &switchOp.getCaseRegions()[index];
          break;
        }
      if (failed(executeCompiledCorePlan(*selectedBody, coreCodeGen, knowledge, resolveWeightSlot)))
        return failure();
      auto yield = cast<mlir::scf::YieldOp>(selectedRegion->front().getTerminator());
      for (auto [result, yielded] : llvm::zip(switchOp.getResults(), yield.getOperands()))
        knowledge.aliases[result] = resolveLoopCarriedAlias(yielded, knowledge);
      continue;
    }

    auto emitBinary = [&](auto op, pim_binary::Opcode opcode) {
      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(),
                                    getVectorByteSizeOrCrash(cast<ShapedType>(op.getInput().getType())),
                                    knowledge,
                                    r2OrImm,
                                    generic1);
    };

    switch (node.opKind) {
    case CompiledCoreOpKind::Load:
      coreCodeGen.codeGenLoadOp(cast<pim::PimMemCopyHostToDevOp>(node.op), knowledge);
      break;
    case CompiledCoreOpKind::Store:
      coreCodeGen.codeGenStoreOp(cast<pim::PimMemCopyDevToHostOp>(node.op), knowledge);
      break;
    case CompiledCoreOpKind::Lmv:     coreCodeGen.codeGenLmvOp(cast<pim::PimMemCopyOp>(node.op), knowledge); break;
    case CompiledCoreOpKind::Receive: coreCodeGen.codeGenReceiveOp(cast<pim::PimReceiveOp>(node.op), knowledge); break;
    case CompiledCoreOpKind::Send:    coreCodeGen.codeGenSendOp(cast<pim::PimSendOp>(node.op), knowledge); break;
    case CompiledCoreOpKind::Concat:  coreCodeGen.codeGenConcatOp(cast<pim::PimConcatOp>(node.op), knowledge); break;
    case CompiledCoreOpKind::Vmm:
      if (auto weightSlot = resolveWeightSlot(cast<pim::PimVMMOp>(node.op), knowledge); succeeded(weightSlot))
        coreCodeGen.codeGenMVMLikeOp<pim::PimVMMOp>(*weightSlot, cast<pim::PimVMMOp>(node.op), true, knowledge);
      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::VSoftmax:
      emitUnary(cast<pim::PimVSoftmaxOp>(node.op), pim_binary::Opcode::vsoftmax, 0, 0);
      break;
    }
  }
  return success();
}

static LogicalResult executeCompiledCoreProgram(
  const CompiledCoreProgram& program,
  PimCodeGen& coreCodeGen,
  const StaticValueKnowledge& initialKnowledge,
  llvm::function_ref<llvm::FailureOr<unsigned>(pim::PimVMMOp, const StaticValueKnowledge&)> resolveWeightSlot) {
  StaticValueKnowledge knowledge = initialKnowledge;
  return executeCompiledCorePlan(program.nodes, coreCodeGen, knowledge, resolveWeightSlot);
}

static SmallDenseMap<memref::GlobalOp, MemEntry, 16>
collectMaterializedHostGlobals(ModuleOp moduleOp, func::FuncOp funcOp, const PimAcceleratorMemory& memory) {
  SmallDenseMap<memref::GlobalOp, MemEntry, 16> materializedHostGlobals;
  funcOp.walk([&](memref::GetGlobalOp getGlobalOp) {
    if (hasWeightAlways(getGlobalOp))
      return;
    auto targetGlobal = lookupGlobalForGetGlobal(moduleOp, getGlobalOp);
    if (!targetGlobal || materializedHostGlobals.contains(targetGlobal))
      return;
    auto it = memory.memEntriesMap.find(getMemoryValueKey(getGlobalOp.getResult()));
    if (it != memory.memEntriesMap.end())
      materializedHostGlobals[targetGlobal] = it->second;
  });
  return materializedHostGlobals;
}

template <typename CoreLikeOpTy>
static void aliasMaterializedHostGlobals(CoreLikeOpTy coreLikeOp,
                                         ModuleOp moduleOp,
                                         const SmallDenseMap<memref::GlobalOp, MemEntry, 16>& materializedHostGlobals,
                                         PimAcceleratorMemory& memory) {
  coreLikeOp.walk([&](memref::GetGlobalOp getGlobalOp) {
    MemoryValueKey key = getMemoryValueKey(getGlobalOp.getResult());
    if (hasWeightAlways(getGlobalOp) || memory.memEntriesMap.contains(key))
      return;

    auto targetGlobal = lookupGlobalForGetGlobal(moduleOp, getGlobalOp);
    if (!targetGlobal)
      return;

    auto it = materializedHostGlobals.find(targetGlobal);
    if (it != materializedHostGlobals.end())
      memory.memEntriesMap[key] = it->second;
  });
}

static OnnxMlirCompilerErrorCodes emitEmptyCoreArtifacts(StringRef outputDirPath, size_t emittedCoreId) {
  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) {
    errs() << "Error while opening core file `" << outputCorePath << "`: " << errorCode.message() << '\n';
    return InvalidOutputFileAccess;
  }

  PimInstructionWriter instructionWriter(coreBinaryStream);
  if (failed(instructionWriter.finalize())) {
    coreBinaryStream.close();
    errs() << "Error while writing core file `" << outputCorePath << "`\n";
    return InvalidOutputFileAccess;
  }
  coreBinaryStream.close();
  if (coreBinaryStream.has_error()) {
    coreBinaryStream.clear_error();
    errs() << "Error while closing core file `" << outputCorePath << "`\n";
    return InvalidOutputFileAccess;
  }

  if (!pimEmitJson.getValue())
    return CompilerSuccess;

  std::string outputCoreJsonPath = (outputDirPath + "/core_" + std::to_string(emittedCoreId) + ".json").str();
  errorCode = std::error_code();
  raw_fd_ostream coreJsonStream(outputCoreJsonPath, errorCode);
  if (errorCode) {
    errs() << "Error while opening core json file `" << outputCoreJsonPath << "`: " << errorCode.message() << '\n';
    return InvalidOutputFileAccess;
  }
  coreJsonStream << "[]";
  coreJsonStream.close();
  return CompilerSuccess;
}

OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::string& outputDirPath) {
  if (!outputDirPath.empty()) {
    if (auto error = sys::fs::create_directory(outputDirPath)) {
      errs() << "Error creating output directory: " << outputDirPath << ": " << error.message() << '\n';
      return InvalidOutputFileAccess;
    }
  }

  auto entryFunc = getPimEntryFunc(moduleOp);
  if (failed(entryFunc))
    return CompilerFailure;
  auto funcOp = *entryFunc;

  PimAcceleratorMemory memory;
  memory.hostMem.allocateHost(moduleOp, funcOp);
  memory.reportHost();

  if (auto err = writeMemoryBinary(moduleOp, funcOp, memory, outputDirPath))
    return err;

  json::Object xbarsPerArrayGroup;
  size_t maxCoreId = 0;
  uint64_t nextBatchReportId = 0;

  SmallVector<Operation*> coreLikeOps = collectTopLevelCoreLikeOps(funcOp);
  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;

  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);
      jobs.push_back(std::move(job));
      continue;
    }

    auto coreBatchOp = cast<pim::PimCoreBatchOp>(op);
    auto batchCoreIds = getBatchCoreIds(coreBatchOp);
    llvm::DenseMap<size_t, SmallVector<unsigned>> lanesByCoreId;
    for (unsigned lane = 0; lane < static_cast<unsigned>(coreBatchOp.getLaneCount()); ++lane)
      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) {
      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;
      jobIndices.push_back(jobs.size());
      jobs.push_back(std::move(job));
    }
    batchJobIndices.push_back(std::move(jobIndices));
    ++nextBatchReportId;
  }

  auto linkCoreWeights =
    [&](size_t coreId, ArrayRef<std::string> weightFiles, json::Array& xbarsPerGroup) -> OnnxMlirCompilerErrorCodes {
    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';
      return InvalidOutputFileAccess;
    }

    for (auto [slot, fileName] : llvm::enumerate(weightFiles)) {
      xbarsPerGroup.push_back(1);
      std::string sourcePath = outputDirPath + "/weights/" + fileName;
      std::string targetPath = coreWeightsDirPath + "/crossbar_" + std::to_string(slot) + ".bin";
      sys::fs::remove(targetPath);
      if (auto error = sys::fs::create_link(sourcePath, targetPath)) {
        errs() << "Error creating link file: " << sourcePath << " to " << targetPath << "\nError:" << error.message()
               << '\n';
        return InvalidOutputFileAccess;
      }
    }

    return CompilerSuccess;
  };

  auto emitJob = [&](const CoreEmissionJob& job) -> CoreEmissionResult {
    CoreEmissionResult result;
    PimAcceleratorMemory jobMemory(memory.memEntriesMap, false);
    llvm::SmallVector<ResolvedWeightView, 8> usedWeights;

    auto resolveWeightSlot = [&](pim::PimVMMOp vmmOp,
                                 const StaticValueKnowledge& knowledge) -> llvm::FailureOr<unsigned> {
      auto weightView = onnx_mlir::resolveWeightView(job.coreLikeOp, vmmOp.getWeight(), knowledge);
      if (failed(weightView)) {
        std::string message;
        llvm::raw_string_ostream os(message);
        os << "requires a statically resolvable dense global weight view during PIM codegen; weight="
           << vmmOp.getWeight() << " type=" << vmmOp.getWeight().getType();
        result.recordDiagnostic(vmmOp, os.str());
        return failure();
      }
      if (weightView->shape.size() != 2) {
        std::string message;
        llvm::raw_string_ostream os(message);
        os << "requires a rank-2 matrix weight view during PIM codegen; resolved shape=[";
        llvm::interleaveComma(weightView->shape, os);
        os << "] weight=" << vmmOp.getWeight() << " type=" << vmmOp.getWeight().getType();
        result.recordDiagnostic(vmmOp, os.str());
        return failure();
      }

      if (auto it = llvm::find(usedWeights, *weightView); it != usedWeights.end())
        return static_cast<unsigned>(std::distance(usedWeights.begin(), it));

      usedWeights.push_back(*weightView);
      return static_cast<unsigned>(usedWeights.size() - 1);
    };

    std::error_code errorCode;
    auto outputCorePath = outputDirPath + "/core_" + std::to_string(job.emittedCoreId) + ".pim";
    raw_fd_ostream coreBinaryStream(outputCorePath, errorCode, sys::fs::OF_None);
    if (errorCode) {
      errs() << "Error while opening core file `" << outputCorePath << "`: " << errorCode.message() << '\n';
      result.status = InvalidOutputFileAccess;
      return result;
    }

    std::unique_ptr<raw_fd_ostream> coreJsonStream;
    if (pimEmitJson.getValue()) {
      std::string outputCoreJsonPath = outputDirPath + "/core_" + std::to_string(job.emittedCoreId) + ".json";
      errorCode = std::error_code();
      coreJsonStream = std::make_unique<raw_fd_ostream>(outputCoreJsonPath, errorCode);
      if (errorCode) {
        errs() << "Error while opening core json file `" << outputCoreJsonPath << "`: " << errorCode.message() << '\n';
        result.status = InvalidOutputFileAccess;
        return result;
      }
      *coreJsonStream << '[';
    }

    PimInstructionWriter instructionWriter(coreBinaryStream);
    PimCodeGen coreCodeGen(jobMemory, instructionWriter, coreJsonStream.get(), emittedCoreIds);
    auto finalizeInstructions = [&]() {
      bool succeeded = mlir::succeeded(instructionWriter.finalize());
      coreBinaryStream.close();
      if (coreBinaryStream.has_error()) {
        coreBinaryStream.clear_error();
        succeeded = false;
      }
      return succeeded;
    };

    if (auto coreOp = dyn_cast<pim::PimCoreOp>(job.coreLikeOp)) {
      aliasMaterializedHostGlobals(coreOp, moduleOp, materializedHostGlobals, jobMemory);
      auto& deviceMemory = jobMemory.getOrCreateDeviceMem(job.emittedCoreId);
      deviceMemory.allocateCore(*job.memoryPlan);

      StaticValueKnowledge knowledge = seedCoreCodegenKnowledge(coreOp);
      if (failed(executeCompiledCoreProgram(*job.program, coreCodeGen, knowledge, resolveWeightSlot))) {
        (void) finalizeInstructions();
        result.status = CompilerFailure;
        return result;
      }
      result.reportRow = deviceMemory.getReportRow();
      result.usedWeights = std::move(usedWeights);
    }
    else {
      auto coreBatchOp = cast<pim::PimCoreBatchOp>(job.coreLikeOp);
      aliasMaterializedHostGlobals(coreBatchOp, moduleOp, materializedHostGlobals, jobMemory);
      auto& deviceMemory = jobMemory.getOrCreateDeviceMem(job.emittedCoreId);

      for (unsigned lane : job.lanes) {
        StaticValueKnowledge knowledge = seedCoreBatchCodegenKnowledge(coreBatchOp, lane);

        deviceMemory.allocateCore(*job.memoryPlan, lane);
        coreCodeGen.setBatchLane(lane);
        if (failed(executeCompiledCoreProgram(*job.program, coreCodeGen, knowledge, resolveWeightSlot))) {
          (void) finalizeInstructions();
          result.status = CompilerFailure;
          return result;
        }
      }

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

    if (!finalizeInstructions()) {
      result.status = InvalidOutputFileAccess;
      return result;
    }

    if (coreJsonStream) {
      coreJsonStream->seek(coreJsonStream->tell() - 1);
      *coreJsonStream << ']';
      coreJsonStream->close();
    }
    return result;
  };

  std::vector<CoreEmissionResult> jobResults(jobs.size());
  mlir::parallelFor(
    moduleOp.getContext(), 0, jobs.size(), [&](size_t index) { jobResults[index] = emitJob(jobs[index]); });

  pim::CappedDiagnosticReporter diagnostics;
  Operation* summaryAnchor = nullptr;
  for (const CoreEmissionResult& result : jobResults) {
    if (!summaryAnchor && !result.diagnostics.empty())
      summaryAnchor = result.diagnostics.front().op;
    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));
  }
  if (diagnostics.hasFailure())
    diagnostics.emitSuppressedSummary(summaryAnchor ? summaryAnchor : moduleOp.getOperation(),
                                      "PIM codegen diagnostic(s)");

  for (size_t jobIndex = 0; jobIndex < jobs.size(); ++jobIndex)
    if (jobResults[jobIndex].status != CompilerSuccess)
      return jobResults[jobIndex].status;

  if (jobs.empty()) {
    if (auto err = emitEmptyCoreArtifacts(outputDirPath, 0))
      return err;
    xbarsPerArrayGroup["core0"] = json::Array {};
    memory.recordCoreReport(0, 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.weights = jobResults[jobIndex].usedWeights;
    weightRequests.push_back(std::move(request));
  }
  auto weightEmission = createAndPopulateWeightFolder(weightRequests, outputDirPath);
  memory.setTotalWeightBytes(weightEmission.totalWeightBytes);
  auto& mapCoreWeightToFileName = weightEmission.mapCoreWeightToFileName;
  for (size_t jobIndex = 0; jobIndex < jobs.size(); ++jobIndex) {
    const CoreEmissionJob& job = jobs[jobIndex];
    const CoreEmissionResult& result = jobResults[jobIndex];
    json::Array xbarsPerGroup;

    if (auto coreOp = dyn_cast<pim::PimCoreOp>(job.coreLikeOp)) {
      if (auto err = linkCoreWeights(job.emittedCoreId, mapCoreWeightToFileName[job.emittedCoreId], xbarsPerGroup))
        return err;
      xbarsPerArrayGroup["core" + std::to_string(job.emittedCoreId)] = std::move(xbarsPerGroup);
      memory.recordCoreReport(job.emittedCoreId, result.reportRow);
      continue;
    }
  }

  for (const SmallVector<size_t>& group : batchJobIndices) {
    SmallVector<int32_t> reportedCoreIds;
    std::optional<MemoryReportRow> batchPerCoreRow;

    for (size_t jobIndex : group) {
      const CoreEmissionJob& job = jobs[jobIndex];
      const CoreEmissionResult& result = jobResults[jobIndex];
      json::Array xbarsPerGroup;
      if (auto err = linkCoreWeights(job.emittedCoreId, mapCoreWeightToFileName[job.emittedCoreId], xbarsPerGroup))
        return err;
      xbarsPerArrayGroup["core" + std::to_string(job.emittedCoreId)] = std::move(xbarsPerGroup);
      reportedCoreIds.push_back(pim::checkedI32OrCrash(job.emittedCoreId, "batch report core id"));
      if (!batchPerCoreRow)
        batchPerCoreRow = result.reportRow;
      else if (!(*batchPerCoreRow == result.reportRow))
        llvm_unreachable("one PIM core batch produced inconsistent per-core memory reports");
    }

    uint64_t batchReportId = jobs[group.front()].batchReportId.value_or(0);
    memory.recordBatchReport(
      batchReportId, reportedCoreIds, batchPerCoreRow.value_or(MemoryReportRow {}));
  }

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

  memory.flushReport();
  return writeConfigJson(funcOp, memory, maxCoreId, std::move(xbarsPerArrayGroup), outputDirPath);
}