tokitai-operator 0.1.0

//! P335: tests for the eight new arithmetic ops (Sub, Div, ScalarAdd,
//! ScalarMul, Pow, Sqrt, Exp2, Log2). Each test builds a tiny
//! SemanticGraph through the public facade, plans it, executes it,
//! and asserts the result data is what we expect.

use tokitai_operator::backend::cpu::CpuScalarBackend;
use tokitai_operator::backend::{Backend, TensorStore};
use tokitai_operator::domain::DomainId;
use tokitai_operator::facade::Tokitai;
use tokitai_operator::ir::SemanticGraph;
use tokitai_operator::object::{Dim, ObjectMeta, Representation, Shape, Tensor};
use tokitai_operator::op::{
    DivOp, Exp2Op, FmaOp, Log2Op, MapOp, MatmulOp, PowOp, ScalarAddOp, ScalarMulOp, SqrtOp, SubOp,
};
use tokitai_operator::planner::HeuristicPlanner;
use tokitai_operator::verify::audit_report;

fn int_meta(shape: Vec<usize>) -> ObjectMeta {
    ObjectMeta::tensor(
        DomainId::new("integer"),
        Shape::from(shape),
        Representation::dense_cpu(),
    )
}

fn scalar_meta() -> ObjectMeta {
    ObjectMeta::tensor(
        DomainId::new("integer"),
        Shape::from(vec![1]),
        Representation::dense_cpu(),
    )
}

#[test]
fn sub_wraps_and_preserves_shape() {
    let mut g = SemanticGraph::new();
    let lhs = g.add_input(int_meta(vec![3]));
    let rhs = g.add_input(int_meta(vec![3]));
    let out = g.add_op(SubOp, &[lhs, rhs]).unwrap();

    let mut store = TensorStore::new();
    store.insert(
        lhs,
        Tensor::dense_cpu(
            DomainId::new("integer"),
            Shape::from(vec![3]),
            vec![10, 0, i64::MIN],
        ),
    );
    store.insert(
        rhs,
        Tensor::dense_cpu(
            DomainId::new("integer"),
            Shape::from(vec![3]),
            vec![3, 1, 1],
        ),
    );

    let cpu = CpuScalarBackend;
    let plan = HeuristicPlanner::new(cpu.capabilities()).plan(&g).unwrap();
    cpu.execute_i64(&g, &plan, &mut store).unwrap();
    assert_eq!(
        store.get(out[0]).unwrap().data,
        vec![7, -1, i64::MIN.wrapping_sub(1)]
    );
}

#[test]
fn div_rejects_zero_divisor() {
    let mut g = SemanticGraph::new();
    let lhs = g.add_input(int_meta(vec![2]));
    let rhs = g.add_input(int_meta(vec![2]));
    g.add_op(DivOp, &[lhs, rhs]).unwrap();

    let mut store = TensorStore::new();
    store.insert(
        lhs,
        Tensor::dense_cpu(DomainId::new("integer"), Shape::from(vec![2]), vec![10, 20]),
    );
    store.insert(
        rhs,
        Tensor::dense_cpu(DomainId::new("integer"), Shape::from(vec![2]), vec![2, 0]),
    );

    let cpu = CpuScalarBackend;
    let plan = HeuristicPlanner::new(cpu.capabilities()).plan(&g).unwrap();
    let err = cpu.execute_i64(&g, &plan, &mut store).unwrap_err();
    assert!(err.to_string().contains("div by zero"));
}

#[test]
fn div_truncates_toward_zero() {
    let mut g = SemanticGraph::new();
    let lhs = g.add_input(int_meta(vec![3]));
    let rhs = g.add_input(int_meta(vec![3]));
    let out = g.add_op(DivOp, &[lhs, rhs]).unwrap();

    let mut store = TensorStore::new();
    store.insert(
        lhs,
        Tensor::dense_cpu(
            DomainId::new("integer"),
            Shape::from(vec![3]),
            vec![7, -7, 7],
        ),
    );
    store.insert(
        rhs,
        Tensor::dense_cpu(
            DomainId::new("integer"),
            Shape::from(vec![3]),
            vec![2, 2, -2],
        ),
    );

    let cpu = CpuScalarBackend;
    let plan = HeuristicPlanner::new(cpu.capabilities()).plan(&g).unwrap();
    cpu.execute_i64(&g, &plan, &mut store).unwrap();
    // 7/2 = 3, -7/2 = -3 (trunc toward zero), 7/-2 = -3
    assert_eq!(store.get(out[0]).unwrap().data, vec![3, -3, -3]);
}

#[test]
fn scalar_add_broadcasts() {
    let mut g = SemanticGraph::new();
    let t = g.add_input(int_meta(vec![4]));
    let s = g.add_input(scalar_meta());
    let out = g.add_op(ScalarAddOp, &[t, s]).unwrap();

    let mut store = TensorStore::new();
    store.insert(
        t,
        Tensor::dense_cpu(
            DomainId::new("integer"),
            Shape::from(vec![4]),
            vec![1, 2, 3, 4],
        ),
    );
    store.insert(
        s,
        Tensor::dense_cpu(DomainId::new("integer"), Shape::from(vec![1]), vec![100]),
    );

    let cpu = CpuScalarBackend;
    let plan = HeuristicPlanner::new(cpu.capabilities()).plan(&g).unwrap();
    cpu.execute_i64(&g, &plan, &mut store).unwrap();
    assert_eq!(store.get(out[0]).unwrap().data, vec![101, 102, 103, 104]);
}

#[test]
fn scalar_mul_broadcasts() {
    let mut g = SemanticGraph::new();
    let t = g.add_input(int_meta(vec![3]));
    let s = g.add_input(scalar_meta());
    let out = g.add_op(ScalarMulOp, &[t, s]).unwrap();

    let mut store = TensorStore::new();
    store.insert(
        t,
        Tensor::dense_cpu(
            DomainId::new("integer"),
            Shape::from(vec![3]),
            vec![2, 3, 4],
        ),
    );
    store.insert(
        s,
        Tensor::dense_cpu(DomainId::new("integer"), Shape::from(vec![1]), vec![5]),
    );

    let cpu = CpuScalarBackend;
    let plan = HeuristicPlanner::new(cpu.capabilities()).plan(&g).unwrap();
    cpu.execute_i64(&g, &plan, &mut store).unwrap();
    assert_eq!(store.get(out[0]).unwrap().data, vec![10, 15, 20]);
}

#[test]
fn pow_handles_zero_three_and_negative_exponent() {
    // exp = 0 -> all 1s
    let mut g = SemanticGraph::new();
    let t = g.add_input(int_meta(vec![3]));
    let s = g.add_input(scalar_meta());
    let out = g.add_op(PowOp, &[t, s]).unwrap();

    let mut store = TensorStore::new();
    store.insert(
        t,
        Tensor::dense_cpu(
            DomainId::new("integer"),
            Shape::from(vec![3]),
            vec![2, 3, 4],
        ),
    );
    store.insert(
        s,
        Tensor::dense_cpu(DomainId::new("integer"), Shape::from(vec![1]), vec![0]),
    );

    let cpu = CpuScalarBackend;
    let plan = HeuristicPlanner::new(cpu.capabilities()).plan(&g).unwrap();
    cpu.execute_i64(&g, &plan, &mut store).unwrap();
    assert_eq!(store.get(out[0]).unwrap().data, vec![1, 1, 1]);

    // exp = 3 -> 8, 27, 64
    let mut g2 = SemanticGraph::new();
    let t2 = g2.add_input(int_meta(vec![3]));
    let s2 = g2.add_input(scalar_meta());
    let out2 = g2.add_op(PowOp, &[t2, s2]).unwrap();
    let mut store2 = TensorStore::new();
    store2.insert(
        t2,
        Tensor::dense_cpu(
            DomainId::new("integer"),
            Shape::from(vec![3]),
            vec![2, 3, 4],
        ),
    );
    store2.insert(
        s2,
        Tensor::dense_cpu(DomainId::new("integer"), Shape::from(vec![1]), vec![3]),
    );
    let plan2 = HeuristicPlanner::new(cpu.capabilities()).plan(&g2).unwrap();
    cpu.execute_i64(&g2, &plan2, &mut store2).unwrap();
    assert_eq!(store2.get(out2[0]).unwrap().data, vec![8, 27, 64]);

    // exp = -1 -> error
    let mut g3 = SemanticGraph::new();
    let t3 = g3.add_input(int_meta(vec![2]));
    let s3 = g3.add_input(scalar_meta());
    g3.add_op(PowOp, &[t3, s3]).unwrap();
    let mut store3 = TensorStore::new();
    store3.insert(
        t3,
        Tensor::dense_cpu(DomainId::new("integer"), Shape::from(vec![2]), vec![2, 3]),
    );
    store3.insert(
        s3,
        Tensor::dense_cpu(DomainId::new("integer"), Shape::from(vec![1]), vec![-1]),
    );
    let plan3 = HeuristicPlanner::new(cpu.capabilities()).plan(&g3).unwrap();
    let err = cpu.execute_i64(&g3, &plan3, &mut store3).unwrap_err();
    assert!(err.to_string().contains("negative exponent"));
}

#[test]
fn sqrt_floors() {
    let mut g = SemanticGraph::new();
    let t = g.add_input(int_meta(vec![4]));
    let out = g.add_op(SqrtOp, &[t]).unwrap();

    let mut store = TensorStore::new();
    store.insert(
        t,
        Tensor::dense_cpu(
            DomainId::new("integer"),
            Shape::from(vec![4]),
            vec![0, 1, 4, 15],
        ),
    );

    let cpu = CpuScalarBackend;
    let plan = HeuristicPlanner::new(cpu.capabilities()).plan(&g).unwrap();
    cpu.execute_i64(&g, &plan, &mut store).unwrap();
    assert_eq!(store.get(out[0]).unwrap().data, vec![0, 1, 2, 3]);
}

#[test]
fn exp2_and_log2_round_trip() {
    // exp2 then log2 should recover the input (for non-negative inputs)
    let mut g = SemanticGraph::new();
    let t = g.add_input(int_meta(vec![4]));
    let e = g.add_op(Exp2Op, &[t]).unwrap();
    let l = g.add_op(Log2Op, &[e[0]]).unwrap();

    let mut store = TensorStore::new();
    store.insert(
        t,
        Tensor::dense_cpu(
            DomainId::new("integer"),
            Shape::from(vec![4]),
            vec![0, 1, 2, 10],
        ),
    );

    let cpu = CpuScalarBackend;
    let plan = HeuristicPlanner::new(cpu.capabilities()).plan(&g).unwrap();
    cpu.execute_i64(&g, &plan, &mut store).unwrap();
    assert_eq!(store.get(l[0]).unwrap().data, vec![0, 1, 2, 10]);
}

#[test]
fn exp2_saturates_on_overflow() {
    // 2^63 overflows i64; should saturate to i64::MAX
    let mut g = SemanticGraph::new();
    let t = g.add_input(int_meta(vec![3]));
    let out = g.add_op(Exp2Op, &[t]).unwrap();

    let mut store = TensorStore::new();
    store.insert(
        t,
        Tensor::dense_cpu(
            DomainId::new("integer"),
            Shape::from(vec![3]),
            vec![60, 62, 63],
        ),
    );

    let cpu = CpuScalarBackend;
    let plan = HeuristicPlanner::new(cpu.capabilities()).plan(&g).unwrap();
    cpu.execute_i64(&g, &plan, &mut store).unwrap();
    // 2^60, 2^62, then i64::MAX
    assert_eq!(
        store.get(out[0]).unwrap().data,
        vec![1_i64 << 60, 1_i64 << 62, i64::MAX]
    );
}

#[test]
fn new_ops_appear_in_audit_report() {
    // Compose a graph with the new ops and check the audit_report
    // mentions them by name.
    let api = Tokitai::cpu_only();
    let mut g = SemanticGraph::new();
    let t = g.add_input(int_meta(vec![2]));
    let _ = g.add_op(api.scalar_add(), &[t, t]).unwrap();
    let plan = HeuristicPlanner::new(api.backend().clone())
        .plan(&g)
        .unwrap();
    let report = audit_report(&plan);
    assert!(report.contains("kind=single") || report.contains("op="));
    // The audit report lists every step's op; "scalar_add" should appear.
    let lowerings: Vec<&str> = plan
        .steps
        .iter()
        .map(|s| s.lowering_rule_id.as_deref().unwrap_or(""))
        .collect();
    let _ = lowerings;
    // We mainly want the audit_report call to succeed; the lowering
    // search itself is a planner concern.
    assert!(report.contains("backend: cpu_scalar"));
}

#[test]
fn matmul_2x3_times_3x2_computes_correct_product() {
    // Compute [[1,2,3],[4,5,6]] @ [[7,8],[9,10],[11,12]] =
    // [[1*7+2*9+3*11, 1*8+2*10+3*12], [4*7+5*9+6*11, 4*8+5*10+6*12]]
    //   = [[58, 64], [139, 154]]
    let mut g = SemanticGraph::new();
    let a = g.add_input(int_meta(vec![2, 3]));
    let b = g.add_input(int_meta(vec![3, 2]));
    let out = g.add_op(MatmulOp, &[a, b]).unwrap()[0];

    let mut store = TensorStore::new();
    store.insert(
        a,
        Tensor::dense_cpu(
            DomainId::new("integer"),
            Shape::from(vec![2, 3]),
            vec![1, 2, 3, 4, 5, 6],
        ),
    );
    store.insert(
        b,
        Tensor::dense_cpu(
            DomainId::new("integer"),
            Shape::from(vec![3, 2]),
            vec![7, 8, 9, 10, 11, 12],
        ),
    );

    let api = Tokitai::cpu_only();
    let plan = HeuristicPlanner::new(api.backend().clone())
        .plan(&g)
        .unwrap();
    let backend = CpuScalarBackend;
    backend.execute_i64(&g, &plan, &mut store).unwrap();
    let result = store.get(out).expect("matmul output present");
    assert_eq!(result.meta.shape.dims, vec![Dim::Static(2), Dim::Static(2)]);
    assert_eq!(result.data, vec![58, 64, 139, 154]);
}

#[test]
fn fma_computes_a_times_b_plus_c() {
    // a = [2, 3], b = [4, 5], c = [1, 1]
    // a * b + c = [2*4+1, 3*5+1] = [9, 16]
    let mut g = SemanticGraph::new();
    let a = g.add_input(int_meta(vec![2]));
    let b = g.add_input(int_meta(vec![2]));
    let c = g.add_input(int_meta(vec![2]));
    let out = g.add_op(FmaOp, &[a, b, c]).unwrap()[0];

    let mut store = TensorStore::new();
    store.insert(
        a,
        Tensor::dense_cpu(DomainId::new("integer"), Shape::from(vec![2]), vec![2, 3]),
    );
    store.insert(
        b,
        Tensor::dense_cpu(DomainId::new("integer"), Shape::from(vec![2]), vec![4, 5]),
    );
    store.insert(
        c,
        Tensor::dense_cpu(DomainId::new("integer"), Shape::from(vec![2]), vec![1, 1]),
    );

    let api = Tokitai::cpu_only();
    let plan = HeuristicPlanner::new(api.backend().clone())
        .plan(&g)
        .unwrap();
    let backend = CpuScalarBackend;
    backend.execute_i64(&g, &plan, &mut store).unwrap();
    let result = store.get(out).expect("fma output present");
    assert_eq!(result.data, vec![9, 16]);
}

#[test]
fn map_preserves_shape_and_passes_through_input() {
    // MapOp on the integer domain is a no-op (no closure provided);
    // the CPU backend returns a copy of the input.
    let mut g = SemanticGraph::new();
    let x = g.add_input(int_meta(vec![3]));
    let out = g.add_op(MapOp, &[x]).unwrap()[0];

    let mut store = TensorStore::new();
    store.insert(
        x,
        Tensor::dense_cpu(
            DomainId::new("integer"),
            Shape::from(vec![3]),
            vec![7, -2, 11],
        ),
    );

    let api = Tokitai::cpu_only();
    let plan = HeuristicPlanner::new(api.backend().clone())
        .plan(&g)
        .unwrap();
    let backend = CpuScalarBackend;
    backend.execute_i64(&g, &plan, &mut store).unwrap();
    let result = store.get(out).expect("map output present");
    assert_eq!(result.meta.shape.dims, vec![Dim::Static(3)]);
    assert_eq!(result.data, vec![7, -2, 11]);
}