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 730585673f..7fe9471032 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) @@ -589,7 +591,7 @@ jobs: 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 test/opt/test_tensor_cores.py --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 --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 @@ -632,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: @@ -827,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 15b5c2f632..b472e30939 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 @@ -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} @@ -26,7 +26,7 @@ export FUSED_SILU_W13=${FUSED_SILU_W13:-1} export SPLIT_W13=${SPLIT_W13:-0} export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0} -export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16" +export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="float32" export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2} export GBS=$((BS * GRADIENT_ACC_STEPS)) 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..bb53cc0ded 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} @@ -26,7 +26,7 @@ export FUSED_SILU_W13=${FUSED_SILU_W13:-1} export SPLIT_W13=${SPLIT_W13:-0} export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0} -export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16" +export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="float32" export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2} export GBS=$((BS * GRADIENT_ACC_STEPS)) 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 4376e1083d..fbdef9f444 100644 --- a/extra/hcq2/ops_amd2.py +++ b/extra/hcq2/ops_amd2.py @@ -288,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) @@ -539,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 @@ -578,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) @@ -691,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/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/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_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/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/helpers.py b/test/helpers.py index ea900d505c..21febed6bc 100644 --- a/test/helpers.py +++ b/test/helpers.py @@ -48,6 +48,8 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te else: assert isinstance(t, UOp), f"can't schedule {t}" linear, var_vals = Tensor(t).linear_with_vars() + # test compiling the linear + compile_linear(linear) kernel_cnt = sum((len(call.device) if isinstance(call.device, tuple) else 1) for call in linear.src if call.src[0].op is Ops.SINK or not filter_sink) if kernel_cnt != allowed: @@ -57,8 +59,6 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te print("kernel", i+1) print(call.src[0]) raise KernelCountException(allowed, kernel_cnt) - # test compiling the linear - compile_linear(linear) return linear, var_vals def assert_kernel_count(expected:int): diff --git a/test/mockgpu/amd/emu.py b/test/mockgpu/amd/emu.py index 266e1f6d84..3f36e331e9 100644 --- a/test/mockgpu/amd/emu.py +++ b/test/mockgpu/amd/emu.py @@ -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) 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_schedule.py b/test/null/test_schedule.py index bcd9487c4f..614e2cf837 100644 --- a/test/null/test_schedule.py +++ b/test/null/test_schedule.py @@ -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 60891d0978..754d713fe2 100644 --- a/test/null/test_simplify_valid_idx.py +++ b/test/null/test_simplify_valid_idx.py @@ -333,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)") @@ -461,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): @@ -480,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_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/opt/test_tensor_cores.py b/test/opt/test_tensor_cores.py index c25eec469a..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): @@ -185,10 +186,10 @@ class TestTensorCores(unittest.TestCase): # 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(64, 64, dtype=tc.dtype_in), Tensor.rand(64, 64, dtype=tc.dtype_in) + 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, 2)] - ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3) + 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 @@ -199,10 +200,10 @@ 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(self): 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(64, 64, dtype=tc.dtype_in), Tensor.rand(64, 64, dtype=tc.dtype_in) + 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, 2)] - ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3) + 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])) @@ -215,10 +216,10 @@ class TestTensorCores(unittest.TestCase): 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 and tc.dtype_in not in dtypes.fp8s][0] - x, y = Tensor.rand(64, 64, dtype=tc.dtype_in), Tensor.rand(64, 64, dtype=tc.dtype_in) + 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, 2)] - ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3) + 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 3abe561c8c..12e55f12d7 100644 --- a/test/unit/test_call.py +++ b/test/unit/test_call.py @@ -359,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_multitensor.py b/test/unit/test_multitensor.py index 2eb12e8db4..139c4588ff 100644 --- a/test/unit/test_multitensor.py +++ b/test/unit/test_multitensor.py @@ -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 diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index e50a8ce3a8..7d8c0dbd37 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -330,10 +330,10 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: sink = graph_rewrite(sink, symbolic_simple+pm_expand_broadcast+pm_add_loads, name="*** expand broadcast / add loads") # devectorize - sink = graph_rewrite(sink, symbolic_simple+pm_fold_cast_const+devectorizer2+indexing_simplify, ctx=ren, name="devectorize2") + sink = graph_rewrite(sink, symbolic_simple+devectorizer2+indexing_simplify, ctx=ren, name="devectorize2") # some coalescing misses without this - sink = graph_rewrite(sink, sym+pm_fold_cast_const, name="early symbolic") + sink = graph_rewrite(sink, sym, name="early symbolic") # do memory coalescing (late) sink = memory_coalescing(sink, ren) @@ -341,11 +341,12 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: 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") diff --git a/tinygrad/codegen/decomp/dtype.py b/tinygrad/codegen/decomp/dtype.py index 5d30df1812..52b8e46304 100644 --- a/tinygrad/codegen/decomp/dtype.py +++ b/tinygrad/codegen/decomp/dtype.py @@ -99,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) 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/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/realize.py b/tinygrad/engine/realize.py index 9c0533625a..92006e87d6 100644 --- a/tinygrad/engine/realize.py +++ b/tinygrad/engine/realize.py @@ -8,7 +8,7 @@ from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer 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 **************** @@ -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 @@ -220,11 +221,16 @@ def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> float|None: 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: cast(Any, Device[d]).synchronize(timeout=ctx.timeout) - 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([ @@ -266,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 @@ -288,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/mixin/elementwise.py b/tinygrad/mixin/elementwise.py index aa746bb454..f06fd5aa52 100644 --- a/tinygrad/mixin/elementwise.py +++ b/tinygrad/mixin/elementwise.py @@ -221,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 380045497e..b916f8af9d 100644 --- a/tinygrad/mixin/movement.py +++ b/tinygrad/mixin/movement.py @@ -268,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/runtime/ops_cpu.py b/tinygrad/runtime/ops_cpu.py index 09a77e6a41..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) diff --git a/tinygrad/runtime/support/hcq2.py b/tinygrad/runtime/support/hcq2.py index 8ff308fb60..9b0372d035 100644 --- a/tinygrad/runtime/support/hcq2.py +++ b/tinygrad/runtime/support/hcq2.py @@ -1,10 +1,11 @@ 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.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 @@ -32,6 +33,7 @@ class HCQInfo: 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 @@ -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 @@ -170,7 +184,7 @@ def _build_finalizers(batch:list[tuple[UOp, tuple[str, ...]]], batch_info:list[t 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, ...]]]) -> list[UOp]: +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 @@ -188,7 +202,7 @@ def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]]) -> list[UOp]: 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: @@ -197,20 +211,27 @@ def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]]) -> list[UOp]: # and make hcq call name, info = get_call_name(call, get_call_arg_uops(call)), HCQInfo(devices, estimate_uop(call)) - q += [call.replace(arg=replace(call.arg, aux=info))] + 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 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 @@ -246,7 +267,7 @@ def merge_queues(linear:UOp) -> UOp: return linear.replace(src=tuple(new_src + [_merged_hcq_call(c) for c in opened_qs.values()])) pm_schedule_and_merge = PatternMatcher([(UPat(Ops.LINEAR, name="l"), - lambda ctx, l: merge_queues(sched_hcq_batches(l).substitute(ctx, walk=True, enter_calls=True)))]) + 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 @@ -395,21 +416,21 @@ 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] = {} -@rewrite_group(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) @@ -516,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: @@ -537,6 +579,7 @@ class HCQ2Compiled(Compiled): 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/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 740c66ebe7..b084b1fb51 100644 --- a/tinygrad/schedule/indexing.py +++ b/tinygrad/schedule/indexing.py @@ -7,27 +7,50 @@ 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 -BUFFER_STATE_OPS: set[Ops] = {Ops.AFTER, Ops.BUFFER, Ops.PARAM, Ops.MSELECT, Ops.MSTACK, Ops.BIND} +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 @@ -43,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) @@ -86,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) @@ -107,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]]) @@ -184,7 +194,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]: 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 04840084ff..1dd683e3c0 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -6,7 +6,7 @@ from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, K 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, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS +from tinygrad.helpers import prod, getenv, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY, 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 @@ -141,7 +141,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 ** @@ -193,6 +193,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 @@ -326,6 +327,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 @@ -562,7 +583,8 @@ def get_kernel_graph(sink:UOp) -> UOp: # convert movement ops to ranges tsink, rctx = run_rangeify(tsink, bool(DEBUG_RANGEIFY)) - tsink = graph_rewrite(tsink, symbolic+pm_fold_cast_const+pm_reduce_simplify+pm_const_buffer_folding+pm_remove_bufferize, + tsink = graph_rewrite(tsink, + symbolic+pm_fold_cast_const+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") @@ -575,4 +597,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 33db031cfd..304d1fe1a3 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -184,7 +184,12 @@ def finalize_after(ctx:AllocCtx, x:UOp): # 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] + # 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) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 772db6b993..9ae599b987 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -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, ...]: @@ -1097,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 @@ -1187,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" 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 3e325dd3da..f3d78ebfb5 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -255,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), @@ -292,8 +293,8 @@ symbolic = symbolic_simple+commutative+PatternMatcher([ (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} else y.src for y in x.src[1:]]))))), - # after with 1 src is just src[0] - (UPat(Ops.AFTER, src=(UPat.var("s"),)), lambda s: s), + # after/end with 1 src is just src[0] + (UPat((Ops.AFTER, Ops.END), src=(UPat.var("s"),)), lambda s: s), ])+div_and_mod_symbolic # ******** we take a small aside to "simplify_valid" to rewrite valids ******** 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;