diff --git a/tinygrad/codegen/linearizer.py b/tinygrad/codegen/linearizer.py index aa629237eb..2406056989 100644 --- a/tinygrad/codegen/linearizer.py +++ b/tinygrad/codegen/linearizer.py @@ -5,11 +5,11 @@ from collections import defaultdict from enum import Enum, auto from dataclasses import dataclass -from tinygrad.helpers import colored, ImageDType, DEBUG, dtypes, DType, prod, PtrDType, getenv, all_same +from tinygrad.helpers import colored, ImageDType, DEBUG, dtypes, DType, prod, PtrDType, getenv, all_same, to_function_name from tinygrad.ops import LazyOp, UnaryOps, ConstBuffer, MemBuffer, BufferOps from tinygrad.ops import ReduceOps, BinaryOps, TernaryOps from tinygrad.shape.shapetracker import ShapeTracker -from tinygrad.shape.symbolic import Variable, NumNode, VariableOrNum, Node, SumNode, MulNode, DivNode, ModNode, LtNode, AndNode, sym_rename +from tinygrad.shape.symbolic import Variable, NumNode, VariableOrNum, Node, SumNode, MulNode, DivNode, ModNode, LtNode, AndNode from tinygrad.codegen.kernel import LocalBuffer, Kernel from tinygrad.lazy import vars_from_ast from tinygrad.features.image import to_image_idx @@ -183,13 +183,12 @@ class Linearizer(Kernel): self.buf_uops.append(self.uop(UOps.DEFINE_LOCAL, PtrDType(dtypes.float32), (), ("temp", self.sts[-1].size()))) # kernel name (before late upcast) - self.function_name = ("r_" if self.reduceop else "E_") + '_'.join([str(x) if isinstance(x, int) else sym_rename(x) for x in self.full_shape]) - self.display_name = ("r_" if self.reduceop else "E_") + colored('_', 'BLACK').join([colored(str(x), c) for x,c in zip(self.full_shape, self.colors())]) + self.name = ("r_" if self.reduceop else "E_") + colored('_', 'BLACK').join([colored(str(x), c) for x,c in zip(self.full_shape, self.colors())]) # name the function something unique - Linearizer.kernel_cnt[self.function_name] += 1 - suffix = f"{'n'+str(Linearizer.kernel_cnt[self.function_name]-1)}" if Linearizer.kernel_cnt[self.function_name] > 1 else "" - self.function_name, self.display_name = self.function_name+suffix, self.display_name+colored(suffix, 'BLACK') + Linearizer.kernel_cnt[(function_name := to_function_name(self.name))] += 1 + suffix = f"{'n'+str(Linearizer.kernel_cnt[function_name]-1)}" if Linearizer.kernel_cnt[function_name] > 1 else "" + self.name = self.name+colored(suffix, 'BLACK') # define indexes global_idxs, loop_global_idxs = get_grouped_dims("gidx", 0, self.full_shape[:self.global_dims], 3 if self.opts.has_local else 0) @@ -213,8 +212,6 @@ class Linearizer(Kernel): self.loop_uops.update({x.expr:self.uop(UOps.SPECIAL, dtypes.int32, (), (len(loop_global_idxs)-1-i, x.expr.replace("gidx", "idx"), x.max+1)) for i,x in enumerate(loop_global_idxs)}) elif self.opts.has_local: self.global_size, self.local_size = [x.max+1 for x in loop_global_idxs][::-1], [x.max+1 for x in loop_local_idxs][::-1] - self.global_size += [1]*(3-len(self.global_size)) - self.local_size += [1]*(3-len(self.local_size)) self.loop_uops.update({x.expr:self.uop(UOps.SPECIAL, dtypes.int32, (), (len(loop_global_idxs)-1-i, x.expr, x.max+1)) for i,x in enumerate(loop_global_idxs)}) self.loop_uops.update({x.expr:self.uop(UOps.SPECIAL, dtypes.int32, (), (len(loop_local_idxs)-1-i, x.expr, x.max+1)) for i,x in enumerate(loop_local_idxs)}) else: diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index 535adc2c1b..3c5d2026d3 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -1,5 +1,5 @@ from __future__ import annotations -import os, functools, platform, time, re, contextlib, operator, hashlib, pickle, sqlite3, cProfile, pstats, tempfile, pathlib +import os, functools, platform, time, re, contextlib, operator, hashlib, pickle, sqlite3, cProfile, pstats, tempfile, pathlib, string import numpy as np from urllib import request from tqdm import tqdm @@ -47,6 +47,8 @@ def get_child(obj, key): else: obj = getattr(obj, k) return obj +@functools.lru_cache(maxsize=None) +def to_function_name(s:str): return ''.join([c if c in (string.ascii_letters+string.digits+'_') else f'{ord(c):02X}' for c in ansistrip(s)]) @functools.lru_cache(maxsize=None) def getenv(key:str, default=0): return type(default)(os.getenv(key, default)) diff --git a/tinygrad/ops.py b/tinygrad/ops.py index 36744ad519..2be284184c 100644 --- a/tinygrad/ops.py +++ b/tinygrad/ops.py @@ -2,7 +2,7 @@ from __future__ import annotations import importlib, inspect, functools, pathlib, time, re from enum import Enum, auto from typing import TYPE_CHECKING, Union, Type, Tuple, Any, List, Optional, Dict, Callable, Mapping -from tinygrad.helpers import ansilen, prod, DEBUG, getenv, GlobalCounters, DType, colored, BEAM, NOOPT, dedup, all_int +from tinygrad.helpers import ansilen, prod, DEBUG, getenv, GlobalCounters, DType, colored, BEAM, NOOPT, dedup, all_int, to_function_name from tinygrad.runtime.lib import RawBuffer from tinygrad.shape.symbolic import Variable, sym_infer, sint from dataclasses import dataclass @@ -228,11 +228,13 @@ def get_interpreted_fxn(fxn_for_op:Dict[Op, Callable], ast:LazyOp) -> Interprete # **************** for Compiled Buffers **************** class CompiledASTRunner(JITRunner): - def __init__(self, ast:Optional[LazyOp], name:str, prg:str, global_size:Optional[List[int]]=None, local_size:Optional[List[int]]=None, op_estimate=0, mem_estimate=0, display_name:Optional[str]=None, runtime_args:Optional[dict]=None): + def __init__(self, ast:Optional[LazyOp], name:str, prg:str, global_size:Optional[List[int]]=None, local_size:Optional[List[int]]=None, runtime_args:Optional[dict]=None): super().__init__() if DEBUG >= 4: print(prg) - self.name, self.prg, self.global_size, self.local_size, self.op_estimate, self.mem_estimate, self.display_name, self.runtime_args = \ - name, prg, global_size, local_size, op_estimate, mem_estimate, display_name, runtime_args if runtime_args is not None else {} + if global_size is not None: global_size = global_size + [1]*(3-len(global_size)) + if local_size is not None: local_size = local_size + [1]*(3-len(local_size)) + self.name, self.display_name, self.prg, self.global_size, self.local_size, self.runtime_args = \ + to_function_name(name), name, prg, global_size, local_size, runtime_args if runtime_args is not None else {} self.vars: List[Variable] = [] if ast: info = get_lazyop_info(ast) @@ -247,8 +249,8 @@ class CompiledASTRunner(JITRunner): return self def launch_dims(self, var_vals): - global_size = ([sym_infer(sz, var_vals) for sz in self.global_size] + [1]*(3-len(self.global_size))) if self.global_size is not None else self.global_size - local_size = ([sym_infer(sz, var_vals) for sz in self.local_size] + [1]*(3-len(self.local_size))) if self.local_size is not None else self.local_size + global_size = [sym_infer(sz, var_vals) for sz in self.global_size] if self.global_size is not None else self.global_size + local_size = [sym_infer(sz, var_vals) for sz in self.local_size] if self.local_size is not None else self.local_size return global_size, local_size def __call__(self, rawbufs:List[RawBuffer], var_vals:Dict[Variable, int], wait=False, jit=False) -> Optional[float]: @@ -262,7 +264,7 @@ class CompiledASTRunner(JITRunner): if global_size: lra['global_size'] = global_size if local_size and 'local_size' not in lra: lra['local_size'] = local_size et = self.clprg(*rawbufs, *[var_vals[k] for k in self.vars], **lra, wait=wait or DEBUG>=2) - update_stats(self.display_name if self.display_name is not None else self.name, self.op_estimate, self.mem_estimate, var_vals, et, len(rawbufs), jit, lra=lra) + update_stats(self.display_name, self.op_estimate, self.mem_estimate, var_vals, et, len(rawbufs), jit, lra=lra) return et class Compiled: @@ -272,9 +274,8 @@ class Compiled: def to_program(self, k:Linearizer) -> CompiledASTRunner: k.linearize() - src, runtime_args = self.renderer(k.function_name, k.uops) - return CompiledASTRunner(k.ast, k.function_name, src, k.global_size, k.local_size, - display_name=k.display_name, runtime_args=runtime_args).build(self.compiler, self.runtime) + src, runtime_args = self.renderer(to_function_name(k.name), k.uops) + return CompiledASTRunner(k.ast, k.name, src, k.global_size, k.local_size, runtime_args).build(self.compiler, self.runtime) def exec_ast(self, ast:LazyOp, output:LazyBuffer, inputs:Tuple[LazyBuffer, ...], var_vals:Dict[Variable, int], **kwargs): # check if we can reuse the output buffer diff --git a/tinygrad/shape/symbolic.py b/tinygrad/shape/symbolic.py index 67ab9331e6..3530c6aa98 100644 --- a/tinygrad/shape/symbolic.py +++ b/tinygrad/shape/symbolic.py @@ -319,8 +319,6 @@ def create_rednode(typ:Type[RedNode], nodes:List[Node]): elif typ == AndNode: ret.min, ret.max = (min([x.min for x in nodes]), max([x.max for x in nodes])) return create_node(ret) -@functools.lru_cache(maxsize=None) -def sym_rename(s) -> str: return f"s{sym_rename.cache_info().currsize}" def sym_render(a: Union[Node, int], ops=None, ctx=None) -> str: return str(a) if isinstance(a, int) else a.render(ops, ctx) def sym_infer(a: Union[Node, int], var_vals: Dict[Variable, int]) -> int: if isinstance(a, (int, float)): return a