forked from tinygrad/tinygrad
do not use Ops.UNROLL and Ops.CONTRACT [PR] (#16870)
This commit is contained in:
chenyu
@@ -76,27 +76,30 @@ def expand_reduce(r:UOp):
|
||||
out_shape = tuple([1 if i in new_axes else s for i,s in enumerate(r.src[0].shape)])
|
||||
return r.src[0].reduce(*range_srcs, arg=(r.arg[0], tuple(new_axes))).reshape(out_shape)
|
||||
|
||||
def do_contract(ctx:dict[int, int], u:UOp):
|
||||
# the context is a mapping from range number (in contract) to axis number
|
||||
permute_tail = [ctx[rn] for rn,_ in u.arg]
|
||||
permute_head = [i for i in range(len(u.src[0].shape)) if i not in permute_tail]
|
||||
out = u.src[0].permute(permute_head+permute_tail)
|
||||
def contract_axis(ctx:dict[int, int], u:UOp, arg):
|
||||
permute_tail = [ctx[rn] for rn,_ in arg]
|
||||
permute_head = [i for i in range(len(u.shape)) if i not in permute_tail]
|
||||
out = u.permute(permute_head+permute_tail)
|
||||
return out.reshape(*out.shape[:len(permute_head)], -1)
|
||||
|
||||
def do_unroll(ctx:dict[int, int], u:UOp):
|
||||
# this is the opposite of contract
|
||||
permute_tail = [ctx[rn] for rn,_ in u.arg]
|
||||
out = u.src[0].reshape(*u.src[0].shape[:-1], *[nm for _,nm in u.arg])
|
||||
def unroll_axis(ctx:dict[int, int], u:UOp, arg):
|
||||
permute_tail = [ctx[rn] for rn,_ in arg]
|
||||
out = u.reshape(*u.shape[:-1], *[nm for _,nm in arg])
|
||||
permute_head = [i for i in range(len(out.shape)) if i not in permute_tail]
|
||||
return out.permute(argsort(permute_head+permute_tail))
|
||||
|
||||
def expand_wmma(ctx:dict[int, int], u:UOp):
|
||||
if u.tag != 1: return None
|
||||
in0, in1, out0 = u.arg[6]
|
||||
wmma = u.replace(src=(contract_axis(ctx, u.src[0], in0), contract_axis(ctx, u.src[1], in1), u.src[2]), tag=None)
|
||||
return unroll_axis(ctx, wmma, out0)
|
||||
|
||||
expander2 = PatternMatcher([
|
||||
(UPat(Ops.REDUCE, name="r"), expand_reduce),
|
||||
(UPat(Ops.RANGE, name="r"),
|
||||
lambda ctx, r: UOp.const(r.dtype, tuple(range(r.vmax+1))) \
|
||||
.reshape(tuple([r.vmax+1 if i == ctx[r.arg[0]] else 1 for i in range(len(ctx))])) if r.arg[0] in ctx else None),
|
||||
(UPat(Ops.CONTRACT, name="u"), do_contract),
|
||||
(UPat(Ops.UNROLL, name="u"), do_unroll),
|
||||
(UPat(Ops.WMMA, name="u"), expand_wmma),
|
||||
])+pm_flatten_range+mop_cleanup
|
||||
|
||||
def broadcast_binary(x:UOp):
|
||||
|
||||
@@ -302,11 +302,8 @@ class Scheduler:
|
||||
# do the reduce_axes always disappear? i think they don't
|
||||
# they need to be moved into the WMMA srcs
|
||||
wmma_arg = (str(tc), tc.dims, tc.dtype_in, tc.dtype_out, self.ren.target.device, tc.threads, tc_upcast_axes, ()) #, tc_reduce_axes)
|
||||
wmma = UOp(Ops.WMMA, dtype=tc.dtype_out.vec(tc.elements_per_thread[2]), src=(
|
||||
UOp(Ops.CONTRACT, dtype=srcs[0].dtype.vec(tc.elements_per_thread[0]), src=(srcs[0],), arg=tc_upcast_axes[0], tag=1),
|
||||
UOp(Ops.CONTRACT, dtype=srcs[1].dtype.vec(tc.elements_per_thread[1]), src=(srcs[1],), arg=tc_upcast_axes[1], tag=1),
|
||||
UOp.const(tc.dtype_out.vec(tc.elements_per_thread[2]), 0.0)), arg=wmma_arg, tag=1)
|
||||
tc_uop = UOp(Ops.UNROLL, tc.dtype_out, (wmma,), arg=tc_upcast_axes[2], tag=1)
|
||||
tc_uop = UOp(Ops.WMMA, dtype=tc.dtype_out, src=(
|
||||
srcs[0], srcs[1], UOp.const(tc.dtype_out.vec(tc.elements_per_thread[2]), 0.0)), arg=wmma_arg, tag=1)
|
||||
|
||||
# preserve extra reduces
|
||||
reduce_ranges = [x for x in UOp.sink(*reduceop.src[1:]).toposort() if x.op is Ops.RANGE and x.arg[0] not in tc_reduce_axes]
|
||||
|
||||
@@ -140,7 +140,7 @@ def param_to_multi(p:UOp):
|
||||
if p.axis is None: return None
|
||||
return UOp.param(p.arg.slot, p.dtype, p.shard_shape, p.device, p.arg.vmin_vmax, p.arg.name, p.arg.addrspace).multi(p.axis)
|
||||
|
||||
# NOTE: this is the same pattern as Ops.UNROLL
|
||||
# NOTE: this is the same pattern as unrolled ranges
|
||||
multi_pm = PatternMatcher([
|
||||
(UPat(Ops.PARAM, name="p"), param_to_multi),
|
||||
(UPat(GroupOp.ALU, name="root", custom_early_reject=set([Ops.MULTI])), alu_multi),
|
||||
|
||||
@@ -83,7 +83,7 @@ pm_mops = PatternMatcher([
|
||||
(UPat(GroupOp.Movement|{Ops.INDEX}, name="r").after(name="a", allow_any_len=True),
|
||||
lambda r,a: UOp(r.op, r.dtype, (a.replace(src=(r.src[0],)+a.src[1:]),)+r.src[1:], r.arg)),
|
||||
(UPat(GroupOp.Movement, name="r").end(name="a", allow_any_len=True), lambda r,a: a.replace(src=(r.src[0],)+a.src[1:])),
|
||||
# lower SHAPED_WMMA to WMMA with CONTRACT/UNROLL
|
||||
# lower SHAPED_WMMA to WMMA
|
||||
(UPat(Ops.SHAPED_WMMA, name="x"), lower_shaped_wmma),
|
||||
])
|
||||
|
||||
|
||||
+4
-7
@@ -261,8 +261,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
|
||||
return self.src[0].shape
|
||||
else:
|
||||
return (len(self.src),) + self.src[0].shape
|
||||
# TODO: contract and unroll should be deleted
|
||||
case Ops.CONST | Ops.CONTRACT | Ops.UNROLL:
|
||||
case Ops.CONST:
|
||||
return (self.dtype.count,) if self.dtype.count > 1 else ()
|
||||
|
||||
# some ops init the shape
|
||||
@@ -385,9 +384,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
|
||||
def ended_ranges(self) -> tuple[UOp, ...]:
|
||||
if self.op in range_start: return self.src[range_start[self.op]:]
|
||||
if self.op is Ops.AFTER: return tuple(flatten([x.ended_ranges for x in self.src[1:]]))
|
||||
if self.op is Ops.CONTRACT:
|
||||
contract_rng_ids = {rng_id for rng_id, _ in self.arg}
|
||||
return tuple(r for r in self.src[0].ranges if r.op is Ops.RANGE and r.arg[0] in contract_rng_ids)
|
||||
return ()
|
||||
|
||||
# determine what ranges this is in
|
||||
@@ -525,7 +521,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
|
||||
def ins(self, arg, **kwargs): return UOp(Ops.INS, kwargs.pop("dtype", self.dtype), kwargs.pop("src", self.src), arg, kwargs.pop("tag", self.tag))
|
||||
def contract(self, *rngs:UOp):
|
||||
assert all(x.arg[-1] == AxisType.UPCAST for x in rngs), "all contract ranges must be upcast"
|
||||
return UOp(Ops.CONTRACT, dtype=self.dtype.vec(prod([x.vmax+1 for x in rngs])), src=(self,), arg=tuple((x.arg[0], x.vmax+1) for x in rngs))
|
||||
return UOp.vectorize(*[self.substitute(dict(zip(rngs, [r.const_like(i) for r,i in zip(rngs, idx)])))
|
||||
for idx in itertools.product(*[range(int(r.vmax)+1) for r in rngs])])
|
||||
def alu(self, op, *src:UOp, **kwargs):
|
||||
all_srcs = (self, *src)
|
||||
# broadcast shaped operands to a common shape (None and () are falsy, so only real shapes participate)
|
||||
@@ -979,7 +976,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
|
||||
if self.op is Ops.PARAM and self.arg.vmin_vmax is not None: return self.arg.vmin_vmax
|
||||
if self.op in (Ops.RANGE, Ops.SPECIAL): return 0, (self.src[0]-1).vmax
|
||||
if self.op is Ops.BIND: return self.src[0]._min_max # ignore the bound value
|
||||
if self.op in {Ops.UNROLL, Ops.STACK}: return min(x.vmin for x in self.src), max(x.vmax for x in self.src)
|
||||
if self.op is Ops.STACK: return min(x.vmin for x in self.src), max(x.vmax for x in self.src)
|
||||
if self.op is Ops.CONST and self.arg is not Invalid: return self.arg, self.arg
|
||||
if self.op is Ops.INDEX and not isinstance(self.src[0].dtype, PtrDType): return self.src[0]._min_max
|
||||
# TODO: CAST to bool/unsigned is not monotone, still some case can be simplified
|
||||
|
||||
@@ -36,7 +36,6 @@ renderer = PatternMatcher([
|
||||
(UPat((Ops.SPECIAL), name="x"), lambda x: x.arg),
|
||||
(UPat(Ops.RANGE, name="x"), lambda x: f"r{range_str(x)}"),
|
||||
(UPat(Ops.CONST, name="x"), lambda x: str(x.arg)),
|
||||
(UPat(Ops.UNROLL, name="x"), lambda ctx,x,u: f"UNROLL({ctx[x.src[0]]}, {u.arg})"),
|
||||
(UPat(Ops.CAST, name="x"), lambda ctx,x: f"({str(x.dtype)[7:]})({ctx[x.src[0]]})"),
|
||||
(UPat(Ops.BIND, name="x"), lambda ctx,x: ctx[x.src[0]]),
|
||||
(UPat(Ops.NEG, name="x"), lambda ctx,x: f"(-{ctx[x.src[0]]})"),
|
||||
|
||||
@@ -3,7 +3,7 @@ from typing import Any
|
||||
from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, AxisType, KernelInfo, ParamArg
|
||||
from tinygrad.uop.render import print_uops, pyrender
|
||||
from tinygrad.dtype import DType, ImageDType, dtypes, PtrDType, AddrSpace, Invalid, ConstFloat
|
||||
from tinygrad.helpers import DEBUG, Context, prod, SPEC, Metadata, panic, CHECK_OOB, all_same
|
||||
from tinygrad.helpers import DEBUG, Context, SPEC, Metadata, panic, CHECK_OOB, all_same
|
||||
|
||||
# ***** uop helpers *****
|
||||
|
||||
@@ -83,8 +83,7 @@ spec_shared = PatternMatcher([
|
||||
isinstance(x.arg, ParamArg) and x.addrspace in (AddrSpace.REG, AddrSpace.LOCAL)),
|
||||
|
||||
# GROUP of stores (or groups, or NOOPs)
|
||||
# TODO: remove UNROLL here, it's for SPEC=2
|
||||
(UPat(Ops.GROUP, dtypes.void, src=UPat((Ops.GROUP, Ops.STORE, Ops.NOOP, Ops.UNROLL, Ops.INS))), lambda: True),
|
||||
(UPat(Ops.GROUP, dtypes.void, src=UPat((Ops.GROUP, Ops.STORE, Ops.NOOP, Ops.INS))), lambda: True),
|
||||
|
||||
# AFTER on Movement Op, PARAM, BUFFER, CONTIGUOUS, or another AFTER
|
||||
(UPat(Ops.AFTER, src=(UPat(GroupOp.Movement.union({Ops.PARAM, Ops.BUFFER, Ops.CONTIGUOUS, Ops.AFTER, Ops.MULTI, Ops.BITCAST, Ops.INS})),),
|
||||
@@ -182,9 +181,6 @@ spec_tensor = PatternMatcher([
|
||||
(UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.LINEAR), UPat(Ops.SOURCE))), lambda: True),
|
||||
(UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.LINEAR), UPat(Ops.SOURCE), UPat(Ops.BINARY))), lambda: True),
|
||||
|
||||
# UNROLL/CONTRACT is used here for WMMA
|
||||
(UPat(Ops.CONTRACT, name="x"), lambda x: x.dtype.count == prod(y[1] for y in x.arg)),
|
||||
(UPat(Ops.UNROLL, name="x"), lambda x: x.src[0].dtype.count == prod(y[1] for y in x.arg)),
|
||||
])+spec_shared
|
||||
|
||||
# these ops can exist in programs but not the tensor spec. example: LOAD
|
||||
@@ -228,9 +224,6 @@ spec_full = PatternMatcher([
|
||||
# allow any AFTER
|
||||
(UPat(Ops.AFTER, src=(UPat(),), allow_any_len=True), lambda: True),
|
||||
|
||||
# expander: unroll/contract
|
||||
(UPat((Ops.UNROLL, Ops.CONTRACT), src=(UPat(),)), lambda: True),
|
||||
|
||||
# all loads/stores
|
||||
(UPat((Ops.LOAD, Ops.STORE)), lambda: True),
|
||||
|
||||
|
||||
@@ -288,7 +288,7 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
|
||||
((UPat.var("x", dtypes.weakint) + UPat.cvar("c")).cast(dtypes.sints, name="cast"), lambda x,c,cast:x.cast(cast.dtype)+c.cast(cast.dtype)),
|
||||
# only RANGE/IF/STORE/KERNEL have side effects
|
||||
(UPat(Ops.AFTER, name="x"), lambda x: x.replace(src=(x.src[0],)+
|
||||
tuple(dedup(flatten([(y,) if y.op in {Ops.RANGE, Ops.STORE, Ops.CALL, Ops.FUNCTION, Ops.BARRIER, Ops.END, Ops.UNROLL, Ops.LINEAR, Ops.STAGE}
|
||||
tuple(dedup(flatten([(y,) if y.op in {Ops.RANGE, Ops.STORE, Ops.CALL, Ops.FUNCTION, Ops.BARRIER, Ops.END, Ops.LINEAR, Ops.STAGE}
|
||||
else y.src for y in x.src[1:]]))))),
|
||||
# after with 1 src is just src[0]
|
||||
(UPat(Ops.AFTER, src=(UPat.var("s"),)), lambda s: s),
|
||||
@@ -422,7 +422,7 @@ pm_simplify_valid = PatternMatcher([
|
||||
])
|
||||
|
||||
# this is symbolic 2.0
|
||||
REMOVE_FROM_SINK_LIKE = {Ops.UNROLL, Ops.NOOP, Ops.STACK, Ops.SINK, Ops.GROUP}
|
||||
REMOVE_FROM_SINK_LIKE = {Ops.NOOP, Ops.STACK, Ops.SINK, Ops.GROUP}
|
||||
pm_clean_up_group_sink = PatternMatcher([
|
||||
# clean up GROUP/SINK
|
||||
(UPat(Ops.GROUP, src=(UPat.var("x"),)), lambda x: x),
|
||||
|
||||
Reference in New Issue
Block a user