diff --git a/extra/gemm/metal_uop_matmul.py b/extra/gemm/metal_uop_matmul.py index a2d619b45e..431a645b32 100644 --- a/extra/gemm/metal_uop_matmul.py +++ b/extra/gemm/metal_uop_matmul.py @@ -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() diff --git a/extra/gemm/mi350x_uop_matmul.py b/extra/gemm/mi350x_uop_matmul.py index b8a82b784a..8d0725d609 100644 --- a/extra/gemm/mi350x_uop_matmul.py +++ b/extra/gemm/mi350x_uop_matmul.py @@ -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)]) diff --git a/extra/thunder/tiny/tk/group.py b/extra/thunder/tiny/tk/group.py index d592a44f9d..cb76fe507d 100644 --- a/extra/thunder/tiny/tk/group.py +++ b/extra/thunder/tiny/tk/group.py @@ -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) diff --git a/test/backend/test_custom_kernel.py b/test/backend/test_custom_kernel.py index 458b3c4c3f..3599bed385 100644 --- a/test/backend/test_custom_kernel.py +++ b/test/backend/test_custom_kernel.py @@ -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): diff --git a/test/null/test_graph_rewrite.py b/test/null/test_graph_rewrite.py index 0cd1fa4281..782e6c770a 100644 --- a/test/null/test_graph_rewrite.py +++ b/test/null/test_graph_rewrite.py @@ -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]) diff --git a/test/null/test_uop_graph.py b/test/null/test_uop_graph.py index 2ff3cc08b9..212a671d12 100644 --- a/test/null/test_uop_graph.py +++ b/test/null/test_uop_graph.py @@ -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) diff --git a/test/null/test_uop_vmin_vmax.py b/test/null/test_uop_vmin_vmax.py index 81b1c02c04..f6c939a8f6 100644 --- a/test/null/test_uop_vmin_vmax.py +++ b/test/null/test_uop_vmin_vmax.py @@ -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) diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index d224f5e34c..e34ad8f487 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -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) ** diff --git a/tinygrad/codegen/decomp/dtype.py b/tinygrad/codegen/decomp/dtype.py index c44bbf0f1f..ec4062024b 100644 --- a/tinygrad/codegen/decomp/dtype.py +++ b/tinygrad/codegen/decomp/dtype.py @@ -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: diff --git a/tinygrad/codegen/late/coalese.py b/tinygrad/codegen/late/coalese.py index 447f60c6da..bcde59364e 100644 --- a/tinygrad/codegen/late/coalese.py +++ b/tinygrad/codegen/late/coalese.py @@ -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: diff --git a/tinygrad/codegen/late/devectorizer.py b/tinygrad/codegen/late/devectorizer.py index 4bb387602f..764a3ebd70 100644 --- a/tinygrad/codegen/late/devectorizer.py +++ b/tinygrad/codegen/late/devectorizer.py @@ -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([ diff --git a/tinygrad/renderer/llvmir.py b/tinygrad/renderer/llvmir.py index 216b712d3c..262da899ac 100644 --- a/tinygrad/renderer/llvmir.py +++ b/tinygrad/renderer/llvmir.py @@ -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), diff --git a/tinygrad/runtime/ops_dsp.py b/tinygrad/runtime/ops_dsp.py index 2daaa54265..12a6a2dc85 100644 --- a/tinygrad/runtime/ops_dsp.py +++ b/tinygrad/runtime/ops_dsp.py @@ -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}))")), ]) diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 20726ab8eb..daa3cbd44e 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -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), diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 020cf54468..06bb9fd48b 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -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 ()) diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index 0533aef5b8..ba25eb1e5a 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -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 diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 8b26cc74f5..4daedd268e 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -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 diff --git a/tinygrad/uop/upat.py b/tinygrad/uop/upat.py index 387a9fd6bf..9e8016624e 100644 --- a/tinygrad/uop/upat.py +++ b/tinygrad/uop/upat.py @@ -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))