fp32 paeak

* add load to the spec

* can

.github/workflows/test.yml

@@ -417,7 +417,7 @@ jobs:
           llvm: 'true'
       - name: Test openpilot model kernel count and gate usage
         run: |
-          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1486 ALLOWED_GATED_READ_IMAGE=18 FLOAT16=1 DEV=CL IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1468 ALLOWED_GATED_READ_IMAGE=18 FLOAT16=1 DEV=CL IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
       - name: Test openpilot CL compile fp16
         run: FLOAT16=1 DEV=CL IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
       - name: Test openpilot CL compile fp32 (test correctness)

README.md

@@ -140,8 +140,8 @@ Documentation along with a quick start guide can be found on the [docs website](
 ```python
 from tinygrad import Tensor
 
-x = Tensor.eye(3, requires_grad=True)
-y = Tensor([[2.0,0,-2.0]], requires_grad=True)
+x = Tensor.eye(3)
+y = Tensor([[2.0,0,-2.0]])
 z = y.matmul(x).sum()
 z.backward()
 

examples/gradaccum_mnist.py

@@ -35,9 +35,8 @@ if __name__ == "__main__":
 
   params = nn.state.get_parameters(model)
 
-  # init params, set requires grad on the ones we need gradients of
+  # init params
   for x in params:
-    if x.requires_grad is None: x.requires_grad_()
     x.replace(x.contiguous())
   Tensor.realize(*params)
 

examples/mlperf/model_train.py

@@ -1458,7 +1458,7 @@ def train_llama3():
     if is_dp: tokens = tokens.to(None).shard(device, 0)
     if is_mp: tokens = tokens.shard(device)
     if not is_sharding: tokens = tokens.to(None)
-    logits:Tensor = model(tokens[:, :-1])
+    logits:Tensor = model(tokens[:, :-1], save=bool(SMALL))
     if getenv("FAST_CE", 0):
       from extra.llama_kernels.fused_ce import fused_ce_loss
       loss = fused_ce_loss(logits.cast(dtypes.bfloat16), tokens[:, 1:], label_smoothing=0.0)

examples/mlperf/models/flat_llama.py

@@ -208,11 +208,12 @@ class FlatTransformer:
     return out, h, amaxs, saves
 
   @function(precompile=True, precompile_backward=True)
-  def run_layer(self, x:Tensor, freqs_cis:Tensor, attn_kwargs:dict, ffn_kwargs:dict):
+  def run_layer(self, x:Tensor, freqs_cis:Tensor, attn_kwargs:dict, ffn_kwargs:dict, save:bool=True):
     attn, attn_amaxs, attn_saves = self.attention(x, freqs_cis, **attn_kwargs)
     ffn, h, ffn_amaxs, ffn_saves = self.feed_forward(x, attn, **ffn_kwargs)
     h = h + ffn
-    return (h, *attn_amaxs, *ffn_amaxs, *attn_saves, *ffn_saves)
+    if save: return (h, *attn_amaxs, *ffn_amaxs, *attn_saves, *ffn_saves)
+    else: return (h, *attn_amaxs, *ffn_amaxs)
 
   def shard(self, device:tuple[str, ...], mp:bool=False):
     from tinygrad.nn.state import get_parameters
@@ -241,7 +242,7 @@ class FlatTransformer:
       for name in self._fp8_inv_scale:
         self._fp8_inv_scale[name] = self._fp8_inv_scale[name].to(device).contiguous().requires_grad_(False)
 
-  def __call__(self, tokens:Tensor):
+  def __call__(self, tokens:Tensor, save:bool=True):
     h = self.tok_embeddings(tokens)
     freqs_cis = self.freqs_cis.cast(h.dtype)[:, :tokens.shape[1], :, :, :]
     a, ga, s = self._fp8_amax, self._fp8_grad_amax, self._fp8_inv_scale
@@ -256,7 +257,7 @@ class FlatTransformer:
                           s_1=s["w1"][i], s_3=s["w3"][i], grad_amax_xw1=ga["xw1"][i], grad_amax_xw3=ga["xw3"][i])
       else:
         ffn_kwargs.update(w13=self.w13[i], amax_x13=a["x13"][i], s_13=s["w13"][i], grad_amax_xw13=ga["xw13"][i])
-      h, *ret = self.run_layer(h, freqs_cis, attn_kwargs, ffn_kwargs)
+      h, *ret = self.run_layer(h, freqs_cis, attn_kwargs, ffn_kwargs, save=save)
       amax_names = ["xqkv", "xo"] + (["x1", "x3"] if SPLIT_W13 else ["x13"]) + ["x2"]
       for name, new_val in zip(amax_names, ret[:len(amax_names)]):
         a[name][i].assign(new_val)
@@ -306,7 +307,7 @@ if __name__ == "__main__":
 
   # preallocate all the grad buffers and zero them out
   grads = {x:Tensor.zeros(x.shape, dtype=x.dtype, device=x.device).contiguous()
-           for x in state.values() if x.requires_grad is None}
+           for x in state.values() if x.requires_grad}
 
   # print model size
   sz = 0

examples/mlperf/models/test_flat_llama.py

@@ -3,7 +3,6 @@ os.environ["WQKV"] = "1"
 import unittest
 import numpy as np
 from tinygrad import Tensor, nn, dtypes
-from tinygrad.nn.state import get_parameters
 from tinygrad.device import is_dtype_supported, Device
 from examples.mlperf.models.llama import Transformer
 from examples.mlperf.models.flat_llama import FlatTransformer
@@ -45,8 +44,6 @@ class TestFlatLlama(unittest.TestCase):
     flat = FlatTransformer(**params)
     copy_weights(flat, ref)
 
-    for p in get_parameters(ref): p.requires_grad_(True)
-    for p in get_parameters(flat): p.requires_grad_(True)
     Tensor.realize(*nn.state.get_state_dict(flat).values())
 
     tokens = Tensor([[1, 50, 100, 999, 2, 10]])

examples/openpilot/compile3.py

@@ -21,6 +21,8 @@ def compile(onnx_file):
   # TODO this seems dumb
   input_types = {k:(dtypes.float32 if v is dtypes.float16 else v) for k,v in input_types.items()}
   Tensor.manual_seed(100)
+  # replace symbolic dimensions (e.g. 'b' for dynamic batch) with 1
+  input_shapes = {k:tuple(s if isinstance(s, int) else 1 for s in shp) for k,shp in input_shapes.items()}
   inputs = {k:Tensor(Tensor.randn(*shp, dtype=input_types[k]).mul(8).realize().numpy(), device='NPY') for k,shp in sorted(input_shapes.items())}
   if not getenv("NPY_IMG"):
     inputs = {k:Tensor(v.numpy(), device=Device.DEFAULT).realize() if 'img' in k else v for k,v in inputs.items()}

extra/mmapeak/qcom_fp16_mad_peak.py

@@ -0,0 +1,110 @@
+#!/usr/bin/env python3
+"""FP16/FP32 MAD peak repro for comparing DEV=CL and DEV=QCOM.
+
+Example:
+  DEV=CL python3 extra/mmapeak/qcom_fp16_mad_peak.py
+  DEV=QCOM python3 extra/mmapeak/qcom_fp16_mad_peak.py --dtype fp32
+"""
+from __future__ import annotations
+
+import argparse
+
+from tinygrad import Device, dtypes
+from tinygrad.device import Buffer
+
+
+MAD_OPS_PER_LOOP = 16
+VEC = 16
+
+
+def kernel_name(dtype:str) -> str:
+  return f"{dtype}_mad_peak"
+
+
+def make_kernel(loops:int, dtype:str="fp16") -> str:
+  assert dtype in {"fp16", "fp32"}
+  scalar = "half" if dtype == "fp16" else "float"
+  vec_type = f"{scalar}{VEC}"
+  prefix = "#pragma OPENCL EXTENSION cl_khr_fp16 : enable\n" if dtype == "fp16" else ""
+  cast = "(half)" if dtype == "fp16" else ""
+  suffix = "f"
+  mad_block = "\n".join([
+    "    x = mad(y, x, y);",
+    "    y = mad(x, y, x);",
+  ] * (MAD_OPS_PER_LOOP // 2))
+
+  x_init = ",\n    ".join(f"bx + {cast}{(i + 1) * 0.001:.3f}{suffix}" for i in range(VEC))
+  y_init = ",\n    ".join(f"by + {cast}{(i + 17) * 0.001:.3f}{suffix}" for i in range(VEC))
+  sum_terms = " + ".join([f"x.s{'0123456789abcdef'[i]}" for i in range(VEC)] +
+                         [f"y.s{'0123456789abcdef'[i]}" for i in range(VEC)])
+  return f"""{prefix}__kernel void {kernel_name(dtype)}(__global {scalar} *out) {{
+  int lid = get_local_id(0);
+  int gid = get_group_id(0);
+  {scalar} bx = {cast}1.0f + {cast}(lid & 15) * {cast}0.001f;
+  {scalar} by = {cast}1.0f + {cast}(gid & 15) * {cast}0.001f;
+  {vec_type} x = ({vec_type})(
+    {x_init});
+  {vec_type} y = ({vec_type})(
+    {y_init});
+
+  for (int i = 0; i < {loops}; i++) {{
+{mad_block}
+  }}
+
+  out[get_global_id(0)] = {sum_terms};
+}}"""
+
+
+def run(args:argparse.Namespace) -> None:
+  dev = Device[Device.DEFAULT]
+  renderer = type(dev.renderer).__name__
+  if renderer == "IR3Renderer":
+    raise SystemExit("This repro uses OpenCL source. Use DEV=QCOM or DEV=CL, not DEV=QCOM:IR3.")
+
+  dtype = args.dtype
+  dt = dtypes.half if dtype == "fp16" else dtypes.float
+  src = make_kernel(args.loops, dtype)
+  if args.print_source: print(src)
+  lib = dev.compiler.compile_cached(src)
+  if args.disasm: dev.compiler.disassemble(lib)
+
+  # Runtime aux mirrors OpenCLRenderer.aux: one __global output pointer at kernel arg 0.
+  global_size = (args.groups, 1, 1)
+  local_size = (args.local, 1, 1)
+  workitems = args.groups * args.local
+  flops = workitems * args.loops * MAD_OPS_PER_LOOP * VEC * 2
+
+  prg = dev.runtime(kernel_name(dtype), lib, (((0, dt.ptr()),),))
+  out = Buffer(dev.device, workitems, dt, preallocate=True)
+
+  for _ in range(args.warmup):
+    prg(out._buf, global_size=global_size, local_size=local_size, wait=True)
+
+  times = [prg(out._buf, global_size=global_size, local_size=local_size, wait=True) for _ in range(args.iters)]
+  best = min(t for t in times if t is not None)
+  out_bits = out.copyout(memoryview(bytearray(out.nbytes))).cast("H" if dtype == "fp16" else "I")[0]
+  out_fmt = "04x" if dtype == "fp16" else "08x"
+
+  print(f"device={dev.device} renderer={renderer} arch={dev.arch}")
+  print(f"dtype={dtype} groups={args.groups} local={args.local} workitems={workitems} loops={args.loops} flops={flops}")
+  print(f"best={best*1e6:.2f} us  {dtype}_mad_peak={flops / best * 1e-9:.2f} GFLOPS  out0=0x{out_bits:{out_fmt}}")
+  if args.show_times:
+    print("times_us=" + ",".join(f"{t*1e6:.2f}" for t in times if t is not None))
+
+
+def main() -> None:
+  parser = argparse.ArgumentParser(description="FP16/FP32 MAD peak repro for DEV=CL vs DEV=QCOM")
+  parser.add_argument("--dtype", choices=("fp16", "fp32"), default="fp16", help="MAD datatype")
+  parser.add_argument("--groups", type=int, default=2048, help="number of workgroups")
+  parser.add_argument("--local", type=int, default=256, help="workitems per workgroup")
+  parser.add_argument("--loops", type=int, default=8, help="inner loop count; default matches clpeak vec16")
+  parser.add_argument("--warmup", type=int, default=2, help="warmup launches")
+  parser.add_argument("--iters", type=int, default=10, help="timed launches")
+  parser.add_argument("--show-times", action="store_true", help="print every timed launch")
+  parser.add_argument("--print-source", action="store_true", help="print generated OpenCL source")
+  parser.add_argument("--disasm", action="store_true", help="call the tinygrad compiler disassembler after compile")
+  run(parser.parse_args())
+
+
+if __name__ == "__main__":
+  main()

extra/remote/serve.py

@@ -84,8 +84,6 @@ def serve(conn:socket.socket):
       conn.sendall(resp_err(str(e)))
 
 if __name__ == "__main__":
-  if not OSX: System.reserve_hugepages(128) # for sysmem allocations
-
   port = int(sys.argv[1]) if len(sys.argv) > 1 else 6667
   server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
   server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

extra/torch_backend/test_kernel_fusion.py

@@ -105,7 +105,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       view = x[1:3]
       view += 1.0
       return x.sum()
-    self._check_kernel_count(fn, 8)
+    self._check_kernel_count(fn, 7)
 
   def test_batchnorm_running_stats_update(self):
     def fn():

extra/viz/kernel_graph.py

@@ -7,15 +7,15 @@ def get_node(graph:dict, key): return graph[str(key)]
 
 if __name__ == "__main__":
   parser = argparse.ArgumentParser(description="print CALL graph from DEBUG=5 tinygrad.viz.cli --json output")
-  parser.add_argument("kernel", type=str, default=None, help="Kernel name to stop at (default: print all kernels)")
+  parser.add_argument("kernel", type=str, nargs="?", default="ALL", metavar="NAME", help="Kernel name to stop at (default: print all kernels)")
   args = parser.parse_args()
   ref:int|None = None
   for line in sys.stdin:
     if not line.strip(): continue
     graph = json.loads(line)
-    if ref is not None and graph.get("ref") == ref:
+    if graph.get("ref") is not None and (args.kernel == "ALL" or graph["ref"] == ref):
       print(graph)
-      if (v:=json.loads(next(sys.stdin)).get("value")): print(v)
+      if (v:=json.loads(next(sys.stdin, "{}")).get("value")): print(v)
     if ref is not None or not isinstance(rec:=next(iter(graph.values()), {}), dict) or "label" not in rec: continue
     for v in graph.values():
       if not v["label"].startswith("CALL"): continue
@@ -39,6 +39,6 @@ if __name__ == "__main__":
           src_str = ["SRC"]+get_node(graph, get_node(graph, s)["src"][0][1])["label"].splitlines()[1:]
           print(" ".join(idx_str+src_str))
         ss += [x[1] for x in get_node(graph, s)["src"]]
-      if args.kernel is not None and args.kernel in ansistrip(v["label"]):
+      if args.kernel != "ALL" and args.kernel in ansistrip(v["label"]):
         ref = v["ref"]
         break

spec/tinyspec.tex

@@ -16,6 +16,7 @@
 \definecolor{elwyellow}{HTML}{F9A825}
 \definecolor{callblue}{HTML}{1565C0}
 \definecolor{assignbrown}{HTML}{795548}
+\definecolor{loadred}{HTML}{c08080}
 \definecolor{multipurple}{HTML}{7B1FA2}
 \definecolor{markerorange}{HTML}{E65100}
 % AxisType colors (from tinygrad)
@@ -48,16 +49,16 @@ All nodes in the tinygrad graph are \textbf{UOps}. A UOp is a tuple $(\mathrm{op
 \toprule
 \textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
 \midrule
+\op{Param} & $(\mathbf{s})$ & slot, dtype, device?, addrspace? &
+  Placeholder with shape $\mathbf{s}$. Substituted in \op{Function}. \\[4pt]
 \op{Buffer} & () & size, dtype, device, addrspace &
   Shape $(n \cdot \textit{size},)$ if device is $n$-tuple, else $(\textit{size},)$. \\
 \op{BufferView} & (buf,) & size, dtype, offset &
   Typed access into a buffer. Zero-copy $(\textit{size},)$ slice at offset; inherits addrspace. \\
-\op{Param} & $(\mathbf{s})$ or $(\mathbf{s}, \text{min}, \text{max})$ & slot, dtype, device? &
-  Placeholder with shape $\mathbf{s}$. Substituted in \op{Function}. \\[4pt]
 \op{Const} & () & value, dtype &
   A scalar constant with shape $(\ )$. \\
-\op{Vconst} & () & values, dtype &
-  A vector constant with shape $(n,)$. \\
+  & & & Form vector consts with \op{Stack} \\
+\op{Binary} & () & data & Raw binary data, has dtype uint8 and shape len($data$)  \\
 \bottomrule
 \end{tabular}
 
@@ -90,7 +91,7 @@ A \op{Buffer}'s \textbf{addrspace} is \texttt{GLOBAL}, \texttt{LOCAL}, or \textt
 \toprule
 \textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
 \midrule
-\op{Reduce} & $(T,)$ & op, axes & Reduce $T$ along axes. Op is \op{Add}, \op{Max}, or \op{Mul}. \\
+\op{Reduce} & ($T$, $r_0$, $r_1$, \ldots) & op, axes & Reduce $T$ along axes or ranges. Op is \op{Add}, \op{Max}, or \op{Mul}. \\
 \bottomrule
 \end{tabular}
 
@@ -109,13 +110,25 @@ A \op{Buffer}'s \textbf{addrspace} is \texttt{GLOBAL}, \texttt{LOCAL}, or \textt
 \end{tabular}
 
 %% ============================================================
-\subsection*{{\color{multipurple}Store Ops} \normalfont\small--- side effects}
+\subsection*{{\color{loadred}Load Ops} \normalfont\small--- can change device or addrspace}
 
 \begin{tabular}{@{}l l l l@{}}
 \toprule
 \textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
 \midrule
-\op{Store} & (buf, val, gate?) & --- & Write val into buf. buf.shape $=$ val.shape. \\
+\op{Load} & (buf, alt?, gate?) & device, addrspace & Read (pull) from buffer into a new anonymous buffer. \\
+  & & & Note: this replaces \op{Copy} and \op{Contiguous}. \\
+\bottomrule
+\end{tabular}
+
+%% ============================================================
+\subsection*{{\color{multipurple}Store Ops} \normalfont\small--- the only op with observable side effects}
+
+\begin{tabular}{@{}l l l l@{}}
+\toprule
+\textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
+\midrule
+\op{Store} & (buf, val, gate?) & --- & Write (push) val into buf. buf.shape $=$ val.shape. \\
   & & & If gate is present, write only when gate is true. Output is void. \\
 \bottomrule
 \end{tabular}
@@ -205,7 +218,6 @@ Ternary & $(P, A, B)$
 \op{Contiguous} & $(T,)$ & --- & Force contiguous memory layout. \\
 \op{ContiguousBackward} & $(T,)$ & --- & Force contiguous in backward pass. \\
 \op{Detach} & $(T,)$ & --- & Stops gradient propagation. \\
-\op{Copy} & $(T,)$ & device & Copy to target device. \\
 \bottomrule
 \end{tabular}
 
@@ -216,8 +228,6 @@ Ternary & $(P, A, B)$
 \toprule
 \textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
 \midrule
-\op{Load} & (idx,alt?,gate?) & --- & Dereference: read element at index from buffer. \\
- & & & All loads will be replaced by \op{Store}. \\
 \op{Barrier} & (deps\ldots) & --- & Synchronize threads within a workgroup. \\
 \op{Ins} & \ldots & \ldots & A single machine instruction (e.g.\ AMD ISA). \\
 \op{Special} & (bound,) & name & GPU thread/workgroup index (e.g.\ \texttt{gidx0}, \texttt{lidx1}). \\

test/backend/test_asm_gemm.py

@@ -12,8 +12,8 @@ def is_cdna4(): return Device[Device.DEFAULT].renderer.target.arch.startswith("g
 
 def run_asm_gemm(a_shape, b_shape, dtype=dtypes.float16, a_shard=None, b_shard=None, gpus:int=1) -> None:
   Tensor.manual_seed(0)
-  a_rand = Tensor.randn(a_shape, dtype=dtypes.float).sub(0.5).cast(dtype)
-  b_rand = Tensor.randn(b_shape, dtype=dtypes.float).sub(0.5).cast(dtype)
+  a_rand = Tensor.randn(a_shape, dtype=dtypes.float, requires_grad=False).sub(0.5).cast(dtype)
+  b_rand = Tensor.randn(b_shape, dtype=dtypes.float, requires_grad=False).sub(0.5).cast(dtype)
   with Context(DEBUG=0):
     Tensor.realize(a_rand, b_rand)
 

test/backend/test_dtype.py

@@ -330,10 +330,6 @@ class TestBitCast(unittest.TestCase):
       # should fail because 3 int8 is 3 bytes but float16 is two and 3 isn't a multiple of 2
       Tensor.empty((3,), dtype=dtypes.int8).bitcast(dtypes.float16)
 
-    with self.assertRaises(RuntimeError):
-      # should fail because backprop through bitcast is undefined
-      Tensor.empty((4,), dtype=dtypes.int8, requires_grad=True).bitcast(dtypes.float16)
-
   def test_bitcast_float_to_int32(self):
     a = Tensor([1.,2,3])
     b = a.bitcast(dtypes.int32)

test/backend/test_rangeify.py

@@ -109,9 +109,9 @@ def fa():
 def fa_bw():
   Tensor.manual_seed(1337)
   with Context(DEBUG=0):
-    q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize().requires_grad_() for _ in range(3)]
+    q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize() for _ in range(3)]
     attn_output = nn.Linear(HEADS*EMB, HEADS*EMB, bias=False)
-    attn_output.weight.requires_grad_().realize()
+    attn_output.weight.realize()
     target = Tensor.rand(BS, SEQLEN, HEADS*EMB).contiguous().realize()
 
   GlobalCounters.reset()

test/backend/test_schedule.py

@@ -238,19 +238,9 @@ class TestSchedule(unittest.TestCase):
       run_linear(*check_schedule(out, 4))
     np.testing.assert_allclose(out.numpy(), (x.numpy() - x.numpy().max(keepdims=True)).max())
 
-  @unittest.skip("these two Tensors are the same")
-  def test_example_matmul(self):
-    x = Tensor.eye(64, requires_grad=True)
-    y = Tensor.eye(64, requires_grad=True)
-    z = y.matmul(x).sum()
-    z.backward()
-    out = x.grad.contiguous()
-    run_linear(*check_schedule(out, 1))
-    np.testing.assert_allclose(out.numpy(), np.ones((64,64)))
-
   def test_example_matmul_contig(self):
-    x = Tensor.eye(64, requires_grad=True).contiguous().realize()
-    y = Tensor.eye(64, requires_grad=True).contiguous().realize()
+    x = Tensor.eye(64).contiguous().realize()
+    y = Tensor.eye(64).contiguous().realize()
     z = y.matmul(x).sum()
     z.backward()
     out = x.grad.contiguous()
@@ -258,7 +248,7 @@ class TestSchedule(unittest.TestCase):
     np.testing.assert_allclose(out.numpy(), np.ones((64,64)))
 
   def test_example_matmul_same(self):
-    x = Tensor.eye(64, requires_grad=True)
+    x = Tensor.eye(64)
     z = x.matmul(x).sum()
     z.backward()
     out = x.grad.contiguous()

test/backend/test_setitem.py

@@ -344,6 +344,28 @@ class TestWithGrad(unittest.TestCase):
     with self.assertRaises(RuntimeError):
       z[:2] = Tensor([0.0, 0.0])
 
+  def test_setitem_raises_with_unrealized_downstream(self):
+    x = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
+    _y = x * 2.0
+    with self.assertRaises(RuntimeError):
+      x[0] = 99.0
+
+  def test_setitem_raises_on_unrealized_compute_base(self):
+    # y has a compute (unrealized) base; tmp is a view of y. eager: tmp would follow y's mutation. lazy: tmp keeps the old MUL graph.
+    x = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
+    y = x * 2.0
+    _tmp = y[:1]
+    with self.assertRaises(RuntimeError):
+      y[0] = 99.0
+
+  def test_setitem_raises_on_aliased_uop(self):
+    # two Tensor objects sharing the exact same unrealized uop. setitem on one updates its uop, the other keeps the stale graph reference.
+    x = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
+    y = x * 2.0
+    _z = Tensor(y.uop)
+    with self.assertRaises(RuntimeError):
+      y[0] = 99.0
+
 class TestSetitemLoop(unittest.TestCase):
   def test_arange(self):
     N = 10

test/backend/test_softmax_fusion.py

@@ -190,7 +190,6 @@ class TestSoftmaxFusion(unittest.TestCase):
 
   def test_softmax_bw(self):
     print("*** softmax bw ***")
-    self.test.requires_grad_()
     with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
       self.test.softmax(-1).sum().backward()
       sg = self.test.grad.realize()

test/backend/test_tensor.py

@@ -179,8 +179,7 @@ class TestTinygrad(unittest.TestCase):
     def test_tinygrad():
       w1 = Tensor(init)
       w2 = Tensor(init)
-      assert w1.requires_grad is None and w2.requires_grad is None
-      # optimizer sets requires_grad=True for params with requires_grad=None
+      assert w1.requires_grad is True and w2.requires_grad is True
       nn.optim.SGD([w1, w2], lr=0.01)
       assert w1.requires_grad is True and w2.requires_grad is True
       out = w1.add(w2)
@@ -599,7 +598,7 @@ class TestMoveTensor(unittest.TestCase):
     assert x is y
 
   def test_to_grad(self):
-    x = Tensor.eye(3, requires_grad=True, device=self.d0)
+    x = Tensor.eye(3, device=self.d0)
     y = Tensor([[2.0,0,-2.0]], requires_grad=True, device=self.d0)
     z = y.matmul(x).to(self.d1).sum()
     z.backward()

test/device/test_metal.py

@@ -59,8 +59,3 @@ kernel void r_5(device int* data0, const device int* data1, uint3 gid [[threadgr
     self.assertEqual(curr:=device.sysdevice.currentAllocatedSize(), before+size, msg=f"{curr=} - {before=}")
     device.allocator.free(buf, buf.size, BufferSpec(nolru=True))
     self.assertEqual(curr:=device.sysdevice.currentAllocatedSize(), before, msg=f"{curr=} - {before=}")
-
-  def test_gpu_family(self):
-    device = Device['METAL']
-    self.assertGreater(device.gpu_family, 0)
-    self.assertLessEqual(device.gpu_family, 15)

test/external/external_test_example.py

@@ -40,7 +40,7 @@ class TestExample(unittest.TestCase):
 
   @multidevice_test
   def test_example_readme(self, device):
-    x = Tensor.eye(3, device=device, requires_grad=True)
+    x = Tensor.eye(3, device=device)
     y = Tensor([[2.0,0,-2.0]], device=device, requires_grad=True)
     z = y.matmul(x).sum()
     z.backward()
@@ -59,8 +59,8 @@ class TestExample(unittest.TestCase):
       print(f"WARNING: {device} test isn't running")
       return
 
-    x = Tensor.eye(8, device=device, requires_grad=True)
-    y = Tensor.eye(8, device=device, requires_grad=True)
+    x = Tensor.eye(8, device=device)
+    y = Tensor.eye(8, device=device)
     z = y.matmul(x).sum()
     z.backward()
 

test/external/external_test_onnx_runner.py

@@ -1,6 +1,6 @@
 import unittest, onnx, tempfile, pathlib
 import numpy as np
-from tinygrad import dtypes, Tensor
+from tinygrad import Tensor
 from tinygrad.uop.ops import Ops
 from tinygrad.device import is_dtype_supported
 from typing import Any
@@ -96,16 +96,7 @@ class TestOnnxRunnerDtypes(unittest.TestCase):
   Internal tensors (initializers, attributes) fallback to default dtype if unsupported by device.
   External tensors (inputs) preserve their original dtype - user must ensure compatibility with device.
   """
-  def _get_expected_dtype(self, onnx_dtype: int, is_input: bool):
-    true_dtype = OnnxDataType(onnx_dtype).to_dtype()
-    # inputs always preserve their true dtype.
-    if is_input:
-      return true_dtype
-    # supported types are always themselves.
-    if onnx_dtype in device_supported_dtypes:
-      return true_dtype
-    # otherwise it's an unsupported dtype that's internal to the ONNX model, which should fallback to default.
-    return dtypes.default_int if dtypes.is_int(true_dtype) else dtypes.default_float
+  def _get_expected_dtype(self, onnx_dtype: int, is_input: bool): return OnnxDataType(onnx_dtype).to_dtype()
 
   @given(onnx_dtype=st.sampled_from(all_dtypes))
   def test_input_dtype(self, onnx_dtype: int):

test/external/external_uop_gc.py

@@ -24,27 +24,27 @@ def two_plus_two_linearize():
 def two_plus_two_realize(): (Tensor([2])+Tensor([2])).realize()
 def two_plus_two_item(): (Tensor([2])+Tensor([2])).item()
 def gradient_test():
-  x = Tensor.eye(3, requires_grad=True)
+  x = Tensor.eye(3)
   y = Tensor([[2.0,0,-2.0]], requires_grad=True)
   z = y.matmul(x).sum()
   z.backward()
 def realized_eye():
-  Tensor.eye(3, requires_grad=True).realize()
+  Tensor.eye(3).realize()
 def realized_list():
   Tensor([[2.0,0,-2.0]], requires_grad=True).realize()
 def kernel_matmul():
-  x = Tensor.eye(3, requires_grad=True)
+  x = Tensor.eye(3)
   y = Tensor([[2.0,0,-2.0]], requires_grad=True)
   z = y.matmul(x)
   linear = z.schedule_linear()
   to_program(linear.src[-1].src[0], Device.default.renderer)
 def realized_matmul():
-  x = Tensor.eye(3, requires_grad=True)
+  x = Tensor.eye(3)
   y = Tensor([[2.0,0,-2.0]], requires_grad=True)
   z = y.matmul(x)
   Tensor.realize(z)
 def realized_gradient():
-  x = Tensor.eye(3, requires_grad=True)
+  x = Tensor.eye(3)
   y = Tensor([[2.0,0,-2.0]], requires_grad=True)
   z = y.matmul(x).sum()
   z.backward()

test/null/test_tensor_metadata.py

@@ -118,7 +118,7 @@ class TestTraceMetaShutdown(unittest.TestCase):
   def test_tracemeta_del_no_shutdown_error(self):
     import subprocess, os
     result = subprocess.run(['python3', '-c', 'from tinygrad import Tensor\n'
-                             'x=Tensor.eye(3,requires_grad=True); (x@x).sum().backward()'],
+                             'x=Tensor.eye(3); (x@x).sum().backward()'],
                             env={**os.environ, "TRACEMETA": "2"}, capture_output=True)
     self.assertEqual(result.returncode, 0)
     self.assertNotIn(b"Exception", result.stderr)

test/test_tiny.py

@@ -141,7 +141,6 @@ class TestTiny(unittest.TestCase):
     Tensor.realize(*[p.replace(Tensor.ones_like(p).contiguous()) for p in nn.state.get_parameters(layers)])
 
     # realize gradients
-    for x in nn.state.get_parameters(layers): x.requires_grad_()
     Tensor.empty(4, 1, 14, 14).sequential(layers).sum().backward()
     Tensor.realize(*[x.grad for x in nn.state.get_parameters(layers) if x.grad is not None])
 

test/testextra/test_hk_fa.py

@@ -142,9 +142,9 @@ class TestFA(unittest.TestCase):
       base_do = Tensor.ones(B, N, H, D, dtype=dtypes.float32).contiguous()
 
     with Context(DEBUG=0):
-      q = base_q.clone().requires_grad_(True).shard(GPUS, axis=0)
-      k = base_k.clone().requires_grad_(True).shard(GPUS, axis=0)
-      v = base_v.clone().requires_grad_(True).shard(GPUS, axis=0)
+      q = base_q.clone().shard(GPUS, axis=0)
+      k = base_k.clone().shard(GPUS, axis=0)
+      v = base_v.clone().shard(GPUS, axis=0)
       Tensor.realize(q, k, v)
 
       do = base_do.clone().shard(GPUS, axis=0)
@@ -157,9 +157,9 @@ class TestFA(unittest.TestCase):
     Tensor.realize(q.grad, k.grad, v.grad)
 
     with Context(DEBUG=0):
-      q_ref = base_q.clone().requires_grad_(True)
-      k_ref = base_k.clone().requires_grad_(True)
-      v_ref = base_v.clone().requires_grad_(True)
+      q_ref = base_q.clone()
+      k_ref = base_k.clone()
+      v_ref = base_v.clone()
       Tensor.realize(q_ref, k_ref, v_ref)
 
       do_ref = base_do.clone()
@@ -189,9 +189,9 @@ class TestFA(unittest.TestCase):
       base_do = Tensor.ones(B, N, H, D, dtype=dtypes.float32).contiguous()
 
     with Context(DEBUG=0):
-      q = base_q.clone().requires_grad_(True).shard(GPUS, axis=2)
-      k = base_k.clone().requires_grad_(True).shard(GPUS, axis=2)
-      v = base_v.clone().requires_grad_(True).shard(GPUS, axis=2)
+      q = base_q.clone().shard(GPUS, axis=2)
+      k = base_k.clone().shard(GPUS, axis=2)
+      v = base_v.clone().shard(GPUS, axis=2)
       Tensor.realize(q, k, v)
 
       do = base_do.clone().shard(GPUS, axis=2)
@@ -204,9 +204,9 @@ class TestFA(unittest.TestCase):
     Tensor.realize(q.grad, k.grad, v.grad)
 
     with Context(DEBUG=0):
-      q_ref = base_q.clone().requires_grad_(True)
-      k_ref = base_k.clone().requires_grad_(True)
-      v_ref = base_v.clone().requires_grad_(True)
+      q_ref = base_q.clone()
+      k_ref = base_k.clone()
+      v_ref = base_v.clone()
       Tensor.realize(q_ref, k_ref, v_ref)
 
       do_ref = base_do.clone()

test/testextra/test_tk.py

@@ -951,9 +951,9 @@ class TestTK(unittest.TestCase):
       base_do = Tensor.ones(B, N, H, D, dtype=dtypes.float32).contiguous()
 
     with Context(DEBUG=0):
-      q = base_q.clone().requires_grad_(True).shard(GPUS, axis=0)
-      k = base_k.clone().requires_grad_(True).shard(GPUS, axis=0)
-      v = base_v.clone().requires_grad_(True).shard(GPUS, axis=0)
+      q = base_q.clone().shard(GPUS, axis=0)
+      k = base_k.clone().shard(GPUS, axis=0)
+      v = base_v.clone().shard(GPUS, axis=0)
       Tensor.realize(q, k, v)
 
       do = base_do.clone().shard(GPUS, axis=0)
@@ -966,9 +966,9 @@ class TestTK(unittest.TestCase):
     Tensor.realize(q.grad, k.grad, v.grad)
 
     with Context(DEBUG=0):
-      q_ref = base_q.clone().requires_grad_(True)
-      k_ref = base_k.clone().requires_grad_(True)
-      v_ref = base_v.clone().requires_grad_(True)
+      q_ref = base_q.clone()
+      k_ref = base_k.clone()
+      v_ref = base_v.clone()
       Tensor.realize(q_ref, k_ref, v_ref)
 
       do_ref = base_do.clone()

test/unit/test_disk_tensor.py

@@ -92,10 +92,6 @@ class TestRawDiskBuffer(unittest.TestCase):
       # should fail because 3 int8 is 3 bytes but float16 is two and 3 isn't a multiple of 2
       Tensor.empty((3,), dtype=dtypes.int8, device=f"DISK:{tmp}").bitcast(dtypes.float16)
 
-    with self.assertRaises(RuntimeError):
-      # should fail because backprop through bitcast is undefined
-      Tensor.empty((4,), dtype=dtypes.int8, requires_grad=True, device=f"DISK:{tmp}").bitcast(dtypes.float16)
-
     pathlib.Path(tmp).unlink()
 
 @unittest.skipUnless(is_dtype_supported(dtypes.uint8), "need uint8")

test/unit/test_gradient.py

@@ -69,20 +69,26 @@ class TestTensorGradient(unittest.TestCase):
     np.testing.assert_allclose(x.grad.numpy(), [2.0+3.0+2*3.0])
     self.assertIs(x.grad, old_grad)
 
-  def test_gradient_through_clone(self):
-    src = Tensor([1.0, 2.0, 3.0, 4.0])
+  def test_gradient_through_clone_from_non_grad_src(self):
+    src = Tensor([1.0, 2.0, 3.0, 4.0], requires_grad=False)
     x = src.clone().requires_grad_(True)
     (x * 2.0).sum().backward()
     np.testing.assert_allclose(x.grad.numpy(), [2.0, 2.0, 2.0, 2.0])
     self.assertIsNone(src.grad)
 
+  def test_gradient_through_clone_from_grad_src(self):
+    # unlike torch, tinygrad accumulates grad on all requires_grad tensors, including non-leaf x
     src = Tensor([1.0, 2.0, 3.0, 4.0], requires_grad=True)
-    x = src.clone().requires_grad_(True)
-    try:
-      (x * 2.0).sum().backward()
-    except RuntimeError:
-      # TODO: this crashes now
-      pass
+    x = src.clone()
+    (x * 2.0).sum().backward()
+    np.testing.assert_allclose(src.grad.numpy(), [2.0, 2.0, 2.0, 2.0])
+    np.testing.assert_allclose(x.grad.numpy(), [2.0, 2.0, 2.0, 2.0])
+
+  def test_setitem_on_grad_used_tensor_raises(self):
+    x = Tensor([1.0, 2.0, 3.0, 4.0], requires_grad=True).realize()
+    _ = (x * 2.0).sum()
+    with self.assertRaises(RuntimeError):
+      x[0] = 99.0
 
   def test_gradient_through_chained_unrealized_setitem(self):
     g1 = Tensor.zeros(4).contiguous()

test/unit/test_linalg.py

@@ -42,7 +42,6 @@ class TestLinAlg(unittest.TestCase):
   def test_svd_nonfull_5_3(self): self._test_svd_nonfull((5,3))
   def test_svd_nonfull_3_5(self): self._test_svd_nonfull((3,5))
   def test_svd_nonfull_2_2_2_2_3(self): self._test_svd_nonfull((2,2,2,2,3))
-  def test_svd_nonfull_5_5(self): self._test_svd_nonfull((5,5))
 
   @unittest.skip("very big. recommend wrapping with TinyJit around inner function")
   def test_svd_large(self):

tinygrad/device.py

@@ -342,7 +342,7 @@ def is_dtype_supported(dtype:DType, target:Target|None=None) -> bool:
   target = target or DEV.target(Device.DEFAULT)
   if dtype == dtypes.bfloat16:
     match target.device:
-      case "METAL": return not CI or BENCHMARKS
+      case "METAL": target.arch.startswith("Apple") and int(target.arch[5:]) >= 6
       case "CUDA": return (not CI or BENCHMARKS) and target.renderer != "PTX"
       case "NV": return (not CI or BENCHMARKS) and target.renderer not in ("PTX", "NAK")
       case "CPU": return (not CI or BENCHMARKS) and platform.machine() in {"arm", "arm64", "aarch64", "x86_64", "amd64"} and target.renderer != "LVP"
@@ -364,7 +364,7 @@ def is_dtype_supported(dtype:DType, target:Target|None=None) -> bool:
   # PYTHON supports half memoryview in 3.12+ https://github.com/python/cpython/issues/90751
   if dtype == dtypes.half:
     match target.device:
-      case "CL": return (not CI or BENCHMARKS) and not OSX
+      case "CL": return "cl_khr_fp16" in target.arch
       case "QCOM": return bool(IMAGE) and bool(FLOAT16) # QCOM compiler is flaky with half
       case "CUDA" | "NV": return not CI or BENCHMARKS or target.renderer == "PYTHON"
       case "CPU" if target.renderer == "LLVM": return OSX

tinygrad/gradient.py

@@ -14,13 +14,6 @@ def reduce_gradient(ctx:UOp, ret:UOp, op:Ops):
     return ((mask/broadcast_to_input(count)) * broadcast_to_input(ctx),)
   if op == Ops.MUL: return (broadcast_to_input(ctx * ret) / ret.src[0],)
 
-def unbroadcast(ctx:UOp, shape:tuple|None) -> UOp:
-  if ctx._shape is None or shape is None or ctx.shape == shape: return ctx
-  if len(shape) > len(ctx.shape): raise RuntimeError(f"can't unbroadcast {ctx.shape} to {shape}")
-  aligned = (1,)*(len(ctx.shape)-len(shape)) + shape
-  axis = tuple(i for i,(s,n) in enumerate(zip(aligned, ctx.shape)) if s != n)
-  return ctx.cast(sum_acc_dtype(ctx.dtype))._rop(Ops.ADD, axis).cast(ctx.dtype).reshape(shape)
-
 def _compact_params(body:UOp, all_args:tuple[UOp, ...]) -> tuple[UOp, tuple[UOp, ...]]:
   """Remove unused PARAMs from body and return compacted (body, args)."""
   used = sorted({p.arg: p for p in body.toposort() if p.op is Ops.PARAM}.items())
@@ -73,7 +66,9 @@ pm_gradient = PatternMatcher([
   (UPat(Ops.CONTIGUOUS), lambda ctx: (ctx,)),
   (UPat(Ops.CONTIGUOUS_BACKWARD), lambda ctx: (ctx.contiguous(),)),
   (UPat(Ops.RESHAPE, name="ret"), lambda ctx, ret: (ctx.reshape(ret.src[0].shape), None)),
-  (UPat(Ops.EXPAND, name="ret"), lambda ctx, ret: (unbroadcast(ctx, ret.src[0]._shape), None)),
+  (UPat(Ops.EXPAND, name="ret"), lambda ctx, ret:
+    (ctx.cast(sum_acc_dtype(ctx.dtype))._rop(Ops.ADD, tuple(i for i,(s,n) in enumerate(zip(ret.src[0].shape, ret.shape)) if s!=n))
+     .cast(ctx.dtype), None)),
   (UPat(Ops.PAD, name="ret"), lambda ctx, ret: (ctx.shrink(tuple([(p[0], s+p[0]) for s,p in zip(ret.src[0].shape, ret.marg)])), None, None)),
   (UPat(Ops.SHRINK, name="ret"), lambda ctx, ret: (ctx.pad(tuple([(p[0], s-p[1]) for s,p in zip(ret.src[0].shape, ret.marg)])), None, None)),
   (UPat(Ops.PERMUTE, name="ret"), lambda ctx, ret: (ctx.permute(argsort(ret.marg)),)),
@@ -83,6 +78,9 @@ pm_gradient = PatternMatcher([
   (UPat(Ops.TUPLE), lambda ctx: ctx.src),
   (UPat(Ops.AFTER, src=(UPat.var("d"), UPat(Ops.CALL, name="k"))), lambda ctx, d, k:
     (ctx, UOp.maketuple(*(ctx if i == k.src.index(d)-1 else UOp(Ops.NOOP) for i in range(len(k.src)-1))))),
+  # clone/assign gradient passes through to val
+  (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE))), lambda ctx: (None, ctx)),
+  (UPat(Ops.STORE, src=(UPat(), UPat())), lambda ctx: (None, ctx)),
   # there's no gradient for bitcast
   (UPat(Ops.BITCAST), lambda: (None,)),
 ])
@@ -119,7 +117,6 @@ def compute_gradient(root:UOp, root_grad:UOp, targets:set[UOp]) -> dict[UOp, UOp
     assert len(lgrads) == len(t0.src), f"got {len(lgrads)} gradient, expected {len(t0.src)}"
     for k,v in zip(t0.src, lgrads):
       if v is None: continue
-      v = unbroadcast(v, k._shape)
       if k in grads and grads[k].op is not Ops.NOOP:
         if v.op is Ops.TUPLE and grads[k].op is Ops.TUPLE:
           grads[k] = UOp.maketuple(*(p + n if (p.op is not Ops.NOOP and n.op is not Ops.NOOP) else

tinygrad/mixin/__init__.py

@@ -5,7 +5,7 @@ from tinygrad.mixin.elementwise import ElementwiseMixin
 from tinygrad.mixin.movement import MovementMixin
 from tinygrad.mixin.reduce import ReduceMixin
 from tinygrad.uop import Ops
-from tinygrad.uop.ops import resolve, smax, smin, identity_element
+from tinygrad.uop.ops import _broadcast_shape, resolve, smax, smin, identity_element
 from tinygrad.dtype import ConstType, DType, DTypeLike, Invalid, InvalidType, PtrDType, PyConst, dtypes, least_upper_dtype, sum_acc_dtype, to_dtype
 from tinygrad.helpers import all_int, argfix, ceildiv, flatten, flat_to_grouped, make_tuple, prod, resolve_pool_pads, round_up
 
@@ -136,6 +136,19 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
 
   @classmethod
   def eye(cls, n:int, m:int|None=None, dtype:DTypeLike|None=None, device:str|tuple[str, ...]|None=None) -> Self:
+    """
+    Returns a 2-D tensor with `n` rows and `m` columns, with ones on the diagonal and zeros elsewhere.
+
+    You can pass in `dtype` and `device` keyword arguments to control the data type and device of the tensor.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor.eye(3).numpy())
+    ```
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor.eye(2, 4).numpy())
+    ```
+    """
     m_ = n if m is None else m
     if n < 0 or m_ < 0: raise ValueError(f"cannot have negative {n=}, {m_=}")
     out_dtype = to_dtype(dtype) if dtype is not None else dtypes.default_float
@@ -306,6 +319,11 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
   def _broadcasted(self, y, reverse=False) -> tuple[Self, Self]:
     if not isinstance(y, type(self)): y = self.ufix(y)
     x, y = (self, y) if not reverse else (y, self)
+    # ValueError: unsized ptr has shape (-1,) which can't broadcast; RuntimeError: shape mismatch
+    try:
+      out_shape = _broadcast_shape(x.shape, y.shape)
+      x, y = x._broadcast_to(out_shape), y._broadcast_to(out_shape)
+    except (RuntimeError, ValueError): pass
     # ptr dtypes aren't in the promo lattice
     if x.dtype == y.dtype or any(isinstance(d, PtrDType) for d in (x.dtype, y.dtype)): return x, y
     return x.cast(out_dtype := least_upper_dtype(x.dtype, y.dtype)), y.cast(out_dtype)

tinygrad/nn/onnx.py

@@ -6,7 +6,7 @@ from tinygrad.tensor import Tensor, _broadcast_shape
 from tinygrad.mixin import ReductionStr
 from tinygrad.helpers import getenv, all_same, prod, flatten, make_tuple, argsort, is_numpy_ndarray, get_single_element, polyN
 from tinygrad.dtype import DType, ConstType, dtypes, _from_np_dtype, truncate, least_upper_dtype, DTYPES_DICT
-from tinygrad.device import is_dtype_supported, Device
+from tinygrad.device import Device
 from tinygrad.uop.ops import sint
 
 # ***** protobuf definitions ******
@@ -36,13 +36,6 @@ class OnnxDataType(enum.IntEnum):
 
   def to_dtype(self) -> DType: return DTYPES_DICT[self.name.lower()]
 
-def dtype_fallback(dtype: DType, fallback_context: str) -> DType:
-  if is_dtype_supported(dtype): return dtype
-  default_dtype = dtypes.default_int if dtypes.is_int(dtype) else dtypes.default_float
-  warnings.warn(f"dtype {dtype} on {Device.DEFAULT} from {fallback_context} is not supported, falling back to {default_dtype}")
-  assert is_dtype_supported(default_dtype), f"dtype {default_dtype} must be supported on {Device.DEFAULT}"
-  return default_dtype
-
 # ***** onnx spec definitions *****
 class Domain(enum.Enum):
   ONNX = "ai.onnx"
@@ -240,21 +233,20 @@ class OnnxPBParser:
       obj["data_location"] = 0
 
     # parse tensor
-    to_dtype = dtype_fallback(true_dtype := OnnxDataType(obj['data_type']).to_dtype(), "buffer parse")
+    dtype = OnnxDataType(obj['data_type']).to_dtype()
     shape = tuple(obj['dims'])
     present_fields = [field for field in ['float_data', 'int32_data', 'int64_data', 'double_data', 'uint64_data', 'raw_data'] if field in obj]
     assert len(present_fields) == 1, f"only 1 data field is allowed from {obj=}"
     data = obj[present_fields[0]]
     if not isinstance(data, Tensor):
-      obj["parsed_tensor"] = Tensor(data, dtype=to_dtype).reshape(shape)
+      obj["parsed_tensor"] = Tensor(data, dtype=dtype).reshape(shape)
       return obj
     assert isinstance(data, Tensor) and data.dtype == dtypes.uint8, data
-    data = data.bitcast(true_dtype).reshape(shape)
-    data = data.to(Device.DEFAULT) if true_dtype is to_dtype else data.to("cpu").cast(to_dtype).to(Device.DEFAULT)
+    data = data.bitcast(dtype).reshape(shape).to(Device.DEFAULT)
     # const folding
     if shape == ():
       if data.dtype == dtypes.float16 and sys.version_info < (3, 12): data = data.cast(dtypes.float32)
-      data = Tensor(data.item(), dtype=to_dtype).reshape(shape)
+      data = Tensor(data.item(), dtype=dtype).reshape(shape)
     obj["parsed_tensor"] = data
     return obj
 
@@ -594,7 +586,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
     raise ValueError(f"pixel_format={pixel_format!r} is not supported.")
 
   def EyeLike(x:Tensor, dtype:int|None=None, k:int=0):
-    ret = Tensor.eye(cast(int, min(x.shape)), dtype=dtype_fallback(OnnxDataType(dtype).to_dtype(), "EyeLike op") if dtype is not None else x.dtype)
+    ret = Tensor.eye(cast(int, min(x.shape)), dtype=OnnxDataType(dtype).to_dtype() if dtype is not None else x.dtype)
     return ret if x.size(0) == x.size(1) else ret.pad(tuple(None if d == ret.size(0) else (k, d-ret.shape[0]-k) for d in x.shape))
 
   def OptionalHasElement(x:Tensor|None=None): return Tensor(x is not None and x.numel() > 0)
@@ -648,7 +640,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
 
   # ***** Casting Ops *****
   # NOTE: saturate only applies to FP8 types
-  def Cast(x:Tensor, to:int, saturate:int=1): return x.cast(dtype_fallback(OnnxDataType(to).to_dtype(), "Cast op"))
+  def Cast(x:Tensor, to:int, saturate:int=1): return x.cast(OnnxDataType(to).to_dtype())
   def CastLike(x:Tensor, target_type:Tensor, saturate:int=1): return x.cast(target_type.dtype)
 
   # ***** Reduce Ops *****
@@ -916,12 +908,13 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
     return x * scale.reshape(1, -1, *[1] * (x.ndim-2)) + bias.reshape(1, -1, *[1] * (x.ndim-2))
   def InstanceNormalization(x:Tensor, scale:Tensor, bias:Tensor, epsilon:float=1e-05):
     return GroupNormalization(x, scale, bias, num_groups=cast(int, x.shape[1]), epsilon=epsilon)
-  def LayerNormalization(x:Tensor, scale:Tensor, bias:Tensor, axis:int=-1, epsilon:float=1e-05, stash_type:int=1):
+  def LayerNormalization(x:Tensor, scale:Tensor, bias:Tensor|None=None, axis:int=-1, epsilon:float=1e-05, stash_type:int=1):
     assert stash_type == 1, "only float32 is supported"
     axes = tuple(i for i in range(axis if axis >= 0 else x.ndim + axis, x.ndim))
     mean = (x32:=x.cast(dtypes.float)).mean(axis=axes, keepdim=True)
     inv_std_dev = (x32.sub(mean)).square().mean(axis=axes, keepdim=True).add(epsilon).rsqrt()
-    return (x32.sub(mean)*inv_std_dev).cast(x.dtype).mul(scale).add(bias), mean, inv_std_dev
+    ret = (x32.sub(mean)*inv_std_dev).cast(x.dtype).mul(scale)
+    return (ret.add(bias) if bias is not None else ret), mean, inv_std_dev
   def SkipLayerNormalization(x:Tensor, skip:Tensor, gamma:Tensor, beta:Tensor|None=None, bias:Tensor|None=None, epsilon:float=1e-12):
     x = x + skip
     if bias is not None: x = x + bias
@@ -986,7 +979,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
     size = int(_resolve_const(size))
     N, n = (size if periodic else size - 1), Tensor.arange(size, requires_grad=False)
     w = a[0] - a[1] * (n * (2 * math.pi / N)).cos() + a[2] * (n * (4 * math.pi / N)).cos()
-    return w.cast(dtype_fallback(OnnxDataType(output_datatype).to_dtype(), "window op"))
+    return w.cast(OnnxDataType(output_datatype).to_dtype())
   def HannWindow(size, output_datatype:int=1, periodic:int=1): return _window(size, output_datatype, periodic, (0.5, 0.5, 0))
   def HammingWindow(size, output_datatype:int=1, periodic:int=1): return _window(size, output_datatype, periodic, (25/46, 21/46, 0))
   def BlackmanWindow(size, output_datatype:int=1, periodic:int=1): return _window(size, output_datatype, periodic, (0.42, 0.5, 0.08))
@@ -1212,7 +1205,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
   # ***** Quantization Ops *****
   def QuantizeLinear(x:Tensor, y_scale:Tensor, y_zero_point:Tensor|int=0, axis:int=1, block_size:int=0, output_dtype:int=0, saturate=1):
     if isinstance(y_zero_point, Tensor): out_dtype = y_zero_point.dtype
-    elif output_dtype != 0: out_dtype = dtype_fallback(OnnxDataType(output_dtype).to_dtype(), "QuantizeLinear op")
+    elif output_dtype != 0: out_dtype = OnnxDataType(output_dtype).to_dtype()
     else: out_dtype = dtypes.uint8
     y_scale, y_zero_point = _prepare_quantize(x, y_scale, y_zero_point, axis, block_size)
     if out_dtype == dtypes.uchar:

tinygrad/nn/optim.py

@@ -10,10 +10,6 @@ class Optimizer:
   """
   def __init__(self, params: list[Tensor], lr: float, device=None, fused=FUSE_OPTIM):
     if lr < 0: raise ValueError(f"Invalid learning rate: {lr}")
-    # if requires_grad is None, but being put into an optimizer, set it to True
-    for x in params:
-      if x.requires_grad is None: x.requires_grad_(True)
-
     self.params: list[Tensor] = dedup([x for x in params if x.requires_grad])
     assert len(self.params) != 0, "optimizer must have at least one param"
     self.buffers: list[Tensor] = dedup([x for x in params if not x.requires_grad])   # buffers are still realized

tinygrad/renderer/cstyle.py

@@ -343,7 +343,7 @@ class MetalRenderer(CStyleLanguage):
   def __init__(self, target:Target):
     super().__init__(target)
     from tinygrad.runtime.ops_metal import MetalCompiler
-    self.compiler, self.tensor_cores = MetalCompiler(), tc.metal if target.arch == "arm64" else []
+    self.compiler, self.tensor_cores = MetalCompiler(), tc.metal if target.arch.startswith("Apple") and int(target.arch[5:]) >= 7 else []
 
   # language options
   kernel_typedef = "kernel void"

tinygrad/runtime/autogen/__init__.py

@@ -155,8 +155,10 @@ def __getattr__(nm):
           *[f"python3 src/compiler/{s}_h.py > gen/{s.split('/')[-1]}.h" for s in ["nir/nir_opcodes", "nir/nir_builder_opcodes"]],
           *[f"python3 src/compiler/nir/nir_{s}_h.py --outdir gen" for s in ["intrinsics", "intrinsics_indices"]]]), cwd=path, shell=True, check=True),
   srcs="https://gitlab.freedesktop.org/mesa/mesa/-/archive/mesa-25.2.7/mesa-25.2.7.tar.gz",
-  dll="([] if DEV.renderer == 'LVP' else ['tinymesa']) + ['tinymesa_cpu']",
-  prolog=["from tinygrad.helpers import DEV", "import gzip, base64"],
+  dll="([] if (_cpu:=DEV.renderer == 'LVP') else ['tinymesa']) + ['tinymesa_cpu'], " \
+      'emsg="not available on this platform" if WIN or (OSX and (platform.machine() != "arm64" or (_mv:=platform.mac_ver()[0][:2]) not in {"14","15","26"})) or (platform.system() == "Linux" and platform.machine() not in {"x86_64", "aarch64"}) else ' \
+      'f"run `sudo curl -fL https://github.com/sirhcm/tinymesa/releases/download/v1/libtinymesa{\'_cpu\'*_cpu}-mesa-25.2.7-{\'macos-\'+_mv if OSX else \'linux\'}-{\'amd64\' if ARCH_X86 else \'arm64\'}.{\'dylib\' if OSX else \'so\'} -o /usr/local/lib/libtinymesa{\'_cpu\'*_cpu}.{\'dylib\' if OSX else \'so\'}`"',
+  prolog=["from tinygrad.helpers import DEV, ARCH_X86, WIN, OSX", "import gzip, base64, platform"],
   epilog=lambda path: [system(f"{root}/extra/mesa/lvp_nir_options.sh {path}")])
     case "libclang":
       return load("libclang",

tinygrad/runtime/autogen/mesa.py

@@ -4,9 +4,9 @@ import ctypes
 from typing import Literal, TypeAlias
 from tinygrad.runtime.support.c import _IO, _IOW, _IOR, _IOWR
 from tinygrad.runtime.support import c
-from tinygrad.helpers import DEV
-import gzip, base64
-dll = c.DLL('mesa', ([] if DEV.renderer == 'LVP' else ['tinymesa']) + ['tinymesa_cpu'])
+from tinygrad.helpers import DEV, ARCH_X86, WIN, OSX
+import gzip, base64, platform
+dll = c.DLL('mesa', ([] if (_cpu:=DEV.renderer == 'LVP') else ['tinymesa']) + ['tinymesa_cpu'], emsg="not available on this platform" if WIN or (OSX and (platform.machine() != "arm64" or (_mv:=platform.mac_ver()[0][:2]) not in {"14","15","26"})) or (platform.system() == "Linux" and platform.machine() not in {"x86_64", "aarch64"}) else f"run `sudo curl -fL https://github.com/sirhcm/tinymesa/releases/download/v1/libtinymesa{'_cpu'*_cpu}-mesa-25.2.7-{'macos-'+_mv if OSX else 'linux'}-{'amd64' if ARCH_X86 else 'arm64'}.{'dylib' if OSX else 'so'} -o /usr/local/lib/libtinymesa{'_cpu'*_cpu}.{'dylib' if OSX else 'so'}`")
 class struct_u_printf_info(c.Struct): pass
 u_printf_info: TypeAlias = struct_u_printf_info
 uint32_t: TypeAlias = ctypes.c_uint32

tinygrad/runtime/graph/metal.py

@@ -23,7 +23,7 @@ class MetalGraph(GraphRunner):
     self.icb = self.dev.sysdevice.newIndirectCommandBufferWithDescriptor_maxCommandCount_options(icb_descriptor, len(self.calls),
                                                                                                  metal.MTLResourceCPUCacheModeDefaultCache)
     if self.icb.value is None: raise GraphException("create indirect command buffer failed, does your system support this?")
-    self.needs_icb_fix = int(self.dev.gpu_family < 9)  # ICB fix not required on M3+ (Apple9+)
+    self.needs_icb_fix = int(not self.dev.arch.startswith("Apple") or int(self.dev.arch[5:]) < 9)  # ICB fix not required on M3+ (Apple9+)
 
     if len(self.vars): self.int_buf = self.dev.allocator.alloc(len(self.vars)*dtypes.int32.itemsize)
 

tinygrad/runtime/ops_cl.py

@@ -112,18 +112,18 @@ class CLDevice(Compiled):
     self.context = checked(cl.clCreateContext(None, 1, self.device_id, CC_CB(), None, status := ctypes.c_int32()), status)
     self.queue = checked(cl.clCreateCommandQueue(self.context, self.device_id, cl.CL_QUEUE_PROFILING_ENABLE, status), status)
     self.pending_copyin: list[memoryview] = []
-    self.device_exts = (cl.clGetDeviceInfo(self.device_id, cl.CL_DEVICE_EXTENSIONS, 4096,
-                                           ctypes.byref(buf := ctypes.create_string_buffer(4096)),
-                                           ctypes.byref(total := ctypes.c_size_t())),
-                                           ctypes.string_at(buf, size=total.value).decode())[1]
+    check(cl.clGetDeviceInfo(self.device_id, cl.CL_DEVICE_EXTENSIONS, 0, None, ctypes.byref(exts_len:=ctypes.c_size_t())))
+    self.device_exts = (cl.clGetDeviceInfo(self.device_id, cl.CL_DEVICE_EXTENSIONS, exts_len.value,
+                                           ctypes.byref(buf := ctypes.create_string_buffer(exts_len.value)), None),
+                                           ctypes.string_at(buf).decode().split())[1]
 
     renderer = IntelRenderer if "cl_intel_subgroup_matrix_multiply_accumulate" in self.device_exts else OpenCLRenderer
     self.cl_compiler = CLCompiler(self, f"{hashlib.md5(self.device_name.encode() + self.driver_version.encode()).hexdigest()}")
 
+    arch = ",".join(self.device_exts)
     if "cl_khr_image2d_from_buffer" in self.device_exts:
       check(cl.clGetDeviceInfo(self.device_id, cl.CL_DEVICE_IMAGE_PITCH_ALIGNMENT, 4, ctypes.byref(ipa := ctypes.c_uint32()), None))
-      arch = f"IMAGE_PITCH_ALIGNMENT={ipa.value}"
-    else: arch = ""
+      arch += f",IMAGE_PITCH_ALIGNMENT={ipa.value}"
     super().__init__(device, CLAllocator(self), [renderer], functools.partial(CLProgram, self), arch=arch)
 
   def count(self) -> int: return len(unwrap(self.device_ids))

tinygrad/runtime/ops_metal.py

@@ -37,12 +37,8 @@ class MetalDevice(Compiled):
     self.timeline_signal = self.sysdevice.newSharedEvent()
     self.timeline_value = 0
 
-    # probe GPU family: Apple9=M3/M4, Apple8=M2, Apple7=M1, etc. values are 1000+N.
-    self.gpu_family = 0
-    for i in range(15, 0, -1):
-      if self.sysdevice.supportsFamily(1000 + i):
-        self.gpu_family = i
-        break
+    # https://developer.apple.com/documentation/metal/mtlgpufamily
+    def check_family(f): return next(filter(self.sysdevice.supportsFamily, reversed([v for v, nm in metal.enum_MTLGPUFamily.items() if f in nm])), 0)
 
     Compiled.profile_events += [ProfileDeviceEvent(device)]
 
@@ -51,7 +47,7 @@ class MetalDevice(Compiled):
     # This can be reproduced locally with any virtualization software (like utm) that can create macOS VMs with apple's own virtualization framework.
     super().__init__(device, MetalAllocator(self), [MetalRenderer],
       functools.partial(MetalProgram, self), MetalGraph if 'virtual' not in from_ns_str(self.sysdevice.name()).lower() else None,
-      arch=platform.machine())
+      arch=metal.enum_MTLGPUFamily[check_family("Apple") or check_family("Mac")][12:])
 
   def synchronize(self):
     for cbuf in self.mtl_buffers_in_flight:

tinygrad/runtime/ops_nv.py

@@ -554,7 +554,6 @@ class PCIIface(PCIIfaceBase):
   def __init__(self, dev, dev_id):
     # PCIIface's MAP_FIXED mmap will overwrite UVM allocations made by NVKIface, so don't try PCIIface if kernel driver was already used.
     if NVKIface.root is not None: raise RuntimeError("Cannot use PCIIface after NVKIface has been initialized (would corrupt UVM memory)")
-    if not OSX: System.reserve_hugepages(64)
     super().__init__(dev, dev_id, vendor=0x10de, devices=((0xff00, (0x2200,0x2400,0x2500,0x2600,0x2700,0x2800,0x2b00,0x2c00,0x2d00,0x2f00)),),
       base_class=0x03, vram_bar=1, va_start=NVMemoryManager.va_allocator.base, va_size=NVMemoryManager.va_allocator.size, dev_impl_t=NVDev)
 

tinygrad/runtime/ops_qcom.py

@@ -137,13 +137,15 @@ class QCOMComputeQueue(HWQueue):
     self.reg(mesa.REG_A6XX_TPL1_DBG_ECO_CNTL, 0)
     self.cmd(mesa.CP_WAIT_FOR_IDLE)
 
+    threadsize = prg.threadsize
+
     self.reg(mesa.REG_A6XX_SP_CS_NDRANGE_0,
              qreg.a6xx_sp_cs_ndrange_0(kerneldim=3, localsizex=local_size[0] - 1, localsizey=local_size[1] - 1, localsizez=local_size[2] - 1),
-             global_size_mp[0], 0, global_size_mp[1], 0, global_size_mp[2], 0, 0xccc0cf, 0xfc | qreg.a6xx_sp_cs_wge_cntl(threadsize=mesa.THREAD64),
+             global_size_mp[0], 0, global_size_mp[1], 0, global_size_mp[2], 0, 0xccc0cf, 0xfc | qreg.a6xx_sp_cs_wge_cntl(threadsize=threadsize),
              cast_int(global_size[0], ceil=True), cast_int(global_size[1], ceil=True), cast_int(global_size[2], ceil=True))
 
     self.reg(mesa.REG_A6XX_SP_CS_CNTL_0,
-             qreg.a6xx_sp_cs_cntl_0(threadsize=mesa.THREAD64, halfregfootprint=prg.hregs, fullregfootprint=prg.fregs, branchstack=prg.brnchstck),
+             qreg.a6xx_sp_cs_cntl_0(threadsize=threadsize, halfregfootprint=prg.hregs, fullregfootprint=prg.fregs, branchstack=prg.brnchstck),
              qreg.a6xx_sp_cs_cntl_1(constantrammode=mesa.CONSTLEN_256, shared_size=prg.shared_size), # should this be CONSTLEN_512?
              0, prg.prg_offset, *data64_le(prg.lib_gpu.va_addr),
              qreg.a6xx_sp_cs_pvt_mem_param(memsizeperitem=prg.pvtmem_size_per_item), *data64_le(prg.dev._stack.va_addr),
@@ -187,7 +189,7 @@ class QCOMComputeQueue(HWQueue):
     if prg.NIR:
       self.reg(mesa.REG_A6XX_SP_CS_CONST_CONFIG_0,
                qreg.a6xx_sp_cs_const_config_0(wgidconstid=prg.wgid, wgsizeconstid=prg.wgsz, wgoffsetconstid=0xfc, localidregid=prg.lid),
-               qreg.a6xx_sp_cs_wge_cntl(linearlocalidregid=0xfc, threadsize=mesa.THREAD64))
+               qreg.a6xx_sp_cs_wge_cntl(linearlocalidregid=0xfc, threadsize=threadsize))
       self.cmd(mesa.CP_EXEC_CS, 0,
                qreg.cp_exec_cs_1(ngroups_x=global_size[0]), qreg.cp_exec_cs_2(ngroups_y=global_size[1]), qreg.cp_exec_cs_3(_ngroups_z=global_size[2]))
     else: self.cmd(mesa.CP_RUN_OPENCL, 0)
@@ -251,6 +253,7 @@ class QCOMProgram(HCQProgram):
 
       self.tex_off, self.ibo_off, self.samp_off = 2048, 2048 + 0x40 * self.tex_cnt, 2048 + 0x40 * (self.tex_cnt + self.ibo_cnt)
       self.fregs, self.hregs = v.info.max_reg + 1, v.info.max_half_reg + 1
+      self.threadsize = mesa.THREAD128 if v.info.double_threadsize else mesa.THREAD64
     else: self._parse_lib(lib)
 
     self.lib_gpu: HCQBuffer = self.dev.allocator.alloc(self.image_size, buf_spec:=BufferSpec(cpu_access=True, nolru=True))
@@ -320,6 +323,10 @@ class QCOMProgram(HCQProgram):
     reg_desc_off = _read_lib(lib, 0x34)
     self.fregs, self.hregs = _read_lib(lib, reg_desc_off + 0x14), _read_lib(lib, reg_desc_off + 0x18)
 
+    # The Qualcomm OpenCL stack dispatches these binaries with 128-thread waves.
+    # THREAD64 leaves half-rate ALU throughput for the same shader image.
+    self.threadsize = mesa.THREAD128 if getenv("THREAD128") else mesa.THREAD64
+
 class QCOMAllocator(HCQAllocatorBase):
   def _alloc(self, size:int, opts:BufferSpec) -> HCQBuffer:
     return self.dev._gpu_map(opts.external_ptr, size) if opts.external_ptr else self.dev._gpu_alloc(size)

tinygrad/runtime/support/nv/ip.py

@@ -163,9 +163,9 @@ class NV_FLCN(NV_IP):
       patched_image[(cmd_off:=self.desc_v3.IMEMLoadSize+dmem.cmd_in_buffer_offset) : cmd_off+len(cmd)] = cmd
       patched_image[(sig_off:=self.desc_v3.IMEMLoadSize+self.desc_v3.PKCDataOffset) : sig_off+0x180] = signature[-0x180:]
 
-      return self.nvdev._alloc_sysmem(len(patched_image), contiguous=True, data=patched_image)
+      return self.nvdev._alloc_boot_mem(len(patched_image), data=patched_image, sysmem=False)
 
-    _, self.frts_image_sysmem = __patch(0x15, bytes(frts_cmd))
+    _, self.frts_image_paddr, _ = __patch(0x15, bytes(frts_cmd))
 
   def prep_booter(self):
     sha = {"ga102":"4497e3eff7e95c774b8a569d17b27c08c9650158d10b229d2be81cdcad9a085b",
@@ -179,14 +179,14 @@ class NV_FLCN(NV_IP):
 
     (patched_image:=bytearray(b[h.data_offset:h.data_offset + h.data_size]))[patch_loc:patch_loc+sig_len] = sig
 
-    _, self.booter_image_sysmem = self.nvdev._alloc_sysmem(len(patched_image), contiguous=True, data=patched_image)
+    _, self.booter_image_paddr, _ = self.nvdev._alloc_boot_mem(len(patched_image), data=patched_image, sysmem=False)
     self.booter_data_off, self.booter_data_sz, self.booter_code_off, self.booter_code_sz = lh.os_data_offset, lh.os_data_size, app.offset, app.size
 
   def init_hw(self):
     self.falcon, self.sec2 = 0x00110000, 0x00840000
 
     self.reset(self.falcon)
-    self.execute_hs(self.falcon, self.frts_image_sysmem[0], code_off=0x0, data_off=self.desc_v3.IMEMLoadSize,
+    self.execute_hs(self.falcon, self.frts_image_paddr, code_off=0x0, data_off=self.desc_v3.IMEMLoadSize,
       imemPa=self.desc_v3.IMEMPhysBase, imemVa=self.desc_v3.IMEMVirtBase, imemSz=self.desc_v3.IMEMLoadSize,
       dmemPa=self.desc_v3.DMEMPhysBase, dmemVa=0x0, dmemSz=self.desc_v3.DMEMLoadSize,
       pkc_off=self.desc_v3.PKCDataOffset, engid=self.desc_v3.EngineIdMask, ucodeid=self.desc_v3.UcodeId)
@@ -195,12 +195,12 @@ class NV_FLCN(NV_IP):
     self.reset(self.falcon, riscv=True)
 
     # set up the mailbox
-    self.nvdev.NV_PGSP_FALCON_MAILBOX0.write(lo32(self.nvdev.gsp.libos_args_sysmem[0]))
-    self.nvdev.NV_PGSP_FALCON_MAILBOX1.write(hi32(self.nvdev.gsp.libos_args_sysmem[0]))
+    self.nvdev.NV_PGSP_FALCON_MAILBOX0.write(lo32(self.nvdev.gsp.libos_args_sysmem))
+    self.nvdev.NV_PGSP_FALCON_MAILBOX1.write(hi32(self.nvdev.gsp.libos_args_sysmem))
 
     # booter
     self.reset(self.sec2)
-    mbx = self.execute_hs(self.sec2, self.booter_image_sysmem[0], code_off=self.booter_code_off, data_off=self.booter_data_off,
+    mbx = self.execute_hs(self.sec2, self.booter_image_paddr, code_off=self.booter_code_off, data_off=self.booter_data_off,
       imemPa=0x0, imemVa=self.booter_code_off, imemSz=self.booter_code_sz, dmemPa=0x0, dmemVa=0x0, dmemSz=self.booter_data_sz,
       pkc_off=0x10, engid=1, ucodeid=3, mailbox=self.nvdev.gsp.wpr_meta_sysmem)
     assert mbx[0] == 0x0, f"Booter failed to execute, mailbox is {mbx[0]:08x}, {mbx[1]:08x}"
@@ -208,11 +208,11 @@ class NV_FLCN(NV_IP):
     self.nvdev.NV_PFALCON_FALCON_OS.with_base(self.falcon).write(0x0)
     assert self.nvdev.NV_PRISCV_RISCV_CPUCTL.with_base(self.falcon).read_bitfields()['active_stat'] == 1, "GSP Core is not active"
 
-  def execute_dma(self, base:int, cmd:int, dest:int, mem_off:int, sysmem:int, size:int):
+  def execute_dma(self, base:int, cmd:int, dest:int, mem_off:int, src:int, size:int):
     wait_cond(lambda: self.nvdev.NV_PFALCON_FALCON_DMATRFCMD.with_base(base).read_bitfields()['full'], value=0, msg="DMA does not progress")
 
-    self.nvdev.NV_PFALCON_FALCON_DMATRFBASE.with_base(base).write(lo32(sysmem >> 8))
-    self.nvdev.NV_PFALCON_FALCON_DMATRFBASE1.with_base(base).write(hi32(sysmem >> 8) & 0x1ff)
+    self.nvdev.NV_PFALCON_FALCON_DMATRFBASE.with_base(base).write(lo32(src >> 8))
+    self.nvdev.NV_PFALCON_FALCON_DMATRFBASE1.with_base(base).write(hi32(src >> 8) & 0x1ff)
 
     xfered = 0
     while xfered < size:
@@ -232,19 +232,19 @@ class NV_FLCN(NV_IP):
 
   def wait_cpu_halted(self, base): wait_cond(lambda: self.nvdev.NV_PFALCON_FALCON_CPUCTL.with_base(base).read_bitfields()['halted'], msg="not halted")
 
-  def execute_hs(self, base, img_sysmem, code_off, data_off, imemPa, imemVa, imemSz, dmemPa, dmemVa, dmemSz, pkc_off, engid, ucodeid, mailbox=None):
+  def execute_hs(self, base, img_paddr, code_off, data_off, imemPa, imemVa, imemSz, dmemPa, dmemVa, dmemSz, pkc_off, engid, ucodeid, mailbox=None):
     self.disable_ctx_req(base)
 
-    self.nvdev.NV_PFALCON_FBIF_TRANSCFG.with_base(base)[ctx_dma:=0].update(target=self.nvdev.NV_PFALCON_FBIF_TRANSCFG_TARGET_COHERENT_SYSMEM,
-      mem_type=self.nvdev.NV_PFALCON_FBIF_TRANSCFG_MEM_TYPE_PHYSICAL)
+    # target=0 is FB (not in published headers)
+    self.nvdev.NV_PFALCON_FBIF_TRANSCFG.with_base(base)[ctx_dma:=0].update(target=0, mem_type=self.nvdev.NV_PFALCON_FBIF_TRANSCFG_MEM_TYPE_PHYSICAL)
 
     cmd = self.nvdev.NV_PFALCON_FALCON_DMATRFCMD.with_base(base).encode(write=0, size=self.nvdev.NV_PFALCON_FALCON_DMATRFCMD_SIZE_256B,
       ctxdma=ctx_dma, imem=1, sec=1)
-    self.execute_dma(base, cmd, dest=imemPa, mem_off=imemVa, sysmem=img_sysmem+code_off-imemVa, size=imemSz)
+    self.execute_dma(base, cmd, dest=imemPa, mem_off=imemVa, src=img_paddr+code_off-imemVa, size=imemSz)
 
     cmd = self.nvdev.NV_PFALCON_FALCON_DMATRFCMD.with_base(base).encode(write=0, size=self.nvdev.NV_PFALCON_FALCON_DMATRFCMD_SIZE_256B,
       ctxdma=ctx_dma, imem=0, sec=0)
-    self.execute_dma(base, cmd, dest=dmemPa, mem_off=dmemVa, sysmem=img_sysmem+data_off-dmemVa, size=dmemSz)
+    self.execute_dma(base, cmd, dest=dmemPa, mem_off=dmemVa, src=img_paddr+data_off-dmemVa, size=dmemSz)
 
     self.nvdev.NV_PFALCON2_FALCON_BROM_PARAADDR.with_base(base)[0].write(pkc_off)
     self.nvdev.NV_PFALCON2_FALCON_BROM_ENGIDMASK.with_base(base).write(engid)
@@ -293,7 +293,9 @@ class NV_FLCN_COT(NV_IP):
     self.nvdev.include("dev_fsp_pri", "gh100")
     self.nvdev.include("dev_bus", "tu102")
 
-    self.fmc_boot_args_view, self.fmc_boot_args_sysmem = self.nvdev._alloc_boot_struct(nv.GSP_FMC_BOOT_PARAMS())
+    self.fmc_boot_args_view, _, fmc_boot_addrs = self.nvdev._alloc_boot_mem(ctypes.sizeof(nv.GSP_FMC_BOOT_PARAMS),
+      data=bytes(nv.GSP_FMC_BOOT_PARAMS()))
+    self.fmc_boot_args_sysmem = fmc_boot_addrs[0]
     self.init_fmc_image()
 
   def init_fmc_image(self):
@@ -302,18 +304,19 @@ class NV_FLCN_COT(NV_IP):
     def _section(s): return next((sh.content for sh in sections if sh.name == s))
     self.fmc_booter_image, self.fmc_booter_hash = _section("image"), memoryview(_section("hash")).cast('I')
     self.fmc_booter_sig, self.fmc_booter_pkey = memoryview(_section("signature")).cast('I'), memoryview(_section("publickey") + b"\x00" * 3).cast('I')
-    _, self.fmc_booter_sysmem = self.nvdev._alloc_sysmem(len(self.fmc_booter_image), contiguous=True, data=self.fmc_booter_image)
+    _, _, fmc_booter_addrs = self.nvdev._alloc_boot_mem(len(self.fmc_booter_image), data=self.fmc_booter_image)
+    self.fmc_booter_bar1 = fmc_booter_addrs[0]
 
   def init_hw(self):
     self.falcon = 0x00110000
 
     boot_args = nv.GSP_ACR_BOOT_GSP_RM_PARAMS(gspRmDescOffset=self.nvdev.gsp.wpr_meta_sysmem,
       gspRmDescSize=ctypes.sizeof(nv.GspFwWprMeta), target=nv.GSP_DMA_TARGET_COHERENT_SYSTEM, bIsGspRmBoot=True)
-    rm_args = nv.GSP_RM_PARAMS(bootArgsOffset=self.nvdev.gsp.libos_args_sysmem[0], target=nv.GSP_DMA_TARGET_COHERENT_SYSTEM)
+    rm_args = nv.GSP_RM_PARAMS(bootArgsOffset=self.nvdev.gsp.libos_args_sysmem, target=nv.GSP_DMA_TARGET_COHERENT_SYSTEM)
     self.fmc_boot_args_view[:ctypes.sizeof(nv.GSP_FMC_BOOT_PARAMS)] = bytes(nv.GSP_FMC_BOOT_PARAMS(bootGspRmParams=boot_args, gspRmParams=rm_args))
 
     cot_payload = nv.NVDM_PAYLOAD_COT(version=0x2, size=ctypes.sizeof(nv.NVDM_PAYLOAD_COT), frtsVidmemOffset=0x1c00000, frtsVidmemSize=0x100000,
-      gspBootArgsSysmemOffset=self.fmc_boot_args_sysmem, gspFmcSysmemOffset=self.fmc_booter_sysmem[0])
+      gspBootArgsSysmemOffset=self.fmc_boot_args_sysmem, gspFmcSysmemOffset=self.fmc_booter_bar1)
     for i,x in enumerate(self.fmc_booter_hash): cot_payload.hash384[i] = x
     for i,x in enumerate(self.fmc_booter_sig): cot_payload.signature[i] = x
     for i,x in enumerate(self.fmc_booter_pkey): cot_payload.publicKey[i] = x
@@ -360,7 +363,7 @@ 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_view, queues_sysmem = self.nvdev._alloc_sysmem(pt_size + queue_size * 2, contiguous=False)
+    queues_view, _, queues_sysmem = self.nvdev._alloc_boot_mem(pt_size + queue_size * 2, sysmem=True)
 
     # Fill up ptes
     for i, sysmem in enumerate(queues_sysmem): queues_view.view(i * 0x8, 0x8, fmt='Q')[0] = sysmem
@@ -368,7 +371,9 @@ class NV_GSP(NV_IP):
     # Fill up arguments
     queue_args = nv.MESSAGE_QUEUE_INIT_ARGUMENTS(sharedMemPhysAddr=queues_sysmem[0], pageTableEntryCount=pte_cnt, cmdQueueOffset=pt_size,
       statQueueOffset=pt_size + queue_size)
-    _, self.rm_args_sysmem = self.nvdev._alloc_boot_struct(nv.GSP_ARGUMENTS_CACHED(bDmemStack=True, messageQueueInitArguments=queue_args))
+    _, _, rm_args_addrs = self.nvdev._alloc_boot_mem(ctypes.sizeof(nv.GSP_ARGUMENTS_CACHED),
+      data=bytes(nv.GSP_ARGUMENTS_CACHED(bDmemStack=True, messageQueueInitArguments=queue_args)))
+    self.rm_args_sysmem = rm_args_addrs[0]
 
     # Build command queue header
     # self.cmd_q_va, self.stat_q_va = queues_view.addr + pt_size, queues_view.addr + pt_size + queue_size
@@ -380,11 +385,12 @@ class NV_GSP(NV_IP):
     self.cmd_q = NVRpcQueue(self, self.cmd_q_view, None)
 
   def init_libos_args(self):
-    _, logbuf_sysmem = self.nvdev._alloc_sysmem((2 << 20), contiguous=True)
-    libos_args_view, self.libos_args_sysmem = self.nvdev._alloc_sysmem(0x1000, contiguous=True)
+    _, _, logbuf_addrs = self.nvdev._alloc_boot_mem(2 << 20)
+    libos_args_view, _, libos_addrs = self.nvdev._alloc_boot_mem(0x1000)
+    self.libos_args_sysmem = libos_addrs[0]
 
     libos_structs = [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)
+        id8=int.from_bytes(bytes(f"LOG{name}", 'utf-8'), 'big'), pa=logbuf_addrs[0] + 0x10000 * i)
         for i, name in enumerate(["INIT", "INTR", "RM", "MNOC", "KRNL"])]
     libos_structs.append(nv.LibosMemoryRegionInitArgument(kind=nv.LIBOS_MEMORY_REGION_CONTIGUOUS, loc=nv.LIBOS_MEMORY_REGION_LOC_SYSMEM, size=0x1000,
         id8=int.from_bytes(bytes("RMARGS", 'utf-8'), 'big'), pa=self.rm_args_sysmem))
@@ -400,7 +406,7 @@ 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_view, self.gsp_radix3_sysmem = self.nvdev._alloc_sysmem(offsets[-1] + len(self.gsp_image), contiguous=False)
+    radix_view, _, self.gsp_radix3_addrs = self.nvdev._alloc_boot_mem(offsets[-1] + len(self.gsp_image))
 
     # Copy image
     radix_view.view(offsets[-1], len(self.gsp_image))[:] = self.gsp_image
@@ -408,10 +414,11 @@ class NV_GSP(NV_IP):
     # Copy level and image pages.
     for i in range(0, 3):
       cur_offset = sum(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]])
+      radix_view.view(offsets[i], npages[i+1] * 8, fmt='Q')[:] = array.array('Q', self.gsp_radix3_addrs[cur_offset:cur_offset+npages[i+1]])
 
     # Copy signature
-    _, self.gsp_signature_sysmem = self.nvdev._alloc_sysmem(len(signature), contiguous=True, data=signature)
+    _, _, gsp_sig_addrs = self.nvdev._alloc_boot_mem(len(signature), data=signature)
+    self.gsp_signature_bar1 = gsp_sig_addrs[0]
 
   def init_boot_binary_image(self):
     sha = {"ga102":"82428f532240727e95bb3083fbaaba9b2cc7b937314323f2d546ce7245f27fad",
@@ -419,15 +426,16 @@ class NV_GSP(NV_IP):
            "gb202":"d40b48e431d1707dc77af3605db358ed7a32ebfc2830eb74de2eddb4d3025071"}[self.nvdev.fw_name]
     h = nv.struct_nvfw_bin_hdr.from_buffer_copy(b:=fetch_fw(f"nvidia/{self.nvdev.fw_name}/gsp", "bootloader-570.144.bin", sha))
     self.booter_image, self.booter_desc = b[h.data_offset:h.data_offset+h.data_size], nv.RM_RISCV_UCODE_DESC.from_buffer_copy(b, h.header_offset)
-    _, self.booter_sysmem = self.nvdev._alloc_sysmem(len(self.booter_image), contiguous=True, data=self.booter_image)
+    _, _, booter_addrs = self.nvdev._alloc_boot_mem(len(self.booter_image), data=self.booter_image)
+    self.booter_bar1 = booter_addrs[0]
 
   def init_wpr_meta(self):
     self.init_gsp_image()
     self.init_boot_binary_image()
 
-    common = {'sizeOfBootloader':(boot_sz:=len(self.booter_image)), 'sysmemAddrOfBootloader':self.booter_sysmem[0],
-      'sizeOfRadix3Elf':(radix3_sz:=len(self.gsp_image)), 'sysmemAddrOfRadix3Elf': self.gsp_radix3_sysmem[0],
-      'sizeOfSignature': 0x1000, 'sysmemAddrOfSignature': self.gsp_signature_sysmem[0],
+    common = {'sizeOfBootloader':(boot_sz:=len(self.booter_image)), 'sysmemAddrOfBootloader':self.booter_bar1,
+      'sizeOfRadix3Elf':(radix3_sz:=len(self.gsp_image)), 'sysmemAddrOfRadix3Elf': self.gsp_radix3_addrs[0],
+      'sizeOfSignature': 0x1000, 'sysmemAddrOfSignature': self.gsp_signature_bar1,
       'bootloaderCodeOffset': self.booter_desc.monitorCodeOffset, 'bootloaderDataOffset': self.booter_desc.monitorDataOffset,
       'bootloaderManifestOffset': self.booter_desc.manifestOffset, 'revision':nv.GSP_FW_WPR_META_REVISION, 'magic':nv.GSP_FW_WPR_META_MAGIC}
 
@@ -441,7 +449,8 @@ class NV_GSP(NV_IP):
         gspFwHeapOffset=(gsp_heap_off:=round_down(gsp_off-gsp_heap_sz, 0x100000)), gspFwWprStart=(wpr_st:=round_down(gsp_heap_off-0x1000, 0x100000)),
         nonWprHeapSize=(non_wpr_sz:=0x100000), nonWprHeapOffset=(non_wpr_off:=round_down(wpr_st-non_wpr_sz, 0x100000)), gspFwRsvdStart=non_wpr_off)
       assert self.nvdev.flcn.frts_offset == m.frtsOffset, f"FRTS mismatch: {self.nvdev.flcn.frts_offset} != {m.frtsOffset}"
-    self.wpr_meta, self.wpr_meta_sysmem = self.nvdev._alloc_boot_struct(m)
+    self.wpr_meta, _, wpr_meta_addrs = self.nvdev._alloc_boot_mem(ctypes.sizeof(type(m)), data=bytes(m))
+    self.wpr_meta_sysmem = wpr_meta_addrs[0]
 
   def promote_ctx(self, client:int, subdevice:int, obj:int, ctxbufs:dict[int, GRBufDesc], bufs=None, virt=None, phys=None):
     res, prom = {}, nv_gpu.NV2080_CTRL_GPU_PROMOTE_CTX_PARAMS(entryCount=len(ctxbufs), engineType=0x1, hChanClient=client, hObject=obj)
@@ -527,8 +536,8 @@ class NV_GSP(NV_IP):
       params.ramfcMem = nv_gpu.NV_MEMORY_DESC_PARAMS(base=ramfc_alloc.paddrs[0][0], size=0x200, addressSpace=2, cacheAttrib=0)
       params.instanceMem = nv_gpu.NV_MEMORY_DESC_PARAMS(base=ramfc_alloc.paddrs[0][0], size=0x1000, addressSpace=2, cacheAttrib=0)
 
-      _, method_sysmem = self.nvdev._alloc_sysmem(0x5000, contiguous=True)
-      params.mthdbufMem = nv_gpu.NV_MEMORY_DESC_PARAMS(base=method_sysmem[0], size=0x5000, addressSpace=1, cacheAttrib=0)
+      _, method_paddr, _ = self.nvdev._alloc_boot_mem(0x5000, sysmem=False)
+      params.mthdbufMem = nv_gpu.NV_MEMORY_DESC_PARAMS(base=method_paddr, size=0x5000, addressSpace=2, cacheAttrib=0)
 
       if client is not None and client != self.priv_root and params.hObjectError != 0:
         params.errorNotifierMem = nv_gpu.NV_MEMORY_DESC_PARAMS(base=0, size=0xecc, addressSpace=0, cacheAttrib=0)
@@ -629,8 +638,8 @@ class NV_GSP(NV_IP):
       elif op == 0x8: # core resume
         self.nvdev.flcn.reset(self.nvdev.flcn.falcon, riscv=True)
 
-        self.nvdev.NV_PGSP_FALCON_MAILBOX0.write(lo32(self.libos_args_sysmem[0]))
-        self.nvdev.NV_PGSP_FALCON_MAILBOX1.write(hi32(self.libos_args_sysmem[0]))
+        self.nvdev.NV_PGSP_FALCON_MAILBOX0.write(lo32(self.libos_args_sysmem))
+        self.nvdev.NV_PGSP_FALCON_MAILBOX1.write(hi32(self.libos_args_sysmem))
 
         self.nvdev.flcn.start_cpu(self.nvdev.flcn.sec2)
         wait_cond(lambda: self.nvdev.NV_PGC6_BSI_SECURE_SCRATCH_14.read_bitfields()['boot_stage_3_handoff'], msg="SEC2 didn't hand off")

tinygrad/runtime/support/nv/nvdev.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
-import ctypes, time, functools, tinygrad.runtime.autogen.nv_regs
-from tinygrad.helpers import getenv, DEBUG, getbits
+import time, functools, tinygrad.runtime.autogen.nv_regs
+from tinygrad.helpers import getenv, DEBUG, getbits, round_up
 from tinygrad.runtime.autogen import pci
 from tinygrad.runtime.support.memory import TLSFAllocator, MemoryManager, AddrSpace
 from tinygrad.runtime.support.nv.ip import NV_FLCN, NV_FLCN_COT, NV_GSP
@@ -145,15 +145,14 @@ class NVDev:
     self.mm = NVMemoryManager(self, self.vram_size - (64 << 20), boot_size=(2 << 20), pt_t=NVPageTableEntry, va_bits=bits, va_shifts=shifts,
       va_base=0, palloc_ranges=[(x, x) for x in [512 << 20, 2 << 20, 4 << 10]], reserve_ptable=not self.large_bar)
 
-  def _alloc_sysmem(self, size:int, vaddr:int=0, contiguous:bool=False, data:bytes|None=None) -> tuple[MMIOInterface, list[int]]:
-    view, paddrs = self.pci_dev.alloc_sysmem(size, vaddr, contiguous=contiguous)
+  def _alloc_boot_mem(self, size:int, data:bytes|None=None, contiguous:bool=False, sysmem:bool|None=None) -> tuple[MMIOInterface, int, list[int]]:
+    sz = round_up(size, 0x1000)
+    if sysmem is True or (sysmem is None and not self.large_bar): view, paddrs = self.pci_dev.alloc_sysmem(size, 0, contiguous=contiguous)
+    else:
+      paddr = self.mm.palloc(sz, boot=False)
+      view, paddrs = self.vram.view(paddr, sz), [self.pci_dev.bar_info(1)[0] + paddr + i * 0x1000 for i in range(sz // 0x1000)]
     if data is not None: view[:size] = data
-    return view, paddrs
-
-  def _alloc_boot_struct(self, struct:ctypes.Structure) -> tuple[MMIOInterface, int]:
-    view, paddrs = self._alloc_sysmem(sz:=ctypes.sizeof(type(struct)), contiguous=True)
-    view[:sz] = bytes(struct)
-    return view, paddrs[0]
+    return view, paddrs[0], paddrs
 
   def include(self, name:str, arch:str):
     for k,v in getattr(getattr(tinygrad.runtime.autogen.nv_regs, name), arch or 'regs').items():

tinygrad/runtime/support/system.py

@@ -38,8 +38,6 @@ class _System:
       return vfio_fd
     except OSError: return None
 
-  def reserve_hugepages(self, cnt): os.system(f"sudo sh -c 'echo {cnt} > /proc/sys/vm/nr_hugepages'")
-
   @functools.cache
   def reserve_va(self, va_start, va_size):
     # cached, runs only once per range. used to not collide with other mappings.

tinygrad/schedule/indexing.py

@@ -144,18 +144,11 @@ def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UO
     case _: raise RuntimeError(f"{op} is not a MovementOp")
   return rngs
 
-pm_do_broadcast = PatternMatcher([
-  (UPat(GroupOp.Broadcastable, name="x"), lambda x: x.replace(src=tuple(y._broadcast_to(x.shape) for y in x.src))),
-])
-
 @profile_matches
 def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
   if debug: print("**************************")
   rctx = IndexingContext()
 
-  # run broadcasting
-  tsink = graph_rewrite(tsink, pm_do_broadcast, name="do broadcast")
-
   # get ops to realize
   graph_rewrite(tsink, pm_generate_realize_map, ctx=rctx.realize_map, name="get realize")
 

tinygrad/tensor.py

@@ -92,7 +92,7 @@ class Tensor(OpMixin):
   training: ClassVar[bool] = False
 
   def __init__(self, data:ConstType|bytes|list|tuple|UOp|'numpy.ndarray'|pathlib.Path|None,
-               device:str|tuple|list|None=None, dtype:DTypeLike|None=None, requires_grad:bool|None=None, _force_unique:bool=False):
+               device:str|tuple|list|None=None, dtype:DTypeLike|None=None, requires_grad:bool=True, _force_unique:bool=False):
     if device is None:
       if isinstance(data, pathlib.Path): device = f"DISK:{data.resolve()}"  # keep it on the disk if device is None
       elif isinstance(data, UOp): device = data._device
@@ -103,9 +103,7 @@ class Tensor(OpMixin):
     # tensors can have gradients if you have called .backward
     self.grad:Tensor|None = None
 
-    # NOTE: this can be in three states. False and None: no gradient, True: gradient
-    # None (the default) will be updated to True if it's put in an optimizer
-    self.requires_grad:bool|None = requires_grad
+    self.requires_grad:bool = requires_grad
 
     # create a UOp from the different types of inputs
     if isinstance(data, UOp):
@@ -115,8 +113,8 @@ class Tensor(OpMixin):
     elif data is None:
       data = UOp.const(_dtype or dtypes.default_float, 0, _device)
     elif isinstance(data, get_args(ConstType)):
-      if _force_unique or requires_grad: data = UOp.unique_const(data, _dtype, _device)
-      else: data = UOp.const(_dtype or dtypes.from_py(data), data, _device)
+      dt = _dtype or dtypes.from_py(data)
+      data = UOp.unique_const(data, dt, _device) if _force_unique or (requires_grad and dtypes.is_float(dt)) else UOp.const(dt, data, _device)
     elif isinstance(data, bytes): data = _frompy(data, _dtype or dtypes.uint8, _device)
     elif isinstance(data, (list, tuple)):
       if _dtype is None:
@@ -151,11 +149,10 @@ class Tensor(OpMixin):
     srcs = (self,)+x
     new_uop: UOp = fxn(*[t.uop for t in srcs], *extra_args, **kwargs)
     if TRACEMETA >= 1 and (metadata:=_METADATA.get()) is not None: all_metadata[new_uop] = (metadata,)
-    needs_input_grad = [t.requires_grad for t in srcs]
     # directly create the Tensor
     ret = Tensor.__new__(Tensor)
     ret.uop, ret.grad = new_uop, None
-    ret.requires_grad = True if any(needs_input_grad) else None if None in needs_input_grad else False
+    ret.requires_grad = any(t.requires_grad for t in srcs)
     # add to all_tensors after construction succeeds
     all_tensors[weakref.ref(ret)] = None
     return ret
@@ -166,7 +163,7 @@ class Tensor(OpMixin):
   @staticmethod
   def unique_const(fill_value:ConstType|UOp, **kwargs) -> Tensor: return Tensor(fill_value, _force_unique=True, **kwargs)
 
-  def requires_grad_(self, requires_grad=True) -> Tensor:
+  def requires_grad_(self, requires_grad:bool=True) -> Tensor:
     # make the UOp unique if it's a CONST to prevent gradient accumulation bugs with cached const UOps
     if requires_grad and self.uop.op is Ops.CONST: self.replace(Tensor(self.uop.arg, device=self.device, dtype=self.dtype, requires_grad=True))
     self.requires_grad = requires_grad
@@ -566,7 +563,7 @@ class Tensor(OpMixin):
     return Tensor._device_seeds[device], low.cat(high)
 
   @staticmethod
-  def rand(*shape, device:str|None=None, dtype:DTypeLike|None=None, requires_grad:bool|None=None, contiguous:bool=True) -> Tensor:
+  def rand(*shape, device:str|None=None, dtype:DTypeLike|None=None, requires_grad:bool=True, contiguous:bool=True) -> Tensor:
     """
     Creates a tensor with the given shape, filled with random values from a uniform distribution over the interval `[0, 1)`.
 
@@ -594,32 +591,14 @@ class Tensor(OpMixin):
 
   # ***** creation helper functions *****
 
-  @classmethod
-  def eye(cls, n:int, m:int|None=None, dtype=None, device=None, requires_grad:bool|None=None) -> Tensor:
-    """
-    Returns a 2-D tensor with `n` rows and `m` columns, with ones on the diagonal and zeros elsewhere.
-
-    You can pass in `dtype` and `device` keyword arguments to control the data type and device of the tensor.
-    Additionally, all other keyword arguments are passed to the constructor of the tensor.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(Tensor.eye(3).numpy())
-    ```
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(Tensor.eye(2, 4).numpy())
-    ```
-    """
-    return super().eye(n, m, dtype, device).requires_grad_(requires_grad)
-
   def _multi_like(self, fxn, *args, **kwargs) -> Tensor:
     dtype = kwargs.pop("dtype", self.dtype)
     if kwargs.get("device") is not None: raise RuntimeError("cannot specify `device` on `*_like` of a multi device tensor")
     if self.uop.axis is None: return fxn(self.shape, *args, dtype=dtype, **kwargs).shard(self.device)
     stacked = UOp.mstack(*[fxn(self.uop.shard_shape, *args, device=d, dtype=dtype, **kwargs).uop for d in self.device])
-    return Tensor(stacked.multi(self.uop.axis), requires_grad=kwargs.get("requires_grad"))
+    return Tensor(stacked.multi(self.uop.axis), requires_grad=kwargs.get("requires_grad", True))
 
-  def full_like(self, fill_value:ConstType, dtype=None, device=None, requires_grad=None) -> Tensor:
+  def full_like(self, fill_value:ConstType, dtype=None, device=None, requires_grad:bool=False) -> Tensor:
     """
     Creates a tensor with the same shape as `self`, filled with the given value.
     If `dtype` is not specified, the dtype of `self` is used.
@@ -631,12 +610,9 @@ class Tensor(OpMixin):
     print(Tensor.full_like(t, 42).numpy())
     ```
     """
-    if device is not None:
-      if isinstance(self.device, tuple): raise RuntimeError("cannot specify `device` on `full_like` of a multi device tensor")
-      return Tensor.full(self.shape, fill_value, dtype=dtype or self.dtype, device=device).requires_grad_(requires_grad)
-    if requires_grad:
-      return Tensor.full(self.shape, fill_value, dtype=dtype or self.dtype, device=self.device).requires_grad_(requires_grad)
-    return super().full_like(fill_value, dtype)
+    if device is None: return super().full_like(fill_value, dtype).requires_grad_(requires_grad)
+    if isinstance(self.device, tuple): raise RuntimeError("cannot specify `device` on `full_like` of a multi device tensor")
+    return Tensor.full(self.shape, fill_value, dtype=dtype or self.dtype, device=device).requires_grad_(requires_grad)
 
   def rand_like(self, **kwargs) -> Tensor:
     """
@@ -655,7 +631,7 @@ class Tensor(OpMixin):
 
   # ***** random functions *****
 
-  def randn_like(self, dtype:DTypeLike|None=None, requires_grad:bool|None=None, **kwargs) -> Tensor:
+  def randn_like(self, dtype:DTypeLike|None=None, requires_grad:bool=True, **kwargs) -> Tensor:
     """
     Creates a tensor with the same shape and sharding as `self`, filled with random values from a normal distribution with mean 0 and variance 1.
 
@@ -672,7 +648,7 @@ class Tensor(OpMixin):
     return (src[0].mul(2*math.pi).cos().mul((1 - src[1]).log().mul(-2).sqrt()).cast(dtype or self.dtype)).requires_grad_(requires_grad)
 
   @staticmethod
-  def randn(*shape, dtype:DTypeLike|None=None, requires_grad:bool|None=None, **kwargs) -> Tensor:
+  def randn(*shape, dtype:DTypeLike|None=None, requires_grad:bool=True, **kwargs) -> Tensor:
     """
     Creates a tensor with the given shape, filled with random values from a normal distribution with mean `0` and standard deviation `1`.
     If `dtype` is not specified, the default type is used.
@@ -707,7 +683,7 @@ class Tensor(OpMixin):
     return Tensor.uniform(*shape, low=low, high=high, dtype=dtype, **kwargs)
 
   @staticmethod
-  def normal(*shape, mean=0.0, std=1.0, requires_grad:bool|None=None, **kwargs) -> Tensor:
+  def normal(*shape, mean=0.0, std=1.0, requires_grad:bool=True, **kwargs) -> Tensor:
     """
     Creates a tensor with the given shape, filled with random values from a normal distribution with the given `mean` and standard deviation `std`.
     Requires `std >= 0`.
@@ -724,7 +700,7 @@ class Tensor(OpMixin):
     return (std * Tensor.randn(*shape, **kwargs) + mean).requires_grad_(requires_grad)
 
   @staticmethod
-  def uniform(*shape, low=0.0, high=1.0, dtype:DTypeLike|None=None, requires_grad:bool|None=None, **kwargs) -> Tensor:
+  def uniform(*shape, low=0.0, high=1.0, dtype:DTypeLike|None=None, requires_grad:bool=True, **kwargs) -> Tensor:
     """
     Creates a tensor with the given shape, filled with random values from a uniform distribution over the interval `[low, high)`.
     Requires `low < high`.
@@ -815,7 +791,7 @@ class Tensor(OpMixin):
     print(Tensor.randperm(6).numpy())
     ```
     """
-    return Tensor.rand(n, device=device, **kwargs).argsort().cast(dtype).requires_grad_(kwargs.get("requires_grad"))
+    return Tensor.rand(n, device=device, **kwargs).argsort().cast(dtype).requires_grad_(kwargs.get("requires_grad", True))
 
   def multinomial(self:Tensor, num_samples:int = 1, replacement:bool = False) -> Tensor:
     """
@@ -883,7 +859,7 @@ class Tensor(OpMixin):
     """
     all_uops = self.uop.toposort()
     tensors_need_grad: list[Tensor] = [t for tref in all_tensors if (t:=tref()) is not None and \
-                                       t.uop in all_uops and t.requires_grad]
+                                       t.uop in all_uops and t.requires_grad and t.is_floating_point()]
     # clear contexts
     for t,g in zip(tensors_need_grad, self.gradient(*tensors_need_grad, gradient=gradient)):
       assert g.shape == t.shape, f"grad shape must match tensor shape, {g.shape!r} != {t.shape!r}"
@@ -1027,12 +1003,13 @@ class Tensor(OpMixin):
 
   def __setitem__(self, indices, v:Tensor|PyConst|list|tuple) -> None:
     if isinstance(v, Tensor) and v.dtype != self.dtype: raise RuntimeError(f"setitem dtype mismatch: {self.dtype=} != {v.dtype=}")
-    if self.requires_grad or (isinstance(v, Tensor) and v.requires_grad):
-      # for +=/-=, v's graph references self.uop through the view — exclude those from the stale-use check
-      v_uop, v_bw = (v.uop, v.uop.backward_slice) if isinstance(v, Tensor) else (None, {})
-      if any(self.uop in t.uop.backward_slice for tref in all_tensors
-             if (t:=tref()) is not None and t is not self and t.uop is not v_uop and t.uop not in v_bw):
-        raise RuntimeError("can't setitem on a tensor that already has other uses and requires grad")
+    # raise if mutation would diverge from eager (allow only pure views of a realized buffer; exclude +=/-= RHS via v_uop/v_bw)
+    v_uop, v_bw = (v.uop, v.uop.backward_slice) if isinstance(v, Tensor) else (None, {})
+    shared = self.uop.base if self.uop.base.is_realized else None
+    if any(self.uop in t.uop.backward_slice_with_self and t.uop.base is not shared for tref in all_tensors
+           if (t:=tref()) is not None and t is not self and t.uop is not v_uop and t.uop not in v_bw):
+      raise RuntimeError("can't setitem on a tensor with other uses")
+    if not self.uop.base.is_realized and self.is_floating_point() and (self.requires_grad or (isinstance(v, Tensor) and v.requires_grad)):
       if not isinstance(v, Tensor): v = Tensor(v, device=self.device, dtype=self.dtype)
       # __iadd__/__isub__ creates AFTER(view, STORE(view, computed)); unwrap to get the computed value
       if v.uop.op is Ops.AFTER and any(s.op is Ops.STORE for s in v.uop.src[1:]): v = v._apply_uop(lambda x: x.src[1].src[1])
@@ -1425,8 +1402,6 @@ class Tensor(OpMixin):
     """
     Bitcasts `self` to the given `dtype` of the same itemsize.
 
-    `self` must not require a gradient.
-
     ```python exec="true" source="above" session="tensor" result="python"
     t = Tensor([-1, 2, 3], dtype=dtypes.int32)
     print(t.dtype, t.numpy())
@@ -1436,7 +1411,6 @@ class Tensor(OpMixin):
     print(t.dtype, t.numpy())
     ```
     """
-    if self.requires_grad: raise RuntimeError("can't backprop through bitcast")
     dt = to_dtype(dtype)
     if (ns:=dt.itemsize) != (os:=self.dtype.itemsize) and (self.shape[-1]*os) % ns != 0: raise RuntimeError("unsupported size in bitcast")
     if (not isinstance(self.device, str) or not self.device.startswith("DISK")) and ns != os:

tinygrad/uop/__init__.py

@@ -118,9 +118,6 @@ class GroupOp:
   # TODO: is BITCAST always Elementwise if it's shape changing?
   Elementwise = set.union(ALU, {Ops.CAST, Ops.BITCAST})
 
-  # all ops that support shape broadcasting
-  Broadcastable = set.union(Elementwise, {Ops.CAST, Ops.GROUP, Ops.STORE})
-
   Defines = {Ops.PARAM, Ops.DEFINE_LOCAL, Ops.DEFINE_REG}
 
   Irreducible = {Ops.CONST, Ops.DEFINE_VAR, Ops.SPECIAL, Ops.RANGE}

tinygrad/uop/ops.py

@@ -51,10 +51,7 @@ def _align_left(*shapes:tuple[sint, ...]) -> tuple[tuple[sint, ...], ...]:
   max_dim = max(len(s) for s in shapes)
   return tuple((1,)*(max_dim-len(s))+s for s in shapes)
 def _broadcast_shape(*shapes:tuple[sint, ...]) -> tuple[sint, ...]:
-  ret = tuple(0 if 0 in nth_dim_sizes else smax(nth_dim_sizes) for nth_dim_sizes in zip(*_align_left(*shapes)))
-  if not all(resolve(s == ns) or resolve(s == 1) for shape in _align_left(*shapes) for s,ns in zip(shape, ret)):
-    raise ValueError(f"shape mismatch: objects cannot be broadcast to a single shape {shapes}")
-  return ret
+  return tuple(0 if 0 in nth_dim_sizes else smax(nth_dim_sizes) for nth_dim_sizes in zip(*_align_left(*shapes)))
 
 def ssimplify(uop:sint): return uop.ssimplify() if isinstance(uop, UOp) else uop
 def sym_infer(uop: UOp|int, var_vals: dict[str, int]) -> int: return uop.sym_infer(var_vals) if isinstance(uop, UOp) else uop
@@ -320,10 +317,13 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
           return tuple(1 if i in axis_arg else s for i,s in enumerate(ps))
 
     # elementwise ops keep the shape the same. all inputs with shape must match
-    if self.op in GroupOp.Broadcastable:
+    if self.op in GroupOp.ALU.union({Ops.CAST, Ops.GROUP, Ops.STORE}):
       input_shapes = [x._shape for x in self.src]
       assert len(self.src) > 0 and all(x is not None for x in input_shapes), f"None input shape not supported for {self.op}"
-      return _broadcast_shape(*input_shapes)
+      # TODO: add broadcasting here
+      if not all_same(input_shapes):
+        raise RuntimeError(f"shape mismatch at {self.op}: {input_shapes} {[x.op for x in self.src]}")
+      return input_shapes[0]
 
     # all Ops must be explicitly handled
     raise NotImplementedError(f"no shape handling for {self.op} with {self.dtype}")
@@ -483,6 +483,10 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     return UOp(Ops.CONTRACT, dtype=self.dtype.vec(prod([x.vmax+1 for x in rngs])), src=(self,), arg=tuple((x.arg[0], x.vmax+1) for x in rngs))
   def alu(self, op, *src:UOp, **kwargs):
     all_srcs = (self, *src)
+    # broadcast shaped operands to a common shape (None and () are falsy, so only real shapes participate)
+    if (shapes := [s for x in all_srcs if (s:=x._shape)]) and not all_same(shapes):
+      out_shape = _broadcast_shape(*shapes)
+      all_srcs = tuple(x._broadcast_to(out_shape) if x._shape else x for x in all_srcs)
     out_dtype = all_srcs[-1].dtype
     if op in {Ops.CMPLT, Ops.CMPNE, Ops.CMPEQ}: out_dtype = dtypes.bool.vec(out_dtype.count) if out_dtype.count > 1 else dtypes.bool
     return UOp(op, out_dtype, all_srcs, **kwargs)