mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-09-02 03:46:07 +00:00
symbolic in ops 2 [pr] (#6895)
* move symbolic to ops, simple [pr] * fix for shapetracker
This commit is contained in:
George Hotz
@@ -1,10 +1,10 @@
|
||||
from __future__ import annotations
|
||||
from typing import Optional, Tuple, Dict, List, Set, cast, TYPE_CHECKING, Any, DefaultDict, Callable
|
||||
import functools, itertools, heapq, math, operator
|
||||
import functools, itertools, heapq, operator
|
||||
from collections import defaultdict
|
||||
from tinygrad.dtype import dtypes, PtrDType, ImageDType, ConstType, DType
|
||||
from tinygrad.ops import UnaryOps, BinaryOps, UOp, UOps, END_FOR_UOP, type_verify, print_uops, identity_element
|
||||
from tinygrad.ops import UPat, PatternMatcher, graph_rewrite, TernaryOps, simple_pm
|
||||
from tinygrad.ops import UPat, PatternMatcher, graph_rewrite, TernaryOps, symbolic_flat, is_irreducible, _get_chain
|
||||
from tinygrad.helpers import DEBUG, getenv, flatten, dedup, TRANSCENDENTAL, AMX, prod, CI, partition, all_same
|
||||
from tinygrad.codegen.transcendental import xexp2, xlog2, xsin, TRANSCENDENTAL_SUPPORTED_DTYPES
|
||||
if TYPE_CHECKING: from tinygrad.renderer import Renderer
|
||||
@@ -75,108 +75,8 @@ float4_folding = PatternMatcher([
|
||||
(UPat((UOps.BARRIER, UOps.SINK), src=UPat(UOps.STORE, src=(UPat.var("buf"), UPat(), UPat()), allow_any_len=True), name="ex"), fold_expanded),
|
||||
])
|
||||
|
||||
# ***** mod *****
|
||||
|
||||
def _get_chain(x:UOp, sep:BinaryOps):
|
||||
if x.op is UOps.ALU and x.arg is sep:
|
||||
for s in x.src: yield from _get_chain(s, sep)
|
||||
else: yield x
|
||||
|
||||
def mod_folding(x:UOp, c:int) -> Optional[UOp]:
|
||||
# simplify x % c, None means no change
|
||||
|
||||
# simple cancel mod case
|
||||
if 0 < c and 0 <= x.vmin and (quotient:=x.vmin//c) == x.vmax//c: return x-quotient*c
|
||||
|
||||
remainder, something_changed = [], False
|
||||
for u in _get_chain(x, BinaryOps.ADD):
|
||||
if (factor:=u.const_factor())%c != factor:
|
||||
divides = u.divides(factor)*(factor%c)
|
||||
assert divides is not None
|
||||
remainder.append(divides)
|
||||
something_changed = True
|
||||
elif u.op is UOps.ALU and u.arg is BinaryOps.MOD and (s1:=u.src[1]).op is UOps.CONST and s1.arg%c == 0:
|
||||
remainder.append(u.src[0])
|
||||
something_changed = True
|
||||
else: remainder.append(u)
|
||||
if not something_changed: return None
|
||||
return functools.reduce(operator.add, remainder)%c if remainder else x.const_like(0)
|
||||
|
||||
def div_folding(x:UOp, c:int) -> Optional[UOp]:
|
||||
# simplify x // c, None means no change
|
||||
|
||||
# simple cancel div case
|
||||
if 0 <= x.vmin and x.vmax < c: return x.const_like(0)
|
||||
|
||||
quotient, remainder, rem_const, something_changed, gcd, divisor = [], [], 0, False, c, 1
|
||||
for u in _get_chain(x, BinaryOps.ADD):
|
||||
if u.op is UOps.CONST:
|
||||
# add all const together first
|
||||
if rem_const != 0: something_changed = True
|
||||
rem_const += u.arg
|
||||
elif (factor:=u.const_factor())%c == 0:
|
||||
if factor:
|
||||
divides = u.divides(c)
|
||||
assert divides is not None
|
||||
quotient.append(divides)
|
||||
something_changed = True
|
||||
else:
|
||||
# divisor is the smallest common divisor of all MULs
|
||||
if u.op is UOps.ALU and u.arg is BinaryOps.MUL and factor > 1 and c % factor == 0 and (divisor == 1 or divisor > factor): divisor = factor
|
||||
remainder.append(u)
|
||||
gcd = math.gcd(gcd, factor)
|
||||
|
||||
# handle the const
|
||||
if rem_const%c != rem_const:
|
||||
something_changed = True
|
||||
quotient.append(x.const_like(rem_const//c))
|
||||
rem_const = rem_const%c
|
||||
if rem_const != 0: remainder.append(x.const_like(rem_const))
|
||||
|
||||
# x // c -> quotient + (remainder // div) // (c // div)
|
||||
div = gcd if gcd > 1 else divisor
|
||||
|
||||
if not something_changed: return newx//(c//div) if 1 < div < c and (newx:=div_folding(x, div)) is not None else None
|
||||
rem:Optional[UOp] = functools.reduce(operator.add, remainder) if remainder else None
|
||||
quo:Optional[UOp] = functools.reduce(operator.add, quotient) if quotient else None
|
||||
if quo is None: return x.const_like(0) if rem is None else cast(UOp, div_folding(rem, div))//(c//div)
|
||||
return quo if rem is None else cast(UOp, div_folding(rem, div))//(c//div)+quo
|
||||
|
||||
def lt_folding(x:UOp, c:int) -> Optional[UOp]:
|
||||
return cast(UOp, x.divides(g)).lt(c//g) if ((g:=math.gcd(x.const_factor(), c)) > 1) else None
|
||||
|
||||
def fold_unrolled_divs(divs:UOp):
|
||||
# div pattern in unrolled arange
|
||||
# example: (x//4+(x+1)//4+(x+2)//4+(x+3)//4 -> x
|
||||
add_chain, seen_const, ans = list(_get_chain(divs, BinaryOps.ADD)), [], None
|
||||
for u in add_chain:
|
||||
if not (u.op is UOps.ALU and u.arg is BinaryOps.IDIV and u.src[1].op is UOps.CONST and u.src[1].arg==len(add_chain)): return None
|
||||
# assumed CONST is the last of an ADD
|
||||
if (s0:=u.src[0]).op is UOps.ALU and s0.arg is BinaryOps.ADD and s0.src[1].op is UOps.CONST and s0.src[1].op is UOps.CONST:
|
||||
seen_const.append(s0.src[1].arg)
|
||||
s0 = s0.src[0]
|
||||
else: seen_const.append(0)
|
||||
if ans is None: ans = s0
|
||||
if ans is not s0: return None
|
||||
return ans if ans is not None and sorted(seen_const)==list(range(len(add_chain))) else None
|
||||
|
||||
# ***** image load valid simplification *****
|
||||
|
||||
def is_irreducible(u:UOp): return u.op in (UOps.DEFINE_VAR, UOps.SPECIAL, UOps.RANGE)
|
||||
|
||||
def canonicalize_simplex(X:UOp) -> Optional[UOp]:
|
||||
# (X := a0*x0 + a1*x1 + ...) > 0 is equivalent to x0 + x1 + ... > 0 if xi >= 0 and ai > 0 for ints.
|
||||
# returns x0 + x1 + ... in such case, or None if not
|
||||
changed, ret = False, []
|
||||
for u in _get_chain(X, BinaryOps.ADD):
|
||||
# assumed the const is the last src of MUL
|
||||
if u.op is UOps.ALU and u.arg is BinaryOps.MUL and u.src[1].op is UOps.CONST and u.src[1].arg > 0:
|
||||
changed = True
|
||||
u = u.src[0]
|
||||
if not (is_irreducible(u) and u.vmin >= 0): return None
|
||||
ret.append(u)
|
||||
return functools.reduce(operator.add, ret) if changed else None
|
||||
|
||||
def is_increasing(f:UOp):
|
||||
# is f a monotonically increasing function regards its input
|
||||
if f.op is UOps.CONST or is_irreducible(f): return True
|
||||
@@ -370,7 +270,7 @@ def no_vectorized_wmma(wmma:UOp):
|
||||
return UOp(UOps.VECTORIZE, wmma.dtype, tuple(wmma_ex))
|
||||
|
||||
# this is symbolic 2.0
|
||||
sym = simple_pm+PatternMatcher([
|
||||
sym = symbolic_flat+PatternMatcher([
|
||||
# self ASSIGN is just self
|
||||
(UPat(UOps.ASSIGN, src=(UPat.var('x'), UPat.var('x'))), lambda x: x),
|
||||
# ASSIGN to global is just self
|
||||
@@ -419,8 +319,6 @@ sym = simple_pm+PatternMatcher([
|
||||
.lt(UPat.cvar("compval")).ne(UPat(UOps.CONST, name="ne", arg=True))
|
||||
.where(UPat.cvar("multconst"), UPat.const(None, 0)), m2 + UPat.var("extra")),),
|
||||
arg=BinaryOps.ADD, name="reduce", allow_any_len=True), loop_collapse),
|
||||
# unrolled arange div folding
|
||||
(UPat(UOps.ALU, name="divs", src=[UPat(), UPat(UOps.ALU, arg=BinaryOps.IDIV)], arg=BinaryOps.ADD), fold_unrolled_divs),
|
||||
# indexing, with cast or where
|
||||
(UPat(UOps.REDUCE, src=(UPat.var("idx").eq(UPat(UOps.RANGE, name="rng")).cast()*
|
||||
UPat(UOps.LOAD, src=(UPat.var("buf"), UPat.any(UPat.var("add")+UPat.var("mul")*UPat(UOps.RANGE, name="rng"), UPat(UOps.RANGE, name="rng"))),
|
||||
@@ -430,28 +328,12 @@ sym = simple_pm+PatternMatcher([
|
||||
name="ld"), UPat.const(None, 0.0)),), arg=BinaryOps.ADD, name="reduce", allow_any_len=True), index_collapse),
|
||||
# GEP/CAST const rules
|
||||
(UPat(UOps.CAST, name="root", src=UPat.cvar("c")), lambda root, c: root.const_like(c.arg)),
|
||||
# ** combine terms (opinionated) **
|
||||
(-1 * (UPat.var("x") + UPat.var("y")), lambda x,y: (-x)+(-y)), # -(x+y) -> -x + -y
|
||||
# (x+y)*c -> x*c+y*c. only for int, float has inf*0=nan issue
|
||||
((UPat.var("x", dtypes.ints) + UPat.var("y")) * UPat.cvar("c"), lambda x,y,c: x*c+y*c),
|
||||
# ** self folding **
|
||||
# cast NOOP (NOTE: it's str to deal with PtrDType)
|
||||
(UPat(UOps.CAST, name="root"), lambda root: root.src[0] if str(root.dtype) == str(root.src[0].dtype) else None),
|
||||
(UPat(UOps.REDUCE, src=(UPat.var("x"),)), lambda x: x), # a REDUCE without ranges is a NOOP
|
||||
# ** load/store folding **
|
||||
(UPat.store(UPat.var("buf"), UPat.var("idx"), UPat.load(UPat.var("buf"), UPat.var("idx"))), lambda buf,idx:UOp(UOps.NOOP)),
|
||||
# *** rules from symbolic ***
|
||||
# generic lt folding
|
||||
(UPat.var("x", dtypes.sints).lt(UPat.cvar("c", vec=False)), lambda x,c: lt_folding(x, c.arg) if 0 < c.arg else None),
|
||||
# canonicalize a simplex with positive coefficients > 0
|
||||
# not x < 1 -> X > 0
|
||||
(UPat.var("x", dtypes.ints).lt(1).ne(True), lambda x: newx.lt(1).ne(True) if (newx:=canonicalize_simplex(x)) is not None else None),
|
||||
# ** div **
|
||||
# # div folding
|
||||
(UPat.var("x", dtypes.sints) // UPat.cvar("c", vec=False), lambda x,c: newx if 0 < c.arg and (newx:=div_folding(x,c.arg)) is not None else None),
|
||||
# ** mod **
|
||||
# mod folding
|
||||
(UPat.var("x") % UPat.cvar("c", vec=False), lambda x,c: newx if 0 < c.arg and (newx:=mod_folding(x,c.arg)) is not None else None),
|
||||
# x!=0 -> (bool)x
|
||||
(UPat.var("x").ne(0), lambda x: x.cast(dtypes.bool.vec(x.dtype.count))),
|
||||
# TODO: can do the invert of this (flip alt/load) when we fix double ops
|
||||
|
||||
+121
-2
@@ -209,7 +209,7 @@ class UOp(MathTrait):
|
||||
def __repr__(self): return pretty_print(self, lambda x: f"{type(self).__name__}({x.op}, {x.dtype}, arg={x.argstr()}, src=(%s))")
|
||||
def argstr(self): return f'({", ".join(map(str, self.arg))})' if self.op is UOps.REDUCE_AXIS else self.arg
|
||||
# *** uop evaluation ***
|
||||
def simplify(self): return graph_rewrite(self, simple_pm)
|
||||
def simplify(self): return graph_rewrite(self, symbolic)
|
||||
def ssimplify(self) -> Union[UOp, ConstType]: return ret.arg if (ret:=self.simplify()).op is UOps.CONST else ret
|
||||
def _eval(self, dtype, expected_type) -> ConstType:
|
||||
assert self.dtype in dtype, f"eval with wrong dtype {self}"
|
||||
@@ -745,7 +745,105 @@ def type_verify(uops:List[UOp]):
|
||||
|
||||
# *** most of symbolic lives here now ***
|
||||
|
||||
simple_pm = PatternMatcher([
|
||||
def _get_chain(x:UOp, sep:BinaryOps):
|
||||
if x.op is UOps.ALU and x.arg is sep:
|
||||
for s in x.src: yield from _get_chain(s, sep)
|
||||
else: yield x
|
||||
|
||||
def mod_folding(x:UOp, c:int) -> Optional[UOp]:
|
||||
# simplify x % c, None means no change
|
||||
|
||||
# simple cancel mod case
|
||||
if 0 < c and 0 <= x.vmin and (quotient:=x.vmin//c) == x.vmax//c: return x-quotient*c
|
||||
|
||||
remainder, something_changed = [], False
|
||||
for u in _get_chain(x, BinaryOps.ADD):
|
||||
if (factor:=u.const_factor())%c != factor:
|
||||
divides = u.divides(factor)*(factor%c)
|
||||
assert divides is not None
|
||||
remainder.append(divides)
|
||||
something_changed = True
|
||||
elif u.op is UOps.ALU and u.arg is BinaryOps.MOD and (s1:=u.src[1]).op is UOps.CONST and s1.arg%c == 0:
|
||||
remainder.append(u.src[0])
|
||||
something_changed = True
|
||||
else: remainder.append(u)
|
||||
if not something_changed: return None
|
||||
return functools.reduce(operator.add, remainder)%c if remainder else x.const_like(0)
|
||||
|
||||
def div_folding(x:UOp, c:int) -> Optional[UOp]:
|
||||
# simplify x // c, None means no change
|
||||
|
||||
# simple cancel div case
|
||||
if 0 <= x.vmin and x.vmax < c: return x.const_like(0)
|
||||
|
||||
quotient, remainder, rem_const, something_changed, gcd, divisor = [], [], 0, False, c, 1
|
||||
for u in _get_chain(x, BinaryOps.ADD):
|
||||
if u.op is UOps.CONST:
|
||||
# add all const together first
|
||||
if rem_const != 0: something_changed = True
|
||||
rem_const += u.arg
|
||||
elif (factor:=u.const_factor())%c == 0:
|
||||
if factor:
|
||||
divides = u.divides(c)
|
||||
assert divides is not None
|
||||
quotient.append(divides)
|
||||
something_changed = True
|
||||
else:
|
||||
# divisor is the smallest common divisor of all MULs
|
||||
if u.op is UOps.ALU and u.arg is BinaryOps.MUL and factor > 1 and c % factor == 0 and (divisor == 1 or divisor > factor): divisor = factor
|
||||
remainder.append(u)
|
||||
gcd = math.gcd(gcd, factor)
|
||||
|
||||
# handle the const
|
||||
if rem_const%c != rem_const:
|
||||
something_changed = True
|
||||
quotient.append(x.const_like(rem_const//c))
|
||||
rem_const = rem_const%c
|
||||
if rem_const != 0: remainder.append(x.const_like(rem_const))
|
||||
|
||||
# x // c -> quotient + (remainder // div) // (c // div)
|
||||
div = gcd if gcd > 1 else divisor
|
||||
|
||||
if not something_changed: return newx//(c//div) if 1 < div < c and (newx:=div_folding(x, div)) is not None else None
|
||||
rem:Optional[UOp] = functools.reduce(operator.add, remainder) if remainder else None
|
||||
quo:Optional[UOp] = functools.reduce(operator.add, quotient) if quotient else None
|
||||
if quo is None: return x.const_like(0) if rem is None else cast(UOp, div_folding(rem, div))//(c//div)
|
||||
return quo if rem is None else cast(UOp, div_folding(rem, div))//(c//div)+quo
|
||||
|
||||
def lt_folding(x:UOp, c:int) -> Optional[UOp]:
|
||||
return cast(UOp, x.divides(g)).lt(c//g) if ((g:=math.gcd(x.const_factor(), c)) > 1) else None
|
||||
|
||||
def fold_unrolled_divs(divs:UOp):
|
||||
# div pattern in unrolled arange
|
||||
# example: (x//4+(x+1)//4+(x+2)//4+(x+3)//4 -> x
|
||||
add_chain, seen_const, ans = list(_get_chain(divs, BinaryOps.ADD)), [], None
|
||||
for u in add_chain:
|
||||
if not (u.op is UOps.ALU and u.arg is BinaryOps.IDIV and u.src[1].op is UOps.CONST and u.src[1].arg==len(add_chain)): return None
|
||||
# assumed CONST is the last of an ADD
|
||||
if (s0:=u.src[0]).op is UOps.ALU and s0.arg is BinaryOps.ADD and s0.src[1].op is UOps.CONST and s0.src[1].op is UOps.CONST:
|
||||
seen_const.append(s0.src[1].arg)
|
||||
s0 = s0.src[0]
|
||||
else: seen_const.append(0)
|
||||
if ans is None: ans = s0
|
||||
if ans is not s0: return None
|
||||
return ans if ans is not None and sorted(seen_const)==list(range(len(add_chain))) else None
|
||||
|
||||
def is_irreducible(u:UOp): return u.op in (UOps.DEFINE_VAR, UOps.SPECIAL, UOps.RANGE)
|
||||
|
||||
def canonicalize_simplex(X:UOp) -> Optional[UOp]:
|
||||
# (X := a0*x0 + a1*x1 + ...) > 0 is equivalent to x0 + x1 + ... > 0 if xi >= 0 and ai > 0 for ints.
|
||||
# returns x0 + x1 + ... in such case, or None if not
|
||||
changed, ret = False, []
|
||||
for u in _get_chain(X, BinaryOps.ADD):
|
||||
# assumed the const is the last src of MUL
|
||||
if u.op is UOps.ALU and u.arg is BinaryOps.MUL and u.src[1].op is UOps.CONST and u.src[1].arg > 0:
|
||||
changed = True
|
||||
u = u.src[0]
|
||||
if not (is_irreducible(u) and u.vmin >= 0): return None
|
||||
ret.append(u)
|
||||
return functools.reduce(operator.add, ret) if changed else None
|
||||
|
||||
symbolic = PatternMatcher([
|
||||
# bool MUL is AND, ADD/MAX is OR. prevents other rules to rewrite bool ADD/MUL incorrectly
|
||||
(UPat.var('x', dtype=dtypes.bool) * UPat.var('y'), lambda x,y: x&y),
|
||||
(UPat.var('x', dtype=dtypes.bool) + UPat.var('y'), lambda x,y: x|y),
|
||||
@@ -816,6 +914,27 @@ simple_pm = PatternMatcher([
|
||||
# ** move mul consts to end (NOTE: this is still happening before constant folding) **
|
||||
(UPat(UOps.ALU, arg=BinaryOps.MUL, src=(UPat.cvar("c1"), UPat.var("x"))), lambda c1,x: x*c1 if x.op not in (UOps.CONST, UOps.VCONST) else None),
|
||||
(UPat(UOps.ALU, arg=BinaryOps.MUL, src=(UPat.var("x"), UPat.cvar("c1"))) * UPat.var("y"), lambda x,c1,y: (x*y)*c1),
|
||||
# *** rules from symbolic ***
|
||||
# unrolled arange div folding
|
||||
(UPat(UOps.ALU, name="divs", src=[UPat(), UPat(UOps.ALU, arg=BinaryOps.IDIV)], arg=BinaryOps.ADD), fold_unrolled_divs),
|
||||
# generic lt folding
|
||||
(UPat.var("x", dtypes.sints).lt(UPat.cvar("c", vec=False)), lambda x,c: lt_folding(x, c.arg) if 0 < c.arg else None),
|
||||
# canonicalize a simplex with positive coefficients > 0
|
||||
# not x < 1 -> X > 0
|
||||
(UPat.var("x", dtypes.ints).lt(1).ne(True), lambda x: newx.lt(1).ne(True) if (newx:=canonicalize_simplex(x)) is not None else None),
|
||||
# ** div **
|
||||
# # div folding
|
||||
(UPat.var("x", dtypes.sints) // UPat.cvar("c", vec=False), lambda x,c: newx if 0 < c.arg and (newx:=div_folding(x,c.arg)) is not None else None),
|
||||
# ** mod **
|
||||
# mod folding
|
||||
(UPat.var("x") % UPat.cvar("c", vec=False), lambda x,c: newx if 0 < c.arg and (newx:=mod_folding(x,c.arg)) is not None else None),
|
||||
])
|
||||
|
||||
symbolic_flat = symbolic+PatternMatcher([
|
||||
# ** combine terms (opinionated) **
|
||||
(-1 * (UPat.var("x") + UPat.var("y")), lambda x,y: (-x)+(-y)), # -(x+y) -> -x + -y
|
||||
# (x+y)*c -> x*c+y*c. only for int, float has inf*0=nan issue
|
||||
((UPat.var("x", dtypes.ints) + UPat.var("y")) * UPat.cvar("c"), lambda x,y,c: x*c+y*c),
|
||||
])
|
||||
|
||||
# for debug
|
||||
|
||||
@@ -6,8 +6,7 @@ from tinygrad.helpers import merge_dicts, getenv
|
||||
from tinygrad.shape.symbolic import Variable, sint
|
||||
from tinygrad.shape.view import View, strides_for_shape
|
||||
from tinygrad.dtype import dtypes
|
||||
from tinygrad.ops import UOp, UOps, BinaryOps, graph_rewrite, resolve
|
||||
from tinygrad.codegen.uopgraph import sym, _get_chain
|
||||
from tinygrad.ops import UOp, UOps, BinaryOps, graph_rewrite, resolve, _get_chain, symbolic_flat
|
||||
|
||||
def variable_to_uop(x, ctx=None) -> UOp: return UOp.const(dtypes.pyint, x) if isinstance(x, int) else x
|
||||
def _uop_view(view:View, idxs:List[UOp], vexpr:UOp) -> Tuple[UOp, UOp]:
|
||||
@@ -90,15 +89,15 @@ class ShapeTracker:
|
||||
if len(self.views) == 1 and self.views[-1].mask is None: return self.views[-1].strides
|
||||
ret: List[Optional[sint]] = [None] * len(self.shape)
|
||||
idx, valid = self.to_indexed_uops()
|
||||
idx = graph_rewrite(idx, pm=sym)
|
||||
idx = graph_rewrite(idx, symbolic_flat)
|
||||
for c in _get_chain(idx, BinaryOps.ADD):
|
||||
if c.op is UOps.RANGE: ret[c.arg] = 1
|
||||
if c.op is UOps.ALU and c.arg is BinaryOps.MUL and c.src[0].op is UOps.RANGE and c.src[1].op is UOps.CONST: ret[c.src[0].arg] = c.src[1].arg
|
||||
if c.op is UOps.ALU and c.arg is BinaryOps.MUL and c.src[1].op is UOps.RANGE and c.src[0].op is UOps.CONST: ret[c.src[1].arg] = c.src[0].arg
|
||||
used_ranges = [x.arg for x in graph_rewrite(idx, pm=sym).sparents if x.op is UOps.RANGE]
|
||||
used_ranges = [x.arg for x in graph_rewrite(idx, symbolic_flat).sparents if x.op is UOps.RANGE]
|
||||
ret = [x if i in used_ranges else 0 for i,x in enumerate(ret)]
|
||||
if not ignore_valid:
|
||||
masked_axis = [x.arg for x in graph_rewrite(valid, pm=sym).sparents if x.op is UOps.RANGE]
|
||||
masked_axis = [x.arg for x in graph_rewrite(valid, symbolic_flat).sparents if x.op is UOps.RANGE]
|
||||
ret = [None if i in masked_axis else x for i,x in enumerate(ret)]
|
||||
return tuple(ret)
|
||||
|
||||
@@ -106,7 +105,7 @@ class ShapeTracker:
|
||||
|
||||
def axis_is_masked(self, axis:int) -> bool:
|
||||
_, valid = self.to_indexed_uops()
|
||||
return axis in [x.arg for x in graph_rewrite(valid, sym).sparents if x.op is UOps.RANGE]
|
||||
return axis in [x.arg for x in graph_rewrite(valid, symbolic_flat).sparents if x.op is UOps.RANGE]
|
||||
|
||||
def simplify(self) -> ShapeTracker:
|
||||
if len(self.views) >= 2 and (new_view := self.views[-2] + self.views[-1]) is not None:
|
||||
|
||||
Reference in New Issue
Block a user