From 37d3ca152ef1ae8d0873103bd34a3a6d0439e7a9 Mon Sep 17 00:00:00 2001 From: Sieds Lykles <93992551+S-Lykles@users.noreply.github.com> Date: Sat, 14 Jun 2025 20:55:51 +0200 Subject: [PATCH] Adapt `>>` for division by power of two to all ints (#10803) * Change divison by power of two to always use shift * Change test to test int instead of uint * simplify condition * add old rule back with comment * remove import * use sresolve instead of simplify * use keyword in simplify instead of sresolve * webgpu cast y to uint * remove comment * explicitly set dtype in wgsl * without simplify * undo simplify kwarg * change test to test both int32 and uint32 --- test/test_uops.py | 19 ++++++++++--------- tinygrad/codegen/devectorizer.py | 7 ++++--- tinygrad/renderer/wgsl.py | 3 ++- 3 files changed, 16 insertions(+), 13 deletions(-) diff --git a/test/test_uops.py b/test/test_uops.py index c5bf5e4183..cff71ff0c4 100644 --- a/test/test_uops.py +++ b/test/test_uops.py @@ -361,15 +361,16 @@ class TestAssembly(unittest.TestCase): self.assertIn(Ops.MUL, ops) def test_division_power_of_two(self): - g = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(), (), 0) - c = UOp(Ops.CONST, dtypes.uint, (), 2) - l = UOp(Ops.LOAD, dtypes.uint, (g.index(c),)) - a = UOp(Ops.IDIV, dtypes.uint, (l, c)) - uops = to_uops_list([a], opts=Device[Device.DEFAULT].renderer) - Device[Device.DEFAULT].renderer.render(uops) - ops = [x.op for x in uops] - self.assertIn(Ops.SHR, ops) - self.assertNotIn(Ops.IDIV, ops) + for dt in (dtypes.int32, dtypes.uint32): + g = UOp(Ops.DEFINE_GLOBAL, dt.ptr(), (), 0) + c = UOp(Ops.CONST, dt, (), 2) + l = UOp(Ops.LOAD, dt, (g.index(c),)) + a = UOp(Ops.IDIV, dt, (l, c)) + uops = to_uops_list([a], opts=Device[Device.DEFAULT].renderer) + Device[Device.DEFAULT].renderer.render(uops) + ops = [x.op for x in uops] + self.assertIn(Ops.SHR, ops, f"For dtype={dt} divison by power of two did not simplify to shift") + self.assertNotIn(Ops.IDIV, ops, f"For dtype={dt} divison by power of two did not simplify to shift") def test_fast_idiv_and_mod(self): g = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(), (), 0) diff --git a/tinygrad/codegen/devectorizer.py b/tinygrad/codegen/devectorizer.py index a6aae9d261..f59286a951 100644 --- a/tinygrad/codegen/devectorizer.py +++ b/tinygrad/codegen/devectorizer.py @@ -4,7 +4,7 @@ from collections import defaultdict from dataclasses import dataclass from tinygrad.device import is_dtype_supported from tinygrad.dtype import dtypes, ImageDType, PtrDType, promo_lattice, DType -from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, resolve, graph_rewrite, GroupOp, identity_element +from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, graph_rewrite, GroupOp, identity_element from tinygrad.uop.symbolic import split_uop, uop_given_valid, parse_valid, simplify_valid, sym, symbolic_flat from tinygrad.helpers import getenv, flatten, AMX, prod, partition from tinygrad.uop.transcendental import xexp2, xlog2, xsin, xpow, TRANSCENDENTAL_SUPPORTED_DTYPES @@ -175,9 +175,10 @@ def get_late_rewrite_patterns(ops, force_transcendental=False): # rewrite MUL/IDIV to SHL+SHR: x*(2**y) -> shl(x,y) and x//(2**y) -> shr(x,y) if Ops.SHL in ops: pat += [(UPat.var("x", dtypes.ints)*UPat.cvar("c"), lambda c,x: x << v if (v:=powers_of_two.get(c.arg, 0)) else None)] if Ops.SHR in ops: - # no reason to check x>=0 for uints + # no reason to check x<0 for uints pat += [(UPat.var("x", dtypes.uints)//UPat.cvar("c"), lambda x,c: x >> v if (v:=powers_of_two.get(c.arg, 0)) else None)] - pat += [(UPat.var("x", dtypes.sints)//UPat.cvar("c"), lambda x,c: x >> v if (v:=powers_of_two.get(c.arg, 0)) and resolve(x>=0,False) else None)] + pat += [(UPat.var("x", dtypes.ints)//UPat.cvar("c"), lambda x,c: (x+(l.const_like(l.vmin) if (l:=(x<0)).vmin==l.vmax else l).where( + c-1, 0)) >> v if (v:=powers_of_two.get(c.arg, 0)) else None)] # (x+(x<0).where(c-1, 0)) >> v if not getenv("DISABLE_FAST_IDIV"): pat += [(UPat.var("x", dtypes.ints)//UPat.cvar("d"), lambda ctx, x, d: fast_idiv(ctx, x, d.arg))] pat += [(UPat.var("x", dtypes.ints)%UPat.cvar("d"), lambda ctx, x, d: x - d*f if (f:=fast_idiv(ctx, x, d.arg)) is not None else None)] diff --git a/tinygrad/renderer/wgsl.py b/tinygrad/renderer/wgsl.py index 1d3147d0ac..3083e50d85 100644 --- a/tinygrad/renderer/wgsl.py +++ b/tinygrad/renderer/wgsl.py @@ -32,7 +32,8 @@ wgsl_matcher = PatternMatcher([ (UPat.load(UPat.var("b"), UPat.cvar("c"), name="l"),lambda l,b,c: packed_load(l,b,l.dtype,c.cast(dtypes.uint32)) if is_packed(l.dtype) else None), (UPat.load(UPat.var("b"), name='l', allow_any_len=True), lambda l,b: packed_load(l, b, l.dtype) if is_packed(l.dtype) else None), (UPat.store(UPat.var("bidx"), UPat.var("var"), allow_any_len=True), lambda bidx,var: packed_store(bidx,var) if is_packed(var.dtype) else None), - (UPat.var("a") << UPat.var("b"),lambda a,b:(a.bitcast(dtypes.uint32)<> UPat.var("y"), lambda x,y: UOp(Ops.SHR, x.dtype, (x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None), ]) + extra_pm class WGSLRenderer(CStyleLanguage):