diff --git a/spy/backend/c/c_ast.py b/spy/backend/c/c_ast.py index 1151c821b..9e28ed31d 100644 --- a/spy/backend/c/c_ast.py +++ b/spy/backend/c/c_ast.py @@ -333,3 +333,33 @@ def precedence(self) -> int: def __str__(self) -> str: return f"({self.type}){self.expr}" + + +@dataclass +class Index(Expr): + expr: Expr + index: Expr + + def precedence(self) -> int: + return 14 + + def __str__(self) -> str: + e = str(self.expr) + if self.expr.precedence() < self.precedence(): + e = f"({e})" + return f"{e}[{self.index}]" + + +@dataclass +class Paren(Expr): + """ + An explicitly parenthesized expression. ``Cast`` does not parenthesize its operand. + """ + + expr: Expr + + def precedence(self) -> int: + return 100 + + def __str__(self) -> str: + return f"({self.expr})" diff --git a/spy/backend/c/cbackend.py b/spy/backend/c/cbackend.py index a7d55f49c..2b98d72eb 100644 --- a/spy/backend/c/cbackend.py +++ b/spy/backend/c/cbackend.py @@ -14,6 +14,7 @@ from spy.highlight import highlight_src from spy.vm.cell import W_Cell from spy.vm.function import W_ASTFunc +from spy.vm.modules.simd import W_SimdType from spy.vm.modules.unsafe.ptr import W_MemLocType from spy.vm.object import W_Object, W_Type from spy.vm.primitive import W_I32 @@ -145,7 +146,7 @@ def foo() -> i32: modname = fqn.modname w_mod = self.vm.modules_w[modname] if w_mod.filepath is None and not isinstance( - w_obj, (W_MemLocType, W_StructType) + w_obj, (W_MemLocType, W_StructType, W_SimdType) ): continue diff --git a/spy/backend/c/cstructwriter.py b/spy/backend/c/cstructwriter.py index 1f2a20aca..169ba3dc7 100644 --- a/spy/backend/c/cstructwriter.py +++ b/spy/backend/c/cstructwriter.py @@ -7,6 +7,7 @@ from spy.backend.c.context import C_Type, Context from spy.fqn import FQN from spy.textbuilder import TextBuilder +from spy.vm.modules.simd import W_SimdType from spy.vm.modules.unsafe.ptr import W_PtrType, W_RefType from spy.vm.object import W_Type from spy.vm.struct import W_StructType @@ -98,6 +99,8 @@ def emit_content(self) -> None: assert fqn == w_type.fqn # sanity check if isinstance(w_type, W_StructType): self.emit_StructType(fqn, w_type) + elif isinstance(w_type, W_SimdType): + self.emit_SimdType(fqn, w_type) elif isinstance(w_type, W_PtrType): self.emit_PtrType(fqn, w_type) elif isinstance(w_type, W_RefType): @@ -155,7 +158,22 @@ def emit_StructType(self, fqn: FQN, w_st: W_StructType) -> None: tb.wl("};") tb.wl("") + def emit_SimdType(self, fqn: FQN, w_simdtype: W_SimdType) -> None: + from spy.vm.modules.unsafe.misc import sizeof + + c_simdtype = C_Type(w_simdtype.fqn.c_name) + c_basetype = self.ctx.w2c(w_simdtype.w_dtype) + nbytes = sizeof(w_simdtype.w_dtype) * w_simdtype.size + human = w_simdtype.fqn.human_name(self.ctx.vm) + # GCC/Clang vector extension: a fixed-size vector + self.tbh_fwdecl.wl( + f"typedef {c_basetype} {c_simdtype} " + f"__attribute__((vector_size({nbytes}))); /* {human} */" + ) + def emit_PtrType(self, fqn: FQN, w_ptrtype: W_PtrType) -> None: + from spy.vm.modules.simd import W_SimdType + c_ptrtype = C_Type(w_ptrtype.fqn.c_name) w_itemT = w_ptrtype.w_itemT c_itemT = self.ctx.w2c(w_itemT) @@ -183,8 +201,33 @@ def emit_PtrType(self, fqn: FQN, w_ptrtype: W_PtrType) -> None: self.tbh_fwdecl.wl() memkind = w_ptrtype.memkind + # SIMD vectors carry natural alignment == their byte size (e.g. 32 for + # SIMD[i64,4]), but the GC/raw allocator returns 16-byte-aligned memory + # on wasm32, so a by-value store/load through ptr[SIMD[...]] traps on + # misalignment for vectors > 16 B. Keep the vector type naturally + # aligned (vector registers, fast) and instead over-align the + # *allocation*: route $alloc through an over-aligning wrapper by giving + # SPY_PTR_FUNCTIONS a synthetic memkind. + if isinstance(w_itemT, W_SimdType): + macro_memkind = f"{memkind}_simd" + guard = f"SPY_{memkind.upper()}_SIMD_ALLOC_DEFINED" + # 64 = max vector byte size (size 8 x 8-byte dtype); over-aligning + # to 64 satisfies the natural alignment of every SIMD type. + self.tbh_ptrs_def.wb(f""" + #ifndef {guard} + #define {guard} + static inline void *spy_{macro_memkind}_alloc(size_t n) {{ + char *raw = (char *)spy_{memkind}_alloc(n + 64); + uintptr_t a = ((uintptr_t)raw + 63) & ~(uintptr_t)63; + return (void *)a; + }} + #endif + """) + else: + macro_memkind = memkind + self.tbh_ptrs_def.wb(f""" - SPY_PTR_FUNCTIONS({memkind}, {c_ptrtype}, {c_itemT}); + SPY_PTR_FUNCTIONS({macro_memkind}, {c_ptrtype}, {c_itemT}); #define {c_ptrtype}$NULL (({c_ptrtype}){{0}}) """) self.tbh_ptrs_def.wl() diff --git a/spy/backend/c/cwriter.py b/spy/backend/c/cwriter.py index dcb17b390..48af68886 100644 --- a/spy/backend/c/cwriter.py +++ b/spy/backend/c/cwriter.py @@ -579,6 +579,18 @@ def fmt_expr_Call(self, call: ast.Call) -> C.Expr: elif irtag.tag == "struct.getfield": return self.fmt_struct_getfield(fqn, call, irtag) + elif irtag.tag == "simd.make": + return self.fmt_simd_make(fqn, call, irtag) + + elif irtag.tag == "simd.getitem": + return self.fmt_simd_getitem(fqn, call) + + elif irtag.tag in ("simd.binop", "simd.cmp"): + return self.fmt_simd_binop(fqn, call, irtag) + + elif irtag.tag == "simd.select": + return self.fmt_simd_select(fqn, call) + elif irtag.tag == "ptr.getfield": return self.fmt_ptr_getfield(fqn, call, irtag) @@ -632,6 +644,65 @@ def fmt_struct_getfield(self, fqn: FQN, call: ast.Call, irtag: IRTag) -> C.Expr: name = irtag.data["name"] return C.Dot(c_struct, name) + def fmt_simd_make(self, fqn: FQN, call: ast.Call, irtag: IRTag) -> C.Expr: + from spy.vm.modules.simd import W_SimdType + + w_func = self.ctx.vm.lookup_global(fqn) + assert isinstance(w_func, W_Func) + w_simdtype = w_func.w_functype.w_restype + assert isinstance(w_simdtype, W_SimdType) + c_simdtype = self.ctx.w2c(w_simdtype) + + c_args = [self.fmt_expr(arg) for arg in call.args] + if irtag.data.get("broadcast"): + # SIMD[T, N](scalar) -> (T){scalar, scalar, ..., scalar} + assert len(c_args) == 1 + s_arg = str(c_args[0]) + strargs = ", ".join([s_arg] * w_simdtype.size) + else: + # SIMD[T, N](v0, ..., v_{N-1}) -> (T){v0, ..., v_{N-1}} + strargs = ", ".join(map(str, c_args)) + return C.Cast(c_simdtype, C.Literal("{ %s }" % strargs)) + + def fmt_simd_getitem(self, fqn: FQN, call: ast.Call) -> C.Expr: + assert len(call.args) == 2 + c_v = self.fmt_expr(call.args[0]) + c_i = self.fmt_expr(call.args[1]) + return C.Index(c_v, c_i) + + def fmt_simd_binop(self, fqn: FQN, call: ast.Call, irtag: IRTag) -> C.Expr: + assert len(call.args) == 2 + l, r = [self.fmt_expr(arg) for arg in call.args] + return C.BinOp(irtag.data["op"], l, r) + + def fmt_simd_select(self, fqn: FQN, call: ast.Call) -> C.Expr: + # call.args = [mask, "select" (blue str, ignored), a, b] + assert len(call.args) == 4 + from spy.vm.modules.simd import W_SimdType + + w_func = self.ctx.vm.lookup_global(fqn) + assert isinstance(w_func, W_Func) + w_ft = w_func.w_functype + w_mask_t = w_ft.params[0].w_T + w_op_t = w_ft.w_restype + assert isinstance(w_mask_t, W_SimdType) + assert isinstance(w_op_t, W_SimdType) + c_mask = self.ctx.w2c(w_mask_t) + c_op = self.ctx.w2c(w_op_t) + c_m = self.fmt_expr(call.args[0]) + c_a = self.fmt_expr(call.args[2]) + c_b = self.fmt_expr(call.args[3]) + # (T)((mask & (M)a) | (~mask & (M)b)) + # The vector ternary `?:` is C++-only, so we use the same-size + # bit-trick blend, valid for every dtype via reinterpret casts (M and + # T share the same total byte size, validated in _is_valid_mask). + blend = C.BinOp( + "|", + C.BinOp("&", c_m, C.Cast(c_mask, c_a)), + C.BinOp("&", C.UnaryOp("~", c_m), C.Cast(c_mask, c_b)), + ) + return C.Cast(c_op, C.Paren(blend)) + def fmt_ptr_getfield(self, fqn: FQN, call: ast.Call, irtag: IRTag) -> C.Expr: assert isinstance(call.args[1], ast.StrLiteral) c_ptr = self.fmt_expr(call.args[0]) diff --git a/spy/tests/compiler/test_simd.py b/spy/tests/compiler/test_simd.py new file mode 100644 index 000000000..68f2ec0ce --- /dev/null +++ b/spy/tests/compiler/test_simd.py @@ -0,0 +1,670 @@ +from spy.errors import SPyError +from spy.tests.support import CompilerTest, expect_errors, only_interp + + +class TestSIMD(CompilerTest): + # === construction & validation (blue-time) === + + def test_valid_sizes(self): + # size must be a positive power of two: 1, 2, 4, 8 are all valid. + mod = self.compile( + """ + from _simd import SIMD + + def s1(x: f32) -> f32: + v = SIMD[f32, 1](x) + return v[0] + + def s2(x: f32) -> f32: + v = SIMD[f32, 2](x) + return v[1] + + def s4(x: f32) -> f32: + v = SIMD[f32, 4](x) + return v[2] + + def s8(x: i32) -> i32: + v = SIMD[i32, 8](x) + return v[7] + """ + ) + assert mod.s1(1.5) == 1.5 + assert mod.s2(2.5) == 2.5 + assert mod.s4(3.5) == 3.5 + assert mod.s8(7) == 7 + + def test_all_dtypes(self): + # every v1 numeric primitive can be used as the lane dtype + mod = self.compile( + """ + from _simd import SIMD + + def f_i8(x: i32) -> i8: + v = SIMD[i8, 4](i8(x)) + return v[0] + + def f_u8(x: i32) -> u8: + v = SIMD[u8, 4](u8(x)) + return v[0] + + def f_i32(x: i32) -> i32: + v = SIMD[i32, 4](x) + return v[0] + + def f_u32(x: i32) -> u32: + v = SIMD[u32, 4](u32(x)) + return v[0] + + def f_i64(x: i32) -> i64: + v = SIMD[i64, 4](i64(x)) + return v[0] + + def f_u64(x: i32) -> u64: + v = SIMD[u64, 4](u64(i64(x))) + return v[0] + + def f_f32(x: f64) -> f32: + v = SIMD[f32, 4](f32(x)) + return v[0] + + def f_f64(x: f64) -> f64: + v = SIMD[f64, 4](x) + return v[0] + """ + ) + assert mod.f_i8(-5) == -5 + assert mod.f_u8(200) == 200 + assert mod.f_i32(42) == 42 + assert mod.f_u32(42) == 42 + assert mod.f_i64(42) == 42 + assert mod.f_u64(42) == 42 + assert mod.f_f32(1.5) == 1.5 + assert mod.f_f64(2.25) == 2.25 + + def test_invalid_size_not_power_of_two(self): + src = """ + from _simd import SIMD + + def bad() -> None: + v = SIMD[f32, 3](1.0) + """ + errors = expect_errors("SIMD size must be a power of two, got 3") + self.compile_raises(src, "bad", errors) + + def test_invalid_size_zero(self): + src = """ + from _simd import SIMD + + def bad() -> None: + v = SIMD[f32, 0](1.0) + """ + errors = expect_errors("SIMD size must be a positive power of two, got 0") + self.compile_raises(src, "bad", errors) + + def test_invalid_size_negative(self): + src = """ + from _simd import SIMD + + def bad() -> None: + v = SIMD[f32, -2](1.0) + """ + errors = expect_errors("SIMD size must be a positive power of two, got -2") + self.compile_raises(src, "bad", errors) + + def test_invalid_dtype_bool(self): + src = """ + from _simd import SIMD + + def bad() -> None: + v = SIMD[bool, 4](True) + """ + errors = expect_errors( + "SIMD element type must be a numeric primitive, got `bool`" + ) + self.compile_raises(src, "bad", errors) + + def test_invalid_dtype_str(self): + src = """ + from _simd import SIMD + + def bad() -> None: + v = SIMD[str, 4]("") + """ + errors = expect_errors( + "SIMD element type must be a numeric primitive, got `str`" + ) + self.compile_raises(src, "bad", errors) + + def test_invalid_dtype_struct(self): + src = """ + from _simd import SIMD + + @struct + class Point: + x: i32 + y: i32 + + def bad() -> None: + v = SIMD[Point, 4](Point(0, 0)) + """ + errors = expect_errors( + "SIMD element type must be a numeric primitive, got `test::Point`" + ) + self.compile_raises(src, "bad", errors) + + # === lane read v[i] (red index, simd.getitem irtag -> C.Index) === + + def test_per_element_read(self): + mod = self.compile( + """ + from _simd import SIMD + + def get_lane(idx: i32) -> f64: + v = SIMD[f64, 4](10.0, 20.0, 30.0, 40.0) + return v[idx] + """ + ) + assert mod.get_lane(0) == 10.0 + assert mod.get_lane(1) == 20.0 + assert mod.get_lane(2) == 30.0 + assert mod.get_lane(3) == 40.0 + + def test_broadcast_all_lanes_equal(self): + mod = self.compile( + """ + from _simd import SIMD + + def lane(a: f32, i: i32) -> f32: + v = SIMD[f32, 4](a) + return v[i] + """ + ) + for i in range(4): + assert mod.lane(1.25, i) == 1.25 + + def test_runtime_index_in_loop(self): + # the index `i` is a red (runtime) value: this is what vector-extension + # subscripting buys us over `tuple` (which requires blue indices). + mod = self.compile( + """ + from _simd import SIMD + + def sum_lanes(a: i32, b: i32, c: i32, d: i32) -> i32: + v = SIMD[i32, 4](a, b, c, d) + s: i32 = 0 + for i in range(4): + s = s + v[i] + return s + """ + ) + assert mod.sum_lanes(1, 2, 3, 4) == 10 + assert mod.sum_lanes(10, 20, 30, 40) == 100 + + @only_interp + def test_index_out_of_bounds(self): + # the interpreter bounds-checks lane access; the C backend lowers + # v[i] to a raw vector-extension subscript (no bounds check), so this + # panic is interp-only. W_PanicError matches the ptr.getitem convention + # (see unsafe/ptr.py::w_GETITEM). + mod = self.compile( + """ + from _simd import SIMD + + def bad(i: i32) -> f32: + v = SIMD[f32, 4](1.0) + return v[i] + """ + ) + with SPyError.raises("W_PanicError", match="SIMD index out of bounds"): + mod.bad(4) + with SPyError.raises("W_PanicError", match="SIMD index out of bounds"): + mod.bad(-1) + + # === whole-vector load/store through gc_ptr[SIMD[...]] === + + def test_store_load_whole_vector(self): + mod = self.compile( + """ + from unsafe import gc_alloc, gc_ptr + from _simd import SIMD + + def roundtrip(a: f32, b: f32, c: f32, d: f32) -> f32: + p: gc_ptr[SIMD[f32, 4]] = gc_alloc[SIMD[f32, 4]](1) + p[0] = SIMD[f32, 4](a, b, c, d) + v = p[0] + return v[2] + """ + ) + assert mod.roundtrip(1.0, 2.0, 3.0, 4.0) == 3.0 + + def test_store_load_multiple_vectors(self): + mod = self.compile( + """ + from unsafe import gc_alloc, gc_ptr + from _simd import SIMD + + def total() -> i32: + p: gc_ptr[SIMD[i32, 2]] = gc_alloc[SIMD[i32, 2]](3) + p[0] = SIMD[i32, 2](1, 2) + p[1] = SIMD[i32, 2](3, 4) + p[2] = SIMD[i32, 2](5, 6) + s: i32 = 0 + for i in range(3): + v = p[i] + s = s + v[0] + v[1] + return s + """ + ) + assert mod.total() == 21 + + def test_overwrite_slot(self): + mod = self.compile( + """ + from unsafe import gc_alloc, gc_ptr + from _simd import SIMD + + def overwrite() -> f32: + p: gc_ptr[SIMD[f32, 2]] = gc_alloc[SIMD[f32, 2]](1) + p[0] = SIMD[f32, 2](1.0, 2.0) + p[0] = SIMD[f32, 2](3.0, 4.0) + v = p[0] + return v[1] + """ + ) + assert mod.overwrite() == 4.0 + + def test_roundtrip_all_dtypes(self): + # one vector of each v1 dtype survives a store/load round-trip + mod = self.compile( + """ + from unsafe import gc_alloc, gc_ptr + from _simd import SIMD + + def rt_i8(x: i32) -> i8: + p: gc_ptr[SIMD[i8, 4]] = gc_alloc[SIMD[i8, 4]](1) + p[0] = SIMD[i8, 4](i8(x), i8(x), i8(x), i8(x)) + return p[0][0] + + def rt_u8(x: i32) -> u8: + p: gc_ptr[SIMD[u8, 4]] = gc_alloc[SIMD[u8, 4]](1) + p[0] = SIMD[u8, 4](u8(x), u8(x), u8(x), u8(x)) + return p[0][0] + + def rt_i32(x: i32) -> i32: + p: gc_ptr[SIMD[i32, 4]] = gc_alloc[SIMD[i32, 4]](1) + p[0] = SIMD[i32, 4](x, x, x, x) + return p[0][0] + + def rt_u32(x: i32) -> u32: + p: gc_ptr[SIMD[u32, 4]] = gc_alloc[SIMD[u32, 4]](1) + p[0] = SIMD[u32, 4](u32(x), u32(x), u32(x), u32(x)) + return p[0][0] + + def rt_i64(x: i32) -> i64: + p: gc_ptr[SIMD[i64, 4]] = gc_alloc[SIMD[i64, 4]](1) + p[0] = SIMD[i64, 4](i64(x), i64(x), i64(x), i64(x)) + return p[0][0] + + def rt_u64(x: i32) -> u64: + p: gc_ptr[SIMD[u64, 4]] = gc_alloc[SIMD[u64, 4]](1) + p[0] = SIMD[u64, 4](u64(i64(x)), u64(i64(x)), u64(i64(x)), u64(i64(x))) + return p[0][0] + + def rt_f32(x: f64) -> f32: + p: gc_ptr[SIMD[f32, 4]] = gc_alloc[SIMD[f32, 4]](1) + p[0] = SIMD[f32, 4](f32(x), f32(x), f32(x), f32(x)) + return p[0][0] + + def rt_f64(x: f64) -> f64: + p: gc_ptr[SIMD[f64, 4]] = gc_alloc[SIMD[f64, 4]](1) + p[0] = SIMD[f64, 4](x, x, x, x) + return p[0][0] + """ + ) + assert mod.rt_i8(-5) == -5 + assert mod.rt_u8(200) == 200 + assert mod.rt_i32(42) == 42 + assert mod.rt_u32(42) == 42 + assert mod.rt_i64(42) == 42 + assert mod.rt_u64(42) == 42 + assert mod.rt_f32(1.5) == 1.5 + assert mod.rt_f64(2.25) == 2.25 + + # === value semantics: immutable bare value, by-value pass/return === + + def test_bare_setitem_rejected(self): + # §4.5: no SIMD.__setitem__ on a bare value, only __getitem__. + # simd.setitem (lane write) is postponed to a later PR, so a bare + # `v[i] = x` is simply not supported. + src = """ + from _simd import SIMD + + def bad() -> None: + v = SIMD[i32, 4](1, 2, 3, 4) + v[0] = 99 + """ + errors = expect_errors( + "type `SIMD[i32, 4]` does not support item assignment", + ("this is `SIMD[i32, 4]`", "v"), + ) + self.compile_raises(src, "bad", errors) + + def test_pass_and_return_by_value(self): + # SIMD passed/returned by value between SPy functions. The exported + # `entry` returns a scalar, so no SIMD value crosses the WASM/Python + # boundary and this runs on all backends. + mod = self.compile( + """ + from _simd import SIMD + + def identity(v: SIMD[f32, 4]) -> SIMD[f32, 4]: + return v + + def first_lane(v: SIMD[f32, 4]) -> f32: + return v[0] + + def entry() -> f32: + v = SIMD[f32, 4](1.0, 2.0, 3.0, 4.0) + w = identity(v) + a = first_lane(v) + # w is a by-value copy: reads back the same lanes + return w[2] + a + """ + ) + assert mod.entry() == 4.0 # w[2] (3.0) + first_lane(v) (1.0) + + # === elementwise arithmetic: simd.binop === + + def test_binop_add_sub_mul_i32(self): + mod = self.compile( + """ + from _simd import SIMD + + def add0(a: i32, b: i32) -> i32: + v = SIMD[i32, 4](a, a, a, a) + w = SIMD[i32, 4](b, b, b, b) + return (v + w)[0] + + def sub0(a: i32, b: i32) -> i32: + v = SIMD[i32, 4](a, a, a, a) + w = SIMD[i32, 4](b, b, b, b) + return (v - w)[0] + + def mul0(a: i32, b: i32) -> i32: + v = SIMD[i32, 4](a, a, a, a) + w = SIMD[i32, 4](b, b, b, b) + return (v * w)[0] + """ + ) + assert mod.add0(3, 4) == 7 + assert mod.sub0(10, 4) == 6 + assert mod.mul0(3, 4) == 12 + + def test_binop_wraparound(self): + # SIMD integer arithmetic follows C wraparound + mod = self.compile( + """ + from _simd import SIMD + + def add_u8(a: i32) -> u8: + v = SIMD[u8, 4](u8(a)) + return (v + v)[0] + + def mul_i8(a: i32) -> i8: + v = SIMD[i8, 4](i8(a)) + w = SIMD[i8, 4](3) + return (v * w)[0] + """ + ) + assert mod.add_u8(200) == 144 # 400 mod 256 + assert mod.mul_i8(100) == 44 # 300 mod 256 = 44, as signed i8 + + def test_binop_add_all_dtypes(self): + mod = self.compile( + """ + from _simd import SIMD + + def f_i8(a: i32) -> i8: + v = SIMD[i8, 4](i8(a)) + return (v + v)[0] + + def f_u8(a: i32) -> u8: + v = SIMD[u8, 4](u8(a)) + return (v + v)[0] + + def f_i32(a: i32) -> i32: + v = SIMD[i32, 4](a) + return (v + v)[0] + + def f_u32(a: i32) -> u32: + v = SIMD[u32, 4](u32(a)) + return (v + v)[0] + + def f_i64(a: i32) -> i64: + v = SIMD[i64, 4](i64(a)) + return (v + v)[0] + + def f_u64(a: i32) -> u64: + v = SIMD[u64, 4](u64(i64(a))) + return (v + v)[0] + + def f_f32(a: f64) -> f32: + v = SIMD[f32, 4](f32(a)) + return (v + v)[0] + + def f_f64(a: f64) -> f64: + v = SIMD[f64, 4](a) + return (v + v)[0] + """ + ) + assert mod.f_i8(5) == 10 + assert mod.f_u8(100) == 200 + assert mod.f_i32(42) == 84 + assert mod.f_u32(42) == 84 + assert mod.f_i64(42) == 84 + assert mod.f_u64(42) == 84 + assert mod.f_f32(1.5) == 3.0 + assert mod.f_f64(2.25) == 4.5 + + def test_binop_float_div(self): + mod = self.compile( + """ + from _simd import SIMD + + def div_f32(a: f32, b: f32) -> f32: + v = SIMD[f32, 4](a, a, a, a) + w = SIMD[f32, 4](b, b, b, b) + return (v / w)[0] + + def div_f64(a: f64, b: f64) -> f64: + v = SIMD[f64, 4](a, a, a, a) + w = SIMD[f64, 4](b, b, b, b) + return (v / w)[0] + """ + ) + assert mod.div_f32(1.0, 4.0) == 0.25 + assert mod.div_f64(1.0, 8.0) == 0.125 + + def test_binop_int_div_error(self): + # integer `/` is deferred in v1 -> NULL -> standard type error. + src = """ + from _simd import SIMD + + def bad() -> None: + v = SIMD[i32, 4](1, 2, 3, 4) + w = SIMD[i32, 4](2, 2, 2, 2) + x = v / w + """ + errors = expect_errors("cannot do `SIMD[i32, 4]` / `SIMD[i32, 4]`") + self.compile_raises(src, "bad", errors) + + def test_binop_per_lane(self): + mod = self.compile( + """ + from _simd import SIMD + + def add0() -> i32: + v = SIMD[i32, 4](1, 2, 3, 4) + w = SIMD[i32, 4](10, 20, 30, 40) + r = v + w + s: i32 = 0 + for i in range(4): + s = s + r[i] + return s + """ + ) + assert mod.add0() == 110 # 11 + 22 + 33 + 44 + + # === elementwise comparison: simd.cmp === + + def test_cmp_i32(self): + mod = self.compile( + """ + from _simd import SIMD + + def lt(a: i32, b: i32) -> i32: + v = SIMD[i32, 4](a, a, a, a) + w = SIMD[i32, 4](b, b, b, b) + return (v < w)[0] + + def eq(a: i32, b: i32) -> i32: + v = SIMD[i32, 4](a, a, a, a) + w = SIMD[i32, 4](b, b, b, b) + return (v == w)[0] + + def ge(a: i32, b: i32) -> i32: + v = SIMD[i32, 4](a, a, a, a) + w = SIMD[i32, 4](b, b, b, b) + return (v >= w)[0] + """ + ) + # mask lanes are -1 (true, all-ones) / 0 (false), as signed i32. + assert mod.lt(3, 5) == -1 + assert mod.lt(5, 3) == 0 + assert mod.lt(3, 3) == 0 + assert mod.eq(3, 3) == -1 + assert mod.eq(3, 4) == 0 + assert mod.ge(3, 3) == -1 + assert mod.ge(2, 3) == 0 + + def test_cmp_mask_dtype_is_signed_int(self): + # f32 comparison yields a SIMD[i32, 4] mask; u8 yields SIMD[i8, 4]. + mod = self.compile( + """ + from _simd import SIMD + + def cmp_f32(a: f32, b: f32) -> i32: + v = SIMD[f32, 4](a, a, a, a) + w = SIMD[f32, 4](b, b, b, b) + return (v < w)[0] + + def cmp_u8(a: i32, b: i32) -> i8: + v = SIMD[u8, 4](u8(a)) + w = SIMD[u8, 4](u8(b)) + return (v < w)[0] + """ + ) + assert mod.cmp_f32(1.0, 2.0) == -1 + assert mod.cmp_f32(2.0, 1.0) == 0 + assert mod.cmp_u8(1, 2) == -1 + assert mod.cmp_u8(2, 1) == 0 + + def test_cmp_per_lane(self): + mod = self.compile( + """ + from _simd import SIMD + + def f() -> i32: + v = SIMD[i32, 4](1, 5, 3, 7) + w = SIMD[i32, 4](2, 4, 3, 0) + m = v < w # T, F, F, F + s: i32 = 0 + for i in range(4): + s = s + m[i] + return s + """ + ) + assert mod.f() == -1 # only lane 0 is true -> one -1 + + # === mask.select(a, b): simd.select === + + def test_select_cmp_mask(self): + mod = self.compile( + """ + from _simd import SIMD + + def sel(a: f32, b: f32) -> f32: + v = SIMD[f32, 4](a, a, a, a) + w = SIMD[f32, 4](b, b, b, b) + m = v < w + r = m.select(v, w) + return r[0] + """ + ) + # a < b => mask true => select picks the first arg (v = a) + assert mod.sel(1.0, 2.0) == 1.0 + # a > b => mask false => picks the second arg (w = b) + assert mod.sel(3.0, 2.0) == 2.0 + # a == b => mask false => picks b + assert mod.sel(2.0, 2.0) == 2.0 + + def test_select_per_lane(self): + mod = self.compile( + """ + from _simd import SIMD + + def sel(idx: i32) -> f32: + v = SIMD[f32, 4](10.0, 20.0, 30.0, 40.0) + w = SIMD[f32, 4](1.0, 2.0, 3.0, 4.0) + m = SIMD[i32, 4](-1, 0, -1, 0) + r = m.select(v, w) + return r[idx] + """ + ) + # canonical mask lanes (-1 picks v, 0 picks w): interp == C. + assert mod.sel(0) == 10.0 + assert mod.sel(1) == 2.0 + assert mod.sel(2) == 30.0 + assert mod.sel(3) == 4.0 + + def test_select_max(self): + # classic blend idiom: per-lane max via (a > b).select(a, b). + mod = self.compile( + """ + from _simd import SIMD + + def vmax0(a: f32, b: f32) -> f32: + v = SIMD[f32, 4](a, a, a, a) + w = SIMD[f32, 4](b, b, b, b) + m = v > w + r = m.select(v, w) + return r[0] + + def vmax_all() -> f32: + v = SIMD[f32, 4](10.0, 2.0, 30.0, 4.0) + w = SIMD[f32, 4](5.0, 8.0, 1.0, 9.0) + m = v > w + r = m.select(v, w) + s: f32 = 0.0 + for i in range(4): + s = s + r[i] + return s + """ + ) + assert mod.vmax0(1.0, 2.0) == 2.0 + assert mod.vmax0(5.0, 2.0) == 5.0 + assert mod.vmax_all() == 57.0 # 10 + 8 + 30 + 9 + + def test_select_wrong_mask_size_error(self): + # a mask whose size does not match the operand is not a valid select. + src = """ + from _simd import SIMD + + def bad() -> None: + v = SIMD[f32, 4](1.0) + w = SIMD[f32, 4](2.0) + m = SIMD[i32, 2](-1, 0) + r = m.select(v, w) + """ + errors = expect_errors("method `SIMD[i32, 2]::select` does not exist") + self.compile_raises(src, "bad", errors) diff --git a/spy/vm/modules/simd.py b/spy/vm/modules/simd.py new file mode 100644 index 000000000..c7d109b51 --- /dev/null +++ b/spy/vm/modules/simd.py @@ -0,0 +1,736 @@ +""" +This module implements the low-level internal ``_simd`` VM module, exposing ``SIMD``. +""" + +import operator +from typing import TYPE_CHECKING, Annotated, Any + +from spy.errors import SPyError +from spy.fqn import FQN +from spy.vm.b import B +from spy.vm.builtin import builtin_method +from spy.vm.irtag import IRTag +from spy.vm.object import W_Object, W_Type +from spy.vm.opspec import W_MetaArg, W_OpSpec +from spy.vm.primitive import ( + W_F32, + W_F64, + W_I8, + W_I32, + W_I64, + W_U8, + W_U32, + W_U64, + W_Dynamic, +) +from spy.vm.registry import ModuleRegistry + +if TYPE_CHECKING: + from spy.vm.vm import SPyVM + + +SIMD = ModuleRegistry("_simd") + + +# The set of numeric primitives which are legal SIMD lane dtypes. +SIMD_DTYPES = ( + B.w_i8, + B.w_u8, + B.w_i32, + B.w_u32, + B.w_f32, + B.w_i64, + B.w_u64, + B.w_f64, +) + + +# the mask dtype for a lane dtype. +# +# GCC/Clang `vector_size` comparison yields a *signed* integer vector of the +# same byte-width as the lane, with lanes -1 (true, all-ones) / 0 (false). So +# the mask dtype is the signed integer of the lane's byte-width: +SIMD_MASK_DTYPE = { + B.w_i8: B.w_i8, + B.w_u8: B.w_i8, + B.w_i32: B.w_i32, + B.w_u32: B.w_i32, + B.w_f32: B.w_i32, + B.w_i64: B.w_i64, + B.w_u64: B.w_i64, + B.w_f64: B.w_i64, +} + + +SIMD_DTYPE_BYTES = { + B.w_i8: 1, + B.w_u8: 1, + B.w_i32: 4, + B.w_u32: 4, + B.w_f32: 4, + B.w_i64: 8, + B.w_u64: 8, + B.w_f64: 8, +} + +# integer lane dtypes: only these may serve as a select mask. +SIMD_INT_DTYPES = frozenset({B.w_i8, B.w_u8, B.w_i32, B.w_u32, B.w_i64, B.w_u64}) + +_W_LANE_CTOR = { + B.w_i8: W_I8, + B.w_u8: W_U8, + B.w_i32: W_I32, + B.w_u32: W_U32, + B.w_f32: W_F32, + B.w_i64: W_I64, + B.w_u64: W_U64, + B.w_f64: W_F64, +} + + +@SIMD.builtin_type("SimdType") +class W_SimdType(W_Type): + """ + The *type* of a SIMD vector, e.g. ``SIMD[f32, 4]``. + + A concrete ``W_SimdType`` instance is created (and cached) by the + ``SIMD`` blue generic (see ``w_SIMD`` below). Each instance carries its + lane ``dtype`` (a primitive ``W_Type``) and its ``size`` (the number of + lanes), and is defined from the ``W_Simd`` value pyclass — which is what + installs ``__getitem__`` / ``__setitem__`` / ``__new__`` into its + ``dict_w``. + + Like ``W_StructType`` / ``W_PtrType``, the *type* lives outside + ``unsafe``; only the memory I/O (``sizeof`` and the + ``generic_mem_read``/``generic_mem_write`` branch) lives in ``unsafe``. + """ + + w_dtype: W_Type + size: int + + def repr_hints(self) -> list[str]: + return super().repr_hints() + ["simd"] + + def is_struct(self, vm: "SPyVM") -> bool: + # SIMD vectors are not structs: ptr[SIMD[...]] loads/stores them + # *by value* (see W_Ptr.w_GETITEM), they never become a ref[T]. + return False + + +@SIMD.builtin_type("Simd") +class W_Simd(W_Object): + """ + A SIMD vector *value*, e.g. an instance of ``SIMD[f32, 4]``. + + Interp-level representation: a plain Python list of ``size`` lane values + (each a ``W_Object`` of ``w_dtype``). This is a *value* type: it is + immutable (only ``__getitem__``, no ``__setitem__``), compares by value, + and is passed/returned by value between SPy functions. Mutation / + addressing of individual lanes exists only through ``ptr[SIMD[...]]``, + mirroring ``struct``. + """ + + __spy_storage_category__ = "value" + + w_simdtype: W_SimdType + lanes_w: list # list[W_Object], length == w_simdtype.size + + def __init__(self, w_simdtype: W_SimdType, lanes_w: list) -> None: + assert len(lanes_w) == w_simdtype.size + self.w_simdtype = w_simdtype + self.lanes_w = lanes_w + + def spy_get_w_type(self, vm: "SPyVM") -> W_Type: + # The app-level type is the concrete W_SimdType (e.g. + # `_simd::SIMD[f32, 4]`), NOT the `Simd` base type registered above. + return self.w_simdtype + + def spy_key(self, vm: "SPyVM") -> Any: + t = self.w_simdtype.spy_key(vm) + lanes = tuple(w_lane.spy_key(vm) for w_lane in self.lanes_w) + return ("simd", t, lanes) + + def __repr__(self) -> str: + fqn = self.w_simdtype.fqn + return f"" + + # ===== construction: SIMD[T, N](...) ===== + # + # Calling a W_SimdType means "instantiate it". W_Type.w_CALL dispatches to + # __new__; here we turn it into a `simd.make` builtin call (compound + # literal in C). Two shapes are supported: + # * broadcast: SIMD[T, N](scalar) -> {scalar, ..., scalar} + # * per-element: SIMD[T, N](v0, ..., vN-1) -> {v0, ..., vN-1} + # + # We build the lowering builtin explicitly (like struct's `_create_w_make`) + # with a fixed-arity W_FuncType, rather than deriving it from a Python + # signature: per-element make needs exactly `size` params, which we + # cannot spell as a static Python signature. + @builtin_method("__new__", color="blue", kind="metafunc") + @staticmethod + def w_NEW(vm: "SPyVM", wam_self: W_MetaArg, *args_wam: W_MetaArg) -> W_OpSpec: + w_simdtype = wam_self.w_blueval + assert isinstance(w_simdtype, W_SimdType) + size = w_simdtype.size + nargs = len(args_wam) + + if nargs == 1: + # broadcast: SIMD[T, N](scalar) + w_make = _get_or_make_simd_make(vm, w_simdtype, broadcast=True) + return W_OpSpec(w_make, [args_wam[0]]) + + elif nargs == size: + # per-element: SIMD[T, N](v0, ..., v_{N-1}) + w_make = _get_or_make_simd_make(vm, w_simdtype, broadcast=False) + return W_OpSpec(w_make, list(args_wam)) + + else: + # Anything else (e.g. 2 args for a 4-wide vector): not supported in + # PR1. Returning NULL yields a clear "cannot call" type error. + return W_OpSpec.NULL + + # ===== lane read: v[i] (red index) -> simd.getitem ===== + @builtin_method("__getitem__", color="blue", kind="metafunc") + @staticmethod + def w_GETITEM(vm: "SPyVM", wam_self: W_MetaArg, wam_i: W_MetaArg) -> W_OpSpec: + w_simdtype = wam_self.w_static_T + assert isinstance(w_simdtype, W_SimdType) + w_dtype = w_simdtype.w_dtype + size = w_simdtype.size + + SIMD_T = Annotated[W_Simd, w_simdtype] + T = Annotated[W_Object, w_dtype] + irtag = IRTag("simd.getitem") + + @vm.register_builtin_func(w_simdtype.fqn, "getitem", irtag=irtag) + def w_simd_getitem(vm: "SPyVM", w_v: SIMD_T, w_i: W_I32) -> T: + i = vm.unwrap_i32(w_i) + if not (0 <= i < size): + raise SPyError("W_PanicError", "SIMD index out of bounds") + return w_v.lanes_w[i] + + return W_OpSpec(w_simd_getitem, [wam_self, wam_i]) + + # ===== lane write: v[i] = x — rejected (§4.5) ===== + # + # SIMD values are immutable: only __getitem__ is provided. A bare + # `v[i] = x` is not supported in PR1 (simd.setitem is postponed to a later + # PR). We implement __setitem__ as a metafunc that raises a precise error + # instead of letting the generic "cannot do `{0}[`{1}`] = ...` message + # through, so the diagnostic matches the value-semantics contract. + @builtin_method("__setitem__", color="blue", kind="metafunc") + @staticmethod + def w_SETITEM( + vm: "SPyVM", wam_self: W_MetaArg, wam_i: W_MetaArg, wam_v: W_MetaArg + ) -> W_OpSpec: + w_simdtype = wam_self.w_static_T + assert isinstance(w_simdtype, W_SimdType) + t = w_simdtype.fqn.human_name(vm) + err = SPyError("W_TypeError", f"type `{t}` does not support item assignment") + err.add("error", f"this is `{t}`", wam_self.loc) + raise err + + # ===== elementwise arithmetic: simd.binop ===== + # + # `a + b`, `a - b`, `a * b`, `a / b` (float only) lower to a single + # `simd.binop` irtag carrying the C operator; the C backend emits + # `C.BinOp(op, l, r)`. Each metafunc builds (once per (W_SimdType, op)) + # a red plain builtin with functype `(T, T) -> T` and returns a SIMPLE + # OpSpec to it, so `typecheck_opspec` rebinds against the live call args. + + @builtin_method("__add__", color="blue", kind="metafunc") + @staticmethod + def w_ADD(vm: "SPyVM", wam_self: W_MetaArg, wam_other: W_MetaArg) -> W_OpSpec: + return _simd_binop_meta( + vm, + wam_self, + wam_other, + dunder="add", + c_op="+", + op_py=operator.add, + ) + + @builtin_method("__sub__", color="blue", kind="metafunc") + @staticmethod + def w_SUB(vm: "SPyVM", wam_self: W_MetaArg, wam_other: W_MetaArg) -> W_OpSpec: + return _simd_binop_meta( + vm, + wam_self, + wam_other, + dunder="sub", + c_op="-", + op_py=operator.sub, + ) + + @builtin_method("__mul__", color="blue", kind="metafunc") + @staticmethod + def w_MUL(vm: "SPyVM", wam_self: W_MetaArg, wam_other: W_MetaArg) -> W_OpSpec: + return _simd_binop_meta( + vm, + wam_self, + wam_other, + dunder="mul", + c_op="*", + op_py=operator.mul, + ) + + @builtin_method("__div__", color="blue", kind="metafunc") + @staticmethod + def w_DIV(vm: "SPyVM", wam_self: W_MetaArg, wam_other: W_MetaArg) -> W_OpSpec: + # `/` is v1 float-only: integer `/` (truncation-vs-floor) is a + # semantic decision deferred to a follow-up. Returning NULL yields + # the standard `cannot do `SIMD[i32,4]` / `SIMD[i32,4]`` type error. + return _simd_binop_meta( + vm, + wam_self, + wam_other, + dunder="div", + c_op="/", + op_py=operator.truediv, + ) + + # ===== elementwise comparison: simd.cmp ===== + # + # All six comparisons lower to the `simd.cmp` irtag carrying the C + # operator; the C result is a *signed* integer vector (the mask type, see + # `get_mask_simdtype`) of the lane's byte-width, with lanes -1 (true) / + # 0 (false) -- matching the GCC/Clang `vector_size` comparison result. + # Defining all six explicitly also overrides the scalar `__eq__`/`__ne__` + # that `W_Type.define._add_eq_ne_maybe` would otherwise auto-generate from + # `spy_key` (W_Simd is a value type with a spy_key). + + @builtin_method("__eq__", color="blue", kind="metafunc") + @staticmethod + def w_EQ(vm: "SPyVM", wam_self: W_MetaArg, wam_other: W_MetaArg) -> W_OpSpec: + return _simd_cmp_meta( + vm, + wam_self, + wam_other, + dunder="eq", + c_op="==", + cmp_py=operator.eq, + ) + + @builtin_method("__ne__", color="blue", kind="metafunc") + @staticmethod + def w_NE(vm: "SPyVM", wam_self: W_MetaArg, wam_other: W_MetaArg) -> W_OpSpec: + return _simd_cmp_meta( + vm, + wam_self, + wam_other, + dunder="ne", + c_op="!=", + cmp_py=operator.ne, + ) + + @builtin_method("__lt__", color="blue", kind="metafunc") + @staticmethod + def w_LT(vm: "SPyVM", wam_self: W_MetaArg, wam_other: W_MetaArg) -> W_OpSpec: + return _simd_cmp_meta( + vm, + wam_self, + wam_other, + dunder="lt", + c_op="<", + cmp_py=operator.lt, + ) + + @builtin_method("__le__", color="blue", kind="metafunc") + @staticmethod + def w_LE(vm: "SPyVM", wam_self: W_MetaArg, wam_other: W_MetaArg) -> W_OpSpec: + return _simd_cmp_meta( + vm, + wam_self, + wam_other, + dunder="le", + c_op="<=", + cmp_py=operator.le, + ) + + @builtin_method("__gt__", color="blue", kind="metafunc") + @staticmethod + def w_GT(vm: "SPyVM", wam_self: W_MetaArg, wam_other: W_MetaArg) -> W_OpSpec: + return _simd_cmp_meta( + vm, + wam_self, + wam_other, + dunder="gt", + c_op=">", + cmp_py=operator.gt, + ) + + @builtin_method("__ge__", color="blue", kind="metafunc") + @staticmethod + def w_GE(vm: "SPyVM", wam_self: W_MetaArg, wam_other: W_MetaArg) -> W_OpSpec: + return _simd_cmp_meta( + vm, + wam_self, + wam_other, + dunder="ge", + c_op=">=", + cmp_py=operator.ge, + ) + + # ===== mask.select(a, b): simd.select ===== + # + # `mask.select(a, b)` is a method call. The default call-method machinery + # (`default_callmethod`) would route the metafunc through `op_METACALL`, + # which wraps every operand with `W_MetaArg.from_w_obj` and *loses* the + # concrete SIMD types (`wam_self.w_static_T` would become `W_MetaArg`). + # We therefore define `__call_method__` on `W_Simd`: `w_CALL_METHOD` + # (callop) dispatches it via `fast_metacall`, passing the *real* red + # MetaArgs, so we can read the mask/operand `W_SimdType`s. + # + # We handle "select" (building the per-(mask, operand) `simd.select` + # builtin) and return `W_OpSpec.NULL` for anything else, which produces + # the `method `...::meth` does not exist` error. + + @builtin_method("__call_method__", color="blue", kind="metafunc") + @staticmethod + def w_CALL_METHOD( + vm: "SPyVM", + wam_self: W_MetaArg, + wam_meth: W_MetaArg, + *args_wam: W_MetaArg, + ) -> W_OpSpec: + if not wam_meth.is_blue(): + return W_OpSpec.NULL + meth = vm.unwrap_str(wam_meth.w_blueval) + if meth == "select" and len(args_wam) == 2: + w_mask_t = wam_self.w_static_T + w_op_t = args_wam[0].w_static_T + if ( + isinstance(w_mask_t, W_SimdType) + and isinstance(w_op_t, W_SimdType) + and _is_valid_mask(w_mask_t, w_op_t) + ): + w_sel = _get_or_make_simd_select(vm, w_mask_t, w_op_t) + return W_OpSpec(w_sel) + return W_OpSpec.NULL + + +def _get_or_make_simd_make( + vm: "SPyVM", w_simdtype: W_SimdType, *, broadcast: bool +) -> "W_BuiltinFunc": # type: ignore[name-defined] + """ + Build (once per (W_SimdType, shape)) and register the red ``simd.make`` + lowering builtin, returning the cached instance on subsequent calls. + + This mirrors struct's ``W_StructType._create_w_make``: we construct the + ``W_BuiltinFunc`` directly with a fixed-arity ``W_FuncType`` rather than + deriving the functype from a Python signature (per-element make needs + exactly ``size`` params, which cannot be spelled statically). + + The broadcast and per-element lowers share the ``simd.make`` irtag (the + C backend dispatches on the tag and inspects ``irtag.data['broadcast']``), + but live at distinct FQNs because they have different arities. + """ + from spy.vm.function import FuncParam, W_BuiltinFunc, W_FuncType + + w_dtype = w_simdtype.w_dtype + size = w_simdtype.size + + if broadcast: + fqn = w_simdtype.fqn.join("__make_broadcast__") + w_functype = W_FuncType.new([FuncParam(w_dtype, "simple")], w_simdtype) + irtag = IRTag("simd.make", broadcast=True) + + def w_make_impl(vm: "SPyVM", w_x: W_Object) -> W_Simd: + return W_Simd(w_simdtype, [w_x] * size) + + else: + fqn = w_simdtype.fqn.join("__make__") + params = [FuncParam(w_dtype, "simple") for _ in range(size)] + w_functype = W_FuncType.new(params, w_simdtype) + irtag = IRTag("simd.make") + + def w_make_impl(vm: "SPyVM", *args_w: W_Object) -> W_Simd: # type: ignore[misc] + assert len(args_w) == size + return W_Simd(w_simdtype, list(args_w)) + + w_existing = vm.lookup_global_maybe(fqn) + if w_existing is not None: + # Already registered by an earlier call site (or a re-typecheck). + # W_FuncType is interned, so the functype is the very same object. + assert isinstance(w_existing, W_BuiltinFunc) + assert w_existing.w_functype is w_functype + return w_existing + + w_func = W_BuiltinFunc(w_functype, fqn, w_make_impl) + vm.add_global(fqn, w_func, irtag=irtag) + return w_func + + +# ===== mask type + binop/cmp/select lowering builtins ===== + + +def _lane_py(w_lane: Any, w_dtype: W_Type) -> Any: + """ + Unwrap a SIMD lane W_Object to a plain Python value for interp + arithmetic. Floats -> python float, ints -> python int (signedness + preserved via ``FixedInt.__int__``). Re-wrapping with the lane ctor + re-applies the width, and hence the C wraparound/narrowing semantics. + """ + if w_dtype is B.w_f32: + return w_lane.value.value + if w_dtype is B.w_f64: + return w_lane.value + return int(w_lane.value) + + +def get_mask_simdtype(vm: "SPyVM", w_simdtype: W_SimdType) -> W_SimdType: + """ + The mask ``W_SimdType`` for a given operand ``W_SimdType``: the + signed-integer SIMD vector of the lane's byte-width and the same size + (e.g. ``SIMD[f32, 4]`` -> ``SIMD[i32, 4]``). Built by calling the ``SIMD`` + blue generic, so it is interned + registered as a global exactly like a + user-written ``SIMD[i32, 4]`` (and emitted as a typedef by the C backend). + """ + w_mask_dtype = SIMD_MASK_DTYPE[w_simdtype.w_dtype] + size = int(w_simdtype.size) + w_mask_simdtype = vm.fast_call(SIMD.w_SIMD, [w_mask_dtype, W_I32(size)]) + assert isinstance(w_mask_simdtype, W_SimdType) + return w_mask_simdtype + + +def _is_valid_mask(w_mask_t: W_SimdType, w_op_t: W_SimdType) -> bool: + """ + A select mask must be an integer SIMD vector of the same size and the + same lane byte-width as the operand, so the C same-size reinterpret + casts in the bit-trick blend are valid. + """ + return ( + w_mask_t.size == w_op_t.size + and w_mask_t.w_dtype in SIMD_INT_DTYPES + and SIMD_DTYPE_BYTES[w_mask_t.w_dtype] == SIMD_DTYPE_BYTES[w_op_t.w_dtype] + ) + + +def _get_or_make_simd_op( + vm: "SPyVM", + w_simdtype: W_SimdType, + *, + dunder: str, + c_op: str, + tag: str, + w_restype: W_Type, + w_impl: Any, +) -> "W_BuiltinFunc": # type: ignore[name-defined] + """ + Build (once per (W_SimdType, op)) and register the red lowering builtin + for a binop/cmp, returning the cached instance on subsequent calls. + Mirrors ``_get_or_make_simd_make``: explicit ``W_FuncType`` (so we can set + the cmp restype to the mask type), ``lookup_global_maybe`` caching, a + distinct FQN per (W_SimdType, op). The C backend dispatches on ``tag`` + and reads ``irtag.data['op']``. + """ + from spy.vm.function import FuncParam, W_BuiltinFunc, W_FuncType + + fqn = w_simdtype.fqn.join(f"__{dunder}__") + w_functype = W_FuncType.new( + [FuncParam(w_simdtype, "simple"), FuncParam(w_simdtype, "simple")], + w_restype, + ) + irtag = IRTag(tag, op=c_op) + + w_existing = vm.lookup_global_maybe(fqn) + if w_existing is not None: + assert isinstance(w_existing, W_BuiltinFunc) + return w_existing + + w_func = W_BuiltinFunc(w_functype, fqn, w_impl) + vm.add_global(fqn, w_func, irtag=irtag) + return w_func + + +def _get_or_make_simd_select( + vm: "SPyVM", w_mask_t: W_SimdType, w_op_t: W_SimdType +) -> "W_BuiltinFunc": # type: ignore[name-defined] + """ + Build (once per (mask, operand)) and register the red ``simd.select`` + lowering builtin. Functype ``(mask, str, T, T) -> T``: the ``str`` param + is the method name carried by ``w_CALL_METHOD``; it is ignored by the + impl and skipped by the C lowering, but keeping it lets us return a + *simple* OpSpec (caching-safe). The C backend reads the mask/operand C + types from this functype. The FQN carries the mask type as a qualifier so + different masks over the same operand get distinct builtins. + """ + from spy.vm.function import FuncParam, W_BuiltinFunc, W_FuncType + + fqn = w_op_t.fqn.join("__select__", qualifiers=[w_mask_t.fqn]) + w_functype = W_FuncType.new( + [ + FuncParam(w_mask_t, "simple"), + FuncParam(B.w_str, "simple"), + FuncParam(w_op_t, "simple"), + FuncParam(w_op_t, "simple"), + ], + w_op_t, + ) + irtag = IRTag("simd.select") + + w_existing = vm.lookup_global_maybe(fqn) + if w_existing is not None: + assert isinstance(w_existing, W_BuiltinFunc) + return w_existing + + def w_impl( + vm: "SPyVM", w_mask: W_Simd, w_meth: W_Object, w_a: W_Simd, w_b: W_Simd + ) -> W_Simd: + # NOTE: interp picks `a[i]`/`b[i]` per `mask[i] != 0`. This matches + # the C bit-trick blend `(T)((mask & (M)a) | (~mask & (M)b))` for + # *canonical* masks (comparison results, lanes 0 / -1), which is the + # v1 contract. Arbitrary integer masks are accepted at the type level + # but may diverge between interp and C; use comparison masks for + # guaranteed cross-backend parity. + lanes = [ + x if int(m.value) != 0 else y + for m, x, y in zip(w_mask.lanes_w, w_a.lanes_w, w_b.lanes_w) + ] + return W_Simd(w_op_t, lanes) + + w_func = W_BuiltinFunc(w_functype, fqn, w_impl) + vm.add_global(fqn, w_func, irtag=irtag) + return w_func + + +def _simd_binop_meta( + vm: "SPyVM", + wam_self: W_MetaArg, + wam_other: W_MetaArg, + *, + dunder: str, + c_op: str, + op_py: Any, +) -> W_OpSpec: + w_simdtype = wam_self.w_static_T + assert isinstance(w_simdtype, W_SimdType) + # `/` is v1 float-only; integer `/` is deferred (NULL -> type error). + if c_op == "/" and not ( + w_simdtype.w_dtype is B.w_f32 or w_simdtype.w_dtype is B.w_f64 + ): + return W_OpSpec.NULL + + w_dtype = w_simdtype.w_dtype + lane_ctor = _W_LANE_CTOR[w_dtype] + + def w_impl(vm: "SPyVM", w_a: W_Simd, w_b: W_Simd) -> W_Simd: + lanes = [ + lane_ctor(op_py(_lane_py(x, w_dtype), _lane_py(y, w_dtype))) + for x, y in zip(w_a.lanes_w, w_b.lanes_w) + ] + return W_Simd(w_simdtype, lanes) + + w_func = _get_or_make_simd_op( + vm, + w_simdtype, + dunder=dunder, + c_op=c_op, + tag="simd.binop", + w_restype=w_simdtype, + w_impl=w_impl, + ) + return W_OpSpec(w_func) + + +def _simd_cmp_meta( + vm: "SPyVM", + wam_self: W_MetaArg, + wam_other: W_MetaArg, + *, + dunder: str, + c_op: str, + cmp_py: Any, +) -> W_OpSpec: + w_simdtype = wam_self.w_static_T + assert isinstance(w_simdtype, W_SimdType) + w_mask_simdtype = get_mask_simdtype(vm, w_simdtype) + w_dtype = w_simdtype.w_dtype + mask_ctor = _W_LANE_CTOR[w_mask_simdtype.w_dtype] + + def w_impl(vm: "SPyVM", w_a: W_Simd, w_b: W_Simd) -> W_Simd: + lanes = [ + mask_ctor(-1 if cmp_py(_lane_py(x, w_dtype), _lane_py(y, w_dtype)) else 0) + for x, y in zip(w_a.lanes_w, w_b.lanes_w) + ] + return W_Simd(w_mask_simdtype, lanes) + + w_func = _get_or_make_simd_op( + vm, + w_simdtype, + dunder=dunder, + c_op=c_op, + tag="simd.cmp", + w_restype=w_mask_simdtype, + w_impl=w_impl, + ) + return W_OpSpec(w_func) + + +@SIMD.builtin_func(color="blue", kind="generic") +def w_SIMD(vm: "SPyVM", w_dtype: W_Type, w_size: W_I32) -> W_Dynamic: + """ + The ``SIMD`` *generic* type constructor. + + ``SIMD[dtype, size]`` is a blue ``getitem`` on the generic ``SIMD`` + function: it calls ``w_SIMD`` with the (blue) ``dtype`` type and the + (blue) ``size`` integer, validates them, and returns — and registers — the + concrete ``W_SimdType`` for that ``(dtype, size)`` pair. + + Validation (blue-time): + + * ``size`` must be a *positive power of two* (1, 2, 4, 8, ...). + - non-positive sizes (0, negative) report + ``"SIMD size must be a positive power of two, got "``; + - positive but non-power-of-two sizes report + ``"SIMD size must be a power of two, got "``. + * ``dtype`` must be one of the v1 numeric primitives + (i8, u8, i32, u32, f32, i64, u64, f64), else + ``"SIMD element type must be a numeric primitive, got ``"``. + """ + size = int(vm.unwrap_i32(w_size)) + + # === validate size === + if size <= 0: + raise SPyError( + "W_TypeError", f"SIMD size must be a positive power of two, got {size}" + ) + if size & (size - 1) != 0: + raise SPyError("W_TypeError", f"SIMD size must be a power of two, got {size}") + + # === validate dtype === + if w_dtype not in SIMD_DTYPES: + t = w_dtype.fqn.human_name(vm) + raise SPyError( + "W_TypeError", + f"SIMD element type must be a numeric primitive, got `{t}`", + ) + + # === register the human alias `_simd::SIMD` -> `SIMD` === + # + # Unlike `list`/`dict`/`tuple`, `SIMD` is NOT re-exported from the builtins + # prelude (PR1 exposes only the low-level `_simd` module), so it does not + # get a seeded human alias. We register one manually so that error + # messages render `SIMD[f32, 4]` instead of `_simd::SIMD[f32, 4]`. + # `_resolve_aliases` reattaches the qualifiers, so `_simd::SIMD[f32, 4]` + # resolves to `SIMD[f32, 4]`. + vm.fqn_human_aliases[FQN("_simd::SIMD")] = FQN("SIMD") + + # === build the concrete W_SimdType === + # + # The FQN carries both the dtype and the size as qualifiers, so that + # `SIMD[f32, 4]` human-renders as `SIMD[f32, 4]` and C-mangles to a stable, + # distinct typedef name `spy__simd$SIMD__f32_4` (one typedef per + # (dtype, size) pair). The size is encoded as a bare FQN qualifier, which + # fqn.c_name renders verbatim. + fqn = FQN("_simd::SIMD").with_qualifiers([w_dtype.fqn, str(size)]) + + # The blue cache memoizes w_SIMD by (dtype spy_key, size spy_key), so + # repeated `SIMD[f32, 4]` evaluations return the *same* W_SimdType + # instance. make_fqn_const then ensures the type is reachable as a global + # (needed by gc_ptr[SIMD[...]] and by the C backend). + w_simdtype = W_SimdType.from_pyclass(fqn, W_Simd) + w_simdtype.w_dtype = w_dtype + w_simdtype.size = size + vm.make_fqn_const(w_simdtype) + return w_simdtype diff --git a/spy/vm/modules/unsafe/mem.py b/spy/vm/modules/unsafe/mem.py index e6fda34f5..bfa9827a0 100644 --- a/spy/vm/modules/unsafe/mem.py +++ b/spy/vm/modules/unsafe/mem.py @@ -535,6 +535,7 @@ def w_mem_write_T(vm: "SPyVM", w_addr: W_I32, w_val: T) -> None: def generic_mem_read(vm: "SPyVM", addr: int, w_T: W_Type) -> W_Object: from spy.vm.modules.posix import POSIX, W__FILE + from spy.vm.modules.simd import W_Simd, W_SimdType if w_T is B.w_i8: return W_I8(vm.ll.mem.read_i8(addr)) @@ -569,6 +570,13 @@ def generic_mem_read(vm: "SPyVM", addr: int, w_T: W_Type) -> W_Object: offset = w_field.offset values_w[fname] = generic_mem_read(vm, addr + offset, w_field.w_T) return W_Struct(w_T, values_w) + elif isinstance(w_T, W_SimdType): + lane_size = sizeof(w_T.w_dtype) + lanes_w = [ + generic_mem_read(vm, addr + i * lane_size, w_T.w_dtype) + for i in range(w_T.size) + ] + return W_Simd(w_T, lanes_w) else: t = w_T.fqn.human_name(vm) raise WIP(f"Cannot read memory of type `{t}`") @@ -576,6 +584,7 @@ def generic_mem_read(vm: "SPyVM", addr: int, w_T: W_Type) -> W_Object: def generic_mem_write(vm: "SPyVM", addr: int, w_T: W_Type, w_val: W_Object) -> None: from spy.vm.modules.posix import POSIX, W__FILE + from spy.vm.modules.simd import W_Simd, W_SimdType if w_T is B.w_i8: assert isinstance(w_val, W_I8) @@ -621,6 +630,11 @@ def generic_mem_write(vm: "SPyVM", addr: int, w_T: W_Type, w_val: W_Object) -> N fname = w_field.name offset = w_field.offset generic_mem_write(vm, addr + offset, w_field.w_T, w_val.values_w[fname]) + elif isinstance(w_T, W_SimdType): + assert isinstance(w_val, W_Simd) + lane_size = sizeof(w_T.w_dtype) + for i in range(w_T.size): + generic_mem_write(vm, addr + i * lane_size, w_T.w_dtype, w_val.lanes_w[i]) else: t = w_T.fqn.human_name(vm) raise WIP(f"Cannot write memory of type `{t}`") diff --git a/spy/vm/modules/unsafe/misc.py b/spy/vm/modules/unsafe/misc.py index 2ac356a61..8daccc162 100644 --- a/spy/vm/modules/unsafe/misc.py +++ b/spy/vm/modules/unsafe/misc.py @@ -5,6 +5,7 @@ def sizeof(w_T: W_Type) -> int: from spy.vm.modules.posix import POSIX + from spy.vm.modules.simd import W_SimdType from spy.vm.modules.unsafe.ptr import W_PtrType from spy.vm.struct import W_StructType @@ -23,6 +24,8 @@ def sizeof(w_T: W_Type) -> int: # but for native it might be 8. Does it mean that we need to # preemptively choose the target platform BEFORE redshifting? return 4 + 4 # in debug mode we store both addr and length + elif isinstance(w_T, W_SimdType): + return sizeof(w_T.w_dtype) * w_T.size elif w_T is POSIX.w__FILE: return 4 # XXX else: diff --git a/spy/vm/primitive.py b/spy/vm/primitive.py index bf42cb5b2..a95913fca 100644 --- a/spy/vm/primitive.py +++ b/spy/vm/primitive.py @@ -123,6 +123,8 @@ def w_NEW(vm: "SPyVM", wam_cls: "W_MetaArg", *args_wam: "W_MetaArg") -> "W_OpSpe wam_arg = args_wam[0] if wam_arg.w_static_T == B.w_str: return W_OpSpec(OP.w_str_to_u32, [wam_arg]) + elif wam_arg.w_static_T == B.w_i32: + return W_OpSpec(OP.w_i32_to_u32, [wam_arg]) return W_OpSpec.NULL def __repr__(self) -> str: diff --git a/spy/vm/vm.py b/spy/vm/vm.py index 2d2d5d88d..ca31238b6 100644 --- a/spy/vm/vm.py +++ b/spy/vm/vm.py @@ -44,6 +44,7 @@ from spy.vm.modules.operator import OPERATOR, convop from spy.vm.modules.posix import POSIX from spy.vm.modules.rawbuffer import RAW_BUFFER +from spy.vm.modules.simd import SIMD from spy.vm.modules.time import TIME from spy.vm.modules.types import TYPES, W_Loc from spy.vm.modules.unsafe import UNSAFE @@ -166,6 +167,7 @@ def __init__( self.make_module(TYPES) self.make_module(MATH) self.make_module(UNSAFE) + self.make_module(SIMD) self.make_module(RAW_BUFFER) self.make_module(JSFFI) self.make_module(POSIX)