sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

remove the .gep method (#16867)

Browse Source
This commit is contained in:
George Hotz
2026-07-04 14:07:09 -07:00
committed by GitHub
parent 971800b46a
commit a145fdce4b
18 changed files with 39 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/gemm/metal_uop_matmul.py
+1 -1
View File
@@ -33,7 +33,7 @@ def hand_spec_tc_cores():
acc_load = UOp.vectorize(acc.after(gk)[0], acc.after(gk)[1])
out = UOp(Ops.WMMA, dtypes.float.vec(2), (a_tc, b_tc, acc_load), arg=wmma_arg)
end_loop = UOp.group(*[acc[i].store(out.gep(i)) for i in range(2)]).end(gk)
end_loop = UOp.group(*[acc[i].store(out.index(i)) for i in range(2)]).end(gk)
sink = UOp.group(*[mat_idx(c.after(end_loop), gx, gy, warp, i).store(acc[i]) for i in range(2)])
return sink.sink(arg=KernelInfo(name="custom_metal_matmul", opts_to_apply=())).simplify()
extra/gemm/mi350x_uop_matmul.py
+2 -2
View File
@@ -180,7 +180,7 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
# store the acc into gmem
cp_i, cp_j = UOp.range(BLOCK_M//TC_M//WARPGROUP_SIZE, 10004), UOp.range(BLOCK_N//TC_N, 10005)
c_load = lambda i: C[gx, cp_i*TC_M*WARPGROUP_SIZE + warpgroup*TC_M + (warp//16)*4+i, gy, cp_j*TC_N + warp%16]
store = UOp.group(*[c_load(i).store(acc[cp_j, cp_i].gep(i)) for i in range(4)])
store = UOp.group(*[c_load(i).store(acc[cp_j, cp_i].index(i)) for i in range(4)])
store = store.end(cp_i, cp_j)
return store.sink(arg=KernelInfo(name="custom_gemm", opts_to_apply=())).simplify()
@@ -197,7 +197,7 @@ wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float,
out = UOp(Ops.WMMA, dtypes.float.vec(4), (A_in, B_in, acc_load), arg=wmma_arg)
# store back the acc
acc = acc.after(UOp.group(*[acc[i].store(out.gep(i)) for i in range(4)]).end(K_loop))
acc = acc.after(UOp.group(*[acc[i].store(out.index(i)) for i in range(4)]).end(K_loop))
# store the acc into gmem
store = UOp.group(*[C[gx, (warp//16)*4+i, gy, warp%16].store(acc[i]) for i in range(4)])
extra/thunder/tiny/tk/group.py
+4 -4
View File
@@ -93,7 +93,7 @@ class Group:
d_in = UOp.vectorize(*[c[height, width, i] for i in range(4)])
out = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in, d_in), arg=wmma_arg)
c_i = [c[height, width, i].store(out.gep(i)) for i in range(4)]
c_i = [c[height, width, i].store(out.index(i)) for i in range(4)]
c_store = UOp.group(*c_i).end(height, width, inner)
self.ker.push_store(c_store, c)
@@ -123,7 +123,7 @@ class Group:
d_in = UOp.vectorize(*[c[height, width, i] for i in range(4)])
out = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in, d_in), arg=wmma_arg)
c_i = [c[height, width, i].store(out.gep(i)) for i in range(4)]
c_i = [c[height, width, i].store(out.index(i)) for i in range(4)]
c_store = UOp.group(*c_i).end(height, width, inner)
self.ker.push_store(c_store, c)
@@ -153,7 +153,7 @@ class Group:
d_in = UOp.vectorize(*[c[height, width, i] for i in range(4)])
out = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in, d_in), arg=wmma_arg)
c_i = [c[height, width, i].store(out.gep(i)) for i in range(4)]
c_i = [c[height, width, i].store(out.index(i)) for i in range(4)]
c_store = UOp.group(*c_i).end(height, width, inner)
self.ker.push_store(c_store, c)
@@ -183,7 +183,7 @@ class Group:
d_in = UOp.vectorize(*[c[height, width, i] for i in range(4)])
out = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in, d_in), arg=wmma_arg)
c_i = [c[height, width, i].store(out.gep(i)) for i in range(4)]
c_i = [c[height, width, i].store(out.index(i)) for i in range(4)]
c_store = UOp.group(*c_i).end(height, width, inner)
self.ker.push_store(c_store, c)
test/backend/test_custom_kernel.py
+1 -1
View File
@@ -51,7 +51,7 @@ def flip_contract_kernel(dest:UOp, src:UOp):
i = UOp.range(dest.shape[0], 0)
j = UOp.range(dest.shape[1], 1, AxisType.UPCAST)
vec = src[i, j].contract(j)
store = UOp.group(*[dest[i, k].store(vec.gep(3-k)) for k in range(4)])
store = UOp.group(*[dest[i, k].store(vec.index(3-k)) for k in range(4)])
return store.end(i, j).sink(arg=KernelInfo(name=f"flip_contract_{dest.numel()}", opts_to_apply=()))
def slice_sum_kernel(dest:UOp, src:UOp):
test/null/test_graph_rewrite.py
+5 -11
View File
@@ -183,33 +183,27 @@ class TestGEPAndVectorizeRewrite(unittest.TestCase):
def test_gep_single_element_extraction(self):
# GEP on a vector dtype to extract a single element
base_vector = UOp.const(dtypes.float32.vec(4), (1.0, 2.0, 3.0, 4.0))
self.assertEqual(apply_rewrite(base_vector.gep(2)).arg, 3.0)
self.assertEqual(apply_rewrite(base_vector.index(2)).arg, 3.0)
def test_gep_tuple_extraction(self):
# GEP on a vector dtype to extract multiple elements as a vector
base_vector = UOp.const(dtypes.float32.vec(4), (1.0, 2.0, 3.0, 4.0))
self.assertEqual(list(apply_rewrite_values(base_vector.gep((2, 3)))), [3.0, 4.0])
self.assertEqual(list(apply_rewrite_values(UOp.vectorize(*[base_vector.index(i) for i in (2, 3)]))), [3.0, 4.0])
def test_gep_on_const_stack(self):
# GEP on a const STACK to extract a single element
const_stack = UOp.const(dtypes.float32.vec(4), (1.0, 2.0, 3.0, 4.0))
self.assertEqual(apply_rewrite(const_stack.gep(2)).arg, 3.0)
self.assertEqual(apply_rewrite(const_stack.index(2)).arg, 3.0)
def test_gep_tuple_on_const_stack(self):
# GEP on a const STACK using a tuple to extract multiple elements
const_stack = UOp.const(dtypes.float32.vec(4), (7.0, 8.0, 9.0, 10.0))
self.assertEqual(list(apply_rewrite_values(const_stack.gep((1, 3)))), [8.0, 10.0])
def test_gep_gep_simplification(self):
# Nested GEP simplification on a vector dtype
base_vector = UOp.const(dtypes.float32.vec(4), (10.0, 20.0, 30.0, 40.0))
gep_inner = base_vector.gep(1) # Extract 2nd element (20.0)
self.assertEqual(apply_rewrite(gep_inner.gep(0)).arg, 20.0)
self.assertEqual(list(apply_rewrite_values(UOp.vectorize(*[const_stack.index(i) for i in (1, 3)]))), [8.0, 10.0])
def test_vectorize_multiple_elements(self):
# Vectorizing multiple elements using GEP
base_vector = UOp.const(dtypes.float32.vec(4), (5.0, 10.0, 15.0, 20.0))
vectorized_uop = UOp(Ops.STACK, dtypes.float32.vec(4), src=(base_vector.gep(0), base_vector.gep(1), base_vector.gep(2), base_vector.gep(3)))
vectorized_uop = UOp(Ops.STACK, dtypes.float32.vec(4), src=tuple(base_vector.index(i) for i in range(4)))
self.assertEqual(list(apply_rewrite_values(vectorized_uop)), [5.0, 10.0, 15.0, 20.0])
test/null/test_uop_graph.py
+9 -16
View File
@@ -21,18 +21,11 @@ def const_values(u:UOp):
class TestGraphRewriteConst(unittest.TestCase):
def test_gep_const(self):
v1 = UOp.const(dtypes.int.vec(3), (0,1,2))
v2 = v1.gep(1)
v2 = v1.index(1)
ret = graph_rewrite(v2, sym)
self.assertEqual(ret.dtype, dtypes.int)
self.assertEqual(ret.arg, 1)
def test_gep_const_single(self):
v1 = UOp.const(dtypes.int.vec(3), 4)
v2 = v1.gep(1)
ret = graph_rewrite(v2, sym)
self.assertEqual(ret.dtype, dtypes.int)
self.assertEqual(ret.arg, 4)
def test_add_const(self):
v1 = UOp.const(dtypes.int, (0,1,2))
v2 = UOp.const(dtypes.int, (5,6,7))
@@ -261,7 +254,7 @@ class TestUOpGraph(unittest.TestCase):
idx = UOp.const(dtypes.int, 0)
ld = d0.load(idx, dtype=dtypes.float.vec(2))
vec = UOp(Ops.STACK, dtypes.float.vec(2), (ld,))
x = vec.gep(0)
x = vec.index(0)
alu = UOp(Ops.SQRT, dtypes.float, (x, ))
out = UOp(Ops.STORE, dtypes.void, (d0, idx, alu))
uops = to_uops_list([out])
@@ -284,27 +277,27 @@ class TestUOpGraph(unittest.TestCase):
# possible
val = d1.index(idx).load(dtype=dtypes.float.vec(4))
xyzw = tuple(val.gep(i) for i in range(4))
xyzw = tuple(val.index(i) for i in range(4))
self.assertIs(_test_vec(xyzw).op, Ops.LOAD)
# unaligned
val = d1.index(idx).load(dtype=dtypes.float.vec(4))
wzyx = tuple(val.gep(i) for i in reversed(range(4)))
wzyx = tuple(val.index(i) for i in reversed(range(4)))
self.assertIs(_test_vec(wzyx).op, Ops.STACK)
# different_size
val = d1.index(idx).load(dtype=dtypes.float.vec(2))
xy = tuple(val.gep(i) for i in range(2))
xy = tuple(val.index(i) for i in range(2))
self.assertIs(_test_vec(xy+xy).op, Ops.STACK)
val = d1.index(idx).load(dtype=dtypes.float.vec(4))
xy = tuple(val.gep(i) for i in range(2))
xy = tuple(val.index(i) for i in range(2))
self.assertIs(_test_vec(xy, count=2).op, Ops.STACK)
# different vals
val1 = d1.index(idx).load(dtype=dtypes.float.vec(2))
val2 = d2.index(idx).load(dtype=dtypes.float.vec(2))
xy1 = tuple(val1.gep(i) for i in range(2))
xy2 = tuple(val2.gep(i) for i in range(2))
xy1 = tuple(val1.index(i) for i in range(2))
xy2 = tuple(val2.index(i) for i in range(2))
self.assertIs(_test_vec(xy1+xy2).op, Ops.STACK)
def test_gep_vec_const_fold(self):
@@ -312,7 +305,7 @@ class TestUOpGraph(unittest.TestCase):
consts = [UOp.const(dtypes.float, float(i)) for i in range(vec_size)]
vec = UOp(Ops.STACK, dtypes.float.vec(vec_size), tuple(consts))
with Context(SPEC=0):
uops = to_uops_list([vec.gep(i) for i in range(vec_size)])
uops = to_uops_list([vec.index(i) for i in range(vec_size)])
for uop, const in zip(uops, consts):
self.assertEqual(uop, const)
test/null/test_uop_vmin_vmax.py
+1 -1
View File
@@ -319,7 +319,7 @@ class TestVminVmaxVConst(unittest.TestCase):
d1 = UOp.param(1, dtypes.int.ptr())
idx = UOp.const(dtypes.int, 0)
val = UOp(Ops.LOAD, dtypes.int.vec(2), (d1.index(idx).cast(dtypes.int.vec(2).ptr()),))
uop = (val // 32).gep(0)
uop = (val // 32).index(0)
self.assertEqual(uop.vmin, -67108864)
self.assertEqual(uop.vmax, 67108863)
tinygrad/codegen/__init__.py
+1 -1
View File
@@ -248,7 +248,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
sink = apply_opts(sink, ren, beam=ast.arg.beam)
# this is new style (TODO: this should all be removed)
sink = graph_rewrite(sink, pm_render, name="pm_render gep/stack")
sink = graph_rewrite(sink, pm_render, name="pm_render stack")
sink = graph_rewrite(sink, pm_remove_vec_dtypes, name="transform to new style")
# ** expander (expand_rewrite) **
tinygrad/codegen/decomp/dtype.py
+1 -1
View File
@@ -116,7 +116,7 @@ def f2f_load(x: UOp, fr:DType, to:DType) -> UOp:
def f2f_store(st, idx, val, fr:DType, to:DType):
if (n:=val.max_numel()) == 1: return st.replace(src=(idx, f2f(val.bitcast(f2f_dt[to]), to, fr)))
return UOp.group(*(st.replace(src=(reindex(idx, i, 1), f2f(val.gep(i).bitcast(f2f_dt[to]), to, fr))) for i in range(n)))
return UOp.group(*(st.replace(src=(reindex(idx, i, 1), f2f(val.index(i).bitcast(f2f_dt[to]), to, fr))) for i in range(n)))
pm_long_decomp = PatternMatcher([
(UPat(GroupOp.Defines, src=(UPat.var("sz"),), name="x"), lambda x,sz:
tinygrad/codegen/late/coalese.py
+1 -1
View File
@@ -38,7 +38,7 @@ def transform_to_image(ctx, buf:UOp, x:UOp) -> UOp|None:
# if no candidates, we don't rewrite
if len(cands) == 0: return None
# and tiebreak with indexing complexity (ie. number of nodes)
h, w, cidx = cands[0] if len(cands) == 1 else min(cands, key=lambda cand: len(cand[2].gep(1).simplify().backward_slice))
h, w, cidx = cands[0] if len(cands) == 1 else min(cands, key=lambda cand: len(cand[2].index(1).simplify().backward_slice))
buf = buf.replace(dtype=(dtypes.imageh if buf.dtype.itemsize == 2 else dtypes.imagef)((h, w, 4)))
shapes[buf.arg.slot] = (h, w)
if valid.op is not Ops.CONST or valid.arg is not True:
tinygrad/codegen/late/devectorizer.py
+2 -2
View File
@@ -18,7 +18,7 @@ def _drop_valid_stmts(valid:UOp, idx:UOp, height:int, width:int) -> list[UOp]:
# for X0 + X1 + ... >= 1, check if it's out of bound when Xi = 0 for all i
if not is_upper_bound and c == 1 and all(u.op in GroupOp.Irreducible and u.vmin == 0 for u in X.split_uop(Ops.ADD)):
testidx = functools.reduce(lambda nowidx,u: nowidx.substitute({u:u.const_like(0)}), X.split_uop(Ops.ADD), idx)
if testidx.gep(0).vmax < 0 or testidx.gep(1).vmax < 0:
if testidx.index(0).vmax < 0 or testidx.index(1).vmax < 0:
drop_stmt.append(stmt)
continue
@@ -45,7 +45,7 @@ def simplify_valid_image_load(buf:UOp, idx_y:UOp, idx_x:UOp, valid:UOp) -> UOp|N
if not drop_stmt and idx is start_idx: return None
new_valid = UOp.uprod(*ss) if (ss:=[s for s in valid.split_uop(Ops.AND) if s not in drop_stmt]) else None
idx_y, idx_x = idx.gep(1), idx.gep(0)
idx_y, idx_x = idx.index(1), idx.index(0)
return buf.index(idx_y.valid(new_valid), idx_x.valid(new_valid), ptr=True) if new_valid is not None else buf.index(idx_y, idx_x, ptr=True)
indexing_simplify = PatternMatcher([
tinygrad/renderer/llvmir.py
+2 -2
View File
@@ -218,9 +218,9 @@ class AMDLLVMRenderer(LLVMRenderer):
]) + base_rewrite
extra_matcher = LLVMRenderer.extra_matcher + create_non_native_float_pats(dtypes.fp8s) + PatternMatcher([
(UPat(Ops.CAST, dtype=dtypes.half.vec(16), src=UPat.var("y", dtypes.half.vec(8))),
lambda y: UOp(Ops.STACK, dtypes.half.vec(16), tuple(y.gep(i // 2) if i % 2 == 0 else UOp.const(dtypes.half, 0.0) for i in range(16)))),
lambda y: UOp(Ops.STACK, dtypes.half.vec(16), tuple(y.index(i // 2) if i % 2 == 0 else UOp.const(dtypes.half, 0.0) for i in range(16)))),
(UPat(Ops.CAST, dtype=dtypes.half.vec(8), src=UPat.var("y", dtypes.half.vec(16))),
lambda y: UOp(Ops.STACK, dtypes.half.vec(8), tuple(y.gep(i * 2) for i in range(8)))),
lambda y: UOp(Ops.STACK, dtypes.half.vec(8), tuple(y.index(i * 2) for i in range(8)))),
# amd llvm intrinsics llvm.log2/llvm.exp2 don't support double
(UPat(Ops.LOG2, dtype=dtypes.double, src=(UPat.var("d"),)), xlog2),
(UPat(Ops.EXP2, dtype=dtypes.double, src=(UPat.var("d"),)), xexp2),
tinygrad/runtime/ops_dsp.py
+1 -1
View File
@@ -13,7 +13,7 @@ from tinygrad.uop.ops import PatternMatcher, UPat
dsp_pm = PatternMatcher([
(((UPat.var('x').maximum(0) ^ -1).maximum(-256) ^ -1).cast(dtypes.uchar.vec(128)),
lambda x: UOp(Ops.CUSTOM, dtypes.uchar.vec(128), src=tuple(x.gep(tuple(range(i, i+32))) for i in range(0, 128, 32)),
lambda x: UOp(Ops.CUSTOM, dtypes.uchar.vec(128), src=tuple(UOp.vectorize(*[x.index(j) for j in range(i, i+32)]) for i in range(0, 128, 32)),
arg="__builtin_HEXAGON_V6_vpackhub_sat_128B(__builtin_HEXAGON_V6_vpackwh_sat_128B({3}, {2}), __builtin_HEXAGON_V6_vpackwh_sat_128B({1}, {0}))")),
])
tinygrad/schedule/rangeify.py
+1 -1
View File
@@ -32,7 +32,7 @@ def lower_shaped_wmma(ctx, x):
wmma_arg = (name, dims, dtype_in, dtype_out, device, threads, tc_upcast_axes, ())
wmma = UOp(Ops.WMMA, dtype_out.vec(x.src[2].shape[-1]), tuple(s[u].contract(u) for s, u in upcasts), arg=wmma_arg)
tmp = UOp.placeholder((x.src[2].shape[-1],), dtype_out, slot=next(ctx), addrspace=AddrSpace.REG)
return tmp.after(UOp.group(*[tmp[e].store(wmma.gep(e)) for e in range(x.src[2].shape[-1])]))
return tmp.after(UOp.group(*[tmp[e].store(wmma.index(e)) for e in range(x.src[2].shape[-1])]))
pm_store_ranges = PatternMatcher([
(UPat(Ops.STORE, name="x"), add_ranges_to_store),
tinygrad/uop/ops.py
+3 -9
View File
@@ -466,7 +466,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
def vectorize(self, *srcs): return self._stack(*srcs)
def index(self, *srcs:UOp|int|None, ptr=False, **kwargs):
new_srcs: list[UOp] = [UOp.const(dtypes.weakint, x) if isinstance(x, int) else x for x in srcs if x is not None]
return UOp(Ops.INDEX, kwargs.pop("dtype", self.dtype if ptr else self.dtype.base), (self,)+tuple(new_srcs), **kwargs)
if not ptr and len(new_srcs) == 1 and new_srcs[0].op is Ops.CONST and self.op is Ops.STACK:
return self.src[new_srcs[0].arg]
return UOp(Ops.INDEX, kwargs.pop("dtype", self.dtype if ptr else self.dtype.base.scalar()), (self,)+tuple(new_srcs), **kwargs)
def __getitem__(self, idx):
# pointers index into INDEX UOps (scalar lookup); everything else uses the shared mixin view path
if not isinstance(self.dtype, PtrDType): return super(UOp, self).__getitem__(idx)
@@ -511,14 +513,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
def bitcast(self, dtype:DTypeLike):
dtype = to_dtype(dtype)
return self if self.dtype == dtype else UOp(Ops.BITCAST, dtype, (self,))
def gep(self, i:tuple[int, ...]|int):
if isinstance(i, tuple) and len(i) == 1: return self.gep(i[0])
if isinstance(i, int):
# NOTE: these are just shortcuts to not have to create and fold later
if self.op is Ops.STACK: return self.src[i]
if self.op is Ops.CONST: return UOp.const(self.dtype.scalar(), self.arg)
return UOp(Ops.INDEX, self.dtype.scalar(), (self, UOp.const(dtypes.int, i)))
return UOp(Ops.STACK, self.dtype.scalar().vec(len(i)), tuple(self.gep(x) for x in i))
def load(self, *src:UOp, **kwargs): return UOp(Ops.LOAD, dtype=kwargs.pop("dtype", self.dtype.base), src=(self,)+src, **kwargs)
def store(self, src:UOp|ConstType, gate:UOp|None=None, **kwargs):
srcs = (self, self.const_like(src) if not isinstance(src, UOp) else src) + ((gate,) if gate is not None else ())
tinygrad/uop/spec.py
+1 -1
View File
@@ -228,7 +228,7 @@ spec_full = PatternMatcher([
# allow any AFTER
(UPat(Ops.AFTER, src=(UPat(),), allow_any_len=True), lambda: True),
# expander: unroll/contract/gep/cat
# expander: unroll/contract
(UPat((Ops.UNROLL, Ops.CONTRACT), src=(UPat(),)), lambda: True),
# all loads/stores
tinygrad/uop/symbolic.py
+1 -1
View File
@@ -182,7 +182,7 @@ symbolic_simple = propagate_invalid + PatternMatcher([
(UPat.cvar("gate").where(UPat.var("c0"), UPat.var("c1")), lambda gate, c0, c1: c0 if gate.arg else c1),
# a.where(b.where(c, d), d) -> (a & b).where(c, d)
(UPat.var("a").where(UPat.var("b").where(UPat.var("c"), UPat.var("d")), UPat.var("d")), lambda a,b,c,d: (a&b).where(c,d)),
# STACK on INDEX CONST (TODO: remove all the GEP crap)
# STACK on INDEX CONST
(UPat(Ops.STACK, src=UPat(Ops.INDEX, src=(UPat.var("src"), UPat(Ops.CONST))), name="stk"),
lambda src,stk: src if stk.shape == src.shape and list(range(len(stk.src))) == [x.src[1].arg for x in stk.src] else None),
# INDEX on STACK
tinygrad/uop/upat.py
+2 -2
View File
@@ -33,7 +33,7 @@ def _get_clause(self:UPat, base:UOp, depth=0) -> UOp:
if self.src is not None:
# single match
if len(self.src) == 1 and isinstance(self.src[0], tuple):
and_clause += [_get_clause(s, base.gep(i), depth) for i,s in enumerate(self.src[0])]
and_clause += [_get_clause(s, base.index(i), depth) for i,s in enumerate(self.src[0])]
# repeat match
elif len(self.src) == 1 and isinstance(self.src[0], itertools.repeat):
it = UOp(Ops.NOOP, arg=f"ituop{depth}")
@@ -41,7 +41,7 @@ def _get_clause(self:UPat, base:UOp, depth=0) -> UOp:
and_clause.append(UOp(Ops.RANGE, src=(match, it, base), arg="all([{0} for {1} in {2}.src])"))
# multi match (fork)
elif len(self.src) > 1 and all(isinstance(x, tuple) for x in self.src):
fork_cond = [UOp(Ops.AND, src=tuple([_get_clause(s, base.gep(i), depth) for i,s in enumerate(ss)])) for ss in self.src]
fork_cond = [UOp(Ops.AND, src=tuple([_get_clause(s, base.index(i), depth) for i,s in enumerate(ss)])) for ss in self.src]
and_clause.append(UOp(Ops.OR, src=tuple(fork_cond)))
else: raise RuntimeError("broken")
return UOp(Ops.AND, src=tuple(and_clause))