SIMD PR2 — arithmetic, comparison and select

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})"

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
 

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()

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])