From 0c4bfaeb486012fffe248405c17ddbc277f6ad32 Mon Sep 17 00:00:00 2001 From: nimlgen <138685161+nimlgen@users.noreply.github.com> Date: Fri, 31 Jul 2026 15:49:56 +0300 Subject: [PATCH] coalesce ints (#17323) * merge ints * fix z3 validation of coalesced loads * fix uint vector names in CUDA and Metal --- test/null/test_validate_oob.py | 8 ++++++++ tinygrad/codegen/late/coalesce.py | 2 +- tinygrad/renderer/cstyle.py | 6 +++--- tinygrad/uop/validate.py | 4 ++-- 4 files changed, 14 insertions(+), 6 deletions(-) diff --git a/test/null/test_validate_oob.py b/test/null/test_validate_oob.py index 4388f2b863..39e509770d 100644 --- a/test/null/test_validate_oob.py +++ b/test/null/test_validate_oob.py @@ -138,6 +138,14 @@ class TestValidateOOB(unittest.TestCase): with self.assertRaises(RuntimeError): to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 64))).load(dtype=dtypes.int)]) # oob + def test_load_from_shrink_as_index(self): + with Context(CHECK_OOB=1, SPEC=2): + buf0 = UOp.param(0, dtypes.int, (16,)) + buf1 = UOp.param(1, dtypes.int, (64,)) + shrink = UOp(Ops.SHRINK, src=(buf0, UOp.const(dtypes.int, 0), UOp.const(dtypes.weakint, 4))) + ld0 = shrink.load(dtype=dtypes.int).index(0) + to_uops_list([buf1.index(ld0.valid((ld0 >= 0) & (ld0 < 64))).load(dtype=dtypes.int)]) + def test_load_bool_as_mask(self): with Context(CHECK_OOB=1, SPEC=2): buf_bool = UOp.param(0, dtypes.bool, (16,)) diff --git a/tinygrad/codegen/late/coalesce.py b/tinygrad/codegen/late/coalesce.py index 8510f639bd..077b4a734b 100644 --- a/tinygrad/codegen/late/coalesce.py +++ b/tinygrad/codegen/late/coalesce.py @@ -127,7 +127,7 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp: if ctx is not None and ctx.target.device == "DSP": lengths = [128,64,32,16,8,4] must_divide = False - elif buf.dtype not in (dtypes.float, dtypes.half, *dtypes.fp8s) and not is_image_shape(buf._shape): + elif buf.dtype not in (dtypes.float, dtypes.half, dtypes.int, dtypes.uint, *dtypes.fp8s) and not is_image_shape(buf._shape): pass elif buf.addrspace == AddrSpace.REG: pass diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index ce6d6330de..1b0287b7de 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -356,7 +356,7 @@ class MetalRenderer(CStyleLanguage): code_for_workitem = {"g": lambda x: f"gid.{chr(120+int(x))}", "l": lambda x: f"lid.{chr(120+int(x))}"} # uint3 used for gid/lid - TODO: this should probably be `ushort3 lid [[thread_position_in_threadgroup]]` extra_args = ['uint3 gid [[threadgroup_position_in_grid]]', 'uint3 lid [[thread_position_in_threadgroup]]'] - type_map = {dtypes.bfloat16: "bfloat"} + type_map = {dtypes.uint32: "uint", dtypes.bfloat16: "bfloat"} # precise::sin code_for_op = {**CStyleLanguage.code_for_op, Ops.SIN: lambda x,dtype: f"precise::sin({x})"} @@ -420,7 +420,7 @@ class CUDARenderer(CStyleLanguage): Ops.EXP2: lambda x,dtype: f"hexp2({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"exp2({x})", Ops.SQRT: lambda x,dtype: f"hsqrt({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"sqrt({x})", Ops.RECIPROCAL: lambda x,dtype: f"hrcp({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"(1/{x})" } - type_map = {dtypes.bfloat16: "nv_bfloat16", dtypes.fp8e4m3: "__nv_fp8_e4m3", dtypes.fp8e5m2: "__nv_fp8_e5m2"} + type_map = {dtypes.uint32: "uint", dtypes.bfloat16: "nv_bfloat16", dtypes.fp8e4m3: "__nv_fp8_e4m3", dtypes.fp8e5m2: "__nv_fp8_e5m2"} extra_matcher = create_non_native_float_pats(dtypes.fp8s, casting=False) + PatternMatcher([ (UPat(Ops.CAST, dtypes.fp8s, UPat.var("x", dtypes.fp8s), name='y'), lambda x,y: x.cast(dtypes.float).cast(y.dtype) if x.dtype!=y.dtype else None), ]) @@ -436,7 +436,7 @@ class CUDARenderer(CStyleLanguage): def render_kernel(self, function_name, kernel, bufs, uops, prefix=None): # TODO: why is dtypes.bfloat16.name == "__bf16"? would be easier not override dtypes.name - prefix = ["#define INFINITY (__int_as_float(0x7f800000))", "#define NAN (__int_as_float(0x7fffffff))", + prefix = ["typedef unsigned int uint;", "#define INFINITY (__int_as_float(0x7f800000))", "#define NAN (__int_as_float(0x7fffffff))", "template __device__ __forceinline__ T tg_bitcast(F v) { union U { F f; T t; }; U u; u.f = v; return u.t; }"] used_dtypes = uops_to_dtypes(uops) if any(dt in dtypes.fp8s for dt, _ in used_dtypes): prefix.append("#include ") diff --git a/tinygrad/uop/validate.py b/tinygrad/uop/validate.py index bcd21c78f1..7c50b76ea4 100644 --- a/tinygrad/uop/validate.py +++ b/tinygrad/uop/validate.py @@ -59,8 +59,8 @@ z3_renderer = PatternMatcher([ ]) def uops_to_z3(solver:z3.Solver, *uops: UOp) -> list[z3.ExprRef]: - # gate on upstream AFTER/BUFFER, but keep INDEX as an unknown LOAD - lst = list(UOp.sink(*uops).toposort(gate=lambda x: x.op not in {Ops.AFTER, Ops.BUFFER} and \ + # gate on upstream memory addressing, but keep INDEX as an unknown LOAD + lst = list(UOp.sink(*uops).toposort(gate=lambda x: x.op not in {Ops.AFTER, Ops.BUFFER, Ops.SHRINK} and \ (x.dtype in dtypes.ints+(dtypes.bool, dtypes.weakint) or x.op is Ops.SINK)))[:-1] z3map: dict[UOp, z3.ExprRef] = {} for u in lst: