simplify priority

.github/workflows/benchmark.yml

@@ -527,7 +527,7 @@ jobs:
     - name: Run 10 CIFAR training steps
       run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=330 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
     - name: Run 10 CIFAR training steps w HALF
-      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=390 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=350 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
     # - name: Run 10 CIFAR training steps w BF16
     #   run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
     # TODO: too slow

examples/whisper.py

@@ -3,7 +3,7 @@
 import sys, base64, multiprocessing, itertools, collections
 from typing import Optional, Union, Literal, List
 
-from tinygrad import Tensor, TinyJit, Variable, nn, dtypes
+from tinygrad import Tensor, TinyJit, Variable, nn
 from tinygrad.nn.state import torch_load, load_state_dict
 from tinygrad.helpers import getenv, fetch
 
@@ -244,16 +244,15 @@ def transcribe_waveform(model: Whisper, enc, waveforms, truncate=False):
 
   log_spec = prep_audio(waveforms, model.batch_size, truncate)
   nsample = model.decoder.max_tokens_to_sample
-  nctx = model.decoder.max_self_attn_cache_len
 
   def inferloop(ctx: Union[np.ndarray, List[np.ndarray]], encoded_audio):
     pos, next_tokens = 0, ctx
-    for i in range(nsample):
-      next_tokens = model.decoder(Tensor(next_tokens, dtype=dtypes.int32), pos, encoded_audio)[:, -1].argmax(axis=-1).numpy().astype(np.int32).reshape(-1, 1)
+    for i in range((nsample-len(start_tokens))*2):
+      next_tokens = model.decoder(Tensor(next_tokens), pos, encoded_audio)[:, -1].argmax(axis=-1).numpy().astype(np.int32).reshape(-1, 1)
       next_tokens[ctx[:, -1] == eot] = eot
       ctx = np.concatenate((ctx, next_tokens), axis=1)
       pos = ctx.shape[-1] - 1
-      if (next_tokens == eot).all() or pos == nctx: break
+      if (next_tokens == eot).all(): break
     return ctx
 
   def gettexttoks(line): return [tok for tok in line if tok < eot or tok > enc._special_tokens["<|notimestamps|>"]][-nsample+len(start_tokens):]

extra/sqtt/test_timing.py

@@ -18,17 +18,15 @@ from extra.sqtt.roc import decode, InstExec, PrgExec
 
 dev = Device["AMD"]
 
-def custom(arg:str, s:UOp|None=None) -> UOp: return UOp(Ops.CUSTOM, src=(s,) if s is not None else (), arg=arg)
-
 def asm_kernel(instrs:list[str], l:int=1, g:int=1) -> Tensor:
   name = sys._getframe(1).f_code.co_name
   def fxn(_):
     L = UOp.special(l, "lidx0")
     G = UOp.special(g, "gidx0")
-    op = custom("asm volatile (")
-    for inst in instrs: op = custom(f'  "{inst}\\n\\t"', op)
-    op = custom(");", op)
-    return UOp.sink(op, L, G, arg=KernelInfo(name=name))
+    ops:list[str] = [UOp(Ops.CUSTOM, arg="asm volatile (")]
+    for inst in instrs: ops.append(UOp(Ops.CUSTOM, src=(ops[-1],), arg=f'  "{inst}\\n\\t"'))
+    ops.append(UOp(Ops.CUSTOM, src=(ops[-1],), arg=");"))
+    return UOp.sink(*ops, L, G, arg=KernelInfo(name=name))
   k = Tensor.custom_kernel(Tensor.empty(1), fxn=fxn)[0]
   return k
 
@@ -89,11 +87,12 @@ class TestTiming(unittest.TestCase):
     n = 1
     def sleep_kernel(data0):
       assert data0.dtype.base == dtypes.ulong
-      op = custom("unsigned long long t0 = __builtin_readcyclecounter();")
-      op = custom(f"__builtin_amdgcn_s_sleep({n});", op)
-      op = custom(f"unsigned long long t1 = __builtin_readcyclecounter();", op)
-      op = custom(f"data0_{data0.size}[0] = t1 - t0;", op)
-      return UOp.sink(data0, op, arg=KernelInfo(name=f"sleep_{n}"))
+      ops:list[UOp] = []
+      ops.append(UOp(Ops.CUSTOM, arg="unsigned long long t0 = __builtin_readcyclecounter();"))
+      ops.append(UOp(Ops.CUSTOM, arg=f"__builtin_amdgcn_s_sleep({n});", src=(ops[-1],)))
+      ops.append(UOp(Ops.CUSTOM, arg="unsigned long long t1 = __builtin_readcyclecounter();", src=(ops[-1],)))
+      ops.append(UOp(Ops.CUSTOM, arg=f"data0_{data0.size}[0] = t1 - t0;", src=(ops[-1],)))
+      return UOp.sink(data0, *ops, arg=KernelInfo(name=f"sleep_{n}"))
     diff_hw_reg = Tensor.empty(1, dtype=dtypes.ulong)
     diff_hw_reg = Tensor.custom_kernel(diff_hw_reg, fxn=sleep_kernel)[0]
     with save_sqtt() as sqtt:

extra/usbgpu/tbgpu/installer/Shared/TinyGPUViewModel.swift

@@ -119,7 +119,14 @@ extension TinyGPUViewModel: OSSystemExtensionRequestDelegate {
 
 		os_log("sysex actionForReplacingExtension: %@ %@", existing, ext)
 
+		// Add appropriate logic here to determine whether to replace the extension
+		// with the new extension. Common things to check for include
+		// testing whether the new extension's version number is newer than
+		// the current version number, or whether the bundleIdentifier is different.
+		// For simplicity, this sample always replaces the current extension
+		// with the new one.
 		replacementAction = .replace
+
 		self.state = .activating
 		return replacementAction
 	}

extra/usbgpu/tbgpu/installer/TinyGPUDriverExtension/TinyGPUDriverUserClient.cpp

@@ -7,48 +7,30 @@
 struct TinyGPUDriverUserClient_IVars
 {
 	OSSharedPtr<TinyGPUDriver> provider = nullptr;
-
-	TinyGPUCreateDMAResp *dmas = nullptr;
-	size_t dmaCount = 0;
-	size_t dmaCap = 0;
-
-	int ensureDMACap(size_t need)
-	{
-		// not thread-safe
-		if (need <= dmaCap) return 0;
-
-		size_t newCap = dmaCap ? dmaCap * 2 : 16;
-		while (newCap < need) newCap *= 2;
-
-		auto *newArr = IONewZero(TinyGPUCreateDMAResp, newCap);
-		if (!newArr) return -kIOReturnNoMemory;
-
-		if (dmas && dmaCount) {
-			memcpy(newArr, dmas, dmaCount * sizeof(TinyGPUCreateDMAResp));
-		}
-
-		IOSafeDeleteNULL(dmas, TinyGPUCreateDMAResp, dmaCap);
-		dmas = newArr;
-		dmaCap = newCap;
-		return 0;
-	}
 };
 
 bool TinyGPUDriverUserClient::init()
 {
-	auto ok = super::init();
-	if (!ok) return false;
+	auto theAnswer = super::init();
+	if (!theAnswer) {
+		return false;
+	}
 
 	ivars = IONewZero(TinyGPUDriverUserClient_IVars, 1);
-	if (!ivars) return false;
+	if (ivars == nullptr) {
+		return false;
+	}
+
 	return true;
 }
 
 void TinyGPUDriverUserClient::free()
 {
-	if (ivars) {
-		IOSafeDeleteNULL(ivars, TinyGPUDriverUserClient_IVars, 1);
+	if (ivars != nullptr) {
+		ivars->provider.reset();
 	}
+
+	IOSafeDeleteNULL(ivars, TinyGPUDriverUserClient_IVars, 1);
 	super::free();
 }
 
@@ -77,22 +59,6 @@ error:
 
 kern_return_t TinyGPUDriverUserClient::Stop_Impl(IOService* in_provider)
 {
-	// release all DMA allocations for this client
-	if (ivars) {
-		for (size_t i = 0; i < ivars->dmaCount; i++) {
-			auto &d = ivars->dmas[i];
-			if (d.dmaCmd) {
-				d.dmaCmd->CompleteDMA(kIODMACommandCompleteDMANoOptions);
-				d.dmaCmd->release();
-				d.dmaCmd = nullptr;
-			}
-		}
-		ivars->dmaCount = 0;
-		IOSafeDeleteNULL(ivars->dmas, TinyGPUCreateDMAResp, ivars->dmaCap);
-		ivars->dmas = nullptr;
-		ivars->provider.reset();
-	}
-
 	return Stop(in_provider, SUPERDISPATCH);
 }
 
@@ -136,26 +102,26 @@ kern_return_t TinyGPUDriverUserClient::ExternalMethod(uint64_t selector, IOUserC
 
 kern_return_t IMPL(TinyGPUDriverUserClient, CopyClientMemoryForType)
 {
-	if (!memory) return kIOReturnBadArgument;
-	if (!ivars->provider.get()) return kIOReturnNotAttached;
+	if (!memory) {
+		return kIOReturnBadArgument;
+	}
+
+	if (ivars->provider.get() == nullptr) {
+		return kIOReturnNotAttached;
+	}
 
-	// bar handling, type is bar num
 	if (type < 6) {
 		uint32_t bar = (uint32_t)type;
 		return ivars->provider->MapBar(bar, memory);
 	}
 
-	// dma handling, type is size
-	if (ivars->ensureDMACap(ivars->dmaCount + 1)) {
-		os_log(OS_LOG_DEFAULT, "tinygpu: cannot grow dma array");
-		return kIOReturnNoMemory;
+	// dma page buffer
+	TinyGPUCreateDMAResp buf;
+	kern_return_t err = ivars->provider->CreateDMA(type, &buf);
+	if (err) {
+		return err;
 	}
 
-	TinyGPUCreateDMAResp buf{};
-	kern_return_t err = ivars->provider->CreateDMA(type, &buf);
-	if (err) return err;
-
-	ivars->dmas[ivars->dmaCount++] = buf;
 	*memory = buf.sharedBuf;
 	return 0;
 }

ruff.toml

@@ -1,6 +1,6 @@
 indent-width = 2
 preview = true
-target-version = "py311"
+target-version = "py310"
 
 lint.select = [
   "F",  # Pyflakes

test/test_custom_kernel.py

@@ -9,11 +9,6 @@ def custom_arange_kernel(C:UOp) -> UOp:
   i = UOp.range(C.size, 0)
   return C[i].store(i.cast(C.dtype.base)).end(i).sink(arg=KernelInfo(name=f"custom_arange_{C.size}"))
 
-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.base)).end(i, j).sink(arg=KernelInfo(name=f"custom_eye_{C.size}"))
-
 def custom_add_one_kernel(B:UOp, A:UOp) -> UOp:
   A,B = A.flatten(), B.flatten()
   assert B.size == A.size
@@ -130,12 +125,6 @@ class TestCustomKernel(unittest.TestCase):
     tst = tst.custom_kernel(fxn=custom_arange_kernel)[0]
     self.assertTrue((ref == tst).all().item())
 
-  def test_eye(self):
-    ref = Tensor.eye(1024).contiguous().realize()
-    tst = Tensor.empty_like(ref)
-    tst = tst.custom_kernel(fxn=custom_eye_kernel)[0]
-    self.assertTrue((ref == tst).all().item())
-
   def test_flip_contract(self):
     a = Tensor.randn(10,4)
     b = Tensor.empty_like(a)

test/test_linearizer.py

@@ -4,7 +4,7 @@ from dataclasses import replace
 
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.codegen.gpudims import get_grouped_dims
-from tinygrad.uop.ops import UOp, Ops, GroupOp, AxisType, PatternMatcher, graph_rewrite, UPat
+from tinygrad.uop.ops import UOp, Ops, GroupOp
 from tinygrad.device import Device, Buffer, is_dtype_supported
 from tinygrad.tensor import Tensor, _to_np_dtype
 from tinygrad.engine.realize import run_schedule, lower_schedule, CompiledRunner, get_program
@@ -38,22 +38,6 @@ class TestLinearizer(unittest.TestCase):
     np.testing.assert_equal(a.numpy(), ta)
     np.testing.assert_equal(b.numpy(), tb)
 
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken on ptx")
-  def test_late_bias_load(self):
-    img = Tensor.empty(1, 3, 16, 16)
-    w = Tensor.empty(16, 3, 3, 3)
-    b = Tensor.empty(16)
-    out = img.conv2d(w, b)
-    ast = helper_linearizer_opt(out)
-    uops = get_program(ast, opts=[]).uops
-    # slice at the last loop end
-    uslice = [i for i,u in enumerate(uops) if u.op == Ops.END][-1]
-    # only valid test if outermost range is the reduce
-    if uops[uslice].src[-1].arg[-1] == AxisType.REDUCE:
-      load_types = [u.src[0].dtype for u in uops[uslice+1:] if u.op == Ops.LOAD]
-      # assert that there is a global load after the reduce ends
-      assert any(dt.addrspace == AddrSpace.GLOBAL for dt in load_types)
-
   def _test_no_nested_ranges(self, lins, skip=None):
     for l in lins:
       range_in_acc = flatten([[x for x in u.src if x.op is Ops.RANGE] for u in l.uops if u.op is Ops.DEFINE_REG])
@@ -278,8 +262,6 @@ class TestLinearizer(unittest.TestCase):
     _assert_grouped_dims("gidx", (65536,), (16,16,256), False, [16,16,256], False)
     #                              2             ->             3
     _assert_grouped_dims("gidx", (128,128), (16,16,256), False, [16,16,64], False)
-    #                              2             ->             2
-    _assert_grouped_dims("gidx", (65536,2), (65535,65535,65535), False, [32768,4], False)
     # test when the only divisor is the square root of dim
     _assert_grouped_dims("gidx", (121,), (12,12,12), False, [11,11], False)
 
@@ -304,27 +286,6 @@ class TestLinearizer(unittest.TestCase):
     with self.assertRaises(RuntimeError):
       get_grouped_dims("gidx", (2,3,4,5,6), (16,16,16))
 
-    # TODO: In the above cases we only test if the shape after reshape is correct, never the indices.
-    # We should check if the returned indices are correct, for all cases.
-    # (65536, 2) -> (32768, 4)
-    dims, expected_limited_dims = (65536,2), (32768, 4)
-    idxs = get_grouped_dims("gidx", dims, (65535,65535,65535))
-    def match_div(): raise RuntimeError("match_div")
-    def match_mod(): raise RuntimeError("match_mod")
-    flat_idx_pattern = UPat(Ops.SPECIAL, arg='gidx0')*expected_limited_dims[1]+UPat(Ops.SPECIAL, arg='gidx1')
-    pm = PatternMatcher([
-      (flat_idx_pattern//dims[1], match_div),
-      (flat_idx_pattern%dims[1], match_mod)
-    ])
-
-    with self.assertRaises(RuntimeError) as error:
-      graph_rewrite(idxs[0], pm)
-    self.assertIn("match_div", str(error.exception))
-
-    with self.assertRaises(RuntimeError) as error:
-      graph_rewrite(idxs[1], pm)
-    self.assertIn("match_mod", str(error.exception))
-
     # # variable too large
     # with self.assertRaises(AssertionError):
     #   get_grouped_dims("gidx", (Variable("start_pos",0,16),3,4), (16,16,16), False,)
@@ -471,8 +432,6 @@ def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]:
   # now all input buffers in s[-1] should be realized
   # create fresh buffers for the outputs
   bufs = [Buffer(x.device, x.size, x.dtype).allocate() if i < len(s[-1].ast.src) else x for i,x in enumerate(s[-1].bufs)]
-  # ensure buffers are allocated
-  for b in bufs: b.ensure_allocated()
   return s[-1].ast, bufs
 
 def helper_linearizer_ast(ast:UOp, inputs:list[Tensor], *args, **kwargs):

test/unit/test_device.py

@@ -30,7 +30,7 @@ class TestDevice(unittest.TestCase):
 
   @unittest.skipIf(WIN and CI, "skipping windows test") # TODO: subproccess causes memory violation?
   def test_env_overwrite_default_compiler(self):
-    expect_failure = "\ntry: assert Device[Device.DEFAULT].compiler is None;\nexcept Exception: pass"
+    expect_failure = "\ntry: assert Device[Device.DEFAULT].compiler is None;\nexcept RuntimeError: pass"
 
     if Device.DEFAULT == "CPU":
       from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, ClangJITCompiler

test/unit/test_uop_symbolic.py

@@ -643,10 +643,6 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(lidx+(gidx//4)*8+2*(gidx%4), 0, 372, "(lidx+(gidx*2))")
     self.helper_test_variable(lidx+2*(gidx%4)+(gidx//4)*8, 0, 372, "(lidx+(gidx*2))")
 
-  def test_div_mod_recombine_partial(self):
-    gidx = Variable("gidx", 0, 15)
-    self.helper_test_variable((gidx//2)%4+(gidx//8)*4, 0, 7, "gidx//2")
-
   def test_div_mod_recombine_folded_mod(self):
     a = Variable("a", 0, 2)
     b = Variable("b", 0, 100)

tinygrad/codegen/gpudims.py

@@ -47,11 +47,6 @@ def get_grouped_dims(prefix, dims:tuple[sint, ...], max_sizes:tuple[int, ...]|No
     if a == 2 and b == 1: ret = [raw_idxs[0] * limited[1] + raw_idxs[1]]
     if a == 3 and b == 1: ret = [raw_idxs[0] * (limited[1] * limited[2]) + raw_idxs[1] * limited[2] + raw_idxs[2]]
     if a == 3 and b == 2: ret = [raw_idxs[0] * limited[1] + raw_idxs[1], raw_idxs[2]]
-  elif limited != dims:
-    # Convert to 1D
-    flat = raw_idxs[0]*limited[1]+raw_idxs[1] if len(dims) == 2 else raw_idxs[0]*(limited[1]*limited[2])+raw_idxs[1]*limited[2]+raw_idxs[2]
-    # Get back original indices from 1D
-    ret = [flat//dims[1], flat%dims[1]] if len(dims) == 2 else [flat//(dims[2]*dims[1]), (flat//dims[2])%dims[1], flat%dims[2]]
   return ret[::-1] if reverse else ret
 
 def add_gpudims(ctx:Renderer, s:UOp):

tinygrad/codegen/late/linearizer.py

@@ -1,59 +1,40 @@
 import heapq
-from typing import Any
 from collections import defaultdict
-from tinygrad.uop.ops import PatternMatcher, UOp, Ops, UPat, multirange_str
-from tinygrad.helpers import prod, getenv, TUPLE_ORDER
+from tinygrad.uop.ops import PatternMatcher, UOp, Ops, UPat
+from tinygrad.helpers import prod
 
-def linearize(sink:UOp) -> list[UOp]:
+def linearize(u:UOp) -> list[UOp]:
   # this is a toposort with priority
-  lst = list(sink.toposort())
+  lst = list(u.toposort())
   consumers: defaultdict[UOp, list[UOp]] = defaultdict(list)
   in_degree:dict[UOp, int] = {}
-  out_degree:dict[UOp, int] = {}
-  priorities:dict[UOp, tuple[int, int, Any]] = {}
+  priorities:dict[UOp, tuple[int, int]] = {}
 
   # get consumers and assign priorities
   # NOTE: this requires the lst be locally toposorted
   for u in reversed(lst):
     for s in u.src: consumers[s].append(u)
     in_degree[u] = len(u.src)
-    out_degree[u] = len(consumers[u])
 
     # we place UOps with higher run_counts later
+    # this will cause ranges to be placed late and ends to be placed early
     run_count = prod([int(r.vmax)+1 for r in u.ranges])
 
-    # simple priority override. this is all bottom up now, smaller numbers will be closer to the top
-    extra = None
-    match u.op:
-      # the order and placement of these defines is important
-      case Ops.DEFINE_GLOBAL: priority, extra = -20, u.arg
-      case Ops.DEFINE_VAR: priority, extra = -19, u.arg
-      case Ops.DEFINE_LOCAL: priority = -18
-      case Ops.DEFINE_REG: priority = -17
-      case Ops.CONST: priority = -10  # early consts
-      case Ops.LOAD: priority = -1    # place loads early
-      case Ops.STORE: priority = 1    # place stores late
-      case Ops.RANGE: priority = 5    # placing RANGE is good
-      case Ops.END: priority = -5     # placing END is bad
-      case _: priority = 0            # everything else has priority 0
-    priorities[u] = (run_count, priority, extra)
+    # simple priority
+    priorities[u] = (run_count, 0)
 
   # number the uops in "ideal" order
-  nkey = {u:i for i,u in enumerate(sorted(lst, key=lambda x: priorities[x]+(x.tuplize if TUPLE_ORDER else ())))}
+  nkey = {u:i for i,u in enumerate(sorted(lst, key=lambda x: (priorities[x],)+x.tuplize))}
 
   # then force then to be toposorted in as close to the ideal order as possible
-  heap = [(-nkey[sink], sink)]
+  heapq.heapify(heap:=[(nkey[u],u) for u in lst if in_degree[u] == 0])
   newlst = []
   while heap:
     newlst.append(u:=heapq.heappop(heap)[1])
-    for v in u.src:
-      out_degree[v] -= 1
-      if out_degree[v] == 0: heapq.heappush(heap, (-nkey[v],v))
-  newlst = newlst[::-1]
-
-  if getenv("DEBUG_LINEARIZE"):
-    for i,u in enumerate(newlst):
-      print(f"{i:4d} {str(u.op):20s} {multirange_str(u.ranges, color=True, pad=10)} {priorities[u]}")
+    for v in consumers[u]:
+      in_degree[v] -= 1
+      if in_degree[v] == 0: heapq.heappush(heap, (nkey[v],v))
+  assert len(newlst) == len(lst), f"len mismatch {len(newlst)} != {len(lst)}"
   return newlst
 
 class CFGContext:

tinygrad/codegen/opt/heuristic.py

@@ -107,7 +107,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
   # potentially do more upcasts of non reduce axes based on a heuristic
   is_dsp = k.ren is not None and k.ren.device == "DSP"
   upcasted_axis: set[int] = set()
-  while resolve(prod(k.output_shape[i] for i in k.upcastable_dims) >= 1024) and (k.upcast_size() < 32):
+  while resolve(prod(k.output_shape[i] for i in k.upcastable_dims) >= 1024):
     xb_choices = []
     # consider all upcastable axes with 3 or 4 upcast (128 on the DSP)
     for axis, upcast_amount in itertools.product(k.upcastable_dims, ([128] if not len(upcasted_axis) else []) if is_dsp else [3,4]):
@@ -135,7 +135,8 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
   # if last reduce dim is small(ish), loop unroll the reduce
   # NOTE: this can fail on multireduce with mismatching dimensions, this is okay
   try:
-    if k.unrollable_dims and (k.upcast_size() <= 4 or not k.axes_of(AxisType.UNROLL)) and (k.upcast_size() < 64):
+    upcast_size = prod(k.full_shape[a] for a in k.axes_of(AxisType.UPCAST, AxisType.UNROLL))
+    if k.unrollable_dims and (upcast_size <= 4 or not k.axes_of(AxisType.UNROLL)) and (upcast_size < 64):
       if (s:=k.full_shape[k.unrollable_dims[-1]]) <= 32:
         k.apply_opt(Opt(OptOps.UNROLL, len(k.unrollable_dims)-1, 0))
         # if it's small, upcast a second reduce dimension too

tinygrad/codegen/opt/postrange.py

@@ -105,8 +105,6 @@ class Scheduler:
   def ranges_of(self, *axis_type:AxisType) -> list[UOp]: return [r for r in self.rngs if r.arg[-1] in axis_type]
   def axes_of(self, *axis_type:AxisType) -> list[int]: return [i for i,t in enumerate(self.axis_types) if t in axis_type]
 
-  def upcast_size(self) -> int: return prod(self.full_shape[a] for a in self.axes_of(AxisType.UPCAST, AxisType.UNROLL))
-
   # copied from kernel.py
   @property
   def upcastable_dims(self) -> list[int]: return [i for i in self.axes_of(AxisType.GLOBAL, AxisType.LOCAL, AxisType.LOOP) \

tinygrad/device.py

@@ -5,7 +5,7 @@ from typing import Any, Generic, TypeVar, Iterator, Sequence, cast, Generator
 import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re, atexit, pickle, decimal
 from tinygrad.helpers import CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored, CPU_LLVM
 from tinygrad.helpers import Context, DISABLE_COMPILER_CACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, dedup
-from tinygrad.helpers import unwrap_class_type, suppress_finalizing, AMD_LLVM, select_first_inited
+from tinygrad.helpers import unwrap_class_type, suppress_finalizing, AMD_LLVM
 from tinygrad.dtype import DType, ImageDType, PtrDType, dtypes, _to_np_dtype
 from tinygrad.renderer import Renderer
 
@@ -291,8 +291,8 @@ class Compiled:
     if len(enable_comps) > 1: raise RuntimeError(f"{self.device}: multiple compilers set in env {enable_comps}")
     for _, comp_pair in disable_comps: self.compilers.remove(comp_pair)
 
-    self.renderer, self.compiler = select_first_inited([list(enable_comps)[0][1]] if len(enable_comps) == 1 else self.compilers,
-                                                       f"No compiler for {self.device} is available")
+    try: self.renderer, self.compiler = next(self._get_available_compilers([list(enable_comps)[0][1]] if len(enable_comps) == 1 else self.compilers))
+    except StopIteration as exc: raise RuntimeError(f"no usable compilers for {self.device}") from exc
 
     if DEBUG >= 1: print(f"{self.device}: using {self.compiler.__class__.__name__}")
 
@@ -300,6 +300,10 @@ class Compiled:
     compiler_name = f"{unwrap_class_type(c).__name__.upper().removesuffix('COMPILER').removeprefix(devname:=self.device.split(':')[0].upper())}"
     return f"{devname}_{compiler_name if len(compiler_name) > 0 else unwrap_class_type(c).__name__.upper()}"
 
+  def _get_available_compilers(self, compilers) -> Iterator[tuple[Renderer, Compiler]]:
+    for renderer, compiler in compilers:
+      with contextlib.suppress(Exception): yield renderer(), compiler()
+
   def synchronize(self):
     """
     Synchronize all pending operations on the device.
@@ -339,7 +343,6 @@ def is_dtype_supported(dtype:DType, device:str|None=None) -> bool:
   # PYTHON supports half memoryview in 3.12+ https://github.com/python/cpython/issues/90751
   if dtype == dtypes.half:
     if device == "CL": return not CI and not OSX
-    if device == "QCOM": return False # QCOM compiler is flaky with half
     if device in ["CUDA", "NV"]: return not CI
     if device == "CPU" and CPU_LLVM: return OSX
     if device == "PYTHON": return sys.version_info >= (3, 12)

tinygrad/helpers.py

@@ -114,13 +114,6 @@ def suppress_finalizing(func):
       if not getattr(sys, 'is_finalizing', lambda: True)(): raise # re-raise if not finalizing
   return wrapper
 
-def select_first_inited(candidates:Sequence[Callable[...,T]|Sequence[Callable[...,T]]], err_msg: str) -> tuple[T,...]|T:
-  excs = []
-  for typ in candidates:
-    try: return tuple([cast(Callable, t)() for t in typ]) if isinstance(typ, Sequence) else cast(Callable, typ)()
-    except Exception as e: excs.append(e)
-  raise ExceptionGroup(err_msg, excs)
-
 def unwrap_class_type(cls_t): return cls_t.func if isinstance(cls_t, functools.partial) else cls_t
 
 def pluralize(st:str, cnt:int): return f"{cnt} {st}"+('' if cnt == 1 else 's')
@@ -186,8 +179,6 @@ SPEC = ContextVar("SPEC", 1)
 IGNORE_OOB = ContextVar("IGNORE_OOB", 1)
 PCONTIG = ContextVar("PCONTIG", 0)  # partial contiguous in rangeify
 DEBUG_RANGEIFY = ContextVar("DEBUG_RANGEIFY", 0)
-# set to 1, this uses tuplize in the linearizer sort order
-TUPLE_ORDER = ContextVar("TUPLE_ORDER", 1)
 
 @dataclass(frozen=True)
 class Metadata:

tinygrad/mixin/movement.py

@@ -3,19 +3,15 @@ import functools
 from typing import TypeAlias, TYPE_CHECKING, Self
 from tinygrad.uop import Ops
 from tinygrad.helpers import prod, argfix, flatten, dedup
-if TYPE_CHECKING: from tinygrad.uop.ops import UOp
-sint: TypeAlias = "UOp | int"
-
-def _align_left(*shapes:tuple[sint, ...]) -> tuple[tuple[sint, ...], ...]:
-  # unsqueeze left to make every shape same length
-  max_dim = max(len(shape) for shape in shapes)
-  return tuple((1,) * (max_dim - len(shape)) + shape for shape in shapes)
+if TYPE_CHECKING:
+  from tinygrad.uop.ops import UOp
+  sint:TypeAlias = UOp|int
 
 class MovementMixin:
   # required to implement
   def _mop(self, op:Ops, arg) -> Self: raise NotImplementedError
   @property
-  def shape(self) -> tuple[sint, ...]: raise NotImplementedError
+  def shape(self) -> tuple["sint", ...]: raise NotImplementedError
 
   # great functions you get!
   @property
@@ -30,7 +26,7 @@ class MovementMixin:
     """
     return len(self.shape)
 
-  def numel(self) -> sint:
+  def numel(self) -> "sint":
     """
     Returns the total number of elements in the tensor.
 
@@ -46,33 +42,6 @@ class MovementMixin:
     if not -max(1, total) <= dim <= max(1, total)-1: raise IndexError(f"{dim=} out of range {[-max(1, total), max(1, total)-1]}")
     return dim + total if dim < 0 else dim
 
-  def _broadcast_to(self, new_shape:tuple[sint, ...]) -> Self:
-    if self.shape == new_shape: return self
-    if self.ndim > len(new_shape): raise ValueError(f"cannot broadcast tensor to fewer dimensions. shape={self.shape} to {new_shape=}")
-    # first unsqueeze left with 1s https://data-apis.org/array-api/latest/API_specification/broadcasting.html
-    shape, _ = _align_left(self.shape, new_shape)
-    # for each dimension, check either dim is 1, or it does not change
-    if not all(s == ns or s == 1 for s,ns in zip(shape, new_shape)):
-      raise ValueError(f"cannot broadcast {self.shape} to {new_shape=}")
-    reshaped = self.reshape(shape)
-    ret = reshaped._mop(Ops.EXPAND, arg=new_shape)
-    return reshaped if ret.shape == reshaped.shape else ret
-
-  def expand(self, shape, *args) -> Self:
-    """
-    Returns a tensor that is expanded to the shape that is specified.
-    Expand can also increase the number of dimensions that a tensor has.
-
-    Passing a `-1` or `None` to a dimension means that its size will not be changed.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    t = Tensor([1, 2, 3])
-    print(t.expand(4, -1).numpy())
-    ```
-    """
-    new_shape = tuple(from_ if to == -1 or to is None else to for from_, to in zip(*(_align_left(self.shape, argfix(shape, *args)))))
-    return self._broadcast_to(new_shape)
-
   def reshape(self, shape, *args) -> Self:
     """
     Returns a tensor with the same data as the original tensor but with a different shape.
@@ -92,7 +61,7 @@ class MovementMixin:
     ret = self._mop(Ops.RESHAPE, arg=new_shape)
     return self if ret.shape == self.shape else ret
 
-  def shrink(self, arg:tuple[tuple[sint, sint]|None, ...]) -> Self:
+  def shrink(self, arg:tuple[tuple["sint", "sint"]|None, ...]) -> Self:
     """
     Returns a tensor that shrinks the each axis based on input arg.
     `arg` must have the same length as `self.ndim`.
@@ -155,9 +124,6 @@ class MovementMixin:
 
   # **** high level ****
 
-  def shrink_to(self, shape, *args) -> Self:
-    return self.shrink(tuple([None if ns is None else (0, ns) for ns in argfix(shape, *args)]))
-
   def view(self, shape, *args) -> Self:
     """`.view` is an alias for `.reshape`."""
     return self.reshape(shape, *args)
@@ -291,38 +257,3 @@ class MovementMixin:
     for i, name in enumerate(lhs): assert (name not in sizes) or sizes[name] == t.shape[i], f"size provided for dimension {name} incorrect"
     t = t.permute([lhs.index(name) for name in rhs])
     return functools.reduce(lambda x, dims: x.flatten(dims[0], dims[1] - 1) if dims[0]<dims[1] else x.unsqueeze(dims[0]), reversed(flatten_dims), t)
-
-  # *** movement ops with expand ***
-
-  def repeat_interleave(self, repeats:int, dim:int|None=None) -> Self:
-    """
-    Repeats elements of a tensor.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    t = Tensor([1, 2, 3])
-    print(t.repeat_interleave(2).numpy())
-    ```
-    """
-    x, dim = (self.flatten(), 0) if dim is None else (self, self._resolve_dim(dim))
-    shp = x.shape
-    return x.reshape(*shp[:dim+1], 1, *shp[dim+1:]).expand(*shp[:dim+1], repeats, *shp[dim+1:]).reshape(*shp[:dim], shp[dim]*repeats, *shp[dim+1:])
-
-  def repeat(self, repeats, *args) -> Self:
-    """
-    Repeats tensor number of times along each dimension specified by `repeats`.
-    `repeats` can be passed as a tuple or as separate arguments.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    t = Tensor([1, 2, 3])
-    print(t.repeat(4, 2).numpy())
-    ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(t.repeat(4, 2, 1).shape)
-    ```
-    """
-    repeats = argfix(repeats, *args)
-    base_shape = _align_left(self.shape, repeats)[0]
-    unsqueezed_shape = flatten([[1, s] for s in base_shape])
-    expanded_shape = flatten([[r, s] for r,s in zip(repeats, base_shape)])
-    final_shape = [r*s for r,s in zip(repeats, base_shape)]
-    return self.reshape(unsqueezed_shape).expand(expanded_shape).reshape(final_shape)

tinygrad/runtime/support/hcq.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
 from typing import cast, Callable, Type, TypeVar, Generic, Any, Sequence
-import contextlib, decimal, statistics, time, ctypes, array, os, struct, collections, functools
+import contextlib, decimal, statistics, time, ctypes, array, os, struct, traceback, collections
 try: import fcntl # windows misses that
 except ImportError: fcntl = None #type:ignore[assignment]
-from tinygrad.helpers import PROFILE, getenv, to_mv, ProfileRangeEvent, select_first_inited
+from tinygrad.helpers import PROFILE, getenv, to_mv, ProfileRangeEvent
 from tinygrad.device import BufferSpec, Compiled, LRUAllocator, ProfileDeviceEvent, ProfileProgramEvent, CompilerPairT
 from tinygrad.uop.ops import sym_infer, sint, UOp
 from tinygrad.runtime.autogen import libc
@@ -437,10 +437,19 @@ class HCQCompiled(Compiled, Generic[SignalType]):
     except MemoryError: buf, realloced = self.allocator.alloc(oldbuf.size if oldbuf is not None else new_size, options=options), False
     return buf, realloced
 
+  def _make_no_iface_error(self, errs:str, err_short:str) -> RuntimeError:
+    # Keep it in a separate function to avoid creating a traceback <-> locals ref cycle
+    e = RuntimeError(f"No interface for {type(self).__name__[:-6]}:{self.device_id} is available")
+    if hasattr(e, "add_note"): e.add_note(errs + err_short)
+    return e
+
   def _select_iface(self, *ifaces:Type):
+    errs, err_short = "", ""
     if val:=getenv(f'{type(self).__name__[:-6].upper()}_IFACE', ""): ifaces = tuple(x for x in ifaces if x.__name__.startswith(val.upper()))
-    return select_first_inited([functools.partial(cast(Callable, iface), self, self.device_id) for iface in ifaces],
-                               f"No interface for {type(self).__name__[:-6]}:{self.device_id} is available")
+    for iface_t in ifaces:
+      try: return iface_t(self, self.device_id)
+      except Exception as e: errs, err_short = errs + f"\n{iface_t.__name__}: {traceback.format_exc()}", err_short + f"\n{iface_t.__name__}: {e}."
+    raise self._make_no_iface_error(errs, err_short)
 
   def _is_cpu(self) -> bool: return hasattr(self, 'device') and self.device.split(":")[0] == "CPU"
 

tinygrad/tensor.py

@@ -10,7 +10,6 @@ from tinygrad.helpers import IMAGE, WINO, Metadata, TRACEMETA, ceildiv, fetch, p
 from tinygrad.helpers import suppress_finalizing
 from tinygrad.gradient import compute_gradient
 from tinygrad.mixin import OpMixin
-from tinygrad.mixin.movement import _align_left
 from tinygrad.uop.ops import smax, smin, resolve, UOp, Ops, sint, identity_element, all_metadata, _index_to_concrete_int, sint_to_uop
 from tinygrad.uop.spec import type_verify, tensor_spec
 from tinygrad.device import Device, Buffer
@@ -80,6 +79,10 @@ def _apply_winograd_matrix(mat, t:Tensor, dims:int) -> Tensor:
   assert isinstance(ret, Tensor), "sum didn't return a Tensor"
   return ret
 
+def _align_left(*shapes:tuple[sint, ...]) -> tuple[tuple[sint, ...], ...]:
+  # unsqueeze left to make every shape same length
+  max_dim = max(len(shape) for shape in shapes)
+  return tuple((1,) * (max_dim - len(shape)) + shape for shape in shapes)
 def _broadcast_shape(*shapes:tuple[sint, ...]) -> tuple[sint, ...]:
   return tuple(0 if 0 in nth_dim_sizes else smax(nth_dim_sizes) for nth_dim_sizes in zip(*_align_left(*shapes)))
 
@@ -1037,6 +1040,21 @@ class Tensor(OpMixin):
 
   def _mop(self, op:Ops, arg) -> Tensor: return self._apply_uop(UOp._mop, extra_args=(op,), arg=arg)
 
+  def expand(self, shape, *args) -> Tensor:
+    """
+    Returns a tensor that is expanded to the shape that is specified.
+    Expand can also increase the number of dimensions that a tensor has.
+
+    Passing a `-1` or `None` to a dimension means that its size will not be changed.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    t = Tensor([1, 2, 3])
+    print(t.expand(4, -1).numpy())
+    ```
+    """
+    new_shape = tuple(from_ if to == -1 or to is None else to for from_, to in zip(*(_align_left(self.shape, argfix(shape, *args)))))
+    return self._broadcast_to(new_shape)
+
   def pad(self, padding:Sequence[sint]|Sequence[tuple[sint, sint]|None], mode:str="constant", value:float=0.0) -> Tensor:
     """
     Returns a tensor with padding applied based on the input `padding`.
@@ -1104,6 +1122,8 @@ class Tensor(OpMixin):
   def pad_to(self, shape, *args):
     if len(new_shape := argfix(shape, *args)) != self.ndim: raise ValueError(f"dim mismatch, cannot pad {self.shape} to {new_shape}")
     return self.pad(tuple([None if ns is None else (0, ns-s) for s,ns in zip(self.shape, new_shape)]))
+  def shrink_to(self, shape, *args):
+    return self.shrink(tuple([None if ns is None else (0, ns) for ns in argfix(shape, *args)]))
 
   # ***** movement high level ops *****
 
@@ -1323,6 +1343,39 @@ class Tensor(OpMixin):
     # checks for shapes and number of dimensions delegated to cat
     return Tensor.cat(*[t.unsqueeze(dim) for t in argfix(self, *args)], dim=dim)
 
+  def repeat_interleave(self, repeats:int, dim:int|None=None) -> Tensor:
+    """
+    Repeats elements of a tensor.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    t = Tensor([1, 2, 3])
+    print(t.repeat_interleave(2).numpy())
+    ```
+    """
+    x, dim = (self.flatten(), 0) if dim is None else (self, self._resolve_dim(dim))
+    shp = x.shape
+    return x.reshape(*shp[:dim+1], 1, *shp[dim+1:]).expand(*shp[:dim+1], repeats, *shp[dim+1:]).reshape(*shp[:dim], shp[dim]*repeats, *shp[dim+1:])
+
+  def repeat(self, repeats, *args) -> Tensor:
+    """
+    Repeats tensor number of times along each dimension specified by `repeats`.
+    `repeats` can be passed as a tuple or as separate arguments.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    t = Tensor([1, 2, 3])
+    print(t.repeat(4, 2).numpy())
+    ```
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(t.repeat(4, 2, 1).shape)
+    ```
+    """
+    repeats = argfix(repeats, *args)
+    base_shape = _align_left(self.shape, repeats)[0]
+    unsqueezed_shape = flatten([[1, s] for s in base_shape])
+    expanded_shape = flatten([[r, s] for r,s in zip(repeats, base_shape)])
+    final_shape = [r*s for r,s in zip(repeats, base_shape)]
+    return self.reshape(unsqueezed_shape).expand(expanded_shape).reshape(final_shape)
+
   def split(self, sizes:int|Sequence[int], dim:int=0) -> tuple[Tensor, ...]:
     """
     Splits the tensor into chunks along the dimension specified by `dim`.
@@ -2100,22 +2153,22 @@ class Tensor(OpMixin):
     noop, i_ = [None] * (self.ndim-len(k_)), self.shape[-len(k_):]
     assert all(resolve(d*(k-1)+1 <= i) for k,d,i in zip(k_,d_,i_)), "kernel size cannot be greater than actual input size"
     o_ = [ceildiv(i-d*(k-1), s) for i,d,k,s in zip(i_,d_,k_,s_)]
-    if getenv("ONE_POOL") or any(resolve(k > s) for k,s in zip(k_,s_)) or any(d != 1 for d in d_):
+    if any(resolve(k > s) for k,s in zip(k_,s_)) or any(d != 1 for d in d_):
       # input size scaling factor to make sure shrink for stride is possible
-      f_ = [smax(1, ceildiv(o*s - d, i)) for o,s,i,d in zip(o_,s_,i_,d_)]
-      # repeats such that we don't need padding
+      f_ = [1 + int(resolve(o*s > (i - d*(k-1)))) for o,s,i,d,k in zip(o_,s_,i_,d_,k_)]
+      # # repeats such that we don't need padding
       x = self.repeat([1]*len(noop) + [ceildiv(k*(i*f+d),i) for k,i,d,f in zip(k_,i_,d_,f_)])
       # handle dilation
-      x = x.shrink_to(noop + [k*(i*f+d) for k,i,d,f in zip(k_,i_,d_,f_)]).reshape(noop + flatten((k,(i*f+d)) for k,i,d,f in zip(k_,i_,d_,f_)))
+      x = x.shrink(tuple(noop + [(0,k*(i*f+d)) for k,i,d,f in zip(k_,i_,d_,f_)])).reshape(noop + flatten((k,(i*f+d)) for k,i,d,f in zip(k_,i_,d_,f_)))
       # handle stride
-      x = x.shrink_to(noop + flatten((k,o*s) for k,o,s in zip(k_,o_,s_))).reshape(noop + flatten((k,o,s) for k,o,s in zip(k_,o_,s_)))
-      x = x.shrink_to(noop + flatten((k,o,1) for k,o in zip(k_,o_))).reshape(noop + flatten((k,o) for k,o in zip(k_,o_)))
+      x = x.shrink(tuple(noop + flatten(((0,k), (0,o*s)) for k,o,s in zip(k_,o_,s_)))).reshape(noop + flatten((k,o,s) for k,o,s in zip(k_,o_,s_)))
+      x = x.shrink(tuple(noop + flatten(((0,k), (0,o), (0,1)) for k,o in zip(k_,o_)))).reshape(noop + flatten((k,o) for k,o in zip(k_,o_)))
       # permute to move reduce to the end
       return x.permute(*range(len(noop)), *[len(noop)+i*2+1 for i in range(len(i_))], *[len(noop)+i*2 for i in range(len(i_))])
     # TODO: once the shapetracker can optimize well, remove this alternative implementation
     x = self.pad(tuple(noop + [(0, max(0,o*s-i)) for i,o,s in zip(i_,o_,s_)])).shrink(tuple(noop + [(0,o*s) for o,s in zip(o_,s_)]))
     x = x.reshape(noop + flatten(((o,s) for o,s in zip(o_,s_))))
-    x = x.shrink_to(noop + flatten((o,k) for o,k in zip(o_,k_)))
+    x = x.shrink(tuple(noop + flatten(((0,o), (0,k)) for o,k in zip(o_,k_))))
     return x.permute(*range(len(noop)), *[len(noop)+i*2 for i in range(len(i_))], *[len(noop)+i*2+1 for i in range(len(i_))])
 
   def _resolve_pool_pads(self, padding:int|Sequence[int], dims:int) -> Sequence[int]:
@@ -3354,8 +3407,18 @@ class Tensor(OpMixin):
     return self / (1 + self.abs())
 
   # ***** broadcasted elementwise ops *****
+  def _broadcast_to(self, new_shape:tuple[sint, ...]) -> Tensor:
+    if self.shape == new_shape: return self
+    if self.ndim > len(new_shape): raise ValueError(f"cannot broadcast tensor to fewer dimensions. shape={self.shape} to {new_shape=}")
+    # first unsqueeze left with 1s https://data-apis.org/array-api/latest/API_specification/broadcasting.html
+    shape, _ = _align_left(self.shape, new_shape)
+    # for each dimension, check either dim is 1, or it does not change
+    if not all(resolve(s == ns) or resolve(s == 1) for s,ns in zip(shape, new_shape)):
+      raise ValueError(f"cannot broadcast {self.shape} to {new_shape=}")
+    # NOTE: this cast is no-op in forward and uses sum_acc_dtype in the backward sum
+    return self.reshape(shape).cast(sum_acc_dtype(self.dtype))._apply_uop(UOp.expand, arg=new_shape).cast(self.dtype)
 
-  def _broadcasted(self, y:Tensor|ConstType|UOp, reverse:bool=False, match_dtype:bool=True, backward_cast:bool=True) -> tuple[Tensor, Tensor]:
+  def _broadcasted(self, y:Tensor|ConstType|UOp, reverse:bool=False, match_dtype:bool=True) -> tuple[Tensor, Tensor]:
     x: Tensor = self
     if not isinstance(y, Tensor):
       # make y a Tensor
@@ -3371,13 +3434,8 @@ class Tensor(OpMixin):
 
     if reverse: x, y = y, x
 
-    # compute the output shape
-    out_shape = _broadcast_shape(x.shape, y.shape)
-
     # broadcast
-    # NOTE: the backward cast is no-op in forward and uses sum_acc_dtype in the backward sum
-    return x.cast(sum_acc_dtype(x.dtype) if backward_cast else x.dtype)._broadcast_to(out_shape).cast(x.dtype), \
-           y.cast(sum_acc_dtype(y.dtype) if backward_cast else y.dtype)._broadcast_to(out_shape).cast(y.dtype)
+    return x._broadcast_to(out_shape:=_broadcast_shape(x.shape, y.shape)), y._broadcast_to(out_shape)
 
   def sub(self, x:Tensor|ConstType, reverse=False) -> Tensor:
     """

tinygrad/uop/__init__.py

@@ -1,5 +1,3 @@
-# flake8: noqa: E702
-# allow semicolons to put multiple ops on one line
 from enum import auto, IntEnum, Enum
 
 # wrapper around IntEnum that preserves Enum.__str__ and makes auto() unique across all FastEnum subclasses
@@ -11,13 +9,16 @@ class FastEnum(IntEnum):
 
 # the order of these Ops controls the order of the toposort
 class Ops(FastEnum):
-  # ** 1 -- defines/special **
+  # ** 1 -- defines/consts **
 
-  # TODO: unify these ops into the levels of the memory hierarchy
-  DEFINE_GLOBAL = auto(); DEFINE_LOCAL = auto(); DEFINE_REG = auto()
+  # TODO: unify these ops into the levels of the memory hierarchy. depends on ASSIGN is STORE
+  DEFINE_GLOBAL = auto(); DEFINE_LOCAL = auto(); DEFINE_REG = auto() # noqa: E702
 
   # this is for symbolic shapes
-  DEFINE_VAR = auto(); BIND = auto()
+  DEFINE_VAR = auto(); BIND = auto() # noqa: E702
+
+  # consts. VCONST is a vectorized const
+  VCONST = auto(); CONST = auto() # noqa: E702
 
   # this is a RANGE for GPU dimensions, similar to symbolic shapes but not exactly
   SPECIAL = auto()
@@ -25,7 +26,8 @@ class Ops(FastEnum):
   # ** 2 -- non op uops **
 
   # uops that aren't rendered
-  NOOP = auto(); SINK = auto(); PRECAST = auto()
+  NOOP = auto(); SINK = auto(); UNIQUE = auto(); DEVICE = auto(); KERNEL = auto(); PRECAST = auto(); REWRITE_ERROR = auto()  # noqa: E702
+  SENTINEL = auto()
 
   # AFTER passes src[0] through and promises in the toposort that any consumers of the AFTER run after src[1:]
   AFTER = auto()
@@ -33,8 +35,24 @@ class Ops(FastEnum):
   # GROUP is a NOOP that just merges things together
   GROUP = auto()
 
-  # vector creation / item selection
-  GEP = auto(); VECTORIZE = auto()
+  # buffer ops
+  COPY = auto(); BUFFER = auto(); BUFFER_VIEW = auto(); MSELECT = auto(); MSTACK = auto() # noqa: E702
+
+  # create buffer
+  BUFFERIZE = auto()
+
+  # ops that adjust the behavior of the scheduler
+  CONTIGUOUS = auto(); CONTIGUOUS_BACKWARD = auto(); DETACH = auto() # noqa: E702
+
+  # movement ops! these only exist in the tensor graph
+  RESHAPE = auto(); PERMUTE = auto(); EXPAND = auto(); PAD = auto(); SHRINK = auto(); FLIP = auto() # noqa: E702
+  MULTI = auto()  # MULTI is really a movement op
+
+  # reduce (movement)
+  REDUCE_AXIS = auto(); REDUCE = auto(); ALLREDUCE = auto() # noqa: E702
+
+  # optimization helper ops
+  UNROLL = auto(); CONTRACT = auto(); GEP = auto(); VECTORIZE = auto(); CAT = auto(); PTRCAT = auto() # noqa: E702
 
   # ** 3 -- load/store **
 
@@ -42,7 +60,8 @@ class Ops(FastEnum):
   INDEX = auto()
 
   # load/store before math
-  LOAD = auto(); STORE = auto()
+  LOAD = auto(); STORE = auto() # noqa: E702
+  ASSIGN = auto()  # TODO: ASSIGN is STORE, remove ASSIGN
 
   # ** 4 -- math **
 
@@ -50,53 +69,24 @@ class Ops(FastEnum):
   WMMA = auto()
 
   # UnaryOps
-  CAST = auto(); BITCAST = auto(); EXP2 = auto(); LOG2 = auto(); SIN = auto()
-  SQRT = auto(); RECIPROCAL = auto(); NEG = auto(); TRUNC = auto()
+  CAST = auto(); BITCAST = auto(); EXP2 = auto(); LOG2 = auto(); SIN = auto(); SQRT = auto(); RECIPROCAL = auto(); NEG = auto(); TRUNC = auto() # noqa: E702
 
   # BinaryOps
-  ADD = auto(); MUL = auto(); SHL = auto(); SHR = auto(); IDIV = auto(); MAX = auto(); MOD = auto()
-  CMPLT = auto(); CMPNE = auto(); CMPEQ = auto()
-  XOR = auto(); OR = auto(); AND = auto()
-  THREEFRY = auto(); SUB = auto(); FDIV = auto(); POW = auto()
+  ADD = auto(); MUL = auto(); SHL = auto(); SHR = auto(); IDIV = auto(); MAX = auto(); MOD = auto() # noqa: E702
+  CMPLT = auto(); CMPNE = auto(); CMPEQ = auto() # noqa: E702
+  XOR = auto(); OR = auto(); AND = auto() # noqa: E702
+  THREEFRY = auto(); SUB = auto(); FDIV = auto(); POW = auto() # noqa: E702
 
   # TernaryOps
-  WHERE = auto(); MULACC = auto()
+  WHERE = auto(); MULACC = auto() # noqa: E702
 
-  # ** 5 -- control flow / consts / custom **
+  # ** 5 -- control flow / other **
 
   # control flow ops
-  BARRIER = auto(); RANGE = auto(); IF = auto(); END = auto(); ENDIF = auto()
-
-  # consts. VCONST is a vectorized const
-  VCONST = auto(); CONST = auto()
+  BARRIER = auto(); RANGE = auto(); IF = auto(); END = auto(); ENDIF = auto() # noqa: E702
 
   # CUSTOM/CUSTOMI are used to output strings into codegen. the I makes the string inline
-  CUSTOM = auto(); CUSTOMI = auto()
-
-  # ** 6 -- ops that don't exist in programs **
-
-  # tensor graph ops
-  UNIQUE = auto(); DEVICE = auto(); KERNEL = auto()
-  ASSIGN = auto()
-
-  # buffer ops
-  BUFFERIZE = auto(); COPY = auto(); BUFFER = auto(); BUFFER_VIEW = auto(); MSELECT = auto(); MSTACK = auto()
-
-  # ops that adjust the behavior of the scheduler
-  CONTIGUOUS = auto(); CONTIGUOUS_BACKWARD = auto(); DETACH = auto()
-
-  # movement ops! these only exist in the tensor graph
-  RESHAPE = auto(); PERMUTE = auto(); EXPAND = auto(); PAD = auto(); SHRINK = auto(); FLIP = auto()
-  MULTI = auto()  # MULTI is really a movement op
-
-  # reduce
-  REDUCE_AXIS = auto(); REDUCE = auto(); ALLREDUCE = auto()
-
-  # errors/placeholders
-  REWRITE_ERROR = auto(); SENTINEL = auto()
-
-  # expander ops
-  UNROLL = auto(); CONTRACT = auto(); CAT = auto(); PTRCAT = auto()
+  CUSTOM = auto(); CUSTOMI = auto() # noqa: E702
 
 class GroupOp:
   Unary = {Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.SQRT, Ops.RECIPROCAL, Ops.NEG, Ops.TRUNC}

tinygrad/uop/ops.py

@@ -48,11 +48,6 @@ def range_str(u:UOp, color=False) -> str:
   ret = '_'.join([str(x) if x >= 0 else "m"+str(-x) for x in u.arg[0:-1]])
   return colored(ret, axis_colors[u.arg[-1]]) if color else ret
 
-def multirange_str(rngs:Iterable[UOp], color=False, pad=None) -> str:
-  ret = ','.join([range_str(x, color=color) for x in sorted(rngs, key=lambda x: x.arg)])
-  if pad is not None: ret += " " * (pad-ansilen(ret))
-  return ret
-
 def consumer_map_from_toposort(lst:Iterable[UOp]):
   ret: dict[UOp, dict[UOp, None]] = {}
   for u in lst:
@@ -559,7 +554,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
   # in these four, if the shape doesn't change we can return self
   def forced_reshape(self, arg:tuple[sint, ...]): return self._mop(Ops.RESHAPE, arg, same_shape_noop=False)
   #def reshape(self, arg:tuple[sint, ...]): return self._mop(Ops.RESHAPE, arg, same_shape_noop=True)
-  #def expand(self, arg:tuple[sint, ...]): return self._mop(Ops.EXPAND, arg, same_shape_noop=True)
+  def expand(self, arg:tuple[sint, ...]): return self._mop(Ops.EXPAND, arg, same_shape_noop=True)
   #def shrink(self, arg:tuple[tuple[sint, sint], ...]): return self._mop(Ops.SHRINK, arg, same_shape_noop=True)
   def pad(self, arg:tuple[tuple[sint, sint], ...]): return self._mop(Ops.PAD, arg, same_shape_noop=True)
 
@@ -789,6 +784,8 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
 
   # *** uop high level syntactic sugar ***
 
+  def shrink_to(self, arg:tuple[sint, ...]): return self.shrink(tuple([(0,x) for x in arg]))
+
   @staticmethod
   def placeholder(shape:tuple[int, ...], dtype:DType, slot:int, addrspace=AddrSpace.GLOBAL):
     lookup = {AddrSpace.GLOBAL: Ops.DEFINE_GLOBAL, AddrSpace.LOCAL: Ops.DEFINE_LOCAL, AddrSpace.REG: Ops.DEFINE_REG}
@@ -854,10 +851,10 @@ def exec_alu(op:Ops, dtype:DType, operands, truncate_output=True):
 # ***** uop helpers *****
 
 def print_uops(uops:list[UOp]):
-  uops_index = {u:i for i,u in enumerate(uops)}
   for i,u in enumerate(uops):
-    formatted_srcs = [(uops_index[x] if x.op is not Ops.CONST else f"{x.arg}") if x in uops else "--" for x in u.src]
-    print(f"{i:4d} {str(u.op):20s}: {multirange_str(u.ranges, color=True, pad=10)} {str(u.dtype):40s} " f"{str(formatted_srcs):32s} {u.arg}")
+    formatted_srcs = [(uops.index(x) if x.op is not Ops.CONST else f"{x.arg}") if x in uops else "--" for x in u.src]
+    formatted_range = ','.join([range_str(r, color=True) for r in sorted(u.ranges, key=lambda x: x.arg)])
+    print(f"{i:4d} {str(u.op):20s}: {(formatted_range)+' '*(10-ansilen(formatted_range))} {str(u.dtype):40s} " f"{str(formatted_srcs):32s} {u.arg}")
 
 # ***** pattern matcher *****
 

tinygrad/uop/spec.py

@@ -134,6 +134,10 @@ shared_codegen_spec = PatternMatcher([
   # WMMA has a <a, b, acc>
   (UPat(Ops.WMMA, src=(UPat(), UPat(), UPat()), name="x"), lambda x: isinstance(x.arg, tuple) and len(x.arg) == 8),
 
+  # UNROLL/CONTRACT is used here for WMMA
+  (UPat(Ops.CONTRACT, name="x"), lambda x: x.dtype.count == prod(y[1] for y in x.arg)),
+  (UPat(Ops.UNROLL, name="x"), lambda x: x.src[0].dtype.count == prod(y[1] for y in x.arg)),
+
   # VECTORIZE/GEP
   (UPat(Ops.VECTORIZE, name="x"), lambda x: len(x.src)>1 and len(x.src) == x.dtype.vcount and all(x.dtype == y.dtype.vec(len(x.src)) for y in x.src)),
   (UPat(Ops.GEP, src=(UPat.var("src"),), name="gep"), lambda gep,src: gep.dtype == src.dtype.scalar()),
@@ -162,10 +166,6 @@ kernel_spec = PatternMatcher([
   # index is allowed here
   (UPat(GroupOp.Elementwise|{Ops.CONST, Ops.RANGE, Ops.DEFINE_VAR}, dtype=dtypes.index), lambda: True),
 
-  # UNROLL/CONTRACT is used here for WMMA
-  (UPat(Ops.CONTRACT, name="x"), lambda x: x.dtype.count == prod(y[1] for y in x.arg)),
-  (UPat(Ops.UNROLL, name="x"), lambda x: x.src[0].dtype.count == prod(y[1] for y in x.arg)),
-
   # END can end multiple axes here
   (UPat(Ops.END, src=(UPat(), UPat()), allow_any_len=True, dtype=dtypes.void), lambda: True),
 

tinygrad/uop/symbolic.py

@@ -48,10 +48,8 @@ symbolic_simple = propagate_invalid + PatternMatcher([
   (UPat.var("x") // 1, lambda x: x),   # x//1 -> x
   (UPat.var("x") // -1, lambda x: -x), # x//-1 -> -x
   ((UPat.var() % UPat.var("y")).named("base") % UPat.var("y"), lambda base,y: base),  # (x%y)%y = -> x%y (rewritten with base for speed)
-  # variations of (x%c)+(x//c)*c = x   TODO: add sorting to remove some variations
+  # 4 variations of (x%c)+(x//c)*c = x   TODO: add sorting to remove some variations
   (UPat.var("x")%UPat.cvar("c")+(UPat.var("x")//UPat.cvar("c"))*UPat.cvar("c"), lambda x,c: x), # (x%c)+(x//c)*c = x
-  ((UPat.var("x")//UPat.cvar("a"))%UPat.cvar("c")+(UPat.var("x")//UPat.cvar("b"))*UPat.cvar("c"),
-    lambda x,a,b,c: x//a if a.arg*c.arg==b.arg else None), # ((x//a)%c)+(x//a*c)*c = x//a. Note if a = 1 it degenerates to the one above
   ((UPat.var("x")//UPat.cvar("c1"))*UPat.cvar("c3")+UPat.var("x")%UPat.cvar("c1")*UPat.cvar("c2"),
     lambda x,c1,c2,c3: x*c2 if c1.arg*c2.arg==c3.arg else None), # (x%c1)*c2+(x//c1)*c3 = x*c2 if c1*c2==c3
   ((UPat.var("y")+(UPat.var("x")//UPat.cvar("c"))*UPat.cvar("c"))+UPat.var("x")%UPat.cvar("c"), lambda y,x,c: y+x),

tinygrad/viz/serve.py

@@ -1,14 +1,14 @@
 #!/usr/bin/env python3
 import multiprocessing, pickle, difflib, os, threading, json, time, sys, webbrowser, socket, argparse, socketserver, functools, codecs, io, struct
 import subprocess, ctypes, pathlib, traceback
-from contextlib import redirect_stdout, redirect_stderr
+from contextlib import redirect_stdout
 from decimal import Decimal
 from http.server import BaseHTTPRequestHandler
 from urllib.parse import parse_qs, urlparse
 from typing import Any, TypedDict, TypeVar, Generator, Callable
 from tinygrad.helpers import colored, getenv, tqdm, unwrap, word_wrap, TRACEMETA, ProfileEvent, ProfileRangeEvent, TracingKey, ProfilePointEvent, temp
 from tinygrad.uop.ops import TrackedGraphRewrite, RewriteTrace, UOp, Ops, printable, GroupOp, srender, sint, sym_infer, range_str, pyrender
-from tinygrad.uop.ops import print_uops, range_start, multirange_str
+from tinygrad.uop.ops import print_uops, range_start
 from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry, Device
 from tinygrad.renderer import ProgramSpec
 from tinygrad.dtype import dtypes
@@ -78,14 +78,11 @@ def uop_to_json(x:UOp) -> dict[int, dict]:
         label += f"\n{x.op.name}{idx} {arg}" + (f" {x.src[0].op}" if len(x.src) else "")
     try:
       if len(rngs:=u.ranges):
-        label += f"\n({multirange_str(rngs, color=True)})"
+        label += f"\n({','.join([range_str(x, color=True) for x in sorted(rngs, key=lambda x: x.arg[0:-1])])})"
       if u.op not in {Ops.BUFFER, Ops.KERNEL, Ops.ASSIGN, Ops.COPY, Ops.SINK, *GroupOp.Buffer} and u._shape is not None:
         label += f"\n{shape_to_str(u.shape)}"
       if u.op in {Ops.INDEX, Ops.BUFFERIZE}:
         label += f"\n{u.render()}"
-        ranges: list[UOp] = []
-        for us in u.src[1:]: ranges += [s for s in us.toposort() if s.op in {Ops.RANGE, Ops.SPECIAL}]
-        if ranges: label += "\n"+' '.join([f"{s.render()}={s.vmax+1}" for s in ranges])
       if u.op in {Ops.END, Ops.REDUCE} and len(trngs:=list(UOp.sink(*u.src[range_start[u.op]:]).ranges)):
         label += "\n"+' '.join([f"{range_str(s, color=True)}({s.vmax+1})" for s in trngs])
     except Exception:
@@ -271,11 +268,8 @@ def get_llvm_mca(asm:str, mtriple:str, mcpu:str) -> dict:
   for i,usage in instr_usage.items(): rows[i].append([[k, v, (v/max_usage)*100] for k,v in usage.items()])
   return {"rows":rows, "cols":["Opcode", "Latency", {"title":"HW Resources", "labels":resource_labels}], "summary":summary}
 
-def get_stdout(f: Callable) -> str:
-  buf = io.StringIO()
-  try:
-    with redirect_stdout(buf), redirect_stderr(buf): f()
-  except Exception: traceback.print_exc(file=buf)
+def get_stdout(f:Callable) -> str:
+  with redirect_stdout(buf:=io.StringIO()): f()
   return buf.getvalue()
 
 def get_render(i:int, j:int, fmt:str) -> dict|None:
@@ -283,8 +277,8 @@ def get_render(i:int, j:int, fmt:str) -> dict|None:
   if not isinstance(prg:=trace.keys[i].ret, ProgramSpec): return None
   if fmt == "uops": return {"src":get_stdout(lambda: print_uops(prg.uops or [])), "lang":"txt"}
   if fmt == "src": return {"src":prg.src, "lang":"cpp"}
-  compiler = Device[prg.device].compiler
-  disasm_str = get_stdout(lambda: compiler.disassemble(compiler.compile(prg.src)))
+  lib = (compiler:=Device[prg.device].compiler).compile(prg.src)
+  disasm_str = get_stdout(lambda: compiler.disassemble(lib))
   from tinygrad.runtime.support.compiler_cpu import llvm, LLVMCompiler
   if isinstance(compiler, LLVMCompiler):
     mtriple = ctypes.string_at(llvm.LLVMGetTargetMachineTriple(tm:=compiler.target_machine)).decode()