diff --git a/justfile b/justfile index d2da8d51..5fd7b398 100644 --- a/justfile +++ b/justfile @@ -132,6 +132,9 @@ aten-emulate nickname *args: test-sliced-layer-builder: python3 transactional_emulator/testbench/test_sliced_layer_builder.py +# Rust emulator timing smoke: unit tests + deterministic simulation-cycle goldens. +timing-smoke: + bash transactional_emulator/testbench/timing_goldens/run_timing_smoke.sh # Unit tests for LUI+ADDI large immediate fix in ASM templates test-large-immediate: @@ -171,4 +174,3 @@ multilayer-decoder-profile model="smolvlm2": # ATen-backed sliced emulator check: PlenaCompiler + ops.* -> emulator -> numerical check test-sliced-aten-emulator model="AICrossSim/clm-60m" seq_len="64" num_layers="1": cd PLENA_Compiler && PYTHONPATH=".:../PLENA_Tools:../transactional_emulator/testbench:..:" python3 -m compiler.aten.sliced_emulator_runner {{model}} --seq-len {{seq_len}} --num-layers {{num_layers}} - diff --git a/transactional_emulator/src/accelerator/dispatch.rs b/transactional_emulator/src/accelerator/dispatch.rs index 12f3b0d3..06abb68c 100644 --- a/transactional_emulator/src/accelerator/dispatch.rs +++ b/transactional_emulator/src/accelerator/dispatch.rs @@ -7,11 +7,12 @@ use half::bf16; use quantize::MxDataType; use crate::runtime_config::{ - HLEN, MATRIX_KV_TYPE, MATRIX_WEIGHT_TYPE, MLEN, PREFETCH_M_AMOUNT, PREFETCH_V_AMOUNT, + BLEN, HLEN, MATRIX_KV_TYPE, MATRIX_WEIGHT_TYPE, MLEN, PREFETCH_M_AMOUNT, PREFETCH_V_AMOUNT, SCALAR_FP_BASIC_CYCLES, SCALAR_FP_EXP_CYCLES, SCALAR_FP_RECI_CYCLES, SCALAR_FP_SQRT_CYCLES, SCALAR_INT_BASIC_CYCLES, STORE_V_AMOUNT, VECTOR_ACTIVATION_TYPE, VECTOR_KV_TYPE, VLEN, }; use crate::stage_profile::{ResourceKind, StageProfiler}; +use crate::timing_overlay::{SramRange, SramSpace, TimingAccess, TimingOverlay}; use crate::{cycle, dma, op}; use runtime::Executor; @@ -54,13 +55,20 @@ impl Accelerator { &mut self, ops: &[op::Opcode], mut stage_profiler: Option<&mut StageProfiler>, + mut timing_overlay: Option<&mut TimingOverlay>, ) { let mut pc: usize = 0; // Program counter while pc < ops.len() { let executed_pc = pc; let op = &ops[pc]; - let profile_start_instant = if stage_profiler.is_some() { + let capture_elapsed = stage_profiler.is_some() || timing_overlay.is_some(); + let timing_access = if timing_overlay.is_some() { + Some(self.timing_access_for_opcode(op)) + } else { + None + }; + let profile_start_instant = if capture_elapsed { Some(Executor::current().now()) } else { None @@ -625,35 +633,118 @@ impl Accelerator { pc += 1; } - if let (Some(start_instant), Some(profiler)) = - (profile_start_instant, stage_profiler.as_deref_mut()) - { + if let Some(start_instant) = profile_start_instant { let elapsed_duration = Executor::current().now() - start_instant; let elapsed_secs = elapsed_duration.as_picos() as f64 / 1_000_000_000_000.0; let elapsed_cycles = StageProfiler::duration_to_cycles(elapsed_duration); - let (hbm_bytes_read, hbm_bytes_written) = if let (Some(before), Some(after)) = - (profile_start_hbm, self.hbm.statistics()) + if let (Some(access), Some(overlay)) = + (timing_access, timing_overlay.as_deref_mut()) { - ( - after - .total_bytes_read - .saturating_sub(before.total_bytes_read), - after - .total_bytes_written - .saturating_sub(before.total_bytes_written), - ) - } else { - (0, 0) - }; - profiler.record( - executed_pc, - elapsed_secs, - elapsed_cycles, - resource_kind_for_opcode(op), - hbm_bytes_read, - hbm_bytes_written, - ); + overlay.record(access, elapsed_cycles); + } + if let Some(profiler) = stage_profiler.as_deref_mut() { + let (hbm_bytes_read, hbm_bytes_written) = if let (Some(before), Some(after)) = + (profile_start_hbm, self.hbm.statistics()) + { + ( + after + .total_bytes_read + .saturating_sub(before.total_bytes_read), + after + .total_bytes_written + .saturating_sub(before.total_bytes_written), + ) + } else { + (0, 0) + }; + profiler.record( + executed_pc, + elapsed_secs, + elapsed_cycles, + resource_kind_for_opcode(op), + hbm_bytes_read, + hbm_bytes_written, + ); + } + } + } + } + + fn timing_access_for_opcode(&self, op: &op::Opcode) -> TimingAccess { + let matrix_tile = *MLEN * *MLEN; + let vector_tile = *VLEN; + let matrix_prefetch = *MLEN * *PREFETCH_M_AMOUNT; + let vector_prefetch = *VLEN * *PREFETCH_V_AMOUNT; + let matrix_vector_batch = *MLEN * *BLEN; + + let matrix = |start: u32, len: u32| SramRange::new(SramSpace::Matrix, start, len); + let vector = |start: u32, len: u32| SramRange::new(SramSpace::Vector, start, len); + let gp = |reg: u8| self.reg_file.read_gp(reg); + + match *op { + op::Opcode::H_PREFETCH_M { rd, .. } => { + TimingAccess::prefetch(vec![matrix(gp(rd), matrix_prefetch)]) + } + op::Opcode::H_PREFETCH_V { rd, .. } => { + TimingAccess::prefetch(vec![vector(gp(rd), vector_prefetch)]) + } + op::Opcode::H_STORE_V { .. } => TimingAccess::Barrier, + + op::Opcode::M_MM { rs1, rs2 } | op::Opcode::M_TMM { rs1, rs2 } => { + TimingAccess::compute(vec![ + matrix(gp(rs1), matrix_tile), + vector(gp(rs2), matrix_vector_batch), + ]) + } + op::Opcode::M_BMM { rs1, rs2 } | op::Opcode::M_BTMM { rs1, rs2 } => { + TimingAccess::compute(vec![ + matrix(gp(rs1), matrix_tile), + vector(gp(rs2), matrix_tile), + ]) } + op::Opcode::M_MV { rs1, rs2 } | op::Opcode::M_TMV { rs1, rs2 } => { + TimingAccess::compute(vec![ + matrix(gp(rs1), matrix_tile), + vector(gp(rs2), vector_tile), + ]) + } + op::Opcode::M_BMV { rs1, rs2, rd } | op::Opcode::M_BTMV { rs1, rs2, rd } => { + TimingAccess::compute(vec![ + matrix(gp(rs1).wrapping_add(gp(rd)), matrix_tile), + vector(gp(rs2), vector_tile), + ]) + } + op::Opcode::M_MM_WO { rd, imm, .. } + | op::Opcode::M_BMM_WO { rd, imm } + | op::Opcode::M_MV_WO { rd, imm } + | op::Opcode::M_BMV_WO { rd, imm } => { + TimingAccess::compute(vec![vector(gp(rd).wrapping_add(imm), vector_tile)]) + } + + op::Opcode::V_ADD_VV { rs1, rs2, .. } + | op::Opcode::V_SUB_VV { rs1, rs2, .. } + | op::Opcode::V_MUL_VV { rs1, rs2, .. } => TimingAccess::compute(vec![ + vector(gp(rs1), vector_tile), + vector(gp(rs2), vector_tile), + ]), + op::Opcode::V_ADD_VF { rs1, .. } + | op::Opcode::V_SUB_VF { rs1, .. } + | op::Opcode::V_MUL_VF { rs1, .. } + | op::Opcode::V_MAX_VF { rs1, .. } + | op::Opcode::V_MIN_VF { rs1, .. } + | op::Opcode::V_EXP_V { rs1, .. } + | op::Opcode::V_RECI_V { rs1, .. } + | op::Opcode::V_RED_SUM { rs1, .. } + | op::Opcode::V_RED_MAX { rs1, .. } + | op::Opcode::V_TOPK { rs1, .. } => { + TimingAccess::compute(vec![vector(gp(rs1), vector_tile)]) + } + op::Opcode::V_SHFT_V { rs1, .. } => { + TimingAccess::compute(vec![vector(gp(rs1), vector_tile)]) + } + + op::Opcode::C_BREAK => TimingAccess::Barrier, + _ => TimingAccess::Other, } } } diff --git a/transactional_emulator/src/cli.rs b/transactional_emulator/src/cli.rs index 6c2eabe7..44b244f6 100644 --- a/transactional_emulator/src/cli.rs +++ b/transactional_emulator/src/cli.rs @@ -173,4 +173,10 @@ pub(crate) struct Opts { #[arg(long)] /// Optional JSON output path for runtime stage profile results. pub(crate) stage_profile_out: Option, + + #[arg(long, help_heading = "Experimental Timing")] + /// Off-by-default exploratory timing overlay that hides independent HBM + /// prefetch cycles behind later matrix/vector compute. This changes only + /// reported simulation cycles, not functional execution or HBM traffic. + pub(crate) experimental_overlap_prefetch_compute: bool, } diff --git a/transactional_emulator/src/main.rs b/transactional_emulator/src/main.rs index 5b74b768..05a5c8f2 100644 --- a/transactional_emulator/src/main.rs +++ b/transactional_emulator/src/main.rs @@ -8,9 +8,10 @@ mod op; mod runner; mod runtime_config; mod stage_profile; +mod timing_overlay; mod vector_machine; -use runtime::{Executor, Instant}; +use runtime::{Duration, Executor, Instant}; #[macro_export] macro_rules! cycle { @@ -30,6 +31,21 @@ async fn main() { let cycles = latency .as_picos() .div_ceil(runtime_config::PERIOD.as_picos().max(1)); + if let Some(adjusted_cycles) = timing_overlay::experimental_report_cycles() { + let adjusted_latency = + Duration::from_picos(runtime_config::PERIOD.as_picos() * adjusted_cycles); + tracing::info!( + "Experimental overlap serial latency {:?} cycles {}", + executor.now(), + cycles + ); + tracing::info!( + "Simulation completed. Latency {:?} cycles {}", + adjusted_latency, + adjusted_cycles + ); + return; + } tracing::info!( "Simulation completed. Latency {:?} cycles {}", executor.now(), diff --git a/transactional_emulator/src/runner.rs b/transactional_emulator/src/runner.rs index 4f2a2f0c..3df4da0b 100644 --- a/transactional_emulator/src/runner.rs +++ b/transactional_emulator/src/runner.rs @@ -16,6 +16,7 @@ use crate::runtime_config::{ VECTOR_SRAM_SIZE, VECTOR_SRAM_TYPE, VLEN, }; use crate::stage_profile::StageProfiler; +use crate::timing_overlay::{self as timing_overlay_state, TimingOverlay}; use crate::vector_machine::VectorMachine; use crate::{cli, op}; @@ -33,6 +34,7 @@ fn dump_to_file(path: &str, bytes: &[u8]) { pub(crate) async fn run_from_cli() { let opts = Opts::parse(); + timing_overlay_state::clear_experimental_report_cycles(); // If --settings is given, set PLENA_SETTINGS_TOML env var BEFORE any // LazyLock access (which triggers load_config()). This ensures the @@ -215,12 +217,36 @@ pub(crate) async fn run_from_cli() { panic!("failed to build stage profile from ASM {:?}: {err}", path) }) }); + let mut timing_overlay = opts + .experimental_overlap_prefetch_compute + .then(TimingOverlay::default); accelerator - .do_ops(&decoded_ops, stage_profiler.as_mut()) + .do_ops( + &decoded_ops, + stage_profiler.as_mut(), + timing_overlay.as_mut(), + ) .await; + let serial_duration = Executor::current().now() - Instant::INIT; + let serial_cycles = StageProfiler::duration_to_cycles(serial_duration); + if let Some(overlay) = timing_overlay.as_ref() { + let summary = overlay.summary(serial_cycles); + timing_overlay_state::set_experimental_report_cycles(summary.adjusted_cycles); + tracing::info!( + serial_cycles = summary.serial_cycles, + adjusted_cycles = summary.adjusted_cycles, + hidden_prefetch_cycles = summary.hidden_prefetch_cycles, + pending_prefetch_cycles = summary.pending_prefetch_cycles, + prefetch_ops = summary.prefetch_ops, + compute_ops = summary.compute_ops, + dependent_prefetch_stalls = summary.dependent_prefetch_stalls, + "Experimental overlap prefetch/compute timing overlay" + ); + } + if let Some(profile) = stage_profiler.as_mut() { - profile.set_total_simulation_duration(Executor::current().now() - Instant::INIT); + profile.set_total_simulation_duration(serial_duration); } if let Some(profile) = stage_profiler.as_ref() { diff --git a/transactional_emulator/src/timing_overlay.rs b/transactional_emulator/src/timing_overlay.rs new file mode 100644 index 00000000..4c4b7205 --- /dev/null +++ b/transactional_emulator/src/timing_overlay.rs @@ -0,0 +1,211 @@ +use std::sync::atomic::{AtomicU64, Ordering}; + +static EXPERIMENTAL_REPORT_CYCLES: AtomicU64 = AtomicU64::new(0); + +pub(crate) fn clear_experimental_report_cycles() { + EXPERIMENTAL_REPORT_CYCLES.store(0, Ordering::Relaxed); +} + +pub(crate) fn set_experimental_report_cycles(cycles: u64) { + EXPERIMENTAL_REPORT_CYCLES.store(cycles.max(1), Ordering::Relaxed); +} + +pub(crate) fn experimental_report_cycles() -> Option { + match EXPERIMENTAL_REPORT_CYCLES.load(Ordering::Relaxed) { + 0 => None, + cycles => Some(cycles), + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum SramSpace { + Matrix, + Vector, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct SramRange { + pub(crate) space: SramSpace, + pub(crate) start: u32, + pub(crate) len: u32, +} + +impl SramRange { + pub(crate) fn new(space: SramSpace, start: u32, len: u32) -> Self { + Self { space, start, len } + } + + fn overlaps(self, other: Self) -> bool { + if self.space != other.space || self.len == 0 || other.len == 0 { + return false; + } + let self_end = self.start.saturating_add(self.len); + let other_end = other.start.saturating_add(other.len); + self.start < other_end && other.start < self_end + } +} + +#[derive(Clone, Debug)] +pub(crate) enum TimingAccess { + Prefetch { write_ranges: Vec }, + Compute { read_ranges: Vec }, + Barrier, + Other, +} + +impl TimingAccess { + pub(crate) fn prefetch(write_ranges: Vec) -> Self { + Self::Prefetch { write_ranges } + } + + pub(crate) fn compute(read_ranges: Vec) -> Self { + Self::Compute { read_ranges } + } +} + +#[derive(Clone, Debug)] +struct PendingPrefetch { + sequence_id: u64, + remaining_cycles: u64, + write_ranges: Vec, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct TimingOverlaySummary { + pub(crate) serial_cycles: u64, + pub(crate) adjusted_cycles: u64, + pub(crate) hidden_prefetch_cycles: u64, + pub(crate) pending_prefetch_cycles: u64, + pub(crate) prefetch_ops: u64, + pub(crate) compute_ops: u64, + pub(crate) dependent_prefetch_stalls: u64, +} + +#[derive(Debug, Default)] +pub(crate) struct TimingOverlay { + next_sequence_id: u64, + pending_prefetches: Vec, + hidden_prefetch_cycles: u64, + prefetch_ops: u64, + compute_ops: u64, + dependent_prefetch_stalls: u64, +} + +impl TimingOverlay { + pub(crate) fn record(&mut self, access: TimingAccess, elapsed_cycles: u64) { + match access { + TimingAccess::Prefetch { write_ranges } => { + if elapsed_cycles == 0 { + return; + } + let sequence_id = self.next_sequence_id; + self.next_sequence_id = self.next_sequence_id.wrapping_add(1); + self.prefetch_ops += 1; + self.pending_prefetches.push(PendingPrefetch { + sequence_id, + remaining_cycles: elapsed_cycles, + write_ranges, + }); + self.pending_prefetches + .sort_by_key(|prefetch| prefetch.sequence_id); + } + TimingAccess::Compute { read_ranges } => { + self.compute_ops += 1; + self.retire_dependent_prefetches(&read_ranges); + self.hide_independent_prefetch_cycles(elapsed_cycles); + } + TimingAccess::Barrier => { + self.pending_prefetches.clear(); + } + TimingAccess::Other => {} + } + } + + pub(crate) fn summary(&self, serial_cycles: u64) -> TimingOverlaySummary { + let hidden = self.hidden_prefetch_cycles.min(serial_cycles); + TimingOverlaySummary { + serial_cycles, + adjusted_cycles: serial_cycles.saturating_sub(hidden), + hidden_prefetch_cycles: hidden, + pending_prefetch_cycles: self.pending_prefetch_cycles(), + prefetch_ops: self.prefetch_ops, + compute_ops: self.compute_ops, + dependent_prefetch_stalls: self.dependent_prefetch_stalls, + } + } + + fn retire_dependent_prefetches(&mut self, read_ranges: &[SramRange]) { + let before = self.pending_prefetches.len(); + self.pending_prefetches.retain(|prefetch| { + !prefetch + .write_ranges + .iter() + .any(|write| read_ranges.iter().any(|read| write.overlaps(*read))) + }); + self.dependent_prefetch_stalls += (before - self.pending_prefetches.len()) as u64; + } + + fn hide_independent_prefetch_cycles(&mut self, compute_cycles: u64) { + let mut remaining_compute_cycles = compute_cycles; + let mut index = 0; + while remaining_compute_cycles > 0 && index < self.pending_prefetches.len() { + let hidden = self.pending_prefetches[index] + .remaining_cycles + .min(remaining_compute_cycles); + self.pending_prefetches[index].remaining_cycles -= hidden; + self.hidden_prefetch_cycles += hidden; + remaining_compute_cycles -= hidden; + if self.pending_prefetches[index].remaining_cycles == 0 { + self.pending_prefetches.remove(index); + } else { + index += 1; + } + } + } + + fn pending_prefetch_cycles(&self) -> u64 { + self.pending_prefetches + .iter() + .map(|prefetch| prefetch.remaining_cycles) + .sum() + } +} + +#[cfg(test)] +mod tests { + use super::{SramRange, SramSpace, TimingAccess, TimingOverlay}; + + #[test] + fn independent_compute_hides_pending_prefetch_cycles() { + let mut overlay = TimingOverlay::default(); + overlay.record( + TimingAccess::prefetch(vec![SramRange::new(SramSpace::Vector, 4096, 256)]), + 40, + ); + overlay.record( + TimingAccess::compute(vec![SramRange::new(SramSpace::Vector, 0, 64)]), + 6400, + ); + let summary = overlay.summary(6440); + assert_eq!(summary.hidden_prefetch_cycles, 40); + assert_eq!(summary.adjusted_cycles, 6400); + assert_eq!(summary.dependent_prefetch_stalls, 0); + } + + #[test] + fn dependent_compute_keeps_prefetch_cycles_serial() { + let mut overlay = TimingOverlay::default(); + overlay.record( + TimingAccess::prefetch(vec![SramRange::new(SramSpace::Vector, 4096, 256)]), + 40, + ); + overlay.record( + TimingAccess::compute(vec![SramRange::new(SramSpace::Vector, 4096, 64)]), + 6400, + ); + let summary = overlay.summary(6440); + assert_eq!(summary.hidden_prefetch_cycles, 0); + assert_eq!(summary.adjusted_cycles, 6440); + assert_eq!(summary.dependent_prefetch_stalls, 1); + } +} diff --git a/transactional_emulator/testbench/emulator_runner.py b/transactional_emulator/testbench/emulator_runner.py index d4759f24..6df8c88a 100644 --- a/transactional_emulator/testbench/emulator_runner.py +++ b/transactional_emulator/testbench/emulator_runner.py @@ -88,6 +88,7 @@ def run_emulator( stage_profile: bool | None = None, stage_profile_out: Path | None = None, run_label: str | None = None, + overlap_prefetch_compute: bool | None = None, ) -> dict: """Run the Rust transactional emulator with build artifacts from build_dir. @@ -110,12 +111,18 @@ def run_emulator( When None, PLENA_EMULATOR_STAGE_PROFILE controls it. run_label: optional filename suffix for repeat/determinism runs. When omitted, output filenames keep their historical names. + overlap_prefetch_compute: when true, pass the off-by-default + --experimental-overlap-prefetch-compute flag. + When None, PLENA_EMULATOR_OVERLAP_PREFETCH_COMPUTE + controls it. """ emulator_dir = Path(__file__).parent.parent # transactional_emulator/ binary = emulator_dir / "target" / "release" / "transactional_emulator" if stage_profile is None: stage_profile = _env_flag("PLENA_EMULATOR_STAGE_PROFILE") + if overlap_prefetch_compute is None: + overlap_prefetch_compute = _env_flag("PLENA_EMULATOR_OVERLAP_PREFETCH_COMPUTE") # Always rebuild before running. `cargo build --release` is a fast no-op when # current, and this prevents false failures from new ASM hitting a stale @@ -191,6 +198,8 @@ def run_emulator( "--stage-profile-out", str(profile_out_path), ] + if overlap_prefetch_compute: + cmd.append("--experimental-overlap-prefetch-compute") # tch's download-libtorch stores libtorch in the Cargo build cache. # The binary needs LD_LIBRARY_PATH to find it at runtime. @@ -233,6 +242,7 @@ def run_emulator( "artifacts": _artifact_summary(build_dir, asm_path, hbm_path), "log_path": str(log_path), "stage_profile_requested": bool(stage_profile), + "experimental_overlap_prefetch_compute": bool(overlap_prefetch_compute), } if run_label: metrics["run_label"] = run_label @@ -344,6 +354,7 @@ def run_emulator_repeat_gate( hbm_size: int | None = None, threads: int | None = None, stage_profile: bool | None = None, + overlap_prefetch_compute: bool | None = None, ) -> dict: """Run the same emulator artifact repeatedly and require identical cycles. @@ -366,6 +377,7 @@ def run_emulator_repeat_gate( threads=threads, stage_profile=stage_profile, run_label=label, + overlap_prefetch_compute=overlap_prefetch_compute, ) ) diff --git a/transactional_emulator/testbench/timing_goldens/check_timing_goldens.py b/transactional_emulator/testbench/timing_goldens/check_timing_goldens.py index 2cb74a59..61387f2b 100755 --- a/transactional_emulator/testbench/timing_goldens/check_timing_goldens.py +++ b/transactional_emulator/testbench/timing_goldens/check_timing_goldens.py @@ -10,14 +10,23 @@ from pathlib import Path from typing import Any -from transactional_emulator.testbench.emulator_runner import run_emulator - - REPO_ROOT = Path(__file__).resolve().parents[3] DEFAULT_PLENA_ROOT = REPO_ROOT.parents[1] DEFAULT_FIXTURE = Path(__file__).with_name("golden_workloads.json") CHECK_KEYS = ("sim_latency_cycles", "hbm_bytes_read", "hbm_bytes_written") +for _path in ( + REPO_ROOT, + REPO_ROOT / "PLENA_Compiler", + REPO_ROOT / "PLENA_Tools", + REPO_ROOT / "transactional_emulator" / "testbench", +): + _text = str(_path) + if _text not in sys.path: + sys.path.insert(0, _text) + +from transactional_emulator.testbench.emulator_runner import run_emulator + def _load_fixture(path: Path) -> dict[str, Any]: with path.open(encoding="utf-8") as f: @@ -83,6 +92,12 @@ def main() -> int: if threads is None: threads = int(workload.get("threads", defaults.get("threads", 1))) stage_profile = bool(workload.get("stage_profile", defaults.get("stage_profile", False))) + overlap_prefetch_compute = bool( + workload.get( + "overlap_prefetch_compute", + defaults.get("overlap_prefetch_compute", False), + ) + ) print(f"\n=== timing golden: {workload_id} ===") metrics = run_emulator( @@ -90,6 +105,7 @@ def main() -> int: threads=threads, stage_profile=stage_profile, run_label=f"golden.{workload_id}", + overlap_prefetch_compute=overlap_prefetch_compute, ) failures.extend(_compare_metrics(workload_id, workload["expected"], metrics)) diff --git a/transactional_emulator/testbench/timing_goldens/overlap_golden_workloads.json b/transactional_emulator/testbench/timing_goldens/overlap_golden_workloads.json new file mode 100644 index 00000000..22b7575c --- /dev/null +++ b/transactional_emulator/testbench/timing_goldens/overlap_golden_workloads.json @@ -0,0 +1,20 @@ +{ + "schema_version": 1, + "description": "Experimental off-by-default prefetch/compute overlap golden. This fixture intentionally requires --experimental-overlap-prefetch-compute and must not be mixed with default timing goldens.", + "defaults": { + "threads": 1, + "stage_profile": false, + "overlap_prefetch_compute": true + }, + "workloads": [ + { + "id": "g2_prefetch_then_compute_overlap_on", + "build_dir": "outputs/window1_p1/timing_gate_runs/g2_prefetch_then_compute_overlap_on", + "expected": { + "sim_latency_cycles": 6408, + "hbm_bytes_read": 512, + "hbm_bytes_written": 0 + } + } + ] +} diff --git a/transactional_emulator/testbench/timing_goldens/run_timing_smoke.sh b/transactional_emulator/testbench/timing_goldens/run_timing_smoke.sh new file mode 100644 index 00000000..a6c24816 --- /dev/null +++ b/transactional_emulator/testbench/timing_goldens/run_timing_smoke.sh @@ -0,0 +1,21 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)" + +nix develop "$repo_root" -c bash -lc "cd '$repo_root/transactional_emulator' && cargo test --workspace" + +cd "$repo_root" +python_bin="${PLENA_PYTHON:-python3}" +if [[ -n "${PLENA_VENV:-}" ]]; then + source "$PLENA_VENV/bin/activate" + python_bin="${PLENA_PYTHON:-python}" +else + default_venv="$repo_root/../../venvs/plena-py311" + if [[ -z "${PLENA_PYTHON:-}" && -d "$default_venv" ]]; then + source "$default_venv/bin/activate" + python_bin="python" + fi +fi +export PYTHONPATH="$repo_root:$repo_root/PLENA_Compiler:$repo_root/PLENA_Tools:$repo_root/transactional_emulator/testbench:${PYTHONPATH:-}" +"$python_bin" transactional_emulator/testbench/timing_goldens/check_timing_goldens.py diff --git a/transactional_emulator/testbench/window1_p1/__init__.py b/transactional_emulator/testbench/window1_p1/__init__.py new file mode 100644 index 00000000..d79ee110 --- /dev/null +++ b/transactional_emulator/testbench/window1_p1/__init__.py @@ -0,0 +1,2 @@ +"""Window 1 P1 timing/trace sidecar tools.""" + diff --git a/transactional_emulator/testbench/window1_p1/create_sample_trace.py b/transactional_emulator/testbench/window1_p1/create_sample_trace.py new file mode 100644 index 00000000..40e49fc2 --- /dev/null +++ b/transactional_emulator/testbench/window1_p1/create_sample_trace.py @@ -0,0 +1,158 @@ +#!/usr/bin/env python3 +"""Create the GPT-OSS layer0 tok1 sample route trace used by Window 1 P1.""" + +from __future__ import annotations + +import argparse +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +from transactional_emulator.testbench.window1_p1.p1_utils import ( + DEFAULT_OUT_ROOT, + gini_from_counts, + load_json, + summarize_run, + write_json, +) +from transactional_emulator.testbench.window1_p1.validate_route_trace import validate_trace + + +def _to_nested_float(values: Any) -> list[list[float]]: + return [[float(x) for x in row] for row in values.tolist()] + + +def _to_nested_int(values: Any) -> list[list[int]]: + return [[int(x) for x in row] for row in values.tolist()] + + +def _routing_stats(indices: list[list[int]], num_experts: int) -> dict[str, Any]: + counts = [0 for _ in range(num_experts)] + for row in indices: + for expert_id in row: + counts[expert_id] += 1 + nonzero = [idx for idx, value in enumerate(counts) if value > 0] + total_assignments = sum(counts) + duplicate_factor = float(total_assignments) / float(len(nonzero)) if nonzero else 0.0 + return { + "expert_counts": counts, + "active_experts": nonzero, + "hot_experts": sorted(nonzero, key=lambda idx: (-counts[idx], idx)), + "cold_experts": [idx for idx, value in enumerate(counts) if value == 0], + "duplicate_factor": duplicate_factor, + "gini": gini_from_counts(counts), + } + + +def create_trace(args: argparse.Namespace) -> dict[str, Any]: + import torch + + reference = torch.load(args.reference_path, map_location="cpu") + if not isinstance(reference, dict): + raise TypeError(f"reference must be a dict-like .pt file: {args.reference_path}") + topk_indices = _to_nested_int(reference["topk_indices"]) + topk_weights = _to_nested_float(reference["topk_weights"]) + x = reference["x"] + rows = int(x.shape[0]) + hidden = int(x.shape[1]) + model = reference.get("metadata", {}).get("model", {}) if isinstance(reference.get("metadata"), dict) else {} + num_experts = int(args.num_experts) + top_k = int(args.top_k or len(topk_indices[0])) + stats = _routing_stats(topk_indices, num_experts) + + fixed_run, fixed_stack = summarize_run("fixed_route_direct", args.fixed_route_build_dir) + device_run, device_stack = summarize_run("device_selected_direct", args.device_selected_build_dir) + + trace = { + "schema_version": 1, + "trace_id": args.trace_id, + "created_by": "transactional_emulator.testbench.window1_p1.create_sample_trace", + "created_at_utc": datetime.now(UTC).isoformat(), + "model": { + "name": args.model_name, + "layer_index": args.layer_index, + "hidden_size": hidden, + "intermediate_size": int(args.intermediate_size), + "num_experts": num_experts, + "top_k": top_k, + "metadata_model": model, + }, + "workload": { + "benchmark": args.benchmark, + "sample_id": args.sample_id, + "phase": args.phase, + "batch_size": args.batch_size, + "seq_len": args.seq_len, + "token_count": rows, + }, + "routing": { + "source": "hf_reference_router_logits_topk", + "topk_indices": topk_indices, + "topk_weights": topk_weights, + **stats, + }, + "logical_bytes": { + "input_hidden_bf16_bytes": rows * hidden * 2, + "routed_hidden_bf16_bytes": rows * top_k * hidden * 2, + "expert_weight_mxfp8_table_note": "Logical expert-weight bytes are model/layout dependent; physical HBM bytes are taken from emulator WithStats.", + }, + "artifacts": { + "reference_pt": str(args.reference_path), + "l1_golden_pt": str(args.l1_golden_path), + "fixed_route_build_dir": str(args.fixed_route_build_dir), + "device_selected_build_dir": str(args.device_selected_build_dir), + }, + "replay": { + "harness_module": "transactional_emulator.testbench.routed_moe.gpt_oss_moe_gather_scatter_test", + "stage": "full_vram", + "mlen": args.mlen, + "blen": args.blen, + "emu_threads": args.emu_threads, + }, + "direct_runs": { + "fixed_route": fixed_run, + "device_selected": device_run, + "routing_tax_cycles": int(device_run["sim_latency_cycles"]) - int(fixed_run["sim_latency_cycles"]), + "routing_tax_hbm_bytes_read": int(device_run["hbm_bytes_read"]) - int(fixed_run["hbm_bytes_read"]), + }, + "direct_stage_stack_rows": fixed_stack + device_stack, + } + errors = validate_trace(trace) + if errors: + raise ValueError("Generated invalid trace:\n" + "\n".join(errors)) + return trace + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--reference-path", type=Path, default=DEFAULT_OUT_ROOT.parent / "real_workload" / "gpt_oss_layer0_tok1" / "reference" / "hf_layer0_moe_reference.pt") + parser.add_argument("--l1-golden-path", type=Path, default=DEFAULT_OUT_ROOT.parent / "real_workload" / "gpt_oss_layer0_tok1" / "emulator" / "golden_output.pt") + parser.add_argument("--fixed-route-build-dir", type=Path, default=DEFAULT_OUT_ROOT / "gpt_oss_layer0_tok1_fixed_route") + parser.add_argument("--device-selected-build-dir", type=Path, default=DEFAULT_OUT_ROOT / "gpt_oss_layer0_tok1_device_selected") + parser.add_argument("--out", type=Path, default=DEFAULT_OUT_ROOT / "route_traces" / "gpt_oss_layer0_tok1_trace.json") + parser.add_argument("--trace-id", default="gpt_oss_layer0_tok1_decode") + parser.add_argument("--model-name", default="gpt-oss-20b") + parser.add_argument("--benchmark", default="local_real_workload") + parser.add_argument("--sample-id", default="gpt_oss_layer0_tok1") + parser.add_argument("--phase", choices=("prefill", "decode", "smoke", "unknown"), default="decode") + parser.add_argument("--batch-size", type=int, default=1) + parser.add_argument("--seq-len", type=int, default=1) + parser.add_argument("--layer-index", type=int, default=0) + parser.add_argument("--intermediate-size", type=int, default=2880) + parser.add_argument("--num-experts", type=int, default=32) + parser.add_argument("--top-k", type=int, default=None) + parser.add_argument("--mlen", type=int, default=64) + parser.add_argument("--blen", type=int, default=4) + parser.add_argument("--emu-threads", type=int, default=1) + args = parser.parse_args() + + trace = create_trace(args) + write_json(args.out, trace) + print(f"wrote route trace: {args.out}") + print(load_json(args.out)["direct_runs"]) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) + diff --git a/transactional_emulator/testbench/window1_p1/export_timing_results.py b/transactional_emulator/testbench/window1_p1/export_timing_results.py new file mode 100644 index 00000000..038ae5bb --- /dev/null +++ b/transactional_emulator/testbench/window1_p1/export_timing_results.py @@ -0,0 +1,101 @@ +#!/usr/bin/env python3 +"""Export Window 1 P1 MoE timing runs to CSV/JSON tables.""" + +from __future__ import annotations + +import argparse +from pathlib import Path +from typing import Any + +from transactional_emulator.testbench.window1_p1.p1_utils import ( + DEFAULT_OUT_ROOT, + load_json, + summarize_run, + write_csv, + write_json, +) + + +def _default_runs(root: Path) -> list[tuple[str, Path]]: + runs = [ + ("fixed_route_direct", root / "gpt_oss_layer0_tok1_fixed_route"), + ("device_selected_direct", root / "gpt_oss_layer0_tok1_device_selected"), + ("trace_replay", root / "trace_replay" / "gpt_oss_layer0_tok1_decode"), + ] + return [(name, path) for name, path in runs if (path / "rust_emulator_run_stats.json").exists()] + + +def export_results(args: argparse.Namespace) -> dict[str, Any]: + root = args.root + run_specs = _default_runs(root) + rows: list[dict[str, Any]] = [] + stage_rows: list[dict[str, Any]] = [] + for run_id, build_dir in run_specs: + row, stack = summarize_run(run_id, build_dir) + rows.append(row) + stage_rows.extend(stack) + + by_id = {row["run_id"]: row for row in rows} + tax_rows = [] + if "fixed_route_direct" in by_id and "device_selected_direct" in by_id: + fixed = by_id["fixed_route_direct"] + device = by_id["device_selected_direct"] + tax_rows.append( + { + "tax_name": "device_selected_minus_fixed_route", + "base_run_id": "fixed_route_direct", + "routed_run_id": "device_selected_direct", + "routing_tax_cycles": int(device["sim_latency_cycles"]) - int(fixed["sim_latency_cycles"]), + "routing_tax_hbm_bytes_read": int(device["hbm_bytes_read"]) - int(fixed["hbm_bytes_read"]), + "routing_tax_hbm_bytes_written": int(device["hbm_bytes_written"]) - int(fixed["hbm_bytes_written"]), + "device_topk_in_emulator": True, + } + ) + + trace_path = root / "route_traces" / "gpt_oss_layer0_tok1_trace.json" + bytes_rows: list[dict[str, Any]] = [] + if trace_path.exists(): + trace = load_json(trace_path) + logical = trace.get("logical_bytes", {}) + for row in rows: + bytes_rows.append( + { + "run_id": row["run_id"], + "logical_input_hidden_bf16_bytes": logical.get("input_hidden_bf16_bytes"), + "logical_routed_hidden_bf16_bytes": logical.get("routed_hidden_bf16_bytes"), + "physical_hbm_bytes_read": row.get("hbm_bytes_read"), + "physical_hbm_bytes_written": row.get("hbm_bytes_written"), + "physical_source": "rust_emulator WithStats 64B transfer counters", + } + ) + + payload = { + "schema_version": 1, + "root": str(root), + "runs": rows, + "routing_tax": tax_rows, + "stage_stack": stage_rows, + "bytes_validation": bytes_rows, + } + write_json(args.out_json, payload) + write_csv(args.out_prefix.with_name(args.out_prefix.name + "_timing.csv"), rows) + write_csv(args.out_prefix.with_name(args.out_prefix.name + "_routing_tax.csv"), tax_rows) + write_csv(args.out_prefix.parent / "gpqa_bfcl_routing_tax.csv", tax_rows) + write_csv(args.out_prefix.with_name(args.out_prefix.name + "_stage_stack.csv"), stage_rows) + write_csv(args.out_prefix.with_name(args.out_prefix.name + "_bytes_validation.csv"), bytes_rows) + return payload + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--root", type=Path, default=DEFAULT_OUT_ROOT) + parser.add_argument("--out-json", type=Path, default=DEFAULT_OUT_ROOT / "gpqa_bfcl_emulation_timing_summary.json") + parser.add_argument("--out-prefix", type=Path, default=DEFAULT_OUT_ROOT / "gpqa_bfcl_emulation") + args = parser.parse_args() + payload = export_results(args) + print(f"exported {len(payload['runs'])} runs and {len(payload['stage_stack'])} stage rows") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/transactional_emulator/testbench/window1_p1/p1_utils.py b/transactional_emulator/testbench/window1_p1/p1_utils.py new file mode 100644 index 00000000..df7253ac --- /dev/null +++ b/transactional_emulator/testbench/window1_p1/p1_utils.py @@ -0,0 +1,146 @@ +from __future__ import annotations + +import csv +import json +import math +import os +import sys +from pathlib import Path +from typing import Any + + +REPO_ROOT = Path(__file__).resolve().parents[3] +PLENA_ROOT = REPO_ROOT.parents[1] +COMPILER_ROOT = REPO_ROOT / "PLENA_Compiler" +TOOLS_ROOT = REPO_ROOT / "PLENA_Tools" +TESTBENCH_ROOT = REPO_ROOT / "transactional_emulator" / "testbench" +DEFAULT_OUT_ROOT = PLENA_ROOT / "outputs" / "window1_p1" + + +def ensure_python_paths() -> None: + for path in (REPO_ROOT, COMPILER_ROOT, TOOLS_ROOT, TESTBENCH_ROOT): + text = str(path) + if text not in sys.path: + sys.path.insert(0, text) + os.environ.setdefault("CUDA_VISIBLE_DEVICES", "") + os.environ["PYTHONPATH"] = ":".join( + [str(REPO_ROOT), str(COMPILER_ROOT), str(TOOLS_ROOT), str(TESTBENCH_ROOT), os.environ.get("PYTHONPATH", "")] + ) + + +def load_json(path: Path) -> dict[str, Any]: + with path.open(encoding="utf-8") as f: + return json.load(f) + + +def write_json(path: Path, data: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(data, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + +def write_csv(path: Path, rows: list[dict[str, Any]]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + if not rows: + path.write_text("", encoding="utf-8") + return + fieldnames: list[str] = [] + for row in rows: + for key in row: + if key not in fieldnames: + fieldnames.append(key) + with path.open("w", newline="", encoding="utf-8") as f: + writer = csv.DictWriter(f, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(rows) + + +def gini_from_counts(counts: list[int]) -> float: + if not counts: + return 0.0 + total = sum(counts) + if total == 0: + return 0.0 + values = sorted(float(x) for x in counts) + n = len(values) + weighted = sum((idx + 1) * value for idx, value in enumerate(values)) + return (2.0 * weighted) / (n * total) - (n + 1.0) / n + + +def finite_number(value: Any) -> bool: + return isinstance(value, (int, float)) and math.isfinite(float(value)) + + +def run_result_path(build_dir: Path) -> Path | None: + for name in ("device_routing_results.json", "gather_scatter_results.json", "real_layer0_results.json"): + path = build_dir / name + if path.exists(): + return path + return None + + +def load_run_bundle(build_dir: Path) -> dict[str, Any]: + stats_path = build_dir / "rust_emulator_run_stats.json" + stage_path = build_dir / "stage_profile.json" + result_path = run_result_path(build_dir) + return { + "build_dir": str(build_dir), + "run_stats_path": str(stats_path) if stats_path.exists() else None, + "stage_profile_path": str(stage_path) if stage_path.exists() else None, + "result_path": str(result_path) if result_path else None, + "run_stats": load_json(stats_path) if stats_path.exists() else {}, + "stage_profile": load_json(stage_path) if stage_path.exists() else {}, + "result": load_json(result_path) if result_path else {}, + } + + +def summarize_run(run_id: str, build_dir: Path) -> tuple[dict[str, Any], list[dict[str, Any]]]: + bundle = load_run_bundle(build_dir) + stats = bundle["run_stats"] + profile = bundle["stage_profile"] + result = bundle["result"] + gate = ( + result.get("device_routing_gate") + or result.get("full_vram_gate") + or result.get("hbm_store_gate") + or {} + ) + row = { + "run_id": run_id, + "build_dir": str(build_dir), + "stage": result.get("stage"), + "sim_latency_cycles": stats.get("sim_latency_cycles"), + "hbm_bytes_read": stats.get("hbm_bytes_read"), + "hbm_bytes_written": stats.get("hbm_bytes_written"), + "stage_total_simulation_cycles": profile.get("total_simulation_cycles"), + "stage_total_wall_cycles": profile.get("total_stage_wall_cycles"), + "cycle_accounting_status": profile.get("cycle_accounting_status"), + "physical_byte_status": profile.get("physical_byte_status"), + "functional_gate_passed": gate.get("passed"), + "result_path": bundle["result_path"], + "stage_profile_path": bundle["stage_profile_path"], + "run_stats_path": bundle["run_stats_path"], + } + stage_rows: list[dict[str, Any]] = [] + stages = profile.get("stages", {}) + if isinstance(stages, dict): + for stage_name, stage_data in stages.items(): + resources = stage_data.get("resource_proxy_cycles", {}) + stage_rows.append( + { + "run_id": run_id, + "stage": stage_name, + "instructions": stage_data.get("instructions"), + "wall_cycles": stage_data.get("wall_cycles"), + "cycle_fraction": stage_data.get("cycle_fraction"), + "physical_hbm_bytes_read": stage_data.get("physical_hbm_bytes_read", stage_data.get("hbm_bytes_read")), + "physical_hbm_bytes_written": stage_data.get("physical_hbm_bytes_written", stage_data.get("hbm_bytes_written")), + "matrix_cycles": resources.get("matrix"), + "vector_cycles": resources.get("vector"), + "scalar_cycles": resources.get("scalar"), + "dma_cycles": resources.get("dma"), + "ramulator_proxy_cycles": resources.get("ramulator_proxy"), + "other_cycles": resources.get("other"), + } + ) + return row, stage_rows + diff --git a/transactional_emulator/testbench/window1_p1/route_trace_schema.json b/transactional_emulator/testbench/window1_p1/route_trace_schema.json new file mode 100644 index 00000000..12018889 --- /dev/null +++ b/transactional_emulator/testbench/window1_p1/route_trace_schema.json @@ -0,0 +1,85 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "plena.route_trace.v1", + "title": "PLENA MoE route trace v1", + "type": "object", + "required": [ + "schema_version", + "trace_id", + "created_by", + "model", + "workload", + "routing", + "artifacts", + "replay" + ], + "properties": { + "schema_version": { "const": 1 }, + "trace_id": { "type": "string" }, + "created_by": { "type": "string" }, + "model": { + "type": "object", + "required": ["name", "layer_index", "hidden_size", "intermediate_size", "num_experts", "top_k"], + "properties": { + "name": { "type": "string" }, + "layer_index": { "type": "integer", "minimum": 0 }, + "hidden_size": { "type": "integer", "minimum": 1 }, + "intermediate_size": { "type": "integer", "minimum": 1 }, + "num_experts": { "type": "integer", "minimum": 1 }, + "top_k": { "type": "integer", "minimum": 1 } + } + }, + "workload": { + "type": "object", + "required": ["benchmark", "sample_id", "phase", "batch_size", "seq_len", "token_count"], + "properties": { + "benchmark": { "type": "string" }, + "sample_id": { "type": "string" }, + "phase": { "enum": ["prefill", "decode", "smoke", "unknown"] }, + "batch_size": { "type": "integer", "minimum": 1 }, + "seq_len": { "type": "integer", "minimum": 1 }, + "token_count": { "type": "integer", "minimum": 1 } + } + }, + "routing": { + "type": "object", + "required": ["source", "topk_indices", "topk_weights", "expert_counts", "duplicate_factor", "gini"], + "properties": { + "source": { "type": "string" }, + "topk_indices": { + "type": "array", + "items": { "type": "array", "items": { "type": "integer", "minimum": 0 } } + }, + "topk_weights": { + "type": "array", + "items": { "type": "array", "items": { "type": "number" } } + }, + "expert_counts": { "type": "array", "items": { "type": "integer", "minimum": 0 } }, + "duplicate_factor": { "type": "number" }, + "gini": { "type": "number" } + } + }, + "artifacts": { + "type": "object", + "required": ["reference_pt", "l1_golden_pt"], + "properties": { + "reference_pt": { "type": "string" }, + "l1_golden_pt": { "type": "string" }, + "fixed_route_build_dir": { "type": "string" }, + "device_selected_build_dir": { "type": "string" } + } + }, + "replay": { + "type": "object", + "required": ["harness_module", "stage", "mlen", "blen", "emu_threads"], + "properties": { + "harness_module": { "type": "string" }, + "stage": { "const": "full_vram" }, + "mlen": { "type": "integer", "minimum": 1 }, + "blen": { "type": "integer", "minimum": 1 }, + "emu_threads": { "type": "integer", "minimum": 1 } + } + } + } +} + diff --git a/transactional_emulator/testbench/window1_p1/run_trace_replay.py b/transactional_emulator/testbench/window1_p1/run_trace_replay.py new file mode 100644 index 00000000..74c25990 --- /dev/null +++ b/transactional_emulator/testbench/window1_p1/run_trace_replay.py @@ -0,0 +1,129 @@ +#!/usr/bin/env python3 +"""Replay a route trace through the existing fixed-route MoE emulator harness.""" + +from __future__ import annotations + +import argparse +import os +import subprocess +import sys +from pathlib import Path +from typing import Any + +from transactional_emulator.testbench.window1_p1.p1_utils import ( + DEFAULT_OUT_ROOT, + ensure_python_paths, + load_json, + summarize_run, + write_json, +) +from transactional_emulator.testbench.window1_p1.validate_route_trace import validate_trace + + +def _make_replay_reference(trace: dict[str, Any], build_dir: Path) -> Path: + import torch + + source_path = Path(trace["artifacts"]["reference_pt"]) + reference = torch.load(source_path, map_location="cpu") + if not isinstance(reference, dict): + raise TypeError(f"reference artifact must be a dict-like .pt file: {source_path}") + replay_reference = dict(reference) + replay_reference["topk_indices"] = torch.tensor(trace["routing"]["topk_indices"], dtype=torch.long) + replay_reference["topk_weights"] = torch.tensor(trace["routing"]["topk_weights"], dtype=torch.bfloat16) + path = build_dir / "replay_reference.pt" + build_dir.mkdir(parents=True, exist_ok=True) + torch.save(replay_reference, path) + return path + + +def replay_trace(args: argparse.Namespace) -> dict[str, Any]: + ensure_python_paths() + trace = load_json(args.trace) + errors = validate_trace(trace) + if errors: + raise ValueError("Invalid route trace:\n" + "\n".join(errors)) + + build_dir = args.build_dir or (DEFAULT_OUT_ROOT / "trace_replay" / trace["trace_id"]) + build_dir.mkdir(parents=True, exist_ok=True) + replay_reference = _make_replay_reference(trace, build_dir) + + replay = trace["replay"] + cmd = [ + sys.executable, + "-m", + replay["harness_module"], + "--stage", + replay["stage"], + "--reference-path", + str(replay_reference), + "--l1-golden-path", + trace["artifacts"]["l1_golden_pt"], + "--build-dir", + str(build_dir), + "--mlen", + str(replay["mlen"]), + "--blen", + str(replay["blen"]), + "--emu-threads", + str(replay["emu_threads"]), + ] + env = os.environ.copy() + if args.stage_profile: + env["PLENA_EMULATOR_STAGE_PROFILE"] = "1" + proc = subprocess.run(cmd, cwd=str(Path(__file__).resolve().parents[3]), env=env, text=True) + if proc.returncode != 0: + raise subprocess.CalledProcessError(proc.returncode, cmd) + + replay_run, replay_stack = summarize_run("trace_replay", build_dir) + direct = trace.get("direct_runs", {}).get("fixed_route", {}) + self_compare = { + "against": "fixed_route_direct", + "direct_cycles": direct.get("sim_latency_cycles"), + "replay_cycles": replay_run.get("sim_latency_cycles"), + "cycles_match": direct.get("sim_latency_cycles") == replay_run.get("sim_latency_cycles"), + "direct_hbm_bytes_read": direct.get("hbm_bytes_read"), + "replay_hbm_bytes_read": replay_run.get("hbm_bytes_read"), + "hbm_read_match": direct.get("hbm_bytes_read") == replay_run.get("hbm_bytes_read"), + "direct_hbm_bytes_written": direct.get("hbm_bytes_written"), + "replay_hbm_bytes_written": replay_run.get("hbm_bytes_written"), + "hbm_written_match": direct.get("hbm_bytes_written") == replay_run.get("hbm_bytes_written"), + } + summary = { + "schema_version": 1, + "trace_path": str(args.trace), + "trace_id": trace["trace_id"], + "build_dir": str(build_dir), + "replay_reference": str(replay_reference), + "command": cmd, + "replay_run": replay_run, + "replay_stage_stack_rows": replay_stack, + "self_compare": self_compare, + "passed": all( + [ + replay_run.get("functional_gate_passed") is True, + self_compare["cycles_match"], + self_compare["hbm_read_match"], + self_compare["hbm_written_match"], + ] + ), + } + write_json(build_dir / "trace_replay_summary.json", summary) + return summary + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("trace", type=Path) + parser.add_argument("--build-dir", type=Path) + parser.add_argument("--no-stage-profile", dest="stage_profile", action="store_false") + parser.set_defaults(stage_profile=True) + args = parser.parse_args() + summary = replay_trace(args) + print(summary["self_compare"]) + if not summary["passed"]: + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/transactional_emulator/testbench/window1_p1/timing_validation_gates.py b/transactional_emulator/testbench/window1_p1/timing_validation_gates.py new file mode 100644 index 00000000..27536da5 --- /dev/null +++ b/transactional_emulator/testbench/window1_p1/timing_validation_gates.py @@ -0,0 +1,347 @@ +#!/usr/bin/env python3 +"""Window 1 P1 timing validation gates G1/G2/G5.""" + +from __future__ import annotations + +import argparse +import random +from pathlib import Path +from typing import Any + +from transactional_emulator.testbench.window1_p1.p1_utils import DEFAULT_OUT_ROOT, ensure_python_paths, load_json, write_json + + +ensure_python_paths() + +from compiler.assembler.assembly_to_binary import AssemblyToBinary # noqa: E402 +from transactional_emulator.testbench.aten.configurable import HardwareConfig # noqa: E402 +from transactional_emulator.testbench.emulator_runner import run_emulator # noqa: E402 + + +MLEN = 64 +VLEN = 64 +BLEN = 4 +HLEN = 16 +BROADCAST_AMOUNT = 4 +HBM_SIZE = 1 << 20 +SCALE_BASE = 65_536 +PREFETCH_COUNT = 64 +COMPUTE_REPS = 100 + + +def write_settings(build_dir: Path) -> Path: + hw = HardwareConfig( + mlen=MLEN, + vlen=VLEN, + blen=BLEN, + hlen=HLEN, + broadcast_amount=BROADCAST_AMOUNT, + dc_en=None, + latency_profile=None, + hbm_m_prefetch_amount=MLEN, + hbm_v_prefetch_amount=BLEN, + hbm_v_writeback_amount=BLEN, + ) + return hw.write_toml(build_dir) + + +def assembler() -> AssemblyToBinary: + repo_root = Path(__file__).resolve().parents[3] + compiler_root = repo_root / "PLENA_Compiler" + return AssemblyToBinary( + str(compiler_root / "doc" / "operation.svh"), + str(compiler_root / "doc" / "configuration.svh"), + ) + + +def setup_prefetch(stride: int, *, dest: int = 4096) -> list[str]: + return [ + f"S_ADDI_INT gp4, gp0, {SCALE_BASE}", + "C_SET_SCALE_REG gp4", + f"S_ADDI_INT gp5, gp0, {stride}", + "C_SET_STRIDE_REG gp5", + f"S_ADDI_INT gp6, gp0, {dest}", + "S_ADDI_INT gp1, gp0, 0", + "C_SET_ADDR_REG gp3, gp0, gp1", + ] + + +def prefetch_offsets_program(name: str, *, stride: int, offsets: list[int]) -> str: + lines = [f"; {name}: H_PREFETCH_V address-pattern microbench"] + lines += setup_prefetch(stride) + for offset in offsets: + lines += [ + f"S_ADDI_INT gp2, gp0, {offset}", + "H_PREFETCH_V gp6, gp2, gp3, 1, 0", + ] + return "\n".join(lines) + + +def prefetch_one_program() -> str: + return "\n".join( + [ + "; G2 prefetch-only component", + *setup_prefetch(VLEN), + "S_ADDI_INT gp2, gp0, 0", + "H_PREFETCH_V gp6, gp2, gp3, 1, 0", + ] + ) + + +def compute_program(reps: int = COMPUTE_REPS) -> str: + return "\n".join(["; G2 compute-only independent M_MV sequence"] + ["M_MV gp0, gp0, gp0"] * reps) + + +def combined_prefetch_compute_program() -> str: + return "\n".join( + [ + "; G2 prefetch then independent compute", + *setup_prefetch(VLEN), + "S_ADDI_INT gp2, gp0, 0", + "H_PREFETCH_V gp6, gp2, gp3, 1, 0", + *["M_MV gp0, gp0, gp0" for _ in range(COMPUTE_REPS)], + ] + ) + + +def emit_case(out_root: Path, name: str, asm: str) -> Path: + build_dir = out_root / "timing_gate_runs" / name + build_dir.mkdir(parents=True, exist_ok=True) + settings = write_settings(build_dir) + (build_dir / "generated_asm_code.asm").write_text(asm.strip() + "\n", encoding="utf-8") + assembler().generate_binary( + build_dir / "generated_asm_code.asm", + build_dir / "generated_machine_code.mem", + ) + (build_dir / "hbm_for_behave_sim.bin").write_bytes(bytes(HBM_SIZE)) + (build_dir / "hbm_size.txt").write_text(f"{HBM_SIZE}\n", encoding="utf-8") + (build_dir / "fp_sram.bin").write_bytes(bytes(2048)) + (build_dir / "int_sram.bin").write_bytes(bytes(4096)) + return settings + + +def run_case( + out_root: Path, + name: str, + asm: str, + *, + overlap_prefetch_compute: bool = False, + stage_profile: bool = True, +) -> dict[str, Any]: + build_dir = out_root / "timing_gate_runs" / name + settings = emit_case(out_root, name, asm) + import os + + old_settings = os.environ.get("PLENA_SETTINGS_TOML") + os.environ["PLENA_SETTINGS_TOML"] = str(settings) + try: + metrics = run_emulator( + build_dir, + hbm_size=HBM_SIZE, + threads=1, + stage_profile=stage_profile, + overlap_prefetch_compute=overlap_prefetch_compute, + ) + finally: + if old_settings is None: + os.environ.pop("PLENA_SETTINGS_TOML", None) + else: + os.environ["PLENA_SETTINGS_TOML"] = old_settings + stage_profile_path = Path(metrics.get("stage_profile_path", build_dir / "stage_profile.json")) + result = { + "name": name, + "build_dir": str(build_dir), + "sim_latency_cycles": metrics.get("sim_latency_cycles"), + "hbm_bytes_read": metrics.get("hbm_bytes_read"), + "hbm_bytes_written": metrics.get("hbm_bytes_written"), + "overlap_prefetch_compute": bool(overlap_prefetch_compute), + "stage_profile_requested": bool(stage_profile), + "run_stats_path": metrics.get("stats_path"), + "stage_profile_path": str(stage_profile_path), + "stage_profile": load_json(stage_profile_path) if stage_profile_path.exists() else {}, + } + write_json(build_dir / "timing_gate_case_result.json", result) + return result + + +def g1_cases() -> dict[str, str]: + continuous_offsets = [idx * (BLEN * VLEN) for idx in range(PREFETCH_COUNT)] + # Keep offsets within S_ADDI_INT's compact immediate range. + sparse_offsets = [idx * 4096 for idx in range(PREFETCH_COUNT)] + rng = random.Random(20260701) + random_offsets = continuous_offsets[:] + rng.shuffle(random_offsets) + row_hit_offsets = [0 for _ in range(PREFETCH_COUNT)] + return { + "g1_continuous_prefetch_v": prefetch_offsets_program( + "continuous", + stride=VLEN, + offsets=continuous_offsets, + ), + "g1_sparse_cross_region_prefetch_v": prefetch_offsets_program( + "sparse-cross-region", + stride=VLEN, + offsets=sparse_offsets, + ), + "g1_random_order_prefetch_v": prefetch_offsets_program( + "random-order-same-address-set", + stride=VLEN, + offsets=random_offsets, + ), + "g1_row_hit_repeat_prefetch_v": prefetch_offsets_program( + "row-hit-repeat", + stride=VLEN, + offsets=row_hit_offsets, + ), + } + + +def evaluate_gates(rows: list[dict[str, Any]]) -> dict[str, Any]: + by_name = {row["name"]: row for row in rows} + g1_candidates = [ + by_name["g1_sparse_cross_region_prefetch_v"], + by_name["g1_random_order_prefetch_v"], + ] + continuous = by_name["g1_continuous_prefetch_v"] + row_hit = by_name["g1_row_hit_repeat_prefetch_v"] + same_bytes = all(row["hbm_bytes_read"] == continuous["hbm_bytes_read"] for row in g1_candidates + [row_hit]) + random_gt_continuous = any(row["sim_latency_cycles"] > continuous["sim_latency_cycles"] for row in g1_candidates) + random_gt_row_hit = any(row["sim_latency_cycles"] > row_hit["sim_latency_cycles"] for row in g1_candidates) + address_sensitive = len({row["sim_latency_cycles"] for row in g1_candidates + [continuous, row_hit]}) > 1 + + prefetch_only = by_name["g2_prefetch_only"] + compute_only = by_name["g2_compute_only"] + combined = by_name["g2_prefetch_then_compute"] + combined_overlap = by_name["g2_prefetch_then_compute_overlap_on"] + prefetch_plus_compute = int(prefetch_only["sim_latency_cycles"]) + int(compute_only["sim_latency_cycles"]) + serial_hidden_cycles = prefetch_plus_compute - int(combined["sim_latency_cycles"]) + overlay_hidden_cycles = int(combined["sim_latency_cycles"]) - int(combined_overlap["sim_latency_cycles"]) + overlay_hbm_bytes_match = ( + combined["hbm_bytes_read"] == combined_overlap["hbm_bytes_read"] + and combined["hbm_bytes_written"] == combined_overlap["hbm_bytes_written"] + ) + + m_mv = by_name["g5_single_m_mv"] + m_mm = by_name["g5_single_m_mm"] + + profiled_rows = [row for row in rows if row.get("stage_profile_requested")] + bytes_match = all( + row["hbm_bytes_read"] == row["stage_profile"].get("total_hbm_bytes_read") + and row["hbm_bytes_written"] == row["stage_profile"].get("total_hbm_bytes_written") + for row in profiled_rows + ) + stage_cycles_match = all( + row["sim_latency_cycles"] == row["stage_profile"].get("total_simulation_cycles") + == row["stage_profile"].get("total_stage_wall_cycles") + for row in profiled_rows + ) + + return { + "g1_ramulator_address_sensitivity": { + "continuous_cycles": continuous["sim_latency_cycles"], + "row_hit_repeat_cycles": row_hit["sim_latency_cycles"], + "candidates": [ + { + "name": row["name"], + "cycles": row["sim_latency_cycles"], + "hbm_bytes_read": row["hbm_bytes_read"], + } + for row in g1_candidates + ], + "same_hbm_read_bytes": same_bytes, + "random_or_sparse_gt_continuous": random_gt_continuous, + "random_or_sparse_gt_row_hit_locality": random_gt_row_hit, + "address_sensitive": address_sensitive, + "pass": same_bytes and random_gt_continuous, + "note": "Strict pass requires random/sparse same-byte access to be slower than the continuous unique baseline. Row-hit locality is reported separately as weaker evidence.", + }, + "g2_prefetch_compute_overlap": { + "prefetch_only_cycles": prefetch_only["sim_latency_cycles"], + "compute_only_cycles": compute_only["sim_latency_cycles"], + "prefetch_plus_compute_cycles": prefetch_plus_compute, + "default_combined_cycles": combined["sim_latency_cycles"], + "overlap_enabled_combined_cycles": combined_overlap["sim_latency_cycles"], + "default_hidden_cycles": serial_hidden_cycles, + "overlay_hidden_cycles": overlay_hidden_cycles, + "default_serial_expected_cycles": 6440, + "overlap_hbm_bytes_match_default": overlay_hbm_bytes_match, + "pass": ( + int(combined["sim_latency_cycles"]) == 6440 + and serial_hidden_cycles == 0 + and overlay_hidden_cycles > 0 + and int(combined_overlap["sim_latency_cycles"]) < int(combined["sim_latency_cycles"]) + and overlay_hbm_bytes_match + ), + "note": "Default must remain serial; only --experimental-overlap-prefetch-compute may reduce cycles for this independent prefetch+compute pattern.", + }, + "g5_matrix_formula_self_consistency": { + "mlen": MLEN, + "expected_m_mv_cycles": MLEN, + "measured_m_mv_cycles": m_mv["sim_latency_cycles"], + "expected_m_mm_cycles": MLEN, + "measured_m_mm_cycles": m_mm["sim_latency_cycles"], + "pass": m_mv["sim_latency_cycles"] == MLEN and m_mm["sim_latency_cycles"] == MLEN, + }, + "stage_profile_accounting": { + "stage_cycles_match_total": stage_cycles_match, + "physical_hbm_bytes_match_run_stats": bytes_match, + "pass": stage_cycles_match and bytes_match, + }, + } + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--out-root", type=Path, default=DEFAULT_OUT_ROOT) + args = parser.parse_args() + + cases: dict[str, str] = {} + cases.update(g1_cases()) + cases.update( + { + "g2_prefetch_only": prefetch_one_program(), + "g2_compute_only": compute_program(), + "g2_prefetch_then_compute": combined_prefetch_compute_program(), + "g5_single_m_mv": "M_MV gp0, gp0, gp0\n", + "g5_single_m_mm": "M_MM gp0, gp0, gp0\n", + } + ) + rows = [run_case(args.out_root, name, asm) for name, asm in cases.items()] + rows.append( + run_case( + args.out_root, + "g2_prefetch_then_compute_overlap_on", + combined_prefetch_compute_program(), + overlap_prefetch_compute=True, + stage_profile=False, + ) + ) + gates = evaluate_gates(rows) + summary = { + "schema_version": 1, + "output_root": str(args.out_root), + "hardware": { + "mlen": MLEN, + "vlen": VLEN, + "blen": BLEN, + "hlen": HLEN, + "broadcast_amount": BROADCAST_AMOUNT, + "hbm_size": HBM_SIZE, + "prefetch_count": PREFETCH_COUNT, + "compute_reps": COMPUTE_REPS, + }, + "cases": [ + {key: value for key, value in row.items() if key != "stage_profile"} + for row in rows + ], + "gates": gates, + } + out_path = args.out_root / "timing_gates_summary.json" + write_json(out_path, summary) + print(out_path) + print(gates) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/transactional_emulator/testbench/window1_p1/validate_route_trace.py b/transactional_emulator/testbench/window1_p1/validate_route_trace.py new file mode 100644 index 00000000..63e0257d --- /dev/null +++ b/transactional_emulator/testbench/window1_p1/validate_route_trace.py @@ -0,0 +1,128 @@ +#!/usr/bin/env python3 +"""Validate a PLENA MoE route trace contract without external dependencies.""" + +from __future__ import annotations + +import argparse +from pathlib import Path +from typing import Any + +from transactional_emulator.testbench.window1_p1.p1_utils import finite_number, load_json, write_json + + +def _require_mapping(parent: dict[str, Any], key: str, errors: list[str]) -> dict[str, Any]: + value = parent.get(key) + if not isinstance(value, dict): + errors.append(f"{key} must be an object") + return {} + return value + + +def validate_trace(trace: dict[str, Any], *, allow_missing_artifacts: bool = False) -> list[str]: + errors: list[str] = [] + if trace.get("schema_version") != 1: + errors.append("schema_version must be 1") + for key in ("trace_id", "created_by"): + if not isinstance(trace.get(key), str) or not trace.get(key): + errors.append(f"{key} must be a non-empty string") + + model = _require_mapping(trace, "model", errors) + workload = _require_mapping(trace, "workload", errors) + routing = _require_mapping(trace, "routing", errors) + artifacts = _require_mapping(trace, "artifacts", errors) + replay = _require_mapping(trace, "replay", errors) + + top_k = model.get("top_k") + num_experts = model.get("num_experts") + token_count = workload.get("token_count") + for scope, fields in ( + (model, ("layer_index", "hidden_size", "intermediate_size", "num_experts", "top_k")), + (workload, ("batch_size", "seq_len", "token_count")), + (replay, ("mlen", "blen", "emu_threads")), + ): + for field in fields: + if not isinstance(scope.get(field), int) or scope[field] <= 0 and field != "layer_index": + errors.append(f"{field} must be a positive integer") + if isinstance(model.get("layer_index"), int) and model["layer_index"] < 0: + errors.append("layer_index must be >= 0") + + topk_indices = routing.get("topk_indices") + topk_weights = routing.get("topk_weights") + if not isinstance(topk_indices, list) or not isinstance(topk_weights, list): + errors.append("routing.topk_indices and routing.topk_weights must be arrays") + topk_indices = [] + topk_weights = [] + if isinstance(token_count, int) and len(topk_indices) != token_count: + errors.append(f"topk_indices row count {len(topk_indices)} != token_count {token_count}") + if isinstance(token_count, int) and len(topk_weights) != token_count: + errors.append(f"topk_weights row count {len(topk_weights)} != token_count {token_count}") + + for row_idx, row in enumerate(topk_indices): + if not isinstance(row, list): + errors.append(f"topk_indices[{row_idx}] must be an array") + continue + if isinstance(top_k, int) and len(row) != top_k: + errors.append(f"topk_indices[{row_idx}] length {len(row)} != top_k {top_k}") + for col_idx, expert_id in enumerate(row): + if not isinstance(expert_id, int): + errors.append(f"topk_indices[{row_idx}][{col_idx}] must be int") + elif isinstance(num_experts, int) and not (0 <= expert_id < num_experts): + errors.append(f"topk_indices[{row_idx}][{col_idx}]={expert_id} outside [0,{num_experts})") + + for row_idx, row in enumerate(topk_weights): + if not isinstance(row, list): + errors.append(f"topk_weights[{row_idx}] must be an array") + continue + if isinstance(top_k, int) and len(row) != top_k: + errors.append(f"topk_weights[{row_idx}] length {len(row)} != top_k {top_k}") + for col_idx, weight in enumerate(row): + if not finite_number(weight): + errors.append(f"topk_weights[{row_idx}][{col_idx}] must be finite number") + + expert_counts = routing.get("expert_counts") + if not isinstance(expert_counts, list): + errors.append("routing.expert_counts must be an array") + elif isinstance(num_experts, int) and len(expert_counts) != num_experts: + errors.append(f"expert_counts length {len(expert_counts)} != num_experts {num_experts}") + elif sum(int(x) for x in expert_counts if isinstance(x, int)) != len(topk_indices) * int(top_k or 0): + errors.append("expert_counts sum must equal token_count * top_k") + + if replay.get("stage") != "full_vram": + errors.append("replay.stage must be full_vram for the current trace replay harness") + + if not allow_missing_artifacts: + for key in ("reference_pt", "l1_golden_pt"): + value = artifacts.get(key) + if not isinstance(value, str) or not Path(value).exists(): + errors.append(f"artifacts.{key} does not exist: {value}") + return errors + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("trace", type=Path) + parser.add_argument("--allow-missing-artifacts", action="store_true") + parser.add_argument("--summary-out", type=Path) + args = parser.parse_args() + + trace = load_json(args.trace) + errors = validate_trace(trace, allow_missing_artifacts=args.allow_missing_artifacts) + summary = { + "schema_version": 1, + "trace_path": str(args.trace), + "valid": not errors, + "errors": errors, + "trace_id": trace.get("trace_id"), + } + if args.summary_out: + write_json(args.summary_out, summary) + if errors: + print("\n".join(errors)) + return 1 + print(f"valid route trace: {args.trace}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) +