experiments with multi being range

fix assert in test_schedule (#12745 )
* fix assert in test_schedule updated kernel counts and some old tests * fix
2025-10-17 14:11:55 +08:00 · 2025-10-16 15:39:50 -04:00 · 2025-10-16 14:11:33 -04:00 · 2025-10-16 12:43:27 -04:00 · 2025-10-16 11:40:36 -04:00 · 2025-10-16 11:17:47 -04:00
41 changed files with 568 additions and 467 deletions
@@ -633,7 +633,7 @@ jobs:
      run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=12 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
    - name: openpilot compile3 Space Lab policy + vision
      run: |
-        PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/22aec22a10ce09384d4a4af2a0bbff08d54af7e0c888503508f356fae4ff0e29
+        PYTHONPATH="." ASSERT_MIN_STEP_TIME=5 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/22aec22a10ce09384d4a4af2a0bbff08d54af7e0c888503508f356fae4ff0e29
        PYTHONPATH="." ASSERT_MIN_STEP_TIME=26 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/c824f68646a3b94f117f01c70dc8316fb466e05fbd42ccdba440b8a8dc86914b
    - name: benchmark MobileNetV2 on DSP
      run: |
@@ -642,7 +642,7 @@ jobs:
        ln -s /data/home/tiny/tinygrad/testsig-*.so .
        PYTHONPATH=. CC=clang-19 CPU=1 CPU_LLVM=0 QUANT=1 CNT=0 python3 examples/test_onnx_imagenet.py https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx /tmp/model.quant.onnx
        # benchmark on DSP with NOOPT=1, the devectorizer has issues
-        PYTHONPATH=. CC=clang-19 DSP=1 DONT_REALIZE_EXPAND=1 NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
+        PYTHONPATH=. CC=clang-19 DSP=1 NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
    - name: Run process replay tests
      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
    - uses: actions/upload-artifact@v4
@@ -279,9 +279,9 @@ generate_llvm() {
    --clang-args="$(llvm-config-14 --cflags)" \
    -o "$BASE/llvm.py"

-  sed -i "s\import ctypes\import ctypes, tinygrad.runtime.support.llvm as llvm_support, tinygrad.helpers as helpers\g" "$BASE/llvm.py"
+  sed -i "s\import ctypes\import ctypes, tinygrad.runtime.support.llvm as llvm_support\g" "$BASE/llvm.py"
  sed -i "s\FIXME_STUB\llvm\g" "$BASE/llvm.py"
-  sed -i "s\FunctionFactoryStub()\ctypes.CDLL(llvm_support.LLVM_PATH, ctypes.RTLD_GLOBAL if helpers.OSX else ctypes.DEFAULT_MODE)\g" "$BASE/llvm.py"
+  sed -i "s\FunctionFactoryStub()\ctypes.CDLL(llvm_support.LLVM_PATH)\g" "$BASE/llvm.py"

  fixup "$BASE/llvm.py"
 }
@@ -232,7 +232,7 @@ if __name__ == "__main__":
  gpt2 = GPT2.build_gguf(args.model_size) if args.model_size.startswith("gpt2_gguf_") else GPT2.build(args.model_size)

  if args.benchmark != -1:
-    gpt2.model(Tensor.rand(args.batch_size, args.benchmark), Variable("a", 0, MAX_CONTEXT).bind(0)).realize()
+    gpt2.model(Tensor.randint(args.batch_size, args.benchmark), Variable("a", 0, MAX_CONTEXT).bind(0)).realize()
  else:
    texts = gpt2.generate(args.prompt, args.count, args.temperature, timing=args.timing, batch_size=args.batch_size)
    if not args.noshow:
@@ -269,7 +269,8 @@ if __name__ == "__main__":

  # load in weights
  with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
-    load_state_dict(model, torch_load(fetch('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', 'sd-v1-4.ckpt'))['state_dict'], verbose=False, strict=False, realize=False)
+    model_bin = fetch('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', 'sd-v1-4.ckpt')
+    load_state_dict(model, torch_load(model_bin)['state_dict'], verbose=False, strict=False, realize=False)

    if args.fp16:
      for k,v in get_state_dict(model).items():
@@ -19,8 +19,8 @@ from tinygrad.helpers import fetch, getenv

 # QUANT=1 python3 examples/test_onnx_imagenet.py
 #   https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx
-# DONT_REALIZE_EXPAND=1 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
-# VIZ=1 DONT_REALIZE_EXPAND=1 python3 examples/benchmark_onnx.py /tmp/model.quant.onnx
+# python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
+# VIZ=1 python3 examples/benchmark_onnx.py /tmp/model.quant.onnx

 def imagenet_dataloader(cnt=0):
  input_mean = Tensor([0.485, 0.456, 0.406]).reshape(1, -1, 1, 1)
@@ -3,8 +3,21 @@ from tinygrad.tensor import _to_np_dtype
 from tinygrad.nn.onnx import OnnxRunner, OnnxValue
 import numpy as np
 import onnxruntime as ort
+ort_options = ort.SessionOptions()
+ort_options.log_severity_level = 3

 def get_example_inputs(graph_inputs:dict[str, OnnxValue], config={}):
+  """
+  Generate example input tensors based on the provided ONNX graph input specifications.
+
+  NOTE: This is not guaranteed to be reliable. It's a best-effort helper
+  that uses heuristics to guess input shapes and values.
+
+  Example:
+    from tinygrad.nn.onnx import OnnxRunner
+    from extra.onnx_helpers import get_example_inputs
+    inputs = get_example_inputs(OnnxRunner(model_path).graph_inputs)
+  """
  def _get_shape(onnx_shape: tuple[str|int]):
    shape = []
    for onnx_dim in onnx_shape:
@@ -44,11 +57,9 @@ def get_example_inputs(graph_inputs:dict[str, OnnxValue], config={}):
    ret.update({name:value})
  return ret

-def validate(onnx_file, inputs, rtol=1e-5, atol=1e-5):
+def _get_tinygrad_and_ort_np_outputs(onnx_file, inputs):
  run_onnx = OnnxRunner(onnx_file)

-  ort_options = ort.SessionOptions()
-  ort_options.log_severity_level = 3
  ort_sess = ort.InferenceSession(onnx_file, ort_options, ["CPUExecutionProvider"])
  np_inputs = {k:v.numpy() if isinstance(v, Tensor) else v for k,v in inputs.items()}
  out_names = list(run_onnx.graph_outputs)
@@ -56,9 +67,101 @@ def validate(onnx_file, inputs, rtol=1e-5, atol=1e-5):
  ort_out = dict(zip(out_names, out_values))

  tinygrad_out = run_onnx(inputs)
+  Tensor.realize(*(x for x in tinygrad_out.values() if x is not None))
+  tinygrad_out = {k:v.numpy() if v is not None else None for k,v in tinygrad_out.items()}
+  return tinygrad_out, ort_out
+
+def validate(onnx_file, inputs, rtol=1e-5, atol=1e-5):
+  """
+  Compares the final output tensors of an onnx model run in tinygrad and onnxruntime.
+  """
+  tinygrad_out, ort_out = _get_tinygrad_and_ort_np_outputs(onnx_file, inputs)

  assert tinygrad_out.keys() == ort_out.keys()
  for k in tinygrad_out.keys():
    tiny_v, onnx_v = tinygrad_out[k], ort_out[k]
    if tiny_v is None: assert onnx_v is None, f"{k}: {tiny_v=}, {onnx_v=}"
-    else: np.testing.assert_allclose(tiny_v.numpy(), onnx_v, rtol=rtol, atol=atol, err_msg=f"For tensor '{k}' in {tinygrad_out.keys()}")
+    else: np.testing.assert_allclose(tiny_v, onnx_v, rtol=rtol, atol=atol, err_msg=f"For tensor '{k}' in {tinygrad_out.keys()}")
+
+def validate_all_intermediates(onnx_file, inputs, rtol=1e-5, atol=1e-5):
+  """
+  Compares all intermediate node output of an onnx model run in tinygrad and onnxruntime.
+  """
+  report = generate_node_output_report(onnx_file, inputs)
+  for i, node in enumerate(report):
+    node_name = node["node"]
+    op = node["op"]
+    outputs = node["outputs"]
+    for output in outputs:
+      output_name = output["name"]
+      tinygrad_out = output["tinygrad"]
+      ort_out = output["onnxruntime"]
+      try:
+        if tinygrad_out is None: assert ort_out is None, f"None outputs are not equal {tinygrad_out=} {ort_out=}"
+        else: np.testing.assert_allclose(tinygrad_out, ort_out, rtol=rtol, atol=atol)
+        print(f"Validated {i}: {op=} {node_name=} {output_name=}")
+      except AssertionError as e:
+        print(f"FAILED {i}: {op=} {node_name=} {output_name=}")
+        print(str(e).strip() + "\n")
+
+def generate_node_output_report(onnx_file, inputs):
+  """
+  Build a report of all ONNX node outputs from tinygrad and onnxruntime
+
+  Returns:
+    A list of dictionaries, where each entry corresponds to one
+    node in the ONNX graph. The structure is as follows:
+    [
+      {
+        "node": str,  # The name of the ONNX node.
+        "op": str,    # The operation type of the ONNX node.
+        "outputs": [
+          {
+            "name": str,                       # The name of the output tensor.
+            "tinygrad": np.ndarray | None,     # The output value from tinygrad.
+            "onnxruntime": np.ndarray | None,  # The output value from onnxruntime.
+          },
+          ...
+        ]
+      },
+      ...
+    ]
+  """
+  import onnx_graphsurgeon as gs
+  import onnx
+  import tempfile
+
+  # rewrite the model to output all the node outputs
+  # `infer_shapes` here tries to fill the shapes and dtypes of intermediate values which graphsurgeon requires when assigning them as outputs
+  inferred_model = onnx.shape_inference.infer_shapes(onnx.load(onnx_file))
+  model = gs.import_onnx(inferred_model)
+  model_nodes = model.nodes
+  node_outputs = [n.outputs for n in model.nodes]
+  model.outputs = [
+      each_output for outputs in node_outputs for each_output in outputs
+      if not (each_output.dtype is None and each_output.shape is None)  # output with None dtype and None shape is likely a `None` value
+  ]
+  rewritten_model = gs.export_onnx(model)
+
+  # TODO: remove this once ORT supports 1.18.0
+  if getattr(rewritten_model, "ir_version", 0) > 10:
+    rewritten_model.ir_version = 10
+
+  with tempfile.NamedTemporaryFile(suffix=".onnx") as f:
+    onnx.save(rewritten_model, f.name)
+    rewritten_model_path = f.name
+    tinygrad_out, ort_out = _get_tinygrad_and_ort_np_outputs(rewritten_model_path, inputs)
+
+  report = []
+  for node in model_nodes:
+    outputs = []
+    for each_output in node.outputs:
+      if each_output.dtype is None and each_output.shape is None:
+        continue
+      name = each_output.name
+      tinygrad_output = tinygrad_out[name]
+      ort_output = ort_out[name]
+      outputs.append({"name": name, "tinygrad": tinygrad_output, "onnxruntime": ort_output})
+    report.append({"node": node.name, "op": node.op, "outputs": outputs})
+
+  return report
@@ -56,7 +56,7 @@ class TinyGPUViewModel: NSObject {
 	}
 #endif

-	private let dextIdentifier: String = Bundle.main.bundleIdentifier! + ".Driver"
+	private let dextIdentifier: String = "org.tinygrad.tinygpu.edriver"

 	public var dextLoadingState: String {
 		switch state {
@@ -12,8 +12,8 @@
 			<string>IOUserService</string>
 			<key>IOMatchCategory</key>
 			<string>TinyGPUDriver</string>
-			<key>IOPCIPrimaryMatch</key>
-			<string>0x70001002&amp;0xF000FFFF</string>
+			<key>IOPCIClassMatch</key>
+			<string>0x03000000</string>
 			<key>IOPCITunnelCompatible</key>
 			<true/>
 			<key>IOProviderClass</key>
@@ -87,7 +87,7 @@ kern_return_t TinyGPUDriver::Start_Impl(IOService* in_provider)
 		}
 		off = next;
 	}
-	ivars->pci->Reset(0);
+	ivars->pci->Reset(kIOPCIDeviceResetTypeHotReset);
 #endif

 	uint16_t commandRegister;
@@ -221,3 +221,39 @@ error:
 	}
 	return err;
 }
+
+kern_return_t TinyGPUDriver::CfgRead(uint32_t off, uint32_t size, uint32_t* outVal)
+{
+  if (!ivars->pci || !outVal) return kIOReturnNotReady;
+
+  if (size == 1) {
+	uint8_t v8 = 0;
+	ivars->pci->ConfigurationRead8(off, &v8);
+	*outVal = v8;
+  } else if (size == 2) {
+	uint16_t v16 = 0;
+	ivars->pci->ConfigurationRead16(off, &v16);
+	*outVal = v16;
+  } else if (size == 4) {
+	uint32_t v32 = 0;
+	ivars->pci->ConfigurationRead32(off, &v32);
+	*outVal = v32;
+  }
+  return 0;
+}
+
+kern_return_t TinyGPUDriver::CfgWrite(uint32_t off, uint32_t size, uint32_t val)
+{
+  if (!ivars->pci) return kIOReturnNotReady;
+  if (size == 1) ivars->pci->ConfigurationWrite8 (off, (uint8_t)val);
+  else if (size == 2) ivars->pci->ConfigurationWrite16(off, (uint16_t)val);
+  else if (size == 4) ivars->pci->ConfigurationWrite32(off, (uint32_t)val);
+  return 0;
+}
+
+kern_return_t TinyGPUDriver::ResetDevice()
+{
+	if (!ivars->pci) return kIOReturnNotReady;
+	ivars->pci->Reset(kIOPCIDeviceResetTypeFunctionReset);
+	return 0;
+}
@@ -6,7 +6,11 @@
 	<array>
 		<dict>
 			<key>IOPCIMatch</key>
-			<string>0x70001002&amp;0xF000FFFF</string>
+			<string>0x00001002&amp;0x0000FFFF</string>
+		</dict>
+		<dict>
+			<key>IOPCIMatch</key>
+			<string>0x000010de&amp;0x0000FFFF</string>
 		</dict>
 	</array>
 	<key>com.apple.developer.driverkit.allow-any-userclient-access</key>
@@ -28,6 +28,11 @@ public:

 	kern_return_t MapBar(uint32_t bar, IOMemoryDescriptor** memory) LOCALONLY;
 	kern_return_t CreateDMA(size_t size, TinyGPUCreateDMAResp* dmaDesc) LOCALONLY;
+
+	kern_return_t CfgRead(uint32_t off, uint32_t size, uint32_t* val) LOCALONLY;
+	kern_return_t CfgWrite(uint32_t off, uint32_t size, uint32_t val) LOCALONLY;
+	kern_return_t ResetDevice() LOCALONLY;
+	kern_return_t BarInfo() LOCALONLY;
 };

 #endif /* TinyGPUDriver_h */
@@ -62,8 +62,41 @@ kern_return_t TinyGPUDriverUserClient::Stop_Impl(IOService* in_provider)
 	return Stop(in_provider, SUPERDISPATCH);
 }

-kern_return_t TinyGPUDriverUserClient::ExternalMethod(uint64_t in_selector, IOUserClientMethodArguments* in_arguments, const IOUserClientMethodDispatch* in_dispatch, OSObject* in_target, void* in_reference)
+kern_return_t TinyGPUDriverUserClient::ExternalMethod(uint64_t selector, IOUserClientMethodArguments* args, const IOUserClientMethodDispatch* in_dispatch, OSObject* in_target, void* in_reference)
 {
+	kern_return_t err = 0;
+
+	os_log(OS_LOG_DEFAULT, "tinygpu: rpc (%llu) in:%d, out:%d", selector, args->scalarInputCount, args->scalarOutputCount);
+
+	if (selector == TinyGPURPC::ReadCfg) {
+		if (args->scalarInputCount != 2 or args->scalarOutputCount < 1) return kIOReturnBadArgument;
+
+		uint32_t off = uint32_t(args->scalarInput[0]);
+		uint32_t size = uint32_t(args->scalarInput[1]);
+
+		uint32_t val = 0;
+		err = ivars->provider->CfgRead(off, size, &val);
+		os_log(OS_LOG_DEFAULT, "tinygpu: read cfg off:%x sz:%d, val:%x", off, size, val);
+
+		if (!err) {
+			args->scalarOutput[0] = val;
+			args->scalarOutputCount = 1;
+		}
+		return err;
+	} else if (selector == TinyGPURPC::WriteCfg) {
+		if (args->scalarInputCount != 3) return kIOReturnBadArgument;
+
+		uint32_t off = uint32_t(args->scalarInput[0]);
+		uint32_t size = uint32_t(args->scalarInput[1]);
+		uint32_t val = uint32_t(args->scalarInput[2]);
+
+		os_log(OS_LOG_DEFAULT, "tinygpu: wr cfg off:%x sz:%d, val:%x", off, size, val);
+		return ivars->provider->CfgWrite(off, size, val);
+	} else if (selector == TinyGPURPC::Reset) {
+		os_log(OS_LOG_DEFAULT, "tinygpu: reset");
+		return ivars->provider->ResetDevice();
+	}
+
 	return kIOReturnUnsupported;
 }

@@ -3,6 +3,13 @@

 #include <DriverKit/IOUserClient.iig>

+enum TinyGPURPC
+{
+	ReadCfg,
+	WriteCfg,
+	Reset
+};
+
 class TinyGPUDriverUserClient : public IOUserClient
 {
 public:
@@ -1,4 +1,4 @@
-from tinygrad import Tensor, dtypes, Context, GlobalCounters
+from tinygrad import Tensor, dtypes, GlobalCounters
 dtypes.default_float = dtypes.float16
 from tinygrad.dtype import to_dtype
 from tinygrad.helpers import getenv
@@ -13,6 +13,5 @@ if __name__ == "__main__":

  # test single kernel softmax
  GlobalCounters.reset()
-  with Context(DONT_GROUP_REDUCES=1):
-    single_kernel_softmax(t, -1, acc_dtype).realize()
+  single_kernel_softmax(t, -1, acc_dtype).realize()

@@ -1,46 +0,0 @@
-# ruff: noqa: E501
-from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps
-from tinygrad.dtype import dtypes
-from tinygrad.engine.realize import CompiledRunner, get_program
-from tinygrad.codegen.opt.search import bufs_from_lin
-from tinygrad.uop.ops import UOp, Ops
-from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad.shape.view import View
-
-ast = UOp(Ops.SINK, dtypes.void, arg=None, src=(
-  UOp(Ops.STORE, dtypes.void, arg=None, src=(
-    UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(), arg=0, src=()),
-    UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(2, 1, 1280, 8, 8, 1, 1, 1), strides=(81920, 0, 64, 8, 1, 0, 0, 0), offset=0, mask=None, contiguous=True),)), src=()),
-    UOp(Ops.ADD, dtypes.half, arg=None, src=(
-      UOp(Ops.ADD, dtypes.half, arg=None, src=(
-        UOp(Ops.CAST, dtypes.half, arg=None, src=(
-          UOp(Ops.REDUCE_AXIS, dtypes.float, arg=(Ops.ADD, (5, 6, 7)), src=(
-            UOp(Ops.CAST, dtypes.float, arg=None, src=(
-              UOp(Ops.MUL, dtypes.half, arg=None, src=(
-                UOp(Ops.LOAD, dtypes.half, arg=None, src=(
-                  UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(), arg=1, src=()),
-                  UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 2, 1, 2560, 4, 10, 4, 10), strides=(0, 163840, 0, 64, 0, 8, 0, 1), offset=-9, mask=((0, 1), (0, 2), (0, 1), (0, 2560), (0, 4), (1, 9), (0, 4), (1, 9)), contiguous=False), View(shape=(2, 1, 1280, 8, 8, 2560, 3, 3), strides=(4096000, 0, 0, 40, 1, 1600, 440, 11), offset=0, mask=None, contiguous=False))), src=()),)),
-                UOp(Ops.LOAD, dtypes.half, arg=None, src=(
-                  UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(), arg=2, src=()),
-                  UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(2, 1, 1280, 8, 8, 2560, 3, 3), strides=(0, 0, 23040, 0, 0, 9, 3, 1), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),)),)),
-        UOp(Ops.LOAD, dtypes.half, arg=None, src=(
-          UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(), arg=3, src=()),
-          x17:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(2, 1, 1280, 8, 8, 1, 1, 1), strides=(0, 0, 1, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
-      UOp(Ops.LOAD, dtypes.half, arg=None, src=(
-        UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(), arg=4, src=()),
-         x17,)),)),)),))
-opts = [Opt(op=OptOps.UPCAST, axis=3, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UNROLL, axis=2, arg=0), Opt(op=OptOps.UNROLL, axis=1, arg=0), Opt(op=OptOps.LOCAL, axis=1, arg=8), Opt(op=OptOps.LOCAL, axis=2, arg=8), Opt(op=OptOps.LOCAL, axis=2, arg=2)]
-
-k = Kernel(ast)
-k.apply_opts(opts)
-bufs = bufs_from_lin(k)
-
-prg = CompiledRunner(get_program(k.ast, k.opts, k.applied_opts))
-
-for i in range(10):
-  speed = prg(bufs, var_vals={}, wait=True)
-  print(f"kernel time: {speed*1e3:.2f} ms")
-
-# on M1 Max
-# 11ms before block 9b0859d71780fef5cf3831e317f74e53f2483229
-# 15ms after block cbcc1c20eb09a1342f6581cfbb99632bade982a8
@@ -1,55 +0,0 @@
-# ruff: noqa: E501
-import unittest
-
-from tinygrad.uop.ops import UOp, Ops
-from .search import Opt, OptOps
-from tinygrad.dtype import dtypes
-from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad.shape.view import View
-from tinygrad.codegen.opt.kernel import Kernel
-
-from test.external.fuzz_linearizer import run_linearizer
-
-class TestTrainGpt2Kernel(unittest.TestCase):
-  def test_1(self):
-    # kernel 244
-    ast = UOp(Ops.SINK, dtypes.void, arg=None, src=(
-      UOp(Ops.STORE, dtypes.void, arg=None, src=(
-        UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(206045184), arg=0, src=()),
-        UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(4, 1024, 50304, 1), strides=(51511296, 50304, 1, 0), offset=0, mask=None, contiguous=True),)), src=()),
-        UOp(Ops.REDUCE_AXIS, dtypes.float, arg=(Ops.ADD, (3,)), src=(
-        UOp(Ops.MUL, dtypes.float, arg=None, src=(
-          UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-          UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(3145728), arg=1, src=()),
-          UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(4, 1024, 50304, 768), strides=(786432, 768, 0, 1), offset=0, mask=None, contiguous=False),)), src=()),)),
-          UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-          UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(38633472), arg=2, src=()),
-          UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(4, 1024, 50304, 768), strides=(0, 0, 768, 1), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),)),))
-
-    opts = [Opt(op=OptOps.LOCAL, axis=0, arg=16), Opt(op=OptOps.UPCAST, axis=1, arg=3), Opt(op=OptOps.LOCAL, axis=0, arg=2)]
-    kernel = Kernel(ast)
-    kernel.apply_opts(opts)
-    run_linearizer(kernel)
-
-  def test_2(self):
-    # kernel 254
-    ast = UOp(Ops.SINK, dtypes.void, arg=None, src=(
-      UOp(Ops.STORE, dtypes.void, arg=None, src=(
-        UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(3145728), arg=0, src=()),
-        UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(4, 1024, 1, 768), strides=(786432, 768, 0, 1), offset=0, mask=None, contiguous=True),)), src=()),
-        UOp(Ops.REDUCE_AXIS, dtypes.float, arg=(Ops.ADD, (2,)), src=(
-        UOp(Ops.MUL, dtypes.float, arg=None, src=(
-          UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-          UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(38633472), arg=1, src=()),
-          UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(4, 1024, 50304, 768), strides=(0, 0, 768, 1), offset=0, mask=None, contiguous=False),)), src=()),)),
-          UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-          UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(205852672), arg=2, src=()),
-          UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(4, 1024, 50304, 768), strides=(51463168, 50257, 1, 0), offset=0, mask=((0, 4), (0, 1024), (0, 50257), (0, 768)), contiguous=False),)), src=()),)),)),)),)),))
-
-    opts = [Opt(op=OptOps.LOCAL, axis=1, arg=16), Opt(op=OptOps.LOCAL, axis=0, arg=8), Opt(op=OptOps.UPCAST, axis=2, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.LOCAL, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=3, arg=4)]
-    kernel = Kernel(ast)
-    kernel.apply_opts(opts)
-    run_linearizer(kernel)
-
-if __name__ == "__main__":
-  unittest.main()
@@ -72,7 +72,7 @@ class TestQuantizeOnnxCPU(unittest.TestCase):
    out_file = get_quantized_model(sz)
    run_onnx = OnnxRunner(out_file)
    inp = Tensor(np.random.uniform(size=(sz, sz)).astype(np.float32))
-    with Context(DONT_REALIZE_EXPAND=1, QUANTIZE=1):
+    with Context(QUANTIZE=1):
      sched = run_onnx({"input":inp})["output"].schedule()
      ei = lower_schedule_item(sched[-2])
      daccs = [u for u in ei.prg.p.uops if u.op is Ops.DEFINE_REG]
@@ -86,8 +86,7 @@ class TestQuantizeOnnx(unittest.TestCase):
    # divide is ~1500-2000 without reduce_range, 750-900 with it
    out_file = get_quantized_model(sz)
    run_onnx_jit, _ = load_onnx_model(out_file)
-    with Context(DONT_REALIZE_EXPAND=1):
-      run_onnx_jit(input=Tensor(np.random.uniform(size=(sz, sz)).astype(np.float32)))
+    run_onnx_jit(input=Tensor(np.random.uniform(size=(sz, sz)).astype(np.float32)))

  def test_prequant_conv2d_1x1(self):
    X = Tensor(np.random.uniform(0, 255, size=(1, 32, 128, 128)).astype(np.uint8))
@@ -109,11 +108,10 @@ class TestQuantizeOnnx(unittest.TestCase):
    N = 512
    X = Tensor(np.random.uniform(0, 255, size=(N,N)).astype(xi))
    W = Tensor(np.random.uniform(0, 255, size=(N,N)).astype(wi))
-    with Context(DONT_REALIZE_EXPAND=1):
-      # this divide is interesting and forces the accumulator to actually be an int
-      out = (X.cast("int").matmul(W.cast("int"))//1000).cast("int8")
-      opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)]
-      sexec(out, opts)
+    # this divide is interesting and forces the accumulator to actually be an int
+    out = (X.cast("int").matmul(W.cast("int"))//1000).cast("int8")
+    opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)]
+    sexec(out, opts)

  def test_prequant_gemm_handcode(self):
    src = """typedef int int128 __attribute__((aligned(512),vector_size(512)));
@@ -203,14 +201,12 @@ class TestQuantizeOnnx(unittest.TestCase):
  def test_prequant_gemm_intacc(self, xi=np.uint8, wi=np.uint8, replace_src=None, N=512, clip=True, opts=None):
    X = Tensor(m1:=(np.random.uniform(0, 255, size=(N,N)).astype(xi))).realize()
    W = Tensor(m2:=(np.random.uniform(0, 255, size=(N,N)).astype(wi))).realize()
-    # ugh, it's so broken with those casts. need DONT_REALIZE_EXPAND=1 python3 test/test_quantize_onnx.py TestQuantizeOnnx.test_prequant
    tg_dtype = dtypes.int8 if xi == np.int8 else dtypes.uint8
-    with Context(DONT_REALIZE_EXPAND=1):
-      out = (X.int().matmul(W.int())//1000)
-      if clip: out = out.clip(dtypes.min(tg_dtype),dtypes.max(tg_dtype))
-      out = out.cast(tg_dtype)
-      opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)] if opts is None else opts
-      sexec(out, opts, replace_src, run_count=1)
+    out = (X.int().matmul(W.int())//1000)
+    if clip: out = out.clip(dtypes.min(tg_dtype),dtypes.max(tg_dtype))
+    out = out.cast(tg_dtype)
+    opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)] if opts is None else opts
+    sexec(out, opts, replace_src, run_count=1)
    tout = out.numpy()
    mout = ((m1.astype(np.int32) @ m2.astype(np.int32)) // 1000)
    if clip: mout = mout.clip(dtypes.min(tg_dtype),dtypes.max(tg_dtype))
@@ -225,7 +221,6 @@ class TestQuantizeOnnx(unittest.TestCase):

  def test_prequant_gemv(self):
    N = 2048
-    # ugh, it's so broken with those casts. need DONT_REALIZE_EXPAND=1 python3 test/test_quantize_onnx.py TestQuantizeOnnx.test_prequant
    X = Tensor(np.random.uniform(0, 255, size=(1,N)).astype(np.uint8)).realize()
    W = Tensor(np.random.uniform(0, 255, size=(N,N)).astype(np.uint8)).realize()
    #out = X.cast(dtypes.int) @ W.cast(dtypes.int)
@@ -1,6 +1,6 @@
 import unittest
 from tinygrad import Tensor, nn
-from tinygrad.helpers import Context, GlobalCounters
+from tinygrad.helpers import Context, GlobalCounters, CI
 from tinygrad.uop.ops import graph_rewrite, PatternMatcher, UPat, Ops

 class TestRangeifyAssign(unittest.TestCase):
@@ -17,8 +17,22 @@ class TestRangeifyAssign(unittest.TestCase):
    self.assertListEqual(lst, lst3)
    self.assertListEqual(lst2, B.permute(1, 0).tolist())

+class TestRangeifyEdgeCase(unittest.TestCase):
+  def test_matmul_relu_cat(self):
+    a = Tensor.ones(100, 512).contiguous().realize()
+    c = Tensor.ones(1, 512).contiguous().realize()
+    cm = Tensor.ones(512, 512)
+    c = c @ cm
+    c = c.relu()
+
+    res = Tensor.cat(a, c, dim=0)
+    self.assertEqual(res.numpy()[-1, :16].tolist(), [512] * 16)
+
+# *** non CI rangeify tests below this line ***
+
 N = 256

+@unittest.skipIf(CI, "useless in CI, doesn't test anything")
 class TestRangeifyOpt(unittest.TestCase):
  def test_randperm(self):
    Tensor.randperm(10000).realize()
@@ -54,6 +68,7 @@ class TestRangeifyOpt(unittest.TestCase):
    A = Tensor.empty(8,8,8,8).permute(1,0,3,2).flatten()
    A.sum().realize()

+@unittest.skipIf(CI, "useless in CI, doesn't test anything")
 class TestRangeify(unittest.TestCase):
  def test_groupnorm(self):
    # ranges 1 and 3 are merging
@@ -280,16 +295,5 @@ class TestRangeifyPM(unittest.TestCase):
    b = self.base.pad(((0,1),(0,0))).pad(((0,0),(0,1)))
    self.assert_same(a, b)

-class TestRangeifyEdgeCase(unittest.TestCase):
-  def test_matmul_relu_cat(self):
-    a = Tensor.ones(100, 512).contiguous().realize()
-    c = Tensor.ones(1, 512).contiguous().realize()
-    cm = Tensor.ones(512, 512)
-    c = c @ cm
-    c = c.relu()
-
-    res = Tensor.cat(a, c, dim=0)
-    self.assertEqual(res.numpy()[-1, :16].tolist(), [512] * 16)
-
 if __name__ == '__main__':
  unittest.main()
@@ -2,9 +2,8 @@
 # schedule confirms the right things are capable of fusing
 # NOTE: this has overlap with external_test_opt.py

-import unittest
+import unittest, functools
 import numpy as np
-import functools
 from typing import cast
 from hypothesis import assume, given, settings, strategies as strat

@@ -31,7 +30,6 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te
  # test lowering all the ScheduleItems to ExecItems
  kernel_cnt = len([si for si,ei in lower_schedule(sched.copy()) if isinstance(ei.prg, CompiledRunner) or not filter_sink])
  if kernel_cnt != allowed:
-    return sched # allow different kernel count, TODO: fix the asserts
    print(f"SCHEDULE ISSUE, expecting {allowed} got {len(sched)}")
    if DEBUG >= 3:
      for i,s in enumerate(sched):
@@ -117,8 +115,7 @@ class TestSchedule(unittest.TestCase):
    c = a+b
    with self.assertRaisesRegex(RuntimeError, "all buffers must be on the same device"): check_schedule(c, 2)

-  @unittest.skipUnless(is_dtype_supported(dtypes.half) and getenv("CAST_AFTER_EXPAND"), "need half and CAST_AFTER_EXPAND=1")
-  @unittest.skip("CAST_AFTER_EXPAND is not supported")
+  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
  def test_expand_buffer_before_cast(self):
    a = Tensor.randn(4, 2, 1).realize().permute((1, 0, 2))
    b = a.cast(dtypes.half).expand((2, 4, 4))+2
@@ -128,7 +125,7 @@ class TestSchedule(unittest.TestCase):
  def test_indexing_scalars_simple(self):
    X = Tensor.randn(2, 2).realize()
    xt = X[Tensor(1)][Tensor(0)]
-    run_schedule(check_schedule(xt, 2))
+    run_schedule(check_schedule(xt, 1))
    np.testing.assert_equal(xt.numpy(), X.numpy()[1][0])

  @unittest.skipIf(CI and Device.DEFAULT == "NV", "crashes on NV CI")
@@ -148,31 +145,30 @@ class TestSchedule(unittest.TestCase):
    assume(a<x and b<y)
    X = Tensor.randn(x, y).realize()
    xt = X[Tensor(a)][Tensor(b)]
-    run_schedule(check_schedule(xt, 2))
+    run_schedule(check_schedule(xt, 1))
    np.testing.assert_equal(xt.numpy(), X.numpy()[a][b])

  def test_push_pads_elementwise(self):
    x = Tensor.full((4,4), 2.).contiguous().realize()
    y = Tensor.full((4,4), 4.).contiguous().realize()
    z = (x.reciprocal()*y).pad((None, (0,1),)).sum()
-    run_schedule(check_schedule(z, 2))
+    run_schedule(check_schedule(z, 1))
    self.assertEqual(z.item(), 32)

  def test_push_pads_contiguous(self):
    x = Tensor.full((4,1), 2.).contiguous()
    y = Tensor.full((4,4), 4.).contiguous()
    z = (x.reciprocal().expand(4,4)*y).pad((None, (0,1),)).sum()
-    run_schedule(check_schedule(z, 2, [x,y]))
+    run_schedule(check_schedule(z, 1, [x,y]))
    self.assertEqual(z.item(), 32)

  def test_rand(self):
    x = Tensor.rand(32)
-    check_schedule(x, 4, [Tensor._device_rng_counters[x.device]])
+    check_schedule(x, 1, [Tensor._device_rng_counters[x.device]])

  def test_rand_recompute_arange(self):
    x = Tensor.rand(32)
-    with Context(DONT_GROUP_REDUCES=1):
-      check_schedule(x, 3, [Tensor._device_rng_counters[x.device]])
+    check_schedule(x, 1, [Tensor._device_rng_counters[x.device]])

  def test_empty_is_not_realized(self):
    a = Tensor.empty(10)
@@ -189,10 +185,7 @@ class TestSchedule(unittest.TestCase):

  def test_simplify_padded_const(self):
    a = Tensor.empty(1022).cummax(axis=0)
-    check_schedule(a, 5)
-    # TODO: what is this testing?
-    #ast = sched[0].ast
-    #self.assertLessEqual(len([u for u in ast.toposort() if u.op is Ops.WHERE]), 6)
+    check_schedule(a, 3)

  def test_basic_binop_fusion(self):
    a = Tensor.empty(10)
@@ -265,18 +258,17 @@ class TestSchedule(unittest.TestCase):
    c = a.sum(axis=0) + b
    check_schedule(c, 1)

-  # not pushing permutes through reduces
  def test_reduce_permute_binop_fusion(self):
    a = Tensor.empty(10,10,10)
    b = Tensor.empty(10,10,1)
    c = a.sum(axis=0, keepdim=True).permute(2,1,0) + b
-    check_schedule(c, 2)
+    check_schedule(c, 1)

  def test_allow_push_permutes(self):
    a = Tensor.randn(10,10,10).realize()
    b = Tensor.randn(10,10,1).realize()
    c = a.sum(axis=0, keepdim=True).permute(2,1,0) + b
-    with Context(DONT_GROUP_REDUCES=1): run_schedule(check_schedule(c, 1))
+    run_schedule(check_schedule(c, 1))
    np.testing.assert_allclose(c.numpy(), np.sum(a.numpy(), axis=0, keepdims=True).transpose(2,1,0)+b.numpy())

  def test_binop_early_reshape_reduce_fusion(self):
@@ -341,7 +333,7 @@ class TestSchedule(unittest.TestCase):
    r1 = (x - r0).sum(axis=0).div(2)
    out0 = r0 + y
    out1 = r1 + y
-    schedule = check_schedule([out0, out1], 2)
+    schedule = check_schedule([out0, out1], 4)
    reduceops = [x for si in schedule for x in si.ast.toposort() if x.op in {Ops.REDUCE_AXIS, Ops.REDUCE}]
    self.assertEqual(len(reduceops), 2) # why is RANGEIFY different?

@@ -374,7 +366,7 @@ class TestSchedule(unittest.TestCase):
    b = Tensor.full((4,), 2.).contiguous()
    first = a.assign(b)
    second = a.assign(b)
-    check_schedule([first, second], 1)
+    check_schedule([first, second], 2) # TODO: 1?

  # NOTE: this is causing "LAZYCACHE=1 incorrectly reuses contiguous const" #4562
  # should contiguous dedup?
@@ -454,7 +446,7 @@ class TestSchedule(unittest.TestCase):
  @unittest.skipUnless(is_dtype_supported(dtypes.ulong), "Needs ulong")
  def test_fold_conv_batchnorm_optim(self):
    # this is too high
-    for optim, cnt in [(nn.optim.Adam, 30), (nn.optim.SGD, 11)]:
+    for optim, cnt in [(nn.optim.Adam, 30), (nn.optim.SGD, 13)]:
      with self.subTest(optim=optim.__name__):
        with Tensor.train():
          img = Tensor.ones(1,3,4,4)
@@ -475,7 +467,7 @@ class TestSchedule(unittest.TestCase):
      fw = bn(x).contiguous_backward().relu().contiguous()
      fw.sum().backward()
      # TODO: this is too many
-      check_schedule([x.grad, bn.weight.grad, bn.bias.grad, fw], 10)
+      check_schedule([x.grad, bn.weight.grad, bn.bias.grad, fw], 9)

  def test_fold_conv_relu(self):
    c1 = nn.Conv2d(3,16,3)
@@ -518,9 +510,8 @@ class TestSchedule(unittest.TestCase):
    img = Tensor.empty(64,64)
    x = (img.sum(0) + img.sum(1))
    out = x.relu()
-    check_schedule(out, 2)
+    check_schedule(out, 1)

-  #@unittest.skip("failing in old lazy")
  def test_push_permute_through_reshape(self):
    a = Tensor.empty(16,16)
    b = Tensor.empty(16,16)
@@ -554,7 +545,7 @@ class TestSchedule(unittest.TestCase):
    c = a+b
    d = a.reshape(10,1)+b.reshape(10,1)
    out = c.sum() + d.sum()
-    check_schedule(out, 2)
+    check_schedule(out, 1)

  def test_children_dont_push(self):
    a = Tensor.empty(10, 10, 1)
@@ -562,7 +553,7 @@ class TestSchedule(unittest.TestCase):
    d = (a+b).expand(10, 10, 10)
    e = (a+b).permute(2,1,0)
    f = d+e
-    check_schedule(f, 2)
+    check_schedule(f, 1)

  # failing in new lazy
  @unittest.skip("always fusing elementwise")
@@ -601,13 +592,13 @@ class TestSchedule(unittest.TestCase):
    e = c[0] * d
    check_schedule(e, 1)

-  def test_expand_nofuse(self):
+  def test_expand_fuse(self):
    a = Tensor.empty(1, 16)
    b = Tensor.empty(1, 16)
    c = a * b
    d = Tensor.empty(8192, 16)
    e = c * d
-    check_schedule(e, 2)
+    check_schedule(e, 1)

  # this is the failing case in openpilot...it's very simple like this
  def test_image_conv_fusion(self):
@@ -625,7 +616,7 @@ class TestSchedule(unittest.TestCase):

    # NOOP, 3 convs, contiguous
    #check_schedule(x, 5)
-    check_schedule(x, 8)
+    check_schedule(x, 7)

  def test_image_conv_fusion_minimal(self):
    b1 = Tensor.empty(16)
@@ -808,13 +799,13 @@ class TestSchedule(unittest.TestCase):
    x = Tensor.empty(32, 32, 32)
    y = Tensor.empty(32, 32)
    out = x.sum(axis=2).T+y
-    check_schedule(out, 2)
+    check_schedule(out, 1)

  def test_two_elus_sum(self):
    x = Tensor.empty(32, 32)
    y = Tensor.empty(32, 32)
    out = x.sum(1).relu().elu() + y.sum(1).relu().elu()
-    check_schedule(out, 2)
+    check_schedule(out, 1)

  @unittest.skipUnless(SPLIT_REDUCEOP, "Testing split reducop requires SPLIT_REDUCEOP")
  def test_preserve_multistage_reduce(self):
@@ -827,7 +818,7 @@ class TestSchedule(unittest.TestCase):
  def test_multistage_reduce(self):
    x = Tensor.empty(32, 32, 32)
    out = x.sum(2).relu().sum(1)
-    check_schedule(out, 2)
+    check_schedule(out, 1)

  def test_multistage_reduce_fork(self):
    x = Tensor.empty(32, 32, 32)
@@ -843,7 +834,7 @@ class TestSchedule(unittest.TestCase):
    z = y.matmul(x).sum()
    z.backward()
    out = x.grad.contiguous()
-    run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 1))
    np.testing.assert_allclose(out.numpy(), np.ones((64,64)))

  def test_example_matmul_contig(self):
@@ -852,7 +843,7 @@ class TestSchedule(unittest.TestCase):
    z = y.matmul(x).sum()
    z.backward()
    out = x.grad.contiguous()
-    run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 1))
    np.testing.assert_allclose(out.numpy(), np.ones((64,64)))

  def test_example_matmul_same(self):
@@ -860,7 +851,7 @@ class TestSchedule(unittest.TestCase):
    z = x.matmul(x).sum()
    z.backward()
    out = x.grad.contiguous()
-    run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 1))
    # NOTE: the gradient flows twice
    np.testing.assert_allclose(out.numpy(), 2*np.ones((64,64)))

@@ -883,8 +874,7 @@ class TestSchedule(unittest.TestCase):
    x = x.sum(1)
    x = x[:16]
    out = x + y
-    # NOTE: this could be 1 kernel if we mask the store?
-    check_schedule(out, 2)
+    check_schedule(out, 1)

  def test_multireduce_shrink(self):
    Tensor.manual_seed(0)
@@ -896,8 +886,7 @@ class TestSchedule(unittest.TestCase):
    b_out = b.sum(1)
    b_out = b_out[:16]
    out = a_out + b_out + c
-    # run_schedule(check_schedule(out, 2))  # TODO: this should be 1 (can we make it 1 with the new linearizer?)
-    run_schedule(check_schedule(out, 3))
+    run_schedule(check_schedule(out, 1))
    np.testing.assert_allclose(out.numpy(), a.numpy().sum(axis=1)[:16] + b.numpy().sum(axis=1)[:16] + c.numpy(), atol=1e-4, rtol=1e-4)

  # broken due to const folding and two contiguous are different kernels
@@ -914,7 +903,7 @@ class TestSchedule(unittest.TestCase):
    out0 = a.sum() + 2
    out1 = a.sum() + 4
    out2 = out0 * out1
-    run_schedule(check_schedule([out0, out1, out2], 1))
+    run_schedule(check_schedule([out0, out1, out2], 3)) # TODO: 1?
    np.testing.assert_allclose(out0.numpy(), out0_np:=a.numpy().sum()+2, atol=1e-4, rtol=1e-6)
    np.testing.assert_allclose(out1.numpy(), out1_np:=a.numpy().sum()+4, atol=1e-4, rtol=1e-6)
    np.testing.assert_allclose(out2.numpy(), out0_np*out1_np, atol=1e-4, rtol=1e-6)
@@ -925,7 +914,7 @@ class TestSchedule(unittest.TestCase):
    out0 = a.sum().exp2()
    # out1 has two paths to a.sum()
    out1 = a.sum() + out0
-    run_schedule(check_schedule([out0, out1], 1))
+    run_schedule(check_schedule([out0, out1], 2)) # TODO: 1?
    np.testing.assert_allclose(out0.numpy(), out0_np:=np.exp2(a.numpy().sum()), atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(out1.numpy(), a.numpy().sum()+out0_np, atol=1e-4, rtol=1e-6)

@@ -938,7 +927,7 @@ class TestSchedule(unittest.TestCase):
    out2 = b.sum().exp2()
    out3 = b.sum() + out2
    # run_schedule(check_schedule([out0, out1, out2, out3], 1))
-    run_schedule(check_schedule([out0, out1, out2, out3], 6))
+    run_schedule(check_schedule([out0, out1, out2, out3], 4))
    np.testing.assert_allclose(out0.numpy(), np_out0:=np.exp2(a.numpy().sum()), atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(out1.numpy(), np_out1:=a.numpy().sum()+np_out0, atol=1e-4, rtol=1e-4)
    np_b = (a.numpy() + np_out0 + np_out1)
@@ -953,7 +942,7 @@ class TestSchedule(unittest.TestCase):
    out0 = a.sum() + b.sum() + 2
    out1 = a.sum() + b.sum() + 4
    # run_schedule(check_schedule([out0, out1], 1))
-    run_schedule(check_schedule([out0, out1], 4))
+    run_schedule(check_schedule([out0, out1], 2))
    np.testing.assert_allclose(out0.numpy(), a.numpy().sum()+b.numpy().sum()+2, atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(out1.numpy(), a.numpy().sum()+b.numpy().sum()+4, atol=1e-4, rtol=1e-4)

@@ -980,7 +969,7 @@ class TestSchedule(unittest.TestCase):
    out1 = b.max() + out0*2
    out2 = a.sum() + out1
    # run_schedule(check_schedule([out0, out1, out2], 1))
-    run_schedule(check_schedule([out0, out1, out2], 4))
+    run_schedule(check_schedule([out0, out1, out2], 3))
    np.testing.assert_allclose(out0.numpy(), out0_np:=a.numpy().sum()+4, atol=1e-4, rtol=1e-6)
    np.testing.assert_allclose(out1.numpy(), out1_np:=b.numpy().max() + out0_np*2, atol=1e-4, rtol=1e-6)
    np.testing.assert_allclose(out2.numpy(), a.numpy().sum() + out1_np, atol=1e-4, rtol=1e-6)
@@ -1017,7 +1006,7 @@ class TestSchedule(unittest.TestCase):
    b = Tensor.empty(10,)
    c = a.sum() + b[0]
    d = a.sum() + 2
-    check_schedule([c, d], 1)
+    check_schedule([c, d], 2) # TODO: 1?

  def test_reduce_multiple_paths_midshrink(self):
    a = Tensor.empty(4, 4)
@@ -1046,7 +1035,7 @@ class TestSchedule(unittest.TestCase):
    k = Tensor.randn(32,8,16,8).realize()
    v = Tensor.randn(32,8,16,8).realize()
    out = Tensor.scaled_dot_product_attention(q,k,v)
-    run_schedule(check_schedule(out, 5))
+    run_schedule(check_schedule(out, 4))
    if getenv("CHECK", 1):
      import torch
      compare = torch.nn.functional.scaled_dot_product_attention(torch.tensor(q.numpy()),torch.tensor(k.numpy()),torch.tensor(v.numpy()))
@@ -1054,7 +1043,7 @@ class TestSchedule(unittest.TestCase):

    with Context(FUSE_ATTENTION=1):
      out = Tensor.scaled_dot_product_attention(q,k,v)
-      run_schedule(check_schedule(out, 1))
+      run_schedule(check_schedule(out, 4)) # TODO: should be 1?
      if getenv("CHECK", 1):
        import torch
        compare = torch.nn.functional.scaled_dot_product_attention(torch.tensor(q.numpy()),torch.tensor(k.numpy()),torch.tensor(v.numpy()))
@@ -1067,7 +1056,7 @@ class TestSchedule(unittest.TestCase):
    c = Tensor.randn(4, 32).realize()
    out = (c * a.sum(-1, keepdim=True)).sum(-1) + (b * a.sum(-1, keepdim=True)).sum(-1) # a.sum has >1 children but should still fuse
    # run_schedule(check_schedule(out, 1))
-    run_schedule(check_schedule(out, 3))
+    run_schedule(check_schedule(out, 2))
    np.testing.assert_allclose(out.numpy(), \
      (c.numpy()*a.numpy().sum(axis=-1,keepdims=True)).sum(-1) + (b.numpy()*a.numpy().sum(axis=-1,keepdims=True)).sum(-1), atol=1e-4, rtol=1e-4)

@@ -1112,8 +1101,7 @@ class TestSchedule(unittest.TestCase):
    x = Tensor.randn(4, 32).realize()
    y = Tensor.randn(4, 32).realize()
    out = y.sum(axis=-1) + x.sum(axis=-1)
-    # run_schedule(check_schedule(out, 1))
-    run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 1))
    np.testing.assert_allclose(out.numpy(), y.numpy().sum(axis=-1) + x.numpy().sum(axis=-1), atol=1e-4, rtol=1e-4)

  def test_multireduce_fusion_sequential(self):
@@ -1130,7 +1118,7 @@ class TestSchedule(unittest.TestCase):
    y = Tensor.randn(4, 32).realize()
    out = x.std(-1) + y.std(-1)
    # run_schedule(check_schedule(out, 1))
-    run_schedule(check_schedule(out, 4))
+    run_schedule(check_schedule(out, 3))
    np.testing.assert_allclose(out.numpy(), x.numpy().std(axis=-1, ddof=1) + y.numpy().std(axis=-1, ddof=1), atol=1e-4, rtol=1e-4)

  def test_multireduce_diffops_sequential(self):
@@ -1146,8 +1134,7 @@ class TestSchedule(unittest.TestCase):
    x = Tensor.randn(4, 32).realize()
    y = Tensor.randn(4, 32).realize()
    out = x.sum(-1) + y.max(-1)
-    # run_schedule(check_schedule(out, 1))
-    run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 1))
    np.testing.assert_allclose(out.numpy(), x.numpy().sum(axis=-1) + y.numpy().max(axis=-1), atol=1e-4, rtol=1e-4)

  def test_multireduce_fusion_sequential_and_parallel(self):
@@ -1159,7 +1146,7 @@ class TestSchedule(unittest.TestCase):
    np_mu = (x.numpy() - x.numpy().max(axis=-1, keepdims=True)).mean(axis=-1, keepdims=True) + \
      (y.numpy() - y.numpy().max(axis=-1, keepdims=True)).mean(axis=-1, keepdims=True)
    # run_schedule(check_schedule(out, 1))
-    run_schedule(check_schedule(out, 6))
+    run_schedule(check_schedule(out, 5))
    np.testing.assert_allclose(out[0].numpy(), np.sqrt(np.square(x.numpy() - np_mu).sum(-1)/x.shape[-1]), atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(out[1].numpy(), np.sqrt(np.square(y.numpy() - np_mu).sum(-1)/y.shape[-1]), atol=1e-4, rtol=1e-4)

@@ -1168,8 +1155,7 @@ class TestSchedule(unittest.TestCase):
    a,b = Tensor.randn(4, 64).realize(), Tensor.rand(64,8).realize()
    c,d = Tensor.randn(4, 64).realize(), Tensor.rand(64,8).realize()
    out = a@b + c@d
-    # run_schedule(check_schedule(out, 1))
-    run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 1))
    np.testing.assert_allclose(out.numpy(), a.numpy()@b.numpy() + c.numpy()@d.numpy(), atol=1e-4, rtol=1e-4)

  def test_softmax_fusion(self):
@@ -1180,17 +1166,15 @@ class TestSchedule(unittest.TestCase):
    expected = (x_exp:=np.exp(x.numpy()-x.numpy().max(-1, keepdims=True)))/x_exp.sum(-1, keepdims=True)
    np.testing.assert_allclose(out.numpy(), expected, atol=1e-4, rtol=1e-4)

-  # TODO: rangeify stores the output in float32
  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
-  @unittest.expectedFailure
  def test_softmax_upcast(self):
    # input half, softmax in float
    Tensor.manual_seed(0)
    x = Tensor.randn(4, 12, 64, 64, dtype=dtypes.half).realize()
    out = x.softmax(dtype=dtypes.float)
    sched = out.schedule()
-    self.assertEqual(len(sched), 2)
-    self.assertEqual(sched[0].bufs[0].dtype, dtypes.half)
+    self.assertEqual(len(sched), 3)
+    self.assertEqual(sched[0].bufs[0].dtype, dtypes.float)

    # input float, softmax in float
    Tensor.manual_seed(0)
@@ -1222,12 +1206,12 @@ class TestSchedule(unittest.TestCase):
  def test_scaled_dot_product_attention_fusion(self):
    x, y, z, m = (Tensor.empty(32, 8, 16, 16) for _ in range(4))
    out = Tensor.scaled_dot_product_attention(x, y, z, attn_mask=m)
-    check_schedule(out, 5)
+    check_schedule(out, 4)

  def test_scaled_dot_product_attention_causal_fusion(self):
    x, y, z = (Tensor.empty(32, 8, 16, 16) for _ in range(3))
    out = Tensor.scaled_dot_product_attention(x, y, z, is_causal=True)
-    check_schedule(out, 5)
+    check_schedule(out, 4)

  def test_adam_step_fusion(self):
    with Tensor.train():
@@ -1257,7 +1241,7 @@ class TestSchedule(unittest.TestCase):
      opt = nn.optim.Adam(nn.state.get_parameters([c1, c2]), lr=1e-4)
      opt.zero_grad()
      c2(c1(img).relu()).relu().sum().backward()
-      check_schedule(opt.schedule_step(), 20)
+      check_schedule(opt.schedule_step(), 18)

  def test_sgd_conv_fuse(self):
    with Tensor.train():
@@ -1267,7 +1251,7 @@ class TestSchedule(unittest.TestCase):
      opt = nn.optim.SGD(nn.state.get_parameters(c1))
      opt.zero_grad()
      c1(img).relu().sum().backward()
-      check_schedule(opt.schedule_step(), 3)
+      check_schedule(opt.schedule_step(), 5) # TODO: 3?

  def test_sgd_2convs_fuse(self):
    with Tensor.train():
@@ -1290,7 +1274,7 @@ class TestSchedule(unittest.TestCase):
      opt = nn.optim.SGD(nn.state.get_parameters([c1, c2]), nesterov=True, momentum=0.9, weight_decay=0.1)
      opt.zero_grad()
      c2(c1(img).relu()).relu().sum().backward()
-      check_schedule(opt.schedule_step(), 13)
+      check_schedule(opt.schedule_step(), 15)

  def test_sgd_4convs_fuse(self):
    with Tensor.train():
@@ -1303,7 +1287,7 @@ class TestSchedule(unittest.TestCase):
      opt = nn.optim.SGD(nn.state.get_parameters([c1, c2, c3, c4]))
      opt.zero_grad()
      c4(c3(c2(c1(img).relu()).relu()).relu()).relu().sum().backward()
-      check_schedule(opt.schedule_step(), 17)
+      check_schedule(opt.schedule_step(), 15)

  def test_sgd_4convs_fuse_conv_bw(self):
    with Tensor.train():
@@ -1316,50 +1300,7 @@ class TestSchedule(unittest.TestCase):
      opt = nn.optim.SGD(nn.state.get_parameters([c1, c2, c3, c4]))
      opt.zero_grad()
      c4(c3(c2(c1(img).relu()).relu()).relu()).relu().sum().backward()
-      check_schedule(opt.schedule_step(), 14)
-
-  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
-  @unittest.expectedFailure
-  def test_prefer_half_buffer(self):
-    x = Tensor.ones(4).contiguous().realize()
-    # y = Tensor.ones(4).contiguous().realize()
-    z = Tensor.ones(4, 4).contiguous().realize()
-
-    # should not create extra kernel if output will be realized anyways
-    dummy = x.sum().half().float()
-    check_schedule(dummy, 1)
-    dummy = x.sum().half().float().contiguous() + 1
-    check_schedule(dummy, 2)
-
-    # shared between two outputs
-    shared = x.sum().half().float()
-    a = shared * 2
-    b = shared * 3
-    sched = check_schedule([a, b], 3)
-    # store reduceop in half
-    self.assertEqual(sched[0].bufs[0].dtype, dtypes.half)
-    # fuse cast with the child kernel
-    self.assertEqual(sched[1].bufs[0].dtype, dtypes.float)
-    self.assertEqual(sched[2].bufs[0].dtype, dtypes.float)
-
-    # reduce
-    a = z.sum(axis=0).half().float().sum(axis=0)
-    sched = check_schedule(a, 2)
-    self.assertEqual(sched[0].bufs[0].dtype, dtypes.half)
-    self.assertEqual(sched[1].bufs[0].dtype, dtypes.float)
-
-    # expand
-    # expand will realize just after the .float(), so requires change to realize-before-expand
-    # normal = (x.sum().half().float().reshape(1) * y).sum()
-    # sched = check_schedule(normal, 2)
-    # for si in sched[:-1]: assert all(out.dtype == dtypes.half for out in si.outputs[:-1])
-
-    # parallel reduce
-    # a = x.sum().half().float() * y.sum().half().float()
-    # b = a + 1
-    # c = a + 2
-    # sched = check_schedule([b, c], 4)
-    # doesn't store either in half because it doesn't chase
+      check_schedule(opt.schedule_step(), 15)

  def test_reduce_simple_chase(self):
    a = Tensor.empty(4, 4, 4)
@@ -1408,7 +1349,7 @@ class TestSchedule(unittest.TestCase):
    c = Tensor.empty(16, )
    r = a.sum(1) + c
    d = r[:4] * b
-    check_schedule(d, 2)
+    check_schedule(d, 1)

  def test_multireduce_push_shrink_chase(self):
    Tensor.manual_seed(0)
@@ -1418,22 +1359,20 @@ class TestSchedule(unittest.TestCase):
    d = Tensor.randn(16, 16).realize()
    r = a.sum(1) + c
    out = r[:4] * b + d.sum(1)[:4]
-    # schedule = check_schedule(out, 2)
-    schedule = check_schedule(out, 3)
+    schedule = check_schedule(out, 1)
    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)
+    check_schedule(b, 1)

  def test_multireduce_midreduce_nochase(self):
    Tensor.manual_seed(0)
    a = Tensor.randn(16, 16).realize()
    b = (a.sum(0)+a.max(0) + a.max(1)+a.sum(1)) + 2
-    # schedule = check_schedule(b, 2)
-    schedule = check_schedule(b, 4)
+    schedule = check_schedule(b, 1)
    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)

@@ -1445,7 +1384,7 @@ class TestSchedule(unittest.TestCase):
    c = a.sum() + 2
    d = (a.sum() - b.sum()) * 4
    # run_schedule(check_schedule([c, d], 1))
-    run_schedule(check_schedule([c, d], 3))
+    run_schedule(check_schedule([c, d], 2))
    np.testing.assert_allclose(c.numpy(), a.numpy().sum()+2, atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(d.numpy(), (a.numpy().sum() - b.numpy().sum()) * 4, atol=1e-4, rtol=1e-4)

@@ -1471,7 +1410,7 @@ class TestSchedule(unittest.TestCase):
    e = c * d
    f = b.sum() - e
    # run_schedule(check_schedule([c, d, e, f], 1))
-    run_schedule(check_schedule([c, d, e, f], 2))
+    run_schedule(check_schedule([c, d, e, f], 4))
    np.testing.assert_allclose(c.numpy(), c_np:=a.numpy().sum()+2, atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(d.numpy(), d_np:=a.numpy().sum()*2, atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(e.numpy(), e_np:=c_np*d_np, atol=1e-4, rtol=1e-4)
@@ -1486,7 +1425,7 @@ class TestSchedule(unittest.TestCase):
    e = c * d
    f = (b - d).sum() - e
    # run_schedule(check_schedule([c, d, e, f], 1))
-    run_schedule(check_schedule([c, d, e, f], 5))
+    run_schedule(check_schedule([c, d, e, f], 4))
    np.testing.assert_allclose(c.numpy(), c_np:=a.numpy().sum()+2, atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(d.numpy(), d_np:=a.numpy().sum()*2, atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(e.numpy(), e_np:=c_np*d_np, atol=1e-4, rtol=1e-4)
@@ -1505,8 +1444,7 @@ class TestSchedule(unittest.TestCase):
    a = Tensor.randn(3, 4, 5).realize()
    b = Tensor.randn(3, 4, 5).realize()
    out = (a.pad(((0, 1), (0, 1), (0, 1)), value=1.0).sum(keepdim=True)+b.pad(((0, 1), (0, 1), (0, 1)), value=1.0).sum()).contiguous()
-    # run_schedule(check_schedule(out, 1))
-    run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 1))
    np.testing.assert_allclose(out.numpy(), np.pad(a.numpy(), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum(keepdims=True) + \
                                                   np.pad(b.numpy(), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum(), atol=1e-4, rtol=1e-4)

@@ -1514,7 +1452,7 @@ class TestSchedule(unittest.TestCase):
    Tensor.manual_seed(0)
    a = Tensor.rand(3, 4, 5).realize()
    out = a.log2().pad(((0, 1), (0, 1), (0, 1)), value=1.0).sum().contiguous()
-    run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 1))
    np.testing.assert_allclose(out.numpy(), np.pad(np.log2(a.numpy()), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum(), atol=1e-5, rtol=1e-6)

  def test_multireduce_pad_reduce_unsafe(self):
@@ -1523,7 +1461,7 @@ class TestSchedule(unittest.TestCase):
    b = Tensor.randn(3, 4, 5).abs().realize()
    out = (a.log2().pad(((0, 1), (0, 1), (0, 1)), value=1.0).sum()+b).abs().log2().pad(((0, 1), (0, 1), (0, 1)), value=1.0).sum().contiguous()
    # run_schedule(check_schedule(out, 1))
-    run_schedule(check_schedule(out, 4))
+    run_schedule(check_schedule(out, 2))
    np.testing.assert_allclose(out.numpy(), np.pad(np.log2(np.abs(np.pad(np.log2(a.numpy()), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum() + \
                                                   b.numpy())), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum(), atol=3e-4, rtol=1e-5)

@@ -1537,7 +1475,7 @@ class TestSchedule(unittest.TestCase):
  def test_shrink_pad_unsafe(self):
    a = Tensor.ones((3, )).contiguous().realize()
    out = a.exp2().shrink(((0, 1),)).pad(((0, 1),)).contiguous()
-    run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 1))
    np.testing.assert_equal(out.numpy(), [2, 0])

  def test_base_change_shrink_pad(self):
@@ -1545,7 +1483,7 @@ class TestSchedule(unittest.TestCase):
    b = a.exp2()
    c = b[:-1, :-1]
    d = c.pad(((0, 1), (0, 1))) * 2
-    run_schedule(check_schedule(d, 2))
+    run_schedule(check_schedule(d, 1))
    np.testing.assert_equal(d.numpy(), np.pad(np.exp2(a.numpy())[:-1, :-1], ((0, 1), (0, 1)))*2)

  def test_base_change_expand_pad(self):
@@ -1553,14 +1491,14 @@ class TestSchedule(unittest.TestCase):
    b = a.exp2()
    c = b[:, None, :]
    d = c.pad(((0, 0), (1, 1), (0, 0))) * 2
-    run_schedule(check_schedule(d, 2))
+    run_schedule(check_schedule(d, 1))
    np.testing.assert_equal(d.numpy(), np.pad(np.exp2(a.numpy())[:, None, :], ((0, 0), (1, 1), (0, 0)))*2)

  def test_fuse_arange_pad_replicate_mode(self):
    x = Tensor.empty(3,3,3,3, requires_grad=True)
    y = x.pad((-1,2,2,-1), mode="replicate")
    dx = y.sum().gradient(x)[0]
-    sched = check_schedule(dx, 3)
+    sched = check_schedule(dx, 1)
    run_schedule(sched)
    np.testing.assert_allclose(dx.numpy(), [[[[0.,3.,9.],[0,1.,3.],[0.,0.,0.]]]*3]*3)

@@ -1570,7 +1508,7 @@ class TestSchedule(unittest.TestCase):
    a = Tensor.ones(4, 4).contiguous().realize()
    b = a.cast(dtypes.half).expand(2, 4, 4)
    c = b.cast(dtypes.int).expand(2, 2, 4, 4)
-    run_schedule(check_schedule(c, 2))
+    run_schedule(check_schedule(c, 1))
    np.testing.assert_equal(c.numpy(), np.ones(((2, 2, 4, 4)), dtype=np.int32))

  def test_base_change_pad_expand(self):
@@ -1578,7 +1516,7 @@ class TestSchedule(unittest.TestCase):
    b = Tensor.full((4, 4), 2.).contiguous().realize()
    c = (a + b).pad(((1, 1), (1, 1)))
    d = c.cast(dtypes.int).expand((2, 6, 6)) * 4
-    run_schedule(check_schedule(d, 2))
+    run_schedule(check_schedule(d, 1))
    c_np = np.pad((np.full((4, 4), 2., dtype=np.float32) + np.full((4, 4), 1., dtype=np.float32)), ((1, 1), (1, 1)), constant_values=0.0)
    np.testing.assert_equal(d.numpy(), np.broadcast_to(c_np.astype(np.half), (2, *c_np.shape)) * 4)

@@ -1677,7 +1615,7 @@ class TestSchedule(unittest.TestCase):
    self._test_fusion([(4, 4), (1, 4)], lambda a,b:a.sum(1).reshape(b.shape)+b, 1)

  def test_late_fusion_post_permute(self):
-    self._test_fusion([(4, 6, 4), (4, 4, 1)], lambda a,b:a.sum(1, keepdim=True).permute((2, 0, 1))+b, 2)
+    self._test_fusion([(4, 6, 4), (4, 4, 1)], lambda a,b:a.sum(1, keepdim=True).permute((2, 0, 1))+b, 1)

  def test_late_fusion_double_transpose(self):
    self._test_fusion([(32, 16, 1)],
@@ -1715,6 +1653,7 @@ class TestSchedule(unittest.TestCase):
    self.assertListEqual(realized_const_view.tolist(), [[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1]])

  @given(strat.sampled_from(dtypes.all), strat.sampled_from(dtypes.all))
+  @unittest.skip("kernel count depends on input")
  def test_cast_padded_const(self, dt1, dt2):
    assume(is_dtype_supported(dt1) and is_dtype_supported(dt2))
    a = Tensor(1, dtype=dt1).reshape(1, 1).pad(((1, 1), None))
@@ -1728,7 +1667,7 @@ class TestSchedule(unittest.TestCase):
    X = Tensor.randn(10, 10).realize()
    idxs = Tensor([0, 2]).realize()
    xt = X[idxs]
-    run_schedule(check_schedule(xt, 2))
+    run_schedule(check_schedule(xt, 1))
    np.testing.assert_equal(xt.numpy(), X.numpy()[idxs.numpy()])

  def test_simple_indexing_alt(self):
@@ -1746,7 +1685,7 @@ class TestSchedule(unittest.TestCase):
  def test_advanced_indexing_alt(self):
    X = Tensor.arange(6).reshape(3, 2)+1
    xt = X[[Tensor([2]), Tensor([1])]]
-    run_schedule(check_schedule(xt, 3))
+    run_schedule(check_schedule(xt, 1))
    np.testing.assert_equal(xt.numpy(), 6)

  def test_advanced_simple_indexing_combined(self):
@@ -1794,7 +1733,7 @@ class TestSchedule(unittest.TestCase):
    x = Tensor.full((2,2), 16)
    y = x.idiv(Tensor.linspace(2, 8, steps=4, dtype=dtypes.int).reshape(2,2)).pad(((1,1), (1,1)))
    out = y.sum(axis=1)
-    run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 1))
    self.assertListEqual(out.tolist(), [0, 12, 4, 0])

  def test_arange_transposed_descendants(self):
@@ -1827,7 +1766,7 @@ class TestSchedule(unittest.TestCase):
    x = Tensor.randn(5, 2).realize()
    a = Tensor.arange(10).contiguous()
    out = (x + a[2]).sum()
-    run_schedule(check_schedule(out, 3))
+    run_schedule(check_schedule(out, 2))
    np.testing.assert_allclose(out.numpy(), (x.numpy()+np.arange(10)[2]).sum(), atol=1e-5, rtol=1e-6)

  def test_arange_index_child(self):
@@ -1843,7 +1782,7 @@ class TestSchedule(unittest.TestCase):
    x = Tensor.randn(5, 2).realize()
    a = (Tensor.arange(10)+1).contiguous()
    out = (x + a[2]).sum()
-    run_schedule(check_schedule(out, 3))
+    run_schedule(check_schedule(out, 2))
    np.testing.assert_allclose(out.numpy(), (x.numpy()+(np.arange(10)+1)[2]).sum(), atol=1e-5, rtol=1e-6)

  @unittest.skip("BUFFER_VIEW no longer supported on non-disk devices")
@@ -1858,10 +1797,10 @@ class TestSchedule(unittest.TestCase):
    from extra.models.llama import precompute_freqs_cis
    args = {"dim":32 if CI else 128, "end":2048 if CI else 8192, "theta":10000}
    fused = precompute_freqs_cis(**args)
-    run_schedule(check_schedule(fused, 3))
+    run_schedule(check_schedule(fused, 1))
    if getenv("CHECK", 1):
      ref = precompute_freqs_cis(**args)
-      run_schedule(check_schedule(ref, 3))
+      run_schedule(check_schedule(ref, 1))
      np.testing.assert_equal(fused.numpy(), ref.numpy())

  def test_fuse_assign_contiguous(self):
@@ -1903,7 +1842,7 @@ class TestSchedule(unittest.TestCase):
    X = Tensor([[0, 2, 3], [1, 2, 3]]).realize()
    Y = Tensor([1, 2]).realize()
    loss = X.sparse_categorical_crossentropy(Y)
-    run_schedule(check_schedule(loss, 4))
+    run_schedule(check_schedule(loss, 3))
    np.testing.assert_allclose(loss.item(), 0.878309, atol=1e-5, rtol=1e-6)

  def test_const_folding_alt(self):
@@ -1924,7 +1863,7 @@ class TestSchedule(unittest.TestCase):
    yt = Tensor.randn(BS, 10).realize()
    with Context(SPLIT_REDUCEOP=0):
      loss = yt.sparse_categorical_crossentropy(Y_train[samples])
-      run_schedule(check_schedule(loss, 6))
+      run_schedule(check_schedule(loss, 5))
      loss_fused = loss.numpy()
    loss_ref = torch.nn.CrossEntropyLoss()(torch.tensor(yt.numpy()), torch.tensor(Y_train.numpy())[torch.tensor(samples.numpy())])
    np.testing.assert_allclose(loss_fused, loss_ref.numpy(), atol=1e-6, rtol=1e-6)
@@ -1934,7 +1873,7 @@ class TestSchedule(unittest.TestCase):
    r = (X+Tensor.arange(16).reshape(4, 4)).sum()
    out0 = r+2
    out1 = r+3
-    run_schedule(check_schedule([out0, out1], 1))
+    run_schedule(check_schedule([out0, out1], 2)) # TODO: 1?
    r_ref = (X.numpy()+np.arange(16).reshape(4, 4)).sum()
    np.testing.assert_allclose(out0.numpy(), r_ref+2, rtol=2e-7)
    np.testing.assert_allclose(out1.numpy(), r_ref+3, rtol=2e-7)
@@ -1976,8 +1915,7 @@ class TestSwizzle(unittest.TestCase):
      a = Tensor.randint(32, 32).realize()
    r = (a+a).sum(1).sum(0)
    # double reduce collapses to a single reduce
-    with Context(DONT_GROUP_REDUCES=1):
-      run_schedule(check_schedule(r, 1))
+    run_schedule(check_schedule(r, 1))
    self.assertEqual(r.numpy(), (a.numpy()+a.numpy()).sum(1).sum(0))

  def test_single_swizzle(self):
@@ -1997,33 +1935,29 @@ class TestSwizzle(unittest.TestCase):
      b = Tensor.randint(4,).realize()
    # parallel reduce!
    add = a.sum(0)+b.sum(0)
-    with Context(DONT_GROUP_REDUCES=1):
-      run_schedule(check_schedule(add, 1))
+    run_schedule(check_schedule(add, 1))
    self.assertEqual(add.numpy(), a.numpy().sum(0)+b.numpy().sum(0))

-  @unittest.skip("TODO: how do we express the norm")
  def test_softmax_one_kernel(self):
    Tensor.manual_seed(0)
    with Context(DEBUG=0, TRACK_MATCH_STATS=0):
      a = Tensor.randn(32, 32).realize()
    t = a.softmax()
-    with Context(DONT_GROUP_REDUCES=1, DONT_REALIZE_EXPAND=1):
-      check_schedule(t, 1)
+    check_schedule(t, 3) # TODO: 1?

  def test_argmax_one_kernel(self):
    Tensor.manual_seed(0)
    with Context(DEBUG=0, TRACK_MATCH_STATS=0):
      a = Tensor.randn(10, 20).realize()
    t = a.argmax(0)
-    with Context(DONT_GROUP_REDUCES=1, DONT_REALIZE_EXPAND=1): t.realize()
+    check_schedule(t, 2) # TODO: 1?

  def test_swizzle_reduceop(self):
    Tensor.manual_seed(0)
    x = Tensor.randn(4,4).realize()
    y = Tensor.randn(4,4,4).realize()
    out = x.reshape(4,4,1).expand(4,4,4).sum(axis=(1,))+y
-    with Context(DONT_REALIZE_EXPAND=1, DONT_GROUP_REDUCES=1):
-      run_schedule(check_schedule(out, 1))
+    run_schedule(check_schedule(out, 2)) # TODO: 1?
    np.testing.assert_allclose(out.numpy(), np.tile(x.numpy().reshape(4,4,1), (1,1,4)).sum(axis=1)+y.numpy())

  def test_permute_rewrite(self):
@@ -2031,7 +1965,7 @@ class TestSwizzle(unittest.TestCase):
    y = Tensor.randn(4, 1, 16).realize()
    z = Tensor.randn(4, 4, 1).realize()
    t = (x*y).sum(axis=(0, 2)).reshape(1, 4, 1).permute(0, 2, 1)+z
-    with Context(DONT_GROUP_REDUCES=1, DONT_REALIZE_EXPAND=1): run_schedule(check_schedule(t, 1))
+    run_schedule(check_schedule(t, 2)) # TODO: 1?
    t_np = (x.numpy()*y.numpy()).sum(axis=(0, 2)).reshape(1, 4, 1).transpose(0, 2, 1)+z.numpy()
    np.testing.assert_allclose(t.numpy(), t_np, atol=1e-6, rtol=1e-3)

@@ -2042,14 +1976,14 @@ class TestSwizzle(unittest.TestCase):
    a_reduce = a.sum(axis=(2,), keepdim=True).sum(axis=(1,))
    b_reduce = b.sum(axis=(0,))
    t = a_reduce+b_reduce
-    with Context(DONT_GROUP_REDUCES=1, DONT_REALIZE_EXPAND=1): run_schedule(check_schedule(t, 1))
+    run_schedule(check_schedule(t, 1))

  def test_parallel_reduce_possible(self):
    Tensor.manual_seed(0)
    x = Tensor.randn(4, 2, 2).realize()
    y = Tensor.randn(4, 2, 2).realize()
    t = x.sum(axis=1)+y.sum(axis=1)
-    with Context(DONT_GROUP_REDUCES=1): run_schedule(check_schedule(t, 1))
+    run_schedule(check_schedule(t, 1))
    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
@@ -2058,7 +1992,7 @@ class TestSwizzle(unittest.TestCase):
    x = Tensor.randn(4, 2, 2).realize()
    y = Tensor.randn(4, 3, 2).realize()
    t = x.sum(axis=1)+y.sum(axis=1)
-    with Context(DONT_GROUP_REDUCES=1): run_schedule(check_schedule(t, 1))
+    run_schedule(check_schedule(t, 1))
    np.testing.assert_allclose(t.numpy(), x.numpy().sum(axis=1)+y.numpy().sum(axis=1), atol=1e-6, rtol=1e-3)

  def test_unsafe_pad(self):
@@ -2151,7 +2085,7 @@ class TestCopyFolding(unittest.TestCase):
    a = Tensor.arange(3).realize()
    zeros = Tensor.zeros(3).realize()
    b = (a*zeros).to("CPU")
-    run_schedule(check_schedule(b, 0, filter_sink=False))
+    run_schedule(check_schedule(b, 2, filter_sink=False)) # TODO: 0?
    self.assertListEqual(b.tolist(), [0, 0, 0])
    self.assertEqual(b.device, "CPU")

@@ -2171,12 +2105,12 @@ class TestCopyFolding(unittest.TestCase):
  def test_copy_to_same_device(self):
    a = Tensor.empty(4).uop
    b = a.copy_to_device(a.device)
-    check_schedule(b, 0, filter_sink=False)
+    check_schedule(b, 1, filter_sink=False) # TODO: 0?

  def test_copy_to_same_device_alt(self):
    a = Tensor.empty(4, 4).uop
    b = a.copy_to_device(a.device)
-    check_schedule(b, 0, filter_sink=False)
+    check_schedule(b, 1, filter_sink=False) # TODO: 0?

  def test_copy_to_same_device_sched(self):
    a = Tensor.ones(4).contiguous().realize().uop.as_buf()
@@ -2191,13 +2125,11 @@ class TestCopyFolding(unittest.TestCase):
    a = Tensor.empty(4)
    check_schedule(a.clone(), 1, filter_sink=False)

-  # NOTE: moving copy before view might change this
  def test_shrink_copy(self):
    a = Tensor.arange(4)
    view = a.shrink(((0, 2),))
    b = view.clone()
-    # NOTE: this was sort of a bug making this 2
-    run_schedule(check_schedule(b, 2, filter_sink=False))
+    run_schedule(check_schedule(b, 1, filter_sink=False))
    self.assertEqual(b.uop.base.buffer.size, 2)
    self.assertEqual(b.uop.size, 2)
    self.assertListEqual(b.tolist(), [0, 1])
@@ -2206,7 +2138,7 @@ class TestCopyFolding(unittest.TestCase):
    a = Tensor.arange(2)
    view = a.reshape(2, 1).expand(2, 2)
    b = view.clone()
-    run_schedule(check_schedule(b, 2, filter_sink=False))
+    run_schedule(check_schedule(b, 1, filter_sink=False))
    self.assertEqual(b.uop.base.buffer.size, 4)
    self.assertEqual(b.uop.size, 4)
    self.assertListEqual(b.tolist(), [[0, 0], [1, 1]])
@@ -2329,7 +2261,7 @@ class TestContiguous(unittest.TestCase):
  def test_double_contiguous_realizes_once(self):
    a = Tensor.empty(4, 1)
    b = a.expand((4, 4)).contiguous().contiguous()
-    check_schedule(b, 1)
+    check_schedule(b, 2) # TODO: should be 1?

  def test_view_does_not_realize(self):
    a = Tensor.empty(4)
@@ -2465,10 +2397,6 @@ class TestUOpBecome(unittest.TestCase):
    c = (a.reshape(1, 1, 4, 4)+0).shrink(((0, 1), (0, 1), (0, 3), (0, 3)))+0
    check_schedule([b, c], 0)
    assert all_same([x.uop.base.realized for x in [a,b,c]])
-    # these movement ops result in the same ShapeTracker
-    assert b.uop.st == c.uop.st
-    assert b.uop is c.uop
-    assert UPat(Ops.VIEW, src=(UPat(Ops.BUFFER),)).match(c.uop, {})

  def test_setitem_becomes_subbuffer(self):
    a = Tensor.full((4,), 2.).contiguous().realize()
@@ -165,8 +165,7 @@ class TestSoftmaxFusion(unittest.TestCase):
      sout.realize()

    print("*** single kernel softmax ***")
-    # NOTE: DONT_GROUP_REDUCES is required here
-    with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2), DONT_GROUP_REDUCES=1):
+    with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
      out = single_kernel_softmax(self.test)
      out.realize()

@@ -186,7 +185,6 @@ class TestSoftmaxFusion(unittest.TestCase):

    np.testing.assert_allclose(sout.numpy(), out.numpy(), atol=3e-7)

-  @unittest.skip("recursion error no longer raised")
  def test_softmax_bw(self):
    print("*** softmax bw ***")
    self.test.requires_grad_()
@@ -197,14 +195,11 @@ class TestSoftmaxFusion(unittest.TestCase):
    self.test.grad = None

    print("*** single kernel softmax bw ***")
-    # NOTE: DONT_GROUP_REDUCES is required here
-    # TODO: fix RecursionError with DONT_GROUP_REDUCES
-    with self.assertRaises(RecursionError):
-      with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2), DONT_GROUP_REDUCES=1):
-        single_kernel_softmax(self.test).sum().backward()
-        g = self.test.grad.realize()
+    with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
+      single_kernel_softmax(self.test).sum().backward()
+      g = self.test.grad.realize()

-      np.testing.assert_allclose(sg.numpy(), g.numpy(), atol=1e-7)
+    np.testing.assert_allclose(sg.numpy(), g.numpy(), atol=1e-7)

 if __name__ == '__main__':
  unittest.main()
@@ -810,6 +810,13 @@ class TestTensorMetadata(unittest.TestCase):
    self.assertEqual(len(si.metadata), 1)
    self.assertEqual(si.metadata[0].name, "relu")

+  def test_assign(self):
+    x = Tensor.empty(10, 10).realize()
+    x.assign(Tensor.ones(10, 10).contiguous())
+    si = x.schedule()[-1]
+    self.assertEqual(len(si.metadata), 1)
+    self.assertEqual(si.metadata[0].name, "assign")
+
  def test_complex(self):
    x = Tensor.rand(3, requires_grad=True)
    y = Tensor.rand(3, requires_grad=True)
@@ -418,5 +418,32 @@ class TestPathTensor(unittest.TestCase):
      Tensor(pathlib.Path(test_file)).tolist()
    os.chmod(test_file, 0o644)
    assert Tensor(pathlib.Path(test_file)).tolist(), list(range(10))
+
+class TestDiskTensorMovement(unittest.TestCase):
+  def setUp(self):
+    self.fn = pathlib.Path(temp("custom_disk_range"))
+    self.fn.unlink(missing_ok=True)
+    Tensor.arange(100, dtype=dtypes.uint8).to(f"disk:{str(self.fn)}").realize()
+
+  def test_simple_read(self):
+    t = Tensor(self.fn)
+    self.assertTrue(Tensor.all(t.to(None) == Tensor.arange(100, dtype=dtypes.uint8)).item())
+
+  def test_slice_read(self):
+    t = Tensor(self.fn)
+    self.assertListEqual(t[16:18].tolist(), [16,17])
+
+  # TODO: fix this! at least assert on it
+  @unittest.expectedFailure
+  def test_slice_read_cat(self):
+    t = Tensor(self.fn)
+    self.assertListEqual(Tensor.cat(t[16:18], t[20:22]).tolist(), [16,17,20,21])
+
+  # TODO: fix this! at least assert on it
+  @unittest.expectedFailure
+  def test_slice_sum(self):
+    t = Tensor(self.fn)
+    self.assertListEqual((t[16:18]+t[20:22]).tolist(), [16+20,17+21])
+
 if __name__ == "__main__":
  unittest.main()
@@ -442,7 +442,7 @@ class TestVizMemoryLayout(BaseTestViz):
    profile_ret = load_profile(Buffer.profile_events)
    ret = profile_ret["layout"][f"{a.device} Memory"]
    self.assertEqual(ret["peak"], 2)
-    self.assertEqual(len(ret["events"]), 2)
+    self.assertEqual(len(ret["events"]), 4)

  def test_del_once(self):
    a = _alloc(1)
@@ -451,7 +451,7 @@ class TestVizMemoryLayout(BaseTestViz):
    profile_ret = load_profile(Buffer.profile_events)
    ret = profile_ret["layout"][f"{b.device} Memory"]
    self.assertEqual(ret["peak"], 1)
-    self.assertEqual(len(ret["events"]), 3)
+    self.assertEqual(len(ret["events"]), 4)

  def test_alloc_free(self):
    a = _alloc(1)
@@ -461,7 +461,7 @@ class TestVizMemoryLayout(BaseTestViz):
    profile_ret = load_profile(Buffer.profile_events)
    ret = profile_ret["layout"][f"{c.device} Memory"]
    self.assertEqual(ret["peak"], 2)
-    self.assertEqual(len(ret["events"]), 4)
+    self.assertEqual(len(ret["events"]), 6)

  def test_free_last(self):
    bufs = []
@@ -480,15 +480,24 @@ class TestVizMemoryLayout(BaseTestViz):
    self.assertEqual(len(profile["markers"]), 6)

  def test_producer_simple(self):
-    a = Tensor.empty(10, device="NULL")
-    Tensor.realize(a.add(1), a.add(2))
-    b = Tensor.empty(10, device="NULL")
-    Tensor.realize(b.add(1))
+    a = Tensor.ones(10, device="NULL")
+    Tensor.realize(a.add(1).contiguous())
+    b = Tensor.ones(10, device="NULL")
+    Tensor.realize(b.add(1).contiguous())
    profile = load_profile(cpu_events+Buffer.profile_events)
    buffers = profile["layout"]["NULL Memory"]["events"]
    programs = profile["layout"]["NULL"]["events"]
    user_cnt = [len(b["arg"]["users"]) for b in buffers if b["arg"].get("users")]
    self.assertEqual(len(user_cnt), len(programs))

+  def test_inflight_buf(self):
+    a = Tensor.empty(1, device="NULL")
+    n = 4
+    for i in range(n): (a+i).realize()
+    profile = load_profile(cpu_events+Buffer.profile_events)
+    buffers = profile["layout"]["NULL Memory"]["events"]
+    user_cnt = [len(b["arg"]["users"]) for b in buffers if b["arg"].get("users")]
+    self.assertEqual(max(user_cnt), n)
+
 if __name__ == "__main__":
  unittest.main()
@@ -42,7 +42,7 @@ def get_rewrites_for_renderer(opts:Renderer, optimize:bool=True, linearizer:bool

@functools.cache
 def _get_rewrites_for_renderer(opts:Renderer, optimize:bool, linearizer:bool, _QUANTIZE, _DEVECTORIZE, _TRANSCENDENTAL) -> list[RewriteStep]:
-  # ** lowerer (rewrite_shapetracker_with_index) **
+  # ** lowerer **
  ret: list[RewriteStep] = []

  if optimize:
@@ -17,9 +17,9 @@ class Opt:
  def __repr__(self): return f"Opt(op={self.op}, axis={self.axis}, arg={self.arg})"

 axis_letters = {AxisType.GLOBAL: "g", AxisType.THREAD: "t", AxisType.LOCAL: "l", AxisType.WARP: "w", AxisType.LOOP: "L", AxisType.UPCAST: "u",
-                AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r"}
+                AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r", AxisType.MULTI: "m"}
 axis_colors = {AxisType.GLOBAL: "blue", AxisType.THREAD: "BLUE", AxisType.LOCAL: "cyan", AxisType.WARP: "CYAN", AxisType.LOOP: "WHITE",
-               AxisType.UPCAST: "yellow", AxisType.GROUP_REDUCE: "RED", AxisType.REDUCE: "red", AxisType.UNROLL: "magenta"}
+               AxisType.UPCAST: "yellow", AxisType.GROUP_REDUCE: "RED", AxisType.REDUCE: "red", AxisType.UNROLL: "magenta", AxisType.MULTI: "GREEN"}

 class KernelOptError(Exception): pass
 def check(cond:bool, msg:str=""):
@@ -13,8 +13,8 @@ from tinygrad.renderer import Renderer
 remove_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])

 # NOTE: LOCAL and GROUP_REDUCE have the same priority. the order here matters
-axis_to_pos = {AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2, AxisType.UPCAST: 3,
-               AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}
+axis_to_pos = {AxisType.MULTI: -2, AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2,
+               AxisType.UPCAST: 3, AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}

 class Scheduler:
  def __init__(self, ast:UOp, opts:Renderer):
@@ -97,7 +97,7 @@ class Scheduler:
    new_rng = UOp.range(amount, self.maxarg+1, new_type) if input_new_rng is None else input_new_rng
    replaced_rng = rng.replace(src=(UOp.const(dtypes.int, old_sz),))
    sub_axis = (new_rng * old_sz + replaced_rng) if top else (replaced_rng * amount + new_rng)
-    self.ast = self.ast.substitute({rng:sub_axis}, name=f"shift {rng.arg[0]} {amount} {str(new_type).split('.')[1].lower()}")
+    self.ast = self.ast.substitute({rng:sub_axis}, name=f"shift {rng.arg[:-1]} {amount} {str(new_type).split('.')[1].lower()}")
    return replaced_rng, new_rng

  def ranges_of(self, *axis_type:AxisType) -> list[UOp]: return [r for r in self.rngs if r.arg[-1] in axis_type]
@@ -200,13 +200,14 @@ class Scheduler:
      self.ast = self.ast.substitute(replaces, f"padto {rng.arg[:-1]} {opt.arg}")
    elif opt.op is OptOps.SWAP:
      try:
-        altrng = self.rngs[opt.arg]
+        altrng:UOp = self.rngs[opt.arg]
      except IndexError:
        raise KernelOptError
      check(rng.arg[-1] == AxisType.GLOBAL and altrng.arg[-1] == AxisType.GLOBAL, "swap only for globals")
      self.ast = self.ast.substitute({rng:rng.replace(arg=(*altrng.arg[0:-1], rng.arg[-1]), tag=1),
-                                      altrng:altrng.replace(arg=(*rng.arg[0:-1], altrng.arg[-1]), tag=1)})
-      self.ast = graph_rewrite(self.ast, remove_tags)
+                                      altrng:altrng.replace(arg=(*rng.arg[0:-1], altrng.arg[-1]), tag=1)},
+                                      name=f"swap {rng.arg[:-1]} {altrng.arg[:-1]}")
+      self.ast = graph_rewrite(self.ast, remove_tags, name="swap remove tags")
    else:
      raise KernelOptError(f"unsupported opt {opt.op}")

@@ -91,20 +91,20 @@ pm_reduce_collapse = PatternMatcher([
  # fold the range
  ((UPat(Ops.RANGE, name="r") < UPat.var("cut")).where(0, UPat.cvar("val")).reduce(UPat.var("r"), arg=Ops.ADD),
   lambda r,cut,val: (r.src[0]-cut).maximum(0).minimum(r.src[0]).cast(val.dtype) * val),
+  (((UPat.var("r")<UPat.var("lower")).logical_not()&(UPat(Ops.RANGE, name="r")<UPat.var("upper"))).where(UPat.cvar("val"), 0).reduce(UPat.var("r"),
+    arg=Ops.ADD), lambda r,lower,upper,val: (upper.minimum(r.src[0])-lower.maximum(0)).maximum(0).minimum(r.src[0]).cast(val.dtype) * val),
  ((UPat(Ops.RANGE, name="r") < UPat.var("cut")).where(UPat.cvar("val"), 0).reduce(UPat.var("r"), arg=Ops.ADD),
   lambda r,cut,val: cut.maximum(0).minimum(r.src[0]).cast(val.dtype) * val),
  # REDUCE on ADD
  ((UPat.var("x")+UPat.var("y")).reduce(arg=Ops.ADD, allow_any_len=True, name="r"),
   lambda x,y,r: x.reduce(*r.src[1:], arg=Ops.ADD) + y.reduce(*r.src[1:],arg=Ops.ADD)),
  # MUL casted bool
-  ((UPat.var("x") * UPat.var("gate", dtype=dtypes.bool).cast().or_broadcasted(name="b")),
-   lambda x,gate,b=None: gate.broadcast(x.dtype.count).where(x, 0) if b is not None else gate.where(x, 0)),
+  ((UPat.var("x") * UPat.var("gate", dtype=dtypes.bool).cast()), lambda x,gate: gate.where(x, 0)),
  # reduce on gated load becomes can substitute the range and remove the reduce
  ((UPat.var("idx")!=(UPat(Ops.RANGE, name="r").or_casted())).where(0, UPat.var("expr")).reduce(UPat.var("r"), arg=Ops.ADD),
   lambda r,idx,expr: (v:=(idx.cast(r.dtype) >= 0) & (idx.cast(r.dtype) < r.src[0])).where(expr.substitute({r:idx.cast(r.dtype).valid(v)}),0)),
  # AND on WHERE
-  ((UPat.any(UPat(Ops.DEFINE_VAR, name="x"), UPat(Ops.DEFINE_VAR).gep(name="x")) & UPat.var("y")) \
-   .where(UPat.cvar("c"), 0).reduce(arg=Ops.ADD, allow_any_len=True, name="r"),
+  ((UPat(Ops.DEFINE_VAR, name="x") & UPat.var("y")).where(UPat.cvar("c"), 0).reduce(arg=Ops.ADD, allow_any_len=True, name="r"),
    lambda x,y,c,r: y.where(c, 0).reduce(*r.src[1:], arg=Ops.ADD)*x.cast(c.dtype)),
  # remove REDUCEs that no longer have a RANGE in the src
  (UPat(Ops.REDUCE, name="red"), reduce_rangeless),
@@ -4,13 +4,13 @@ from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, all_metadata
 from tinygrad.helpers import argsort

 def reduce_gradient(ctx:UOp, ret:UOp):
-  def to_inp_shape(x): return x.reshape(x.shape+(1,)*(len(ret.src[0].shape)-len(x.shape))).expand(ret.src[0].shape)
-  if ret.arg[0] == Ops.ADD: return (to_inp_shape(ctx),)
+  def broadcast_to_input(x): return x.reshape(x.shape+(1,)*(len(ret.src[0].shape)-len(x.shape))).expand(ret.src[0].shape)
+  if ret.arg[0] == Ops.ADD: return (broadcast_to_input(ctx),)
  if ret.arg[0] == Ops.MAX:
-    max_is_1s = ret.src[0].eq(to_inp_shape(ret)).cast(ctx.dtype)
-    div = to_inp_shape(max_is_1s.r(Ops.ADD, ret.arg[1]))
-    return ((max_is_1s/div) * to_inp_shape(ctx),)
-  if ret.arg[0] == Ops.MUL: return (to_inp_shape(ctx * ret) / ret.src[0],)
+    mask = ret.src[0].eq(broadcast_to_input(ret)).cast(ctx.dtype)
+    count = mask.r(Ops.ADD, ret.arg[1])
+    return ((mask/broadcast_to_input(count)) * broadcast_to_input(ctx),)
+  if ret.arg[0] == Ops.MUL: return (broadcast_to_input(ctx * ret) / ret.src[0],)

 # ctx is grad_output
 pm_gradient = PatternMatcher([
@@ -60,5 +60,9 @@ def compute_gradient(root:UOp, root_grad:UOp, targets:set[UOp]) -> dict[UOp, UOp
      if v is None: continue
      if k in grads: grads[k] = grads[k] + v
      else: grads[k] = v
-      if len(forward_metadata:=all_metadata.get(t0, ())): all_metadata[v] = tuple(dataclasses.replace(x, backward=True) for x in forward_metadata)
+      if len(forward_metadata:=all_metadata.get(t0, ())):
+        backward_metadata = tuple(dataclasses.replace(x, backward=True) for x in forward_metadata)
+        # we add the backward metadata to everything new in the graph
+        for bw_uop in v.toposort(lambda x: x not in (t0, *t0.src, grads[t0])):
+          all_metadata[bw_uop] = all_metadata.get(bw_uop, ())+backward_metadata
  return grads
@@ -158,7 +158,6 @@ SPLIT_REDUCEOP, NO_MEMORY_PLANNER, RING = ContextVar("SPLIT_REDUCEOP", 1), Conte
 PICKLE_BUFFERS, LRU = ContextVar("PICKLE_BUFFERS", 1), ContextVar("LRU", 1)
 CACHELEVEL, IGNORE_BEAM_CACHE, DEVECTORIZE = ContextVar("CACHELEVEL", 2), ContextVar("IGNORE_BEAM_CACHE", 0), ContextVar("DEVECTORIZE", 1)
 DISABLE_COMPILER_CACHE, BLOCK_REORDER = ContextVar("DISABLE_COMPILER_CACHE", 0), ContextVar("BLOCK_REORDER", 1)
-DONT_REALIZE_EXPAND, DONT_GROUP_REDUCES = ContextVar("DONT_REALIZE_EXPAND", 0), ContextVar("DONT_GROUP_REDUCES", 0)
 QUANTIZE, VALIDATE_WITH_CPU, DISABLE_FAST_IDIV = ContextVar("QUANTIZE", 0), ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("DISABLE_FAST_IDIV", 0)
 CORRECT_DIVMOD_FOLDING, FUSE_OPTIM = ContextVar("CORRECT_DIVMOD_FOLDING", 0), ContextVar("FUSE_OPTIM", 0)
 ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE = ContextVar("ALLOW_DEVICE_USAGE", 1), ContextVar("MAX_BUFFER_SIZE", 0)
@@ -50,7 +50,7 @@ class Optimizer:
    if self.fused:
      # optimizer fusion just concatenates all the buffers, runs the _step, then splits them back up
      out, extra = self._step([Tensor.cat(*[t.flatten() for t in self.params], dim=0)],
-                              [Tensor.cat(*[unwrap(t.grad).flatten() for t in self.params], dim=0)])
+                              [Tensor.cat(*[unwrap(t.grad).contiguous().flatten() for t in self.params], dim=0)])
      updated_params = [out[0][self.pos_params[i]:self.pos_params[i+1]].reshape(tt.shape) for i, tt in enumerate(self.params)]
    else:
      updated_params, extra = self._step(self.params, [unwrap(t.grad) for t in self.params])
@@ -1,6 +1,6 @@
 import json, pathlib, zipfile, pickle, tarfile, struct, functools, io
 from collections import OrderedDict
-from typing import Any, Callable, BinaryIO, Iterable
+from typing import Any, Callable, BinaryIO, Iterable, cast
 from tinygrad.tensor import Tensor
 from tinygrad.dtype import dtypes
 from tinygrad.helpers import prod, argsort, DEBUG, Timing, CI, unwrap, GlobalCounters, tqdm, round_up, T, strides_for_shape
@@ -209,7 +209,7 @@ def torch_load(t:Tensor) -> dict[str, Tensor]:
      assert tuple([shape_strides[i][1] for i in argsort(permute_indexes)]) == strides_for_shape(intermediate_shape), "nonpermutable strides"
      if DEBUG >= 3: print(f"WARNING: this torch load is slow. to permute {intermediate_shape} with {permute_indexes}")
      assert storage[1] != dtypes.bfloat16, "can't permute BF16"
-      # TODO: find a nice way to support all shapetracker on disktensors
+      # TODO: find a nice way to support all movement ops on disktensors
      ret = ret.to(None).reshape(intermediate_shape).permute(permute_indexes)

    return ret.reshape(size)
@@ -237,11 +237,18 @@ def torch_load(t:Tensor) -> dict[str, Tensor]:

  if passthrough_reset(zipfile.is_zipfile(fobj)): # NOTE: passthrough_reset required to support python < 3.14
    myzip = zipfile.ZipFile(fobj, 'r')
-    base_name = myzip.namelist()[0].split('/', 1)[0]
-    for n in myzip.namelist():
-      if n.startswith(f'{base_name}/data/'):
-        with myzip.open(n) as myfile:
-          offsets[n.split("/")[-1]] = myfile._orig_compress_start # type: ignore
+    base_name = None
+    header_offsets = {}
+    for zi in myzip.filelist:
+      if base_name is None: base_name = zi.filename.split('/', 1)[0]
+      if zi.filename.startswith(f'{base_name}/data/'): header_offsets[zi.filename.split("/")[-1]] = zi.header_offset
+    # sadly there's no way to get the start of the file in the zip without reading the header
+    # at least here we read them in parallel
+    header_contents = [t[v+26:v+30].bitcast(dtypes.uint16).to('CPU') for v in header_offsets.values()]
+    Tensor.realize(*header_contents)
+    for (n,o),c in zip(header_offsets.items(), header_contents):
+      # header_offset + sizeFileHeader + File name length + Extra field length : https://en.wikipedia.org/wiki/ZIP_(file_format)
+      offsets[n] = o+30+sum(cast(list[int], c.tolist()))
    with myzip.open(f'{base_name}/data.pkl') as myfile:
      return TorchPickle(myfile).load()
  elif passthrough_reset(tarfile.is_tarfile(fobj)): # NOTE: passthrough_reset required to support python < 3.11
@@ -6,7 +6,7 @@
 # POINTER_SIZE is: 8
 # LONGDOUBLE_SIZE is: 16
 #
-import ctypes, tinygrad.runtime.support.llvm as llvm_support, tinygrad.helpers as helpers
+import ctypes, tinygrad.runtime.support.llvm as llvm_support


 class AsDictMixin:
@@ -146,7 +146,7 @@ class FunctionFactoryStub:
 # You can either re-run clan2py with -l /path/to/library.so
 # Or manually fix this by comment the ctypes.CDLL loading
 _libraries = {}
-_libraries['llvm'] = ctypes.CDLL(llvm_support.LLVM_PATH, ctypes.RTLD_GLOBAL if helpers.OSX else ctypes.DEFAULT_MODE) #  ctypes.CDLL('llvm')
+_libraries['llvm'] = ctypes.CDLL(llvm_support.LLVM_PATH) #  ctypes.CDLL('llvm')
 c_int128 = ctypes.c_ubyte*16
 c_uint128 = c_int128
 void = None
@@ -111,7 +111,7 @@ class NVCommandQueue(HWQueue[HCQSignal, 'NVDevice', 'NVProgram', 'NVArgsState'])
  def _submit_to_gpfifo(self, dev:NVDevice, gpfifo:GPFifo):
    if dev == self.binded_device: cmdq_addr = self.hw_page.va_addr
    else:
-      cmdq_addr = dev.cmdq_allocator.alloc(len(self._q) * 4)
+      cmdq_addr = dev.cmdq_allocator.alloc(len(self._q) * 4, 16)
      cmdq_wptr = (cmdq_addr - dev.cmdq_page.va_addr) // 4
      dev.cmdq[cmdq_wptr : cmdq_wptr + len(self._q)] = array.array('I', self._q)

@@ -156,10 +156,14 @@ class NVComputeQueue(NVCommandQueue):
      for i in range(2):
        if self.active_qmd.read(f'release{i}_enable') == 0:
          self.active_qmd.write(**{f'release{i}_enable': 1})
-          self.bind_sints_to_mem(signal.value_addr, mem=self.active_qmd_buf.cpu_view(), fmt='Q', mask=0xfffffffff,
-            offset=self.active_qmd.field_offset(f'release{i}_address_lower' if self.active_qmd.ver<4 else f'release_semaphore{i}_addr_lower'))
-          self.bind_sints_to_mem(value, mem=self.active_qmd_buf.cpu_view(), fmt='Q',
-            offset=self.active_qmd.field_offset(f'release{i}_payload_lower' if self.active_qmd.ver<4 else f'release_semaphore{i}_payload_lower'))
+
+          addr_off = self.active_qmd.field_offset(f'release{i}_address_lower' if self.active_qmd.ver<4 else f'release_semaphore{i}_addr_lower')
+          self.bind_sints_to_mem(signal.value_addr & 0xffffffff, mem=self.active_qmd_buf.cpu_view(), fmt='I', offset=addr_off)
+          self.bind_sints_to_mem(signal.value_addr >> 32, mem=self.active_qmd_buf.cpu_view(), fmt='I', mask=0xf, offset=addr_off+4)
+
+          val_off = self.active_qmd.field_offset(f'release{i}_payload_lower' if self.active_qmd.ver<4 else f'release_semaphore{i}_payload_lower')
+          self.bind_sints_to_mem(value & 0xffffffff, mem=self.active_qmd_buf.cpu_view(), fmt='I', offset=val_off)
+          self.bind_sints_to_mem(value >> 32, mem=self.active_qmd_buf.cpu_view(), fmt='I', offset=val_off+4)
          return self

    self.nvm(0, nv_gpu.NVC56F_SEM_ADDR_LO, *data64_le(signal.value_addr), *data64_le(value),
@@ -384,7 +388,7 @@ class NVKIface:
    if made.params.status != 0: raise RuntimeError(f"_gpu_map_to_cpu returned {get_error_str(made.params.status)}")
    return fd_dev.mmap(target, size, mmap.PROT_READ|mmap.PROT_WRITE, mmap.MAP_SHARED | (MAP_FIXED if target is not None else 0), 0)

-  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, map_flags=0, cpu_addr=None) -> HCQBuffer:
+  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, map_flags=0, cpu_addr=None, **kwargs) -> HCQBuffer:
    # Uncached memory is "system". Use huge pages only for gpu memory.
    page_size = (4 << (12 if OSX else 10)) if uncached or host else ((2 << 20) if size >= (8 << 20) else (4 << (12 if OSX else 10)))
    size = round_up(size, page_size)
@@ -455,13 +459,13 @@ class PCIIface(PCIIfaceBase):
  def __init__(self, dev, dev_id):
    super().__init__(dev, dev_id, vendor=0x10de, devices=[0x2204, 0x2684, 0x2b85], bars=[0, 1], vram_bar=1,
      va_start=NVMemoryManager.va_allocator.base, va_size=NVMemoryManager.va_allocator.size)
-    System.reserve_hugepages(64)
+    if not OSX: System.reserve_hugepages(64)

    self.pci_dev.write_config(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
    self.dev_impl:NVDev = NVDev(self.pci_dev.pcibus, self.pci_dev.map_bar(0, fmt='I'), self.pci_dev.map_bar(1),
      self.pci_dev.read_config(pci.PCI_VENDOR_ID, 4), self.pci_dev.read_config(pci.PCI_SUBSYSTEM_VENDOR_ID, 4),
      self.pci_dev.read_config(pci.PCI_REVISION_ID, 1), self.pci_dev.bar_info)
-    self.root, self.gpu_instance, self.p2p_base_addr = 0xc1000000, 0, self.pci_dev.bar_info[1][0]
+    self.root, self.gpu_instance = 0xc1000000, 0
    self.rm_alloc(0, nv_gpu.NV01_ROOT, nv_gpu.NV0000_ALLOC_PARAMETERS())

    # Setup classes for the GPU
@@ -508,7 +512,7 @@ class NVDevice(HCQCompiled[HCQSignal]):
    channel_params = nv_gpu.NV_CHANNEL_GROUP_ALLOCATION_PARAMETERS(engineType=nv_gpu.NV2080_ENGINE_TYPE_GRAPHICS)
    channel_group = self.iface.rm_alloc(self.nvdevice, nv_gpu.KEPLER_CHANNEL_GROUP_A, channel_params)

-    gpfifo_area = self.iface.alloc(0x200000, contiguous=True, cpu_access=True, map_flags=0x10d0000)
+    gpfifo_area = self.iface.alloc(0x200000, contiguous=True, cpu_access=True, force_devmem=True, map_flags=0x10d0000)

    ctxshare_params = nv_gpu.NV_CTXSHARE_ALLOCATION_PARAMETERS(hVASpace=vaspace, flags=nv_gpu.NV_CTXSHARE_ALLOCATION_FLAGS_SUBCONTEXT_ASYNC)
    ctxshare = self.iface.rm_alloc(channel_group, nv_gpu.FERMI_CONTEXT_SHARE_A, ctxshare_params)
@@ -139,7 +139,7 @@ class NV_FLCN(NV_IP):

      return System.alloc_sysmem(len(patched_image), contiguous=True, data=patched_image)

-    self.frts_image_va, self.frts_image_sysmem = __patch(0x15, bytes(frts_cmd))
+    _, self.frts_image_sysmem = __patch(0x15, bytes(frts_cmd))

  def prep_booter(self):
    image = self.nvdev.extract_fw("kgspBinArchiveBooterLoadUcode", "image_prod_data")
@@ -150,7 +150,7 @@ class NV_FLCN(NV_IP):

    patched_image = bytearray(image)
    patched_image[patch_loc:patch_loc+sig_len] = sig[:sig_len]
-    self.booter_image_va, self.booter_image_sysmem = System.alloc_sysmem(len(patched_image), contiguous=True, data=patched_image)
+    _, self.booter_image_sysmem = System.alloc_sysmem(len(patched_image), contiguous=True, data=patched_image)
    _, _, self.booter_data_off, self.booter_data_sz, _, self.booter_code_off, self.booter_code_sz, _, _ = struct.unpack("9I", header)

  def init_hw(self):
@@ -327,10 +327,10 @@ class NV_GSP(NV_IP):
    # Alloc queues
    pte_cnt = ((queue_pte_cnt:=(queue_size * 2) // 0x1000)) + round_up(queue_pte_cnt * 8, 0x1000) // 0x1000
    pt_size = round_up(pte_cnt * 8, 0x1000)
-    queues_va, queues_sysmem = System.alloc_sysmem(pt_size + queue_size * 2, contiguous=False)
+    queues_view, queues_sysmem = System.alloc_sysmem(pt_size + queue_size * 2, contiguous=False)

    # Fill up ptes
-    for i, sysmem in enumerate(queues_sysmem): to_mv(queues_va + i * 0x8, 0x8).cast('Q')[0] = sysmem
+    for i, sysmem in enumerate(queues_sysmem): queues_view.view(i * 0x8, 0x8, fmt='Q')[0] = sysmem

    # Fill up arguments
    queue_args = nv.MESSAGE_QUEUE_INIT_ARGUMENTS(sharedMemPhysAddr=queues_sysmem[0], pageTableEntryCount=pte_cnt, cmdQueueOffset=pt_size,
@@ -338,7 +338,7 @@ class NV_GSP(NV_IP):
    _, self.rm_args_sysmem = self.nvdev._alloc_boot_struct(nv.GSP_ARGUMENTS_CACHED(bDmemStack=True, messageQueueInitArguments=queue_args))

    # Build command queue header
-    self.cmd_q_va, self.stat_q_va = queues_va + pt_size, queues_va + pt_size + queue_size
+    self.cmd_q_va, self.stat_q_va = queues_view.addr + pt_size, queues_view.addr + pt_size + queue_size

    cmd_q_tx = nv.msgqTxHeader(version=0, size=queue_size, entryOff=0x1000, msgSize=0x1000, msgCount=(queue_size - 0x1000) // 0x1000,
      writePtr=0, flags=1, rxHdrOff=ctypes.sizeof(nv.msgqTxHeader))
@@ -348,9 +348,9 @@ class NV_GSP(NV_IP):

  def init_libos_args(self):
    _, logbuf_sysmem = System.alloc_sysmem((2 << 20), contiguous=True)
-    libos_args_va, self.libos_args_sysmem = System.alloc_sysmem(0x1000, contiguous=True)
+    libos_args_view, self.libos_args_sysmem = System.alloc_sysmem(0x1000, contiguous=True)

-    libos_structs = (nv.LibosMemoryRegionInitArgument * 6).from_address(libos_args_va)
+    libos_structs = (nv.LibosMemoryRegionInitArgument * 6).from_address(libos_args_view.addr)
    for i, name in enumerate(["INIT", "INTR", "RM", "MNOC", "KRNL"]):
      libos_structs[i] = nv.LibosMemoryRegionInitArgument(kind=nv.LIBOS_MEMORY_REGION_CONTIGUOUS, loc=nv.LIBOS_MEMORY_REGION_LOC_SYSMEM, size=0x10000,
        id8=int.from_bytes(bytes(f"LOG{name}", 'utf-8'), 'big'), pa=logbuf_sysmem[0] + 0x10000 * i)
@@ -370,18 +370,18 @@ class NV_GSP(NV_IP):
    for i in range(3, 0, -1): npages[i-1] = ((npages[i] - 1) >> (nv.LIBOS_MEMORY_REGION_RADIX_PAGE_LOG2 - 3)) + 1

    offsets = [sum(npages[:i]) * 0x1000 for i in range(4)]
-    radix_va, self.gsp_radix3_sysmem = System.alloc_sysmem(offsets[-1] + len(self.gsp_image), contiguous=False)
+    radix_view, self.gsp_radix3_sysmem = System.alloc_sysmem(offsets[-1] + len(self.gsp_image), contiguous=False)

    # Copy image
-    to_mv(radix_va + offsets[-1], len(self.gsp_image))[:] = self.gsp_image
+    radix_view.view(offsets[-1], len(self.gsp_image))[:] = self.gsp_image

    # Copy level and image pages.
    for i in range(0, 3):
      cur_offset = sum(npages[:i+1])
-      to_mv(radix_va + offsets[i], npages[i+1] * 8).cast('Q')[:] = array.array('Q', self.gsp_radix3_sysmem[cur_offset:cur_offset+npages[i+1]])
+      radix_view.view(offsets[i], npages[i+1] * 8, fmt='Q')[:] = array.array('Q', self.gsp_radix3_sysmem[cur_offset:cur_offset+npages[i+1]])

    # Copy signature
-    self.gsp_signature_va, self.gsp_signature_sysmem = System.alloc_sysmem(len(signature), contiguous=True, data=signature)
+    _, self.gsp_signature_sysmem = System.alloc_sysmem(len(signature), contiguous=True, data=signature)

  def init_boot_binary_image(self):
    self.booter_image = self.nvdev.extract_fw("kgspBinArchiveGspRmBoot", "ucode_image_prod_data")
@@ -522,7 +522,7 @@ class NV_GSP(NV_IP):
    self.stat_q.wait_resp(nv.NV_VGPU_MSG_FUNCTION_SET_PAGE_DIRECTORY)

  def rpc_set_gsp_system_info(self):
-    def bdf_as_int(s): return (int(s[5:7],16)<<8) | (int(s[8:10],16)<<3) | int(s[-1],16)
+    def bdf_as_int(s): return 0x000 if s.startswith("usb") else (int(s[5:7],16)<<8) | (int(s[8:10],16)<<3) | int(s[-1],16)

    data = nv.GspSystemInfo(gpuPhysAddr=self.nvdev.bars[0][0], gpuPhysFbAddr=self.nvdev.bars[1][0], gpuPhysInstAddr=self.nvdev.bars[3][0],
      pciConfigMirrorBase=[0x88000, 0x92000][self.nvdev.fmc_boot], pciConfigMirrorSize=0x1000, nvDomainBusDeviceFunc=bdf_as_int(self.nvdev.devfmt),
@@ -1,6 +1,6 @@
 from __future__ import annotations
 import ctypes, time, functools, re, gzip, struct
-from tinygrad.helpers import getenv, DEBUG, fetch, getbits, to_mv
+from tinygrad.helpers import getenv, DEBUG, fetch, getbits
 from tinygrad.runtime.support.hcq import MMIOInterface
 from tinygrad.runtime.support.memory import TLSFAllocator, MemoryManager
 from tinygrad.runtime.support.nv.ip import NV_FLCN, NV_FLCN_COT, NV_GSP
@@ -137,9 +137,9 @@ class NVDev(PCIDevImplBase):
    self.large_bar = self.vram.nbytes >= self.vram_size

  def _alloc_boot_struct(self, struct:ctypes.Structure) -> tuple[ctypes.Structure, int]:
-    va, paddrs = System.alloc_sysmem(sz:=ctypes.sizeof(type(struct)), contiguous=True)
-    to_mv(va, sz)[:] = bytes(struct)
-    return type(struct).from_address(va), paddrs[0]
+    view, paddrs = System.alloc_sysmem(sz:=ctypes.sizeof(type(struct)), contiguous=True)
+    view[:sz] = bytes(struct)
+    return type(struct).from_address(view.addr), paddrs[0]

  def _download(self, file:str) -> str:
    url = f"https://raw.githubusercontent.com/NVIDIA/open-gpu-kernel-modules/8ec351aeb96a93a4bb69ccc12a542bf8a8df2b6f/{file}"
@@ -1,6 +1,6 @@
 import os, mmap, array, functools, ctypes, select, contextlib, dataclasses, sys, errno, itertools
 from typing import cast, ClassVar
-from tinygrad.helpers import round_up, to_mv, getenv, OSX, temp
+from tinygrad.helpers import round_up, getenv, OSX, temp
 from tinygrad.runtime.autogen import libc, vfio
 from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface, HCQBuffer
 from tinygrad.runtime.support.memory import MemoryManager, VirtMapping
@@ -49,6 +49,18 @@ class _System:
      raise RuntimeError("IOServiceOpen failed")
    return conn

+  def iokit_pci_memmap(self, typ:int):
+    if self.iokit.IOConnectMapMemory64(self.macos_tinygpu_conn, ctypes.c_uint32(typ), System.mach_task_self,
+      ctypes.byref(addr:=ctypes.c_uint64(0)), ctypes.byref(size:=ctypes.c_uint64(0)), 0x1): raise RuntimeError(f"IOConnectMapMemory64({typ=}) failed")
+    return MMIOInterface(addr.value, size.value)
+
+  def iokit_pci_rpc(self, sel:int, *args:int):
+    in_scalars = (ctypes.c_uint64 * len(args))(*args) if args else ctypes.POINTER(ctypes.c_uint64)()
+    if (self.iokit.IOConnectCallMethod(self.macos_tinygpu_conn, sel, in_scalars, len(args), None, ctypes.c_size_t(0),
+        out_scalars:=(ctypes.c_uint64*16)(), ctypes.byref(outcnt:=ctypes.c_uint32(16)), None, ctypes.byref(ctypes.c_size_t(0)))):
+      raise RuntimeError(f"IOConnectCallMethod({sel=}, {args=}) failed")
+    return out_scalars[:outcnt.value]
+
  def reserve_hugepages(self, cnt): os.system(f"sudo sh -c 'echo {cnt} > /proc/sys/vm/nr_hugepages'")

  def memory_barrier(self): lib.atomic_thread_fence(__ATOMIC_SEQ_CST:=5) if (lib:=self.libsys if OSX else self.atomic_lib) is not None else None
@@ -60,15 +72,25 @@ class _System:
    self.pagemap.seek(vaddr // mmap.PAGESIZE * 8)
    return [(x & ((1<<55) - 1)) * mmap.PAGESIZE for x in array.array('Q', self.pagemap.read(size//mmap.PAGESIZE*8, binary=True))]

-  def alloc_sysmem(self, size:int, vaddr:int=0, contiguous:bool=False, data:bytes|None=None) -> tuple[int, list[int]]:
-    assert not contiguous or size <= (2 << 20), "Contiguous allocation is only supported for sizes up to 2MB"
-    flags = (libc.MAP_HUGETLB if contiguous and (size:=round_up(size, mmap.PAGESIZE)) > 0x1000 else 0) | (MAP_FIXED if vaddr else 0)
-    va = FileIOInterface.anon_mmap(vaddr, size, mmap.PROT_READ|mmap.PROT_WRITE, mmap.MAP_SHARED|mmap.MAP_ANONYMOUS|MAP_POPULATE|MAP_LOCKED|flags, 0)
+  def alloc_sysmem(self, size:int, vaddr:int=0, contiguous:bool=False, data:bytes|None=None) -> tuple[MMIOInterface, list[int]]:
+    if OSX:
+      sysmem_view = System.iokit_pci_memmap(round_up(size, mmap.PAGESIZE))
+      paddrs = list(itertools.takewhile(lambda p: p[1] != 0, zip(sysmem_view.view(fmt='Q')[0::2], sysmem_view.view(fmt='Q')[1::2])))
+      assert not contiguous or len(paddrs) == 1, "not contiguous, but required"
+      paged_paddrs = [p + i for p, sz in paddrs for i in range(0, sz, 0x1000)][:round_up(size, 0x1000)//0x1000]
+    else:
+      assert not contiguous or size <= (2 << 20), "Contiguous allocation is only supported for sizes up to 2MB"
+      flags = (libc.MAP_HUGETLB if contiguous and (size:=round_up(size, mmap.PAGESIZE)) > 0x1000 else 0) | (MAP_FIXED if vaddr else 0)
+      va = FileIOInterface.anon_mmap(vaddr, size, mmap.PROT_READ|mmap.PROT_WRITE, mmap.MAP_SHARED|mmap.MAP_ANONYMOUS|MAP_POPULATE|MAP_LOCKED|flags, 0)
+      sysmem_view, paged_paddrs = MMIOInterface(va, size), self.system_paddrs(va, size)

-    if data is not None: to_mv(va, len(data))[:] = data
-    return va, self.system_paddrs(va, size)
+    if data is not None: sysmem_view[:len(data)] = data
+    return sysmem_view, paged_paddrs
+
+  def pci_reset(self, gpu):
+    if OSX: System.iokit_pci_rpc(__TinyGPURPCReset:=2)
+    else: os.system(f"sudo sh -c 'echo 1 > /sys/bus/pci/devices/{gpu}/reset'")

-  def pci_reset(self, gpu): os.system(f"sudo sh -c 'echo 1 > /sys/bus/pci/devices/{gpu}/reset'")
  def pci_scan_bus(self, target_vendor:int, target_devices:list[int]) -> list[str]:
    result = []
    for pcibus in FileIOInterface("/sys/bus/pci/devices").listdir():
@@ -143,14 +165,12 @@ class PCIDevice:
    return MMIOInterface(loc, sz, fmt=fmt)

 class APLPCIDevice(PCIDevice):
-  def __init__(self, pcibus:str, bars:list[int], resize_bars:list[int]|None=None): self.pcibus, self.bars = pcibus, {b: self.map_mem(b) for b in bars}
-  def map_mem(self, typ:int) -> MMIOInterface:
-    if System.iokit.IOConnectMapMemory64(System.macos_tinygpu_conn, ctypes.c_uint32(typ), System.mach_task_self,
-      ctypes.byref(addr:=ctypes.c_uint64(0)), ctypes.byref(size:=ctypes.c_uint64(0)), 0x1): raise RuntimeError(f"IOConnectMapMemory64({typ=}) failed")
-    return MMIOInterface(addr.value, size.value)
+  def __init__(self, pcibus:str, bars:list[int], resize_bars:list[int]|None=None):
+    self.pcibus, self.bars = pcibus, {b: System.iokit_pci_memmap(b) for b in bars}
+    self.bar_info = {b:(0, self.bars[b].nbytes-1 if b in self.bars else 0, 0) for b in range(6)} # NOTE: fake bar info for nv.
  def map_bar(self, bar:int, off:int=0, addr:int=0, size:int|None=None, fmt='B') -> MMIOInterface: return self.bars[bar].view(off, size, fmt)
-  def read_config(self, offset:int, size:int): return 0
-  def write_config(self, offset:int, value:int, size:int): pass
+  def read_config(self, offset:int, size:int): return System.iokit_pci_rpc(__TinyGPURPCReadCfg:=0, offset, size)[0]
+  def write_config(self, offset:int, value:int, size:int): System.iokit_pci_rpc(__TinyGPURPCWriteCfg:=1, offset, size, value)

 class PCIDevImplBase:
  mm: MemoryManager
@@ -173,23 +193,23 @@ class LNXPCIIfaceBase:
    self.pci_dev, self.dev, self.vram_bar = PCIDevice(cls.gpus[dev_id], bars=bars, resize_bars=[vram_bar]), dev, vram_bar
    self.p2p_base_addr = self.pci_dev.bar_info[vram_bar][0]

-  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, **kwargs) -> HCQBuffer:
-    if host or (uncached and cpu_access): # host or gtt-like memory.
+  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, force_devmem=False, **kwargs) -> HCQBuffer:
+    # NOTE: logic on macos is different, since bar is small
+    should_use_sysmem = host or (((uncached or cpu_access) if OSX else (uncached and cpu_access)) and not force_devmem)
+    if should_use_sysmem:
      vaddr = self.dev_impl.mm.alloc_vaddr(size:=round_up(size, mmap.PAGESIZE), align=mmap.PAGESIZE)
-      paddrs = [(paddr, mmap.PAGESIZE) for paddr in System.alloc_sysmem(size, vaddr=vaddr, contiguous=contiguous)[1]]
-      mapping = self.dev_impl.mm.map_range(vaddr, size, paddrs, system=True, snooped=True, uncached=True)
-      return HCQBuffer(vaddr, size, meta=PCIAllocationMeta(mapping, has_cpu_mapping=True, hMemory=paddrs[0][0]),
-        view=MMIOInterface(mapping.va_addr, size, fmt='B'), owner=self.dev)
+      memview, paddrs = System.alloc_sysmem(size, vaddr=vaddr, contiguous=contiguous)
+      mapping = self.dev_impl.mm.map_range(vaddr, size, [(paddr, 0x1000) for paddr in paddrs], system=True, snooped=True, uncached=True)
+      return HCQBuffer(vaddr, size, meta=PCIAllocationMeta(mapping, has_cpu_mapping=True, hMemory=paddrs[0]), view=memview, owner=self.dev)

    mapping = self.dev_impl.mm.valloc(size:=round_up(size, 4 << 10), uncached=uncached, contiguous=cpu_access)
-    if cpu_access: self.pci_dev.map_bar(bar=self.vram_bar, off=mapping.paddrs[0][0], addr=mapping.va_addr, size=mapping.size)
-    return HCQBuffer(mapping.va_addr, size, view=MMIOInterface(mapping.va_addr, size, fmt='B') if cpu_access else None,
-      meta=PCIAllocationMeta(mapping, has_cpu_mapping=cpu_access, hMemory=mapping.paddrs[0][0]), owner=self.dev)
+    barview = self.pci_dev.map_bar(bar=self.vram_bar, off=mapping.paddrs[0][0], size=mapping.size) if cpu_access else None
+    return HCQBuffer(mapping.va_addr, size, view=barview, meta=PCIAllocationMeta(mapping, cpu_access, hMemory=mapping.paddrs[0][0]), owner=self.dev)

  def free(self, b:HCQBuffer):
    for dev in b.mapped_devs[1:]: dev.iface.dev_impl.mm.unmap_range(b.va_addr, b.size)
    if not b.meta.mapping.system: self.dev_impl.mm.vfree(b.meta.mapping)
-    if b.owner == self.dev and b.meta.has_cpu_mapping: FileIOInterface.munmap(b.va_addr, b.size)
+    if b.owner == self.dev and b.meta.has_cpu_mapping and not OSX: FileIOInterface.munmap(b.va_addr, b.size)

  def map(self, b:HCQBuffer):
    if b.owner is not None and b.owner._is_cpu():
@@ -205,21 +225,6 @@ class LNXPCIIfaceBase:
 class APLPCIIfaceBase(LNXPCIIfaceBase):
  def __init__(self, dev, dev_id, vendor, devices, bars, vram_bar, va_start, va_size):
    self.pci_dev, self.dev, self.vram_bar = APLPCIDevice(pcibus=f'usb4:{dev_id}', bars=bars), dev, vram_bar
-
-  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, **kwargs) -> HCQBuffer:
-    if host or uncached or cpu_access: # cpu access memory goes here, since bar is small.
-      vaddr = self.dev_impl.mm.alloc_vaddr(size:=round_up(size, mmap.PAGESIZE), align=mmap.PAGESIZE)
-      assert size >= mmap.PAGESIZE, "Size must be at least one page"
-
-      sysmem = cast(APLPCIDevice, self.pci_dev).map_mem(size).view(fmt='Q')
-      paddrs = list(itertools.takewhile(lambda p: p[1] != 0, zip(sysmem[0::2], sysmem[1::2])))
-
-      mapping = self.dev_impl.mm.map_range(vaddr, size, paddrs, system=True, snooped=True, uncached=True)
-      return HCQBuffer(vaddr, size, meta=PCIAllocationMeta(mapping, has_cpu_mapping=True), view=sysmem.view(fmt='B'), owner=self.dev)
-
-    mapping = self.dev_impl.mm.valloc(size:=round_up(size, 4 << 10), uncached=uncached, contiguous=cpu_access)
-    return HCQBuffer(mapping.va_addr, size, view=None, meta=PCIAllocationMeta(mapping, has_cpu_mapping=False), owner=self.dev)
-
  def map(self, b:HCQBuffer): raise RuntimeError(f"map failed: {b.owner} -> {self.dev}")

 PCIIfaceBase:type = APLPCIIfaceBase if OSX else LNXPCIIfaceBase
@@ -108,6 +108,18 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
   (UPat(Ops.ASSIGN, src=(UPat.var("a"), UPat.var("b")), name="assign"), find_permutes),
 ])

+# *****************
+
+pm_where_is_multi = PatternMatcher([
+  # move *0 through where
+  (UPat.var("gate").where(UPat.var("a"), 0) * UPat.var("b"), lambda gate,a,b: gate.where(a*b, 0)),
+  # move *0 through unary op
+  (UPat(Ops.CONTIGUOUS, src=(UPat.var("gate").where(UPat.var("a"), 0),), name="u"), lambda gate,a,u: gate.where(u.replace(src=(a,)), 0)),
+  # move where 0 through reduce if the reduce ranges are not in the gate
+  (UPat(Ops.REDUCE, src=(UPat.var("gate").where(UPat.var("a"), 0),), name="red", allow_any_len=True),
+   lambda gate,a,red: gate.where(red.replace(src=(a,)+red.src[1:]), 0) if all(r not in gate.ranges for r in red.src[1:]) else None),
+])
+
 # *****************
 # 3.5 cleanups

@@ -272,7 +284,7 @@ def bufferize_to_store(x:UOp):
    assert assign_target.op is Ops.INDEX, f"{assign_target.op} is not index"
    # in assign, this is the buffer size, not the bufferize size
    # TODO: assign_mops here
-    ret = assign_target.replace(dtype=sdtype).store(assign_src, *rngs, dtype=x.dtype)
+    ret = assign_target.replace(dtype=sdtype).store(assign_src, *rngs, dtype=x.dtype).replace(tag=x.tag)
    mops = []
    walk = assign_mops
    while walk is not assign_mops.base:
@@ -284,7 +296,7 @@ def bufferize_to_store(x:UOp):
  # NOTE: the DEFINE_LOCAL needs to be disambiguated here
  if sdtype.addrspace == AddrSpace.GLOBAL:
    buf = UOp.new_buffer(x.arg.device, size, x.dtype)
-    ret = buf.reshape(shape).index(*rngs, dtype=sdtype).store(x.src[0], *rngs, dtype=x.dtype)
+    ret = buf.reshape(shape).index(*rngs, dtype=sdtype).store(x.src[0], *rngs, dtype=x.dtype).replace(tag=x.tag)
    ret = ret.forced_reshape(shape)
    # TODO: is this right? what if it's offset
    if any(r.op is Ops.RANGE and r.src[0].op is not Ops.CONST for r in rngs):
@@ -340,6 +352,7 @@ def handle_assign(ctx:LocalAddBufferContext, assign:UOp):

 def renumber_range(ctx:LocalAddBufferContext, r:UOp):
  if r.tag is not None: return None
+  if r.arg[-1] is AxisType.MULTI: return None
  ret = r.replace(arg=(ctx.range,)+r.arg[1:], tag=())
  ctx.range += 1
  return ret
@@ -414,7 +427,7 @@ class Kernel:
    return f"<Kernel {len(list(self.ast.toposort()))} {ast_rep} {self.metadata}>"

 def split_store(ctx:list[UOp], x:UOp):
-  if len(x.ranges): return None
+  if len([r for r in x.ranges if r.arg[-1] != AxisType.MULTI]): return None
  if x.src[0].ptrdtype.addrspace is AddrSpace.LOCAL: return None

  # local kernel rewrite
@@ -496,6 +509,9 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:

  # NOTE: sym (vs symbolic_simple) breaks things here because ranges with len 1 aren't handled right
  tsink = graph_rewrite(tsink, symbolic_simple+pm_reduce_unparented, name="symbolic")  # this supports const folding
+
+  tsink = graph_rewrite(tsink, pm_where_is_multi, name="where_is_multi")
+
  tsink = graph_rewrite(tsink, pm_cleanups, bottom_up=True, name="remove costly buffers")
  # TODO: can you substitute and remove costly buffers at the same time?
  tsink = graph_rewrite(tsink, pm_substitute_recurse, bottom_up=True, name="run substitutes")
@@ -231,9 +231,9 @@ class Tensor(MathTrait):
    # verify Tensors match the spec
    if __debug__: type_verify(list(big_sink.toposort()), tensor_uop_spec)

-    if any(isinstance(x._device, tuple) for x in big_sink.toposort()):
-      _apply_map_to_tensors(get_multi_map(big_sink), "Apply Multi Map")
-      big_sink = UOp.sink(*flatten([x.uop.src if x.uop.op is Ops.MULTI else [x.uop] for x in (self,)+lst]))
+    #if any(isinstance(x._device, tuple) for x in big_sink.toposort()):
+    #  _apply_map_to_tensors(get_multi_map(big_sink), "Apply Multi Map")
+    #  big_sink = UOp.sink(*flatten([x.uop.src if x.uop.op is Ops.MULTI else [x.uop] for x in (self,)+lst]))

    becomes_map = get_rangeify_map(big_sink)
    _apply_map_to_tensors(becomes_map, name="Apply Kernelize Map")
@@ -256,7 +256,7 @@ class Tensor(MathTrait):
    # create the schedule
    schedule, var_vals = create_schedule_with_vars(sink)
    schedule = memory_planner(schedule)
-    if DEBUG >= 1 and len(schedule) > 1: print(f"scheduled {len(schedule)} kernels in {(time.perf_counter()-st)*1000:.2f} ms")
+    if (DEBUG >= 1 and len(schedule) > 1) or DEBUG >= 3: print(f"scheduled {len(schedule)} kernels in {(time.perf_counter()-st)*1000:.2f} ms")
    return schedule, var_vals

  def schedule(self, *lst:Tensor) -> list[ScheduleItem]:
@@ -267,7 +267,8 @@ class Tensor(MathTrait):

  def realize(self, *lst:Tensor, do_update_stats=True) -> Tensor:
    """Triggers the computation needed to create these Tensor(s)."""
-    run_schedule(*self.schedule_with_vars(*lst), do_update_stats=do_update_stats)
+    if len(to_realize:=[x for x in (self,)+lst if not x.uop.is_contiguous()]):
+      run_schedule(*Tensor.schedule_with_vars(*to_realize), do_update_stats=do_update_stats)
    return self

  def replace(self, x:Tensor, allow_shape_mismatch=False) -> Tensor:
@@ -293,8 +294,7 @@ class Tensor(MathTrait):
    assert self.shape == x.shape, f"assign shape mismatch {self.shape} != {x.shape}"
    assert self.device == x.device, f"assign device mismatch {self.device} != {x.device}"
    assert self.dtype == x.dtype, f"assign dtype mismatch {self.dtype} != {x.dtype}"
-    self.uop = self.uop.assign(x.uop)
-    return self
+    return self.replace(self._apply_uop(UOp.assign, x))

  def detach(self) -> Tensor:
    """
@@ -15,7 +15,7 @@ if TYPE_CHECKING:
 class AxisType(Enum):
  def __repr__(self): return str(self)
  GLOBAL = auto(); WARP = auto(); LOCAL = auto(); LOOP = auto(); GROUP_REDUCE = auto(); REDUCE = auto(); UPCAST = auto(); UNROLL = auto() # noqa: E702
-  THREAD = auto()
+  THREAD = auto(); MULTI = auto() # noqa: E702

 range_start = {Ops.BUFFERIZE: 1, Ops.REDUCE: 1, Ops.STORE: 2, Ops.WMMA: 3}

@@ -114,7 +114,9 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
  @functools.cached_property
  def key(self) -> bytes:
    return hashlib.sha256(str((self.op, self.dtype, self.arg)).encode() + b"".join([s.key for s in self.src])).digest()
-  def __repr__(self): return pretty_print(self, lambda x: f"{type(self).__name__}({x.op}, {x.dtype}, arg={x.argstr()}{x.tagstr()}, src=(%s))")
+  def __repr__(self):
+    if self.dtype == dtypes.index: return srender(self)  # makes shapes print nicely
+    return pretty_print(self, lambda x: f"{type(self).__name__}({x.op}, {x.dtype}, arg={x.argstr()}{x.tagstr()}, src=(%s))")
  def argstr(self): return f'({", ".join(map(str, self.arg))})' if self.op is Ops.REDUCE_AXIS else repr(self.arg)
  def tagstr(self): return f", tag={self.tag}" if self.tag is not None else ""

@@ -220,7 +222,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
      match self.op:
        case Ops.RESHAPE:
          if not all(x >= 0 for x in self.marg): raise ValueError(f"shape can't contain negative numbers {self.marg}")
-          if prod(ps) != prod(self.marg): raise ValueError(f"bad reshape: {ps} -> {self.marg}")
+          #if prod(ps) != prod(self.marg): raise ValueError(f"bad reshape: {ps} -> {self.marg}")
          return self.marg
        case Ops.EXPAND:
          if len(ps) != len(self.marg) or not all(s==ns or (s==1 and ns>=0) for s,ns in zip(ps, self.marg)):
@@ -370,15 +372,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
    return UOp(Ops.RANGE, dtype=dtypes.index, src=(sint_to_uop(end),), arg=arg)
  def r(self, op:Ops, axis:tuple[int, ...]):
    axis = tuple(sorted([x for x in axis if resolve(self.shape[x] != 1)]))
-    if len(axis) == 0: return self
-    # move any non reduce axis before the first reduce axis
-    move_early, rest = partition(range(axis[0], len(self.shape)), lambda i: i not in axis and resolve(self.shape[i] != 1))
-    permaxis = tuple(range(axis[0])) + tuple(move_early) + tuple(rest)
-    ret = self.permute(permaxis)
-    new_axis = tuple([x for x in range(axis[0]+len(move_early), len(self.shape)) if resolve(ret.shape[x] != 1)])
-    assert len(axis) == len(new_axis)
-    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)]))
+    return UOp(Ops.REDUCE_AXIS, self.dtype, (self,), (op, axis)) if len(axis) else self
  @staticmethod
  def invalid(count=1): return UOp(Ops.CONST, dtypes.index.vec(count), src=(), arg=Invalid)
  def valid(self, cond): return self if cond.op is Ops.WHERE and cond.arg else cond.where(self, UOp.invalid(self.dtype.count))
@@ -389,7 +383,15 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
    assert self.dtype.scalar() is dtypes.index, "Can only call get_valid on index dtype"
    return self.src[0] if self.op is Ops.WHERE and self.src[2].arg is Invalid else UOp.const(dtypes.bool, self.arg is not Invalid)
  def reduce(self, *src:UOp, **kwargs): return UOp(Ops.REDUCE, kwargs.pop('dtype', self.dtype), src=(self,)+src, **kwargs)
-  def contiguous(self, *args, **kwargs): return UOp(Ops.CONTIGUOUS, dtype=self.dtype, src=(self,)+args, **kwargs)
+
+  def is_contiguous(self):
+    # TODO: this is is_realized
+    if self.op is Ops.RESHAPE: return self.src[0].is_contiguous()
+    return self.op is Ops.BUFFER
+
+  def contiguous(self, *args, **kwargs):
+    if self.is_contiguous(): return self
+    return UOp(Ops.CONTIGUOUS, dtype=self.dtype, src=(self,)+args, **kwargs)
  def contiguous_backward(self): return self.alu(Ops.CONTIGUOUS_BACKWARD)
  def bufferize(self, *args, **kwargs): return UOp(Ops.BUFFERIZE, dtype=self.dtype, src=(self,)+args, **kwargs)
  def fuse(self): return self.alu(Ops.FUSE)
@@ -436,16 +438,30 @@ class UOp(MathTrait, metaclass=UOpMetaClass):

  def _unshard(self, axis:int) -> UOp:
    bsz, dcount = self.shape[axis], len(self.device)
-    dnum = UOp.variable("_device_num", 0, dcount-1)
+    #dnum = UOp.variable("_device_num", 0, dcount-1)
+    dnum = UOp.range(dcount, -10, AxisType.MULTI)
    return self.pad(tuple((0,0) if a != axis else (bsz*dnum, bsz*(dcount-1) - bsz*dnum) for a in range(len(self.shape))))

  def _shard(self, axis:int) -> UOp:
    dcount = len(self.device)
-    dnum = UOp.variable("_device_num", 0, dcount-1)
+    dnum = UOp.range(dcount, -10, AxisType.MULTI)
    if self.shape[axis] % dcount != 0: raise RuntimeError(f"multi axis uneven: {self.shape[axis]=} {axis=} {dcount=}")
    sz = self.shape[axis] // dcount
-    return self.shrink(tuple((0,s) if i != axis else (dnum*sz,dnum*sz+sz) for i,s in enumerate(self.shape)))
-  def shard(self, devices:tuple[str, ...], axis:int) -> UOp: return self.copy_to_device(devices)._shard(axis).multi(axis)
+    #ret = self.reshape(tuple(s if i != axis else dnum*sz for i,s in enumerate(self.shape)))
+    #return ret
+
+    #flatten([[s] if i != axis else [dcount, sz] for i,s in enumerate(self.shape)])))
+
+    #ret = self.shrink(tuple((0,s) if i != axis else (dnum*sz,dnum*sz+sz) for i,s in enumerate(self.shape)))
+    #print(ret.shape)
+    #print(dnum)
+    #dnum = UOp.variable("_device_num", 0, dcount-1)
+    # TODO: 0 isn't correct here
+    ret = self.shrink(tuple((0,s) if i != axis else (dnum*sz,dnum*sz+sz) for i,s in enumerate(self.shape)))
+    ret = ret.pad(tuple((0,0) if a != axis else (sz*dnum, sz*(dcount-1) - sz*dnum) for a in range(len(self.shape))))
+    return ret
+
+  def shard(self, devices:tuple[str, ...], axis:int) -> UOp: return self.copy_to_device(devices)._shard(axis) #.multi(axis)

  # *** from LazyBuffer ***

@@ -497,10 +513,6 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
    if ret.shape == self.shape and same_shape_noop: return self
    return ret

-  def is_contiguous(self):
-    if self.op is Ops.RESHAPE: return self.src[0].is_contiguous()
-    return self.op is Ops.BUFFER
-
  # in these four, if the shape doesn't change we can return self
  def forced_reshape(self, arg:tuple[sint, ...]): return self._mop(Ops.RESHAPE, arg, same_shape_noop=False)
  def reshape(self, arg:tuple[sint, ...]): return self._mop(Ops.RESHAPE, arg, same_shape_noop=True)
@@ -272,10 +272,6 @@ async function renderProfiler() {
          }
        }
      }
-      for (const [_, v] of temp) {
-        v.x.push(x);
-        v.y.push(v.y.at(-1));
-      }
      timestamps.push(dur);
      const height = heightScale(peak);
      const yscale = d3.scaleLinear().domain([0, peak]).range([height, 0]);
@@ -152,12 +152,17 @@ def timeline_layout(dev_events:list[tuple[int, int, float, DevEvent]], start_ts:
    events.append(struct.pack("<IIIfI", enum_str(name, scache), option(ref), st-start_ts, dur, enum_str(info or "", scache)))
  return struct.pack("<BI", 0, len(events))+b"".join(events) if events else None

+def encode_mem_free(key:int, ts:int, execs:list[ProfilePointEvent], scache:dict) -> bytes:
+  kernel_names = [enum_str(ei.key, scache) for ei in execs]
+  return struct.pack(f"<BIII{len(kernel_names)}I", 0, ts, key, len(kernel_names), *kernel_names)
+
 def mem_layout(dev_events:list[tuple[int, int, float, DevEvent]], start_ts:int, end_ts:int, peaks:list[int], dtype_size:dict[str, int],
               scache:dict[str, int]) -> bytes|None:
  peak, mem = 0, 0
  temp:dict[int, int] = {}
  events:list[bytes] = []
  buf_ei:dict[int, list[ProfilePointEvent]] = {}
+
  for st,_,_,e in dev_events:
    if not isinstance(e, ProfilePointEvent): continue
    if e.name == "alloc":
@@ -170,9 +175,9 @@ def mem_layout(dev_events:list[tuple[int, int, float, DevEvent]], start_ts:int,
    if e.name == "exec" and e.arg["bufs"]:
      for b in e.arg["bufs"]: buf_ei.setdefault(b, []).append(e)
    if e.name == "free":
-      kernel_names = [enum_str(ei.key, scache) for ei in buf_ei.pop(e.key, [])]
-      events.append(struct.pack(f"<BIII{len(kernel_names)}I", 0, int(e.ts) - start_ts, e.key, len(kernel_names), *kernel_names))
+      events.append(encode_mem_free(e.key, int(e.ts) - start_ts, buf_ei.pop(e.key, []), scache))
      mem -= temp.pop(e.key)
+  for t in temp: events.append(encode_mem_free(t, end_ts-start_ts, buf_ei.pop(t, []), scache))
  peaks.append(peak)
  return struct.pack("<BIQ", 1, len(events), peak)+b"".join(events) if events else None
Author	SHA1	Message	Date
geohot	978502be46	experiments with multi being range	2025-10-17 14:11:55 +08:00
chenyuandGitHub	9561803cb0	fix assert in test_schedule (#12745 ) * fix assert in test_schedule updated kernel counts and some old tests * fix	2025-10-16 15:39:50 -04:00
chenyuandGitHub	285534ce64	delete DONT_REALIZE_EXPAND and DONT_GROUP_REDUCES (#12744 ) does nothing now	2025-10-16 14:11:33 -04:00
chenyuandGitHub	98239f1156	few shapetracker cleanups (#12741 )	2025-10-16 12:43:27 -04:00
chenyuandGitHub	53478c741d	relax ASSERT_MIN_STEP_TIME for space lab policy (#12742 )	2025-10-16 11:40:36 -04:00
geohotstanandGitHub	5d209ee7ec	onnx helper intermediate node output validation (#12740 ) * start * update comments * good * add comments and better printing * done	2025-10-16 11:17:47 -04:00
sirhcmandGitHub	bce2bc0465	Revert "use RTLD_GLOBAL on macos" (#12738 ) This reverts commit `89fe3e574d`.	2025-10-16 10:07:21 -04:00
chenyuandGitHub	f34f26bca0	fix gpt2 with benchmark (#12736 ) `CPU=1 python3 examples/gpt2.py --benchmark 128` works now	2025-10-16 09:55:20 -04:00
Sieds LyklesandGitHub	55db1b0e0e	reduce where that is cut from two sides (#12733 ) * better rule * correct pattern * shorten line	2025-10-16 15:25:15 +02:00
nimlgenandGitHub	cf9baeea61	Revert "nv: check if jitlink is avail (#12731 )" (#12735 ) This reverts commit `a069a45d14`.	2025-10-16 20:41:49 +08:00
George HotzandGitHub	8be7844b2e	use apply uop for assign to fix assign metadata (#12732 ) * use apply uop for assign * fix metadata for assign * fix backward metadata * those aren't real tests	2025-10-16 20:34:12 +08:00
nimlgenandGitHub	3aa2277b8f	nv: usb4 (#12696 ) * hackish * prog * match * l * simpler * refactor * not osx * apple things * tiny changes * fix mask * match fix * nn	2025-10-16 20:11:19 +08:00
nimlgenandGitHub	a069a45d14	nv: check if jitlink is avail (#12731 )	2025-10-16 19:58:50 +08:00
George HotzandGitHub	a498ec9c18	cleanup names of postrange + fast FUSE_OPTIM (#12730 ) * cleanup names of postrange * make FUSE_OPTIM not slow * delete junk in def r	2025-10-16 19:38:31 +08:00
Sieds LyklesandGitHub	8f740e07ff	no broadcasting/vectors in reduce collapse (#12729 )	2025-10-16 13:22:57 +02:00
qazalandGitHub	533f18b22c	viz: add trace data for inflight buffers (#12728 ) * viz: add trace data for inflight buffers * add test_inflight_buf * temp stores the keys * update tests / use Tensor.ones	2025-10-16 19:15:03 +08:00
George HotzandGitHub	af4479c169	faster stable diffusion load (#12725 ) * faster stable diffusion load * failing tests	2025-10-16 18:31:59 +08:00
nimlgenandGitHub	e7c057d5dc	system: alloc_sysmem return view (#12724 ) * system: alloc_sysmem return view * e	2025-10-16 17:55:01 +08:00