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

add multiple_of to ParamArg (#17101)

Browse Source
This commit is contained in:
sirhcm
2026-07-20 20:11:54 -04:00
committed by GitHub
parent f7ce7f330d
commit 980748ccfc
7 changed files with 59 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/null/test_uop_vmin_vmax.py
+14
View File
@@ -371,6 +371,10 @@ class TestConstFactor(unittest.TestCase):
uop = (x * 3) * 5
self.assertEqual(uop.const_factor(), 15) # Constant multipliers are combined (3 * 5 = 15)
def test_const_factor_variable_multiple_of(self):
x = UOp.variable('x', 16, 32, multiple_of=4)
self.assertEqual(x.const_factor(), 4)
class TestDivides(unittest.TestCase):
def test_divides_constant_exact(self):
# Divides a constant by an exact divisor
@@ -409,5 +413,15 @@ class TestDivides(unittest.TestCase):
result = uop.divides(3)
self.assertIsNone(result) # Cannot divide by 3, since 4 is not divisible by 3
def test_divides_variable_multiple_of_exact(self):
x = UOp.variable('x', 16, 32, multiple_of=4)
result = x.divides(4)
self.assertIsNotNone(result)
def test_divides_variable_multiple_of_factor(self):
x = UOp.variable('x', 16, 32, multiple_of=4)
result = x.divides(2)
self.assertIsNotNone(result)
if __name__ == '__main__':
unittest.main()
test/opt/test_gen_float4.py
+27 -1
View File
@@ -1,5 +1,5 @@
import unittest
from tinygrad import Device, Tensor, dtypes
from tinygrad import Device, Tensor, Variable, dtypes
from tinygrad.uop.ops import UOp, Ops
from tinygrad.codegen import to_program
from tinygrad.codegen.opt import Opt, OptOps
@@ -114,5 +114,31 @@ class TestFloat4(unittest.TestCase):
assert TestFloat4.count_float4(uops) == (1, 1)
def test_float4_aligned_variable(self):
x = Variable('x', 0, 4, multiple_of=4).bind(4)
a = Tensor.empty(4).realize()
b = Tensor.empty(12).realize().shrink(((x, x+4),))
c = a + b
# should float4 both
s = c.linear_with_vars()[0].src[0]
uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).src[1].src)
assert TestFloat4.count_float4(uops) == (2, 1)
def test_float4_unaligned_variable(self):
x = Variable('x', 0, 4, multiple_of=2).bind(4)
a = Tensor.empty(4).realize()
b = Tensor.empty(12).realize().shrink(((x, x+4),))
c = a + b
# should float4 a but not b
s = c.linear_with_vars()[0].src[0]
uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).src[1].src)
assert TestFloat4.count_float4(uops) == (1, 1)
if __name__ == '__main__':
unittest.main()
tinygrad/callify.py
+3 -2
View File
@@ -180,8 +180,9 @@ def finalize_after(ctx:AllocCtx, x:UOp):
def replace_input_buffer(ctx:AllocCtx, b:UOp):
ctx.replacements.append(b)
return UOp.param(len(ctx.replacements)-1, b.dtype, b.shape, b.device,
b._min_max if b.op is Ops.BIND else None, b.src[0].expr if b.op is Ops.BIND else None,
b.addrspace if b.addrspace is not None else AddrSpace.GLOBAL)
b._min_max if b.op is Ops.BIND else None, name=b.src[0].expr if b.op is Ops.BIND else None,
addrspace=b.addrspace if b.addrspace is not None else AddrSpace.GLOBAL,
multiple_of=b.src[0].arg.multiple_of if b.op is Ops.BIND else None)
pm_finalize_call = PatternMatcher([
(UPat(Ops.AFTER, name="x"), finalize_after),
tinygrad/schedule/multi.py
+1 -1
View File
@@ -148,7 +148,7 @@ def rewrite_into_function(call:UOp):
def param_to_multi(p:UOp):
if p.axis is None: return None
return UOp.param(p.arg.slot, p.dtype, p.shard_shape, p.device, p.arg.vmin_vmax, p.arg.name, p.arg.addrspace).multi(p.axis)
return UOp.param(p.arg.slot, p.dtype, p.shard_shape, p.device, p.arg.vmin_vmax, p.arg.multiple_of, p.arg.name, p.arg.addrspace).multi(p.axis)
# NOTE: this is the same pattern as unrolled ranges
multi_pm = PatternMatcher([
tinygrad/schedule/rangeify.py
+1 -1
View File
@@ -470,7 +470,7 @@ to_define_global = PatternMatcher([
(UPat(Ops.STORE, name="x"), find_bufs),
(UPat((Ops.BUFFER, Ops.MSTACK, Ops.MSELECT), name="buf"), debuf),
(UPat(Ops.PARAM, name="v"), lambda v:
UOp.variable(v.arg.name, v.arg.vmin_vmax[0], v.arg.vmin_vmax[1], v.dtype)
UOp.variable(v.arg.name, v.arg.vmin_vmax[0], v.arg.vmin_vmax[1], v.dtype, multiple_of=v.arg.multiple_of)
if v.arg.name is not None and v.arg.vmin_vmax is not None else None),
# this renumbers the params
tinygrad/uop/divandmod.py
+2
View File
@@ -11,6 +11,8 @@ def fold_divmod_general(d: UOp) -> UOp|None:
if y.vmin==y.vmax==0: raise ZeroDivisionError(f"{'Division' if d.op is Ops.FLOORDIV else 'Mod'} by zero trying to rewrite {x.alu(d.op, y)}")
# x//y is constant
if (xdiv:=x//y).vmin == xdiv.vmax: return x - xdiv.vmin*y if d.op is Ops.FLOORMOD else xdiv.const_like(xdiv.vmin)
# PARAM // c is irreducible
if x.op is Ops.PARAM and y.op is Ops.CONST and x.arg.multiple_of % y.arg == 0: return d.const_like(0) if d.op is Ops.FLOORMOD else None
# split uops for the rest of the processing
x_peeled, const = x.pop_const()
tinygrad/uop/ops.py
+11 -8
View File
@@ -24,12 +24,13 @@ class ParamArg:
slot: int
dtype: DType
vmin_vmax: tuple[PyConst, PyConst]|None = None
multiple_of: int|None = None
name: str|None = None
addrspace: AddrSpace|None = AddrSpace.GLOBAL
axis: int|None = None
device: str|tuple[str, ...]|None = None
def __repr__(self):
fields = (("vmin_vmax", None), ("name", None), ("addrspace", AddrSpace.GLOBAL), ("axis", None), ("device", None))
fields = (("vmin_vmax", None), ("multiple_of", None), ("name", None), ("addrspace", AddrSpace.GLOBAL), ("axis", None), ("device", None))
args = [repr(self.slot), repr(self.dtype)] + [f"{k}={v!r}" for k,default in fields if (v:=getattr(self, k)) != default]
return f"ParamArg({', '.join(args)})"
axis_letters = {AxisType.GLOBAL: "g", AxisType.THREAD: "t", AxisType.LOCAL: "l", AxisType.WARP: "w", AxisType.LOOP: "L", AxisType.UPCAST: "u",
@@ -915,9 +916,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
# *** uop Variable stuff ***
@staticmethod
def variable(name:str, min_val:PyConst, max_val:PyConst, dtype:DType=dtypes.index) -> UOp:
def variable(name:str, min_val:PyConst, max_val:PyConst, dtype:DType=dtypes.index, multiple_of:int=1) -> UOp:
return UOp(Ops.PARAM, src=(shape_to_shape_arg(()),),
arg=ParamArg(-1, dtype, name=name, vmin_vmax=(min_val, max_val), addrspace=AddrSpace.ALU))
arg=ParamArg(-1, dtype, name=name, vmin_vmax=(min_val, max_val), multiple_of=multiple_of, addrspace=AddrSpace.ALU))
@property
def expr(self) -> str:
assert self.op is Ops.PARAM
@@ -926,6 +927,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
assert self.op is Ops.PARAM and self.addrspace is AddrSpace.ALU, f"op is {self.op}, need PARAM"
uval = self.const_like(val) if isinstance(val, int) else val
assert self.vmin <= uval.vmin and uval.vmax <= self.vmax, f"bind {val} not in range [{self.vmin}, {self.vmax}]"
assert uval.divides(self.arg.multiple_of) is not None, f"bind {val} not divisible by {self.arg.multiple_of}"
return UOp(Ops.BIND, src=(self, uval))
def unbind(self) -> tuple[Variable, int]:
assert self.op is Ops.BIND and self.src[0].op is Ops.PARAM and self.src[1].op is Ops.CONST, f"can't unbind {self}"
@@ -948,6 +950,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
if self.op is Ops.STACK: return math.gcd(*[x.const_factor() for x in self.src])
if self.op is Ops.ADD: return math.gcd(self.src[0].const_factor(), self.src[1].const_factor())
if self.op is Ops.MUL: return self.src[0].arg if self.src[0].op is Ops.CONST else self.src[1].arg if self.src[1].op is Ops.CONST else 1
if self.op is Ops.PARAM and self.arg.multiple_of is not None: return self.arg.multiple_of
return 1
def divides(self, v:int) -> UOp|None:
if v==1: return self
@@ -959,6 +962,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
if self.op is Ops.MUL:
if (d0:=self.src[0].divides(v)) is not None: return d0 * self.src[1]
if (d1:=self.src[1].divides(v)) is not None: return self.src[0] * d1
if self.op is Ops.PARAM and self.arg.multiple_of is not None: return self // v if self.arg.multiple_of%v == 0 else None
return None # generic None if we aren't sure
def pop_const(self, op=Ops.ADD) -> tuple[UOp, PyConst]: # NOTE: assume Invalid ALU is resolved
return (self.src[0], self.src[1].arg) if self.op is op and self.src[1].op is Ops.CONST else (self, identity_element(op, self.dtype))
@@ -1087,17 +1091,16 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
# TODO: this should replace placeholder
@staticmethod
def param(slot:int, dtype:DType, shape:tuple[sint, ...]|None=None, device=None, vmin_vmax:tuple[PyConst, PyConst]|None=None, name=None,
addrspace=AddrSpace.GLOBAL, axis:int|None=None):
def param(slot:int, dtype:DType, shape:tuple[sint, ...]|None=None, device=None, vmin_vmax:tuple[PyConst, PyConst]|None=None,
multiple_of:int|None=None, name=None, addrspace=AddrSpace.GLOBAL, axis:int|None=None):
if dtype in dtypes.weaks: raise RuntimeError(f"cannot create param for weak dtype {dtype}")
if shape is not None and axis is not None and isinstance(device, tuple):
shape = tuple(s*len(device) if i == axis else s for i,s in enumerate(shape))
src: tuple[UOp, ...] = (UOp(Ops.NOOP) if shape is None else shape_to_shape_arg(shape),)
return UOp(Ops.PARAM, src=src, arg=ParamArg(slot, dtype, vmin_vmax, name, addrspace, axis, device))
return UOp(Ops.PARAM, src=src, arg=ParamArg(slot, dtype, vmin_vmax, multiple_of, name, addrspace, axis, device))
def param_like(self, slot:int):
addrspace = self.addrspace if self.addrspace is not None else AddrSpace.GLOBAL
if self.op is Ops.BIND:
return UOp.param(slot, self.dtype, self._shape, self.device, cast(tuple[int, int], self._min_max), self.src[0].expr, addrspace)
if self.op is Ops.BIND: return self.src[0].replace(arg=replace(self.src[0].arg, slot=slot, addrspace=addrspace))
return UOp.param(slot, self.dtype, self.shard_shape if self.axis is not None else self._shape, self.device, addrspace=addrspace, axis=self.axis)
@staticmethod