geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

clean up STACK with a const [PR] (#17539)

Browse Source
This commit is contained in:
chenyu
2026-08-14 21:07:01 -04:00
committed by GitHub
parent 6ea665ed66
commit 64ccbde3bb
2 changed files with 3 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/uop/ops.py
+1 -3
View File
@@ -796,11 +796,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
case Ops.PAD | Ops.SHRINK: src_args = list(zip(*arg))
case Ops.PERMUTE | Ops.FLIP: src_args = []
case Ops.STACK:
# arg is the other srcs; all are cast to the promoted dtype, spec requires STACK srcs to match its dtype
srcs = (self,)+tuple(arg)
dtype = cast(DType, dtype_from_uop(Ops.STACK, srcs, None))
# TODO: why cast here?
return UOp(Ops.STACK, dtype, tuple(u if u.base.is_invalid else u.cast(dtype) for u in srcs))
return UOp(Ops.STACK, dtype, tuple(u if u.base.is_invalid else UOp.const(u.val, dtype) if u.op is Ops.CONST else u.cast(dtype) for u in srcs))
case _: raise RuntimeError(f"{op} is not a MovementOp")
usrcs = [shape_to_shape_arg(arg) for arg in src_args]
if len(usrcs) == 0: return UOp(op, src=(self,), arg=arg)
tinygrad/uop/weak.py
+2 -2
View File
@@ -12,7 +12,7 @@ def lower_weak_node(u:UOp) -> UOp|None:
if src == u.src or any(s.dtype in dtypes.weaks for s in src[start:]): return None
dt = strong_dtype(least_upper_dtype(select_dtype(u), *(s.dtype for s in src)) if u.op in GroupOp.Binary
else unwrap(dtype_from_uop(u.op, src, u.arg)))
return u.replace(dtype=None, src=src[:start]+tuple(s if s.base.is_invalid else s.cast(dt) for s in src[start:])).cast(u.dtype)
return u.replace(dtype=None, src=src[:start]+tuple(s if s.base.is_invalid else commit_weak(s, dt) for s in src[start:])).cast(u.dtype)
pm_lower_weak = PatternMatcher([
(UPat(Ops.CONST, dtype=dtypes.weaks, name="u"), lambda u: UOp.const(u.val, select_dtype(u)).cast(u.dtype)),
@@ -40,7 +40,7 @@ def lower_weak_srcs(ctx:dict[UOp, UOp]|None, u:UOp) -> UOp|None:
return None if ret is u else ret
def commit_weak(s:UOp, dt:DType) -> UOp:
# a bare weak CONST commits directly (the value stays mathematical, emission truncates), a weak non-const src takes the demand cast
# a CONST commits directly at dt (the value stays mathematical, emission truncates), a non-const src takes the cast
return UOp.const(s.val, dt) if s.op is Ops.CONST else s.cast(dt)
def commit_weak_srcs(u:UOp) -> UOp|None: