From df000116ea6d2da1a8dd840a0cd951553922a889 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Tue, 2 Jun 2026 13:02:28 -0700 Subject: [PATCH] more renderer cleanups --- tinygrad/dtype.py | 4 ++-- tinygrad/renderer/__init__.py | 12 ++++++------ 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/tinygrad/dtype.py b/tinygrad/dtype.py index a76c8b6da7..73925ed683 100644 --- a/tinygrad/dtype.py +++ b/tinygrad/dtype.py @@ -56,7 +56,7 @@ class AddrSpace(Enum): @dataclass(frozen=True, eq=False) class DType(metaclass=DTypeMetaClass): priority: int # this determines when things get upcasted - bitsize: int + bitsize: int # this is the bitsize of the base dtype name: str fmt: FmtStr|None count: int @@ -76,7 +76,7 @@ class DType(metaclass=DTypeMetaClass): def vec(self, sz:int) -> DType: assert self.count == 1, f"can't vectorize {self} with size {sz}" if sz == 1 or self == dtypes.void: return self # void doesn't vectorize, and sz=1 is scalar - return DType(self.priority, self.bitsize*sz, f"{INVERSE_DTYPES_DICT[self.name]}{sz}", None, sz, self) + return DType(self.priority, self.bitsize, f"{INVERSE_DTYPES_DICT[self.name]}{sz}", None, sz, self) def ptr(self, size=-1, addrspace=AddrSpace.GLOBAL) -> PtrDType: return PtrDType(self.priority, self.bitsize, self.name, self.fmt, self.count, None, self, addrspace, 1, size) def scalar(self) -> DType: return self._scalar if self._scalar is not None else self diff --git a/tinygrad/renderer/__init__.py b/tinygrad/renderer/__init__.py index 9552b27700..f8f15c146c 100644 --- a/tinygrad/renderer/__init__.py +++ b/tinygrad/renderer/__init__.py @@ -3,7 +3,7 @@ from typing import Callable, cast from dataclasses import dataclass from tinygrad.helpers import prod, Target, EMULATED_DTYPES from tinygrad.uop.ops import Ops, UOp, sint, ssimplify, smin, GroupOp, PatternMatcher -from tinygrad.dtype import AddrSpace, PtrDType, DType, dtypes +from tinygrad.dtype import AddrSpace, DType, dtypes from tinygrad.codegen.opt.tc import TensorCore from tinygrad.device import Compiler @@ -41,7 +41,7 @@ class Estimates: while len(buf.src) and buf.op is not Ops.PARAM: buf = buf.src[0] if buf.op is Ops.PARAM: # u.src[0] is INDEX, cap at buffer size for re-reads (e.g. matmul) - accessed = mem.get((buf, u.op), 0) + u.src[0].dtype.base.itemsize * mults + accessed = mem.get((buf, u.op), 0) + u.max_numel() * u.src[0].dtype.itemsize * mults mem[(buf, u.op)] = smin(accessed, buf.max_numel() * buf.dtype.itemsize) if u.op is Ops.RANGE: mult_stack.append(mults) @@ -51,10 +51,10 @@ class Estimates: elif u.op is Ops.END: mults = mult_stack.pop(-1) elif u.op is Ops.SPECIAL: mults *= cast(sint, u.src[0].ssimplify()) # NOTE: we don't push to the mult_stack here, you can't end these elif u.op is Ops.DEFINE_VAR and u.arg[0] == 'core_id': mults *= u.arg[2] + 1 - 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 is Ops.LOAD and u.src[0].addrspace != AddrSpace.REG: + lds += u.max_numel() * u.dtype.itemsize * mults + elif u.op is Ops.STORE and u.src[0].addrspace != AddrSpace.REG: + lds += u.max_numel() * 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, sum(mem.values()))