From 74411984fc1cd59c2dec46107daef9e93f98d718 Mon Sep 17 00:00:00 2001 From: Sieds Lykles <93992551+S-Lykles@users.noreply.github.com> Date: Fri, 26 Sep 2025 07:21:02 +0200 Subject: [PATCH] Rangeify IMAGE (#12304) * add imagedtype to rangeify * enable some image tests * move the tests * image upcast before locals * add if statement * rangeify image_dtype test * decrease read_image count --- .github/workflows/test.yml | 8 +++-- test/test_image_dtype.py | 4 +-- tinygrad/codegen/opt/heuristic.py | 53 ++++++++++++++++--------------- tinygrad/schedule/rangeify.py | 3 ++ 4 files changed, 38 insertions(+), 30 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 9254c27bd4..283ac7b81c 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -312,7 +312,7 @@ jobs: testopenclimage: name: CL IMAGE Tests runs-on: ubuntu-22.04 - timeout-minutes: 10 + timeout-minutes: 15 steps: - name: Checkout Code uses: actions/checkout@v4 @@ -326,6 +326,10 @@ jobs: run: | CL=1 IMAGE=2 python -m pytest -n=auto test/test_ops.py --durations=20 CL=1 IMAGE=2 python test/models/test_end2end.py TestEnd2End.test_linear_mnist + - name: Test CL IMAGE=2 ops + training (rangeify) + run: | + RANGEIFY=1 CL=1 IMAGE=2 python -m pytest -n=auto test/test_ops.py --durations=20 + RANGEIFY=1 CL=1 IMAGE=2 python test/models/test_end2end.py TestEnd2End.test_linear_mnist - name: Run process replay tests uses: ./.github/actions/process-replay @@ -370,7 +374,7 @@ jobs: llvm: 'true' - name: Test openpilot model kernel count and gate usage run: | - ALLOWED_KERNEL_COUNT=208 ALLOWED_READ_IMAGE=2175 ALLOWED_GATED_READ_IMAGE=16 FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx + ALLOWED_KERNEL_COUNT=208 ALLOWED_READ_IMAGE=2160 ALLOWED_GATED_READ_IMAGE=16 FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx - name: Test openpilot alt model correctness (float32) run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/3799fe46b3a629e491d4b8498b8ae83e4c88c304/selfdrive/modeld/models/supercombo.onnx - name: Test openpilot fastvits model correctness (float32) diff --git a/test/test_image_dtype.py b/test/test_image_dtype.py index 6adab73e51..394babf87a 100644 --- a/test/test_image_dtype.py +++ b/test/test_image_dtype.py @@ -4,7 +4,7 @@ from tinygrad import Device, dtypes, Tensor, Context from tinygrad.device import LRUAllocator, is_dtype_supported from tinygrad.dtype import ImageDType from tinygrad.engine.realize import lower_schedule -from tinygrad.helpers import prod, unwrap +from tinygrad.helpers import prod, unwrap, RANGEIFY from test.helpers import REAL_DEV IMAGE_SUPPORTED_DEVICES = ("QCOM", "CL") @@ -139,7 +139,7 @@ class TestImageDType(unittest.TestCase): # NOTE: the w1 grad must realize to a seperate kernel assert w1.grad.uop.is_realized, f"never realized {w1.grad}" self.assertEqual(w1.grad.uop.base.buffer.dtype, dtypes.float32) - self.assertEqual(len(sched), 10) + self.assertEqual(len(sched), 8 if RANGEIFY else 10) @unittest.skipUnless(REAL_DEV in IMAGE_SUPPORTED_DEVICES, "Images not supported") class TestImageRealization(unittest.TestCase): diff --git a/tinygrad/codegen/opt/heuristic.py b/tinygrad/codegen/opt/heuristic.py index a73234c4c7..c1c69ef498 100644 --- a/tinygrad/codegen/opt/heuristic.py +++ b/tinygrad/codegen/opt/heuristic.py @@ -48,32 +48,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler: # make a copy so it does not mutate the input k = k.copy() - # should use matvec - TODO: adjust/tune based on the wide vs tall/large vs small mat - MV_BLOCKSIZE, MV_THREADS_PER_ROW, MV_ROWS_PER_THREAD = getenv("MV_BLOCKSIZE", 4), getenv("MV_THREADS_PER_ROW", 8), getenv("MV_ROWS_PER_THREAD", 4) - if k.opts.has_local and getenv("MV",1) != 0 and (MV_BLOCKSIZE > 1 or MV_THREADS_PER_ROW > 1 or MV_ROWS_PER_THREAD > 1) and \ - k.reduceop is not None and k.reduceop.arg[0] is Ops.ADD and len(k.full_shape) >= 2 and k.opts.has_shared and \ - (mulop:=k.reduceop.src[0]).op is Ops.MUL and mulop.src[0].op is Ops.LOAD and mulop.src[1].op is Ops.LOAD: - idx0, idx1 = mulop.src[0].src[0].src[1].get_idx(), mulop.src[1].src[0].src[1].get_idx() - first_reduce_rng = k.ranges_of(AxisType.REDUCE)[0] - if any(u is first_reduce_rng for u in idx0.split_uop(Ops.ADD)) and all(r in idx1.ranges for r in idx0.ranges): - for global_idx in k.axes_of(AxisType.GLOBAL): - if first_reduce_rng.src[0].divides(MV_THREADS_PER_ROW) is not None and k.full_shape[global_idx]%(MV_BLOCKSIZE*MV_ROWS_PER_THREAD) == 0: - if DEBUG >= 3: - print(f"MATVEC: {k.full_shape=} {first_reduce_rng.render()} {MV_BLOCKSIZE=} {MV_THREADS_PER_ROW=} {MV_ROWS_PER_THREAD=}") - if MV_THREADS_PER_ROW > 1: k.apply_opt(Opt(OptOps.GROUP, 0, MV_THREADS_PER_ROW)) - if MV_BLOCKSIZE > 1: k.apply_opt(Opt(OptOps.LOCAL, global_idx, MV_BLOCKSIZE)) - if MV_ROWS_PER_THREAD > 1: k.apply_opt(Opt(OptOps.UPCAST, global_idx, MV_ROWS_PER_THREAD)) - return k - - # are we grouping? (requires local shape support) - if resolve(prod(k.output_shape[i] for i in k.upcastable_dims) <= 2048, False): - for sz in [16]: - try: - k.apply_opt(Opt(OptOps.GROUPTOP, 0, sz)) - break - except KernelOptError: pass - - # upcast float4 images + # upcast float4 images, this must be early so we don't accidentally add locals before the upcast for buf_index,buf in enumerate(k.bufs): if isinstance(buf.src[0].dtype, ImageDType): # part of real_strides @@ -85,6 +60,32 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler: elif axis in k.unrollable_dims: k.apply_opt(Opt(OptOps.UNROLL, k.unrollable_dims.index(axis), 4)) + # should use matvec - TODO: adjust/tune based on the wide vs tall/large vs small mat + MV_BLOCKSIZE, MV_THREADS_PER_ROW, MV_ROWS_PER_THREAD = getenv("MV_BLOCKSIZE", 4), getenv("MV_THREADS_PER_ROW", 8), getenv("MV_ROWS_PER_THREAD", 4) + if k.opts.has_local and getenv("MV",1) != 0 and (MV_BLOCKSIZE > 1 or MV_THREADS_PER_ROW > 1 or MV_ROWS_PER_THREAD > 1) and \ + k.reduceop is not None and k.reduceop.arg[0] is Ops.ADD and len(k.full_shape) >= 2 and k.opts.has_shared and \ + (mulop:=k.reduceop.src[0]).op is Ops.MUL and mulop.src[0].op is Ops.LOAD and mulop.src[1].op is Ops.LOAD: + idx0, idx1 = mulop.src[0].src[0].src[1].get_idx(), mulop.src[1].src[0].src[1].get_idx() + if k.ranges_of(AxisType.REDUCE): + first_reduce_rng = k.ranges_of(AxisType.REDUCE)[0] + if any(u is first_reduce_rng for u in idx0.split_uop(Ops.ADD)) and all(r in idx1.ranges for r in idx0.ranges): + for global_idx in k.axes_of(AxisType.GLOBAL): + if first_reduce_rng.src[0].divides(MV_THREADS_PER_ROW) is not None and k.full_shape[global_idx]%(MV_BLOCKSIZE*MV_ROWS_PER_THREAD) == 0: + if DEBUG >= 3: + print(f"MATVEC: {k.full_shape=} {first_reduce_rng.render()} {MV_BLOCKSIZE=} {MV_THREADS_PER_ROW=} {MV_ROWS_PER_THREAD=}") + if MV_THREADS_PER_ROW > 1: k.apply_opt(Opt(OptOps.GROUP, 0, MV_THREADS_PER_ROW)) + if MV_BLOCKSIZE > 1: k.apply_opt(Opt(OptOps.LOCAL, global_idx, MV_BLOCKSIZE)) + if MV_ROWS_PER_THREAD > 1: k.apply_opt(Opt(OptOps.UPCAST, global_idx, MV_ROWS_PER_THREAD)) + return k + + # are we grouping? (requires local shape support) + if resolve(prod(k.output_shape[i] for i in k.upcastable_dims) <= 2048, False): + for sz in [16]: + try: + k.apply_opt(Opt(OptOps.GROUPTOP, 0, sz)) + break + except KernelOptError: pass + # no more opt if we are grouping if k.group_for_reduces: return k diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index dacf9f4f52..e834108cee 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -388,6 +388,9 @@ def pre_bufferize(b:UOp, x:UOp, copy:UOp): pm_cleanups = double_reshape+pm_mops+PatternMatcher([ #(UPat(Ops.BUFFERIZE, name="b"), cleanup_dead_axes), + (UPat(GroupOp.All-{Ops.BUFFERIZE, Ops.BUFFER}, name="x"), lambda x: x.replace(dtype=x.dtype.base) if isinstance(x.dtype, ImageDType) else None), + (UPat((Ops.BUFFERIZE), name="x"), lambda x: x.replace(dtype=x.dtype.base) if isinstance(x.dtype, ImageDType) + and (resolve(prod(x.dtype.shape)!=prod(x.shape)) or x.shape[-1]%4!=0) else None), # remove noop buffers. if we look at the next index we can remove even more of these # NOTE: this is mostly the same case as below, but if there's no INDEX this gets more (UPat(Ops.INDEX, name="idx").f(Ops.BUFFERIZE, allow_any_len=True, name="b2"),