diff --git a/tinygrad/codegen/rangeify.py b/tinygrad/codegen/rangeify.py index 8c931e1192..c3f809d64e 100644 --- a/tinygrad/codegen/rangeify.py +++ b/tinygrad/codegen/rangeify.py @@ -1,4 +1,4 @@ -from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, KernelInfo +from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, KernelInfo, GroupOp from tinygrad.dtype import dtypes from tinygrad.helpers import argsort @@ -61,4 +61,7 @@ pm_rangeify = PatternMatcher([ (UPat(Ops.INDEX, src=(UPat(Ops.EXPAND),), allow_any_len=True, name="x"), map_expand), # TODO: CONST shouldn't have src (UPat(Ops.INDEX, src=(UPat(Ops.CONST, name="c"),)), lambda c: c.replace(src=())), + # move MAP through elementwise ALU + (UPat(Ops.INDEX, src=(UPat(GroupOp.Elementwise.union({Ops.LOAD})),), allow_any_len=True, name="x"), + lambda x: x.src[0].replace(src=tuple([UOp(Ops.INDEX, dtype=s.dtype, src=(s,)+x.src[1:]) for s in x.src[0].src]))), ]) \ No newline at end of file diff --git a/tinygrad/uop/__init__.py b/tinygrad/uop/__init__.py index 2e6304fa17..3be36aff14 100644 --- a/tinygrad/uop/__init__.py +++ b/tinygrad/uop/__init__.py @@ -82,6 +82,7 @@ class GroupOp: Ops.XOR, Ops.SHL, Ops.SHR, Ops.OR, Ops.AND, Ops.THREEFRY, Ops.SUB, Ops.FDIV, Ops.POW} Ternary = {Ops.WHERE, Ops.MULACC} ALU = set.union(Unary, Binary, Ternary) + Elementwise = set.union(ALU, {Ops.CAST, Ops.BITCAST}) Defines = {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG}