diff --git a/tinygrad/codegen/linearizer.py b/tinygrad/codegen/linearizer.py index 6423b58a6d..81480b55ed 100644 --- a/tinygrad/codegen/linearizer.py +++ b/tinygrad/codegen/linearizer.py @@ -96,17 +96,23 @@ class Linearizer(OptimizedKernel): if len(upcast_dim) == 1 and len(expanded_nodes[upcast_dim[0]]) in [4,2]: dim, amt = upcast_dim[0], len(expanded_nodes[upcast_dim[0]]) + # calculate expr_idxs using placeholder variables + fake_idxs = [idx if isinstance(idx, NumNode) else Variable(f"_uidx{i}", idx.min, idx.max) for i, idx in enumerate(idxs)] + g_idx, g_valid = self.sts[i].expr_idxs(fake_idxs) + ret = [] invalid_value = 0 if dtypes.is_int(self.bufs[i].dtype) else 0.0 for _idx in _idxs: + substitute: Dict[VariableOrNum, Node] = {a: b for a, b in zip(fake_idxs, _idx) if isinstance(a, Variable)} if amt > 1: - idx, valid = self.sts[i].expr_idxs((_idx[:dim] + (expanded_nodes[dim][0],) + _idx[dim+1:])) + float4_substitute = {**substitute, fake_idxs[dim]: expanded_nodes[dim][0]} + idx, valid = g_idx.substitute(float4_substitute), g_valid.substitute(float4_substitute) localtype = dtypes._float4 if amt == 4 else dtypes._float2 if idx.render() != ((idx//amt)*amt).render(): - idx, valid = self.sts[i].expr_idxs(_idx) + idx, valid = g_idx.substitute(substitute), g_valid.substitute(substitute) localtype = dtypes.float32 else: - idx, valid = self.sts[i].expr_idxs(_idx) + idx, valid = g_idx.substitute(substitute), g_valid.substitute(substitute) localtype = dtypes.float32 this_const, idx, valid = (invalid_value, Variable.num(0), Variable.num(1)) if valid.max == 0 else (const, idx, valid) key = f"{acc}{localtype}{this_const if this_const is not None and acc is None else self.get_buffer_name(i)}{idx.render()}{valid.render()}" diff --git a/tinygrad/shape/shapetracker.py b/tinygrad/shape/shapetracker.py index 7df6e2253a..7dd95eb55b 100644 --- a/tinygrad/shape/shapetracker.py +++ b/tinygrad/shape/shapetracker.py @@ -48,8 +48,9 @@ class View(ViewInternal): if self.mask is not None: acc = 1 for ns,(x,y) in reversed(list(zip(self.shape, self.mask))): - base = ((idx//acc) % ns) - expr += [base >= x, base < y] + if x != 0 or y != ns: + base = ((idx//acc) % ns) + expr += [base >= x, base < y] acc *= ns return Variable.ands(expr) @@ -172,6 +173,7 @@ class ShapeTracker: def _expr_idx(self, idx, valid) -> Tuple[Node, Node]: for v in reversed(self.views[0:-1]): + if valid.max == 0: return Variable.num(-1), valid valid = v.expr_node_mask(idx, valid) idx = v.expr_node(idx) return idx, valid diff --git a/tinygrad/shape/symbolic.py b/tinygrad/shape/symbolic.py index 153c73ccf0..6e49b22d3a 100644 --- a/tinygrad/shape/symbolic.py +++ b/tinygrad/shape/symbolic.py @@ -23,7 +23,7 @@ class Node(ABC): # expand a Node into List[Node] that enumerates the underlying Variables from min to max def expand(self) -> List[Node]: raise NotImplementedError(self.__class__.__name__) # infer the value of a Node given Variable values in var_vals - def infer(self, var_vals: Dict[Variable, int]) -> int: raise NotImplementedError(self.__class__.__name__) + def substitute(self, var_vals: Dict[Variable, Node]) -> Node: raise NotImplementedError(self.__class__.__name__) @functools.cached_property def key(self) -> str: return self.render(ctx="DEBUG") @functools.cached_property @@ -43,6 +43,7 @@ class Node(ABC): def __gt__(self, b:Union[Node,int]): return (-self) < (-b) def __ge__(self, b:Union[Node,int]): return (-self) < (-b+1) def __lt__(self, b:Union[Node,int]): + #if self.min >= (b.max if isinstance(b, Node) else b): return Variable.num(0) lhs = self if isinstance(lhs, SumNode) and isinstance(b, int): muls, others = partition(lhs.nodes, lambda x: isinstance(x, MulNode) and x.b > 0 and x.max >= b) @@ -153,7 +154,7 @@ class Variable(Node): self.expr, self.min, self.max = expr, nmin, nmax def vars(self): return [self] def expand(self) -> List[Node]: return [self] if self.expr is not None else [Variable.num(j) for j in range(self.min, self.max+1)] - def infer(self, var_vals: Dict[Variable, int]) -> int: return var_vals[self] + def substitute(self, var_vals: Dict[Variable, Node]) -> Node: return var_vals[self] if self in var_vals else self class NumNode(Node): def __init__(self, num:int): @@ -165,7 +166,7 @@ class NumNode(Node): def __eq__(self, other): return self.b == other def __hash__(self): return self.hash # needed with __eq__ override def expand(self) -> List[Node]: return [self] - def infer(self, var_vals: Dict[Variable, int]) -> int: return self.b + def substitute(self, var_vals: Dict[Variable, Node]) -> Node: return self def create_node(ret:Node): assert ret.min <= ret.max, f"min greater than max! {ret.min} {ret.max} when creating {type(ret)} {ret}" @@ -186,6 +187,7 @@ class LtNode(OpNode): def get_bounds(self) -> Tuple[int, int]: if isinstance(self.b, int): return int(self.a.max < self.b), int(self.a.min < self.b) return (1, 1) if self.a.max < self.b.min else (0, 0) if self.a.min > self.b.max else (0, 1) + def substitute(self, var_vals: Dict[Variable, Node]) -> Node: return self.a.substitute(var_vals) < (self.b if isinstance(self.b, int) else self.b.substitute(var_vals)) class MulNode(OpNode): def __mul__(self, b: Union[Node, int]): return self.a*(self.b*b) # two muls in one mul @@ -199,7 +201,7 @@ class MulNode(OpNode): def get_bounds(self) -> Tuple[int, int]: return (self.a.min*self.b, self.a.max*self.b) if self.b >= 0 else (self.a.max*self.b, self.a.min*self.b) def expand(self) -> List[Node]: return [x*self.b for x in self.a.expand()] - def infer(self, var_vals: Dict[Variable, int]) -> int: return self.a.infer(var_vals) * sym_infer(self.b, var_vals) + def substitute(self, var_vals: Dict[Variable, Node]) -> Node: return self.a.substitute(var_vals) * (self.b if isinstance(self.b, int) else self.b.substitute(var_vals)) class DivNode(OpNode): def __floordiv__(self, b: Union[Node, int], _=False): return self.a//(self.b*b) # two divs is one div @@ -207,6 +209,7 @@ class DivNode(OpNode): assert self.a.min >= 0 and isinstance(self.b, int) return self.a.min//self.b, self.a.max//self.b def expand(self) -> List[Node]: return [x//self.b for x in self.a.expand()] + def substitute(self, var_vals: Dict[Variable, Node]) -> Node: return self.a.substitute(var_vals) // (self.b if isinstance(self.b, int) else self.b.substitute(var_vals)) class ModNode(OpNode): def __floordiv__(self, b: Union[Node, int], factoring_allowed=True): @@ -216,13 +219,16 @@ class ModNode(OpNode): assert self.a.min >= 0 and isinstance(self.b, int) return (0, self.b-1) if self.a.max - self.a.min >= self.b or (self.a.min != self.a.max and self.a.min%self.b >= self.a.max%self.b) else (self.a.min%self.b, self.a.max%self.b) def expand(self) -> List[Node]: return [x%self.b for x in self.a.expand()] + def substitute(self, var_vals: Dict[Variable, Node]) -> Node: return self.a.substitute(var_vals) % self.b class RedNode(Node): def __init__(self, nodes:List[Node]): self.nodes = nodes def vars(self): return functools.reduce(lambda l,x: l+x.vars(), self.nodes, []) class SumNode(RedNode): + @functools.lru_cache(maxsize=None) # pylint: disable=method-cache-max-size-none def __mul__(self, b: Union[Node, int]): return Node.sum([x*b for x in self.nodes]) # distribute mul into sum + @functools.lru_cache(maxsize=None) # pylint: disable=method-cache-max-size-none def __floordiv__(self, b: Union[Node, int], factoring_allowed=True): fully_divided: List[Node] = [] rest: List[Node] = [] @@ -255,6 +261,7 @@ class SumNode(RedNode): if divisor > 1: return Node.sum(fully_divided) + Node.sum(rest).__floordiv__(divisor) // (b//divisor) return Node.sum(fully_divided) + Node.__floordiv__(Node.sum(rest), b) + @functools.lru_cache(maxsize=None) # pylint: disable=method-cache-max-size-none def __mod__(self, b: Union[Node, int]): if isinstance(b, SumNode): nu_num = sum(node.b for node in self.flat_components if node.__class__ is NumNode) @@ -279,7 +286,7 @@ class SumNode(RedNode): return Node.__lt__(self, b) def expand(self) -> List[Node]: return [Variable.sum(list(it)) for it in itertools.product(*[x.expand() for x in self.nodes])] - def infer(self, var_vals: Dict[Variable, int]) -> int: return sum([node.infer(var_vals) for node in self.nodes]) + def substitute(self, var_vals: Dict[Variable, Node]) -> Node: return Variable.sum([node.substitute(var_vals) for node in self.nodes]) @property def flat_components(self): # recursively expand sumnode components @@ -290,6 +297,7 @@ class SumNode(RedNode): class AndNode(RedNode): def __mul__(self, b: Union[Node, int]): Variable.ands([x*b for x in self.nodes]) def __floordiv__(self, b: Union[Node, int], _=True): return Variable.ands([x//b for x in self.nodes]) + def substitute(self, var_vals: Dict[Variable, Node]) -> Node: return Variable.ands([node.substitute(var_vals) for node in self.nodes]) def create_rednode(typ:Type[RedNode], nodes:List[Node]): ret = typ(nodes) @@ -300,7 +308,10 @@ def create_rednode(typ:Type[RedNode], nodes:List[Node]): @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: return a if isinstance(a, int) else a.infer(var_vals) +def sym_infer(a: Union[Node, int], var_vals: Dict[Variable, int]) -> int: + ret = (Variable.num(a) if isinstance(a, int) else a).substitute({k:Variable.num(v) for k, v in var_vals.items()}) + assert isinstance(ret, NumNode) + return ret.b render_python: Dict[Type, Callable] = { Variable: lambda self,ops,ctx: f"{self.expr}[{self.min}-{self.max}]" if ctx == "DEBUG" else f"{self.expr}",