diff --git a/tinygrad/codegen/gpudims.py b/tinygrad/codegen/gpudims.py index 8be324a6e6..fd4e1d844e 100644 --- a/tinygrad/codegen/gpudims.py +++ b/tinygrad/codegen/gpudims.py @@ -91,7 +91,8 @@ def add_gpudims(ctx:Renderer, s:UOp): def fix_reduce_unroll(x:UOp): reduce_range, reduce_expand = partition(x.src[1:], lambda y: y.op is Ops.RANGE) if len(reduce_expand) == 0: return None - assert all(x.op is Ops.UNROLL for x in reduce_expand), f"not all UNROLLS in {reduce_expand} for {x.axis_arg}" + reduce_expand = [x for x in reduce_expand if x.op is not Ops.CONST] + assert all(x.op is Ops.UNROLL for x in reduce_expand), f"not all UNROLLS in {reduce_expand}" ret = x.src[0] if len(contract_axis:=flatten(x.arg for x in reduce_expand)): ret = UOp(Ops.CONTRACT, x.dtype.vec(prod(x[1] for x in contract_axis)), (ret,), tuple(contract_axis), tag=1) diff --git a/tinygrad/codegen/opt/__init__.py b/tinygrad/codegen/opt/__init__.py index 54d9a2869d..aa4c0c1c83 100644 --- a/tinygrad/codegen/opt/__init__.py +++ b/tinygrad/codegen/opt/__init__.py @@ -39,13 +39,22 @@ def apply_opt(ast:UOp, renderer:Renderer, cls:type[Kernel]): k = cls(ast, opts=renderer) k.apply_opts(ast.arg.opts_to_apply) ret = k.get_optimized_ast() - if __debug__: type_verify(list(ret.toposort())) + if __debug__ and cls == Kernel: type_verify(list(ret.toposort())) return ret pm_do_optimize = PatternMatcher([ (UPat(Ops.SINK, name="ast"), lambda ctx,ast: apply_opt(ast, ctx, Kernel) if ast.arg is not None and ast.arg.opts_to_apply is not None else None), ]) +def flatten_range(r:UOp): + off = 2 if r.op is Ops.STORE else 1 + rngs = r.src[off:] + if not len(rngs): return None + new_rngs = [x for x in UOp.sink(*rngs).toposort() if x.op is Ops.RANGE] + return r.replace(src=r.src[:off]+tuple(new_rngs)) + pm_postrange_opt = PatternMatcher([ (UPat(Ops.SINK, name="ast"), lambda ctx,ast: apply_opt(ast, ctx, RKernel) if ast.arg is not None and ast.arg.opts_to_apply is not None else None), + # real ranges only + (UPat((Ops.REDUCE, Ops.STORE), name="r"), flatten_range), ]) diff --git a/tinygrad/codegen/opt/kernel.py b/tinygrad/codegen/opt/kernel.py index 507039996a..d3d704f9c1 100644 --- a/tinygrad/codegen/opt/kernel.py +++ b/tinygrad/codegen/opt/kernel.py @@ -86,7 +86,7 @@ class Kernel: # group simplifies self.simplify_ones() - self.simplify_merge_adjacent() + #self.simplify_merge_adjacent() # axis types global_loops = AxisType.GLOBAL if self.opts.has_local else AxisType.LOOP @@ -122,7 +122,7 @@ class Kernel: @property def output_shape(self) -> tuple[sint, ...]: return self.sts[0].shape @property - def shape_len(self) -> int: return len(self.sts[0].shape) + def shape_len(self) -> int: return len(self.output_shape) def axes_of(self, *axis_type:AxisType) -> list[int]: return [i for i,t in enumerate(self.axis_types) if t in argfix(axis_type)] @property diff --git a/tinygrad/codegen/opt/postrange.py b/tinygrad/codegen/opt/postrange.py index a5e87b4e25..ab277e2982 100644 --- a/tinygrad/codegen/opt/postrange.py +++ b/tinygrad/codegen/opt/postrange.py @@ -1,7 +1,52 @@ -from tinygrad.uop.ops import UOp, Ops +from dataclasses import replace +from tinygrad.uop.ops import UOp, Ops, sint, ssimplify, AxisType from tinygrad.codegen.opt.kernel import Kernel from tinygrad.renderer import Renderer +from tinygrad.dtype import dtypes class RKernel(Kernel): def __init__(self, ast:UOp, opts:Renderer|None=None): + self.rng = sorted([u for u in ast.toposort() if u.op is Ops.RANGE and u.vmax > 0], key=lambda x: x.arg) super().__init__(ast, opts) + self.sts.clear() + + # convert LOOP to GLOBAL + self.replaces = {} + if self.opts.has_local: + rng = [x.replace(arg=(x.arg[0], AxisType.GLOBAL)) if x.arg[1] == AxisType.LOOP else x for x in self.rng] + self.replaces.update(dict(zip(self.rng, rng))) + self.rng = rng + + def shift_to(self, axis:int, amount:int, new_type:AxisType, top:bool=False, insert_at:int|None=None): + old_sz = self.rng[axis].src[0].arg // amount + assert old_sz > 0, f"bad old_sz on {axis} {amount} {self.rng[axis]}" + + maxarg = max([x.arg[0] for x in self.rng]) + new_rng = UOp.range(dtypes.int, amount, maxarg+1, new_type) + + if old_sz == 1: + self.replaces[self.rng[axis]] = new_rng + self.rng.insert(insert_at if insert_at is not None else len(self.rng), new_rng) + del self.rng[axis] + else: + replaced_rng = self.rng[axis].replace(src=(UOp.const(dtypes.int, old_sz),)) + self.replaces[self.rng[axis]] = replaced_rng * amount + new_rng + self.rng[axis] = replaced_rng + self.rng.insert(insert_at if insert_at is not None else len(self.rng), new_rng) + + @property + def axis_types(self) -> list[AxisType]: return [x.arg[1] for x in self.rng] + @property + def shape_len(self): return len(self.rng) + + @property + def full_shape(self) -> tuple[sint, ...]: return tuple([ssimplify(x.src[0]) for x in self.rng]) + @property + def output_shape(self) -> tuple[sint, ...]: return tuple([ssimplify(x.src[0]) for x in self.ast.src[0].src[2:]]) + + def get_optimized_ast(self, name_override:str|None=None) -> UOp: + return self.ast.substitute(self.replaces).replace(arg=replace(self.ast.arg, name=self.name, opts_to_apply=None)) + + # does nothing + @axis_types.setter + def axis_types(self, value): pass