move that

.github/workflows/benchmark.yml

@@ -749,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
@@ -870,28 +871,28 @@ jobs:
       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
+    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)

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/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

@@ -3,6 +3,7 @@ import unittest
 from dataclasses import replace
 
 from tinygrad.opt.kernel import Opt, OptOps, KernelOptError, Kernel, AxisType
+from tinygrad.codegen import rewrites_for_views, apply_rewrites
 from tinygrad.codegen.gpudims import get_grouped_dims
 from tinygrad.uop.ops import UOp, Ops, GroupOp, KernelInfo
 from tinygrad.device import Device, Buffer, is_dtype_supported
@@ -22,7 +23,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 s[-1].ast, bufs
+  return apply_rewrites(s[-1].ast, rewrites_for_views), 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)
@@ -214,7 +215,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?"
@@ -225,7 +228,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"
 

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_softmax_fusion.py

@@ -111,13 +111,6 @@ class TestFuse(unittest.TestCase):
     with Context(NOOPT=1):
       self._test_fuse(Tensor.scaled_dot_product_attention, q, k, v, atol=1e-5)
 
-  @unittest.expectedFailure
-  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_uops_stats.py

@@ -173,7 +173,8 @@ 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)
 
@@ -188,30 +189,42 @@ class TestStatsOptimized(unittest.TestCase):
   # 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)

tinygrad/codegen/__init__.py

@@ -16,6 +16,7 @@ from tinygrad.codegen.devectorizer import load_store_folding, load_store_indexin
   ReduceContext, correct_load_store, pm_render
 from tinygrad.codegen.optional import get_late_rewrite_patterns
 from tinygrad.codegen.linearize import block_create, pm_blockend_merge, block_merge, pm_finalize, BlockContext
+from tinygrad.opt.swizzler import view_left, view_right, cleanup_pm
 from tinygrad.opt import pm_optimize
 
 @dataclass
@@ -29,6 +30,12 @@ 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="view left"),
+  RewriteStep(view_right, name="view right"),
+  RewriteStep(cleanup_pm, name="cleanup view"),
+]
+
 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"),
@@ -44,6 +51,9 @@ def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVEC
   # ** lowerer (rewrite_shapetracker_with_index) **
   ret: list[RewriteStep] = []
 
+  # this used to be in schedule
+  ret.extend(rewrites_for_views)
+
   # 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/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
@@ -73,7 +73,7 @@ class Kernel:
       self.sts.append(unwrap(x.src[0].st))
 
     # add a shapetracker to the end to track the full shape, with 0 strides so it can merge
-    full_shape = ast.full_shape
+    full_shape = self.ast.full_shape
     self.sts.append(ShapeTracker.from_shape(full_shape, (0,)*len(full_shape)))
 
     # parameters for optimization
@@ -451,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
@@ -504,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

@@ -1,9 +1,9 @@
 from tinygrad.uop.ops import UOp, Ops, GroupOp, PatternMatcher, UPat, graph_rewrite, resolve, sint
 from tinygrad.helpers import all_same, prod, unwrap, colored
+from tinygrad.dtype import dtypes
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View, strides_for_shape, get_contraction_with_reduce
 from tinygrad.schedule.grouper import ALWAYS_CONTIGUOUS
-from tinygrad.dtype import ImageDType, dtypes
 
 merge_views = PatternMatcher([
   # merge adjacent views
@@ -15,6 +15,9 @@ merge_views = PatternMatcher([
    lambda x,view: x if x.st is not None and x.op not in GroupOp.Defines and view.st.contiguous and view.shape == x.shape else None),
   (UPat(GroupOp.All-{Ops.DEFINE_GLOBAL}).view(name="view"),
    lambda view: view.const_like(0) if (mask:=view.st.views[-1].mask) is not None and any((x[1]-x[0]) == 0 for x in mask) else None),
+  # only unmaksed VIEW on CONST replaces the ShapeTracker
+  (UPat(Ops.VIEW, src=(UPat((Ops.CONST, Ops.DEFINE_VAR), name="x"),), name="view"),
+   lambda x,view: x.replace(src=(x.src[0].replace(arg=x.st+view.st),)) if all(v.mask is None for v in (x.st+view.st).views) else None),
 ])
 
 def reduce_push_add_ones(src:UOp, r:UOp, view:UOp):
@@ -41,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.
@@ -103,8 +100,6 @@ view_right = merge_views+PatternMatcher([
   # 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):
@@ -117,18 +112,14 @@ 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)),
-  # only unmaksed VIEW on CONST replaces the ShapeTracker
-  (UPat(Ops.VIEW, src=(UPat((Ops.CONST, Ops.DEFINE_VAR), name="x"),), name="view"),
-   lambda x,view: x.replace(src=(UOp(Ops.VIEW, arg=x.st+view.st),)) if all(v.mask is None for v in (x.st+view.st).views) else None),
+cleanup_pm = PatternMatcher([
+  # add VIEW to any DEFINE_GLOBAL that somehow lost its view
+  (UPat(Ops.STORE, src=(UPat(Ops.DEFINE_GLOBAL, name="d"),), name="x", allow_any_len=True), lambda d,x: x.replace(src=(d.view(d.st),)+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)),
-  # 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),
+  # VIEW on SINK is SINK
+  (UPat(Ops.VIEW, name="v").sink(), lambda v: v.src[0].sink()),
   # if this kernel also assigns to the loaded buffer, ensure we can index it correctly
   (UPat(Ops.LOAD, src=(UPat.var("glbl").view(name="view"),)), check_load_st),
 ])

tinygrad/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,7 +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.LOAD}
+                     Ops.CONST, Ops.BIND, Ops.DEVICE, Ops.MSELECT, Ops.MSTACK, Ops.DEFINE_GLOBAL}
 
 # **** Grouper decides which of the UOps realize
 

tinygrad/schedule/kernelize.py

@@ -7,7 +7,7 @@ from tinygrad.helpers import Metadata, all_int, all_same, prod, dedup, unwrap, g
 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, view_left, view_right, fix_kernel_ops, apply_swizzle, swizzle_reduceop
+from tinygrad.opt.swizzler import merge_views, apply_swizzle, swizzle_reduceop
 
 # creation can recurse a lot
 import sys
@@ -149,28 +149,36 @@ create_kernels = PatternMatcher([
 
 # **** fix kernel AST
 
-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]),
+add_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()),
-  # no SINK for meta ops
-  (UPat(Ops.SINK, src=(UPat(Ops.CONTIGUOUS, src=(UPat(GroupOp.Meta, name="x"),),))), lambda x:x),
+  (UPat(Ops.BUFFER, name="x"), lambda ctx,x: UOp.load(UOp(Ops.DEFINE_GLOBAL, x.dtype.ptr(x.size), (), ctx.index(x)).view(x.st),)),
   # STORE (except for meta ops)
+  (UPat(Ops.SINK, src=(UPat(Ops.CONTIGUOUS, src=(UPat(GroupOp.Meta, name="x"),),))), lambda x:x),
   (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)),
+   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)])),
+  # passthrough ASSIGN
+  (UPat(Ops.ASSIGN, name="x"), lambda x: x.src[1]),
+])
+
+fix_kernel_ops = PatternMatcher([
   # 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]),
+  # 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),
+])
+
+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:
   if k.arg.ast.op in GroupOp.Meta or all(s.op is Ops.STORE for s in k.arg.ast.src): return None
+  # replace global memory ops with the BUFFER they write to
+  ast = graph_rewrite(k.arg.ast, replace_globals, bottom_up=True, name="replace globals")
   # replace buffer with define_global + add load/store last
   bufs = []
   for s in k.src:
@@ -178,14 +186,9 @@ def fix_kernel_ast(k:UOp) -> UOp|None:
     # 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
-  ast = graph_rewrite(k.arg.ast, replace_buffers, bufs, bottom_up=True, name="replace buffers")
+  ast = graph_rewrite(ast, merge_views+add_buffer_ops+fix_kernel_ops, bufs, bottom_up=True, name="replace buffer")
   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)}")
-  # TODO: move these to codegen
-  ast = graph_rewrite(ast, view_left, name="Main View Left")
-  ast = graph_rewrite(ast, view_right, name="Main View Right")
-  ast = graph_rewrite(ast, view_left+fix_kernel_ops, bottom_up=True, name="Finalize Kernel")
   return k.replace(arg=Kernel(ast, k.arg.metadata))
 
 create_ast = PatternMatcher([(UPat(Ops.KERNEL, name="k"), fix_kernel_ast),])

tinygrad/uop/ops.py

@@ -255,8 +255,11 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     if isinstance(b, UOp): return b.unbind()[0] if b.op is Ops.BIND else b
     if isinstance(b, tuple) and all_same(b): b = b[0]  # doesn't have to be a VCONST if they are all the same
     ret = UOp(Ops.VCONST if isinstance(b, tuple) else Ops.CONST, dtype, arg=dtypes.as_const(b, dtype))
-    if device is not None: ret = ret.replace(src=(UOp(Ops.DEVICE, arg=device),))
-    if shape is not None: ret = ret.reshape((1,)*len(shape)).expand(shape)
+    if shape is not None:
+      from tinygrad.shape.shapetracker import ShapeTracker
+      ret = ret.replace(src=(UOp(Ops.VIEW, dtypes.void, (), ShapeTracker.from_shape(shape, (0,)*len(shape))),))
+    if device is not None:
+      ret = ret.replace(src=(UOp(Ops.DEVICE, arg=device).view(unwrap(ret.st)),))
     return ret
   @staticmethod
   def range(dtype:DType, end:sint, idx:int): return UOp(Ops.RANGE, dtype=dtype, src=(sint_to_uop(end),), arg=idx)
@@ -437,7 +440,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     all_vars = set([x for x in self.toposort() if x.op is Ops.DEFINE_VAR])
     return bound_vars.union(set([x for x in all_vars if x not in bound_var_base]))
   def variables(self) -> list[Variable]:
-    st_vars: list[set[Variable]] = [x.arg.vars() for x in self.toposort() if x.op is Ops.VIEW]
+    st_vars: list[set[Variable]] = [x.st_arg.vars() for x in self.toposort() if x.op in GroupOp.Buffer]
     return sorted(set.union(*st_vars, set([x.unbind()[0] if x.op is not Ops.DEFINE_VAR else x for x in self.vars()])), key=lambda v: v.arg)
 
   # *** uop symbolic stuff ***

tinygrad/viz/js/index.js

@@ -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