diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index 6027ed90ea..739f7242c0 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -94,6 +94,7 @@ jobs: shell: bash -e -o pipefail {0} env: DEV: ${{ matrix.dev }} + HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }} if: github.repository_owner == 'tinygrad' steps: - name: Checkout Code @@ -104,6 +105,9 @@ jobs: ./extra/amdpci/setup_python_cap.sh ./extra/hcq/hcq_smi.py amd rmmod ./extra/hcq/hcq_smi.py amd kill_pids + - name: Setup (NV) + if: ${{ matrix.dev == 'NV' }} + run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; } - name: Symlink models and datasets run: | mkdir -p weights @@ -145,6 +149,7 @@ jobs: shell: bash -e -o pipefail {0} env: DEV: ${{ matrix.dev }} + HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }} if: github.repository_owner == 'tinygrad' steps: - name: Checkout Code @@ -155,6 +160,9 @@ jobs: ./extra/amdpci/setup_python_cap.sh ./extra/hcq/hcq_smi.py amd rmmod ./extra/hcq/hcq_smi.py amd kill_pids + - name: Setup (NV) + if: ${{ matrix.dev == 'NV' }} + run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; } - name: setup staging db if: github.ref == 'refs/heads/update_benchmark_staging' run: | @@ -194,6 +202,7 @@ jobs: shell: bash -e -o pipefail {0} env: DEV: ${{ matrix.dev }} + HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }} if: github.repository_owner == 'tinygrad' steps: - name: Checkout Code @@ -204,6 +213,9 @@ jobs: ./extra/amdpci/setup_python_cap.sh ./extra/hcq/hcq_smi.py amd rmmod ./extra/hcq/hcq_smi.py amd kill_pids + - name: Setup (NV) + if: ${{ matrix.dev == 'NV' }} + run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; } - name: Symlink models and datasets run: | mkdir -p extra/datasets @@ -240,6 +252,7 @@ jobs: shell: bash -e -o pipefail {0} env: DEV: ${{ matrix.dev }} + HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }} if: github.repository_owner == 'tinygrad' steps: - name: Checkout Code @@ -250,6 +263,9 @@ jobs: ./extra/amdpci/setup_python_cap.sh ./extra/hcq/hcq_smi.py amd rmmod ./extra/hcq/hcq_smi.py amd kill_pids + - name: Setup (NV) + if: ${{ matrix.dev == 'NV' }} + run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; } - name: setup staging db if: github.ref == 'refs/heads/update_benchmark_staging' run: | @@ -292,6 +308,9 @@ jobs: ./extra/amdpci/setup_python_cap.sh ./extra/hcq/hcq_smi.py amd rmmod ./extra/hcq/hcq_smi.py amd kill_pids + - name: Setup (NV) + if: ${{ matrix.dev == 'NV' }} + run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; } - name: setup staging db if: github.ref == 'refs/heads/update_benchmark_staging' run: | diff --git a/.github/workflows/platform.yml b/.github/workflows/platform.yml new file mode 100644 index 0000000000..bb27f446d1 --- /dev/null +++ b/.github/workflows/platform.yml @@ -0,0 +1,213 @@ +name: Platform Tests +env: + # increment this when downloads substantially change to avoid the internet + CACHE_VERSION: '19' + 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 }} + CHECK_OOB: 1 + +on: + push: + branches: + - master + pull_request: + workflow_dispatch: + +concurrency: + group: platform-${{ github.event_name }}-${{ github.event_name == 'pull_request' && github.event.pull_request.number || github.run_id }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + +jobs: + +# ****** OSX Tests ****** + + unittestmacos: + name: MacOS (unit) + runs-on: macos-26 + timeout-minutes: 20 + steps: + - name: Checkout Code + uses: actions/checkout@v6 + - name: Setup Environment + uses: ./.github/actions/setup-tinygrad + with: + key: unittest-macos + deps: testing_unit + - name: Run unit tests + run: DEV=METAL python -m pytest -n=auto test/unit/ --durations=20 + - name: Test tensor core ops (fake) + run: DEV=METAL DEBUG=3 TC=2 python test/backend/test_ops.py TestOps.test_gemm + - name: Test tensor core ops (real) + run: DEV=METAL DEBUG=3 python test/backend/test_ops.py TestOps.test_big_gemm + - name: Test Beam Search + run: DEV=METAL IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py + - name: Test Device Specific + run: DEV=METAL python3 -m pytest test/device/test_metal.py + #- name: Fuzz Test linearizer + # run: DEV=METAL DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py + - name: Run process replay tests + uses: ./.github/actions/process-replay + + unittestmacosmock: + name: MacOS (unit, mock) + runs-on: macos-26 + timeout-minutes: 20 + steps: + - name: Checkout Code + uses: actions/checkout@v6 + - name: Setup Environment + uses: ./.github/actions/setup-tinygrad + with: + key: unittest-macos-mock + deps: testing_unit + amd: 'true' + ocelot: 'true' + - name: Run NULL backend tests + run: SPEC=2 DEV=NULL python -m pytest -n=auto test/null/ --durations=20 + - name: Run pytest (amd) + env: + DEV: MOCKKFD+AMD + FORWARD_ONLY: 1 + run: | + python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20 + - name: Run pytest (ptx) + env: + DEV: "MOCK+NV:PTX" + FORWARD_ONLY: 1 + # TODO: failing due to library loading error + CAPTURE_PROCESS_REPLAY: 0 + run: | + python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py \ + test/testextra/test_hevc.py::TestHevc::test_hevc_decode_compile --durations=20 + - name: Run process replay tests + uses: ./.github/actions/process-replay + + testmetal: + strategy: + fail-fast: false + matrix: + group: [1, 2] + name: MacOS (DEV=METAL) (${{ matrix.group }}) + runs-on: macos-26 + timeout-minutes: 20 + env: + DEV: METAL + steps: + - name: Checkout Code + uses: actions/checkout@v6 + - name: Setup Environment + uses: ./.github/actions/setup-tinygrad + with: + key: macos-metal + deps: testing_unit + - name: Check Device.DEFAULT and print some source + run: | + python -c "from tinygrad import Device; assert Device.DEFAULT == 'METAL'" + DEBUG=4 python test/test_tiny.py TestTiny.test_plus + - name: Run backend tests + run: python -m pytest -n=auto test/backend --durations=20 --splits 2 --group ${{ matrix.group }} + - name: Run process replay tests + uses: ./.github/actions/process-replay + + testmacos: + strategy: + fail-fast: false + matrix: + dev: + - 'CPU:CLANG' + - 'CPU:LLVM' + - 'CPU:LVP' + - 'WEBGPU' + + name: MacOS (DEV=${{ matrix.dev }}) + runs-on: macos-26 + timeout-minutes: 20 + steps: + - name: Checkout Code + uses: actions/checkout@v6 + - name: Setup Environment + uses: ./.github/actions/setup-tinygrad + with: + key: macos-${{ matrix.dev }} + deps: "testing_unit${{ contains(matrix.dev, 'LVP') && ' mesa' || '' }}" + llvm: ${{ contains(matrix.dev, 'LLVM') || contains(matrix.dev, 'LVP') }} + webgpu: ${{ matrix.dev == 'WEBGPU' }} + - name: Set env + run: printf "DEV=${{ matrix.dev }}${{ matrix.dev == 'CPU:CLANG' && '\nCPU_COUNT=2' || '' }}" >> $GITHUB_ENV + - name: Check Device.DEFAULT and print some source + run: | + python -c "from tinygrad import Device; from tinygrad.helpers import Target; assert Device.DEFAULT == Target.parse('${{ matrix.dev }}').device" + DEBUG=4 python test/test_tiny.py TestTiny.test_plus + - name: Run test_tiny + run: python -m pytest -n=auto test/test_tiny.py --durations=20 + - name: Run process replay tests + uses: ./.github/actions/process-replay + +# ****** Windows Tests ****** + + testwindows: + strategy: + fail-fast: false + matrix: + dev: + - 'CPU:CLANG' + - 'CPU:LLVM' + - 'CPU:X86' + - 'WEBGPU' + + name: Windows (DEV=${{ matrix.dev }}) + runs-on: windows-2025 + timeout-minutes: 15 + steps: + - name: Checkout Code + uses: actions/checkout@v6 + - name: Setup Environment + uses: ./.github/actions/setup-tinygrad + with: + key: windows-${{ matrix.dev }}-minimal + deps: testing_unit + pydeps: ${{ matrix.dev == 'WEBGPU' && 'dawn-python' || '' }} + - name: Set env + shell: bash + run: printf "DEV=${{ matrix.dev }}${{ matrix.dev == 'CPU:CLANG' && '\nCPU_COUNT=2' || '' }}" >> $GITHUB_ENV + - name: Check Device.DEFAULT and print some source + shell: bash + run: | + python -c "from tinygrad import Device; from tinygrad.helpers import Target; assert Device.DEFAULT == Target.parse('${{ matrix.dev }}').device" + DEBUG=4 python test/test_tiny.py TestTiny.test_plus + - name: Run test_tiny + shell: bash + run: python -m pytest -n=auto test/test_tiny.py --durations=20 + + + qcomclcompiletests: + name: Compile-only (QCOM CL) + runs-on: ubuntu-24.04-arm + timeout-minutes: 15 + steps: + - name: Checkout Code + uses: actions/checkout@v6 + - name: Setup Environment + uses: ./.github/actions/setup-tinygrad + with: + key: compile-qcomcl + deps: testing_unit + tinydreno: 'true' + - name: Set env + shell: bash + run: printf "DEV=NULL:QCOMCL:a630\nNULL_ALLOW_COPYOUT=1" >> $GITHUB_ENV + - name: Run test_ops + shell: bash + run: | + python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'" + DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add + python -m pytest -n=auto test/backend/test_ops.py --durations=20 + - name: Run test_ops (IMAGE) + shell: bash + env: + IMAGE: 1 + DEV: "NULL:QCOMCL:a630,IMAGE_PITCH_ALIGNMENT=64" + run: | + DEBUG=4 python test/backend/test_ops.py TestOps.test_gemm | grep read_imagef + python -m pytest -n=auto test/backend/test_ops.py --durations=20 diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 4d272ca048..7bd3063e41 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -294,7 +294,7 @@ jobs: llvm: 'true' - name: Test openpilot model kernel count and gate usage run: | - ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1361 ALLOWED_GATED_READ_IMAGE=54 FLOAT16=1 DEV="CL::IMAGE_PITCH_ALIGNMENT=64" IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916 + ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1361 ALLOWED_GATED_READ_IMAGE=38 FLOAT16=1 DEV="CL::IMAGE_PITCH_ALIGNMENT=64" IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916 # IMAGE_PITCH_ALIGNMENT=64 matches adreno 630 - name: Test openpilot CL compile fp32 (test correctness) run: | @@ -527,6 +527,8 @@ jobs: TestMultiTensor.test_backward_sum TestMultiTensor.test_matmul_shard_0_0 - name: Run HCQ2 JIT tests run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/unit/test_jit.py + - name: Run HCQ2 unit tests + run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python -m pytest test/device/test_hcq2.py testmockam: name: Linux (am) @@ -585,8 +587,11 @@ jobs: run: | python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['AMD'], Device.DEFAULT" DEBUG=5 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus + - name: Run MXFP4 Llama training on NULL backend + if: ${{ matrix.backend == 'amd' && matrix.arch == 'gfx950' }} + run: PYTHONPATH=. DEV=NULL:HIP:gfx950 MXFP4=1 LLAMA_LAYERS=2 BENCHMARK=3 NULL_ALLOW_COPYOUT=1 NO_HIPCC=1 ROCM_PATH=/opt/rocm JITBEAM=0 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/profile.sh - name: Run pytest (amd) - run: python -m pytest -n=auto test/backend/test_ops.py test/backend/test_dtype.py test/backend/test_dtype_alu.py test/backend/test_linearizer.py test/backend/test_randomness.py test/backend/test_jit.py test/backend/test_graph.py test/backend/test_multitensor.py test/device/test_hcq.py test/external/external_test_am.py test/backend/test_asm_gemm.py::TestAsmGEMM --durations=20 + run: python -m pytest -n=auto test/backend/test_ops.py test/backend/test_dtype.py test/backend/test_dtype_alu.py test/backend/test_linearizer.py test/backend/test_randomness.py test/backend/test_jit.py test/backend/test_graph.py test/backend/test_multitensor.py test/device/test_hcq.py test/external/external_test_am.py test/backend/test_asm_gemm.py::TestAsmGEMM test/opt/test_tensor_cores.py --durations=20 - name: Run disk copy tests run: python -m pytest test/unit/test_disk_tensor.py -k test_copy_from_disk - name: Run TRANSCENDENTAL math @@ -629,165 +634,6 @@ jobs: - name: Run process replay tests uses: ./.github/actions/process-replay -# ****** OSX Tests ****** - - unittestmacos: - name: MacOS (unit) - runs-on: macos-26 - timeout-minutes: 20 - steps: - - name: Checkout Code - uses: actions/checkout@v6 - - name: Setup Environment - uses: ./.github/actions/setup-tinygrad - with: - key: unittest-macos - deps: testing_unit - - name: Run unit tests - run: DEV=METAL python -m pytest -n=auto test/unit/ --durations=20 - - name: Test tensor core ops (fake) - run: DEV=METAL DEBUG=3 TC=2 python test/backend/test_ops.py TestOps.test_gemm - - name: Test tensor core ops (real) - run: DEV=METAL DEBUG=3 python test/backend/test_ops.py TestOps.test_big_gemm - - name: Test Beam Search - run: DEV=METAL IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py - - name: Test Device Specific - run: DEV=METAL python3 -m pytest test/device/test_metal.py - #- name: Fuzz Test linearizer - # run: DEV=METAL DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py - - name: Run process replay tests - uses: ./.github/actions/process-replay - - unittestmacosmock: - name: MacOS (unit, mock) - runs-on: macos-26 - timeout-minutes: 20 - steps: - - name: Checkout Code - uses: actions/checkout@v6 - - name: Setup Environment - uses: ./.github/actions/setup-tinygrad - with: - key: unittest-macos-mock - deps: testing_unit - amd: 'true' - ocelot: 'true' - - name: Run NULL backend tests - run: SPEC=2 DEV=NULL python -m pytest -n=auto test/null/ --durations=20 - - name: Run pytest (amd) - env: - DEV: MOCKKFD+AMD - FORWARD_ONLY: 1 - run: | - python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20 - - name: Run pytest (ptx) - env: - DEV: "MOCK+NV:PTX" - FORWARD_ONLY: 1 - # TODO: failing due to library loading error - CAPTURE_PROCESS_REPLAY: 0 - run: | - python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20 - - name: Run process replay tests - uses: ./.github/actions/process-replay - - testmetal: - strategy: - fail-fast: false - matrix: - group: [1, 2] - name: MacOS (DEV=METAL) (${{ matrix.group }}) - runs-on: macos-26 - timeout-minutes: 20 - env: - DEV: METAL - steps: - - name: Checkout Code - uses: actions/checkout@v6 - - name: Setup Environment - uses: ./.github/actions/setup-tinygrad - with: - key: macos-metal - deps: testing_unit - - name: Check Device.DEFAULT and print some source - run: | - python -c "from tinygrad import Device; assert Device.DEFAULT == 'METAL'" - DEBUG=4 python test/test_tiny.py TestTiny.test_plus - - name: Run backend tests - run: python -m pytest -n=auto test/backend --durations=20 --splits 2 --group ${{ matrix.group }} - - name: Run process replay tests - uses: ./.github/actions/process-replay - - testmacos: - strategy: - fail-fast: false - matrix: - dev: - - 'CPU:CLANG' - - 'CPU:LLVM' - - 'CPU:LVP' - - 'WEBGPU' - - name: MacOS (DEV=${{ matrix.dev }}) - runs-on: macos-26 - timeout-minutes: 20 - steps: - - name: Checkout Code - uses: actions/checkout@v6 - - name: Setup Environment - uses: ./.github/actions/setup-tinygrad - with: - key: macos-${{ matrix.dev }} - deps: "testing_unit${{ contains(matrix.dev, 'LVP') && ' mesa' || '' }}" - llvm: ${{ contains(matrix.dev, 'LLVM') || contains(matrix.dev, 'LVP') }} - webgpu: ${{ matrix.dev == 'WEBGPU' }} - - name: Set env - run: printf "DEV=${{ matrix.dev }}${{ matrix.dev == 'CPU:CLANG' && '\nCPU_COUNT=2' || '' }}" >> $GITHUB_ENV - - name: Check Device.DEFAULT and print some source - run: | - python -c "from tinygrad import Device; from tinygrad.helpers import Target; assert Device.DEFAULT == Target.parse('${{ matrix.dev }}').device" - DEBUG=4 python test/test_tiny.py TestTiny.test_plus - - name: Run test_tiny - run: python -m pytest -n=auto test/test_tiny.py --durations=20 - - name: Run process replay tests - uses: ./.github/actions/process-replay - -# ****** Windows Tests ****** - - testwindows: - strategy: - fail-fast: false - matrix: - dev: - - 'CPU:CLANG' - - 'CPU:LLVM' - - 'CPU:X86' - - 'WEBGPU' - - name: Windows (DEV=${{ matrix.dev }}) - runs-on: windows-2025 - timeout-minutes: 15 - steps: - - name: Checkout Code - uses: actions/checkout@v6 - - name: Setup Environment - uses: ./.github/actions/setup-tinygrad - with: - key: windows-${{ matrix.dev }}-minimal - deps: testing_unit - pydeps: ${{ matrix.dev == 'WEBGPU' && 'dawn-python' || '' }} - - name: Set env - shell: bash - run: printf "DEV=${{ matrix.dev }}${{ matrix.dev == 'CPU:CLANG' && '\nCPU_COUNT=2' || '' }}" >> $GITHUB_ENV - - name: Check Device.DEFAULT and print some source - shell: bash - run: | - python -c "from tinygrad import Device; from tinygrad.helpers import Target; assert Device.DEFAULT == Target.parse('${{ matrix.dev }}').device" - DEBUG=4 python test/test_tiny.py TestTiny.test_plus - - name: Run test_tiny - shell: bash - run: python -m pytest -n=auto test/test_tiny.py --durations=20 - # ****** Compile-only Tests ****** compiletests: @@ -824,33 +670,3 @@ jobs: run: | DEBUG=4 python3 test/backend/test_ops.py TestOps.test_gemm | grep image_load python -m pytest -n=auto test/backend/test_ops.py --durations=20 - qcomclcompiletests: - name: Compile-only (QCOM CL) - runs-on: ubuntu-24.04-arm - timeout-minutes: 15 - steps: - - name: Checkout Code - uses: actions/checkout@v6 - - name: Setup Environment - uses: ./.github/actions/setup-tinygrad - with: - key: compile-qcomcl - deps: testing_unit - tinydreno: 'true' - - name: Set env - shell: bash - run: printf "DEV=NULL:QCOMCL:a630\nNULL_ALLOW_COPYOUT=1" >> $GITHUB_ENV - - name: Run test_ops - shell: bash - run: | - python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'" - DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add - python -m pytest -n=auto test/backend/test_ops.py --durations=20 - - name: Run test_ops (IMAGE) - shell: bash - env: - IMAGE: 1 - DEV: "NULL:QCOMCL:a630,IMAGE_PITCH_ALIGNMENT=64" - run: | - DEBUG=4 python test/backend/test_ops.py TestOps.test_gemm | grep read_imagef - python -m pytest -n=auto test/backend/test_ops.py --durations=20 diff --git a/examples/mlperf/model_train.py b/examples/mlperf/model_train.py index e9d8086aab..47c4977e39 100644 --- a/examples/mlperf/model_train.py +++ b/examples/mlperf/model_train.py @@ -1458,7 +1458,8 @@ def train_llama3(): # realize everything here if optim.master_params: Tensor.realize(*optim.master_params) - Tensor.realize(*optim.params, *fp8_inv_scales, *fp8_amax, *fp8_next_amax, *fp8_grad_amax, *fp8_next_grad_amax) + loss_acc = Tensor.zeros(1, dtype=dtypes.float32, device=device) + Tensor.realize(loss_acc, *optim.params, *fp8_inv_scales, *fp8_amax, *fp8_next_amax, *fp8_grad_amax, *fp8_next_grad_amax) @TinyJit def minibatch(tokens:Tensor): @@ -1476,8 +1477,8 @@ def train_llama3(): for g, new_g in zip(grads, loss.gradient(*optim.params)): apply_grad(g, new_g.uop) - loss_cpu = loss.flatten().float().to("CPU") - return loss_cpu.realize(*grads, *fp8_amax, *fp8_next_amax, *fp8_grad_amax, *fp8_next_grad_amax) + loss_acc.assign(loss_acc + loss.flatten().float()) + return loss_acc.realize(*grads, *fp8_amax, *fp8_next_amax, *fp8_grad_amax, *fp8_next_grad_amax) @TinyJit def optim_step(): @@ -1490,9 +1491,10 @@ def train_llama3(): lr_cpu = optim.lr.float().to("CPU") grad_norm_cpu = grad_norm.float().to("CPU") - Tensor.realize(lr_cpu, grad_norm_cpu, *grads, *fp8_inv_scales, *fp8_amax, *fp8_grad_amax) + loss_cpu = loss_acc.to("CPU") + Tensor.realize(lr_cpu, grad_norm_cpu, loss_cpu, loss_acc.assign(0), *grads, *fp8_inv_scales, *fp8_amax, *fp8_grad_amax) - return lr_cpu, grad_norm_cpu + return lr_cpu, grad_norm_cpu, loss_cpu @TinyJit @Context(TRAINING=0) @@ -1547,8 +1549,8 @@ def train_llama3(): st = time.perf_counter() stopped = False - losses, data_time, dev_time = [], 0, 0 - for _ in range(grad_acc if i >= 2 else 1): + data_time, dev_time = 0, 0 + for _ in range(accum_steps:=grad_acc if i >= 2 else 1): ist = time.perf_counter() try: tokens = next(train_iter) except StopIteration: @@ -1556,16 +1558,15 @@ def train_llama3(): break mst = time.perf_counter() data_time += mst - ist - losses.append(minibatch(tokens).item()) + minibatch(tokens) dev_time += time.perf_counter() - mst if stopped: break gt = time.perf_counter() ret = optim_step() - lr, grad_norm = ret[0].item(), ret[1].item() + lr, grad_norm, loss = ret[0].item(), ret[1].item(), ret[2].item() / accum_steps et = time.perf_counter() - loss = sum(losses) / len(losses) optim_time = et - gt dev_time += optim_time step_time = et - st diff --git a/examples/mlperf/models/flat_llama.py b/examples/mlperf/models/flat_llama.py index d237823687..4e3a2244ce 100644 --- a/examples/mlperf/models/flat_llama.py +++ b/examples/mlperf/models/flat_llama.py @@ -114,6 +114,11 @@ def silu_w13_quantize_matmul(x_w13:Tensor, w2:Tensor, s_2:Tensor, amax_x2:Tensor|None, next_amax_x2:Tensor|None, grad_amax_xw13:Tensor|None, next_grad_amax_xw13:Tensor|None, grad_amax_xout:Tensor|None, next_grad_amax_xout:Tensor|None): + if FUSED_SILU_W13 and MXFP4: + from extra.llama_kernels.swiglu import swiglu + out, *ret = matmul(swiglu(x_w13), w2, amax_x=amax_x2, w_inv_scale=s_2, grad_amax_state=grad_amax_xout, + next_grad_amax_state=next_grad_amax_xout, next_amax_x=next_amax_x2) + return out, ret if FUSED_SILU_W13 and not MXFP4: from extra.llama_kernels.cast_amax import fused_quantize_fp8_w13 x2_fp8 = fused_quantize_fp8_w13(x_w13, amax_x2, FP8_DTYPE, grad_amax_state=grad_amax_xw13, diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh index dc07e6a04b..db5daa5f2f 100755 --- a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh @@ -1,8 +1,8 @@ #!/usr/bin/env bash export PYTHONPATH="." -export PATH="/opt/rocm-7.1.1/bin:$PATH" -export ROCM_PATH="/opt/rocm-7.1.1" +export ROCM_PATH=${ROCM_PATH:-/opt/rocm-7.1.1} +export PATH="$ROCM_PATH/bin:$PATH" export DEV=${DEV:-AMD} export CHECK_OOB=0 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000 @@ -16,7 +16,7 @@ export USE_ATOMICS=${USE_ATOMICS:-1} export ASM_GEMM=${ASM_GEMM:-1} export WQKV=${WQKV:-1} export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1} -export FP8=${FP8:-1} +export MXFP4=${MXFP4:-1} export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1} export FAST_CE=${FAST_CE:-1} export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1} diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh index 30eb5c5116..0bfe3ae827 100755 --- a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh @@ -16,7 +16,7 @@ export USE_ATOMICS=${USE_ATOMICS:-1} export ASM_GEMM=${ASM_GEMM:-1} export WQKV=${WQKV:-1} export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1} -export FP8=${FP8:-1} +export MXFP4=${MXFP4:-1} export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1} export FAST_CE=${FAST_CE:-1} export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1} diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/profile.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/profile.sh index a2f4f207d7..125bf5f106 100755 --- a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/profile.sh +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/profile.sh @@ -1,4 +1,6 @@ #!/bin/bash +set -e + export BENCHMARK=${BENCHMARK:-5} export EVAL_BS=0 VIZ=${VIZ:--1} FULL_LAYERS=1 DEBUG=${DEBUG:--0} examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh diff --git a/extra/benchmark_llm.py b/extra/benchmark_llm.py new file mode 100644 index 0000000000..aa320c5b2b --- /dev/null +++ b/extra/benchmark_llm.py @@ -0,0 +1,31 @@ +import argparse, time +from tinygrad.llm.model import Transformer + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--model", required=True, help="path to gguf model") + parser.add_argument("--max-context", type=int, default=8192, help="max context length (default: %(default)s)") + parser.add_argument("--prompt-tokens", type=int, default=1024, help="number of prompt tokens (default: %(default)s)") + parser.add_argument("--decode-tokens", type=int, default=16, help="number of tokens to decode (default: %(default)s)") + parser.add_argument("--chunk-size", type=int, default=32, help="chunk size for prefill (default: %(default)s)") + args = parser.parse_args() + + st = time.perf_counter() + model, _ = Transformer.from_gguf(args.model, args.max_context) + print(f"load {time.perf_counter()-st:.3f}s", flush=True) + + st = time.perf_counter() + model.warmup() + print(f"warm {time.perf_counter()-st:.3f}s", flush=True) + + prompt = [257] + [1000+i%1000 for i in range(args.prompt_tokens-1)] + gen = model.generate(prompt, chunk_size=args.chunk_size) + st = time.perf_counter() + # first token is time-to-first-token; counted as part of prefill + output = [next(gen)] + pt = time.perf_counter() + print(f"prefill {args.prompt_tokens/(pt-st):.3f} tok/s", flush=True) + + for _ in range(args.decode_tokens): output.append(next(gen)) + et = time.perf_counter() + print(f"decode {args.decode_tokens/(et-pt):.3f} tok/s output {output}", flush=True) diff --git a/extra/gemm/cdna_asm_gemm.py b/extra/gemm/cdna_asm_gemm.py index ef3d35fd50..e9cb78e98a 100644 --- a/extra/gemm/cdna_asm_gemm.py +++ b/extra/gemm/cdna_asm_gemm.py @@ -6,6 +6,7 @@ from tinygrad.renderer import Estimates from tinygrad.helpers import getenv, all_same, DEBUG, ceildiv from tinygrad.runtime.support.compiler_amd import HIPCCCompiler from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8 +from extra.llama_kernels.quantize_mxfp4 import quantize_mxfp4 TILE_M, TILE_N, TILE_K = 256, 256, 64 @@ -125,6 +126,25 @@ def custom_mxfp4_gemm(C:UOp, A:UOp, B:UOp, scale_a:UOp, scale_b:UOp, *extra:UOp, insts = build_kernel(M, N, K, tile_m, tile_n) return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=x) for x in insts)))) +def _mxfp4_gemm_quantized(a_q:Tensor, b_q:Tensor, scale_a:Tensor, scale_b:Tensor) -> Tensor: + M, half_k = a_q.shape + N, half_k_b = b_q.shape + assert half_k == half_k_b + is_multi = isinstance(a_q.device, tuple) + reduce_out = is_multi and (a_q.uop.axis == 1 or b_q.uop.axis == 1) + if not is_multi: out = Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a_q.device) + elif reduce_out: out = Tensor(Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a_q.device).uop.unshard(0), device=a_q.device) + elif a_q.uop.axis == 0: + out = Tensor(Tensor.invalids(1, M//len(a_q.device), N, dtype=dtypes.bfloat16, device=a_q.device).uop.unshard(1), device=a_q.device) + elif b_q.uop.axis == 0: + out = Tensor(Tensor.invalids(1, M, N//len(a_q.device), dtype=dtypes.bfloat16, device=a_q.device).uop.unshard(2), device=a_q.device) + else: out = Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a_q.device) + tile_m, tile_n = next((tm, tn) for tm, tn in ((256, 256), (192, 256), (128, 512)) if M % tm == N % tn == 0) + out = Tensor.custom_kernel(out, a_q, b_q, scale_a, scale_b, + fxn=functools.partial(custom_mxfp4_gemm, tile_m=tile_m, tile_n=tile_n))[0] + if reduce_out: out = out.sum(0) + return out.squeeze(0) + def quantize_mxfp8(x:Tensor) -> tuple[Tensor, Tensor, Tensor]: # 1x32 block scaling along the last axis *batch, K = x.shape @@ -137,50 +157,6 @@ def quantize_mxfp8(x:Tensor) -> tuple[Tensor, Tensor, Tensor]: packed = mx_pack(e8) if len(batch) == 1 and scale_K % 4 == 0 else None return x_clamped.cast(FP8_DTYPE), e8, packed -def _mxfp4_shuffle_weight(x:Tensor) -> Tensor: - # shuffle_weight(x, layout=(16, 16)) on the packed uint8 buffer. - if x.ndim == 3: - ndev, rows, half_k = x.shape - return x.reshape(ndev, rows//16, 16, half_k//32, 2, 16).permute(0, 1, 3, 4, 2, 5).reshape(ndev, rows, half_k).contiguous() - rows, half_k = x.shape - return x.reshape(rows//16, 16, half_k//32, 2, 16).permute(0, 2, 3, 1, 4).reshape(rows, half_k).contiguous() - -def _mxfp4_shuffle_scales(x:Tensor) -> Tensor: - # e8m0_shuffle: each 256x8 scale tile is arranged for the raw MFMA scale loads. - if x.ndim == 3: - ndev, rows, scale_k = x.shape - return x.reshape(ndev, rows//32, 2, 16, scale_k//8, 2, 4).permute(0, 1, 4, 6, 3, 5, 2).reshape(ndev, rows, scale_k).contiguous() - rows, scale_k = x.shape - return x.reshape(rows//32, 2, 16, scale_k//8, 2, 4).permute(0, 3, 5, 2, 4, 1).reshape(rows, scale_k).contiguous() - -def quantize_mxfp4(x:Tensor) -> tuple[Tensor, Tensor, Tensor]: - # OCP MXFP4: 1x32 blocks, E2M1 values packed low-nibble first, and E8M0 scales. - *batch, K = x.shape - rows = math.prod(batch) - assert x.ndim >= 2 and K % 256 == 0 and rows % 32 == 0, \ - f"mxfp4 quantization needs rows%32 and K%256, got {x.shape}" - xb = x.float().reshape(*batch, K//32, 32) - amax = xb.abs().max(axis=-1) - - # even scale rounding: round the fp32 significand before choosing 2^(floor(log2)-2). - amax_rounded = ((amax.bitcast(dtypes.uint32) + 0x200000) & 0xFF800000).bitcast(dtypes.float32) - scale_exp = (amax_rounded.maximum(2**-126).log2().floor() - 2).clamp(-127, 127) - e8 = (scale_exp + 127).cast(dtypes.uint8) - scaled = xb * (-scale_exp).exp2().reshape(*batch, K//32, 1) - - mag = scaled.abs() - code = sum(x.cast(dtypes.uint8) for x in - (mag > .25, mag >= .75, mag > 1.25, mag >= 1.75, mag > 2.5, mag >= 3.5, mag > 5.0)) - code = code | ((scaled < 0).cast(dtypes.uint8) << 3) - code = code.reshape(*batch, K) - packed = code[..., 0::2] | (code[..., 1::2] << 4) - if isinstance(x.device, tuple) and x.uop.axis == x.ndim-2 and x.shape[x.uop.axis] == len(x.device): - axis = x.uop.axis - order = (axis, *range(axis), *range(axis+1, e8.ndim)) - e8_local = e8.permute(order) - return packed, e8, _mxfp4_shuffle_scales(e8_local.reshape(e8_local.shape[0], -1, K//32)) - return packed, e8, _mxfp4_shuffle_scales(e8.reshape(rows, K//32)) - def mx_pack(e8:Tensor) -> Tensor: rows, scale_K = e8.shape return e8.reshape(rows, scale_K // 4, 4).bitcast(dtypes.uint32).reshape(rows, scale_K // 4).permute(1, 0).contiguous() @@ -405,15 +381,16 @@ def custom_mx_gemm_bw(gradient:UOp, kernel:UOp, has_w_post:bool, w_stored:bool=F # ** mxfp4 gemm backward def custom_mxfp4_gemm_bw(gradient:UOp, kernel:UOp): - # The raw kernel consumes quantized buffers, while the final two inputs retain the BF16 operands for STE gradients. - inputs = kernel.src[1:] # (out, a_q, b_q, scale_a, scale_b, a, w) - assert len(inputs) == 7 + inputs = kernel.src[1:] # out, row operands/scales, BF16 operands, column operands/scales + assert len(inputs) == 11 a, w = Tensor(inputs[5], device=inputs[5].device), Tensor(inputs[6], device=inputs[6].device) + a_col, scale_a_col = Tensor(inputs[7], device=a.device), Tensor(inputs[8], device=a.device) + w_col, scale_w_col = Tensor(inputs[9], device=a.device), Tensor(inputs[10], device=a.device) g = Tensor(gradient, device=a.device)[:a.shape[0]].cast(dtypes.bfloat16) - grad_a = asm_gemm(g, w, mxfp4=True) - a_flat, g_flat = a.reshape(-1, a.shape[-1]), g.reshape(-1, g.shape[-1]) - grad_w = asm_gemm(g_flat.T, a_flat, mxfp4=True) - return (None, None, None, None, None, grad_a.uop, grad_w.uop) + g_row, scale_g_row, g_col, scale_g_col = quantize_mxfp4(g, flatten_row=True) + grad_a = _mxfp4_gemm_quantized(g_row, w_col, scale_g_row, scale_w_col).reshape(*a.shape[:-1], w.shape[-1]) + grad_w = _mxfp4_gemm_quantized(g_col, a_col, scale_g_col, scale_a_col).reshape(w.shape) + return (None, None, None, None, None, grad_a.uop, grad_w.uop, None, None, None, None) # ** main gemm function @@ -459,16 +436,10 @@ def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=N tile_m, tile_n = next((tm, tn) for tm, tn in ((256, 256), (192, 256), (128, 512)) if (batch*M) % tm == N % tn == 0) fxn = functools.partial(custom_mxfp4_gemm, tile_m=tile_m, tile_n=tile_n) w = b.T - if k_sharded: - ndev = len(a.device) - a_q, _, scale_a = quantize_mxfp4(a.reshape(batch, M, ndev, K)) - b_q, _, scale_b = quantize_mxfp4(w.reshape(w.shape[0], ndev, K)) - b_q = _mxfp4_shuffle_weight(b_q.permute(1, 0, 2)) - else: - a_q, _, scale_a = quantize_mxfp4(a.reshape(batch*M, K)) - b_q, _, scale_b = quantize_mxfp4(w) - a_q, b_q = a_q.reshape(batch, M, K//2).contiguous(), _mxfp4_shuffle_weight(b_q) - out = Tensor.custom_kernel(out, a_q, b_q, scale_a, scale_b, a, w, fxn=fxn, grad_fxn=custom_mxfp4_gemm_bw)[0] + a_q, scale_a, a_col, scale_a_col = quantize_mxfp4(a, shuffle_col=True) + b_q, scale_b, b_col, scale_b_col = quantize_mxfp4(w, shuffle_row=True, shuffle_col=True) + out = Tensor.custom_kernel(out, a_q, b_q, scale_a, scale_b, a, w, + a_col, scale_a_col, b_col, scale_b_col, fxn=fxn, grad_fxn=custom_mxfp4_gemm_bw)[0] elif mx: # mxfp8 1x32 block scaling if mx_scales is not None: diff --git a/extra/gemm/moe_routing.py b/extra/gemm/moe_routing.py index e683b83067..30ba001c3f 100644 --- a/extra/gemm/moe_routing.py +++ b/extra/gemm/moe_routing.py @@ -5,7 +5,8 @@ BLOCK_ROW = 256 def _sharded_invalids(shape:tuple[int, ...], dtype, device) -> Tensor: if isinstance(device, tuple): - return Tensor.invalids(*shape, dtype=dtype, device=device[0]).shard(device, axis=0) + per = Tensor.invalids(shape[0]//len(device), *shape[1:], dtype=dtype, device=device) + return Tensor(per.uop.unshard(0), device=device) return Tensor.invalids(*shape, dtype=dtype, device=device) def _atomic_add(device:str) -> str: diff --git a/extra/hcq2/ops_amd2.py b/extra/hcq2/ops_amd2.py index 788d40f495..fbdef9f444 100644 --- a/extra/hcq2/ops_amd2.py +++ b/extra/hcq2/ops_amd2.py @@ -182,7 +182,8 @@ def sdma_copy(ctx, call): src_addr, dst_addr = call.src[2].getaddr(ctx.devs), call.src[1].getaddr(ctx.devs) return call.ins(SDMAOps.COPY, src=tuple(UOp.const(x, dtypes.uint32) for off in range(0, sz, ctx.max_copy_size) for x in ( ctx.sdma.SDMA_OP_COPY | ctx.sdma.SDMA_PKT_COPY_LINEAR_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_COPY_LINEAR), - ctx.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz-off, ctx.max_copy_size)-1), 0, *data64_le(src_addr+off), *data64_le(dst_addr+off)))) + ctx.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz-off, ctx.max_copy_size)-1), 0, + *data64_le(src_addr+UOp.const(off, dtypes.uint64)), *data64_le(dst_addr+UOp.const(off, dtypes.uint64))))) def sdma_wait(ctx, ins, dst, val): op = ctx.sdma.SDMA_OP_POLL_REGMEM | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_FUNC(WAIT_REG_MEM_FUNCTION_GEQ) \ @@ -287,10 +288,10 @@ def amd_build_program(prg:UOp) -> UOp: class AMDAllocator(HCQAllocator['AMDDevice']): def __init__(self, dev:AMDDevice): - super().__init__(dev, supports_copy_from_disk=dev.has_sdma_queue, supports_transfer=dev.has_sdma_queue and not dev.is_usb()) + super().__init__(dev, supports_copy_from_disk=dev.has_copy_queue, supports_transfer=dev.has_copy_queue and not dev.is_usb()) def _alloc(self, size:int, options:BufferSpec) -> HCQ2Buffer: - return self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access or not self.dev.has_sdma_queue) + return self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access or not self.dev.has_copy_queue) def _do_free(self, opaque, options:BufferSpec): self.dev.iface.free(opaque) @@ -507,11 +508,12 @@ class PCIIface(PCIIfaceBase): if drain_only: d.iface.dev_impl.ih.drain() else: d.iface.dev_impl.ih.interrupt_handler() - if reset and d.iface.dev_impl.recover(): + if reset and d.iface.dev_impl.recover(force=True): cq = d.compute_queue for b in (cq.put_value, cq.read_ptr, cq.write_ptr): b._buf.view.view(fmt='Q')[0] = 0 d.iface.dev_impl.gfx.setup_ring(*cq.params) - d.signal('timeline')._buf.cpu_view().mv.cast('Q')[0] = d.signal('value', 1).as_memoryview(force_zero_copy=True).cast('Q')[0] - 1 + d.signal('timeline')._buf.cpu_view().mv.cast('Q')[0] = \ + d.signal('value', 1).as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] - 1 def sleep(self, timeout): if hasattr(self.pci_dev, 'irq_poller') and self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))): @@ -537,9 +539,12 @@ class AMDDevice(HCQ2Compiled): ]) timestamp_divider = 100.0 # AMD GPU clock: ticks/us + max_scratch_psize = 0 ifaces = [KFDIface, PCIIface, _mock(KFDIface, "MOCKIface"), _mock(KFDIface), _mock(PCIIface)] + def device_props(self): return self.iface.props + def is_am(self) -> bool: return isinstance(self.iface, (PCIIface,)) def is_usb(self) -> bool: return False @@ -576,7 +581,7 @@ class AMDDevice(HCQ2Compiled): self.max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000 self.sdma_queues:dict = {} - self.has_sdma_queue = True # self.sdma_queue(0) is not None, TODO: think of this + self.has_copy_queue = not getenv("AMD_DISABLE_SDMA") super().__init__(device, AMDAllocator(self), [HIPRenderer, AMDLLVMRenderer, HIPCCRenderer], None, can_recover=self.is_am(), arch=self.arch) @@ -689,7 +694,7 @@ class AMDDevice(HCQ2Compiled): return tmpring def scratch_buffer(self, private_segment_size): - private_segment_size = max(private_segment_size, 128) + AMDDevice.max_scratch_psize = private_segment_size = max(private_segment_size, 128, AMDDevice.max_scratch_psize) if self.max_private_segment_size < private_segment_size: lanes_per_wave = 64 # wave64 mem_alignment_size = 256 if self.target[0] != 9 else 1024 diff --git a/extra/llama_kernels/quantize_mxfp4/__init__.py b/extra/llama_kernels/quantize_mxfp4/__init__.py new file mode 100644 index 0000000000..6cdd983c6f --- /dev/null +++ b/extra/llama_kernels/quantize_mxfp4/__init__.py @@ -0,0 +1,36 @@ +import functools, math, pathlib +from tinygrad import Tensor, dtypes +from tinygrad.uop.ops import UOp, Ops, KernelInfo +from tinygrad.renderer import Estimates +from extra.llama_kernels import alloc_like, compile_hip + +@functools.cache +def _custom_quantize_mxfp4(row_fp4:UOp, row_scale:UOp, col_fp4:UOp, col_scale:UOp, x:UOp, *, shuffle_row:bool, shuffle_col:bool) -> UOp: + M, N = math.prod(x.shape[:-1]), x.shape[-1] + assert M % 256 == 0 and N % 256 == 0, f"MXFP4 quantization requires multiples of 256, got {x.shape}" + name = f"quantize_mxfp4_{int(shuffle_row)}_{int(shuffle_col)}_{M}_{N}" + mem = M*N*2 + M*N + M*N//16 # read bf16, write row+col fp4 + e8m0 + outputs = (row_fp4, row_scale, col_fp4, col_scale) + sink = UOp.sink(*(o.base for o in outputs), x.base, + *(UOp(Ops.CUSTOM, dtypes.void, (o.base.index(0),), arg="") for o in outputs), + UOp.special(256, "lidx0"), UOp.special(M//128, "gidx0"), UOp.special(N//64, "gidx1"), + arg=KernelInfo(name, estimates=Estimates(ops=12*M*N, mem=mem))) + src = (pathlib.Path(__file__).parent/"quantize_mxfp4.cpp").read_text() + return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src), + UOp(Ops.BINARY, arg=compile_hip(src, [f"-DKERNEL_NAME={name}", f"-DM_DIM={M}", f"-DN_DIM={N}", + f"-DSHUFFLE_ROWWISE_FP4_VALUE={int(shuffle_row)}", + f"-DSHUFFLE_COLWISE_FP4_VALUE={int(shuffle_col)}"])))) + +def quantize_mxfp4(x:Tensor, *, shuffle_row:bool=False, shuffle_col:bool=False, flatten_row:bool=False) -> tuple[Tensor, Tensor, Tensor, Tensor]: + assert x.dtype == dtypes.bfloat16 and x.ndim >= 2, f"expected BF16 matrix, got {x.dtype} {x.shape}" + M, N = math.prod(x.shape[:-1]), x.shape[-1] + assert M % 256 == 0 and N % 256 == 0, f"MXFP4 quantization requires multiples of 256, got {x.shape}" + axis = x.uop.axis if isinstance(x.device, tuple) else None + row_axis = 0 if flatten_row and axis is not None else axis + col_axis = None if axis is None else (0 if axis == x.ndim-1 else 1) + outputs = (alloc_like((M, N//2) if flatten_row else (*x.shape[:-1], N//2), dtypes.uint8, x.device, row_axis), + alloc_like((M, N//32) if flatten_row else (*x.shape[:-1], N//32), dtypes.uint8, x.device, row_axis), + alloc_like((N, M//2), dtypes.uint8, x.device, col_axis), + alloc_like((N, M//32), dtypes.uint8, x.device, col_axis)) + fxn = functools.partial(_custom_quantize_mxfp4, shuffle_row=shuffle_row, shuffle_col=shuffle_col) + return tuple(Tensor.custom_kernel(*outputs, x, fxn=fxn)[:4]) diff --git a/extra/llama_kernels/quantize_mxfp4/quantize_mxfp4.cpp b/extra/llama_kernels/quantize_mxfp4/quantize_mxfp4.cpp new file mode 100644 index 0000000000..e3a977bcc5 --- /dev/null +++ b/extra/llama_kernels/quantize_mxfp4/quantize_mxfp4.cpp @@ -0,0 +1,226 @@ +// Copyright (c) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT + +#include +#include + +#if !defined(KERNEL_NAME) || !defined(M_DIM) || !defined(N_DIM) || !defined(SHUFFLE_ROWWISE_FP4_VALUE) || \ + !defined(SHUFFLE_COLWISE_FP4_VALUE) +#error kernel dimensions and layouts must be defined +#endif + +namespace { + +constexpr int BLOCK = 32; +constexpr int TILE_M = 128; +constexpr int TILE_N = 64; +constexpr int THREADS = 256; +constexpr int THREADS_PER_ROW = 8; +constexpr int VALUES_PER_THREAD = 4; +constexpr int SMEM_STRIDE = BLOCK + 2; +constexpr int M = M_DIM; +constexpr int N = N_DIM; +constexpr int M_PACKED = M / 2; +constexpr int N_PACKED = N / 2; +constexpr int M_SCALES = M / BLOCK; +constexpr int N_SCALES = N / BLOCK; +constexpr bool SHUFFLE_ROWWISE_FP4 = SHUFFLE_ROWWISE_FP4_VALUE; +constexpr bool SHUFFLE_COLWISE_FP4 = SHUFFLE_COLWISE_FP4_VALUE; + +static_assert(M % 256 == 0 && N % 256 == 0); + +struct Quantized4 { + uint16_t fp4; + uint8_t scale; +}; + +__device__ __forceinline__ float swizzle_xor1(float value) { + float result; + asm volatile("ds_swizzle_b32 %0, %1 offset:0x041f\n\ts_waitcnt lgkmcnt(0)" : "=v"(result) : "v"(value)); + return result; +} + +__device__ __forceinline__ float swizzle_xor2(float value) { + float result; + asm volatile("ds_swizzle_b32 %0, %1 offset:0x081f\n\ts_waitcnt lgkmcnt(0)" : "=v"(result) : "v"(value)); + return result; +} + +__device__ __forceinline__ float swizzle_xor4(float value) { + float result; + asm volatile("ds_swizzle_b32 %0, %1 offset:0x101f\n\ts_waitcnt lgkmcnt(0)" : "=v"(result) : "v"(value)); + return result; +} + +__device__ __forceinline__ float max8(float value) { + value = fmaxf(value, swizzle_xor4(value)); + value = fmaxf(value, swizzle_xor2(value)); + return fmaxf(value, swizzle_xor1(value)); +} + +__device__ __forceinline__ float4 load_bf16x4(const uint16_t* values) { + const uint32_t lo = *reinterpret_cast(values); + const uint32_t hi = *reinterpret_cast(values + 2); + return make_float4(__uint_as_float(lo << 16), __uint_as_float(lo & 0xffff0000u), + __uint_as_float(hi << 16), __uint_as_float(hi & 0xffff0000u)); +} + +__device__ __forceinline__ void hadamard16(float4& value, int lane) { + const float a0 = value.x + value.y, a1 = value.x - value.y; + const float a2 = value.z + value.w, a3 = value.z - value.w; + value = make_float4(a0 + a2, a1 + a3, a0 - a2, a1 - a3); + + const float4 xor1 = make_float4(swizzle_xor1(value.x), swizzle_xor1(value.y), swizzle_xor1(value.z), swizzle_xor1(value.w)); + value = lane & 1 ? make_float4(xor1.x - value.x, xor1.y - value.y, xor1.z - value.z, xor1.w - value.w) + : make_float4(xor1.x + value.x, xor1.y + value.y, xor1.z + value.z, xor1.w + value.w); + + const float4 xor2 = make_float4(swizzle_xor2(value.x), swizzle_xor2(value.y), swizzle_xor2(value.z), swizzle_xor2(value.w)); + value = lane & 2 ? make_float4(xor2.x - value.x, xor2.y - value.y, xor2.z - value.z, xor2.w - value.w) + : make_float4(xor2.x + value.x, xor2.y + value.y, xor2.z + value.z, xor2.w + value.w); + value.x *= 0.25f; + value.y *= 0.25f; + value.z *= 0.25f; + value.w *= 0.25f; +} + +__device__ __forceinline__ uint8_t e8m0_scale(float amax, float& scale) { + if (amax == 0.0f) { + scale = 1.0f; + return 127; + } + + const uint32_t rounded = (__float_as_uint(amax) + 0x200000u) & 0xff800000u; + int exponent = static_cast((rounded >> 23) & 0xff) - 129; + exponent = exponent < -127 ? -127 : exponent > 127 ? 127 : exponent; + scale = exponent == -127 ? __uint_as_float(0x00400000u) : __uint_as_float(static_cast(exponent + 127) << 23); + return static_cast(exponent + 127); +} + +__device__ __forceinline__ uint16_t pack_fp4(float4 value, float scale) { + uint32_t lo = 0, hi = 0; + asm volatile("v_cvt_scalef32_pk_fp4_f32 %0, %1, %2, %3" : "+v"(lo) : "v"(value.x), "v"(value.y), "v"(scale)); + asm volatile("v_cvt_scalef32_pk_fp4_f32 %0, %1, %2, %3" : "+v"(hi) : "v"(value.z), "v"(value.w), "v"(scale)); + return static_cast(lo | (hi << 8)); +} + +__device__ __forceinline__ Quantized4 quantize(float4 value, int lane) { + hadamard16(value, lane); + const float local_max = fmaxf(fmaxf(fabsf(value.x), fabsf(value.y)), fmaxf(fabsf(value.z), fabsf(value.w))); + float scale; + const uint8_t e8m0 = e8m0_scale(max8(local_max), scale); + return {pack_fp4(value, scale), e8m0}; +} + +__device__ __forceinline__ void store_scale(uint8_t* output, int row, int col, int cols, uint8_t value) { + const int tile = ((row >> 5) * (cols >> 3) + (col >> 3)) << 8; + const int offset = ((col & 3) << 6) + ((row & 15) << 2) + (((col >> 2) & 1) << 1) + ((row >> 4) & 1); + output[tile + offset] = value; +} + +template +__device__ __forceinline__ void store_fp4(uint8_t* output, int row, int col, int packed_cols, uint16_t value) { + int index = row * packed_cols + col; + if constexpr (Shuffled) { + const int tile = (row >> 4) * (packed_cols << 4) + (col >> 5) * 512; + const int offset = ((col >> 4) & 1) * 256 + (row & 15) * 16 + (col & 15); + index = tile + offset; + } + *reinterpret_cast(output + index) = value; +} + +__device__ __forceinline__ void load_tile(uint16_t* tile, const uint16_t* input, int tile_m, int tile_n) { + const int row = threadIdx.x / THREADS_PER_ROW; + const int col = threadIdx.x % THREADS_PER_ROW * VALUES_PER_THREAD; + const uint64_t packed = *reinterpret_cast(input + (tile_m + row) * N + tile_n + col); + *reinterpret_cast(tile + row * SMEM_STRIDE + col) = static_cast(packed); + *reinterpret_cast(tile + row * SMEM_STRIDE + col + 2) = static_cast(packed >> 32); +} + +__device__ __forceinline__ void quantize_row(uint16_t* tile, uint8_t* fp4_output, uint8_t* scale_output, + int tile_m, int tile_n, int local_row, int lane) { + const int row = tile_m + local_row; + const int col = lane * VALUES_PER_THREAD; + const Quantized4 result = quantize(load_bf16x4(tile + local_row * SMEM_STRIDE + col), lane); + store_fp4(fp4_output, row, (tile_n + col) / 2, N_PACKED, result.fp4); + if (lane == 0) store_scale(scale_output, row, tile_n / BLOCK, N_SCALES, result.scale); +} + +__device__ __forceinline__ Quantized4 quantize_col(uint16_t* tile, int col, int lane) { + const int row = lane * VALUES_PER_THREAD; + return quantize(make_float4( + __uint_as_float(static_cast(tile[(row + 0) * SMEM_STRIDE + col]) << 16), + __uint_as_float(static_cast(tile[(row + 1) * SMEM_STRIDE + col]) << 16), + __uint_as_float(static_cast(tile[(row + 2) * SMEM_STRIDE + col]) << 16), + __uint_as_float(static_cast(tile[(row + 3) * SMEM_STRIDE + col]) << 16)), lane); +} + +} // namespace + +extern "C" __global__ __launch_bounds__(THREADS, 8) +void KERNEL_NAME(uint8_t* __restrict__ rowwise_fp4, uint8_t* __restrict__ rowwise_scale, + uint8_t* __restrict__ colwise_fp4, uint8_t* __restrict__ colwise_scale, + const uint16_t* __restrict__ input) { + __shared__ uint16_t tile[BLOCK * SMEM_STRIDE]; + const int tid = threadIdx.x; + const int line = tid / THREADS_PER_ROW; + const int lane = tid % THREADS_PER_ROW; + const int block_m = blockIdx.x * TILE_M; + const int block_n = blockIdx.y * TILE_N; + + if constexpr (!SHUFFLE_COLWISE_FP4) { + uint16_t col_fp4[TILE_N / BLOCK][TILE_M / BLOCK]; + uint8_t col_scale[TILE_N / BLOCK][TILE_M / BLOCK]; + + for (int chunk_m = 0; chunk_m < TILE_M / BLOCK; chunk_m++) { + for (int chunk_n = 0; chunk_n < TILE_N / BLOCK; chunk_n++) { + const int tile_m = block_m + chunk_m * BLOCK; + const int tile_n = block_n + chunk_n * BLOCK; + load_tile(tile, input, tile_m, tile_n); + __syncthreads(); + + quantize_row(tile, rowwise_fp4, rowwise_scale, tile_m, tile_n, line, lane); + const Quantized4 result = quantize_col(tile, line, lane); + col_fp4[chunk_n][chunk_m] = result.fp4; + col_scale[chunk_n][chunk_m] = result.scale; + __syncthreads(); + } + } + + for (int chunk_n = 0; chunk_n < TILE_N / BLOCK; chunk_n++) { + for (int chunk_m = 0; chunk_m < TILE_M / BLOCK; chunk_m++) + tile[line * BLOCK + chunk_m * THREADS_PER_ROW + lane] = col_fp4[chunk_n][chunk_m]; + __syncthreads(); + + for (int round = 0; round < BLOCK / THREADS_PER_ROW; round++) { + const int col = round * THREADS_PER_ROW + tid / BLOCK; + const int row_pair = tid % BLOCK; + *reinterpret_cast(colwise_fp4 + (block_n + chunk_n * BLOCK + col) * M_PACKED + block_m / 2 + row_pair * 2) = + tile[col * BLOCK + row_pair]; + } + + if (lane == 0) { + const int col = block_n + chunk_n * BLOCK + line; + for (int chunk_m = 0; chunk_m < TILE_M / BLOCK; chunk_m++) + store_scale(colwise_scale, col, block_m / BLOCK + chunk_m, M_SCALES, col_scale[chunk_n][chunk_m]); + } + __syncthreads(); + } + } else { + for (int chunk_m = 0; chunk_m < TILE_M / BLOCK; chunk_m++) { + for (int chunk_n = 0; chunk_n < TILE_N / BLOCK; chunk_n++) { + const int tile_m = block_m + chunk_m * BLOCK; + const int tile_n = block_n + chunk_n * BLOCK; + load_tile(tile, input, tile_m, tile_n); + __syncthreads(); + + quantize_row(tile, rowwise_fp4, rowwise_scale, tile_m, tile_n, line, lane); + const int row = lane * VALUES_PER_THREAD; + const int col = tile_n + line; + const Quantized4 result = quantize_col(tile, line, lane); + store_fp4(colwise_fp4, col, (tile_m + row) / 2, M_PACKED, result.fp4); + if (lane == 0) store_scale(colwise_scale, col, tile_m / BLOCK, M_SCALES, result.scale); + __syncthreads(); + } + } + } +} diff --git a/extra/llama_kernels/swiglu/__init__.py b/extra/llama_kernels/swiglu/__init__.py new file mode 100644 index 0000000000..810cb6d8ec --- /dev/null +++ b/extra/llama_kernels/swiglu/__init__.py @@ -0,0 +1,49 @@ +import functools, math +from tinygrad import Tensor, dtypes +from tinygrad.uop.ops import UOp, KernelInfo +from tinygrad.renderer import Estimates +from extra.llama_kernels import alloc_like + +LOG2E = 1.4426950408889634 + +@functools.cache +def _custom_swiglu(out:UOp, x_w13:UOp) -> UOp: + rows, hidden = math.prod(x_w13.shape[:-1]), x_w13.shape[-1]//2 + n_elems = rows * hidden + out, x_w13 = out.reshape(n_elems), x_w13.reshape(rows, 2*hidden) + i = UOp.range(n_elems, 0) + row, col = i // hidden, i % hidden + act, gate = x_w13[row, col].cast(dtypes.float), x_w13[row, hidden+col].cast(dtypes.float) + sigmoid = (1.0 + (-LOG2E * act).exp2()).reciprocal() + store = out[i].store((act * sigmoid * gate).cast(out.dtype)) + return store.end(i).sink(arg=KernelInfo(f"swiglu_fwd_{n_elems}", estimates=Estimates(ops=5*n_elems, mem=6*n_elems))) + +@functools.cache +def _custom_swiglu_bwd(grad_out:UOp, x_w13:UOp, grad_act:UOp) -> UOp: + rows, hidden = math.prod(x_w13.shape[:-1]), x_w13.shape[-1]//2 + n_elems = rows * hidden + grad_out, x_w13, grad_act = grad_out.reshape(rows, 2*hidden), x_w13.reshape(rows, 2*hidden), grad_act.reshape(n_elems) + i = UOp.range(n_elems, 0) + row, col = i // hidden, i % hidden + act, gate = x_w13[row, col].cast(dtypes.float), x_w13[row, hidden+col].cast(dtypes.float) + grad = grad_act[i].cast(dtypes.float) + sigmoid = (1.0 + (-LOG2E * act).exp2()).reciprocal() + silu = act * sigmoid + dact = grad_out[row, col].store((grad * (sigmoid + silu * (1.0 - sigmoid)) * gate).cast(grad_out.dtype)) + dgate = grad_out.after(dact)[row, hidden+col].store((grad * silu).cast(grad_out.dtype)) + return dgate.end(i).sink(arg=KernelInfo(f"swiglu_bwd_{n_elems}", estimates=Estimates(ops=10*n_elems, mem=10*n_elems))) + +def _swiglu_bwd(gradient:UOp, kernel:UOp): + _, x_w13 = kernel.src[1:] + axis = x_w13.axis if isinstance(x_w13.device, tuple) else None + grad_out = alloc_like(x_w13.shape, dtypes.bfloat16, x_w13.device, axis) + grad_out, *_ = Tensor.custom_kernel(grad_out, Tensor(x_w13, device=x_w13.device), Tensor(gradient, device=x_w13.device), + fxn=_custom_swiglu_bwd) + return (None, grad_out.uop) + +def swiglu(x_w13:Tensor) -> Tensor: + assert x_w13.dtype == dtypes.bfloat16 and x_w13.ndim >= 2 and x_w13.shape[-1] % 32 == 0 + *prefix, two_k = x_w13.shape + axis = x_w13.uop.axis if isinstance(x_w13.device, tuple) else None + out = alloc_like((*prefix, two_k//2), dtypes.bfloat16, x_w13.device, axis) + return Tensor.custom_kernel(out, x_w13, fxn=_custom_swiglu, grad_fxn=_swiglu_bwd)[0] diff --git a/extra/thunder/amd/fa.py b/extra/thunder/amd/fa.py index c273cfca90..bbb4d46857 100644 --- a/extra/thunder/amd/fa.py +++ b/extra/thunder/amd/fa.py @@ -2,7 +2,7 @@ import math, pathlib, functools, struct from tinygrad import Device, Tensor from tinygrad.dtype import DTypeLike, dtypes -from tinygrad.helpers import DEBUG +from tinygrad.helpers import DEBUG, getenv from tinygrad.renderer import Estimates from tinygrad.runtime.support.compiler_amd import HIPCCCompiler from tinygrad.runtime.support.elf import elf_loader @@ -206,10 +206,11 @@ def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, sinks:UOp|None=None arg=KernelInfo(name="custom_fa_forward", estimates=estimates)) lib = HIPCCCompiler(arch, compile_args).compile_cached(code) - lib = bytearray(lib) - rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata") - struct.pack_into('0: sliding-window backward (query i sees keys in [i-WINDOW+1, i]) +#ifndef WINDOW +#define WINDOW 0 +#endif #define NUM_WARPS 4 #define NUM_THREADS (kittens::WARP_THREADS * NUM_WARPS) @@ -88,7 +92,12 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr const int k_start_min = j_min * WARP_SIZE_KV; // first Q step that can overlap this K_span: const int first_step = max(0, k_start_min / STEP_QO); +#if WINDOW + // cap the Q loop, padded by 2 masked steps: the epilogue's deferred dq path miscomputes in-window tail queries + const int num_steps_per_head = min(total_steps_per_head - first_step, (BLOCK_SIZE_KV + WINDOW) / STEP_QO + 2); +#else const int num_steps_per_head = total_steps_per_head - first_step; +#endif const int num_steps = num_steps_per_head * HEADS_PER_WG; const int k_pos = j * WARP_SIZE_KV; @@ -380,6 +389,13 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr mov<0, 1, neg_inf_v>(P_ij); mov<0, 2, neg_inf_v>(P_ij); mov<0, 3, neg_inf_v>(P_ij); +#if WINDOW + // window lower boundary, mirror of the causal edge + } else if (q_pos - k_pos == WINDOW) { + make_window<0, 0, neg_inf_v>(P_ij, P_ij); + } else if (q_pos - k_pos > WINDOW) { + mov(P_ij); +#endif } } mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR); @@ -640,6 +656,13 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr make_causal<0, 1, neg_inf_v>(P_ij, P_ij); mov<0, 2, neg_inf_v>(P_ij); mov<0, 3, neg_inf_v>(P_ij); +#if WINDOW + } else if (q_pos - k_pos == WINDOW) { + mov<0, 0, neg_inf_v>(P_ij); + make_window<0, 1, neg_inf_v>(P_ij, P_ij); + } else if (q_pos - k_pos > WINDOW) { + mov(P_ij); +#endif } } mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR); @@ -899,6 +922,14 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr // Apply the causal mask to [0, 2] and set [0, 3:4] to -inf make_causal<0, 2, neg_inf_v>(P_ij, P_ij); mov<0, 3, neg_inf_v>(P_ij); +#if WINDOW + } else if (q_pos - k_pos == WINDOW) { + mov<0, 0, neg_inf_v>(P_ij); + mov<0, 1, neg_inf_v>(P_ij); + make_window<0, 2, neg_inf_v>(P_ij, P_ij); + } else if (q_pos - k_pos > WINDOW) { + mov(P_ij); +#endif } } mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR); @@ -1157,6 +1188,15 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr } else if (q_pos == k_pos) { // Apply the causal mask to [0, 3] make_causal<0, 3, neg_inf_v>(P_ij, P_ij); +#if WINDOW + } else if (q_pos - k_pos == WINDOW) { + mov<0, 0, neg_inf_v>(P_ij); + mov<0, 1, neg_inf_v>(P_ij); + mov<0, 2, neg_inf_v>(P_ij); + make_window<0, 3, neg_inf_v>(P_ij, P_ij); + } else if (q_pos - k_pos > WINDOW) { + mov(P_ij); +#endif } } mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR); @@ -1436,6 +1476,13 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr mov<0, 1, neg_inf_v>(P_ij); mov<0, 2, neg_inf_v>(P_ij); mov<0, 3, neg_inf_v>(P_ij); +#if WINDOW + // window lower boundary, mirror of the causal edge + } else if (q_pos - k_pos == WINDOW) { + make_window<0, 0, neg_inf_v>(P_ij, P_ij); + } else if (q_pos - k_pos > WINDOW) { + mov(P_ij); +#endif } } mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR); @@ -1699,6 +1746,13 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr make_causal<0, 1, neg_inf_v>(P_ij, P_ij); mov<0, 2, neg_inf_v>(P_ij); mov<0, 3, neg_inf_v>(P_ij); +#if WINDOW + } else if (q_pos - k_pos == WINDOW) { + mov<0, 0, neg_inf_v>(P_ij); + make_window<0, 1, neg_inf_v>(P_ij, P_ij); + } else if (q_pos - k_pos > WINDOW) { + mov(P_ij); +#endif } } mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR); @@ -1958,6 +2012,14 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr // Apply the causal mask to [0, 2] and set [0, 3:4] to -inf make_causal<0, 2, neg_inf_v>(P_ij, P_ij); mov<0, 3, neg_inf_v>(P_ij); +#if WINDOW + } else if (q_pos - k_pos == WINDOW) { + mov<0, 0, neg_inf_v>(P_ij); + mov<0, 1, neg_inf_v>(P_ij); + make_window<0, 2, neg_inf_v>(P_ij, P_ij); + } else if (q_pos - k_pos > WINDOW) { + mov(P_ij); +#endif } } mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR); @@ -2216,6 +2278,15 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr } else if (q_pos == k_pos) { // Apply the causal mask to [0, 3] make_causal<0, 3, neg_inf_v>(P_ij, P_ij); +#if WINDOW + } else if (q_pos - k_pos == WINDOW) { + mov<0, 0, neg_inf_v>(P_ij); + mov<0, 1, neg_inf_v>(P_ij); + mov<0, 2, neg_inf_v>(P_ij); + make_window<0, 3, neg_inf_v>(P_ij, P_ij); + } else if (q_pos - k_pos > WINDOW) { + mov(P_ij); +#endif } } mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR); @@ -2487,6 +2558,12 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr mov<0, 1, neg_inf_v>(P_ij); mov<0, 2, neg_inf_v>(P_ij); mov<0, 3, neg_inf_v>(P_ij); +#if WINDOW + } else if (q_pos - k_pos == WINDOW) { + make_window<0, 0, neg_inf_v>(P_ij, P_ij); + } else if (q_pos - k_pos > WINDOW) { + mov(P_ij); +#endif } } mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR); @@ -2748,6 +2825,13 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr make_causal<0, 1, neg_inf_v>(P_ij, P_ij); mov<0, 2, neg_inf_v>(P_ij); mov<0, 3, neg_inf_v>(P_ij); +#if WINDOW + } else if (q_pos - k_pos == WINDOW) { + mov<0, 0, neg_inf_v>(P_ij); + make_window<0, 1, neg_inf_v>(P_ij, P_ij); + } else if (q_pos - k_pos > WINDOW) { + mov(P_ij); +#endif } } mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR); @@ -3004,6 +3088,14 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr // Apply the causal mask to [0, 2] and set [0, 3:4] to -inf make_causal<0, 2, neg_inf_v>(P_ij, P_ij); mov<0, 3, neg_inf_v>(P_ij); +#if WINDOW + } else if (q_pos - k_pos == WINDOW) { + mov<0, 0, neg_inf_v>(P_ij); + mov<0, 1, neg_inf_v>(P_ij); + make_window<0, 2, neg_inf_v>(P_ij, P_ij); + } else if (q_pos - k_pos > WINDOW) { + mov(P_ij); +#endif } } mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR); @@ -3260,6 +3352,15 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr } else if (q_pos == k_pos) { // Apply the causal mask to [0, 3] make_causal<0, 3, neg_inf_v>(P_ij, P_ij); +#if WINDOW + } else if (q_pos - k_pos == WINDOW) { + mov<0, 0, neg_inf_v>(P_ij); + mov<0, 1, neg_inf_v>(P_ij); + mov<0, 2, neg_inf_v>(P_ij); + make_window<0, 3, neg_inf_v>(P_ij, P_ij); + } else if (q_pos - k_pos > WINDOW) { + mov(P_ij); +#endif } } mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR); diff --git a/extra/thunder/amd/fa_fwd_causal.cpp b/extra/thunder/amd/fa_fwd_causal.cpp index 21e3c3b2eb..3deb93bf46 100644 --- a/extra/thunder/amd/fa_fwd_causal.cpp +++ b/extra/thunder/amd/fa_fwd_causal.cpp @@ -34,6 +34,10 @@ constexpr int ATTN_D = 128; // dimension constexpr int Q_BLOCK_SIZE = 32; // q block size constexpr int KV_BLOCK_SIZE = 64; // kv block size constexpr bool causal = true; +// WINDOW>0: sliding-window attention, query i attends keys in [i-WINDOW+1, i] +#ifndef WINDOW +#define WINDOW 0 +#endif #define NUM_WARPS 8 #define NUM_THREADS (kittens::WARP_THREADS * NUM_WARPS) @@ -82,11 +86,26 @@ template using attn /**********************************************************/ template -__device__ inline void mask_vec2_imm(uint32_t rel_vgpr, uint32_t neg_inf_vgpr, +__device__ inline void mask_vec2_imm(uint32_t rel_vgpr, uint32_t rel_hi_vgpr, uint32_t neg_inf_vgpr, uint32_t& x_ref, uint32_t& y_ref) { uint64_t x_mask, y_mask; - // uint32_t ox, oy; +#if WINDOW + // causal+window in one asm block to not disturb register allocation + asm volatile( + "v_cmp_lt_i32_e64 %0, %4, %5\n\t" + "v_cmp_lt_i32_e64 %1, %4, %7\n\t" + "v_cndmask_b32_e64 %2, %2, %6, %0\n\t" + "v_cndmask_b32_e64 %3, %3, %6, %1\n\t" + "v_cmp_ge_i32_e64 %0, %8, %5\n\t" + "v_cmp_ge_i32_e64 %1, %8, %7\n\t" + "v_cndmask_b32_e64 %2, %2, %6, %0\n\t" + "v_cndmask_b32_e64 %3, %3, %6, %1\n\t" + : "=s"(x_mask), "=s"(y_mask), "+v"(x_ref), "+v"(y_ref) + : "v"(rel_vgpr), "n"(THR_X), "v"(neg_inf_vgpr), "n"(THR_Y), "v"(rel_hi_vgpr) + : "vcc" + ); +#else asm volatile( // x: rel < THR_X ? "v_cmp_lt_i32_e64 %0, %6, %7\n\t" @@ -99,7 +118,7 @@ __device__ inline void mask_vec2_imm(uint32_t rel_vgpr, uint32_t neg_inf_vgpr, "n"(THR_X), "v"(neg_inf_vgpr), "n"(THR_Y) : "vcc" ); - // x_ref = ox; y_ref = oy; +#endif } template @@ -122,6 +141,8 @@ __device__ inline void mask_kv_tile(RT &dst, int q_abs, int k_abs, uint32_t neg_ // (smaller rel ⇒ more "future" keys that must be -inf) const int rel0 = q_pos - (k_base + row_base); const uint32_t rel = static_cast(rel0); + // rel-WINDOW keeps THR within the inline-constant range + const uint32_t rel_hi = static_cast(rel0 - WINDOW); #pragma unroll for (int j = 0; j < dst.width; ++j) { @@ -145,14 +166,14 @@ __device__ inline void mask_kv_tile(RT &dst, int q_abs, int k_abs, uint32_t neg_ // - reuse a single neg_inf register // - keep VCC live across the pair // - avoid reloading -inf or recomputing rel - mask_vec2_imm< 0, 1 >(rel, neg_inf_v, d0x, d0y); - mask_vec2_imm< 2, 3 >(rel, neg_inf_v, d1x, d1y); - mask_vec2_imm< 8, 9 >(rel, neg_inf_v, d2x, d2y); - mask_vec2_imm<10,11 >(rel, neg_inf_v, d3x, d3y); - mask_vec2_imm<16,17 >(rel, neg_inf_v, d4x, d4y); - mask_vec2_imm<18,19 >(rel, neg_inf_v, d5x, d5y); - mask_vec2_imm<24,25 >(rel, neg_inf_v, d6x, d6y); - mask_vec2_imm<26,27 >(rel, neg_inf_v, d7x, d7y); + mask_vec2_imm< 0, 1 >(rel, rel_hi, neg_inf_v, d0x, d0y); + mask_vec2_imm< 2, 3 >(rel, rel_hi, neg_inf_v, d1x, d1y); + mask_vec2_imm< 8, 9 >(rel, rel_hi, neg_inf_v, d2x, d2y); + mask_vec2_imm<10,11 >(rel, rel_hi, neg_inf_v, d3x, d3y); + mask_vec2_imm<16,17 >(rel, rel_hi, neg_inf_v, d4x, d4y); + mask_vec2_imm<18,19 >(rel, rel_hi, neg_inf_v, d5x, d5y); + mask_vec2_imm<24,25 >(rel, rel_hi, neg_inf_v, d6x, d6y); + mask_vec2_imm<26,27 >(rel, rel_hi, neg_inf_v, d7x, d7y); } } } @@ -201,6 +222,16 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p else max_num_tiles = num_tiles; const int q_start_pos = tile_idx * Q_BLOCK_SIZE; +#if WINDOW + // start at the first in-window tile; clamp keeps >=4 tiles for the pipeline unroll + const int block_min_q = block_tile_idx * NUM_WARPS * Q_BLOCK_SIZE; + int min_tile = (block_min_q - WINDOW + 1) / KV_BLOCK_SIZE; + if (min_tile < 0) min_tile = 0; + if (min_tile > max_num_tiles - 4) min_tile = max(0, max_num_tiles - 4); +#else + constexpr int min_tile = 0; +#endif + constexpr float TEMPERATURE_SCALE = (D == 128) ? 0.08838834764f*1.44269504089f : 0.125f*1.44269504089f; uint32_t neg_inf_v = 0xff800000; @@ -231,7 +262,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p G::prefill_swizzled_offsets<1, false>(k_smem[0], g.Kg, swizzled_offsets_K); G::prefill_swizzled_offsets<1, false>(v_smem[0], g.Vg, swizzled_offsets_V); - G::load<1, false>(k_smem[0], g.Kg, {batch_idx, 0, head_idx_kv, 0}, swizzled_offsets_K); + G::load<1, false>(k_smem[0], g.Kg, {batch_idx, min_tile, head_idx_kv, 0}, swizzled_offsets_K); __builtin_amdgcn_s_waitcnt(0); __builtin_amdgcn_sched_barrier(0); __builtin_amdgcn_s_barrier(); @@ -243,9 +274,9 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p transpose(q_reg_transposed, q_reg); // All warps then collaboratively load in the first slice of V (V0) and the second slice of K (K1) into shared memory - G::load<1, false>(k_smem[1], g.Kg, {batch_idx, 1, head_idx_kv, 0}, swizzled_offsets_K); + G::load<1, false>(k_smem[1], g.Kg, {batch_idx, min_tile + 1, head_idx_kv, 0}, swizzled_offsets_K); // All warps then load in the first slice of K (K0) - G::load<1, false>(v_smem[0], g.Vg, {batch_idx, 0, head_idx_kv, 0}, swizzled_offsets_V); + G::load<1, false>(v_smem[0], g.Vg, {batch_idx, min_tile, head_idx_kv, 0}, swizzled_offsets_V); load(k_reg, k_smem[0]); __builtin_amdgcn_sched_barrier(0); asm volatile("s_waitcnt lgkmcnt(0)"); @@ -259,13 +290,20 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p mma_AtB(att_block[0], k_reg_transposed, q_reg_transposed, att_block[0]); __builtin_amdgcn_sched_barrier(0); if constexpr (causal) { - const int kv_end_pos = (1) * KV_BLOCK_SIZE; - if (__builtin_expect(q_start_pos < kv_end_pos, 0)) { // Only mask if needed - mask_kv_tile(att_block[0], tile_idx, 0, neg_inf_v, lane); + const int kv_end_pos = (min_tile + 1) * KV_BLOCK_SIZE; + if (__builtin_expect(WINDOW || q_start_pos < kv_end_pos, WINDOW ? 1 : 0)) { + mask_kv_tile(att_block[0], tile_idx, min_tile, neg_inf_v, lane); } } // Each warp performs a partial softmax of QK0 (i.e. some of the online softmax up until but not including the second exponential scaling of the attention block likely) +#if WINDOW + // floor the max: min_tile can be fully masked, which would NaN via exp2(-inf - -inf) + zero(max_vec_prev); + add(max_vec_prev, max_vec_prev, -1e4f); + col_max(max_vec, att_block[0], max_vec_prev); +#else col_max(max_vec, att_block[0]); +#endif copy(max_vec_prev, max_vec); exp2(scale_vec, scale_vec); @@ -284,21 +322,25 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p // All warps then load in the second slice of K (K1) load(k_reg, k_smem[1]); // All warps then collaboratively load in the third slice of K (K2) into shared memory - G::load<1, false>(k_smem[0], g.Kg, {batch_idx, 2, head_idx_kv, 0}, swizzled_offsets_K); + G::load<1, false>(k_smem[0], g.Kg, {batch_idx, min_tile + 2, head_idx_kv, 0}, swizzled_offsets_K); // All warps then collaboratively load in the second slice of V (V1) into shared memory - G::load<1, false>(v_smem[1], g.Vg, {batch_idx, 1, head_idx_kv, 0}, swizzled_offsets_V); + G::load<1, false>(v_smem[1], g.Vg, {batch_idx, min_tile + 1, head_idx_kv, 0}, swizzled_offsets_V); asm volatile("s_waitcnt lgkmcnt(0)"); asm volatile("s_waitcnt vmcnt(" FA_VM4 ")"); __builtin_amdgcn_sched_barrier(0); __builtin_amdgcn_s_barrier(); // hot loop - for (int j = 3; j < max_num_tiles - 1; j += 2) { + for (int j = min_tile + 3; j < max_num_tiles - 1; j += 2) { // Cluster 0: // QK1 zero(att_block[1]); transpose(k_reg_transposed, k_reg); mma_AtB(att_block[1], k_reg_transposed, q_reg_transposed, att_block[1]); +#if WINDOW + // window masks interior tiles that causal skips + mask_kv_tile(att_block[1], tile_idx, j - 2, neg_inf_v, lane); +#endif // Finish softmax for QK0 exp2(att_block[0].tiles[1][0], att_block[0].tiles[1][0]); mul(norm_vec, norm_vec, scale_vec); @@ -379,7 +421,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p load(v_reg, v_smem[1]); if constexpr (causal) { const int kv_end_pos = (j) * KV_BLOCK_SIZE; - if (q_start_pos < kv_end_pos) { // Only mask if needed + if (WINDOW || q_start_pos < kv_end_pos) { mask_kv_tile(att_block[0], tile_idx, j - 1, neg_inf_v, lane); } } @@ -447,7 +489,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p load(v_reg, v_smem[0]); if constexpr (causal) { const int kv_end_pos = (max_num_tiles - 2) * KV_BLOCK_SIZE; - if (__builtin_expect(q_start_pos < kv_end_pos, 0)) { // Only mask if needed + if (__builtin_expect(WINDOW || q_start_pos < kv_end_pos, WINDOW ? 1 : 0)) { mask_kv_tile(att_block[1], tile_idx, max_num_tiles - 3, neg_inf_v, lane); } } @@ -510,7 +552,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p load(v_reg, v_smem[1]); if constexpr (causal) { const int kv_end_pos = (max_num_tiles - 1) * KV_BLOCK_SIZE; - if (__builtin_expect(q_start_pos < kv_end_pos, 1)) { // Only mask if needed + if (__builtin_expect(WINDOW || q_start_pos < kv_end_pos, 1)) { mask_kv_tile(att_block[0], tile_idx, max_num_tiles - 2, neg_inf_v, lane); } } @@ -572,7 +614,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p load(v_reg, v_smem[0]); if constexpr (causal) { const int kv_end_pos = (max_num_tiles) * KV_BLOCK_SIZE; - if (__builtin_expect(q_start_pos < kv_end_pos, 1)) { // Only mask if needed + if (__builtin_expect(WINDOW || q_start_pos < kv_end_pos, 1)) { mask_kv_tile(att_block[1], tile_idx, max_num_tiles - 1, neg_inf_v, lane); } } diff --git a/extra/thunder/amd/include/utils.cpp b/extra/thunder/amd/include/utils.cpp index 594e5e2749..f2174f2bd8 100644 --- a/extra/thunder/amd/include/utils.cpp +++ b/extra/thunder/amd/include/utils.cpp @@ -97,4 +97,33 @@ __device__ inline static void atomic_pk_add_bf16_with_warpid(const GL &dst, cons }(std::make_index_sequence{}); }.template operator()(), ...); }(std::make_index_sequence{}); -} \ No newline at end of file +} +// make_window: complement of make_causal for the window lower boundary (q_pos-k_pos == WINDOW). masks = ~(causal masks) +template +__device__ static inline void make_window(T0 &dst, const T1 &src) { + static_assert(std::is_same_v && std::is_same_v, "Only float to float window mask is supported"); + static_assert(std::is_same_v, "Only same layout is supported"); + static_assert(std::is_same_v, "Only same shape is supported"); + + if constexpr (std::is_same_v && std::is_same_v) { + using range_type_T0 = ducks::art::get_nth_range_t; + using registers_T0 = ducks::art::split_many_t, 1>; + using range_type_T1 = ducks::art::get_nth_range_t; + using registers_T1 = ducks::art::split_many_t, 1>; + static_assert(registers_T0::size == registers_T1::size); + + uint64_t window_mask = 0x1FFF01FF001F0001; + macros::v_cndmask_b32_e64::lo, ducks::art::get_nth_range_t::lo, GPR>(window_mask); + + window_mask = 0x3FFF03FF003F0003; + macros::v_cndmask_b32_e64::lo, ducks::art::get_nth_range_t::lo, GPR>(window_mask); + + window_mask = 0x7FFF07FF007F0007; + macros::v_cndmask_b32_e64::lo, ducks::art::get_nth_range_t::lo, GPR>(window_mask); + + window_mask = 0xFFFF0FFF00FF000F; + macros::v_cndmask_b32_e64::lo, ducks::art::get_nth_range_t::lo, GPR>(window_mask); + } else { + static_assert(false, "Unsupported window mask"); + } +} diff --git a/test/amd/hw/test_vopc.py b/test/amd/hw/test_vopc.py index 312244004c..5a7ed5e938 100644 --- a/test/amd/hw/test_vopc.py +++ b/test/amd/hw/test_vopc.py @@ -471,6 +471,20 @@ class TestCmpFloat(unittest.TestCase): st = run_program(instructions, n_lanes=1) self.assertEqual(st.vcc & 1, 1, "Expected vcc=1 (1.0 != 2.0)") + def test_v_cmp_eq_f16_src0_hi(self): + """v_cmp_eq_f16 with src0 from high half (true16 384+n encoding).""" + cmp = v_cmp_eq_f16_e32(v[0], v[1]) + cmp._raw += 128 # src0 v[0] -> v[0].h, the dsl can't encode hi-half src0 yet + instructions = [ + s_mov_b32(s[0], 0x42003c00), # hi=3.0, lo=1.0 + v_mov_b32_e32(v[0], s[0]), + s_mov_b32(s[0], 0x47004200), # hi=7.0, lo=3.0 + v_mov_b32_e32(v[1], s[0]), + cmp, + ] + st = run_program(instructions, n_lanes=1) + self.assertEqual(st.vcc & 1, 1, "Expected vcc=1 (v0.hi 3.0 == v1.lo 3.0)") + def test_v_cmp_nge_f16_inf_self(self): """v_cmp_nge_f16 comparing -inf with itself (unordered less than). diff --git a/test/backend/test_arange.py b/test/backend/test_arange.py index 16e7d22e29..d58e86e5f2 100644 --- a/test/backend/test_arange.py +++ b/test/backend/test_arange.py @@ -4,13 +4,13 @@ from tinygrad import Tensor, GlobalCounters, dtypes, nn, Device, Variable from tinygrad.helpers import Context, getenv, DEV from tinygrad.engine.realize import run_linear, estimate_uop, compile_linear from tinygrad.renderer.ptx import PTXRenderer -from test.helpers import needs_second_gpu, check_schedule, assert_kernel_count +from test.helpers import needs_second_gpu, check_schedule, assert_kernel_count, KernelCountException class TestArange(unittest.TestCase): def _get_flops(self, tensor, desired): GlobalCounters.reset() linear = compile_linear(tensor.schedule_linear()) - self.assertEqual(len(linear.src), 1) + if len(linear.src) != 1: raise KernelCountException(1, len(linear.src)) run_linear(linear) np.testing.assert_equal(tensor.numpy(), desired) return estimate_uop(linear.src[-1]).ops @@ -253,7 +253,7 @@ class TestIndexing(unittest.TestCase): xq_rope, _ = apply_rotary_emb(xq, xq, freqs_cis) xq_rope.sum().backward() linear = compile_linear(wq.grad.schedule_linear()) - assert len(linear.src) == 1, f"expected one kernel for backward, got: {len(linear.src)}" + if len(linear.src) != 1: raise KernelCountException(1, len(linear.src)) bwd_ops = estimate_uop(linear.src[0]).ops expected_ops = bs*seqlen*dim*dim*ops_scale print(f"rope matmul bwd ({dtype}): {GlobalCounters.kernel_count} kernels, {bwd_ops:,} ops") diff --git a/test/backend/test_asm_gemm.py b/test/backend/test_asm_gemm.py index 53c148c862..e18458f808 100644 --- a/test/backend/test_asm_gemm.py +++ b/test/backend/test_asm_gemm.py @@ -1,7 +1,7 @@ import unittest from tinygrad import Tensor, Device, dtypes, Context from tinygrad.helpers import getenv, system, DEV -from extra.gemm.cdna_asm_gemm import asm_gemm, hk_bf16_atb_gemm, quantize_mxfp4 +from extra.gemm.cdna_asm_gemm import asm_gemm, hk_bf16_atb_gemm from test.helpers import needs_second_gpu from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8, FP8_MAX @@ -157,13 +157,20 @@ class TestMXFP4(unittest.TestCase): def test_quantize(self): import numpy as np - block = np.array([0, .26, .74, .75, 1.26, 1.75, 2.51, 3.5, 5.1, 6, -6] + [0] * 21, dtype=np.float32) - x = Tensor(np.tile(block, (32, 8)), dtype=dtypes.bfloat16) - packed, scale, _ = quantize_mxfp4(x) - p = packed.numpy() - codes = np.stack((p & 0xF, p >> 4), axis=-1).reshape(32, 256) - np.testing.assert_array_equal(codes[0, :11], [0, 1, 1, 2, 3, 4, 5, 6, 7, 7, 15]) - np.testing.assert_array_equal(scale.numpy(), np.full((32, 8), 127, dtype=np.uint8)) + from extra.llama_kernels.quantize_mxfp4 import quantize_mxfp4 + rng = np.random.default_rng(0) + x = np.triu(rng.standard_normal((256, 256), dtype=np.float32)) + x += np.triu(x, 1).T + x[:32, :32] = 0 + row, row_scale, col, col_scale = quantize_mxfp4(Tensor(x, dtype=dtypes.bfloat16)) + Tensor.realize(row, row_scale, col, col_scale) + row, row_scale = row.numpy(), row_scale.numpy() + col, col_scale = col.numpy(), col_scale.numpy() + np.testing.assert_array_equal(row, col) + np.testing.assert_array_equal(row_scale, col_scale) + self.assertTrue(row.any()) + self.assertTrue((row_scale == 127).any()) + self.assertTrue((row_scale != 127).any()) def test_correctness(self): import numpy as np @@ -171,17 +178,9 @@ class TestMXFP4(unittest.TestCase): rng = np.random.default_rng(1) a = Tensor(rng.standard_normal((M, K), dtype=np.float32), dtype=dtypes.bfloat16) b = Tensor(rng.standard_normal((N, K), dtype=np.float32), dtype=dtypes.bfloat16) - out = asm_gemm(a, b.T, mxfp4=True).realize() - # reference gemm - a_packed, scale_a, _ = quantize_mxfp4(a) - b_packed, scale_b, _ = quantize_mxfp4(b) - def unpack(x): return np.stack((x & 0xF, x >> 4), axis=-1).reshape(x.shape[0], -1) - code_a, code_b = unpack(a_packed.numpy()), unpack(b_packed.numpy()) - lut = np.array([0, .5, 1, 1.5, 2, 3, 4, 6, -0., -.5, -1, -1.5, -2, -3, -4, -6], dtype=np.float32) - a_dequant = lut[code_a] * np.repeat(np.exp2(scale_a.numpy().astype(np.int16)-127), 32, axis=1) - b_dequant = lut[code_b] * np.repeat(np.exp2(scale_b.numpy().astype(np.int16)-127), 32, axis=1) - ref = Tensor(a_dequant @ b_dequant.T, dtype=dtypes.bfloat16).realize().numpy() - np.testing.assert_array_equal(out.numpy(), ref) + out = asm_gemm(a, b.T, mxfp4=True).realize().numpy().astype(np.float32) + ref = a.numpy().astype(np.float32) @ b.numpy().astype(np.float32).T + self.assertLess(np.linalg.norm(out-ref) / np.linalg.norm(ref), 0.2) def test_empty(self): M, N, K = getenv("M", 16384), getenv("N", 4096), getenv("K", 14336) diff --git a/test/backend/test_custom_kernel.py b/test/backend/test_custom_kernel.py index 68e866a35f..f92a253a7d 100644 --- a/test/backend/test_custom_kernel.py +++ b/test/backend/test_custom_kernel.py @@ -190,6 +190,12 @@ class TestCustomKernel(unittest.TestCase): b = Tensor.custom_kernel(tst, a, fxn=custom_sum)[0] self.assertEqual(b.item(), 15) + def test_sum_outside(self): + a = Tensor([1.0, 2, 3, 4, 5])+1 + tst = Tensor.empty(1) + b = Tensor.custom_kernel(tst, a, fxn=custom_sum)[0] + self.assertEqual(b.item(), 20) + def test_sum_int(self): a = Tensor([1, 2, 3, 4, 5]) tst = Tensor.empty(1, dtype=a.dtype) @@ -287,7 +293,7 @@ class TestCustomKernel(unittest.TestCase): GlobalCounters.reset() c.realize() assert all(i == 3. for i in c.flatten().tolist()), f"all 3 {c.tolist()}" - assert_kernel_count(3) + assert_kernel_count(2) def test_multi_after_schedule_order(self): """Test correct scheduling order when custom_kernel has multiple outputs. @@ -405,10 +411,8 @@ class TestCustomKernel(unittest.TestCase): assert_kernel_count(2) self.assertEqual(z.tolist(), x.add(2).tolist()) - @unittest.expectedFailure def test_custom_kernel_sched_copy(self): self.test_custom_kernel_sched(use_custom=True) - @unittest.expectedFailure def test_sliced_buffer_function(self): x = Tensor.arange(32).reshape(8, 4).clone().realize() from tinygrad import function @@ -419,7 +423,8 @@ class TestCustomKernel(unittest.TestCase): GlobalCounters.reset() y = run(x[0]).realize() # it's copying the input and the output - assert_kernel_count(1) + # TODO: subbuffer usage has runtime specific behavior, this will be fixed after the removal of SLICE. + assert_kernel_count(2 if y.device in ("CL", "WEBGPU") else 1) self.assertEqual(y.tolist(), [1, 2, 3, 4]) @Context(DEV="CPU") @@ -429,12 +434,28 @@ class TestCustomKernel(unittest.TestCase): # TODO: it currently requires a compiler for Ops.BINARY from tinygrad.device import Device binary = Device[a.device].renderer.compiler.compile(src) - def custom_src_kernel(A:UOp) -> UOp: + def custom_src_kernel(A:UOp, B:UOp) -> UOp: sink = UOp.sink(A, arg=KernelInfo(name="test_src")) return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(sink.toposort())), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=binary))) - a = Tensor.custom_kernel(a.reshape(2, 2).T, fxn=custom_src_kernel)[0] - self.assertEqual(a.tolist(), [[1, 2], [1, 3]]) + a = Tensor.custom_kernel(a.reshape(2, 2).clone(), a.reshape(2, 2).T, fxn=custom_src_kernel)[0] + self.assertEqual(a.tolist(), [[1, 1], [2, 3]]) + @Context(DEV="CPU") + def test_simple_from_source_alt(self): + a = Tensor.arange(4).clone().realize() + src = "void copy(int* restrict out, int* restrict in) { for (int i = 0; i < 4; i++) out[i] = in[i]; }" + from tinygrad.device import Device + binary = Device[a.device].renderer.compiler.compile(src) + def custom_src_kernel(out:UOp, inp:UOp) -> UOp: + sink = UOp.sink(out, inp, arg=KernelInfo(name="copy")) + return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(sink.toposort())), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=binary))) + out = Tensor.custom_kernel(Tensor.empty_like(a), a+1, fxn=custom_src_kernel)[0] + GlobalCounters.reset() + out.realize() + assert_kernel_count(2) + self.assertEqual(out.tolist(), [1, 2, 3, 4]) + + @unittest.skip("this shouldn't be expected to work") def test_inplace_transpose(self): def custom_assign_row_max_kernel(A:UOp) -> UOp: row = UOp.range(A.shape[0], 0) @@ -471,8 +492,8 @@ class TestCustomKernelInput(unittest.TestCase): def test_reshape(self): self._test_mop(lambda x: x.reshape(16, 2), max_kernels=2) def test_permute(self): self._test_mop(lambda x: x.reshape(4, 8).T, max_kernels=3) - def test_double_permute(self): self._test_mop(lambda x: x.reshape(4, 8).T.T, max_kernels=3) - def test_shrink(self): self._test_mop(lambda x: x[:4], max_kernels=2) + def test_double_permute(self): self._test_mop(lambda x: x.reshape(4, 8).T.T, max_kernels=2) + def test_shrink(self): self._test_mop(lambda x: x[:4], max_kernels=1) def test_pad(self): self._test_mop(lambda x: x[:4].pad(((0, 4),)), max_kernels=2) def test_flip(self): self._test_mop(lambda x: x.flip(0), max_kernels=2) def test_offset_shrink(self): self._test_mop(lambda x: x[4:8], max_kernels=2) diff --git a/test/backend/test_dtype.py b/test/backend/test_dtype.py index dc122ffcb4..594ebd0e23 100644 --- a/test/backend/test_dtype.py +++ b/test/backend/test_dtype.py @@ -169,6 +169,13 @@ class TestFp8sConversions(unittest.TestCase): def test_fp8e5m2fnuz_to_float(self, x): np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e5m2fnuz), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e5m2fnuz).float().item()) + def test_fp8e5m2fnuz_to_float_smallest_normals(self): + # fnuz bias exceeds half's, so exp-1 normals land below half's normal range: they flush to zero like denormals + if dtypes.half not in supported_dtypes or dtypes.half in EMULATED_DTYPES.tolist(dtypes) or dtypes.fp8e5m2fnuz in supported_dtypes: + self.skipTest("needs the emulated fp8 with a native half intermediate") + vals = Tensor([0x04, 0x05, 0x06, 0x07], dtype=dtypes.uint8).bitcast(dtypes.fp8e5m2fnuz).float().numpy() + np.testing.assert_equal(vals, [0., 0., 0., 0.]) + class TestBFloat16DType(unittest.TestCase): def test_bf16_to_float(self): _test_cast(Tensor([100000], dtype=dtypes.bfloat16), dtypes.float32) diff --git a/test/backend/test_dtype_alu.py b/test/backend/test_dtype_alu.py index 0e4b60faeb..6a182616c8 100644 --- a/test/backend/test_dtype_alu.py +++ b/test/backend/test_dtype_alu.py @@ -399,9 +399,10 @@ class TestDTypeALU(unittest.TestCase): if float_dtype not in supported_dtypes: float_dtype = dtypes.float32 universal_test_cast(a, float_dtype, unsigned_dtype) - @unittest.expectedFailure - def test_unsafe_cast_float_to_int_failure(self): - val = float(dtypes.int32.max - 1) + def test_unsafe_cast_float_to_int(self): + # the value is off the float32 grid but rounds in-range: the buffer and const-fold paths must agree + # (out-of-range float->int cast stays undefined: hardware may saturate where the fold wraps) + val = 2147483000.0 t1 = Tensor([val], dtype=dtypes.float32).cast(dtypes.int32) t2 = Tensor(val, dtype=dtypes.float32).cast(dtypes.int32) np.testing.assert_equal(t1.item(), t2.item()) diff --git a/test/backend/test_linearizer.py b/test/backend/test_linearizer.py index 958a9be840..5a43d68011 100644 --- a/test/backend/test_linearizer.py +++ b/test/backend/test_linearizer.py @@ -437,7 +437,7 @@ def reset_bufs(bufs:list[Buffer]): for buf in bufs: buf.copy_from(Buffer("PYTHON", buf.size, buf.dtype, opaque=memoryview(bytearray(buf.nbytes)))) def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[], - apply_tc=False, atol=1e-4, rtol=1e-4, color_sizes=[], wanna_output=[]): + apply_tc=False, atol=1e-4, rtol=1e-4, color_sizes=[], wanna_output=[], check_default_opt=True): outbufs = real_bufs[:len(realized_ast.src)] wanna_output = [np.array(x).flatten() for x in wanna_output] buf_uops = [UOp.new_buffer(b.device, b.size, b.dtype) for b in real_bufs] @@ -459,9 +459,7 @@ def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[] for buf,want in zip(copyout_outputs(outbufs), wanna_output): np.testing.assert_allclose(buf, want, atol=atol, rtol=rtol) # Check correctness of handcoded optimiztions. - reset_bufs(outbufs) - run_prg(opts=None) - for buf,want in zip(copyout_outputs(outbufs), wanna_output): np.testing.assert_allclose(buf, want, atol=atol, rtol=rtol) + if check_default_opt: check_opt(None) for x in opts: # Check custom transformations if any. check_opt(([Opt(OptOps.TC, 0, (TC_SELECT.value, TC_OPT.value, 1))] if apply_tc else [])+x) diff --git a/test/backend/test_llama_kernels.py b/test/backend/test_llama_kernels.py index 4542d66a6f..f03a57b393 100644 --- a/test/backend/test_llama_kernels.py +++ b/test/backend/test_llama_kernels.py @@ -5,6 +5,7 @@ from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8 from extra.llama_kernels.fused_ce import fused_ce_loss from extra.llama_kernels import local_abs_max from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed, quantize_fp8_scalar +from extra.llama_kernels.swiglu import swiglu from extra.models.llama import apply_rotary_emb, precompute_freqs_cis from extra.thunder.amd.fa import custom_fused_qkv_rope_backward, fused_qkv_rope from test.helpers import needs_second_gpu, assert_kernel_count @@ -161,5 +162,31 @@ class TestFusedQKVRoPE(unittest.TestCase): ref = Tensor.cat(dq_ref, dk_ref, dv_ref, dim=3).reshape(*dx.shape).realize() with Context(DEBUG=0): self.assertTrue(dx.allclose(ref, atol=2e-2, rtol=2e-2).item(), "backward mismatch") +def run_swiglu(test:unittest.TestCase, shape:tuple[int, ...]) -> None: + Tensor.manual_seed(0) + x = (Tensor.randn(*shape) * 2).cast(dtypes.bfloat16).realize() + hidden = x.shape[-1] // 2 + out, ref = swiglu(x), x[..., :hidden].silu() * x[..., hidden:] + Tensor.realize(out, ref) + with Context(DEBUG=0): test.assertTrue(out.allclose(ref, atol=2.5e-1, rtol=3e-2).item(), "SwiGLU forward mismatch") + + grad = (Tensor.randn(*out.shape) * 2).cast(dtypes.bfloat16).realize() + grad_x, grad_ref = out.gradient(x, gradient=grad)[0], ref.gradient(x, gradient=grad)[0] + Tensor.realize(grad_x, grad_ref) + test.assertEqual(grad_x.shape, shape) + test.assertEqual(grad_x.dtype, dtypes.bfloat16) + with Context(DEBUG=0): test.assertTrue(grad_x.allclose(grad_ref, atol=2.5e-1, rtol=3e-2).item(), "SwiGLU backward mismatch") + +class TestSwiGLU(unittest.TestCase): + def setUp(self): + if dtypes.bfloat16 not in Device[Device.DEFAULT].renderer.supported_dtypes(): self.skipTest("need bfloat16") + + def test_simple(self): run_swiglu(self, (2, 32, 64)) + + def test_llama_shape(self): + if Device.DEFAULT != "AMD" or not Device[Device.DEFAULT].renderer.target.arch.startswith("gfx950"): + self.skipTest("only run on real machine for speed") + run_swiglu(self, (2, 8192, 28672)) + if __name__ == '__main__': unittest.main() diff --git a/test/backend/test_ops.py b/test/backend/test_ops.py index ae2c81ee15..9eacc5b524 100644 --- a/test/backend/test_ops.py +++ b/test/backend/test_ops.py @@ -1535,6 +1535,8 @@ class TestOps(unittest.TestCase): def test_prod(self): helper_test_op(None, lambda x: x.prod(), vals=[[1.0, 2.0, 3.0]]) + helper_test_op(None, lambda x: x.prod(), vals=[[0.0, 2.0, 3.0]]) + helper_test_op(None, lambda x: x.prod(), vals=[[0.0, 0.0, 3.0]]) with Context(NOOPT=1): helper_test_op(None, lambda x: x.prod(), vals=[[1.0, 2.0, 3.0]]) helper_test_op([(3,4,5,6)], lambda x: x.prod(dim=3), lambda x: x.prod(axis=3)) helper_test_op([(3,4,5,6)], lambda x: x.prod(dim=1), lambda x: x.prod(axis=1)) diff --git a/test/backend/test_symbolic_ops.py b/test/backend/test_symbolic_ops.py index 8e7020dc2c..344bd54493 100644 --- a/test/backend/test_symbolic_ops.py +++ b/test/backend/test_symbolic_ops.py @@ -6,6 +6,15 @@ from examples.gpt2 import Attention import numpy as np class TestSymbolicOps(unittest.TestCase): + def test_negative_slice(self): + a = Tensor.rand(3, 10, 4) + for i in range(3, 10): + vi = Variable("i", 1, 10).bind(i) + # negative int bounds against a symbolic dim must resolve against the size, like slice.indices + np.testing.assert_allclose(a[:, :vi][:, -3:-1].numpy(), a[:, :i][:, -3:-1].numpy(), atol=1e-6, rtol=1e-6) + np.testing.assert_allclose(a[:, :vi][:, -1:].numpy(), a[:, :i][:, -1:].numpy(), atol=1e-6, rtol=1e-6) + np.testing.assert_allclose(a[:, :vi][:, -1].numpy(), a[:, :i][:, -1].numpy(), atol=1e-6, rtol=1e-6) + def test_plus1(self): def f(a): return (a+1).realize() a = Tensor.rand(3, 10) diff --git a/test/device/test_hcq2.py b/test/device/test_hcq2.py new file mode 100644 index 0000000000..19d72d7bb8 --- /dev/null +++ b/test/device/test_hcq2.py @@ -0,0 +1,14 @@ +import unittest, numpy as np +from unittest.mock import patch +from tinygrad import Device, Tensor +from tinygrad.helpers import getenv +from tinygrad.runtime.support.hcq2 import HCQ_DEVS, all_devices_in + +@unittest.skipUnless(getenv("HCQ2") and all_devices_in(Device.DEFAULT, HCQ_DEVS), "hcq2 device required") +class TestHCQ2(unittest.TestCase): + def test_copy_without_copy_queue(self): + with patch.object(Device[Device.DEFAULT], "has_copy_queue", False): + np.testing.assert_equal(Tensor(np.arange(61, dtype=np.float32)).to(Device.DEFAULT).contiguous().realize().numpy(), np.arange(61)) + +if __name__ == "__main__": + unittest.main() diff --git a/test/mockgpu/amd/emu.py b/test/mockgpu/amd/emu.py index 46925b876b..3f36e331e9 100644 --- a/test/mockgpu/amd/emu.py +++ b/test/mockgpu/amd/emu.py @@ -109,7 +109,7 @@ def _init_sqtt_encoder(): _SMEM = (ir3.SMEM, ir4.SMEM, irc.SMEM) _VALU = (ir3.VOP1, ir3.VOP2, ir3.VOP3, ir3.VOP3P, ir3.VOPC, ir3.VOPD, ir3.VOP3SD, ir3.VOP3_SDST, ir3.VOP1_SDST, ir4.VOP1, ir4.VOP2, ir4.VOP3, ir4.VOP3P, ir4.VOPC, ir4.VOPD, ir4.VOP3SD, ir4.VOP3_SDST, ir4.VOP1_SDST, - irc.VOP1, irc.VOP2, irc.VOP3, irc.VOP3P, irc.VOPC, irc.VOP3SD, irc.VOP3_SDST) + irc.VOP1, irc.VOP2, irc.VOP3, irc.VOP3P, irc.VOP3PX2, irc.VOPC, irc.VOP3SD, irc.VOP3_SDST) _DS = (ir3.DS, ir4.DS, irc.DS) _GLOBAL = (ir3.GLOBAL, ir4.VGLOBAL, irc.GLOBAL) _FLAT = (ir3.FLAT, ir4.VFLAT, irc.FLAT) @@ -1150,6 +1150,9 @@ def _compile_vopc(inst: ir3.VOPC|ir3.VOPC_DPP16|ir3.VOP3|ir4.VOPC|ir4.VOPC_DPP16 def get_cmp_bit(lane) -> UOp: lc = lane.cast(dtypes.int) if isinstance(lane, UOp) else _c(lane, 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) + s0 = src0_hi.where(_hi16(ctx.rvgpr_dyn(src0_hi.where(src0_off - _c(384), _c(0)), lc)), s0) 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) @@ -1323,7 +1326,7 @@ def _compile_vop3sd(inst: ir3.VOP3SD | ir4.VOP3SD | irc.VOP3SD, ctx: _Ctx) -> UO else: return ctx.compile_vop_pcode(inst.op, srcs, lane, vdst_reg, exec_mask, sdst_reg=inst.sdst.offset) -def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp: +def _compile_mfma(inst: irc.VOP3P|irc.VOP3PX2, ctx: _Ctx) -> UOp: """CDNA MFMA matrix multiply-accumulate emulation. Uses local temp arrays to cache inputs, avoiding aliasing issues when vdst overlaps src0/src1. @@ -1349,6 +1352,25 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp: src0_is_vgpr = src0_off >= _c(256) src1_is_vgpr = src1_off >= _c(256) + scaled = isinstance(inst, irc.VOP3PX2) + if scaled: + assert isinstance(inst, irc.VOP3PX2) + # F8F6F4 input formats: 0=FP8(E4M3), 1=BF8(E5M2). FP6/FP4 (2-4) not emulated. + src0_fmt, src1_fmt = int(inst.cbsz), int(inst.blgp) + if src0_fmt > 1 or src1_fmt > 1: raise RuntimeError(f"unsupported scaled MFMA formats cbsz={src0_fmt} blgp={src1_fmt}") + # scale_src0/scale_src1 are source operands pointing at 32-bit registers holding 4 packed E8M0 scale exponents. + # The 2-bit opsel/opsel_hi select which byte applies to A/B for this instruction. + scale0_off = ctx.inst_field(type(inst).scale_src0) + scale1_off = ctx.inst_field(type(inst).scale_src1) + sel0, sel1 = int(inst.opsel) & 3, int(inst.opsel_hi) & 3 + def _scale_exp(off: UOp, sel: int, lane: UOp) -> UOp: + sv = ctx.rsrc_dyn(off, lane, 32) + byte = (sv >> UOp.const(sel * 8, dtypes.uint32)) & UOp.const(0xFF, dtypes.uint32) + return byte.cast(dtypes.int32) - UOp.const(127, dtypes.int32) + # combined A*B scale for this lane: 2^(ea-127) * 2^(eb-127) + def scale_factor(lane: UOp) -> UOp: + return UOp.exp2((_scale_exp(scale0_off, sel0, lane) + _scale_exp(scale1_off, sel1, lane)).cast(dtypes.float32)) + m = _re.search(r'(\d+)X(\d+)X(\d+)', op_name) if m is None: raise ValueError(f"could not parse MFMA dimensions from {op_name}") M, N, K = int(m.group(1)), int(m.group(2)), int(m.group(3)) @@ -1404,7 +1426,18 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp: # The optimizer folds bitcast(uint32→float32) stores to float32 arrays, losing the conversion. tmp = UOp.placeholder((n_a_elems + n_b_elems,), dtypes.uint32, slot=0, addrspace=AddrSpace.LOCAL) - def cvt_elem(raw: UOp, sub_idx: int) -> UOp: + # Per-operand fp8 format ("fp8"=E4M3, "bf8"=E5M2) for A and B + if 'F8F6F4' in op_name: + assert isinstance(inst, (irc.VOP3P_MFMA, irc.VOP3PX2)) + _fmts = {0: "fp8", 1: "bf8"} + a_fmt, b_fmt = _fmts.get(int(inst.cbsz), "fp8"), _fmts.get(int(inst.blgp), "fp8") + elif is_fp8: + # A/B formats from name suffix, e.g. V_MFMA_F32_16X16X32_BF8_FP8 + suffixes = op_name.rsplit('_', 2)[-2:] + a_fmt, b_fmt = ("bf8" if sfx == "BF8" else "fp8" for sfx in suffixes) + else: a_fmt = b_fmt = "fp8" + + def cvt_elem(raw: UOp, sub_idx: int, fp8_fmt: str = "fp8") -> UOp: if is_i8: # Extract i8, sign-extend to i32 byte_val = (raw >> UOp.const(sub_idx * 8, dtypes.uint32)) & UOp.const(0xFF, dtypes.uint32) @@ -1412,7 +1445,7 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp: elif is_f32_src: return raw # already uint32 (f32 bit pattern) elif is_fp8: - return ((raw >> UOp.const(sub_idx * 8, dtypes.uint32)) & UOp.const(0xFF, dtypes.uint32)).cast(dtypes.uint32) + return _FUNCS[f"{fp8_fmt}_to_f32"](raw >> UOp.const(sub_idx * 8, dtypes.uint32)).bitcast(dtypes.uint32) elif is_bf16: # bf16→f32 bits: just shift left by 16 (bf16 is upper 16 bits of f32) return ((raw >> UOp.const(sub_idx * 16, dtypes.uint32)) & UOp.const(0xFFFF, dtypes.uint32)) << UOp.const(16, dtypes.uint32) @@ -1454,7 +1487,7 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp: # Read A/B sources. Use rsrc_dyn for inline constants/SGPRs (src_off < 256), rvgpr_dyn for VGPRs (src_off >= 256). a_raw = src0_is_vgpr.where(ctx.rvgpr_dyn(src0_r + _c(reg_idx), read_lane), ctx.rsrc_dyn(src0_off, _c(0, dtypes.int), 32)) - a_val = cvt_elem(a_raw, sub_idx) + a_val = cvt_elem(a_raw, sub_idx, a_fmt) if M == 4: a_idx = grp_idx * UOp.const(M * K, dtypes.int) + mn_idx * UOp.const(K, dtypes.int) + UOp.const(kl, dtypes.int) else: @@ -1463,7 +1496,7 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp: b_raw = src1_is_vgpr.where(ctx.rvgpr_dyn(src1_r + _c(reg_idx), read_lane), ctx.rsrc_dyn(src1_off, _c(0, dtypes.int), 32)) - b_val = cvt_elem(b_raw, sub_idx) + b_val = cvt_elem(b_raw, sub_idx, b_fmt) if M == 4: b_idx = b_off + grp_idx * UOp.const(N * K, dtypes.int) + mn_idx * UOp.const(K, dtypes.int) + UOp.const(kl, dtypes.int) else: @@ -1480,6 +1513,17 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp: # Actually: 16 ACCVGPRs per lane, organized as 4 groups (l//32 gives half, each half has 2 sub-groups) of 4 rows tmp2 = tmp.after(read_phase) + 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]. For scaled MFMA, only the dot product 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) + if not scaled: + for k in range(K): acc = acc + prod(k) + return acc + dot = prod(0) + for k in range(1, K): dot = dot + prod(k) + return acc + dot * scale_factor(lane) + compute_lane = ctx.range() compute_stores = [] @@ -1510,10 +1554,7 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp: else: acc_v = acc_v.bitcast(dtypes.float32) acc = src2_is_vgpr.where(acc_v, acc_scalar) - for k in range(K): - a_val = tmp2.index(m_base * UOp.const(K, dtypes.int) + UOp.const(k, dtypes.int)).bitcast(acc_dt) - b_val = tmp2.index(b_off + n_idx * UOp.const(K, dtypes.int) + UOp.const(k, dtypes.int)).bitcast(acc_dt) - acc = acc + a_val * b_val + acc = _dot_accum(acc, m_base * UOp.const(K, dtypes.int), b_off + n_idx * UOp.const(K, dtypes.int), compute_lane) if is_int_out: compute_stores.append((ctx.waccvgpr_dyn if use_acc else ctx.wvgpr_dyn)( @@ -1535,17 +1576,13 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp: if M == 4: # 4x4: each group is independent. A/B indexed per-group. m_base = c_grp * UOp.const(M * K, dtypes.int) + UOp.const(out_reg * K, dtypes.int) - for k in range(K): - a_val = tmp2.index(m_base + UOp.const(k, dtypes.int)).bitcast(acc_dt) - b_val = tmp2.index(b_off + c_grp * UOp.const(N*K, dtypes.int) + n_idx * UOp.const(K, dtypes.int)+UOp.const(k, dtypes.int)).bitcast(acc_dt) - acc = acc + a_val * b_val + b_base = b_off + c_grp * UOp.const(N * K, dtypes.int) + n_idx * UOp.const(K, dtypes.int) else: # 16x16: K is split across groups. Shared MxK/NxK arrays. m_base = c_grp * UOp.const(out_per_lane, dtypes.int) + UOp.const(out_reg, dtypes.int) - for k in range(K): - a_val = tmp2.index(m_base * UOp.const(K, dtypes.int) + UOp.const(k, dtypes.int)).bitcast(acc_dt) - b_val = tmp2.index(b_off + n_idx * UOp.const(K, dtypes.int) + UOp.const(k, dtypes.int)).bitcast(acc_dt) - acc = acc + a_val * b_val + b_base = b_off + n_idx * UOp.const(K, dtypes.int) + + acc = _dot_accum(acc, m_base if M == 4 else m_base * UOp.const(K, dtypes.int), b_base, compute_lane) if is_int_out: compute_stores.append((ctx.waccvgpr_dyn if use_acc else ctx.wvgpr_dyn)( @@ -1563,33 +1600,41 @@ def _compile_wmma(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P, ctx: _Ctx) -> UOp: vdst_reg = ctx.inst_field(type(inst).vdst) src0_r = ctx.inst_field(type(inst).src0) - _c(256) src1_r = ctx.inst_field(type(inst).src1) - _c(256) - src2_r = ctx.inst_field(type(inst).src2) - _c(256) - is_f16_output = 'F16_16X16X16_F16' in op_name or 'BF16_16X16X16_BF16' in op_name # F16/BF16 output vs F32 output + src2_r = ctx.inst_field(type(inst).src2) + src2_r = (src2_r >= 256).where(src2_r - _c(256), src2_r) + output_type = op_name.split("WMMA_", 1)[1].split("_", 1)[0] is_bf16 = 'BF16' in op_name cvt = _FUNCS['bf16_to_f32'] if is_bf16 else _FUNCS['f16_to_f32'] is_rdna4 = isinstance(inst, ir4.VOP3P) - # read 16x16 F16/BF16 matrix from VGPRs → flat f32 array[row*16+k] - def read_f16_val(src, lane, vgpr, half): + sz = 8 if "8" in op_name else 16 + # read matrix from VGPRs → flat f32/i32 array[row*16+k] + def gval(src, lane, vgpr, ridx): v = ctx.rvgpr_dyn(src + _c(vgpr), UOp.const(lane, dtypes.int)) - return cvt((v >> UOp.const(16, dtypes.uint32)) if half else (v & UOp.const(0xFFFF, dtypes.uint32))) + pkd = v >> UOp.const(ridx * sz, dtypes.uint32) if ridx > 0 else v + pkd = pkd & UOp.const((1 << sz) - 1, dtypes.uint32) + if "F" in output_type: return cvt(pkd) + return (pkd << _c(24, dtypes.uint)).bitcast(dtypes.int32) >> _c(24, dtypes.int32) # sign extend - # RDNA3: 16 lanes × 8 VGPRs × 2 halves, k maps linearly - # RDNA4: 32 lanes × 4 VGPRs × 2 halves, k bits are scrambled (k[2] goes to lane bit 4) - def read_f16_mat(src): - # (row, k) → (lane, vgpr, half) + # RDNA3 f16/bf16: 16 lanes × 8 VGPRs × 2 halves, k maps linearly + # RDNA3 iu8: 16 lanes × 4 VGPRs × 4 quarters, k maps linearly + # RDNA4: 32 lanes x 4 VGPRS x 2 halves, k bits are scrambled (k[2] goes to lane bit 4) + def read_mat(src): + n = 32 // sz # values per vgpr + # (row, k) → (lane, vgpr, row index) def ab_map(i, k): elem, lane = ((k & 3) | ((k >> 1) & 4), i + ((k >> 2) & 1) * 16) if is_rdna4 else (k, i) - return lane, elem // 2, elem % 2 - return [read_f16_val(src, *ab_map(row, k)) for row in range(16) for k in range(16)] - mat_a, mat_b = read_f16_mat(src0_r), read_f16_mat(src1_r) + return lane, elem // n, elem % n + return [gval(src, *ab_map(row, k)) for row in range(16) for k in range(16)] + + mat_a, mat_b = read_mat(src0_r), read_mat(src1_r) # (row, col) -> (lane, vgpr) def d_map(m, n): lane_bit, vgpr = (m >> 3, m & 7) if is_rdna4 else (m & 1, m >> 1) return n + lane_bit * 16, vgpr - if is_f16_output: + if output_type in ["F16", "BF16"]: # read accumulator C with f16 layout: for RDNA4, pairs of f32 vgprs pack into one f16 vgpr # for RDNA3, same layout as f32 but only lo 16 bits used - mat_c = [read_f16_val(src2_r, *((lane, vgpr // 2, vgpr % 2) if is_rdna4 else (lane, vgpr, 0))) + mat_c = [gval(src2_r, *((lane, vgpr // 2, vgpr % 2) if is_rdna4 else (lane, vgpr, 0))) for m in range(16) for n in range(16) for lane, vgpr in [d_map(m, n)]] mat_d = [sum(mat_a[r*16+k] * mat_b[c*16+k] for k in range(16)) + mat_c[r*16+c] for r in range(16) for c in range(16)] def f32_to_f16_bits(v: UOp) -> UOp: return v.cast(dtypes.half).bitcast(dtypes.uint16).cast(dtypes.uint32) @@ -1602,18 +1647,22 @@ def _compile_wmma(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P, ctx: _Ctx) -> UOp: else: # (rdna3) 1 f16 per VGPR (lo half only) stores = [ctx.wvgpr_dyn(vdst_reg + _c(d_map(m, n)[1]), UOp.const(d_map(m, n)[0], dtypes.int), out_cvt(mat_d[m*16+n]), exec_mask) for m in range(16) for n in range(16)] - else: # f32 - mat_c = [ctx.rvgpr_dyn(src2_r + _c(d_map(m, n)[1]), UOp.const(d_map(m, n)[0], dtypes.int)).bitcast(dtypes.float32) + else: # f32/i32 + out_dt = dtypes.float32 if output_type == "F32" else dtypes.int32 + mat_c = [ctx.rvgpr_dyn(src2_r + _c(d_map(m, n)[1]), UOp.const(d_map(m, n)[0], dtypes.int)).bitcast(out_dt) for m in range(16) for n in range(16)] mat_d = [sum(mat_a[r*16+k] * mat_b[c*16+k] for k in range(16)) + mat_c[r*16+c] for r in range(16) for c in range(16)] stores = [ctx.wvgpr_dyn(vdst_reg + _c(d_map(m, n)[1]), UOp.const(d_map(m, n)[0], dtypes.int), mat_d[m*16+n].bitcast(dtypes.uint32), exec_mask) for m in range(16) for n in range(16)] return UOp.sink(*stores, *ctx.inc_pc()) -def _compile_vop3p(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P, ctx: _Ctx) -> UOp: +def _compile_vop3p(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P | irc.VOP3PX2, ctx: _Ctx) -> UOp: op_name = _op_name(inst) - if 'WMMA' in op_name and ('16X16X16_F16' in op_name or '16X16X16_BF16' in op_name): return _compile_wmma(inst, ctx) - if 'MFMA' in op_name and any(f'{s}X{s}X' in op_name for s in ('4', '16', '32')) and isinstance(inst, irc.VOP3P): return _compile_mfma(inst, ctx) + if 'WMMA' in op_name: + assert not isinstance(inst, irc.VOP3PX2) + return _compile_wmma(inst, ctx) + if 'MFMA' in op_name and any(f'{s}X{s}X' in op_name for s in ('4', '16', '32')) and isinstance(inst, (irc.VOP3P, irc.VOP3PX2)): + return _compile_mfma(inst, ctx) # ACCVGPR_WRITE/READ/MOV: copies between VGPR and ACCVGPR register files # Detect by checking operand types for ACCVGPR involvement @@ -2044,7 +2093,7 @@ _INST_HANDLERS: dict[type, Callable[..., UOp]] = { irc.SOPP: _compile_sopp, irc.SMEM: _compile_smem, irc.SOP1: _compile_sop, irc.SOP2: _compile_sop, irc.SOPC: _compile_sop, irc.SOPK: _compile_sop, irc.VOP1: _compile_vop12, irc.VOP1_DPP16: _compile_vop12, irc.VOP2: _compile_vop12, irc.VOP2_DPP16: _compile_vop12, irc.VOPC: _compile_vopc, irc.VOP3: _compile_vop3, - irc.VOP3_SDST: _compile_vop3, irc.VOP3SD: _compile_vop3sd, irc.VOP3P: _compile_vop3p, + irc.VOP3_SDST: _compile_vop3, irc.VOP3SD: _compile_vop3sd, irc.VOP3P: _compile_vop3p, irc.VOP3PX2: _compile_vop3p, irc.VOP1_SDWA: _compile_sdwa, irc.VOP2_SDWA: _compile_sdwa, irc.VOP2_SDWA_SDST: _compile_sdwa, irc.VOPC_SDWA_SDST: _compile_sdwa, irc.DS: _compile_mem_op, irc.FLAT: _compile_mem_op, irc.GLOBAL: _compile_mem_op, irc.SCRATCH: _compile_mem_op, irc.MUBUF: _compile_mubuf, diff --git a/test/null/test_const_folding.py b/test/null/test_const_folding.py index 0f217e8fa5..5e365c4233 100644 --- a/test/null/test_const_folding.py +++ b/test/null/test_const_folding.py @@ -1,6 +1,6 @@ import unittest, itertools, math from tinygrad import Tensor, dtypes, Context -from tinygrad.dtype import DType, ConstType +from tinygrad.dtype import DType, ConstType, truncate from tinygrad.uop.ops import Ops, UOp from test.helpers import full_rewrite import numpy as np @@ -51,6 +51,17 @@ class TestWeakConstFolding(unittest.TestCase): def test_invalid_poison(self): self.assertTrue(UOp.invalid().alu(Ops.CDIV, UOp.const(0)).simplify().is_invalid) + def test_cast_commits_to_dtype_grid(self): + # committing a weak const to a stated width puts the value on that width's grid, same as storage packing and native compilers + v = 1/123008 # not representable in float16 + out = UOp.const(v).cast(dtypes.half).simplify() + self.assertEqual((out.op, out.dtype, out.val), (Ops.CONST, dtypes.half, truncate[dtypes.half](v))) + self.assertNotEqual(out.val, v) + # the grid commit preserves the sign of zero + self.assertEqual(math.copysign(1, UOp.const(-0.0).cast(dtypes.half).simplify().val), -1) + # observable at tensor level: the const-folded comparison agrees with the committed value + self.assertTrue((Tensor(-3.2).cast(dtypes.float32) <= truncate[dtypes.float32](-3.2)).item()) + class TestBinaryOpsConstFolding(unittest.TestCase): def test_add_literal_zero(self): _check_ast_count(0, Tensor([1.0, 2, 3, 4]) + 0) diff --git a/test/null/test_graph_rewrite.py b/test/null/test_graph_rewrite.py index a10d13ea8b..9194cc01bd 100644 --- a/test/null/test_graph_rewrite.py +++ b/test/null/test_graph_rewrite.py @@ -208,7 +208,7 @@ class TestGEPAndVectorizeRewrite(unittest.TestCase): import inspect -from tinygrad.uop.ops import graph_rewrite, _substitute, track_rewrites +from tinygrad.uop.ops import graph_rewrite, _substitute, rewrite_group from tinygrad.uop.symbolic import symbolic_simple class TestBottomUpRewrite(unittest.TestCase): @@ -220,7 +220,7 @@ class TestBottomUpRewrite(unittest.TestCase): self.assertIs(gt, ret) # normally .substitute would be fine, but it's not tracked -@track_rewrites() +@rewrite_group() def named_substitute(name:str, uop:UOp, rel:dict[UOp, UOp]): return graph_rewrite(uop, _substitute, rel, bottom_up=True) def substitute(uop:UOp, rel:dict[UOp, UOp]): return named_substitute(inspect.stack()[1].function, uop, rel) diff --git a/test/null/test_schedule.py b/test/null/test_schedule.py index f0a5b5c4d3..614e2cf837 100644 --- a/test/null/test_schedule.py +++ b/test/null/test_schedule.py @@ -6,7 +6,7 @@ from tinygrad.uop.ops import UOp, Ops, GroupOp, UPat, KernelInfo, AxisType from tinygrad.helpers import GlobalCounters, Context from tinygrad.engine.realize import run_linear, compile_linear from tinygrad.codegen import to_program, full_rewrite_to_sink -from test.helpers import check_schedule, assert_kernel_count +from test.helpers import check_schedule, assert_kernel_count, KernelCountException def _realize_weights(m): for p in nn.state.get_parameters(m): p.realize() @@ -592,9 +592,7 @@ class TestSchedule(unittest.TestCase): img = Tensor.randn(BS, CIN, 64, 64).realize() w = Tensor.uniform(16, CIN, 3, 3).realize() ret = Tensor.conv2d(img, w).relu().mean().backward() - linear, var_vals = Tensor.linear_with_vars(ret, img.grad, w.grad) - cnt = len([call for call in linear.src if call.src[0].op is Ops.SINK]) - assert cnt == allowed, f"expected {allowed} kernels, got {cnt}" + check_schedule([ret, img.grad, w.grad], allowed) def test_conv2d_half(self): self.test_conv2d(4, dtype=dtypes.half) @@ -615,7 +613,8 @@ class TestSchedule(unittest.TestCase): return len([call for call in linear.src if call.src[0].op is Ops.PROGRAM]) with Context(IMAGE=1): - self.assertEqual(cnt(), 5) + got = cnt() + if got != 5: raise KernelCountException(5, got) def test_image_f16_residual_fusion(self): with Context(FLOAT16=1, OPENPILOT_HACKS=1): @@ -630,7 +629,8 @@ class TestSchedule(unittest.TestCase): return len([call for call in linear.src if call.src[0].op is Ops.PROGRAM]) with Context(IMAGE=1): - self.assertEqual(cnt(), 9) + got = cnt() + if got != 9: raise KernelCountException(9, got) def _test_fusion(self, shapes, f, cnt): with Context(DEBUG=0, TRACK_MATCH_STATS=0): @@ -858,6 +858,65 @@ class TestSchedule(unittest.TestCase): x = Tensor.rand(32) check_schedule(x, 1, [Tensor._device_rng_counters[x.device]]) + # **** custom kernel realize tests + + @staticmethod + def _copy_fxn(name:str="copy"): + def copy_kernel(out:UOp, inp:UOp) -> UOp: + i = UOp.range(inp.numel(), 0) + return UOp.group(out[i].store(inp[i])).end(i).sink(arg=KernelInfo(name=name)) + return copy_kernel + + def _copy_call(self, out:Tensor, expr:Tensor, name:str="copy") -> Tensor: + # forge a custom kernel call with params and call args, like llm/kernels does (no Tensor.custom_kernel contiguous) + params = tuple(UOp.placeholder_like(u, slot=i) for i,u in enumerate((out.uop, expr.uop))) + return Tensor(out.uop.after(self._copy_fxn(name)(*params).call(out.uop, expr.uop))) + + def test_custom_kernel_buffer_src(self): + # custom kernels need buffers: a buffer input must never add a realize kernel + y = Tensor.ones(64).contiguous().realize() + out = Tensor.empty_like(y) + check_schedule(self._copy_call(out, y), 1) + + def test_custom_kernel_view_src(self): + # a RESHAPE over a buffer resolves to the buffer state (RESHAPEs on call args are stripped), no realize kernel + y = Tensor.ones(64).contiguous().realize() + out = Tensor.empty_like(y) + check_schedule(self._copy_call(out, y.reshape(8, 8).reshape(64)), 1) + + def test_custom_kernel_elementwise_src(self): + # a computed input is not a buffer state: the call args are unwrapped to their base buffer, + # so the compute would be silently dropped. this must raise instead of producing wrong results + y = Tensor.ones(64).contiguous().realize() + out = Tensor.empty_like(y) + check_schedule(self._copy_call(out, y + y), 2) + + def test_custom_kernel_lazy_const_src(self): + # a lazy const expression above the call has no buffer at all. this used to crash rangeify with a KeyError + x = Tensor.linspace(-1.0, 1.0, 64) + out = Tensor.empty_like(x) + check_schedule(self._copy_call(out, x), 2) + + def test_custom_kernel_offset_view_src(self): + # a SHRINK with an offset over a buffer is not a buffer state either, the offset would be silently dropped + y = Tensor.ones(128).contiguous().realize() + out = Tensor.empty(64) + check_schedule(self._copy_call(out, y[16:80]), 2) + + def test_custom_kernel_computed_src_api(self): + # the supported way to pass computed inputs: Tensor.custom_kernel makes inputs contiguous (one realize kernel) + y = Tensor.ones(64).contiguous().realize() + out = Tensor.empty_like(y) + check_schedule(Tensor.custom_kernel(out, y + y, fxn=self._copy_fxn())[0], 2) + + def test_custom_kernel_on_custom_kernel(self): + # the output of a custom kernel is a buffer state, chaining custom kernels must not add kernels + y = Tensor.ones(64).contiguous().realize() + k1 = self._copy_call(Tensor.empty_like(y), y, name="k1") + k2 = self._copy_call(Tensor.empty_like(y), k1, name="k2") + sched, _ = check_schedule(k2, 2) + self.assertEqual([call.src[0].arg.name for call in sched.src], ["k1", "k2"]) + def test_empty_is_not_realized(self): a = Tensor.empty(10) child = a+2 diff --git a/test/null/test_simplify_valid_idx.py b/test/null/test_simplify_valid_idx.py index 15a9ce9171..754d713fe2 100644 --- a/test/null/test_simplify_valid_idx.py +++ b/test/null/test_simplify_valid_idx.py @@ -2,7 +2,8 @@ import unittest, itertools from tinygrad.codegen.late.coalesce import indexing_simplify from tinygrad.dtype import dtypes -from tinygrad.uop.ops import UOp, Ops, graph_rewrite, pm_lower_index_dtype +from tinygrad.uop.ops import UOp, Ops, graph_rewrite +from tinygrad.uop.weak import pm_lower_index_dtype from tinygrad.uop.symbolic import simplify_valid, sym, pm_move_where_on_load from tinygrad.helpers import Context from test.helpers import full_rewrite @@ -332,7 +333,7 @@ class TestImageSimplification(unittest.TestCase): load = get_load_image_uop(shape, valid, idx) self.check(load, - "((((idx2*2)+r0)<11)&((((idx1*8)+r1)<3)!=True))", + "(((idx2*2)+r0)<11)", "(idx0+(idx1*512+r1*64)+-192)", "((((idx2*2)+r0)+(((idx1+((r1+5)//8))+1)//2))+-4)") @@ -460,7 +461,7 @@ class TestImageSimplification(unittest.TestCase): self.check(load, None, "(gidx0+lidx0*1024+r0*1024+lidx1*128+-3168)", "0") except AssertionError: # TODO: fold valid - self.check(load, "(((lidx1<1)!=True)&(((lidx0+r0)<3)!=True)&((lidx0+r0)<19))", + self.check(load, "(((lidx1<1)!=True)&((lidx0+r0)<19))", "(gidx0+lidx1*128+(lidx0*1024+r0*1024)+-3168)", "0") def test_simplify10(self): @@ -479,7 +480,7 @@ class TestImageSimplification(unittest.TestCase): self.check(load, None, "(lidx2+gidx0*4+lidx0*1024+r0*1024+lidx1*256+-3264)", "0") except AssertionError: # TODO: fold valid - self.check(load, "(((lidx1<1)!=True)&(((lidx0+r0)<3)!=True)&((lidx0+r0)<11))", + self.check(load, "(((lidx1<1)!=True)&((lidx0+r0)<11))", "(lidx2+gidx0*4+lidx1*256+(lidx0*1024+r0*1024)+-3264)", "0") def test_drop_non_monotonic_window(self): diff --git a/test/null/test_uop_graph.py b/test/null/test_uop_graph.py index e889df8611..b7ff3fbebb 100644 --- a/test/null/test_uop_graph.py +++ b/test/null/test_uop_graph.py @@ -1,10 +1,11 @@ import unittest, pytest -from tinygrad import dtypes, Variable +from tinygrad import dtypes, Variable, Device from tinygrad.dtype import AddrSpace from tinygrad.helpers import DEBUG, Context -from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, graph_rewrite, GroupOp, AxisType, broadcast_axes +from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, graph_rewrite, GroupOp, AxisType, broadcast_axes, KernelInfo from tinygrad.uop.symbolic import sym from test.helpers import to_uops_list +from tinygrad.codegen import full_rewrite_to_sink simple_pm = PatternMatcher([ (UPat.cvar('x', dtypes.weakint), lambda x: UOp.const(1.0) + UOp.const(2.0)), @@ -536,6 +537,15 @@ class TestReduceCollapse(unittest.TestCase): # Should become add of two separate reduces self.assertEqual(result.op, Ops.ADD) + def test_reduce_shapeless_const_unroll(self): + """a REDUCE over a shapeless CONST (e.g. x*0 folded late in codegen) must collapse before the expander""" + out = UOp.param(0, dtypes.float, (1,)) + red = UOp.const(3.0).cast(dtypes.float).reduce(UOp.range(4, 0, AxisType.UNROLL), arg=(Ops.ADD, 0)) + ast = UOp.sink(out.index(UOp.const(0)).store(red)).replace(arg=KernelInfo()) + uops = full_rewrite_to_sink(ast, Device["CPU"].renderer, optimize=False).toposort() + self.assertNotIn(Ops.REDUCE, [u.op for u in uops]) + self.assertIn(12.0, [u.val for u in uops if u.op is Ops.CONST]) + class TestMovementOps(unittest.TestCase): def test_pm_mops_partial_reshape_index_removes_reshape(self): from tinygrad.schedule.rangeify import pm_mops diff --git a/test/null/test_uop_symbolic.py b/test/null/test_uop_symbolic.py index 36629e025a..575fe9a55a 100644 --- a/test/null/test_uop_symbolic.py +++ b/test/null/test_uop_symbolic.py @@ -6,7 +6,7 @@ from tinygrad.dtype import dtypes, ConstType, DType, Invalid from test.helpers import get_uops from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer from tinygrad.uop.spec import spec_shared, type_verify -from tinygrad.uop.symbolic import sym, commutative, pm_simplify_valid, pm_move_where_on_load +from tinygrad.uop.symbolic import sym, pm_fold_cast_const, commutative, pm_simplify_valid, pm_move_where_on_load from tinygrad.uop.validate import uops_to_z3 def check_uop_against_string(self, v:UOp, s:str): @@ -35,7 +35,7 @@ class TestSymbolic(unittest.TestCase): self.assertEqual(solver.check(expr1 != expr2), z3.unsat, "simplified expression not equal to original") def helper_test_variable(self, v, n, m, s, test_z3:bool=True): - v_simplified = graph_rewrite(v, sym, name="simplify symbolic uop") + v_simplified = graph_rewrite(v, sym+pm_fold_cast_const, name="simplify symbolic uop") if test_z3: self.check_equal_z3(v, v_simplified) nmin, nmax = v_simplified.vmin, v_simplified.vmax check_uop_against_string(self, v_simplified, s) diff --git a/test/null/test_uop_vmin_vmax.py b/test/null/test_uop_vmin_vmax.py index fd5c3a7ab2..2dc3f27c35 100644 --- a/test/null/test_uop_vmin_vmax.py +++ b/test/null/test_uop_vmin_vmax.py @@ -1,6 +1,6 @@ import unittest, math from tinygrad.uop.ops import UOp, Ops -from tinygrad.dtype import dtypes, Invalid +from tinygrad.dtype import dtypes, Invalid, truncate class TestVminVmaxProperties(unittest.TestCase): def test_vmin_vmax_constant(self): @@ -168,6 +168,10 @@ class TestVminVmaxProperties(unittest.TestCase): x = UOp.const(4.5).cast(dtypes.float) self.assertIs(x.ne(x.cast(dtypes.int).cast(dtypes.float)).simplify().arg, True) + def test_vmin_vmax_cast_int_to_float_grid(self): + # a cast to float only takes values on the float grid, so its bounds are the source bounds rounded at the destination + self.assertEqual(UOp.variable('x', 0, 16777219, dtypes.int).cast(dtypes.float)._min_max, (0.0, 16777220.0)) + def test_vmin_vmax_invalid(self): i = UOp.invalid() self.assertNotEqual(i.vmin, i.vmax) @@ -317,8 +321,8 @@ class TestVminVmaxVConst(unittest.TestCase): def test_vmin_vmax_vconst_with_floats(self): # vmin and vmax for a vector constant of float values uop = UOp.const((1.5, -3.2, 0.0)) - self.assertEqual(uop.vmin, -3.2) - self.assertEqual(uop.vmax, 1.5) + self.assertEqual(uop.vmin, truncate[dtypes.default_float](-3.2)) + self.assertEqual(uop.vmax, truncate[dtypes.default_float](1.5)) def test_vmin_vmax_vconst_with_bools(self): # vmin and vmax for a vector constant of bool values diff --git a/test/null/test_uops.py b/test/null/test_uops.py index 247c611250..c1d6931fd6 100644 --- a/test/null/test_uops.py +++ b/test/null/test_uops.py @@ -5,7 +5,8 @@ from tinygrad.tensor import Tensor from tinygrad.helpers import Timing, Context, cdiv from tinygrad.dtype import dtypes, AddrSpace, ConstFloat, Invalid # noqa: F401 from tinygrad.device import Device -from tinygrad.uop.ops import Ops, ParamArg, PatternMatcher, UOp, UPat, dtype_from_uop, exec_alu, graph_rewrite, pm_lower_index_dtype # noqa: F401 # ParamArg used by eval(str(uop)) roundtrip tests +from tinygrad.uop.ops import Ops, ParamArg, PatternMatcher, UOp, UPat, dtype_from_uop, exec_alu, graph_rewrite # noqa: F401 # ParamArg used by eval(str(uop)) roundtrip tests +from tinygrad.uop.weak import pm_lower_index_dtype from tinygrad.uop.spec import spec_program, spec_shared, type_verify from tinygrad.uop.symbolic import sym, pm_remove_invalid from test.helpers import eval_uop, to_uops_list @@ -55,11 +56,11 @@ class TestDTypeFromUOp(unittest.TestCase): invalid = UOp.invalid() self.assertIs(invalid.dtype, dtypes.bool) self.assertIs(UOp.const(Invalid, dtypes.float32), invalid) - self.assertIs((moved:=invalid.reshape((1,))).cast(dtypes.float32), moved) scratch = Tensor.invalids(4, dtype=dtypes.float32) self.assertEqual((scratch.dtype, next(u.dtype for u in scratch.uop.toposort() if u.op is Ops.BUFFER), next(u.dtype for u in scratch.uop.toposort() if u.is_invalid)), (dtypes.float32, dtypes.float32, dtypes.bool)) invalid, value = UOp.invalid(), UOp.const(1, dtypes.float32) + for u in (UOp.param(0, dtypes.bool, ()).where(value, invalid), value+invalid, UOp.stack(value, invalid)): self.assertIs(u.src[-1], invalid) for u in (UOp(Ops.STACK, dtypes.float32, src=(value, invalid)), UOp(Ops.ADD, dtypes.float32, src=(value, invalid)), UOp.const(True).where(value, invalid), UOp(Ops.CMPLT, src=(invalid, value)), UOp(Ops.CMPLT, src=(value, invalid)), UOp.param(0, dtypes.float32, (4,)).index(invalid)): type_verify(u, spec_shared) diff --git a/test/null/test_upat_compile.py b/test/null/test_upat_compile.py index b59f5f03c6..eeea86efad 100644 --- a/test/null/test_upat_compile.py +++ b/test/null/test_upat_compile.py @@ -1,11 +1,11 @@ import unittest from tinygrad.helpers import DEBUG, Context from tinygrad.dtype import dtypes -from tinygrad.uop.ops import UPat, track_rewrites, GroupOp, Ops +from tinygrad.uop.ops import UPat, rewrite_group, GroupOp, Ops from tinygrad.uop.upat import _get_code, upat_compile import dis -@track_rewrites() +@rewrite_group() def do_compile(up): print("\n***** COMPILE", up) match_code = _get_code(up, False) diff --git a/test/null/test_viz.py b/test/null/test_viz.py index 954a4f2252..c6cdb6213c 100644 --- a/test/null/test_viz.py +++ b/test/null/test_viz.py @@ -3,7 +3,7 @@ from pathlib import Path from dataclasses import dataclass from typing import Generator -from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, TrackedPatternMatcher, graph_rewrite, track_rewrites, profile_matches +from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, TrackedPatternMatcher, graph_rewrite, rewrite_group from tinygrad.uop.symbolic import sym from tinygrad.dtype import dtypes, AddrSpace from tinygrad.helpers import colored, ansistrip, flatten, TracingKey, ProfileRangeEvent, ProfileEvent, Context, cpu_events, profile_marker @@ -14,7 +14,7 @@ from tinygrad.uop.ops import tracked_keys, tracked_ctxs, uop_fields, active_rewr from tinygrad.viz.serve import load_rewrites, get_full_rewrite, uop_to_json, VizData, get_render, addrspace_colors from tinygrad.codegen import do_to_program -@track_rewrites(name=True) +@rewrite_group(name=True) def exec_rewrite(sink:UOp, pm_lst:list[PatternMatcher], names:None|list[str]=None) -> UOp: for i,pm in enumerate(pm_lst): sink = graph_rewrite(sink, TrackedPatternMatcher(pm.patterns), name=names[i] if names else None) @@ -109,7 +109,7 @@ class TestViz(unittest.TestCase): def test_default_name(self): with save_viz() as viz: a = UOp.variable("a", 1, 10) - @track_rewrites() + @rewrite_group() def name_default(): return graph_rewrite(a, PatternMatcher([])) name_default() lst = viz.list_items() @@ -118,7 +118,7 @@ class TestViz(unittest.TestCase): # name can also come from a function that returns a string def test_dyn_name_fxn(self): with save_viz() as viz: - @track_rewrites(name=lambda *args,ret,**kwargs: ret.render()) + @rewrite_group(name=lambda *args,ret,**kwargs: ret.render()) def name_from_fxn(s:UOp, arg:list|None=None): return graph_rewrite(s, PatternMatcher([])) name_from_fxn(UOp.variable("a", 1, 10)+1, arg=["test"]) lst = viz.list_items() @@ -128,18 +128,18 @@ class TestViz(unittest.TestCase): # name can also come from a function that returns a TracingKey def test_tracing_key(self): with save_viz() as viz: - @track_rewrites(name=lambda inp,ret: TracingKey("custom_name", (inp,))) + @rewrite_group(name=lambda inp,ret: TracingKey("custom_name", (inp,))) def test(s:UOp): return graph_rewrite(s, PatternMatcher([])) test(UOp.variable("a", 1, 10)+1) lst = viz.list_items() # NOTE: names from TracingKey do not get deduped self.assertEqual(lst[0]["name"], "custom_name") - def test_nested_track_rewrites(self): + def test_nested_rewrite_group(self): with save_viz() as viz: - @track_rewrites(name=lambda x,ret: TracingKey(f"inner fxn for {x.render()}", (ret,))) + @rewrite_group(name=lambda x,ret: TracingKey(f"inner fxn for {x.render()}", (ret,))) def inner(x:UOp): return graph_rewrite(x, PatternMatcher([]), name="each") - @track_rewrites(name=lambda *args,ret: f"outer rewrite of {len(args)} inputs") + @rewrite_group(name=lambda *args,ret: f"outer rewrite of {len(args)} inputs") def outer(*xs:tuple[UOp, ...]): return graph_rewrite(UOp.sink(*[inner(x) for x in xs]), PatternMatcher([]), name="all") items = ["a", "b", "c"] outer(*[UOp.variable(x, 1, 10) for x in items]) @@ -156,13 +156,13 @@ class TestViz(unittest.TestCase): self.assertEqual(len(steps), 1) self.assertEqual(steps[0]["name"], "each") - def test_profile_matches(self): + def test_rewrite_group_nested(self): with save_viz() as viz: - @profile_matches + @rewrite_group(new_ctx=False) def nested_function(u:UOp): for i in range(2): graph_rewrite(u, PatternMatcher([]), name=f"step {i+1}") - @track_rewrites() + @rewrite_group() def main_rewrite(u:UOp): graph_rewrite(u, PatternMatcher([]), name="init") nested_function(u) @@ -173,9 +173,9 @@ class TestViz(unittest.TestCase): self.assertEqual(steps[1]["name"], "nested_function") self.assertEqual(len(steps), 4) - def test_profile_matches_invalid_arg(self): + def test_rewrite_group_invalid_arg(self): with save_viz(): - @profile_matches + @rewrite_group(new_ctx=False) def invalid_fxn(arg:str): return graph_rewrite(UOp(Ops.SINK), PatternMatcher([])) with self.assertRaisesRegex(AssertionError, "invalid match tracing input"): invalid_fxn("test") @@ -395,7 +395,7 @@ class TestVizIntegration(unittest.TestCase): graph = next(viz.get_details(0, 0))["graph"] self.assertEqual(len([n for n in graph.values() if repr(metadata) in n["label"]]), 1) - # tracing also works without a track_rewrites context + # tracing also works without a rewrite_group context # all graph_rewrites get put into the default group def test_default_tracing(self): with save_viz() as viz: @@ -407,11 +407,11 @@ class TestVizIntegration(unittest.TestCase): self.assertEqual(len(ls), 1) self.assertEqual(ls[0]["name"], "default graph_rewrite") - # using @track_rewrites organizes function calls into groups + # using @rewrite_group organizes function calls into groups # and nicely counts function calls. def test_group_traces(self): with save_viz() as viz: - @track_rewrites() + @rewrite_group() def test(root): return graph_rewrite(root, sym) test(c:=UOp.const(1)) @@ -420,11 +420,11 @@ class TestVizIntegration(unittest.TestCase): self.assertEqual(len(ls), 2) for i in range(2): self.assertEqual(ls[i]["name"], f"test n{i+1}") - # @track_rewrites always starts a new group. + # @rewrite_group always starts a new group. def test_group_combined(self): with save_viz() as viz: def default_test(root): return graph_rewrite(root, sym) - tracked_test = track_rewrites()(default_test) + tracked_test = rewrite_group()(default_test) c = UOp.const(1) default_test(c+1) # goes to the default group tracked_test(c) # all rewrites after this go inside the second group. diff --git a/test/opt/test_tensor_cores.py b/test/opt/test_tensor_cores.py index 81eb7c6a42..0d9b5db3d1 100644 --- a/test/opt/test_tensor_cores.py +++ b/test/opt/test_tensor_cores.py @@ -79,7 +79,8 @@ class TestTensorCores(unittest.TestCase): @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores") def test_tensor_cores(self): for tc in Device[Device.DEFAULT].renderer.tensor_cores: - helper_tc_allclose(tc.dims[0], tc.dims[1], tc.dims[2], tc.dtype_in, tc.dtype_out, axis=0, tc_opt=0) + with self.subTest(tc=tc): + helper_tc_allclose(tc.dims[0], tc.dims[1], tc.dims[2], tc.dtype_in, tc.dtype_out, axis=0, tc_opt=0) @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores") def test_tensor_cores_nested_reduce(self): @@ -101,7 +102,8 @@ class TestTensorCores(unittest.TestCase): if Device.DEFAULT == "CPU" and DEV.renderer == "LLVM": assert "0x201000" in prg.src[2].arg elif Device.DEFAULT == "AMD" and DEV.renderer == "LLVM": - assert "@llvm.amdgcn.wmma" in prg.src[2].arg + # RDNA emits wmma intrinsics, CDNA emits mfma intrinsics + assert ("@llvm.amdgcn.wmma" in prg.src[2].arg) or ("@llvm.amdgcn.mfma" in prg.src[2].arg) elif Device[Device.DEFAULT].renderer.suffix == "PTX": assert "mma.sync.aligned" in prg.src[2].arg else: @@ -181,11 +183,13 @@ class TestTensorCores(unittest.TestCase): @unittest.skipIf(Device.DEFAULT == "PYTHON", "slow on EMULATED device") @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores") def test_tensor_cores_unroll_phi(self): - tc = Device[Device.DEFAULT].renderer.tensor_cores[0] - x, y = Tensor.rand(128, 128, dtype=tc.dtype_in), Tensor.rand(128, 128, dtype=tc.dtype_in) + # skip fp8 tcs: the unoptimized ALU baseline quantizes products to fp8 (JAX promotion), which legitimately + # differs from the MFMA path (f32 accumulation), so the baseline-vs-TC numerical gate can't hold for fp8. + tc = next(tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in not in dtypes.fp8s) + x, y = Tensor.rand(16, 64, dtype=tc.dtype_in), Tensor.rand(64, 16, dtype=tc.dtype_in) r = x.matmul(y, dtype=tc.dtype_out) - opts = [Opt(OptOps.UNROLL, 0, 4)] - ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3) + opts = [Opt(OptOps.UNROLL, 0, 2)] + ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3, check_default_opt=False) for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[1].src): if u.op is Ops.WMMA: assert u.src[-1].src[0].op != Ops.STORE @@ -195,11 +199,11 @@ class TestTensorCores(unittest.TestCase): @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores") @unittest.skipIf(Device.DEFAULT in {"CPU"}, "CPU does not support using a different type for accumulation") def test_tensor_cores_unroll_casted_phi(self): - tc = [tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in != tc.dtype_out][0] - x, y = Tensor.rand(128, 128, dtype=tc.dtype_in), Tensor.rand(128, 128, dtype=tc.dtype_in) + tc = [tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in != tc.dtype_out and tc.dtype_in not in dtypes.fp8s][0] + x, y = Tensor.rand(16, 64, dtype=tc.dtype_in), Tensor.rand(64, 16, dtype=tc.dtype_in) r = x.matmul(y, dtype=tc.dtype_out) - opts = [Opt(OptOps.UNROLL, 0, 4)] - ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3) + opts = [Opt(OptOps.UNROLL, 0, 2)] + ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3, check_default_opt=False) for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[1].src): if u.op is Ops.WMMA: #assert u.src[-1].dtype == dtypes.float.vec(prod(tc.thread_local_sizes[2])) @@ -211,11 +215,11 @@ class TestTensorCores(unittest.TestCase): @unittest.skipIf(Device.DEFAULT in {"CPU"}, "CPU does not support using a different type for accumulation") def test_tensor_cores_unroll_casted_phi_with_children(self): # all STORE children are outside the loop - tc = [tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in != tc.dtype_out][0] - x, y = Tensor.rand(128, 128, dtype=tc.dtype_in), Tensor.rand(128, 128, dtype=tc.dtype_in) + tc = [tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in != tc.dtype_out and tc.dtype_in not in dtypes.fp8s][0] + x, y = Tensor.rand(16, 64, dtype=tc.dtype_in), Tensor.rand(64, 16, dtype=tc.dtype_in) r = x.matmul(y, dtype=tc.dtype_out).relu() - opts = [Opt(OptOps.UNROLL, 0, 4)] - ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3) + opts = [Opt(OptOps.UNROLL, 0, 2)] + ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3, check_default_opt=False) for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[1].src): if u.op is Ops.WMMA: #assert u.src[-1].dtype == dtypes.float.vec(prod(tc.thread_local_sizes[2])) diff --git a/test/testextra/test_hevc.py b/test/testextra/test_hevc.py index 058b237f42..174813b8a9 100644 --- a/test/testextra/test_hevc.py +++ b/test/testextra/test_hevc.py @@ -1,7 +1,9 @@ import unittest -from tinygrad import Tensor, Device, dtypes -from tinygrad.helpers import fetch, round_up +from tinygrad import Tensor, Device, Variable, dtypes +from tinygrad.helpers import DEV, fetch, round_up +from tinygrad.engine.realize import compile_linear +from tinygrad.uop.ops import Ops from extra.hevc.hevc import parse_hevc_file_headers, nv_gpu from extra.hevc.decode import hevc_decode @@ -63,7 +65,7 @@ class TestHevc(unittest.TestCase): self.assertEqual(list(frame3.initreflistidxl1), [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) self.assertEqual(list(frame3.RefDiffPicOrderCnts), [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) - @unittest.skipUnless(Device.DEFAULT == "NV", "NV only") + @unittest.skipUnless(Device.DEFAULT == "NV" and not DEV.interface.startswith("MOCK"), "real NV only") def test_hevc_decode(self): url = "https://github.com/haraschax/filedump/raw/09a497959f7fa6fd8dba501a25f2cdb3a41ecb12/comma_video.hevc" dat = fetch(url, headers={"Range": f"bytes=0-{512<<10}"}).read_bytes() @@ -83,5 +85,22 @@ class TestHevc(unittest.TestCase): self.assertEqual(f.dtype, dtypes.uint8) self.assertEqual(f.device, "NV") + @unittest.skipUnless(Device.DEFAULT == "NV", "NV only") + def test_hevc_decode_compile(self): + url = "https://github.com/haraschax/filedump/raw/09a497959f7fa6fd8dba501a25f2cdb3a41ecb12/comma_video.hevc" + dat = fetch(url, headers={"Range": f"bytes=0-{512<<10}"}).read_bytes() + + opaque, frame_info, _, _, luma_w, luma_h, _ = parse_hevc_file_headers(dat) + offset, sz, frame_pos, max_hist, _ = frame_info[1] + out_image_size = luma_h + (luma_h + 1) // 2, round_up(luma_w, 64) + history = [Tensor.empty(*out_image_size, dtype=dtypes.uint8, device="NV") for _ in range(max_hist)] + decoded = Tensor(dat, device="NV")[offset:offset+sz].decode_hevc_frame( + Variable("pos", 0, max_hist + 1).bind(frame_pos), out_image_size, opaque[1], history) + + compiled = compile_linear(decoded.linear_with_vars()[0]) + self.assertTrue(any(call.src[0].op is Ops.PROGRAM for call in compiled.src)) + encdec_calls = [call for call in compiled.src if call.src[0].op is Ops.CUSTOM_FUNCTION and call.src[0].arg == "encdec"] + self.assertEqual(len(encdec_calls), 1) + if __name__ == "__main__": unittest.main() diff --git a/test/unit/test_call.py b/test/unit/test_call.py index 02a1be038e..12e55f12d7 100644 --- a/test/unit/test_call.py +++ b/test/unit/test_call.py @@ -212,6 +212,18 @@ class TestCallSchedule(unittest.TestCase): out = f(a, v.bind(5)) np.testing.assert_allclose(out.numpy(), [5., 10., 15.]) + def test_precompile_scoped_bind_arg(self): + @function(precompile=True) + def f(x:Tensor, scale:UOp) -> Tensor: return x * scale + a = Tensor.ones(3) + x = f(a, UOp.variable("scale_a", 1, 100).bind(2)) + y = f(a, UOp.variable("scale_b", 1, 100).bind(3)) + fx = next(u for u in x.uop.toposort() if u.op is Ops.FUNCTION) + fy = next(u for u in y.uop.toposort() if u.op is Ops.FUNCTION) + self.assertEqual(fx.src[0].key, fy.src[0].key) + np.testing.assert_equal(x.numpy(), [2, 2, 2]) + np.testing.assert_equal(y.numpy(), [3, 3, 3]) + def test_precompile_schedule_cache_hit(self): """two instances of the same @function should produce identical function body keys (schedule cache hit)""" @function(precompile=True) @@ -347,5 +359,15 @@ class TestCallMultiSharded(unittest.TestCase): np.testing.assert_allclose(a.grad.numpy(), b.numpy(), rtol=1e-5) np.testing.assert_allclose(b.grad.numpy(), a.numpy(), rtol=1e-5) + def test_symbolic_reshape_shard_axis(self): + toks = UOp.variable("toks", 1, 2).bind(2) + devs = ("CPU:0", "CPU:1") + x = Tensor(np.arange(16, dtype=np.float32).reshape(1, 2, 8)).shard(devs, axis=2).realize() + @function + def f(x:Tensor) -> Tensor: return x.reshape(1, x.shape[1], 2, 4) + out = f(x[:, :toks]).realize() + self.assertEqual(out.uop.axis, 2) + np.testing.assert_equal(out[:1, :2].to(devs[0]).numpy(), np.arange(16, dtype=np.float32).reshape(1, 2, 2, 4)) + if __name__ == '__main__': unittest.main() diff --git a/test/unit/test_dtype_spec.py b/test/unit/test_dtype_spec.py index cbb3bc5ba6..2a07b73488 100644 --- a/test/unit/test_dtype_spec.py +++ b/test/unit/test_dtype_spec.py @@ -222,6 +222,12 @@ class TestAutoCastType(unittest.TestCase): t.square().mean().backward() np.testing.assert_allclose(t.grad.numpy().flatten(), [60000 * 2 / (N*N)] * N*N) + @unittest.skipUnless(dtypes.half in supported_dtypes, "need half") + def test_var_half_precision_large_n(self): + # the element count (70000) exceeds half max (65504): the denominator must not be materialized in half + t = Tensor([[0.0, 1.0]], dtype=dtypes.half).expand(35000, 2).contiguous() + np.testing.assert_allclose(t.var().numpy(), 0.25, rtol=1e-3) + @unittest.skipIf(Device.DEFAULT == "WEBGPU", "Precision error") @unittest.skipUnless(dtypes.half in supported_dtypes, "need half") def test_softmax_dtype(self): diff --git a/test/unit/test_dtype_weak.py b/test/unit/test_dtype_weak.py index 7178aa5e6f..a4e9740eb9 100644 --- a/test/unit/test_dtype_weak.py +++ b/test/unit/test_dtype_weak.py @@ -3,7 +3,8 @@ import tempfile, unittest, math from tinygrad import Tensor, dtypes, TinyJit from tinygrad.helpers import Context from tinygrad.dtype import least_upper_float -from tinygrad.uop.ops import UOp, Ops, dtype_from_uop, graph_rewrite, pm_lower_index_dtype, pm_commit_weak +from tinygrad.uop.ops import UOp, Ops, dtype_from_uop, graph_rewrite +from tinygrad.uop.weak import pm_lower_index_dtype, pm_commit_weak from tinygrad.uop.symbolic import symbolic_simple from tinygrad.uop.spec import spec_shared, type_verify from tinygrad.engine.jit import JitError @@ -62,6 +63,31 @@ class TestWeakPromotion(unittest.TestCase): self.assertEqual((x._uop.base.op, x._uop.base.val, x.dtype, x.shape, y.dtype), (Ops.CONST, 1, dtypes.weakfloat, (1,), dtypes.float32)) + def test_weak_expression_anchors_at_strong_lub(self): + # regression test for the HALF bert nan (#17408, reverted in #17409): lub(int32, weakfloat)==weakfloat makes + # `loss_mask.sum() + 1e-5` a weakfloat EXPRESSION. Meeting a strong float in a binop must pin it at the lub + denom = (Tensor.zeros(912, dtype=dtypes.int32) != Tensor.zeros(912, dtype=dtypes.float32)).sum() + 1e-5 + self.assertIs(denom.dtype, dtypes.weakfloat) # the setup: the denominator expression itself is weak + x, y = Tensor([2048.0], dtype=dtypes.float32)._broadcasted(denom) + self.assertIs(y.dtype, dtypes.float32) + recips = [u for u in (x / y)._uop.toposort() if u.op is Ops.RECIPROCAL] + self.assertEqual([(u.dtype, u.src[0].dtype) for u in recips], [(dtypes.float32, dtypes.float32)]) + with Context(DEFAULT_FLOAT=dtypes.float16): + committed = graph_rewrite((UOp.const(1).cast(dtypes.int32) + UOp.const(1.0)).cast(dtypes.float32), pm_lower_index_dtype, ctx={}) + self.assertEqual([u.dtype for u in committed.toposort() if u.op is Ops.ADD], [dtypes.float32]) + + def test_cast_weak_expression_commits_at_cast_floor(self): + # the floor never narrows: a cast BELOW the default does not pull the compute width down with it + with Context(DEFAULT_FLOAT=dtypes.float32): + narrowed = graph_rewrite((UOp.const(1.0) + UOp.const(2.0)).cast(dtypes.float16), pm_lower_index_dtype, ctx={}) + self.assertEqual((narrowed.dtype, narrowed.src[0].dtype), (dtypes.float16, dtypes.float32)) + + def test_cast_weak_expression_value_uses_cast_floor(self): + with Context(DEFAULT_FLOAT=dtypes.float16): + denom = Tensor.ones(1, dtype=dtypes.int32, device="CPU").sum() * 70000 + 1e-5 + out = Tensor(1.0, dtype=dtypes.float32, device="CPU") / denom + self.assertAlmostEqual(out.item(), 1 / (70000 + 1e-5), places=10) + def test_uop_scalar_const_lifts_kind(self): for dtype, value, out_dtype, const_dtype in ((dtypes.weakint, 1, dtypes.weakint, dtypes.weakint), (dtypes.int32, 1, dtypes.int32, dtypes.weakint), diff --git a/test/unit/test_multitensor.py b/test/unit/test_multitensor.py index 96b90ff345..139c4588ff 100644 --- a/test/unit/test_multitensor.py +++ b/test/unit/test_multitensor.py @@ -2,7 +2,7 @@ import unittest, numpy as np from tinygrad import Tensor, Variable, Context, Device, TinyJit, GlobalCounters, dtypes, UOp, nn, getenv from tinygrad.nn.state import get_parameters, get_state_dict from tinygrad.uop.ops import Ops -from test.helpers import not_support_multi_device, needs_second_gpu, slow, assert_kernel_count +from test.helpers import not_support_multi_device, needs_second_gpu, slow, assert_kernel_count, KernelCountException from hypothesis import given, strategies as strat, settings settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False)) @@ -384,6 +384,12 @@ class TestMultiTensor(unittest.TestCase): np.testing.assert_allclose(r.numpy(), np.ones(256)+np.ones(256), atol=1e-4, rtol=1e-5) assert jf.captured is not None + def test_symbolic_broadcast_copy(self): + rows = Variable("rows", 1, 4).bind(3) + out = Tensor.ones(rows, 8).to(devices_2).realize() + self.assertEqual(out.shape, (rows, 8)) + np.testing.assert_equal(out[:3].to(Device.DEFAULT).numpy(), np.ones((3, 8))) + def test_multitensor_jit_in_list(self): # test MULTI tensor inside a list container - exercises the container unpacking + MULTI unpacking @TinyJit @@ -583,7 +589,7 @@ class TestMultiTensor(unittest.TestCase): zeros = Tensor.zeros(3).realize() b = a.to(devices_2)*zeros.to(devices_2) sched = b.schedule_linear().src - self.assertEqual(len(sched), 0) + if len(sched) != 0: raise KernelCountException(0, len(sched)) self.assertListEqual(b.tolist(), [0, 0, 0]) @unittest.skipIf(not_support_multi_device(), "no multi") diff --git a/tinygrad/callify.py b/tinygrad/callify.py deleted file mode 100644 index b7bb8d4972..0000000000 --- a/tinygrad/callify.py +++ /dev/null @@ -1,221 +0,0 @@ -from dataclasses import dataclass, field -from tinygrad.dtype import dtypes, AddrSpace -from tinygrad.uop.ops import UOp, UPat, PatternMatcher, Ops, GroupOp, ParamArg, graph_rewrite, track_rewrites -from tinygrad.helpers import VIZ, pluralize, all_int - -@dataclass -class AllocCtx: - uop_list: list[UOp] = field(default_factory=list) - buffer_map: dict[UOp, UOp] = field(default_factory=dict) - bases: set[UOp] = field(default_factory=set) - assigns: list[UOp] = field(default_factory=list) - replacements: list[UOp] = field(default_factory=list) - -def tag_uop(ctx:AllocCtx, x:UOp): - if x.tag is not None: return None - ctx.uop_list.append(x) - return x.replace(tag=(len(ctx.uop_list)-1,)) - -def disk_like(u:UOp): return isinstance(u.device, str) and u.device.startswith(("DISK", "TINYFS")) - -def disk_copy_is_buffer(ctx:AllocCtx, u:UOp): - # copies to disk are replaced with the disk buffer - if disk_like(u) and u.tag is None: - ctx.buffer_map[u] = u.empty_like() - return u.rtag(()) - # all copies from disk/numpy are realized into a real buffer - from_creation = isinstance(u.src[0].device, str) and u.src[0].device.startswith(("NPY", "DISK", "PYTHON", "TINYFS")) - if from_creation: return tag_uop(ctx, u) - -# CONTIGUOUS and AFTER + parents are the only nodes that get updated -add_tags = PatternMatcher([ - (UPat(Ops.COPY, name="u"), disk_copy_is_buffer), - # no tag on copies that are assigned via STORE+AFTER — merge COPY tag into AFTER - (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE, src=(UPat(name="dest"), UPat(Ops.COPY, name="c")))), name="a"), - lambda a,c,dest: a.replace(src=(a.src[0], a.src[1].replace(src=(dest, c.rtag(())))), tag=a.tag+c.tag) if a.tag and c.tag else None), - (UPat((Ops.CONTIGUOUS, Ops.AFTER), name="x"), tag_uop), - (UPat(GroupOp.All, name="x"), lambda ctx,x: tag_uop(ctx,x) if x in ctx.bases else None), -]) - -def replace_contig_with_store_after(u:UOp): - # can't allocate a buffer for a virtual value - if u.is_virtual: return None - # if size is 0, remove the contig - if 0 in u.shape: return u.src[0] - # no real contig for DISK/TINYFS tensors, they are left alone - if disk_like(u): return u.rtag(None) - buf = u.empty_like() - return buf.after(buf.store(u.src[0])).rtag(u.tag) - -def replace_store_after_with_contig(u:UOp, src:UOp): - assigned_to = u - while assigned_to.op in {Ops.BITCAST, Ops.AFTER, Ops.UNSHARD}: assigned_to = assigned_to.src[0].base - if assigned_to.op not in {Ops.BUFFER, Ops.SLICE}: return src.contiguous(tag=u.tag) - -def _make_buffer_view(src:UOp) -> UOp|None: - """If movement ops on src collapse to a contiguous range, return SLICE. Otherwise None.""" - if (offset := src.contiguous_view_offset()) is None: return None - buf = src.base - if buf.op is Ops.SLICE: - byte_offset = buf.src[1].val * buf.src[0].dtype.itemsize + offset * src.dtype.itemsize - buf = buf.src[0] - if byte_offset % buf.dtype.itemsize != 0: return None - offset = byte_offset // buf.dtype.itemsize - return UOp(Ops.SLICE, src.dtype, (buf, UOp.const(offset)), src.numel()) - -def contiguous_mops_to_view(c:UOp, src:UOp): - """MOPS(BUFFER) → SLICE when movement ops collapse to a contiguous range.""" - buf = src.base - if buf.op not in {Ops.BUFFER, Ops.SLICE, Ops.UNSHARD}: return None - if src.op is Ops.RESHAPE and src.src[0].op in {Ops.BUFFER, Ops.SLICE} and c.op is not Ops.BITCAST: return None - if c.op is not Ops.BITCAST and src.op is Ops.BUFFER: return None - - # no symbolic shape - if not all_int(c.shape): return None - - if buf.op is not Ops.UNSHARD and (view := _make_buffer_view(src)) is not None: - view = (view.replace(dtype=c.dtype, arg=c.numel()) if c.op is Ops.BITCAST else view).reshape(c.shape) - return c.replace(src=(view,)) if c.op is Ops.COPY else view - - # for UNSHARD tensors, use multi_pm to resolve per-shard movement ops, then create SLICE on the resolved result - if not isinstance(c.device, str): - from tinygrad.schedule.multi import multi_pm - resolved = graph_rewrite(src, multi_pm, name="multi_buffer_view") - if resolved.op is not Ops.UNSHARD: return None - if (view := _make_buffer_view(resolved.src[0])) is None: return None - return view.reshape(resolved.src[0].shape).unshard(resolved.arg, resolved.src[1:]).contiguous(tag=c.tag) - - return None - -def _precompiled_output_redirect(s:UOp, t:UOp) -> UOp|None: - # how output s lands in the caller's buffer t, or None if it must be copied into t - # materialize straight into t - if s.op is Ops.CONTIGUOUS: return t.after(t.store(s.src[0])) - # rebind output storage to t - if s.op in {Ops.BUFFER, Ops.UNSHARD} and s.has_buffer_identity(): return t - return None - -def transform_precompiled_call(c:UOp) -> UOp|None: - if not c.arg.precompile: return None - assert c.src[0].op is Ops.TUPLE, f"expected TUPLE body for precompiled FUNCTION, got {c.src[0].op}" - input_buffers = tuple(x.contiguous() if x.op not in {Ops.AFTER, Ops.BIND} else x for x in c.src[1:]) - - # add the outputs to the call - srcs = c.src[0].src - resolved = [c.gettuple(i) for i in range(len(srcs))] - outs = tuple(r.empty_like() for r in resolved) - targets = [o.param_like(len(c.src)-1+i).shrink_to(s.shape) for i,(o,s) in enumerate(zip(outs, srcs))] - - subs:dict[UOp, UOp] = {} - items:list[UOp] = [] - for s, t in zip(srcs, targets): - after_deps:list[UOp] = [] - while s.op is Ops.AFTER: - after_deps.extend(s.src[1:]) - s = s.src[0] - if (placed := _precompiled_output_redirect(s, t)) is not None and s not in subs: - subs[s] = placed - items.append(s.after(*after_deps) if after_deps else s) - else: - items.append(t.after(t.store(s.after(*after_deps)))) - fxn = UOp.sink(*(x.substitute(subs) for x in items)) - - # body switches from TUPLE to SINK, so the node becomes an opaque CALL (not FUNCTION) - new_call = UOp(Ops.CALL, src=(fxn, *input_buffers, *outs), arg=c.arg) - rets = tuple(o.after(new_call) for o in outs) - - # if the CALL has symbolic shapes, shrink the max-sized output to the actual symbolic shape - # NOTE: must use resolved shapes from the FUNCTION (which substitutes PARAMs with external args), not raw body shapes - rets = tuple(r.shrink_to(rs.shape) for r,rs in zip(rets, resolved)) - - return UOp.maketuple(*rets) - -# NOTE: adding rules to here is bad. these all need to run before the schedule cache -pm_early_transform_tensor_graph = PatternMatcher([ - # transform precompiled FUNCTIONs into CALLs (body becomes SINK with stores) - (UPat(Ops.FUNCTION, name="c"), transform_precompiled_call), - - # resolve TUPLE+GETTUPLE (for precompiled calls) - (UPat(Ops.GETTUPLE, src=(UPat(Ops.TUPLE, name="t"),), name="g"), lambda g,t: t.src[g.arg]), - - # fold MOPS+BITCAST over BUFFER/SLICE into SLICE when movement ops collapse to contiguous range - (UPat((Ops.BITCAST, Ops.COPY, Ops.CONTIGUOUS), src=(UPat(GroupOp.Movement|{Ops.BUFFER}, name="src"),), name="c"), contiguous_mops_to_view), - - # remove contiguous on movement ops before a copy on disk - (UPat(GroupOp.Movement-{Ops.SHRINK, Ops.RESHAPE}, name="x").f(Ops.CONTIGUOUS).f(Ops.COPY, name="copy"), lambda x,copy: - copy.replace(src=(x,), tag=None) if isinstance(x.device, str) and x.device.startswith("DISK") else None), - # push copy past movement ops to disk - (UPat(GroupOp.Movement-{Ops.SHRINK, Ops.RESHAPE}, name="x").f(Ops.COPY, name="copy"), lambda x,copy: - x.replace(src=(copy.replace(src=(x.src[0],), tag=None),)+x.src[1:]) \ - if isinstance(x.device, str) and x.device.startswith("DISK") else None), - - # add CONTIGUOUS to tagged UOps - (UPat(GroupOp.All-{Ops.CONTIGUOUS, Ops.AFTER, Ops.STORE}, name="x"), - lambda x: None if x.tag is None else x.rtag(None).contiguous(tag=x.tag) if x.tag else x.replace(tag=None)), - # remove extra CONTIGUOUS on AFTER (only when target is contiguous) - (UPat(Ops.CONTIGUOUS, src=(UPat(Ops.AFTER, name="a"),), name="c"), - lambda a,c: a.replace(tag=(a.tag or ())+(c.tag or ())) if a.src[0].has_buffer_identity() else None), - # replace AFTER+STORE with CONTIGUOUS when target is not a buffer - (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE, src=(UPat(), UPat(name="src")))), name="u"), replace_store_after_with_contig), - # replace CONTIGUOUS with STORE+AFTER - (UPat(Ops.CONTIGUOUS, name="u"), replace_contig_with_store_after), - # remove DETACH/CONTIGUOUS_BACKWARD (allows more contiguous removal) - (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD), name="x"), lambda x: x.src[0]), -]) - -def finalize_after(ctx:AllocCtx, x:UOp): - # untagged: record as an assign for the call body - if x.tag is None: - ctx.assigns.append(x) - return None - # tagged: untag and map each original pre-rewrite UOp to the stripped buffer; the untagged result is reprocessed as untagged - ret = x.replace(tag=None) - replace_uop = ret - while replace_uop.op is Ops.AFTER: replace_uop = replace_uop.src[0] - for t in x.tag: - original_uop: UOp = ctx.uop_list[t] - ctx.buffer_map[original_uop] = replace_uop.shrink_to(original_uop.shape) - return ret - -def replace_input_buffer(ctx:AllocCtx, b:UOp): - ctx.replacements.append(b) - if b.op is Ops.BIND: return b.param_like(len(ctx.replacements)-1) - return UOp.param(len(ctx.replacements)-1, b.dtype, b.shape, b.device, - addrspace=b.addrspace if b.addrspace is not None else AddrSpace.GLOBAL) - -pm_finalize_call = PatternMatcher([ - (UPat(Ops.AFTER, name="x"), finalize_after), - (UPat(Ops.COPY, name="x"), lambda ctx,x: ctx.assigns.append(x) if isinstance(x.device, str) and x.device.startswith(("DISK", "TINYFS")) else None), -]) - -pm_replace_buf = PatternMatcher([ - # replace BUFFER with PARAM for cache key normalization - (UPat(Ops.BUFFER, src=(UPat(),), name="b"), lambda ctx,b: - replace_input_buffer(ctx, b) if isinstance(b.arg, ParamArg) and b.addrspace is AddrSpace.GLOBAL else None), - # replace SLICE with PARAM. this rewrite is bottom up so BUFFERs we don't need won't be in the input - (UPat(Ops.SLICE, src=(UPat(Ops.BUFFER), UPat(Ops.CONST, dtype=dtypes.weakint)), name="b"), replace_input_buffer), - # strip value from BIND for cache key normalization, so different values hit same cache - (UPat(Ops.BIND, src=(UPat(Ops.PARAM), UPat(Ops.CONST)), name="b"), replace_input_buffer), -]) - -@track_rewrites(lambda _,ret: f"Callify {pluralize('Buffer', len(ret[1]))}") -def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]: - if VIZ: graph_rewrite(big_sink, PatternMatcher([]), name="View Tensor Graph") - # uop list is a list in the original_sink graph and we can map to the tags later - # same predicate as Tensor.realize - ctx = AllocCtx(bases={base for x in big_sink.src if not (base:=x.base).is_virtual and not base.has_buffer_identity() - and base.op is not Ops.AFTER and base.addrspace is not AddrSpace.ALU}) - - # this rewrite is "read-only", it adds simple things to buffer_map and may sink things on big_sink, bottom_up - # this is the only one where we have to be careful to not break the tensor graph - big_sink = graph_rewrite(big_sink, add_tags, ctx=ctx, bottom_up=True, name="number the uops") - - # here we can break the tensor graph. this is the only place you need to maintain numbered tags - big_sink = graph_rewrite(big_sink, pm_early_transform_tensor_graph, name="early transform tensor graph") - - # here we construct the final buffer_map: as-built nodes -> their final storage. values are never keys - graph_rewrite(big_sink, pm_finalize_call, ctx=ctx, name="finalize call") - ret = graph_rewrite(UOp.sink(*ctx.assigns), pm_replace_buf, ctx=ctx, bottom_up=True, name="replace bufs").call(*ctx.replacements) - assert not any(x in ctx.buffer_map for x in ctx.buffer_map.values()) - if VIZ: graph_rewrite(ret, PatternMatcher([]), name="View Call") - return ret, ctx.buffer_map diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 58fceb0882..7d8c0dbd37 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -2,8 +2,8 @@ from dataclasses import replace, dataclass import itertools, functools from tinygrad.helpers import DISABLE_FAST_IDIV, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC from tinygrad.helpers import ALLOW_TF32, DEFAULT_FLOAT, DEFAULT_INT, TracingKey, Context, panic -from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, ProgramInfo, GroupOp -from tinygrad.uop.ops import AxisType, pm_commit_weak, pm_cast_weak +from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, Ops, UPat, rewrite_group, KernelInfo, ProgramInfo, GroupOp, AxisType +from tinygrad.uop.weak import pm_lower_index_dtype, pm_commit_weak, pm_cast_weak from tinygrad.uop.render import pyrender from tinygrad.uop.spec import type_verify, spec_tensor, spec_program from tinygrad.renderer import Renderer, Estimates @@ -12,7 +12,7 @@ from tinygrad.dtype import dtypes, AddrSpace # import all pattern matchers here from tinygrad.codegen.gpudims import pm_add_gpudims -from tinygrad.uop.symbolic import sym, symbolic_simple, symbolic, pm_move_where_on_load, pm_clean_up_group_sink, pm_remove_invalid +from tinygrad.uop.symbolic import sym, symbolic_simple, symbolic, pm_fold_cast_const, pm_move_where_on_load, pm_clean_up_group_sink, pm_remove_invalid from tinygrad.uop.movement import mop_cleanup from tinygrad.codegen.decomp.dtype import pm_dtype_decomps from tinygrad.codegen.decomp.op import get_late_rewrite_patterns, get_simplifying_rewrite_patterns @@ -20,7 +20,7 @@ from tinygrad.codegen.decomp.transcendental import get_transcendental_patterns from tinygrad.codegen.late.coalesce import indexing_simplify from tinygrad.codegen.opt.postrange import apply_opts from tinygrad.codegen.late.gater import pm_move_gates_from_index -from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse +from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse, pm_reduce_unparented from tinygrad.schedule.multi import multi_pm from tinygrad.schedule.rangeify import pm_mops from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize @@ -301,7 +301,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: sink = graph_rewrite(sink, pm_split_ranges+pm_flatten_range, ctx={}, name="split ranges") # symbolic (NOTE: this is a requirement for pm_simplify_ranges to be correct) - sink = graph_rewrite(sink, sym+pm_flatten_range, name="initial symbolic") + sink = graph_rewrite(sink, sym+pm_fold_cast_const+pm_flatten_range, name="initial symbolic") # optimize (schedule) the AST sink = graph_rewrite(sink, pm_flatten_range+pm_simplify_ranges, ctx={}, name="simplify ranges") @@ -310,7 +310,8 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: sink = apply_opts(sink, ren, beam=ast.arg.beam) # ** expander (expand_rewrite) ** - sink = graph_rewrite(sink, sym+pm_move_where_on_load+pm_flatten_range, name="postopt symbolic") + # reduce_unparented: a REDUCE whose src folded to a CONST (e.g. x*0) has no parented ranges, collapse it before the expander + sink = graph_rewrite(sink, sym+pm_move_where_on_load+pm_flatten_range+pm_reduce_unparented, name="postopt symbolic") # expand sink = graph_rewrite(sink, expander2, ctx=build_range_map(sink), name="expander") @@ -336,14 +337,16 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: # do memory coalescing (late) sink = memory_coalescing(sink, ren) - sink = graph_rewrite(sink, symbolic_simple+ew_devectorizer+pm_simplify_add_image, name="add images", ctx=({}, ren), bottom_up=True) + sink = graph_rewrite(sink, symbolic_simple+ew_devectorizer+pm_simplify_add_image, + name="add images", ctx=({}, ren), bottom_up=True) # extra symbolic before decomp. crashes without this? - sink = graph_rewrite(sink, sym, name="extra symbolic") + # NOTE: also run indexing_simplify here, while the index is still weakint and (x+y)*c -> x*c+y*c applies + sink = graph_rewrite(sink, sym+indexing_simplify, name="extra symbolic") # lower index dtype # NOTE: we need indexing_simplify to remove the cast to long using the Invalid - sink = graph_rewrite(sink, pm_lower_index_dtype+indexing_simplify, ctx={}, name="lower all index dtypes") + sink = graph_rewrite(sink, symbolic_simple+pm_fold_cast_const+pm_lower_index_dtype+indexing_simplify, ctx={}, name="lower all index dtypes") # final symbolic before decomp sink = graph_rewrite(sink, symbolic, name="final symbolic") @@ -354,7 +357,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: # floordiv+mod / dtype decomp (early) supported_ops = tuple(ren.code_for_op.keys()) - pm_decomp = symbolic_simple+get_simplifying_rewrite_patterns(supported_ops) + pm_decomp = symbolic_simple+pm_fold_cast_const+get_simplifying_rewrite_patterns(supported_ops) sink = graph_rewrite(sink, pm_decomp, name="early decompositions") # late decomps + move gates from unrenderable INVALID where @@ -448,7 +451,7 @@ pm_to_program = PatternMatcher([ (UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.LINEAR), UPat(Ops.SOURCE, name="source")), name="prg"), do_compile), ]) -@track_rewrites(name=lambda ast,renderer,ret,**kwargs: TracingKey(ret.src[0].arg.name,(ret.src[0].arg.function_name, ast), ret=renderer), replay=True) +@rewrite_group(name=lambda ast,renderer,ret,**kwargs: TracingKey(ret.src[0].arg.name,(ret.src[0].arg.function_name, ast), ret=renderer), replay=True) @Context(ALLOW_DEVICE_USAGE=0) def do_to_program(ast:UOp, renderer:Renderer) -> UOp: """ diff --git a/tinygrad/codegen/decomp/dtype.py b/tinygrad/codegen/decomp/dtype.py index cbb324ace0..52b8e46304 100644 --- a/tinygrad/codegen/decomp/dtype.py +++ b/tinygrad/codegen/decomp/dtype.py @@ -78,9 +78,11 @@ def l2i(op: Ops, dt: DType, *uops:UOp): case Ops.MAX: return l2i(Ops.WHERE, dt, l2i(Ops.CMPLT, dt, *uops), b0, b1, a0, a1) case _: raise NotImplementedError(f"long decomposition of {op} unsupported") -def split_l2i(op: Ops, dt: DType, *uops:UOp): - # l2i does arithmetic on its inputs; rules enter here to split them to 32-bit words first, l2i recurses on itself - return l2i(op, dt, *graph_rewrite(UOp.sink(*uops), pm_long_decomp, bottom_up=True).src) +def split_l2i(ctx:dict, op: Ops, dt: DType, *uops:UOp): + # l2i does arithmetic on its inputs; rules enter here to split them to 32-bit words first, l2i recurses on itself. + # both word halves of a node ask for the same split, so ctx memos it for the pass + if (key:=(op, dt, uops)) not in ctx: ctx[key] = l2i(op, dt, *graph_rewrite(UOp.sink(*uops), pm_long_decomp, ctx=ctx, bottom_up=True).src) + return ctx[key] # ***** floats ***** f2f_dt = { f:getattr(dtypes, f"uint{f.bitsize}") for f in dtypes.floats } @@ -97,7 +99,8 @@ def f2f(v, fr:DType, to:DType, sat=True): if fr in dtypes.fp8_fnuz: fnuz_nan = sign.ne(0) & nosign.eq(0) qnan = shl(shl(1, te) - 1, tm) | shl(1, tm - 1) - return fnuz_nan.where(qnan, sign | exp.eq(0).where(0, norm)).bitcast(to) + # the fnuz bias can exceed the target's: exp in [1, fb-tb] is normal in fr but lands below to's normal range, so it flushes like a denormal + return fnuz_nan.where(qnan, sign | (exp < max(fb - tb, 0) + 1).where(0, norm)).bitcast(to) # fp8e4m3 has only one nan is_nan = (nosign.eq(shl(1, fm + fe) - 1) if fr == dtypes.fp8e4m3 else exp.eq(shl(1, fe) - 1)) return (sign | exp.eq(0).where(0, is_nan.where(nan, norm))).bitcast(to) @@ -139,21 +142,21 @@ pm_long_decomp = PatternMatcher([ (UPat(Ops.STORE, src=(UPat.var('idx', tuple(l2i_dt.keys())), UPat.var('val')), name='st'), lambda st,idx,val: st.replace(src=(idx.rtag((0, dt:=l2i_dt[idx.dtype])), val.rtag((0, dt)))).group( st.replace(src=(idx.rtag((1, dt)), val.rtag((1, dt))))) if val.tag is None else None), - (UPat(GroupOp.Comparison, src=[UPat.var('a', tuple(l2i_dt.keys())), UPat()], name="x"), lambda a,x: - split_l2i(x.op, dt:=l2i_dt[a.dtype], *flatten((s.rtag((0, dt)), s.rtag((1, dt))) for s in x.src))), - (UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a', tuple(l2i_dt.keys())),), name="x"), lambda a,x: - split_l2i(Ops.BITCAST, l2i_dt[x.dtype], a.rtag((0, dt:=l2i_dt[a.dtype])), a.rtag((1, dt)))[x.tag[0]]), - (UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a'),), name="x"), lambda a,x: - split_l2i(x.op, x.dtype, a)[x.tag[0]] if x.tag is not None else None), - (UPat(Ops.CAST, src=(UPat.var('a', tuple(l2i_dt.keys())),), name="x"), lambda a,x: - split_l2i(x.op, x.dtype, a.rtag((0, dt:=l2i_dt[a.dtype])), a.rtag((1, dt))) if x.dtype not in l2i_dt and a.tag is None else None), - (UPat((Ops.SHL, Ops.SHR), tuple(l2i_dt.keys()), src=(UPat.var('a'), UPat.var('b')), name="x"), lambda a,b,x: - split_l2i(x.op, dt:=l2i_dt[x.dtype], a.rtag((0, dt)), a.rtag((1, dt)), b.rtag((0, dt)))[x.tag[0]] if x.tag is not None else None), - (UPat(Ops.WHERE, tuple(l2i_dt.keys()), src=(UPat.var('c'), UPat.var('a'), UPat.var('b')), name="x"), lambda a,b,c,x: - split_l2i(x.op, dt:=l2i_dt[x.dtype], c, a.rtag((0, dt)), a.rtag((1, dt)), b.rtag((0, dt)), b.rtag((1, dt)))[x.tag[0]] + (UPat(GroupOp.Comparison, src=[UPat.var('a', tuple(l2i_dt.keys())), UPat()], name="x"), lambda ctx,a,x: + split_l2i(ctx, x.op, dt:=l2i_dt[a.dtype], *flatten((s.rtag((0, dt)), s.rtag((1, dt))) for s in x.src))), + (UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a', tuple(l2i_dt.keys())),), name="x"), lambda ctx,a,x: + split_l2i(ctx, Ops.BITCAST, l2i_dt[x.dtype], a.rtag((0, dt:=l2i_dt[a.dtype])), a.rtag((1, dt)))[x.tag[0]]), + (UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a'),), name="x"), lambda ctx,a,x: + split_l2i(ctx, x.op, x.dtype, a)[x.tag[0]] if x.tag is not None else None), + (UPat(Ops.CAST, src=(UPat.var('a', tuple(l2i_dt.keys())),), name="x"), lambda ctx,a,x: + split_l2i(ctx, x.op, x.dtype, a.rtag((0, dt:=l2i_dt[a.dtype])), a.rtag((1, dt))) if x.dtype not in l2i_dt and a.tag is None else None), + (UPat((Ops.SHL, Ops.SHR), tuple(l2i_dt.keys()), src=(UPat.var('a'), UPat.var('b')), name="x"), lambda ctx,a,b,x: + split_l2i(ctx, x.op, dt:=l2i_dt[x.dtype], a.rtag((0, dt)), a.rtag((1, dt)), b.rtag((0, dt)))[x.tag[0]] if x.tag is not None else None), + (UPat(Ops.WHERE, tuple(l2i_dt.keys()), src=(UPat.var('c'), UPat.var('a'), UPat.var('b')), name="x"), lambda ctx,a,b,c,x: + split_l2i(ctx, x.op, dt:=l2i_dt[x.dtype], c, a.rtag((0, dt)), a.rtag((1, dt)), b.rtag((0, dt)), b.rtag((1, dt)))[x.tag[0]] if x.tag is not None else None), - (UPat((*(GroupOp.ALU - GroupOp.Comparison - {Ops.SHL, Ops.SHR, Ops.WHERE}), Ops.BITCAST), tuple(l2i_dt.keys()), name="x"), lambda x: - split_l2i(x.op, l2i_dt[x.dtype], *flatten((a.rtag((0, l2i_dt[x.dtype])), a.rtag((1, l2i_dt[x.dtype]))) for a in x.src))[x.tag[0]] + (UPat((*(GroupOp.ALU - GroupOp.Comparison - {Ops.SHL, Ops.SHR, Ops.WHERE}), Ops.BITCAST), tuple(l2i_dt.keys()), name="x"), lambda ctx,x: + split_l2i(ctx, x.op, l2i_dt[x.dtype], *flatten((a.rtag((0, l2i_dt[x.dtype])), a.rtag((1, l2i_dt[x.dtype]))) for a in x.src))[x.tag[0]] if x.tag is not None else None), (UPat(Ops.LOAD, tuple(l2i_dt.keys()), src=(UPat.var('idx'),), name='x'), lambda x,idx: x.replace(dtype=l2i_dt[x.dtype], src=(reindex(idx, x.tag[0]).replace(dtype=l2i_dt[x.dtype], tag=None),), tag=None) if x.tag is not None else None), @@ -197,7 +200,7 @@ def do_dtype_decomps(sink:UOp, ctx:tuple[set[DType], Renderer]) -> UOp: to = dtypes.int if fr == dtypes.long else dtypes.half if not _should_emulate(dtypes.half) and fr in dtypes.fp8s else dtypes.float if DEBUG >= 2: print(f"emulating {fr} as {to}") pm = pm_float_decomp if fr in dtypes.floats else pm_long_decomp - sink = graph_rewrite(sink, pm, name=f"decomp {fr} -> {to}", ctx=(fr, to), bottom_up=True) + sink = graph_rewrite(sink, pm, name=f"decomp {fr} -> {to}", ctx={} if pm is pm_long_decomp else (fr, to), bottom_up=True) ctx[0].clear() return sink diff --git a/tinygrad/codegen/late/coalesce.py b/tinygrad/codegen/late/coalesce.py index 47ec328017..f72c951c16 100644 --- a/tinygrad/codegen/late/coalesce.py +++ b/tinygrad/codegen/late/coalesce.py @@ -1,8 +1,8 @@ import itertools, functools from collections import defaultdict from tinygrad.dtype import dtypes, AddrSpace, Invalid, DType -from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat, GroupOp, shape_to_shape_arg -from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate +from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat, GroupOp, shape_to_shape_arg, graph_rewrite +from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate, sym from tinygrad.helpers import getenv, IMAGE, OSX, ceildiv, is_image_shape from tinygrad.renderer import Renderer @@ -27,11 +27,14 @@ def _drop_valid_stmts(valid:UOp, idx:UOp, height:int, width:int) -> list[UOp]: lo, hi = (c + 1, X.vmax) if is_upper_bound else (X.vmin, c - 1) if lo <= hi: fake = UOp.variable(f"fake{i}", lo, hi, X.dtype) - for coord,b in zip(idx.src, (width, height)): - rw = coord.substitute({X:fake}).simplify() - if rw.vmin >= b or rw.vmax < 0: - drop_stmt.append(stmt) - break + subs = [{X: fake}] + # idx may not have X itself, so also substitute a term of X: v -> fake - (X - v) + terms = list(X.split_uop(Ops.ADD)) + v = next((u for u in terms if u.op in GroupOp.Irreducible and u.op is not Ops.CONST), None) + if v is not None and (rest:=[u for u in terms if u is not v]): subs.append({v: fake - UOp.usum(*rest)}) + if any((testidx:=graph_rewrite(coord.substitute(sub), sym)).vmin >= b or testidx.vmax < 0 + for sub in subs for coord,b in zip(idx.src, (width, height))): + drop_stmt.append(stmt) return drop_stmt def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None: diff --git a/tinygrad/codegen/opt/postrange.py b/tinygrad/codegen/opt/postrange.py index 987aea76ba..f5f682c270 100644 --- a/tinygrad/codegen/opt/postrange.py +++ b/tinygrad/codegen/opt/postrange.py @@ -332,9 +332,9 @@ class Scheduler: @property def group_for_reduces(self) -> int: return len(self.axes_of(AxisType.GROUP_REDUCE)) -def bufs_from_ast(ast:UOp, dname:str) -> list[Buffer]: +def args_from_ast(ast:UOp, dname:str) -> tuple[list[Buffer], dict[str, int]]: glbls = sorted([x for x in ast.backward_slice if x.op is Ops.PARAM and x.arg.slot >= 0], key=lambda x: x.arg.slot) - return [Buffer(dname, x.max_numel(), x.dtype) for x in glbls] + return [Buffer(dname, x.max_numel(), x.dtype) for x in glbls], {k.expr:int(k.vmax+k.vmin)//2 for k in ast.variables()} def apply_opts(ast:UOp, ren:Renderer, beam:int=0) -> UOp: if ast.tag is not None: return ast @@ -344,10 +344,10 @@ def apply_opts(ast:UOp, ren:Renderer, beam:int=0) -> UOp: for opt in ast.arg.opts_to_apply: k.apply_opt(opt) elif beam >= 1: from tinygrad.codegen.opt.search import beam_search - rawbufs = bufs_from_ast(ast, ren.target.device) + rawbufs, var_vals = args_from_ast(ast, ren.target.device) # beam search may open devices with Context(ALLOW_DEVICE_USAGE=1): - k = beam_search(k, rawbufs, beam, bool(getenv("BEAM_ESTIMATE", 1))) + k = beam_search(k, rawbufs, var_vals, beam, bool(getenv("BEAM_ESTIMATE", 1))) elif not NOOPT and (ast.arg is None or ast.arg.applied_opts == ()): from tinygrad.codegen.opt.heuristic import hand_coded_optimizations # NOTE: hand_coded_optimizations doesn't support multiblock opts yet diff --git a/tinygrad/codegen/opt/search.py b/tinygrad/codegen/opt/search.py index cf85cdebf3..c8c7a3680a 100644 --- a/tinygrad/codegen/opt/search.py +++ b/tinygrad/codegen/opt/search.py @@ -1,6 +1,6 @@ import math, time, multiprocessing, traceback, signal, atexit from dataclasses import replace -from tinygrad.uop.ops import sym_infer, AxisType, UOp +from tinygrad.uop.ops import sym_infer, AxisType, UOp, Ops from tinygrad.uop.render import pyrender from tinygrad.device import Device, Buffer from tinygrad.helpers import prod, flatten, DEBUG, CACHELEVEL, diskcache_get, diskcache_put, getenv, Context, colored, time_to_str @@ -62,7 +62,8 @@ def _try_compile(x:tuple[int,Scheduler]) -> tuple[int, tuple[UOp, float]|None]: ret = None try: st = time.perf_counter() - prg = to_program(x[1].copy().get_optimized_ast(name_override="test"), x[1].ren) + ast, dev = x[1].copy().get_optimized_ast(name_override="test"), x[1].ren.target.device + prg = to_program(ast.substitute({p: p.replace(arg=replace(p.arg, device=dev)) for p in ast.toposort() if p.op is Ops.PARAM}), x[1].ren) et = time.perf_counter() - st uops = prg.src[1].src if len(uops) >= (uops_max:=getenv("BEAM_UOPS_MAX", 3000)) > 0: @@ -111,7 +112,7 @@ def get_kernel_actions(s:Scheduler, include_0=True, max_up:int|None=None) -> dic return acted beam_pool, BEAM_DEBUG = None, getenv("BEAM_DEBUG") -def beam_search(s:Scheduler, rawbufs:list[Buffer], amt:int, allow_test_size=True, disable_cache=IGNORE_BEAM_CACHE.value): +def beam_search(s:Scheduler, rawbufs:list[Buffer], var_vals:dict[str,int], amt:int, allow_test_size=True, disable_cache=IGNORE_BEAM_CACHE.value): global beam_pool key = {"ast": s.ast.key, "amt": amt, "allow_test_size": allow_test_size, "device": s.ren.target.device, "suffix": s.ren.suffix} if not disable_cache and CACHELEVEL >= 1 and (val:=diskcache_get("beam_search", key)) is not None: @@ -136,7 +137,6 @@ def beam_search(s:Scheduler, rawbufs:list[Buffer], amt:int, allow_test_size=True try: rawbufs = _ensure_buffer_alloc(rawbufs) - var_vals: dict[str, int] = {k.expr:int(k.vmax+k.vmin)//2 for k in s.ast.variables()} exiting, st = False, time.perf_counter() dev = Device[s.ren.target.device] while not exiting: diff --git a/tinygrad/codegen/simplify.py b/tinygrad/codegen/simplify.py index 4aadeb6d6f..879dd0e6c9 100644 --- a/tinygrad/codegen/simplify.py +++ b/tinygrad/codegen/simplify.py @@ -1,7 +1,7 @@ import itertools from typing import Callable from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, graph_rewrite, _substitute, range_start, AxisType -from tinygrad.uop.symbolic import symbolic, invalid_gate +from tinygrad.uop.symbolic import symbolic, pm_fold_cast_const, invalid_gate from tinygrad.helpers import partition from tinygrad.dtype import dtypes @@ -32,7 +32,7 @@ def simplify_merge_adjacent(u:UOp) -> UOp|None: 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}, + nidx = graph_rewrite(u, _substitute+symbolic+pm_fold_cast_const+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1}, name=f"check_merge_{r0.arg[0]}_{r1.arg[0]}") # check if it simplifies diff --git a/tinygrad/device.py b/tinygrad/device.py index 570cd79813..6f796e07a6 100644 --- a/tinygrad/device.py +++ b/tinygrad/device.py @@ -336,6 +336,8 @@ class Compiled: pm_lower:Any = None pm_bufferize:Any = None + has_copy_queue:bool = True + def __init__(self, device:str, allocator:Allocator, renderers:list[type[Renderer]], runtime:type[Program[Self]]|None, graph=None, arch=None): from tinygrad.renderer import Renderer self.device, self.allocator, self.runtime_t, self.graph, self.renderers = device, allocator, runtime, graph, renderers or [Renderer] diff --git a/tinygrad/dtype.py b/tinygrad/dtype.py index 62464e02ea..9c0b336a58 100644 --- a/tinygrad/dtype.py +++ b/tinygrad/dtype.py @@ -80,7 +80,7 @@ class DType(metaclass=DTypeMetaClass): # NOTE: float('nan') != float('nan'), so we canonicalize here if isinstance(val, float) and math.isnan(val): val = math.nan # int is the default. wrap floats in ConstFloat to distinguish -0.0 from 0.0 in cache - return ConstFloat(float(val)) if dtypes.is_float(self) else bool(val) if dtypes.is_bool(self) else int(val) + return ConstFloat(truncate.get(self, float)(float(val))) if dtypes.is_float(self) else bool(val) if dtypes.is_bool(self) else int(val) class DTypes: diff --git a/tinygrad/engine/jit.py b/tinygrad/engine/jit.py index 97074388ff..fe11805c17 100644 --- a/tinygrad/engine/jit.py +++ b/tinygrad/engine/jit.py @@ -4,7 +4,7 @@ from tinygrad.tensor import Tensor, all_tensors from tinygrad.helpers import flatten, merge_dicts, DEBUG, Context, BEAM, getenv, JIT, JIT_BATCH_SIZE, dedup, pluralize, VIZ, disable_gc from tinygrad.device import Buffer, Compiled, Device, MultiBuffer, DepsTracker from tinygrad.dtype import DType -from tinygrad.uop.ops import UOp, PatternMatcher, Variable, sym_infer, Ops, buffers, track_rewrites, graph_rewrite +from tinygrad.uop.ops import UOp, PatternMatcher, Variable, sym_infer, Ops, buffers, rewrite_group, graph_rewrite from tinygrad.renderer import Estimates from tinygrad.engine.realize import capturing, compile_linear, link_linear, run_linear, graph_cache, estimate_uop, get_runtime from tinygrad.engine.realize import unwrap_multi, resolve_params, get_call_arg_uops, get_call_outs_ins @@ -64,7 +64,7 @@ def _copy_input(u:UOp) -> UOp: run_linear(UOp(Ops.LINEAR, src=(u.copy_to_device(u.device).call(new:=UOp.new_buffer(u.device, u.max_numel(), u.dtype), u),))) return new -@track_rewrites(lambda linear,held_bufs,input_uops,ret=(): f"JIT {pluralize('call', len(linear.src))}") +@rewrite_group(lambda linear,held_bufs,input_uops,ret=(): f"JIT {pluralize('call', len(linear.src))}") def jit_lower(linear:UOp, held_bufs:set[UOp], input_uops:list[UOp]) -> UOp: if VIZ: graph_rewrite(linear, PatternMatcher([]), name="View captured linear") diff --git a/tinygrad/engine/realize.py b/tinygrad/engine/realize.py index 38d5c2831d..92006e87d6 100644 --- a/tinygrad/engine/realize.py +++ b/tinygrad/engine/realize.py @@ -3,12 +3,12 @@ from typing import cast, Iterator, Any, Sequence import time, random, itertools, math, contextlib, weakref, array from dataclasses import dataclass, replace, field from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansilen, all_int, prod, flatten, Context, getenv, to_tuple -from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events +from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events, wait_cond from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, buffers, graph_rewrite from tinygrad.device import Device, Buffer, MultiBuffer from tinygrad.renderer import Estimates from tinygrad.codegen import to_program -from tinygrad.codegen.opt.postrange import bufs_from_ast +from tinygrad.codegen.opt.postrange import args_from_ast # **************** Helpers **************** @@ -33,7 +33,7 @@ def get_call_name(call:UOp, bufs:Sequence[Buffer|UOp], var_vals:dict[str, int]|N if ast.op is Ops.COPY: return colored(f"copy {_uop_sz_to_str(arg_uops[0]):>10}, {_dev_str(bufs[0]):>7s} <- {_dev_str(bufs[1]):7s}", "yellow") if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec": return colored(f"enc/dec {_uop_sz_to_str(arg_uops[0])}", "yellow") if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph": return colored(f"batched {len(ast.src[0].src)}", "cyan") - if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "hcq": return call.arg.aux.name + if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "hcq": return cast(str, call.arg.name) raise NotImplementedError("get_call_name is not implemented") # **************** Stat **************** @@ -90,12 +90,13 @@ def optimize_local_size(call:UOp, prg:UOp) -> UOp|None: if (local_size:=local_size_cache.get(prg.key)) is None: # reuse one loaded runtime across candidates, only launch dims vary - bufs, runtime = [b.allocate() for b in bufs_from_ast(prg.src[0], device)], get_runtime(device, prg, cache=False) + (bufs, var_vals), runtime = args_from_ast(prg.src[0], device), get_runtime(device, prg, cache=False) + bufs = [b.allocate() for b in bufs] def try_exec(local_size): try: new_gs = tuple(g//l if g%l == 0 else g/l for g,l in zip(prg.arg.global_size, local_size)) return runtime(*[bufs[i].get_buf(device) for i in prg.arg.globals], global_size=new_gs, local_size=(*local_size,), - vals=prg.arg.vals({}), wait=True) + vals=prg.arg.vals(var_vals), wait=True) except Exception: return float('inf') MAX_WORKGROUP = 1024 @@ -214,16 +215,22 @@ def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> float|None: table = call.src[1+inputs].buffer for j,dev in enumerate(call.arg.aux.device): addrs = array.array('Q', [(b.bufs[j] if isinstance(b, MultiBuffer) else b).get_buf(dev).va_addr for b in bufs]) - buf = table.bufs[j] if isinstance(table, MultiBuffer) else table - buf.ensure_allocated()._buf.cpu_view().view(fmt='Q')[:len(addrs)] = addrs + mv = (table.bufs[j] if isinstance(table, MultiBuffer) else table).ensure_allocated()._buf.cpu_view().view(fmt='Q') + wait_cond(lambda: mv[0], value=0, timeout_ms=ctx.timeout or getenv("HCQDEV_WAIT_TIMEOUT_MS", 30000), msg=f"{dev} hang detected") + mv[:len(addrs)] = addrs exec_kernel(replace(ctx, update_stats=False), call, ast) - st = time.perf_counter() - for d in call.arg.aux.device: - with track_stats(ctx, call, d, [], ctx.var_vals): - if ctx.wait: Device[d].synchronize() - return time.perf_counter() - st + tms:list[float|None] = [] + for e in (aux:=call.arg.aux).prof: cast(Any, Device[e.device]).prof_ents[e.st_id] = e + for d in [cast(Any, Device[x]) for x in aux.device]: + with track_stats(ctx, call, d.device, [], ctx.var_vals) as et: + if ctx.wait: + d.synchronize(timeout=ctx.timeout) + ts = [d.signal(i)._buf.cpu_view().view(fmt='Q')[0] for e in aux.prof if e.device == d.device for i in (e.st_id, e.en_id)] + if ts: et[0] = float(max(ts)-min(ts))/d.timestamp_divider/1e6 + tms += et + return tms[0] # flatten LINEAR-in-LINEAR: any nested LINEAR child gets inlined into its parent's src pm_flatten_linear = PatternMatcher([ @@ -265,11 +272,11 @@ pm_exec = PatternMatcher([ if getenv("HCQ2"): from tinygrad.runtime.support.hcq2 import hcq_compile, hcq_link # noqa: E402 # down here, hcq2 imports the helpers above -def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:list[UOp]|None=None) -> UOp: +def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:list[UOp]|None=None, profile:bool|None=None) -> UOp: if validate: linear = graph_rewrite(linear, pm_validate, name="validate", walk=True) if (beam_val:=BEAM.value if beam is None else beam) >= 1: linear = graph_rewrite(linear, pm_beam, ctx=beam_val, walk=True) linear = graph_rewrite(linear, pm_compile, name="precompile kernels", walk=True) - if getenv("HCQ2"): linear = hcq_compile(linear, input_uops) + if getenv("HCQ2"): linear = hcq_compile(linear, input_uops, bool(PROFILE) if profile is None else profile) return graph_rewrite(linear, pm_optimize_local_size, name="optimize local size", walk=True) def link_linear(linear:UOp, cache=True) -> UOp: return hcq_link(linear, cache=cache) if getenv("HCQ2") else linear @@ -287,5 +294,5 @@ def time_call(call:UOp, var_vals:dict[str, int]|None=None, timeout:int|None=None from tinygrad.tensor import Tensor with Context(DEBUG=0, BEAM=0, CAPTURING=0, TRACK_MATCH_STATS=0): Tensor.ones(1024, 1024).contiguous().realize(do_update_stats=False) ctx = ExecContext(var_vals or {}, update_stats=False, wait=True, timeout=timeout, cache=False) - linear = link_linear(compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0), cache=ctx.cache) + linear = link_linear(compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0, profile=True), cache=ctx.cache) return max(pm_exec.rewrite(c, ctx) or 0.0 for c in linear.src) diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index 86d0136b18..fab620deaa 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -251,7 +251,7 @@ DEFAULT_FLOAT, DEFAULT_INT = ContextVar("DEFAULT_FLOAT", "float32"), ContextVar( CAPTURE_PROCESS_REPLAY = ContextVar("CAPTURE_PROCESS_REPLAY", 0) def _get_cpu_count() -> int: # os.process_cpu_count (3.13+) respects cgroup limits - if hasattr(os, "process_cpu_count"): return max(1, os.process_cpu_count()) + if hasattr(os, "process_cpu_count"): return max(1, os.process_cpu_count() or 1) # cgroup v2 (containers with --cpus=N) try: with open("/sys/fs/cgroup/cpu.max") as f: diff --git a/tinygrad/llm/chat.html b/tinygrad/llm/chat.html index 2ea21ff279..c5eb80ce76 100644 --- a/tinygrad/llm/chat.html +++ b/tinygrad/llm/chat.html @@ -21,7 +21,7 @@ const d = document.createElement('div'); d.className = 'msg'; chat.appendChild(d); const r = await fetch('/v1/chat/completions', {method: 'POST', headers: {'Content-Type': 'application/json'}, body: JSON.stringify({model: 'llama', messages: msgs, stream: true, temperature: 0.7})}); - let buf = ''; + let buf = '', txt = '', rsn = ''; for (const rd = r.body.getReader(), dec = new TextDecoder();;) { const {done, value} = await rd.read(); if (done) break; @@ -30,9 +30,13 @@ buf = lines.pop(); for (const ln of lines) if (ln.startsWith('data: ') && !ln.includes('[DONE]')) - try { d.textContent += JSON.parse(ln.slice(6)).choices[0]?.delta?.content || '' } catch {} + try { const dl = JSON.parse(ln.slice(6)).choices[0]?.delta; + if (dl?.reasoning_content) { const s = document.createElement('span'); s.style.color = '#888'; + s.textContent = dl.reasoning_content; rsn += dl.reasoning_content; d.appendChild(s) } + if (dl?.content) { const s = document.createElement('span'); + s.textContent = dl.content; txt += dl.content; d.appendChild(s) } } catch {} chat.scrollTop = chat.scrollHeight; } - msgs.push({role: 'assistant', content: d.textContent}); + const m = {role:'assistant', content:txt}; if (rsn) m.reasoning_content = rsn; msgs.push(m); } diff --git a/tinygrad/mixin/dtype.py b/tinygrad/mixin/dtype.py index 08d112da74..b46bac675b 100644 --- a/tinygrad/mixin/dtype.py +++ b/tinygrad/mixin/dtype.py @@ -30,7 +30,7 @@ class DTypeMixin: print(t.dtype, t.numpy()) ``` """ - return self if self.dtype == (dt:=to_dtype(dtype)) or self._uop.base.is_invalid else self._wrap_uop(self._uop.alu(Ops.CAST, arg=dt)) + return self if self.dtype == (dt:=to_dtype(dtype)) else self._wrap_uop(self._uop.alu(Ops.CAST, arg=dt)) def bitcast(self, dtype:DTypeLike) -> Self: """ diff --git a/tinygrad/mixin/elementwise.py b/tinygrad/mixin/elementwise.py index b6f19facfd..f06fd5aa52 100644 --- a/tinygrad/mixin/elementwise.py +++ b/tinygrad/mixin/elementwise.py @@ -24,6 +24,7 @@ class ElementwiseMixin(CreationMixin): out_dtype = least_upper_dtype(x.dtype, y.dtype) # keep weak CONST weak, might lift weakint -> weakfloat def promote(t): + if t._uop.base.is_invalid: return t # invalid bool is weak const if t.dtype in dtypes.weaks and t._uop.base.op is Ops.CONST: return t._wrap_uop(t._uop.const_like(t._uop.base.val, weak_dtype(out_dtype))) return t.cast(out_dtype) return promote(x), promote(y) @@ -220,7 +221,7 @@ class ElementwiseMixin(CreationMixin): if dtypes.is_int(a.dtype) and dtypes.is_int(b.dtype): return a.alu(Ops.CMOD, b) return a - a.div(b, rounding_mode="trunc") * b - def div(self, x: Self | ConstType, reverse: bool = False, rounding_mode: Literal["trunc", "floor"] | None = None) -> Self: + def div(self, x: 'Self|ConstType|UOp', reverse: bool = False, rounding_mode: Literal["trunc", "floor"] | None = None) -> Self: """ Divides `self` by `x`. Equivalent to `self / x`. diff --git a/tinygrad/mixin/gradient.py b/tinygrad/mixin/gradient.py index 93cc63843b..14fa98c00b 100644 --- a/tinygrad/mixin/gradient.py +++ b/tinygrad/mixin/gradient.py @@ -7,7 +7,11 @@ from tinygrad.dtype import sum_acc_dtype def reduce_gradient(ctx:UOp, ret:UOp, op:Ops): if op == Ops.ADD: return (ctx._broadcast_to(ret.src[0].shape),) if op == Ops.MAX: return (((mask:=ret.src[0].eq(ret).cast(ctx.dtype))/mask._rop(Ops.ADD, tuple(range(ret.arg[1])))) * ctx,) - if op == Ops.MUL: return (ctx * ret / ret.src[0],) + if op == Ops.MUL: + # d(prod x)/dx_j = prod_{i!=j} x_i: ret/x_j whenever x_j != 0 (any zero makes ret 0), else the product of the others + safe_x, axes = (is_zero:=(x:=ret.src[0]).eq(0)).where(1, x), tuple(range(ret.arg[1])) + zero_count = is_zero.cast(sum_acc_dtype(is_zero.dtype))._rop(Ops.ADD, axes) + return (ctx * is_zero.where(zero_count.eq(1).where(safe_x._rop(Ops.MUL, axes), 0), ret/safe_x),) def _compact_params(body:UOp, all_args:tuple[UOp, ...]) -> tuple[UOp, tuple[UOp, ...]]: """Remove unused PARAMs from body and return compacted (body, args).""" diff --git a/tinygrad/mixin/movement.py b/tinygrad/mixin/movement.py index 14632c3804..b916f8af9d 100644 --- a/tinygrad/mixin/movement.py +++ b/tinygrad/mixin/movement.py @@ -90,8 +90,11 @@ class MovementMixin: if resolve(index.step == 0, False): raise ValueError(f"{index=} cannot have 0 as step") start, stop = 0 if index.start is None else index.start, size if index.stop is None else index.stop step = 1 if index.step is None else index.step + # resolve negative int bounds against the (possibly symbolic) size, like slice.indices + if isinstance(start, int) and start < 0: start = start + size + if isinstance(stop, int) and stop < 0: stop = stop + size if all_int((start, stop, step)): - # handle int slicing (resolve negative bounds, clamp, stride) + # handle int slicing (clamp, stride) *bound, stride = index.indices(int(size.vmax) if isinstance(size, UOp) else size) bound = [0, 0] if stride * (bound[1] - bound[0]) < 0 else ([bound[1]+1, bound[0]+1] if stride < 0 else bound) return {"size":ceildiv(bound[1]-bound[0], abs(stride)), "boundary":tuple(bound), "stride":stride, "collapse_dim":False} @@ -265,7 +268,8 @@ class MovementMixin: return self.shrink(tuple([None if ns is None else (0, ns) for ns in argfix(shape, *args)])) def pad_to(self, shape, *args) -> Self: - return self._mop(Ops.PAD, tuple((0, s if ns is None else ns) for s,ns in zip(self.shape, argfix(shape, *args), strict=True))) + ret = self._mop(Ops.PAD, tuple((0, s if ns is None else ns) for s,ns in zip(self.shape, argfix(shape, *args), strict=True))) + return self if ret.shape == self.shape else ret def view(self, shape, *args) -> Self: """`.view` is an alias for `.reshape`.""" diff --git a/tinygrad/mixin/op.py b/tinygrad/mixin/op.py index b501faac9a..15c9c0d1ed 100644 --- a/tinygrad/mixin/op.py +++ b/tinygrad/mixin/op.py @@ -514,7 +514,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin): output_dtype = self.dtype if dtypes.is_float(self.dtype) else dtypes.float32 numerator = self.cast(sum_acc_dtype(self.dtype)).sum(axis=axis, keepdim=keepdim) denominator = prod([si for si, so in zip(self.shape, self.sum(axis=axis, keepdim=True).shape) if resolve(si != so)]) - return numerator.div(denominator).cast(output_dtype) # type: ignore[arg-type] + return numerator.div(denominator).cast(output_dtype) def var(self, axis:int|Sequence[int]|None=None, keepdim=False, correction=1) -> Self: """ @@ -538,12 +538,11 @@ class OpMixin(ElementwiseMixin, ReduceMixin): print(t.var(axis=1).numpy()) ``` """ + output_dtype = self.dtype if dtypes.is_float(self.dtype) else dtypes.float32 squares = (self - self.mean(axis=axis, keepdim=True)).square() n = prod([si for si, so in zip(self.shape, squares.sum(axis=axis, keepdim=True).shape) if resolve(si != so)]) - reduced = squares.sum(axis=axis, keepdim=keepdim) - denominator = reduced.const_like(n) - correction # type: ignore[arg-type] - # TODO: remove relu? - return reduced.div(denominator.relu()) + numerator = squares.cast(sum_acc_dtype(self.dtype)).sum(axis=axis, keepdim=keepdim) + return numerator.div(smax(n - correction, 0)).cast(output_dtype) def var_mean(self, axis:int|Sequence[int]|None=None, keepdim=False, correction=1) -> tuple[Self, Self]: """ @@ -1057,14 +1056,16 @@ class OpMixin(ElementwiseMixin, ReduceMixin): assert not (align_corners and mode != "linear"), "align_corners option can only be set with the interpolating mode linear" x, expand = self, list(self.shape) for i in range(-1,-len(size)-1,-1): - scale = (int(self.shape[i]) - int(align_corners)) / (size[i] - int(align_corners)) - arr, reshape = type(self).arange(size[i], dtype=dtypes.float32), [1] * self.ndim + in_sz, reshape = int(self.shape[i]), [1] * self.ndim reshape[i] = expand[i] = size[i] if mode == "linear": - index = (scale*arr if align_corners else (scale*(arr+0.5))-0.5).clip(0, self.shape[i]-1) - low, high, perc = [y.reshape(reshape).expand(expand) for y in (index.floor().int(), index.ceil().int(), index - index.floor())] + arr = type(self).arange(size[i]) + num, den = (arr*(in_sz-1), size[i]-1) if align_corners else ((arr*2+1)*in_sz - size[i], size[i]*2) + num = num.clip(0, (in_sz-1)*den) + low, high, perc = [y.reshape(reshape).expand(expand) for y in (num//den, (num+den-1)//den, (num % den).cast(dtypes.float32)/den)] x = x.gather(i, low).lerp(x.gather(i, high), perc) else: + scale, arr = in_sz / size[i], type(self).arange(size[i], dtype=dtypes.float32) index = (scale*(arr+0.5) if mode=="nearest-exact" else scale*arr).cast(dtypes.int32).reshape(reshape).expand(expand) x = x.gather(i, index) return x.cast(self.dtype) diff --git a/tinygrad/renderer/llvmir.py b/tinygrad/renderer/llvmir.py index 56b3b38027..99da4d1345 100644 --- a/tinygrad/renderer/llvmir.py +++ b/tinygrad/renderer/llvmir.py @@ -35,21 +35,36 @@ def lcast(input_type:DType, output_type:DType): if dtypes.is_int(output_type): return 'trunc' if output_type.itemsize < input_type.itemsize else 'sext' raise NotImplementedError(f"cast from {input_type} -> {output_type} not implemented") -def render_wmma_amd(ctx, wmma: UOp, cdna=False) -> str: +def render_wmma_amd(ctx, wmma: UOp, cdna=False, rdna4=False) -> str: dt_map = {dtypes.half: "f16", dtypes.float: "f32", dtypes.ushort: "bf16.1k" if cdna else "bf16", dtypes.bfloat16: "bf16.1k" if cdna else "bf16", dtypes.fp8e4m3: ".fp8.fp8", dtypes.fp8e5m2: ".bf8.bf8", dtypes.int8: "iu8", dtypes.int32: "i32"} # https://github.com/llvm/llvm-project/blob/main/clang/test/CodeGenOpenCL/builtins-amdgcn-mfma.cl N,M,K = wmma.arg[0] if cdna: if K == 32: dt_map.update({dtypes.half: ".f16", dtypes.bfloat16: ".bf16"}) - return f" {ctx[wmma]} = call {ldt(wmma.dtype, wmma.max_numel())} @llvm.amdgcn.mfma.{dt_map[wmma.src[-1].dtype]}" + \ - f".{N}x{M}x{K}{dt_map[wmma.arg[1]]}(" + ", ".join([f"{ldt(w.dtype, w.max_numel())} {ctx[w]}" for w in wmma.src]) + ", i32 0, i32 0, i32 0)" + scaled = K == 128 + args = [f"{ldt(w.dtype, w.max_numel())} {ctx[w]}" for w in wmma.src] + # scaled mfma call require E8M0 scale args, byte = 0x7F = 127, scale = 2^(127 - 127) = 1.0 + if scaled: + _fmt = { dtypes.fp8e5m2:1, dtypes.fp8e4m3:0 } + # (a_fp8_fmt, b_fp8_fmt, opsel, scale_a, opsel, scale_b) + args.extend([f"i32 {_fmt[wmma.arg[1]]}", f"i32 {_fmt[wmma.arg[1]]}", "i32 0", "i32 127", "i32 0", "i32 127"]) + else: args.extend(["i32 0", "i32 0", "i32 0"]) # (cbsz, blgp, ?) + + scale = "scale." if scaled else "" + dt_in = dt_map[wmma.arg[1]] if not scaled else ".f8f6f4" + return f" {ctx[wmma]} = call {ldt(wmma.dtype, wmma.max_numel())} @llvm.amdgcn.mfma.{scale}{dt_map[wmma.src[-1].dtype]}" + \ + f".{N}x{M}x{K}{dt_in}(" + ", ".join(args) + ")" # https://github.com/llvm/llvm-project/blob/main/llvm/test/CodeGen/AMDGPU/GlobalISel/llvm.amdgcn.wmma_32.ll # example: %wmma0 = call <8 x float> @llvm.amdgcn.wmma.f32.16x16x16.f16(<16 x half> %v99,<16 x half> %v100,<8 x float> %v101) args = [f"{ldt(w.dtype, w.max_numel())} {ctx[w]}" for w in wmma.src] if wmma.arg[1] == dtypes.int8: args = ["i1 true", args[0], "i1 true", args[1], args[2]] # iu8 flags A/B signed - return f" {ctx[wmma]} = call {ldt(wmma.dtype, wmma.max_numel())} @llvm.amdgcn.wmma.{dt_map[wmma.src[-1].dtype]}.16x16x16." + \ - f"{dt_map[wmma.arg[1]]}(" + ", ".join(args) + (", i1 false)" if wmma.dtype != dtypes.float else ")") + if wmma.dtype != dtypes.float: args.append("i1 false") # opsel + def _bf16(dt:DType): return dtypes.ushort if dt is dtypes.bfloat16 else dt + suffix = f".v{wmma.max_numel()}{dt_map[_bf16(wmma.dtype)]}.v{wmma.src[0].max_numel()}{dt_map[_bf16(wmma.arg[1])]}" if rdna4 else "" + # bfloat treated as i16 in LLVM call + return f" {ctx[wmma]} = call {ldt(_bf16(wmma.dtype), wmma.max_numel())} @llvm.amdgcn.wmma.{dt_map[wmma.src[-1].dtype]}.16x16x16." + \ + f"{dt_map[wmma.arg[1]]}{suffix}(" + ", ".join(args) + ")" # llvm ops, lop[][] unsigned_lop = { Ops.ADD: "add", Ops.MUL: "mul", Ops.CDIV: "udiv", Ops.CMOD: "urem", @@ -254,13 +269,21 @@ exit: %packed = phi i32 [%packed_bf8, %do_bf8], [%packed_fp8, %do_fp8]\n %trunc attributes = ["alwaysinline", "nounwind", '"no-builtins"', f'"amdgpu-flat-work-group-size"="1,{requiredMaxThreadsPerBlock}"', '"no-trapping-math"="true"'] return 'attributes #0 = { ' + ' '.join(attributes) + ' }' + @staticmethod + def is_rdna4(arch): return arch.split(':')[0] in {'gfx1200', 'gfx1201'} def __init__(self, target:Target): super().__init__(target) from tinygrad.runtime.support.compiler_llvm import AMDLLVMCompiler self.compiler, self.tensor_cores, self.is_cdna = AMDLLVMCompiler(target.arch), tc.get_amd(target.arch), HIPRenderer.is_cdna(target.arch) - self.string_rewrite += PatternMatcher([(UPat(Ops.WMMA, name="wmma"), lambda ctx, wmma, cdna=self.is_cdna: render_wmma_amd(ctx, wmma, cdna))]) + self.string_rewrite += PatternMatcher([ + (UPat(Ops.WMMA, name="wmma"), lambda ctx, wmma, rdna4=AMDLLVMRenderer.is_rdna4(target.arch), cdna=self.is_cdna: + render_wmma_amd(ctx, wmma, cdna, rdna4)) + ]) if self.is_cdna: self.extra_matcher += PatternMatcher([ + (UPat(Ops.WMMA, name="x", dtype=dtypes.float), + lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint32), x.src[1].bitcast(dtypes.uint32), x.src[2])) + if x.arg[0][2] == 128 and x.src[0].dtype.itemsize <= 8 else None), (UPat(Ops.WMMA, name="x", dtype=dtypes.float), lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2])) if x.max_numel() == 4 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 4 else None), @@ -274,9 +297,10 @@ exit: %packed = phi i32 [%packed_bf8, %do_bf8], [%packed_fp8, %do_fp8]\n %trunc src=(x.src[0].bitcast(dtypes.uint32), x.src[1].bitcast(dtypes.uint32), x.src[2])) if x.src[0].dtype == dtypes.int8 and x.src[0].max_numel() == 16 else None), (UPat(Ops.WMMA, name="x", dtype=dtypes.half), lambda x: UOp(Ops.STACK, src=tuple(x.replace( - src=(x.src[0], x.src[1], UOp(Ops.STACK, src=tuple(x.src[2].index(j//2) if j%2 == 0 else UOp.const(0.0, x.src[2].dtype) + src=(x.src[0], x.src[1], UOp(Ops.STACK, src=tuple(x.src[2].index(UOp.const(j//2, dtypes.int16)) + if j%2 == 0 else UOp.const(0.0, x.src[2].dtype) for j in range(x.max_numel()*2)))), - arg=(*x.arg[:4], None)).index(i*2) + arg=(*x.arg[:4], None)).index(UOp.const(i*2, dtypes.int16)) for i in range(x.max_numel()))) if x.max_numel() == 8 else None), (UPat(Ops.WMMA, name="x"), lambda x: x.replace( src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2])) @@ -285,6 +309,7 @@ exit: %packed = phi i32 [%packed_bf8, %do_bf8], [%packed_fp8, %do_fp8]\n %trunc if target.arch in {"gfx1200", "gfx1201"}: self.extra_matcher += PatternMatcher([ (UPat(Ops.WMMA, name="x", dtype=dtypes.bfloat16), lambda x: x.replace( + dtype=dtypes.uint16, src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2].bitcast(dtypes.uint16))) .bitcast(dtypes.bfloat16) if x.max_numel() == 8 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 8 else None), (UPat(Ops.WMMA, name="x", dtype=dtypes.float), diff --git a/tinygrad/runtime/ops_cpu.py b/tinygrad/runtime/ops_cpu.py index c74363e659..3ecb49a54e 100644 --- a/tinygrad/runtime/ops_cpu.py +++ b/tinygrad/runtime/ops_cpu.py @@ -50,7 +50,7 @@ def worker_prog(): # spin on windows, sem_wait to sleep on posix if WIN: ready = (v:=wait.after(lw:=UOp.loop(1), cur)[0].load()).end(lw, v <= cur) - else: ready = wait.after(cur)[0].load().call(sem.after(cur)[0], ret_dtype=dtypes.void) + else: ready = (rv:=wait.after(lw:=UOp.loop(1), cur)[0].load().call(sem.after(cur)[0], ret_dtype=dtypes.int)).end(lw, rv != 0) entry = [ring.after(ready).index((cur % RING_SLOTS) * CMD_SIZE + i).load() for i in range(CMD_SIZE)] return entry[0].call(*entry[1:], ret_dtype=dtypes.void).end(cur) @@ -167,10 +167,11 @@ class CPUDevice(HCQCompiled): (UPat(Ops.PARAM, tag="sentinel_signal"), lambda ctx: ctx[0].signal("sentinel", (1 << 64) - 1)), (UPat(Ops.PARAM, tag="timeline_signal"), lambda ctx: ctx[0].signal("timeline")), (UPat(Ops.PARAM, tag="timeline_value"), lambda ctx: ctx[0].signal("value", 1)), + (UPat(Ops.PARAM, tag="signal", name="b"), lambda ctx, b: ctx[0].signal(b.arg.slot)), ]) @functools.cache - def signal(self, name:str, init_value:int=0) -> Buffer: + def signal(self, name:str|int, init_value:int=0) -> Buffer: (buf:=Buffer(self.device, 1, dtypes.uint64, preallocate=True)).as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] = init_value return buf diff --git a/tinygrad/runtime/support/compiler_amd.py b/tinygrad/runtime/support/compiler_amd.py index 7df98d7a80..eef5e07f88 100644 --- a/tinygrad/runtime/support/compiler_amd.py +++ b/tinygrad/runtime/support/compiler_amd.py @@ -89,9 +89,11 @@ class HIPCompiler(Compiler): class HIPCCCompiler(Compiler): def __init__(self, arch:str, extra_options:list[str]=[]): - self.arch, self.extra_options = arch, extra_options - super().__init__(f"compile_hipcc_{self.arch}_{hashlib.sha256(' '.join(extra_options).encode()).hexdigest()[:8]}") + self.arch, self.extra_options, self.no_hipcc = arch, extra_options, getenv("NO_HIPCC") + super().__init__(f"compile_hipcc_{self.arch}_{hashlib.sha256(' '.join(extra_options).encode()).hexdigest()[:8]}"+ + ("_nohipcc" if self.no_hipcc else "")) def compile(self, src:str) -> bytes: + if self.no_hipcc: return b"" with tempfile.NamedTemporaryFile(suffix=".cpp") as srcf, tempfile.NamedTemporaryFile(suffix=".bc") as bcf: with tempfile.NamedTemporaryFile(suffix=".hsaco") as libf: srcf.write(src.encode()) diff --git a/tinygrad/runtime/support/hcq2.py b/tinygrad/runtime/support/hcq2.py index e01dab947f..9b0372d035 100644 --- a/tinygrad/runtime/support/hcq2.py +++ b/tinygrad/runtime/support/hcq2.py @@ -1,12 +1,13 @@ from __future__ import annotations from typing import cast, Callable, TypeVar, Generic, Any, Sequence -import struct, functools, time, collections, itertools +import struct, functools, time, collections, itertools, decimal, statistics from dataclasses import replace, dataclass -from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, dedup, pluralize, JIT_BATCH_SIZE, unwrap -from tinygrad.helpers import to_tuple, round_up, partition, data64_le, panic, ContextVar +from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, dedup, pluralize, JIT_BATCH_SIZE, unwrap, PROFILE +from tinygrad.helpers import to_tuple, round_up, partition, data64_le, panic, ContextVar, perf_counter_us, Context from tinygrad.device import Device, Buffer, BufferSpec, Compiled, LRUAllocator, MultiBuffer, DepsTracker -from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, graph_rewrite, track_rewrites, GroupOp -from tinygrad.uop.symbolic import symbolic +from tinygrad.device import ProfileDeviceEvent, ProfileGraphEntry, ProfileGraphEvent +from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, graph_rewrite, rewrite_group, GroupOp +from tinygrad.uop.symbolic import symbolic, pm_fold_cast_const from tinygrad.dtype import dtypes, truncate from tinygrad.runtime.support.hcq import MMIOInterface from tinygrad.runtime.support.memory import BumpAllocator @@ -27,13 +28,12 @@ HCQ_CACHE_TAGS = frozenset(("program", "systems", "template")) @dataclass(frozen=True) class HCQInfo: - name:str - estimates:Estimates device:tuple[str, ...] - queue:str + estimates:Estimates = Estimates() input_idxs:tuple[int, ...] = () # indexes into input_uops used by this call inputs:int|None = None + prof:tuple[ProfileGraphEntry, ...] = () # st_id/en_id are timestamp signal slots until collect def all_devices_in(d:Any, c:frozenset[str]) -> bool: return {x.split(":")[0] for x in to_tuple(d)} <= c @@ -73,6 +73,8 @@ def make_submit(*cmds, devs:str|tuple[str, ...], queue:str) -> UOp: return UOp.custom_function("submit_cmdbuf", UOp(Ops.LINEAR, src=tuple(cmds), arg=(to_tuple(devs), queue))) def get_submit(ast:UOp) -> UOp: return next(u for u in ast.toposort() if u.op is Ops.CUSTOM_FUNCTION and u.arg == "submit_cmdbuf") +def make_call(name:str, body:UOp, info:HCQInfo) -> UOp: return UOp.custom_function("hcq", body).call(name=name, aux=info) + def encode_kernargs_clike(call:UOp, prg:UOp, devs:str|tuple[str, ...]) -> UOp: data, info = prg.arg buf = UOp.placeholder((data.kernargs_alloc_size // 4,), dtypes.uint32, next(UOp.unique_num), device=devs).rtag("kernargs") @@ -94,12 +96,24 @@ pm_replace_buffers = PatternMatcher([(UPat(Ops.CALL, name="call"), replace_call_ def _need_staging(a, b): return all_devices_in(a.device, HCQ_DEVS) and not all_devices_in(b.device, HCQ_P2P_DEVS) +def hcq_call_devs(call:UOp) -> Any|None: return next((b.device for b in call.src[1:] if all_devices_in(b.device, HCQ_DEVS)), None) + def stage_copy(dst:UOp, src:UOp) -> UOp|None: if not (_need_staging(src, dst) or _need_staging(dst, src)): return None stage = UOp.new_buffer("CPU", src.max_numel() * src.dtype.itemsize, dtypes.uint8) return UOp(Ops.LINEAR, src=(src.copy_to_device("CPU").call(stage, src), stage.copy_to_device(dst.device).call(dst, stage))) -pm_insert_copy_staging = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), stage_copy)]) + +def kernel_copy(call:UOp, dst:UOp, src:UOp) -> UOp|None: + if (devs:=hcq_call_devs(call)) is None or Device[(dev:=to_tuple(devs)[0])].has_copy_queue: return None + d, s = (UOp.param(i, dst.dtype, (n:=dst.max_numel(),), device=devs) for i in range(2)) + ast = d.index(r:=UOp.range(n, 0)).store(s.index(r).load()).end(r).sink(arg=KernelInfo(name="copy"), tag=1) + return call.replace(src=(to_program(ast, Device[dev].renderer), dst, src)) + +pm_insert_copy_staging = PatternMatcher([ + (UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), stage_copy), + (UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src")), name="call"), kernel_copy) +]) # ***************** # 2. deps @@ -137,12 +151,6 @@ def _build_wait_cmds(slots:dict[str, int], dep_lanes:list[tuple[tuple, int, int] waits.append(UOp(Ops.INS, arg="wait", src=(sig, UOp.const(dtag + 1, dtypes.uint64)))) return waits, {dtag for _, _, dtag in deps} -def make_fence(timeline:UOp, prev:UOp, sigs:list[UOp]) -> UOp: - free = (cur:=timeline.after(loop:=UOp.loop(0)).index(0).load()).end(loop, cur < prev.index(0).load()) - return UOp.sink(*[s.after(free).index(0).store(0) for s in sigs]) - -def _hcq_call(devs, name:str, body:UOp) -> UOp: return UOp.custom_function("hcq", body).call(aux=HCQInfo(name, Estimates(), devs, "COMPUTE:0")) - def _build_finalizers(batch:list[tuple[UOp, tuple[str, ...]]], batch_info:list[tuple[tuple[str, ...], str]], tracker:HCQDepsTracker, slots:dict[str, int]) -> tuple[list[UOp], list[UOp], set[int]]: # collect all buffers which belong to devices @@ -151,48 +159,50 @@ def _build_finalizers(batch:list[tuple[UOp, tuple[str, ...]]], batch_info:list[t for b in itertools.chain.from_iterable(_get_call_bufs_by_lane(call, devices)): for bd in to_tuple(b.device): dev_bufs[bd][id(b)] = b - n, fences, fins, waited = len(batch_info), [], [], set() + n, fences, fins, signal_tags = len(batch_info), [], [], set() for _, devgroup in itertools.groupby(sorted(dev_bufs), key=lambda d: d.split(":")[0]): devs = tuple(devgroup) # to finalize the batch, sync all accesses from other devices to buffers that belong to this device fin_deps = [dl for dl in _get_deps(tracker, [list(dev_bufs[d].values()) for d in devs], None, key=(devs, "COMPUTE:0", n)) if dl[0][2] < n] - waits, cur_waited = _build_wait_cmds(slots, fin_deps, devs, "COMPUTE:0") - waited |= cur_waited + waits, cur_signal_tags = _build_wait_cmds(slots, fin_deps, devs, "COMPUTE:0") + signal_tags |= cur_signal_tags # wait the syncs and signal the device epoch, then bump the timeline on the host - timeline, tl = make_signal(devs, tag="timeline_signal"), make_signal(devs, tag="timeline_value") - submit = make_submit(*waits, UOp(Ops.INS, arg="store", src=(timeline, tl.index(0))), devs=devs, queue="COMPUTE:0") - cur = (bump:=tl.after(submit).index(0)).load() - bumps = [bump.store(cur + 1)] + tl_signal, tl_value = make_signal(devs, tag="timeline_signal"), make_signal(devs, tag="timeline_value") + fin_submit = make_submit(*waits, UOp(Ops.INS, arg="store", src=(tl_signal, tl_value.index(0))), devs=devs, queue="COMPUTE:0") + epoch = (epoch_slot:=tl_value.after(fin_submit).index(0)).load() - # devices running the batch reset their queue signals before each run, fencing on the epoch kept from the previous one - if qs:=dedup([qn for bdevs, qn in batch_info if set(bdevs) & set(devs)]): - prev = make_signal(devs, next(UOp.unique_num)) - fences.append(_hcq_call(devs, "hcq_fence", make_fence(timeline, prev, [make_signal(devs, slots[q]) for q in qs]))) - bumps.append(prev.after(submit).index(0).store(cur)) - fins.append(_hcq_call(devs, "hcq_finalizer", UOp.sink(*bumps))) - return fences, fins, waited + # fence once per device group on this schedule's previous epoch, then reset any queue signals used by the group + qs = dedup([qn for bdevs, qn in batch_info if set(bdevs) & set(devs)]) + sched_epoch = make_signal(devs, next(UOp.unique_num)) -def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]]) -> list[UOp]: + wait_device_epoch = (done:=tl_signal.after(loop:=UOp.loop(0)).index(0).load()).end(loop, done < sched_epoch.index(0).load()) + resets = [make_signal(devs, slots[q]).after(wait_device_epoch).index(0).store(0) for q in qs] + + fences.append(make_call("hcq_fence", UOp.sink(*(resets or [wait_device_epoch])), HCQInfo(devs))) + fins.append(make_call("hcq_finalizer", UOp.sink(epoch_slot.store(epoch + 1), sched_epoch.after(fin_submit).index(0).store(epoch)), HCQInfo(devs))) + return fences, fins, signal_tags + +def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]], profile:bool) -> list[UOp]: batch_info = [(devices, "COMPUTE:0" if call.src[0].op is Ops.PROGRAM else "COPY:0") for call, devices in batch] # schedule deps - waited:set[int] = set() + signal_tags:set[int] = set() slots:dict[str, int] = collections.defaultdict(lambda: next(UOp.unique_num)) deps_tracker = HCQDepsTracker() call_waits:list[list[UOp]] = [] for tag, ((call, _), (devices, queue)) in enumerate(zip(batch, batch_info)): deps = _get_deps(deps_tracker, _get_call_bufs_by_lane(call, devices), get_call_outs_ins(call)[0], key=(devices, queue, tag)) - cmds, cur_waited = _build_wait_cmds(slots, deps, devices, queue) + cmds, cur_signal_tags = _build_wait_cmds(slots, deps, devices, queue) call_waits.append(cmds) - waited |= cur_waited + signal_tags |= cur_signal_tags # build fences and finalizers - fences, finalizers, finalizer_waited = _build_finalizers(batch, batch_info, deps_tracker, slots) - waited |= finalizer_waited + fences, finalizers, finalizer_signal_tags = _build_finalizers(batch, batch_info, deps_tracker, slots) + signal_tags |= finalizer_signal_tags - src = [] + src, prof = [], [] for tag, ((call, _), (devices, queue), q) in enumerate(zip(batch, batch_info, call_waits)): # first queue use, sync prior device work with the device timeline if batch_info.index((devices, queue)) == tag: @@ -200,31 +210,38 @@ def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]]) -> list[UOp]: q = [UOp(Ops.INS, arg="barrier", src=()), UOp(Ops.INS, arg="wait", src=(make_signal(devices, tag="timeline_signal"), epoch))] + q # and make hcq call - info = HCQInfo(get_call_name(call, get_call_arg_uops(call)), estimate_uop(call), devices, queue) - q += [call.replace(arg=replace(call.arg, aux=info))] + name, info = get_call_name(call, get_call_arg_uops(call)), HCQInfo(devices, estimate_uop(call)) + ts_ids = [next(UOp.unique_num) for _ in range(2)] if profile else [] + prof += [ProfileGraphEntry(d, name, *ts_ids) for d in devices if ts_ids] + + ts_ins = [UOp(Ops.INS, arg="timestamp", src=(make_signal(devices, s),)) for s in ts_ids] + q += ts_ins[:1] + [call.replace(arg=replace(call.arg, aux=info))] + ts_ins[1:] # signal the queue if someone waits for us - if tag in waited: q += [UOp(Ops.INS, arg="store", src=(make_signal(devices, slots[queue]), UOp.const(tag + 1, dtypes.uint64)))] - src.append(UOp.custom_function("hcq", make_submit(*q, devs=devices, queue=queue).sink()).call(name="hcq", aux=info)) + if tag in signal_tags: q += [UOp(Ops.INS, arg="store", src=(make_signal(devices, slots[queue]), UOp.const(tag + 1, dtypes.uint64)))] + src.append(make_call(name, make_submit(*q, devs=devices, queue=queue).sink(), info)) + + # append batch timestamps to finalizers + finalizers = [f.replace(arg=replace(f.arg, aux=replace(a:=f.arg.aux, prof=tuple(e for e in prof if e.device in a.device)))) for f in finalizers] return fences + src + finalizers -def sched_hcq_batches(l:UOp) -> UOp: +def sched_hcq_batches(l:UOp, profile:bool) -> UOp: srcs:list[UOp] = [] batch:list[tuple[UOp, tuple[str, ...]]] = [] for call in l.src: - if (devs:=next((b.device for b in call.src[1:] if all_devices_in(b.device, HCQ_DEVS)), None)) is not None: batch.append((call, to_tuple(devs))) - else: srcs, batch = srcs + _finalize_batch(batch) + [call], [] - return l.replace(src=tuple(srcs + _finalize_batch(batch))) + if (devs:=hcq_call_devs(call)) is not None: batch.append((call, to_tuple(devs))) + else: srcs, batch = srcs + _finalize_batch(batch, profile) + [call], [] + return l.replace(src=tuple(srcs + _finalize_batch(batch, profile))) # ***************** # 3. merge into queues def _merged_hcq_call(calls:list[UOp]) -> UOp: # TODO: simplify? if len(calls) == 1: return calls[0] - info = replace(calls[0].arg.aux, name=f"submit {calls[0].arg.aux.queue} ({len(calls)})", - estimates=sum((c.arg.aux.estimates for c in calls), start=Estimates())) - cmds = [cmd for c in calls for cmd in get_submit(c).src[0].src] - return UOp.custom_function("hcq", make_submit(*cmds, devs=info.device, queue=info.queue).sink()).call(name="hcq", aux=info) + devs, queue = get_submit(calls[0]).src[0].arg + body = make_submit(*[cmd for c in calls for cmd in get_submit(c).src[0].src], devs=devs, queue=queue).sink() + return make_call(f"submit {queue} ({len(calls)})", body, + replace(calls[0].arg.aux, estimates=sum((c.arg.aux.estimates for c in calls), start=Estimates()))) def merge_queues(linear:UOp) -> UOp: new_src:list[UOp] = [] @@ -232,24 +249,25 @@ def merge_queues(linear:UOp) -> UOp: limits:dict[tuple[tuple[str, ...], str], int] = collections.defaultdict(lambda: JIT_BATCH_SIZE.value) for call in linear.src: - if not isinstance(info:=call.arg.aux, HCQInfo) or info.name.startswith("hcq_"): # non-hcq call, fence or finalizer: close all open queues + # non-hcq call, fence or finalizer: close all open queues + if not isinstance(call.arg.aux, HCQInfo) or (call.arg.name or "").startswith("hcq_"): new_src += [_merged_hcq_call(opened_qs.pop(k)) for k in list(opened_qs)] + [call] continue - if (old:=opened_qs.pop(key:=(info.device, info.queue), None)) is not None: + devs, queue = get_submit(call).src[0].arg + if (old:=opened_qs.pop(key:=(devs, queue), None)) is not None: if limits[key] and len(old) >= limits[key]: new_src, old, limits[key] = new_src + [_merged_hcq_call(old)], [], limits[key] * 2 new_rec = old + [call] else: # no such queue opened: close every open submit on this queue that shares a device, so submit order is kept - closing = [k for k in opened_qs if k[1] == info.queue and set(k[0]) & set(info.device)] + closing = [k for k in opened_qs if k[1] == queue and set(k[0]) & set(devs)] new_src += [_merged_hcq_call(opened_qs.pop(k)) for k in closing] new_rec = [call] - opened_qs[(info.device, info.queue)] = new_rec + opened_qs[(devs, queue)] = new_rec return linear.replace(src=tuple(new_src + [_merged_hcq_call(c) for c in opened_qs.values()])) -def schedule_and_merge(ctx:dict[UOp, UOp], linear:UOp) -> UOp: - return merge_queues(sched_hcq_batches(linear).substitute(ctx, walk=True, enter_calls=True)) -pm_schedule_and_merge = PatternMatcher([(UPat(Ops.LINEAR, name="linear"), schedule_and_merge)]) +pm_schedule_and_merge = PatternMatcher([(UPat(Ops.LINEAR, name="l"), + lambda ctx, l: merge_queues(sched_hcq_batches(l, ctx[1]).substitute(ctx[0], walk=True, enter_calls=True)))]) # ***************** # 4.2. hcq lowering: ops to ir @@ -285,21 +303,26 @@ def make_addr_table(call:UOp, gaddrs:list[UOp], name:str) -> tuple[UOp, dict[UOp fills = (table.after(*make_patches(table, [(i*table.dtype.itemsize, addr) for addr, i in slots.items()])),) if slots else () return table, reads, fills, {g:slots[bare[g]] for g in gaddrs} -def make_scatter_loop(patches:list[UOp], inputs_table:tuple, lt_patches:list[UOp]) -> dict[UOp, UOp]: - (table, _, _, slots), dst, data, subs = inputs_table, patches[0].buf_uop, [], {} - for p in patches: - words = [(off, val, get_getaddrs(val)) for off,val in zip(p.src[0].src[1].src, p.src[1].src)] - data += [off.val << 32 | slots[gaddrs[0]] for off,_,gaddrs in words if gaddrs][::2] - scalars = [(off.val*dst.dtype.itemsize, val) for off,val,gaddrs in words if not gaddrs] - subs[p] = UOp.group(*make_patches(dst, scalars)) if scalars else UOp(Ops.NOOP) +def is_bare_addr(val:UOp) -> bool: return val.op is Ops.CAST and val.src[0].op in (Ops.AND, Ops.SHR) and val.src[0].src[0].op is Ops.GETADDR - # plan entry: dst word offset << 32 | addr table slot - plan = UOp.placeholder((len(data),), dtypes.uint64, next(UOp.unique_num), device=dst.device).rtag("systems") - entry = plan.index(ridx:=UOp.range(len(data), next(UOp.unique_num), dtype=dtypes.int, src=(plan, dst))).load() - slot, widx = ((entry & 0xffffffff) % table.max_numel()).cast(dtypes.int), ((entry >> 32) % (dst.max_numel()-1)).cast(dtypes.int) # CHECK_OOB bounds - loop = UOp.group(*[dst.index(widx+i).store((table.index(slot).load() >> 32*i).cast(dtypes.uint32)) for i in range(2)]).end(ridx) - lt_patches.append(make_binary_patch(plan, struct.pack(f'<{len(data)}Q', *data))) - subs[patches[0]] = UOp.group(loop, subs[patches[0]]) +def make_scatter_loops(patches:list[UOp], inputs_table:tuple, lt_patches:list[UOp]) -> dict[UOp, UOp]: + table, _, _, slots = inputs_table + subs, by_dst = {}, collections.defaultdict(list) + for p in patches: by_dst[p.buf_uop].append(p) + for dst, patches in by_dst.items(): + data = [] + for p in patches: + words = [(off, val, get_getaddrs(val)) for off,val in zip(p.src[0].src[1].src, p.src[1].src)] + data += [(off.val, slots[gaddrs[0]]) for off,_,gaddrs in words if gaddrs][::2] + scalars = [(off.val*dst.dtype.itemsize, val) for off,val,gaddrs in words if not gaddrs] + subs[p] = UOp.group(*make_patches(dst, scalars)) if scalars else UOp(Ops.NOOP) + + word_table, slot_table = (UOp.placeholder((len(data),), dtypes.uint32, next(UOp.unique_num), device=dst.device).rtag("systems") for _ in range(2)) + ridx = UOp.range(len(data), next(UOp.unique_num), dtype=dtypes.int, src=(word_table, slot_table, dst)) + widx, slot = ((p.index(ridx).load() % bound).cast(dtypes.int) for p,bound in ((word_table, dst.max_numel()-1), (slot_table, table.max_numel()))) + loop = UOp.group(*[dst.index(widx+i).store((table.index(slot).load() >> 32*i).cast(dtypes.uint32)) for i in range(2)]).end(ridx) + lt_patches += [make_binary_patch(buf, struct.pack(f'<{len(data)}I', *vals)) for buf,vals in zip((word_table, slot_table), zip(*data))] + subs[patches[0]] = UOp.group(loop, subs[patches[0]]) return subs def is_input_addr(g:UOp) -> bool: return all(x.op is Ops.PARAM and x.tag is None for x in unwrap_mstack(g.buf_uop)) @@ -307,17 +330,22 @@ def is_input_addr(g:UOp) -> bool: return all(x.op is Ops.PARAM and x.tag is None def split_patches(call:UOp) -> UOp|None: rt_patches:list[UOp] = [] lt_patches:list[UOp] = [] - body = graph_rewrite(call.src[0], pm_trim_link_patches, ctx=(rt_patches, lt_patches), name=f"trim link-time patches ({call.arg.aux.name})") + body = graph_rewrite(call.src[0], pm_trim_link_patches, ctx=(rt_patches, lt_patches), name=f"trim link-time patches ({call.arg.name})") # split patches inputs, internals = partition(dedup(g for p in rt_patches for g in get_getaddrs(p)), is_input_addr) runtimes, systems = partition(internals, lambda g: any(x.tag in {"program", "kernargs", "cmdbuf"} for x in unwrap_mstack(g.buf_uop))) tables = [make_addr_table(call, gs, n) for gs,n in ((inputs, "inputs"), (runtimes, "runtime"), (systems, "systems"))] reads, fills = {k:v for _,r,_,_ in tables for k,v in r.items()}, [f for t in tables[1:] for f in t[2]] # inputs table is filled by exec - input_patches = [p for p in rt_patches if (gs:=get_getaddrs(p)) and all(map(is_input_addr, gs))] - scatter = make_scatter_loop(input_patches, tables[0], lt_patches) if input_patches else {} + input_patches = [p for p in rt_patches if (gs:=get_getaddrs(p)) and all(map(is_input_addr, gs)) + and all(is_bare_addr(v) for v in p.src[1].src if get_getaddrs(v))] + scatter = make_scatter_loops(input_patches, tables[0], lt_patches) body = body.substitute({p:p.substitute(scatter | reads) for p in rt_patches}) + if inputs: # fence inputs + fills.append((t:=tables[0][0]).after(make_binary_patch(t, bytes(t.max_numel() * 8)))) # zeroed at link, slot 0 is the host fence + body = body.replace(src=(UOp.sink(*body.src[0].src, t.after(*body.src[0].src).index(0).store(0)),)) # open it once consumed + lt_srcs = collections.defaultdict(list) for p in lt_patches: lt_srcs[p.buf_uop].append(p) return call.replace(src=(body, *call.src[1:], *[b.after(*ps) for b,ps in lt_srcs.items()], *fills), @@ -341,7 +369,7 @@ def replace_params(call:UOp) -> UOp|None: sub = {(b:=u.without_after): UOp.param(i, u.dtype, shape=b.shape, device=HCQ_RUNTIME_DEV.value, volatile=b.op is Ops.PARAM and b.arg.volatile) for i,u in enumerate(c_args)} | {v: v.replace(arg=replace(v.arg, slot=-1)) for v in variables if v.op is Ops.PARAM} - info = replace(call.arg.aux, inputs=next((i for i,u in enumerate(c_args) if u.tag == "inputs"), None)) + info = replace(call.arg.aux, inputs=next((i for i,u in enumerate(c_args) if u.without_after.tag == "inputs"), None)) return call.replace(src=(body.substitute(sub).replace(arg="hcq_args"), *c_args, *refhold), arg=replace(call.arg, aux=info)) # TODO: call.after(*refhold)? pm_replace_params = PatternMatcher([ @@ -388,27 +416,27 @@ def callify_hcq(call:UOp, cf:UOp) -> UOp: pm_callify_hcq = PatternMatcher([(UPat(Ops.CALL, src=( UPat(Ops.CUSTOM_FUNCTION, arg="hcq_args", src=(UPat(Ops.SINK),), name="cf"),), name="call", allow_any_len=True), callify_hcq)]) -hcq_compile_cache:dict[bytes, UOp] = {} +hcq_compile_cache:dict[tuple[bytes, bool], UOp] = {} -@track_rewrites(lambda linear,input_uops,ret: f"HCQ Compile {pluralize('Kernel', len(ret.src))}") -def hcq_compile(linear:UOp, input_uops:list[UOp]|None=None) -> UOp: +@rewrite_group(lambda linear,input_uops,profile,ret: f"HCQ Compile {pluralize('Kernel', len(ret.src))}") +def hcq_compile(linear:UOp, input_uops:list[UOp]|None, profile:bool) -> UOp: if input_uops is not None: slots = {u:i for i,u in reversed(tuple(enumerate(input_uops)))} linear = graph_rewrite(linear, pm_replace_buffers, ctx=(input_uops, slots), walk=True, name="replace buffer") - if (final_linear:=(hcq_compile_cache.get(cache_key:=linear.key))) is None: + if (final_linear:=(hcq_compile_cache.get(cache_key:=(linear.key, profile)))) is None: # prep linear = linear.substitute(back_map:={s.param_like(i): s for i,s in enumerate(input_uops)} if input_uops is not None else {}, walk=True) linear = graph_rewrite(linear, pm_insert_copy_staging+pm_flatten_linear, name="insert copy staging") # schedule - linear = graph_rewrite(linear, pm_schedule_and_merge, ctx={s:p for p,s in back_map.items()}, walk=True, name="schedule and merge hcq") + linear = graph_rewrite(linear, pm_schedule_and_merge, ctx=({s:p for p,s in back_map.items()}, profile), walk=True, name="schedule and merge hcq") # lowering to hcq ir linear = graph_rewrite(linear, pm_encode_cmdbufs+pm_pack_placeholders, walk=True, name="encode and pack", enter_calls=True) # patches and runtime uops - linear = graph_rewrite(linear, pm_early_simplify+symbolic, bottom_up=False, name="simplify patches", enter_calls=True) + linear = graph_rewrite(linear, pm_early_simplify+symbolic+pm_fold_cast_const, bottom_up=False, name="simplify patches", enter_calls=True) linear = graph_rewrite(linear, pm_split_patches, walk=True, name="split patches") # and compile it @@ -474,14 +502,14 @@ def link_buf_key(a:UOp): return a.key, to_tuple(a.device) link_buf_cache:dict[tuple[bytes, tuple[str, ...]], UOp] = {} link_linear_cache:dict[bytes, UOp] = {} -@track_rewrites(lambda _,cache,ret: f"HCQ Link {pluralize('Kernel', len(ret.src))}") +@rewrite_group(lambda _,cache,ret: f"HCQ Link {pluralize('Kernel', len(ret.src))}") def hcq_link(linear:UOp, cache=True) -> UOp: if (linked:=link_linear_cache.get(linear_key:=linear.key)) is not None: return linked bufs = {(j,i):a for j,c in enumerate(linear.src) for i,a in enumerate(c.src[1:], 1) if a.op is Ops.AFTER and unwrap_mstack(a.src[0])[0].tag in HCQ_CACHE_TAGS} linear = linear.substitute({x:link_buf_cache[k] for a in bufs.values() if (k:=link_buf_key(a)) in link_buf_cache for x in (a, a.src[0])}, walk=True) - linear = graph_rewrite(linear, pm_resolve_patches+symbolic+pm_assert_no_afters, bpm=pm_bufferize, ctx=cache, bottom_up=False, + linear = graph_rewrite(linear, pm_resolve_patches+symbolic+pm_fold_cast_const+pm_assert_no_afters, bpm=pm_bufferize, ctx=cache, bottom_up=False, name="resolve patches") for (j,i),a in bufs.items(): link_buf_cache.setdefault(link_buf_key(a), linear.src[j].src[i]) if cache: link_linear_cache[linear_key] = linear @@ -495,6 +523,7 @@ class HCQ2Compiled(Compiled): def __init__(self, device:str, allocator:HCQAllocator, compilers:list[type[Renderer]], runtime, can_recover:bool=False, arch=None): self.device_id:int = int(device.split(":")[1]) if ":" in device else 0 + self.can_recover = can_recover self.pm_bufferize = PatternMatcher([ (UPat(Ops.PARAM, tag="sentinel_signal"), lambda ctx: ctx[0].signal("sentinel", (1 << 64) - 1)), @@ -508,6 +537,27 @@ class HCQ2Compiled(Compiled): self.rt_buffer = Buffer(self.device, 64 << 20, dtypes.uint8, options=BufferSpec(uncached=True, cpu_access=True)) self.rt_allocator = BumpAllocator(64 << 20) + self.prof_ents:dict[int, ProfileGraphEntry] = {} + + def collect_prof(self): + if PROFILE: + es = list(self.prof_ents.values()) + sigs = [self.signal(i)._buf.cpu_view().view(fmt='Q')[0]/decimal.Decimal(self.timestamp_divider) for e in es for i in (e.st_id, e.en_id)] + Compiled.profile_events.append(ProfileGraphEvent([replace(e, st_id=2*i, en_id=2*i+1) for i,e in enumerate(es)], [], sigs)) + self.prof_ents.clear() + + def _at_profile_finalize(self): + from tinygrad.tensor import Tensor + tdiffs = [] + for _ in range(5): + with Context(DEBUG=0, BEAM=0, TRACK_MATCH_STATS=0): Tensor.ones(1, device=self.device).contiguous().realize() + if not (ents:=list(self.prof_ents.values())): return + self.prof_ents.clear() + st = perf_counter_us() + self.synchronize() + gpu = max(self.signal(e.en_id)._buf.cpu_view().view(fmt='Q')[0] for e in ents)/decimal.Decimal(self.timestamp_divider) + tdiffs.append((st+perf_counter_us())/2 - gpu) + Compiled.profile_events.append(ProfileDeviceEvent(self.device, statistics.median(tdiffs), self.device_props())) def new_buffer(self, b:UOp, cache:bool) -> Buffer: if cache or b.tag in HCQ_CACHE_TAGS: @@ -521,12 +571,15 @@ class HCQ2Compiled(Compiled): return buf def synchronize(self, timeout:int|None=None): - if not hasattr(self, 'iface'): return + if HCQ_RUNTIME_DEV.value != self.device: Device[HCQ_RUNTIME_DEV.value].synchronize() + sig = self.signal("timeline").as_memoryview(force_zero_copy=True, no_sync=True).cast('Q') tl = self.signal("value", 1).as_memoryview(force_zero_copy=True, no_sync=True).cast('Q') + timeout = timeout if timeout is not None and self.can_recover else None st = time.perf_counter() while sig[0] < tl[0] - 1: if time.perf_counter() - st > (timeout or 3000) / 1000: self.on_device_hang() + if self.prof_ents: self.collect_prof() def on_device_hang(self): raise RuntimeError(f"{self.device} hang detected") diff --git a/tinygrad/schedule/__init__.py b/tinygrad/schedule/__init__.py index 99d228945f..a98e97b7c4 100644 --- a/tinygrad/schedule/__init__.py +++ b/tinygrad/schedule/__init__.py @@ -1,6 +1,6 @@ import time, inspect from collections import deque -from tinygrad.uop.ops import UOp, Ops, UOpMetaClass, track_rewrites, graph_rewrite, gate_kernel_sink, KernelInfo +from tinygrad.uop.ops import UOp, Ops, UOpMetaClass, rewrite_group, graph_rewrite, gate_kernel_sink, KernelInfo from tinygrad.uop.spec import type_verify, spec_tensor from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, pluralize, SCACHE, BASEDIR, partition, dedup @@ -98,10 +98,14 @@ pm_post_sched_cache = PatternMatcher([ create_new_buffer(ctx, b) if isinstance(b.arg, ParamArg) and b.addrspace is AddrSpace.GLOBAL else None), ]) +def resolve_linear_call(linear_call:UOp): + linear = graph_rewrite(linear_call.src[0], pm_post_sched_cache, ctx=({}, linear_call.src[1:]), walk=True, name="params to buffers") + binds = {f"p{i}":x.src[0] for i,x in enumerate(linear_call.src[1:]) if x.op is Ops.BIND} + return linear.substitute({v:binds[v.expr] for v in linear.variables() if v.expr in binds}, enter_calls=True, name="resolve scalar params") + pm_resolve_linear_call = PatternMatcher([ # call LINEAR is resolved here - (UPat(Ops.CALL, src=(UPat(Ops.LINEAR),), name="linear_call", allow_any_len=True), lambda linear_call: - graph_rewrite(linear_call.src[0], pm_post_sched_cache, ctx=({}, linear_call.src[1:]), walk=True, name="params to buffers")), + (UPat(Ops.CALL, src=(UPat(Ops.LINEAR),), name="linear_call", allow_any_len=True), resolve_linear_call), ])+pm_flatten_linear schedule_cache: dict[bytes, UOp] = {} @@ -167,7 +171,7 @@ pm_copy_from_store = PatternMatcher([ (UPat(Ops.CALL, src=(UPat(Ops.SINK, name="ast"),), allow_any_len=True), assert_all_same_devices), ]) -@track_rewrites(lambda _,ret: f"Schedule {pluralize('Kernel', len(ret[0].src))}") +@rewrite_group(lambda _,ret: f"Schedule {pluralize('Kernel', len(ret[0].src))}") def create_linear_with_vars(big_sink:UOp) -> tuple[UOp, dict[str, int]]: # big_sink srcs are all the Tensors linear_call = graph_rewrite(big_sink, pm_schedule, name="schedule to linear", enter_calls=True) diff --git a/tinygrad/schedule/allreduce.py b/tinygrad/schedule/allreduce.py index f5cc4c8e95..b48ad1e06a 100644 --- a/tinygrad/schedule/allreduce.py +++ b/tinygrad/schedule/allreduce.py @@ -15,14 +15,13 @@ def handle_allreduce(buf:UOp, red:UOp) -> UOp|None: use_ring = concrete and not use_all2all and (RING >= 2 or (ndev > 2 and numel > getenv("RING_ALLREDUCE_THRESHOLD", 256_000) and RING >= 1)) if DEBUG >= 2: print(f"{'ALL2ALL' if use_all2all else 'RING' if use_ring else 'NAIVE'} ALLREDUCE {ndev}x{numel} | {buf.dtype}") - if not concrete: buf = buf.pad_to(buf.max_shape) + buf = buf.pad_to(buf.max_shape) # contiguous before we copy it buf = buf.contiguous() # naive: copy to all devices. if you shrink later, that'll be handled if not use_ring and not use_all2all: - out = functools.reduce(lambda x,y: x.alu(op, y), [buf.mselect(i).copy_to_device(device) for i in range(ndev)]) - return out if concrete else out.shrink_to(shape) + return functools.reduce(lambda x,y: x.alu(op, y), [buf.mselect(i).copy_to_device(device) for i in range(ndev)]).shrink_to(shape) # chunk data into ndev pieces assert isinstance(numel, int) diff --git a/tinygrad/schedule/indexing.py b/tinygrad/schedule/indexing.py index b3b6d6da68..b084b1fb51 100644 --- a/tinygrad/schedule/indexing.py +++ b/tinygrad/schedule/indexing.py @@ -2,30 +2,55 @@ from typing import Iterator import functools, itertools from dataclasses import dataclass, field, replace from tinygrad.dtype import dtypes, AddrSpace -from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType, profile_matches, broadcast_axes +from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType, rewrite_group, broadcast_axes from tinygrad.uop.ops import gate_kernel_sink from tinygrad.uop.symbolic import symbolic, pm_simplify_valid, pm_drop_and_clauses from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored, Context, SPEC +@dataclass +class IndexingContext: + realize_map: dict[UOp, None|list[int]] = field(default_factory=dict) + non_removable: dict[UOp, None] = field(default_factory=dict) + range_map: dict[UOp, tuple[tuple[UOp, ...], tuple[UOp, ...]]] = field(default_factory=dict) + # loads reachable from each UOp memoized across matches + buf_cache: dict[UOp, frozenset[UOp]] = field(default_factory=dict) + + # create ranges + range_idx: Iterator[int] = field(default_factory=itertools.count) + def new_range(self, s:sint, axistype:AxisType=AxisType.WEAK) -> UOp: + if isinstance(s, UOp) and s.op is Ops.RANGE: return s + # if a range has a 1 src, it's the same as UOp.const(0) + return UOp.range(s, next(self.range_idx), axistype) if resolve(s!=1) else UOp.const(0) + + ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.AFTER, Ops.BUFFER, Ops.SLICE, Ops.CONST, Ops.BIND, Ops.MSELECT, Ops.MSTACK, Ops.PARAM, Ops.LOAD, Ops.CALL, Ops.FUNCTION} -def realize(ctx:dict[UOp, None], tr:UOp) -> None: ctx[tr] = None +def realize(ctx:IndexingContext, tr:UOp) -> None: ctx.realize_map[tr] = None -def realize_srcs(ctx:dict[UOp, None], rb:UOp) -> None: +def realize_srcs(ctx:IndexingContext, rb:UOp) -> None: for s in rb.src: - if s.base.op not in ALWAYS_CONTIGUOUS: ctx[s] = None + if s.base.op not in ALWAYS_CONTIGUOUS: ctx.realize_map[s] = None -def realize_store_after_src(ctx:dict[UOp, None], dest:UOp, src:UOp): +def realize_store_after_src(ctx:IndexingContext, dest:UOp, src:UOp): # don't realize SLICE when it's the direct source of STORE+AFTER — the target buffer is the output - if src.op is Ops.SLICE and src in ctx \ + if src.op is Ops.SLICE and src in ctx.realize_map \ and not dest.op_in_backward_slice_with_self(Ops.SHRINK, Ops.PERMUTE, Ops.FLIP, Ops.PAD): - del ctx[src] + del ctx.realize_map[src] # you don't usually have to do this for assign unless there's a WAR hazard like TestAssign.test_assign_double_diamond_reduce - if dest.base in src.backward_slice_with_self: ctx[src] = None + if dest.base in src.backward_slice_with_self: ctx.realize_map[src] = None + +def realize_custom_kernel_srcs(ctx:IndexingContext, c:UOp) -> None: + for s in c.src[1:]: + while s.op is Ops.RESHAPE: s = s.src[0] + if s.op not in ALWAYS_CONTIGUOUS: + ctx.realize_map[s] = None + ctx.non_removable[s] = None pm_generate_realize_map = PatternMatcher([ + # realize the inputs of custom kernel calls + (UPat(Ops.CALL, src=(UPat((Ops.SINK, Ops.PROGRAM)),), name="c", allow_any_len=True), realize_custom_kernel_srcs), # always realize (UPat({Ops.CONTIGUOUS, Ops.STORE}, name="tr"), realize), # realize srcs of these @@ -41,20 +66,6 @@ class BufferizeOpts: addrspace: AddrSpace = AddrSpace.GLOBAL removable: bool = True -@dataclass -class IndexingContext: - realize_map: dict[UOp, None|list[int]] = field(default_factory=dict) - range_map: dict[UOp, tuple[tuple[UOp, ...], tuple[UOp, ...]]] = field(default_factory=dict) - # loads reachable from each UOp memoized across matches - buf_cache: dict[UOp, frozenset[UOp]] = field(default_factory=dict) - - # create ranges - range_idx: Iterator[int] = field(default_factory=itertools.count) - def new_range(self, s:sint, axistype:AxisType=AxisType.WEAK) -> UOp: - if isinstance(s, UOp) and s.op is Ops.RANGE: return s - # if a range has a 1 src, it's the same as UOp.const(0) - return UOp.range(s, next(self.range_idx), axistype) if resolve(s!=1) else UOp.const(0) - def broadcast_rngs(x:UOp, src:UOp, rngs:tuple[UOp, ...]) -> tuple[UOp, ...]: if x.op not in GroupOp.Broadcastable: return rngs baxes, nleft = broadcast_axes(src.shape, x.shape), len(x.shape)-len(src.shape) @@ -84,7 +95,7 @@ def create_bufferize_and_index_srcs(ctx:IndexingContext, x:UOp) -> list[UOp]: new_src = s.end(*[r for r in closed_ranges if r.op is Ops.RANGE]) del ctx.realize_map[s] else: - removable = s.op not in ALWAYS_CONTIGUOUS + removable = s.op not in ALWAYS_CONTIGUOUS and s not in ctx.non_removable # LOCAL: None in the device assigns it a number later opts = BufferizeOpts(device=s.device, removable=removable) if len(ctx.range_map[s][1]) == len(realized_ranges) else \ BufferizeOpts(device=s.device, addrspace=AddrSpace.LOCAL, removable=removable) @@ -105,6 +116,7 @@ def convert_pad_to_where_to_keep_behavior_local(ctx:IndexingContext, x:UOp): def convert_reduce_to_reduce_with_ranges(ctx:IndexingContext, x:UOp): if x.arg[1] == 0: return None + if x not in ctx.range_map: raise RuntimeError("REDUCE has no ranges in rangeify, UOp verification failed") bx = create_bufferize_and_index_based_on_ranges(ctx, x) # input ranges new_ranges = list(ctx.range_map[x][0][:x.arg[1]]) @@ -176,13 +188,13 @@ def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UO case _: raise RuntimeError(f"{op} is not a MovementOp") return rngs -@profile_matches +@rewrite_group(new_ctx=False) def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]: if debug: print("**************************") rctx = IndexingContext() # get ops to realize - graph_rewrite(tsink, pm_generate_realize_map, ctx=rctx.realize_map, name="get realize") + graph_rewrite(tsink, pm_generate_realize_map, ctx=rctx, name="get realize") # get the consumer map with cpu_profile("consumer map in rangeify", "TINY"): diff --git a/tinygrad/schedule/multi.py b/tinygrad/schedule/multi.py index 17110647f4..14d4126812 100644 --- a/tinygrad/schedule/multi.py +++ b/tinygrad/schedule/multi.py @@ -126,7 +126,7 @@ def reshape_multi(root:UOp, multi:UOp): new_shardings = [] for ax, rng in multi.sharding: count = int(rng.vmax)+1 - target = prod(multi.shape[:ax]) + target = ssimplify(prod(multi.shape[:ax])) if target not in arg_acc: raise RuntimeError(f"reshape {multi.shape} -> {new_shape} moved items between shards") new_ax = len(arg_acc) - arg_acc[::-1].index(target) - 1 if new_shape[new_ax] % count != 0: raise RuntimeError(f"reshape {multi.shape} -> {new_shape} moved items between shards") diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 9544c8336f..3a3fcc37c0 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -3,10 +3,10 @@ from typing import cast import itertools from tinygrad.dtype import dtypes, AddrSpace, Invalid, to_dtype, strong_dtype from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, KernelInfo, ParamArg, shape_to_shape_arg -from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, profile_matches, identity_element -from tinygrad.uop.symbolic import symbolic +from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, rewrite_group, identity_element +from tinygrad.uop.symbolic import symbolic, pm_fold_cast_const from tinygrad.uop.movement import mop_cleanup -from tinygrad.helpers import prod, getenv, dedup, all_int, DEBUG, SPLIT_REDUCEOP, VIZ, MAX_KERNEL_BUFFERS +from tinygrad.helpers import prod, getenv, dedup, all_int, DEBUG, SPLIT_REDUCEOP, VIZ, MAX_KERNEL_BUFFERS, SPEC from tinygrad.helpers import PCONTIG, FLOAT16, OPENPILOT_HACKS, argsort, partition, get_single_element from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_simplify from tinygrad.codegen.opt import Opt @@ -150,7 +150,7 @@ earliest_rewrites = mop_cleanup+PatternMatcher([ (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD), name="x"), lambda x: x.src[0]), # SINK only ever references the base - (UPat(Ops.SINK, name="x"), lambda x: x.replace(src=tuple(y.base for y in x.src))), + (UPat(Ops.SINK, name="x"), lambda x: x.replace(src=tuple(y.unsharded_base for y in x.src))), # ** copy rules ** @@ -202,6 +202,7 @@ ALWAYS_RUN_OPS = {Ops.CONTIGUOUS, Ops.NOOP} # you don't know in the first pass if axes are going to die, this happens if there's an EXPAND to the left def cleanup_dead_axes(b:UOp): + if not b.arg.removable: return None # don't optimize ALWAYS_RUN_OPS or AFTER (AFTER is a buffer identity — ranges define consumer access, not computation) if b.src[0].op in ALWAYS_RUN_OPS or b.src[0].op is Ops.AFTER: return None @@ -335,6 +336,26 @@ pm_remove_bufferize = PatternMatcher([ (UPat(Ops.END, src=(UPat(Ops.NOOP, name="x"),), allow_any_len=True), lambda x: x), ]) +def no_indexing_calls(u:UOp): + new_srcs = [] + for x in u.src: + if x.op is Ops.INDEX: + # sometimes if call srcs have children the call will get an INDEX. we remove it here. + # TODO: we should add safety checks here for contiguous + new_srcs.append(x.src[0]) + elif x.op is Ops.SHRINK: + # SHRINK with offset 0 is fine + # TODO: check offset + new_srcs.append(x.src[0]) + else: + # everything else we pass through + new_srcs.append(x) + return u.replace(src=tuple(new_srcs)) + +pm_no_indexing_calls = PatternMatcher([ + (UPat(Ops.CALL, name="u"), no_indexing_calls), +]) + DEVICE_MAX_BUFS = {"METAL": 31, "WEBGPU": 8, "CPU": 31} # TODO: get from device? def limit_bufs(ctx:IndexingContext, root:UOp): if (device:=root.device) is None: return None # no device, index related calculations @@ -652,7 +673,7 @@ pm_range_creation = PatternMatcher([ (UPat(Ops.STAGE, name="x"), rangeify_on_stage), ]) -@profile_matches +@rewrite_group(new_ctx=False) def get_kernel_graph(sink:UOp) -> UOp: tsink = graph_rewrite(sink, multi_pm, name="multi_pm") if OPENPILOT_HACKS: tsink = graph_rewrite(tsink, pm_fold_moved_after, ctx={}, name="fold moved afters") @@ -711,7 +732,7 @@ def get_kernel_graph(sink:UOp) -> UOp: tsink = tsink.substitute(subs) tsink = graph_rewrite(tsink, pm_simple_rangeify, bottom_up=True, name=f"merge kernels ({len(subs)})") - tsink = graph_rewrite(tsink, symbolic+pm_reduce_simplify+pm_const_buffer_folding+pm_remove_bufferize, name="symbolic+reduce_collapse+debuf") + tsink = graph_rewrite(tsink, symbolic+pm_reduce_simplify+pm_const_buffer_folding+pm_remove_bufferize+pm_no_indexing_calls, name="symbolic+reduce_collapse+debuf") #tsink = graph_rewrite(tsink, pm_limit_bufs, ctx=rctx, name="limit buffers") if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Rangeify") @@ -723,4 +744,8 @@ def get_kernel_graph(sink:UOp) -> UOp: tsink = graph_rewrite(tsink, split_kernels, bottom_up=True, name="split kernels") if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Kernel Graph") + if SPEC: + # validate the kernel graph + from tinygrad.uop.spec import type_verify, spec_kernel_graph + type_verify(tsink, spec_kernel_graph, enter_calls=False) return tsink diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 71b4fa8f9d..304d1fe1a3 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -1,17 +1,242 @@ # inspired by https://github.com/karpathy/micrograd/blob/master/micrograd/engine.py from __future__ import annotations import time, functools, sys, inspect, pathlib, hashlib, weakref +from dataclasses import dataclass, field from typing import Any, Callable, cast, get_args, ParamSpec, TypeGuard, TypeVar, Generic, TYPE_CHECKING if TYPE_CHECKING: import numpy -from tinygrad.dtype import DType, DTypeLike, dtypes, ConstType, least_upper_dtype, to_dtype, strong_dtype, _from_np_dtype, _to_np_dtype, PyConst +from tinygrad.dtype import DType, DTypeLike, dtypes, ConstType, least_upper_dtype, to_dtype, strong_dtype, \ + _from_np_dtype, _to_np_dtype, PyConst, AddrSpace from tinygrad.helpers import all_int, getenv, fetch, Metadata, TRACEMETA, TracingKey -from tinygrad.helpers import cpu_profile, suppress_finalizing, disable_gc -from tinygrad.uop.ops import UOp, Ops, sint, all_metadata, Variable, ConstLike +from tinygrad.helpers import cpu_profile, suppress_finalizing, disable_gc, VIZ, pluralize +from tinygrad.uop.ops import UOp, Ops, sint, all_metadata, Variable, ConstLike, UPat, PatternMatcher, GroupOp, ParamArg, graph_rewrite, rewrite_group from tinygrad.mixin.rand import RandMixin from tinygrad.schedule import create_linear_with_vars from tinygrad.device import Buffer, canonicalize_device from tinygrad.engine.realize import run_linear -from tinygrad.callify import transform_to_call + +# *** callify: transform a tensor graph into a CALL UOp such that all state is properly scoped *** + +@dataclass +class AllocCtx: + uop_list: list[UOp] = field(default_factory=list) + buffer_map: dict[UOp, UOp] = field(default_factory=dict) + bases: set[UOp] = field(default_factory=set) + assigns: list[UOp] = field(default_factory=list) + replacements: list[UOp] = field(default_factory=list) + +def tag_uop(ctx:AllocCtx, x:UOp): + if x.tag is not None: return None + ctx.uop_list.append(x) + return x.replace(tag=(len(ctx.uop_list)-1,)) + +def disk_like(u:UOp): return isinstance(u.device, str) and u.device.startswith(("DISK", "TINYFS")) + +def disk_copy_is_buffer(ctx:AllocCtx, u:UOp): + # copies to disk are replaced with the disk buffer + if disk_like(u) and u.tag is None: + ctx.buffer_map[u] = u.empty_like() + return u.rtag(()) + # all copies from disk/numpy are realized into a real buffer + from_creation = isinstance(u.src[0].device, str) and u.src[0].device.startswith(("NPY", "DISK", "PYTHON", "TINYFS")) + if from_creation: return tag_uop(ctx, u) + +# CONTIGUOUS and AFTER + parents are the only nodes that get updated +add_tags = PatternMatcher([ + (UPat(Ops.COPY, name="u"), disk_copy_is_buffer), + # no tag on copies that are assigned via STORE+AFTER — merge COPY tag into AFTER + (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE, src=(UPat(name="dest"), UPat(Ops.COPY, name="c")))), name="a"), + lambda a,c,dest: a.replace(src=(a.src[0], a.src[1].replace(src=(dest, c.rtag(())))), tag=a.tag+c.tag) if a.tag and c.tag else None), + (UPat((Ops.CONTIGUOUS, Ops.AFTER), name="x"), tag_uop), + (UPat(GroupOp.All, name="x"), lambda ctx,x: tag_uop(ctx,x) if x in ctx.bases else None), +]) + +def replace_contig_with_store_after(u:UOp): + # can't allocate a buffer for a virtual value + if u.is_virtual: return None + # if size is 0, remove the contig + if 0 in u.shape: return u.src[0] + # no real contig for DISK/TINYFS tensors, they are left alone + if disk_like(u): return u.rtag(None) + buf = u.empty_like() + return buf.after(buf.store(u.src[0])).rtag(u.tag) + +def replace_store_after_with_contig(u:UOp, src:UOp): + assigned_to = u + while assigned_to.op in {Ops.BITCAST, Ops.AFTER, Ops.UNSHARD}: assigned_to = assigned_to.src[0].base + if assigned_to.op not in {Ops.BUFFER, Ops.SLICE}: return src.contiguous(tag=u.tag) + +def _make_buffer_view(src:UOp) -> UOp|None: + """If movement ops on src collapse to a contiguous range, return SLICE. Otherwise None.""" + if (offset := src.contiguous_view_offset()) is None: return None + buf = src.base + if buf.op is Ops.SLICE: + byte_offset = buf.src[1].val * buf.src[0].dtype.itemsize + offset * src.dtype.itemsize + buf = buf.src[0] + if byte_offset % buf.dtype.itemsize != 0: return None + offset = byte_offset // buf.dtype.itemsize + return UOp(Ops.SLICE, src.dtype, (buf, UOp.const(offset)), src.numel()) + +def contiguous_mops_to_view(c:UOp, src:UOp): + """MOPS(BUFFER) → SLICE when movement ops collapse to a contiguous range.""" + buf = src.base + if buf.op not in {Ops.BUFFER, Ops.SLICE, Ops.UNSHARD}: return None + if src.op is Ops.RESHAPE and src.src[0].op in {Ops.BUFFER, Ops.SLICE} and c.op is not Ops.BITCAST: return None + if c.op is not Ops.BITCAST and src.op is Ops.BUFFER: return None + + # no symbolic shape + if not all_int(c.shape): return None + + if buf.op is not Ops.UNSHARD and (view := _make_buffer_view(src)) is not None: + view = (view.replace(dtype=c.dtype, arg=c.numel()) if c.op is Ops.BITCAST else view).reshape(c.shape) + return c.replace(src=(view,)) if c.op is Ops.COPY else view + + # for UNSHARD tensors, use multi_pm to resolve per-shard movement ops, then create SLICE on the resolved result + if not isinstance(c.device, str): + from tinygrad.schedule.multi import multi_pm + resolved = graph_rewrite(src, multi_pm, name="multi_buffer_view") + if resolved.op is not Ops.UNSHARD: return None + if (view := _make_buffer_view(resolved.src[0])) is None: return None + return view.reshape(resolved.src[0].shape).unshard(resolved.arg, resolved.src[1:]).contiguous(tag=c.tag) + + return None + +def _precompiled_output_redirect(s:UOp, t:UOp) -> UOp|None: + # how output s lands in the caller's buffer t, or None if it must be copied into t + # materialize straight into t + if s.op is Ops.CONTIGUOUS: return t.after(t.store(s.src[0])) + # rebind output storage to t + if s.op in {Ops.BUFFER, Ops.UNSHARD} and s.has_buffer_identity(): return t + return None + +def transform_precompiled_call(c:UOp) -> UOp|None: + if not c.arg.precompile: return None + assert c.src[0].op is Ops.TUPLE, f"expected TUPLE body for precompiled FUNCTION, got {c.src[0].op}" + input_buffers = tuple(x.contiguous() if x.op not in {Ops.AFTER, Ops.BIND} else x for x in c.src[1:]) + + # add the outputs to the call + srcs = c.src[0].src + resolved = [c.gettuple(i) for i in range(len(srcs))] + outs = tuple(r.empty_like() for r in resolved) + targets = [o.param_like(len(c.src)-1+i).shrink_to(s.shape) for i,(o,s) in enumerate(zip(outs, srcs))] + + subs:dict[UOp, UOp] = {} + items:list[UOp] = [] + for s, t in zip(srcs, targets): + after_deps:list[UOp] = [] + while s.op is Ops.AFTER: + after_deps.extend(s.src[1:]) + s = s.src[0] + if (placed := _precompiled_output_redirect(s, t)) is not None and s not in subs: + subs[s] = placed + items.append(s.after(*after_deps) if after_deps else s) + else: + items.append(t.after(t.store(s.after(*after_deps)))) + fxn = UOp.sink(*(x.substitute(subs) for x in items)) + + # body switches from TUPLE to SINK, so the node becomes an opaque CALL (not FUNCTION) + new_call = UOp(Ops.CALL, src=(fxn, *input_buffers, *outs), arg=c.arg) + rets = tuple(o.after(new_call) for o in outs) + + # if the CALL has symbolic shapes, shrink the max-sized output to the actual symbolic shape + # NOTE: must use resolved shapes from the FUNCTION (which substitutes PARAMs with external args), not raw body shapes + rets = tuple(r.shrink_to(rs.shape) for r,rs in zip(rets, resolved)) + + return UOp.maketuple(*rets) + +# NOTE: adding rules to here is bad. these all need to run before the schedule cache +pm_early_transform_tensor_graph = PatternMatcher([ + # transform precompiled FUNCTIONs into CALLs (body becomes SINK with stores) + (UPat(Ops.FUNCTION, name="c"), transform_precompiled_call), + + # resolve TUPLE+GETTUPLE (for precompiled calls) + (UPat(Ops.GETTUPLE, src=(UPat(Ops.TUPLE, name="t"),), name="g"), lambda g,t: t.src[g.arg]), + + # fold MOPS+BITCAST over BUFFER/SLICE into SLICE when movement ops collapse to contiguous range + (UPat((Ops.BITCAST, Ops.COPY, Ops.CONTIGUOUS), src=(UPat(GroupOp.Movement|{Ops.BUFFER}, name="src"),), name="c"), contiguous_mops_to_view), + + # remove contiguous on movement ops before a copy on disk + (UPat(GroupOp.Movement-{Ops.SHRINK, Ops.RESHAPE}, name="x").f(Ops.CONTIGUOUS).f(Ops.COPY, name="copy"), lambda x,copy: + copy.replace(src=(x,), tag=None) if isinstance(x.device, str) and x.device.startswith("DISK") else None), + # push copy past movement ops to disk + (UPat(GroupOp.Movement-{Ops.SHRINK, Ops.RESHAPE}, name="x").f(Ops.COPY, name="copy"), lambda x,copy: + x.replace(src=(copy.replace(src=(x.src[0],), tag=None),)+x.src[1:]) \ + if isinstance(x.device, str) and x.device.startswith("DISK") else None), + + # add CONTIGUOUS to tagged UOps + (UPat(GroupOp.All-{Ops.CONTIGUOUS, Ops.AFTER, Ops.STORE}, name="x"), + lambda x: None if x.tag is None else x.rtag(None).contiguous(tag=x.tag) if x.tag else x.replace(tag=None)), + # remove extra CONTIGUOUS on AFTER (only when target is contiguous) + (UPat(Ops.CONTIGUOUS, src=(UPat(Ops.AFTER, name="a"),), name="c"), + lambda a,c: a.replace(tag=(a.tag or ())+(c.tag or ())) if a.src[0].has_buffer_identity() else None), + # replace AFTER+STORE with CONTIGUOUS when target is not a buffer + (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE, src=(UPat(), UPat(name="src")))), name="u"), replace_store_after_with_contig), + # replace CONTIGUOUS with STORE+AFTER + (UPat(Ops.CONTIGUOUS, name="u"), replace_contig_with_store_after), + # remove DETACH/CONTIGUOUS_BACKWARD (allows more contiguous removal) + (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD), name="x"), lambda x: x.src[0]), +]) + +def finalize_after(ctx:AllocCtx, x:UOp): + # untagged: record as an assign for the call body + if x.tag is None: + ctx.assigns.append(x) + return None + # tagged: untag and map each original pre-rewrite UOp to the stripped buffer; the untagged result is reprocessed as untagged + ret = x.replace(tag=None) + replace_uop = ret + # then, add views back + views:list[UOp] = [] + while replace_uop.op in GroupOp.Movement|{Ops.UNSHARD, Ops.BITCAST, Ops.AFTER}: + if replace_uop.op is not Ops.AFTER: views.append(replace_uop) + replace_uop = replace_uop.src[0] + for v in reversed(views): replace_uop = v.replace(src=(replace_uop,)+v.src[1:]) + for t in x.tag: + original_uop: UOp = ctx.uop_list[t] + ctx.buffer_map[original_uop] = replace_uop.shrink_to(original_uop.shape) + return ret + +def replace_input_buffer(ctx:AllocCtx, b:UOp): + ctx.replacements.append(b) + if b.op is Ops.BIND: return b.param_like(len(ctx.replacements)-1) + return UOp.param(len(ctx.replacements)-1, b.dtype, b.shape, b.device, + addrspace=b.addrspace if b.addrspace is not None else AddrSpace.GLOBAL) + +pm_finalize_call = PatternMatcher([ + (UPat(Ops.AFTER, name="x"), finalize_after), + (UPat(Ops.COPY, name="x"), lambda ctx,x: ctx.assigns.append(x) if isinstance(x.device, str) and x.device.startswith(("DISK", "TINYFS")) else None), +]) + +pm_replace_buf = PatternMatcher([ + # replace BUFFER with PARAM for cache key normalization + (UPat(Ops.BUFFER, src=(UPat(),), name="b"), lambda ctx,b: + replace_input_buffer(ctx, b) if isinstance(b.arg, ParamArg) and b.addrspace is AddrSpace.GLOBAL else None), + # replace SLICE with PARAM. this rewrite is bottom up so BUFFERs we don't need won't be in the input + (UPat(Ops.SLICE, src=(UPat(Ops.BUFFER), UPat(Ops.CONST, dtype=dtypes.weakint)), name="b"), replace_input_buffer), + # strip value from BIND for cache key normalization, so different values hit same cache + (UPat(Ops.BIND, src=(UPat(Ops.PARAM), UPat(Ops.CONST)), name="b"), replace_input_buffer), +]) + +@rewrite_group(lambda _,ret: f"Callify {pluralize('Buffer', len(ret[1]))}") +def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]: + if VIZ: graph_rewrite(big_sink, PatternMatcher([]), name="View Tensor Graph") + # uop list is a list in the original_sink graph and we can map to the tags later + # same predicate as Tensor.realize + ctx = AllocCtx(bases={base for x in big_sink.src if not (base:=x.base).is_virtual and not base.has_buffer_identity() + and base.op is not Ops.AFTER and base.addrspace is not AddrSpace.ALU}) + + # this rewrite is "read-only", it adds simple things to buffer_map and may sink things on big_sink, bottom_up + # this is the only one where we have to be careful to not break the tensor graph + big_sink = graph_rewrite(big_sink, add_tags, ctx=ctx, bottom_up=True, name="number the uops") + + # here we can break the tensor graph. this is the only place you need to maintain numbered tags + big_sink = graph_rewrite(big_sink, pm_early_transform_tensor_graph, name="early transform tensor graph") + + # here we construct the final buffer_map: as-built nodes -> their final storage. values are never keys + graph_rewrite(big_sink, pm_finalize_call, ctx=ctx, name="finalize call") + ret = graph_rewrite(UOp.sink(*ctx.assigns), pm_replace_buf, ctx=ctx, bottom_up=True, name="replace bufs").call(*ctx.replacements) + assert not any(x in ctx.buffer_map for x in ctx.buffer_map.values()) + if VIZ: graph_rewrite(ret, PatternMatcher([]), name="View Call") + return ret, ctx.buffer_map # *** all in scope Tensors are here. this gets relevant UOps *** diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 4f8eabd954..9ae599b987 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -528,9 +528,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass): if self.op is Ops.CONST: return self if self.op is Ops.SINK and all(s.op is Ops.CONST or (s.op is Ops.STACK and len(s.src) == 0) for s in self.src): return self # late import! - from tinygrad.uop.symbolic import symbolic + from tinygrad.uop.symbolic import symbolic, pm_fold_cast_const with Context(TRACK_MATCH_STATS=0 if not tracked else TRACK_MATCH_STATS.value): - return graph_rewrite(self, symbolic, name="simplify") + return graph_rewrite(self, symbolic+pm_fold_cast_const, name="simplify") def ssimplify(self) -> UOp|ConstType: return ret.val if (ret:=self.simplify()).op is Ops.CONST else ret def _eval(self, dtype, expected_type:Type[T]) -> T: assert self.dtype in dtype, f"eval with wrong dtype {self}" @@ -758,7 +758,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): assert arg is None or isinstance(self.device, tuple) inp = self if arg is None else UOp(Ops.MSELECT, src=(self,), arg=arg) if inp.dtype in dtypes.weaks: raise RuntimeError(f"cannot create storage for weak dtype {inp.dtype}") - return UOp(Ops.COPY, src=(inp,), arg=device) + return UOp(Ops.COPY, src=(inp.pad_to(inp.max_shape),), arg=device).shrink_to(inp.shape) def mselect(self, arg:int) -> UOp: return UOp(Ops.MSELECT, src=(self,), arg=arg) def mstack(self, *srcs: UOp) -> UOp: return UOp(Ops.MSTACK, src=(self,)+srcs) if len(srcs) else self @property @@ -772,6 +772,15 @@ class UOp(RandMixin, metaclass=UOpMetaClass): if self.op is Ops.DETACH: return self.src[0].base # DETACH can't change base return self + # base with UNSHARD + @property + def unsharded_base(self) -> UOp: + if self.op in GroupOp.Movement: return self.src[0].base + if self.op is Ops.DETACH: return self.src[0].base # DETACH can't change base + # TODO: why can't this be in normal base? + if self.op is Ops.UNSHARD: return self.src[0].base + return self + # cached property here makes external_uop_gc fail, why? @property def as_shape(self) -> tuple[sint, ...]: @@ -801,7 +810,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass): # arg is the other srcs; all are cast to the promoted dtype, spec requires STACK srcs to match its dtype srcs = (self,)+tuple(arg) dtype = cast(DType, dtype_from_uop(Ops.STACK, srcs, None)) - return UOp(Ops.STACK, dtype, tuple(u.cast(dtype) for u in srcs)) + # TODO: why cast here? + return UOp(Ops.STACK, dtype, tuple(u if u.base.is_invalid else u.cast(dtype) for u in srcs)) case _: raise RuntimeError(f"{op} is not a MovementOp") usrcs = [shape_to_shape_arg(arg) for arg in src_args] if len(usrcs) == 0: return UOp(op, src=(self,), arg=arg) @@ -1096,12 +1106,10 @@ class UOp(RandMixin, metaclass=UOpMetaClass): if self.op is Ops.CONST and self.val is not Invalid: return self.val, self.val if self.op is Ops.INDEX: return self.src[0]._min_max if self.op is Ops.CAST: - # an int destination truncates a float source toward zero. trunc is monotone + # rounding is monotone (truncation toward zero into an int, to-nearest onto the value grid into a float) smin, smax = self.src[0]._min_max - if dtypes.is_int(self.dtype) and dtypes.is_float(self.src[0].dtype) and all(math.isfinite(v) for v in (smin, smax)): - smin, smax = math.trunc(smin), math.trunc(smax) - # a cast to unsigned keeps exact bounds when the source fits - # TODO: can do more based on new dtype window + trunc = truncate.get(self.dtype) if dtypes.is_float(self.dtype) else math.trunc if dtypes.is_int(self.dtype) else None + if trunc is not None and all(math.isfinite(v) for v in (smin, smax)): smin, smax = trunc(smin), trunc(smax) if dtypes.is_unsigned(self.dtype) and 0 <= smin and smax <= self.dtype.max: return smin, smax if self.dtype in dtypes.floats+dtypes.sints+(dtypes.weakint,): return max(self.dtype.min, smin), min(smax, self.dtype.max) return self.dtype.min, self.dtype.max @@ -1165,8 +1173,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass): src: tuple[UOp, ...] = (UOp(Ops.NOOP) if shape is None else shape_to_shape_arg(shape),) return UOp(Ops.PARAM, src=src, arg=ParamArg(slot, dtype, vmin_vmax, multiple_of, name, addrspace, axis, device, volatile)) def param_like(self, slot:int): + if self.op is Ops.BIND: return self.src[0].replace(arg=replace(self.src[0].arg, slot=slot, name=f"p{slot}")) addrspace = self.addrspace if self.addrspace is not None else AddrSpace.GLOBAL - if self.op is Ops.BIND: return self.src[0].replace(arg=replace(self.src[0].arg, slot=slot, addrspace=addrspace)) return UOp.param(slot, self.dtype, self.shard_shape if self.axis is not None else self._shape, self.device, addrspace=addrspace, axis=self.axis) @staticmethod @@ -1186,10 +1194,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass): body = self if self.op is Ops.TUPLE else UOp.maketuple(self) return UOp(Ops.FUNCTION, src=(body,)+srcs, arg=CallInfo(grad_fxn, name, precompile, precompile_backward, aux)) def custom_kernel(*srcs:UOp, fxn:Callable, grad_fxn:Callable|None=None) -> list[UOp]: - contig_srcs = tuple(x.contiguous() if x.op is not Ops.AFTER else x for x in srcs) - placeholders = [UOp.placeholder_like(s, slot=i) for i,s in enumerate(contig_srcs)] - kernel = fxn(*placeholders).call(*contig_srcs, grad_fxn=grad_fxn) - return [s.after(kernel) for s in contig_srcs] + placeholders = [UOp.placeholder_like(s, slot=i) for i,s in enumerate(srcs)] + kernel = fxn(*placeholders).call(*srcs, grad_fxn=grad_fxn) + return [s.after(kernel) for s in srcs] def to_elf(self) -> TinyELF: assert self.op is Ops.PROGRAM and isinstance(self.arg, ProgramInfo), "to_elf should only be called on a PROGRAM ast" @@ -1514,55 +1521,52 @@ def add_trace_group(kt:TracingKey) -> None: tracked_ctxs.append([]) active_group:list[int] = [] -def track_rewrites(name:Callable[..., str|TracingKey]|bool=True, replay:bool=False): +active_rewrites:list[TrackedGraphRewrite] = [] +def rewrite_group(name:Callable[..., str|TracingKey]|bool=True, replay:bool=False, new_ctx:bool=True): + if not new_ctx: assert not callable(name) and not replay, "name fxn and replay are only supported for new_ctx groups" def _decorator(func): def __wrapper(*args, **kwargs): + # without tracking, we just call the function (unless top-level, which always profiles) + if TRACK_MATCH_STATS < 2 and not new_ctx: return func(*args, **kwargs) fn = key = func.__name__ idx = -1 if TRACK_MATCH_STATS >= 2: - add_trace_group(key:=TracingKey(n:=f"{fn} n{next(_name_cnt.setdefault(fn, itertools.count(1)))}", (n,))) - active_group.append(idx:=len(tracked_keys)-1) + if new_ctx: + add_trace_group(key:=TracingKey(n:=f"{fn} n{next(_name_cnt.setdefault(fn, itertools.count(1)))}", (n,))) + active_group.append(idx:=len(tracked_keys)-1) + else: + rewrite_name = str(kwargs.get("name", None) or fn) + assert args and isinstance(args[0], UOp), f"invalid match tracing inputs for {rewrite_name} with {args}" + loc = ((frm:=sys._getframe(1)).f_code.co_filename, frm.f_lineno) + depth = len(active_rewrites) + if not tracked_ctxs: add_trace_group(TracingKey(f"default {fn}")) + dest_group = active_group[-1] if active_group else len(tracked_ctxs)-1 + tracked_ctxs[dest_group].append(ctx:=TrackedGraphRewrite(loc, args[0].trace_num, [], rewrite_name, depth, kwargs.get("bottom_up", False), + kwargs.get("walk", False), kwargs.get("enter_calls", False))) + active_rewrites.append(ctx) + key = rewrite_name # profile spans are named after the rewrite step with cpu_profile(key, "TINY") as e: ret = func(*args, **kwargs) - if TRACK_MATCH_STATS >= 2: active_group.pop() - if TRACK_MATCH_STATS >= 2 and callable(name): - name_ret = name(*args, **kwargs, ret=ret) - assert isinstance(name_ret, (TracingKey, str)), f"name function returned {type(name_ret)}" - tracked_keys[idx] = k = TracingKey(n:=tracked_keys[idx].display_name.replace(fn, name_ret), (n,)) if isinstance(name_ret, str) else name_ret - e.name = TracingKey(k.display_name if isinstance(name_ret, str) else f"{fn} for {k.display_name}", k.keys) + if TRACK_MATCH_STATS >= 2: + if new_ctx: active_group.pop() + else: active_rewrites.pop() + if callable(name): + name_ret = name(*args, **kwargs, ret=ret) + assert isinstance(name_ret, (TracingKey, str)), f"name function returned {type(name_ret)}" + tracked_keys[idx] = k = TracingKey(n:=tracked_keys[idx].display_name.replace(fn, name_ret), (n,)) if isinstance(name_ret, str) else name_ret + e.name = TracingKey(k.display_name if isinstance(name_ret, str) else f"{fn} for {k.display_name}", k.keys) if CAPTURE_PROCESS_REPLAY and replay: # find the unittest frame we're capturing in frm = sys._getframe(1) while (f_back:=frm.f_back) is not None and "unittest" not in f_back.f_code.co_filename: frm = f_back - loc = f"{frm.f_code.co_filename.split('/')[-1]}:{frm.f_lineno} {frm.f_code.co_name}" + replay_loc = f"{frm.f_code.co_filename.split('/')[-1]}:{frm.f_lineno} {frm.f_code.co_name}" # capture global context vars and all the args passed in inputs = (fn, args, kwargs, ContextVar._cache) - replay_capture.append(pickle.dumps(inputs+(loc, ret))) + replay_capture.append(pickle.dumps(inputs+(replay_loc, ret))) return ret return __wrapper return _decorator -active_rewrites:list[TrackedGraphRewrite] = [] -def profile_matches(fxn:Callable): - def wrap_profile_matches(*args, **kwargs): - if TRACK_MATCH_STATS >= 2: - name = str(kwargs.get("name", None) or fxn.__name__) - assert args and isinstance(args[0], UOp), f"invalid match tracing inputs for {name} with {args}" - loc = ((frm:=sys._getframe(1)).f_code.co_filename, frm.f_lineno) - depth = len(active_rewrites) - if not tracked_ctxs: add_trace_group(TracingKey(f"default {fxn.__name__}")) - dest_group = active_group[-1] if active_group else len(tracked_ctxs)-1 - tracked_ctxs[dest_group].append(ctx:=TrackedGraphRewrite(loc, args[0].trace_num, [], name, depth, kwargs.get("bottom_up", False), - kwargs.get("walk", False), kwargs.get("enter_calls", False))) - active_rewrites.append(ctx) - with cpu_profile(name, "TINY"): - ret = fxn(*args, **kwargs) - active_rewrites.pop() - return ret - # without tracking, we just call the function - return fxn(*args, **kwargs) - return wrap_profile_matches - class TrackedPatternMatcher(PatternMatcher): def rewrite(self, uop:UOp, ctx=None): if len(pats:=self.pdict.get(uop.op, [])): @@ -1741,7 +1745,7 @@ class RewriteContext: if n in waitlist: stack.extend(waitlist.pop(n)) return self.replace[root] -@profile_matches +@rewrite_group(new_ctx=False) def graph_rewrite(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False, name=None, bpm=None, walk=False, enter_calls=False) -> UOp: rewrite_ctx = RewriteContext(pm if not bottom_up else None, pm if bottom_up else bpm, ctx, enter_calls) return rewrite_ctx.walk_rewrite(sink) if walk else rewrite_ctx.unified_rewrite(sink) @@ -1754,73 +1758,6 @@ def _rebuild_dtype(n:UOp, new_src:tuple[UOp,...]) -> DType: def sint_to_uop(x:sint, dtype=dtypes.weakint) -> UOp: return UOp.const(x, dtype) def to_max_shape(shape:tuple[sint, ...]) -> tuple[int, ...]: return tuple(int(x.vmax) if isinstance(x, UOp) else x for x in shape) -def select_dtype(u:UOp): - if u.dtype is dtypes.weakfloat: return dtypes.default_float - return dtypes.long if u.overflows(dtypes.int32) else dtypes.int -def lower_weak_node(u:UOp) -> UOp|None: - start, src = (1 if u.op is Ops.WHERE else 0), tuple(s.src[0] if s.op is Ops.CAST and s.dtype in dtypes.weaks else s for s in u.src) - if src == u.src or any(s.dtype in dtypes.weaks for s in src[start:]): return None - dt = strong_dtype(least_upper_dtype(select_dtype(u), *(s.dtype for s in src)) if u.op in GroupOp.Binary - else unwrap(dtype_from_uop(u.op, src, u.arg))) - return u.replace(dtype=None, src=src[:start]+tuple(s.cast(dt) for s in src[start:])).cast(u.dtype) -pm_lower_weak = PatternMatcher([ - (UPat(Ops.CONST, dtype=dtypes.weaks, name="u"), lambda u: UOp.const(u.val, select_dtype(u)).cast(u.dtype)), - # two stacked weak casts are a weakint value used as weakfloat (or vice versa): resolve the inner one at the outer kind's default. - # a SINGLE weak cast is never rewritten here, each consumer absorbs it on its own edge (see lower_weak_srcs) - (UPat(Ops.CAST, dtype=dtypes.weaks, src=(UPat(Ops.CAST, dtype=dtypes.weaks, src=(UPat.var("x"),)),), name="u"), - lambda u,x: x.cast(select_dtype(u)).cast(u.dtype) if x.dtype not in dtypes.weaks else None), - # Binary can widen from the bounds, all other nodes derive from the lowered sources. - # a weakfloat Unary (sin/exp2/...) must resolve here, before the transcendental decomposition - (UPat(GroupOp.Binary|GroupOp.Unary|{Ops.WHERE, Ops.RANGE, Ops.STACK, Ops.SPECIAL}, name="u"), lower_weak_node), - (UPat(Ops.PARAM, dtype=dtypes.weakint, name="u"), - lambda u: u.replace(dtype=None, arg=replace(u.arg, dtype=select_dtype(u))).cast(dtypes.weakint) if u.addrspace == AddrSpace.ALU else None), -]) -def lower_weak_srcs(ctx:dict[UOp, UOp]|None, u:UOp) -> UOp|None: - if ctx is None: ctx = {} - def lower(s:UOp) -> UOp: - if (r:=ctx.get(s)) is None: - r = graph_rewrite(s, pm_lower_weak) - # the consumer absorbs the cast on its own edge - ctx[s] = r = r.src[0] if r.op is Ops.CAST and r.dtype in dtypes.weaks else r - return r - # a comparison demands a common operand width: lower it whole so the Binary rule unifies its operands - ret = lower(u) if u.op in GroupOp.Comparison else u.replace(src=tuple(lower(s) if s.dtype in dtypes.weaks else s for s in u.src)) - return None if ret is u else ret - -def commit_weak(s:UOp, dt:DType) -> UOp: - # a bare weak CONST commits directly (its number must fit), a weak non-const src takes the demand cast - return UOp.const(s.val, dt) if s.op is Ops.CONST else s.cast(dt) - -def commit_weak_srcs(u:UOp) -> UOp|None: - if (dt:=least_upper_dtype(*(s.dtype for s in u.src))) in dtypes.weaks: return None - # the root re-derives: a shift's dtype is its lhs's, so committing the lhs commits the node too - return u.replace(dtype=None, src=tuple(commit_weak(s, dt) if s.dtype in dtypes.weaks else s for s in u.src)) - -# runs in index lowering and in the decomps: a rule that mints a weak const commits it in the same rewrite, so none reaches the renderer -pm_commit_weak = PatternMatcher([ - (UPat(GroupOp.Broadcastable, name="u"), commit_weak_srcs), - # demand from the destination: a STORE's weak value commits at the destination's dtype - (UPat(Ops.STORE, src=(UPat(), UPat(dtype=dtypes.weaks)), allow_any_len=True, name="u"), - lambda u: u.replace(src=(u.src[0], commit_weak(u.src[1], u.src[0].dtype), *u.src[2:]))), -]) - -# push cast to weak src -pm_cast_weak = PatternMatcher([ - (UPat(Ops.CAST, name="c", src=(UPat(GroupOp.Broadcastable, dtype=dtypes.weaks, name="u"),)), - lambda c,u: u.replace(dtype=None, src=tuple(commit_weak(s, c.dtype) if s.dtype in dtypes.weaks else s for s in u.src)).cast(c.dtype) - if c.dtype not in dtypes.weaks else None), -]) - -pm_lower_index_dtype = pm_commit_weak+PatternMatcher([ - (UPat(GroupOp.All, name="u"), - lambda ctx,u: lower_weak_srcs(ctx, u) if u.dtype not in dtypes.weaks and any(s.dtype in dtypes.weaks for s in u.src) else None), - # a valid index into an n-element buffer lives in [0,n): a gated long index narrows when n-1 fits int32 (out-of-gate wraps, discarded) - # TODO: more generic - (UPat((Ops.INDEX, Ops.SHRINK), src=(UPat.var("buf"), UPat.var("gate").where(UPat.var("idx", dtypes.long), UPat(Ops.CONST, arg=Invalid))), - allow_any_len=True, name="u"), - lambda u,buf,gate,idx: u.replace(src=(buf, idx.cast(dtypes.int).valid(gate))+u.src[2:]) if buf.max_numel()-1 <= dtypes.int32.max else None), -]) - _substitute = PatternMatcher([(UPat(tuple(Ops), name="x"), lambda ctx,x: ctx.get(x,None))]) _pm_resolve_params = PatternMatcher([(UPat(Ops.PARAM, name="p"), lambda ctx,p: ctx[p.arg.slot])]) remove_all_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)]) diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index a6af821b61..654f491cd7 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -32,8 +32,8 @@ def validate_index(uidx:UOp, gate:UOp|None=None): from tinygrad.uop.validate import validate_index_with_z3 return validate_index_with_z3(sz, idx, gate) -def type_verify(ast:UOp|list[UOp], check_spec:PatternMatcher): - lst = list(ast.toposort()) if isinstance(ast, UOp) else ast +def type_verify(ast:UOp|list[UOp], check_spec:PatternMatcher, enter_calls=True): + lst = list(ast.toposort(enter_calls=enter_calls)) if isinstance(ast, UOp) else ast if SPEC > 1: test_pyrender(lst[-1]) # assume this is the sink with Context(TRACK_MATCH_STATS=0): @@ -253,15 +253,44 @@ spec_full = PatternMatcher([ (UPat(Ops.BIND, (dtypes.int, dtypes.weakint), (UPat(), UPat()), arg=None), lambda: True), ])+spec_tensor+spec_program+spec_hcq +# ***** kernel graph spec ***** + +spec_kernel_graph = PatternMatcher([ + # sink + (UPat(Ops.SINK, dtypes.void), lambda: True), + # bind + (UPat(Ops.BIND), lambda: True), + # const + stack to make vconsts + (UPat(Ops.CONST, src=()), lambda: True), + (UPat(Ops.STACK, src=()), lambda: True), + (UPat(Ops.STACK, src=UPat((Ops.CONST, Ops.BIND, Ops.PARAM))), lambda: True), + # linear for more kernels (TODO: we should enter non sink calls) + #(UPat(Ops.LINEAR), lambda: True), + # param is outside buffer, buffer is local buffer + (UPat(Ops.PARAM, name="x"), lambda x: isinstance(x.arg, ParamArg)), + (UPat(Ops.BUFFER, name="x"), lambda x: isinstance(x.arg, ParamArg) and x.addrspace == AddrSpace.GLOBAL), + # RESHAPE/BITCAST are NOOPs in the kernel graph (do we need them?) + (UPat((Ops.RESHAPE, Ops.BITCAST)), lambda: True), + # mstack/mselect + (UPat(Ops.MSTACK, name="x"), lambda x: all(isinstance(s.device, str) for s in x.src) or (all_same(x.src) and x.src[0].device is None)), + (UPat(Ops.MSELECT, name="x"), lambda x: isinstance(x.src[0].device, tuple) and x.arg < len(x.src[0].device)), + # all calls are on various sinks + (UPat(Ops.CALL, src=(UPat((Ops.SINK, Ops.LINEAR, Ops.PROGRAM, Ops.CUSTOM_FUNCTION)),), allow_any_len=True), lambda: True), + # after on PARAM or AFTER + (UPat(Ops.AFTER, src=(UPat(GroupOp.Movement.union({Ops.PARAM, Ops.AFTER, Ops.BUFFER, Ops.MSTACK, Ops.MSELECT, Ops.BITCAST, Ops.RESHAPE})),), + allow_any_len=True, name="x"), lambda x: matches_dtype(x.src[0], x.dtype)), +]) + # **** pyrender (move this) **** # late imports to avoid circular import from tinygrad.codegen.opt import Opt, OptOps from tinygrad.schedule.rangeify import BufferizeOpts +from tinygrad.renderer import Estimates glbls:dict[str, Any] = {"inf": math.inf, "nan": math.nan, "KernelInfo": KernelInfo, "Metadata": Metadata, "UOp": UOp, "dtypes": dtypes, "Ops": Ops, "AxisType": AxisType, "Invalid": Invalid, "Opt": Opt, "OptOps": OptOps, "BufferizeOpts": BufferizeOpts, "AddrSpace": AddrSpace, "panic": panic, - "ConstFloat": ConstFloat, "ParamArg": ParamArg} + "ConstFloat": ConstFloat, "ParamArg": ParamArg, "Estimates": Estimates} def eval_pyrender(code:str) -> UOp: lcls:dict[str, Any] = {} exec(code, glbls, lcls) diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 1dce0bed20..f3d78ebfb5 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -68,7 +68,7 @@ invalid_pat = UPat(Ops.CONST, arg=Invalid, name="i") invalid_gate = UPat.var("cond").where(UPat.var("x"), invalid_pat) pm_data_invalid = PatternMatcher([ (invalid_pat.broadcast(), lambda i: i), - (UPat(GroupOp.Unary|{Ops.BITCAST}, src=(invalid_pat,)), lambda i: i), + (UPat(GroupOp.Unary|{Ops.CAST, Ops.BITCAST}, src=(invalid_pat,)), lambda i: i), (UPat(GroupOp.Unary|{Ops.CAST, Ops.BITCAST}, src=(invalid_gate,), name="op"), lambda cond,x,op,i: cond.where(op.replace(src=(x,)), i)), # binary ops move inside the gate, with Invalid in the false branch @@ -96,6 +96,10 @@ pm_remove_invalid = PatternMatcher([ if any(x.is_invalid for x in s.src) else None), ]) +# the one rule that collapses the pair CAST(dt, CONST(v)) into a typed CONST +# TODO: delete this once CONST has no dtype +pm_fold_cast_const = PatternMatcher([(UPat(Ops.CAST, name="root", src=(UPat.cvar("c"),)), lambda root, c: root.const_like(c.val))]) + symbolic_simple = pm_data_invalid + PatternMatcher([ # ** self folding ** (UPat.var("x") + 0, lambda x: x), # x+0 -> x @@ -152,8 +156,6 @@ symbolic_simple = pm_data_invalid + PatternMatcher([ (UPat.var("x") * 0, lambda x: x.const_like(float("nan") if x.op is Ops.CONST and isinstance(x.val, float) and (math.isnan(x.val) or math.isinf(x.val)) else 0)), # *** cast/bitcast *** - # TODO: delete this once CONST has no dtype - (UPat(Ops.CAST, name="root", src=(UPat.cvar("c"),)), lambda root, c: root.const_like(c.val)), (UPat((Ops.CAST, Ops.BITCAST), name="root"), lambda root: root.src[0] if root.dtype == root.src[0].dtype else None), (UPat(Ops.BITCAST, name="root", src=(UPat.cvar("c"),)), fold_bitcast), # b.cast(a).cast(b) -> b if a preserves all values in b @@ -253,10 +255,11 @@ symbolic = symbolic_simple+commutative+PatternMatcher([ # ** two stage ALU folding ** *((UPat.var("x").alu(op, UPat.cvar("c1")).alu(op, UPat.cvar("c2")).named("f"), lambda f,x,c1,c2: x.alu(f.op,c1.alu(f.op,c2))) for op in GroupOp.Associative), - ((UPat.cvar("c0") + UPat.var("x")) < UPat.cvar("c1"), lambda x,c0,c1: x<(c1-c0)), # c0 + x < c1 -> x < c1 - c0 # (x//c1)//c2 -> x//(c1*c2) for c2>0 ((UPat.var("x") // UPat.cvar("c1")) // UPat.cvar("c2"), lambda x,c1,c2: x//(c1*c2) if c2.vmin>0 else None), # ** lt ** + # c0+x x < c1-c0 + ((UPat.cvar("c0") + UPat.var("x", dtype=dtypes.ints+(dtypes.weakint,))) < UPat.cvar("c1"), lambda x,c0,c1: x<(c1-c0)), # c0*x sign(c0)*x < ceil(c1/abs(c0)) ((UPat.cvar("c0")*UPat.var("x", dtype=dtypes.weakint)) 0 else -x)<-(-c1.val//abs(c0.val)) if abs(c0.val) > 1 else None), @@ -285,7 +288,7 @@ symbolic = symbolic_simple+commutative+PatternMatcher([ (UOp.const(x.val) if x.op is Ops.CONST else x.cast(dtypes.int)).alu(u.op, UOp.const(y.val) if y.op is Ops.CONST else y.cast(dtypes.int)).cast(u.dtype) if not any(v.overflows(dtypes.int) for v in (u,x,y)) else None), - ((UPat.var("x", dtypes.weakint) + UPat.cvar("c")).cast(dtypes.sints, name="cast"), lambda x,c,cast:x.cast(cast.dtype)+c.cast(cast.dtype)), + ((UPat.var("x", dtypes.weakint) + UPat.cvar("c")).cast(dtypes.sints, name="cast"), lambda x,c,cast:x.cast(cast.dtype)+cast.const_like(c.val)), # only RANGE/IF/STORE/KERNEL have side effects (UPat(Ops.AFTER, name="x"), lambda x: x.replace(src=(x.src[0],)+ tuple(dedup(flatten([(y,) if y.op in {Ops.RANGE, Ops.STORE, Ops.CALL, Ops.FUNCTION, Ops.BARRIER, Ops.END, Ops.LINEAR, Ops.STAGE} diff --git a/tinygrad/uop/weak.py b/tinygrad/uop/weak.py new file mode 100644 index 0000000000..0b0e578cae --- /dev/null +++ b/tinygrad/uop/weak.py @@ -0,0 +1,78 @@ +from dataclasses import replace +from tinygrad.dtype import dtypes, DType, AddrSpace, Invalid, least_upper_dtype, strong_dtype, weak_dtype +from tinygrad.helpers import unwrap +from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, GroupOp, graph_rewrite, dtype_from_uop + +def select_dtype(u:UOp): + if u.dtype is dtypes.weakfloat: return dtypes.default_float + return dtypes.long if u.overflows(dtypes.int32) else dtypes.int + +def lower_weak_node(u:UOp) -> UOp|None: + start, src = (1 if u.op is Ops.WHERE else 0), tuple(s.src[0] if s.op is Ops.CAST and s.dtype in dtypes.weaks else s for s in u.src) + if src == u.src or any(s.dtype in dtypes.weaks for s in src[start:]): return None + dt = strong_dtype(least_upper_dtype(select_dtype(u), *(s.dtype for s in src)) if u.op in GroupOp.Binary + else unwrap(dtype_from_uop(u.op, src, u.arg))) + return u.replace(dtype=None, src=src[:start]+tuple(s if s.base.is_invalid else s.cast(dt) for s in src[start:])).cast(u.dtype) + +pm_lower_weak = PatternMatcher([ + (UPat(Ops.CONST, dtype=dtypes.weaks, name="u"), lambda u: UOp.const(u.val, select_dtype(u)).cast(u.dtype)), + # two stacked weak casts are a weakint value used as weakfloat (or vice versa): resolve the inner one at the outer kind's default. + # a SINGLE weak cast is never rewritten here, each consumer absorbs it on its own edge (see lower_weak_srcs) + (UPat(Ops.CAST, dtype=dtypes.weaks, src=(UPat(Ops.CAST, dtype=dtypes.weaks, src=(UPat.var("x"),)),), name="u"), + lambda u,x: x.cast(select_dtype(u)).cast(u.dtype) if x.dtype not in dtypes.weaks else None), + # Binary can widen from the bounds, all other nodes derive from the lowered sources. + # a weakfloat Unary (sin/exp2/...) must resolve here, before the transcendental decomposition + (UPat(GroupOp.Binary|GroupOp.Unary|{Ops.WHERE, Ops.RANGE, Ops.STACK, Ops.SPECIAL}, name="u"), lower_weak_node), + (UPat(Ops.PARAM, dtype=dtypes.weakint, name="u"), + lambda u: u.replace(dtype=None, arg=replace(u.arg, dtype=select_dtype(u))).cast(dtypes.weakint) if u.addrspace == AddrSpace.ALU else None), +]) + +def lower_weak_srcs(ctx:dict[UOp, UOp]|None, u:UOp) -> UOp|None: + if ctx is None: ctx = {} + def lower(s:UOp) -> UOp: + if (r:=ctx.get(s)) is None: + r = graph_rewrite(s, pm_lower_weak) + # the consumer absorbs the cast on its own edge + ctx[s] = r = r.src[0] if r.op is Ops.CAST and r.dtype in dtypes.weaks else r + return r + # a comparison demands a common operand width: lower it whole so the Binary rule unifies its operands + ret = lower(u) if u.op in GroupOp.Comparison else u.replace(src=tuple(lower(s) if s.dtype in dtypes.weaks else s for s in u.src)) + return None if ret is u else ret + +def commit_weak(s:UOp, dt:DType) -> UOp: + # a bare weak CONST commits directly (the value stays mathematical, emission truncates), a weak non-const src takes the demand cast + return UOp.const(s.val, dt) if s.op is Ops.CONST else s.cast(dt) + +def commit_weak_srcs(u:UOp) -> UOp|None: + if not any(s.dtype in dtypes.weaks for s in u.src): return None + if (dt:=least_upper_dtype(*(s.dtype for s in u.src))) in dtypes.weaks: return None + # the root re-derives: a shift's dtype is its lhs's, so committing the lhs commits the node too + return u.replace(dtype=None, src=tuple(commit_weak(s, dt) if s.dtype in dtypes.weaks else s for s in u.src)) + +# runs in index lowering and in the decomps: a rule that mints a weak const commits it in the same rewrite, so none reaches the renderer +pm_commit_weak = PatternMatcher([ + (UPat(GroupOp.Broadcastable, name="u"), commit_weak_srcs), + # demand from the destination: a STORE's weak value commits at the destination's dtype + (UPat(Ops.STORE, src=(UPat(), UPat(dtype=dtypes.weaks)), allow_any_len=True, name="u"), + lambda u: u.replace(src=(u.src[0], commit_weak(u.src[1], u.src[0].dtype), *u.src[2:]))), +]) + +# a concrete CAST over a weak node states the width the value will live at. that width is a floor, never a narrowing +def cast_weak_srcs(c:UOp, u:UOp) -> UOp|None: + if c.dtype in dtypes.weaks or weak_dtype(c.dtype) is not u.dtype: return None + dt = least_upper_dtype(c.dtype, select_dtype(u)) + return u.replace(dtype=None, src=tuple(commit_weak(s, dt) if s.dtype in dtypes.weaks else s for s in u.src)).cast(c.dtype) + +pm_cast_weak = PatternMatcher([ + (UPat(Ops.CAST, name="c", src=(UPat(GroupOp.ALU, dtype=dtypes.weaks, name="u"),)), cast_weak_srcs), +]) + +pm_lower_index_dtype = pm_commit_weak+pm_cast_weak+PatternMatcher([ + (UPat(GroupOp.All, name="u"), + lambda ctx,u: lower_weak_srcs(ctx, u) if u.dtype not in dtypes.weaks and any(s.dtype in dtypes.weaks for s in u.src) else None), + # a valid index into an n-element buffer lives in [0,n): a gated long index narrows when n-1 fits int32 (out-of-gate wraps, discarded) + # TODO: more generic + (UPat((Ops.INDEX, Ops.SHRINK), src=(UPat.var("buf"), UPat.var("gate").where(UPat.var("idx", dtypes.long), UPat(Ops.CONST, arg=Invalid))), + allow_any_len=True, name="u"), + lambda u,buf,gate,idx: u.replace(src=(buf, idx.cast(dtypes.int).valid(gate))+u.src[2:]) if buf.max_numel()-1 <= dtypes.int32.max else None), +]) diff --git a/tinygrad/viz/cli.py b/tinygrad/viz/cli.py index 50449843be..0e62ac451b 100755 --- a/tinygrad/viz/cli.py +++ b/tinygrad/viz/cli.py @@ -80,14 +80,14 @@ def main(args) -> None: def emit(val, to_str=str) -> str: return json.dumps(val if isinstance(val, dict) else {"value":val}) if args.json else to_str(val) def print_step(step:dict, print_graph=False, reconstruct_matches=False) -> None: - data = viz.get_render(viz_data, step["query"]) + data = viz.get_render(viz_data, step["query"], update_sink=False) if isinstance(data.get("value"), Iterator): for m in data["value"]: if print_graph and "graph" in m and not args.json: for k,v in m["graph"].items(): print(f"[{k}] {' '.join((lines:=v['label'].splitlines())[:5])}{'...' if len(lines) > 5 else ''}"+(f" tag={v['tag']}" if v['tag'] else '')) if v["src"]: - print(" src: "+", ".join([f"{i}->[{x}]" for i,x in v["src"][:5]])+(f", ... and {len(v['src'])-5} more" if len(v["src"]) > 5 else "")) + print(" src: "+", ".join([f"{i}->[{x}]" for i,x in v["src"]])) elif "uop" in m: print(emit(m["graph"] if print_graph else m["uop"])) if not reconstruct_matches: return None if m.get("diff"): diff --git a/tinygrad/viz/js/index.js b/tinygrad/viz/js/index.js index 0d1293f0d6..8ea2ef384a 100644 --- a/tinygrad/viz/js/index.js +++ b/tinygrad/viz/js/index.js @@ -120,9 +120,9 @@ const drawGraph = (data) => { .attr("transform", d => `translate(${d.width/2-8}, ${-d.height/2+8})`).datum(e => ({ rect:true, width:10, height:10, fill:e.addrspace, stroke:"none" }))); const CALL_TAG_WIDTH = 14; addTags(nodes.selectAll("g.type").data(d => d.collapsible ? [d] : []).join("g").attr("class", d => `tag clickable ${d.collapsed ? 'collapsed' : 'expanded'}`) - .attr("transform", d => d.callNode ? `translate(${CALL_TAG_WIDTH/2-d.width/2}, ${0})` : `translate(${-d.width/2}, ${0})`) - .datum(d => ({ ...d, text:d.collapsed ? "+" : "−", fill:d.callNode ? null : d.color, - ...(d.callNode && { rect:true, width:CALL_TAG_WIDTH }) })).on("click", (e,d) => { + .attr("transform", d => d.collapsePorts != null ? `translate(${CALL_TAG_WIDTH/2-d.width/2}, ${0})` : `translate(${-d.width/2}, ${0})`) + .datum(d => ({ ...d, text:d.collapsed ? "+" : "−", fill:d.collapsePorts != null ? null : d.color, + ...(d.collapsePorts != null && { rect:true, width:CALL_TAG_WIDTH }) })).on("click", (e,d) => { e.stopPropagation(); const t = d3.zoomTransform(document.getElementById("graph-svg")); const [x, y] = t.apply([d.x, d.y]); diff --git a/tinygrad/viz/js/worker.js b/tinygrad/viz/js/worker.js index 5a73d1b64e..00de46c5ff 100644 --- a/tinygrad/viz/js/worker.js +++ b/tinygrad/viz/js/worker.js @@ -54,15 +54,17 @@ const layoutUOp = (g, { graph, change }, opts) => { width = Math.max(width, ctx.measureText(line).width); height += lineHeight; } - const callNode = label.startsWith("CALL\n") || label.startsWith("FUNCTION\n"); + const op = label.split("\n", 1)[0]; + const callNode = op === "CALL" || op === "FUNCTION", programNode = op === "PROGRAM"; + const collapsePorts = callNode ? [0] : programNode ? [0, 1] : null; if (callNode) callCount++; - g.setNode(k, {...rectDims(width, height), label, labelX:0, ref, id:k, color, tag, callNode, exclude, addrspace, + g.setNode(k, {...rectDims(width, height), label, labelX:0, ref, id:k, color, tag, callNode, collapsePorts, exclude, addrspace, className:label.startsWith("REWRITE_ERROR") ? "err" : null}); // add edges const edgeCounts = {}; for (const [_, s] of src) edgeCounts[s] = (edgeCounts[s] || 0)+1; for (const [port, s] of src) g.setEdge(s, k, { label: edgeCounts[s] > 1 ? {type:"tag", text:edgeCounts[s]} : {type:"port", text:port}, - ...(callNode && port === 0 && {color:"#a0a1b8"})}); + ...(collapsePorts?.includes(port) && {color:"#a0a1b8"})}); if (change?.includes(parseInt(k))) g.setParent(k, "overlay"); } // optionally hide nodes from the layout @@ -87,11 +89,11 @@ const layoutUOp = (g, { graph, change }, opts) => { const consumer = g.node(consumerId); // add +- toggle if this consumer has collapsible sources const edge = g.edge(n, consumerId); - const collapsible = consumer.callNode ? edge?.label?.text === 0 : node.exclude; + const collapsible = consumer.collapsePorts != null ? consumer.collapsePorts.includes(edge?.label?.text) : node.exclude; if (!collapsible) continue; consumer.collapsible = true; - // increase width of call/function nodes to make space for a toggle - if (consumer.callNode) { consumer.width = consumer.labelWidth+NODE_PADDING*2+CALL_TAG_WIDTH; consumer.labelX = CALL_TAG_WIDTH/2; } + // increase width of call/function/program nodes to make space for a toggle + if (consumer.collapsePorts != null) { consumer.width = consumer.labelWidth+NODE_PADDING*2+CALL_TAG_WIDTH; consumer.labelX = CALL_TAG_WIDTH/2; } // make sources invisible if UI has toggled it off const collapsed = consumer.callNode ? opts.showCallSrc === opts.callSrcMask.has(consumerId) : !opts.expandedNodes.has(consumerId); if (!collapsed) continue; diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index 7e41d8a8c3..506fcac00c 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -177,14 +177,14 @@ def _reconstruct(data:VizData, a:int, depth:int|None=None): if depth is None: data.all_uops[a] = ret return ret -def get_full_rewrite(data:VizData, ctx:TrackedGraphRewrite, depth:int|None=None) -> Generator[GraphRewriteDetails, None, None]: +def get_full_rewrite(data:VizData, ctx:TrackedGraphRewrite, depth:int|None=None, update_sink=True) -> Generator[GraphRewriteDetails, None, None]: next_sink, err = _reconstruct(data, ctx.sink, depth=depth), False yield {"graph":uop_to_json(data, next_sink), "uop":pystr(next_sink), "change":None, "diff":None, "upat":None, "_sink":next_sink} replaces: dict[UOp, UOp] = {} for u0_num,u1_num,upat_loc,dur in ctx.matches: if err: break replaces[u0:=_reconstruct(data, u0_num, depth=depth)] = u1 = _reconstruct(data, u1_num, depth=depth) - try: new_sink = next_sink.substitute(replaces, walk=ctx.walk, enter_calls=ctx.enter_calls) + try: new_sink = next_sink.substitute(replaces, walk=ctx.walk, enter_calls=ctx.enter_calls) if update_sink else next_sink except RuntimeError: new_sink, err = UOp(Ops.REWRITE_ERROR, arg=traceback.format_exc()), True match_repr = f"# {dur*1e6:.2f} us\n"+printable(upat_loc) yield {"graph":(sink_json:=uop_to_json(data, new_sink)), "uop":pystr(new_sink), "change":[id(x) for x in u1.toposort() if id(x) in sink_json], @@ -611,11 +611,11 @@ def amdgpu_cfg(lib:bytes, target:str) -> dict: # ** Main render function to get the complete details about a trace event -def get_render(viz_data:VizData, query:str) -> dict: +def get_render(viz_data:VizData, query:str, **kwargs) -> dict: url = urlparse(query) i, j, fmt = get_int(qs:=parse_qs(url.query), "ctx"), get_int(qs, "step"), url.path.lstrip("/") data = viz_data.ctxs[i]["steps"][j]["_data"] - if fmt == "graph-rewrites": return {"value":get_full_rewrite(viz_data, viz_data.trace.rewrites[i][j]), "content_type":"text/event-stream"} + if fmt == "graph-rewrites": return {"value":get_full_rewrite(viz_data, viz_data.trace.rewrites[i][j], **kwargs), "content_type":"text/event-stream"} if fmt == "uops": if (sink:=get_sink_at(("do_linearize",), viz_data, i, data)) is None: return {"src":"No linear found"} return {"src":sink.arg} if sink.op is Ops.REWRITE_ERROR else {"src":get_stdout(lambda: print_uops(list(unwrap(sink).src[1].src)))}