tokitai-operator 0.1.0

//! ROCm/HIP fp16 GEMM forward pilot (gated on `rocm-hip`).
//!
//! fp16 GEMM forward kernel. The largest HIP pilot; integrates
//! with the 0.7B MoE training project. Source/compiler fingerprint
//! and CPU oracle comparison; no ROCm/HIP bitcode verification.
//!
use std::collections::hash_map::DefaultHasher;
use std::fs;
use std::hash::{Hash, Hasher};
use std::path::PathBuf;

use crate::backend::hip_dense::{
    hipcc_compile_executable, hipcc_compiler_fingerprint, hipcc_recheck_artifact,
};
use crate::backend::kernel_server;
use crate::backend::rocm::{RocmHipCapabilityReport, detect_local_rocm_hip};
use crate::{Error, Result};

pub const ROCM_HIP_GEMM_F16_BACKEND: &str = "rocm_hip_gemm_f16_pilot";
pub const ROCM_HIP_GEMM_F16_LOWERING_ID: &str = "hip.gemm.fp16_f32";

/// Kernel-type label used by the persistent `KernelServer` pool to
/// route calls to the correct long-lived child. Keep this stable —
/// changing it forces a fresh spawn of the child on first call
/// after the change.
const GEMM_F16_KERNEL_TYPE: &str = "hip-gemm-f16-fwd";

pub const HIP_GEMM_F16_KERNEL: &str = r#"
#include <hip/hip_runtime.h>
#include <hip/hip_fp16.h>
#include <cstdint>
#include <cstdlib>
#include <iostream>
#include <sstream>
#include <string>
#include <vector>

// Naive fp16 GEMM: C = A (M x K) * B (K x N), M and N must be multiples of 16.
// grid = (M / 16, N / 16), block = (16, 16). Each thread computes a 1x1 output
// element with a K-loop accumulation in fp32. Boundary conversions use
// __half2float and __float2half_rn; the accumulator stays in fp32.
__global__ void gemm_fp16_f32_kernel(
    const __half* A,
    const __half* B,
    float* C,
    int M,
    int N,
    int K) {
    int row = blockIdx.x * blockDim.x + threadIdx.x;
    int col = blockIdx.y * blockDim.y + threadIdx.y;
    if (row >= M || col >= N) {
        return;
    }
    float acc = 0.0f;
    for (int kk = 0; kk < K; ++kk) {
        float a = __half2float(A[static_cast<int>(row) * K + kk]);
        float b = __half2float(B[kk * static_cast<int>(N) + col]);
        acc += a * b;
    }
    C[static_cast<int>(row) * N + col] = acc;
}

static void check(hipError_t status, const char* label) {
    if (status != hipSuccess) {
        std::cerr << "HIP_ERROR " << label << "=" << hipGetErrorString(status) << "\n";
        std::exit(10);
    }
}

// Forward declaration of the existing main() body, extracted into
// a static helper so the server-mode loop can call it on each
// request. The default `main()` also routes through this helper so
// the one-shot and server code paths share the same compute logic.
static int run_one_shot_from_main_body();

// Persistent server-mode protocol (used when this kernel is invoked
// with `--server` as argv[1]). In server mode the host writes a
// little-endian u32 payload_len followed by `payload_len` bytes of the
// existing text payload, then reads back a little-endian u32
// response_len followed by `response_len` bytes of the existing text
// response. This avoids the per-call process spawn + HIP runtime init
// cost (~400-500ms) that dominates the small-kernel training loop.
//
// The one-shot text mode (no --server) is unchanged so unit tests and
// the KernelServer::oneshot() path can still drive the binary directly.
static int run_server_mode() {
    while (true) {
        uint32_t payload_len = 0;
        std::cin.read(reinterpret_cast<char*>(&payload_len), 4);
        if (!std::cin || std::cin.gcount() == 0) {
            return 0;  // clean EOF
        }
        if (std::cin.gcount() != 4) {
            std::cerr << "server_mode: short read on payload_len (got "
                      << std::cin.gcount() << " bytes)\n";
            return 20;
        }
        std::vector<char> payload(payload_len);
        if (payload_len > 0) {
            std::cin.read(payload.data(), payload_len);
            if (static_cast<uint32_t>(std::cin.gcount()) != payload_len) {
                std::cerr << "server_mode: short read on payload (got "
                          << std::cin.gcount() << " of " << payload_len << ")\n";
                return 21;
            }
        }
        // Drive the existing one-shot logic by reconstructing a fake
        // stdin via a stringstream. This is the cleanest way to reuse
        // the existing kernel compute code (which reads from std::cin)
        // without refactoring it to take an arbitrary input stream.
        std::string payload_str(payload.begin(), payload.end());
        std::istringstream fake_stdin(payload_str);
        std::streambuf* old_buf = std::cin.rdbuf(fake_stdin.rdbuf());
        std::ostringstream captured;
        std::streambuf* old_cout = std::cout.rdbuf(captured.rdbuf());
        std::ostringstream captured_err;
        std::streambuf* old_cerr = std::cerr.rdbuf(captured_err.rdbuf());
        int rc = run_one_shot_from_main_body();
        std::cin.rdbuf(old_buf);
        std::cout.rdbuf(old_cout);
        std::cerr.rdbuf(old_cerr);
        std::string response = captured.str();
        if (rc != 0) {
            // Forward the captured stderr so the host can see why.
            std::string err_str = captured_err.str();
            response += err_str;
        }
        uint32_t response_len = static_cast<uint32_t>(response.size());
        std::cout.write(reinterpret_cast<const char*>(&response_len), 4);
        if (response_len > 0) {
            std::cout.write(response.data(), response_len);
        }
        std::cout.flush();
        if (rc != 0) {
            return rc;
        }
    }
}

int main(int argc, char** argv) {
    if (argc > 1 && std::string(argv[1]) == "--server") {
        return run_server_mode();
    }
    return run_one_shot_from_main_body();
}

static int run_one_shot_from_main_body() {
    int M = 0;
    int N = 0;
    int K = 0;
    if (!(std::cin >> M >> N >> K)) {
        std::cerr << "usage: stdin payload is \"M N K\\n<A_bits> <A_bits> ...\\n<B_bits> <B_bits> ...\\n\"\n";
        return 2;
    }
    if (M <= 0 || N <= 0 || K <= 0) {
        std::cerr << "M N K must all be positive\n";
        return 3;
    }
    if (M % 16 != 0 || N % 16 != 0) {
        std::cerr << "M=" << M << " and N=" << N << " must be multiples of 16 for the 16x16 tile design\n";
        return 4;
    }
    std::size_t a_count = static_cast<std::size_t>(M) * static_cast<std::size_t>(K);
    std::size_t b_count = static_cast<std::size_t>(K) * static_cast<std::size_t>(N);
    std::size_t c_count = static_cast<std::size_t>(M) * static_cast<std::size_t>(N);

    std::vector<uint16_t> a_bits(a_count);
    std::vector<uint16_t> b_bits(b_count);
    for (std::size_t i = 0; i < a_count; ++i) {
        if (!(std::cin >> a_bits[i])) {
            std::cerr << "failed to read A element " << i << "\n";
            return 5;
        }
    }
    for (std::size_t i = 0; i < b_count; ++i) {
        if (!(std::cin >> b_bits[i])) {
            std::cerr << "failed to read B element " << i << "\n";
            return 6;
        }
    }

    int device = 0;
    check(hipSetDevice(device), "hipSetDevice");
    hipDeviceProp_t props;
    check(hipGetDeviceProperties(&props, device), "hipGetDeviceProperties");

    __half* d_A = nullptr;
    __half* d_B = nullptr;
    float* d_C = nullptr;
    std::size_t a_bytes = a_count * sizeof(__half);
    std::size_t b_bytes = b_count * sizeof(__half);
    std::size_t c_bytes = c_count * sizeof(float);
    check(hipMalloc(&d_A, a_bytes), "hipMalloc(A)");
    check(hipMalloc(&d_B, b_bytes), "hipMalloc(B)");
    check(hipMalloc(&d_C, c_bytes), "hipMalloc(C)");

    check(hipMemcpy(d_A, a_bits.data(), a_bytes, hipMemcpyHostToDevice), "hipMemcpy(A)");
    check(hipMemcpy(d_B, b_bits.data(), b_bytes, hipMemcpyHostToDevice), "hipMemcpy(B)");

    dim3 block(16, 16);
    dim3 grid(M / 16, N / 16);

    hipEvent_t start;
    hipEvent_t stop;
    check(hipEventCreate(&start), "hipEventCreate(start)");
    check(hipEventCreate(&stop), "hipEventCreate(stop)");
    check(hipEventRecord(start), "hipEventRecord(start)");
    hipLaunchKernelGGL(gemm_fp16_f32_kernel, grid, block, 0, 0, d_A, d_B, d_C, M, N, K);
    check(hipGetLastError(), "hipLaunchKernelGGL");
    check(hipEventRecord(stop), "hipEventRecord(stop)");
    check(hipEventSynchronize(stop), "hipEventSynchronize");
    float kernel_time_ms = 0.0f;
    check(hipEventElapsedTime(&kernel_time_ms, start, stop), "hipEventElapsedTime");
    check(hipEventDestroy(start), "hipEventDestroy(start)");
    check(hipEventDestroy(stop), "hipEventDestroy(stop)");

    std::vector<float> c_out(c_count);
    check(hipMemcpy(c_out.data(), d_C, c_bytes, hipMemcpyDeviceToHost), "hipMemcpy(C)");

    check(hipFree(d_A), "hipFree(A)");
    check(hipFree(d_B), "hipFree(B)");
    check(hipFree(d_C), "hipFree(C)");

    std::cout << "DEVICE_NAME=" << props.name << "\n";
    std::cout << "GFX=" << props.gcnArchName << "\n";
    std::cout << "M=" << M << "\n";
    std::cout << "N=" << N << "\n";
    std::cout << "K=" << K << "\n";
    std::cout << "GRID_X=" << grid.x << "\n";
    std::cout << "GRID_Y=" << grid.y << "\n";
    std::cout << "BLOCK_X=" << block.x << "\n";
    std::cout << "BLOCK_Y=" << block.y << "\n";
    std::cout << "KERNEL_TIME_MS=" << kernel_time_ms << "\n";
    std::cout << "RESULTS=";
    for (std::size_t i = 0; i < c_out.size(); ++i) {
        if (i != 0) {
            std::cout << " ";
        }
        std::cout << c_out[i];
    }
    std::cout << "\n";
    return 0;
}
"#;

#[derive(Debug, Clone, PartialEq)]
pub struct RocmHipGemmF16Report {
    pub m: usize,
    pub n: usize,
    pub k: usize,
    pub outputs: Vec<f32>,
    pub cpu_oracle_outputs: Vec<f32>,
    pub max_abs_error: f32,
    pub within_tolerance: bool,
    pub kernel_time_ms: f32,
    pub kernel_source_fingerprint: String,
    pub compiler_fingerprint: String,
    pub build_command: String,
    pub executable_path: String,
    pub device_evidence: RocmHipCapabilityReport,
    pub evidence: Vec<String>,
    pub non_claims: Vec<String>,
}

impl RocmHipGemmF16Report {
    pub fn to_markdown(&self) -> String {
        let mut lines = vec![
            "# ROCm/HIP fp16 GEMM Pilot".to_string(),
            String::new(),
            format!("backend: {}", ROCM_HIP_GEMM_F16_BACKEND),
            format!("m: {}", self.m),
            format!("n: {}", self.n),
            format!("k: {}", self.k),
            format!("max_abs_error: {}", self.max_abs_error),
            format!("within_tolerance: {}", self.within_tolerance),
            format!("kernel_time_ms: {}", self.kernel_time_ms),
            format!(
                "kernel_source_fingerprint: {}",
                self.kernel_source_fingerprint
            ),
            format!("compiler_fingerprint: {}", self.compiler_fingerprint),
            String::new(),
            "## Evidence".to_string(),
        ];
        for item in &self.evidence {
            lines.push(format!("- {item}"));
        }
        lines.push(String::new());
        lines.push("## Non-Claims".to_string());
        for item in &self.non_claims {
            lines.push(format!("- {item}"));
        }
        lines.join("\n")
    }
}

pub fn run_rocm_hip_gemm_f16(
    a: &[u16],
    b: &[u16],
    m: usize,
    n: usize,
    k: usize,
) -> Result<RocmHipGemmF16Report> {
    if a.len() != m * k {
        return Err(Error::backend(format!(
            "fp16 GEMM A length {} does not match m*k={}",
            a.len(),
            m * k
        )));
    }
    if b.len() != k * n {
        return Err(Error::backend(format!(
            "fp16 GEMM B length {} does not match k*n={}",
            b.len(),
            k * n
        )));
    }
    if m == 0 || n == 0 || k == 0 {
        return Err(Error::backend("fp16 GEMM dimensions must all be positive"));
    }
    if m % 16 != 0 || n % 16 != 0 {
        return Err(Error::backend(format!(
            "fp16 GEMM m={} and n={} must both be multiples of 16 for the 16x16 tile design",
            m, n
        )));
    }

    let device_evidence = detect_local_rocm_hip();
    if !device_evidence.available {
        return Err(Error::backend(
            "ROCm/HIP is unavailable; fp16 GEMM pilot remains inadmissible",
        ));
    }

    let source_fingerprint = hip_gemm_f16_kernel_source_fingerprint();
    let cache_dir = PathBuf::from("target/rocm-hip-cache");
    fs::create_dir_all(&cache_dir)
        .map_err(|err| Error::backend(format!("failed to create HIP cache directory: {err}")))?;
    let source_path = cache_dir.join(format!("{source_fingerprint}.cpp"));
    let executable_path = cache_dir.join(format!("{source_fingerprint}-gemm-fp16-f32"));
    fs::write(&source_path, HIP_GEMM_F16_KERNEL)
        .map_err(|err| Error::backend(format!("failed to write HIP kernel source: {err}")))?;

    let hipcc = "/opt/rocm/bin/hipcc";
    // Note: the spec asked for `-target gfx1101`, but on this hipcc that
    // argument is forwarded to clang as a CPU triple (which fails with
    // `unknown target triple 'gfx1101'`). hipcc auto-derives
    // `--offload-arch=gfx1101 --offload-arch=gfx1103` from it, so passing
    // `--offload-arch=gfx1101` directly is the equivalent + correct form.
    let compiler_fingerprint = hipcc_compiler_fingerprint(hipcc)?;
    let build_command =
        hipcc_compile_executable(hipcc, &source_path, &executable_path, Some("gfx1101"))?;

    // Build stdin payload: header line + A bits + B bits, space-separated.
    let mut payload = String::with_capacity((a.len() + b.len()) * 8);
    payload.push_str(&format!("{m} {n} {k}\n"));
    for (i, v) in a.iter().enumerate() {
        if i != 0 {
            payload.push(' ');
        }
        payload.push_str(&v.to_string());
    }
    payload.push('\n');
    for (i, v) in b.iter().enumerate() {
        if i != 0 {
            payload.push(' ');
        }
        payload.push_str(&v.to_string());
    }
    payload.push('\n');

    hipcc_recheck_artifact(hipcc, &source_path, &executable_path, Some("gfx1101"))?;

    // Send the payload through the persistent kernel server pool
    // (one long-lived child per kernel_type). This is the same text
    // payload the legacy one-shot text mode would have read from
    // stdin; the server wraps it in a length-prefixed binary
    // protocol so the child can demux multiple requests without a
    // per-call process spawn.
    let stdout = kernel_server::run_persistent(GEMM_F16_KERNEL_TYPE, &executable_path, &payload)?;
    let outputs = parse_gemm_results(&stdout)?;
    let kernel_time_ms = parse_gemm_f32_line(&stdout, "KERNEL_TIME_MS=")
        .ok_or_else(|| Error::backend("HIP fp16 GEMM did not print KERNEL_TIME_MS marker"))?;
    let cpu_oracle_outputs = cpu_gemm_f16(a, b, m, n, k);

    let mut max_abs_error = 0.0f32;
    for (g, c) in outputs.iter().zip(cpu_oracle_outputs.iter()) {
        let err = (g - c).abs();
        if err > max_abs_error {
            max_abs_error = err;
        }
    }
    let within_tolerance = max_abs_error < 1e-2;

    Ok(RocmHipGemmF16Report {
        m,
        n,
        k,
        outputs,
        cpu_oracle_outputs,
        max_abs_error,
        within_tolerance,
        kernel_time_ms,
        kernel_source_fingerprint: source_fingerprint,
        compiler_fingerprint,
        build_command,
        executable_path: executable_path.display().to_string(),
        device_evidence,
        evidence: vec![
            "compiled HIP kernel with /opt/rocm/bin/hipcc -O2 --offload-arch=gfx1101".to_string(),
            "shipped A and B to the kernel via stdin (Stdio::piped)".to_string(),
            "launched gemm_fp16_f32_kernel with grid=(M/16,N/16) block=(16,16)".to_string(),
            "captured kernel time with hipEventRecord/hipEventSynchronize".to_string(),
            "compared every output element against the CPU fp16 oracle within 1e-2".to_string(),
        ],
        non_claims: vec![
            "not production speedup evidence".to_string(),
            "not optimized GEMM (no shared-memory tiling, no vectorized loads)".to_string(),
            "not general fp16 tensor contraction (no batched/strided variants)".to_string(),
            "not machine-code verification".to_string(),
        ],
    })
}

pub fn hip_gemm_f16_kernel_source_fingerprint() -> String {
    fingerprint("hip-gemm-f16-source", HIP_GEMM_F16_KERNEL)
}

/// CPU oracle: convert each fp16 bit pattern to f32, compute the matmul in
/// plain f32 with the same K-loop summation order as the kernel, then round
/// the result through fp16 and back to f32 to mirror the precision that
/// fp16 inputs admit.
pub fn cpu_gemm_f16(a: &[u16], b: &[u16], m: usize, n: usize, k: usize) -> Vec<f32> {
    let a_f32: Vec<f32> = a.iter().copied().map(f16_to_f32).collect();
    let b_f32: Vec<f32> = b.iter().copied().map(f16_to_f32).collect();
    let mut c = vec![0.0f32; m * n];
    for i in 0..m {
        for j in 0..n {
            let mut acc = 0.0f32;
            for kk in 0..k {
                acc += a_f32[i * k + kk] * b_f32[kk * n + j];
            }
            c[i * n + j] = f16_to_f32(f32_to_f16(acc));
        }
    }
    c
}

// Re-export the canonical IEEE 754 binary16 <-> binary32 conversion
// helpers from `tokitai_operator::backend::f16_convert`. Keeping
// the helpers in one place avoids drift: the subnormal encoding
// path is subtle and had a long-standing bug in earlier inline
// copies (zeroed any value in `(0, 2^-14)` — see the regression
// tests in `f16_convert::tests`). All call sites in this file
// (CPU oracle, the `*_f32` helpers, and the C++ wire side) reach
// the helpers through this re-export.
pub use crate::backend::f16_convert::{f16_to_f32, f32_to_f16};

fn parse_gemm_results(stdout: &str) -> Result<Vec<f32>> {
    let line = stdout
        .lines()
        .find_map(|line| line.strip_prefix("RESULTS="))
        .ok_or_else(|| Error::backend("HIP fp16 GEMM did not print RESULTS marker"))?;
    if line.trim().is_empty() {
        return Ok(Vec::new());
    }
    line.split_whitespace()
        .map(|value| {
            value.trim().parse::<f32>().map_err(|err| {
                Error::backend(format!(
                    "invalid HIP fp16 GEMM output value {value:?}: {err}"
                ))
            })
        })
        .collect()
}

fn parse_gemm_f32_line(stdout: &str, prefix: &str) -> Option<f32> {
    stdout
        .lines()
        .find_map(|line| line.strip_prefix(prefix))
        .and_then(|value| value.trim().parse::<f32>().ok())
}

fn fingerprint(label: &str, value: &str) -> String {
    let mut hasher = DefaultHasher::new();
    label.hash(&mut hasher);
    value.hash(&mut hasher);
    format!("{label}-{:016x}", hasher.finish())
}