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

fix shapes to match vecless (#15866)

Browse Source 
* fix shapes

* need to simplify shapes
This commit is contained in:
George Hotz
2026-04-22 18:27:46 +08:00
committed by GitHub
parent de8f58899e
commit 2d7fa58e61
2 changed files with 8 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/schedule/indexing.py
+4 -10
View File
@@ -177,21 +177,15 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
# 2. from the single consumer if this op only has one consumer
# 3. potentially new if this op has 2+ consumers
shape = x._shape
if x.op is Ops.STORE:
# TODO: TestTensorVariable.test_symbolic_var_sum_alt_name fails with this, fix canonicalize on variables.
#assert x.src[0].shape == x.src[1].shape, f"STORE must have matching input shapes, {x.src[0].shape} != {x.src[1].shape}"
shape = x.src[0].shape
consumer_rngs = [rctx.range_map[c][0] for c in consumer_map[x] if c in rctx.range_map]
if x in rctx.realize_map:
# if this is in the realize_map, we create new ranges (at the output)
out_rngs = tuple(rctx.new_range(s) for s in shape)
out_rngs = tuple(rctx.new_range(s) for s in x.shape)
# all ranges are ended now
ending_ranges[x] = []
# mark all ranges as ended
assert rctx.realize_map[x] is None
rctx.realize_map[x] = list(range(len(shape)))
rctx.realize_map[x] = list(range(len(x.shape)))
elif len(consumer_rngs) == 0:
# if no consumers have ranges and this isn't realized, this doesn't have ranges either.
continue
@@ -217,7 +211,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
minimum_valid = UOp.const(dtypes.bool, False).usum(valids)
_out_rngs.append(graph_rewrite(minimum_valid.where(local_rngs[0], UOp.invalid()), symbolic, name="minimum_valid"))
else:
_out_rngs.append(rctx.new_range(shape[i]))
_out_rngs.append(rctx.new_range(x.shape[i]))
_realize_axis.append(i)
out_rngs = tuple(_out_rngs)
@@ -234,7 +228,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
ending_ranges[x] = []
if len(_realize_axis):
rctx.realize_map[x] = _realize_axis
out_rngs = tuple([(rctx.new_range(shape[i]) if i in _realize_axis else r) for i,r in enumerate(out_rngs)])
out_rngs = tuple([(rctx.new_range(x.shape[i]) if i in _realize_axis else r) for i,r in enumerate(out_rngs)])
# TODO: some ops don't have shape, enable this after the `.st` property is removed
#assert len(out_rngs) == len(x.shape), \
tinygrad/uop/ops.py
+4 -4
View File
@@ -211,7 +211,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
def _shape(self) -> tuple[sint, ...]|None:
match self.op:
# late ops don't have shape
case Ops.UNIQUE | Ops.LUNIQUE | Ops.DEVICE | Ops.LOAD | Ops.STORE | Ops.IF | Ops.BARRIER | Ops.CUSTOM | Ops.CUSTOMI | \
case Ops.UNIQUE | Ops.LUNIQUE | Ops.DEVICE | Ops.IF | Ops.BARRIER | Ops.CUSTOM | Ops.CUSTOMI | \
Ops.VECTORIZE | Ops.GEP | Ops.UNROLL | Ops.CONTRACT | Ops.SINK | \
Ops.LINEAR | Ops.PROGRAM | Ops.SOURCE | Ops.BINARY | Ops.INS | Ops.TUPLE | Ops.CALL | Ops.FUNCTION:
return None
@@ -259,7 +259,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
case Ops.SHAPED_WMMA: return self.src[2]._shape
# passthrough ops
case Ops.REDUCE | Ops.MSTACK | Ops.MSELECT | Ops.DETACH | Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.AFTER | Ops.END:
case Ops.REDUCE | Ops.MSTACK | Ops.MSELECT | Ops.DETACH | Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.AFTER | Ops.END | Ops.LOAD:
return self.src[0]._shape
# TODO: disallow shape changing bitcast
@@ -312,7 +312,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
return tuple(1 if i in axis_arg else s for i,s in enumerate(ps))
# elementwise ops keep the shape the same. all inputs with shape must match
if self.op in GroupOp.ALU.union({Ops.CAST, Ops.COPY, Ops.NOOP, Ops.GROUP, Ops.SINK, Ops.ALLREDUCE}):
if self.op in GroupOp.ALU.union({Ops.CAST, Ops.COPY, Ops.NOOP, Ops.GROUP, Ops.SINK, Ops.ALLREDUCE, Ops.STORE}):
input_shapes = [x._shape for x in self.src if x._shape is not None]
if len(input_shapes) == 0: return None
if not all_same(input_shapes): raise RuntimeError(f"shape mismatch at {self.op}: {input_shapes}")
@@ -624,7 +624,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
@functools.cached_property
def marg(self):
match self.op:
case Ops.RESHAPE | Ops.EXPAND: return tuple(self.src[1].sgep(i) for i in range(self.src[1].dtype.count))
case Ops.RESHAPE | Ops.EXPAND: return tuple(ssimplify(self.src[1].sgep(i)) for i in range(self.src[1].dtype.count))
case Ops.PAD | Ops.SHRINK: return tuple((self.src[1].sgep(i), self.src[2].sgep(i)) for i in range(self.src[1].dtype.count))
case Ops.PERMUTE | Ops.FLIP: return self.arg
case _: raise RuntimeError(f"{self.op} is not a MovementOp")