Merge branch 'master' into mac_pytest

.github/workflows/benchmark.yml

@@ -44,7 +44,7 @@ jobs:
     - name: Run pytest -nauto
       run: |
         source /tmp/tinygrad_pytest_ci/bin/activate
-        pytest -nauto --durations=20
+        pytest -nauto
 
   testmacbenchmark:
     name: Mac Benchmark

.github/workflows/test.yml

@@ -266,19 +266,17 @@ jobs:
       run: python -c "from tinygrad import Device; assert Device.DEFAULT == 'CPU', Device.DEFAULT"
     - name: Run unit tests
       run: |
-        CPU=1 python test/null/test_device.py TestRunAsModule.test_module_runs
-        CPU=1 python -m pytest -n=auto test/unit/ --durations=20
-    - name: Run NULL backend tests
-      run: NULL=1 python -m pytest -n=auto test/null/ --durations=20
+        CPU=1 python test/unit/test_device.py TestRunAsModule.test_module_runs
+        CPU=1 python -m pytest -n=auto test/unit/ --durations=20 --deselect=test/unit/test_device.py::TestRunAsModule::test_module_runs
     - name: Run targetted tests on NULL backend
-      run: NULL=1 python3 -m unittest test.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step
+      run: NULL=1 python3 -m unittest test.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step test/device/test_null.py
     # TODO: too slow
     # - name: Run SDXL on NULL backend
     #   run: NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
     - name: Run Clip tests for SD MLPerf on NULL backend
       run: NULL=1 python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
     - name: Run AMD emulated BERT training on NULL backend
-      run: EMULATE=AMD_RDNA4 NULL=1 NULL_ALLOW_COPYOUT=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+      run: EMULATE=AMD_RDNA4 NULL=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
     # TODO: support fake weights
     #- name: Run LLaMA 7B on 4 fake devices
     #  run: NULL=1 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
@@ -316,7 +314,7 @@ jobs:
         deps: testing_unit
         python-version: '3.14'
     - name: Test SPEC=2
-      run: SPEC=2 pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore=test/null --ignore test/test_custom_kernel.py --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
+      run: SPEC=2 pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore test/test_custom_kernel.py --ignore test/unit/test_hashing.py --ignore test/unit/test_autogen.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
 
   fuzzing:
     name: Fuzzing
@@ -467,11 +465,11 @@ jobs:
       - name: Test MLPerf stuff
         run: CL=1 python -m pytest -n=auto test/external/external_test_optim.py test/external/external_test_losses.py test/external/external_test_metrics.py test/external/external_test_datasets.py --durations=20
       - name: NULL=1 beautiful_mnist_multigpu
-        run: NULL=1 NULL_ALLOW_COPYOUT=1 python examples/beautiful_mnist_multigpu.py
+        run: NULL=1 python examples/beautiful_mnist_multigpu.py
       - name: Test Bert training
-        run: NULL=1 NULL_ALLOW_COPYOUT=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+        run: NULL=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
       - name: Test llama 3 training
-        run: NULL=1 NULL_ALLOW_COPYOUT=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
+        run: NULL=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
@@ -611,7 +609,7 @@ jobs:
         WEBGPU=1 DEBUG=4 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
     - name: Run selected webgpu tests
       run: |
-          WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Vulkan" python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore=test/null --durations=20
+          WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Vulkan" python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
 
@@ -739,7 +737,7 @@ jobs:
           DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
       - name: Run pytest (cuda)
         # skip multitensor because it's slow
-        run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore=test/null --ignore test/test_gc.py --ignore test/test_multitensor.py --durations=20
+        run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore test/test_gc.py --ignore test/test_multitensor.py --durations=20
       - name: Run TestOps.test_add with PMA
         run: VIZ=-1 PMA=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
       - name: Run process replay tests
@@ -772,7 +770,7 @@ jobs:
           python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CPU','CL'], Device.DEFAULT"
           DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
       - name: Run pytest (${{ matrix.backend }})
-        run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore=test/null --durations=20
+        run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
       - name: Run TRANSCENDENTAL math
         run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
       - name: Run process replay tests
@@ -799,8 +797,6 @@ jobs:
         llvm: 'true'
     - name: Run unit tests
       run: METAL=1 python -m pytest -n=auto test/unit/ --durations=20
-    - name: Run NULL backend tests
-      run: NULL=1 python -m pytest -n=auto test/null/ --durations=20
     - name: Run ONNX
       run: METAL=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
     - name: Test tensor core ops (fake)
@@ -902,7 +898,7 @@ jobs:
           python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU','LVP':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
           DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
       - name: Run pytest (${{ matrix.backend }})
-        run: python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore=test/null --durations=20
+        run: python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
       - name: Run macOS-specific unit test
@@ -935,11 +931,7 @@ jobs:
       - name: Run unit tests
         if: matrix.backend=='llvm'
         # test_newton_schulz hits RecursionError
-        run: python -m pytest -n=auto test/unit/ --ignore=test/unit/test_disk_tensor.py --ignore=test/unit/test_tar.py --ignore=test/unit/test_linalg.py --durations=20
-      - name: Run NULL backend tests
-        if: matrix.backend=='llvm'
-        shell: bash
-        run: CPU=0 CPU_LLVM=0 NULL=1 python -m pytest -n=auto test/null/ --ignore=test/null/test_elf.py --durations=20
+        run: python -m pytest -n=auto test/unit/ --ignore=test/unit/test_disk_tensor.py --ignore=test/unit/test_elf.py --ignore=test/unit/test_tar.py --ignore=test/unit/test_linalg.py --durations=20
       - name: Run pytest (${{ matrix.backend }})
         shell: bash
         run: |
@@ -965,10 +957,10 @@ jobs:
           key: compile-${{ matrix.backend }}
           deps: testing_unit
           mesa: ${{ (matrix.backend == 'ir3' || matrix.backend == 'nak') && 'true' }}
-          python-version: '3.12'
+          python-version: '3.14'
       - name: Set env
         shell: bash
-        run: printf "NULL=1\nNULL_ALLOW_COPYOUT=1\n${{ matrix.backend == 'ir3' && 'NULL_IR3=1' || matrix.backend == 'nak' && 'NULL_NAK=1' }}" >> $GITHUB_ENV
+        run: printf "NULL=1\n${{ matrix.backend == 'ir3' && 'NULL_IR3=1' || matrix.backend == 'nak' && 'NULL_NAK=1' }}" >> $GITHUB_ENV
       - name: Run test_ops
         shell: bash
         run: |

.pre-commit-config.yaml

@@ -28,7 +28,7 @@ repos:
         pass_filenames: false
       - id: tests
         name: comprehensive test suite
-        entry: env OMP_NUM_THREADS=1 SKIP_SLOW_TEST=1 PYTHONPATH="." python3 -m pytest -n=6 test/test_ops.py test/test_schedule.py test/unit/test_assign.py test/test_tensor.py test/test_jit.py test/unit/test_schedule_cache.py test/null/test_pattern_matcher.py test/null/test_uop_symbolic.py test/unit/test_helpers.py
+        entry: env OMP_NUM_THREADS=1 SKIP_SLOW_TEST=1 PYTHONPATH="." python3 -m pytest -n=6 test/test_ops.py test/test_schedule.py test/unit/test_assign.py test/test_tensor.py test/test_jit.py test/unit/test_schedule_cache.py test/unit/test_pattern_matcher.py test/unit/test_uop_symbolic.py test/unit/test_helpers.py
         language: system
         always_run: true
         pass_filenames: false

examples/anthropic_challenge.py

@@ -72,7 +72,7 @@ vliw_prepare = PatternMatcher([
   # cast is fake
   (UPat(Ops.CAST, name="c"), lambda c: c.src[0]),
   # rewrites to hardcode the addresses in memory
-  (UPat(Ops.PARAM, name="dg"), lambda dg: UOp.const(dtypes.uint, global_addrs[dg.arg])),
+  (UPat(Ops.DEFINE_GLOBAL, name="dg"), lambda dg: UOp.const(dtypes.uint, global_addrs[dg.arg])),
   # INDEX is just plus
   (UPat(Ops.INDEX, name="i"), lambda i: i.src[0]+i.src[1]),
 ])+symbolic

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_beam.sh

@@ -8,11 +8,9 @@ export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
 export DEBUG=${DEBUG:-2}
 export FLASH_ATTENTION=${FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}
-export USE_ATOMICS=${USE_ATOMICS:-1}
-export ASM_GEMM=${ASM_GEMM:-1}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=8 BS=16 EVAL_BS=8 GRADIENT_ACC_STEPS=1
+export DP=8 BS=8 EVAL_BS=8 GRADIENT_ACC_STEPS=2
 export GBS=$((BS * GRADIENT_ACC_STEPS))
 
 export MODEL="llama3"

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh

@@ -13,7 +13,7 @@ export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=${DP:-8} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
+export DP=${DP:-8} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
 export GBS=$((BS * GRADIENT_ACC_STEPS))
 
 export MODEL="llama3"

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/startup_walltime.sh

@@ -2,9 +2,8 @@
 export BENCHMARK=5
 export EVAL_BS=0
 export FAKEDATA=1
-export NULL_ALLOW_COPYOUT=1
 export HIP_VISIBLE_DEVICES=""
 export DEV=NULL
 export JITBEAM=0
 export LLAMA_LAYERS=${LLAMA_LAYERS:-"2"}
-time examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
+time examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh

extra/assembly/amd/emu.py

@@ -265,11 +265,11 @@ def _collect_data_slices(assigns: list[tuple[str, UOp]], data_prefix: str, pcode
 class _Ctx:
   """Context for instruction compilation - holds buffers and helpers."""
   __slots__ = ('inst_size', 'dyn_fields', '_axis_id')
-  sgpr = UOp(Ops.PARAM, dtypes.uint32.ptr(SGPR_COUNT), arg=0)
-  vgpr = UOp(Ops.PARAM, dtypes.uint32.ptr(VGPR_SIZE), arg=1)
-  vmem = UOp(Ops.PARAM, dtypes.uint32.ptr(1 << 46), arg=2)
-  lds = UOp(Ops.PARAM, dtypes.uint32.ptr(16384), arg=3)
-  scratch = UOp(Ops.PARAM, dtypes.uint8.ptr(1 << 30), arg=4)
+  sgpr = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(SGPR_COUNT), arg=0)
+  vgpr = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(VGPR_SIZE), arg=1)
+  vmem = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(1 << 46), arg=2)
+  lds = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(16384), arg=3)
+  scratch = UOp(Ops.DEFINE_GLOBAL, dtypes.uint8.ptr(1 << 30), arg=4)
 
   def __init__(self, inst_size: int):
     self.inst_size, self._axis_id = inst_size, 0

extra/assembly/amd/test/test_emu2_pcode.py

@@ -180,7 +180,7 @@ class TestDSPcodePatterns(unittest.TestCase):
   def test_mem_read_parsing(self):
     """Test MEM[addr].type read expression parsing."""
     # Create a mock LDS buffer
-    lds = UOp(Ops.PARAM, dtypes.uint32.ptr(16384), arg=3)
+    lds = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(16384), arg=3)
     addr = UOp.const(dtypes.uint32, 0)
     vars = {'_lds': lds, 'ADDR': addr, 'OFFSET': UOp.const(dtypes.uint32, 0)}
 
@@ -213,7 +213,7 @@ class TestDSPcodePatterns(unittest.TestCase):
     """Test DS_LOAD_2ADDR_B32 pcode parsing produces RETURN_DATA assignments."""
     pcode = PCODE.get(DSOp.DS_LOAD_2ADDR_B32)
     self.assertIsNotNone(pcode)
-    lds = UOp(Ops.PARAM, dtypes.uint32.ptr(16384), arg=3)
+    lds = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(16384), arg=3)
     srcs = {
       'ADDR': UOp.const(dtypes.uint32, 0),
       'OFFSET0': UOp.const(dtypes.uint32, 0),
@@ -286,7 +286,7 @@ class TestAllPcode(unittest.TestCase):
   def _make_srcs(self):
     """Create dummy source variables for pcode parsing."""
     u32, u64 = lambda v=0: UOp.const(dtypes.uint32, v), lambda v=0: UOp.const(dtypes.uint64, v)
-    lds = UOp(Ops.PARAM, dtypes.uint32.ptr(16384), arg=3)
+    lds = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(16384), arg=3)
     return {'laneId': u32(), 'laneID': u32(), 'S0': u32(), 'S1': u32(), 'S2': u32(), 'S3': u32(), 'SRC0': u32(),
             'D0': u32(), 'D1': u32(), 'DST': u32(), 'VDST': u32(), 'SDST': u32(),
             'VCC': u64(), 'VCCZ': u32(), 'EXEC': u64(), 'EXEC_LO': u32(), 'EXECZ': u32(), 'SCC': u32(),

extra/optimization/get_action_space.py

@@ -0,0 +1,31 @@
+import random
+from extra.optimization.helpers import load_worlds, ast_str_to_lin
+from tinygrad.codegen.opt.search import actions
+from tinygrad.codegen.opt.kernel import Kernel
+from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
+from tinygrad.helpers import tqdm
+
+tactions = set()
+def test_rebuild(lin):
+  linr = Kernel(lin.ast)
+  for o in lin.applied_opts:
+    assert o in actions, f"{o} is not in actions"
+    tactions.add(o)
+    linr.apply_opt(o)
+
+  assert len(lin.sts) == len(linr.sts)
+  for st1,st2 in zip(lin.sts, linr.sts):
+    assert st1 == st2, f"{st1} != {st2}"
+
+if __name__ == "__main__":
+  ast_strs = load_worlds(False, False, False)
+  random.shuffle(ast_strs)
+  ast_strs = ast_strs[:2000]
+  for ast_str in tqdm(ast_strs):
+    lin = ast_str_to_lin(ast_str)
+    #if not lin.apply_tensor_cores():
+    lin.apply_opts(hand_coded_optimizations(lin))
+    test_rebuild(lin)
+
+  print(len(tactions), len(actions))
+  print(sorted(list(tactions)))

extra/optimization/rl.py

@@ -0,0 +1,76 @@
+import os
+import numpy as np
+import math, random
+from tinygrad.tensor import Tensor
+from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
+from tinygrad.codegen.opt.search import actions, bufs_from_lin, get_kernel_actions
+from tinygrad.nn.optim import Adam
+from extra.optimization.extract_policynet import PolicyNet
+from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats, time_linearizer
+
+if __name__ == "__main__":
+  net = PolicyNet()
+  if os.path.isfile("/tmp/policynet.safetensors"): load_state_dict(net, safe_load("/tmp/policynet.safetensors"))
+  optim = Adam(get_parameters(net))
+
+  ast_strs = load_worlds()
+
+  # select a world
+  all_feats, all_acts, all_rews = [], [], []
+  while 1:
+    Tensor.training = False
+    lin = ast_str_to_lin(random.choice(ast_strs))
+    rawbufs = bufs_from_lin(lin)
+    tm = last_tm = base_tm = time_linearizer(lin, rawbufs)
+
+    # take actions
+    feats, acts, rews = [], [], []
+    while 1:
+      feat = lin_to_feats(lin)
+      feats.append(feat)
+      probs = net(Tensor([feat])).exp()[0].numpy()
+
+      # mask valid actions
+      valid_action_mask = np.zeros((len(actions)+1), dtype=np.float32)
+      for x in get_kernel_actions(lin): valid_action_mask[x] = 1
+      probs *= valid_action_mask
+      probs /= sum(probs)
+
+      act = np.random.choice(len(probs), p=probs)
+      acts.append(act)
+      if act == 0:
+        rews.append(0)
+        break
+      try:
+        lin.apply_opt(actions[act-1])
+        tm = time_linearizer(lin, rawbufs)
+        if math.isinf(tm): raise Exception("failed")
+        rews.append(((last_tm-tm)/base_tm))
+        last_tm = tm
+      except Exception:
+        rews.append(-0.5)
+        break
+      #print(f"{tm*1e6:10.2f}", lin.colored_shape())
+
+    assert len(feats) == len(acts) and len(acts) == len(rews)
+    #print(rews)
+    print(f"***** EPISODE {len(rews)} steps, {sum(rews):5.2f} reward, {base_tm*1e6:12.2f} -> {tm*1e6:12.2f} : {lin.colored_shape()}")
+    all_feats += feats
+    all_acts += acts
+    # rewards to go
+    for i in range(len(rews)-2, -1, -1): rews[i] += rews[i+1]
+    all_rews += rews
+
+    BS = 32
+    if len(all_feats) >= BS:
+      Tensor.training = True
+      x = Tensor(all_feats[:BS])
+      mask = np.zeros((BS, len(actions)+1), dtype=np.float32)
+      mask[range(BS), all_acts[:BS]] = all_rews[:BS]
+      loss = -(net(x) * Tensor(mask)).mean()
+      optim.zero_grad()
+      loss.backward()
+      optim.step()
+      all_feats = all_feats[BS:]
+      all_acts = all_acts[BS:]
+      all_rews = all_rews[BS:]

extra/optimization/run_qnet.py

@@ -0,0 +1,32 @@
+from typing import List, Tuple
+from tinygrad.codegen.opt.kernel import Kernel
+from tinygrad.codegen.opt.search import get_kernel_actions, actions
+
+_net = None
+def beam_q_estimate(beam:List[Tuple[Kernel, float]]) -> List[Tuple[Kernel, float]]:
+  global _net
+  if _net is None:
+    from tinygrad.nn.state import load_state_dict, safe_load
+    from extra.optimization.pretrain_valuenet import ValueNet
+    _net = ValueNet(1021+len(actions), 2)
+    load_state_dict(_net, safe_load("/tmp/qnet.safetensors"), verbose=False)
+  from tinygrad.tensor import Tensor
+  from tinygrad.helpers import Context
+  from extra.optimization.helpers import lin_to_feats
+  import numpy as np
+  feats = []
+  lins = []
+  base_tms = []
+  for lin,tm in beam:
+    lin_feats = lin_to_feats(lin)
+    for a,v in get_kernel_actions(lin, include_0=False).items():
+      acts = np.zeros(len(actions))
+      acts[a-1] = 1.0
+      feats.append(np.concatenate([lin_feats, acts]))
+      lins.append(v)
+      base_tms.append(tm)
+  with Context(BEAM=0):
+    with Tensor.train(False):
+      preds = _net(Tensor(feats)).numpy()
+  pred_time = np.array(base_tms) / np.exp(preds[:, 0])
+  return sorted(zip(lins, pred_time), key=lambda x: x[1])

extra/optimization/search.py

@@ -0,0 +1,34 @@
+import argparse
+from extra.optimization.helpers import ast_str_to_lin, time_linearizer
+
+from tinygrad import dtypes
+from tinygrad.helpers import BEAM, getenv
+from tinygrad.device import Device, Compiled
+from tinygrad.codegen.opt.kernel import Kernel
+from tinygrad.codegen.opt.search import beam_search, bufs_from_lin
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser(description="Run a search for the optimal opts for a kernel", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+  parser.add_argument("--ast", type=str, default=None, help="the ast for the kernel to be optimized")
+  parser.add_argument("--file", type=str, default=None, help="a file containing asts to be optimized, one per line")
+  args = parser.parse_args()
+
+  device: Compiled = Device[Device.DEFAULT]
+  print(f"optimizing for {Device.DEFAULT}")
+
+  if args.ast is not None:
+    ast_strs = [args.ast]
+  elif args.file is not None:
+    with open(args.file, 'r') as file:
+      ast_strs = file.readlines()
+
+  for i, ast_str in enumerate(ast_strs):
+    print(f"optimizing {i}/{len(ast_strs)}\nast={ast_str}")
+    lin = ast_str_to_lin(ast_str, opts=device.renderer)
+    rawbufs = bufs_from_lin(lin)
+    lin = beam_search(lin, rawbufs, getenv("BEAM", 8), bool(getenv("BEAM_ESTIMATE", 1)))
+
+    tm = time_linearizer(lin, rawbufs, allow_test_size=False, cnt=10)
+    print(f"final time {tm*1e6:9.0f} us: {lin.colored_shape()}")
+    print(lin.applied_opts)

extra/optimization/test_helpers.py

@@ -0,0 +1,19 @@
+import unittest
+
+from extra.optimization.helpers import load_worlds
+
+class TestKernelDataset(unittest.TestCase):
+  def test_load_worlds_filters(self):
+    all_kernels = load_worlds(filter_reduce=False, filter_noimage=False, filter_novariable=False)
+
+    reduce_kernels = load_worlds(filter_reduce=True, filter_noimage=False, filter_novariable=False)
+    self.assertGreater(len(all_kernels), len(reduce_kernels))
+
+    image_kernels = load_worlds(filter_reduce=False, filter_noimage=True, filter_novariable=False)
+    self.assertGreater(len(all_kernels), len(image_kernels))
+
+    variable_kernels = load_worlds(filter_reduce=False, filter_noimage=False, filter_novariable=True)
+    self.assertGreater(len(all_kernels), len(variable_kernels))
+
+if __name__ == '__main__':
+  unittest.main()

extra/optimization/test_net.py

@@ -0,0 +1,67 @@
+import numpy as np
+import math
+import random
+np.set_printoptions(suppress=True)
+from copy import deepcopy
+from tinygrad.helpers import getenv, colored
+from tinygrad.tensor import Tensor
+from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
+from tinygrad.codegen.opt.search import bufs_from_lin, actions, get_kernel_actions
+from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
+from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats, time_linearizer
+from extra.optimization.extract_policynet import PolicyNet
+from extra.optimization.pretrain_valuenet import ValueNet
+
+VALUE = getenv("VALUE")
+
+if __name__ == "__main__":
+  if VALUE:
+    net = ValueNet()
+    load_state_dict(net, safe_load("/tmp/valuenet.safetensors"))
+  else:
+    net = PolicyNet()
+    load_state_dict(net, safe_load("/tmp/policynet.safetensors"))
+
+  ast_strs = load_worlds()
+
+  # real randomness
+  random.seed()
+  random.shuffle(ast_strs)
+
+  wins = 0
+  for ep_num,ast_str in enumerate(ast_strs):
+    print("\nEPISODE", ep_num, f"win {wins*100/max(1,ep_num):.2f}%")
+    lin = ast_str_to_lin(ast_str)
+    rawbufs = bufs_from_lin(lin)
+
+    linhc = deepcopy(lin)
+    linhc.applied_opts(hand_coded_optimizations(linhc))
+    tmhc = time_linearizer(linhc, rawbufs)
+    print(f"{tmhc*1e6:10.2f}     HC    ", linhc.colored_shape())
+
+    pred_time = float('nan')
+    tm = float('inf')
+    while 1:
+      if VALUE:
+        acts,feats = [], []
+        for k,v in get_kernel_actions(lin).items():
+          acts.append(k)
+          feats.append(lin_to_feats(v))
+        preds = net(Tensor(feats))
+        pred_time = math.exp(preds.numpy().min())
+        act = acts[preds.numpy().argmin()]
+      else:
+        probs = net(Tensor([lin_to_feats(lin)]))
+        dist = probs.exp().numpy()
+        act = dist.argmax()
+      if act == 0: break
+      try:
+        lin.apply_opt(actions[act-1])
+      except Exception:
+        print("FAILED")
+        break
+      tm = time_linearizer(lin, rawbufs)
+      print(f"{tm*1e6:10.2f} {pred_time*1e6:10.2f}", lin.colored_shape())
+
+    print(f"{colored('BEAT', 'green') if tm < tmhc else colored('lost', 'red')} hand coded {tmhc/tm:5.2f}x")
+    wins += int(tm < tmhc)

extra/optimization/test_time_linearizer.py

@@ -0,0 +1,21 @@
+from extra.optimization.helpers import load_worlds, ast_str_to_lin, time_linearizer
+from tinygrad.codegen.opt.search import bufs_from_lin, get_kernel_actions
+
+if __name__ == "__main__":
+  ast_strs = load_worlds()
+  for i, ast_str in enumerate(ast_strs):
+    lin = ast_str_to_lin(ast_str)
+    rawbufs = bufs_from_lin(lin)
+    test_tm = time_linearizer(lin, rawbufs)
+    if test_tm < 1e-2: continue
+    print(f"EXAMPLE {i}")
+    acted_lins = get_kernel_actions(lin)
+    ok_avg, short_avg = 0, 0
+    for k,v in acted_lins.items():
+      tm1 = time_linearizer(v, rawbufs)
+      tm2 = time_linearizer(v, rawbufs)
+      tm3 = time_linearizer(v, rawbufs, False)
+      print(v.colored_shape(50), f"{tm1*1e3:10.2f} {tm2*1e3:10.2f} {tm3*1e3:10.2f}  :  {((tm1-tm2)/tm1)*100:5.2f}%  vs  {((tm1-tm3)/tm1)*100:5.2f}%")
+      ok_avg += (tm1-tm2)/tm1
+      short_avg += (tm1-tm3)/tm1
+    print(f"{ok_avg/len(acted_lins)*100:5.2f}% vs {short_avg/len(acted_lins)*100:5.2f}%")

extra/thunder/tiny/fa.py

@@ -2,7 +2,7 @@ import math
 
 from tinygrad import Tensor, dtypes
 from tinygrad.helpers import DEBUG
-from tinygrad.uop.ops import UOp, Ops
+from tinygrad.uop.ops import UOp
 
 from extra.thunder.tiny.tk import WARP_THREADS
 from extra.thunder.tiny.tk.kernel import Kernel
@@ -43,12 +43,11 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
   B_local = B // num_devices
   if DEBUG >= 2: print(f"Flash Attention {B=} {B_local=} {N=} {H=} {D=} {H_KV=} {GROUP_SIZE=}")
 
-  def _custom_forward_impl(ou:UOp, l_vecu:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp|None) -> UOp:
+  def custom_forward(ou:UOp, l_vecu:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp) -> UOp:
     with Kernel("fa_custom_forward", (H, N // (Q_BLOCK_SIZE*NUM_WORKERS), B_local), NUM_WORKERS * WARP_THREADS) as ker:
       warp = ker.warp
 
-      o, q, k, v, l_vec = GL(ou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker), GL(l_vecu, ker)
-      mask = GL(masku, ker) if masku is not None else None
+      o, q, k, v, mask, l_vec = GL(ou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker), GL(masku, ker), GL(l_vecu, ker)
 
       head = ker.blockIdx_x
       head_kv = head // GROUP_SIZE
@@ -87,8 +86,7 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
       q_reg = warp.copy(q_reg, q_reg_fl)
       q_reg_transposed = warp.transpose(q_reg_transposed, q_reg)
 
-      num_kv_blocks = (q_seq + 1) if is_causal else (N // KV_BLOCK_SIZE)
-      for kv_idx in ker.range(num_kv_blocks):
+      for kv_idx in ker.range(N // KV_BLOCK_SIZE):
         k_smem = warp.load(k_smem, k, (), (batch, kv_idx, head_kv, 0), axis=1)
         v_smem = warp.load(v_smem, v, (), (batch, kv_idx, head_kv, 0), axis=1)
 
@@ -101,16 +99,9 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
         att_block = warp.mma_AtB(att_block, k_reg_transposed, q_reg_transposed)
 
         # apply attention mask
-        if is_causal:
-          bs_rows, bs_cols, bs_stride = att_block.base_shape.rows, att_block.base_shape.cols, att_block.base_shape.stride
-          q_base = q_seq * Q_BLOCK_SIZE + (warp.laneid % bs_cols)
-          kv_base = kv_idx * KV_BLOCK_SIZE + (warp.laneid // bs_cols) * bs_stride
-          att_block = warp.map(att_block,
-            lambda x, idx: ((kv_base + idx[0]*bs_rows + idx[2]) > (q_base + idx[1]*bs_cols)).alu(Ops.WHERE, UOp.ufix(x._uop, -math.inf), x))
-        elif mask is not None:
-          mask_reg = warp.load(mask_reg, mask, (), (batch, 0, q_seq, kv_idx), axis=2)
-          mask_reg_transposed = warp.transpose(mask_reg_transposed, mask_reg)
-          att_block += mask_reg_transposed
+        mask_reg = warp.load(mask_reg, mask, (), (batch, 0, q_seq, kv_idx), axis=2)
+        mask_reg_transposed = warp.transpose(mask_reg_transposed, mask_reg)
+        att_block += mask_reg_transposed
 
         # softmax
         max_vec_last = warp.copy(max_vec_last.after(kv_idx), max_vec)
@@ -150,18 +141,11 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
 
       return ker.finish()
 
-  def custom_forward_causal(ou:UOp, l_vecu:UOp, qu:UOp, ku:UOp, vu:UOp) -> UOp:
-    return _custom_forward_impl(ou, l_vecu, qu, ku, vu, None)
-
-  def custom_forward_masked(ou:UOp, l_vecu:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp) -> UOp:
-    return _custom_forward_impl(ou, l_vecu, qu, ku, vu, masku)
-
-  def _custom_backward_q_impl(dqu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp|None, l_vecu:UOp, delta_vecu:UOp) -> UOp:
+  def custom_backward_q(dqu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp, l_vecu:UOp, delta_vecu:UOp) -> UOp:
     with Kernel("fa_custom_backward_q", (H, N // (Q_BLOCK_SIZE*NUM_WORKERS), B_local), NUM_WORKERS * WARP_THREADS) as ker:
       warp = ker.warp
 
-      dq, do, q, k, v = GL(dqu, ker), GL(dou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker)
-      mask = GL(masku, ker) if masku is not None else None
+      dq, do, q, k, v, mask = GL(dqu, ker), GL(dou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker), GL(masku, ker)
       l_vec, delta_vec = GL(l_vecu, ker), GL(delta_vecu, ker)
 
       head = ker.blockIdx_x
@@ -210,8 +194,7 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
       l_vec_reg *= 1.0 / math.log(2)
       delta_vec_reg = warp.load(delta_vec_reg, delta_vec, (), (batch, head, 0, q_seq), axis=2)
 
-      num_kv_blocks = (q_seq + 1) if is_causal else (N // KV_BLOCK_SIZE)
-      for kv_idx in ker.range(num_kv_blocks):
+      for kv_idx in ker.range(N // KV_BLOCK_SIZE):
         k_smem = warp.load(k_smem, k, (), (batch, kv_idx, head_kv, 0), axis=1)
         v_smem = warp.load(v_smem, v, (), (batch, kv_idx, head_kv, 0), axis=1)
 
@@ -226,16 +209,9 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
         att_block = warp.mma_AtB(att_block, k_reg_t, q_reg_t)
 
         # apply attention mask
-        if is_causal:
-          bs_rows, bs_cols, bs_stride = att_block.base_shape.rows, att_block.base_shape.cols, att_block.base_shape.stride
-          q_base = q_seq * Q_BLOCK_SIZE + (warp.laneid % bs_cols)
-          kv_base = kv_idx * KV_BLOCK_SIZE + (warp.laneid // bs_cols) * bs_stride
-          att_block = warp.map(att_block,
-            lambda x, idx: ((kv_base + idx[0]*bs_rows + idx[2]) > (q_base + idx[1]*bs_cols)).alu(Ops.WHERE, UOp.ufix(x._uop, -math.inf), x))
-        elif mask is not None:
-          mask_reg = warp.load(mask_reg, mask, (), (batch, 0, q_seq, kv_idx), axis=2)
-          mask_reg_transposed = warp.transpose(mask_reg_transposed, mask_reg)
-          att_block += mask_reg_transposed
+        mask_reg = warp.load(mask_reg, mask, (), (batch, 0, q_seq, kv_idx), axis=2)
+        mask_reg_transposed = warp.transpose(mask_reg_transposed, mask_reg)
+        att_block += mask_reg_transposed
 
         att_block -= l_vec_reg
         att_block = att_block.exp2()
@@ -255,18 +231,11 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
 
       return ker.finish()
 
-  def custom_backward_q_causal(dqu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, l_vecu:UOp, delta_vecu:UOp) -> UOp:
-    return _custom_backward_q_impl(dqu, dou, qu, ku, vu, None, l_vecu, delta_vecu)
-
-  def custom_backward_q_masked(dqu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp, l_vecu:UOp, delta_vecu:UOp) -> UOp:
-    return _custom_backward_q_impl(dqu, dou, qu, ku, vu, masku, l_vecu, delta_vecu)
-
-  def _custom_backward_kv_impl(dku:UOp, dvu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp|None, l_vecu:UOp, delta_vecu:UOp):
+  def custom_backward_kv(dku:UOp, dvu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp, l_vecu:UOp, delta_vecu:UOp):
     with Kernel("fa_custom_backward_kv", (H_KV, N // (KV_BLOCK_SIZE*NUM_WORKERS), B_local), NUM_WORKERS * WARP_THREADS) as ker:
       warp = ker.warp
 
-      dk, dv, do, q, k, v = GL(dku, ker), GL(dvu, ker), GL(dou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker)
-      mask = GL(masku, ker) if masku is not None else None
+      dk, dv, do, q, k, v, mask = GL(dku, ker), GL(dvu, ker), GL(dou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker), GL(masku, ker)
       l_vec, delta_vec = GL(l_vecu, ker), GL(delta_vecu, ker)
 
       head_kv = ker.blockIdx_x
@@ -333,16 +302,9 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
           att_block *= (1.0 / math.sqrt(D)) * (1.0 / math.log(2))
 
           # apply attention mask
-          if is_causal:
-            bs_rows, bs_cols, bs_stride = att_block.base_shape.rows, att_block.base_shape.cols, att_block.base_shape.stride
-            q_base = q_idx * Q_BLOCK_SIZE + (warp.laneid % bs_cols)
-            kv_base = kv_seq * KV_BLOCK_SIZE + (warp.laneid // bs_cols) * bs_stride
-            att_block = warp.map(att_block,
-              lambda x, idx: ((kv_base + idx[0]*bs_rows + idx[2]) > (q_base + idx[1]*bs_cols)).alu(Ops.WHERE, UOp.ufix(x._uop, -math.inf), x))
-          elif mask is not None:
-            mask_reg = warp.load(mask_reg, mask, (), (batch, 0, q_idx, kv_seq), axis=2)
-            mask_reg_transposed = warp.transpose(mask_reg_transposed, mask_reg)
-            att_block += mask_reg_transposed
+          mask_reg = warp.load(mask_reg, mask, (), (batch, 0, q_idx, kv_seq), axis=2)
+          mask_reg_transposed = warp.transpose(mask_reg_transposed, mask_reg)
+          att_block += mask_reg_transposed
 
           att_block -= l_vec_reg
           att_block = att_block.exp2()
@@ -374,31 +336,24 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
 
       return ker.finish(2)
 
-  def custom_backward_kv_causal(dku:UOp, dvu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, l_vecu:UOp, delta_vecu:UOp):
-    return _custom_backward_kv_impl(dku, dvu, dou, qu, ku, vu, None, l_vecu, delta_vecu)
-
-  def custom_backward_kv_masked(dku:UOp, dvu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp, l_vecu:UOp, delta_vecu:UOp):
-    return _custom_backward_kv_impl(dku, dvu, dou, qu, ku, vu, masku, l_vecu, delta_vecu)
-
   single_device = xq.device[0] if isinstance(xq.device, tuple) else xq.device
 
   if is_causal:
     if attn_mask is not None: raise RuntimeError("cannot set attn_mask when is_causal=True")
-  elif attn_mask is not None:
+    attn_mask = Tensor.ones((B, 1, N, N), requires_grad=False, device=single_device, dtype=dtypes.bool).tril()
+  if attn_mask is not None:
     if attn_mask.dtype == dtypes.bool: attn_mask = attn_mask.where(0, -float("inf"))
-    if attn_mask.shape != (B, 1, N, N):
-      attn_mask = attn_mask.expand(B, 1, N, N)
-    if isinstance(xq.device, tuple) and not isinstance(attn_mask.device, tuple):
-      attn_mask = attn_mask.shard(xq.device, axis=0)
   else:
     attn_mask = Tensor.zeros((B, 1, N, N), requires_grad=False, device=single_device, dtype=dtypes.float32)
-    if isinstance(xq.device, tuple):
-      attn_mask = attn_mask.shard(xq.device, axis=0)
+  if attn_mask.shape != (B, 1, N, N):
+    attn_mask = attn_mask.expand(B, 1, N, N)
+  if isinstance(xq.device, tuple) and not isinstance(attn_mask.device, tuple):
+    attn_mask = attn_mask.shard(xq.device, axis=0)
 
   attn = _sharded_empty_like(xq, axis=0)
   l_vec = _sharded_empty((B, H, 1, N), xq, axis=0)
 
-  def grad_causal(gradu:UOp, _) -> tuple[None, None, UOp, UOp, UOp]:
+  def grad(gradu:UOp, _) -> tuple[None, None, UOp, UOp, UOp, None]:
     grad = Tensor(gradu, device=gradu.device)
     grad_q = _sharded_empty_like(xq, axis=0)
     grad_k = _sharded_empty_like(xk, axis=0)
@@ -406,26 +361,11 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
 
     delta_vec = (grad * attn).sum(-1, dtype=dtypes.float32).transpose(1, 2).unsqueeze(-2).detach()
 
-    grad_q = Tensor.custom_kernel(grad_q, grad, xq, xk, xv, l_vec, delta_vec, fxn=custom_backward_q_causal)[0]
-    grad_k, grad_v = Tensor.custom_kernel(grad_k, grad_v, grad, xq, xk, xv, l_vec, delta_vec, fxn=custom_backward_kv_causal)[:2]
-    return (None, None, grad_q.uop, grad_k.uop, grad_v.uop)
-
-  def grad_masked(gradu:UOp, _) -> tuple[None, None, UOp, UOp, UOp, None]:
-    grad = Tensor(gradu, device=gradu.device)
-    grad_q = _sharded_empty_like(xq, axis=0)
-    grad_k = _sharded_empty_like(xk, axis=0)
-    grad_v = _sharded_empty_like(xv, axis=0)
-
-    delta_vec = (grad * attn).sum(-1, dtype=dtypes.float32).transpose(1, 2).unsqueeze(-2).detach()
-
-    grad_q = Tensor.custom_kernel(grad_q, grad, xq, xk, xv, attn_mask, l_vec, delta_vec, fxn=custom_backward_q_masked)[0]
-    grad_k, grad_v = Tensor.custom_kernel(grad_k, grad_v, grad, xq, xk, xv, attn_mask, l_vec, delta_vec, fxn=custom_backward_kv_masked)[:2]
+    grad_q = Tensor.custom_kernel(grad_q, grad, xq, xk, xv, attn_mask, l_vec, delta_vec, fxn=custom_backward_q)[0]
+    grad_k, grad_v = Tensor.custom_kernel(grad_k, grad_v, grad, xq, xk, xv, attn_mask, l_vec, delta_vec, fxn=custom_backward_kv)[:2]
     return (None, None, grad_q.uop, grad_k.uop, grad_v.uop, None)
 
-  if is_causal:
-    attn, l_vec = Tensor.custom_kernel(attn, l_vec, xq, xk, xv, fxn=custom_forward_causal, grad_fxn=grad_causal)[:2]
-  else:
-    attn, l_vec = Tensor.custom_kernel(attn, l_vec, xq, xk, xv, attn_mask, fxn=custom_forward_masked, grad_fxn=grad_masked)[:2]
+  attn, l_vec = Tensor.custom_kernel(attn, l_vec, xq, xk, xv, attn_mask, fxn=custom_forward, grad_fxn=grad)[:2]
   attn_ = attn[:, :N_, :, :D_]
 
   return attn_.transpose(1, 2).cast(odtype)

test/external/external_benchmark_hcopt.py

@@ -0,0 +1,43 @@
+import random
+from tinygrad.helpers import getenv
+from tinygrad.codegen.opt.search import beam_search, bufs_from_lin
+from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
+from extra.optimization.helpers import load_worlds, ast_str_to_lin, time_linearizer
+
+def optimize_kernel(k):
+  # TODO: update this
+  return hand_coded_optimizations(k)
+
+if __name__ == '__main__':
+  hcopt_wins = beam_wins = tie = 0
+  hcopt_total = beam_total = 0.0
+
+  worlds = load_worlds(filter_reduce=False, filter_noimage=True, filter_novariable=False)
+  random.seed(0)
+  random.shuffle(worlds)
+
+  for world in worlds[:500]:
+    k = ast_str_to_lin(world)
+    rawbufs = bufs_from_lin(k)
+
+    k_hcopt = k.copy()
+    k_hcopt.apply_opts(optimize_kernel(k_hcopt))
+    k_beam = beam_search(k.copy(), rawbufs, getenv("BEAM", 2))
+
+    disable_cache = bool(getenv("NOCACHE", 0))
+    t_hcopt = time_linearizer(k_hcopt, rawbufs, allow_test_size=False, cnt=10, disable_cache=disable_cache, clear_l2=True) * 1e6
+    t_beam = time_linearizer(k_beam, rawbufs, allow_test_size=False, cnt=10, disable_cache=disable_cache, clear_l2=True) * 1e6
+
+    if t_hcopt == t_beam: tie += 1
+    elif t_hcopt < t_beam: hcopt_wins += 1
+    else: beam_wins += 1
+    hcopt_total += t_hcopt
+    beam_total += t_beam
+
+    print(f"{t_hcopt=:5.2f} {k_hcopt.applied_opts=}")
+    print("")
+    print(f"{t_beam=:5.2f} {k_beam.applied_opts=}")
+    print("*"*20)
+
+  print(f"{hcopt_wins=}, {beam_wins=}, {tie=}")
+  print(f"{hcopt_total=:.2f}, {beam_total=:.2f}")

test/external/external_benchmark_op_conv.py

@@ -11,7 +11,7 @@ from tinygrad.dtype import ImageDType, Invalid
 # PYTHONPATH="." DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
 
 def vision_conv_143():
-  c0 = UOp(Ops.PARAM, dtypes.imageh((16, 1024, 4)), (), 0)
+  c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((16, 1024, 4)), (), 0)
   c2 = UOp.range(32, 3, AxisType.LOOP)
   c5 = UOp.range(128, 4, AxisType.LOOP)
   c8 = UOp.range(16, 2, AxisType.LOOP)
@@ -21,13 +21,13 @@ def vision_conv_143():
   c26 = UOp.range(7, 1, AxisType.REDUCE)
   c27 = c2*2+c26
   c32 = ((c27<3)!=True)&(c27<67)
-  c34 = UOp(Ops.PARAM, dtypes.imageh((32, 1024, 4)), (), 1)
+  c34 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 1024, 4)), (), 1)
   c38 = c5//2
   c45 = (c32&c24).where((c27*64+c38+c17*4096+-12480), UOp.const(dtypes.index, Invalid))
   c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
-  c49 = UOp(Ops.PARAM, dtypes.imageh((64, 49, 4)), (), 2)
+  c49 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((64, 49, 4)), (), 2)
   c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
-  c63 = UOp(Ops.PARAM, dtypes.float.ptr(128), (), 3)
+  c63 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(128), (), 3)
   c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
   c67 = c0.index((c2*128+c5+c8*4096), ptr=True).store(c65).end(c8, c2, c5)
 
@@ -37,7 +37,7 @@ def vision_conv_143():
   return c67.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))
 
 def vision_conv_153():
-  c0 = UOp(Ops.PARAM, dtypes.imageh((8, 1024, 4)), (), 0)
+  c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((8, 1024, 4)), (), 0)
   c2 = UOp.range(16, 3, AxisType.LOOP)
   c5 = UOp.range(256, 4, AxisType.LOOP)
   c8 = UOp.range(8, 2, AxisType.LOOP)
@@ -47,13 +47,13 @@ def vision_conv_153():
   c26 = UOp.range(7, 1, AxisType.REDUCE)
   c27 = c2*2+c26
   c32 = ((c27<3)!=True)&(c27<35)
-  c34 = UOp(Ops.PARAM, dtypes.imageh((16, 1024, 4)), (), 1)
+  c34 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((16, 1024, 4)), (), 1)
   c38 = c5//2
   c45 = (c32&c24).where((c27*128+c38+c17*4096+-12672), UOp.const(dtypes.index, Invalid))
   c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
-  c49 = UOp(Ops.PARAM, dtypes.imageh((128, 49, 4)), (), 2)
+  c49 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((128, 49, 4)), (), 2)
   c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
-  c63 = UOp(Ops.PARAM, dtypes.float.ptr(256), (), 3)
+  c63 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(256), (), 3)
   c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
   c67 = c0.index((c2*256+c5+c8*4096), ptr=True).store(c65).end(c8, c2, c5)
 
@@ -63,16 +63,16 @@ def vision_conv_153():
   return c67.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))
 
 def dm_conv_172():
-  c0 = UOp(Ops.PARAM, dtypes.imageh((1, 240, 4)), (), 0)
+  c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 240, 4)), (), 0)
   c2 = UOp.range(960, 4, AxisType.LOOP)
-  c5 = UOp(Ops.PARAM, dtypes.imageh((8, 384, 4)), (), 1)
+  c5 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((8, 384, 4)), (), 1)
   c7 = UOp.range(32, 0, AxisType.REDUCE)
   c10 = UOp.range(4, 1, AxisType.REDUCE)
   c13 = UOp.range(12, 3, AxisType.REDUCE)
   c18 = UOp.range(8, 2, AxisType.REDUCE)
-  c23 = UOp(Ops.PARAM, dtypes.imageh((240, 128, 4)), (), 2)
+  c23 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((240, 128, 4)), (), 2)
   c35 = c5.index((c7*4+c10+c13*128+c18*1536))*c23.index((c10*4+c2%4+c7*16+c2//4*512))
-  c37 = UOp(Ops.PARAM, dtypes.float.ptr(960), (), 3)
+  c37 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(960), (), 3)
   c39 = c35.reduce(c7, c10, arg=Ops.ADD)+c37.index(c2)
   c50 = (1.0+((c39+0.044708251953125*(c39*(c39*c39)))*-2.3021129851685216).exp2()).reciprocal()*c39
   c53 = c50.reduce(c18, c13, arg=Ops.ADD)*0.010416666666666666
@@ -91,7 +91,7 @@ allocator = Device.default.allocator
 ps = get_program(ast, renderer)
 cr = CompiledRunner(replace(ps, device=Device.DEFAULT))
 
-gs = sorted(dedup([u for u in ast.toposort() if u.op is Ops.PARAM]), key=lambda u: u.arg)
+gs = sorted(dedup([u for u in ast.toposort() if u.op is Ops.DEFINE_GLOBAL]), key=lambda u: u.arg)
 # print(len(gs))
 # print([g.dtype for g in gs])
 bufs = [Buffer(ps.device, g.size, g.dtype if isinstance(g.dtype, ImageDType) else g.dtype._base).ensure_allocated() for g in gs]

test/external/external_test_gpu_crash.py

@@ -1,126 +0,0 @@
-# ruff: noqa: F405
-"""Tests for GPU crash scenarios using AMD assembly to trigger invalid operations.
-
-These tests intentionally cause GPU faults to verify error handling.
-Run with: AMD=1 python -m pytest test/external/external_test_gpu_crash.py -v
-"""
-import unittest, re
-from tinygrad.device import Device
-from extra.assembly.amd.autogen.rdna3.ins import *  # noqa: F403
-from extra.assembly.amd.dsl import s, v, Inst, NULL
-
-def assemble(code:str, name:str="test") -> str:
-  kd = {"next_free_vgpr": 8, "next_free_sgpr": 8, "wavefront_size32": 1, "user_sgpr_kernarg_segment_ptr": 1, "kernarg_size": 8}
-  return f".text\n.globl {name}\n.p2align 8\n.type {name},@function\n{name}:\n{code}\n.rodata\n.p2align 6\n.amdhsa_kernel {name}\n" + \
-         "\n".join(f".amdhsa_{k} {v}" for k,v in kd.items()) + "\n.end_amdhsa_kernel"
-
-@unittest.skipIf(Device.DEFAULT != "AMD", "AMD required")
-class TestGPUCrash(unittest.TestCase):
-  @classmethod
-  def setUpClass(cls):
-    from tinygrad.runtime.support.compiler_amd import HIPCompiler
-    cls.dev = Device["AMD"]
-    cls.compiler = HIPCompiler(cls.dev.arch)
-
-  def setUp(self):
-    # Verify device works before each test
-    from tinygrad import Tensor
-    try:
-      t = Tensor([1.0, 2.0], device="AMD").realize()
-      assert (t + 1).numpy().tolist() == [2.0, 3.0]
-    except Exception:
-      self.fail("Device not working before test")
-
-  def _run(self, code: str):
-    from tinygrad.runtime.ops_amd import AMDProgram
-    prg = AMDProgram(self.dev, "test", self.compiler.compile(assemble(code)))
-    prg(self.dev.allocator.alloc(64), global_size=(1,1,1), local_size=(1,1,1), wait=True)
-
-  def _run_insts(self, insts: list[Inst]): self._run("\n".join(i.disasm() for i in insts))
-
-  def _assert_gpu_fault(self, func):
-    """Assert that func raises a RuntimeError indicating a GPU fault (not a setup error)."""
-    with self.assertRaises(RuntimeError) as cm:
-      func()
-    err_msg = str(cm.exception).lower()
-    # Verify it's a GPU fault, not a setup/device initialization error
-    self.assertTrue(
-      re.search(r'fault|hang|timeout|illegal|memviol', err_msg),
-      f"Expected GPU fault error, got: {cm.exception}"
-    )
-
-
-class TestOutOfBoundsMemoryAccess(TestGPUCrash):
-  """Tests for out-of-bounds memory accesses."""
-
-  def test_global_load_null_ptr(self):
-    """Global load from NULL pointer."""
-    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0),
-             global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
-    self._assert_gpu_fault(lambda: self._run_insts(insts))
-
-  def test_global_store_null_ptr(self):
-    """Global store to NULL pointer."""
-    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), v_mov_b32_e32(v[2], 0xDEADBEEF),
-             global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
-    self._assert_gpu_fault(lambda: self._run_insts(insts))
-
-  def test_global_load_unmapped_high_address(self):
-    """Global load from high unmapped address (0xDEAD00000000)."""
-    insts = [v_mov_b32_e32(v[0], 0x00000000), v_mov_b32_e32(v[1], 0xDEAD),
-             global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
-    self._assert_gpu_fault(lambda: self._run_insts(insts))
-
-  def test_global_store_unmapped_high_address(self):
-    """Global store to high unmapped address."""
-    insts = [v_mov_b32_e32(v[0], 0x00000000), v_mov_b32_e32(v[1], 0xDEAD), v_mov_b32_e32(v[2], 0x12345678),
-             global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
-    self._assert_gpu_fault(lambda: self._run_insts(insts))
-
-  def test_global_atomic_unmapped(self):
-    """Atomic operation on unmapped memory."""
-    insts = [v_mov_b32_e32(v[0], 0xBEEF0000), v_mov_b32_e32(v[1], 0xDEAD), v_mov_b32_e32(v[2], 1),
-             global_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
-    self._assert_gpu_fault(lambda: self._run_insts(insts))
-
-
-class TestSMEMFaults(TestGPUCrash):
-  """Tests for scalar memory (SMEM) faults."""
-
-  def test_smem_load_null(self):
-    """SMEM load from NULL base."""
-    insts = [s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
-             s_load_b32(s[4], s[2:3], 0, soffset=NULL), s_waitcnt(0), s_endpgm()]
-    self._assert_gpu_fault(lambda: self._run_insts(insts))
-
-  def test_smem_load_unmapped(self):
-    """SMEM load from unmapped address."""
-    insts = [s_mov_b32(s[2], 0xBEEF0000), s_mov_b32(s[3], 0xDEAD),
-             s_load_b32(s[4], s[2:3], 0, soffset=NULL), s_waitcnt(0), s_endpgm()]
-    self._assert_gpu_fault(lambda: self._run_insts(insts))
-
-
-class TestFlatMemoryFaults(TestGPUCrash):
-  """Tests for FLAT memory instruction faults."""
-
-  def test_flat_load_null(self):
-    """FLAT load from NULL address."""
-    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0),
-             flat_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
-    self._assert_gpu_fault(lambda: self._run_insts(insts))
-
-  def test_flat_store_null(self):
-    """FLAT store to NULL address."""
-    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), v_mov_b32_e32(v[2], 0xDEADBEEF),
-             flat_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
-    self._assert_gpu_fault(lambda: self._run_insts(insts))
-
-  def test_flat_atomic_null(self):
-    """FLAT atomic on NULL address."""
-    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), v_mov_b32_e32(v[2], 1),
-             flat_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
-    self._assert_gpu_fault(lambda: self._run_insts(insts))
-
-
-if __name__ == "__main__":
-  unittest.main()

test/external/fuzz_linearizer.py

@@ -0,0 +1,339 @@
+import random, traceback, ctypes, argparse, os
+from typing import Any
+import numpy as np
+from collections import defaultdict
+from extra.optimization.helpers import load_worlds, ast_str_to_lin, kern_str_to_lin
+from tinygrad.engine.realize import get_program
+
+# We need to insert ioctl before opening devices.
+if os.getenv("VALIDATE_HCQ", 0) != 0:
+  try:
+    import extra.nv_gpu_driver.nv_ioctl
+    from tinygrad import Device
+    _, _ = Device["NV"], Device["CUDA"]
+  except Exception: pass
+
+  try:
+    import extra.qcom_gpu_driver.opencl_ioctl
+    from tinygrad import Device
+    _, _ = Device["QCOM"], Device["CL"]
+  except Exception: pass
+
+from tinygrad import Tensor, Device, dtypes
+from tinygrad.tensor import _to_np_dtype
+from tinygrad.codegen.opt.kernel import Kernel
+from tinygrad.codegen.opt import Opt, OptOps
+from tinygrad.codegen.opt.search import get_kernel_actions, bufs_from_lin
+from tinygrad.engine.realize import CompiledRunner
+from tinygrad.helpers import getenv, from_mv, prod, colored, Context, DEBUG, Timing
+from tinygrad.uop.ops import UOp, Ops
+from tinygrad.device import is_dtype_supported
+
+def on_linearizer_will_run(): pass
+def on_linearizer_did_run(): pass
+def compare_states(x, y): return (True, "")
+
+if getenv("VALIDATE_HCQ"):
+  if Device.DEFAULT == "NV":
+    print("VALIDATE_HCQ: Comparing NV to CUDA")
+    import extra.nv_gpu_driver.nv_ioctl
+    validate_device = Device["CUDA"]
+    on_linearizer_will_run = extra.nv_gpu_driver.nv_ioctl.before_launch
+    on_linearizer_did_run = extra.nv_gpu_driver.nv_ioctl.collect_last_launch_state
+    compare_states = extra.nv_gpu_driver.nv_ioctl.compare_launch_state
+  elif Device.DEFAULT == "QCOM":
+    print("VALIDATE_HCQ: Comparing QCOM to CL")
+    import extra.qcom_gpu_driver.opencl_ioctl
+    validate_device = Device["CL"]
+    on_linearizer_will_run = extra.qcom_gpu_driver.opencl_ioctl.before_launch
+    on_linearizer_did_run = extra.qcom_gpu_driver.opencl_ioctl.collect_last_launch_state
+    compare_states = extra.qcom_gpu_driver.opencl_ioctl.compare_launch_state
+  else:
+    print(colored("VALIDATE_HCQ options is ignored", 'red'))
+
+def tuplize_uops(uops:list[UOp]) -> tuple:
+  return tuple([(x.op, x.dtype, tuple(uops.index(x) for x in x.src), x.arg) for x in uops])
+
+def get_fuzz_rawbufs(lin):
+  rawbufs = bufs_from_lin(lin)
+
+  # Reallocate output buffer with additional area to detect out-of-bounds writes.
+  RED_AREA_SIZE = 1024
+  # setting output  # TODO: multi-output kernel
+  rawbufs[0] = get_fuzz_rawbuf_like(rawbufs[0], zero=True, size=rawbufs[0].size+RED_AREA_SIZE)
+  # setting inputs
+  with Context(DEBUG=0):
+    for rawbuf in rawbufs[1:]:
+      if dtypes.is_unsigned(rawbuf.dtype):
+        data = np.random.randint(0, 100, size=rawbuf.size, dtype=_to_np_dtype(rawbuf.dtype))
+      elif dtypes.is_int(rawbuf.dtype):
+        data = np.random.randint(-100, 100, size=rawbuf.size, dtype=_to_np_dtype(rawbuf.dtype))
+      elif rawbuf.dtype == dtypes.bool:
+        data = np.random.choice([True, False], size=rawbuf.size)
+      elif rawbuf.dtype == dtypes.half:
+        data = np.random.uniform(-1, 1, size=rawbuf.size).astype(dtype=_to_np_dtype(rawbuf.dtype))
+      else:
+        data = np.random.uniform(-10, 10, size=rawbuf.size).astype(dtype=_to_np_dtype(rawbuf.dtype))
+      rawbuf.copyin(Tensor(data, device=lin.opts.device).realize().uop.base.realized.as_buffer())
+  return rawbufs
+
+def get_fuzz_rawbuf_like(old_rawbuf, zero=False, copy=False, size=None, force_device=None):
+  rawbuf = type(old_rawbuf)(force_device or old_rawbuf.device, old_rawbuf.size if size is None else size, old_rawbuf.dtype).allocate()
+  if copy:
+    with Context(DEBUG=0): rawbuf.copyin(old_rawbuf.as_buffer())
+  elif zero:
+    with Context(DEBUG=0):
+      mv = memoryview(bytearray(rawbuf.size * rawbuf.dtype.itemsize))
+      ctypes.memset(from_mv(mv), 0, len(mv))
+      rawbuf.copyin(mv)
+  return rawbuf
+
+def run_linearizer(lin: Kernel, rawbufs=None, var_vals=None) -> tuple[str, Any]: # (error msg, run state)
+  if rawbufs is None: rawbufs = bufs_from_lin(lin)
+  if var_vals is None: var_vals = {v.expr: v.min for v in lin.vars}
+
+  # TODO: images needs required_optimization
+  try:
+    prg = CompiledRunner(get_program(lin.get_optimized_ast(), lin.opts))
+  except KeyboardInterrupt: raise
+  except Exception:
+    traceback.print_exc()
+    return "COMPILE_ERROR", None
+
+  if getenv("VALIDATE_HCQ"): on_linearizer_will_run()
+  try:
+    prg(rawbufs, var_vals, wait=True)
+  except KeyboardInterrupt: raise
+  except Exception:
+    traceback.print_exc()
+    return "EXEC_ERROR", None
+
+  if getenv("VALIDATE_HCQ"): run_state = on_linearizer_did_run()
+  else: run_state = None
+
+  return "PASS", run_state
+
+def compare_linearizer(lin: Kernel, rawbufs=None, var_vals=None, ground_truth=None, rtol=1e-2, atol=1e-2):
+  # TODO: for bfloat16 it compiles linearizer, but it does not run because numpy cannot generate bf16 buffer.
+  has_bf16 = any(b.dtype.base == dtypes.bfloat16 for b in lin.bufs)
+
+  # TODO: raise specific fuzzing errors instead of str, and propagate the error message
+  try:
+    if rawbufs is None:
+      rawbufs = get_fuzz_rawbufs(lin)
+    else:
+      rawbufs[0] = get_fuzz_rawbuf_like(rawbufs[0], zero=True) # get a new output buffer
+  except KeyboardInterrupt: raise
+  except BaseException:
+    return ("RAWBUFS_ERROR", rawbufs, var_vals, ground_truth, None)
+
+  if var_vals is None:
+    # TODO: handle symbolic max case
+    var_vals = {v.expr: random.randint(v.vmin, v.vmax) for v in lin.ast.variables()}
+
+  if ground_truth is None and not has_bf16:
+    unoptimized = Kernel(lin.ast)
+    if run_linearizer(unoptimized, rawbufs, var_vals)[0] != "PASS":
+      return ("BASELINE_ERROR", rawbufs, var_vals, ground_truth, None)
+    ground_truth = np.frombuffer(rawbufs[0].as_buffer(), _to_np_dtype(rawbufs[0].dtype)).copy()
+
+  rawbufs[0] = get_fuzz_rawbuf_like(rawbufs[0], zero=True) # get a new output buffer
+  run_msg, run_state = run_linearizer(lin, rawbufs, var_vals)
+  if run_msg != "PASS": return (run_msg, rawbufs, var_vals, ground_truth, run_state)
+
+  try:
+    if not has_bf16:
+      result = np.frombuffer(rawbufs[0].as_buffer(), _to_np_dtype(rawbufs[0].dtype))
+      np.testing.assert_allclose(result, ground_truth, rtol=rtol, atol=atol)
+  except KeyboardInterrupt: raise
+  except AssertionError as e:
+    if DEBUG >= 2:
+      print(f"COMPARE_ERROR details: {e}")
+      if getenv("DEBUG_VALUES") > 0:
+        mismatch_indices = np.where(~np.isclose(result, ground_truth, rtol=rtol, atol=atol))
+        mismatched_result = result[mismatch_indices]
+        mismatched_ground_truth = ground_truth[mismatch_indices]
+        for i, idx in enumerate(mismatch_indices[0]):
+          print(f"mismatch at {idx=}: result={mismatched_result[i]} <> ground_truth={mismatched_ground_truth[i]}")
+    return ("COMPARE_ERROR", rawbufs, var_vals, ground_truth, run_state)
+
+  return ("PASS", rawbufs, var_vals, ground_truth, run_state)
+
+def fuzz_linearizer(lin: Kernel, rtol=1e-2, atol=1e-2, opts_list=None):
+  SEED = getenv("SEED", 42)
+  random.seed(SEED)
+  np.random.seed(SEED)
+  print(lin.ast)
+  print(lin.colored_shape())
+  seen_uops = {}
+  last_lins = [lin]
+  failures:defaultdict[str, list[tuple[tuple[UOp, ...], list[Opt]]]] = defaultdict(list)
+  rawbufs, var_vals, ground_truth, validate_rawbufs = None, None, None, None
+
+  FUZZ_ALL_ACTIONS = getenv("FUZZ_ALL_ACTIONS", 0)
+  FUZZ_MAX_SIZE = getenv("FUZZ_MAX_SIZE", 0)
+  FUZZ_IGNORE_SIMPLE_OPS = getenv("FUZZ_IGNORE_SIMPLE_OPS", 1)
+
+  if FUZZ_MAX_SIZE > 0 and prod(lin.full_shape) > FUZZ_MAX_SIZE:
+    print("skipping large kernel")
+    return failures
+  if FUZZ_IGNORE_SIMPLE_OPS and _is_simple(lin):
+    print("skipping simple kernel")
+    return failures
+
+  test_depth = 1 if opts_list is not None else getenv("DEPTH", 1 if FUZZ_ALL_ACTIONS else 10)
+  for depth in range(test_depth):
+    next_lins = []
+    for lin in last_lins:
+      if opts_list is None: actions = get_kernel_actions(lin, include_0=False)
+      else:
+        actions = {}
+        for oi,opts in enumerate(opts_list):
+          lin2 = lin.copy()
+          for o in opts: lin2.apply_opt(o)
+          actions[oi] = lin2
+
+      if not actions: continue
+      if depth == 0 and getenv("FUZZ_REQUIRE_TC", 0):
+        tc_acts = {i: k for k in actions.values() if k.applied_opts[0].op == OptOps.TC}
+        if len(tc_acts) == 0: return failures
+        else: actions = tc_acts
+
+      test_lins = list(actions.values())
+      if FUZZ_ALL_ACTIONS: print(f"testing {lin.applied_opts=} with {len(actions)} actions")
+      elif opts_list is None: test_lins = [random.choice(test_lins)]
+
+      for test_lin in test_lins:
+        if not FUZZ_ALL_ACTIONS and test_lin.applied_opts: print(f"applied opts: {test_lin.applied_opts}")
+
+        # stop if kernel uops repeat
+        try: tuops = tuplize_uops(get_program(test_lin.get_optimized_ast(), test_lin.ren).uops)
+        except KeyboardInterrupt: raise
+        except BaseException as e:
+          print(test_lin.ast)
+          print(test_lin.applied_opts)
+          print(e)
+          failures["LINEARIZE_ERROR"].append((test_lin.ast, test_lin.applied_opts))
+          continue
+
+        if tuops in seen_uops: continue
+        seen_uops[tuops] = tuple(test_lin.applied_opts)
+
+        if not FUZZ_ALL_ACTIONS: print(test_lin.colored_shape())
+
+        (msg, rawbufs, var_vals, ground_truth, state1) = compare_linearizer(test_lin, rawbufs, var_vals, ground_truth, rtol=rtol, atol=atol)
+        if state1 is not None and validate_device is not None:
+          validate_lin = test_lin.copy()
+          validate_lin.ren = validate_device.renderer
+          if validate_rawbufs is None:
+            validate_rawbufs = [get_fuzz_rawbuf_like(x, copy=True, force_device=validate_device.device) for x in rawbufs]
+          (_msg, _, _, _, state2) = compare_linearizer(validate_lin, validate_rawbufs, var_vals, ground_truth, rtol=rtol, atol=atol)
+
+          if _msg != "PASS": failures[f"VALIDATE_DEV_{_msg}"].append((validate_lin.ast, validate_lin.applied_opts))
+
+          ok, err_msg = compare_states(state1, state2)
+          if not ok: failures["HCQ_COMPARE_FAILURE"].append((err_msg, test_lin.ast, test_lin.applied_opts, state1, state2))
+
+        if msg != "PASS":
+          print(test_lin.ast)
+          print(test_lin.applied_opts)
+          print(msg)
+          failures[msg].append((test_lin.ast, test_lin.applied_opts))
+          continue
+
+        next_lins.append(test_lin)
+
+    last_lins = next_lins
+    if FUZZ_ALL_ACTIONS: print(f"depth={depth} total_lins={len(last_lins)} {failures=}")
+  return failures
+
+def _is_simple(lin: Kernel) -> bool:
+  if len(lin.ast.src) > 1: return False
+  ast:UOp = lin.ast.src[0]
+  if ast.src[0].op is Ops.CAST and ast.src[0].src[0].op is Ops.LOAD: return True
+  return False
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser(description="Run a fuzz testing on one or more kernels", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+  parser.add_argument("--ast", type=str, default=None, help="the ast for the kernel to be optimized")
+  parser.add_argument("--file", type=str, default=None, help="a file containing asts to be optimized, one per line")
+  parser.add_argument("--beamreplay", type=str, default=None, help="replay asts and opts got from beam with CAPTURE_BEAM")
+  parser.add_argument("--logfile", type=str, default=None, help="a file containing a tuple of ast and applied_opts, one per line")
+  parser.add_argument("--expected-failures", type=int, default=0, help="the number of expected failed kernels")
+  parser.add_argument("--rtol", type=float, default=1e-2, help="relative tolerance for numerical comparison")
+  parser.add_argument("--atol", type=float, default=1e-2, help="absolute tolerance for numerical comparison")
+  args = parser.parse_args()
+
+  opts_list = None
+  if args.ast is not None:
+    print("loaded AST from CLI")
+    ast_strs = [args.ast]
+  elif args.file is not None:
+    print(f"loading ASTs from file '{args.file}'")
+    with open(args.file, 'r') as file:
+      ast_strs = file.readlines()
+  elif args.beamreplay is not None:
+    print(f"loading BEAM replay from file '{args.beamreplay}'")
+    with open(args.beamreplay, 'r') as file: fdata = file.readlines()
+    ast_strs, opts_list = [x.split(' :: ')[0] for x in fdata if not x.startswith("#")], [x.split(' :: ')[1] for x in fdata if not x.startswith("#")]
+
+    # dedup ast_strs and opts_list
+    dct = defaultdict(list)
+    for i in range(len(ast_strs)): dct[ast_strs[i]].append(eval(opts_list[i]))
+    ast_strs_items = list(dct.keys())
+    opts_list = [dct[c] for c in ast_strs_items]
+  elif args.logfile is not None:
+    print(f"loading ASTs from LOGKERNS file '{args.file}'")
+    with open(args.logfile, 'r') as file:
+      kern_strs = file.readlines()
+      test_lins = [kern_str_to_lin(kern_str) for kern_str in kern_strs]
+      ast_strs = [f"{lin.ast}" for lin in test_lins]
+  else:
+    print("loading ASTs from world")
+    ast_strs = load_worlds(filter_reduce=False, filter_novariable=False)
+
+  print(f"{len(ast_strs)=}")
+  tested = 0
+  failed_ids = []
+  failures = defaultdict(list)
+  seen_ast_strs = set()
+
+  try:
+    for i, ast in enumerate(ast_strs[:getenv("FUZZ_N", len(ast_strs))]):
+      if (nth := getenv("FUZZ_NTH", -1)) != -1 and i != nth: continue
+      if getenv("FUZZ_IMAGEONLY") and "dtypes.image" not in ast: continue
+      if "dtypes.image" in ast and Device.DEFAULT not in {"CL", "QCOM"}: continue  # IMAGE is only for CL
+      if ast in seen_ast_strs: continue
+      seen_ast_strs.add(ast)
+
+      lin = ast_str_to_lin(ast)
+      if not all(is_dtype_supported(buf.dtype) for buf in lin.bufs):
+        print("skipping kernel due to not supported dtype")
+        continue
+
+      with Timing(f"tested ast {i}: "):
+        tested += 1
+        fuzz_failures = fuzz_linearizer(lin, rtol=args.rtol, atol=args.atol, opts_list=(opts_list[i] if opts_list else None))
+        if fuzz_failures: failed_ids.append(i)
+        for k, v in fuzz_failures.items():
+          for f in v:
+            failures[k].append(f)
+  except KeyboardInterrupt: print(colored("STOPPING...", 'red'))
+
+  for msg, errors in failures.items():
+    for i, payload in enumerate(errors):
+      print(f"{msg} {i} kernel: {payload}") # easier to use with output with verify_kernel.py
+
+  print(f"{tested=}")
+  if failures:
+    print(f"{failed_ids=}")
+    for msg, errors in failures.items():
+      print(f"{msg}: {len(errors)}")
+    if len(failed_ids) == args.expected_failures:
+      print(colored(f"{len(failed_ids)} failed as expected", "yellow"))
+  if len(failed_ids) != args.expected_failures:
+    print(colored(f"failed on {len(failed_ids)} kernels, expected {args.expected_failures}", "red"))
+    # TODO: fix this
+    # raise RuntimeError(f"failed on {len(failed_ids)} kernels, expected {args.expected_failures}")
+  else:
+    print(colored("all passed", "green"))

test/external/speed_compare_amd_am.py

@@ -0,0 +1,157 @@
+from tinygrad import Device, dtypes
+from tinygrad.helpers import getenv, colorize_float, DEBUG
+from extra.optimization.helpers import load_worlds, ast_str_to_lin
+from test.external.fuzz_linearizer import get_fuzz_rawbufs
+from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
+from tinygrad.codegen.opt.search import bufs_from_lin
+from tinygrad.engine.realize import CompiledRunner, get_program
+from tinygrad.tensor import _to_np_dtype
+from tinygrad.runtime.ops_amd import AMDDevice
+from contextlib import contextmanager
+import numpy as np
+import os, random, statistics
+
+am_signal_pages, am_signal_pool, am_devices = [], [], []
+amd_signal_pages, amd_signal_pool, amd_devices = [], [], []
+
+def rebind_vfio(pcibus="0000:44:00.0"):
+  print("rebind ", pcibus)
+  os.system("sudo rmmod amdgpu")
+  os.system("sudo modprobe vfio-pci")
+
+  base = f"/sys/bus/pci/devices/{pcibus}"
+  if os.path.exists(f"{base}/driver"):
+    with open(f"{base}/driver/unbind", "w") as f: f.write(pcibus)
+  with open(f"{base}/driver_override", "w") as f: f.write("vfio-pci")
+  with open("/sys/bus/pci/drivers_probe", "w") as f: f.write(pcibus)
+
+  os.system("sudo modprobe amdgpu")
+  os.system("rocm-smi --setprofile compute")
+  os.system("rocm-smi --setperflevel high")
+
+@contextmanager
+def run_amd():
+  global amd_signal_pages, amd_signal_pool, amd_devices
+  AMDDevice.driverless = False
+  AMDDevice.signal_pages, AMDDevice.signal_pool, AMDDevice.devices = amd_signal_pages, amd_signal_pool, amd_devices
+  yield
+  amd_signal_pages, amd_signal_pool, amd_devices = AMDDevice.signal_pages, AMDDevice.signal_pool, AMDDevice.devices
+  AMDDevice.signal_pages, AMDDevice.signal_pool, AMDDevice.devices = [], [], []
+
+@contextmanager
+def run_am():
+  global am_signal_pages, am_signal_pool, am_devices
+  AMDDevice.driverless = True
+  AMDDevice.signal_pages, AMDDevice.signal_pool, AMDDevice.devices = am_signal_pages, am_signal_pool, am_devices
+  yield
+  am_signal_pages, am_signal_pool, am_devices = AMDDevice.signal_pages, AMDDevice.signal_pool, AMDDevice.devices
+  AMDDevice.signal_pages, AMDDevice.signal_pool, AMDDevice.devices = [], [], []
+
+if __name__ == "__main__":
+  CHECK_CPU = getenv("CHECK_CPU", 0)
+  SEED = getenv("SEED", 42)
+  CNT = getenv("CNT", 7)
+  random.seed(SEED)
+  np.random.seed(SEED)
+
+  # TODO: NUM=780 is super slow
+  # NUM=1907 is broken on AMD and AM have some mismatches (0 vs 1)
+  # kfd feels so bad when taking gpu out while it's running... Need hacks to rebind it before running.
+  rebind_vfio(pcibus="0000:44:00.0")
+
+  ast_strs = load_worlds(filter_reduce=False, filter_novariable=True)
+
+  with run_am():
+    amdev = Device["AMD:1"]
+
+  with run_amd():
+    amddev = Device["AMD"]
+
+  if CHECK_CPU: cpudev = Device["CPU"]
+
+  single = getenv("NUM", -1)
+  if single != -1: ast_strs = ast_strs[single:single+1]
+
+  average_tm_amd, average_tm_am = 0, 0
+  for num,ast in enumerate(ast_strs):
+    with run_amd():
+      amdlin = ast_str_to_lin(ast, opts=amddev.renderer)
+      amdlin.apply_opts(hand_coded_optimizations(amdlin))
+      has_bf16 = any(b.dtype == dtypes.bfloat16 for b in amdlin.bufs)
+
+      amd_prg = CompiledRunner(get_program(amdlin.get_optimized_ast(), amdlin.opts))
+      amdbufs = bufs_from_lin(amdlin)
+      test_amdbufs = get_fuzz_rawbufs(amdlin) if not has_bf16 else amdbufs
+      if not has_bf16: contents = [buf.as_buffer() for buf in test_amdbufs]
+
+    with run_am():
+      rdr = amdev.renderer
+      rdr.device = "AMD:1"
+      amlin = ast_str_to_lin(ast, opts=amdev.renderer)
+      amlin.apply_opts(hand_coded_optimizations(amlin))
+      am_prg = CompiledRunner(get_program(amlin.get_optimized_ast(), amlin.opts))
+      ambufs = bufs_from_lin(amlin)
+      test_ambufs = get_fuzz_rawbufs(amlin) if not has_bf16 else ambufs
+      if not has_bf16:
+        for i,rawbuf in enumerate(test_ambufs): rawbuf.copyin(contents[i])
+
+    if CHECK_CPU:
+      cpu_rdr = cpudev.renderer
+      cpu_rdr.device = "CPU"
+      cpulin = ast_str_to_lin(ast, opts=cpu_rdr)
+      cpulin.apply_opts(hand_coded_optimizations(cpulin))
+      cpu_prg = CompiledRunner(get_program(cpulin.get_optimized_ast(), cpulin.opts))
+      cpubufs = bufs_from_lin(cpulin)
+      test_cpubufs = get_fuzz_rawbufs(cpulin) if not has_bf16 else ambufs
+      if not has_bf16:
+        for i,rawbuf in enumerate(test_cpubufs): rawbuf.copyin(contents[i])
+
+    # warmup
+    tm_amd, tm_am, failed = [], [], False
+    with run_amd():
+      try:
+        amd_prg(test_amdbufs, {}, wait=True)
+        for i in range(CNT): tm_amd.append(amd_prg(amdbufs, {}, wait=True))
+      except RuntimeError:
+        print("AMD FAILED")
+        tm_amd = [1e9]
+        failed = True
+
+    with run_am():
+      try:
+        am_prg(test_ambufs, {}, wait=True)
+        for i in range(CNT): tm_am.append(am_prg(ambufs, {}, wait=True))
+      except RuntimeError:
+        print("AM FAILED")
+        tm_am = [1e9]
+        failed = True
+
+    if CHECK_CPU:
+      cpu_prg(test_cpubufs, {}, wait=True)
+      for i in range(1): cpu_prg(cpubufs, {}, wait=True)
+
+    if not failed and not has_bf16:
+      with run_amd():
+        curesult = np.frombuffer(test_amdbufs[0].as_buffer(), _to_np_dtype(test_amdbufs[0].dtype))
+
+      with run_am():
+        amresult = np.frombuffer(test_ambufs[0].as_buffer(), _to_np_dtype(test_ambufs[0].dtype))
+
+      if CHECK_CPU:
+        cpuresult = np.frombuffer(test_cpubufs[0].as_buffer(), _to_np_dtype(test_cpubufs[0].dtype))
+        np.testing.assert_allclose(amresult, cpuresult, rtol=1e-2, atol=1e-2)
+        np.testing.assert_allclose(curesult, cpuresult, rtol=1e-2, atol=1e-2)
+
+      try:
+        np.testing.assert_allclose(curesult, amresult, rtol=1e-2, atol=1e-2)
+      except AssertionError as e:
+        print("AM and AMD results do not match")
+        print(e)
+
+    bam = statistics.median(tm_am)
+    bamd = statistics.median(tm_amd)
+    average_tm_amd += bamd
+    average_tm_am += bam
+    ratio = bam/bamd
+    print(f"{average_tm_am/average_tm_amd:5.2f}x -- {num:4d} {colorize_float(ratio)} {bam*1e6:7.2f} vs {bamd*1e6:7.2f} us", amlin.name)
+    if DEBUG > 3 and ratio > 1.04: print(f"AM slower {ratio}", amlin.ast, amlin.applied_opts)

test/external/speed_compare_cuda_nv.py

@@ -0,0 +1,69 @@
+from tinygrad import Device, dtypes
+from tinygrad.helpers import getenv, colorize_float
+from extra.optimization.helpers import load_worlds, ast_str_to_lin
+from test.external.fuzz_linearizer import get_fuzz_rawbufs
+from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
+from tinygrad.codegen.opt.search import bufs_from_lin
+from tinygrad.engine.realize import CompiledRunner, get_program
+from tinygrad.tensor import _to_np_dtype
+import numpy as np
+
+if __name__ == "__main__":
+  ast_strs = load_worlds(filter_reduce=False, filter_novariable=True)
+  cudev = Device["CUDA"]
+  nvdev = Device["NV"]
+
+  # NUM=112 python3 test/external/speed_compare_cuda_nv.py
+
+  single = getenv("NUM", -1)
+  if single != -1: ast_strs = ast_strs[single:single+1]
+
+  average_tm_cuda, average_tm_nv = 0, 0
+  for num,ast in enumerate(ast_strs):
+    # cuda compile
+    culin = ast_str_to_lin(ast, opts=cudev.renderer)
+    culin.apply_opts(hand_coded_optimizations(culin))
+    has_bf16 = any(b.dtype == dtypes.bfloat16 for b in culin.bufs)
+
+    cuda_prg = CompiledRunner(get_program(culin.get_optimized_ast(), culin.opts))
+    cubufs = bufs_from_lin(culin)
+    test_cubufs = get_fuzz_rawbufs(culin) if not has_bf16 else cubufs
+
+    rdr = nvdev.renderer
+    rdr.device = "NV"
+    nvlin = ast_str_to_lin(ast, opts=rdr)
+    nvlin.apply_opts(hand_coded_optimizations(nvlin))
+    nv_prg = CompiledRunner(get_program(nvlin.get_optimized_ast(), nvlin.opts))
+    nvbufs = bufs_from_lin(nvlin)
+    test_nvbufs = get_fuzz_rawbufs(nvlin) if not has_bf16 else nvbufs
+    if not has_bf16:
+      for i,rawbuf in enumerate(test_nvbufs): rawbuf.copyin(test_cubufs[i].as_buffer())
+
+    # warmup
+    tm_cuda, tm_nv, failed = [], [], False
+    try:
+      cuda_prg(test_cubufs, {}, wait=True)
+      for i in range(5): tm_cuda.append(cuda_prg(cubufs, {}, wait=True))
+    except RuntimeError:
+      print("CUDA FAILED")
+      tm_cuda = [1e9]
+      failed = True
+
+    try:
+      nv_prg(test_nvbufs, {}, wait=True)
+      for i in range(5): tm_nv.append(nv_prg(nvbufs, {}, wait=True))
+    except RuntimeError:
+      print("NV FAILED")
+      tm_nv = [1e9]
+      failed = True
+
+    if not failed and not has_bf16:
+      curesult = np.frombuffer(test_cubufs[0].as_buffer(), _to_np_dtype(test_cubufs[0].dtype))
+      nvresult = np.frombuffer(test_nvbufs[0].as_buffer(), _to_np_dtype(test_nvbufs[0].dtype))
+      np.testing.assert_allclose(curesult, nvresult, rtol=1e-2, atol=1e-2)
+
+    average_tm_cuda += min(tm_cuda)
+    average_tm_nv += min(tm_nv)
+    ratio = min(tm_nv)/min(tm_cuda)
+    print(f"{average_tm_nv/average_tm_cuda:5.2f}x -- {num:4d} {colorize_float(ratio)}  {min(tm_nv)*1e6:7.2f} us", nvlin.name)
+    if ratio > 1.04: print(f"NV slower {ratio}", nvlin.ast, nvlin.applied_opts)

test/external/speed_compare_cuda_ptx.py

@@ -0,0 +1,60 @@
+import itertools
+from tinygrad import Device
+from tinygrad.engine.realize import CompiledRunner, get_program
+from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
+from tinygrad.helpers import getenv, colorize_float
+from extra.optimization.helpers import load_worlds, ast_str_to_lin
+from tinygrad.codegen.opt.search import bufs_from_lin
+from tinygrad.runtime.ops_cuda import PTXCompiler, PTXRenderer, CUDACompiler
+
+if __name__ == "__main__":
+  ast_strs = load_worlds(filter_reduce=False, filter_novariable=True)
+  # no bfloat16 for ptx at the moment
+  ast_strs = [x for x in ast_strs if "dtypes.bfloat16" not in x]
+  dev = Device["CUDA"]
+  ptx = PTXRenderer(dev.arch)
+
+  # NUM=112 python3 test/external/speed_compare_cuda_ptx.py
+
+  single = getenv("NUM", -1)
+  if single != -1: ast_strs = ast_strs[single:single+1]
+
+  average_tm_cuda, average_tm_ptx = 0, 0
+  for num,ast in enumerate(ast_strs):
+    # cuda compile
+    dev.compiler = CUDACompiler(dev.arch)
+    lin = ast_str_to_lin(ast, opts=dev.renderer)
+    lin.apply_opts(hand_coded_optimizations(lin))
+    cuda_prg = CompiledRunner(get_program(lin.get_optimized_ast(), lin.opts))
+
+    bufs = bufs_from_lin(lin)
+
+    # ptx compile
+    dev.compiler = PTXCompiler(dev.arch)
+    lin = ast_str_to_lin(ast, opts=ptx)
+    lin.apply_opts(hand_coded_optimizations(lin))
+    ptx_prg = CompiledRunner(get_program(lin.get_optimized_ast(), lin.opts))
+
+    # warmup
+    try:
+      cuda_prg(bufs, {}, wait=True)
+    except RuntimeError:
+      print("cuda failed ast:", num)
+      continue
+    ptx_prg(bufs, {}, wait=True)
+
+    tm_cuda, tm_ptx = [], []
+    for i in range(5):
+      tm_cuda.append(cuda_prg(bufs, {}, wait=True))
+      tm_ptx.append(ptx_prg(bufs, {}, wait=True))
+    average_tm_cuda += min(tm_cuda)
+    average_tm_ptx += min(tm_ptx)
+    ratio = min(tm_ptx)/min(tm_cuda)
+    print(f"{average_tm_ptx/average_tm_cuda:5.2f}x -- {num:4d} {colorize_float(ratio)}  {min(tm_ptx)*1e6:7.2f} us", lin.name)
+    if ratio > 1.5:
+      def fix(x): return x.replace('\t', ' ').strip()
+      ll1, ll2 = cuda_prg.lib.decode().split('\n'), ptx_prg.lib.decode().split('\n')
+      if single != -1:
+        for ln, (l1, l2) in enumerate(itertools.zip_longest(ll1, ll2, fillvalue='')):
+          print(f"{ln:5d} | {fix(l1):80s} | {fix(l2):80s}")
+      print(len(ll1), len(ll2), "RATIO", ratio, "us", min(tm_ptx)*1e6)

test/external/verify_kernel.py

@@ -0,0 +1,78 @@
+import argparse
+from collections import defaultdict
+from extra.optimization.helpers import kern_str_to_lin, time_linearizer
+from test.external.fuzz_linearizer import compare_linearizer
+from tinygrad.helpers import colored
+from tinygrad.codegen.opt.kernel import Kernel
+
+# Use this with the LOGKERNS options to verify that all executed kernels are valid and evaluate to the same ground truth results
+
+# Example for GPT2:
+# 1) Run the model to log all kernels: `PYTHONPATH=. LOGKERNS=/tmp/gpt2_kerns.txt JIT=1 HALF=1 BEAM=2 CACHELEVEL=0 python3 examples/gpt2.py --count 10 --temperature 0 --timing`   # noqa: E501
+# 2) Validate the kernel correctness: `PYTHONPATH=. python3 ./test/external/verify_kernel.py --file /tmp/gpt2_kerns.txt`
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser(description="Verify the correctness of one or more kernel", formatter_class=argparse.ArgumentDefaultsHelpFormatter)    # noqa: E501
+  parser.add_argument("--kernel", type=str, default=None, help="a string of a tuple of (ast, applied_opts,)")
+  parser.add_argument("--file", type=str, default=None, help="a file containing a tuple of ast and applied_opts, one per line")
+  parser.add_argument("--pkl", type=str, default=None, help="a pickle file containing a single tuple of ast and applied_opts")
+  parser.add_argument("--rtol", type=float, default=1e-2, help="relative tolerance for numerical comparison")
+  parser.add_argument("--atol", type=float, default=1e-2, help="absolute tolerance for numerical comparison")
+  parser.add_argument("--timing", action='store_true', help="show final timing for the kernel")
+  parser.add_argument("--expected-failures", type=int, default=0, help="the number of expected failed kernels")
+  args = parser.parse_args()
+
+  if args.kernel is not None:
+    print("loading kernel from args")
+    test_lins = [kern_str_to_lin(args.kernel)]
+  elif args.file is not None:
+    print(f"loading kernel from file '{args.file}'")
+    with open(args.file, 'r') as file:
+      kern_strs = file.readlines()
+      test_lins = [kern_str_to_lin(kern_str) for kern_str in kern_strs]
+  elif args.pkl is not None:
+    print(f"loading kernel from pickle file '{args.file}'")
+    import pickle
+    with open(args.pkl, 'rb') as file:
+      (ast, applied_opts,) = pickle.load(file)
+      lin = Kernel(ast)
+      lin.apply_opts(applied_opts)
+      test_lins = [lin]
+
+  else:
+    raise RuntimeError("no kernel specified; use --kernel, --file, or --pkl options")
+
+  print(f"verifying {len(test_lins)} kernels")
+
+  failed_ids = []
+  failures = defaultdict(list)
+  for i, test_lin in enumerate(test_lins):
+    print(f"testing kernel {i}")
+    print(test_lin.ast)
+    print(test_lin.applied_opts)
+    unoptimized_lin = Kernel(test_lin.ast)
+    print(f"{unoptimized_lin.colored_shape()} -> {test_lin.colored_shape()}")
+    (msg,rb,vv,gt) = compare_linearizer(test_lin, None, None, None, rtol=args.rtol, atol=args.atol)
+    if msg != "PASS":
+      failed_ids.append(i)
+      failures[msg].append((test_lin.ast, test_lin.applied_opts))
+    if args.timing:
+      tm = time_linearizer(test_lin, rb, allow_test_size=False, cnt=10)
+      print(f"final time {tm*1e6:9.0f} us")
+
+  for msg, errors in failures.items():
+    for i, (ast, opts) in enumerate(errors):
+      print(f"{msg} {i} AST: {ast}")
+      print(f"{msg} {i} OPTS: {opts}\n")
+
+  print(f"tested {len(test_lins)} kernels")
+  if failures:
+    print(f"{failed_ids=}")
+    for msg, errors in failures.items():
+      print(f"{msg}: {len(errors)}")
+    if len(failed_ids) == args.expected_failures:
+      print(colored(f"{len(failed_ids)} failed as expected", "yellow"))
+  if len(failed_ids) != args.expected_failures:
+    raise RuntimeError(f"failed on {len(failed_ids)} kernels, expected {args.expected_failures}")
+  else:
+    print(colored("all passed", "green"))

test/helpers.py

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

test/test_linearizer.py

@@ -188,8 +188,9 @@ class TestLinearizer(unittest.TestCase):
     assert any(x.op is Ops.DEFINE_LOCAL for x in stores[0].toposort())
     # the second store is to gds with no upcasts
     assert stores[1].src[1].dtype == dtypes.float
-    assert any(x.op is Ops.PARAM for x in stores[1].toposort())
+    assert any(x.op is Ops.DEFINE_GLOBAL for x in stores[1].toposort())
 
+  @unittest.skipIf(Device.DEFAULT=="CPU", "CPU splits the cat so cant upcast")
   def test_zero_fold(self):
     a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize()
     r = Tensor.stack(a, b)
@@ -449,7 +450,7 @@ class TestLinearizer(unittest.TestCase):
     def get_recursive(uop): return set.union(set(uop.src), [uop], *[get_recursive(v) for v in uop.src])
     uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=opt).uops
     local_stores = [u for u in uops if u.op is Ops.STORE and any(x.op is Ops.DEFINE_LOCAL for x in get_recursive(u.src[0]))]
-    global_stores = [u for u in uops if u.op is Ops.STORE and any(x.op is Ops.PARAM for x in get_recursive(u.src[0]))]
+    global_stores = [u for u in uops if u.op is Ops.STORE and any(x.op is Ops.DEFINE_GLOBAL for x in get_recursive(u.src[0]))]
     barrier = [u for u in uops if u.op is Ops.BARRIER]
     assert len(barrier) == 1
     # check that the float4 cast collapses for all stores

test/test_linearizer_dumb.py

@@ -11,16 +11,16 @@ from tinygrad.engine.realize import get_program
 class TestLinearizerFailure(unittest.TestCase):
   @unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL")
   def test_failure_beam_mnist(self):
-    c0 = UOp(Ops.PARAM, dtypes.uchar.ptr(4014080), arg=0, src=())
+    c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.uchar.ptr(4014080), arg=0, src=())
     c1 = UOp.range(UOp.const(dtypes.index, 512), 0, AxisType.GLOBAL)
     c2 = UOp.range(UOp.const(dtypes.index, 784), 1, AxisType.GLOBAL)
     c3 = UOp.range(UOp.const(dtypes.index, 10), 3, AxisType.GLOBAL)
-    c4 = UOp(Ops.PARAM, dtypes.int.ptr(512), arg=1, src=())
+    c4 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(512), arg=1, src=())
     c5 = c4.index(c1.valid(UOp.const(dtypes.bool, True)))
     c6 = UOp.range(UOp.const(dtypes.index, 6000), 1004, AxisType.REDUCE)
     c7 = UOp.range(UOp.const(dtypes.index, 3750), 2006, AxisType.REDUCE)
     c8 = UOp.range(UOp.const(dtypes.index, 16), 2007, AxisType.GROUP_REDUCE)
-    c9 = UOp(Ops.PARAM, dtypes.uchar.ptr(47040000), arg=2, src=())
+    c9 = UOp(Ops.DEFINE_GLOBAL, dtypes.uchar.ptr(47040000), arg=2, src=())
     c10 = c9.index((((c3*UOp.const(dtypes.index, 4704000))+c2)+(c6*UOp.const(dtypes.index, 784))).valid(UOp.const(dtypes.bool, True)))
     c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(dtypes.index, 6000))+c6)+((c7*UOp.const(dtypes.index, 16))+c8)).alu(Ops.CMPLT, UOp.const(dtypes.index, 59999)).where(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(dtypes.int, -1))).where(UOp.const(dtypes.uchar, 0), c10).reduce(c6, arg=Ops.ADD)
     c12 = c0.index((((c1*UOp.const(dtypes.index, 7840))+(c2*UOp.const(dtypes.index, 10)))+c3).valid(UOp.const(dtypes.bool, True))).store(c11).end(c1, c2, c3)

test/test_linearizer_failures.py

@@ -7,14 +7,14 @@ from tinygrad.device import Device
 
 class TestLinearizerFailures(unittest.TestCase):
   def test_fail_1(self):
-    c0 = UOp(Ops.PARAM, dtypes.float.ptr(64), arg=0, src=())
+    c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(64), arg=0, src=())
     c1 = UOp.range(UOp.const(dtypes.index, 2), 1, AxisType.LOOP)
     c2 = UOp.range(UOp.const(dtypes.index, 32), 2, AxisType.LOOP)
     c3 = ((c1*UOp.const(dtypes.index, 32))+c2)
-    c4 = UOp(Ops.PARAM, dtypes.float.ptr(163840), arg=1, src=())
+    c4 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(163840), arg=1, src=())
     c5 = UOp.range(UOp.const(dtypes.index, 2560), 0, AxisType.REDUCE)
     c6 = c4.index(((((((c5//UOp.const(dtypes.index, 8))%UOp.const(dtypes.index, 8))*UOp.const(dtypes.index, 8))+(c5%UOp.const(dtypes.index, 8)))+(((c2*UOp.const(dtypes.index, 40))+(c5//UOp.const(dtypes.index, 64)))*UOp.const(dtypes.index, 64)))+(c1*UOp.const(dtypes.index, 81920))))
-    c7 = UOp(Ops.PARAM, dtypes.float.ptr(64), arg=2, src=())
+    c7 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(64), arg=2, src=())
     c8 = c7.index(c3)
     c9 = ((((c6+(c8*UOp.const(dtypes.float, -1.0)))*(c6+(c8*UOp.const(dtypes.float, -1.0)))).reduce(c5, arg=Ops.ADD)*UOp.const(dtypes.float, 0.000390625))+UOp.const(dtypes.float, 1e-05)).sqrt().reciprocal()
     c10 = c0.index(c3).store(c9).end(c1, c2)

test/test_renderer_failures.py

@@ -26,8 +26,8 @@ def _test_uop_result(inputs:list[Tensor], prg, local_size=None):
 
 def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp):
   dtype = alu_src_uops[0].dtype
-  a = UOp(Ops.PARAM, dtype.ptr(), (), 0)
-  b = UOp(Ops.PARAM, dtype.ptr(), (), 1)
+  a = UOp(Ops.DEFINE_GLOBAL, dtype.ptr(), (), 0)
+  b = UOp(Ops.DEFINE_GLOBAL, dtype.ptr(), (), 1)
   idx = UOp.const(dtypes.int, 0)
   ld = b.index(idx)
   alu = ld.alu(alu_op, *alu_src_uops)
@@ -39,7 +39,7 @@ def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp):
 class TestRendererFailures(unittest.TestCase):
   @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
   def test_gated_store_with_alu(self):
-    a = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
+    a = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
     gate_alu = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0)
     gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index(lidx0.valid(gate_alu)), UOp.const(dtypes.int, 1)))
     sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,))
@@ -49,7 +49,7 @@ class TestRendererFailures(unittest.TestCase):
 
   @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
   def test_gated_store_with_alu_2d(self):
-    a = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
+    a = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
     gate_alu_0 = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0)
     gate_alu_1 = (lidx1:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 2),), 'lidx1')).ne(0)
     gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index((lidx0+lidx1*4).valid(gate_alu_0&gate_alu_1)), UOp.const(dtypes.int, 1)))
@@ -94,7 +94,7 @@ class TestWGSLFailures(unittest.TestCase):
 class TestPTXFailures(unittest.TestCase):
   @unittest.skip("INDEX can only have a gate ALU parent, not an IF")
   def test_gated_store_with_if(self):
-    a = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
+    a = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
     gate_alu = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0)
     val = UOp.const(dtypes.int, 1)
     if_uop = UOp(Ops.IF, dtypes.void, (gate_alu,))

test/test_schedule.py

@@ -1092,18 +1092,6 @@ class TestSchedule(unittest.TestCase):
     np.testing.assert_allclose(out[0].numpy(), np.sqrt(np.square(x.numpy() - np_mu).sum(-1)/x.shape[-1]), atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(out[1].numpy(), np.sqrt(np.square(y.numpy() - np_mu).sum(-1)/y.shape[-1]), atol=1e-4, rtol=1e-4)
 
-  def test_cumsum_parallel_reduce_fused(self):
-    # two-stage cumsum + ops triggers parallel REDUCEs in one kernel that must share an END
-    step, num_steps = 513, 10
-    t = Tensor.arange(step).float().realize()
-    phase = t.cumsum()
-    tiled = phase.repeat((num_steps,)).reshape(num_steps, step)
-    pattern = Tensor([1,0,0,1,0,0,0,0,1,0]).reshape(num_steps, 1)
-    out = (tiled * pattern).flatten()
-    expected = np.tile(np.arange(step).astype(np.float32).cumsum(), num_steps).reshape(num_steps, step)
-    expected = (expected * np.array([1,0,0,1,0,0,0,0,1,0]).reshape(num_steps, 1)).flatten()
-    np.testing.assert_allclose(out.numpy(), expected, atol=1e-4, rtol=1e-4)
-
   def test_multimatmul_fusion(self):
     Tensor.manual_seed(0)
     a,b = Tensor.randn(4, 64).realize(), Tensor.rand(64,8).realize()

test/test_uops.py

@@ -36,8 +36,8 @@ def uop(uops:list[UOp], op:Ops, dtype:Optional[DType], src:tuple[UOp, ...], arg:
 def _test_single_value(vals, op, dts):
   uops = []
   output_dtype = dtypes.bool if op in (Ops.CMPLT, Ops.CMPNE) else dts[-1]
-  buf_store = uop(uops, Ops.PARAM, output_dtype.ptr(), (), 0)
-  buf_loads = [uop(uops, Ops.PARAM, dtype.ptr(), (), i+1) for i,dtype in enumerate(dts)]
+  buf_store = uop(uops, Ops.DEFINE_GLOBAL, output_dtype.ptr(), (), 0)
+  buf_loads = [uop(uops, Ops.DEFINE_GLOBAL, dtype.ptr(), (), i+1) for i,dtype in enumerate(dts)]
   loads = (buf_loads[i].index(uop(uops, Ops.CONST, dtypes.int32, (), 0)) for i, dtype in enumerate(dts))
   alu = uop(uops, op, output_dtype, loads)
   out = uop(uops, Ops.STORE, dtypes.void, (buf_store.index(uop(uops, Ops.CONST, dtypes.int32, (), 0), ptr=True), alu))
@@ -52,7 +52,7 @@ def _test_single_value(vals, op, dts):
 def _test_single_value_const(vals, op, dts):
   uops = []
   output_dtype = dtypes.bool if op in (Ops.CMPLT, Ops.CMPNE) else dts[-1]
-  buf_store = uop(uops, Ops.PARAM, output_dtype.ptr(), (), 0)
+  buf_store = uop(uops, Ops.DEFINE_GLOBAL, output_dtype.ptr(), (), 0)
   loads = (uop(uops, Ops.CONST, dtype, [], a) for a,dtype in zip(vals, dts))
   alu = uop(uops, op, output_dtype, loads)
   out = buf_store[UOp.const(dtypes.int32, 0)].store(alu)
@@ -65,7 +65,7 @@ def _test_single_value_const(vals, op, dts):
 
 def _test_uops_result(output_dtype, uops, res):
   # uops = []
-  buf_store = uop(uops, Ops.PARAM, output_dtype.ptr(), (), 0)
+  buf_store = uop(uops, Ops.DEFINE_GLOBAL, output_dtype.ptr(), (), 0)
   # res = output_fn(uops)
   out = uop(uops, Ops.STORE, dtypes.void, (buf_store.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), res))
   buf = Buffer(Device.DEFAULT, 1, output_dtype).allocate()
@@ -273,7 +273,7 @@ class TestConstantFolding(unittest.TestCase):
 
 class TestGatedStoreRewrite(unittest.TestCase):
   def test_tiny_gate_store(self):
-    gmem = UOp(Ops.PARAM, dtypes.float.ptr(), (), 0)
+    gmem = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
     gidx0 = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'gidx0')
     gate = gidx0<UOp.const(dtypes.int, 1)
     idx = UOp(Ops.INDEX, dtypes.float.ptr(), (gmem, (gidx0 * UOp.const(dtypes.int, 2)).valid(gate)))
@@ -289,8 +289,8 @@ class TestGatedStoreRewrite(unittest.TestCase):
     self.assertIs(gated_uops[-1].op, Ops.STORE)
 
   def test_gate_some_stores(self):
-    gmem0 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 0)
-    gmem1 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 1)
+    gmem0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
+    gmem1 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 1)
     gidx0 = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'gidx0')
     idx = gidx0 * UOp.const(dtypes.int, 2)
     idx0 = UOp(Ops.INDEX, dtypes.float.ptr(), (gmem0, idx.valid(gidx0<UOp.const(dtypes.int, 1))))
@@ -309,8 +309,8 @@ class TestGatedStoreRewrite(unittest.TestCase):
   # scaled down version of TestLinearizerDumb.test_unmerged_ifs
   @unittest.skip("we don't merge ifs anymore")
   def test_merge_ifs_alt(self):
-    gmem0 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 0)
-    gmem1 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 1)
+    gmem0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
+    gmem1 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 1)
     gidx0 = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'gidx0')
     idx = gidx0*UOp.const(dtypes.int, 2)
     gate = gidx0<UOp.const(dtypes.int, 1)
@@ -380,7 +380,7 @@ class TestLocalAccess(unittest.TestCase):
 class TestFastIdiv(unittest.TestCase):
   def test_division_power_of_two(self):
     for dt in (dtypes.int32, dtypes.uint32):
-      g = UOp(Ops.PARAM, dt.ptr(), (), 0)
+      g = UOp(Ops.DEFINE_GLOBAL, dt.ptr(), (), 0)
       c = UOp.const(dt, 2)
       l = g.index(c)
       a = UOp(Ops.IDIV, dt, (l, c))
@@ -392,7 +392,7 @@ class TestFastIdiv(unittest.TestCase):
 
   @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support long")
   def test_fast_idiv_and_mod(self):
-    g = UOp(Ops.PARAM, dtypes.uint32.ptr(), (), 0)
+    g = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(), (), 0)
     c = UOp.const(dtypes.uint, 3)
     l = g.index(c)
     a = UOp(Ops.IDIV, dtypes.uint, (l, c))
@@ -420,7 +420,7 @@ class TestFastIdiv(unittest.TestCase):
   @unittest.expectedFailure
   def test_fast_idiv_overflow(self):
     # This will be possible with a slightly different method for fast_idiv
-    g = UOp(Ops.PARAM, dtypes.uint32.ptr(), (), 0)
+    g = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(), (), 0)
     c = UOp.const(dtypes.uint, 7)
     l = UOp(Ops.LOAD, dtypes.uint, (g.index(c),))
     a = UOp(Ops.IDIV, dtypes.uint, (l, c))
@@ -431,7 +431,7 @@ class TestFastIdiv(unittest.TestCase):
     self.assertNotIn(Ops.IDIV, ops)
 
   def test_disable_fast_idiv(self):
-    g = UOp(Ops.PARAM, dtypes.uint32.ptr(), (), 0)
+    g = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(), (), 0)
     c = UOp.const(dtypes.uint, 3)
     l = g.index(c)
     a = UOp(Ops.IDIV, dtypes.uint, (l, c))
@@ -445,7 +445,7 @@ class TestFastIdiv(unittest.TestCase):
 @unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "This only tests assembly backends")
 class TestAssembly(unittest.TestCase):
   def test_bitshift_left(self):
-    g1 = UOp(Ops.PARAM, dtypes.int32.ptr(), (), 0)
+    g1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 0)
     c1 = UOp.const(dtypes.int, 2)
     c2 = UOp.const(dtypes.int, 3)
     l1 = g1.index(c1)
@@ -471,7 +471,7 @@ class TestAssembly(unittest.TestCase):
     self.assertEqual(len([x.op for x in uops if x.op is Ops.MULACC]), 4)
 
   def test_use_cmpeq(self):
-    g = UOp(Ops.PARAM, dtypes.uint32.ptr(), (), 0)
+    g = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(), (), 0)
     c = UOp.const(dtypes.uint, 7)
     comp = g.index(c).ne(c).ne(True)
     uops = to_uops_list([comp], ren=Device[Device.DEFAULT].renderer)
@@ -507,7 +507,7 @@ class TestUOpMethod(unittest.TestCase):
     self.assertEqual((gidx0*3+1).const_factor(), 1)
 
   def test_replace(self):
-    x = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
+    x = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
     self.assertIs(x.replace(arg=None).arg, None)
     with self.assertRaises(AssertionError): x.replace(field="a")
 

test/test_uops_stats.py

@@ -137,7 +137,7 @@ class TestUOpsStats(unittest.TestCase):
 
   #MULACC should have the same stats as MUL + ADD
   def test_mulacc(self):
-    globl = UOp(Ops.PARAM, dtypes.int.ptr(), tuple())
+    globl = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), tuple())
     o1 = UOp(Ops.CONST, dtypes.int, tuple(), 1)
     o2 = UOp(Ops.CONST, dtypes.int, tuple(), 2)
     u1 = globl.index(o1)
@@ -147,7 +147,7 @@ class TestUOpsStats(unittest.TestCase):
     u5 = UOp(Ops.ADD, dtypes.int, (u4,u3))
     uops = list(u5.toposort())
 
-    globl = UOp(Ops.PARAM, dtypes.int.ptr(), tuple())
+    globl = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), tuple())
     o1 = UOp(Ops.CONST, dtypes.int, tuple(), 1)
     o2 = UOp(Ops.CONST, dtypes.int, tuple(), 2)
     u1 = globl.index(o1)

test/testextra/test_tk.py

@@ -750,35 +750,6 @@ class TestTK(unittest.TestCase):
 
     fa_jitted = TinyJit(flash_attention)
 
-    for _ in range(10):
-      st = time.perf_counter()
-      out = fa_jitted(q, k, v, is_causal=False)
-      et = time.perf_counter() - st
-      attn_flops = 2 * B * H * N * N * D + \
-                   4 * B * H * N * N + \
-                   2 * B * H * N * N * D
-      print(f"{attn_flops/(et*1e9):2f} GFLOPS")
-    out = out.float().transpose(1, 2)
-
-    ref = q.scaled_dot_product_attention(k, v, is_causal=False, enable_gqa=True).float().transpose(1, 2)
-
-    np.testing.assert_allclose(out.numpy(), ref.numpy(), atol=2e-2, rtol=2e-2)
-
-  def test_fast_fa_causal(self):
-    from extra.thunder.tiny.fa import flash_attention
-
-    B, N, H, H_KV, D = 2, 8192, 32, 8, 128
-
-    with Context(DEBUG=0):
-      q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16).contiguous()
-      k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
-      v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
-      Tensor.realize(q, k, v)
-
-    q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
-
-    fa_jitted = TinyJit(flash_attention)
-
     for _ in range(10):
       st = time.perf_counter()
       out = fa_jitted(q, k, v, is_causal=True)
@@ -867,7 +838,7 @@ class TestTK(unittest.TestCase):
 
     np.testing.assert_allclose(q.grad.numpy(), q_ref.grad.numpy(), atol=2e-2, rtol=2e-2)
     np.testing.assert_allclose(v.grad.numpy(), v_ref.grad.numpy(), atol=2e-2, rtol=2e-2)
-    np.testing.assert_allclose(k.grad.numpy(), k_ref.grad.numpy(), atol=6e-2, rtol=2e-2)
+    np.testing.assert_allclose(k.grad.numpy(), k_ref.grad.numpy(), atol=5e-2, rtol=2e-2)
 
   def test_fast_fa_bwd_causal_jitted(self):
     from extra.thunder.tiny.fa import flash_attention

test/unit/test_call.py

@@ -36,22 +36,6 @@ class TestCall(unittest.TestCase):
     np.testing.assert_allclose(a.grad.numpy(), gt_a_grad, rtol=1e-5)
     np.testing.assert_allclose(b.grad.numpy(), gt_b_grad, rtol=1e-5)
 
-  def test_call_plus_backward_auto(self):
-    a = Tensor.ones(10, 10, requires_grad=True)
-    b = Tensor.ones(10, 10, requires_grad=True)
-
-    (a+b).mean().backward()
-    gt_a_grad = a.grad.numpy()
-    gt_b_grad = b.grad.numpy()
-    a.grad, b.grad = None, None
-
-    plus_fxn = UOp.param(0, dtypes.float, (10,10)) + UOp.param(1, dtypes.float, (10,10))
-    c = Tensor.call(a, b, fxn=plus_fxn)
-    c.mean().backward()
-
-    np.testing.assert_allclose(a.grad.numpy(), gt_a_grad, rtol=1e-5)
-    np.testing.assert_allclose(b.grad.numpy(), gt_b_grad, rtol=1e-5)
-
   def test_call_gemm(self):
     M, K, N = 4, 8, 4
     a = Tensor.randn(M, K)
@@ -74,23 +58,5 @@ class TestCall(unittest.TestCase):
 
     np.testing.assert_allclose(c.numpy(), a.numpy() @ b.numpy(), rtol=1e-5, atol=1e-6)
 
-  def test_call_complex_backward_auto(self):
-    # complex chain: (a*b + a).exp2() * b.reciprocal() - tests mul, add, exp2, reciprocal, param reuse
-    a = Tensor.randn(10, 10, requires_grad=True)
-    b = Tensor.randn(10, 10, requires_grad=True) + 2  # avoid div by zero
-    Tensor.realize(a, b)
-
-    ((a*b + a).exp2() * b.reciprocal()).mean().backward()
-    gt_a_grad, gt_b_grad = a.grad.numpy(), b.grad.numpy()
-    a.grad, b.grad = None, None
-
-    p0, p1 = UOp.param(0, dtypes.float, (10,10)), UOp.param(1, dtypes.float, (10,10))
-    complex_fxn = (p0*p1 + p0).exp2() * p1.reciprocal()
-    c = Tensor.call(a, b, fxn=complex_fxn)
-    c.mean().backward()
-
-    np.testing.assert_allclose(a.grad.numpy(), gt_a_grad, rtol=1e-5)
-    np.testing.assert_allclose(b.grad.numpy(), gt_b_grad, rtol=1e-5)
-
 if __name__ == '__main__':
   unittest.main()

test/unit/test_simplify_valid_idx.py

@@ -5,17 +5,17 @@ from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.uop.symbolic import simplify_valid
 from tinygrad.helpers import Context
-from test.null.test_uop_symbolic import check_uop_against_string
+from test.unit.test_uop_symbolic import check_uop_against_string
 
 def get_gated_load_uop(valid:UOp, idx:UOp):
   return UOp(Ops.LOAD, dtypes.float, (
-    UOp(Ops.PARAM, dtypes.float.ptr(), arg=0).index(idx.valid(valid), ptr=True),
+    UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=0).index(idx.valid(valid), ptr=True),
     UOp.const(dtypes.float, 0.0)
   ))
 
 def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UOp]):
   return UOp(Ops.LOAD, dtypes.float.vec(4), (
-    UOp(Ops.PARAM, dtypes.imagef(image_shape), arg=0).index(UOp(Ops.VECTORIZE, dtypes.index.vec(2), idx).valid(valid), ptr=True),
+    UOp(Ops.DEFINE_GLOBAL, dtypes.imagef(image_shape), arg=0).index(UOp(Ops.VECTORIZE, dtypes.index.vec(2), idx).valid(valid), ptr=True),
     UOp(Ops.VECTORIZE, dtypes.float.vec(4), src=(UOp.const(dtypes.float, 0.0),) * 4)
   ))
 
@@ -461,13 +461,13 @@ class TestUnfoldableImageChannelSelection(unittest.TestCase):
   def test_bounded_channel_no_nan(self):
     # unfoldable image load with bounded idx % 4 range [0,1] -> no NAN fallback needed
     lidx = Special("lidx", 2)
-    load = UOp(Ops.LOAD, dtypes.float, (UOp(Ops.PARAM, dtypes.imagef((10, 10, 4)), arg=0).index(lidx, ptr=True), UOp.const(dtypes.float, 0)))
+    load = UOp(Ops.LOAD, dtypes.float, (UOp(Ops.DEFINE_GLOBAL, dtypes.imagef((10, 10, 4)), arg=0).index(lidx, ptr=True), UOp.const(dtypes.float, 0)))
     self.assertEqual(self._count_nans(load), 0)
 
   def test_unbounded_channel_has_nan(self):
     # variable with negative range -> x % 4 can be negative -> needs NAN fallback
     x = Variable("x", -10, 10)
-    load = UOp(Ops.LOAD, dtypes.float, (UOp(Ops.PARAM, dtypes.imagef((10, 10, 4)), arg=0).index(x, ptr=True), UOp.const(dtypes.float, 0)))
+    load = UOp(Ops.LOAD, dtypes.float, (UOp(Ops.DEFINE_GLOBAL, dtypes.imagef((10, 10, 4)), arg=0).index(x, ptr=True), UOp.const(dtypes.float, 0)))
     self.assertEqual(self._count_nans(load), 1)
 
 class TestDropTrueGate(unittest.TestCase):
@@ -475,7 +475,7 @@ class TestDropTrueGate(unittest.TestCase):
     # test that INDEX with a constant True gate gets simplified to drop the gate
     from tinygrad.codegen.late.devectorizer import load_store_indexing
     from tinygrad.uop.ops import graph_rewrite
-    buf = UOp(Ops.PARAM, dtypes.int.ptr(), arg=0)
+    buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), arg=0)
     idx = UOp.const(dtypes.index, 0)
     true_gate = UOp.const(dtypes.bool, True)
     index_with_gate = UOp(Ops.INDEX, dtypes.int.ptr(), (buf, idx, true_gate))

test/unit/test_transcendental_helpers.py

@@ -10,7 +10,7 @@ class TestTranscendentalFunctions(unittest.TestCase):
   def test_payne_hanek_reduction(self):
     # TODO: Test constant input when constant folding is fixed (or maybe test both variants)
     # Load input value from a buffer to prevent constant folding
-    input_buf = UOp(Ops.PARAM, dtypes.double.ptr(), arg=1, src=())
+    input_buf = UOp(Ops.DEFINE_GLOBAL, dtypes.double.ptr(), arg=1, src=())
     loaded_value = input_buf.index(UOp.const(dtypes.int, 0))
     def eval_payne_hanek_reduction(v:float) -> tuple[float, int]:
       return tuple(eval_uop(u, [(dtypes.float64, [v])]) for u in payne_hanek_reduction(loaded_value))

test/unit/test_uop_graph.py

@@ -253,7 +253,7 @@ class TestUOpGraph(unittest.TestCase):
 
   @unittest.skip("this test isn't valid uops")
   def test_noop_vectorize_fold(self):
-    d0 = UOp(Ops.PARAM, dtypes.float.ptr(), arg=0)
+    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=0)
     idx = UOp.const(dtypes.int, 0)
     ld = UOp(Ops.LOAD, dtypes.float.vec(2), (d0, idx))
     vec = UOp(Ops.VECTORIZE, dtypes.float.vec(2), (ld,))
@@ -265,9 +265,9 @@ class TestUOpGraph(unittest.TestCase):
 
   @unittest.skip("this test isn't valid uops")
   def test_gep_vec_fold(self):
-    d0 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 0)
-    d1 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 1)
-    d2 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 2)
+    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
+    d1 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 1)
+    d2 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 2)
     idx = UOp.const(dtypes.int, 0)
     def _test_vec(geps, count=4):
       vec = UOp(Ops.VECTORIZE, dtypes.float.vec(count), geps)
@@ -373,8 +373,8 @@ class TestUOpGraph(unittest.TestCase):
       self.assertEqual(uops[-2], wmma)  # -2 to skip SINK
 
   def test_cast_alu_fold(self):
-    d0 = UOp(Ops.PARAM, dtypes.bool.ptr(), arg=0)
-    d1 = UOp(Ops.PARAM, dtypes.int.ptr(), arg=1)
+    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.bool.ptr(), arg=0)
+    d1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), arg=1)
     idx = UOp.const(dtypes.int, 0)
     ld = d1.index(idx)
     alu = (ld<1).cast(dtypes.bool)
@@ -383,8 +383,8 @@ class TestUOpGraph(unittest.TestCase):
     self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 0)
 
   def test_double_cast_fold(self):
-    d0 = UOp(Ops.PARAM, dtypes.float.ptr(), arg=0)
-    d1 = UOp(Ops.PARAM, dtypes.int.ptr(), arg=1)
+    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=0)
+    d1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), arg=1)
     idx = UOp.const(dtypes.int, 0)
     ld = d1.index(idx)
     alu = ld.cast(dtypes.float).cast(dtypes.float)
@@ -407,7 +407,7 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_bitcast_to_same_dtype_fold(self):
     for dt in dtypes.ints + dtypes.floats + (dtypes.bool,):
-      d0 = UOp(Ops.PARAM, dt.ptr(), arg=0)
+      d0 = UOp(Ops.DEFINE_GLOBAL, dt.ptr(), arg=0)
       v = d0.index(UOp.const(dtypes.int, 0))
       uops = to_uops_list([v.bitcast(dt)])
       self.assertEqual(len([x for x in uops if x.op is Ops.BITCAST]), 0, f"dtype = {dt}")
@@ -420,7 +420,7 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_where_on_gated_load_fold(self):
     ridx0 = UOp.range(100, 0)
-    d0 = UOp(Ops.PARAM, dtypes.long.ptr(), (), 0)
+    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 0)
     ld = d0.index(ridx0.valid(ridx0<50))
     w = (ridx0<50).where(ld, 5)
     uops = to_uops_list([w])
@@ -430,7 +430,7 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_where_on_gated_load_folds_swapped_branches(self):
     ridx0 = UOp.range(100, 0)
-    d0 = UOp(Ops.PARAM, dtypes.long.ptr(), (), 0)
+    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 0)
     ld = d0.index(ridx0.valid((ridx0<50).logical_not()))
     w = (ridx0<50).where(5, ld)
     uops = to_uops_list([w])
@@ -440,7 +440,7 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_where_on_gated_load_with_cast(self):
     ridx0 = UOp.range(100, 0)
-    d0 = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
+    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
     gate_idx = ridx0.valid((ridx0<50))
     ld = d0.index(gate_idx).cast(dtypes.float)
     w = (ridx0<50).where(ld, 5.0)
@@ -451,7 +451,7 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_where_in_store_becomes_gate(self):
     ridx0 = UOp.range(100, 0)
-    d0 = UOp(Ops.PARAM, dtypes.long.ptr(), (), 0)
+    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 0)
     idx = d0.index(ridx0)
     ld = idx.load()
     val = (ridx0<50).where(5, ld)
@@ -464,14 +464,14 @@ class TestUOpGraph(unittest.TestCase):
   def test_load_idx_becomes_int(self):
     # mnist indexing with split reduceop
     # Make sure we are not doign math on the loaded index, which would promote it to long
-    c0 = UOp(Ops.PARAM, dtypes.uchar.ptr(128000), arg=0, src=())
+    c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.uchar.ptr(128000), arg=0, src=())
     c1 = UOp.range(UOp.const(dtypes.index, 512), 1, AxisType.LOOP)
     c2 = UOp.range(UOp.const(dtypes.index, 250), 2, AxisType.LOOP)
-    c3 = UOp(Ops.PARAM, dtypes.int.ptr(512), arg=1, src=())
+    c3 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(512), arg=1, src=())
     c4 = c3.index(c1)
     c5 = UOp.range(UOp.const(dtypes.index, 240), 0, AxisType.REDUCE)
     c6 = ((c2*UOp.const(dtypes.index, 240))+c5)
-    c7 = UOp(Ops.PARAM, dtypes.uchar.ptr(60000), arg=2, src=())
+    c7 = UOp(Ops.DEFINE_GLOBAL, dtypes.uchar.ptr(60000), arg=2, src=())
     c8 = c7.index(c6)
     c9 = ((c4<0).where((c4+60000), c4)!=c6.cast(dtypes.int)).where(0, c8.cast(dtypes.uint).cast(dtypes.uchar)).reduce(c5, arg=Ops.ADD)
     c10 = c0.index(((c1*UOp.const(dtypes.index, 250))+c2)).store(c9).end(c1, c2)
@@ -481,14 +481,14 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_load_idx_no_math_on_loaded(self):
     # test the (x+y)<c pattern where x has loads - we shouldn't do math on loaded indices
-    c0 = UOp(Ops.PARAM, dtypes.uchar.ptr(128000), arg=0, src=())
+    c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.uchar.ptr(128000), arg=0, src=())
     c1 = UOp.range(UOp.const(dtypes.index, 512), 1, AxisType.LOOP)
     c2 = UOp.range(UOp.const(dtypes.index, 250), 2, AxisType.LOOP)
-    c3 = UOp(Ops.PARAM, dtypes.int.ptr(512), arg=1, src=())
+    c3 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(512), arg=1, src=())
     c4 = c3.index(c1)  # c4 is a load
     c5 = UOp.range(UOp.const(dtypes.index, 240), 0, AxisType.REDUCE)
     c6 = ((c2*UOp.const(dtypes.index, 240))+c5)
-    c7 = UOp(Ops.PARAM, dtypes.uchar.ptr(60000), arg=2, src=())
+    c7 = UOp(Ops.DEFINE_GLOBAL, dtypes.uchar.ptr(60000), arg=2, src=())
     c8 = c7.index(c6)
     # (loaded + range) < const pattern - loaded value shouldn't be promoted to long
     loaded_idx = c4.cast(dtypes.index)
@@ -500,9 +500,9 @@ class TestUOpGraph(unittest.TestCase):
       self.assertNotEqual(u.dtype, dtypes.long)
 
   def test_fold_gated_load(self):
-    glbl0 = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
-    glbl1 = UOp(Ops.PARAM, dtypes.int.ptr(), (), 1)
-    glbl2 = UOp(Ops.PARAM, dtypes.int.ptr(), (), 2)
+    glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
+    glbl1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 1)
+    glbl2 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 2)
     idx = UOp.const(dtypes.int, 0)
     ld0 = glbl1.index(UOp.invalid())
     ld1 = glbl2.index(idx.valid(UOp.const(dtypes.bool, True)))
@@ -512,7 +512,7 @@ class TestUOpGraph(unittest.TestCase):
     self.assertEqual(ld0, UOp.load(glbl2.index(idx, ptr=True), dtype=dtypes.int))
 
   def test_fold_gated_load_local(self):
-    glbl0 = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
+    glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
     smem = UOp(Ops.DEFINE_LOCAL, dtypes.int.ptr(size=18, addrspace=AddrSpace.LOCAL), (), "temp")
     lidx = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 16),), "lidx0")
     st = UOp(Ops.STORE, dtypes.void, (smem.index(lidx, ptr=True), glbl0.index(lidx, ptr=True).load()))
@@ -526,7 +526,7 @@ class TestUOpGraph(unittest.TestCase):
     self.assertEqual(ld0.src[0], smem.after(barrier).index(lidx+2, ptr=True))
 
   def test_fold_gated_store(self):
-    glbl = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
+    glbl = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
     idx0 = UOp.const(dtypes.int, 0)
     idx1 = UOp.const(dtypes.int, 0)
     val = UOp.const(dtypes.int, 42)
@@ -539,7 +539,7 @@ class TestUOpGraph(unittest.TestCase):
 
   @unittest.skip("this is a uop type error")
   def test_asserts_bad_gate(self):
-    glbl0 = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
+    glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
     idx = UOp.const(dtypes.int, 0)
     bad_gate = UOp.const(dtypes.int, 1)
     with self.assertRaises(AssertionError): to_uops_list([UOp(Ops.STORE, dtypes.void, (glbl0, idx, UOp.const(dtypes.int, 42), bad_gate))])
@@ -727,7 +727,7 @@ class TestLoadStoreFolding(unittest.TestCase):
   def test_gated_load_gep_preserves_alt(self):
     """Test that LOAD(GEP, alt) preserves alt value after rewrite"""
     from tinygrad.codegen.late.devectorizer import load_store_folding
-    buf = UOp(Ops.PARAM, dtypes.float.vec(4).ptr(), (), 0)
+    buf = UOp(Ops.DEFINE_GLOBAL, dtypes.float.vec(4).ptr(), (), 0)
     idx = UOp.const(dtypes.int, 0)
     gate = UOp.const(dtypes.bool, True)
     gated_index = buf.index(idx, gate)
@@ -745,8 +745,8 @@ class TestLoadStoreFolding(unittest.TestCase):
   def test_gated_load_ptrcat_preserves_alt(self):
     """Test that LOAD(PTRCAT, alt) preserves alt value after rewrite"""
     from tinygrad.codegen.late.devectorizer import load_store_folding
-    buf1 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 0)
-    buf2 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 1)
+    buf1 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
+    buf2 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 1)
     idx = UOp.const(dtypes.int, 0)
     idx1 = buf1.index(idx)
     idx2 = buf2.index(idx)

test/unit/test_uop_symbolic.py

@@ -746,7 +746,7 @@ class TestSymbolic(unittest.TestCase):
     expr = cond.where(a, b).cast(dtypes.half)
 
     # TODO: copied from render, render does not support cast
-    glbl = UOp(Ops.PARAM, dtypes.int.ptr(), arg=0)
+    glbl = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), arg=0)
     uops = get_uops(UOp(Ops.STORE, dtypes.void, (glbl.index(UOp.const(dtypes.int, 0)), expr)).sink())
     rewritten_uop = [uop for uop in uops if uop.op is Ops.STORE][0].src[1]
 
@@ -1028,7 +1028,7 @@ class TestStoreLoadFolding(unittest.TestCase):
   """Tests for store(index, load(index)) -> NOOP rule. This rule matches patterns that EMERGE during simplification."""
   def test_store_load_folding(self):
     # store(idx, load(idx)) -> NOOP, including emergent patterns like store(idx, load(idx) + 0)
-    buf = UOp(Ops.PARAM, dtypes.int.ptr(), arg=0)
+    buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), arg=0)
     index = buf.index(UOp.const(dtypes.index, 0))
     # Direct: store(idx, load(idx)) -> NOOP
     self.assertEqual(graph_rewrite(index.store(index.load()), sym).op, Ops.NOOP)
@@ -1080,7 +1080,7 @@ class TestRangeSplitting(unittest.TestCase):
     from tinygrad.codegen.simplify import pm_split_ranges, pm_flatten_range
     r0 = UOp.range(uconst(8), 0)
     # create a simple expression using the range with mod: store range%2 to a buffer
-    buf = UOp(Ops.PARAM, dtypes.int.ptr(), arg=0)
+    buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), arg=0)
     val = (r0 % uconst(2)).cast(dtypes.int)
     store = UOp(Ops.STORE, dtypes.void, (buf.index(uconst(0)), val))
     sink = UOp(Ops.SINK, dtypes.void, (UOp(Ops.END, dtypes.void, (store, r0)),))

test/unit/test_uop_vmin_vmax.py

@@ -82,7 +82,7 @@ class TestVminVmaxProperties(unittest.TestCase):
   def test_vmin_vmax_multiplication_0_inf(self):
     # vmin and vmax for multiplication with a variable
     x = UOp.const(dtypes.float, 0.0)
-    y = UOp.load(UOp(Ops.PARAM, dtypes.float.ptr(), (), 0), UOp.const(dtypes.int, 0), dtype=dtypes.float)
+    y = UOp.load(UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0), UOp.const(dtypes.int, 0), dtype=dtypes.float)
     uop = x * y
     # TODO: these should be 0, but definitely should not be nan
     self.assertEqual(uop.vmin, -math.inf)
@@ -279,7 +279,7 @@ class TestVminVmaxVConst(unittest.TestCase):
 
   def test_vmin_vmax_vector_with_gep(self):
     # vmin and vmax for a vector constant of bool values
-    d1 = UOp(Ops.PARAM, dtypes.int.ptr(), (), 1)
+    d1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 1)
     idx = UOp.const(dtypes.int, 0)
     val = UOp(Ops.LOAD, dtypes.int.vec(2), (d1.index(idx),))
     uop = (val // 32).gep(0)

test/unit/test_validate_oob.py

@@ -11,7 +11,7 @@ class TestValidateOOB(unittest.TestCase):
   # basic index patterns
   def test_const_index(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
+      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
       to_uops_list([buf.index(UOp.const(dtypes.int, 0), ptr=True).load(dtype=dtypes.int)])  # valid
       to_uops_list([buf.index(UOp.const(dtypes.int, 15), ptr=True).load(dtype=dtypes.int)])  # valid (last element)
       with self.assertRaises(RuntimeError):
@@ -21,7 +21,7 @@ class TestValidateOOB(unittest.TestCase):
 
   def test_variable_index(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
+      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
       to_uops_list([buf.index(Variable("i", 0, 15), ptr=True).load(dtype=dtypes.int)])  # valid
       with self.assertRaises(RuntimeError):
         to_uops_list([buf.index(Variable("i", 0, 20), ptr=True).load(dtype=dtypes.int)])  # oob
@@ -30,7 +30,7 @@ class TestValidateOOB(unittest.TestCase):
 
   def test_range_with_mask(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
+      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
       r = UOp.range(42, 0, AxisType.GLOBAL)
       to_uops_list([buf.index(r.valid(r < 16), ptr=True).load(dtype=dtypes.int)])  # valid
       with self.assertRaises(RuntimeError):
@@ -38,7 +38,7 @@ class TestValidateOOB(unittest.TestCase):
 
   def test_variable_with_mask(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
+      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
       v = Variable("v", -5, 80)
       to_uops_list([buf.index(v.valid((v >= 0) & (v < 16)), ptr=True).load(dtype=dtypes.int)])  # valid
       with self.assertRaises(RuntimeError):
@@ -46,7 +46,7 @@ class TestValidateOOB(unittest.TestCase):
 
   def test_gated_store(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
+      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
       v = Variable("v", 0, 20)
       to_uops_list([buf.index(v.valid(v < 16)).store(0)])  # valid
       with self.assertRaises(RuntimeError):
@@ -55,14 +55,14 @@ class TestValidateOOB(unittest.TestCase):
   # ALU ops in index
   def test_idiv(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
+      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
       to_uops_list([buf.index(UOp.range(32, 0, AxisType.GLOBAL) // 2, ptr=True).load(dtype=dtypes.int)])  # 0..15 valid
       with self.assertRaises(RuntimeError):
         to_uops_list([buf.index(UOp.range(34, 0, AxisType.GLOBAL) // 2, ptr=True).load(dtype=dtypes.int)])  # 0..16 oob
 
   def test_mod(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
+      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
       r = UOp.range(100, 0, AxisType.GLOBAL)
       to_uops_list([buf.index(r % 16, ptr=True).load(dtype=dtypes.int)])  # 0..15 valid
       with self.assertRaises(RuntimeError):
@@ -70,14 +70,14 @@ class TestValidateOOB(unittest.TestCase):
 
   def test_shr(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
+      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
       to_uops_list([buf.index(UOp.range(64, 0, AxisType.GLOBAL) >> 2, ptr=True).load(dtype=dtypes.int)])  # 0..15 valid
       with self.assertRaises(RuntimeError):
         to_uops_list([buf.index(UOp.range(128, 0, AxisType.GLOBAL) >> 2, ptr=True).load(dtype=dtypes.int)])  # 0..31 oob
 
   def test_shl(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.PARAM, dtypes.int.ptr(64), (), 0)
+      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(64), (), 0)
       r = UOp.range(8, 0, AxisType.GLOBAL)
       to_uops_list([buf.index(r << 2, ptr=True).load(dtype=dtypes.int)])  # 0..28 valid
       with self.assertRaises(RuntimeError):
@@ -85,7 +85,7 @@ class TestValidateOOB(unittest.TestCase):
 
   def test_and(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
+      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
       r = UOp.range(100, 0, AxisType.GLOBAL)
       to_uops_list([buf.index(r & 15, ptr=True).load(dtype=dtypes.int)])  # 0..15 valid
       with self.assertRaises(RuntimeError):
@@ -93,14 +93,14 @@ class TestValidateOOB(unittest.TestCase):
 
   def test_max(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
+      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
       to_uops_list([buf.index(Variable("v", -10, 15).maximum(0), ptr=True).load(dtype=dtypes.int)])  # 0..15 valid
       with self.assertRaises(RuntimeError):
         to_uops_list([buf.index(Variable("v2", -10, 20).maximum(0), ptr=True).load(dtype=dtypes.int)])  # 0..20 oob
 
   def test_xor_in_mask(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
+      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
       r = UOp.range(32, 0, AxisType.GLOBAL)
       to_uops_list([buf.index(r.valid((r < 8) ^ ((r >= 8) & (r < 16))), ptr=True).load(dtype=dtypes.int)])  # 0..15 valid
       with self.assertRaises(RuntimeError):
@@ -109,22 +109,22 @@ class TestValidateOOB(unittest.TestCase):
   # cast patterns
   def test_float_cast_in_index(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
+      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
       r = UOp.range(20, 0)
       i = (r.cast(dtypes.float) * 0.68).trunc().cast(dtypes.int)
       to_uops_list([buf.index(i.valid((i >= 0) & (i < 16)), ptr=True).load(dtype=dtypes.int)])
 
   def test_bool_cast_in_mask(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.PARAM, dtypes.int.ptr(1), (), 0)
+      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1), (), 0)
       r = UOp.range(20, 0)
       to_uops_list([buf.index(r.valid(r.cast(dtypes.bool).logical_not()), ptr=True).load(dtype=dtypes.int)])  # only r=0 valid
 
   # load result as index/mask
   def test_load_as_index(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf0 = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
-      buf1 = UOp(Ops.PARAM, dtypes.int.ptr(64), (), 1)
+      buf0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
+      buf1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(64), (), 1)
       r = UOp.range(42, 0, AxisType.GLOBAL)
       ld0 = buf0.index(r.valid(r < 8), ptr=True).load(dtype=dtypes.int).cast(dtypes.index)
       to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 32)), ptr=True).load(dtype=dtypes.int)])  # valid
@@ -133,8 +133,8 @@ class TestValidateOOB(unittest.TestCase):
 
   def test_load_bool_as_mask(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf_bool = UOp(Ops.PARAM, dtypes.bool.ptr(16), (), 0)
-      buf_int = UOp(Ops.PARAM, dtypes.int.ptr(8), (), 1)
+      buf_bool = UOp(Ops.DEFINE_GLOBAL, dtypes.bool.ptr(16), (), 0)
+      buf_int = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(8), (), 1)
       gidx = UOp(Ops.SPECIAL, dtypes.index, (UOp.const(dtypes.index, 16),), "gidx0")
       ld_bool = buf_bool.index(gidx, ptr=True).load()
       with self.assertRaises(RuntimeError):
@@ -145,7 +145,7 @@ class TestValidateOOB(unittest.TestCase):
   def test_in_bounds_access_gated_local(self):
     with Context(CHECK_OOB=1):
       # Define buffers
-      gbuf = UOp(Ops.PARAM, dtypes.uint.ptr(400), (), 0)
+      gbuf = UOp(Ops.DEFINE_GLOBAL, dtypes.uint.ptr(400), (), 0)
       sbuf = UOp(Ops.DEFINE_LOCAL, dtypes.uint.ptr(8, addrspace=AddrSpace.LOCAL), (), "temp0")
 
       # Define indices, valids and barrier
@@ -169,8 +169,8 @@ class TestValidateOOB(unittest.TestCase):
   @unittest.skip("Bool load is not supported yet")
   def test_load_mask(self):
     with Context(CHECK_OOB=1):
-      glbl0 = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
-      mask = UOp(Ops.PARAM, dtypes.bool.ptr(16), (), 0)
+      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
+      mask = UOp(Ops.DEFINE_GLOBAL, dtypes.bool.ptr(16), (), 0)
       ridx = UOp.range(20, 0)
       ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(ridx, ridx<16&mask), ptr=True)))
       to_uops_list([ld0])

tinygrad/codegen/__init__.py

@@ -16,7 +16,7 @@ from tinygrad.codegen.late.expander import expander, pm_pre_expander, pm_group_f
 from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \
   ReduceContext, correct_load_store, pm_render, pm_add_loads
 from tinygrad.codegen.opt.postrange import apply_opts, make_images
-from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse
+from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse, pm_split_store
 from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen, pm_mops
 from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize
 
@@ -50,6 +50,9 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -
     # optimize (schedule) the AST
     sink = graph_rewrite(sink, pm_simplify_ranges, name="simplify ranges")
 
+    # split store range (only on CPU for now)
+    sink = graph_rewrite(sink, pm_split_store, ctx=ren.device, name="cut store ranges")
+
     # create image buffers
     sink = make_images(sink, ren)
 

tinygrad/codegen/late/devectorizer.py

@@ -1,7 +1,7 @@
 from typing import Any, cast
 import functools, operator, itertools
 from collections import defaultdict
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from tinygrad.dtype import dtypes, ImageDType, DType, AddrSpace, Invalid, PtrDType
 from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, GroupOp, identity_element
 from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate
@@ -299,8 +299,6 @@ pm_render = PatternMatcher([
 @dataclass
 class ReduceContext:
   acc_num: int = 0
-  # track ENDs by range for merging parallel reduces
-  range_to_ends: dict[tuple[UOp, ...], list[UOp]] = field(default_factory=dict)
 
 def horizontal_reduce(inp:UOp, out_dtype:DType) -> list[UOp]:
   # if this has a horizontal reduction component, do that first
@@ -326,19 +324,11 @@ def reduce_to_acc(ctx:ReduceContext, red:UOp):
     ctx.acc_num += 1
   ret = functools.reduce(lambda x,y: x.alu(red.arg, y), lst)
   if len(reduce_range) == 0: return ret
-  end = acc.index(UOp.const(dtypes.int, 0)).store(ret).end(*reduce_range)
-  ctx.range_to_ends.setdefault(reduce_range, []).append(end)
-  return acc.after(end).index(UOp.const(dtypes.int, 0))
-
-def merge_reduce_ends(ctx:ReduceContext, sink:UOp):
-  # merge ENDs that share the same range
-  subs = {e: UOp.group(*(e.src[0] for e in ends)).end(*r) for r, ends in ctx.range_to_ends.items() if len(ends) > 1 for e in ends}
-  return sink.substitute(subs) if subs else None
+  return acc.after(acc.index(UOp.const(dtypes.int, 0)).store(ret).end(*reduce_range)).index(UOp.const(dtypes.int, 0))
 
 pm_reduce = PatternMatcher([
-  # REDUCE -> DEFINE_ACC+ASSIGN, then merge ENDs with same range
+  # REDUCE -> DEFINE_ACC+ASSIGN
   (UPat(Ops.REDUCE, name="red"), reduce_to_acc),
-  (UPat(Ops.SINK, name="sink"), merge_reduce_ends),
   # tensor core built in accumulate
   (UPat(Ops.WMMA, name="wmma") + UPat.var("add"),
     lambda add, wmma: UOp(wmma.op, wmma.dtype, (wmma.src[0], wmma.src[1], wmma.src[2]+add), wmma.arg)),

tinygrad/codegen/late/linearizer.py

@@ -26,7 +26,7 @@ def linearize(sink:UOp) -> list[UOp]:
     extra = None
     match u.op:
       # the order and placement of these defines is important
-      case Ops.PARAM: priority, extra = -20, u.arg
+      case Ops.DEFINE_GLOBAL: priority, extra = -20, u.arg
       case Ops.DEFINE_VAR: priority, extra = -19, u.arg
       case Ops.DEFINE_LOCAL: priority = -18
       case Ops.DEFINE_REG: priority = -17

tinygrad/codegen/opt/postrange.py

@@ -331,7 +331,7 @@ class Scheduler:
   def group_for_reduces(self) -> int: return len(self.axes_of(AxisType.GROUP_REDUCE))
 
 def bufs_from_ast(ast:UOp, dname:str) -> list[Buffer]:
-  glbls = sorted([x for x in ast.backward_slice if x.op is Ops.PARAM], key=lambda x: x.arg)
+  glbls = sorted([x for x in ast.backward_slice if x.op is Ops.DEFINE_GLOBAL], key=lambda x: x.arg)
   return [Buffer(dname, x.ptrdtype.size, x.dtype.base if not isinstance(x.dtype, ImageDType) else x.dtype) for x in glbls]
 
 def apply_opts(ast:UOp, ren:Renderer) -> UOp:
@@ -362,7 +362,7 @@ def make_images(ast:UOp, ren:Renderer) -> UOp:
         ctx[dg.arg] = dt
         return dg.replace(dtype=dtypes.imagef((1, dt.size // 4, 4), dt.nbytes()))
 
-    ast = graph_rewrite(ast, PatternMatcher([(UPat(Ops.PARAM, name="dg"), make_image)]), ctx=dg_types, name="create image buffers")
+    ast = graph_rewrite(ast, PatternMatcher([(UPat(Ops.DEFINE_GLOBAL, name="dg"), make_image)]), ctx=dg_types, name="create image buffers")
 
     # undo unfoldable stores
     def undo_image_store(ctx, st, idx, dg):
@@ -370,6 +370,6 @@ def make_images(ast:UOp, ren:Renderer) -> UOp:
         return st.replace(src=(idx.replace(src=(dg.replace(dtype=ctx[dg.arg]),)+idx.src[1:]),)+st.src[1:])
 
     ast = graph_rewrite(ast, PatternMatcher([
-      (UPat(Ops.PARAM, name="dg").index(UPat(), name="idx").store(UPat(), name="st"), undo_image_store)
+      (UPat(Ops.DEFINE_GLOBAL, name="dg").index(UPat(), name="idx").store(UPat(), name="st"), undo_image_store)
     ]), ctx=dg_types, name="remove unfoldable image stores")
   return ast

tinygrad/codegen/simplify.py

@@ -1,7 +1,7 @@
 import itertools
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, graph_rewrite, _substitute, range_start
 from tinygrad.uop.symbolic import symbolic
-from tinygrad.helpers import partition
+from tinygrad.helpers import partition, dedup
 from tinygrad.dtype import dtypes, ImageDType
 
 def flatten_range(r:UOp) -> UOp|None:
@@ -126,7 +126,7 @@ def reduce_collapse(red:UOp, u:UOp, pm:PatternMatcher=pm_reduce_collapse) -> UOp
     replaces: dict[UOp, UOp] = {}
     for u in included:
       for s in u.src:
-        if s in included or s in replaces or s.op in {Ops.CONST, Ops.VCONST, Ops.PARAM, Ops.DEFINE_LOCAL, Ops.DEFINE_VAR}: continue
+        if s in included or s in replaces or s.op in {Ops.CONST, Ops.VCONST, Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_VAR}: continue
         replaces[s] = UOp(Ops.DEFINE_VAR, dtype=s.dtype, arg=(f'in{len(replaces)}', s.vmin, s.vmax))
     collapse_fxn = u.substitute(replaces).reduce(r, arg=Ops.ADD)
     sink = graph_rewrite(collapse_fxn, pm, name="reduce_collapse")
@@ -147,3 +147,16 @@ pm_load_collapse = PatternMatcher([
   # we want to make sure we dont do math on a loaded index since that can cause overflow, this undoes the rule in pm_reduce_load_collapse
   ((UPat.var("x", dtypes.index)+UPat.var("y"))<UPat.var("c"), lambda x,y,c: x < c-y if no_load(y) and no_load(c) and not no_load(x) else None),
 ])
+
+def cut_store_range(ctx:str, store:UOp, r:UOp) -> UOp|None:
+  # only cut ranges on CPU for now
+  if r.src[0].op is not Ops.CONST or ctx!="CPU": return None
+  if not (cuts:=[c.src[1].arg for c in store.get_consumer_map()[r] if c.op is Ops.CMPLT and r is c.src[0] and c.src[1].op is Ops.CONST]): return None
+  cuts = sorted(dedup([0] + cuts + [r.src[0].arg]))
+  ranges = [UOp.range((end-start), *(r.arg[0:-1]+(i,r.arg[-1]))) for i,(start,end) in enumerate(zip(cuts[:-1], cuts[1:]))]
+
+  return UOp.group(*[store.substitute({r: new_r+start}).end(new_r) for new_r, start in zip(ranges, cuts[:-1])])
+
+pm_split_store = pm_flatten_range+PatternMatcher([
+  (UPat(Ops.END, src=(UPat(Ops.STORE, name="store"), UPat.var("r"))), cut_store_range),
+])

tinygrad/engine/schedule.py

@@ -2,7 +2,7 @@ import time
 from typing import cast
 from collections import deque
 from tinygrad.uop.ops import UOp, Ops, buffers, UOpMetaClass, track_rewrites, PatternMatcher, UPat, graph_rewrite, graph_rewrite_map, Kernel
-from tinygrad.uop.spec import type_verify, tensor_spec, kernel_spec
+from tinygrad.uop.spec import type_verify, tensor_spec
 from tinygrad.device import Buffer, MultiBuffer
 from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, flatten, pluralize, SCACHE, Metadata
 from tinygrad.engine.realize import ExecItem
@@ -144,7 +144,7 @@ def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[dict[UOp, UOp], li
 
   if not SCACHE or (sc_ret:=schedule_cache.get(sched_cache_key, None)) is None:
     # verify Tensors match the spec (on big_sink, we only need to do this if cache misses)
-    if SPEC: type_verify(big_sink, tensor_spec+kernel_spec)
+    if SPEC: type_verify(big_sink, tensor_spec)
 
     # hack to preserve metadata
     graph_rewrite_map(big_sink, pm_pre_sched_cache, ctx=({}, {}), name="preserve metadata")

tinygrad/gradient.py

@@ -13,15 +13,6 @@ def reduce_gradient(ctx:UOp, ret:UOp, op:Ops):
     return ((mask/broadcast_to_input(count)) * broadcast_to_input(ctx),)
   if op == Ops.MUL: return (broadcast_to_input(ctx * ret) / ret.src[0],)
 
-def call_gradient(ctx:UOp, k:UOp):
-  if k.arg is not None: return (None,) + k.arg(ctx, k)
-  # auto-differentiate the function
-  fxn, args = k.src[0], k.src[1:]
-  params = sorted([x for x in fxn.toposort() if x.op == Ops.PARAM], key=lambda x: x.arg)
-  grads = compute_gradient(fxn, ctx, set(params))
-  subst = dict(zip(params, args))
-  return (None,) + tuple(grads[p].substitute(subst) if p in grads else None for p in params)
-
 # ctx is grad_output
 pm_gradient = PatternMatcher([
   (UPat(Ops.CAST, name="ret"), lambda ctx, ret: (ctx.cast(ret.src[0].dtype),)),
@@ -53,8 +44,8 @@ pm_gradient = PatternMatcher([
   # NOTE: this is only correct when the KERNEL has a single output
   (UPat(Ops.AFTER), lambda ctx: (ctx, ctx)),
   (UPat(Ops.CUSTOM_KERNEL, name="k"), lambda ctx, k: k.arg.grad_fxn(ctx, k)),
-  # gradient on CALL: use provided grad_fxn or auto-differentiate
-  (UPat(Ops.CALL, name="k"), call_gradient),
+  # gradient on CALL is a custom function
+  (UPat(Ops.CALL, name="k"), lambda ctx, k: (None,)+k.arg(ctx, k)),
   # there's no gradient for bitcast
   (UPat(Ops.BITCAST), lambda: (None,)),
 ])

tinygrad/helpers.py

@@ -185,7 +185,7 @@ CPU_COUNT = ContextVar("CPU_COUNT", max(1, len(os.sched_getaffinity(0)) if hasat
 # Compilers
 CPU_LLVM, CPU_LVP, AMD_LLVM = ContextVar("CPU_LLVM", 0), ContextVar("CPU_LVP", 0), ContextVar("AMD_LLVM", 0)
 NV_PTX, CUDA_PTX, NV_NAK, QCOM_IR3 = ContextVar("NV_PTX", 0), ContextVar("CUDA_PTX", 0), ContextVar("NV_NAK", 0), ContextVar("QCOM_IR3", 0)
-NULL_IR3, NULL_NAK, NULL_ALLOW_COPYOUT = ContextVar("NULL_IR3", 0), ContextVar("NULL_NAK", 0), ContextVar("NULL_ALLOW_COPYOUT", 0)
+NULL_IR3, NULL_NAK = ContextVar("NULL_IR3", 0), ContextVar("NULL_NAK", 0)
 AMD_CC, CPU_CC, NV_CC, CUDA_CC = ContextVar("AMD_CC", ""), ContextVar("CPU_CC", ""), ContextVar("NV_CC", ""), ContextVar("CUDA_CC", "")
 QCOM_CC = ContextVar("QCOM_CC", "")
 # VIZ implies PROFILE, but you can run PROFILE without VIZ

tinygrad/renderer/__init__.py

@@ -42,7 +42,7 @@ class Estimates:
       if u.op in {Ops.LOAD, Ops.STORE}:
         buf = u
         while len(buf.src): buf = buf.src[0]
-        if buf.op is Ops.PARAM: # assume all DEFINE_GLOBAL memory is accessed
+        if buf.op is Ops.DEFINE_GLOBAL: # assume all DEFINE_GLOBAL memory is accessed
           mem[(buf, u.op)] = buf.ptrdtype.size * buf.dtype.itemsize
       if u.op is Ops.RANGE:
         mult_stack.append(mults)
@@ -118,10 +118,10 @@ class ProgramSpec:
     local_size: list[int]|None = [1, 1, 1]
     for u in uops:
       if u.op is Ops.DEFINE_VAR: _vars.append(u)
-      if u.op is Ops.PARAM: _globals.append(u.arg)
+      if u.op is Ops.DEFINE_GLOBAL: _globals.append(u.arg)
       if u.op in (Ops.STORE, Ops.LOAD):
         if (idx:=u.src[0]).op is Ops.INDEX or (u.src[0].op is Ops.CAST and (idx:=u.src[0].src[0]).op is Ops.INDEX):
-          if (buf:=idx.src[0]).op is Ops.PARAM: (outs if u.op is Ops.STORE else ins).append(buf.arg)
+          if (buf:=idx.src[0]).op is Ops.DEFINE_GLOBAL: (outs if u.op is Ops.STORE else ins).append(buf.arg)
         # TODO: can else happen?
       if u.op is Ops.SPECIAL:
         if u.arg[0] == 'i': local_size = None

tinygrad/renderer/cstyle.py

@@ -177,15 +177,15 @@ class CStyleLanguage(Renderer):
       if u.op is Ops.SINK:
         if u.arg is not None: name = u.arg.function_name
         continue
-      if u.op in (Ops.PARAM, Ops.DEFINE_VAR):
-        r[u] = (f"data{u.arg}_{sz}" if (sz:=u.ptrdtype.size) > 0 else f"data{u.arg}") if u.op is Ops.PARAM else u.arg[0]
+      if u.op in (Ops.DEFINE_GLOBAL, Ops.DEFINE_VAR):
+        r[u] = (f"data{u.arg}_{sz}" if (sz:=u.ptrdtype.size) > 0 else f"data{u.arg}") if u.op is Ops.DEFINE_GLOBAL else u.arg[0]
         bufs[u] = (r[u], (u.dtype, False))
         continue
 
       # mark buffers that we store to writable
       if u.op is Ops.STORE:
         for up in u.src[0].toposort():
-          if up.op is Ops.PARAM: bufs[up] = (bufs[up][0], (bufs[up][1][0], True))
+          if up.op is Ops.DEFINE_GLOBAL: bufs[up] = (bufs[up][0], (bufs[up][1][0], True))
 
       # naming
       prefix = None
@@ -318,7 +318,7 @@ class OpenCLRenderer(CStyleLanguage):
     if any(uop.dtype.base == dtypes.half for uop in uops): prefix = (["#pragma OPENCL EXTENSION cl_khr_fp16 : enable"] + (prefix or []))
     return super().render_kernel(function_name, kernel, bufs, uops, prefix)
 
-  def aux(self, uops:list[UOp]): return (tuple(u.dtype for u in uops if u.op == Ops.PARAM),)
+  def aux(self, uops:list[UOp]): return (tuple(u.dtype for u in uops if u.op == Ops.DEFINE_GLOBAL),)
 
 class IntelRenderer(OpenCLRenderer):
   device, suffix, kernel_typedef = "CL", "INTEL", "__attribute__((intel_reqd_sub_group_size(8)))\n" + "__kernel void"

tinygrad/renderer/llvmir.py

@@ -168,8 +168,8 @@ class LLVMRenderer(Renderer):
       if u.op is Ops.SINK:
         if u.arg is not None: name = u.arg.function_name
         continue
-      if u.op in (Ops.PARAM, Ops.DEFINE_VAR):
-        r[u] = f"%data{u.arg}" if u.op is Ops.PARAM else f"%{u.arg[0]}"
+      if u.op in (Ops.DEFINE_GLOBAL, Ops.DEFINE_VAR):
+        r[u] = f"%data{u.arg}" if u.op is Ops.DEFINE_GLOBAL else f"%{u.arg[0]}"
         args.append((r[u], u.dtype))
       elif u.op in (Ops.DEFINE_LOCAL, Ops.DEFINE_REG):
         r[u] = f"%{'local' if u.op is Ops.DEFINE_LOCAL else 'reg'}_{str(u.arg).replace('(', '').replace(')', '').replace(',', '_').replace(' ', '')}"

tinygrad/renderer/nir.py

@@ -139,7 +139,7 @@ class NIRRenderer(Renderer):
 
   def_rewrite = PatternMatcher([
     (UPat(Ops.CONST, name="x"), lambda ctx,x: nimm(ctx.b, x.arg, x.dtype)),
-    (UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx.param(ctx.b, x, 8)),
+    (UPat(Ops.DEFINE_GLOBAL, name="x"), lambda ctx,x: ctx.param(ctx.b, x, 8)),
     (UPat(Ops.DEFINE_VAR, name="x"), lambda ctx,x: ctx.param(ctx.b, x, 4)),
     (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: nchannel(ctx.b, {'g':ngid, 'l':nlid, 'i': nid}[x.arg[0]](ctx.b), int(x.arg[-1]))),
     (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat.var("buf"),UPat.var("off")), allow_any_len=True), UPat.var("val")), allow_any_len=True, name="x"),
@@ -246,7 +246,7 @@ class LVPRenderer(NIRRenderer):
 
   def prerender(self, uops:list[UOp]):
     super().prerender(uops)
-    self.param_sz = sum([8 if u.op == Ops.PARAM else u.dtype.itemsize for u in uops if u.op in (Ops.PARAM, Ops.DEFINE_VAR)])
+    self.param_sz = sum([8 if u.op == Ops.DEFINE_GLOBAL else u.dtype.itemsize for u in uops if u.op in (Ops.DEFINE_GLOBAL, Ops.DEFINE_VAR)])
 
 # FIXME: this should be a rewrite rule
 def tovec(b, coord): return nalu(b, "vec4", nchannel(b, coord, 0), nchannel(b, coord, 1), nundef(b, dtypes.int), nundef(b, dtypes.int))
@@ -262,7 +262,6 @@ _nload_img = nir_instr(intrins=lambda dtype:{'IMAGE_DIM':mesa.GLSL_SAMPLER_DIM_2
 
 class IR3Renderer(NIRRenderer):
   device = "QCOM"
-  has_aux = True
 
   def nload_img(ctx,img,coord):
     ctx.texs.add(img)
@@ -287,13 +286,11 @@ class IR3Renderer(NIRRenderer):
     super().prerender(uops)
     self.texs:set[UOp] = set()
     self.uops, self.ibo_idx, self.img_idx = uops, 0, 0
-    self.param_sz = sum([8 if u.op == Ops.PARAM else u.dtype.itemsize for u in uops if u.op in (Ops.PARAM, Ops.DEFINE_VAR)])
+    self.param_sz = sum([8 if u.op == Ops.DEFINE_GLOBAL else u.dtype.itemsize for u in uops if u.op in (Ops.DEFINE_GLOBAL, Ops.DEFINE_VAR)])
 
   def postrender(self, uops:list[UOp]):
-    bufs, texs, imgs = [u for u in uops if u.op == Ops.PARAM], itertools.count().__next__, itertools.count().__next__
+    bufs, texs, imgs = [u for u in uops if u.op == Ops.DEFINE_GLOBAL], itertools.count().__next__, itertools.count().__next__
     for b in filter(lambda b: isinstance(b.dtype, ImageDType), bufs): nimm_set(self.r[b], texs() if b in self.texs else imgs(), dtypes.int)
 
     self.b.shader.contents.info.num_ubos = len([u for u in bufs if not isinstance(u.dtype, ImageDType)])
     self.b.shader.contents.info.num_images = texs() + imgs()
-
-  def aux(self, uops:list[UOp]): return (tuple(u.dtype for u in uops if u.op == Ops.PARAM),)

tinygrad/renderer/ptx.py

@@ -64,7 +64,7 @@ def mem_type(x:UOp) -> str:
   match x.op:
     case Ops.AFTER: return mem_type(x.src[0])
     case Ops.DEFINE_LOCAL: return 'shared'
-    case Ops.PARAM: return 'global'
+    case Ops.DEFINE_GLOBAL: return 'global'
     case _: raise RuntimeError(f"{x.op} needs to be memory")
 
 def render_wmma(ctx: "PTXRenderer", wmma: UOp):
@@ -91,7 +91,7 @@ string_rewrite = PatternMatcher([
   (UPat.cvar("x", dtypes.bool), lambda ctx, x: f"setp.ne.s16 {ctx.r[x]}, {render_val(x.arg, x.dtype)}, 0;"),
   (UPat.cvar("x"), lambda ctx, x: f"mov.b{ctx.types[x.dtype][1:]} {ctx.r[x]}, {render_val(x.arg, x.dtype)};"),
   (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: f"mov.u32 %{x.arg}, %{'ctaid' if x.arg[0] == 'g' else 'tid'}.{chr(120+int(x.arg[-1]))};"),
-  (UPat(Ops.PARAM, name="x"), lambda ctx, x: f"ld.param.{ctx.types[dtypes.ulong]} {ctx.r[x]}, [data{x.arg}+0];"),
+  (UPat(Ops.DEFINE_GLOBAL, name="x"), lambda ctx, x: f"ld.param.{ctx.types[dtypes.ulong]} {ctx.r[x]}, [data{x.arg}+0];"),
   (UPat((Ops.CMPLT, Ops.CMPNE, Ops.CMPEQ), name="x", allow_any_len=True, src=(UPat.var("src0"),)),
     lambda ctx, x, src0: ctx.code_for_op[x.op](ctx.r[x], *[ctx.r[v] for v in x.src], src0.dtype, ctx.types[src0.dtype])),
   (UPat(GroupOp.ALU, name="x"), lambda ctx, x: ctx.code_for_op[x.op](ctx.r[x], *[ctx.r[v] for v in x.src], x.dtype, ctx.types[x.dtype])),
@@ -222,7 +222,7 @@ class PTXRenderer(Renderer):
       elif u.op is Ops.LOAD:
         assert u.src[0].dtype == dtypes.int64, "load isn't int64"
         r[u] = [ssa('val', dtype=self.types[u.dtype.scalar()]) for _ in range(u.dtype.count)] if u.dtype.count > 1 else ssa('val', u)
-      elif u.op is Ops.PARAM: bufs.append((f"data{u.arg}", u.dtype))
+      elif u.op is Ops.DEFINE_GLOBAL: bufs.append((f"data{u.arg}", u.dtype))
       elif u.op is Ops.WMMA:
         # registers for packing/unpacking input and acc
         self.wmma_r = [[ssa("wmma_in", dtype="b32") for _ in range(0, len(r[u.src[0]]), 4 // u.src[0].dtype.scalar().itemsize)],
@@ -231,7 +231,7 @@ class PTXRenderer(Renderer):
         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.END: ("pred", "pred"), Ops.RANGE: ("ridx", None),
         Ops.DEFINE_VAR: ("dat", None), Ops.CONST: ("const", None), Ops.DEFINE_LOCAL: ("local", self.types[dtypes.ulong]),
-        Ops.PARAM: ("dat", self.types[dtypes.ulong]), **{op: ("alu", None) for op in GroupOp.ALU}}.get(u.op, (None, None))
+        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)
 
       if (l:=cast(str|list[str], string_rewrite.rewrite(u, ctx=self))) is None:

tinygrad/runtime/ops_amd.py

@@ -644,7 +644,6 @@ class AMDQueueDesc:
   write_ptr: MMIOInterface
   doorbell: MMIOInterface
   put_value: int = 0
-  params: tuple|None = None  # setup_ring params for recovery
 
   def signal_doorbell(self, dev, doorbell_value:int|None=None):
     try:
@@ -847,15 +846,16 @@ class PCIIface(PCIIfaceBase):
                    xcc_id=0, idx=0):
     assert cwsr_buffer is None, "no cwsr buffer for am"
 
-    rcvr_params: tuple
     if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA:
-      pv, doorbell_index = self.dev_impl.sdma.setup_ring(*(rcvr_params:=(ring.va_addr, ring.size, gart.va_addr+rptr, gart.va_addr+wptr, idx)))
+      pv, doorbell_index = 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, idx=idx)
     else:
-      pv, doorbell_index = self.dev_impl.gfx.setup_ring(*(rcvr_params:=(ring.va_addr, ring.size, gart.va_addr+rptr, gart.va_addr+wptr,
-        eop_buffer.va_addr, eop_buffer.size, is_aql:=(queue_type==kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL), is_aql)))
+      pv, doorbell_index = 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, eop_addr=eop_buffer.va_addr, eop_size=eop_buffer.size,
+        idx=int(is_aql:=(queue_type==kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL)), aql=is_aql)
 
-    return AMDQueueDesc(ring=ring.cpu_view().view(fmt='I'), doorbell=self.dev_impl.doorbell64.view(doorbell_index * 8, 8, fmt='Q'), put_value=pv,
-      read_ptr=gart.cpu_view().view(offset=rptr, size=8, fmt='Q'), write_ptr=gart.cpu_view().view(offset=wptr, size=8, fmt='Q'), params=rcvr_params)
+    return AMDQueueDesc(ring=ring.cpu_view().view(fmt='I'), doorbell=self.dev_impl.doorbell64.view(doorbell_index * 8, 8, fmt='Q'),
+      read_ptr=gart.cpu_view().view(offset=rptr, size=8, fmt='Q'), write_ptr=gart.cpu_view().view(offset=wptr, size=8, fmt='Q'), put_value=pv)
 
   def sleep(self, timeout) -> bool:
     if hasattr(self.pci_dev, 'irq_poller') and self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))):
@@ -867,12 +867,6 @@ class PCIIface(PCIIfaceBase):
     devs:list[AMDDevice] = [d for pg in HCQCompiled.peer_groups.values() for d in pg if isinstance(d, AMDDevice) and d.is_am()]
     for d in devs: d.iface.dev_impl.ih.interrupt_handler()
     faults = [f for d in devs if (f:=d.iface.dev_impl.gmc.check_fault())]
-    for d in devs:
-      if d.iface.dev_impl.recover():
-        d.compute_queue.put_value, _ = d.iface.dev_impl.gfx.setup_ring(*d.compute_queue.params)
-        d.compute_queue.read_ptr[0] = d.compute_queue.write_ptr[0] = d.compute_queue.put_value
-        d.timeline_signal.value = d.timeline_value - 1
-        d.error_state = None
     raise RuntimeError(f"Device hang detected: {'; '.join(faults)}" if faults else "Device hang detected")
 
   def device_fini(self): self.dev_impl.fini()

tinygrad/runtime/ops_null.py

@@ -3,7 +3,7 @@ from tinygrad.device import Compiled, Compiler, Allocator, CompilerSet, Compiler
 from tinygrad.engine.jit import MultiGraphRunner
 from tinygrad.renderer.cstyle import Renderer, CStyleLanguage, AMDHIPRenderer
 from tinygrad.uop.ops import Ops
-from tinygrad.helpers import cpu_profile, EMULATE, NULL_IR3, NULL_NAK, NULL_ALLOW_COPYOUT
+from tinygrad.helpers import cpu_profile, EMULATE, NULL_IR3, NULL_NAK
 from tinygrad.renderer.nir import IR3Renderer, NAKRenderer
 
 class NullRenderer(CStyleLanguage):
@@ -14,15 +14,14 @@ class NullRenderer(CStyleLanguage):
   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:
-  def __init__(self, device:str, name:str, lib:bytes, *args, **kwargs): self.device, self.name = device, name
+  def __init__(self, device:str, name:str, lib:bytes, **kwargs): self.device, self.name = device, name
   def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
     with cpu_profile(self.name, self.device): return 1e-3
 
 class NullAllocator(Allocator['NullDevice']):
   def _alloc(self, size, options): pass
   def _copyin(self, dest, src:memoryview): pass
-  def _copyout(self, dest:memoryview, src):
-    if not NULL_ALLOW_COPYOUT: raise RuntimeError("no copyout on NULL")
+  def _copyout(self, dest:memoryview, src): pass
   def _transfer(self, dest, src, sz:int, src_dev, dest_dev):
     with cpu_profile(f"{src_dev.device} -> {dest_dev.device}", self.dev.device): pass
   def _offset(self, buf, offset:int, size:int): pass

tinygrad/runtime/ops_python.py

@@ -85,7 +85,7 @@ class PythonProgram:
           i += 1
           continue
         if uop is Ops.AFTER: values[i] = src_values[0]
-        elif uop in {Ops.PARAM, Ops.DEFINE_LOCAL, Ops.DEFINE_REG}:
+        elif uop in {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG}:
           assert isinstance(dtype, PtrDType), dtype
           storage_fmt = storage_fmt_for_dtype(dtype.base.scalar())
           if storage_fmt is None: raise RuntimeError(f"{dtype=} is not supported")
@@ -94,7 +94,7 @@ class PythonProgram:
             # REGs are per thread
             values[i] = [memoryview(bytearray(dtype.size*dtype.itemsize)).cast(storage_fmt) for _ in range(warp_size)]
           else:
-            buf = memoryview(bytearray(dtype.size*dtype.itemsize)) if uop is not Ops.PARAM else pbufs.pop(0)
+            buf = memoryview(bytearray(dtype.size*dtype.itemsize)) if uop is not Ops.DEFINE_GLOBAL else pbufs.pop(0)
             values[i] = [buf.cast(storage_fmt)] * warp_size
         elif uop is Ops.DEFINE_VAR:
           values[i] = [pvals.pop(0)] * warp_size

tinygrad/runtime/support/am/amdev.py

@@ -225,15 +225,6 @@ class AMDev(PCIDevImplBase):
     self.ih.interrupt_handler()
     self.reg("regSCRATCH_REG6").write(self.is_err_state) # set finalized state.
 
-  def recover(self) -> bool:
-    if self.is_hive() or not self.is_err_state: return False # TODO: support mi300
-    if DEBUG >= 2: print(f"am {self.devfmt}: Start recovery")
-    self.ih.interrupt_handler()
-    self.gfx.reset_mec()
-    self.is_err_state = False
-    if DEBUG >= 2: print(f"am {self.devfmt}: Recovery complete")
-    return True
-
   def is_hive(self) -> bool: return self.gmc.xgmi_seg_sz > 0
 
   def paddr2mc(self, paddr:int) -> int: return self.gmc.mc_base + paddr

tinygrad/runtime/support/am/ip.py

@@ -240,7 +240,7 @@ class AM_GFX(AM_IP):
     self.adev.gmc.init_hub("GC", inst_cnt=self.xccs)
     if self.adev.partial_boot: return
 
-    self._config_mec()
+    self._config_gfx_rs64()
 
     # NOTE: Golden reg for gfx11. No values for this reg provided. The kernel just ors 0x20000000 to this reg.
     for xcc in range(self.xccs): self.adev.regTCP_CNTL.write(self.adev.regTCP_CNTL.read() | 0x20000000, inst=xcc)
@@ -276,17 +276,22 @@ class AM_GFX(AM_IP):
       self.adev.regCP_MEC_DOORBELL_RANGE_LOWER.write(0x100 * xcc, inst=xcc)
       self.adev.regCP_MEC_DOORBELL_RANGE_UPPER.write(0x100 * xcc + 0xf8, inst=xcc)
 
-    self._enable_mec()
+      # Enable MEC
+      if self.adev.ip_ver[am.GC_HWIP] < (10,0,0): self.adev.regCP_MEC_CNTL.write(0x0, inst=xcc)
+      else: self.adev.regCP_MEC_RS64_CNTL.update(mec_invalidate_icache=0, mec_pipe0_reset=0, mec_pipe0_active=1, mec_halt=0, inst=xcc)
+    # NOTE: Wait for MEC to be ready. The kernel does udelay here as well.
+    time.sleep(0.05)
 
     # Set 1 partition
     if self.xccs > 1 and not self.adev.partial_boot: self.adev.psp._spatial_partition_cmd(1)
 
-  def fini_hw(self): self._dequeue_hqds()
-
-  def reset_mec(self):
-    self._dequeue_hqds(reset=True)
-    self._config_mec()
-    self._enable_mec()
+  def fini_hw(self):
+    # NOTE: For aqls with xccs (queue=1), will continue from the saved state.
+    for q in range(2 if self.xccs == 1 else 1):
+      for xcc in range(self.xccs):
+        self._grbm_select(me=1, pipe=0, queue=q, inst=xcc)
+        if self.adev.regCP_HQD_ACTIVE.read(inst=xcc) & 1: self.adev.regCP_HQD_DEQUEUE_REQUEST.write(0x2, inst=xcc) # 1 - DRAIN_PIPE; 2 - RESET_WAVES
+        self._grbm_select(inst=xcc)
 
   def setup_ring(self, ring_addr:int, ring_size:int, rptr_addr:int, wptr_addr:int, eop_addr:int, eop_size:int, idx:int, aql:bool) -> tuple[int, int]:
     pipe, queue, doorbell = idx // 4, idx % 4, am.AMDGPU_NAVI10_DOORBELL_MEC_RING0
@@ -358,13 +363,7 @@ class AM_GFX(AM_IP):
   def _grbm_select(self, me=0, pipe=0, queue=0, vmid=0, inst=0):
     self.adev.regGRBM_GFX_CNTL.write(meid=me, pipeid=pipe, vmid=vmid, queueid=queue, inst=inst)
 
-  def _enable_mec(self):
-    for xcc in range(self.xccs):
-      if self.adev.ip_ver[am.GC_HWIP] >= (10,0,0): self.adev.regCP_MEC_RS64_CNTL.update(mec_pipe0_reset=0, mec_pipe0_active=1, mec_halt=0, inst=xcc)
-      else: self.adev.regCP_MEC_CNTL.write(0x0, inst=xcc)
-    time.sleep(0.05)  # Wait for MEC to be ready
-
-  def _config_mec(self):
+  def _config_gfx_rs64(self):
     def _config_helper(eng_name, cntl_reg, eng_reg, pipe_cnt, me=0, xcc=0):
       for pipe in range(pipe_cnt):
         self._grbm_select(me=me, pipe=pipe, inst=xcc)
@@ -382,17 +381,6 @@ class AM_GFX(AM_IP):
       if self.adev.ip_ver[am.GC_HWIP] >= (10,0,0):
         _config_helper(eng_name="MEC", cntl_reg="MEC_RS64", eng_reg="MEC_RS64", pipe_cnt=1, me=1, xcc=xcc)
 
-  def _dequeue_hqds(self, reset=False):
-    # NOTE: For aqls with xccs (queue=1), will continue from the saved state.
-    for q in range(2 if self.xccs == 1 else 1):
-      for xcc in range(self.xccs):
-        self._grbm_select(me=1, pipe=0, queue=q, inst=xcc)
-        if self.adev.regCP_HQD_ACTIVE.read(inst=xcc) & 1:
-          self.adev.regCP_HQD_DEQUEUE_REQUEST.write(0x2, inst=xcc) # 1 - DRAIN_PIPE; 2 - RESET_WAVES
-          if reset: self.adev.regSPI_COMPUTE_QUEUE_RESET.write(1, inst=xcc)
-          else: wait_cond(lambda: self.adev.regCP_HQD_ACTIVE.read(inst=xcc) & 1, value=0, msg="HQD dequeue timeout")
-    self._grbm_select()
-
 class AM_IH(AM_IP):
   def init_sw(self):
     self.ring_size = 256 << 10

tinygrad/runtime/support/compiler_mesa.py

@@ -1,4 +1,4 @@
-import base64, ctypes, pathlib, tempfile, hashlib
+import base64, ctypes, pathlib, tempfile, hashlib, sys
 from tinygrad.device import Compiler
 from tinygrad.helpers import cpu_objdump, system, data64
 from tinygrad.runtime.autogen import mesa, llvm
@@ -90,6 +90,7 @@ def disas_adreno(lib:bytes, gpu_id=630):
 
 class IR3Compiler(Compiler):
   def __init__(self, chip_id, cache_key="ir3"):
+    assert sys.version_info >= (3,14), "IR3 requires python 3.14's bitfield fixes"
     self.dev_id = mesa.struct_fd_dev_id(((chip_id >> 24) & 0xFF) * 100 + ((chip_id >> 16) & 0xFF) * 10 + ((chip_id >>  8) & 0xFF), chip_id)
     self.cc = mesa.ir3_compiler_create(None, self.dev_id, mesa.fd_dev_info(self.dev_id),
                                        mesa.struct_ir3_compiler_options(disable_cache=True)).contents

tinygrad/schedule/indexing.py

@@ -8,7 +8,7 @@ from tinygrad.uop.symbolic import symbolic, pm_simplify_valid, pm_drop_and_claus
 from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored
 
 ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.ASSIGN, Ops.COPY, Ops.BUFFER, Ops.BUFFER_VIEW,
-                     Ops.CONST, Ops.BIND, Ops.DEVICE, Ops.MSELECT, Ops.MSTACK, Ops.PARAM,
+                     Ops.CONST, Ops.BIND, Ops.DEVICE, Ops.MSELECT, Ops.MSTACK, Ops.DEFINE_GLOBAL,
                      Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.LOAD, Ops.KERNEL, Ops.ENCDEC}
 
 def realize(ctx:dict[UOp, None], tr:UOp) -> None: ctx[tr] = None

tinygrad/schedule/rangeify.py

@@ -68,15 +68,7 @@ def resolve_custom_kernel(ck:UOp) -> UOp:
   placeholders = [UOp.placeholder_like(s, slot=i) for i,s in enumerate(ck.src)]
   return UOp(Ops.KERNEL, src=ck.src, arg=Kernel(ck.arg.fxn(*placeholders)))
 
-param_to_ptr = PatternMatcher([
-  (UPat(Ops.PARAM, name="x"), lambda x:
-   None if isinstance(x.dtype, PtrDType) else x.replace(src=(), dtype=x.dtype.ptr(size=x.size)).reshape(x.shape)),
-])
-
 def resolve_call(c:UOp) -> UOp:
-  if c.src[0].op in {Ops.SINK, Ops.PROGRAM}:
-    # CALL is KERNEL...sort of
-    return UOp(Ops.KERNEL, src=c.src[1:], arg=Kernel(graph_rewrite(c.src[0], param_to_ptr)))
   params = sorted([x for x in c.src[0].toposort() if x.op == Ops.PARAM], key=lambda x: x.arg)
   args = c.src[1:]
   # TODO: this check belongs in spec, not here
@@ -150,7 +142,8 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
 # *****************
 # 3.5 cleanups
 
-ALWAYS_RUN_OPS = {Ops.CONTIGUOUS, Ops.COPY, Ops.ASSIGN, Ops.ENCDEC}
+# Ops.NOOP happens when we have a COPY to the device the Tensor is already on. We treat it like COPY here for MSTACK.
+ALWAYS_RUN_OPS = {Ops.CONTIGUOUS, Ops.COPY, Ops.ASSIGN, Ops.ENCDEC, Ops.NOOP}
 
 # you don't know in the first pass if axes are going to die, this happens if there's an EXPAND to the left
 def cleanup_dead_axes(b:UOp):
@@ -265,7 +258,7 @@ pm_const_buffer_folding = pm_mops+PatternMatcher([
   # copy on CONST is CONST
   (UPat(Ops.COPY, src=(UPat.cvar("x"), UPat()), name="copy"), lambda copy,x: copy.const_like(x.arg)),
   # hack if a noop turned to a const
-  (UPat(Ops.NOOP, src=(UPat.cvar("c"),), name="noop"), lambda c,noop: c.rtag(noop.tag)),
+  (UPat.cvar("c").f(Ops.NOOP).f(Ops.BUFFERIZE, allow_any_len=True, name="buf"), lambda c,buf: buf.replace(src=(c,)+buf.src[1:])),
   # mstack on CONST is CONST
   (UPat(Ops.MSTACK, src=(UPat.var("s"),), allow_any_len=True).f(Ops.INDEX, allow_any_len=True),
    lambda s: UOp.const(c.dtype, c.arg) if (c:=s.base).op is Ops.CONST else None),
@@ -416,7 +409,7 @@ class LocalAddBufferContext:
   opts:tuple|None = None
 
 def debuf(ctx:LocalAddBufferContext, buf:UOp):
-  ret = UOp(Ops.PARAM, buf.dtype.ptr(buf.arg), arg=ctx.dg)
+  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

tinygrad/tensor.py

@@ -248,10 +248,7 @@ class Tensor(OpMixin):
 
     This API is alpha and may change.
     """
-    contig_srcs = tuple(x.contiguous() if x.uop.op is not Ops.AFTER else x for x in ((self,)+lst))
-    params = [x.as_param(i) for i,x in enumerate(contig_srcs)]
-    kernel = UOp.call(*[x.uop for x in contig_srcs], fxn=fxn(*[x.uop for x in params]), arg=grad_fxn)
-    return [Tensor(s.uop.after(kernel), device=s.device) for s in contig_srcs]
+    return [Tensor(u, device=u.device) for u in UOp.custom_kernel(*[t.uop for t in (self,)+lst], fxn=fxn, grad_fxn=grad_fxn)]
 
   def schedule_with_vars(self, *lst:Tensor) -> tuple[list[ExecItem], dict[str, int]]:
     """

tinygrad/uop/__init__.py

@@ -28,6 +28,9 @@ class Ops(FastEnum):
   NOOP = auto(); REWRITE_ERROR = auto()
   PARAM = auto(); CALL = auto()
 
+  # TODO: remove this alias, DEFINE_GLOBAL is PARAM now
+  DEFINE_GLOBAL = PARAM
+
   # renderer
   # LINEAR is a list of UOps, SOURCE has a str arg that's human readable, BINARY has bytes arg that's compiled
   PROGRAM = auto(); LINEAR = auto(); SOURCE = auto(); BINARY = auto()
@@ -111,7 +114,7 @@ class GroupOp:
   # TODO: is BITCAST always Elementwise if it's shape changing?
   Elementwise = set.union(ALU, {Ops.CAST, Ops.BITCAST})
 
-  Defines = {Ops.PARAM, Ops.DEFINE_LOCAL, Ops.DEFINE_REG}
+  Defines = {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG}
 
   Irreducible = {Ops.CONST, Ops.DEFINE_VAR, Ops.SPECIAL, Ops.RANGE}
   Movement = {Ops.RESHAPE, Ops.EXPAND, Ops.PERMUTE, Ops.PAD, Ops.SHRINK, Ops.FLIP}

tinygrad/uop/ops.py

@@ -626,7 +626,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
   def buf_target(self) -> UOp:
     # the buffer that's being loaded from or store to
     match self.op:
-      case Ops.PARAM | Ops.DEFINE_LOCAL | Ops.DEFINE_REG: return self
+      case Ops.DEFINE_GLOBAL | Ops.DEFINE_LOCAL | Ops.DEFINE_REG: return self
       case Ops.AFTER | Ops.INDEX | Ops.STORE | Ops.LOAD: return self.src[0].buf_target()
       case Ops.VECTORIZE:
         assert all_same(self.src)
@@ -812,7 +812,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
 
   @staticmethod
   def placeholder(shape:tuple[int, ...], dtype:DType, slot:int, addrspace=AddrSpace.GLOBAL):
-    lookup = {AddrSpace.GLOBAL: Ops.PARAM, AddrSpace.LOCAL: Ops.DEFINE_LOCAL, AddrSpace.REG: Ops.DEFINE_REG}
+    lookup = {AddrSpace.GLOBAL: Ops.DEFINE_GLOBAL, AddrSpace.LOCAL: Ops.DEFINE_LOCAL, AddrSpace.REG: Ops.DEFINE_REG}
     ret = UOp(lookup[addrspace], dtype.ptr(prod(shape), addrspace), arg=slot)
     if len(shape) > 1: ret = ret.reshape(shape)
     return ret
@@ -1426,7 +1426,7 @@ def pyrender(ast:UOp) -> str:
 
   cmap = consumer_map_from_toposort(lst)
   not_rendered = {Ops.CONST, Ops.VCONST, Ops.DEVICE}
-  always_rendered = {Ops.PARAM, Ops.LOAD, Ops.SPECIAL, Ops.RANGE, Ops.CONTIGUOUS, Ops.VECTORIZE,
+  always_rendered = {Ops.DEFINE_GLOBAL, Ops.LOAD, Ops.SPECIAL, Ops.RANGE, Ops.CONTIGUOUS, Ops.VECTORIZE,
                      Ops.BUFFER, Ops.COPY, Ops.KERNEL, Ops.WHERE, Ops.END, Ops.ASSIGN}
 
   to_render: set[UOp] = {ast}

tinygrad/uop/spec.py

@@ -58,17 +58,6 @@ shared_spec = PatternMatcher([
 
   # RANGE/SPECIAL define loops, END closes them
   (UPat(Ops.END, src=(UPat(), UPat(Ops.RANGE))), lambda: True),
-
-  # codegen: standalone LINEAR/SOURCE/BINARY
-  (UPat(Ops.LINEAR, dtypes.void), lambda: True),
-  (UPat(Ops.SOURCE, dtypes.void, src=()), lambda: True),
-  (UPat(Ops.BINARY, dtypes.void, src=()), lambda: True),
-
-  # codegen: PROGRAM with progressive sources through the pipeline (SINK, DEVICE, LINEAR?, SOURCE?, BINARY?)
-  (UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.DEVICE))), lambda: True),
-  (UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.DEVICE), UPat(Ops.LINEAR))), lambda: True),
-  (UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.DEVICE), UPat(Ops.LINEAR), UPat(Ops.SOURCE))), lambda: True),
-  (UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.DEVICE), UPat(Ops.LINEAR), UPat(Ops.SOURCE), UPat(Ops.BINARY))), lambda: True),
 ])
 
 # ***** UOp spec in the Tensor graph *****
@@ -161,7 +150,7 @@ tensor_spec = PatternMatcher([
 
 shared_codegen_spec = PatternMatcher([
   # DEFINEs
-  (UPat(Ops.PARAM, name="x"), lambda x: isinstance(x.dtype, (PtrDType, ImageDType)) and x.dtype.addrspace == AddrSpace.GLOBAL),
+  (UPat(Ops.DEFINE_GLOBAL, name="x"), lambda x: isinstance(x.dtype, (PtrDType, ImageDType)) and x.dtype.addrspace == AddrSpace.GLOBAL),
   (UPat(Ops.DEFINE_LOCAL, name="x"), lambda x: isinstance(x.dtype, PtrDType) and x.dtype.addrspace == AddrSpace.LOCAL),
   (UPat(Ops.DEFINE_REG, src=(), name="x"), lambda x: isinstance(x.arg, int)),
 
@@ -277,6 +266,16 @@ full_spec = PatternMatcher([
   # in progress MSTACK may lose device
   (UPat((Ops.MSELECT, Ops.MSTACK), name="x"), lambda x: True),
 
+  # codegen: PROGRAM with progressive sources through the pipeline (SINK, DEVICE, LINEAR?, SOURCE?, BINARY?)
+  (UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.DEVICE))), lambda: True),
+  (UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.DEVICE), UPat(Ops.LINEAR))), lambda: True),
+  (UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.DEVICE), UPat(Ops.LINEAR), UPat(Ops.SOURCE))), lambda: True),
+  (UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.DEVICE), UPat(Ops.LINEAR), UPat(Ops.SOURCE), UPat(Ops.BINARY))), lambda: True),
+  # codegen: standalone LINEAR/SOURCE/BINARY
+  (UPat(Ops.LINEAR, dtypes.void), lambda: True),
+  (UPat(Ops.SOURCE, dtypes.void, src=()), lambda: True),
+  (UPat(Ops.BINARY, dtypes.void, src=()), lambda: True),
+
   # temp VECTORIZE/INDEX during rewrite have the wrong dtype
   (UPat(Ops.VECTORIZE), lambda: True),
   (UPat(Ops.INDEX), lambda: True),

tinygrad/viz/serve.py

@@ -43,7 +43,7 @@ from tinygrad.renderer import ProgramSpec
 from tinygrad.dtype import dtypes
 
 uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0", Ops.VCONST: "#e0e0e0", Ops.REDUCE: "#FF5B5B",
-               Ops.PARAM:"#cb9037", **{x:"#f2cb91" for x in {Ops.DEFINE_LOCAL, Ops.DEFINE_REG}}, Ops.REDUCE_AXIS: "#FF6B6B",
+               Ops.DEFINE_GLOBAL:"#cb9037", **{x:"#f2cb91" for x in {Ops.DEFINE_LOCAL, Ops.DEFINE_REG}}, Ops.REDUCE_AXIS: "#FF6B6B",
                Ops.RANGE: "#c8a0e0", Ops.ASSIGN: "#909090", Ops.BARRIER: "#ff8080", Ops.IF: "#c8b0c0", Ops.SPECIAL: "#c0c0ff",
                Ops.INDEX: "#cef263", Ops.WMMA: "#efefc0", Ops.MULTI: "#f6ccff", Ops.KERNEL: "#3e7f55", Ops.CUSTOM_KERNEL: "#3ebf55",
                **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80",
@@ -109,7 +109,6 @@ def uop_to_json(x:UOp) -> dict[int, dict]:
     if u.op is Ops.KERNEL:
       ast_str = f"SINK{tuple(s.op for s in u.arg.ast.src)}" if u.arg.ast.op is Ops.SINK else repr(u.arg.ast.op)
       argst = f"<Kernel {len(list(u.arg.ast.toposort()))} {ast_str} {[str(m) for m in u.arg.metadata]}>"
-    if u.op is Ops.BINARY: argst = f"<{len(u.arg)} bytes>"
     label = f"{str(u.op).split('.')[1]}{(chr(10)+word_wrap(argst.replace(':', ''))) if u.arg is not None else ''}"
     if u.dtype != dtypes.void: label += f"\n{u.dtype}"
     for idx,x in enumerate(u.src[:1] if u.op in {Ops.BUFFERIZE, Ops.INDEX} else (u.src if u.op is not Ops.END else [])):