start disabling rules there's no tests for

.github/workflows/test.yml

@@ -343,8 +343,6 @@ jobs:
       run: |
         python -m mypy --strict-equality --lineprecision-report .
         cat lineprecision.txt
-    - name: Run TYPED=1
-      run: TYPED=1 python -c "import tinygrad"
 
   unittest:
     name: Unit Tests
@@ -382,8 +380,8 @@ jobs:
         PYTHONPATH=. python extra/optimization/extract_dataset.py
         gzip -c /tmp/sops > extra/datasets/sops.gz
         DEBUG=1 MIN_ASTS=1 PYTHONPATH=. python extra/optimization/get_action_space.py
-    - name: Repo line count < 17500 lines
-      run: MAX_LINE_COUNT=17500 python sz.py
+    - name: Repo line count < 17000 lines
+      run: MAX_LINE_COUNT=17000 python sz.py
 
   fuzzing:
     name: Fuzzing
@@ -593,33 +591,6 @@ jobs:
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
-  testrangeify:
-    name: Linux (rangeify)
-    runs-on: ubuntu-24.04
-    timeout-minutes: 15
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v4
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: rangeify-minimal-llvm
-        deps: testing_minimal
-        llvm: "true"
-    - name: Test CPU=1 RANGEIFY=1
-      # TODO: add more passing tests here
-      # test_symbolic_arange_sym_step is passing now
-      # test_threefry_doesnt_use_long is because there's a contig after the long now
-      run: |
-        CPU=1 RANGEIFY=1 python3 -m pytest -n auto --durations 20 \
-          -k "not test_symbolic_arange_sym_step and not test_threefry_doesnt_use_long" \
-          test/test_tiny.py test/test_rangeify.py test/test_ops.py test/test_tensor_variable.py \
-          test/test_outerworld_range.py test/test_sample.py test/test_randomness.py test/test_tensor_data.py
-    - name: Test CPU=1 RANGEIFY=2
-      run: CPU=1 RANGEIFY=2 python3 -m pytest -n auto test/test_tiny.py test/test_rangeify.py test/test_ops.py --durations 20
-    - name: Test LLVM=1 RANGEIFY=1 (slow tests)
-      run: LLVM=1 RANGEIFY=1 python3 -m pytest -n auto test/models/test_mnist.py --durations 20
-
   testdevectorize:
     name: Linux (devectorize)
     runs-on: ubuntu-24.04

docs/tensor/elementwise.md

@@ -78,7 +78,6 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
 ::: tinygrad.Tensor.minimum
 ::: tinygrad.Tensor.where
 ::: tinygrad.Tensor.copysign
-::: tinygrad.Tensor.logaddexp
 
 ## Casting Ops
 

docs/tinybox.md

@@ -6,7 +6,7 @@ If you don't have a tinybox and you want one, see [tinygrad.org](https://tinygra
 
 ## Welcome
 
-Welcome to your tinybox! The tinybox is the universal system purpose-built for all AI infrastructure and workloads, from training to inference. The red box includes six 7900XTX GPUs, the green box includes six 4090 GPUs, and the green v2 box includes four 5090 GPUs. Whether you bought a red one or a green one, we want you to love it.
+Welcome to your tinybox! The tinybox is the universal system purpose-built for all AI infrastructure and workloads, from training to inference. The red box includes six 7900XTX GPUs, and the green box includes six 4090 GPUs. Whether you bought a red one or a green one, we want you to love it.
 
 We don't have a stupid cloud service, you don't have to create a tiny account to set it up, and we aren't tracking how you use the box. We're just happy you bought one. This petaflop is your petaflop.
 

examples/hlb_cifar10.py

@@ -118,7 +118,7 @@ class SpeedyResNet:
 # hyper-parameters were exactly the same as the original repo
 bias_scaler = 58
 hyp = {
-  'seed' : 201,
+  'seed' : 200,
   'opt': {
     'bias_lr':            1.76 * bias_scaler/512,
     'non_bias_lr':        1.76 / 512,

examples/mlperf/model_train.py

@@ -1297,9 +1297,6 @@ def train_llama3():
   SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
   TRAIN_ON_VAL       = config["TRAIN_ON_VAL"]           = getenv("TRAIN_ON_VAL", 0)
   SAMPLES            = config["SAMPLES"]                = getenv("SAMPLES", 5_760 if TRAIN_ON_VAL else 1_200_000 * 1152)
-  EVAL_FREQ          = config["EVAL_FREQ"]              = getenv("EVAL_FREQ", 46080)
-  EVAL_BS            = config["EVAL_BS"]                = getenv("EVAL_BS", 16)
-  EVAL_TARGET        = config["EVAL_TARGET"]            = getenv("EVAL_TARGET", 5.6)
 
   # LR=1e-4 TRAIN_ON_VAL=1 DEFAULT_FLOAT=bfloat16 FUSE_ARANGE=1 JITBEAM=2 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B WARMUP_STEPS=36 DECAY_STEPS=360 SEQLEN=512 PYTHONPATH=. AMD=1 AMD_LLVM=0 MODEL=llama3 python3 examples/mlperf/model_train.py
   # trains to 7
@@ -1378,7 +1375,7 @@ def train_llama3():
         total_norm += p.grad.float().square().sum()
       total_norm = total_norm.sqrt().contiguous()
       for p in optim.params:
-        p.grad = p.grad * (opt_gradient_clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)
+        p.grad = p.grad * opt_gradient_clip_norm / (total_norm + 1e-6)
 
     optim.step()
     scheduler.step()
@@ -1387,40 +1384,16 @@ def train_llama3():
     loss.realize(lr)
     return loss, lr
 
-  @TinyJit
-  @Tensor.train(False)
-  def eval_step(model, tokens:Tensor):
-    if (DP := getenv("DP", 1)) > 1:
-      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
-      tokens = tokens.shard(device, 0)
-    if (MP := getenv("MP", 1)) > 1:
-      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
-      tokens = tokens.shard(device)
-    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
-    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
-    return loss.flatten().float()
+  if getenv("FAKEDATA", 0):
+    def fake_data():
+      for _ in range(SAMPLES // GBS):
+        yield Tensor.randint(GBS, SEQLEN + 1, low=0, high=32000, dtype=dtypes.int32, device=Device.DEFAULT)
+    iter = fake_data()
+  else:
+    from examples.mlperf.dataloader import batch_load_llama3
+    iter = batch_load_llama3(GBS, SAMPLES, SEQLEN, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=SEED, val=bool(TRAIN_ON_VAL))
 
-  # ** data iters **
-  def fake_data(bs, samples):
-    for _ in range(samples // bs):
-      yield Tensor.randint(bs, SEQLEN + 1, low=0, high=32000, dtype=dtypes.int32, device=Device.DEFAULT)
-
-  def get_train_iter():
-    if getenv("FAKEDATA", 0):
-      return fake_data(GBS, SAMPLES)
-    else:
-      from examples.mlperf.dataloader import batch_load_llama3
-      return batch_load_llama3(GBS, SAMPLES, SEQLEN, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=SEED, val=bool(TRAIN_ON_VAL))
-
-  def get_eval_iter():
-    if getenv("FAKEDATA", 0):
-      return fake_data(EVAL_BS, 5760)
-    else:
-      from examples.mlperf.dataloader import batch_load_llama3
-      return batch_load_llama3(EVAL_BS, 5760, SEQLEN, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=SEED, val=True)
-
-  iter = get_train_iter()
-  i, sequences_seen = 0, 0
+  i = 0
   for tokens in tqdm(iter, total=SAMPLES//GBS):
     t = time.perf_counter()
     GlobalCounters.reset()
@@ -1435,33 +1408,9 @@ def train_llama3():
     if getenv("CKPT") and (i % 200 == 0 or i == 10):
       tqdm.write("saving checkpoint")
       if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
-      fn = f"{ckpt_dir}/llama3_{i}.safe"
+      fn = f"{ckpt_dir}/{i}.safe"
       safe_save(get_state_dict(model), fn)
-
     i += 1
-    sequences_seen += tokens.shape[0]
-
-    if sequences_seen % EVAL_FREQ == 0 and (i != 1 or EVAL_FREQ == 1):
-      tqdm.write(f"evaluating after {sequences_seen} sequences")
-
-      # run eval
-      eval_losses = []
-      eval_iter = get_eval_iter()
-      tqdm.write(f"evaluating {5760//EVAL_BS} batches of {EVAL_BS} sequences")
-
-      for tokens in tqdm(eval_iter, total=5760//EVAL_BS):
-        eval_losses += eval_step(model, tokens).tolist()
-      log_perplexity = Tensor(eval_losses).mean().float().item()
-
-      tqdm.write(f"eval log perplexity: {log_perplexity:.4f}")
-
-      if log_perplexity < EVAL_TARGET:
-        tqdm.write(f"target achieved after {sequences_seen} sequences")
-        if getenv("CKPT"):
-          if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
-          fn = f"{ckpt_dir}/llama3.safe"
-          safe_save(get_state_dict(model), fn)
-        break
 
 if __name__ == "__main__":
   multiprocessing.set_start_method('spawn')

extra/gemm/amd_uop_matmul.py

@@ -5,7 +5,7 @@ from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import getenv, colored, prod, unwrap
 from tinygrad.shape.shapetracker import ShapeTracker, View
 from tinygrad.shape.view import strides_for_shape
-from tinygrad.codegen.opt.kernel import axis_colors, Opt, OptOps
+from tinygrad.codegen.opt.kernel import axis_colors
 from tinygrad.codegen.opt.swizzler import merge_views, view_left
 
 def to_colored(full_shape, axis_types): return '_'.join([colored(str(s), axis_colors[at]) for s,at in zip(full_shape, axis_types)])
@@ -44,21 +44,6 @@ pm = PatternMatcher([
   (UPat(Ops.VIEW, src=(UPat(Ops.REDUCE_AXIS, src=(UPat.var("src"),), name="r"),), name="view"), swizzle_reduceop),
 ])
 
-def rangeify_kernel3():
-  a = Tensor.empty(N,N)
-  b = Tensor.empty(N,N)
-  c = a@b
-  #c = c.reshape((32,2,16,4,32,2,16,4)).contiguous()
-  with Context(RANGEIFY=1):
-    sink = c.schedule()[-1].ast
-  #print(sink)
-
-  opts  = [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.LOCAL, 0, 16), Opt(OptOps.UPCAST, 0, 2)]
-  opts += [Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.LOCAL, 1, 16), Opt(OptOps.UPCAST, 1, 2)]
-  opts += [Opt(OptOps.UNROLL, 0, 8)]
-
-  return sink.replace(arg=KernelInfo(opts_to_apply=tuple(opts)))
-
 def top_spec_kernel3():
   a = Tensor.empty(N,N)
   b = Tensor.empty(N,N)
@@ -324,15 +309,10 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
 
 if __name__ == "__main__":
   HL = getenv("HL")
-  if HL == 3: hprg = rangeify_kernel3()
-  elif HL == 2: hprg = top_spec_kernel3()
+  if HL == 2: hprg = top_spec_kernel3()
   elif HL == 1: hprg = hl_spec_kernel3()
   else: hprg = hand_spec_kernel3()
-  if HL == 3:
-    with Context(RANGEIFY=1, BLOCK_REORDER=0):
-      prg = get_program(hprg, Device.default.renderer)
-  else:
-    prg = get_program(hprg, Device.default.renderer)
+  prg = get_program(hprg, Device.default.renderer)
   print(prg.src)
   if getenv("SRC"): exit(0)
   hrunner = CompiledRunner(prg)

extra/torch_backend/backend.py

@@ -381,7 +381,6 @@ decomps = [
   aten.elu,  # elu has a scale + input_scale param
   aten.elu_backward,
   aten.softplus,
-  aten.logaddexp,
   aten.threshold,
   aten.nll_loss_forward,
   aten.nll_loss_backward,

ruff.toml

@@ -35,7 +35,6 @@ lint.select = [
 line-length = 150
 
 exclude = [
-  ".git/",
   "docs/",
   "extra/",
   "tinygrad/runtime/autogen",

setup.py

@@ -29,7 +29,6 @@ setup(name='tinygrad',
         'tinygrad.apps',
         'tinygrad.codegen',
         'tinygrad.codegen.opt',
-        'tinygrad.codegen.late',
         'tinygrad.engine',
         'tinygrad.frontend',
         'tinygrad.nn',
@@ -64,7 +63,6 @@ setup(name='tinygrad',
             "pre-commit",
             "ruff",
             "numpy",
-            "typeguard",
         ],
         #'mlperf': ["mlperf-logging @ git+https://github.com/mlperf/[email protected]"],
         'testing_minimal': testing_minimal,

test/external/external_test_onnx_backend.py

@@ -134,6 +134,7 @@ backend_test.exclude('test_simple_rnn_*')
 
 # no control flow
 # control flow uses AttributeProto.GRAPH
+backend_test.exclude('test_if_*')
 backend_test.exclude('test_loop*')
 backend_test.exclude('test_range_float_type_positive_delta_expanded_cpu') # requires loop
 backend_test.exclude('test_affine_grid_2d_align_corners_expanded_cpu')
@@ -182,8 +183,6 @@ backend_test.exclude('test_resize_downsample_scales_cubic_antialias_cpu') # anti
 backend_test.exclude('test_resize_downsample_sizes_cubic_antialias_cpu') # antialias not implemented
 backend_test.exclude('test_ai_onnx_ml_label_encoder_tensor_value_only_mapping_cpu') # bad data type string
 backend_test.exclude('test_ai_onnx_ml_label_encoder_tensor_mapping_cpu') # bad data type string
-backend_test.exclude('test_if_opt_cpu') # ValueError: 13 is not a valid AttributeType
-backend_test.exclude('test_if_seq_cpu') # NotImplementedError: op='SequenceConstruct' is not supported
 
 backend_test.exclude('test_scatternd_min_cpu') # min not yet supported
 backend_test.exclude('test_scatternd_max_cpu') # max not yet supported

test/external/external_test_onnx_ops.py

@@ -100,25 +100,6 @@ class TestMainOnnxOps(TestOnnxOps):
     self._test_resize_scales([0.01, 0.25, 0.5, 0.51, 0.6, 1.0, 1.5, 2.0, 3.5, 20.0], mode="cubic", exclude_outside=1)
     self._test_resize_scales([0.01, 0.25, 0.5, 0.51, 0.6, 1.0, 1.5, 2.0, 3.5, 20.0], mode="cubic", exclude_outside=0)
 
-  def _test_if(self, then_value, else_value):
-    then_out = onnx.helper.make_tensor_value_info("res", onnx.TensorProto.FLOAT, then_value.shape)
-    else_out = onnx.helper.make_tensor_value_info("res", onnx.TensorProto.FLOAT, else_value.shape)
-
-    then_const_node = onnx.helper.make_node("Constant", inputs=[], outputs=["res"], value=onnx.numpy_helper.from_array(then_value))
-    else_const_node = onnx.helper.make_node("Constant", inputs=[], outputs=["res"], value=onnx.numpy_helper.from_array(else_value))
-
-    then_body = onnx.helper.make_graph([then_const_node], "then_body", [], [then_out])
-    else_body = onnx.helper.make_graph([else_const_node], "else_body", [], [else_out])
-
-    self.helper_test_single_op("If", {"cond": np.array(False).astype(bool)}, {"then_branch": then_body, "else_branch": else_body}, ["res"])
-    self.helper_test_single_op("If", {"cond": np.array(True).astype(bool)}, {"then_branch": then_body, "else_branch": else_body}, ["res"])
-
-  def test_if_different_shapes_broadcastable(self):
-    self._test_if(np.array([[1], [2]]).astype(np.float32), np.array([[6, 5, 4, 3, 2, 1]]).astype(np.float32))
-
-  def test_if_different_shapes_not_broadcastable(self):
-    self._test_if(np.array([[1, 2, 3], [4, 5, 6]]).astype(np.float32), np.array([[6, 5, 4, 3, 2, 1]]).astype(np.float32))
-
   def test_resize_downsample_scales_linear_align_corners(self):
     # https://github.com/onnx/onnx/blob/main/docs/Operators.md#examples-131
     X = np.array([[[[1, 2, 3, 4], [5, 6, 7, 8]]]], dtype=np.float32)

test/external/fuzz_fast_idiv.py

@@ -1,8 +1,8 @@
 import random
 import z3
 from tinygrad import dtypes
-from tinygrad.uop.spec import uops_to_z3, z3_cdiv
-from tinygrad.uop.ops import UOp
+from tinygrad.uop.spec import z3_renderer, z3_cdiv
+from tinygrad.uop.ops import UOp, graph_rewrite
 from tinygrad.uop.decompositions import fast_idiv
 random.seed(42)
 
@@ -19,7 +19,8 @@ if __name__ == "__main__":
     if expr is None: continue
 
     solver = z3.Solver()
-    z3_expr, x =uops_to_z3(solver, expr, u)
+    z3_sink = graph_rewrite(expr.sink(u), z3_renderer, ctx=(solver, {}))
+    z3_expr, x = z3_sink.src[0].arg, z3_sink.src[1].arg
 
     if solver.check(z3_expr != z3_cdiv(x, d)) == z3.sat:
       assert False, f"Failed: {expr.render()} != x//{d} at x={solver.model()}\nx={u}\nd={d}\n{z3_expr=}\n{x/d=}"

test/external/fuzz_symbolic.py

@@ -1,8 +1,8 @@
 import random, operator
 import z3
 from tinygrad import Variable, dtypes
-from tinygrad.uop.ops import UOp
-from tinygrad.uop.spec import uops_to_z3
+from tinygrad.uop.ops import UOp, graph_rewrite
+from tinygrad.uop.spec import z3_renderer
 from tinygrad.helpers import DEBUG, Context
 
 seed = random.randint(0, 100)
@@ -57,7 +57,8 @@ if __name__ == "__main__":
 
     solver = z3.Solver()
     solver.set(timeout=5000)  # some expressions take very long verify, but its very unlikely they actually return sat
-    z3_expr, z3_simplified_expr, v1, v2, v3 = uops_to_z3(solver, expr, simplified_expr, u1, u2, u3)
+    z3_sink = graph_rewrite(expr.sink(simplified_expr, u1, u2, u3), z3_renderer, ctx=(solver, {}))
+    z3_expr, z3_simplified_expr = z3_sink.src[0].arg, z3_sink.src[1].arg
     check = solver.check(z3_simplified_expr != z3_expr)
     if check == z3.unknown and DEBUG>=1:
       skipped += 1
@@ -68,6 +69,7 @@ if __name__ == "__main__":
             f"expr = {expr.render(simplify=False)}\n")
     elif check == z3.sat:
       m = solver.model()
+      v1, v2, v3 = z3_sink.src[2].arg, z3_sink.src[3].arg, z3_sink.src[4].arg
       n1, n2, n3 = m[v1], m[v2], m[v3]
       u1_val, u2_val, u3_val = u1.const_like(n1.as_long()), u2.const_like(n2.as_long()), u3.const_like(n3.as_long())
       with Context(CORRECT_DIVMOD_FOLDING=1):

test/test_const_folding.py

@@ -1,10 +1,11 @@
 import unittest, itertools, math
+from typing import Any
 from tinygrad import Tensor, Device, dtypes
-from tinygrad.dtype import DType, ConstType
+from tinygrad.dtype import DType
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.codegen import full_rewrite_to_sink
-from tinygrad.device import is_dtype_supported
 import numpy as np
+from tinygrad.device import is_dtype_supported
 from test.helpers import not_support_multi_device
 
 def _check_ast_count(desired_count:int, t:Tensor):
@@ -24,7 +25,7 @@ class TestUnaryOpsConstFolding(unittest.TestCase):
     _check_ast_count(0, Tensor.ones(4).cast(dtypes.int16))
     _check_ast_count(0, Tensor.full(4, fill_value=-1).cast(dtypes.uint16))
 
-  @unittest.expectedFailure  # no two level fold
+  @unittest.expectedFailure  # no two level fold at lazybuffer
   def test_neg_folding(self):
     _check_ast_count(0, Tensor([1, 2, 3]).mul(-1).neg())
     _check_ast_count(0, Tensor([1, 2, 3]).neg().mul(-1))
@@ -103,7 +104,7 @@ class TestBinaryOpsConstFolding(unittest.TestCase):
 
 class TestBitcastConstFolding(unittest.TestCase):
   def test_scalar_bitcast(self):
-    def t(cases: dict[DType, ConstType]):
+    def t(cases: dict[DType, Any]):
       for (from_dt, from_v), (to_dt, to_v) in itertools.product(cases.items(), cases.items()):
         if not math.isnan(from_v):
           r = full_rewrite_to_sink(UOp.const(from_dt, from_v).bitcast(to_dt).sink()).src[0]
@@ -164,6 +165,7 @@ class TestMovedConstFolding(unittest.TestCase):
     _check_ast_count(1, Tensor([1.0, 2, 3, 4]) * Tensor.ones(2).pad(((1, 1),)))
 
   def test_cast_padded(self):
+    # NOTE: this is folded due to CAST_BEFORE_VIEW
     if is_dtype_supported(dtypes.int16):
       _check_ast_count(0, Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int16))
       np.testing.assert_equal(Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int16).numpy(), [0, 1, 1, 1, 1, 0])

test/test_define_reg.py

@@ -0,0 +1,32 @@
+import unittest
+from tinygrad import dtypes, Device, Tensor, Context
+from tinygrad.dtype import AddrSpace
+from tinygrad.helpers import getenv
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
+from tinygrad.shape.shapetracker import ShapeTracker
+from tinygrad.engine.realize import get_program, ExecItem, CompiledRunner
+
+class TestDefineReg(unittest.TestCase):
+  def test_simple(self, at=AxisType.UPCAST):
+    N = 16
+    bout = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=0).view(ShapeTracker.from_shape((N,N)))
+    a = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=1).view(ShapeTracker.from_shape((N,N)))
+    a_col = UOp(Ops.DEFINE_REG, dtypes.float.ptr(N, AddrSpace.REG), arg=0).view(ShapeTracker.from_shape((N,N), (0,1)))
+
+    out = a_col.load(a_col.store(a.load()))
+    sink = bout.store(out).sink(arg=KernelInfo(name="regcopy", axis_types=(AxisType.LOOP, at)))
+    prg = get_program(sink, Device.default.renderer)
+
+    with Context(DEBUG=0):
+      a = Tensor.randn(N, N).realize()
+      b = Tensor.empty(N, N).realize()
+    hrunner = CompiledRunner(prg)
+    ExecItem(hrunner, [b.uop.buffer, a.uop.buffer]).run(wait=True)
+    with Context(DEBUG=0):
+      self.assertEqual((b-a).mean().item(), 0.0)
+
+  @unittest.skipIf(getenv("PTX"), "ptx needs regs to be unrolled")
+  def test_simple_loop(self): self.test_simple(AxisType.LOOP)
+
+if __name__ == '__main__':
+  unittest.main()

test/test_dtype.py

@@ -414,11 +414,11 @@ class TestDtypeUsage(unittest.TestCase):
         t = Tensor([[1, 2], [3, 4]], dtype=d)
         (t*t).max().item()
 
-@unittest.skipUnless(is_dtype_supported(dtypes.bfloat16) or Device.DEFAULT == "PYTHON", f"no bfloat16 on {Device.DEFAULT}")
+@unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), f"no bfloat16 on {Device.DEFAULT}")
 class TestOpsBFloat16(unittest.TestCase):
   def test_cast(self):
     # TODO: helper_test_op breaks in unrelated part
-    # TODO: wrong output with GPU=1 on mac
+    # TODO: wrong output with GPU=1 / PYTHON=1 on mac
     data = [60000.0, 70000.0, 80000.0]
     np.testing.assert_allclose(Tensor(data).cast("bfloat16").numpy(), torch.tensor(data).type(torch.bfloat16).float().numpy())
 

test/test_dtype_alu.py

@@ -1,13 +1,16 @@
-import unittest, operator, math
+import unittest
+
 from tinygrad import Tensor, dtypes, Device
+import operator
+import numpy as np
+from hypothesis import given, strategies as strat, settings, HealthCheck
 from tinygrad.dtype import DType
 from tinygrad.helpers import CI, getenv
+from tinygrad.engine.realize import run_schedule
+from tinygrad.uop.ops import GroupOp
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
-import numpy as np
-import pytest
-from hypothesis import given, strategies as strat, settings, HealthCheck
-
+import pytest, math
 pytestmark = pytest.mark.filterwarnings("ignore")
 
 settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
@@ -60,21 +63,25 @@ def universal_test(a, b, dtype, op):
   ta, tb = Tensor([a], dtype=dtype), Tensor([b], dtype=dtype)
   tensor_value = (op[0](ta, tb)).numpy()
   numpy_value = op[1](ta.numpy(), tb.numpy())
-  if dtype in dtypes.floats:
-    atol, rtol = {dtypes.bfloat16:(1e-3, 1e-2)}.get(dtype, (1e-10, 1e-7))
-    np.testing.assert_allclose(tensor_value, numpy_value, atol=atol, rtol=rtol)
+  if dtype == dtypes.bfloat16: np.testing.assert_allclose(tensor_value, numpy_value, atol=1e-3, rtol=1e-2)
+  elif dtype in dtypes_float: np.testing.assert_allclose(tensor_value, numpy_value, atol=1e-10)
   else: np.testing.assert_equal(tensor_value, numpy_value)
 
 def universal_test_unary(a, dtype, op):
   if not isinstance(op, tuple): op = (op, op)
   ta = Tensor([a], dtype=dtype)
   out: Tensor = op[0](ta)
+  sched = out.schedule()
+  ast = sched[-1].ast
+  run_schedule(sched)
   tensor_value = out.numpy()
   numpy_value = op[1](ta.numpy())
-  if dtype in dtypes.floats:
-    atol, rtol = {dtypes.float16:(1e-3, 1e-2), dtypes.bfloat16:(1e-3, 1e-2)}.get(dtype, (1e-6, 1e-5))
-    np.testing.assert_allclose(tensor_value, numpy_value, atol=atol, rtol=rtol)
+  if dtype in (dtypes.float16, dtypes.bfloat16): np.testing.assert_allclose(tensor_value, numpy_value, atol=1e-3, rtol=1e-2)
+  elif dtype in dtypes_float: np.testing.assert_allclose(tensor_value, numpy_value, atol=1e-6, rtol=1e-5)
   else: np.testing.assert_equal(tensor_value, numpy_value)
+  if op[0] != Tensor.reciprocal: # reciprocal is not supported in most backends
+    op = [x for x in ast.toposort() if x.op in GroupOp.Unary][0]
+    assert op.dtype == dtype
 
 def universal_test_cast(a, in_dtype, dtype):
   tensor_value = Tensor([a], dtype=in_dtype).cast(dtype)
@@ -92,44 +99,45 @@ def universal_test_midcast(a, b, c, op1, op2, d1:DType, d2:DType):
   np.testing.assert_allclose(tensor_value, numpy_value, rtol=1e-6 if getenv("PTX") else 1e-7)
 
 class TestDTypeALU(unittest.TestCase):
-  @unittest.skipUnless(is_dtype_supported(dtypes.float64), f"no float64 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.float64, Device.DEFAULT), f"no float64 on {Device.DEFAULT}")
   @given(ht.float64, ht.float64, strat.sampled_from(binary_operations))
   def test_float64(self, a, b, op): universal_test(a, b, dtypes.float64, op)
 
   @given(ht.float32, ht.float32, strat.sampled_from(binary_operations))
   def test_float32(self, a, b, op): universal_test(a, b, dtypes.float32, op)
 
-  @unittest.skipUnless(is_dtype_supported(dtypes.float16), f"no float16 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.float16, Device.DEFAULT), f"no float16 on {Device.DEFAULT}")
   @given(ht.float16, ht.float16, strat.sampled_from(binary_operations))
   def test_float16(self, a, b, op): universal_test(a, b, dtypes.float16, op)
 
-  @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), f"no bfloat16 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16, Device.DEFAULT), f"no bfloat16 on {Device.DEFAULT}")
   @given(ht.bfloat16, ht.bfloat16, strat.sampled_from(binary_operations))
   def test_bfloat16(self, a, b, op): universal_test(a, b, dtypes.bfloat16, op)
 
   @given(ht.float32, strat.sampled_from(unary_operations))
   def test_float32_unary(self, a, op): universal_test_unary(a, dtypes.float32, op)
 
-  @unittest.skipUnless(is_dtype_supported(dtypes.float16), f"no float16 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.float16, Device.DEFAULT), f"no float16 on {Device.DEFAULT}")
   @given(ht.float16, strat.sampled_from(unary_operations))
   def test_float16_unary(self, a, op): universal_test_unary(a, dtypes.float16, op)
 
-  @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), f"no bfloat16 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16, Device.DEFAULT), f"no bfloat16 on {Device.DEFAULT}")
   @given(ht.bfloat16, strat.sampled_from(unary_operations))
+  @unittest.skipIf(Device.DEFAULT in ["AMD"], "broken on AMD?")
   def test_bfloat16_unary(self, a, op): universal_test_unary(a, dtypes.bfloat16, op)
 
   @given(ht.uint8, ht.uint8, strat.sampled_from(integer_binary_operations))
   def test_uint8(self, a, b, op): universal_test(a, b, dtypes.uint8, op)
 
-  @unittest.skipUnless(is_dtype_supported(dtypes.uint16), f"no uint16 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.uint16, Device.DEFAULT), f"no uint16 on {Device.DEFAULT}")
   @given(ht.uint16, ht.uint16, strat.sampled_from(integer_binary_operations))
   def test_uint16(self, a, b, op): universal_test(a, b, dtypes.uint16, op)
 
-  @unittest.skipUnless(is_dtype_supported(dtypes.uint32), f"no uint32 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.uint32, Device.DEFAULT), f"no uint32 on {Device.DEFAULT}")
   @given(ht.uint32, ht.uint32, strat.sampled_from(integer_binary_operations))
   def test_uint32(self, a, b, op): universal_test(a, b, dtypes.uint32, op)
 
-  @unittest.skipUnless(is_dtype_supported(dtypes.uint64), f"no uint64 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.uint64, Device.DEFAULT), f"no uint64 on {Device.DEFAULT}")
   @given(ht.uint64, ht.uint64, strat.sampled_from(integer_binary_operations))
   def test_uint64(self, a, b, op): universal_test(a, b, dtypes.uint64, op)
 
@@ -142,7 +150,7 @@ class TestDTypeALU(unittest.TestCase):
   @given(ht.int32, ht.int32, strat.sampled_from(integer_binary_operations))
   def test_int32(self, a, b, op): universal_test(a, b, dtypes.int32, op)
 
-  @unittest.skipUnless(is_dtype_supported(dtypes.int64), f"no int64 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.int64, Device.DEFAULT), f"no int64 on {Device.DEFAULT}")
   @given(ht.int64, ht.int64, strat.sampled_from(integer_binary_operations))
   def test_int64(self, a, b, op): universal_test(a, b, dtypes.int64, op)
 
@@ -172,7 +180,7 @@ class TestDTypeALU(unittest.TestCase):
   @settings(suppress_health_check=[HealthCheck.filter_too_much])
   @given(strat.data(), strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
   def test_float_cast_to_unsigned(self, a, float_dtype, unsigned_dtype):
-    if not is_dtype_supported(float_dtype): float_dtype = dtypes.float32
+    if not is_dtype_supported(float_dtype, Device.DEFAULT): float_dtype = dtypes.float32
     float_strat = {dtypes.float16: ht.float16, dtypes.float32: ht.float32, dtypes.float64: ht.float64}[float_dtype]
     float_strat = float_strat.filter(lambda x: 0 < x < dtypes.max(unsigned_dtype))
     universal_test_cast(a.draw(float_strat), float_dtype, unsigned_dtype)
@@ -180,7 +188,7 @@ class TestDTypeALU(unittest.TestCase):
   @settings(suppress_health_check=[HealthCheck.filter_too_much])
   @given(strat.data(), strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
   def test_float_cast_to_unsigned_overflow(self, a, float_dtype, unsigned_dtype):
-    if not is_dtype_supported(float_dtype): float_dtype = dtypes.float32
+    if not is_dtype_supported(float_dtype, Device.DEFAULT): float_dtype = dtypes.float32
     float_strat = {dtypes.float16: ht.float16, dtypes.float32: ht.float32, dtypes.float64: ht.float64}[float_dtype]
     overflow_strat = float_strat.filter(lambda x: x > dtypes.max(unsigned_dtype) and x <= dtypes.max(dtypes.int32))
     universal_test_cast(a.draw(overflow_strat), float_dtype, unsigned_dtype)
@@ -188,7 +196,7 @@ class TestDTypeALU(unittest.TestCase):
   @settings(suppress_health_check=[HealthCheck.filter_too_much])
   @given(strat.data(), strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
   def test_float_cast_to_unsigned_underflow(self, a, float_dtype, unsigned_dtype):
-    if not is_dtype_supported(float_dtype): float_dtype = dtypes.float32
+    if not is_dtype_supported(float_dtype, Device.DEFAULT): float_dtype = dtypes.float32
     float_strat = {dtypes.float16: ht.float16, dtypes.float32: ht.float32, dtypes.float64: ht.float64}[float_dtype]
     underflow_strat = float_strat.filter(lambda x: x < 0 and x >= dtypes.min(dtypes.int32))
     universal_test_cast(a.draw(underflow_strat), float_dtype, unsigned_dtype)

test/test_linearizer.py

@@ -117,7 +117,6 @@ class TestLinearizer(unittest.TestCase):
         if skip and i in skip: continue
         assert ranges[i-1] != u, f"multireduce nested the ranges! {ranges[i-1], {u}}"
 
-  @unittest.skip("broken. should not depends on push_views and implementation details of getitem")
   @unittest.skipIf(CI and Device.DEFAULT in {"PTX", "AMD", "NV"}, "very slow")
   def test_indexing_multireduce(self):
     dataset = Tensor.rand(16384, 256).realize()

test/test_memory_planner.py

@@ -120,19 +120,5 @@ class TestMemoryPlanner(unittest.TestCase):
     ]
     check_assign(bs)
 
-  def test_very_small_buffers(self):
-    bs = [
-      [b(0, pin=True), b(1, size=32)],
-      [b(3, size=4), b(4, size=6)],
-    ]
-    check_assign(bs)
-
-  def test_very_big_buffers(self):
-    bs = [
-      [b(0, pin=True), b(1, size=34359738368000)],
-      [b(3, size=1 << 128), b(4, size=1 << 64)],
-    ]
-    check_assign(bs)
-
 if __name__ == "__main__":
   unittest.main()

test/test_multitensor.py

@@ -1128,7 +1128,6 @@ class TestMultiRamUsage(unittest.TestCase):
     self.assertUsed(self.N*self.N*4) # sharding should not increase total ram usage
   def test_zeros_shard_self(self): self.test_zeros_shard((d0, d1))
 
-  @unittest.skip("flaky")
   def test_zeros_contiguous_shard(self):
     _ = Tensor.zeros(self.N, self.N).contiguous().shard(devices_2, axis=0).contiguous().realize()
     self.assertUsed(self.N*self.N*4) # sharding should not increase total ram usage

test/test_nn.py

@@ -210,27 +210,6 @@ class TestNN(unittest.TestCase):
       np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
       np.testing.assert_allclose(layer.bias.grad.numpy(), torch_layer.bias.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
 
-  def test_layernorm_forward(self):
-    N, C, H, W = 20, 5, 10, 10
-
-    # create in torch
-    torch_layer = torch.nn.LayerNorm([H, W]).eval()
-
-    # create in tinygrad
-    layer = LayerNorm([H, W])
-    layer.weight = Tensor(torch_layer.weight.detach().numpy(), requires_grad=True)
-    layer.bias = Tensor(torch_layer.bias.detach().numpy(), requires_grad=True)
-
-    x = Tensor.empty(N, C, H, W, requires_grad=True)
-    z = layer(x)
-    z.realize()
-
-    torch_x = torch.tensor(x.numpy(), requires_grad=True)
-    torch_z = torch_layer(torch_x)
-    torch_z.sum().backward()
-
-    np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-6, rtol=5e-6)
-
   def test_layernorm(self):
     N, C, H, W = 20, 5, 10, 10
 

test/test_ops.py

@@ -928,12 +928,6 @@ class TestOps(unittest.TestCase):
       for j in [-1., 0., 1.]:
         helper_test_op(None, torch.copysign, Tensor.copysign, vals=[[i], [j]])
 
-  def test_logaddexp(self):
-    helper_test_op([(45,65), (45,65)], torch.logaddexp, Tensor.logaddexp)
-    helper_test_op(None, torch.logaddexp, Tensor.logaddexp, vals=[[-1.], [-1.0, 2, 3]])
-    helper_test_op(None, torch.logaddexp, Tensor.logaddexp, vals=[[-100.0, -200, -300], [-1.0, 2, 3]])
-    helper_test_op(None, torch.logaddexp, Tensor.logaddexp, vals=[[1.0, 2000, 30000], [-1.0, 2, 3]])
-
   def test_softsign(self):
     helper_test_op([(45,65)], torch.nn.functional.softsign, Tensor.softsign)
     helper_test_op([()], torch.nn.functional.softsign, Tensor.softsign)
@@ -971,6 +965,8 @@ class TestOps(unittest.TestCase):
     helper_test_op([(45,65)], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6)
     helper_test_op([(45,65)], lambda t: torch.nn.functional.softplus(t, beta=3), lambda t: Tensor.softplus(t, beta=3), grad_atol=1e-6)
     helper_test_op([(45,65)], lambda t: torch.nn.functional.softplus(t, beta=1/3), lambda t: Tensor.softplus(t, beta=1/3), grad_atol=1e-6)
+    helper_test_op([(45,65)], lambda t: torch.nn.functional.softplus(t, beta=3, threshold=0.5),
+                              lambda t: Tensor.softplus(t, beta=3, threshold=0.5), grad_atol=1e-6)
     helper_test_op([(45,65)], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6, low=300, high=400)
     helper_test_op([(45,65)], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6, low=-400, high=-300)
     helper_test_op([()], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6)
@@ -2465,20 +2461,6 @@ class TestOps(unittest.TestCase):
       lambda x: Tensor.max_unpool2d(*Tensor.max_pool2d(x, kernel_size=(2,2), return_indices=True),
                                     kernel_size=(2,2), output_size=(99,99,7,6)), forward_only=True)
 
-  def test_max_unpool2d_inf(self):
-    data = [[[[math.inf, -math.inf, math.nan], [1.0, 2.0, 3.0]]]]
-    ksz = (2,2)
-    helper_test_op((),
-      lambda: torch.nn.functional.max_unpool2d(
-        *torch.nn.functional.max_pool2d(torch.tensor(data), kernel_size=ksz, return_indices=True),
-        kernel_size=ksz
-      ),
-      lambda: Tensor.max_unpool2d(
-        *Tensor.max_pool2d(Tensor(data), kernel_size=ksz, return_indices=True),
-        kernel_size=ksz
-      ),
-      forward_only=True)
-
   def test_avg_pool2d(self):
     shape = (32,2,111,28)
     for ksz in [(2,2), (3,3), (3,2), (5,5), (5,1)]:
@@ -2712,10 +2694,6 @@ class TestOps(unittest.TestCase):
     i, j, k, o, p = [Tensor(tor.detach().cpu().numpy().astype(np.int32), requires_grad=False) for tor in [a,b,c,d,e]]
     return a,b,c,d,e,i,j,k,o,p
 
-  def test_fancy_indexing_inf(self):
-    data = [math.inf, -math.inf, math.nan]
-    helper_test_op((), lambda: torch.tensor(data)[torch.tensor([0, 1, 2])], lambda: Tensor(data)[Tensor([0, 1, 2])])
-
   def test_slice_fancy_indexing_no_dim_collapse(self):
     a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
     # no dim collapse from int or dim injection from None
@@ -2826,7 +2804,11 @@ class TestOps(unittest.TestCase):
     helper_test_op(None, lambda x: x.gather(dim=0, index=torch.tensor([2, 1, 0, 1, 2], requires_grad=False)),
                          lambda x: x.gather(dim=0, index=Tensor([2, 1, 0, 1, 2])),
                          vals=[[1., 2., 3.]])
-    # gather with inf values
+
+  @unittest.expectedFailure
+  @unittest.skipIf(torch._C._get_privateuse1_backend_name() == "tiny", 'results in a success instead of a failure')
+  def test_gather_failure(self):
+    # gather with inf values do not work, other values results in nan
     helper_test_op(None, lambda x: x.gather(dim=0, index=torch.tensor([2, 1, 0, 1, 2], requires_grad=False)),
                          lambda x: x.gather(dim=0, index=Tensor([2, 1, 0, 1, 2])),
                          vals=[[-float("inf"), 2., 3.]])

test/test_pickle.py

@@ -20,7 +20,7 @@ class TestPickle(unittest.TestCase):
     self.assertEqual(pm2.rewrite(sink).key, tt.key)
 
   def test_pickle_main_pattern_matcher(self):
-    from tinygrad.codegen.late.devectorizer import sym
+    from tinygrad.codegen.devectorizer import sym
     ssym = pickle.dumps(sym)
     dsym = pickle.loads(ssym)
     self.assertEqual(dsym.patterns[0][0].location, sym.patterns[0][0].location)

test/test_profiler.py

@@ -1,6 +1,6 @@
 import unittest, struct, contextlib, statistics, time, gc
 from tinygrad import Device, Tensor, dtypes, TinyJit
-from tinygrad.helpers import CI, getenv, Context, ProfileRangeEvent, cpu_profile, cpu_events, ProfilePointEvent, dedup
+from tinygrad.helpers import CI, getenv, Context, ProfileRangeEvent, cpu_profile, cpu_events
 from tinygrad.device import Buffer, BufferSpec, Compiled, ProfileDeviceEvent, ProfileGraphEvent
 from tinygrad.runtime.support.hcq import HCQCompiled
 from tinygrad.engine.realize import get_runner
@@ -209,18 +209,5 @@ class TestProfiler(unittest.TestCase):
     for ge in graphs:
       self.assertEqual(len(ge.ents), len(graphs))
 
-  def test_trace_metadata(self):
-    with Context(TRACEMETA=1):
-      a = Tensor.empty(1)+2
-      b = Tensor.empty(1)+2
-      with helper_collect_profile(TestProfiler.d0) as profile:
-        Tensor.realize(a, b)
-    profile, _ = helper_profile_filter_device(profile, TestProfiler.d0.device)
-    exec_points = [e for e in profile if isinstance(e, ProfilePointEvent) and e.name == "exec"]
-    range_events = [e for e in profile if isinstance(e, ProfileRangeEvent)]
-    self.assertEqual(len(exec_points), len(range_events), 2)
-    self.assertEqual(len(dedup(e.key for e in exec_points)), 1)
-    self.assertEqual(len(dedup(e.arg['metadata'] for e in exec_points)), 1)
-
 if __name__ == "__main__":
   unittest.main()

test/test_rangeify.py

@@ -1,198 +0,0 @@
-import unittest
-from tinygrad import Tensor
-from tinygrad.helpers import RANGEIFY, Context, GlobalCounters
-
-N = 256
-
-@unittest.skipIf(RANGEIFY<1, "tests only for RANGEIFY")
-class TestRangeify(unittest.TestCase):
-  def test_expand_children(self):
-    A = Tensor.empty(N, N).sum(axis=1)
-    ba = A.expand(N, N)
-    ((ba+1).sum(axis=1) + (ba+2).sum(axis=0)).realize()
-
-  def test_partial_contig(self):
-    A = Tensor.empty(64, 64, 64)
-    ret = A.sum(axis=2).contiguous(arg=(1,)).sum(axis=1)
-    ret.realize()
-
-  def test_double_gemm_real(self):
-    def go():
-      with Context(DEBUG=0):
-        Tensor.manual_seed(1337)
-        A,B,C = [Tensor.randn(N, N) for _ in range(3)]
-        Tensor.realize(A, B, C)
-      GlobalCounters.reset()
-      return (A@B@C).realize()
-    rng = go()
-    with Context(RANGEIFY=0, DEBUG=2):
-      ref = go()
-      mse = ((rng-ref)**2).sum().item()
-    print(f"mse: {mse}")
-    self.assertLessEqual(mse, 1e-2)
-
-  def test_double_gemm(self):
-    A = Tensor.empty(N, N)
-    B = Tensor.empty(N, N)
-    C = Tensor.empty(N, N)
-    (A@B@C).realize()
-
-  def test_double_gemm_exp(self):
-    A = Tensor.empty(N, N)
-    B = Tensor.empty(N, N)
-    C = Tensor.empty(N, N)
-    (((A@B).exp()@C).exp()).realize()
-
-  def test_double_gemm_relu(self):
-    A = Tensor.empty(N, N)
-    B = Tensor.empty(N, N)
-    C = Tensor.empty(N, N)
-    (((A@B).relu()@C).relu()).realize()
-
-  def test_double_gemm_relu_half_contig(self):
-    A = Tensor.empty(N, N)
-    B = Tensor.empty(N, N)
-    C = Tensor.empty(N, N)
-    (((A@B).relu().contiguous(arg=(1,))@C).relu()).realize()
-
-  def test_double_gemm_half_contig(self):
-    A = Tensor.empty(N, N)
-    B = Tensor.empty(N, N)
-    C = Tensor.empty(N, N)
-    ((A@B).contiguous(arg=(1,))@C).realize()
-
-  def test_double_gemm_contig(self):
-    A = Tensor.empty(N, N)
-    B = Tensor.empty(N, N)
-    C = Tensor.empty(N, N)
-    ((A@B).contiguous()@C).realize()
-
-  def test_many_gemm(self):
-    A = Tensor.empty(N, N)
-    B = Tensor.empty(N, N)
-    C = Tensor.empty(N, N)
-    D = Tensor.empty(N, N)
-    E = Tensor.empty(N, N)
-    F = Tensor.empty(N, N)
-    (A@B@C@D@E@F).realize()
-
-  def test_conv2d(self):
-    x = Tensor.empty(1, 4, 32, 32)
-    w1 = Tensor.empty(8, 4, 3, 3)
-    x.conv2d(w1).realize()
-
-  def test_conv2d_t(self):
-    x = Tensor.empty(1, 4, 32, 32)
-    w1 = Tensor.empty(8, 4, 3, 3)
-    (x*2).conv2d(w1).realize()
-
-  def test_double_conv2d(self):
-    x = Tensor.empty(1, 4, 32, 32)
-    w1 = Tensor.empty(8, 4, 3, 3)
-    w2 = Tensor.empty(12, 8, 3, 3)
-    x.conv2d(w1).conv2d(w2).realize()
-
-  def test_double_conv2d_half_contig(self):
-    x = Tensor.empty(1, 4, 32, 32)
-    w1 = Tensor.empty(8, 4, 3, 3)
-    w2 = Tensor.empty(12, 8, 3, 3)
-    # NOTE: this contiguous doesn't help
-    x.conv2d(w1).contiguous(arg=(1,)).conv2d(w2).permute(0,2,3,1).contiguous().realize()
-
-  def test_double_conv2d_contig(self):
-    x = Tensor.empty(1, 4, 32, 32)
-    w1 = Tensor.empty(8, 4, 3, 3)
-    w2 = Tensor.empty(12, 8, 3, 3)
-    x.conv2d(w1).contiguous().conv2d(w2).realize()
-
-  def test_transformer_ffn(self):
-    from tinygrad.apps.llm import TransformerBlock
-    from tinygrad import nn
-    blk = TransformerBlock(1024, 4096, 1, 1, 1e-5)
-    for p in nn.state.get_parameters(blk): p.replace(Tensor.empty(p.shape))
-
-    x = Tensor.empty(128, 1024)
-    out = blk._feed_forward(x)
-    out.realize()
-
-  def test_flash_attention(self):
-    #BS, HEADS, SEQLEN, EMB = 4, 2, 16, 8
-
-    # bigger
-    BS, HEADS, SEQLEN, EMB = 4, 32, 1024, 64
-
-    # llama 8B
-    #BS, HEADS, SEQLEN, EMB = 4, 32, 2048, 128
-
-    def fa():
-      Tensor.manual_seed(1337)
-      with Context(DEBUG=0): q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize() for _ in range(3)]
-      return q.scaled_dot_product_attention(k, v).realize()
-
-    with Context(DEBUG=4):
-      GlobalCounters.reset()
-      ret = fa()
-    with Context(RANGEIFY=0):
-      with Context(DEBUG=2):
-        GlobalCounters.reset()
-        cmp = fa()
-      with Context(DEBUG=0):
-        mse = ((cmp-ret)**2).sum().item()
-    print(f"mse: {mse}")
-    self.assertLessEqual(mse, 1e-6)
-
-from tinygrad import dtypes
-from tinygrad.uop.ops import UOp
-
-# contiguous + reduce can support ranges?
-
-@unittest.skipIf(RANGEIFY<1, "tests only for RANGEIFY")
-class TestOuterworld(unittest.TestCase):
-  def test_passthrough_range(self):
-    t = Tensor.rand(10, 10).realize()
-
-    # passthrough ranges
-    a = UOp.range(dtypes.int, 10, -1)
-    sel = t[a]
-    cpy = sel.contiguous(a).realize()
-
-    self.assertTrue((t==cpy).all().item())
-
-  def test_flip_range(self):
-    t = Tensor.rand(10, 10).realize()
-
-    # passthrough ranges
-    a = UOp.range(dtypes.int, 10, -1)
-    sel = t[9-a]
-    cpy = sel.contiguous(a).realize()
-
-    self.assertTrue((t.flip(0)==cpy).all().item())
-
-  def test_vmap(self):
-    def f(x): return x.sum(axis=0)*2
-
-    x = Tensor.ones(3, 10, 2).contiguous()
-
-    # vmap across axis 0
-    a = UOp.range(dtypes.int, 3, -1)
-    out = f(x[a])
-    out = out.contiguous(a)
-
-    # 3x2 grid of 20
-    out.realize()
-    print(out.numpy())
-
-  def test_triple_gemm(self):
-    x = Tensor.rand(1, 16).realize()
-    W = Tensor.rand(3, 16, 16).realize()
-
-    manual = (x @ W[0] @ W[1] @ W[2]).contiguous().realize()
-
-    a = UOp.range(dtypes.int, 3, -1)
-    x = x.assign(x @ W[a])
-    out = x.contiguous(a)[-1].contiguous().realize()
-
-    self.assertTrue((manual==out).all().item())
-
-if __name__ == '__main__':
-  unittest.main()

test/test_schedule.py

@@ -1050,14 +1050,6 @@ class TestSchedule(unittest.TestCase):
       compare = torch.nn.functional.scaled_dot_product_attention(torch.tensor(q.numpy()),torch.tensor(k.numpy()),torch.tensor(v.numpy()))
       np.testing.assert_allclose(out.numpy(), compare.numpy(), atol=1e-6, rtol=1e-3)
 
-    with Context(FUSE_ATTENTION=1):
-      out = Tensor.scaled_dot_product_attention(q,k,v)
-      run_schedule(check_schedule(out, 1))
-      if getenv("CHECK", 1):
-        import torch
-        compare = torch.nn.functional.scaled_dot_product_attention(torch.tensor(q.numpy()),torch.tensor(k.numpy()),torch.tensor(v.numpy()))
-        np.testing.assert_allclose(out.numpy(), compare.numpy(), atol=1e-6, rtol=1e-3)
-
   def test_ugly_reduceop_pairing(self):
     Tensor.manual_seed(0)
     a = Tensor.randn(4, 32).realize()

test/test_softmax_fusion.py

@@ -41,7 +41,7 @@ class TestFuse(unittest.TestCase):
 
   def test_fuse_norm(self):
     a = Tensor.rand(50,50).realize()
-    self._test_fuse(lambda a: a / a.mean(axis=1), a)
+    self._test_fuse(lambda a: a / a.mean(axis=1), a, atol=1e-6)
 
   def test_fuse_argmax(self):
     a = Tensor.rand(50,50).realize()
@@ -163,7 +163,7 @@ class TestSoftmaxFusion(unittest.TestCase):
       out = single_kernel_softmax(self.test)
       out.realize()
 
-    np.testing.assert_allclose(sout.numpy(), out.numpy(), atol=3e-7)
+    np.testing.assert_allclose(sout.numpy(), out.numpy())
 
   def test_auto_softmax(self):
     print("*** softmax ***")
@@ -176,7 +176,7 @@ class TestSoftmaxFusion(unittest.TestCase):
       out = self.test.contiguous().softmax(-1).fuse()
       run_one_schedule_item(out)
 
-    np.testing.assert_allclose(sout.numpy(), out.numpy(), atol=3e-7)
+    np.testing.assert_allclose(sout.numpy(), out.numpy())
 
   @unittest.skip("recursion error no longer raised")
   def test_softmax_bw(self):

test/test_symbolic_ops.py

@@ -229,12 +229,12 @@ class TestSymbolicOps(unittest.TestCase):
           np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_var(self):
-    a = Tensor.rand(10, 3)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
       for axis in [None, 0, 1]:
-        expected = a[:i, :].var(axis).numpy()
-        symbolic = a[:vi, :].var(axis).reshape(expected.shape).numpy()
+        a = Tensor.rand(i, 3)
+        expected = a.var(axis).numpy()
+        symbolic = a.reshape(vi, 3).var(axis).reshape(expected.shape).numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_var_2d(self):

test/test_tensor.py

@@ -415,21 +415,6 @@ class TestTinygrad(unittest.TestCase):
         data = _generate_data(depth)
         np.testing.assert_allclose(Tensor(data).numpy(), np.array(data))
 
-  def test_tensor_list_implicit_cast(self):
-    data = [True, False]
-    np.testing.assert_equal(Tensor(data, dtype=dtypes.int).numpy(), torch.tensor(data, dtype=torch.int).numpy())
-    np.testing.assert_equal(Tensor(data, dtype=dtypes.uint8).numpy(), torch.tensor(data, dtype=torch.uint8).numpy())
-    np.testing.assert_equal(Tensor(data, dtype=dtypes.float).numpy(), torch.tensor(data, dtype=torch.float).numpy())
-    data = [-1, 0, 1, 2, 3]
-    np.testing.assert_equal(Tensor(data, dtype=dtypes.int).numpy(), torch.tensor(data, dtype=torch.int).numpy())
-    np.testing.assert_equal(Tensor(data, dtype=dtypes.uint8).numpy(), torch.tensor(data, dtype=torch.uint8).numpy())
-    np.testing.assert_equal(Tensor(data, dtype=dtypes.float).numpy(), torch.tensor(data, dtype=torch.float).numpy())
-    data = [-3.5, -2.5, -1.5, 0, 1.5, 2.5, 3.5]
-    np.testing.assert_equal(Tensor(data, dtype=dtypes.int).numpy(), torch.tensor(data, dtype=torch.int).numpy())
-    # NOTE: torch and jax raise OverflowError: Python integer -3 out of bounds for uint8
-    # np.testing.assert_equal(Tensor(data, dtype=dtypes.uint8).numpy(), torch.tensor(data, dtype=torch.uint8).numpy())
-    np.testing.assert_equal(Tensor(data, dtype=dtypes.float).numpy(), torch.tensor(data, dtype=torch.float).numpy())
-
   def test_tensor_list_special_values(self):
     if is_dtype_supported(dtypes.float16):
       data = [math.nan, -math.inf, 65504, 65519, 65519.999, 65520, 65520.1]

test/test_tiny.py

@@ -30,10 +30,7 @@ class TestTiny(unittest.TestCase):
   def test_gemm(self, N=64, out_dtype=dtypes.float):
     a = Tensor.ones(N,N).contiguous()
     b = Tensor.eye(N).contiguous()
-    lst = (out:=a@b).tolist()
-    for y in range(N):
-      for x in range(N):
-        self.assertEqual(lst[y][x], 1.0, msg=f"mismatch at ({y},{x})")
+    self.assertListEqual((out:=a@b).flatten().tolist(), [1.0]*(N*N))
     if IMAGE < 2: self.assertEqual(out.dtype, out_dtype)
 
   # *** randomness ***
@@ -76,17 +73,17 @@ class TestTiny(unittest.TestCase):
 
   def test_symbolic(self):
     i = Variable('i', 1, 10)
-    ones = Tensor.ones(10).contiguous()
-    for s in [2,5]:
-      ret = ones[:i.bind(s)] + 1
-      self.assertListEqual(ret.contiguous().reshape(s).tolist(), [2.0]*s)
+    with Context(IGNORE_OOB=1):
+      for s in [2,5]:
+        ret = Tensor.ones(s).contiguous().reshape(i.bind(s)) + 1
+        self.assertListEqual(ret.reshape(s).tolist(), [2.0]*s)
 
   def test_symbolic_reduce(self):
     i = Variable('i', 1, 10)
-    ones = Tensor.ones(10).contiguous()
-    for s in [2,5]:
-      ret = ones[:i.bind(s)].sum()
-      self.assertEqual(ret.item(), s)
+    with Context(IGNORE_OOB=1):
+      for s in [2,5]:
+        ret = Tensor.ones(s).contiguous().reshape(i.bind(s)).sum()
+        self.assertEqual(ret.item(), s)
 
   # *** a model ***
 

test/test_uop_graph.py

@@ -6,7 +6,7 @@ from tinygrad.helpers import DEBUG, Context
 from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, track_rewrites, graph_rewrite, GroupOp
 from tinygrad.uop.symbolic import sym
 from tinygrad.codegen import full_rewrite, full_rewrite_to_sink
-from tinygrad.codegen.late.expander import expander
+from tinygrad.codegen.expander import expander
 
 simple_pm = PatternMatcher([
   (UPat.cvar('x', dtypes.int), lambda x: UOp.const(dtypes.float, 1.0) + UOp.const(dtypes.float, 2.0)),
@@ -441,16 +441,18 @@ class TestUOpGraph(unittest.TestCase):
       ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(Variable("i", 0, 20)),))
       with self.assertRaises(RuntimeError): to_uops_list([ld0])
 
+  @unittest.skip("outdated")
   def test_in_out_of_bounds_access_gated_store(self):
     with Context(IGNORE_OOB=0):
-      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), src=(), arg=0)
+      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
       v = Variable("v", 0, 20)
-      st0 = UOp(Ops.STORE, dtypes.void, src=(glbl0.index(v), UOp.const(dtypes.int, 0), UOp(Ops.IF, src=(v<16,))))
+      st0 = UOp(Ops.STORE, dtypes.void, (glbl0.index(v), UOp.const(dtypes.int, 0), v<16))
       to_uops_list([st0])
 
       st1 = UOp(Ops.STORE, dtypes.void, (glbl0.index(v), v, v<20))
       with self.assertRaises(RuntimeError): to_uops_list([st1])
 
+  @unittest.skip("outdated")
   def test_in_bounds_access_gated_local(self):
     with Context(IGNORE_OOB=0):
       # Define buffers
@@ -463,7 +465,7 @@ class TestUOpGraph(unittest.TestCase):
 
       gate = (gidx<400) & (lidx<8)
 
-      local_store = UOp(Ops.STORE, dtypes.void, (sbuf.index(lidx), UOp.const(dtypes.uint, 1), UOp(Ops.IF, src=(lidx<8,))))
+      local_store = UOp(Ops.STORE, dtypes.void, (sbuf.index(lidx), UOp.const(dtypes.uint, 1), lidx<8))
 
       barrier = UOp(Ops.BARRIER, dtypes.void, (local_store,))
       if_barrier = UOp(Ops.IF, dtypes.void, (gate, barrier))
@@ -475,34 +477,6 @@ class TestUOpGraph(unittest.TestCase):
       global_store = UOp(Ops.STORE, dtypes.void, (gbuf.index(gidx), local_load))
       to_uops_list([global_store])
 
-  def test_load_with_float_in_index(self):
-    with Context(IGNORE_OOB=0):
-      ridx = UOp.range(dtypes.int, 20, 0)
-      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
-      i = (ridx.cast(dtypes.float)*0.68).trunc().cast(dtypes.int)
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(i, ((0<=i)&(i<16))),))
-      to_uops_list([ld0])
-      glblfloat = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(20), (), 0)
-      ldfloat = UOp(Ops.LOAD, dtypes.float, (glblfloat.index(ridx),))
-      i = (ldfloat+3.14).cast(dtypes.int)
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(i, ((0<=i)&(i<16))),))
-
-  def test_load_cast_to_bool(self):
-    with Context(IGNORE_OOB=0):
-      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1), (), 0)
-      ridx = UOp.range(dtypes.int, 20, 0)
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(ridx, ridx.cast(dtypes.bool).logical_not()),))
-      to_uops_list([ld0])
-
-  @unittest.skip("Bool load is not supported yet")
-  def test_load_mask(self):
-    with Context(IGNORE_OOB=0):
-      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
-      mask = UOp(Ops.DEFINE_GLOBAL, dtypes.bool.ptr(16), (), 0)
-      ridx = UOp.range(dtypes.int, 20, 0)
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(ridx, ridx<16&mask),)))
-      to_uops_list([ld0])
-
   def test_out_of_bounds_off_by_one_access(self):
     with Context(IGNORE_OOB=0):
       glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
@@ -591,9 +565,10 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_switched_range_order(self):
     glbl = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
+    c2 = UOp.const(dtypes.int, 2)
     cf = UOp.const(dtypes.float, 0.0)
-    r1 = UOp.range(dtypes.int, 2, 0)
-    r2 = UOp.range(dtypes.int, 2, 1)
+    r1 = UOp(Ops.RANGE, dtypes.int, (c2,), 0)
+    r2 = UOp(Ops.RANGE, dtypes.int, (c2,), 1)
     alu = UOp(Ops.MUL, dtypes.int, (r2, r1))
     store = UOp(Ops.STORE, dtypes.void, (glbl.index(alu), cf))
     uops = to_uops_list([store])

test/test_uops.py

@@ -402,14 +402,6 @@ class TestAssembly(unittest.TestCase):
     self.assertIn(Ops.SHR, ops)
     self.assertNotIn(Ops.IDIV, ops)
 
-  def test_fast_idiv_remove_powers_of_two(self):
-    ridx = UOp.range(dtypes.int, 2**20, 0)
-    uops = to_uops_list([ridx//(7*64)], opts=Device[Device.DEFAULT].renderer)
-    ops = [x.op for x in uops]
-    # this requires shifting out the powers of two before doing fast_idiv
-    # (((ridx0>>6)*18725)>>17) instead of (int)((((long)(ridx0)*1198373)>>29))
-    self.assertNotIn(Ops.CAST, ops)
-
   def test_mulacc_unrolled(self):
     # test that     acc = acc + a0*b0 + a1*b1 + a2*b2 + a3*b3
     # is not        acc = acc + (a0*b0 + a1*b1 + a2*b2 + a3*b3)

test/unit/test_block_reorder.py

@@ -1,7 +1,7 @@
 import unittest, random
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import print_uops, UOp, Ops
-from tinygrad.codegen.late.linearize import block_reorder
+from tinygrad.codegen.linearize import block_reorder
 from tinygrad.renderer.cstyle import OpenCLRenderer
 
 def is_toposorted(lst:list[UOp]):

test/unit/test_dtype.py

@@ -56,7 +56,6 @@ class TestCastConvenienceMethod(unittest.TestCase):
 class TestDtypeTolist(unittest.TestCase):
   def test_bfloat16(self):
     self.assertEqual(Tensor([-60000, 1.5, 3.1, 60000], device="PYTHON", dtype=dtypes.bfloat16).tolist(), [-59904.0, 1.5, 3.09375, 59904.0])
-  def test_fp8(self):
     # 448
     self.assertEqual(Tensor([-30000, 1.5, 3.1, 30000], device="PYTHON", dtype=dtypes.fp8e4m3).tolist(), [-448.0, 1.5, 3.0, 448.0])
     # 57344

test/unit/test_dtype_spec.py

@@ -1,6 +1,6 @@
-import unittest, math, operator, subprocess, struct
+import unittest, math, operator, subprocess
 from tinygrad.tensor import Tensor, dtypes, Device
-from tinygrad.dtype import DType, DTYPES_DICT, truncate, truncate_fp16, float_to_bf16, _to_np_dtype, least_upper_dtype, least_upper_float
+from tinygrad.dtype import DType, DTYPES_DICT, truncate, truncate_fp16, truncate_bf16, _to_np_dtype, least_upper_dtype, least_upper_float
 from tinygrad.device import is_dtype_supported
 from tinygrad.helpers import getenv, CI, DEBUG
 from hypothesis import given, settings, strategies as strat
@@ -26,9 +26,6 @@ def _assert_eq(tensor:Tensor, target_dtype:DType, target, tol_target_dtype:float
   except AssertionError as e:
     raise AssertionError(f"\ntensor {tensor.numpy()} dtype {tensor.dtype} does not match target {target} with dtype {target_dtype}") from e
 
-def u32_to_f32(u): return struct.unpack('f', struct.pack('I', u))[0]
-def f32_to_u32(f): return struct.unpack('I', struct.pack('f', f))[0]
-
 class TestHelpers(unittest.TestCase):
   signed_ints = (dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64)
   uints = (dtypes.uint8, dtypes.uint16, dtypes.uint32, dtypes.uint64)
@@ -105,79 +102,18 @@ class TestHelpers(unittest.TestCase):
     self.assertEqual(truncate_fp16(65504), 65504)
     self.assertEqual(truncate_fp16(65519.999), 65504)
     self.assertEqual(truncate_fp16(65520), math.inf)
-    self.assertEqual(truncate_fp16(1e-8), 0.0)
-    self.assertEqual(truncate_fp16(-65504), -65504)
-    self.assertEqual(truncate_fp16(-65519.999), -65504)
-    self.assertEqual(truncate_fp16(-65520), -math.inf)
-    self.assertTrue(math.isnan(truncate_fp16(math.nan)))
 
-  def test_float_to_bf16(self):
-    # TODO: fuzz this better
+  def test_truncate_bf16(self):
+    self.assertEqual(truncate_bf16(1), 1)
+    self.assertAlmostEqual(truncate_bf16(1.1), 1.09375, places=7)
+    for a in [1234, 23456, -777.777]:
+      self.assertEqual(truncate_bf16(a), torch.tensor([a], dtype=torch.bfloat16).item())
+    # TODO: torch bfloat 1.1 gives 1.1015625 instead of 1.09375
     max_bf16 = torch.finfo(torch.bfloat16).max
-    for a in [1, 1.1, 1234, 23456, -777.777, max_bf16, max_bf16 * 1.00001, -max_bf16, -max_bf16 * 1.00001, math.inf, -math.inf]:
-      self.assertEqual(float_to_bf16(a), torch.tensor([a], dtype=torch.bfloat16).item())
-    self.assertTrue(math.isnan(float_to_bf16(math.nan)))
-
-  def test_float_to_bf16_nan(self):
-    # In f32, NaN = exp 0xFF and mantissa ≠ 0. Quiet-vs-signaling is bit 22 of the mantissa: 1 = qNaN, 0 = sNaN.
-    # qNaN(+/-), sNaN(+/-) overflow(+/-)
-    patterns = [0x7FC00001, 0xFFC00001, 0x7F800001, 0xFF800001, 0x7FFFFFFF, 0xFFFFFFFF]
-    for u in patterns:
-      x = u32_to_f32(u)
-      y = float_to_bf16(x)
-      t = torch.tensor([x], dtype=torch.bfloat16).item()
-      self.assertTrue(math.isnan(y))
-      self.assertTrue(math.isnan(t))
-
-  def test_float_to_bf16_round(self):
-    # round_to_nearest_even
-    uppers = [0x3f800000, 0x41230000, 0xC1460000] # 1.0, 10.1875, -12.375
-    for upper in uppers:
-      base = upper & 0xFFFF0000
-      base_f32 = u32_to_f32(base)
-      base_f32_round_up = u32_to_f32(base + 0x00010000)
-
-      # low < 0x8000(0.5ULP) -> round down
-      x = u32_to_f32(base | 0x00007000)
-      self.assertEqual(float_to_bf16(x), base_f32)
-      self.assertEqual(torch.tensor([x], dtype=torch.bfloat16).item(), base_f32)
-
-      # low > 0x8000(0.5ULP) -> round up
-      x = u32_to_f32(base | 0x0000C000)
-      self.assertEqual(float_to_bf16(x), base_f32_round_up)
-      self.assertEqual(torch.tensor([x], dtype=torch.bfloat16).item(), base_f32_round_up)
-
-      # low == 0x8000(0.5ULP) and LSB even -> round down
-      if ((upper >> 16) & 1) == 0:
-        x = u32_to_f32(base | 0x00008000)
-        self.assertEqual(float_to_bf16(x), base_f32)
-        self.assertEqual(torch.tensor([x], dtype=torch.bfloat16).item(), base_f32)
-      # low == 0x8000(0.5ULP) and LSB odd -> round up
-      else:
-        x = u32_to_f32(base | 0x00008000)
-        self.assertEqual(float_to_bf16(x), base_f32_round_up)
-        self.assertEqual(torch.tensor([x], dtype=torch.bfloat16).item(), base_f32_round_up)
-
-  def test_float_to_bf16_boundary(self):
-    # bf16 max finite: exp=0xFE, faction=0x7F => 0x7F7F0000(f32)
-    # bf16 inf(+/-):   exp=0xFF
-    base = 0x7F7F0000
-    inf_u32 = 0x7F800000
-
-    # low < 0.5ULP
-    x = u32_to_f32(base | 0x00007FFF)
-    self.assertEqual(f32_to_u32(float_to_bf16(x)), base)
-    self.assertEqual(f32_to_u32(torch.tensor([x], dtype=torch.bfloat16).item()), base)
-
-    # low > 0.5ULP -> overflows to +inf
-    x = u32_to_f32(base | 0x0000C000)
-    self.assertEqual(f32_to_u32(float_to_bf16(x)), inf_u32)
-    self.assertEqual(f32_to_u32(torch.tensor([x], dtype=torch.bfloat16).item()), inf_u32)
-
-    # low == 0.5ULP and LSB odd -> overflows to +inf
-    x = u32_to_f32(base | 0x00008000)
-    self.assertEqual(f32_to_u32(float_to_bf16(x)), inf_u32)
-    self.assertEqual(f32_to_u32(torch.tensor([x], dtype=torch.bfloat16).item()), inf_u32)
+    self.assertEqual(truncate_bf16(max_bf16), max_bf16)
+    self.assertEqual(truncate_bf16(min_bf16:=-max_bf16), min_bf16)
+    self.assertEqual(truncate_bf16(max_bf16 * 1.00001), math.inf)
+    self.assertEqual(truncate_bf16(min_bf16 * 1.00001), -math.inf)
 
   @given(strat.floats(width=32, allow_subnormal=True, allow_nan=True, allow_infinity=True))
   def test_truncate_fp8e4m3(self, x):

test/unit/test_gguf.py

@@ -53,37 +53,11 @@ class TestGGUF(unittest.TestCase):
   def test_load_tinyllama_q4_0(self): self._test_gguf_load("https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories15M-q4_0.gguf?download=true")
   def test_load_gpt2_q4_1(self): self._test_gguf_load("https://huggingface.co/PrunaAI/gpt2-GGUF-smashed/resolve/main/gpt2.Q4_1.gguf?download=true")
   def test_load_sample_q6_k(self): self._test_gguf_load("https://huggingface.co/Isotr0py/test-gguf-sample/resolve/main/Quant_Q6_K_1024.gguf?download=true")
-  def test_load_sample_mxfp4(self): self._test_gguf_load("https://huggingface.co/ngxson/boring-testing-tiny/resolve/main/stories260K-mxfp4.gguf?download=true")
 
   def test_dequantization_q4_0(self): self._test_dequantization(ggml.GGML_TYPE_Q4_0)
   def test_dequantization_q4_1(self): self._test_dequantization(ggml.GGML_TYPE_Q4_1)
   def test_dequantization_q8_0(self): self._test_dequantization(ggml.GGML_TYPE_Q8_0)
   def test_dequantization_q6_k(self): self._test_dequantization(ggml.GGML_TYPE_Q6_K)
-  def test_dequantization_mxfp4(self):
-    MXFP4 = 39
-
-    def encode(nibbles, E):
-      packed = [(low & 0xF) | ((high & 0xF) << 4) for low, high in zip(nibbles[:16], nibbles[16:])]
-      return np.array([E] + packed, dtype=np.uint8)
-
-    def decode(code, E):
-      sign = -1.0 if code * 0b1000 else 1.0
-      exp = (code >> 1) & 0b11
-      mant = code & 0b1
-      val = (1.0 + 0.5 * mant) * np.exp2(exp - 1) if exp else 0.5 * mant
-      scale = np.exp2(E - 128) if E >= 2 else np.exp2(-127 if E == 1 else -128)
-      return sign * val * scale
-
-    blocks, expected = [], []
-    rng = np.random.default_rng(42)
-    for _ in range(4):
-      E = rng.integers(0, 256)
-      codes = rng.integers(0, 16, size=32, dtype=np.uint8)
-      blocks.append(encode(codes, E))
-      expected.extend(decode(c, E) for c in codes)
-    tensor = Tensor(np.concatenate(blocks))
-    out = ggml_data_to_tensor(tensor, len(expected), MXFP4)
-    self.assertListEqual(out.numpy().tolist(), np.array(expected, dtype=np.float32).tolist())
 
   def test_expected_failure_unknown_type(self):
     with self.assertRaises(ValueError):

test/unit/test_shapetracker.py

@@ -6,7 +6,7 @@ from tinygrad.helpers import prod
 from tinygrad.shape.shapetracker import ShapeTracker, View
 from tinygrad import Variable
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite
-from tinygrad.codegen.late.devectorizer import sym
+from tinygrad.codegen.devectorizer import sym
 from itertools import product
 
 def shapetracker_getitem(st:ShapeTracker, val:int):

test/unit/test_simplify_valid_idx.py

@@ -19,7 +19,7 @@ def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UO
 
 def Special(expr, nmax): return UOp(Ops.SPECIAL, dtypes.int, (), (expr, nmax))
 def Variable(expr, nmin, nmax): return UOp.variable(expr, nmin, nmax)
-def Range(n, nmax): return UOp.range(dtypes.int, nmax, n)
+def Range(n, nmax): return UOp(Ops.RANGE, dtypes.int, arg=n, src=(UOp.const(dtypes.int, nmax),))
 
 class TestHelpers(unittest.TestCase):
   def test_is_increasing(self):

test/unit/test_transcendental_helpers.py

@@ -2,7 +2,7 @@ import unittest, math
 import numpy as np
 from tinygrad import dtypes
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.uop.decompositions import TRANSCENDENTAL_DTYPES, payne_hanek_reduction, cody_waite_reduction
+from tinygrad.uop.decompositions import TRANSCENDENTAL_SUPPORTED_DTYPES, payne_hanek_reduction, cody_waite_reduction
 from tinygrad.uop.decompositions import frexp, rintk, xpow, xexp2, xlog2, trig_poly, pow2if
 from test.helpers import eval_uop
 
@@ -89,7 +89,7 @@ class TestTranscendentalVectorizedFunctions(unittest.TestCase):
       assert u1.op == u2.op, f'expected {u1.op=} but got {u2.op=} for UOps\n{u1=}\n{u2}'
     [self._check_uops_match(x1, x2) for x1, x2 in zip((u1 if isinstance(u1, tuple) else u1.src), (u2 if isinstance(u2, tuple) else u2.src))]
 
-  def _test_vectorized(self, fxn, scalar_dtypes=TRANSCENDENTAL_DTYPES, vals=[-2,1.3,194], vcounts=[1,4,19]):
+  def _test_vectorized(self, fxn, scalar_dtypes=TRANSCENDENTAL_SUPPORTED_DTYPES, vals=[-2,1.3,194], vcounts=[1,4,19]):
     for scalar_dtype in scalar_dtypes:
       for val in vals:
         for vcount in vcounts:

test/unit/test_uop_spec.py

@@ -81,16 +81,5 @@ class TestUOpSpec(unittest.TestCase):
     with self.assertRaisesRegex(RuntimeError, "UOp verification failed"):
       type_verify([a], tensor_uop_spec)
 
-class TestUOpSink(unittest.TestCase):
-  def test_0(self):
-    s = UOp.sink()
-    self.assertEqual(len(s.src), 0)
-
-  def test_1(self):
-    a = UOp.const(dtypes.int, 0)
-    s1 = UOp.sink(a)
-    s2 = a.sink()
-    self.assertIs(s1, s2)
-
 if __name__ == '__main__':
   unittest.main()

test/unit/test_uop_symbolic.py

@@ -4,11 +4,11 @@ import z3
 
 from tinygrad.dtype import dtypes, ConstType
 from tinygrad.codegen import full_rewrite
-from tinygrad.codegen.late.devectorizer import sym
+from tinygrad.codegen.devectorizer import sym
 from tinygrad.helpers import Context
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer
 from tinygrad import Variable
-from tinygrad.uop.spec import uops_to_z3
+from tinygrad.uop.spec import z3_renderer
 
 def render(self) -> tuple[str, ConstType, ConstType]:
   # NOTE: we need STORE so the ALU op has children
@@ -32,8 +32,9 @@ class TestSymbolic(unittest.TestCase):
   def helper_test_variable(self, v, n, m, s, test_z3:bool=True):
     if test_z3:
       solver = z3.Solver()
-      expr, expr_simplified = uops_to_z3(solver, v, v.simplify())
-      self.assertEqual(solver.check(expr != expr_simplified), z3.unsat, "simplified expression not equal to original")
+      z3_sink = graph_rewrite(v.sink(v.simplify()), z3_renderer, ctx=(solver, {}))
+      expr, epxr_simplified = z3_sink.src[0].arg, z3_sink.src[1].arg
+      self.assertEqual(solver.check(expr != epxr_simplified), z3.unsat, "simplified expression not equal to original")
     rendered, nmin, nmax = render(v)
     if isinstance(s, tuple): self.assertIn(rendered, s)
     else: self.assertEqual(rendered, s)
@@ -127,8 +128,6 @@ class TestSymbolic(unittest.TestCase):
     b = Variable("b", 0, 8)
     self.helper_test_variable(a+a, 0, 16, "(a*2)")
     self.helper_test_variable((a+b)+b, 0, 24, "(a+(b*2))")
-    self.helper_test_variable((a*3+b)+a, 0, 40, "(b+(a*4))")
-    self.helper_test_variable((a+b)+a*3, 0, 40, "(b+(a*4))")
 
   def test_sub_self(self):
     a = Variable("a", 0, 8)
@@ -163,6 +162,10 @@ class TestSymbolic(unittest.TestCase):
   def test_div_remove(self):
     self.helper_test_variable(Variable("a", 0, 7) // 20, 0, 0, "0")
 
+  def test_div_min_max(self):
+    self.helper_test_variable(Variable("a", 1, 7) // 2, 0, 3, "(a//2)")
+    self.helper_test_variable(Variable("a", 0, 6) // 2, 0, 3, "(a//2)")
+
   def test_div_neg_min_max(self):
     self.helper_test_variable(Variable("a", 1, 7) // -2, -3, 0, "((a//2)*-1)")
     self.helper_test_variable(Variable("a", 0, 6) // -2, -3, 0, "((a//2)*-1)")
@@ -208,18 +211,6 @@ class TestSymbolic(unittest.TestCase):
     self.assertEqual((Variable("x", -10, 0)%Variable("y", -10, -1))._min_max, (-9, 0))
     self.assertEqual((Variable("x", -10, 0)%Variable("y", 1, 10))._min_max, (-9, 0))
 
-  def test_div_min_max(self):
-    self.helper_test_variable(Variable("a", 2, 7) // 2, 1, 3, "(a//2)")
-    self.helper_test_variable(Variable("a", 0, 6) // 2, 0, 3, "(a//2)")
-
-    self.helper_test_variable(Variable("x", 0, 10)//Variable("y", 1, 10), 0, 10, "(x//y)")
-    self.helper_test_variable(Variable("x", -10, 0)//Variable("y", 1, 10), -10, 0, "(((x*-1)//y)*-1)")
-    self.helper_test_variable(Variable("x", 0, 10)//Variable("y", -10, -1), -10, 0, "((x//(y*-1))*-1)")
-    self.helper_test_variable(Variable("x", -10, 0)//Variable("y", -10, -1), 0, 10, "((x*-1)//(y*-1))")
-
-    self.helper_test_variable(Variable("x", -10, 10)//Variable("y", 1, 10), -10, 10, "(x//y)")
-    self.helper_test_variable(Variable("x", -10, 10)//Variable("y", -10, -1), -10, 10, "((x//(y*-1))*-1)")
-
   def test_mod_factor(self):
     self.helper_test_variable(usum([Variable("a", 0, 7)*100, Variable("b", 0, 3)*50]) % 100, 0, 50, "((b%2)*50)")
 
@@ -449,8 +440,7 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable((-Variable("a", 10, 10))%7, -3, -3, "-3")
 
   def test_div_numerator_negative(self):
-    with Context(CORRECT_DIVMOD_FOLDING=1):
-      self.helper_test_variable((Variable("idx", 0, 9)*-10)//11, -8, 0, "(((idx*10)//11)*-1)")
+    self.helper_test_variable((Variable("idx", 0, 9)*-10)//11, -8, 0, "(((idx*10)//11)*-1)")
 
   def test_nest_div_negative_factor(self):
     ridx0=UOp.variable("ridx0", 0, 9)
@@ -639,16 +629,15 @@ class TestSymbolic(unittest.TestCase):
     cond = Variable("x", 0, 3) < 2
     a = Variable("a", 0, 3)
     b = Variable("b", 0, 3)
-    c = Variable("c", 0, 3)
     aa = cond.where(a, a.ufix(0))
     bb = cond.where(b, b.ufix(1))
     self.helper_test_variable(aa, 0, 3, "(a if (x<2) else 0)")
     self.helper_test_variable(bb, 0, 3, "(b if (x<2) else 1)")
     self.helper_test_variable(aa+bb, 0, 6, "((a+b) if (x<2) else 1)")
     self.helper_test_variable(aa.maximum(bb), 0, 3, "(max(a, b) if (x<2) else 1)")
-    self.helper_test_variable((c+aa)+bb, 0, 9, "(c+((a+b) if (x<2) else 1))")
 
     # not combining because it increased total ALU
+    c = Variable("c", 0, 3)
     cc = cond.where(c, c+1)
     self.helper_test_variable(bb+cc, 0, 7, "((b if (x<2) else 1)+(c if (x<2) else (c+1)))")
 

test/unit/test_viz.py

@@ -1,11 +1,11 @@
-import unittest, decimal, json, struct
+import unittest, decimal, json
 from dataclasses import dataclass
 
 from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, TrackedPatternMatcher
 from tinygrad.uop.ops import graph_rewrite, track_rewrites, TRACK_MATCH_STATS
 from tinygrad.uop.symbolic import sym
 from tinygrad.dtype import dtypes
-from tinygrad.helpers import PROFILE, colored, ansistrip, flatten, TracingKey, ProfileRangeEvent, ProfileEvent, Context
+from tinygrad.helpers import PROFILE, colored, ansistrip, flatten, TracingKey, ProfileRangeEvent
 from tinygrad.device import Buffer
 
 @track_rewrites(name=True)
@@ -240,59 +240,15 @@ class TestVizIntegration(BaseTestViz):
     self.assertEqual(lst[0]["name"], "Schedule 1 Kernel n1")
     self.assertEqual(lst[1]["name"], prg.name)
 
-  def test_metadata_tracing(self):
-    with Context(TRACEMETA=2):
-      a = Tensor.empty(1)
-      b = Tensor.empty(1)
-      metadata = (alu:=a+b).uop.metadata
-      alu.kernelize()
-      graph = next(get_details(tracked_ctxs[0][0]))["graph"]
-    self.assertEqual(len([n for n in graph.values() if repr(metadata) in n["label"]]), 1)
-
 from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry
 from tinygrad.viz.serve import get_profile
 
-class TinyUnpacker:
-  def __init__(self, buf): self.buf, self.offset = buf, 0
-  def __call__(self, fmt:str) -> tuple:
-    ret = struct.unpack_from(fmt, self.buf, self.offset)
-    self.offset += struct.calcsize(fmt)
-    return ret
-
-# 0 means None, otherwise it's an enum value
-def option(i:int) -> int|None: return None if i == 0 else i-1
-
-def load_profile(lst:list[ProfileEvent]) -> dict:
-  ret = get_profile(lst)
-  u = TinyUnpacker(ret)
-  dur, global_peak, index_len, layout_len = u("<IQII")
-  strings, dtypes = json.loads(ret[u.offset:u.offset+index_len]).values()
-  u.offset += index_len
-  layout:dict[str, dict] = {}
-  for _ in range(layout_len):
-    klen = u("<B")[0]
-    k = ret[u.offset:u.offset+klen].decode()
-    u.offset += klen
-    layout[k] = v = {"shapes":[]}
-    event_type, event_count = u("<BI")
-    if event_type == 0:
-      for _ in range(event_count):
-        name, ref, st, dur, cat, _ = u("<IIIfBI")
-        v["shapes"].append({"name":strings[name], "ref":option(ref), "st":st, "dur":dur, "cat":option(cat)})
-    else:
-      v["peak"] = u("<Q")[0]
-      for _ in range(event_count):
-        alloc, ts, key = u("<BII")
-        if alloc: v["shapes"].append({"event":"alloc", "ts":ts, "key":key, "arg": {"dtype":strings[u("<I")[0]], "sz":u("<Q")[0]}})
-        else: v["shapes"].append({"event":"free", "ts":ts, "key":key})
-  return {"dur":dur, "peak":global_peak, "layout":layout}
-
 class TestVizProfiler(unittest.TestCase):
   def test_perfetto_node(self):
     prof = [ProfileRangeEvent(device='NV', name='E_2', st=decimal.Decimal(1000), en=decimal.Decimal(1010), is_copy=False),
             ProfileDeviceEvent(device='NV', comp_tdiff=decimal.Decimal(-1000), copy_tdiff=decimal.Decimal(-100))]
 
-    j = load_profile(prof)
+    j = json.loads(get_profile(prof))
 
     dev_events = j['layout']['NV']['shapes']
     self.assertEqual(len(dev_events), 1)
@@ -300,24 +256,18 @@ class TestVizProfiler(unittest.TestCase):
     self.assertEqual(event['name'], 'E_2')
     self.assertEqual(event['st'], 0)
     self.assertEqual(event['dur'], 10)
-    assert event['ref'] is None
 
   def test_perfetto_copy_node(self):
     prof = [ProfileRangeEvent(device='NV', name='COPYxx', st=decimal.Decimal(1000), en=decimal.Decimal(1010), is_copy=True),
-            ProfileRangeEvent(device='NV:2', name='COPYxx', st=decimal.Decimal(1000), en=decimal.Decimal(1010), is_copy=True),
-            ProfileDeviceEvent(device='NV', comp_tdiff=decimal.Decimal(-1000), copy_tdiff=decimal.Decimal(-100)),
-            ProfileDeviceEvent(device='NV:2', comp_tdiff=decimal.Decimal(-800), copy_tdiff=decimal.Decimal(-80))]
+            ProfileDeviceEvent(device='NV', comp_tdiff=decimal.Decimal(-1000), copy_tdiff=decimal.Decimal(-100))]
 
-    j = load_profile(prof)
+    j = json.loads(get_profile(prof))
 
     event = j['layout']['NV']['shapes'][0]
     self.assertEqual(event['name'], 'COPYxx')
-    self.assertEqual(event['st'], 0)   # first event
+    self.assertEqual(event['st'], 900) # diff clock
     self.assertEqual(event['dur'], 10)
 
-    event2 = j['layout']['NV:2']['shapes'][0]
-    self.assertEqual(event2['st'], 20) # second event, diff clock
-
   def test_perfetto_graph(self):
     prof = [ProfileDeviceEvent(device='NV', comp_tdiff=decimal.Decimal(-1000), copy_tdiff=decimal.Decimal(-100)),
             ProfileDeviceEvent(device='NV:1', comp_tdiff=decimal.Decimal(-500), copy_tdiff=decimal.Decimal(-50)),
@@ -326,12 +276,12 @@ class TestVizProfiler(unittest.TestCase):
                               deps=[[], [0]],
                               sigs=[decimal.Decimal(1000), decimal.Decimal(1002), decimal.Decimal(1004), decimal.Decimal(1008)])]
 
-    j = load_profile(prof)
+    j = json.loads(get_profile(prof))
 
     tracks = list(j['layout'])
     self.assertEqual(tracks[0], 'NV Graph')
-    self.assertEqual(tracks[1], 'NV')
-    self.assertEqual(tracks[2], 'NV:1')
+    self.assertEqual(tracks[2], 'NV')
+    self.assertEqual(tracks[4], 'NV:1')
 
     nv_events = j['layout']['NV']['shapes']
     self.assertEqual(nv_events[0]['name'], 'E_25_4n2')
@@ -348,22 +298,6 @@ class TestVizProfiler(unittest.TestCase):
     self.assertEqual(graph_events[0]['st'], nv_events[0]['st'])
     self.assertEqual(graph_events[0]['st']+graph_events[0]['dur'], nv1_events[0]['st']+nv1_events[0]['dur'])
 
-  def test_bytes_per_kernel(self):
-    step = 10
-    n_events = 1_000
-    prof = [ProfileRangeEvent("CPU", name="k_test", st=decimal.Decimal(ts:=i*step), en=decimal.Decimal(ts)+step) for i in range(n_events)]
-    sz = len(get_profile(prof))
-    self.assertLessEqual(sz/n_events, 26)
-
-  # can pack up to 1hr 11 min of trace events
-  def test_trace_duration(self):
-    dur_mins = 72
-    n_events = 1_000
-    step = decimal.Decimal(dur_mins*60*1e6//n_events)
-    prof = [ProfileRangeEvent("CPU", name="k_test", st=decimal.Decimal(ts:=i*step), en=decimal.Decimal(ts)+step) for i in range(n_events)]
-    with self.assertRaises(struct.error):
-      get_profile(prof)
-
 def _alloc(b:int):
   a = Tensor.empty(b, device="NULL", dtype=dtypes.char)
   a.uop.buffer.allocate()
@@ -373,29 +307,38 @@ class TestVizMemoryLayout(BaseTestViz):
   def test_double_alloc(self):
     a = _alloc(1)
     _b = _alloc(1)
-    profile_ret = load_profile(Buffer.profile_events)
+    profile_ret = json.loads(get_profile(Buffer.profile_events))
     ret = profile_ret["layout"][f"{a.device} Memory"]
     self.assertEqual(ret["peak"], 2)
-    self.assertEqual(len(ret["shapes"]), 2)
+    self.assertEqual(ret["shapes"][0]["x"], [0, 2])
+    self.assertEqual(ret["shapes"][1]["x"], [1, 2])
 
   def test_del_once(self):
     a = _alloc(1)
     del a
     b = _alloc(1)
-    profile_ret = load_profile(Buffer.profile_events)
+    profile_ret = json.loads(get_profile(Buffer.profile_events))
     ret = profile_ret["layout"][f"{b.device} Memory"]
     self.assertEqual(ret["peak"], 1)
-    self.assertEqual(len(ret["shapes"]), 3)
+    self.assertEqual(ret["shapes"][0]["x"], [0, 2])
+    self.assertEqual(ret["shapes"][1]["x"], [2, 3])
+    self.assertEqual(ret["shapes"][0]["y"], [0, 0])
+    self.assertEqual(ret["shapes"][1]["y"], [0, 0])
 
   def test_alloc_free(self):
     a = _alloc(1)
     _b = _alloc(1)
     del a
     c = _alloc(1)
-    profile_ret = load_profile(Buffer.profile_events)
+    profile_ret = json.loads(get_profile(Buffer.profile_events))
     ret = profile_ret["layout"][f"{c.device} Memory"]
     self.assertEqual(ret["peak"], 2)
-    self.assertEqual(len(ret["shapes"]), 4)
+    self.assertEqual(ret["shapes"][0]["x"], [0, 3])
+    self.assertEqual(ret["shapes"][1]["x"], [1, 3, 3, 4])
+    self.assertEqual(ret["shapes"][0]["y"], [0, 0])
+    self.assertEqual(ret["shapes"][1]["y"], [1, 1, 0, 0])
+    self.assertEqual(ret["shapes"][2]["x"], [3, 4])
+    self.assertEqual(ret["shapes"][2]["y"], [1, 1])
 
 if __name__ == "__main__":
   unittest.main()

tinygrad/codegen/__init__.py

@@ -1,7 +1,7 @@
 from typing import Any, Callable
 import functools
 from dataclasses import dataclass
-from tinygrad.helpers import QUANTIZE, DEVECTORIZE, TRANSCENDENTAL, RANGEIFY, POSTOPT
+from tinygrad.helpers import QUANTIZE, DEVECTORIZE, TRANSCENDENTAL
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp
 from tinygrad.uop.spec import type_verify
 from tinygrad.renderer import Renderer
@@ -9,18 +9,15 @@ from tinygrad.renderer import Renderer
 # import all pattern matchers here
 from tinygrad.codegen.lowerer import pm_lowerer, get_index
 from tinygrad.codegen.quantize import pm_quant
-from tinygrad.codegen.gpudims import pm_add_gpudims, pm_tensor_cores, pm_group_for_reduce, pm_fix_locals, pm_bufferize_loop
+from tinygrad.codegen.gpudims import pm_add_gpudims
 from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing
 from tinygrad.uop.decompositions import get_late_rewrite_patterns
-from tinygrad.codegen.late.expander import migrate_indexing, expander, pm_pre_expander
-from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \
+from tinygrad.codegen.expander import migrate_indexing, expander
+from tinygrad.codegen.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \
   ReduceContext, correct_load_store, pm_render
-from tinygrad.codegen.late.linearize import block_create, pm_blockend_merge, block_merge, pm_finalize, BlockContext
-from tinygrad.codegen.opt import pm_get_optimization, pm_do_optimize, pm_postrange_opt
+from tinygrad.codegen.linearize import block_create, pm_blockend_merge, block_merge, pm_finalize, BlockContext
+from tinygrad.codegen.opt import pm_optimize
 from tinygrad.codegen.opt.swizzler import view_left, view_right, fix_kernel_ops
-from tinygrad.codegen.opt.postrange import pm_postrange_opt
-from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen
-from tinygrad.codegen.opt.postrange import pm_flatten_range
 
 @dataclass
 class RewriteStep:
@@ -47,10 +44,10 @@ rewrites_for_linearizer = [
 
 def get_rewrites_for_renderer(opts:Renderer, linearizer:bool=True) -> list[RewriteStep]:
   # cache with the values of the context vars
-  return _get_rewrites_for_renderer(opts, linearizer, QUANTIZE.value, DEVECTORIZE.value, TRANSCENDENTAL.value, RANGEIFY.value, POSTOPT.value)
+  return _get_rewrites_for_renderer(opts, linearizer, QUANTIZE.value, DEVECTORIZE.value, TRANSCENDENTAL.value)
 
 @functools.cache
-def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVECTORIZE, _TRANSCENDENTAL, _RANGEIFY, _POSTOPT) -> list[RewriteStep]:
+def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVECTORIZE, _TRANSCENDENTAL) -> list[RewriteStep]:
   # ** lowerer (rewrite_shapetracker_with_index) **
   ret: list[RewriteStep] = []
 
@@ -58,39 +55,22 @@ def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVEC
   ret.extend(rewrites_for_views)
 
   # this is kernel.py
-  if not _RANGEIFY: ret.append(RewriteStep(pm_get_optimization, ctx=lambda _: opts, name="get optimization"))
-
-  if not _POSTOPT and not _RANGEIFY: ret.append(RewriteStep(pm_do_optimize, ctx=lambda _: opts, name="optimize ast"))
+  ret.append(RewriteStep(pm_optimize, ctx=lambda _: opts, name="optimize ast"))
 
   if _QUANTIZE and opts.device in {"CPU", "DSP"}: ret.append(RewriteStep(pm_quant, name="quantize"))
   ret.append(RewriteStep(pm_lowerer, get_index, name="lowerer", bottom_up=True))
-  
-  # add tensor cores
-  if _RANGEIFY:
-    ret.append(RewriteStep(pm_bufferize_loop, name="bufferize loop"))
-    ret.append(RewriteStep(pm_tensor_cores, lambda _: ({}, opts), name="tensor cores", bottom_up=True))
-
-  if _POSTOPT or _RANGEIFY: ret.append(RewriteStep(pm_postrange_opt, ctx=lambda _: opts, name="post optimize ast"))
 
   # ** expander (expand_rewrite) **
   ret.append(RewriteStep(sym+migrate_indexing, name="initial symbolic"))
 
-  ret.append(RewriteStep(pm_fix_locals, name="fix locals"))
-
-  # add locals
-  ret.append(RewriteStep(pm_flatten_range+pm_add_buffers_local+rangeify_codegen+pm_group_for_reduce, name="add local buffers"))
-
-  # add gpu dims (late). this also handles UNROLL range
-  ret.append(RewriteStep(pm_add_gpudims, lambda _: opts, name="add gpudims"))
-
   # expand
-  ret.append(RewriteStep(sym+pm_pre_expander+expander, name="expander"))
+  ret.append(RewriteStep(sym+expander, name="expander"))
 
   # ** devectorizer (full_graph_rewrite) **
   # remove reduce
   ret.append(RewriteStep(pm_reduce+gep_pushing, lambda _: ReduceContext(), name="remove_reduce"))
 
-  # add gpu dims (late). this works after devectorize, but it's faster here
+  # add gpu dims (late)
   ret.append(RewriteStep(pm_add_gpudims, lambda _: opts, name="add gpudims"))
 
   # devectorize (TODO: does this need opts?)

tinygrad/codegen/devectorizer.py

@@ -232,21 +232,17 @@ def no_vectorized_alu(alu:UOp):
   alus = tuple(UOp(alu.op, alu.dtype.scalar(), tuple(s.gep(i) for s in alu.src), alu.arg) for i in range(alu.dtype.vcount))
   return UOp(Ops.VECTORIZE, alu.dtype, alus)
 
-def no_vectorized_buf(buf:UOp):
-  dtype = cast(PtrDType, buf.dtype)
-  return buf.replace(dtype=dtype.base.scalar().ptr(dtype.size*dtype.count, dtype.addrspace)).cast(dtype)
-
-def no_vectorized_index(buf:UOp, cast:UOp, idx:UOp):
-  cnt = cast.dtype.count
-  assert idx.dtype.count == 1, f"idx dtype must be 1 {idx.dtype}"
-  return buf.broadcast(cnt).index(idx.broadcast(cnt)*cnt+UOp.const(dtypes.int.vec(cnt), tuple(range(cnt))))
+def no_vectorized_acc(acc:UOp, c:UOp):
+  if acc.dtype.count == 1: return None
+  assert c.arg == 0, "this only supports index 0"
+  new_acc = acc.replace(dtype=acc.dtype.base.scalar().ptr(acc.dtype.count, cast(PtrDType, acc.dtype).addrspace))
+  return UOp(Ops.PTRCAT, acc.dtype, tuple([new_acc.index(UOp.const(dtypes.int, i)) for i in range(acc.dtype.count)]))
 
 devectorize = PatternMatcher([
   # no ALU on vectorized dtypes
   (UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST), name="alu"), no_vectorized_alu),
   (UPat(Ops.WMMA, name="wmma"), no_vectorized_wmma),
-  (UPat((Ops.DEFINE_LOCAL, Ops.DEFINE_REG), name="buf"), no_vectorized_buf),
-  (UPat((Ops.DEFINE_LOCAL, Ops.DEFINE_REG), name="buf").cast(name="cast").index(UPat.var("idx")), no_vectorized_index),
+  (UPat(Ops.DEFINE_REG, name="acc").index(UPat.cvar("c")), no_vectorized_acc),
 ])
 
 pm_render = PatternMatcher([

tinygrad/codegen/expander.py

@@ -1,9 +1,8 @@
 # this converts a lowerer program into a vectorized program
 
 import functools, itertools, operator
-from tinygrad.dtype import dtypes, PtrDType, AddrSpace
-from tinygrad.helpers import AMX, dedup, flatten, all_same, prod, partition
-from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, GroupOp, AxisType
+from tinygrad.helpers import AMX, dedup, flatten, all_same, prod
+from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, GroupOp
 
 def _expand_arg_to_idx(args:tuple[tuple[int, int], ...], rpk:dict[int, int]) -> int:
   idx, mul = 0, 1
@@ -47,13 +46,11 @@ def do_expand(root:UOp):
         new_srcs.append(src.src[0].gep(tuple(lst)))
     else:
       # non-UNROLL input
-      if root.op is Ops.IF or src.op is Ops.IF:
+      if root.op is Ops.IF:
         # for the first arg of IF, just pass them through ignoring UNROLLS
         new_srcs.append(src)
-      elif (root.op is Ops.STORE and i >= 2) or (root.op in {Ops.REDUCE, Ops.BUFFERIZE} and i >= 1) or (root.op is Ops.WMMA and i >= 3):
-        # for any range args of STORE/REDUCE, pass them through
-        new_srcs.append(src)
-      elif root.op is Ops.INDEX and i >= 1 and not isinstance(root.dtype, PtrDType):
+      elif root.op in {Ops.REDUCE, Ops.STORE} and src.op is Ops.RANGE:
+        # for any range args of REDUCE, pass them through
         new_srcs.append(src)
       elif src.dtype.count > 1:
         # put any input dtype > 1 grouped together
@@ -75,7 +72,7 @@ def do_contract(con:UOp):
   # CONTRACT without UNROLL repeats the element VECTORIZED
   if ex.op is not Ops.UNROLL: return UOp(Ops.VECTORIZE, con.dtype, con.src*con.dtype.count)
   # CONTRACT may remove several axes from UNROLL
-  assert con.dtype == dtypes.void or con.dtype.count == prod([x[1] for x in con.arg]), "dtype is wrong"
+  assert con.dtype.count == prod([x[1] for x in con.arg]), "dtype is wrong"
   idxs = []
   for rpk in _choices_from_args(new_ex_args:=tuple(x for x in ex.arg if x not in con.arg)):
     idxs += [_expand_arg_to_idx(ex.arg, {**rpk, **lrpk}) for lrpk in _choices_from_args(con.arg)]
@@ -86,7 +83,7 @@ expander = PatternMatcher([
   (UPat(Ops.UNROLL, name="outer", src=(UPat(Ops.UNROLL, name="inner"),)),
    lambda outer, inner: UOp(Ops.UNROLL, outer.dtype, (inner.src[0],), inner.arg+outer.arg)),
   # do expansion
-  (UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST, Ops.GEP, Ops.WMMA, Ops.LOAD, Ops.STORE, Ops.INDEX, Ops.BUFFERIZE,
+  (UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST, Ops.GEP, Ops.WMMA, Ops.LOAD, Ops.STORE, Ops.INDEX,
          Ops.VECTORIZE, Ops.IF, Ops.REDUCE), name="root", custom_early_reject=set([Ops.UNROLL])), do_expand),
   (UPat(Ops.CONTRACT, name="con"), do_contract),
   # BARRIERs aren't actually expanded
@@ -114,49 +111,3 @@ migrate_indexing = PatternMatcher([
   # create gate MUST BE BEFORE expander
   (UPat(Ops.STORE, name="root"), create_gate),
 ])
-
-# ****
-
-def fix_reduce_unroll(x:UOp):
-  reduce_range, reduce_expand = partition(x.src[1:], lambda y: y.op is Ops.RANGE)
-  if len(reduce_expand) == 0: return None
-  reduce_expand = [x for x in reduce_expand if x.op is not Ops.CONST]
-  assert all(x.op is Ops.UNROLL for x in reduce_expand), f"not all UNROLLS in {reduce_expand}"
-  ret = x.src[0]
-  if len(contract_axis:=flatten(x.arg for x in reduce_expand)):
-    ret = UOp(Ops.CONTRACT, x.dtype.vec(prod(x[1] for x in contract_axis)), (ret,), tuple(contract_axis), tag=1)
-  # REDUCE supports both "horizontal" reduction and range reduction. the horizontal elements are taken in the nearest group
-  return x.replace(src=(ret,)+tuple(reduce_range))
-
-def fix_store_unroll(x:UOp):
-  store_expand, store_range = partition(x.src[2:], lambda y: y.op is Ops.UNROLL)
-  if len(store_expand) == 0: return None
-  return UOp(Ops.CONTRACT, dtypes.void, (x.replace(src=x.src[:2]+tuple(store_range)),), tuple(flatten(x.arg for x in store_expand)), tag=1)
-
-def fix_group_for_reduce(x:UOp):
-  reduce_gfr, reduce_r = partition(x.src[1:], lambda u: u.op is Ops.RANGE and u.arg[1] == AxisType.GROUP_REDUCE)
-  if len(reduce_gfr) == 0: return None
-
-  # NOTE: if there's other locals here, we need them in the buffer too
-  upstream_locals = [u for u in x.toposort() if u.op is Ops.RANGE and u.arg[1] == AxisType.LOCAL]
-
-  # do only the non grouped reduces early
-  ret = x.replace(src=(x.src[0],)+tuple(reduce_r))
-  reduce_loop = [x.replace(arg=(x.arg[0]+100, AxisType.REDUCE)) for x in reduce_gfr]
-  buf = ret.bufferize(*upstream_locals, *reduce_gfr, arg=(AddrSpace.LOCAL, reduce_gfr[0].arg[0])).index(*upstream_locals, *reduce_loop)
-
-  # gate with an if on the store + do the final reduce
-  buf = UOp(Ops.IF, dtype=buf.dtype, src=(functools.reduce(operator.and_, [x.eq(0) for x in reduce_gfr]), buf))
-  return buf.reduce(*reduce_loop, arg=x.arg)
-
-pm_pre_expander = PatternMatcher([
-  # rewrite UPCAST/UNROLL range to something to be expanded
-  (UPat(Ops.RANGE, name="r"),
-   lambda r: UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(s:=r.vmax+1), tuple(range(s))),), ((r.arg[0],s),)) \
-    if r.arg[1] in {AxisType.UNROLL, AxisType.UPCAST} else None),
-  # fix REDUCEs with UNROLLs
-  (UPat(Ops.REDUCE, name="x"), fix_reduce_unroll),
-  (UPat(Ops.STORE, name="x"), fix_store_unroll),
-  # fix group for reduce
-  (UPat(Ops.REDUCE, name="x"), fix_group_for_reduce),
-])

tinygrad/codegen/gpudims.py

@@ -1,10 +1,9 @@
-import math, functools, operator
-from tinygrad.uop.ops import UOp, Ops, sint, PatternMatcher, UPat, KernelInfo, ssimplify, AxisType, graph_rewrite
-from tinygrad.helpers import all_int, partition, flatten, prod, dedup, USE_TC, DEBUG
-from tinygrad.dtype import dtypes, AddrSpace
+import math
+from tinygrad.uop.ops import UOp, Ops, sint, PatternMatcher, UPat, KernelInfo, ssimplify, AxisType
+from tinygrad.helpers import all_int
+from tinygrad.dtype import dtypes
 from tinygrad.shape.view import get_contraction
 from tinygrad.renderer import Renderer
-from tinygrad.codegen.opt.tc import TensorCore
 
 def _group_dims(dims:tuple[sint, ...], max_sizes:tuple[int, ...]):
   # TODO: symbolic shape
@@ -53,24 +52,20 @@ def get_grouped_dims(prefix, dims:tuple[sint, ...], max_sizes:tuple[int, ...]|No
 
 def add_gpudims(ctx:Renderer, s:UOp):
   if s.arg is None: return None
+  ki: KernelInfo = s.arg
+  global_dims = [i for i,x in enumerate(ki.axis_types) if x is AxisType.GLOBAL]
+  local_dims = [i for i,x in enumerate(ki.axis_types) if x in (AxisType.LOCAL, AxisType.GROUP_REDUCE)]
+  if not global_dims and not local_dims: return None
   s_topo = list(s.toposort())
   if any(x.op is Ops.SPECIAL for x in s_topo): return None
 
-  # get ranges
-  all_ranges = {x.arg[0:-1]:x for x in s_topo if x.op is Ops.RANGE}
-
-  # extract global/local dims
-  global_dims = sorted(dedup([x.arg[0:-1] for x in all_ranges.values() if x.arg[-1] is AxisType.GLOBAL]))
-  local_dims = sorted(dedup([x.arg[0:-1] for x in all_ranges.values() if x.arg[-1] in (AxisType.LOCAL, AxisType.GROUP_REDUCE)]))
-  if not global_dims and not local_dims: return None
-
   # get global and local shape
+  all_ranges = {x.arg%1000:x for x in s_topo if x.op is Ops.RANGE}
   ranges = [all_ranges[r] for r in global_dims+local_dims if r in all_ranges]
-  global_shape = tuple([ssimplify(r.src[0]) for r in ranges if r.arg[0:-1] in global_dims])
-  local_shape = tuple([ssimplify(r.src[0]) for r in ranges if r.arg[0:-1] in local_dims])
+  global_shape = tuple([ssimplify(r.src[0]) for r in ranges if r.arg%1000 in global_dims])
+  local_shape = tuple([ssimplify(r.src[0]) for r in ranges if r.arg%1000 in local_dims])
 
   # get the idxs
-  ki: KernelInfo = s.arg
   if ki.dont_use_locals:
     assert not local_dims, "can't use locals if there's no local dims"
     idxs = get_grouped_dims("idx", global_shape, ctx.global_max, reverse=True)
@@ -83,187 +78,12 @@ def add_gpudims(ctx:Renderer, s:UOp):
   for r in s_topo:
     if r.op is not Ops.RANGE: continue
     try:
-      ii = (global_dims+local_dims).index(r.arg[0:-1])
-      if r.arg[1] == AxisType.REDUCE: continue
+      ii = (global_dims+local_dims).index(r.arg%1000)
+      if r.arg < 2000 and ki.axis_types[r.arg%1000] == AxisType.GROUP_REDUCE: continue
       subs[r] = idxs[ii]
     except ValueError: continue
   return s.substitute(subs)
 
-def fix_reduce_unroll(x:UOp):
-  reduce_range, reduce_expand = partition(x.src[1:], lambda y: y.op is Ops.RANGE)
-  if len(reduce_expand) == 0: return None
-  reduce_expand = [x for x in reduce_expand if x.op is not Ops.CONST]
-  assert all(x.op is Ops.UNROLL for x in reduce_expand), f"not all UNROLLS in {reduce_expand}"
-  ret = x.src[0]
-  if len(contract_axis:=flatten(x.arg for x in reduce_expand)):
-    ret = UOp(Ops.CONTRACT, x.dtype.vec(prod(x[1] for x in contract_axis)), (ret,), tuple(contract_axis), tag=1)
-  # REDUCE supports both "horizontal" reduction and range reduction. the horizontal elements are taken in the nearest group
-  return x.replace(src=(ret,)+tuple(reduce_range))
-
-def fix_store_unroll(x:UOp):
-  store_expand, store_range = partition(x.src[2:], lambda y: y.op is Ops.UNROLL)
-  if len(store_expand) == 0: return None
-  return UOp(Ops.CONTRACT, dtypes.void, (x.replace(src=x.src[:2]+tuple(store_range)),), tuple(flatten(x.arg for x in store_expand)), tag=1)
-
-def fix_group_for_reduce(x:UOp):
-  reduce_gfr, reduce_r = partition(x.src[1:], lambda u: u.op is Ops.RANGE and u.arg[1] == AxisType.GROUP_REDUCE)
-  if len(reduce_gfr) == 0: return None
-
-  # NOTE: if there's other locals here, we need them in the buffer too
-  upstream_locals = [u for u in x.toposort() if u.op is Ops.RANGE and u.arg[1] == AxisType.LOCAL]
-
-  # do only the non grouped reduces early
-  ret = x.replace(src=(x.src[0],)+tuple(reduce_r))
-  reduce_loop = [x.replace(arg=(x.arg[0]+100, AxisType.REDUCE)) for x in reduce_gfr]
-  buf = ret.bufferize(*upstream_locals, *reduce_gfr, arg=AddrSpace.LOCAL).index(*upstream_locals, *reduce_loop)
-
-  # gate with an if on the store + do the final reduce
-  buf = UOp(Ops.IF, dtype=buf.dtype, src=(functools.reduce(operator.and_, [x.eq(0) for x in reduce_gfr]), buf))
-  return buf.reduce(*reduce_loop, arg=x.arg)
-
-pm_group_for_reduce = PatternMatcher([
-  # fix group for reduce
-  (UPat(Ops.REDUCE, name="x"), fix_group_for_reduce),
-])
-
 pm_add_gpudims = PatternMatcher([
-  # add gpudims must be last
   (UPat(Ops.SINK, name="s"), add_gpudims),
-  # rewrite UPCAST/UNROLL range to something to be expanded
-  (UPat(Ops.RANGE, name="r"),
-   lambda r: UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(s:=r.vmax+1), tuple(range(s))),), ((r.arg[0],s),)) \
-    if r.arg[1] in {AxisType.UNROLL, AxisType.UPCAST} else None),
-  # fix REDUCEs with UNROLLs
-  (UPat(Ops.REDUCE, name="x"), fix_reduce_unroll),
-  (UPat(Ops.STORE, name="x"), fix_store_unroll),
 ])
-
-def apply_tensor_cores(ctx:tuple[dict, Renderer], in0:UOp, in1:UOp, r_range:UOp, reduceop:UOp):
-  if not USE_TC: return None
-  # tensor cores have three ranges. X, Y, and REDUCE
-  in0_ranges = sorted([u for u in in0.ranges if u not in in1.ranges], key=lambda x: x.arg[0])
-  in1_ranges = sorted([u for u in in1.ranges if u not in in0.ranges], key=lambda x: x.arg[0])
-  #print(len(in0_ranges), len(in1_ranges))
-  if not len(in0_ranges) or not len(in1_ranges): return None
-  in0_range, in1_range = in0_ranges[0], in1_ranges[0]
-  if DEBUG >= 2: print('TC', in0_range.arg, in1_range.arg, r_range.arg)
-
-  # confirm the dtype and size is good
-  tc_opts: list[TensorCore] = []
-  for tc in ctx[1].tensor_cores:
-    if reduceop.dtype == tc.dtype_out and in0.dtype == tc.dtype_in and in1.dtype == tc.dtype_in:
-      if all(i <= j for i,j in zip(tc.dims, [in0_range.vmax+1, in1_range.vmax+1, r_range.vmax+1])):
-        tc_opts.append(tc)
-  if len(tc_opts) == 0: return None
-  tc = tc_opts[0]
-
-  # create the new ranges as speced by the tensor core
-  old_range = [in0_range, in1_range, r_range]
-  new_range = [r.replace(src=(r.src[0]//tc.dims[i],)) for i,r in enumerate(old_range)]
-  new_reduce_range = new_range[2]
-  tc_range = 9050 #+ r_range.arg[0]*100
-  red_ranges = []
-
-  ne: list[UOp] = []
-  for o in tc.opts:
-    lrange = UOp.range(dtypes.int, 2, tc_range, AxisType.UPCAST if o[0] == "u" else AxisType.LOCAL)
-    ne.append(lrange)
-    tc_range += 1
-    new_range[1-int(o[1])] = (2 * new_range[1-int(o[1])]) + lrange
-  for _, amt in tc.get_reduce_axes():
-    lrange = UOp.range(dtypes.int, amt, tc_range, AxisType.UNROLL)
-    ne.append(lrange)
-    red_ranges.append(lrange)
-    tc_range += 1
-    new_range[2] = (amt * new_range[2]) + lrange
-  tne = [x.replace(tag=1) for x in ne]
-
-  # replace ranges in other parts of the graph
-  for x,y in zip(old_range, new_range): ctx[0][x] = y
-
-  # apply the swizzled ranges to the srcs
-  srcs = [s.substitute(dict(zip(old_range, new_range))).substitute(dict(zip(ne, tne))) for s in (in0, in1)]
-  srcs = [x.substitute(dict(zip(tne, [ne[i] for i in p]))) for x,p in zip(srcs, tc.permutes_for_shape_str(tc.base_shape_str()))]
-
-  ned = dict(zip(tc.base_shape_str(), ne))
-  tc_reduce_axes = tuple([ned[f"r{i}"].arg[0] for i in range(len(tc.get_reduce_axes()))])
-  base_upcast_axes = tuple([(ned[s].arg[0], 2) for s in tc.base_upcast_axes()])
-  tc_upcast_axes = tuple([base_upcast_axes[:int(math.log2(tc.elements_per_thread[i]))] for i in range(3)])
-
-  # construct the op
-  # TODO: remove tc_upcast_axes from the arg
-  wmma_arg = (str(tc), tc.dims, tc.dtype_in, tc.dtype_out, ctx[1].device, tc.threads, tc_upcast_axes, tc_reduce_axes)
-  wmma = UOp(Ops.WMMA, dtype=tc.dtype_out.vec(tc.elements_per_thread[2]), src=(
-    UOp(Ops.CONTRACT, dtype=srcs[0].dtype.vec(tc.elements_per_thread[0]), src=(srcs[0],), arg=tc_upcast_axes[0]),
-    UOp(Ops.CONTRACT, dtype=srcs[1].dtype.vec(tc.elements_per_thread[1]), src=(srcs[1],), arg=tc_upcast_axes[1]),
-    UOp.const(tc.dtype_out.vec(tc.elements_per_thread[2]), 0.0)), arg=wmma_arg)
-  tc_uop = UOp(Ops.UNROLL, tc.dtype_out, (wmma,), arg=tc_upcast_axes[2])
-  ret = tc_uop.reduce(new_reduce_range, arg=Ops.ADD)
-  # confirm the UNROLLs aren't actually used, these need to be broadcast MUL
-  assert all(u not in red_ranges for u in ret.toposort()), "UNROLLs in TC"
-  return ret
-
-from tinygrad.codegen.opt.postrange import pm_flatten_range
-
-pm_tensor_cores = PatternMatcher([
-  ((UPat.var("in0")*UPat.var("in1")).reduce(UPat(Ops.RANGE, name="r_range"), name="reduceop", arg=Ops.ADD), apply_tensor_cores),
-
-  # replace range
-  #(UPat(Ops.RANGE, name="r"), lambda ctx,r: ctx[0].get(r, None)),
-  (UPat(Ops.SINK, name="s"), lambda ctx,s: graph_rewrite(s.substitute(ctx[0]), pm_flatten_range, name="flatten")),
-])
-
-def fix_bufferize(x:UOp):
-  if x.arg != AddrSpace.LOCAL: return None
-  locals_left = [r for r in x.ranges if r.arg[1] == AxisType.LOCAL]
-  if not len(locals_left): return None
-  acc = x.size
-  st = []
-  for l in locals_left:
-    st.append(l*acc)
-    acc *= l.vmax+1
-  return x.replace(src=(x.src[0],) + tuple(locals_left[::-1]) + x.src[1:]).index(sum(st))
-
-
-pm_double_index = PatternMatcher([
-  (UPat(Ops.INDEX, src=(UPat(Ops.INDEX, src=(UPat.var("b"), UPat.var("x"))), UPat.var("y"))), lambda b,x,y: b.index(x+y)),
-])
-
-pm_fix_locals = pm_double_index+PatternMatcher([
-  (UPat(Ops.BUFFERIZE, name="x"), fix_bufferize),
-])
-
-B = 8
-
-def loop_store(x:UOp):
-  r_maybe = [r for r in x.src[0].ranges if r.arg[0] == 2 and r.tag is None]
-  ur_maybe = [r for r in x.ranges if r.arg[0] < 0]
-  #print("store", len(r_maybe), len(ur_maybe))
-  if not len(r_maybe) or not len(ur_maybe): return None
-
-  r = r_maybe[0]
-  ur = ur_maybe[0]
-  rr = r.replace(src=(r.src[0]//B,), tag=1)
-  return x.substitute({r:rr*B+ur})
-
-def loop_bufferize(x:UOp):
-  if x.arg != AddrSpace.LOCAL: return None
-  r_maybe = [r for r in x.ranges if r.arg[0] == 2 and r.tag is None]
-  ur_maybe = [r for r in x.ranges if r.arg[0] < 0]
-
-  if len(ur_maybe):
-    ur = ur_maybe[0]
-  else:
-    if len(r_maybe) == 0: return None
-    r = r_maybe[0]
-    ur = UOp.range(dtypes.int, B, -2010)
-    rr = r.replace(src=(r.src[0]//B,), tag=1)
-    x = x.substitute({r:rr*B + ur})
-
-  ur1 = UOp.range(dtypes.int, B, ur.arg[0]+1)
-  return x.replace(src=(x.src[0],ur)+x.src[1:]).index(x.size*ur1)
-
-pm_bufferize_loop = pm_double_index+PatternMatcher([
-  (UPat(Ops.BUFFERIZE, name="x"), loop_bufferize),
-  (UPat(Ops.STORE, name="x"), loop_store),
-])

tinygrad/codegen/linearize.py

@@ -3,7 +3,7 @@ import heapq
 from collections import defaultdict
 from dataclasses import dataclass, replace
 from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat, GroupOp
-from tinygrad.helpers import dedup, all_same, flatten, BLOCK_REORDER
+from tinygrad.helpers import dedup, all_same, flatten, getenv
 
 # NOTE: any toposort should be valid here, unlike last time this isn't required, it's just for speed
 def block_reorder(lst:list[UOp]) -> list[UOp]:
@@ -150,7 +150,7 @@ def make_block_bottom_up(ctx:BlockContext, x:UOp):
     srcs.append(add_blockends(base_block, new_ctx, current_ctx))
 
   lst = lst[::-1]
-  if BLOCK_REORDER: lst = block_reorder(lst)
+  if getenv("BLOCK_REORDER", 1): lst = block_reorder(lst)
   bb = BasicBlock(tuple(lst), ctx=current_ctx, cnt=child_count, child_ctx=child_ctx)
   return UOp(Ops.BLOCK, src=tuple(srcs), arg=bb)
 

tinygrad/codegen/lowerer.py

@@ -1,7 +1,10 @@
 # the job of the lowerer is to do indexing
+import functools, operator
+from typing import cast
 from dataclasses import dataclass
-from tinygrad.dtype import dtypes
-from tinygrad.uop.ops import KernelInfo, UOp, Ops, PatternMatcher, UPat, sint_to_uop, AxisType, graph_rewrite, resolve
+from tinygrad.dtype import dtypes, AddrSpace, PtrDType
+from tinygrad.uop.ops import KernelInfo, UOp, Ops, PatternMatcher, UPat, sint_to_uop, AxisType, graph_rewrite
+from tinygrad.helpers import prod, partition, flatten
 
 # ***** indexing *****
 
@@ -12,12 +15,20 @@ class IndexContext:
   start: int = 0
 
 def shape_to_idx(s, axis_types, start=0):
-  return [UOp.range(dtypes.int, sint_to_uop(s), start+i, at) for i, (s, at) in enumerate(zip(s, axis_types))]
+  # indexes
+  idxs = []
+  for i, (s, at) in enumerate(zip(s, axis_types)):
+    if at in (AxisType.UPCAST, AxisType.UNROLL):
+      assert isinstance(s, int), "needs to be int to upcast/unroll"
+      idxs.append(UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(s), tuple(range(s))),), ((i,s),), tag=1))
+    else:
+      # all others are RANGES
+      idxs.append(UOp(Ops.RANGE, dtypes.int, (sint_to_uop(s),), start+i))
+  return idxs
 
 def get_index(ast:UOp) -> IndexContext:
   axis_types = ast.arg.axis_types if isinstance(ast.arg, KernelInfo) else ()
-  if len(ast.full_shape) != len(axis_types):
-    axis_types = tuple([AxisType.REDUCE if resolve(s != fs) else AxisType.LOOP for s,fs in zip(ast.shape, ast.full_shape)])
+  if len(ast.full_shape) != len(axis_types): axis_types = (AxisType.LOOP,)*len(ast.full_shape)
   return IndexContext(axis_types, [], 0)
 
 # ***** lowering (given index) *****
@@ -31,8 +42,16 @@ def lower_reduce_axis(ctx: IndexContext, x: UOp):
   new_idxs = shape_to_idx(x.src[0].shape, ctx.axis_types, ctx.start)
   full_new_idx = list(ctx.idxs)
   for a in x.axis_arg: full_new_idx[a] = new_idxs[a]
+
   ret = subblock(ctx, full_new_idx, x.src[0])
-  return UOp(Ops.REDUCE, x.dtype, (ret,)+tuple([full_new_idx[i] for i in x.axis_arg]), x.arg[0])
+
+  # NOTE: always using ridxs is fine here
+  reduce_range, reduce_expand = partition([full_new_idx[i] for i in x.axis_arg], lambda y: y.op is Ops.RANGE)
+  assert all(x.op is Ops.UNROLL for x in reduce_expand), f"not all UNROLLS in {reduce_expand} for {x.axis_arg}"
+  if len(contract_axis:=flatten(x.arg for x in reduce_expand)):
+    ret = UOp(Ops.CONTRACT, x.dtype.vec(prod(x[1] for x in contract_axis)), (ret,), tuple(contract_axis), tag=1)
+  # REDUCE supports both "horizontal" reduction and range reduction. the horizontal elements are taken in the nearest group
+  return UOp(Ops.REDUCE, x.dtype, (ret,)+tuple(reduce_range), x.arg[0])
 
 def lower_store(ctx: IndexContext, x: UOp, buf: UOp):
   # TODO: reenable after REDUCE_AXIS is fixed
@@ -48,7 +67,15 @@ def lower_store(ctx: IndexContext, x: UOp, buf: UOp):
 
   stored = subblock(ctx, real_new_idxs, x.src[1])
   used_ranges = [x for x in used_idxs if x.op is Ops.RANGE]
-  return buf.index(idx, valid).store(stored, *used_ranges)
+  ret = buf.index(idx, valid).store(stored, *used_ranges)
+
+  # insert BARRIER if we are ending a LOCAL, IF if we are ending a GROUP_REDUCE
+  if cast(PtrDType, buf.dtype).addrspace == AddrSpace.LOCAL and \
+      any(ctx.axis_types[x.arg%1000] in {AxisType.GROUP_REDUCE, AxisType.LOCAL} for x in used_ranges):
+    ret = ret.barrier()
+    range_gates = [x.eq(0) for x in used_ranges if ctx.axis_types[x.arg%1000] == AxisType.GROUP_REDUCE]
+    if len(range_gates): ret = UOp(Ops.IF, src=(functools.reduce(operator.and_, range_gates), ret))
+  return ret
 
 def fixup_wmma(ctx:IndexContext, x:UOp):
   if x.tag is not None: return None
@@ -59,8 +86,8 @@ def fixup_wmma(ctx:IndexContext, x:UOp):
   srcs = subblock(ctx, full_new_idx, UOp.sink(*x.src)).src
 
   # NOTE: this assumes these are expanded. which now shouldn't change anything
-  new_x_arg_m2 = tuple([tuple([(full_new_idx[a].arg[0], sz) for a,sz in v]) for v in x.arg[-2]])
-  new_x_arg_m1 = tuple([full_new_idx[a].arg[0] for a in x.arg[-1]])
+  new_x_arg_m2 = tuple([tuple([(full_new_idx[a].arg[0][0], sz) for a,sz in v]) for v in x.arg[-2]])
+  new_x_arg_m1 = tuple([full_new_idx[a].arg[0][0] for a in x.arg[-1]])
   return x.replace(src=srcs, arg=x.arg[:-2]+(new_x_arg_m2, new_x_arg_m1), tag=1)
 
 pm_lowerer = PatternMatcher([
@@ -83,5 +110,5 @@ pm_lowerer = PatternMatcher([
 
   # axis fixups for WMMA
   (UPat((Ops.CONTRACT, Ops.UNROLL), name="x"),
-   lambda ctx,x: x.replace(tag=1, arg=tuple([(ctx.idxs[a].arg[0], sz) for a,sz in x.arg])) if x.tag is None else None),
+   lambda ctx,x: x.replace(tag=1, arg=tuple([(ctx.idxs[a].arg[0][0], sz) for a,sz in x.arg])) if x.tag is None else None),
 ])

tinygrad/codegen/opt/__init__.py

@@ -1,9 +1,8 @@
 # opt opinionatedly transforms an ast into an optimized ast using either heuristics or beam search
 
 from tinygrad.codegen.opt.kernel import Kernel
-from tinygrad.codegen.opt.postrange import RKernel, pm_flatten_range
 from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
-from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, KernelInfo
+from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops
 from tinygrad.helpers import NOOPT, BEAM, USE_TC, getenv
 from tinygrad.renderer import Renderer
 from tinygrad.uop.spec import type_verify
@@ -20,40 +19,20 @@ def get_optimized_ast(ast:UOp, renderer:Renderer) -> UOp:
     The Ops.SINK rooted AST transformed to apply the opts and with a KernelInfo in the arg.
   """
 
-  assert ast.arg is None, "no opt if there's an arg"
   k = Kernel(ast, opts=renderer)
-  if not NOOPT:
+  if ast.arg is not None and ast.arg.opts_to_apply is not None: k.apply_opts(ast.arg.opts_to_apply)
+  elif not NOOPT:
     if not k.apply_tensor_cores(USE_TC.value): k.apply_opts(hand_coded_optimizations(k))
     if BEAM >= 1:
       from tinygrad.codegen.opt.search import beam_search, bufs_from_lin
       kb = Kernel(ast, opts=renderer)
       rawbufs = bufs_from_lin(kb, allocate=False)
       k = beam_search(kb, rawbufs, BEAM.value, bool(getenv("BEAM_ESTIMATE", 1)))
-  return ast.replace(arg=KernelInfo(opts_to_apply=tuple(k.applied_opts)))
-
-pm_get_optimization = PatternMatcher([
-  (UPat(Ops.SINK, name="ast"), lambda ctx,ast: get_optimized_ast(ast, ctx) if ast.arg is None and ast.src[0].st is not None else None),
-])
-
-def apply_opt(ast:UOp, renderer:Renderer):
-  k = Kernel(ast, opts=renderer)
-  k.apply_opts(ast.arg.opts_to_apply)
   ret = k.get_optimized_ast()
   if __debug__: type_verify(list(ret.toposort()))
   return ret
 
-pm_do_optimize = PatternMatcher([
-  (UPat(Ops.SINK, name="ast"), lambda ctx,ast: apply_opt(ast, ctx) if ast.arg is not None and ast.arg.opts_to_apply is not None else None),
-])
-
-# ** postrange **
-
-def apply_ropt(ast:UOp, renderer:Renderer):
-  k = RKernel(ast, opts=renderer)
-  if ast.arg is not None: k.apply_opts(ast.arg.opts_to_apply)
-  return k.get_optimized_ast()
-
-pm_postrange_opt = pm_flatten_range+PatternMatcher([
-  (UPat(Ops.SINK, name="ast"), lambda ctx,ast: apply_ropt(ast, ctx) if ast.arg is None or \
-   (ast.arg is not None and ast.arg.opts_to_apply is not None) else None),
+pm_optimize = PatternMatcher([
+  (UPat(Ops.SINK, name="ast"), lambda ctx,ast:
+   get_optimized_ast(ast, ctx) if (ast.arg is None or ast.arg.opts_to_apply is not None) and ast.src[0].st is not None else None),
 ])

tinygrad/codegen/opt/heuristic.py

@@ -28,7 +28,7 @@ def hand_coded_optimizations(k:Kernel) -> list[Opt]:
           return k.applied_opts
 
   # are we grouping? (requires local shape support)
-  if resolve(prod(k.output_shape[i] for i in k.upcastable_dims) <= 2048, False):
+  if resolve(prod(k.sts[0].shape[i] for i in k.upcastable_dims) <= 2048, False):
     for sz in [16]:
       try:
         k.apply_opt(Opt(OptOps.GROUPTOP, 0, sz))
@@ -62,7 +62,7 @@ def hand_coded_optimizations(k:Kernel) -> list[Opt]:
   # potentially do more upcasts of non reduce axes based on a heuristic
   is_dsp = k.opts is not None and k.opts.device == "DSP"
   upcasted_axis: set[int] = set()
-  while resolve(prod(k.output_shape[i] for i in k.upcastable_dims) >= 1024):
+  while resolve(prod(k.sts[0].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]):

tinygrad/codegen/opt/kernel.py

@@ -10,7 +10,7 @@ from tinygrad.uop.spec import type_verify, ast_spec
 from tinygrad.device import Device
 from tinygrad.codegen.opt.tc import TensorCore
 from tinygrad.renderer import Renderer
-from tinygrad.dtype import ImageDType
+from tinygrad.dtype import ImageDType, AddrSpace
 from tinygrad.helpers import all_same, colored, ansilen, dedup, prod, round_up, to_function_name, unwrap, argfix, DEBUG, TC_SELECT, TC_OPT, AMX
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import strides_for_shape, get_contraction
@@ -60,7 +60,7 @@ class Kernel:
 
     self.vars: list[Variable] = self.ast.variables()
     # NOTE: this requires a specific order with the [::-1], this is likely a bug
-    self.bufs: list[UOp] = [x for x in self.ast.toposort() if x.op in GroupOp.Buffer and x.st is not None][::-1]
+    self.bufs: list[UOp] = [x for x in self.ast.toposort() if x.op in GroupOp.Buffer][::-1]
 
     # create new shapetrackers inside this kernel, we will permute them
     self.sts: list[ShapeTracker] = [x.st_arg for x in self.bufs]
@@ -92,6 +92,10 @@ class Kernel:
     global_loops = AxisType.GLOBAL if self.opts.has_local else AxisType.LOOP
     self.axis_types: list[AxisType] = [AxisType.REDUCE if resolve(x!=y) else global_loops for x,y in zip(self.output_shape, self.full_shape)]
 
+    # confirm all reduce axes are at the end
+    if (final_reduces := [x for x in self.axis_types if x == AxisType.REDUCE]) and final_reduces != self.axis_types[-len(final_reduces):]:
+      raise RuntimeError(f"reduces are not at the end of the shape {self.full_shape} -> {self.output_shape}")
+
   def copy(self):
     ret = type(self).__new__(type(self))
 
@@ -118,7 +122,7 @@ class Kernel:
   @property
   def output_shape(self) -> tuple[sint, ...]: return self.sts[0].shape
   @property
-  def shape_len(self) -> int: return len(self.full_shape)
+  def shape_len(self) -> int: return len(self.sts[0].shape)
 
   def axes_of(self, *axis_type:AxisType) -> list[int]: return [i for i,t in enumerate(self.axis_types) if t in argfix(axis_type)]
   @property
@@ -170,7 +174,7 @@ class Kernel:
   # amount : the amount to take
   # top : if you want to pull that amount from the top
   # insert_at : place to insert the new stuff
-  def shift_to(self, axis:int, amount:int, new_type:AxisType, top:bool=False, insert_at:int|None=None) -> int:
+  def shift_to(self, axis:int, amount:int, new_type:AxisType, top:bool=False, insert_at:int|None=None):
     if insert_at is None: insert_at = self.shape_len
     self.axis_types.insert(insert_at, new_type)
     move_axis = axis if top else axis+1
@@ -179,7 +183,6 @@ class Kernel:
     new_axes = [i for i in range(insert_at) if i != move_axis]+[move_axis]+[i for i in range(insert_at, self.shape_len+1) if i != move_axis]
     self.reshape(new_shape_fxn)
     self.permute(new_axes)
-    return insert_at
 
   # ******************** complex simplifiers ********************
 
@@ -241,11 +244,11 @@ class Kernel:
       if axis is None: return -1
       if op is OptOps.UNROLL: return self.unrollable_dims[axis]
       if op in {OptOps.GROUP, OptOps.GROUPTOP}: return self.axes_of(AxisType.REDUCE)[axis]
-      check(axis < self.shape_len, f"invalid axis on {axis=} {op=} {self.shape_len=}")
+      check(axis < self.shape_len, "invalid axis")
       return axis
     except IndexError as e: raise KernelOptError from e
 
-  def apply_opt(self, opt:Opt, append_opt:bool=True) -> int|None:
+  def apply_opt(self, opt:Opt, append_opt:bool=True):
     if self.finalized: raise RuntimeError("can't optimize Kernel after it's finalized")
     if self.dont_use_locals: check(opt.op not in {OptOps.LOCAL, OptOps.GROUP, OptOps.GROUPTOP}, "not using locals")
 
@@ -259,7 +262,7 @@ class Kernel:
       check(0 < (use_tensor_cores:=cast(tuple, opt.arg)[2]) <= 2, "use_tensor_cores value is not valid")
       check(self._apply_tc_opt(use_tensor_cores, cast(int, opt.axis), tc_select, tc_opt), "no tensor core available")
       self.applied_opts.append(opt)
-      return None
+      return
 
     axis = self.real_axis(opt.op, opt.axis)
 
@@ -282,30 +285,28 @@ class Kernel:
       smem_sz = amt*acc_sz*upcast_sz*local_sz
       check(smem_sz <= self.opts.shared_max, f"exceeds maximum shared memory size: needs {smem_sz}, max {self.opts.shared_max}")
 
-    new_axis = None
     if opt.op is OptOps.LOCAL:    # cyan
       # NOTE: LLVM/CPU can use locals too, but they are treated the same as globals (still helpful for L1 cache)
       # it's disabled for now since it makes BEAM slow for little gain
       check(self.opts.has_local, "target does not support local")
       check(self.axis_types[axis] is AxisType.GLOBAL, "local is for globals")
-      new_axis = self.shift_to(axis, amt, AxisType.LOCAL, insert_at=max(self.axes_of(AxisType.GLOBAL, AxisType.LOCAL))+1)
+      self.shift_to(axis, amt, AxisType.LOCAL, insert_at=max(self.axes_of(AxisType.GLOBAL, AxisType.LOCAL))+1)
     elif opt.op in {OptOps.GROUP, OptOps.GROUPTOP}:   # green
       check(self.opts.has_local and self.opts.has_shared, "target does not support local or shared mem")
       check(self.axis_types[axis] is AxisType.REDUCE, "must be reduce axis to group")
       check(not self.tensor_core, "can't group with tensor cores")
       check(len(reduce_axes:=[i for r in self.reduceops for i in r.axis_arg]) == len(set(reduce_axes)), "can't group with parallel reduces")
-      new_axis = self.shift_to(axis, amt, AxisType.GROUP_REDUCE, top=(opt.op is OptOps.GROUPTOP), insert_at=min(self.axes_of(AxisType.REDUCE)))
+      self.shift_to(axis, amt, AxisType.GROUP_REDUCE, top=(opt.op is OptOps.GROUPTOP), insert_at=min(self.axes_of(AxisType.REDUCE)))
     elif opt.op is OptOps.UNROLL:                     # purple
       check(self.axis_types[axis] not in (AxisType.UPCAST, AxisType.UNROLL), "can't upcasted already upcasted")
       check(amt <= 32, "don't unroll more than 32")
-      new_axis = self.shift_to(axis, amt, AxisType.UNROLL, insert_at=None)
+      self.shift_to(axis, amt, AxisType.UNROLL, insert_at=None)
     elif opt.op is OptOps.UPCAST:                     # yellow
       check(axis in self.upcastable_dims, f"{axis=} not in {self.upcastable_dims=}")
       # NOTE: assume the first get_local_axes() LOCAL are for TC
       check(not (self.tensor_core and axis in self.axes_of(AxisType.LOCAL)[:len(self.tensor_core.get_local_axes())]), "can't upcast TC locals")
       check((self.opts is not None and self.opts.device == "DSP") or amt <= 16, "don't upcast more than 16")
-      new_axis = self.shift_to(axis, amt, AxisType.UPCAST,
-                               insert_at=max(self.axes_of(AxisType.GLOBAL, AxisType.LOCAL, AxisType.LOOP, AxisType.UPCAST))+1)
+      self.shift_to(axis, amt, AxisType.UPCAST, insert_at=max(self.axes_of(AxisType.GLOBAL, AxisType.LOCAL, AxisType.LOOP, AxisType.UPCAST))+1)
     elif opt.op is OptOps.NOLOCALS:
       check(self.opts.has_local and not self.dont_use_locals, "NOLOCALS is meaningless if target does not support local or already not using locals")
       check(AxisType.LOCAL not in self.axis_types and self.group_for_reduces == 0, "can't have no locals with locals")
@@ -335,7 +336,6 @@ class Kernel:
     if append_opt: self.applied_opts.append(opt)
     if self.simplify_ones() and self.tensor_core_opts:
       self.tensor_core_opts.fix_axes(axis) # fix up axes in TC opts if required after simplify_ones()
-    return new_axis
 
   def apply_opts(self, opts:Sequence[Opt]) -> Kernel:
     for opt in opts: self.apply_opt(opt)
@@ -460,7 +460,8 @@ class Kernel:
       if op.op is Ops.REDUCE_AXIS:
         reduce_idx = len(self.bufs) + self.reduceops.index(op) * 2
         changed = tuple(i for i in range(self.shape_len) if resolve(self.sts[reduce_idx].shape[i] != self.sts[reduce_idx + 1].shape[i]))
-        axes = tuple(i for i in self.axes_of(AxisType.REDUCE, AxisType.GROUP_REDUCE, AxisType.UNROLL) if i in changed)
+        axes = tuple(i for i in self.axes_of(AxisType.REDUCE, AxisType.UNROLL) if i in changed)
+        grouped_axes = tuple(i for i in self.axes_of(AxisType.GROUP_REDUCE) if i in changed)
         if (tc := self.tensor_core) and self.use_tensor_cores == 1:
           # get reduce/upcast axes for the tensor cores
           tc_reduce_axes = self.shape_str_to_axis([f"r{i}" for i in range(len(tc.get_reduce_axes()))])
@@ -485,6 +486,23 @@ class Kernel:
           return ret.replace(src=(tc_uop,), arg=(Ops.ADD, new_axes)) if (new_axes := tuple(i for i in axes if i not in tc_reduce_axes)) else tc_uop
 
         ret = ret.replace(arg = (op.arg[0], axes))
+        if self.group_for_reduces and grouped_axes:
+          local_axes = tuple([i for i,t in enumerate(self.axis_types) if t in (AxisType.LOCAL, AxisType.UPCAST) or i in grouped_axes])
+          slocal, supcast, sgroup = sorted(self.axes_of(AxisType.LOCAL)), sorted(self.axes_of(AxisType.UPCAST)), sorted(grouped_axes)
+          # NOTE: start with UPCAST at the end so it has stride 1 and can merge
+          base_shape = tuple([self.full_shape[i] for i in slocal] + [self.full_shape[i] for i in sgroup] + [self.full_shape[i] for i in supcast])
+          permute_axes = tuple([local_axes.index(i) for i in slocal+sgroup+supcast])
+          local_shape = tuple([s if i in local_axes else 1 for i,s in enumerate(self.full_shape)])
+          local_src_shape = tuple([self.full_shape[i] if i in self.axes_of(AxisType.GLOBAL) else s for i,s in enumerate(local_shape)])
+          st = ShapeTracker.from_shape(base_shape).permute(permute_axes).reshape(local_shape).expand(local_src_shape)
+          local_size = st.real_size()
+          local_buffer = UOp(Ops.DEFINE_LOCAL, op.dtype.ptr(local_size, addrspace=AddrSpace.LOCAL), (), f"temp{self.reduceops.index(op)}")
+          local_load = local_buffer.view(st).load(local_buffer.view(st).store(ret))
+          grouped_reduce = UOp(Ops.REDUCE_AXIS, op.dtype, (local_load,), arg=(op.arg[0], grouped_axes))
+          if op is self.reduceops[-1]: return grouped_reduce
+          st = ShapeTracker.from_shape(tuple([1 if i in grouped_axes else s for i,s in enumerate(local_shape)]))
+          return local_buffer.view(st).load(local_buffer.view(st).store(grouped_reduce))
+
       return ret
     self.finalized = True
     fixed_ast = fixup_ast(self.ast)

tinygrad/codegen/opt/postrange.py

@@ -1,183 +0,0 @@
-import math
-from tinygrad.dtype import AddrSpace
-from tinygrad.uop.ops import UOp, Ops, sint, ssimplify, AxisType, KernelInfo, PatternMatcher, UPat, graph_rewrite, _substitute
-from tinygrad.uop.symbolic import symbolic
-from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps
-from tinygrad.renderer import Renderer
-from tinygrad.dtype import dtypes
-
-def flatten_range(r:UOp):
-  off = 2 if r.op is Ops.STORE else 1
-  rngs = r.src[off:]
-  if not len(rngs): return None
-  new_rngs = [x for x in UOp.sink(*rngs).toposort() if x.op is Ops.RANGE]
-  return r.replace(src=r.src[:off]+tuple(new_rngs))
-
-pm_flatten_range = PatternMatcher([
-  # real ranges only
-  (UPat((Ops.REDUCE, Ops.STORE), name="r"), flatten_range),
-])
-
-def count_divmod(x:UOp): return len([u for u in x.toposort() if u.op in {Ops.IDIV, Ops.MOD}])
-
-class RKernel(Kernel):
-  def __init__(self, ast:UOp, opts:Renderer|None=None):
-    self.rng = sorted([u for u in ast.toposort() if u.op is Ops.RANGE and u.vmax > 0], key=lambda x: x.arg)
-    super().__init__(ast, opts)
-    self.sts.clear()
-
-    # convert LOOP to GLOBAL
-    self.replaces = {}
-    if self.opts.has_local:
-      store_rngs = self.ast.src[0].src[2:]
-
-      # filter any not in local stores
-      local_store_rngs = [x.ranges for x in self.ast.toposort() if (x.op is Ops.STORE and x.src[0].dtype.addrspace == AddrSpace.LOCAL) \
-                          or (x.op is Ops.BUFFERIZE and x.arg == AddrSpace.LOCAL)]
-      for ls in local_store_rngs: store_rngs = [x for x in store_rngs if x in ls]
-
-      store_rng = [x for x in UOp.sink(*store_rngs).toposort() if x.op is Ops.RANGE] if store_rngs else []
-      rng = [x.replace(arg=(x.arg[0], AxisType.GLOBAL)) if x.arg[1] == AxisType.LOOP and x in store_rng else x for x in self.rng]
-      self.replaces.update(dict(zip(self.rng, rng)))
-      self.rng = rng
-
-  def simplify_merge_adjacent(self):
-    return
-    # NOTE: this one is better than the one in kernel.py, which is kind of a problem
-    terminators = [u for u in self.ast.toposort() if u.op in {Ops.REDUCE, Ops.STORE}]
-    termination = {}
-    for t in terminators:
-      for u in t.src[1 if t.op is Ops.REDUCE else 2:]: termination[u] = t
-
-    replaces = {}
-    i = 0
-    while i < len(self.rng)-1:
-      r0, r1 = self.rng[i], self.rng[i+1]
-      # same axistype and same termination
-      if r0.arg[1] == r1.arg[1] and termination[r0] == termination[r1]:
-        s0, s1 = r0.src[0], r1.src[0]
-        new_range = r0.replace(src=(s0*s1,)).simplify()
-        # this checks the legality of a merge
-        oidx = self.ast.simplify()
-        nidx = graph_rewrite(oidx, _substitute+symbolic+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1}, name=f"check_merge_{i}_{i+1}")
-        # it simplifies
-        if count_divmod(nidx) <= count_divmod(oidx):
-          # it is correct
-          midx = graph_rewrite(nidx, _substitute+symbolic+pm_flatten_range, ctx={new_range:r0*s1+r1}, name=f"correct_merge_{i}_{i+1}")
-          if oidx is midx:
-            termination[new_range] = termination[r0]
-            replaces[r0] = new_range//s1
-            replaces[r1] = new_range%s1
-            self.rng[i] = new_range
-            del self.rng[i+1]
-            continue
-      i += 1
-    self.ast = self.ast.substitute(replaces, name="simplify_merge_adjacent")
-
-  def shift_to(self, axis:int, amount:int, new_type:AxisType, top:bool=False, insert_at:int|None=None):
-    old_sz = self.rng[axis].src[0].arg // amount
-    assert old_sz > 0, f"bad old_sz on {axis} {amount} {self.rng[axis]}"
-
-    maxarg = max([x.arg[0] for x in self.rng])
-    new_rng = UOp.range(dtypes.int, amount, maxarg+1, new_type)
-
-    if old_sz == 1:
-      self.replaces[self.rng[axis]] = new_rng
-      self.rng.insert(insert_at if insert_at is not None else len(self.rng), new_rng)
-      del self.rng[axis]
-    else:
-      replaced_rng = self.rng[axis].replace(src=(UOp.const(dtypes.int, old_sz),))
-      self.replaces[self.rng[axis]] = (new_rng * old_sz + replaced_rng) if top else (replaced_rng * amount + new_rng)
-      self.rng[axis] = replaced_rng
-      self.rng.insert(insert_at if insert_at is not None else len(self.rng), new_rng)
-    return new_rng
-
-  @property
-  def axis_types(self) -> list[AxisType]: return [x.arg[1] for x in self.rng]
-  @property
-  def shape_len(self): return len(self.rng)
-
-  @property
-  def full_shape(self) -> tuple[sint, ...]: return tuple([ssimplify(x.src[0]) for x in self.rng])
-  @property
-  def output_shape(self) -> tuple[sint, ...]: return tuple([ssimplify(x.src[0]) for x in self.ast.src[0].src[2:]])
-
-  def get_optimized_ast(self, name_override:str|None=None) -> UOp:
-    ret = self.ast
-    kernel_name = ret.arg.name if ret.arg is not None and ret.arg.name != "test" else self.name if name_override is None else name_override
-    rarg = KernelInfo(kernel_name, tuple(self.axis_types), self.dont_use_locals, tuple(self.applied_opts))
-    return ret.substitute(self.replaces).replace(arg=rarg)
-
-  # does nothing
-  @axis_types.setter
-  def axis_types(self, value): pass
-
-  def _apply_tc_opt(self, use_tensor_cores:int, axis:int, tc_select:int, opt_level:int) -> bool:
-    reduceop = [x for x in self.ast.toposort() if x.op is Ops.REDUCE][0]
-    if use_tensor_cores and reduceop is not None and reduceop.arg is Ops.ADD:
-      tensor_cores = self.opts.tensor_cores if tc_select == -1 else [self.opts.tensor_cores[tc_select]]
-      for tc in tensor_cores:
-        if tc.dtype_in == dtypes.float and tc.dtype_out == dtypes.float:
-          axes = [1,0]
-
-          # do optimizations and save the ranges
-          ne: list[UOp] = []
-          for opt in tc.opts:
-            ne.append(self.apply_opt(Opt({"u":OptOps.UPCAST, "l":OptOps.LOCAL}[opt[0]], axes[int(opt[1])], 2), append_opt=False))
-          for _, amt in tc.get_reduce_axes():
-            ne.append(self.apply_opt(Opt(OptOps.UNROLL, 0, amt), append_opt=False)) # TODO: this should be the reduce, not 0
-
-          # early realize for TC
-          self.ast = self.ast.substitute(self.replaces)
-          self.replaces = {}
-
-          # fix the srcs
-          reduceop = [x for x in self.ast.toposort() if x.op is Ops.REDUCE][0]
-          tne = [x.replace(tag=1) for x in ne]
-          ret = reduceop.substitute(dict(zip(ne, tne)))
-          srcs = list((ret.src[0] if ret.src[0].op is not Ops.CAST else ret.src[0].src[0]).src)
-          srcs = [x.substitute(dict(zip(tne, [ne[i] for i in p]))) for x,p in zip(srcs, tc.permutes_for_shape_str(tc.base_shape_str()))]
-
-          # get reduce/upcast axes for the tensor cores
-          tc_reduce_axes = self.shape_str_to_axis([f"r{i}" for i in range(len(tc.get_reduce_axes()))])
-          base_upcast_axes = tuple([(s,2) for s in self.shape_str_to_axis(tc.base_upcast_axes())])
-          tc_upcast_axes = tuple([base_upcast_axes[:int(math.log2(tc.elements_per_thread[i]))] for i in range(3)])
-
-          # axes to range number (was done in lowerer)
-          tc_upcast_axes = tuple([tuple([(self.rng[a].arg[0], sz) for a,sz in v]) for v in tc_upcast_axes])
-          tc_reduce_axes = tuple([self.rng[a].arg[0] for a in tc_reduce_axes])
-
-          # construct the op
-          # TODO: remove tc_upcast_axes from the arg
-          wmma_arg = (str(tc), tc.dims, tc.dtype_in, tc.dtype_out, self.opts.device, tc.threads, tc_upcast_axes, tc_reduce_axes)
-          wmma = UOp(Ops.WMMA, dtype=tc.dtype_out.vec(tc.elements_per_thread[2]), src=(
-            UOp(Ops.CONTRACT, dtype=srcs[0].dtype.vec(tc.elements_per_thread[0]), src=(srcs[0],), arg=tc_upcast_axes[0]),
-            UOp(Ops.CONTRACT, dtype=srcs[1].dtype.vec(tc.elements_per_thread[1]), src=(srcs[1],), arg=tc_upcast_axes[1]),
-            UOp.const(tc.dtype_out.vec(tc.elements_per_thread[2]), 0.0)), arg=wmma_arg)
-          tc_uop = UOp(Ops.UNROLL, tc.dtype_out, (wmma,), arg=tc_upcast_axes[2])
-
-          # preserve extra reduces
-          reduce_ranges = [x for x in UOp.sink(*reduceop.src[1:]).toposort() if x.op is Ops.RANGE and x.arg[0] not in tc_reduce_axes]
-          if len(reduce_ranges): tc_uop = UOp(Ops.REDUCE, tc_uop.dtype, (tc_uop,)+tuple(reduce_ranges), Ops.ADD)
-          self.ast = self.ast.substitute({reduceop: tc_uop})
-          return True
-    return False
-
-from dataclasses import replace
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo
-from tinygrad.helpers import colored
-from tinygrad.codegen.opt.kernel import axis_colors
-
-def rename_sink(s:UOp):
-  if s.arg is not None and s.arg.name != "test": return None
-
-  # get all ranges (sorted)
-  rngs = sorted([u for u in s.parents if u.op is Ops.RANGE], key=lambda x: x.arg[0:-1])
-
-  # add name to kernel
-  name = "k" + colored('_', 'BLACK').join(['']+[colored(x.src[0].render(), axis_colors[x.arg[-1]]) for x in rngs])
-  return s.replace(arg=KernelInfo(name=name) if s.arg is None else replace(s.arg, name=name))
-
-pm_postrange_opt = PatternMatcher([
-  (UPat(Ops.SINK, name="s"), rename_sink),
-])

tinygrad/codegen/opt/swizzler.py

@@ -128,8 +128,7 @@ fix_kernel_ops = view_left_through_load+PatternMatcher([
   (UPat(Ops.VIEW, src=(UPat.cvar(),), name="self"),
    lambda self: UOp.where(UOp(Ops.VALID, dtypes.bool, (UOp(Ops.VIEW, arg=self.st),)), self.const_like(self.base.arg), 0)),
   # no ImageDType after index
-  (UPat(GroupOp.All-{Ops.DEFINE_GLOBAL, Ops.VIEW, Ops.INDEX}, name="x"),
-   lambda x: x.replace(dtype=x.dtype.base) if isinstance(x.dtype, ImageDType) else None),
+  (UPat(GroupOp.All-{Ops.DEFINE_GLOBAL, Ops.VIEW}, name="x"), lambda x: x.replace(dtype=x.dtype.base) if isinstance(x.dtype, ImageDType) else None),
   # if this kernel also assigns to the loaded buffer, ensure we can index it correctly
   (UPat(Ops.LOAD, src=(UPat.var("glbl").view(name="view"),)), check_load_st),
 ])

tinygrad/codegen/opt/tc.py

@@ -22,15 +22,6 @@ class TensorCore: # D = A * B + C, A is (M x K), B is (K x N), C and D are (M x
   def permutes_for_shape_str(self, shape_str:list[str]) -> tuple[tuple[int, ...], tuple[int, ...]]:
     ret = [[shape_str.index(remap[ss]) if ss in remap else i for i,ss in enumerate(shape_str)] for remap in self._remaps()]
     return tuple(ret[0]), tuple(ret[1])
-  @functools.cache  # pylint: disable=method-cache-max-size-none
-  def base_shape_str(self) -> list[str]:
-    ret = []
-    cnt = {'u': 0, 'l': 0}
-    for opt in self.opts:
-      ret.append(f"{opt[0]}{cnt[opt[0]]}")
-      cnt[opt[0]] += 1
-    # assumes you do the UNROLL after the opts
-    return ret + [f"r{i}" for i in range(len(self.get_reduce_axes()))]
   def get_reduce_axes(self): return [(i, 2) for i in range(int(math.log2(self.dims[2])))]
   def get_upcast_axes(self): return [opt for opt in self.opts if opt[0] == "u"]
   def get_local_axes(self): return [opt for opt in self.opts if opt[0] == "l"]

tinygrad/device.py

@@ -139,7 +139,7 @@ class Buffer:
       if PROFILE:
         self._prof_num = num = len(Buffer.profile_events)
         ts = decimal.Decimal(time.perf_counter_ns())/1000
-        Buffer.profile_events.append(ProfilePointEvent(self.device, "alloc", ts, num, {"dtype":self.dtype, "sz":self.size}))
+        Buffer.profile_events.append(ProfilePointEvent(self.device, "alloc", ts, num, {"dtype":str(self.dtype),"sz":self.size,"nbytes":self.nbytes}))
     return self
   def deallocate(self):
     assert hasattr(self, '_buf'), "buffer must be allocated to deallocate"

tinygrad/dtype.py

@@ -108,6 +108,7 @@ class dtypes:
     if isinstance(val, tuple):
       assert len(val) == dtype.count, f"mismatch {val} {dtype}"
       return tuple(dtypes.as_const(x, dtype) for x in val)
+    # TODO: should truncate here
     return int(val) if dtypes.is_int(dtype) else float(val) if dtypes.is_float(dtype) else bool(val)
   @staticmethod
   @functools.cache
@@ -214,14 +215,15 @@ def sum_acc_dtype(dt:DType):
   return least_upper_dtype(dt, to_dtype(getenv("SUM_DTYPE", "float32")))
 
 def truncate_fp16(x):
-  try: return struct.unpack('e', struct.pack('e', float(x)))[0]
+  try: return struct.unpack("@e", struct.pack("@e", float(x)))[0]
   except OverflowError: return math.copysign(math.inf, x)
 
-def float_to_bf16(x):
-  if not math.isfinite(x): return x
-  u = struct.unpack('I', struct.pack('f', x))[0]
-  u = (u + 0x7FFF + ((u >> 16) & 1)) & 0xFFFF0000
-  return struct.unpack('f', struct.pack('I', u))[0]
+def truncate_bf16(x):
+  max_bf16 = struct.unpack('f', struct.pack('I', 0x7f7f0000))[0]
+  if abs(x) > max_bf16: return math.copysign(math.inf, x)
+  f32_int = struct.unpack('I', struct.pack('f', x))[0]
+  bf = struct.unpack('f', struct.pack('I', f32_int & 0xFFFF0000))[0]
+  return bf
 
 # fp8-float conversions based on https://gitlab.com/nvidia/headers/cuda-individual/cudart/-/blob/main/cuda_fp8.hpp
 def float_to_fp8(x: float, dtype: DType) -> int:
@@ -286,7 +288,7 @@ def fp8_to_float(x: int, dtype: DType) -> float:
   return float(float32_val)
 
 truncate: dict[DType, Callable] = {dtypes.bool: bool,
-  dtypes.float16: truncate_fp16, dtypes.bfloat16: lambda x: float_to_bf16(float(x)),
+  dtypes.float16: truncate_fp16, dtypes.bfloat16: truncate_bf16,
   **{fp8: (lambda x, dtype=fp8: fp8_to_float(float_to_fp8(x, dtype), dtype)) for fp8 in dtypes.fp8s},
   dtypes.float32: lambda x: ctypes.c_float(x).value, dtypes.float64: lambda x: ctypes.c_double(x).value,
   dtypes.uint8: lambda x: ctypes.c_uint8(x).value, dtypes.uint16: lambda x: ctypes.c_uint16(x).value,

tinygrad/engine/memory.py

@@ -23,13 +23,12 @@ def _internal_memory_planner(buffers:list[list[Buffer]], noopt_buffers=None, ign
   # Sort buffer operations in timeline order. Two events: buffer is allocated or buffer is freed.
   buffer_requests = sorted([((first_appearance[buf], True), buf) for buf in first_appearance.keys()] + \
                            [((last_appearance[buf] + 1, False), buf) for buf in first_appearance.keys()], key=lambda x: x[0])
-  total_memory = sum(round_up(buf.nbytes, min_block_size:=0x1000) for buf in first_appearance.keys()) * 2 # *2 for fragmentation (which is about 15%)
 
   # Try to suballocate from a shared buffer managed by global_planner using TLSFAllocator.
   # Also track buffer replacements for buffers that do not support suballocation.
   buffer_replace:dict[Buffer, tuple[Buffer|None, int|None]] = {}
   reuse_buffers:dict[tuple, list[Buffer]] = defaultdict(list)
-  global_planner:dict[str, tuple[int, TLSFAllocator]] = defaultdict(lambda: (0, TLSFAllocator(total_memory, block_size=min_block_size, lv2_cnt=32)))
+  global_planner:dict[str, tuple[int, TLSFAllocator]] = defaultdict(lambda: (0, TLSFAllocator(1 << 44, block_size=0x1000, lv2_cnt=32)))
   for (_, is_open_ev), buf in buffer_requests:
     # Check if suballocation is possible for the given buffer and device.
     if hasattr(Device[buf.device].allocator, "_offset") and not isinstance(buf.dtype, ImageDType):

tinygrad/engine/realize.py

@@ -1,8 +1,8 @@
 from typing import cast, Generator
-import time, pprint, decimal
+import time, pprint
 from dataclasses import dataclass, replace, field
 from tinygrad.helpers import all_same, colored, DEBUG, GlobalCounters, ansilen, BEAM, NOOPT, all_int, CAPTURING, Metadata, TRACEMETA, TracingKey
-from tinygrad.helpers import DEVECTORIZE, time_to_str, VALIDATE_WITH_CPU, getenv, cpu_profile, PROFILE, ProfilePointEvent, cpu_events
+from tinygrad.helpers import DEVECTORIZE, time_to_str, VALIDATE_WITH_CPU, getenv, cpu_profile
 from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, Variable, sym_infer, graph_rewrite, print_uops, track_rewrites, KernelInfo
 from tinygrad.device import Device, Buffer
 from tinygrad.renderer import Renderer, ProgramSpec, Estimates
@@ -149,8 +149,6 @@ class ExecItem:
   def run(self, _var_vals:dict[Variable, int]|None=None, wait=False, jit=False, do_update_stats=True) -> float|None:
     var_vals = self.fixedvars if _var_vals is None else (_var_vals|self.fixedvars)
     bufs = [cast(Buffer, x) for x in self.bufs] if jit else [cast(Buffer, x).ensure_allocated() for x in self.bufs]
-    if PROFILE: cpu_events.append(ProfilePointEvent(self.prg.device, "exec", decimal.Decimal(time.perf_counter_ns())/1000, self.prg.display_name,
-                                                    {"metadata":self.metadata, "var_vals":var_vals}))
     et = self.prg(bufs, var_vals, wait=wait or DEBUG >= 2)
     if do_update_stats:
       GlobalCounters.kernel_count += 1
@@ -160,15 +158,10 @@ class ExecItem:
       if DEBUG >= 2:
         lds_est = sym_infer(self.prg.estimates.lds, var_vals)
         mem_est = min(mem_est, lds_est)   # there can't be more memory accessed than loads/stores. remove this when symbolic is fixed
-        header_color = 'magenta' if jit else ('green' if self.prg.first_run else None)
         ptm = colored(time_to_str(et, w=9), "yellow" if et > 0.01 else None) if et is not None else ""
-        flops, membw, ldsbw = op_est/(et or 1e-20), mem_est/(et or 1e-20), lds_est/(et or 1e-20)
-        flops_str = f"{flops*1e-9:9.2f} GFLOPS" if flops < 1e14 else colored(f"{flops*1e-12:9.2f} TFLOPS", 'green')
-        mem_str = f"{membw*1e-9:6.1f}|{ldsbw*1e-9:<7.1f} GB/s" if membw < 1e13 else colored(f"{membw*1e-12:6.1f}|{ldsbw*1e-12:<7.1f} TB/s", 'green')
-        print(f"{colored(f'*** {self.prg.device[:7]:7s} {GlobalCounters.kernel_count:4d}', header_color)}"+
-          f" {self.prg.display_name+' '*(44-ansilen(self.prg.display_name))} arg {len(bufs):2d} mem {GlobalCounters.mem_used/1e9:5.2f} GB"+
-          ("" if et is None else f" tm {ptm}/{GlobalCounters.time_sum_s*1e3:9.2f}ms ({flops_str} {mem_str})")+
-          f" {[repr(m) if TRACEMETA >= 2 else str(m) for m in self.metadata] if self.metadata else ''}")
+        print(f"{colored(f'*** {self.prg.device[:7]:7s} {GlobalCounters.kernel_count:4d}', 'magenta' if jit else ('green' if self.prg.first_run else None))} {self.prg.display_name+' '*(44-ansilen(self.prg.display_name))} arg {len(bufs):2d} mem {GlobalCounters.mem_used/1e9:5.2f} GB " +  # noqa: E501
+              (str() if et is None else f"tm {ptm}/{GlobalCounters.time_sum_s*1e3:9.2f}ms ({op_est/((et or 1e-20)*1e9):9.2f} GFLOPS {mem_est/((et or 1e-20)*1e9):6.1f}|{lds_est/((et or 1e-20)*1e9):<7.1f} GB/s)" +  # noqa: E501
+               f" {[repr(m) if TRACEMETA >= 2 else str(m) for m in self.metadata] if self.metadata else ''}"))
       self.prg.first_run = False
     return et
 

tinygrad/engine/schedule.py

@@ -33,7 +33,7 @@ def create_schedule_with_vars(sched_sink:UOp) -> tuple[list[ScheduleItem], dict[
         for ss in s.src:
           if ss.op is Ops.MSELECT: ss = ss.src[0]
           if ss.op is not Ops.BUFFER:
-            assert ss.op is Ops.ASSIGN, f"ss.op is not ASSIGN, it's {ss.op}"
+            assert ss.op is Ops.ASSIGN
             children[ss.src[1]].append(k)
             in_degree[k] += 1
       elif s.op is Ops.BUFFER:

tinygrad/frontend/onnx.py

@@ -21,9 +21,9 @@ class AttributeType(enum.IntEnum):
   ONNX attribute type identifiers.
   Reference: https://github.com/onnx/onnx/blob/rel-1.18.0/onnx/onnx.proto3#L128-L145
   """
-  FLOAT = 1; INT = 2; STRING = 3; TENSOR = 4; GRAPH = 5; FLOATS = 6; INTS = 7; STRINGS = 8 # noqa: E702
+  FLOAT = 1; INT = 2; STRING = 3; TENSOR = 4; FLOATS = 6; INTS = 7; STRINGS = 8 # noqa: E702
 
-  def to_field_name(self) -> str: return {1: "f", 2: "i", 3: "s", 4: "t", 5: "g", 6: "floats", 7: "ints", 8: "strings"}[self.value]
+  def to_field_name(self) -> str: return {1: "f", 2: "i", 3: "s", 4: "t", 6: "floats", 7: "ints", 8: "strings"}[self.value]
 
 class OnnxDataType(enum.IntEnum):
   """
@@ -266,7 +266,6 @@ class OnnxPBParser:
         case 3: obj["i"] = self.reader.read_int64()
         case 4: obj["s"] = self.reader.read_bytes().data().tobytes().decode("utf8")
         case 5: obj["t"] = self._parse_TensorProto()['parsed_tensor']
-        case 6: obj["g"] = OnnxRunner._from_subgraph(self._parse_GraphProto())
         case 7: obj["floats"].append(self.reader.read_float())
         case 8: obj["ints"].append(self.reader.read_int64())
         case 9: obj["strings"].append(self.reader.read_bytes().data().tobytes().decode("utf8"))
@@ -402,11 +401,8 @@ class OnnxRunner:
   """
   def __init__(self, model_path: Tensor | str | pathlib.Path):
     model = OnnxPBParser(model_path, load_external_data=True).parse()
-    self._init_from_graph(model["graph"])
-
-  def _init_from_graph(self, graph: dict, is_subgraph: bool = False):
+    graph = model["graph"]
     self.is_training = any(n['parsed_node'].opset_id.domain in {Domain.AI_ONNX_TRAINING, Domain.AI_ONNX_PREVIEW_TRAINING} for n in graph["node"])
-    self.graph_name = graph["name"] if is_subgraph else ""
     self.graph_values = {"": None, **{i["name"]: i["parsed_tensor"] for i in graph["initializer"]}}
     self.graph_inputs = {i["name"]: i["parsed_type"] for i in graph["input"] if i["name"] not in self.graph_values}
     self.graph_outputs = tuple(o["name"] for o in graph["output"])
@@ -418,12 +414,6 @@ class OnnxRunner:
     self.variable_dims: dict[str, int] = {}
     self.onnx_ops = onnx_ops
 
-  @classmethod
-  def _from_subgraph(cls, graph: dict) -> "OnnxRunner":
-    subgraph = cls.__new__(cls)
-    subgraph._init_from_graph(graph, is_subgraph=True)
-    return subgraph
-
   def _parse_input(self, name: str, value: Any, spec: OnnxValue):
     if spec.is_optional and value is None: return None
     if spec.is_sequence:
@@ -455,10 +445,9 @@ class OnnxRunner:
     return {name:Tensor.empty(*spec.shape, device=device, dtype=dtype or spec.dtype) for name, spec in self.graph_inputs.items()}
 
   def to(self, device:str|None):
-    self.graph_values = {k: (v.to(device) if isinstance(v, Tensor) else v) for k,v in self.graph_values.items()}
+    self.graph_values = {k:v.to(device) if isinstance(v, Tensor) else v for k,v in self.graph_values.items()}
     self.graph_nodes = tuple(OnnxNode(n.op, n.opset_id, tuple(n.inputs), tuple(n.outputs),
-                                      {k: (v.to(device) if isinstance(v, (Tensor, OnnxRunner)) else v) for k,v in n.opts.items()})
-                                      for n in self.graph_nodes)
+                                      {k:v.to(device) if isinstance(v, Tensor) else v for k,v in n.opts.items()}) for n in self.graph_nodes)
     return self
 
   def __call__(self, inputs:dict[str, Any], debug=debug):
@@ -472,9 +461,9 @@ class OnnxRunner:
 
       # provide additional opts
       if node.op == "Split" and 'num_outputs' not in opts: opts['num_outputs'] = len(node.outputs)
-      if node.op in {"Gradient", "If"}: opts['intermediate_tensors'] = self.graph_values
+      if node.op == "Gradient": opts['intermediate_tensors'] = self.graph_values
 
-      if debug >= 1: print((f"[{self.graph_name}] " if self.graph_name else "") + f"{num}: op '{node.op}' opt {opts}")
+      if debug >= 1: print(f"{num}: op '{node.op}' opt {opts}")
       if debug >= 2 and node.inputs: print("\tinputs:\n" + "\n".join(f"\t\t{x} - {i!r}" for x,i in zip(node.inputs, inps)))
       ret = self._select_op(node.op, node.opset_id)(*inps, **opts)
       ret = ret if isinstance(ret, tuple) else (ret,)
@@ -554,23 +543,6 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
     return __decorator
 
   # ***** Property/Graph Ops *****
-  def If(condition:Tensor, else_branch:OnnxRunner, then_branch:OnnxRunner, intermediate_tensors:dict[str, Tensor]):
-    def run_branch(branch:OnnxRunner):
-      branch.graph_values.update(intermediate_tensors)
-      out = branch({k:intermediate_tensors[k] for k in branch.graph_inputs.keys()})
-      # dereference intermediate tensors so Buffer can be deallocated
-      for k in intermediate_tensors: del branch.graph_values[k]
-      return out
-    # both branch must be ran before the condition can be evaluated
-    else_out, then_out = run_branch(else_branch), run_branch(then_branch)
-    assert len(else_out) == len(then_out), f"else_out and then_out must have the same number of outputs: {len(else_out)} != {len(then_out)}"
-    # can use where op when output shape is the same
-    if all(t.shape == e.shape for t,e in zip(then_out.values(), else_out.values())):
-      return tuple(condition.where(t,e) for t,e in zip(then_out.values(), else_out.values()))
-    # otherwise, use condition to select the output in python
-    cond = _resolve_const(_cached_to_python_const(condition))
-    return tuple(t if cond else e for t,e in zip(then_out.values(), else_out.values()))
-
   def Identity(x:Tensor): return x
   def Constant(sparse_value:Tensor|None=None, value:Tensor|None=None, value_float:float|None=None, value_floats:list[float]|None=None,
               value_int:int|None=None, value_ints:list[int]|None=None, value_string:str|None=None, value_strings:list[str]|None=None):

tinygrad/gradient.py

@@ -22,10 +22,11 @@ pm_gradient = PatternMatcher([
   (UPat(Ops.SQRT, name="ret"), lambda ctx, ret: (ctx / (ret*2),)),
   (UPat((Ops.CMPLT, Ops.CMPNE)), lambda: (None, None)),
   (UPat(Ops.ADD), lambda ctx: (ctx, ctx)),
-  (UPat(Ops.POW, name="ret", src=(UPat.var("b"), UPat.var("e"))), lambda ctx, ret, b, e:
-    (ctx * (b.eq(0)&e.eq(0)).where(e, e*b.pow(e-1)), ctx * b.eq(0).where((e<0).where(ret.const_like(-math.inf), 0), ret*b.log2()*math.log(2.0)))),
-  (UPat(Ops.MAX, name="ret", src=(UPat.var("x"), UPat.var("y"))), lambda ctx, ret, x, y:
-    ((x>y).where(ctx, (x.eq(y)).where(ctx * 0.5, 0)), (x<y).where(ctx, (x.eq(y)).where(ctx * 0.5, 0)))),
+  (UPat(Ops.POW, name="ret"), lambda ctx, ret:
+   (ctx*(ret.src[0].eq(0) & ret.src[1].eq(0)).where(ret.src[1], ret.src[1]*ret.src[0].pow(ret.src[1]-1)),
+    ctx*ret.src[0].eq(0).where((ret.src[1]<0).where(ret.const_like(-math.inf), ret.const_like(0)), ret*ret.src[0].log2()*math.log(2.0)))),
+  (UPat(Ops.MAX, name="ret"), lambda ctx, ret: ((ret.src[0]>ret.src[1]).where(ctx, (ret.src[0]!=ret.src[1]).where(ctx.const_like(0), ctx * 0.5)),
+                                                (ret.src[0]<ret.src[1]).where(ctx, (ret.src[0]!=ret.src[1]).where(ctx.const_like(0), ctx * 0.5)))),
   (UPat(Ops.MUL, name="ret"), lambda ctx, ret: (ret.src[1]*ctx, ret.src[0]*ctx)),
   (UPat(Ops.WHERE, name="ret"), lambda ctx, ret: (None, ret.src[0].where(ctx, ctx.const_like(0)), ret.src[0].where(ctx.const_like(0), ctx))),
   (UPat(Ops.REDUCE_AXIS, name="ret"), reduce_gradient),

tinygrad/helpers.py

@@ -135,12 +135,11 @@ FUSE_ARANGE, FUSE_CONV_BW = ContextVar("FUSE_ARANGE", 1), ContextVar("FUSE_CONV_
 SPLIT_REDUCEOP, NO_MEMORY_PLANNER, RING = ContextVar("SPLIT_REDUCEOP", 1), ContextVar("NO_MEMORY_PLANNER", 0), ContextVar("RING", 1)
 PICKLE_BUFFERS, PROFILE, LRU = ContextVar("PICKLE_BUFFERS", 1), ContextVar("PROFILE", getenv("VIZ")), ContextVar("LRU", 1)
 CACHELEVEL, IGNORE_BEAM_CACHE, DEVECTORIZE = ContextVar("CACHELEVEL", 2), ContextVar("IGNORE_BEAM_CACHE", 0), ContextVar("DEVECTORIZE", 1)
-DISABLE_COMPILER_CACHE, BLOCK_REORDER = ContextVar("DISABLE_COMPILER_CACHE", 0), ContextVar("BLOCK_REORDER", 1)
+DISABLE_COMPILER_CACHE = ContextVar("DISABLE_COMPILER_CACHE", 0)
 DONT_REALIZE_EXPAND, DONT_GROUP_REDUCES = ContextVar("DONT_REALIZE_EXPAND", 0), ContextVar("DONT_GROUP_REDUCES", 0)
 QUANTIZE, VALIDATE_WITH_CPU, DISABLE_FAST_IDIV = ContextVar("QUANTIZE", 0), ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("DISABLE_FAST_IDIV", 0)
 CORRECT_DIVMOD_FOLDING, FUSE_OPTIM = ContextVar("CORRECT_DIVMOD_FOLDING", 0), ContextVar("FUSE_OPTIM", 0)
 ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, AMD_LLVM = ContextVar("ALLOW_DEVICE_USAGE", 1), ContextVar("MAX_BUFFER_SIZE", 0), ContextVar("AMD_LLVM", 1)
-RANGEIFY, POSTOPT, FUSE_ATTENTION = ContextVar("RANGEIFY", 0), ContextVar("POSTOPT", 0), ContextVar("FUSE_ATTENTION", 0)
 
 @dataclass(frozen=True)
 class Metadata:
@@ -196,7 +195,7 @@ class Profiling(contextlib.ContextDecorator):
 @dataclass(frozen=True)
 class TracingKey:
   display_name:str                       # display name of this trace event
-  keys:tuple[Any, ...]=()                # optional keys to search for related traces
+  keys:tuple[str, ...]=()                # optional keys to search for related traces
   cat:str|None=None                      # optional category to color this by
   ret:Any=None
 
@@ -206,7 +205,7 @@ class ProfileEvent: pass
 class ProfileRangeEvent(ProfileEvent): device:str; name:str|TracingKey; st:decimal.Decimal; en:decimal.Decimal|None=None; is_copy:bool=False # noqa: E702
 
 @dataclass(frozen=True)
-class ProfilePointEvent(ProfileEvent): device:str; name:str; ts:decimal.Decimal; key:Any; arg:dict=field(default_factory=dict) # noqa: E702
+class ProfilePointEvent(ProfileEvent): device:str; name:str; ts:decimal.Decimal; key:int; arg:dict=field(default_factory=dict) # noqa: E702
 
 cpu_events:list[ProfileEvent] = []
 @contextlib.contextmanager

tinygrad/nn/state.py

@@ -274,9 +274,9 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
   Converts ggml tensor data to a tinygrad tensor.
 
   Supported native types: float32 (id: 0), float16 (id: 1), int8 (id: 16), int16 (id: 17), int32 (id: 18)
-  Supported quantized types: Q4_0 (id: 2), Q4_1 (id: 3), Q8_0 (id: 8), Q6_K (id: 14), MXFP4 (id: 39)
+  Supported quantized types: Q4_0 (id: 2), Q4_1 (id: 3), Q8_0 (id: 8), Q6_K (id: 14)
   """
-  # https://github.com/ggerganov/ggml/blob/323951f1bdcdfbd5b5ff3a9a7c3770e63b1a560e/include/ggml.h#L356
+  # https://github.com/ggerganov/ggml/blob/6dccc647264f5429df2624f36138f601e7ce23e5/include/ggml.h#L356
 
   # native types
   if (dtype := { 0: dtypes.float32, 1: dtypes.float16, 16: dtypes.int8, 17: dtypes.int16, 18: dtypes.int32 }.get(ggml_type)) is not None:
@@ -288,7 +288,7 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
     return t.unsqueeze(-1).expand((*t.shape,8//b)).idiv(shift_tensor).bitwise_and(bitmask).transpose(-1, -2).flatten(-2)
 
   # map to (number of elements, number of bytes)
-  if (nelements_nbytes := { 2: (32, 18), 3: (32, 20), 14: (256, 210), 8: (32, 34), 39: (32, 17) }.get(ggml_type)) is not None:
+  if (nelements_nbytes := { 2: (32, 18), 3: (32, 20), 14: (256, 210), 8: (32, 34) }.get(ggml_type)) is not None:
     blocks = t[:(n//nelements_nbytes[0])*nelements_nbytes[1]].reshape((-1, nelements_nbytes[1]))
     if ggml_type == 2: return (q_to_uint8(blocks[:,2:], 4).bitcast(dtypes.int8) - 8) * blocks[:,:2].bitcast(dtypes.float16).cast(dtypes.float32)
     if ggml_type == 3:
@@ -300,17 +300,6 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
       scales = blocks[:,192:208].bitcast(dtypes.int8).unsqueeze(-1).expand((-1, 16, 16)).reshape((-1, 256))
       d = blocks[:,-2:].bitcast(dtypes.float16).cast(dtypes.float32).expand((-1, 256))
       return d * (xl.bitwise_or(xh).bitcast(dtypes.int8) - 32).flatten(-2) * scales
-    if ggml_type == 39:
-      e_int = blocks[:, 0].cast(dtypes.int32)
-      d = ((e_int >= 2).cast(dtypes.float32) * (e_int.cast(dtypes.float32) - 128).exp2() +
-           (e_int == 1).cast(dtypes.float32) * 2.0**(-127) +
-           (e_int == 0).cast(dtypes.float32) * 2.0**(-128)).unsqueeze(-1)
-      codes = q_to_uint8(blocks[:, 1:17], 4)
-      sign = 1.0 - codes.rshift(3).cast(dtypes.float32) * 2.0
-      exp, mant = codes.rshift(1).bitwise_and(0x3).cast(dtypes.float32), codes.bitwise_and(0x1).cast(dtypes.float32)
-      fp4_val = sign * ((exp != 0).cast(dtypes.float32) * (1.0 + 0.5 * mant) * (exp - 1.0).exp2() +
-                        (exp == 0).cast(dtypes.float32) * 0.5 * mant)
-      return (fp4_val * d).flatten(-2)[:n]
   raise ValueError(f"GGML type '{ggml_type}' is not supported!")
 
 @accept_filename

tinygrad/renderer/cstyle.py

@@ -6,7 +6,7 @@ from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat
 from tinygrad.helpers import strip_parens, getenv, prod, dedup, AMX
 from tinygrad.dtype import ImageDType, dtypes, DType, PtrDType, AddrSpace, truncate
 from tinygrad.renderer import Renderer
-from tinygrad.codegen.late.devectorizer import no_vectorized_alu
+from tinygrad.codegen.devectorizer import no_vectorized_alu
 
 base_rewrite = PatternMatcher([
   (UPat(Ops.DEFINE_REG, name="x"), lambda ctx,x: f"{ctx.render_dtype(x.dtype.base)} {ctx[x]}[{x.dtype.size}];"),
@@ -157,7 +157,7 @@ class CStyleLanguage(Renderer):
       # naming
       prefix = None
       if u.op is Ops.SPECIAL: r[u] = u.arg[0]
-      elif u.op is Ops.RANGE: r[u] = "ridx"+'_'.join([str(x) if x >= 0 else "m"+str(-x) for x in u.arg[0:-1]])
+      elif u.op is Ops.RANGE: r[u] = f"ridx{u.arg}"
       else:
         prefix = {Ops.WMMA: "wmma", Ops.DEFINE_LOCAL: "temp", Ops.CONST: "const",
                   Ops.CAST: "cast", Ops.BITCAST: "cast", Ops.GEP: "gep", Ops.VECTORIZE: "cast", Ops.PRECAST: "precast",
@@ -199,13 +199,12 @@ class ClangRenderer(CStyleLanguage):
   # language options
   buffer_suffix = " restrict"
   type_map = {dtypes.bool:"_Bool", dtypes.half:"__fp16"}
-  code_for_op = {**({k:v for k,v in CStyleLanguage.code_for_op.items() if k not in [Ops.EXP2, Ops.SIN, Ops.LOG2, Ops.TRUNC, Ops.RECIP]}),
+  code_for_op = {**({k:v for k,v in CStyleLanguage.code_for_op.items() if k not in [Ops.EXP2, Ops.SIN, Ops.LOG2, Ops.TRUNC]}),
                  Ops.SQRT: lambda x,dtype: f"__builtin_sqrt({x})" if dtype == dtypes.float64 else f"__builtin_sqrtf({x})",
-                 Ops.TRUNC: lambda x,dtype: f"__builtin_trunc({x})" if dtype == dtypes.float64 else f"__builtin_truncf({x})",
-                 Ops.FDIV: lambda a,b,dtype: f"({a}/{b})"}
+                 Ops.TRUNC: lambda x,dtype: f"__builtin_trunc({x})" if dtype == dtypes.float64 else f"__builtin_truncf({x})"}
   # LLVM legalizes double => half cast on systems that don't support it natively (like x86 cpus without AVX512-FP16) into a compiler-rt libcall.
   extra_matcher = PatternMatcher([(UPat.var("x", dtypes.float64).cast(dtypes.float16), lambda x: x.cast(dtypes.float32).cast(dtypes.float16)),
-    (UPat((Ops.SQRT, Ops.TRUNC), name="alu"), no_vectorized_alu)]) + CStyleLanguage.extra_matcher
+    (UPat((Ops.SQRT, Ops.TRUNC), name="alu"), no_vectorized_alu),]) + CStyleLanguage.extra_matcher
 
   if sys.platform == 'win32':
     kernel_typedef = "__attribute__((ms_abi)) void"
@@ -316,9 +315,7 @@ class MetalRenderer(CStyleLanguage):
 
   def render_kernel(self, function_name, kernel, bufs, uops, prefix=None):
     prefix = ["#include <metal_stdlib>","using namespace metal;"]
-    wargs = wmma_args(uops)
-    if len(wargs) > 0: wargs = wargs[0:1]
-    for name, _, dtype_in, dtype_out, _, _, _, _ in wargs: prefix.append(
+    for name, _, dtype_in, dtype_out, _, _, _, _ in wmma_args(uops): prefix.append(
   f"""{(dstr_out:=self.render_dtype(dtype_out.vec(2)))} __{name}({(dstr_in:=self.render_dtype(dtype_in.vec(2)))} a, {dstr_in} b, {dstr_out} c){{
   simdgroup_{self.render_dtype(dtype_in)}8x8 mat_a, mat_b; simdgroup_{self.render_dtype(dtype_out)}8x8 mat_c;
   mat_a.thread_elements()[0] = a[0]; mat_b.thread_elements()[0] = b[0]; mat_c.thread_elements()[0] = c[0];

tinygrad/renderer/llvmir.py

@@ -45,10 +45,10 @@ def render_wmma_amx(ctx, wmma: UOp) -> str:
       f'  call void asm sideeffect "nop\\0Anop\\0Anop\\0A.word ({0x201000 + (17 << 5) + 1})", "~{{memory}}"() #0; AMX clr',             # clr
       f'  {ctx[wmma]} = load {ldt(wmma.dtype)}, ptr {ctx[wmma]}_amx2, align {wmma.dtype.itemsize}'])
 
-def render_wmma_amd(ctx, wmma: UOp, cdna=False) -> str:
-  dt_map = {dtypes.half: "f16", dtypes.float: "f32", dtypes.ushort: "bf16.1k" if cdna else "bf16", dtypes.bfloat16: "bf16.1k" if cdna else "bf16"}
+def render_wmma_amd(ctx, wmma: UOp, arch: str) -> str:
+  dt_map = {dtypes.half: "f16", dtypes.float: "f32", dtypes.bfloat16: "bf16", dtypes.ushort: "bf16"}
   # https://github.com/llvm/llvm-project/blob/main/clang/test/CodeGenOpenCL/builtins-amdgcn-mfma.cl
-  if cdna:
+  if arch.split(":")[0] in {"gfx942", "gfx950"}:
     return f"  {ctx[wmma]} = call {ldt(wmma.dtype)} @llvm.amdgcn.mfma.{dt_map[wmma.src[-1].dtype.scalar()]}" + \
            f".16x16x16{dt_map[wmma.src[0].dtype.scalar()]}(" + ", ".join([f"{ldt(w.dtype)} {ctx[w]}" for w in wmma.src]) + ", i32 0, i32 0, i32 0)"
   # https://github.com/llvm/llvm-project/blob/main/llvm/test/CodeGen/AMDGPU/GlobalISel/llvm.amdgcn.wmma_32.ll
@@ -101,13 +101,13 @@ base_rewrite = PatternMatcher([
 
   # range
   (UPat(Ops.RANGE, name="x"), lambda ctx,x:
-   f"  br label %loop_entry_{x.arg[0]}\nloop_entry_{x.arg[0]}:\n"
-   f"  br label %loop_body_{x.arg[0]}\nloop_body_{x.arg[0]}:\n"
-   f"  {ctx[x]} = phi {ldt(x.dtype)} [ 0, %loop_entry_{x.arg[0]} ], [ {ctx[x]}phi, %loop_latch_{x.arg[0]} ]"),
+   f"  br label %loop_entry_{x.arg}\nloop_entry_{x.arg}:\n"
+   f"  br label %loop_body_{x.arg}\nloop_body_{x.arg}:\n"
+   f"  {ctx[x]} = phi {ldt(x.dtype)} [ 0, %loop_entry_{x.arg} ], [ {ctx[x]}phi, %loop_latch_{x.arg} ]"),
   (UPat(Ops.ENDRANGE, name="x"), lambda ctx,x:
-   f"  br label %loop_latch_{x.src[0].arg[0]}\nloop_latch_{x.src[0].arg[0]}:\n"
+   f"  br label %loop_latch_{x.src[0].arg}\nloop_latch_{x.src[0].arg}:\n"
    f"  {ctx[x.src[0]]}phi = add i32 {ctx[x.src[0]]}, 1\n  {ctx[x]} = icmp ult i32 {ctx[x.src[0]]}phi, {ctx[x.src[0].src[0]]}\n"
-   f"  br i1 {ctx[x]}, label %loop_body_{x.src[0].arg[0]}, label %loop_exit_{x.src[0].arg[0]}\nloop_exit_{x.src[0].arg[0]}:"),
+   f"  br i1 {ctx[x]}, label %loop_body_{x.src[0].arg}, label %loop_exit_{x.src[0].arg}\nloop_exit_{x.src[0].arg}:"),
 
   # if
   (UPat(Ops.IF, name="x"), lambda ctx,x: f"  br i1 {ctx[x.src[0]]}, label %ifbody_{ctx[x][1:]}, label %ifskip_{ctx[x][1:]}\nifbody_{ctx[x][1:]}:"),
@@ -123,10 +123,11 @@ class LLVMRenderer(Renderer):
   has_local = False
   global_max: tuple[int, ...] | None = None
   string_rewrite = base_rewrite + PatternMatcher([(UPat(Ops.WMMA, name="wmma"), render_wmma_amx)])
-  code_for_op = {Ops.FDIV: lambda: None}
   if AMX: tensor_cores = tc.amx
 
   extra_matcher = PatternMatcher([
+    # rewrite RECIP with FDIV
+    (UPat(Ops.RECIP, name="x"), lambda x: UOp(Ops.FDIV, x.dtype, (x.const_like(1), x.src[0]))),
     # rewrite cast to bool to CMPNE 0
     (UPat(Ops.CAST, dtype=dtypes.bool, name="x"), lambda x: x.src[0] != x.src[0].const_like(0)),
     # rewrite MAX to CMPLT + WHERE
@@ -221,14 +222,7 @@ class AMDLLVMRenderer(LLVMRenderer):
   def __init__(self, arch:str):
     self.arch = arch
     self.tensor_cores = AMDRenderer.get_tensor_cores(arch)
-    self.is_cdna = arch.split(":")[0] in {"gfx942", "gfx950"}
-    self.string_rewrite += PatternMatcher([(UPat(Ops.WMMA, name="wmma"), lambda ctx, wmma, cdna=self.is_cdna: render_wmma_amd(ctx, wmma, cdna))])
-    if self.is_cdna:
-      self.extra_matcher += PatternMatcher([
-        (UPat(Ops.WMMA, name="x", dtype=dtypes.float.vec(4)),
-          lambda x: UOp(Ops.WMMA, dtypes.float.vec(4), (x.src[0].bitcast(dtypes.uint16.vec(4)), x.src[1].bitcast(dtypes.uint16.vec(4)),
-            x.src[2]), (*x.arg,)) if x.src[0].dtype == dtypes.bfloat16.vec(4) else None)
-      ])
+    self.string_rewrite += PatternMatcher([(UPat(Ops.WMMA, name="wmma"), lambda ctx, wmma, arch=arch: render_wmma_amd(ctx, wmma, arch))])
     if self.arch.split(":")[0] == "gfx1100":
       self.extra_matcher += PatternMatcher([
         (UPat(Ops.WMMA, name="x", dtype=dtypes.half.vec(8)),

tinygrad/renderer/ptx.py

@@ -119,7 +119,7 @@ string_rewrite = PatternMatcher([
     ctx.code_for_op[Ops.CMPLT](ctx.r[x], ctx.r[x.src[0]], ctx.r[src0.src[0]], dtypes.int, ctx.types[dtypes.int]),
     f"@{ctx.r[x]} bra LOOP_{ctx.r[src0][1:]};"]),
   (UPat(Ops.DEFINE_LOCAL, name="x"),
-   lambda ctx, x: [f".shared .align 16 .b8 local{x.arg}[{x.dtype.size*x.dtype.itemsize}];", f"mov.u64 {ctx.r[x]}, local{x.arg}[0];"]),
+   lambda ctx, x: [f".shared .align 16 .b8 {x.arg}[{x.dtype.size*x.dtype.itemsize}];", f"mov.u64 {ctx.r[x]}, {x.arg}[0];"]),
   (UPat(Ops.IF, name="x"), lambda ctx, x: f"@!{ctx.r[x.src[0]]} bra IF_{ctx.r[x.src[0]][1:]}_{ctx.uops.index(x)};"),
   (UPat(Ops.ENDIF, name="x"), lambda ctx, x: f"IF_{ctx.r[x.src[0].src[0]][1:]}_{ctx.uops.index(x.src[0])}:"),
   (UPat(Ops.WMMA, name="x"), lambda ctx, x: list(render_wmma(ctx, x))),
@@ -215,7 +215,7 @@ class PTXRenderer(Renderer):
                        [ssa("wmma_acc", dtype="b32") for _ in range(0, len(r[u.src[2]]), 4 // u.dtype.scalar().itemsize)]]
         r[u] = [ssa("wmma", dtype=self.types[u.dtype.scalar()]) for _ in range(u.dtype.count)]
       prefix, dtype = {Ops.CAST: ("cast", None), Ops.BITCAST: ("cast", None), Ops.ENDRANGE: ("pred", "pred"), Ops.RANGE: ("ridx", None),
-        Ops.DEFINE_VAR: ("dat", None), Ops.CONST: ("const", None), Ops.DEFINE_LOCAL: ("local",self.types[dtypes.ulong]),
+        Ops.DEFINE_VAR: ("dat", None), Ops.CONST: ("const", None), Ops.DEFINE_LOCAL:("local",self.types[dtypes.ulong]),
         Ops.DEFINE_GLOBAL: ("dat", self.types[dtypes.ulong]), **{op: ("alu", None) for op in GroupOp.ALU}}.get(u.op, (None, None))
       if prefix: r[u] = ssa(prefix, u, dtype)
 

tinygrad/runtime/ops_amd.py

@@ -4,11 +4,10 @@ import os, ctypes, ctypes.util, struct, hashlib, functools, importlib, mmap, err
 assert sys.platform != 'win32'
 from dataclasses import dataclass
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, HWQueue, CLikeArgsState, HCQSignal, HCQProgram, FileIOInterface
-from tinygrad.runtime.support.hcq import MMIOInterface, BumpAllocator
+from tinygrad.runtime.support.hcq import MMIOInterface
 from tinygrad.uop.ops import sint
 from tinygrad.device import Compiled, DMAFdRef, BufferSpec
 from tinygrad.helpers import getenv, to_mv, round_up, data64_le, all_same, flatten, DEBUG, AMD_LLVM, PROFILE, ProfileEvent, suppress_finalizing
-from tinygrad.helpers import lo32, hi32
 from tinygrad.renderer.cstyle import AMDRenderer
 from tinygrad.renderer.llvmir import AMDLLVMRenderer
 from tinygrad.runtime.autogen import kfd, hsa, pci, sqtt
@@ -25,8 +24,6 @@ EVENT_INDEX_PARTIAL_FLUSH = 4 # based on a comment in nvd.h
 WAIT_REG_MEM_FUNCTION_EQ  = 3 # ==
 WAIT_REG_MEM_FUNCTION_NEQ = 4 # !=
 WAIT_REG_MEM_FUNCTION_GEQ = 5 # >=
-AQL_HDR = (1 << hsa.HSA_PACKET_HEADER_BARRIER) | (hsa.HSA_FENCE_SCOPE_SYSTEM << hsa.HSA_PACKET_HEADER_SCACQUIRE_FENCE_SCOPE) \
-        | (hsa.HSA_FENCE_SCOPE_SYSTEM << hsa.HSA_PACKET_HEADER_SCRELEASE_FENCE_SCOPE)
 
 class AMDSignal(HCQSignal):
   def __init__(self, *args, **kwargs): super().__init__(*args, **{**kwargs, 'timestamp_divider': 100})
@@ -287,7 +284,7 @@ class AMDComputeQueue(HWQueue):
 
   def wait(self, signal:AMDSignal, value:sint=0):
     self.wait_reg_mem(mem=signal.value_addr, value=value, mask=0xffffffff)
-    if self.dev.xccs > 1 and not self.dev.is_aql: self.xcc_barrier()
+    if self.dev.xccs > 1: self.xcc_barrier()
     return self
 
   def timestamp(self, signal:AMDSignal):
@@ -332,41 +329,6 @@ class AMDComputeQueue(HWQueue):
     dev.compute_queue.put_value += len(cmds)
     dev.compute_queue.signal_doorbell(dev)
 
-class AMDComputeAQLQueue(AMDComputeQueue):
-  def exec(self, prg:AMDProgram, args_state:CLikeArgsState, global_size:tuple[sint, ...], local_size:tuple[sint, ...]):
-    self.bind_args_state(args_state)
-    self._q.append(pkt:=hsa.hsa_kernel_dispatch_packet_t(header=AQL_HDR | (hsa.HSA_PACKET_TYPE_KERNEL_DISPATCH << hsa.HSA_PACKET_HEADER_TYPE),
-      setup=3<<hsa.HSA_KERNEL_DISPATCH_PACKET_SETUP_DIMENSIONS, private_segment_size=prg.private_segment_size,
-      group_segment_size=prg.group_segment_size, kernel_object=prg.aql_prog_addr, kernarg_address=args_state.buf.va_addr))
-    self.bind_sints_to_mem(*local_size, mem=(pkt_view:=MMIOInterface(addr=ctypes.addressof(pkt), nbytes=ctypes.sizeof(pkt))), fmt='H', offset=4)
-    self.bind_sints_to_mem(*[l * g for l,g in zip(local_size, global_size)], mem=pkt_view, fmt='I', offset=12)
-
-  def bind(self, dev:AMDDevice): pass # not supported
-  def _submit(self, dev:AMDDevice):
-    pm4_batch:list[int] = []
-    aql_bytes = bytes()
-
-    def flush_pm4_batch():
-      nonlocal pm4_batch
-      if not pm4_batch: return bytes()
-      dev.pm4_ibs.cpu_view().view(off:=dev.pm4_ib_alloc.alloc(len(pm4_batch) * 4), fmt='I')[:len(pm4_batch)] = array.array('I', pm4_batch)
-      pkt = [AQL_HDR | (hsa.HSA_PACKET_TYPE_VENDOR_SPECIFIC << hsa.HSA_PACKET_HEADER_TYPE) | (1 << 16),
-        self.pm4.PACKET3(self.pm4.PACKET3_INDIRECT_BUFFER, 2), *data64_le(dev.pm4_ibs.va_addr+off), len(pm4_batch)|self.pm4.INDIRECT_BUFFER_VALID, 10]
-      pm4_batch.clear()
-      return bytes(array.array('I', pkt + [0] * 10))
-
-    for cmd in self._q:
-      if isinstance(cmd, hsa.hsa_kernel_dispatch_packet_t): aql_bytes += flush_pm4_batch() + bytes(cmd)
-      else: pm4_batch.append(cmd)
-    aql_bytes += flush_pm4_batch()
-
-    assert len(aql_bytes) < dev.compute_queue.ring.nbytes, "submit is too large for the queue"
-    cp_bytes = min(len(aql_bytes), (dev.compute_queue.ring.nbytes - (dev.compute_queue.put_value * 64) % dev.compute_queue.ring.nbytes))
-    dev.compute_queue.ring.view(offset=(dev.compute_queue.put_value * 64) % dev.compute_queue.ring.nbytes, fmt='B')[:cp_bytes] = aql_bytes[:cp_bytes]
-    if (tail_bytes:=(len(aql_bytes) - cp_bytes)) > 0: dev.compute_queue.ring.view(offset=0, fmt='B')[:tail_bytes] = aql_bytes[cp_bytes:]
-    dev.compute_queue.put_value += len(aql_bytes) // 64
-    dev.compute_queue.signal_doorbell(dev, doorbell_value=dev.compute_queue.put_value-1)
-
 class AMDCopyQueue(HWQueue):
   def __init__(self, dev, max_copy_size=0x40000000):
     self.dev, self.sdma, self.internal_cmd_sizes, self.max_copy_size = dev, dev.sdma, [], max_copy_size
@@ -464,19 +426,14 @@ class AMDProgram(HCQProgram):
     # TODO; this API needs the type signature of the function and global_size/local_size
     self.dev, self.name, self.lib = dev, name, lib
 
-    image, sections, relocs = elf_loader(self.lib)
-
-    rodata_entry = next((sh.header.sh_addr for sh in sections if sh.name == ".rodata"), -1)
-    assert rodata_entry >= 0, ".rodata section not found"
-
-    for apply_image_offset, rel_sym_offset, typ, addent in relocs:
-      if typ == 5: image[apply_image_offset:apply_image_offset+8] = struct.pack('<q', rel_sym_offset - apply_image_offset + addent) # R_AMDGPU_REL64
-      else: raise RuntimeError(f"unknown AMD reloc {typ}")
-
+    image, sections, _ = elf_loader(self.lib)
     self.lib_gpu = self.dev.allocator.alloc(round_up(image.nbytes, 0x1000), buf_spec:=BufferSpec(cpu_access=True, nolru=True))
     self.dev.allocator._copyin(self.lib_gpu, image)
     self.dev.synchronize()
 
+    rodata_entry = next((sh.header.sh_addr for sh in sections if sh.name == ".rodata"), -1)
+    text_entry = next((sh.header.sh_addr for sh in sections if sh.name == ".text"), -1)
+    assert rodata_entry >= 0 and text_entry >= 0, ".text or .rodata section not found"
     self.group_segment_size = image[rodata_entry:rodata_entry+4].cast("I")[0]
     self.private_segment_size = image[rodata_entry+4:rodata_entry+8].cast("I")[0]
     self.kernargs_segment_size = image[rodata_entry+8:rodata_entry+12].cast("I")[0]
@@ -494,8 +451,8 @@ class AMDProgram(HCQProgram):
     self.rsrc1: int = code.compute_pgm_rsrc1 | ((1 << 20) if (11,0,0) <= self.dev.target < (12,0,0) else 0)
     self.rsrc2: int = code.compute_pgm_rsrc2 | (lds_size << 15)
     self.rsrc3: int = image[rodata_entry+44:rodata_entry+48].cast("I")[0] # NOTE: kernel descriptor, not in amd_kernel_code_t struct
-    self.aql_prog_addr: int = self.lib_gpu.va_addr + rodata_entry
     self.prog_addr: int = self.lib_gpu.va_addr + rodata_entry + code.kernel_code_entry_byte_offset
+    if code.kernel_code_entry_byte_offset == 0: self.prog_addr = self.lib_gpu.va_addr + text_entry
     # Some programs use hsa_kernel_dispatch_packet_t to read workgroup sizes during execution.
     # The packet is represented as a pointer and set up in SGPRs. Space for the packet is allocated as part of the kernel arguments.
     self.enable_dispatch_ptr: int = code.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_DISPATCH_PTR
@@ -544,7 +501,7 @@ class AMDQueueDesc:
     return cls(ring=queues[0].ring, put_value=queues[0].put_value, doorbells=flatten(q.doorbells for q in queues),
                read_ptrs=flatten(q.read_ptrs for q in queues), write_ptrs=flatten(q.write_ptrs for q in queues))
 
-  def signal_doorbell(self, dev, doorbell_value:int|None=None):
+  def signal_doorbell(self, dev):
     for write_ptr in self.write_ptrs: write_ptr[0] = self.put_value
 
     # Ensure all prior writes are visible to the GPU.
@@ -552,7 +509,7 @@ class AMDQueueDesc:
 
     # Flush hdp if queue is in dev mem.
     if dev.is_am() and not dev.is_usb(): dev.iface.dev_impl.gmc.flush_hdp()
-    for doorbell in self.doorbells: doorbell[0] = self.put_value if doorbell_value is None else doorbell_value
+    for doorbell in self.doorbells: doorbell[0] = self.put_value
 
 class KFDIface:
   kfd:FileIOInterface|None = None
@@ -655,12 +612,12 @@ class KFDIface:
     stm = kfd.AMDKFD_IOC_MAP_MEMORY_TO_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(c_gpus), n_devices=1)
     assert stm.n_success == 1
 
-  def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0, xcc_id=0):
+  def create_queue(self, queue_type, ring, gart, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0, xcc_id=0):
     queue = kfd.AMDKFD_IOC_CREATE_QUEUE(KFDIface.kfd, ring_base_address=ring.va_addr, ring_size=ring.size, gpu_id=self.gpu_id,
       queue_type=queue_type, queue_percentage=kfd.KFD_MAX_QUEUE_PERCENTAGE|(xcc_id<<8), queue_priority=kfd.KFD_MAX_QUEUE_PRIORITY,
       eop_buffer_address=eop_buffer.va_addr if eop_buffer else 0, eop_buffer_size=eop_buffer.size if eop_buffer else 0, ctl_stack_size=ctl_stack_size,
       ctx_save_restore_address=cwsr_buffer.va_addr if cwsr_buffer else 0, ctx_save_restore_size=ctx_save_restore_size,
-      write_pointer_address=gart.va_addr+wptr, read_pointer_address=gart.va_addr+rptr+8*xcc_id)
+      write_pointer_address=gart.va_addr, read_pointer_address=gart.va_addr + 8 * (xcc_id + 1))
 
     if not hasattr(self, 'doorbells'):
       self.doorbells_base = queue.doorbell_offset & (~0x1fff) # doorbell is two pages
@@ -705,19 +662,18 @@ class PCIIface(PCIIfaceBase):
       'max_slots_scratch_cu': self.dev_impl.gc_info.gc_max_scratch_slots_per_cu, 'max_waves_per_simd': self.dev_impl.gc_info.gc_max_waves_per_simd,
       'simd_arrays_per_engine': self.dev_impl.gc_info.gc_num_sa_per_se, 'lds_size_in_kb': self.dev_impl.gc_info.gc_lds_size}
 
-  def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0, xcc_id=0):
+  def create_queue(self, queue_type, ring, gart, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0, xcc_id=0):
     assert cwsr_buffer is None, "no cwsr buffer for am"
-    assert queue_type != kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL, "no AQL queues for am"
 
     if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA:
-      self.dev_impl.sdma.setup_ring(ring_addr=ring.va_addr, ring_size=ring.size, rptr_addr=gart.va_addr+rptr, wptr_addr=gart.va_addr+wptr,
+      self.dev_impl.sdma.setup_ring(ring_addr=ring.va_addr, ring_size=ring.size, rptr_addr=gart.va_addr, wptr_addr=gart.va_addr+0x10,
                                     doorbell=(doorbell_index:=am.AMDGPU_NAVI10_DOORBELL_sDMA_ENGINE0), pipe=0, queue=0)
     else:
-      self.dev_impl.gfx.setup_ring(ring_addr=ring.va_addr, ring_size=ring.size, rptr_addr=gart.va_addr+rptr, wptr_addr=gart.va_addr+wptr,
+      self.dev_impl.gfx.setup_ring(ring_addr=ring.va_addr, ring_size=ring.size, rptr_addr=gart.va_addr, wptr_addr=gart.va_addr+0x10,
         eop_addr=eop_buffer.va_addr, eop_size=eop_buffer.size, doorbell=(doorbell_index:=am.AMDGPU_NAVI10_DOORBELL_MEC_RING0), pipe=0, queue=0)
 
     return AMDQueueDesc(ring=ring.cpu_view().view(fmt='I'), doorbells=[self.dev_impl.doorbell64.view(doorbell_index * 8, 8, fmt='Q')],
-      read_ptrs=[gart.cpu_view().view(offset=rptr, size=8, fmt='Q')], write_ptrs=[gart.cpu_view().view(offset=wptr, size=8, fmt='Q')])
+      read_ptrs=[gart.cpu_view().view(size=8, fmt='Q')], write_ptrs=[gart.cpu_view().view(offset=0x10, size=8, fmt='Q')])
 
   def sleep(self, timeout):
     if self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))):
@@ -759,9 +715,9 @@ class USBIface(PCIIface):
     return HCQBuffer(am_mapping.va_addr, size, meta=PCIAllocationMeta(am_mapping, has_cpu_mapping=False),
       view=USBMMIOInterface(self.usb, self.bars[0][0] + am_mapping.paddrs[0][0], size, fmt='B') if cpu_access else None, owner=self.dev)
 
-  def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0, xcc_id=0):
+  def create_queue(self, queue_type, ring, gart, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0, xcc_id=0):
     if queue_type == kfd.KFD_IOC_QUEUE_TYPE_COMPUTE: self.usb._pci_cacheable += [(ring.cpu_view().addr, ring.size)]
-    return super().create_queue(queue_type, ring, gart, rptr, wptr, eop_buffer, cwsr_buffer, ctl_stack_size, ctx_save_restore_size, xcc_id)
+    return super().create_queue(queue_type, ring, gart, eop_buffer, cwsr_buffer, ctl_stack_size, ctx_save_restore_size, xcc_id)
 
   def sleep(self, timeout): pass
 
@@ -807,13 +763,7 @@ class AMDDevice(HCQCompiled):
     nbio_pad = (0,) if self.target[0] == 9 else ()
     self.nbio = AMDIP(nbio_name, self.iface.ip_versions[am.NBIF_HWIP], {i:nbio_pad+x for i,x in self.iface.ip_offsets[am.NBIF_HWIP].items()})
 
-    self.is_aql = getenv("AMD_AQL", self.xccs > 1)
-    if self.is_aql:
-      self.pm4_ibs = self.iface.alloc(0x2000 if self.is_usb() else (16 << 20), uncached=True, cpu_access=True)
-      self.pm4_ib_alloc = BumpAllocator(self.pm4_ibs.size, wrap=True)
-
-    self.compute_queue = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL if self.is_aql else kfd.KFD_IOC_QUEUE_TYPE_COMPUTE,
-      0x2000 if self.is_usb() else (16 << 20), eop_buffer_size=0x1000,
+    self.compute_queue = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_COMPUTE, 0x2000 if self.is_usb() else (16 << 20), eop_buffer_size=0x1000,
       ctx_save_restore_size=0 if self.is_am() else wg_data_size + ctl_stack_size, ctl_stack_size=ctl_stack_size, debug_memory_size=debug_memory_size)
 
     max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000
@@ -821,8 +771,7 @@ class AMDDevice(HCQCompiled):
 
     super().__init__(device, AMDAllocator(self), AMDLLVMRenderer(self.arch) if AMD_LLVM else AMDRenderer(self.arch),
                      AMDLLVMCompiler(self.arch) if AMD_LLVM else HIPCompiler(self.arch), functools.partial(AMDProgram, self),
-                     AMDSignal, functools.partial(AMDComputeAQLQueue if self.is_aql else AMDComputeQueue, self),
-                     functools.partial(AMDCopyQueue, self, max_copy_size=max_copy_size),
+                     AMDSignal, functools.partial(AMDComputeQueue, self), functools.partial(AMDCopyQueue, self, max_copy_size=max_copy_size),
                      kernargs_size=(8 << 10) if self.is_usb() else (16 << 20), sigalloc_size=0x100 if self.is_usb() else 0x1000)
 
     # Scratch setup
@@ -831,10 +780,10 @@ class AMDDevice(HCQCompiled):
 
     # XCC setup
     self.xcc_sync: tuple[AMDSignal, AMDSignal]|None = None
-    if self.xccs > 1 and not self.is_aql:
+    if self.xccs > 1:
       self.xcc_sync_area = self.allocator.alloc(0x1000, BufferSpec(nolru=True, cpu_access=True))
       self.xcc_sync = (AMDSignal(base_buf=self.xcc_sync_area), AMDSignal(base_buf=self.xcc_sync_area.offset(256)))
-      cast(AMDComputeQueue, self.hw_compute_queue_t()).xcc_config().submit(self)
+      AMDComputeQueue(self).xcc_config().submit(self)
 
     # SQTT is disabled by default because of runtime overhead and big file sizes (~200mb to Tensor.full() two 4096x4096 tensors and matmul them)
     self.sqtt_enabled = PROFILE and bool(getenv("SQTT", 0))
@@ -849,26 +798,18 @@ class AMDDevice(HCQCompiled):
       self.sqtt_buffers = [self.allocator.alloc(SQTT_BUFFER_SIZE*1024*1024, BufferSpec(cpu_access=True, nolru=True)) for _ in range(SQTT_NUM)]
       self.sqtt_itrace_se_mask = getenv("SQTT_ITRACE_SE_MASK", 2) # -1 enable all, 0 disable all, >0 bitmask for where to enable instruction tracing
       self.cmd_id = 0
-      cast(AMDComputeQueue, self.hw_compute_queue_t()).sqtt_start(self.sqtt_buffers, self.sqtt_itrace_se_mask).submit(self)
+      AMDComputeQueue(self).sqtt_start(self.sqtt_buffers, self.sqtt_itrace_se_mask).submit(self)
 
   def create_queue(self, queue_type, ring_size, ctx_save_restore_size=0, eop_buffer_size=0, ctl_stack_size=0, debug_memory_size=0):
     ring = self.iface.alloc(ring_size, uncached=True, cpu_access=True)
     gart = self.iface.alloc(0x100, uncached=True, cpu_access=True)
 
-    if queue_type == kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL:
-      aql_desc = hsa.amd_queue_t(queue_properties=hsa.AMD_QUEUE_PROPERTIES_IS_PTR64 | hsa.AMD_QUEUE_PROPERTIES_ENABLE_PROFILING,
-        read_dispatch_id_field_base_byte_offset=getattr(hsa.amd_queue_t, 'read_dispatch_id').offset,
-        max_cu_id=self.max_cu_id, max_wave_id=self.max_wave_id)
-      gart.cpu_view().view(fmt='B')[:ctypes.sizeof(aql_desc)] = bytes(aql_desc)
-      self.aql_desc = hsa.amd_queue_t.from_address(gart.va_addr)
-
     cwsr_buffer_size = round_up((ctx_save_restore_size + debug_memory_size) * self.iface.props.get('num_xcc', 1), mmap.PAGESIZE)
     cwsr_buffer = self.iface.alloc(cwsr_buffer_size) if ctx_save_restore_size else None
     eop_buffer = self.iface.alloc(eop_buffer_size) if eop_buffer_size else None
 
-    return AMDQueueDesc.multi(*(self.iface.create_queue(queue_type, ring, gart, rptr=getattr(hsa.amd_queue_t, 'read_dispatch_id').offset,
-                                wptr=getattr(hsa.amd_queue_t, 'write_dispatch_id').offset, eop_buffer=eop_buffer, cwsr_buffer=cwsr_buffer,
-                                xcc_id=xcc_id, ctx_save_restore_size=ctx_save_restore_size, ctl_stack_size=ctl_stack_size)
+    return AMDQueueDesc.multi(*(self.iface.create_queue(queue_type, ring, gart, eop_buffer=eop_buffer, cwsr_buffer=cwsr_buffer, xcc_id=xcc_id,
+                                                            ctx_save_restore_size=ctx_save_restore_size, ctl_stack_size=ctl_stack_size)
                                 for xcc_id in range(self.xccs if queue_type == kfd.KFD_IOC_QUEUE_TYPE_COMPUTE else 1)))
 
   def _ensure_has_local_memory(self, required):
@@ -887,16 +828,8 @@ class AMDDevice(HCQCompiled):
       self.tmpring_size = waves << 12 | wavesize
       self.max_private_segment_size = required
 
-      if hasattr(self, 'aql_desc'):
-        self.aql_desc.scratch_backing_memory_location = self.scratch.va_addr
-        self.aql_desc.scratch_backing_memory_byte_size = self.scratch.size
-        self.aql_desc.scratch_wave64_lane_byte_size = self.max_private_segment_size * (self.aql_desc.max_wave_id + 1) // 64
-        self.aql_desc.scratch_resource_descriptor[:] = [lo32(self.scratch.va_addr), hi32(self.scratch.va_addr) | (1 << 30), lo32(self.scratch.size),
-          0x20814fac] # FORMAT=BUF_FORMAT_32_UINT,OOB_SELECT=2,ADD_TID_ENABLE=1,TYPE=SQ_RSRC_BUF,SQ_SELs
-        self.aql_desc.compute_tmpring_size = self.tmpring_size
-
   def invalidate_caches(self):
-    self.hw_compute_queue_t().memory_barrier().signal(self.timeline_signal, self.next_timeline()).submit(self)
+    AMDComputeQueue(self).memory_barrier().signal(self.timeline_signal, self.next_timeline()).submit(self)
     self.synchronize()
 
   def on_device_hang(self): self.iface.on_device_hang()
@@ -905,8 +838,7 @@ class AMDDevice(HCQCompiled):
     if self.sqtt_enabled:
       wptrs_buf = self.allocator.alloc(round_up(len(self.sqtt_buffers), 0x1000), BufferSpec(cpu_access=True, nolru=True))
       wptrs = to_mv(wptrs_buf.va_addr, wptrs_buf.size)
-      cast(AMDComputeQueue, self.hw_compute_queue_t()).sqtt_stop(len(self.sqtt_buffers), wptrs_buf) \
-                                                      .signal(self.timeline_signal, self.next_timeline()).submit(self)
+      AMDComputeQueue(self).sqtt_stop(len(self.sqtt_buffers), wptrs_buf).signal(self.timeline_signal, self.next_timeline()).submit(self)
       self.synchronize()
       if DEBUG>=2: print('Saving SQTT in profile...')
       for i,buf0 in enumerate(self.sqtt_buffers):

tinygrad/runtime/ops_null.py

@@ -7,7 +7,6 @@ class NullRenderer(CStyleLanguage):
   device = "NULL"
   has_local = False
   float4 = "float4"
-  barrier = "// BARRIER"
   code_for_op = {**CStyleLanguage.code_for_op, Ops.THREEFRY: lambda a,b,dtype: f"threefry({a},{b})", Ops.MAX: lambda a,b,dtype: f"max({a},{b})"}
 
 class NullProgram:

tinygrad/runtime/ops_remote.py

@@ -306,10 +306,7 @@ class RemoteHandler:
           case ProgramAlloc():
             lib = dev.compiler.compile_cached(req._h[c.datahash].decode())
             session.programs[(c.name, c.datahash)] = dev.runtime(c.name, lib)
-          case ProgramFree():
-            key = (c.name, c.datahash)
-            # WORKAROUND: should be unconditional once the protocol supports proper exception handling
-            if key in session.programs: del session.programs[key]
+          case ProgramFree(): del session.programs[(c.name, c.datahash)]
           case ProgramExec():
             bufs = [session.buffers[x]._buf for x in c.bufs]
             extra_args = {k:v for k,v in [("global_size", c.global_size), ("local_size", c.local_size)] if v is not None}
@@ -424,24 +421,19 @@ class RemoteConnection:
     conns = RemoteConnection.all.keys()
     datas = {conn: conn.req.serialize() for conn in conns}
     reqs, hashes, hash_datas = sum(len(c.req._q) for c in conns), sum(len(c.req._h) for c in conns), sum(len(data) for data in datas.values())
-    resps = []
     with Timing(f"*** send {reqs:-3d} requests {hashes:-3d} hashes with len {hash_datas/1024:.2f} kB in ", enabled=DEBUG>=3):
       for conn,data in datas.items(): conn.conn.request("POST", "/batch", data)
       for conn in datas.keys():
-        resp = conn.conn.getresponse()
-        body = resp.read()
-        resps.append((conn, resp, body))
-        conn.req = BatchRequest()
-    if take_q: RemoteConnection.q_lock.release()
-    for conn,resp,body in resps:
-      match resp.status:
-        case http.HTTPStatus.OK: pass
-        case http.HTTPStatus.INTERNAL_SERVER_ERROR:
-          exc_wrapper = safe_eval(ast.parse(body.decode(), mode="eval").body)
+        response = conn.conn.getresponse()
+        resp = response.read()
+        conn.req = BatchRequest() # no matter what response, reset conn
+        if response.status == http.HTTPStatus.INTERNAL_SERVER_ERROR:
+          exc_wrapper = safe_eval(ast.parse(resp.decode(), mode="eval").body)
           exc_wrapper.exc.add_note(exc_wrapper.trace)
           raise exc_wrapper.exc
-        case code: raise RuntimeError(f"POST /batch failed with {code}: {body.decode()}")
-      if conn == self: ret = body
+        assert response.status == http.HTTPStatus.OK, f"POST /batch failed: {resp.decode()}"
+        if conn == self: ret = resp
+    if take_q: RemoteConnection.q_lock.release()
     return ret
 
 def parse_hosts(hs:str) -> list[tuple[str, int]]|LazySeq[tuple[str, int]]:

tinygrad/runtime/support/am/amdev.py

@@ -104,7 +104,7 @@ class AMPageTableEntry:
   def entry(self, entry_id:int) -> int: return self.entries[entry_id]
   def valid(self, entry_id:int) -> bool: return (self.entries[entry_id] & am.AMDGPU_PTE_VALID) != 0
   def address(self, entry_id:int) -> int: return self.entries[entry_id] & 0x0000FFFFFFFFF000
-  def is_page(self, entry_id:int) -> bool: return self.lv == am.AMDGPU_VM_PTB or self.adev.gmc.is_pte_huge_page(self.entries[entry_id])
+  def is_huge_page(self, entry_id:int) -> bool: return self.lv == am.AMDGPU_VM_PTB or self.adev.gmc.is_pte_huge_page(self.entries[entry_id])
   def supports_huge_page(self, paddr:int): return self.lv >= am.AMDGPU_VM_PDB2
 
 class AMMemoryManager(MemoryManager):
@@ -239,7 +239,7 @@ class AMDev(PCIDevImplBase):
       ip_offset = ctypes.addressof(self.bhdr) + ctypes.sizeof(dhdr) + ihdr.die_info[num_die].die_offset
       for _ in range(dhdr.num_ips):
         ip = am.struct_ip_v4.from_address(ip_offset)
-        ba = ((ctypes.c_uint64 if ihdr.base_addr_64_bit else ctypes.c_uint32) * ip.num_base_address).from_address(ip_offset + 8)
+        ba = (ctypes.c_uint32 * ip.num_base_address).from_address(ip_offset + 8)
         for hw_ip in range(1, am.MAX_HWIP):
           if hw_ip in hw_id_map and hw_id_map[hw_ip] == ip.hw_id:
             self.regs_offset[hw_ip][ip.instance_number] = tuple(list(ba))

tinygrad/runtime/support/hcq.py

@@ -438,13 +438,12 @@ class HCQCompiled(Compiled, Generic[SignalType]):
     return buf, realloced
 
   def _select_iface(self, *ifaces:Type):
-    errs, err_short = "", ""
+    errs:str = ""
     if val:=getenv(f'{type(self).__name__[:-6].upper()}_IFACE', ""): ifaces = tuple(x for x in ifaces if x.__name__.startswith(val.upper()))
     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 RuntimeError(f"{errs}\nNo interface for {type(self).__name__[:-6]}:{self.device_id} is available:{err_short}\n" \
-                       f"\nForce an interface with {type(self).__name__[:-6].upper()}_IFACE={('|'.join(x.__name__[:-5] for x in ifaces))}.")
+      except Exception: errs += f"\n{iface_t.__name__}: {traceback.format_exc()}"
+    raise RuntimeError(f"Cannot find a usable interface for {type(self).__name__[:-6]}:{self.device_id}:\n{errs}")
 
   def _is_cpu(self) -> bool: return hasattr(self, 'device') and self.device.split(":")[0] in ("CPU", "LLVM")
 

tinygrad/runtime/support/memory.py

@@ -77,10 +77,11 @@ class TLSFAllocator:
       if self.lv1_entries[l1] == 0: continue
       for l2 in range(self.lv2(size) if l1 == size.bit_length() else 0, (1 << self.l2_cnt)):
         if len(self.storage[l1][l2]) > 0:
+          nsize = self.blocks[self.storage[l1][l2][0]][0]
+          assert nsize >= size, "block must be larger"
+
           # Block start address.
           start = self.storage[l1][l2][0]
-          nsize = self.blocks[start][0]
-          assert nsize >= size, "block must be larger"
 
           # If request contains alignment, split the block into two parts.
           if (new_start:=round_up(start, align)) != start:
@@ -117,7 +118,7 @@ class PageTableTraverseContext:
       assert self.create_pts, "Not allowed to create new page table"
       pt.set_entry(pte_idx, self.dev.mm.palloc(0x1000, zero=True, boot=self.boot), table=True, valid=True)
 
-    assert not pt.is_page(pte_idx), f"Must be table pt={pt.paddr:#x}, {pt.lv=} {pte_idx=} {pt.read_fields(pte_idx)}"
+    assert not pt.is_huge_page(pte_idx), f"Must be table pt={pt.paddr:#x}, {pt.lv=} {pte_idx=} {pt.read_fields(pte_idx)}"
     child_page_table = self.dev.mm.pt_t(self.dev, pt.address(pte_idx), lv=pt.lv+1)
 
     self.pt_stack.append((child_page_table, self._pt_pte_idx(child_page_table, self.vaddr), self._pt_pte_size(child_page_table)))
@@ -144,7 +145,7 @@ class PageTableTraverseContext:
         assert paddr is not None, "paddr must be provided when allocating new page tables"
         while pte_covers > size or not pt.supports_huge_page(paddr+off) or self.vaddr&(pte_covers-1) != 0: pt, pte_idx, pte_covers = self.level_down()
       else:
-        while not pt.is_page(pte_idx): pt, pte_idx, pte_covers = self.level_down()
+        while not pt.is_huge_page(pte_idx): pt, pte_idx, pte_covers = self.level_down()
 
       entries = min(size // pte_covers, self._pt_pte_cnt(pt.lv) - pte_idx)
       assert entries > 0, f"Invalid entries {size=:#x}, {pte_covers=:#x}"

tinygrad/runtime/support/nv/nvdev.py

@@ -51,14 +51,14 @@ class NVPageTableEntry:
     return (self.entries[2*entry_id+1]<<64) | self.entries[2*entry_id] if self._is_dual_pde() else self.entries[entry_id]
 
   def read_fields(self, entry_id:int) -> dict:
-    if self.is_page(entry_id): return self.nvdev.pte_t.decode(self.entry(entry_id))
+    if self.is_huge_page(entry_id): return self.nvdev.pte_t.decode(self.entry(entry_id))
     return (self.nvdev.dual_pde_t if self._is_dual_pde() else self.nvdev.pde_t).decode(self.entry(entry_id))
 
-  def is_page(self, entry_id) -> bool: return (self.entry(entry_id) & 1 == 1) if self.lv < self.nvdev.mm.level_cnt - 1 else True
+  def is_huge_page(self, entry_id) -> bool: return (self.entry(entry_id) & 1 == 1) if self.lv < self.nvdev.mm.level_cnt - 1 else True
   def supports_huge_page(self, paddr:int): return self.lv >= self.nvdev.mm.level_cnt - 3 and paddr % self.nvdev.mm.pte_covers[self.lv] == 0
 
   def valid(self, entry_id):
-    if self.is_page(entry_id): return self.read_fields(entry_id)['valid']
+    if self.is_huge_page(entry_id): return self.read_fields(entry_id)['valid']
     return self.read_fields(entry_id)['aperture_small' if self._is_dual_pde() else 'aperture'] != 0
 
   def address(self, entry_id:int) -> int:

tinygrad/runtime/support/system.py

@@ -1,4 +1,4 @@
-import os, mmap, array, functools, ctypes, select, contextlib, dataclasses, sys, errno
+import os, mmap, array, functools, ctypes, select, contextlib, dataclasses, sys
 from typing import cast, ClassVar
 from tinygrad.helpers import round_up, to_mv, getenv, OSX, temp
 from tinygrad.runtime.autogen import libc, vfio
@@ -84,11 +84,7 @@ class PCIDevice:
     for i in resize_bars or []:
       if FileIOInterface.exists(rpath:=f"/sys/bus/pci/devices/{self.pcibus}/resource{i}_resize"):
         try: FileIOInterface(rpath, os.O_RDWR).write(str(int(FileIOInterface(rpath, os.O_RDONLY).read(), 16).bit_length() - 1))
-        except OSError as e:
-          if e.errno in {errno.EPERM, errno.EACCES}:
-            raise RuntimeError(f"Cannot resize BAR {i}: {e}. Permission error: run `extra/amdpci/setup_python_cap.sh`"
-                                " to allow python accessing device or run with sudo") from e
-          raise RuntimeError(f"Cannot resize BAR {i}: {e}. Ensure the resizable BAR option is enabled on your system.") from e
+        except OSError as e: raise RuntimeError(f"Cannot resize BAR {i}: {e}. Ensure the resizable BAR option is enabled on your system.") from e
 
     if getenv("VFIO", 0) and (vfio_fd:=System.vfio()) is not None:
       FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/driver_override", os.O_WRONLY).write("vfio-pci")

tinygrad/schedule/kernelize.py

@@ -91,8 +91,8 @@ kernelize_sym = symbolic_simple+PatternMatcher([
   (UPat(Ops.CAST, name="cast", src=(UPat(Ops.VIEW, name="vm"),)), lambda cast,vm: vm.base.cast(cast.dtype).view(vm.st)
     if cast.dtype.itemsize <= vm.dtype.itemsize and resolve(prod(vm.shape) > vm.st.real_size()) else None),
   # put UnaryOps before EXPANDs, if it can fuse with the input
-  (UPat(GroupOp.Unary, src=(UPat(Ops.VIEW, src=(UPat(GroupOp.All-ALWAYS_CONTIGUOUS, name="inp"),), name="v"),), name="alu"),
-   lambda inp,v,alu: inp.alu(alu.op).view(v.st) if resolve(prod(alu.shape) > v.st.real_size()) else None),
+  #(UPat(GroupOp.Unary, src=(UPat(Ops.VIEW, src=(UPat(GroupOp.All-ALWAYS_CONTIGUOUS, name="inp"),), name="v"),), name="alu"),
+  # lambda inp,v,alu: inp.alu(alu.op).view(v.st) if resolve(prod(alu.shape) > v.st.real_size()) else None),
 ])
 
 # support for using a contiguous permuted view instead of the parent view if one exists

tinygrad/schedule/rangeify.py

@@ -1,464 +0,0 @@
-from typing import Any
-from dataclasses import dataclass, field
-from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, RewriteNotReady, _substitute
-from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, RANGEIFY
-from tinygrad.schedule.multi import multi_pm
-
-from tinygrad.schedule.kernelize import Kernel
-from tinygrad.uop.ops import track_rewrites, graph_rewrite_map, graph_rewrite, identity_element, sint, AxisType
-
-# 0. do some cleanup rewrites, mostly copied from the old stuff
-
-double_reshape = PatternMatcher([
-  # RESHAPE on RESHAPE is the second reshape
-  (UPat(Ops.RESHAPE, src=(UPat(Ops.RESHAPE),), name="x"), lambda x: x.replace(src=(x.src[0].src[0],))),
-])
-
-earliest_rewrites = double_reshape+PatternMatcher([
-  # UOp with size 0 is zero
-  (UPat(GroupOp.All-{Ops.SINK}, name="root"), lambda root: root.const_like(0) if root.base.st is not None and root.size == 0 else None),
-  # DETACH and CONTIGUOUS_BACKWARD are NOOPs here, so is FUSE
-  (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD, Ops.FUSE), name="x"), lambda x: x.src[0]),
-  # reduce of size 0 is the identity element
-  (UPat(Ops.REDUCE_AXIS, name="reduce", src=(UPat.var("x"),)),
-   lambda reduce,x: reduce.const_like(identity_element(reduce.arg[0], reduce.dtype)) if x.size == 0 and reduce.size != 0 else None),
-  # non shape changing RESHAPE is NOOP
-  (UPat(Ops.RESHAPE, name="x"), lambda x: x.src[0] if x.src[0].shape == x.arg else None),
-  # RESHAPE after COPY
-  (UPat(Ops.COPY, src=(UPat(Ops.RESHAPE, name="r"),UPat(name="d")), name="c"), lambda c,r,d: c.replace(src=(r.src[0],d)).reshape(r.arg)),
-  # TODO: this should be BUFFER_VIEW
-  (UPat(Ops.COPY, src=(UPat(Ops.SHRINK, name="r"),UPat(name="d")), name="c"), lambda c,r,d: c.replace(src=(r.src[0],d)).shrink(r.arg)),
-  # const hacks
-  (UPat(Ops.CONST, name="x"), lambda x:
-   x.replace(src=(x.src[0].src[0],)).reshape((1,)*len(x.shape)).expand(x.shape) if \
-    len(x.src) and x.src[0].op is Ops.VIEW and not any(s == 0 for s in x.shape) else None),
-  # assign only to buffer
-  (UPat(Ops.ASSIGN, src=(UPat(GroupOp.All-{Ops.BUFFER}, name="target"), UPat(name="x"))),
-   lambda x,target: x if target.base.op is not Ops.BUFFER else None),
-  # contiguous/buffer/copy/assign is already contiguous
-  (UPat(Ops.CONTIGUOUS, name="root", src=(UPat((Ops.CONTIGUOUS, Ops.BUFFER, Ops.COPY, Ops.ASSIGN)),)), lambda root: root.src[0]),
-])
-
-# 1. add contiguous where we have to
-
-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.DEFINE_GLOBAL,
-                     Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.LOAD}
-
-def realize(ctx:dict[UOp, None], tr:UOp) -> None: ctx[tr] = None
-
-def realize_parents(ctx:dict[UOp, None], rb:UOp) -> None:
-  for s in rb.src:
-    if s.op not in ALWAYS_CONTIGUOUS: ctx[s] = 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
-
-do_realize = PatternMatcher([
-  # always realize SINK parents
-  (UPat(Ops.SINK, name="s"), lambda ctx,s: ctx.update((x.base, None) for x in s.src if x.base.op not in ALWAYS_CONTIGUOUS)),
-  # always realize ASSIGN/COPY/BUFFER_VIEW
-  (UPat({Ops.ASSIGN, Ops.COPY, Ops.BUFFER_VIEW}, name="tr"), realize),
-  # realize parents of COPY, MSELECT, MSTACK
-  (UPat((Ops.COPY, Ops.MSELECT, Ops.MSTACK), name="rb"), realize_parents),
-  # realize input to assign (might be optimized out)
-  (UPat(Ops.ASSIGN, name="a"), realize_assign),
-])
-
-add_contiguous = PatternMatcher([
-  (UPat(GroupOp.All-{Ops.CONTIGUOUS}, name="x"), lambda ctx,x: x.replace(tag=1).contiguous() if x in ctx and x.tag is None else None),
-])
-remove_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])
-
-# 2. mark all children
-
-@dataclass
-class ChildrenContext: children: dict[UOp, list[UOp]]|None = None
-def extract_children(ctx:ChildrenContext, x:UOp):
-  if ctx.children is not None: return
-  children_map = x.get_children_map()
-  ctx.children = {}
-  for k,v in children_map.items():
-    non_sink_children = [u for u in v if u.op is not Ops.SINK]
-    if len(non_sink_children) <= 1: continue
-    # NOTE: this gate shouldn't be here
-    if any(x.op is Ops.REDUCE_AXIS for x in k.toposort()) and any(x.op in {Ops.BUFFER, Ops.CONTIGUOUS} for x in k.toposort()):
-      ctx.children[k] = non_sink_children
-
-def mark_children(ctx:ChildrenContext, x:UOp):
-  assert ctx.children is not None
-  new_srcs = [(UOp(Ops.CHILD, s.dtype, src=(UOp(Ops.CHILDREN, s.dtype, (s,), arg=len(ctx.children[s])),),
-                   arg=(ctx.children[s].index(x), len(ctx.children[s]))) if s in ctx.children else s) for s in x.src]
-  return x.replace(src=tuple(new_srcs))
-
-pm_children = PatternMatcher([
-  (UPat(Ops.SINK, name="x"), extract_children),
-  (UPat(GroupOp.All-{Ops.CHILD, Ops.CHILDREN}, name="x"), mark_children),
-])
-
-# 3. rangeify
-
-@dataclass
-class RangeifyContext:
-  # block on parent until all children have been seen
-  seen_children: dict[UOp, dict[int, UOp]] = field(default_factory=dict)
-  seen_child: dict[UOp, Any] = field(default_factory=dict)
-  progress: int = 0
-
-  # create ranges
-  range_idx: int = 0
-  def new_range(self, s:sint, axistype:AxisType=AxisType.LOOP):
-    ret = UOp.range(dtypes.int, s, self.range_idx, axistype)
-    self.range_idx += 1
-    return ret
-
-def map_reshape(idx:UOp, r:UOp):
-  acc = 1
-  to_sum = []
-  for s,src in list(zip(idx.shape, idx.src[1:]))[::-1]:
-    to_sum.append(acc*src)
-    acc *= s
-  mish = sum(to_sum, start=UOp.const(dtypes.int, 0))
-  ret:list[UOp] = []
-  for s in r.src[0].shape[::-1]:
-    ret.append(mish % s) # NOTE: simplify will turn this to CONST
-    mish //= s
-  tret = ret[0].sink(*ret[1:]).simplify().src[::-1] if len(ret) else ()
-  return r.src[0].index(*tret, dtype=idx.dtype, arg=idx.arg)
-
-def map_pad(idx:UOp, r:UOp):
-  ret = list(idx.src[1:])
-  bigwhere = UOp.const(dtypes.bool, True)
-  for i,(sh,(s,e)) in enumerate(zip(r.shape, r.arg)):
-    if s == 0 and e == 0: continue
-    where = UOp.const(dtypes.bool, True)
-    if resolve(e > 0): where = where & (ret[i] < (sh-e))
-    if resolve(s > 0): where = where & (ret[i] >= s)
-    bigwhere = bigwhere & where
-    # this is safe but dumb
-    # TODO (S-Lykles): switch to mixed index/valid
-    ret[i] = (ret[i] - s).maximum(0).minimum(r.src[0].shape[i]-1)
-  # PAD is with 0
-  return bigwhere.simplify().where(r.src[0].index(*ret, dtype=idx.dtype, arg=idx.arg), UOp.const(r.dtype, 0))
-
-def map_expand(r:UOp, idx:UOp):
-  new_rngs = []
-  ending_ranges = []
-  non_ending_ranges = []
-  for a,x,y in zip(idx.src[1:], r.src[0].shape, r.shape):
-    axis_to_range = [u for u in a.toposort() if u.op is Ops.RANGE]
-    if resolve(x!=y, False):
-      ending_ranges.extend(axis_to_range)
-      new_rngs.append(a.const_like(0))
-    else:
-      non_ending_ranges.extend(axis_to_range)
-      new_rngs.append(a)
-  ending_ranges = [x.arg for x in ending_ranges if x not in non_ending_ranges]
-  if idx.arg is not None: ending_ranges.append(idx.arg)
-  return r.src[0].index(*new_rngs, arg=min(ending_ranges) if ending_ranges else None)
-
-pm_mops = PatternMatcher([
-  # this is like the definitions of these
-  (UPat(Ops.SHRINK, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"),
-   lambda r,idx: r.src[0].index(*[a+ss if resolve(ss != 0) else a for a,(ss,_) in zip(idx.src[1:], r.arg)], dtype=idx.dtype, arg=idx.arg)),
-  (UPat(Ops.PERMUTE, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"),
-   lambda r,idx: r.src[0].index(*[idx.src[1+p] for p in argsort(idx.src[0].arg)], dtype=idx.dtype, arg=idx.arg)),
-  (UPat(Ops.FLIP, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"),
-   lambda r,idx: r.src[0].index(*[((s-1)-a) if f else a for a,s,f in zip(idx.src[1:], r.shape, r.arg)], dtype=idx.dtype, arg=idx.arg)),
-  # expand needs to end ranges
-  (UPat(Ops.EXPAND, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"), map_expand),
-  # reshape does a lot of symbolic stuff
-  (UPat(Ops.RESHAPE, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"), map_reshape),
-  # pad adds min and max
-  (UPat(Ops.PAD, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"), map_pad),
-])
-
-def map_partial_contiguous(ctx:RangeifyContext, x:UOp, idx:UOp):
-  if x.arg is None: return None  # map_contiguous can handle this
-  # NOTE: all partial contiguous can safely be replaced by full contiguous. we should be able to match old functionality like this
-  if not (RANGEIFY > 1): return idx.replace(src=(x.replace(arg=None),)+idx.src[1:])
-  ranges = []
-  new_ranges = []
-  passthrough_idx = []
-  for i,s in enumerate(x.shape):
-    if i not in x.arg:
-      ranges.append(idx.src[1+i])
-      continue
-    passthrough_idx.append(idx.src[1+i])
-    ranges.append(ctx.new_range(s) if resolve(s!=1) else UOp.const(dtypes.int, 0))
-    new_ranges.append(ranges[-1])
-  ret = x.src[0].index(*ranges).bufferize(*[x for x in new_ranges if x.op is not Ops.CONST], arg=x.device)
-  if len(ret.ranges): ret = ret.replace(arg=AddrSpace.LOCAL) # if some ranges are still open, this has to be LOCAL
-  return ret.index(*passthrough_idx)
-
-def map_contiguous(ctx:RangeifyContext, x:UOp):
-  if x.arg is not None: return None
-  ranges = []
-  for s in x.shape[len(x.src)-1:]:
-    ranges.append(ctx.new_range(s) if resolve(s!=1) else UOp.const(dtypes.int, 0))
-  return x.src[0].index(*ranges).bufferize(*x.src[1:], *[x for x in ranges if x.op is not Ops.CONST], arg=x.device).forced_reshape(x.shape)
-
-def map_reduce(ctx:RangeifyContext, idx:UOp, red:UOp):
-  rngs = list(idx.src[1:])
-  new_ranges = []
-  for i,s in enumerate(red.src[0].shape):
-    if i in red.arg[1]:
-      rngs[i] = ctx.new_range(s, axistype=AxisType.REDUCE)
-      new_ranges.append(rngs[i])
-  return UOp(Ops.REDUCE, red.dtype, src=(red.src[0].index(*rngs),)+tuple(new_ranges), arg=red.arg[0])
-
-def index_child(ctx:RangeifyContext, c:UOp, x:UOp, idx:UOp):
-  if c not in ctx.seen_children: ctx.seen_children[c] = {}
-  # wait here until we have seen all the children
-  if len(ctx.seen_children[c]) != x.arg[1]:
-    ctx.progress += 1
-    if ctx.progress > 10000: raise RuntimeError("children not making progress")
-    # NOTE: we mark this here
-    ctx.seen_children[c][x.arg[0]] = idx
-    raise RewriteNotReady
-  ctx.progress = 0
-
-  if c not in ctx.seen_child:
-    all_rngs = zip(*[ch.src[1:] for ch in ctx.seen_children[c].values()])
-    out_rngs = []
-    end_ranges = []
-    idx_ranges = []
-    for i,r in enumerate(all_rngs):
-      if all_same(r):
-        out_rngs.append(r[0])
-      else:
-        out_rngs.append(ctx.new_range(c.shape[i]))
-        end_ranges.append(out_rngs[-1])
-        idx_ranges.append(i)
-    ctx.seen_child[c] = (idx_ranges, end_ranges)
-  else:
-    out_rngs = list(idx.src[1:])
-    idx_ranges, end_ranges = ctx.seen_child[c]
-    for i,nr in zip(idx_ranges, end_ranges): out_rngs[i] = nr
-  # index based on the shared ranges
-  ret = c.index(*out_rngs)
-  # if all ranges aren't the same between children, we have to bufferize
-  if len(idx_ranges) > 0:
-    ret = ret.bufferize(*end_ranges, arg=x.device)
-    if len(ret.ranges): ret = ret.replace(arg=AddrSpace.LOCAL) # if some ranges are still open, this has to be LOCAL
-    ret = ret.index(*[idx.src[1+i] for i in idx_ranges])
-  return ret
-
-def children_gate(ctx:RangeifyContext, idx:UOp, c:UOp):
-  if len(ctx.seen_children[c]) != c.arg: raise RuntimeError("all children should have been seen by now")
-  return idx.replace(src=(idx.src[0].src[0],)+idx.src[1:])
-
-def might_end_axis(idx:UOp):
-  if idx.arg is None: return None
-  # TODO: write a proper cost function here
-  if all(x.op not in {Ops.BUFFER, Ops.CONTIGUOUS, Ops.BUFFERIZE} for x in idx.toposort()): return None
-  if all(x.op not in {Ops.REDUCE_AXIS} for x in idx.toposort()): return None
-  to_end_axis = []
-  for i,a in enumerate(idx.src[1:]):
-    if any(x.arg > idx.arg for x in a.toposort() if x.op is Ops.RANGE):
-      to_end_axis.append(i)
-  if to_end_axis: return idx.replace(src=(idx.src[0].contiguous(arg=tuple(to_end_axis)),)+idx.src[1:], arg=None)
-  return idx.replace(arg=None)
-
-pm_rangeify = pm_mops+PatternMatcher([
-  # sink contigs to kick it off
-  (UPat(Ops.CONTIGUOUS, src=(UPat(),), name="x", allow_any_len=True), map_contiguous),
-  # if there's an INDEX it can support partial contig
-  (UPat(Ops.INDEX, src=(UPat(Ops.CONTIGUOUS, src=(UPat(),), name="x"),), allow_any_len=True, name="idx"), map_partial_contiguous),
-
-  # if there are new ended children, tag the SINK
-  (UPat(Ops.INDEX, src=(UPat(Ops.CHILD, src=(UPat(name="c"), ), name="x"),), allow_any_len=True, name="idx"), index_child),
-  (UPat(Ops.INDEX, src=(UPat(Ops.CHILDREN, name="c"),), allow_any_len=True, name="idx"), children_gate),
-
-  # if we come across this, remove it. it was a CHILD unused in an INDEX
-  (UPat(Ops.CHILD, src=(UPat(Ops.CHILDREN, src=(UPat.var("x"),)),)), lambda x: x),
-
-  # CONST (or DEFINE_VAR) can't have axes. remove srcs when we INDEX it
-  (UPat(Ops.INDEX, src=(UPat((Ops.CONST, Ops.DEFINE_VAR), name="c"),)), lambda c: c.replace(src=())),
-
-  # handle arg on any op with weight. old endrange stuff
-  (UPat(Ops.INDEX, src=(UPat(GroupOp.Elementwise.union({Ops.REDUCE_AXIS})),), allow_any_len=True, name="idx"), might_end_axis),
-
-  # move MAP through elementwise ALU / reduce. these are the items with cost
-  (UPat(Ops.INDEX, src=(UPat(GroupOp.Elementwise.union({Ops.STORE, Ops.ASSIGN, Ops.COPY, Ops.DEVICE, Ops.BIND})),), allow_any_len=True, name="x"),
-   lambda x: x.src[0].replace(src=tuple([s.index(*x.src[1:]) for s in x.src[0].src]))),
-  (UPat(Ops.INDEX, src=(UPat(Ops.REDUCE_AXIS, name="red"),), allow_any_len=True, name="idx"), map_reduce),
-])
-
-# 3.5 cleanups
-
-# you don't know in the first pass if axes are going to die, this happens if there's an EXPAND to the left
-# TODO: figure out how to reenable this
-def cleanup_dead_axes(b:UOp):
-  parents = b.src[0].toposort()
-  new_rng = []
-  hit = False
-  reshape: list[sint] = []
-  for s,rng in zip(b.shape, b.src[1:]):
-    if rng not in parents and rng.op is Ops.RANGE:
-      reshape.append(1)
-      hit = True
-    else:
-      reshape.append(s)
-      new_rng.append(rng)
-  if hit:
-    return b.replace(src=b.src[0:1]+tuple(new_rng)).reshape(tuple(reshape)).expand(b.shape)
-
-# if a buffer is being stored just for permutes or something, remove it
-# we want to reexpress the indexes of idx2 in terms of the implied b1
-def remove_bufferize(b2:UOp, idx2:UOp):
-  # HACK
-  if len(b2.src) != len(idx2.src): return None
-  assert len(b2.src) == len(idx2.src)
-  assert all(x.op is Ops.RANGE for x in b2.src[1:])
-  return b2.src[0].substitute(dict(zip(b2.src[1:], idx2.src[1:])))
-
-pm_cleanups = double_reshape+pm_mops+PatternMatcher([
-  #(UPat(Ops.BUFFERIZE, name="b"), cleanup_dead_axes),
-  # remove noop buffers. if we look at the next index we can remove even more of these
-  # NOTE: this is mostly the same case as below, but if there's no INDEX this gets more
-  #(UPat(Ops.INDEX, name="idx").f(Ops.BUFFERIZE, allow_any_len=True, name="b2"),
-  # lambda idx,b2: idx.src[0] if idx.src[1:] == b2.src[1:] else None),
-  # remove reindexing
-  (UPat(Ops.INDEX).f(Ops.BUFFERIZE, allow_any_len=True, name="b2").f(Ops.INDEX, allow_any_len=True, name="idx2"), remove_bufferize),
-  # no buffers for const
-  #(UPat(Ops.CONST, name='c').f(Ops.BUFFERIZE, allow_any_len=True, name="b"), lambda c,b: c.reshape((1,)*len(b.shape)).expand(b.shape)),
-])
-
-# 4. put in buffers for bufferize
-# TODO: should BUFFERIZE look a lot more like STORE
-# BUFFERIZE has device in arg
-# BUFFERIZE doesn't have indexing, that's implied by the ranges it closes
-# BUFFERIZE returns the BUFFER ready for INDEXing (doing this will make splitting a lot easier)
-# NOTE: this has been fixed up a bit
-
-def bufferize_to_store(x:UOp, locals_allowed=False):
-  rngs = x.src[1:]
-  shape = tuple([int(r.vmax+1) for r in rngs])
-  size = prod(shape)
-  assert size > 0, f"no zero sized buffers {shape}"
-  sdtype = x.dtype.ptr(size=size, addrspace=AddrSpace.GLOBAL if not isinstance(x.arg, AddrSpace) else x.arg)
-  if x.src[0].op is Ops.ASSIGN:
-    assign_target, assign_src = x.src[0].src
-    assert assign_target.op is Ops.INDEX
-    return assign_target.replace(dtype=sdtype).store(assign_src, *rngs, dtype=sdtype)
-  # NOTE: the DEFINE_LOCAL needs to be disambiguated here
-  if sdtype.addrspace == AddrSpace.GLOBAL:
-    buf = UOp.new_buffer(x.arg, size, x.dtype)
-  else:
-    if not locals_allowed: return None
-    buf = UOp(Ops.DEFINE_LOCAL, sdtype, arg=0) #UOp.unique().arg)
-  return buf.reshape(shape).index(*rngs, dtype=sdtype).store(x.src[0], *rngs, dtype=sdtype).forced_reshape(shape, dtype=x.dtype)
-
-pm_add_buffers_local = pm_mops+PatternMatcher([
-  (UPat(Ops.BUFFERIZE, name="x"), lambda x: bufferize_to_store(x, True)),
-])
-
-pm_add_buffers = pm_mops+PatternMatcher([
-  (UPat(Ops.BUFFERIZE, name="x"), bufferize_to_store),
-
-  # move RESHAPEs through MSELECT/MSTACK
-  (UPat((Ops.MSELECT, Ops.MSTACK), src=UPat(Ops.RESHAPE), name="m"),
-   lambda m: m.replace(src=tuple([x.src[0] for x in m.src])).reshape(m.src[0].arg)),
-])
-
-# 5. split into kernels
-
-@dataclass
-class LocalAddBufferContext:
-  dg:int = 0
-  map:dict = field(default_factory=dict)
-  vars:dict = field(default_factory=dict)
-
-def debuf(ctx:LocalAddBufferContext, buf:UOp):
-  ret = UOp(Ops.DEFINE_GLOBAL, buf.dtype.ptr(buf.arg), arg=ctx.dg)
-  if buf not in ctx.map: ctx.map[buf] = buf
-  ctx.dg += 1
-  return ret
-
-def unbind_kernel(ctx:LocalAddBufferContext, b:UOp):
-  ctx.vars[b] = None
-  return b.src[0]
-
-def handle_assign(ctx:LocalAddBufferContext, assign:UOp):
-  buf = assign.as_buf()
-  # HACK to put the buffer in the MAP instead of MSTACK/MSELECT
-  if buf.op in {Ops.MSTACK, Ops.MSELECT}: buf = buf.src[0]
-  assert buf not in ctx.map
-  ctx.map[buf] = assign
-  return buf
-
-to_define_global = PatternMatcher([
-  (UPat(Ops.BUFFER, name="buf"), debuf),
-  (UPat(Ops.BIND, name="b"), unbind_kernel),
-  (UPat((Ops.ASSIGN, Ops.MSTACK, Ops.MSELECT), name="assign"), handle_assign),
-
-  # HACK in case any CONSTs were replaced
-  # this is only needed if you are using symbolic
-  #(UPat(Ops.CONST, name="c"), lambda c: c.replace(src=()) if len(c.src) else None),
-])
-
-rangeify_codegen = PatternMatcher([
-  # add loads to non ptr indexes
-  # TODO: this can be moved into codegen?
-  (UPat((Ops.DEFINE_GLOBAL, Ops.STORE), name="dg").f(Ops.INDEX, name="idx", allow_any_len=True),
-   lambda dg,idx: None if isinstance(idx.dtype, (PtrDType, ImageDType)) else idx.replace(dtype=dg.dtype, arg=None).load()),
-
-  # TODO: this can be moved into codegen
-  (UPat(Ops.STORE, name="store").f(Ops.INDEX, allow_any_len=True, name="idx").f(Ops.LOAD),
-    lambda store,idx: idx.replace(src=(store.as_buf(),)+idx.src[1:]).load(store if idx.dtype.addrspace != AddrSpace.LOCAL else store.barrier())),
-
-  # TODO: hack for group for reduce
-  (UPat(Ops.IF, src=(UPat.var("gate"), UPat(Ops.LOAD, src=(UPat.var("src"), UPat.var("barrier"))),)),
-   lambda src, barrier, gate: src.load(UOp(Ops.IF, src=(gate, barrier)))),
-])
-
-def split_store(x:UOp):
-  if len(x.ranges): return None
-  ctx = LocalAddBufferContext()
-  ret = graph_rewrite(x, to_define_global+rangeify_codegen, ctx=ctx, name="kernel split", bottom_up=True)
-
-  # NOTE: the hack for COPY is here
-  ret = ret.sink() if ret.src[1].op is not Ops.COPY else ret.src[1]
-  kernel = UOp(Ops.KERNEL, src=tuple(ctx.map.values())+tuple(ctx.vars.keys()), arg=Kernel(ret,()))
-  return x.as_buf().assign(kernel)
-
-split_kernels = PatternMatcher([
-  (UPat(Ops.STORE, name="x"), split_store),
-])
-
-@track_rewrites(name=lambda sink,ret: f"Schedule {pluralize('Kernel',len([u for u in ret[sink].toposort() if u.op is Ops.KERNEL]))}", replay=True)
-def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
-  tensor_map = graph_rewrite_map(sink, multi_pm+earliest_rewrites, name="earliest")
-  realize_map: dict[UOp, UOp] = {}
-  graph_rewrite(tensor_map[sink], do_realize, ctx=realize_map, name="Input Graph")
-  tensor_map = graph_rewrite_map(tensor_map[sink], add_contiguous, ctx=realize_map, bottom_up=True, input_map=tensor_map, name="add contiguous")
-  tensor_map = graph_rewrite_map(tensor_map[sink], remove_tags, input_map=tensor_map, name="cleanup")
-  tensor_map = graph_rewrite_map(tensor_map[sink], pm_children, ctx=ChildrenContext(), bottom_up=True, input_map=tensor_map, name="children")
-  tensor_map = graph_rewrite_map(tensor_map[sink], pm_rangeify, ctx=RangeifyContext(), bottom_up=True, input_map=tensor_map, name="rangeify")
-  # NOTE: running symbolic can break the graph, leaving RANGE/INDEX/BUFFERIZE in the final graph
-  #tensor_map = graph_rewrite_map(tensor_map[sink], symbolic_simple, input_map=tensor_map, name="symbolic")
-  tensor_map = graph_rewrite_map(tensor_map[sink], pm_cleanups, bottom_up=True, input_map=tensor_map, name="cleanups")
-  if getenv("VIZ"): graph_rewrite(tensor_map[sink], PatternMatcher([]), name="View Rangeify Graph")
-
-  tensor_map = graph_rewrite_map(tensor_map[sink], pm_add_buffers, bottom_up=True, input_map=tensor_map, name="add buffers")
-  tensor_map = graph_rewrite_map(tensor_map[sink], split_kernels, input_map=tensor_map, name="split kernels")
-
-  # 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] = {}
-  assign_rep: dict[UOp, UOp] = {}
-  for u in tensor_map[sink].toposort():
-    if u.op is not Ops.ASSIGN: continue
-    kernel_assign[u.buf_uop] = u
-    for s in u.src[1].src:
-      # TODO: this is probably broken for MSELECT/MSTACK
-      if s.op is not Ops.BUFFER or s is u.buf_uop or (a:=kernel_assign.get(s)) is None: continue
-      if any(x.op is Ops.ASSIGN and x.buf_uop is s for x in u.toposort()):
-        raise RuntimeError(f"cycle detected in graph, kernel for {u.buf_uop} must either depend on ASSIGN or BUFFER")
-      assign_rep[a] = kernel_assign[s] = a.replace(src=a.src+(u,))
-  if assign_rep:
-    tensor_map = graph_rewrite_map(tensor_map[sink], _substitute, ctx=assign_rep, bottom_up=True, input_map=tensor_map, name="fix_assign")
-
-  if getenv("VIZ"): graph_rewrite(tensor_map[sink], PatternMatcher([]), name="View Kernel Graph")
-  return tensor_map

tinygrad/shape/shapetracker.py

@@ -44,13 +44,13 @@ def views_to_real_strides(views: tuple[View, ...], ignore_valid=False) -> tuple[
   ret: list[sint|None] = [None] * len(views[-1].shape)
   idx, valid = views_to_indexed_uops(views)
   for c in split_uop(idx, Ops.ADD):
-    if c.op is Ops.RANGE: ret[c.arg[0]] = 1
-    if c.op is Ops.MUL and c.src[0].op is Ops.RANGE and c.src[1].op is Ops.CONST: ret[c.src[0].arg[0]] = c.src[1].arg
-    if c.op is Ops.MUL and c.src[1].op is Ops.RANGE and c.src[0].op is Ops.CONST: ret[c.src[1].arg[0]] = c.src[0].arg
-  used_ranges = [x.arg[0] for x in idx.toposort() if x.op is Ops.RANGE]
+    if c.op is Ops.RANGE: ret[c.arg] = 1
+    if c.op is Ops.MUL and c.src[0].op is Ops.RANGE and c.src[1].op is Ops.CONST: ret[c.src[0].arg] = c.src[1].arg
+    if c.op is Ops.MUL and c.src[1].op is Ops.RANGE and c.src[0].op is Ops.CONST: ret[c.src[1].arg] = c.src[0].arg
+  used_ranges = [x.arg for x in idx.toposort() if x.op is Ops.RANGE]
   ret = [x if i in used_ranges else 0 for i,x in enumerate(ret)]
   if not ignore_valid:
-    for masked_axis in [x.arg[0] for x in valid.toposort() if x.op is Ops.RANGE]: ret[masked_axis] = None
+    for masked_axis in [x.arg for x in valid.toposort() if x.op is Ops.RANGE]: ret[masked_axis] = None
   return tuple(ret)
 
 @dataclass(frozen=True, order=True)
@@ -112,7 +112,7 @@ class ShapeTracker:
   def axis_is_masked(self, axis:int) -> bool:
     with Context(TRACK_MATCH_STATS=0):
       _, valid = self.to_indexed_uops()
-      return axis in [x.arg[0] for x in graph_rewrite(valid, symbolic_flat).toposort() if x.op is Ops.RANGE]
+      return axis in [x.arg for x in graph_rewrite(valid, symbolic_flat).toposort() if x.op is Ops.RANGE]
 
   def simplify(self) -> ShapeTracker:
     if len(self.views) >= 2 and (new_view := self.views[-2] + self.views[-1]) is not None:

tinygrad/tensor.py

@@ -6,7 +6,7 @@ from typing import Callable, ClassVar, Sequence, cast, get_args, Literal, Suppor
 from tinygrad.dtype import DType, DTypeLike, dtypes, ImageDType, ConstType, least_upper_float, least_upper_dtype, sum_acc_dtype, to_dtype, truncate
 from tinygrad.dtype import _from_np_dtype, _to_np_dtype
 from tinygrad.helpers import argfix, make_tuple, flatten, prod, all_int, round_up, merge_dicts, argsort, getenv, all_same, fully_flatten, dedup
-from tinygrad.helpers import IMAGE, WINO, Metadata, TRACEMETA, ceildiv, fetch, polyN, unwrap, DEBUG, is_numpy_ndarray, RANGEIFY, FUSE_ATTENTION
+from tinygrad.helpers import IMAGE, WINO, Metadata, TRACEMETA, ceildiv, fetch, polyN, unwrap, DEBUG, is_numpy_ndarray
 from tinygrad.gradient import compute_gradient
 from tinygrad.uop.ops import smax, smin, resolve, UOp, Ops, sint, Variable, MathTrait, identity_element, all_metadata
 from tinygrad.uop.spec import tensor_uop_spec, type_verify
@@ -14,7 +14,6 @@ from tinygrad.device import Device, Buffer
 from tinygrad.engine.realize import run_schedule
 from tinygrad.engine.memory import memory_planner
 from tinygrad.engine.schedule import ScheduleItem, create_schedule_with_vars
-from tinygrad.schedule.rangeify import get_rangeify_map
 from tinygrad.schedule.kernelize import get_kernelize_map
 
 # *** all in scope Tensors are here. this gets relevant UOps ***
@@ -40,9 +39,6 @@ def _apply_map_to_tensors(applied_map:dict[UOp, UOp], name:str|None=None) -> Non
     sink = UOp.sink(*[t.uop for t in fixed_tensors])
     new_sink = sink.substitute(applied_map, name=name)
 
-    # NOTE: you can check the Tensor graph early here
-    #if __debug__: type_verify(list(new_sink.toposort()), tensor_uop_spec)
-
     # set the relevant uop to the realized UOps
     for t,s,ns in zip(fixed_tensors, sink.src, new_sink.src):
       if s is ns: continue
@@ -68,7 +64,7 @@ def _frompy(x:list|tuple|bytes, dtype:DType) -> UOp:
     ret = UOp.new_buffer("PYTHON", prod(shape:=get_shape(x)), dtype).reshape(shape)
     assert dtype.fmt is not None, f"{dtype=} has None fmt"
     truncate_function = truncate[dtype]
-    data = struct.pack(f"{ret.size}{dtype.fmt}", *[truncate_function(dtypes.as_const(xi, dtype)) for xi in fully_flatten(x)])
+    data = struct.pack(f"@{ret.size}{dtype.fmt}", *[truncate_function(xi) for xi in fully_flatten(x)])
   # fake realize
   ret.buffer.allocate(memoryview(data if Device.DEFAULT != "PYTHON" else bytearray(data)))
   return ret
@@ -177,8 +173,8 @@ class Tensor(MathTrait):
     all_tensors[weakref.ref(self)] = None
   def __del__(self): all_tensors.pop(weakref.ref(self), None)
 
-  def _apply_uop(self, fxn:Callable, *x:Tensor, extra_args=(), **kwargs) -> Tensor:
-    new_uop: UOp = fxn(*[t.uop for t in (self,)+x], *extra_args, **kwargs)
+  def _apply_uop(self, fxn:Callable, *x:Tensor, **kwargs) -> Tensor:
+    new_uop: UOp = fxn(*[t.uop for t in (self,)+x], **kwargs)
     if (metadata:=_METADATA.get()) is not None: all_metadata[new_uop] = (metadata,)
     needs_input_grad = [t.requires_grad for t in (self,)+x]
     return Tensor(new_uop, device=new_uop.device, requires_grad=True if any(needs_input_grad) else None if None in needs_input_grad else False)
@@ -235,7 +231,7 @@ class Tensor(MathTrait):
     # verify Tensors match the spec
     if __debug__: type_verify(list(big_sink.toposort()), tensor_uop_spec)
 
-    becomes_map = get_rangeify_map(big_sink) if RANGEIFY else get_kernelize_map(big_sink)
+    becomes_map = get_kernelize_map(big_sink)
     _apply_map_to_tensors(becomes_map, name="Apply Kernelize Map")
     return self
 
@@ -349,8 +345,7 @@ class Tensor(MathTrait):
     print(t.tolist())
     ```
     """
-    # TODO: remove half once minimum python supports it
-    if self.dtype in (dtypes.half, dtypes.bfloat16, *dtypes.fp8s): return self.cast(dtypes.float32).tolist()
+    if self.dtype in (dtypes.bfloat16, *dtypes.fp8s): return self.cast(dtypes.float32).tolist()
     return self.data().tolist()
 
   def numpy(self) -> 'np.ndarray':  # type: ignore [name-defined] # noqa: F821
@@ -1192,7 +1187,7 @@ class Tensor(MathTrait):
         x = x.shrink(tuple(flatten(((0, s), (0, 1)) for s in x.shape[::2]))).reshape(x.shape[::2])
 
     # dim injection from None by including None dim size (which is 1) and dim collapse by skipping int dim size
-    x = x.reshape(tuple(index['size'] for index in indices_parsed if not isinstance(index['index'], sint)))
+    x = x.reshape(tuple(index['size'] for index in indices_parsed if not isinstance(index['index'], (int, UOp))))
 
     # tensor indexing
     if tops := [(d,i) for d,i in enumerate(i_ for i_ in indices_parsed if not isinstance(i_['index'], int)) if isinstance(i['index'], Tensor)]:
@@ -1212,7 +1207,7 @@ class Tensor(MathTrait):
       # inject 1's for the extra dims added in create masks
       reshape_arg = x.shape[:dims[0]] + (1,) * len(big_shape) + x.shape[dims[0]:]
       # sum reduce the extra dims introduced in create masks
-      x = (mask.where(x.reshape(reshape_arg), 0)).sum(sum_axis:=tuple(d + len(big_shape) for d in dims), dtype=x.dtype)
+      x = (x.reshape(reshape_arg) * mask).sum(sum_axis:=tuple(d + len(big_shape) for d in dims), dtype=x.dtype)
 
       # special permute case
       if dims[0] != 0 and len(dims) != 1 and tuple(dims) != tuple(range(dims[0], dims[-1]+1)):
@@ -1303,7 +1298,7 @@ class Tensor(MathTrait):
     assert all(s >= i for d,(s,i) in enumerate(zip(self.shape, index.shape)) if d != dim), "requires self.shape[d] >= index.shape[d] for all d != dim"
     index = index.to(self.device)
     x = self.shrink(tuple((0, i) if d != dim else None for d,i in enumerate(index.shape))).unsqueeze(-1).transpose(-1, dim)
-    return (index.unsqueeze(-1)._one_hot_along_dim(self.shape[dim]).where(x, 0)).sum(-1, dtype=self.dtype)
+    return (x * index.unsqueeze(-1)._one_hot_along_dim(self.shape[dim])).sum(-1, dtype=self.dtype)
 
   def cat(self:Tensor, *args:Tensor, dim:int=0) -> Tensor:
     """
@@ -2437,7 +2432,7 @@ class Tensor(MathTrait):
       # https://arxiv.org/pdf/1603.07285 inverse of relationship 15 in section 5.1.
       output_size = tuple((i-1)*s - (pB+pA) + (d*(k-1)+1) for i,k,d,s,(pA,pB) in zip(spatial_shape,k_,d_,s_,p_))
     else: output_size = output_size[-len(spatial_shape):]
-    ret = (indices.reshape(bs,c,1,-1)._one_hot_along_dim(prod(output_size), 2).where(self.reshape(bs,c,1,-1), 0)).sum(3)
+    ret = (indices.reshape(bs,c,1,-1)._one_hot_along_dim(prod(output_size), 2) * self.reshape(bs,c,1,-1)).sum(3)
     return ret.reshape(bs,c,*output_size)
 
   def conv2d(self, weight:Tensor, bias:Tensor|None=None, groups=1, stride=1, dilation=1, padding:int|tuple[int, ...]=0,
@@ -2941,11 +2936,11 @@ class Tensor(MathTrait):
     """
     return self*-1 if self.dtype != dtypes.bool else self.logical_not()
 
-  def contiguous(self, *args, **kwargs) -> Tensor:
+  def contiguous(self, **kwargs) -> Tensor:
     """
     Returns a contiguous tensor.
     """
-    return self._apply_uop(UOp.contiguous, extra_args=args, **kwargs)
+    return self._apply_uop(UOp.contiguous, **kwargs)
 
   def fuse(self) -> Tensor:
     """
@@ -2996,9 +2991,6 @@ class Tensor(MathTrait):
     print(Tensor([0., 1., 2., 3.]).exp().numpy())
     ```
     """
-    # TODO: make it generic, and same thing to log and cos
-    if self.is_floating_point(): return self.cast(least_upper_dtype(self.dtype, dtypes.float32)).mul(1/math.log(2)).exp2().cast(self.dtype)
-    # TODO: behavior when DEFAULT_FLOAT is bfloat16 and input is int32?
     return self.mul(1/math.log(2)).exp2()
 
   def exp2(self) -> Tensor:
@@ -3518,15 +3510,16 @@ class Tensor(MathTrait):
     """
     return self * self.softplus().tanh()
 
-  def softplus(self, beta=1.0) -> Tensor:
+  def softplus(self, beta=1.0, threshold=20.0) -> Tensor:
     """
     Applies the Softplus function element-wise.
+    For numerical stability, the implementation folds into identity function when `self * beta > threshold`.
 
     ```python exec="true" source="above" session="tensor" result="python"
     print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).softplus().numpy())
     ```
     """
-    return (1/beta) * (self*beta).logaddexp(0.0)
+    return (self * beta > threshold).where(self, (1/beta) * (1 + (self*beta).exp()).log())
 
   def softsign(self) -> Tensor:
     """
@@ -3758,13 +3751,6 @@ class Tensor(MathTrait):
     # TODO: remove other*0?
     return (other < 0).where(-self.abs(), self.abs()) + other*0
 
-  def logaddexp(self, other) -> Tensor:
-    """
-    Calculates (self.exp()+other.exp()).log(), elementwise.
-    """
-    m = self.maximum(other)
-    return ((self-m).exp() + (self._broadcasted(other)[1]-m).exp()).log() + m
-
   # ***** op wrappers *****
 
   def __invert__(self) -> Tensor: return self.bitwise_not()
@@ -3930,11 +3916,7 @@ class Tensor(MathTrait):
     if enable_gqa:
       key = key.repeat_interleave(self.shape[-3] // key.shape[-3], dim=-3)
       value = value.repeat_interleave(self.shape[-3] // value.shape[-3], dim=-3)
-
-    if FUSE_ATTENTION: q, key, value = self.contiguous(), key.contiguous(), value.contiguous()
-    else: q = self
-
-    qk = q.matmul(key.transpose(-2,-1), dtype=least_upper_dtype(q.dtype, key.dtype, dtypes.float32)) / math.sqrt(q.shape[-1])
+    qk = self.matmul(key.transpose(-2,-1), dtype=least_upper_dtype(self.dtype, key.dtype, dtypes.float32)) / math.sqrt(self.shape[-1])
     # handle attention mask
     if is_causal:
       if attn_mask is not None: raise RuntimeError("cannot set attn_mask when is_causal=True")
@@ -3942,8 +3924,7 @@ class Tensor(MathTrait):
     if attn_mask is not None:
       if attn_mask.dtype == dtypes.bool: attn_mask = attn_mask.where(0, -float("inf"))
       qk = qk + attn_mask
-    attn = qk.cast(self.dtype).softmax(-1).dropout(dropout_p) @ value
-    return attn.fuse() if FUSE_ATTENTION else attn
+    return qk.cast(self.dtype).softmax(-1).dropout(dropout_p) @ value
 
   def _do_reduction(self, reduction:ReductionStr="mean") -> Tensor:
     if reduction not in get_args(ReductionStr): raise ValueError(f"{reduction=} must be one of {get_args(ReductionStr)}")

tinygrad/uop/__init__.py

@@ -109,9 +109,6 @@ class GroupOp:
   # BinaryOps that satisfy f(x,x)=x see https://en.wikipedia.org/wiki/Idempotence
   Idempotent = {Ops.OR, Ops.AND, Ops.MAX}
 
-  # These can change the dtype to bool
-  Comparison = {Ops.CMPLT, Ops.CMPNE, Ops.CMPEQ}
-
   # do not preserve f(0) = 0
   UnsafePad = {Ops.RECIP, Ops.LOG2, Ops.EXP2, Ops.IDIV, Ops.POW}
 

tinygrad/uop/decompositions.py

@@ -5,7 +5,7 @@ from tinygrad.device import is_dtype_supported
 from tinygrad.helpers import polyN, DISABLE_FAST_IDIV
 from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher
 
-TRANSCENDENTAL_DTYPES = (dtypes.float16, dtypes.float32, dtypes.float64)
+TRANSCENDENTAL_SUPPORTED_DTYPES = (dtypes.float16, dtypes.float32, dtypes.float64)
 
 def _lazy_map_numbers(x:UOp, inf:UOp, _inf:UOp, nan:UOp, ratio:UOp):
   """replace inf -> inf, -inf -> _inf, nan -> nan, otherwise -> ratio"""
@@ -32,14 +32,14 @@ def pow2if(q:UOp, float_dtype:DType):
 
 def ilogb2k(d:UOp) -> UOp:
   """calculate the integer part of log2(d), where d is normalized fp value in the range of [0, +inf)."""
-  assert d.dtype.scalar() in TRANSCENDENTAL_DTYPES
+  assert d.dtype.scalar() in TRANSCENDENTAL_SUPPORTED_DTYPES
   dint = d.bitcast({dtypes.float64: dtypes.int64, dtypes.float32: dtypes.int32, dtypes.float16: dtypes.int16}[d.dtype.scalar()].vec(d.dtype.vcount))
   # -1 <= ilog2bk(d) <= 128
   return (shr(dint, mantissa_bits(d.dtype)) & exponent_mask(d.dtype)) - exponent_bias(d.dtype)
 
 def ldexp3k(d:UOp, e:UOp) -> UOp:
   """d*2^e. e is a number obtained by casting an integer in the range [-127, 127] to a float. d is any float number."""
-  assert d.dtype.scalar() in TRANSCENDENTAL_DTYPES and e.dtype.scalar() in TRANSCENDENTAL_DTYPES
+  assert d.dtype.scalar() in TRANSCENDENTAL_SUPPORTED_DTYPES and e.dtype.scalar() in TRANSCENDENTAL_SUPPORTED_DTYPES
   dtype = {dtypes.float64: dtypes.int64, dtypes.float32: dtypes.int32, dtypes.float16: dtypes.int16}[d.dtype.scalar()].vec(d.dtype.count)
   m1 = d.bitcast(dtype)
   m2 = shl(e.cast(dtype), mantissa_bits(d.dtype))
@@ -47,12 +47,12 @@ def ldexp3k(d:UOp, e:UOp) -> UOp:
 
 def ldexp2k(d:UOp, e:UOp) -> UOp:
   """d*2^e. much faster than ldexp3k but risky. d > 0 and d is not denormal."""
-  assert d.dtype.scalar() in TRANSCENDENTAL_DTYPES and e.dtype.scalar() in (dtypes.int16, dtypes.int32, dtypes.int64)
+  assert d.dtype.scalar() in TRANSCENDENTAL_SUPPORTED_DTYPES and e.dtype.scalar() in (dtypes.int16, dtypes.int32, dtypes.int64)
   return (d * pow2if(shr(e, 1), d.dtype)) * pow2if(e - shr(e, 1), d.dtype)
 
 def frexp(v:UOp) -> tuple[UOp, UOp]:
   """frexp(v) -> (mantissa, exponent) assuming v != 0"""
-  assert v.dtype.scalar() in TRANSCENDENTAL_DTYPES
+  assert v.dtype.scalar() in TRANSCENDENTAL_SUPPORTED_DTYPES
   # m1 = masks for mantissa, m2 = masks to normalize the mantissa.
   m1 = {dtypes.float64: 0x000FFFFFFFFFFFFF, dtypes.float32: 0x807FFFFF, dtypes.float16: 0x83FF}[v.dtype.scalar()]
   m2 = {dtypes.float64: 0x3FE0000000000000, dtypes.float32: 0x3F000000, dtypes.float16: 0x3800}[v.dtype.scalar()]
@@ -72,7 +72,7 @@ def payne_hanek_reduction(d:UOp) -> tuple[UOp, UOp]:
   - `r`[d.dtype] is the reminder value corresponding to `round_to_nearest(x % pi/2)`.
   - `q`[int32] is an integer, and q % 4 is corresponding to the quadrant of the original angle `d`.
   """
-  assert d.dtype.scalar() in TRANSCENDENTAL_DTYPES
+  assert d.dtype.scalar() in TRANSCENDENTAL_SUPPORTED_DTYPES
   # https://stackoverflow.com/questions/30463616/payne-hanek-algorithm-implementation-in-c/30465751#30465751
   # 190 bits of 2/pi for Payne-Hanek style argument reduction
   two_over_pi_f = [0x00000000, 0x28be60db, 0x9391054a, 0x7f09d5f4, 0x7d4d3770, 0x36d8a566, 0x4f10e410]
@@ -174,7 +174,7 @@ def xsin(d:UOp, fast:bool=False, switch_over:float=30.0) -> UOp:
   - fast=True assumes x <= switch_over.
   - switch_over is the threshold for switching to payne_hanek_reduction.
   """
-  assert d.dtype.scalar() in TRANSCENDENTAL_DTYPES
+  assert d.dtype.scalar() in TRANSCENDENTAL_SUPPORTED_DTYPES
   # mask +-inf/nan as zero
   x = _lazy_map_numbers(d, d.const_like(0.0), d.const_like(0.0), d.const_like(0.0), d)
   # x_sign = sign(x)
@@ -196,7 +196,7 @@ def xexp2(d:UOp) -> UOp:
   Implements a 1.0 ULP approximation for Ops.EXP2
   - Paper: https://arxiv.org/pdf/2001.09258
   """
-  assert d.dtype.scalar() in TRANSCENDENTAL_DTYPES
+  assert d.dtype.scalar() in TRANSCENDENTAL_SUPPORTED_DTYPES
   # mask +=inf/nan as zero.
   x = _lazy_map_numbers(d, d.const_like(0.0), d.const_like(0.0), d.const_like(0.0), d)
   q = rintk(x)
@@ -222,7 +222,7 @@ def xlog2(d:UOp) -> UOp:
   Implements a 1.0 ULP approximation for Ops.LOG2
   Paper: https://arxiv.org/pdf/2001.09258 5.5
   """
-  assert d.dtype.scalar() in TRANSCENDENTAL_DTYPES
+  assert d.dtype.scalar() in TRANSCENDENTAL_SUPPORTED_DTYPES
   # TODO: float16 denormal need float32 to achieve precision
   if d.dtype.scalar() == dtypes.float16: return xlog2(d.cast(dtypes.float32)).cast(dtypes.float16)
   FLT_MIN = d.const_like(1e-6 if d.dtype.scalar() == dtypes.float16 else 1e-4)
@@ -280,7 +280,7 @@ def magicgu(vmax:int, d:int) -> tuple[int,int]:
       return m, s
   assert False
 
-def fast_idiv(device: str, x: UOp, d: int, dont_cast=False) -> UOp|None:
+def fast_idiv(device: str, x: UOp, d: int) -> UOp|None:
   # If d is a power of two this is not valid for signed ints!
   is_unsigned = True if x.vmin>=0 or x.dtype in dtypes.uints else False
   assert d>0, "Sign should have been taken out of divisor"
@@ -288,10 +288,6 @@ def fast_idiv(device: str, x: UOp, d: int, dont_cast=False) -> UOp|None:
   m,s = magicgu(max(vmax, abs(vmin)), d)
   if m*vmin >= dtypes.min(x.dtype) and m*vmax <= dtypes.max(x.dtype):
     return ((x*m) >> s) if is_unsigned else ((x*m) >> s) + (x<0).where(x.ufix(1), 0)
-  # before we try casting to a larger dtype (slow), we see if there are powers of two in d we can shift to make x smaller
-  if (largest_factor_of_two_in_d := (d & -d)) > 1:
-    if (ret:=fast_idiv(device, x//largest_factor_of_two_in_d, d//largest_factor_of_two_in_d, dont_cast=True)) is not None: return ret
-  if dont_cast: return None
   # promo_lattice needs to return an unsigned type if the type is unsigned
   if dtypes.is_int(next_dtype := promo_lattice[x.dtype][-1]) and is_dtype_supported(next_dtype, None if device=='' else device):
     if m*vmin >= dtypes.min(next_dtype) and m*vmax <= dtypes.max(next_dtype):
@@ -319,12 +315,8 @@ def threefry2x32(x: UOp, key: UOp):
 powers_of_two = {2**i:i for i in range(64)}
 @functools.cache
 def get_late_rewrite_patterns(ops:tuple[Ops, ...], force_transcendental=False):
-  pat: list[tuple[UPat, Callable]] = []
-  for op,f in ((Ops.EXP2, xexp2), (Ops.LOG2, xlog2), (Ops.SIN, xsin)):
-    if op not in ops or force_transcendental:
-      pat += [(UPat(op, dtype=TRANSCENDENTAL_DTYPES, src=(UPat.var("d"),)), f),
-              (UPat(op, dtype=tuple(dt for dt in dtypes.floats if dt not in TRANSCENDENTAL_DTYPES), src=(UPat.var("d"),), name="x"),
-                lambda x,d: d.cast(dtypes.float32).alu(x.op).cast(x.dtype))]
+  pat: list[tuple[UPat, Callable]] = [(UPat(op, dtype=TRANSCENDENTAL_SUPPORTED_DTYPES, src=(UPat.var("d"),)), f) for op,f in \
+           ((Ops.EXP2, xexp2), (Ops.LOG2, xlog2), (Ops.SIN, xsin)) if op not in ops or force_transcendental]
   # no real hardware supports THREEFRY, but NullRenderer does
   if Ops.THREEFRY not in ops: pat.append((UPat(Ops.THREEFRY, dtype=dtypes.uint64, src=(UPat.var("x"), UPat.var("key"))), threefry2x32))
   # MAX can be rewritten as CMPLT + WHERE (max function is annoying on many cstyle backends)
@@ -333,8 +325,6 @@ def get_late_rewrite_patterns(ops:tuple[Ops, ...], force_transcendental=False):
   if Ops.SQRT not in ops: pat.append((UPat(Ops.SQRT, src=UPat.var("d")), lambda d: xpow(d, d.const_like(0.5))))
   # rewrite MOD to AND (which should always be supported, but not for generic in tests): x % (2**y) -> x & (2**y-1)
   if Ops.AND in ops: pat += [(UPat.var("x", dtypes.ints)%UPat.cvar("c"), lambda x,c: x & (c.arg-1) if c.arg in powers_of_two else None)]
-  if Ops.OR in ops: pat += [(UPat.var("x", dtypes.bool).logical_not()&UPat.var("y", dtypes.bool).logical_not(),
-    lambda x,y: (x | y).logical_not())]
   # rewrite MUL/IDIV to SHL+SHR: x*(2**y) -> shl(x,y) and x//(2**y) -> shr(x,y)
   if Ops.SHL in ops: pat += [(UPat.var("x", dtypes.ints)*UPat.cvar("c"), lambda c,x: x << v if (v:=powers_of_two.get(c.arg, 0)) else None)]
   if Ops.SHR in ops:
@@ -360,8 +350,4 @@ def get_late_rewrite_patterns(ops:tuple[Ops, ...], force_transcendental=False):
     ]
   if Ops.CMPEQ in ops: pat += [(UPat.var('x').ne(UPat.var('y')).logical_not(), lambda x,y: x.alu(Ops.CMPEQ, y))]
   if Ops.MULACC in ops: pat += [(UPat.var('a')*UPat.var('b')+UPat.var('c'), lambda a,b,c: a.alu(Ops.MULACC, b, c))]
-  # some backends emit FDIV for RECIP, in that case: a*(1/b) -> a/b
-  if Ops.FDIV in ops:
-    pat += [(UPat.var("x").reciprocal(), lambda x: x.const_like(1).alu(Ops.FDIV, x))]
-    pat += [(UPat.var("a", dtypes.floats) * UPat.const(dtypes.floats, 1).alu(Ops.FDIV, UPat.var("b")), lambda a,b: a.alu(Ops.FDIV, b))]
   return PatternMatcher(pat)

tinygrad/uop/ops.py

@@ -6,15 +6,12 @@ from enum import Enum, auto
 from tinygrad.uop import Ops, GroupOp
 from tinygrad.uop.mathtraits import MathTrait
 from tinygrad.dtype import ConstType, ImageDType, dtypes, DType, truncate, PtrDType
-from tinygrad.helpers import ContextVar, all_int, prod, getenv, all_same, Context, partition, temp, unwrap, T, argfix, Metadata, flatten, TRACEMETA
+from tinygrad.helpers import ContextVar, all_int, prod, getenv, all_same, Context, partition, temp, unwrap, T, argfix, Metadata, flatten
 from tinygrad.helpers import PICKLE_BUFFERS, PROFILE, dedup, cdiv, cmod, diskcache_put, to_function_name, cpu_profile, TracingKey
 if TYPE_CHECKING:
   from tinygrad.shape.shapetracker import ShapeTracker
   from tinygrad.device import Buffer, MultiBuffer
 
-class AxisType(Enum):
-  GLOBAL = auto(); LOCAL = auto(); LOOP = auto(); GROUP_REDUCE = auto(); REDUCE = auto(); UPCAST = auto(); UNROLL = auto()  # noqa: E702
-
 # https://en.wikipedia.org/wiki/Identity_element
 def identity_element(op:Ops, dt:DType) -> ConstType: return dtypes.as_const({Ops.ADD:0, Ops.MUL:1, Ops.MAX:dtypes.min(dt)}[op], dt)
 
@@ -139,22 +136,15 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
 
   @functools.cached_property
   def st(self) -> ShapeTracker|None:
-    if self.op is Ops.INDEX and self.src[0].op in {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG,
-                                                   Ops.BUFFER, Ops.BUFFERIZE, Ops.VECTORIZE, Ops.STORE}:
-      return None
-    if self.op is Ops.BARRIER: return None
-    if self.op in GroupOp.Block: return None
+    if self.op in GroupOp.Block or self.op is Ops.INDEX: return None
     from tinygrad.shape.shapetracker import ShapeTracker
     # VIEW and MovementOps define a new ShapeTracker from the arg
     if self.op is Ops.VIEW: return self.arg
-    if self.op is Ops.BUFFERIZE: return ShapeTracker.from_shape((prod(tuple([int(r.vmax+1) for r in self.src[1:]])),))
-    #if self.op is Ops.BUFFERIZE: return ShapeTracker.from_shape(tuple([r.vmax+1 for r in self.src[1:]]))
     # allow reshape from nothing
     if self.op is Ops.RESHAPE and self.src[0].st is None: return ShapeTracker.from_shape(self.arg)
     if self.op in GroupOp.Movement: return unwrap(self.src[0].st).mop(self.op, self.arg)
     # CONST with a DEVICE has a shape of ()
     if self.op is Ops.CONST and len(self.src) and self.src[0].op is Ops.DEVICE: return ShapeTracker.from_shape(())
-    if self.op is Ops.STORE and isinstance(self.dtype, PtrDType): return ShapeTracker.from_shape((self.dtype.size,))
     # BufferOps and ASSIGN flow ShapeTracker from a direct edge
     if self.op in {Ops.STORE, Ops.ASSIGN, Ops.LOAD}: return self.src[0].st
     if self.op in GroupOp.Buffer: return views[0] if (views:=[x.st for x in self.src if x.op is Ops.VIEW]) else None
@@ -166,11 +156,6 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
       sz = cast(PtrDType, self.dtype).size
       return ShapeTracker.from_shape((sz,)) if sz > 0 else None
 
-    # CONTIGUOUS with RANGE
-    # TODO: how are these not RANGE?
-    if self.op is Ops.CONTIGUOUS and len(self.src) > 1 and all(x.op is Ops.RANGE for x in self.src[1:]):
-      return ShapeTracker.from_shape((tuple([int(x.vmax+1) for x in self.src[1:]])+self.src[0].shape))
-
     # hack for PTX, CASTing the ptr loses the shape
     if self.op is Ops.CAST and self.src[0].op is Ops.DEFINE_GLOBAL: return None
 
@@ -203,25 +188,27 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
   @functools.cached_property
   def ranges(self) -> dict[UOp, None]:
     if self.op is Ops.RANGE: return {self:None}
-    range_start = {Ops.BUFFERIZE: 1, Ops.REDUCE: 1, Ops.STORE: 2, Ops.WMMA: 3}
-    ret: dict[UOp, None] = {}
-    if self.op in range_start.keys():
-      for s in self.src[:range_start[self.op]]: ret.update(s.ranges)
-      delete_ranges = self.src[range_start[self.op]:]
-      if len(delete_ranges):
-        for s in UOp.sink(*delete_ranges).ranges:
-          if s in ret: del ret[s]
+    if self.op in {Ops.BUFFERIZE, Ops.REDUCE}:
+      ret = self.src[0].ranges.copy()
+      for s in self.src[1:]:
+        if s in ret: del ret[s]
+    elif self.op in {Ops.STORE}:
+      ret = self.src[0].ranges.copy()
+      ret.update(self.src[1].ranges)
+      for s in self.src[2:]:
+        if s in ret: del ret[s]
     else:
+      ret = {}
       for s in self.src: ret.update(s.ranges)
     return ret
 
   # *** uop evaluation ***
 
-  def simplify(self, tracked=False):
+  def simplify(self):
     # late import!
     from tinygrad.uop.symbolic import symbolic
-    with Context(TRACK_MATCH_STATS=0 if not tracked else TRACK_MATCH_STATS.value):
-      return graph_rewrite(self, symbolic, name="simplify")
+    with Context(TRACK_MATCH_STATS=0):
+      return graph_rewrite(self, symbolic)
   def ssimplify(self) -> UOp|ConstType: return ret.arg if (ret:=self.simplify()).op is Ops.CONST else ret
   def _eval(self, dtype, expected_type:Type[T]) -> T:
     assert self.dtype in dtype, f"eval with wrong dtype {self}"
@@ -250,8 +237,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     ret = self.arg[1] if self.op is Ops.REDUCE_AXIS else self.arg[7]
     assert isinstance(ret, tuple) and all(isinstance(x, int) for x in ret), f"axis_arg trying to return {ret}"
     return ret
-  def sink(*srcs:UOp|None, **kwargs):  # pylint: disable=no-self-argument
-    return UOp(Ops.SINK, dtypes.void, tuple([x for x in srcs if x is not None]), **kwargs)
+  def sink(self, *srcs:UOp|None, **kwargs): return UOp(Ops.SINK, dtypes.void, (self,)+tuple([x for x in srcs if x is not None]), **kwargs)
   def detach(self): return UOp(Ops.DETACH, self.dtype, (self,))
   def index(self, *srcs:UOp|None, **kwargs):
     return UOp(Ops.INDEX, kwargs.pop("dtype", self.dtype), (self,)+tuple([x for x in srcs if x is not None]), **kwargs)
@@ -279,7 +265,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
       i = (i,)
     return UOp(Ops.GEP, self.dtype.scalar().vec(len(i)) if len(i) > 1 else self.dtype.scalar(), (self,), i)
   def load(self, *src:UOp, **kwargs): return UOp(Ops.LOAD, dtype=kwargs.pop("dtype", self.dtype.base), src=(self,)+src, **kwargs)
-  def store(self, *src:UOp, **kwargs): return UOp(Ops.STORE, kwargs.pop("dtype", dtypes.void), (self,)+src, **kwargs)
+  def store(self, *src:UOp, **kwargs): return UOp(Ops.STORE, dtypes.void, (self,)+src, **kwargs)
   def assign(self, x:UOp): return UOp(Ops.ASSIGN, self.dtype, (self, x))
   def barrier(self, *src:UOp): return UOp(Ops.BARRIER, src=(self,)+src)
   def alu(self, op, *src:UOp, **kwargs):
@@ -299,10 +285,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
       else: ret = ret.replace(src=(UOp(Ops.DEVICE, arg=device),))
     return ret
   @staticmethod
-  def range(dtype:DType, end:sint, *arg):
-    if len(arg) == 0: raise RuntimeError("range needs an arg")
-    if len(arg) == 1: arg = arg+(AxisType.LOOP,)
-    return UOp(Ops.RANGE, dtype=dtype, src=(sint_to_uop(end),), arg=arg)
+  def range(dtype:DType, end:sint, idx:int): return UOp(Ops.RANGE, dtype=dtype, src=(sint_to_uop(end),), arg=idx)
   def r(self, op:Ops, axis:tuple[int, ...]):
     axis = tuple(sorted([x for x in axis if resolve(self.shape[x] != 1)]))
     if len(axis) == 0: return self
@@ -389,13 +372,12 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     if self.st == ret.st: return self  # ignore NOOPs, also check ret.st
     return ret
 
-  def forced_reshape(self, arg:tuple[sint, ...], **kwargs): return UOp(Ops.RESHAPE, kwargs.pop("dtype", self.dtype), src=(self,), arg=arg)
-
+  def forced_reshape(self, arg:tuple[sint, ...]): return UOp(Ops.RESHAPE, self.dtype, src=(self,), arg=arg)
   def reshape(self, arg:tuple[sint, ...]): return self._mop(Ops.RESHAPE, arg)
-  def expand(self, arg:tuple[sint, ...]): return self._mop(Ops.EXPAND, arg)
-  def shrink(self, arg:tuple[tuple[sint, sint], ...]): return self._mop(Ops.SHRINK, arg)
   def pad(self, arg:tuple[tuple[sint, sint], ...]): return self._mop(Ops.PAD, arg)
-  def permute(self, arg:tuple[int, ...]): return self._mop(Ops.PERMUTE, arg)
+  def expand(self, arg:tuple[sint, ...]): return self._mop(Ops.EXPAND, arg)
+  def permute(self, arg:tuple[sint, ...]): return self._mop(Ops.PERMUTE, arg)
+  def shrink(self, arg:tuple[tuple[sint, sint], ...]): return self._mop(Ops.SHRINK, arg)
   def flip(self, arg:tuple[bool, ...]): return self._mop(Ops.FLIP, arg)
 
   # *** uop UNIQUE ***
@@ -427,15 +409,6 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     if self.op is Ops.MSTACK: return UOp(Ops.MSTACK, self.dtype, src=tuple(x.buf_uop for x in self.src))
     assert self.op is Ops.ASSIGN, f"must be ASSIGN {self.op}"
     return self.src[0].base
-
-  def as_buf(self) -> UOp:
-    if self.op is Ops.MSELECT: return self.src[0].as_buf().mselect(self.arg)
-    if self.op is Ops.MSTACK: return UOp(Ops.MSTACK, self.dtype, src=tuple(x.as_buf() for x in self.src))
-    # TODO: this should be the only one of these. this is the one RANGEIFY uses
-    s = self
-    while len(s.src) and s.op is not Ops.BUFFER: s = s.src[0]
-    return s
-
   @property
   def buffer(self) -> Buffer|MultiBuffer:
     from tinygrad.device import Buffer, MultiBuffer
@@ -541,8 +514,13 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
         if s1_vmax < 0: return (0, -s1_vmin-1) if s0_vmin >= 0 else (-(-s1_vmin-1), 0) if s0_vmax <= 0 else (-(-s1_vmin-1), -s1_vmin-1)
       if self.op is Ops.IDIV:
         assert isinstance(s0_vmin, int) and isinstance(s0_vmax, int) and isinstance(s1_vmin, int) and isinstance(s1_vmax, int)
-        if s1_vmin*s1_vmax>0:
-          return min(vals:=(cdiv(s0_vmin, s1_vmin), cdiv(s0_vmin, s1_vmax), cdiv(s0_vmax, s1_vmin), cdiv(s0_vmax, s1_vmax))), max(vals)
+        if (c:=s1_vmin) == s1_vmax:  # s1 is a const
+          if c > 0: return cdiv(s0_vmin, c), cdiv(s0_vmax, c)
+          if c < 0: return cdiv(s0_vmax, c), cdiv(s0_vmin, c)
+        if (s0_vmax <= 0 and s1_vmax < 0): return cdiv(s0_vmax, s1_vmin), cdiv(s0_vmin, s1_vmax)
+        if (s0_vmin >= 0 and s1_vmin > 0): return cdiv(s0_vmin, s1_vmax), cdiv(s0_vmax, s1_vmin)
+        if (s0_vmax <= 0 and s1_vmin > 0): return cdiv(s0_vmin, s1_vmin), cdiv(s0_vmax, s1_vmax)
+        if (s0_vmin >= 0 and s1_vmax < 0): return cdiv(s0_vmax, s1_vmax), cdiv(s0_vmin, s1_vmin)
       if self.op is Ops.MAX: return max(s0_vmin, s1_vmin), max(s0_vmax, s1_vmax)
       if self.op is Ops.CMPLT: return (s0_vmax<s1_vmin, s0_vmin<s1_vmax)
       if self.op is Ops.CMPNE: return ((s0_vmax < s1_vmin) or (s1_vmax < s0_vmin), not (s0_vmin == s0_vmax == s1_vmin == s1_vmax))
@@ -577,10 +555,12 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     return fxn(**{k.arg[0]:v for k,v in var_vals.items() if k.arg[0] in varnames})
 
   def render(self, simplify=True, pm:PatternMatcher|None=None) -> str:
-    with Context(TRACK_MATCH_STATS=0):
-      ret = graph_rewrite(self.simplify() if simplify else self, renderer if pm is None else pm)
+    ret = graph_rewrite(self.simplify() if simplify else self, renderer if pm is None else pm)
     return ret.arg if ret.op is Ops.NOOP else str(ret)
 
+class AxisType(Enum):
+  GLOBAL = auto(); LOCAL = auto(); LOOP = auto(); GROUP_REDUCE = auto(); REDUCE = auto(); UPCAST = auto(); UNROLL = auto()  # noqa: E702
+
 @dataclass(frozen=True)
 class KernelInfo:
   name: str = "test"            # name of the kernel
@@ -683,8 +663,7 @@ class UPat(MathTrait):
   def var(name:str|None=None, dtype:DType|tuple[DType, ...]|None=None): return UPat(dtype=dtype, name=name)
   @staticmethod
   @functools.cache
-  def cvar(name:str|None=None, dtype:DType|tuple[DType, ...]|None=None, vec=True):
-    return UPat((Ops.CONST,Ops.VCONST) if vec else Ops.CONST, dtype, name=name)
+  def cvar(name:str|None=None, dtype:DType|None=None, vec=True): return UPat((Ops.CONST,Ops.VCONST) if vec else Ops.CONST, dtype, name=name)
   @staticmethod
   def const(dtype:DType|tuple[DType, ...]|None, b:ConstType): return UPat(Ops.CONST, dtype=dtype, arg=b)
 
@@ -781,7 +760,7 @@ class PatternMatcher:
   def __reduce__(self): return PatternMatcher, ([(x,deconstruct_function(fxn) if fxn.__name__ == "<lambda>" else fxn) for x,fxn in self.patterns],)
 
   @functools.cache  # pylint: disable=method-cache-max-size-none
-  def __add__(self, more:PatternMatcher) -> PatternMatcher: return PatternMatcher(self.patterns+more.patterns)
+  def __add__(self, more:PatternMatcher): return PatternMatcher(self.patterns+more.patterns)
 
   def rewrite(self, uop:UOp, ctx=None) -> UOp|None:
     ler = {u.op for u in uop.src}
@@ -800,7 +779,7 @@ def track_uop(u:UOp):
   uop_number[u] = num = next(ucount)
   # KERNEL also has a UOp in the arg
   arg = type(u.arg)(track_uop(u.arg.ast), u.arg.metadata) if u.op is Ops.KERNEL else u.arg
-  uop_fields[num] = (u.op, u.dtype, tuple(track_uop(s) for s in u.src), arg, u.tag)+((u.metadata,) if TRACEMETA>=2 else ())
+  uop_fields[num] = (u.op, u.dtype, tuple(track_uop(s) for s in u.src), arg, u.tag)
   return num
 
 # *** tracking pattern matcher ***
@@ -1014,7 +993,7 @@ syms = { Ops.ADD: "+", Ops.SUB: "-", Ops.IDIV: "//", Ops.MOD: "%", Ops.SHL: "<<"
          Ops.MUL: "*", Ops.CMPLT: "<", Ops.CMPNE: "!=", Ops.AND: "&", Ops.OR: "|", Ops.XOR: "^"}
 renderer = PatternMatcher([
   (UPat((Ops.DEFINE_VAR, Ops.SPECIAL), name="x"), lambda x: UOp(Ops.NOOP, arg=x.arg[0])),
-  (UPat(Ops.RANGE, name="x"), lambda x: UOp(Ops.NOOP, arg=f"ridx{x.arg[0]}" if x.arg[0] >= 0 else f"ridxm{-x.arg[0]}")),
+  (UPat(Ops.RANGE, name="x"), lambda x: UOp(Ops.NOOP, arg=f"ridx{x.arg}")),
   (UPat((Ops.CONST, Ops.VCONST), name="x"), lambda x: UOp(Ops.NOOP, arg=str(x.arg))),
   (UPat(Ops.UNROLL, name="x"), lambda x: UOp(Ops.NOOP, arg=f"UNROLL({x.src[0].arg}, {x.arg})")),
   (UPat(Ops.CAST, name="x"), lambda x: UOp(Ops.NOOP, arg=f"({str(x.dtype)[7:]})({x.src[0].arg})")),

tinygrad/uop/spec.py

@@ -10,46 +10,29 @@ try:
   def z3_cdiv(a, b):return z3.If((a<0), z3.If(0<b, (a+(b-1))/b, (a-(b+1))/b), a/b)
   z3_alu: dict[Ops, Callable] = python_alu | {Ops.MOD: lambda a,b: a-z3_cdiv(a,b)*b, Ops.IDIV: z3_cdiv, Ops.SHR: lambda a,b: a/(2**b.as_long()),
     Ops.SHL: lambda a,b: a*(2**b.as_long()), Ops.AND: lambda a,b: a%(b+1) if isinstance(b, z3.ArithRef) else a&b, Ops.WHERE: z3.If,
-    Ops.MAX: lambda a,b: z3.If(a<b, b, a), Ops.TRUNC: lambda a: a if a.is_int() else z3.ToReal(z3.If(a >= 0, z3.ToInt(a), -z3.ToInt(-a)))}
+    Ops.MAX: lambda a,b: z3.If(a<b, b, a)}
   def create_bounded(name:str, vmin, vmax, solver:z3.Solver) -> z3.ArithRef:
     s = z3.Int(name, ctx=solver.ctx)
     solver.add(vmin <= s, s <= vmax)
     return s
 
   # ctx is (solver, load_number_dict)
-  # each uop gets rewritten to NOOP(arg=(solver, z3_object)), the arg has the solver first due to UOpMetaClass caching. z3 objects from different
-  # contexts can have the same hash but error on comparison
   z3_renderer = PatternMatcher([
     # Ops.SPECIAL can have symbolic arg but it wont be in the toposort beacuse its not a src, we need to add it manually
     (UPat(Ops.SPECIAL, src=(), name="x"), lambda x: UOp(Ops.SPECIAL, arg=x.arg[0], src=(x.ufix(x.arg[1]),))),
-    (UPat(Ops.SPECIAL, src=UPat(Ops.NOOP), name="x"), lambda x,ctx: UOp(Ops.NOOP, arg=(ctx[0],create_bounded(x.arg, 0, x.src[0].arg[1]-1, ctx[0])))),
-    (UPat(Ops.DEFINE_VAR, name="x"), lambda x,ctx: UOp(Ops.NOOP, arg=(ctx[0],create_bounded(x.arg[0], x.arg[1], x.arg[2], ctx[0])))),
-    (UPat(Ops.RANGE, name="x"), lambda x,ctx: UOp(Ops.NOOP, arg=(ctx[0],create_bounded(f"ridx{x.arg}", 0, x.src[0].arg[1]-1, ctx[0])))),
-    # float loads only become a variable when they get cast to int/bool
+    (UPat(Ops.SPECIAL, src=UPat(Ops.NOOP), name="x"), lambda x,ctx: UOp(Ops.NOOP, arg=create_bounded(x.arg, 0, x.src[0].arg-1, ctx[0]))),
+    (UPat(Ops.DEFINE_VAR, name="x"), lambda x,ctx: UOp(Ops.NOOP, arg=create_bounded(x.arg[0], x.arg[1], x.arg[2], ctx[0]))),
+    (UPat(Ops.RANGE, name="x"), lambda x,ctx: UOp(Ops.NOOP, arg=create_bounded(f"ridx{x.arg}", 0, x.src[0].arg-1, ctx[0]))),
     (UPat(Ops.LOAD, dtypes.ints, name="x"),
-      lambda x,ctx: UOp(Ops.NOOP, arg=(ctx[0],create_bounded(f"load{ctx[1].setdefault(x, len(ctx[1]))}", x.dtype.min, x.dtype.max, ctx[0])))),
-    (UPat(Ops.CONST, dtype=dtypes.ints+(dtypes.bool,), name="x"),
-      lambda x,ctx: UOp(Ops.NOOP, arg=(ctx[0],(z3.BoolVal if dtypes.is_bool(x.dtype) else z3.IntVal)(x.arg, ctx=ctx[0].ctx)))),
-    # z3 can cast from bool to int automatically
-    (UPat(Ops.CAST, dtype=dtypes.ints, src=UPat(Ops.NOOP), name="x"), lambda x: x.src[0]),
-    (UPat(Ops.CAST, dtype=dtypes.bool, src=UPat(Ops.NOOP), name="x"), lambda x,ctx: UOp(Ops.NOOP, arg=(ctx[0], x.src[0].arg[1]!=0))),
-    # if the source of the cast is not a noop it means that it is a float and so we create a new variable
-    (UPat(Ops.CAST, dtype=dtypes.ints, name="x"), lambda x,ctx:
-      UOp(Ops.NOOP, arg=(ctx[0], create_bounded(f"cast{ctx[1].setdefault(x, len(ctx[1]))}", x.dtype.min, x.dtype.max, ctx[0])))),
-    (UPat(Ops.CAST, dtype=dtypes.bool, name="x"), lambda x,ctx:
-      UOp(Ops.NOOP, arg=(ctx[0], z3.Bool(f"cast{ctx[1].setdefault(x, len(ctx[1]))}",ctx=ctx[0].ctx)))),
+      lambda x,ctx: UOp(Ops.NOOP, arg=create_bounded(f"load{ctx[1].setdefault(x, len(ctx[1]))}", x.vmin, x.vmax, ctx[0]))),
+    (UPat(Ops.CONST, name="x"), lambda x,ctx: UOp(Ops.NOOP, arg=(z3.BoolVal if dtypes.is_bool(x.dtype) else z3.IntVal)(x.arg, ctx=ctx[0].ctx))),
+    (UPat(Ops.CAST, dtype=dtypes.ints+(dtypes.bool,), src=UPat(Ops.NOOP), name="x"), lambda x: x.src[0]),
+    (UPat(Ops.CAST, name="x"), lambda x,ctx: UOp(Ops.NOOP, arg=create_bounded(f"cast{ctx[1].setdefault(x, len(ctx[1]))}", x.vmin, x.vmax, ctx[0]))),
     (UPat(Ops.XOR, src=UPat(Ops.NOOP), name="x"),
-      lambda x,ctx: UOp(Ops.NOOP, arg=(ctx[0], z3.BV2Int(z3_alu[x.op](*(z3.Int2BV(s.arg[1], x.dtype.itemsize*8) for s in x.src)))))),
-    (UPat(GroupOp.ALU, src=UPat(Ops.NOOP), name="x"), lambda x,ctx: UOp(Ops.NOOP, arg=(ctx[0], z3_alu[x.op](*(s.arg[1] for s in x.src))))),
-    # A comparison between floats introduces a new bool variable
-    (UPat(GroupOp.Comparison, src=UPat(dtype=dtypes.floats), name="x"), lambda x,ctx:
-      UOp(Ops.NOOP, arg=(ctx[0], z3.Bool(f"float_cmp{ctx[1].setdefault(x, len(ctx[1]))}",ctx=ctx[0].ctx)))),
+      lambda x: UOp(Ops.NOOP, arg=z3.BV2Int(z3_alu[x.op](*(z3.Int2BV(s.arg, x.dtype.itemsize*8) for s in x.src))))),
+    (UPat(GroupOp.ALU, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=z3_alu[x.op](*(s.arg for s in x.src)))),
   ])
 
-  def uops_to_z3(solver, *uops: UOp) -> 'list[z3.ExprRef]':
-    with Context(TRACK_MATCH_STATS=0):  # cant pickle z3 objects
-      return [s.arg[1] for s in graph_rewrite(uops[0].sink(*uops[1:]), z3_renderer, ctx=(solver, {})).src]
-
   z3_imported = True
 except (ImportError, AttributeError): z3_imported = False
 
@@ -107,10 +90,6 @@ tensor_uop_spec = buffer_spec+assign_spec+PatternMatcher([
   (UPat((Ops.DETACH, Ops.CONTIGUOUS, Ops.CONTIGUOUS_BACKWARD, Ops.FUSE), name="root", src=(UPat.var("x"),), arg=None),
    lambda root,x: root.dtype == x.dtype),
 
-  # CONTIGUOUS with a range
-  (UPat(Ops.CONTIGUOUS, name="root", src=(UPat.var("x"),), allow_any_len=True, arg=None),
-   lambda root,x: root.dtype == x.dtype and all(u.op is Ops.RANGE for u in root.src[1:])),
-
   # COPY/ALLREDUCE/MULTI
   (UPat(Ops.COPY, name="copy", src=(UPat.var("x"), UPat(Ops.DEVICE)), arg=None), lambda copy,x: copy.dtype == x.dtype),
   (UPat(Ops.ALLREDUCE, name="red", src=(UPat.var("x"), UPat(Ops.DEVICE))), lambda red,x: red.dtype == x.dtype and isinstance(red.arg, Ops)),
@@ -130,8 +109,9 @@ def validate_index(idx:UOp, gate:UOp=UOp.const(dtypes.bool, True)):
 
   if not z3_imported: raise ImportError("z3 is required for bounds checking, try IGNORE_OOB=0 or \"pip install z3-solver\"")
   solver = z3.Solver(ctx=z3.Context())
-  z3_idx, z3_mask = uops_to_z3(solver, idx.src[1], mask)
-  solver.add(z3_mask)
+  z3_sink = graph_rewrite(idx.src[1].sink(mask), z3_renderer, ctx=(solver, {}))
+  z3_idx = z3_sink.src[0].arg
+  solver.add(z3_sink.src[1].arg)
   if solver.check((z3_idx<0)|(sz<=z3_idx)) == z3.sat:
     print(f"idx={idx.src[1].render(simplify=False)}")
     print(f"mask & gate={mask.render(simplify=False)}")
@@ -156,7 +136,7 @@ spec = PatternMatcher([
   (UPat(Ops.DEFINE_REG, src=()), lambda: True),
   (UPat(Ops.DEFINE_VAR, name="x"), lambda x: isinstance(x.arg[1], int) and isinstance(x.arg[2], int)),
 
-  (UPat(Ops.RANGE, src=(UPat.var("x"),), name="rng"), lambda rng,x: rng.dtype == x.dtype and isinstance(rng.arg, tuple)),
+  (UPat(Ops.RANGE, src=(UPat.var("x"),), name="rng"), lambda rng,x: rng.dtype == x.dtype and isinstance(rng.arg, int)),
   (UPat(Ops.SPECIAL, src=()), lambda: True),
 
   (UPat(Ops.VIEW, dtypes.void, src=(), name="x"), lambda x: isinstance(x.arg, ShapeTracker)),
@@ -188,8 +168,8 @@ spec = PatternMatcher([
   (UPat(Ops.LOAD, src=(index_pat,), allow_any_len=True), validate_index),
 
   # STORE takes a <bufidx, val, gate?>
-  (UPat(Ops.STORE, src=(index_pat, UPat(name="val"), UPat(Ops.IF, name="gate")), allow_any_len=True), validate_store),
-  (UPat(Ops.STORE, src=(index_pat, UPat(name="val")), allow_any_len=True), validate_store),
+  (UPat(Ops.STORE, dtype=dtypes.void, src=(index_pat, UPat(name="val"), UPat(Ops.IF, name="gate")), allow_any_len=True), validate_store),
+  (UPat(Ops.STORE, dtype=dtypes.void, src=(index_pat, UPat(name="val")), allow_any_len=True), validate_store),
 
   # most ALUs have all matching dtypes, except CMPLT, CMPNE, and WHERE
   (UPat(Ops.WHERE, name="w", src=(UPat(dtype=dtypes.bool), UPat.var("x"), UPat.var("y"))), lambda w,x,y: w.dtype == x.dtype == y.dtype),

tinygrad/uop/symbolic.py

@@ -1,5 +1,5 @@
 # all of symbolic lives here now
-from typing import cast
+from typing import Any, cast
 import math, operator, struct, functools
 from collections import defaultdict
 from tinygrad.uop.ops import Ops, PatternMatcher, UPat, UOp, GroupOp, exec_alu
@@ -19,7 +19,7 @@ def simplify_pow(x:UOp, c:UOp) -> UOp|None:
 def fold_bitcast(root:UOp, c:UOp) -> UOp|None:
   if (from_fmt:=c.dtype.scalar().fmt) is None or (to_fmt:=root.dtype.scalar().fmt) is None: return None
   if c.dtype.itemsize != root.dtype.itemsize: return None
-  def convert(v:ConstType): return struct.unpack(to_fmt, struct.pack(from_fmt, v))[0]
+  def convert(v:Any): return struct.unpack(to_fmt, struct.pack(from_fmt, v))[0]
   return root.const_like(convert(c.arg) if root.dtype.count == 1 else tuple(map(convert, c.arg)))
 
 symbolic_simple = PatternMatcher([
@@ -183,9 +183,10 @@ def fold_divmod_congruence(d: UOp, x: UOp, y: UOp) -> UOp|None:
   terms, factors = zip(*[(u.divides(f:=u.const_factor()),f) for u in split_uop(x, Ops.ADD)])
   # a//c = (a-a%c)/c, if we can fold a%c, we can fold a//c
   rems = [min((r:=f%c), r-c, key=abs) for f in factors]
-  if (rem:=sum(r*v for r,v in zip(rems,terms))+const%c).vmin//c!=rem.vmax//c: return None
-  if d.op is Ops.MOD: return rem - rem.vmin//c*c
-  return sum((f-r)//c * v for f,r,v in zip(factors,rems,terms)) + (const-const%c+rem.vmin//c*c)//c
+  if (rem:=sum(r*v for r,v in zip(rems,terms))+const%c).vmin//c==rem.vmax//c and all(f > 0 for f in factors):
+    if d.op is Ops.MOD: return rem - rem.vmin//c*c
+    return sum((f-r)//c * v for f,r,v in zip(factors,rems,terms)) + (const-const%c+rem.vmin//c*c)//c
+  return None
 
 def divide_by_gcd(d: UOp, x: UOp, y: UOp) -> UOp|None:
   # x//y -> (x//gcd)//(y//gcd) or x%y -> gcd*(x//gcd)%(y//gcd)
@@ -279,7 +280,6 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
   ((UPat.var("y") + UPat.var("x") * UPat.cvar("c0")) + UPat.var("x") * UPat.cvar("c1"), lambda x,y,c0,c1: y+x*(c0+c1)),
   (UPat.var("x") + UPat.var("x") * UPat.cvar("c"), lambda x,c: x*(c+1)), # (x+x*c)-> x*(c+1)
   ((UPat.var("y") + UPat.var("x")) + UPat.var("x") * UPat.cvar("c"), lambda x,y,c: y+x*(c+1)),
-  ((UPat.var("y") + UPat.var("x") * UPat.cvar("c")) + UPat.var("x"), lambda x,y,c: y+x*(c+1)),
   (UPat.var("x") + UPat.var("x"), lambda x: x*2), # (x+x)-> x*2
   ((UPat.var("y") + UPat.var("x")) + UPat.var("x"), lambda y,x: y+x*2),
   ((UPat.var("x") / UPat.var("x2")) / UPat.var("x3"), lambda x,x2,x3: x/(x2*x3) if x2 is not x3 else None), # (x/x2)/x3 -> x/(x2*x3)
@@ -291,9 +291,6 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
   # alu of two where with same conds can combine, only do if true branch or false branch is const
   (UPat(GroupOp.Binary, name="alu", src=(UPat.var("c").where(UPat.var("t"), UPat.var("f")), UPat.var("c").where(UPat.var("tt"), UPat.var("ff")))), \
    lambda alu,c,t,tt,f,ff: c.where(t.alu(alu.op, tt), f.alu(alu.op, ff)) if t.op == tt.op == Ops.CONST or f.op == ff.op == Ops.CONST else None),
-  # if its a plus we add the associative variation too
-  ((UPat.var("y")+UPat.var("c").where(UPat.var("t"), UPat.var("f"))) + UPat.var("c").where(UPat.var("tt"), UPat.var("ff")), \
-   lambda y,c,t,tt,f,ff: y+c.where(t+tt, f+ff) if t.op == tt.op == Ops.CONST or f.op == ff.op == Ops.CONST else None),
   # ALU/variable min==max -> CONST (slow!)
   (UPat(GroupOp.ALU|{Ops.DEFINE_VAR, Ops.SPECIAL, Ops.RANGE}, name="x"), lambda x: x.const_like(x.vmin) if x.vmin == x.vmax else None),
   # max folding

tinygrad/viz/index.html

@@ -75,14 +75,9 @@
   g.tag circle {
     fill: #FFD700;
     stroke: #B8860B;
-  }
-  g.port circle {
-    fill: #b3dcc2;
-  }
-  g.tag circle, #edge-labels circle {
     stroke-width: 0.8;
   }
-  g.tag text, #edge-labels text {
+  g.tag text {
     text-anchor: middle;
     font-size: 6px;
     fill: #08090e;
@@ -90,30 +85,11 @@
   .label :is(text, p) {
     font-weight: 350;
   }
-  rect.node {
-    stroke-width: 1.4;
-    stroke: #4a4b57;
-  }
-  rect.overlay {
-    fill: rgba(26, 27, 38, 0.5);
-  }
   .edgePath {
     stroke: #4a4b57;
     fill: none;
     stroke-width: 1.4px;
   }
-  .highlight rect, .edgePath.highlight, g.port circle {
-    stroke: #89C9A2;
-  }
-  #edge-labels g.port.highlight {
-    display: block
-  }
-  #edge-labels g.port {
-    display: none
-  }
-  #arrowhead {
-    fill: #4a4b57;
-  }
   .main-container {
     display: flex;
     width: 100%;
@@ -355,7 +331,7 @@
         </g>
         <defs>
           <marker id="arrowhead" viewBox="0 -5 10 10" refX="10" refY="0" markerWidth="6" markerHeight="6" orient="auto">
-            <path d="M0,-5L10,0L0,5" fill="context-stroke"></path>
+            <path d="M0,-5L10,0L0,5" fill="#4a4b57"></path>
           </marker>
         </defs>
       </svg>

tinygrad/viz/js/index.js

@@ -4,15 +4,6 @@ const displayGraph = (cls) => {
   for (const e of document.getElementsByClassName("view")) e.style.display = e.classList.contains(cls) ? "flex" : "none";
 }
 
-const darkenHex = (h, p = 0) =>
-  `#${(
-    c = parseInt(h.slice(1), 16),
-    f = 1 - p / 100,
-    ((c >> 16 & 255) * f | 0) << 16 |
-    ((c >>  8 & 255) * f | 0) <<  8 |
-    ((c       & 255) * f | 0)
-  ).toString(16).padStart(6, '0')}`;
-
 const ANSI_COLORS = ["#b3b3b3", "#ff6666", "#66b366", "#ffff66", "#6666ff", "#ff66ff", "#66ffff", "#ffffff"];
 const parseColors = (name, defaultColor="#ffffff") => Array.from(name.matchAll(/(?:\u001b\[(\d+)m([\s\S]*?)\u001b\[0m)|([^\u001b]+)/g),
   ([_, code, colored_st, st]) => ({ st: colored_st ?? st, color: code != null ? ANSI_COLORS[(parseInt(code)-30+60)%60] : defaultColor }));
@@ -65,23 +56,11 @@ async function renderDag(graph, additions, recenter=false) {
     const g = dagre.graphlib.json.read(e.data);
     // draw nodes
     const STROKE_WIDTH = 1.4;
-    d3.select("#graph-svg").on("click", () => d3.selectAll(".highlight").classed("highlight", false));
     const nodes = d3.select("#nodes").selectAll("g").data(g.nodes().map(id => g.node(id)), d => d).join("g")
-      .attr("transform", d => `translate(${d.x},${d.y})`).classed("clickable", d => d.ref != null).on("click", (e,d) => {
-        if (d.ref != null) return setCtxWithHistory(d.ref);
-        const parents = g.predecessors(d.id);
-        if (parents == null) return;
-        const src = [...parents, d.id];
-        nodes.classed("highlight", n => src.includes(n.id));
-        d3.select("#edges").selectAll("path.edgePath").classed("highlight", e => src.includes(e.v) && e.w===d.id);
-        d3.select("#edge-labels").selectAll("g.port").classed("highlight", (_, i, nodes) => {
-          const [v, w] = nodes[i].id.split("-");
-          return src.includes(v) && w===d.id;
-        });
-        e.stopPropagation();
-      });
+      .attr("transform", d => `translate(${d.x},${d.y})`).classed("clickable", d => d.ref != null)
+      .on("click", (_,d) => setCtxWithHistory(d.ref));
     nodes.selectAll("rect").data(d => [d]).join("rect").attr("width", d => d.width).attr("height", d => d.height).attr("fill", d => d.color)
-      .attr("x", d => -d.width/2).attr("y", d => -d.height/2).attr("class", d => d.className ?? "node");
+      .attr("x", d => -d.width/2).attr("y", d => -d.height/2).attr("style", d => d.style ?? `stroke:#4a4b57; stroke-width:${STROKE_WIDTH}px;`);
     nodes.selectAll("g.label").data(d => [d]).join("g").attr("class", "label").attr("transform", d => {
       const x = (d.width-d.padding*2)/2;
       const y = (d.height-d.padding*2)/2+STROKE_WIDTH;
@@ -96,19 +75,19 @@ async function renderDag(graph, additions, recenter=false) {
       }
       return [ret];
     }).join("text").selectAll("tspan").data(d => d).join("tspan").attr("x", "0").attr("dy", 14).selectAll("tspan").data(d => d).join("tspan")
-      .attr("fill", d => darkenHex(d.color, 25)).text(d => d.st).attr("xml:space", "preserve");
+      .attr("fill", d => d.color).text(d => d.st).attr("xml:space", "preserve");
     addTags(nodes.selectAll("g.tag").data(d => d.tag != null ? [d] : []).join("g").attr("class", "tag")
       .attr("transform", d => `translate(${-d.width/2+8}, ${-d.height/2+8})`).datum(e => e.tag));
     // draw edges
-    const line = d3.line().x(d => d.x).y(d => d.y).curve(d3.curveBasis), edges = g.edges();
-    d3.select("#edges").selectAll("path.edgePath").data(edges).join("path").attr("class", "edgePath").attr("d", (e) => {
+    const line = d3.line().x(d => d.x).y(d => d.y).curve(d3.curveBasis);
+    d3.select("#edges").selectAll("path.edgePath").data(g.edges()).join("path").attr("class", "edgePath").attr("d", (e) => {
       const edge = g.edge(e);
       const points = edge.points.slice(1, edge.points.length-1);
       points.unshift(intersectRect(g.node(e.v), points[0]));
       points.push(intersectRect(g.node(e.w), points[points.length-1]));
       return line(points);
     }).attr("marker-end", "url(#arrowhead)");
-    addTags(d3.select("#edge-labels").selectAll("g").data(edges).join("g").attr("transform", (e) => {
+    addTags(d3.select("#edge-labels").selectAll("g").data(g.edges().filter(e => g.edge(e).label != null)).join("g").attr("transform", (e) => {
       // get a point near the end
       const [p1, p2] = g.edge(e).points.slice(-2);
       const dx = p2.x-p1.x;
@@ -122,7 +101,7 @@ async function renderDag(graph, additions, recenter=false) {
       const x = p2.x - ux * offset;
       const y = p2.y - uy * offset;
       return `translate(${x}, ${y})`
-    }).attr("class", e => g.edge(e).label.type).attr("id", e => `${e.v}-${e.w}`).datum(e => g.edge(e).label.text));
+    }).attr("class", "tag").datum(e => g.edge(e).label));
     if (recenter) document.getElementById("zoom-to-fit-btn").click();
   };
 
@@ -172,17 +151,7 @@ async function renderProfiler() {
   // layout once!
   if (data != null) return;
   const profiler = d3.select(".profiler").html("");
-  const buf = await (await fetch("/get_profile")).arrayBuffer();
-  const view = new DataView(buf);
-  let offset = 0;
-  const u8 = () => { const ret = view.getUint8(offset); offset += 1; return ret; }
-  const u32 = () => { const ret = view.getUint32(offset, true); offset += 4; return ret; }
-  const u64 = () => { const ret = new Number(view.getBigUint64(offset, true)); offset += 8; return ret; }
-  const f32 = () => { const ret = view.getFloat32(offset, true); offset += 4; return ret; }
-  const optional = (i) => i === 0 ? null : i-1;
-  const dur = u32(), peak = u64(), indexLen = u32(), layoutsLen = u32();
-  const textDecoder = new TextDecoder("utf-8");
-  const { strings, dtypeSize }  = JSON.parse(textDecoder.decode(new Uint8Array(buf, offset, indexLen))); offset += indexLen;
+  const { layout, st, et } = await (await fetch("/get_profile")).json();
   // place devices on the y axis and set vertical positions
   const [tickSize, padding] = [10, 8];
   const deviceList = profiler.append("div").attr("id", "device-list").style("padding-top", tickSize+padding+"px");
@@ -193,34 +162,22 @@ async function renderProfiler() {
   const canvasTop = rect(canvas).top;
   // color by key (name/category/device)
   const colorMap = new Map();
-  data = {tracks:new Map(), axes:{}};
-  const heightScale = d3.scaleLinear().domain([0, peak]).range([4,maxheight=100]);
-  for (let i=0; i<layoutsLen; i++) {
-    const nameLen = view.getUint8(offset, true); offset += 1;
-    const k = textDecoder.decode(new Uint8Array(buf, offset, nameLen)); offset += nameLen;
+  data = {tracks:new Map(), axes:{}, st, et};
+  const heightScale = d3.scaleLinear().domain([0, Object.entries(layout).reduce((peak, [_,d]) => Math.max(peak, d.peak||0), 0)]).range([4,maxheight=100]);
+  for (const [k, v] of Object.entries(layout)) {
+    if (v.shapes.length === 0) continue;
     const div = deviceList.append("div").attr("id", k).text(k).style("padding", padding+"px");
     const { y:baseY, height:baseHeight } = rect(div.node());
     const offsetY = baseY-canvasTop+padding/2;
-    const shapes = [];
-    const EventTypes = {TIMELINE:0, MEMORY:1};
-    const eventType = u8(), eventsLen = u32();
-    if (eventType === EventTypes.TIMELINE) {
+    if (v.shapes[0].dur != null) {
       const levelHeight = baseHeight-padding;
-      const levels = [];
+      const shapes = [];
       data.tracks.set(k, { shapes, offsetY });
       let colorKey, ref;
-      for (let j=0; j<eventsLen; j++) {
-        const e = {name:strings[u32()], ref:optional(u32()), st:u32(), dur:f32(), cat:optional(u8()), info:strings[u32()] || null};
-        // find a free level to put the event
-        let depth = levels.findIndex(levelEt => e.st >= levelEt);
-        const et = e.st+Math.trunc(e.dur);
-        if (depth === -1) {
-          depth = levels.length;
-          levels.push(et);
-        } else levels[depth] = et;
-        if (depth === 0) colorKey = e.cat ?? e.name;
+      for (const e of v.shapes) {
+        if (e.depth === 0) colorKey = e.cat ?? e.name;
         if (!colorMap.has(colorKey)) colorMap.set(colorKey, cycleColors(colorScheme[k] ?? colorScheme.DEFAULT, colorMap.size));
-        const fillColor = d3.color(colorMap.get(colorKey)).brighter(depth).toString();
+        const fillColor = d3.color(colorMap.get(colorKey)).brighter(e.depth).toString();
         const label = parseColors(e.name).map(({ color, st }) => ({ color, st, width:ctx.measureText(st).width }));
         if (e.ref != null) ref = {ctx:e.ref, step:0};
         else if (ref != null) {
@@ -230,49 +187,19 @@ async function renderProfiler() {
         }
         const arg = { tooltipText:formatTime(e.dur)+(e.info != null ? "\n"+e.info : ""), ...ref };
         // offset y by depth
-        shapes.push({x:e.st, y:levelHeight*depth, width:e.dur, height:levelHeight, arg, label, fillColor });
+        shapes.push({x:e.st-st, y:levelHeight*e.depth, width:e.dur, height:levelHeight, arg, label, fillColor });
       }
-      div.style("height", levelHeight*levels.length+padding+"px").style("pointerEvents", "none");
+      div.style("height", levelHeight*v.maxDepth+padding+"px").style("pointerEvents", "none");
     } else {
-      const peak = u64();
-      const height = heightScale(peak);
-      const yscale = d3.scaleLinear().domain([0, peak]).range([height, 0]);
-      let x = 0, y = 0;
-      const buf_shapes = new Map(), temp = new Map();
-      const timestamps = [];
-      for (let j=0; j<eventsLen; j++) {
-        const alloc = u8(), ts = u32(), key = u32();
-        if (alloc) {
-          const dtype = strings[u32()], sz = u64(), nbytes = dtypeSize[dtype]*sz;
-          const shape = {x:[x], y:[y], dtype, sz, nbytes, key};
-          buf_shapes.set(key, shape); temp.set(key, shape);
-          timestamps.push(ts);
-          x += 1; y += nbytes;
-        } else {
-          const free = buf_shapes.get(key);
-          timestamps.push(ts);
-          x += 1; y -= free.nbytes;
-          free.x.push(x);
-          free.y.push(free.y.at(-1));
-          temp.delete(key);
-          for (const [k, v] of temp) {
-            if (k <= key) continue;
-            v.x.push(x, x);
-            v.y.push(v.y.at(-1), v.y.at(-1)-free.nbytes);
-          }
-        }
+      const height = heightScale(v.peak);
+      const yscale = d3.scaleLinear().domain([0, v.peak]).range([height, 0]);
+      const shapes = [];
+      for (const [i,e] of v.shapes.entries()) {
+        const x = e.x.map(tsIdx => v.timestamps[tsIdx]-st);
+        const arg = {tooltipText:`${e.arg.dtype} len:${formatUnit(e.arg.sz)}\n${formatUnit(e.arg.nbytes, "B")}`};
+        shapes.push({ x, y0:e.y.map(yscale), y1:e.y.map(y => yscale(y+e.arg.nbytes)), arg, fillColor:cycleColors(colorScheme.BUFFER, i) });
       }
-      for (const [_, v] of temp) {
-        v.x.push(x);
-        v.y.push(v.y.at(-1));
-      }
-      timestamps.push(dur);
-      for (const [_, {dtype, sz, nbytes, y, x:steps}] of buf_shapes) {
-        const x = steps.map(s => timestamps[s]);
-        const arg = {tooltipText:`${dtype} len:${formatUnit(sz)}\n${formatUnit(nbytes, "B")}`};
-        shapes.push({ x, y0:y.map(yscale), y1:y.map(y0 => yscale(y0+nbytes)), arg, fillColor:cycleColors(colorScheme.BUFFER, shapes.length) });
-      }
-      data.tracks.set(k, { shapes, offsetY, height, peak, scaleFactor:maxheight*4/height });
+      data.tracks.set(k, { shapes, offsetY, height, peak:v.peak, scaleFactor:maxheight*4/height });
       div.style("height", height+padding+"px").style("cursor", "pointer").on("click", (e) => {
         const newFocus = e.currentTarget.id === focusedDevice ? null : e.currentTarget.id;
         let offset = 0;
@@ -298,7 +225,7 @@ async function renderProfiler() {
     ctx.save();
     ctx.clearRect(0, 0, canvas.clientWidth, canvas.clientHeight);
     // rescale to match current zoom
-    const xscale = d3.scaleLinear().domain([0, dur]).range([0, canvas.clientWidth]);
+    const xscale = d3.scaleLinear().domain([0, et-st]).range([0, canvas.clientWidth]);
     xscale.domain(xscale.range().map(zoomLevel.invertX, zoomLevel).map(xscale.invert, xscale));
     const zoomDomain = transform != null ? xscale.domain() : null;
     let yscale = null;
@@ -362,7 +289,7 @@ async function renderProfiler() {
       // tick label
       ctx.textBaseline = "top";
       ctx.textAlign = "left";
-      ctx.fillText(formatTime(tick, dur), x+ctx.lineWidth+2, tickSize);
+      ctx.fillText(formatTime(tick, et-st), x+ctx.lineWidth+2, tickSize);
     }
     if (yscale != null) {
       drawLine(ctx, [0, 0], yscale.range());
@@ -390,7 +317,8 @@ async function renderProfiler() {
     d3.select(canvas).call(canvasZoom.transform, zoomLevel);
   }
 
-  canvasZoom = d3.zoom().filter(vizZoomFilter).scaleExtent([1, Infinity]).translateExtent([[0,0], [Infinity,0]]).on("zoom", e => render(e.transform));
+  canvasZoom = d3.zoom().filter(e => (!e.ctrlKey || e.type === 'wheel' || e.type === 'mousedown') && !e.button)
+    .scaleExtent([1, Infinity]).translateExtent([[0,0], [Infinity,0]]).on("zoom", e => render(e.transform));
   d3.select(canvas).call(canvasZoom);
   document.addEventListener("contextmenu", e => e.ctrlKey && e.preventDefault());
 
@@ -427,8 +355,7 @@ async function renderProfiler() {
 
 // ** zoom and recentering
 
-const vizZoomFilter = e => (!e.ctrlKey || e.type === 'wheel' || e.type === 'mousedown') && !e.button && e.type !== 'dblclick';
-const svgZoom = d3.zoom().filter(vizZoomFilter).on("zoom", (e) => d3.select("#render").attr("transform", e.transform));
+const svgZoom = d3.zoom().on("zoom", (e) => d3.select("#render").attr("transform", e.transform));
 d3.select("#graph-svg").call(svgZoom);
 
 // zoom to fit into view
@@ -532,6 +459,7 @@ function setState(ns) {
 
 // set a new context and keep the old one in browser history
 function setCtxWithHistory(newCtx, step=0) {
+  if (newCtx == null) return;
   // NOTE: browser does a structured clone, passing a mutable object is safe.
   history.replaceState(state, "");
   history.pushState(state, "");

tinygrad/viz/js/worker.js

@@ -8,7 +8,7 @@ onmessage = (e) => {
   const { graph, additions, ctxs } = e.data;
   const g = new dagre.graphlib.Graph({ compound: true });
   g.setGraph({ rankdir: "LR" }).setDefaultEdgeLabel(function() { return {}; });
-  if (additions.length !== 0) g.setNode("addition", {label:"", className:"overlay", padding:0});
+  if (additions.length !== 0) g.setNode("addition", {label:"", style:"fill: rgba(26, 27, 38, 0.5);", padding:0});
   for (let [k, {label, src, ref, ...rest }] of Object.entries(graph)) {
     // adjust node dims by label size (excluding escape codes) + add padding
     let [width, height] = [0, 0];
@@ -16,11 +16,11 @@ onmessage = (e) => {
       width = Math.max(width, ctx.measureText(line).width);
       height += LINE_HEIGHT;
     }
-    g.setNode(k, {width:width+NODE_PADDING*2, height:height+NODE_PADDING*2, padding:NODE_PADDING, label, ref, id:k, ...rest});
+    g.setNode(k, {width:width+NODE_PADDING*2, height:height+NODE_PADDING*2, padding:NODE_PADDING, label, ref, ...rest});
     // add edges
     const edgeCounts = {}
-    for (const [_, s] of src) edgeCounts[s] = (edgeCounts[s] || 0)+1;
-    for (const [port, s] of src) g.setEdge(s, k, { label: edgeCounts[s] > 1 ? {type:"tag", text:edgeCounts[s]} : {type:"port", text:port}});
+    for (const s of src) edgeCounts[s] = (edgeCounts[s] || 0)+1;
+    for (const s of src) g.setEdge(s, k, { label: edgeCounts[s] > 1 ? edgeCounts[s] : null });
     if (additions.includes(parseInt(k))) g.setParent(k, "addition");
   }
   dagre.layout(g);

tinygrad/viz/serve.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-import multiprocessing, pickle, difflib, os, threading, json, time, sys, webbrowser, socket, argparse, socketserver, functools, codecs, io, struct
+import multiprocessing, pickle, difflib, os, threading, json, time, sys, webbrowser, socket, argparse, socketserver, functools, codecs, io
 import subprocess, ctypes
 from contextlib import redirect_stdout
 from decimal import Decimal
@@ -7,11 +7,10 @@ from http.server import BaseHTTPRequestHandler
 from urllib.parse import parse_qs, urlparse
 from typing import Any, TypedDict, Generator
 from tinygrad.helpers import colored, getenv, tqdm, unwrap, word_wrap, TRACEMETA, ProfileEvent, ProfileRangeEvent, TracingKey, ProfilePointEvent
-from tinygrad.uop.ops import TrackedGraphRewrite, UOp, Ops, printable, GroupOp, srender, sint, sym_infer
+from tinygrad.uop.ops import TrackedGraphRewrite, UOp, Ops, printable, GroupOp, srender, sint
 from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry, Device
 from tinygrad.renderer import ProgramSpec
 from tinygrad.dtype import dtypes
-from tinygrad.codegen.opt.kernel import axis_colors
 
 uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0", Ops.VCONST: "#e0e0e0", Ops.REDUCE: "#FF5B5B",
                Ops.DEFINE_GLOBAL: "#ffe0b0", Ops.DEFINE_LOCAL: "#ffe0d0", Ops.DEFINE_REG: "#f0ffe0", Ops.REDUCE_AXIS: "#FF6B6B",
@@ -80,21 +79,21 @@ def uop_to_json(x:UOp) -> dict[int, dict]:
       if u.op not in {Ops.VIEW, Ops.BUFFER, Ops.KERNEL, Ops.ASSIGN, Ops.COPY, Ops.SINK, *GroupOp.Buffer} and u.st is not None:
         label += f"\n{shape_to_str(u.shape)}"
       elif len(rngs:=u.ranges):
-        label += f"\n({','.join([colored(str(x.arg[0]), axis_colors[x.arg[-1]]) for x in sorted(rngs, key=lambda x: x.arg[0:-1])])})"
+        label += f"\n{str(sorted([x.arg for x in rngs]))}"
     except Exception:
       label += "\n<ISSUE GETTING LABEL>"
     if (ref:=ref_map.get(u.arg.ast) if u.op is Ops.KERNEL else None) is not None: label += f"\ncodegen@{ctxs[ref]['name']}"
     # NOTE: kernel already has metadata in arg
     if TRACEMETA >= 2 and u.metadata is not None and u.op is not Ops.KERNEL: label += "\n"+repr(u.metadata)
-    graph[id(u)] = {"label":label, "src":[(i,id(x)) for i,x in enumerate(u.src) if x not in excluded], "color":uops_colors.get(u.op, "#ffffff"),
+    graph[id(u)] = {"label":label, "src":[id(x) for x in u.src if x not in excluded], "color":uops_colors.get(u.op, "#ffffff"),
                     "ref":ref, "tag":u.tag}
   return graph
 
 @functools.cache
 def _reconstruct(a:int):
-  op, dtype, src, arg, *rest = contexts[2][a]
+  op, dtype, src, arg, tag = contexts[2][a]
   arg = type(arg)(_reconstruct(arg.ast), arg.metadata) if op is Ops.KERNEL else arg
-  return UOp(op, dtype, tuple(_reconstruct(s) for s in src), arg, *rest)
+  return UOp(op, dtype, tuple(_reconstruct(s) for s in src), arg, tag)
 
 def get_details(ctx:TrackedGraphRewrite) -> Generator[GraphRewriteDetails, None, None]:
   yield {"graph":uop_to_json(next_sink:=_reconstruct(ctx.sink)), "uop":str(next_sink), "changed_nodes":None, "diff":None, "upat":None}
@@ -107,15 +106,6 @@ def get_details(ctx:TrackedGraphRewrite) -> Generator[GraphRewriteDetails, None,
            "diff":list(difflib.unified_diff(str(u0).splitlines(), str(u1).splitlines())), "upat":(upat_loc, printable(upat_loc))}
     if not ctx.bottom_up: next_sink = new_sink
 
-# encoder helpers
-
-def enum_str(s, cache:dict[str, int]) -> int:
-  if (cret:=cache.get(s)) is not None: return cret
-  cache[s] = ret = len(cache)
-  return ret
-
-def option(s:int|None) -> int: return 0 if s is None else s+1
-
 # Profiler API
 
 device_ts_diffs:dict[str, tuple[Decimal, Decimal]] = {}
@@ -132,45 +122,56 @@ def flatten_events(profile:list[ProfileEvent]) -> Generator[tuple[Decimal, Decim
       yield (st:=min(cpu_ts)), (et:=max(cpu_ts)), ProfileRangeEvent(f"{e.ents[0].device.split(':')[0]} Graph", f"batched {len(e.ents)}", st, et)
       for i,ent in enumerate(e.ents): yield (cpu_ts[i*2], cpu_ts[i*2+1], ent)
 
-# normalize event timestamps and attach kernel metadata
-def timeline_layout(dev_events:list[tuple[int, int, float, DevEvent]], start_ts:int, scache:dict[str, int]) -> bytes|None:
-  events:list[bytes] = []
-  exec_points:dict[str, dict] = {}
-  category_enum:dict[str, int] = {}
-  for st,et,dur,e in dev_events:
-    if isinstance(e, ProfilePointEvent) and e.name == "exec": exec_points[e.key] = e.arg
+# timeline layout stacks events in a contiguous block. When a late starter finishes late, there is whitespace in the higher levels.
+def timeline_layout(events:list[tuple[int, int, float, DevEvent]]) -> dict:
+  shapes:list[dict] = []
+  levels:list[int] = []
+  for st,et,dur,e in events:
     if dur == 0: continue
+    # find a free level to put the event
+    depth = next((i for i,level_et in enumerate(levels) if st>=level_et), len(levels))
+    if depth < len(levels): levels[depth] = et
+    else: levels.append(et)
     name, cat, info = e.name, None, None
     if (ref:=ref_map.get(name)) is not None:
       name = ctxs[ref]["name"]
-      if isinstance(p:=contexts[0][ref].ret, ProgramSpec) and (ei:=exec_points.get(p.name)) is not None:
-        info = f"{sym_infer(p.estimates.ops, ei['var_vals'])/(t:=dur*1e3):.2f} GFLOPS {sym_infer(p.estimates.mem, ei['var_vals'])/t:4.1f}"+ \
-               f"|{sym_infer(p.estimates.lds,ei['var_vals'])/t:.1f} GB/s\n{ei['metadata']}"
+      # TODO: support symbolic by capturing var_vals in profile events
+      if isinstance(p:=contexts[0][ref].ret, ProgramSpec) and all(isinstance(es,int) for es in [p.estimates.ops, p.estimates.mem, p.estimates.lds]):
+        info = f"{p.estimates.ops/(t:=dur*1e3):.2f} GFLOPS {p.estimates.mem/t:4.1f}|{p.estimates.lds/t:.1f} GB/s"
     elif isinstance(e.name, TracingKey):
       name, cat = e.name.display_name, e.name.cat
       ref = next((v for k in e.name.keys if (v:=ref_map.get(k)) is not None), None)
-    events.append(struct.pack("<IIIfBI", enum_str(name, scache), option(ref), st-start_ts, dur,
-                              option(None if cat is None else enum_str(cat, category_enum)), enum_str(info or "", scache)))
-  return struct.pack("<BI", 0, len(events))+b"".join(events) if events else None
+    shapes.append({"name":name, "ref":ref, "st":st, "dur":dur, "depth":depth, "cat":cat, "info":info})
+  return {"shapes":shapes, "maxDepth":len(levels)}
 
-def mem_layout(events:list[tuple[int, int, float, DevEvent]], start_ts:int, end_ts:int, peaks:list[int], dtype_size:dict[str, int],
-               scache:dict[str, int]) -> bytes|None:
-  peak, mem = 0, 0
-  temp:dict[int, int] = {}
-  bufs:list[bytes] = []
+def mem_layout(events:list[tuple[int, int, float, DevEvent]], max_ts:int) -> dict:
+  step, peak, mem = 0, 0, 0
+  shps:dict[int, dict] = {}
+  temp:dict[int, dict] = {}
+  timestamps:list[int] = []
   for st,_,_,e in events:
     if not isinstance(e, ProfilePointEvent): continue
     if e.name == "alloc":
-      bufs.append(struct.pack("<BIIIQ", 1, int(e.ts)-start_ts, e.key, enum_str(e.arg["dtype"].name, scache), e.arg["sz"]))
-      dtype_size.setdefault(e.arg["dtype"].name, e.arg["dtype"].itemsize)
-      temp[e.key] = nbytes = e.arg["sz"]*e.arg["dtype"].itemsize
-      mem += nbytes
+      shps[e.key] = temp[e.key] = {"x":[step], "y":[mem], "arg":e.arg}
+      timestamps.append(int(e.ts))
+      step += 1
+      mem += e.arg["nbytes"]
       if mem > peak: peak = mem
     if e.name == "free":
-      bufs.append(struct.pack("<BII", 0, int(e.ts)-start_ts, e.key))
-      mem -= temp.pop(e.key)
-  peaks.append(peak)
-  return struct.pack("<BIQ", 1, len(bufs), peak)+b"".join(bufs) if bufs else None
+      timestamps.append(int(e.ts))
+      step += 1
+      mem -= (removed:=temp.pop(e.key))["arg"]["nbytes"]
+      removed["x"].append(step)
+      removed["y"].append(removed["y"][-1])
+      for k,v in temp.items():
+        if k > e.key:
+          v["x"] += [step, step]
+          v["y"] += [v["y"][-1], v["y"][-1]-removed["arg"]["nbytes"]]
+  for v in temp.values():
+    v["x"].append(step)
+    v["y"].append(v["y"][-1])
+  timestamps.append(max_ts)
+  return {"shapes":list(shps.values()), "peak":peak, "timestamps":timestamps}
 
 def get_profile(profile:list[ProfileEvent]) -> bytes|None:
   # start by getting the time diffs
@@ -178,25 +179,20 @@ def get_profile(profile:list[ProfileEvent]) -> bytes|None:
     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)
   # map events per device
   dev_events:dict[str, list[tuple[int, int, float, DevEvent]]] = {}
-  start_ts:int|None = None
-  end_ts:int|None = None
+  min_ts:int|None = None
+  max_ts:int|None = None
   for ts,en,e in flatten_events(profile):
     dev_events.setdefault(e.device,[]).append((st:=int(ts), et:=int(en), float(en-ts), e))
-    if start_ts is None or st < start_ts: start_ts = st
-    if end_ts is None or et > end_ts: end_ts = et
-  if start_ts is None: return None
+    if min_ts is None or st < min_ts: min_ts = st
+    if max_ts is None or et > max_ts: max_ts = et
+  if min_ts is None: return None
   # return layout of per device events
-  layout:dict[str, bytes|None] = {}
-  scache:dict[str, int] = {}
-  peaks:list[int] = []
-  dtype_size:dict[str, int] = {}
+  layout:dict[str, dict] = {}
   for k,v in dev_events.items():
     v.sort(key=lambda e:e[0])
-    layout[k] = timeline_layout(v, start_ts, scache)
-    layout[f"{k} Memory"] = mem_layout(v, start_ts, unwrap(end_ts), peaks, dtype_size, scache)
-  ret = [b"".join([struct.pack("<B", len(k)), k.encode(), v]) for k,v in layout.items() if v is not None]
-  index = json.dumps({"strings":list(scache), "dtypeSize":dtype_size}).encode()
-  return struct.pack("<IQII", unwrap(end_ts)-start_ts, max(peaks,default=0), len(index), len(ret))+index+b"".join(ret)
+    layout[k] = timeline_layout(v)
+    layout[f"{k} Memory"] = mem_layout(v, unwrap(max_ts))
+  return json.dumps({"layout":layout, "st":min_ts, "et":max_ts}).encode("utf-8")
 
 def get_runtime_stats(key) -> list[dict]:
   ret:list[dict] = []
@@ -258,7 +254,7 @@ class Handler(BaseHTTPRequestHandler):
       if url.path == "/disasm": ret, content_type = get_disassembly(**query), "application/json"
       else: return self.stream_json(get_details(contexts[1][int(query["ctx"][0])][int(query["idx"][0])]))
     elif url.path == "/ctxs": ret, content_type = json.dumps(ctxs).encode(), "application/json"
-    elif url.path == "/get_profile" and profile_ret: ret, content_type = profile_ret, "application/octet-stream"
+    elif url.path == "/get_profile" and profile_ret is not None: ret, content_type = profile_ret, "application/json"
     else: status_code = 404
 
     # send response
@@ -291,8 +287,8 @@ def reloader():
       os.execv(sys.executable, [sys.executable] + sys.argv)
     time.sleep(0.1)
 
-def load_pickle(path:str|None) -> list:
-  if path is None or not os.path.exists(path): return []
+def load_pickle(path:str):
+  if path is None or not os.path.exists(path): return None
   with open(path, "rb") as f: return pickle.load(f)
 
 # NOTE: using HTTPServer forces a potentially slow socket.getfqdn
@@ -315,16 +311,16 @@ if __name__ == "__main__":
   contexts, profile = load_pickle(args.kernels), load_pickle(args.profile)
 
   # NOTE: this context is a tuple of list[keys] and list[values]
-  ctxs = get_metadata(*contexts[:2]) if contexts else []
+  ctxs = get_metadata(*contexts[:2]) if contexts is not None else []
 
-  profile_ret = get_profile(profile)
+  profile_ret = get_profile(profile) if profile is not None else None
 
   server = TCPServerWithReuse(('', PORT), Handler)
   reloader_thread = threading.Thread(target=reloader)
   reloader_thread.start()
   print(f"*** started viz on {HOST}:{PORT}")
   print(colored(f"*** ready in {(time.perf_counter()-st)*1e3:4.2f}ms", "green"), flush=True)
-  if len(getenv("BROWSER", "")) > 0: webbrowser.open(f"{HOST}:{PORT}")
+  if len(getenv("BROWSER", "")) > 0: webbrowser.open(f"{HOST}:{PORT}{'/profiler' if contexts is None else ''}")
   try: server.serve_forever()
   except KeyboardInterrupt:
     print("*** viz is shutting down...")