From 980748ccfc4e3900ac652d8451e2ead9bfb4d09a Mon Sep 17 00:00:00 2001 From: Christopher Milan Date: Mon, 20 Jul 2026 20:11:54 -0400 Subject: [PATCH] add multiple_of to ParamArg (#17101) --- test/null/test_uop_vmin_vmax.py | 14 ++++++++++++++ test/opt/test_gen_float4.py | 28 +++++++++++++++++++++++++++- tinygrad/callify.py | 5 +++-- tinygrad/schedule/multi.py | 2 +- tinygrad/schedule/rangeify.py | 2 +- tinygrad/uop/divandmod.py | 2 ++ tinygrad/uop/ops.py | 19 +++++++++++-------- 7 files changed, 59 insertions(+), 13 deletions(-) diff --git a/test/null/test_uop_vmin_vmax.py b/test/null/test_uop_vmin_vmax.py index 248d002ff1..dc2b2ae37f 100644 --- a/test/null/test_uop_vmin_vmax.py +++ b/test/null/test_uop_vmin_vmax.py @@ -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() diff --git a/test/opt/test_gen_float4.py b/test/opt/test_gen_float4.py index 9300d46e85..b946cb6bd1 100644 --- a/test/opt/test_gen_float4.py +++ b/test/opt/test_gen_float4.py @@ -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() diff --git a/tinygrad/callify.py b/tinygrad/callify.py index 4299e19c48..438598b06f 100644 --- a/tinygrad/callify.py +++ b/tinygrad/callify.py @@ -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), diff --git a/tinygrad/schedule/multi.py b/tinygrad/schedule/multi.py index 0b6fe5e6b0..a67ada0b16 100644 --- a/tinygrad/schedule/multi.py +++ b/tinygrad/schedule/multi.py @@ -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([ diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index f36381794a..eca1e6887d 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -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 diff --git a/tinygrad/uop/divandmod.py b/tinygrad/uop/divandmod.py index e60a9b9c85..afd305cb3d 100644 --- a/tinygrad/uop/divandmod.py +++ b/tinygrad/uop/divandmod.py @@ -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() diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index ff9992de7e..4e73be3c97 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -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