nx_sparse_tensor_test.nx source

1// nx_sparse_tensor_test.nx -- algo-led correctness: sparse matmul 2// must match dense matmul bit-exact across the same inputs. Oracle- 3// verified per the cardinal: every algorithmic variant ships through 4// nx_numeric_oracle vs the dense reference. 5 6import "nx_syscalls.nx" 7import "nx_tier.nx" 8import "nx_tensor.nx" 9import "nx_blas_i64.nx" 10import "nx_sparse_tensor.nx" 11import "nx_numeric_oracle.nx" 12 13func main() -> nx_int { 14 let err: *i64 = (sys_mmap(8)) as *i64 15 16 // ===== Build a sparse 4x6 matrix (mostly zero) =============== 17 // 18 // A = [ 0 3 0 0 0 0 ] 19 // [ 0 0 0 5 0 0 ] 20 // [ 0 0 0 0 0 0 ] <- entirely zero row 21 // [ 7 0 0 0 9 0 ] 22 // 23 // NNZ = 4. Sparsity = 4 / 24 = 16.7% non-zero, ~83% sparse. 24 let s_a: *i64 = (sys_mmap(16)) as *i64 25 s_a[0] = 4; s_a[1] = 6 26 let a_dense: *NxTensor = nx_t_alloc(NX_DT_I64, s_a, 2, err) 27 nx_t_fill_zero(a_dense) 28 let idx: *i64 = (sys_mmap(64)) as *i64 29 idx[0] = 0; idx[1] = 1; nx_t_set_i64(a_dense, idx, 3) 30 idx[0] = 1; idx[1] = 3; nx_t_set_i64(a_dense, idx, 5) 31 idx[0] = 3; idx[1] = 0; nx_t_set_i64(a_dense, idx, 7) 32 idx[0] = 3; idx[1] = 4; nx_t_set_i64(a_dense, idx, 9) 33 34 // ===== Convert to CSR ======================================== 35 let a_csr: *NxSparseTensor = nx_sp_alloc(4, 6, 16, NX_DT_I64) 36 let cv: nx_int = nx_sp_dense_to_csr(a_dense, a_csr) 37 if cv != NX_SP_OK { return 1 } 38 if a_csr.nnz != 4 { return 2 } 39 40 // row_ptr should be [0, 1, 2, 2, 4] 41 if a_csr.row_ptr[0] != 0 { return 3 } 42 if a_csr.row_ptr[1] != 1 { return 4 } 43 if a_csr.row_ptr[2] != 2 { return 5 } 44 if a_csr.row_ptr[3] != 2 { return 6 } // empty row 45 if a_csr.row_ptr[4] != 4 { return 7 } 46 47 // col_idx in order: 1, 3, 0, 4 48 if a_csr.col_idx[0] != 1 { return 10 } 49 if a_csr.col_idx[1] != 3 { return 11 } 50 if a_csr.col_idx[2] != 0 { return 12 } 51 if a_csr.col_idx[3] != 4 { return 13 } 52 // values: 3, 5, 7, 9 53 if a_csr.values[0] != 3 { return 14 } 54 if a_csr.values[1] != 5 { return 15 } 55 if a_csr.values[2] != 7 { return 16 } 56 if a_csr.values[3] != 9 { return 17 } 57 58 // ===== Round-trip CSR -> dense -> oracle ===================== 59 let s_back: *i64 = (sys_mmap(16)) as *i64 60 s_back[0] = 4; s_back[1] = 6 61 let a_back: *NxTensor = nx_t_alloc(NX_DT_I64, s_back, 2, err) 62 nx_sp_csr_to_dense(a_csr, a_back) 63 let witness: *i64 = (sys_mmap(NX_NO_WITNESS_FIELDS * 8)) as *i64 64 if nx_no_check_bit_exact_i64(a_dense, a_back, witness) 65 != NX_NO_VERDICT_EQUAL { return 20 } 66 67 // ===== Build a dense B (6x3) ================================= 68 let s_b: *i64 = (sys_mmap(16)) as *i64 69 s_b[0] = 6; s_b[1] = 3 70 let b_dense: *NxTensor = nx_t_alloc(NX_DT_I64, s_b, 2, err) 71 nx_t_fill_zero(b_dense) 72 // Fill with simple incrementing values for verification 73 let pb: *i64 = b_dense.storage as *i64 74 var i: nx_int = 0 75 while i < b_dense.numel { 76 pb[i] = i + 1 // 1, 2, 3, ..., 18 77 i = i + 1 78 } 79 80 // ===== Compute via DENSE matmul (reference) ================== 81 let s_c: *i64 = (sys_mmap(16)) as *i64 82 s_c[0] = 4; s_c[1] = 3 83 let c_dense_ref: *NxTensor = nx_t_alloc(NX_DT_I64, s_c, 2, err) 84 nx_t_fill_zero(c_dense_ref) 85 let dv: nx_int = nx_blas_matmul(a_dense, b_dense, c_dense_ref) 86 if dv != NX_BLAS_OK { return 30 } 87 88 // ===== Compute via SPARSE matmul (variant) ================== 89 let c_sparse_out: *NxTensor = nx_t_alloc(NX_DT_I64, s_c, 2, err) 90 let sv: nx_int = nx_sp_csr_matmul_dense(a_csr, b_dense, c_sparse_out) 91 if sv != NX_SP_OK { return 31 } 92 93 // ===== Oracle: sparse output MUST equal dense reference ===== 94 let v_eq: nx_int = nx_no_check_bit_exact_i64(c_sparse_out, c_dense_ref, 95 witness) 96 if v_eq != NX_NO_VERDICT_EQUAL { 97 // Return 40 + the verdict byte so smoke pinpoints 98 return 40 + v_eq 99 } 100 101 // ===== Spot-check expected values ============================ 102 // 103 // C[0, *] = A[0, *] * B = [0 3 0 0 0 0] * B 104 // = 3 * B[1, *] = 3 * [4, 5, 6] = [12, 15, 18] 105 idx[0] = 0; idx[1] = 0 106 if nx_t_get_i64(c_sparse_out, idx) != 12 { return 50 } 107 idx[0] = 0; idx[1] = 1 108 if nx_t_get_i64(c_sparse_out, idx) != 15 { return 51 } 109 idx[0] = 0; idx[1] = 2 110 if nx_t_get_i64(c_sparse_out, idx) != 18 { return 52 } 111 // 112 // C[2, *] = entirely zero (empty row) 113 idx[0] = 2; idx[1] = 0 114 if nx_t_get_i64(c_sparse_out, idx) != 0 { return 53 } 115 // 116 // C[3, *] = 7*B[0,*] + 9*B[4,*] = 7*[1,2,3] + 9*[13,14,15] 117 // = [7+117, 14+126, 21+135] = [124, 140, 156] 118 idx[0] = 3; idx[1] = 0 119 if nx_t_get_i64(c_sparse_out, idx) != 124 { return 54 } 120 idx[0] = 3; idx[1] = 1 121 if nx_t_get_i64(c_sparse_out, idx) != 140 { return 55 } 122 idx[0] = 3; idx[1] = 2 123 if nx_t_get_i64(c_sparse_out, idx) != 156 { return 56 } 124 125 // ===== Sparsity + speedup measurements ======================= 126 // 127 // NNZ=4, total=24 -> ratio_q10 = 4 * 1024 / 24 = 170 (16.6%) 128 let ratio: nx_int = nx_sp_sparsity_ratio_q10(a_csr) 129 if ratio != 170 { return 60 } 130 131 // Speedup = (4 * 6 * 1024) / 4 = 6144 Q10 = 6.0x 132 let su: nx_int = nx_sp_estimated_speedup_q10(a_csr) 133 if su != 6144 { return 61 } 134 135 // ===== High-sparsity case (attention-shape) ================= 136 // 137 // 16x16 with only 1 non-zero -> ratio = 4 Q10 (~0.4% dense) 138 let s_big: *i64 = (sys_mmap(16)) as *i64 139 s_big[0] = 16; s_big[1] = 16 140 let big_dense: *NxTensor = nx_t_alloc(NX_DT_I64, s_big, 2, err) 141 nx_t_fill_zero(big_dense) 142 idx[0] = 7; idx[1] = 11; nx_t_set_i64(big_dense, idx, 42) 143 let big_csr: *NxSparseTensor = nx_sp_alloc(16, 16, 8, NX_DT_I64) 144 nx_sp_dense_to_csr(big_dense, big_csr) 145 if big_csr.nnz != 1 { return 70 } 146 // 1 / 256 in Q10 = 4 147 if nx_sp_sparsity_ratio_q10(big_csr) != 4 { return 71 } 148 // Speedup = 16 * 16 * 1024 / 1 = 262144 Q10 = 256x 149 if nx_sp_estimated_speedup_q10(big_csr) != 262144 { return 72 } 150 151 // ===== Verdict enum coverage ================================= 152 if nx_sp_fmt_is_valid(NX_SP_FMT_CSR) != 1 { return 80 } 153 if nx_sp_fmt_is_valid(NX_SP_FMT_N_KINDS) != 0 { return 81 } 154 if nx_sp_verdict_is_valid(NX_SP_OK) != 1 { return 82 } 155 if nx_sp_verdict_is_valid(NX_SP_N_VERDICTS) != 0 { return 83 } 156 157 return 0 158}