Function flatten_bufferize

tinygrad/schedule/rangeify.py:426–434

(x:UOp)

# collapse any BUFFERIZE to single input BUFFERIZE
def flatten_bufferize(x:UOp):
  if len(x.src) == 2: return None
  ret = x.replace(src=(x.src[0], get_single_element(apply_movement_op(Ops.RESHAPE, (prod(x.shape),), x.shape, x.src[1:]))))
  rngs = x.src[1:]
  ret = ret.reshape(x.shape)
  if any(r.op is Ops.RANGE and r.src[0].op is not Ops.CONST for r in rngs):
    sym_shape = tuple([r.src[0] if r.op is not Ops.CONST else 1 for r in rngs])
    ret = ret.shrink(tuple([(0,x) for x in sym_shape]))
  return ret
pm_flatten_bufferize = PatternMatcher([(UPat(Ops.STAGE, name="x"), flatten_bufferize)])

pm_add_buffers = pm_mops+pm_flatten_bufferize+to_bufferview+PatternMatcher([