#include "TestCheck.hpp"
#include <limits>
#include "src/Accelerators/PIM/Compiler/PimPeriodicAnalysis.hpp"
#include "src/Accelerators/PIM/Compiler/PimInstructionTiming.hpp"
#include "src/Accelerators/PIM/Compiler/PimSpatialCost.hpp"
#include "src/Accelerators/PIM/Compiler/PimBinaryFormat.hpp"
#include "src/Accelerators/PIM/Compiler/PimDependencyTokens.hpp"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/IR/Verifier.h"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"

using namespace onnx_mlir;
using namespace mlir;
using E = PimExecutionEvent;
using O = pim_binary::Opcode;

static PimInstructionTiming hardware(int64_t period) {
  PimInstructionTiming target;
  target.processorCount = 2;
  target.processorPeriodNs = period;
  target.localMemoryPeriodNs = 3;
  target.peerPermissionThenSourceRead = true;
  target.signalIssueNs = period;
  target.controlMessageBytes = 1;
  target.sharedMemoryWidthBytes = 64;
  target.sharedMemoryReadLatencyNs = 80;
  target.sharedMemoryWriteLatencyNs = 120;
  target.sharedMemoryRequestLatencyNs = {2, 3};
  target.sharedMemoryResponseLatencyNs = {2, 3};
  target.interProcessorLatencyNs = {0, 2, 2, 0};
  target.synchronizationRegisterCount = 4;
  return target;
}

static std::string binary(const std::vector<pim_binary::InstructionRecord> &records) {
  std::string result(12, '\0');
  result.replace(0, 4, "PIMB");
  llvm::support::endian::write32le(result.data() + 4, 1);
  llvm::support::endian::write32le(result.data() + 8, records.size());
  for (const auto &record : records) {
    auto encoded = pim_binary::encodeInstructionRecord(record);
    result.append(encoded.data(), encoded.size());
  }
  return result;
}

static void testSearchGuidance(PimInstructionTiming target) {
  PimExecutionProgram program{{{E::Work, 29}, {E::Load, 65}}};
  PimPeriodicOptions options;
  options.verifyPeriodicState = true;
  auto plain = evaluatePimExecution(program, target, options);
  options.searchGuidance = true;
  auto guided = evaluatePimExecution(program, target, options);
  TEST_CHECK(plain.cost.primary == guided.cost.primary);
  TEST_CHECK(plain.cost.alternative == guided.cost.alternative);
  TEST_CHECK(plain.cost.convergenceProved() == guided.cost.convergenceProved());
  TEST_CHECK(guided.cost.processorTaskWorkNs.size() == target.processorCount);
  TEST_CHECK(guided.cost.processorTaskWorkNs.front() == 29);
  TEST_CHECK(guided.cost.processorCriticalDelayNs.size() == target.processorCount);
  TEST_CHECK(guided.cost.processorCriticalDelayNs.front() > 0);
  program.front()[1].constant = true;
  auto constants = evaluatePimExecution(program, target, options);
  TEST_CHECK(constants.cost.primary == plain.cost.primary);
  TEST_CHECK(constants.cost.alternative == plain.cost.alternative);
  TEST_CHECK(constants.cost.processorTaskWorkNs.front() > 29);
}

static void testTypedResult(PimInstructionTiming target) {
  PimPeriodicOptions options;
  options.verifyPeriodicState = true;
  PimExecutionProgram program{{{E::Work, 29}, {E::Load, 65}}};
  auto estimate = evaluatePimExecution(program, target, options);
  TEST_CHECK(estimate.cost.available() && estimate.cost.convergenceProved() == true);
  TEST_CHECK(estimate.report.getString("status") == "estimated");
  TEST_CHECK(estimate.cost.loweringSupport == spatial::StaticLoweringSupport::Unknown);
  TEST_CHECK(!estimate.cost.eligible() && !estimate.cost.selectionMargin);
  auto *queue = estimate.report.getObject("memory_controller");
  for (size_t policy = 0; policy < 2; ++policy) {
    auto *state = queue->getObject(policy ? "alternative_periodic_state" : "periodic_state");
    const auto &proof = estimate.cost.recurrence[policy];
    TEST_CHECK(proof.proved() && state->getString("status") == "proved");
    TEST_CHECK(proof.intervalNs == state->getNumber("interval_ns"));
    TEST_CHECK(proof.periodGenerations == state->getInteger("period_generations"));
  }
  auto primary = estimate.cost.primary, alternative = estimate.cost.alternative;
  estimate.report = llvm::json::Object{{"status", "unavailable"}, {"convergence_proved", false}};
  TEST_CHECK(estimate.cost.primary == primary && estimate.cost.alternative == alternative);
  TEST_CHECK(estimate.cost.convergenceProved() == true);
  options.verifyPeriodicState = false;
  estimate = evaluatePimExecution(program, target, options);
  TEST_CHECK(estimate.cost.available() && !estimate.cost.convergenceProved().has_value());
  TEST_CHECK(estimate.report.getString("status") == "estimated");
  TEST_CHECK(!estimate.report.getObject("memory_controller")->getObject("periodic_state"));
  options.intervalWindowEventBudget = 0;
  for (bool detailed : {false, true}) {
    options.detailed = detailed;
    estimate = evaluatePimExecution(program, target, options);
    TEST_CHECK(estimate.queue && !estimate.cost.available());
    TEST_CHECK(!estimate.cost.convergenceProved().has_value());
    TEST_CHECK(estimate.report.getString("status") == "unavailable");
    TEST_CHECK(estimate.report.getString("reason") == estimate.cost.reason);
    TEST_CHECK(estimate.report.getObject("memory_controller")->getString("status") == "unavailable");
    TEST_CHECK(!estimate.report.getObject("memory_controller")->getObject("periodic_state"));
  }
  options.intervalWindowEventBudget.reset();
  options.detailed = false;
  options.verifyPeriodicState = true;
  estimate = evaluatePimExecution({{{E::Wait, 0, 0, 0, 1}, {E::Load, 65}}}, target, options);
  TEST_CHECK(!estimate.cost.available() && estimate.cost.convergenceProved() == false);
  TEST_CHECK(!estimate.cost.reason.empty() && !estimate.cost.eligible());
  target.sharedMemoryWidthBytes = 0;
  estimate = evaluatePimExecution(program, target, options);
  TEST_CHECK(!estimate.cost.available() && !estimate.cost.reason.empty());
  TEST_CHECK(estimate.report.getString("status") == "unavailable");
  TEST_CHECK(estimate.report.getString("reason") == estimate.cost.reason);
}

static void testLateRecurrenceAuthority(PimInstructionTiming target) {
  PimPeriodicOptions options;
  options.verifyPeriodicState = true;
  options.intervalWindowEventBudget = 1;
  PimExecutionProgram program{{{E::Work, 29}, {E::Load, 65}}};
  auto late = evaluatePimExecution(program, target, options);
  TEST_CHECK(late.queue && late.cost.available() && late.cost.convergenceProved() == true);
  for (const auto &policy : late.queue->policies) {
    TEST_CHECK(policy.legacy.termination == PimQueueTermination::EventBudget);
    TEST_CHECK(!policy.legacy.intervalNs && policy.recurrence.proved());
    TEST_CHECK(policy.proofProcessedEvents > policy.legacy.progress.processedEvents);
    TEST_CHECK(policy.intervalNs() == policy.recurrence.intervalNs);
    TEST_CHECK(policy.termination == PimQueueTermination::RecurrenceProved);
  }
  auto first = serializePimQueueEvaluation(*late.queue);
  auto second = serializePimQueueEvaluation(*late.queue);
  TEST_CHECK(first.getString("status") == "estimated" && !first.getString("reason"));
  TEST_CHECK(first.getString("alternative_tie_status") == "estimated");
  TEST_CHECK(first.getNumber("interval_ns") == late.cost.primary);
  TEST_CHECK(first.getNumber("alternative_tie_interval_ns") == late.cost.alternative);
  TEST_CHECK(first.getObject("legacy_interval_window")->getString("status") == "unavailable");
  TEST_CHECK(first.getObject("legacy_interval_window")->getString("termination") == "event-budget");
  TEST_CHECK(llvm::json::Value(std::move(first)) == llvm::json::Value(std::move(second)));

  auto blocked = evaluatePimExecution({{{E::Wait, 0, 0, 0, 1}, {E::Signal, 0, 1}},
    {{E::Wait, 0, 0, 0, 1}, {E::Signal, 0, 0}}}, target, options);
  TEST_CHECK(blocked.queue && blocked.cost.convergenceProved() == false);
  TEST_CHECK(blocked.queue->policies[1].termination == PimQueueTermination::NoReadyEvents);
  for (size_t failedPolicy = 0; failedPolicy < 2; ++failedPolicy) {
    auto asymmetric = *late.queue;
    asymmetric.policies[failedPolicy] = blocked.queue->policies[failedPolicy];
    auto cost = asymmetric.cost();
    cost.loweringSupport = spatial::StaticLoweringSupport::Supported;
    cost.selectionMargin = 0.021;
    TEST_CHECK(cost.primary.has_value() == (failedPolicy != 0));
    TEST_CHECK(cost.alternative.has_value() == (failedPolicy != 1));
    TEST_CHECK(cost.convergenceProved() == false && !cost.eligible());
    TEST_CHECK(cost.reason == blocked.queue->policies[failedPolicy].recurrence.reason);
    auto report = serializePimQueueEvaluation(asymmetric);
    TEST_CHECK(report.getString("status") == (failedPolicy == 0 ? "unavailable" : "estimated"));
    TEST_CHECK(report.getString("alternative_tie_status") == (failedPolicy == 1 ? "unavailable" : "estimated"));
    TEST_CHECK(report.getObject(failedPolicy == 0 ? "periodic_state" : "alternative_periodic_state")
      ->getString("status") == "unproved");
  }

  options.intervalWindowEventBudget.reset();
  auto normal = evaluatePimExecution(program, target, options);
  options.workBreakdown = true;
  auto detailed = evaluatePimExecution(program, target, options);
  TEST_CHECK(late.cost.primary == normal.cost.primary && normal.cost.primary == detailed.cost.primary);
  TEST_CHECK(late.cost.alternative == normal.cost.alternative && normal.cost.alternative == detailed.cost.alternative);
  TEST_CHECK(normal.queue->policies[0].progress.processedEvents
    <= detailed.queue->policies[0].progress.processedEvents);
}

static void testSharedEngine(PimInstructionTiming target) {
  auto record = [](O op, int32_t peer = 0, int32_t reg = 0, int32_t count = 0, int32_t bytes = 0) {
    pim_binary::InstructionRecord r;
    r.opcode = op; r.r2OrImm = peer; r.generic1 = reg; r.generic2 = count; r.generic3 = bytes;
    return r;
  };
  auto vectorWork = getPimVectorWorkNs(target, 65, 1, 1, true);
  PimInstructionProgram instructions;
  instructions.instructions = {
    binary({record(O::setbw, 0, 8, 8), record(O::ld, 0, 0, 0, 65),
      record(O::vvadd, 0, 0, 0, 65), record(O::send, 1, 0, 0, 65), record(O::sync, 1, 0)}),
    binary({record(O::recv, 0, 0, 0, 65), record(O::wait, 0, 0, 1), record(O::st, 0, 0, 0, 33)})};
  instructions.arrayGroups.resize(2);
  PimExecutionProgram abstract(2);
  abstract[0] = {{E::Work, target.processorPeriodNs}, {E::Load, 65}, {E::Work, vectorWork},
    {E::Send, 65, 1}, {E::Signal, 0, 1, 0}};
  abstract[1] = {{E::Receive, 65, 0}, {E::Wait, 0, 0, 0, 1}, {E::Store, 33}};
  PimPeriodicOptions details;
  details.workBreakdown = true;
  PimExecutionProgram trace;
  details.executionTrace = &trace;
  auto reference = analyzePimInstructions(instructions, target, 2, details);
  TEST_CHECK(reference.getString("status") == "estimated");
  TEST_CHECK(!reference.getObject("memory_controller")->getObject("periodic_state"));
  details.executionTrace = nullptr;
  auto roundTrip = analyzePimExecution(trace, target, details);
  for (const char *field : {"interval_ns", "alternative_tie_interval_ns"})
    TEST_CHECK(roundTrip.getObject("memory_controller")->getNumber(field)
      == reference.getObject("memory_controller")->getNumber(field));
  for (size_t resource = 0; resource < 7; ++resource)
    TEST_CHECK((*roundTrip.getArray("resource_service_ns"))[resource].getAsInteger()
      == (*reference.getArray("resource_service_ns"))[resource].getAsInteger());
  details.verifyPeriodicState = true;
  auto proof = analyzePimExecution(trace, target, details);
  TEST_CHECK(proof.getObject("memory_controller")->getString("status")
    == reference.getObject("memory_controller")->getString("status"));
  TEST_CHECK(proof.getObject("memory_controller")->getObject("legacy_interval_window")->getNumber("interval_ns")
    == reference.getObject("memory_controller")->getNumber("interval_ns"));
  TEST_CHECK(proof.getObject("memory_controller")->getNumber("interval_ns")
    == proof.getObject("memory_controller")->getObject("periodic_state")->getNumber("interval_ns"));
  TEST_CHECK(proof.getObject("memory_controller")->getObject("periodic_state"));
  details.verifyPeriodicState = false;
  auto actual = analyzePimExecution(abstract, target, details);
  for (size_t core = 0; core < 2; ++core)
    TEST_CHECK((*actual.getArray("request_order"))[core].getAsObject()->getString("identity")
      == (*reference.getArray("request_order"))[core].getAsObject()->getString("identity"));
  for (size_t core = 0; core < 2; ++core)
    TEST_CHECK((*actual.getArray("request_order"))[core].getAsObject()->getString("blocking_identity")
      == (*reference.getArray("request_order"))[core].getAsObject()->getString("blocking_identity"));
  for (const char *field : {"global_read_bytes", "global_write_bytes", "global_read_requests",
      "global_write_requests", "global_memory_service_ns"})
    TEST_CHECK(actual.getInteger(field) == reference.getInteger(field));
  auto *a = actual.getObject("memory_controller"), *b = reference.getObject("memory_controller");
  TEST_CHECK(a && b && a->getString("status") == "estimated");
  TEST_CHECK(a->getNumber("interval_ns") == b->getNumber("interval_ns"));
  TEST_CHECK(a->getNumber("alternative_tie_interval_ns") == b->getNumber("alternative_tie_interval_ns"));
  abstract[1].clear();
  abstract[0] = {{E::Load, 1}, {E::Load, 65}};
  auto ordered = analyzePimExecution(abstract, target, details);
  abstract[0] = {{E::Load, 65}, {E::Load, 1}};
  auto reversed = analyzePimExecution(abstract, target, details);
  TEST_CHECK(ordered.getInteger("global_read_bytes") == reversed.getInteger("global_read_bytes"));
  TEST_CHECK((*ordered.getArray("request_order"))[0].getAsObject()->getString("identity")
    != (*reversed.getArray("request_order"))[0].getAsObject()->getString("identity"));
  // Preserve repeated request order, not a single larger request.
  abstract[0] = {{E::Repeat, 2, 0, 0, 0, {{E::Load, 1}, {E::Work, 4}}}};
  abstract[1].clear();
  actual = analyzePimExecution(abstract, target);
  TEST_CHECK(actual.getInteger("global_read_requests") == 2);
  TEST_CHECK(actual.getInteger("global_memory_service_ns") == 160);
  abstract[0] = {{E::Repeat, 1000000000000ull, 0, 0, 0, {{E::Load, 1}}}};
  actual = analyzePimExecution(abstract, target);
  TEST_CHECK(actual.getString("status") == "unavailable");
  abstract[0] = {{E::Repeat, 1000000000000ull, 0, 0, 0, {{E::Work, 4}}}};
  actual = analyzePimExecution(abstract, target);
  TEST_CHECK(actual.getObject("memory_controller")->getNumber("interval_ns") == 4000000000000.0);
  abstract[0] = {{E::Signal, 0, int64_t(1) << 32}};
  TEST_CHECK(analyzePimExecution(abstract, target).getString("status") == "unavailable");
  TEST_CHECK(getPimMatrixOutputBytes(target, 2, pim_binary::kMatrixWeightPrecisionBits, 3, 1)
    == target.matrixColumns * 2 / pim_binary::kMatrixWeightPrecisionBits * 3);
  E copy{E::Work, 4000};
  copy.workClass = E::LocalCopy;
  abstract[0] = {copy, {E::Load, 1}};
  auto charged = analyzePimExecution(abstract, target);
  PimPeriodicOptions sensitivity;
  sensitivity.ignoreLocalCopies = true;
  auto omitted = analyzePimExecution(abstract, target, sensitivity);
  TEST_CHECK(*charged.getObject("memory_controller")->getNumber("interval_ns")
    - *omitted.getObject("memory_controller")->getNumber("interval_ns") == 4000);
}

static void testAccountingAndOverflow(PimInstructionTiming target) {
  auto record = [](O op, int64_t bytes) {
    pim_binary::InstructionRecord r;
    r.opcode = op; r.generic3 = bytes;
    return r;
  };
  PimPeriodicOptions options;
  options.scalarIssue = false;
  options.workBreakdown = true;
  auto setbw = record(O::setbw, 0);
  setbw.generic1 = setbw.generic2 = 8;
  auto matrixInstruction = record(O::mvmul, 0);
  matrixInstruction.r2OrImm = pim_binary::kMatrixWeightPrecisionBits;
  PimInstructionProgram reference;
  reference.instructions = {binary({setbw, matrixInstruction, record(O::vvadd, 65), record(O::vmv, 3), record(O::lmv, 65), record(O::lldi, 65)})};
  reference.arrayGroups = {{3}};
  E matrix{E::Work, getPimMatrixWorkNs(target, 8,
    getPimMatrixOutputBytes(target, 2, pim_binary::kMatrixWeightPrecisionBits, 3, 1))};
  matrix.workClass = E::Matrix;
  E vector{E::Work, getPimVectorWorkNs(target, 65, 1, 1, true)};
  vector.workClass = E::Vector;
  E vectorCopy{E::Work, getPimVectorWorkNs(target, 3, 1, 1, false)};
  vectorCopy.workClass = E::VectorCopy;
  E copy{E::Work, target.getLocalMemoryServiceNs(65, false)
    + target.getLocalMemoryServiceNs(65, true) + target.processorPeriodNs};
  copy.workClass = E::LocalCopy;
  E fill{E::Work, target.getLocalMemoryServiceNs(65, true) + target.processorPeriodNs};
  fill.workClass = E::Fill;
  auto exact = analyzePimInstructions(reference, target, 2, options);
  auto abstract = analyzePimExecution({{matrix, vector, vectorCopy, copy, fill}}, target, options);
  TEST_CHECK(*abstract.getArray("resource_service_ns") == *exact.getArray("resource_service_ns"));
  TEST_CHECK(abstract.getObject("memory_controller")->getNumber("interval_ns")
    == exact.getObject("memory_controller")->getNumber("interval_ns"));
  auto repeated = analyzePimExecution({{{E::Repeat, 7, 0, 0, 0, {matrix, vector, vectorCopy, copy, fill}}}}, target, options);
  for (size_t i = 0; i < 7; ++i)
    TEST_CHECK(*(*repeated.getArray("resource_service_ns"))[i].getAsInteger()
      == 7 * *(*abstract.getArray("resource_service_ns"))[i].getAsInteger());
  auto nested = analyzePimExecution({{{E::Repeat, 1000000, 0, 0, 0,
    {{E::Repeat, 1000000, 0, 0, 0, {fill}}}}}}, target, options);
  TEST_CHECK(*(*nested.getArray("resource_service_ns"))[2].getAsInteger()
    == static_cast<int64_t>(1000000000000ull * fill.amount));
  TEST_CHECK(nested.getInteger("expanded_graph_events").value() < 10);
  PimExecutionProgram pattern(2), expanded(2);
  pattern[0] = {{E::Repeat, 3, 0, 0, 0, {{E::Load, 65},
    {E::Repeat, 5, 0, 0, 0, {matrix, vector, copy, fill}}}}};
  for (size_t repeat = 0; repeat < 3; ++repeat) {
    expanded[0].push_back({E::Load, 65});
    for (size_t work = 0; work < 5; ++work)
      expanded[0].insert(expanded[0].end(), {matrix, vector, copy, fill});
  }
  pattern[1] = expanded[1] = {{E::Work, 311}, {E::Load, 257}};
  auto compactCost = analyzePimExecution(pattern, target, options);
  auto expandedCost = analyzePimExecution(expanded, target, options);
  for (const char *policy : {"interval_ns", "alternative_tie_interval_ns"})
    TEST_CHECK(compactCost.getObject("memory_controller")->getNumber(policy)
      == expandedCost.getObject("memory_controller")->getNumber(policy));
  TEST_CHECK(*compactCost.getArray("resource_service_ns") == *expandedCost.getArray("resource_service_ns"));
  std::vector<E> body{{E::Load, 1}};
  for (size_t i = 0; i < 10000; ++i) {
    E event{E::Work, 1}; event.workClass = i % 2 ? E::Vector : E::Fill;
    body.push_back(event);
  }
  auto cachedWork = analyzePimExecution({{{E::Repeat, 10000, 0, 0, 0, body}}}, target, options);
  TEST_CHECK(cachedWork.getInteger("global_read_requests") == 10000);
  TEST_CHECK(cachedWork.getInteger("expanded_graph_events").value() < 100000);
  TEST_CHECK((*cachedWork.getArray("resource_service_ns"))[1].getAsInteger() == 50000000);
  TEST_CHECK((*cachedWork.getArray("resource_service_ns"))[2].getAsInteger() == 50000000);
  uint64_t maximum = std::numeric_limits<uint64_t>::max();
  for (auto program : std::vector<PimExecutionProgram>{
        {{{E::Work, maximum}}}, {{{E::Load, maximum}}},
        {{{E::Repeat, maximum, 0, 0, 0, {fill}}}},
        {{{E::Work, uint64_t(1) << 62}, {E::Work, uint64_t(1) << 62}}}})
    TEST_CHECK(analyzePimExecution(program, target).getString("status") == "unavailable");
  target.sharedMemoryReadLatencyNs = maximum;
  TEST_CHECK(analyzePimExecution({{{E::Load, 65}}}, target).getString("status") == "unavailable");
}

static PimPeriodicOptions spatialDetails() {
  PimPeriodicOptions options;
  options.workBreakdown = true;
  return options;
}

static void testCoreAssignments() {
  MLIRContext context;
  context.loadDialect<func::FuncDialect, spatial::SpatialDialect>();
  OwningOpRef<ModuleOp> module = ModuleOp::create(UnknownLoc::get(&context));
  OpBuilder builder(module->getBodyRegion());
  auto loc = module->getLoc();
  auto function = func::FuncOp::create(builder, loc, "test", builder.getFunctionType({}, {}));
  builder.setInsertionPointToStart(function.addEntryBlock());
  auto empty = [&](std::optional<int32_t> id) {
    auto compute = spatial::SpatScheduledCompute::create(builder, loc, TypeRange{}, ValueRange{}, ValueRange{});
    if (id) compute->setAttr("coreId", builder.getI32IntegerAttr(*id));
    compute.getBody().push_back(new Block);
    builder.setInsertionPointToStart(&compute.getBody().front());
    spatial::SpatYieldOp::create(builder, loc, ValueRange{});
    builder.setInsertionPointAfter(compute);
    return compute;
  };
  auto first = empty(0);
  auto second = empty(1);
  func::ReturnOp::create(builder, loc);
  TEST_CHECK(succeeded(verify(*module)));
  auto target = hardware(1);
  auto reason = [&]() {
    auto report = analyzePimSpatialSchedule(*module, target);
    return report.getString("reason").value_or("").str();
  };
  TEST_CHECK(reason().find("physical core") == std::string::npos);
  second->setAttr("coreId", builder.getI32IntegerAttr(0));
  TEST_CHECK(reason() == "duplicate physical core assignment: 0");
  TEST_CHECK(evaluatePimSpatialSchedule(*module, target).cost.loweringSupport
    == spatial::StaticLoweringSupport::Unsupported);
  second->removeAttr("coreId");
  TEST_CHECK(reason() == "missing physical core");
  TEST_CHECK(evaluatePimSpatialSchedule(*module, target).cost.loweringSupport
    == spatial::StaticLoweringSupport::Unknown);
  for (int32_t id : {-1, 2}) {
    second->setAttr("coreId", builder.getI32IntegerAttr(id));
    TEST_CHECK(reason() == "physical core ID is outside the target: " + std::to_string(id));
  }
  second.erase();
  builder.setInsertionPoint(function.getBody().front().getTerminator());
  auto batch = spatial::SpatScheduledComputeBatch::create(builder, loc, TypeRange{},
    builder.getI32IntegerAttr(2), ValueRange{}, ValueRange{});
  batch.getBody().push_back(new Block);
  batch.getBody().front().addArgument(builder.getIndexType(), loc);
  builder.setInsertionPointToStart(&batch.getBody().front());
  spatial::SpatYieldOp::create(builder, loc, ValueRange{});
  TEST_CHECK(reason() == "missing physical batch cores");
  batch->setAttr("coreIds", builder.getDenseI32ArrayAttr({1, 1}));
  TEST_CHECK(reason() == "duplicate physical core assignment: 1");
  batch->setAttr("coreIds", builder.getDenseI32ArrayAttr({1, 0}));
  TEST_CHECK(reason() == "duplicate physical core assignment: 0");
  TEST_CHECK(evaluatePimSpatialSchedule(*module, target).cost.loweringSupport
    == spatial::StaticLoweringSupport::Unsupported);
  first.erase();
  batch->setAttr("coreIds", builder.getDenseI32ArrayAttr({1, 2}));
  TEST_CHECK(reason() == "physical core ID is outside the target: 2");
}

static void testDependencyTokens() {
  PimDependencyTokens fresh, drained;
  TEST_CHECK(fresh.empty() && fresh.remaining().empty());
  drained.initialize(1);
  TEST_CHECK(drained.consume() == 0);
  TEST_CHECK(drained.empty() && drained.remaining().empty());
  // Empty queues have identical future behavior, regardless of their capacity.
  for (uint64_t time : {0, 7, 7, 123}) {
    fresh.push(time); drained.push(time);
    TEST_CHECK(fresh.remaining() == drained.remaining());
    TEST_CHECK(fresh.consume() == drained.consume());
    TEST_CHECK(fresh.empty() && drained.empty());
  }
  // Consumed prefixes and compaction cannot change FIFO contents or readiness.
  drained.initialize(3);
  TEST_CHECK(drained.consume() == 0);
  drained.push(11);
  TEST_CHECK(drained.consume() == 0);
  TEST_CHECK(drained.consume() == 0);
  drained.push(11);
  drained.push(12);
  for (uint64_t time : {11, 11, 12}) TEST_CHECK(drained.consume() == time);
  TEST_CHECK(drained.empty());
  drained.initialize(0);
  TEST_CHECK(drained.empty() && drained.remaining().empty());
}

static void testReadOnlySpatial(PimInstructionTiming target, bool batched) {
  MLIRContext context;
  context.loadDialect<arith::ArithDialect, func::FuncDialect, scf::SCFDialect,
    tensor::TensorDialect, linalg::LinalgDialect, spatial::SpatialDialect>();
  auto module = ModuleOp::create(UnknownLoc::get(&context));
  OpBuilder builder(module.getBodyRegion());
  auto loc = module.getLoc();
  auto type = RankedTensorType::get({1, 64}, builder.getF32Type());
  auto function = func::FuncOp::create(builder, loc, "test", builder.getFunctionType({}, {}));
  builder.setInsertionPointToStart(function.addEntryBlock());
  Value constant = arith::ConstantOp::create(builder, loc,
    DenseElementsAttr::get(type, builder.getF32FloatAttr(2)));
  Operation *compute;
  if (batched) {
    compute = spatial::SpatScheduledComputeBatch::create(builder, loc, TypeRange{},
      builder.getI32IntegerAttr(2), ValueRange{}, ValueRange{constant});
    compute->setAttr("coreIds", builder.getDenseI32ArrayAttr({0, 1}));
  } else {
    compute = spatial::SpatScheduledCompute::create(builder, loc, TypeRange{}, ValueRange{}, ValueRange{constant});
    compute->setAttr("coreId", builder.getI32IntegerAttr(0));
  }
  Block *body = new Block;
  compute->getRegion(0).push_back(body);
  if (batched) body->addArgument(builder.getIndexType(), loc);
  Value input = body->addArgument(type, loc);
  builder.setInsertionPointToStart(body);
  Value zero = arith::ConstantIndexOp::create(builder, loc, 0);
  Value upper = arith::ConstantIndexOp::create(builder, loc, 28);
  Value one = arith::ConstantIndexOp::create(builder, loc, 1);
  Value local = tensor::EmptyOp::create(builder, loc, type.getShape(), type.getElementType());
  auto loop = scf::ForOp::create(builder, loc, zero, upper, one);
  builder.setInsertionPointToStart(loop.getBody());
  auto add = spatial::SpatVAddOp::create(builder, loc, type, local, input);
  builder.setInsertionPointAfter(loop);
  spatial::SpatWaitOp::create(builder, loc, zero, zero);
  spatial::SpatYieldOp::create(builder, loc, ValueRange{});
  builder.setInsertionPointAfter(compute);
  func::ReturnOp::create(builder, loc);
  TEST_CHECK(succeeded(verify(module)));
  std::string before, after;
  llvm::raw_string_ostream os(before); module.print(os);
  auto report = analyzePimSpatialSchedule(module, target, spatialDetails());
  auto fast = analyzePimSpatialSchedule(module, target);
  PimPeriodicOptions proofOptions;
  proofOptions.verifyPeriodicState = true;
  auto typed = evaluatePimSpatialSchedule(module, target, proofOptions);
  TEST_CHECK(typed.cost.available() && typed.cost.convergenceProved() == true);
  TEST_CHECK(typed.cost.loweringSupport == spatial::StaticLoweringSupport::Supported);
  TEST_CHECK(typed.report.getBoolean("lowering_supported") == typed.cost.loweringSupported());
  TEST_CHECK(typed.report.getBoolean("convergence_proved") == typed.cost.convergenceProved());
  TEST_CHECK(!fast.getArray("request_order") && !fast.getArray("resource_service_ns"));
  TEST_CHECK(!fast.getObject("without_local_copy_work") && !fast.getArray("work_class_ns"));
  for (const char *field : {"interval_ns", "alternative_tie_interval_ns"})
    TEST_CHECK(fast.getObject("memory_controller")->getNumber(field)
      == report.getObject("memory_controller")->getNumber(field));
  TEST_CHECK(report.getBoolean("safe_for_selection") == false);
  TEST_CHECK(report.getInteger("global_read_requests") == (batched ? 2 : 28));
  TEST_CHECK(report.getInteger("global_read_bytes") == (batched ? 2 : 28) * 64);
  TEST_CHECK(report.getInteger("constant_read_requests") == (batched ? 2 : 28));
  llvm::raw_string_ostream out(after); module.print(out);
  TEST_CHECK(before == after);
  add->setAttr("test.annotation", builder.getUnitAttr());
  report = analyzePimSpatialSchedule(module, target, spatialDetails());
  TEST_CHECK(report.getInteger("global_read_requests") == (batched ? 2 : 28));
  builder.setInsertionPoint(body->getTerminator());
  Value source = tensor::EmptyOp::create(builder, loc, ArrayRef<int64_t>{1, 3, 1, 224}, builder.getF32Type());
  Value destination = tensor::EmptyOp::create(builder, loc, ArrayRef<int64_t>{1, 1, 224, 3}, builder.getF32Type());
  linalg::TransposeOp::create(builder, loc, source, destination, ArrayRef<int64_t>{0, 2, 3, 1});
  TEST_CHECK(succeeded(verify(module)));
  report = analyzePimSpatialSchedule(module, target, spatialDetails());
  auto expectedCopyWork = 224 * getPimVectorWorkNs(target, 3, 1, 1, false) * (batched ? 2 : 1);
  auto *classes = report.getArray("work_class_ns");
  auto *columns = report.getArray("work_class_columns");
  TEST_CHECK(classes && columns && classes->size() == 6 && columns->size() == 6);
  const char *names[] = {"matrix", "vector", "local_copy", "fill", "vector_copy", "scalar_issue"};
  uint64_t totalWork = 0;
  for (size_t i = 0; i < classes->size(); ++i) {
    TEST_CHECK((*columns)[i].getAsString() == names[i]);
    totalWork += *(*classes)[i].getAsInteger();
  }
  TEST_CHECK((*classes)[E::VectorCopy].getAsInteger() == expectedCopyWork);
  TEST_CHECK((*classes)[E::LocalCopy].getAsInteger() == 0);
  TEST_CHECK((*classes)[E::ScalarIssue].getAsInteger() == 0);
  TEST_CHECK(totalWork == report.getInteger("compute_and_local_work_ns"));
  auto softmax = spatial::SpatSoftmaxOp::create(builder, loc, type, local);
  auto secondSoftmax = spatial::SpatSoftmaxOp::create(builder, loc, type, softmax.getOutput());
  TEST_CHECK(succeeded(verify(module)));
  auto softmaxReport = analyzePimSpatialSchedule(module, target, spatialDetails());
  TEST_CHECK(softmaxReport.getString("status") == "estimated");
  TEST_CHECK(softmaxReport.getInteger("compute_and_local_work_ns") == totalWork);
  TEST_CHECK(softmaxReport.getInteger("global_read_requests") == report.getInteger("global_read_requests"));
  secondSoftmax.erase(); softmax.erase();
  source = tensor::EmptyOp::create(builder, loc, ArrayRef<int64_t>{1, 0, 1, 224}, builder.getF32Type());
  destination = tensor::EmptyOp::create(builder, loc, ArrayRef<int64_t>{1, 1, 224, 0}, builder.getF32Type());
  linalg::TransposeOp::create(builder, loc, source, destination, ArrayRef<int64_t>{0, 2, 3, 1});
  TEST_CHECK(succeeded(verify(module)));
  report = analyzePimSpatialSchedule(module, target, spatialDetails());
  TEST_CHECK(report.getString("reason") == "zero-element transpose has no abstract timing contract");
  module.erase();
}

static void testHostSlice(PimInstructionTiming target) {
  MLIRContext context;
  context.loadDialect<arith::ArithDialect, func::FuncDialect, scf::SCFDialect,
    tensor::TensorDialect, spatial::SpatialDialect>();
  auto module = ModuleOp::create(UnknownLoc::get(&context));
  OpBuilder builder(module.getBodyRegion());
  auto loc = module.getLoc();
  auto type = RankedTensorType::get({1, 3, 4, 4}, builder.getF32Type());
  auto function = func::FuncOp::create(builder, loc, "test", builder.getFunctionType({type}, {}));
  builder.setInsertionPointToStart(function.addEntryBlock());
  auto batch = spatial::SpatScheduledComputeBatch::create(builder, loc, TypeRange{},
    builder.getI32IntegerAttr(2), ValueRange{}, ValueRange{function.getArgument(0)});
  batch->setAttr("coreIds", builder.getDenseI32ArrayAttr({0, 1}));
  Block *body = new Block;
  batch.getBody().push_back(body);
  body->addArgument(builder.getIndexType(), loc);
  Value input = body->addArgument(type, loc);
  builder.setInsertionPointToStart(body);
  auto offsets = getAsOpFoldResult(builder.getI64ArrayAttr({0, 0, 1, 0}));
  auto sizes = getAsOpFoldResult(builder.getI64ArrayAttr({1, 3, 1, 4}));
  auto strides = getAsOpFoldResult(builder.getI64ArrayAttr({1, 1, 1, 1}));
  auto slice = tensor::ExtractSliceOp::create(builder, loc, input, offsets, sizes, strides);
  spatial::SpatReluOp::create(builder, loc, slice.getType(), slice.getResult());
  spatial::SpatYieldOp::create(builder, loc, ValueRange{});
  builder.setInsertionPointAfter(batch);
  func::ReturnOp::create(builder, loc);
  TEST_CHECK(succeeded(verify(module)));
  auto report = analyzePimSpatialSchedule(module, target, spatialDetails());
  TEST_CHECK(report.getInteger("global_read_requests") == 6);
  TEST_CHECK(report.getInteger("global_read_bytes") == 24);
  module.erase();
}

static void testCapturedVectorWeights(PimInstructionTiming target) {
  MLIRContext context;
  context.loadDialect<arith::ArithDialect, func::FuncDialect, tensor::TensorDialect,
    spatial::SpatialDialect>();
  auto module = ModuleOp::create(UnknownLoc::get(&context));
  OpBuilder builder(module.getBodyRegion());
  auto loc = module.getLoc();
  auto type = RankedTensorType::get({1, 64}, builder.getF32Type());
  auto function = func::FuncOp::create(builder, loc, "test", builder.getFunctionType({}, {}));
  builder.setInsertionPointToStart(function.addEntryBlock());
  Value zero = arith::ConstantOp::create(builder, loc,
    DenseElementsAttr::get(type, builder.getF32FloatAttr(0)));
  auto compute = spatial::SpatScheduledCompute::create(builder, loc, TypeRange{},
    ValueRange{zero}, ValueRange{});
  compute->setAttr("coreId", builder.getI32IntegerAttr(0));
  Block *body = new Block;
  compute.getBody().push_back(body);
  Value weight = body->addArgument(type, loc);
  builder.setInsertionPointToStart(body);
  Value local = tensor::EmptyOp::create(builder, loc, type.getShape(), type.getElementType());
  spatial::SpatVAddOp::create(builder, loc, type, local, weight);
  spatial::SpatYieldOp::create(builder, loc, ValueRange{});
  builder.setInsertionPointAfter(compute);
  func::ReturnOp::create(builder, loc);
  TEST_CHECK(succeeded(verify(module)));
  auto report = analyzePimSpatialSchedule(module, target, spatialDetails());
  // Captured weight operands stay block arguments; codegen cannot fold this load to LLDI.
  TEST_CHECK(report.getInteger("global_read_requests") == 1);
  TEST_CHECK(report.getInteger("constant_read_bytes") == 64);
  module.erase();
}

static void testPublicationAndOutput(PimInstructionTiming target) {
  MLIRContext context;
  context.loadDialect<arith::ArithDialect, func::FuncDialect, tensor::TensorDialect,
    spatial::SpatialDialect>();
  auto module = ModuleOp::create(UnknownLoc::get(&context));
  OpBuilder builder(module.getBodyRegion());
  auto loc = module.getLoc();
  auto inputType = RankedTensorType::get({1, 2, 3, 5}, builder.getF32Type());
  auto outputType = RankedTensorType::get({1, 4, 3, 5}, builder.getF32Type());
  auto function = func::FuncOp::create(builder, loc, "test", builder.getFunctionType({}, {outputType}));
  builder.setInsertionPointToStart(function.addEntryBlock());
  auto compute = spatial::SpatScheduledCompute::create(builder, loc, TypeRange{outputType},
    ValueRange{}, ValueRange{});
  compute->setAttr("coreId", builder.getI32IntegerAttr(0));
  compute.getBody().push_back(new Block);
  builder.setInsertionPointToStart(&compute.getBody().front());
  Value zero = arith::ConstantIndexOp::create(builder, loc, 0);
  Value one = arith::ConstantIndexOp::create(builder, loc, 1);
  Value input = tensor::EmptyOp::create(builder, loc, inputType.getShape(), inputType.getElementType());
  auto concat = spatial::SpatConcatOp::create(builder, loc, outputType,
    builder.getI64IntegerAttr(1), ValueRange{input, input});
  auto slice = tensor::ExtractSliceOp::create(builder, loc, concat.getOutput(),
    getAsOpFoldResult(builder.getI64ArrayAttr({0, 0, 1, 0})),
    getAsOpFoldResult(builder.getI64ArrayAttr({1, 4, 1, 5})),
    getAsOpFoldResult(builder.getI64ArrayAttr({1, 1, 1, 1})));
  spatial::SpatHostStoreSyncOp::create(builder, loc, zero, one, zero, zero, slice.getResult());
  spatial::SpatWaitOp::create(builder, loc, one, one);
  spatial::SpatYieldOp::create(builder, loc, ValueRange{concat.getOutput()});
  builder.setInsertionPointAfter(compute);
  auto reader = spatial::SpatScheduledCompute::create(builder, loc, TypeRange{}, ValueRange{}, ValueRange{});
  reader->setAttr("coreId", builder.getI32IntegerAttr(1));
  reader.getBody().push_back(new Block);
  builder.setInsertionPointToStart(&reader.getBody().front());
  zero = arith::ConstantIndexOp::create(builder, loc, 0);
  one = arith::ConstantIndexOp::create(builder, loc, 1);
  spatial::SpatWaitOp::create(builder, loc, zero, one);
  spatial::SpatSyncOp::create(builder, loc, zero, one);
  spatial::SpatYieldOp::create(builder, loc, ValueRange{});
  builder.setInsertionPointAfter(reader);
  func::ReturnOp::create(builder, loc, compute.getResults());
  TEST_CHECK(succeeded(verify(module)));
  auto report = analyzePimSpatialSchedule(module, target, spatialDetails());
  TEST_CHECK(report.getInteger("global_write_requests") == 5);
  TEST_CHECK(report.getInteger("global_write_bytes") == 80);
  auto *runs = (*report.getArray("request_order"))[0].getAsObject()->getArray("runs");
  TEST_CHECK(runs->size() == 2);
  TEST_CHECK((*(*runs)[0].getAsArray())[1].getAsInteger() == 60);
  TEST_CHECK((*(*runs)[1].getAsArray())[1].getAsInteger() == 5);
  TEST_CHECK((*(*runs)[1].getAsArray())[2].getAsInteger() == 4);

  // Ablation removes standalone and fused synchronization, not host accesses.
  builder.setInsertionPoint(compute.getBody().front().getTerminator());
  Value readerId = arith::ConstantIndexOp::create(builder, loc, 1);
  Value eventRegister = arith::ConstantIndexOp::create(builder, loc, 2);
  spatial::SpatSyncOp::create(builder, loc, readerId, eventRegister);
  builder.setInsertionPoint(reader.getBody().front().getTerminator());
  eventRegister = arith::ConstantIndexOp::create(builder, loc, 2);
  spatial::SpatHostWaitLoadOp::create(builder, loc, slice.getType(),
    zero, one, zero, eventRegister, one);
  TEST_CHECK(succeeded(verify(module)));
  PimExecutionProgram reference, omitted;
  auto options = spatialDetails();
  options.executionTrace = &reference;
  auto synchronized = evaluatePimSpatialSchedule(module, target, options);
  TEST_CHECK(synchronized.cost.available());
  options.omitSynchronization = true;
  options.executionTrace = &omitted;
  auto ablation = evaluatePimSpatialSchedule(module, target, options);
  TEST_CHECK(ablation.cost.available());
  TEST_CHECK(ablation.report.getObject("event_counts")->getInteger("sync") == 0);
  TEST_CHECK(ablation.report.getObject("event_counts")->getInteger("wait") == 0);
  for (const char *metric : {"global_read_requests", "global_read_bytes",
         "global_write_requests", "global_write_bytes", "global_memory_service_ns"})
    TEST_CHECK(ablation.report.getInteger(metric) == synchronized.report.getInteger(metric));
  for (auto &sequence : reference)
    llvm::erase_if(sequence, [](const E &event) { return event.kind == E::Signal || event.kind == E::Wait; });
  for (const auto &sequence : omitted)
    for (const E &event : sequence)
      TEST_CHECK(event.kind != E::Signal && event.kind != E::Wait);
  options.executionTrace = nullptr;
  auto expected = evaluatePimExecution(reference, target, options);
  TEST_CHECK(ablation.cost.primary == expected.cost.primary);
  TEST_CHECK(ablation.cost.alternative == expected.cost.alternative);
  module.erase();
}

static void testUnrepresentableSpatial(PimInstructionTiming target) {
  MLIRContext context;
  context.loadDialect<arith::ArithDialect, func::FuncDialect, tensor::TensorDialect,
    scf::SCFDialect, spatial::SpatialDialect>();
  for (bool loopOverflow : {false, true}) {
    auto module = ModuleOp::create(UnknownLoc::get(&context));
    OpBuilder builder(module.getBodyRegion());
    auto loc = module.getLoc();
    auto function = func::FuncOp::create(builder, loc, "test", builder.getFunctionType({}, {}));
    builder.setInsertionPointToStart(function.addEntryBlock());
    auto compute = spatial::SpatScheduledCompute::create(builder, loc, TypeRange{}, ValueRange{}, ValueRange{});
    compute->setAttr("coreId", builder.getI32IntegerAttr(0));
    compute.getBody().push_back(new Block);
    builder.setInsertionPointToStart(&compute.getBody().front());
    if (loopOverflow) {
      Value lower = arith::ConstantIndexOp::create(builder, loc, std::numeric_limits<int64_t>::min());
      Value upper = arith::ConstantIndexOp::create(builder, loc, std::numeric_limits<int64_t>::max());
      Value step = arith::ConstantIndexOp::create(builder, loc, 1);
      scf::ForOp::create(builder, loc, lower, upper, step);
    } else {
      auto type = RankedTensorType::get({std::numeric_limits<int64_t>::max(), 2}, builder.getF32Type());
      Value input = tensor::EmptyOp::create(builder, loc, type.getShape(), type.getElementType());
      spatial::SpatReluOp::create(builder, loc, type, input);
    }
    spatial::SpatYieldOp::create(builder, loc, ValueRange{});
    builder.setInsertionPointAfter(compute);
    func::ReturnOp::create(builder, loc);
    TEST_CHECK(succeeded(verify(module)));
    auto report = analyzePimSpatialSchedule(module, target, spatialDetails());
    TEST_CHECK(report.getString("status") == "unavailable");
    TEST_CHECK(!report.getBoolean("safe_for_selection").value());
    module.erase();
  }
}


static void testMatrixGeometry(PimInstructionTiming target) {
  MLIRContext context;
  context.loadDialect<arith::ArithDialect, func::FuncDialect, tensor::TensorDialect,
    spatial::SpatialDialect>();
  target.matrixCellBits = 2;
  for (bool encoded : {false, true})
  for (int64_t rows : {64, 128, 192})
  for (int64_t columns : {64, 128, 192, 256}) {
    auto module = ModuleOp::create(UnknownLoc::get(&context));
    OpBuilder builder(module.getBodyRegion());
    auto loc = module.getLoc();
    auto inputType = RankedTensorType::get({1, rows}, builder.getF32Type());
    auto weightType = RankedTensorType::get({rows, columns}, builder.getF32Type(),
      encoded ? builder.getStringAttr("unsupported-layout") : Attribute{});
    auto resultType = RankedTensorType::get({1, columns}, builder.getF32Type());
    auto function = func::FuncOp::create(builder, loc, "test", builder.getFunctionType({}, {}));
    builder.setInsertionPointToStart(function.addEntryBlock());
    Value weight = arith::ConstantOp::create(builder, loc,
      DenseElementsAttr::get(weightType, builder.getF32FloatAttr(1)));
    auto compute = spatial::SpatScheduledCompute::create(builder, loc, TypeRange{},
      ValueRange{weight}, ValueRange{});
    compute->setAttr("coreId", builder.getI32IntegerAttr(0));
    compute.getBody().push_back(new Block);
    Value captured = compute.getBody().front().addArgument(weightType, loc);
    builder.setInsertionPointToStart(&compute.getBody().front());
    Value input = tensor::EmptyOp::create(builder, loc, inputType.getShape(), inputType.getElementType());
    spatial::SpatVMMOp::create(builder, loc, resultType, captured, input);
    spatial::SpatYieldOp::create(builder, loc, ValueRange{});
    builder.setInsertionPointAfter(compute);
    func::ReturnOp::create(builder, loc);
    TEST_CHECK(succeeded(verify(module)));
    auto report = analyzePimSpatialSchedule(module, target);
    if (encoded || columns % target.matrixColumns || static_cast<uint64_t>(rows) > target.matrixRows)
      TEST_CHECK(report.getString("reason") == "unavailable physical matrix weight geometry");
    else TEST_CHECK(report.getString("status") != "unavailable");
    module.erase();
  }
}

static void testOutputMultiplicity(PimInstructionTiming target) {
  MLIRContext context;
  context.loadDialect<arith::ArithDialect, func::FuncDialect, tensor::TensorDialect,
    spatial::SpatialDialect>();
  for (bool repeatedResult : {false, true}) {
    auto module = ModuleOp::create(UnknownLoc::get(&context));
    OpBuilder builder(module.getBodyRegion());
    auto loc = module.getLoc();
    auto type = RankedTensorType::get({1, 64}, builder.getF32Type());
    auto function = func::FuncOp::create(builder, loc, "test", builder.getFunctionType({}, {type, type}));
    builder.setInsertionPointToStart(function.addEntryBlock());
    SmallVector<Type> results(repeatedResult ? 1 : 2, type);
    auto compute = spatial::SpatScheduledCompute::create(builder, loc, results, ValueRange{}, ValueRange{});
    compute->setAttr("coreId", builder.getI32IntegerAttr(0));
    compute.getBody().push_back(new Block);
    builder.setInsertionPointToStart(&compute.getBody().front());
    Value input = tensor::EmptyOp::create(builder, loc, type.getShape(), type.getElementType());
    Value output = spatial::SpatReluOp::create(builder, loc, type, input);
    SmallVector<Value> yielded(results.size(), output);
    spatial::SpatYieldOp::create(builder, loc, yielded);
    builder.setInsertionPointAfter(compute);
    func::ReturnOp::create(builder, loc, ValueRange{compute.getResult(0), compute.getResult(repeatedResult ? 0 : 1)});
    TEST_CHECK(succeeded(verify(module)));
    auto report = analyzePimSpatialSchedule(module, target, spatialDetails());
    if (repeatedResult)
      TEST_CHECK(report.getString("reason") == "unavailable indirect compute output collection");
    else {
      TEST_CHECK(report.getInteger("global_write_requests") == 2);
      TEST_CHECK(report.getInteger("global_write_bytes") == 128);
      auto *runs = (*report.getArray("request_order"))[0].getAsObject()->getArray("runs");
      TEST_CHECK((*(*runs)[0].getAsArray())[2].getAsInteger() == 2);
    }
    module.erase();
  }
}

static void testPeriodicStateProof() {
  auto target = hardware(1);
  target.processorCount = 3;
  target.sharedMemoryReadReturnsPayload = false;
  target.localMemoryPeriodNs = 1;
  target.localMemoryReadLatencyCycles = 20;
  target.localMemoryWriteLatencyCycles = 30;
  target.transferWidthBytes = target.sharedMemoryWidthBytes = 32;
  target.sharedMemoryReadLatencyNs = 40;
  target.sharedMemoryWriteLatencyNs = 50;
  target.sharedMemoryRequestLatencyNs = target.sharedMemoryResponseLatencyNs = {6, 6, 6};
  target.interProcessorLatencyNs = {0, 3, 4, 3, 0, 3, 4, 3, 0};
  PimExecutionProgram program{
    {{E::Work, 377}, {E::Store, 875}, {E::Work, 545}, {E::Load, 367},
      {E::Work, 333}, {E::Store, 80}, {E::Signal, 0, 1}, {E::Wait, 0, 0, 0, 1}},
    {{E::Work, 924}, {E::Load, 312}, {E::Work, 70}, {E::Store, 434},
      {E::Work, 337}, {E::Load, 275}, {E::Signal, 0, 2}, {E::Wait, 0, 0, 0, 1}},
    {{E::Work, 158}, {E::Store, 565}, {E::Work, 126}, {E::Load, 245},
      {E::Work, 1124}, {E::Store, 367}, {E::Signal, 0, 0}, {E::Wait, 0, 0, 0, 1}}};
  auto baseline = analyzePimExecution(program, target);
  PimPeriodicOptions options;
  options.verifyPeriodicState = true;
  options.workBreakdown = true;
  auto verified = analyzePimExecution(program, target, options);
  auto *queue = verified.getObject("memory_controller");
  auto *legacy = queue->getObject("legacy_interval_window");
  TEST_CHECK(legacy);
  for (const char *field : {"interval_ns", "alternative_tie_interval_ns"})
    TEST_CHECK(legacy->getNumber(field) == baseline.getObject("memory_controller")->getNumber(field));
  auto *state = queue->getObject("periodic_state");
  TEST_CHECK(state && state->getString("status") == "proved");
  TEST_CHECK(state->getInteger("period_generations") == 3);
  TEST_CHECK(state->getNumber("interval_ns") == 17485.0 / 3);
  TEST_CHECK(queue->getNumber("interval_ns") == state->getNumber("interval_ns"));
  TEST_CHECK(legacy->getNumber("interval_ns") != state->getNumber("interval_ns"));
  // A fast writer builds signal tokens. Repeated local intervals are insufficient.
  program = {{{E::Signal, 0, 1}, {E::Work, 10}},
    {{E::Wait, 0, 0, 0, 1}, {E::Work, 100}}};
  verified = analyzePimExecution(program, target, options);
  state = verified.getObject("memory_controller")->getObject("periodic_state");
  TEST_CHECK(state && state->getString("status") == "unproved");
}

static void testCompleteSchedulerState() {
  auto target = hardware(1);
  PimPeriodicOptions options;
  options.verifyPeriodicState = true;
  auto verify = [&](const PimExecutionProgram &program, bool expected) {
    auto baseline = analyzePimExecution(program, target);
    auto typed = evaluatePimExecution(program, target, options);
    TEST_CHECK(typed.cost.convergenceProved() == expected);
    auto detailed = std::move(typed.report);
    auto *queue = detailed.getObject("memory_controller");
    TEST_CHECK(queue);
    TEST_CHECK(detailed.getString("status") == (typed.cost.available() ? "estimated" : "unavailable"));
    if (!typed.cost.available()) TEST_CHECK(detailed.getString("reason") == typed.cost.reason);
    if (expected) {
      TEST_CHECK(queue->getNumber("interval_ns") == typed.cost.primary);
      TEST_CHECK(queue->getNumber("alternative_tie_interval_ns") == typed.cost.alternative);
    } else {
      for (const char *field : {"interval_ns", "alternative_tie_interval_ns"})
        TEST_CHECK(queue->getNumber(field) == baseline.getObject("memory_controller")->getNumber(field));
    }
    size_t policy = 0;
    for (const char *field : {"periodic_state", "alternative_periodic_state"}) {
      auto *state = queue->getObject(field);
      TEST_CHECK(state && (state->getString("status") == "proved") == expected);
      const auto &evidence = typed.cost.recurrence[policy++];
      TEST_CHECK(evidence.proved() == expected);
      TEST_CHECK(evidence.intervalNs == state->getNumber("interval_ns"));
      if (expected) TEST_CHECK(state->getBoolean("complete_scheduler_state") == true);
    }
    return detailed;
  };
  verify({{{E::Work, 0}}}, false);
  auto steady = verify({{{E::Work, 100}}}, true);
  auto *state = steady.getObject("memory_controller")->getObject("periodic_state");
  TEST_CHECK(state->getInteger("period_generations") == 1);
  TEST_CHECK(state->getNumber("interval_ns") == 100);
  verify({{{E::Work, uint64_t{1} << 60}}}, true);
  // Simultaneous signals, drained token queues, and controller arrivals must
  // recur under each deterministic arbitration policy.
  verify({{{E::Work, 10}, {E::Signal, 0, 1}, {E::Wait, 0, 0, 0, 1}},
    {{E::Work, 10}, {E::Signal, 0, 0}, {E::Wait, 0, 0, 0, 1}}}, true);
  verify({{{E::Load, 65}, {E::Signal, 0, 1}, {E::Wait, 0, 0, 0, 1}},
    {{E::Load, 65}, {E::Signal, 0, 0}, {E::Wait, 0, 0, 0, 1}}}, true);
  // Repeated output intervals do not prove uniform progress on every core.
  auto unequal = verify({{{E::Work, 7}}, {{E::Work, 11}}}, false);
  TEST_CHECK(unequal.getObject("memory_controller")->getString("status") == "estimated");
  // The complete state retains growing dependency-token queues.
  auto backlog = verify({{{E::Signal, 0, 1}, {E::Work, 10}},
    {{E::Wait, 0, 0, 0, 1}, {E::Work, 100}}}, false);
  TEST_CHECK(backlog.getObject("memory_controller")->getString("status") == "estimated");
  // Both endpoints exist, but neither WAIT can receive its signal.
  auto deadlock = verify({{{E::Wait, 0, 0, 0, 1}, {E::Signal, 0, 1}},
    {{E::Wait, 0, 0, 0, 1}, {E::Signal, 0, 0}}}, false);
  TEST_CHECK(deadlock.getObject("memory_controller")->getString("status") == "unavailable");
  verify({{{E::Send, 65, 1}, {E::Receive, 65, 1}},
    {{E::Send, 65, 0}, {E::Receive, 65, 0}}}, false);
  // An independent active core must not hide an inactive dependency cycle.
  verify({{{E::Work, 10}}, {{E::Wait, 0, 0, 0, 1}, {E::Signal, 0, 1}}}, false);
  auto overflow = analyzePimExecution({{{E::Work, uint64_t{1} << 62}}}, target, options);
  TEST_CHECK(overflow.getString("status") == "unavailable");
}

int main() {
  testDependencyTokens();
  testCoreAssignments();
  testPeriodicStateProof();
  testCompleteSchedulerState();
  for (int64_t period : {1, 5}) {
    testLateRecurrenceAuthority(hardware(period));
    testTypedResult(hardware(period));
    testSearchGuidance(hardware(period));
    testSharedEngine(hardware(period));
    testAccountingAndOverflow(hardware(period));
    testMatrixGeometry(hardware(period));
    testOutputMultiplicity(hardware(period));
    testHostSlice(hardware(period));
    testCapturedVectorWeights(hardware(period));
    testPublicationAndOutput(hardware(period));
    testUnrepresentableSpatial(hardware(period));
    testReadOnlySpatial(hardware(period), false);
    testReadOnlySpatial(hardware(period), true);
  }
}