diff --git a/openpilot/starpilot/models/tests/test_modeld_jit_copies.py b/openpilot/starpilot/models/tests/test_modeld_jit_copies.py new file mode 100644 index 0000000000..e91d7d6f2f --- /dev/null +++ b/openpilot/starpilot/models/tests/test_modeld_jit_copies.py @@ -0,0 +1,20 @@ +"""The production modern copy helper must see fresh NumPy storage on every call.""" +import __future__ +import ast +from pathlib import Path + +import numpy as np + + +def test_modern_copy_jit_reloads_mutated_host_inputs(): + from tinygrad.tensor import Tensor + path = Path(__file__).resolve().parents[3] / 'selfdrive/modeld/modeld.py' + function = next(n for n in ast.parse(path.read_text()).body if isinstance(n, ast.FunctionDef) and n.name == 'compile_copy') + namespace = {} + exec(compile(ast.Module(body=[function], type_ignores=[]), str(path), 'exec', flags=__future__.annotations.compiler_flag), namespace) + source = np.arange(32, dtype=np.float32) + tensor = Tensor(source, device='NPY').to('CPU') + copy = namespace['compile_copy'](tensor) + for iteration in range(5): + source[:] = np.arange(32, dtype=np.float32) + iteration * 0.25 + np.testing.assert_array_equal(copy().numpy(), source) diff --git a/openpilot/starpilot/models/tests/test_stock_model_error_state.py b/openpilot/starpilot/models/tests/test_stock_model_error_state.py index fcff7d3e1c..ff248f61cc 100644 --- a/openpilot/starpilot/models/tests/test_stock_model_error_state.py +++ b/openpilot/starpilot/models/tests/test_stock_model_error_state.py @@ -105,6 +105,8 @@ class TestStockModelErrorState(unittest.TestCase): model.run_warp, model.run_model = run_warp, run_model model.outputs = {'outputs': Output()} + model.copy_inputs = lambda: model.input_device.copy_from(model.input_host) + model.copy_outputs = lambda: model.outputs['outputs'] model.output_slices = {'fixture': slice(0, 2)} parsed = {'fixture': np.array([[1., 2.]], dtype=np.float32)} diff --git a/tinygrad_repo/.github/actions/setup-tinygrad/action.yml b/tinygrad_repo/.github/actions/setup-tinygrad/action.yml index 5a14bcacf0..06c96c680a 100644 --- a/tinygrad_repo/.github/actions/setup-tinygrad/action.yml +++ b/tinygrad_repo/.github/actions/setup-tinygrad/action.yml @@ -74,17 +74,22 @@ runs: with: enable-cache: 'false' # see below for manual caching - # **** Caching packages **** + # **** Caching **** + # NOTE: self-hosted CI (ie. collaborator PRs) should never cache anything, instead, either: + # * bake into the CI image: https://git.tinygrad.win/tinygrad/images, OR + # * handle in CI runner HTTP cache + # Because CI for master and non-collaborator PRs are github-hosted, they need caching to workaround the flaky internet. + # Additionally, the actions/cache/restore pattern should be used to avoid PRs polluting the limited github cache storage. - name: Cache Python packages (PR) - if: github.event_name == 'pull_request' && env.RUNNER_ENVIRONMENT == 'github-hosted' + if: github.event_name == 'pull_request' && runner.environment == 'github-hosted' id: restore-venv-pr uses: actions/cache/restore@v5 with: path: /tmp/.uv-cache key: uv-${{ runner.os }}-${{ runner.arch }}-python-${{ inputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }} - name: Cache Python packages - if: github.event_name != 'pull_request' + if: github.event_name != 'pull_request' && runner.environment == 'github-hosted' id: restore-venv uses: actions/cache@v5 with: @@ -94,13 +99,13 @@ runs: # **** Caching downloads **** - name: Cache downloads (PR) - if: inputs.key != '' && github.event_name == 'pull_request' && env.RUNNER_ENVIRONMENT == 'github-hosted' + if: inputs.key != '' && github.event_name == 'pull_request' && runner.environment == 'github-hosted' uses: actions/cache/restore@v5 with: path: ${{ runner.os == 'Linux' && '~/.cache/tinygrad/downloads/' || '~/Library/Caches/tinygrad/downloads/' }} key: downloads-${{ github.job }}-${{ inputs.key }}-${{ env.CACHE_VERSION }} - name: Cache downloads - if: inputs.key != '' && github.event_name != 'pull_request' + if: inputs.key != '' && github.event_name != 'pull_request' && runner.environment == 'github-hosted' uses: actions/cache@v5 with: path: ${{ runner.os == 'Linux' && '~/.cache/tinygrad/downloads/' || '~/Library/Caches/tinygrad/downloads/' }} @@ -205,13 +210,13 @@ runs: echo "installed=$installed" >> "$GITHUB_OUTPUT" - name: Cache apt (PR) - if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true' || inputs.autogen == 'true') && github.event_name == 'pull_request' && steps.apt-pkgs.outputs.installed == 'false' && env.RUNNER_ENVIRONMENT == 'github-hosted' + if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true' || inputs.autogen == 'true') && github.event_name == 'pull_request' && steps.apt-pkgs.outputs.installed == 'false' && runner.environment == 'github-hosted' uses: actions/cache/restore@v5 with: path: /var/cache/apt/archives/ key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }} - name: Cache apt - if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true' || inputs.autogen == 'true') && github.event_name != 'pull_request' && steps.apt-pkgs.outputs.installed == 'false' + if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true' || inputs.autogen == 'true') && github.event_name != 'pull_request' && steps.apt-pkgs.outputs.installed == 'false' && runner.environment == 'github-hosted' uses: actions/cache@v5 with: path: /var/cache/apt/archives/ diff --git a/tinygrad_repo/.github/workflows/autogen.yml b/tinygrad_repo/.github/workflows/autogen.yml index 262d44d971..d7cd209698 100644 --- a/tinygrad_repo/.github/workflows/autogen.yml +++ b/tinygrad_repo/.github/workflows/autogen.yml @@ -1,7 +1,7 @@ name: Autogen env: # increment this when downloads substantially change to avoid the internet - CACHE_VERSION: '14' + CACHE_VERSION: '15' CAPTURE_PROCESS_REPLAY: 1 GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} PYTHONPATH: ${{ github.workspace }} @@ -53,7 +53,7 @@ jobs: python3 -c "from tinygrad.runtime.autogen import mesa" python3 -c "from tinygrad.runtime.autogen import avcodec" python3 -c "from tinygrad.runtime.autogen import llvm_qcom" - python3 -c "from tinygrad.runtime.autogen import mlx5, bnxt" + python3 -c "from tinygrad.runtime.autogen import mlx5, bnxt, nvme" python3 -c "from tinygrad.runtime.autogen import ggml_common" REGEN=1 python3 -c "from tinygrad.runtime.autogen import libclang" - name: Check for differences diff --git a/tinygrad_repo/.github/workflows/test.yml b/tinygrad_repo/.github/workflows/test.yml index b8542bcff0..d28688c9ba 100644 --- a/tinygrad_repo/.github/workflows/test.yml +++ b/tinygrad_repo/.github/workflows/test.yml @@ -1,7 +1,7 @@ name: Unit Tests env: # increment this when downloads substantially change to avoid the internet - CACHE_VERSION: '20' + CACHE_VERSION: '21' CAPTURE_PROCESS_REPLAY: ${{ github.event_name == 'pull_request' && contains(github.event.pull_request.title, '[pr]') && '1' || '0' }} GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} PYTHONPATH: ${{ github.workspace }} diff --git a/tinygrad_repo/extra/export_model.py b/tinygrad_repo/extra/export_model.py index 2b0574977c..76f1a837f2 100644 --- a/tinygrad_repo/extra/export_model.py +++ b/tinygrad_repo/extra/export_model.py @@ -66,13 +66,16 @@ def export_model_clang(functions:Dict[str,str], statements:Dict[str,Tuple[str,in inputs = [(name, dtype_map[bufs[name][1]], bufs[name][0]) for name in input_names + list(symbolic_vars.values())] outputs = [(name, dtype_map[bufs[name][1]], bufs[name][0]) for name in output_names] forward_args = ",".join(f"{dtype}{'*' if name not in symbolic_vars.values() else ''} {name}" for name,dtype,_ in (outputs+inputs if wasm else inputs+outputs)) + def kernel_call(name:str, args:List[str]) -> str: + fields = " ".join(f"{dtype_map[bufs[a][1]] if a in symbolic_vars.values() else 'void*'} a{i};" for i,a in enumerate(args)) + return f"{name}((unsigned char*)&(struct {{ {fields} }}){{{', '.join(args)}}});" if not wasm: for name,cl in bufs_to_save.items(): weight = ''.join(["\\x%02X"%x for x in cl.as_memoryview()]) cprog.append(f"unsigned char {name}_data[] = \"{weight}\";") cprog += [f"{dtype_map[dtype]} {name}[{len}];" if name not in bufs_to_save else f"{dtype_map[dtype]} *{name} = ({dtype_map[dtype]} *){name}_data;" for name,(len,dtype,_key) in bufs.items() if name not in input_names+output_names] - cprog += [f"void net({forward_args}) {{"] + [f"{name}(&(struct {name}_args){{{', '.join(args)}}});" for (name, args, _global_size, _local_size) in statements] + ["}"] + cprog += [f"void net({forward_args}) {{"] + [kernel_call(name, args) for (name, args, _global_size, _local_size) in statements] + ["}"] return '\n'.join(headers + cprog) else: if bufs_to_save: @@ -88,7 +91,7 @@ def export_model_clang(functions:Dict[str,str], statements:Dict[str,Tuple[str,in cprog += [f"void net({forward_args})"] + ["{"] get_weight_ptr = lambda x: f"({dtype_map[bufs_to_save[x][1]]} *)bufs[{buf_to_name[x]['idx']}]" if x in bufs_to_save else x - cprog += [f" {name}(&(struct {name}_args){{{', '.join(map(get_weight_ptr, args))}}});" for (name, args, _global_size, _local_size) in statements] + ["}"] + cprog += [" " + kernel_call(name, list(map(get_weight_ptr, args))) for (name, args, _global_size, _local_size) in statements] + ["}"] weightMapping = "" if not bufs_to_save else f"""\nconst weightNames = [{", ".join([f'"{weight_name}"' for weight_name in [v["name"] for v in buf_to_name.values()]])}]; const {model_name}_name_to_id = Object.fromEntries(weightNames.map((name, index) => [name, index]));\n""" top = f"""import {model_name}Module from './{model_name}.js'{weightMapping}""" diff --git a/tinygrad_repo/extra/nvme_driver/bench_nvme.py b/tinygrad_repo/extra/nvme_driver/bench_nvme.py new file mode 100644 index 0000000000..c499e6b686 --- /dev/null +++ b/tinygrad_repo/extra/nvme_driver/bench_nvme.py @@ -0,0 +1,23 @@ +#!/usr/bin/env python3 +# GPUS=4 DRIVES=4 GB=2 python3 extra/nvme_driver/bench_nvme.py: each gpu reads GB GiB from its DRIVES drives at once. never writes +import time +from tinygrad import Tensor, Device, TinyJit, dtypes +from tinygrad.helpers import getenv + +GPUS, DRIVES, N = getenv("GPUS", 4), getenv("DRIVES", 4), getenv("GB", 2) << 30 +gpus = [f"AMD:{g}" for g in range(GPUS)] +srcs = [Tensor.empty(sz:=Device[f"NVME:{d}"].nvme.size, dtype=dtypes.uint8, device=f"NVME:{d}")[sz // 2:sz // 2 + N] for d in range(GPUS * DRIVES)] +dsts = [Tensor.empty(N, dtype=dtypes.uint8, device=gpus[d // DRIVES]).realize() for d in range(GPUS * DRIVES)] +@TinyJit +def read() -> list[Tensor]: return [dst.assign(src.to(dst.device)).realize() for dst, src in zip(dsts, srcs)] +def run(n:int) -> float: + for g in gpus: Device[g].synchronize() + st = time.perf_counter() + for _ in range(n): read() + for g in gpus: Device[g].synchronize() + return time.perf_counter() - st +run(3) +for n in (1, getenv("ITERS", 8)): + dt = run(n) + print(f"{len(srcs)} drives -> {GPUS} gpus, {n} x {len(srcs) * N / 2**30:.0f} GiB in {dt * 1e3:.0f} ms: " + f"{n * len(srcs) * N / dt / 1e9:.1f} GB/s, {n * DRIVES * N / dt / 1e9:.1f} GB/s a gpu") diff --git a/tinygrad_repo/extra/nvme_driver/test_nvme.py b/tinygrad_repo/extra/nvme_driver/test_nvme.py new file mode 100644 index 0000000000..0fbe809357 --- /dev/null +++ b/tinygrad_repo/extra/nvme_driver/test_nvme.py @@ -0,0 +1,28 @@ +#!/usr/bin/env python3 +# NVME=0 GPU=AMD python3 extra/nvme_driver/test_nvme.py: reads of the last 256 MiB of a drive (which must hold data) into gpu memory. never writes +import unittest +import numpy as np +from tinygrad import Tensor, Device, TinyJit, dtypes +from tinygrad.helpers import getenv + +NVME, GPU, N = f"NVME:{getenv('NVME', 0)}", getenv("GPU", "AMD"), 256 << 20 + +class TestNVME(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.region = Tensor.empty(size:=Device[NVME].nvme.size, dtype=dtypes.uint8, device=NVME)[size - N:] + cls.ref = cls.region.to(GPU).numpy() + assert np.count_nonzero(np.bincount(cls.ref)) > 64, "the last 256 MiB of the drive hold no data to compare" + def check(self, reads:list[tuple[int, int]]): + Tensor.realize(*(ts:=[self.region[off:off + n].to(GPU) for off, n in reads])) + for t, (off, n) in zip(ts, reads): np.testing.assert_equal(t.numpy(), self.ref[off:off + n], err_msg=f"{n} bytes at {off}") + + def test_reads(self): + for read in [(0, 4096), (4096, 8192), (0, 3 * 4096), (0, 1 << 20), (12288, (1 << 20) + 4096), (4096, N - 8192)]: self.check([read]) + def test_reads_one_submit(self): self.check([(0, 1 << 20), (8 << 20, 3 << 20)]) + def test_jit(self): + @TinyJit + def f(t:Tensor) -> Tensor: return (t.to(GPU) + 1).realize() + for _ in range(5): np.testing.assert_equal(f(self.region[1 << 20:3 << 20]).numpy(), self.ref[1 << 20:3 << 20] + 1) + +if __name__ == "__main__": unittest.main() diff --git a/tinygrad_repo/test/amd/hw/helpers.py b/tinygrad_repo/test/amd/hw/helpers.py index 3477494544..67ef758f5c 100644 --- a/tinygrad_repo/test/amd/hw/helpers.py +++ b/tinygrad_repo/test/amd/hw/helpers.py @@ -136,6 +136,7 @@ def get_prologue_epilogue(n_lanes: int) -> tuple[list, list]: for i in range(N_SGPRS): prologue.append(s_mov_b32(s[i], 0)) prologue.append(s_mov_b32(VCC_LO, 0)) + prologue.append(s_cmp_eq_u32(0, 1)) # initialize SCC independently of the emulator/runtime's entry state epilogue = [ s_mov_b32(s[90], VCC_LO), diff --git a/tinygrad_repo/test/amd/hw/test_cdna_sdwa.py b/tinygrad_repo/test/amd/hw/test_cdna_sdwa.py index 7901ba8a0f..78b0c44f50 100644 --- a/tinygrad_repo/test/amd/hw/test_cdna_sdwa.py +++ b/tinygrad_repo/test/amd/hw/test_cdna_sdwa.py @@ -1,11 +1,12 @@ #!/usr/bin/env python3 import unittest import tinygrad.runtime.autogen.amd.cdna.ins as cdna -from test.amd.hw.test_cdna_vop3 import run_cdna +from test.amd.hw.test_cdna_vop3 import run_cdna, run_cdna_lanes class TestCDNASDWA(unittest.TestCase): def test_v_add_co_u32_e32_writes_vcc(self): out = run_cdna([ + cdna.s_mov_b64(cdna.VCC, 0), cdna.s_mov_b32(cdna.s[0], 0xffffffff), cdna.v_mov_b32_e32(cdna.v[0], cdna.s[0]), cdna.v_mov_b32_e32(cdna.v[13], 1), @@ -17,5 +18,17 @@ class TestCDNASDWA(unittest.TestCase): ]) self.assertEqual(out, 0x80000000) + def test_v_add_co_u32_e32_vcc_multiple_lanes(self): + out = run_cdna_lanes([ + cdna.s_mov_b32(cdna.VCC_LO, 0), + cdna.s_mov_b32(cdna.s[0], 0xffffffff), + cdna.v_mov_b32_e32(cdna.v[0], cdna.s[0]), + cdna.v_mov_b32_e32(cdna.v[13], 1), + cdna.v_add_co_u32_e32(cdna.v[0], cdna.SDWA, cdna.v[13], vsrc0=cdna.v[0], dst_sel=6, src0_sel=6), + cdna.s_nop(1), + cdna.v_mov_b32_e32(cdna.v[2], cdna.VCC_LO), + ], lanes=2) + self.assertEqual(out, [3, 3]) + if __name__ == "__main__": unittest.main() diff --git a/tinygrad_repo/test/amd/hw/test_sop.py b/tinygrad_repo/test/amd/hw/test_sop.py index f1bda8059a..c50f1cde60 100644 --- a/tinygrad_repo/test/amd/hw/test_sop.py +++ b/tinygrad_repo/test/amd/hw/test_sop.py @@ -267,6 +267,12 @@ class Test64BitLiterals(unittest.TestCase): class TestSCCBehavior(unittest.TestCase): """Tests for SCC condition code behavior.""" + def test_scc_initial(self): + """The shared prologue clears SCC before the test instructions run.""" + st = run_program([s_cselect_b32(s[0], 1, 0)], n_lanes=1) + self.assertEqual(st.sgpr[0], 0) + self.assertEqual(st.scc, 0) + def test_scc_from_s_cmp(self): """SCC should be set by scalar compare.""" instructions = [ diff --git a/tinygrad_repo/test/amd/hw/test_vop1.py b/tinygrad_repo/test/amd/hw/test_vop1.py index f1767b9509..1cb7c99319 100644 --- a/tinygrad_repo/test/amd/hw/test_vop1.py +++ b/tinygrad_repo/test/amd/hw/test_vop1.py @@ -1258,7 +1258,7 @@ class TestF16EdgeCases(unittest.TestCase): class TestExpEdgeCases(unittest.TestCase): - """Additional edge cases for V_EXP_F32.""" + """Tests for V_EXP_F32 and V_EXP_F16.""" def test_v_exp_f32_zero(self): """V_EXP_F32 of 0.0 returns 1.0 (2^0 = 1).""" @@ -1287,6 +1287,19 @@ class TestExpEdgeCases(unittest.TestCase): st = run_program(instructions, n_lanes=1) self.assertAlmostEqual(i2f(st.vgpr[0][1]), 0.5, places=5) + def test_v_exp_f16(self): + """V_EXP_F16 computes 2^x in half precision: 2^3 = 8, 2^-1 = 0.5.""" + instructions = [ + s_mov_b32(s[0], f32_to_f16(3.0)), + s_mov_b32(s[1], f32_to_f16(-1.0)), + v_mov_b32_e32(v[0], s[0]), + v_mov_b32_e32(v[2], s[1]), + v_exp_f16_e32(v[1], v[0]), + v_exp_f16_e32(v[3], v[2]), + ] + st = run_program(instructions, n_lanes=1) + self.assertEqual((f16(st.vgpr[0][1]), f16(st.vgpr[0][3])), (8.0, 0.5)) + def test_v_exp_f32_inf(self): """V_EXP_F32 of +inf returns +inf.""" import math diff --git a/tinygrad_repo/test/amd/hw/test_vop3.py b/tinygrad_repo/test/amd/hw/test_vop3.py index 8c2f172bf2..0212d2c3ed 100644 --- a/tinygrad_repo/test/amd/hw/test_vop3.py +++ b/tinygrad_repo/test/amd/hw/test_vop3.py @@ -4,7 +4,6 @@ Includes: v_fma_f32, v_div_scale_f32, v_div_fmas_f32, v_div_fixup_f32, v_alignbit_b32, v_bfe_i32, v_mad_u64_u32, v_readlane_b32, v_writelane_b32 """ import unittest -from tinygrad.helpers import OSX from test.amd.hw.helpers import * class TestFMA(unittest.TestCase): @@ -3721,7 +3720,6 @@ class TestClampLdExpRegressions(unittest.TestCase): class TestNaNPropagationRegressions(unittest.TestCase): """Regression tests: float arithmetic propagates a NaN from the FIRST NaN operand, quieted with its own sign/payload.""" - @unittest.skipIf(OSX, "broken on mac, TODO: why?") def test_mul_nan_priority(self): # first NaN operand wins (sign+payload), not x86's second-source propagation for a, b, want in [(0x7FC00001, 0x7F800003, 0x7FC00001), (0xFFC00005, 0x7F800003, 0xFFC00005), diff --git a/tinygrad_repo/test/device/dsp/test_quantize_onnx.py b/tinygrad_repo/test/device/dsp/test_quantize_onnx.py index 057ecad32e..23969674e8 100644 --- a/tinygrad_repo/test/device/dsp/test_quantize_onnx.py +++ b/tinygrad_repo/test/device/dsp/test_quantize_onnx.py @@ -115,11 +115,9 @@ class TestQuantizeOnnx(unittest.TestCase): unsigned_char128 hi128; }; }; - struct fxn_args { unsigned char* data0; unsigned char* data1; signed char* data2; }; - __attribute__((noinline)) void fxn(const struct fxn_args* args) { - unsigned char* restrict __attribute__((align_value(128))) data0 = args->data0; - unsigned char* restrict __attribute__((align_value(128))) data1 = args->data1; - signed char* restrict __attribute__((align_value(128))) data2 = args->data2; + __attribute__((noinline)) void fxn(unsigned char* restrict __attribute__((align_value(128))) data0, + unsigned char* restrict __attribute__((align_value(128))) data1, + signed char* restrict __attribute__((align_value(128))) data2) { for (int ridx0 = 0; ridx0 < 512; ridx0++) { int alu0 = (ridx0<<9); for (int ridx1 = 0; ridx1 < 4; ridx1++) { diff --git a/tinygrad_repo/test/external/external_test_usb_asm24.py b/tinygrad_repo/test/external/external_test_usb_asm24.py index 3fc98762e6..984a7cc444 100644 --- a/tinygrad_repo/test/external/external_test_usb_asm24.py +++ b/tinygrad_repo/test/external/external_test_usb_asm24.py @@ -1,7 +1,13 @@ import unittest -from tinygrad.helpers import Timing, getenv -from tinygrad import Tensor, Device, TinyJit +from tinygrad.helpers import Timing, getenv, unwrap, wait_cond +from tinygrad import Tensor, Device, TinyJit, dtypes +from tinygrad.device import Buffer +from tinygrad.engine.realize import run_linear +from tinygrad.uop.ops import UOp +from tinygrad.runtime.ops_amd import AMDSDMAQueue +from tinygrad.runtime.support.hcq2 import make_submit from tinygrad.runtime.support.usb import HALF, CHUNK, SLOT +from test.null.test_hcq2 import lower_hcq import numpy as np class USBTestCase(unittest.TestCase): @@ -81,10 +87,26 @@ class TestUSBIntegrity(USBTestCase): np.testing.assert_array_equal(a, (Tensor(a, device="NPY").to(Device.DEFAULT) * 1).numpy()) for size in (31, CHUNK + 31, 2 * CHUNK + 31): self.roundtrip(np.tile(tags, size // 32 + 1).view(np.uint8)[:size]) - def testRingWrap(self): # 64MB of chunks: the sdma ring (1MB on usb) wraps within the copy + def testRingWrap(self): + self.dev.synchronize() + q = unwrap(self.dev.sdma_queue(0)) + put, read = [b.host.view(fmt='Q') for b in (q.put_value, q.read_ptr)] + ring_size = q.ring.nbytes + # advance with NOPs, leave one dword so the next eager copy forces wrapping + target = (put[0] // ring_size + 1) * ring_size - 4 + if nbytes := target - put[0]: + nops = UOp.from_buffer(Buffer("CPU", nbytes, initial_value=bytes(nbytes)), dtypes.uint8) + hq = AMDSDMAQueue(make_submit(dev=self.dev.device, queue="COPY:0", deps=(nops,))) + run_linear(lower_hcq(hq.submit(nops))) + wait_cond(lambda: read[0], value=target) + self.assertEqual(put[0], target) + a = self.rng.integers(0, 256, 64 << 20, dtype=np.uint8) t = Tensor(a, device="NPY").to(Device.DEFAULT).realize() + self.assertGreater(put[0] // ring_size, target // ring_size) np.testing.assert_array_equal(a, t.numpy()) + self.dev.synchronize() + wait_cond(lambda: read[0], value=put[0]) if __name__ == "__main__": unittest.main() diff --git a/tinygrad_repo/test/mockgpu/amd/amdgpu.py b/tinygrad_repo/test/mockgpu/amd/amdgpu.py index 93ec1273c6..eb0a6251b1 100644 --- a/tinygrad_repo/test/mockgpu/amd/amdgpu.py +++ b/tinygrad_repo/test/mockgpu/amd/amdgpu.py @@ -193,12 +193,6 @@ class PM4Executor(AMDQueue): args_addr = self.gpu.regs[regCOMPUTE_USER_DATA_0] + (self.gpu.regs[regCOMPUTE_USER_DATA_0 + 1] << 32) lc = [self.gpu.regs[i] for i in range(regCOMPUTE_NUM_THREAD_X, regCOMPUTE_NUM_THREAD_X+3)] rsrc2 = self.gpu.regs[regCOMPUTE_PGM_RSRC2] - # Read all user data registers (hardware loads these directly into s[0:N]) - user_sgpr_count = (rsrc2 >> 1) & 0x1F # USER_SGPR_COUNT is bits 1:5 - user_data = [] - for i in range(user_sgpr_count): - try: user_data.append(self.gpu.regs[regCOMPUTE_USER_DATA_0 + i]) - except KeyError: user_data.append(0) prg_sz = 0 for st,sz in sorted(self.gpu.mapped_ranges): @@ -216,7 +210,7 @@ class PM4Executor(AMDQueue): assert 0 < self.kernel_text_size <= prg_sz, "Invalid prg ptr or text size" from test.mockgpu.amd.emu import run_asm err = run_asm(prg_addr, self.kernel_text_size, gl[0], gl[1], gl[2], lc[0], lc[1], lc[2], args_addr, - rsrc2, scratch_size, self.gpu.arch, user_data, self.kernel_name) + rsrc2, scratch_size, self.gpu.arch, self.kernel_name) if err != 0: raise RuntimeError("remu does not support the new instruction introduced in this kernel") def _exec_indirect_buffer(self, n): diff --git a/tinygrad_repo/test/mockgpu/amd/emu.py b/tinygrad_repo/test/mockgpu/amd/emu.py index 2010e01b0e..92b46874b1 100644 --- a/tinygrad_repo/test/mockgpu/amd/emu.py +++ b/tinygrad_repo/test/mockgpu/amd/emu.py @@ -1,11 +1,11 @@ # RDNA3 emulator v2 - compiles pcode to UOps executed via tinygrad CPU backend # Each instruction is compiled to a kernel that operates on buffers: -# arg=0: sgpr - sgpr[0-127], inline constants[128-255], PC_LO=256, PC_HI=257, SCC=258, SCRATCH_STRIDE=259 +# arg=0: sgpr - sgpr[0-127], inline constants[128-255], SCC=258 # arg=1: vgpr - vgpr[reg * 32 + lane] # arg=2: vmem - base address 0, INDEX offsets directly to host memory # arg=3: lds - local data share # arg=4: scratch - per-lane scratch memory -import ctypes, functools, itertools, re, struct, platform, subprocess, tempfile, traceback +import ctypes, functools, itertools, re, platform, subprocess, tempfile, traceback from typing import Callable # Set/restore DAZ+FTZ (denormals-are-zero + flush-to-zero) to match RDNA3 default float mode @@ -50,7 +50,7 @@ class _MXCSRContext: if lib is None or not hasattr(self, '_saved'): return lib.set_fpcr(self._saved) -from tinygrad.uop.ops import UOp, sint, Ops, KernelInfo, uopfunc, AxisType, rewrite_group +from tinygrad.uop.ops import UOp, Ops, KernelInfo, uopfunc, AxisType, rewrite_group from tinygrad.dtype import dtypes, AddrSpace from tinygrad.device import Buffer, BufferSpec, Device from tinygrad.engine.realize import run_linear @@ -75,12 +75,12 @@ def _c(val, dtype=dtypes.uint32): return UOp.const(val, dtype) def _u64(lo: UOp, hi: UOp) -> UOp: """Combine two 32-bit UOps into a 64-bit UOp.""" - return lo.cast(dtypes.uint64) | (hi.cast(dtypes.uint64) << UOp.const(32, dtypes.uint64)) + return lo.cast(dtypes.uint64) | (hi.cast(dtypes.uint64) << _c(32, dtypes.uint64)) def _split64(val: UOp) -> tuple[UOp, UOp]: """Split a 64-bit value into (lo, hi) 32-bit values.""" - v64 = val.bitcast(dtypes.uint64) if val.dtype == dtypes.float64 else val.cast(dtypes.uint64) if val.dtype != dtypes.uint64 else val - return v64.cast(dtypes.uint32), (v64 >> UOp.const(32, dtypes.uint64)).cast(dtypes.uint32) + v64 = val.bitcast(dtypes.uint64) if val.dtype == dtypes.float64 else val.cast(dtypes.uint64) + return v64.cast(dtypes.uint32), (v64 >> _c(32, dtypes.uint64)).cast(dtypes.uint32) _SRC_MOD_TYPES = {16: (dtypes.uint16, dtypes.half, 0x7FFF), 32: (dtypes.uint32, dtypes.float32, 0x7FFFFFFF), 64: (dtypes.uint64, dtypes.float64, 0x7FFFFFFFFFFFFFFF)} @@ -89,9 +89,9 @@ def _apply_src_mods(val: UOp, mod_bit: int, abs_bits: int, neg_bits: int, bits: if not (abs_bits & (1 << mod_bit)) and not (neg_bits & (1 << mod_bit)): return val ut, ft, mask = _SRC_MOD_TYPES[bits] fv = val.cast(ut).bitcast(ft) if bits == 16 else val.bitcast(ft) if val.dtype == ut else val - if abs_bits & (1 << mod_bit): fv = (fv.bitcast(ut) & UOp.const(mask, ut)).bitcast(ft) + if abs_bits & (1 << mod_bit): fv = (fv.bitcast(ut) & _c(mask, ut)).bitcast(ft) # neg modifier is a pure sign-bit toggle (preserves NaN payloads), not an arithmetic negate - if neg_bits & (1 << mod_bit): fv = (fv.bitcast(ut) ^ UOp.const((mask + 1) & (1 << (bits - 1)), ut)).bitcast(ft) + if neg_bits & (1 << mod_bit): fv = (fv.bitcast(ut) ^ _c((mask + 1) & (1 << (bits - 1)), ut)).bitcast(ft) return fv.bitcast(ut).cast(dtypes.uint32) if bits == 16 else fv.bitcast(ut) # Map VOPD ops to VOP2/VOP1 ops for pcode lookup (both RDNA3 and RDNA4 share these targets) @@ -113,17 +113,15 @@ def _iattr(inst, name: str, default: int = 0) -> int: """Optional integer attribute of a decoded instruction (None/missing -> default).""" v = getattr(inst, name, default) return default if v is None else v -# Special registers stored after inline constants (256-259) -SCRATCH_STRIDE_IDX = 259 -# SGPR buffer: 0-127 = SGPRs, 128-255 = inline constants, 256-259 = special registers -SGPR_COUNT = 260 +# SGPR buffer: 0-127 = SGPRs, 128-255 = inline constants +SGPR_COUNT = 256 def _op_name(inst) -> str: if hasattr(inst, 'opx'): return f"{inst.opx.name}_{inst.opy.name}" # VOPD has opx/opy not op - return inst.op.name if hasattr(inst.op, 'name') else str(inst.op) + return inst.op.name def _lane_active(exec_mask: UOp, lane: UOp) -> UOp: - if exec_mask.dtype == dtypes.uint64: return ((exec_mask >> lane.cast(dtypes.uint64)) & UOp.const(1, dtypes.uint64)).ne(UOp.const(0, dtypes.uint64)) + if exec_mask.dtype == dtypes.uint64: return ((exec_mask >> lane.cast(dtypes.uint64)) & _c(1, dtypes.uint64)).ne(_c(0, dtypes.uint64)) return ((exec_mask >> lane.cast(dtypes.uint32)) & _c(1)).ne(_c(0)) def _hi16(v: UOp) -> UOp: return (v >> _c(16)) & _c(0xFFFF) def _cond(cond, if_true, if_false): @@ -134,7 +132,6 @@ def _apply_opsel(val: UOp, sel_bit: int, opsel: int) -> UOp: return _hi16(val) i def _val_to_u32(val: UOp) -> UOp: """Convert any value to uint32 for storage (bitcast floats, cast ints).""" - if val.dtype == dtypes.uint32: return val if val.dtype == dtypes.float32: return val.bitcast(dtypes.uint32) if val.dtype == dtypes.half: return val.bitcast(dtypes.uint16).cast(dtypes.uint32) return val.cast(dtypes.uint32) @@ -168,6 +165,7 @@ _pcode_fixes = { 'V_CMP_CLASS_F32': ('64\'F(abs(S0.f32)) > 0.0', '(64\'U(S0.u32 & 0x7FFFFFFF) != 0)'), 'V_CMP_CLASS_F16': ('64\'F(abs(S0.f16)) > 0.0', '(64\'U(S0.u32 & 0x7FFF) != 0)'), 'V_CMP_CLASS_F64': ('64\'F(abs(S0.f64)) > 0.0', '(64\'U(S0.u64 & 0x7FFFFFFFFFFFFFFF) != 0)'), + 'V_MUL_F32': ('D0.f32 = S0.f32 * S1.f32', 'D0.f32 = isNAN(S0.f32) ? cvtToQuietNAN(S0.f32) : isNAN(S1.f32) ? cvtToQuietNAN(S1.f32) : S0.f32*S1.f32'), } def _get_pcode_dict(op) -> dict: @@ -223,7 +221,7 @@ def _write_64bit(val: UOp, wfn, reg_or_addr, is_mem: bool, *args) -> list[UOp]: """Write a 64-bit value as two 32-bit writes. args passed to wfn after reg/addr and lo/hi value.""" lo, hi = _split64(val) incr = 4 if is_mem else 1 # 4 bytes for memory addresses, 1 for register indices - return [wfn(reg_or_addr, lo, *args), wfn(reg_or_addr + (UOp.const(incr, reg_or_addr.dtype) if isinstance(reg_or_addr, UOp) else incr), hi, *args)] + return [wfn(reg_or_addr, lo, *args), wfn(reg_or_addr + (_c(incr, reg_or_addr.dtype) if isinstance(reg_or_addr, UOp) else incr), hi, *args)] def _write_val(bits: int, val: UOp, wfn, reg_or_addr, *args, is_mem: bool = False) -> list[UOp]: """Write value, splitting 64-bit if needed. bits=64 for 64-bit writes, otherwise 32-bit.""" @@ -232,10 +230,10 @@ def _write_val(bits: int, val: UOp, wfn, reg_or_addr, *args, is_mem: bool = Fals def _mem_store_bytes(mem: UOp, addr: UOp, val: UOp, active: UOp, data_bits: int = 32) -> list[UOp]: """Store to byte-addressable memory. addr is byte offset, mem is uint8 buffer.""" stores = [] - val_u32 = val.cast(dtypes.uint32) if val.dtype != dtypes.uint32 else val + val_u32 = val.cast(dtypes.uint32) for i in range(data_bits // 8): - byte_val = (val_u32 >> UOp.const(i * 8, dtypes.uint32)) & UOp.const(0xFF, dtypes.uint32) - stores.append(mem.index((addr + UOp.const(i, dtypes.uint64)).valid(active)).store(byte_val.cast(dtypes.uint8))) + byte_val = (val_u32 >> _c(i * 8)) & _c(0xFF) + stores.append(mem.index((addr + _c(i, dtypes.uint64)).valid(active)).store(byte_val.cast(dtypes.uint8))) return stores def _collect_data_slices(assigns: list[tuple[str, UOp]], data_prefix: str, pcode_vars: dict | None = None, op_name: str = "") -> dict[int, UOp]: @@ -291,11 +289,12 @@ def _int_clamp(op_name: str, srcs: dict) -> UOp | None: class _Ctx: """Context for instruction compilation - holds buffers and helpers.""" - __slots__ = ('dyn_fields', '_axis_id', 'wave_size', 'vgpr', 'accvgpr', 'branch_cond') + __slots__ = ('dyn_fields', '_axis_id', 'wave_size', 'vgpr', 'accvgpr', 'branch_cond', 'sgpr_reads') sgpr = UOp.param(0, dtypes.uint32, SGPR_COUNT, name="sgpr") vmem = UOp.param(2, dtypes.uint32, 1 << 46, name="vmem") lds = UOp.param(3, dtypes.uint32, 16384, name="lds") scratch = UOp.param(4, dtypes.uint8, 1 << 30, name="scratch") + scratch_stride = UOp.param(9, dtypes.uint32, name="scratch_stride", addrspace=AddrSpace.ALU) def __init__(self, wave_size: int = 32): self._axis_id, self.wave_size = 0, wave_size @@ -303,6 +302,7 @@ class _Ctx: self.vgpr = UOp.param(1, dtypes.uint32, 256 * wave_size, name="vgpr") self.accvgpr = UOp.param(5, dtypes.uint32, 256 * wave_size, name="accvgpr") if wave_size == 64 else self.vgpr self.branch_cond: UOp | None = None + self.sgpr_reads: dict[UOp, UOp] = {} def range(self, n: int | None = None) -> UOp: """Create a lane range UOp with unique axis ID.""" @@ -331,15 +331,15 @@ class _Ctx: word = self.inst_word(dword_idx) if lo // 32 == hi // 32: # Same dword mask = (1 << (hi - lo + 1)) - 1 - shifted = word if lo_in_dword == 0 else word >> UOp.const(lo_in_dword, dtypes.uint32) - return shifted & UOp.const(mask, dtypes.uint32) + shifted = word if lo_in_dword == 0 else word >> _c(lo_in_dword) + return shifted & _c(mask) else: # Spans two dwords lo_bits = 32 - lo_in_dword lo_mask = (1 << lo_bits) - 1 hi_mask = (1 << (hi_in_dword + 1)) - 1 - lo_part = (word >> UOp.const(lo_in_dword, dtypes.uint32)) & UOp.const(lo_mask, dtypes.uint32) - hi_part = self.inst_word(dword_idx + 1) & UOp.const(hi_mask, dtypes.uint32) - return lo_part | (hi_part << UOp.const(lo_bits, dtypes.uint32)) + lo_part = (word >> _c(lo_in_dword)) & _c(lo_mask) + hi_part = self.inst_word(dword_idx + 1) & _c(hi_mask) + return lo_part | (hi_part << _c(lo_bits)) def optional_field(self, inst: Inst, name: str) -> UOp | None: """Extract a field that only exists on some instruction formats (e.g. 'literal', 'saddr'). None if absent.""" @@ -370,15 +370,14 @@ class _Ctx: # Dynamic register access (takes UOp index instead of int) def rsgpr_dyn(self, reg: UOp, valid: UOp | None = None) -> UOp: - """Read SGPR with dynamic register index.""" - if valid is not None: return self.sgpr.index(reg.valid(valid)).load() - return self.sgpr.index(reg).load() + index = reg.valid(valid) if valid is not None else reg + if index not in self.sgpr_reads: self.sgpr_reads[index] = UOp.param(SGPR_COUNT+len(self.sgpr_reads), dtypes.uint32, addrspace=AddrSpace.ALU) + return self.sgpr_reads[index] def wsgpr_dyn(self, reg: UOp, val: UOp) -> UOp: """Write SGPR with dynamic register index. On RDNA, index 124 = NULL (writes discarded). On CDNA, index 124 = M0 (read/write).""" # RDNA: NULL (124) discards writes. CDNA: M0 (124) is writable. - valid = None if self.wave_size == 64 else reg.ne(_c(124)) - return self.sgpr.index(reg.valid(valid) if valid is not None else reg).store(val.cast(dtypes.uint32)) + return self.sgpr.index(reg if self.wave_size == 64 else reg.valid(reg.ne(_c(124)))).store(val.cast(dtypes.uint32)) def wmask(self, reg: UOp, val: UOp) -> list[UOp]: """Write a lane mask (VCC/EXEC). Splits into lo/hi for wave64.""" @@ -451,7 +450,7 @@ class _Ctx: inline = is_float_const.where(float_inline.bitcast(dtypes.uint64), int_inline.bitcast(dtypes.uint64)) # Literal handling: F64 VOP puts literal in high 32 bits; B64/I64/U64 VOP and SOP zero-extend if literal is not None: - lit_val = literal.cast(dtypes.uint64) << UOp.const(32, dtypes.uint64) if is_f64 else literal.cast(dtypes.uint64) + lit_val = literal.cast(dtypes.uint64) << _c(32, dtypes.uint64) if is_f64 else literal.cast(dtypes.uint64) inline = off.eq(_c(255)).where(lit_val, inline) scalar_val = (off < _c(128)).where(sgpr_val, inline) else: @@ -492,16 +491,16 @@ class _Ctx: def compile_lane_pcode(self, op, inst) -> UOp: """Compile cross-lane ops (READLANE/WRITELANE/PERMLANE) using pcode parser.""" pcode = get_pcode(op) - op_name = op.name if hasattr(op, 'name') else str(op) + op_name = op.name src0_off, vdst_off = self.inst_field(type(inst).src0), self.inst_field(type(inst).vdst) src0_reg = (src0_off >= _c(256)).where(src0_off - _c(256), _c(0)) # VGPR index or 0 src1_off, src2_off = self.optional_field(inst, 'src1'), self.optional_field(inst, 'src2') src1_reg = (src1_off >= _c(256)).where(src1_off - _c(256), src1_off) if src1_off is not None else _c(0) src2_reg = (src2_off >= _c(256)).where(src2_off - _c(256), src2_off) if src2_off is not None else _c(0) exec_val = self.rexec() - exec_lo = exec_val.cast(dtypes.uint32) if exec_val.dtype == dtypes.uint64 else exec_val + exec_lo = exec_val.cast(dtypes.uint32) srcs = { - 'SRC0': src0_reg, 'VDST': vdst_off, 'EXEC_LO': exec_lo, 'EXEC': exec_val if exec_val.dtype == dtypes.uint64 else exec_val.cast(dtypes.uint64), + 'SRC0': src0_reg, 'VDST': vdst_off, 'EXEC_LO': exec_lo, 'EXEC': exec_val.cast(dtypes.uint64), '_vgpr': self.vgpr, '_wave_size': self.wave_size, 'SRC1': src1_reg, 'SRC2': src2_reg, 'S0': self.rsrc_dyn(src0_off, _c(0, dtypes.int)) if 'WRITELANE' in op_name else src0_reg, 'S1': self.rsrc_dyn(src1_off, _c(0, dtypes.int)) if src1_off is not None else _c(0), @@ -531,7 +530,7 @@ class _Ctx: pcode = get_pcode(op) vcc_reg = sdst_reg if sdst_reg is not None else VCC_LO.offset if 'VCC' not in srcs: srcs['VCC'] = self.rmask(_c(vcc_reg)) - srcs.update(self.base_srcs(exec_mask, lane), VDST=vdst_reg, MAX_FLOAT_F32=UOp.const(3.4028234663852886e38, dtypes.float32)) + srcs.update(self.base_srcs(exec_mask, lane), VDST=vdst_reg, MAX_FLOAT_F32=_c(3.4028234663852886e38, dtypes.float32)) # f32 min/max/median ops flush denormal inputs to signed zero (select-style ops: results propagate inputs bitwise) # (RDNA4 calls them _NUM_: V_MIN_NUM_F32 etc.) if any(p in op.name for p in ('MIN_F32', 'MAX_F32', 'MIN3_F32', 'MAX3_F32', 'MED3_F32', 'MIN_NUM_F32', 'MAX_NUM_F32')): @@ -554,23 +553,22 @@ class _Ctx: elif dest.startswith('D0'): if (dest_suffix := re.match(r'D0\.(\w+)', dest)) is not None: target_dt = {'u16': dtypes.uint16, 'i16': dtypes.int16, 'f16': dtypes.half}.get(dest_suffix.group(1)) - if target_dt is not None and val.dtype != target_dt: val = val.cast(target_dt) + if target_dt is not None: val = val.cast(target_dt) if (slice_match := re.match(r'D0\[(\d+)\s*:\s*(\d+)\]', dest)) and (int(slice_match.group(1)), int(slice_match.group(2))) != (31, 0): slice_stores.append((int(slice_match.group(2)), int(slice_match.group(1)) - int(slice_match.group(2)) + 1, _val_to_bits(val))) continue if int_saturate is not None: val = int_saturate elif clmp and val.dtype in (dtypes.float32, dtypes.half, dtypes.float64): # hardware clamp: -0 becomes +0 and NaN becomes 0 (hardware verified) - val = (val > UOp.const(0.0, val.dtype)).where(val.minimum(UOp.const(1.0, val.dtype)), UOp.const(0.0, val.dtype)) + val = (val > _c(0.0, val.dtype)).where(val.minimum(_c(1.0, val.dtype)), _c(0.0, val.dtype)) if val.dtype in (dtypes.uint64, dtypes.int64, dtypes.float64): lo, hi = _split64(val) lane_stores.extend([self.wvgpr_dyn(vdst_reg, lane, lo, exec_mask), self.wvgpr_dyn(vdst_reg + _c(1), lane, hi, exec_mask)]) elif val.dtype in (dtypes.half, dtypes.uint16, dtypes.int16): result, old_val = _val_to_u32(val), self.rvgpr_dyn(vdst_reg, lane) - hi_result = (old_val & UOp.const(0xFFFF, dtypes.uint32)) | (result << UOp.const(16, dtypes.uint32)) + hi_result = (old_val & _c(0xFFFF)) | (result << _c(16, dtypes.uint32)) # GFX9/CDNA zeroes upper 16 bits on lo-half write; RDNA preserves them - lo_result = (result & UOp.const(0xFFFF, dtypes.uint32)) if self.wave_size == 64 else \ - (old_val & UOp.const(0xFFFF0000, dtypes.uint32)) | (result & UOp.const(0xFFFF, dtypes.uint32)) + lo_result = (result & _c(0xFFFF)) if self.wave_size == 64 else (old_val & _c(0xFFFF0000)) | (result & _c(0xFFFF)) result = opsel_dst_hi.where(hi_result, lo_result) if isinstance(opsel_dst_hi, UOp) else hi_result if opsel_dst_hi else lo_result lane_stores.append(self.wvgpr_dyn(vdst_reg, lane, result, exec_mask)) else: lane_stores.append(self.wvgpr_dyn(vdst_reg, lane, _val_to_u32(val), exec_mask)) @@ -615,7 +613,7 @@ def _compile_sopp(inst: ir3.SOPP | ir4.SOPP, ctx: _Ctx) -> UOp: pcode = get_pcode(inst.op) vcc, exec_val = ctx.rmask(_c(VCC_LO.offset)), ctx.rexec() srcs = {'PC': _c(0, dtypes.uint32), 'SIMM16': simm16, 'SCC': ctx.rsgpr_dyn(_c(SCC.offset)), 'VCC': vcc, - 'VCCZ': vcc.eq(UOp.const(0, vcc.dtype)).cast(dtypes.uint32), 'EXECZ': exec_val.eq(UOp.const(0, exec_val.dtype)).cast(dtypes.uint32)} + 'VCCZ': vcc.eq(_c(0, vcc.dtype)).cast(dtypes.uint32), 'EXECZ': exec_val.eq(_c(0, exec_val.dtype)).cast(dtypes.uint32)} for dest, val in parse_pcode(pcode, srcs)[1]: if dest.startswith('PC') and val.op is Ops.WHERE: ctx.branch_cond = val.src[0] return UOp.sink() @@ -761,7 +759,7 @@ def _load_dpp16_src0(ctx: _Ctx, inst, lane: UOp, fallback: UOp) -> UOp: src_lane, enabled, valid = _dpp16_ctrl(lane, _iattr(inst, 'dpp'), _iattr(inst, 'row_mask', 0xf), _iattr(inst, 'bank_mask', 0xf), ctx.wave_size) safe_src_lane = (enabled & valid).where(src_lane, _c(0, dtypes.int)) swizzled = ctx.rvgpr_dyn(ctx.inst_field(type(inst).vsrc0), safe_src_lane) - invalid = UOp.const(0, fallback.dtype) if _iattr(inst, 'bc') else fallback + invalid = _c(0, fallback.dtype) if _iattr(inst, 'bc') else fallback return enabled.where(valid.where(swizzled, invalid), fallback) def _compile_sdwa(inst: irc.VOP1_SDWA | irc.VOP2_SDWA | irc.VOP2_SDWA_SDST | irc.VOPC_SDWA_SDST, ctx: _Ctx) -> UOp: @@ -785,17 +783,16 @@ def _compile_sdwa(inst: irc.VOP1_SDWA | irc.VOP2_SDWA | irc.VOP2_SDWA_SDST | irc # For VOPC: use unroll_lanes to build the bitmask from scratch (no read-modify-write on stale data) if is_vopc: - def get_cmp_bit(lane) -> UOp: - lc = lane.cast(dtypes.int) if isinstance(lane, UOp) else _c(lane, dtypes.int) + def get_cmp_bit(lane: UOp) -> UOp: + lc = lane.cast(dtypes.int) s0_raw = ctx.rsgpr_dyn(vsrc0_reg) if inst.s0 else ctx.rvgpr_dyn(vsrc0_reg, lc) s0 = _sdwa_select(s0_raw, src0_sel, src0_sext) s1_raw = ctx.rsgpr_dyn(vsrc1_reg) if inst.s1 else ctx.rvgpr_dyn(vsrc1_reg, lc) s1 = _sdwa_select(s1_raw, src1_sel, src1_sext) - srcs = {'S0': s0, 'S1': s1, 'laneId': lc} - for dest, val in parse_pcode(pcode, srcs)[1]: + for dest, val in parse_pcode(pcode, {'S0': s0, 'S1': s1, 'laneId': lc})[1]: if '[laneId]' in dest and ('D0' in dest or 'EXEC' in dest): return val.cast(dtypes.uint32) return _c(0) - new_result = ctx.unroll_lanes(get_cmp_bit, exec_mask, apply_exec=False) & exec_mask + new_result = ctx.unroll_lanes(get_cmp_bit, exec_mask) stores = ctx.wmask(sdst_off, new_result) return UOp.sink(*stores) @@ -828,15 +825,10 @@ def _compile_sdwa(inst: irc.VOP1_SDWA | irc.VOP2_SDWA | irc.VOP2_SDWA_SDST | irc old = ctx.rvgpr_dyn(vdst_reg, lane) result = _sdwa_write(old, result, dst_sel, dst_unused) stores.append(ctx.wvgpr_dyn(vdst_reg, lane, result, exec_mask)) - elif dest.startswith('VCC'): - stores.extend(ctx.wmask_lane_bit(_c(VCC_LO.offset), lane, val, exec_mask)) + elif dest.startswith('VCC'): vcc_val = val if vcc_val is not None: - # Initialize sdst to 0 before lane loop (old value may be unrelated data), then set lane bits in loop - init_stores = [ctx.wsgpr_dyn(sdst_off, _c(0)), ctx.wsgpr_dyn(sdst_off + _c(1), _c(0))] - stores.extend(ctx.wmask_lane_bit(sdst_off, lane, vcc_val, exec_mask)) - if stores: - return UOp.sink(*init_stores, UOp.sink(*stores).end(lane)) - return UOp.sink(*init_stores) + mask = ctx.unroll_lanes(lambda l: (_to_u32(vcc_val.substitute({lane:l})) & _c(1)).cast(dtypes.uint32), exec_mask) + return UOp.sink(*ctx.wmask(sdst_off, mask), UOp.sink(*stores).end(lane)) if stores: return UOp.sink(UOp.sink(*stores).end(lane)) return UOp.sink() @@ -878,7 +870,6 @@ def _compile_vop12(inst: ir3.VOP1 | ir3.VOP1_SDST | ir3.VOP1_DPP16 | ir3.VOP2 | vdst_reg = ctx.inst_field(type(inst).vdst) write_hi_half = bits['d'] == 16 and not isinstance(inst, (irc.VOP1, irc.VOP2)) and (vdst_reg >= _c(128)) if isinstance(write_hi_half, UOp): vdst_reg = write_hi_half.where(vdst_reg - _c(128), vdst_reg) - elif write_hi_half: vdst_reg -= 128 if isinstance(inst, (ir3.VOP1, ir4.VOP1, irc.VOP1)): d0 = _cond_hi16(write_hi_half, ctx.rvgpr_dyn(vdst_reg, lane)) s0, src0_off = _load_vsrc0(ctx, inst, lane, bits, literal, is_f64, is_float, d0) @@ -909,8 +900,8 @@ def _compile_vopc(inst: ir3.VOPC|ir3.VOPC_DPP16|ir3.VOP3|ir4.VOPC|ir4.VOPC_DPP16 is_dpp16 = hasattr(type(inst), 'dpp') and hasattr(type(inst), 'vsrc0') # Handle both VOPC (vsrc1) and VOP3 (src1) instruction formats - read operands dynamically + src0_off = ctx.inst_field(type(inst).src0) if is_vopc: - src0_off = ctx.inst_field(type(inst).src0) vsrc1_off = ctx.inst_field(type(inst).vsrc1) # type: ignore[union-attr] # For 16-bit ops, vsrc1 >= 128 means hi-half of v[vsrc1-128] if bits['s0'] == 16 and not isinstance(inst, irc.VOPC): @@ -920,15 +911,14 @@ def _compile_vopc(inst: ir3.VOPC|ir3.VOPC_DPP16|ir3.VOP3|ir4.VOPC|ir4.VOPC_DPP16 vsrc1_hi = False src1_off = _c(256) + vsrc1_off else: - src0_off = ctx.inst_field(type(inst).src0) src1_off = ctx.inst_field(type(inst).src1) # type: ignore[union-attr] dst_off = ctx.inst_field(type(inst).vdst) # type: ignore[union-attr] vsrc1_hi = False literal = ctx.optional_field(inst, 'literal') is_float, is_f64, pcode = any(x in op_name for x in ('_F32', '_F64', '_F16')), '_F64' in op_name, get_pcode(inst.op) - def get_cmp_bit(lane) -> UOp: - lc = lane.cast(dtypes.int) if isinstance(lane, UOp) else _c(lane, dtypes.int) + def get_cmp_bit(lane: UOp) -> UOp: + lc = lane.cast(dtypes.int) s0 = _load_dpp16_src0(ctx, inst, lc, _c(0)) if is_dpp16 else ctx.rsrc_dyn(src0_off, lc, bits['s0'], literal, is_f64) if is_vopc and not isinstance(inst, irc.VOPC) and bits['s0'] == 16 and not is_dpp16: src0_hi = src0_off >= _c(384) @@ -936,22 +926,17 @@ def _compile_vopc(inst: ir3.VOPC|ir3.VOPC_DPP16|ir3.VOP3|ir4.VOPC|ir4.VOPC_DPP16 s1 = _cond_hi16(vsrc1_hi, ctx.rsrc_dyn(src1_off, lc, bits['s1'], literal, is_f64)) if bits['s0'] == 16 \ else ctx.rsrc_dyn(src1_off, lc, bits['s1'], literal, is_f64) if bits['s0'] == 16 and opsel: s0, s1 = _apply_opsel(s0, 0, opsel), _apply_opsel(s1, 1, opsel) - if is_float: - if is_dpp16: - s0 = _apply_src_mods(s0, 0, 1 if _iattr(inst, 'src0_abs') else 0, 1 if _iattr(inst, 'src0_neg') else 0, bits['s0']) - s1 = _apply_src_mods(s1, 0, 1 if _iattr(inst, 'src1_abs') else 0, 1 if _iattr(inst, 'src1_neg') else 0, bits['s1']) - s0 = _apply_src_mods(s0, 0, abs_bits, neg_bits, bits['s0']) - s1 = _apply_src_mods(s1, 1, abs_bits, neg_bits, bits['s1']) - elif abs_bits or neg_bits: # int compares also honor abs/neg, as bit-level sign clear/flip (not integer abs/negate) - s0 = _apply_src_mods(s0, 0, abs_bits, neg_bits, bits['s0']) - s1 = _apply_src_mods(s1, 1, abs_bits, neg_bits, bits['s1']) - for dest, val in parse_pcode(pcode, {'S0': s0, 'S1': s1, 'laneId': lc, 'D0': UOp.const(0, dtypes.uint64)})[1]: + if is_float and is_dpp16: + s0 = _apply_src_mods(s0, 0, 1 if _iattr(inst, 'src0_abs') else 0, 1 if _iattr(inst, 'src0_neg') else 0, bits['s0']) + s1 = _apply_src_mods(s1, 0, 1 if _iattr(inst, 'src1_abs') else 0, 1 if _iattr(inst, 'src1_neg') else 0, bits['s1']) + s0 = _apply_src_mods(s0, 0, abs_bits, neg_bits, bits['s0']) + s1 = _apply_src_mods(s1, 1, abs_bits, neg_bits, bits['s1']) + for dest, val in parse_pcode(pcode, {'S0': s0, 'S1': s1, 'laneId': lc, 'D0': _c(0, dtypes.uint64)})[1]: if '[laneId]' in dest and ('D0' in dest or 'EXEC' in dest): return val.cast(dtypes.uint32) return _c(0) - new_bits = ctx.unroll_lanes(get_cmp_bit, exec_mask, apply_exec=False) # Both VOPC and VOP3 clear inactive lane bits (hardware verified) - new_result = new_bits & exec_mask + new_result = ctx.unroll_lanes(get_cmp_bit, exec_mask) # CMPX writes EXEC only (hardware verified: e64 CMPX does not write SDST); non-CMPX writes SDST/VCC if is_cmpx: stores = ctx.wmask(_c(EXEC_LO.offset), new_result) @@ -972,8 +957,8 @@ def _compile_bitop3(inst, ctx: _Ctx, exec_mask: UOp, bits: dict, op_name: str) - is_16 = 'B16' in op_name dt, mask = (dtypes.uint16, 0xFFFF) if is_16 else (dtypes.uint32, 0xFFFFFFFF) s0, s1, s2 = src0.cast(dt), src1.cast(dt), src2.cast(dt) - def bnot(v): return v ^ UOp.const(mask, dt) - result = UOp.const(0, dt) + def bnot(v): return v ^ _c(mask, dt) + result = _c(0, dt) for i in range(8): if not (ttbl & (1 << i)): continue result = result | ((s0 if i & 4 else bnot(s0)) & (s1 if i & 2 else bnot(s1)) & (s2 if i & 1 else bnot(s2))) @@ -993,7 +978,7 @@ def _compile_vop3(inst: ir3.VOP3 | ir4.VOP3 | irc.VOP3, ctx: _Ctx) -> UOp: return ctx.compile_lane_pcode(inst.op, inst) # VOP3 VOPC (v_cmp_*_e64) - delegate to unified VOPC handler - if 'V_CMP' in op_name or 'V_CMPX' in op_name: + if 'V_CMP' in op_name: return _compile_vopc(inst, ctx, opsel=opsel, abs_bits=_iattr(inst, 'abs'), neg_bits=_iattr(inst, 'neg')) # BITOP3: abs/neg/omod encode truth table, not source modifiers @@ -1025,8 +1010,8 @@ def _compile_vop3(inst: ir3.VOP3 | ir4.VOP3 | irc.VOP3, ctx: _Ctx) -> UOp: src0 = _apply_src_mods(src0, 0, abs_bits, neg_bits, bits['s0']) src1 = _apply_src_mods(src1, 1, abs_bits, neg_bits, bits['s1']) src2 = _apply_src_mods(src2, 2, abs_bits, neg_bits, bits['s2']) - srcs = {'S0': src0, 'S1': src1, 'S2': src2, 'OPSEL': UOp.const(opsel, dtypes.uint32)} - if 'CNDMASK' in op_name and src2 is not None: srcs['VCC'] = src2 + srcs = {'S0': src0, 'S1': src1, 'S2': src2, 'OPSEL': _c(opsel, dtypes.uint32)} + if 'CNDMASK' in op_name: srcs['VCC'] = src2 # FMAC instructions need D0 (accumulator) from destination register if 'FMAC' in op_name: srcs['D0'] = ctx.rvgpr_dyn(vdst_reg, lane) opsel_dst_hi = bool(opsel & 0b1000) and bits['d'] == 16 @@ -1127,8 +1112,8 @@ def _compile_mfma(inst: irc.VOP3P|irc.VOP3PX2, ctx: _Ctx) -> UOp: scale0 = ctx.inst_field(type(inst).scale_src0), _iattr(inst, 'opsel') & 3 scale1 = ctx.inst_field(type(inst).scale_src1), _iattr(inst, 'opsel_hi') & 3 def _scale_exp(off_sel: tuple[UOp, int], lane: UOp) -> UOp: - byte = (ctx.rsrc_dyn(off_sel[0], lane, 32) >> UOp.const(off_sel[1] * 8, dtypes.uint32)) & UOp.const(0xFF, dtypes.uint32) - return byte.cast(dtypes.int32) - UOp.const(127, dtypes.int32) + byte = (ctx.rsrc_dyn(off_sel[0], lane, 32) >> _c(off_sel[1] * 8, dtypes.uint32)) & _c(0xFF, dtypes.uint32) + return byte.cast(dtypes.int32) - _c(127, dtypes.int32) def scale_factor(lane: UOp) -> UOp: # combined A*B scale for this lane: 2^(ea-127) * 2^(eb-127) return UOp.exp2((_scale_exp(scale0, lane) + _scale_exp(scale1, lane)).cast(dtypes.float32)) @@ -1163,26 +1148,26 @@ def _compile_mfma(inst: irc.VOP3P|irc.VOP3PX2, ctx: _Ctx) -> UOp: def cvt_elem(raw: UOp, sub_idx: int, fp8_fmt: str) -> UOp: if is_i8: # extract i8, sign-extend to i32 - byte = (raw >> UOp.const(sub_idx * 8, dtypes.uint32)) & UOp.const(0xFF, dtypes.uint32) - return (byte.cast(dtypes.int32) ^ UOp.const(0x80, dtypes.int32)) - UOp.const(0x80, dtypes.int32) + byte = (raw >> _c(sub_idx * 8)) & _c(0xFF, dtypes.uint32) + return (byte.cast(dtypes.int32) ^ _c(0x80, dtypes.int32)) - _c(0x80, dtypes.int32) if is_f32_src: return raw # already uint32 (f32 bit pattern) - if is_fp8: return _FUNCS[f"{fp8_fmt}_to_f32"](raw >> UOp.const(sub_idx * 8, dtypes.uint32)).bitcast(dtypes.uint32) - h = (raw >> UOp.const(sub_idx * 16, dtypes.uint32)) & UOp.const(0xFFFF, dtypes.uint32) - if is_bf16: return h << UOp.const(16, dtypes.uint32) # bf16 is the upper 16 bits of f32 + if is_fp8: return _FUNCS[f"{fp8_fmt}_to_f32"](raw >> _c(sub_idx * 8, dtypes.uint32)).bitcast(dtypes.uint32) + h = (raw >> _c(sub_idx * 16)) & _c(0xFFFF) + if is_bf16: return h << _c(16) # bf16 is the upper 16 bits of f32 return _FUNCS['f16_to_f32'](h).bitcast(dtypes.uint32) def mn_idx(lane: UOp) -> UOp: # M/N matrix index held by a lane if M == 32: # (lane%32)/16 selects the 16-wide block, (lane%32)%16 the index within it - return ((lane // UOp.const(16, dtypes.int)) % UOp.const(2, dtypes.int)) * UOp.const(16, dtypes.int) + lane % UOp.const(16, dtypes.int) - return lane % UOp.const(grp_size, dtypes.int) + return ((lane // _c(16, dtypes.int)) % _c(2, dtypes.int)) * _c(16, dtypes.int) + lane % _c(16, dtypes.int) + return lane % _c(grp_size, dtypes.int) - def grp_idx(lane: UOp) -> UOp: return lane // UOp.const(grp_size, dtypes.int) + def grp_idx(lane: UOp) -> UOp: return lane // _c(grp_size, dtypes.int) read_lane = ctx.range() mn, grp = mn_idx(read_lane), grp_idx(read_lane) def mat_idx(base: int, dim: int, kl: int) -> UOp: # tmp index of element (mn, kl) of this lane's group - if M == 4: return UOp.const(base + kl, dtypes.int) + grp * UOp.const(dim * K, dtypes.int) + mn * UOp.const(K, dtypes.int) - return UOp.const(base + kl, dtypes.int) + mn * UOp.const(K, dtypes.int) + grp * UOp.const(k_per_grp, dtypes.int) + if M == 4: return _c(base + kl, dtypes.int) + grp * _c(dim * K, dtypes.int) + mn * _c(K, dtypes.int) + return _c(base + kl, dtypes.int) + mn * _c(K, dtypes.int) + grp * _c(k_per_grp, dtypes.int) read_stores = [] for kl in range(k_per_grp): @@ -1203,7 +1188,7 @@ def _compile_mfma(inst: irc.VOP3P|irc.VOP3PX2, ctx: _Ctx) -> UOp: def _dot_accum(acc: UOp, a_row: UOp, b_row: UOp, lane: UOp) -> UOp: """acc += sum_k A[a_row+k] * B[b_row+k] in order (FP-associativity matters). For scaled MFMA only the dot is scaled: D = dot*scale + C.""" def prod(k: int) -> UOp: - return tmp2.index(a_row + UOp.const(k, dtypes.int)).bitcast(acc_dt) * tmp2.index(b_row + UOp.const(k, dtypes.int)).bitcast(acc_dt) + return tmp2.index(a_row + _c(k, dtypes.int)).bitcast(acc_dt) * tmp2.index(b_row + _c(k, dtypes.int)).bitcast(acc_dt) if not scaled: for k in range(K): acc = acc + prod(k) return acc @@ -1213,16 +1198,16 @@ def _compile_mfma(inst: irc.VOP3P|irc.VOP3PX2, ctx: _Ctx) -> UOp: compute_lane = ctx.range() c_mn, c_grp = mn_idx(compute_lane), grp_idx(compute_lane) - b_off = UOp.const(n_a_elems, dtypes.int) + b_off = _c(n_a_elems, dtypes.int) def out_ab(out_reg: int) -> tuple[UOp, UOp]: # A/B tmp base indices for one output element of this lane if M == 32: # 16 outputs per lane: rows (lane//32)*16 + blocks of 4 within the wave - m_base = c_grp * UOp.const(16, dtypes.int) + UOp.const((out_reg // 4) * 4 + out_reg % 4, dtypes.int) - return m_base * UOp.const(K, dtypes.int), b_off + c_mn * UOp.const(K, dtypes.int) + m_base = c_grp * _c(16, dtypes.int) + _c(out_reg, dtypes.int) + return m_base * _c(K, dtypes.int), b_off + c_mn * _c(K, dtypes.int) if M == 4: # each group of 4 lanes computes an independent 4x4 block - a_base = c_grp * UOp.const(M * K, dtypes.int) + UOp.const(out_reg * K, dtypes.int) - return a_base, b_off + (c_grp * UOp.const(N, dtypes.int) + c_mn) * UOp.const(K, dtypes.int) - m_base = c_grp * UOp.const(out_per_lane, dtypes.int) + UOp.const(out_reg, dtypes.int) # 16x16: 4 outputs, one row each - return m_base * UOp.const(K, dtypes.int), b_off + c_mn * UOp.const(K, dtypes.int) + a_base = c_grp * _c(M * K, dtypes.int) + _c(out_reg * K, dtypes.int) + return a_base, b_off + (c_grp * _c(N, dtypes.int) + c_mn) * _c(K, dtypes.int) + m_base = c_grp * _c(out_per_lane, dtypes.int) + _c(out_reg, dtypes.int) # 16x16: 4 outputs, one row each + return m_base * _c(K, dtypes.int), b_off + c_mn * _c(K, dtypes.int) compute_stores = [] for out_reg in range(out_per_lane): @@ -1250,9 +1235,8 @@ def _compile_wmma(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P, ctx: _Ctx) -> UOp: # read a source element from VGPRs: (src, lane, vgpr, element-in-vgpr) -> f32/i32 def gval(src, lane, vgpr, ridx, *, fmt): - v = ctx.rvgpr_dyn(src + _c(vgpr), UOp.const(lane, dtypes.int)) - pkd = v >> UOp.const(ridx * sz, dtypes.uint32) - pkd = pkd & UOp.const((1 << sz) - 1, dtypes.uint32) + v = ctx.rvgpr_dyn(src + _c(vgpr), _c(lane, dtypes.int)) + pkd = (v >> _c(ridx * sz)) & _c((1 << sz) - 1) if not fmt.startswith('IU'): return _FUNCS[f'{fmt.lower()}_to_f32'](pkd) return (pkd << _c(24, dtypes.uint)).bitcast(dtypes.int32) >> _c(24, dtypes.int32) # sign extend @@ -1274,8 +1258,7 @@ def _compile_wmma(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P, ctx: _Ctx) -> UOp: # src2 may be a VGPR or an inline/scalar constant (128 = int 0, the usual ", 0" C form); the runner must handle both dynamically out_dt = dtypes.int32 if output_type == "I32" else dtypes.float32 cbits = ctx.rsrc_dyn(src2_r, None, 32) - cval_const = _FUNCS[f'{output_type.lower()}_to_f32'](cbits & UOp.const(0xFFFF, dtypes.uint32)) \ - if output_type in ("F16", "BF16") else cbits.bitcast(out_dt) + cval_const = _FUNCS[f'{output_type.lower()}_to_f32'](cbits & _c(0xFFFF)) if output_type in ("F16", "BF16") else cbits.bitcast(out_dt) elem = ctx.range(256) row, col = elem // 16, elem % 16 lane, vgpr = d_map(row, col) @@ -1291,13 +1274,13 @@ def _compile_wmma(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P, ctx: _Ctx) -> UOp: def w_store(m: int, n: UOp, val: UOp, vgpr_off: int) -> UOp: # store one output element to its (lane, vgpr) slot lane_i, _ = d_map(m, n) - return ctx.wvgpr_dyn(vdst_reg + _c(vgpr_off), UOp.const(lane_i, dtypes.int), val, exec_mask) + return ctx.wvgpr_dyn(vdst_reg + _c(vgpr_off), _c(lane_i, dtypes.int), val, exec_mask) if output_type in ("F16", "BF16"): def to_bits(v: UOp) -> UOp: # f32 result -> 16 output bits - return ((v.bitcast(dtypes.uint32) >> UOp.const(16, dtypes.uint32)) & UOp.const(0xFFFF, dtypes.uint32)) if output_type == "BF16" \ + return ((v.bitcast(dtypes.uint32) >> _c(16, dtypes.uint32)) & _c(0xFFFF)) if output_type == "BF16" \ else v.cast(dtypes.half).bitcast(dtypes.uint16).cast(dtypes.uint32) if is_rdna4: # pack 2 outputs per VGPR (adjacent m values share a VGPR) - stores = [w_store(m, n, to_bits(mat_d[m]) | (to_bits(mat_d[m+1]) << UOp.const(16, dtypes.uint32)), d_map(m, n)[1] // 2) + stores = [w_store(m, n, to_bits(mat_d[m]) | (to_bits(mat_d[m+1]) << _c(16, dtypes.uint32)), d_map(m, n)[1] // 2) for m in range(0, 16, 2)] else: # one output per VGPR (lo half) stores = [w_store(m, n, to_bits(mat_d[m]), d_map(m, n)[1]) for m in range(16)] @@ -1323,11 +1306,11 @@ def _compile_vop3p(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P | irc.VOP3PX2, ctx: _ exec_mask = ctx.rexec() vdst_reg = ctx.inst_field(type(inst).vdst) src0_off = ctx.inst_field(type(inst).src0) - if src0_is_acc and not vdst_is_acc: + if not vdst_is_acc: # v_accvgpr_read: VGPR[vdst] = ACCVGPR[src0] val = ctx.raccvgpr_dyn(src0_off - _c(256), lane) return UOp.sink(ctx.wvgpr_dyn(vdst_reg, lane, val, exec_mask).end(lane)) - elif vdst_is_acc and not src0_is_acc: + elif not src0_is_acc: # v_accvgpr_write: ACCVGPR[vdst] = src0 (src0 can be VGPR or SGPR/const) src0 = ctx.rsrc_dyn(src0_off, lane, 32) return UOp.sink(ctx.waccvgpr_dyn(vdst_reg, lane, src0, exec_mask).end(lane)) @@ -1362,8 +1345,7 @@ def _compile_vop3p(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P | irc.VOP3PX2, ctx: _ # v_pk_mov_b32: D[lo] = src0[opsel_bit0 ? hi : lo], D[hi] = src1[opsel_bit1 ? hi : lo] lo_val = _pk_sel(src0, ctx.inst_field(type(inst).src0), opsel & 1) hi_val = _pk_sel(src1, ctx.inst_field(type(inst).src1), opsel & 2) - lo_out, hi_out = _split64(_u64(lo_val, hi_val)) - stores = [ctx.wvgpr_dyn(vdst_reg, lane, lo_out, exec_mask), ctx.wvgpr_dyn(vdst_reg + _c(1), lane, hi_out, exec_mask)] + stores = [ctx.wvgpr_dyn(vdst_reg, lane, lo_val, exec_mask), ctx.wvgpr_dyn(vdst_reg + _c(1), lane, hi_val, exec_mask)] return UOp.sink(UOp.group(*stores).end(lane)) srcs: dict[str, UOp | int] = {} @@ -1373,8 +1355,8 @@ def _compile_vop3p(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P | irc.VOP3PX2, ctx: _ hi_bits = (opsel_hi & 1, opsel_hi & 2, 1 if opsel_hi2 else 0) for i, (bit, s0) in enumerate(zip((1, 2, 4), (src0, src1, src2))): lo, hi = _pk_sel(s0, src_offs[i], opsel & bit), _pk_sel(s0, src_offs[i], hi_bits[i]) - if neg & bit: lo = lo ^ UOp.const(0x80000000, dtypes.uint32) - if neg_hi & bit: hi = hi ^ UOp.const(0x80000000, dtypes.uint32) + if neg & bit: lo = lo ^ _c(0x80000000, dtypes.uint32) + if neg_hi & bit: hi = hi ^ _c(0x80000000, dtypes.uint32) srcs[f'S{i}'] = _u64(lo, hi) elif 'FMA_MIX' in op_name or 'MAD_MIX' in op_name: combined_opsel_hi = (opsel_hi & 0x3) | ((opsel_hi2 & 0x1) << 2) @@ -1382,35 +1364,35 @@ def _compile_vop3p(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P | irc.VOP3PX2, ctx: _ def apply_abs(v, bit, opsel_hi_bit, opsel_bit): if not (neg_hi & bit): return v # Apply abs based on whether source is f32 or f16 - if not (combined_opsel_hi & opsel_hi_bit): return v & UOp.const(0x7FFFFFFF, dtypes.uint32) # f32 abs - if opsel & opsel_bit: return v & UOp.const(0x7FFF0000, dtypes.uint32) # f16 hi abs (preserve lo) - return v & UOp.const(0xFFFF7FFF, dtypes.uint32) # f16 lo abs (preserve hi) + if not (combined_opsel_hi & opsel_hi_bit): return v & _c(0x7FFFFFFF, dtypes.uint32) # f32 abs + if opsel & opsel_bit: return v & _c(0x7FFF0000, dtypes.uint32) # f16 hi abs (preserve lo) + return v & _c(0xFFFF7FFF, dtypes.uint32) # f16 lo abs (preserve hi) def apply_neg_mix(v, bit, opsel_hi_bit, opsel_bit): if not (neg & bit): return v - if not (combined_opsel_hi & opsel_hi_bit): return v ^ UOp.const(0x80000000, dtypes.uint32) # f32 neg - if opsel & opsel_bit: return v ^ UOp.const(0x80000000, dtypes.uint32) # f16 hi neg - return v ^ UOp.const(0x00008000, dtypes.uint32) # f16 lo neg + if not (combined_opsel_hi & opsel_hi_bit): return v ^ _c(0x80000000, dtypes.uint32) # f32 neg + if opsel & opsel_bit: return v ^ _c(0x80000000, dtypes.uint32) # f16 hi neg + return v ^ _c(0x00008000, dtypes.uint32) # f16 lo neg s0_mod = apply_neg_mix(apply_abs(src0, 1, 1, 1), 1, 1, 1) s1_mod = apply_neg_mix(apply_abs(src1, 2, 2, 2), 2, 2, 2) s2_mod = apply_neg_mix(apply_abs(src2, 4, 4, 4), 4, 4, 4) srcs = {'S@0': s0_mod, 'S@1': s1_mod, 'S@2': s2_mod, - 'OPSEL_HI': UOp.const(combined_opsel_hi, dtypes.uint32), 'OPSEL': UOp.const(opsel, dtypes.uint32)} + 'OPSEL_HI': _c(combined_opsel_hi, dtypes.uint32), 'OPSEL': _c(opsel, dtypes.uint32)} else: def get_half_bits(val: UOp, use_hi: bool, apply_neg: bool = False) -> UOp: - bits = ((val >> UOp.const(16, dtypes.uint32)) if use_hi else val) & UOp.const(0xFFFF, dtypes.uint32) + bits = ((val >> _c(16, dtypes.uint32)) if use_hi else val) & _c(0xFFFF) if apply_neg: bits = bits.cast(dtypes.uint16).bitcast(dtypes.half).neg().bitcast(dtypes.uint16).cast(dtypes.uint32) return bits def build_remapped_src(src: UOp, opsel_lo_bit: int, opsel_hi_bit: int, neg_lo_bit: int, neg_hi_bit: int) -> UOp: lo = get_half_bits(src, bool(opsel_lo_bit), bool(neg_lo_bit)) hi = get_half_bits(src, bool(opsel_hi_bit), bool(neg_hi_bit)) - return lo | (hi << UOp.const(16, dtypes.uint32)) + return lo | (hi << _c(16, dtypes.uint32)) # DOT IU instructions use NEG bits for signed/unsigned selection, not fp16 negation is_dot_iu = 'DOT' in op_name and 'IU' in op_name n0, n1, n2, nh0, nh1, nh2 = (0, 0, 0, 0, 0, 0) if is_dot_iu else (neg & 1, neg & 2, neg & 4, neg_hi & 1, neg_hi & 2, neg_hi & 4) srcs = {'S0': build_remapped_src(src0, opsel & 1, opsel_hi & 1, n0, nh0), 'S1': build_remapped_src(src1, opsel & 2, opsel_hi & 2, n1, nh1), 'S2': build_remapped_src(src2, opsel & 4, 1 if opsel_hi2 else 0, n2, nh2)} - if is_dot_iu: srcs['NEG'] = UOp.const(neg, dtypes.uint32) + if is_dot_iu: srcs['NEG'] = _c(neg, dtypes.uint32) return ctx.compile_vop_pcode(inst.op, srcs, lane, vdst_reg, exec_mask) def _compile_vopd(inst: ir3.VOPD | ir4.VOPD, ctx: _Ctx) -> UOp: @@ -1441,7 +1423,6 @@ def _compile_vopd(inst: ir3.VOPD | ir4.VOPD, ctx: _Ctx) -> UOp: if vop in (ir3.VOP2Op.V_FMAAK_F32_E32, ir3.VOP2Op.V_FMAMK_F32_E32): assert literal is not None srcs['SIMM32'] = literal - if op in (ir3.VOPDOp.V_DUAL_CNDMASK_B32, ir4.VOPDOp.V_DUAL_CNDMASK_B32): srcs['VCC'] = ctx.rmask(_c(VCC_LO.offset)) pcode = get_pcode(vop) srcs.update(ctx.base_srcs(exec_mask, lane), VCC=ctx.rmask(_c(VCC_LO.offset))) for dest, val in parse_pcode(pcode, srcs)[1]: @@ -1463,7 +1444,7 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA # CDNA acc bit: when set, VGPR operands (vdst/vdata) target ACCVGPR file instead of VGPR use_acc = bool(_iattr(inst, 'acc')) mem = ctx.lds if is_lds else ctx.scratch if is_scratch else ctx.vmem - addr_shift = UOp.const(2, dtypes.uint32 if is_lds else dtypes.uint64) + addr_shift = _c(2, dtypes.uint32 if is_lds else dtypes.uint64) # Field names differ per format: DS (addr/data0/offset0+offset1), RDNA4 V* (vaddr/vsrc/ioffset), RDNA3+CDNA (addr/data/offset) if is_lds: addr_field, data_field = 'addr', 'data0' @@ -1483,7 +1464,6 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA # Data width from canonical_op_bits (32/64/96/128), default to 32 for untyped ops data_bits_mem = inst.canonical_op_bits.get('data', 32) is_atomic, glc = 'ATOMIC' in op_name, _iattr(inst, 'glc') - has_data1 = is_lds and hasattr(inst, 'data1') and inst.data1 is not None data1_reg = ctx.inst_field(type(inst).data1) if is_lds else _c(0) # type: ignore[union-attr] if is_lds and op_name == 'DS_SWIZZLE_B32': @@ -1516,18 +1496,17 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA # Dynamic saddr check: saddr < 124 means valid SGPR, otherwise use VGPR pair for address use_saddr = (saddr_reg < _c(124)) if saddr_reg is not None else UOp.const(False) if is_scratch: - scratch_stride = ctx.rsgpr_dyn(_c(SCRATCH_STRIDE_IDX)).cast(dtypes.uint64) - base = lane.cast(dtypes.uint64) * scratch_stride + base = lane.cast(dtypes.uint64) * ctx.scratch_stride.cast(dtypes.uint64) # SVE (Scratch VGPR Enable): when SVE=1, VADDR is used as offset; when SVE=0, VADDR is ignored sve = _iattr(inst, 'sve') vaddr = ctx.rvgpr_dyn(addr_reg, lane).cast(dtypes.uint64) - addr_offset = vaddr if sve == 1 else UOp.const(0, dtypes.uint64) + addr_offset = vaddr if sve == 1 else _c(0, dtypes.uint64) # Add saddr value only if use_saddr is true (saddr < 124) - saddr_contrib = use_saddr.where(ctx.rsgpr_dyn(saddr_reg).cast(dtypes.uint64), UOp.const(0, dtypes.uint64)) \ - if saddr_reg is not None else UOp.const(0, dtypes.uint64) + saddr_contrib = use_saddr.where(ctx.rsgpr_dyn(saddr_reg).cast(dtypes.uint64), _c(0, dtypes.uint64)) \ + if saddr_reg is not None else _c(0, dtypes.uint64) return base + addr_offset + saddr_contrib + offset64 # GLOBAL can use an SGPR base plus a VGPR offset; FLAT always uses the full VGPR pair. - saddr_base = _u64(ctx.rsgpr_dyn(saddr_reg), ctx.rsgpr_dyn(saddr_reg + _c(1))) if saddr_reg is not None else UOp.const(0, dtypes.uint64) + saddr_base = _u64(ctx.rsgpr_dyn(saddr_reg), ctx.rsgpr_dyn(saddr_reg + _c(1))) if saddr_reg is not None else _c(0, dtypes.uint64) vaddr_base = _u64(ctx.rvgpr_dyn(addr_reg, lane), ctx.rvgpr_dyn(addr_reg + _c(1), lane)) # When saddr is valid: base = saddr pair, vaddr is 32-bit offset; otherwise: base = 0, vaddr is 64-bit address base_addr = use_saddr.where(saddr_base + ctx.rvgpr_dyn(addr_reg, lane).cast(dtypes.uint64), vaddr_base) @@ -1553,10 +1532,10 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA _rvdata = ctx.raccvgpr_dyn if use_acc else ctx.rvgpr_dyn if is_lds: if data_bits_mem <= 32: - data = {'DATA': _rvdata(vdata_reg, lane), 'DATA2': _rvdata(data1_reg, lane) if has_data1 else UOp.const(0, dtypes.uint32)} + data = {'DATA': _rvdata(vdata_reg, lane), 'DATA2': _rvdata(data1_reg, lane)} elif data_bits_mem == 64: # DATA/DATA2 are the 64-bit input registers, formed from register pairs data = {'DATA': _u64(_rvdata(vdata_reg, lane), _rvdata(vdata_reg + _c(1), lane)), - 'DATA2': _u64(_rvdata(data1_reg, lane), _rvdata(data1_reg + _c(1), lane)) if has_data1 else UOp.const(0, dtypes.uint64)} + 'DATA2': _u64(_rvdata(data1_reg, lane), _rvdata(data1_reg + _c(1), lane))} else: # 96/128-bit: one register per dword data = {'DATA': _rvdata(vdata_reg, lane), **{f'DATA{i}': _rvdata(vdata_reg + _c(i), lane) for i in range(1, data_bits_mem // 32)}} if 'D16' in op_name: data['RETURN_DATA'] = _rvdata(vdst_reg, lane) @@ -1567,8 +1546,8 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA active = _lane_active(exec_mask, lane) # saddr < 124 means valid SGPR pair, otherwise use 0 (NULL means no saddr contribution) use_saddr = (saddr_reg < _c(124)) if saddr_reg is not None else UOp.const(False) - saddr_raw = _u64(ctx.rsgpr_dyn(saddr_reg), ctx.rsgpr_dyn(saddr_reg + _c(1))) if saddr_reg is not None else UOp.const(0, dtypes.uint64) - saddr_base = use_saddr.where(saddr_raw, UOp.const(0, dtypes.uint64)) + saddr_raw = _u64(ctx.rsgpr_dyn(saddr_reg), ctx.rsgpr_dyn(saddr_reg + _c(1))) if saddr_reg is not None else _c(0, dtypes.uint64) + saddr_base = use_saddr.where(saddr_raw, _c(0, dtypes.uint64)) # Sign-extend offset to 64-bit for the final address calculation ioffset64 = offset.cast(dtypes.int64).cast(dtypes.uint64) # v_addr for CalcGlobalAddr: when saddr valid, use low 32 bits as offset; otherwise full 64-bit address. Include ioffset. @@ -1580,13 +1559,13 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA if data_bits_mem == 64 else ctx.rvgpr_dyn(vdata_reg, lane) return {'ADDR': addr, 'DATA': atomic_data, '_vmem': mem, '_active': active, 'laneId': lane, 'v_addr': vaddr_base, 's_saddr': saddr_base} vdata = _rvdata(vdata_reg, lane).cast(dtypes.uint64) if 'STORE' in op_name \ - else _rvdata(vdst_reg, lane) if 'D16' in op_name else UOp.const(0, dtypes.uint32) + else _rvdata(vdst_reg, lane) if 'D16' in op_name else _c(0, dtypes.uint32) if 'STORE' in op_name and data_bits_mem >= 64: - vdata = vdata | (_rvdata(vdata_reg + _c(1), lane).cast(dtypes.uint64) << UOp.const(32, dtypes.uint64)) + vdata = vdata | (_rvdata(vdata_reg + _c(1), lane).cast(dtypes.uint64) << _c(32, dtypes.uint64)) srcs = {'ADDR': addr, 'VDATA': vdata, '_vmem': mem, '_active': active, 'laneId': lane, 'v_addr': vaddr_base, 's_saddr': saddr_base, 'SADDR': saddr_base, 'OFFSET': offset} for i in range(data_bits_mem // 32): - srcs[f'VDATA{i}'] = _rvdata(vdata_reg + _c(i), lane) if 'STORE' in op_name else UOp.const(0, dtypes.uint32) + srcs[f'VDATA{i}'] = _rvdata(vdata_reg + _c(i), lane) if 'STORE' in op_name else _c(0, dtypes.uint32) return srcs def make_stores(dest: str, val: UOp, lane: UOp, active: UOp, writes_return_data: bool) -> list[UOp]: @@ -1664,7 +1643,7 @@ def _compile_mubuf(inst: irc.MUBUF, ctx: _Ctx) -> UOp: offset, offen, idxen = ctx.inst_field(type(inst).offset), ctx.inst_field(type(inst).offen), ctx.inst_field(type(inst).idxen) # V# descriptor: base[0:1], num_records[2], stride=word3[13:0] - base = _u64(ctx.rsgpr_dyn(srsrc), ctx.rsgpr_dyn(srsrc + _c(1))) & UOp.const(0xFFFFFFFFFFFF, dtypes.uint64) + base = _u64(ctx.rsgpr_dyn(srsrc), ctx.rsgpr_dyn(srsrc + _c(1))) & _c(0xFFFFFFFFFFFF, dtypes.uint64) num_records = ctx.rsgpr_dyn(srsrc + _c(2)) stride = (ctx.rsgpr_dyn(srsrc + _c(3)) & _c(0x3FFF)).cast(dtypes.uint64) @@ -1681,7 +1660,7 @@ def _compile_mubuf(inst: irc.MUBUF, ctx: _Ctx) -> UOp: buffer_offset = (stride * index + voff + offset.cast(dtypes.uint64)).cast(dtypes.uint32) in_bounds = active & buffer_offset.__lt__(num_records) addr = base + soff + buffer_offset.cast(dtypes.uint64) - addr = in_bounds.where(addr, UOp.const(0, dtypes.uint64)) # safe address when OOB + addr = in_bounds.where(addr, _c(0, dtypes.uint64)) # safe address when OOB mem = ctx.vmem stores: list[UOp] = [] @@ -1690,20 +1669,20 @@ def _compile_mubuf(inst: irc.MUBUF, ctx: _Ctx) -> UOp: lds_base = ctx.rsgpr_dyn(_c(124)) & _c(0x3FFFF) lds_addr = lds_base + lane.cast(dtypes.uint32) * _c(n_dwords * 4) for i in range(n_dwords): - word_addr = (addr + UOp.const(i * 4, dtypes.uint64)) >> UOp.const(2, dtypes.uint64) + word_addr = (addr + _c(i * 4, dtypes.uint64)) >> _c(2, dtypes.uint64) val = in_bounds.where(mem.index(word_addr.cast(dtypes.int64)).load(), _c(0)) lds_idx = (lds_addr + _c(i * 4)) >> _c(2) lds_slot = ctx.lds.index(lds_idx.valid(active)) stores.append(lds_slot.store(active.where(val, lds_slot))) elif is_store: for i in range(n_dwords): - word_addr = (addr + UOp.const(i * 4, dtypes.uint64)) >> UOp.const(2, dtypes.uint64) + word_addr = (addr + _c(i * 4, dtypes.uint64)) >> _c(2, dtypes.uint64) idx = mem.index(word_addr.cast(dtypes.int64).valid(in_bounds)) val = (ctx.raccvgpr_dyn if use_acc else ctx.rvgpr_dyn)(vdata + _c(i), lane) stores.append(idx.store(in_bounds.where(_to_u32(val), idx))) else: for i in range(n_dwords): - word_addr = (addr + UOp.const(i * 4, dtypes.uint64)) >> UOp.const(2, dtypes.uint64) + word_addr = (addr + _c(i * 4, dtypes.uint64)) >> _c(2, dtypes.uint64) val = in_bounds.where(mem.index(word_addr.cast(dtypes.int64).valid(in_bounds)).load(), _c(0)) stores.append((ctx.waccvgpr_dyn if use_acc else ctx.wvgpr_dyn)(vdata + _c(i), lane, val, exec_mask)) return UOp.sink(UOp.group(*stores).end(lane)) @@ -1741,49 +1720,19 @@ def _get_handler(inst: Inst) -> Callable[..., UOp]: # PROGRAM DECODE AND COMPILATION # ═══════════════════════════════════════════════════════════════════════════════ -_BARRIER_OPS = {ir3.SOPPOp.S_BARRIER, irc.SOPPOp.S_BARRIER} -if hasattr(ir4.SOPPOp, 'S_BARRIER_WAIT'): _BARRIER_OPS.add(ir4.SOPPOp.S_BARRIER_WAIT) -_BARRIER_SOP1_OPS: set = set() -if hasattr(ir4.SOP1Op, 'S_BARRIER_SIGNAL'): _BARRIER_SOP1_OPS.add(ir4.SOP1Op.S_BARRIER_SIGNAL) +_BARRIER_OPS = {ir3.SOPPOp.S_BARRIER, irc.SOPPOp.S_BARRIER, ir4.SOPPOp.S_BARRIER_WAIT} +_BARRIER_SOP1_OPS = {ir4.SOP1Op.S_BARRIER_SIGNAL} # Inline float constants (as bit patterns) for GPU instructions F32_INLINE = {240: 0x3f000000, 241: 0xbf000000, 242: 0x3f800000, 243: 0xbf800000, # 0.5, -0.5, 1.0, -1.0 244: 0x40000000, 245: 0xc0000000, 246: 0x40800000, 247: 0xc0800000, 248: 0x3e22f983} # 2.0, -2.0, 4.0, -4.0, 1/(2*pi) -instruction_cache:dict[tuple, list[tuple[int, int, tuple[UOp|None, UOp|None]]]] = {} - -# this is meant to replace the old emulator @uopfunc -def init_wave(wg:UOp, wave:UOp, sgpr:UOp, vgpr:UOp, args_ptr:UOp, gx:int, gy:int, lx:int, ly:int, total_threads:int, wave_size:int, - scratch_size:int, rsrc2:int, arch:str="rdna3", user_data_count:int=0): +def store_consts(sgpr:UOp, wave_size:int): ci = UOp.range(81, 2, AxisType.LOOP, dtype=dtypes.int) init_ints = sgpr.index(ci+128).store((ci<65).where(ci, 64-ci).cast(dtypes.uint32)).end(ci) - fi = UOp.range(len(F32_INLINE), 3, AxisType.LOOP, dtype=dtypes.int) - floats_buf = UOp(Ops.BINARY, arg=struct.pack(f"<{len(F32_INLINE)}I", *F32_INLINE.values())).bitcast(dtypes.uint32) - init_floats = sgpr.index(fi+240).store(floats_buf[fi]).end(fi) - # set RANGE registers - gidx, gidy, gidz = wg%gx, (wg//gx)%gy, wg//(gx*gy) - n_lanes = (total_threads-wave*wave_size).minimum(wave_size) - initial:list[tuple[int, sint]] = [(124, 0)] if wave_size == 32 else [] - if user_data_count: initial += [(i, UOp.variable(f"user_data_{i}", 0, dtypes.uint32.max, dtypes.uint32)) for i in range(user_data_count)] - else: initial += [(0, args_ptr.cast(dtypes.uint32)), (1, (args_ptr>>32).cast(dtypes.uint32))] - if arch == "rdna4": initial += [(ttmp[7].offset, (gidy&0xFFFF)|((gidz&0xFFFF)<<16)), (ttmp[9].offset, gidx)] - else: - sgpr_id = (rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_USER_SGPR_COUNT) >> hsa.AMD_COMPUTE_PGM_RSRC_TWO_USER_SGPR_COUNT_SHIFT - for enabled, gid in [(hsa.AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_X, gidx), - (hsa.AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_Y, gidy), - (hsa.AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_Z, gidz)]: - if rsrc2 & enabled: - initial.append((sgpr_id, gid)) - sgpr_id += 1 - initial += [(EXEC_LO.offset, ((UOp.const(1, dtypes.uint64)< UOp: body = sink.substitute({b:b.param_like(i, name=b.arg.name) for i,b in enumerate(bufs)}, walk=True) return body.call(*bufs, name=name) +instruction_cache:dict[tuple, list[tuple[int, int, tuple[UOp|None, UOp|None, dict[UOp, UOp]]]]] = {} + def instruction_call(inst:Inst, arch:str, renderer:Renderer) -> tuple[UOp|None, UOp|None]: entries = instruction_cache.setdefault((arch, type(inst), inst.size(), type(renderer), renderer.target, *(x.value for x in to_program_config)), []) inst_bytes = inst.to_bytes() @@ -1806,26 +1757,32 @@ def instruction_call(inst:Inst, arch:str, renderer:Renderer) -> tuple[UOp|None, sink = _get_handler(inst)(inst, ctx) base, mask = ctx.canonical_mask(inst_bytes) name = f"{_op_name(inst).lower()}_{base.to_bytes(len(inst_bytes), 'little').hex()}" - ret = lower_body(sink, name) if sink.src else None, ctx.branch_cond + call = lower_body(sink, name) if sink.src else None + used = UOp.sink(call, ctx.branch_cond).toposort() + ret = call, ctx.branch_cond, {p:index for index,p in ctx.sgpr_reads.items() if p in used} entries.append((base, mask, ret)) - call, cond = ret - words = {_Ctx.inst_word(i):_c((bits >> (32*i)) & 0xffffffff) for i in range(inst.size()//4)} - if call is not None: call = call.replace(src=(call.body, *(words.get(b, b) for b in call.src[1:]))) - return call, cond.substitute(words) if cond is not None else None + call, cond, reads = ret + words = {f"inst_w{i}":(bits >> (32*i)) & 0xffffffff for i in range(inst.size()//4)} + bindings = {_Ctx.inst_word(i):_c(words[f"inst_w{i}"]) for i in range(inst.size()//4)} + bindings.update({p:_Ctx.sgpr.index(_c(index.get_idx().sym_infer(words))).load() if index.get_valid().sym_infer(words) else _c(0) + for p,index in reads.items()}) + if call is not None: call = call.replace(src=(call.body, *(bindings.get(b, b) for b in call.src[1:]))) + return call, cond.substitute(bindings) if cond is not None else None def instruction_block(instructions:list[UOp]) -> tuple[UOp, list[UOp]]: - bufs = sorted({b for call in instructions for b in call.src[1:] if b.op is Ops.PARAM}, key=lambda b:b.arg.slot) + bufs = sorted({buf for call in instructions for b in call.src[1:] if (buf:=b.buf_uop).op is Ops.PARAM}, key=lambda b: b.arg.slot) params = {b:b.param_like(i, name=b.arg.name) for i,b in enumerate(bufs)} linear = list(UOp.sink(*params.values()).toposort())[:-1] - seen = set(linear) + seen: set[UOp] = set() for call in instructions: args = [params.get(b, b) for b in call.src[1:]] inputs = UOp.sink(*args) - nodes = inputs.toposort(lambda u:u not in seen) + nodes = inputs.toposort(lambda u:u.op is not Ops.PARAM and u not in seen) nodes.pop(inputs) seen.update(nodes) linear.extend(nodes) linear.append(call.replace(src=(call.body, *args))) + if _Ctx.sgpr in call.src[1:]: seen.clear() return UOp(Ops.PROGRAM, src=(sink:=UOp.sink(*linear, arg=KernelInfo("asm_block", optimized=True)), UOp(Ops.LINEAR, src=(*linear, sink)))), bufs def thread_branches(insts, blocks:dict[int, list[int]], paths:dict[int, dict[int, int]]): @@ -1923,8 +1880,8 @@ def lift(code:bytes, arch:str, renderer:Renderer, dispatch:tuple, name:str="") - possible = (possible-{pc}) | (set(cfg["paths"][pc]) or {-1}) routed = len(unconditional) != len(regions) axes = itertools.count(sum(inst.size() for inst in insts.values())) - gx, gy, gz, lx, ly, lz, rsrc2, scratch_size, user_data_count = dispatch - # construct CALL graph + gx, gy, gz, lx, ly, lz, rsrc2, scratch_size = dispatch + # init ranges wave_size, total_threads = _wave_size(arch), lx*ly*lz n_waves = (total_threads+wave_size-1)//wave_size wg = UOp.range(gx*gy*gz, 0, AxisType.LOOP, dtype=dtypes.uint64) @@ -1946,11 +1903,30 @@ def lift(code:bytes, arch:str, renderer:Renderer, dispatch:tuple, name:str="") - route = UOp.param(6, dtypes.int, (1,), name="route") taken, fallthrough = [UOp.param(i, dtypes.int, name=name, addrspace=AddrSpace.ALU) for i,name in enumerate(("taken", "fallthrough"), 7)] buffers[route] = (routes, 1) - init = init_wave(wg, wave, sgpr.shrink(((wave*SGPR_COUNT, (wave+1)*SGPR_COUNT),)), vgpr.shrink(((wave*256*wave_size, (wave+1)*256*wave_size),)), - args_ptr, gx, gy, lx, ly, total_threads, wave_size, scratch_size, rsrc2, arch, user_data_count) - user_vars = [u for u in init.body.toposort() if u.is_variable] - init = init.replace(src=(init.body.substitute({u:UOp.param(len(init.src)-1+i, u.dtype, name=u.arg.name, addrspace=AddrSpace.ALU) - for i,u in enumerate(user_vars)}, walk=True), *init.src[1:], *user_vars)) + # set per wave registers + wave_sgpr = sgpr.shrink(((wave*SGPR_COUNT, (wave+1)*SGPR_COUNT),)).after(wg) + wave_vgpr = vgpr.shrink(((wave*256*wave_size, (wave+1)*256*wave_size),)).after(wg) + consts = store_consts(wave_sgpr, wave_size) + gidx, gidy, gidz = wg%gx, (wg//gx)%gy, wg//(gx*gy) + n_lanes = (total_threads-wave*wave_size).minimum(wave_size) + runtime_sgprs = [(0, args_ptr.cast(dtypes.uint32)), (1, (args_ptr>>32).cast(dtypes.uint32))] + if arch == "rdna4": runtime_sgprs += [(ttmp[7].offset, (gidy&0xFFFF)|((gidz&0xFFFF)<<16)), (ttmp[9].offset, gidx)] + else: + sgpr_id = (rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_USER_SGPR_COUNT) >> hsa.AMD_COMPUTE_PGM_RSRC_TWO_USER_SGPR_COUNT_SHIFT + for enabled, gid in [(hsa.AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_X, gidx), + (hsa.AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_Y, gidy), + (hsa.AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_Z, gidz)]: + if rsrc2 & enabled: + runtime_sgprs.append((sgpr_id, gid)) + sgpr_id += 1 + runtime_sgprs += [(EXEC_LO.offset, ((_c(1, dtypes.uint64)< UOp: cond = UOp.const(True) @@ -1962,19 +1938,19 @@ def lift(code:bytes, arch:str, renderer:Renderer, dispatch:tuple, name:str="") - call, branch_cond = instruction_call(inst, arch, renderer) if branch_cond is not None: cond = branch_cond if call is not None: instructions.append(call) - values = {} + values = {ctx.scratch_stride:_c(scratch_size)} if routed: dsts = {kind:dst for dst,kind in cfg["paths"][block_pc].items()} target = cond.where(taken, fallthrough) if len(dsts) == 2 else taken - values = {taken:UOp.const(dsts[COND_TAKEN], dtypes.int), fallthrough:UOp.const(dsts[COND_NOT_TAKEN], dtypes.int)} if len(dsts) == 2 \ - else {taken:UOp.const(next(iter(dsts.values()), -1), dtypes.int)} + values.update({taken:_c(dsts[COND_TAKEN], dtypes.int), fallthrough:_c(dsts[COND_NOT_TAKEN], dtypes.int)} if len(dsts) == 2 + else {taken:_c(next(iter(dsts.values()), -1), dtypes.int)}) instructions.append(lower_body(route[0].store(target).sink(), "asm_branch")) if not instructions: return end body, bufs = instruction_block(instructions) args = [values[b] if b.addrspace is AddrSpace.ALU else buffers[b][0] for b in bufs] - args += [UOp.const(block_pc, dtypes.int).after(end, *scopes)] + args += [_c(block_pc, dtypes.int).after(end, *scopes)] params = [UOp.param(i, a.dtype, a.shape, addrspace=AddrSpace.ALU if i >= len(bufs) or bufs[i].addrspace is AddrSpace.ALU - else AddrSpace.GLOBAL) for i,a in enumerate(args)] + else AddrSpace.GLOBAL, name=bufs[i].arg.name if i < len(bufs) else None) for i,a in enumerate(args)] # lower the wave loop with the block, so repeated blocks share its indexing and route check too. w = UOp.range(n_waves, 0, AxisType.LOOP, dtype=dtypes.int) gates = () if block_pc in unconditional else (UOp.range(params[bufs.index(route)][w].load().eq(params[-1]).cast(dtypes.int), @@ -1989,7 +1965,7 @@ def lift(code:bytes, arch:str, renderer:Renderer, dispatch:tuple, name:str="") - loop = UOp.loop(next(axes)).replace(src=(UOp(Ops.NOOP).after(end),)) end = emit_dispatch(body, scopes+(loop,), loop) # waves that exited keep their route while the remaining waves repeat the loop. - check = loop_active(routes.after(end, *scopes, loop), active, UOp.const(pc, dtypes.int), UOp.const(n_waves, dtypes.int)) + check = loop_active(routes.after(end, *scopes, loop), active, _c(pc, dtypes.int), _c(n_waves, dtypes.int)) end = end.backedge(loop, active.after(check, *scopes, loop)[0].load()) else: # complete the block for every participating wave before entering its successors, including barrier successors. @@ -2004,12 +1980,12 @@ RELIFT = ContextVar("RELIFT", 0) # lift and recompile to assembly N times REPLAY_EMU = ContextVar("REPLAY_EMU", 0) def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int, lz: int, args_ptr: int, rsrc2: int = 0x19c, - scratch_size: int = 0, arch: str = "rdna3", user_data: list[int]|None = None, name: str = "asm") -> int: + scratch_size: int = 0, arch: str = "rdna3", name: str = "asm") -> int: """Execute AMD assembly program. scratch_size is private_segment_fixed_size from kernel descriptor (per-lane).""" with Context(DEV=DEV.target('CPU', renderer='LLVM'), NOOPT=1, CHECK_OOB=0, TUPLE_ORDER=0, EMULATED_DTYPES="", CAPTURE_PROCESS_REPLAY=REPLAY_EMU.value, DEFAULT_FLOAT="float32", DEFAULT_INT="int32", TRANSCENDENTAL=1): code = bytes(to_mv(lib, lib_sz)) - dispatch = (gx, gy, gz, lx, ly, lz, rsrc2, scratch_size, len(user_data or [])) + dispatch = (gx, gy, gz, lx, ly, lz, rsrc2, scratch_size) renderer = Device['CPU'].renderer if (key:=(code, arch, dispatch, name, type(renderer), renderer.target, *(x.value for x in to_program_config))) not in dispatch_cache: dispatch_cache[key] = to_program(lift(code, arch, renderer, dispatch, f"emu {name}"), renderer) @@ -2022,11 +1998,10 @@ def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int, break relift_lib = get_elf_section(ret.src[3].arg, ".text").content prg = dispatch_cache[key] - vals = {"args_ptr":args_ptr, **{f"user_data_{i}":v for i,v in enumerate(user_data or [])}} + vmem_buf = UOp.from_buffer(Buffer('CPU', 1 << 48, options=BufferSpec(external_ptr=0), opaque=0), dtypes.uint32) + args = [vmem_buf] * (max(p.arg.slot for p in prg.kernel_params) + 1) + for p in prg.kernel_params: + if p.addrspace is AddrSpace.ALU: args[p.arg.slot] = p.bind(args_ptr) with _MXCSRContext(): - vmem_buf = UOp.from_buffer(Buffer('CPU', 1 << 48, options=BufferSpec(external_ptr=0), opaque=0), dtypes.uint32) - args = [vmem_buf] * (max(p.arg.slot for p in prg.kernel_params) + 1) - for p in prg.kernel_params: - if p.addrspace is AddrSpace.ALU: args[p.arg.slot] = p.bind(vals[p.arg.name]) run_linear(UOp(Ops.LINEAR, src=(prg.call(*args),)), update_stats=False) return 0 diff --git a/tinygrad_repo/test/mockgpu/amd/pcode.py b/tinygrad_repo/test/mockgpu/amd/pcode.py index 9d989f77d7..336859c591 100644 --- a/tinygrad_repo/test/mockgpu/amd/pcode.py +++ b/tinygrad_repo/test/mockgpu/amd/pcode.py @@ -364,7 +364,8 @@ _FUNCS: dict[str, Callable[..., UOp]] = { 'isEven': lambda a: (UOp(Ops.TRUNC, src=(a,)).cast(dtypes.int) & _const(dtypes.int, 1)).eq(_const(dtypes.int, 0)), 'max': lambda a, b: UOp(Ops.MAX, src=(a, b)), 'min': lambda a, b: UOp(Ops.MAX, src=(a.neg(), b.neg())).neg(), - 'pow': lambda a, b: UOp(Ops.EXP2, src=(b.bitcast(dtypes.float32),)), + # pow(2.0, x): interpret the operand's float bits at its original width. + 'pow': lambda a, b: UOp(Ops.EXP2, src=(b.bitcast({2: dtypes.half, 4: dtypes.float32, 8: dtypes.float64}[b.dtype.itemsize]),)), 'fma': lambda a, b, c: a * b + c, 'i32_to_f32': lambda a: a.cast(dtypes.int).cast(dtypes.float32), 'u32_to_f32': lambda a: a.cast(dtypes.uint32).cast(dtypes.float32), diff --git a/tinygrad_repo/test/null/test_hcq2.py b/tinygrad_repo/test/null/test_hcq2.py index 3ee9c7c8b1..0b0dfea4b1 100644 --- a/tinygrad_repo/test/null/test_hcq2.py +++ b/tinygrad_repo/test/null/test_hcq2.py @@ -93,7 +93,7 @@ class TestHCQ2Deps(unittest.TestCase): def test_dependencies_through_selected_slices(self): b = UOp.param(0, dtypes.float32, 64, device=("NULL", "NULL:1")) for view in [b.mselect(0).shrink(((8, 16),)), b.shrink(((8, 16),)).mselect(0), b.shrink(((4, 32),)).mselect(0).shrink(((4, 12),))]: - tracker = hcq2.DepsTracker() + tracker = hcq2.DepsTracker(range(5)) tracker.access_resources([(view, True)], 0) self.assertEqual(tracker.access_resources([(b.mselect(1), False)], 1), []) self.assertEqual(tracker.access_resources([(b.mselect(0).shrink(((16, 24),)), False)], 2), []) @@ -102,7 +102,7 @@ class TestHCQ2Deps(unittest.TestCase): def test_disjoint_write_preserves_dependencies(self): b = UOp.param(0, dtypes.uint8, 16, device="NULL") for write in (False, True): - tracker = hcq2.DepsTracker() + tracker = hcq2.DepsTracker(range(5)) tracker.access_resources([(b.shrink(((0, 4),)), write)], 0) self.assertEqual(tracker.access_resources([(b.shrink(((4, 8),)), True)], 1), []) self.assertEqual(tracker.access_resources([(b.shrink(((0, 4),)), True)], 2), [0]) @@ -110,7 +110,7 @@ class TestHCQ2Deps(unittest.TestCase): def test_partial_write_preserves_dependencies(self): b = UOp.param(0, dtypes.uint8, 16, device="NULL") for write in (False, True): - tracker = hcq2.DepsTracker() + tracker = hcq2.DepsTracker(range(5)) tracker.access_resources([(b, write)], 0) self.assertEqual(tracker.access_resources([(b.shrink(((4, 12),)), True)], 1), [0]) self.assertEqual(tracker.access_resources([(b.shrink(((0, 4),)), True)], 2), [0]) @@ -119,16 +119,21 @@ class TestHCQ2Deps(unittest.TestCase): def test_write_waits_for_all_readers(self): b = UOp.param(0, dtypes.uint8, 16, device="NULL") - tracker = hcq2.DepsTracker() + tracker = hcq2.DepsTracker(range(5)) self.assertEqual(tracker.access_resources([(b, False)], 0), []) self.assertEqual(tracker.access_resources([(b, False)], 1), []) self.assertEqual(tracker.access_resources([(b, True)], 2), [0, 1]) self.assertEqual(tracker.access_resources([(b, False)], 3), [2]) + def test_write_waits_for_latest_reader_per_queue(self): + b, tracker = UOp.param(0, dtypes.uint8, 16, device="NULL"), hcq2.DepsTracker("AABAA") + for tag, view in enumerate([b, b, b, b.shrink(((0, 8),))]): self.assertEqual(tracker.access_resources([(view, False)], tag), []) + self.assertEqual(tracker.access_resources([(b, True)], 4), [1, 2, 3]) + def test_aliases_do_not_wait_on_themselves(self): b = UOp.param(0, dtypes.uint8, 16, device="NULL") for whole, part in ((True, False), (False, True), (True, True)): - tracker = hcq2.DepsTracker() + tracker = hcq2.DepsTracker(range(5)) tracker.access_resources([(b, True)], 0) self.assertEqual(tracker.access_resources([(b, whole), (b.shrink(((4, 12),)), part)], 1), [0]) self.assertEqual(tracker.access_resources([(b.shrink(((4, 12),)), True)], 2), [1]) diff --git a/tinygrad_repo/test/null/test_llm_server.py b/tinygrad_repo/test/null/test_llm_server.py index 9656b98fe8..fbae36adf7 100644 --- a/tinygrad_repo/test/null/test_llm_server.py +++ b/tinygrad_repo/test/null/test_llm_server.py @@ -44,6 +44,7 @@ class TestLLMServer(unittest.TestCase): cls.mock_tok.encode = Mock(return_value=[200, 201, 202]) cls.mock_tok.decode = Mock(return_value="Hello") cls.mock_tok.stream_decoder = Mock(return_value=lambda tid=None: "Hello" if tid is not None else "") + cls.mock_tok._special_tokens = {} cls.mock_tok.preset = "llama3" cls.mock_tok.bos_id = 1 cls.mock_tok.eos_id = 999 @@ -245,6 +246,7 @@ class TestLLMToolCalls(unittest.TestCase): cls.mock_tok = Mock() cls.mock_tok.encode = Mock(return_value=[200, 201, 202]) cls.mock_tok.decode = Mock(return_value="") + cls.mock_tok._special_tokens = {} cls.mock_tok.preset = "qwen2" cls.mock_tok.bos_id, cls.mock_tok.eos_id, cls.mock_tok.eot_id = None, 999, None cls.mock_tok.is_end = Mock(return_value=False) @@ -347,6 +349,20 @@ class TestLLMToolCalls(unittest.TestCase): self.assertEqual(second.choices[0].message.content, "done") self.assertEqual(second.choices[0].finish_reason, "stop") + def test_kimi_tool_call(self): + self.set_output('a b<|close|>think<|sep|><|open|>response<|sep|>c {} d<|close|>response<|sep|>' + '<|open|>tools<|sep|><|open|>call tool="read" index="1"<|sep|><|open|>argument key="x"y" type="string"<|sep|>\n\n' + '<|close|>argument<|sep|><|open|>argument key="n" type="number"<|sep|>3<|close|>argument<|sep|><|close|>call<|sep|><|close|>tools<|sep|>') + with (patch.object(self.mock_tok, "_special_tokens", {"<|close|>": 0}), + patch.object(self.server.template, "render", return_value='<|open|>message role="assistant"<|sep|><|open|>think<|sep|>')): + chunks = list(self.client.chat.completions.create(model="tool-model", messages=[{"role":"user", "content":"hi"}], tools=self.tools(), + stream=True)) + self.assertEqual(''.join(getattr(c.choices[0].delta, 'reasoning_content', '') for c in chunks if c.choices), 'a b') + self.assertEqual(''.join(c.choices[0].delta.content or '' for c in chunks if c.choices), 'c {} d') + calls = [tc for c in chunks if c.choices for tc in c.choices[0].delta.tool_calls or []] + self.assertEqual([(tc.function.name, json.loads(tc.function.arguments)) for tc in calls], [('read', {'x"y':'\n\n', 'n':3})]) + self.assertEqual(chunks[-1].choices[0].finish_reason, 'tool_calls') + class TestTransformerGenerate(unittest.TestCase): def test_warmup(self): model, calls = Transformer(TEST_CONFIG), [] @@ -357,7 +373,7 @@ class TestTransformerGenerate(unittest.TestCase): self.assertEqual(calls, [[0], [0]]) def test_template_starts_reasoning(self): - router = StreamRouter(reasoning=True) + router = StreamRouter({}, reasoning=True) self.assertEqual(list(router.route("reasoninganswer")), [("reasoning_content", "reasoning"), ("content", "answer")]) diff --git a/tinygrad_repo/test/null/test_llm_tokenizer.py b/tinygrad_repo/test/null/test_llm_tokenizer.py index fa7f42f10d..f2819985be 100644 --- a/tinygrad_repo/test/null/test_llm_tokenizer.py +++ b/tinygrad_repo/test/null/test_llm_tokenizer.py @@ -162,6 +162,15 @@ class TestLLMTokenizer(unittest.TestCase): self.assertEqual(split["kimi-k2"](s), go, f"kimi-k2 {s!r}") self.assertEqual(split["kimi-k2"]("a.\n/b"), ["a", ".\n", "/b"]) + def test_joyai_llm_split(self): + tok = SimpleTokenizer({}, {}, "joyai-llm") + tok._encode_word = lambda word: [word.decode()] + self.assertEqual(tok.encode("é café"), ["é", " café"]) + self.assertEqual(tok.encode(" 12345"), [" ", "123", "45"]) + self.assertEqual(tok.encode("it's I'M don'T"), ["it", "'s", " I", "'M", " don", "'T"]) + self.assertEqual(tok.encode(" 你好hello"), [" ", "你好", "hello"]) + self.assertEqual(tok.encode("a\u200b"), ["a", "\u200b"]) + def test_stream_decoder(self): """stream_decoder buffers incomplete UTF-8: token 25677 has 3/4 of emoji, token 138 completes it.""" bs = [*range(33, 127), *range(161, 173), *range(174, 256)] diff --git a/tinygrad_repo/test/null/test_tensor.py b/tinygrad_repo/test/null/test_tensor.py index c51d5a57ba..0b45dc00b2 100644 --- a/tinygrad_repo/test/null/test_tensor.py +++ b/tinygrad_repo/test/null/test_tensor.py @@ -7,7 +7,8 @@ from tinygrad.dtype import DTYPES_DICT, DType from tinygrad.uop.ops import UOp, Ops from tinygrad.renderer.ptx import PTXRenderer from tinygrad.renderer.nir import NIRRenderer -from tinygrad.codegen import to_program +from tinygrad.codegen import full_rewrite_to_sink +from tinygrad.codegen.late.linearizer import linearize x_init = np.random.randn(1,3).astype(np.float32) W_init = np.random.randn(3,3).astype(np.float32) @@ -56,9 +57,7 @@ class TestIdxUpcast(unittest.TestCase): for si in linear.src: ast = si.src[0] if ast.op is Ops.SINK: - renderer = Device[si.src[1].buffer.device].renderer - prg = to_program(ast, renderer) - return tuple(prg.src[1].src) + return tuple(linearize(full_rewrite_to_sink(ast, Device[si.src[1].buffer.device].renderer))) def _assert(self, dtype: DType, a: Tensor): uops = self._schedule_render(a) diff --git a/tinygrad_repo/test/null/test_uops.py b/tinygrad_repo/test/null/test_uops.py index 9065d549de..504d9cc7df 100644 --- a/tinygrad_repo/test/null/test_uops.py +++ b/tinygrad_repo/test/null/test_uops.py @@ -73,6 +73,26 @@ class TestDTypeFromUOp(unittest.TestCase): type_verify(out.sink(), spec_program) class TestMemoryCoalescing(unittest.TestCase): + def test_store_deps_ignore_reg_stores(self): + buf, reg = UOp.param(0, dtypes.int, 2), UOp.placeholder((2,), dtypes.int, 1, addrspace=AddrSpace.REG) + lanes = [buf.index(i).store(UOp.const(i, dtypes.int).after(reg.index(i).store(i))) for i in range(2)] + sink = memory_coalescing(UOp.sink(*lanes), Renderer(Target())) + self.assertEqual(sum(u.op is Ops.STORE and u.src[0].addrspace is AddrSpace.GLOBAL for u in sink.toposort()), 1) + + def test_store_deps_ignore_local_stores(self): + buf, local = UOp.param(0, dtypes.int, 4), UOp.placeholder((4,), dtypes.int, 1, addrspace=AddrSpace.LOCAL) + lanes = [buf.index(0).store(0)]+[buf.index(i).store(UOp.const(i, dtypes.int).after(local.index(i).store(i))) for i in range(1, 4)] + sink = memory_coalescing(UOp.sink(*lanes), Renderer(Target())) + self.assertEqual(sum(u.op is Ops.STORE and u.src[0].addrspace is AddrSpace.GLOBAL for u in sink.toposort()), 1) + + def test_store_deps_local_waits_on_global_stores(self): + buf, local = UOp.param(0, dtypes.int, 2), UOp.placeholder((2,), dtypes.int, 1, addrspace=AddrSpace.LOCAL) + # the GLOBAL stores coalesce, so the LOCAL stores must not: t1 waits on s1 and s0 waits on t0 + s1, t0 = buf.index(1).store(5), local.index(0).store(7) + s0, t1 = buf.index(0).store(local.after(t0).index(0).load()), local.index(1).store(buf.after(s1).index(1).load()) + sink = memory_coalescing(UOp.sink(s0, s1, t0, t1), Renderer(Target())) + self.assertEqual(sum(u.op is Ops.STORE and u.src[0].addrspace is AddrSpace.LOCAL for u in sink.toposort()), 2) + def test_volatile_view_not_coalesced(self): buf = UOp.param(0, dtypes.uint32, 4, volatile=True).bitcast(dtypes.int32) sink = memory_coalescing(UOp.sink(*(buf.index(i).load() for i in range(4))), Renderer(Target())) @@ -107,6 +127,20 @@ class TestSafeCast(unittest.TestCase): self.assertEqual(a.cast(dtypes.int64).cast(dtypes.uint8).simplify(), a) self.assertEqual(a.cast(dtypes.uint32).cast(dtypes.uint8).simplify(), a) + def test_const_float_to_int(self): + for dtype in dtypes.ints: + for value, expected in ((math.nan, 0), (math.inf, dtype.max), (-math.inf, dtype.min), (1.9, 1)): + with self.subTest(dtype=dtype, value=value): + self.assertEqual(UOp.const(value).cast(dtype).ssimplify(), expected) + + def test_const_float_to_int_wraps(self): + for value, dtype, expected in ((128.0, dtypes.int8, -128), (256.0, dtypes.uint8, 0), (-1.0, dtypes.uint8, 255)): + with self.subTest(dtype=dtype, value=value): + self.assertEqual(UOp.const(value).cast(dtype).ssimplify(), expected) + + def test_const_nan_to_int_tensor(self): + self.assertEqual(Tensor(math.nan).cast(dtypes.int32).item(), 0) + def test_remove_intermediate_cast(self): a = UOp.variable("a", 0., 100., dtype=dtypes.half) self.assertEqual(a.cast(dtypes.double).cast(dtypes.float).simplify(), a.cast(dtypes.float)) diff --git a/tinygrad_repo/test/null/test_uops_stats.py b/tinygrad_repo/test/null/test_uops_stats.py index 89221f86a2..2192c289bc 100644 --- a/tinygrad_repo/test/null/test_uops_stats.py +++ b/tinygrad_repo/test/null/test_uops_stats.py @@ -107,6 +107,13 @@ class TestUOpsStats(unittest.TestCase): self.assertEqual(prg.src[0].arg.estimates.mem, 4) self.assertEqual(prg.src[0].arg.estimates.lds, 4) + def test_triangular_loop_estimate(self): + buf = UOp.param(0, dtypes.int32, 6) + i = UOp.range(3, 0) + j = UOp.range(i+1, 1) + prg = to_program(buf.index(i*(i+1)//2+j).store(j).end(j, i).sink(arg=KernelInfo()), X86Renderer(Target("CPU", arch="x86_64"))) + self.assertEqual(prg.src[0].arg.estimates.lds, 3*3*4) + def test_simple_add(self): a = Tensor.empty(100,100) b = Tensor.empty(100,100) diff --git a/tinygrad_repo/test/runtime/test_arange.py b/tinygrad_repo/test/runtime/test_arange.py index bd46d8546f..1642a85111 100644 --- a/tinygrad_repo/test/runtime/test_arange.py +++ b/tinygrad_repo/test/runtime/test_arange.py @@ -39,6 +39,10 @@ class TestArange(unittest.TestCase): DSET, DDIM = 2048, 32 class TestIndexing(unittest.TestCase): + def test_index_compared_with_offset(self): + a, b = Tensor([1, 2, 3, 4]), Tensor([0, 2, 2, 4]) + self.assertEqual((a == b + 1).float()[Tensor([0, 2, 1])].tolist(), [1.0, 1.0, 0.0]) + def test_arange_2_reduce(self): needle = Tensor.zeros(16384, dtype=dtypes.int).contiguous() needle[1337] = 1 @@ -121,6 +125,18 @@ class TestIndexing(unittest.TestCase): self.assertEqual(((x*x) == 2.).where(data, 0.).tolist(), [0.]*4) self.assertEqual(((x*x) == 2.).reshape(2, 2).where(data.reshape(2, 2), 0.).sum(axis=1).tolist(), [0.]*2) + def test_index_collapse_float_idx(self): + x = Tensor.arange(4).cast(dtypes.float32) + self.assertEqual((x == 1.5).where(x+1, 0.).sum().item(), 0.) + self.assertEqual(((Tensor.arange(4)+1).cast(dtypes.float32) == 2.5).where(x+1, 0.).sum().item(), 0.) + self.assertEqual(((Tensor.arange(8).cast(dtypes.float32)+1e8) < 100000008.).where(1, 0).sum().item(), 5) + + def test_index_collapse_cast_wraps(self): + r, y = Tensor.arange(8, dtype=dtypes.int32), Tensor([254], dtype=dtypes.int32).realize() + self.assertEqual(((r+y).cast(dtypes.uint8) == Tensor([1], dtype=dtypes.uint8).realize()).where(r+1, 0).sum().item(), 4) + self.assertEqual(((r+y).cast(dtypes.int64) == Tensor([2**32+254], dtype=dtypes.int64).realize()).where(r+1, 0).sum().item(), 0) + self.assertEqual((Tensor.arange(260).cast(dtypes.uint8) == 1).where(Tensor.arange(260)+1, 0).sum().item(), 260) + def test_index_mnist(self, noopt=1, op_limit=512*784*13, split_reduceop=0): # WEBGPU generates more ops due to bitpacking of < 4-byte dtypes if Device.DEFAULT == "WEBGPU": op_limit *= 15 diff --git a/tinygrad_repo/test/runtime/test_assign.py b/tinygrad_repo/test/runtime/test_assign.py index 6fc3c7fb11..c6ed5f64c0 100644 --- a/tinygrad_repo/test/runtime/test_assign.py +++ b/tinygrad_repo/test/runtime/test_assign.py @@ -257,6 +257,17 @@ class TestAssign(unittest.TestCase): np.testing.assert_equal(b0.numpy(), 128) np.testing.assert_equal(b1.numpy(), 608) + def test_assign_reads_behind_other_after(self): + a = Tensor([1, 2]).contiguous().realize() + b = Tensor([0, 0]).contiguous().realize() + b.assign(a + 1) + a.assign(b.sum().expand(2)) + GlobalCounters.reset() + Tensor.realize(a, b) + assert_kernel_count(2) + self.assertEqual(a.tolist(), [5, 5]) + self.assertEqual(b.tolist(), [2, 3]) + def test_crossunder_assign(self): a = Tensor.full((4,), 2).contiguous().realize() b = Tensor.full((4,), 3).contiguous().realize() diff --git a/tinygrad_repo/test/runtime/test_attention.py b/tinygrad_repo/test/runtime/test_attention.py index ee12332321..e7d511e284 100644 --- a/tinygrad_repo/test/runtime/test_attention.py +++ b/tinygrad_repo/test/runtime/test_attention.py @@ -3,7 +3,7 @@ import numpy as np from tinygrad import Tensor, dtypes, nn from tinygrad.llm.model import ( GatedDeltaNetBlock, SSMConfig, TransformerBlock, TransformerConfig, MLATransformerBlock, Indexer, IndexerConfig, - HyperConnection, HCConfig, + HyperConnection, HCConfig, AttentionResidual, apply_rope as apply_rope_new, precompute_freqs_cis, pairwise_topk, ) from test.helpers import slow @@ -83,6 +83,14 @@ class TestAttention(unittest.TestCase): self.assertIs(returned_mask, mask) np.testing.assert_allclose(actual.numpy(), expected, rtol=1e-6, atol=1e-7) + def test_attention_residual(self): + # rmsnorm makes the streams [1, -3] into scores [log 3, -log 3], softmax [9/10, 1/10] mixes them to 0.6, the sublayer squares it + for new_block, expected in ((True, [1, -3, 0.36]), (False, [1, -3 + 0.36])): + res = AttentionResidual(self._make_config(dim=1), new_block) + res.weight = Tensor([np.log(3)]) + actual, _ = res(Tensor([[[[1.], [-3.]]]]), lambda x: (x.square(), None)) + np.testing.assert_allclose(actual.numpy().flatten(), expected, rtol=1e-5) + def test_indexer(self): indexer = Indexer(self._make_config(dim=2, q_lora_rank=2, indexer=IndexerConfig(n_heads=1, head_dim=2, top_k=2, kpool=2, norm_eps=1e-6))) indexer.attn_q_b.weight, indexer.attn_k.weight = Tensor.eye(2), Tensor.eye(2) @@ -131,7 +139,8 @@ class TestGatedDeltaNetBlock(unittest.TestCase): return TransformerConfig(**({"num_blocks":1, "dim":8, "hidden_dim":16, "n_heads":1, "n_kv_heads":1, "norm_eps":1e-5, "vocab_size":32, "head_dim":8, "rope_theta":10000.0, "rope_dim":8, "v_head_dim":8, "max_context":4, "ssm_layers":(True,), - "ssm":SSMConfig(conv_kernel=2, state_size=4, group_count=1, time_step_rank=1, inner_size=4)} | kwargs)) + "ssm":SSMConfig(conv_kernel=2, state_size=4, group_count=1, time_step_rank=1, inner_size=4, + silu_gate=True)} | kwargs)) def _make_block(self, config:TransformerConfig) -> GatedDeltaNetBlock: block = GatedDeltaNetBlock(config, config.ssm) @@ -319,7 +328,7 @@ class TestGatedDeltaNetBlock(unittest.TestCase): # full prefill is proven equivalent to decode by test_gatedeltanet_reference_and_reset (delta rule) and # test_kda_prefill_matches_decode (kda), so use it as the baseline and only exercise multi-chunk handoffs here for kda in (False, True): - ssm = SSMConfig(conv_kernel=2, state_size=4, group_count=1, time_step_rank=1, inner_size=4, kda=kda) + ssm = SSMConfig(conv_kernel=2, state_size=4, group_count=1, time_step_rank=1, inner_size=4, kda=kda, silu_gate=not kda) config = self._make_config(ssm=ssm) if kda: block = GatedDeltaNetBlock(config, config.ssm) diff --git a/tinygrad_repo/test/runtime/test_custom_kernel.py b/tinygrad_repo/test/runtime/test_custom_kernel.py index 868bed705c..b3b3f73601 100644 --- a/tinygrad_repo/test/runtime/test_custom_kernel.py +++ b/tinygrad_repo/test/runtime/test_custom_kernel.py @@ -4,7 +4,7 @@ from tinygrad.engine.realize import run_linear, get_runtime from tinygrad.codegen import to_program import numpy as np from tinygrad.dtype import AddrSpace, dtypes, Invalid -from tinygrad.helpers import getenv, DEV +from tinygrad.helpers import getenv from tinygrad.schedule.rangeify import BufferizeOpts from tinygrad.uop.ops import KernelInfo, AxisType, Ops, uopfunc from tinygrad.codegen.opt import Opt, OptOps @@ -289,6 +289,69 @@ class TestCustomKernel(unittest.TestCase): tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0] self.assertTrue(tst.allclose(a@b, atol=1e-3).item()) + def test_unused_reduce_dependency(self): + for value in (UOp.const(2), UOp.const((1, 2, 3, 4)).reshape((2, 2))): + with self.subTest(shape=value.shape): + def kernel(C:UOp) -> UOp: + i = UOp.range(3, 0, AxisType.LOOP) + cell = C.after(i)[0] + v = value.after(cell.store(cell + 1)) + return C.after(v.reduce(i, arg=Ops.ADD)).sink(arg=KernelInfo(opts_to_apply=())) + self.assertEqual(Tensor.zeros(1, dtype=dtypes.int).contiguous().custom_kernel(fxn=kernel)[0].tolist(), [3]) + + def test_range_merge_keeps_after_domain(self): + # C is read once per i or once per o, so it counts to 4, and merging o and i must not make that 16 + o, i = UOp.range(4, 0, AxisType.LOOP), UOp.range(4, 1, AxisType.LOOP) + stores = {"after_i": lambda C,A: (c:=C.after(i)[0]).store(c+1), + "after_o": lambda C,A: UOp.group((c:=C.after(o)[0]).store(c+1), A.after(o, i)[0].store(1)), + "load_after_o": lambda C,A: C.after(o, i)[0].store(C.after(o)[0].load()+1)} + for name, store in stores.items(): + with self.subTest(name): + def kernel(C:UOp, A:UOp, store=store) -> UOp: return store(C, A).end(o, i).sink(arg=KernelInfo(opts_to_apply=())) + self.assertEqual(Tensor.zeros(1).contiguous().custom_kernel(Tensor.empty(1), fxn=kernel)[0].item(), 4) + + def test_range_merge_keeps_index_domain(self): + # C[c] is indexed by o alone, so each cell counts to 2 while the B store merges o and i + o, i = UOp.range(4, 0, AxisType.LOOP), UOp.range(4, 1, AxisType.LOOP) + for name, c in {"div": o//2, "cmp": (o<2).cast(dtypes.int)}.items(): + with self.subTest(name): + def kernel(C:UOp, B:UOp, c=c) -> UOp: + return UOp.group(C[c].store(C[c]+1), B[o*4+i].store(1)).end(o, i).sink(arg=KernelInfo(opts_to_apply=())) + self.assertEqual(Tensor.custom_kernel(Tensor.zeros(2).contiguous(), Tensor.zeros(16).contiguous(), fxn=kernel)[0].tolist(), [2, 2]) + + def test_range_merge_keeps_inner_reduce(self): + # the inner reduce ends r1 alone, so r0 and r1 don't merge. the inner sum is 4, the outer adds it 4*4*4 times + def kernel(C:UOp, A:UOp) -> UOp: + r0, r1, r2 = [UOp.range(4, n, AxisType.UPCAST) for n in range(3)] + return C[0].store(A[r2].reduce(r1, arg=Ops.ADD).reduce(r0, r1, r2, arg=Ops.ADD)).sink(arg=KernelInfo(opts_to_apply=())) + self.assertEqual(Tensor.custom_kernel(Tensor.empty(1), Tensor.ones(4).contiguous(), fxn=kernel)[0].item(), 256) + + def test_range_merge_keeps_reduce_order(self): + # the reduce body stores, so it loops a, k, b in order, and the digits of S are the A[a*2+b] it visits + def kernel(O:UOp, S:UOp, A:UOp) -> UOp: + a, k, b = [UOp.range(2, n, AxisType.LOOP) for n in range(3)] + v = A[a*2+b].after((cell:=S.after(a, k, b)[0]).store(cell*10+A[a*2+b])) + return O[0].store(v.reduce(a, k, b, arg=Ops.ADD)).sink(arg=KernelInfo(opts_to_apply=())) + o, s = Tensor.custom_kernel(Tensor.empty(1, dtype=dtypes.int), Tensor.zeros(1, dtype=dtypes.int).contiguous(), + Tensor([1, 2, 3, 4], dtype=dtypes.int), fxn=kernel)[:2] + Tensor.realize(o, s) + self.assertEqual(s.item(), 12123434) + + def test_range_merge_keeps_loop_order(self): + # loops nest by axis id whatever order END lists them, so x is outer and the digits of S are A[0], A[2], A[1], A[3] + def kernel(S:UOp, A:UOp) -> UOp: + x, y = UOp.range(2, 0, AxisType.LOOP), UOp.range(2, 1, AxisType.LOOP) + return (cell:=S.after(x, y)[0]).store(cell*10+A[y*2+x]).end(y, x).sink(arg=KernelInfo(opts_to_apply=())) + s = Tensor.custom_kernel(Tensor.zeros(1, dtype=dtypes.int).contiguous(), Tensor([1, 2, 3, 4], dtype=dtypes.int), fxn=kernel)[0] + self.assertEqual(s.item(), 1324) + + def test_range_merge_then_shrink(self): + # r0<2 shrinks the merged range of 8 to 4, which is the same UOp as r0, so it isn't shrunk again to 2 + def kernel(C:UOp) -> UOp: + r0, r1 = UOp.range(4, 0, AxisType.LOOP), UOp.range(2, 1, AxisType.LOOP) + return (cell:=C[(r0*2+r1).valid(r0<2)]).store(cell+1).end(r0, r1).sink(arg=KernelInfo(opts_to_apply=())) + self.assertEqual(Tensor.zeros(8, dtype=dtypes.int).contiguous().custom_kernel(fxn=kernel)[0].tolist(), [1, 1, 1, 1, 0, 0, 0, 0]) + @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "test requires locals") def test_group_reduce_split_range(self): # j%2 splits j into two ranges, both are still LOCAL @@ -458,6 +521,14 @@ class TestCustomKernel(unittest.TestCase): out = Tensor.custom_kernel(Tensor.zeros(2, dtype=dtypes.int).contiguous().realize(), fxn=functools.partial(kernel, loop=loop))[0] self.assertEqual(out.tolist(), [7, 9]) + def test_store_value_behind_coalesced_store(self): + def c(v): return UOp.const(v, dtypes.int) + for stores, want in ((lambda A: (A[0].store(7), A[1].store(A.after(A[0].store(7))[0]+5)), [7, 12, 5, 0]), + (lambda A: (A[1].store(c(1).after(A[0].store(7)) + A[2].load()),), [7, 6, 5, 0]), + (lambda A: (A[1].store(c(1).after(A[2].store(2))), A[3].store(c(3).after(A[0].store(10)))), [10, 1, 2, 3])): + out = Tensor([0, 0, 5, 0], dtype=dtypes.int).contiguous().custom_kernel(fxn=lambda A: UOp.sink(*stores(A), arg=KernelInfo(opts_to_apply=())))[0] + self.assertEqual(out.tolist(), want) + def test_deleted_store_keeps_deps(self): for f in (lambda a, s: s[1].store(s[1].load()), lambda a, s: s[UOp.invalid()].store(5), lambda a, s: s[1].store(UOp.invalid()), lambda a, s: a[UOp.invalid()].store(s[1].load())): @@ -499,6 +570,42 @@ class TestCustomKernel(unittest.TestCase): a = Tensor.arange(16).reshape(4, 4).float().contiguous().realize() self.assertEqual(Tensor.custom_kernel(Tensor.empty(4, 4), a, fxn=kernel)[0].tolist(), a[:, :1].pad(((0, 0), (0, 3))).tolist()) + def test_shrink_range_keeps_ungated_index(self): + # q is ungated in C[q], so q keeps its 4 iterations after r shrinks to its gate r<2 + q, r = UOp.range(4, 0, AxisType.LOOP), UOp.range(4, 1, AxisType.LOOP) + def kernel(C:UOp, D:UOp) -> UOp: + return UOp.group(C[q].store(1), D[(q*4+r).valid((q<2)&(r<2))].store(1)).end(q, r).sink(arg=KernelInfo(opts_to_apply=())) + c, d = Tensor.custom_kernel(Tensor.zeros(4, dtype=dtypes.int).contiguous(), Tensor.zeros(16, dtype=dtypes.int).contiguous(), fxn=kernel)[:2] + Tensor.realize(c, d) + self.assertEqual(c.tolist(), [1, 1, 1, 1]) + self.assertEqual(d.tolist(), [1, 1, 0, 0, 1, 1, 0, 0] + [0]*8) + + def test_shrink_range_keeps_ungated_store(self): + # C[0] stores r every iteration, so r keeps its 4 iterations though D gates it to r<2 + r = UOp.range(4, 0, AxisType.LOOP) + def kernel(C:UOp, D:UOp) -> UOp: + return UOp.group(C[0].store(r.cast(dtypes.int)), D[r.valid(r<2)].store(1)).end(r).sink(arg=KernelInfo(opts_to_apply=())) + c, d = Tensor.custom_kernel(Tensor.zeros(1, dtype=dtypes.int).contiguous(), Tensor.zeros(4, dtype=dtypes.int).contiguous(), fxn=kernel)[:2] + Tensor.realize(c, d) + self.assertEqual((c.tolist(), d.tolist()), ([3], [1, 1, 0, 0])) + + def test_shrink_range_keeps_store_gated_otherwise(self): + # C[0] is gated by r>=2, not r UOp: + return UOp.group(C[UOp.const(0).valid(r>=2)].store(r.cast(dtypes.int)), D[r.valid(r<2)].store(1)).end(r).sink(arg=KernelInfo(opts_to_apply=())) + c, d = Tensor.custom_kernel(Tensor.zeros(1, dtype=dtypes.int).contiguous(), Tensor.zeros(4, dtype=dtypes.int).contiguous(), fxn=kernel)[:2] + Tensor.realize(c, d) + self.assertEqual((c.tolist(), d.tolist()), ([3], [1, 1, 0, 0])) + + def test_shrink_range_keeps_symbolic_end(self): + # D gates r to r<4, but r only runs to n=2 + def kernel(D:UOp, N:UOp) -> UOp: + r = UOp.range(N, 0, AxisType.LOOP) + return D[r.valid(r<4)].store(1).end(r).sink(arg=KernelInfo(opts_to_apply=())) + d = Tensor.custom_kernel(Tensor.zeros(8, dtype=dtypes.int).contiguous(), Tensor(Variable("n", 1, 8, dtypes.int).bind(2)), fxn=kernel)[0] + self.assertEqual(d.tolist(), [1, 1, 0, 0, 0, 0, 0, 0]) + @unittest.skipIf(Device.DEFAULT == "WEBGPU", "kernel timing not supported") def test_invalids_into_custom_kernel_with_beam(self): a = Tensor.full((4, 4), 3.).contiguous() @@ -592,16 +699,17 @@ class TestCustomKernel(unittest.TestCase): self.assertEqual(a.flatten().tolist(), [2, 2, 3, 3]) self.assertEqual(a.shape, (2, 2)) - @unittest.skipUnless((isinstance(Device[Device.DEFAULT].renderer, CStyleLanguage) or Device.DEFAULT == "PYTHON") and + @unittest.skipUnless((isinstance(Device[Device.DEFAULT].renderer, (CStyleLanguage, LLVMRenderer)) or Device.DEFAULT == "PYTHON") and Device.DEFAULT != "WEBGPU", "binary not supported on this backend") - @unittest.skipIf(Device.DEFAULT == "AMD" and DEV.interface.startswith("MOCK"), "AMD emulator doesn't support direct pc read and write") def test_binary(self): payload = bytes(range(256)) + def view(x): return x.reshape((16, 16)).permute((1, 0)).reshape((256,)) def kernel(out:UOp): i = UOp.range(len(payload), 0) - data = UOp(Ops.BINARY, arg=payload) + data = view(UOp(Ops.BINARY, arg=payload)) return out[i].store(data[i]).end(i).sink(arg=KernelInfo(name="binary", opts_to_apply=())) - self.assertEqual(Tensor.empty(len(payload), dtype=dtypes.uint8).custom_kernel(fxn=kernel)[0].tolist(), list(payload)) + self.assertEqual(Tensor.empty(len(payload), dtype=dtypes.uint8).custom_kernel(fxn=kernel)[0].tolist(), + view(Tensor(payload, dtype=dtypes.uint8)).tolist()) def test_one_lane_stack_runtime_index(self): def kernel(A:UOp): return A[0].store(UOp.stack(UOp.const(42, dtypes.int)).index(A[1].load())).sink(arg=KernelInfo(opts_to_apply=())) @@ -638,6 +746,18 @@ class TestCallInKernel(unittest.TestCase): out2 = Tensor.custom_kernel(Tensor.empty(1, dtype=x.dtype), x, fxn=main)[0] self.assertEqual(out2.tolist(), [N*(N+3)//2]) + def test_reduce_over_call_keeps_loop_order(self): + # a call can hide stores, so the reduce loops a, k, b in order, and the digits of S are the A[a*2+b] it visits + @uopfunc + def update(S:UOp, A:UOp) -> UOp: return S[0].store(S[0]*10+A[0]).sink() + def kernel(O:UOp, S:UOp, A:UOp) -> UOp: + a, k, b = [UOp.range(2, n, AxisType.LOOP) for n in range(3)] + return O[0].store(A[a*2+b].after(update(S.after(k), A[a*2+b])).reduce(a, k, b, arg=Ops.ADD)).sink(arg=KernelInfo(opts_to_apply=())) + o, s = Tensor.custom_kernel(Tensor.empty(1, dtype=dtypes.int), Tensor.zeros(1, dtype=dtypes.int).contiguous(), + Tensor([1, 2, 3, 4], dtype=dtypes.int), fxn=kernel)[:2] + Tensor.realize(o, s) + self.assertEqual(s.item(), 12123434) + def test_nested_call(self): @uopfunc def incr(out:UOp, A:UOp): diff --git a/tinygrad_repo/test/runtime/test_dtype.py b/tinygrad_repo/test/runtime/test_dtype.py index 73b7ede869..53493abb84 100644 --- a/tinygrad_repo/test/runtime/test_dtype.py +++ b/tinygrad_repo/test/runtime/test_dtype.py @@ -338,6 +338,9 @@ class TestUint8DType(TestDType): @unittest.skipIf(Device.DEFAULT == "CUDA" or isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "cuda saturation works differently") def test_uint8_to_int8_overflow(self): _test_op(lambda: Tensor([255, 254, 253, 252], dtype=dtypes.uint8).cast(dtypes.int8), dtypes.int8, [-1, -2, -3, -4]) + @unittest.skipIf(Device.DEFAULT == "CUDA" or isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "cuda saturation works differently") + def test_uint8_to_int8_overflow_compare(self): + self.assertEqual((Tensor([128, 5], dtype=dtypes.uint8).cast(dtypes.int8) < 0).tolist(), [True, False]) class TestBitCast(unittest.TestCase): def test_subword_slice(self): diff --git a/tinygrad_repo/test/runtime/test_dtype_alu.py b/tinygrad_repo/test/runtime/test_dtype_alu.py index 1af2409b77..1c005fc5ed 100644 --- a/tinygrad_repo/test/runtime/test_dtype_alu.py +++ b/tinygrad_repo/test/runtime/test_dtype_alu.py @@ -5,6 +5,7 @@ from tinygrad.helpers import EMULATED_DTYPES, DEV, getenv from tinygrad.tensor import _to_np_dtype from tinygrad.runtime.ops_python import from_storage_scalar from tinygrad.renderer.ptx import PTXRenderer +from tinygrad.renderer.cstyle import CUDARenderer from tinygrad.renderer.nir import NIRRenderer from tinygrad.renderer.llvmir import CPULLVMRenderer from tinygrad.renderer.x86 import X86Renderer @@ -334,6 +335,28 @@ class TestDTypeALU(unittest.TestCase): self.assertEqual(((t+t*(2**31-1))//2).tolist(), [-2**30, 0]) self.assertEqual((((u+t*(2**31-1))+t)//2).tolist(), [-2**30, 0]) + @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, CUDARenderer), "nvrtc 11.5 reassociates (x+c)+y> 32).tolist(), [-1, 0]) + def test_mod_long_of_int_min(self): + a, b = Tensor([-2**31, 5], dtype=dtypes.int32).cast(dtypes.int64), Tensor([-1, -2], dtype=dtypes.int8).minimum(-1).cast(dtypes.int64) + self.assertEqual((a % b).tolist(), [0, -1]) @Context(EMULATED_DTYPES="long") def test_emulated_shl(self): self._test_shl() diff --git a/tinygrad_repo/test/runtime/test_linearizer.py b/tinygrad_repo/test/runtime/test_linearizer.py index 8b7b24fe07..e9b7a91059 100644 --- a/tinygrad_repo/test/runtime/test_linearizer.py +++ b/tinygrad_repo/test/runtime/test_linearizer.py @@ -109,7 +109,7 @@ class TestLinearizer(unittest.TestCase): uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src) ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE] # LOAD -> RANGE -> LOAD -> STORE - assert len([x for x in uops[:ranges[0]] if x.op is Ops.LOAD]) == 1 + assert len([x for x in uops[:ranges[0]] if x.op is Ops.LOAD and x.src[0].buf_uop.arg.name != "args"]) == 1 def test_range_outer_op_before_phi_nested_range(self): a = Tensor.randn(2, ).realize() diff --git a/tinygrad_repo/test/runtime/test_llm_tp.py b/tinygrad_repo/test/runtime/test_llm_tp.py index 610d33b18c..00206bddcc 100644 --- a/tinygrad_repo/test/runtime/test_llm_tp.py +++ b/tinygrad_repo/test/runtime/test_llm_tp.py @@ -105,14 +105,16 @@ class TestTensorParallel(unittest.TestCase): typ = Q.Q4_K data = random_data(shape, typ).reshape(-1) w = ggml_data_to_tensor(Tensor(data, device='CPU'), E*O*I, typ).reshape(shape).float().numpy() - sel = np.array([[[3, 1]]], dtype=np.int32) - for axis in (1, 2): + sel, T = np.array([[[3, 1], [0, 2]]], dtype=np.int32), Variable('T', 1, 2).bind(1) + for axis in (0, 1, 2): with self.subTest(axis=axis): layer = ExpertWeights(E, I, O) layer.weight = gguf_shard({'w': (Tensor(data, device='CPU'), shape, typ)}, DEVICES, {'w': axis})['w'].half() - x = np.random.default_rng(1).normal(size=(1, 1, 1 if axis == 1 else 2, I)).astype(np.float16) + x = np.random.default_rng(1).normal(size=(1, 2, 1 if axis == 1 else 2, I)).astype(np.float16) ref = np.stack([w[e] @ x[0, 0, j % x.shape[2]].astype(np.float32) for j, e in enumerate(sel[0, 0])]) - got = layer(Tensor(sel).shard(DEVICES), Tensor(x).shard(DEVICES, 3 if axis == 2 else None)).float().to(Device.DEFAULT).numpy() + got = layer(Tensor(sel).shard(DEVICES)[:, :T], Tensor(x).shard(DEVICES, 3 if axis == 2 else None)[:, :T]) + got = got.pad_to(got.max_shape).float().to(Device.DEFAULT).numpy() + if axis == 0: got = got.sum(0) self.assertEqual(layer.ggml_type, typ) np.testing.assert_allclose(got[0, 0], ref, rtol=2e-2, atol=2e-2*np.abs(ref).max()) @@ -121,7 +123,7 @@ class TestTensorParallel(unittest.TestCase): rng = np.random.default_rng(42) # Keep GQA, partial RoPE, gated attention and replicated SSM state; custom kernels have separate tests. config = TransformerConfig(num_blocks=2, dim=8, hidden_dim=8, n_heads=4, n_kv_heads=2, norm_eps=1e-5, vocab_size=16, head_dim=4, - v_head_dim=4, rope_theta=10000, rope_dim=2, max_context=4, qk_norm=4, attn_output_gate=True, ssm=SSMConfig(2, 2, 2, 4, 8), + v_head_dim=4, rope_theta=10000, rope_dim=2, max_context=4, qk_norm=4, attn_output_gate=True, ssm=SSMConfig(2, 2, 2, 4, 8, silu_gate=True), ssm_layers=(True, False)) with tempfile.TemporaryDirectory() as folder: writer = GGUFWriter(path:=pathlib.Path(folder)/'model.gguf', 'qwen35') diff --git a/tinygrad_repo/tinygrad/codegen/__init__.py b/tinygrad_repo/tinygrad/codegen/__init__.py index e78c95196d..f4535a4f7d 100644 --- a/tinygrad_repo/tinygrad/codegen/__init__.py +++ b/tinygrad_repo/tinygrad/codegen/__init__.py @@ -27,6 +27,7 @@ from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_c from tinygrad.codegen.late.regalloc import LinearScanRegallocContext, pm_regalloc_rewrite from tinygrad.codegen.late.coalesce import memory_coalescing, pm_simplify_add_image from tinygrad.helpers import all_same, argsort, partition, to_function_name +from tinygrad.device import TinyELF from tinygrad.uop.ops import _broadcast_shape, identity_element from tinygrad.schedule.rangeify import BufferizeOpts @@ -406,7 +407,7 @@ class LowerCallsCtx: def lower_call(ctx:LowerCallsCtx, call:UOp, b:UOp) -> UOp|None: if b.op is Ops.PROGRAM and b.src[-1].op is Ops.BINARY: return None bodies = ctx.names.setdefault(name:=to_function_name(call.arg.name), {}) - ki = KernelInfo(name + (f"_{i}" if (i:=bodies.setdefault(b, len(bodies))) else ""), optimized=True) + ki = KernelInfo(name + (f"_{i}" if (i:=bodies.setdefault(b, len(bodies))) else ""), optimized=True, kernel=False) if ctx.renderer.inline_prefix: b = graph_rewrite(b, pm_lower_calls, ctx=ctx, name="lower calls", walk=True) b = UOp(Ops.LINEAR, src=tuple(line_rewrite(list(b.src[1].src) if b.op is Ops.PROGRAM \ @@ -418,6 +419,13 @@ pm_lower_calls = PatternMatcher([ (UPat(Ops.CALL, src=(UPat((Ops.SINK, Ops.PROGRAM), name="b"),), allow_any_len=True, name="call"), lower_call), ]) +def pack_args(sink:UOp, params:tuple[UOp, ...]) -> UOp: + sig = TinyELF.iter_sig(tuple((None, p.addrspace, p.dtype, ()) for p in params)) + args = UOp.param(0, dtypes.uint8, len(params) * 8, name="args") + loads = [args.bitcast(dt).index(UOp.cconst(off // dt.itemsize, dtypes.int)).load() for off, dt in sig] + reads = {p: v if p.addrspace is AddrSpace.ALU else UOp(Ops.REFPARAM, src=(p.src[0], v), arg=p.arg) for p, v in zip(params, loads)} + return sink.substitute(reads, name="pack args") + def do_linearize(ctx:Renderer, prg:UOp, sink:UOp) -> UOp: if DEBUG >= 3 and sink.arg.applied_opts: print(f"{sink.arg.function_name:<25} opts: {sink.arg.applied_opts}") lst = line_rewrite(linearize(sink), pm_linearize_cleanups+pm_alloc_to_buf) @@ -484,11 +492,11 @@ def do_to_program(ast:UOp, renderer:Renderer) -> UOp: assert isinstance(ast.arg, KernelInfo), "requires KernelInfo on arg to to_program" full_sink = full_rewrite_to_sink(ast, renderer, optimize=not ast.arg.optimized) prog_info = ProgramInfo.from_sink(full_sink, renderer.target) + if (isinstance(renderer, ISARenderer) or renderer.packed_args) and full_sink.arg.estimates is None: + full_sink = full_sink.replace(arg=replace(full_sink.arg, estimates=Estimates.from_uops(tuple(linearize(full_sink)), ignore_indexing=True))) + if renderer.packed_args and full_sink.arg.kernel: full_sink = pack_args(full_sink, prog_info.params) # instruction selection if isinstance(renderer, ISARenderer): - # Instruction selection replaces LOAD/STORE/ALU with INS, so estimate while their meaning is still available. - if full_sink.arg.estimates is None: - full_sink = full_sink.replace(arg=replace(full_sink.arg, estimates=Estimates.from_uops(tuple(linearize(full_sink)), ignore_indexing=True))) full_sink = graph_rewrite(full_sink, renderer.pre_isel_matcher, ctx=itertools.count(-1, -1), name="pre instruction selection", bottom_up=True) full_sink = graph_rewrite(full_sink, renderer.isel_matcher, ctx=IselContext(full_sink), name="instruction selection", bottom_up=True) prg = UOp(Ops.PROGRAM, src=(full_sink,), arg=prog_info) diff --git a/tinygrad_repo/tinygrad/codegen/gpudims.py b/tinygrad_repo/tinygrad/codegen/gpudims.py index bc20ef7215..48f88612fb 100644 --- a/tinygrad_repo/tinygrad/codegen/gpudims.py +++ b/tinygrad_repo/tinygrad/codegen/gpudims.py @@ -92,9 +92,6 @@ pm_device_to_var = PatternMatcher([ # the DEVICE axis is not a program axis, it's bound per device at launch. lower it to the _device_num variable (like SPECIAL for devices) (UPat(Ops.RANGE, name="r"), lambda r: UOp.variable("_device_num", 0, r.vmax, dtype=r.dtype) if r.axis_type is AxisType.DEVICE else None), - # ENDs that closed a DEVICE range no longer close it - (UPat(Ops.END, name="e"), lambda e: e.replace(src=(e.src[0],)+tuple(s for s in e.src[1:] if s.op is not Ops.PARAM)) - if any(s.op is Ops.PARAM and s.arg.name == '_device_num' for s in e.src[1:]) else None), ]) # Run once: grouping creates new GLOBAL/LOCAL ranges, which must not be grouped again. diff --git a/tinygrad_repo/tinygrad/codegen/late/coalesce.py b/tinygrad_repo/tinygrad/codegen/late/coalesce.py index dc1dc23a51..65fb294914 100644 --- a/tinygrad_repo/tinygrad/codegen/late/coalesce.py +++ b/tinygrad_repo/tinygrad/codegen/late/coalesce.py @@ -107,9 +107,12 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp: if getenv("DMC"): return sink # collect - memory: defaultdict[tuple[Ops, UOp, UOp|str, UOp, object], dict[int, list[UOp]]] = defaultdict(dict) + memory: defaultdict[tuple[Ops, UOp, UOp|None, UOp, object, int|tuple[int, int]], dict[int, list[UOp]]] = defaultdict(dict) + # depth[u] is the lexicographically largest (GLOBAL stores, LOCAL stores) count on any chain u waits on, u included + depth: dict[UOp, tuple[int, int]] = {} for u in sink.toposort(): - # TODO: this should handle images too, it's just memory coalescing + depth[u] = waits = max([depth[s] for s in u.src_without_body], default=(0, 0)) + if u.op is Ops.STORE: depth[u] = (waits[0]+(u.src[0].addrspace is AddrSpace.GLOBAL), waits[1]+(u.src[0].addrspace is AddrSpace.LOCAL)) if u.op in {Ops.LOAD, Ops.STORE}: assert len(u.src) == (2 if u.op is Ops.STORE else 1), "memory coalescing does not support gated loads/stores" assert u.src[0].op is Ops.INDEX, f"memory coalescing should be on INDEX, not {u.src[0].op}" @@ -118,48 +121,47 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp: if buf.buf_uop.op is Ops.PARAM and buf.buf_uop.arg.volatile: continue # volatile accesses never merge idx, valid = idx_u.get_idx(), idx_u.get_valid() if idx.is_invalid: continue - root_src: UOp|str + root_src: UOp|None if idx.op is Ops.ADD and idx.src[1].op is Ops.CONST: root_src, arg = idx.src[0], idx.src[1].val elif idx.op is Ops.ADD and idx.src[0].op is Ops.CONST: root_src, arg = idx.src[1], idx.src[0].val - elif idx.op is Ops.CONST: root_src, arg = "CONST", idx.val + elif idx.op is Ops.CONST: root_src, arg = None, idx.val else: root_src, arg = idx, 0 # loads/stores only coalesce with others carrying the same arg (e.g. the nontemporal flag) - memory[(u.op, buf, root_src, valid, u.arg)].setdefault(arg, []).append(u) + # and the same depth (GLOBAL count only for GLOBAL accesses), so no merge closes a dependency cycle + memory[(u.op, buf, root_src, valid, u.arg, waits[0] if buf.addrspace is AddrSpace.GLOBAL else waits)].setdefault(arg, []).append(u) # build replacements replacements = {} - for (op,buf,_,_,ld_arg),offsets in memory.items(): + for (op,buf,_,_,ld_arg,_),offsets in memory.items(): # allowed lengths (copied in) lengths = [] must_divide = True if ctx.target.device == "DSP": lengths = [128,64,32,16,8,4] must_divide = False - elif is_image_shape(buf._shape): - lengths = [4] elif buf.dtype in (dtypes.float, dtypes.half, dtypes.int, dtypes.uint, *dtypes.fp8s) and ctx.supports_float4: # TODO: a better way to get this than ctx lengths = [8,4,2] if buf.dtype == dtypes.half and getenv("ALLOW_HALF8") else [4,2] lengths.append(1) # worst case, it's not folded # do the grouping - grouped_offsets = [[x for _,x in group] for _,group in itertools.groupby(enumerate(sorted(offsets.keys())), lambda x: x[1]-x[0])] - for full_grp in grouped_offsets: + for _,group in itertools.groupby(enumerate(sorted(offsets.keys())), lambda x: x[1]-x[0]): + full_grp = [offsets[x] for _,x in group] while len(full_grp): # the group starts at its first access's index, valid included - offset = offsets[full_grp[0]][0].src[0].src[1] + offset = full_grp[0][0].src[0].src[1] length = [l for l in lengths if l <= len(full_grp) and (not must_divide or offset.get_idx().divides(l) is not None)][0] grp, full_grp = full_grp[:length], full_grp[length:] - assert op is not Ops.STORE or all(len(offsets[g]) == 1 for g in grp), "attempting multiple stores" + assert op is not Ops.STORE or all(len(g) == 1 for g in grp), "attempting multiple stores" if length == 1: continue idx = UOp(Ops.SHRINK, src=(buf, offset, UOp.const(length))) if op == Ops.STORE: - datas = [offsets[g][0].src[1] for g in grp] + datas = [g[0].src[1] for g in grp] store = idx.store(UOp.stack(*datas)) - for g in grp: replacements[offsets[g][0]] = store + for g in grp: replacements[g[0]] = store else: ld = idx.load(arg=ld_arg) for i,g in enumerate(grp): - for oo in offsets[g]: replacements[oo] = ld.index(i) + for oo in g: replacements[oo] = ld.index(i) # apply return sink.substitute(replacements, name="memory coalescing") diff --git a/tinygrad_repo/tinygrad/codegen/simplify.py b/tinygrad_repo/tinygrad/codegen/simplify.py index 73cd405311..36de3a7f23 100644 --- a/tinygrad_repo/tinygrad/codegen/simplify.py +++ b/tinygrad_repo/tinygrad/codegen/simplify.py @@ -15,48 +15,50 @@ pm_flatten_range = PatternMatcher([ (UPat((Ops.REDUCE, Ops.END), name="r"), flatten_range), ]) -# index/range arithmetic uses DIV/MOD prior to late rewrite -def count_divmod(x:UOp) -> int: return sum(u.op in {Ops.DIV, Ops.MOD} for u in x.backward_slice) -def simplify_merge_adjacent(u:UOp) -> UOp|None: - reduce_ranges = [x.ranges for x in u.backward_slice_with_self if x.op is Ops.REDUCE] - # on END we only want to merge adjacent ranges, on REDUCE we want to try all combinations - for r0, r1 in (zip(u.ended_ranges, u.ended_ranges[1:]) if u.op is Ops.END else itertools.permutations(u.ended_ranges, 2)): - # check same type and neither bound depends on the other range - if r0.axis_type == r1.axis_type and r0 not in r1.ranges and r1 not in r0.ranges: - # check if the ranges to merge are in the same reduces - if all((r0 in rngs) == (r1 in rngs) for rngs in reduce_ranges): - s0, s1 = r0.src[0], r1.src[0] - # do the merge - new_range = r0.replace(src=(s0*s1,)) - nidx = graph_rewrite(u, _substitute+symbolic+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1}, - name=f"check_merge_{r0.axis_id}_{r1.axis_id}") +# TODO: something in GroupOp +pure_ops = GroupOp.Elementwise|GroupOp.Irreducible|{Ops.INDEX, Ops.LOAD, Ops.STACK} - # check if it simplifies. return after one merge so the next rewrite uses the new ranges, - # rather than continuing with stale pairs from the original ended_ranges. - if count_divmod(nidx) <= count_divmod(u): return nidx +def mergeable_pairs(x:UOp) -> list[tuple[UOp, UOp]]: + # an END or a REDUCE with effects runs its ranges as loops nested by axis id, so only neighbors merge + if x.op is Ops.REDUCE and all(y.op in pure_ops|{Ops.REDUCE} for y in x.backward_slice): return list(itertools.permutations(x.ended_ranges, 2)) + return list(itertools.pairwise(sorted(x.ended_ranges, key=lambda r: r.arg))) if x.op in {Ops.END, Ops.REDUCE} else [] + +def merge_ranges(u:UOp) -> UOp|None: + enders = {x:mergeable_pairs(x) for x in u.backward_slice if x.op in range_start} + impure = [x for x in u.backward_slice if x.op not in pure_ops] + for r0, r1 in dedup(p for pairs in enders.values() for p in pairs): + if r0.axis_type is not r1.axis_type or r0 in r1.ranges or r1 in r0.ranges: continue + # every op that ends r0 or r1 must merge them, and an op that isn't pure must run as often as before + if any((r0, r1) not in pairs for x, pairs in enders.items() if {r0, r1} & set(x.ended_ranges)): continue + if any((r0 in x.ranges) != (r1 in x.ranges) for x in impure): continue + s1 = r1.src[0] + n = r0.replace(src=((r0.src[0]*s1).simplify(),)) + sub = {r0:n//s1, r1:n%s1} + ret = graph_rewrite(u, _substitute+symbolic+pm_flatten_range, ctx=sub, name=f"check_merge_{r0.axis_id}_{r1.axis_id}") + if not any(v in ret.backward_slice for v in sub.values()): return ret return None -def mark_gated(ctx, idx): +def mark_gated(ctx, idx, store=None): if len(idx.src) > 1 and idx.src[1].op is Ops.WHERE: x, cond = idx.src[1].get_idx(), idx.src[1].get_valid() # get all ranges r with guards "r < c" for some const c guards = {r:c for v in cond.split_uop(Ops.AND) if v.op is Ops.CMPLT and (r:=v.src[0]).op is Ops.RANGE and (c:=v.src[1]).op is Ops.CONST} else: x, guards = idx, {} # ensure that we choose max(c_i) for all i where r < c_i - ctx |= {r:c for r,c in guards.items() if (r not in ctx or ctx[r].vmax < c.val)} - # but if a range is ever ungated, we cannot shrink it - ctx |= {r:r.src[0] for r in x.ranges if r not in guards} + ctx |= {r:c.minimum(r.src[0]) for r,c in guards.items() if (r not in ctx or ctx[r].vmax < c.val)} + # but if a range is ever ungated, we cannot shrink it, and a STORE runs every iteration of a range its index doesn't gate + ctx |= {r:r.src[0] for r in (x if store is None else store).ranges if r not in guards} def do_substitute(ctx:dict, x: UOp, sub_fxn:Callable[[UOp, UOp], UOp]) -> UOp|None: # Only the kernel root: rewriting a nested SINK would leave its enclosing END's binders unchanged. if x.arg is None: return None - ret = x.substitute({k:sub_fxn(k,v) for k,v in ctx.items() if v is not None}) - ctx.clear() + ret = x.substitute({k:sub_fxn(k,v) for k,v in ctx.items() if v is not None and k in x.backward_slice}) return None if ret is x else ret.simplify() pm_simplify_ranges = PatternMatcher([ - (UPat((Ops.END, Ops.REDUCE), name="u"), simplify_merge_adjacent), + (UPat(Ops.SINK, name="u"), merge_ranges), (UPat(Ops.INDEX, name="idx"), mark_gated), + (UPat(Ops.STORE, src=(UPat(Ops.INDEX, name="idx"), UPat()), name="store"), mark_gated), # reduce ranges can't be shrunk (UPat(Ops.REDUCE, name="red"), lambda ctx, red: ctx.update({r:r.src[0] for r in red.src[1:]})), (UPat(Ops.SINK, name="x"), lambda ctx, x: do_substitute(ctx, x, lambda r,c: r.replace(src=(c,)))), @@ -94,7 +96,7 @@ pm_reduce_unparented = PatternMatcher([ pm_reduce_collapse = pm_reduce_unparented + PatternMatcher([ # lift x+y out of reduce on lt - (UPat(Ops.ADD, dtypes.sints+dtypes.floats+dtypes.weaks, [UPat.var("x"), UPat.var("y")]) < UPat.var("c"), + (UPat(Ops.ADD, dtypes.sints+(dtypes.weakint,), [UPat.var("x"), UPat.var("y")]) < UPat.var("c"), lambda x,y,c: (x < (c-y)) if no_range(y) and no_range(c) else None), # lift x*y out of reduce (UPat(Ops.MUL, dtypes.sints+(dtypes.weakint,), [UPat.var("x"), UPat.var("y")]) < UPat.var("c"), @@ -118,10 +120,14 @@ pm_reduce_collapse = pm_reduce_unparented + PatternMatcher([ pm_reduce_load_collapse = pm_reduce_collapse + PatternMatcher([ # lift x+y out of reduce on eq - ((UPat.var("x")+UPat.var("y")).or_casted().eq(UPat.var("c")), lambda x,y,c: x.eq(c.cast(y.dtype)-y) if no_range(y) and no_range(c) else None), + ((UPat.var("x", dtypes.wints)+UPat.var("y")).or_casted("a").eq(UPat.var("c", dtypes.wints)), + lambda x,y,a,c: x.eq(c.cast(y.dtype)-y) if no_range(y) and no_range(c) and not no_range(x) and + (a.op is not Ops.CAST or not a.src[0].overflows(a.dtype)) and not c.overflows(y.dtype) else None), # reduce on gated load becomes can substitute the range and remove the reduce - (UPat.var("idx").eq(UPat(Ops.RANGE, name="r").or_casted()).where(UPat.var("expr"), 0).reduce(UPat.var("r"), arg=Ops.ADD), - lambda r,idx,expr: (v:=(idx.cast(r.dtype) >= 0) & (idx.cast(r.dtype) < r.src[0])).where(expr.substitute({r:idx.cast(r.dtype).valid(v)}),0)), + (UPat.var("idx", dtypes.wints).eq(UPat(Ops.RANGE, name="r").or_casted("rc")).where(UPat.var("expr"), 0) + .reduce(UPat.var("r"), arg=Ops.ADD), + lambda r,rc,idx,expr: None if r.overflows(rc.dtype) else + (v:=(idx.cast(r.dtype) >= 0) & (idx.cast(r.dtype) < r.src[0])).where(expr.substitute({r:idx.cast(r.dtype).valid(v)}),0)), ]) def reduce_collapse(red:UOp, u:UOp, pm:PatternMatcher=pm_reduce_collapse) -> UOp|None: diff --git a/tinygrad_repo/tinygrad/device.py b/tinygrad_repo/tinygrad/device.py index a4a4b137f6..f30c81f8d8 100644 --- a/tinygrad_repo/tinygrad/device.py +++ b/tinygrad_repo/tinygrad/device.py @@ -69,7 +69,7 @@ def canonicalize_device(device:str|tuple|list|None) -> str|tuple[str, ...]: return canonical[0] if len(canonical:=tuple(Device.canonicalize(d) for d in device)) == 1 else canonical def is_disk_device(device:str|tuple[str, ...]) -> bool: - return any(d.split(":", 1)[0].upper() == "DISK" for d in to_tuple(device)) + return any(d.split(":", 1)[0].upper() in ("DISK", "NVME") for d in to_tuple(device)) # **************** Profile **************** @@ -182,7 +182,7 @@ class Buffer: else: storage = self.allocator.alloc(self.nbytes, self.options) storage = replace(storage, host=storage.host.view(self.offset, self.nbytes, fmt='B') if storage.host is not None else None) if self._base is None: - if not self.device.startswith("DISK") and self.options.external_ptr is None: + if not is_disk_device(self.device) and self.options.external_ptr is None: GlobalCounters.mem_used += self.nbytes GlobalCounters.mem_used_per_device[self.device] += self.nbytes if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "alloc", self.trace_num, {"nbytes":self.nbytes})) @@ -193,7 +193,7 @@ class Buffer: assert self._storage is not None, "buffer must be allocated to deallocate" if DEBUG is not None and DEBUG >= 7: print(f"buffer: deallocate {self.nbytes} bytes on {self.device}") if self._base is None: - if GlobalCounters is not None and not self.device.startswith("DISK") and self.options.external_ptr is None: + if GlobalCounters is not None and not is_disk_device(self.device) and self.options.external_ptr is None: GlobalCounters.mem_used -= self.nbytes GlobalCounters.mem_used_per_device[self.device] -= self.nbytes if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "free", self.trace_num)) diff --git a/tinygrad_repo/tinygrad/dtype.py b/tinygrad_repo/tinygrad/dtype.py index e404abc694..625b9e65f4 100644 --- a/tinygrad_repo/tinygrad/dtype.py +++ b/tinygrad_repo/tinygrad/dtype.py @@ -78,6 +78,8 @@ class DType(metaclass=DTypeMetaClass): return float("inf") if dtypes.is_float(self) else True def const(self, val: ConstType): if isinstance(val, InvalidType): return val + if self in dtypes.ints and isinstance(val, float) and not math.isfinite(val): + val = 0 if math.isnan(val) else self.max if val > 0 else self.min if self in dtypes.ints: return truncate[self](int(val)) # NOTE: float('nan') != float('nan'), so we canonicalize here if isinstance(val, float) and math.isnan(val): val = math.nan @@ -91,7 +93,7 @@ class DTypes: def is_float(x: DType) -> bool: return x in (dtypes.floats + (dtypes.weakfloat,)) @staticmethod # static methods on top, or bool in the type info will refer to dtypes.bool @functools.cache - def is_int(x: DType) -> bool: return x in (dtypes.ints + (dtypes.weakint,)) + def is_int(x: DType) -> bool: return x in dtypes.wints @staticmethod @functools.cache def is_unsigned(x: DType) -> bool: return x in dtypes.uints @@ -158,6 +160,7 @@ class DTypes: uints = (uint8, uint16, uint32, uint64) sints = (int8, int16, int32, int64) ints = uints + sints + wints = ints + (weakint,) weaks = (weakint, weakfloat) all = floats + ints + (bool,) # noqa: A003 @@ -313,12 +316,9 @@ def _from_np_dtype(npdtype:'np.dtype') -> DType: # type: ignore [name-defined] # def _to_torch_dtype(dtype:DType) -> 'torch.dtype'|None: # type: ignore [name-defined] # noqa: F821 import numpy as np, torch dtype = strong_dtype(dtype) - if dtype == dtypes.uint64: return torch.uint64 if dtype == dtypes.bfloat16: return torch.bfloat16 if dtype in dtypes.fp8s: return torch.uint8 - # NOTE: torch doesn't expose this mapping with a stable API - try: return torch.from_numpy(np.array([], dtype=_to_np_dtype(dtype))).dtype - except TypeError: return None + return getattr(torch, np.dtype(dtype.fmt).name, None) if dtype.fmt is not None else None @functools.cache def _from_torch_dtype(torchdtype:'torch.dtype') -> DType: # type: ignore [name-defined] # noqa: F821 return {v:k for k in DTYPES_DICT.values() if (v:=_to_torch_dtype(k)) is not None}[torchdtype] diff --git a/tinygrad_repo/tinygrad/llm/chat.html b/tinygrad_repo/tinygrad/llm/chat.html index 1c91cfb69c..8b53640a9d 100644 --- a/tinygrad_repo/tinygrad/llm/chat.html +++ b/tinygrad_repo/tinygrad/llm/chat.html @@ -1,38 +1,63 @@ -tinygrad chat
+
- +
+