tinygrad 0.10.2__py3-none-any.whl → 0.11.0__py3-none-any.whl

tinygrad/nn/optim.py

@@ -1,5 +1,6 @@
 # sorted in order of increasing complexity
-from tinygrad.helpers import dedup, flatten, getenv, unwrap
+import itertools
+from tinygrad.helpers import dedup, flatten, getenv, unwrap, FUSE_OPTIM
 from tinygrad.tensor import Tensor
 from tinygrad.dtype import dtypes, least_upper_dtype
 
@@ -7,7 +8,7 @@ class Optimizer:
   """
   Base class for all optimizers.
   """
-  def __init__(self, params: list[Tensor], lr: float):
+  def __init__(self, params: list[Tensor], lr: float, fused=FUSE_OPTIM):
     # if it's 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
@@ -16,9 +17,16 @@ class Optimizer:
     assert len(self.params) != 0, "optimizer must have at least one param"
     self.device = self.params[0].device
     self.buffers: list[Tensor] = dedup([x for x in params if not x.requires_grad])   # buffers are still realized
+    self.fused = fused
     # store lr in at least float32 precision
     self.lr = Tensor(lr if getenv("CONST_LR") else [lr], requires_grad=False, device=self.device,
                      dtype=least_upper_dtype(dtypes.default_float, dtypes.float32))
+    if self.fused: self.pos_params = list(itertools.accumulate(self.params, lambda x,y: x+y.numel(), initial=0))
+
+  def _new_optim_param(self) -> list[Tensor]:
+    param_dtype = getenv("OPTIM_DTYPE", "float32")
+    if self.fused: return [Tensor.zeros(self.pos_params[-1], dtype=param_dtype, device=self.device, requires_grad=False).contiguous()]
+    return [Tensor.zeros(*t.shape, dtype=param_dtype, device=t.device, requires_grad=False).contiguous() for t in self.params]
 
   def zero_grad(self):
     """
@@ -39,9 +47,17 @@ class Optimizer:
     if not Tensor.training: raise RuntimeError(
             f"""Tensor.training={Tensor.training}, Tensor.training must be enabled to use the optimizer.
                 - help: Consider setting Tensor.training=True before calling Optimizer.step().""")
-    return self.schedule_step_with_grads([unwrap(t.grad) for t in self.params])+self.params+self.buffers
-
-  def schedule_step_with_grads(self, grads:list[Tensor]) -> list[Tensor]: raise NotImplementedError
+    if self.fused:
+      # optimizer fusion just concatenates all the buffers, runs the _step, then splits them back up
+      out, extra = self._step([Tensor.cat(*[t.flatten() for t in self.params], dim=0)],
+                              [Tensor.cat(*[unwrap(t.grad).flatten() for t in self.params], dim=0)])
+      updated_params = [out[0][self.pos_params[i]:self.pos_params[i+1]].reshape(tt.shape) for i, tt in enumerate(self.params)]
+    else:
+      updated_params, extra = self._step(self.params, [unwrap(t.grad) for t in self.params])
+    for i, tt in enumerate(self.params): tt.assign(updated_params[i])
+    return extra+self.params+self.buffers
+
+  def _step(self, params:list[Tensor], grads:list[Tensor]) -> tuple[list[Tensor], list[Tensor]]: raise NotImplementedError
 
 class OptimizerGroup(Optimizer):
   """
@@ -54,93 +70,108 @@ class OptimizerGroup(Optimizer):
   def zero_grad(self): [o.zero_grad() for o in self.optimizers]
   def schedule_step(self) -> list[Tensor]: return [x for o in self.optimizers for x in o.schedule_step()]
 
-# LARS is essentially just trust ratio to SGD so if we just set the trust coeff 0.0 its just standard SGD.
-def SGD(params: list[Tensor], lr=0.001, momentum=0.0, weight_decay=0.0, nesterov=False, classic=False):
+# LARS is essentially just trust ratio to SGD so if we just set the trust coeff 0.0 it's just standard SGD.
+def SGD(params: list[Tensor], lr=0.001, momentum=0.0, weight_decay=0.0, nesterov=False, classic=False, fused=FUSE_OPTIM):
   """
   Stochastic Gradient Descent (SGD) optimizer with optional momentum and weight decay.
 
   `classic` is a boolean flag that determines whether to use the popular momentum update rule or the classic momentum update rule.
+  """
+  return LARS(params, lr, momentum, weight_decay, 0, None, nesterov, classic=classic, pre_wd=True, tcoef=0.0, fused=fused)
 
-  - Described: https://paperswithcode.com/method/sgd
+# Muon applies the newton schulz algorithm on gradient. also can include momentum, nesterov, and weight decay
+def Muon(params: list[Tensor], lr=0.02, momentum=0.95, weight_decay=0.0, ns_steps=5, ns_params=(3.4445, -4.775, 2.0315),
+         nesterov=True, fused=FUSE_OPTIM):
   """
-  return LARS(params, lr, momentum, weight_decay, nesterov, classic, tcoef=0.0)
+  SGD with newton-schulz iteration and post momentum weight decay.
+
+  - Described: https://kellerjordan.github.io/posts/muon/
+  - Paper: https://arxiv.org/pdf/2502.16982
+  """
+  assert not fused, "FUSE_OPTIM not allowed for Muon optimizer"
+  return LARS(params, lr, momentum, weight_decay, ns_steps, ns_params, nesterov, classic=False, pre_wd=False, tcoef=0.0, fused=fused)
 
 class LARS(Optimizer):
   """
   Layer-wise Adaptive Rate Scaling (LARS) optimizer with optional momentum and weight decay.
 
-  - Described: https://paperswithcode.com/method/lars
   - Paper: https://arxiv.org/abs/1708.03888v3
   """
-  def __init__(self, params:list[Tensor], lr=0.001, momentum=0.9, weight_decay=1e-4, nesterov=False, classic=True, tcoef=0.001):
-    super().__init__(params, lr)
-    self.momentum, self.wd, self.nesterov, self.classic, self.tcoef = momentum, weight_decay, nesterov, classic, tcoef
-    self.b = [Tensor.zeros(*t.shape, dtype=t.dtype, device=t.device, requires_grad=False) for t in self.params] if self.momentum else []
-
-  def schedule_step_with_grads(self, grads:list[Tensor]) -> list[Tensor]:
-    for i, (t, g) in enumerate(zip(self.params, grads)):
+  def __init__(self, params:list[Tensor], lr=0.001, momentum=0.9, weight_decay=1e-4, ns_steps=0, ns_params=None,
+               nesterov=False, classic=True, pre_wd=True, tcoef=0.001, fused=FUSE_OPTIM):
+    super().__init__(params, lr, fused)
+    self.momentum, self.wd, self.ns_steps, self.ns_params  = momentum, weight_decay, ns_steps, ns_params
+    self.nesterov, self.classic, self.pre_wd, self.tcoef = nesterov, classic, pre_wd, tcoef
+    self.b = self._new_optim_param() if self.momentum else []
+
+  def _step(self, params:list[Tensor], grads:list[Tensor]) -> tuple[list[Tensor], list[Tensor]]:
+    ret = []
+    for i, (t, g) in enumerate(zip(params, grads)):
       if self.tcoef != 0:
         r1 = t.detach().square().sum().sqrt()
         r2 = g.square().sum().sqrt()
-        r = (r1 > 0).where((r2 > 0).where(self.tcoef * r1 / (r2 + self.wd * r1), 1.0), 1.0)
+        r:Tensor|float = (r1 > 0).where((r2 > 0).where(self.tcoef * r1 / (r2 + self.wd * r1), 1.0), 1.0)
       else: r = 1.0
-      g = g + self.wd * t.detach()
+      if self.pre_wd and self.wd > 0: g = g + self.wd * t.detach()
       # classic momentum does post learning rate update
       if self.classic: g = g * r * self.lr
       if self.momentum:
-        self.b[i].assign(self.momentum * self.b[i] + g)  # NOTE: self.b[i] is zero on the first run, no if required
+        # TODO: this contiguous is required for correctness because self.b[i] becomes a non contiguous view
+        # the scheduler should detect this and just insert contiguous
+        self.b[i].assign(self.momentum * self.b[i].contiguous() + g)  # NOTE: self.b[i] is zero on the first run, no if required
         g = (g + self.momentum * self.b[i]) if self.nesterov else self.b[i]
+      if self.ns_params: g = g.reshape(g.shape[0], -1).newton_schulz(self.ns_steps, self.ns_params).reshape(g.shape)
+      # muon does post momentum weight decay
+      if not self.pre_wd and self.wd > 0: t = t.detach() * (1.0 - self.wd * self.lr)
       # popular momentum does pre learning rate update
       if not self.classic: g = g * r * self.lr
-      t.assign((t.detach() - g).cast(t.dtype))
-    return self.b
+      ret.append((t.detach() - g).cast(t.dtype))
+    return ret, self.b
 
-# LAMB is essentially just the trust ratio part of LARS applied to Adam/W so if we just set the trust ratio to 1.0 its just Adam/W.
-def AdamW(params: list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-8, weight_decay=0.01):
+# LAMB is essentially just the trust ratio part of LARS applied to Adam/W so if we just set the trust ratio to 1.0 it's just Adam/W.
+def AdamW(params: list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-8, weight_decay=0.01, fused=FUSE_OPTIM):
   """
   AdamW optimizer with optional weight decay.
 
-  - Described: https://paperswithcode.com/method/adamw
   - Paper: https://arxiv.org/abs/1711.05101v3
   """
-  return LAMB(params, lr, b1, b2, eps, weight_decay, adam=True)
-def Adam(params: list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-8):
+  return LAMB(params, lr, b1, b2, eps, weight_decay, adam=True, fused=fused)
+def Adam(params: list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-8, fused=FUSE_OPTIM):
   """
   Adam optimizer.
 
-  - Described: https://paperswithcode.com/method/adam
   - Paper: https://arxiv.org/abs/1412.6980
   """
-  return LAMB(params, lr, b1, b2, eps, 0.0, adam=True)
+  return LAMB(params, lr, b1, b2, eps, 0.0, adam=True, fused=fused)
 
 class LAMB(Optimizer):
   """
   LAMB optimizer with optional weight decay.
 
-  - Described: https://paperswithcode.com/method/lamb
   - Paper: https://arxiv.org/abs/1904.00962
   """
-  def __init__(self, params: list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-6, weight_decay=0.0, adam=False):
-    super().__init__(params, lr)
+  def __init__(self, params: list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-6, weight_decay=0.0, adam=False, fused=FUSE_OPTIM):
+    super().__init__(params, lr, fused)
     self.b1, self.b2, self.eps, self.wd, self.adam = b1, b2, eps, weight_decay, adam
     self.b1_t, self.b2_t = (Tensor.ones((1,), dtype=dtypes.float32, device=self.device, requires_grad=False).contiguous() for _ in [b1, b2])
-    self.m = [Tensor.zeros(*t.shape, dtype=dtypes.float32, device=t.device, requires_grad=False).contiguous() for t in self.params]
-    self.v = [Tensor.zeros(*t.shape, dtype=dtypes.float32, device=t.device, requires_grad=False).contiguous() for t in self.params]
+    self.m = self._new_optim_param()
+    self.v = self._new_optim_param()
 
-  def schedule_step_with_grads(self, grads:list[Tensor]) -> list[Tensor]:
+  def _step(self, params:list[Tensor], grads:list[Tensor]) -> tuple[list[Tensor], list[Tensor]]:
+    ret = []
     self.b1_t *= self.b1
     self.b2_t *= self.b2
-    for i, (t, g) in enumerate(zip(self.params, grads)):
-      self.m[i].assign(self.b1 * self.m[i] + (1.0 - self.b1) * g)
-      self.v[i].assign(self.b2 * self.v[i] + (1.0 - self.b2) * (g * g))
+    for i, (t, g) in enumerate(zip(params, grads)):
+      self.m[i].assign((self.b1 * self.m[i] + (1.0 - self.b1) * g).cast(self.m[i].dtype))
+      self.v[i].assign((self.b2 * self.v[i] + (1.0 - self.b2) * (g * g)).cast(self.v[i].dtype))
       m_hat = self.m[i] / (1.0 - self.b1_t)
       v_hat = self.v[i] / (1.0 - self.b2_t)
       up = (m_hat / (v_hat.sqrt() + self.eps)) + self.wd * t.detach()
       if not self.adam:
         r1 = t.detach().square().sum().sqrt()
         r2 = up.square().sum().sqrt()
-        r = Tensor.where(r1 > 0, Tensor.where(r2 > 0, r1 / r2, 1.0), 1.0)
+        r: Tensor|float = Tensor.where(r1 > 0, Tensor.where(r2 > 0, r1 / r2, 1.0), 1.0)
       else:
         r = 1.0
-      t.assign((t.detach() - self.lr * r * up).cast(t.dtype))
-    return [self.b1_t, self.b2_t] + self.m + self.v
+      ret.append((t.detach() - self.lr * r * up).cast(t.dtype))
+    return ret, [self.b1_t, self.b2_t] + self.m + self.v

tinygrad/nn/state.py

@@ -1,6 +1,6 @@
 import json, pathlib, zipfile, pickle, tarfile, struct, functools, io
 from collections import OrderedDict
-from typing import Union, Optional, Any, Callable, BinaryIO, Iterable
+from typing import Any, Callable, BinaryIO, Iterable
 from tinygrad.tensor import Tensor
 from tinygrad.dtype import dtypes
 from tinygrad.helpers import prod, argsort, DEBUG, Timing, CI, unwrap, GlobalCounters, tqdm, round_up, T
@@ -35,22 +35,22 @@ safe_dtypes = {"BOOL":dtypes.bool, "I8":dtypes.int8, "U8":dtypes.uint8, "I16":dt
                "I64":dtypes.int64, "U64":dtypes.uint64, "F16":dtypes.float16, "BF16":dtypes.bfloat16, "F32":dtypes.float32, "F64":dtypes.float64}
 inverse_safe_dtypes = {v:k for k,v in safe_dtypes.items()}
 
-def accept_filename(func: Callable[[Tensor], T]) -> Callable[[Union[Tensor, str, pathlib.Path]], T]:
+def accept_filename(func: Callable[[Tensor], T]) -> Callable[[Tensor|str|pathlib.Path], T]:
   @functools.wraps(func)
-  def wrapper(fn: Union[Tensor, str, pathlib.Path]) -> T: return func(Tensor(pathlib.Path(fn)) if not isinstance(fn, Tensor) else fn)
+  def wrapper(fn: Tensor|str|pathlib.Path) -> T: return func(Tensor(pathlib.Path(fn)) if not isinstance(fn, Tensor) else fn)
   return wrapper
 
 @accept_filename
 def safe_load_metadata(t:Tensor) -> tuple[Tensor, int, dict[str, Any]]:
   """
-  Loads a .safetensor file from disk, returning the data, metadata length, and metadata.
+  Loads a .safetensor file, returning the source tensor, data start position, and metadata.
   """
   data_start = int.from_bytes(t[0:8].data(), "little") + 8
   return t, data_start, json.loads(t[8:data_start].data().tobytes())
 
-def safe_load(fn:Union[Tensor, str, pathlib.Path]) -> dict[str, Tensor]:
+def safe_load(fn:Tensor|str|pathlib.Path) -> dict[str, Tensor]:
   """
-  Loads a .safetensor file from disk, returning the state_dict.
+  Loads a .safetensor file, returning the `state_dict`.
 
   ```python
   state_dict = nn.state.safe_load("test.safetensor")
@@ -61,9 +61,9 @@ def safe_load(fn:Union[Tensor, str, pathlib.Path]) -> dict[str, Tensor]:
   return { k: data[v['data_offsets'][0]:v['data_offsets'][1]].bitcast(safe_dtypes[v['dtype']]).reshape(v['shape'])
           for k, v in metadata.items() if k != "__metadata__" }
 
-def safe_save(tensors:dict[str, Tensor], fn:str, metadata:Optional[dict[str, Any]]=None):
+def safe_save(tensors:dict[str, Tensor], fn:str, metadata:dict[str, Any]|None=None):
   """
-  Saves a state_dict to disk in a .safetensor file with optional metadata.
+  Saves a `state_dict` to disk in a .safetensor file with optional metadata.
 
   ```python
   t = Tensor([1, 2, 3])
@@ -87,7 +87,7 @@ def safe_save(tensors:dict[str, Tensor], fn:str, metadata:Optional[dict[str, Any
 
 def get_state_dict(obj, prefix:str='', tensor_type=Tensor) -> dict[str, Tensor]:
   """
-  Returns a state_dict of the object, with optional prefix.
+  Returns a `state_dict` of the object, with optional prefix.
 
   ```python exec="true" source="above" session="tensor" result="python"
   class Net:
@@ -124,9 +124,9 @@ def get_parameters(obj) -> list[Tensor]:
   """
   return list(get_state_dict(obj).values())
 
-def load_state_dict(model, state_dict:dict[str, Tensor], strict=True, verbose=True, consume=False) -> None:
+def load_state_dict(model, state_dict:dict[str, Tensor], strict=True, verbose=True, consume=False, realize=True) -> list[Tensor]:
   """
-  Loads a state_dict into a model.
+  Loads a `state_dict` into a model. Return the loaded Tensors.
 
   ```python
   class Net:
@@ -140,7 +140,9 @@ def load_state_dict(model, state_dict:dict[str, Tensor], strict=True, verbose=Tr
   ```
   """
   start_mem_used = GlobalCounters.mem_used
-  with Timing("loaded weights in ", lambda et_ns: f", {(B:=(GlobalCounters.mem_used-start_mem_used))/1e9:.2f} GB loaded at {B/et_ns:.2f} GB/s"):
+  ret = []
+  with Timing("loaded weights in ",
+              lambda et_ns: f", {(B:=(GlobalCounters.mem_used-start_mem_used))/1e9:.2f} GB loaded at {B/et_ns:.2f} GB/s", enabled=verbose):
     model_state_dict = get_state_dict(model)
     if DEBUG >= 1 and len(state_dict) > len(model_state_dict):
       print("WARNING: unused weights in state_dict", sorted(list(state_dict.keys() - model_state_dict.keys())))
@@ -152,15 +154,22 @@ def load_state_dict(model, state_dict:dict[str, Tensor], strict=True, verbose=Tr
       if v.shape != state_dict[k].shape:
         raise ValueError(f'Shape mismatch in layer `{k}`: Expected shape {v.shape}, but found {state_dict[k].shape} in state dict.')
       if isinstance(v.device, tuple):
-        if isinstance(state_dict[k].device, tuple): v.replace(state_dict[k]).realize()
-        else: v.replace(state_dict[k].shard(v.device, v.lazydata.axis)).realize()
-      else: v.replace(state_dict[k].to(v.device)).realize()
+        if isinstance(state_dict[k].device, tuple): v.replace(state_dict[k])
+        else: v.replace(state_dict[k].shard(v.device, v.uop.axis))
+      else: v.replace(state_dict[k].to(v.device))
+      if realize: v.realize()
       if consume: del state_dict[k]
+      ret.append(v)
+  return ret
 
 @accept_filename
 def tar_extract(t: Tensor) -> dict[str, Tensor]:
   """
-  Extracts files from a tar archive and returns them as dictionary of names (keys) and tensors (values).
+  ```python
+  tar_extract(fn: Tensor | str | Path) -> dict[str, Tensor]
+  ```
+
+  Extracts files from a tar archive and returns them as a dictionary of names (keys) and tensors (values).
 
   ```python
   tensors = nn.state.tar_extract(Tensor(pathlib.Path("archive.tar")))
@@ -174,14 +183,18 @@ def tar_extract(t: Tensor) -> dict[str, Tensor]:
 @accept_filename
 def torch_load(t:Tensor) -> dict[str, Tensor]:
   """
-  Loads a torch .pth file from disk.
+  ```python
+  torch_load(fn: Tensor | str | Path) -> dict[str, Tensor]
+  ```
+
+  Loads a torch .pth file, returning the `state_dict`.
 
   ```python
   state_dict = nn.state.torch_load("test.pth")
   ```
   """
-  offsets: dict[Union[str, int], int] = {}
-  lens: dict[Union[str, int], int] = {}
+  offsets: dict[str|int, int] = {}
+  lens: dict[str|int, int] = {}
   def _rebuild_tensor_v2(storage, storage_offset, size, stride, requires_grad=None, backward_hooks=None, metadata=None):
     #print(storage, storage_offset, size, stride, requires_grad, backward_hooks, metadata)
     lens[storage[2]] = storage[4] * storage[1].itemsize
@@ -292,13 +305,14 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
 @accept_filename
 def gguf_load(tensor: Tensor) -> tuple[dict, dict[str, Tensor]]:
   """
-  Loads a gguf file from a tensor.
+  Loads a .gguf file, returning the `kv_data` and `state_dict`.
 
   ```python
-  fn = "Meta-Llama-3-8B-Instruct.Q4_0.gguf"
-  gguf_tensor = Tensor.empty(os.stat(fn).st_size, dtype=dtypes.uint8, device=f"disk:{fn}").to(Device.DEFAULT)
-  kv_data, state_dict = gguf_load(gguf_tensor)
+  gguf_tensor = Tensor(pathlib.Path("Meta-Llama-3-8B-Instruct.Q4_0.gguf")).to(Device.DEFAULT)
+  kv_data, state_dict = nn.state.gguf_load(gguf_tensor)
   ```
+
+  NOTE: The provided tensor must be on a device that supports execution.
   """
   reader, kv_data, state_dict = io.BufferedReader(TensorIO(tensor), 1_000_000), {}, {}
   def read_unpack(fmt: str, n: int): return struct.unpack(fmt, reader.read(n))[0]

tinygrad/renderer/__init__.py

@@ -1,46 +1,13 @@
 from __future__ import annotations
-from typing import Optional, Callable
-import functools, math
-from enum import Enum, auto
+from typing import Callable, cast, TYPE_CHECKING
+import functools, itertools
 from dataclasses import dataclass, field, replace
 from tinygrad.helpers import to_function_name, dedup, prod
-from tinygrad.ops import Ops, UOp, sym_infer, sint, Variable, ssimplify, GroupOp, PatternMatcher
-from tinygrad.dtype import DType
-
-class OptOps(Enum):
-  TC = auto(); UPCAST = auto(); UNROLL = auto(); LOCAL = auto() # noqa: E702
-  GROUP = auto(); GROUPTOP = auto(); NOLOCALS = auto(); PADTO = auto(); SWAP = auto() # noqa: E702
-  def __lt__(self, x:OptOps): return self.value < x.value
-
-@dataclass(frozen=True, order=True)
-class Opt:
-  op: OptOps
-  axis: Optional[int] = None
-  arg: Optional[int | tuple] = None
-  def __repr__(self): return f"Opt(op={self.op}, axis={self.axis}, arg={self.arg})"
-
-@dataclass(frozen=True)
-class TensorCore: # D = A * B + C, A is (M x K), B is (K x N), C and D are (M x N)
-  dims: tuple[int,int,int] # N, M, K
-  threads: int # number of threads that construct the warp
-  elements_per_thread: tuple[int, int, int] # elements per-thread to load/store from A/B/C
-  dtype_in: DType # dtype for A and B
-  dtype_out: DType # dtype for C and D
-  opts: tuple[str, ...] # ordered tuple of "ux" or "lx" specifing kernel opts to perform. "ux" upcasts dim x and "lx" localizes dim x
-  swizzle: tuple[Optional[tuple[tuple[int, ...], tuple[int, ...]]], Optional[tuple[tuple[int, ...], tuple[int, ...]]]] = (None, None)
-  def get_reduce_axes(self): return [(i, 2) for i in range(int(math.log2(self.dims[2])))]
-  def get_upcast_axes(self): return [opt for opt in self.opts if opt[0] == "u"]
-  def get_local_axes(self): return [opt for opt in self.opts if opt[0] == "l"]
-  def __str__(self): return "_".join(["WMMA"] + list(map(str, self.dims)) + [self.dtype_in.name, self.dtype_out.name])
-  def __post_init__(self):
-    local_axes, upcast_axes, reduce_axes = len(self.get_local_axes()), len(self.get_upcast_axes()), len(self.get_reduce_axes())
-    assert self.dims[0] * self.dims[1] == 2**(local_axes + upcast_axes), (
-      f"N({self.dims[0]}) x M({self.dims[1]}) != local({2**local_axes}) x upcast({2**upcast_axes}) with opts({self.opts})")
-    assert 2**local_axes == self.threads, f"{self.threads} threads construct the warp but found {2**local_axes} in {self.opts}"
-    assert 2**upcast_axes == self.elements_per_thread[2], (
-      f"{self.elements_per_thread[2]} elements from C are processed per thread but found {2**upcast_axes} in {self.opts}")
-    assert all(len(perm[0]) == local_axes and len(perm[1]) == reduce_axes + upcast_axes for perm in self.swizzle if perm), (
-      f"swizzle perm should be of len (({local_axes})({reduce_axes + upcast_axes}))")
+from tinygrad.uop.ops import Ops, UOp, sym_infer, sint, Variable, ssimplify, GroupOp, PatternMatcher
+from tinygrad.dtype import AddrSpace, PtrDType
+if TYPE_CHECKING:
+  from tinygrad.codegen.opt.tc import TensorCore
+  from tinygrad.codegen.opt.kernel import Opt
 
 @dataclass(frozen=True)
 class Estimates:
@@ -61,19 +28,23 @@ class Estimates:
     dont_count: set[UOp] = set()
     if ignore_indexing:
       for u in uops:
-        if u.op in {Ops.LOAD, Ops.STORE}:
-          dont_count = dont_count.union(u.src[0].toposort)
-          if len(u.src) > 2: dont_count = dont_count.union(u.src[2].toposort)
+        if u.op in {Ops.LOAD, Ops.STORE} and (not isinstance(u.src[0].dtype, PtrDType) or u.src[0].dtype.addrspace != AddrSpace.REG):
+          dont_count = dont_count.union(u.src[0].toposort())
+          if len(u.src) > 2: dont_count = dont_count.union(u.src[2].toposort())
         elif u.op is Ops.IF:
-          dont_count = dont_count.union(u.src[0].toposort)
+          dont_count = dont_count.union(u.src[0].toposort())
     for u in uops:
       if u.op is Ops.RANGE:
         mult_stack.append(mults)
-        mults *= (u.src[1] - u.src[0]).ssimplify()
+        mults *= cast(sint, u.src[0].ssimplify())
+        # SPECIAL are already counted in mults
+        mults = mults.substitute({x:x.const_like(0) for x in mults.toposort() if x.op is Ops.SPECIAL}) if isinstance(mults, UOp) else mults
       elif u.op is Ops.ENDRANGE: mults = mult_stack.pop(-1)
       elif u.op is Ops.SPECIAL: mults *= u.arg[1] # NOTE: we don't push to the mult_stack here, you can't end these
-      elif u.op is Ops.LOAD: lds += u.dtype.itemsize * mults
-      elif u.op is Ops.STORE: lds += u.src[1].dtype.itemsize * mults
+      elif u.op is Ops.LOAD and (not isinstance(u.src[0].dtype, PtrDType) or u.src[0].dtype.addrspace != AddrSpace.REG):
+        lds += u.dtype.itemsize * mults
+      elif u.op is Ops.STORE and (not isinstance(u.src[0].dtype, PtrDType) or u.src[0].dtype.addrspace != AddrSpace.REG):
+        lds += u.src[1].dtype.itemsize * mults
       elif u.op in GroupOp.ALU and u not in dont_count: flops += (mults * (2 if u.op is Ops.MULACC else 1)) * u.dtype.count
       elif u.op is Ops.WMMA and u not in dont_count: flops += 2 * prod(u.arg[1]) // u.arg[5] * mults
     return Estimates(flops, lds, lds) # TODO: properly track memory, lds is always a high estimate
@@ -84,13 +55,11 @@ class ProgramSpec:
   src:str
   device:str
   ast:UOp  # save the base ast (this is method cache key)
-  uops:Optional[list[UOp]]=None
-  applied_opts:Optional[list[Opt]]=None
-  mem_estimate:sint=0  # TODO: get this from the load/store uops once min/max are good
+  uops:list[UOp]|None=None
 
   # filled in from uops (if we have uops)
-  global_size:Optional[list[int]]=None
-  local_size:Optional[list[int]]=None
+  global_size:list[int]|None=None
+  local_size:list[int]|None=None
   vars:list[Variable]=field(default_factory=list)
   globals:list[int]=field(default_factory=list)
   outs:list[int]=field(default_factory=list)
@@ -103,19 +72,26 @@ class ProgramSpec:
       for u in self.uops:
         if u.op is Ops.DEFINE_VAR: self.vars.append(u)
         if u.op is Ops.DEFINE_GLOBAL: self.globals.append(u.arg)
-        if u.op is Ops.STORE: self.outs.extend([x.arg for x in u.src[0].toposort if x.op is Ops.DEFINE_GLOBAL])
-        if u.op is Ops.LOAD: self.ins.extend([x.arg for x in u.src[0].toposort if x.op is Ops.DEFINE_GLOBAL])
+        if u.op is Ops.STORE: self.outs.extend([x.arg for x in u.src[0].toposort() if x.op is Ops.DEFINE_GLOBAL])
+        if u.op is Ops.LOAD: self.ins.extend([x.arg for x in u.src[0].toposort() if x.op is Ops.DEFINE_GLOBAL])
         if u.op is Ops.SPECIAL:
           # NOTE: you have to set local_size and global_size to the base [1,1,1] outside this
           if u.arg[0][0] == 'i': self.local_size = None
           special_size = self.local_size if u.arg[0][0] == 'l' else self.global_size
-          assert special_size is not None
-          special_size[int(u.arg[0][-1])] = u.arg[1]
+          if special_size is not None: special_size[int(u.arg[0][-1])] = u.arg[1]
       self.vars = sorted(self.vars, key=lambda v: v.arg)
       self.outs = sorted(dedup(self.outs))
       self.ins = sorted(dedup(self.ins))
       self._ran_post_init = True
 
+  @functools.cached_property
+  def mem_estimate(self) -> sint:
+    # group non-local bufs by the op type (LOAD or STORE) and the buffer arg. take the max access of that buffer in bytes
+    # TODO: these max and min don't work on symbolic, and results are very wrong.
+    return sum(max(x.src[0].dtype.nbytes() for x in group)
+      for _, group in itertools.groupby([x for x in self.ast.toposort() if x.op in {Ops.LOAD, Ops.STORE} and x.src[0].base.op is Ops.DEFINE_GLOBAL],
+                        key=lambda x: (x.op, x.src[0].base.arg)))
+
   @functools.cached_property
   def estimates(self) -> Estimates:
     return replace(Estimates() if self.uops is None else Estimates.from_uops(self.uops, ignore_indexing=True), mem=self.mem_estimate)
@@ -123,6 +99,10 @@ class ProgramSpec:
   @functools.cached_property
   def function_name(self) -> str: return to_function_name(self.name)
 
+  @property
+  def applied_opts(self) -> tuple[Opt, ...]|None: return self.uops[-1].arg.applied_opts if \
+    self.uops is not None and self.uops[-1].op is Ops.SINK and self.uops[-1].arg is not None else None
+
   def launch_dims(self, var_vals:dict[Variable, int]):
     global_size = [sym_infer(sz, var_vals) for sz in self.global_size] if self.global_size is not None else None
     local_size = [sym_infer(sz, var_vals) for sz in self.local_size] if self.local_size is not None else None
@@ -136,12 +116,12 @@ class Renderer:
   has_local: bool = True
   has_shared: bool = True
   # NOTE: these two should be in (x,y,z) order to match the max_sizes argument in get_grouped_dims
-  global_max: Optional[tuple[int, ...]] = (0x8FFFFFFF,) * (3) # TODO: Ops.SPECIAL int32 indexes right now
-  local_max: Optional[tuple[int, ...]] = (0x8FFFFFFF,) * (3) # TODO: Ops.SPECIAL int32 indexes right now
+  global_max: tuple[int, ...]|None = (0x8FFFFFFF,) * (3) # TODO: Ops.SPECIAL int32 indexes right now
+  local_max: tuple[int, ...]|None = (0x8FFFFFFF,) * (3) # TODO: Ops.SPECIAL int32 indexes right now
   shared_max: int = 32768
   tensor_cores: list[TensorCore] = []
-  pre_matcher: Optional[PatternMatcher] = None
-  extra_matcher: Optional[PatternMatcher] = None
+  pre_matcher: PatternMatcher|None = None
+  extra_matcher: PatternMatcher|None = None
   code_for_op: dict[Ops, Callable] = {}
 
   def __reduce__(self): return self.__class__, ()