import unittest, functools
from tinygrad import Tensor, UOp, GlobalCounters, Context, Device, Variable, TinyJit
from tinygrad.engine.realize import run_linear, get_runtime
from tinygrad.codegen import to_program
import numpy as np
from tinygrad.dtype import AddrSpace, dtypes, Invalid
from tinygrad.helpers import getenv
from tinygrad.schedule.rangeify import BufferizeOpts
from tinygrad.uop.ops import KernelInfo, AxisType, Ops, uopfunc
from tinygrad.codegen.opt import Opt, OptOps
from tinygrad.renderer.ptx import PTXRenderer
from tinygrad.renderer.cstyle import CStyleLanguage
from tinygrad.renderer.llvmir import LLVMRenderer
from test.helpers import assert_kernel_count
from test.null.test_custom_kernel import custom_elementwise_add_kernel, custom_elementwise_addmul_kernel, custom_gemm

# **** kernels ****

def custom_arange_kernel(C:UOp) -> UOp:
  i = UOp.range(C.shape[0], 0)
  return C[i].store(i.cast(C.dtype)).end(i).sink(arg=KernelInfo(name=f"custom_arange_{C.shape[0]}"))

def custom_eye_kernel(C:UOp) -> UOp:
  i = UOp.range(C.shape[0], 0)
  j = UOp.range(C.shape[1], 1)
  return C[i, j].store((i.eq(j)).cast(C.dtype)).end(i, j).sink(arg=KernelInfo(name=f"custom_eye_{C.numel()}"))

def custom_add_one_kernel(B:UOp, A:UOp) -> UOp:
  A,B = A.flatten(), B.flatten()
  assert B.numel() == A.numel()
  i = UOp.range(A.numel(), 0)
  return B[i].store(A[i] + 1).end(i).sink(arg=KernelInfo(name=f"add_one_{A.numel()}"))

def custom_ignore_first_kernel(C:UOp, A:UOp, B:UOp) -> UOp:
  # A is unused on purpose: the kernel takes call buffers 0 and 2, not 0, 1, 2
  C, B = C.flatten(), B.flatten()
  i = UOp.range(C.numel(), 0)
  return C[i].store(B[i] + 1).end(i).sink(arg=KernelInfo(name=f"ignore_first_{C.numel()}"))

def custom_add_var_kernel(*srcs:UOp, n_slot:int) -> UOp:
  C, *Bs = [s.flatten() for i,s in enumerate(srcs) if i != n_slot]
  i = UOp.range(C.numel(), 0)
  out = functools.reduce(lambda a,b: a+b, [B[i] for B in Bs]) + srcs[n_slot]
  return C[i].store(out).end(i).sink(arg=KernelInfo(name=f"add_var_{n_slot}_{len(Bs)}"))

def run_kernel_args(out:Tensor, name:str) -> list[str|None]:
  linear, var_vals = out.linear_with_vars()
  sink = [u for u in linear.toposort(enter_calls=True) if u.op is Ops.SINK and isinstance(u.arg, KernelInfo) and u.arg.name == name][0]
  sig = to_program(sink, Device[Device.DEFAULT].renderer).to_elf().signature
  run_linear(linear, var_vals)
  return [nm for nm, *_ in sig]

def custom_ignore_first_var_kernel(C:UOp, A:UOp, B:UOp, n:UOp) -> UOp:
  C, B = C.flatten(), B.flatten()
  k = UOp.variable("k", 0, 100, dtype=dtypes.int)
  i = UOp.range(C.numel(), 0)
  return C[i].store(B[i] + n + k).end(i).sink(arg=KernelInfo(name=f"ignore_first_var_{C.numel()}"))

def custom_sum(B:UOp, A:UOp) -> UOp:
  i = UOp.range(A.shape[0], 0, axis_type=AxisType.LOOP)
  B = B[0].set(0.0)
  B = B[0].set(B.after(i)[0] + A[i], end=i)
  return B.sink(arg=KernelInfo(name=f"custom_sum_{A.shape[0]}", opts_to_apply=()))

def flip_contract_kernel(dest:UOp, src:UOp):
  i = UOp.range(dest.shape[0], 0)
  j = UOp.range(dest.shape[1], 1, AxisType.UPCAST)
  vec = src[i, j].contract(j)
  store = UOp.group(*[dest[i, k].store(vec.index(3-k)) for k in range(4)])
  return store.end(i, j).sink(arg=KernelInfo(name=f"flip_contract_{dest.numel()}", opts_to_apply=()))

def slice_sum_kernel(dest:UOp, src:UOp):
  G = UOp.range(src.shape[0], 0, dtype=dtypes.int)
  slice_src = src[G, :]
  reg = UOp.placeholder((1,), dest.dtype, 0, addrspace=AddrSpace.REG)
  reg = reg.after(G)[0].set(0)
  R = UOp.range(src.shape[1], 1, AxisType.LOOP)
  reg = reg[0].set(reg.after(R)[0] + slice_src[R], end=R)
  ast = dest[G].set(reg[0], end=G)
  return ast.sink(arg=KernelInfo(name=f"slice_sum_{src.shape[0]}_{src.shape[1]}", opts_to_apply=()))

def simple_qkv_kernel(O:UOp, Q:UOp, K:UOp, V:UOp) -> UOp:
  # attention without softmax
  N, d = Q.shape[0], Q.shape[1]

  i = UOp.range(N, 0)  # output row
  d_out = UOp.range(d, 1)  # output column
  j = UOp.range(N, 2, axis_type=AxisType.LOOP)

  k_inner = UOp.range(d, 3, axis_type=AxisType.LOOP)
  qk_acc = UOp.placeholder((1,), Q.dtype, 0, addrspace=AddrSpace.REG)
  qk_acc = qk_acc.after(i, j)[0].set(0.0)
  qk_acc = qk_acc[0].set(qk_acc.after(k_inner)[0] + Q[i, k_inner] * K[j, k_inner], end=k_inner)
  qk_score = qk_acc[0] / (d ** 0.5)

  out_acc = UOp.placeholder((1,), Q.dtype, 1, addrspace=AddrSpace.REG)
  out_acc = out_acc.after(i, d_out)[0].set(0.0)
  out_acc = out_acc[0].set(out_acc.after(j)[0] + qk_score * V[j, d_out], end=j)

  store = O[i, d_out].store(out_acc[0])
  return store.end(d_out).end(i).sink(arg=KernelInfo(name=f"simple_qkv_{N}_{d}", opts_to_apply=()))

# **** backward callbacks ****

def backward_gemm(gradient:UOp, kernel:UOp) -> tuple[UOp, UOp]:
  out, a, b = kernel.src[1:]
  grad_a = (Tensor(gradient) @ Tensor(b).T).uop
  grad_b = (Tensor(a).T @ Tensor(gradient)).uop
  return (None, grad_a, grad_b)

def backward_gemm_custom(gradient:UOp, kernel:UOp) -> tuple[UOp, UOp]:
  out, a, b = kernel.src[1:]
  grad_a = Tensor.empty_like(Tensor(a)).custom_kernel(Tensor(gradient), Tensor(b).T, fxn=custom_gemm)[0].uop
  grad_b = Tensor.empty_like(Tensor(b)).custom_kernel(Tensor(a).T, Tensor(gradient), fxn=custom_gemm)[0].uop
  return (None, grad_a, grad_b)

# **** tests ****

class TestCustomKernel(unittest.TestCase):
  def test_empty(self):
    a = Tensor.empty(1)
    a = Tensor.custom_kernel(a, fxn=lambda _: UOp.sink(arg=KernelInfo()))[0]
    a.realize()

  def test_simple(self):
    a = Tensor.ones(16, 16).contiguous()
    b = Tensor.ones(16, 16).contiguous()
    c = Tensor.empty(16, 16)

    c = Tensor.custom_kernel(c,a,b, fxn=custom_elementwise_add_kernel)[0]

    out = c.flatten().tolist()
    assert all(x == 2 for x in out), "all 2"

  def test_duplicate_call_arg(self):
    x = Tensor.arange(4).clone().realize()
    x = Tensor.custom_kernel(x, x, fxn=custom_add_one_kernel)[0]
    # webgpu silently errors when a kernel has duplicate buffer args, so the list stays the same.
    # https://gpuweb.github.io/gpuweb/#abstract-opdef-encoder-bind-groups-alias-a-writable-resource
    self.assertEqual(x.tolist(), [1, 2, 3, 4] if Device.DEFAULT != "WEBGPU" else [0, 1, 2, 3])

  def test_simple_sharded(self):
    devs = ("CPU:0", "CPU:1")

    a = Tensor.ones(16, 16).contiguous().shard(devs, axis=0)
    b = Tensor.ones(16, 16).contiguous().shard(devs, axis=0)
    # ugly construction to get a sharded empty tensor
    c = Tensor(Tensor.empty(8, 16, device=devs).uop.unshard(0), device=devs)
    c = Tensor.custom_kernel(c,a,b, fxn=custom_elementwise_add_kernel)[0]
    out = c.flatten().tolist()
    assert all(x == 2 for x in out), "all 2"

  def test_sharded_add_one(self):
    # PYTHON backend explicitly checks for OOB access for wrong multi shape regression
    devs = ("PYTHON:0", "PYTHON:1")
    a = Tensor.ones(4, 4).contiguous().shard(devs, axis=0)
    c = Tensor(Tensor.empty(2, 4, device=devs).uop.unshard(0), device=devs)
    c = Tensor.custom_kernel(c, a, fxn=custom_add_one_kernel)[0]
    assert (c == 2).all().item()

  def test_multioutput(self):
    a = Tensor.full((16, 16), 3.).contiguous()
    b = Tensor.full((16, 16), 3.).contiguous()
    c = Tensor.empty(16, 16)
    d = Tensor.empty(16, 16)

    c,d = Tensor.custom_kernel(c,d,a,b, fxn=custom_elementwise_addmul_kernel)[:2]
    Tensor.realize(c,d)

    assert all(x == 6 for x in c.flatten().tolist()), "all 6"
    assert all(x == 9 for x in d.flatten().tolist()), "all 9"

  def test_arange(self):
    ref = Tensor.arange(100)
    tst = Tensor.empty_like(ref)
    tst = tst.custom_kernel(fxn=custom_arange_kernel)[0]
    self.assertTrue((ref == tst).all().item())

  def test_eye(self):
    ref = Tensor.eye(1024).clone().realize()
    tst = Tensor.empty_like(ref)
    tst = tst.custom_kernel(fxn=custom_eye_kernel)[0]
    self.assertTrue((ref == tst).all().item())

  @unittest.skip("contract shouldn't be supported here")
  def test_flip_contract(self):
    a = Tensor.randn(10,4)
    b = Tensor.empty_like(a)
    b = b.custom_kernel(a, fxn=flip_contract_kernel)[0]
    self.assertTrue((a.flip(1) == b).all().item())

  def test_noncontig(self):
    a = Tensor.ones(16, 16).contiguous()
    tst = Tensor.empty_like(a)
    b = a+1
    b_p1 = Tensor.custom_kernel(tst, b, fxn=custom_add_one_kernel)[0]
    self.assertTrue((b_p1 == 3).all().item())

  def test_unused_buffer_arg(self):
    a, b = Tensor([100.0, 200, 300, 400]), Tensor([1.0, 2, 3, 4])
    out = Tensor.custom_kernel(Tensor.empty(4), a, b, fxn=custom_ignore_first_kernel)[0]
    self.assertEqual(out.tolist(), [2, 3, 4, 5])

  def test_scalar_arg_any_position(self):
    for n_slot in range(3):
      with self.subTest(n_slot=n_slot):
        srcs = [Tensor.empty(4, dtype=dtypes.int), Tensor([1, 2, 3, 4], dtype=dtypes.int).realize()]
        srcs.insert(n_slot, Tensor(Variable("n", 0, 100, dtypes.int).bind(5)))
        out = Tensor.custom_kernel(*srcs, fxn=functools.partial(custom_add_var_kernel, n_slot=n_slot))[1 if n_slot == 0 else 0]
        self.assertEqual(run_kernel_args(out, f"add_var_{n_slot}_1"), ["n" if i == n_slot else None for i in range(3)])
        self.assertEqual(out.tolist(), [6, 7, 8, 9])

  @unittest.skipIf(Device.DEFAULT == "PYTHON", "needs a device that isnt the default")
  def test_scalar_arg_first_other_device(self):
    srcs = [Tensor(Variable("n", 0, 100, dtypes.int).bind(5), device="PYTHON"), Tensor.empty(4, dtype=dtypes.int, device="PYTHON"),
            Tensor([1, 2, 3, 4], dtype=dtypes.int, device="PYTHON").realize()]
    out = Tensor.custom_kernel(*srcs, fxn=functools.partial(custom_add_var_kernel, n_slot=0))[1]
    self.assertEqual(out.tolist(), [6, 7, 8, 9])

  def test_scalar_arg_with_other_variable(self):
    other = (Tensor([1, 2, 3, 4], dtype=dtypes.int) + Tensor(Variable("m", 0, 100, dtypes.int).bind(70))).contiguous()
    srcs = [Tensor.empty(4, dtype=dtypes.int), Tensor([1, 2, 3, 4], dtype=dtypes.int).realize(), Tensor(Variable("n", 0, 100, dtypes.int).bind(5))]
    out = Tensor.custom_kernel(*srcs, fxn=functools.partial(custom_add_var_kernel, n_slot=2))[0]
    Tensor.realize(other, out)
    self.assertEqual(out.tolist(), [6, 7, 8, 9])

  def test_scalar_arg_past_registers(self):
    # 9 inputs with the scalar between pointers, x86 passes the ones after the 6th (4th on windows) on the stack. every buffer differs
    srcs = [Tensor.empty(4, dtype=dtypes.int)] + [Tensor([k, 2*k, 3*k, 4*k], dtype=dtypes.int).realize() for k in range(1, 8)]
    srcs.insert(4, Tensor(Variable("n", 0, 100, dtypes.int).bind(5)))
    out = Tensor.custom_kernel(*srcs, fxn=functools.partial(custom_add_var_kernel, n_slot=4))[0]
    self.assertEqual(run_kernel_args(out, "add_var_4_7"), [None]*4 + ["n"] + [None]*4)
    self.assertEqual(out.tolist(), [33, 61, 89, 117])

  def test_unused_arg_then_var(self):
    srcs = (Tensor.empty(4, dtype=dtypes.int), Tensor.empty(4, dtype=dtypes.int), Tensor([1, 2, 3, 4], dtype=dtypes.int).realize(),
            Tensor(Variable("n", 0, 100, dtypes.int).bind(5)))
    out = Tensor.custom_kernel(*srcs, fxn=custom_ignore_first_var_kernel)[0]
    linear, var_vals = out.linear_with_vars()
    run_linear(linear, {**var_vals, "k": 10})
    self.assertEqual(out.tolist(), [16, 17, 18, 19])

  def test_scalar_arg_jit(self):
    # the scalar changes every call, a replay must not keep the value it was captured with
    kernel = functools.partial(custom_add_var_kernel, n_slot=0)
    jf = TinyJit(lambda B, n: Tensor.custom_kernel(Tensor(n), Tensor.empty(4, dtype=dtypes.int), B, fxn=kernel)[1].realize())
    B, n = Tensor([1, 2, 3, 4], dtype=dtypes.int).realize(), Variable("n", 0, 100, dtypes.int)
    for k in range(1, 6): self.assertEqual(jf(B, n.bind(k)).tolist(), [1+k, 2+k, 3+k, 4+k])

  def test_scalar_arg_jit_fixed_bind(self):
    # m is bound inside the jitted function so it's not an input, a replay keeps m=10 while the input n changes
    kernel = functools.partial(custom_add_var_kernel, n_slot=0)
    def f(B, n):
      Bm = Tensor.custom_kernel(Tensor(Variable("m", 0, 100, dtypes.int).bind(10)), Tensor.empty(4, dtype=dtypes.int), B, fxn=kernel)[1]
      return Tensor.custom_kernel(Tensor(n), Tensor.empty(4, dtype=dtypes.int), Bm, fxn=kernel)[1].realize()
    jf, B, n = TinyJit(f), Tensor([1, 2, 3, 4], dtype=dtypes.int).realize(), Variable("n", 0, 100, dtypes.int)
    for k in range(1, 6): self.assertEqual(jf(B, n.bind(k)).tolist(), [11+k, 12+k, 13+k, 14+k])

  def test_sum(self):
    a = Tensor([1.0, 2, 3, 4, 5])
    tst = Tensor.empty(1)
    b = Tensor.custom_kernel(tst, a, fxn=custom_sum)[0]
    self.assertEqual(b.item(), 15)

  def test_sum_outside(self):
    a = Tensor([1.0, 2, 3, 4, 5])+1
    tst = Tensor.empty(1)
    b = Tensor.custom_kernel(tst, a, fxn=custom_sum)[0]
    self.assertEqual(b.item(), 20)

  def test_sum_int(self):
    a = Tensor([1, 2, 3, 4, 5])
    tst = Tensor.empty(1, dtype=a.dtype)
    b = Tensor.custom_kernel(tst, a, fxn=custom_sum)[0]
    self.assertEqual(b.item(), 15)

  def test_slice_sum(self):
    A = Tensor.randn(16, 16).contiguous()
    B = Tensor.empty(16)
    B = Tensor.custom_kernel(B, A, fxn=slice_sum_kernel)[0]
    self.assertTrue(B.allclose(A.sum(1)).item())

  def test_gemm(self):
    N = 16
    a = Tensor.randn(N, N)
    b = Tensor.randn(N, N)
    c = Tensor.empty(N, N)

    tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
    self.assertTrue(tst.allclose(a@b, atol=1e-3).item())

  def test_unused_reduce_dependency(self):
    for value in (UOp.const(2), UOp.const((1, 2, 3, 4)).reshape((2, 2))):
      with self.subTest(shape=value.shape):
        def kernel(C:UOp) -> UOp:
          i = UOp.range(3, 0, AxisType.LOOP)
          cell = C.after(i)[0]
          v = value.after(cell.store(cell + 1))
          return C.after(v.reduce(i, arg=Ops.ADD)).sink(arg=KernelInfo(opts_to_apply=()))
        self.assertEqual(Tensor.zeros(1, dtype=dtypes.int).contiguous().custom_kernel(fxn=kernel)[0].tolist(), [3])

  def test_range_merge_keeps_after_domain(self):
    # C is read once per i or once per o, so it counts to 4, and merging o and i must not make that 16
    o, i = UOp.range(4, 0, AxisType.LOOP), UOp.range(4, 1, AxisType.LOOP)
    stores = {"after_i": lambda C,A: (c:=C.after(i)[0]).store(c+1),
              "after_o": lambda C,A: UOp.group((c:=C.after(o)[0]).store(c+1), A.after(o, i)[0].store(1)),
              "load_after_o": lambda C,A: C.after(o, i)[0].store(C.after(o)[0].load()+1)}
    for name, store in stores.items():
      with self.subTest(name):
        def kernel(C:UOp, A:UOp, store=store) -> UOp: return store(C, A).end(o, i).sink(arg=KernelInfo(opts_to_apply=()))
        self.assertEqual(Tensor.zeros(1).contiguous().custom_kernel(Tensor.empty(1), fxn=kernel)[0].item(), 4)

  def test_range_merge_keeps_index_domain(self):
    # C[c] is indexed by o alone, so each cell counts to 2 while the B store merges o and i
    o, i = UOp.range(4, 0, AxisType.LOOP), UOp.range(4, 1, AxisType.LOOP)
    for name, c in {"div": o//2, "cmp": (o<2).cast(dtypes.int)}.items():
      with self.subTest(name):
        def kernel(C:UOp, B:UOp, c=c) -> UOp:
          return UOp.group(C[c].store(C[c]+1), B[o*4+i].store(1)).end(o, i).sink(arg=KernelInfo(opts_to_apply=()))
        self.assertEqual(Tensor.custom_kernel(Tensor.zeros(2).contiguous(), Tensor.zeros(16).contiguous(), fxn=kernel)[0].tolist(), [2, 2])

  def test_range_merge_keeps_inner_reduce(self):
    # the inner reduce ends r1 alone, so r0 and r1 don't merge. the inner sum is 4, the outer adds it 4*4*4 times
    def kernel(C:UOp, A:UOp) -> UOp:
      r0, r1, r2 = [UOp.range(4, n, AxisType.UPCAST) for n in range(3)]
      return C[0].store(A[r2].reduce(r1, arg=Ops.ADD).reduce(r0, r1, r2, arg=Ops.ADD)).sink(arg=KernelInfo(opts_to_apply=()))
    self.assertEqual(Tensor.custom_kernel(Tensor.empty(1), Tensor.ones(4).contiguous(), fxn=kernel)[0].item(), 256)

  def test_range_merge_keeps_reduce_order(self):
    # the reduce body stores, so it loops a, k, b in order, and the digits of S are the A[a*2+b] it visits
    def kernel(O:UOp, S:UOp, A:UOp) -> UOp:
      a, k, b = [UOp.range(2, n, AxisType.LOOP) for n in range(3)]
      v = A[a*2+b].after((cell:=S.after(a, k, b)[0]).store(cell*10+A[a*2+b]))
      return O[0].store(v.reduce(a, k, b, arg=Ops.ADD)).sink(arg=KernelInfo(opts_to_apply=()))
    o, s = Tensor.custom_kernel(Tensor.empty(1, dtype=dtypes.int), Tensor.zeros(1, dtype=dtypes.int).contiguous(),
                                Tensor([1, 2, 3, 4], dtype=dtypes.int), fxn=kernel)[:2]
    Tensor.realize(o, s)
    self.assertEqual(s.item(), 12123434)

  def test_range_merge_keeps_loop_order(self):
    # loops nest by axis id whatever order END lists them, so x is outer and the digits of S are A[0], A[2], A[1], A[3]
    def kernel(S:UOp, A:UOp) -> UOp:
      x, y = UOp.range(2, 0, AxisType.LOOP), UOp.range(2, 1, AxisType.LOOP)
      return (cell:=S.after(x, y)[0]).store(cell*10+A[y*2+x]).end(y, x).sink(arg=KernelInfo(opts_to_apply=()))
    s = Tensor.custom_kernel(Tensor.zeros(1, dtype=dtypes.int).contiguous(), Tensor([1, 2, 3, 4], dtype=dtypes.int), fxn=kernel)[0]
    self.assertEqual(s.item(), 1324)

  def test_range_merge_then_shrink(self):
    # r0<2 shrinks the merged range of 8 to 4, which is the same UOp as r0, so it isn't shrunk again to 2
    def kernel(C:UOp) -> UOp:
      r0, r1 = UOp.range(4, 0, AxisType.LOOP), UOp.range(2, 1, AxisType.LOOP)
      return (cell:=C[(r0*2+r1).valid(r0<2)]).store(cell+1).end(r0, r1).sink(arg=KernelInfo(opts_to_apply=()))
    self.assertEqual(Tensor.zeros(8, dtype=dtypes.int).contiguous().custom_kernel(fxn=kernel)[0].tolist(), [1, 1, 1, 1, 0, 0, 0, 0])

  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "test requires locals")
  def test_group_reduce_split_range(self):
    # j%2 splits j into two ranges, both are still LOCAL
    def kernel(C:UOp, A:UOp) -> UOp:
      i, j = UOp.range(4, 0), UOp.range(8, 1, AxisType.LOCAL)
      return C[i].store((A[i, j] * (j%2).cast(A.dtype)).reduce(j, arg=Ops.ADD)).end(i).sink(arg=KernelInfo(opts_to_apply=()))
    a = Tensor.arange(32).reshape(4, 8).float().contiguous().realize()
    self.assertEqual(Tensor.custom_kernel(Tensor.empty(4), a, fxn=kernel)[0].tolist(), a[:, 1::2].sum(1).tolist())

  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "test requires locals")
  def test_nested_group_reduce(self):
    # the inner group's stage is indexed by the outer group's range, which is live at the inner reduce
    def kernel(C:UOp, B:UOp) -> UOp:
      i, g1, g2 = UOp.range(4, 0), UOp.range(4, 1, AxisType.LOCAL), UOp.range(8, 2, AxisType.LOCAL)
      return C[i].store(B[i, g1, g2].reduce(g2, arg=Ops.ADD).reduce(g1, arg=Ops.ADD)).end(i).sink(arg=KernelInfo(opts_to_apply=()))
    b = Tensor.arange(128).reshape(4, 4, 8).float().contiguous().realize()
    self.assertEqual(Tensor.custom_kernel(Tensor.empty(4), b, fxn=kernel)[0].tolist(), b.sum((1, 2)).tolist())

  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "test requires locals")
  def test_local_reduce(self):
    # a reduce over a thread range combines across threads
    a = Tensor.arange(32).reshape(4, 8).float().contiguous().realize()
    for at in (AxisType.LOCAL, AxisType.WARP):
      def kernel(C:UOp, A:UOp) -> UOp:
        i, j = UOp.range(4, 0), UOp.range(8, 1, at)
        return C[i].store(A[i, j].reduce(j, arg=Ops.ADD)).end(i).sink(arg=KernelInfo(opts_to_apply=()))
      self.assertEqual(Tensor.custom_kernel(Tensor.empty(4), a, fxn=kernel)[0].tolist(), a.sum(1).tolist())

  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_shared, "LOCAL STAGE needs shared memory")
  def test_stage_then_reduce(self):
    # the STAGE ends j, so the accumulator of the reduce over jj is initialized before the jj loop, not inside it
    def kernel(C:UOp, A:UOp) -> UOp:
      i, j, jj = UOp.range(4, 0), UOp.range(8, 1, AxisType.LOOP), UOp.range(8, 2, AxisType.LOOP)
      stage = (A[i, j] * 2).bufferize(j, arg=BufferizeOpts(None, AddrSpace.LOCAL))
      return C[i].store(stage.index(jj).reduce(jj, arg=Ops.ADD)).end(i).sink(arg=KernelInfo(opts_to_apply=()))
    a = Tensor.arange(32).reshape(4, 8).float().contiguous().realize()
    self.assertEqual(Tensor.custom_kernel(Tensor.empty(4), a, fxn=kernel)[0].tolist(), (a*2).sum(1).tolist())

  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "test requires locals")
  def test_stage_in_thread_range(self):
    # the STAGE is inside the thread range i, so every thread stages its own row
    def kernel(C:UOp, A:UOp) -> UOp:
      i, j, jj = UOp.range(4, 0, AxisType.LOCAL), UOp.range(8, 1, AxisType.LOOP), UOp.range(8, 2, AxisType.LOOP)
      stage = (A[i, j] * 2).bufferize(j, arg=BufferizeOpts(None, AddrSpace.LOCAL))
      return C[i].store(stage.index(jj).reduce(jj, arg=Ops.ADD)).end(i).sink(arg=KernelInfo(opts_to_apply=()))
    a = Tensor.arange(32).reshape(4, 8).float().contiguous().realize()
    self.assertEqual(Tensor.custom_kernel(Tensor.empty(4), a, fxn=kernel)[0].tolist(), (a*2).sum(1).tolist())

  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does not support dynamic register indexing")
  def test_reg_stage_then_reduce(self):
    # the REG buffer of the STAGE and the accumulator of the reduce are different buffers
    def kernel(C:UOp, A:UOp) -> UOp:
      i, j, jj = UOp.range(4, 0), UOp.range(8, 1, AxisType.LOOP), UOp.range(8, 2, AxisType.LOOP)
      stage = (A[i, j] * 2).bufferize(j, arg=BufferizeOpts(None, AddrSpace.REG))
      return C[i].store(stage.index(jj).reduce(jj, arg=Ops.ADD)).end(i).sink(arg=KernelInfo(opts_to_apply=()))
    a = Tensor.arange(32).reshape(4, 8).float().contiguous().realize()
    self.assertEqual(Tensor.custom_kernel(Tensor.empty(4), a, fxn=kernel)[0].tolist(), (a*2).sum(1).tolist())

  def test_reg_placeholder_then_reduce(self):
    # the accumulator of the reduce does not reuse the slot of a REG placeholder in the kernel
    def kernel(C:UOp, A:UOp) -> UOp:
      i, j = UOp.range(4, 0), UOp.range(8, 1)
      reg = UOp.placeholder((1,), dtypes.float, 0, addrspace=AddrSpace.REG)
      reg = reg.after(i)[0].set(A[i, 0])
      return C[i].store(A[i, j].reduce(j, arg=Ops.ADD) + reg[0]).end(i).sink(arg=KernelInfo(opts_to_apply=()))
    a = Tensor.arange(32).reshape(4, 8).float().contiguous().realize()
    self.assertEqual(Tensor.custom_kernel(Tensor.empty(4), a, fxn=kernel)[0].tolist(), (a.sum(1) + a[:, 0]).tolist())

  def test_split_range_id_free_of_loop(self):
    # the UPCAST range minted by the split gets a fresh id, the while loop's id 1 is taken
    def kernel(C:UOp, A:UOp) -> UOp:
      r, l = UOp.range(4, 0), UOp.loop(1)
      cnt = UOp.placeholder((1,), dtypes.int, slot=0, addrspace=AddrSpace.REG)
      cnt = cnt.after(r)[0].set(0)
      cnt = cnt.after(cnt[0].store(nxt:=cnt.after(l)[0] + 1).backedge(l, nxt < 3))
      return C[r].set(A[r] + cnt[0].cast(C.dtype), end=r).sink(arg=KernelInfo(opts_to_apply=(Opt(OptOps.SPLIT, 0, (2, AxisType.UPCAST)),)))
    a = Tensor([1., 2, 3, 4])
    self.assertEqual(Tensor.custom_kernel(Tensor.empty(4), a, fxn=kernel)[0].tolist(), [4., 5, 6, 7])

  def test_range_bound_after_outer_range(self):
    def kernel(C:UOp) -> UOp:
      i = UOp.range(3, 0, AxisType.LOOP)
      j = UOp.range(UOp.const(2, dtypes.weakint).after(i), 1, AxisType.LOOP)
      return C[0].set(C.after(i, j)[0] + 1, end=(j, i)).sink(arg=KernelInfo())
    self.assertEqual(Tensor.custom_kernel(Tensor.zeros(1, dtype=dtypes.int).contiguous(), fxn=kernel)[0].tolist(), [6])

  def test_gemm_multi(self):
    devs = ("CPU:0", "CPU:1")
    N = 16
    a = Tensor.randn(N, N).shard_(devs, axis=0)
    b = Tensor.randn(N, N).to(devs)
    c = Tensor(Tensor.empty(N//2, N, device=devs).uop.unshard(0), device=devs)
    tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
    self.assertTrue(tst.allclose(a@b, atol=1e-3).item())

  def test_gemm_backward_custom(self): self.test_gemm_backward(True)
  def test_gemm_backward(self, custom_backward_gemm=False):
    N = 4
    a_rand = Tensor.randn(N, 8)
    b_rand = Tensor.randn(8, N)
    Tensor.realize(a_rand, b_rand)

    a, b = Tensor(a_rand.numpy()), Tensor(b_rand.numpy())
    c = Tensor.empty(N, N)
    tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm, grad_fxn=backward_gemm_custom if custom_backward_gemm else backward_gemm)[0]
    tst.sum().backward()
    grad_a, grad_b = a.grad, b.grad
    Tensor.realize(tst, grad_a, grad_b)

    a, b = Tensor(a_rand.numpy()), Tensor(b_rand.numpy())
    ref = (a@b)
    ref.sum().backward()
    real_grad_a, real_grad_b = a.grad, b.grad
    Tensor.realize(ref, real_grad_a, real_grad_b)

    self.assertTrue(tst.allclose(ref, atol=1e-3).item())
    self.assertTrue(grad_a.allclose(real_grad_a, atol=1e-3).item())
    self.assertTrue(grad_b.allclose(real_grad_b, atol=1e-3).item())

  def test_simple_qkv(self):
    N, d = 8, 4
    Q = Tensor.randn(N, d)
    K = Tensor.randn(N, d)
    V = Tensor.randn(N, d)
    O = Tensor.empty(N, d)

    O_custom = Tensor.custom_kernel(O, Q, K, V, fxn=lambda o,q,k,v: simple_qkv_kernel(o,q,k,v))[0]
    O_ref = ((Q @ K.T) / (d ** 0.5)) @ V

    Tensor.realize(O_custom, O_ref)
    self.assertTrue(O_custom.allclose(O_ref, atol=1e-3).item())

  def test_simple_reshape(self):
    a = Tensor.ones(2,3,4).realize()
    b = Tensor.custom_kernel(Tensor.empty_like(a), a, fxn=custom_add_one_kernel)[0]
    b2 = b.reshape(2,12)
    c = Tensor.custom_kernel(Tensor.empty_like(b2), b2, fxn=custom_add_one_kernel)[0]
    GlobalCounters.reset()
    c.realize()
    assert all(i == 3. for i in c.flatten().tolist()), f"all 3 {c.tolist()}"
    assert_kernel_count(2)

  def test_multi_invalids_custom_kernel_no_copy(self):
    devs = ("CPU:0", "CPU:1")
    a = Tensor.ones(4, 4).shard(devs, axis=0).realize()
    c = Tensor(Tensor.invalids(2, 4, dtype=dtypes.float, device=devs).uop.unshard(0), device=devs)
    c = Tensor.custom_kernel(c, a, fxn=custom_add_one_kernel)[0]
    GlobalCounters.reset()
    c.realize()
    assert_kernel_count(len(devs))
    self.assertTrue((c == 2).all().item())

  def test_partial_invalid_store_keeps_uncovered_reads(self):
    x = Tensor([10., 20., 30., 40.]).realize()
    after = x.uop.after(x.uop.shrink(((0, 2),)).store(Invalid))
    self.assertEqual(Tensor(after).contiguous().tolist(), [10., 20., 30., 40.])

  def test_multi_after_invalid_store_dep_removed(self):
    x = Tensor.empty(4).uop
    self.assertEqual(Tensor(x.after(x.store(5), x.store(Invalid))).tolist(), [5]*4)

  def test_store_behind_value_after_kept(self):
    def kernel(A:UOp, loop:bool) -> UOp:
      x = UOp.const(1, dtypes.weakint).after(A[0].store(7))
      return (A[(r:=UOp.range(x, 0))+1].store(9).end(r) if loop else A[x].store(9)).sink(arg=KernelInfo(opts_to_apply=()))
    for loop in (False, True):
      out = Tensor.custom_kernel(Tensor.zeros(2, dtype=dtypes.int).contiguous().realize(), fxn=functools.partial(kernel, loop=loop))[0]
      self.assertEqual(out.tolist(), [7, 9])

  def test_store_value_behind_coalesced_store(self):
    def c(v): return UOp.const(v, dtypes.int)
    for stores, want in ((lambda A: (A[0].store(7), A[1].store(A.after(A[0].store(7))[0]+5)), [7, 12, 5, 0]),
                         (lambda A: (A[1].store(c(1).after(A[0].store(7)) + A[2].load()),), [7, 6, 5, 0]),
                         (lambda A: (A[1].store(c(1).after(A[2].store(2))), A[3].store(c(3).after(A[0].store(10)))), [10, 1, 2, 3])):
      out = Tensor([0, 0, 5, 0], dtype=dtypes.int).contiguous().custom_kernel(fxn=lambda A: UOp.sink(*stores(A), arg=KernelInfo(opts_to_apply=())))[0]
      self.assertEqual(out.tolist(), want)

  def test_deleted_store_keeps_deps(self):
    for f in (lambda a, s: s[1].store(s[1].load()), lambda a, s: s[UOp.invalid()].store(5),
              lambda a, s: s[1].store(UOp.invalid()), lambda a, s: a[UOp.invalid()].store(s[1].load())):
      out = Tensor.custom_kernel(Tensor.zeros(3, dtype=dtypes.int).contiguous(),
                                 fxn=lambda a: f(a, a.after(a[0].store(7))).sink(arg=KernelInfo(opts_to_apply=())))[0]
      self.assertEqual(out.tolist(), [7, 0, 0])

  def test_expand_view_invalid_assign_keeps_uncovered_reads(self):
    x = Tensor([[10., 11., 12., 13.], [20., 21., 22., 23.], [30., 31., 32., 33.], [40., 41., 42., 43.]]).realize()
    v = x[:1, :].expand(4, 4)
    v.assign(Tensor.invalids(4, 4, dtype=dtypes.float))
    self.assertEqual(v.contiguous().tolist(), [[10., 11., 12., 13.]]*4)

  @unittest.expectedFailure
  def test_gated_store_2d(self):
    # TODO: broken now, the valid is dropped. the valid on one index of the 2d C gates the whole store, only j == 0 writes C[i, 0]
    def kernel(C:UOp) -> UOp:
      i, j = UOp.range(4, 0), UOp.range(4, 1, AxisType.LOOP)
      return C[i.valid(j.eq(0)), 0].store((j+1).cast(C.dtype)).end(i, j).sink(arg=KernelInfo(opts_to_apply=()))
    self.assertEqual(Tensor.custom_kernel(Tensor.empty(4, 4), fxn=kernel)[0][:, 0].tolist(), [1.]*4)

  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_shared, "LOCAL buffer needs shared memory")
  @unittest.expectedFailure
  def test_gated_local_store_2d(self):
    # TODO: broken now, the valid is dropped. the valid on one index of the 2d LOCAL tmp gates the whole store, only j == 0 writes tmp[i, 0]
    def kernel(C:UOp) -> UOp:
      i, j = UOp.range(4, 0), UOp.range(4, 1, AxisType.LOOP)
      tmp = UOp.placeholder((4, 4), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL)
      st = tmp[i.valid(j.eq(0)), 0].store((j+1).cast(dtypes.float)).end(j)
      return C[i].store(tmp.after(st)[i, 0]).end(i).sink(arg=KernelInfo(opts_to_apply=()))
    self.assertEqual(Tensor.custom_kernel(Tensor.empty(4), fxn=kernel)[0].tolist(), [1.]*4)

  @unittest.expectedFailure
  def test_gated_load_2d(self):
    # TODO: broken now, the valid is dropped. the valid on one index of the 2d A gates the whole load, it is 0 where j != 0
    def kernel(C:UOp, A:UOp) -> UOp:
      i, j = UOp.range(4, 0), UOp.range(4, 1)
      return C[i, j].store(A[i.valid(j.eq(0)), 0]).end(i, j).sink(arg=KernelInfo(opts_to_apply=()))
    a = Tensor.arange(16).reshape(4, 4).float().contiguous().realize()
    self.assertEqual(Tensor.custom_kernel(Tensor.empty(4, 4), a, fxn=kernel)[0].tolist(), a[:, :1].pad(((0, 0), (0, 3))).tolist())

  def test_shrink_range_keeps_ungated_index(self):
    # q is ungated in C[q], so q keeps its 4 iterations after r shrinks to its gate r<2
    q, r = UOp.range(4, 0, AxisType.LOOP), UOp.range(4, 1, AxisType.LOOP)
    def kernel(C:UOp, D:UOp) -> UOp:
      return UOp.group(C[q].store(1), D[(q*4+r).valid((q<2)&(r<2))].store(1)).end(q, r).sink(arg=KernelInfo(opts_to_apply=()))
    c, d = Tensor.custom_kernel(Tensor.zeros(4, dtype=dtypes.int).contiguous(), Tensor.zeros(16, dtype=dtypes.int).contiguous(), fxn=kernel)[:2]
    Tensor.realize(c, d)
    self.assertEqual(c.tolist(), [1, 1, 1, 1])
    self.assertEqual(d.tolist(), [1, 1, 0, 0, 1, 1, 0, 0] + [0]*8)

  def test_shrink_range_keeps_ungated_store(self):
    # C[0] stores r every iteration, so r keeps its 4 iterations though D gates it to r<2
    r = UOp.range(4, 0, AxisType.LOOP)
    def kernel(C:UOp, D:UOp) -> UOp:
      return UOp.group(C[0].store(r.cast(dtypes.int)), D[r.valid(r<2)].store(1)).end(r).sink(arg=KernelInfo(opts_to_apply=()))
    c, d = Tensor.custom_kernel(Tensor.zeros(1, dtype=dtypes.int).contiguous(), Tensor.zeros(4, dtype=dtypes.int).contiguous(), fxn=kernel)[:2]
    Tensor.realize(c, d)
    self.assertEqual((c.tolist(), d.tolist()), ([3], [1, 1, 0, 0]))

  def test_shrink_range_keeps_store_gated_otherwise(self):
    # C[0] is gated by r>=2, not r<c, so r keeps its 4 iterations though D gates it to r<2
    r = UOp.range(4, 0, AxisType.LOOP)
    def kernel(C:UOp, D:UOp) -> UOp:
      return UOp.group(C[UOp.const(0).valid(r>=2)].store(r.cast(dtypes.int)), D[r.valid(r<2)].store(1)).end(r).sink(arg=KernelInfo(opts_to_apply=()))
    c, d = Tensor.custom_kernel(Tensor.zeros(1, dtype=dtypes.int).contiguous(), Tensor.zeros(4, dtype=dtypes.int).contiguous(), fxn=kernel)[:2]
    Tensor.realize(c, d)
    self.assertEqual((c.tolist(), d.tolist()), ([3], [1, 1, 0, 0]))

  def test_shrink_range_keeps_symbolic_end(self):
    # D gates r to r<4, but r only runs to n=2
    def kernel(D:UOp, N:UOp) -> UOp:
      r = UOp.range(N, 0, AxisType.LOOP)
      return D[r.valid(r<4)].store(1).end(r).sink(arg=KernelInfo(opts_to_apply=()))
    d = Tensor.custom_kernel(Tensor.zeros(8, dtype=dtypes.int).contiguous(), Tensor(Variable("n", 1, 8, dtypes.int).bind(2)), fxn=kernel)[0]
    self.assertEqual(d.tolist(), [1, 1, 0, 0, 0, 0, 0, 0])

  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "kernel timing not supported")
  def test_invalids_into_custom_kernel_with_beam(self):
    a = Tensor.full((4, 4), 3.).contiguous()
    b = Tensor.full((4, 4), 2.).contiguous()
    Tensor.realize(a, b)
    with Context(BEAM=1, IGNORE_BEAM_CACHE=1):
      out = Tensor.invalids(*a.shape, dtype=a.dtype)
      out, *_ = Tensor.custom_kernel(out, a, b, fxn=custom_elementwise_add_kernel)
      result = out.flatten().tolist()
    self.assertTrue(all(x == 5 for x in result), f"expected all 5.0, got {result}")

  @unittest.skip("what are anonymous buffers?")
  def test_anonymous_buffers_in_function(self):
    """Test that custom kernels with anonymous output buffers work inside @function."""
    a = Tensor.full((4, 4), 3.).contiguous()
    b = Tensor.full((4, 4), 2.).contiguous()
    Tensor.realize(a, b)

    def custom_add_with_tmp(o1:UOp, o2:UOp, A:UOp, B:UOp) -> UOp:
      o1,o2,A,B = o1.flatten(), o2.flatten(), A.flatten(), B.flatten()
      i = UOp.range(o1.numel(), 0)
      store_o1 = o1[i].store(A[i]+B[i])
      store_o2 = o2[i].store(A[i]+B[i]+2)
      return UOp.group(store_o1, store_o2).end(i).sink(arg=KernelInfo(name=f"add_with_tmp_{o1.numel()}")).simplify()

    from tinygrad import function
    @function(precompile=True)
    def run(x:Tensor, w:Tensor) -> Tensor:
      out = Tensor.invalids(*x.shape, dtype=x.dtype)
      tmp = Tensor.invalids(*x.shape, dtype=x.dtype)
      out, tmp = Tensor.custom_kernel(out, tmp, x, w, fxn=custom_add_with_tmp)[:2]
      return out+tmp

    result = run(a, b).flatten().tolist()
    expected = (3+2)*2+2
    assert all(x == expected for x in result), f"expected all {expected}, got {result}"

  def test_custom_kernel_sched(self, use_custom=False):
    x = Tensor.arange(32).reshape(8, 4).clone().realize()
    y = Tensor.empty_like(x)
    y = Tensor.custom_kernel(y, x, fxn=custom_add_one_kernel)[0]
    if use_custom:
      z = Tensor.empty_like(x)
      z = Tensor.custom_kernel(z, y.T.T, fxn=custom_add_one_kernel)[0]
    else: z = y.T.T+1
    GlobalCounters.reset()
    z.realize()
    assert_kernel_count(2)
    self.assertEqual(z.tolist(), x.add(2).tolist())

  def test_custom_kernel_sched_copy(self): self.test_custom_kernel_sched(use_custom=True)

  @unittest.skipIf(Device.DEFAULT == "CPU", "test needs to copy from CPU to another device")
  def test_call_to_program(self):
    t = Tensor([-1]).realize()
    # render a source string
    from tinygrad.codegen import do_to_program
    source:str = do_to_program(custom_add_one_kernel(t.uop.param_like(0), t.uop.param_like(1)), Device[t.device].renderer).src[2].arg
    assert isinstance(source, str)
    # call to PROGRAM
    def custom_source(out:UOp, inp:UOp) -> UOp:
      sink = UOp.sink(out.base, inp.base, arg=KernelInfo("add_one_1"))
      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(sink.toposort())), UOp(Ops.SOURCE, arg=source)))
    cpu_src = Tensor([2], device="CPU").realize()
    t = Tensor.custom_kernel(t, cpu_src.to(t.device), fxn=custom_source)[0]
    cp = t.to("CPU").realize()
    self.assertEqual(t.tolist(), [3])
    self.assertEqual(cp.tolist(), [3])

  def test_sliced_buffer_function(self):
    x = Tensor.arange(32).reshape(8, 4).clone().realize()
    from tinygrad import function
    @function(precompile=True)
    def run(x:Tensor) -> Tensor:
      y = Tensor.invalids(*x.shape, dtype=x.dtype)
      return Tensor.custom_kernel(y, x, fxn=custom_add_one_kernel)[0]
    GlobalCounters.reset()
    y = run(x[0]).realize()
    # backends that support contiguous views don't launch extra kernels
    assert_kernel_count(2 if x[0].uop.contiguous_view() is None else 1)
    self.assertEqual(y.tolist(), [1, 2, 3, 4])

  def test_inplace_transpose(self):
    def custom_assign_row_max_kernel(A:UOp) -> UOp:
      row = UOp.range(A.shape[0], 0)
      col = UOp.range(A.shape[1], 1)
      return A[row, col].store(A[row].max(axis=0)).end(col).end(row).sink(arg=KernelInfo(name=f"assign_row_max_{A.numel()}", opts_to_apply=()))
    a = Tensor.arange(4).reshape(2, 2).clone().realize()
    a = Tensor.custom_kernel(a.T, fxn=custom_assign_row_max_kernel)[0]
    with self.assertRaises(AssertionError):
      self.assertEqual(a.flatten().tolist(), [2, 2, 3, 3])
    self.assertEqual(a.shape, (2, 2))

  @unittest.skipUnless((isinstance(Device[Device.DEFAULT].renderer, (CStyleLanguage, LLVMRenderer)) or Device.DEFAULT == "PYTHON") and
                       Device.DEFAULT != "WEBGPU", "binary not supported on this backend")
  def test_binary(self):
    payload = bytes(range(256))
    def view(x): return x.reshape((16, 16)).permute((1, 0)).reshape((256,))
    def kernel(out:UOp):
      i = UOp.range(len(payload), 0)
      data = view(UOp(Ops.BINARY, arg=payload))
      return out[i].store(data[i]).end(i).sink(arg=KernelInfo(name="binary", opts_to_apply=()))
    self.assertEqual(Tensor.empty(len(payload), dtype=dtypes.uint8).custom_kernel(fxn=kernel)[0].tolist(),
                     view(Tensor(payload, dtype=dtypes.uint8)).tolist())

  def test_one_lane_stack_runtime_index(self):
    def kernel(A:UOp): return A[0].store(UOp.stack(UOp.const(42, dtypes.int)).index(A[1].load())).sink(arg=KernelInfo(opts_to_apply=()))
    self.assertEqual(Tensor.zeros(2, dtype=dtypes.int).contiguous().custom_kernel(fxn=kernel)[0].tolist(), [42, 0])

@unittest.skipUnless((Device.DEFAULT == "CPU" and isinstance(Device[Device.DEFAULT].renderer, (CStyleLanguage, LLVMRenderer)))
                     or Device.DEFAULT == "AMD" and isinstance(Device[Device.DEFAULT].renderer, CStyleLanguage),
                     "calls in kernel not implemented in this renderer")
class TestCallInKernel(unittest.TestCase):
  def test_simple(self):
    N = getenv("N", 4)

    @uopfunc
    def add(out:UOp, x:UOp) -> UOp:
      i = UOp.range(x.numel(), 0)
      return out[i].store(x[i] + 1).end(i).sink()

    @uopfunc
    def sum_fn(out:UOp, x:UOp) -> UOp:
      i = UOp.range(x.numel(), 0)
      return out[0].store(x[i].reduce(i, arg=Ops.ADD)).sink()

    def main(out:UOp, x:UOp) -> UOp:
      tmp0, tmp1 = (UOp.alloc_like(x, addrspace=AddrSpace.REG) for _ in range(2))
      call0 = add(tmp0, x)
      call1 = add(tmp1, tmp0.after(call0))
      return sum_fn(out, tmp1.after(call1)).sink(arg=KernelInfo(name="test_simple"))

    x = Tensor.arange(N).clone().realize()
    out = Tensor.custom_kernel(Tensor.empty_like(x), x, fxn=lambda out,x: add(out, x).sink(arg=KernelInfo(name="add_kernel")))[0]
    out = Tensor.custom_kernel(Tensor.empty(1, dtype=x.dtype), out,
                               fxn=lambda out,x: sum_fn(out, x).sink(arg=KernelInfo(name="sum_kernel")))[0]
    self.assertEqual(out.tolist(), [N*(N+1)//2])
    out2 = Tensor.custom_kernel(Tensor.empty(1, dtype=x.dtype), x, fxn=main)[0]
    self.assertEqual(out2.tolist(), [N*(N+3)//2])

  def test_reduce_over_call_keeps_loop_order(self):
    # a call can hide stores, so the reduce loops a, k, b in order, and the digits of S are the A[a*2+b] it visits
    @uopfunc
    def update(S:UOp, A:UOp) -> UOp: return S[0].store(S[0]*10+A[0]).sink()
    def kernel(O:UOp, S:UOp, A:UOp) -> UOp:
      a, k, b = [UOp.range(2, n, AxisType.LOOP) for n in range(3)]
      return O[0].store(A[a*2+b].after(update(S.after(k), A[a*2+b])).reduce(a, k, b, arg=Ops.ADD)).sink(arg=KernelInfo(opts_to_apply=()))
    o, s = Tensor.custom_kernel(Tensor.empty(1, dtype=dtypes.int), Tensor.zeros(1, dtype=dtypes.int).contiguous(),
                                Tensor([1, 2, 3, 4], dtype=dtypes.int), fxn=kernel)[:2]
    Tensor.realize(o, s)
    self.assertEqual(s.item(), 12123434)

  def test_nested_call(self):
    @uopfunc
    def incr(out:UOp, A:UOp):
      i = UOp.range(A.shape[0], 0)
      return out[i].store(A[i]+1).end(i).sink()

    @uopfunc
    def square(out:UOp, A:UOp):
      tmp0, tmp1 = (UOp.alloc_like(A, addrspace=AddrSpace.REG) for _ in range(2))
      call0 = incr(tmp0, A)
      call1 = incr(tmp1, A.after(call0))
      assert call0.body is call1.body
      i = UOp.range(A.shape[0], 1)
      return out[i].store(tmp1.after(call1)[i]**2).end(i).sink()

    def kernel(C:UOp, A:UOp):
      tmp = UOp.alloc_like(A, addrspace=AddrSpace.REG)
      call = square(tmp, A)
      i = UOp.range(A.shape[0], 2)
      return C[i].store(tmp.after(call)[i]+10).end(i).sink(arg=KernelInfo(name="nested_calls"))

    a = Tensor.arange(4).clone().realize()
    out = Tensor.custom_kernel(Tensor.empty_like(a), a, fxn=kernel)[0]
    self.assertEqual(out.tolist(), [11, 14, 19, 26])

  def test_call_loop_mini(self, opts=()):
    @uopfunc
    def mul(p:UOp, q:UOp): return p[0].store(q[0]*3).sink()

    def kernel(C:UOp, A:UOp):
      i = UOp.range(4, 0)
      call = mul(C[i], A[i])
      return call.end(i).sink(arg=KernelInfo(name="call_loop_mini", opts_to_apply=opts))
    a = Tensor([1, 2, 3, 4], dtype=dtypes.int).realize()
    out = Tensor.custom_kernel(Tensor.zeros(4, dtype=dtypes.int).clone().realize(), a, fxn=kernel)[0]
    self.assertEqual(out.tolist(), [3, 6, 9, 12])

  def test_call_scalar_arg(self):
    @uopfunc
    def add(out:UOp, src:UOp, val:UOp) -> UOp:
      return out[0].store(src[0]+val).sink()

    def kernel(out:UOp, src:UOp) -> UOp:
      i = UOp.range(4, 0)
      return add(out[i], src[i], i).end(i).sink(arg=KernelInfo(name="call_scalar_arg", opts_to_apply=()))

    x = Tensor([1, 2, 3, 4], dtype=dtypes.int).realize()
    self.assertEqual(Tensor.empty_like(x).custom_kernel(x, fxn=kernel)[0].tolist(), [1, 3, 5, 7])

  def test_call_with_gated_store(self):
    @uopfunc
    def gated_store(out:UOp, idx:UOp):
      i = idx[0]
      return out[i.valid((i >= 0) & (i < out.shape[0]))].store(7).sink()

    def kernel(C:UOp, A:UOp):
      tmp = UOp.alloc_like(C, addrspace=AddrSpace.REG)
      i = UOp.range(C.shape[0], 0)
      init = tmp[i].store(0).end(i)
      call = gated_store(tmp.after(init), A)
      j = UOp.range(C.shape[0], 1)
      return C[j].store(tmp.after(call)[j] + 1).end(j).sink(arg=KernelInfo(name="call_with_gated_store", opts_to_apply=()))

    a = Tensor([2], dtype=dtypes.int).realize()
    out = Tensor.custom_kernel(Tensor.zeros(4, dtype=dtypes.int).clone().realize(), a, fxn=kernel)[0]
    self.assertEqual(out.tolist(), [1, 1, 8, 1])

  def test_call_body_range_is_not_ours(self):
    @uopfunc
    def mul(out:UOp, A:UOp):
      k = UOp.range(4, 0, AxisType.LOOP)
      return out[k].store(A[k]*3).end(k).sink()

    def kernel(C:UOp, A:UOp):
      m, i, call = UOp.range(3, 1, AxisType.LOOP), UOp.range(4, 0, AxisType.LOOP), mul(C, A)
      self.assertIn(i, call.body.toposort())
      end_m = C.after(call)[m].store(A[m]).end(m)
      return C.after(end_m)[i].store(A[i]+1).end(i).sink(arg=KernelInfo(name="call_body_range", opts_to_apply=()))

    a = Tensor([1, 2, 3, 4], dtype=dtypes.int).realize()
    out = Tensor.custom_kernel(Tensor.zeros(4, dtype=dtypes.int).clone().realize(), a, fxn=kernel)[0]
    self.assertEqual(out.tolist(), [2, 3, 4, 5])

  @unittest.expectedFailure
  def test_call_loop_mini_opts(self): self.test_call_loop_mini(opts=None)

  @unittest.expectedFailure
  def test_call_loop_split(self): self.test_call_loop_mini((Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST)),))

  @unittest.expectedFailure
  def test_call_loop_pad(self):
    @uopfunc
    def add(p:UOp, q:UOp): return p[0].store(p[0]+q[0]).sink()

    def kernel(C:UOp, A:UOp):
      i = UOp.range(3, 0)
      call = add(C[i], A[i])
      return call.end(i).sink(arg=KernelInfo(name="call_loop_padto", opts_to_apply=(Opt(OptOps.PADTO, 0, 4),)))

    a = Tensor([1, 2, 3, 4], dtype=dtypes.int).realize()
    out = Tensor.custom_kernel(Tensor.zeros(4, dtype=dtypes.int).clone().realize(), a, fxn=kernel)[0]
    self.assertEqual(out.tolist(), [1, 2, 3, 0])

  def test_uopfunc(self):
    @uopfunc
    def axpy(out:UOp, x:UOp, a:int) -> UOp: # a is a trace-time constant: two bodies, two C functions
      i = UOp.range(x.shape[0], 0, AxisType.LOOP)
      return out[i].store(out[i] + x[i] * a).end(i).sink()

    @uopfunc
    def triple(out:UOp, x:UOp): return axpy(out, x, 3).sink()

    def kernel(C:UOp, A:UOp) -> UOp:
      C = C.after(triple(C.after(axpy(C, A, 2)), A))
      i = UOp.range(A.shape[0], 1, AxisType.LOOP) # a loop after the calls
      return C[i].store(C[i] + 1).end(i).sink(arg=KernelInfo(name="uopfunc", opts_to_apply=()))

    a = Tensor([1, 2, 3, 4], dtype=dtypes.int).realize()
    out = Tensor.custom_kernel(Tensor.zeros(4, dtype=dtypes.int).contiguous().realize(), a, fxn=kernel)[0]
    self.assertEqual(out.tolist(), [6, 11, 16, 21])

  def test_copy_offset_slices(self):
    @uopfunc
    def copy_block(dst:UOp, src:UOp) -> UOp:
      i = UOp.range(4, 0, AxisType.UPCAST)
      return dst[i].store(src[i]).end(i).sink()

    def kernel(dst:UOp, src:UOp) -> UOp:
      return copy_block(dst[1:5], src[1:5]).sink(arg=KernelInfo(opts_to_apply=()))

    src = Tensor([0., 1., 2., 3., 4.]).realize()
    dst = Tensor([-1.] * 6).realize()
    out = dst.custom_kernel(src, fxn=kernel)[0]
    self.assertEqual(out.tolist(), [-1., 1., 2., 3., 4., -1.])

  @unittest.skipIf(Device.DEFAULT == "CPU", "test needs to copy from CPU to another device")
  def test_call_to_copy(self):
    @uopfunc
    def copy(out:UOp, A:UOp):
      i = UOp.range(out.shape[0], 0, AxisType.LOOP)
      return out[i].store(A[i]).end(i).sink()

    def forward(out1:UOp, out2:UOp, A:UOp):
      call = copy(out1, A)
      i = UOp.range(out2.shape[0], 0, AxisType.LOOP)
      inline = out2[i].store(A[i]).end(i).sink()
      return UOp.group(call, inline).sink(arg=KernelInfo(name="forward", opts_to_apply=()))

    x = Tensor([1, 2, 3, 4], device="CPU").realize().to(Device.DEFAULT)
    y, z = x.empty_like(), x.empty_like()
    y, z, _ = Tensor.custom_kernel(y, z, x, fxn=forward)
    y, z = y.to("CPU"), z.to("CPU")
    Tensor.realize(y, z)

    self.assertEqual(y.tolist(), z.tolist())
    self.assertEqual(x.tolist(), z.tolist())

  @unittest.skipIf(Device[Device.DEFAULT].renderer.inline_prefix is not None, "backend inlines calls instead of linking")
  def test_linker(self):
    @uopfunc
    def plus2(out:UOp, A:UOp): return out[0].store(A[0]+2).sink()

    @uopfunc
    def copy(out:UOp, A:UOp): return out[0].store(A[0]).sink()

    def kernel(A:UOp, extra=False):
      A = A.flatten()
      i = UOp.range(A.shape[0], 0)
      tmp = UOp.alloc_like(A, addrspace=AddrSpace.REG)
      if extra: tmp = tmp.after(copy(tmp[i], A[i]))
      return A[i].store(tmp[i].after(plus2(tmp[i], A[i]))).end(i).sink(arg=KernelInfo(name="test_kernel", opts_to_apply=()))

    x = Tensor.custom_kernel(Tensor.arange(4).clone().realize(), fxn=kernel)[0]
    run_linear(lin:=x.schedule_linear())

    prg = to_program(lin.src[0].body, Device[Device.DEFAULT].renderer)
    assert "+2" not in prg.src[2].arg, f"found callee body in caller kernel {prg.src[2].arg}"
    call = next(u for u in prg.src[1].src if u.op is Ops.CALL)
    first_binary = call.body.src[3]
    call_runtime = get_runtime(Device.DEFAULT, call.body)
    self.assertListEqual(x.tolist(), Tensor.arange(4).add(2).tolist())

    # launch a different caller with the same callee
    x = Tensor.custom_kernel(Tensor.arange(5).clone().realize(), fxn=lambda A: kernel(A, extra=True))[0]
    run_linear(lin:=x.schedule_linear())
    prg2 = to_program(lin.src[0].body, Device[Device.DEFAULT].renderer).src[1]
    self.assertIs(call.body, next(u.body for u in prg2.src if u.op is Ops.CALL and u.arg.name == "plus2"))
    self.assertIs(first_binary, next(u.body.src[3] for u in prg2.src if u.op is Ops.CALL and u.arg.name == "plus2"))
    self.assertIs(call_runtime, get_runtime(Device.DEFAULT, call.body))
    self.assertListEqual(x.tolist(), Tensor.arange(5).add(2).tolist())

class TestCustomKernelInput(unittest.TestCase):
  def _test_mop(self, mop_fxn, max_kernels):
    # default: input is BUFFER
    x = mop_fxn(Tensor.arange(32).clone().realize())
    y = Tensor.custom_kernel(Tensor.empty_like(x), x, fxn=custom_add_one_kernel)[0]
    GlobalCounters.reset()
    y.realize()
    assert_kernel_count(max_kernels)
    self.assertEqual(y.tolist(), x.add(1).tolist())
    # same test with @function, input is PARAM
    from tinygrad import function
    x0 = Tensor.arange(32).clone().realize()
    @function(precompile=True)
    def run(a:Tensor) -> Tensor:
      xv = mop_fxn(a)
      y = Tensor.invalids(*xv.shape, dtype=xv.dtype, device=a.device)
      return Tensor.custom_kernel(y, xv, fxn=custom_add_one_kernel)[0]
    GlobalCounters.reset()
    y = run(x0).realize()
    assert_kernel_count(max_kernels)
    self.assertEqual(y.tolist(), mop_fxn(x0).add(1).tolist())

  def test_reshape(self): self._test_mop(lambda x: x.reshape(16, 2), max_kernels=1)
  def test_permute(self): self._test_mop(lambda x: x.reshape(4, 8).T, max_kernels=2)
  def test_double_permute(self): self._test_mop(lambda x: x.reshape(4, 8).T.T, max_kernels=1)
  def test_shrink(self): self._test_mop(lambda x: x[:4], max_kernels=1)
  def test_pad(self): self._test_mop(lambda x: x[:4].pad(((0, 4),)), max_kernels=2)
  def test_flip(self): self._test_mop(lambda x: x.flip(0), max_kernels=2)
  def test_offset_shrink(self): self._test_mop(lambda x: x[4:8], max_kernels=2)
  def test_2d_shrink(self): self._test_mop(lambda x: x.reshape(4, 8)[:, 2:6], max_kernels=2)
  def test_expand(self): self._test_mop(lambda x: x.reshape(16, 2)[:, :1].expand(16, 2), max_kernels=2)

class TestUnshardIndex(unittest.TestCase):
  """Regression tests for INDEX on UNSHARD (fragment) resolution in schedule/multi.py.

  A fragment is a per-thread REG buffer wrapped in UNSHARD over LOCAL thread ranges.
  index_multi must resolve an INDEX on the UNSHARD view into an INDEX on the per-thread
  shard. Two ownership patterns must work:
    contiguous: idx = rng*shard_sz + local   (thread rng owns [rng*shard_sz, ...))
    strided:    idx = rng + ir*shard_sz      (thread rng owns {rng, rng+shard_sz, ...})
  """
  def _run(self, kernel, shape=(8, 8)):
    c = Tensor.empty(*shape)
    out = Tensor.custom_kernel(c, fxn=kernel)[0]
    try: return out.numpy()
    except RuntimeError as e:
      if isinstance(Device[Device.DEFAULT].renderer, PTXRenderer) and "dynamic register indexing" in str(e):
        self.skipTest("PTX does not support dynamic register indexing")
      raise

  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
  def test_contiguous_fragment_index(self):
    # thread ty owns rows [ty*8, ty*8+8) of a 64-row fragment -- contiguous ownership.
    # This is the pre-existing case that index_multi always handled.
    def kernel(C:UOp) -> UOp:
      ty = UOp.range(8, 0, AxisType.LOCAL)
      ir = UOp.range(8, 1, AxisType.LOOP)
      j = UOp.range(8, 2, AxisType.LOOP)
      # 8x8 fragment, 8 threads -> 64x8 full tile. thread ty owns rows [ty*8, ty*8+8).
      frag = UOp.placeholder((8, 8), dtypes.float32, 0, AddrSpace.REG).unshard((0,), (ty,))
      return C[ty*8 + ir, j].store(frag[ty*8 + ir, j]).end(j, ir, ty).sink(arg=KernelInfo(name="contig_frag"))
    out = self._run(kernel, (64, 8))
    assert out.shape == (64, 8)

  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
  def test_strided_fragment_index(self):
    # thread ty owns rows {ty, ty+8, ty+16, ty+24, ..., ty+56} of a 64-row fragment --
    # strided ownership. idx = ty + ir*8 where shard_sz=8 (8 threads, shard rows=8).
    # The contiguous check (idx - rng*shard_sz) fails; the strided check
    # (idx-rng) % shard_sz == 0 must succeed. This is the pattern the index_multi fix adds.
    def kernel(C:UOp) -> UOp:
      ty = UOp.range(8, 0, AxisType.LOCAL)
      ir = UOp.range(8, 1, AxisType.LOOP)
      j = UOp.range(8, 2, AxisType.LOOP)
      # 8x8 fragment, 8 threads -> 64x8 full tile. thread ty owns rows {ty, ty+8, ..., ty+56}.
      frag = UOp.placeholder((8, 8), dtypes.float32, 0, AddrSpace.REG).unshard((0,), (ty,))
      return C[ty + ir*8, j].store(frag[ty + ir*8, j]).end(j, ir, ty).sink(arg=KernelInfo(name="strided_frag"))
    out = self._run(kernel, (64, 8))
    assert out.shape == (64, 8)

  def test_fragment_index_cannot_shard(self):
    # thread ty indexing rows [ty, ty+8) overlaps with other threads' rows -- this matches neither
    # the contiguous nor the strided ownership pattern, so index_multi must raise.
    def kernel(C:UOp) -> UOp:
      ty = UOp.range(8, 0, AxisType.LOCAL)
      ir = UOp.range(8, 1, AxisType.LOOP)
      j = UOp.range(8, 2, AxisType.LOOP)
      frag = UOp.placeholder((8, 8), dtypes.float32, 0, AddrSpace.REG).unshard((0,), (ty,))
      return C[ty + ir, j].store(frag[ty + ir, j]).end(j, ir, ty).sink(arg=KernelInfo(name="bad_frag"))
    with self.assertRaisesRegex(RuntimeError, "cannot shard index"):
      self._run(kernel, (64, 8))

def _run_fragment_kernel(testcase, kernel, out_shape, inputs=()):
  c = Tensor.empty(*out_shape)
  out = Tensor.custom_kernel(c, *inputs, fxn=kernel)[0]
  try: return out.numpy()
  except RuntimeError as e:
    if isinstance(Device[Device.DEFAULT].renderer, PTXRenderer) and "dynamic register indexing" in str(e):
      testcase.skipTest("PTX does not support dynamic register indexing")
    raise

class TestUnshardAlu(unittest.TestCase):
  """Tests for ALU on (fragment) UNSHARD values in schedule/multi.py's alu_multi.

  An ALU with UNSHARD srcs lowers to per-shard ops when every src is one of:
    same sharding:  peel the UNSHARD, keep the layout
    scalar:         broadcast to every shard
    whole unsharded same-shape value: takes its per-shard sub-view (shard_subview)
  """
  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
  def test_alu_scalar_broadcast(self):
    # scalar srcs broadcast to every shard: frag*2.0 where frag is 1.5 per thread -> 3.0 everywhere
    def kernel(C:UOp) -> UOp:
      ty = UOp.range(8, 0, AxisType.LOCAL)
      # 8 values per thread, 8 threads -> 64-value full view
      frag = UOp.placeholder((8,), dtypes.float32, 0, AddrSpace.LOCAL).unshard((0,), (ty,))
      v = frag.after(frag.store(1.5)) * 2.0
      return C.store(v).end(ty).sink(arg=KernelInfo(name="alu_scalar", opts_to_apply=()))
    out = _run_fragment_kernel(self, kernel, (64,))
    np.testing.assert_allclose(out, 3.0)

  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
  def test_alu_whole_value_subview(self):
    # UNSHARD + whole unsharded same-shape value: each shard adds its own sub-view of A.
    def kernel(C:UOp, A:UOp) -> UOp:
      ty = UOp.range(8, 0, AxisType.LOCAL)
      frag = UOp.placeholder((8,), dtypes.float32, 0, AddrSpace.LOCAL).unshard((0,), (ty,))
      v = frag.after(frag.store(0.0)) + A
      return C.store(v).end(ty).sink(arg=KernelInfo(name="alu_subview", opts_to_apply=()))
    a = Tensor(np.arange(64, dtype=np.float32))
    out = _run_fragment_kernel(self, kernel, (64,), inputs=(a,))
    np.testing.assert_allclose(out, a.numpy(), atol=1e-4)

class TestUnshardStore(unittest.TestCase):
  """Tests for STORE of a sharded value into an unsharded dest (store_value_multi in schedule/multi.py).

  Every shard stores its value into its own contiguous sub-view of the dest, one SHRINK per sharded axis.
  """
  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
  def test_store_unshard_value(self):
    # single-axis: 8 threads each own 8 values of the 64-value output tile
    def kernel(C:UOp) -> UOp:
      ty = UOp.range(8, 0, AxisType.LOCAL)
      frag = UOp.placeholder((8,), dtypes.float32, 0, AddrSpace.LOCAL).unshard((0,), (ty,))
      v = frag.after(frag.store(0.0)) + 2.5
      return C.store(v).end(ty).sink(arg=KernelInfo(name="store_unshard", opts_to_apply=()))
    out = _run_fragment_kernel(self, kernel, (64,))
    np.testing.assert_allclose(out, 2.5)

  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
  def test_store_unshard_value_2axis(self):
    # two sharded axes (the gemm fragment layout): thread (ty, tx) owns the (2, 1, 1, 2) sub-view of the
    # (2, 4, 2, 2) output tile; the store must SHRINK dest on both sharded axes
    def kernel(C:UOp, A:UOp) -> UOp:
      ty = UOp.range(4, 0, AxisType.LOCAL)
      tx = UOp.range(2, 1, AxisType.LOCAL)
      frag = UOp.placeholder((2, 1, 1, 2), dtypes.float32, 0, AddrSpace.REG).unshard((1, 2), (ty, tx))
      v = frag.after(frag.store(0.0)) + A
      return C.store(v).end(tx, ty).sink(arg=KernelInfo(name="store_unshard_2axis", opts_to_apply=()))
    a = Tensor(np.arange(32, dtype=np.float32).reshape(2, 4, 2, 2))
    out = _run_fragment_kernel(self, kernel, (2, 4, 2, 2), inputs=(a,))
    np.testing.assert_allclose(out, a.numpy(), atol=1e-4)

  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
  def test_store_unshard_value_2axis_reshape(self):
    # the reshape keeps both sharded axes, each with its own shard count (4 and 2)
    def kernel(C:UOp, A:UOp) -> UOp:
      ty = UOp.range(4, 0, AxisType.LOCAL)
      tx = UOp.range(2, 1, AxisType.LOCAL)
      frag = UOp.placeholder((2, 1, 1, 2), dtypes.float32, 0, AddrSpace.REG).unshard((1, 2), (ty, tx))
      v = (frag.after(frag.store(0.0)) + A).reshape(2, 4, 2, 2, 1)
      return C.store(v).end(tx, ty).sink(arg=KernelInfo(name="store_unshard_2axis_reshape", opts_to_apply=()))
    a = Tensor(np.arange(32, dtype=np.float32).reshape(2, 4, 2, 2))
    out = _run_fragment_kernel(self, kernel, (2, 4, 2, 2, 1), inputs=(a,))
    np.testing.assert_allclose(out, a.numpy().reshape(2, 4, 2, 2, 1), atol=1e-4)

  def _test_store_load_fragment(self, addrspace:AddrSpace):
    # thread ty stores A[ty*8:ty*8+8] into its fragment, then reads it back into the same slice of C
    def kernel(C:UOp, A:UOp) -> UOp:
      ty = UOp.range(8, 0, AxisType.LOCAL)
      frag = UOp.placeholder((8,), dtypes.float32, 0, addrspace).unshard((0,), (ty,))
      return C.store(frag.after(frag.store(A))).end(ty).sink(arg=KernelInfo(name="store_load_fragment", opts_to_apply=()))
    a = Tensor.arange(64, dtype=dtypes.float32)
    out = _run_fragment_kernel(self, kernel, (64,), inputs=(a,))
    np.testing.assert_equal(out, a.numpy())

  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
  def test_store_load_reg_fragment(self): self._test_store_load_fragment(AddrSpace.REG)

  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
  @unittest.expectedFailure  # TODO: should not fail silently
  def test_store_load_local_fragment(self): self._test_store_load_fragment(AddrSpace.LOCAL)

class TestUOpReduce(unittest.TestCase):
  def test_uop_sum(self):
    a = Tensor([1.0, 2, 3, 4, 5])
    self.assertAlmostEqual(Tensor(a.uop.sum(axis=0)).item(), 15.0)

  def test_uop_sum_2d(self):
    a = Tensor.arange(6).reshape(2, 3).float()
    result = Tensor(a.uop.sum(axis=1)).numpy()
    assert result[0] == 3 and result[1] == 12

  def test_uop_sum_all(self):
    a = Tensor.arange(6).reshape(2, 3).float()
    self.assertAlmostEqual(Tensor(a.uop.sum()).item(), 15.0)

  def test_uop_sum_negative_axis(self):
    a = Tensor.arange(6).reshape(2, 3).float()
    result = Tensor(a.uop.sum(axis=-1)).numpy()
    assert result[0] == 3 and result[1] == 12

  def test_uop_sum_multi_axis(self):
    a = Tensor.arange(24).reshape(2, 3, 4).float()
    ref = a.sum(axis=(0, 2)).numpy()
    result = Tensor(a.uop.sum(axis=(0, 2))).numpy()
    for i in range(3): self.assertAlmostEqual(result[i], ref[i])

  def test_uop_sum_dtype(self):
    a = Tensor([1.0, 2, 3], dtype=dtypes.float16)
    result = Tensor(a.uop.sum(axis=0, dtype=dtypes.float32))
    self.assertEqual(result.dtype, dtypes.float)
    self.assertAlmostEqual(result.item(), 6.0, places=2)

  def test_uop_prod(self):
    a = Tensor([1.0, 2, 3, 4, 5])
    self.assertAlmostEqual(Tensor(a.uop.prod(axis=0)).item(), 120.0)

  def test_uop_max(self):
    a = Tensor([1.0, 5, 3, 2, 4])
    self.assertAlmostEqual(Tensor(a.uop.max(axis=0)).item(), 5.0)

  def test_uop_max_2d(self):
    a = Tensor([[1, 5, 3], [4, 2, 6]]).float()
    result = Tensor(a.uop.max(axis=0)).numpy()
    assert result[0] == 4 and result[1] == 5 and result[2] == 6

  def test_uop_std(self):
    a = Tensor([2.0, 4, 4, 4, 5, 5, 7, 9])
    self.assertAlmostEqual(Tensor(a.uop.std()).item(), a.std().item(), places=5)

class TestUOpWhere(unittest.TestCase):
  def test_uop_where_both_const(self):
    cond = Tensor([True, False, True])
    result = Tensor(cond.uop.where(1, 0))
    self.assertEqual(result.tolist(), [1, 0, 1])

    result = Tensor(cond.uop.where(1.5, 0))
    self.assertEqual(result.tolist(), [1.5, 0, 1.5])

if __name__ == '__main__':
  unittest.main()