Merge branch 'master' into kernel_is_call

extra/gemm/asm/cdna/asm.py

@@ -9,8 +9,7 @@ GEMM_ARGS = {
   (8192, 4096, 4096): (256, 64, 32768),
   (8192, 14336, 4096): (256, 64, 114688),
   (8192, 4096, 14336): (256, 224, 114688),
-  # TODO: get a fast gemm for this shape
-  #(8192, 128256, 4096): (16032, 64, 1026048),
+  (8192, 128256, 4096): (16032, 64, 1026048),
   (8192, 8192, 8192): (256, 128, 131072),
   (4096, 4096, 4096): (256, 64, 16384),
   (4096, 14336, 4096): (256, 64, 57344),

extra/models/llama.py

@@ -55,7 +55,7 @@ class Attention:
       xqkv = x @ self.wqkv.T
       xq, xk, xv = xqkv.split([self.wq.weight.shape[0], self.wk.weight.shape[0], self.wv.weight.shape[0]], dim=2)
     else:
-      xq, xk, xv = self.wq(x), self.wk(x.contiguous_backward()), self.wv(x)
+      xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
 
     if self.q_norm is not None and self.k_norm is not None:
       xq = self.q_norm(xq)

test/helpers.py

@@ -62,7 +62,7 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None):
     bufs.append(buf:=allocator.alloc(len(data) * buf_dt.itemsize))
     allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *data)))
   g = UOp(Ops.PARAM, uop.dtype.ptr(), arg=0, src=())
-  prg = get_program(UOp.store(g.index(UOp.const(dtypes.int, 0)), uop).sink(arg=KernelInfo()), PythonRenderer())
+  prg = get_program(UOp.store(g.index(UOp.const(dtypes.int, 0)), uop).sink(), PythonRenderer())
   prog = PythonProgram("run", PythonCompiler().compile(prg.src))
   prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *bufs)
   return out_buf.cast(uop.dtype.fmt or "").tolist()[0]

test/null/test_linearizer_failures.py

@@ -1,6 +1,6 @@
 # ruff: noqa: E501
 import unittest
-from tinygrad.uop.ops import UOp, Ops, AxisType, KernelInfo
+from tinygrad.uop.ops import UOp, Ops, AxisType
 from tinygrad.dtype import dtypes
 from tinygrad.engine.realize import get_program
 from tinygrad.device import Device
@@ -18,7 +18,7 @@ class TestLinearizerFailures(unittest.TestCase):
     c8 = c7.index(c3)
     c9 = ((((c6+(c8*UOp.const(dtypes.float, -1.0)))*(c6+(c8*UOp.const(dtypes.float, -1.0)))).reduce(c5, arg=Ops.ADD)*UOp.const(dtypes.float, 0.000390625))+UOp.const(dtypes.float, 1e-05)).sqrt().reciprocal()
     c10 = c0.index(c3).store(c9).end(c1, c2)
-    ast = c10.sink(arg=KernelInfo())
+    ast = c10.sink()
     get_program(ast, renderer=Device[Device.DEFAULT].renderer)
 
 if __name__ == '__main__':

test/null/test_linearizer_rewrite.py

@@ -35,7 +35,7 @@ class TestLinearizerRewrite(unittest.TestCase):
     prg = get_program(ast, Device["CPU"].renderer)
     assert prg.applied_opts == (), f"expected no opts, got {prg}"
 
-    prg = get_program(ast.replace(arg=KernelInfo()), Device["CPU"].renderer)
+    prg = get_program(ast.replace(arg=None), Device["CPU"].renderer)
     assert prg.applied_opts != (), f"expected opts to apply, got {prg.applied_opts}"
 
     prg = get_program(ast.replace(arg=KernelInfo(name="custom")), Device["CPU"].renderer)

test/null/test_tinyfs.py

@@ -5,22 +5,18 @@ class TestLoadStore(unittest.TestCase):
   def test_load_shape(self):
     t = Tensor(bytes(16)).fs_load(1024)
     assert t.shape == (1024,), t.shape
-    t.schedule()
 
   def test_store_shape(self):
     t = Tensor.zeros(1024).fs_store()
     assert t.shape == (16,), t.shape
-    t.schedule()
 
   def test_load_large_shape(self):
     t = Tensor(bytes(16)).fs_load(10_000_000)
     assert t.shape == (10_000_000,), t.shape
-    t.schedule()
 
   def test_store_large_shape(self):
     t = Tensor.zeros(10_000_000).fs_store()
     assert t.shape == (16,), t.shape
-    t.schedule()
 
 if __name__ == "__main__":
   unittest.main()

test/test_custom_kernel.py

@@ -52,7 +52,7 @@ def flip_contract_kernel(dest:UOp, src:UOp):
   j = UOp.range(dest.shape[1], 1, AxisType.UPCAST)
   vec = src[i, j].contract(j)
   store = UOp.group(*[dest[i, k].store(vec.gep(3-k)) for k in range(4)])
-  return store.end(i, j).sink(arg=KernelInfo(name=f"flip_contract_{dest.size}", opts_to_apply=()))
+  return store.end(i).sink(arg=KernelInfo(name=f"flip_contract_{dest.size}", opts_to_apply=()))
 
 def slice_sum_kernel(dest:UOp, src:UOp):
   G = UOp.range(src.shape[0], 0)
@@ -104,7 +104,7 @@ def backward_gemm_custom(gradient:UOp, kernel:UOp) -> tuple[UOp, UOp]:
 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 = Tensor.custom_kernel(a, fxn=lambda _: UOp.sink())[0]
     a.realize()
 
   def test_simple(self):

test/test_profiler.py

@@ -178,7 +178,6 @@ class TestProfiler(unittest.TestCase):
 
     print("pairwise clock jitter matrix (us):\n" + '\n'.join([''.join([f'{float(item):8.3f}' for item in row]) for row in jitter_matrix]))
 
-  @unittest.skip("this test is flaky")
   def test_cpu_profile(self):
     def test_fxn(err=False):
       time.sleep(0.1)

test/test_renderer_failures.py

@@ -9,7 +9,7 @@ from tinygrad.renderer.cstyle import CStyleLanguage
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.wgsl import WGSLRenderer
 from tinygrad.runtime.ops_python import PythonRenderer
-from tinygrad.uop.ops import UOp, Ops, KernelInfo, python_alu
+from tinygrad.uop.ops import UOp, Ops, python_alu
 from tinygrad.tensor import Tensor, _to_np_dtype
 
 def _test_uop_result(inputs:list[Tensor], prg, local_size=None):
@@ -32,7 +32,7 @@ def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp):
   ld = b.index(idx)
   alu = ld.alu(alu_op, *alu_src_uops)
   store = UOp.store(a.index(idx), alu)
-  sink = UOp(Ops.SINK, dtypes.void, (store,), arg=KernelInfo())
+  sink = UOp(Ops.SINK, dtypes.void, (store,))
   prg = get_program(sink, Device[Device.DEFAULT].renderer)
   return _test_uop_result([Tensor([input_val])], prg)[0]
 
@@ -42,7 +42,7 @@ class TestRendererFailures(unittest.TestCase):
     a = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
     gate_alu = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0)
     gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index(lidx0.valid(gate_alu)), UOp.const(dtypes.int, 1)))
-    sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,), arg=KernelInfo())
+    sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,))
     prg = get_program(sink, Device[Device.DEFAULT].renderer)
     ret = _test_uop_result([], prg, local_size=[4, 1, 1])[0]
     np.testing.assert_equal(ret, [0, 1, 1, 1])
@@ -53,7 +53,7 @@ class TestRendererFailures(unittest.TestCase):
     gate_alu_0 = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0)
     gate_alu_1 = (lidx1:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 2),), 'lidx1')).ne(0)
     gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index((lidx0+lidx1*4).valid(gate_alu_0&gate_alu_1)), UOp.const(dtypes.int, 1)))
-    sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,), arg=KernelInfo())
+    sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,))
     prg = get_program(sink, Device[Device.DEFAULT].renderer)
     ret = _test_uop_result([], prg, local_size=[4, 2, 1])[0]
     np.testing.assert_equal(ret, [0, 0, 0, 0, 0, 1, 1, 1])
@@ -99,7 +99,7 @@ class TestPTXFailures(unittest.TestCase):
     val = UOp.const(dtypes.int, 1)
     if_uop = UOp(Ops.IF, dtypes.void, (gate_alu,))
     gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index(lidx0, if_uop), val))
-    sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,), arg=KernelInfo())
+    sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,))
     prg = get_program(sink, Device[Device.DEFAULT].renderer)
     ret = _test_uop_result([], prg, local_size=[4, 1, 1])[0]
     np.testing.assert_equal(ret, [0, 1, 1, 1])

test/test_schedule.py

@@ -10,8 +10,9 @@ from hypothesis import assume, given, settings, strategies as strat
 from tinygrad import nn, dtypes, Device, Tensor, Variable
 from tinygrad.device import is_dtype_supported
 from tinygrad.dtype import DType, ImageDType
-from tinygrad.uop.ops import UOp, Ops, GroupOp, UPat, Kernel
+from tinygrad.uop.ops import UOp, Ops, GroupOp, UPat
 from tinygrad.helpers import CI, DEBUG, SPLIT_REDUCEOP, GlobalCounters, Context, getenv, all_same, temp
+from tinygrad.schedule.rangeify import Kernel
 from tinygrad.engine.realize import CompiledRunner, run_schedule
 
 class KernelCountException(Exception): pass

test/test_uops.py

@@ -23,7 +23,7 @@ def to_uops_list(u:list[UOp], ren=None) -> list[UOp]:
   return ret
 
 def _uops_to_prg(uops_list):
-  prg = get_program(UOp.sink(*uops_list, arg=KernelInfo()), Device[Device.DEFAULT].renderer)
+  prg = get_program(UOp.sink(*uops_list), Device[Device.DEFAULT].renderer)
   return CompiledRunner(replace(prg, device=Device.DEFAULT))
 
 def uop(uops:list[UOp], op:Ops, dtype:Optional[DType], src:tuple[UOp, ...], arg:Any=None) -> UOp:
@@ -281,7 +281,7 @@ class TestUOpPrograms(unittest.TestCase):
     ptr = UOp.placeholder(out.shape, out.dtype, slot=0)
     i, j = UOp.range(10, axis_id=0), UOp.range(10, axis_id=1)
     prog = ptr[i,j].set(42).end(i,j)
-    self._run(prog.sink(arg=KernelInfo()), out)
+    self._run(prog.sink(), out)
 
     with Context(DEBUG=0): self.assertTrue((out == 42).all().item())
 

test/testextra/test_asm_gemm.py

@@ -3,7 +3,6 @@ from tinygrad import Tensor, Device, dtypes, Context
 from tinygrad.device import is_dtype_supported
 from tinygrad.helpers import getenv, CI
 from extra.gemm.asm.cdna.gemm import asm_gemm
-from test.helpers import needs_second_gpu
 
 def verify_asm_gemm(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=1) -> None:
   Tensor.manual_seed(0)
@@ -38,7 +37,6 @@ class TestGemm(unittest.TestCase):
   def test_simple(self): verify_asm_gemm(1, N:=(getenv("N", 4096)//SCALE), N, N, dtype=dtypes.half)
   def test_gemm(self): verify_asm_gemm(1, 8192//SCALE, 4096//SCALE, 14336//SCALE)
   def test_gemm_batched(self): verify_asm_gemm(2, 8192//SCALE, 4096//SCALE, 4096//SCALE)
-  @needs_second_gpu
   def test_gemm_multi(self): verify_asm_gemm(2, 8192//SCALE, 4096//SCALE, 4096//SCALE, gpus=2)
 
 class TestGemmLarge(unittest.TestCase):
@@ -47,13 +45,11 @@ class TestGemmLarge(unittest.TestCase):
       self.skipTest("very slow on non mi350x")
 
   def test_gemm1(self): verify_asm_gemm(8, 8192, 4096, 14336, dtype=dtypes.bfloat16, gpus=8)
-  @unittest.skip("disabled, asm in this shape is slower than tinygrad")
   def test_gemm2(self): verify_asm_gemm(8, 8192, 128256, 4096, dtype=dtypes.bfloat16, gpus=8)
   def test_gemm3(self): verify_asm_gemm(8, 8192, 14336, 4096, dtype=dtypes.bfloat16, gpus=8)
   def test_gemm4(self): verify_asm_gemm(8, 4096, 14336, 4096, dtype=dtypes.bfloat16, gpus=8)
   def test_gemm5(self): verify_asm_gemm(8, 4096, 4096, 14336, dtype=dtypes.bfloat16, gpus=8)
   def test_gemm6(self): verify_asm_gemm(16, 4096, 4096, 14336, dtype=dtypes.bfloat16, gpus=8)
-  @unittest.skip("disabled, asm in this shape is slower than tinygrad")
   def test_gemm7(self): verify_asm_gemm(1, 8192, 128256, 4096)
   def test_gemm_unsupported(self):
     with self.assertRaisesRegex(AssertionError, "shape not supported"):

test/unit/test_assign.py

@@ -259,15 +259,15 @@ class TestAssign(unittest.TestCase):
     np.testing.assert_allclose(out, [1.,1.,1.,1.,1.,1.,0.,0.,1.,1.,1.,1.,1.,1.,0.,0.])
 
   def test_assign_contiguous(self):
-    b = Tensor.arange(16).reshape(4,4).contiguous().realize()
-    a = (Tensor.arange(16).reshape(4,4).contiguous().realize() + 1)
+    b = Tensor.rand(4,4).realize()
+    a = (Tensor.rand(4,4).realize() + 1)
     kc = GlobalCounters.kernel_count
     b.assign(a.contiguous()).realize()
     assert GlobalCounters.kernel_count - kc == 2
 
   def test_assign_contiguous_permute(self):
-    b = Tensor.arange(16).reshape(4,4).contiguous().realize()
-    a = (Tensor.arange(16).reshape(4,4).contiguous().realize() + 1).permute((1,0))
+    b = Tensor.rand(4,4).realize()
+    a = (Tensor.rand(4,4).realize() + 1).permute((1,0))
     kc = GlobalCounters.kernel_count
     b.assign(a.contiguous()).realize()
     assert GlobalCounters.kernel_count - kc == 2
@@ -333,7 +333,7 @@ class TestAssign(unittest.TestCase):
 
   @unittest.skip("multi output not supported anymore")
   def test_simple_assignment_multioutput(self):
-    a = Tensor.arange(32*32).reshape(32, 32).contiguous().realize()
+    a = Tensor.randn(32, 32).realize()
     b = Tensor.full((32, ), 1.).contiguous().realize()
     c = Tensor.full((32, ), 2.).contiguous().realize()
     d = Tensor.full((32, ), 3.).contiguous().realize()
@@ -361,16 +361,16 @@ class TestAssign(unittest.TestCase):
     np.testing.assert_equal(a.numpy(), np.arange(4 * 4).reshape(4, 4).transpose(1, 0) + np.arange(4 * 4).reshape(4, 4))
 
   def test_permuted_reduceop_child_dual_use(self):
-    a = Tensor.arange(32*32*32).reshape(32, 32, 32).contiguous().realize()
-    b = Tensor.ones(32, 32, dtype=dtypes.int).contiguous().realize()
+    a = Tensor.randn(32, 32, 32).realize()
+    b = Tensor.full((32, 32), 1.).contiguous().realize()
     r = a.sum(axis=1)
     b.assign(r + b.permute(1, 0))
     b.realize()
-    np.testing.assert_equal(b.numpy(), a.numpy().sum(axis=1)+np.ones((32, 32), dtype=np.int32).transpose(1, 0))
+    np.testing.assert_allclose(b.numpy(), a.numpy().sum(axis=1)+np.ones((32, 32)).transpose(1, 0), atol=1e-6, rtol=1e-3)
 
   @unittest.skip("multi output not supported anymore")
   def test_permuted_reduceop_multioutput_dual_use(self):
-    a = Tensor.arange(32*32*32).reshape(32, 32, 32).contiguous().realize()
+    a = Tensor.randn(32, 32, 32).realize()
     b = Tensor.full((32, 32), 1.).contiguous().realize()
     c = Tensor.full((32, 32), 2.).contiguous().realize()
 
@@ -383,7 +383,7 @@ class TestAssign(unittest.TestCase):
 
   @unittest.skip("multi output not supported anymore")
   def test_permuted_reduceop_multioutput_dual_use_possible(self):
-    a = Tensor.arange(32*32*32).reshape(32, 32, 32).contiguous().realize()
+    a = Tensor.randn(32, 32, 32, dtype=dtypes.int).realize()
     b = Tensor.arange(32 * 32).reshape(32, 32).realize()
     c = Tensor.arange(32 * 32).reshape(32, 32).realize()
 
@@ -529,7 +529,6 @@ class TestAssign(unittest.TestCase):
     a = Tensor.empty(5, device=f"disk:{temp('disk_assignment')}").assign(Tensor.ones(5)).numpy()
     np.testing.assert_equal(a, np.ones(5))
 
-  @unittest.skip("this test is crashing!")
   def test_assign_slice_then_read(self):
     """Assign to slice then read from buffer - read should see the assigned values.
     This is the KV cache pattern from llm.py.

test/unit/test_disk_tensor.py

@@ -4,17 +4,9 @@ from tinygrad import Tensor, Device, dtypes
 from tinygrad.device import is_dtype_supported
 from tinygrad.dtype import DType, DTYPES_DICT
 from tinygrad.nn.state import safe_load, safe_save, get_state_dict, torch_load
-from tinygrad.helpers import Timing, fetch, OSX, dedup
+from tinygrad.helpers import Timing, fetch, temp, OSX
 from test.helpers import slow
 
-class TempDirTestCase(unittest.TestCase):
-  def setUp(self):
-    self.temp_dir = tempfile.TemporaryDirectory()
-  def tearDown(self):
-    self.temp_dir.cleanup()
-  def tmp(self, name:str) -> str:
-    return (pathlib.Path(self.temp_dir.name) / name).as_posix()
-
 def compare_weights_both(url):
   import torch
   fn = fetch(url)
@@ -92,7 +84,7 @@ class TestRawDiskBuffer(unittest.TestCase):
     pathlib.Path(tmp).unlink()
 
 @unittest.skipUnless(is_dtype_supported(dtypes.uint8), "need uint8")
-class TestSafetensors(TempDirTestCase):
+class TestSafetensors(unittest.TestCase):
   def test_real_safetensors(self):
     import torch
     from safetensors.torch import save_file
@@ -103,19 +95,19 @@ class TestSafetensors(TempDirTestCase):
       "weight3": torch.arange(0, 17, dtype=torch.int32).reshape(17,1,1),
       "weight4": torch.arange(0, 2, dtype=torch.uint8),
     }
-    save_file(tensors, self.tmp("real.safetensors"))
+    save_file(tensors, temp("real.safetensors"))
 
-    ret = safe_load(self.tmp("real.safetensors"))
+    ret = safe_load(temp("real.safetensors"))
     for k,v in tensors.items(): np.testing.assert_array_equal(ret[k].numpy(), v.numpy())
-    safe_save(ret, self.tmp("real.safetensors_alt"))
-    with open(self.tmp("real.safetensors"), "rb") as f:
-      with open(self.tmp("real.safetensors_alt"), "rb") as g:
+    safe_save(ret, temp("real.safetensors_alt"))
+    with open(temp("real.safetensors"), "rb") as f:
+      with open(temp("real.safetensors_alt"), "rb") as g:
         assert f.read() == g.read()
-    ret2 = safe_load(self.tmp("real.safetensors_alt"))
+    ret2 = safe_load(temp("real.safetensors_alt"))
     for k,v in tensors.items(): np.testing.assert_array_equal(ret2[k].numpy(), v.numpy())
 
   def test_real_safetensors_open(self):
-    fn = self.tmp("real_safe")
+    fn = temp("real_safe")
     state_dict = {"tmp": Tensor.rand(10,10)}
     safe_save(state_dict, fn)
     import os
@@ -131,15 +123,15 @@ class TestSafetensors(TempDirTestCase):
     from extra.models.efficientnet import EfficientNet
     model = EfficientNet(0)
     state_dict = get_state_dict(model)
-    safe_save(state_dict, self.tmp("eff0"))
-    state_dict_loaded = safe_load(self.tmp("eff0"))
+    safe_save(state_dict, temp("eff0"))
+    state_dict_loaded = safe_load(temp("eff0"))
     assert sorted(state_dict_loaded.keys()) == sorted(state_dict.keys())
     for k,v in state_dict.items():
       np.testing.assert_array_equal(v.numpy(), state_dict_loaded[k].numpy())
 
     # load with the real safetensors
     from safetensors import safe_open
-    with safe_open(self.tmp("eff0"), framework="pt", device="cpu") as f:
+    with safe_open(temp("eff0"), framework="pt", device="cpu") as f:
       assert sorted(f.keys()) == sorted(state_dict.keys())
       for k in f.keys():
         np.testing.assert_array_equal(f.get_tensor(k).numpy(), state_dict[k].numpy())
@@ -163,19 +155,19 @@ class TestSafetensors(TempDirTestCase):
 
   def test_metadata(self):
     metadata = {"hello": "world"}
-    safe_save({}, self.tmp('metadata.safetensors'), metadata)
+    safe_save({}, temp('metadata.safetensors'), metadata)
     import struct
-    with open(self.tmp('metadata.safetensors'), 'rb') as f:
+    with open(temp('metadata.safetensors'), 'rb') as f:
       dat = f.read()
     sz = struct.unpack(">Q", dat[0:8])[0]
     import json
     assert json.loads(dat[8:8+sz])['__metadata__']['hello'] == 'world'
 
   def test_save_all_dtypes(self):
-    for dtype in dedup(DTYPES_DICT.values()):
+    for dtype in DTYPES_DICT.values():
       if dtype in [dtypes.bfloat16]: continue # not supported in numpy
       if not is_dtype_supported(dtype): continue
-      path = self.tmp(f"ones.{dtype}.safetensors")
+      path = temp(f"ones.{dtype}.safetensors")
       ones = Tensor(np.random.rand(10,10), dtype=dtype)
       safe_save(get_state_dict(ones), path)
       np.testing.assert_equal(ones.numpy(), list(safe_load(path).values())[0].numpy())
@@ -197,9 +189,9 @@ class TestSafetensors(TempDirTestCase):
       "weight_I16": torch.tensor([127, 64], dtype=torch.short),
       "weight_BF16": torch.randn((2, 2), dtype=torch.bfloat16),
     }
-    save_file(tensors, self.tmp("dtypes.safetensors"))
+    save_file(tensors, temp("dtypes.safetensors"))
 
-    loaded = safe_load(self.tmp("dtypes.safetensors"))
+    loaded = safe_load(temp("dtypes.safetensors"))
     for k,v in loaded.items():
       if v.dtype != dtypes.bfloat16:
         assert v.numpy().dtype == tensors[k].numpy().dtype
@@ -211,52 +203,57 @@ class TestSafetensors(TempDirTestCase):
       "weight_U32": np.array([1, 2, 3], dtype=np.uint32),
       "weight_U64": np.array([1, 2, 3], dtype=np.uint64),
     }
-    np_save_file(tensors, self.tmp("dtypes.safetensors"))
+    np_save_file(tensors, temp("dtypes.safetensors"))
 
-    loaded = safe_load(self.tmp("dtypes.safetensors"))
+    loaded = safe_load(temp("dtypes.safetensors"))
     for k,v in loaded.items():
       assert v.numpy().dtype == tensors[k].dtype
       np.testing.assert_allclose(v.numpy(), tensors[k])
 
-def helper_test_disk_tensor(tmp, fn, data, np_fxn, tinygrad_fxn=None):
+def helper_test_disk_tensor(fn, data, np_fxn, tinygrad_fxn=None):
   if tinygrad_fxn is None: tinygrad_fxn = np_fxn
-  pathlib.Path(tmp(fn)).unlink(missing_ok=True)
-  tinygrad_tensor = Tensor(data, device="CPU").to(f"disk:{tmp(fn)}")
+  pathlib.Path(temp(fn)).unlink(missing_ok=True)
+  tinygrad_tensor = Tensor(data, device="CPU").to(f"disk:{temp(fn)}")
   numpy_arr = np.array(data)
   tinygrad_fxn(tinygrad_tensor)
   np_fxn(numpy_arr)
   np.testing.assert_allclose(tinygrad_tensor.numpy(), numpy_arr)
 
-class TestDiskTensor(TempDirTestCase):
+class TestDiskTensor(unittest.TestCase):
   def test_empty(self):
-    Tensor.empty(100, 100, device=f"disk:{self.tmp('dt_empty')}")
+    pathlib.Path(temp("dt_empty")).unlink(missing_ok=True)
+    Tensor.empty(100, 100, device=f"disk:{temp('dt_empty')}")
 
   def test_simple_read(self):
-    fn = pathlib.Path(self.tmp("dt_simple_read"))
+    fn = pathlib.Path(temp("dt_simple_read"))
+    fn.unlink(missing_ok=True)
     fn.write_bytes(bytes(range(256)))
-    t = Tensor.empty(16, 16, device=f"disk:{self.tmp('dt_simple_read')}", dtype=dtypes.uint8)
+    t = Tensor.empty(16, 16, device=f"disk:{temp('dt_simple_read')}", dtype=dtypes.uint8)
     out = t[1].to(Device.DEFAULT).tolist()
     assert out == list(range(16, 32))
 
   def test_simple_read_bitcast(self):
-    fn = pathlib.Path(self.tmp("dt_simple_read_bitcast"))
+    fn = pathlib.Path(temp("dt_simple_read_bitcast"))
+    fn.unlink(missing_ok=True)
     fn.write_bytes(bytes(range(256))*2)
-    t = Tensor.empty(16, 16*2, device=f"disk:{self.tmp('dt_simple_read_bitcast')}", dtype=dtypes.uint8)
+    t = Tensor.empty(16, 16*2, device=f"disk:{temp('dt_simple_read_bitcast')}", dtype=dtypes.uint8)
     out = t[1].bitcast(dtypes.uint16).to(Device.DEFAULT).tolist()
     tout = [(x//256, x%256) for x in out]
     assert tout == list([(x+1,x) for x in range(32,64,2)])
 
   def test_simple_read_bitcast_alt(self):
-    fn = pathlib.Path(self.tmp("dt_simple_read_bitcast_alt"))
+    fn = pathlib.Path(temp("dt_simple_read_bitcast_alt"))
+    fn.unlink(missing_ok=True)
     fn.write_bytes(bytes(range(256))*2)
-    t = Tensor.empty(16, 16*2, device=f"disk:{self.tmp('dt_simple_read_bitcast_alt')}", dtype=dtypes.uint8)
+    t = Tensor.empty(16, 16*2, device=f"disk:{temp('dt_simple_read_bitcast_alt')}", dtype=dtypes.uint8)
     out = t.bitcast(dtypes.uint16)[1].to(Device.DEFAULT).tolist()
     tout = [(x//256, x%256) for x in out]
     assert tout == list([(x+1,x) for x in range(32,64,2)])
 
   def test_strided_read(self):
     # test non-contiguous (strided) read - should read elements at indices 0, 2, 4
-    dt = Tensor([0, 1, 2, 3, 4, 5]).to(f"disk:{self.tmp('dt_strided_read')}")
+    pathlib.Path(temp(fn:="dt_strided_read")).unlink(missing_ok=True)
+    dt = Tensor([0, 1, 2, 3, 4, 5]).to(f"disk:{temp(fn)}")
     result = dt[::2].tolist()
     # TODO: dt[::2] selects indices 0, 2, 4, so result should be [0, 2, 4]
     # self.assertEqual(result, [0, 2, 4])
@@ -264,38 +261,43 @@ class TestDiskTensor(TempDirTestCase):
 
   def test_permuted_read(self):
     # test non-contiguous (permuted) read - should read transposed
-    dt = Tensor([[0, 1, 2], [3, 4, 5]]).to(f"disk:{self.tmp('dt_permuted_read')}")
+    pathlib.Path(temp(fn:="dt_permuted_read")).unlink(missing_ok=True)
+    dt = Tensor([[0, 1, 2], [3, 4, 5]]).to(f"disk:{temp(fn)}")
     result = dt.T.tolist()
     # TODO: transpose should give [[0, 3], [1, 4], [2, 5]]
     # self.assertEqual(result, [[0, 3], [1, 4], [2, 5]])
     self.assertEqual(result, [[0, 1], [2, 3], [4, 5]])  # wrong!
 
   def test_write_ones(self):
+    pathlib.Path(temp("dt_write_ones")).unlink(missing_ok=True)
+
     out = Tensor.ones(10, 10, device="CPU").contiguous()
-    outdisk = out.to(f"disk:{self.tmp('dt_write_ones')}")
+    outdisk = out.to(f"disk:{temp('dt_write_ones')}")
     print(outdisk)
     outdisk.realize()
     del out, outdisk
 
     import struct
     # test file
-    with open(self.tmp("dt_write_ones"), "rb") as f:
+    with open(temp("dt_write_ones"), "rb") as f:
       assert f.read() == struct.pack('<f', 1.0) * 100 == b"\x00\x00\x80\x3F" * 100
 
     # test load alt
-    reloaded = Tensor.empty(10, 10, device=f"disk:{self.tmp('dt_write_ones')}")
+    reloaded = Tensor.empty(10, 10, device=f"disk:{temp('dt_write_ones')}")
     np.testing.assert_almost_equal(reloaded.numpy(), np.ones((10, 10)))
 
   def test_simple_setitem(self):
+    pathlib.Path(temp(fn:="dt_simple_setitem")).unlink(missing_ok=True)
     data = [[1],[2]]
     src = Tensor(data)
-    dt = src.to(f"disk:{self.tmp('dt_simple_setitem')}")
+    dt = src.to(f"disk:{temp(fn)}")
     dt[1] = [3]
     self.assertEqual(dt.tolist(), [[1], [3]])
 
   def test_strided_setitem(self):
     # test non-contiguous (strided) setitem - should set elements at indices 0, 2, 4
-    dt = Tensor([1, 2, 3, 4, 5, 6]).to(f"disk:{self.tmp('dt_strided_setitem')}")
+    pathlib.Path(temp(fn:="dt_strided_setitem")).unlink(missing_ok=True)
+    dt = Tensor([1, 2, 3, 4, 5, 6]).to(f"disk:{temp(fn)}")
     dt[::2] = Tensor([10, 20, 30])
     # TODO: dt[::2] selects indices 0, 2, 4, so result should be [10, 2, 20, 4, 30, 6]
     # self.assertEqual(dt.tolist(), [10, 2, 20, 4, 30, 6])
@@ -303,35 +305,39 @@ class TestDiskTensor(TempDirTestCase):
 
   def test_assign_const_to_disk(self):
     # assign from CONST (Tensor.full) to disk - source has no buffer, needs contiguous first
-    dt = Tensor.empty(4, device=f"disk:{self.tmp('dt_assign_const')}", dtype=dtypes.int32)
+    pathlib.Path(temp(fn:="dt_assign_const")).unlink(missing_ok=True)
+    dt = Tensor.empty(4, device=f"disk:{temp(fn)}", dtype=dtypes.int32)
     dt.assign(Tensor.full((4,), 42, dtype=dtypes.int32)).realize()
     np.testing.assert_array_equal(dt.numpy(), [42, 42, 42, 42])
 
   def test_assign_slice_from_const(self):
     # slice assign from CONST to disk - tests size calculation when no RANGE ops
-    dt = Tensor([0, 1, 2, 3], dtype=dtypes.int32).to(f"disk:{self.tmp('dt_slice_const')}")
+    pathlib.Path(temp(fn:="dt_slice_const")).unlink(missing_ok=True)
+    dt = Tensor([0, 1, 2, 3], dtype=dtypes.int32).to(f"disk:{temp(fn)}")
     dt[1:3].assign(Tensor.full((2,), 99, dtype=dtypes.int32)).realize()
     np.testing.assert_array_equal(dt.numpy(), [0, 99, 99, 3])
 
   def test_disk_to_disk_copy(self):
     # disk-to-disk copy needs to go through CPU
-    src = Tensor([1, 2, 3, 4], dtype=dtypes.int32).to(f"disk:{self.tmp('dt_d2d_src')}")
-    dst = Tensor.empty(4, device=f"disk:{self.tmp('dt_d2d_dst')}", dtype=dtypes.int32)
+    pathlib.Path(temp(fn1:="dt_d2d_src")).unlink(missing_ok=True)
+    pathlib.Path(temp(fn2:="dt_d2d_dst")).unlink(missing_ok=True)
+    src = Tensor([1, 2, 3, 4], dtype=dtypes.int32).to(f"disk:{temp(fn1)}")
+    dst = Tensor.empty(4, device=f"disk:{temp(fn2)}", dtype=dtypes.int32)
     dst.assign(src.to("CPU")).realize()
     np.testing.assert_array_equal(dst.numpy(), [1, 2, 3, 4])
 
   def test_assign_slice(self):
     def assign(x,s,y): x[s] = y
-    helper_test_disk_tensor(self.tmp, "dt_assign_slice_1", [0,1,2,3], lambda x: assign(x, slice(0,2), [13, 12]))
-    helper_test_disk_tensor(self.tmp, "dt_assign_slice_2", [[0,1,2,3],[4,5,6,7]], lambda x: assign(x, slice(0,1), [[13, 12, 11, 10]]))
+    helper_test_disk_tensor("dt_assign_slice_1", [0,1,2,3], lambda x: assign(x, slice(0,2), [13, 12]))
+    helper_test_disk_tensor("dt_assign_slice_2", [[0,1,2,3],[4,5,6,7]], lambda x: assign(x, slice(0,1), [[13, 12, 11, 10]]))
 
   def test_reshape(self):
-    helper_test_disk_tensor(self.tmp, "dt_reshape_1", [1,2,3,4,5], lambda x: x.reshape((1,5)))
-    helper_test_disk_tensor(self.tmp, "dt_reshape_2", [1,2,3,4], lambda x: x.reshape((2,2)))
+    helper_test_disk_tensor("dt_reshape_1", [1,2,3,4,5], lambda x: x.reshape((1,5)))
+    helper_test_disk_tensor("dt_reshape_2", [1,2,3,4], lambda x: x.reshape((2,2)))
 
   def test_assign_to_different_dtype(self):
     # NOTE: this is similar to Y_train in fetch_cifar
-    t = Tensor.empty(10, device=f'disk:{self.tmp("dt_assign_to_different_dtype")}', dtype=dtypes.int64)
+    t = Tensor.empty(10, device=f'disk:{temp("dt_assign_to_different_dtype")}', dtype=dtypes.int64)
 
     for i in range(5):
       data = np.array([3, 3])
@@ -343,7 +349,8 @@ class TestDiskTensor(TempDirTestCase):
   def test_assign_with_bitcast(self):
     # bitcast assign is used in safe_save for writing header length
     # bitcast on source side works, bitcast on target side raises
-    t = Tensor.empty(16, device=f"disk:{self.tmp('dt_assign_bitcast')}", dtype=dtypes.uint8)
+    pathlib.Path(temp(fn:="dt_assign_bitcast")).unlink(missing_ok=True)
+    t = Tensor.empty(16, device=f"disk:{temp(fn)}", dtype=dtypes.uint8)
     # correct way: bitcast the source to match target dtype
     t[0:8].assign(Tensor([12345], dtype=dtypes.int64, device="CPU").bitcast(dtypes.uint8))
     val = int.from_bytes(t[0:8].data(), 'little')
@@ -354,7 +361,8 @@ class TestDiskTensor(TempDirTestCase):
 
   def test_assign_to_bitcast_view(self):
     # assign float values to a float32 view of a uint8 disk buffer (used by safe_save)
-    t = Tensor.empty(32, device=f"disk:{self.tmp('dt_bitcast_view_assign')}", dtype=dtypes.uint8)
+    pathlib.Path(temp(fn:="dt_bitcast_view_assign")).unlink(missing_ok=True)
+    t = Tensor.empty(32, device=f"disk:{temp(fn)}", dtype=dtypes.uint8)
     # create float32 view of bytes 8-24 (4 floats)
     float_view = t[8:24].bitcast(dtypes.float32)
     float_view.assign(Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32, device="CPU"))
@@ -362,20 +370,21 @@ class TestDiskTensor(TempDirTestCase):
 
   def test_assign_cross_device(self):
     # disk assign allows cross-device (source on GPU/CPU, target on disk)
-    t = Tensor.empty(4, device=f"disk:{self.tmp('dt_assign_cross')}", dtype=dtypes.float32)
+    pathlib.Path(temp(fn:="dt_assign_cross")).unlink(missing_ok=True)
+    t = Tensor.empty(4, device=f"disk:{temp(fn)}", dtype=dtypes.float32)
     src = Tensor([1.0, 2.0, 3.0, 4.0])  # on default device
     t.assign(src)
     np.testing.assert_array_equal(t.numpy(), [1.0, 2.0, 3.0, 4.0])
 
   def test_bitcast(self):
-    with open(self.tmp('dt_bitcast'), "wb") as f: f.write(bytes(range(10,20)))
-    t = Tensor.empty(5, dtype=dtypes.int16, device=f"disk:{self.tmp('dt_bitcast')}")
+    with open(temp('dt_bitcast'), "wb") as f: f.write(bytes(range(10,20)))
+    t = Tensor.empty(5, dtype=dtypes.int16, device=f"disk:{temp('dt_bitcast')}")
     ret = t.to("CPU").bitcast(dtypes.uint16) + 1
     assert ret.tolist() == [2827, 3341, 3855, 4369, 4883]
 
   def test_bitcast_view(self):
-    with open(self.tmp('dt_bitcast_view'), "wb") as f: f.write(bytes(range(10, 24)))
-    t = Tensor.empty(3, dtype=dtypes.uint, device=f"disk:{self.tmp('dt_bitcast_view')}").shrink([(0, 2)])
+    with open(temp('dt_bitcast_view'), "wb") as f: f.write(bytes(range(10, 24)))
+    t = Tensor.empty(3, dtype=dtypes.uint, device=f"disk:{temp('dt_bitcast_view')}").shrink([(0, 2)])
     ret = t.bitcast(dtypes.uint16).to("CPU") + 1
     assert ret.tolist() == [2827, 3341, 3855, 4369]
 
@@ -383,55 +392,59 @@ class TestDiskTensor(TempDirTestCase):
   @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), "bfloat16 not supported")
   def test_bf16_disk_write_read(self):
     t = Tensor([10000, -1, -1000, -10000, 20], dtype=dtypes.float32)
-    t.to(f"disk:{self.tmp('dt_bf16_disk_write_read_f32')}").realize()
+    t.to(f"disk:{temp('dt_bf16_disk_write_read_f32')}").realize()
 
     # hack to "cast" f32 -> bf16
-    with open(self.tmp('dt_bf16_disk_write_read_f32'), "rb") as f: dat = f.read()
+    with open(temp('dt_bf16_disk_write_read_f32'), "rb") as f: dat = f.read()
     adat = b''.join([dat[i+2:i+4] for i in range(0, len(dat), 4)])
-    with open(self.tmp('dt_bf16_disk_write_read_bf16'), "wb") as f: f.write(adat)
+    with open(temp('dt_bf16_disk_write_read_bf16'), "wb") as f: f.write(adat)
 
-    t = Tensor.empty(5, dtype=dtypes.bfloat16, device=f"disk:{self.tmp('dt_bf16_disk_write_read_bf16')}")
+    t = Tensor.empty(5, dtype=dtypes.bfloat16, device=f"disk:{temp('dt_bf16_disk_write_read_bf16')}")
     ct = t.to(Device.DEFAULT).cast(dtypes.float)
     assert ct.numpy().tolist() == [9984., -1, -1000, -9984, 20]
 
   def test_copy_from_disk(self):
-    fn = pathlib.Path(self.tmp("dt_copy_from_disk"))
+    fn = pathlib.Path(temp("dt_copy_from_disk"))
+    fn.unlink(missing_ok=True)
     fn.write_bytes(bytes(range(256))*1024)
 
-    t = Tensor.empty(256*1024, device=f"disk:{self.tmp('dt_copy_from_disk')}", dtype=dtypes.uint8)
+    t = Tensor.empty(256*1024, device=f"disk:{temp('dt_copy_from_disk')}", dtype=dtypes.uint8)
     on_dev = t.to(Device.DEFAULT).realize()
     np.testing.assert_equal(on_dev.numpy(), t.numpy())
 
   def test_copy_from_disk_offset(self):
-    fn = pathlib.Path(self.tmp("dt_copy_from_disk_offset"))
+    fn = pathlib.Path(temp("dt_copy_from_disk_offset"))
+    fn.unlink(missing_ok=True)
     fn.write_bytes(bytes(range(256))*1024)
 
     for off in [314, 991, 2048, 4096]:
-      t = Tensor.empty(256*1024, device=f"disk:{self.tmp('dt_copy_from_disk_offset')}", dtype=dtypes.uint8)[off:]
+      t = Tensor.empty(256*1024, device=f"disk:{temp('dt_copy_from_disk_offset')}", dtype=dtypes.uint8)[off:]
       on_dev = t.to(Device.DEFAULT).realize()
       np.testing.assert_equal(on_dev.numpy(), t.numpy())
 
   @slow
   def test_copy_from_disk_huge(self):
-    fn = pathlib.Path(self.tmp("dt_copy_from_disk_huge"))
+
+    fn = pathlib.Path(temp("dt_copy_from_disk_huge"))
+    fn.unlink(missing_ok=True)
     fn.write_bytes(bytes(range(256))*1024*256)
 
     for off in [0, 551]:
-      t = Tensor.empty(256*1024*256, device=f"disk:{self.tmp('dt_copy_from_disk_huge')}", dtype=dtypes.uint8)[off:]
+      t = Tensor.empty(256*1024*256, device=f"disk:{temp('dt_copy_from_disk_huge')}", dtype=dtypes.uint8)[off:]
       on_dev = t.to(Device.DEFAULT).realize()
       np.testing.assert_equal(on_dev.numpy(), t.numpy())
 
   @unittest.skip("this allocates a lot of RAM")
   @unittest.skipUnless(OSX, "seems to only be an issue on macOS with file size >2 GiB")
   def test_copy_to_cpu_not_truncated(self):
-    fn = self.tmp("dt_copy_to_cpu_not_truncated")
-    with open(fn, "wb") as f: f.write(b'\x01' * (size := int(2 * 1024**3)) + (test := b"test"))
+    with open((fn:=temp("dt_copy_to_cpu_not_truncated")), "wb") as f: f.write(b'\x01' * (size := int(2 * 1024**3)) + (test := b"test"))
     x = Tensor.empty(size + len(test), dtype=dtypes.uint8, device=f"disk:{fn}").to("CPU").realize()
     assert x[size:].data().tobytes() == test
 
   def test_disk_device_reuse(self):
     from tinygrad.runtime.ops_disk import DiskDevice
-    fn = pathlib.Path(self.tmp("dt_device_reuse"))
+    fn = pathlib.Path(temp("dt_device_reuse"))
+    fn.unlink(missing_ok=True)
     fn.write_bytes(bytes(range(256)))
     # create first tensor and realize it
     t1 = Tensor.empty(128, device=f"disk:{fn}", dtype=dtypes.uint8)
@@ -453,7 +466,8 @@ class TestDiskTensor(TempDirTestCase):
 
   def test_disk_open_failure_state(self):
     from tinygrad.runtime.ops_disk import DiskDevice
-    fn = pathlib.Path(self.tmp("dt_open_failure"))
+    fn = pathlib.Path(temp("dt_open_failure"))
+    fn.unlink(missing_ok=True)
     fn.write_bytes(bytes(range(256)))
     os.chmod(fn, 0o000)
     try:
@@ -472,7 +486,8 @@ class TestDiskTensor(TempDirTestCase):
     assert disk_device.size == 200
 
   def test_disk_permission_error(self):
-    fn = pathlib.Path(self.tmp("dt_permission"))
+    fn = pathlib.Path(temp("dt_permission"))
+    fn.unlink(missing_ok=True)
     fn.write_bytes(bytes(range(256)))
     os.chmod(fn, 0o000)
     try:
@@ -481,14 +496,17 @@ class TestDiskTensor(TempDirTestCase):
     finally:
       os.chmod(fn, 0o644)
 
-class TestPathTensor(TempDirTestCase):
+class TestPathTensor(unittest.TestCase):
   def setUp(self):
-    super().setUp()
+    self.temp_dir = tempfile.TemporaryDirectory()
     self.test_file = pathlib.Path(self.temp_dir.name) / "test_file.bin"
     self.test_data = np.arange(100, dtype=np.uint8).tobytes()
     with open(self.test_file, "wb") as f:
       f.write(self.test_data)
 
+  def tearDown(self):
+    self.temp_dir.cleanup()
+
   def test_path_tensor_no_device(self):
     t = Tensor(self.test_file)
     self.assertEqual(t.shape, (100,))
@@ -539,10 +557,10 @@ class TestPathTensor(TempDirTestCase):
     os.chmod(test_file, 0o644)
     assert Tensor(pathlib.Path(test_file)).tolist(), list(range(10))
 
-class TestDiskTensorMovement(TempDirTestCase):
+class TestDiskTensorMovement(unittest.TestCase):
   def setUp(self):
-    super().setUp()
-    self.fn = pathlib.Path(self.tmp("custom_disk_range"))
+    self.fn = pathlib.Path(temp("custom_disk_range"))
+    self.fn.unlink(missing_ok=True)
     Tensor.arange(100, dtype=dtypes.uint8).to(f"disk:{str(self.fn)}").realize()
 
   def test_simple_read(self):

tinygrad/codegen/__init__.py

@@ -1,5 +1,4 @@
 from typing import cast
-from dataclasses import replace
 import itertools
 from tinygrad.helpers import DISABLE_FAST_IDIV, EMULATED_DTYPES, DEVECTORIZE, TRANSCENDENTAL, SPEC, DEBUG, getenv, TracingKey, Context
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, pyrender
@@ -165,19 +164,17 @@ def get_program(ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> Program
     The ProgramSpec of the program.
   """
 
+  # fix up KernelInfo
+  if opts is not None:
+    assert ast.arg is None, "can't apply opts if sink has an arg"
+    ast = ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts)))
+  if ast.arg is None and ast.op is Ops.SINK: ast = ast.replace(arg=KernelInfo())
+
+  # rewrite to prg
   if ast.op is Ops.PROGRAM: prg = ast
-  elif ast.op is Ops.SINK:
-    # rewrite to prg
-    assert isinstance(ast.arg, KernelInfo), "requires KernelInfo on arg to get_program"
-    if opts is not None:
-      # TODO: should this be here?
-      assert ast.arg.opts_to_apply is None, "can't apply opts if there's already opts to apply"
-      ast = ast.replace(arg=replace(ast.arg, opts_to_apply=tuple(opts)))
+  else:
     full_sink = full_rewrite_to_sink(ast, renderer, optimize=ast.tag is None)
     prg = UOp(Ops.PROGRAM, src=(full_sink, UOp(Ops.DEVICE, arg=renderer.device)))
-  else:
-    raise RuntimeError(f"can't call get_program on {ast.op}")
-
   prg = graph_rewrite(prg, pm_to_program, ctx=renderer, name="linearize/render")
 
   # create the ProgramSpec

tinygrad/engine/schedule.py

@@ -1,7 +1,7 @@
 import time
 from typing import cast
 from collections import deque
-from tinygrad.uop.ops import UOp, Ops, buffers, UOpMetaClass, track_rewrites, PatternMatcher, UPat, graph_rewrite, graph_rewrite_map, Kernel
+from tinygrad.uop.ops import UOp, Ops, buffers, UOpMetaClass, track_rewrites, PatternMatcher, UPat, graph_rewrite, graph_rewrite_map, CallInfo
 from tinygrad.uop.spec import type_verify, tensor_spec
 from tinygrad.device import Buffer, MultiBuffer
 from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, flatten, pluralize, SCACHE, Metadata
@@ -28,8 +28,7 @@ def create_schedule(sched_sink:UOp) -> tuple[list[ExecItem], UOp]:
       if (k:=u.src[1]).op is Ops.RANGE: continue  # RANGEs are scheduled directly, not through dependency graph
       assert k.op in {Ops.KERNEL, Ops.END}, f"AFTER src[1] should be KERNEL or END, not {k.op}"
       in_degree.setdefault(k, 0)
-      if k.op is Ops.END: assert k.src[0].op is Ops.KERNEL, f"END src[0] should be KERNEL, not {k.src[0].op}"
-      for s in k.src[0].src if k.op is Ops.END else k.src:
+      for s in k.src[0].src if k.op is Ops.END else k.src[1:]:
         match (s := _unwrap_src(s)).op:
           case Ops.AFTER:
             children.setdefault(s.src[1], []).append(k)
@@ -54,14 +53,13 @@ def create_schedule(sched_sink:UOp) -> tuple[list[ExecItem], UOp]:
     while len(queue):
       k = rk = queue.popleft()
       if k.op is Ops.END: k = k.src[0]
-      assert k.op in {Ops.RANGE, Ops.KERNEL}, f"unexpected op in queue: {k.op}"
+      assert k.op in {Ops.RANGE, Ops.CALL}, f"unexpected op in queue: {k.op}"
       if k.op is Ops.RANGE: schedule.append(k)
-      elif k.op is Ops.KERNEL:
-        ast = (kernel:=cast(Kernel, k.arg)).ast
-        buf_uops = tuple(_unwrap_src(s).buf_uop for s in k.src if s.op is not Ops.BIND)
-        bound_ranges = tuple(s for s in k.src if s.op is Ops.BIND and len(s.src) > 1 and s.src[1].op is Ops.RANGE)
-        sched_item[k] = (ast, buf_uops, kernel.metadata, bound_ranges)
-        schedule.append(k)
+      elif k.op is Ops.CALL:
+        ast = k.src[0]
+        buf_uops = tuple(_unwrap_src(s).buf_uop for s in k.src[1:] if s.op is not Ops.BIND)
+        bound_ranges = tuple(s for s in k.src[1:] if s.op is Ops.BIND and len(s.src) > 1 and s.src[1].op is Ops.RANGE)
+        schedule.append((ast, buf_uops, cast(CallInfo, k.arg).metadata, {}, bound_ranges))
         if rk.op is Ops.END: schedule.append(rk)
       for x in children.get(rk, []):
         in_degree[x] -= 1

tinygrad/schedule/indexing.py

@@ -9,7 +9,7 @@ from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored
 
 ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.ASSIGN, Ops.COPY, Ops.BUFFER, Ops.BUFFER_VIEW,
                      Ops.CONST, Ops.BIND, Ops.DEVICE, Ops.MSELECT, Ops.MSTACK, Ops.PARAM,
-                     Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.LOAD, Ops.KERNEL, Ops.ENCDEC}
+                     Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.LOAD, Ops.CALL, Ops.ENCDEC}
 
 def realize(ctx:dict[UOp, None], tr:UOp) -> None: ctx[tr] = None
 
@@ -19,7 +19,8 @@ def realize_srcs(ctx:dict[UOp, None], rb:UOp) -> None:
 
 def realize_assign(ctx:dict[UOp, None], a:UOp) -> None:
   if a.src[1].op not in ALWAYS_CONTIGUOUS: ctx[a.src[1]] = None
-  ctx[a] = None
+  # if it's a kernel, we don't realize it
+  if a.src[1].op is not Ops.CALL: ctx[a] = None
 
 pm_generate_realize_map = PatternMatcher([
   # always realize SINK src
@@ -98,7 +99,7 @@ def remove_movement_op_after_rangeify(ctx:IndexingContext, x:UOp):
   if x in ctx.range_map or x.src[0].op is Ops.INDEX: return x.src[0]
 
 def add_third_op_to_assign_to_track_shape(ctx:IndexingContext, assign:UOp):
-  if assign.src[1].op is Ops.KERNEL: return None
+  if assign.src[1].op is Ops.CALL: return None
   to_mop = graph_rewrite(assign.src[0], PatternMatcher([(UPat(GroupOp.Movement, name="x"), lambda x: x.replace(tag=()))]))
   ret = assign.replace(src=assign.src+(to_mop,))
   ctx.range_map[ret] = ctx.range_map[assign]
@@ -161,7 +162,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
   rctx = IndexingContext()
 
   # get ops to realize
-  graph_rewrite(tsink, pm_generate_realize_map, ctx=rctx.realize_map, bottom_up=True, name="get realize")
+  graph_rewrite(tsink, pm_generate_realize_map, ctx=rctx.realize_map, name="get realize")
 
   # get the consumer map
   with cpu_profile("consumer map in rangeify", "TINY"):
@@ -173,7 +174,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
     if x.op in {Ops.DEVICE, Ops.UNIQUE}: continue
 
     # no ranges on kernels, they are internal
-    if x.op is Ops.KERNEL: continue
+    if x.op is Ops.CALL: continue
 
     if x.dtype.scalar() == dtypes.index: continue  # TODO: why do I need this?
     ending_ranges[x] = sum([ending_ranges.get(u, []) for u in consumer_map[x]], [])

tinygrad/schedule/rangeify.py

@@ -2,10 +2,10 @@ from dataclasses import dataclass, field, replace
 import itertools
 from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, _substitute, KernelInfo
-from tinygrad.uop.ops import graph_rewrite, identity_element, sint, AxisType, BottomUpGate, Kernel, _remove_all_tags, range_str
+from tinygrad.uop.ops import graph_rewrite, identity_element, sint, AxisType, BottomUpGate, _remove_all_tags, range_str
 from tinygrad.uop.symbolic import symbolic
 from tinygrad.helpers import argsort, prod, all_same, getenv, flatten, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY
-from tinygrad.helpers import PCONTIG, partition, get_single_element
+from tinygrad.helpers import PCONTIG, partition, get_single_element, panic
 from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_simplify
 from tinygrad.codegen.opt import Opt
 from tinygrad.schedule.indexing import run_rangeify, BufferizeOpts, ALWAYS_CONTIGUOUS, IndexingContext, apply_movement_op
@@ -66,12 +66,11 @@ mop_cleanup = PatternMatcher([
 
 def resolve_custom_kernel(ck:UOp) -> UOp:
   placeholders = [UOp.placeholder_like(s, slot=i) for i,s in enumerate(ck.src)]
-  return UOp(Ops.KERNEL, src=ck.src, arg=Kernel(ck.arg.fxn(*placeholders)))
+  return ck.arg.fxn(*placeholders).call(*ck.src)
 
 def resolve_call(c:UOp) -> UOp|None:
-  # don't resolve real kernel calls, sink or program
-  if c.src[0].op is Ops.SINK and isinstance(c.src[0].arg, KernelInfo): return None
-  if c.src[0].op is Ops.PROGRAM: return None
+  # don't resolve CALLs with SINK - those are kernel calls from split_store/custom_kernel
+  if c.src[0].op is Ops.SINK: return None
   params = sorted([x for x in c.src[0].toposort() if x.op == Ops.PARAM], key=lambda x: x.arg)
   args = c.src[1:]
   # TODO: this check belongs in spec, not here
@@ -86,7 +85,7 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
   # just removing it works...
   (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD), name="x"), lambda x: x.src[0]),
 
-  # resolve calls
+  # resolve calls (but not CALLs with SINK - those are kernel calls from split_store)
   (UPat(Ops.CALL, name="c"), resolve_call),
 
   # resolve custom kernels
@@ -274,14 +273,15 @@ def late_buffer_view(t:UOp, b:UOp):
   size = prod(shape)
 
   # walk up for the INDEX
+  # NOTE: even though we allow RESHAPE and SHRINK, they can combine to form non-contiguous access patterns (e.g. t[::2])
   x = t
-  while not any(u.op is Ops.INDEX for u in x.src):
-    assert x.op not in GroupOp.Elementwise, "can't buffer view elementwise"
+  while x.op is not Ops.INDEX:
+    assert x.op in {Ops.BITCAST, Ops.CONTIGUOUS, Ops.SHRINK, Ops.RESHAPE}, f"unexpected op {x.op} in buffer view walk"
     x = x.src[0]
-  x = next(u for u in x.src if u.op is Ops.INDEX)
 
   if len(shape) == 0: offset = x.src[1].arg
-  else: offset = max(sum(idx.vmin for idx in x.src[1:]), 0)
+  else: offset = sum(idx.vmin for idx in x.src[1:])
+  if offset < 0: raise RuntimeError(f"negative offset {offset} in buffer view")
 
   return b.replace(src=(UOp(Ops.BUFFER_VIEW, t.dtype, (x.base,), (size, offset), tag=t.tag), b.src[1]))
 
@@ -386,8 +386,8 @@ pm_add_buffers = pm_mops+pm_flatten_bufferize+to_bufferview+PatternMatcher([
   (UPat((Ops.MSELECT, Ops.MSTACK), src=UPat(Ops.RESHAPE), name="m"),
    lambda m: m.replace(src=tuple([x.src[0].base for x in m.src]), tag=None).reshape(m.shape).rtag(m.tag)),
 
-  # remove any RESHAPEs on KERNEL
-  (UPat(Ops.KERNEL, name="k"), lambda k: k.replace(src=tuple(x.src[0] if x.op is Ops.RESHAPE else x for x in k.src))),
+  # remove any RESHAPEs on CALL
+  (UPat(Ops.CALL, name="k"), lambda k: k.replace(src=(k.src[0],)+tuple(x.src[0] if x.op is Ops.RESHAPE else x for x in k.src[1:]))),
 ])
 
 pm_add_buffers_local = pm_mops+pm_flatten_bufferize+to_bufferview+PatternMatcher([
@@ -516,17 +516,18 @@ def split_store(ctx:list[UOp], x:UOp) -> UOp|None:
   elif ret.op is Ops.END and ret.src[0].op is Ops.STORE: stored = ret.src[0].src[1]
   else: raise RuntimeError(f"unknown kernel type {ret.op}")
   if stored.op in {Ops.COPY, Ops.BUFFER_VIEW, Ops.ENCDEC}: ret = stored
-  else: ret = ret.sink(arg=KernelInfo(opts_to_apply=lctx.opts))
-
+  else:
+    ret = ret.sink(arg=KernelInfo(opts_to_apply=lctx.opts) if lctx.opts is not None else KernelInfo())
   metadata = tuple(dedup(flatten([x for x in metadatas if x is not None])))[::-1]
-  kernel_arg = Kernel(ret, metadata)
-  kernel = UOp(Ops.KERNEL, src=tuple(lctx.map.values())+tuple(lctx.vars.keys()), arg=kernel_arg)
-  if ret.op is Ops.SINK and not all_same([x.device for x in kernel.src if x.op is not Ops.BIND]):
-    raise RuntimeError(f"all buffers must be on the same device: {tuple(b.buf_uop for b in kernel.src)}")
+  kernel = ret.call(*lctx.map.values(), *lctx.vars.keys(), metadata=metadata)
+  if ret.op is Ops.SINK and not all_same([x.device for x in kernel.src[1:] if x.op is not Ops.BIND]):
+    raise RuntimeError(f"all buffers must be on the same device: {tuple(b.buf_uop for b in kernel.src[1:])}")
   return kernel
 
 split_kernels = PatternMatcher([
   (UPat((Ops.STORE, Ops.END), name="x"), split_store),
+  # if it's a Kernel, stop
+  (UPat(Ops.SINK, name="sink"), lambda sink: panic(BottomUpGate()) if isinstance(sink.arg, KernelInfo) else None),
 ])
 
 def tag_uop(ctx:tuple[list[UOp], set[UOp]], x:UOp):
@@ -539,7 +540,7 @@ def tag_uop(ctx:tuple[list[UOp], set[UOp]], x:UOp):
   return x.replace(tag=(len(ctx[0])-1,))
 add_tags = PatternMatcher([
   # don't tag BUFFERs, they are global
-  (UPat(GroupOp.All-{Ops.BUFFER, Ops.CONST, Ops.DEVICE, Ops.UNIQUE, Ops.LUNIQUE, Ops.DEFINE_VAR, Ops.BIND, Ops.KERNEL, Ops.END,
+  (UPat(GroupOp.All-{Ops.BUFFER, Ops.CONST, Ops.DEVICE, Ops.UNIQUE, Ops.LUNIQUE, Ops.DEFINE_VAR, Ops.BIND, Ops.CALL, Ops.END,
                      Ops.MSTACK, Ops.MSELECT, Ops.RANGE}.union(GroupOp.Movement), name="x"), tag_uop),
   (UPat({Ops.MSTACK, Ops.MSELECT}, name="x"), lambda ctx,x: None if all(s.op is Ops.BUFFER for s in x.src) else tag_uop(ctx, x)),
 ])
@@ -583,7 +584,7 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
   # bufferize -> store
   lunique_start: int = max([-1]+[x.arg for x in tsink.toposort() if x.op is Ops.LUNIQUE]) + 1
   tsink = graph_rewrite(tsink, pm_add_buffers+pm_add_range_tags, ctx=itertools.count(lunique_start), bottom_up=True, name="bufferize to store")
-  tsink = graph_rewrite(tsink, split_kernels, ctx=uop_list, bottom_up=True, name="split kernels")
+  tsink = graph_rewrite(tsink, split_kernels, ctx=uop_list, name="split kernels", bottom_up=True)
 
   # if a kernel depends on a buffer, and that buffer is later assigned to, make the assign depend on the kernel's assign
   kernel_assign: dict[UOp, UOp] = {}

tinygrad/uop/ops.py

@@ -67,8 +67,7 @@ def consumer_map_from_toposort(lst:Iterable[UOp]):
   ret: dict[UOp, dict[UOp, None]] = {}
   for u in lst:
     ret[u] = {}
-    for s in u.src:
-      if s in ret: ret[s][u] = None
+    for s in u.src: ret[s][u] = None
   return ret
 
 def pretty_print(x:UOp, cache=None, d=0)->str:
@@ -311,8 +310,6 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
   def ended_ranges(self):
     if self.op in range_start: return self.src[range_start[self.op]:]
     if self.op is Ops.AFTER: return tuple(flatten([x.ended_ranges for x in self.src[1:]]))
-    # TODO: copy isn't using range properly and isn't ending the range it uses, remove this
-    if self.op in {Ops.COPY, Ops.BUFFER_VIEW}: return self.src[0].ranges
     return ()
 
   # determine what ranges this is in
@@ -367,7 +364,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
   @recursive_property
   def trace_num(self):
     num = next(ucount)
-    # KERNEL also has a UOp in the arg
+    # KERNEL has a UOp in the arg, CALL has it in src[0] so no special handling needed
     arg = type(self.arg)(self.arg.ast.trace_num, self.arg.metadata) if self.op is Ops.KERNEL else self.arg
     uop_fields[num] = (self.op, self.dtype, tuple(s.trace_num for s in self.src), arg, self.tag)+((self.metadata,) if TRACEMETA>=2 else ())
     return num
@@ -822,7 +819,6 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     return UOp(Ops.PARAM, dtype, src, arg=slot)
 
   def call(self, *srcs:UOp, grad_fxn:Callable|None=None, metadata:tuple[Metadata, ...]=()) -> UOp:
-    assert len(self.ranges) == 0, f"ranges {self.ranges} are leaking out of the call"
     return UOp(Ops.CALL, self.dtype, (self,)+srcs, CallInfo(grad_fxn, metadata))
   def custom_kernel(*srcs:UOp, fxn:Callable, grad_fxn:Callable|None=None) -> list[UOp]:
     contig_srcs = tuple(x.contiguous() if x.op is not Ops.AFTER else x for x in srcs)
@@ -1418,7 +1414,7 @@ pm_pyrender_extra = PatternMatcher([
 
 # NOTE: you can remove pm_pyrender_extra and it'll still be correct
 pm_pyrender = pm_pyrender_extra+PatternMatcher([
-  (UPat(Ops.KERNEL, name="u"), lambda ctx,u: f"UOp(Ops.KERNEL, src={srcs(ctx,u.src)}, arg=Kernel({ctx[u.arg.ast]}(), {u.arg.metadata}))"),
+  (UPat(Ops.CALL, name="u"), lambda ctx,u: f"{ctx[u.src[0]]}.call({', '.join(ctx[s] for s in u.src[1:])}, metadata={u.arg.metadata})"),
   (UPat(GroupOp.All, name="u"), lambda ctx,u: f"UOp({u.op}, {u.dtype}, {srcs(ctx,u.src)}"+(f", {repr(u.arg)})" if u.arg is not None else ")")),
 ])
 
@@ -1428,7 +1424,7 @@ def pyrender(ast:UOp) -> str:
   cmap = consumer_map_from_toposort(lst)
   not_rendered = {Ops.CONST, Ops.VCONST, Ops.DEVICE}
   always_rendered = {Ops.PARAM, Ops.LOAD, Ops.SPECIAL, Ops.RANGE, Ops.CONTIGUOUS, Ops.VECTORIZE,
-                     Ops.BUFFER, Ops.COPY, Ops.KERNEL, Ops.WHERE, Ops.END, Ops.ASSIGN}
+                     Ops.BUFFER, Ops.COPY, Ops.CALL, Ops.WHERE, Ops.END, Ops.ASSIGN}
 
   to_render: set[UOp] = {ast}
   for u in lst:

tinygrad/uop/spec.py

@@ -87,8 +87,9 @@ _tensor_spec = PatternMatcher([
   (UPat(Ops.BUFFER, src=(UPat((Ops.LUNIQUE, Ops.UNIQUE)), UPat(Ops.DEVICE)), name="buf"),
    lambda buf: isinstance(buf.arg, int) and isinstance(buf.dtype, (DType, ImageDType))),
 
-  # KERNEL can attach to an AFTER to describe the compute required to realize a BUFFER
-  (UPat(Ops.KERNEL, src=UPat((Ops.BUFFER, Ops.AFTER, Ops.MSELECT, Ops.MSTACK, Ops.BIND))), lambda: True),
+  # CALL can attach to an AFTER to describe the compute required to realize a BUFFER
+  # src[0] is the function (SINK), src[1:] are buffers/bindings
+  (UPat(Ops.CALL, src=(UPat(Ops.SINK), UPat((Ops.BUFFER, Ops.AFTER, Ops.MSELECT, Ops.MSTACK, Ops.BIND))), allow_any_len=True), lambda: True),
 
   # ASSIGN has a target and a value. It can also optionally depend on other assigns
   (UPat(Ops.ASSIGN, name="x"), lambda x: len(x.src) >= 2 and all(s.op is Ops.ASSIGN for s in x.src[2:])),
@@ -249,8 +250,8 @@ full_spec = PatternMatcher([
   # vectorized index
   (UPat(Ops.INDEX, src=(UPat((Ops.VECTORIZE, Ops.CAST)), UPat())), lambda: True),
 
-  # linearizer: outputs + intermediate KERNELs
-  (UPat(Ops.KERNEL, dtype=dtypes.void), lambda: True),
+  # linearizer: outputs + intermediate CALLs
+  (UPat(Ops.CALL, dtype=dtypes.void), lambda: True),
 
   # Invalid must have type Index
   (UPat(Ops.CONST, arg=Invalid, name="x"), lambda x: x.dtype.scalar() == dtypes.index),

tinygrad/uop/symbolic.py

@@ -251,7 +251,7 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
   ((UPat.var("x", dtypes.index) + UPat.cvar("c")).cast(dtypes.sints, name="cast"), lambda x,c,cast:x.cast(cast.dtype)+c.cast(cast.dtype)),
   # only RANGE/IF/STORE/KERNEL have side effects
   (UPat(Ops.AFTER, name="x"), lambda x: x.replace(src=(x.src[0],)+
-    tuple(flatten([(y,) if y.op in {Ops.RANGE, Ops.STORE, Ops.KERNEL, Ops.BARRIER, Ops.END, Ops.UNROLL} else y.src for y in x.src[1:]])))),
+    tuple(flatten([(y,) if y.op in {Ops.RANGE, Ops.STORE, Ops.CALL, Ops.BARRIER, Ops.END, Ops.UNROLL} else y.src for y in x.src[1:]])))),
   # after with 1 src is just src[0]
   (UPat(Ops.AFTER, src=(UPat.var("s"),)), lambda s: s),
   # VECTORIZE/CONST

tinygrad/viz/serve.py

@@ -171,7 +171,7 @@ def option(s:int|None) -> int: return 0 if s is None else s+1
 device_ts_diffs:dict[str, tuple[Decimal, Decimal]] = {}
 def cpu_ts_diff(device:str, thread=0) -> Decimal: return device_ts_diffs.get(device, (Decimal(0),))[thread]
 
-amdgpu_targets:dict[str, int] = {}
+device_props:dict[str, dict] = {}
 
 DevEvent = ProfileRangeEvent|ProfileGraphEntry|ProfilePointEvent
 def flatten_events(profile:list[ProfileEvent]) -> Generator[tuple[Decimal, Decimal, DevEvent], None, None]:
@@ -308,7 +308,7 @@ def load_counters(profile:list[ProfileEvent]) -> None:
     if (sqtt:=v.get(ProfileSQTTEvent)):
       for e in sqtt:
         if e.itrace: steps.append(create_step(f"PKTS SE:{e.se}", (f"/prg-pkts-{e.se}", len(ctxs), len(steps)),
-                                              data=(e.blob, prg_events[k].lib, amdgpu_targets[e.device])))
+                                              data=(e.blob, prg_events[k].lib, device_props[e.device]["gfx_target_version"])))
       steps.append(create_step("SQTT", ("/prg-sqtt", len(ctxs), len(steps)), ((k, tag), sqtt, prg_events[k])))
     ctxs.append({"name":f"Exec {name}"+(f" n{run_number[k]}" if run_number[k] > 1 else ""), "steps":steps})
 
@@ -348,7 +348,7 @@ def unpack_sqtt(key:tuple[str, int], data:list, p:ProfileProgramEvent) -> tuple[
   # * init decoder
   from extra.sqtt.roc import decode
   base = unwrap(p.base)
-  addr_table = amd_decode(unwrap(p.lib), amdgpu_targets[p.device])
+  addr_table = amd_decode(unwrap(p.lib), device_props[p.device]["gfx_target_version"])
   disasm:dict[int, tuple[str, int]] = {addr+base:(inst.disasm(), inst.size()) for addr, inst in addr_table.items()}
   rctx = decode(data, {p.tag:disasm})
   cu_events:dict[str, list[ProfileEvent]] = {}
@@ -385,9 +385,8 @@ def get_profile(profile:list[ProfileEvent], sort_fn:Callable[[str], Any]=device_
   for ev in profile:
     if isinstance(ev, ProfileDeviceEvent):
       device_ts_diffs[ev.device] = (ev.comp_tdiff,ev.copy_tdiff if ev.copy_tdiff is not None else ev.comp_tdiff)
-      if (d:=ev.device.split(":")[0]) == "AMD":
-        device_decoders[d] = load_counters
-        amdgpu_targets[d] = unwrap(ev.props)["gfx_target_version"]
+      if ev.props is not None: device_props[ev.device] = ev.props
+      if (d:=ev.device.split(":")[0]) == "AMD": device_decoders[d] = load_counters
   # load device specific counters
   for fxn in device_decoders.values(): fxn(profile)
   # map events per device
@@ -511,7 +510,7 @@ def get_render(query:str) -> dict:
   if fmt == "asm":
     ret:dict = {"metadata":[]}
     if data.device.startswith("AMD") and data.lib is not None:
-      with soft_err(lambda err: ret.update(err)): ret.update(amdgpu_cfg(data.lib, amdgpu_targets[data.device]))
+      with soft_err(lambda err: ret.update(err)): ret.update(amdgpu_cfg(data.lib, device_props[data.device]["gfx_target_version"]))
       with soft_err(lambda err: ret["metadata"].append(err)): ret["metadata"].append(amd_readelf(data.lib))
     else: ret["src"] = get_stdout(lambda: (compiler:=Device[data.device].compiler).disassemble(compiler.compile(data.src)))
     return ret