nx_simd_i8x32_test.nx source

1// nx_simd_i8x32_test.nx -- exercise the i8x32 SIMD intrinsics. 2// 3// 32 lanes of 8-bit ints in 256-bit vectors. Widest practical SIMD. 4// Useful for crypto (AES rounds), string ops, INT8 quantized ML 5// inference. No VMUL because AVX2 lacks vpmullb -- callers that 6// need byte-mul split to i16x16 first. 7 8import "nx_kernel_v2.nx" 9import "nx_log.nx" 10 11// Pack 8 i8 values (low byte each) into one i64. 12func pack8(b0: i64, b1: i64, b2: i64, b3: i64, 13 b4: i64, b5: i64, b6: i64, b7: i64) -> i64 { 14 return (b0 & 0xFF) 15 | ((b1 & 0xFF) << 8) 16 | ((b2 & 0xFF) << 16) 17 | ((b3 & 0xFF) << 24) 18 | ((b4 & 0xFF) << 32) 19 | ((b5 & 0xFF) << 40) 20 | ((b6 & 0xFF) << 48) 21 | ((b7 & 0xFF) << 56) 22} 23 24func t1_load_store() -> i64 { 25 let a_raw: *u8 = sys_mmap(64) 26 let r_raw: *u8 = sys_mmap(64) 27 let a: *i64 = a_raw as *i64 28 let r: *i64 = r_raw as *i64 29 a[0] = pack8(1, 2, 3, 4, 5, 6, 7, 8) 30 a[1] = pack8(9, 10, 11, 12, 13, 14, 15, 16) 31 a[2] = pack8(17, 18, 19, 20, 21, 22, 23, 24) 32 a[3] = pack8(25, 26, 27, 28, 29, 30, 31, 32) 33 let v: i64 = __simd_vload_i8_x32(a as *i64) 34 __simd_vstore_i8_x32(v, r as *i64) 35 if r[0] != a[0] { return 1 } 36 if r[1] != a[1] { return 2 } 37 if r[2] != a[2] { return 3 } 38 if r[3] != a[3] { return 4 } 39 return 0 40} 41 42func t2_vadd() -> i64 { 43 let a_raw: *u8 = sys_mmap(64) 44 let b_raw: *u8 = sys_mmap(64) 45 let r_raw: *u8 = sys_mmap(64) 46 let a: *i64 = a_raw as *i64 47 let b: *i64 = b_raw as *i64 48 let r: *i64 = r_raw as *i64 49 a[0] = pack8(10, 20, 30, 40, 50, 60, 70, 80) 50 a[1] = pack8(0, 0, 0, 0, 0, 0, 0, 0) 51 a[2] = pack8(0, 0, 0, 0, 0, 0, 0, 0) 52 a[3] = pack8(0, 0, 0, 0, 0, 0, 0, 0) 53 b[0] = pack8(1, 2, 3, 4, 5, 6, 7, 8) 54 b[1] = pack8(0, 0, 0, 0, 0, 0, 0, 0) 55 b[2] = pack8(0, 0, 0, 0, 0, 0, 0, 0) 56 b[3] = pack8(0, 0, 0, 0, 0, 0, 0, 0) 57 let va: i64 = __simd_vload_i8_x32(a as *i64) 58 let vb: i64 = __simd_vload_i8_x32(b as *i64) 59 let vsum: i64 = __simd_vadd_i8_x32(va, vb) 60 __simd_vstore_i8_x32(vsum, r as *i64) 61 // First word: (11, 22, 33, 44, 55, 66, 77, 88) 62 if r[0] != pack8(11, 22, 33, 44, 55, 66, 77, 88) { return 10 } 63 return 0 64} 65 66func t3_vbroadcast() -> i64 { 67 let r_raw: *u8 = sys_mmap(64) 68 let r: *i64 = r_raw as *i64 69 let v: i64 = __simd_vbroadcast_i8_x32(0x42) 70 __simd_vstore_i8_x32(v, r as *i64) 71 let expected: i64 = pack8(0x42, 0x42, 0x42, 0x42, 0x42, 0x42, 0x42, 0x42) 72 if r[0] != expected { return 20 } 73 if r[1] != expected { return 21 } 74 if r[2] != expected { return 22 } 75 if r[3] != expected { return 23 } 76 return 0 77} 78 79func t4_vreduce_sum() -> i64 { 80 let a_raw: *u8 = sys_mmap(64) 81 let a: *i64 = a_raw as *i64 82 // Fill all 32 bytes with the value 5; sum = 5*32 = 160. 83 a[0] = pack8(5, 5, 5, 5, 5, 5, 5, 5) 84 a[1] = pack8(5, 5, 5, 5, 5, 5, 5, 5) 85 a[2] = pack8(5, 5, 5, 5, 5, 5, 5, 5) 86 a[3] = pack8(5, 5, 5, 5, 5, 5, 5, 5) 87 let v: i64 = __simd_vload_i8_x32(a as *i64) 88 let s: i64 = __simd_vreduce_sum_i8_x32(v) 89 if s != 160 { return 30 } 90 return 0 91} 92 93func main() -> nx_exit { 94 println("=== nx_simd_i8x32 smoke (32 lanes of 8-bit) ===" as *u8) 95 if t1_load_store() != 0 { println("FAIL T1" as *u8); return 1 } 96 println("T1 vload+vstore PASS 32-lane i8 round-trip" as *u8) 97 if t2_vadd() != 0 { println("FAIL T2" as *u8); return 2 } 98 println("T2 vadd PASS lanes (11..88) bit-exact" as *u8) 99 if t3_vbroadcast() != 0 { println("FAIL T3" as *u8); return 3 } 100 println("T3 vbroadcast PASS all 32 lanes = 0x42" as *u8) 101 if t4_vreduce_sum() != 0 { println("FAIL T4" as *u8); return 4 } 102 println("T4 vreduce_sum PASS 5*32 = 160" as *u8) 103 println("" as *u8) 104 println("All i8x32 ops bit-exact; SIMD width spectrum complete (4/8/16/32)." as *u8) 105 return 0 106}