Writing kernels

@sf.jit                                                # Triton style
def axpy(a, x_ptr, y_ptr, out_ptr, n, BLOCK: sf.constexpr):
    offs = sf.program_id(0) * BLOCK + sf.arange(0, BLOCK)
    mask = offs < n
    x = sf.load(x_ptr + offs, mask=mask)
    y = sf.load(y_ptr + offs, mask=mask)
    sf.store(out_ptr + offs, sf.fma(a, x, y), mask=mask)

dev = sf.device("atreides-sim")
x, y, out = dev.upload(xs), dev.upload(ys), dev.alloc(n)
r = axpy[(sf.cdiv(n, 256),)](sf.sf16(0.75), x, y, out, n, BLOCK=256)
r.status, r.cycles                                     # exact events, simulated cycles