tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 09:36:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

use weak 0 in convert_pad_to_where_to_keep_behavior_local [pr] (#17398)

Browse Source
This commit is contained in:
chenyu
2026-08-04 19:44:34 -04:00
committed by GitHub
parent d79772f057
commit f295f9fc99
3 changed files with 5 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/schedule/indexing.py
+1 -1
View File
@@ -101,7 +101,7 @@ def convert_pad_to_where_to_keep_behavior_local(ctx:IndexingContext, x:UOp):
if x not in ctx.range_map: return None
bx = create_bufferize_and_index_based_on_ranges(ctx, x)
valid: UOp = UOp.const(True).uprod([r.get_valid() for r in ctx.range_map[x][0]])
return valid.where(bx.src[0], UOp.const(0, x.dtype))
return valid.where(bx.src[0], UOp.const(x.dtype.const(0)))
def convert_reduce_to_reduce_with_ranges(ctx:IndexingContext, x:UOp):
if x.arg[1] == 0: return None
tinygrad/uop/ops.py
+1
View File
@@ -1138,6 +1138,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
@staticmethod
def placeholder(shape:tuple[int, ...], dtype:DType, slot:int, addrspace=AddrSpace.GLOBAL, device=None, volatile=False):
dtype = strong_dtype(dtype) # storage is never weak: a placeholder commits the width of what's put in it
if addrspace is AddrSpace.GLOBAL:
ret = UOp(Ops.PARAM, src=(shape_to_shape_arg((prod(shape),)),), arg=ParamArg(slot, dtype, addrspace=addrspace, device=device,volatile=volatile))
else:
tinygrad/uop/render.py
+3 -1
View File
@@ -104,8 +104,10 @@ pm_pyrender_extra = PatternMatcher([
(UPat(Ops.CMOD, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.alu(Ops.CMOD, {ctx[x.src[1]]})"),
# `.where` re-promotes its operands, so render WHERE via .alu() too
(UPat(Ops.WHERE, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.alu(Ops.WHERE, {ctx[x.src[1]]}, {ctx[x.src[2]]})"),
# the binary operators re-promote their operands (a weak src meeting a strong one gets a cast), render those via .alu() too
(UPat(set(syms.keys())-{Ops.SUB, Ops.CDIV, Ops.CMOD}, name="x"), lambda ctx,x:
strip_binary_parens(x, ctx[x.src[0]], ctx[x.src[1]], lambda a,b: f"({a}{syms[x.op]}{b})")),
strip_binary_parens(x, ctx[x.src[0]], ctx[x.src[1]], lambda a,b: f"({a}{syms[x.op]}{b})")
if x.src[0]._broadcasted(x.src[1]) == x.src else f"{ctx[x.src[0]]}.alu({x.op}, {ctx[x.src[1]]})"),
(UPat(sugar, src=(), name="x"), lambda x: f"UOp.{x.op.name.lower()}("+', '.join(([f'arg={repr(x.arg)}'] if x.arg is not None else []))+")"),
(UPat(sugar, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}("+', '.join([ctx[y] for y in x.src[1:]] + \
([f'arg={repr(x.arg)}'] if x.arg is not None else []))+")"),