// Copyright (c) 2026 Advanced Micro Devices, Inc.
// SPDX-License-Identifier: MIT
//
// FlatBuffers schema for simulation checkpoint/restore.
// Captures the full state of a running AMDGPU simulation.

include "simulation_config.fbs";

namespace rocjitsu.fb;

/// Serialized state of a single wavefront.
table WavefrontState {
  wf_id: uint32;
  wg_id: uint32;
  pc: uint64;
  exec: uint64;
  vcc: uint64;
  m0: uint32;
  halted: bool;
  status: uint32;
  sgprs: [uint32]; // Scalar register file (NUM_SGPRS entries).
  vgprs: [ubyte];  // Vector register file (raw bytes).
  mode: uint32;     // MODE register.
  wave_sched_mode: uint32; // WAVE_SCHED_MODE register.
}

/// Serialized state of a compute unit.
table ComputeUnitState {
  name: string;
  wavefronts: [WavefrontState];
  next_wf: uint64; // Round-robin scheduling index.
}

/// Serialized state of a command processor.
table CommandProcessorState {
  name: string;
  dispatched: uint64; // Number of packets already processed.
  next_cu: uint64;    // Round-robin CU index.
}

/// A single 4KB memory page.
table MemoryPage {
  address: uint64; // Page-aligned base address.
  data: [ubyte];   // 4096 bytes of page data.
}

/// Serialized state of GPU memory.
table GpuMemoryState {
  pages: [MemoryPage];
}

/// Complete simulation checkpoint.
table SimulationCheckpoint {
  tick: uint64;
  config: SimulationConfig;
  compute_units: [ComputeUnitState];
  command_processor: CommandProcessorState;
  memory: GpuMemoryState;
}

root_type SimulationCheckpoint;