fix st_vars

.github/workflows/benchmark.yml

@@ -692,9 +692,8 @@ jobs:
       run: AMD=1 TESTFILE=/raid/downloads/llama3-8b-sfr/model-00001-of-00004.safetensors python3 test/external/external_benchmark_disk_raw.py
     - name: Run full CIFAR training w 1 GPU
       run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
-    # TODO: enable
-    # - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-    #   run: BENCHMARK_LOG=resnet_10steps AMD=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee am_train_resnet_one_gpu.txt
+    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
+      run: BENCHMARK_LOG=resnet_10steps AMD=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee am_train_resnet_one_gpu.txt
     - name: Run 10 MLPerf Bert training steps (1 gpu)
       # TODO: remove BERT_LAYERS once scheduler is fast
       run: BENCHMARK_LOG=bert_10steps AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 FUSE_ARANGE=1 FUSE_ARANGE_UINT=0 MODEL=bert python3 examples/mlperf/model_train.py | tee am_train_bert_one_gpu.txt
@@ -750,7 +749,7 @@ jobs:
     - name: Test LLAMA-3
       run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --benchmark --temperature 0 | tee nv_llama3_beam.txt
     - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee nv_train_cifar_one_gpu.txt
+      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee nv_train_cifar_one_gpu.txt
     - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
       run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee nv_train_resnet_one_gpu.txt
     - name: Run 10 MLPerf Bert training steps (1 gpu)

.github/workflows/test.yml

@@ -329,7 +329,7 @@ jobs:
       run: |
         pip3 install --upgrade --force-reinstall ruff==0.11.0
         python3 -m ruff check .
-        python3 -m ruff check extra/onnx.py
+        python3 -m ruff check extra/onnx.py extra/onnx_parser.py
         python3 -m ruff check examples/mlperf/ --ignore E501
     - name: Lint tinygrad with pylint
       run: python -m pylint tinygrad/
@@ -337,6 +337,7 @@ jobs:
       run: |
         python -m mypy --strict-equality --lineprecision-report .
         cat lineprecision.txt
+        python -m mypy --strict-equality extra/onnx_parser.py
         python -m mypy --strict-equality extra/onnx.py
 
   unittest:
@@ -375,8 +376,8 @@ jobs:
         PYTHONPATH=. python extra/optimization/extract_dataset.py
         gzip -c /tmp/sops > extra/datasets/sops.gz
         DEBUG=1 MIN_ASTS=1 PYTHONPATH=. python extra/optimization/get_action_space.py
-    - name: Repo line count < 16000 lines
-      run: MAX_LINE_COUNT=16000 python sz.py
+    - name: Repo line count < 15500 lines
+      run: MAX_LINE_COUNT=15500 python sz.py
 
   fuzzing:
     name: Fuzzing
@@ -542,6 +543,8 @@ jobs:
         run: PYTHONPATH="." GPU=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
       - name: Test MLPerf stuff
         run: GPU=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: Run handcode_opt
+        run: PYTHONPATH=. MODEL=resnet GPU=1 DEBUG=1 BS=4 HALF=0 python3 examples/handcode_opt.py
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
@@ -624,7 +627,7 @@ jobs:
     - name: Test LLVM=1 DEVECTORIZE=0 for model
       run: PYTHONPATH="." LLVM=1 DEVECTORIZE=0 python3 test/models/test_efficientnet.py
     - name: Test CPU=1 DEVECTORIZE=0
-      run: CPU=1 DEVECTORIZE=0 FUSE_ARANGE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"
+      run: CPU=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"
 
   testwebgpu:
     name: Linux (WebGPU)
@@ -867,29 +870,29 @@ jobs:
     - name: Test ONNX Runner (WEBGPU)
       run: WEBGPU=1 PYTHONPATH=. python3 test/external/external_test_onnx_runner.py
 
-  osxremote:
-   name: MacOS (remote metal)
-   runs-on: macos-15
-   timeout-minutes: 10
-   env:
-     REMOTE: 1
-     REMOTEDEV: METAL
-   steps:
-     - name: Checkout Code
-       uses: actions/checkout@v4
-     - name: Setup Environment
-       uses: ./.github/actions/setup-tinygrad
-       with:
-         key: macos-remote
-         deps: testing_minimal
-     - name: Check Device.DEFAULT and print some source
-       run: |
-         python -c "from tinygrad import Device; assert Device.DEFAULT == 'REMOTE', Device.DEFAULT"
-         python -c "from tinygrad import Device; assert Device.default.properties.real_device == 'METAL', Device.default.properties.real_device"
-         DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
-     - name: Run REMOTE=1 Test
-       run: |
-         python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_tensor_variable.py
+  #osxremote:
+  #  name: MacOS (remote metal)
+  #  runs-on: macos-15
+  #  timeout-minutes: 10
+  #  env:
+  #    REMOTE: 1
+  #    REMOTEDEV: METAL
+  #  steps:
+  #    - name: Checkout Code
+  #      uses: actions/checkout@v4
+  #    - name: Setup Environment
+  #      uses: ./.github/actions/setup-tinygrad
+  #      with:
+  #        key: macos-remote
+  #        deps: testing_minimal
+  #    - name: Check Device.DEFAULT and print some source
+  #      run: |
+  #        python -c "from tinygrad import Device; assert Device.DEFAULT == 'REMOTE', Device.DEFAULT"
+  #        python -c "from tinygrad import Device; assert Device.default.properties.real_device == 'METAL', Device.default.properties.real_device"
+  #        DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
+  #    - name: Run REMOTE=1 Test
+  #      run: |
+  #        python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_tensor_variable.py
 
   amdremote:
     name: Linux (remote)

examples/handcode_opt.py

@@ -0,0 +1,134 @@
+from extra.models.resnet import ResNet50
+from extra.mcts_search import mcts_search
+from examples.mlperf.helpers import get_mlperf_bert_model
+from tinygrad import Tensor, Device, dtypes, nn
+from tinygrad.opt.kernel import Kernel
+from tinygrad.opt.heuristic import hand_coded_optimizations
+from tinygrad.uop.ops import Ops, sym_infer
+from tinygrad.device import Compiled
+from tinygrad.opt.search import beam_search, bufs_from_lin
+from tinygrad.helpers import DEBUG, ansilen, getenv, colored, TRACEMETA
+from extra.optimization.helpers import time_linearizer
+from tinygrad.engine.realize import get_program
+
+def get_sched_resnet():
+  mdl = ResNet50()
+  optim = (nn.optim.LARS if getenv("LARS") else nn.optim.SGD)(nn.state.get_parameters(mdl))
+  BS = getenv("BS", 64)
+
+  # run model twice to get only what changes, these are the kernels of the model
+  for _ in range(2):
+    out = mdl(Tensor.empty(BS, 3, 224, 224))
+    targets = [out]
+    if getenv("BACKWARD"):
+      optim.zero_grad()
+      out.sparse_categorical_crossentropy(Tensor.empty(BS, dtype=dtypes.int)).backward()
+      targets += [x for x in optim.schedule_step()]
+    sched = Tensor.schedule(*targets)
+    print(f"schedule length {len(sched)}")
+  return sched
+
+def get_sched_bert():
+  mdl = get_mlperf_bert_model()
+  optim = nn.optim.LAMB(nn.state.get_parameters(mdl))
+
+  # fake data
+  BS = getenv("BS", 9)
+  input_ids = Tensor.empty((BS, 512), dtype=dtypes.float32)
+  segment_ids = Tensor.empty((BS, 512), dtype=dtypes.float32)
+  attention_mask = Tensor.empty((BS, 512), dtype=dtypes.default_float)
+  masked_positions = Tensor.empty((BS, 76), dtype=dtypes.float32)
+  masked_lm_ids = Tensor.empty((BS, 76), dtype=dtypes.float32)
+  masked_lm_weights = Tensor.empty((BS, 76), dtype=dtypes.float32)
+  next_sentence_labels = Tensor.empty((BS, 1), dtype=dtypes.float32)
+
+  # run model twice to get only what changes, these are the kernels of the model
+  for _ in range(2):
+    lm_logits, seq_relationship_logits = mdl(input_ids, attention_mask, masked_positions, segment_ids)
+    targets = [lm_logits, seq_relationship_logits]
+    if getenv("BACKWARD"):
+      optim.zero_grad()
+      loss = mdl.loss(lm_logits, seq_relationship_logits, masked_lm_ids, masked_lm_weights, next_sentence_labels)
+      # ignore grad norm and loss scaler for now
+      loss.backward()
+      targets += [x for x in optim.schedule_step()]
+    sched = Tensor.schedule(*targets)
+    print(f"schedule length {len(sched)}")
+  return sched
+
+if __name__ == "__main__":
+  if getenv("HALF", 1):
+    dtypes.default_float = dtypes.half
+
+  # the device we are optimizing for
+  device: Compiled = Device[Device.DEFAULT]
+  if getenv("BACKWARD"): Tensor.training = True
+  print(f"optimizing for {Device.DEFAULT}")
+
+  sched = globals()[f"get_sched_{getenv('MODEL', 'resnet')}"]()
+  sched = [x for x in sched if x.ast.op is Ops.SINK]
+
+  # focus on one kernel
+  if getenv("KERNEL", -1) >= 0: sched = sched[getenv("KERNEL", -1):getenv("KERNEL", -1)+1]
+
+  # work with the schedule
+  total_tm = 0
+  running_gflops = 0
+  usage = {}
+  for i,si in enumerate(sched):
+    if DEBUG >= 3: print(si.ast)
+
+    rawbufs = bufs_from_lin(Kernel(si.ast))
+
+    # "linearize" the op into uops in different ways
+    lins: list[tuple[Kernel, str]] = []
+
+    # always try hand coded opt
+    lin = Kernel(si.ast, opts=device.renderer)
+    lin.apply_opts(hand_coded_optimizations(lin))
+    lins.append((lin, "HC"))
+
+    # maybe try tensor cores
+    lin = Kernel(si.ast, opts=device.renderer)
+    if lin.apply_tensor_cores():
+      lins.append((lin, "TC"))
+
+    # try a beam search
+    if beam:=getenv("BEAM"):
+      lin = Kernel(si.ast, opts=device.renderer)
+      lin = beam_search(lin, rawbufs, beam, bool(getenv("BEAM_ESTIMATE", 1)))
+      lins.append((lin, "BEAM"))
+
+    # try MCTS
+    if mcts:=getenv("MCTS"):
+      lin = Kernel(si.ast, opts=device.renderer)
+      lin = mcts_search(lin, rawbufs, mcts)
+      lins.append((lin, "MCTS"))
+
+    # benchmark the programs
+    choices = []
+    for lin, nm in lins:
+      tm = time_linearizer(lin, rawbufs, allow_test_size=False, cnt=10, disable_cache=True)
+      ops = (prg:=get_program(lin.get_optimized_ast(), lin.opts)).estimates.ops
+      gflops = sym_infer(ops, {k:k.min for k in lin.ast.variables()})*1e-9/tm
+      choices.append((tm, gflops, lin, prg, nm))
+
+    sorted_choices = sorted(choices, key=lambda x: x[0])
+    if DEBUG >= 1: # print all kernels
+      for tm, gflops, lin, prg, nm in choices:
+        print(f"                 kernel {i:2d} {lin.name+' '*(37-ansilen(lin.name))} {str(prg.global_size):18s} {str(prg.local_size):12s} takes {tm*1000:7.2f} ms, {gflops:6.0f} GFLOPS -- {colored(nm, 'green') if lin is sorted_choices[0][2] else nm}")
+
+    tm, gflops, lin, prg, nm = sorted_choices[0]
+    if getenv("SRC"):
+      print(si.ast)
+      print(lin.applied_opts)
+      print(get_program(lin.get_optimized_ast(), lin.opts).src)
+    total_tm += tm
+    running_gflops += gflops * tm
+    if (key := str([str(m) for m in si.metadata])) not in usage: usage[key] = (0, 0)
+    usage[key] = (usage[key][0] + tm, usage[key][1] + 1)
+    print(f"*** {total_tm*1000:7.2f} ms : kernel {i:2d} {lin.name+' '*(37-ansilen(lin.name))} {str(prg.global_size):18s} {str(prg.local_size):12s} takes {tm*1000:7.2f} ms, {gflops:6.0f} GFLOPS {[repr(m) if TRACEMETA >= 2 else str(m) for m in si.metadata]}")
+  print(f"******* total {total_tm*1000:.2f} ms, {running_gflops/total_tm:6.0f} GFLOPS")
+  print("usage:")
+  for k in sorted(usage, key=lambda x: -usage[x][0])[:10]:
+    print(f"{usage[k][0]*1000:.2f} ms: {k} ({usage[k][1]} times)")

examples/mlperf/dataloader.py

@@ -1,4 +1,6 @@
-import os, random, pickle, queue, struct, math, functools, hashlib, time
+import functools
+import hashlib
+import os, random, pickle, queue, struct, math
 from typing import List
 from pathlib import Path
 from multiprocessing import Queue, Process, shared_memory, connection, Lock, cpu_count
@@ -530,21 +532,21 @@ class BinIdxDataset:
 
     start = self.idx.tell()
     end = start + self.count * dtypes.int32.itemsize
-    self.sizes = self.idx_t[start:end].bitcast(dtypes.int32).numpy()
+    self.sizes = self.idx_t[start:end].bitcast(dtypes.int32)
 
     start = end
     end = start + self.count * dtypes.int64.itemsize
-    self.pointers = self.idx_t[start:end].bitcast(dtypes.int64).numpy()
+    self.pointers = self.idx_t[start:end].bitcast(dtypes.int64)
 
     start = end
     end = start + doc_count * dtypes.int64.itemsize
-    self.doc_idx = self.idx_t[start:end].bitcast(dtypes.int64).numpy()
+    self.doc_idx = self.idx_t[start:end].bitcast(dtypes.int64)
 
     # bin file
     self.bin_t = Tensor(base_path.with_name(f"{base_path.name}.bin"))
 
   def _index(self, idx) -> tuple[int, int]:
-    return int(self.pointers[idx]), int(self.sizes[idx])
+    return self.pointers[idx].item(), self.sizes[idx].item()
 
   def get(self, idx, offset:int=0, length:int|None=None):
     ptr, size = self._index(idx)
@@ -564,13 +566,10 @@ class GPTDataset:
     # check for cache
     cache_hash = hashlib.sha256(f"{samples}:{seqlen}:{seed}:{shuffle}".encode()).hexdigest()
     cache_path = base_path.with_name(f"{base_path.name}.{cache_hash}.index_cache")
-    print(f"try loading GPTDataset from {cache_path}...")
     if cache_path.exists():
-      print("cache found, loading...")
       with open(cache_path, "rb") as f:
         self.doc_idx, self.sample_idx, self.shuffle_idx = pickle.load(f)
     else:
-      print("cache not found, building index...")
       self.doc_idx = self._build_doc_idx()
       self.sample_idx = self._build_sample_idx()
       self.shuffle_idx = self._build_shuffle_idx()
@@ -629,20 +628,14 @@ class GPTDataset:
 
   # https://github.com/NVIDIA/Megatron-LM/blob/94bd476bd840c2fd4c3ebfc7448c2af220f4832b/megatron/core/datasets/gpt_dataset.py#L558
   def _build_doc_idx(self):
-    print(f"building doc_idx for {self.num_epochs=}, {self.indexed_dataset.count=}")
-    st = time.perf_counter()
-    # doc_idx = np.mgrid[:self.num_epochs, :self.indexed_dataset.count][1]
-    doc_idx = np.arange(self.indexed_dataset.count).reshape(1, -1).repeat(self.num_epochs, axis=0).flatten()
+    doc_idx = np.mgrid[:self.num_epochs, :self.indexed_dataset.count][1]
+    doc_idx = doc_idx.reshape(-1)
     doc_idx = doc_idx.astype(np.int32)
-    at = time.perf_counter()
     if self.shuffle: self.rng.shuffle(doc_idx)
-    print(f"doc_idx built in {at - st:.3f}s, shuffled in {time.perf_counter() - at:.3f}s")
     return doc_idx
 
   def _build_sample_idx(self):
-    print(f"building sample_idx for {self.samples=}, {self.seqlen=}, {self.doc_idx.shape[0]=}")
-    sample_idx_max = max(self.doc_idx.shape[0], self.indexed_dataset.sizes.max())
-    sample_idx = np.empty((self.samples + 1, 2), dtype=np.int64 if sample_idx_max > dtypes.int32.max else np.int32)
+    sample_idx = np.empty((self.samples + 1, 2), dtype=np.int32)
 
     sample_idx_idx, doc_idx_idx, doc_offset = 0, 0, 0
     sample_idx[sample_idx_idx, 0], sample_idx[sample_idx_idx, 1] = doc_idx_idx, doc_offset
@@ -652,7 +645,7 @@ class GPTDataset:
       remaining_seqlen = self.seqlen + 1
       while remaining_seqlen > 0:
         doc_idx = int(self.doc_idx[doc_idx_idx])
-        doc_len = int(self.indexed_dataset.sizes[doc_idx]) - doc_offset
+        doc_len = self.indexed_dataset.sizes[doc_idx].item() - doc_offset
         remaining_seqlen -= doc_len
         if remaining_seqlen <= 0:
           doc_offset += remaining_seqlen + doc_len - 1
@@ -661,7 +654,7 @@ class GPTDataset:
           if doc_idx_idx == len(self.doc_idx) - 1:
             assert sample_idx_idx == self.samples
             doc_idx = int(self.doc_idx[doc_idx_idx])
-            doc_offset = int(self.indexed_dataset.sizes[doc_idx]) - 1
+            doc_offset = self.indexed_dataset.sizes[doc_idx].item() - 1
             break
           doc_idx_idx += 1
           doc_offset = 0
@@ -672,18 +665,13 @@ class GPTDataset:
     return sample_idx
 
   def _build_shuffle_idx(self):
-    print(f"building shuffle_idx for {self.samples=}")
-    st = time.perf_counter()
     shuffle_idx = np.arange(self.samples, dtype=np.int32)
-    at = time.perf_counter()
     if self.shuffle: self.rng.shuffle(shuffle_idx)
-    print(f"shuffle_idx built in {at - st:.3f}s, shuffled in {time.perf_counter() - at:.3f}s")
     return shuffle_idx
 
 class BlendedGPTDataset:
   def __init__(self, paths:list[Path], weights:list[float], samples:int, seqlen:int, seed:int, shuffle:bool):
-    self.shuffle = shuffle
-    self.rng = np.random.RandomState(seed)
+    self.seed = seed
 
     # normalize weights
     total_weight = sum(weights)
@@ -695,47 +683,10 @@ class BlendedGPTDataset:
 
     self.datasets = [GPTDataset(path, samples_per_blend[i], seqlen, seed + i, shuffle) for i,path in enumerate(paths)]
 
-    # check for cache
-    cache_hash = hashlib.sha256(f"{samples}:{seqlen}:{seed}:{shuffle}".encode()).hexdigest()
-    cache_path = paths[0].with_name(f"{paths[0].name}.{cache_hash}.blend_cache")
-    print(f"try loading BlendedGPTDataset from {cache_path}...")
-    if cache_path.exists():
-      print("cache found, loading...")
-      with open(cache_path, "rb") as f:
-        self.dataset_idx, self.dataset_sample_idx = pickle.load(f)
-    else:
-      print("cache not found, building index...")
-      self.dataset_idx, self.dataset_sample_idx = self._build_blend_idx()
-      # save cache
-      with open(cache_path, "wb") as f:
-        pickle.dump((self.dataset_idx, self.dataset_sample_idx), f)
-
   def get(self, idx:int):
-    tokens = self.datasets[self.dataset_idx[idx]][self.dataset_sample_idx[idx]]
+    tokens = self.datasets[0][idx]
     return tokens
 
-  def _build_blend_idx(self):
-    dataset_idx = np.zeros(self.samples, dtype=np.int16)
-    dataset_sample_idx = np.zeros(self.samples, dtype=np.int64)
-
-    unspent_datasets = set(range(len(self.datasets)))
-    dataset_sample_counts = [0] * len(self.datasets)
-
-    for i in tqdm(range(self.samples)):
-      error_argmax, error_max = 0, 0.0
-      for di in unspent_datasets:
-        error = self.weights[di] * max(i, 1) - dataset_sample_counts[di]
-        if error > error_max:
-          error_max = error
-          error_argmax = di
-
-      dataset_idx[i] = error_argmax
-      dataset_sample_idx[i] = dataset_sample_counts[error_argmax]
-
-      dataset_sample_counts[error_argmax] += 1
-
-    return dataset_idx, dataset_sample_idx
-
 def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True):
   if val:
     dataset = BlendedGPTDataset([
@@ -788,8 +739,8 @@ if __name__ == "__main__":
 
   def load_llama3(val):
     bs = 24
-    samples = 5760 if val else 1_200_000 * 1152
-    seqlen = 8192
+    samples = 5760 if val else 1_200_000
+    seqlen = 512
 
     max_, min_ = 0, math.inf
     for tokens in tqdm(batch_load_llama3(bs, samples, seqlen, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=5760, val=bool(val)), total=samples//bs):

examples/mlperf/model_train.py

@@ -1296,7 +1296,7 @@ def train_llama3():
   SEED               = config["SEED"]                   = getenv("SEED", 5760)
   SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
   TRAIN_ON_VAL       = config["TRAIN_ON_VAL"]           = getenv("TRAIN_ON_VAL", 0)
-  SAMPLES            = config["SAMPLES"]                = getenv("SAMPLES", 5_760 if TRAIN_ON_VAL else 1_200_000 * 1152)
+  SAMPLES            = config["SAMPLES"]                = getenv("SAMPLES", 5_760 if TRAIN_ON_VAL else 1_200_000)
 
   # LR=1e-4 TRAIN_ON_VAL=1 DEFAULT_FLOAT=bfloat16 FUSE_ARANGE=1 JITBEAM=2 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B WARMUP_STEPS=36 DECAY_STEPS=360 SEQLEN=512 PYTHONPATH=. AMD=1 AMD_LLVM=0 MODEL=llama3 python3 examples/mlperf/model_train.py
   # trains to 7

extra/onnx.py

@@ -1,213 +1,12 @@
 # mypy: disable-error-code="misc, list-item, assignment, operator, index, arg-type"
 from types import SimpleNamespace
-from io import BufferedReader
 from typing import Any, Sequence, cast, Literal, Callable, get_args, NamedTuple
-import dataclasses, functools, io, math, types, warnings, pathlib, sys, enum, os, struct
-from tinygrad.nn.state import TensorIO
+import dataclasses, functools, io, math, types, warnings, pathlib, sys, enum
 from tinygrad.tensor import Tensor, _broadcast_shape, ReductionStr
 from tinygrad.helpers import getenv, DEBUG, all_same, prod, flatten, make_tuple, argsort, is_numpy_ndarray, get_single_element
 from tinygrad.dtype import DType, ConstType, dtypes, _from_np_dtype
 from tinygrad.device import is_dtype_supported, Device
-
-# Protobuf Wire Types
-WIRETYPE_VARINT = 0; WIRETYPE_FIXED64 = 1; WIRETYPE_LENGTH_DELIMITED = 2; WIRETYPE_START_GROUP = 3; WIRETYPE_END_GROUP = 4; WIRETYPE_FIXED32 = 5 # noqa: E702
-
-# TensorProto.DataType
-class TensorDataType:
-  UNDEFINED = 0; FLOAT = 1; UINT8 = 2; INT8 = 3; UINT16 = 4; INT16 = 5; INT32 = 6; INT64 = 7 # noqa: E702
-  STRING = 8; BOOL = 9; FLOAT16 = 10; DOUBLE = 11; UINT32 = 12; UINT64 = 13; COMPLEX64 = 14; COMPLEX128 = 15; BFLOAT16 = 16 # noqa: E702
-
-# AttributeProto.AttributeType
-class AttributeType:
-  UNDEFINED = 0; FLOAT = 1; INT = 2; STRING = 3; TENSOR = 4; GRAPH = 5; SPARSE_TENSOR = 11; TYPE_PROTO = 13; FLOATS = 6; INTS = 7 # noqa: E702
-  STRINGS = 8; TENSORS = 9; GRAPHS = 10; SPARSE_TENSORS = 12; TYPE_PROTOS = 14 # noqa: E702
-
-class PBType: FLOAT = 1; INT = 2; STRING = 3; FLOATS = 4; INTS = 5; STRINGS = 6; BYTES = 7; SUB = 8 # noqa: E702
-
-PB_INFOS: dict[str, dict] = {
-  "OperatorSetIdProto": {1: ("domain", PBType.STRING), 2: ("version", PBType.INT)},
-  "StringStringEntryProto": {1: ("key", PBType.STRING), 2: ("value", PBType.STRING)},
-  "TensorProto": {1: ("dims", PBType.INT, True), 2: ("data_type", PBType.INT), 4: ("float_data", PBType.FLOATS),
-    13: ("external_data", PBType.SUB, True, "StringStringEntryProto"), 14: ("data_location", PBType.INT),
-    5: ("int32_data", PBType.INTS), 7: ("int64_data", PBType.INTS), 8: ("name", PBType.STRING), 9: ("raw_data", PBType.BYTES),
-    10: ("double_data", PBType.FLOATS), 11: ("uint64_data", PBType.INTS)},
-  "TensorShapeProtoDimension": {1: ("dim_value", PBType.INT), 2: ("dim_param", PBType.STRING)},
-  "TensorShapeProto": {1: ("dim", PBType.SUB, True, "TensorShapeProtoDimension")},
-  "ModelProto": {1: ("ir_version", PBType.INT), 5: ("model_version", PBType.INT),
-    2: ("producer_name", PBType.STRING), 3: ("producer_version", PBType.STRING), 4: ("domain", PBType.STRING), 6: ("doc_string", PBType.STRING),
-    7: ("graph", PBType.SUB, False, ("GraphProto", lambda: {"node": [], "initializer": [], "input": [], "output": [], "value_info": []})),
-    8: ("opset_import",PBType.SUB, True, "OperatorSetIdProto")},
-  "GraphProto": {2: ("name", PBType.STRING), 10: ("doc_string", PBType.STRING),
-    1: ("node", PBType.SUB, True, ("NodeProto", lambda: {"input": [], "output": [], "attribute": [], "domain": None})),
-    5: ("initializer", PBType.SUB, True, ("TensorProto", lambda: {"dims": [], "float_data": None, "int32_data": None, "string_data": None,
-                                                                  "int64_data": None, "double_data": None, "uint64_data": None, "raw_data": None})),
-    11: ("input", PBType.SUB, True, "ValueInfoProto"), 12: ("output", PBType.SUB, True, "ValueInfoProto")},
-  "NodeProto": { 1: ("input", PBType.STRING, True), 2: ("output", PBType.STRING, True), 3: ("name", PBType.STRING),
-    4: ("op_type", PBType.STRING), 6: ("doc_string", PBType.STRING), 7: ("domain", PBType.STRING),
-    5: ("attribute", PBType.SUB, True, ("AttributeProto", lambda: {"floats": [], "ints": [], "strings": []}))},
-  "AttributeProto": {1: ("name", PBType.STRING), 20: ("type", PBType.INT), 3: ("i", PBType.INT), 8: ("ints", PBType.INT, True),
-    2: ("f", PBType.FLOAT), 7: ("floats", PBType.FLOAT, True), 4: ("s", PBType.BYTES), 9: ("strings", PBType.BYTES, True),
-    5:("t", PBType.SUB, False, ("TensorProto", lambda: {"dims": [], "float_data": None, "int32_data": None, "string_data": None, "int64_data": None,
-                                                        "double_data": None, "uint64_data": None, "raw_data": None}))},
-  "ValueInfoProto": {1: ("name", PBType.STRING), 2: ("type", PBType.SUB, False, "TypeProto"), 3: ("doc_string", PBType.STRING)},
-  "TypeProto": {1: ("tensor_type", PBType.SUB, False, "TypeProtoTensor"), 4: ("sequence_type", PBType.SUB, False, "TypeProtoSequence"),
-    9: ("optional_type", PBType.SUB, False, "TypeProtoOptional"), 6: ("denotation", PBType.STRING)},
-  "TypeProtoSequence": {1: ("elem_type", PBType.SUB, False, "TypeProto")},
-  "TypeProtoOptional": {1: ("elem_type", PBType.SUB, False, "TypeProto")},
-  "TypeProtoTensor": {1: ("elem_type", PBType.INT), 2: ("shape", PBType.SUB, False, ("TensorShapeProto", lambda: {"dim": []}))},
-}
-
-def onnx_load(fn: Tensor|str|pathlib.Path, load_external_data: bool=True):
-  parser = OnnxParser(fn, load_external_data)
-  onnx_model = parser.parse()
-  model = dict_to_namespace(onnx_model)
-  return model
-
-def gen_result(obj: dict, key_name, val, repeated: bool):
-  if repeated: obj.setdefault(key_name, []).append(val)
-  else: obj[key_name] = val
-
-def dict_to_namespace(d):
-  if isinstance(d, dict): return SimpleNamespace(**{k: dict_to_namespace(v) for k, v in d.items()})
-  elif isinstance(d, list): return [dict_to_namespace(i) for i in d]
-  return d
-
-class OnnxParser:
-  def __init__(self, inp: Tensor|str|pathlib.Path, load_external_data: bool=True):
-    self.file_path: pathlib.Path|None = None
-    self.load_external_data = load_external_data
-    if not isinstance(inp, Tensor):
-      self.file_path = pathlib.Path(inp)
-      self.tensor = Tensor(self.file_path)
-    else: self.tensor = inp
-    self.attr_func_dict = { PBType.BYTES: self._handle_bytes, PBType.SUB: self._handle_sub_message, PBType.FLOATS: self._handle_packed_floats,
-      PBType.INT: self._handle_int64, PBType.INTS: self._handle_packed_int64s, PBType.STRING: self._handle_string, PBType.FLOAT: self._handle_float}
-    self.registered_handles = {}
-    for pb_name in PB_INFOS:
-      res = {}
-      for fid, config in PB_INFOS[pb_name].items():
-        parser_fn, repeated = None, False
-        if len(config) == 2: name, attr = config
-        elif len(config) == 3: name, attr, repeated = config
-        elif len(config) == 4: name, attr, repeated, parser_fn = config
-        handler_fn = self.attr_func_dict[attr]
-        def _wrapper_handler(obj, reader, wt, h=handler_fn, n=name, p=parser_fn, r=repeated): return h(obj, n, reader, wt, parser_func=p, repeated=r)
-        res[fid] = _wrapper_handler
-      self.registered_handles[pb_name] = res
-
-  def parse(self):
-    reader = BufferedReader(TensorIO(self.tensor))
-    return self._parse_message(reader, "ModelProto", lambda: {"opset_import": [], "domain": None, "graph": None})
-
-  def decode_varint(self, reader: BufferedReader) -> int:
-    result = 0
-    shift = 0
-    while True:
-      data = reader.read(1)
-      if data == b"": raise EOFError("decode_varint EOF")
-      result |= (data[0] & 0x7F) << shift
-      if not (data[0] & 0x80): return result
-      shift += 7
-      if shift >= 70: raise ValueError("Varint too long")
-
-  def skip_field_value(self, reader: BufferedReader, wire_type):
-    if wire_type == WIRETYPE_VARINT: self.decode_varint(reader)
-    elif wire_type == WIRETYPE_FIXED64: reader.seek(8, os.SEEK_CUR)
-    elif wire_type == WIRETYPE_FIXED32: reader.seek(4, os.SEEK_CUR)
-    elif wire_type == WIRETYPE_LENGTH_DELIMITED: reader.seek(self.decode_varint(reader), os.SEEK_CUR)
-    else: raise ValueError(f"Unknown wire type: {wire_type}")
-
-  def _parse_message(self, reader, message_field_handlers_name, initial_obj_factory=lambda: {}):
-    message_field_handlers = self.registered_handles[message_field_handlers_name]
-    obj = initial_obj_factory()
-    while True:
-      try:
-        tag_val = self.decode_varint(reader)
-        field_number = tag_val >> 3
-        wire_type = tag_val & 0x07
-        if handler := message_field_handlers.get(field_number):
-          handler(obj, reader, wire_type)
-        else: self.skip_field_value(reader, wire_type)
-      except EOFError: break
-    if message_field_handlers_name == "TensorProto" and self.load_external_data and obj.get("data_location", 0) == 1: self._parse_external_data(obj)
-    return obj
-
-  def _handle_delimited(self, reader:BufferedReader, use_tensor=False) -> Tensor|bytes:
-    str_len = self.decode_varint(reader)
-    if not use_tensor: return reader.read(str_len)
-    raw = reader.raw
-    assert isinstance(raw, TensorIO)
-    res = raw._tensor[reader.tell():(reader.tell()+str_len)]
-    reader.seek(str_len, os.SEEK_CUR)
-    return res
-
-  def _handle_string(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
-    if wire_type != WIRETYPE_LENGTH_DELIMITED: raise ValueError(f"Expected length-delimited for string field '{key_name}'")
-    value = self._handle_delimited(reader)
-    assert isinstance(value, bytes)
-    gen_result(obj, key_name, value.decode("utf-8"), repeated)
-
-  def _handle_bytes(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
-    if wire_type != WIRETYPE_LENGTH_DELIMITED: raise ValueError(f"Expected length-delimited for bytes field '{key_name}'")
-    value = self._handle_delimited(reader, use_tensor=True)
-    gen_result(obj, key_name, value, repeated)
-
-  def _handle_int64(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
-    if wire_type != WIRETYPE_VARINT: raise ValueError(f"Expected varint for int64 field '{key_name}'")
-    val = self.decode_varint(reader)
-    gen_result(obj, key_name, val - 2**64 if val & (1 << 63) else val, repeated)
-
-  def _handle_float(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
-    if wire_type != WIRETYPE_FIXED32: raise ValueError(f"Expected fixed32 for float field '{key_name}'")
-    val, = struct.unpack("<f", reader.read(4))
-    gen_result(obj, key_name, val, repeated)
-
-  def _handle_packed_int64s(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
-    if wire_type != WIRETYPE_LENGTH_DELIMITED: raise ValueError("Packed int64s expected length_delimited")
-    total_bytes_len = self.decode_varint(reader)
-    old_pos = reader.tell()
-    values = []
-    while reader.tell() < total_bytes_len + old_pos:
-      val = self.decode_varint(reader) # need copy here because packed ints are varint
-      values.append(val - 2**64 if val & (1 << 63) else val)
-    obj[key_name] = values
-
-  def _handle_packed_floats(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
-    if wire_type != WIRETYPE_LENGTH_DELIMITED: raise ValueError("Packed floats expected length_delimited")
-    value = self._handle_delimited(reader, use_tensor=True)
-    obj[key_name] = value
-
-  def _handle_sub_message(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
-    if wire_type != WIRETYPE_LENGTH_DELIMITED: raise ValueError(f"Expected length-delimited for sub-message field '{key_name}'")
-    value = self._handle_delimited(reader, use_tensor=True)
-    assert isinstance(value, Tensor)
-    if isinstance(parser_func, str): sub_obj = self._parse_message(BufferedReader(TensorIO(value)), parser_func)
-    elif isinstance(parser_func, tuple): sub_obj = self._parse_message(BufferedReader(TensorIO(value)), parser_func[0], parser_func[1])
-    else: sub_obj = parser_func(BufferedReader(TensorIO(value)))
-    gen_result(obj, key_name, sub_obj, repeated)
-
-  def _parse_external_data(self, obj):
-    if "external_data" not in obj: raise ValueError("no external_data")
-    location = None
-    length = None
-    offset = 0
-    for kv in obj["external_data"]:
-      if kv["key"] == "location": location = kv["value"]
-      if kv["key"] == "offset": offset = int(kv["value"])
-      if kv["key"] == "length": length = int(kv["value"])
-    if location is None: raise ValueError("no location in external_data")
-    if self.file_path is None:
-      # get onnx file path from Tensor
-      if isinstance(self.tensor.device, str) and self.tensor.device.startswith("DISK:"):
-        self.file_path = pathlib.Path(self.tensor.device[5:])
-        if not (ext_path := self.file_path.parent.joinpath(location)).exists():
-          raise Exception(f"external location not exists: {ext_path}, may caused by symbolic link, try passing onnx file path to onnx_load")
-      else: raise Exception("onnx external_data need the origin file path, try passing onnx file path to onnx_load")
-    ext_path = self.file_path.parent.joinpath(location)
-    if not ext_path.exists(): raise Exception(f"external location not exists: {ext_path}")
-    ext_tensor = Tensor(ext_path)
-    obj["raw_data"] = ext_tensor[offset:offset+length] if length is not None else ext_tensor[offset:]
-    obj["data_location"] = 0
+from extra.onnx_parser import onnx_load
 
 # https://github.com/onnx/onnx/blob/rel-1.17.0/onnx/onnx.proto3#L500-L544
 data_types: dict[int, DType] = {

extra/onnx_parser.py

@@ -0,0 +1,207 @@
+# https://github.com/onnx/onnx/blob/main/onnx/onnx.proto3
+
+import os, pathlib, struct
+from io import BufferedReader
+from types import SimpleNamespace
+from tinygrad.nn.state import TensorIO
+from tinygrad.tensor import Tensor
+
+# Protobuf Wire Types
+WIRETYPE_VARINT = 0; WIRETYPE_FIXED64 = 1; WIRETYPE_LENGTH_DELIMITED = 2; WIRETYPE_START_GROUP = 3; WIRETYPE_END_GROUP = 4; WIRETYPE_FIXED32 = 5 # noqa: E702
+
+# TensorProto.DataType
+class TensorDataType:
+  UNDEFINED = 0; FLOAT = 1; UINT8 = 2; INT8 = 3; UINT16 = 4; INT16 = 5; INT32 = 6; INT64 = 7 # noqa: E702
+  STRING = 8; BOOL = 9; FLOAT16 = 10; DOUBLE = 11; UINT32 = 12; UINT64 = 13; COMPLEX64 = 14; COMPLEX128 = 15; BFLOAT16 = 16 # noqa: E702
+
+# AttributeProto.AttributeType
+class AttributeType:
+  UNDEFINED = 0; FLOAT = 1; INT = 2; STRING = 3; TENSOR = 4; GRAPH = 5; SPARSE_TENSOR = 11; TYPE_PROTO = 13; FLOATS = 6; INTS = 7 # noqa: E702
+  STRINGS = 8; TENSORS = 9; GRAPHS = 10; SPARSE_TENSORS = 12; TYPE_PROTOS = 14 # noqa: E702
+
+class PBType: FLOAT = 1; INT = 2; STRING = 3; FLOATS = 4; INTS = 5; STRINGS = 6; BYTES = 7; SUB = 8 # noqa: E702
+
+PB_INFOS: dict[str, dict] = {
+  "OperatorSetIdProto": {1: ("domain", PBType.STRING), 2: ("version", PBType.INT)},
+  "StringStringEntryProto": {1: ("key", PBType.STRING), 2: ("value", PBType.STRING)},
+  "TensorProto": {1: ("dims", PBType.INT, True), 2: ("data_type", PBType.INT), 4: ("float_data", PBType.FLOATS),
+    13: ("external_data", PBType.SUB, True, "StringStringEntryProto"), 14: ("data_location", PBType.INT),
+    5: ("int32_data", PBType.INTS), 7: ("int64_data", PBType.INTS), 8: ("name", PBType.STRING), 9: ("raw_data", PBType.BYTES),
+    10: ("double_data", PBType.FLOATS), 11: ("uint64_data", PBType.INTS)},
+  "TensorShapeProtoDimension": {1: ("dim_value", PBType.INT), 2: ("dim_param", PBType.STRING)},
+  "TensorShapeProto": {1: ("dim", PBType.SUB, True, "TensorShapeProtoDimension")},
+  "ModelProto": {1: ("ir_version", PBType.INT), 5: ("model_version", PBType.INT),
+    2: ("producer_name", PBType.STRING), 3: ("producer_version", PBType.STRING), 4: ("domain", PBType.STRING), 6: ("doc_string", PBType.STRING),
+    7: ("graph", PBType.SUB, False, ("GraphProto", lambda: {"node": [], "initializer": [], "input": [], "output": [], "value_info": []})),
+    8: ("opset_import",PBType.SUB, True, "OperatorSetIdProto")},
+  "GraphProto": {2: ("name", PBType.STRING), 10: ("doc_string", PBType.STRING),
+    1: ("node", PBType.SUB, True, ("NodeProto", lambda: {"input": [], "output": [], "attribute": [], "domain": None})),
+    5: ("initializer", PBType.SUB, True, ("TensorProto", lambda: {"dims": [], "float_data": None, "int32_data": None, "string_data": None,
+                                                                  "int64_data": None, "double_data": None, "uint64_data": None, "raw_data": None})),
+    11: ("input", PBType.SUB, True, "ValueInfoProto"), 12: ("output", PBType.SUB, True, "ValueInfoProto")},
+  "NodeProto": { 1: ("input", PBType.STRING, True), 2: ("output", PBType.STRING, True), 3: ("name", PBType.STRING),
+    4: ("op_type", PBType.STRING), 6: ("doc_string", PBType.STRING), 7: ("domain", PBType.STRING),
+    5: ("attribute", PBType.SUB, True, ("AttributeProto", lambda: {"floats": [], "ints": [], "strings": []}))},
+  "AttributeProto": {1: ("name", PBType.STRING), 20: ("type", PBType.INT), 3: ("i", PBType.INT), 8: ("ints", PBType.INT, True),
+    2: ("f", PBType.FLOAT), 7: ("floats", PBType.FLOAT, True), 4: ("s", PBType.BYTES), 9: ("strings", PBType.BYTES, True),
+    5:("t", PBType.SUB, False, ("TensorProto", lambda: {"dims": [], "float_data": None, "int32_data": None, "string_data": None, "int64_data": None,
+                                                        "double_data": None, "uint64_data": None, "raw_data": None}))},
+  "ValueInfoProto": {1: ("name", PBType.STRING), 2: ("type", PBType.SUB, False, "TypeProto"), 3: ("doc_string", PBType.STRING)},
+  "TypeProto": {1: ("tensor_type", PBType.SUB, False, "TypeProtoTensor"), 4: ("sequence_type", PBType.SUB, False, "TypeProtoSequence"),
+    9: ("optional_type", PBType.SUB, False, "TypeProtoOptional"), 6: ("denotation", PBType.STRING)},
+  "TypeProtoSequence": {1: ("elem_type", PBType.SUB, False, "TypeProto")},
+  "TypeProtoOptional": {1: ("elem_type", PBType.SUB, False, "TypeProto")},
+  "TypeProtoTensor": {1: ("elem_type", PBType.INT), 2: ("shape", PBType.SUB, False, ("TensorShapeProto", lambda: {"dim": []}))},
+}
+
+def onnx_load(fn: Tensor|str|pathlib.Path, load_external_data: bool=True):
+  parser = OnnxParser(fn, load_external_data)
+  onnx_model = parser.parse()
+  model = dict_to_namespace(onnx_model)
+  return model
+
+def gen_result(obj: dict, key_name, val, repeated: bool):
+  if repeated: obj.setdefault(key_name, []).append(val)
+  else: obj[key_name] = val
+
+def dict_to_namespace(d):
+  if isinstance(d, dict): return SimpleNamespace(**{k: dict_to_namespace(v) for k, v in d.items()})
+  elif isinstance(d, list): return [dict_to_namespace(i) for i in d]
+  return d
+
+class OnnxParser:
+  def __init__(self, inp: Tensor|str|pathlib.Path, load_external_data: bool=True):
+    self.file_path: pathlib.Path|None = None
+    self.load_external_data = load_external_data
+    if not isinstance(inp, Tensor):
+      self.file_path = pathlib.Path(inp)
+      self.tensor = Tensor(self.file_path)
+    else: self.tensor = inp
+    self.attr_func_dict = { PBType.BYTES: self._handle_bytes, PBType.SUB: self._handle_sub_message, PBType.FLOATS: self._handle_packed_floats,
+      PBType.INT: self._handle_int64, PBType.INTS: self._handle_packed_int64s, PBType.STRING: self._handle_string, PBType.FLOAT: self._handle_float}
+    self.registered_handles = {}
+    for pb_name in PB_INFOS:
+      res = {}
+      for fid, config in PB_INFOS[pb_name].items():
+        parser_fn, repeated = None, False
+        if len(config) == 2: name, attr = config
+        elif len(config) == 3: name, attr, repeated = config
+        elif len(config) == 4: name, attr, repeated, parser_fn = config
+        handler_fn = self.attr_func_dict[attr]
+        def _wrapper_handler(obj, reader, wt, h=handler_fn, n=name, p=parser_fn, r=repeated): return h(obj, n, reader, wt, parser_func=p, repeated=r)
+        res[fid] = _wrapper_handler
+      self.registered_handles[pb_name] = res
+
+  def parse(self):
+    reader = BufferedReader(TensorIO(self.tensor))
+    return self._parse_message(reader, "ModelProto", lambda: {"opset_import": [], "domain": None, "graph": None})
+
+  def decode_varint(self, reader: BufferedReader) -> int:
+    result = 0
+    shift = 0
+    while True:
+      data = reader.read(1)
+      if data == b"": raise EOFError("decode_varint EOF")
+      result |= (data[0] & 0x7F) << shift
+      if not (data[0] & 0x80): return result
+      shift += 7
+      if shift >= 70: raise ValueError("Varint too long")
+
+  def skip_field_value(self, reader: BufferedReader, wire_type):
+    if wire_type == WIRETYPE_VARINT: self.decode_varint(reader)
+    elif wire_type == WIRETYPE_FIXED64: reader.seek(8, os.SEEK_CUR)
+    elif wire_type == WIRETYPE_FIXED32: reader.seek(4, os.SEEK_CUR)
+    elif wire_type == WIRETYPE_LENGTH_DELIMITED: reader.seek(self.decode_varint(reader), os.SEEK_CUR)
+    else: raise ValueError(f"Unknown wire type: {wire_type}")
+
+  def _parse_message(self, reader, message_field_handlers_name, initial_obj_factory=lambda: {}):
+    message_field_handlers = self.registered_handles[message_field_handlers_name]
+    obj = initial_obj_factory()
+    while True:
+      try:
+        tag_val = self.decode_varint(reader)
+        field_number = tag_val >> 3
+        wire_type = tag_val & 0x07
+        if handler := message_field_handlers.get(field_number):
+          handler(obj, reader, wire_type)
+        else: self.skip_field_value(reader, wire_type)
+      except EOFError: break
+    if message_field_handlers_name == "TensorProto" and self.load_external_data and obj.get("data_location", 0) == 1: self._parse_external_data(obj)
+    return obj
+
+  def _handle_delimited(self, reader:BufferedReader, use_tensor=False) -> Tensor|bytes:
+    str_len = self.decode_varint(reader)
+    if not use_tensor: return reader.read(str_len)
+    raw = reader.raw
+    assert isinstance(raw, TensorIO)
+    res = raw._tensor[reader.tell():(reader.tell()+str_len)]
+    reader.seek(str_len, os.SEEK_CUR)
+    return res
+
+  def _handle_string(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
+    if wire_type != WIRETYPE_LENGTH_DELIMITED: raise ValueError(f"Expected length-delimited for string field '{key_name}'")
+    value = self._handle_delimited(reader)
+    assert isinstance(value, bytes)
+    gen_result(obj, key_name, value.decode("utf-8"), repeated)
+
+  def _handle_bytes(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
+    if wire_type != WIRETYPE_LENGTH_DELIMITED: raise ValueError(f"Expected length-delimited for bytes field '{key_name}'")
+    value = self._handle_delimited(reader, use_tensor=True)
+    gen_result(obj, key_name, value, repeated)
+
+  def _handle_int64(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
+    if wire_type != WIRETYPE_VARINT: raise ValueError(f"Expected varint for int64 field '{key_name}'")
+    val = self.decode_varint(reader)
+    gen_result(obj, key_name, val - 2**64 if val & (1 << 63) else val, repeated)
+
+  def _handle_float(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
+    if wire_type != WIRETYPE_FIXED32: raise ValueError(f"Expected fixed32 for float field '{key_name}'")
+    val, = struct.unpack("<f", reader.read(4))
+    gen_result(obj, key_name, val, repeated)
+
+  def _handle_packed_int64s(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
+    if wire_type != WIRETYPE_LENGTH_DELIMITED: raise ValueError("Packed int64s expected length_delimited")
+    total_bytes_len = self.decode_varint(reader)
+    old_pos = reader.tell()
+    values = []
+    while reader.tell() < total_bytes_len + old_pos:
+      val = self.decode_varint(reader) # need copy here because packed ints are varint
+      values.append(val - 2**64 if val & (1 << 63) else val)
+    obj[key_name] = values
+
+  def _handle_packed_floats(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
+    if wire_type != WIRETYPE_LENGTH_DELIMITED: raise ValueError("Packed floats expected length_delimited")
+    value = self._handle_delimited(reader, use_tensor=True)
+    obj[key_name] = value
+
+  def _handle_sub_message(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
+    if wire_type != WIRETYPE_LENGTH_DELIMITED: raise ValueError(f"Expected length-delimited for sub-message field '{key_name}'")
+    value = self._handle_delimited(reader, use_tensor=True)
+    assert isinstance(value, Tensor)
+    if isinstance(parser_func, str): sub_obj = self._parse_message(BufferedReader(TensorIO(value)), parser_func)
+    elif isinstance(parser_func, tuple): sub_obj = self._parse_message(BufferedReader(TensorIO(value)), parser_func[0], parser_func[1])
+    else: sub_obj = parser_func(BufferedReader(TensorIO(value)))
+    gen_result(obj, key_name, sub_obj, repeated)
+
+  def _parse_external_data(self, obj):
+    if "external_data" not in obj: raise ValueError("no external_data")
+    location = None
+    length = None
+    offset = 0
+    for kv in obj["external_data"]:
+      if kv["key"] == "location": location = kv["value"]
+      if kv["key"] == "offset": offset = int(kv["value"])
+      if kv["key"] == "length": length = int(kv["value"])
+    if location is None: raise ValueError("no location in external_data")
+    if self.file_path is None:
+      # get onnx file path from Tensor
+      if isinstance(self.tensor.device, str) and self.tensor.device.startswith("DISK:"):
+        self.file_path = pathlib.Path(self.tensor.device[5:])
+        if not (ext_path := self.file_path.parent.joinpath(location)).exists():
+          raise Exception(f"external location not exists: {ext_path}, may caused by symbolic link, try passing onnx file path to onnx_load")
+      else: raise Exception("onnx external_data need the origin file path, try passing onnx file path to onnx_load")
+    ext_path = self.file_path.parent.joinpath(location)
+    if not ext_path.exists(): raise Exception(f"external location not exists: {ext_path}")
+    ext_tensor = Tensor(ext_path)
+    obj["raw_data"] = ext_tensor[offset:offset+length] if length is not None else ext_tensor[offset:]
+    obj["data_location"] = 0

extra/perfetto/to_perfetto.py

@@ -1,6 +1,6 @@
 import sys, pickle, decimal, json
 from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent
-from tinygrad.helpers import tqdm, temp, ProfileEvent, ProfileRangeEvent, TracingKey
+from tinygrad.helpers import tqdm, temp, ProfileEvent, ProfileRangeEvent
 
 devices:dict[str, tuple[decimal.Decimal, decimal.Decimal, int]] = {}
 def prep_ts(device:str, ts:decimal.Decimal, is_copy): return int(decimal.Decimal(ts) + devices[device][is_copy])
@@ -11,14 +11,12 @@ def dev_ev_to_perfetto_json(ev:ProfileDeviceEvent):
           {"name": "thread_name", "ph": "M", "pid": dev_to_pid(ev.device)['pid'], "tid": 0, "args": {"name": "COMPUTE"}},
           {"name": "thread_name", "ph": "M", "pid": dev_to_pid(ev.device)['pid'], "tid": 1, "args": {"name": "COPY"}}]
 def range_ev_to_perfetto_json(ev:ProfileRangeEvent):
-  name = ev.name.display_name if isinstance(ev.name, TracingKey) else ev.name
-  return [{"name": name, "ph": "X", "ts": prep_ts(ev.device, ev.st, ev.is_copy), "dur": float(ev.en-ev.st), **dev_to_pid(ev.device, ev.is_copy)}]
+  return [{"name": ev.name, "ph": "X", "ts": prep_ts(ev.device, ev.st, ev.is_copy), "dur": float(ev.en-ev.st), **dev_to_pid(ev.device, ev.is_copy)}]
 def graph_ev_to_perfetto_json(ev:ProfileGraphEvent, reccnt):
   ret = []
   for i,e in enumerate(ev.ents):
     st, en = ev.sigs[e.st_id], ev.sigs[e.en_id]
-    name = e.name.display_name if isinstance(e.name, TracingKey) else e.name
-    ret += [{"name": name, "ph": "X", "ts": prep_ts(e.device, st, e.is_copy), "dur": float(en-st), **dev_to_pid(e.device, e.is_copy)}]
+    ret += [{"name": e.name, "ph": "X", "ts": prep_ts(e.device, st, e.is_copy), "dur": float(en-st), **dev_to_pid(e.device, e.is_copy)}]
     for dep in ev.deps[i]:
       d = ev.ents[dep]
       ret += [{"ph": "s", **dev_to_pid(d.device, d.is_copy), "id": reccnt+len(ret), "ts": prep_ts(d.device, ev.sigs[d.en_id], d.is_copy), "bp": "e"}]
@@ -26,8 +24,6 @@ def graph_ev_to_perfetto_json(ev:ProfileGraphEvent, reccnt):
   return ret
 def to_perfetto(profile:list[ProfileEvent]):
   # Start json with devices.
-  profile += [ProfileDeviceEvent("TINY")]
-
   prof_json = [x for ev in profile if isinstance(ev, ProfileDeviceEvent) for x in dev_ev_to_perfetto_json(ev)]
   for ev in tqdm(profile, desc="preparing profile"):
     if isinstance(ev, ProfileRangeEvent): prof_json += range_ev_to_perfetto_json(ev)

test/external/external_benchmark_schedule.py

@@ -1,8 +1,12 @@
+from typing import List
 from extra.models.resnet import ResNet50
-from tinygrad import Tensor, nn, Device
-from tinygrad.helpers import Profiling, Timing, getenv
+from tinygrad import Tensor, nn
+from tinygrad.helpers import Profiling, Timing, getenv, BEAM, NOOPT, DEBUG, Context, ansilen
 from tinygrad.uop.ops import Ops
+from tinygrad.opt.kernel import Kernel
+from tinygrad.opt.heuristic import hand_coded_optimizations
 from tinygrad.codegen import get_rewrites_for_renderer, apply_rewrites, rewrites_for_linearizer
+from tinygrad.opt.search import beam_search, bufs_from_lin
 from tinygrad.uop.spec import type_verify
 
 if __name__ == "__main__":
@@ -27,13 +31,26 @@ if __name__ == "__main__":
       if not SCHEDULE_ONLY:
         asts = list({x.ast.key:x.ast for x in sched if x.ast.op is Ops.SINK}.values())
         if (restrict_kernel := getenv("RESTRICT_KERNEL", -1)) != -1: asts = asts[restrict_kernel:restrict_kernel+1]
+        kernels: List[Kernel] = []
+        with Timing(f"***** model opts({len(asts):2d}) in  "):
+          with Profiling(PROFILE >= 3):
+            for ast in asts:
+              k = Kernel(ast)
+              if BEAM:
+                with Context(DEBUG=max(2, DEBUG.value)): k = beam_search(k, bufs_from_lin(k), BEAM.value)
+              elif NOOPT: pass
+              else: k.apply_opts(hand_coded_optimizations(k))
+              kernels.append(k)
+
+        with Timing("***** model prep in      "):
+          kernels = [(k, k.get_optimized_ast(), get_rewrites_for_renderer(k.opts, linearizer=False)) for k in kernels]
 
-        rewrites = get_rewrites_for_renderer(Device.default.renderer, linearizer=False)
         with Profiling(PROFILE, fn="/tmp/rewrite.prof"):
           with Timing("***** model rewrite in   "):
             rewritten_uops = []
-            for u in asts:
-              rewritten_uops.append(apply_rewrites(u, rewrites))
+            for i,(k,u,rewrites) in enumerate(kernels):
+              with Timing(f"rewrite {i:2d} {k.name}{' '*(50-ansilen(k.name))}", enabled=getenv("VERBOSE", 0)):
+                rewritten_uops.append(apply_rewrites(u, rewrites))
 
         if LINEARIZE:
           with Timing("***** model linearize in "):

test/test_arange.py

@@ -7,7 +7,6 @@ from tinygrad.opt.kernel import Opt, OptOps, Kernel, KernelOptError
 from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
 from tinygrad.opt.search import get_kernel_actions
 from tinygrad.uop.ops import Ops
-from tinygrad.codegen import apply_rewrites, rewrites_for_views
 
 class TestArange(unittest.TestCase):
   def _get_flops(self, N, opts=None):
@@ -15,7 +14,10 @@ class TestArange(unittest.TestCase):
     tt = Tensor.arange(N)
     sched = tt.schedule()
     self.assertEqual(len(sched), 1)
-    p = get_program(sched[-1].ast, opts=opts)
+    k = Kernel(sched[-1].ast)
+    if opts is not None:
+      for o in opts: k.apply_opt(o)
+    p = get_program(k.get_optimized_ast(), k.opts)
     print(p.name)
     #print(p.src)
     ExecItem(CompiledRunner(p), [tt.uop.buffer]).run()
@@ -50,11 +52,11 @@ class TestArange(unittest.TestCase):
     def test_complexity_w_local_and_padto(self): return self.test_complexity([Opt(OptOps.LOCAL, 0, 16), Opt(OptOps.PADTO, axis=1, arg=32)])
 
   def test_all_opts(self, opts=None, exclude=None):
-    k = Kernel(apply_rewrites(Tensor.arange(256).schedule()[-1].ast, rewrites_for_views))
+    k = Kernel(Tensor.arange(256).schedule()[-1].ast)
     if opts is not None:
       for o in opts: k.apply_opt(o)
     all_opts_256 = [kk.applied_opts for kk in get_kernel_actions(k, include_0=False).values()]
-    k = Kernel(apply_rewrites(Tensor.arange(2560).schedule()[-1].ast, rewrites_for_views))
+    k = Kernel(Tensor.arange(2560).schedule()[-1].ast)
     if opts is not None:
       for o in opts: k.apply_opt(o)
     all_opts_2560 = [kk.applied_opts for kk in get_kernel_actions(k, include_0=False).values()]

test/test_const_folding.py

@@ -139,9 +139,10 @@ class TestBitcastConstFolding(unittest.TestCase):
 class TestIndexingConstFolding(unittest.TestCase):
   def test_scalar_index(self):
     t = Tensor.arange(16).float().reshape(1,1,4,4).realize()
-    _check_ast_count(1, t[:,:,Tensor(1),:])
-    _check_ast_count(1, t[:,:,Tensor(1)+2,:])
-    _check_ast_count(1, t[:,:,Tensor(1),Tensor(0)])
+    # TODO: fold these
+    _check_ast_count(2, t[:,:,Tensor(1),:])
+    _check_ast_count(2, t[:,:,Tensor(1)+2,:])
+    _check_ast_count(2, t[:,:,Tensor(1),Tensor(0)])
 
   @unittest.expectedFailure
   def test_const_tensor_index(self):

test/test_hcq.py

@@ -6,7 +6,7 @@ from tinygrad.runtime.support.hcq import HCQCompiled, HCQBuffer
 from tinygrad.runtime.autogen import libc
 from tinygrad.runtime.support.system import PCIIfaceBase
 from tinygrad.engine.realize import get_runner, CompiledRunner, get_program
-from tinygrad.opt.kernel import Opt, OptOps
+from tinygrad.opt.kernel import Kernel, Opt, OptOps
 from tinygrad import Variable
 
 MOCKGPU = getenv("MOCKGPU")
@@ -163,8 +163,10 @@ class TestHCQ(unittest.TestCase):
     a = Tensor.randint((3, 3, 3), dtype=dtypes.int, device=Device.DEFAULT).realize()
     b = a + 1
     si = b.schedule()[-1]
+    k = Kernel(si.ast, opts=TestHCQ.d0.renderer)
+    for i in range(3): k.apply_opt(Opt(op=OptOps.LOCAL, axis=0, arg=3))
 
-    runner = CompiledRunner(get_program(si.ast, TestHCQ.d0.renderer, opts=[Opt(op=OptOps.LOCAL, axis=0, arg=3) for _ in range(3)]))
+    runner = CompiledRunner(get_program(k.get_optimized_ast(), k.opts))
 
     zb = Buffer(Device.DEFAULT, 3 * 3 * 3, dtypes.int, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
     zt = Buffer(Device.DEFAULT, 3 * 3 * 3, dtypes.int, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
@@ -336,7 +338,7 @@ class TestHCQ(unittest.TestCase):
     et = float(sig_en.timestamp - sig_st.timestamp)
 
     print(f"exec kernel time: {et:.2f} us")
-    assert 0.1 <= et <= (100000 if MOCKGPU or Device.DEFAULT in {"CPU", "LLVM"} else 100)
+    assert 0.1 <= et <= (15000 if MOCKGPU or Device.DEFAULT in {"CPU", "LLVM"} else 100)
 
   def test_speed_copy_bandwidth(self):
     if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")

test/test_jit.py

@@ -792,7 +792,6 @@ class TestJitGraphSplit(unittest.TestCase):
       multigraph=[self.ji_graph(5)],
       hcqgraph=[self.ji_graph(5)])
 
-  @unittest.skip("flaky")
   def test_jit_multidev_xfer(self):
     if Device.DEFAULT in {"CPU", "LLVM"}: raise unittest.SkipTest("CPU/LLVM is not a valid default device for this test (zero-copies)")
 

test/test_linearizer.py

@@ -13,9 +13,6 @@ from tinygrad.engine.realize import run_schedule, lower_schedule, CompiledRunner
 from tinygrad.opt.heuristic import hand_coded_optimizations
 from tinygrad.helpers import prod, Context, getenv, CI, flatten, dedup, AMX, AMD_LLVM
 from tinygrad.dtype import DType, dtypes, AddrSpace
-from tinygrad.codegen import apply_rewrites, rewrites_for_views
-
-def push_views(ast): return apply_rewrites(ast, rewrites_for_views)
 
 def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]:
   if isinstance(r, Tensor): r = [r]
@@ -25,7 +22,7 @@ def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]:
   # now all input buffers in s[-1] should be realized
   # create fresh buffers for the outputs
   bufs = [Buffer((x).device, x.size, x.dtype).allocate() if i < len(s[-1].ast.src) else x for i,x in enumerate(s[-1].bufs)]
-  return push_views(s[-1].ast), bufs
+  return s[-1].ast, bufs
 
 def helper_tc_allclose(N:int, M:int, K:int, dtype_in:DType, dtype_out:DType, axis:int=0, tc_select:int=-1, tc_opt:int=0, use_tensor_cores:int=1):
   a, b = Tensor.rand(M, K, dtype=dtype_in), Tensor.rand(K, N, dtype=dtype_in)
@@ -124,7 +121,7 @@ class TestLinearizer(unittest.TestCase):
     with Context(FUSE_ARANGE=1):
       sink = dataset[idxs].contiguous().kernelize().uop.base.src[1].arg.ast
     real_index = dataset.numpy()[idxs.numpy()].reshape(4, 256, 1, 1)
-    helper_linearizer_ast(push_views(sink), [dataset, idxs], wanna_output=[real_index])
+    helper_linearizer_ast(sink, [dataset, idxs], wanna_output=[real_index])
 
   def test_two_nested_range(self):
     a = Tensor.randn(2, ).realize()
@@ -217,7 +214,9 @@ class TestLinearizer(unittest.TestCase):
     # these are of size 3 to avoid float4 coalesce
     r = a[:-1] + a[1:]
 
-    uops = get_program(r.schedule()[-1].ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0)]).uops
+    k = Kernel(r.schedule()[-1].ast)
+    k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=0))
+    uops = get_program(k.get_optimized_ast(), k.opts).uops
     num_loads = len([uop for uop in uops if uop.op is Ops.LOAD])
     assert num_loads <= 4, "more load uops than needed"
     assert num_loads >= 4, "unexpected number of uops, maybe this test needs updating?"
@@ -228,7 +227,9 @@ class TestLinearizer(unittest.TestCase):
     a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize()
     r = a.expand([2]) + b.expand([2])
 
-    uops = get_program(r.schedule()[-1].ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0)]).uops
+    k = Kernel(r.schedule()[-1].ast)
+    k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=0))
+    uops = get_program(k.get_optimized_ast(), k.opts).uops
     num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
     assert num_ops <= 1, "more alu uops than needed"
 
@@ -237,7 +238,10 @@ class TestLinearizer(unittest.TestCase):
     x, w = Tensor.randn((1,1,3)).realize(), Tensor.randn((1,1,2)).realize()
     r = Tensor.conv2d(x,w,padding=1).relu()
 
-    uops = get_program(r.schedule()[-1].ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]).uops
+    k = Kernel(r.schedule()[-1].ast)
+    k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=0))
+    k.apply_opt(Opt(op=OptOps.UNROLL, axis=0, arg=0))
+    uops = get_program(k.get_optimized_ast(), k.opts).uops
     accs = [u for u in uops if u.op is Ops.DEFINE_REG]
     stores = [u for u in uops if u.op is Ops.STORE]
     assert len(accs) == 0  # it's removed now
@@ -249,7 +253,9 @@ class TestLinearizer(unittest.TestCase):
   @unittest.skipUnless(Device.DEFAULT == "CPU", "test only for CPU")
   def test_upcast_with_locals_cpu(self):
     out = Tensor.ones(64,64).contiguous() @ Tensor.ones(64,64).contiguous()
-    prg = get_program(out.schedule()[-1].ast, opts=[Opt(OptOps.LOCAL, axis=0, arg=4)]).uops
+    k = Kernel(out.schedule()[-1].ast)
+    k.apply_opt(Opt(OptOps.LOCAL, axis=0, arg=4))
+    prg = get_program(k.get_optimized_ast(), k.opts)
     self.assertEqual(len(prg.src.split("for")), 5)
 
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@@ -259,8 +265,10 @@ class TestLinearizer(unittest.TestCase):
   def test_upcast_with_locals(self):
     x, y = Tensor.rand(1,128), Tensor.rand(128, 128)
     r = (x@y).relu()
+    realized_ast = r.schedule()[-1].ast
     opts_to_apply = [Opt(op=OptOps.GROUP, axis=0, arg=8), Opt(op=OptOps.LOCAL, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=4)]
-    program = get_program(r.schedule()[-1].ast, opts=opts_to_apply)
+    realized_ast = realized_ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply)))
+    program = get_program(realized_ast, Device[Device.DEFAULT].renderer)
 
     stores = [u for u in program.uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]
 
@@ -274,7 +282,10 @@ class TestLinearizer(unittest.TestCase):
   def test_zero_fold(self):
     a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize()
     r = Tensor.stack(a, b)
-    uops = get_program(r.schedule()[-1].ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0)]).uops
+
+    k = Kernel(r.schedule()[-1].ast)
+    k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=0))
+    uops = get_program(k.get_optimized_ast(), k.opts).uops
     num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
     assert num_ops == 0, "more alu uops than needed"
 
@@ -284,14 +295,16 @@ class TestLinearizer(unittest.TestCase):
       if is_dtype_supported(tensor_dtype) and is_dtype_supported(acc_dtype):
         a = Tensor([1, 2, 3], dtype=tensor_dtype).sum()
         realized_ast = a.schedule()[-1].ast
-        program = get_program(realized_ast, opts=[])
+        realized_ast = realized_ast.replace(arg=KernelInfo(opts_to_apply=tuple()))
+        program = get_program(realized_ast, Device[Device.DEFAULT].renderer)
         local = [uop for uop in program.uops if uop.op is Ops.DEFINE_REG]
         assert local[0].dtype.base == acc_dtype
 
   def test_arg_acc_dtype(self):
     def helper_arg_acc_dtype(c: Tensor, expected_dtype:DType):
       realized_ast = c.schedule()[-1].ast
-      program = get_program(realized_ast, opts=[])
+      realized_ast = realized_ast.replace(arg=KernelInfo(opts_to_apply=tuple()))
+      program = get_program(realized_ast, Device[Device.DEFAULT].renderer)
       local = [uop for uop in program.uops if uop.op is Ops.DEFINE_REG]
       self.assertEqual(local[0].dtype.base, expected_dtype)
 
@@ -328,7 +341,7 @@ class TestLinearizer(unittest.TestCase):
       a, b = Tensor.rand(m, k, dtype=tc.dtype_in), Tensor.rand(k, n, dtype=tc.dtype_in)
       r = a.matmul(b, dtype=tc.dtype_out)
       sched = r.schedule()
-      realized_ast = push_views(sched[-1].ast)
+      realized_ast = sched[-1].ast
       kernel = Kernel(realized_ast)
       kernel.apply_tensor_cores(1, axis=0, tc_select=-1, tc_opt=2)
       prg = get_program(kernel.get_optimized_ast(), kernel.opts)
@@ -749,7 +762,11 @@ class TestFloat4(unittest.TestCase):
     c = a + b
 
     s = c.schedule()[0]
-    uops = get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=2)]).uops
+    k = Kernel(s.ast)
+    k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=4))
+    k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=2))
+    uops = get_program(k.get_optimized_ast(), k.opts).uops
+
     assert TestFloat4.count_float4(uops) == (4, 2)
 
   @unittest.skipUnless(Device.DEFAULT in {"CPU", "LLVM"} and AMX, "Only CPU with AMX upcasts float up to size 16")
@@ -760,7 +777,10 @@ class TestFloat4(unittest.TestCase):
       c = a + b
 
       s = c.schedule()[0]
-      return get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=shift)]).uops
+      k = Kernel(s.ast)
+      k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=4))
+      k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=shift))
+      return get_program(k.get_optimized_ast(), k.opts).uops
 
     sizes = [12, 8, 16]
     shifts = [3, 2, 4]
@@ -790,7 +810,10 @@ class TestFloat4(unittest.TestCase):
     c = a + b
 
     s = c.schedule()[0]
-    uops = get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=2)]).uops
+    k = Kernel(s.ast)
+    k.apply_opt(Opt(op=OptOps.UPCAST, axis=1, arg=4))
+    k.apply_opt(Opt(op=OptOps.UPCAST, axis=1, arg=2))
+    uops = get_program(k.get_optimized_ast(), k.opts).uops
 
     assert TestFloat4.count_float4(uops) == (0, 2)
 
@@ -802,7 +825,10 @@ class TestFloat4(unittest.TestCase):
       c = a + b
 
       s = c.schedule()[0]
-      return get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=shift)]).uops
+      k = Kernel(s.ast)
+      k.shift_to(1, 4, AxisType.UPCAST)  # manual trigger float4 dim
+      k.shift_to(1, shift, AxisType.UPCAST, insert_at=k.shape_len-1)
+      return get_program(k.get_optimized_ast(), k.opts).uops
 
     sizes = [13, 9, 17]
     shifts = [3, 2, 4]
@@ -820,7 +846,9 @@ class TestFloat4(unittest.TestCase):
     # float4 should be emitted (the reduce axis of size 4 is the float4 axis here)
 
     s = c.schedule()[0]
-    uops = get_program(s.ast, opts=[Opt(op=OptOps.UNROLL, axis=0, arg=4)]).uops
+    k = Kernel(s.ast)
+    k.apply_opt(Opt(op=OptOps.UNROLL, axis=0, arg=4))
+    uops = get_program(k.get_optimized_ast(), k.opts).uops
 
     assert TestFloat4.count_float4(uops) == (0, 0)
 
@@ -834,7 +862,10 @@ class TestFloat4(unittest.TestCase):
     # UPDATE: now we do this fusion
 
     s = c.schedule()[0]
-    uops = get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]).uops
+    k = Kernel(s.ast)
+    k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=0))
+    k.apply_opt(Opt(op=OptOps.UNROLL, axis=0, arg=0))
+    uops = get_program(k.get_optimized_ast(), k.opts).uops
 
     assert TestFloat4.count_float4(uops) in {(0,1), (1,1)}
 
@@ -847,7 +878,9 @@ class TestFloat4(unittest.TestCase):
     # since the top axis is not contiguous.
 
     s = c.schedule()[0]
-    uops = get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4)]).uops
+    k = Kernel(s.ast)
+    k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=4))
+    uops = get_program(k.get_optimized_ast(), k.opts).uops
 
     assert TestFloat4.count_float4(uops) == (0, 1)
 
@@ -859,7 +892,9 @@ class TestFloat4(unittest.TestCase):
     # should float4 b but not a
 
     s = c.schedule()[0]
-    uops = get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4)]).uops
+    k = Kernel(s.ast)
+    k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=4))
+    uops = get_program(k.get_optimized_ast(), k.opts).uops
 
     assert TestFloat4.count_float4(uops) == (1, 1)
 
@@ -948,7 +983,7 @@ class TestHandCodedOpts(unittest.TestCase):
     layer_2 = Tensor.cat(layer_1.unsqueeze(0), Tensor.empty(6, 20))
 
     s = layer_2.schedule()[-1]
-    k = Kernel(push_views(s.ast))
+    k = Kernel(s.ast)
     k.apply_opts(hand_coded_optimizations(k))
     assert len(k.bufs) == 6  # make sure all ops are done in one kernel
     # masked upcast should upcast masked axis of size 7
@@ -961,7 +996,7 @@ class TestHandCodedOpts(unittest.TestCase):
     monster = Tensor.stack(*[Tensor.stack(*[Tensor.empty(16) for _ in range(6)]) for _ in range(6)])
 
     s = monster.schedule()[-1]
-    k = Kernel(push_views(s.ast))
+    k = Kernel(s.ast)
     k.apply_opts(hand_coded_optimizations(k))
     assert len(k.bufs) == 37  # make sure all ops are done in one kernel
     # should upcast the two Tensor.stacks
@@ -977,7 +1012,7 @@ class TestHandCodedOpts(unittest.TestCase):
       wino_schedule = out.schedule()
       # collect upcasts of tile transform kernels
       for i, si in enumerate(wino_schedule):
-        k = Kernel(push_views(si.ast))
+        k = Kernel(si.ast)
         k.apply_opts(hand_coded_optimizations(k))
         if k.reduceop is not None: continue  # not a tile transform kernel (there is a gemm reduce kernel)
         if len(k.bufs) < 22: continue  # not a tile transform kernel (there's a permute kernel at the end)
@@ -989,7 +1024,7 @@ class TestHandCodedOpts(unittest.TestCase):
 
       backward_schedule = Tensor.schedule(x.grad, w.grad)
       for si in backward_schedule:
-        k = Kernel(push_views(si.ast))
+        k = Kernel(si.ast)
         k.apply_opts(hand_coded_optimizations(k))
         if len(k.bufs) < 20: continue  # not a tile transform kernel
         # heuristic number to make sure that at least some upcasts but not too many upcasts are being done

test/test_nn.py

@@ -401,7 +401,7 @@ class TestNN(unittest.TestCase):
       torch_z = torch_layer(torch_x)
       np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=1e-8, rtol=1e-8)
 
-  def test_embedding_one_kernel(self, ops=612000, kcount=2):
+  def test_embedding_one_kernel(self, ops=41410, kcount=3):
     GlobalCounters.reset()
     layer = Embedding(20, 30)
     layer.weight = Tensor.zeros_like(layer.weight).contiguous()
@@ -409,7 +409,7 @@ class TestNN(unittest.TestCase):
                 [12, 19, 8, 1]])
     result = layer(a)
     schedule = result.schedule()
-    self.assertEqual(len([item for item in schedule if item.ast.op is Ops.SINK]), kcount, "first run realizes weight and embedding")
+    self.assertEqual(kcount, len([item for item in schedule if item.ast.op is Ops.SINK]), "first run realizes weight and embedding")
     run_schedule(schedule)
 
     b = Tensor([[1, 2, 3],

test/test_ops.py

@@ -957,9 +957,8 @@ class TestOps(unittest.TestCase):
     helper_test_op([(45,65)], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6)
     helper_test_op([(45,65)], lambda t: torch.nn.functional.softplus(t, beta=3), lambda t: Tensor.softplus(t, beta=3), grad_atol=1e-6)
     helper_test_op([(45,65)], lambda t: torch.nn.functional.softplus(t, beta=1/3), lambda t: Tensor.softplus(t, beta=1/3), grad_atol=1e-6)
-    helper_test_op([(45,65)], lambda t: torch.nn.functional.softplus(t, beta=3, threshold=0.5),
-                              lambda t: Tensor.softplus(t, beta=3, threshold=0.5), grad_atol=1e-6)
-    helper_test_op([(45,65)], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6, low=300, high=400)
+    # # TODO: support threshold and enable this
+    # helper_test_op([(45,65)], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6, low=300, high=400)
     helper_test_op([(45,65)], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6, low=-400, high=-300)
     helper_test_op([()], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6)
 

test/test_opt_gemm.py

@@ -2,7 +2,7 @@ import numpy as np
 import unittest
 from tinygrad import Tensor
 from tinygrad.helpers import get_single_element
-from tinygrad.opt.kernel import Opt, OptOps
+from tinygrad.opt.kernel import Kernel, Opt, OptOps
 from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
 
 class TestOptGemm(unittest.TestCase):
@@ -17,7 +17,9 @@ class TestOptGemm(unittest.TestCase):
     t = self.a.T @ self.b.T
     # TODO: this should be a generic test helper
     si = get_single_element(t.schedule())
-    run = CompiledRunner(get_program(si.ast, opts=opts))
+    k = Kernel(si.ast)
+    k.apply_opts(opts)
+    run = CompiledRunner(get_program(k.get_optimized_ast(), k.opts))
     ExecItem(run, si.bufs).run()
     test = si.bufs[0].numpy().reshape(self.res.shape)
     np.testing.assert_allclose(self.res, test, atol=1e-4)

test/test_quantize_onnx.py

@@ -40,7 +40,10 @@ def create_gemm_model(model_path:str, batch_size=N, in_size=N, out_size=N, bias=
 
 def sexec(out:Tensor, opts:list[Opt], replace_src=None, run_count=3):
   si = out.schedule()[-1]
-  prg = get_program(si.ast, opts=opts)
+  k = Kernel(si.ast, opts=Device[Device.DEFAULT].renderer)
+  #opts = [Opt(op=OptOps.UPCAST, axis=0, arg=128)] #, Opt(op=OptOps.UNROLL, axis=0, arg=4)]
+  k.apply_opts(opts)
+  prg = get_program(k.get_optimized_ast(), k.opts)
   if replace_src is not None:
     old_name = prg.src.split("__attribute__((noinline)) void ")[1].split("(")[0]
     prg = replace(prg, src=replace_src + "/* DSP boilerplate */" + prg.src.split("/* DSP boilerplate */")[1].replace(old_name, "fxn"))
@@ -294,7 +297,10 @@ class TestDSPCache(unittest.TestCase):
             x41,)),)),)),))""")
     opts = [Opt(op=OptOps.UNROLL, axis=0, arg=8), Opt(op=OptOps.UPCAST, axis=1, arg=32), Opt(op=OptOps.UPCAST, axis=0, arg=4)]
     with Context(DEVECTORIZE=0, QUANTIZE=1):
-      prg = get_program(ast, opts=opts)
+      k = Kernel(ast, opts=Device[Device.DEFAULT].renderer)
+      k.apply_opts(opts)
+      prg = get_program(k.get_optimized_ast(), k.opts)
+      #print(prg.src)
 
     new_src = """
 typedef int int32 __attribute__((aligned(128),vector_size(128)));
@@ -356,7 +362,7 @@ __attribute__((noinline)) void r_196_32_4_24_8(unsigned char* restrict __attribu
     prg = replace(prg, src=new_src+prg.src.split("/* DSP boilerplate */ ")[1])
     rt = CompiledRunner(prg)
     #Device.default.compiler.disassemble(rt.lib)
-    ei = ExecItem(rt, bufs_from_lin(Kernel(ast)))
+    ei = ExecItem(rt, bufs_from_lin(k))
     tm = ei.run(wait=True)
     print(f"final time {tm*1e6:.2f} us")
 

test/test_schedule.py

@@ -6,7 +6,7 @@ import unittest
 import numpy as np
 import functools
 from typing import cast
-from hypothesis import assume, given, settings, strategies as strat
+from hypothesis import assume, given, strategies as strat
 
 from tinygrad import nn, dtypes, Device, Tensor
 from tinygrad.device import is_dtype_supported
@@ -16,7 +16,7 @@ from tinygrad.uop.ops import PatternMatcher, UOp, Ops, GroupOp, UPat, graph_rewr
 from tinygrad.uop.symbolic import symbolic_simple
 from tinygrad.helpers import CI, DEBUG, FUSE_ARANGE, SPLIT_REDUCEOP, GlobalCounters, Context, getenv, all_same, temp
 from tinygrad.schedule.kernelize import merge_views, get_kernelize_map, Kernel
-from tinygrad.engine.schedule import create_schedule_with_vars
+from tinygrad.engine.schedule import ScheduleItem, create_schedule_with_vars
 from tinygrad.engine.realize import CompiledRunner, run_schedule, lower_schedule
 
 class KernelCountException(Exception): pass
@@ -70,7 +70,7 @@ def _test_conv2d(allowed:int, dtype:DType=dtypes.float, **kwargs):
 def schedule_graph_rewrite(big_sink:UOp): return get_kernelize_map(big_sink)[big_sink]
 
 class TestSchedule(unittest.TestCase):
-  def test_arange_avgpool2d(self, kcount=1):
+  def test_arange_avgpool2d(self, kcount=2):
     x = Tensor.arange(25).reshape(1,1,5,5).cast(dtypes.float32)
     t = x.avg_pool2d(padding=1)
     sched = t.schedule()
@@ -151,7 +151,6 @@ class TestSchedule(unittest.TestCase):
       self.assertEqual(root.item(), sum(range(N)))
 
   @given(strat.sampled_from(range(2,4)), strat.sampled_from(range(2,4)), strat.sampled_from(range(0,4)), strat.sampled_from(range(0,4)))
-  @settings(deadline=None)
   def test_indexing_scalars(self, x, y, a, b):
     assume(a<x and b<y)
     X = Tensor.randn(x, y).realize()
@@ -1028,14 +1027,14 @@ class TestSchedule(unittest.TestCase):
     Tensor.manual_seed(0)
     x = Tensor.randn(4, 32).realize()
     out = x.argmin(-1)
-    run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 3))
     np.testing.assert_equal(out.numpy(), x.numpy().argmin(axis=-1))
 
   def test_argmax_multireduce_fusion(self):
     Tensor.manual_seed(0)
     x = Tensor.randn(4, 32).realize()
     out = x.argmax(-1)
-    run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 3))
     np.testing.assert_equal(out.numpy(), x.numpy().argmax(axis=-1))
 
   def test_scaled_dot_product_attention_multireduce_fusion(self):
@@ -1354,7 +1353,8 @@ class TestSchedule(unittest.TestCase):
     r = a.sum(0) + 6
     b = r.sum(0) * 4
     c = r.sum(1) * 2
-    check_schedule([b, c], 3)
+    schedule = check_schedule([b, c], 3)
+    self.assertIs(store_val(schedule[0]).op, Ops.ADD)
 
   def test_multireduce_simple_chase(self):
     Tensor.manual_seed(0)
@@ -1376,7 +1376,8 @@ class TestSchedule(unittest.TestCase):
     r = a.sum(2) + b
     d = r.T * 4
     e = r * d
-    check_schedule([d, e], 3)
+    schedule = check_schedule([d, e], 3)
+    self.assertIs(store_val(schedule[0]).op, Ops.ADD)
 
   def test_multireduce_push_permute_chase(self):
     Tensor.manual_seed(0)
@@ -1386,6 +1387,7 @@ class TestSchedule(unittest.TestCase):
     d = r.T * 4
     e = r * (d + a).sum(2)
     schedule = check_schedule([d, e], 3) # make sure it doesn't fuse
+    self.assertIs(store_val(schedule[0]).op, Ops.ADD)
     run_schedule(schedule)
     np.testing.assert_allclose(d.numpy(), (a.numpy().sum(2) + b.numpy()).T * 4, atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(e.numpy(), (a.numpy().sum(2) + b.numpy()) * (d.numpy() + a.numpy()).sum(2), atol=1e-4, rtol=1e-4)
@@ -1396,7 +1398,8 @@ class TestSchedule(unittest.TestCase):
     c = Tensor.empty(16, )
     r = a.sum(1) + c
     d = r[:4] * b
-    check_schedule(d, 2)
+    schedule = check_schedule(d, 2)
+    self.assertIs(store_val(schedule[0]).op, Ops.ADD)
 
   def test_multireduce_push_shrink_chase(self):
     Tensor.manual_seed(0)
@@ -1408,13 +1411,15 @@ class TestSchedule(unittest.TestCase):
     out = r[:4] * b + d.sum(1)[:4]
     # schedule = check_schedule(out, 2)
     schedule = check_schedule(out, 3)
+    self.assertIs(store_val(schedule[0]).op, Ops.ADD)
     run_schedule(schedule)
     np.testing.assert_allclose(out.numpy(), (a.numpy().sum(1) + c.numpy())[:4] * b.numpy() + d.numpy().sum(1)[:4], atol=1e-4, rtol=1e-4)
 
   def test_midreduce_nochase(self):
     a = Tensor.empty(16, 16)
     b = (a.sum(0) + a.max(1)) + 2
-    check_schedule(b, 2)
+    schedule = check_schedule(b, 2)
+    self.assertIs(store_val(schedule[0]).op, Ops.REDUCE_AXIS)
 
   def test_multireduce_midreduce_nochase(self):
     Tensor.manual_seed(0)
@@ -1422,6 +1427,7 @@ class TestSchedule(unittest.TestCase):
     b = (a.sum(0)+a.max(0) + a.max(1)+a.sum(1)) + 2
     # schedule = check_schedule(b, 2)
     schedule = check_schedule(b, 4)
+    self.assertIs(store_val(schedule[0]).op, Ops.REDUCE_AXIS)
     run_schedule(schedule)
     np.testing.assert_allclose(b.numpy(), a.numpy().sum(0)+a.numpy().max(0) + a.numpy().max(1)+a.numpy().sum(1)+2, atol=1e-4, rtol=1e-4)
 
@@ -1613,7 +1619,7 @@ class TestSchedule(unittest.TestCase):
     Tensor.manual_seed(0)
     x = Tensor.randn(10, 20).realize()
     out = x.argmax(1)
-    run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 3)) # TODO: push a reduceop through a reshape
 
   def test_conv2d(self): _test_conv2d(7)
   def test_conv2d_fused(self): _test_conv2d(5, FUSE_CONV_BW=1)
@@ -1727,8 +1733,7 @@ class TestIndexing(unittest.TestCase):
       s = Tensor.schedule(*lst)
       lowered = [x[1] for x in lower_schedule(s.copy())]
       kernels = [ei for ei in list(lowered) if isinstance(ei.prg, CompiledRunner)]
-      if FUSE_ARANGE and len(kernels) != cnt:
-        raise KernelCountException(f"{len(kernels)} != {cnt}")
+      if FUSE_ARANGE: self.assertEqual(len(kernels), cnt)
       for ei in lowered: ei.run(do_update_stats=True)
     return s
 
@@ -1742,7 +1747,7 @@ class TestIndexing(unittest.TestCase):
   def test_simple_indexing_alt(self):
     X = Tensor.arange(16).reshape(4, 4)
     xt = X[[1, 2], [1, 2]]
-    self.check_schedule(xt, 3)
+    self.check_schedule(xt, 5)
     np.testing.assert_equal(xt.numpy(), (np.arange(16).reshape(4, 4))[[1, 2], [1, 2]])
 
   def test_advanced_indexing(self):
@@ -1754,13 +1759,13 @@ class TestIndexing(unittest.TestCase):
   def test_advanced_indexing_alt(self):
     X = Tensor.arange(6).reshape(3, 2)+1
     xt = X[[Tensor([2]), Tensor([1])]]
-    self.check_schedule(xt, 3)
+    self.check_schedule(xt, 6)
     np.testing.assert_equal(xt.numpy(), 6)
 
   def test_advanced_simple_indexing_combined(self):
     X = Tensor.arange(16).reshape(4, 4)
     xt = X[1:2, [1, 2]]
-    self.check_schedule(xt, 2)
+    self.check_schedule(xt, 4)
 
   def test_push_through_reshape(self):
     Tensor.manual_seed(0)
@@ -1897,6 +1902,8 @@ class TestIndexing(unittest.TestCase):
     a = Tensor.arange(4).reshape(2, 2, 1).expand(2, 2, 2).contiguous().to("CPU")
     sched = self.check_schedule(a, 2) # NOTE: there is a contiguous between REDUCE_AXIS and COPY
     self.assertIs(sched[2].ast.op, Ops.COPY)
+    self.assertIs(store_val(sched[1]).op, Ops.LOAD)
+    self.assertIs(store_val(sched[0]).op, Ops.ADD)
     np.testing.assert_equal(a.numpy(), [[[0, 0], [1, 1]], [[2, 2], [3, 3]]])
 
   @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
@@ -1980,6 +1987,24 @@ class TestIndexing(unittest.TestCase):
     self.assertEqual(new_uop.st, ShapeTracker.from_shape((4,)).reshape((4, 1)))
     self.assertEqual(swizzle_cnt(new_uop), 0)
 
+  def test_no_rewrite_elementwise(self):
+    a = Tensor.empty(32, 32)
+    b = Tensor.empty(32, 32)
+    sink = (a+b).schedule()[0].ast
+    self.assertEqual(swizzle_cnt(sink), 0)
+
+  def test_simple_store_reshape(self):
+    a = Tensor.empty(32, 32).sum(axis=1)+Tensor.empty(1,32)
+    ast = a.schedule()[0].ast
+    self.assertEqual(ast.shape, (32, 1))
+    self.assertEqual(a.uop.shape, (1, 32))
+
+  def test_no_reshape_reduceop(self):
+    a = Tensor.empty(32, 32).sum(axis=(1,)).contiguous()
+    ast = a.schedule()[0].ast
+    self.assertEqual(ast.shape, (32, 1))
+    self.assertEqual(a.uop.shape, (32,))
+
 def swizzle_cnt(u:UOp) -> int:
   return len([x for x in u.toposort() if x.op is Ops.VIEW and len(x.src) != 0 and x.src[0].op not in {Ops.BUFFER, Ops.DEFINE_GLOBAL, Ops.ASSIGN}])
 
@@ -2067,6 +2092,7 @@ class TestSwizzle(unittest.TestCase):
     np.testing.assert_allclose(t.numpy(), x.numpy().sum(axis=1)+y.numpy().sum(axis=1), atol=1e-6, rtol=1e-3)
 
   # kernels can only have 1 or n in each dim
+  @unittest.expectedFailure
   def test_dont_parallelize_different_n(self):
     Tensor.manual_seed(0)
     x = Tensor.randn(4, 2, 2).realize()
@@ -2082,6 +2108,7 @@ class TestSwizzle(unittest.TestCase):
     run_schedule(check_schedule(t, 3))
     np.testing.assert_equal(t.numpy(), [[0.5, 0.5], [0.5, 0.5], [0., 0.]])
 
+def store_val(si:ScheduleItem): return si.ast.src[0].src[1]
 zero_pm = UPat(Ops.CONST, arg=0)
 class TestView(unittest.TestCase):
   def test_all_masked_out(self):
@@ -2090,6 +2117,7 @@ class TestView(unittest.TestCase):
     # all masked out, degrades to const 0
     b = a.pad(((0, 10), None))[10:]
     sched = check_schedule(b.contiguous(), 1)
+    assert zero_pm.match(store_val(sched[-1]), {})
     run_schedule(sched)
     np.testing.assert_equal(b.numpy(), 0)
 
@@ -2100,6 +2128,7 @@ class TestView(unittest.TestCase):
     assert b.shape == (10, 10)
     sched = check_schedule(b.contiguous(), 1)
     self.assertEqual(sched[-1].ast.full_shape, (10, 10))
+    assert zero_pm.match(store_val(sched[-1]), {})
     run_schedule(sched)
     np.testing.assert_equal(b.numpy(), 0)
 
@@ -2114,6 +2143,8 @@ class TestView(unittest.TestCase):
     b = a.pad(((0, 5), None))[5:]
     assert b.shape == (10, 10)
     sched = check_schedule(b.contiguous(), 1)
+    self.assertEqual(store_val(sched[-1]).op, Ops.LOAD)
+    self.assertEqual(store_val(sched[-1]).st_arg, b.uop.st)
     run_schedule(sched)
     np.testing.assert_allclose(b.numpy(), np.pad(a.numpy(), ((0, 5), (0, 0)))[5:])
 
@@ -2229,6 +2260,24 @@ class TestConst(unittest.TestCase):
     sched = a.schedule()
     self.assertEqual(len(sched), 1)
 
+  def test_const_ast(self):
+    a = Tensor.ones((4,)).pad((1, 1)).contiguous()
+    sched = a.schedule()
+    print(sched[0].ast)
+    const_ast_pattern = UPat(Ops.SINK, src=(UPat.store(UPat(), UPat.where(UPat(Ops.VALID), UPat.cvar("x"), UPat(Ops.CONST, arg=0))),))
+    self.assertEqual(len(const_ast_pattern.match(sched[0].ast, {})), 1)
+    run_schedule(sched)
+    self.assertListEqual(a.tolist(), [0, 1, 1, 1, 1, 0])
+
+  def test_unmasked_const_ast(self):
+    a = Tensor.ones((4,)).contiguous()
+    sched = a.schedule()
+    print(sched[0].ast)
+    const_ast_pattern = UPat(Ops.SINK, src=(UPat.store(UPat(), UPat(Ops.CONST)),))
+    self.assertEqual(len(const_ast_pattern.match(sched[0].ast, {})), 1)
+    run_schedule(sched)
+    self.assertListEqual(a.tolist(), [1, 1, 1, 1])
+
   # ** part 2: scheduler behavior when const folding happens later
 
   def test_const_folding_no_realize(self):

test/test_search.py

@@ -131,11 +131,10 @@ class TestBEAM(unittest.TestCase):
     assert tm
 
   def test_beam_unnamed_kernels(self):
-    from test.test_linearizer import push_views
     a = Tensor.rand(100)
     b = Tensor.rand(100)
     si = (a+b).schedule()[-1]
-    lin = Kernel(push_views(si.ast))
+    lin = Kernel(si.ast)
     bufs = bufs_from_lin(lin)
     # TODO: beam should have better instrumentation so we don't have to check this indirect thing
     kcount = len(Kernel.kernel_cnt)

test/test_softmax_fusion.py

@@ -111,12 +111,6 @@ class TestFuse(unittest.TestCase):
     with Context(NOOPT=1):
       self._test_fuse(Tensor.scaled_dot_product_attention, q, k, v, atol=1e-5)
 
-  def test_mismatch_reduce(self):
-    a = Tensor.ones(16, 10).contiguous().realize()
-    b = Tensor.ones(16, 20).contiguous().realize()
-    c = (a.sum(axis=1) + b.sum(axis=1)).fuse()
-    self.assertListEqual(c.tolist(), [30]*16)
-
 class TestSoftmaxFusion(unittest.TestCase):
   @classmethod
   def setUpClass(cls):

test/test_stunning.py

@@ -25,7 +25,7 @@ class TestStunning(unittest.TestCase):
     nv = a[12].cat(a[76]).tolist()
 
     vi = Variable('i', 0, a.shape[0]-1)
-    with self.assertRaisesRegex(AssertionError, "bind mismatch on"):
+    with self.assertRaisesRegex(AssertionError, "different values for the same key"):
       wv = a[vi.bind(12)].cat(a[vi.bind(76)]).tolist()
       self.assertListEqual(nv, wv)
 

test/test_uops.py

@@ -508,6 +508,19 @@ class TestShapeSpec(unittest.TestCase):
     a = Tensor.ones((4, 4)).uop
     self.assertEqual(a.st, ShapeTracker.from_shape(()).reshape((1,1)).expand((4,4)))
 
+  def test_padded_const(self):
+    a = Tensor.ones((1, 1)).pad(((1, 1), (1, 1)))
+    ast = a.contiguous().schedule()[0].ast
+    valid_pattern = UPat(Ops.WHERE, src=(UPat(Ops.VALID), UPat.cvar(), UPat.cvar()))
+    valid_ternary = [x for x in ast.toposort() if valid_pattern.match(x, {})][0]
+    # the WHERE outputs a contiguous (3, 3)
+    self.assertEqual(valid_ternary.st, ShapeTracker.from_shape((3, 3)))
+    valid, x, y = valid_ternary.src
+    # very notably, only the first source is padded
+    self.assertIsNotNone(valid.st.views[-1].mask)
+    assert x.st.views[-1].mask is y.st.views[-1].mask is None
+    assert all(s.shape == (3, 3) for s in valid_ternary.src)
+
   # NOTE: CONST ShapeTracker comes from its source
   def test_scalar_const(self):
     a = Tensor(0).uop

test/test_uops_stats.py

@@ -6,7 +6,7 @@ from tinygrad.renderer import Estimates
 from tinygrad.codegen import full_rewrite
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.dtype import dtypes
-from tinygrad.opt.kernel import Opt, OptOps, KernelOptError
+from tinygrad.opt.kernel import Kernel, Opt, OptOps, KernelOptError
 from tinygrad.device import Device
 
 def flops_mem(uops, ignore_indexing=False):
@@ -173,60 +173,76 @@ class TestStatsOptimized(unittest.TestCase):
     self.assertEqual(p.estimates.mem, 3*N*N*4) # 3 NxN mats with floats
 
   def test_gemm(self):
-    p = get_program(self.ast_gemm, opts=[])
+    k = Kernel(self.ast_gemm)
+    p = get_program(k.get_optimized_ast(), k.opts)
     self.check_gemm(p)
     self.assertEqual(p.estimates.lds, 2*N*N*N*4 + 4*N*N)
 
   def test_gemm_tc_unroll(self):
-    try:
-      p = get_program(self.ast_gemm, opts=[Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.UNROLL, 0, 2)])
-    except KernelOptError:
-      raise unittest.SkipTest("no tensor cores")
+    k = Kernel(self.ast_gemm)
+    if not k.apply_tensor_cores(): self.skipTest("no tensor cores")
+    k.apply_opt(Opt(OptOps.UNROLL, 0, 2))
+    p = get_program(k.get_optimized_ast(), k.opts)
     print(p.src)
     self.check_gemm(p)
 
   # this is a good lesson about why UPCASTing is a good idea
 
   def test_gemm_one_upcasted(self):
-    p = get_program(self.ast_gemm, opts=[Opt(OptOps.UPCAST, 0, 4)])
+    k = Kernel(self.ast_gemm)
+    k.apply_opt(Opt(OptOps.UPCAST, 0, 4))
+    p = get_program(k.get_optimized_ast(), k.opts)
     self.check_gemm(p)
     self.assertEqual(p.estimates.lds, N*N*N*4 + N*N*N*4//4 + 4*N*N)
 
   def test_gemm_upcasted(self):
-    p = get_program(self.ast_gemm, opts=[Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4)])
+    k = Kernel(self.ast_gemm)
+    k.apply_opt(Opt(OptOps.UPCAST, 0, 4))
+    k.apply_opt(Opt(OptOps.UPCAST, 1, 4))
+    k.apply_opt(Opt(OptOps.UNROLL, 0, 4))
+    p = get_program(k.get_optimized_ast(), k.opts)
     self.check_gemm(p)
     self.assertEqual(p.estimates.lds, 2*N*N*N*4//4 + 4*N*N)
 
   def test_gemm_upcasted_locals(self):
+    k = Kernel(self.ast_gemm)
+    k.apply_opt(Opt(OptOps.UPCAST, 0, 4))
+    k.apply_opt(Opt(OptOps.UPCAST, 1, 4))
     try:
-      p = get_program(self.ast_gemm, opts=[Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4),
-                                           Opt(OptOps.LOCAL, 0, 4),  Opt(OptOps.LOCAL, 1, 4)])
+      k.apply_opt(Opt(OptOps.LOCAL, 0, 5))
+      k.apply_opt(Opt(OptOps.LOCAL, 1, 5))
     except KernelOptError:
       raise unittest.SkipTest("no locals")
+    p = get_program(k.get_optimized_ast(), k.opts)
     self.check_gemm(p)
     self.assertEqual(p.estimates.lds, 2*N*N*N*4//4 + 4*N*N)
 
   def test_gemm_group(self):
+    k = Kernel(self.ast_gemm)
     try:
-      p = get_program(self.ast_gemm, opts=[Opt(OptOps.GROUP, 0, 4)])
+      k.apply_opt(Opt(OptOps.GROUP, 0, 4))
     except KernelOptError:
       raise unittest.SkipTest("no locals")
     SZ = N*N*4
+    p = get_program(k.get_optimized_ast(), k.opts)
     # NOTE: these are sort of wrong. they aren't honoring the IF statement
     self.check_gemm(p, extra_flops=SZ*4)
     self.assertEqual(p.estimates.lds, 2*N*N*N*4 + SZ*4 + (SZ*4 + 4*N*N)*4)
 
   def test_reduce(self):
-    p = get_program(self.ast_reduce, opts=[])
+    k = Kernel(self.ast_reduce)
+    p = get_program(k.get_optimized_ast(), k.opts)
     print(p.name, p.estimates.ops, p.estimates.mem, p.estimates.lds)
     self.assertEqual(p.estimates.ops, N*N)
     self.assertEqual(p.estimates.mem, N*N*4 + 4)
 
   def test_reduce_group(self):
+    k = Kernel(self.ast_reduce)
     try:
-      p = get_program(self.ast_reduce, opts=[Opt(OptOps.GROUP, 0, 50)])
+      k.apply_opt(Opt(OptOps.GROUP, 0, 50))
     except KernelOptError:
       raise unittest.SkipTest("no locals")
+    p = get_program(k.get_optimized_ast(), k.opts)
     # NOTE: these are wrong, they don't respect the if statement
     print(p.name, p.estimates.ops, p.estimates.mem, p.estimates.lds)
 

test/test_winograd.py

@@ -1,7 +1,10 @@
 import unittest
 import numpy as np
 from tinygrad import Tensor, GlobalCounters, dtypes, Context, nn
-from tinygrad.helpers import CI, Profiling, WINO, getenv
+from tinygrad.uop.ops import Ops
+from tinygrad.helpers import Timing, CI, Profiling, WINO, DEBUG, getenv
+from tinygrad.opt.kernel import Kernel
+from tinygrad.opt.heuristic import hand_coded_optimizations
 
 class TestWinogradClose(unittest.TestCase):
   def test_close(self):
@@ -25,6 +28,30 @@ class TestWinograd(unittest.TestCase):
   def tearDown(self):
     WINO.value = self.old
 
+  def test_speed(self):
+    x = Tensor.empty(1,4,9,9)
+    w = Tensor.empty(4,4,3,3)
+
+    with Timing("running conv: "):
+      out = Tensor.conv2d(x, w)
+
+    with Timing("scheduling: "):
+      sched = out.schedule()
+
+    for i,s in enumerate(sched):
+      if s.ast.op is not Ops.SINK: continue
+      ops = s.ast.toposort()
+      with Timing(f"linearize {i} with {len(ops):4d} ops: "):
+        l = Kernel(s.ast)
+        l.apply_opts(hand_coded_optimizations(l))
+      assert len(l.sts) <= 256  # just the current value to prevent regression
+      if DEBUG >= 2: print(f"{len(l.sts):4d} shapetrackers with max {max(len(x.views) for x in l.sts)} views")
+      for st in l.sts:
+        assert len(st.views) <= 2, "too many views in winograd"
+        if DEBUG >= 3:
+          print(f"{len(st.views):3d} views")
+          for v in st.views: print(v)
+
   def test_profile(self):
     x,w = Tensor.rand(1,4,9,9).realize(), Tensor.rand(4,4,3,3).realize()
     with Profiling(enabled=not CI, sort='time'):

test/unit/test_conv_shapetracker.py

@@ -0,0 +1,55 @@
+#!/usr/bin/env python
+import unittest
+from tinygrad.uop.ops import Ops
+from tinygrad.tensor import Tensor
+from tinygrad.nn import Conv2d
+from tinygrad.shape.shapetracker import ShapeTracker, View
+from tinygrad.helpers import prod
+from test.unit.test_shapetracker import shapetracker_getitem
+
+class TestConvShapetracker(unittest.TestCase):
+  def test_conv_3x3_one_view(self):
+    conv = Conv2d(16, 32, (3, 3))
+    # first run to init the weights, they are scheduled.
+    conv(Tensor.empty(1, 16, 10, 10)).schedule()
+    # run it again to get the kernels
+    sched = [si for si in conv(Tensor.empty(1, 16, 10, 10)).schedule() if si.ast.op is Ops.SINK]
+    assert len(sched) == 1, f"conv should only have one kernel, getting {len(sched)}"
+    for st in [x.st_arg for x in sched[0].ast.toposort() if x.op is Ops.LOAD]:
+      assert len(st.views) == 1
+
+  def test_conv_2x2_backward_one_view(self):
+    X = Tensor.rand(1, 1, 3, 3, requires_grad=True)
+    conv = Conv2d(1, 1, (2, 2), bias=False)
+    conv(X).mean().backward()
+    si = X.grad.schedule()[-1]
+    print(si)
+    ldb = [x for x in si.ast.toposort() if x.op is Ops.LOAD][0]
+    st: ShapeTracker = ldb.st_arg.simplify()
+    print(si.bufs[1].size)
+    self.assertEqual(si.bufs[1].size, st.real_size())
+    for v in st.views: print(v)
+
+    # same st
+    test_st = ShapeTracker((
+      View(shape=(1, 1, 2, 4, 2, 4), strides=(0, 0, 2, 8, 1, 4), offset=0, mask=((0, 1), (0, 1), (0, 2), (0, 2), (0, 2), (0, 2)), contiguous=False),
+      View(shape=(1, 1, 1, 1, 3, 3, 3, 3), strides=(0, 0, 0, 0, 24, 8, 3, 1), offset=0,
+           mask=((0, 1), (0, 1), (0, 1), (0, 1), (0, 2), (0, 3), (0, 2), (0, 3)), contiguous=False)))
+    #test_st = ShapeTracker((
+    #  View(shape=(2,4), strides=(1,4), offset=0, mask=None, contiguous=False),
+    #)).simplify()
+      #View(shape=(1, 1, 2, 4, 2, 4), strides=(0, 0, 2, 8, 1, 4), offset=0, mask=((0, 1), (0, 1), (0, 2), (0, 2), (0, 2), (0, 2)), contiguous=False),
+      #View(shape=(1, 1, 1, 1, 3, 3, 3, 3), strides=(0, 0, 0, 0, 24, 8, 3, 1), offset=0,
+      #     mask=((0, 1), (0, 1), (0, 1), (0, 1), (0, 2), (0, 3), (0, 2), (0, 3)), contiguous=False))).simplify()
+    print("*** new ***")
+    for v in test_st.views: print(v)
+    for i in range(prod(st.shape)):
+      i1, i2 = shapetracker_getitem(st, i), shapetracker_getitem(test_st, i)
+      print(i, i1, i2, si.bufs[1].size, i1==i2)
+      #self.assertEqual(i1, i2)
+
+    with self.assertRaises(AssertionError):
+      assert len(st.views) <= 2
+
+if __name__ == '__main__':
+  unittest.main()

test/unit/test_llm_tokenizer.py

@@ -1,8 +1,7 @@
-import unittest, base64, functools, sys
+import unittest, base64, functools
 from tinygrad.apps.llm import SimpleTokenizer, get_llama_re
 from tinygrad.helpers import fetch
 
-@unittest.skipIf(sys.platform == 'win32', "fetch race condition on Windows")
 class TestLLMTokenizer(unittest.TestCase):
   @functools.cached_property
   def basic_tok(self): return SimpleTokenizer(".*", { b"a": 0, b"b": 1, b"c": 2, b"ab": 3, b"bc": 4 }, { "<x>": 5, "<y>": 6, "<z>": 7 })

test/unit/test_search.py

@@ -5,7 +5,6 @@ from tinygrad.device import Buffer
 from tinygrad.opt.search import get_test_global_size, bufs_from_lin
 from tinygrad.helpers import GlobalCounters
 from extra.optimization.helpers import time_linearizer
-from test.test_linearizer import push_views
 
 class TestSearchUtil(unittest.TestCase):
   def test_get_test_global_size(self):
@@ -26,7 +25,7 @@ class TestSearchUtil(unittest.TestCase):
     a = Tensor.randn(4, 4).realize()
     b = a+a[0]
     si = b.schedule()[0]
-    rawbufs = bufs_from_lin(Kernel(push_views(si.ast)))
+    rawbufs = bufs_from_lin(Kernel(si.ast))
     assert len(rawbufs) == 2
     assert all(r is not None for r in rawbufs)
     assert all(isinstance(r, Buffer) for r in rawbufs)
@@ -39,13 +38,13 @@ class TestTimeLinearizer(unittest.TestCase):
     si = (a+1).schedule()[0]
     # create fresh empty buffers
     rawbufs = [Buffer(b.device, b.size, b.dtype).allocate() for b in si.bufs]
-    tm = time_linearizer(Kernel(push_views(si.ast)), rawbufs, allow_test_size=False, cnt=10, disable_cache=True)
+    tm = time_linearizer(Kernel(si.ast), rawbufs, allow_test_size=False, cnt=10, disable_cache=True)
     assert tm > 0 and tm != float('inf')
 
   # Ensure that the kernel count is not incremented by time_linearizer when clearing l2
   def test_kernel_count(self):
     ast = Tensor.zeros(16).contiguous().kernelize().uop.src[1].arg.ast
-    lin = Kernel(push_views(ast))
+    lin = Kernel(ast)
     bufs = bufs_from_lin(lin)
 
     kernel_count = GlobalCounters.kernel_count

test/unit/test_uop_spec.py

@@ -34,6 +34,17 @@ class TestUOpSpec(unittest.TestCase):
     store = UOp(Ops.STORE, dtypes.void, (buf_0.view(ShapeTracker.from_shape((32, 1))), a+b))
     helper_test_verify_ast(store)
 
+  def test_exactly_one_full_shape(self):
+    dtype = dtypes.int
+    bufs = [UOp(Ops.DEFINE_GLOBAL, dtype.ptr(), (), i) for i in range(6)]
+    a = UOp(Ops.LOAD, dtype, (bufs[2].view(ShapeTracker.from_shape((32, 1))),))
+    b = UOp(Ops.LOAD, dtype, (bufs[3].view(ShapeTracker.from_shape((32, 1))),))
+    st0 = UOp.store(bufs[0].view(ShapeTracker.from_shape((32, 1))), a+b)
+    a = UOp(Ops.LOAD, dtype, (bufs[4].view(ShapeTracker.from_shape((32, 32))),))
+    b = UOp(Ops.LOAD, dtype, (bufs[5].view(ShapeTracker.from_shape((32, 32))),))
+    st1 = UOp.store(bufs[1].view(ShapeTracker.from_shape((32, 32))), a+b)
+    with self.assertRaises(InvalidASTException): helper_test_verify_ast(st0, st1)
+
   def test_no_implicit_broadcasting(self):
     bufs = [UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), i) for i in range(2)]
     a = UOp(Ops.LOAD, dtypes.float, (bufs[1].view(ShapeTracker.from_shape((4, 32))),))
@@ -62,6 +73,14 @@ class TestUOpSpec(unittest.TestCase):
     st = UOp.store(bufs[0].view(ShapeTracker.from_shape((32, 1))), r+a)
     with self.assertRaises(InvalidASTException): helper_test_verify_ast(st)
 
+  def test_buffer_uops_st(self):
+    a = Tensor.randn(4, 4)+2
+    helper_test_verify_ast(ast:=a.schedule()[-1].ast)
+    store_st = [u.st for u in ast.toposort() if u.op is Ops.STORE][0]
+    self.assertEqual(store_st, ShapeTracker.from_shape((4, 4)))
+    const_st = [u.st for u in ast.toposort() if u.op is Ops.CONST][0]
+    self.assertEqual(const_st, ShapeTracker.from_shape((1, 1)).expand((4, 4)))
+
   def test_assert_swizzle(self):
     buf = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
     a = UOp(Ops.LOAD, dtypes.float, (buf.view(ShapeTracker.from_shape((32, 1))),))

tinygrad/codegen/__init__.py

@@ -30,12 +30,6 @@ class RewriteStep:
 
 def apply_rewrites(sink:UOp, rewrites:list[RewriteStep]): return functools.reduce(lambda x,f: f(x), rewrites, sink)
 
-rewrites_for_views = [
-  RewriteStep(view_left, name="Main View Left"),
-  RewriteStep(view_right, name="Main View Right"),
-  RewriteStep(view_left+fix_kernel_ops, bottom_up=True, name="Finalize Kernel"),
-]
-
 rewrites_for_linearizer = [
   RewriteStep(block_create, ctx=BlockContext.from_sink, name="Linearizer: Create Blocks", bottom_up=True),
   RewriteStep(pm_blockend_merge, name="Linearizer: Merge Blockends"),
@@ -51,8 +45,10 @@ def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVEC
   # ** lowerer (rewrite_shapetracker_with_index) **
   ret: list[RewriteStep] = []
 
-  # view pushing
-  ret.extend(rewrites_for_views)
+  # TODO: move these to codegen
+  ret.append(RewriteStep(view_left, name="Main View Left"))
+  ret.append(RewriteStep(view_right, name="Main View Right"))
+  ret.append(RewriteStep(view_left+fix_kernel_ops, bottom_up=True, name="replace buffer"))
 
   # this is kernel.py
   ret.append(RewriteStep(pm_optimize, ctx=lambda _: opts, name="optimize ast"))

tinygrad/engine/realize.py

@@ -3,17 +3,16 @@ import time, pprint
 from dataclasses import dataclass, replace, field
 from tinygrad.helpers import all_same, colored, DEBUG, GlobalCounters, ansilen, BEAM, NOOPT, all_int, CAPTURING, Metadata, TRACEMETA, TracingKey
 from tinygrad.helpers import DEVECTORIZE, time_to_str, VALIDATE_WITH_CPU, getenv, cpu_profile
-from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, Variable, sym_infer, graph_rewrite, print_uops, track_rewrites, KernelInfo
+from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, Variable, sym_infer, graph_rewrite, print_uops, track_rewrites
 from tinygrad.device import Device, Buffer
 from tinygrad.renderer import Renderer, ProgramSpec, Estimates
 from tinygrad.engine.schedule import ScheduleItem
 from tinygrad.codegen import full_rewrite
-from tinygrad.opt.kernel import Opt
 
 # **************** Program Creation ****************
 
 @track_rewrites(name=lambda _ast,_renderer,ret: TracingKey(ret.name, (ret.function_name, ret.ast), ret=ret))
-def get_program(ast:UOp, renderer:Renderer|None=None, opts:list[Opt]|None=None) -> ProgramSpec:
+def get_program(ast:UOp, renderer:Renderer) -> ProgramSpec:
   """
   Transform an AST into a ProgramSpec. May trigger BEAM search.
 
@@ -28,10 +27,6 @@ def get_program(ast:UOp, renderer:Renderer|None=None, opts:list[Opt]|None=None)
   if getenv("VIZ"): graph_rewrite(ast, PatternMatcher([]), name="View Base AST")
 
   # linearize
-  if renderer is None: renderer = Device.default.renderer
-  if opts is not None:
-    assert ast.arg is None, "can't apply opts if sink has an arg"
-    ast = ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts)))
   try:
     uops = full_rewrite(ast, renderer)
   except RuntimeError:
@@ -44,7 +39,7 @@ def get_program(ast:UOp, renderer:Renderer|None=None, opts:list[Opt]|None=None)
   if DEBUG >= 6: print_uops(uops)
   src = renderer.render(uops)
 
-  return ProgramSpec(uops[-1].arg.name if uops[-1].arg is not None else "test", src, renderer.device, ast, uops,
+  return ProgramSpec(uops[-1].arg.name, src, renderer.device, ast, uops,
                      global_size=[1,1,1] if renderer.has_local else None, local_size=[1,1,1] if renderer.has_local else None)
 
 # **************** Runners ****************

tinygrad/engine/schedule.py

@@ -1,9 +1,9 @@
 from typing import cast
 from dataclasses import dataclass, field
 from collections import deque, defaultdict
-from tinygrad.uop.ops import UOp, Variable, Ops, buffers
+from tinygrad.uop.ops import UOp, Variable, Ops, UPat, PatternMatcher, graph_rewrite, buffers
 from tinygrad.device import Device, Buffer, MultiBuffer
-from tinygrad.helpers import Metadata, all_same
+from tinygrad.helpers import Metadata, unwrap, all_same, merge_dicts
 
 # **** ScheduleItem return type
 
@@ -14,13 +14,32 @@ class ScheduleItem:
   metadata: tuple[Metadata, ...] = ()
   fixedvars: dict[Variable, int] = field(default_factory=dict)
 
+# **** unbind Variables
+
+def unbind_view(ctx:list[dict[Variable, int]], x:UOp):
+  st = unwrap(x.st).simplify()
+  if any(x.op is Ops.BIND for x in st.vars()):
+    st, var_vals = st.unbind()
+    ctx.append(var_vals)
+    return x.replace(arg=st)
+  return None
+
+def unbind_bind(ctx:list[dict[Variable, int]], x:UOp):
+  var, val = x.unbind()
+  ctx.append({var.replace(src=()):val})
+  return var
+
+pm_unbind = PatternMatcher([
+  (UPat(Ops.VIEW, name="x"), unbind_view),
+  (UPat(Ops.BIND, name="x"), unbind_bind),
+])
+
 # **** schedule linearizer
 
 def create_schedule_with_vars(sched_sink:UOp) -> tuple[list[ScheduleItem], dict[Variable, int]]:
   # construct the KERNEL children graph based on assigns
   children: defaultdict[UOp, list[UOp]] = defaultdict(list)
   in_degree: dict[UOp, int] = {}
-  var_vals: dict[Variable, int] = {}
   for u in sched_sink.toposort():
     if u.op is not Ops.ASSIGN: continue  # anything that's not an ASSIGN doesn't write a kernel, so we can skip
     k = u.src[1]
@@ -38,10 +57,6 @@ def create_schedule_with_vars(sched_sink:UOp) -> tuple[list[ScheduleItem], dict[
             in_degree[k] += 1
       elif s.op is Ops.BUFFER:
         pass  # a BUFFER is already realized, nothing to do here
-      elif s.op is Ops.BIND:
-        var, val = s.unbind()
-        assert var not in var_vals or var_vals[var] == val, f"bind mismatch on {var}, {var_vals[var]} != {val}"
-        var_vals[var] = val
       else:
         raise RuntimeError(f"input to kernel must be ASSIGN or BUFFER, not {s.op}")
 
@@ -58,16 +73,20 @@ def create_schedule_with_vars(sched_sink:UOp) -> tuple[list[ScheduleItem], dict[
     if v == 0: queues[_heuristic(k)].append(k)
 
   schedule: list[ScheduleItem] = []
+  var_vals: dict[Variable, int] = {}
   while last_queue or any(queues.values()):
     if not last_queue: last_heuristic, last_queue = min((it for it in queues.items() if it[1]), key=lambda x: abs(x[0]-last_heuristic))
     k = last_queue.popleft()
-    ast = k.arg.ast
+    # unbind var_vals from the kernel
+    local_var_vals: list[dict[Variable, int]] = []
+    ast = graph_rewrite(k.arg.ast, pm_unbind, ctx=local_var_vals, name="unbind vars")
+    var_vals = merge_dicts([var_vals, *local_var_vals])
     # create subbuffers if needed
     if ast.op is Ops.BUFFER_VIEW:
       base = k.src[1].buf_uop.buffer
       assert isinstance(base, Buffer), "base can't be MultiBuffer"
       buffers[k.src[0]] = base.view(k.size, ast.dtype, ast.arg[1]*base.dtype.itemsize)
-    ubufs = tuple(s.buf_uop.buffer for s in k.src if s.op is not Ops.BIND)
+    ubufs = tuple(s.buf_uop.buffer for s in k.src)
     if any(isinstance(x, MultiBuffer) for x in ubufs):
       assert all(isinstance(x, MultiBuffer) for x in ubufs), "kernel must all be multibuffer"
       dnums = [x for x in ast.variables() if x.arg[0] == '_device_num']

tinygrad/helpers.py

@@ -131,7 +131,7 @@ JIT_BATCH_SIZE = ContextVar("JIT_BATCH_SIZE", 32)
 WINO, CAPTURING, TRACEMETA = ContextVar("WINO", 0), ContextVar("CAPTURING", 1), ContextVar("TRACEMETA", 1)
 USE_TC, TC_SELECT, TC_OPT, AMX = ContextVar("TC", 1), ContextVar("TC_SELECT", -1), ContextVar("TC_OPT", 0), ContextVar("AMX", 0)
 TRANSCENDENTAL, TC_SEARCH_OVER_SHAPE, NOLOCALS = ContextVar("TRANSCENDENTAL", 1), ContextVar("TC_SEARCH_OVER_SHAPE", 1), ContextVar("NOLOCALS", 0)
-FUSE_ARANGE, FUSE_CONV_BW = ContextVar("FUSE_ARANGE", 1), ContextVar("FUSE_CONV_BW", 0)
+FUSE_ARANGE, FUSE_CONV_BW = ContextVar("FUSE_ARANGE", 0), ContextVar("FUSE_CONV_BW", 0)
 SPLIT_REDUCEOP, NO_MEMORY_PLANNER, RING = ContextVar("SPLIT_REDUCEOP", 1), ContextVar("NO_MEMORY_PLANNER", 0), ContextVar("RING", 1)
 PICKLE_BUFFERS, PROFILE, LRU = ContextVar("PICKLE_BUFFERS", 1), ContextVar("PROFILE", getenv("VIZ")), ContextVar("LRU", 1)
 CACHELEVEL, IGNORE_BEAM_CACHE, DEVECTORIZE = ContextVar("CACHELEVEL", 2), ContextVar("IGNORE_BEAM_CACHE", 0), ContextVar("DEVECTORIZE", 1)

tinygrad/opt/heuristic.py

@@ -80,21 +80,18 @@ def hand_coded_optimizations(k:Kernel) -> list[Opt]:
     else: break
 
   # if last reduce dim is small(ish), loop unroll the reduce
-  # NOTE: this can fail on multireduce with mismatching dimensions, this is okay
-  try:
-    upcast_size = prod(k.full_shape[a] for a in k.axes_of(AxisType.UPCAST, AxisType.UNROLL))
-    if k.unrollable_dims and (upcast_size <= 4 or not k.axes_of(AxisType.UNROLL)) and (upcast_size < 64):
-      if (s:=k.full_shape[k.unrollable_dims[-1]]) <= 32:
+  upcast_size = prod(k.full_shape[a] for a in k.axes_of(AxisType.UPCAST, AxisType.UNROLL))
+  if k.unrollable_dims and (upcast_size <= 4 or not k.axes_of(AxisType.UNROLL)) and (upcast_size < 64):
+    if (s:=k.full_shape[k.unrollable_dims[-1]]) <= 32:
+      k.apply_opt(Opt(OptOps.UNROLL, len(k.unrollable_dims)-1, 0))
+      # if it's small, upcast a second reduce dimension too
+      if k.unrollable_dims and s <= 3 and k.full_shape[k.unrollable_dims[-1]] <= 3:
         k.apply_opt(Opt(OptOps.UNROLL, len(k.unrollable_dims)-1, 0))
-        # if it's small, upcast a second reduce dimension too
-        if k.unrollable_dims and s <= 3 and k.full_shape[k.unrollable_dims[-1]] <= 3:
-          k.apply_opt(Opt(OptOps.UNROLL, len(k.unrollable_dims)-1, 0))
-      else:
-        for splits in [4]:
-          if k.full_shape[axis:=k.unrollable_dims[-1]]%splits == 0:
-            k.apply_opt(Opt(OptOps.UNROLL, len(k.unrollable_dims)-1, splits))
-            break
-  except KernelOptError: pass
+    else:
+      for splits in [4]:
+        if k.full_shape[axis:=k.unrollable_dims[-1]]%splits == 0:
+          k.apply_opt(Opt(OptOps.UNROLL, len(k.unrollable_dims)-1, splits))
+          break
 
   # if nothing at all is upcasted and it's easy to, do an upcast
   for splits in [4]:

tinygrad/opt/kernel.py

@@ -14,7 +14,7 @@ from tinygrad.dtype import ImageDType, AddrSpace
 from tinygrad.helpers import all_same, colored, ansilen, dedup, prod, round_up, to_function_name, unwrap, argfix, DEBUG, TC_SELECT, TC_OPT, AMX
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import strides_for_shape, get_contraction
-from tinygrad.opt.swizzler import view_left, view_left_through_load
+from tinygrad.opt.swizzler import view_left
 
 class OptOps(Enum):
   TC = auto(); UPCAST = auto(); UNROLL = auto(); LOCAL = auto() # noqa: E702
@@ -271,10 +271,7 @@ class Kernel:
       check(isinstance(opt.arg, int), "arg should be int")
       amt = arg if (arg:=cast(int, opt.arg)) != 0 else self.full_shape[axis]
       check(isinstance(amt, int) and amt != 1, f"shift/padto of {amt=}, 1 or symbolic amount is meaningless")
-      if opt.op is not OptOps.PADTO:
-        # we check both the full_shape and each shape
-        check(self.full_shape[axis] % amt == 0, f"no longer valid shift {self.full_shape[axis]=}, {amt=}")
-        for st in self.sts: check(st.shape[axis] == 1 or st.shape[axis] % amt == 0, f"no longer valid shift {st.shape[axis]=}, {amt=}")
+      if opt.op is not OptOps.PADTO: check(self.full_shape[axis] % amt == 0, f"no longer valid shift {self.full_shape[axis]=}, {amt=}")
     else: amt = -1
 
     if self.reduceop is not None and (opt.op in {OptOps.GROUP, OptOps.GROUPTOP} or \
@@ -454,8 +451,7 @@ class Kernel:
         return ret.replace(src=(ret.src[0].replace(arg=st),)+ret.src[1:])
       if op.op is Ops.SINK:
         # NOTE: should group_for_reduces be added to the local_dims?
-        # TODO: arg.name should be able to be None
-        kernel_name = ret.arg.name if ret.arg is not None and ret.arg.name != "test" else self.name if name_override is None else name_override
+        kernel_name = ret.arg.name if ret.arg is not None else self.name if name_override is None else name_override
         return ret.replace(arg=KernelInfo(kernel_name, tuple(self.axis_types), self.dont_use_locals, tuple(self.applied_opts)))
       if op.op is Ops.REDUCE_AXIS:
         reduce_idx = len(self.bufs) + self.reduceops.index(op) * 2
@@ -507,4 +503,4 @@ class Kernel:
     self.finalized = True
     fixed_ast = fixup_ast(self.ast)
     del fixup_ast
-    return graph_rewrite(fixed_ast, view_left+view_left_through_load, name="fixup optimized AST")
+    return graph_rewrite(fixed_ast, view_left, name="fixup optimized AST")

tinygrad/opt/swizzler.py

@@ -44,18 +44,12 @@ def reduce_push_add_ones(src:UOp, r:UOp, view:UOp):
 
 view_left = merge_views+PatternMatcher([
   # view before elementwise and buffer ops
-  (UPat(Ops.VIEW, src=(UPat({*GroupOp.ALU, Ops.CAST, Ops.BITCAST, Ops.BIND, Ops.STORE, Ops.VALID, Ops.SINK}, name="e"),), name="view"),
+  (UPat(Ops.VIEW, src=(UPat({*GroupOp.ALU, Ops.CAST, Ops.BITCAST, Ops.BIND, Ops.LOAD, Ops.STORE, Ops.VALID, Ops.SINK}, name="e"),), name="view"),
    lambda e,view: e.replace(src=tuple(s.view(view.st) for s in e.src))),
   # if there's ones added after reduce, put this before the reduce
   (UPat(Ops.VIEW, src=(UPat(Ops.REDUCE_AXIS, src=(UPat.var("src"),), name="r"),), name="view"), reduce_push_add_ones),
 ])
 
-view_left_through_load = PatternMatcher([
-  # view before load
-  (UPat(Ops.VIEW, src=(UPat(Ops.LOAD, name="e"),), name="view"),
-   lambda e,view: e.replace(src=tuple(s.view(view.st) for s in e.src))),
-])
-
 def apply_swizzle(u:UOp) -> UOp: return graph_rewrite(u, view_left, name="Sub View Left")
 
 # change reduceop axes and input ShapeTrackers, view gets replaced with a reshape.
@@ -102,12 +96,10 @@ view_right = merge_views+PatternMatcher([
   # apply view after reduceops
   (UPat(Ops.REDUCE_AXIS, src=(UPat(Ops.VIEW, src=(UPat(GroupOp.All-ALWAYS_CONTIGUOUS, name="src"),), name="v"),), name="r"), reduceop_view_right),
   # apply view after elementwise ops
-  (UPat(GroupOp.All-{Ops.SINK, Ops.REDUCE_AXIS}, name="root"), elementwise_view_right),
+  (UPat(GroupOp.All-{Ops.SINK, Ops.REDUCE_AXIS, Ops.LOAD, Ops.STORE}, name="root"), elementwise_view_right),
   # merge axes for double reduce (invert of SPLIT_REDUCEOP=1)
   (UPat(Ops.REDUCE_AXIS, src=(UPat(Ops.REDUCE_AXIS, name="r1"),), name="r2"),
    lambda r1,r2: r1.replace(arg=(r1.arg[0], r2.arg[1]+r1.arg[1])) if r1.arg[0] is r2.arg[0] else None),
-  # remove view from sink
-  (UPat(Ops.VIEW, name="v").sink(name="sink"), lambda v,sink: v.src[0].sink(arg=sink.arg)),
 ])
 
 def check_load_st(glbl:UOp, view:UOp):
@@ -120,13 +112,20 @@ def check_load_st(glbl:UOp, view:UOp):
   raise RuntimeError("self operand of augmented assign must be contiguous.\nhelp: consider using .contiguous():\n"
                      +colored("   - a += a.T\n", "red")+colored("   + a += a.T.contiguous()", "green"))
 
-fix_kernel_ops = view_left_through_load+PatternMatcher([
-  # add view to LOAD and STORE
-  (UPat(Ops.DEFINE_GLOBAL, name="g").load(), lambda g: g.view(g.st).load()),
-  (UPat(Ops.DEFINE_GLOBAL, name="g").store(UPat.var('x')), lambda g,x: g.view(g.st).store(x)),
+fix_kernel_ops = PatternMatcher([
+  # add the LOAD
+  (UPat(Ops.DEFINE_GLOBAL, name="x"), lambda x: x.replace(tag=None).view(x.st).load() if x.tag is not None else None),
+  # STORE (except for meta ops)
+  (UPat(Ops.SINK, src=UPat(GroupOp.All-{Ops.STORE}), name="sink"), lambda sink:
+   UOp.sink(*[UOp.store(UOp(Ops.DEFINE_GLOBAL, (s:=x.base).dtype.ptr(s.st.real_size()), (), i).view(s.st), s) for i,x in enumerate(sink.src)])),
+  # passthrough ASSIGN
+  (UPat(Ops.ASSIGN, name="x"), lambda x: x.src[1]),
   # VALID
   (UPat(Ops.VIEW, src=(UPat.cvar(),), name="self"),
    lambda self: UOp.where(UOp(Ops.VALID, dtypes.bool, (UOp(Ops.VIEW, arg=self.st),)), self.const_like(self.base.arg), 0)),
+  # remove CONTIGUOUS/DEVICE from kernel AST
+  (UPat((Ops.CONTIGUOUS, Ops.MSELECT), src=(UPat.var("x"),)), lambda x: x),
+  (UPat(Ops.VIEW, src=(UPat(Ops.DEVICE),), name="view"), lambda view: view.replace(src=())),
   # no ImageDType after index
   (UPat(GroupOp.All-{Ops.DEFINE_GLOBAL, Ops.VIEW}, name="x"), lambda x: x.replace(dtype=x.dtype.base) if isinstance(x.dtype, ImageDType) else None),
   # if this kernel also assigns to the loaded buffer, ensure we can index it correctly

tinygrad/runtime/graph/hcq.py

@@ -58,8 +58,8 @@ class HCQGraph(MultiGraphRunner):
 
     # When profiling allocate 2 signals for each jit item to measure speed. The jth jit item have signals at 2*j and 2*j+1.
     # TODO: This logic might allocate a few extra signals...
-    self.prof_signals: list[HCQSignal] = []
-    self.prof_graph_deps: list[list[int]] = []
+    self.prof_signals: list[HCQSignal] = [self.devices[0].new_signal() for i in range(len(jit_cache) * 2)] if PROFILE else []
+    self.prog_graph_deps: list[list[int]] = []
     self.prof_graph_entries: list[ProfileGraphEntry] = []
 
     last_j: dict[HWQueue, int|None] = collections.defaultdict(lambda: None)
@@ -127,12 +127,12 @@ class HCQGraph(MultiGraphRunner):
         prof_ji_desc = ji.prg._prg.name if is_exec_prg else f"{ji.bufs[1].device} -> {ji.bufs[0].device}" # type: ignore
 
         self.prof_graph_entries.append(ProfileGraphEntry(enqueue_dev.device, prof_ji_desc, sig_st, j * 2 + 1, is_copy=not is_exec_prg))
-        self.prof_graph_deps.append([d - 1 for _, d in rdeps])
+        self.prog_graph_deps.append([d - 1 for _, d in rdeps])
 
       last_j[enqueue_queue] = j
 
     # Check which signals are used in the profile graph.
-    self.prof_signal_is_used = [any(ent.st_id == j or ent.en_id == j for ent in self.prof_graph_entries) for j in range(len(jit_cache) * 2)]
+    self.prof_signal_is_used = [any(ent.st_id == j or ent.en_id == j for ent in self.prof_graph_entries) for j in range(len(self.prof_signals))]
 
     # Build hardware queues.
     self.copy_to_devs: dict[HCQCompiled, set[HCQCompiled]] = {dev: set() for dev in self.devices}
@@ -149,9 +149,6 @@ class HCQGraph(MultiGraphRunner):
     for j,ji in enumerate(jit_cache):
       enqueue_dev, enqueue_queue, sync_signals, deps, signal, signal_val = self.ji_schedule[j]
 
-      # Lazy allocate signals
-      if PROFILE: self.prof_signals += [enqueue_dev.new_signal(value=0) for _ in range(2)]
-
       for sig, val in sync_signals + deps: enqueue_queue.wait(sig, val)
 
       # Encode waits and start profile timestamp (if needed).
@@ -216,7 +213,7 @@ class HCQGraph(MultiGraphRunner):
 
   def collect_timestamps(self):
     # NOTE: Append to any device is fine...
-    self.devices[0].profile_events += [ProfileGraphEvent(self.prof_graph_entries, self.prof_graph_deps, [s.timestamp for s in self.prof_signals])]
+    self.devices[0].profile_events += [ProfileGraphEvent(self.prof_graph_entries, self.prog_graph_deps, [s.timestamp for s in self.prof_signals])]
 
   def dev_name(self, dev) -> str: return dev.device.replace(":", "_")
 

tinygrad/schedule/grouper.py

@@ -3,8 +3,7 @@ from tinygrad.helpers import all_int, prod, unwrap, dedup, DONT_REALIZE_EXPAND,
 from tinygrad.shape.shapetracker import ShapeTracker
 
 ALWAYS_CONTIGUOUS = {Ops.CONTIGUOUS, Ops.ASSIGN, Ops.COPY, Ops.BUFFER, Ops.BUFFER_VIEW,
-                     Ops.CONST, Ops.BIND, Ops.DEVICE, Ops.MSELECT, Ops.MSTACK, Ops.DEFINE_GLOBAL,
-                     Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.LOAD}
+                     Ops.CONST, Ops.BIND, Ops.DEVICE, Ops.MSELECT, Ops.MSTACK, Ops.DEFINE_GLOBAL}
 
 # **** Grouper decides which of the UOps realize
 
@@ -26,8 +25,8 @@ def realize_before_view(ctx:dict[UOp, None], view:UOp, tr:UOp) -> None:
 do_realize = PatternMatcher([
   # always realize SINK parents
   (UPat(Ops.SINK, name="s"), lambda ctx,s: ctx.update((x.base, None) for x in s.src if x.base.op not in ALWAYS_CONTIGUOUS)),
-  # always realize ASSIGN/CONTIGUOUS/COPY/BUFFER_VIEW
-  (UPat({Ops.ASSIGN, Ops.CONTIGUOUS, Ops.COPY, Ops.BUFFER_VIEW}, name="tr"), realize),
+  # always realize ASSIGN/CONTIGUOUS/GroupOp.Meta
+  (UPat({Ops.ASSIGN, Ops.CONTIGUOUS, *GroupOp.Meta}, name="tr"), realize),
   # realize before expand or unsafe pad ops
   (UPat(Ops.VIEW, src=(UPat(GroupOp.All-ALWAYS_CONTIGUOUS, name="tr"),), name="view"), realize_before_view),
   # realize parents of COPY, MSELECT, MSTACK

tinygrad/schedule/kernelize.py

@@ -3,11 +3,11 @@ from tinygrad.uop.ops import UOp, Ops, GroupOp, PatternMatcher, UPat, graph_rewr
 from tinygrad.uop.ops import track_rewrites, _substitute
 from tinygrad.uop.spec import type_verify, tensor_uop_spec
 from tinygrad.uop.symbolic import symbolic_simple
-from tinygrad.helpers import Metadata, all_int, all_same, prod, dedup, unwrap, getenv, pluralize, DEBUG, SPLIT_REDUCEOP, flatten
+from tinygrad.helpers import Metadata, all_int, all_same, prod, dedup, unwrap, getenv, pluralize, FUSE_ARANGE, DEBUG, SPLIT_REDUCEOP
 from tinygrad.dtype import ImageDType
 from tinygrad.schedule.multi import multi_pm
 from tinygrad.schedule.grouper import group_realizes, ALWAYS_CONTIGUOUS
-from tinygrad.opt.swizzler import merge_views
+from tinygrad.opt.swizzler import merge_views, apply_swizzle, swizzle_reduceop
 
 # creation can recurse a lot
 import sys
@@ -121,7 +121,7 @@ def create_kernel(x:UOp, b:UOp|None=None):
   buffer = b.base if b.size == b.base.size else UOp(Ops.BUFFER_VIEW, b.dtype, (b.base,), (b.size, b.arg.views[0].offset))
   return buffer.assign(kernel).reshape(x.shape)
 
-DONT_PLACE_IN_KERNEL = {Ops.KERNEL, Ops.ASSIGN, Ops.BUFFER, Ops.MSELECT, Ops.MSTACK, Ops.MULTI, Ops.BIND}
+DONT_PLACE_IN_KERNEL = {Ops.KERNEL, Ops.ASSIGN, Ops.BUFFER, Ops.MSELECT, Ops.MSTACK, Ops.MULTI}
 def append_to_kernel(x:UOp):
   new_srcs: list[UOp] = []
   metadata = x.arg.metadata
@@ -149,34 +149,18 @@ create_kernels = PatternMatcher([
 
 # **** fix kernel AST
 
-def unbind_view(x:UOp):
-  if any(x.op is Ops.BIND for x in x.arg.vars()): return x.replace(arg=x.arg.unbind()[0])
-  return None
-
-replace_buffers = PatternMatcher([
-  # replace ASSIGN with the target BUFFER
-  (UPat(Ops.ASSIGN, src=(UPat((Ops.BUFFER, Ops.LOAD)), UPat(Ops.KERNEL)), name="assign", allow_any_len=True), lambda assign: assign.src[0]),
-  # HACK: select the 0 branch of MSTACK (the device is wrong after this, is that okay?)
-  (UPat(Ops.MSTACK, name="x"), lambda x: x.src[0]),
+early_buffer_ops = PatternMatcher([
   # LOAD
-  (UPat(Ops.BUFFER, name="x"), lambda ctx,x: UOp(Ops.DEFINE_GLOBAL, x.dtype.ptr(x.size), (), ctx.index(x)).load()),
+  (UPat(Ops.BUFFER, name="x"), lambda ctx,x: UOp(Ops.DEFINE_GLOBAL, x.dtype.ptr(x.size), (), ctx.index(x), tag=1)),
   # no SINK for meta ops
   (UPat(Ops.SINK, src=(UPat(Ops.CONTIGUOUS, src=(UPat(GroupOp.Meta, name="x"),),))), lambda x:x),
-  # STORE (except for meta ops)
-  (UPat(Ops.SINK, src=UPat(GroupOp.All-{Ops.STORE}), name="sink"), lambda ctx,sink:
-   UOp.sink(*[UOp.store(UOp(Ops.DEFINE_GLOBAL, (s:=x.base).dtype.ptr(ctx[i].size), (), i).view(s.st), s) for i,x in enumerate(sink.src)],
-            arg=sink.arg)),
-  # remove CONTIGUOUS/DEVICE from kernel AST
-  (UPat((Ops.CONTIGUOUS, Ops.MSELECT), src=(UPat.var("x"),)), lambda x: x),
-  (UPat(Ops.VIEW, src=(UPat(Ops.DEVICE),), name="view"), lambda view: view.replace(src=())),
-  # passthrough ASSIGN (but let MSTACK process first)
-  (UPat(Ops.ASSIGN, src=(UPat(GroupOp.All-{Ops.MSTACK}), UPat()), name="x"), lambda x: x.src[1]),
-  # remove any BINDs from VIEWS
-  (UPat(Ops.VIEW, src=(UPat(), UPat((Ops.BIND, Ops.DEFINE_VAR))), allow_any_len=True, name="x"), lambda x: x.replace(src=x.src[0:1])),
-  # remove any BINDs from DEFINE_VARs
-  (UPat(Ops.BIND, name="x"), lambda x: x.src[0]),
-  # remove BINDs from ShapeTrackers
-  (UPat(Ops.VIEW, name="x"), unbind_view),
+])
+
+replace_globals = PatternMatcher([
+  # replace ASSIGN with the target BUFFER
+  (UPat(Ops.ASSIGN, src=(UPat(Ops.BUFFER), UPat(Ops.KERNEL)), name="assign", allow_any_len=True), lambda assign: assign.src[0]),
+  # HACK: select the 0 branch of MSTACK (the device is wrong after this, is that okay?)
+  (UPat(Ops.MSTACK, name="x"), lambda x: x.src[0]),
 ])
 
 def fix_kernel_ast(k:UOp) -> UOp|None:
@@ -184,22 +168,18 @@ def fix_kernel_ast(k:UOp) -> UOp|None:
   # replace buffer with define_global + add load/store last
   bufs = []
   for s in k.src:
-    if s.op is Ops.BIND: continue
     s = s.buf_uop
     # traverse back through MSELECT and MSTACK. HACK: 0 branch of MSTACK only
     while s.op in {Ops.MSELECT, Ops.MSTACK}: s = s.src[0]
     bufs.append(s)
   # replace global memory ops with the BUFFER they write to
-  # NOTE: merge_views is needed to unbind the reshapes
-  ast = graph_rewrite(k.arg.ast, merge_views+replace_buffers, bufs, bottom_up=True, name="replace buffers")
-  if ast.op is Ops.SINK and not all_same([x.device for x in k.src if x.op is not Ops.BIND]):
+  ast = graph_rewrite(k.arg.ast, replace_globals, bottom_up=True, name="replace globals")
+  ast = graph_rewrite(ast, early_buffer_ops, bufs, bottom_up=True, name="replace buffer early")
+  if ast.op is Ops.SINK and not all_same([x.device for x in k.src]):
     raise RuntimeError(f"all buffers must be on the same device: {tuple(b.buf_uop.buffer for b in k.src)}")
   return k.replace(arg=Kernel(ast, k.arg.metadata))
 
-create_ast = PatternMatcher([
-  (UPat(Ops.KERNEL, name="k"), fix_kernel_ast),
-  (UPat(Ops.DEFINE_VAR, src=(UPat(),), allow_any_len=True, name="x"), lambda x: x.replace(src=())),
-])
+create_ast = PatternMatcher([(UPat(Ops.KERNEL, name="k"), fix_kernel_ast),])
 
 # ** add metadata of KERNEL outputs
 
@@ -209,8 +189,77 @@ def append_metadata(root:UOp, k:UOp):
 
 replace_metadata = PatternMatcher([(UPat(Ops.ASSIGN, src=(UPat(), UPat(Ops.KERNEL, name="k")), name="root", allow_any_len=True), append_metadata),])
 
+pm_fuse = PatternMatcher([
+  # FUSE on CONTIGUOUS removes FUSE
+  (UPat(Ops.CONTIGUOUS, name="c").fuse(), lambda c: c),
+
+  # FUSE triggers swizzle on reduceop
+  (UPat(Ops.VIEW, src=(UPat(Ops.REDUCE_AXIS, src=(UPat.var("src"),), name="r").or_casted(),), name="view").fuse(),
+   lambda r,src,view: ret.cast(view.dtype) if (ret:=swizzle_reduceop(r, src, view, fuse=True)) is not None else None),
+
+  # FUSE on reduce (without view) adds fuse marker to grouper
+  (UPat(Ops.REDUCE_AXIS, name="r").fuse(),
+   lambda r: r.replace(src=(r.src[0].fuse(),), arg=r.arg+(True,)) if len(r.arg) == 2 else None),
+
+  # remove FUSE and insert CONTIGUOUS if it's an unsafe pad
+  (UPat(Ops.VIEW, src=(UPat(GroupOp.UnsafePad, name="alu"),), name="view").fuse(),
+   lambda alu, view: alu.contiguous().view(view.st) if any(v.mask is not None for v in view.st.views) else None),
+
+  # FUSE elementwise.
+  (UPat(Ops.VIEW, src=(UPat({*GroupOp.ALU, Ops.CAST}, name="alu"),), name="view").fuse(),
+   lambda alu, view: alu.replace(src=tuple(apply_swizzle(x.view(view.arg)).fuse() for x in alu.src))),
+
+  # push FUSE through to srcs
+  (UPat(Ops.FUSE, name="x"), lambda x: x.src[0].replace(src=tuple(y.fuse() for y in x.src[0].src))),
+])
+
+def do_fusion(x:UOp):
+  found_contiguous = {}
+  def gate_contiguous(x):
+    if is_contiguous:=(x.op is Ops.CONTIGUOUS): found_contiguous[x] = x.replace(src=(UOp(Ops.VIEW, arg=x.st), UOp.unique()))
+    return not is_contiguous
+  x.toposort(gate=gate_contiguous)
+  del gate_contiguous
+  return graph_rewrite(x.substitute(found_contiguous), pm_fuse, name="local fusion").substitute({v:k for k,v in found_contiguous.items()})
+
+def fuse_arange(root:UOp):
+  # skip if root is arange
+  if not FUSE_ARANGE or root.src[0].base.op is Ops.CONST: return None
+  # gather all local aranges (including any fused ones)
+  local_arange: list[UOp] = []
+  def gate_reduce(u):
+    if u.op is Ops.REDUCE_AXIS and u.src[0].base.op is Ops.CONST: local_arange.append(u)
+    return u.op not in {*ALWAYS_CONTIGUOUS, Ops.REDUCE_AXIS} or u is root
+  toposort = root.toposort(gate=gate_reduce)
+  if not local_arange: return None
+  # fuse the nearest expand child of arange
+  local_children: dict[UOp, list[UOp]] = {}
+  for u in toposort:
+    for s in u.src: local_children.setdefault(s, []).append(u)
+  fuse_rep: dict[UOp, UOp] = {}
+  # skip if root depends on aranges with different ndims. This can be improved
+  if any(len(set(dims)) > 1 for dims in zip(*[r.src[0].shape for r in local_arange])): return
+  for r in local_arange:
+    # skip if already fused
+    if len(r.arg) > 2: continue
+    q = list(local_children[r])
+    while q:
+      u = q.pop()
+      if not (curr_children:=local_children.get(u, [])): continue
+      for child in curr_children:
+        other_paths = {s for s in child.toposort() if s.op in {Ops.REDUCE_AXIS, Ops.BUFFER} and s not in {root, r}}
+        fuse_rep[child] = child.replace(src=tuple(s.fuse() if s is u else s for s in child.src))
+        if other_paths: break
+      else: q.extend(curr_children)
+  return root.substitute(fuse_rep, name="fuse_arange") if fuse_rep else None
+
+do_fuse = PatternMatcher([
+  (UPat(Ops.FUSE, name="x"), do_fusion),
+  (UPat(Ops.REDUCE_AXIS, name="root"), fuse_arange),
+])
+
 add_contiguous = PatternMatcher([(UPat(GroupOp.All-{Ops.CONTIGUOUS, Ops.ASSIGN}, name="x"),
-                                lambda ctx,x: x.replace(tag=1).contiguous(tag=3 if x in ctx[1] else 2) if x in ctx[0] and x.tag is None else None)])
+                                lambda ctx,x: x.replace(tag=1).contiguous() if x in ctx and x.tag is None else None)])
 
 # TODO: get this from the device through GrouperOpts
 DEVICE_MAX_BUFS = {"METAL":32, "WEBGPU":8}
@@ -229,21 +278,6 @@ def limit_bufs(root:UOp):
   if len(bufs)>=MAX_BUFS-1:
     return root.replace(src=tuple(s if s.base in bufs else s.replace(tag=1).contiguous() for s in root.src))
 
-def view_add_srcs(x:UOp):
-  if len(avars:=x.arg.vars()) and len(x.src) == 1:
-    return x.replace(src=x.src+tuple(avars))
-  return None
-
-new_fusion = PatternMatcher([
-  # FUSE removes CONTIGUOUS tag=2, dies to CONTIGUOUS w/o tag,
-  (UPat(Ops.FUSE, src=(UPat(Ops.CONTIGUOUS, name="c"),)), lambda c: c.src[0].replace(tag=None).fuse() if c.tag == 2 else c),
-  (UPat(Ops.FUSE, src=(UPat(name="s"),)), lambda s: s.replace(src=tuple([y.fuse() for y in s.src]))),
-  # remove CONTIGUOUS if there's no BUFFER upsteam
-  (UPat(Ops.CONTIGUOUS, name="c"),
-   lambda c: None if c.tag != 2 or c.src[0].op is Ops.COPY or
-     any(x.op in GroupOp.UnsafePad.union({Ops.BUFFER}) for x in c.toposort()) else c.src[0].replace(tag=None)),
-])
-
 finalize_contiguous = PatternMatcher([
   # if an op takes more than one input, check combined LOADs don't exceed device limits
   (UPat(set.union(GroupOp.Binary, GroupOp.Ternary), name="root"), limit_bufs),
@@ -251,8 +285,6 @@ finalize_contiguous = PatternMatcher([
   (UPat(Ops.CONTIGUOUS, src=(UPat(Ops.CONTIGUOUS),), name="x"), lambda x: x.src[0]),
   # simplify views
   (UPat(Ops.VIEW, src=(UPat.var('x')), name="v"), lambda x,v: x.view(new_st) if (new_st:=v.arg.simplify()) != v.arg else None),
-  # vars to views srcs
-  (UPat(Ops.VIEW, name="x"), view_add_srcs),
 ])
 
 remove_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])
@@ -270,17 +302,14 @@ def get_kernelize_map(sink:UOp) -> dict[UOp, UOp]:
   """
 
   # multi + merge_views + simplify
-  tensor_map = graph_rewrite_map(sink, multi_pm+merge_views+sym+replace_contiguous, ctx={}, name="merge_views")
+  tensor_map = graph_rewrite_map(sink, multi_pm+do_fuse+merge_views+sym+replace_contiguous, ctx={}, name="merge_views")
 
   # display the cleaned up tensor graph
   if getenv("VIZ"): graph_rewrite(tensor_map[sink], PatternMatcher([]), name="View Tensor Graph")
 
   # insert contiguous in places determined by the realize map
-  forced_realize = flatten([x.base.src if x.base.op is Ops.MSTACK else [x.base] for x in tensor_map[sink].src])
   realize_map = group_realizes(tensor_map[sink])
-  tensor_map = graph_rewrite_map(tensor_map[sink], add_contiguous, ctx=(realize_map, forced_realize),
-                                 bottom_up=True, input_map=tensor_map, name="add_contiguous")
-  tensor_map = graph_rewrite_map(tensor_map[sink], new_fusion, input_map=tensor_map, name="new_fusion")
+  tensor_map = graph_rewrite_map(tensor_map[sink], add_contiguous, ctx=realize_map, bottom_up=True, input_map=tensor_map, name="add_contiguous")
   tensor_map = graph_rewrite_map(tensor_map[sink], finalize_contiguous+remove_tags, input_map=tensor_map, name="finalize_contiguous")
 
   # group into kernels (this is context-free)

tinygrad/tensor.py

@@ -3507,16 +3507,15 @@ class Tensor(MathTrait):
     """
     return self * self.softplus().tanh()
 
-  def softplus(self, beta=1.0, threshold=20.0) -> Tensor:
+  def softplus(self, beta=1) -> Tensor:
     """
     Applies the Softplus function element-wise.
-    For numerical stability, the implementation folds into identity function when `self * beta > threshold`.
 
     ```python exec="true" source="above" session="tensor" result="python"
     print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).softplus().numpy())
     ```
     """
-    return (self * beta > threshold).where(self, (1/beta) * (1 + (self*beta).exp()).log())
+    return (1/beta) * (1 + (self*beta).exp()).log()
 
   def softsign(self) -> Tensor:
     """

tinygrad/uop/ops.py

@@ -154,8 +154,8 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
       sz = cast(PtrDType, self.dtype).size
       return ShapeTracker.from_shape((sz,)) if sz > 0 else None
 
-    # hack for PTX, CASTing the ptr loses the shape
-    if self.op is Ops.CAST and self.src[0].op is Ops.DEFINE_GLOBAL: return None
+    # hack for PTX, CASTing the ptr loses the shape. even worse hack with tag
+    if self.op is Ops.CAST and self.src[0].op is Ops.DEFINE_GLOBAL and self.src[0].tag is None: return None
 
     # otherwise we get the shape from sources
     if not (src_sts := [x.st for x in self.src if x.st is not None]): return None
@@ -275,7 +275,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     ret = UOp(Ops.REDUCE_AXIS, self.dtype, (ret,), (op, new_axis))
     return ret.reshape(tuple([x if i not in axis else 1 for i,x in enumerate(self.shape)]))
   def reduce(self, *src:UOp, **kwargs): return UOp(Ops.REDUCE, kwargs.pop('dtype', self.dtype), src=(self,)+src, **kwargs)
-  def contiguous(self, **kwargs): return self.alu(Ops.CONTIGUOUS, **kwargs)
+  def contiguous(self): return self.alu(Ops.CONTIGUOUS)
   def contiguous_backward(self): return self.alu(Ops.CONTIGUOUS_BACKWARD)
   def fuse(self): return self.alu(Ops.FUSE)
   def allreduce(self, op, device:str|tuple[str, ...]|UOp):

tinygrad/uop/spec.py

@@ -1,5 +1,5 @@
 from typing import cast, Callable
-from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, print_uops, python_alu, graph_rewrite
+from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, print_uops, python_alu, graph_rewrite, resolve
 from tinygrad.dtype import DType, ImageDType, dtypes, PtrDType, AddrSpace
 from tinygrad.helpers import all_same, prod, DEBUG, ContextVar, Context
 from tinygrad.shape.shapetracker import ShapeTracker
@@ -52,7 +52,7 @@ buffer_spec = PatternMatcher([
 
 assign_spec = PatternMatcher([
   # KERNEL can attach to an ASSIGN to describe the compute required to realize a BUFFER
-  (UPat(Ops.KERNEL, src=UPat((Ops.BUFFER, Ops.BUFFER_VIEW, Ops.ASSIGN, Ops.MSELECT, Ops.MSTACK, Ops.BIND))), lambda: True),
+  (UPat(Ops.KERNEL, src=UPat((Ops.BUFFER, Ops.BUFFER_VIEW, Ops.ASSIGN, Ops.MSELECT, Ops.MSTACK))), lambda: True),
 
   # ASSIGN has a target and a value. It can also optionally depend on other assigns
   (UPat(Ops.ASSIGN, name="x"), lambda x: len(x.src) >= 2 and all(s.op is Ops.ASSIGN for s in x.src[2:])),
@@ -207,7 +207,16 @@ spec = PatternMatcher([
 
 # *** this is the UOp AST spec ***
 
+def verify_sink_dims(sink:UOp):
+  if not all_same([s.shape for s in sink.src]): return False
+  for dims in zip(*[x.shape for x in sink.toposort() if x.op is Ops.VIEW]):
+    if len(n_dims:={s for s in dims if resolve(s!=1)}) > 1:
+      print(f"# INVALID KERNEL DIMS: can only have 1 or n in each dimension: {n_dims}")
+      return False
+
 ast_spec = PatternMatcher([
+  # shapes must have either 1 or n in each dimension
+  (UPat(Ops.SINK, src=UPat(Ops.STORE), name="sink"), verify_sink_dims),
   # VIEW can only exist in the edges
   (UPat(Ops.VIEW, src=(UPat((Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL),))), lambda: True),
   (UPat(Ops.VIEW, name="view"), lambda view: len(view.src) == 0),

tinygrad/viz/index.html

@@ -55,11 +55,6 @@
     opacity: 0.4;
     pointer-events: none;
   }
-  label {
-    display: inline-flex;
-    align-items: center;
-    gap: 4px;
-  }
   .graph svg {
     width: 100%;
     height: 100%;

tinygrad/viz/js/index.js

@@ -109,7 +109,7 @@ function formatTime(ts, dur=ts) {
 }
 const formatUnit = (d, unit="") => d3.format(".3~s")(d)+unit;
 
-const colorScheme = {TINY:["#1b5745", "#354f52", "#354f52", "#1d2e62", "#63b0cd"],
+const colorScheme = {TINY:["#1b5745", "#354f52", "#354f52", "#46acc2", "#1d2e62", "#63b0cd"],
   DEFAULT:["#2b2e39", "#2c2f3a", "#31343f", "#323544", "#2d303a", "#2e313c", "#343746", "#353847", "#3c4050", "#404459", "#444862", "#4a4e65"],
   BUFFER:["#3A57B7","#5066C1","#6277CD","#7488D8","#8A9BE3","#A3B4F2"],
   CATEGORICAL:["#ff8080", "#F4A261", "#C8F9D4", "#8D99AE", "#F4A261", "#ffffa2", "#ffffc0", "#87CEEB"],}
@@ -221,13 +221,7 @@ async function renderProfiler() {
         ctx.closePath();
         ctx.fill();
         // NOTE: y coordinates are in reverse order
-        for (let i = 0; i < x.length - 1; i++) {
-          let tooltipText = e.arg.tooltipText;
-          if (yscale != null && ((yaxisVal=yscale.invert(e.y1[i]))>0)) {
-            tooltipText += `\nTotal: ${formatUnit(yaxisVal, data.axes.y.fmt)}`;
-          }
-          rectLst.push({ x0:x[i], x1:x[i+1], y0:e.y1[i], y1:e.y0[i], arg:{...e.arg, tooltipText} });
-        }
+        for (let i = 0; i < x.length - 1; i++) rectLst.push({ x0:x[i], x1:x[i+1], y0:e.y1[i], y1:e.y0[i], arg:e.arg });
         continue;
       }
       // contiguous rect