diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index f734435c0e..54ac1a9485 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -39,7 +39,16 @@ pm_fold_moved_after = PatternMatcher([ def _mop_index(r:UOp, idx:UOp): idxs = idx.src[1:] if len(idxs) == len(r.shape): - return r.src[0].index(*apply_movement_op(r.op, r.src[0].shape, r.marg, idxs), dtype=idx.dtype, arg=idx.arg) + ret = r.src[0].index(*apply_movement_op(r.op, r.src[0].shape, r.marg, idxs), dtype=idx.dtype, arg=idx.arg) + if r.op is Ops.PAD: + # insert 0 for PAD with where + # TODO: does this need simplify + a = UOp.const(True) + for s in ret.src[1:]: + if s.op is Ops.WHERE and s.src[2].op is Ops.CONST and s.src[2].arg == Invalid: + a = a & s.src[0] + ret = a.where(ret, 0) + return ret if r.op is Ops.RESHAPE: src_prefix = len(r.src[0].shape) - len(r.shape[len(idxs):]) if src_prefix >= 0 and r.src[0].shape[src_prefix:] == r.shape[len(idxs):]: @@ -577,7 +586,7 @@ def rangeify_on_store(ctx, x:UOp): rngs = [UOp.range(s, next(ctx)) for s in x.shape] return x.src[0].index(*rngs).store(x.src[1].index(*rngs)).end(*rngs) -pm_simple_rangeify = pm_mops+PatternMatcher([ +pm_simple_rangeify = PatternMatcher([ # INDEX without src is nothing (TODO: this should be in mop_cleanup) (UPat(Ops.INDEX, src=(UPat.var('x'),)), lambda x: x), # handle movement ops on INDEX diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index d9c32c0a31..6950cbb5cd 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -47,7 +47,7 @@ from tinygrad.dtype import dtypes, AddrSpace uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0", Ops.REDUCE: "#FF5B5B", Ops.RANGE: "#c8a0e0", Ops.BARRIER: "#ff8080", Ops.IF: "#c8b0c0", Ops.SPECIAL: "#c0c0ff", - Ops.INDEX: "#D8F9E4", Ops.STACK: "#D8F9E4", + Ops.INDEX: "#CEF9B7", Ops.STACK: "#D8F9E4", Ops.WMMA: "#efefc0", Ops.UNSHARD: "#f6ccff", Ops.INS: "#eec4ff", **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80", Ops.SLICE: "#E5EAFF", Ops.BUFFER: "#B0BDFF", Ops.GETADDR: "#9DB1F0", Ops.COPY: "#a040a0", Ops.CUSTOM_FUNCTION: "#bf71b6",