diff --git a/.github/actions/process-replay/action.yml b/.github/actions/process-replay/action.yml index b2d1374d8ca83..a5cec6ecc241b 100644 --- a/.github/actions/process-replay/action.yml +++ b/.github/actions/process-replay/action.yml @@ -5,6 +5,7 @@ runs: steps: - name: Run process replay tests shell: bash + if: env.CAPTURE_PROCESS_REPLAY == '1' run: | export PR_TITLE=$(jq -r .pull_request.title "$GITHUB_EVENT_PATH") export CURRENT_SHA=${{ github.event.pull_request && github.event.pull_request.head.sha || github.sha }} diff --git a/.github/actions/setup-tinygrad/action.yml b/.github/actions/setup-tinygrad/action.yml index a22f06f004c76..3d061fcb92712 100644 --- a/.github/actions/setup-tinygrad/action.yml +++ b/.github/actions/setup-tinygrad/action.yml @@ -42,7 +42,7 @@ inputs: required: false default: 'false' mesa: - description: "Install mesa" + description: "Install mesa (true, false, cpu)" required: false default: 'false' tinydreno: @@ -228,6 +228,11 @@ runs: sudo chown -R $USER:$USER /var/cache/apt/archives/ + - name: Add clang to PATH (Linux) + if: inputs.llvm == 'true' && runner.os == 'Linux' + shell: bash + run: echo "/usr/lib/llvm-20/bin" >> "$GITHUB_PATH" + # **** AMD **** - name: Setup AMD (Linux) if: inputs.amd == 'true' && runner.os == 'Linux' @@ -281,7 +286,7 @@ runs: cache-name: cache-gpuocelot-build-1 with: path: ${{ github.workspace }}/gpuocelot/ocelot - key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.CACHE_VERSION }} + key: ${{ runner.os }}-gpuocelot-f463259669c69abce7b3a0567b6c284f348d0f32-rebuild-${{ env.CACHE_VERSION }} - name: Cache gpuocelot if: inputs.ocelot == 'true' && github.event_name != 'pull_request' id: cache-build @@ -290,14 +295,14 @@ runs: cache-name: cache-gpuocelot-build-1 with: path: ${{ github.workspace }}/gpuocelot/ocelot - key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.CACHE_VERSION }} + key: ${{ runner.os }}-gpuocelot-f463259669c69abce7b3a0567b6c284f348d0f32-rebuild-${{ env.CACHE_VERSION }} - name: Clone/compile gpuocelot if: inputs.ocelot == 'true' && steps.cache-build-pr.outputs.cache-hit != 'true' && steps.cache-build.outputs.cache-hit != 'true' shell: bash run: | - git clone --recurse-submodules https://github.com/gpuocelot/gpuocelot.git ${{ github.workspace }}/gpuocelot + git clone --recurse-submodules https://github.com/tinygrad/gpuocelot.git ${{ github.workspace }}/gpuocelot cd ${{ github.workspace }}/gpuocelot/ocelot - git checkout b16039dc940dc6bc4ea0a98380495769ff35ed99 + git checkout f463259669c69abce7b3a0567b6c284f348d0f32 mkdir build cd build @@ -306,10 +311,7 @@ runs: sudo xcode-select -s /Applications/Xcode_16.2.app/Contents/Developer CMAKE_ARGS="$CMAKE_ARGS -DBoost_INCLUDE_DIR=$(brew --prefix boost)/include -DBoost_LIBRARY_DIR=$(brew --prefix boost)/lib" else - curl -fL https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvcc/linux-x86_64/cuda_nvcc-linux-x86_64-11.5.119-archive.tar.xz \ - | sudo tar -xJ -C /usr/ --strip-components=1 - curl -fL https://developer.download.nvidia.com/compute/cuda/redist/cuda_cudart/linux-x86_64/cuda_cudart-linux-x86_64-11.5.117-archive.tar.xz \ - | sudo tar -xJ -C /usr/ --strip-components=1 + CMAKE_ARGS="$CMAKE_ARGS -DLLVM_DIR=$(llvm-config-15 --cmakedir)" fi cmake .. $CMAKE_ARGS @@ -345,13 +347,13 @@ runs: # **** mesa **** - name: Install mesa (linux) - if: inputs.mesa == 'true' && runner.os == 'Linux' + if: inputs.mesa != 'false' && runner.os == 'Linux' shell: bash - run: sudo curl -fL https://github.com/sirhcm/tinymesa/releases/download/v1/libtinymesa_cpu-mesa-25.2.7-linux-amd64.so -o /usr/lib/libtinymesa_cpu.so + run: sudo curl -fL https://github.com/sirhcm/tinymesa/releases/download/v1/libtinymesa${{ inputs.mesa == 'cpu' && '_cpu' || '' }}-mesa-25.2.7-linux-amd64.so -o /usr/lib/libtinymesa${{ inputs.mesa == 'cpu' && '_cpu' || '' }}.so - name: Install mesa (macOS) - if: inputs.mesa == 'true' && runner.os == 'macOS' + if: inputs.mesa != 'false' && runner.os == 'macOS' shell: bash - run: brew install sirhcm/tinymesa/tinymesa_cpu + run: brew install sirhcm/tinymesa/tinymesa${{ inputs.mesa == 'cpu' && '_cpu' || '' }} # *** tinydreno *** - name: Install tinydreno (linux) diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index 9207c25317ded..62bd91e069444 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -25,7 +25,7 @@ jobs: CI: "" CAPTURE_PROCESS_REPLAY: "0" runs-on: [self-hosted, macOS] - timeout-minutes: 3 + timeout-minutes: 4 defaults: run: shell: bash -e -o pipefail {0} @@ -625,7 +625,7 @@ jobs: - name: IR3 openpilot compile3 0.11.0 driving_vision run: BENCHMARK_LOG=ir3_openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM:IR3 FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx - name: openpilot compile3 0.11.0 driving_policy - run: BENCHMARK_LOG=openpilot_0_11_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_policy.onnx + run: BENCHMARK_LOG=openpilot_0_11_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3.2 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_policy.onnx - name: openpilot compile3 0.11.0 dmonitoring run: BENCHMARK_LOG=openpilot_0_11_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/dmonitoring_model.onnx - name: DEBUG=2 openpilot compile3 0.10.1 driving_vision @@ -633,7 +633,7 @@ jobs: - name: openpilot compile3 0.10.1 driving_vision run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx - name: openpilot compile3 0.10.1 driving_policy - run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx + run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3.2 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx - name: openpilot compile3 0.10.1 dmonitoring run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx - name: benchmark MobileNetV2 on DSP @@ -806,3 +806,16 @@ jobs: pkill -f 'extra/remote/serve.py' || true - name: Run process replay tests uses: ./.github/actions/process-replay + + llvmspeed: + name: LLVM Speed + runs-on: [self-hosted, Linux, tinyboxrandom] + timeout-minutes: 20 + if: github.repository_owner == 'tinygrad' + steps: + - name: Checkout Code + uses: actions/checkout@v6 + - name: Speed Test + run: DEV=CPU:LLVM THREADS=0 python3 test/speed/external_test_speed_v_torch.py + - name: Speed Test (BEAM=2) + run: BEAM=2 DEV=CPU:LLVM THREADS=0 python3 test/speed/external_test_speed_v_torch.py diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 9070b58ceaa80..524af917b400c 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -2,7 +2,7 @@ name: Unit Tests env: # increment this when downloads substantially change to avoid the internet CACHE_VERSION: '19' - CAPTURE_PROCESS_REPLAY: 1 + 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 @@ -14,28 +14,14 @@ on: pull_request: workflow_dispatch: -jobs: - llvmspeed: - name: LLVM Speed - runs-on: ubuntu-24.04 - timeout-minutes: 20 - steps: - - name: Checkout Code - uses: actions/checkout@v6 - - name: Setup Environment - uses: ./.github/actions/setup-tinygrad - with: - key: llvm-speed - deps: testing_unit - llvm: 'true' - - name: Speed Test - run: DEV=CPU:LLVM THREADS=0 python3 test/speed/external_test_speed_v_torch.py - - name: Speed Test (BEAM=2) - run: BEAM=2 DEV=CPU:LLVM THREADS=0 python3 test/speed/external_test_speed_v_torch.py +concurrency: + group: test-${{ 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: docs: name: Docs - runs-on: ubuntu-22.04 + runs-on: &linux ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }} timeout-minutes: 10 env: CHECK_OOB: 0 @@ -89,7 +75,7 @@ jobs: torchbackend: name: Torch Backend Tests - runs-on: ubuntu-latest + runs-on: *linux timeout-minutes: 15 steps: - name: Checkout Code @@ -125,7 +111,7 @@ jobs: torchbackendmore: name: Torch Backend Tests More - runs-on: ubuntu-latest + runs-on: *linux timeout-minutes: 15 steps: - name: Checkout Code @@ -147,7 +133,7 @@ jobs: bepython: name: Python Backend - runs-on: ubuntu-latest + runs-on: *linux timeout-minutes: 15 steps: - name: Checkout Code @@ -215,7 +201,7 @@ jobs: linter: name: Linters - runs-on: ubuntu-latest + runs-on: *linux timeout-minutes: 10 steps: @@ -246,7 +232,7 @@ jobs: nulltest: name: Null Tests - runs-on: ubuntu-latest + runs-on: *linux timeout-minutes: 15 steps: @@ -277,7 +263,7 @@ jobs: unittest: name: Unit Tests - runs-on: ubuntu-latest + runs-on: *linux timeout-minutes: 15 steps: @@ -290,7 +276,6 @@ jobs: pydeps: "pillow ftfy regex pre-commit" deps: testing_unit llvm: 'true' - amd: 'true' - name: Run pre-commit test hooks run: SKIP=ruff,mypy pre-commit run --all-files - name: Check Device.DEFAULT @@ -312,8 +297,8 @@ jobs: python extra/optimization/extract_dataset.py gzip -c /tmp/sops > extra/datasets/sops.gz #DEBUG=1 MIN_ASTS=1 python extra/optimization/get_action_space.py - - name: Repo line count < 24000 lines - run: MAX_LINE_COUNT=24000 python sz.py + - name: Repo line count < 25000 lines + run: MAX_LINE_COUNT=25000 python sz.py spec: strategy: @@ -321,7 +306,7 @@ jobs: matrix: group: [1, 2] name: SPEC=2 (${{ matrix.group }}) - runs-on: ubuntu-latest + runs-on: *linux timeout-minutes: 15 steps: - name: Checkout Code @@ -332,12 +317,13 @@ jobs: key: spec-unit deps: testing_unit python-version: '3.14' + llvm: 'true' - name: Test SPEC=2 run: SPEC=2 pytest --maxfail=10 -n auto --durations=30 test/unit test/backend test/opt --ignore test/backend/test_custom_kernel.py --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" -k "not test_conv2d_ceildiv_edge_case" --splits 2 --group ${{ matrix.group }} fuzzing: name: Fuzzing - runs-on: ubuntu-latest + runs-on: *linux timeout-minutes: 10 steps: - name: Checkout Code @@ -358,7 +344,7 @@ jobs: testopenclimage: name: CL IMAGE Tests - runs-on: ubuntu-22.04 + runs-on: *linux timeout-minutes: 15 steps: - name: Checkout Code @@ -378,7 +364,7 @@ jobs: testgpumisc: name: CL Misc tests - runs-on: ubuntu-22.04 + runs-on: *linux timeout-minutes: 10 steps: - name: Checkout Code @@ -403,7 +389,7 @@ jobs: testopenpilot: name: openpilot Compile Tests - runs-on: ubuntu-22.04 + runs-on: *linux timeout-minutes: 15 steps: - name: Checkout Code @@ -418,10 +404,10 @@ jobs: - name: Test openpilot model kernel count and gate usage run: | ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1468 ALLOWED_GATED_READ_IMAGE=18 FLOAT16=1 DEV=CL IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916 - - name: Test openpilot CL compile fp16 - run: FLOAT16=1 DEV=CL IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916 - name: Test openpilot CL compile fp32 (test correctness) - run: DEV=CL IMAGE=1 SELFTEST=1 python examples/openpilot/compile3.py https://github.com/haraschax/filedump/raw/refs/heads/master/driving_vision_fp32.onnx + run: | + DEV=CL IMAGE=1 SELFTEST=1 python examples/openpilot/compile3.py https://github.com/haraschax/filedump/raw/refs/heads/master/driving_vision_fp32.onnx + DEV=CL IMAGE=1 SELFTEST=1 RUN_PICKLE=1 python examples/openpilot/compile3.py https://github.com/haraschax/filedump/raw/refs/heads/master/driving_vision_fp32.onnx - name: Test openpilot LLVM compile fp16 run: IMAGE=1 FLOAT16=1 DEV=CPU:LLVM python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916 - name: Run process replay tests @@ -431,7 +417,7 @@ jobs: testonnxcpu: name: ONNX (CPU) Tests - runs-on: ubuntu-22.04 + runs-on: *linux timeout-minutes: 20 steps: @@ -459,7 +445,7 @@ jobs: testopencl: name: ONNX (CL)+Optimization Tests - runs-on: ubuntu-22.04 + runs-on: *linux timeout-minutes: 20 steps: - name: Checkout Code @@ -493,7 +479,7 @@ jobs: testllm: name: Test LLM - runs-on: ubuntu-24.04 + runs-on: *linux timeout-minutes: 15 env: CHECK_OOB: 0 @@ -518,7 +504,7 @@ jobs: testmodels: name: Models (llvm+cpu+gpu) - runs-on: ubuntu-22.04 + runs-on: *linux timeout-minutes: 15 steps: - name: Checkout Code @@ -559,30 +545,9 @@ jobs: # ****** Feature Tests ****** - testdevectorize: - name: Linux (devectorize) - runs-on: ubuntu-24.04 - timeout-minutes: 15 - steps: - - name: Checkout Code - uses: actions/checkout@v6 - - name: Setup Environment - uses: ./.github/actions/setup-tinygrad - with: - key: devectorize-minimal - deps: testing_unit - pydeps: "pillow" - llvm: "true" - - name: Test LLVM=1 DEVECTORIZE=0 - run: DEV=CPU:LLVM DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/backend/test_ops.py - - name: Test LLVM=1 DEVECTORIZE=0 for model - run: DEV=CPU:LLVM DEVECTORIZE=0 python3 test/models/test_efficientnet.py - - name: Test DEV=CPU DEVECTORIZE=0 - run: DEV=CPU DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/backend/test_ops.py - testdsp: name: Linux (DSP) - runs-on: ubuntu-24.04 + runs-on: *linux timeout-minutes: 15 steps: - name: Checkout Code @@ -591,8 +556,7 @@ jobs: uses: ./.github/actions/setup-tinygrad with: key: dsp-minimal - deps: testing_unit - pydeps: "onnx==1.18.0 onnxruntime ml_dtypes" + deps: testing llvm: "true" qemu: "true" - name: Set MOCKDSP env @@ -600,13 +564,24 @@ jobs: - name: Run test_tiny on DSP run: DEBUG=2 DEV=DSP python test/test_tiny.py - name: Test transcendentals - run: CC=clang-20 DEBUG=2 DEV=DSP python test/backend/test_transcendental.py TestTranscendentalVectorized + run: DEBUG=2 DEV=DSP python test/backend/test_transcendental.py TestTranscendentalVectorized - name: Test quantize onnx run: DEBUG=2 DEV=DSP python3 test/backend/test_quantize_onnx.py - testwebgpu: - name: Linux (WebGPU) - runs-on: ubuntu-22.04 + testlinux: + strategy: + fail-fast: false + matrix: + dev: + - 'CPU:CLANG' + - 'CPU:LLVM' + - 'CPU:LVP' + - 'CPU:X86' + - 'CL' + - 'WEBGPU' + + name: Linux (DEV=${{ matrix.dev }}) + runs-on: *linux timeout-minutes: 20 steps: - name: Checkout Code @@ -614,23 +589,27 @@ jobs: - name: Setup Environment uses: ./.github/actions/setup-tinygrad with: - key: webgpu-minimal + key: linux-${{ matrix.dev }} deps: testing_unit python-version: '3.12' - webgpu: 'true' - - name: Check Device.DEFAULT (WEBGPU) and print some source - run: | - DEV=WEBGPU python -c "from tinygrad import Device; assert Device.DEFAULT == 'WEBGPU', Device.DEFAULT" - DEV=WEBGPU DEBUG=4 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus - - name: Run selected webgpu tests + llvm: ${{ contains(matrix.dev, 'LLVM') || contains(matrix.dev, 'LVP') || contains(matrix.dev, 'CLANG') }} + mesa: ${{ contains(matrix.dev, 'LVP') && 'cpu' || 'false' }} + webgpu: ${{ matrix.dev == 'WEBGPU' }} + opencl: ${{ matrix.dev == 'CL' }} + - 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: | - DEV=WEBGPU WEBGPU_BACKEND="WGPUBackendType_Vulkan" python3 -m pytest -n=auto test/backend --durations=20 + 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 backend tests + run: python -m pytest -n=auto test/backend --durations=20 - name: Run process replay tests uses: ./.github/actions/process-replay testamdasm: name: AMD ASM IDE - runs-on: ubuntu-24.04 + runs-on: *linux timeout-minutes: 20 env: DEV: MOCKKFD+AMD @@ -677,7 +656,7 @@ jobs: testmockam: name: Linux (am) - runs-on: ubuntu-24.04 + runs-on: *linux timeout-minutes: 15 env: DEV: MOCKPCI+AMD @@ -713,7 +692,7 @@ jobs: arch: [gfx1100, gfx1201, gfx950] name: Linux (${{ matrix.backend }} ${{ matrix.arch }}) - runs-on: ubuntu-22.04 + runs-on: *linux timeout-minutes: 15 env: DEV: MOCKKFD+AMD:${{ matrix.backend == 'amdllvm' && 'LLVM' || '' }}:${{ matrix.arch }} @@ -748,7 +727,7 @@ jobs: backend: [ptx, nv] name: Linux (${{ matrix.backend }}) - runs-on: ubuntu-22.04 + runs-on: *linux timeout-minutes: 20 env: FORWARD_ONLY: 1 @@ -776,39 +755,6 @@ jobs: - name: Run process replay tests uses: ./.github/actions/process-replay - testcpuopencl: - strategy: - fail-fast: false - matrix: - backend: [llvm, cpu, opencl, lvp] - - name: Linux (${{ matrix.backend }}) - runs-on: ubuntu-22.04 - timeout-minutes: 20 - steps: - - name: Checkout Code - uses: actions/checkout@v6 - - name: Setup Environment - uses: ./.github/actions/setup-tinygrad - with: - key: ${{ matrix.backend }}-minimal - deps: testing_unit - opencl: ${{ matrix.backend == 'opencl' && 'true' }} - llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }} - mesa: ${{ matrix.backend == 'lvp' && 'true' }} - - name: Set env - run: printf "${{ matrix.backend == 'llvm' && 'DEV=CPU:LLVM' || matrix.backend == 'cpu' && 'DEV=CPU\nCPU_COUNT=2' || matrix.backend == 'opencl' && 'DEV=CL' || matrix.backend == 'lvp' && 'DEV=CPU:LVP' }}" >> $GITHUB_ENV - - name: Check Device.DEFAULT and print some source - run: | - python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CPU','CL'], Device.DEFAULT" - DEBUG=5 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus - - name: Run pytest (${{ matrix.backend }}) - run: python -m pytest -n=auto test/backend --durations=20 - - name: Run TRANSCENDENTAL math - run: TRANSCENDENTAL=2 python -m pytest -n=auto test/backend/test_ops.py::TestOps::test_sin test/backend/test_ops.py::TestOps::test_cos test/backend/test_ops.py::TestOps::test_tan test/backend/test_ops.py::TestOps::test_exp test/backend/test_ops.py::TestOps::test_log --durations=20 - - name: Run process replay tests - uses: ./.github/actions/process-replay - # ****** OSX Tests ****** testmetal: @@ -868,84 +814,56 @@ jobs: - name: Run process replay tests uses: ./.github/actions/process-replay - osxwebgpu: - name: MacOS (WebGPU) - runs-on: macos-14 - timeout-minutes: 10 - steps: - - name: Checkout Code - uses: actions/checkout@v6 - - name: Setup Environment - uses: ./.github/actions/setup-tinygrad - with: - key: osx-webgpu - deps: testing - webgpu: 'true' - - name: Build WEBGPU Efficientnet - run: DEV=WEBGPU WEBGPU_BACKEND="WGPUBackendType_Metal" python3 -m examples.compile_efficientnet - - name: Run selected webgpu tests - run: DEV=WEBGPU WEBGPU_BACKEND="WGPUBackendType_Metal" python3 -m pytest -n=auto test/backend --durations=20 - #- name: Clean npm cache - # run: npm cache clean --force - #- name: Install Puppeteer - # run: npm install puppeteer - # this is also flaky - #- name: Run WEBGPU Efficientnet - # run: node test/web/test_webgpu.js - # this is flaky - #- name: Run VIZ tests as external package - # run: | - # mkdir $GITHUB_WORKSPACE/test_dir - # cd $GITHUB_WORKSPACE/test_dir - # python -m venv venv - # source venv/bin/activate - # pip install $GITHUB_WORKSPACE - # cp $GITHUB_WORKSPACE/test/web/test_viz.js . - # node test_viz.js - - name: Test ONNX Runner (WEBGPU) - run: DEV=WEBGPU python3 test/external/external_test_onnx_runner.py - - osxtests: + testmacos: strategy: fail-fast: false matrix: - backend: [metal, llvm, cpu, lvp] - name: MacOS (${{ matrix.backend }}) + dev: + - 'CPU:CLANG' + - 'CPU:LLVM' + - 'CPU:LVP' + - 'METAL' + - 'WEBGPU' + + name: MacOS (DEV=${{ matrix.dev }}) runs-on: macos-15 timeout-minutes: 20 steps: - - name: Checkout Code - uses: actions/checkout@v6 - - name: Setup Environment - uses: ./.github/actions/setup-tinygrad - with: - key: macos-${{ matrix.backend }}-minimal - deps: testing_unit - llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }} - mesa: ${{ matrix.backend == 'lvp' && 'true' }} - - name: Set env - run: printf "${{ matrix.backend == 'llvm' && 'DEV=CPU:LLVM' || matrix.backend == 'cpu' && 'DEV=CPU\nCPU_COUNT=2' || matrix.backend == 'metal' && 'DEV=METAL' || matrix.backend == 'lvp' && 'DEV=CPU:LVP' }}" >> $GITHUB_ENV - - name: Check Device.DEFAULT and print some source - run: | - python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU','LVP':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT" - DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus - - name: Run pytest (${{ matrix.backend }}) - run: python3 -m pytest -n=auto test/backend --durations=20 - - name: Run process replay tests - uses: ./.github/actions/process-replay - - name: Run macOS-specific unit test - if: matrix.backend == 'llvm' - run: python3 -m pytest test/unit/test_disk_tensor.py::TestDiskTensor::test_copy_to_cpu_not_truncated test/unit/test_cpu.py + - name: Checkout Code + uses: actions/checkout@v6 + - name: Setup Environment + uses: ./.github/actions/setup-tinygrad + with: + key: macos-${{ matrix.dev }} + deps: testing_unit + python-version: '3.12' + llvm: ${{ contains(matrix.dev, 'LLVM') || contains(matrix.dev, 'LVP') }} + mesa: ${{ contains(matrix.dev, 'LVP') && 'cpu' || 'false' }} + 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 backend tests + run: python -m pytest -n=auto test/backend --durations=20 + - name: Run process replay tests + uses: ./.github/actions/process-replay # ****** Windows Tests ****** - wintests: + testwindows: strategy: fail-fast: false matrix: - backend: [llvm, cpu, webgpu] + dev: + - 'CPU:CLANG' + - 'CPU:LLVM' + - 'CPU:X86' + - 'WEBGPU' - name: Windows (${{ matrix.backend }}) + name: Windows (DEV=${{ matrix.dev }}) runs-on: windows-latest timeout-minutes: 15 steps: @@ -954,25 +872,20 @@ jobs: - name: Setup Environment uses: ./.github/actions/setup-tinygrad with: - key: windows-${{ matrix.backend }}-minimal + key: windows-${{ matrix.dev }}-minimal deps: testing_unit - pydeps: ${{ matrix.backend == 'webgpu' && 'dawn-python' || '' }} + pydeps: ${{ matrix.dev == 'WEBGPU' && 'dawn-python' || '' }} - name: Set env shell: bash - run: printf "${{ matrix.backend == 'llvm' && 'DEV=CPU:LLVM' || matrix.backend == 'cpu' && 'DEV=CPU\nCPU_COUNT=2' || matrix.backend == 'webgpu' && 'DEV=WEBGPU'}}" >> $GITHUB_ENV - - name: Run unit tests - if: matrix.backend=='llvm' - # test_newton_schulz hits RecursionError - run: python -m pytest -n=auto test/unit/ --ignore=test/unit/test_disk_tensor.py --ignore=test/unit/test_tar.py --ignore=test/unit/test_linalg.py --durations=20 - - name: Run NULL backend tests - if: matrix.backend=='llvm' - shell: bash - run: DEV=NULL python -m pytest -n=auto test/null/ --ignore=test/null/test_elf.py --durations=20 - - name: Run pytest (${{ matrix.backend }}) + 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; assert Device.DEFAULT == {'LLVM':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT" - python -m pytest -n=auto test/test_tiny.py test/backend/test_ops.py --durations=20 + 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 ****** @@ -982,7 +895,7 @@ jobs: matrix: backend: [ir3, nak] name: Compile-only (${{ matrix.backend }}) - runs-on: ubuntu-24.04 + runs-on: *linux timeout-minutes: 15 steps: - name: Checkout Code @@ -1014,7 +927,7 @@ jobs: python -m pytest -n=auto test/backend/test_ops.py --durations=20 qcomclcompiletests: name: Compile-only (QCOM CL) - runs-on: ubuntu-24.04-arm + runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'MEMBER' && 'namespace-profile-tinygrad-arm64' || 'ubuntu-24.04-arm' }} timeout-minutes: 15 steps: - name: Checkout Code diff --git a/README.md b/README.md index c6d999aea88f6..4d378470e5b20 100644 --- a/README.md +++ b/README.md @@ -140,8 +140,8 @@ Documentation along with a quick start guide can be found on the [docs website]( ```python from tinygrad import Tensor -x = Tensor.eye(3, requires_grad=True) -y = Tensor([[2.0,0,-2.0]], requires_grad=True) +x = Tensor.eye(3) +y = Tensor([[2.0,0,-2.0]]) z = y.matmul(x).sum() z.backward() diff --git a/docs/quickstart.md b/docs/quickstart.md index 0153cd5126f82..57546abce72d4 100644 --- a/docs/quickstart.md +++ b/docs/quickstart.md @@ -133,7 +133,7 @@ For our loss function we will be using sparse categorical cross entropy loss. Th ```python def sparse_categorical_crossentropy(self, Y, ignore_index=-1) -> Tensor: loss_mask = Y != ignore_index - y_counter = Tensor.arange(self.shape[-1], dtype=dtypes.int32, requires_grad=False, device=self.device).unsqueeze(0).expand(Y.numel(), self.shape[-1]) + y_counter = Tensor.arange(self.shape[-1], dtype=dtypes.int32, device=self.device).unsqueeze(0).expand(Y.numel(), self.shape[-1]) y = ((y_counter == Y.flatten().reshape(-1, 1)).where(-1.0, 0) * loss_mask.reshape(-1, 1)).reshape(*Y.shape, self.shape[-1]) return self.log_softmax().mul(y).sum() / loss_mask.sum() ``` @@ -175,7 +175,7 @@ with Tensor.train(): for step in range(1000): # random sample a batch samp = np.random.randint(0, X_train.shape[0], size=(64)) - batch = Tensor(X_train[samp], requires_grad=False) + batch = Tensor(X_train[samp]) # get the corresponding labels labels = Tensor(Y_train[samp]) @@ -213,7 +213,7 @@ with Timing("Time: "): for step in range(1000): # random sample a batch samp = np.random.randint(0, X_test.shape[0], size=(64)) - batch = Tensor(X_test[samp], requires_grad=False) + batch = Tensor(X_test[samp]) # get the corresponding labels labels = Y_test[samp] @@ -257,7 +257,7 @@ with Timing("Time: "): for step in range(1000): # random sample a batch samp = np.random.randint(0, X_test.shape[0], size=(64)) - batch = Tensor(X_test[samp], requires_grad=False) + batch = Tensor(X_test[samp]) # get the corresponding labels labels = Y_test[samp] diff --git a/examples/anthropic_challenge.py b/examples/anthropic_challenge.py index 0dafd21992412..a9adea334cf52 100644 --- a/examples/anthropic_challenge.py +++ b/examples/anthropic_challenge.py @@ -174,7 +174,7 @@ def render(self, uops:list[UOp]): # *** render to device *** from tinygrad.codegen import to_program - with Context(PCONTIG=2, DEVECTORIZE=2, SPEC=0): + with Context(PCONTIG=2, SPEC=0): out = tree_traversal(forest_t, val_t, height, rounds) sink = out.schedule_linear().src[-1].src[0] prg = to_program(sink, VLIWRenderer()) diff --git a/examples/beautiful_cifar.py b/examples/beautiful_cifar.py index 5bc2fc87c3799..85f737e058dbb 100644 --- a/examples/beautiful_cifar.py +++ b/examples/beautiful_cifar.py @@ -67,8 +67,8 @@ def __init__(self, channels_in, channels_out): self.conv2 = nn.Conv2d(channels_out, channels_out, kernel_size=3, padding=1, bias=False) self.norm1 = nn.BatchNorm(channels_out, track_running_stats=False, eps=1e-12, momentum=hyp['net']['batch_norm_momentum']) self.norm2 = nn.BatchNorm(channels_out, track_running_stats=False, eps=1e-12, momentum=hyp['net']['batch_norm_momentum']) - cast(Tensor, self.norm1.weight).requires_grad = False - cast(Tensor, self.norm2.weight).requires_grad = False + cast(Tensor, self.norm1.weight).is_param_(False) + cast(Tensor, self.norm2.weight).is_param_(False) def __call__(self, x:Tensor) -> Tensor: x = self.norm1(self.conv1(x).max_pool2d().float()).cast(dtypes.default_float).quick_gelu() return self.norm2(self.conv2(x).float()).cast(dtypes.default_float).quick_gelu() + x diff --git a/examples/gradaccum_mnist.py b/examples/gradaccum_mnist.py index 2a0ac6f14311c..daf38956786e4 100644 --- a/examples/gradaccum_mnist.py +++ b/examples/gradaccum_mnist.py @@ -41,15 +41,15 @@ def functional_adam(g:Tensor, m:Tensor, v:Tensor, b1_t:Tensor, b2_t:Tensor, lr=0 Tensor.realize(*params) # split params (with grads) and buffers (without) - params, buffers = partition(params, lambda x: x.requires_grad) + params, buffers = partition(params, lambda x: x.is_param) print(f"params: {len(params)} buffers: {len(buffers)}") # optim params pos_params = list(itertools.accumulate(params, lambda x,y: x+y.numel(), initial=0)) adam_m = Tensor.zeros(pos_params[-1], device="CPU").contiguous() adam_v = Tensor.zeros(pos_params[-1], device="CPU").contiguous() - adam_b1_t = Tensor.ones((1,), dtype=dtypes.float32, device="CPU", requires_grad=False).contiguous() - adam_b2_t = Tensor.ones((1,), dtype=dtypes.float32, device="CPU", requires_grad=False).contiguous() + adam_b1_t = Tensor.ones((1,), dtype=dtypes.float32, device="CPU").contiguous() + adam_b2_t = Tensor.ones((1,), dtype=dtypes.float32, device="CPU").contiguous() adam_params = [adam_m, adam_v, adam_b1_t, adam_b2_t] # create loss and grads. init all state so the JIT works on microbatch diff --git a/examples/hlb_cifar10.py b/examples/hlb_cifar10.py index 6fc6d08fd1d2e..d08d596d6830b 100644 --- a/examples/hlb_cifar10.py +++ b/examples/hlb_cifar10.py @@ -30,9 +30,9 @@ def __init__(self, sz:int, eps=1e-5, affine=True, track_running_stats=True, mome if affine: self.weight, self.bias = Tensor.ones(sz, dtype=dtypes.float32), Tensor.zeros(sz, dtype=dtypes.float32) else: self.weight, self.bias = None, None - self.running_mean = Tensor.zeros(num_devices, sz, dtype=dtypes.float32, requires_grad=False) - self.running_var = Tensor.ones(num_devices, sz, dtype=dtypes.float32, requires_grad=False) - self.num_batches_tracked = Tensor.zeros(1, dtype=dtypes.int, requires_grad=False) + self.running_mean = Tensor.zeros(num_devices, sz, dtype=dtypes.float32).is_param_(False) + self.running_var = Tensor.ones(num_devices, sz, dtype=dtypes.float32).is_param_(False) + self.num_batches_tracked = Tensor.zeros(1, dtype=dtypes.int).is_param_(False) def __call__(self, x:Tensor): xr = x.reshape(self.num_devices, -1, *x.shape[1:]).cast(dtypes.float32) @@ -68,8 +68,7 @@ def calc_stats(self, x:Tensor): class BatchNorm(nn.BatchNorm2d if getenv("SYNCBN") else UnsyncedBatchNorm): def __init__(self, num_features): super().__init__(num_features, track_running_stats=False, eps=1e-12, momentum=0.85, affine=True) - self.weight.requires_grad = False - self.bias.requires_grad = True + self.weight.is_param_(False) class ConvGroup: def __init__(self, channels_in, channels_out): @@ -172,7 +171,7 @@ def _eigens(patches): Λ, V = _eigens(_patches(X.float().numpy())) W = V/np.sqrt(Λ+1e-2)[:,None,None,None] - return Tensor(W.astype(np.float32), requires_grad=False).cast(dtypes.default_float) + return Tensor(W.astype(np.float32)).cast(dtypes.default_float).is_param_(False) # ========== Loss ========== def cross_entropy(x:Tensor, y:Tensor, reduction:str='mean', label_smoothing:float=0.0) -> Tensor: @@ -264,7 +263,6 @@ def __init__(self, w, net): # self.model_ema = copy.deepcopy(net) # won't work for opencl due to unpickeable pyopencl._cl.Buffer self.net_ema = SpeedyResNet(w) for net_ema_param, net_param in zip(get_state_dict(self.net_ema).values(), get_state_dict(net).values()): - net_ema_param.requires_grad = False net_ema_param.assign(net_param.numpy()) @TinyJit @@ -307,7 +305,7 @@ def update(self, net, decay): params_bias = [] params_non_bias = [] for params in params_dict: - if params_dict[params].requires_grad is not False: + if params_dict[params].is_param: if 'bias' in params: params_bias.append(params_dict[params]) else: diff --git a/examples/llama3.py b/examples/llama3.py index c7476c50b80dd..f55fdb3273609 100644 --- a/examples/llama3.py +++ b/examples/llama3.py @@ -102,7 +102,7 @@ def __init__(self, vocab_size:int, embed_size:int): self.weight, self.scale = Tensor.ones(vocab_size, embed_size, dtype=dtypes.int8), Tensor.ones(vocab_size, dtype=dtypes.half) def __call__(self, idx:Tensor) -> Tensor: - if not hasattr(self, 'arange'): self.arange = Tensor.arange(self.vocab_sz, requires_grad=False, device=self.weight.device).unsqueeze(-1) + if not hasattr(self, 'arange'): self.arange = Tensor.arange(self.vocab_sz, device=self.weight.device).unsqueeze(-1) big_shp = idx.shape+(self.vocab_sz, self.embed_sz) arange, idx, vals = self.arange.expand(big_shp), idx.reshape(idx.shape+(1, 1)).expand(big_shp), (self.weight.cast(self.scale.dtype).T*self.scale).T return (arange == idx).mul(vals).sum(-2, dtype=vals.dtype) diff --git a/examples/llm.c/train_gpt2.py b/examples/llm.c/train_gpt2.py index e259cf446608b..e0ba635a94876 100755 --- a/examples/llm.c/train_gpt2.py +++ b/examples/llm.c/train_gpt2.py @@ -25,7 +25,7 @@ def __init__(self, config:GPTConfig): self.n_embd = config.n_embd # not really a 'bias', more of a mask, but following the OpenAI/HF naming though self.bias = Tensor.ones(1, 1, config.block_size, config.block_size).tril() - self.bias.requires_grad = False + self.bias.is_param_(False) def __call__(self, x:Tensor): B, T, C = x.shape diff --git a/examples/mixtral.py b/examples/mixtral.py index c621d409e6f4d..62f400f9a6c62 100644 --- a/examples/mixtral.py +++ b/examples/mixtral.py @@ -1,6 +1,6 @@ import functools, argparse, pathlib from tinygrad import Tensor, nn, Device, GlobalCounters, Variable -from tinygrad.helpers import Timing, Profiling, CI, tqdm +from tinygrad.helpers import Timing, Profiling, tqdm from tinygrad.nn.state import torch_load, get_state_dict from extra.models.llama import FeedForward, Transformer from extra.bench_log import BenchEvent, WallTimeEvent @@ -36,7 +36,7 @@ def __call__(self, x:Tensor) -> Tensor: model = Transformer(n_layers=32, dim=4096, hidden_dim=14336, n_heads=32, n_kv_heads=8, norm_eps=1e-5, vocab_size=32000, feed_forward=functools.partial(MixtureFeedForward, 8), jit=False) model_state_dict = get_state_dict(model) - for k in (t := tqdm(state, disable=CI)): + for k in (t := tqdm(state, disable=None)): if 'feed_forward.experts.' in k: expert_no = int(k.split('feed_forward.experts.')[1].split('.')[0]) device = Device.DEFAULT + ":" + str((expert_no//2)+1) @@ -44,7 +44,7 @@ def __call__(self, x:Tensor) -> Tensor: device = Device.DEFAULT t.set_description(f"ram used: {GlobalCounters.mem_used/1e9:5.2f} GB, loading {k} to {device}") model_state_dict[k].replace(state[k].to(device).half()).realize() - if CI: print(f"ram used: {GlobalCounters.mem_used/1e9:5.2f} GB") + if t.disable: print(f"ram used: {GlobalCounters.mem_used/1e9:5.2f} GB") from sentencepiece import SentencePieceProcessor spp = SentencePieceProcessor(model_file=args.weights + "/tokenizer.model") diff --git a/examples/mlperf/initializers.py b/examples/mlperf/initializers.py index d10792d9178c8..338d283d89a62 100644 --- a/examples/mlperf/initializers.py +++ b/examples/mlperf/initializers.py @@ -57,7 +57,7 @@ def __init__(self, vocab_size:int, embed_size:int, std=0.02): def __call__(self, idx:Tensor) -> Tensor: if idx.numel() == 0: return Tensor.empty(idx.shape+(self.embed_sz,), dtype=self.weight.dtype, device=self.weight.device) arange_shp, weight_shp, big_shp = (1, 1, self.vocab_sz, 1), (1, 1, self.vocab_sz, self.embed_sz), idx.shape+(self.vocab_sz, self.embed_sz,) - if not hasattr(self, 'arange'): self.arange = Tensor.arange(self.vocab_sz, requires_grad=False, device=self.weight.device).reshape(arange_shp) + if not hasattr(self, 'arange'): self.arange = Tensor.arange(self.vocab_sz, device=self.weight.device).reshape(arange_shp) arange, idx, vals = self.arange.expand(big_shp), idx.reshape(idx.shape+(1, 1,)).expand(big_shp), self.weight.cast(dtypes.default_float).reshape(weight_shp).expand(big_shp) return (arange == idx).where(vals, 0).sum(2, dtype=vals.dtype) @@ -77,11 +77,11 @@ class FrozenBatchNorm2dRetinaNet(nn.BatchNorm2d): def __init__(self, sz:int, eps=1e-5, affine=True, track_running_stats=True, momentum=0.1): self.eps, self.track_running_stats, self.momentum = eps, track_running_stats, momentum - self.weight = Tensor.ones(sz, dtype=dtypes.float32, requires_grad=False) if affine else None - self.bias = Tensor.zeros(sz, dtype=dtypes.float32, requires_grad=False) if affine else None + self.weight = Tensor.ones(sz, dtype=dtypes.float32).is_param_(False) if affine else None + self.bias = Tensor.zeros(sz, dtype=dtypes.float32).is_param_(False) if affine else None - if track_running_stats: self.running_mean, self.running_var = Tensor.zeros(sz, dtype=dtypes.float32, requires_grad=False), Tensor.ones(sz, dtype=dtypes.float32, requires_grad=False) - self.num_batches_tracked = Tensor.zeros(1, dtype=dtypes.long, requires_grad=False) + if track_running_stats: self.running_mean, self.running_var = Tensor.zeros(sz, dtype=dtypes.float32).is_param_(False), Tensor.ones(sz, dtype=dtypes.float32).is_param_(False) + self.num_batches_tracked = Tensor.zeros(1, dtype=dtypes.long).is_param_(False) def __call__(self, x:Tensor) -> Tensor: batch_mean, batch_var = super().calc_stats(x.cast(dtypes.float32)) diff --git a/examples/mlperf/model_train.py b/examples/mlperf/model_train.py index e9442791552e6..d2dab878145ca 100644 --- a/examples/mlperf/model_train.py +++ b/examples/mlperf/model_train.py @@ -180,11 +180,11 @@ def eval_step(X, Y): def fake_data_get(batch_size): x = Tensor.zeros(batch_size, 224, 224, 3, dtype=dtypes.uchar).contiguous() y = [0] * batch_size - return x.shard(GPUS, axis=0).realize(), Tensor(y, requires_grad=False).shard(GPUS, axis=0), y, None + return x.shard(GPUS, axis=0).realize(), Tensor(y).shard(GPUS, axis=0), y, None def data_get(it): x, y, cookie = next(it) - return x.shard(GPUS, axis=0).realize(), Tensor(y, requires_grad=False).shard(GPUS, axis=0), y, cookie + return x.shard(GPUS, axis=0).realize(), Tensor(y).shard(GPUS, axis=0), y, cookie # ** epoch loop ** step_times = [] @@ -413,7 +413,7 @@ def _freeze_backbone_layers(backbone:resnet.ResNet, trainable_layers:int): layers_to_train = ["layer4", "layer3", "layer2", "layer1", "conv1"][:trainable_layers] for k, v in get_state_dict(backbone).items(): if all([not k.startswith(layer) for layer in layers_to_train]): - v.requires_grad = False + v.is_param_(False) def _data_get(it:Iterator[tuple[Tensor, ...]], val:bool=False): if val: @@ -798,7 +798,7 @@ def train_step(model, x, y): @Tensor.train(mode=False) def eval_step(model, x, y): y_hat, y = sliding_window_inference(model, x, y, gpus=GPUS) - y_hat, y = Tensor(y_hat), Tensor(y, requires_grad=False) + y_hat, y = Tensor(y_hat), Tensor(y) loss = dice_ce_loss(y_hat, y) score = dice_score(y_hat, y) return loss.realize(), score.realize() @@ -1419,10 +1419,7 @@ def train_llama3(): for p in optim.params: grad_dtype = dtypes.bfloat16 if p.dtype == FP8_DTYPE else p.dtype - if isinstance(p.device, tuple) and p.uop.axis is not None: - p.grad = Tensor.zeros(p.shape, dtype=grad_dtype, device=p.device[0]).shard_(p.device, axis=p.uop.axis).contiguous() - else: - p.grad = Tensor.zeros(p.shape, dtype=grad_dtype, device=p.device).contiguous() + p.grad = p.zeros_like(dtype=grad_dtype).contiguous() grads = [p.grad for p in optim.params] scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps) @@ -1438,16 +1435,19 @@ def train_llama3(): fp8_amax = [t for ts in model._fp8_amax.values() for t in ts] fp8_grad_amax = [t for ts in model._fp8_grad_amax.values() for t in ts] if hasattr(model, "_fp8_grad_amax") else [] - fp8_inv_scales = list(model._fp8_inv_scale.values()) + fp8_inv_scales = list(model._fp8_inv_scale.values()) + list(model._fp8_next_inv_scale.values()) from tinygrad.nn.state import get_state_dict model_state = get_state_dict(model) for wname in model._fp8_inv_scale: w = model_state[wname] w._inv_scale = model._fp8_inv_scale[wname] + w._next_inv_scale = model._fp8_next_inv_scale[wname] if optim.master_params: idx = next(j for j, p in enumerate(optim.params) if p is w) - optim.master_params[idx].assign((optim.master_params[idx] * w._inv_scale.reshape(-1, *([1]*(w.ndim-1)))).contiguous()) + master = optim.master_params[idx] + inv = w._inv_scale if w._inv_scale.device == master.device else w._inv_scale.to(master.device) + master.assign((master * inv.reshape(-1, *([1]*(w.ndim-1)))).contiguous()) # realize everything here if optim.master_params: Tensor.realize(*optim.master_params) @@ -1476,7 +1476,7 @@ def optim_step(): grad_norm = optim.fstep(grads) scheduler.step() - for g in grads: g.assign(g.zeros_like()) + for g in grads: g.assign(g.const_like(0)) lr_cpu = optim.lr.float().to("CPU") grad_norm_cpu = grad_norm.float().to("CPU") @@ -1498,7 +1498,7 @@ def eval_step(tokens:Tensor): def fake_data(bs, samples): import numpy as np for _ in range(samples // bs): - fake_data_np = np.random.randint(0, model_params["vocab_size"], size=(bs, SEQLEN + 1), dtype=np.int32) + fake_data_np = np.random.randint(0, real_vocab_size, size=(bs, SEQLEN + 1), dtype=np.int32) yield Tensor(fake_data_np, device="NPY") def get_train_iter(): diff --git a/examples/mlperf/models/flat_llama.py b/examples/mlperf/models/flat_llama.py index 018d61f7aad28..12240d068ca12 100644 --- a/examples/mlperf/models/flat_llama.py +++ b/examples/mlperf/models/flat_llama.py @@ -2,9 +2,8 @@ if __name__ == "__main__": os.environ["DEFAULT_FLOAT"] = "bfloat16" os.environ["OPTIM_DTYPE"] = "bfloat16" - if "DEV" not in os.environ: os.environ["DEV"] = "NULL" + if "DEV" not in os.environ: os.environ["DEV"] = "NULL::gfx950" # CDNA - os.environ["EMULATE"] = "AMD_CDNA4" os.environ["DEVICE_IN_FUNCTION_BUG"] = "1" os.environ["ALL2ALL"] = "1" os.environ["USE_ATOMICS"] = "1" @@ -13,7 +12,7 @@ if "ASM_GEMM" not in os.environ: os.environ["ASM_GEMM"] = "1" from tinygrad import Tensor, nn, function, getenv, dtypes, TinyJit -from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker +from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker, round_up from tinygrad.uop.ops import Ops, UOp from extra.models.llama import apply_rotary_emb, precompute_freqs_cis from extra.llama_kernels.rmsnorm import rmsnorm @@ -53,8 +52,8 @@ def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_sca if ASM_GEMM: from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm if can_use_asm_gemm(x_fp8, w.T): - return asm_gemm(x_fp8, w.T, x_scale=x_scale, w_scale=w_inv_scale, grad_amax_state=grad_amax_state), x_new_amax, x_fp8, w - return (x_fp8.dot(w.T, dtype=dtypes.float) * x_scale * w_inv_scale).cast(dtypes.bfloat16), x_new_amax, x_fp8, w + return asm_gemm(x_fp8, w.T, x_scale=x_scale, w_scale=w_inv_scale, grad_amax_state=grad_amax_state), x_new_amax, x_fp8 + return (x_fp8.dot(w.T, dtype=dtypes.float) * x_scale * w_inv_scale).cast(dtypes.bfloat16), x_new_amax, x_fp8 def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor, grad_amax_state:Tensor): if FUSED_ADD_NORM_MUL_QUANTIZE: @@ -125,9 +124,9 @@ def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_layers:int, norm_eps: self.tok_embeddings = nn.Embedding(vocab_size, dim) self.tok_embeddings.weight = Tensor.normal(vocab_size, dim, mean=0.0, std=0.02, dtype=dtypes.bfloat16) self.output = Tensor.normal(1, vocab_size, dim, mean=0.0, std=0.02, dtype=dtypes.bfloat16) - self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).contiguous().requires_grad_(False) + self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).contiguous().is_param_(False) - def _amax(): return Tensor.full((), FP8_MAX, dtype=dtypes.float32).contiguous().requires_grad_(False) + def _amax(): return Tensor.full((), FP8_MAX, dtype=dtypes.float32).contiguous().is_param_(False) names = ["xqkv", "xo", "x2"] names += ["x1", "x3"] if SPLIT_W13 else ["x13"] self._fp8_amax = {name: [_amax() for _ in range(n_layers)] for name in names} @@ -136,7 +135,8 @@ def _amax(): return Tensor.full((), FP8_MAX, dtype=dtypes.float32).contiguous(). self._fp8_grad_amax = {name: [_amax() for _ in range(n_layers)] for name in grad_names} w_scales = [("wqkv", s_qkv), ("wo", s_o), ("w2", s_2)] w_scales += [("w1", s_1), ("w3", s_3)] if SPLIT_W13 else [("w13", s_13)] - self._fp8_inv_scale = {name: s.float().contiguous().requires_grad_(False) for name, s in w_scales} + self._fp8_inv_scale = {name: s.float().contiguous().is_param_(False) for name, s in w_scales} + self._fp8_next_inv_scale = {name: s.float().contiguous().is_param_(False) for name, s in w_scales} def lin_per_layer(self, in_features:int, out_features:int, std:float=0.02): if getenv("ZEROS"): w = Tensor.zeros(self.n_layers, out_features, in_features) @@ -212,8 +212,9 @@ def run_layer(self, x:Tensor, freqs_cis:Tensor, attn_kwargs:dict, ffn_kwargs:dic attn, attn_amaxs, attn_saves = self.attention(x, freqs_cis, **attn_kwargs) ffn, h, ffn_amaxs, ffn_saves = self.feed_forward(x, attn, **ffn_kwargs) h = h + ffn - if save: return (h, *attn_amaxs, *ffn_amaxs, *attn_saves, *ffn_saves) - else: return (h, *attn_amaxs, *ffn_amaxs) + amaxs = tuple(a.detach() for a in (*attn_amaxs, *ffn_amaxs)) + if save: return (h, *amaxs, *attn_saves, *ffn_saves) + else: return (h, *amaxs) def shard(self, device:tuple[str, ...], mp:bool=False): from tinygrad.nn.state import get_parameters @@ -221,14 +222,19 @@ def shard(self, device:tuple[str, ...], mp:bool=False): for v in get_parameters(self): v.shard_(device, axis=None) else: # flat per-layer weights: axis 0 is n_layers, so shard axes are +1 vs per-layer Transformer - self.wqkv.shard_(device, axis=1).realize() # (n_layers, out, dim) shard out - self.wo.shard_(device, axis=2).realize() # (n_layers, dim, in) shard in + def _shard_fp8(name:str, axis:int): + getattr(self, name).shard_(device, axis=axis) + self._fp8_inv_scale[name] = self._fp8_inv_scale[name].to(device).contiguous().is_param_(False) + self._fp8_next_inv_scale[name] = self._fp8_next_inv_scale[name].to(device).contiguous().is_param_(False) + Tensor.realize(getattr(self, name), self._fp8_inv_scale[name], self._fp8_next_inv_scale[name]) + _shard_fp8("wqkv", 1) # (n_layers, out, dim) shard out + _shard_fp8("wo", 2) # (n_layers, dim, in) shard in if SPLIT_W13: - self.w1.shard_(device, axis=1).realize() - self.w3.shard_(device, axis=1).realize() + _shard_fp8("w1", 1) + _shard_fp8("w3", 1) else: - self.w13.shard_(device, axis=1).realize() # (n_layers, hidden*2, dim) shard out - self.w2.shard_(device, axis=2).realize() # (n_layers, dim, hidden) shard in + _shard_fp8("w13", 1) # (n_layers, hidden*2, dim) shard out + _shard_fp8("w2", 2) # (n_layers, dim, hidden) shard in self.attention_norm.shard_(device, axis=None).realize() self.ffn_norm.shard_(device, axis=None).realize() self.norm.weight.shard_(device, axis=None).realize() @@ -238,9 +244,7 @@ def shard(self, device:tuple[str, ...], mp:bool=False): for amax_dict in (self._fp8_amax, self._fp8_grad_amax): for name in amax_dict: for i in range(len(amax_dict[name])): - amax_dict[name][i] = amax_dict[name][i].to(device).contiguous().requires_grad_(False) - for name in self._fp8_inv_scale: - self._fp8_inv_scale[name] = self._fp8_inv_scale[name].to(device).contiguous().requires_grad_(False) + amax_dict[name][i] = amax_dict[name][i].to(device).contiguous().is_param_(False) def __call__(self, tokens:Tensor, save:bool=True): h = self.tok_embeddings(tokens) @@ -273,41 +277,59 @@ def apply_grad(grad_buf:Tensor, new_grad:UOp): pads = _get_pads(new_grad) if len(pads) <= 1: new_grad = new_grad.cast(grad_buf.dtype) - store = grad_buf.uop.store(grad_buf.uop + new_grad) - grad_buf.uop = grad_buf.uop.after(store) + grad_buf.uop = grad_buf.uop.after(grad_buf.uop.store(grad_buf.uop + new_grad)) return - sorted_pads = sorted(pads, key=lambda p: p.marg[0][0] if p.op == Ops.PAD else 0) - inners_raw = [Tensor(p.src[0] if p.op == Ops.PAD else p, device=grad_buf.device) for p in sorted_pads] - if getenv("FUSED_PAD_GRAD_ACCUM", 0): - from extra.llama_kernels.fused_pad_grad_accum import fused_pad_grad_accum, can_fused_pad_grad_accum - if can_fused_pad_grad_accum(grad_buf, inners_raw): - grad_buf.uop = fused_pad_grad_accum(grad_buf, inners_raw).uop - return - inners = [t.cast(grad_buf.dtype) for t in inners_raw] - grad_buf.assign(grad_buf + inners[0].cat(*inners[1:], dim=0)) + cur = grad_buf.uop + for pad in sorted(pads, key=lambda p: p.marg[0][0] if p.op == Ops.PAD else 0, reverse=True): + if pad.op == Ops.PAD: + grad_shrink = tuple([(p[0], s+p[0]) for s,p in zip(pad.src[0].shape, pad.marg)]) + buf_slice = cur.shrink(grad_shrink) + cur = cur.after(buf_slice.store(buf_slice + pad.src[0].cast(cur.dtype))) + else: + cur = cur.after(cur.store(cur + pad.cast(cur.dtype))) + grad_buf.uop = cur if __name__ == "__main__": config = {} BS = config["BS"] = getenv("BS", 16) SEQLEN = config["SEQLEN"] = getenv("SEQLEN", 8192) + SMALL = config["SMALL"] = getenv("SMALL", 0) from examples.llama3 import MODEL_PARAMS - model_params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"] - if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params['n_layers'] = llama_layers + model_params = MODEL_PARAMS[llama_size:=getenv("LLAMA3_SIZE", "8B")]["args"] + # vocab_size from mixtral tokenizer + if not SMALL: model_params |= {"vocab_size": 32000} + real_vocab_size = model_params['vocab_size'] + if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params["n_layers"] = llama_layers + + # pad vocab + if (MP := getenv("MP", 1)) > 1: model_params["vocab_size"] = round_up(model_params["vocab_size"], 256 * MP) + vocab_mask:Tensor = Tensor.arange(model_params["vocab_size"]).reshape(1, 1, -1) >= real_vocab_size + model = FlatTransformer(**model_params, max_context=SEQLEN) + state = nn.state.get_state_dict(model) print("tensor count:", len(state)) # shard the model from tinygrad import Device - if (DP := getenv("DP", 1)) > 1: - model.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))) - if (MP := getenv("MP", 1)) > 1: - model.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(MP)), mp=True) + is_dp = (DP := getenv("DP", 1)) > 1 + is_mp = (MP := getenv("MP", 1)) > 1 + is_sharding = is_dp or is_mp + device_count = max(DP, MP) + device = tuple(f"{Device.DEFAULT}:{i}" for i in range(device_count)) + + model.shard(device, is_mp) + + if is_dp: vocab_mask.shard_(device, axis=None).realize() + if is_mp: vocab_mask.shard_(device, axis=2).realize() # preallocate all the grad buffers and zero them out - grads = {x:Tensor.zeros(x.shape, dtype=x.dtype, device=x.device).contiguous() - for x in state.values() if x.requires_grad} + grad_dtype = lambda x: dtypes.bfloat16 if x.dtype in dtypes.fp8s else x.dtype + grads = {x:x.zeros_like(dtype=grad_dtype(x)).contiguous() for x in state.values() if x.is_param} + + fp8_amax = [t for ts in model._fp8_amax.values() for t in ts] + fp8_grad_amax = [t for ts in model._fp8_grad_amax.values() for t in ts] # print model size sz = 0 @@ -316,23 +338,31 @@ def apply_grad(grad_buf:Tensor, new_grad:UOp): sz += v.nbytes() print(f"total sz: {sz/1e9:.2f} GB") - with Timing("fake data: "): tokens = Tensor.randint(BS, SEQLEN+1, low=0, high=model.vocab_size, dtype=dtypes.int) + with Timing("fake data: "): tokens = Tensor.randint(BS, SEQLEN+1, low=0, high=real_vocab_size, dtype=dtypes.int) with Timing("realize weights/grads/data: "): Tensor.realize(*state.values(), *grads.values(), tokens) print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items()))) if DP > 1: tokens = tokens.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(DP)), axis=0) if MP > 1: tokens = tokens.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))) @TinyJit - def jit_step(tokens:Tensor): - with Timing("python forward: "): loss = model(tokens[:, :-1]).sparse_categorical_crossentropy(tokens[:, 1:]) + def fwd_bwd(tokens:Tensor): + with Timing("python forward: "): + logits = model(tokens[:, :-1], save=llama_size=="8B") + loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:]) with Timing("python backward: "): for t,g in zip(grads, loss.gradient(*grads)): apply_grad(grads[t], g.uop) - with Timing("run step: "): loss.realize(*grads.values()) + with Timing("run fwd_bwd: "): loss.realize(*grads.values(), *fp8_amax, *fp8_grad_amax) + + @TinyJit + def optim_step(): + for g in grads.values(): g.assign(g.zeros_like()) + Tensor.realize(*grads.values()) for i in range(6): GlobalCounters.reset() profile_marker(f"step {i}") with Timing(colored(f"*** step {i}: ", "red")): - jit_step(tokens) + fwd_bwd(tokens) + optim_step() print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items()))) diff --git a/examples/mlperf/models/test_apply_grad.py b/examples/mlperf/models/test_apply_grad.py new file mode 100644 index 0000000000000..6af56bd48d31b --- /dev/null +++ b/examples/mlperf/models/test_apply_grad.py @@ -0,0 +1,68 @@ +import unittest +from tinygrad import Tensor, TinyJit +from tinygrad.nn.state import get_parameters +from examples.mlperf.models.flat_llama import apply_grad + +class FlatModel: + def __init__(self, n_layers:int, dim:int, hidden:int): + self.n_layers = n_layers + self.w1 = Tensor.uniform(n_layers, dim, hidden, low=-0.1, high=0.1) + self.w2 = Tensor.uniform(n_layers, hidden, dim, low=-0.1, high=0.1) + self.scale = Tensor.uniform(dim, low=0.9, high=1.1) + self.bias = Tensor.zeros(dim).contiguous() + + def __call__(self, x:Tensor) -> Tensor: + h = x + for i in range(self.n_layers): + h = (h @ self.w1[i]).relu() @ self.w2[i] + h + return (h * self.scale + self.bias).sum() + +class TestApplyGradE2E(unittest.TestCase): + def _run_with_apply_grad(self, model, xs): + grads = {p: Tensor.zeros(p.shape, dtype=p.dtype).contiguous().realize() for p in get_parameters(model)} + for x in xs: + loss = model(x) + for p, g in zip(grads, loss.gradient(*grads)): + apply_grad(grads[p], g.uop) + Tensor.realize(loss, *grads.values()) + return [grads[p] for p in get_parameters(model)] + + def _run_reference(self, model, xs): + for x in xs: model(x).backward() + return [p.grad for p in get_parameters(model)] + + def _assert_close(self, got, expected, atol, rtol): + for g, e in zip(got, expected): + self.assertTrue(g.allclose(e, atol=atol, rtol=rtol).item(), f"grad mismatch (max abs diff {(g - e).abs().max().item()})") + + def _assert_match(self, model, xs, atol, rtol): + self._assert_close(self._run_with_apply_grad(model, xs), self._run_reference(model, xs), atol, rtol) + + def test_e2e_single_step(self): + model = FlatModel(n_layers=3, dim=8, hidden=16) + Tensor.realize(*get_parameters(model)) + self._assert_match(model, [Tensor.randn(2, 8).realize()], atol=1e-4, rtol=1e-4) + + def test_e2e_multi_step_accumulation(self): + model = FlatModel(n_layers=4, dim=8, hidden=16) + Tensor.realize(*get_parameters(model)) + self._assert_match(model, [Tensor.randn(2, 8).realize() for _ in range(3)], atol=1e-4, rtol=1e-4) + + def test_e2e_jit(self): + model = FlatModel(n_layers=3, dim=8, hidden=16) + Tensor.realize(*get_parameters(model)) + grads = {p: Tensor.zeros(p.shape, dtype=p.dtype).contiguous().realize() for p in get_parameters(model)} + + @TinyJit + def fwd_bwd(x:Tensor): + loss = model(x) + for p, g in zip(grads, loss.gradient(*grads)): apply_grad(grads[p], g.uop) + Tensor.realize(loss, *grads.values()) + + xs = [Tensor.randn(2, 8).realize() for _ in range(3)] + for x in xs: fwd_bwd(x) + self._assert_close([grads[p] for p in get_parameters(model)], self._run_reference(model, xs), atol=1e-3, rtol=1e-3) + + +if __name__ == "__main__": + unittest.main() diff --git a/examples/mlperf/models/test_flat_llama.py b/examples/mlperf/models/test_flat_llama.py index 98542ddad2a5b..600288efe2c75 100644 --- a/examples/mlperf/models/test_flat_llama.py +++ b/examples/mlperf/models/test_flat_llama.py @@ -3,8 +3,7 @@ import unittest import numpy as np from tinygrad import Tensor, nn, dtypes -from tinygrad.nn.state import get_parameters -from tinygrad.device import is_dtype_supported, Device +from tinygrad.device import Device from examples.mlperf.models.llama import Transformer from examples.mlperf.models.flat_llama import FlatTransformer @@ -45,8 +44,6 @@ def test_backward_match(self): flat = FlatTransformer(**params) copy_weights(flat, ref) - for p in get_parameters(ref): p.requires_grad_(True) - for p in get_parameters(flat): p.requires_grad_(True) Tensor.realize(*nn.state.get_state_dict(flat).values()) tokens = Tensor([[1, 50, 100, 999, 2, 10]]) @@ -114,7 +111,7 @@ def test_forward_match_dp(self): self.assertEqual(ref_logits.shape, flat_logits.shape) np.testing.assert_allclose(flat_logits, ref_logits, atol=1e-4, rtol=1e-4) - @unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3), "fp8 not supported on this device") + @unittest.skipUnless(dtypes.fp8e4m3 in Device[Device.DEFAULT].renderer.supported_dtypes(), "fp8 not supported on this device") def test_forward_fp8(self): import examples.mlperf.models.flat_llama as flat_llama_mod old_fp8 = flat_llama_mod.FP8 diff --git a/examples/mlperf/optim.py b/examples/mlperf/optim.py index 1148a331d8159..d2fe057167fe8 100644 --- a/examples/mlperf/optim.py +++ b/examples/mlperf/optim.py @@ -6,6 +6,7 @@ STOCHASTIC_ROUND = getenv("STOCHASTIC_ROUND", 0) MASTER_WEIGHTS = getenv("MASTER_WEIGHTS", 0) +FP8_AMAX_MARGIN = getenv("FP8_AMAX_MARGIN", 1.1) def stochastic_round_bf16(x:Tensor) -> Tensor: bits = x.bitcast(dtypes.uint32) @@ -21,11 +22,14 @@ class GradAccClipAdamW(Optimizer): def __init__(self, params:list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-6, weight_decay=0.0, grad_acc=1, clip_norm=1.0, device=None, fused=FUSE_OPTIM): super().__init__(params, lr, device, fused) self.b1, self.b2, self.eps, self.wd = b1, b2, eps, weight_decay - self.b1_t, self.b2_t = (Tensor.ones((1,), dtype=dtypes.float32, device=self.device, requires_grad=False) for _ in [b1, b2]) + self.b1_t, self.b2_t = (Tensor.ones((1,), dtype=dtypes.float32, device=self.device) for _ in [b1, b2]) self.m = self._new_optim_param() self.v = self._new_optim_param() self.grad_acc, self.clip_norm = grad_acc, clip_norm - self.master_params:list[Tensor]|None = [p.float().contiguous() for p in self.params] if MASTER_WEIGHTS and self.params[0].dtype != dtypes.float32 else None + if MASTER_WEIGHTS and self.params[0].dtype != dtypes.float32: + self.master_params:list[Tensor]|None = [p.to(self.device).float().contiguous() for p in self.params] + else: + self.master_params = None def fstep(self, grads:list[Tensor]): if self.fused: @@ -36,7 +40,8 @@ def fstep(self, grads:list[Tensor]): for i, tt in enumerate(self.params): tt.assign(self._apply_update(tt, updates[i], self.master_params[i] if self.master_params else None)) # collect inv_scale tensors attached to fp8 params (set by _apply_update) fp8_inv_scales = [tt._inv_scale for tt in self.params if hasattr(tt, '_inv_scale')] - to_realize = extra+self.params+self.buffers+(self.master_params or [])+fp8_inv_scales + fp8_next_inv_scales = [tt._next_inv_scale for tt in self.params if hasattr(tt, '_next_inv_scale')] + to_realize = extra+self.params+self.buffers+(self.master_params or [])+fp8_inv_scales+fp8_next_inv_scales Tensor.realize(*to_realize) return extra[-1] @@ -78,13 +83,23 @@ def _apply_update(self, t:Tensor, up:Tensor, master:Tensor|None=None) -> Tensor: up = up.float().shard_like(w) + self.lr.to(w.device) * wd * w.detach() new_w = w.detach() - up if master is not None: master.assign(new_w) - if STOCHASTIC_ROUND and t.dtype == dtypes.bfloat16: return stochastic_round_bf16(new_w) + # when master is offloaded to a different device than the param, results are resharded back onto the param's (sharded) device + offloaded = master is not None and master.device != t.device + if STOCHASTIC_ROUND and t.dtype == dtypes.bfloat16: + out = stochastic_round_bf16(new_w) + return out.shard_like(t) if offloaded else out if t.dtype in dtypes.fp8s: from examples.mlperf.models.flat_llama import FP8_MAX - amax = new_w.float().abs().max(axis=tuple(range(1, new_w.ndim))).detach() # per-layer amax for (n_layers, out, in) - scale = FP8_MAX / (amax + 1e-8) - fp8_w = (new_w * scale.reshape(-1, *([1]*(new_w.ndim-1)))).clamp(-FP8_MAX, FP8_MAX).cast(t.dtype) - if hasattr(t, '_inv_scale'): - t._inv_scale.assign(((amax + 1e-8) / FP8_MAX).cast(t._inv_scale.dtype)) - return fp8_w - return new_w.cast(t.dtype) + # delayed scaling: reuse previous step's inv_scale + t._inv_scale.assign(t._next_inv_scale) + inv_scale = t._inv_scale.to(new_w.device) if offloaded else t._inv_scale + scale = inv_scale.reciprocal().reshape(-1, *([1]*(new_w.ndim-1))) + scaled = (new_w * scale).clamp(-FP8_MAX, FP8_MAX) + ret = scaled.cast(t.dtype) + # update inv_scale for next step from quantized result + new_amax = (ret.float().abs().max(axis=tuple(range(1, ret.ndim))) * inv_scale * FP8_AMAX_MARGIN).detach() + new_inv = ((new_amax + 1e-8) / FP8_MAX).cast(t._inv_scale.dtype) + t._next_inv_scale.assign(new_inv.shard_like(t._next_inv_scale) if offloaded else new_inv) + return ret.shard_like(t) if offloaded else ret + out = new_w.cast(t.dtype) + return out.shard_like(t) if offloaded else out diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_405b/implementations/tinybox_8xMI350X/dev_beam.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_405b/implementations/tinybox_8xMI350X/dev_beam.sh index 07afe1469b8a8..026f52a8f4ddc 100755 --- a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_405b/implementations/tinybox_8xMI350X/dev_beam.sh +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_405b/implementations/tinybox_8xMI350X/dev_beam.sh @@ -1,6 +1,8 @@ #!/usr/bin/env bash export PYTHONPATH="." +export PATH="/opt/rocm-7.1.1/bin:$PATH" +export ROCM_PATH="/opt/rocm-7.1.1" export DEV=${DEV:-AMD} export CHECK_OOB=0 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000 @@ -16,10 +18,12 @@ export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1} export FP8=${FP8:-1} export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1} export FAST_CE=${FAST_CE:-0} -export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1} -export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1} -export FUSED_SILU_W13=${FUSED_SILU_W13:-1} -export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-1} +export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0} +export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0} +export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0} +export FUSED_SILU_W13=${FUSED_SILU_W13:-0} +export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-0} +export SPLIT_W13=${SPLIT_W13:-1} export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1} export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16" diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_405b/implementations/tinybox_8xMI350X/dev_run.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_405b/implementations/tinybox_8xMI350X/dev_run.sh index dffd5783fdde7..1abdac9666662 100755 --- a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_405b/implementations/tinybox_8xMI350X/dev_run.sh +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_405b/implementations/tinybox_8xMI350X/dev_run.sh @@ -1,22 +1,34 @@ #!/usr/bin/env bash export PYTHONPATH="." +export PATH="/opt/rocm-7.1.1/bin:$PATH" +export ROCM_PATH="/opt/rocm-7.1.1" export DEV=${DEV:-AMD} -export EMULATE="AMD_CDNA4" export CHECK_OOB=0 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000 +export DEVICE_IN_FUNCTION_BUG=1 export DEBUG=${DEBUG:-0} export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1} export ALL2ALL=${ALL2ALL:-1} -export USE_ATOMICS=${USE_ATOMICS:-0} +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 ALLREDUCE_CAST=${ALLREDUCE_CAST:-1} +export FAST_CE=${FAST_CE:-0} +export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0} +export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0} +export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0} +export FUSED_SILU_W13=${FUSED_SILU_W13:-0} +export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-0} +export SPLIT_W13=${SPLIT_W13:-1} export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1} export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16" -export DP=${DP:-1} MP=${MP:-8} -export BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1152} +export DP=${DP:-1} MP=${MP:-8} BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1152} +export GBS=$((BS * GRADIENT_ACC_STEPS)) export MODEL="llama3" export BASEDIR="/raid/datasets/c4/" 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 52dcfa32a82d5..aac9ba92d2169 100755 --- a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh @@ -1,6 +1,8 @@ #!/usr/bin/env bash export PYTHONPATH="." +export PATH="/opt/rocm-7.1.1/bin:$PATH" +export ROCM_PATH="/opt/rocm-7.1.1" export DEV=${DEV:-AMD} export CHECK_OOB=0 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000 @@ -21,7 +23,7 @@ export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1} export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-1} export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1} export FUSED_SILU_W13=${FUSED_SILU_W13:-1} -export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-1} +export SPLIT_W13=${SPLIT_W13:-0} export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16" export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2} diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam_mp.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam_mp.sh index 7369f12199205..037f2b2613034 100755 --- a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam_mp.sh +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam_mp.sh @@ -1,6 +1,8 @@ #!/usr/bin/env bash export PYTHONPATH="." +export PATH="/opt/rocm-7.1.1/bin:$PATH" +export ROCM_PATH="/opt/rocm-7.1.1" export DEV=${DEV:-AMD} export CHECK_OOB=0 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000 @@ -20,7 +22,6 @@ export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0} export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0} export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0} export FUSED_SILU_W13=${FUSED_SILU_W13:-0} -export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-0} export SPLIT_W13=${SPLIT_W13:-1} export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1} diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh index 3589721923e4f..6a29fbc44f6c5 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 @@ -1,6 +1,8 @@ #!/usr/bin/env bash export PYTHONPATH="." +export PATH="/opt/rocm-7.1.1/bin:$PATH" +export ROCM_PATH="/opt/rocm-7.1.1" export DEV=${DEV:-AMD} export CHECK_OOB=0 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000 @@ -21,7 +23,7 @@ export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1} export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-1} export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1} export FUSED_SILU_W13=${FUSED_SILU_W13:-1} -export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-1} +export SPLIT_W13=${SPLIT_W13:-0} export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16" export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2} diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run_mp.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run_mp.sh index 3703f0ab69b6e..0e5620540cf12 100755 --- a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run_mp.sh +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run_mp.sh @@ -1,8 +1,9 @@ #!/usr/bin/env bash export PYTHONPATH="." +export PATH="/opt/rocm-7.1.1/bin:$PATH" +export ROCM_PATH="/opt/rocm-7.1.1" export DEV=${DEV:-AMD} -export EMULATE="AMD_CDNA4" export CHECK_OOB=0 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000 export DEVICE_IN_FUNCTION_BUG=1 @@ -21,7 +22,6 @@ export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0} export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0} export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0} export FUSED_SILU_W13=${FUSED_SILU_W13:-0} -export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-0} export SPLIT_W13=${SPLIT_W13:-1} export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1} diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/profile.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/profile.sh index 17884d26270e7..a98e06ff7dd9a 100755 --- a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/profile.sh +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/profile.sh @@ -1,6 +1,6 @@ #!/bin/bash export BENCHMARK=5 export EVAL_BS=0 -VIZ=${VIZ:--1} FULL_LAYERS=1 DEBUG=0 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_beam.sh +VIZ=${VIZ:--1} FULL_LAYERS=1 DEBUG=0 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh SRC="AMD"; [[ $DEV == NULL* ]] && SRC="NULL" -python -m tinygrad.viz.cli -s "$SRC" -t +python -m tinygrad.viz.cli -s "$SRC" -t --interval "train @ 2" "train @ 3" diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/run_and_time.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/run_and_time.sh index d8ee6a053822f..e1d5a86734b75 100755 --- a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/run_and_time.sh +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/run_and_time.sh @@ -3,6 +3,8 @@ set -e # Exit on any error set -o pipefail # Make pipeline fail if any command fails export PYTHONPATH="." +export PATH="/opt/rocm-7.1.1/bin:$PATH" +export ROCM_PATH="/opt/rocm-7.1.1" export DEV=AMD export CHECK_OOB=0 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000 @@ -22,7 +24,7 @@ export FUSED_INPUT_QUANTIZE=1 export FUSED_GRAD_QUANTIZE=1 export FUSED_ADD_NORM_MUL_QUANTIZE=1 export FUSED_SILU_W13=1 -export FUSED_PAD_GRAD_ACCUM=1 +export SPLIT_W13=0 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16" export DP=8 MP=1 BS=16 EVAL_BS=8 GRADIENT_ACC_STEPS=2 diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/startup_walltime.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/startup_walltime.sh index b20cf20af8de2..71a8216c9be0e 100755 --- a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/startup_walltime.sh +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/startup_walltime.sh @@ -4,7 +4,7 @@ export EVAL_BS=0 export FAKEDATA=1 export NULL_ALLOW_COPYOUT=1 export HIP_VISIBLE_DEVICES="" -export DEV=NULL +export DEV=NULL:HIP:gfx950 export JITBEAM=0 export LLAMA_LAYERS=${LLAMA_LAYERS:-"2"} -time examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh +time examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh diff --git a/examples/mnist_gan.py b/examples/mnist_gan.py index 3e1ebe0ef0791..640d485032659 100644 --- a/examples/mnist_gan.py +++ b/examples/mnist_gan.py @@ -71,7 +71,7 @@ def train_generator(optimizer, data_fake): if __name__ == "__main__": # data for training and validation X_train, _, _, _ = mnist() - ds_noise = Tensor.randn(64, 128, requires_grad=False) + ds_noise = Tensor.randn(64, 128) # parameters epochs, batch_size, k = 300, 512, 1 sample_interval = epochs // 10 diff --git a/examples/openpilot/compile3.py b/examples/openpilot/compile3.py index bb604fa3dd6e7..031088f6c82e9 100644 --- a/examples/openpilot/compile3.py +++ b/examples/openpilot/compile3.py @@ -25,6 +25,8 @@ def compile(onnx_file): # TODO this seems dumb input_types = {k:(dtypes.float32 if v is dtypes.float16 else v) for k,v in input_types.items()} Tensor.manual_seed(100) + # replace symbolic dimensions (e.g. 'b' for dynamic batch) with 1 + input_shapes = {k:tuple(s if isinstance(s, int) else 1 for s in shp) for k,shp in input_shapes.items()} inputs = {k:Tensor(Tensor.randn(*shp, dtype=input_types[k]).mul(8).realize().numpy(), device='NPY') for k,shp in sorted(input_shapes.items())} if not getenv("NPY_IMG"): inputs = {k:Tensor(v.numpy(), device=Device.DEFAULT).realize() if 'img' in k else v for k,v in inputs.items()} @@ -89,7 +91,7 @@ def test_vs_compile(run, inputs, test_val=None): step_times.append((et-st)*1e3) print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {step_times[-1]:6.2f} ms") - if (assert_time:=getenv("ASSERT_MIN_STEP_TIME")): + if (assert_time:=getenv("ASSERT_MIN_STEP_TIME", 0.0)): min_time = min(step_times) assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms" @@ -106,7 +108,7 @@ def test_vs_compile(run, inputs, test_val=None): def test_vs_onnx(new_inputs, test_val, onnx_file, tol): import onnx import onnxruntime as ort - + onnx_inputs = {k:v.numpy() for k,v in new_inputs.items()} onnx_model = onnx.load(onnx_file) @@ -139,7 +141,7 @@ def bench(run, inputs): if __name__ == "__main__": if getenv("RUN_PICKLE"): with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f) - inputs = {name: Tensor(Tensor.randn(*[int(s) for s in view.src[1].arg], dtype=dtype).numpy(), device=device) + inputs = {name: Tensor(Tensor.randn(*view.shape, dtype=dtype).numpy(), device=device) for name, (view, _vars, dtype, device) in zip(pickle_loaded.captured.expected_names, pickle_loaded.captured.expected_input_info)} test_vs_compile(pickle_loaded, inputs) else: diff --git a/examples/vgg7.py b/examples/vgg7.py index 1b01fd1aaf337..0aa49c30b9005 100644 --- a/examples/vgg7.py +++ b/examples/vgg7.py @@ -164,8 +164,8 @@ def load_and_save(path, save): x_img = image_load(samples_base + "/" + str(sample_idx) + "a.png") y_img = image_load(samples_base + "/" + str(sample_idx) + "b.png") - sample_x = Tensor(x_img, requires_grad = False) - sample_y = Tensor(y_img, requires_grad = False) + sample_x = Tensor(x_img) + sample_y = Tensor(y_img) # magic code roughly from readme example # An explanation, in case anyone else has to go down this path: diff --git a/extra/export_model.py b/extra/export_model.py index 4ff53218a67b9..843abcfe770ad 100644 --- a/extra/export_model.py +++ b/extra/export_model.py @@ -23,7 +23,7 @@ def compile_net(linear:UOp, output_bufs:List[Buffer]) -> Tuple[Dict[str,str], Li def name_of(bu:UOp, is_out:bool) -> str: nonlocal n - if bu.op is Ops.PARAM: key, name, size = ("in", bu.arg), f"input{bu.arg}", prod(bu.shape)*bu.dtype.itemsize + if bu.op is Ops.PARAM: key, name, size = ("in", bu.arg.slot), f"input{bu.arg.slot}", prod(bu.shape)*bu.dtype.itemsize else: b = bu.buffer key, size = (id(b.base), b.offset, b.size, b.dtype), b.size*b.dtype.itemsize diff --git a/extra/gemm/amd_uop_matmul.py b/extra/gemm/amd_uop_matmul.py index 9deb336eac21a..e83233b04a1d8 100644 --- a/extra/gemm/amd_uop_matmul.py +++ b/extra/gemm/amd_uop_matmul.py @@ -122,7 +122,7 @@ def eval_custom_matmul(fxn, dt=dtypes.float): with Context(DEBUG=0): Tensor.realize(a, b) ets = [] - with Context(DEBUG=max(2, DEBUG.value), DEVECTORIZE=2 if dt == dtypes.half else 0): + with Context(DEBUG=max(2, DEBUG.value)): for _ in range(NUM_RUNS): GlobalCounters.reset() tst = Tensor.custom_kernel(c, a, b, fxn=fxn)[0].realize() diff --git a/extra/gemm/mi350x_uop_matmul.py b/extra/gemm/mi350x_uop_matmul.py index 8aba22eb19ac7..b8a82b784ab95 100644 --- a/extra/gemm/mi350x_uop_matmul.py +++ b/extra/gemm/mi350x_uop_matmul.py @@ -218,7 +218,7 @@ def compute_on_locals(acc:UOp, Asl:UOp, Bsl:UOp, rng:int, afters:tuple[UOp, ...] ref.realize() GlobalCounters.reset() - with Context(DEBUG=max(2, DEBUG.value), DEVECTORIZE=2): + with Context(DEBUG=max(2, DEBUG.value)): tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0] tst.realize() print(f"{(N*M*K*2 / GlobalCounters.time_sum_s)*1e-12:.2f} REAL TFLOPS") diff --git a/extra/gemm/mi350x_uop_matmul_2.py b/extra/gemm/mi350x_uop_matmul_2.py index b42c6d40e0eae..41c51f6e74010 100644 --- a/extra/gemm/mi350x_uop_matmul_2.py +++ b/extra/gemm/mi350x_uop_matmul_2.py @@ -127,7 +127,7 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp: GlobalCounters.reset() - with Context(DEBUG=max(2, DEBUG.value), DEVECTORIZE=2): + with Context(DEBUG=max(2, DEBUG.value)): tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0] tst.realize() print(f"{(N*M*K*2 / GlobalCounters.time_sum_s)*1e-12:.2f} REAL TFLOPS") diff --git a/extra/hcq2/graph/hcq.py b/extra/hcq2/graph/hcq.py index ea7395d22b95c..9ef809e4468fd 100644 --- a/extra/hcq2/graph/hcq.py +++ b/extra/hcq2/graph/hcq.py @@ -5,9 +5,13 @@ from tinygrad.dtype import dtypes from tinygrad.engine.jit import GraphRunner from tinygrad.engine.realize import get_call_outs_ins, get_runtime +from tinygrad.helpers import round_up, ceildiv +from tinygrad.runtime.support.memory import BumpAllocator from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, graph_rewrite -from extra.hcq2.hcq2 import HCQ2Compiled, HCQ2DeviceCtx, HCQ2LowerCtx, prep_runtime, pm_lower_kernargs, pm_lower_ops -from extra.hcq2.hcq2 import pm_split_into_queues, pm_add_barriers, pm_add_signals, build_host_program +from extra.hcq2.hcq2 import HCQ2Compiled, HCQ2DeviceCtx, HCQ2LowerCtx, pm_prep_runtime, pm_lower_ops +from extra.hcq2.hcq2 import pm_split_into_queues, pm_add_barriers, pm_add_signals +from extra.hcq2.hcq2 import pm_bufferize, pm_lift_patches_to_cmdbuf, pm_resolve_patches, pm_parametrize_host_buffers +from extra.hcq2.hcq2 import pm_add_timeline_inc, pm_callify, pm_calc_kernargs_sizes # **************** insert deps **************** @@ -22,18 +26,20 @@ def insert_deps(ctx:HCQ2Graph, linear:UOp) -> UOp: return linear.replace(src=tuple(src)) pm_insert_deps = PatternMatcher([(UPat(Ops.LINEAR, name="linear"), insert_deps)]) -def replace_params(ctx:HCQ2Graph, call:UOp) -> UOp|None: - if not any(x.op is Ops.PARAM for x in call.src[1:]): return None - return call.replace(src=tuple(ctx.input_addrs_uop[x.arg] if x.op is Ops.PARAM else x for x in call.src)) -pm_replace_params = PatternMatcher([(UPat(Ops.CALL, name="call", allow_any_len=True), replace_params)]) +pm_replace_params = PatternMatcher([ + (UPat(Ops.PARAM, name="p"), lambda ctx, p: ctx.input_addrs_uop.index(UOp.const(dtypes.int, p.arg))), + (UPat(Ops.SLICE, src=(UPat(Ops.INDEX, name="addr"), UPat(Ops.CONST, dtype=dtypes.weakint, name="off")), name="bv"), + lambda ctx, bv, addr, off: addr.cast(dtypes.uint64) + UOp.const(dtypes.uint64, off.arg * ctx.input_uops[addr.src[1].arg].dtype.itemsize)), +]) # **************** graph-only passes **************** def alloc_queue_sig(ctx:HCQ2Graph, q:UOp) -> None: if q.arg in ctx.queue_sigs: return None - buf = Buffer(q.arg[0], 0x100, dtypes.uint8, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True) + dev = q.arg[0][0] # TODO: multi device + buf = Buffer(dev, 0x100, dtypes.uint8, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True) ctx.queue_sig_bufs.append(buf) - ctx.queue_sigs[q.arg] = UOp.from_buffer(buf, q.arg[0]) + ctx.queue_sigs[q.arg] = UOp.from_buffer(buf, dev) return None pm_alloc_queue_sigs = PatternMatcher([(UPat(Ops.LINEAR, src=UPat({Ops.PROGRAM, Ops.COPY}), name="q"), alloc_queue_sig)]) @@ -69,12 +75,13 @@ def drop_dead_stores(ctx:HCQ2Graph, outer:UOp) -> UOp: return outer.replace(src=tuple(q.replace(src=tuple(x for x in q.src if x.op is not Ops.STORE or x in live)) for q in outer.src)) pm_drop_dead_stores = PatternMatcher([(UPat(Ops.LINEAR, src=UPat(Ops.LINEAR), name="outer"), drop_dead_stores)]) -def add_queue_sig_resets(ctx:HCQ2Graph, outer:UOp) -> UOp|None: - if not ctx.queue_sig_bufs: return None - resets = tuple(ctx.hcq_ctx.host_param(sig).index(UOp.const(dtypes.int, 0), ptr=True).cast(dtypes.uint64.ptr()) - .store(UOp.const(dtypes.uint64, 0)) for sig in ctx.queue_sig_bufs) - return outer.replace(src=tuple(c.replace(src=c.src + resets) if c.op is Ops.AFTER else c.after(*resets) for c in outer.src)) -pm_add_queue_sig_resets = PatternMatcher([(UPat(Ops.LINEAR, name="outer"), add_queue_sig_resets)]) +def add_queue_sig_resets(ctx:HCQ2Graph, x:UOp, cmdbuf:UOp) -> UOp|None: + if not ctx.queue_sig_bufs or cmdbuf.tag not in ("compute", "copy"): return None + resets = tuple((b:=UOp.from_buffer(sig)).index(UOp.const(dtypes.int, 0), dtype=b.dtype.ptr()) + .cast(dtypes.uint64.ptr()).store(UOp.const(dtypes.uint64, 0)) for sig in ctx.queue_sig_bufs) + return x.replace(src=x.src + resets) +pm_add_queue_sig_resets = PatternMatcher([(UPat(Ops.AFTER, src=(UPat(Ops.BUFFER, name="cmdbuf"),), allow_any_len=True, name="x"), + add_queue_sig_resets)]) # **************** Graph **************** @@ -85,16 +92,12 @@ def __init__(self, linear:UOp, input_uops:tuple[UOp, ...]=()): self.hcq_ctx = HCQ2LowerCtx(name="hcq_graph") self.input_addrs = Buffer("CPU", max(len(input_uops), 1), dtypes.uint64, preallocate=True) - self.input_addrs_uop = self.hcq_ctx.host_param(self.input_addrs) + self.input_addrs_uop = UOp.from_buffer(self.input_addrs, "CPU") self.linear = graph_rewrite(self.linear, pm_insert_deps, ctx=self, name="hcq: insert deps", walk=True) - self.linear, sizes = prep_runtime(self.hcq_ctx, self.linear) - for dev_name, sz in sizes.items(): - buf = Buffer(dev_name, sz, dtypes.uint8, options=BufferSpec(cpu_access=True), preallocate=True) - self.hcq_ctx.devs[dev_name] = HCQ2DeviceCtx(dev_name, UOp.from_buffer(buf, dev_name), UOp.const(dtypes.uint64, buf._buf.va_addr)) - self.linear = graph_rewrite(self.linear, pm_replace_params, ctx=self, name="hcq: replace params", walk=True) - self.linear = graph_rewrite(self.linear, pm_lower_kernargs + pm_lower_ops, ctx=self.hcq_ctx, name="hcq: lower ops") + self.linear = graph_rewrite(self.linear, pm_prep_runtime, ctx=self.hcq_ctx, name="hcq: prepare runtime") + self.linear = graph_rewrite(self.linear, pm_lower_ops, ctx=self.hcq_ctx, name="hcq: lower ops") # per-queue signal state — populated as a side-effect by pm_alloc_queue_sigs walking the lowered linear. self.queue_sig_bufs:list[Buffer] = [] @@ -107,9 +110,20 @@ def __init__(self, linear:UOp, input_uops:tuple[UOp, ...]=()): self.linear = graph_rewrite(self.linear, pm_optimize_queue_deps, ctx=self, name="hcq: optimize queue deps", walk=True) self.linear = graph_rewrite(self.linear, pm_drop_dead_stores, ctx=self, name="hcq: drop dead stores") self.linear = graph_rewrite(self.linear, pm_add_signals, ctx=self.hcq_ctx, name="hcq: add signals", walk=True) + self.linear = graph_rewrite(self.linear, pm_add_timeline_inc, ctx=self.hcq_ctx, name="hcq: add submit", walk=True) self.linear = graph_rewrite(self.linear, self.dev.pm_lower, ctx=self.hcq_ctx, name=f"hcq: encode cmdbuf {self.dev.device}", walk=True) + + graph_rewrite(self.linear, pm_calc_kernargs_sizes, ctx=(sizes:={}), name=None) + for dev_name, sz in sizes.items(): + buf = Buffer(dev_name, sz, dtypes.uint8, options=BufferSpec(cpu_access=True), preallocate=True) + self.hcq_ctx.dev_ctx[dev_name] = HCQ2DeviceCtx(dev_name, UOp.from_buffer(buf, dev_name), UOp.const(dtypes.uint64, buf._buf.va_addr)) + + self.linear = graph_rewrite(self.linear, pm_bufferize, ctx=self.hcq_ctx, bottom_up=True, name="realize binaries") + self.linear = graph_rewrite(self.linear, pm_lift_patches_to_cmdbuf, ctx=self.hcq_ctx, bottom_up=False, name="lift patches to cmdbuf") + self.linear = graph_rewrite(self.linear, pm_resolve_patches, ctx=self.hcq_ctx, bottom_up=False, name="simplify patches") self.linear = graph_rewrite(self.linear, pm_add_queue_sig_resets, ctx=self, name="hcq: add queue sig resets", walk=True) - self.host_call = build_host_program(self.hcq_ctx, self.linear, None, self.dev) + self.linear = graph_rewrite(self.linear, pm_parametrize_host_buffers, ctx=self.hcq_ctx, bottom_up=True, name="parametrize host buffers") + self.host_call = graph_rewrite(self.linear, pm_callify, ctx=self.hcq_ctx, name="hcq: callify") self.host_rt, self.host_globals = get_runtime("CPU", self.host_call.src[0]), self.host_call.src[0].arg.globals diff --git a/extra/hcq2/hcq2.py b/extra/hcq2/hcq2.py index bef477c170117..7ccd56df90b5a 100644 --- a/extra/hcq2/hcq2.py +++ b/extra/hcq2/hcq2.py @@ -1,39 +1,35 @@ from __future__ import annotations from typing import cast, Callable, TypeVar, Generic, Any, TYPE_CHECKING -import struct, functools, time, collections +import struct, functools, time, collections, importlib, itertools from dataclasses import replace if TYPE_CHECKING: from tinygrad.engine.realize import ExecContext -from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, mv_address, round_up, DEBUG, dedup -from tinygrad.device import Device, Buffer, BufferSpec, Compiled, LRUAllocator -from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, graph_rewrite, track_rewrites -from tinygrad.dtype import dtypes +from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, mv_address, round_up, DEBUG, dedup, pluralize +from tinygrad.device import Device, Buffer, BufferSpec, Compiled, LRUAllocator, MultiBuffer +from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, graph_rewrite, track_rewrites, GroupOp +from tinygrad.uop.symbolic import symbolic_simple, symbolic +from tinygrad.dtype import dtypes, DType from dataclasses import dataclass, field from tinygrad.runtime.support.memory import BumpAllocator from tinygrad.runtime.support.hcq import MMIOInterface from tinygrad.renderer import Renderer, Estimates -from tinygrad.engine.realize import to_program, track_stats, get_call_arg_uops, resolve_params +from tinygrad.engine.realize import to_program, track_stats, get_call_arg_uops, resolve_params, pm_flatten_linear HCQDeviceType = TypeVar('HCQDeviceType', bound='HCQ2Compiled') class HCQ2Compiled(Compiled): - """ - A base class for devices compatible with the HCQ (Hardware Command Queue) API. - """ - timestamp_divider: float = 1000.0 # GPU timestamp counter ticks per microsecond; override per device + timestamp_divider: float = 1000.0 # GPU timestamp counter ticks per microsecond; override per device - def __init__(self, device:str, allocator:'HCQAllocator', compilers:list[type[Renderer]], runtime, - kernargs_size=(16 << 20), can_recover:bool=False, arch=None): + def __init__(self, device:str, allocator:'HCQAllocator', compilers:list[type[Renderer]], runtime, can_recover:bool=False, arch=None): self.device_id:int = int(device.split(":")[1]) if ":" in device else 0 - from extra.hcq2.graph.hcq import HCQ2Graph - super().__init__(device, allocator, compilers, lambda *a, **kw: None, HCQ2Graph, arch=arch) + # default pm bufferize + self.pm_bufferize = PatternMatcher([ + (UPat(Ops.BUFFER, tag="timeline_signal"), lambda ctx: ctx.timeline_signal), + (UPat(Ops.BUFFER, tag="timeline_value"), lambda ctx: ctx.timeline_value), + (UPat(Ops.BUFFER, name="b"), lambda ctx, b: Buffer(ctx.device, b.arg, b.dtype, options=BufferSpec(host=True, uncached=True, cpu_access=True))), + ]) - self.kernargs_size = kernargs_size - self.kernargs_offset_allocator:BumpAllocator = BumpAllocator(kernargs_size, wrap=True) - - @functools.cached_property - def kernargs_buf(self) -> Buffer: - return Buffer(self.device, self.kernargs_size, dtypes.uint8, options=BufferSpec(cpu_access=True), preallocate=True) + super().__init__(device, allocator, compilers, lambda *a, **kw: None, None, arch=arch) @functools.cached_property def timeline_signal(self) -> Buffer: @@ -59,14 +55,6 @@ def synchronize(self, timeout:int|None=None): def device_props(self) -> dict[str,Any]: return {} # to be overridden if needed. dict keys are backend dependent. - def _realloc(self, oldbuf:HCQ2Buffer|None, new_size:int, options:BufferSpec|None=None, force=False) -> tuple[HCQ2Buffer, bool]: - if oldbuf is not None: self.allocator.free(oldbuf, oldbuf.size, options=options) - try: buf, realloced = self.allocator.alloc(new_size, options=options), True - except MemoryError: - if force: raise - buf, realloced = self.allocator.alloc(oldbuf.size if oldbuf is not None else new_size, options=options), False - return buf, realloced - def count(self) -> int: return self.iface.count if hasattr(self, 'iface') else 1 def _select_iface(self): @@ -125,7 +113,7 @@ def _wrap(self, dev:str, sz:int, opaque:HCQ2Buffer) -> Buffer: def _copy(self, dst:Buffer, src:Buffer): from tinygrad.engine.realize import run_linear su = UOp.from_buffer(src) - run_linear(UOp(Ops.LINEAR, dtypes.void, (su.copy_to_device(dst.device).call(UOp.from_buffer(dst), su),)), jit=True, update_stats=False) + run_linear(UOp(Ops.LINEAR, dtypes.void, (su.copy_to_device(dst.device).call(UOp.from_buffer(dst), su),)), update_stats=False) def _copyin(self, dest:HCQ2Buffer, src:memoryview): s = Buffer(self.dev.device, len(src), dtypes.uint8, options=BufferSpec(host=True), preallocate=True) @@ -138,240 +126,270 @@ def _copyout(self, dest:memoryview, src:HCQ2Buffer): self.dev.synchronize() dest[:] = d._buf.cpu_view()[:len(dest)] - def _as_buffer(self, buf): return buf.cpu_view().mv - -# **************** lower context **************** + # def _as_buffer(self, buf): return buf.cpu_view().mv -@dataclass -class HCQ2DeviceCtx: - device:str # device name; resolve to instance via Device[device] - kernargs_host:UOp # UOp whose .buffer is dev.kernargs_buf (BUFFER UOp in runtime, PARAM in graph) - kernargs_gpu:UOp # va_addr const of dev.kernargs_buf - kernargs_allocator:BumpAllocator = field(default_factory=lambda: BumpAllocator(2 << 20, wrap=False)) - -@dataclass -class HCQ2LowerCtx: - name:str - inputs:list[Buffer] = field(default_factory=list) - holds:list[UOp] = field(default_factory=list) - devs:dict[str, HCQ2DeviceCtx] = field(default_factory=dict) - - def host_param(self, buf:Buffer) -> UOp: - if buf not in self.inputs: self.inputs.append(buf) - return UOp.placeholder((buf.size,), buf.dtype, self.inputs.index(buf)) +def unwrap_after(uop): + while uop.op is Ops.AFTER: uop = uop.src[0] + return uop class HCQEncoder: - def __init__(self, ctx:HCQ2LowerCtx, dev:HCQ2Compiled): self.ctx, self.dev, self.blob, self.patches, self.deps = ctx, dev, b'', [], [] - - @property - def src(self) -> tuple[UOp, ...]: return tuple(self.patches + dedup(self.deps)) + def __init__(self): self.blob, self.patches = b'', [] - def get_dev_addr(self, uop:UOp) -> sint|UOp: - while uop.op is Ops.AFTER: - self.deps.extend(uop.src[1:]) - uop = uop.src[0] - if isinstance(val:=uop.ssimplify(), UOp): self.deps.append(uop) - return uop.buffer.get_buf(self.dev.device).va_addr if uop.op in (Ops.BUFFER, Ops.BUFFER_VIEW) else val + def get_dev_addr(self, uop:UOp) -> UOp: + if unwrap_after(uop).op not in (Ops.BUFFER, Ops.SLICE, Ops.BINARY, Ops.MSTACK, Ops.MSELECT): return uop + return UOp(Ops.GETADDR, dtypes.uint64, src=(uop, UOp(Ops.DEVICE, arg=self.dev.device))) def append(self, *data, dtype=dtypes.uint32): for d in data: if isinstance(d, int): self.blob += struct.pack(f'<{dtype.fmt}', d) - elif d.op is Ops.CONST: self.blob += struct.pack(f'<{dtype.fmt}', d.arg) else: - self.patches.append(UOp(Ops.PATCH, dtype, src=(d,), arg=len(self.blob))) + self.patches.append((len(self.blob), self.get_dev_addr(d), dtype)) self.blob += struct.pack(f'<{dtype.fmt}', 0) def q(self, *values): self.append(*values) -# **************** prep runtime **************** + def uop(self, dev:str|tuple[str, ...], tag:str|None=None) -> UOp: + buf = UOp.new_buffer(dev, len(self.blob), dtypes.uint8) + if tag: buf = buf.rtag(tag) + blob_uop = UOp(Ops.BINARY, dtypes.void, src=(), arg=self.blob) + stores = [buf.index(UOp.const(dtypes.int, off), dtype=buf.dtype.ptr()).cast(dt.ptr()).store(val.cast(dt)) for off, val, dt in self.patches] + return buf.after(buf.store(blob_uop), *stores) -pm_prep_runtime = PatternMatcher([ - # device-specific lowering of the program - (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.DEVICE), UPat(), UPat(), UPat(Ops.BINARY)), name="p"),), name="c", allow_any_len=True), - lambda ctx, c, p: c.replace(src=(Device[p.src[1].arg].pm_lower.rewrite(p, ctx),) + c.src[1:])), -]) +# ***************** +# 0. helpers -def calc_kernargs_sizes(ctx:dict[str,int], u:UOp) -> None: - d = u.src[0].buffer.device - ctx[d] = ctx.get(d, 0) + round_up(u.arg[0].kernargs_alloc_size, 16) -pm_calc_kernargs_sizes = PatternMatcher([(UPat(Ops.PROGRAM, name="u"), calc_kernargs_sizes)]) +HCQ_DEVS = frozenset(("AMD",)) +HCQ_P2P_DEVS = HCQ_DEVS | frozenset(("CPU",)) -# **************** lower kernargs **************** +def to_tuple(d): return d if isinstance(d, tuple) else (d,) -def lower_kernargs(ctx:HCQ2LowerCtx, call:UOp, prg:UOp) -> UOp: - data, info = prg.arg - # after amd_build_program, prg.src is (BUFFER_lib_gpu,); the buffer's device names the device - dctx = ctx.devs[prg.src[0].buffer.device] +def all_devices_in(d:Any, c:frozenset[str]) -> bool: return {x.split(":")[0] for x in to_tuple(d)} <= c - enc = HCQEncoder(ctx, Device[dctx.device]) - for gi in info.globals: enc.append(enc.get_dev_addr(call.src[1+gi]), dtype=dtypes.uint64) - for v in info.vars: enc.append(v, dtype=dtypes.uint32) +# ***************** +# 1.1. prep runtimes: staging copies - args_off = dctx.kernargs_allocator.alloc(data.kernargs_alloc_size, 16) - dctx.kernargs_host.buffer.view(len(enc.blob), dtypes.uint8, args_off).ensure_allocated().as_memoryview(force_zero_copy=True)[:] = enc.blob +def _need_staging(a, b): return all_devices_in(a.device, HCQ_DEVS) and not all_devices_in(b.device, HCQ_P2P_DEVS) - args_uop = (dctx.kernargs_gpu + args_off).after(dctx.kernargs_host.after(*tuple(p.replace(arg=p.arg+args_off) for p in enc.patches))) - return call.replace(src=(prg.replace(src=prg.src + (args_uop,), arg=(data, info)),) + call.src[1:]) +def stage_copy(dst:UOp, src:UOp) -> UOp|None: + if not (_need_staging(src, dst) or _need_staging(dst, src)): return None -pm_lower_kernargs = PatternMatcher([ - (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(Ops.BUFFER),), name="prg"),), name="call", allow_any_len=True), lower_kernargs), -]) + stage = UOp.new_buffer("CPU", src.buffer.nbytes, dtypes.uint8) + return UOp(Ops.LINEAR, dtypes.void, (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)]) -# **************** lower ops **************** +# ***************** +# 1.2. prep runtimes: programs/kernargs -def lower_program(ctx:HCQ2LowerCtx, call:UOp, prg:UOp) -> UOp: - q = UOp(Ops.LINEAR, dtypes.void, (prg,), arg=(prg.src[0].buffer.device, "COMPUTE")) - return UOp(Ops.LINEAR, dtypes.void, (q,), tag=call.tag) +@functools.cache +def get_pm_prep_program(name:str) -> PatternMatcher|None: + try: + importlib.import_module(f'tinygrad.runtime.ops_{name.lower()}') # TODO: remove that + return importlib.import_module(f'extra.hcq2.ops_{name.lower()}2').pm_prep_program + except ImportError: return None -def lower_copy(ctx:HCQ2LowerCtx, call:UOp, copy:UOp) -> UOp: - dst, src = call.src[1], call.src[2] - q = UOp(Ops.LINEAR, dtypes.void, (UOp(Ops.COPY, dtypes.void, src=(dst, src), arg=src.buffer.nbytes),), arg=(dst.buffer.device, "COPY")) - return UOp(Ops.LINEAR, dtypes.void, (q,), tag=call.tag) +def prep_program(call:UOp, prg:UOp) -> UOp|None: + dev = call.src[1].device + if (pm:=get_pm_prep_program(to_tuple(dev)[0].split(":")[0])) is None or (lowered:=pm.rewrite(prg)) is None: return None -pm_lower_ops = PatternMatcher([ - (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(Ops.BUFFER), UPat()), name="prg"),), name="call", allow_any_len=True), lower_program), - (UPat(Ops.CALL, src=(UPat(Ops.COPY, name="copy"),), name="call", allow_any_len=True), lower_copy), -]) + data, image_bytes = lowered + buf = UOp.new_buffer(dev, len(image_bytes), dtypes.uint8).rtag("program") + blob = UOp(Ops.BINARY, dtypes.void, src=(), arg=image_bytes) + return call.replace(src=(prg.replace(src=(buf.after(buf.store(blob)),), arg=(data, prg.arg)),) + call.src[1:]) -# **************** split into queues **************** +def prep_kernargs(call:UOp, prg:UOp) -> UOp: + data, info = prg.arg + patches = [(i*dtypes.uint64.itemsize, UOp(Ops.GETADDR, dtypes.uint64, src=(call.src[1+gi], UOp(Ops.DEVICE, arg=call.src[1+gi].device))), + dtypes.uint64) for i,gi in enumerate(info.globals)] \ + + [(len(info.globals)*dtypes.uint64.itemsize + i*dtypes.uint32.itemsize, v, dtypes.uint32) for i,v in enumerate(info.vars)] + + buf = UOp.new_buffer(call.src[1].device, data.kernargs_alloc_size, dtypes.uint8).rtag("kernargs") + kernargs = buf.after(*tuple(buf.index(UOp.const(dtypes.int, o), dtype=buf.dtype.ptr()).cast(dt.ptr()).store(val.cast(dt)) for o, val, dt in patches)) + + return call.replace(src=(prg.replace(src=prg.src + (kernargs,), arg=(data, info)),) + call.src[1:]) + +pm_prep_runtime = PatternMatcher([ + # bind generic PROGRAM device to the call's actual dev(s), then run device-specific lowering + (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(), UPat(), UPat(), UPat(), UPat(Ops.BINARY)), name="prg"),), + name="call", allow_any_len=True), prep_program), -def split_into_queues(ctx:HCQ2LowerCtx, outer:UOp) -> UOp: - groups:dict[tuple, list[UOp]] = collections.defaultdict(list) - for child in outer.src: - wrapper = child.src[0] if child.op is Ops.AFTER else child - for q in wrapper.src: groups[q.arg].extend(q.src) - return outer.replace(src=tuple(UOp(Ops.LINEAR, dtypes.void, tuple(cmds), arg=k) for k, cmds in groups.items())) -pm_split_into_queues = PatternMatcher([(UPat(Ops.LINEAR, src=UPat(Ops.LINEAR, src=UPat(Ops.LINEAR)).or_after(), name="outer"), split_into_queues)]) + # lower kernargs (PROGRAM.src[0] is now AFTER(BUFFER, COPY) — the lowered program image) + (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(Ops.AFTER),), name="prg"),), name="call", allow_any_len=True), prep_kernargs), +]) -# **************** add signals (runtime) **************** +# ***************** +# 2.1. lowering to hcq ir -def add_signals(ctx:HCQ2LowerCtx, outer:UOp) -> UOp: - def wrap(q:UOp) -> UOp: - (dev_name, qname), devs = q.arg, {q.arg[0]} | {u.buffer.device for u in q.toposort() if u.op in (Ops.BUFFER, Ops.BUFFER_VIEW)} - sigs_tls = [(UOp.from_buffer(Device[d].timeline_signal), ctx.host_param(Device[d].timeline_value)) for d in sorted(devs) if d.startswith("AMD")] - return q.replace(src=(*(s.wait(t[0]-1) for s,t in sigs_tls), *q.src, *(s.store(t[0]) for s,t in sigs_tls)), arg=qname) - return outer.replace(src=tuple(wrap(q) for q in outer.src)) +def lower_program(call:UOp, prg:UOp) -> UOp: + q = UOp(Ops.LINEAR, dtypes.void, (prg,), arg=(call.src[1].device, "COMPUTE")) + return call.replace(src=(q,) + call.src[1:]).rtag('hcq') -pm_add_barriers = PatternMatcher([(UPat(Ops.LINEAR, src=UPat(Ops.LINEAR), name="outer"), - lambda ctx, outer: outer.replace(src=tuple(q.replace(src=(UOp(Ops.BARRIER, dtypes.void), *q.src)) for q in outer.src)))]) +def lower_copy(call:UOp, copy:UOp) -> UOp|None: + dst, src = call.src[1], call.src[2] + if (hcq_dev:=next((b.device for b in (dst, src) if b.device.split(":")[0] in HCQ_DEVS), None)) is None: return None -pm_add_signals = PatternMatcher([(UPat(Ops.LINEAR, src=UPat(Ops.LINEAR), name="outer"), add_signals)]) + q = UOp(Ops.LINEAR, dtypes.void, (UOp(Ops.COPY, dtypes.void, src=(dst, src), arg=src.buffer.nbytes),), arg=(hcq_dev, "COPY")) + return call.replace(src=(q,) + call.src[1:]).rtag('hcq') -# **************** build host program **************** +pm_lower_ops = PatternMatcher([ + (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(Ops.AFTER), UPat(Ops.AFTER)), name="prg"),), name="call", allow_any_len=True), lower_program), + (UPat(Ops.CALL, src=(UPat(Ops.COPY, name="copy"),), name="call", allow_any_len=True), lower_copy), +]) -def resolve_cmdbuf(ctx:HCQ2LowerCtx, blob:UOp) -> UOp: - inner = blob.src[0] if blob.op is Ops.AFTER else blob - dev_name, qtype = inner.tag +# ***************** +# 2.2. queue split + +# def split_into_queues(linear:UOp) -> UOp: +# out = [] +# for k, grp in itertools.groupby(linear.src, lambda c: c.src[0].arg if c.op is Ops.CALL and c.src[0].op is Ops.LINEAR else None): +# if k is None: out.extend(grp) +# else: +# calls = list(grp) +# items = tuple(x for c in calls for x in c.src[0].src) +# args = tuple(a for c in calls for a in c.src[1:]) +# out.append(calls[0].replace(src=(UOp(Ops.LINEAR, dtypes.void, items, arg=k),) + args)) +# return linear.replace(src=tuple(out)) +# pm_split_into_queues = PatternMatcher([(UPat(Ops.LINEAR, name="linear"), split_into_queues)]) + +# ***************** +# 2.3. barriers / signals / timeline inc + +def add_barriers(call:UOp, q:UOp) -> UOp: + return call.replace(src=(q.replace(src=(UOp(Ops.BARRIER, dtypes.void), *q.src)),) + call.src[1:]) +pm_add_barriers = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.LINEAR, name="q"),), name="call", allow_any_len=True), add_barriers)]) + +def add_signals(call:UOp, q:UOp) -> UOp: + sig = UOp.new_buffer(q.arg[0], 0x100, dtypes.uint8).rtag("timeline_signal") + tl = UOp.new_buffer(q.arg[0], 1, dtypes.uint64).rtag("timeline_value").index(UOp.const(dtypes.int, 0)) + return call.replace(src=(q.replace(src=(sig.wait(tl-1), *q.src, sig.store(tl)), arg=q.arg),) + call.src[1:]) +pm_add_signals = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.LINEAR, name="q"),), name="call", allow_any_len=True), add_signals)]) + +# ***************** +# 3.1. encode cmdbufs + +@functools.cache +def get_pm_lower(name:str) -> PatternMatcher|None: + try: + importlib.import_module(f'tinygrad.runtime.ops_{name.lower()}') # TODO: remove that + return importlib.import_module(f'extra.hcq2.ops_{name.lower()}2').pm_lower + except ImportError: return None + +def encode_cmdbuf(call:UOp, q:UOp) -> UOp|None: + if (pm:=get_pm_lower(to_tuple(q.arg[0])[0].split(":")[0])) is None or (encoded:=pm.rewrite(q)) is None: return None + return call.replace(src=(encoded,) + call.src[1:]) +pm_encode_cmdbufs = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.LINEAR, name="q"),), name="call", allow_any_len=True), encode_cmdbuf)]) + +# ***************** +# 3.2. add timeline inc + +def add_timeline_inc(call:UOp, s:UOp) -> UOp: + tl = UOp.new_buffer(s.device, 1, dtypes.uint64).rtag("timeline_value") + return call.replace(src=(tl.after(s).index(UOp.const(dtypes.int, 0), dtype=tl.dtype.ptr()).store(tl.index(UOp.const(dtypes.int, 0)) + 1),) + call.src[1:]) +pm_add_timeline_inc = PatternMatcher([(UPat(Ops.CALL, tag="hcq", src=(UPat(name="s"),), name="call", allow_any_len=True), add_timeline_inc)]) + +# ***************** +# 3.3. lift patches to the command buffer (root) + +def lift_patches_to_cmdbuf(cmdbuf:UOp) -> UOp|None: + if not (patches:=dedup(u for store in cmdbuf.src[1:] for u in store.toposort() if u.op is Ops.AFTER)): return None + deps = tuple(d for p in patches for d in p.src[1:]) + return cmdbuf.replace(src=cmdbuf.src + deps).substitute({p: p.src[0] for p in patches}) +pm_lift_patches_to_cmdbuf = PatternMatcher([ + (UPat(Ops.AFTER, src=(UPat(Ops.BUFFER, tag={"compute", "copy"}),), allow_any_len=True, name="cmdbuf"), lift_patches_to_cmdbuf), +]) - # prepare the cmdbuf and make it a param - bb = Buffer("CPU", len(inner.arg)//4, dtypes.uint32, preallocate=True) - bb.copyin(memoryview(bytearray(inner.arg))) - bb_param = ctx.host_param(bb) +# ***************** +# 4. bufferize placeholders: replace placeholders with real buffers. - submit_cf = UOp(Ops.CUSTOM_FUNCTION, dtypes.void, src=(bb_param.after(*(blob.src[1:] if blob.op is Ops.AFTER else ())),), - arg=f"submit_{qtype.lower()}", tag=dev_name) +def bufferize_buf(buf:UOp) -> UOp|None: + if buf.tag is None: return None + uops = tuple(UOp.from_buffer((dv:=Device[dev]).pm_bufferize.rewrite(buf, ctx=dv), dev) for dev in to_tuple(buf.src[1].arg)) + return uops[0] if len(uops) == 1 else UOp(Ops.MSTACK, uops[0].dtype, uops) +pm_bufferize = PatternMatcher([(UPat(Ops.BUFFER, name="buf"), bufferize_buf)]) - # increment the timeline value - tl = ctx.host_param(Device[dev_name].timeline_value) - return tl.after(UOp(Ops.BARRIER, dtypes.void, src=(submit_cf,))).index(UOp.const(dtypes.int, 0), ptr=True).store(tl[0] + 1) +# ***************** +# 5.1. capture buffers reachable from each hcq call as BIND, so we don't drop their refs -def resolve_patches(ctx:HCQ2LowerCtx, buf:UOp) -> UOp|None: - inner = buf.src[0] +def hold_call_buffers(call:UOp) -> UOp|None: + if not (bufs:=tuple(dedup(u for u in call.src[0].toposort() if u.op is Ops.BUFFER and u not in call.src))): return None + return call.replace(src=call.src + (UOp(Ops.BIND, dtypes.void, src=bufs),)) +pm_hold_call_buffers = PatternMatcher([(UPat(Ops.CALL, tag="hcq", name="call"), hold_call_buffers)]) - # buffer is accessed from the launcher, so transform it to a host param - if inner.op is Ops.BUFFER: inner = ctx.host_param(inner.buffer) +# ***************** +# 5.2. resolve patches - return inner.after(*(inner.index(UOp.const(dtypes.int, p.arg//inner.dtype.base.itemsize), ptr=True).cast(p.dtype.ptr()).store(p.src[0].cast(p.dtype)) - if p.op is Ops.PATCH else p for p in buf.src[1:])) +def push_stack(op, s): return UOp(Ops.STACK, op.dtype.scalar().vec(len(s.src)), + tuple(op.replace(dtype=op.dtype.scalar(), src=tuple(x if y is s else y for y in op.src)) for x in s.src)) -def resolve_ref_buffers(ctx:HCQ2LowerCtx, buf:UOp) -> UOp: - if buf not in ctx.holds: ctx.holds.append(buf) +def fold_blob_store(buf:UOp, blob:UOp) -> UOp: + for b in (buf.src if buf.op is Ops.MSTACK else (buf,)): b.buffer.ensure_allocated()._buf.cpu_view().mv.cast('B')[:len(blob.arg)] = blob.arg return UOp(Ops.NOOP) -def hcq_callify(ctx:HCQ2LowerCtx, sink:UOp) -> UOp: - call = to_program(sink, Device["CPU"].renderer).call(*[UOp.from_buffer(b, "CPU") if isinstance(b, Buffer) else b for b in ctx.inputs]) - return call.replace(src=call.src + (UOp(Ops.BIND, dtypes.void, src=tuple(ctx.holds)),)) if ctx.holds else call - -pm_create_host_sink = PatternMatcher([ - (UPat(Ops.LINEAR, name="l", allow_any_len=True), lambda ctx, l: UOp.sink(*l.src, arg=KernelInfo(name=ctx.name, estimates=Estimates()), tag=1)) -]) +def fold_const_store(buf:UOp, off:UOp, val:UOp) -> UOp: + for b, v in zip((buf.src if buf.op is Ops.MSTACK else (buf,)), (val.src if val.op is Ops.STACK else (val,))): + struct.pack_into(f'<{v.dtype.fmt}', b.buffer.ensure_allocated()._buf.cpu_view().mv.cast('B'), off.arg * b.dtype.base.itemsize, v.arg) + return UOp(Ops.NOOP) -# lower cmdbuf submits -pm_lower_cmdbufs = PatternMatcher([ - (UPat(Ops.AFTER, src=(UPat(Ops.BINARY),), name="blob", allow_any_len=True), resolve_cmdbuf), - (UPat(Ops.BINARY, name="blob"), resolve_cmdbuf), -]) +def resolve_getaddr(buf:UOp, g:UOp) -> UOp: + if isinstance(b:=buf.buffer, Buffer): return UOp.const(dtypes.uint64, b.get_buf(g.src[1].arg).va_addr) + return UOp(Ops.STACK, dtypes.uint64.vec(len(b.bufs)), tuple(UOp.const(dtypes.uint64, x.ensure_allocated()._buf.va_addr) for x in b.bufs)) -# transform patches attached to buffers and params pm_resolve_patches = PatternMatcher([ - (UPat(Ops.AFTER, src=(UPat((Ops.BUFFER, Ops.PARAM)),), name="buf", allow_any_len=True), resolve_patches) + # multi + (UPat(GroupOp.ALU, src=[UPat(Ops.STACK, name="s"), UPat(Ops.CONST)], name="op"), push_stack), + (UPat(Ops.CAST, src=(UPat(Ops.STACK, name="s"),), name="op"), push_stack), + + # getaddr + (UPat(Ops.GETADDR, src=(UPat(Ops.SLICE, name="bv"), UPat(Ops.DEVICE, name="dev"))), # getaddr(slice(x)) -> offset+getaddr(x) + lambda bv, dev: UOp(Ops.GETADDR, dtypes.uint64, src=(bv.src[0], dev)) + UOp.const(dtypes.uint64, bv.src[1].arg * bv.src[0].dtype.itemsize)), + (UPat(Ops.GETADDR, src=(UPat({Ops.BUFFER, Ops.MSTACK, Ops.MSELECT}, name="buf"), UPat(Ops.DEVICE)), name="g"), resolve_getaddr), + + # folders + (UPat({Ops.BUFFER, Ops.MSTACK}, name="buf").store(UPat(Ops.BINARY, name="blob")), fold_blob_store), + (UPat({Ops.BUFFER, Ops.MSTACK}, name="buf").index(UPat.cvar("off")).or_casted().store(UPat.any(UPat.cvar("val"), UPat(Ops.STACK, name="val"))), + fold_const_store), +]) + symbolic_simple + +# ***************** +# 6. callify hcq programs + +pm_fixup = PatternMatcher([ # TODO: this should gone? + (UPat(Ops.CONST, name="c"), lambda c: c.replace(src=()) if len(c.src) else None), ]) -# replace referenced buffers with noops -pm_resolve_ref_buffers = PatternMatcher([(UPat((Ops.BUFFER, Ops.BUFFER_VIEW), name="buf"), resolve_ref_buffers)]) - -pm_callify = PatternMatcher([(UPat(Ops.SINK, name="sink"), hcq_callify)]) - -# **************** schedule **************** - -def prep_runtime(ctx:HCQ2LowerCtx, linear:UOp) -> tuple[UOp, dict[str,int]]: - linear = graph_rewrite(linear, pm_prep_runtime, ctx=ctx, name="hcq: prepare runtime") - graph_rewrite(linear, pm_calc_kernargs_sizes, ctx=(sizes:={}), enter_calls=True) - return linear, sizes - -def build_host_program(ctx:HCQ2LowerCtx, linear:UOp, ast:UOp, dev:HCQ2Compiled) -> UOp: - sink = graph_rewrite(linear, pm_create_host_sink, ctx=ctx, name="hcq: create host sink", walk=True) - sink = graph_rewrite(sink, pm_lower_cmdbufs, ctx=ctx, bottom_up=True, name="hcq: lower cmdbufs") - sink = graph_rewrite(sink, pm_resolve_patches, ctx=ctx, bottom_up=True, name="hcq: resolve patches") - sink = graph_rewrite(sink, pm_resolve_ref_buffers, ctx=ctx, bottom_up=True, name="hcq: resolve ref buffers") - sink = graph_rewrite(sink, dev.pm_lower, ctx=ctx, name=f"hcq: device lower {dev.device}", walk=True) - return graph_rewrite(sink, pm_callify, ctx=ctx, name="hcq: callify") - -@track_rewrites(name=lambda ctx,linear,ast,dev,**kw: f"hcq schedule {getattr(ast.arg, 'name', ast.op.name.lower())}") -def hcq_schedule(ctx:HCQ2LowerCtx, linear:UOp, ast:UOp, dev:HCQ2Compiled) -> UOp: - linear, sizes = prep_runtime(ctx, linear) - for dev_name, sz in sizes.items(): - off = dev.kernargs_offset_allocator.alloc(sz, 16) - ctx.devs[dev_name] = HCQ2DeviceCtx(dev_name, UOp.from_buffer(dev.kernargs_buf.view(sz, dtypes.uint8, off), dev_name), - UOp.const(dtypes.uint64, dev.kernargs_buf.get_buf(dev_name).va_addr + off)) - linear = graph_rewrite(linear, pm_lower_kernargs + pm_lower_ops, ctx=ctx, name="hcq: lower ops") - linear = graph_rewrite(linear, pm_split_into_queues, ctx=ctx, name="hcq: split into queues") - linear = graph_rewrite(linear, pm_add_barriers, ctx=ctx, name="hcq: add barriers", walk=True) - linear = graph_rewrite(linear, pm_add_signals, ctx=ctx, name="hcq: add signals", walk=True) - linear = graph_rewrite(linear, dev.pm_lower, ctx=ctx, name=f"hcq: encode cmdbuf {dev.device}", walk=True) - return build_host_program(ctx, linear, ast, dev) - -def ensure_accessible(ctx:HCQ2LowerCtx, call:UOp, copy:UOp) -> UOp|None: - src_buf = call.src[2].buffer # TODO: cleanup - dev = call.src[1].buffer.device - try: src_buf.get_buf(dev) - except Exception: - (cpubuf := Buffer("CPU", src_buf.nbytes, dtypes.uint8, preallocate=True)).copyin(src_buf.ensure_allocated().as_memoryview()) - ctx.holds.append(buf_uop:=UOp.from_buffer(cpubuf, dev)) - return call.replace(src=call.src[:2] + (buf_uop,) + call.src[3:]) -pm_ensure_bufs_accessible = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.COPY, name="copy"),), name="call", allow_any_len=True), ensure_accessible)]) - -def hcq_exec(ctx:ExecContext, call:UOp, ast:UOp) -> float|None: - from tinygrad.engine.realize import run_linear - - if ast.src[1].arg.split(":")[0] != "AMD": return None - - # TODO: this mess should gone - resolved_call = call.replace(src=(ast,) + tuple(resolve_params(call, ctx.input_uops)) + tuple(s for s in call.src[1:] if s.op is Ops.BIND)) - bufs = [cast(Buffer, resolved_call.src[1+gi].buffer) for gi in ast.arg.globals] if ast.op is Ops.PROGRAM \ - else [cast(Buffer, resolved_call.src[i].buffer) for i in range(1, len(resolved_call.src))] - dev = cast(HCQ2Compiled, Device[bufs[0].device]) - hcq_ctx = HCQ2LowerCtx(name="submit") - linear = graph_rewrite(UOp(Ops.LINEAR, dtypes.void, (resolved_call,)), pm_ensure_bufs_accessible, ctx=hcq_ctx) - host_call = hcq_schedule(hcq_ctx, linear, ast, dev) - with track_stats(ctx, call, dev.device, bufs, ctx.var_vals) as tm: - st = time.perf_counter() if ctx.wait else 0.0 - run_linear(UOp(Ops.LINEAR, dtypes.void, (host_call,)), var_vals=ctx.var_vals, jit=True, update_stats=DEBUG>=3) - if ctx.wait: - dev.synchronize() - tm[0] = time.perf_counter() - st - return tm[0] if tm[0] is not None else 0.0 - -pm_hcq_exec = PatternMatcher([ - (UPat(Ops.CALL, src=(UPat({Ops.PROGRAM, Ops.COPY}, name="ast"),), name="call", allow_any_len=True), hcq_exec), -]) +def to_param(bufs:list[UOp], ref:UOp) -> UOp: + bufs.append(ref) + return UOp.placeholder((ref.buffer.size,), ref.dtype, len(bufs)-1) +pm_to_param = PatternMatcher([(UPat({Ops.MSELECT, Ops.MSTACK, Ops.BUFFER}, name="r"), lambda ctx, r: to_param(ctx, r))]) + +def parametrize_host_buffers(call:UOp) -> UOp: + body = graph_rewrite(call.src[0], pm_to_param, ctx=(bufs:=[]), bottom_up=True, name="parametrize host buffers") + return call.replace(src=(body, *bufs) + call.src[1:], tag="hcq_param") +pm_parametrize_host_buffers = PatternMatcher([(UPat(Ops.CALL, tag="hcq", name="call"), parametrize_host_buffers)]) + +def callify_hcq(call:UOp) -> UOp: + sink = UOp.sink(call.src[0], arg=KernelInfo(name="hcq_submit", estimates=Estimates()), tag=1) + return to_program(sink, Device["CPU"].renderer).call(*call.src[1:]) +pm_callify_hcq = PatternMatcher([(UPat(Ops.CALL, tag="hcq_param", name="call"), callify_hcq)]) + +@track_rewrites(lambda _,ret: f"HCQ Schedule {pluralize('Kernel', len(ret.src))}") +def hcq_schedule(linear:UOp) -> UOp: + linear = graph_rewrite(linear, pm_insert_copy_staging + pm_flatten_linear, name="insert copy staging") + linear = graph_rewrite(linear, pm_prep_runtime, name="prepare runtime") + + linear = graph_rewrite(linear, pm_lower_ops, name="lower ops into hcq ir") + # linear = graph_rewrite(linear, pm_split_into_queues, name="split into queues") + linear = graph_rewrite(linear, pm_add_barriers, walk=True, name="add barriers") + linear = graph_rewrite(linear, pm_add_signals, walk=True, name="add signals") + linear = graph_rewrite(linear, pm_encode_cmdbufs, walk=True, name="encode cmdbufs") + linear = graph_rewrite(linear, pm_add_timeline_inc, walk=True, name="add timeline inc") + linear = graph_rewrite(linear, pm_lift_patches_to_cmdbuf, name="lift patches to cmdbuf", enter_calls=True) + + # realize starts from here + linear = graph_rewrite(linear, pm_bufferize, bottom_up=True, name="bufferize placeholders", enter_calls=True) + linear = graph_rewrite(linear, pm_hold_call_buffers, walk=True, name="hold call buffers") + linear = graph_rewrite(linear, pm_resolve_patches, bottom_up=False, name="simplify patches", enter_calls=True) + linear = graph_rewrite(linear, pm_fixup, bottom_up=False, name="fixup", enter_calls=True) + linear = graph_rewrite(linear, pm_parametrize_host_buffers, name="parametrize host buffers") + linear = graph_rewrite(linear, pm_callify_hcq, name="callify hcq") + + return linear diff --git a/extra/hcq2/ops_amd2.py b/extra/hcq2/ops_amd2.py index ba0e3c7920ff4..7c096471c7fd7 100644 --- a/extra/hcq2/ops_amd2.py +++ b/extra/hcq2/ops_amd2.py @@ -23,14 +23,14 @@ from tinygrad.runtime.ops_amd import EVENT_INDEX_PARTIAL_FLUSH, WAIT_REG_MEM_FUNCTION_EQ, WAIT_REG_MEM_FUNCTION_NEQ, WAIT_REG_MEM_FUNCTION_GEQ if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl # noqa: F401 # pylint: disable=unused-import -from extra.hcq2.hcq2 import HCQ2LowerCtx from tinygrad.engine.realize import get_runtime from tinygrad.uop.ops import Ops, UPat, PatternMatcher, graph_rewrite class AMDComputeQueue(HCQEncoder): - def __init__(self, ctx:HCQ2LowerCtx, dev:AMDDevice): - super().__init__(ctx, dev) - self.pm4, self.gc, self.nbio, self.soc = self.dev.pm4, self.dev.gc, self.dev.nbio, self.dev.soc + def __init__(self, dev:AMDDevice, devs:tuple[str, ...]|None=None): + super().__init__() + self.dev, self.devs = dev, devs or (dev.device,) + self.pm4, self.gc, self.nbio, self.soc = dev.pm4, dev.gc, dev.nbio, dev.soc def pkt3(self, cmd, *vals): self.q(self.pm4.PACKET3(cmd, len(vals) - 1), *vals) @@ -105,10 +105,12 @@ def program(self, x): self.acquire_mem(gli=0, gl2=0) + scratch_addr = self.get_dev_addr(UOp.new_buffer(self.devs, data.private_segment_size, dtypes.uint8).rtag("scratch")) args_addr = self.get_dev_addr(args) + user_regs = [] if data.enable_private_segment_sgpr: - scratch_hilo = data64_le(self.dev.scratch.va_addr) + scratch_hilo = data64_le(scratch_addr) user_regs = [scratch_hilo[0], scratch_hilo[1] | 1 << 31, 0xffffffff, 0x20c14000] if data.enable_dispatch_ptr: user_regs += [*data64_le(args_addr + data.kernargs_segment_size)] user_regs += [*data64_le(args_addr)] @@ -116,10 +118,10 @@ def program(self, x): self.wreg(self.gc.regCOMPUTE_PGM_LO, *data64_le(prog_addr >> 8)) self.wreg(self.gc.regCOMPUTE_PGM_RSRC1, data.rsrc1, data.rsrc2) self.wreg(self.gc.regCOMPUTE_PGM_RSRC3, data.rsrc3) - self.wreg(self.gc.regCOMPUTE_TMPRING_SIZE, self.dev.tmpring_size) + self.wreg(self.gc.regCOMPUTE_TMPRING_SIZE, self.dev.tmpring_size(data.private_segment_size)) for xcc_id in range(self.dev.xccs): - scratch_base = self.dev.scratch.va_addr + (self.dev.scratch.size // self.dev.xccs * xcc_id) + scratch_base = scratch_addr + (data.private_segment_size // self.dev.xccs * xcc_id) self.wreg(self.gc.regCOMPUTE_DISPATCH_SCRATCH_BASE_LO, *data64_le(scratch_base >> 8)) self.wreg(self.gc.regCOMPUTE_RESTART_X, 0, 0, 0) @@ -133,42 +135,47 @@ def program(self, x): self.pkt3(self.pm4.PACKET3_EVENT_WRITE, self.pm4.EVENT_TYPE(self.soc.CS_PARTIAL_FLUSH) | self.pm4.EVENT_INDEX(EVENT_INDEX_PARTIAL_FLUSH)) amd_inner_pm = PatternMatcher([ - (UPat(Ops.WAIT, name="x"), lambda ctx, x: ctx.wait(x)), - (UPat(Ops.BARRIER, name="x"), lambda ctx, x: ctx.barrier(x)), - (UPat(Ops.PROGRAM, name="x"), lambda ctx, x: ctx.program(x)), - (UPat(Ops.CUSTOM_FUNCTION, arg="timestamp", name="x"), lambda ctx, x: ctx.timestamp(x)), - (UPat(Ops.STORE, src=(UPat((Ops.BUFFER, Ops.PARAM)), UPat()), name="x"), lambda ctx, x: ctx.store(x)), + (UPat(Ops.LINEAR, src=(UPat(Ops.WAIT, name="x"),)), lambda ctx, x: ctx.wait(x)), + (UPat(Ops.LINEAR, src=(UPat(Ops.BARRIER, name="x"),)), lambda ctx, x: ctx.barrier(x)), + (UPat(Ops.LINEAR, src=(UPat(Ops.PROGRAM, name="x"),)), lambda ctx, x: ctx.program(x)), + (UPat(Ops.LINEAR, src=(UPat(Ops.CUSTOM_FUNCTION, arg="timestamp", name="x"),)), lambda ctx, x: ctx.timestamp(x)), + (UPat(Ops.LINEAR, src=(UPat(Ops.STORE, src=(UPat((Ops.BUFFER, Ops.PARAM)), UPat()), name="x"),)), lambda ctx, x: ctx.store(x)), ]) -def amd_lower_pm4(ctx, linear): - prg = next(s for s in linear.src if s.op is Ops.PROGRAM) - dev = Device[prg.src[1].arg] - enc = AMDComputeQueue(ctx, dev) - graph_rewrite(linear, amd_inner_pm, ctx=enc, name="amd: encode") - return UOp(Ops.BINARY, dtypes.void, arg=enc.blob).rtag((dev.device, "COMPUTE")).after(*enc.src) - -def amd_submit_pm4(ctx, cf): - dev = Device[cf.tag] - bb_param = cf.src[0] - q = dev.compute_queue - ring, wptr, doorbell, put_ptr = (ctx.host_param(b) for b in (q.ring, q.write_ptr, q.doorbell, q.put_value)) - size, ring_dwords = UOp.const(dtypes.uint32, bb_param.dtype.size), q.ring.size - - put = put_ptr[0] - i = UOp.range(size, 0, dtype=dtypes.int) +def amd_lower_pm4(linear, devs): + enc = AMDComputeQueue(Device[devs[0]], devs) + graph_rewrite(linear.replace(src=tuple(UOp(Ops.LINEAR, dtypes.void, (cmd,)) for cmd in linear.src)), amd_inner_pm, ctx=enc, name="amd: encode") + return enc.uop(dev=devs if len(devs) > 1 else devs[0], tag="compute") + +def amd_submit_pm4(cmdbuf, devs): + size, zero = UOp.const(dtypes.uint32, cmdbuf.src[0].arg // dtypes.uint32.itemsize), UOp.const(dtypes.int, 0) + + # the compute queue's ring and its host-side ring/write/put pointers (placeholders, resolved in pm_bufferize) + q = Device['AMD'].compute_queue + ring, wptr, doorbell, put_ptr = (UOp.new_buffer(devs, b.size, b.dtype).rtag(("COMPUTE:0", name)) + for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value))) + + # place the cmdbuf at the ring's write offset, wrapping the ring + put = put_ptr.index(zero) next_put = put + size.cast(put.dtype) - ring_idx = ((put + i.cast(put.dtype)) % ring_dwords).cast(dtypes.int) + i = UOp.range(size, 0, dtype=dtypes.int, src=(cmdbuf,)) + ring_idx = ((put + i.cast(put.dtype)) % q.ring.size).cast(dtypes.int) + + # copy the cmdbuf into the ring and advance the put/write pointers + copy_to_ring = ring.index(ring_idx, dtype=ring.dtype.ptr()).store( + cmdbuf.index(i*4, dtype=cmdbuf.dtype.ptr()).cast(dtypes.uint32.ptr()).load()).end(i) + bump_put_ptr = put_ptr.index(zero, dtype=put_ptr.dtype.ptr()).store(next_put) + bump_wptr = wptr.index(zero, dtype=wptr.dtype.ptr()).store(next_put) - copy_to_ring = ring[ring_idx].store(bb_param[i]).end(i) - bump_put_ptr = put_ptr[0].store(next_put) - bump_wptr = wptr[0].store(next_put) + # ring the doorbell once the copy and pointer bumps have landed flush = UOp.barrier(copy_to_ring, bump_put_ptr, bump_wptr) - return doorbell.after(flush)[0].store(next_put) + return doorbell.after(flush).index(zero, dtype=doorbell.dtype.ptr()).store(next_put) class AMDCopyQueue(HCQEncoder): - def __init__(self, ctx:HCQ2LowerCtx, dev:AMDDevice, queue_idx=0): - super().__init__(ctx, dev) - self.sdma, self.queue_idx, self.max_copy_size = self.dev.sdma, queue_idx, self.dev.max_copy_size + def __init__(self, dev:AMDDevice, queue_idx=0): + super().__init__() + self.dev = dev + self.sdma, self.queue_idx, self.max_copy_size = dev.sdma, queue_idx, dev.max_copy_size def copy(self, x): dest, src, copy_size = self.get_dev_addr(x.src[0]), self.get_dev_addr(x.src[1]), x.arg @@ -193,89 +200,93 @@ def timestamp(self, x): self.q(self.sdma.SDMA_OP_TIMESTAMP | self.sdma.SDMA_PKT_TIMESTAMP_GET_HEADER_SUB_OP(self.sdma.SDMA_SUBOP_TIMESTAMP_GET_GLOBAL), *data64_le(self.get_dev_addr(x.src[0]))) -def amd_lower_sdma(ctx, linear): - copy = next(s for s in linear.src if s.op is Ops.COPY) - dev = Device[copy.src[0].buffer.device] - enc = AMDCopyQueue(ctx, dev) - graph_rewrite(linear, amd_inner_sdma_pm, ctx=enc, name="amd: encode sdma") - return UOp(Ops.BINARY, dtypes.void, arg=enc.blob).rtag((dev.device, "COPY")).after(*enc.src) +def amd_lower_sdma(linear, devs): + enc = AMDCopyQueue(Device[devs[0]]) + graph_rewrite(linear.replace(src=tuple(UOp(Ops.LINEAR, dtypes.void, (cmd,)) for cmd in linear.src)), amd_inner_sdma_pm, ctx=enc, name="amd: encode sdma") + return enc.uop(dev=devs if len(devs) > 1 else devs[0], tag="copy") amd_inner_sdma_pm = PatternMatcher([ - (UPat(Ops.WAIT, name="x"), lambda ctx, x: ctx.wait(x)), - (UPat(Ops.BARRIER, name="x"), lambda ctx, x: None), - (UPat(Ops.COPY, name="x"), lambda ctx, x: ctx.copy(x)), - (UPat(Ops.CUSTOM_FUNCTION, arg="timestamp", name="x"), lambda ctx, x: ctx.timestamp(x)), - (UPat(Ops.STORE, src=(UPat((Ops.BUFFER, Ops.PARAM)), UPat()), name="x"), lambda ctx, x: ctx.store(x)), + (UPat(Ops.LINEAR, src=(UPat(Ops.WAIT, name="x"),)), lambda ctx, x: ctx.wait(x)), + (UPat(Ops.LINEAR, src=(UPat(Ops.BARRIER, name="x"),)), lambda ctx, x: None), + (UPat(Ops.LINEAR, src=(UPat(Ops.COPY, name="x"),)), lambda ctx, x: ctx.copy(x)), + (UPat(Ops.LINEAR, src=(UPat(Ops.CUSTOM_FUNCTION, arg="timestamp", name="x"),)), lambda ctx, x: ctx.timestamp(x)), + (UPat(Ops.LINEAR, src=(UPat(Ops.STORE, src=(UPat((Ops.BUFFER, Ops.PARAM)), UPat()), name="x"),)), lambda ctx, x: ctx.store(x)), ]) -def amd_submit_sdma(ctx, cf): - dev = Device[cf.tag] - bb_param = cf.src[0] - q = dev.sdma_queue(0) - ring, wptr, doorbell, put_ptr = (ctx.host_param(b) for b in (q.ring, q.write_ptr, q.doorbell, q.put_value)) - size_dw, ring_bytes = bb_param.dtype.size, q.ring.size * 4 +def amd_submit_sdma(cmdbuf, devs): + # the cmdbuf to submit + the patch writes that fill it + size_dw, zero = cmdbuf.src[0].arg // dtypes.uint32.itemsize, UOp.const(dtypes.int, 0) + + # the sdma queue's ring and its host-side ring/write/put pointers + q = Device['AMD'].sdma_queue(0) + ring, wptr, doorbell, put_ptr = (UOp.new_buffer(devs, b.size, b.dtype).rtag(("SDMA:0", name)) + for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value))) - put_b = put_ptr[0] - tail_off_dw = ((put_b % ring_bytes) // 4).cast(dtypes.int) + # sdma needs the cmdbuf contiguous: if it won't fit before the ring end, restart at 0 and zero the tail + put_b = put_ptr.index(zero) + tail_off_dw = ((put_b % (q.ring.size * 4)) // 4).cast(dtypes.int) fits = (size_dw <= q.ring.size - tail_off_dw).cast(dtypes.int) start_dw = fits * tail_off_dw zero_amt_dw = (1 - fits) * (q.ring.size - tail_off_dw) - zi = UOp.range(zero_amt_dw, 0, dtype=dtypes.int) - zero_tail = ring[tail_off_dw + zi].store(UOp.const(dtypes.uint32, 0)).end(zi) - - i = UOp.range(UOp.const(dtypes.int, size_dw), 0, dtype=dtypes.int) - copy_to_ring = ring[start_dw + i].store(bb_param[i]).end(i) + # zero the wrapped tail, then copy the cmdbuf into the ring + zi = UOp.range(zero_amt_dw, 0, dtype=dtypes.int, src=(cmdbuf,)) + zero_tail = ring.index(tail_off_dw + zi, dtype=ring.dtype.ptr()).store(UOp.const(dtypes.uint32, 0)).end(zi) + i = UOp.range(UOp.const(dtypes.int, size_dw), 0, dtype=dtypes.int, src=(cmdbuf,)) + copy_to_ring = ring.index(start_dw + i, dtype=ring.dtype.ptr()).store( + cmdbuf.index(i*4, dtype=cmdbuf.dtype.ptr()).cast(dtypes.uint32.ptr()).load()).end(i) + # advance the put/write pointers past the zeroed tail and the cmdbuf next_put_b = put_b + ((zero_amt_dw + size_dw) * 4).cast(put_b.dtype) - bump_put_ptr = put_ptr[0].store(next_put_b) - bump_wptr = wptr[0].store(next_put_b) + bump_put_ptr = put_ptr.index(zero, dtype=put_ptr.dtype.ptr()).store(next_put_b) + bump_wptr = wptr.index(zero, dtype=wptr.dtype.ptr()).store(next_put_b) + + # ring the doorbell once the writes have landed flush = UOp.barrier(zero_tail, copy_to_ring, bump_put_ptr, bump_wptr) - return doorbell.after(flush)[0].store(next_put_b) + return doorbell.after(flush).index(zero, dtype=doorbell.dtype.ptr()).store(next_put_b) @dataclass(frozen=True) class AMDProgramData: entry_point_offset:int; rsrc1:int; rsrc2:int; rsrc3:int; wave32:bool - kernargs_segment_size:int; kernargs_alloc_size:int + private_segment_size:int; kernargs_segment_size:int; kernargs_alloc_size:int enable_dispatch_ptr:int; enable_private_segment_sgpr:int -_amd_program_cache:dict[tuple[bytes,str], tuple[AMDProgramData,Buffer]] = {} +_amd_program_cache:dict[tuple[bytes,str], tuple[AMDProgramData,bytes]] = {} -def amd_build_program(ctx:HCQ2LowerCtx, prg:UOp) -> UOp: - dev = Device[prg.src[1].arg] +def amd_build_program(prg:UOp) -> UOp: + dev = Device[prg.src[1].arg] # TODO: rm this if (cached:=_amd_program_cache.get(key:=(lib:=prg.src[4].arg, dev.device))) is None: image, sections, relocs = elf_loader(lib) rodata = next(sh.header.sh_addr for sh in sections if sh.name == ".rodata") for off, sym, typ, addent in relocs: assert typ == 5, f"unknown AMD reloc {typ}" # R_AMDGPU_REL64 image[off:off+8] = struct.pack(' (dev.iface.props['lds_size_in_kb']*1024)//512: raise RuntimeError("Too many resources requested: group_segment_size") - dev._ensure_has_local_memory(desc.private_segment_fixed_size) edp = desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_DISPATCH_PTR cached = _amd_program_cache[key] = (AMDProgramData( entry_point_offset=rodata + desc.kernel_code_entry_byte_offset, rsrc1=desc.compute_pgm_rsrc1 | ((1<<20) if dev.target[0]==11 else 0), # priv=1 on gfx11 for cwsr rsrc2=desc.compute_pgm_rsrc2 | (lds<<15), rsrc3=desc.compute_pgm_rsrc3, wave32=bool(desc.kernel_code_properties & 0x400), + private_segment_size=desc.private_segment_fixed_size, kernargs_segment_size=desc.kernarg_size, kernargs_alloc_size=desc.kernarg_size + (ctypes.sizeof(hsa.hsa_kernel_dispatch_packet_t) if edp else 0), enable_dispatch_ptr=edp, - enable_private_segment_sgpr=desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER, - ), lib_gpu) - data, lib_gpu = cached - return prg.replace(src=(UOp.from_buffer(lib_gpu, dev.device),), arg=(data, prg.arg)) + enable_private_segment_sgpr=desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER), bytes(image)) + return cached + +pm_prep_program = PatternMatcher([ + (UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.DEVICE, arg="AMD"), UPat(), UPat(), UPat(Ops.BINARY)), name="prg"), amd_build_program), +]) 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()) def _alloc(self, size:int, options:BufferSpec) -> HCQ2Buffer: - return self.dev.iface.alloc(size, host=True, uncached=options.uncached, cpu_access=True) + return self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access or not self.dev.has_sdma_queue) def _do_free(self, opaque, options:BufferSpec): self.dev.iface.free(opaque) @@ -362,17 +373,18 @@ def device_fini(self): self.dev_impl.fini() def _mock(iface, name=None): return type(name or f"MOCK{iface.__name__}", (iface,), {}) +def encode_queue(q:UOp) -> UOp|None: + if not (isinstance(q.arg, tuple) and len(q.arg) == 2 and q.arg[1] in ("COMPUTE", "COPY")): return None + devs = (q.arg[0],) if isinstance(q.arg[0], str) else q.arg[0] # TODO: make this prettier + return amd_submit_pm4(amd_lower_pm4(q, devs), devs) if q.arg[1] == "COMPUTE" else amd_submit_sdma(amd_lower_sdma(q, devs), devs) + +pm_lower = PatternMatcher([ + (UPat(Ops.LINEAR, name="q"), encode_queue), +]) + class AMDDevice(HCQ2Compiled): timestamp_divider = 100.0 # AMD GPU clock: ticks/us - pm_lower = PatternMatcher([ - (UPat(Ops.PROGRAM, src=(UPat(), UPat(), UPat(), UPat(), UPat(Ops.BINARY)), name="prg"), amd_build_program), - (UPat(Ops.LINEAR, arg="COMPUTE", name="linear"), amd_lower_pm4), - (UPat(Ops.LINEAR, arg="COPY", name="linear"), amd_lower_sdma), - (UPat(Ops.CUSTOM_FUNCTION, arg="submit_compute", name="cf"), amd_submit_pm4), - (UPat(Ops.CUSTOM_FUNCTION, arg="submit_copy", name="cf"), amd_submit_sdma), - ]) - ifaces = [PCIIface] def is_am(self) -> bool: return isinstance(self.iface, (PCIIface,)) @@ -411,14 +423,13 @@ def __init__(self, device:str=""): 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 = self.sdma_queue(0) is not None + self.has_sdma_queue = True # self.sdma_queue(0) is not None, TODO: think of this - super().__init__(device, AMDAllocator(self), [HIPRenderer, AMDLLVMRenderer, HIPCCRenderer], None, - kernargs_size=16 << 20, can_recover=self.is_am(), arch=self.arch) + super().__init__(device, AMDAllocator(self), [HIPRenderer, AMDLLVMRenderer, HIPCCRenderer], None, can_recover=self.is_am(), arch=self.arch) # Scratch setup self.max_private_segment_size = 0 - self._ensure_has_local_memory(128) # set default scratch size to 128 bytes per thread + self.pm_bufferize = PatternMatcher([(UPat(Ops.BUFFER, tag="scratch", name="b"), lambda ctx, b: ctx.scratch_buffer(b.arg))]) + self.pm_bufferize self.pmc_enabled:bool = PROFILE > 0 and PMC > 0 if self.pmc_enabled: @@ -445,19 +456,6 @@ def __init__(self, device:str=""): self.sqtt_wptrs = self.allocator.alloc(round_up(self.se_cnt * self.xccs * 4, 0x1000), BufferSpec(cpu_access=True, nolru=True)) self.sqtt_next_cmd_id = itertools.count(0) - @functools.cached_property - def compute_queue(self) -> AMDQueueDesc: - # https://gitlab.freedesktop.org/agd5f/linux/-/blob/a1fc9f584c4aaf8bc1ebfa459fc57a3f26a290d8/drivers/gpu/drm/amd/amdkfd/kfd_queue.c#L391 - sgrp_size_per_cu, hwreg_size_per_cu = 0x4000, 0x1000 - lds_size_per_cu = self.iface.props["lds_size_in_kb"] << 10 if self.target[:2] == (9,5) else 0x10000 - vgpr_size_per_cu = 0x60000 if self.target in {(11,0,0), (11,0,1), (11,5,1), (12,0,0), (12,0,1)} else 0x80000 if self.target[0] == 9 else 0x40000 - wg_data_size = round_up((vgpr_size_per_cu + sgrp_size_per_cu + lds_size_per_cu + hwreg_size_per_cu) * self.cu_cnt, mmap.PAGESIZE) - ctl_stack_size = round_up((12 if self.target[0] != 9 else 8) * self.wave_cnt + 8 + 40, mmap.PAGESIZE) - return self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL if self.is_aql else kfd.KFD_IOC_QUEUE_TYPE_COMPUTE, - 0x2000 if self.is_usb() else (16 << 20), eop_buffer_size=0x1000, - ctx_save_restore_size=0 if self.is_am() else wg_data_size + ctl_stack_size, ctl_stack_size=ctl_stack_size, - debug_memory_size=round_up(self.wave_cnt * 32, 64)) - def create_queue(self, queue_type, ring_size, ctx_save_restore_size=0, eop_buffer_size=0, ctl_stack_size=0, debug_memory_size=0, idx=0): ring = self.iface.alloc(ring_size, uncached=True, cpu_access=True) gart = self.iface.alloc(0x100, uncached=True, cpu_access=True) @@ -473,9 +471,29 @@ def create_queue(self, queue_type, ring_size, ctx_save_restore_size=0, eop_buffe cwsr_buffer = self.iface.alloc(cwsr_buffer_size) if ctx_save_restore_size else None eop_buffer = self.iface.alloc(eop_buffer_size) if eop_buffer_size else None - return (self.iface.create_queue(queue_type, ring, gart, rptr=getattr(hsa.amd_queue_t, 'read_dispatch_id').offset, - wptr=getattr(hsa.amd_queue_t, 'write_dispatch_id').offset, eop_buffer=eop_buffer, cwsr_buffer=cwsr_buffer, - ctx_save_restore_size=ctx_save_restore_size, ctl_stack_size=ctl_stack_size, idx=idx)) + queue = (self.iface.create_queue(queue_type, ring, gart, rptr=getattr(hsa.amd_queue_t, 'read_dispatch_id').offset, + wptr=getattr(hsa.amd_queue_t, 'write_dispatch_id').offset, eop_buffer=eop_buffer, cwsr_buffer=cwsr_buffer, + ctx_save_restore_size=ctx_save_restore_size, ctl_stack_size=ctl_stack_size, idx=idx)) + + qname = f"{'SDMA' if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA else 'COMPUTE'}:{idx}" + self.pm_bufferize = PatternMatcher([ + (UPat(Ops.BUFFER, tag={(qname, name)}), lambda ctx, b=getattr(queue, name): b) for name in ["ring", "write_ptr", "doorbell", "put_value"] + ]) + self.pm_bufferize + + return queue + + @functools.cached_property + def compute_queue(self) -> AMDQueueDesc: + # https://gitlab.freedesktop.org/agd5f/linux/-/blob/a1fc9f584c4aaf8bc1ebfa459fc57a3f26a290d8/drivers/gpu/drm/amd/amdkfd/kfd_queue.c#L391 + sgrp_size_per_cu, hwreg_size_per_cu = 0x4000, 0x1000 + lds_size_per_cu = self.iface.props["lds_size_in_kb"] << 10 if self.target[:2] == (9,5) else 0x10000 + vgpr_size_per_cu = 0x60000 if self.target in {(11,0,0), (11,0,1), (11,5,1), (12,0,0), (12,0,1)} else 0x80000 if self.target[0] == 9 else 0x40000 + wg_data_size = round_up((vgpr_size_per_cu + sgrp_size_per_cu + lds_size_per_cu + hwreg_size_per_cu) * self.cu_cnt, mmap.PAGESIZE) + ctl_stack_size = round_up((12 if self.target[0] != 9 else 8) * self.wave_cnt + 8 + 40, mmap.PAGESIZE) + return self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL if self.is_aql else kfd.KFD_IOC_QUEUE_TYPE_COMPUTE, + 0x2000 if self.is_usb() else (16 << 20), eop_buffer_size=0x1000, + ctx_save_restore_size=0 if self.is_am() else wg_data_size + ctl_stack_size, ctl_stack_size=ctl_stack_size, + debug_memory_size=round_up(self.wave_cnt * 32, 64)) def sdma_queue(self, idx:int): if getenv("AMD_DISABLE_SDMA"): return None @@ -484,38 +502,49 @@ def sdma_queue(self, idx:int): self.sdma_queues[idx] = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x200 if self.is_usb() else (16 << 20), idx=idx) return self.sdma_queues.get(idx, None) - def _ensure_has_local_memory(self, private_segment_size): - if self.max_private_segment_size >= private_segment_size: return + def tmpring_size(self, private_segment_size): + private_segment_size = max(private_segment_size, 128) lanes_per_wave = 64 # wave64 mem_alignment_size = 256 if self.target[0] != 9 else 1024 size_per_thread = round_up(private_segment_size, mem_alignment_size // lanes_per_wave) size_per_xcc = size_per_thread * lanes_per_wave * self.iface.props['max_slots_scratch_cu'] * self.cu_cnt - self.scratch, ok = self._realloc(getattr(self, 'scratch', None), size_per_xcc * self.xccs) - if ok: - # NOTE: xcc logic is correct only for GFX9. - max_scratch_waves = self.cu_cnt * self.iface.props['max_slots_scratch_cu'] * self.xccs - wave_scratch = ceildiv(lanes_per_wave * size_per_thread, mem_alignment_size) - num_waves = (size_per_xcc // (wave_scratch * mem_alignment_size)) // (self.se_cnt if self.target[0] != 9 else 1) - - tmpring_t = getattr(hsa, f'union_COMPUTE_TMPRING_SIZE{"_GFX"+str(self.target[0]) if self.target[0] != 9 else ""}_bitfields') - self.tmpring_size = int.from_bytes(tmpring_t(WAVES=min(num_waves, max_scratch_waves), WAVESIZE=wave_scratch), 'little') - self.max_private_segment_size = private_segment_size - if hasattr(self, 'aql_desc'): - gfx9_rsrc = {'NUM_FORMAT':hsa.BUF_NUM_FORMAT_UINT, 'DATA_FORMAT':hsa.BUF_DATA_FORMAT_32, 'ELEMENT_SIZE':1, 'INDEX_STRIDE':3} - rsrc = {'DST_SEL_X':hsa.SQ_SEL_X, 'DST_SEL_Y':hsa.SQ_SEL_Y, 'DST_SEL_Z':hsa.SQ_SEL_Z, 'DST_SEL_W':hsa.SQ_SEL_W, 'ADD_TID_ENABLE':1, - 'TYPE':hsa.SQ_RSRC_BUF, **(gfx9_rsrc if self.target[0] == 9 else {'FORMAT':hsa.BUF_FORMAT_32_UINT, 'OOB_SELECT':2})} - rsrc1_t = getattr(hsa, f'union_SQ_BUF_RSRC_WORD1{"_GFX11" if self.target[0] != 9 else ""}_bitfields') - rsrc3_t = getattr(hsa, f'union_SQ_BUF_RSRC_WORD3{"_GFX"+str(self.target[0]) if self.target[0] != 9 else ""}_bitfields') - - self.aql_desc.scratch_backing_memory_location = int(self.scratch.va_addr) - self.aql_desc.scratch_wave64_lane_byte_size = self.max_private_segment_size * lanes_per_wave // 64 - self.aql_desc.scratch_resource_descriptor[:] = [lo32(self.scratch.va_addr), - int.from_bytes(rsrc1_t(BASE_ADDRESS_HI=hi32(self.scratch.va_addr), SWIZZLE_ENABLE=1), 'little'), - lo32(size_per_xcc), int.from_bytes(bytes(rsrc3_t(**rsrc)), 'little')] - self.aql_desc.compute_tmpring_size = self.tmpring_size - self.aql_gart.cpu_view()[:ctypes.sizeof(self.aql_desc)] = bytes(self.aql_desc) + # NOTE: xcc logic is correct only for GFX9. + max_scratch_waves = self.cu_cnt * self.iface.props['max_slots_scratch_cu'] * self.xccs + wave_scratch = ceildiv(lanes_per_wave * size_per_thread, mem_alignment_size) + num_waves = (size_per_xcc // (wave_scratch * mem_alignment_size)) // (self.se_cnt if self.target[0] != 9 else 1) + + tmpring_t = getattr(hsa, f'union_COMPUTE_TMPRING_SIZE{"_GFX"+str(self.target[0]) if self.target[0] != 9 else ""}_bitfields') + tmpring = int.from_bytes(tmpring_t(WAVES=min(num_waves, max_scratch_waves), WAVESIZE=wave_scratch), 'little') + + if hasattr(self, 'aql_desc'): + gfx9_rsrc = {'NUM_FORMAT':hsa.BUF_NUM_FORMAT_UINT, 'DATA_FORMAT':hsa.BUF_DATA_FORMAT_32, 'ELEMENT_SIZE':1, 'INDEX_STRIDE':3} + rsrc = {'DST_SEL_X':hsa.SQ_SEL_X, 'DST_SEL_Y':hsa.SQ_SEL_Y, 'DST_SEL_Z':hsa.SQ_SEL_Z, 'DST_SEL_W':hsa.SQ_SEL_W, 'ADD_TID_ENABLE':1, + 'TYPE':hsa.SQ_RSRC_BUF, **(gfx9_rsrc if self.target[0] == 9 else {'FORMAT':hsa.BUF_FORMAT_32_UINT, 'OOB_SELECT':2})} + rsrc1_t = getattr(hsa, f'union_SQ_BUF_RSRC_WORD1{"_GFX11" if self.target[0] != 9 else ""}_bitfields') + rsrc3_t = getattr(hsa, f'union_SQ_BUF_RSRC_WORD3{"_GFX"+str(self.target[0]) if self.target[0] != 9 else ""}_bitfields') + + self.aql_desc.scratch_backing_memory_location = int(self.scratch.get_buf().va_addr) + self.aql_desc.scratch_wave64_lane_byte_size = self.max_private_segment_size * lanes_per_wave // 64 + self.aql_desc.scratch_resource_descriptor[:] = [lo32(self.scratch.get_buf().va_addr), + int.from_bytes(rsrc1_t(BASE_ADDRESS_HI=hi32(self.scratch.get_buf().va_addr), SWIZZLE_ENABLE=1), 'little'), + lo32(size_per_xcc), int.from_bytes(bytes(rsrc3_t(**rsrc)), 'little')] + self.aql_desc.compute_tmpring_size = tmpring + self.aql_gart.cpu_view()[:ctypes.sizeof(self.aql_desc)] = bytes(self.aql_desc) + + return tmpring + + def scratch_buffer(self, private_segment_size): + private_segment_size = max(private_segment_size, 128) + 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 + size_per_thread = round_up(private_segment_size, mem_alignment_size // lanes_per_wave) + size_per_xcc = size_per_thread * lanes_per_wave * self.iface.props['max_slots_scratch_cu'] * self.cu_cnt + self.scratch = Buffer(self.device, size_per_xcc * self.xccs, dtypes.uint8, options=BufferSpec(nolru=True), preallocate=True) + self.max_private_segment_size = private_segment_size + return self.scratch def on_device_hang(self): self.iface.on_device_hang() diff --git a/extra/llama_kernels/__init__.py b/extra/llama_kernels/__init__.py index 3f18b3a5519bd..c033bd9b0252b 100644 --- a/extra/llama_kernels/__init__.py +++ b/extra/llama_kernels/__init__.py @@ -1,7 +1,8 @@ from __future__ import annotations import functools, pathlib +from dataclasses import replace from tinygrad import Tensor, dtypes -from tinygrad.uop.ops import Ops +from tinygrad.uop.ops import shape_to_shape_arg from tinygrad.runtime.support.compiler_amd import HIPCCCompiler FP8_MAX = 448.0 @@ -11,7 +12,7 @@ @functools.cache def _local_abs_max_fxn(x_p, device): x = Tensor(x_p, device=device) - inner = Tensor(x.uop.src[0]) if x.uop.op is Ops.MULTI else x + inner = Tensor(x.uop.replace(src=(shape_to_shape_arg(x.uop.shard_shape),), arg=replace(x.uop.arg, axis=None))) if x.uop.axis is not None else x return (inner.abs().max(),) def local_abs_max(x:Tensor) -> Tensor: diff --git a/extra/llama_kernels/cast_amax/__init__.py b/extra/llama_kernels/cast_amax/__init__.py index 45161bc8a1e8f..567e8987ccc5b 100644 --- a/extra/llama_kernels/cast_amax/__init__.py +++ b/extra/llama_kernels/cast_amax/__init__.py @@ -5,19 +5,19 @@ from tinygrad.renderer import Estimates from extra.llama_kernels import FP8_MAX, NUM_WG, THREADS_PER_WG, compile_cpp, alloc_like, alloc_local, scalar_amax, dname_of -# module-level mailbox: grad_xw13 UOp -> (grad_xw13_fp8 UOp, inv_scale UOp, new_amax UOp, store_effect) +# module-level mailbox: grad_xw13 UOp -> (grad_xw13_fp8 UOp, inv_scale UOp) # lets cdna_asm_gemm's bwd reuse the fp8 companion produced by the fused silu_mul bwd kernel # instead of doing a redundant bf16 -> fp8 quantize. -_grad_fp8_mailbox:dict = {} +_grad_fp8_mailbox:dict[UOp, tuple[UOp, UOp]] = {} @functools.cache -def _custom_fused_bwd_w13(grad_xw13:UOp, grad_xw13_fp8:UOp, grad_amax_buf:UOp, +def _custom_fused_bwd_w13(grad_xw13_fp8:UOp, grad_amax_buf:UOp, xw13:UOp, grad_x2:UOp, amax_state:UOp, grad_amax_state:UOp, dname:str) -> UOp: hidden = xw13.shape[2] // 2 n_elems = xw13.shape[0] * xw13.shape[1] * hidden threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0") - mem = n_elems * 2 * 5 + n_elems * 2 + NUM_WG * 4 + 4 - sink = UOp.sink(grad_xw13.base, grad_xw13_fp8.base, grad_amax_buf.base, + mem = n_elems * 2 * 3 + n_elems * 2 + NUM_WG * 4 + 4 + sink = UOp.sink(grad_xw13_fp8.base, grad_amax_buf.base, xw13.base, grad_x2.base, amax_state.base, grad_amax_state.base, threads, workgroups, arg=KernelInfo(f"fused_silu_mul_bwd_w13_{n_elems}", estimates=Estimates(ops=10*n_elems, mem=mem))) src, lib = compile_cpp(pathlib.Path(__file__).parent, "cast_amax_bwd_w13.cpp", n_elems, hidden) @@ -41,23 +41,23 @@ def _fused_quantize_bwd_w13(gradient:UOp, kernel:UOp): _, _, xw13, amax_state, grad_amax_state = kernel.src[1:] device = xw13.device axis = xw13.axis if isinstance(device, tuple) else None - grad_xw13 = alloc_like(xw13.shape, dtypes.bfloat16, device, axis) grad_xw13_fp8 = alloc_like(xw13.shape, dtypes.fp8e4m3, device, axis) grad_amax_buf = alloc_local((NUM_WG,), dtypes.float32, device, axis) grad_amax_state_t = Tensor(grad_amax_state, device=device) fxn = functools.partial(_custom_fused_bwd_w13, dname=dname_of(device)) - grad_xw13, grad_xw13_fp8, grad_amax_buf, *_ = Tensor.custom_kernel( - grad_xw13, grad_xw13_fp8, grad_amax_buf, + grad_xw13_fp8, grad_amax_buf, *_ = Tensor.custom_kernel( + grad_xw13_fp8, grad_amax_buf, Tensor(xw13, device=device), Tensor(gradient, device=device).cast(dtypes.bfloat16), Tensor(amax_state, device=device), grad_amax_state_t, fxn=fxn) + grad_xw13_uop = grad_xw13_fp8.uop.cast(dtypes.bfloat16) inv_scale = (grad_amax_state_t.float() + 1e-8) / FP8_MAX new_grad_amax = scalar_amax(grad_amax_buf) store_effect = grad_amax_state_t.uop.store(new_grad_amax.uop) assert grad_xw13_fp8.uop.op is Ops.AFTER, f"expected AFTER, got {grad_xw13_fp8.uop.op}" grad_xw13_fp8_uop = grad_xw13_fp8.uop.replace(src=grad_xw13_fp8.uop.src + (store_effect,)) # Stash fp8 companion for cdna_asm_gemm's bwd to attach to grad_a. - _grad_fp8_mailbox[grad_xw13.uop] = (grad_xw13_fp8_uop, inv_scale.uop) - return (None, None, grad_xw13.uop, None, None) + _grad_fp8_mailbox[grad_xw13_uop] = (grad_xw13_fp8_uop, inv_scale.uop) + return (None, None, grad_xw13_uop, None, None) def fused_quantize_fp8_w13(xw13:Tensor, amax_state:Tensor, fp8_dtype, grad_amax_state:Tensor) -> tuple[Tensor, Tensor, Tensor]: # NOTE: silu(xw1)*xw3 -> fp8 + amax over fused xw13 layout. Returns (fp8, inv_scale, new_amax) diff --git a/extra/llama_kernels/cast_amax/cast_amax_bwd_w13.cpp b/extra/llama_kernels/cast_amax/cast_amax_bwd_w13.cpp index 93c2d8f0c6cb2..c787fea4b05b3 100644 --- a/extra/llama_kernels/cast_amax/cast_amax_bwd_w13.cpp +++ b/extra/llama_kernels/cast_amax/cast_amax_bwd_w13.cpp @@ -21,15 +21,13 @@ constexpr float FP8_MAX = 448.0f; static_assert(N_ELEMS % VEC == 0, "N_ELEMS must be divisible by VEC"); static_assert(HIDDEN % VEC == 0, "HIDDEN must be divisible by VEC"); -// fused silu*mul backward, three outputs in a single HBM pass: -// 1) bf16 grad_xw13 — consumed by downstream bf16 autograd chain -// 2) fp8 grad_xw13_fp8 — delayed-scale quantize using grad_amax_state (mailbox to matmul bwd) -// 3) fp32 grad_amax_buf — per-WG partial |grad_xw13|, reduced into next step's grad_amax_state +// fused silu*mul backward, two outputs in a single HBM pass: +// 1) fp8 grad_xw13_fp8 — delayed-scale quantize using grad_amax_state (mailbox to matmul bwd) +// 2) fp32 grad_amax_buf — per-WG partial |grad_xw13|, reduced into next step's grad_amax_state // grad_amax_state is read for the fp8 scale. The store of new_grad_amax into grad_amax_state's // buffer is built in Python as a separate effect and threaded into grad_a via .after(store). extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void fused_silu_mul_bwd_w13( - __hip_bfloat16* __restrict__ grad_xw13_out, // bf16, 2*N_ELEMS __hip_fp8_storage_t* __restrict__ grad_xw13_fp8_out, // fp8, 2*N_ELEMS float* __restrict__ grad_amax_buf, // fp32, NUM_WG per-WG partials const __hip_bfloat16* __restrict__ xw13, // bf16, 2*N_ELEMS @@ -62,7 +60,6 @@ fused_silu_mul_bwd_w13( const __hip_bfloat16 *x3 = reinterpret_cast(&x3_raw); const __hip_bfloat16 *gv = reinterpret_cast(&g_raw); - __hip_bfloat16 out1[VEC], out3[VEC]; __hip_fp8_storage_t fp8_1[VEC], fp8_3[VEC]; #pragma unroll for (int i = 0; i < VEC; i++) { @@ -75,15 +72,11 @@ fused_silu_mul_bwd_w13( const float gs = fg * scale; const float g1 = gs * silu_prime * f3; const float g3 = gs * silu; - out1[i] = static_cast<__hip_bfloat16>(g1); - out3[i] = static_cast<__hip_bfloat16>(g3); local_max = fmaxf(local_max, fmaxf(fabsf(g1), fabsf(g3))); fp8_1[i] = __hip_cvt_float_to_fp8(fmaxf(-FP8_MAX, fminf(FP8_MAX, g1 * g_scale)), __HIP_SATFINITE, __HIP_E4M3); fp8_3[i] = __hip_cvt_float_to_fp8(fmaxf(-FP8_MAX, fminf(FP8_MAX, g3 * g_scale)), __HIP_SATFINITE, __HIP_E4M3); } - *reinterpret_cast(&grad_xw13_out[xw1_off]) = *reinterpret_cast(out1); - *reinterpret_cast(&grad_xw13_out[xw3_off]) = *reinterpret_cast(out3); *reinterpret_cast(&grad_xw13_fp8_out[xw1_off]) = *reinterpret_cast(fp8_1); *reinterpret_cast(&grad_xw13_fp8_out[xw3_off]) = *reinterpret_cast(fp8_3); } diff --git a/extra/llama_kernels/fused_ce/__init__.py b/extra/llama_kernels/fused_ce/__init__.py index 76e6de8e98286..3aba728a28c7c 100644 --- a/extra/llama_kernels/fused_ce/__init__.py +++ b/extra/llama_kernels/fused_ce/__init__.py @@ -1,41 +1,38 @@ -from __future__ import annotations -import functools, pathlib +import functools from tinygrad import Tensor, dtypes -from tinygrad.uop.ops import UOp, Ops, KernelInfo -from tinygrad.renderer import Estimates -from tinygrad.runtime.support.compiler_amd import HIPCCCompiler - -THREADS_PER_WG = 256 +from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType @functools.cache def _custom_fused_ce_loss_fwd(loss_out:UOp, max_out:UOp, lse_out:UOp, logits:UOp, targets:UOp, - dname:str, vocab:int, rows:int, label_smoothing:float) -> UOp: - threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(rows, "gidx0") - mem = rows * vocab * 2 + rows * 12 + rows * 4 - sink = UOp.sink(loss_out.base, max_out.base, lse_out.base, logits.base, targets.base, - threads, workgroups, - arg=KernelInfo(f"fused_ce_loss_fwd", estimates=Estimates(ops=6*rows*vocab, mem=mem))) - src = (pathlib.Path(__file__).parent/"fused_ce_loss.cpp").read_text() - defines = [f"-DVOCAB={vocab}", f"-DTHREADS_PER_WG={THREADS_PER_WG}", - f"-DLABEL_SMOOTHING={label_smoothing}f"] - lib = HIPCCCompiler("gfx950", ["-std=c++20", "-ffast-math", *defines]).compile_cached(src) - return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), - UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib))) + vocab:int, rows:int, label_smoothing:float) -> UOp: + row = UOp.range(rows, 0) + + v_max = UOp.range(vocab, 1, axis_type=AxisType.REDUCE) + row_max = logits[row, v_max].cast(dtypes.float).reduce(v_max, arg=Ops.MAX) + + v_lse = UOp.range(vocab, 2, axis_type=AxisType.REDUCE) + row_lse = (logits[row, v_lse].cast(dtypes.float) - row_max).exp().reduce(v_lse, arg=Ops.ADD).log() + row_max + + v_smooth = UOp.range(vocab, 3, axis_type=AxisType.REDUCE) + target = logits[row, targets[row].cast(dtypes.weakint)].cast(dtypes.float) + mean_logits = logits[row, v_smooth].cast(dtypes.float).reduce(v_smooth, arg=Ops.ADD) / vocab + loss = row_lse - (1.0 - label_smoothing) * target - label_smoothing * mean_logits + stores = UOp.group(loss_out[row].store(loss), max_out[row].store(row_max), lse_out[row].store(row_lse)) + + return stores.end(row).sink(arg=KernelInfo(f"fused_ce_loss_fwd_{rows}_{vocab}")) @functools.cache def _custom_fused_ce_loss_bwd(d_logits:UOp, logits:UOp, lse:UOp, targets:UOp, scale:UOp, - dname:str, vocab:int, rows:int, label_smoothing:float) -> UOp: - threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(rows, "gidx0") - mem = rows * vocab * 4 + rows * 8 + 4 - sink = UOp.sink(d_logits.base, logits.base, lse.base, targets.base, scale.base, - threads, workgroups, - arg=KernelInfo(f"fused_ce_loss_bwd", estimates=Estimates(ops=4*rows*vocab, mem=mem))) - src = (pathlib.Path(__file__).parent/"fused_ce_loss_bwd.cpp").read_text() - defines = [f"-DVOCAB={vocab}", f"-DTHREADS_PER_WG={THREADS_PER_WG}", - f"-DLABEL_SMOOTHING={label_smoothing}f"] - lib = HIPCCCompiler("gfx950", ["-std=c++20", "-ffast-math", *defines]).compile_cached(src) - return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), - UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib))) + vocab:int, rows:int, label_smoothing:float) -> UOp: + row = UOp.range(rows, 0) + v = UOp.range(vocab, 1) + + prob = (logits[row, v].cast(dtypes.float) - lse[row]).exp() + target = v.eq(targets[row].cast(dtypes.weakint)).where(1.0 - label_smoothing, 0.0) + smooth = label_smoothing / vocab + grad = (prob - target - smooth) * scale[0] + + return d_logits[row, v].store(grad.cast(d_logits.dtype.base)).end(v, row).sink(arg=KernelInfo(f"fused_ce_loss_bwd_{rows}_{vocab}")) def _fused_ce_loss_bwd(gradient:UOp, kernel:UOp, label_smoothing:float): # NOTE: forward inputs are (loss_out, max_out, lse_out, logits, targets) @@ -47,18 +44,16 @@ def _fused_ce_loss_bwd(gradient:UOp, kernel:UOp, label_smoothing:float): axis = logits_u.axis ndev = len(device) d_logits = Tensor(Tensor.invalids(rows // ndev, VOCAB, dtype=dtypes.bfloat16, device=device).uop.multi(axis), device=device) - dname = device[0].split(":")[0] rows_per_dev = rows // ndev else: d_logits = Tensor.invalids(rows, VOCAB, dtype=dtypes.bfloat16, device=device) - dname = device.split(":")[0] if isinstance(device, str) else device rows_per_dev = rows # NOTE: .mean() backward gives same grad per row (1/N), so broadcast is safe; take scalar scale = Tensor(gradient, device=device).float().reshape(-1)[0:1].contiguous() logits_t = Tensor(logits_u.after(kernel), device=device) lse_t = Tensor(lse_u.after(kernel), device=device) targets_t = Tensor(targets_u, device=device) - fxn = functools.partial(_custom_fused_ce_loss_bwd, dname=dname, vocab=VOCAB, rows=rows_per_dev, label_smoothing=label_smoothing) + fxn = functools.partial(_custom_fused_ce_loss_bwd, vocab=VOCAB, rows=rows_per_dev, label_smoothing=label_smoothing) d_logits, *_ = Tensor.custom_kernel(d_logits, logits_t, lse_t, targets_t, scale, fxn=fxn) return (None, None, None, d_logits.uop, None) @@ -78,17 +73,15 @@ def fused_ce_loss(logits:Tensor, targets:Tensor, label_smoothing:float=0.1) -> T device=logits.device) lse_out = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.multi(0), device=logits.device) - dname = logits.device[0].split(":")[0] rows_per_dev = rows // ndev else: loss_out = Tensor.invalids(rows, dtype=dtypes.float32, device=logits.device) max_out = Tensor.invalids(rows, dtype=dtypes.float32, device=logits.device) lse_out = Tensor.invalids(rows, dtype=dtypes.float32, device=logits.device) - dname = logits.device.split(":")[0] if isinstance(logits.device, str) else logits.device rows_per_dev = rows logits_flat = logits.reshape(rows, VOCAB) targets_flat = targets.reshape(-1).cast(dtypes.int32) - fxn = functools.partial(_custom_fused_ce_loss_fwd, dname=dname, vocab=VOCAB, rows=rows_per_dev, + fxn = functools.partial(_custom_fused_ce_loss_fwd, vocab=VOCAB, rows=rows_per_dev, label_smoothing=label_smoothing) loss_out, max_out, lse_out, *_ = Tensor.custom_kernel( loss_out, max_out, lse_out, logits_flat, targets_flat, diff --git a/extra/llama_kernels/fused_ce/fused_ce_loss.cpp b/extra/llama_kernels/fused_ce/fused_ce_loss.cpp deleted file mode 100644 index cdbd2d630a49f..0000000000000 --- a/extra/llama_kernels/fused_ce/fused_ce_loss.cpp +++ /dev/null @@ -1,104 +0,0 @@ -#include -#include - -// Fused forward sparse-CE with label smoothing. -// SINGLE-PASS online softmax + vectorized 8-wide bf16 loads for HBM coalescing. - -#ifndef VOCAB -#define VOCAB 128256 -#endif -#ifndef THREADS_PER_WG -#define THREADS_PER_WG 256 -#endif -#ifndef LABEL_SMOOTHING -#define LABEL_SMOOTHING 0.1f -#endif - -constexpr int VEC = 8; - -extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void -fused_ce_loss_fwd( - float* __restrict__ loss_out, // out: fp32, ROWS - float* __restrict__ max_out, // out: fp32, ROWS - float* __restrict__ lse_out, // out: fp32, ROWS - const __hip_bfloat16* __restrict__ logits, // in: bf16, ROWS*VOCAB - const int* __restrict__ targets) // in: int32, ROWS -{ - __shared__ float sdata_m[THREADS_PER_WG]; - __shared__ float sdata_s[THREADS_PER_WG]; - __shared__ float sdata_sumx[THREADS_PER_WG]; - __shared__ float sdata_tgt[THREADS_PER_WG]; - - const int tid = threadIdx.x; - const int row = blockIdx.x; - const int target = targets[row]; - const __hip_bfloat16* row_logits = logits + (size_t)row * VOCAB; - - float m = -INFINITY; - float s = 0.0f; - float sum_x = 0.0f; - float target_logit = 0.0f; - constexpr bool needs_sum_x = (LABEL_SMOOTHING != 0.0f); - - // Vectorized stride: each iter loads 8 bf16 = 16 bytes. Warp loads 32*16 = 512 bytes (4 cache lines). - const int VOCAB_VEC = VOCAB & ~(VEC - 1); // round down to multiple of VEC - for (int i = tid * VEC; i < VOCAB_VEC; i += THREADS_PER_WG * VEC) { - float4 raw = *reinterpret_cast(&row_logits[i]); - const __hip_bfloat16* xi = reinterpret_cast(&raw); - #pragma unroll - for (int k = 0; k < VEC; k++) { - const float x = static_cast(xi[k]); - if constexpr (needs_sum_x) sum_x += x; - if (i + k == target) target_logit = x; - if (x > m) { - s = s * __expf(m - x) + 1.0f; - m = x; - } else { - s += __expf(x - m); - } - } - } - // tail (VOCAB not divisible by VEC): - for (int i = VOCAB_VEC + tid; i < VOCAB; i += THREADS_PER_WG) { - const float x = static_cast(row_logits[i]); - if constexpr (needs_sum_x) sum_x += x; - if (i == target) target_logit = x; - if (x > m) { s = s * __expf(m - x) + 1.0f; m = x; } - else { s += __expf(x - m); } - } - - sdata_m[tid] = m; - sdata_s[tid] = s; - sdata_sumx[tid] = sum_x; - sdata_tgt[tid] = target_logit; - __syncthreads(); - - for (int step = THREADS_PER_WG / 2; step > 0; step >>= 1) { - if (tid < step) { - const float m1 = sdata_m[tid]; - const float m2 = sdata_m[tid + step]; - const float s1 = sdata_s[tid]; - const float s2 = sdata_s[tid + step]; - const float m_new = fmaxf(m1, m2); - const float s_new = s1 * __expf(m1 - m_new) + s2 * __expf(m2 - m_new); - sdata_m[tid] = m_new; - sdata_s[tid] = s_new; - sdata_sumx[tid] += sdata_sumx[tid + step]; - sdata_tgt[tid] += sdata_tgt[tid + step]; - } - __syncthreads(); - } - - if (tid == 0) { - const float row_max = sdata_m[0]; - const float row_sum_exp = sdata_s[0]; - const float row_sum_x = sdata_sumx[0]; - const float tgt = sdata_tgt[0]; - const float row_lse = logf(row_sum_exp) + row_max; - const float mean_logits = row_sum_x / static_cast(VOCAB); - const float loss = row_lse - (1.0f - LABEL_SMOOTHING) * tgt - LABEL_SMOOTHING * mean_logits; - loss_out[row] = loss; - max_out[row] = row_max; - lse_out[row] = row_lse; - } -} diff --git a/extra/llama_kernels/fused_ce/fused_ce_loss_bwd.cpp b/extra/llama_kernels/fused_ce/fused_ce_loss_bwd.cpp deleted file mode 100644 index 70e70cb7260a0..0000000000000 --- a/extra/llama_kernels/fused_ce/fused_ce_loss_bwd.cpp +++ /dev/null @@ -1,58 +0,0 @@ -#include -#include - -// Vectorized CE bwd: 8-wide bf16 loads + stores. - -#ifndef VOCAB -#define VOCAB 128256 -#endif -#ifndef THREADS_PER_WG -#define THREADS_PER_WG 256 -#endif -#ifndef LABEL_SMOOTHING -#define LABEL_SMOOTHING 0.1f -#endif - -constexpr int VEC = 8; - -extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void -fused_ce_loss_bwd( - __hip_bfloat16* __restrict__ d_logits, - const __hip_bfloat16* __restrict__ logits, - const float* __restrict__ lse, - const int* __restrict__ targets, - const float* __restrict__ scale_in) -{ - const int tid = threadIdx.x; - const int row = blockIdx.x; - const int target = targets[row]; - const float lse_r = lse[row]; - const __hip_bfloat16* row_logits = logits + (size_t)row * VOCAB; - __hip_bfloat16* row_dlogits = d_logits + (size_t)row * VOCAB; - const float inv_vocab = 1.0f / static_cast(VOCAB); - const float scale = *scale_in; - const float ls_term = LABEL_SMOOTHING * inv_vocab; - - const int VOCAB_VEC = VOCAB & ~(VEC - 1); - for (int i = tid * VEC; i < VOCAB_VEC; i += THREADS_PER_WG * VEC) { - float4 raw = *reinterpret_cast(&row_logits[i]); - const __hip_bfloat16* xi = reinterpret_cast(&raw); - __hip_bfloat16 out[VEC]; - #pragma unroll - for (int k = 0; k < VEC; k++) { - const float x = static_cast(xi[k]); - float g = __expf(x - lse_r); - if (i + k == target) g -= (1.0f - LABEL_SMOOTHING); - g -= ls_term; - out[k] = static_cast<__hip_bfloat16>(g * scale); - } - *reinterpret_cast(&row_dlogits[i]) = *reinterpret_cast(out); - } - for (int i = VOCAB_VEC + tid; i < VOCAB; i += THREADS_PER_WG) { - const float x = static_cast(row_logits[i]); - float g = __expf(x - lse_r); - if (i == target) g -= (1.0f - LABEL_SMOOTHING); - g -= ls_term; - row_dlogits[i] = static_cast<__hip_bfloat16>(g * scale); - } -} diff --git a/extra/llama_kernels/fused_pad_grad_accum/__init__.py b/extra/llama_kernels/fused_pad_grad_accum/__init__.py deleted file mode 100644 index 897d7ddee3d61..0000000000000 --- a/extra/llama_kernels/fused_pad_grad_accum/__init__.py +++ /dev/null @@ -1,55 +0,0 @@ -from __future__ import annotations -import functools, pathlib -from tinygrad import Tensor, dtypes -from tinygrad.uop.ops import UOp, Ops, KernelInfo -from tinygrad.renderer import Estimates -from extra.llama_kernels import THREADS_PER_WG, dname_of, compile_hip - -ELEMS_PER_THREAD = 8 # vectorized 16-byte load (uint4 = 8 bf16) - -def _build_src(n_chunks:int) -> str: - template = (pathlib.Path(__file__).parent/"fused_pad_grad_accum.cpp").read_text() - params = "".join(f",\n const __hip_bfloat16* __restrict__ chunk{i}" for i in range(n_chunks)) - dispatch = "\n ".join(f"case {i}: chunk_ptr = chunk{i}; break;" for i in range(n_chunks)) - return (template.replace("__FUSED_PAD_GRAD_ACCUM_PARAMS", params) - .replace("__FUSED_PAD_GRAD_ACCUM_DISPATCH", dispatch)) - -@functools.cache -def _custom_fused_pad_grad_accum(grad_buf:UOp, *chunk_uops, dname:str, n_chunks:int, chunk_size:int) -> UOp: - total = n_chunks * chunk_size - elems_per_block = THREADS_PER_WG * ELEMS_PER_THREAD - assert chunk_size % elems_per_block == 0, f"chunk_size {chunk_size} must be multiple of {elems_per_block}" - num_wg = total // elems_per_block - threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(num_wg, "gidx0") - mem = total * 2 * 3 - sink = UOp.sink(grad_buf.base, *(c.base for c in chunk_uops), threads, workgroups, - arg=KernelInfo(f"fused_pad_grad_accum_n{n_chunks}_c{chunk_size}", - estimates=Estimates(ops=2*total, mem=mem))) - src = _build_src(n_chunks) - defines = [f"-DCHUNK_SIZE={chunk_size}", f"-DTHREADS_PER_WG={THREADS_PER_WG}", f"-DELEMS_PER_THREAD={ELEMS_PER_THREAD}"] - return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), - UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines)))) - -def can_fused_pad_grad_accum(grad_buf:Tensor, chunks:list[Tensor]) -> bool: - if not chunks or grad_buf.dtype != dtypes.bfloat16: return False - if any(c.dtype != dtypes.bfloat16 for c in chunks): return False - chunk_shape = chunks[0].shape - if any(c.shape != chunk_shape for c in chunks): return False - chunk_size, total = 1, 1 - for d in chunk_shape: chunk_size *= d - for d in grad_buf.shape: total *= d - return total == len(chunks) * chunk_size and chunk_size % (THREADS_PER_WG * ELEMS_PER_THREAD) == 0 - -def fused_pad_grad_accum(grad_buf:Tensor, chunks:list[Tensor]) -> Tensor: - # NOTE: grad_buf += cat(*chunks, dim=0) in one HBM pass (in-place add). Returns new grad_buf Tensor. - # Requires uniform chunk shapes and chunk_size % (THREADS_PER_WG*ELEMS_PER_THREAD) == 0. - assert chunks and grad_buf.dtype == dtypes.bfloat16 - for c in chunks: assert c.dtype == dtypes.bfloat16, f"chunk dtype must be bf16, got {c.dtype}" - chunk_size, total = 1, 1 - for d in chunks[0].shape: chunk_size *= d - for d in grad_buf.shape: total *= d - assert total == len(chunks) * chunk_size, f"grad_buf size {total} != n_chunks {len(chunks)} * chunk_size {chunk_size}" - fxn = functools.partial(_custom_fused_pad_grad_accum, dname=dname_of(grad_buf.device), - n_chunks=len(chunks), chunk_size=chunk_size) - out, *_ = Tensor.custom_kernel(grad_buf, *chunks, fxn=fxn) - return out diff --git a/extra/llama_kernels/fused_pad_grad_accum/fused_pad_grad_accum.cpp b/extra/llama_kernels/fused_pad_grad_accum/fused_pad_grad_accum.cpp deleted file mode 100644 index d4df9c8b33a49..0000000000000 --- a/extra/llama_kernels/fused_pad_grad_accum/fused_pad_grad_accum.cpp +++ /dev/null @@ -1,63 +0,0 @@ -// Fused custom kernel: grad_buf += cat(*chunks, dim=0) in one HBM pass. -// -// Template source — chunk parameter list and switch dispatch are filled by codegen -// in cast_amax.py:_build_fused_pad_grad_accum_src to support arbitrary N. -// -// Defines required at compile time: -// CHUNK_SIZE elements per chunk (must be multiple of THREADS_PER_WG * ELEMS_PER_THREAD) -// THREADS_PER_WG -// ELEMS_PER_THREAD (8 = one uint4 per thread = 16-byte vectorized load) -// -// Layout: one block-per-(slice-of-chunk) — blockIdx.x / BLOCKS_PER_CHUNK selects the chunk. -// All threads in a block read the same chunk → switch is uniform → no warp divergence. - -#include -#include - -#ifndef THREADS_PER_WG -#define THREADS_PER_WG 256 -#endif -#ifndef ELEMS_PER_THREAD -#define ELEMS_PER_THREAD 8 -#endif - -#define ELEMS_PER_BLOCK (THREADS_PER_WG * ELEMS_PER_THREAD) -#define BLOCKS_PER_CHUNK (CHUNK_SIZE / ELEMS_PER_BLOCK) - -extern "C" __attribute__((global)) -__attribute__((amdgpu_flat_work_group_size(1, THREADS_PER_WG))) -void fused_pad_grad_accum( - __hip_bfloat16* __restrict__ grad_buf - __FUSED_PAD_GRAD_ACCUM_PARAMS -) { - const int bid = blockIdx.x; - const int chunk_idx = bid / BLOCKS_PER_CHUNK; - const int block_in_chunk = bid - chunk_idx * BLOCKS_PER_CHUNK; - const int tid = threadIdx.x; - - const __hip_bfloat16* chunk_ptr; - switch (chunk_idx) { - __FUSED_PAD_GRAD_ACCUM_DISPATCH - default: chunk_ptr = (const __hip_bfloat16*)0; break; // unreachable - } - - // int64 for global_offset: at 32 chunks × 117M elements = 3.6B, int32 overflows → MEMVIOL. - const int local_offset = block_in_chunk * ELEMS_PER_BLOCK + tid * ELEMS_PER_THREAD; - const long long global_offset = (long long)chunk_idx * (long long)CHUNK_SIZE + (long long)local_offset; - - // Vectorized 16-byte load (uint4 = 8 bf16). Requires CHUNK_SIZE % 8 == 0 and 16-byte alignment. - const uint4 chunk_v = *reinterpret_cast(&chunk_ptr[local_offset]); - const uint4 grad_v = *reinterpret_cast(&grad_buf[global_offset]); - uint4 out_v; - - const __hip_bfloat16* chunk_bf = reinterpret_cast(&chunk_v); - const __hip_bfloat16* grad_bf = reinterpret_cast(&grad_v); - __hip_bfloat16* out_bf = reinterpret_cast<__hip_bfloat16*>(&out_v); - - #pragma unroll - for (int i = 0; i < ELEMS_PER_THREAD; i++) { - out_bf[i] = (__hip_bfloat16)((float)grad_bf[i] + (float)chunk_bf[i]); - } - - *reinterpret_cast(&grad_buf[global_offset]) = out_v; -} diff --git a/extra/llama_kernels/quantize_fp8_delayed/__init__.py b/extra/llama_kernels/quantize_fp8_delayed/__init__.py index 60d7060347056..f51f2b49b5e85 100644 --- a/extra/llama_kernels/quantize_fp8_delayed/__init__.py +++ b/extra/llama_kernels/quantize_fp8_delayed/__init__.py @@ -1,35 +1,64 @@ -from __future__ import annotations -import functools, pathlib +import functools from tinygrad import Tensor, dtypes -from tinygrad.uop.ops import UOp, Ops, KernelInfo -from tinygrad.renderer import Estimates -from extra.llama_kernels import FP8_MAX, NUM_WG, THREADS_PER_WG, alloc_like, alloc_local, scalar_amax, dname_of, compile_hip +from tinygrad.dtype import AddrSpace +from tinygrad.helpers import prod +from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType +from extra.llama_kernels import FP8_MAX, NUM_WG, THREADS_PER_WG, alloc_like, alloc_local, scalar_amax @functools.cache -def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_partial:UOp, x:UOp, amax_state:UOp, dname:str) -> UOp: - n_elems = 1 - for d in x.shape: n_elems *= d - threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0") - mem = n_elems * 2 + n_elems + 4 + NUM_WG * 4 - sink = UOp.sink(fp8_out.base, amax_partial.base, x.base, amax_state.base, threads, workgroups, - arg=KernelInfo(f"quantize_fp8_with_amax_{n_elems}", estimates=Estimates(ops=3*n_elems, mem=mem))) - src = (pathlib.Path(__file__).parent/"quantize_fp8_with_amax.cpp").read_text() - defines = [f"-DN_ELEMS={n_elems}", f"-DNUM_WG={NUM_WG}", f"-DTHREADS_PER_WG={THREADS_PER_WG}"] - return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), - UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines)))) +def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_partial:UOp, x:UOp, amax_state:UOp) -> UOp: + VEC = 8 + n_elems = prod(x.shape) + assert n_elems % (NUM_WG * THREADS_PER_WG * VEC) == 0 + assert amax_partial.shape[0] == NUM_WG + + x = x.reshape(n_elems) + fp8_out = fp8_out.reshape(n_elems) + + wg = UOp.range(NUM_WG, 0, AxisType.GLOBAL) + tid = UOp.range(THREADS_PER_WG, 1, AxisType.LOCAL) + it = UOp.range((n_elems // VEC) // (NUM_WG * THREADS_PER_WG), 2, AxisType.LOOP) + lane = UOp.range(VEC, 3, AxisType.UNROLL) + + idx = (((it * NUM_WG + wg) * THREADS_PER_WG + tid) * VEC) + lane + + scale = FP8_MAX / (amax_state[0].cast(dtypes.float) + 1e-8) + x_f = x[idx].cast(dtypes.float) + abs_x = (x_f < 0.0).where(-x_f, x_f) + scaled = (x_f * scale).maximum(-FP8_MAX).minimum(FP8_MAX) + + fp8_store = fp8_out[idx].store(scaled.cast(fp8_out.dtype.base)).end(lane) + lane_max = abs_x.reduce(lane, arg=Ops.MAX) + + lmax = UOp.placeholder((1,), dtypes.float, slot=1, addrspace=AddrSpace.REG) + lmax_init = lmax.after(wg, tid)[0].store(0.0) + lmax_prev = lmax.after(lmax_init, it)[0] + lmax_store = lmax.after(fp8_store)[0].store(lmax_prev.maximum(lane_max)) + lmax_val = lmax.after(lmax_store.end(it))[0] + + lds = UOp.placeholder((THREADS_PER_WG,), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL) + lds = lds.after(lds[tid].store(lmax_val).barrier()) + + step = THREADS_PER_WG // 2 + while step: + active = tid < step + other = lds[tid + step].load(UOp.const(dtypes.float, 0.0), active) + lds = lds.after(lds[tid].store(lds[tid].maximum(other), gate=active).barrier()) + step //= 2 + + amax_store = amax_partial[tid.eq(0).where(wg, UOp.invalid())].store(lds[0]) + return amax_store.end(tid, wg).sink(arg=KernelInfo(f"quantize_fp8_with_amax_{n_elems}", opts_to_apply=())) @functools.cache -def _custom_quantize_fp8_scalar(fp8_out:UOp, x:UOp, amax_state:UOp, dname:str) -> UOp: - n_elems = 1 - for d in x.shape: n_elems *= d - threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0") - mem = n_elems * 2 + n_elems - sink = UOp.sink(fp8_out.base, x.base, amax_state.base, threads, workgroups, - arg=KernelInfo(f"quantize_fp8_scalar_{n_elems}", estimates=Estimates(ops=2*n_elems, mem=mem))) - src = (pathlib.Path(__file__).parent/"quantize_fp8_scalar.cpp").read_text() - defines = [f"-DN_ELEMS={n_elems}", f"-DNUM_WG={NUM_WG}", f"-DTHREADS_PER_WG={THREADS_PER_WG}"] - return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), - UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines)))) +def _custom_quantize_fp8_scalar(fp8_out:UOp, x:UOp, amax_state:UOp) -> UOp: + n_elems = prod(x.shape) + i = UOp.range(n_elems, 0) + + x_f = x.reshape(n_elems)[i].cast(dtypes.float) + scale = FP8_MAX / (amax_state[0].cast(dtypes.float) + 1e-8) + store = fp8_out.reshape(n_elems)[i].store((x_f * scale).cast(fp8_out.dtype.base)) + + return store.end(i).sink(arg=KernelInfo(f"quantize_fp8_scalar_{n_elems}")) def _quantize_fp8_delayed_bwd(gradient:UOp, kernel:UOp): # NOTE: STE-equivalent backward — grad_x = grad_fp8 * scale, scale = FP8_MAX / amax_state. @@ -49,8 +78,10 @@ def quantize_fp8_delayed(x:Tensor, amax_state:Tensor, fp8_dtype=dtypes.fp8e4m3) assert x.dtype == dtypes.bfloat16, f"expected bf16, got {x.dtype}" axis = x.uop.axis if isinstance(x.device, tuple) else None fp8_out = alloc_like(x.shape, fp8_dtype, x.device, axis) + n_elems = prod(x.uop.shard_shape) + assert n_elems % NUM_WG == 0, f"{n_elems=} must divide over {NUM_WG=}" amax_partial = alloc_local((NUM_WG,), dtypes.float32, x.device, axis) - fxn = functools.partial(_custom_quantize_fp8_with_amax, dname=dname_of(x.device)) + fxn = _custom_quantize_fp8_with_amax fp8_out, amax_partial, *_ = Tensor.custom_kernel(fp8_out, amax_partial, x, amax_state, fxn=fxn, grad_fxn=_quantize_fp8_delayed_bwd) new_amax = scalar_amax(amax_partial) @@ -62,6 +93,6 @@ def quantize_fp8_scalar(x:Tensor, amax_state:Tensor, fp8_dtype=dtypes.fp8e4m3) - # NOTE: pure one-pass bf16 -> fp8 quantize with delayed scalar scale. No amax computation. axis = x.uop.axis if isinstance(x.device, tuple) else None fp8_out = alloc_like(x.shape, fp8_dtype, x.device, axis) - fxn = functools.partial(_custom_quantize_fp8_scalar, dname=dname_of(x.device)) + fxn = _custom_quantize_fp8_scalar fp8_out, *_ = Tensor.custom_kernel(fp8_out, x, amax_state, fxn=fxn) return fp8_out diff --git a/extra/llama_kernels/quantize_fp8_delayed/quantize_fp8_scalar.cpp b/extra/llama_kernels/quantize_fp8_delayed/quantize_fp8_scalar.cpp deleted file mode 100644 index 33c1636f6e49d..0000000000000 --- a/extra/llama_kernels/quantize_fp8_delayed/quantize_fp8_scalar.cpp +++ /dev/null @@ -1,48 +0,0 @@ -#include -#include -#include - -// Pure one-pass bf16 -> fp8 quantize with delayed scalar scale. No amax computation. - -#ifndef N_ELEMS -#define N_ELEMS 67108864 -#endif -#ifndef NUM_WG -#define NUM_WG 1024 -#endif -#ifndef THREADS_PER_WG -#define THREADS_PER_WG 256 -#endif - -constexpr int VEC = 8; -constexpr float FP8_MAX = 448.0f; - -static_assert(N_ELEMS % VEC == 0, "N_ELEMS must be divisible by VEC"); - -extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void -quantize_fp8_scalar( - __hip_fp8_storage_t* __restrict__ fp8_out, // fp8, N_ELEMS - const __hip_bfloat16* __restrict__ x, // bf16, N_ELEMS - const float* __restrict__ amax_state) // fp32 scalar (delayed) -{ - const int tid = threadIdx.x; - const int wg = blockIdx.x; - const int gid = wg * THREADS_PER_WG + tid; - const int stride_elems = NUM_WG * THREADS_PER_WG * VEC; - - const float scale = FP8_MAX / (static_cast(*amax_state) + 1e-8f); - - for (int base = gid * VEC; base < N_ELEMS; base += stride_elems) { - float4 x_raw = *reinterpret_cast(&x[base]); - const __hip_bfloat16 *xi = reinterpret_cast(&x_raw); - - __hip_fp8_storage_t out[VEC]; - #pragma unroll - for (int i = 0; i < VEC; i++) { - const float v = static_cast(xi[i]); - const float scaled = fmaxf(-FP8_MAX, fminf(FP8_MAX, v * scale)); - out[i] = __hip_cvt_float_to_fp8(scaled, __HIP_SATFINITE, __HIP_E4M3); - } - *reinterpret_cast(&fp8_out[base]) = *reinterpret_cast(out); - } -} diff --git a/extra/llama_kernels/quantize_fp8_delayed/quantize_fp8_with_amax.cpp b/extra/llama_kernels/quantize_fp8_delayed/quantize_fp8_with_amax.cpp deleted file mode 100644 index 360e4c1af49fd..0000000000000 --- a/extra/llama_kernels/quantize_fp8_delayed/quantize_fp8_with_amax.cpp +++ /dev/null @@ -1,63 +0,0 @@ -#include -#include -#include - -// One-pass bf16 -> fp8 quantize using a scalar delayed amax state, -// AND simultaneously computes per-WG |x| max partials for the next step's amax state. -// Saves one full HBM pass over the grad tensor vs. doing quantize + separate abs().max(). - -#ifndef N_ELEMS -#define N_ELEMS 67108864 -#endif -#ifndef NUM_WG -#define NUM_WG 1024 -#endif -#ifndef THREADS_PER_WG -#define THREADS_PER_WG 256 -#endif - -constexpr int VEC = 8; -constexpr float FP8_MAX = 448.0f; - -static_assert(N_ELEMS % VEC == 0, "N_ELEMS must be divisible by VEC"); - -extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void -quantize_fp8_with_amax( - __hip_fp8_storage_t* __restrict__ fp8_out, // out: fp8, N_ELEMS - float* __restrict__ amax_partial, // out: fp32, NUM_WG per-WG partials - const __hip_bfloat16* __restrict__ x, // in: bf16, N_ELEMS - const float* __restrict__ amax_state) // in: fp32 scalar (delayed) -{ - __shared__ float sdata[THREADS_PER_WG]; - - const int tid = threadIdx.x; - const int wg = blockIdx.x; - const int gid = wg * THREADS_PER_WG + tid; - const int stride_elems = NUM_WG * THREADS_PER_WG * VEC; - - const float scale = FP8_MAX / (static_cast(*amax_state) + 1e-8f); - float local_max = 0.0f; - - for (int base = gid * VEC; base < N_ELEMS; base += stride_elems) { - float4 x_raw = *reinterpret_cast(&x[base]); - const __hip_bfloat16 *xi = reinterpret_cast(&x_raw); - - __hip_fp8_storage_t out[VEC]; - #pragma unroll - for (int i = 0; i < VEC; i++) { - const float v = static_cast(xi[i]); - local_max = fmaxf(local_max, fabsf(v)); - const float scaled = fmaxf(-FP8_MAX, fminf(FP8_MAX, v * scale)); - out[i] = __hip_cvt_float_to_fp8(scaled, __HIP_SATFINITE, __HIP_E4M3); - } - *reinterpret_cast(&fp8_out[base]) = *reinterpret_cast(out); - } - - sdata[tid] = local_max; - __syncthreads(); - for (int s = THREADS_PER_WG / 2; s > 0; s >>= 1) { - if (tid < s) sdata[tid] = fmaxf(sdata[tid], sdata[tid + s]); - __syncthreads(); - } - if (tid == 0) amax_partial[wg] = sdata[0]; -} diff --git a/extra/lr_scheduler.py b/extra/lr_scheduler.py index 87ff077a40c08..094155a71274c 100644 --- a/extra/lr_scheduler.py +++ b/extra/lr_scheduler.py @@ -6,7 +6,7 @@ class LR_Scheduler: def __init__(self, optimizer: Optimizer): self.optimizer = optimizer - self.epoch_counter = Tensor([0], requires_grad=False, device=self.optimizer.device) + self.epoch_counter = Tensor([0], device=self.optimizer.device) def get_lr(self): pass diff --git a/extra/models/bert.py b/extra/models/bert.py index 4528be8920b3b..df619724f0bf6 100644 --- a/extra/models/bert.py +++ b/extra/models/bert.py @@ -52,7 +52,7 @@ def __call__(self, input_ids:Tensor, attention_mask:Tensor, masked_lm_positions: # Reference has residual on denominator: https://github.com/mlcommons/training/blob/master/language_model/tensorflow/bert/run_pretraining.py#L315 def sparse_categorical_crossentropy(self, predictions:Tensor, labels:Tensor, ignore_index=-1): log_probs, loss_mask = predictions.log_softmax(dtype=dtypes.float), (labels != ignore_index) - y_counter = Tensor.arange(predictions.shape[-1], requires_grad=False, device=predictions.device).unsqueeze(0).expand(labels.numel(), predictions.shape[-1]) + y_counter = Tensor.arange(predictions.shape[-1], device=predictions.device).unsqueeze(0).expand(labels.numel(), predictions.shape[-1]) y = ((y_counter == labels.flatten().reshape(-1, 1)) * loss_mask.reshape(-1, 1)).reshape(*labels.shape, predictions.shape[-1]) return -((log_probs * y).sum()) / (loss_mask.sum() + 1e-5) # Small constant to avoid division by zero @@ -159,7 +159,7 @@ def __call__(self, hidden_states:Tensor): return self.dense(hidden_states[:, 0]).tanh() def gather(prediction_logits:Tensor, masked_lm_positions:Tensor): - counter = Tensor.arange(prediction_logits.shape[1], device=prediction_logits.device, requires_grad=False).reshape(1, 1, prediction_logits.shape[1]).expand(*masked_lm_positions.shape, prediction_logits.shape[1]) + counter = Tensor.arange(prediction_logits.shape[1], device=prediction_logits.device).reshape(1, 1, prediction_logits.shape[1]).expand(*masked_lm_positions.shape, prediction_logits.shape[1]) onehot = counter == masked_lm_positions.unsqueeze(2).expand(*masked_lm_positions.shape, prediction_logits.shape[1]) return onehot @ prediction_logits @@ -189,7 +189,7 @@ def __call__(self, input_ids, token_type_ids): input_shape = input_ids.shape seq_length = input_shape[1] - position_ids = Tensor.arange(seq_length, requires_grad=False, device=input_ids.device).unsqueeze(0).expand(*input_shape) + position_ids = Tensor.arange(seq_length, device=input_ids.device).unsqueeze(0).expand(*input_shape) words_embeddings = self.word_embeddings(input_ids) position_embeddings = self.position_embeddings(position_ids) token_type_embeddings = self.token_type_embeddings(token_type_ids) diff --git a/extra/models/llama.py b/extra/models/llama.py index f8f8f52fa8970..abb9d63a066ec 100644 --- a/extra/models/llama.py +++ b/extra/models/llama.py @@ -201,7 +201,7 @@ def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_layers:int, norm_eps: self.tok_embeddings = embedding(vocab_size, dim) self.output = nn.Linear(dim, vocab_size, bias=False) if embedding == nn.Embedding else linear(dim, vocab_size, bias=False) self.max_context = max_context - self.freqs_cis = precompute_freqs_cis(dim // n_heads, self.max_context * 2, rope_theta).contiguous().requires_grad_(False) + self.freqs_cis = precompute_freqs_cis(dim // n_heads, self.max_context * 2, rope_theta).contiguous().is_param_(False) self.forward_jit = TinyJit(self.forward) if jit else None def forward(self, tokens:Tensor, start_pos:Union[Variable,int], temperature:float, top_k:int, top_p:float, alpha_f:float, alpha_p:float): diff --git a/extra/models/mask_rcnn.py b/extra/models/mask_rcnn.py index 052ccd2328b60..d61b811e75fd7 100644 --- a/extra/models/mask_rcnn.py +++ b/extra/models/mask_rcnn.py @@ -78,7 +78,7 @@ def tensor_getitem(tensor, *keys): # for gather with indicies only on axis=0 def tensor_gather(tensor, indices): if not isinstance(indices, Tensor): - indices = Tensor(indices, requires_grad=False) + indices = Tensor(indices) if len(tensor.shape) > 2: rem_shape = list(tensor.shape)[1:] tensor = tensor.reshape(tensor.shape[0], -1) diff --git a/extra/models/rnnt.py b/extra/models/rnnt.py index e7ad0f54b9a84..83218cd087ffe 100644 --- a/extra/models/rnnt.py +++ b/extra/models/rnnt.py @@ -15,7 +15,7 @@ def __init__(self, input_features=240, vocab_size=29, enc_hidden_size=1024, pred @TinyJit def __call__(self, x, y, hc=None): f, _ = self.encoder(x, None) - g, _ = self.prediction(y, hc, Tensor.ones(1, requires_grad=False)) + g, _ = self.prediction(y, hc, Tensor.ones(1)) out = self.joint(f, g) return out.realize() @@ -30,10 +30,10 @@ def decode(self, x, x_lens): return outputs def _greedy_decode(self, logits, logit_len): - hc = Tensor.zeros(self.prediction.rnn.layers, 2, self.prediction.hidden_size, requires_grad=False) + hc = Tensor.zeros(self.prediction.rnn.layers, 2, self.prediction.hidden_size) labels = [] - label = Tensor.zeros(1, 1, requires_grad=False) - mask = Tensor.zeros(1, requires_grad=False) + label = Tensor.zeros(1, 1) + mask = Tensor.zeros(1) for time_idx in range(logit_len): logit = logits[time_idx, :, :].unsqueeze(0) not_blank = True @@ -41,7 +41,7 @@ def _greedy_decode(self, logits, logit_len): while not_blank and added < 30: if len(labels) > 0: mask = (mask + 1).clip(0, 1) - label = Tensor([[labels[-1] if labels[-1] <= 28 else labels[-1] - 1]], requires_grad=False) + 1 - 1 + label = Tensor([[labels[-1] if labels[-1] <= 28 else labels[-1] - 1]]) + 1 - 1 jhc = self._pred_joint(Tensor(logit.numpy()), label, hc, mask) k = jhc[0, 0, :29].argmax(axis=0).numpy() not_blank = k != 28 @@ -129,7 +129,7 @@ def _do_step(x_, hc_): return self.do_step(x_, hc_) if hc is None: - hc = Tensor.zeros(self.layers, 2 * x.shape[1], self.hidden_size, requires_grad=False).contiguous().realize() + hc = Tensor.zeros(self.layers, 2 * x.shape[1], self.hidden_size).contiguous().realize() output = None for t in range(x.shape[0]): diff --git a/extra/models/transformer.py b/extra/models/transformer.py index 4753e50731042..8e0ada930d4c8 100644 --- a/extra/models/transformer.py +++ b/extra/models/transformer.py @@ -41,7 +41,7 @@ def __call__(self, x): class Transformer: def __init__(self, syms, maxlen, layers, embed_dim, num_heads, ff_dim): self.maxlen, self.syms = maxlen, syms - self.embed = Tensor.scaled_uniform(maxlen+syms, embed_dim, requires_grad=False) + self.embed = Tensor.scaled_uniform(maxlen+syms, embed_dim).is_param_(False) self.tbs = [TransformerBlock(embed_dim, num_heads, ff_dim) for _ in range(layers)] self.final = Tensor.scaled_uniform(embed_dim, syms) diff --git a/extra/models/unet.py b/extra/models/unet.py index 1d39d820e5564..88be107981b18 100644 --- a/extra/models/unet.py +++ b/extra/models/unet.py @@ -1,5 +1,4 @@ -from tinygrad import Tensor, dtypes, nn -from tinygrad.device import is_dtype_supported +from tinygrad import Tensor, Device, dtypes, nn from typing import Optional, Union, List, Any, Tuple, Callable import math @@ -13,7 +12,7 @@ def timestep_embedding(timesteps:Tensor, dim:int, max_period=10000): freqs = (-math.log(max_period) * Tensor.arange(half, device=timesteps.device) / half).exp() args = timesteps.unsqueeze(1) * freqs.unsqueeze(0) out = Tensor.cat(args.cos(), args.sin(), dim=-1) - return out.cast(mixed_precision_dtype) if is_dtype_supported(mixed_precision_dtype) else out + return out.cast(mixed_precision_dtype) if mixed_precision_dtype in Device[Device.DEFAULT].renderer.supported_dtypes() else out class ResBlock: def __init__(self, channels:int, emb_channels:int, out_channels:int, num_groups:int=32): @@ -238,7 +237,7 @@ def __call__(self, x:Tensor, tms:Tensor, ctx:Tensor, y:Optional[Tensor]=None) -> assert y.shape[0] == x.shape[0] emb = emb + y.sequential(self.label_emb[0]) - if is_dtype_supported(mixed_precision_dtype): + if mixed_precision_dtype in Device[Device.DEFAULT].renderer.supported_dtypes(): emb = emb.cast(mixed_precision_dtype) ctx = ctx.cast(mixed_precision_dtype) x = x .cast(mixed_precision_dtype) diff --git a/extra/optimization/test_beam_search.py b/extra/optimization/test_beam_search.py index a7cc18f83f71f..133c779960148 100644 --- a/extra/optimization/test_beam_search.py +++ b/extra/optimization/test_beam_search.py @@ -1,7 +1,7 @@ import unittest import numpy as np -from tinygrad.helpers import BEAM, Timing, CI, prod +from tinygrad.helpers import BEAM, Timing, prod from tinygrad import Variable, Device, Tensor from tinygrad.nn import Conv2d from tinygrad.uop.ops import AxisType, Ops @@ -64,7 +64,7 @@ def test_no_mutate_rawbuffers(self): actual = a.numpy() np.testing.assert_allclose(actual, desired) - @unittest.skipIf(CI, "flaky. CL_OUT_OF_RESOURCES") + @unittest.skip("flaky. CL_OUT_OF_RESOURCES") def test_conv_beam(self): c = Conv2d(3, 16, (3,3)) x = rand(1,3,32,32) diff --git a/extra/remote/serve.py b/extra/remote/serve.py index 0f9496ffdc5cc..a1f30cbd79ddf 100644 --- a/extra/remote/serve.py +++ b/extra/remote/serve.py @@ -84,8 +84,6 @@ def serve(conn:socket.socket): conn.sendall(resp_err(str(e))) if __name__ == "__main__": - if not OSX: System.reserve_hugepages(128) # for sysmem allocations - port = int(sys.argv[1]) if len(sys.argv) > 1 else 6667 server = socket.socket(socket.AF_INET, socket.SOCK_STREAM) server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) diff --git a/extra/setup_hipcc_osx.sh b/extra/setup_hipcc_osx.sh index e41a3629ab71f..48b04f9467369 100755 --- a/extra/setup_hipcc_osx.sh +++ b/extra/setup_hipcc_osx.sh @@ -1,16 +1,7 @@ #!/bin/sh install_loc="$HOME/.local/bin" -docker build --platform=linux/amd64 -t rocm-hipcc:7.2 - <<'EOF' -FROM ubuntu:22.04 -ENV DEBIAN_FRONTEND=noninteractive -ENV TZ=Etc/UTC -RUN apt-get update && apt-get install -y --no-install-recommends wget ca-certificates gnupg tzdata && \ - wget https://repo.radeon.com/amdgpu-install/7.2/ubuntu/jammy/amdgpu-install_7.2.70200-1_all.deb && \ - apt-get install -y ./amdgpu-install_7.2.70200-1_all.deb && \ - amdgpu-install -y --usecase=rocm --no-dkms --no-32 && \ - rm -rf /var/lib/apt/lists/* -ENV PATH=/opt/rocm/bin:$PATH -EOF +docker pull --platform=linux/amd64 rocm/dev-ubuntu-22.04:7.1.1 +docker tag rocm/dev-ubuntu-22.04:7.1.1 rocm-hipcc:7.1.1 mkdir -p "$install_loc" tee "$install_loc/hipccshim" >/dev/null <<'EOF' @@ -21,7 +12,7 @@ if ! docker inspect --format='{{.State.Running}}' "$cname" 2>/dev/null | grep -q docker rm -f "$cname" 2>/dev/null || true docker run -d --platform=linux/amd64 --name "$cname" \ -v /var/folders:/var/folders -v "$HOME":"$HOME" \ - rocm-hipcc:7.2 sleep 300 >/dev/null + rocm-hipcc:7.1.1 sleep 300 >/dev/null fi exec docker exec "$cname" "$(basename "$0")" "$@" EOF diff --git a/extra/testsig/generate_testsig.py b/extra/testsig/generate_testsig.py new file mode 100644 index 0000000000000..ad7f78d345087 --- /dev/null +++ b/extra/testsig/generate_testsig.py @@ -0,0 +1,294 @@ +#!/usr/bin/env python3 +""" +Self-contained Qualcomm Hexagon **testsig** generator. + +Replicates: python2 elfsigner.py -t 0x67489311 -o . +Dependencies: standard library + cryptography (pip install cryptography). +Multiple serial numbers: use -t multiple times. +""" + +import argparse, base64, hashlib, os, struct +from datetime import datetime, timedelta, timezone +from cryptography import x509 +from cryptography.hazmat.primitives import hashes, serialization +from cryptography.hazmat.primitives.asymmetric import rsa +from cryptography.x509.name import _ASN1Type +from cryptography.x509.oid import ExtensionOID, NameOID, ObjectIdentifier + +# Embedded assets (raw base64) --- no external files needed +# Compact test_elf_nop.so contents used for signing. Only bytes covered by +# program headers are needed; section headers and alignment gaps are not. +ORIG_PHDRS = [ + {'t': 1, 'o': 0x0000, 'v': 0x0000, 'p': 0x0000, 'fs': 0x02fc, 'ms': 0x02fc, 'fl': 0x4, 'al': 0x1000}, + {'t': 1, 'o': 0x1000, 'v': 0x1000, 'p': 0x1000, 'fs': 0x0104, 'ms': 0x0104, 'fl': 0x5, 'al': 0x1000}, + {'t': 1, 'o': 0x2000, 'v': 0x2000, 'p': 0x2000, 'fs': 0x0004, 'ms': 0x0004, 'fl': 0x4, 'al': 0x1000}, + {'t': 1, 'o': 0x3000, 'v': 0x4000, 'p': 0x4000, 'fs': 0x00d0, 'ms': 0x0100, 'fl': 0x6, 'al': 0x1000}, + {'t': 2, 'o': 0x3010, 'v': 0x4010, 'p': 0x4010, 'fs': 0x00a8, 'ms': 0x00a8, 'fl': 0x6, 'al': 0x4}, +] +ORIG_SEGS = { + 0x0000: base64.b64decode(""" +f0VMRgEBAQAAAAAAAAAAAAMApAABAAAAsBAAADQAAACIMQAAAwAAADQAIAAFACgAFQASAAEAAAAAAAAAAAAAAAAAAAD8AgAA/AIA +AAQAAAAAEAAAAQAAAAAQAAAAEAAAABAAAAQBAAAEAQAABQAAAAAQAAABAAAAACAAAAAgAAAAIAAABAAAAAQAAAAEAAAAABAAAAEA +AAAAMAAAAEAAAABAAADQAAAAAAEAAAYAAAAAEAAAAgAAABAwAAAQQAAAEEAAAKgAAACoAAAABgAAAAQAAAADAAAAEwAAABIAAAAR +AAAADgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAACgAAAAAAAAAJAAAACwAAAAwAAAANAAAA +DwAAABAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAQAAAAAAAAAwAFAAAAAACwEAAAAAAAAAMABwAAAAAAwBAAAAAAAAADAAgAAAAA +AAAgAAAAAAAAAwAJAAAAAAAAQAAAAAAAAAMACgAAAAAACEAAAAAAAAADAAsAAAAAAMxAAAAAAAAAAwAPAAAAAAAAQQAAAAAAAAMA +EAAPAAAAABAAAGwAAAASAAUAFQAAAAAAAAAAAAAAEgAAAGkAAAAAQQAAAAAAABAA8f9KAAAAzEAAAAQAAAARAA8AMQAAAARAAAAA +AAAAEAAKAAEAAAAIQAAAAAAAABAACwBcAAAAwBAAAEQAAAASAAgAYgAAANBAAAAAAAAAEADx/3UAAAAAQQAAAAAAABAA8f9GAAAA +sBAAAAQAAAASAAcAAF9fRFRPUl9MSVNUX18AX2luaXQAX19yZWdpc3Rlcl9mcmFtZV9pbmZvX2Jhc2VzAF9fQ1RPUl9FTkRfXwBs +aWJjLnNvAG5vcABub3BfdmFyAGxpYmdjYy5zbwBfZmluaQBfZWRhdGEAX19ic3Nfc3RhcnQAX2VuZAB0ZXN0X2VsZl9ub3Auc28A +AADIQAAAIgoAAAAAAAA= +"""), + 0x1000: base64.b64decode(""" +AcCdoADbnaEB2J2hGMAJalTP6nH//+pyGNgq8///4HJI3+BxAMAY8wHAgJEIwAEQAkAAeAEoAyg0wABa///7ckz/+3Eb2xjzm//7 +vwDAm5EGwAAQAMCgUPj//1k4wJ2RG0CdkR7AHpAAwJ9SAAAAAMFAAAAcxElqDkKc4k9AnJE8wJyRDkIOjADAnFIAAAAAAAAAAAAA +AAAAAAAAAAAAAMBAAAAO1ElqHMCOkQDAnFIAwJ9SAAAAAAAAAAAAAAAAAcCdoADbnaEPwAlqENDqcf//6nIPzyrz///7clD/+3Eb +2w/zm8AbsADAm5EGwAAQAMCgUPj//1kbQJ2RHsAekADAn1I= +"""), + 0x2000: base64.b64decode("AAAAAA=="), + 0x3000: base64.b64decode(""" +AAAAAAAAAAAAAAAAAAAAAAEAAAA+AAAAAQAAAFIAAAAOAAAAegAAAAwAAAAAEAAADQAAAMAQAAAEAAAA1AAAAAUAAABkAgAABgAA +ADQBAAAKAAAAigAAAAsAAAAQAAAAAwAAALhAAAACAAAADAAAABQAAAAHAAAAFwAAAPACAAABAABwAwAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABBAAAAAAAAAAAAAAAAAAABwEAAACgAAAA== +"""), +} +ORIG_EHDR = ORIG_SEGS[0][:0x34] + +ATTESTCA_KEY = base64.b64decode("""LS0tLS1CRUdJTiBSU0EgUFJJVkFURSBLRVktLS0tLQpNSUlFb3dJQkFBS0NBUUVBeGpLNkRZbVZxbGIvZHFtb1BtY1FsNkV6Vmx3S3dIUnFCQmdsZzBkQnd6djJMR1NKCi9xWnlpbkFWZlloaEFkKzZwbm1kYk1vMVMzaXdjNXUrWUtRYlFOSGpVMUZrZDN4c29weHA3YTZHNnZUbSsyNVEKVnRHRkdLRHk3dm9mRWRKNTE2NVpQU0JtZ0thcWFQcFIxRlNHQnRDaEV3bzhkVEFZdjBaeXFYeGlmeDU4ejcvbgozT1BBbDJQdkpjZWpzd1J2ekRvcVRDR0xlMjVMbzFiYlQ1RGM4bFhZeWFmcjJVcGpXc1ZBaGlHQVNaK2dIRzYrCk1xWDFUdUlvM3FJTGNOcnRNUVh6K2NMamZaSnYzZExqNmRFSlZ0c3MvT3lMdUQ2bVozMXZUdzMxakhXUGRLeDIKNHh3dXVrbkp5OUVrUElQem5ORTZGcFZ4dmdHQWN4aEhKY0pLdndJQkF3S0NBUUVBaENIUlhsdTVIRG4vcEhFYQoxRVMxdW10M2p1Z0hLdmhHcldWdVY0VFdnaWY1Y3UyeFZHNzNCdlZqcVFXV0FUL1J4RkVUbmR3amg2WEs5NzBwCmxjSzgxZUZDTjR1WVQ2aEliR2hHbm5SWjhmaVovUFExanpaWXV4WDNTZndVdG94Uk9uUTdmaFdacXhuRzhLYmgKT0RoWldlQnJZZ2JTK01xN0tpNzNHNmhCcWhNbU1xcFJoK2dUNUVZbnBHZmJUbVRlUUJveXBrTEtkZm1PUytsWQpDM0lFaXJ1LzRhTXhpc0FDNzMvY3VWenQ1T2k5S3BaVVNhWTJkK1pINWtyWWtEMWhMVDFFRTJxbDhuTS9YZ2NlCk81bDkxeVMxZXg5Z2pzZ0lwTnc5elF0U05qMW80SUFraWRnWDhCak5MdmJSYThWdzBibDgxZXp2UEFoSE9tM0wKazkxTGV3S0JnUUR3MkYrUUZxK2xWcEtzL3pFNTlyUjNCdHBBTFVIVEFqWW1rbTVNUHRLeWtYVUQyQ2lNV3lYVgpxNGNaRnBaV2tUY1YzKzN2Q2FQcnNWdFBlcFZsaTZvM3VyMmR2VVFyQ2IxQUh0d2E4R21Bb2VseUJFSzZSeUdPCjh0dW16aTZqZDEyM0xQQTgvQ1JzL01ncmtjM2QxWTVSSDhzS3F1OGZCdlRncjNId2JTYkRPd0tCZ1FEU3EyRGQKZWxmOWFtZ0dyL3FneHJpclpUaWVKWHVJeUVIUG53cUsvNU1qa01jMUh3cVZMZ1NKUi81Y1dYL0U0UXR1Zko5UQp1b2kzNmExdGlweGhETlNqc214d0gxNi9pQ2hRc2V3Ym9BRkw4anFxdi9yQW12RUtmVmRyelU1V3c3dTY5dndECkdET2QxVnhxYzVFTVJVMTVRbS8wS3NMcjNRQTA3QUdsV3MrV1RRS0JnUUNna0QrMVpIVVk1R0hJcWlEUlR5TDYKQkpHQUhpdmlBWGx2REVtSUtlSE1ZUGl0T3NXeTUyNk9jbG9RdWJtUEMzb09sVWxLQm0xSHk1STAvR09aQjhiUAowZE8rZmkxeUJuNHF2ejFuU3ZFQWEvRDJyWUhSaE1FSjl6MFozc25DVDVQUGMwclRVc0x6VXpBZEM5NlQ0N1EyCkZUSUhISjlxQktOQWRQYWdTTVNDSndLQmdRQ01ja0NUcHVWVG5FVlp5cWNWMmRCeVEzc1VHUDBGMnRhS2FnY0gKVlF6Q1lJVE9GTEc0eVZoYmhWUTlrUC9ZbGdlZS9iK0xKd1hQOFI1SkJ4THJYZU1YekVoS3Y1Ui9zQnJnZHAxbgp3QURkVENjY2YveUFaMHRjVTQrZE00bVBMU2ZSK2YxWFpYZStqajJjVFF0ZGc0ajdnWi80SElIeWsxVjRuVlp1ClBJcGtNd0tCZ0FrZ0ljYlh0VHlDbGxJWnFWVTJseUUrcng3Wm5RNGs5ZTBaSENhaXFBT3JDNVBseUtwQ0hRZEgKWGYxcVhtRTBPeFhBeVEycU8wbWJWMzgza2ovU3E3b0p0RHRveS9Bc3ZiRG1vNjZjTzJRSXRSRmZPazJ6Q3UycQpDMFpraE9nazNGWUo4aXloV3pPV3VDWExKck9QVVVqekFkSDJTYWloVy9KY0hVUmhNTnBTCi0tLS0tRU5EIFJTQSBQUklWQVRFIEtFWS0tLS0tCg==""") +ATTESTCA_CERT = base64.b64decode("""MIIEIDCCAwigAwIBAgIBBTANBgkqhkiG9w0BAQsFADCBsjELMAkGA1UEBhMCVVMxEzARBgNVBAgTCkNhbGlmb3JuaWExEjAQBgNVBAcTCVNhbiBEaWVnbzEwMC4GA1UECxMnR2VuZXJhbCBVc2UgVGVzdCBLZXkgKGZvciB0ZXN0aW5nIG9ubHkpMRowGAYDVQQLExFDRE1BIFRlY2hub2xvZ2llczERMA8GA1UEChMIUVVBTENPTU0xGTAXBgNVBAMTEFFQU0EgU0hBMjU2IFJvb3QwHhcNMTMwNDIyMjIwNjE4WhcNMzMwNDE3MjIwNjE4WjB8MQswCQYDVQQGEwJVUzELMAkGA1UECBMCQ0ExEjAQBgNVBAcTCVNhbiBEaWVnbzEaMBgGA1UECxMRQ0RNQSBUZWNobm9sb2dpZXMxETAPBgNVBAoTCFFVQUxDT01NMR0wGwYDVQQDExRRUFNBIE9QRU5EU1AgVEVTVCBDQTCCASAwDQYJKoZIhvcNAQEBBQADggENADCCAQgCggEBAMYyug2JlapW/3apqD5nEJehM1ZcCsB0agQYJYNHQcM79ixkif6mcopwFX2IYQHfuqZ5nWzKNUt4sHObvmCkG0DR41NRZHd8bKKcae2uhur05vtuUFbRhRig8u76HxHSedeuWT0gZoCmqmj6UdRUhgbQoRMKPHUwGL9Gcql8Yn8efM+/59zjwJdj7yXHo7MEb8w6Kkwhi3tuS6NW20+Q3PJV2Mmn69lKY1rFQIYhgEmfoBxuvjKl9U7iKN6iC3Da7TEF8/nC432Sb93S4+nRCVbbLPzsi7g+pmd9b08N9Yx1j3SsduMcLrpJycvRJDyD85zROhaVcb4BgHMYRyXCSr8CAQOjeDB2MB8GA1UdIwQYMBaAFElk8+VAE1VZc2dnWT99Qreru/tXMB0GA1UdDgQWBBQnxAfEeRhBNAnuLkunmI4I+aSyOzAPBgNVHRMECDAGAQH/AgEAMAsGA1UdDwQEAwIBBjAWBgorBgEEAYspCQYDBAgAAeJAAAn78TANBgkqhkiG9w0BAQsFAAOCAQEAYHPAAlh+ezXdqUDIptraYfoiVxw2YsX++Ytg2eJ69YFVlCo33bLJFwQMj+zTMauRgvLew2cZTK47ghVV7130M13E53aN49p/DTOe3u5OFGA+z+ZLrqhraUPT+UhaAuVO9Yu9eOLudsPvgJTeD1a7RaC6PmPsUFPxLUlmlJn3lSXjlYe98+hittLnJ9gTnjdTVH/PgEJhMvUcjjyBWdRsog54VpyqesqLJedC4OF7fHJZ4S7rxDAINI15aDBQrOW/LD6HsBdr4WikS5Lnmecaw+2Um/ge/3Jl/kFBgh8EyORmSzaN4q1OoPYykxTGxenP3Z6D9WJurPd0d0fnuf+bNw==""") +ROOTCA_CERT = base64.b64decode("""MIIEGzCCAwOgAwIBAgIBATANBgkqhkiG9w0BAQsFADCBsjELMAkGA1UEBhMCVVMxEzARBgNVBAgTCkNhbGlmb3JuaWExEjAQBgNVBAcTCVNhbiBEaWVnbzEwMC4GA1UECxMnR2VuZXJhbCBVc2UgVGVzdCBLZXkgKGZvciB0ZXN0aW5nIG9ubHkpMRowGAYDVQQLExFDRE1BIFRlY2hub2xvZ2llczERMA8GA1UEChMIUVVBTENPTU0xGTAXBgNVBAMTEFFQU0EgU0hBMjU2IFJvb3QwHhcNMTMwMzI4MjMxOTA4WhcNMzMwMzIzMjMxOTA4WjCBsjELMAkGA1UEBhMCVVMxEzARBgNVBAgTCkNhbGlmb3JuaWExEjAQBgNVBAcTCVNhbiBEaWVnbzEwMC4GA1UECxMnR2VuZXJhbCBVc2UgVGVzdCBLZXkgKGZvciB0ZXN0aW5nIG9ubHkpMRowGAYDVQQLExFDRE1BIFRlY2hub2xvZ2llczERMA8GA1UEChMIUVVBTENPTU0xGTAXBgNVBAMTEFFQU0EgU0hBMjU2IFJvb3QwggEgMA0GCSqGSIb3DQEBAQUAA4IBDQAwggEIAoIBAQC3mmlyc5XmZ4nQcUf8gXoHX82fCU12SW6VJdlz5IyKOJzl+IeYs2ArpkDHXaF2NwYvS4cJVBHtvx5TLbsBMAF9teFORqSs6wl+r+3nQwCogNOn/8JZrcPdxkjA8cVAkydxSK0jPxAdAGf8vGXD7tKDWWZyHquPoWqNVG/P4OyHAWMKCYg/w7/5MYTOcV1TXW2BraH7dztGkS4ey2hsOPlJzxP74cN1WyXjLPkn5CZWkx+95CKN5i+T9S+FeKD/1zbxuNlwv4x3x1Ohw9KBJYQzrB/wP9wrsVEnh2K9jy7rapKFFWOgQj8omg1EbIMqdOHuSZYcexFvAqN233xxluDBAgEDozwwOjAdBgNVHQ4EFgQUSWTz5UATVVlzZ2dZP31Ct6u7+1cwDAYDVR0TBAUwAwEB/zALBgNVHQ8EBAMCAQYwDQYJKoZIhvcNAQELBQADggEBAJ31kC2nKcTK1XrArhpkmnAX8zCPZkl+Azm7qF/Mr09h8FJiXJ7hBBoLHD+/+DifgUkLF4EjLOUnTTUPLPPKS5KgvuDkgJKvAMvv/GXxuGabdl4EebYCbJvnEgwkyG4pYVd5pGHQ0z2Md4nI6aMYco2X44bIjdqxFJwWOgPhioef1KbO/6CzykG0cPvpJB0XkWR8QGztFt9HofF+uVgpY2t1dL4/SuM/wJSeS8rdqstk0BYq/EDFFs99r1wP2R6hjJgCPkMvo7WiSE2yRrEkMNhgyEKrrD4pr7YWtsFPkYeTMXOvYoY16aOfvcrw0FfH+yATdn/OdAQ2saJISmilWh4=""") + +PT_NULL = 0 +PF_OS_PHDR = 0x07000000 +PF_OS_HASH = 0x02200000 +MBN_V3 = 3 +SIG_SIZE = 256 +CC_SIZE = 0x1800 +PAD_FF = b'\xff' +IPAD = 0x3636363636363636 +OPAD = 0x5C5C5C5C5C5C5C5C + + +def _pad(data, size, pad=PAD_FF): + if len(data) > size: + raise ValueError("data too large: %d > %d" % (len(data), size)) + return data + pad * (size - len(data)) + + +def _orig_segment_data(ph): + off, size = ph['o'], ph['fs'] + for base_off, data in ORIG_SEGS.items(): + rel = off - base_off + if 0 <= rel and rel + size <= len(data): + return data[rel:rel + size] + raise KeyError("missing original segment at 0x%x" % off) + + +def _build_ehdr(base_ehdr, num_phdrs): + e = bytearray(base_ehdr) + struct.pack_into('>y, "aten.__irshift__.Scalar": lambda x,y: x>>y, # inplace ops using replace for fusion - "aten.zero_": lambda x: x.zeros_like(), - "aten.fill_.Scalar": lambda x, y: x.full_like(y), + "aten.zero_": lambda x: x.const_like(0), + "aten.fill_.Scalar": lambda x, y: x.const_like(y), "aten.add_.Tensor": lambda self, other, alpha=1.0: self + other * alpha, "aten.add_.Scalar": lambda self, other, alpha=1.0: self + other * alpha, "aten.mul_.Tensor": lambda self, other: self * other, @@ -617,7 +617,7 @@ def _inplace_op(t, new_value): "aten.asinh": Tensor.asinh, "aten.mul": Tensor.mul, "aten.atanh": Tensor.atanh, - "aten.fill_.Tensor": lambda self, value: Tensor.full(self.shape, value.reshape(()).item(), device=self.device, dtype=self.dtype), + "aten.fill_.Tensor": lambda self, value: self.const_like(value.reshape(()).item()), "aten.flip": Tensor.flip, "aten.scatter_reduce.two": Tensor.scatter_reduce, "aten.squeeze_.dim": Tensor.squeeze, @@ -709,7 +709,7 @@ def nf(*args, **kwargs): views = derived_views(base) if views: old_base = Tensor(base.uop, device=base.device) - old_base.requires_grad = base.requires_grad + old_base.is_param = base.is_param old_base._views = getattr(base, "_views", set()) for v in views: v._view_base = old_base base._views = set() diff --git a/extra/torch_backend/test.py b/extra/torch_backend/test.py index ef3043569affd..f3a1bcc93f73d 100644 --- a/extra/torch_backend/test.py +++ b/extra/torch_backend/test.py @@ -240,9 +240,9 @@ def test_floor_div(self): np.testing.assert_equal(result.cpu().numpy(), [3., 3., 2.]) def test_mnist_index(self): + # from tinygrad.nn.datasets import mnist + X_train, Y_train = Tensor.randint(60000, 1, 28, 28, dtype='uchar').realize(), Tensor.randint(60000, dtype='uchar').realize() GlobalCounters.reset() - from tinygrad.nn.datasets import mnist - X_train, Y_train, _, _ = mnist() X_train = torch.tensor(X_train.float().numpy(), device=device) Y_train = torch.tensor(Y_train.cast('int64').numpy(), device=device) samples = torch.randint(0, X_train.shape[0], (32,)) diff --git a/extra/training.py b/extra/training.py index d134e68bf2ef4..7d0830ea8c999 100644 --- a/extra/training.py +++ b/extra/training.py @@ -1,6 +1,6 @@ import numpy as np from tinygrad.tensor import Tensor -from tinygrad.helpers import CI, trange +from tinygrad.helpers import trange from tinygrad.engine.jit import TinyJit @@ -24,9 +24,9 @@ def train_step(x, y): with Tensor.train(): losses, accuracies = [], [] - for i in (t := trange(steps, disable=CI)): + for i in (t := trange(steps, disable=None)): samp = np.random.randint(0, X_train.shape[0], size=(BS)) - x = Tensor(transform(X_train[samp]), requires_grad=False) + x = Tensor(transform(X_train[samp])) y = Tensor(target_transform(Y_train[samp])) loss, accuracy = train_step(x, y) # printing @@ -43,7 +43,7 @@ def evaluate(model, X_test, Y_test, num_classes=None, BS=128, return_predict=Fal Tensor.training = False def numpy_eval(Y_test, num_classes): Y_test_preds_out = np.zeros(list(Y_test.shape)+[num_classes]) - for i in trange((len(Y_test)-1)//BS+1, disable=CI): + for i in trange((len(Y_test)-1)//BS+1, disable=None): x = Tensor(transform(X_test[i*BS:(i+1)*BS])) out = model.forward(x) if hasattr(model, 'forward') else model(x) Y_test_preds_out[i*BS:(i+1)*BS] = out.numpy() diff --git a/extra/viz/kernel_graph.py b/extra/viz/kernel_graph.py index cbcc4157d9d55..8fd09f18bbb55 100755 --- a/extra/viz/kernel_graph.py +++ b/extra/viz/kernel_graph.py @@ -7,15 +7,15 @@ def get_node(graph:dict, key): return graph[str(key)] if __name__ == "__main__": parser = argparse.ArgumentParser(description="print CALL graph from DEBUG=5 tinygrad.viz.cli --json output") - parser.add_argument("kernel", type=str, default=None, help="Kernel name to stop at (default: print all kernels)") + parser.add_argument("kernel", type=str, nargs="?", default="ALL", metavar="NAME", help="Kernel name to stop at (default: print all kernels)") args = parser.parse_args() ref:int|None = None for line in sys.stdin: if not line.strip(): continue graph = json.loads(line) - if ref is not None and graph.get("ref") == ref: + if graph.get("ref") is not None and (args.kernel == "ALL" or graph["ref"] == ref): print(graph) - if (v:=json.loads(next(sys.stdin)).get("value")): print(v) + if (v:=json.loads(next(sys.stdin, "{}")).get("value")): print(v) if ref is not None or not isinstance(rec:=next(iter(graph.values()), {}), dict) or "label" not in rec: continue for v in graph.values(): if not v["label"].startswith("CALL"): continue @@ -39,6 +39,6 @@ def get_node(graph:dict, key): return graph[str(key)] src_str = ["SRC"]+get_node(graph, get_node(graph, s)["src"][0][1])["label"].splitlines()[1:] print(" ".join(idx_str+src_str)) ss += [x[1] for x in get_node(graph, s)["src"]] - if args.kernel is not None and args.kernel in ansistrip(v["label"]): + if args.kernel != "ALL" and args.kernel in ansistrip(v["label"]): ref = v["ref"] break diff --git a/pyproject.toml b/pyproject.toml index c0653ece757a6..ea8f1c15dac24 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "tinygrad" -version = "0.12.0" +version = "0.13.0" description = "You like pytorch? You like micrograd? You love tinygrad! <3" authors = [{ name = "George Hotz" }] @@ -28,6 +28,7 @@ packages = [ 'tinygrad.nn', 'tinygrad.renderer', 'tinygrad.renderer.amd', + 'tinygrad.renderer.isa', 'tinygrad.runtime', 'tinygrad.runtime.autogen', 'tinygrad.runtime.autogen.am', @@ -35,6 +36,7 @@ packages = [ 'tinygrad.runtime.autogen.amd.rdna3', 'tinygrad.runtime.autogen.amd.rdna4', 'tinygrad.runtime.autogen.amd.cdna', + 'tinygrad.runtime.autogen.nv_regs', 'tinygrad.runtime.graph', 'tinygrad.runtime.support', 'tinygrad.runtime.support.am', @@ -47,6 +49,7 @@ packages = [ [tool.setuptools.package-data] tinygrad = ["py.typed"] +"tinygrad.llm" = ["chat.html"] "tinygrad.viz" = ["index.html", "assets/**/*", "js/*"] @@ -133,7 +136,7 @@ debug = true [tool.mypy] warn_unused_configs = true -files = ["tinygrad"] +files = ["tinygrad", "test/mockgpu"] ignore_missing_imports = true check_untyped_defs = true explicit_package_bases = true diff --git a/spec/tinyspec.pdf b/spec/tinyspec.pdf index 249befa073e27..e2420c6ffd28e 100644 Binary files a/spec/tinyspec.pdf and b/spec/tinyspec.pdf differ diff --git a/spec/tinyspec.tex b/spec/tinyspec.tex index 28dc5a9e86614..7f44b38163be0 100644 --- a/spec/tinyspec.tex +++ b/spec/tinyspec.tex @@ -16,6 +16,7 @@ \definecolor{elwyellow}{HTML}{F9A825} \definecolor{callblue}{HTML}{1565C0} \definecolor{assignbrown}{HTML}{795548} +\definecolor{loadred}{HTML}{c08080} \definecolor{multipurple}{HTML}{7B1FA2} \definecolor{markerorange}{HTML}{E65100} % AxisType colors (from tinygrad) @@ -48,16 +49,14 @@ \subsection*{Source Ops \normalfont\small--- leaf nodes} \toprule \textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\ \midrule +\op{Param} & $(\mathbf{s})$ & slot, dtype, device?, addrspace? & + Placeholder with shape $\mathbf{s}$. Substituted in \op{Function}. \\[4pt] \op{Buffer} & () & size, dtype, device, addrspace & Shape $(n \cdot \textit{size},)$ if device is $n$-tuple, else $(\textit{size},)$. \\ -\op{BufferView} & (buf,) & size, dtype, offset & - Typed access into a buffer. Zero-copy $(\textit{size},)$ slice at offset; inherits addrspace. \\ -\op{Param} & $(\mathbf{s})$ or $(\mathbf{s}, \text{min}, \text{max})$ & slot, dtype, device? & - Placeholder with shape $\mathbf{s}$. Substituted in \op{Function}. \\[4pt] \op{Const} & () & value, dtype & A scalar constant with shape $(\ )$. \\ -\op{Vconst} & () & values, dtype & - A vector constant with shape $(n,)$. \\ + & & & Form vector consts with \op{Stack} \\ +\op{Binary} & () & data & Raw binary data, has dtype uint8 and shape len($data$) \\ \bottomrule \end{tabular} @@ -65,7 +64,7 @@ \subsection*{Source Ops \normalfont\small--- leaf nodes} A \op{Buffer}'s \textbf{addrspace} is \texttt{GLOBAL}, \texttt{LOCAL}, or \texttt{REG}. %% ============================================================ -\subsection*{{\color{movgreen}Movement Ops} \normalfont\small--- no arithmetic, shapes are $(k,)$-shaped UOps with dtype \texttt{index} in src} +\subsection*{{\color{movgreen}Movement Ops} \normalfont\small--- no arithmetic; view, indexing, and reinterpretation only} \begin{tabular}{@{}l l l l@{}} \toprule @@ -75,11 +74,13 @@ \subsection*{{\color{movgreen}Movement Ops} \normalfont\small--- no arithmetic, \op{Flip} & $(T,)$ & bools $\mathbf{f}$ & Reverse along flagged axes. \\ \op{Reshape} & $(T, \mathbf{s'})$ & --- & Reinterpret in row-major order. $\prod s_k = \prod s'_k$. \\ \op{Expand} & $(T, \mathbf{s'})$ & --- & Broadcast size-1 axes. $s_k \in \{1, s'_k\}$. \\ -\op{Pad} & $(T, \mathbf{b}, \mathbf{e})$ & --- & Pad with $0$s: $b_k$ before, $e_k$ after each axis. \\ -\op{Shrink} & $(T, \mathbf{b}, \mathbf{e})$ & --- & Keep $[b_k, e_k)$ per axis. Inverse of \op{Pad}. \\ +\op{Pad} & $(T, \mathbf{o}, \mathbf{s'})$ & --- & Place $T$ at offset $o_k$ in a zero-filled output of shape $s'_k$. \\ +\op{Shrink} & $(T, \mathbf{o}, \mathbf{s'})$ & --- & Keep $s'_k$ elements starting at offset $o_k$ per axis. Inverse of \op{Pad}. \\ \op{Index} & $(T, i_0, i_1, \ldots)$ & --- & Index from left. $()$-shaped $i$ removes dim; $(k,)$-shaped makes it $k$. \\ \op{Stack} & $(T_0, T_1, \ldots)$ & --- & Join along a newly created leading axis. All shapes must match. \\ \op{Replicated} & $(T,)$ & axes & Mark $T$ as replicated along axes. Collapse axes to $1$. \\ +\op{Slice} & $(T, \mathrm{offset})$ & size, dtype & Zero-copy \textit{size} elems of dtype; offset is elems of $T$ dtype. \\ +\op{Bitcast} & $(T,)$ & dtype & Reinterpret storage as target dtype; preserve total bytes. \\ \bottomrule \end{tabular} @@ -90,7 +91,7 @@ \subsection*{{\color{reducered}Reduce Ops} \normalfont\small--- collapse axes to \toprule \textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\ \midrule -\op{Reduce} & $(T,)$ & op, axes & Reduce $T$ along axes. Op is \op{Add}, \op{Max}, or \op{Mul}. \\ +\op{Reduce} & ($T$, $r_0$, $r_1$, \ldots) & op, axes & Reduce $T$ along axes or ranges. Op is \op{Add}, \op{Max}, or \op{Mul}. \\ \bottomrule \end{tabular} @@ -109,13 +110,25 @@ \subsection*{{\color{callblue}Call Ops} \normalfont\small--- function abstractio \end{tabular} %% ============================================================ -\subsection*{{\color{multipurple}Store Ops} \normalfont\small--- side effects} +\subsection*{{\color{loadred}Load Ops} \normalfont\small--- can change device or addrspace} + +\begin{tabular}{@{}l l l l@{}} +\toprule +\textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\ +\midrule +\op{Load} & (buf, alt?, gate?) & device, addrspace & Read (pull) from buffer into a new anonymous buffer. \\ + & & & Note: this replaces \op{Copy} and \op{Contiguous}. \\ +\bottomrule +\end{tabular} + +%% ============================================================ +\subsection*{{\color{multipurple}Store Ops} \normalfont\small--- the only op with observable side effects} \begin{tabular}{@{}l l l l@{}} \toprule \textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\ \midrule -\op{Store} & (buf, val, gate?) & --- & Write val into buf. buf.shape $=$ val.shape. \\ +\op{Store} & (buf, val, gate?) & --- & Write (push) val into buf. buf.shape $=$ val.shape. \\ & & & If gate is present, write only when gate is true. Output is void. \\ \bottomrule \end{tabular} @@ -153,8 +166,7 @@ \subsection*{{\color{elwyellow}Elementwise Ops} \normalfont\small--- all inputs & $\mathrm{trunc}(x)$: round toward zero. \\ & & \op{Cast} & Convert to target dtype (specified in arg). \\ -& & \op{Bitcast} - & Reinterpret bits as target dtype. Must be same size. \\[4pt] + \\[4pt] Binary & $(A, B)$ & \op{Add}, \op{Mul}, \op{Max}, \op{Mod}, \op{Idiv} & $a+b$, $a \cdot b$, $\max(a,b)$, $a \bmod b$, $\lfloor a/b \rfloor$ \\ @@ -205,7 +217,6 @@ \subsection*{{\color{markerorange}Marker Ops} \normalfont\small--- identity on d \op{Contiguous} & $(T,)$ & --- & Force contiguous memory layout. \\ \op{ContiguousBackward} & $(T,)$ & --- & Force contiguous in backward pass. \\ \op{Detach} & $(T,)$ & --- & Stops gradient propagation. \\ -\op{Copy} & $(T,)$ & device & Copy to target device. \\ \bottomrule \end{tabular} @@ -216,8 +227,6 @@ \subsection*{Codegen Ops \normalfont\small--- generated code primitives, these d \toprule \textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\ \midrule -\op{Load} & (idx,alt?,gate?) & --- & Dereference: read element at index from buffer. \\ - & & & All loads will be replaced by \op{Store}. \\ \op{Barrier} & (deps\ldots) & --- & Synchronize threads within a workgroup. \\ \op{Ins} & \ldots & \ldots & A single machine instruction (e.g.\ AMD ISA). \\ \op{Special} & (bound,) & name & GPU thread/workgroup index (e.g.\ \texttt{gidx0}, \texttt{lidx1}). \\ diff --git a/test/amd/hw/test_cdna_sdwa.py b/test/amd/hw/test_cdna_sdwa.py new file mode 100644 index 0000000000000..61b95a5ba3324 --- /dev/null +++ b/test/amd/hw/test_cdna_sdwa.py @@ -0,0 +1,20 @@ +#!/usr/bin/env python3 +import unittest +import tinygrad.runtime.autogen.amd.cdna.ins as cdna +from test.amd.hw.test_cdna_vop3 import run_cdna + +class TestCDNASDWA(unittest.TestCase): + def test_v_add_co_u32_e32_writes_vcc(self): + out = run_cdna([ + cdna.s_mov_b32(cdna.s[0], 0xffffffff), + cdna.v_mov_b32_e32(cdna.v[0], cdna.s[0]), + cdna.v_mov_b32_e32(cdna.v[13], 1), + cdna.v_add_co_u32_e32(cdna.v[0], cdna.SDWA, cdna.v[13], vsrc0=cdna.v[0], dst_sel=6, src0_sel=6), + cdna.v_mov_b32_e32(cdna.v[2], cdna.VCC_LO), + cdna.v_lshlrev_b32_e32(cdna.v[2], 31, cdna.v[2]), + cdna.v_or_b32_e32(cdna.v[2], cdna.v[2], cdna.v[0]), + ]) + self.assertEqual(out, 0x80000000) + +if __name__ == "__main__": + unittest.main() diff --git a/test/amd/hw/test_cdna_vop3.py b/test/amd/hw/test_cdna_vop3.py new file mode 100644 index 0000000000000..6d634f6425502 --- /dev/null +++ b/test/amd/hw/test_cdna_vop3.py @@ -0,0 +1,130 @@ +"""CDNA VOP3 instruction coverage. + +Exercises generated CDNA pcode end-to-end in the emulator and compares against +gfx950 hardware when USE_HW=1. +""" +import ctypes, struct, unittest +import tinygrad.runtime.autogen.amd.cdna.ins as cdna +from tinygrad.helpers import flat_mv +from tinygrad.renderer.amd.dsl import NULL +from test.amd.hw.helpers import USE_HW, assemble +from test.mockgpu.amd.emu import run_asm + +LANES = 1 + +def _code(instructions: list, out_reg: int = 2, out_addr: int | None = None) -> bytes: + load_out_addr = [ + cdna.s_mov_b32(cdna.s[92], out_addr & 0xffffffff), + cdna.s_mov_b32(cdna.s[93], out_addr >> 32), + ] if out_addr is not None else [ + cdna.s_load_dwordx2(cdna.s[92:93], cdna.s[80:81], 0, soffset=NULL), + cdna.s_waitcnt(0), + ] + return assemble([ + cdna.s_mov_b32(cdna.s[80], cdna.s[0]), + cdna.s_mov_b32(cdna.s[81], cdna.s[1]), + cdna.v_mov_b32_e32(cdna.v[255], cdna.v[0]), + *instructions, + *load_out_addr, + cdna.v_lshlrev_b32_e32(cdna.v[240], 2, cdna.v[255]), + cdna.global_store_dword(addr=cdna.v[240], data=cdna.v[out_reg], saddr=cdna.s[92:93], offset=0), + cdna.s_endpgm(), + ]) + +def _run_emu(instructions: list, out_reg: int = 2) -> int: + out_buf = (ctypes.c_uint32 * LANES)(*([0] * LANES)) + args = (ctypes.c_uint64 * 1)(ctypes.addressof(out_buf)) + code = _code(instructions, out_reg) + kernel_buf = (ctypes.c_char * len(code)).from_buffer_copy(code) + result = run_asm(ctypes.addressof(kernel_buf), len(code), 1, 1, 1, LANES, 1, 1, ctypes.addressof(args), + 0x19c | (128 << 15), 0x10000, arch="cdna") + assert result == 0, f"run_asm failed with {result}" + return out_buf[0] + +def _run_hw(instructions: list, out_reg: int = 2) -> int: + from tinygrad.device import Device + from tinygrad.runtime.ops_amd import AMDProgram + from tinygrad.runtime.support.compiler_amd import HIPCompiler + + dev = Device["AMD"] + if dev.arch != "gfx950": raise unittest.SkipTest("requires gfx950 hardware") + out_gpu = dev.allocator.alloc(LANES * 4) + code = _code(instructions, out_reg, out_gpu.va_addr) + byte_str = ", ".join(f"0x{b:02x}" for b in code) + asm_src = f""".text +.globl test +.p2align 8 +.type test,@function +test: +.byte {byte_str} + +.rodata +.p2align 6 +.amdhsa_kernel test + .amdhsa_next_free_vgpr 256 + .amdhsa_next_free_sgpr 96 + .amdhsa_accum_offset 256 + .amdhsa_kernarg_size 0 +.end_amdhsa_kernel + +.amdgpu_metadata +--- +amdhsa.version: + - 1 + - 0 +amdhsa.kernels: + - .name: test + .symbol: test.kd + .kernarg_segment_size: 0 + .group_segment_fixed_size: 0 + .private_segment_fixed_size: 0 + .kernarg_segment_align: 8 + .wavefront_size: 64 + .sgpr_count: 96 + .vgpr_count: 256 + .max_flat_workgroup_size: 1024 +... +.end_amdgpu_metadata +""" + prg = AMDProgram(dev, "test", HIPCompiler(dev.arch).compile(asm_src)) + prg(global_size=(1, 1, 1), local_size=(LANES, 1, 1), wait=True) + out = bytearray(LANES * 4) + dev.allocator._copyout(flat_mv(memoryview(out)), out_gpu) + return struct.unpack(" int: + emu = _run_emu(instructions, out_reg) + if not USE_HW: return emu + hw = _run_hw(instructions, out_reg) + if emu != hw: raise AssertionError(f"Emulator vs Hardware mismatch: emu=0x{emu:08x} hw=0x{hw:08x}") + return hw + +class TestCDNAVOP3(unittest.TestCase): + def test_cvt_pk_fp8_f32_preserves_upper_half(self): + """V_CVT_PK_FP8_F32 with OPSEL[3]=0 writes only D[15:0].""" + out = run_cdna([ + cdna.s_mov_b32(cdna.s[0], 0xdeadbeef), + cdna.v_mov_b32_e32(cdna.v[2], cdna.s[0]), + cdna.v_mov_b32_e32(cdna.v[0], 1.0), + cdna.v_mov_b32_e32(cdna.v[1], 2.0), + cdna.v_cvt_pk_fp8_f32(cdna.v[2], cdna.v[0], cdna.v[1]), + ]) + self.assertEqual(out, 0xdead4038) + + def test_cvt_pk_bf8_f32_overflow_and_inf(self): + """V_CVT_PK_BF8_F32 converts finite overflow and infinities to E5M2 infinities.""" + for name, bits, expected in [ + ("finite_overflow", 0x47700000, 0x7c), + ("pos_inf", 0x7f800000, 0x7c), + ("neg_inf", 0xff800000, 0xfc), + ]: + with self.subTest(name=name): + out = run_cdna([ + cdna.s_mov_b32(cdna.s[0], 0xdeadbeef), + cdna.v_mov_b32_e32(cdna.v[2], cdna.s[0]), + cdna.s_mov_b32(cdna.s[0], bits), + cdna.v_mov_b32_e32(cdna.v[0], cdna.s[0]), + cdna.v_mov_b32_e32(cdna.v[1], 1.0), + cdna.v_cvt_pk_bf8_f32(cdna.v[2], cdna.v[0], cdna.v[1]), + ]) + self.assertEqual(out, 0xdead3c00 | expected) diff --git a/test/amd/test_emu2_pcode.py b/test/amd/test_emu2_pcode.py index c0086b4c72fa5..e8f77d748a708 100644 --- a/test/amd/test_emu2_pcode.py +++ b/test/amd/test_emu2_pcode.py @@ -167,7 +167,7 @@ class TestDSPcodePatterns(unittest.TestCase): def test_global_atomic_add_f32_parsing(self): """Test GLOBAL_ATOMIC_ADD_F32 keeps memory values in float dtype.""" - vmem = UOp(Ops.PARAM, dtypes.uint32.ptr(1024), arg=2) + vmem = UOp.param(2, dtypes.uint32.ptr(1024)) srcs = { 'ADDR': UOp.const(dtypes.uint64, 0), 'DATA': UOp.const(dtypes.uint32, 0x3f800000), @@ -198,7 +198,7 @@ def test_ds_store_b32_pcode(self): def test_mem_read_parsing(self): """Test MEM[addr].type read expression parsing.""" # Create a mock LDS buffer - lds = UOp(Ops.PARAM, dtypes.uint32.ptr(16384), arg=3) + lds = UOp.param(3, dtypes.uint32.ptr(16384)) addr = UOp.const(dtypes.uint32, 0) vrs = {'_lds': lds, 'ADDR': addr, 'OFFSET': UOp.const(dtypes.uint32, 0)} @@ -233,7 +233,7 @@ def test_ds_load_2addr_b32_parsing(self): pcode = PCODE.get(DSOp.DS_LOAD_2ADDR_B32) self.assertIsNotNone(pcode) assert pcode is not None - lds = UOp(Ops.PARAM, dtypes.uint32.ptr(16384), arg=3) + lds = UOp.param(3, dtypes.uint32.ptr(16384)) srcs = { 'ADDR': UOp.const(dtypes.uint32, 0), 'OFFSET0': UOp.const(dtypes.uint32, 0), @@ -314,7 +314,7 @@ def test_permlane64_altlane_concat(self): self.assertEqual(parsed.simplify().arg, expected) def test_permlane64_wave64_pcode_indices(self): - vgpr = UOp(Ops.PARAM, dtypes.uint32.ptr(256), arg=0) + vgpr = UOp.param(0, dtypes.uint32.ptr(256)) srcs = { 'SRC0': UOp.const(dtypes.uint32, 0), 'VDST': UOp.const(dtypes.uint32, 1), @@ -347,7 +347,7 @@ class TestAllPcode(unittest.TestCase): def _make_srcs(self): """Create dummy source variables for pcode parsing.""" u32, u64 = lambda v=0: UOp.const(dtypes.uint32, v), lambda v=0: UOp.const(dtypes.uint64, v) - lds = UOp(Ops.PARAM, dtypes.uint32.ptr(16384), arg=3) + lds = UOp.param(3, dtypes.uint32.ptr(16384)) return {'laneId': u32(), 'laneID': u32(), 'S0': u32(), 'S1': u32(), 'S2': u32(), 'S3': u32(), 'SRC0': u32(), 'D0': u32(), 'D1': u32(), 'DST': u32(), 'VDST': u32(), 'SDST': u32(), 'VCC': u64(), 'VCCZ': u32(), 'EXEC': u64(), 'EXEC_LO': u32(), 'EXECZ': u32(), 'SCC': u32(), diff --git a/test/backend/test_arange.py b/test/backend/test_arange.py index 096fc22f6efb0..1e223f797d1e4 100644 --- a/test/backend/test_arange.py +++ b/test/backend/test_arange.py @@ -19,7 +19,7 @@ def test_arange_complexity(self): self.assertEqual(self._get_flops(Tensor.arange(256), np.arange(256)), 0) self.assertEqual(self._get_flops(Tensor.arange(2560), np.arange(2560)), 0) - @unittest.skipIf(Device.DEFAULT == "CL", "TODO: fails on CI CL") + @unittest.skipIf(Device.DEFAULT == "CL", "flaky in CI") def test_arange_cumsum(self): np.testing.assert_equal(Tensor.arange(513).cumsum(0).numpy(), np.arange(513).cumsum()) @@ -64,7 +64,7 @@ def test_manual_index(self): rng = Tensor.arange(DSET, dtype=dtypes.int).reshape(1, 1, DSET, 1).expand(4, DDIM, DSET, 1) idxs = idxs.reshape(4,1,1,1).expand(4, DDIM, DSET, 1) reshape_dataset = dataset.T.reshape(1, DDIM, DSET, 1).expand(4, DDIM, DSET, 1) - full = (rng==idxs).where(reshape_dataset, Tensor.zeros(4, DDIM, DSET, 1)) + full = (rng==idxs).where(reshape_dataset, Tensor.zeros(4, DDIM, DSET, 1, buffer=False)) X = full.sum(axis=(2,3)) linear, var_vals = X.linear_with_vars() self.assertEqual(len(linear.src), 1) @@ -125,8 +125,8 @@ def test_index_fused_out_of_bounds(self): def test_index_mnist(self, noopt=1, op_limit=512*784*13, split_reduceop=0): # WEBGPU generates more ops due to bitpacking of < 4-byte dtypes if Device.DEFAULT == "WEBGPU": op_limit *= 15 - from tinygrad.nn.datasets import mnist - X_train, Y_train, _, _ = mnist() + # from tinygrad.nn.datasets import mnist + X_train, Y_train = Tensor.randint(DSET, 1, 28, 28, dtype='uchar').realize(), Tensor.randint(DSET, dtype='uchar').realize() with Context(NOOPT=noopt, SPLIT_REDUCEOP=split_reduceop): samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0]).realize() GlobalCounters.reset() @@ -172,7 +172,7 @@ def test_llama_8b_embedding_backward(self): bs, seqlen = 4, 256 idx = Tensor.randint(bs, seqlen, high=vocab_size) emb = nn.Embedding(vocab_size, embed_size) - emb.weight = Tensor.ones(vocab_size, embed_size, requires_grad=True) + emb.weight = Tensor.ones(vocab_size, embed_size) gt = Tensor.zeros(bs, seqlen, embed_size) Tensor.realize(idx, emb.weight, gt) GlobalCounters.reset() @@ -198,14 +198,14 @@ def test_embedding_backward_vocab_sharded(self): bs, seqlen = 4, 256 idx = Tensor.randint(bs, seqlen, high=vocab_size) emb = nn.Embedding(vocab_size, embed_size) - emb.weight = Tensor.ones(vocab_size, embed_size, requires_grad=True) + emb.weight = Tensor.ones(vocab_size, embed_size) gt = Tensor.zeros(bs, seqlen, embed_size) Tensor.realize(idx, emb.weight, gt) # compute expected grad on single device expected_grad = np.zeros((vocab_size, embed_size), dtype=np.float32) for i in idx.flatten().numpy(): expected_grad[i] += 2 # now shard the embedding weight on vocab axis and recompute - emb.weight = Tensor.ones(vocab_size, embed_size, requires_grad=True) + emb.weight = Tensor.ones(vocab_size, embed_size) emb.weight.shard_(devices, axis=0) idx = idx.shard(devices, axis=None) gt = gt.shard(devices, axis=None) @@ -215,12 +215,12 @@ def test_embedding_backward_vocab_sharded(self): np.testing.assert_allclose(emb.weight.grad.numpy(), expected_grad, rtol=1e-5, atol=1e-5) @unittest.skipUnless(Device.DEFAULT == "AMD" or (Device.DEFAULT == "NULL" and DEV.arch.startswith("gfx")), "tests AMD bf16 cast overhead") - def base_test_llama_8b_rope_backward(self, dtype): + def base_test_llama_8b_rope_backward(self, dtype, ops_scale=1): from extra.models.llama import precompute_freqs_cis, apply_rotary_emb bs, seqlen, dim, n_heads = 1, 512, 256, 4 head_dim = dim // n_heads x = Tensor.randn(bs, seqlen, dim, dtype=dtype) - wq = Tensor.randn(dim, dim, dtype=dtype, requires_grad=True) + wq = Tensor.randn(dim, dim, dtype=dtype) freqs_cis = precompute_freqs_cis(head_dim, seqlen).cast(dtype) Tensor.realize(x, wq, freqs_cis) xq = (x @ wq.T) @@ -232,15 +232,15 @@ def base_test_llama_8b_rope_backward(self, dtype): linear = compile_linear(wq.grad.schedule_linear()) assert len(linear.src) == 1, f"expected one kernel for backward, got: {len(linear.src)}" bwd_ops = estimate_uop(linear.src[0]).ops - # bfloat16 on non CDNA4 has ~10x ops overhead because of the software emulation - if dtype == dtypes.bfloat16 and not Device[Device.DEFAULT].renderer.target.arch.startswith("gfx950"): ops_scale = 10 - else: ops_scale = 1 expected_ops = bs*seqlen*dim*dim*ops_scale print(f"rope matmul bwd ({dtype}): {GlobalCounters.kernel_count} kernels, {bwd_ops:,} ops") self.assertLess(bwd_ops, expected_ops, f"rope bwd ops {bwd_ops:,} should be < {ops_scale} per (got {bwd_ops/(bs*seqlen*dim*dim):.1f})") - def test_llama_8b_rope_backward_f16(self): self.base_test_llama_8b_rope_backward(dtypes.float16) - def test_llama_8b_rope_backward_bf16(self): self.base_test_llama_8b_rope_backward(dtypes.bfloat16) + def test_llama_8b_rope_backward_f16(self): + self.base_test_llama_8b_rope_backward(dtypes.float16, ops_scale=2) + # bfloat16 on non CDNA4 has ~10x ops overhead because of the software emulation + def test_llama_8b_rope_backward_bf16(self): + self.base_test_llama_8b_rope_backward(dtypes.bfloat16, ops_scale=2 if Device[Device.DEFAULT].renderer.target.arch.startswith("gfx950") else 25) if __name__ == "__main__": unittest.main() diff --git a/test/backend/test_asm_gemm.py b/test/backend/test_asm_gemm.py index 6792040a6e576..18ac8cef7f250 100644 --- a/test/backend/test_asm_gemm.py +++ b/test/backend/test_asm_gemm.py @@ -1,10 +1,9 @@ import unittest from tinygrad import Tensor, Device, dtypes, Context -from tinygrad.device import is_dtype_supported from tinygrad.helpers import getenv, system, DEV from extra.gemm.cdna_asm_gemm import asm_gemm from test.helpers import needs_second_gpu -from examples.mlperf.models.flat_llama import FP8_DTYPE +from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8, FP8_MAX # On non CDNA4 it will only validate the Tensor.custom_kernel integration # Use DEV=NULL:HIP:gfx950 to also test the assembly @@ -12,33 +11,48 @@ def is_cdna4(): return Device[Device.DEFAULT].renderer.target.arch.startswith("g def run_asm_gemm(a_shape, b_shape, dtype=dtypes.float16, a_shard=None, b_shard=None, gpus:int=1) -> None: Tensor.manual_seed(0) - a_rand = Tensor.randn(a_shape, dtype=dtypes.float, requires_grad=False).sub(0.5).cast(dtype) - b_rand = Tensor.randn(b_shape, dtype=dtypes.float, requires_grad=False).sub(0.5).cast(dtype) + input_dtype = dtypes.bfloat16 if dtype == FP8_DTYPE else dtype + a_rand = Tensor.randn(a_shape, dtype=dtypes.float).sub(0.5).cast(input_dtype) + b_rand = Tensor.randn(b_shape, dtype=dtypes.float).sub(0.5).cast(input_dtype) with Context(DEBUG=0): Tensor.realize(a_rand, b_rand) devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(gpus)) if (multi:=gpus>1) else None - a, b = a_rand.clone().requires_grad_(), b_rand.clone().requires_grad_() + if dtype == FP8_DTYPE: + a_rand, x_scale, _ = quantize_fp8(a_rand) + b_rand, w_scale, _ = quantize_fp8(b_rand) + grad_amax_state = Tensor.full((), FP8_MAX, dtype=dtypes.float32, device=devs).contiguous() + with Context(DEBUG=0): + Tensor.realize(a_rand, x_scale, b_rand, w_scale, grad_amax_state) + + # clone all inputs before any backward: a clone copies the source's current .grad + a, b = a_rand.clone(), b_rand.clone() + if dtype == FP8_DTYPE: + a_ref, b_ref = a_rand.detach().cast(dtypes.bfloat16), b_rand.detach().cast(dtypes.bfloat16) + else: + a_ref, b_ref = a_rand.clone(), b_rand.clone() if multi: a, b = a.shard(devs, axis=a_shard), b.shard(devs, axis=b_shard) - tst = asm_gemm(a, b) + if dtype == FP8_DTYPE: + tst = asm_gemm(a, b, x_scale=x_scale, w_scale=w_scale, grad_amax_state=grad_amax_state) + else: + tst = asm_gemm(a, b) tst.sum().backward() Tensor.realize(tst, a.grad, b.grad) - a_ref, b_ref = a_rand.clone().requires_grad_(), b_rand.clone().requires_grad_() - # do reference gemm in bf16 for fp8, adjusting atol for quantization effects - if a_ref.dtype == FP8_DTYPE: - a_ref = a_ref.cast(dtypes.bfloat16) - b_ref = b_ref.cast(dtypes.bfloat16) if multi: a_ref, b_ref = a_ref.shard(devs, axis=a_shard), b_ref.shard(devs, axis=b_shard) - ref = a_ref @ b_ref + if dtype == FP8_DTYPE: + ref = ((a_ref @ b_ref) * x_scale * w_scale).cast(dtypes.bfloat16) + else: + ref = a_ref @ b_ref ref.sum().backward() Tensor.realize(ref, a_ref.grad, b_ref.grad) # no validation on the NULL device if a_rand.device.startswith("NULL"): return None atol, rtol = (2e-1, 1e-2) if dtype == dtypes.bfloat16 else (256, 1e-2) if dtype == FP8_DTYPE else (1e-2, 1e-3) - grad_atol, grad_rtol = (16895, 0.125) if dtype == FP8_DTYPE else (atol, rtol) + # allow more rtol for multi because of ALLREDUCE_CAST + grad_atol, grad_rtol = (16895, 0.125) if dtype == FP8_DTYPE else (atol, 2e-2 if multi else rtol) with Context(DEBUG=0): # enable for debugging, slow for larger gemms if getenv("USE_NPY"): @@ -70,7 +84,7 @@ def verify_asm_gemm_k_sharded_3d(batch:int, M:int, N:int, K:int, dtype=dtypes.fl # 128x smaller than usual # uses the UOp GEMM, runs on non CDNA4 and CI -@unittest.skipUnless(is_dtype_supported(dtypes.half), "need half") +@unittest.skipUnless(dtypes.half in Device[Device.DEFAULT].renderer.supported_dtypes(), "need half") class TestGemm(unittest.TestCase): def setUp(self): if is_cdna4(): self.skipTest("shapes are too small for the assembly GEMM") @@ -137,13 +151,19 @@ def setUp(self): def test_empty(self): asm_gemm(Tensor.empty(N:=getenv("N", 4096), N, dtype=self.dtype), Tensor.empty(N, N, dtype=self.dtype)).realize() def test_empty_bw(self): - x = Tensor.empty(1, N:=getenv("N", 4096), N, dtype=self.dtype, requires_grad=True) - y = Tensor.empty((N, N), dtype=self.dtype, requires_grad=True) - z = asm_gemm(x, y) + x = Tensor.empty(1, N:=getenv("N", 4096), N, dtype=self.dtype) + y = Tensor.empty((N, N), dtype=self.dtype) + if self.dtype == FP8_DTYPE: + x_scale = Tensor.empty((), dtype=dtypes.float32) + w_scale = Tensor.empty((), dtype=dtypes.float32) + grad_amax_state = Tensor.empty((), dtype=dtypes.float32).contiguous() + z = asm_gemm(x, y, x_scale=x_scale, w_scale=w_scale, grad_amax_state=grad_amax_state) + else: + z = asm_gemm(x, y) z.sum().backward() Tensor.realize(z, x.grad, y.grad) - # FP8 forward output is bf16, gradients use fp8e5m2 (aka bf8) - grad_dtype = dtypes.fp8e5m2 if self.dtype == FP8_DTYPE else self.dtype + # FP8 GEMM stores bf16 output and its backward produces bf16 gradients. + grad_dtype = dtypes.bfloat16 if self.dtype == FP8_DTYPE else self.dtype assert z.dtype == dtypes.bfloat16 assert x.grad.dtype == y.grad.dtype == grad_dtype diff --git a/test/backend/test_const_folding.py b/test/backend/test_const_folding.py index 1d62e6813eb0f..7c1f8aef92fc8 100644 --- a/test/backend/test_const_folding.py +++ b/test/backend/test_const_folding.py @@ -2,7 +2,6 @@ from tinygrad import Tensor, Device, dtypes from tinygrad.dtype import DTYPES_DICT from tinygrad.uop.ops import Ops, UOp -from tinygrad.device import is_dtype_supported import numpy as np from test.helpers import not_support_multi_device @@ -15,6 +14,11 @@ def _check_ast_count(desired_count:int, t:Tensor): #assert len(asts) == desired_count, f"{len(asts)} != {desired_count}" class TestMovedConstFolding(unittest.TestCase): + def test_contiguous_deviceless_const(self): + t = Tensor(UOp.const(dtypes.float, 2.0)).contiguous() + self.assertIs(t.uop.op, Ops.CONST) + self.assertIsNone(t.uop.device) + def test_add_shrunk_zero(self): _check_ast_count(0, Tensor([1.0, 2, 3, 4]) + Tensor.zeros(6).shrink(((1, 5),))) @@ -28,18 +32,16 @@ def test_add_padded_one(self): _check_ast_count(1, Tensor([1.0, 2, 3, 4]) * Tensor.ones(2).pad(((1, 1),))) def test_copy_padded_const(self): - schedule = Tensor.ones(4, device="CPU:0").pad(((1, 1),)).to("CPU:1").schedule_linear() + schedule = Tensor.ones(4, device="CPU:0", buffer=False).pad(((1, 1),)).to("CPU:1").schedule_linear() assert not any(si.src[0].op is Ops.COPY for si in schedule.src), "const copy should be folded" - np.testing.assert_equal(Tensor.ones(4, device="CPU:0").pad(((1, 1),)).to("CPU:1").numpy(), [0, 1, 1, 1, 1, 0]) + np.testing.assert_equal(Tensor.ones(4, device="CPU:0", buffer=False).pad(((1, 1),)).to("CPU:1").numpy(), [0, 1, 1, 1, 1, 0]) def test_cast_padded(self): # NOTE: it's always 1 kernel when calling .numpy, limitation of _check_ast_count - if is_dtype_supported(dtypes.int16): - _check_ast_count(1, Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int16)) - np.testing.assert_equal(Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int16).numpy(), [0, 1, 1, 1, 1, 0]) - if is_dtype_supported(dtypes.uint16): - _check_ast_count(1, Tensor.full(4, fill_value=-1).pad(((1, 1),)).cast(dtypes.uint16)) - np.testing.assert_equal(Tensor.full(4, fill_value=-1).pad(((1, 1),)).cast(dtypes.uint16).numpy(), [0, 65535, 65535, 65535, 65535, 0]) + _check_ast_count(1, Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int16)) + np.testing.assert_equal(Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int16).numpy(), [0, 1, 1, 1, 1, 0]) + _check_ast_count(1, Tensor.full(4, fill_value=-1).pad(((1, 1),)).cast(dtypes.uint16)) + np.testing.assert_equal(Tensor.full(4, fill_value=-1).pad(((1, 1),)).cast(dtypes.uint16).numpy(), [0, 65535, 65535, 65535, 65535, 0]) # folded _check_ast_count(1, Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int64)) np.testing.assert_equal(Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int64).numpy(), [0, 1, 1, 1, 1, 0]) @@ -109,7 +111,7 @@ def test_const_max(self): def test_sum_output_dtype(self): # sum output dtype can be different from input for dt in DTYPES_DICT.values(): - if is_dtype_supported(dt): + if dt in Device[Device.DEFAULT].renderer.supported_dtypes(): t = Tensor.ones(16, dtype=dt).reshape(4, 4) assert t.sum().dtype == t.contiguous().sum().dtype @@ -117,7 +119,7 @@ def test_sum_output_dtype(self): class TestMultiConstFolding(unittest.TestCase): def test_multi_const_folding_literal(self): ds = tuple(f"{Device.DEFAULT}:{i}" for i in range(4)) - t = Tensor.arange(16).float().to(ds).realize() + t = Tensor.arange(16).float().clone().to(ds).realize() # non const folding case creates one ast on each shard _check_ast_count(4, t + 1) @@ -142,7 +144,7 @@ def test_multi_const_folding_literal(self): def test_multi_const_folding_tensor(self): ds = tuple(f"{Device.DEFAULT}:{i}" for i in range(4)) - t = Tensor.arange(16).float().to(ds).realize() + t = Tensor.arange(16).float().clone().to(ds).realize() zero = Tensor.zeros(16).to(ds).realize() one = Tensor.ones(16).to(ds).realize() diff --git a/test/backend/test_custom_kernel.py b/test/backend/test_custom_kernel.py index 29328fe2179e8..1e68428d41257 100644 --- a/test/backend/test_custom_kernel.py +++ b/test/backend/test_custom_kernel.py @@ -1,5 +1,5 @@ import unittest -from tinygrad import Tensor, UOp, GlobalCounters, Context +from tinygrad import Tensor, UOp, GlobalCounters, Context, Device from tinygrad.dtype import AddrSpace, dtypes from tinygrad.uop.ops import KernelInfo, AxisType, Ops @@ -155,7 +155,7 @@ def test_arange(self): self.assertTrue((ref == tst).all().item()) def test_eye(self): - ref = Tensor.eye(1024).contiguous().realize() + ref = Tensor.eye(1024).clone().realize() tst = Tensor.empty_like(ref) tst = tst.custom_kernel(fxn=custom_eye_kernel)[0] self.assertTrue((ref == tst).all().item()) @@ -220,14 +220,14 @@ def test_gemm_backward(self, custom_backward_gemm=False): b_rand = Tensor.randn(8, N) Tensor.realize(a_rand, b_rand) - a, b = Tensor(a_rand.numpy(), requires_grad=True), Tensor(b_rand.numpy(), requires_grad=True) + a, b = Tensor(a_rand.numpy()), Tensor(b_rand.numpy()) c = Tensor.empty(N, N) tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm, grad_fxn=backward_gemm_custom if custom_backward_gemm else backward_gemm)[0] tst.sum().backward() grad_a, grad_b = a.grad, b.grad Tensor.realize(tst, grad_a, grad_b) - a, b = Tensor(a_rand.numpy(), requires_grad=True), Tensor(b_rand.numpy(), requires_grad=True) + a, b = Tensor(a_rand.numpy()), Tensor(b_rand.numpy()) ref = (a@b) ref.sum().backward() real_grad_a, real_grad_b = a.grad, b.grad @@ -284,6 +284,30 @@ def test_multi_after_schedule_order(self): self.assertIsNotNone(custom_idx, "custom_addmul kernel not found in schedule") self.assertEqual(custom_idx, 3, f"custom_addmul should be at index 3, got {custom_idx}") + def test_invalids_into_custom_kernel_no_empty_kernel(self): + from tinygrad.engine.realize import compile_linear + a = Tensor.full((4, 4), 3.).contiguous() + b = Tensor.full((4, 4), 2.).contiguous() + Tensor.realize(a, b) + out = Tensor.invalids(*a.shape, dtype=a.dtype) + out, *_ = Tensor.custom_kernel(out, a, b, fxn=custom_elementwise_add_kernel) + compiled = compile_linear(out.schedule_linear()) + for call in compiled.src: + prg = call.src[0] + if prg.op is not Ops.PROGRAM: continue + self.assertTrue(len(prg.arg.globals) > 0, f"empty kernel compiled (no globals): name={prg.arg.name}") + + @unittest.skipIf(Device.DEFAULT == "WEBGPU", "kernel timing not supported") + def test_invalids_into_custom_kernel_with_beam(self): + a = Tensor.full((4, 4), 3.).contiguous() + b = Tensor.full((4, 4), 2.).contiguous() + Tensor.realize(a, b) + with Context(BEAM=1, IGNORE_BEAM_CACHE=1): + out = Tensor.invalids(*a.shape, dtype=a.dtype) + out, *_ = Tensor.custom_kernel(out, a, b, fxn=custom_elementwise_add_kernel) + result = out.flatten().tolist() + self.assertTrue(all(x == 5 for x in result), f"expected all 5.0, got {result}") + @unittest.skip("what are anonymous buffers?") def test_anonymous_buffers_in_function(self): """Test that custom kernels with anonymous output buffers work inside @function.""" @@ -311,7 +335,7 @@ def run(x:Tensor, w:Tensor) -> Tensor: assert all(x == expected for x in result), f"expected all {expected}, got {result}" def test_custom_kernel_sched(self, use_custom=False): - x = Tensor.arange(32).reshape(8, 4).realize() + x = Tensor.arange(32).reshape(8, 4).clone().realize() y = Tensor.empty_like(x) y = Tensor.custom_kernel(y, x, fxn=custom_add_one_kernel)[0] if use_custom: @@ -328,7 +352,7 @@ def test_custom_kernel_sched_copy(self): self.test_custom_kernel_sched(use_custo @unittest.expectedFailure def test_sliced_buffer_function(self): - x = Tensor.arange(32).reshape(8, 4).realize() + x = Tensor.arange(32).reshape(8, 4).clone().realize() from tinygrad import function @function(precompile=True) def run(x:Tensor) -> Tensor: diff --git a/test/backend/test_dtype.py b/test/backend/test_dtype.py index 430bfa5db2405..cd63186864a75 100644 --- a/test/backend/test_dtype.py +++ b/test/backend/test_dtype.py @@ -2,8 +2,7 @@ import numpy as np import torch from typing import Any, List -from tinygrad.device import is_dtype_supported -from tinygrad.helpers import getenv, DEBUG, CI, EMULATED_DTYPES +from tinygrad.helpers import getenv, DEBUG, EMULATED_DTYPES, DEV from tinygrad.dtype import DType, DTYPES_DICT, least_upper_dtype, fp8_to_float, float_to_fp8, _to_np_dtype, _to_torch_dtype, truncate from tinygrad.renderer.ptx import PTXRenderer from tinygrad.renderer.nir import NIRRenderer @@ -17,11 +16,13 @@ settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False)) settings.load_profile("my_profile") +supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes() + def get_available_cast_dtypes(dtype: DType) -> List[DType]: - dts = [v for k, v in DTYPES_DICT.items() if v != dtype and is_dtype_supported(v) or v in dtypes.fp8s+(dtypes.half,dtypes.bfloat16,dtypes.long)] - if dtype in (dtypes.long, dtypes.ulong) and (not is_dtype_supported(dtype) or dtypes.long in EMULATED_DTYPES.tolist(dtypes)): + dts = [v for k, v in DTYPES_DICT.items() if v != dtype and v in supported_dtypes or v in dtypes.fp8s+(dtypes.half,dtypes.bfloat16,dtypes.long)] + if dtype in (dtypes.long, dtypes.ulong) and (dtype not in supported_dtypes or dtypes.long in EMULATED_DTYPES.tolist(dtypes)): return [dt for dt in dts if dt != dtypes.double] # can't bitcast with no 64-bit support - if not is_dtype_supported(dtype) and dtype not in dtypes.fp8s+(dtypes.half,dtypes.bfloat16): return [] + if dtype not in supported_dtypes and dtype not in dtypes.fp8s+(dtypes.half,dtypes.bfloat16): return [] return dts def _to_torch_storage_type(dtype:DType): @@ -60,7 +61,7 @@ class TestDType(unittest.TestCase): @classmethod def setUpClass(cls): if cls.DTYPE is None: raise unittest.SkipTest("base class") - cls.DATA = rand_for_dtype(cls.DTYPE, 0x10, allow_subnormal=is_dtype_supported(cls.DTYPE)) + cls.DATA = rand_for_dtype(cls.DTYPE, 0x10, allow_subnormal=cls.DTYPE in supported_dtypes) def test_to_np(self): _test_to_np(Tensor(self.DATA, dtype=self.DTYPE), _to_np_dtype(self.DTYPE), np.array(self.DATA, dtype=_to_np_dtype(self.DTYPE))) @@ -102,31 +103,6 @@ def test_uint_overflow(self): _test_to_np(Tensor(v, dtype=self.DTYPE)+2, _to_np_dtype(self.DTYPE), np.array(v, dtype=_to_np_dtype(self.DTYPE))+2) _test_to_np(Tensor(v, dtype=self.DTYPE)*2, _to_np_dtype(self.DTYPE), np.array(v, dtype=_to_np_dtype(self.DTYPE))*2) - def test_dtypes_DTYPES_DICT(self): - self.assertIn("float", DTYPES_DICT) - self.assertIn("float32", DTYPES_DICT) - self.assertEqual(len(DTYPES_DICT), 28) - self.assertTrue(all(isinstance(value, DType) for value in DTYPES_DICT.values())) - self.assertTrue(all(issubclass(_to_np_dtype(value), np.generic) for value in DTYPES_DICT.values() if _to_np_dtype(value) is not None)) - - def test_resulting_and_init_dtypes_match(self): - dtypes = list(map(np.dtype, ["bool", "uint8", "int8", "int16", "int32", "int64", "float32", "float64"])) - data = [1., 2., 0., 0.5, -1.5, 5.25] - for dt in dtypes: - arr = np.asarray(data).astype(dt) - tensor = Tensor(arr) - if not is_dtype_supported(tensor.dtype): continue - tin = tensor.numpy() - tor = torch.as_tensor(arr).detach().numpy() - assert dt == tin.dtype == tor.dtype, f"dtype mismatch: expected={dt} | tinygrad={tin.dtype} | torch={tor.dtype}" - np.testing.assert_allclose(tin, tor, atol=1e-6, rtol=1e-3) - - def test_finfo(self): - if self.DTYPE not in [dtypes.float16, dtypes.float32, dtypes.float64]: return - info = np.finfo(_to_np_dtype(self.DTYPE)) - self.assertEqual(info.bits, self.DTYPE.bitsize) - self.assertEqual((info.nexp, info.nmant), dtypes.finfo(self.DTYPE)) - def _test_ops(a_dtype:DType, b_dtype:DType, target_dtype=None): target_dtype = target_dtype or least_upper_dtype(a_dtype, b_dtype) if a_dtype == dtypes.bool or b_dtype == dtypes.bool: return @@ -134,13 +110,9 @@ def _test_ops(a_dtype:DType, b_dtype:DType, target_dtype=None): _assert_eq((Tensor([1], dtype=a_dtype).cast(b_dtype)+Tensor([1], dtype=a_dtype).cast(b_dtype)).cast(a_dtype), a_dtype, [2]) _assert_eq(Tensor([1,2,3,4], dtype=a_dtype)*Tensor([1,2,3,4], dtype=b_dtype), target_dtype, [1,4,9,16]) _assert_eq(Tensor([[1,2],[3,4]], dtype=a_dtype)@Tensor.eye(2, dtype=b_dtype), target_dtype, [[1,2],[3,4]]) - _assert_eq(Tensor([1,1,1,1], dtype=a_dtype)+Tensor.ones((4,4), dtype=b_dtype), target_dtype, 2*Tensor.ones(4,4).numpy()) - -class TestFp8s(unittest.TestCase): - def test_fp8e4m3_creation(self): assert Tensor([-1, 1, 2], dtype=dtypes.fp8e4m3).dtype == dtypes.fp8e4m3 - def test_fp8e5m2_creation(self): assert Tensor([-1, 1, 2], dtype=dtypes.fp8e5m2).dtype == dtypes.fp8e5m2 - def test_fp8e4m3fnuz_creation(self): assert Tensor([-1, 1, 2], dtype=dtypes.fp8e4m3fnuz).dtype == dtypes.fp8e4m3fnuz - def test_fp8e5m2fnuz_creation(self): assert Tensor([-1, 1, 2], dtype=dtypes.fp8e5m2fnuz).dtype == dtypes.fp8e5m2fnuz + _assert_eq(Tensor([1,1,1,1], dtype=a_dtype)+Tensor.ones((4,4), dtype=b_dtype), target_dtype, 2*np.ones((4,4))) + _assert_eq(Tensor([1,1,1,1], dtype=a_dtype)+Tensor.ones((4,4), dtype=b_dtype).clone(), target_dtype, 2*np.ones((4,4))) + _assert_eq(Tensor.ones((4,4), dtype=b_dtype).clone(), b_dtype, np.ones((4,4))) class TestFp8sConversions(unittest.TestCase): @given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E4M3_MAX, max_value=FP8E4M3_MAX)) @@ -191,25 +163,6 @@ def test_fp8e4m3fnuz_to_float(self, x): 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()) -class TestBFloat16(unittest.TestCase): - def test_bf16_creation_numpy(self): - data = [-1, 1, 2] - t = Tensor(data, dtype=dtypes.bfloat16) - assert t.dtype == dtypes.bfloat16 - tnp = t.numpy() - assert tnp.dtype == np.float32 - np.testing.assert_allclose(tnp, np.array(data)) - - def test_bf16_ones(self): - t = Tensor.ones(3, 5, dtype=dtypes.bfloat16) - assert t.dtype == dtypes.bfloat16 - np.testing.assert_allclose(t.numpy(), np.ones((3, 5))) - - def test_bf16_eye(self): - t = Tensor.eye(3, dtype=dtypes.bfloat16) - assert t.dtype == dtypes.bfloat16 - np.testing.assert_allclose(t.numpy(), np.eye(3)) - class TestBFloat16DType(unittest.TestCase): def test_bf16_to_float(self): _test_cast(Tensor([100000], dtype=dtypes.bfloat16), dtypes.float32) @@ -271,10 +224,10 @@ def test_float_to_uint(self): _test_op(lambda: Tensor([-0.9, -0.3, 1.2], dtype=dtypes.float32).cast(dtypes.uint32), dtypes.uint32, [0, 0, 1]) -@unittest.skipUnless(is_dtype_supported(dtypes.double), f"no double on {Device.DEFAULT}") +@unittest.skipUnless(dtypes.double in supported_dtypes, f"no double on {Device.DEFAULT}") class TestDoubleDType(TestDType): DTYPE = dtypes.double - @unittest.skipIf((CI and Device.DEFAULT in {"CUDA", "NV"}) or \ + @unittest.skipIf((DEV.interface.startswith("MOCK") and Device.DEFAULT in {"CUDA", "NV"}) or \ isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)), "conversion not supported on CI CUDA, PTX, and NIR") # TODO: why not? def test_float64_increased_precision(self): for func in [ @@ -418,48 +371,6 @@ def setUpClass(cls): @classmethod def tearDownClass(cls): cls.stack.close() -class TestPtrDType(unittest.TestCase): - def test_vec_double(self): - dt1 = dtypes.float.vec(4).ptr().vec(4) - dt2 = dtypes.float.vec(4).ptr().vec(4) - self.assertEqual(dt1, dt2) - self.assertEqual(str(dt1), str(dt2)) - - def test_scalar(self): - dt = dtypes.float.vec(4).ptr().scalar() - self.assertEqual(dt.base, dtypes.float.vec(4)) - - dt = dtypes.float.vec(4).ptr().vec(4).scalar() - self.assertEqual(dt.base, dtypes.float.vec(4)) - - dt = dtypes.float.vec(4).scalar() - self.assertEqual(dt, dtypes.float) - - def test_serialize(self): - dt = dtypes.float.vec(4).ptr().vec(4) - self.assertEqual(dt, eval(str(dt))) - - def test_vec_ptr_sz(self): - dt = dtypes.float.ptr(1024).vec(4) - self.assertEqual(dt, eval(str(dt))) - self.assertEqual(str(dt), "dtypes.float.ptr(1024).vec(4)") - - def test_vcount(self): - dt = dtypes.float.ptr().vec(4) - self.assertEqual(dt.vcount, 4) - self.assertEqual(dt.v, 4) - self.assertEqual(dt.count, 1) - - dt = dtypes.float.vec(4).ptr() - self.assertEqual(dt.vcount, 1) - self.assertEqual(dt.v, 1) - self.assertEqual(dt.count, 4) - - dt = dtypes.float.vec(4).ptr().vec(4) - self.assertEqual(dt.vcount, 4) - self.assertEqual(dt.v, 4) - self.assertEqual(dt.count, 4) - class TestImplicitFunctionTypeChange(unittest.TestCase): def test_functions(self): result = [] @@ -478,7 +389,7 @@ def test_functions(self): class TestTensorMethod(unittest.TestCase): @given(strat.sampled_from(core_dtypes)) def test_abs_diff(self, dt): - if dt == dtypes.bool or not is_dtype_supported(dt): return + if dt == dtypes.bool or dt not in supported_dtypes: return a, b = Tensor([2], dtype=dt), Tensor([1], dtype=dt) ret = (a - b).abs() np.testing.assert_allclose(ret.numpy(), np.abs(a.numpy()-b.numpy())) @@ -486,11 +397,11 @@ def test_abs_diff(self, dt): class TestDtypeUsage(unittest.TestCase): def test_max_w_alu(self): for d in dtypes.ints: - if is_dtype_supported(d): + if d in supported_dtypes: t = Tensor([[1, 2], [3, 4]], dtype=d) (t*t).max().item() -@unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), f"no bfloat16 on {Device.DEFAULT}") +@unittest.skipUnless(dtypes.bfloat16 in supported_dtypes, f"no bfloat16 on {Device.DEFAULT}") class TestOpsBFloat16(unittest.TestCase): def test_cast(self): # TODO: helper_test_op breaks in unrelated part diff --git a/test/backend/test_dtype_alu.py b/test/backend/test_dtype_alu.py index 8d61cdcd86ce5..71679a37a6894 100644 --- a/test/backend/test_dtype_alu.py +++ b/test/backend/test_dtype_alu.py @@ -1,12 +1,12 @@ import unittest, operator, math from tinygrad import Context, Tensor, dtypes, Device from tinygrad.dtype import DType, truncate, fp8_to_float -from tinygrad.helpers import CI, EMULATED_DTYPES, DEV, getenv +from tinygrad.helpers import EMULATED_DTYPES, DEV, getenv from tinygrad.tensor import _to_np_dtype -from tinygrad.device import is_dtype_supported from tinygrad.runtime.ops_python import from_storage_scalar from tinygrad.renderer.ptx import PTXRenderer from tinygrad.renderer.nir import NIRRenderer +from tinygrad.uop import Ops import numpy as np import pytest from hypothesis import assume, given, strategies as strat, settings @@ -37,6 +37,11 @@ unary_operations.remove((Tensor.sin, np.sin)) unary_operations.remove((Tensor.cos, np.cos)) +# transcendental isn't accurate enough +if Ops.SQRT not in Device[Device.DEFAULT].renderer.code_for_op: unary_operations.remove((Tensor.sqrt, np.sqrt)) + +supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes() + class ht: float64 = strat.floats(width=64, allow_subnormal=False) float32 = strat.floats(width=32, allow_subnormal=False) @@ -67,7 +72,7 @@ def universal_test(a, b, dtype, op): numpy_value = truncate[dtype](op[1](ta.numpy(), tb.numpy()).item()) else: tensor_value, numpy_value = (op[0](ta, tb)).numpy(), op[1](ta.numpy(), tb.numpy()) if dtype in dtypes.floats: - if not is_dtype_supported(dtype) or dtype in EMULATED_DTYPES.tolist(dtypes): # denormals are zero + if dtype not in supported_dtypes or dtype in EMULATED_DTYPES.tolist(dtypes): # denormals are zero fe, fm = dtypes.finfo(dtype) atol, rtol = 2 ** (2 - (1 << (fe - 1))), 2 ** (-fm) else: atol, rtol = {dtypes.bfloat16:(1e-3, 1e-2), dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2:(1.0, 5e-1), @@ -83,7 +88,8 @@ def universal_test_unary(a, dtype, op): if op[0] == Tensor.log and a <= 0: return if dtype in dtypes.fp8s: # denormals are zero - if dtype in EMULATED_DTYPES.tolist(dtypes) or not is_dtype_supported(dtype) and abs(ta.numpy().item()) < 0.015625: return + if (dtype in EMULATED_DTYPES.tolist(dtypes) or dtype not in supported_dtypes + and abs(ta.numpy().item()) < 0.015625): return tensor_value = fp8_to_float(op[0](ta.realize()).bitcast(dtypes.uint8).item(), dtype) numpy_value = truncate[dtype](v:=op[1](ta.numpy()).item()) # cuda cast f32 inf to f8 MAX, amd cast it to nan(E4M3)/inf(E5M2) @@ -105,6 +111,7 @@ def universal_test_cast(a, in_dtype, dtype): def universal_test_midcast(a, b, c, op1, op2, d1:DType, d2:DType): if not isinstance(op1, tuple): op1 = (op1, op1) if not isinstance(op2, tuple): op2 = (op2, op2) + if op1[0] == operator.mod and b == 0: return # lt and max with nan is undefined in tinygrad if op1[0] in (operator.lt, Tensor.maximum) and (math.isnan(a) or math.isnan(b)): return if op2[0] in (operator.lt, Tensor.maximum) and math.isnan(c): return @@ -115,14 +122,14 @@ def universal_test_midcast(a, b, c, op1, op2, d1:DType, d2:DType): np.testing.assert_allclose(tensor_value, numpy_value, rtol=1e-6 if isinstance(Device[Device.DEFAULT].renderer, PTXRenderer) else 1e-7) class TestDTypeALU(unittest.TestCase): - @unittest.skipUnless(is_dtype_supported(dtypes.float64), f"no float64 on {Device.DEFAULT}") + @unittest.skipUnless(dtypes.float64 in supported_dtypes, f"no float64 on {Device.DEFAULT}") @given(ht.float64, ht.float64, strat.sampled_from(binary_operations)) def test_float64(self, a, b, op): universal_test(a, b, dtypes.float64, op) @given(ht.float32, ht.float32, strat.sampled_from(binary_operations)) def test_float32(self, a, b, op): universal_test(a, b, dtypes.float32, op) - @unittest.skipUnless(is_dtype_supported(dtypes.float16), f"no float16 on {Device.DEFAULT}") + @unittest.skipUnless(dtypes.float16 in supported_dtypes, f"no float16 on {Device.DEFAULT}") @given(ht.float16, ht.float16, strat.sampled_from(binary_operations)) def test_float16(self, a, b, op): universal_test(a, b, dtypes.float16, op) @@ -130,17 +137,17 @@ def test_float16(self, a, b, op): universal_test(a, b, dtypes.float16, op) @Context(EMULATED_DTYPES="half") def test_emulated_float16(self, a, b, op): universal_test(a, b, dtypes.float16, op) - @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), f"no bfloat16 on {Device.DEFAULT}") + @unittest.skipUnless(dtypes.bfloat16 in supported_dtypes, f"no bfloat16 on {Device.DEFAULT}") @given(ht.bfloat16, ht.bfloat16, strat.sampled_from(binary_operations)) def test_bfloat16(self, a, b, op): - universal_test(from_storage_scalar(a, dtypes.bfloat16), from_storage_scalar(a, dtypes.bfloat16), dtypes.bfloat16, op) + universal_test(from_storage_scalar(a, dtypes.bfloat16), from_storage_scalar(b, dtypes.bfloat16), dtypes.bfloat16, op) @given(ht.bfloat16, ht.bfloat16, strat.sampled_from(binary_operations)) @Context(EMULATED_DTYPES="bfloat16") def test_emulated_bfloat16(self, a, b, op): - universal_test(from_storage_scalar(a, dtypes.bfloat16), from_storage_scalar(a, dtypes.bfloat16), dtypes.bfloat16, op) + universal_test(from_storage_scalar(a, dtypes.bfloat16), from_storage_scalar(b, dtypes.bfloat16), dtypes.bfloat16, op) - @unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3), f"no fp8e4m3 on {Device.DEFAULT}") + @unittest.skipUnless(dtypes.fp8e4m3 in supported_dtypes, f"no fp8e4m3 on {Device.DEFAULT}") @given(ht.fp8e4m3, ht.fp8e4m3, strat.sampled_from(binary_operations)) def test_fp8e4m3(self, a, b, op): universal_test(from_storage_scalar(a, dtypes.fp8e4m3), from_storage_scalar(b, dtypes.fp8e4m3), dtypes.fp8e4m3, op) @@ -150,7 +157,7 @@ def test_fp8e4m3(self, a, b, op): def test_emulated_fp8e4m3(self, a, b, op): universal_test(from_storage_scalar(a, dtypes.fp8e4m3), from_storage_scalar(b, dtypes.fp8e4m3), dtypes.fp8e4m3, op) - @unittest.skipUnless(is_dtype_supported(dtypes.fp8e5m2), f"no fp8e5m2 on {Device.DEFAULT}") + @unittest.skipUnless(dtypes.fp8e5m2 in supported_dtypes, f"no fp8e5m2 on {Device.DEFAULT}") @given(ht.fp8e5m2, ht.fp8e5m2, strat.sampled_from(binary_operations)) def test_fp8e5m2(self, a, b, op): universal_test(from_storage_scalar(a, dtypes.fp8e5m2), from_storage_scalar(b, dtypes.fp8e5m2), dtypes.fp8e5m2, op) @@ -160,12 +167,12 @@ def test_fp8e5m2(self, a, b, op): def test_emulated_fp8e5m2(self, a, b, op): universal_test(from_storage_scalar(a, dtypes.fp8e5m2), from_storage_scalar(b, dtypes.fp8e5m2), dtypes.fp8e5m2, op) - @unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3fnuz), f"no fp8e4m3fnuz on {Device.DEFAULT}") + @unittest.skipUnless(dtypes.fp8e4m3fnuz in supported_dtypes, f"no fp8e4m3fnuz on {Device.DEFAULT}") @given(ht.fp8e4m3fnuz, ht.fp8e4m3fnuz, strat.sampled_from(binary_operations)) def test_fp8e4m3fnuz(self, a, b, op): universal_test(from_storage_scalar(a, dtypes.fp8e4m3fnuz), from_storage_scalar(b, dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op) - @unittest.skipUnless(is_dtype_supported(dtypes.fp8e5m2fnuz), f"no fp8e5m2fnuz on {Device.DEFAULT}") + @unittest.skipUnless(dtypes.fp8e5m2fnuz in supported_dtypes, f"no fp8e5m2fnuz on {Device.DEFAULT}") @given(ht.fp8e5m2fnuz, ht.fp8e5m2fnuz, strat.sampled_from(binary_operations)) def test_fp8e5m2fnuz(self, a, b, op): universal_test(from_storage_scalar(a, dtypes.fp8e5m2fnuz), from_storage_scalar(b, dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op) @@ -183,7 +190,7 @@ def test_emulated_fp8e5m2fnuz(self, a, b, op): @given(ht.float32, strat.sampled_from(unary_operations)) def test_float32_unary(self, a, op): universal_test_unary(a, dtypes.float32, op) - @unittest.skipUnless(is_dtype_supported(dtypes.float16), f"no float16 on {Device.DEFAULT}") + @unittest.skipUnless(dtypes.float16 in supported_dtypes, f"no float16 on {Device.DEFAULT}") @given(ht.float16, strat.sampled_from(unary_operations)) def test_float16_unary(self, a, op): universal_test_unary(a, dtypes.float16, op) @@ -191,7 +198,7 @@ def test_float16_unary(self, a, op): universal_test_unary(a, dtypes.float16, op) @Context(EMULATED_DTYPES="half") def test_emulated_float16_unary(self, a, op): universal_test_unary(a, dtypes.float16, op) - @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), f"no bfloat16 on {Device.DEFAULT}") + @unittest.skipUnless(dtypes.bfloat16 in supported_dtypes, f"no bfloat16 on {Device.DEFAULT}") @given(ht.bfloat16, strat.sampled_from(unary_operations)) def test_bfloat16_unary(self, a, op): universal_test_unary(from_storage_scalar(a, dtypes.bfloat16), dtypes.bfloat16, op) @@ -199,7 +206,7 @@ def test_bfloat16_unary(self, a, op): universal_test_unary(from_storage_scalar(a @Context(EMULATED_DTYPES="bfloat16") def test_emulated_bfloat16_unary(self, a, op): universal_test_unary(from_storage_scalar(a, dtypes.bfloat16), dtypes.bfloat16, op) - @unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3), f"no fp8e4m3 on {Device.DEFAULT}") + @unittest.skipUnless(dtypes.fp8e4m3 in supported_dtypes, f"no fp8e4m3 on {Device.DEFAULT}") @given(ht.fp8e4m3, strat.sampled_from(unary_operations)) def test_fp8e4m3_unary(self, a, op): if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e4m3) != 0.0) @@ -211,7 +218,7 @@ def test_emulated_fp8e4m3_unary(self, a, op): if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e4m3) != 0.0) universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e4m3), dtypes.fp8e4m3, op) - @unittest.skipUnless(is_dtype_supported(dtypes.fp8e5m2), f"no fp8e5m2 on {Device.DEFAULT}") + @unittest.skipUnless(dtypes.fp8e5m2 in supported_dtypes, f"no fp8e5m2 on {Device.DEFAULT}") @given(ht.fp8e5m2, strat.sampled_from(unary_operations)) def test_fp8e5m2_unary(self, a, op): if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2) != 0.0) @@ -223,13 +230,13 @@ def test_emulated_fp8e5m2_unary(self, a, op): if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2) != 0.0) universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2), dtypes.fp8e5m2, op) - @unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3fnuz), f"no fp8e4m3fnuz on {Device.DEFAULT}") + @unittest.skipUnless(dtypes.fp8e4m3fnuz in supported_dtypes, f"no fp8e4m3fnuz on {Device.DEFAULT}") @given(ht.fp8e4m3fnuz, strat.sampled_from(unary_operations)) def test_fp8e4m3fnuz_unary(self, a, op): if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz) != 0.0) universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op) - @unittest.skipUnless(is_dtype_supported(dtypes.fp8e5m2fnuz), f"no fp8e5m2fnuz on {Device.DEFAULT}") + @unittest.skipUnless(dtypes.fp8e5m2fnuz in supported_dtypes, f"no fp8e5m2fnuz on {Device.DEFAULT}") @given(ht.fp8e5m2fnuz, strat.sampled_from(unary_operations)) def test_fp8e5m2fnuz_unary(self, a, op): if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz) != 0.0) @@ -250,15 +257,15 @@ def test_emulated_fp8e5m2fnuz_unary(self, a, op): @given(ht.uint8, ht.uint8, strat.sampled_from(integer_binary_operations)) def test_uint8(self, a, b, op): universal_test(a, b, dtypes.uint8, op) - @unittest.skipUnless(is_dtype_supported(dtypes.uint16), f"no uint16 on {Device.DEFAULT}") + @unittest.skipUnless(dtypes.uint16 in supported_dtypes, f"no uint16 on {Device.DEFAULT}") @given(ht.uint16, ht.uint16, strat.sampled_from(integer_binary_operations)) def test_uint16(self, a, b, op): universal_test(a, b, dtypes.uint16, op) - @unittest.skipUnless(is_dtype_supported(dtypes.uint32), f"no uint32 on {Device.DEFAULT}") + @unittest.skipUnless(dtypes.uint32 in supported_dtypes, f"no uint32 on {Device.DEFAULT}") @given(ht.uint32, ht.uint32, strat.sampled_from(integer_binary_operations)) def test_uint32(self, a, b, op): universal_test(a, b, dtypes.uint32, op) - @unittest.skipUnless(is_dtype_supported(dtypes.uint64), f"no uint64 on {Device.DEFAULT}") + @unittest.skipUnless(dtypes.uint64 in supported_dtypes, f"no uint64 on {Device.DEFAULT}") @given(ht.uint64, ht.uint64, strat.sampled_from(integer_binary_operations)) def test_uint64(self, a, b, op): universal_test(a, b, dtypes.uint64, op) @@ -287,15 +294,15 @@ def test_emulated_int64(self, a, b, op): universal_test(a, b, dtypes.int64, op) @given(ht.uint8, strat.sampled_from(integer_unary_operations)) def test_uint8_unary(self, a, op): universal_test_unary(a, dtypes.uint8, op) - @unittest.skipUnless(is_dtype_supported(dtypes.uint16), f"no uint16 on {Device.DEFAULT}") + @unittest.skipUnless(dtypes.uint16 in supported_dtypes, f"no uint16 on {Device.DEFAULT}") @given(ht.uint16, strat.sampled_from(integer_unary_operations)) def test_uint16_unary(self, a, op): universal_test_unary(a, dtypes.uint16, op) - @unittest.skipUnless(is_dtype_supported(dtypes.uint32), f"no uint32 on {Device.DEFAULT}") + @unittest.skipUnless(dtypes.uint32 in supported_dtypes, f"no uint32 on {Device.DEFAULT}") @given(ht.uint32, strat.sampled_from(integer_unary_operations)) def test_uint32_unary(self, a, op): universal_test_unary(a, dtypes.uint32, op) - @unittest.skipUnless(is_dtype_supported(dtypes.uint64), f"no uint64 on {Device.DEFAULT}") + @unittest.skipUnless(dtypes.uint64 in supported_dtypes, f"no uint64 on {Device.DEFAULT}") @given(ht.uint64, strat.sampled_from(integer_unary_operations)) def test_uint64_unary(self, a, op): universal_test_unary(a, dtypes.uint64, op) @@ -324,12 +331,12 @@ def test_emulated_int64_unary(self, a, op): universal_test_unary(a, dtypes.int64 @given(ht.bool, ht.bool, strat.sampled_from(((operator.add, operator.add), (operator.mul, operator.mul)))) def test_bool(self, a, b, op): universal_test(a, b, dtypes.bool, op) - @unittest.skipIf(not CI and Device.DEFAULT == "METAL", "broken on local M3") @given(ht.int32, ht.int32, ht.float32, strat.sampled_from(integer_binary_operations), strat.sampled_from(binary_operations)) def test_int32_midcast_float(self, a, b, c, op1, op2): universal_test_midcast(a, b, c, op1, op2, dtypes.int32, dtypes.float32) - # Metal and CUDA and HIP and NIR behave differently than numpy in CI for overflows - skip_overflow = (CI and Device.DEFAULT in {"AMD", "NV", "CUDA"}) or isinstance(Device[Device.DEFAULT].renderer, NIRRenderer) + # Metal and (MOCK)CUDA and HIP and NIR behave differently than numpy for overflows + skip_overflow = ((DEV.interface.startswith("MOCK") and Device.DEFAULT in {"AMD", "NV", "CUDA"}) + or isinstance(Device[Device.DEFAULT].renderer, NIRRenderer)) @given(strat.floats(width=32, min_value=0, max_value=10.0) if skip_overflow else ht.float32, strat.floats(width=32, min_value=0, max_value=10.0) if skip_overflow else ht.float32, ht.int32, strat.sampled_from(binary_operations), strat.sampled_from(integer_binary_operations)) @@ -348,21 +355,21 @@ def test_int32_cast(self, a, dtype): universal_test_cast(a, dtypes.int32, dtype) @given(strat.floats(width=32, min_value=1.0, max_value=254.0, allow_subnormal=False), strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16))) def test_float_cast_to_unsigned(self, a, float_dtype, unsigned_dtype): - if not is_dtype_supported(float_dtype): float_dtype = dtypes.float32 + if float_dtype not in supported_dtypes: float_dtype = dtypes.float32 universal_test_cast(a, float_dtype, unsigned_dtype) @unittest.skip("relied on hacks") @given(strat.floats(width=32, min_value=256.0, max_value=65000.0, allow_subnormal=False), strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16))) def test_float_cast_to_unsigned_overflow(self, a, float_dtype, unsigned_dtype): - if not is_dtype_supported(float_dtype): float_dtype = dtypes.float32 + if float_dtype not in supported_dtypes: float_dtype = dtypes.float32 universal_test_cast(a, float_dtype, unsigned_dtype) @unittest.skip("relied on hacks") @given(strat.floats(width=32, min_value=-65000.0, max_value=-1.0, allow_subnormal=False), strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16))) def test_float_cast_to_unsigned_underflow(self, a, float_dtype, unsigned_dtype): - if not is_dtype_supported(float_dtype): float_dtype = dtypes.float32 + if float_dtype not in supported_dtypes: float_dtype = dtypes.float32 universal_test_cast(a, float_dtype, unsigned_dtype) @unittest.expectedFailure diff --git a/test/backend/test_edgecases.py b/test/backend/test_edgecases.py index 53b56392b167e..89bd245ae3d56 100644 --- a/test/backend/test_edgecases.py +++ b/test/backend/test_edgecases.py @@ -127,19 +127,19 @@ def test_negative_weight_decay(self): with self.assertRaises(ValueError): torch.optim.AdamW([torch.tensor([1.], requires_grad=True)], lr=0.1, weight_decay=-0.1) with self.assertRaises(ValueError): - nn.optim.AdamW([Tensor([1.], requires_grad=True)], lr=0.1, weight_decay=-0.1) + nn.optim.AdamW([Tensor([1.])], lr=0.1, weight_decay=-0.1) def test_negative_lr(self): with self.assertRaises(ValueError): torch.optim.SGD([torch.tensor([1.], requires_grad=True)], lr=-0.1) with self.assertRaises(ValueError): - nn.optim.SGD([Tensor([1.], requires_grad=True)], lr=-0.1) + nn.optim.SGD([Tensor([1.])], lr=-0.1) def test_negative_momentum(self): with self.assertRaises(ValueError): torch.optim.SGD([torch.tensor([1.], requires_grad=True)], lr=0.1, momentum=-0.1) with self.assertRaises(ValueError): - nn.optim.SGD([Tensor([1.], requires_grad=True)], lr=0.1, momentum=-0.1) + nn.optim.SGD([Tensor([1.])], lr=0.1, momentum=-0.1) class TestZeroFolding(unittest.TestCase): # we don't need more of these @@ -219,7 +219,6 @@ def test_float_floordiv_tensor(self): class TestEdgeCases(unittest.TestCase): # add tests exposing new and diverse kinds of bugs that might impact real users here - @unittest.expectedFailure def test_circular_pad_negative(self): # negative pads with circular mode should wrap like PyTorch arr = np.arange(9).reshape(1, 1, 3, 3).astype(np.float32) diff --git a/test/backend/test_encodings.py b/test/backend/test_encodings.py new file mode 100644 index 0000000000000..2c9a939c1f7b0 --- /dev/null +++ b/test/backend/test_encodings.py @@ -0,0 +1,149 @@ +import unittest +from tinygrad import Device +from tinygrad.uop.ops import UOp, Ops +from tinygrad.dtype import dtypes +from tinygrad.renderer.isa.x86 import X86Ops, X86Renderer, RBP, RDI, RSP, RSI, RAX, RDX, XMM, GPR, imm, def_reg + +def ins(op, dt, src, tag=None): return UOp(Ops.INS, arg=op, dtype=dt, src=src, tag=tag) + +@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "only on x86") +class TestEncodingsX86(unittest.TestCase): + # NOTE: x86 supports a single displacement as memory address and index without base memory address + # these have no use cases so they aren't supported + def encode(self, u:UOp): return Device[Device.DEFAULT].renderer.render([u]) + + # displacement of 0 isn't emitted + def test_base_address(self): + load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RDI), UOp(Ops.NOOP), imm(dtypes.int8, 0)), RDI) + # mov edi, dword ptr [rdi] + self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 3F")) + + # rsp/r12 require a sib byte when used as base memory address + def test_rsp_base_address(self): + load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RSP), UOp(Ops.NOOP), imm(dtypes.int8, 0)), RSP) + # mov esp, dword ptr [rsp] + self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 24 24")) + + # rbp/r13 require a displacement when used as base memory address + def test_rbp_base_address(self): + load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RBP), UOp(Ops.NOOP), imm(dtypes.int8, 0)), RBP) + # mov ebp, dword ptr [rbp + 0] + self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 6D 00")) + + # test [base + index*scale] + def test_base_index_address(self): + load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RAX), def_reg(dtypes.int32, RDX), imm(dtypes.int8, 0)), RAX) + # mov eax, dword ptr [rax + rdx*4] + self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 04 90")) + + # rsp as index means no index + def test_rsp_index_address(self): + load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RAX), def_reg(dtypes.int32, RSP), imm(dtypes.int8, 0)), RAX) + # mov eax, dword ptr [rax] + self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 00")) + + # however r12 is a valid index + def test_r12_index_address(self): + load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RAX), def_reg(dtypes.int32, GPR[12]), imm(dtypes.int8, 0)), RAX) + # mov eax, dword ptr [rax + r12*4] + self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("42 8B 04 A0")) + + # test [base + index*scale + 8bit disp] + def test_complex_address_8bit_disp(self): + load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int8, 10)), RDI) + # mov edi, dword ptr [rdi + rsi*4 + 0xa] + self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 7C B7 0A")) + + # test [base + index*scale + 32bit disp] + def test_complex_address_32bit_disp(self): + load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int32, 10000)), RDI) + # mov edi, dword ptr [rdi + rsi*4 + 0x2710] + self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B BC B7 10 27 00 00")) + + # 8bit variants of legacy instructions subtract 1 from opcode + def test_8bit_legacy_encoding(self): + cast = ins(X86Ops.MOVSX, dtypes.int32, (def_reg(dtypes.int8, RDX),), RAX) + # movsx eax, dl + self.assertEqual(bytes.fromhex(self.encode(cast)), bytes.fromhex("0F BE C2")) + + # accessing lower 8 bits of rsp, rbp, rsi, rdi requires rex prefix + def test_lower_8bits_reg(self): + cast = ins(X86Ops.MOVSX, dtypes.int32, (def_reg(dtypes.int8, RDI),), RAX) + # movsx eax, dil + self.assertEqual(bytes.fromhex(self.encode(cast)), bytes.fromhex("40 0F BE C7")) + + # test 16 bit variant of legacy instruction + def test_16bit_legacy_encoding(self): + cast = ins(X86Ops.MOVSX, dtypes.int16, (def_reg(dtypes.int8, RDX),), RAX) + # movsx ax, dl + self.assertEqual(bytes.fromhex(self.encode(cast)), bytes.fromhex("66 0F BE C2")) + + # test 64 bit variant of legacy instruction + def test_64bit_legacy_encoding(self): + cast = ins(X86Ops.MOVSX, dtypes.int64, (def_reg(dtypes.int8, RDX),), RAX) + # movsx rax, dl + self.assertEqual(bytes.fromhex(self.encode(cast)), bytes.fromhex("48 0F BE C2")) + + # test compact vex encoding + def test_compact_vex_encoding(self): + xmm0, xmm1 = def_reg(dtypes.float32, XMM[0]), def_reg(dtypes.float32, XMM[1]) + add = ins(X86Ops.VADDSS, dtypes.float32, (xmm0, xmm1), XMM[0]) + # vaddss xmm0, xmm0, xmm1 + self.assertEqual(bytes.fromhex(self.encode(add)), bytes.fromhex("C5 FA 58 C1")) + + # test long vex encoding + def test_long_vex_encoding(self): + xmm0, xmm8 = def_reg(dtypes.float32, XMM[0]), def_reg(dtypes.float32, XMM[8]) + add = ins(X86Ops.VADDSS, dtypes.float32, (xmm0, xmm8), XMM[0]) + # vaddss xmm0, xmm0, xmm8 + self.assertEqual(bytes.fromhex(self.encode(add)), bytes.fromhex("C4 C1 7A 58 C0")) + + # test ymm encoding + def test_ymm_encoding(self): + xmm0, xmm1 = def_reg(dtypes.float32.vec(8), XMM[0]), def_reg(dtypes.float32.vec(8), XMM[1]) + add = ins(X86Ops.VADDPS, dtypes.float32.vec(8), (xmm0, xmm1), XMM[0]) + # vaddps ymm0, ymm0, ymm1 + self.assertEqual(bytes.fromhex(self.encode(add)), bytes.fromhex("C5 FC 58 C1")) + + # test encoding where register is in the immediate field + def test_reg_in_imm_field(self): + xmm0, xmm1, xmm2 = def_reg(dtypes.float32, XMM[0]), def_reg(dtypes.float32, XMM[1]), def_reg(dtypes.float32, XMM[2]) + blend = ins(X86Ops.VBLENDVPS, dtypes.float32, (xmm0, xmm1, xmm2), XMM[0]) + # vblendvps xmm0, xmm0, xmm1, xmm2 + self.assertEqual(bytes.fromhex(self.encode(blend)), bytes.fromhex("C4 E3 79 4A C1 20")) + + # when writting to mem the uop takes the store form where dtype is void and there's no definition + def test_write_mem(self): + base, index, disp = def_reg(dtypes.int32.ptr(), RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int8, 10) + xmm0 = def_reg(dtypes.float32, XMM[0]) + extr = ins(X86Ops.VPEXTRD, dtypes.void, (base, index, disp, xmm0, imm(dtypes.uint8, 0))) + # vpextrd dword ptr [rdi + rsi*4 + 0xa], xmm0, 0 + self.assertEqual(bytes.fromhex(self.encode(extr)), bytes.fromhex("C4 E3 79 16 44 B7 0A 00")) + + # test two address instruction with fused load works + def test_two_address_load(self): + base, index, disp = def_reg(dtypes.int32.ptr(), RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int8, 10) + cmove = ins(X86Ops.CMOVE, dtypes.int32, (base, index, disp), RAX) + # cmove eax, dword ptr [rdi + rsi*4 + 0xa] + self.assertEqual(bytes.fromhex(self.encode(cmove)), bytes.fromhex("0F 44 44 B7 0A")) + + # test instruction where displacement and imm have the same value + def test_disp_imm_same_value(self): + base, index, disp = def_reg(dtypes.int8.ptr(), RDI), def_reg(dtypes.int8, RSI), imm(dtypes.int8, 10) + mov = ins(X86Ops.MOVi, dtypes.void, (base, index, disp, disp)) + # mov byte ptr [rdi + rsi + 0xa], 0xa + self.assertEqual(bytes.fromhex(self.encode(mov)), bytes.fromhex("40 C6 44 37 0A 0A")) + + base, index, disp = def_reg(dtypes.int32.ptr(), RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int32, 10) + imul = ins(X86Ops.IMULi, dtypes.int32, (base, index, disp) + (imm(dtypes.int32, 10),), RDI) + # imul edi, dword ptr [rdi + rsi*4 + 0xa], 0xa + self.assertEqual(bytes.fromhex(self.encode(imul)), bytes.fromhex("69 BC B7 0A 00 00 00 0A 00 00 00")) + + # cmoves have the cmp as the last src even though it is not explicitly used, the cmp doesn't define a reg and is ignored in the encoding + def test_cmove_ignore_cmp(self): + cmove = ins(X86Ops.CMOVE, dtypes.int32, (def_reg(dtypes.int32, RAX), UOp(Ops.INS, arg=X86Ops.CMP)), RDX) + # cmove edx, eax + self.assertEqual(bytes.fromhex(self.encode(cmove)), bytes.fromhex("0F 44 D0")) + +if __name__ == "__main__": + unittest.main() \ No newline at end of file diff --git a/test/backend/test_interop.py b/test/backend/test_interop.py index acf5072958fc3..18346212cf57e 100644 --- a/test/backend/test_interop.py +++ b/test/backend/test_interop.py @@ -1,9 +1,9 @@ #!/usr/bin/env python -import unittest +import unittest, os import torch import numpy as np -from tinygrad.helpers import CI, DEV +from tinygrad.helpers import DEV from tinygrad.tensor import Tensor from tinygrad.device import Device from tinygrad.dtype import _from_torch_dtype, _to_torch_dtype @@ -27,7 +27,7 @@ def test_torch_interop(self): tg_out = tg_data[:, :, 0] * 0.2989 + tg_data[:, :, 1] * 0.5870 + tg_data[:, :, 2] * 0.1140 tg_res = tg_out.numpy() - if self.torch_device == "mps" and CI: + if self.torch_device == "mps" and os.getenv("CI", "") != "": # MPS backend out of memory: https://discuss.pytorch.org/t/mps-back-end-out-of-memory-on-github-action/189773 # Calculate expected value on cpu. inp = inp.cpu() diff --git a/test/backend/test_isel.py b/test/backend/test_isel.py new file mode 100644 index 0000000000000..2c3cede802583 --- /dev/null +++ b/test/backend/test_isel.py @@ -0,0 +1,147 @@ +import unittest +from typing import cast +from tinygrad import Device +from tinygrad.uop import Ops +from tinygrad.uop.ops import UOp, dtypes, graph_rewrite +from tinygrad.renderer.isa.x86 import X86Renderer, X86Ops +from tinygrad.renderer.isa import IselContext + +@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "only x86") +class TestIselX86(unittest.TestCase): + def isel_rewrite(self, x:UOp): + return graph_rewrite(x, cast(X86Renderer, Device[Device.DEFAULT].renderer).isel_matcher, IselContext(x), bottom_up=True) + + def _check_op(self, dt_op, expr): + nargs = expr.__code__.co_argcount + for dt,op in dt_op: + with self.subTest(dtype=dt): + v = [UOp.variable(str(i), 0, 0, dt) for i in range(nargs)] + n = self.isel_rewrite(expr(*v)) + self.assertIs(n.arg, op) + + def test_cmove(self): + a = UOp.variable("a", 0, 0, dtypes.int32) + b = UOp.variable("b", 0, 0, dtypes.int32) + c = (a < b).where(a, b) + d = (a != b).where(a, b) + f = c + d + n = self.isel_rewrite(f) + self.assertTrue(n.src[0].arg is X86Ops.CMOVL and n.src[1].arg is X86Ops.CMOVNE) + # both comparisons become the same instruction + self.assertTrue(n.src[0].src[2] == n.src[1].src[2] and n.src[0].src[2].arg is X86Ops.CMP) + + def test_vmax(self): + dt_op = [(dtypes.float32, X86Ops.VMAXSS), (dtypes.float64, X86Ops.VMAXSD), + (dtypes.float32.vec(4), X86Ops.VMAXPS), (dtypes.float64.vec(4), X86Ops.VMAXPD)] + self._check_op(dt_op, lambda a,b: (a < b).where(b, a)) + + def test_vmin(self): + dt_op = [(dtypes.float32, X86Ops.VMINSS), (dtypes.float64, X86Ops.VMINSD), + (dtypes.float32.vec(4), X86Ops.VMINPS), (dtypes.float64.vec(4), X86Ops.VMINPD)] + self._check_op(dt_op, lambda a,b: (a < b).where(a, b)) + + def test_vfmadd(self): + dt_op = [(dtypes.float32, X86Ops.VFMADD213SS), (dtypes.float64, X86Ops.VFMADD213SD), + (dtypes.float32.vec(4), X86Ops.VFMADD213PS), (dtypes.float64.vec(4), X86Ops.VFMADD213PD)] + self._check_op(dt_op, lambda a,b,c: a * b + c) + + # don't use fmadd if op being fused (mul) is used multiple times + def test_no_vfmadd(self): + dt_op = [(dtypes.float32, X86Ops.VADDSS), (dtypes.float64, X86Ops.VADDSD), + (dtypes.float32.vec(4), X86Ops.VADDPS), (dtypes.float64.vec(4), X86Ops.VADDPD)] + self._check_op(dt_op, lambda a,b: a * b + a * b) + + def test_vpbroadcast(self): + a = UOp.variable("a", 0, 0, dtypes.int32) + n = self.isel_rewrite(a.broadcast(4)) + # need to move src from gpr to xmm before broadcasting + self.assertTrue(n.arg is X86Ops.VPBROADCASTD and n.src[0].arg is X86Ops.VMOVD) + # if we can fuse a load we can skip the move and access memory directly + load = UOp.param(0, dtypes.int32.ptr()).index(UOp.const(dtypes.int32, 0), ptr=True).load() + n = self.isel_rewrite(load.broadcast(4)) + self.assertTrue(n.arg is X86Ops.VPBROADCASTD and len(n.src) == 3) + + def test_vbroadcastss(self): + a = UOp.variable("a", 0, 0, dtypes.float32) + valid = [UOp.vectorize(a, a, a, a), UOp.vectorize(a, a, a, a, a, a, a, a)] + for shuf in valid: self.assertIs(self.isel_rewrite(shuf).arg, X86Ops.VBROADCASTSS) + + def test_vshufps(self): + a = UOp.variable("a", 0, 0, dtypes.float32.vec(8)) + b = UOp.variable("b", 0, 0, dtypes.float32.vec(8)) + c = UOp.variable("c", 0, 0, dtypes.float32) + d = UOp.variable("d", 0, 0, dtypes.float32) + + valid = [UOp.vectorize(c, c, d, d), + UOp.vectorize(a.gep(0), a.gep(1), c, c), + UOp.vectorize(a.gep(0), a.gep(1), b.gep(2), b.gep(3)), + UOp.vectorize(a.gep(1), a.gep(2), a.gep(3), a.gep(0)), + UOp.vectorize(a.gep(3), a.gep(2), a.gep(1), a.gep(0), a.gep(7), a.gep(6), a.gep(5), a.gep(4)), + UOp.vectorize(a.gep(0), a.gep(0), b.gep(1), b.gep(1), a.gep(4), a.gep(4), b.gep(5), b.gep(5))] + for shuf in valid: self.assertIs(self.isel_rewrite(shuf).arg, X86Ops.VSHUFPS) + + invalid = [UOp.vectorize(a.gep(0), a.gep(1), b.gep(4), b.gep(5)), + UOp.vectorize(a.gep(0), a.gep(5), b.gep(2), b.gep(3)), + UOp.vectorize(a.gep(0), a.gep(0), a.gep(0), a.gep(0), a.gep(4), a.gep(4), a.gep(4), a.gep(5)), + UOp.vectorize(a.gep(0), a.gep(0), b.gep(0), b.gep(0), a.gep(4), a.gep(4), b.gep(4), a.gep(4))] + for shuf in invalid: self.assertIsNot(self.isel_rewrite(shuf).arg, X86Ops.VSHUFPS) + + def test_vshufpd(self): + a = UOp.variable("a", 0, 0, dtypes.float64.vec(4)) + b = UOp.variable("b", 0, 0, dtypes.float64.vec(4)) + c = UOp.variable("c", 0, 0, dtypes.float64) + d = UOp.variable("d", 0, 0, dtypes.float64) + + valid = [UOp.vectorize(c, d), + UOp.vectorize(a.gep(0), c), + UOp.vectorize(a.gep(1), b.gep(1)), + UOp.vectorize(a.gep(0), b.gep(1), a.gep(2), b.gep(3)), + UOp.vectorize(a.gep(1), a.gep(1), a.gep(3), a.gep(3))] + for shuf in valid: self.assertIs(self.isel_rewrite(shuf).arg, X86Ops.VSHUFPD) + + invalid = [UOp.vectorize(c, c, c, c), + UOp.vectorize(a.gep(0), a.gep(1), b.gep(2), b.gep(3)), + UOp.vectorize(a.gep(2), b.gep(3), a.gep(2), b.gep(3)), + UOp.vectorize(a.gep(0), b.gep(1), a.gep(0), b.gep(1))] + for shuf in invalid: self.assertIsNot(self.isel_rewrite(shuf).arg, X86Ops.VSHUFPD) + + def test_vinsertps(self): + a = UOp.variable("a", 0, 0, dtypes.float32.vec(4)) + b = UOp.variable("b", 0, 0, dtypes.float32.vec(4)) + c = UOp.variable("c", 0, 0, dtypes.float32.vec(4)) + d = UOp.variable("e", 0, 0, dtypes.float32) + # moving 0th element to position 0 does nothing so only 1 vinsertps is generated + n = self.isel_rewrite(UOp.vectorize(a.gep(0), d)) + self.assertIs(n.arg, X86Ops.VINSERTPS) + self.assertIsNot(n.src[0].arg, X86Ops.VINSERTPS) + + valid = [UOp.vectorize(a.gep(0), b.gep(1), a.gep(2), b.gep(3)), + UOp.vectorize(a.gep(3), b.gep(2), c.gep(1), d)] + for shuf in valid: self.assertIs(self.isel_rewrite(shuf).arg, X86Ops.VINSERTPS) + + # complex address is [base + index*scale + displacement] + def test_complex_address(self): + a = UOp.variable("a", 0, 0, dtypes.int32) + load = UOp.param(0, dtypes.int32.ptr()).index(a + 1, ptr=True).load() + n = self.isel_rewrite(load) + # displacement is the constant in "a" scaled to the buffer element size, dtype is int8 when the value fits otherwise int32 + self.assertTrue(n.src[2].op is Ops.CONST and n.src[2].dtype is dtypes.int8 and n.src[2].arg == 4) + + def test_fold_load(self): + load1 = UOp.param(0, dtypes.int32.ptr()).index(UOp.const(dtypes.int32, 0), ptr=True).load() + load2 = UOp.param(0, dtypes.int32.ptr()).index(UOp.const(dtypes.int32, 1), ptr=True).load() + n = self.isel_rewrite(load1 + load2) + self.assertTrue(len(n.src) == 4) + + # don't fold when used multiple times + def test_dont_fold_load(self): + load = UOp.param(0, dtypes.int32.ptr()).index(UOp.const(dtypes.int32, 0), ptr=True).load() + # used by multiple users + n = self.isel_rewrite(load + 1 + load) + self.assertTrue(len(n.src) == 2) + # used mutiple times by same user + n = self.isel_rewrite(load * load) + self.assertTrue(len(n.src) == 2) + +if __name__ == "__main__": + unittest.main() diff --git a/test/backend/test_jit.py b/test/backend/test_jit.py index d955a77544b2a..733fbc4e00701 100644 --- a/test/backend/test_jit.py +++ b/test/backend/test_jit.py @@ -10,8 +10,9 @@ from tinygrad.device import Device from tinygrad.helpers import Context, JIT, DEV, GlobalCounters from tinygrad.dtype import dtypes -from tinygrad.uop.ops import Ops +from tinygrad.uop.ops import Ops, UOp from extra.models.unet import ResBlock +from tinygrad.renderer.isa.x86 import X86Renderer def _simple_test(add, extract=lambda x: x, N=10): for _ in range(5): @@ -48,7 +49,7 @@ def f(x, st): y = (x + 1).contiguous().realize() z = x.shrink(((st, st + N),)).contiguous().realize() return y, z - x = Tensor.arange(2*N).contiguous().realize() + x = Tensor.arange(2*N).clone().realize() for _ in range(3): y, z = f(x, Variable("a", 0, N).bind(0)) self.assertEqual(y.shape, (2*N,)) self.assertEqual(z.shape, (N,)) @@ -91,7 +92,7 @@ def test_jit_input_view(self): @TinyJit def f(x): return (x[2:5].contiguous() + 1).realize() for i in range(5): - x = (Tensor.arange(10).float() + i * 10).contiguous().realize() + x = (Tensor.arange(10).float() + i * 10).clone().realize() np.testing.assert_allclose(f(x).numpy(), x.numpy()[2:5] + 1) def test_jit_multiple_outputs(self): @@ -106,6 +107,7 @@ def f(a, b): return (a+b).realize(), (a-b).realize(), (a*b).realize() np.testing.assert_allclose(e.numpy(), a.numpy()*b.numpy(), atol=1e-4, rtol=1e-5) assert_jit_cache_len(f, 3) + @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "estimates are wrong for x86") def test_global_counters_jit(self): @TinyJit def f(a, b): @@ -533,6 +535,18 @@ def f(x:Tensor) -> Tensor: return (x + 1).realize() f(Tensor(2.0)).item() # self.assertEqual(f(Tensor([2.0])).item(), 1.0) # TODO: wrong output, should be 3.0. currently depends on empty value + def test_jit_const_input(self): + @TinyJit + def f(x:Tensor) -> Tensor: return (x + 1).realize() + with self.assertRaises(JitError): + f(Tensor(UOp.const(dtypes.float, 2.0))).item() + + def test_jit_deviceless_compute_input(self): + @TinyJit + def f(x:Tensor) -> Tensor: return (x + 1).realize() + with self.assertRaises(JitError): + f(Tensor(UOp.const(dtypes.float, 2.0) + UOp.const(dtypes.float, 1.0))).item() + def test_jit_init_empty_alt(self): @TinyJit def f(a:Tensor, b:Tensor) -> Tensor: return b.assign(a+1) diff --git a/test/backend/test_kernel_cache.py b/test/backend/test_kernel_cache.py index a4f0f2193a912..d4dc23d3d6263 100644 --- a/test/backend/test_kernel_cache.py +++ b/test/backend/test_kernel_cache.py @@ -8,10 +8,10 @@ def test_kernel_cache_in_action(self): if Device.DEFAULT not in ["CPU"]: self.skipTest("No custom kernel cache is implemented") - unique_const = 0.6765677269 + const_value = 0.6765677269 a = Tensor.rand(4,4).realize() b = Tensor.rand(4,4).realize() - x = a + b + unique_const + x = a + b + const_value x.realize() a1 = Tensor.rand(4,4).realize() @@ -20,7 +20,7 @@ def test_kernel_cache_in_action(self): Device['CPU'].compiler.compile_cached = None # making it not callable try: - x1 = a1 + b1 + unique_const + x1 = a1 + b1 + const_value x1.realize() # Same kernel should be from cache. finally: Device['CPU'].compiler.compile_cached = orig_compile_func diff --git a/test/backend/test_linearizer.py b/test/backend/test_linearizer.py index 89132f3ea265a..b89c1f94da9de 100644 --- a/test/backend/test_linearizer.py +++ b/test/backend/test_linearizer.py @@ -3,7 +3,7 @@ from tinygrad.codegen.opt import Opt, OptOps from tinygrad.uop.ops import UOp, Ops, GroupOp, AxisType, buffers -from tinygrad.device import Device, Buffer, is_dtype_supported +from tinygrad.device import Device, Buffer from tinygrad.tensor import Tensor, _to_np_dtype from tinygrad.engine.realize import run_linear from tinygrad.codegen import to_program @@ -11,11 +11,13 @@ from tinygrad.dtype import DType, dtypes, PtrDType, AddrSpace from tinygrad.renderer.ptx import PTXRenderer from tinygrad.renderer.cstyle import CUDARenderer +from tinygrad.renderer.isa import ISARenderer from test.helpers import replace_opts MOCKGPU = DEV.interface.startswith("MOCK") from tinygrad.uop.render import print_uops # noqa: F401 # pylint: disable=unused-import +@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, ISARenderer), "isa backends don't preserve the op spec when lowering") class TestLinearizer(unittest.TestCase): def test_arg_dedup(self): # NOTE: this realize exists because Tensor.numpy calls .contiguous() internally @@ -204,7 +206,7 @@ def test_zero_fold(self): def test_sum_acc_dtype(self): for tensor_dtype, acc_dtype in ( (dtypes.bool, dtypes.int), (dtypes.int16, dtypes.int), (dtypes.float16, dtypes.float), (dtypes.bfloat16, dtypes.float)): - if is_dtype_supported(tensor_dtype) and is_dtype_supported(acc_dtype): + if tensor_dtype in (dts:=Device[Device.DEFAULT].renderer.supported_dtypes()) and acc_dtype in dts: a = Tensor([1, 2, 3], dtype=tensor_dtype).sum() realized_ast = a.schedule_linear().src[-1].src[0] program = to_program(replace_opts(realized_ast, []), renderer=Device[Device.DEFAULT].renderer) @@ -227,7 +229,7 @@ def helper_arg_acc_dtype(c: Tensor, expected_dtype:DType): (dtypes.float, dtypes.float16, dtypes.float16), ) for tensor_dtype, acc_dtype, expected_dtype in tests: - if is_dtype_supported(tensor_dtype) and is_dtype_supported(acc_dtype) and is_dtype_supported(expected_dtype): + if tensor_dtype in (dts:=Device[Device.DEFAULT].renderer.supported_dtypes()) and acc_dtype in dts and expected_dtype in dts: a, b = Tensor.rand(8, 8, dtype=tensor_dtype), Tensor.rand(8, 8, dtype=tensor_dtype) helper_arg_acc_dtype(a.sum(dtype=acc_dtype), expected_dtype) helper_arg_acc_dtype(a.matmul(b, dtype=acc_dtype), expected_dtype) @@ -316,6 +318,7 @@ def helper(t, max_ops=0): #helper(Tensor.arange(256), max_ops=2) helper(Tensor.arange(255), max_ops=2) + @unittest.skip("test implicitly depends on certain optimizations") @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4") @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken on ptx for some reason") def test_grouped_store_phis(self): @@ -338,6 +341,7 @@ def test_grouped_store_phis(self): for val in store_vals: assert val.dtype == dtypes.float.vec(4) # and val.op is not Ops.VECTORIZE + @unittest.skip("test implicitly depends on certain optimizations") @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4") def test_grouped_store_values(self): x = Tensor.randn((4,3,6,6)).realize() @@ -369,6 +373,7 @@ def get_recursive(uop): return set.union(set(uop.src), [uop], *[get_recursive(v) # assert barrier.src == tuple(local_stores) assert len([u for u in uops if u.op is Ops.IF]) + @unittest.skip("test implicitly depends on certain optimizations") @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals") @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared") @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4") diff --git a/test/backend/test_linearizer_dumb.py b/test/backend/test_linearizer_dumb.py index da5160d4e6845..fa4651679d56b 100644 --- a/test/backend/test_linearizer_dumb.py +++ b/test/backend/test_linearizer_dumb.py @@ -11,16 +11,16 @@ class TestLinearizerFailure(unittest.TestCase): @unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL") def test_failure_beam_mnist(self): - c0 = UOp(Ops.PARAM, dtypes.uchar.ptr(4014080), arg=0, src=()) + c0 = UOp.param(0, dtypes.uchar.ptr(4014080)) c1 = UOp.range(UOp.const(dtypes.weakint, 512), 0, AxisType.GLOBAL) c2 = UOp.range(UOp.const(dtypes.weakint, 784), 1, AxisType.GLOBAL) c3 = UOp.range(UOp.const(dtypes.weakint, 10), 3, AxisType.GLOBAL) - c4 = UOp(Ops.PARAM, dtypes.int.ptr(512), arg=1, src=()) + c4 = UOp.param(1, dtypes.int.ptr(512)) c5 = c4.index(c1.valid(UOp.const(dtypes.bool, True))) c6 = UOp.range(UOp.const(dtypes.weakint, 6000), 1004, AxisType.REDUCE) c7 = UOp.range(UOp.const(dtypes.weakint, 3750), 2006, AxisType.REDUCE) c8 = UOp.range(UOp.const(dtypes.weakint, 16), 2007, AxisType.GROUP_REDUCE) - c9 = UOp(Ops.PARAM, dtypes.uchar.ptr(47040000), arg=2, src=()) + c9 = UOp.param(2, dtypes.uchar.ptr(47040000)) c10 = c9.index((((c3*UOp.const(dtypes.weakint, 4704000))+c2)+(c6*UOp.const(dtypes.weakint, 784))).valid(UOp.const(dtypes.bool, True))) c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(dtypes.weakint, 6000))+c6)+((c7*UOp.const(dtypes.weakint, 16))+c8)).alu(Ops.CMPLT, UOp.const(dtypes.weakint, 59999)).where(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(dtypes.int, -1))).where(UOp.const(dtypes.uchar, 0), c10).reduce(c6, arg=Ops.ADD) c12 = c0.index((((c1*UOp.const(dtypes.weakint, 7840))+(c2*UOp.const(dtypes.weakint, 10)))+c3).valid(UOp.const(dtypes.bool, True))).store(c11).end(c1, c2, c3) diff --git a/test/backend/test_llama_kernels.py b/test/backend/test_llama_kernels.py new file mode 100644 index 0000000000000..06a27a9e4a208 --- /dev/null +++ b/test/backend/test_llama_kernels.py @@ -0,0 +1,96 @@ +import unittest +from tinygrad import Tensor, Device, dtypes, Context, GlobalCounters +from tinygrad.helpers import getenv +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 test.helpers import needs_second_gpu + +def run_fused_ce(bs:int, seqlen:int, vocab:int, label_smoothing:float=0.0) -> None: + Tensor.manual_seed(0) + logits_rand = Tensor.randn(bs, seqlen, vocab).cast(dtypes.bfloat16) + targets = Tensor.randint(bs, seqlen, high=vocab, dtype=dtypes.int32) + logits, logits_ref = logits_rand.clone(), logits_rand.detach().float().contiguous() + with Context(DEBUG=0): + Tensor.realize(logits, logits_ref, targets) + + loss = fused_ce_loss(logits, targets, label_smoothing=label_smoothing) + loss.backward() + Tensor.realize(loss, logits.grad) + + ref = logits_ref.sparse_categorical_crossentropy(targets, label_smoothing=label_smoothing) + ref.backward() + Tensor.realize(ref, logits_ref.grad) + + assert logits.grad.shape == (bs, seqlen, vocab) + with Context(DEBUG=0): + assert loss.allclose(ref, atol=2e-3, rtol=2e-3).item(), "forward mismatch" + assert logits.grad.allclose(logits_ref.grad, atol=2e-3, rtol=2e-3).item(), "grad mismatch" + +class TestFusedCE(unittest.TestCase): + def setUp(self): + if dtypes.bfloat16 not in Device[Device.DEFAULT].renderer.supported_dtypes(): self.skipTest("need bfloat16") + + def test_fused_ce_1_2_16(self): run_fused_ce(1, 2, 16, label_smoothing=0.2) + def test_fused_ce_2_16_128(self): run_fused_ce(2, 16, 128) + def test_fused_ce_4_128_1024(self): run_fused_ce(4, 128, 1024, label_smoothing=0.2) + + # note: this is the shape used in llama 8b + #def test_fused_ce_smoothing_16_1024_128256(self): run_fused_ce(16, 1024, 128256, label_smoothing=0.2) + +def run_quantize_fp8(shape:tuple[int, ...], delayed:bool=True) -> None: + Tensor.manual_seed(0) + x = Tensor.randn(*shape).cast(dtypes.bfloat16).contiguous() + amax_state = Tensor.full((), 2.0, dtype=dtypes.float32).contiguous() + with Context(DEBUG=0): Tensor.realize(x, amax_state) + + if delayed: + fp8, inv_scale, new_amax, _ = quantize_fp8_delayed(x, amax_state, FP8_DTYPE) + ref_fp8, ref_inv_scale, ref_new_amax = quantize_fp8(x, amax_state=amax_state) + Tensor.realize(fp8, inv_scale, new_amax) + Tensor.realize(ref_fp8, ref_inv_scale, ref_new_amax) + else: + fp8 = quantize_fp8_scalar(x, amax_state, FP8_DTYPE) + ref_fp8, _, _ = quantize_fp8(x, amax_state=amax_state) + Tensor.realize(fp8) + Tensor.realize(ref_fp8) + + with Context(DEBUG=0): + assert fp8.cast(dtypes.float).allclose(ref_fp8.cast(dtypes.float), atol=0, rtol=0).item(), "fp8 mismatch" + if delayed: + assert inv_scale.allclose(ref_inv_scale, atol=0, rtol=0).item(), "inv_scale mismatch" + assert new_amax.allclose(ref_new_amax, atol=0, rtol=0).item(), \ + f"amax mismatch: got={new_amax.item()} ref={ref_new_amax.item()} diff={abs(new_amax.item()-ref_new_amax.item())}" + +class TestQuantizeFP8(unittest.TestCase): + def setUp(self): + ren = Device[Device.DEFAULT].renderer + if dtypes.bfloat16 not in ren.supported_dtypes(): self.skipTest("need bfloat16") + if not ren.has_local or not ren.has_shared: self.skipTest("need local/shared") + + def test_scalar(self): run_quantize_fp8((getenv("N", 1024), 32), delayed=False) + def test_delayed(self): run_quantize_fp8((getenv("N", 2048), 1024)) + + @needs_second_gpu + def test_multi(self): + devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(8)) + x = Tensor.empty(2048*8, 1024, dtype=dtypes.bfloat16, device=devs).uop.multi(0) + x = Tensor(x, device=devs) + amax_state = Tensor.full((), 2.0, dtype=dtypes.float32, device=devs).contiguous() + fp8, _, new_amax, _ = quantize_fp8_delayed(x, amax_state, FP8_DTYPE) + Tensor.realize(fp8, new_amax) + assert fp8.uop.shape == x.uop.shape + assert new_amax.shape == () + +class TestLocalAmax(unittest.TestCase): + def test_multi_tensor_local_shard_amax(self): + devices = ("CPU:0", "CPU:1") + x = Tensor.arange(16, device=devices[0]).reshape(4, 4).cast(dtypes.float).contiguous().realize().shard(devices, axis=0).realize() + GlobalCounters.reset() + out = (x * local_abs_max(x)).contiguous().realize() + self.assertEqual(GlobalCounters.kernel_count, 4) + self.assertEqual(out.tolist(), [[0., 7., 14., 21.], [28., 35., 42., 49.], [120., 135., 150., 165.], [180., 195., 210., 225.]]) + +if __name__ == '__main__': + unittest.main() diff --git a/test/backend/test_multitensor.py b/test/backend/test_multitensor.py index c8a580fff0a7e..dd03d04cbddc9 100644 --- a/test/backend/test_multitensor.py +++ b/test/backend/test_multitensor.py @@ -1,6 +1,5 @@ import unittest, random from tinygrad import Tensor, Device, nn, GlobalCounters, TinyJit, dtypes, Variable -from tinygrad.device import is_dtype_supported from tinygrad.uop.ops import Ops, UOp from tinygrad.helpers import getenv, prod, Context from tinygrad.nn.state import get_parameters, get_state_dict @@ -98,6 +97,11 @@ def _test_shard_op(self, op, out, n=4): self.assertEqual(r.tolist(), out) def test_shard_reshape(self): self._test_shard_op(lambda t:t.reshape(2, 2), [[1.,1.],[1.,1.]]) def test_shard_elementwise(self): self._test_shard_op(lambda t:(t+t).reshape(2, 2), [[2.,2.],[2.,2.]]) + def test_alu_deviceless_const(self): + s = Tensor([1.0, 2, 3, 4]).shard((f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1"), axis=0) + np.testing.assert_equal((s + Tensor(UOp.const(dtypes.float, 1.0))).numpy(), [2, 3, 4, 5]) + np.testing.assert_equal((s + Tensor(UOp.const(dtypes.float, 1.0)).reshape((1,)).expand((4,))).numpy(), [2, 3, 4, 5]) + def test_shard_reduce(self): self._test_shard_op(lambda t:t.reshape(2, 3).sum(axis=1), [3.,3.], n=6) self._test_shard_op(lambda t:t.reshape(2, 3).sum(axis=0), [2.,2.,2.], n=6) @@ -117,7 +121,7 @@ def test_tensor_from_multi(self): _ = Tensor(X.uop, dtype=dtypes.float) def test_sharded_arange(self): - sharded_arange = Tensor.arange(1000).shard(devices_2, 0) + sharded_arange = Tensor.arange(1000).clone().shard(devices_2, 0) sharded_arange.realize() np.testing.assert_equal(sharded_arange.numpy(), np.arange(1000)) @@ -183,7 +187,7 @@ def test_elementwise_dtype(self): np.testing.assert_allclose(O.numpy(), X.numpy()[0:2]*W.numpy()[0:2] < 2) def test_shrink_on_shard_axis(self): - X = Tensor.arange(4*4).reshape(4,4).realize() + X = Tensor.arange(4*4).reshape(4,4).clone().realize() X_np = X.numpy() X.shard_(devices_2, 0) # only shrink on the device that owns the shard, this is enabled by the mselect simplifier @@ -230,12 +234,33 @@ def test_multiple_to_single_device(self): for ring in (0, 2): GlobalCounters.reset() with Context(RING=ring, SCACHE=0): - t = Tensor.arange(32).contiguous().shard(devices_4, 0).to(Device.DEFAULT) + t = Tensor.arange(32).clone().shard(devices_4, 0).to(Device.DEFAULT) t.realize() kernel_counts[ring] = GlobalCounters.kernel_count self.assertEqual(t.device, Device.DEFAULT) np.testing.assert_equal(t.numpy(), np.arange(32)) - self.assertNotEqual(kernel_counts[0], kernel_counts[2]) + self.assertEqual(kernel_counts[0], kernel_counts[2]) + + def test_to_single_device_gather_memory(self): + nrows, ncols = 64, 1024 + nbytes = nrows*ncols*4 + for devs in (devices_2, devices_4): + ndev = len(devs) + for axis in (0, 1): + sh = Tensor.arange(nrows*ncols).reshape(nrows, ncols).clone().shard(devs, axis).realize() + kernels, mem = {}, {} + for ring in (0, 2): + GlobalCounters.reset() + with Context(RING=ring, SCACHE=0): + t = sh.to(Device.DEFAULT) + t.realize() + kernels[ring], mem[ring] = GlobalCounters.kernel_count, GlobalCounters.global_mem + self.assertEqual(t.device, Device.DEFAULT) + np.testing.assert_equal(t.numpy(), np.arange(nrows*ncols).reshape(nrows, ncols)) + self.assertEqual(kernels[0], kernels[2]) + self.assertEqual(mem[0], mem[2]) + self.assertLess(kernels[0], 2*ndev) + self.assertLessEqual(mem[0], 4*nbytes) def test_allreduce_all2all(self): with Context(ALL2ALL=2): @@ -268,7 +293,7 @@ def test_multitensor_jit_input(self): @TinyJit def f(x): return (x+1).contiguous().sum() for _ in range(5): - tt = Tensor.arange(0, 4).contiguous().realize().shard((d1,d2), 0).realize() + tt = Tensor.arange(0, 4).clone().realize().shard((d1,d2), 0).realize() out = f(tt) assert out.item() == 1+2+3+4 @@ -284,7 +309,7 @@ def test_multitensor_inside_jit(self): @TinyJit def f(x): return (x.shard((d1,d2), 0)+1).contiguous().sum() for _ in range(5): - tt = Tensor.arange(0, 4).contiguous().realize() + tt = Tensor.arange(0, 4).clone().realize() out = f(tt) assert out.item() == 1+2+3+4 @@ -372,7 +397,7 @@ def test_backprop_conv_wino(self): def test_backward_sum(self): x = Tensor([[1.,2,3,4], [5,6,7,8]]).shard(devices_2, axis=0) - w = Tensor([1.,2,3,4], requires_grad=True).shard(devices_2) + w = Tensor([1.,2,3,4]).shard(devices_2) out = x * w out.mean().backward() tst = w.grad.numpy() @@ -404,7 +429,6 @@ def test_embedding_backward(self, shard_weight_axis=None): B, T, embed_size, vocab_size = 4, 10, 20, 28 layer = nn.Embedding(vocab_size, embed_size) - layer.weight.requires_grad = True x = Tensor(np.random.randint(0, vocab_size, (B, T), dtype=np.int32)) z = layer(x) z.sum().backward() @@ -412,7 +436,6 @@ def test_embedding_backward(self, shard_weight_axis=None): layer_sharded = nn.Embedding(vocab_size, embed_size) layer_sharded.weight.replace(layer.weight.shard(devices_2, axis=shard_weight_axis)).realize() - layer_sharded.weight.requires_grad = True x_sharded = x.shard(devices_2, axis=None) z_shard = layer_sharded(x_sharded) z_shard.sum().backward() @@ -475,7 +498,7 @@ def test_data_parallel_resnet(self): def _test_model_train_step(self, m, fake_image, labels): from tinygrad.nn.optim import LARS - optimizer = LARS(get_parameters(m), 0.1) # set requires_grad for all params + optimizer = LARS(get_parameters(m), 0.1) optimizer.zero_grad() m.load_from_pretrained() @@ -554,7 +577,7 @@ def jf(a, b) -> Tensor: def test_multi_tensor_jit_graph_assign_updates_each_shard(self): @TinyJit def jf(out: Tensor) -> Tensor: - tmp = (Tensor.arange(4, dtype=dtypes.float).shard(devices_2, 0) + 1).contiguous().realize() + tmp = (Tensor.arange(4, dtype=dtypes.float).clone().shard(devices_2, 0) + 1).contiguous().realize() out.assign((tmp + 1).contiguous()).realize() return out @@ -746,11 +769,6 @@ def test_rand_like_on_shard(self, axis=None): t2.realize() def test_rand_like_on_shard_axis(self): self.test_rand_like_on_shard(0) - def test_rand_like_on_shard_axis_requires_grad(self): - t = Tensor.empty((16, 16)).shard(devices_2, axis=0) - self.assertIs(t.rand_like(requires_grad=True).requires_grad, True) - self.assertIs(t.rand_like(requires_grad=False).requires_grad, False) - def test_rand_like_from_alu(self): a = Tensor.ones(4, 4).shard(devices_4, axis=0) aa = a + a @@ -803,9 +821,9 @@ def test_full_like_on_shard(self, axis=None): t2.realize() def test_full_like_on_shard_axis(self): self.test_full_like_on_shard(0) - def test_full_like_shrink_on_shard_axis(self): + def test_const_like_shrink_on_shard_axis(self): t = Tensor.ones(16, 16, dtype=dtypes.int).shard(devices_2, axis=0) - out = Tensor.full_like(t, 2)[:, :8] + out = t.const_like(2)[:, :8] linear, var_vals = out.linear_with_vars() self.assertEqual(len(linear.src), 0) run_linear(linear, var_vals) @@ -817,7 +835,7 @@ def test_dropout_on_shard(self): output = X.dropout(0.5).numpy() unique, counts = np.unique(output, return_counts=True) assert set(unique) == {0, 2}, unique - assert 100 < counts[0] < 156, counts[0] + assert 96 < counts[0] < 160, counts[0] def test_dropout_on_shard_axis(self): with Tensor.train(): @@ -825,7 +843,7 @@ def test_dropout_on_shard_axis(self): output = X.dropout(0.5).numpy() unique, counts = np.unique(output, return_counts=True) assert set(unique) == {0, 2}, unique - assert 200 < counts[0] < 312, counts[0] + assert 192 < counts[0] < 320, counts[0] @unittest.skip("TODO: this requires forced_realize to be deleted.") def test_shard_memory(self): @@ -836,7 +854,7 @@ def test_shard_memory(self): def test_clone(self): for axis in (None, 0): - t = Tensor.arange(16).reshape(4, 4).shard(devices_2, axis=axis).contiguous().realize() + t = Tensor.arange(16).reshape(4, 4).clone().shard(devices_2, axis=axis).contiguous().realize() t_clone = t.clone().realize() self.assertEqual(t_clone.device, t.device) self.assertEqual(t_clone.uop.axis, axis) @@ -847,7 +865,7 @@ def test_clone(self): @unittest.skip("RANGEIFY doesn't support multi const folding") def test_multi_const_folding(self): with Context(TRACK_MATCH_STATS=0): - a = Tensor.arange(3).realize() + a = Tensor.arange(3).clone().realize() zeros = Tensor.zeros(3).realize() b = a.to(devices_2)*zeros.to(devices_2) sched = b.schedule_linear().src @@ -886,7 +904,7 @@ def setUp(self): pass # shrink a multitensor on sharded axis def test_shrink_bad_args(self): - t = Tensor.arange(64).reshape(8, 8).contiguous().realize() + t = Tensor.arange(64).reshape(8, 8).clone().realize() t.shard_([f"{Device.DEFAULT}:{i}" for i in range(4)], axis=0) with self.assertRaises(AssertionError): @@ -908,8 +926,8 @@ def test_shrink_bad_args(self): @given(strat.sampled_from([dtypes.float, dtypes.int, dtypes.int64, dtypes.int16])) def test_ops(self, dtype): - if not is_dtype_supported(dtype): return - t = Tensor.arange(64).reshape(8, 8).contiguous().realize() + if dtype not in Device[Device.DEFAULT].renderer.supported_dtypes(): return + t = Tensor.arange(64).reshape(8, 8).clone().realize() t.shard_([f"{Device.DEFAULT}:{i}" for i in range(4)], axis=0) for i in range(4): print(f"{i=}") @@ -953,7 +971,7 @@ def test_ops(self, dtype): np.testing.assert_allclose(a.flip(-1).numpy(), b.flip(-1).numpy(), rtol=1e-7, atol=1e-3) def test_add_two_partitions(self): - t = Tensor.arange(64).reshape(8, 8).contiguous().realize() + t = Tensor.arange(64).reshape(8, 8).clone().realize() t.shard_([f"{Device.DEFAULT}:{i}" for i in range(4)], axis=0) a = t.shrink(((2, 4), None)) @@ -970,7 +988,7 @@ def test_add_two_partitions(self): def test_add_different_tensors(self): devices = [f"{Device.DEFAULT}:{i}" for i in range(4)] - x = Tensor.arange(64).reshape(8, 8).contiguous().realize().shard(devices, axis=0) + x = Tensor.arange(64).reshape(8, 8).clone().realize().shard(devices, axis=0) to_add = [] for i in range(len(devices)): @@ -1080,7 +1098,7 @@ def test_unsynced_backprop_sync_weights(self): @given(strat.sampled_from((False, True))) def test_batchnorm(self, is_training): devices = [f"{Device.DEFAULT}:{i}" for i in range(4)] - x = Tensor.arange(4096).reshape(8, 8, 8, 8).contiguous().realize().shard(devices, axis=0) + x = Tensor.arange(4096).reshape(8, 8, 8, 8).clone().realize().shard(devices, axis=0) with Tensor.train(is_training): bns = [] @@ -1088,8 +1106,6 @@ def test_batchnorm(self, is_training): bn = nn.BatchNorm2d(8) for p in get_parameters(bn): p.shard_(devices) - bn.weight.requires_grad = True - bn.bias.requires_grad = True bns.append(bn) bn_ts = [] @@ -1168,28 +1184,28 @@ def _check(self, a_ref:Tensor, a_multi:Tensor, view_fn): @unittest.skip("flaky on LLVM") def test_shrink_non_shard_axis(self): - ref = Tensor.arange(8*4*10).reshape(8, 4, 10).contiguous().realize() - a = Tensor.arange(8*4*10).reshape(8, 4, 10).contiguous().shard(devices_2, axis=1).realize() + ref = Tensor.arange(8*4*10).reshape(8, 4, 10).clone().realize() + a = Tensor.arange(8*4*10).reshape(8, 4, 10).clone().shard(devices_2, axis=1).realize() self._check(ref, a, lambda t: t[3]) def test_shrink_2d(self): - ref = Tensor.arange(6*4).reshape(6, 4).contiguous().realize() - a = Tensor.arange(6*4).reshape(6, 4).contiguous().shard(devices_2, axis=1).realize() + ref = Tensor.arange(6*4).reshape(6, 4).clone().realize() + a = Tensor.arange(6*4).reshape(6, 4).clone().shard(devices_2, axis=1).realize() self._check(ref, a, lambda t: t.shrink(((1, 4), None))) def test_reshape_then_shrink(self): - ref = Tensor.arange(8*6).reshape(8, 6).contiguous().realize() - a = Tensor.arange(8*6).reshape(8, 6).contiguous().shard(devices_2, axis=1).realize() + ref = Tensor.arange(8*6).reshape(8, 6).clone().realize() + a = Tensor.arange(8*6).reshape(8, 6).clone().shard(devices_2, axis=1).realize() self._check(ref, a, lambda t: t.reshape(4, 2, 6)[1]) def test_chained_shrink(self): - ref = Tensor.arange(10*8).reshape(10, 8).contiguous().realize() - a = Tensor.arange(10*8).reshape(10, 8).contiguous().shard(devices_2, axis=1).realize() + ref = Tensor.arange(10*8).reshape(10, 8).clone().realize() + a = Tensor.arange(10*8).reshape(10, 8).clone().shard(devices_2, axis=1).realize() self._check(ref, a, lambda t: t.shrink(((2, 8), None)).shrink(((1, 4), None))) def test_4_devices(self): - ref = Tensor.arange(8*12).reshape(8, 12).contiguous().realize() - a = Tensor.arange(8*12).reshape(8, 12).contiguous().shard(devices_4, axis=1).realize() + ref = Tensor.arange(8*12).reshape(8, 12).clone().realize() + a = Tensor.arange(8*12).reshape(8, 12).clone().shard(devices_4, axis=1).realize() out = a[5].contiguous() linear, var_vals = out.linear_with_vars() if all(hasattr(Device[d].allocator, "_offset") for d in out.device): @@ -1280,7 +1296,7 @@ class TestMultiSetitem(unittest.TestCase): @needs_second_gpu def setUp(self): pass - def _t(self, axis): return Tensor.arange(16).contiguous().realize().shard(self.device, axis=axis) + def _t(self, axis): return Tensor.arange(16).clone().realize().shard(self.device, axis=axis) def test_setitem_scalar_axis0(self): t = self._t(0) diff --git a/test/backend/test_nn.py b/test/backend/test_nn.py index 263ee5760e832..d6a84b06c4b59 100644 --- a/test/backend/test_nn.py +++ b/test/backend/test_nn.py @@ -149,8 +149,6 @@ def test_conv2d_winograd(self): # create in tinygrad layer = Conv2d(C1, C2, kernel_size=K, stride=S, padding=P) - layer.weight.requires_grad = True - layer.bias.requires_grad = True # create in torch torch_layer = torch.nn.Conv2d(C1, C2, kernel_size=K, stride=S, padding=P).eval() @@ -158,7 +156,7 @@ def test_conv2d_winograd(self): torch_layer.bias = torch.nn.Parameter(torch.tensor(layer.bias.numpy(), dtype=torch.float32)) # test - x = Tensor.uniform(BS, C1, H, W, requires_grad=True) + x = Tensor.uniform(BS, C1, H, W) with Context(WINO=1): z = layer(x) @@ -192,12 +190,12 @@ def test_groupnorm(self): # create in tinygrad layer = GroupNorm(G, C) - layer.weight = Tensor(torch_layer.weight.detach().numpy(), requires_grad=True) - layer.bias = Tensor(torch_layer.bias.detach().numpy(), requires_grad=True) + layer.weight = Tensor(torch_layer.weight.detach().numpy()) + layer.bias = Tensor(torch_layer.bias.detach().numpy()) for _ in range(10): # forward - x = Tensor.randn(BS, C, H, W, requires_grad=True) + x = Tensor.randn(BS, C, H, W) z = layer(x) z.sum().backward() @@ -218,10 +216,10 @@ def test_layernorm_forward(self): # create in tinygrad layer = LayerNorm([H, W]) - layer.weight = Tensor(torch_layer.weight.detach().numpy(), requires_grad=True) - layer.bias = Tensor(torch_layer.bias.detach().numpy(), requires_grad=True) + layer.weight = Tensor(torch_layer.weight.detach().numpy()) + layer.bias = Tensor(torch_layer.bias.detach().numpy()) - x = Tensor.empty(N, C, H, W, requires_grad=True) + x = Tensor.empty(N, C, H, W) z = layer(x) z.realize() @@ -240,12 +238,12 @@ def test_layernorm(self): # create in tinygrad layer = LayerNorm([H, W]) - layer.weight = Tensor(torch_layer.weight.detach().numpy(), requires_grad=True) - layer.bias = Tensor(torch_layer.bias.detach().numpy(), requires_grad=True) + layer.weight = Tensor(torch_layer.weight.detach().numpy()) + layer.bias = Tensor(torch_layer.bias.detach().numpy()) for _ in range(10): # forward - x = Tensor.randn(N, C, H, W, requires_grad=True) + x = Tensor.randn(N, C, H, W) z = layer(x) z.sum().backward() @@ -266,12 +264,12 @@ def test_layernorm_2d(self): # create in tinygrad layer = LayerNorm2d(C) - layer.weight = Tensor(torch_layer.weight.detach().numpy(), requires_grad=True) - layer.bias = Tensor(torch_layer.bias.detach().numpy(), requires_grad=True) + layer.weight = Tensor(torch_layer.weight.detach().numpy()) + layer.bias = Tensor(torch_layer.bias.detach().numpy()) for _ in range(10): # forward - x = Tensor.randn(N, C, H, W, requires_grad=True) + x = Tensor.randn(N, C, H, W) z = layer(x) z.sum().backward() @@ -292,12 +290,12 @@ def test_instancenorm_2d(self): # create in tinygrad layer = InstanceNorm(C) - layer.weight = Tensor(torch_layer.weight.detach().numpy(), requires_grad=True) - layer.bias = Tensor(torch_layer.bias.detach().numpy(), requires_grad=True) + layer.weight = Tensor(torch_layer.weight.detach().numpy()) + layer.bias = Tensor(torch_layer.bias.detach().numpy()) for _ in range(10): # forward - x = Tensor.randn(N, C, H, W, requires_grad=True) + x = Tensor.randn(N, C, H, W) z = layer(x) z.sum().backward() @@ -318,12 +316,12 @@ def test_instancenorm_3d(self): # create in tinygrad layer = InstanceNorm(C) - layer.weight = Tensor(torch_layer.weight.detach().numpy(), requires_grad=True) - layer.bias = Tensor(torch_layer.bias.detach().numpy(), requires_grad=True) + layer.weight = Tensor(torch_layer.weight.detach().numpy()) + layer.bias = Tensor(torch_layer.bias.detach().numpy()) for _ in range(10): # forward - x = Tensor.randn(N, C, D, H, W, requires_grad=True) + x = Tensor.randn(N, C, D, H, W) z = layer(x) z.sum().backward() @@ -356,11 +354,10 @@ def forward(self, x): B, T, embed_size = 4, 10, 20 torch_layer = TorchRMSNorm(embed_size) layer = RMSNorm(embed_size) - layer.weight.requires_grad = True for _ in range(10): # forward - x = Tensor.randn(B, T, embed_size, requires_grad=True) + x = Tensor.randn(B, T, embed_size) z = layer(x) z.sum().backward() @@ -377,7 +374,7 @@ def forward(self, x): for _ in range(10): # forward - x = Tensor.randn(B, T, embed_size, requires_grad=True) + x = Tensor.randn(B, T, embed_size) z = layer(x) z.sum().backward() diff --git a/test/backend/test_ops.py b/test/backend/test_ops.py index 9409697807b36..7c04506b7c68e 100644 --- a/test/backend/test_ops.py +++ b/test/backend/test_ops.py @@ -1,11 +1,10 @@ -import time, math, unittest, functools, platform, warnings +import time, math, unittest, functools, platform, warnings, sys import numpy as np from typing import List, Callable import torch -from tinygrad.helpers import getenv, CI, DEBUG, DEV, IMAGE, Context +from tinygrad.helpers import getenv, DEBUG, DEV, IMAGE, Context from tinygrad import Tensor, Device, dtypes from tinygrad.tensor import _to_np_dtype -from tinygrad.device import is_dtype_supported from tinygrad.renderer.cstyle import QCOMCLRenderer from tinygrad.renderer.nir import NIRRenderer @@ -74,7 +73,7 @@ def compare(s, tinygrad_output, torch_output, atol, rtol): for i, (t, torch_grad) in enumerate(zip(tiny_grads, torch_grads)): compare(f"backward pass tensor {i}", t.numpy(), torch_grad.detach().cpu().numpy(), atol=grad_atol, rtol=grad_rtol) - if not CI: + if sys.stdout.isatty(): print("\ntesting %40r torch/tinygrad fp: %.2f / %.2f ms bp: %.2f / %.2f ms " % \ (shps, torch_fp*1000, tinygrad_fp*1000, torch_fbp*1000, tinygrad_fbp*1000), end="") @@ -88,7 +87,7 @@ def prepare_test_op(low, high, shps, vals, forward_only=False): for i in range(len(ts)): # NOTE: torch default int64 for python ints input if ts[i].dtype == torch.int64: ts[i] = ts[i].type(torch.int32) - tst = [Tensor(x.detach().cpu().numpy(), requires_grad=(not forward_only and not FORWARD_ONLY)) for x in ts] + tst = [Tensor(x.detach().cpu().numpy()) for x in ts] return ts, tst class TestOps(unittest.TestCase): @@ -103,7 +102,7 @@ def helper_test_exception(self, shps, torch_fxn, tinygrad_fxn=None, expected=Non with self.assertRaises(expected) as tinygrad_cm: tinygrad_fxn(*tst) if exact: self.assertEqual(str(torch_cm.exception), str(tinygrad_cm.exception)) - if not CI: print("\ntesting %40r torch/tinygrad exception: %s / %s" % (shps, torch_cm.exception, tinygrad_cm.exception), end="") + if sys.stdout.isatty(): print("\ntesting %40r torch/tinygrad exception: %s / %s" % (shps, torch_cm.exception, tinygrad_cm.exception), end="") def test_full_like(self): a = Tensor([[1,2,3],[4,5,6]], dtype=dtypes.float32) @@ -249,9 +248,9 @@ def test_unfold(self): self.helper_test_exception([(8,)], lambda x: x.unfold(0, 1, -1), expected=RuntimeError) def test_meshgrid(self): - x, xt = torch.tensor([0.,1.,2.], requires_grad=True), Tensor([0.,1.,2.], requires_grad=True) - y, yt = torch.tensor([3.,4.,5.,6.], requires_grad=True), Tensor([3.,4.,5.,6.], requires_grad=True) - z, zt = torch.tensor([7.,8.,9.], requires_grad=True), Tensor([7.,8.,9.], requires_grad=True) + x, xt = torch.tensor([0.,1.,2.], requires_grad=True), Tensor([0.,1.,2.]) + y, yt = torch.tensor([3.,4.,5.,6.], requires_grad=True), Tensor([3.,4.,5.,6.]) + z, zt = torch.tensor([7.,8.,9.], requires_grad=True), Tensor([7.,8.,9.]) for indexing in ("ij", "xy"): tor = torch.meshgrid(x, indexing=indexing) ten = xt.meshgrid(indexing=indexing) @@ -264,7 +263,7 @@ def test_meshgrid(self): for tor_i, ten_i in zip(tor, ten): helper_test_op([], lambda: tor_i, lambda: ten_i) tor = torch.meshgrid(x, torch.tensor(10., requires_grad=True), y, z, indexing=indexing) - ten = xt.meshgrid(Tensor(10., requires_grad=True), yt, zt, indexing=indexing) + ten = xt.meshgrid(Tensor(10.), yt, zt, indexing=indexing) self.assertEqual(len(tor), len(ten)) for tor_i, ten_i in zip(tor, ten): helper_test_op([], lambda: tor_i, lambda: ten_i) @@ -386,11 +385,11 @@ def test_cmp_ne_backwards(self): t1 = torch.ones(4, requires_grad=True) t2 = torch.ones(4, requires_grad=True) self.assertRaises(RuntimeError, (t1 != t2).sum().backward) - tt1 = Tensor.ones(4, requires_grad=True) - tt2 = Tensor.ones(4, requires_grad=True) + tt1 = Tensor.ones(4) + tt2 = Tensor.ones(4) self.assertRaises(RuntimeError, (tt1 != tt2).sum().backward) """ - tt = Tensor.randn(4, requires_grad=True) + tt = Tensor.randn(4) (tt*(tt != 0)).sum().backward() t = torch.tensor(tt.numpy(), requires_grad=True) (t*(t != 0)).sum().backward() @@ -402,11 +401,11 @@ def test_cmp_lt_backwards(self): t1 = torch.ones(4, requires_grad=True) t2 = torch.ones(4, requires_grad=True) self.assertRaises(RuntimeError, (t1 < t2).sum().backward) - tt1 = Tensor.ones(4, requires_grad=True) - tt2 = Tensor.ones(4, requires_grad=True) + tt1 = Tensor.ones(4) + tt2 = Tensor.ones(4) self.assertRaises(RuntimeError, (tt1 < tt2).sum().backward) """ - tt = Tensor.randn(4, requires_grad=True) + tt = Tensor.randn(4) (tt*(tt < 0)).sum().backward() t = torch.tensor(tt.numpy(), requires_grad=True) (t*(t < 0)).sum().backward() @@ -683,7 +682,7 @@ def test_pow(self): helper_test_op([(45,65)], lambda x: x**1.2, low=-30, high=-27) helper_test_op([()], lambda x: x**0.2, low=-30, high=-27) helper_test_op([()], lambda x: x**1.2, low=-30, high=-27) - a, b = Tensor([0.0], requires_grad=True), torch.tensor([0.0], requires_grad=True) + a, b = Tensor([0.0]), torch.tensor([0.0], requires_grad=True) helper_test_op([], lambda: b**1.1, lambda: a**1.1) def test_pow_const(self): @@ -1449,9 +1448,7 @@ def test_small_gemm_range(self): np.arange(64,128,dtype=np.float32).reshape(8,8)]) def test_small_gemm_eye(self): helper_test_op(None, lambda x,y: x.matmul(y), lambda x,y: x@y, vals=[np.eye(8).astype(np.float32), np.eye(8).astype(np.float32)]) - @unittest.skipIf(CI and Device.DEFAULT in ["NV", "CL", "CUDA"] or (Device.DEFAULT == "CPU" and DEV.renderer == "LLVM") or IMAGE - or (Device.DEFAULT == "WEBGPU" and platform.system() == "Windows"), "not supported on these in CI/IMAGE") - @unittest.skipIf(Device.DEFAULT == "QCOM", "not precise enough") + @unittest.skipUnless(dtypes.half in Device[Device.DEFAULT].renderer.supported_dtypes(), "not precise enough when emulating") def test_gemm_fp16(self): helper_test_op([(64,64), (64,64)], lambda x,y: x.half().matmul(y.half()), atol=5e-3, rtol=5e-3, grad_atol=5e-3, grad_rtol=5e-3) def test_gemm(self): @@ -1508,7 +1505,8 @@ def test_sum(self): def test_sum_dtype_arg(self): helper_test_op([(45,3)], lambda x: x.sum(), lambda x: x.sum(dtype=dtypes.float32)) - if is_dtype_supported(dtypes.float64): helper_test_op([(45,3)], lambda x: x.sum(dtype=torch.float64), lambda x: x.sum(dtype=dtypes.float64)) + if dtypes.float64 in Device[Device.DEFAULT].renderer.supported_dtypes(): + helper_test_op([(45,3)], lambda x: x.sum(dtype=torch.float64), lambda x: x.sum(dtype=dtypes.float64)) with self.assertRaises(AttributeError): Tensor([1.0, 2.0]).sum(dtype="") @@ -1991,9 +1989,7 @@ def test_pad_circular_mode(self): self.helper_test_exception([(1,1,5,5)], lambda x: torch.nn.functional.pad(x, (3,6,0,0), mode="circular"), lambda x: x.pad((3,6,0,0), mode="circular"), expected=(RuntimeError, ValueError)) - with self.assertRaises(NotImplementedError): - # negative pads with circular pads is not supported - Tensor.randn(1,1,5,5).pad((3,-5,1,-5), mode="circular") + helper_test_op([(1,1,5,5)], lambda x: torch.nn.functional.pad(x, (1,-2,2,-1), mode="circular"), lambda x: x.pad((1,-2,2,-1), mode="circular")) def test_pad_reshape(self): helper_test_op([(1, 2)], @@ -2449,7 +2445,7 @@ def test_strided_conv2d_simple(self): @unittest.skipUnless(Device.DEFAULT == "CPU" and DEV.renderer == "LLVM", "DEVECTORIZE=0 only for LLVM") def test_strided_conv2d_simple_vec(self): - with Context(DEVECTORIZE=0): self.test_strided_conv2d_simple() + self.test_strided_conv2d_simple() @slow_test def test_strided_conv2d(self): @@ -2847,6 +2843,7 @@ def test_multicat(self): for dim in range(-1, 2): helper_test_op([(45,65), (45,65), (45,65)], lambda x,y,z: torch.cat((x,y,z), dim), lambda x,y,z: x.cat(y, z, dim=dim)) + @unittest.skipIf(COMPILE_ONLY, "test requires runtime") def test_stack(self): for dim in range(-1, 3): helper_test_op([(5,6,3), (5,6,3), (5,6,3)], lambda x, y, z: torch.stack((x, y, z), dim), lambda x, y, z: Tensor.stack(x, y, z, dim=dim)) @@ -2857,8 +2854,7 @@ def test_stack(self): with self.assertRaises(ValueError): Tensor.stack((Tensor([1, 2]), Tensor([3, 4])), Tensor([5, 6])) - a = Tensor(3.14) - np.testing.assert_allclose(Tensor.stack(a, a).numpy(), Tensor([3.14, 3.14]).numpy()) + np.testing.assert_allclose(Tensor.stack(Tensor(3.14), Tensor(3.14)).numpy(), np.array([3.14, 3.14])) def test_stack_max(self): helper_test_op(None, lambda x, y: torch.stack((x, y)).max(axis=0)[0], lambda x, y: Tensor.stack(x, y).max(axis=0), vals=[[1.], [2.]]) @@ -2926,7 +2922,7 @@ def _get_index_randoms(self): c = torch.randint(low=-5, high=5, size=(1,1,4,1,1,1), dtype=torch.int64, requires_grad=False) d = torch.randint(high=4, size=(2,1,1,5,1,1), dtype=torch.int64, requires_grad=False) e = torch.randint(high=1, size=(1,1,1,1,6,1), dtype=torch.int64, requires_grad=False) - i, j, k, o, p = [Tensor(tor.detach().cpu().numpy().astype(np.int32), requires_grad=False) for tor in [a,b,c,d,e]] + i, j, k, o, p = [Tensor(tor.detach().cpu().numpy().astype(np.int32)) for tor in [a,b,c,d,e]] return a,b,c,d,e,i,j,k,o,p def test_fancy_indexing_inf(self): @@ -3037,7 +3033,7 @@ def test_gather(self): # indices cannot have gradient # indices cannot be negative (torch gather) b = torch.randint(3, size=[3,4,5], dtype=torch.int64, requires_grad=False) - a = Tensor(b.detach().cpu().numpy().astype(np.int32), dtype=dtypes.int32, requires_grad=False) + a = Tensor(b.detach().cpu().numpy().astype(np.int32), dtype=dtypes.int32) helper_test_op([(4,5,6)], lambda x: x.gather(dim=0, index=b), lambda x: x.gather(dim=0, index=a)) helper_test_op([(4,5,6)], lambda x: x.gather(dim=1, index=b), lambda x: x.gather(dim=1, index=a)) helper_test_op([(4,5,6)], lambda x: x.gather(dim=2, index=b), lambda x: x.gather(dim=2, index=a)) @@ -3059,7 +3055,7 @@ def test_gather(self): def test_scatter(self): b = torch.randint(3, size=[3,4,5], dtype=torch.int64, requires_grad=False) - a = Tensor(b.detach().cpu().numpy().astype(np.int32), dtype=dtypes.int32, requires_grad=False) + a = Tensor(b.detach().cpu().numpy().astype(np.int32), dtype=dtypes.int32) for dim in (0,1,2,-1,-2,-3): helper_test_op([(4,5,6), (4,5,6)], lambda x,src: x.scatter(dim=dim, index=b, src=src), lambda x,src: x.scatter(dim=dim, index=a, src=src), forward_only=True) @@ -3084,7 +3080,7 @@ def test_scatter(self): # overlapping indices with 0s b = torch.tensor([0,0], requires_grad=False) - a = Tensor(b.detach().cpu().numpy().astype(np.int32), dtype=dtypes.int32, requires_grad=False) + a = Tensor(b.detach().cpu().numpy().astype(np.int32), dtype=dtypes.int32) helper_test_op(None, lambda x,src: x.scatter(0, b, src), lambda x,src: x.scatter(0, a, src), forward_only=True, @@ -3092,7 +3088,7 @@ def test_scatter(self): def test_scatter_add(self): b = torch.randint(3, size=[3,4,5], dtype=torch.int64, requires_grad=False) - a = Tensor(b.detach().cpu().numpy().astype(np.int32), dtype=dtypes.int32, requires_grad=False) + a = Tensor(b.detach().cpu().numpy().astype(np.int32), dtype=dtypes.int32) helper_test_op([(4,5,6)], lambda x: x.scatter(dim=1, index=b, value=float("inf"), reduce="add"), lambda x: x.scatter(dim=1, index=a, src=float("inf"), reduce="add"), forward_only=True) @@ -3104,7 +3100,7 @@ def test_scatter_add(self): def test_scatter_mul(self): b = torch.randint(3, size=[3,4,5], dtype=torch.int64, requires_grad=False) - a = Tensor(b.detach().cpu().numpy().astype(np.int32), dtype=dtypes.int32, requires_grad=False) + a = Tensor(b.detach().cpu().numpy().astype(np.int32), dtype=dtypes.int32) helper_test_op([(4,5,6)], lambda x: x.scatter(dim=1, index=b, value=float("inf"), reduce="multiply"), lambda x: x.scatter(dim=1, index=a, src=float("inf"), reduce="multiply"), forward_only=True) @@ -3121,7 +3117,7 @@ def test_scatter_no_reduce_tensor_src(self): @slow_test def test_scatter_reduce(self): b = torch.randint(3, size=[3,4,5], dtype=torch.int64, requires_grad=False) - a = Tensor(b.detach().cpu().numpy().astype(np.int32), requires_grad=False) + a = Tensor(b.detach().cpu().numpy().astype(np.int32)) for reduce in ("sum", "prod", "mean", "amin", "amax"): for dim in (-1,1,-3): helper_test_op([(3,4,5), (3,4,5)], @@ -3133,7 +3129,7 @@ def test_scatter_reduce(self): def test_scatter_reduce_prod_zeros(self): b = torch.randint(3, size=[3,4,5], dtype=torch.int64, requires_grad=False) - a = Tensor(b.detach().cpu().numpy().astype(np.int32), dtype=dtypes.int32, requires_grad=False) + a = Tensor(b.detach().cpu().numpy().astype(np.int32), dtype=dtypes.int32) x = Tensor.zeros([4,5,6]).float() y = torch.zeros([4,5,6]).float() helper_test_op([(4,5,6)], @@ -3142,7 +3138,7 @@ def test_scatter_reduce_prod_zeros(self): def test_scatter_reduce_errors(self): b = torch.randint(3, size=[3,4,5], dtype=torch.int64, requires_grad=False) - a = Tensor(b.detach().cpu().numpy().astype(np.int32), dtype=dtypes.int32, requires_grad=False) + a = Tensor(b.detach().cpu().numpy().astype(np.int32), dtype=dtypes.int32) # invalid reduce arg self.helper_test_exception([(4,5,6), (4,5,6)], lambda x,src: x.scatter_reduce(dim=0, index=b, src=src, reduce="INVALID"), @@ -3186,7 +3182,7 @@ def test_scaled_dot_product_attention_gqa_errors(self): self.helper_test_exception([(32,31,16,64), (32,8,16,64), (32,8,16,64)], lambda x,y,z: torch.nn.functional.scaled_dot_product_attention(x,y,z), lambda x,y,z: Tensor.scaled_dot_product_attention(x,y,z,enable_gqa=True), - expected=(AssertionError, RuntimeError, ValueError)) + expected=(AssertionError, RuntimeError, ValueError, IndexError)) def test_binary_crossentropy(self): helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.binary_cross_entropy(x.sigmoid(),y.clip(0,1)), @@ -3348,6 +3344,7 @@ def test_masked_select_size(self): # fill_value must not alter output dtype self.assertEqual(Tensor([1.0, 2.0]).masked_select(Tensor([True, False]), size=3, fill_value=-1).dtype, dtypes.default_float) + @slow_test def test_nonzero(self): helper_test_op([(32, 10)], lambda x: (x>0.5).nonzero().int(), lambda x: (x>0.5).nonzero(), forward_only=True) helper_test_op([(20,)], lambda x: (x>0.5).nonzero().int(), lambda x: (x>0.5).nonzero(), forward_only=True) @@ -3379,7 +3376,6 @@ def test_int_or(self): t = (Tensor([0], dtype='int') | 0xFFFFFFFF).item() if not COMPILE_ONLY: assert t == -1 -@unittest.skipUnless(is_dtype_supported(dtypes.uchar), f"no uint8 on {Device.DEFAULT}") class TestOpsUint8(unittest.TestCase): def test_cast(self): helper_test_op([(2,3,64,64)], lambda x: x.type(torch.uint8), lambda x: x.cast('uint8'), forward_only=True, low=0, high=255) diff --git a/test/backend/test_optim.py b/test/backend/test_optim.py index 4d27b60b6e6a6..06514b2a7c494 100644 --- a/test/backend/test_optim.py +++ b/test/backend/test_optim.py @@ -3,7 +3,6 @@ import unittest from tinygrad import Tensor, Device, dtypes from tinygrad.nn.optim import Adam, SGD, AdamW, Muon, LAMB -from tinygrad.device import is_dtype_supported from test.helpers import needs_second_gpu, slow np.random.seed(1337) @@ -11,17 +10,20 @@ W_init = np.random.randn(4,4).astype(np.float32) m_init = np.random.randn(1,4).astype(np.float32) +def _param(tensor, val): + return tensor(val, requires_grad=True) if tensor is torch.tensor else tensor(val) + class TeenyNet: def __init__(self, tensor): - self.x = tensor(x_init.copy(), requires_grad=True) - self.W = tensor(W_init.copy(), requires_grad=True) + self.x = _param(tensor, x_init.copy()) + self.W = _param(tensor, W_init.copy()) def forward(self): return (self.x * self.W).sum() class TinyNet: def __init__(self, tensor): - self.x = tensor(x_init.copy(), requires_grad=True) - self.W = tensor(W_init.copy(), requires_grad=True) + self.x = _param(tensor, x_init.copy()) + self.W = _param(tensor, W_init.copy()) self.m = tensor(m_init.copy()) def forward(self): @@ -142,7 +144,7 @@ def test_duped_weights(self): np.testing.assert_allclose(losses[0], losses[1], atol=1e-4, rtol=0) - @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half") + @unittest.skipUnless(dtypes.half in Device[Device.DEFAULT].renderer.supported_dtypes(), "need half") def test_mixed_precision(self): old_default_float, dtypes.default_float = dtypes.default_float, dtypes.half # weight update would overflow without upcasting @@ -152,7 +154,7 @@ def test_mixed_precision(self): dtypes.default_float = old_default_float def test_assert_tensor_train(self): - t = Tensor.ones((1,1), requires_grad=True) + t = Tensor.ones((1,1)) optimizer = Adam([t]) optimizer.zero_grad() old_state = Tensor.training @@ -165,7 +167,7 @@ def test_assert_tensor_train(self): def test_lamb_cpu_offload(self): # test that LAMB works when optimizer params (m, v, b1_t, b2_t) are moved to CPU - t = Tensor(x_init.copy(), requires_grad=True) + t = Tensor(x_init.copy()) opt = LAMB([t]) # move optimizer state to CPU for p in opt.m + opt.v + [opt.b1_t, opt.b2_t]: p.to_("CPU") @@ -178,7 +180,7 @@ def test_lamb_cpu_offload(self): @needs_second_gpu def test_lamb_cpu_offload_multi(self): ds = tuple(f"{Device.DEFAULT}:{i}" for i in range(2)) - t = Tensor(x_init.copy(), requires_grad=True).shard(ds, axis=1) + t = Tensor(x_init.copy()).shard(ds, axis=1) ds = t.device opt = LAMB([t]) # move optimizer state to CPU diff --git a/test/backend/test_pickle.py b/test/backend/test_pickle.py index d2db7ec37d8db..25f29b6639deb 100644 --- a/test/backend/test_pickle.py +++ b/test/backend/test_pickle.py @@ -67,7 +67,7 @@ def test_pickle_realized_tensor_alt2(self): # NOTE: currently Buffer exists on the uop, not tensor def test_pickle_buffer_uop(self): - t = Tensor.arange(4).realize() + t = Tensor.arange(4).clone().realize() a = t.uop assert a.is_realized self.assertIsNotNone(buffer:=a.base.realized) @@ -95,7 +95,7 @@ def test_pickle_variable(self): np.testing.assert_equal(vt2.numpy(), 20) def test_pickle_buffer_view(self): - t = Tensor.arange(10, device="CPU").contiguous().realize() + t = Tensor.arange(10).clone(device="CPU").realize() vt = t[3:5].contiguous().realize() assert hasattr(vt.uop.buffer, 'base') ref_value = vt.tolist() diff --git a/test/backend/test_profiler.py b/test/backend/test_profiler.py index e14536a897473..2091f85ac4785 100644 --- a/test/backend/test_profiler.py +++ b/test/backend/test_profiler.py @@ -1,6 +1,6 @@ import unittest, struct, contextlib, statistics, gc from tinygrad import Device, Tensor, dtypes, TinyJit -from tinygrad.helpers import CI, DEV, Context, ProfileRangeEvent, cpu_profile, cpu_events, ProfilePointEvent, dedup +from tinygrad.helpers import DEV, Context, ProfileRangeEvent, cpu_profile, cpu_events, ProfilePointEvent, dedup from tinygrad.device import Buffer, BufferSpec, Compiled, ProfileDeviceEvent, ProfileGraphEvent from tinygrad.runtime.support.hcq import HCQCompiled from tinygrad.engine.realize import get_runtime @@ -144,7 +144,8 @@ def f(a): assert len(graph_evs) == 2, "2 graph events are expected" assert len(graph_evs[0].ents) == 2, "two entities are expected" - @unittest.skipIf(CI or not issubclass(type(Device[Device.DEFAULT]), HCQCompiled), "skip CI") + @unittest.skipIf(MOCKGPU, "skip MOCKGPU") + @unittest.skipUnless(issubclass(type(Device[Device.DEFAULT]), HCQCompiled), "must be HCQ") def test_dev_jitter_matrix(self): dev_cnt = 6 try: devs = [Device[f"{Device.DEFAULT}:{i}"] for i in range(dev_cnt)] diff --git a/test/backend/test_randomness.py b/test/backend/test_randomness.py index 5ac503b8123fc..834eb9ddca6cc 100644 --- a/test/backend/test_randomness.py +++ b/test/backend/test_randomness.py @@ -1,15 +1,15 @@ import unittest, math -from functools import partial -from tinygrad import nn, dtypes, Tensor, Device, TinyJit, Variable -from tinygrad.helpers import getenv, CI, OSX -from tinygrad.device import is_dtype_supported +from tinygrad import dtypes, Tensor, Device +from tinygrad.helpers import getenv, DEV from tinygrad.codegen import to_program from tinygrad.uop.ops import Ops from tinygrad.renderer.ptx import PTXRenderer from tinygrad.renderer.nir import NIRRenderer +from tinygrad.renderer.isa.x86 import X86Renderer from test.helpers import not_support_multi_device, needs_second_gpu +from test.unit.test_randomness import equal_distribution, normal_test import numpy as np import torch @@ -18,54 +18,6 @@ settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False)) settings.load_profile("my_profile") -# https://gist.github.com/devries/11405101 -def ksprob(a): - fac, total, termbf = 2.0, 0.0, 0.0 - a2 = -2.0 * a * a - for j in range(1, 101): - term = fac * math.exp(a2 * j * j) - total += term - if math.fabs(term) <= 0.001 * termbf or math.fabs(term) <= 1e-8 * total: - return total - fac = -fac - termbf = math.fabs(term) - return 1.0 - -def kstest(l1, l2): - n1, n2 = len(l1), len(l2) - l1.sort() - l2.sort() - j1, j2, d, fn1, fn2 = 0, 0, 0.0, 0.0, 0.0 - while j1 < n1 and j2 < n2: - d1, d2 = l1[j1], l2[j2] - if d1 <= d2: - fn1 = (float(j1) + 1.0) / float(n1) - j1 += 1 - if d2 <= d1: - fn2 = (float(j2) + 1.0) / float(n2) - j2 += 1 - dtemp = math.fabs(fn2 - fn1) - if dtemp > d: - d = dtemp - ne = float(n1 * n2) / float(n1 + n2) - nesq = math.sqrt(ne) - prob = ksprob((nesq + 0.12 + 0.11 / nesq) * d) - return prob - -def equal_distribution(tiny_func, torch_func=None, numpy_func=None, shape=(40, 43), alpha=0.04): - Tensor.manual_seed(1337) - torch.manual_seed(1337) - np.random.seed(1337) - assert not (torch_func is None and numpy_func is None), "no function to compare with" - x1 = tiny_func(*shape).numpy().flatten() - x2 = tiny_func(shape).numpy().flatten() - if numpy_func is not None: y = numpy_func(shape).flatten() - if torch_func is not None: z = torch_func(shape).numpy().flatten() - return (numpy_func is None or (kstest(x1, y) >= alpha and kstest(x2, y) >= alpha)) and \ - (torch_func is None or (kstest(x1, z) >= alpha and kstest(x2, z) >= alpha)) - -def normal_test(func, shape=(20, 45), alpha=0.05): return equal_distribution(func, numpy_func=lambda x: np.random.randn(*x), shape=shape, alpha=alpha) - class TestRandomness(unittest.TestCase): def test_rand(self): self.assertFalse(normal_test(Tensor.rand)) @@ -85,7 +37,7 @@ def test_rand_is_lazy(self): self.assertTrue(r1.uop.is_realized, "tensor should be realized after .realize()") self.assertTrue(r2.uop.is_realized, "tensor should be realized after .realize()") - @unittest.skipUnless(is_dtype_supported(dtypes.float16), "need float16 support") + @unittest.skipUnless(dtypes.float16 in Device[Device.DEFAULT].renderer.supported_dtypes(), "need float16 support") def test_rand_float16(self): N = 128 x = Tensor.rand((2, N, N), dtype=dtypes.float16) @@ -96,7 +48,7 @@ def test_rand_float16(self): assert nx[nx == 0].size > 0 equal_distribution(lambda *x: Tensor.rand(*x, dtype=dtypes.float16), torch.rand, lambda x: np.random.rand(*x), shape=(2, N, N)) - @unittest.skipIf(CI and Device.DEFAULT in {"NV", "CUDA"}, "gpuocelot doesn't support certain ops needed for threefry") + @unittest.skipIf(DEV.interface.startswith("MOCK") and Device.DEFAULT in {"NV", "CUDA"}, "gpuocelot doesn't support certain ops needed for threefry") def test_threefry_against_reference(self): Tensor.manual_seed(1337) @@ -118,6 +70,7 @@ def test_threefry_against_reference(self): np.testing.assert_allclose(jr, r) @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (NIRRenderer, PTXRenderer)), "PTX and NIR use pointer arithmetic") + @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "X86 callee saved registers have ulong dtype") def test_threefry_doesnt_use_long(self): linear = Tensor.rand(20).schedule_linear() for call in linear.src: @@ -209,7 +162,7 @@ def test_threefry_same_kernels(self): if not (x.src[0] == y.src[0]): print(f"{x.src[0]} != {y.src[0]}") - @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), "need bfloat16 support") + @unittest.skipUnless(dtypes.bfloat16 in Device[Device.DEFAULT].renderer.supported_dtypes(), "need bfloat16 support") def test_rand_bfloat16(self): N = 128 x = Tensor.rand((2, N, N), dtype=dtypes.bfloat16) @@ -273,17 +226,12 @@ def test_randn_like_dtype(self): assert rand.dtype == dtypes.float16 assert rand.device == empty.device - def test_randn(self): - self.assertEqual(Tensor.randn(3,3,dtype=dtypes.half).dtype, dtypes.half) - self.assertTrue(normal_test(Tensor.randn)) - self.assertTrue(equal_distribution(Tensor.randn, torch.randn, lambda x: np.random.randn(*x))) - def test_randn_device(self): self.assertEqual(Tensor.randn(3,3,device="CPU").device, "CPU") @given(strat.sampled_from([dtypes.float, dtypes.float16, dtypes.bfloat16])) def test_randn_finite(self, default_float): - if not is_dtype_supported(default_float): return + if default_float not in Device[Device.DEFAULT].renderer.supported_dtypes(): return old_default_float = dtypes.default_float # low precision can result in inf from randn dtypes.default_float = default_float @@ -295,128 +243,6 @@ def test_randn_finite(self, default_float): assert math.isfinite(mn), mn dtypes.default_float = old_default_float - def test_randint(self): - self.assertFalse(normal_test(Tensor.randint)) - self.assertTrue(equal_distribution(partial(Tensor.randint, low=-2, high=5), - numpy_func=lambda x: np.random.randint(low=-2, high=5, size=x))) - self.assertTrue(equal_distribution(partial(Tensor.randint, low=-2, high=5, dtype="int32"), - numpy_func=lambda x: np.random.randint(low=-2, high=5, size=x))) - self.assertTrue(Tensor.randint(1, device="CPU").device=="CPU") - # check types of args - with self.assertRaises(TypeError): Tensor.randint((3, 4), low=0.1, high=3) - with self.assertRaises(TypeError): Tensor.randint((3, 4), low=0, high=3.5) - with self.assertRaises(TypeError): Tensor.randint((3, 4), low=1, high=3, dtype="float") - with self.assertRaises(TypeError): Tensor.randint((3, 4), low=0, high=3, dtype=dtypes.float32) - # check low < high - with self.assertRaises(ValueError): Tensor.randint((3, 4), low=10, high=5) - with self.assertRaises(ValueError): Tensor.randint((3, 4), low=10, high=10) - np.testing.assert_array_equal(Tensor.randint(16, low=5, high=6).numpy(), 5) - - def test_normal(self): - self.assertTrue(normal_test(Tensor.normal)) - self.assertTrue(equal_distribution(Tensor.normal, lambda x: torch.nn.init.normal_(torch.empty(x), mean=0, std=1), - lambda x: np.random.normal(loc=0, scale=1, size=x))) - # check std >= 0 - with self.assertRaises(ValueError): Tensor.normal((3, 4), mean=0, std=-1) - - def test_uniform(self): - self.assertFalse(normal_test(Tensor.uniform)) - self.assertTrue(equal_distribution(Tensor.uniform, lambda x: torch.nn.init.uniform_(torch.empty(x)), lambda x: np.random.uniform(size=x))) - self.assertTrue(equal_distribution(partial(Tensor.uniform, low=-100, high=100, dtype=dtypes.int32), - numpy_func=lambda x: np.random.randint(low=-100, high=100, size=x))) - # check low < high - with self.assertRaises(ValueError): Tensor.uniform((3, 4), low=5.0, high=3.0) - with self.assertRaises(ValueError): Tensor.uniform((3, 4), low=1.0, high=1.0) - - def test_scaled_uniform(self): - self.assertFalse(normal_test(Tensor.scaled_uniform)) - self.assertTrue(equal_distribution(Tensor.scaled_uniform, lambda x: torch.nn.init.uniform_(torch.empty(x), a=-1, b=1) / math.sqrt(math.prod(x)), - lambda x: np.random.uniform(-1, 1, size=x) / math.sqrt(math.prod(x)))) - - def test_glorot_uniform(self): - self.assertFalse(normal_test(Tensor.glorot_uniform)) - self.assertTrue(equal_distribution(Tensor.glorot_uniform, lambda x: torch.nn.init.xavier_uniform_(torch.empty(x)), - lambda x: np.random.uniform(-1, 1, size=x) * math.sqrt(6 / (x[0] + math.prod(x[1:]))))) - - def test_kaiming_uniform(self): - for shape in [(32, 16, 3, 3), (20, 44), (5, 15, 35)]: - self.assertTrue(equal_distribution(Tensor.kaiming_uniform, lambda x: torch.nn.init.kaiming_uniform_(torch.empty(x)), shape=shape)) - - def test_kaiming_normal(self): - for shape in [(32, 16, 3, 3), (20, 44), (3, 15, 35)]: - self.assertTrue(equal_distribution(Tensor.kaiming_normal, lambda x: torch.nn.init.kaiming_normal_(torch.empty(x)), shape=shape)) - - def test_multinomial(self): - self.assertRaises(AssertionError, lambda: Tensor(2).multinomial(1, replacement=False)) - self.assertRaises(AssertionError, lambda: Tensor([1, 9]).multinomial(0, replacement=False)) - def _check_with_torch(w, num_samples, replacement): - tiny_res = Tensor(w).multinomial(num_samples, replacement=replacement) - torch_res = torch.tensor(w).multinomial(num_samples, replacement=replacement) - self.assertEqual(tiny_res.shape, torch_res.shape) - if torch_res.ndim == 1: - tiny_res = tiny_res.unsqueeze(0) - torch_res = torch_res.unsqueeze(0) - for i in range(torch_res.shape[0]): - self.assertTrue(equal_distribution(lambda *_: tiny_res[i], lambda _: torch_res[i])) - _check_with_torch(w=[0.231, 0., 1., 0.5], num_samples=300, replacement=True) - _check_with_torch(w=[[0.2, 0.8]], num_samples=300, replacement=True) # 2D but only 1 row - _check_with_torch(w=[[0.453, 0., 1., 0.81], [0.1, 0.8, 0., 0.1]], num_samples=300, replacement=True) - # no-replacement - w = [0.1, 0.9] - self.assertRaises(AssertionError, lambda: Tensor(w).multinomial(100, replacement=False)) - - @TinyJit - def sample_one(): return Tensor(w).multinomial(1, replacement=False).realize() - - tiny_samples = [sample_one().item() for _ in range(1000)] - torch_samples = [torch.tensor(w).multinomial(1, replacement=False).item() for _ in range(1000)] - self.assertTrue(equal_distribution(lambda *_: Tensor(tiny_samples), lambda _: torch.tensor(torch_samples))) - - w = list(range(32)) - s1 = Tensor(w).multinomial(5, replacement=False).numpy() - self.assertEqual(len(set(s1.tolist())), 5) - s2 = Tensor(w).multinomial(5, replacement=False).numpy() - self.assertFalse(np.array_equal(s1, s2)) - full = Tensor(w).multinomial(len(w), replacement=False).numpy() - self.assertEqual(sorted(full.tolist()), w) - - w = [0.1, 0.2, 0.3, 0.4] - @TinyJit - def sample_three(): return Tensor(w).multinomial(3, replacement=False).realize() - - tiny_draws = np.array([sample_three().numpy() for _ in range(1000)]) - torch_draws = np.array([torch.tensor(w).multinomial(3, replacement=False).numpy() for _ in range(1000)]) - for pos in range(3): - self.assertTrue(equal_distribution(lambda *_: Tensor(tiny_draws[:, pos]), lambda _: torch.tensor(torch_draws[:, pos]))) - - @unittest.skip("this test is flaky") - def test_multinomial_counterexample(self): - tiny_res = Tensor([0.3, 0.6, 0.1]).multinomial(4000, replacement=True) - torch_res = torch.tensor([0.3, 0.6, 0.1]).multinomial(4000, replacement=True) - self.assertTrue(equal_distribution(lambda *_: tiny_res, lambda _: torch_res)) - torch_res = torch.tensor([0.2, 0.7, 0.1]).multinomial(4000, replacement=True) - self.assertFalse(equal_distribution(lambda *_: tiny_res, lambda _: torch_res)) - - def test_conv2d_init(self): - params = (128, 256, (3,3)) - assert equal_distribution(lambda *_: nn.Conv2d(*params).weight, lambda _: torch.nn.Conv2d(*params).weight.detach()) - assert equal_distribution(lambda *_: nn.Conv2d(*params).bias, lambda _: torch.nn.Conv2d(*params).bias.detach()) - - def test_linear_init(self): - params = (64, 256) - assert equal_distribution(lambda *_: nn.Linear(*params).weight, lambda _: torch.nn.Linear(*params).weight.detach()) - assert equal_distribution(lambda *_: nn.Linear(*params).bias, lambda _: torch.nn.Linear(*params).bias.detach()) - - def test_bn_init(self): - params = (64,) - assert equal_distribution(lambda *_: nn.BatchNorm2d(*params).weight, lambda _: torch.nn.BatchNorm2d(*params).weight.detach()) - assert equal_distribution(lambda *_: nn.BatchNorm2d(*params).bias, lambda _: torch.nn.BatchNorm2d(*params).bias.detach()) - - def test_rand_chain(self): - # NOTE: this fails if property propagates deeper than stack limit - for _ in range(833): Tensor.rand(1) - Tensor.rand(1).realize() - def test_random_counter_overflow(self): device = Device.DEFAULT Tensor.manual_seed(1337) @@ -432,20 +258,5 @@ def test_random_counter_overflow(self): c = Tensor._device_rng_counters[device].numpy() np.testing.assert_allclose(c, [14, 1]) -# TODO: still fails with MAX_KERNEL_BUFFERS -@unittest.skipIf(Device.DEFAULT == "WEBGPU" and not OSX, "WEBGPU Vulkan can only run kernels with up to 10 buffers") -class TestSample(unittest.TestCase): - def test_sample(self): - X = Tensor.rand(1000, 50).realize() - BS = 16 - idxs = np.random.randint(0, X.shape[0], size=(BS)) - # this uncovered a bug with arg sort order - batch = [Variable(f'idx{i}', 0, X.shape[0]-1).bind(s) for i,s in enumerate(idxs.tolist())] - x = Tensor.cat(*[X.shrink(((batch[i], batch[i]+1), None)) for i in range(BS)]) - print(idxs) - ret = x.numpy() - base = X.numpy()[idxs] - np.testing.assert_equal(ret, base) - if __name__ == "__main__": unittest.main() diff --git a/test/backend/test_rangeify.py b/test/backend/test_rangeify.py index b4ec084a1bb1b..9e0ca0d3583ac 100644 --- a/test/backend/test_rangeify.py +++ b/test/backend/test_rangeify.py @@ -109,9 +109,9 @@ def fa(): def fa_bw(): Tensor.manual_seed(1337) with Context(DEBUG=0): - q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize().requires_grad_() for _ in range(3)] + q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize() for _ in range(3)] attn_output = nn.Linear(HEADS*EMB, HEADS*EMB, bias=False) - attn_output.weight.requires_grad_().realize() + attn_output.weight.realize() target = Tensor.rand(BS, SEQLEN, HEADS*EMB).contiguous().realize() GlobalCounters.reset() diff --git a/test/backend/test_renderer_failures.py b/test/backend/test_renderer_failures.py index 383d54e0b6915..5b7945c922f8b 100644 --- a/test/backend/test_renderer_failures.py +++ b/test/backend/test_renderer_failures.py @@ -1,6 +1,6 @@ import unittest import numpy as np -from tinygrad.device import Device, is_dtype_supported +from tinygrad.device import Device from tinygrad.dtype import dtypes, ConstType from tinygrad.engine.realize import run_linear from tinygrad.codegen import to_program @@ -22,8 +22,8 @@ def _test_uop_result(inputs:list[Tensor], sink:UOp, local_size=None): def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp): dtype = alu_src_uops[0].dtype - a = UOp(Ops.PARAM, dtype.ptr(), (), 0) - b = UOp(Ops.PARAM, dtype.ptr(), (), 1) + a = UOp.param(0, dtype.ptr()) + b = UOp.param(1, dtype.ptr()) idx = UOp.const(dtypes.int, 0) ld = b.index(idx) alu = ld.alu(alu_op, *alu_src_uops) @@ -33,7 +33,7 @@ def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp): class TestRendererFailures(unittest.TestCase): @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer") def test_gated_store_with_alu(self): - a = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0) + a = UOp.param(0, dtypes.int.ptr()) gate_alu = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0) gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index(lidx0.valid(gate_alu)), UOp.const(dtypes.int, 1))) sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,), arg=KernelInfo()) @@ -42,7 +42,7 @@ def test_gated_store_with_alu(self): @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer") def test_gated_store_with_alu_2d(self): - a = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0) + a = UOp.param(0, dtypes.int.ptr()) gate_alu_0 = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0) gate_alu_1 = (lidx1:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 2),), 'lidx1')).ne(0) gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index((lidx0+lidx1*4).valid(gate_alu_0&gate_alu_1)), UOp.const(dtypes.int, 1))) @@ -87,7 +87,7 @@ def test_multiply_infinity(self): class TestPTXFailures(unittest.TestCase): @unittest.skip("INDEX can only have a gate ALU parent, not an IF") def test_gated_store_with_if(self): - a = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0) + a = UOp.param(0, dtypes.int.ptr()) gate_alu = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0) val = UOp.const(dtypes.int, 1) if_uop = UOp(Ops.IF, dtypes.void, (gate_alu,)) @@ -96,7 +96,7 @@ def test_gated_store_with_if(self): ret = _test_uop_result([], sink, local_size=[4, 1, 1])[0] np.testing.assert_equal(ret, [0, 1, 1, 1]) - @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half") + @unittest.skipUnless(dtypes.half in Device[Device.DEFAULT].renderer.supported_dtypes(), "need half") def test_gated_define_acc_with_half_dtype(self): a = Tensor.randn(32, 32, dtype=dtypes.half).realize() b = Tensor.randn(34, 32, dtype=dtypes.half).realize() diff --git a/test/backend/test_schedule.py b/test/backend/test_schedule.py index 6cf47b26f6049..3f1fffafe03cc 100644 --- a/test/backend/test_schedule.py +++ b/test/backend/test_schedule.py @@ -5,15 +5,16 @@ import gc, unittest, functools import numpy as np from typing import cast -from hypothesis import assume, given, settings, strategies as strat +from hypothesis import assume, given, strategies as strat from tinygrad import nn, dtypes, Device, Tensor, Variable -from tinygrad.device import is_dtype_supported from tinygrad.dtype import DType from tinygrad.uop.ops import UOp, Ops, UPat -from tinygrad.helpers import CI, DEBUG, OSX, GlobalCounters, Context, getenv, all_same, temp +from tinygrad.helpers import DEBUG, DEV, OSX, GlobalCounters, Context, getenv, all_same, temp from tinygrad.engine.realize import compile_linear, run_linear +supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes() + class KernelCountException(Exception): pass def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Tensor]|None=None, filter_sink=True): if to_prerealize: @@ -44,8 +45,8 @@ def _test_conv2d(allowed:int, dtype:DType=dtypes.float): dtypes.default_float = dtype Tensor.manual_seed(0) BS, CIN = 2, 3 - img = Tensor.randn(BS, CIN, 64, 64, requires_grad=True).realize() - w = Tensor.uniform(16, CIN, 3, 3, requires_grad=True).realize() + img = Tensor.randn(BS, CIN, 64, 64).realize() + w = Tensor.uniform(16, CIN, 3, 3).realize() ret = Tensor.conv2d(img, w).relu().mean().backward() dtypes.default_float = old_default_float linear, var_vals = Tensor.linear_with_vars(ret, img.grad, w.grad) @@ -81,13 +82,6 @@ def test_arange_avgpool2d(self, kcount=1): def test_arange_avgpool2d_fused_noopt(self): with Context(NOOPT=1): self.test_arange_avgpool2d(kcount=1) - # linearizer error - @unittest.skip("recursion error no longer raised") - @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "needs supports_float4 to fail") - def test_arange_avgpool2d_fused(self): - with self.assertRaises(RecursionError): - with Context(NOOPT=0): self.test_arange_avgpool2d(kcount=1) - # when we're fusing a reduce, all ReduceOps must have the same N in the dimensions # all permutes, reshapes, expands and shrinks push through the reduce def test_arange_sum(self): @@ -105,7 +99,7 @@ def test_permute_arange(self): run_linear(*check_schedule(a, 1)) self.assertListEqual(a.tolist(), [[15]]) - @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half") + @unittest.skipUnless(dtypes.half in supported_dtypes, "need half") @unittest.skipIf(Device.DEFAULT == "WEBGPU" and OSX, "WEBGPU Metal backend is not accurate enough") def test_expand_buffer_before_cast(self): a = Tensor.randn(4, 2, 1).realize().permute((1, 0, 2)) @@ -113,13 +107,6 @@ def test_expand_buffer_before_cast(self): run_linear(*check_schedule(b, 1)) np.testing.assert_allclose(b.numpy(), np.broadcast_to(a.numpy().astype(np.float16), (2, 4, 4))+2, rtol=1e-3) - def test_indexing_scalars_simple(self): - X = Tensor.randn(2, 2).realize() - xt = X[Tensor(1)][Tensor(0)] - run_linear(*check_schedule(xt, 1)) - np.testing.assert_equal(xt.numpy(), X.numpy()[1][0]) - - @unittest.skipIf(CI and Device.DEFAULT == "NV", "crashes on NV CI") def test_add_chain_buffers(self): N = 31 with Context(TRACK_MATCH_STATS=0, DEBUG=0): @@ -130,14 +117,14 @@ def test_add_chain_buffers(self): root = root + functools.reduce(lambda a,b:a+b, bufs[i:i+X]) self.assertEqual(root.item(), sum(range(N))) - @given(strat.sampled_from(range(2,4)), strat.sampled_from(range(2,4)), strat.sampled_from(range(0,4)), strat.sampled_from(range(0,4))) - @settings(deadline=None) - def test_indexing_scalars(self, x, y, a, b): - assume(a Tensor: # only support axis =-1 @@ -62,7 +61,7 @@ def test_fuse_gemm_softmax(self): b = Tensor.rand(50,50).realize() self._test_fuse(lambda a,b: ((a@b).relu()+a).contiguous().softmax(axis=-1), a,b, allow_multiple=True) - @unittest.skipUnless(is_dtype_supported(dtypes.float16), f"no float16 on {Device.DEFAULT}") + @unittest.skipUnless(dtypes.float16 in Device[Device.DEFAULT].renderer.supported_dtypes(), f"no float16 on {Device.DEFAULT}") @unittest.skip("needs RANGEIFY>1") def test_fuse_softmax_dtype(self): a = Tensor.rand(50,50).realize() @@ -190,7 +189,6 @@ def test_auto_softmax(self): def test_softmax_bw(self): print("*** softmax bw ***") - self.test.requires_grad_() with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)): self.test.softmax(-1).sum().backward() sg = self.test.grad.realize() diff --git a/test/backend/test_subbuffer.py b/test/backend/test_subbuffer.py index 6755cf4a4a6b2..d6d866597adc9 100644 --- a/test/backend/test_subbuffer.py +++ b/test/backend/test_subbuffer.py @@ -36,7 +36,7 @@ def test_subbuffer_len(self): assert len(mv) == 5 def test_subbuffer_used(self): - t = Tensor.arange(0, 10, dtype=dtypes.uint8).realize() + t = Tensor.arange(0, 10, dtype=dtypes.uint8).clone().realize() vt = t[2:4].realize() out = (vt + 100).tolist() assert out == [102, 103] @@ -44,7 +44,7 @@ def test_subbuffer_used(self): @needs_second_gpu @unittest.skipIf(Device.DEFAULT not in {"CUDA", "NV", "AMD"} or DEV.interface.startswith("MOCK"), "only NV, AMD, CUDA") def test_subbuffer_transfer(self): - t = Tensor.arange(0, 10, dtype=dtypes.uint8).realize() + t = Tensor.arange(0, 10, dtype=dtypes.uint8).clone().realize() vt = t[2:5].contiguous().realize() out = vt.to(f"{Device.DEFAULT}:1").realize().tolist() assert out == [2, 3, 4] diff --git a/test/backend/test_tensor.py b/test/backend/test_tensor.py index 4135217f58b2d..4494da0423e49 100644 --- a/test/backend/test_tensor.py +++ b/test/backend/test_tensor.py @@ -5,8 +5,8 @@ from tinygrad.helpers import getenv, temp, mv_address from extra.gradcheck import numerical_jacobian, jacobian, gradcheck from hypothesis import given, settings, strategies as strat -from tinygrad.device import is_dtype_supported from tinygrad.dtype import DTYPES_DICT +from tinygrad.uop.ops import UOp settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False)) settings.load_profile("my_profile") @@ -23,6 +23,19 @@ def test_zerodim_initialization(self): self.assertEqual(Tensor(55).shape, ()) self.assertEqual(Tensor(3.14).shape, ()) + def test_deviceless_const_construct_device_repr(self): + t = Tensor(UOp.const(dtypes.float, 2.0)) + self.assertIsNone(t.uop.device) + self.assertIsNone(t.device) + self.assertIn(" UOp: if op is Ops.CONST: uops.append(UOp.const(dtype, arg)) + elif op is Ops.PARAM: uops.append(UOp.param(arg, dtype).replace(src=())) else: uops.append(UOp(op, dtype, tuple(src), arg)) return uops[-1] @@ -141,9 +141,7 @@ def test_mod_int32(self): lambda a,b: abs(int(a))%abs(int(b))*(1,-1)[a<0], (dtypes.int32, dtypes.int32), no_b_zero=True) def test_cmplt_int32(self): self._test_bop_fxn(Ops.CMPLT, lambda a,b: int(a) bool: backend_test.exclude('test_resize_downsample_scales_cubic_align_corners_cpu') # about different dtypes -if not is_dtype_supported(dtypes.float64): +if dtypes.float64 not in Device[Device.DEFAULT].renderer.supported_dtypes(): backend_test.exclude('float64') backend_test.exclude('DOUBLE') # these have float64 inputs @@ -80,7 +79,7 @@ def supports_device(cls, device: str) -> bool: backend_test.exclude('test_einsum_*') backend_test.exclude('test_cumsum_*') -if not is_dtype_supported(dtypes.float16): +if dtypes.float16 not in Device[Device.DEFAULT].renderer.supported_dtypes(): backend_test.exclude('float16') backend_test.exclude('FLOAT16') diff --git a/test/external/external_test_onnx_ops.py b/test/external/external_test_onnx_ops.py index 9f0f39912f69b..81f8279da4f87 100644 --- a/test/external/external_test_onnx_ops.py +++ b/test/external/external_test_onnx_ops.py @@ -4,7 +4,7 @@ from typing import Any import unittest, onnx, tempfile -from tinygrad import dtypes, Tensor +from tinygrad import dtypes, Tensor, Context from tinygrad.nn.onnx import OnnxRunner import numpy as np from extra.onnx_helpers import validate @@ -284,22 +284,23 @@ def test_dynamic_quantize_linear(self): def test_qlinear_conv(self): for dtype, zero_point in [(np.uint8, 128), (np.int8, 0)]: for b in (np.ones([32], dtype=np.int32), np.zeros([32], dtype=np.int32)): - with self.subTest(dtype=dtype, zero_point=zero_point): - dtype_min, dtype_max = np.iinfo(dtype).min, np.iinfo(dtype).max - inputs = { - "x": np.random.randint(dtype_min, dtype_max + 1, [1, 3, 224, 224], dtype=dtype), - "x_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32), - "x_zero_point": np.array(zero_point, dtype=dtype), - "w": np.random.randint(dtype_min, dtype_max + 1, [32, 3, 3, 3], dtype=dtype), - "w_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32), - "w_zero_point": np.array(zero_point, dtype=dtype), - "y_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32), - "y_zero_point": np.array(zero_point, dtype=dtype), - "b": b - } - attributes = {'auto_pad': 'NOTSET', 'dilations': (1, 1), 'group': 1, 'kernel_shape': (3, 3), 'pads': (1, 1, 1, 1), 'strides': (2, 2)} - outputs = ["out"] - self.helper_test_single_op("QLinearConv", inputs, attributes, outputs, atol=1) # occasionally inaccurate + for channel_shape in [(), (32,)]: + with self.subTest(dtype=dtype, zero_point=zero_point, channel_shape=channel_shape): + dtype_min, dtype_max = np.iinfo(dtype).min, np.iinfo(dtype).max + inputs = { + "x": np.random.randint(dtype_min, dtype_max + 1, [1, 3, 224, 224], dtype=dtype), + "x_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32), + "x_zero_point": np.array(zero_point, dtype=dtype), + "w": np.random.randint(dtype_min, dtype_max + 1, [32, 3, 3, 3], dtype=dtype), + "w_scale": np.random.uniform(0.01, 0.1, channel_shape).astype(np.float32), + "w_zero_point": np.full(channel_shape, zero_point, dtype=dtype), + "y_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32), + "y_zero_point": np.array(zero_point, dtype=dtype), + "b": b + } + attributes = {'auto_pad': 'NOTSET', 'dilations': (1, 1), 'group': 1, 'kernel_shape': (3, 3), 'pads': (1, 1, 1, 1), 'strides': (2, 2)} + outputs = ["out"] + self.helper_test_single_op("QLinearConv", inputs, attributes, outputs, atol=1) # occasionally inaccurate def test_qlinear_matmul(self): for dtype, zero_point in [(np.uint8, 128), (np.int8, 0)]: @@ -364,6 +365,19 @@ def test_reduce_l2_half(self): inputs = {"data": np.random.randn(1, 1, 32, 32, 32).astype(np.half)*100} self.helper_test_single_op("ReduceL2", inputs, {}, ["reduced"]) + def test_same_device_as_input(self): + from tinygrad.nn.onnx import onnx_ops + EyeLike = onnx_ops["EyeLike"] + Shape = onnx_ops["Shape"] + Compress = onnx_ops["Compress"] + with Context(DEV="CPU"): + x = Tensor.arange(4, device="PYTHON").reshape(2,2) + self.assertEqual(EyeLike(x).device, x.device) + self.assertEqual(Shape(x).device, x.device) + out = Compress(x, [True, False, True, False]) + self.assertEqual(out.device, x.device) + self.assertEqual(out.tolist(), [0, 2]) + class TestTrainingOnnxOps(TestOnnxOps): # NOTE: ORT doesn't actually support training ops on cpu so we test using functions provided by onnx DOMAIN = AI_ONNX_PREVIEW_TRAINING_DOMAIN @@ -581,5 +595,28 @@ def test_qlinear_global_average_pool(self): outputs = ["C"] self.helper_test_single_op("QLinearGlobalAveragePool", inputs, attributes, outputs) + def test_same_device_as_input(self): + from tinygrad.nn.onnx import onnx_ops, OpSetId, Domain + EmbedLayerNormalization = onnx_ops["EmbedLayerNormalization"] + Attention = onnx_ops["Attention"] + with Context(DEV="CPU"): + input_ids = Tensor([[1, 2]], device="PYTHON", dtype=dtypes.int32) + segment_ids = Tensor([[0, 0]], device="PYTHON", dtype=dtypes.int32) + word = Tensor.ones(4, 3, device="PYTHON") + pos = Tensor.ones(5, 3, device="PYTHON") + seg = Tensor.ones(1, 3, device="PYTHON") + gamma, beta = Tensor.ones(3, device="PYTHON"), Tensor.zeros(3, device="PYTHON") + out, _, _ = EmbedLayerNormalization(input_ids, segment_ids, word, pos, seg, gamma, beta) + self.assertEqual(out.device, input_ids.device) + out.realize() + + attn = Attention[OpSetId(Domain.MICROSOFT_CONTRIB_OPS, 1)] + x = Tensor.ones(1, 2, 4, device="PYTHON") + w = Tensor.ones(4, 12, device="PYTHON") + mask = Tensor([2, 0], device="PYTHON", dtype=dtypes.int32) + out, _ = attn(x, w, mask_index=mask, num_heads=1, unidirectional=1) + self.assertEqual(out.device, x.device) + out.realize() + if __name__ == "__main__": unittest.main() diff --git a/test/external/external_test_onnx_runner.py b/test/external/external_test_onnx_runner.py index 03a29bf605497..193a30b882500 100644 --- a/test/external/external_test_onnx_runner.py +++ b/test/external/external_test_onnx_runner.py @@ -2,7 +2,6 @@ import numpy as np from tinygrad import Tensor from tinygrad.uop.ops import Ops -from tinygrad.device import is_dtype_supported from typing import Any from tinygrad.nn.onnx import OnnxRunner, OnnxPBParser, OnnxDataType from hypothesis import given, strategies as st @@ -89,7 +88,6 @@ def test_external_data_loading(self): np.testing.assert_equal(output.numpy(), weights + 1) all_dtypes = list(OnnxDataType) -device_supported_dtypes = {odt for odt in OnnxDataType if is_dtype_supported(odt.to_dtype())} class TestOnnxRunnerDtypes(unittest.TestCase): """ diff --git a/test/external/external_test_optim.py b/test/external/external_test_optim.py index 6e809f1b957ba..087ec04681da9 100644 --- a/test/external/external_test_optim.py +++ b/test/external/external_test_optim.py @@ -21,8 +21,8 @@ class TinyNet: def __init__(self): - self.x = Tensor(x_init.copy(), requires_grad=True) - self.W = Tensor(W_init.copy(), requires_grad=True) + self.x = Tensor(x_init.copy()) + self.W = Tensor(W_init.copy()) self.m = Tensor(m_init.copy()) def forward(self): diff --git a/test/external/external_uop_gc.py b/test/external/external_uop_gc.py index 10ffa84de864f..e52500fc195b4 100644 --- a/test/external/external_uop_gc.py +++ b/test/external/external_uop_gc.py @@ -24,28 +24,28 @@ def two_plus_two_linearize(): def two_plus_two_realize(): (Tensor([2])+Tensor([2])).realize() def two_plus_two_item(): (Tensor([2])+Tensor([2])).item() def gradient_test(): - x = Tensor.eye(3, requires_grad=True) - y = Tensor([[2.0,0,-2.0]], requires_grad=True) + x = Tensor.eye(3) + y = Tensor([[2.0,0,-2.0]]) z = y.matmul(x).sum() z.backward() def realized_eye(): - Tensor.eye(3, requires_grad=True).realize() + Tensor.eye(3).clone().realize() def realized_list(): - Tensor([[2.0,0,-2.0]], requires_grad=True).realize() + Tensor([[2.0,0,-2.0]]).realize() def kernel_matmul(): - x = Tensor.eye(3, requires_grad=True) - y = Tensor([[2.0,0,-2.0]], requires_grad=True) + x = Tensor.eye(3) + y = Tensor([[2.0,0,-2.0]]) z = y.matmul(x) linear = z.schedule_linear() to_program(linear.src[-1].src[0], Device.default.renderer) def realized_matmul(): - x = Tensor.eye(3, requires_grad=True) - y = Tensor([[2.0,0,-2.0]], requires_grad=True) + x = Tensor.eye(3) + y = Tensor([[2.0,0,-2.0]]) z = y.matmul(x) Tensor.realize(z) def realized_gradient(): - x = Tensor.eye(3, requires_grad=True) - y = Tensor([[2.0,0,-2.0]], requires_grad=True) + x = Tensor.eye(3) + y = Tensor([[2.0,0,-2.0]]) z = y.matmul(x).sum() z.backward() Tensor.realize(x, y, z, x.grad, y.grad) diff --git a/test/external/fuzz_fast_idiv.py b/test/external/fuzz_fast_idiv.py index 546b8e5447ef9..88c45111d3740 100644 --- a/test/external/fuzz_fast_idiv.py +++ b/test/external/fuzz_fast_idiv.py @@ -1,7 +1,6 @@ import random import z3 from tinygrad import dtypes, Device -from tinygrad.helpers import DEV from tinygrad.uop.validate import uops_to_z3, z3_cdiv from tinygrad.uop.ops import UOp from tinygrad.uop.decompositions import fast_idiv @@ -16,7 +15,7 @@ u = UOp.variable('x', random.randint(dt.min, 0), random.randint(1, dt.max), dtype=dt) d = random.randint(1, max(1, u.arg[2])*2) if d in powers_of_two: continue - expr = fast_idiv(DEV.target(Device.DEFAULT), u, d) + expr = fast_idiv(Device[Device.DEFAULT].renderer, u, d) if expr is None: continue solver = z3.Solver() diff --git a/test/external/mlperf_bert/external_benchmark_bert.py b/test/external/mlperf_bert/external_benchmark_bert.py index 67dff59c4417a..01d6f83f43f9d 100644 --- a/test/external/mlperf_bert/external_benchmark_bert.py +++ b/test/external/mlperf_bert/external_benchmark_bert.py @@ -59,7 +59,7 @@ def step(inputs): best_tm = None flops, mem_used, mem, kernels = None, None, None, None for _ in range(CNT): - with Context(TRACK_MATCH_STATS=0): inputs = [Tensor.randn(*shape, requires_grad=False).realize() for shape in input_shapes] + with Context(TRACK_MATCH_STATS=0): inputs = [Tensor.randn(*shape).realize() for shape in input_shapes] GlobalCounters.reset() st = time.perf_counter() diff --git a/test/external/process_replay/process_replay.py b/test/external/process_replay/process_replay.py index 7d63217d73977..df34060b1343a 100755 --- a/test/external/process_replay/process_replay.py +++ b/test/external/process_replay/process_replay.py @@ -4,7 +4,7 @@ from dataclasses import replace from typing import Callable, Any -ASSERT_DIFF = int((flag:="[pr]") in os.getenv("COMMIT_MESSAGE", flag) or flag in os.getenv("PR_TITLE", flag)) +ASSERT_DIFF = int((flag:="[PR]") in os.getenv("COMMIT_MESSAGE", flag) or flag in os.getenv("PR_TITLE", flag)) if not int(os.getenv("ASSERT_PROCESS_REPLAY", "1")): ASSERT_DIFF = 0 try: diff --git a/test/helpers.py b/test/helpers.py index c7874c2409953..c93bd451faf6f 100644 --- a/test/helpers.py +++ b/test/helpers.py @@ -8,7 +8,7 @@ from tinygrad.codegen import to_program from tinygrad.dtype import DType from tinygrad.nn.state import get_parameters -from tinygrad.helpers import T, CI, Target +from tinygrad.helpers import T, Target, DEV from tinygrad.renderer import Renderer from tinygrad.codegen import full_rewrite_to_sink, line_rewrite, pm_linearize_cleanups from tinygrad.codegen.late.linearizer import linearize @@ -82,7 +82,7 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None, vals:tuple for buf_dt, data in inputs or []: bufs.append(buf:=allocator.alloc(len(data) * buf_dt.itemsize)) allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *data))) - g = UOp(Ops.PARAM, uop.dtype.ptr(), arg=0, src=()) + g = UOp.param(0, uop.dtype.ptr()) prg = to_program(UOp.store(g.index(UOp.const(dtypes.int, 0)), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON"))) prog = PythonProgram("run", PythonCompiler().compile(prg.src[3].arg)) prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *bufs, vals=vals) @@ -97,7 +97,7 @@ def to_uops_list(u:list[UOp], ren=None) -> list[UOp]: def not_support_multi_device(): # CL and CUDA don't support multi device if in CI - return CI and Device.DEFAULT in ("CL", "CUDA") + return (Device.DEFAULT == "CL" and Device[Device.DEFAULT].count() < 2) or (Device.DEFAULT == "CUDA" and DEV.interface.startswith("MOCK")) def needs_second_gpu(fn): @functools.wraps(fn) diff --git a/test/mockgpu/amd/emu.py b/test/mockgpu/amd/emu.py index d8cbc20bfb248..b23c8021e83bd 100644 --- a/test/mockgpu/amd/emu.py +++ b/test/mockgpu/amd/emu.py @@ -69,7 +69,7 @@ def __exit__(self, *args): from tinygrad.renderer.amd.dsl import VCC_LO, EXEC_LO, SCC, ttmp from tinygrad.runtime.autogen.amd.common import Fmt, OpType from test.amd.helpers import decode_dpp16 -from test.mockgpu.amd.pcode import parse_block, _FUNCS, _set_bits, _val_to_bits +from test.mockgpu.amd.pcode import parse_block, _FUNCS, _set_bits, _to_bool, _val_to_bits MASK32 = 0xFFFFFFFF @@ -340,7 +340,7 @@ def get_pcode(op) -> str: pcode = pcode.replace('VCC = 0x0LL', 'VCC.u64[laneId] = 0').replace('VCC = 0x1LL', 'VCC.u64[laneId] = 1') return pcode -def parse_pcode(pcode: str, srcs: dict[str, UOp | int] | None = None) -> tuple[dict, list[tuple[str, UOp]]]: +def parse_pcode(pcode: str, srcs: dict[str, UOp | int] | None = None) -> tuple[dict, list]: env: dict = srcs.copy() if srcs else {} assigns: list[tuple[str, UOp]] = [] raw_lines = [l.strip().rstrip(';') for l in pcode.split('\n') if l.strip() and not l.strip().startswith('//')] @@ -423,10 +423,10 @@ def _collect_data_slices(assigns: list[tuple[str, UOp]], data_prefix: str, pcode class _Ctx: """Context for instruction compilation - holds buffers and helpers.""" __slots__ = ('inst_size', 'dyn_fields', '_axis_id', 'wave_size', 'vgpr', 'accvgpr') - sgpr = UOp(Ops.PARAM, dtypes.uint32.ptr(SGPR_COUNT), arg=0) - vmem = UOp(Ops.PARAM, dtypes.uint32.ptr(1 << 46), arg=2) - lds = UOp(Ops.PARAM, dtypes.uint32.ptr(16384), arg=3) - scratch = UOp(Ops.PARAM, dtypes.uint8.ptr(1 << 30), arg=4) + sgpr = UOp.param(0, dtypes.uint32.ptr(SGPR_COUNT)) + vmem = UOp.param(2, dtypes.uint32.ptr(1 << 46)) + lds = UOp.param(3, dtypes.uint32.ptr(16384)) + scratch = UOp.param(4, dtypes.uint8.ptr(1 << 30)) # Cache PARAM UOps by wave_size so all _Ctx instances with same wave_size share identical UOp references _vgpr_cache: dict[int, UOp] = {} _accvgpr_cache: dict[int, UOp] = {} @@ -434,10 +434,10 @@ class _Ctx: def __init__(self, inst_size: int, wave_size: int = 32): self.inst_size, self._axis_id, self.wave_size = inst_size, 0, wave_size self.dyn_fields: list[tuple[int, int]] = [] # (lo, hi) of fields read dynamically - if wave_size not in _Ctx._vgpr_cache: _Ctx._vgpr_cache[wave_size] = UOp(Ops.PARAM, dtypes.uint32.ptr(256 * wave_size), arg=1) + if wave_size not in _Ctx._vgpr_cache: _Ctx._vgpr_cache[wave_size] = UOp.param(1, dtypes.uint32.ptr(256 * wave_size)) self.vgpr = _Ctx._vgpr_cache[wave_size] if wave_size == 64: - if wave_size not in _Ctx._accvgpr_cache: _Ctx._accvgpr_cache[wave_size] = UOp(Ops.PARAM, dtypes.uint32.ptr(256 * wave_size), arg=5) + if wave_size not in _Ctx._accvgpr_cache: _Ctx._accvgpr_cache[wave_size] = UOp.param(5, dtypes.uint32.ptr(256 * wave_size)) self.accvgpr = _Ctx._accvgpr_cache[wave_size] else: self.accvgpr = self.vgpr @@ -523,7 +523,7 @@ def wsgpr_dyn(self, reg: UOp, val: UOp) -> UOp: """Write SGPR with dynamic register index. On RDNA, index 124 = NULL (writes discarded). On CDNA, index 124 = M0 (read/write).""" # RDNA: NULL (124) discards writes. CDNA: M0 (124) is writable. valid = None if self.wave_size == 64 else reg.ne(_c(124)) - return self.sgpr.index(reg.cast(dtypes.int).valid(valid) if valid is not None else reg.cast(dtypes.int)).store(val.cast(dtypes.uint32)) + return self.sgpr.index(reg.cast(dtypes.int).valid(valid) if valid is not None else reg.cast(dtypes.int), ptr=True).store(val.cast(dtypes.uint32)) def wmask(self, reg: UOp, val: UOp) -> list[UOp]: """Write a lane mask (VCC/EXEC). Splits into lo/hi for wave64.""" @@ -598,12 +598,13 @@ def rsrc_dyn(self, off: UOp, lane: UOp | None, bits: int = 32, literal: UOp | No def rpc(self) -> UOp: """Read PC as 64-bit byte address.""" # Index at PC_LO, then cast to uint64 ptr and load - return self.sgpr.index(_c(PC_LO_IDX, dtypes.int), ptr=True).cast(dtypes.uint64.ptr(SGPR_COUNT // 2)).load() + return _u64(self.rsgpr_dyn(_c(PC_LO_IDX)), self.rsgpr_dyn(_c(PC_HI_IDX))) def inc_pc(self) -> list[UOp]: """Increment PC by instruction size in bytes. Returns [store].""" new_pc = self.rpc() + UOp.const(dtypes.uint64, self.inst_size) - return [self.sgpr.index(_c(PC_LO_IDX, dtypes.int), ptr=True).cast(dtypes.uint64.ptr(SGPR_COUNT // 2)).store(new_pc)] + lo, hi = _split64(new_pc) + return [self.wsgpr_dyn(_c(PC_LO_IDX), lo), self.wsgpr_dyn(_c(PC_HI_IDX), hi)] def scalar_stores(self, assigns: list[tuple[str, UOp]], sdst_reg: UOp, sdst_size: int = 1) -> list[UOp]: """Generate stores for scalar assigns with dynamic destination register (D0, SCC, EXEC, VCC).""" @@ -713,6 +714,7 @@ def to_wide(x): return (x.bitcast(narrow_dt) if x.dtype.itemsize == narrow_dt.it old = self.vgpr.index(val[0].cast(dtypes.int), ptr=True).load() new_val = _set_bits(old, _val_to_bits(val[1]), width, lo_bit).cast(dtypes.uint32) active = _lane_active(exec_mask, lane) + if len(val) > 4: active = active & _to_bool(val[4]) raw_stores.append(('vgpr_direct', self.vgpr.index(val[0].cast(dtypes.int).valid(active)).store(new_val))) continue if 'D0' in dest and '[laneId]' in dest: diff --git a/test/mockgpu/amd/pcode.py b/test/mockgpu/amd/pcode.py index 3a97f6885febe..686f00623325d 100644 --- a/test/mockgpu/amd/pcode.py +++ b/test/mockgpu/amd/pcode.py @@ -105,9 +105,9 @@ def _bf8_to_f32(v: UOp) -> UOp: return is_sub.where(sub_f32.bitcast(dtypes.float32), normal) def _f32_to_fp8(v: UOp) -> UOp: - return f2f((v.bitcast(dtypes.float32) if v.dtype != dtypes.float32 else v).bitcast(dtypes.uint32), dtypes.float32, dtypes.fp8e4m3) + return f2f((v.bitcast(dtypes.float32) if v.dtype != dtypes.float32 else v).bitcast(dtypes.uint32), dtypes.float32, dtypes.fp8e4m3, sat=False) def _f32_to_bf8(v: UOp) -> UOp: - return f2f((v.bitcast(dtypes.float32) if v.dtype != dtypes.float32 else v).bitcast(dtypes.uint32), dtypes.float32, dtypes.fp8e5m2) + return f2f((v.bitcast(dtypes.float32) if v.dtype != dtypes.float32 else v).bitcast(dtypes.uint32), dtypes.float32, dtypes.fp8e5m2, sat=False) def _f32_to_bf16(v: UOp) -> UOp: """Convert f32 to bf16 with round-to-nearest-even. BF16 is the upper 16 bits of F32 with rounding.""" bits = (v.bitcast(dtypes.float32) if v.dtype != dtypes.float32 else v).bitcast(dtypes.uint32) diff --git a/test/models/test_end2end.py b/test/models/test_end2end.py index 7d545e539f61d..452e0d0704e7f 100644 --- a/test/models/test_end2end.py +++ b/test/models/test_end2end.py @@ -1,19 +1,18 @@ import torch from torch import nn -import unittest +import unittest, sys import numpy as np from tinygrad.nn.state import get_parameters, get_state_dict from tinygrad.nn import optim, Linear, Conv2d, BatchNorm2d from tinygrad.tensor import Tensor from extra.datasets import fetch_mnist -from tinygrad.helpers import CI def compare_tiny_torch(model, model_torch, X, Y): with Tensor.train(): model_torch.train() model_state_dict = get_state_dict(model) for k,v in model_torch.named_parameters(): - if not CI: print(f"initting {k} from torch") + if sys.stdout.isatty(): print(f"initting {k} from torch") model_state_dict[k].assign(Tensor(v.detach().numpy())).realize() optimizer = optim.SGD(get_parameters(model), lr=0.001) @@ -35,14 +34,14 @@ def compare_tiny_torch(model, model_torch, X, Y): loss_torch.backward() # assert losses match - if not CI: print(loss.realize().numpy()) - if not CI: print(loss_torch.detach().numpy()) + if sys.stdout.isatty(): print(loss.realize().numpy()) + if sys.stdout.isatty(): print(loss_torch.detach().numpy()) np.testing.assert_allclose(loss.realize().numpy(), loss_torch.detach().numpy(), atol=1e-4) for k,v in list(model_torch.named_parameters())[::-1]: g = model_state_dict[k].grad.numpy() gt = v.grad.detach().numpy() - if not CI: print("testing grads", k, model_state_dict[k].grad.dtype) + if sys.stdout.isatty(): print("testing grads", k, model_state_dict[k].grad.dtype) np.testing.assert_allclose(g, gt, atol=1e-3, err_msg=f'grad mismatch {k}') # take the steps @@ -51,7 +50,7 @@ def compare_tiny_torch(model, model_torch, X, Y): # assert weights match for k,v in model_torch.named_parameters(): - if not CI: print("testing weight", k, model_state_dict[k].dtype) + if sys.stdout.isatty(): print("testing weight", k, model_state_dict[k].dtype) np.testing.assert_allclose(model_state_dict[k].numpy(), v.detach().numpy(), atol=1e-3, err_msg=f'weight mismatch {k}') def get_mnist_data(): diff --git a/test/models/test_whisper.py b/test/models/test_whisper.py index d6b21e4d4363c..818994dc86677 100644 --- a/test/models/test_whisper.py +++ b/test/models/test_whisper.py @@ -6,7 +6,6 @@ from tinygrad.helpers import fetch from test.helpers import slow from tinygrad import Tensor, Device, dtypes -from tinygrad.device import is_dtype_supported import numpy as np # Audio generated with the command on MacOS: @@ -53,7 +52,6 @@ def wer_helper(result: str, reference: str)->float: return wer @unittest.skipIf(Device.DEFAULT in ["CPU"], "slow") -@unittest.skipUnless(is_dtype_supported(dtypes.float16), "need float16 support") # TODO: WEBGPU GPU dispatch dimensions limit @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU GPU dispatch dimensions limit") class TestWhisper(unittest.TestCase): diff --git a/test/null/test_compile_failures.py b/test/null/test_compile_failures.py index 3d4a5c3b34877..fcb5d75240e25 100644 --- a/test/null/test_compile_failures.py +++ b/test/null/test_compile_failures.py @@ -1,8 +1,7 @@ import unittest, io from contextlib import redirect_stdout from tinygrad import Tensor, dtypes, Device -from tinygrad.helpers import OSX, DEV -from tinygrad.device import is_dtype_supported +from tinygrad.helpers import OSX from tinygrad.engine.realize import compile_linear from tinygrad.codegen import to_program @@ -10,7 +9,6 @@ class TestCompileFailures(unittest.TestCase): def compile(self, out:Tensor): compile_linear(out.schedule_linear()) - @unittest.skipUnless(is_dtype_supported(dtypes.uchar), f"no uint8 on {Device.DEFAULT}") def test_interpolate_atari(self): self.compile(Tensor.empty(210, 160, dtype='uint8').interpolate((64, 64))) @@ -18,8 +16,7 @@ def test_add_max_uchar(self): self.compile((Tensor.empty(1024, dtype='uint8') + Tensor.empty(1024, dtype='uint8')).max()) class TestDisassembly(unittest.TestCase): - # TODO: fails on llvm. llvm.LLVMGetHostCPUName() returns "generic" - @unittest.skipUnless(Device.DEFAULT in ("CPU",) and DEV.renderer not in ("LLVM", "LVP") and OSX, "m series cpus support fp16 arithmetic") + @unittest.skipUnless(Device.DEFAULT == "CPU" and OSX, "m series cpus support fp16 arithmetic") def test_float16_alu(self): c = Tensor([1], dtype=dtypes.float16) + Tensor([1], dtype=dtypes.float16) s = c.schedule_linear().src[-1] diff --git a/test/null/test_const_folding.py b/test/null/test_const_folding.py index 40e6b843f1adf..a6557e0dad783 100644 --- a/test/null/test_const_folding.py +++ b/test/null/test_const_folding.py @@ -127,22 +127,21 @@ def t(cases: dict[DType, ConstType]): def test_vec_bitcast(self): with Context(SPEC=0): - r = full_rewrite(UOp.const(dtypes.int32.vec(3), (-1, -2**31, 75)).bitcast(dtypes.uint32.vec(3)).sink()).src[0] - self.assertEqual(r.op, Ops.STACK) - self.assertEqual(r.dtype, dtypes.uint32.vec(3)) - self.assertEqual(tuple(x.arg for x in r.src), (2**32-1, 2**31, 75)) + srcs = full_rewrite(UOp.const(dtypes.int32.vec(3), (-1, -2**31, 75)).bitcast(dtypes.uint32.vec(3)).sink()).src + self.assertTrue(all(r.op is Ops.CONST and r.dtype == dtypes.uint32 for r in srcs)) + self.assertEqual(tuple(x.arg for x in srcs), (2**32-1, 2**31, 75)) # folds advance indexing into basic indexing class TestIndexingConstFolding(unittest.TestCase): def test_scalar_index(self): - t = Tensor.arange(16).float().reshape(1,1,4,4).realize() + t = Tensor.arange(16).float().reshape(1,1,4,4).clone().realize() _check_ast_count(1, t[:,:,Tensor(1),:]) _check_ast_count(1, t[:,:,Tensor(1)+2,:]) _check_ast_count(1, t[:,:,Tensor(1),Tensor(0)]) def test_const_tensor_index(self): # TODO: these can be 0, implement const tensor folded indexing - t = Tensor.arange(16).float().reshape(1,1,4,4).realize() + t = Tensor.arange(16).float().reshape(1,1,4,4).clone().realize() _check_ast_count(1, t[:,:,Tensor.ones(2,1,dtype=dtypes.int),:]) _check_ast_count(1, t[:,:,Tensor.ones(1,2,dtype=dtypes.int)+2,:]) _check_ast_count(1, t[:,:,Tensor.ones(1,1,dtype=dtypes.int),Tensor.zeros(2,1,2,dtype=dtypes.int)]) diff --git a/test/null/test_device.py b/test/null/test_device.py index e798d653f18ca..ec62d5eaa61bc 100644 --- a/test/null/test_device.py +++ b/test/null/test_device.py @@ -3,7 +3,7 @@ from unittest.mock import patch from tinygrad import Tensor from tinygrad.device import Device, Compiler, enumerate_devices_str -from tinygrad.helpers import diskcache_get, diskcache_put, getenv, Context, Target, WIN, CI, OSX, DEV +from tinygrad.helpers import diskcache_get, diskcache_put, getenv, Context, Target, WIN, OSX, DEV from tinygrad.runtime.support.c import DLL class TestDevice(unittest.TestCase): @@ -28,8 +28,8 @@ def test_getitem_not_exist(self): def test_nonexistent_renderer(self): with self.assertRaisesRegex(RuntimeError, "has no renderer"): with Context(DEV="CPU:TYPO"): Device[Device.DEFAULT].renderer - with self.assertRaisesRegex(RuntimeError, "did you mean: 'CLANGJIT'"): - with Context(DEV="CPU:CLANG"): Device[Device.DEFAULT].renderer + with self.assertRaisesRegex(RuntimeError, "did you mean: 'CLANG'"): + with Context(DEV="CPU:CLANGJIT"): Device[Device.DEFAULT].renderer @unittest.skipIf(Device.DEFAULT != "AMD", "only run on AMD") def test_nonexistent_iface(self): @@ -66,20 +66,20 @@ def test_old_renderer_env_raises(self): self.assertNotEqual(result.returncode, 0) self.assertIn(b"deprecated", result.stderr) - @unittest.skipIf(WIN and CI, "skipping windows test") # TODO: subprocess causes memory violation? + @unittest.skipIf(WIN, "skipping windows test") # TODO: subprocess causes memory violation? def test_env_overwrite_default_compiler(self): if Device.DEFAULT == "CPU": - from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, ClangJITCompiler - try: _, _ = CPULLVMCompiler(), ClangJITCompiler() + from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, ClangCompiler + try: _, _ = CPULLVMCompiler(), ClangCompiler() except Exception as e: self.skipTest(f"skipping compiler test: not all compilers: {e}") - imports = "from tinygrad import Device; from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, ClangJITCompiler" + imports = "from tinygrad import Device; from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, ClangCompiler" subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, CPULLVMCompiler)"'], shell=True, check=True, env={**os.environ, "DEV": "CPU:LLVM"}) - subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, ClangJITCompiler)"'], + subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, ClangCompiler)"'], shell=True, check=True, env={**os.environ, "DEV": "CPU"}) - subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, ClangJITCompiler)"'], - shell=True, check=True, env={**os.environ, "DEV": "CPU:CLANGJIT"}) + subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, ClangCompiler)"'], + shell=True, check=True, env={**os.environ, "DEV": "CPU:CLANG"}) elif Device.DEFAULT == "AMD": from tinygrad.runtime.support.compiler_amd import HIPCompiler, AMDLLVMCompiler try: _, _ = HIPCompiler(Device[Device.DEFAULT].arch), AMDLLVMCompiler(Device[Device.DEFAULT].arch) @@ -94,17 +94,17 @@ def test_env_overwrite_default_compiler(self): shell=True, check=True, env={**os.environ, "DEV": "AMD:HIP"}) else: self.skipTest("only run on CPU/AMD") - @unittest.skipIf(WIN and CI, "skipping windows test") + @unittest.skipIf(WIN, "skipping windows test") def test_env_online(self): - from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, ClangJITCompiler - try: _, _ = CPULLVMCompiler(), ClangJITCompiler() + from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, ClangCompiler + try: _, _ = CPULLVMCompiler(), ClangCompiler() except Exception as e: self.skipTest(f"skipping compiler test: not all compilers: {e}") with Context(DEV="CPU:LLVM"): inst = Device["CPU"].compiler self.assertIsInstance(Device["CPU"].compiler, CPULLVMCompiler) with Context(DEV="CPU"): - self.assertIsInstance(Device["CPU"].compiler, ClangJITCompiler) + self.assertIsInstance(Device["CPU"].compiler, ClangCompiler) with Context(DEV="CPU:LLVM"): self.assertIsInstance(Device["CPU"].compiler, CPULLVMCompiler) assert inst is Device["CPU"].compiler # cached @@ -118,7 +118,7 @@ def test_compiler_autodetect_fallback(self): dev = Device["CPU"] dev.cached_renderer.clear() - with patch("tinygrad.renderer.cstyle.ClangJITRenderer.__init__", side_effect=RuntimeError("broken")): + with patch("tinygrad.renderer.cstyle.ClangRenderer.__init__", side_effect=RuntimeError("broken")): self.assertIsInstance(dev.renderer.compiler, CPULLVMCompiler) def test_dev_contextvar(self): @@ -181,6 +181,7 @@ def test_device_compile(self): a = Tensor([0.,1.], device=Device.DEFAULT).realize() (a + 1).realize() +@unittest.skip("this test is broken if you have tinymesa installed") @unittest.skipIf(OSX and 'libclang' in DLL._loaded_, "MTLCompiler can't be loaded after libclang on OSX") class TestRunAsModule(unittest.TestCase): def test_module_runs(self): diff --git a/test/null/test_dtype.py b/test/null/test_dtype.py index e5db4af55e955..7d54a83d302f3 100644 --- a/test/null/test_dtype.py +++ b/test/null/test_dtype.py @@ -10,6 +10,48 @@ def test_image_vec(self): assert dtypes.imagef((10,10)).base.vec(4) == dtypes.float32.vec(4) assert dtypes.imageh((10,10)).base.vec(4) == dtypes.float32.vec(4) +class TestPtrDType(unittest.TestCase): + def test_vec_double(self): + dt1 = dtypes.float.vec(4).ptr().vec(4) + dt2 = dtypes.float.vec(4).ptr().vec(4) + self.assertEqual(dt1, dt2) + self.assertEqual(str(dt1), str(dt2)) + + def test_scalar(self): + dt = dtypes.float.vec(4).ptr().scalar() + self.assertEqual(dt.base, dtypes.float.vec(4)) + + dt = dtypes.float.vec(4).ptr().vec(4).scalar() + self.assertEqual(dt.base, dtypes.float.vec(4)) + + dt = dtypes.float.vec(4).scalar() + self.assertEqual(dt, dtypes.float) + + def test_serialize(self): + dt = dtypes.float.vec(4).ptr().vec(4) + self.assertEqual(dt, eval(str(dt))) + + def test_vec_ptr_sz(self): + dt = dtypes.float.ptr(1024).vec(4) + self.assertEqual(dt, eval(str(dt))) + self.assertEqual(str(dt), "dtypes.float.ptr(1024).vec(4)") + + def test_vcount(self): + dt = dtypes.float.ptr().vec(4) + self.assertEqual(dt.vcount, 4) + self.assertEqual(dt.v, 4) + self.assertEqual(dt.count, 1) + + dt = dtypes.float.vec(4).ptr() + self.assertEqual(dt.vcount, 1) + self.assertEqual(dt.v, 1) + self.assertEqual(dt.count, 4) + + dt = dtypes.float.vec(4).ptr().vec(4) + self.assertEqual(dt.vcount, 4) + self.assertEqual(dt.v, 4) + self.assertEqual(dt.count, 4) + class TestEqStrDType(unittest.TestCase): def test_image_ne(self): if ImageDType is None: raise unittest.SkipTest("no ImageDType support") diff --git a/test/null/test_dtype_spec.py b/test/null/test_dtype_spec.py index 1116433319a1b..6a947fe8609fb 100644 --- a/test/null/test_dtype_spec.py +++ b/test/null/test_dtype_spec.py @@ -1,7 +1,6 @@ import unittest, math, struct, operator -from tinygrad.tensor import Tensor, dtypes -from tinygrad.dtype import DTYPES_DICT, truncate, float_to_fp16, float_to_bf16, _to_np_dtype, least_upper_dtype, least_upper_float -from tinygrad.device import is_dtype_supported +from tinygrad import Tensor, Device +from tinygrad.dtype import DTYPES_DICT, dtypes, truncate, float_to_fp16, float_to_bf16, _to_np_dtype, least_upper_dtype, least_upper_float from tinygrad.helpers import getenv from hypothesis import given, settings, strategies as strat @@ -12,8 +11,8 @@ settings.load_profile("my_profile") core_dtypes = list(DTYPES_DICT.values()) -dtype_ints = [dt for dt in core_dtypes if dtypes.is_int(dt) and is_dtype_supported(dt)] -dtype_floats = [dt for dt in core_dtypes if dtypes.is_float(dt) and is_dtype_supported(dt)] +dtype_ints = [dt for dt in core_dtypes if dtypes.is_int(dt) and dt in Device[Device.DEFAULT].renderer.supported_dtypes()] +dtype_floats = [dt for dt in core_dtypes if dtypes.is_float(dt) and dt in Device[Device.DEFAULT].renderer.supported_dtypes()] FP8E4M3_MAX = 448.0 FP8E5M2_MAX = 57344.0 @@ -179,6 +178,12 @@ def test_truncate_fp8e5m2(self, x): elif x < -FP8E5M2_MAX: np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), -FP8E5M2_MAX) else: np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), torch.tensor(x, dtype=torch.float8_e5m2).float().item()) + def test_finfo(self): + for dt in [dtypes.float16, dtypes.float32, dtypes.float64]: + info = np.finfo(_to_np_dtype(dt)) + self.assertEqual(info.bits, dt.bitsize) + self.assertEqual((info.nexp, info.nmant), dtypes.finfo(dt)) + class TestTypePromotion(unittest.TestCase): @given(strat.sampled_from(core_dtypes)) def test_self_promo_to_self(self, dtype): diff --git a/test/null/test_elf.py b/test/null/test_elf.py index f7d350bd34e36..0f9a3d48f0678 100644 --- a/test/null/test_elf.py +++ b/test/null/test_elf.py @@ -1,5 +1,5 @@ import unittest, subprocess, platform -from tinygrad.runtime.support.compiler_cpu import ClangJITCompiler +from tinygrad.runtime.support.compiler_cpu import ClangCompiler from tinygrad.runtime.support.elf import elf_loader class TestElfLoader(unittest.TestCase): @@ -23,7 +23,7 @@ def test_clang_jit_compiler_external_raise(self): } ''' with self.assertRaisesRegex(RuntimeError, 'evil_external_function'): - ClangJITCompiler([{'AMD64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine(), m), "native"]).compile(src) + ClangCompiler([{'AMD64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine(), m), "native"]).compile(src) def test_link(self): src = ''' float powf(float, float); // from libm diff --git a/test/null/test_gc.py b/test/null/test_gc.py index 21682e395faef..9799967310c9d 100644 --- a/test/null/test_gc.py +++ b/test/null/test_gc.py @@ -29,8 +29,8 @@ class TestGC(unittest.TestCase): def test_gc(self): Tensor.manual_seed(0) base = tensors_allocated() - a = Tensor.rand(4, 4, requires_grad=True) - b = Tensor.zeros(4, 4, requires_grad=True) + a = Tensor.rand(4, 4) + b = Tensor.zeros(4, 4) (a*b).mean().backward() assert (tensors_allocated()-base > 0) del a,b @@ -40,14 +40,14 @@ def test_gc(self): def test_gc_complex(self): Tensor.manual_seed(0) base = tensors_allocated() - a = Tensor(np.zeros((4, 4), dtype=np.float32), requires_grad=True) - b = Tensor.rand(4, 4, requires_grad=True) + a = Tensor(np.zeros((4, 4), dtype=np.float32)) + b = Tensor.rand(4, 4) assert (tensors_allocated()-base == 4) (a*b).mean().backward() assert (tensors_allocated()-base == 6) del b assert (tensors_allocated()-base == 4) - b = Tensor(np.zeros((4, 4), dtype=np.float32), requires_grad=True) + b = Tensor(np.zeros((4, 4), dtype=np.float32)) print(tensors_allocated()) (a*b).mean().backward() print(tensors_allocated()) diff --git a/test/null/test_gpudims.py b/test/null/test_gpudims.py index 0a45dd1cb90c5..0a6fb27037856 100644 --- a/test/null/test_gpudims.py +++ b/test/null/test_gpudims.py @@ -96,7 +96,7 @@ def test_grouped_direct_dims_are_special(self): def test_global_prod_max(self): g, l = UOp.range(256, 0, AxisType.GLOBAL), UOp.range(256, 1, AxisType.LOCAL) - sink = UOp(Ops.PARAM, dtypes.float.ptr(), (), 0).index(g + l).store(UOp.const(dtypes.float, 1.0)).end(g, l).sink(arg=KernelInfo()) + sink = UOp.param(0, dtypes.float.ptr()).index(g + l).store(UOp.const(dtypes.float, 1.0)).end(g, l).sink(arg=KernelInfo()) class R(Renderer): global_max, local_max, global_prod_max = (256, 256, 256), (128, 128, 128), (128, 128, 128) specials = [u for u in add_gpudims(R(Target()), sink).toposort() if u.op is Ops.SPECIAL] self.assertGreater(len([s for s in specials if "lidx" in s.arg]), 1) diff --git a/test/null/test_gradient.py b/test/null/test_gradient.py index 1727608f08b39..36ca157811720 100644 --- a/test/null/test_gradient.py +++ b/test/null/test_gradient.py @@ -64,18 +64,18 @@ def test_where(self): self._test_two_input_function(lambda x,y: (x 0 and kernels_used == GlobalCounters.kernel_count or (kernels_used <= GlobalCounters.kernel_count and getattr(Device[Device.DEFAULT], "graph", None)), f"only {kernels_used} out of {GlobalCounters.kernel_count} were jitted" # noqa: E501 +supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes() + class TestRealWorld(unittest.TestCase): def setUp(self): gc.collect() @@ -53,7 +54,7 @@ def tearDown(self): dtypes.default_float = self.old_float @slow - @unittest.skipUnless(is_dtype_supported(dtypes.float16), "need dtypes.float16") + @unittest.skipUnless(dtypes.float16 in supported_dtypes, "need dtypes.float16") def test_stable_diffusion(self): params = unet_params params["model_ch"] = 8 @@ -78,7 +79,7 @@ def test(t, t2): exp_mem = 0.00037 if Device.DEFAULT == "CL" else 0.0002 helper_test("test_unet_resblock", lambda: (Tensor.empty(4, 16, 8, 8), Tensor.empty(1, 24)), test, exp_mem, 37) - @unittest.skipUnless(is_dtype_supported(dtypes.float16), "need dtypes.float16") + @unittest.skipUnless(dtypes.float16 in supported_dtypes, "need dtypes.float16") def test_llama(self): dtypes.default_float = dtypes.float16 @@ -90,7 +91,7 @@ def test(t): return model(t, 0).realize() # TODO: test first token vs rest properly helper_test("test_llama", lambda: (Tensor([[1,2,3,4]]),), test, 0.23, 118, all_jitted=True) - @unittest.skipUnless(is_dtype_supported(dtypes.float16), "need dtypes.float16") + @unittest.skipUnless(dtypes.float16 in supported_dtypes, "need dtypes.float16") def test_gpt2(self): dtypes.default_float = dtypes.float16 @@ -147,7 +148,7 @@ def train(X): helper_test("train_cifar", lambda: (Tensor.randn(BS, 3, 32, 32),), train, 0.12, 126) - @unittest.skipUnless(is_dtype_supported(dtypes.float16), "need dtypes.float16") + @unittest.skipUnless(dtypes.float16 in supported_dtypes, "need dtypes.float16") def test_train_cifar_hyp(self): dtypes.default_float = dtypes.float16 with Tensor.train(): diff --git a/test/null/test_schedule.py b/test/null/test_schedule.py index e4a919c948843..5981954e331e0 100644 --- a/test/null/test_schedule.py +++ b/test/null/test_schedule.py @@ -118,7 +118,7 @@ def test_non_contiguous_buffer_view(self): def test_size_change_buffer_view(self): a = Tensor.empty(4) b = a.reshape((1, 1, 4)).shrink(((0, 1), (0, 1), (0, 3))).contiguous() - check_schedule(b, 0) # contiguous shrink of a realized buffer is a zero-copy BUFFER_VIEW + check_schedule(b, 0) # contiguous shrink of a realized buffer is a zero-copy SLICE def test_double_contiguous_realizes_once(self): a = Tensor.empty(4, 1) @@ -414,7 +414,7 @@ def test_dedup_outputs(self): check_schedule([a+b, a+b], 1) def test_const_realize(self): - t = Tensor.ones(2) + t = Tensor.ones(2, buffer=False) check_schedule(t[0], 0) check_schedule(t[1], 0) @@ -429,7 +429,7 @@ def test_fold_batchnorm(self): img = Tensor.empty(1,32,4,4) bn = nn.BatchNorm2d(32, track_running_stats=False) out = bn(img) - check_schedule(out, 3) + check_schedule(out, 3, nn.state.get_parameters(bn)) def test_fold_conv_batchnorm_notrain(self): with Tensor.train(False): @@ -437,7 +437,7 @@ def test_fold_conv_batchnorm_notrain(self): c1 = nn.Conv2d(3,32,3) bn = nn.BatchNorm2d(32, track_running_stats=True) out = bn(c1(img)).relu() - check_schedule(out, 1, [c1.weight, c1.bias]) + check_schedule(out, 1, [c1.weight, c1.bias, *nn.state.get_parameters(bn)]) def test_fold_conv_batchnorm_notrain_no_running_stats(self): with Tensor.train(False): @@ -445,7 +445,7 @@ def test_fold_conv_batchnorm_notrain_no_running_stats(self): c1 = nn.Conv2d(3,32,3) bn = nn.BatchNorm2d(32, track_running_stats=False) out = bn(c1(img)).relu() - check_schedule(out, 4, [c1.weight, c1.bias]) + check_schedule(out, 4, [c1.weight, c1.bias, *nn.state.get_parameters(bn)]) def test_fold_conv_batchnorm(self): with Tensor.train(): @@ -453,17 +453,17 @@ def test_fold_conv_batchnorm(self): c1 = nn.Conv2d(3,32,3) bn = nn.BatchNorm2d(32, track_running_stats=False) out = bn(c1(img)).relu() - check_schedule(out, 4, [c1.weight, c1.bias]) + check_schedule(out, 4, [c1.weight, c1.bias, *nn.state.get_parameters(bn)]) def test_fold_conv_batchnorm_optim(self, adam=False): - # 2 is too low? - optim, cnt = (nn.optim.Adam, 16) if adam else (nn.optim.SGD, 2) + optim, cnt = (nn.optim.Adam, 29) if adam else (nn.optim.SGD, 15) with Tensor.train(): img = Tensor.ones(1,3,4,4) c1 = nn.Conv2d(3,32,3) bn = nn.BatchNorm2d(32, track_running_stats=False) _realize_weights([c1, bn]) opt = optim(nn.state.get_parameters([c1, bn])) + Tensor.realize(img, *nn.state.get_parameters(opt)) img_bn = bn(c1(img)).elu().sum() opt.zero_grad() img_bn.backward() @@ -474,18 +474,17 @@ def test_fold_batchnorm_backward(self): with Tensor.train(): x = Tensor.empty((2, 16, 8, 8)).contiguous() bn = nn.BatchNorm2d(16) - bn.weight.requires_grad = bn.bias.requires_grad = x.requires_grad = True fw = bn(x).contiguous_backward().relu().contiguous() fw.sum().backward() # TODO: this is too many - check_schedule([x.grad, bn.weight.grad, bn.bias.grad, fw], 9) + check_schedule([x.grad, bn.weight.grad, bn.bias.grad, fw], 10, nn.state.get_parameters(bn)) def test_fold_conv_relu(self): c1 = nn.Conv2d(3,16,3) # run img = Tensor.ones(2,3,64,64) out = c1(img).relu() - check_schedule(out, 1, [c1.weight, c1.bias]) + check_schedule(out, 1, [c1.weight, c1.bias, img]) def test_fold_conv_relu_alt(self): img = Tensor.ones(1,4,8,8) @@ -767,12 +766,6 @@ def test_reduce_shrink(self): out = x + y check_schedule(out, 1) - def test_const_no_recompute(self): - x = Tensor(2) + Tensor(2) - y = Tensor(2) + Tensor(2) - out = x.contiguous() + y.contiguous() - check_schedule(out, 2, filter_sink=False) - def test_reduce_shrink_child(self): a = Tensor.empty(100, 100) b = Tensor.empty(10,) @@ -808,7 +801,7 @@ def test_softmax_upcast(self): def test_softmax_backward(self): Tensor.manual_seed(0) - x = Tensor.randn(4, 12, 64, 64, requires_grad=True).realize() + x = Tensor.randn(4, 12, 64, 64).realize() x.softmax().sum().backward() run_linear(*check_schedule(x.grad, 4)) @@ -828,6 +821,7 @@ def test_adam_step_fusion(self): layer = nn.Linear(32, 32*4) _realize_weights(layer) opt = nn.optim.Adam(nn.state.get_parameters(layer), lr=1e-4) + Tensor.realize(*nn.state.get_parameters(opt)) layer(x).relu().sum().backward() check_schedule(opt.schedule_step(), 13) @@ -837,6 +831,7 @@ def test_adam_conv_fuse(self): c1 = nn.Conv2d(3,32,3) _realize_weights(c1) opt = nn.optim.Adam(nn.state.get_parameters(c1), lr=1e-4) + Tensor.realize(*nn.state.get_parameters(opt)) opt.zero_grad() c1(img).relu().sum().backward() check_schedule(opt.schedule_step(), 13) @@ -848,6 +843,7 @@ def test_adam_2convs_fuse(self): c2 = nn.Conv2d(16,32,2,bias=False) _realize_weights([c1, c2]) opt = nn.optim.Adam(nn.state.get_parameters([c1, c2]), lr=1e-4) + Tensor.realize(*nn.state.get_parameters(opt)) opt.zero_grad() c2(c1(img).relu()).relu().sum().backward() check_schedule(opt.schedule_step(), 15) @@ -880,6 +876,7 @@ def test_fold_2convs_sgd_nesterov_momentum_wd(self): c2 = nn.Conv2d(16,32,2,bias=False) _realize_weights([c1, c2]) opt = nn.optim.SGD(nn.state.get_parameters([c1, c2]), nesterov=True, momentum=0.9, weight_decay=0.1) + Tensor.realize(*nn.state.get_parameters(opt)) opt.zero_grad() c2(c1(img).relu()).relu().sum().backward() check_schedule(opt.schedule_step(), 11) @@ -994,7 +991,7 @@ def test_fuse_arange_avg_pool2d_ceil_mode(self): def test_fuse_arange_pad_circular_mode_bw(self): x = Tensor.empty(1,1,5,5,5) out = x.pad((1,2,3,5,1,2), mode="circular") - g = out.sum().gradient(x)[0] + g = out.sum().gradient(x)[0].clone() linear, _ = check_schedule(g, 1) self.assertEqual(len([x for x in linear.src[0].src[0].backward_slice_with_self if x.op is Ops.REDUCE]), 0) @@ -1009,7 +1006,7 @@ def test_resnet_block(self): out = bn1(conv1(x)).relu() out = bn2(conv2(out)) out = (out + x).relu() - run_linear(*check_schedule(out, 2, [conv1.weight, conv2.weight])) + run_linear(*check_schedule(out, 2, [conv1.weight, conv2.weight, *nn.state.get_parameters(bn1), *nn.state.get_parameters(bn2)])) class TestSwizzle(unittest.TestCase): def test_softmax_one_kernel(self): @@ -1213,118 +1210,118 @@ def test_recursive_reshape(self): self.assertEqual(len(linear.src), 1) self.assertLess(time.perf_counter()-st, 2.0) -# NOTE: the NULL backend supports BUFFER_VIEW +# NOTE: the NULL backend supports SLICE class TestBufferView(unittest.TestCase): def test_shrink_contiguous_is_buffer_view(self): - # simple 1D shrink of a realized buffer should be BUFFER_VIEW, not a copy kernel - a = Tensor.arange(100).contiguous().realize() + # simple 1D shrink of a realized buffer should be SLICE, not a copy kernel + a = Tensor.arange(100).clone().realize() b = a.shrink(((10, 50),)).contiguous() run_linear(*check_schedule(b, 0)) def test_shrink_2d_contiguous_is_buffer_view(self): - a = Tensor.arange(100).reshape(10,10).contiguous().realize() + a = Tensor.arange(100).reshape(10,10).clone().realize() b = a.shrink(((1, 5),None)).contiguous() run_linear(*check_schedule(b, 0)) def test_chained_shrink_is_buffer_view(self): - a = Tensor.arange(1000).contiguous().realize() + a = Tensor.arange(1000).clone().realize() b = a.shrink(((200, 800),)).shrink(((0, 300),)).reshape((30, 10)).shrink(((20, 25), (0, 10))).contiguous() run_linear(*check_schedule(b, 0)) def test_shrink_non_shard_axis_is_buffer_view_multi(self): - # indexing a non-shard axis of a realized sharded tensor should be BUFFER_VIEW on each device, not copy kernels + # indexing a non-shard axis of a realized sharded tensor should be SLICE on each device, not copy kernels # this is the flat_llama pattern: weight[layer_idx] where weight is (n_layers, out, dim) sharded on axis=1 devices = ("NULL:1", "NULL:2") - a = Tensor.arange(8*4*10).reshape(8, 4, 10).contiguous().shard(devices, axis=1).realize() + a = Tensor.arange(8*4*10).reshape(8, 4, 10).clone().shard(devices, axis=1).realize() run_linear(*check_schedule(a[3].contiguous(), 0)) def test_shrink_2d_non_shard_axis_multi(self): devices = ("NULL:1", "NULL:2") - a = Tensor.arange(6*4).reshape(6, 4).contiguous().shard(devices, axis=1).realize() + a = Tensor.arange(6*4).reshape(6, 4).clone().shard(devices, axis=1).realize() run_linear(*check_schedule(a.shrink(((1, 4), None)).contiguous(), 0)) def test_shrink_shard_axis_0_multi(self): # shrinking a middle dim is not contiguous per shard, so this needs copy kernels devices = ("NULL:1", "NULL:2") - a = Tensor.arange(4*6*2).reshape(4, 6, 2).contiguous().shard(devices, axis=0).realize() + a = Tensor.arange(4*6*2).reshape(4, 6, 2).clone().shard(devices, axis=0).realize() run_linear(*check_schedule(a.shrink((None, (2, 5), None)).contiguous(), 2)) def test_reshape_then_shrink_multi(self): devices = ("NULL:1", "NULL:2") - a = Tensor.arange(8*6).reshape(8, 6).contiguous().shard(devices, axis=1).realize() + a = Tensor.arange(8*6).reshape(8, 6).clone().shard(devices, axis=1).realize() run_linear(*check_schedule(a.reshape(4, 2, 6)[1].contiguous(), 0)) def test_permute_then_shrink_multi(self): # permute makes per-shard view non-contiguous, needs copy kernels devices = ("NULL:1", "NULL:2") - a = Tensor.arange(4*6*2).reshape(4, 6, 2).contiguous().shard(devices, axis=1).realize() + a = Tensor.arange(4*6*2).reshape(4, 6, 2).clone().shard(devices, axis=1).realize() run_linear(*check_schedule(a.permute(1, 0, 2).shrink(((0, 6), (1, 3), None)).contiguous(), 2)) def test_multi_buffer_view_4_devices(self): devices = tuple(f"NULL:{i}" for i in range(4)) - a = Tensor.arange(8*12).reshape(8, 12).contiguous().shard(devices, axis=1).realize() + a = Tensor.arange(8*12).reshape(8, 12).clone().shard(devices, axis=1).realize() run_linear(*check_schedule(a[5].contiguous(), 0)) def test_chained_shrink_multi(self): devices = ("NULL:1", "NULL:2") - a = Tensor.arange(10*8).reshape(10, 8).contiguous().shard(devices, axis=1).realize() + a = Tensor.arange(10*8).reshape(10, 8).clone().shard(devices, axis=1).realize() run_linear(*check_schedule(a.shrink(((2, 8), None)).shrink(((1, 4), None)).contiguous(), 0)) # negative tests: these should NOT become BUFFER_VIEW (non-contiguous per shard) def test_expand_multi_not_buffer_view(self): devices = ("NULL:1", "NULL:2") - a = Tensor.arange(4*2).reshape(4, 1, 2).contiguous().shard(devices, axis=2).realize() + a = Tensor.arange(4*2).reshape(4, 1, 2).clone().shard(devices, axis=2).realize() run_linear(*check_schedule(a.expand(4, 3, 2).contiguous(), 2)) def test_pad_multi_not_buffer_view(self): devices = ("NULL:1", "NULL:2") - a = Tensor.arange(4*2).reshape(4, 2).contiguous().shard(devices, axis=1).realize() + a = Tensor.arange(4*2).reshape(4, 2).clone().shard(devices, axis=1).realize() run_linear(*check_schedule(a.pad(((1, 1), (0, 0))).contiguous(), 2)) def test_flip_multi_not_buffer_view(self): devices = ("NULL:1", "NULL:2") - a = Tensor.arange(4*2).reshape(4, 2).contiguous().shard(devices, axis=1).realize() + a = Tensor.arange(4*2).reshape(4, 2).clone().shard(devices, axis=1).realize() run_linear(*check_schedule(a.flip(0).contiguous(), 2)) def test_replicated_reshape_is_buffer_view(self): devices = ("NULL:1", "NULL:2") - a = Tensor.arange(24).contiguous().to(devices).realize() + a = Tensor.arange(24).clone().to(devices).realize() run_linear(*check_schedule(a.reshape(4, 6).contiguous(), 0)) def test_replicated_shrink_is_buffer_view(self): # DP pattern: replicated weight[layer_idx] devices = ("NULL:1", "NULL:2") - a = Tensor.arange(8*10).reshape(8, 10).contiguous().to(devices).realize() + a = Tensor.arange(8*10).reshape(8, 10).clone().to(devices).realize() run_linear(*check_schedule(a[3].contiguous(), 0)) def test_replicated_chained_mops_is_buffer_view(self): devices = ("NULL:1", "NULL:2") - a = Tensor.arange(100).contiguous().to(devices).realize() + a = Tensor.arange(100).clone().to(devices).realize() run_linear(*check_schedule(a.reshape(10, 10).shrink(((2, 7), None)).contiguous(), 0)) def test_replicated_shard_none_is_buffer_view(self): devices = ("NULL:1", "NULL:2") - a = Tensor.arange(24).contiguous().shard(devices, axis=None).realize() + a = Tensor.arange(24).clone().shard(devices, axis=None).realize() run_linear(*check_schedule(a.reshape(4, 6).contiguous(), 0)) def test_replicated_4_devices_is_buffer_view(self): devices = tuple(f"NULL:{i}" for i in range(4)) - a = Tensor.arange(8*10).reshape(8, 10).contiguous().to(devices).realize() + a = Tensor.arange(8*10).reshape(8, 10).clone().to(devices).realize() run_linear(*check_schedule(a[3].contiguous(), 0)) def test_replicated_expand_not_buffer_view(self): devices = ("NULL:1", "NULL:2") - a = Tensor.arange(12).reshape(4, 1, 3).contiguous().to(devices).realize() + a = Tensor.arange(12).reshape(4, 1, 3).clone().to(devices).realize() run_linear(*check_schedule(a.expand(4, 3, 3).contiguous(), 2)) def test_replicated_permute_not_buffer_view(self): devices = ("NULL:1", "NULL:2") - a = Tensor.arange(24).reshape(4, 6).contiguous().to(devices).realize() + a = Tensor.arange(24).reshape(4, 6).clone().to(devices).realize() run_linear(*check_schedule(a.permute(1, 0).contiguous(), 2)) def test_replicated_flip_not_buffer_view(self): devices = ("NULL:1", "NULL:2") - a = Tensor.arange(24).reshape(4, 6).contiguous().to(devices).realize() + a = Tensor.arange(24).reshape(4, 6).clone().to(devices).realize() run_linear(*check_schedule(a.flip(0).contiguous(), 2)) class TestInvalidTensor(unittest.TestCase): diff --git a/test/null/test_simplify_valid_idx.py b/test/null/test_simplify_valid_idx.py index 1fd781651af08..a537a51c72d96 100644 --- a/test/null/test_simplify_valid_idx.py +++ b/test/null/test_simplify_valid_idx.py @@ -15,13 +15,13 @@ def simplify_image_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move def get_gated_load_uop(valid:UOp, idx:UOp): return UOp(Ops.LOAD, dtypes.float, ( - UOp(Ops.PARAM, dtypes.float.ptr(), arg=0).index(idx.valid(valid), ptr=True), + UOp.param(0, dtypes.float.ptr()).index(idx.valid(valid), ptr=True), UOp.const(dtypes.float, 0.0) )) def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UOp]): return UOp(Ops.LOAD, dtypes.float.vec(4), ( - UOp(Ops.PARAM, dtypes.imagef(image_shape), arg=0).index(UOp(Ops.STACK, dtypes.weakint.vec(2), idx).valid(valid), ptr=True), + UOp.param(0, dtypes.imagef(image_shape)).index(idx[1].valid(valid), idx[0].valid(valid), ptr=True), UOp(Ops.STACK, dtypes.float.vec(4), src=(UOp.const(dtypes.float, 0.0),) * 4) )) @@ -222,17 +222,16 @@ def test_valid_with_non_const_rhs(self): class TestImageSimplification(unittest.TestCase): def check(self, load, svalid, sidx0, sidx1): load = simplify_image_idx(load.sink()).src[0] - off = load.src[0].src[1] - idx = off.get_idx() - self.assertEqual(idx.op, Ops.STACK) - self.assertEqual(len(idx.src), 2) - idx0, idx1 = idx.src[0], idx.src[1] + off = load.src[0] + self.assertEqual(len(off.src), 3) + idx0, idx1 = off.src[2].get_idx(), off.src[1].get_idx() check_uop_against_string(self, idx0, sidx0) check_uop_against_string(self, idx1, sidx1) + self.assertEqual(off.src[1].get_valid(), off.src[2].get_valid()) if svalid is not None: - check_uop_against_string(self, off.get_valid(), svalid) + check_uop_against_string(self, off.src[1].get_valid(), svalid) else: - self.assertEqual(off.get_valid(), UOp.const(dtypes.bool, True), "svalid is None but valid is not True") + self.assertEqual(off.src[1].get_valid(), UOp.const(dtypes.bool, True), "svalid is None but valid is not True") def test_idx_gt_c(self): # (idx1 < c+1).ne(True) ? (..., idx1-1+c) : 0 can drop the valid @@ -508,21 +507,13 @@ def test_simplify10(self): self.check(load, "(((lidx1<1)!=True)&(((lidx0+r0)<3)!=True)&((lidx0+r0)<11))", "(lidx2+gidx0*4+lidx1*256+(lidx0*1024+r0*1024)+-3264)", "0") -class TestUnfoldableImage(unittest.TestCase): - def test_unfoldable_becomes_buffer(self): - with Context(SPEC=0): - lidx = Special("lidx", 2) - load = UOp(Ops.LOAD, dtypes.float, (UOp(Ops.PARAM, dtypes.imagef((10, 10, 4)), arg=0).index(lidx, ptr=True), UOp.const(dtypes.float, 0))) - res = full_rewrite(load.sink()).src[0] - self.assertEqual(res.src[0].src[0].dtype, dtypes.float.ptr(400)) - class TestDropTrueGate(unittest.TestCase): def test_drop_true_gate_on_index(self): # test that INDEX with a constant True valid gets simplified to drop the valid from tinygrad.codegen.late.devectorizer import load_store_indexing from tinygrad.uop.ops import graph_rewrite from tinygrad.uop.symbolic import sym - buf = UOp(Ops.PARAM, dtypes.int.ptr(), arg=0) + buf = UOp.param(0, dtypes.int.ptr()) idx = UOp.const(dtypes.weakint, 0) true_gate = UOp.const(dtypes.bool, True) index_with_gate = UOp(Ops.INDEX, dtypes.int.ptr(), (buf, idx.valid(true_gate))) @@ -566,7 +557,7 @@ def test_range_no_shrink_when_unguarded_elsewhere(self): # one load guards r < 4, but another load uses r without a gate -> no shrink r = Range(0, 204) load1 = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r) - load2 = UOp(Ops.LOAD, dtypes.float, (UOp(Ops.PARAM, dtypes.float.ptr(), arg=1).index(r, ptr=True),)) + load2 = UOp(Ops.LOAD, dtypes.float, (UOp.param(1, dtypes.float.ptr()).index(r, ptr=True),)) ranges = self.get_ranges(UOp.sink(load1, load2)) self.assertEqual(len(ranges), 1) self.assertEqual(ranges[0].src[0].arg, 204) @@ -592,7 +583,7 @@ def test_range_shrink_store_where_invalid(self): from tinygrad.dtype import Invalid r = Range(0, 204) x = (r < 4).where(UOp.const(dtypes.float, 1), Invalid) - ranges = self.get_ranges(UOp(Ops.PARAM, dtypes.float.ptr(), arg=0).index(r).store((r < 4).where(x, 0)).sink()) + ranges = self.get_ranges(UOp.param(0, dtypes.float.ptr()).index(r).store((r < 4).where(x, 0)).sink()) self.assertEqual(len(ranges), 1) self.assertEqual(ranges[0].src[0].arg, 4) @@ -601,7 +592,7 @@ def test_range_shrink_store_where_invalid_flipped(self): from tinygrad.dtype import Invalid r = Range(0, 204) x = (r < 4).where(UOp.const(dtypes.float, 1), Invalid) - ranges = self.get_ranges(UOp(Ops.PARAM, dtypes.float.ptr(), arg=0).index(r).store((r < 4).where(0, x)).sink()) + ranges = self.get_ranges(UOp.param(0, dtypes.float.ptr()).index(r).store((r < 4).where(0, x)).sink()) self.assertEqual(len(ranges), 1) self.assertEqual(ranges[0].src[0].arg, 4) diff --git a/test/null/test_tensor.py b/test/null/test_tensor.py index 209d122e4235b..cf6a2001d0ab5 100644 --- a/test/null/test_tensor.py +++ b/test/null/test_tensor.py @@ -2,7 +2,6 @@ import numpy as np import unittest from tinygrad import Tensor, Device, dtypes -from tinygrad.device import is_dtype_supported from tinygrad.uop.ops import Ops, UOp from tinygrad.renderer.ptx import PTXRenderer from tinygrad.renderer.nir import NIRRenderer @@ -24,9 +23,9 @@ def f(): class TestInferenceMode(unittest.TestCase): def test_inference(self): - x = Tensor(x_init, requires_grad=True) - m = Tensor(m_init, requires_grad=True) - W = Tensor(W_init, requires_grad=True) + x = Tensor(x_init) + m = Tensor(m_init) + W = Tensor(W_init) tmp = x.mul(m) mm = tmp.matmul(W) out = mm.relu() @@ -37,12 +36,11 @@ def test_inference(self): assert tmp.grad is None assert mm.grad is None assert W.grad is None - assert W.requires_grad def test_no_grad_mode_context_manager(self): - x = Tensor(x_init, requires_grad=True) - m = Tensor(m_init, requires_grad=True) - W = Tensor(W_init, requires_grad=True) + x = Tensor(x_init) + m = Tensor(m_init) + W = Tensor(W_init) def f(x, m, W): tmp = x.mul(m) mm = tmp.matmul(W) @@ -80,18 +78,18 @@ def _assert(self, dtype: DType, a: Tensor): if not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)): assert idx.op is Ops.INDEX idx_val = idx.src[1] - assert idx_val.dtype is dtype + self.assertIs(idx_val.dtype, dtype) # use expand to generate kernel that uses large idx def do_op_then_assert(self, dtype: DType, dim1, dim2, dim3): self._assert(dtype, Tensor.empty(dim1, dim2, 1).expand(-1, -1, dim3).contiguous()) - @unittest.skipUnless(is_dtype_supported(dtypes.long), "int64 is supported") + @unittest.skipUnless(dtypes.long in Device[Device.DEFAULT].renderer.supported_dtypes(), "int64 is supported") def test_overflow(self): # 2**11, 2**11, 2**11 -> 2**33 will overflow when indexed self.do_op_then_assert(dtypes.long, 2048, 2048, 2048) - @unittest.skipUnless(is_dtype_supported(dtypes.long), "int64 is supported") + @unittest.skipUnless(dtypes.long in Device[Device.DEFAULT].renderer.supported_dtypes(), "int64 is supported") def test_overflow_sym(self): self.do_op_then_assert(dtypes.long, 2048, 2048, UOp.variable("dim3", 1, 2048).bind(32)) @@ -99,7 +97,7 @@ def test_regular(self): self.do_op_then_assert(dtypes.int, 64, 64, 64) def test_regular_sym(self): - self.do_op_then_assert(dtypes.int, 2048, 2048, UOp.variable("dim3", 1, 64).bind(32)) + self.do_op_then_assert(dtypes.int, 256, 256, UOp.variable("dim3", 1, 64).bind(32)) @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)), "PTX and NIR always converts Ops.INDEX to int64") def test_symfold(self): @@ -108,12 +106,12 @@ def test_symfold(self): uops = self._schedule_render(a) assert all(uop.dtype is not dtypes.long for uop in uops) - @unittest.skipIf(is_dtype_supported(dtypes.long), "int64 is supported") + @unittest.skipIf(dtypes.long in Device[Device.DEFAULT].renderer.supported_dtypes(), "int64 is supported") def test_int64_unsupported_overflow_sym(self): with self.assertRaises((KeyError, RuntimeError)): self.do_op_then_assert(dtypes.long, 2048, 2048, UOp.variable("dim3", 1, 2048).bind(32)) - @unittest.skipIf(is_dtype_supported(dtypes.long), "int64 is supported") + @unittest.skipIf(dtypes.long in Device[Device.DEFAULT].renderer.supported_dtypes(), "int64 is supported") @unittest.expectedFailure # bug in gpu dims limiting def test_int64_unsupported_overflow(self): with self.assertRaises((KeyError, RuntimeError)): @@ -146,12 +144,6 @@ def test_zeros_bufs_unique(self): Tensor.realize(a,b) self.assertIsNot(a.uop.buffer, b.uop.buffer) - def test_eye_bufs_unique(self): - a = Tensor.eye(10).contiguous() - b = Tensor.eye(10).contiguous() - Tensor.realize(a,b) - self.assertIsNot(a.uop.buffer, b.uop.buffer) - def test_times_2_not_unique(self): a = Tensor.zeros(10, 10).contiguous() b = a * 2 diff --git a/test/null/test_tensor_metadata.py b/test/null/test_tensor_metadata.py index 0ffa9ee8b40d8..dfd0d4ca59ca4 100644 --- a/test/null/test_tensor_metadata.py +++ b/test/null/test_tensor_metadata.py @@ -29,8 +29,8 @@ def test_exclude_const_metadata(self): self.assertEqual([m.name for m in sched.src[1].arg.metadata], ["contiguous"]) def test_matmul(self): - x = Tensor.rand(3, requires_grad=True) - W = Tensor.rand(3, 3, requires_grad=True) + x = Tensor.rand(3) + W = Tensor.rand(3, 3) out = x.matmul(W) self.assertEqual(out.uop.metadata[0].name, "matmul") si = out.schedule_linear().src[-1] @@ -38,7 +38,7 @@ def test_matmul(self): self.assertEqual(si.arg.metadata[0].name, "matmul") def test_relu(self): - x = Tensor.rand(3, requires_grad=True) + x = Tensor.rand(3) out = x.relu() self.assertEqual(out.uop.metadata[0].name, "relu") si = out.schedule_linear().src[-1] @@ -54,8 +54,8 @@ def test_assign(self): self.assertEqual(si.arg.metadata[0].name, "assign") def test_complex(self): - x = Tensor.rand(3, requires_grad=True) - y = Tensor.rand(3, requires_grad=True) + x = Tensor.rand(3) + y = Tensor.rand(3) out = x.relu() * y.sigmoid() self.assertEqual(out.uop.metadata[0].name, "__mul__") self.assertEqual(out.uop.src[0].metadata[0].name, "relu") @@ -66,8 +66,8 @@ def test_complex(self): @unittest.skip("flaky") def test_complex_backward(self): - x = Tensor.rand(3, requires_grad=True).realize() - y = Tensor.rand(3, requires_grad=True).realize() + x = Tensor.rand(3).realize() + y = Tensor.rand(3).realize() out = (x.relu() * y.sigmoid()).sum() self.assertEqual(out.uop.metadata[0].name, "sum") out.backward() @@ -85,8 +85,8 @@ def test_complex_backward(self): def test_tracemeta_0(self): with Context(TRACEMETA=0): - x = Tensor.rand(3, requires_grad=True) - y = Tensor.rand(3, requires_grad=True) + x = Tensor.rand(3) + y = Tensor.rand(3) out = (x.relu() * y.sigmoid()).sum() self.assertIsNone(out.uop.metadata) self.assertIsNone(out.uop.src[0].metadata) @@ -118,7 +118,7 @@ class TestTraceMetaShutdown(unittest.TestCase): def test_tracemeta_del_no_shutdown_error(self): import subprocess, os result = subprocess.run(['python3', '-c', 'from tinygrad import Tensor\n' - 'x=Tensor.eye(3,requires_grad=True); (x@x).sum().backward()'], + 'x=Tensor.eye(3); (x@x).sum().backward()'], env={**os.environ, "TRACEMETA": "2"}, capture_output=True) self.assertEqual(result.returncode, 0) self.assertNotIn(b"Exception", result.stderr) diff --git a/test/null/test_tensor_uop_mixin.py b/test/null/test_tensor_uop_mixin.py index 906f9608435da..5a4e881be5f24 100644 --- a/test/null/test_tensor_uop_mixin.py +++ b/test/null/test_tensor_uop_mixin.py @@ -2,7 +2,9 @@ from tinygrad import Tensor, dtypes from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, graph_rewrite -_strip_unique_pm = PatternMatcher([(UPat(Ops.CONST, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE, name="d")), name="b"), lambda b,d: b.replace(src=(d,))),]) +_strip_unique_pm = PatternMatcher([ + (UPat((Ops.UNIQUE, Ops.LUNIQUE), name="u"), lambda u: u.replace(arg=0) if u.arg != 0 else None), +]) def _strip_unique(u: UOp) -> UOp: return graph_rewrite(u, _strip_unique_pm) def _t(*shape): @@ -49,6 +51,14 @@ def test_floordiv_bool(self): _check(self, _t(4).cast(dtypes.bool), lambda x: x def test_mod_bool(self): _check(self, _t(4).cast(dtypes.bool), lambda x: x % True) def test_fmod_bool(self): _check(self, _t(4).cast(dtypes.bool), lambda x: x.fmod(True)) +class TestTensorUOpClone(unittest.TestCase): + def test_clone(self): + t = _t(3, 4).float() + self.assertIs(_strip_unique(t.clone().uop), _strip_unique(t.uop.clone())) + def test_clone_deviceless_const(self): + u = UOp.const(dtypes.float, 2.0) + self.assertIs(_strip_unique(Tensor(u).clone().uop), _strip_unique(u.clone())) + class TestTensorUOpGetitem(unittest.TestCase): # ---- pure slice patterns ---- def test_slice_full(self): _check(self, _t(4), lambda x: x[slice(None)]) @@ -384,10 +394,8 @@ def test_full_kwargs(self): self.assertIs(_strip_unique(Tensor.full((2, 3), 42, dtype=dtypes.int8, device="NULL").uop), _strip_unique(UOp.full((2, 3), 42, dtype=dtypes.int8, device="NULL"))) def test_full_symbolic_fill(self): - # bound symbolic variable — flows through Tensor.__init__'s UOp branch, no UNIQUE added t = Tensor.full((2, 3), UOp.variable("x", 1, 10).bind(5)) self.assertEqual(t.shape, (2, 3)) - self.assertFalse(t.uop.op_in_backward_slice_with_self(Ops.UNIQUE)) def test_zeros(self): self.assertIs(_strip_unique(Tensor.zeros(2, 3).uop), _strip_unique(UOp.zeros(2, 3))) def test_ones(self): diff --git a/test/null/test_transcendental_helpers.py b/test/null/test_transcendental_helpers.py index 55dbc8f395554..b6e95ee9328b9 100644 --- a/test/null/test_transcendental_helpers.py +++ b/test/null/test_transcendental_helpers.py @@ -1,7 +1,7 @@ import unittest, math import numpy as np from tinygrad import dtypes -from tinygrad.uop.ops import UOp, Ops +from tinygrad.uop.ops import UOp from tinygrad.uop.decompositions import TRANSCENDENTAL_DTYPES, payne_hanek_reduction, cody_waite_reduction from tinygrad.uop.decompositions import frexp, rintk, xpow, xexp2, xlog2, trig_poly, pow2if from test.helpers import eval_uop @@ -10,7 +10,7 @@ class TestTranscendentalFunctions(unittest.TestCase): def test_payne_hanek_reduction(self): # TODO: Test constant input when constant folding is fixed (or maybe test both variants) # Load input value from a buffer to prevent constant folding - input_buf = UOp(Ops.PARAM, dtypes.double.ptr(), arg=1, src=()) + input_buf = UOp.param(1, dtypes.double.ptr()) loaded_value = input_buf.index(UOp.const(dtypes.int, 0)) def eval_payne_hanek_reduction(v:float) -> tuple[float, int]: return tuple(eval_uop(u, [(dtypes.float64, [v])]) for u in payne_hanek_reduction(loaded_value)) diff --git a/test/null/test_uop_graph.py b/test/null/test_uop_graph.py index 56153b6f21477..2fda5df0f656f 100644 --- a/test/null/test_uop_graph.py +++ b/test/null/test_uop_graph.py @@ -14,6 +14,11 @@ ((UPat.var('x') + UPat.cvar('c1')) + UPat.cvar('c2'), lambda x,c1,c2: x + (c1.arg+c2.arg)), ]) +def const_values(u:UOp): + if u.op is Ops.CONST: return (u.arg,)*u.dtype.count + if u.op is Ops.STACK: return tuple(x.arg for x in u.src) + raise AssertionError(f"expected const-like UOp, got {u.op}") + class TestGraphRewriteConst(unittest.TestCase): def test_gep_const(self): v1 = UOp.const(dtypes.int.vec(3), (0,1,2)) @@ -33,9 +38,9 @@ def test_add_const(self): v1 = UOp.const(dtypes.int.vec(3), (0,1,2)) v2 = UOp.const(dtypes.int.vec(3), (5,6,7)) ret = graph_rewrite(v1+v2, sym) - self.assertEqual(ret.op, Ops.VCONST) + self.assertEqual(ret.op, Ops.STACK) self.assertEqual(ret.dtype, dtypes.int.vec(3)) - self.assertEqual(ret.arg, (5,7,9)) + self.assertEqual(const_values(ret), (5,7,9)) def test_add_const_lose_v(self): v1 = UOp.const(dtypes.int.vec(3), (0,1,2)) @@ -255,7 +260,7 @@ def test_const_shape_change_bitcast(self): @unittest.skip("this test isn't valid uops") def test_noop_vectorize_fold(self): - d0 = UOp(Ops.PARAM, dtypes.float.ptr(), arg=0) + d0 = UOp.param(0, dtypes.float.ptr()) idx = UOp.const(dtypes.int, 0) ld = UOp(Ops.LOAD, dtypes.float.vec(2), (d0, idx)) vec = UOp(Ops.STACK, dtypes.float.vec(2), (ld,)) @@ -267,9 +272,9 @@ def test_noop_vectorize_fold(self): @unittest.skip("this test isn't valid uops") def test_gep_vec_fold(self): - d0 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 0) - d1 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 1) - d2 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 2) + d0 = UOp.param(0, dtypes.float.ptr()) + d1 = UOp.param(1, dtypes.float.ptr()) + d2 = UOp.param(2, dtypes.float.ptr()) idx = UOp.const(dtypes.int, 0) def _test_vec(geps, count=4): vec = UOp(Ops.STACK, dtypes.float.vec(count), geps) @@ -375,8 +380,8 @@ def test_wmma_vectorize_no_fold(self): self.assertEqual(uops[-2], wmma) # -2 to skip SINK def test_cast_alu_fold(self): - d0 = UOp(Ops.PARAM, dtypes.bool.ptr(), arg=0) - d1 = UOp(Ops.PARAM, dtypes.int.ptr(), arg=1) + d0 = UOp.param(0, dtypes.bool.ptr()) + d1 = UOp.param(1, dtypes.int.ptr()) idx = UOp.const(dtypes.int, 0) ld = d1.index(idx) alu = (ld<1).cast(dtypes.bool) @@ -385,8 +390,8 @@ def test_cast_alu_fold(self): self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 0) def test_double_cast_fold(self): - d0 = UOp(Ops.PARAM, dtypes.float.ptr(), arg=0) - d1 = UOp(Ops.PARAM, dtypes.int.ptr(), arg=1) + d0 = UOp.param(0, dtypes.float.ptr()) + d1 = UOp.param(1, dtypes.int.ptr()) idx = UOp.const(dtypes.int, 0) ld = d1.index(idx) alu = ld.cast(dtypes.float).cast(dtypes.float) @@ -409,7 +414,7 @@ def test_depth_2_const_fold(self): def test_bitcast_to_same_dtype_fold(self): for dt in dtypes.ints + dtypes.floats + (dtypes.bool,): - d0 = UOp(Ops.PARAM, dt.ptr(), arg=0) + d0 = UOp.param(0, dt.ptr()) v = d0.index(UOp.const(dtypes.int, 0)) uops = to_uops_list([v.bitcast(dt)]) self.assertEqual(len([x for x in uops if x.op is Ops.BITCAST and x.dtype is dt]), 0, f"dtype = {dt}") @@ -422,10 +427,10 @@ def test_sub_with_cast_folds(self): def test_where_on_gated_load_fold(self): ridx0 = UOp.range(100, 0) - d0 = UOp(Ops.PARAM, dtypes.long.ptr(), (), 0) + d0 = UOp.param(0, dtypes.long.ptr()) ld = d0.index(ridx0.valid(ridx0<50)) w = (ridx0<50).where(ld, 5) - out = UOp(Ops.PARAM, dtypes.long.ptr(), (), 1) + out = UOp.param(1, dtypes.long.ptr()) uops = to_uops_list([out.index(ridx0).store(w)]) for u in uops: assert u.op is not Ops.WHERE @@ -433,7 +438,7 @@ def test_where_on_gated_load_fold(self): def test_where_on_gated_load_folds_swapped_branches(self): ridx0 = UOp.range(100, 0) - d0 = UOp(Ops.PARAM, dtypes.long.ptr(), (), 0) + d0 = UOp.param(0, dtypes.long.ptr()) ld = d0.index(ridx0.valid((ridx0<50).logical_not())) w = (ridx0<50).where(5, ld) uops = to_uops_list([w]) @@ -443,11 +448,11 @@ def test_where_on_gated_load_folds_swapped_branches(self): def test_where_on_gated_load_with_cast(self): ridx0 = UOp.range(100, 0) - d0 = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0) + d0 = UOp.param(0, dtypes.int.ptr()) gate_idx = ridx0.valid((ridx0<50)) ld = d0.index(gate_idx).cast(dtypes.float) w = (ridx0<50).where(ld, 5.0) - out = UOp(Ops.PARAM, dtypes.float.ptr(), (), 1) + out = UOp.param(1, dtypes.float.ptr()) uops = to_uops_list([out.index(ridx0).store(w)]) for u in uops: assert u.op is not Ops.WHERE @@ -455,27 +460,27 @@ def test_where_on_gated_load_with_cast(self): def test_where_on_casted_gated_load_extra_cond(self): ridx0 = UOp.range(100, 0) - d0 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 0) + d0 = UOp.param(0, dtypes.float.ptr()) ld = d0.index(ridx0.valid(ridx0<50)) w = ((ridx0<50) & (ridx0>30)).where(ld, UOp.const(dtypes.float, 0)).cast(dtypes.half) - out = UOp(Ops.PARAM, dtypes.half.ptr(), (), 1) + out = UOp.param(1, dtypes.half.ptr()) uops = to_uops_list([out.index(ridx0).store(w)]) for u in uops: assert u.op is not Ops.WHERE def test_where_on_casted_gated_load_extra_cond_swapped(self): ridx0 = UOp.range(100, 0) - d0 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 0) + d0 = UOp.param(0, dtypes.float.ptr()) ld = d0.index(ridx0.valid(ridx0<50)) w = ((ridx0<50) & (ridx0>30)).where(UOp.const(dtypes.float, 0), ld).cast(dtypes.half) - out = UOp(Ops.PARAM, dtypes.half.ptr(), (), 1) + out = UOp.param(1, dtypes.half.ptr()) uops = to_uops_list([out.index(ridx0).store(w)]) for u in uops: assert u.op is not Ops.WHERE def test_where_in_store_becomes_gate(self): ridx0 = UOp.range(100, 0) - d0 = UOp(Ops.PARAM, dtypes.long.ptr(), (), 0) + d0 = UOp.param(0, dtypes.long.ptr()) idx = d0.index(ridx0) ld = idx.load() val = (ridx0<50).where(5, ld) @@ -488,14 +493,14 @@ def test_where_in_store_becomes_gate(self): def test_load_idx_becomes_int(self): # mnist indexing with split reduceop # Make sure we are not doign math on the loaded index, which would promote it to long - c0 = UOp(Ops.PARAM, dtypes.uchar.ptr(128000), arg=0, src=()) + c0 = UOp.param(0, dtypes.uchar.ptr(128000)) c1 = UOp.range(UOp.const(dtypes.weakint, 512), 1, AxisType.LOOP) c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.LOOP) - c3 = UOp(Ops.PARAM, dtypes.int.ptr(512), arg=1, src=()) + c3 = UOp.param(1, dtypes.int.ptr(512)) c4 = c3.index(c1) c5 = UOp.range(UOp.const(dtypes.weakint, 240), 0, AxisType.REDUCE) c6 = ((c2*UOp.const(dtypes.weakint, 240))+c5) - c7 = UOp(Ops.PARAM, dtypes.uchar.ptr(60000), arg=2, src=()) + c7 = UOp.param(2, dtypes.uchar.ptr(60000)) c8 = c7.index(c6) c9 = ((c4<0).where((c4+60000), c4)!=c6.cast(dtypes.int)).where(0, c8.cast(dtypes.uint).cast(dtypes.uchar)).reduce(c5, arg=Ops.ADD) c10 = c0.index(((c1*UOp.const(dtypes.weakint, 250))+c2)).store(c9).end(c1, c2) @@ -505,14 +510,14 @@ def test_load_idx_becomes_int(self): def test_load_idx_no_math_on_loaded(self): # test the (x+y) NOOP rule. This rule matches patterns that EMERGE during simplification.""" def test_store_load_folding(self): # store(idx, load(idx)) -> NOOP, including emergent patterns like store(idx, load(idx) + 0) - buf = UOp(Ops.PARAM, dtypes.int.ptr(), arg=0) + buf = UOp.param(0, dtypes.int.ptr()) index = buf.index(UOp.const(dtypes.weakint, 0)) # Direct: store(idx, load(idx)) -> NOOP self.assertEqual(graph_rewrite(index.store(index.load()), sym).op, Ops.NOOP) @@ -1340,7 +1340,7 @@ def test_range_split_on_mod(self): from tinygrad.codegen.simplify import pm_split_ranges, pm_flatten_range r0 = UOp.range(uconst(8), 0) # create a simple expression using the range with mod: store range%2 to a buffer - buf = UOp(Ops.PARAM, dtypes.int.ptr(), arg=0) + buf = UOp.param(0, dtypes.int.ptr()) val = (r0 % uconst(2)).cast(dtypes.int) store = UOp(Ops.STORE, dtypes.void, (buf.index(uconst(0)), val)) sink = UOp(Ops.SINK, dtypes.void, (UOp(Ops.END, dtypes.void, (store, r0)),)) diff --git a/test/null/test_uop_vmin_vmax.py b/test/null/test_uop_vmin_vmax.py index 0e4c1ed9b93e4..c7f794e95400d 100644 --- a/test/null/test_uop_vmin_vmax.py +++ b/test/null/test_uop_vmin_vmax.py @@ -82,7 +82,7 @@ def test_vmin_vmax_variable_inside_special(self): def test_vmin_vmax_multiplication_0_inf(self): # vmin and vmax for multiplication with a variable x = UOp.const(dtypes.float, 0.0) - y = UOp.load(UOp(Ops.PARAM, dtypes.float.ptr(), (), 0), UOp.const(dtypes.int, 0), dtype=dtypes.float) + y = UOp.load(UOp.param(0, dtypes.float.ptr(1)), UOp.const(dtypes.int, 0), dtype=dtypes.float) uop = x * y # TODO: these should be 0, but definitely should not be nan self.assertEqual(uop.vmin, -math.inf) @@ -316,7 +316,7 @@ def test_vmin_vmax_vconst_with_bools(self): def test_vmin_vmax_vector_with_gep(self): # vmin and vmax for a vector constant of bool values - d1 = UOp(Ops.PARAM, dtypes.int.ptr(), (), 1) + d1 = UOp.param(1, dtypes.int.ptr()) idx = UOp.const(dtypes.int, 0) val = UOp(Ops.LOAD, dtypes.int.vec(2), (d1.index(idx).cast(dtypes.int.vec(2).ptr()),)) uop = (val // 32).gep(0) diff --git a/test/null/test_uops.py b/test/null/test_uops.py index 2fc533ad85abc..18ff80efc2032 100644 --- a/test/null/test_uops.py +++ b/test/null/test_uops.py @@ -110,7 +110,7 @@ def test_overflow(self): class TestGatedStoreRewrite(unittest.TestCase): def test_tiny_gate_store(self): - gmem = UOp(Ops.PARAM, dtypes.float.ptr(), (), 0) + gmem = UOp.param(0, dtypes.float.ptr()) gidx0 = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'gidx0') gate = gidx0= 0) & (v < 16)), ptr=True).load(dtype=dtypes.int)]) # valid with self.assertRaises(RuntimeError): @@ -46,7 +46,7 @@ def test_variable_with_mask(self): def test_gated_store(self): with Context(CHECK_OOB=1, SPEC=2): - buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0) + buf = UOp.param(0, dtypes.int.ptr(16)) v = Variable("v", 0, 20) to_uops_list([buf.index(v.valid(v < 16), ptr=True).store(0)]) # valid with self.assertRaises(RuntimeError): @@ -55,14 +55,14 @@ def test_gated_store(self): # ALU ops in index def test_floordiv(self): with Context(CHECK_OOB=1, SPEC=2): - buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0) + buf = UOp.param(0, dtypes.int.ptr(16)) to_uops_list([buf.index(UOp.range(32, 0, AxisType.GLOBAL) // 2, ptr=True).load(dtype=dtypes.int)]) # 0..15 valid with self.assertRaises(RuntimeError): to_uops_list([buf.index(UOp.range(34, 0, AxisType.GLOBAL) // 2, ptr=True).load(dtype=dtypes.int)]) # 0..16 oob def test_mod(self): with Context(CHECK_OOB=1, SPEC=2): - buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0) + buf = UOp.param(0, dtypes.int.ptr(16)) r = UOp.range(100, 0, AxisType.GLOBAL) to_uops_list([buf.index(r % 16, ptr=True).load(dtype=dtypes.int)]) # 0..15 valid with self.assertRaises(RuntimeError): @@ -70,14 +70,14 @@ def test_mod(self): def test_shr(self): with Context(CHECK_OOB=1, SPEC=2): - buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0) + buf = UOp.param(0, dtypes.int.ptr(16)) to_uops_list([buf.index(UOp.range(64, 0, AxisType.GLOBAL) >> 2, ptr=True).load(dtype=dtypes.int)]) # 0..15 valid with self.assertRaises(RuntimeError): to_uops_list([buf.index(UOp.range(128, 0, AxisType.GLOBAL) >> 2, ptr=True).load(dtype=dtypes.int)]) # 0..31 oob def test_shl(self): with Context(CHECK_OOB=1, SPEC=2): - buf = UOp(Ops.PARAM, dtypes.int.ptr(64), (), 0) + buf = UOp.param(0, dtypes.int.ptr(64)) r = UOp.range(8, 0, AxisType.GLOBAL) to_uops_list([buf.index(r << 2, ptr=True).load(dtype=dtypes.int)]) # 0..28 valid with self.assertRaises(RuntimeError): @@ -85,7 +85,7 @@ def test_shl(self): def test_and(self): with Context(CHECK_OOB=1, SPEC=2): - buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0) + buf = UOp.param(0, dtypes.int.ptr(16)) r = UOp.range(100, 0, AxisType.GLOBAL) to_uops_list([buf.index(r & 15, ptr=True).load(dtype=dtypes.int)]) # 0..15 valid with self.assertRaises(RuntimeError): @@ -93,14 +93,14 @@ def test_and(self): def test_max(self): with Context(CHECK_OOB=1, SPEC=2): - buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0) + buf = UOp.param(0, dtypes.int.ptr(16)) to_uops_list([buf.index(Variable("v", -10, 15).maximum(0), ptr=True).load(dtype=dtypes.int)]) # 0..15 valid with self.assertRaises(RuntimeError): to_uops_list([buf.index(Variable("v2", -10, 20).maximum(0), ptr=True).load(dtype=dtypes.int)]) # 0..20 oob def test_xor_in_mask(self): with Context(CHECK_OOB=1, SPEC=2): - buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0) + buf = UOp.param(0, dtypes.int.ptr(16)) r = UOp.range(32, 0, AxisType.GLOBAL) to_uops_list([buf.index(r.valid((r < 8) ^ ((r >= 8) & (r < 16))), ptr=True).load(dtype=dtypes.int)]) # 0..15 valid with self.assertRaises(RuntimeError): @@ -109,22 +109,22 @@ def test_xor_in_mask(self): # cast patterns def test_float_cast_in_index(self): with Context(CHECK_OOB=1, SPEC=2): - buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0) + buf = UOp.param(0, dtypes.int.ptr(16)) r = UOp.range(20, 0) i = (r.cast(dtypes.float) * 0.68).trunc().cast(dtypes.int) to_uops_list([buf.index(i.valid((i >= 0) & (i < 16)), ptr=True).load(dtype=dtypes.int)]) def test_bool_cast_in_mask(self): with Context(CHECK_OOB=1, SPEC=2): - buf = UOp(Ops.PARAM, dtypes.int.ptr(1), (), 0) + buf = UOp.param(0, dtypes.int.ptr(1)) r = UOp.range(20, 0) to_uops_list([buf.index(r.valid(r.cast(dtypes.bool).logical_not()), ptr=True).load(dtype=dtypes.int)]) # only r=0 valid # load result as index/mask def test_load_as_index(self): with Context(CHECK_OOB=1, SPEC=2): - buf0 = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0) - buf1 = UOp(Ops.PARAM, dtypes.int.ptr(64), (), 1) + buf0 = UOp.param(0, dtypes.int.ptr(16)) + buf1 = UOp.param(1, dtypes.int.ptr(64)) r = UOp.range(42, 0, AxisType.GLOBAL) ld0 = buf0.index(r.valid(r < 8), ptr=True).load(dtype=dtypes.int).cast(dtypes.weakint) to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 32)), ptr=True).load(dtype=dtypes.int)]) # valid @@ -133,8 +133,8 @@ def test_load_as_index(self): def test_load_bool_as_mask(self): with Context(CHECK_OOB=1, SPEC=2): - buf_bool = UOp(Ops.PARAM, dtypes.bool.ptr(16), (), 0) - buf_int = UOp(Ops.PARAM, dtypes.int.ptr(8), (), 1) + buf_bool = UOp.param(0, dtypes.bool.ptr(16)) + buf_int = UOp.param(1, dtypes.int.ptr(8)) gidx = UOp(Ops.SPECIAL, dtypes.weakint, (UOp.const(dtypes.weakint, 16),), "gidx0") ld_bool = buf_bool.index(gidx, ptr=True).load() with self.assertRaises(RuntimeError): @@ -145,7 +145,7 @@ def test_load_bool_as_mask(self): def test_in_bounds_access_gated_local(self): with Context(CHECK_OOB=1): # Define buffers - gbuf = UOp(Ops.PARAM, dtypes.uint.ptr(400), (), 0) + gbuf = UOp.param(0, dtypes.uint.ptr(400)) sbuf = UOp(Ops.DEFINE_LOCAL, dtypes.uint.ptr(8, addrspace=AddrSpace.LOCAL), (), "temp0") # Define indices, valids and barrier @@ -169,8 +169,8 @@ def test_in_bounds_access_gated_local(self): @unittest.skip("Bool load is not supported yet") def test_load_mask(self): with Context(CHECK_OOB=1): - glbl0 = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0) - mask = UOp(Ops.PARAM, dtypes.bool.ptr(16), (), 0) + glbl0 = UOp.param(0, dtypes.int.ptr(16)) + mask = UOp.param(0, dtypes.bool.ptr(16)) ridx = UOp.range(20, 0) ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(ridx, ridx<16&mask), ptr=True))) to_uops_list([ld0]) diff --git a/test/null/test_viz.py b/test/null/test_viz.py index 01c0222dbc44c..d3727a09d06ec 100644 --- a/test/null/test_viz.py +++ b/test/null/test_viz.py @@ -1,19 +1,18 @@ -import unittest, decimal, sys, json, contextlib, tempfile, pickle, io +import unittest, decimal, sys, json, contextlib, tempfile, pickle, io, math from pathlib import Path from dataclasses import dataclass from typing import Generator from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, TrackedPatternMatcher, graph_rewrite, track_rewrites, profile_matches from tinygrad.uop.symbolic import sym -from tinygrad.dtype import dtypes +from tinygrad.dtype import dtypes, AddrSpace from tinygrad.helpers import colored, ansistrip, flatten, TracingKey, ProfileRangeEvent, ProfileEvent, Context, cpu_events, profile_marker from tinygrad.helpers import cpu_profile, ProfilePointEvent, unwrap from tinygrad.device import Buffer from tinygrad.uop.ops import tracked_keys, tracked_ctxs, uop_fields, active_rewrites, active_group, _name_cnt, RewriteTrace -from tinygrad.viz.serve import load_rewrites, get_full_rewrite, uop_to_json, VizData, get_render -from tinygrad.codegen import to_program_cache -from tinygrad.codegen import to_program +from tinygrad.viz.serve import load_rewrites, get_full_rewrite, uop_to_json, VizData, get_render, addrspace_colors +from tinygrad.codegen import do_to_program @track_rewrites(name=True) def exec_rewrite(sink:UOp, pm_lst:list[PatternMatcher], names:None|list[str]=None) -> UOp: @@ -41,7 +40,6 @@ def get_details(self, rewrite_idx:int, step:int) -> Generator[dict, None, None]: @contextlib.contextmanager def save_viz(): for lst in [tracked_keys, tracked_ctxs, active_rewrites, active_group, _name_cnt]: lst.clear() - to_program_cache.clear() Buffer.profile_events.clear() cpu_events.clear() viz = VizTrace() @@ -219,32 +217,47 @@ def test_const_node_visibility(self): with save_viz() as viz: a = UOp.variable("a", 0, 10, dtype=dtypes.int) z = UOp.const(a.dtype, 0) + y = UOp.const(dtypes.float, math.pi) alu = a*z - exec_rewrite(alu, [sym]) + ret = exec_rewrite(sink:=UOp.sink(alu, y), [sym]) lst = viz.list_items() self.assertEqual(len(lst), 1) graphs = [x["graph"] for x in viz.get_details(0, 0)] - # embed const in the parent node when possible - self.assertEqual(list(graphs[0]), [id(a), id(alu)]) - self.assertEqual(list(graphs[1]), [id(z)]) - - # TODO: DEFINE_VAR (shape ()) now gets wrapped in RESHAPE+EXPAND when broadcast against a shaped operand - # (due to shared OpMixin._binop using _broadcasted). Either extend viz to fold RESHAPE/EXPAND around - # DEFINE_VAR/RANGE/SPECIAL the way it does for CONST, or redesign scalar-compiler-op broadcasting. - @unittest.expectedFailure + # const is always in the graph, client side hides exclude=True nodes by default + self.assertEqual(list(graphs[0]), [id(a), id(z), id(alu), id(y), id(sink)]) + self.assertTrue(graphs[0][id(z)]["exclude"]) + self.assertTrue(graphs[0][id(y)]["exclude"]) + self.assertFalse(graphs[0][id(alu)]["exclude"]) + self.assertEqual(graphs[0][id(y)]["label"].split("\n")[:2], ["CONST", "3.14159"]) + self.assertEqual(list(graphs[1]), [id(z), id(y), id(ret)]) + def test_const_reshape_expand_folded(self): # CONST->RESHAPE->EXPAND should be folded into the ALU node, not shown as separate RESHAPE/EXPAND nodes c = UOp.const(dtypes.float, 1.0, device="CPU", shape=(3,4)) # creates CONST->RESHAPE->EXPAND chain a = UOp(Ops.DEFINE_VAR, dtypes.float, arg=("a", 0.0, 10.0)) alu = a + c - graph = uop_to_json(VizData(), alu) - # the RESHAPE and EXPAND nodes from the const should not appear in the graph - labels = {v["label"].split("\n")[0] for v in graph.values()} - self.assertNotIn("RESHAPE", labels) - self.assertNotIn("EXPAND", labels) - # the CONST should be inlined into the ALU node's label - alu_label = graph[id(alu)]["label"] - self.assertIn("CONST", alu_label) + with save_viz() as viz: + graph_rewrite(alu, PatternMatcher([])) + graph = [x["graph"] for x in viz.get_details(0, 0)][0] + excluded_nodes = {v["label"].split("\n")[0] for v in graph.values() if v["exclude"]} + self.assertIn("CONST", excluded_nodes) + self.assertIn("STACK", excluded_nodes) + self.assertIn("RESHAPE", excluded_nodes) + self.assertIn("EXPAND", excluded_nodes) + self.assertIn("CONST1 1 Ops.DEVICE", graph[id(alu)]["label"]) + + def test_stack_movement_not_folded_unless_all_const(self): + a = UOp.variable("a", 0, 10, dtype=dtypes.int) + c = UOp.const(dtypes.int, 1) + stack = a.vectorize(c) + reshaped = stack.reshape((1, 2)) + graph = uop_to_json(VizData(), reshaped) + self.assertFalse(graph[id(stack)]["exclude"]) + + const_stack = c.vectorize(UOp.const(dtypes.int, 2)) + const_reshaped = const_stack.reshape((1, 2)) + const_graph = uop_to_json(VizData(), const_reshaped) + self.assertTrue(const_graph[id(const_stack)]["exclude"]) # VIZ displays nested graph_rewrites in a tree view @@ -327,12 +340,15 @@ class TestVizIntegration(unittest.TestCase): def test_codegen_tracing(self): with save_viz() as viz: ast = (Tensor.empty(4)+Tensor.empty(4)).schedule_linear().src[0].src[0] - prg = to_program(ast, Device[Device.DEFAULT].renderer) + prg = do_to_program(ast, Device[Device.DEFAULT].renderer) lst = viz.list_items() self.assertEqual(len(lst), 3) self.assertEqual(lst[0]["name"], "Callify 1 Buffer n1") self.assertEqual(lst[1]["name"], "Schedule 1 Kernel n1") self.assertEqual(lst[2]["name"], prg.arg.name) + input_ast = next(viz.get_details(2, 0))["graph"].values() + for u in input_ast: + if u["label"].startswith("PARAM\n"): self.assertEqual(u["addrspace"], addrspace_colors[AddrSpace.GLOBAL]) # schedule graph CALL nodes have a link to jump to codegen def test_link_sched_codegen(self): @@ -344,7 +360,7 @@ def test_link_sched_codegen(self): from tinygrad.engine.realize import compile_linear sched = compile_linear(sched) with Context(NO_COLOR=0): - prgs = [to_program(si.src[0], Device[c1.device].renderer).arg.name for si in sched.src] + prgs = [do_to_program(si.src[0], Device[c1.device].renderer).arg.name for si in sched.src] lst = viz.list_items() sched_idx = next(i for i,l in enumerate(lst) if l["name"].startswith("Schedule")) viz_kernel = next(i for i,s in enumerate(lst[sched_idx]["steps"]) if s["name"] == "View Kernel Graph") @@ -406,9 +422,16 @@ def default_test(root): return graph_rewrite(root, sym) default_test(c+2) ls = viz.list_items() self.assertEqual(len(ls), 2) - self.assertEqual(list(next(viz.get_details(0, 0))["graph"]), [id(c+1)]) + graph = next(viz.get_details(0, 0))["graph"] + self.assertEqual(list(graph), [id(c), id(c+1)]) + self.assertTrue(graph[id(c)]["exclude"]) + self.assertFalse(graph[id(c+1)]["exclude"]) self.assertEqual(list(next(viz.get_details(1, 0))["graph"]), [id(c)]) - self.assertEqual(list(next(viz.get_details(1, 1))["graph"]), [id(c+2)]) + graph = next(viz.get_details(1, 1))["graph"] + self.assertEqual(list(graph), [id(c), id(c.const_like(2)), id(c+2)]) + self.assertTrue(graph[id(c)]["exclude"]) + self.assertTrue(graph[id(c.const_like(2))]["exclude"]) + self.assertFalse(graph[id(c+2)]["exclude"]) def test_recurse(self): with save_viz() as viz: @@ -744,7 +767,7 @@ def fxn(out:UOp) -> UOp: with save_viz() as viz: with Context(DEV=f"NULL::{self.arch}"): out = Tensor.custom_kernel(Tensor.empty(1), fxn=fxn)[0] - _ = to_program(out.schedule_linear().src[-1].src[0], Device[out.device].renderer) + _ = do_to_program(out.schedule_linear().src[-1].src[0], Device[out.device].renderer) codegen_rewrites = next(s for s in viz.list_items() if s["name"] == name) disasm = next(s for s in codegen_rewrites["steps"] if s["name"] == "View Disassembly") return get_render(viz.data, disasm["query"]) @@ -981,8 +1004,9 @@ def f(a, b): return (a@a.T), (b@b.T) def test_dedup(self): with save_viz() as viz: for _ in range(CNT:=4): - Tensor.empty(4, device="NULL").add(1).realize() - Tensor.empty(8, device="NULL").add(1).realize() + # use kernel names unique to this test + Tensor.custom_kernel(Tensor.empty(4, device="NULL"), fxn=lambda _: UOp.sink(arg=KernelInfo("k1_test_viz_dedup")))[0].realize() + Tensor.custom_kernel(Tensor.empty(8, device="NULL"), fxn=lambda _: UOp.sink(arg=KernelInfo("k2_test_viz_dedup")))[0].realize() with write_files(viz) as files, Context(NO_COLOR=1): name = run_cli(*files, "-s", "NULL")[0]["name"] with Context(DEBUG=3): @@ -1008,5 +1032,22 @@ def f(x): for i,n in enumerate(call_nodes): assert prgs[i] in n["label"], f"CALL must contain kernel name, got {n['label']}" + def test_interval(self): + def emit_kernel(name:str): Tensor.custom_kernel(Tensor.empty(1, device="NULL"), fxn=lambda _: UOp.sink(arg=KernelInfo(name=name)))[0].realize() + with save_viz() as viz: + emit_kernel("pre_1") + emit_kernel("pre_2") + profile_marker("interval_start") + emit_kernel("target_1") + emit_kernel("target_2") + profile_marker("interval_end") + emit_kernel("post_1") + emit_kernel("post_2") + with write_files(viz) as files, Context(NO_COLOR=1): + flat = run_cli(*files, "-s", "NULL", "--interval", "interval_start", "interval_end") + aggregate = run_cli(*files, "-s", "NULL", "--interval", "interval_start", "interval_end", "-t") + self.assertEqual([s["name"] for s in flat], ["interval_start", "target_1", "target_2", "interval_end"]) + self.assertEqual(sorted(s["name"] for s in aggregate), ["target_1", "target_2"]) + if __name__ == "__main__": unittest.main() diff --git a/test/null/test_winograd.py b/test/null/test_winograd.py index 8e3402e1f2dd1..e9ac04b8bedbb 100644 --- a/test/null/test_winograd.py +++ b/test/null/test_winograd.py @@ -1,6 +1,6 @@ import unittest, sys from tinygrad import Tensor, GlobalCounters, dtypes, Context -from tinygrad.helpers import CI, Profiling, WINO +from tinygrad.helpers import WINO @unittest.skipIf(sys.platform.startswith("win"), "flaky on Windows") class TestWinograd(unittest.TestCase): @@ -10,23 +10,19 @@ def setUp(self): def tearDown(self): WINO.value = self.old - def test_profile(self): - x,w = Tensor.rand(1,4,9,9).realize(), Tensor.rand(4,4,3,3).realize() - with Profiling(enabled=not CI, sort='time'): - Tensor.conv2d(x,w).realize() - def test_forward_kernels(self): x,w = Tensor.rand(1,4,9,9).realize(), Tensor.rand(4,4,3,3).realize() out = Tensor.conv2d(x,w) self.assertEqual(len(out.schedule_linear().src), 2) def test_backward_kernels(self): - x,w = Tensor.empty(1,4,9,9,requires_grad=True).realize(), Tensor.empty(4,4,3,3,requires_grad=True).realize() + x,w = Tensor.empty(1,4,9,9).realize(), Tensor.empty(4,4,3,3).realize() out = Tensor.conv2d(x,w, padding=1) out.mean().backward() backward_schedule = x.grad.schedule_linear(w.grad) self.assertEqual(len(backward_schedule.src), 4) + @unittest.skip("this requires optimizations") def test_counters(self): IC, OC, X, Y = 4,4,9,9 x,w = Tensor.rand(1,IC,Y,X).realize(), Tensor.rand(OC,IC,3,3).realize() diff --git a/test/opt/test_tensor_cores.py b/test/opt/test_tensor_cores.py index 957bed5efca17..9d6998b5764d5 100644 --- a/test/opt/test_tensor_cores.py +++ b/test/opt/test_tensor_cores.py @@ -5,7 +5,7 @@ from tinygrad.tensor import _to_np_dtype from tinygrad.uop.ops import Ops, UOp, buffers from tinygrad.dtype import DType -from tinygrad.device import Buffer, is_dtype_supported +from tinygrad.device import Buffer from tinygrad.helpers import DEV, Context from test.helpers import slow, replace_opts from tinygrad.engine.realize import run_linear @@ -69,7 +69,6 @@ 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: - if not is_dtype_supported(tc.dtype_in) or not is_dtype_supported(tc.dtype_out): continue # for AMX, tc.dims[2] == 1 so reduceop is None thus tensor_cores are not triggered helper_tc_allclose(tc.dims[0], tc.dims[1], 2 if AMX else tc.dims[2], tc.dtype_in, tc.dtype_out, axis=0, tc_opt=0) @@ -78,7 +77,6 @@ def test_tensor_cores(self): @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores") def test_tensor_cores_codegen(self): for tc in Device[Device.DEFAULT].renderer.tensor_cores: - if not is_dtype_supported(tc.dtype_in) or not is_dtype_supported(tc.dtype_out): continue n, m, k = tc.dims[0], tc.dims[1], 2 if AMX else tc.dims[2] a, b = Tensor.rand(m, k, dtype=tc.dtype_in), Tensor.rand(k, n, dtype=tc.dtype_in) r = a.matmul(b, dtype=tc.dtype_out) @@ -98,7 +96,6 @@ def test_tensor_cores_codegen(self): @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores") def test_tensor_cores_padded(self): for tc in Device[Device.DEFAULT].renderer.tensor_cores: - if not is_dtype_supported(tc.dtype_in) or not is_dtype_supported(tc.dtype_out): continue helper_tc_allclose(tc.dims[0]+(pad:=1), tc.dims[1]+pad, tc.dims[2]+pad, tc.dtype_in, tc.dtype_out, tc_opt=2) # AMD compiler bug: AMD miscompiles non-zero padded tc kernels with -O3, producing wrong results, nans or hang (see #9606) @@ -109,7 +106,6 @@ def test_tensor_cores_padded(self): @unittest.skip("warp elements not duplicated properly across lanes") def test_tensor_cores_padded_amd(self): for tc in Device[Device.DEFAULT].renderer.tensor_cores: - if not is_dtype_supported(tc.dtype_in) or not is_dtype_supported(tc.dtype_out): continue helper_tc_allclose(tc.dims[0]+(pad:=1), tc.dims[1]+pad, tc.dims[2]+pad, tc.dtype_in, tc.dtype_out, tc_opt=2) @Context(ALLOW_TF32=1) @@ -140,7 +136,6 @@ def test_tensor_cores_padded_uops(self): @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores") def test_tensor_cores_multi_reduce(self): for tc in Device[Device.DEFAULT].renderer.tensor_cores: - if not is_dtype_supported(tc.dtype_in) or not is_dtype_supported(tc.dtype_out): continue if tc.dtype_in is dtypes.bfloat16: continue # <-- broken with numpy # this will be a M=G16, N=G32, M=G16, M=G16, K=R16, K=R16, K=R16 with 9 choices of TC MNK axes golden_result = None diff --git a/test/speed/external_test_copy_speed.py b/test/speed/external_test_copy_speed.py index 359a2499ace48..e34b0fc1fe7ff 100644 --- a/test/speed/external_test_copy_speed.py +++ b/test/speed/external_test_copy_speed.py @@ -1,6 +1,6 @@ -import unittest, numpy as np +import unittest, numpy as np, os from tinygrad import Tensor, Device, TinyJit -from tinygrad.helpers import Timing, CI, OSX, getenv +from tinygrad.helpers import Timing, getenv import multiprocessing.shared_memory as shared_memory N = getenv("NSZ", 256) @@ -11,7 +11,7 @@ def setUpClass(cls): Device[Device.DEFAULT].synchronize() def testCopySHMtoDefault(self): s = shared_memory.SharedMemory(name="test_X", create=True, size=N*N*4) s.close() - if CI and not OSX: + if os.path.exists("/dev/shm"): t = Tensor.empty(N, N, device="disk:/dev/shm/test_X").realize() else: t = Tensor.empty(N, N, device="disk:shm:test_X").realize() @@ -76,11 +76,8 @@ def _do_copy(x): return x.to(Device.DEFAULT).realize() Device[Device.DEFAULT].synchronize() np.testing.assert_equal(t.numpy(), x.numpy()) - @unittest.skipIf(CI, "CI doesn't have 6 GPUs") - @unittest.skipIf(Device.DEFAULT != "CL", "only test this on CL") + @unittest.skipIf(Device.DEFAULT != "CL" or Device[Device.DEFAULT].count() != 6, "only test this on CL, with 6 gpus") def testCopyCPUto6GPUs(self): - from tinygrad.runtime.ops_cl import CLDevice - if len(CLDevice.device_ids) != 6: raise unittest.SkipTest("computer doesn't have 6 GPUs") t = Tensor.ones(N, N, device="CPU").contiguous().realize() print(f"buffer: {t.nbytes()*1e-9:.2f} GB") for _ in range(3): diff --git a/test/speed/external_test_specific_conv.py b/test/speed/external_test_specific_conv.py index ed6492b5da98b..190a1ccee9d72 100644 --- a/test/speed/external_test_specific_conv.py +++ b/test/speed/external_test_specific_conv.py @@ -1,10 +1,9 @@ import unittest -from tinygrad.helpers import CI from tinygrad import Tensor, Device, dtypes -from tinygrad.device import is_dtype_supported +from tinygrad.helpers import DEV # similar to test/external/external_test_gpu_ast.py, but universal -@unittest.skipIf(Device.DEFAULT in {"CUDA", "NV"} and CI, "slow on CUDA CI") +@unittest.skipIf(Device.DEFAULT in {"CUDA", "NV"} and DEV.interface.startswith("MOCK"), "slow on ocelot") class TestSpecific(unittest.TestCase): # from openpilot @@ -20,7 +19,7 @@ def test_vec_mul(self): w = Tensor.randn(2048, 512) (x @ w).reshape(1, 128, 4).contiguous().realize() - @unittest.skipUnless(is_dtype_supported(dtypes.float16), "need float16 support") + @unittest.skipUnless(dtypes.float16 in Device[Device.DEFAULT].renderer.supported_dtypes(), "need float16 support") def test_big_vec_mul(self): # from LLaMA # 0 buffer<4096, dtypes.float> [View((1024, 1, 1, 4), (4, 0, 0, 1), 0, None)] diff --git a/test/speed/external_test_speed_v_torch.py b/test/speed/external_test_speed_v_torch.py index 5694c933d0016..bd3520c056088 100644 --- a/test/speed/external_test_speed_v_torch.py +++ b/test/speed/external_test_speed_v_torch.py @@ -9,10 +9,11 @@ torch.set_num_threads(1) import time import numpy as np +import sys np.set_printoptions(linewidth=160) from tinygrad import Tensor, Device, GlobalCounters, TinyJit from tinygrad.nn import Conv2d -from tinygrad.helpers import colorize_float, getenv, CI, DEV +from tinygrad.helpers import colorize_float, getenv, DEV IN_CHANS = [int(x) for x in getenv("IN_CHANS", "4,16,64").split(",")] @@ -95,7 +96,7 @@ def helper_test_generic(name, f1, f1_args, f2, f2_args): desc = "faster" if et_torch > et_tinygrad else "slower" flops = save_ops*1e-6 mem = save_mem*1e-6 - print(("\r" if not CI else "")+f"{name:42s} {et_torch:7.2f} ms ({flops/et_torch:9.2f} GFLOPS {mem/et_torch:7.2f} GB/s) in torch, {et_tinygrad:7.2f} ms ({flops/et_tinygrad:9.2f} GFLOPS {mem/et_tinygrad:7.2f} GB/s) in tinygrad, {colorize_float(et_tinygrad/et_torch)} {desc} {flops:10.2f} MOPS {mem:8.2f} MB") # noqa: E501 + print(("\r" if sys.stdout.isatty() else "")+f"{name:42s} {et_torch:7.2f} ms ({flops/et_torch:9.2f} GFLOPS {mem/et_torch:7.2f} GB/s) in torch, {et_tinygrad:7.2f} ms ({flops/et_tinygrad:9.2f} GFLOPS {mem/et_tinygrad:7.2f} GB/s) in tinygrad, {colorize_float(et_tinygrad/et_torch)} {desc} {flops:10.2f} MOPS {mem:8.2f} MB") # noqa: E501 atol, rtol = (1e-2, 1e-2) if torch_dt == torch.float16 else (1e-3, 1e-3) np.testing.assert_allclose(val_tinygrad, val_torch, atol=atol, rtol=rtol) diff --git a/test/test_tiny.py b/test/test_tiny.py index 2e221853b9fa7..10ad7428dfe78 100644 --- a/test/test_tiny.py +++ b/test/test_tiny.py @@ -1,7 +1,7 @@ # basic self-contained tests of the external functionality of tinygrad import unittest, random from tinygrad import Tensor, Context, Variable, TinyJit, dtypes, Device, nn -from tinygrad.helpers import CI, getenv +from tinygrad.helpers import getenv class TestTiny(unittest.TestCase): @@ -111,7 +111,7 @@ def test_symbolic_reduce(self): # *** a model *** # TODO: this is failing because of how swizzling rewrites the ShapeTracker of the final STORE - @unittest.skipIf(CI and Device.DEFAULT == "DSP", "failing because of make things that can't be images not images") + @unittest.skipIf(Device.DEFAULT == "DSP", "failing because of make things that can't be images not images") def test_mnist(self): layers = [ nn.Conv2d(1, 32, 5), Tensor.relu, @@ -130,7 +130,7 @@ def test_mnist(self): self.assertEqual(len(probs[0]), 10) # TODO: this is failing because of how swizzling rewrites the ShapeTracker of the final STORE - @unittest.skipIf(CI and Device.DEFAULT == "DSP", "failing because of make things that can't be images not images") + @unittest.skipIf(Device.DEFAULT == "DSP", "failing because of make things that can't be images not images") def test_mnist_backward(self): # NOTE: we don't have the whole model here for speed layers = [ @@ -141,7 +141,6 @@ def test_mnist_backward(self): Tensor.realize(*[p.replace(Tensor.ones_like(p).contiguous()) for p in nn.state.get_parameters(layers)]) # realize gradients - for x in nn.state.get_parameters(layers): x.requires_grad_() Tensor.empty(4, 1, 14, 14).sequential(layers).sum().backward() Tensor.realize(*[x.grad for x in nn.state.get_parameters(layers) if x.grad is not None]) diff --git a/test/testextra/test_bench_log.py b/test/testextra/test_bench_log.py index 1934227d8e922..975bef2fe1833 100644 --- a/test/testextra/test_bench_log.py +++ b/test/testextra/test_bench_log.py @@ -2,11 +2,12 @@ from unittest.case import skipIf from extra.bench_log import BenchEvent, InstantBenchEvent, WallTimeEvent, KernelTimeEvent, log_event_instant, _events, clear_events -from tinygrad.helpers import Context, CI +from tinygrad.helpers import Context, DEV from tinygrad.tensor import Tensor from tinygrad.device import Device -_SKIP_KERNEL_TIMING = Device.DEFAULT == "WEBGPU" # WEBGPU kernel timing not supported +# WEBGPU kernel timing not supported, ocelot CUDA is inaccurate +_SKIP_KERNEL_TIMING = Device.DEFAULT == "WEBGPU" or (Device.DEFAULT == "CUDA" and DEV.interface.startswith("MOCK")) class TestBenchLog(unittest.TestCase): def setUp(self): @@ -37,7 +38,7 @@ def test_log_double_wall_time(self): self.assertGreater(_events[event]["wall"][0], 0) self.assertGreater(_events[event]["wall"][1], 0) - @skipIf(CI or _SKIP_KERNEL_TIMING, "ci timing is not accurate") + @skipIf(_SKIP_KERNEL_TIMING, "ci timing is not accurate") def test_log_single_kernel_time(self): wall_times = [] @@ -54,7 +55,7 @@ def test_log_single_kernel_time(self): self.assertLess(_events[event]["kernel"][0], wall_times[0]) self.assertGreater(_events[event]["kernel"][0], 0) - @skipIf((CI and Device.DEFAULT == "CUDA") or _SKIP_KERNEL_TIMING, "ci cuda timing is not accurate") + @skipIf(_SKIP_KERNEL_TIMING, "ci cuda timing is not accurate") def test_interleaved_wall_kernel_time(self): wall_times = [] with Context(DEBUG=2): @@ -76,7 +77,7 @@ def test_interleaved_wall_kernel_time(self): self.assertLess(_events[event]["kernel"][0], wall_times[0]) self.assertGreater(_events[event]["kernel"][0], 0) - @skipIf((CI and Device.DEFAULT == "CUDA") or _SKIP_KERNEL_TIMING, "ci cuda timing is not accurate") + @skipIf(_SKIP_KERNEL_TIMING, "ci cuda timing is not accurate") def test_stacked_wall_kernel_time(self): with Context(DEBUG=2): for event in BenchEvent: diff --git a/test/testextra/test_f16_decompress.py b/test/testextra/test_f16_decompress.py index 68068770343f1..3f3487b3b0a9f 100644 --- a/test/testextra/test_f16_decompress.py +++ b/test/testextra/test_f16_decompress.py @@ -1,12 +1,10 @@ import unittest from extra.f16_decompress import u32_to_f16 from tinygrad.tensor import Tensor -from tinygrad.device import is_dtype_supported from tinygrad import dtypes import numpy as np class TestF16Decompression(unittest.TestCase): - @unittest.skipUnless(is_dtype_supported(dtypes.float16), "need float16") def test_u32_to_f16(self): a = Tensor.randn(50, dtype=dtypes.float16) f16_as_u32 = a.bitcast(dtypes.uint32) diff --git a/test/testextra/test_fp8_linear.py b/test/testextra/test_fp8_linear.py index 34522c82c98db..50abeab076bda 100644 --- a/test/testextra/test_fp8_linear.py +++ b/test/testextra/test_fp8_linear.py @@ -4,12 +4,11 @@ from tinygrad import Tensor, dtypes, Device from tinygrad.nn import Linear from extra.fp8.fp8_linear import FP8Linear, convert_to_float8_training -from tinygrad.device import is_dtype_supported from test.helpers import not_support_multi_device, needs_second_gpu BS, T, in_dim, out_dim = 16, 4, 128, 128 -@unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3), f"no fp8e4m3 on {Device.DEFAULT}") +@unittest.skipUnless(dtypes.fp8e4m3 in Device[Device.DEFAULT].renderer.supported_dtypes(), f"no fp8e4m3 on {Device.DEFAULT}") class TestFP8Linear(unittest.TestCase): def setUp(self): Tensor.manual_seed(42) @@ -34,9 +33,8 @@ def _test_backward(self, shape, in_features, out_features): bias = Tensor.randn(out_features, dtype=dtypes.float32) * 0.2 fp8_layer.weight, normal_layer.weight = weight.detach(), weight.detach() fp8_layer.bias, normal_layer.bias = bias.detach(), bias.detach() - fp8_layer.weight.requires_grad = normal_layer.weight.requires_grad = True - x_fp8 = Tensor.randn(*shape, dtype=dtypes.float32, requires_grad=True) * 0.2 - x_normal = x_fp8.detach().requires_grad_(True) + x_fp8 = Tensor.randn(*shape, dtype=dtypes.float32) * 0.2 + x_normal = x_fp8.detach() fp8_layer(x_fp8).sum().backward() normal_layer(x_normal).sum().backward() np.testing.assert_allclose(x_fp8.grad.numpy(), x_normal.grad.numpy(), rtol=1.0, atol=0.1) diff --git a/test/testextra/test_hk_fa.py b/test/testextra/test_hk_fa.py index 8481de1698a34..58f676664b55f 100644 --- a/test/testextra/test_hk_fa.py +++ b/test/testextra/test_hk_fa.py @@ -49,9 +49,9 @@ def test_fast_fa_bwd_causal(self): B, N, H, H_KV, D = 1, 8192, 32, 8, 128 with Context(DEBUG=0): - q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous() - k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous() - v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous() + q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16).contiguous() + k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous() + v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous() Tensor.realize(q, k, v) do = Tensor.ones(B, N, H, D, dtype=dtypes.float32).contiguous() @@ -64,9 +64,9 @@ def test_fast_fa_bwd_causal(self): Tensor.realize(q.grad, k.grad, v.grad) with Context(DEBUG=0): - q_ref = q.detach().clone().requires_grad_(True) - k_ref = k.detach().clone().requires_grad_(True) - v_ref = v.detach().clone().requires_grad_(True) + q_ref = q.detach().clone() + k_ref = k.detach().clone() + v_ref = v.detach().clone() Tensor.realize(q_ref, k_ref, v_ref) q_ref_, k_ref_, v_ref_ = q_ref.transpose(1, 2), k_ref.transpose(1, 2), v_ref.transpose(1, 2) @@ -85,9 +85,9 @@ def test_fast_fa_bwd_causal_jitted(self): B, N, H, H_KV, D = 1, 8192, 32, 8, 128 with Context(DEBUG=0): - q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous() - k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous() - v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous() + q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16).contiguous() + k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous() + v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous() Tensor.realize(q, k, v) do = Tensor.ones(B, N, H, D, dtype=dtypes.float32).contiguous() @@ -104,18 +104,18 @@ def fn(q, k, v, do): fn_jitted = TinyJit(fn) for _ in range(10): - q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous() - k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous() - v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous() + q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16).contiguous() + k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous() + v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous() Tensor.realize(q, k, v) do = Tensor.ones(B, N, H, D, dtype=dtypes.float32).contiguous() Tensor.realize(do) q.grad, k.grad, v.grad = fn_jitted(q, k, v, do) with Context(DEBUG=0): - q_ref = q.detach().clone().requires_grad_(True) - k_ref = k.detach().clone().requires_grad_(True) - v_ref = v.detach().clone().requires_grad_(True) + q_ref = q.detach().clone() + k_ref = k.detach().clone() + v_ref = v.detach().clone() Tensor.realize(q_ref, k_ref, v_ref) q_ref_, k_ref_, v_ref_ = q_ref.transpose(1, 2), k_ref.transpose(1, 2), v_ref.transpose(1, 2) @@ -135,16 +135,16 @@ def test_fast_fa_bwd_dp(self): GPUS = tuple(f"AMD:{i}" for i in range(B)) with Context(DEBUG=0): - base_q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous() - base_k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous() - base_v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous() + base_q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16).contiguous() + base_k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous() + base_v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous() base_do = Tensor.ones(B, N, H, D, dtype=dtypes.float32).contiguous() with Context(DEBUG=0): - q = base_q.clone().requires_grad_(True).shard(GPUS, axis=0) - k = base_k.clone().requires_grad_(True).shard(GPUS, axis=0) - v = base_v.clone().requires_grad_(True).shard(GPUS, axis=0) + q = base_q.clone().shard(GPUS, axis=0) + k = base_k.clone().shard(GPUS, axis=0) + v = base_v.clone().shard(GPUS, axis=0) Tensor.realize(q, k, v) do = base_do.clone().shard(GPUS, axis=0) @@ -157,9 +157,9 @@ def test_fast_fa_bwd_dp(self): Tensor.realize(q.grad, k.grad, v.grad) with Context(DEBUG=0): - q_ref = base_q.clone().requires_grad_(True) - k_ref = base_k.clone().requires_grad_(True) - v_ref = base_v.clone().requires_grad_(True) + q_ref = base_q.clone() + k_ref = base_k.clone() + v_ref = base_v.clone() Tensor.realize(q_ref, k_ref, v_ref) do_ref = base_do.clone() @@ -182,16 +182,16 @@ def test_fast_fa_bwd_mp(self): GPUS = tuple(f"AMD:{i}" for i in range(B)) with Context(DEBUG=0): - base_q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous() - base_k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous() - base_v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous() + base_q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16).contiguous() + base_k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous() + base_v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous() base_do = Tensor.ones(B, N, H, D, dtype=dtypes.float32).contiguous() with Context(DEBUG=0): - q = base_q.clone().requires_grad_(True).shard(GPUS, axis=2) - k = base_k.clone().requires_grad_(True).shard(GPUS, axis=2) - v = base_v.clone().requires_grad_(True).shard(GPUS, axis=2) + q = base_q.clone().shard(GPUS, axis=2) + k = base_k.clone().shard(GPUS, axis=2) + v = base_v.clone().shard(GPUS, axis=2) Tensor.realize(q, k, v) do = base_do.clone().shard(GPUS, axis=2) @@ -204,9 +204,9 @@ def test_fast_fa_bwd_mp(self): Tensor.realize(q.grad, k.grad, v.grad) with Context(DEBUG=0): - q_ref = base_q.clone().requires_grad_(True) - k_ref = base_k.clone().requires_grad_(True) - v_ref = base_v.clone().requires_grad_(True) + q_ref = base_q.clone() + k_ref = base_k.clone() + v_ref = base_v.clone() Tensor.realize(q_ref, k_ref, v_ref) do_ref = base_do.clone() diff --git a/test/testextra/test_lr_scheduler.py b/test/testextra/test_lr_scheduler.py index ce7464ecf67ab..aaa916683afcd 100644 --- a/test/testextra/test_lr_scheduler.py +++ b/test/testextra/test_lr_scheduler.py @@ -61,7 +61,7 @@ def tearDown(self): def _test_lr_scheduler(self, tinygrad_sched, torch_sched, epochs, opts, atol, rtol, adam=True): accs = opts.pop('accs', None) - test_tensor = Tensor([0.], requires_grad=True) # NOTE: optimizers are broken on 0-dim tensors because it broadcasts to [lr] + test_tensor = Tensor([0.]) # NOTE: optimizers are broken on 0-dim tensors because it broadcasts to [lr] test_tensor.mean().backward() if adam: tinygrad_optim, torch_optim = Adam([test_tensor], lr=0.01), torch.optim.Adam([torch.tensor([0.], requires_grad=True)], lr=0.01) diff --git a/test/testextra/test_tk.py b/test/testextra/test_tk.py index ac387facc24b7..a572a9b395252 100644 --- a/test/testextra/test_tk.py +++ b/test/testextra/test_tk.py @@ -816,9 +816,9 @@ def test_fast_fa_bwd(self): B, N, H, H_KV, D = 1, 32, 2, 1, 32 with Context(DEBUG=0): - q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous() - k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous() - v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous() + q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16).contiguous() + k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous() + v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous() Tensor.realize(q, k, v) do = Tensor.ones(B, N, H, D, dtype=dtypes.float32).contiguous() @@ -831,9 +831,9 @@ def test_fast_fa_bwd(self): Tensor.realize(q.grad, k.grad, v.grad) with Context(DEBUG=0): - q_ref = q.detach().clone().requires_grad_(True) - k_ref = k.detach().clone().requires_grad_(True) - v_ref = v.detach().clone().requires_grad_(True) + q_ref = q.detach().clone() + k_ref = k.detach().clone() + v_ref = v.detach().clone() Tensor.realize(q_ref, k_ref, v_ref) q_ref_, k_ref_, v_ref_ = q_ref.transpose(1, 2), k_ref.transpose(1, 2), v_ref.transpose(1, 2) @@ -854,9 +854,9 @@ def test_fast_fa_bwd_causal(self): B, N, H, H_KV, D = 1, 8192, 32, 32, 128 with Context(DEBUG=0): - q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous() - k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous() - v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous() + q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16).contiguous() + k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous() + v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous() Tensor.realize(q, k, v) do = Tensor.ones(B, N, H, D, dtype=dtypes.float32).contiguous() @@ -869,9 +869,9 @@ def test_fast_fa_bwd_causal(self): Tensor.realize(q.grad, k.grad, v.grad) with Context(DEBUG=0): - q_ref = q.detach().clone().requires_grad_(True) - k_ref = k.detach().clone().requires_grad_(True) - v_ref = v.detach().clone().requires_grad_(True) + q_ref = q.detach().clone() + k_ref = k.detach().clone() + v_ref = v.detach().clone() Tensor.realize(q_ref, k_ref, v_ref) q_ref_, k_ref_, v_ref_ = q_ref.transpose(1, 2), k_ref.transpose(1, 2), v_ref.transpose(1, 2) @@ -892,9 +892,9 @@ def test_fast_fa_bwd_causal_jitted(self): B, N, H, H_KV, D = 1, 8192, 32, 32, 128 with Context(DEBUG=0): - q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous() - k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous() - v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous() + q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16).contiguous() + k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous() + v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous() Tensor.realize(q, k, v) do = Tensor.ones(B, N, H, D, dtype=dtypes.float32).contiguous() @@ -911,18 +911,18 @@ def fn(q, k, v, do): fn_jitted = TinyJit(fn) for _ in range(10): - q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous() - k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous() - v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous() + q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16).contiguous() + k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous() + v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous() Tensor.realize(q, k, v) do = Tensor.ones(B, N, H, D, dtype=dtypes.float32).contiguous() Tensor.realize(do) q.grad, k.grad, v.grad = fn_jitted(q, k, v, do) with Context(DEBUG=0): - q_ref = q.detach().clone().requires_grad_(True) - k_ref = k.detach().clone().requires_grad_(True) - v_ref = v.detach().clone().requires_grad_(True) + q_ref = q.detach().clone() + k_ref = k.detach().clone() + v_ref = v.detach().clone() Tensor.realize(q_ref, k_ref, v_ref) q_ref_, k_ref_, v_ref_ = q_ref.transpose(1, 2), k_ref.transpose(1, 2), v_ref.transpose(1, 2) @@ -944,16 +944,16 @@ def test_fast_fa_bwd_multidevice(self): GPUS = tuple(f"{Device.DEFAULT}:{i}" for i in range(B)) with Context(DEBUG=0): - base_q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous() - base_k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous() - base_v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous() + base_q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16).contiguous() + base_k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous() + base_v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous() base_do = Tensor.ones(B, N, H, D, dtype=dtypes.float32).contiguous() with Context(DEBUG=0): - q = base_q.clone().requires_grad_(True).shard(GPUS, axis=0) - k = base_k.clone().requires_grad_(True).shard(GPUS, axis=0) - v = base_v.clone().requires_grad_(True).shard(GPUS, axis=0) + q = base_q.clone().shard(GPUS, axis=0) + k = base_k.clone().shard(GPUS, axis=0) + v = base_v.clone().shard(GPUS, axis=0) Tensor.realize(q, k, v) do = base_do.clone().shard(GPUS, axis=0) @@ -966,9 +966,9 @@ def test_fast_fa_bwd_multidevice(self): Tensor.realize(q.grad, k.grad, v.grad) with Context(DEBUG=0): - q_ref = base_q.clone().requires_grad_(True) - k_ref = base_k.clone().requires_grad_(True) - v_ref = base_v.clone().requires_grad_(True) + q_ref = base_q.clone() + k_ref = base_k.clone() + v_ref = base_v.clone() Tensor.realize(q_ref, k_ref, v_ref) do_ref = base_do.clone() diff --git a/test/unit/test_assign.py b/test/unit/test_assign.py index d7434489ee98a..955ab56c8b354 100644 --- a/test/unit/test_assign.py +++ b/test/unit/test_assign.py @@ -2,9 +2,8 @@ import unittest import numpy as np from tinygrad import dtypes, Tensor, TinyJit, GlobalCounters, Variable -from tinygrad.uop.ops import Ops -from tinygrad.device import is_dtype_supported -from tinygrad.helpers import temp, CI, DEV, Context +from tinygrad.uop.ops import Ops, UOp +from tinygrad.helpers import temp, DEV, Context N = 200 # has to be bigger than the cache to fail @@ -140,7 +139,7 @@ def test_assign_changes_alt(self, realize=False): def test_assign_changes_realized_alt(self): return self.test_assign_changes_alt(realize=True) def test_assign_changes_buffer_alt(self): - a, b = [Tensor(Tensor(0).contiguous().realize().uop.buf_uop) for _ in range(2)] + a, b = [Tensor(Tensor([0]).realize().uop.buf_uop) for _ in range(2)] Tensor.realize(a.contiguous().assign(1), b.contiguous().assign(2)) self.assertEqual((a + b).item(), 3) @@ -189,7 +188,7 @@ def test_assign_diamond_alt(self): new = a + times_a np.testing.assert_allclose(new.numpy(), 8) - @unittest.skipIf(CI and DEV.renderer == "LVP", "flaky in CI") + @unittest.skipIf(DEV.renderer == "LVP", "flaky in CI") def test_double_assign(self): a = Tensor.ones(4).contiguous().realize() a += 1 @@ -282,15 +281,15 @@ def test_assign_after_target_chain(self): np.testing.assert_equal(t.numpy(), [[100, 104, 108, 112], [101, 105, 109, 113], [102, 106, 110, 114], [103, 107, 111, 115]]) def test_assign_contiguous(self): - b = Tensor.arange(16).reshape(4,4).contiguous().realize() - a = (Tensor.arange(16).reshape(4,4).contiguous().realize() + 1) + b = Tensor.arange(16).reshape(4,4).clone().realize() + a = (Tensor.arange(16).reshape(4,4).clone().realize() + 1) GlobalCounters.reset() b.assign(a.contiguous()).realize() self.assertEqual(GlobalCounters.kernel_count, 2) def test_assign_contiguous_permute(self): - b = Tensor.arange(16).reshape(4,4).contiguous().realize() - a = (Tensor.arange(16).reshape(4,4).contiguous().realize() + 1).permute((1,0)) + b = Tensor.arange(16).reshape(4,4).clone().realize() + a = (Tensor.arange(16).reshape(4,4).clone().realize() + 1).permute((1,0)) GlobalCounters.reset() b.assign(a.contiguous()).realize() self.assertEqual(GlobalCounters.kernel_count, 2) @@ -326,29 +325,29 @@ def test_post_permuted_assignment(self): np.testing.assert_allclose(a.numpy(), np.arange(N*N).reshape((N,N)) + np.arange(N*N).reshape((N,N)).transpose(1,0)) def test_post_permuted_assignment_alt(self): - a = Tensor.arange(N*N).reshape(N,N).contiguous().realize() - b = Tensor.arange(N*N).reshape(N,N).contiguous().realize() + a = Tensor.arange(N*N).reshape(N,N).clone().realize() + b = Tensor.arange(N*N).reshape(N,N).clone().realize() new_a = (a.T+b).numpy() a.assign(a.T+b) np.testing.assert_allclose(a.numpy(), new_a) def test_post_flipped_assignment(self): - a = Tensor.arange(N*N).reshape(N,N).contiguous().realize() - b = Tensor.arange(N*N).reshape(N,N).contiguous().realize() + a = Tensor.arange(N*N).reshape(N,N).clone().realize() + b = Tensor.arange(N*N).reshape(N,N).clone().realize() new_a = (a.flip(0)+b).numpy() a.assign(a.flip(0)+b) np.testing.assert_allclose(a.numpy(), new_a) def test_post_flipped_assignment_axis1(self): - a = Tensor.arange(N*N).reshape(N,N).contiguous().realize() - b = Tensor.arange(N*N).reshape(N,N).contiguous().realize() + a = Tensor.arange(N*N).reshape(N,N).clone().realize() + b = Tensor.arange(N*N).reshape(N,N).clone().realize() new_a = (a.flip(1)+b).numpy() a.assign(a.flip(1)+b) np.testing.assert_allclose(a.numpy(), new_a) def test_post_reshape_assignment_fine(self): - a = Tensor.arange(N*N).reshape(N, N).contiguous().realize() - b = Tensor.arange(N*N).reshape(N, N).contiguous().realize() + a = Tensor.arange(N*N).reshape(N, N).clone().realize() + b = Tensor.arange(N*N).reshape(N, N).clone().realize() rhs = a.reshape(-1).reshape(N, N) new_a = (rhs+b).numpy() a.assign(rhs+b) # self-assign with reshape view is fine @@ -356,7 +355,7 @@ def test_post_reshape_assignment_fine(self): @unittest.skip("multi output not supported anymore") def test_simple_assignment_multioutput(self): - a = Tensor.arange(32*32).reshape(32, 32).contiguous().realize() + a = Tensor.arange(32*32).reshape(32, 32).clone().realize() b = Tensor.full((32, ), 1.).contiguous().realize() c = Tensor.full((32, ), 2.).contiguous().realize() d = Tensor.full((32, ), 3.).contiguous().realize() @@ -376,15 +375,15 @@ def test_simple_assignment_multioutput(self): # NOTE: if the assign target is read/write in a single kernel, it should be contiguous def test_permuted_assignment_correct(self): - a = Tensor.arange(4 * 4).reshape(4, 4).contiguous().realize() - b = Tensor.arange(4 * 4).reshape(4, 4).contiguous().realize() + a = Tensor.arange(4 * 4).reshape(4, 4).clone().realize() + b = Tensor.arange(4 * 4).reshape(4, 4).clone().realize() a = a.permute(1, 0) new_val = a + b a.assign(new_val) np.testing.assert_equal(a.numpy(), np.arange(4 * 4).reshape(4, 4).transpose(1, 0) + np.arange(4 * 4).reshape(4, 4)) def test_permuted_reduceop_child_dual_use(self): - a = Tensor.arange(32*32*32).reshape(32, 32, 32).contiguous().realize() + a = Tensor.arange(32*32*32).reshape(32, 32, 32).clone().realize() b = Tensor.ones(32, 32, dtype=dtypes.int).contiguous().realize() r = a.sum(axis=1) b.assign(r + b.permute(1, 0)) @@ -393,7 +392,7 @@ def test_permuted_reduceop_child_dual_use(self): @unittest.skip("multi output not supported anymore") def test_permuted_reduceop_multioutput_dual_use(self): - a = Tensor.arange(32*32*32).reshape(32, 32, 32).contiguous().realize() + a = Tensor.arange(32*32*32).reshape(32, 32, 32).clone().realize() b = Tensor.full((32, 32), 1.).contiguous().realize() c = Tensor.full((32, 32), 2.).contiguous().realize() @@ -406,9 +405,9 @@ def test_permuted_reduceop_multioutput_dual_use(self): @unittest.skip("multi output not supported anymore") def test_permuted_reduceop_multioutput_dual_use_possible(self): - a = Tensor.arange(32*32*32).reshape(32, 32, 32).contiguous().realize() - b = Tensor.arange(32 * 32).reshape(32, 32).realize() - c = Tensor.arange(32 * 32).reshape(32, 32).realize() + a = Tensor.arange(32*32*32).reshape(32, 32, 32).clone().realize() + b = Tensor.arange(32 * 32).reshape(32, 32).clone().realize() + c = Tensor.arange(32 * 32).reshape(32, 32).clone().realize() GlobalCounters.reset() r = a.sum(axis=1) @@ -442,7 +441,7 @@ def test_overlapping_shrink_assignment_forward(self): # Forward shift: read index > write index in overlap N = 100000 shift = 1000 - a = Tensor.arange(N).float().contiguous().realize() + a = Tensor.arange(N).float().clone().realize() expected = np.arange(N, dtype=np.float32) expected[:N-shift] = expected[shift:].copy() with Context(NOOPT=1): a[0:N-shift].assign(a[shift:N]).realize() @@ -452,7 +451,7 @@ def test_overlapping_shrink_assignment_reverse(self): # Reverse shift: write index > read index in overlap N = 100000 shift = 1000 - a = Tensor.arange(N).float().contiguous().realize() + a = Tensor.arange(N).float().clone().realize() expected = np.arange(N, dtype=np.float32) expected[shift:] = expected[:N-shift].copy() with Context(NOOPT=1): a[shift:N].assign(a[0:N-shift]).realize() @@ -460,7 +459,7 @@ def test_overlapping_shrink_assignment_reverse(self): def test_nonoverlapping_shrink_assignment(self): # TODO: non-overlapping shrinks don't actually need contiguous, could be 1 kernel with smarter range analysis - a = Tensor.arange(100).float().contiguous().realize() + a = Tensor.arange(100).float().clone().realize() expected = np.arange(100, dtype=np.float32) expected[0:10] = expected[50:60].copy() GlobalCounters.reset() @@ -468,7 +467,6 @@ def test_nonoverlapping_shrink_assignment(self): self.assertEqual(GlobalCounters.kernel_count, 2) # currently conservative, forces contiguous np.testing.assert_allclose(a.numpy(), expected) - @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half") def test_setitem_half(self): a = Tensor.full((8,), 1.0, dtype=dtypes.half).contiguous().realize() b = Tensor.full((4,), 2.0, dtype=dtypes.half).contiguous().realize() @@ -578,7 +576,7 @@ def test_chained_assign_kernel_count(self): """Chained pending assigns must not produce excessive kernels (tests recursive transitive processing).""" D, N = 4, 5 caches = [Tensor.zeros(8, D).contiguous().realize() for _ in range(N)] - caches[0][0:1].assign(Tensor.ones(1, D) * 10) + caches[0][0:1].assign(Tensor.ones(1, D, buffer=False) * 10) x = caches[0][:1].sum(0, keepdim=True) for i in range(1, N): caches[i][0:1].assign(x) @@ -608,13 +606,18 @@ def test_shared_computation_assign_kernel_count(self): def test_double_assign_from_const(self): a = Tensor.empty(2) - a.assign(Tensor.ones(2)) - a.assign(Tensor.ones(2)) + a.assign(Tensor.ones(2, buffer=False)) + a.assign(Tensor.ones(2, buffer=False)) GlobalCounters.reset() a.realize() self.assertEqual(GlobalCounters.kernel_count, 1) self.assertEqual(a.tolist(), [1.,1.]) + def test_assign_deviceless_const(self): + s = Tensor.empty(4, device="CPU:1", dtype=dtypes.float) + s.assign(Tensor(UOp.const(dtypes.float, 2.0))) + np.testing.assert_equal(s.numpy(), [2, 2, 2, 2]) + def test_nested_after_contiguous_store(self): # Mirrors the nested contiguous-write-then-assign-back shape from torch backend view updates. base = Tensor.empty(3, dtype=dtypes.int64) diff --git a/test/unit/test_call.py b/test/unit/test_call.py index 1cecd3bb3829d..4bda7003f70aa 100644 --- a/test/unit/test_call.py +++ b/test/unit/test_call.py @@ -17,8 +17,8 @@ def test_call_plus(self): np.testing.assert_equal(c.numpy(), (a+b).numpy()) def test_call_plus_backward(self): - a = Tensor.ones(10, 10, requires_grad=True) - b = Tensor.ones(10, 10, requires_grad=True) + a = Tensor.ones(10, 10) + b = Tensor.ones(10, 10) (a+b).mean().backward() gt_a_grad = a.grad.numpy() @@ -37,8 +37,8 @@ def grad_fxn(grad:UOp, call:UOp): return (grad, grad) np.testing.assert_allclose(b.grad.numpy(), gt_b_grad, rtol=1e-5) def test_call_plus_backward_auto(self): - a = Tensor.ones(10, 10, requires_grad=True) - b = Tensor.ones(10, 10, requires_grad=True) + a = Tensor.ones(10, 10) + b = Tensor.ones(10, 10) (a+b).mean().backward() gt_a_grad = a.grad.numpy() @@ -75,8 +75,8 @@ def test_call_gemm_uop(self): def test_call_complex_backward_auto(self): # complex chain: (a*b + a).exp2() * b.reciprocal() - tests mul, add, exp2, reciprocal, param reuse - a = Tensor.randn(10, 10, requires_grad=True) - b = Tensor.randn(10, 10, requires_grad=True) + 2 # avoid div by zero + a = Tensor.randn(10, 10) + b = Tensor.randn(10, 10) + 2 # avoid div by zero Tensor.realize(a, b) ((a*b + a).exp2() * b.reciprocal()).mean().backward() @@ -162,9 +162,9 @@ def s(x): return x*2 s(s(a).contiguous()).realize() def test_call_double_gemm(self): - a = Tensor.randn(4, 8, requires_grad=True) - b = Tensor.randn(8, 12, requires_grad=True) - c = Tensor.randn(12, 16, requires_grad=True) + a = Tensor.randn(4, 8) + b = Tensor.randn(8, 12) + c = Tensor.randn(12, 16) ref = Tensor.randn(4, 16) Tensor.realize(a,b,c,ref) @function(precompile=True) @@ -230,7 +230,7 @@ def test_precompile_symbolic_2d(self): @function(precompile=True) def f(x:Tensor) -> Tensor: return x * 2 + 1 sz = UOp.variable("sz", 1, 16) - a = Tensor.arange(16*4).reshape(16, 4).float()[:sz.bind(5)] + a = Tensor.arange(16*4).reshape(16, 4).float().clone()[:sz.bind(5)] out = f(a) # result shape should have the symbolic dim, not the max self.assertIsInstance(out.shape[0], UOp) @@ -240,7 +240,7 @@ def test_precompile_multi_sharded(self): @function(precompile=True) def f(x:Tensor) -> Tensor: return x + 1 devs = ("CPU:0", "CPU:1") - a = Tensor.arange(8).reshape(4, 2).float().shard(devs, axis=0) + a = Tensor.arange(8).reshape(4, 2).float().clone().shard(devs, axis=0) out = f(a) + 2 np.testing.assert_allclose(out.numpy(), np.arange(8, dtype=np.float32).reshape(4, 2) + 3) @@ -251,7 +251,7 @@ def test_tuple_sharded(self): devs = ("CPU:0", "CPU:1") @function def f(x:Tensor): return (x + 1, x * 2) - a = Tensor.arange(8).reshape(4, 2).float().shard(devs, axis=0) + a = Tensor.arange(8).reshape(4, 2).float().clone().shard(devs, axis=0) t1, t2 = f(a) ref = np.arange(8, dtype=np.float32).reshape(4, 2) np.testing.assert_allclose(t1.numpy(), ref + 1) @@ -262,7 +262,7 @@ def test_tuple_sharded_precompile(self): devs = ("CPU:0", "CPU:1") @function(precompile=True) def f(x:Tensor): return (x + 1, x * 2) - a = Tensor.arange(8).reshape(4, 2).float().shard(devs, axis=0) + a = Tensor.arange(8).reshape(4, 2).float().clone().shard(devs, axis=0) t1, t2 = f(a) ref = np.arange(8, dtype=np.float32).reshape(4, 2) np.testing.assert_allclose(t1.numpy(), ref + 1) @@ -273,7 +273,7 @@ def test_tuple_sharded_different_axis(self): devs = ("CPU:0", "CPU:1") @function def f(x:Tensor): return (x.sum(axis=0), x.sum(axis=1)) - a = Tensor.arange(8).reshape(4, 2).float().shard(devs, axis=0) + a = Tensor.arange(8).reshape(4, 2).float().clone().shard(devs, axis=0) t1, t2 = f(a) ref = np.arange(8, dtype=np.float32).reshape(4, 2) np.testing.assert_allclose(t1.numpy(), ref.sum(axis=0)) @@ -284,8 +284,8 @@ def test_tuple_sharded_different_ops(self): devs = ("CPU:0", "CPU:1") @function def f(x:Tensor, y:Tensor): return (x + y, x * y) - a = Tensor.arange(8).reshape(4, 2).float().shard(devs, axis=0) - b = Tensor.arange(8).reshape(4, 2).float().shard(devs, axis=0) + 1 + a = Tensor.arange(8).reshape(4, 2).float().clone().shard(devs, axis=0) + b = Tensor.arange(8).reshape(4, 2).float().clone().shard(devs, axis=0) + 1 t1, t2 = f(a, b) ref_a = np.arange(8, dtype=np.float32).reshape(4, 2) ref_b = ref_a + 1 @@ -297,7 +297,7 @@ def test_tuple_sharded_mixed_use(self): devs = ("CPU:0", "CPU:1") @function def f(x:Tensor): return (x + 1, x * 2) - a = Tensor.arange(8).reshape(4, 2).float().shard(devs, axis=0) + a = Tensor.arange(8).reshape(4, 2).float().clone().shard(devs, axis=0) t1, t2 = f(a) out = (t1 + t2).sum() ref = np.arange(8, dtype=np.float32).reshape(4, 2) @@ -308,8 +308,8 @@ def test_tuple_sharded_outputs_different_axis(self): devs = ("CPU:0", "CPU:1") @function def f(x:Tensor, y:Tensor): return (x + 1, y + 2) - a = Tensor.arange(8).reshape(4, 2).float().shard(devs, axis=0) - b = Tensor.arange(8).reshape(4, 2).float().shard(devs, axis=1) + a = Tensor.arange(8).reshape(4, 2).float().clone().shard(devs, axis=0) + b = Tensor.arange(8).reshape(4, 2).float().clone().shard(devs, axis=1) t1, t2 = f(a, b) ref_a = np.arange(8, dtype=np.float32).reshape(4, 2) ref_b = np.arange(8, dtype=np.float32).reshape(4, 2) @@ -333,8 +333,8 @@ def test_call_reduce_sharded_mixed_args(self): def test_call_reduce_sharded_backward(self): devs = ("CPU:0", "CPU:1") - a = Tensor.randn(10, 10, requires_grad=True).shard(devs, axis=0) - b = Tensor.randn(10, 10, requires_grad=True).shard(devs, axis=0) + a = Tensor.randn(10, 10).shard(devs, axis=0) + b = Tensor.randn(10, 10).shard(devs, axis=0) Tensor.realize(a, b) def grad_fxn(grad, call): diff --git a/test/unit/test_cpu.py b/test/unit/test_cpu.py index 2921bb9c93d71..7aaef1730c34d 100644 --- a/test/unit/test_cpu.py +++ b/test/unit/test_cpu.py @@ -3,6 +3,7 @@ from tinygrad import Tensor, Device from tinygrad.helpers import Target from tinygrad.renderer.nir import LVPRenderer +from tinygrad.renderer.isa.x86 import X86Renderer from tinygrad.codegen import to_program @unittest.skipIf(Device.DEFAULT != "CPU", "only run on CPU") @@ -12,6 +13,7 @@ def test_arch_feats(self): for ren in Device[Device.DEFAULT].renderers: for arch, expect_vmov in [("x86_64,x86-64,avx", True), ("x86_64,x86-64,-avx", False)]: with self.subTest(arch=arch): + if ren is X86Renderer: continue # X86 requires avx support if ren is LVPRenderer: continue # LVP does not play nice with cross compilation r = ren(Target(device="CPU", arch=arch)) p = to_program(ast, r) diff --git a/test/unit/test_disk_tensor.py b/test/unit/test_disk_tensor.py index 3db019a1384b9..85dcf6f9aa88c 100644 --- a/test/unit/test_disk_tensor.py +++ b/test/unit/test_disk_tensor.py @@ -1,7 +1,6 @@ import os, pathlib, tempfile, unittest import numpy as np from tinygrad import Tensor, Device, dtypes -from tinygrad.device import is_dtype_supported from tinygrad.dtype import DType, DTYPES_DICT from tinygrad.nn.state import safe_load, safe_save, get_state_dict, torch_load from tinygrad.helpers import Timing, fetch, OSX, dedup @@ -36,7 +35,6 @@ def test_load_enet_alt(self): compare_weights_both("https://download.pytorch.org # pytorch zip format def test_load_convnext(self): compare_weights_both('https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth') - @unittest.skipUnless(is_dtype_supported(dtypes.float16), "need float16 support") def test_load_llama2bfloat(self): compare_weights_both("https://huggingface.co/qazalin/bf16-lightweight/resolve/main/consolidated.00.pth?download=true") # pytorch tar format @@ -94,7 +92,6 @@ def test_bitcasts_on_disk(self): pathlib.Path(tmp).unlink() -@unittest.skipUnless(is_dtype_supported(dtypes.uint8), "need uint8") class TestSafetensors(TempDirTestCase): def test_real_safetensors(self): import torch @@ -184,7 +181,6 @@ def test_safe_save_only_copy(self): def test_save_all_dtypes(self): for dtype in dedup(DTYPES_DICT.values()): if dtype in [dtypes.bfloat16]: continue # not supported in numpy - if not is_dtype_supported(dtype): continue path = self.tmp(f"ones.{dtype}.safetensors") ones = Tensor(np.random.rand(10,10), dtype=dtype) safe_save(get_state_dict(ones), path) @@ -306,7 +302,7 @@ def test_strided_setitem(self): dt[::2] = Tensor([10, 20, 30]) def test_advanced_setitem_not_supported(self): - dt = Tensor.arange(12).reshape(3, 4).to(f"disk:{self.tmp('dt_advanced_setitem')}") + dt = Tensor.arange(12).reshape(3, 4).clone().to(f"disk:{self.tmp('dt_advanced_setitem')}") with self.assertRaises(RuntimeError, msg="advanced setitem is not supported for DISK tensors"): dt[Tensor([0, 2]), Tensor([1, 3])] = 99 @@ -384,7 +380,6 @@ def test_bitcast_view(self): assert ret.tolist() == [2827, 3341, 3855, 4369] @unittest.skipIf(OSX or Device.DEFAULT == "CL", "new LLVM has an issue on OSX, DEV=CL gives the wrong output") - @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), "bfloat16 not supported") def test_bf16_disk_write_read(self): t = Tensor([10000, -1, -1000, -10000, 20], dtype=dtypes.float32) t.to(f"disk:{self.tmp('dt_bf16_disk_write_read_f32')}").realize() @@ -549,7 +544,7 @@ class TestDiskTensorMovement(TempDirTestCase): def setUp(self): super().setUp() self.fn = pathlib.Path(self.tmp("custom_disk_range")) - Tensor.arange(100, dtype=dtypes.uint8).to(f"disk:{str(self.fn)}").realize() + Tensor.arange(100, dtype=dtypes.uint8).clone().to(f"disk:{str(self.fn)}").realize() def test_simple_read(self): t = Tensor(self.fn) diff --git a/test/unit/test_dtype_spec.py b/test/unit/test_dtype_spec.py index 477faf22bb2f8..c74e8e1583c5f 100644 --- a/test/unit/test_dtype_spec.py +++ b/test/unit/test_dtype_spec.py @@ -1,7 +1,7 @@ import unittest, math, subprocess from tinygrad.tensor import Tensor from tinygrad.dtype import dtypes, DType, DTYPES_DICT -from tinygrad.device import Device, is_dtype_supported +from tinygrad.device import Device from tinygrad.helpers import getenv, DEBUG, EMULATED_DTYPES from test.helpers import slow from hypothesis import given, settings, strategies as strat @@ -11,9 +11,10 @@ settings.register_profile("my_profile", max_examples=50, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False)) settings.load_profile("my_profile") +supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes() core_dtypes = list(DTYPES_DICT.values()) -dtype_ints = [dt for dt in core_dtypes if dtypes.is_int(dt) and is_dtype_supported(dt)] -dtype_floats = [dt for dt in core_dtypes if dtypes.is_float(dt) and is_dtype_supported(dt)] +dtype_ints = [dt for dt in core_dtypes if dtypes.is_int(dt) and dt in supported_dtypes] +dtype_floats = [dt for dt in core_dtypes if dtypes.is_float(dt) and dt in supported_dtypes] FP8E4M3_MAX = 448.0 FP8E5M2_MAX = 57344.0 @@ -25,7 +26,7 @@ def _assert_eq(tensor:Tensor, target_dtype:DType, target, tol_target_dtype:float try: assert tensor.dtype == target_dtype # denormals are zero - if target_dtype in dtypes.floats and (not is_dtype_supported(target_dtype) or target_dtype in EMULATED_DTYPES.tolist(dtypes)): + if target_dtype in dtypes.floats and (target_dtype not in supported_dtypes or target_dtype in EMULATED_DTYPES.tolist(dtypes)): fe, fm = dtypes.finfo(target_dtype) kwargs = {"atol":2 ** (2 - (1 << (fe - 1))), "rtol": 2 ** (-fm)} else: kwargs = {"rtol": {dtypes.float16:1e-3, dtypes.bfloat16:1e-2, dtypes.fp8e4m3:1e-1, dtypes.fp8e5m2:5e-1, @@ -58,7 +59,6 @@ def test_env_set_default_float(self): subprocess.run(['DEFAULT_FLOAT=TYPO python3 -c "from tinygrad import dtypes"'], shell=True, check=True) - @unittest.skipUnless(is_dtype_supported(dtypes.int8), f"no int8 on {Device.DEFAULT}") def test_dtype_str_arg(self): n = np.random.normal(0, 1, (10, 10)).astype(np.float32) tested = 0 @@ -91,7 +91,7 @@ def test_creation(self, default_int, default_float): _assert_eq(Tensor.eye(0), dtypes.default_float, np.eye(0)) _assert_eq(Tensor.eye(3), dtypes.default_float, np.eye(3)) _assert_eq(Tensor.eye(3, dtype=dtypes.int64), dtypes.int64, np.eye(3)) - if is_dtype_supported(dtypes.float16): + if dtypes.float16 in supported_dtypes: _assert_eq(Tensor.eye(3, dtype=dtypes.float16), dtypes.float16, np.eye(3)) @given(strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats)) @@ -100,12 +100,12 @@ def test_full(self, default_int, default_float): _assert_eq(Tensor.zeros((2, 3)), dtypes.default_float, np.zeros((2, 3))) _assert_eq(Tensor.zeros((2, 3), dtype=dtypes.int64), dtypes.int64, np.zeros((2, 3))) - if is_dtype_supported(dtypes.float16): + if dtypes.float16 in supported_dtypes: _assert_eq(Tensor.zeros((2, 3), dtype=dtypes.float16), dtypes.float16, np.zeros((2, 3))) _assert_eq(Tensor.ones((2, 3)), dtypes.default_float, np.ones((2, 3))) _assert_eq(Tensor.ones((2, 3), dtype=dtypes.int64), dtypes.int64, np.ones((2, 3))) - if is_dtype_supported(dtypes.float16): + if dtypes.float16 in supported_dtypes: _assert_eq(Tensor.ones((2, 3), dtype=dtypes.float16), dtypes.float16, np.ones((2, 3))) _assert_eq(Tensor.full((2, 3), 3.0), dtypes.default_float, np.full((2, 3), 3.0)) @@ -113,7 +113,7 @@ def test_full(self, default_int, default_float): _assert_eq(Tensor.full((2, 3), True), dtypes.bool, np.full((2, 3), True)) _assert_eq(Tensor.full((2, 3), 3, dtype=dtypes.int64), dtypes.int64, np.full((2, 3), 3)) _assert_eq(Tensor.full((2, 3), 3.0, dtype=dtypes.int64), dtypes.int64, np.full((2, 3), 3)) - if is_dtype_supported(dtypes.float16): + if dtypes.float16 in supported_dtypes: _assert_eq(Tensor.full((2, 3), 3, dtype=dtypes.float16), dtypes.float16, np.full((2, 3), 3)) _assert_eq(Tensor.full((2, 3), 3.0, dtype=dtypes.float16), dtypes.float16, np.full((2, 3), 3)) @@ -132,10 +132,10 @@ def test_arange(self, default_int, default_float): _assert_eq(Tensor.arange(5), dtypes.default_int, np.arange(5)) _assert_eq(Tensor.arange(120), dtypes.default_int, np.arange(120)) _assert_eq(Tensor.arange(5.0), dtypes.default_float, np.arange(5)) - if is_dtype_supported(dtypes.int16): + if dtypes.int16 in supported_dtypes: _assert_eq(Tensor.arange(5, dtype=dtypes.int16), dtypes.int16, np.arange(5)) _assert_eq(Tensor.arange(5, dtype=dtypes.int64), dtypes.int64, np.arange(5)) - if is_dtype_supported(dtypes.float16): + if dtypes.float16 in supported_dtypes: _assert_eq(Tensor.arange(5, dtype=dtypes.float16), dtypes.float16, np.arange(5)) _assert_eq(Tensor.arange(3, 9, 0.7), dtypes.default_float, np.arange(3, 9, 0.7), 1e-6 if Device.DEFAULT == "WEBGPU" else 1e-7) _assert_eq(Tensor.arange(3, 8.5, 3), dtypes.default_float, np.arange(3, 8.5, 3)) @@ -149,7 +149,7 @@ def setUp(self): def tearDown(self): dtypes.default_int, dtypes.default_float = self.old_default_int, self.old_default_float - @given(strat.sampled_from([d for d in core_dtypes if dtypes.is_int(d) and is_dtype_supported(d)])) + @given(strat.sampled_from([d for d in core_dtypes if dtypes.is_int(d) and d in supported_dtypes])) def test_int_to_float_unary_func(self, dtype): for func in [ lambda t: t.exp(), @@ -167,7 +167,7 @@ def test_int_to_float_unary_func(self, dtype): # float16 can have larger precision errors np.testing.assert_allclose(func(Tensor(a, dtype=dtype)).numpy(), func(torch.tensor(a)), rtol=1e-3, atol=1e-3) - @unittest.skipUnless(is_dtype_supported(dtypes.float16), "need float16") + @unittest.skipUnless(dtypes.float16 in supported_dtypes, "need float16") def test_sum_dtype_arg(self): t = Tensor([40000, 40000], dtype=dtypes.float16) # default float16 sum returns in float16, overflowed in this case @@ -188,13 +188,13 @@ def test_gradient_dtype(self): old_default_float = dtypes.default_float for default_dtype in dtypes.floats: - if not is_dtype_supported(default_dtype): continue + if default_dtype not in supported_dtypes: continue dtypes.default_float = default_dtype - for dtype in dtypes.floats: - if not is_dtype_supported(dtype): continue + for dtype in dtypes.floats: + if dtype not in supported_dtypes: continue if DEBUG >= 2: print(f"testing {default_dtype=}, {dtype=}") - a = Tensor([1, 2, 3], dtype=dtype, requires_grad=True) + a = Tensor([1, 2, 3], dtype=dtype) b = (a * 5).sum() b.backward() # if there is dtype mismatch, lazy should assert assert a.grad.dtype == a.dtype @@ -205,24 +205,24 @@ def test_gradient_dtype(self): @unittest.skipIf(Device.DEFAULT == "PYTHON", "very slow") @slow @unittest.skipIf(Device.DEFAULT == "WEBGPU", "Binding size is larger than the maximum storage buffer binding size") - @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half") + @unittest.skipUnless(dtypes.half in supported_dtypes, "need half") def test_mean_half_precision_underflow(self): N = 10000 x = 0.001 - t = Tensor([[x]], dtype=dtypes.half, requires_grad=True).expand(N, N).contiguous() + t = Tensor([[x]], dtype=dtypes.half).expand(N, N).contiguous() np.testing.assert_allclose(t.mean(axis=1).numpy(), np.array([x] * N, dtype=np.float16), rtol=1e-3) @unittest.skip("this test only works with SPLIT_REDUCEOP=1") - @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half") + @unittest.skipUnless(dtypes.half in supported_dtypes, "need half") def test_mean_half_precision_overflow(self): N = 256 - t = Tensor([60000] * N*N, dtype=dtypes.half, requires_grad=True).reshape(N, N) + t = Tensor([60000] * N*N, dtype=dtypes.half).reshape(N, N) np.testing.assert_allclose(t.mean().numpy(), 60000) t.square().mean().backward() np.testing.assert_allclose(t.grad.numpy().flatten(), [60000 * 2 / (N*N)] * N*N) @unittest.skipIf(Device.DEFAULT == "WEBGPU", "Precision error") - @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half") + @unittest.skipUnless(dtypes.half in supported_dtypes, "need half") def test_softmax_dtype(self): data = [1, 2, 3] t = Tensor(data, dtype=dtypes.half) diff --git a/test/unit/test_function.py b/test/unit/test_function.py index 677dae3492977..a779aa6e021d9 100644 --- a/test/unit/test_function.py +++ b/test/unit/test_function.py @@ -91,8 +91,8 @@ def test_grad_gemm(self): @function def f(a:Tensor, b:Tensor) -> Tensor: return a @ b - a = Tensor([[1.,2.],[3.,4.]], requires_grad=True) - b = Tensor([[5.,6.],[7.,8.]], requires_grad=True) + a = Tensor([[1.,2.],[3.,4.]]) + b = Tensor([[5.,6.],[7.,8.]]) (f(a, b).contiguous() * b).sum().backward() Tensor.realize(a, b, a.grad, b.grad) # L = sum((a@b) * b), dL/d(a@b) = b, dL/da = b @ b^T, dL/db = a^T @ b + (a@b) @@ -101,7 +101,7 @@ def f(a:Tensor, b:Tensor) -> Tensor: return a @ b np.testing.assert_allclose(b.grad.numpy(), na.T @ nb + na @ nb) def test_grad_implicit(self): - w = Tensor([1., 2., 3.], requires_grad=True) + w = Tensor([1., 2., 3.]) w.realize() # TODO: this is required @function(allow_implicit=True) def f(x:Tensor) -> Tensor: return x * w @@ -144,8 +144,8 @@ def inner(x:Tensor) -> Tensor: return x + w @function(allow_implicit=True) def outer(a:Tensor, b:Tensor) -> Tensor: return inner(a.reshape(1,2) + b.reshape(1,2)) - a = Tensor([1., 2.], requires_grad=True) - b = Tensor([3., 4.], requires_grad=True) + a = Tensor([1., 2.]) + b = Tensor([3., 4.]) outer(a, b).sum().backward() np.testing.assert_allclose(a.grad.numpy(), [2., 2.]) np.testing.assert_allclose(b.grad.numpy(), [2., 2.]) @@ -154,9 +154,9 @@ def test_unused_param_backward(self): @function def f(a:Tensor, b:Tensor, c:Tensor) -> Tensor: return a + c # b is unused - a = Tensor([1., 2., 3.], requires_grad=True) - b = Tensor([4., 5., 6.], requires_grad=True) - c = Tensor([7., 8., 9.], requires_grad=True) + a = Tensor([1., 2., 3.]) + b = Tensor([4., 5., 6.]) + c = Tensor([7., 8., 9.]) f(a, b, c).sum().backward() np.testing.assert_allclose(a.grad.numpy(), [1., 1., 1.]) np.testing.assert_allclose(b.grad.numpy(), [0., 0., 0.]) @@ -279,7 +279,7 @@ def f(x:Tensor, w:Tensor) -> Tensor: return x @ w np.testing.assert_allclose(f(x, w).numpy(), [[1.,2.],[3.,4.],[5.,6.],[7.,8.]]) def test_grad_implicit_multi(self): - w = Tensor([1., 2., 3., 4.], requires_grad=True).shard(self.devices_2, axis=None) + w = Tensor([1., 2., 3., 4.]).shard(self.devices_2, axis=None) w.realize() @function(allow_implicit=True) def f(x:Tensor) -> Tensor: return x * w @@ -315,8 +315,8 @@ def test_data_parallel_backward(self): @function def f(x:Tensor, w:Tensor) -> Tensor: return x @ w - x = Tensor([[1.,0.],[0.,1.],[1.,1.],[0.,0.]], requires_grad=True).shard(self.devices_2, axis=0) - w = Tensor([[1.,2.],[3.,4.]], requires_grad=True).shard(self.devices_2, axis=None) + x = Tensor([[1.,0.],[0.,1.],[1.,1.],[0.,0.]]).shard(self.devices_2, axis=0) + w = Tensor([[1.,2.],[3.,4.]]).shard(self.devices_2, axis=None) w.realize() f(x, w).sum().backward() # d/dx = ones @ w^T = [[1,3],[1,3],[1,3],[1,3]], but sum so ones(4,2) @ w^T? no: @@ -328,26 +328,26 @@ def test_data_parallel_backward_4(self): @function def f(x:Tensor, w:Tensor) -> Tensor: return x @ w - x = Tensor(np.arange(16).reshape(8,2).astype(np.float32), requires_grad=True).shard(devices_4, axis=0) - w = Tensor([[1.,2.],[3.,4.]], requires_grad=True).shard(devices_4, axis=None) + x = Tensor(np.arange(16).reshape(8,2).astype(np.float32)).shard(devices_4, axis=0) + w = Tensor([[1.,2.],[3.,4.]]).shard(devices_4, axis=None) w.realize() f(x, w).sum().backward() np.testing.assert_allclose(x.grad.numpy(), np.ones((8,2)) @ np.array([[1,3],[2,4]])) def test_data_parallel_backward_implicit(self): devices_4 = tuple(f"CPU:{i}" for i in range(4)) - w = Tensor([[1.,2.],[3.,4.]], requires_grad=True).shard(devices_4, axis=None) + w = Tensor([[1.,2.],[3.,4.]]).shard(devices_4, axis=None) w.realize() @function(allow_implicit=True) def f(x:Tensor) -> Tensor: return x @ w - x = Tensor(np.arange(16).reshape(8,2).astype(np.float32), requires_grad=True).shard(devices_4, axis=0) + x = Tensor(np.arange(16).reshape(8,2).astype(np.float32)).shard(devices_4, axis=0) f(x).sum().backward() np.testing.assert_allclose(x.grad.numpy(), np.ones((8,2)) @ np.array([[1,3],[2,4]])) def test_data_parallel_backward_twice(self): devices_4 = tuple(f"CPU:{i}" for i in range(4)) - w = Tensor([[1.,2.],[3.,4.]], requires_grad=True).shard(devices_4, axis=None) + w = Tensor([[1.,2.],[3.,4.]]).shard(devices_4, axis=None) w.realize() # pre-init grads like the training loop does w.grad = w.zeros_like().contiguous().realize() @@ -356,7 +356,7 @@ def f(x:Tensor) -> Tensor: return x @ w expected = np.ones((8,2)) @ np.array([[1,3],[2,4]]) for _ in range(2): - x = Tensor(np.arange(16).reshape(8,2).astype(np.float32), requires_grad=True).shard(devices_4, axis=0) + x = Tensor(np.arange(16).reshape(8,2).astype(np.float32)).shard(devices_4, axis=0) f(x).sum().backward() np.testing.assert_allclose(x.grad.numpy(), expected) @@ -373,8 +373,8 @@ def f(t:Tensor): return (t+1, t+2) def test_tuple_precompile(self): self.test_tuple(True) def test_grad_tuple(self, precompile=False): - x = Tensor.ones(3, requires_grad=True).contiguous() - y = Tensor.ones(3, requires_grad=True).contiguous() + x = Tensor.ones(3).contiguous() + y = Tensor.ones(3).contiguous() @function(precompile=precompile) def f(u1:Tensor, u2:Tensor): return (u1+1, u2+2) t1, t2 = f(x,y) @@ -389,8 +389,8 @@ def grad_fxn(d_out0:UOp, d_out1:UOp, call:UOp): # df/du1 = d_out0, df/du2 = d_out1 return (d_out0, d_out1) - x = Tensor.ones(3, requires_grad=True).contiguous() - y = Tensor.ones(3, requires_grad=True).contiguous() + x = Tensor.ones(3).contiguous() + y = Tensor.ones(3).contiguous() @function(grad_fxn=grad_fxn) def f(u1:Tensor, u2:Tensor): return (u1+1, u2+2) t1, t2 = f(x, y) @@ -406,8 +406,8 @@ def f(x:Tensor, w:Tensor): b = (x @ w) * 2 # shares x@w with a return (a, b) - x = Tensor([[1., 2.], [3., 4.]], requires_grad=True).contiguous() - w = Tensor([[1., 0.], [0., 1.]], requires_grad=True).contiguous() + x = Tensor([[1., 2.], [3., 4.]]).contiguous() + w = Tensor([[1., 0.], [0., 1.]]).contiguous() Tensor.realize(x, w) t1, _ = f(x, w) t1.sum().backward() @@ -435,7 +435,7 @@ def f(a:Tensor): c, d = Tensor.custom_kernel(c, d, a, fxn=my_kernel, grad_fxn=my_grad)[:2] return c, d - a = Tensor([1., 2., 3., 4.], requires_grad=True).contiguous() + a = Tensor([1., 2., 3., 4.]).contiguous() Tensor.realize(a) c, _ = f(a) c.sum().backward() @@ -459,7 +459,7 @@ def f(a:Tensor): c, d = Tensor.custom_kernel(c, d, a, fxn=my_kernel, grad_fxn=my_grad)[:2] return (c, d) - a = Tensor([1., 2., 3., 4.], requires_grad=True).contiguous() + a = Tensor([1., 2., 3., 4.]).contiguous() Tensor.realize(a) c, d = f(a) (c.sum() + d.sum()).backward() # dL/da = (1 + 1) since grad_fxn passes d_combined through @@ -485,7 +485,7 @@ def count_kernels(t:Tensor): return sum((len(call.device) if isinstance(call.device, tuple) else 1) for call in linear.src if call.src[0].op is Ops.SINK) - a = Tensor([1., 2., 3., 4.], requires_grad=True).contiguous() + a = Tensor([1., 2., 3., 4.]).contiguous() Tensor.realize(a) c = f(a) @@ -495,6 +495,25 @@ def count_kernels(t:Tensor): Tensor.realize(a.grad) np.testing.assert_allclose(a.grad.numpy(), [2., 2., 2., 2.]) + def test_custom_kernel_precompile_multidevice(self): + # a custom_kernel output placeholder (invalids) under multi-device @function(precompile=True) must return the + # kernel's computed result. read it back through .numpy() so the cross-device gather reads the output buffer + devs = ("CPU:0", "CPU:1") + def double_kernel(C:UOp, A:UOp) -> UOp: + C, A = C.flatten(), A.flatten() + i = UOp.range(A.numel(), 0) + return C[i].store(A[i] * 2.0).end(i).sink(arg=KernelInfo(name="double_kernel")) + def double_grad(d_c:UOp, call:UOp): return (None, (Tensor(d_c) * 2.0).uop) + + @function(precompile=True, precompile_backward=True) + def f(a:Tensor): + c = Tensor(Tensor.invalids(a.shape[0]//len(devs), a.shape[1], dtype=a.dtype, device=devs).uop.multi(0), device=devs) + return Tensor.custom_kernel(c, a, fxn=double_kernel, grad_fxn=double_grad)[0] + + a = Tensor.full((4, 4), 7.0).contiguous().shard(devs, axis=0) + Tensor.realize(a) + np.testing.assert_allclose(f(a).numpy(), 14.0) + def test_custom_kernel_precompile_further_compute(self): def my_kernel(C:UOp, A:UOp) -> UOp: i = UOp.range(A.shape[0], 0) @@ -513,9 +532,9 @@ class TestFunctionGrad(unittest.TestCase): def test_function_grad_ops(self, precompile=False, precompile_backward=False): N = 64 x = Tensor.ones(N,N).contiguous() - w1 = Tensor.ones(N,N, requires_grad=True).contiguous() - w2 = Tensor.ones(N,N, requires_grad=True).contiguous() - w3 = Tensor.ones(N,N, requires_grad=True).contiguous() + w1 = Tensor.ones(N,N).contiguous() + w2 = Tensor.ones(N,N).contiguous() + w3 = Tensor.ones(N,N).contiguous() ref = Tensor.ones(N,N).contiguous() Tensor.realize(x, w1, w2, w3, ref) @function(precompile=precompile, precompile_backward=precompile_backward) diff --git a/test/unit/test_gguf.py b/test/unit/test_gguf.py index 890f14eea96c1..8b74d88aa516d 100644 --- a/test/unit/test_gguf.py +++ b/test/unit/test_gguf.py @@ -3,12 +3,12 @@ from tinygrad.helpers import disable_gc from tinygrad.llm.gguf import _ggml_iq_grid, ggml_data_to_tensor, gguf_load from tinygrad.runtime.autogen import ggml_common as _ggml -from tinygrad.device import is_dtype_supported import numpy as np from gguf import GGUFReader, GGUFValueType, GGMLQuantizationType, GGML_QUANT_SIZES, dequantize, quantize from gguf.quants import IQ2_S, IQ3_S, IQ3_XXS ggml_test_block_count = 4 +supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes() class TestGGUFTables(unittest.TestCase): def test_iq2_s_grid_matches_gguf_py(self): @@ -26,7 +26,7 @@ def test_iq3_s_grid_matches_gguf_py(self): grid = _ggml_iq_grid(Device.DEFAULT, _ggml.iq3s_grid, (512, 4)).numpy() np.testing.assert_equal(grid, IQ3_S.grid.reshape(512, 4)) -@unittest.skipIf(any(not is_dtype_supported(t) for t in [ dtypes.uint8, dtypes.half ]), "Backend must support uint8 and half") +@unittest.skipUnless(dtypes.uint8 in supported_dtypes and dtypes.half in supported_dtypes, "Backend must support uint8 and half") class TestGGUF(unittest.TestCase): def test_load_tinyllama_q8_0(self): self._test_gguf_load("https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories15M-q8_0.gguf?download=true") def test_load_tinyllama_q4_0(self): self._test_gguf_load("https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories15M-q4_0.gguf?download=true") @@ -60,7 +60,7 @@ def test_dequantization_iq3_s(self): self._test_dequantization(GGMLQuantizationT def test_dequantization_iq2_s(self): self._test_dequantization(GGMLQuantizationType.IQ2_S) def test_dequantization_iq4_xs(self): self._test_dequantization(GGMLQuantizationType.IQ4_XS) def test_dequantization_mxfp4(self): self._test_dequantization(GGMLQuantizationType.MXFP4) - @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), "Backend must support bfloat16") + @unittest.skipUnless(dtypes.bfloat16 in supported_dtypes, "Backend must support bfloat16") def test_dequantization_bf16(self): self._test_dequantization(GGMLQuantizationType.BF16) def test_dequantization_mxfp4_old(self): def encode(nibbles, E): @@ -229,7 +229,7 @@ def _test_gguf_gemv(self, qtype: GGMLQuantizationType): x = rng.standard_normal(cols).astype(np.float32) with np.errstate(all='ignore'): np.testing.assert_allclose((tensors["weight"] @ Tensor(x)).numpy(), ref @ x, atol=1e-2, rtol=1e-2) - if qtype == GGMLQuantizationType.BF16 or is_dtype_supported(dtypes.half): np.testing.assert_equal(tensors["weight"].numpy(), ref) + if qtype == GGMLQuantizationType.BF16 or dtypes.half in supported_dtypes: np.testing.assert_equal(tensors["weight"].numpy(), ref) assert np.isfinite(ref).all() and np.isfinite(tensors["weight"].numpy()).all(), f"{qtype.name} has NaN/Inf" def test_gguf_gemv_q8_0(self): self._test_gguf_gemv(GGMLQuantizationType.Q8_0) @@ -243,7 +243,7 @@ def test_gguf_gemv_iq3_s(self): self._test_gguf_gemv(GGMLQuantizationType.IQ3_S) def test_gguf_gemv_iq2_s(self): self._test_gguf_gemv(GGMLQuantizationType.IQ2_S) def test_gguf_gemv_iq4_xs(self): self._test_gguf_gemv(GGMLQuantizationType.IQ4_XS) def test_gguf_gemv_mxfp4(self): self._test_gguf_gemv(GGMLQuantizationType.MXFP4) - @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), "Backend must support bfloat16") + @unittest.skipUnless(dtypes.bfloat16 in supported_dtypes, "Backend must support bfloat16") def test_gguf_gemv_bf16(self): self._test_gguf_gemv(GGMLQuantizationType.BF16) class TestGGUFGC(unittest.TestCase): diff --git a/test/unit/test_gradient.py b/test/unit/test_gradient.py index a52af55ed7366..47774c712b152 100644 --- a/test/unit/test_gradient.py +++ b/test/unit/test_gradient.py @@ -52,13 +52,13 @@ def test_non_float_tensor_raise(self): with self.assertRaises(RuntimeError): x.float().sum().gradient(x) def test_copy_to_device_gradient(self): - t = Tensor([1.0, 2, 3], requires_grad=True).realize() + t = Tensor([1.0, 2, 3]).realize() t.to("CPU:1").square().sum().backward() self.assertEqual(t.grad.device, t.device) self.assertListEqual(t.grad.tolist(), [2.0, 4.0, 6.0]) def test_multiple_backward(self): - x = Tensor([3.], requires_grad=True) + x = Tensor([3.]) (x*2)[0].backward() np.testing.assert_allclose(x.grad.numpy(), [2.0]) old_grad = x.grad @@ -69,23 +69,23 @@ def test_multiple_backward(self): np.testing.assert_allclose(x.grad.numpy(), [2.0+3.0+2*3.0]) self.assertIs(x.grad, old_grad) - def test_gradient_through_clone_from_non_grad_src(self): - src = Tensor([1.0, 2.0, 3.0, 4.0], requires_grad=False) - x = src.clone().requires_grad_(True) - (x * 2.0).sum().backward() - np.testing.assert_allclose(x.grad.numpy(), [2.0, 2.0, 2.0, 2.0]) - self.assertIsNone(src.grad) - def test_gradient_through_clone_from_grad_src(self): - # unlike torch, tinygrad accumulates grad on all requires_grad tensors, including non-leaf x - src = Tensor([1.0, 2.0, 3.0, 4.0], requires_grad=True) - x = src.clone().requires_grad_(True) + # unlike torch, tinygrad accumulates grad on every tensor in the graph, including non-leaf x + src = Tensor([1.0, 2.0, 3.0, 4.0]) + x = src.clone() (x * 2.0).sum().backward() np.testing.assert_allclose(src.grad.numpy(), [2.0, 2.0, 2.0, 2.0]) np.testing.assert_allclose(x.grad.numpy(), [2.0, 2.0, 2.0, 2.0]) + def test_gradient_through_clone_from_detached_src(self): + base = Tensor([1.0, 2.0, 3.0, 4.0]) + x = base.detach().clone() + (x * 2.0).sum().backward() + np.testing.assert_allclose(x.grad.numpy(), [2.0, 2.0, 2.0, 2.0]) # gradient flows through clone + np.testing.assert_allclose(base.grad.numpy(), [0.0, 0.0, 0.0, 0.0]) # ...but detach blocks it from base + def test_setitem_on_grad_used_tensor_raises(self): - x = Tensor([1.0, 2.0, 3.0, 4.0], requires_grad=True).realize() + x = Tensor([1.0, 2.0, 3.0, 4.0]).realize() _ = (x * 2.0).sum() with self.assertRaises(RuntimeError): x[0] = 99.0 @@ -98,6 +98,12 @@ def test_gradient_through_chained_unrealized_setitem(self): x = Tensor.randn(4, 4) np.testing.assert_allclose(x.pad(((1,0),(0,0))).gradient(x, gradient=g2)[0].numpy(), np.zeros((4, 4))) + def test_bare_const_skipped_by_backward(self): + Tensor.manual_seed(0) + w = Tensor(1.0) + (Tensor.rand(()) + w).backward() + self.assertIsNone(w.grad) + class TestMultiOutputGradient(unittest.TestCase): @staticmethod def addmul_kernel(C:UOp, D:UOp, A:UOp, B:UOp) -> UOp: @@ -115,10 +121,10 @@ def backward_addmul(grad_c, grad_d, call): def test_custom_kernel_multi_output_backward(self): a_np, b_np = np.random.randn(4, 4).astype(np.float32), np.random.randn(4, 4).astype(np.float32) - a_ref, b_ref = Tensor(a_np, requires_grad=True), Tensor(b_np, requires_grad=True) + a_ref, b_ref = Tensor(a_np), Tensor(b_np) ((a_ref + b_ref).sum() + (a_ref * b_ref).sum()).backward() - a, b = Tensor(a_np, requires_grad=True), Tensor(b_np, requires_grad=True) + a, b = Tensor(a_np), Tensor(b_np) Tensor.realize(a, b) c, d, _, _ = Tensor.custom_kernel(Tensor.empty(4, 4), Tensor.empty(4, 4), a, b, fxn=self.addmul_kernel, grad_fxn=self.backward_addmul) (c.sum() + d.sum()).backward() @@ -127,10 +133,10 @@ def test_custom_kernel_multi_output_backward(self): def test_custom_kernel_multi_output_backward_interacting(self): a_np, b_np = np.random.randn(4, 4).astype(np.float32), np.random.randn(4, 4).astype(np.float32) - a_ref, b_ref = Tensor(a_np, requires_grad=True), Tensor(b_np, requires_grad=True) + a_ref, b_ref = Tensor(a_np), Tensor(b_np) ((a_ref + b_ref) * (a_ref * b_ref)).sum().backward() - a, b = Tensor(a_np, requires_grad=True), Tensor(b_np, requires_grad=True) + a, b = Tensor(a_np), Tensor(b_np) Tensor.realize(a, b) c, d, _, _ = Tensor.custom_kernel(Tensor.empty(4, 4), Tensor.empty(4, 4), a, b, fxn=self.addmul_kernel, grad_fxn=self.backward_addmul) (c * d).sum().backward() @@ -152,21 +158,21 @@ def backward_addmulsub(grad_c, grad_d, grad_e, call): return (None, None, None, grad_a, grad_b) a_np, b_np = np.random.randn(4, 4).astype(np.float32), np.random.randn(4, 4).astype(np.float32) - a_ref, b_ref = Tensor(a_np, requires_grad=True), Tensor(b_np, requires_grad=True) + a_ref, b_ref = Tensor(a_np), Tensor(b_np) ((a_ref + b_ref).sum() + (a_ref * b_ref).sum() + (a_ref - b_ref).sum()).backward() - a, b = Tensor(a_np, requires_grad=True), Tensor(b_np, requires_grad=True) + a, b = Tensor(a_np), Tensor(b_np) Tensor.realize(a, b) c, d, e, _, _ = Tensor.custom_kernel(Tensor.empty(4, 4), Tensor.empty(4, 4), Tensor.empty(4, 4), a, b, fxn=addmulsub_kernel, grad_fxn=backward_addmulsub) (c.sum() + d.sum() + e.sum()).backward() - np.testing.assert_allclose(a.grad.numpy(), a_ref.grad.numpy(), rtol=1e-5) - np.testing.assert_allclose(b.grad.numpy(), b_ref.grad.numpy(), rtol=1e-5) + np.testing.assert_allclose(a.grad.numpy(), a_ref.grad.numpy(), atol=1e-6, rtol=1e-5) + np.testing.assert_allclose(b.grad.numpy(), b_ref.grad.numpy(), atol=1e-6, rtol=1e-5) class TestViewGradient(unittest.TestCase): def test_expand(self): x = Tensor.randn(5,2) - a = Tensor([3.], requires_grad=True) + a = Tensor([3.]) aex = a.expand(10) (aex.reshape(5,2) * x).sum().backward() np.testing.assert_allclose(aex.grad.numpy(), x.reshape(10).numpy()) diff --git a/test/unit/test_hashing.py b/test/unit/test_hashing.py index 62a73d72dc88c..78df77115ba8e 100644 --- a/test/unit/test_hashing.py +++ b/test/unit/test_hashing.py @@ -3,11 +3,12 @@ from tinygrad import Tensor, Device, dtypes from tinygrad.helpers import DEV from test.helpers import slow -from tinygrad.device import is_dtype_supported from tinygrad.uop.ops import UOp from tinygrad.engine.jit import TinyJit -@unittest.skipUnless(is_dtype_supported(dtypes.uint8) and is_dtype_supported(dtypes.uint64), "Device must support uint8 and uint64") +supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes() + +@unittest.skipUnless(dtypes.uint8 in supported_dtypes and dtypes.uint64 in supported_dtypes, "Device must support uint8 and uint64") @unittest.skipIf(DEV.interface.startswith("MOCK") and Device.DEFAULT == "NV", "crashes in NV CI") class TestHashing(unittest.TestCase): def _python_hash_1mb(self, data:bytes): @@ -21,7 +22,7 @@ def test_abc(self): out = Tensor(b"abc").hash() self.assertEqual(bytes(out.data()), expected) -@unittest.skipUnless(is_dtype_supported(dtypes.uint8) and is_dtype_supported(dtypes.uint64), "Device must support uint8 and uint64") +@unittest.skipUnless(dtypes.uint8 in supported_dtypes and dtypes.uint64 in supported_dtypes, "Device must support uint8 and uint64") @unittest.skipIf(DEV.interface.startswith("MOCK") and Device.DEFAULT == "NV", "crashes in NV CI") class TestKeccak(unittest.TestCase): def setUp(self) -> None: random.seed(1337) diff --git a/test/unit/test_indexing.py b/test/unit/test_indexing.py index 05066fa71bf54..db1c962519080 100644 --- a/test/unit/test_indexing.py +++ b/test/unit/test_indexing.py @@ -153,7 +153,7 @@ def test_set_item_to_scalar_tensor(self): n = random.randint(1, 10) z = Tensor.randn([m, n]) a = 1.0 - w = Tensor(a, requires_grad=True) + w = Tensor([a]) z[:, 0] = w z.sum().backward() numpy_testing_assert_equal_helper(w.grad, m * a) @@ -952,11 +952,9 @@ def assert_set_eq(tensor: Tensor, indexer, val): # NOTE: torch initiates the gradients using g0cpu (rand as gradients) def assert_backward_eq(tensor: Tensor, indexer): cpu = clone(tensor.float()) - cpu.requires_grad = True outcpu = cpu[indexer].sum() outcpu.backward() dev = cpu.detach() - dev.requires_grad = True outdev = dev[indexer].sum() outdev.backward() numpy_testing_assert_equal_helper(cpu.grad, dev.grad) diff --git a/test/unit/test_invalid_tensor.py b/test/unit/test_invalid_tensor.py index 6eb2e083a65fe..2e051cbac9711 100644 --- a/test/unit/test_invalid_tensor.py +++ b/test/unit/test_invalid_tensor.py @@ -82,17 +82,17 @@ def test_where_reduce_always_true(self): def test_invalid_unary(self): mask = Tensor.arange(4) < 2 - out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor.full((4,), Invalid, dtype=dtypes.float).sqrt()) + out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor.full((4,), Invalid, dtype=dtypes.float, buffer=False).sqrt()) self._invalid_test_helper(out, [1.0, 2.0, None, None]) def test_invalid_binary(self): mask = Tensor.arange(4) < 2 - out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor.full((4,), Invalid, dtype=dtypes.float) + 2) + out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor.full((4,), Invalid, dtype=dtypes.float, buffer=False) + 2) self._invalid_test_helper(out, [1.0, 2.0, None, None]) def test_invalid_binary_left(self): mask = Tensor.arange(4) < 2 - out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), 2 + Tensor.full((4,), Invalid, dtype=dtypes.float)) + out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), 2 + Tensor.full((4,), Invalid, dtype=dtypes.float, buffer=False)) self._invalid_test_helper(out, [1.0, 2.0, None, None]) def test_invalid_reshape(self): @@ -102,17 +102,17 @@ def test_invalid_reshape(self): def test_invalid_cast(self): mask = Tensor.arange(4) < 2 - out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor.full((4,), Invalid, dtype=dtypes.int).cast(dtypes.float)) + out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor.full((4,), Invalid, dtype=dtypes.int, buffer=False).cast(dtypes.float)) self._invalid_test_helper(out, [1.0, 2.0, None, None]) def test_invalid_bitcast(self): mask = Tensor.arange(4) < 2 - out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor.full((4,), Invalid, dtype=dtypes.int).bitcast(dtypes.float)) + out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor.full((4,), Invalid, dtype=dtypes.int, buffer=False).bitcast(dtypes.float)) self._invalid_test_helper(out, [1.0, 2.0, None, None]) def test_where_bitcast(self): mask = Tensor.arange(4) < 2 - out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor.full((4,), Invalid, dtype=dtypes.int)).bitcast(dtypes.int) + out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor.full((4,), Invalid, dtype=dtypes.int, buffer=False)).bitcast(dtypes.int) self._invalid_test_helper(out, [0x3f800000, 0x40000000, None, None]) # tensor indexing uses reduce, so the entire result becomes invalid diff --git a/test/unit/test_linalg.py b/test/unit/test_linalg.py index 1bf0f36d9a9fb..29eb996865911 100644 --- a/test/unit/test_linalg.py +++ b/test/unit/test_linalg.py @@ -42,7 +42,6 @@ def test_svd_nonfull_2_2(self): self._test_svd_nonfull((2,2)) def test_svd_nonfull_5_3(self): self._test_svd_nonfull((5,3)) def test_svd_nonfull_3_5(self): self._test_svd_nonfull((3,5)) def test_svd_nonfull_2_2_2_2_3(self): self._test_svd_nonfull((2,2,2,2,3)) - def test_svd_nonfull_5_5(self): self._test_svd_nonfull((5,5)) @unittest.skip("very big. recommend wrapping with TinyJit around inner function") def test_svd_large(self): diff --git a/test/unit/test_metal_graph.py b/test/unit/test_metal_graph.py index 3038dcfb345d7..6f2951cfeefac 100644 --- a/test/unit/test_metal_graph.py +++ b/test/unit/test_metal_graph.py @@ -1,7 +1,7 @@ import unittest from unittest.mock import MagicMock from tinygrad import Device -from tinygrad.uop.ops import Ops +from tinygrad.uop.ops import Ops, UOp from tinygrad.dtype import dtypes @unittest.skipUnless(Device.DEFAULT == "METAL", "Metal device required to run") @@ -14,8 +14,10 @@ def setUp(self): def metal_buf(self, offset): buf = MagicMock() if offset > 0: - buf.op = Ops.BUFFER_VIEW - buf.arg = (None, offset) + buf.op = Ops.SLICE + src = MagicMock() + src.dtype = dtypes.uint8 + buf.src = (src, UOp.const(dtypes.weakint, offset)) buf.dtype = dtypes.uint8 else: buf.op = Ops.BUFFER @@ -34,7 +36,7 @@ def test_supports_uop_overflow_offset(self): assert self.MetalGraph.supports_uop([self.dev], self.call(self.metal_buf(0), self.metal_buf(0x100000000))) is False def test_supports_uop_nonmetal_buf(self): - # non-BUFFER_VIEW ops should not be checked for offset + # non-SLICE ops should not be checked for offset buf = MagicMock() buf.op = Ops.BUFFER buf.device = Device.DEFAULT diff --git a/test/unit/test_randomness.py b/test/unit/test_randomness.py new file mode 100644 index 0000000000000..67211e472278b --- /dev/null +++ b/test/unit/test_randomness.py @@ -0,0 +1,194 @@ +import unittest, math, torch +import numpy as np +from functools import partial +from tinygrad import nn, dtypes, Tensor, Device, TinyJit, Variable +from tinygrad.helpers import OSX + +# https://gist.github.com/devries/11405101 +def ksprob(a): + fac, total, termbf = 2.0, 0.0, 0.0 + a2 = -2.0 * a * a + for j in range(1, 101): + term = fac * math.exp(a2 * j * j) + total += term + if math.fabs(term) <= 0.001 * termbf or math.fabs(term) <= 1e-8 * total: + return total + fac = -fac + termbf = math.fabs(term) + return 1.0 + +def kstest(l1, l2): + n1, n2 = len(l1), len(l2) + l1.sort() + l2.sort() + j1, j2, d, fn1, fn2 = 0, 0, 0.0, 0.0, 0.0 + while j1 < n1 and j2 < n2: + d1, d2 = l1[j1], l2[j2] + if d1 <= d2: + fn1 = (float(j1) + 1.0) / float(n1) + j1 += 1 + if d2 <= d1: + fn2 = (float(j2) + 1.0) / float(n2) + j2 += 1 + dtemp = math.fabs(fn2 - fn1) + if dtemp > d: + d = dtemp + ne = float(n1 * n2) / float(n1 + n2) + nesq = math.sqrt(ne) + prob = ksprob((nesq + 0.12 + 0.11 / nesq) * d) + return prob + +def equal_distribution(tiny_func, torch_func=None, numpy_func=None, shape=(40, 43), alpha=0.04): + Tensor.manual_seed(1337) + torch.manual_seed(1337) + np.random.seed(1337) + assert not (torch_func is None and numpy_func is None), "no function to compare with" + x1 = tiny_func(*shape).numpy().flatten() + x2 = tiny_func(shape).numpy().flatten() + if numpy_func is not None: y = numpy_func(shape).flatten() + if torch_func is not None: z = torch_func(shape).numpy().flatten() + return (numpy_func is None or (kstest(x1, y) >= alpha and kstest(x2, y) >= alpha)) and \ + (torch_func is None or (kstest(x1, z) >= alpha and kstest(x2, z) >= alpha)) + +def normal_test(func, shape=(20, 45), alpha=0.05): return equal_distribution(func, numpy_func=lambda x: np.random.randn(*x), shape=shape, alpha=alpha) + +class TestRandomness(unittest.TestCase): + def test_randn(self): + self.assertEqual(Tensor.randn(3,3,dtype=dtypes.half).dtype, dtypes.half) + self.assertTrue(normal_test(Tensor.randn)) + self.assertTrue(equal_distribution(Tensor.randn, torch.randn, lambda x: np.random.randn(*x))) + + def test_randint(self): + self.assertFalse(normal_test(Tensor.randint)) + self.assertTrue(equal_distribution(partial(Tensor.randint, low=-2, high=5), + numpy_func=lambda x: np.random.randint(low=-2, high=5, size=x))) + self.assertTrue(equal_distribution(partial(Tensor.randint, low=-2, high=5, dtype="int32"), + numpy_func=lambda x: np.random.randint(low=-2, high=5, size=x))) + self.assertTrue(Tensor.randint(1, device="CPU").device=="CPU") + # check types of args + with self.assertRaises(TypeError): Tensor.randint((3, 4), low=0.1, high=3) + with self.assertRaises(TypeError): Tensor.randint((3, 4), low=0, high=3.5) + with self.assertRaises(TypeError): Tensor.randint((3, 4), low=1, high=3, dtype="float") + with self.assertRaises(TypeError): Tensor.randint((3, 4), low=0, high=3, dtype=dtypes.float32) + # check low < high + with self.assertRaises(ValueError): Tensor.randint((3, 4), low=10, high=5) + with self.assertRaises(ValueError): Tensor.randint((3, 4), low=10, high=10) + np.testing.assert_array_equal(Tensor.randint(16, low=5, high=6).numpy(), 5) + + def test_normal(self): + self.assertTrue(normal_test(Tensor.normal)) + self.assertTrue(equal_distribution(Tensor.normal, lambda x: torch.nn.init.normal_(torch.empty(x), mean=0, std=1), + lambda x: np.random.normal(loc=0, scale=1, size=x))) + # check std >= 0 + with self.assertRaises(ValueError): Tensor.normal((3, 4), mean=0, std=-1) + + def test_uniform(self): + self.assertFalse(normal_test(Tensor.uniform)) + self.assertTrue(equal_distribution(Tensor.uniform, lambda x: torch.nn.init.uniform_(torch.empty(x)), lambda x: np.random.uniform(size=x))) + self.assertTrue(equal_distribution(partial(Tensor.uniform, low=-100, high=100, dtype=dtypes.int32), + numpy_func=lambda x: np.random.randint(low=-100, high=100, size=x))) + # check low < high + with self.assertRaises(ValueError): Tensor.uniform((3, 4), low=5.0, high=3.0) + with self.assertRaises(ValueError): Tensor.uniform((3, 4), low=1.0, high=1.0) + + def test_scaled_uniform(self): + self.assertFalse(normal_test(Tensor.scaled_uniform)) + self.assertTrue(equal_distribution(Tensor.scaled_uniform, lambda x: torch.nn.init.uniform_(torch.empty(x), a=-1, b=1) / math.sqrt(math.prod(x)), + lambda x: np.random.uniform(-1, 1, size=x) / math.sqrt(math.prod(x)))) + + def test_glorot_uniform(self): + self.assertFalse(normal_test(Tensor.glorot_uniform)) + self.assertTrue(equal_distribution(Tensor.glorot_uniform, lambda x: torch.nn.init.xavier_uniform_(torch.empty(x)), + lambda x: np.random.uniform(-1, 1, size=x) * math.sqrt(6 / (x[0] + math.prod(x[1:]))))) + + def test_kaiming_uniform(self): + for shape in [(32, 16, 3, 3), (20, 44), (5, 15, 35)]: + self.assertTrue(equal_distribution(Tensor.kaiming_uniform, lambda x: torch.nn.init.kaiming_uniform_(torch.empty(x)), shape=shape)) + + def test_kaiming_normal(self): + for shape in [(32, 16, 3, 3), (20, 44), (3, 15, 35)]: + self.assertTrue(equal_distribution(Tensor.kaiming_normal, lambda x: torch.nn.init.kaiming_normal_(torch.empty(x)), shape=shape)) + + def test_multinomial(self): + self.assertRaises(AssertionError, lambda: Tensor(2).multinomial(1, replacement=False)) + self.assertRaises(AssertionError, lambda: Tensor([1, 9]).multinomial(0, replacement=False)) + def _check_with_torch(w, num_samples, replacement): + tiny_res = Tensor(w).multinomial(num_samples, replacement=replacement) + torch_res = torch.tensor(w).multinomial(num_samples, replacement=replacement) + self.assertEqual(tiny_res.shape, torch_res.shape) + if torch_res.ndim == 1: + tiny_res = tiny_res.unsqueeze(0) + torch_res = torch_res.unsqueeze(0) + for i in range(torch_res.shape[0]): + self.assertTrue(equal_distribution(lambda *_: tiny_res[i], lambda _: torch_res[i])) + _check_with_torch(w=[0.231, 0., 1., 0.5], num_samples=300, replacement=True) + _check_with_torch(w=[[0.2, 0.8]], num_samples=300, replacement=True) # 2D but only 1 row + _check_with_torch(w=[[0.453, 0., 1., 0.81], [0.1, 0.8, 0., 0.1]], num_samples=300, replacement=True) + # no-replacement + w = [0.1, 0.9] + self.assertRaises(AssertionError, lambda: Tensor(w).multinomial(100, replacement=False)) + + @TinyJit + def sample_one(): return Tensor(w).multinomial(1, replacement=False).realize() + + tiny_samples = [sample_one().item() for _ in range(400)] + torch_samples = [torch.tensor(w).multinomial(1, replacement=False).item() for _ in range(400)] + self.assertTrue(equal_distribution(lambda *_: Tensor(tiny_samples), lambda _: torch.tensor(torch_samples))) + + w = list(range(32)) + s1 = Tensor(w).multinomial(5, replacement=False).numpy() + self.assertEqual(len(set(s1.tolist())), 5) + s2 = Tensor(w).multinomial(5, replacement=False).numpy() + self.assertFalse(np.array_equal(s1, s2)) + full = Tensor(w).multinomial(len(w), replacement=False).numpy() + self.assertEqual(sorted(full.tolist()), w) + + w = [0.1, 0.2, 0.3, 0.4] + @TinyJit + def sample_three(): return Tensor(w).multinomial(3, replacement=False).realize() + + tiny_draws = np.array([sample_three().numpy() for _ in range(400)]) + torch_draws = np.array([torch.tensor(w).multinomial(3, replacement=False).numpy() for _ in range(400)]) + for pos in range(3): + self.assertTrue(equal_distribution(lambda *_: Tensor(tiny_draws[:, pos]), lambda _: torch.tensor(torch_draws[:, pos]))) + + @unittest.skip("this test is flaky") + def test_multinomial_counterexample(self): + tiny_res = Tensor([0.3, 0.6, 0.1]).multinomial(4000, replacement=True) + torch_res = torch.tensor([0.3, 0.6, 0.1]).multinomial(4000, replacement=True) + self.assertTrue(equal_distribution(lambda *_: tiny_res, lambda _: torch_res)) + torch_res = torch.tensor([0.2, 0.7, 0.1]).multinomial(4000, replacement=True) + self.assertFalse(equal_distribution(lambda *_: tiny_res, lambda _: torch_res)) + + def test_conv2d_init(self): + params = (128, 256, (3,3)) + assert equal_distribution(lambda *_: nn.Conv2d(*params).weight, lambda _: torch.nn.Conv2d(*params).weight.detach()) + assert equal_distribution(lambda *_: nn.Conv2d(*params).bias, lambda _: torch.nn.Conv2d(*params).bias.detach()) + + def test_linear_init(self): + params = (64, 256) + assert equal_distribution(lambda *_: nn.Linear(*params).weight, lambda _: torch.nn.Linear(*params).weight.detach()) + assert equal_distribution(lambda *_: nn.Linear(*params).bias, lambda _: torch.nn.Linear(*params).bias.detach()) + + def test_bn_init(self): + params = (64,) + assert equal_distribution(lambda *_: nn.BatchNorm2d(*params).weight, lambda _: torch.nn.BatchNorm2d(*params).weight.detach()) + assert equal_distribution(lambda *_: nn.BatchNorm2d(*params).bias, lambda _: torch.nn.BatchNorm2d(*params).bias.detach()) + +# TODO: still fails with MAX_KERNEL_BUFFERS +@unittest.skipIf(Device.DEFAULT == "WEBGPU" and not OSX, "WEBGPU Vulkan can only run kernels with up to 10 buffers") +class TestSample(unittest.TestCase): + def test_sample(self): + X = Tensor.rand(1000, 50).realize() + BS = 16 + idxs = np.random.randint(0, X.shape[0], size=(BS)) + # this uncovered a bug with arg sort order + batch = [Variable(f'idx{i}', 0, X.shape[0]-1).bind(s) for i,s in enumerate(idxs.tolist())] + x = Tensor.cat(*[X.shrink(((batch[i], batch[i]+1), None)) for i in range(BS)]) + print(idxs) + ret = x.numpy() + base = X.numpy()[idxs] + np.testing.assert_equal(ret, base) + +if __name__ == "__main__": + unittest.main() diff --git a/test/unit/test_realize_is_realize.py b/test/unit/test_realize_is_realize.py index 5bf6581269518..b2e141a6cd42c 100644 --- a/test/unit/test_realize_is_realize.py +++ b/test/unit/test_realize_is_realize.py @@ -15,6 +15,10 @@ def test_contiguous(self): t = Tensor.zeros(10).contiguous().realize() assert t.uop.is_realized + def test_ones(self): + t = Tensor.ones(4, 4).realize() + assert t.uop.is_realized + def test_bytes(self): t = Tensor(b'\x01\x02\x03').realize() assert t.uop.is_realized @@ -51,10 +55,6 @@ def test_const_not_realized(self): t = Tensor(3.14).realize() assert not t.uop.is_realized - def test_ones_not_realized(self): - t = Tensor.ones(4, 4).realize() - assert not t.uop.is_realized - def test_none_not_realized(self): t = Tensor(None).realize() assert not t.uop.is_realized diff --git a/test/unit/test_setitem_schedule.py b/test/unit/test_setitem_schedule.py index 34d9be4b18b78..e69ac4b018660 100644 --- a/test/unit/test_setitem_schedule.py +++ b/test/unit/test_setitem_schedule.py @@ -50,7 +50,7 @@ def test_setitem_into_empty_alu(self): self.assertEqual(GlobalCounters.kernel_count, 0) t.realize() self.assertEqual(GlobalCounters.kernel_count, 1) - self.assertEqual(GlobalCounters.global_mem, 4*(3*2+1)) # 3 elements had +1, 1 is assigned directly + self.assertLessEqual(GlobalCounters.global_mem, 32) t[1].realize() t.realize() self.assertEqual(GlobalCounters.kernel_count, 1) @@ -75,15 +75,15 @@ def test_setitem_into_tensor_alu(self): self.assertEqual(GlobalCounters.kernel_count, 0) t[1].realize() self.assertEqual(GlobalCounters.kernel_count, 1) - self.assertEqual(GlobalCounters.global_mem, 4*(3*2+1)) # 3 elements had +1, 1 is assigned directly + self.assertLessEqual(GlobalCounters.global_mem, 32) t[1].realize() t.realize() self.assertEqual(GlobalCounters.kernel_count, 1) self.assertListEqual(t.tolist(), [2, 5, 4, 5]) - def test_setitem_into_cont(self): + def test_setitem_into_const(self): GlobalCounters.reset() - t = Tensor.ones(4, dtype=dtypes.int32) + t = Tensor.ones(4, dtype=dtypes.int32, buffer=False) t[1] = 5 self.assertEqual(GlobalCounters.kernel_count, 0) t.realize() @@ -96,7 +96,7 @@ def test_setitem_into_cont(self): def test_setitem_into_const_alu(self): GlobalCounters.reset() - t = Tensor.ones(4, dtype=dtypes.int32) + 1 + t = Tensor.ones(4, dtype=dtypes.int32, buffer=False) + 1 t[1] = 5 self.assertEqual(GlobalCounters.kernel_count, 0) t.realize() @@ -110,7 +110,9 @@ def test_setitem_into_const_alu(self): def test_setitem_into_arange(self): # NOTE: arange has no real buffer, but assigning to it is fine GlobalCounters.reset() + other = Tensor.arange(4, dtype=dtypes.int32) t = Tensor.arange(4, dtype=dtypes.int32) + self.assertIs(other.uop, t.uop) t[1] = 5 self.assertEqual(GlobalCounters.kernel_count, 0) t.realize() diff --git a/test/unit/test_shm_tensor.py b/test/unit/test_shm_tensor.py index 71d27cd1a5a4f..19b6c9ff9dba1 100644 --- a/test/unit/test_shm_tensor.py +++ b/test/unit/test_shm_tensor.py @@ -1,11 +1,11 @@ import unittest import multiprocessing.shared_memory as shared_memory -from tinygrad.helpers import CI, WIN +from tinygrad.helpers import WIN from tinygrad import Tensor, Device import numpy as np class TestRawShmBuffer(unittest.TestCase): - @unittest.skipIf(WIN and CI, "only fails on CI windows instance") + @unittest.skipIf(WIN, "only fails on CI windows instance") def test_e2e(self): t = Tensor.randn(2, 2, 2).realize() diff --git a/test/unit/test_tensor_data.py b/test/unit/test_tensor_data.py index 8d9a0036da4b8..8c8f4eec3e3ae 100644 --- a/test/unit/test_tensor_data.py +++ b/test/unit/test_tensor_data.py @@ -65,6 +65,10 @@ def test_data_const(self): assert dat.tolist() == 3 assert dat.shape == () + def test_const_dtype_for_uop(self): + self.assertEqual(Tensor.const(dtypes.int8, UOp.const(dtypes.float32, 1.0)).dtype, dtypes.int8) + self.assertEqual(Tensor.const(dtypes.int32, UOp.variable("x", 1, 10).bind(5)).item(), 5) + def test_data_float32(self): a = Tensor([[1,2.5],[3,4]], dtype=dtypes.float32) dat = a.data() diff --git a/tinygrad/callify.py b/tinygrad/callify.py index 85b1e11c6e466..02f866974b037 100644 --- a/tinygrad/callify.py +++ b/tinygrad/callify.py @@ -1,4 +1,5 @@ from dataclasses import dataclass, field +from tinygrad.dtype import dtypes, AddrSpace, PtrDType, ImageDType from tinygrad.uop.ops import UOp, UPat, PatternMatcher, Ops, GroupOp, graph_rewrite, track_rewrites from tinygrad.helpers import VIZ, pluralize, all_int @@ -17,10 +18,10 @@ def tag_uop(ctx:AllocCtx, x:UOp): def disk_copy_is_buffer(ctx:AllocCtx, u:UOp): # copies to disk are replaced with the disk buffer - to_disk = isinstance(u._device, str) and u._device.startswith(("DISK", "TINYFS")) + to_disk = isinstance(u.device, str) and u.device.startswith(("DISK", "TINYFS")) if to_disk: ctx.buffer_map[u] = u.empty_like() # all copies from disk/numpy are realized into a real buffer - from_creation = isinstance(u.src[0]._device, str) and any(u.src[0]._device.startswith(x) for x in ["NPY", "DISK", "PYTHON", "TINYFS"]) + from_creation = isinstance(u.src[0].device, str) and any(u.src[0].device.startswith(x) for x in ["NPY", "DISK", "PYTHON", "TINYFS"]) if from_creation: return tag_uop(ctx, u) def apply_after(ctx:AllocCtx, u:UOp): @@ -42,11 +43,11 @@ def apply_after(ctx:AllocCtx, u:UOp): def replace_contig_with_store_after(u:UOp): # can't allocate a buffer without a device (e.g., inside a CALL function body with only PARAMs) - if u._device is None: return None + if u.device is None: return None # if size is 0, remove the contig if 0 in u.shape: return u.src[0] # no real contig for DISK/TINYFS tensors, they are left alone - if isinstance(u._device, str) and u._device.startswith(("DISK", "TINYFS")): return u.rtag(None) + if isinstance(u.device, str) and u.device.startswith(("DISK", "TINYFS")): return u.rtag(None) buf = u.empty_like() return buf.after(buf.store(u.src[0])).rtag(u.tag) @@ -56,17 +57,21 @@ def replace_store_after_with_contig(u:UOp, src:UOp): if assigned_to.op is not Ops.BUFFER: return src.contiguous(tag=u.tag) def _make_buffer_view(src:UOp) -> UOp|None: - """If movement ops on src collapse to a contiguous range, return BUFFER_VIEW.reshape(src.shape). Otherwise None.""" + """If movement ops on src collapse to a contiguous range, return SLICE.reshape(src.shape). Otherwise None.""" if (offset := src.contiguous_view_offset()) is None: return None buf = src.base - if buf.op is Ops.BUFFER_VIEW: offset, buf = offset + buf.arg[1], buf.src[0] - return UOp(Ops.BUFFER_VIEW, src.dtype, (buf,), (src.numel(), offset)).reshape(src.shape) + if buf.op is Ops.SLICE: + byte_offset = buf.src[1].arg * buf.src[0].dtype.itemsize + offset * src.dtype.itemsize + buf = buf.src[0] + if byte_offset % buf.dtype.itemsize != 0: return None + offset = byte_offset // buf.dtype.itemsize + return UOp(Ops.SLICE, src.dtype, (buf, UOp.const(dtypes.weakint, offset)), src.numel()).reshape(src.shape) def contiguous_mops_to_view(c:UOp, src:UOp): - """CONTIGUOUS(MOPS(BUFFER)) → CONTIGUOUS(BUFFER_VIEW) when movement ops collapse to a contiguous range.""" + """CONTIGUOUS(MOPS(BUFFER)) → CONTIGUOUS(SLICE) when movement ops collapse to a contiguous range.""" buf = src.base - if buf.op not in {Ops.BUFFER, Ops.BUFFER_VIEW}: return None - if src.op is Ops.RESHAPE and src.src[0].op in {Ops.BUFFER, Ops.BUFFER_VIEW}: return None + if buf.op not in {Ops.BUFFER, Ops.SLICE}: return None + if src.op is Ops.RESHAPE and src.src[0].op in {Ops.BUFFER, Ops.SLICE}: return None # no symbolic shape if not all_int(c.shape): return None @@ -79,11 +84,11 @@ def contiguous_mops_to_view(c:UOp, src:UOp): x = src while x.op in GroupOp.Movement: x = x.src[0] - # NOTE: this contiguous is removed because this BUFFER_VIEW/RESHAPE has_buffer_identity + # NOTE: this contiguous is removed because this SLICE/RESHAPE has_buffer_identity if x.op is not Ops.MULTI and (view := _make_buffer_view(src)) is not None: return view.contiguous(tag=c.tag) - # for MULTI tensors, use multi_pm to resolve per-shard movement ops, then create BUFFER_VIEW on the resolved result + # for MULTI tensors, use multi_pm to resolve per-shard movement ops, then create SLICE on the resolved result if not isinstance(c.device, str): from tinygrad.schedule.multi import multi_pm resolved = graph_rewrite(src, multi_pm, name="multi_buffer_view") @@ -107,13 +112,16 @@ def transform_precompiled_call(c:UOp) -> UOp|None: subs:dict[UOp, UOp] = {} items:list[UOp] = [] for s, t in zip(srcs, targets): - while s.op is Ops.AFTER: s = s.src[0] + after_deps:list[UOp] = [] + while s.op is Ops.AFTER: + after_deps.extend(s.src[1:]) + s = s.src[0] base = s.base if base.op in {Ops.CONTIGUOUS, Ops.BUFFER} and base.shape == t.shape and base not in subs: subs[base] = t.after(t.store(base.src[0])) if base.op is Ops.CONTIGUOUS else t - items.append(s) + items.append(s.after(*after_deps) if after_deps else s) else: - items.append(t.after(t.store(s))) + items.append(t.after(t.store(s), *after_deps)) fxn = UOp.sink(*(x.substitute(subs) for x in items)) # body switches from TUPLE to SINK, so the node becomes an opaque CALL (not FUNCTION) @@ -134,7 +142,7 @@ def transform_precompiled_call(c:UOp) -> UOp|None: # resolve TUPLE+GETTUPLE (for precompiled calls) (UPat(Ops.GETTUPLE, src=(UPat(Ops.TUPLE, name="t"),), name="g"), lambda g,t: t.src[g.arg]), - # CONTIGUOUS(MOPS(BUFFER/BUFFER_VIEW)) → CONTIGUOUS(BUFFER_VIEW) when movement ops collapse to contiguous range + # CONTIGUOUS(MOPS(BUFFER/SLICE)) → CONTIGUOUS(SLICE) when movement ops collapse to contiguous range (UPat(Ops.CONTIGUOUS, src=(UPat(GroupOp.Movement, name="src"),), name="c"), contiguous_mops_to_view), # add CONTIGUOUS to tagged UOps @@ -167,8 +175,9 @@ def finalize_after(ctx:AllocCtx, x:UOp): def replace_input_buffer(ctx:AllocCtx, b:UOp): ctx.replacements.append(b) - return UOp.param(len(ctx.replacements)-1, b.dtype, b.shape, b._device, - b._min_max if b.op is Ops.BIND else None, b.src[0].arg[0] if b.op is Ops.BIND else None) + return UOp.param(len(ctx.replacements)-1, b.dtype, b.shape, b.device, + b._min_max if b.op is Ops.BIND else None, b.src[0].arg[0] if b.op is Ops.BIND else None, + b.addrspace if isinstance(b.dtype, (PtrDType, ImageDType)) else AddrSpace.GLOBAL) pm_finalize_call = PatternMatcher([ (UPat(Ops.AFTER, name="x"), finalize_after), @@ -180,8 +189,8 @@ def replace_input_buffer(ctx:AllocCtx, b:UOp): pm_replace_buf = PatternMatcher([ # replace BUFFER with PARAM for cache key normalization (UPat(Ops.BUFFER, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE)), name="b"), replace_input_buffer), - # replace BUFFER_VIEW with PARAM. this rewrite is bottom up so BUFFERs we don't need won't be in the input - (UPat(Ops.BUFFER_VIEW, src=(UPat(Ops.BUFFER),), name="b"), replace_input_buffer), + # replace SLICE with PARAM. this rewrite is bottom up so BUFFERs we don't need won't be in the input + (UPat(Ops.SLICE, src=(UPat(Ops.BUFFER), UPat(Ops.CONST, dtype=dtypes.weakint)), name="b"), replace_input_buffer), # strip value from BIND for cache key normalization, so different values hit same cache (UPat(Ops.BIND, src=(UPat(Ops.DEFINE_VAR), UPat(Ops.CONST)), name="b"), replace_input_buffer), ]) diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index beddee011359c..f2b3a28568e43 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -1,12 +1,13 @@ from typing import cast from dataclasses import replace import itertools -from tinygrad.helpers import DISABLE_FAST_IDIV, DEVECTORIZE, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC +from tinygrad.helpers import DISABLE_FAST_IDIV, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC from tinygrad.helpers import ALLOW_TF32, TracingKey, Context, panic from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, ProgramInfo from tinygrad.uop.render import pyrender from tinygrad.uop.spec import type_verify, spec_tensor, spec_program from tinygrad.renderer import Renderer, Estimates +from tinygrad.renderer.isa import ISARenderer, IselContext, PreRegAllocContext from tinygrad.dtype import dtypes # import all pattern matchers here @@ -14,13 +15,14 @@ from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing, symbolic, pm_move_where_on_load from tinygrad.uop.decompositions import get_late_rewrite_patterns, get_transcendental_patterns, pm_dtype_decomps from tinygrad.codegen.late.expander import expander, pm_pre_expander, pm_group_for_reduce -from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \ +from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_indexing, devectorize_buf_and_index, devectorize_alu, pm_reduce, \ ReduceContext, correct_load_store, pm_render, pm_add_loads, pm_make_images from tinygrad.codegen.opt.postrange import apply_opts from tinygrad.codegen.late.gater import pm_move_gates_from_index from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen, pm_mops, pm_syntactic_sugar, pm_store_ranges from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize +from tinygrad.codegen.late.regalloc import LinearScanRegallocContext, pm_regalloc_rewrite def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: if VIZ: graph_rewrite(ast, PatternMatcher([]), name="View Base AST") @@ -72,11 +74,9 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: if IMAGE and ren.target.device in {"QCOM", "CL", "PYTHON", "NULL"}: sink = graph_rewrite(sink, pm_make_images, name="create image buffers", bottom_up=True, ctx=ren.target.arch) - # devectorize (TODO: does this need opts?) - if DEVECTORIZE >= 2: pm_devectorize = sym+load_store_folding+load_store_indexing - elif DEVECTORIZE: pm_devectorize = sym+devectorize+load_store_folding+correct_load_store+load_store_indexing - else: pm_devectorize = sym+load_store_folding+correct_load_store+load_store_indexing - if DEVECTORIZE >= 0: sink = graph_rewrite(sink, pm_devectorize, ctx=ren, name="devectorize") + # devectorize + sink = graph_rewrite(sink, sym+devectorize_alu+devectorize_buf_and_index+load_store_folding+correct_load_store+load_store_indexing, + ctx=ren, name="devectorize") # lower the index dtype to a concrete int sink = graph_rewrite(sink, pm_lower_index_dtype+load_store_indexing+gep_pushing, name="lower all index dtypes") @@ -89,8 +89,8 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: supported_ops = tuple(ren.code_for_op.keys()) pm_decomp = symbolic_simple+get_late_rewrite_patterns(supported_ops, bool(DISABLE_FAST_IDIV)) pm_transcendental = symbolic_simple+get_transcendental_patterns(supported_ops, TRANSCENDENTAL>=2) - sink = graph_rewrite(sink, pm_decomp, ctx=ren.target, name="decompositions") - sink = graph_rewrite(sink, pm_dtype_decomps, ctx=(set(), ren.target), name="decomp dtypes") + sink = graph_rewrite(sink, pm_decomp, ctx=ren, name="decompositions") + sink = graph_rewrite(sink, pm_dtype_decomps, ctx=(set(), ren), name="decomp dtypes") sink = graph_rewrite(sink, pm_transcendental, name="transcendental") # move gates from unrenderable INVALID where @@ -99,12 +99,13 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: # final rules for the renderer (without sym) extra_matcher = ren.extra_matcher if ren.extra_matcher is not None else PatternMatcher([]) pm_final_rewrite = pm_decomp+pm_render+extra_matcher+pm_split_ends - sink = graph_rewrite(sink, pm_final_rewrite, ctx=ren.target, name="final rewrite") + sink = graph_rewrite(sink, pm_final_rewrite, ctx=ren, name="final rewrite") # this was the linearizer sink = graph_rewrite(sink, pm_add_control_flow, ctx=CFGContext(sink), name="add control flow", bottom_up=True) if VIZ: graph_rewrite(sink, PatternMatcher([]), name="View Output AST") + if SPEC: type_verify(sink, spec_program) # return the rewritten sink return sink @@ -115,23 +116,30 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: (UPat((Ops.IF, Ops.ENDIF)), lambda: panic(RuntimeError, "if not allowed in graph")), # gated STORE becomes IF-STORE-ENDIF. this is the only use of IF-ENDIF (UPat(Ops.STORE, name="u", src=(UPat(Ops.INDEX).or_casted(), UPat(), UPat(name="gate", dtype=dtypes.bool))), - lambda u, gate: (u, [mif:=UOp(Ops.IF, src=(gate, u.src[0])), u, UOp(Ops.ENDIF, src=(mif,))])) + lambda u, gate: ((st:=u.replace(src=u.src[0:2])), [mif:=UOp(Ops.IF, src=(gate, u.src[0])), st, UOp(Ops.ENDIF, src=(mif,))])) ]) # requires lst be toposorted. like graph rewrite, but for lines -def line_rewrite(lst:list[UOp], pm:PatternMatcher) -> list[UOp]: +def line_rewrite(lst:list[UOp], pm:PatternMatcher, ctx=None) -> list[UOp]: newlst = [] replaced: dict[UOp, UOp] = {} for u in lst: - nu = u.replace(src=tuple([replaced[x] for x in u.src])) - ret: tuple[UOp, list[UOp]] = cast(tuple[UOp, list[UOp]]|None, pm.rewrite(nu)) or (nu, [nu]) + nu = u.replace(src=tuple([replaced.get(x, x) for x in u.src])) + ret: tuple[UOp, list[UOp]] = cast(tuple[UOp, list[UOp]]|None, pm.rewrite(nu, ctx)) or (nu, [nu]) replaced[u] = ret[0] newlst.extend(ret[1]) return newlst -def do_linearize(prg:UOp, sink:UOp) -> UOp: +def do_linearize(ctx:Renderer, prg:UOp, sink:UOp) -> UOp: + if DEBUG >= 3 and sink.arg.applied_opts: print(f"{sink.arg.function_name:<25} opts: {sink.arg.applied_opts}") lst = line_rewrite(linearize(sink), pm_linearize_cleanups) - if SPEC: type_verify(lst, spec_program) + # isa renderers need to allocate registers + if isinstance(ctx, ISARenderer): + if ctx.pre_regalloc_matcher is not None: lst = line_rewrite(lst, ctx.pre_regalloc_matcher, PreRegAllocContext()) + regalloc_ctx = LinearScanRegallocContext(lst, ctx) + lst = line_rewrite(lst, pm_regalloc_rewrite, regalloc_ctx) + lst = line_rewrite(lst, ctx.post_regalloc_matcher, regalloc_ctx) + if DEBUG >= 4: print(ctx.asm_str(lst, sink.arg.function_name)) return prg.replace(src=prg.src + (UOp(Ops.LINEAR, src=tuple(lst)),)) def do_estimates(prg:UOp, sink:UOp, lin:UOp) -> UOp|None: @@ -139,8 +147,9 @@ def do_estimates(prg:UOp, sink:UOp, lin:UOp) -> UOp|None: return prg.replace(src=(sink.replace(arg=replace(sink.arg, estimates=Estimates.from_uops(lin.src, ignore_indexing=True))),)+prg.src[1:]) def do_assemble(ctx:Renderer, prg:UOp, lin:UOp) -> UOp: - binary = ctx.asm(prg, lin) src = "\n".join(str(u.arg) for u in lin.src) + if DEBUG >= 4: print(src) + binary = ctx.asm(prg, lin) return prg.replace(src=prg.src[:3]+(UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=binary))) def do_render(ctx:Renderer, prg:UOp, lin:UOp) -> UOp: @@ -149,7 +158,9 @@ def do_render(ctx:Renderer, prg:UOp, lin:UOp) -> UOp: return prg.replace(src=prg.src + (UOp(Ops.SOURCE, arg=src),), arg=new_arg) def do_compile(ctx:Renderer, prg:UOp, source:UOp) -> UOp|None: + if DEBUG >= 4: print(source.arg) lib = ctx.compiler.compile_cached(source.arg) + if DEBUG >= 7: ctx.compiler.disassemble(lib) return prg.replace(src=prg.src + (UOp(Ops.BINARY, arg=lib),)) pm_to_program = PatternMatcher([ @@ -177,7 +188,12 @@ def do_to_program(ast:UOp, renderer:Renderer) -> UOp: elif ast.op is Ops.SINK: assert isinstance(ast.arg, KernelInfo), "requires KernelInfo on arg to to_program" full_sink = full_rewrite_to_sink(ast, renderer, optimize=ast.tag is None) - prg = UOp(Ops.PROGRAM, src=(full_sink, UOp(Ops.DEVICE, arg=renderer.target.device)), arg=ProgramInfo.from_sink(full_sink)) + prog_info = ProgramInfo.from_sink(full_sink) + # instruction selection + if isinstance(renderer, ISARenderer): + full_sink = graph_rewrite(full_sink, renderer.pre_isel_matcher, ctx=itertools.count(-1, -1), name="pre instruction selection", bottom_up=True) + full_sink = graph_rewrite(full_sink, renderer.isel_matcher, ctx=IselContext(full_sink), name="instruction selection", bottom_up=True) + prg = UOp(Ops.PROGRAM, src=(full_sink, UOp(Ops.DEVICE, arg=renderer.target.device)), arg=prog_info) else: raise RuntimeError(f"can't call to_program on {ast.op}") if not isinstance(prg.arg, ProgramInfo): prg = prg.replace(arg=ProgramInfo.from_sink(prg.src[0])) prg = graph_rewrite(prg, pm_to_program, ctx=renderer, name="linearize/render") @@ -186,7 +202,7 @@ def do_to_program(ast:UOp, renderer:Renderer) -> UOp: to_program_cache: dict[tuple, UOp] = {} def to_program(ast:UOp, renderer:Renderer) -> UOp: - config = (NOOPT, DEVECTORIZE, EMULATED_DTYPES, NOLOCALS, USE_TC, IMAGE, DISABLE_FAST_IDIV, TRANSCENDENTAL, ALLOW_TF32) + config = (NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC, IMAGE, DISABLE_FAST_IDIV, TRANSCENDENTAL, ALLOW_TF32) key = (ast.key, type(renderer), renderer.target, *[x.value for x in config]) if (prg:=to_program_cache.get(key)) is None: to_program_cache[key] = prg = do_to_program(ast, renderer) return prg diff --git a/tinygrad/codegen/gpudims.py b/tinygrad/codegen/gpudims.py index 60b28c83975ad..7ffd44c2ccebf 100644 --- a/tinygrad/codegen/gpudims.py +++ b/tinygrad/codegen/gpudims.py @@ -91,7 +91,7 @@ def add_gpudims(ctx:Renderer, s:UOp): subs = {} for r in s_topo: # look for local INDEXes that are not used in the GLOBAL store, then add them as an INVALID - if r.op is Ops.STORE and (idx := r.src[0]).src[0].ptrdtype.addrspace == AddrSpace.GLOBAL: + if r.op is Ops.STORE and (idx := r.src[0]).src[0].addrspace == AddrSpace.GLOBAL: missing_locals = [all_ranges[rng] for rng in local_dims if all_ranges[rng] not in idx.ranges] if len(missing_locals): assert len(idx.src) == 2, "index has 2 sources" diff --git a/tinygrad/codegen/late/__init__.py b/tinygrad/codegen/late/__init__.py new file mode 100644 index 0000000000000..e69de29bb2d1d diff --git a/tinygrad/codegen/late/devectorizer.py b/tinygrad/codegen/late/devectorizer.py index 209b55e0b0567..8268317392b77 100644 --- a/tinygrad/codegen/late/devectorizer.py +++ b/tinygrad/codegen/late/devectorizer.py @@ -38,21 +38,24 @@ def _drop_valid_stmts(valid:UOp, idx:UOp, height:int, width:int) -> list[UOp]: def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None: idx = uop_given_valid(valid, start_idx) - if not isinstance(buf.dtype, ImageDType): return None if idx is start_idx else buf.index(idx.valid(valid), ptr=True) - - # wait for it to be image indexed before running simplification - if start_idx.dtype.count != 2: return None + return None if idx is start_idx else buf.index(idx.valid(valid), ptr=True) +def simplify_valid_image_load(buf:UOp, idx_y:UOp, idx_x:UOp, valid:UOp) -> UOp|None: + if not isinstance(buf.dtype, ImageDType): return None + start_idx = UOp.vectorize(idx_x, idx_y) + idx = uop_given_valid(valid, start_idx) drop_stmt = _drop_valid_stmts(valid, idx, buf.dtype.shape[0], buf.dtype.shape[1]) if not drop_stmt and idx is start_idx: return None new_valid = UOp.uprod(*ss) if (ss:=[s for s in valid.split_uop(Ops.AND) if s not in drop_stmt]) else None - return buf.index(idx.valid(new_valid) if new_valid is not None else idx, ptr=True) - + idx_y, idx_x = idx.gep(1), idx.gep(0) + return buf.index(idx_y.valid(new_valid), idx_x.valid(new_valid), ptr=True) if new_valid is not None else buf.index(idx_y, idx_x, ptr=True) load_store_indexing = PatternMatcher([ # image load valid idx simplification (UPat(Ops.INDEX, src=(UPat.var("buf"), invalid_gate)), lambda buf,x,i,cond: simplify_valid_load(buf, x, cond)), + (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("valid").where(UPat.var("idx_y"), UPat(arg=Invalid)), + UPat.var("valid").where(UPat.var("idx_x"), UPat(arg=Invalid)))), simplify_valid_image_load), ]) # ***** load/store grouping ***** @@ -69,7 +72,8 @@ def expand_index(ctx, buf:UOp, vec:UOp): elif dropped == best_drop: cands.append((ch, cw, cidx)) # and tiebreak with indexing complexity (ie. number of nodes) h, w, _ = cands[0] if len(cands) == 1 else min(cands, key=lambda cand: len(cand[2].gep(1).simplify().backward_slice)) - buf = buf.replace(dtype=(dtypes.imageh if dt.itemsize == 2 else dtypes.imagef)((h, w, 4))) + assert buf.op is Ops.RESHAPE + buf = buf.src[0].replace(dtype=(dtypes.imageh if dt.itemsize == 2 else dtypes.imagef)((h, w, 4))).flatten() if getenv("UNSAFE_DISABLE_MASK", 0): vec = vec.get_idx() # generate the individual indexes return UOp(Ops.STACK, buf.dtype, tuple(buf.index(vec.gep(i), ptr=True) for i in range(vec.dtype.count))) @@ -100,7 +104,7 @@ def fold_expanded_index(midx:UOp): for grp in grouped_offsets: # get the index offset for this element. using [0] is okay, because they are the same lidx = midx.src[offsets[grp[0]][0]] - if len(grp) > 1: lidx = lidx.cast(buf.ptrdtype.base.vec(len(grp)).ptr(size=buf.ptrdtype.size, addrspace=buf.ptrdtype.addrspace)) + if len(grp) > 1: lidx = lidx.cast(buf.ptrdtype.base.vec(len(grp)).ptr(size=buf.max_numel(), addrspace=buf.addrspace)) # set the idxs of the output for i,g in enumerate(grp): for oo in offsets[g]: idxs[oo] = global_offset+i @@ -109,7 +113,7 @@ def fold_expanded_index(midx:UOp): global_offset += len(grp) assert None not in idxs, f"some idxs are missing {idxs}" # this base thing is for image, we want the CAT to be a normal pointer - post_cat = UOp(Ops.PTRCAT, buf.ptrdtype.base.ptr(size=buf.ptrdtype.size, addrspace=buf.ptrdtype.addrspace).vec(global_offset), tuple(ret)) + post_cat = UOp(Ops.PTRCAT, buf.ptrdtype.base.ptr(size=buf.max_numel(), addrspace=buf.addrspace).vec(global_offset), tuple(ret)) return post_cat.gep(tuple(cast(list[int], idxs))) def cat_after_store(cat:UOp, data:UOp): @@ -130,7 +134,7 @@ def gep_on_store(gep:UOp, st:UOp): return gep.src[0].store(st.gep(new_arg)) load_store_folding = PatternMatcher([ - (UPat(Ops.INDEX, src=(UPat(Ops.STACK, src=UPat(GroupOp.Defines).or_after(name="buf")), UPat.var("vec"))), expand_index), + (UPat(Ops.INDEX, src=(UPat(Ops.STACK, src=UPat(name="buf")), UPat.var("vec"))), expand_index), (UPat(Ops.STACK, src=UPat(Ops.INDEX), name="midx"), fold_expanded_index), # GEP after LOAD (UPat(Ops.LOAD, src=(UPat(Ops.GEP, name="gep"),), name="ld", allow_any_len=True), @@ -161,7 +165,7 @@ def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp): must_divide = False elif buf.dtype.base not in (dtypes.float, dtypes.half, *dtypes.fp8s) and not isinstance(buf.dtype, ImageDType): pass - elif buf.ptrdtype.addrspace == AddrSpace.REG: + elif buf.addrspace == AddrSpace.REG: pass elif isinstance(buf.dtype, ImageDType): lengths = [4] @@ -182,7 +186,7 @@ def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp): for fold_length in lengths: if global_offset+fold_length > sz: continue lidx = buf.index((offset + global_offset).valid(mask), ptr=True) - if fold_length > 1: lidx = lidx.cast(buf.ptrdtype.base.vec(fold_length).ptr(size=buf.ptrdtype.size, addrspace=buf.ptrdtype.addrspace)) + if fold_length > 1: lidx = lidx.cast(buf.ptrdtype.base.vec(fold_length).ptr(size=buf.max_numel(), addrspace=buf.addrspace)) if ls.op is Ops.STORE: ret.append(ls.replace(src=(lidx,ls.src[1].gep(tuple(range(global_offset, global_offset+fold_length)))))) else: ret.append(ls.replace(src=(lidx,)+ls.src[1:], dtype=ls.dtype.scalar().vec(fold_length))) global_offset += fold_length @@ -195,7 +199,8 @@ def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp): def get_image_idx(idx:UOp, width:int): x, valid = idx.src[1].get_idx(), idx.src[1].get_valid() idx_x, idx_y = (x // 4) % width, x // (4*width) - return idx.replace(src=(idx.src[0], UOp.vectorize(idx_x, idx_y).valid(valid))) + assert idx.src[0].op is Ops.RESHAPE, "image idx must be on reshape" + return idx.replace(src=(idx.src[0].src[0], idx_y.valid(valid), idx_x.valid(valid))) def image_fixup(ls:UOp): # normal image load or store, with the CAST from expand_index @@ -204,8 +209,10 @@ def image_fixup(ls:UOp): return ls.replace(src=(get_image_idx(ls.src[0].src[0], dt.shape[1]),)+ls.src[1:]) # this is an unprocessed image without a cast, we should just make it a buffer - if isinstance(dt, ImageDType) and (off:=ls.src[0].src[1]).get_idx().dtype != dtypes.weakint.vec(2): - idx = ls.src[0].src[0].replace(dtype=(new_dt:=dtypes.half if dt.itemsize == 2 else dtypes.float).ptr(dt.size)).index(off) + if isinstance(dt, ImageDType) and len(ls.src[0].src) == 2: + off = ls.src[0].src[1] + assert ls.src[0].src[0].op is Ops.RESHAPE, "image idx must be on reshape" + idx = ls.src[0].src[0].src[0].replace(dtype=(new_dt:=dtypes.half if dt.itemsize == 2 else dtypes.float).ptr(dt.size)).index(off) return ls.replace(src=(idx,), dtype=new_dt).cast(dtypes.float) if ls.op is Ops.LOAD else ls.replace(src=(idx, ls.src[1].cast(new_dt))) correct_load_store = PatternMatcher([ @@ -236,7 +243,9 @@ def no_vectorized_alu(alu:UOp): return UOp(Ops.STACK, alu.dtype, alus) def no_vectorized_buf(buf:UOp): - return buf.replace(dtype=buf.ptrdtype.base.scalar().ptr(buf.ptrdtype.size*buf.ptrdtype.count, buf.ptrdtype.addrspace)).cast(buf.dtype) + # TODO: this fails on regs + #assert buf.max_numel() == buf.ptrdtype.size + return buf.replace(dtype=buf.ptrdtype.base.scalar().ptr(buf.ptrdtype.size*buf.ptrdtype.count, buf.addrspace)).cast(buf.dtype) def no_vectorized_index(buf:UOp, cast:UOp, idx:UOp, bcast:UOp|None=None): cnt = cast.dtype.count @@ -261,19 +270,18 @@ def no_vectorized_index(buf:UOp, cast:UOp, idx:UOp, bcast:UOp|None=None): no_vectorized_index), ]) -devectorize = PatternMatcher([ +devectorize_alu = PatternMatcher([ # CAST after AFTER (UPat(Ops.CAST, name="c").f(Ops.AFTER, allow_any_len=True, name="a"), lambda c,a: c.src[0].after(*a.src[1:]).cast(c.dtype)), # no ALU on vectorized dtypes (UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST), name="alu"), no_vectorized_alu), (UPat(Ops.WMMA, name="wmma"), no_vectorized_wmma), -])+devectorize_buf_and_index +]) pm_render = PatternMatcher([ # for rendering, we use explicit VECTORIZE (UPat(Ops.CONST, name='c'), lambda c: UOp(Ops.STACK, c.dtype, (UOp.const(c.dtype.scalar(), c.arg),)*c.dtype.vcount) if c.dtype.vcount > 1 else None), - (UPat(Ops.VCONST, name='c'), lambda c: UOp(Ops.STACK, c.dtype, tuple(UOp.const(c.dtype.scalar(), x) for x in c.arg))), (UPat(Ops.GEP, name='gep'), lambda gep: UOp(Ops.STACK, gep.dtype, tuple(gep.src[0].gep(x) for x in gep.arg)) if len(gep.arg) > 1 else None), (UPat(Ops.GEP, name='gep'), lambda gep: gep.src[0] if gep.src[0].dtype.vcount == 1 and gep.arg == (0,) else None), (UPat(Ops.STACK, src=(UPat(name='x'),)), lambda x: x), @@ -352,7 +360,8 @@ def add_load(idx:UOp): # add loads to non ptr index (UPat(Ops.INDEX, name="idx"), add_load), # remove loads from stores - (UPat(Ops.STORE, src=(UPat(Ops.LOAD), UPat(name="val")), name="s"), lambda s,val: s.replace(src=(s.src[0].src[0], val))), + (UPat(Ops.STORE, src=(UPat(Ops.LOAD),), allow_any_len=True, name="s"), lambda s: s.replace(src=(s.src[0].src[0],)+s.src[1:])), + (UPat(Ops.LOAD, src=(UPat(Ops.LOAD),), allow_any_len=True, name="l"), lambda l: l.replace(src=(l.src[0].src[0],)+l.src[1:])), ]) # make images @@ -369,7 +378,7 @@ def add_load(idx:UOp): def make_image(ctx, ls, buf, off): if (vcount:=buf.dtype.vcount) != 1: buf = buf.src[0] if buf.op == Ops.PARAM and not isinstance(dt:=buf.dtype, ImageDType) and (dims:=ImageDType.valid_dims(dt, ctx)): - buf = buf.replace(dtype=(dtypes.imageh if dt.base == dtypes.half else dtypes.imagef)((*dims[0], 4))) + buf = buf.replace(dtype=(dtypes.imageh if dt.base == dtypes.half else dtypes.imagef)((*dims[0], 4))).flatten() if vcount != 1: buf = UOp.vectorize(*([buf] * vcount)) if ls.op is Ops.LOAD: return ls.replace(src=(buf.index(off, ptr=True),), dtype=dtypes.float.vec(ls.dtype.vcount)).cast(dt.base) return buf.index(off, ptr=True).store(pm_imageh_store.rewrite(ls.src[1]) if dt.base == dtypes.half else ls.src[1]) diff --git a/tinygrad/codegen/late/gater.py b/tinygrad/codegen/late/gater.py index 13d0cec76256f..2f2195bd96339 100644 --- a/tinygrad/codegen/late/gater.py +++ b/tinygrad/codegen/late/gater.py @@ -9,13 +9,17 @@ (UPat.var("buf").index(UPat.var("gate").where(UPat.var("idx"), UPat(arg=Invalid))).or_casted(name="cast").store(UPat.var("data")), lambda buf,gate,idx,cast,data: buf.index(idx, ptr=True).cast(cast.dtype).store(data, gate)), + # for image idx + (UPat.var("buf").index(UPat.var("gate").where(UPat.var("idx_y"), UPat(arg=Invalid)), + UPat.var("gate").where(UPat.var("idx_x"), UPat(arg=Invalid))).or_casted(name="cast").load(name="l"), + lambda buf,gate,idx_y,idx_x,cast,l: buf.index(idx_y, idx_x, ptr=True).cast(cast.dtype).load(l.const_like(0), gate, dtype=l.dtype)), + (UPat.var("buf").index(UPat.var("gate").where(UPat.var("idx_y"), UPat(arg=Invalid)), + UPat.var("gate").where(UPat.var("idx_x"), UPat(arg=Invalid))).or_casted(name="cast").store(UPat.var("data")), + lambda buf,gate,idx_y,idx_x,cast,data: buf.index(idx_y, idx_x, ptr=True).cast(cast.dtype).store(data, gate)), + # Where after gated load becomes alt value (UPat.var("gate").where(UPat().load(UPat(), UPat.var("gate", dtype=dtypes.bool), name="l").or_casted(), UPat.var("a")), lambda gate,l,a: l.replace(src=(l.src[0], a.src[0] if a.op is Ops.CAST and a.src[0].dtype == l.dtype else a.cast(l.dtype), l.src[2])).cast(a.dtype)), (UPat.var("gate").where(UPat.var("a"), UPat().load(UPat(), ~UPat.var("gate", dtype=dtypes.bool), name="l").or_casted()), lambda gate,l,a: l.replace(src=(l.src[0], a.src[0] if a.op is Ops.CAST and a.src[0].dtype == l.dtype else a.cast(l.dtype), l.src[2])).cast(a.dtype)), - - # images use 2D INDEX now (y,x) - (UPat(Ops.INDEX, src=(UPat(), UPat((Ops.CONST, Ops.VCONST, Ops.STACK), name="vec")), name="idx"), - lambda idx,vec: idx.replace(src=(idx.src[0], vec.gep(1).cast(dtypes.int), vec.gep(0).cast(dtypes.int))) if vec.dtype.count == 2 else None), ]) diff --git a/tinygrad/codegen/late/linearizer.py b/tinygrad/codegen/late/linearizer.py index de96bba04f41f..75908503eb690 100644 --- a/tinygrad/codegen/late/linearizer.py +++ b/tinygrad/codegen/late/linearizer.py @@ -22,10 +22,10 @@ def linearize(sink:UOp) -> list[UOp]: extra = None match u.op: # the order and placement of these defines is important - case Ops.PARAM: priority, extra = -20, u.arg + case Ops.PARAM: priority, extra = -20, u.arg.slot case Ops.DEFINE_VAR: priority, extra = -19, u.arg - case Ops.DEFINE_LOCAL: priority = -18 - case Ops.DEFINE_REG: priority = -17 + case Ops.DEFINE_REG: priority = -18 + case Ops.DEFINE_LOCAL: priority = -17 case Ops.LOAD: priority = -1 # place loads early case Ops.STORE: priority = 1 # place stores late case Ops.RANGE: priority = 5 # placing RANGE is good @@ -93,4 +93,4 @@ def do_split_ends(e:UOp): pm_split_ends = PatternMatcher([ # split the ends (UPat(Ops.END, name="e"), do_split_ends), -]) \ No newline at end of file +]) diff --git a/tinygrad/codegen/late/regalloc.py b/tinygrad/codegen/late/regalloc.py new file mode 100644 index 0000000000000..deb993671a9d5 --- /dev/null +++ b/tinygrad/codegen/late/regalloc.py @@ -0,0 +1,137 @@ +import itertools +from tinygrad.helpers import dedup +from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat +from tinygrad.renderer.isa import ISARenderer, Register +from tinygrad.dtype import dtypes, PtrDType + +PSEUDO_OPS = {Ops.CONST, Ops.NOOP, Ops.AFTER, Ops.BARRIER, Ops.GROUP} + +class LinearScanRegallocContext: + # returns the uop that defines the virtual register + def vdef(self, v:Register) -> UOp: return self.uops[self.live_range[v][0]] + def __init__(self, uops:list[UOp], ren:ISARenderer): + self.uops = uops + self.ren = ren + self.idx = itertools.count() + # the label associated with each loop NOTE: this is only used post regalloc and should be removed + self.loop_label: dict[UOp, str] = {} + + # compute live ranges + self.live_range: dict[Register, list[int]] = {} + lr = self.live_range + ranges: list[Register] = [] + for i,u in enumerate(reversed(uops)): + if u.op in PSEUDO_OPS: continue + defs = u.tag if isinstance(u.tag, tuple) else () + for v in defs + tuple(s.reg for s in dedup(u.src)): + if isinstance(v, Register): lr.setdefault(v, []).insert(0, len(uops) - 1 - i) + for v in defs: + if v in lr and (n:=max((lr[rng][-1] for rng in ranges if lr[rng][0] <= lr[v][-1] < lr[rng][-1]), default=None)): lr[v].append(n) + if u.op is Ops.RANGE: ranges.append(u.reg) + + # allocate registers + self.stack_size: int = 0 + self.locals: dict[UOp, UOp] = {} + self.spills: dict[Register, UOp] = {} # mapping from virtual to stack slot + self.reals: dict[int, dict[Register, Register]] = {} # mapping from virtual to real at each program point + self.insert_before: dict[int, list[tuple[Register, Register]]] = {} # fills to be inserted at each program point + live: dict[Register, Register] = {} # mapping from virtual to real that's currently assigned to it + live_ins: list[dict[Register, Register]] = [] # mapping from virtual to real at loop entry + + def alloc(cons:tuple[Register, ...], i:int) -> Register: + live_inv = {v:k for k,v in live.items()} + # allocate the best register. Registers not in live or not used again are free and have priority, + # otherwise pick the one with the furthest next use. Regs that appear first in cons have priority in case of a tie + reg,vreg = max(((r,live_inv.get(r)) for r in cons), + key=lambda rv: next((j-i for j in ([] if rv[1] is None else lr[rv[1]]) if j >= i), len(uops))) + return live.pop(vreg) if vreg is not None else reg + + # assign register to spilled virtual and record load to be emitted before current uop, also assign it a stack slot + def fill(v:Register, i:int, cons:tuple[Register, ...]|None=None) -> Register: + if v not in self.spills: + dt = self.vdef(v).dtype + sz = dt.scalar().itemsize * dt.count if not isinstance(dt, PtrDType) else 8 + offset = self.stack_size + (sz - self.stack_size % sz) % sz + self.spills[v] = UOp.const(dtypes.int32, offset) + self.stack_size = offset + sz + r = alloc(cons if cons is not None else v.cons, i) + self.insert_before.setdefault(i, []).append((v, r)) + return r + + for i,u in enumerate(uops): + if u.op in PSEUDO_OPS: continue + # allocate uses + for s in u.src: + # HACK: cause of later hacks to lower range + if u.op is Ops.END: continue + if not isinstance(v:=s.reg, Register): continue + if v not in live: live[v] = fill(v, i) + self.reals.setdefault(i, {})[v] = live[v] + + # allocate defs + if isinstance(u.tag, tuple): + for j,v in enumerate(u.tag): + # register should only be defined once + assert isinstance(v, Register) and lr[v][0] == i + cons = v.cons + # two address instructions (src is reused by def) can only coalesce reused src. reused src goes first to get priority in case of a tiebreak + if ren.is_two_address(u) and j == 0: + uses = tuple(live.get(s.reg) for s in u.src) + cons = ((uses[0],) if uses[0] in cons else ()) + tuple(r for r in cons if r not in uses) + # HACK: cause the range is missing the comparison + live[v] = alloc(cons, i+1 if u.op is not Ops.RANGE else i) + self.reals.setdefault(i, {})[v] = live[v] + + # allocate stack array + if u.op is Ops.DEFINE_LOCAL: + self.locals[u] = UOp.const(dtypes.int32, self.stack_size) + self.stack_size += u.dtype.nbytes() + + # loop prologue, avoid loading inside the loop + if u.op is Ops.RANGE: + # we move to registers vars used in the loop sorted by next use, vars not used in the loop will not be reloaded in the epilogue + used_in_loop = [v for v in live.keys() | self.spills.keys() if any(i <= l < lr[u.reg][-1] for l in lr[v])] + sorted_uses = sorted(used_in_loop, key=lambda k: (next(l-i for l in lr[k] if l >= i), lr[k][0], k.name, k.index)) + live_in: dict[Register, Register] = {} + for v in sorted_uses: + # if all the possible registers are already in live_in there's no space for this var + if set(v.cons).issubset(live_in.values()): continue + if v not in live: live[v] = fill(v, i) + live_in[v] = live[v] + live_ins.append(live_in) + + # loop epilogue, reload registers that were live at loop entry + if u.op is Ops.END: + # TODO: if a uop is in a different reg in live out vs live in move between registers instead of loading + # TODO: don't reload if first use in loop is a load + for v,r in live_ins.pop().items(): + if v not in live or live[v] != r: live[v] = fill(v, i, (r,)) + +def regalloc_rewrite(ctx:LinearScanRegallocContext, x:UOp): + i = next(ctx.idx) + if x.op in PSEUDO_OPS: return None + nsrc = [] + for j,s in enumerate(x.src): + # v here is the virtual defined by the original s as s is the rewritten version + if i in ctx.reals and (v:=ctx.uops[i].src[j].reg) in ctx.spills: nsrc.append(ctx.ren.fill(ctx.spills[v], ctx.vdef(v), ctx.reals[i][v])) + else: nsrc.append(s) + ndefs = tuple(ctx.reals[i][v] for v in x.tag) if isinstance(x.tag, tuple) else x.tag + if x.op is Ops.DEFINE_LOCAL: nx = ctx.ren.isel_matcher.rewrite(ctx.ren.stack_pointer().index(ctx.locals[x], dtype=x.dtype, tag=ndefs)) + else: nx = x.replace(src=tuple(nsrc), tag=ndefs) + + before = [ctx.ren.fill(ctx.spills[v], ctx.vdef(v), r) for v,r in ctx.insert_before.get(i, [])] + after = [ctx.ren.spill(ctx.spills[v], nx) for v in x.tag if v in ctx.spills] if isinstance(x.tag, tuple) else [] + + # alloc/dealloc stack + if ctx.stack_size > 0: + sp = ctx.ren.stack_pointer() + offset = UOp(Ops.CONST, sp.dtype, arg=ctx.stack_size) + if i == 0: before = [ctx.ren.isel_matcher.rewrite(UOp(Ops.SUB, sp.dtype, (sp, offset), tag=sp.tag))] + before + elif i == len(ctx.uops) - 2: before += [ctx.ren.isel_matcher.rewrite(UOp(Ops.ADD, sp.dtype, (sp, offset), tag=sp.tag))] + + return nx, before + [nx] + after + +pm_regalloc_rewrite = PatternMatcher([ + (UPat({Ops.INS, Ops.RANGE, Ops.END, Ops.DEFINE_REG, Ops.DEFINE_LOCAL, Ops.PARAM, Ops.DEFINE_VAR, Ops.SPECIAL} | PSEUDO_OPS, name="x"), + regalloc_rewrite), +]) diff --git a/tinygrad/codegen/opt/postrange.py b/tinygrad/codegen/opt/postrange.py index 893f3d3b24b1a..21c2025e3200c 100644 --- a/tinygrad/codegen/opt/postrange.py +++ b/tinygrad/codegen/opt/postrange.py @@ -95,7 +95,7 @@ def shift_to(self, rng:UOp, amount:int, new_type:AxisType, top:bool=False, input if (old_sz:=rng.src[0].divides(amount)) is None: raise KernelOptError(f"{amount} can't divide {rng.src[0]} in {self.colored_shape()}") new_rng = UOp.range(amount, next(self.opt_range), new_type) if input_new_rng is None else input_new_rng - replaced_rng = rng.replace(src=(UOp.const(dtypes.int, old_sz),)) + replaced_rng = rng.replace(src=(old_sz,)) sub_axis = (new_rng * old_sz + replaced_rng) if top else (replaced_rng * amount + new_rng) self.ast = self.ast.substitute({rng:sub_axis}, name=f"shift {rng.arg[:-1]} {amount} {str(new_type).split('.')[1].lower()}") return replaced_rng, new_rng @@ -329,8 +329,8 @@ def upcasted(self) -> int: return len(self.axes_of(AxisType.UPCAST, AxisType.UNR def group_for_reduces(self) -> int: return len(self.axes_of(AxisType.GROUP_REDUCE)) def bufs_from_ast(ast:UOp, dname:str) -> list[Buffer]: - glbls = sorted([x for x in ast.backward_slice if x.op is Ops.PARAM], key=lambda x: x.arg) - return [Buffer(dname, x.ptrdtype.size, x.dtype.base) for x in glbls] + glbls = sorted([x for x in ast.backward_slice if x.op is Ops.PARAM], key=lambda x: x.arg.slot) + return [Buffer(dname, x.max_numel(), x.dtype.base) for x in glbls] def apply_opts(ast:UOp, ren:Renderer, beam:int=0) -> UOp: if ast.tag is not None: return ast diff --git a/tinygrad/codegen/simplify.py b/tinygrad/codegen/simplify.py index 29ca29f40315e..76b668272f88d 100644 --- a/tinygrad/codegen/simplify.py +++ b/tinygrad/codegen/simplify.py @@ -134,7 +134,7 @@ def reduce_collapse(red:UOp, u:UOp, pm:PatternMatcher=pm_reduce_collapse) -> UOp replaces: dict[UOp, UOp] = {} for u in included: for s in u.src: - if s in included or s in replaces or s.op in {Ops.CONST, Ops.VCONST, Ops.PARAM, Ops.DEFINE_LOCAL, Ops.DEFINE_VAR}: continue + if s in included or s in replaces or s.op in {Ops.CONST, Ops.PARAM, Ops.DEFINE_LOCAL, Ops.DEFINE_VAR}: continue replaces[s] = UOp.variable(f'in{len(replaces)}', s.vmin, s.vmax, s.dtype) collapse_fxn = u.substitute(replaces).reduce(r, arg=Ops.ADD) sink = graph_rewrite(collapse_fxn, pm, name="reduce_collapse") diff --git a/tinygrad/device.py b/tinygrad/device.py index 82f0acbe772a3..6fb5e12d732bb 100644 --- a/tinygrad/device.py +++ b/tinygrad/device.py @@ -2,12 +2,11 @@ from dataclasses import dataclass, replace from collections import defaultdict from typing import Any, Generic, TypeVar, Iterator, Generator, TYPE_CHECKING -import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re, atexit, pickle, decimal -from tinygrad.helpers import BENCHMARKS, CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored +import importlib, inspect, functools, pathlib, os, contextlib, re, atexit, pickle, decimal +from tinygrad.helpers import LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, suppress_finalizing -from tinygrad.helpers import select_by_name, select_first_inited, DEV, EMULATED_DTYPES, IMAGE, FLOAT16, TracingKey, size_to_str, Target -from tinygrad.helpers import pluralize -from tinygrad.dtype import DType, PtrDType, dtypes, _to_np_dtype +from tinygrad.helpers import select_by_name, select_first_inited, DEV, TracingKey, size_to_str, pluralize +from tinygrad.dtype import DType, PtrDType, _to_np_dtype if TYPE_CHECKING: from tinygrad.renderer import Renderer # **************** Device **************** @@ -336,47 +335,6 @@ def finalize(self): """ # override this in your device implementation -# TODO: move this to each Device -# this only tracks if the dtype is natively supported, it may be supported in the frontend using decomps -def is_dtype_supported(dtype:DType, target:Target|None=None) -> bool: - target = target or DEV.target(Device.DEFAULT) - if dtype == dtypes.bfloat16: - match target.device: - case "METAL": return not CI or BENCHMARKS - case "CUDA": return (not CI or BENCHMARKS) and target.renderer != "PTX" - case "NV": return (not CI or BENCHMARKS) and target.renderer not in ("PTX", "NAK") - case "CPU": return (not CI or BENCHMARKS) and platform.machine() in {"arm", "arm64", "aarch64", "x86_64", "amd64"} and target.renderer != "LVP" - case "AMD" | "CL" | "PYTHON" | "NULL": return True - case _: return False - if dtype in dtypes.fp8_ocp: - match target.device: - case "CUDA": return (not CI or BENCHMARKS) and target.renderer != "PTX" - case "NV": return (not CI or BENCHMARKS) and target.renderer not in ("PTX", "NAK") - case "AMD": return (not CI or BENCHMARKS) and target.arch == "gfx950" - case "PYTHON" | "NULL": return True - case _: return False - if dtype in dtypes.fp8_fnuz: return target.device in {"PYTHON", "NULL"} - if target.device == "WEBGPU": return dtype in [dtypes.bool, dtypes.char, dtypes.uchar, dtypes.short, - dtypes.ushort, dtypes.float, dtypes.int32, dtypes.uint32, dtypes.half] - # for CI GPU and OSX, cl_khr_fp16 isn't supported - # for CI LLVM, it segfaults because it can't link to the casting function - # CI CUDA architecture is sm_35 but we need at least sm_70 to run fp16 ALUs - # PYTHON supports half memoryview in 3.12+ https://github.com/python/cpython/issues/90751 - if dtype == dtypes.half: - match target.device: - case "CL": return (not CI or BENCHMARKS) and not OSX - case "QCOM": return bool(IMAGE) and bool(FLOAT16) # QCOM compiler is flaky with half - case "CUDA" | "NV": return not CI or BENCHMARKS or target.renderer == "PYTHON" - case "CPU" if target.renderer == "LLVM": return OSX - case "PYTHON": return sys.version_info >= (3, 12) - if dtype == dtypes.float64: - match target.device: - case _ if dtypes.long in EMULATED_DTYPES.tolist(dtypes): return False # double can't be bitcast to anything without long support - case "CL": return not OSX - case "NULL": return target.renderer not in ("IR3", "QCOMCL") - case "METAL" | "QCOM": return False - return True - if PROFILE: @atexit.register def finalize_profile(): diff --git a/tinygrad/engine/jit.py b/tinygrad/engine/jit.py index 8a1ae02860e62..42502bae7f9bb 100644 --- a/tinygrad/engine/jit.py +++ b/tinygrad/engine/jit.py @@ -43,7 +43,7 @@ def flush_batch(): current_batch, current_batch_devs = [], [] for si in linear.src: - if si.src[0].op is Ops.BUFFER_VIEW: continue + if si.src[0].op is Ops.SLICE: continue devs = dedup([Device[x] for b in si.src[1:] if b.op is not Ops.BIND for x in (b.device if isinstance(b.device, tuple) else (b.device,))]) graph_t = graph_class(devs[0]) if devs[0].graph is not None else None @@ -94,7 +94,7 @@ def __init__(self, linear:UOp, input_uops:tuple[UOp, ...]=()): self.runtimes: list[Any|None] = [] self.uop_replace: list[list[tuple[int, int]]] = [] for call in self.linear.src: - replace = [(p, b.arg) for p, b in enumerate(get_call_arg_uops(call)) if b.op is Ops.PARAM] + replace = [(p, b.arg.slot) for p, b in enumerate(get_call_arg_uops(call)) if b.op is Ops.PARAM] for dev_idx, (bufs, device_vars) in enumerate(unwrap_multi(call, resolve_params(call, input_uops))): self.calls.append((dev_idx, call.src[0], [b.ensure_allocated() for b in bufs], device_vars)) self.runtimes.append(get_runtime(bufs[0].device, call.src[0]) if call.src[0].op is Ops.PROGRAM else None) @@ -193,7 +193,7 @@ def _written_uops(self) -> set[UOp]: if call.op is not Ops.CALL: continue arg_uops = get_call_arg_uops(call) outs, ins = get_call_outs_ins(call) - out |= {arg_uops[k] for k in set(outs) - set(ins) if arg_uops[k].op in (Ops.BUFFER, Ops.BUFFER_VIEW)} + out |= {arg_uops[k] for k in set(outs) - set(ins) if arg_uops[k].op in (Ops.BUFFER, Ops.SLICE)} return out def __call__(self, input_uops:list[UOp], var_vals:dict[str, int]) -> ReturnType: @@ -219,10 +219,11 @@ def _prepare_jit_inputs(args, kwargs): for x in args + tuple(kwargs.values()): it = x if isinstance(x, (tuple,list)) else x.values() if isinstance(x, dict) else [] tensors += [t for t in it if t.__class__ is Tensor and not any(t is y for y in tensors)] + def get_input_uops() -> list[UOp]: return flatten([t.uop.src if t.uop.op is Ops.MULTI else [t.uop] for t in tensors]) + # TODO: drop the CONST branch once all CONST are deviceless + if any(u.device is None or u.base.op is Ops.CONST for u in get_input_uops()): raise JitError("JIT inputs must be real buffers; use .clone()") if len(unrealized_tensors := [x for x in tensors if not x.uop.is_realized]): Tensor.realize(*unrealized_tensors) - input_uops: list[UOp] = flatten([t.uop.src if t.uop.op is Ops.MULTI else [t.uop] for t in tensors]) - if any(u.base.op is Ops.CONST for u in input_uops): - raise JitError("JIT inputs cannot be const, create a buffer with .contiguous()") + input_uops = get_input_uops() # collect buffer UOps (including MultiBuffer) input_buf_uops: list[UOp] = [u.base for u in input_uops if u.base.realized is not None] if len(set(input_buf_uops)) != len(input_buf_uops): raise JitError("duplicate inputs to JIT") diff --git a/tinygrad/engine/realize.py b/tinygrad/engine/realize.py index 18927da39dfaa..b952df66aaaad 100644 --- a/tinygrad/engine/realize.py +++ b/tinygrad/engine/realize.py @@ -18,7 +18,7 @@ def get_call_arg_uops(call:UOp) -> tuple[UOp, ...]: return tuple(s for s in call def get_call_outs_ins(call:UOp) -> tuple[tuple[int, ...], tuple[int, ...]]: ast = call.src[0] if ast.op is Ops.PROGRAM: return tuple(ast.arg.outs), tuple(ast.arg.ins) - if ast.op in (Ops.COPY, Ops.BUFFER_VIEW): return (0,), (1,) + if ast.op in (Ops.COPY, Ops.SLICE): return (0,), (1,) if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec": return (0,), tuple(range(1, len(get_call_arg_uops(call)))) return (), () @@ -27,7 +27,9 @@ def _uop_sz_to_str(uop:UOp) -> str: return size_to_str(sym_infer(prod(uop.shape) ast, arg_uops = call.src[0], get_call_arg_uops(call) if ast.op is Ops.PROGRAM: return ast.arg.name - if ast.op is Ops.BUFFER_VIEW: return colored(f"view {_uop_sz_to_str(arg_uops[0]):>10} @ {ast.arg[1] * arg_uops[1].dtype.itemsize:<10d}", "yellow") + if ast.op is Ops.SLICE: + offset = ast.src[1].arg * arg_uops[1].dtype.itemsize + return colored(f"view {_uop_sz_to_str(arg_uops[0]):>10} @ {offset:<10d}", "yellow") if ast.op is Ops.COPY: return colored(f"copy {_uop_sz_to_str(arg_uops[0]):>10}, {bufs[0].device[:7]:>7s} <- {bufs[1].device[:7]:7s}", "yellow") if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec": return colored(f"enc/dec {_uop_sz_to_str(arg_uops[0])}", "yellow") if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph": return colored(f"batched {len(ast.src[0].src)}", "cyan") @@ -109,9 +111,6 @@ def try_exec(local_size): def get_runtime(device:str, ast:UOp, cache=True): assert ast.op is Ops.PROGRAM and isinstance(ast.arg, ProgramInfo), "get_runtime should only be called with a PROGRAM ast" if (runtime:=runtime_cache.get(key:=(ast.key, device))) is None: - if DEBUG >= 3 and ast.src[0].arg.applied_opts: print(ast.src[0].arg.applied_opts) - if DEBUG >= 4: print(ast.src[3].arg) - if DEBUG >= 7: Device[device].compiler.disassemble(ast.src[4].arg) runtime = Device[device].runtime(ast.arg.function_name, ast.src[4].arg, *ast.arg.aux, runtimevars=ast.arg.runtimevars, prg=ast) if cache: runtime_cache[key] = runtime return runtime @@ -138,8 +137,8 @@ class ExecContext: cache: bool = True def _resolve(b:UOp, inputs:tuple[UOp, ...]) -> UOp: - if b.op in (Ops.BUFFER_VIEW, Ops.MSELECT) and b.src[0].op is Ops.PARAM: return b.replace(src=(inputs[b.src[0].arg], *b.src[1:])) - return inputs[b.arg] if b.op is Ops.PARAM else b + if b.op in (Ops.SLICE, Ops.MSELECT) and b.src[0].op is Ops.PARAM: return b.replace(src=(inputs[b.src[0].arg.slot], *b.src[1:])) + return inputs[b.arg.slot] if b.op is Ops.PARAM else b def resolve_params(call:UOp, inputs:tuple[UOp, ...]) -> list[UOp]: return [_resolve(b, inputs) for b in get_call_arg_uops(call)] def unwrap_multi(call:UOp, resolved:list[UOp]) -> Iterator[tuple[list[Buffer], dict[str, int]]]: @@ -152,7 +151,7 @@ def unwrap_multi(call:UOp, resolved:list[UOp]) -> Iterator[tuple[list[Buffer], d def exec_view(ctx:ExecContext, call:UOp, ast:UOp) -> float|None: resolved = resolve_params(call, ctx.input_uops) bufs = [cast(Buffer, b.buffer) for b in resolved] - bv = bufs[1].view(resolved[0].arg, ast.dtype, ast.arg[1]*bufs[1].dtype.itemsize) + bv = bufs[1].view(resolved[0].arg, ast.dtype, ast.src[1].arg*bufs[1].dtype.itemsize) with track_stats(ctx, call, bv.device, [bv, bufs[1]], ctx.var_vals): buffers[resolved[0]] = bv return None @@ -178,7 +177,7 @@ def exec_kernel(ctx:ExecContext, call:UOp, ast:UOp) -> float|None: rt = get_runtime(device:=bufs[0].device, ast, cache=ctx.cache) global_size, local_size = ast.arg.launch_dims(var_vals) with track_stats(ctx, call, device, prg_bufs, var_vals) as tm: - et = tm[0] = rt(*[b._buf for b in prg_bufs], global_size=global_size, local_size=local_size, vals=ast.arg.vals(var_vals), + et = tm[0] = rt(*[b.get_buf(device) for b in prg_bufs], global_size=global_size, local_size=local_size, vals=ast.arg.vals(var_vals), wait=ctx.wait, timeout=ctx.timeout) return et @@ -234,7 +233,7 @@ def _validate(call:UOp, sink:UOp) -> UOp: ]) pm_exec = PatternMatcher([ - (UPat(Ops.CALL, src=(UPat(Ops.BUFFER_VIEW, name="ast"),), name="call", allow_any_len=True), exec_view), + (UPat(Ops.CALL, src=(UPat(Ops.SLICE, name="ast"),), name="call", allow_any_len=True), exec_view), (UPat(Ops.CALL, src=(UPat(Ops.COPY, name="ast"),), name="call", allow_any_len=True), exec_copy), (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="ast"),), name="call", allow_any_len=True), exec_kernel), (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="encdec", name="ast"),), name="call", allow_any_len=True), exec_encdec), @@ -242,14 +241,13 @@ def _validate(call:UOp, sink:UOp) -> UOp: (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="validate", name="ast"),), name="call", allow_any_len=True), exec_validate), ]) -if getenv("HCQ2"): - from extra.hcq2.hcq2 import pm_hcq_exec - pm_exec = pm_hcq_exec + pm_exec - -def compile_linear(linear:UOp, beam=0, validate=False) -> UOp: +def compile_linear(linear:UOp, beam:int|None=None, validate=False) -> UOp: if validate: linear = graph_rewrite(linear, pm_validate, name="validate", walk=True) - if (beam_val:=(beam or BEAM.value)) >= 1: linear = graph_rewrite(linear, pm_beam, ctx=beam_val, 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"): + from extra.hcq2.hcq2 import hcq_schedule + linear = hcq_schedule(linear) return graph_rewrite(linear, pm_optimize_local_size, name="optimize local size", walk=True) def run_linear(linear:UOp, var_vals:dict[str, int]|None=None, input_uops:tuple[UOp, ...]=(), update_stats=True, jit=False, wait=False): @@ -263,4 +261,5 @@ def time_call(call:UOp, var_vals:dict[str, int]|None=None, timeout:int|None=None else: 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) + call = compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0).src[0] return cast(float, pm_exec.rewrite(call, ExecContext(var_vals or {}, update_stats=False, wait=True, timeout=timeout, cache=False))) diff --git a/tinygrad/function.py b/tinygrad/function.py index ec8c1154c810d..0f71b909233be 100644 --- a/tinygrad/function.py +++ b/tinygrad/function.py @@ -40,7 +40,7 @@ def __call__(self, *args, **kwargs) -> ReturnType: params = get_state_dict((args, kwargs), tensor_type=(Tensor, UOp)).values() # deduplicate input_uops, keeping the first occurrence index for each unique uop - call_uops: list[UOp] = dedup([(t.uop if isinstance(t, Tensor) else t) for t in params]) + call_uops: list[UOp] = dedup([u for t in params if not ((u:=(t.uop if isinstance(t, Tensor) else t)).base.op is Ops.CONST and u.device is None)]) # disable realize/schedule while this is running # run it and do surgery later @@ -84,7 +84,7 @@ def __call__(self, *args, **kwargs) -> ReturnType: precompile_backward=self.precompile_backward) if DEBUG >= 2: - #signature = [(x._shape, x.dtype, x._device) for x in call_uops] + #signature = [(x._shape, x.dtype, x.device) for x in call_uops] print(" "*_function.depth+f"function {uret.key.hex()[:8]} in {(time.perf_counter()-st)*1000:8.2f} ms: {name}") # with sig {signature}") if isinstance(ret, tuple): diff --git a/tinygrad/gradient.py b/tinygrad/gradient.py index 4612e770c9744..c971f7725cb7f 100644 --- a/tinygrad/gradient.py +++ b/tinygrad/gradient.py @@ -16,20 +16,24 @@ def broadcast_to_input(x): return x.reshape(x.shape+(1,)*(len(ret.src[0].shape)- def _compact_params(body:UOp, all_args:tuple[UOp, ...]) -> tuple[UOp, tuple[UOp, ...]]: """Remove unused PARAMs from body and return compacted (body, args).""" - used = sorted({p.arg: p for p in body.toposort() if p.op is Ops.PARAM}.items()) - return body.substitute({p: p.replace(arg=j) for j,(_, p) in enumerate(used)}, walk=True), tuple(all_args[i] for i,_ in used) + used = sorted({p.arg.slot: p for p in body.toposort() if p.op is Ops.PARAM}.items()) + body = body.substitute({p: p.replace(arg=dataclasses.replace(p.arg, slot=j)) for j,(_, p) in enumerate(used)}, walk=True) + return body, tuple(all_args[i] for i,_ in used) def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]: fxn, args = k.src[0], k.src[1:] if k.arg.grad_fxn is not None: + # put const on a device, also TODO why do we still have NOOP... + def on_dev(g, i): return g.clone(device=args[i].device if k.op is Ops.CALL else k.device) if g.device is None else g if ctx.op is Ops.TUPLE: - real = [g for g in ctx.src if g.op is not Ops.NOOP] + real = [on_dev(g, i) for i,g in enumerate(ctx.src) if g.op is not Ops.NOOP] return (None,) + (k.arg.grad_fxn(*real, call=k) if len(real) > 1 else k.arg.grad_fxn(real[0], k)) - return (None,) + k.arg.grad_fxn(ctx, k) + return (None,) + k.arg.grad_fxn(on_dev(ctx, 0), k) assert fxn.op is Ops.TUPLE, f"expected TUPLE body for gradient, got {fxn.op}" - params = {x.arg:x for x in fxn.toposort(enter_calls=False) if x.op == Ops.PARAM} + params = {x.arg.slot:x for x in fxn.toposort(enter_calls=False) if x.op == Ops.PARAM} grad_args = ctx.src - root_grad = UOp(Ops.TUPLE, src=tuple(UOp(Ops.NOOP) if g.op is Ops.NOOP else g.param_like(len(args)+i) for i,g in enumerate(grad_args))) + root_grad = UOp(Ops.TUPLE, src=tuple(UOp(Ops.NOOP) if g.op is Ops.NOOP else + g if g.base.op is Ops.CONST and g.device is None else g.param_like(len(args)+i) for i,g in enumerate(grad_args))) grads = compute_gradient(fxn, root_grad, set(params.values())) # for precompiled calls, substitute forward outputs with params so intermediates aren't recomputed fwd_subs = {src: src.param_like(len(args)+len(grad_args)+i) for i, src in enumerate(fxn.src)} if k.arg.precompile else {} @@ -70,7 +74,7 @@ def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]: (ctx.cast(sum_acc_dtype(ctx.dtype))._rop(Ops.ADD, tuple(i for i,(s,n) in enumerate(zip(ret.src[0].shape, ret.shape)) if s!=n)) .cast(ctx.dtype), None)), (UPat(Ops.PAD, name="ret"), lambda ctx, ret: (ctx.shrink(tuple([(p[0], s+p[0]) for s,p in zip(ret.src[0].shape, ret.marg)])), None, None)), - (UPat(Ops.SHRINK, name="ret"), lambda ctx, ret: (ctx.pad(tuple([(p[0], s-p[1]) for s,p in zip(ret.src[0].shape, ret.marg)])), None, None)), + (UPat(Ops.SHRINK, name="ret"), lambda ctx, ret: (ctx.pad(tuple([(p[0], s-p[0]-p[1]) for s,p in zip(ret.src[0].shape, ret.marg)])), None, None)), (UPat(Ops.PERMUTE, name="ret"), lambda ctx, ret: (ctx.permute(argsort(ret.marg)),)), (UPat(Ops.FLIP, name="ret"), lambda ctx, ret: (ctx.flip([i for i,x in enumerate(ret.marg) if x]),)), (UPat(Ops.COPY, name="ret"), lambda ctx, ret: (ctx.copy_to_device(ret.src[0].device), None)), diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index db2e462d7128f..7b515e9421e2e 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -14,7 +14,6 @@ def prod(x:Iterable[T]) -> T|int: return functools.reduce(operator.mul, x, 1) # NOTE: helpers is not allowed to import from anything else in tinygrad OSX, WIN = platform.system() == "Darwin", sys.platform == "win32" -CI, BENCHMARKS = os.getenv("CI", "") != "", os.getenv("RUNNER_ENVIRONMENT", "") == "self-hosted" ARCH_X86 = any(x in platform.processor() for x in ("Intel", "i386", "x86_64")) BASEDIR = pathlib.Path(__file__).parent @@ -99,8 +98,8 @@ def get_child(obj, key): def word_wrap(x, wrap=80): if len(ansistrip(x)) <= wrap: return x if len(lines:=x.splitlines()) > 1: return "\n".join(word_wrap(line, wrap) for line in lines) - i = 0 - while len(ansistrip(x[:i])) < wrap and i < len(x): i += 1 + i = vis = 0 + while vis < wrap and i < len(x): i, vis = (i + m.end(), vis) if (m:=re.match('\x1b\\[(K|.*?m)', x[i:])) is not None else (i+1, vis+1) return x[:i] + "\n" + word_wrap(x[i:], wrap) def pad_bytes(b:bytes, align:int) -> bytes: return b + b'\x00' * ((align - (len(b) % align)) % align) @@ -120,7 +119,7 @@ def strides_for_shape(shape:tuple[T, ...]) -> tuple[T, ...]: # returns the axes to create new_shape if new_shape can be created by combining axis from old_shape def get_contraction(old_shape:tuple[T, ...], new_shape:tuple[T, ...]) -> list[list[int]]|None: # T is sint acc_old, acc_new = list(itertools.accumulate(old_shape, operator.mul)), list(itertools.accumulate(new_shape, operator.mul)) - try: split = [acc_old.index(acc)+1 if acc != 1 else 0 for acc in acc_new] + try: split = [0 if isinstance(acc, int) and acc == 1 else acc_old.index(acc)+1 for acc in acc_new] except ValueError: return None return [list(range(st,ed)) for st,ed in zip([0]+split[:-1], split[:-1]+[len(old_shape)])] @@ -239,7 +238,7 @@ def target(self, dev:str, **kwargs) -> Target: TRANSCENDENTAL, NOLOCALS = ContextVar("TRANSCENDENTAL", 1), ContextVar("NOLOCALS", 0) SPLIT_REDUCEOP, NO_MEMORY_PLANNER, LRU = ContextVar("SPLIT_REDUCEOP", 1), ContextVar("NO_MEMORY_PLANNER", 0), ContextVar("LRU", 1) RING, ALL2ALL, ALLREDUCE_CAST = ContextVar("RING", 1), ContextVar("ALL2ALL", 0), ContextVar("ALLREDUCE_CAST", 1) -CACHELEVEL, IGNORE_BEAM_CACHE, DEVECTORIZE = ContextVar("CACHELEVEL", 2), ContextVar("IGNORE_BEAM_CACHE", 0), ContextVar("DEVECTORIZE", 1) +CACHELEVEL, IGNORE_BEAM_CACHE = ContextVar("CACHELEVEL", 2), ContextVar("IGNORE_BEAM_CACHE", 0) VALIDATE_WITH_CPU = ContextVar("VALIDATE_WITH_CPU", 0) # TODO: this is broken for some indexing DISABLE_FAST_IDIV = ContextVar("DISABLE_FAST_IDIV", 1) @@ -268,6 +267,8 @@ def target(self, dev:str, **kwargs) -> Target: SCACHE = ContextVar("SCACHE", 1) # allow use of atomics for embedding backward USE_ATOMICS = ContextVar("USE_ATOMICS", 0) +# don't allow broadcast +DISALLOW_BROADCAST = ContextVar("DISALLOW_BROADCAST", 1) @dataclass(frozen=True) class Metadata: @@ -456,11 +457,11 @@ def fetch(url:str, name:pathlib.Path|str|None=None, subdir:str|None=None, gunzip fp = _ensure_downloads_dir() / (subdir or "") / ((name or hashlib.md5(url.encode('utf-8')).hexdigest()) + hh + (".gunzip" if gunzip else "")) if not fp.is_file() or not allow_caching or (sha256 and hashlib.sha256(fp.read_bytes()).hexdigest() != sha256): (_dir := fp.parent).mkdir(parents=True, exist_ok=True) - with urllib.request.urlopen(urllib.request.Request(url, headers={"User-Agent": "tinygrad 0.12.0", **headers}), timeout=10) as r: + with urllib.request.urlopen(urllib.request.Request(url, headers={"User-Agent": "tinygrad 0.13.0", **headers}), timeout=10) as r: assert r.status in {200, 206}, r.status length = int(r.headers.get('content-length', 0)) if not gunzip else None readfile = gzip.GzipFile(fileobj=r) if gunzip else r - progress_bar:tqdm = tqdm(total=length, unit='B', unit_scale=True, desc=f"{url}", disable=CI) + progress_bar:tqdm = tqdm(total=length, unit='B', unit_scale=True, desc=f"{url}") h = hashlib.sha256() if sha256 else None with tempfile.NamedTemporaryFile(dir=_dir, delete=False) as f: while chunk := readfile.read(16384): @@ -529,9 +530,10 @@ def flat_mv(mv:memoryview): return mv if len(mv) == 0 else mv.cast("B", shape=(m # *** tqdm class tqdm(Generic[T]): - def __init__(self, iterable:Iterable[T]|None=None, desc:str='', disable:bool=False, + def __init__(self, iterable:Iterable[T]|None=None, desc:str='', disable:bool|None=False, unit:str='it', unit_scale=False, total:int|None=None, rate:int=100): - self.iterable, self.disable, self.unit, self.unit_scale, self.rate = iterable, disable, unit, unit_scale, rate + self.disable = not sys.stderr.isatty() if disable is None else disable + self.iterable, self.unit, self.unit_scale, self.rate = iterable, unit, unit_scale, rate self.st, self.i, self.n, self.skip, self.t = time.perf_counter(), -1, 0, 1, getattr(iterable, "__len__", lambda:0)() if total is None else total self.set_description(desc) self.update(0) diff --git a/tinygrad/llm/model.py b/tinygrad/llm/model.py index 6278e91befc26..fb8cfe7646ac8 100644 --- a/tinygrad/llm/model.py +++ b/tinygrad/llm/model.py @@ -30,7 +30,7 @@ def pairwise_topk(x: Tensor, k: int) -> tuple[Tensor, Tensor]: vals = Tensor.arange(n, device=x.device).reshape(1,1,n).cast(x.dtype).expand(x.shape) cmp = (x.unsqueeze(-1) > x.unsqueeze(-2)) | ((x.unsqueeze(-1) == x.unsqueeze(-2)) & \ (Tensor.arange(n, device=x.device).reshape(1,1,n,1) < Tensor.arange(n, device=x.device).reshape(1,1,1,n))) - sel = Tensor.zeros_like(x).scatter(-1, cmp.sum(axis=-1).cast('int32'), vals)[:,:,n-k:].cast('int32') + sel = x.const_like(0).scatter(-1, cmp.sum(axis=-1).cast('int32'), vals)[:,:,n-k:].cast('int32') return x.gather(-1, sel), sel @dataclass(frozen=True) @@ -177,7 +177,8 @@ def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor: # NOTE: this mask is causal_lower_right, not the causal_upper_left generated by is_casual = True # TODO: this if statement should be removed and it shouldn't generate extra kernels - mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, device=x.device).triu(start_pos+1) if resolve(T != 1) else None + mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, device=x.device, buffer=False).triu(start_pos+1) \ + if resolve(T != 1) else None attn = q.scaled_dot_product_attention(k, v, attn_mask=mask, enable_gqa=True) # (B,H,T,Hd) attn = attn.transpose(1, 2).reshape(B, T, -1) # back to (B,T,D) return self.attn_output(attn if not self.config.attn_output_gate else (attn * gate.sigmoid())) @@ -222,7 +223,8 @@ def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor: k = Tensor(self.cache_k.uop.after(self.cache_k[:, :, start_pos:start_pos+T, :].uop.store(k_store.uop)))[:, :, 0:start_pos+T, :] v = k[..., :self.config.kv_lora_rank] - mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, device=x.device).triu(start_pos+1) if resolve(T != 1) else None + mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, device=x.device, buffer=False).triu(start_pos+1) \ + if resolve(T != 1) else None attn = q @ k.transpose(-1, -2) * (1.0 / self.config.head_dim ** 0.5) if mask is not None: attn = attn + mask attn = attn.softmax(-1) @@ -282,8 +284,8 @@ def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor: # recurrent state can't be partially reused after divergence, force a full rebuild def _state_reset_ops(self): - return [self.conv_state.assign(Tensor.zeros_like(self.conv_state)), - self.recurrent_state.assign(Tensor.zeros_like(self.recurrent_state))] if hasattr(self, "conv_state") else [] + return [self.conv_state.assign(self.conv_state.const_like(0)), + self.recurrent_state.assign(self.recurrent_state.const_like(0))] if hasattr(self, "conv_state") else [] def _reusable_prefix_len(self, prefix_len:int, cached_len:int) -> int: return 0 if prefix_len != cached_len else prefix_len def _init_state(self, x): @@ -396,7 +398,7 @@ def generate(self, tokens:list[int], chunk_size:int=32, temperature:float=0.0): v_start_pos = UOp.variable("start_pos", 0, self.max_context-1) v_toks = UOp.variable("toks", 1, chunk_size) # TODO: use UOp.variable for temperature once float variables are supported - temp = Tensor(temperature).contiguous() + temp = Tensor([temperature]) # assign all input tokens once, then slice from start_pos for the model call t = Tensor(tokens + [0] * (self.max_context - len(tokens)), dtype="int32").reshape(1, self.max_context) # recompute start_pos from what's currently valid in the caches diff --git a/tinygrad/mixin/__init__.py b/tinygrad/mixin/__init__.py index 8ce45cda23a05..0e5b79e86ed8b 100644 --- a/tinygrad/mixin/__init__.py +++ b/tinygrad/mixin/__init__.py @@ -6,11 +6,12 @@ from tinygrad.mixin.reduce import ReduceMixin from tinygrad.uop import Ops from tinygrad.uop.ops import _broadcast_shape, resolve, smax, smin, identity_element +from tinygrad.device import canonicalize_device from tinygrad.dtype import ConstType, DType, DTypeLike, Invalid, InvalidType, PtrDType, PyConst, dtypes, least_upper_dtype, sum_acc_dtype, to_dtype from tinygrad.helpers import all_int, argfix, ceildiv, flatten, flat_to_grouped, make_tuple, prod, resolve_pool_pads, round_up if TYPE_CHECKING: - from tinygrad.uop.ops import sint + from tinygrad.uop.ops import sint, UOp ReductionStr = Literal["mean", "sum", "none"] @@ -18,14 +19,17 @@ class OpMixin(ElementwiseMixin, ReduceMixin): @staticmethod def unique_const(fill_value:ConstType, **kwargs): raise NotImplementedError("creation helpers are only supported on Tensor and UOp") + @staticmethod + def const(dtype, b, device=None): raise NotImplementedError("creation helpers are only supported on Tensor and UOp") @classmethod - def full(cls, shape:tuple[sint, ...], fill_value:ConstType, **kwargs) -> Self: + def full(cls, shape:tuple[sint, ...], fill_value:ConstType|UOp, dtype:DTypeLike|None=None, + device:str|tuple[str, ...]|None=None, buffer=True) -> Self: """ Creates a tensor with the given shape, filled with the given value. You can pass in `dtype` and `device` keyword arguments to control the data type and device of the tensor. - Additionally, all other keyword arguments are passed to the constructor of the tensor. + Pass `buffer=False` to get a broadcast const value instead of a materialized buffer. ```python exec="true" source="above" session="tensor" result="python" print(Tensor.full((2, 3), 42).numpy()) @@ -34,7 +38,13 @@ def full(cls, shape:tuple[sint, ...], fill_value:ConstType, **kwargs) -> Self: print(Tensor.full((2, 3), False).numpy()) ``` """ - return cls.unique_const(fill_value, **kwargs).reshape((1,)*len(new_shape := argfix(shape))).expand(new_shape) + from tinygrad.uop.ops import UOp + new_shape = argfix(shape) + dt = to_dtype(dtype) if dtype is not None else None + if isinstance(fill_value, UOp): val = cls.const(dt or fill_value.dtype, fill_value) + else: val = cls.const(dt or dtypes.from_py(fill_value), fill_value, None if buffer else canonicalize_device(device)) + val = val.reshape((1,)*len(new_shape)).expand(new_shape) + return val.clone(device=device) if buffer else val @classmethod def invalids(cls, *shape, **kwargs) -> Self: @@ -45,7 +55,8 @@ def invalids(cls, *shape, **kwargs) -> Self: Eventually Tensor.empty will be replaced by this. """ - return cls.full(argfix(*shape), Invalid, **kwargs) + new_shape = argfix(*shape) + return cls.unique_const(Invalid, **kwargs).reshape((1,)*len(new_shape)).expand(new_shape) @classmethod def zeros(cls, *shape, **kwargs) -> Self: @@ -111,8 +122,8 @@ def arange(cls, start, stop=None, step=1, **kwargs) -> Self: lo, hi = (start, stop-step) if step > 0 else (stop-step, start) if lo < (dt:=to_dtype(dtype)).min or dt.max < hi: raise OverflowError(f"arange [{start}, {stop}) is not representable in dtype {dtype}") # NOTE: this matches numpy, torch raises RuntimeError if stop-start and step have different signs - if (output_len:=ceildiv(stop-start, step)) <= 0: return cls.full((0,), 0, dtype=dtype, **kwargs) - return (cls.full((output_len,), step, dtype=dtype, **kwargs)._cumalu(0, Ops.ADD) + (start - step)).cast(dtype) + if (output_len:=ceildiv(stop-start, step)) <= 0: return cls.full((0,), 0, dtype=dtype, buffer=False, **kwargs) + return (cls.full((output_len,), step, dtype=dtype, buffer=False, **kwargs)._cumalu(0, Ops.ADD) + (start - step)).cast(dtype) @classmethod def linspace(cls, start:int|float, stop:int|float, steps:int, **kwargs) -> Self: @@ -131,11 +142,24 @@ def linspace(cls, start:int|float, stop:int|float, steps:int, **kwargs) -> Self: """ if steps < 0: raise ValueError("number of steps must be non-negative") if (dtype := to_dtype(kwargs.pop("dtype", dtypes.default_float))) == dtypes.bool: raise ValueError("linspace with bool dtype is not supported") - if steps == 1: return cls.full((1,), start, dtype=dtype, **kwargs) + if steps == 1: return cls.full((1,), start, dtype=dtype, buffer=False, **kwargs) return (start + cls.arange(steps, dtype=dtypes.default_float, **kwargs) * ((stop - start) / (steps - 1))).cast(dtype) @classmethod def eye(cls, n:int, m:int|None=None, dtype:DTypeLike|None=None, device:str|tuple[str, ...]|None=None) -> Self: + """ + Returns a 2-D tensor with `n` rows and `m` columns, with ones on the diagonal and zeros elsewhere. + + You can pass in `dtype` and `device` keyword arguments to control the data type and device of the tensor. + + ```python exec="true" source="above" session="tensor" result="python" + print(Tensor.eye(3).numpy()) + ``` + + ```python exec="true" source="above" session="tensor" result="python" + print(Tensor.eye(2, 4).numpy()) + ``` + """ m_ = n if m is None else m if n < 0 or m_ < 0: raise ValueError(f"cannot have negative {n=}, {m_=}") out_dtype = to_dtype(dtype) if dtype is not None else dtypes.default_float @@ -166,7 +190,7 @@ def triu(self, diagonal:sint=0) -> Self: print(t.triu(diagonal=-1).numpy()) ``` """ - return self._tri(self.shape[-2], self.shape[-1], diagonal, self.device).where(self, self.zeros_like()) + return self._tri(self.shape[-2], self.shape[-1], diagonal, self.device).where(self, self.const_like(0)) def tril(self, diagonal:sint=0) -> Self: """ @@ -189,7 +213,7 @@ def tril(self, diagonal:sint=0) -> Self: print(t.tril(diagonal=-1).numpy()) ``` """ - return self._tri(self.shape[-2], self.shape[-1], diagonal+1, self.device).where(self.zeros_like(), self) + return self._tri(self.shape[-2], self.shape[-1], diagonal+1, self.device).where(self.const_like(0), self) # ***** random ***** @@ -218,7 +242,7 @@ def _bits_to_rand(bits, shape:tuple[int, ...], dtype:DType): _, nmant = dtypes.finfo(dtype) uint_dtype = {1: dtypes.uint8, 2: dtypes.uint16, 4: dtypes.uint32, 8: dtypes.uint64}[dtype.itemsize] uint_bits = bits.bitcast(uint_dtype) - float_one_bits = uint_bits.ones_like(dtype=dtype).bitcast(uint_dtype) + float_one_bits = uint_bits.const_like(1).cast(dtype).bitcast(uint_dtype) return uint_bits.rshift(dtype.bitsize - nmant).bitwise_or(float_one_bits).bitcast(dtype)[:prod(shape)].sub(1).reshape(shape) def _pad_constant(self, pX, value:ConstType) -> Self: @@ -230,12 +254,14 @@ def _pad_constant(self, pX, value:ConstType) -> Self: base = MovementMixin.pad(X, pads) if value == 0: return base base = base.cast(least_upper_dtype(base.dtype, dtypes.from_py(value))) - return MovementMixin.pad(X.ones_like(dtype=dtypes.bool), pads).where(base, base.full_like(value)) + return MovementMixin.pad(X.const_like(1).cast(dtypes.bool), pads).where(base, base.const_like(value)) def _pad_circular(self, pX:tuple[tuple[sint, sint], ...]) -> Self: - if any(pB>sh or pA>sh for (pB,pA),sh in zip(pX, self.shape)): raise ValueError('Padding value causes wrapping around more than once.') - if any(pB<0 or pA<0 for pB,pA in pX): raise NotImplementedError("Negative pads with circular pads is not supported") - orig_shape, X = self.shape, self.repeat(tuple(1 + bool(pB) + bool(pA) for pB,pA in pX)) + # shrink first for negative pads, then wrap the non-negative remainder + X = self.shrink(tuple((-smin(pB,0), smin(pA+sh,sh)) for (pB,pA),sh in zip(pX, self.shape))) + pX = tuple((smax(pB,0), smax(pA,0)) for pB,pA in pX) + if any(pB>sh or pA>sh for (pB,pA),sh in zip(pX, X.shape)): raise ValueError('Padding value causes wrapping around more than once.') + orig_shape, X = X.shape, X.repeat(tuple(1 + bool(pB) + bool(pA) for pB,pA in pX)) return X.shrink(tuple((0 if pB == 0 else osh-pB, xsh if pA == 0 else xsh-osh+pA) for (pB,pA),osh,xsh in zip(pX, orig_shape, X.shape))) def _pad_reflect_replicate(self, pX:tuple[tuple[sint, sint], ...], mode:str) -> Self: @@ -695,10 +721,10 @@ def cummax(self, axis:int=0) -> tuple[Self, Self]: print(indices.numpy()) ``` """ - if self.ndim == 0: return self._split_cumalu(axis, Ops.MAX), type(self).zeros(self.shape, dtype=dtypes.int32, device=self.device) + if self.ndim == 0: return self._split_cumalu(axis, Ops.MAX), type(self).zeros(self.shape, dtype=dtypes.int32, device=self.device, buffer=False) values, n = self._split_cumalu(axis, Ops.MAX), int(self.shape[axis]) x, values_t = self.transpose(axis, -1), values.transpose(axis, -1) - match = x.unsqueeze(-1).eq(values_t.unsqueeze(-2)) * type(self).ones(n, n, device=self.device).triu() + match = x.unsqueeze(-1).eq(values_t.unsqueeze(-2)) * type(self).ones(n, n, device=self.device, buffer=False).triu() idx = (-(match * type(self).arange(n, 0, -1, device=self.device).reshape(n, 1)).max(-2) + n).cast(dtypes.int32) return values, idx.transpose(-1, axis) @@ -745,7 +771,7 @@ def logcumsumexp(self, axis=0) -> Self: last_dim_size = x.shape[-1] x_unsqueezed = x.unsqueeze(-2).expand((None,)*(self.ndim-1)+(last_dim_size, None)) x_cummax, _ = x.cummax(-1) - mask = type(self).ones(last_dim_size, last_dim_size, device=self.device).tril() + mask = type(self).ones(last_dim_size, last_dim_size, device=self.device, buffer=False).tril() ret = mask.where(x_unsqueezed - x_cummax.unsqueeze(-1), self.dtype.min).exp().sum(-1).log() + x_cummax return ret.transpose(-1, axis) @@ -818,7 +844,7 @@ def sort(self, dim:int=-1, descending:bool=False) -> tuple[Self, Self]: ``` """ x, dim = self, self._resolve_dim(dim) - if (orig_len := int(x.shape[dim])) <= 1: return x, x.zeros_like(dtype=dtypes.default_int) + if (orig_len := int(x.shape[dim])) <= 1: return x, x.const_like(0).cast(dtypes.default_int) # pad to power of 2 n_stages = (orig_len-1).bit_length() pads = tuple((0, 2**n_stages - orig_len) if i == dim else None for i in range(x.ndim)) @@ -842,7 +868,8 @@ def sort(self, dim:int=-1, descending:bool=False) -> tuple[Self, Self]: x = blue_box.cat(flipped_green_box.flip(flip_dims), dim=crossover_dim) x = x.flatten(dim, dim+n_stages-1).shrink_to(self.shape) # compute indices for sorted values - mask = type(self).ones(orig_len, orig_len, dtype=dtypes.bool, device=self.device).tril().reshape((None, None) + (1,)*(self.ndim-dim-1)) + mask = type(self).ones(orig_len, orig_len, dtype=dtypes.bool, device=self.device, buffer=False).tril() + mask = mask.reshape((None, None) + (1,)*(self.ndim-dim-1)) def compute_counts(t:Self): return (mask & t.unsqueeze(dim).eq(t.unsqueeze(dim+1))).sum(dim+1) count_orig, count_sorted = compute_counts(self), compute_counts(x) cond = self.unsqueeze(dim+1).eq(x.unsqueeze(dim)) & count_orig.unsqueeze(dim+1).eq(count_sorted.unsqueeze(dim)) @@ -922,7 +949,8 @@ def gather(self, dim:int, index:Self) -> Self: print(t.gather(1, Tensor([[0, 0], [1, 0]])).numpy()) ``` """ - if index.device != self.device: raise RuntimeError(f"expected index and self on the same device, {index.device=}, {self.device=}") + if index.device is not None and self.device is not None and index.device != self.device: + raise RuntimeError(f"expected index and self on the same device, {index.device=}, {self.device=}") if index.ndim != self.ndim: raise RuntimeError(f"self.ndim must equal index.ndim, {self.ndim=}, {index.ndim=}") dim = self._resolve_dim(dim) assert all(s >= i for d,(s,i) in enumerate(zip(self.shape, index.shape)) if d != dim), "requires self.shape[d] >= index.shape[d] for all d != dim" @@ -962,8 +990,10 @@ def interpolate(self, size:tuple[int, ...], mode:str="linear", align_corners:boo return x.cast(self.dtype) def _pre_scatter(self, dim:int, index:Self, src:Self) -> tuple[Self, Self]: - if index.device != self.device: raise RuntimeError(f"expected index and self on the same device, {index.device=}, {self.device=}") - if src.device != self.device: raise RuntimeError(f"expected src and self on the same device, {src.device=}, {self.device=}") + if index.device is not None and self.device is not None and index.device != self.device: + raise RuntimeError(f"expected index and self on the same device, {index.device=}, {self.device=}") + if src.device is not None and self.device is not None and src.device != self.device: + raise RuntimeError(f"expected src and self on the same device, {src.device=}, {self.device=}") dim = self._resolve_dim(dim) assert index.ndim == self.ndim == src.ndim, f"self.ndim, index.ndim and src.ndim must all equal, {self.ndim=} {index.ndim=} {src.ndim=}" assert all((d == dim or self_ >= index_) and src_ >= index_ for d,(self_,index_,src_) in enumerate(zip(self.shape, index.shape, src.shape))), \ @@ -1045,7 +1075,7 @@ def scatter(self, dim:int, index:Self, src:Self|PyConst, reduce:Literal['multipl ``` """ if reduce not in {None, "add", "multiply"}: raise TypeError(f"{reduce=} must be one of None, 'multiply', or 'add'") - if isinstance(src, (int, float, bool)): src = type(self).full(index.shape, src, dtype=self.dtype, device=self.device) + if isinstance(src, (int, float, bool)): src = type(self).full(index.shape, src, dtype=self.dtype, device=self.device, buffer=False) elif reduce: raise TypeError("non-scalar src is not supported with reduce arg. use scatter_reduce") if reduce == "add": return self.scatter_reduce(dim, index, src, "sum", include_self=True) if reduce == "multiply": return self.scatter_reduce(dim, index, src, "prod", include_self=True) @@ -1150,9 +1180,9 @@ def pool(x:Self, padding_:Sequence[int]) -> Self: reg_pads = resolve_pool_pads(padding, len(k_)) pads = self._apply_ceil_mode(reg_pads, k_, s_, dilation) if ceil_mode else reg_pads if not count_include_pad: - return pool(self, pads).sum(axis) / pool(self.ones_like(), pads).sum(axis) + return pool(self, pads).sum(axis) / pool(self.const_like(1), pads).sum(axis) if not ceil_mode: return pool(self, pads).mean(axis) - return pool(self, pads).sum(axis) / pool(self._pad_constant(((0,0),)*(self.ndim-len(k_)) + flat_to_grouped(reg_pads), 0.0).ones_like(), + return pool(self, pads).sum(axis) / pool(self._pad_constant(((0,0),)*(self.ndim-len(k_)) + flat_to_grouped(reg_pads), 0.0).const_like(1), tuple(cp-rp for cp,rp in zip(pads, reg_pads))).sum(axis) def max_pool2d(self, kernel_size:tuple[int, ...]=(2,2), stride=None, dilation=1, padding:int|tuple[int, ...]=0, @@ -1380,9 +1410,10 @@ def sparse_categorical_crossentropy(self, Y:Self, ignore_index:int=-1, label_smo ``` """ assert 0.0 <= label_smoothing <= 1.0, "label_smoothing must be in [0.0, 1.0]" - if Y.device != self.device: raise RuntimeError(f"expected Y and self on the same device, {Y.device=}, {self.device=}") + if Y.device is not None and self.device is not None and Y.device != self.device: + raise RuntimeError(f"expected Y and self on the same device, {Y.device=}, {self.device=}") log_probs = self.log_softmax() - loss_mask = Y.ne(ignore_index) if ignore_index != -1 else Y.ones_like(dtype=dtypes.bool) + loss_mask = Y.ne(ignore_index) if ignore_index != -1 else Y.const_like(1).cast(dtypes.bool) y = Y.unsqueeze(-1)._one_hot_along_dim(self.shape[-1], dim=-1) * loss_mask.unsqueeze(-1) smoothing = label_smoothing * (log_probs.mean(-1) * loss_mask) unreduced = ((1 - label_smoothing) * (log_probs * y).sum(-1) + smoothing) @@ -1434,7 +1465,7 @@ def nll_loss(self, Y:Self, weight:Self|None=None, ignore_index:int|None=None, re print(t.log_softmax().nll_loss(Y, reduction='none').numpy()) ``` """ - weight = Y.ones_like() if weight is None else weight.gather(0, Y.flatten()).reshape(Y.shape) + weight = Y.const_like(1) if weight is None else weight.gather(0, Y.flatten()).reshape(Y.shape) masked_weight = weight if ignore_index is None else weight * Y.ne(ignore_index) nll = -self.gather(1, Y.unsqueeze(1)).squeeze(1) * masked_weight return nll.sum() / masked_weight.sum() if reduction == "mean" else nll._do_reduction(reduction) diff --git a/tinygrad/mixin/elementwise.py b/tinygrad/mixin/elementwise.py index e61b070c5b904..0a4fd08779003 100644 --- a/tinygrad/mixin/elementwise.py +++ b/tinygrad/mixin/elementwise.py @@ -25,8 +25,10 @@ def _binop(self, op: Ops, x: Self | ConstType, reverse: bool) -> Self: def usum(self, *uops) -> Self: return functools.reduce(operator.or_ if self.dtype is dtypes.bool else operator.add, argfix(*uops), self) def uprod(self, *uops) -> Self: return functools.reduce(operator.and_ if self.dtype is dtypes.bool else operator.mul, argfix(*uops), self) - # NOTE: Tensor overrides this to also set requires_grad=False def detach(self) -> Self: + """ + Returns a new tensor with the same data as this tensor, but detached from the autograd graph. + """ return self.alu(Ops.DETACH) def logical_not(self) -> Self: @@ -133,7 +135,8 @@ def bitwise_not(self) -> Self: ``` """ self._check_dtype() - return self.logical_not() if self.dtype == dtypes.bool else self ^ -1 + if self.dtype == dtypes.bool: return self.logical_not() + return (self ^ self.dtype.max) if dtypes.is_unsigned(self.dtype) else (self ^ -1) def bitwise_and(self, x: Self | ConstType, reverse: bool = False) -> Self: """ diff --git a/tinygrad/mixin/movement.py b/tinygrad/mixin/movement.py index 2f4853ee07db4..dc03177ec9d4b 100644 --- a/tinygrad/mixin/movement.py +++ b/tinygrad/mixin/movement.py @@ -19,7 +19,7 @@ def shape(self) -> tuple[sint, ...]: raise NotImplementedError @property - def device(self) -> str|tuple[str, ...]: + def device(self) -> str|tuple[str, ...]|None: raise NotImplementedError # great functions you get! @@ -178,7 +178,7 @@ def reshape(self, shape, *args) -> Self: def pad(self, arg:tuple[tuple[sint, sint] | None, ...]) -> Self: if self.ndim != len(arg): raise ValueError(f"{self.ndim=} != {len(arg)=}") - ret = self._mop(Ops.PAD, tuple(x if x is not None else (0, 0) for x in arg)) + ret = self._mop(Ops.PAD, tuple((x[0], s+x[0]+x[1]) if x is not None else (0, s) for x, s in zip(arg, self.shape))) return self if ret.shape == self.shape else ret def shrink(self, arg: tuple[tuple[sint, sint] | None, ...]) -> Self: @@ -200,7 +200,7 @@ def shrink(self, arg: tuple[tuple[sint, sint] | None, ...]) -> Self: """ if self.ndim != len(arg): raise ValueError(f"{self.ndim=} != {len(arg)=}") - ret = self._mop(Ops.SHRINK, arg=[x if x is not None else (0, s) for x, s in zip(arg, self.shape)]) + ret = self._mop(Ops.SHRINK, arg=[(x[0], x[1]-x[0]) if x is not None else (0, s) for x, s in zip(arg, self.shape)]) return self if ret.shape == self.shape else ret def permute(self, order, *args) -> Self: @@ -251,7 +251,7 @@ def shrink_to(self, shape, *args) -> Self: 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, 0 if ns is None else ns-s) for s,ns in zip(self.shape, argfix(shape, *args), strict=True)])) + 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))) def view(self, shape, *args) -> Self: """`.view` is an alias for `.reshape`.""" diff --git a/tinygrad/nn/__init__.py b/tinygrad/nn/__init__.py index 21d1b41b645e6..d8c120466d595 100644 --- a/tinygrad/nn/__init__.py +++ b/tinygrad/nn/__init__.py @@ -35,8 +35,8 @@ def __init__(self, sz:int, eps=1e-5, affine=True, track_running_stats=True, mome self.weight: Tensor|None = Tensor.ones(sz) if affine else None self.bias: Tensor|None = Tensor.zeros(sz) if affine else None - self.num_batches_tracked = Tensor.zeros(dtype='long', requires_grad=False) - if track_running_stats: self.running_mean, self.running_var = Tensor.zeros(sz, requires_grad=False), Tensor.ones(sz, requires_grad=False) + self.num_batches_tracked = Tensor.zeros(dtype='long').is_param_(False) + if track_running_stats: self.running_mean, self.running_var = Tensor.zeros(sz).is_param_(False), Tensor.ones(sz).is_param_(False) def calc_stats(self, x:Tensor) -> tuple[Tensor, Tensor]: shape_mask: list[int] = [1, -1, *([1]*(x.ndim-2))] @@ -366,7 +366,7 @@ def _embedding_bwd_kernel(grad_weight:UOp, grad_emb:UOp, idx:UOp) -> UOp: return (grad_weight_uop.cast(weight.dtype), None) def _embedding_fwd(weight:Tensor, idx:Tensor) -> Tensor: - arange = Tensor.arange(weight.shape[0], requires_grad=False, device=weight.device) + arange = Tensor.arange(weight.shape[0], device=weight.device) return (arange == idx.unsqueeze(-1)).unsqueeze(-1).where(weight, 0).sum(-2, dtype=weight.dtype) @functools.cache @@ -411,7 +411,7 @@ def __init__(self, input_size:int, hidden_size:int, bias:bool=True): self.bias_hh: Tensor|None = Tensor.zeros(hidden_size*4) if bias else None def __call__(self, x:Tensor, hc:tuple[Tensor, Tensor]|None=None) -> tuple[Tensor, Tensor]: - if hc is None: hc = (Tensor.zeros(x.size(0), self.weight_hh.size(1), dtype=x.dtype, device=x.device),)*2 + if hc is None: hc = (Tensor.zeros(x.size(0), self.weight_hh.size(1), dtype=x.dtype, device=x.device, buffer=False),)*2 gates = x.linear(self.weight_ih.T, self.bias_ih) + hc[0].linear(self.weight_hh.T, self.bias_hh) i, f, g, o = gates.chunk(4, dim=1) i, f, g, o = i.sigmoid(), f.sigmoid(), g.tanh(), o.sigmoid() diff --git a/tinygrad/nn/onnx.py b/tinygrad/nn/onnx.py index 8850a525a54fe..e6cd3fba1f84f 100644 --- a/tinygrad/nn/onnx.py +++ b/tinygrad/nn/onnx.py @@ -403,12 +403,12 @@ def _parse_input(self, name: str, value: Any, spec: OnnxValue): if spec.is_optional and value is None: return None if spec.is_sequence: if not isinstance(value, Sequence): raise RuntimeError(f"input {name} received {value}, expected a sequence type") - sequence = [Tensor(v, dtype=spec.dtype, requires_grad=self.is_training) if not isinstance(v, Tensor) else v for v in value] + sequence = [Tensor(v, dtype=spec.dtype) if not isinstance(v, Tensor) else v for v in value] if not all_same(tuple(t.shape for t in sequence)): raise RuntimeError(f"Shapes for input {name} sequence must be homogeneous") if not all(t.dtype is spec.dtype for t in sequence): warnings.warn(f"Dtypes for input {name} sequence aren't all {spec.dtype}") return sequence dtype = _from_np_dtype(value.dtype) if is_numpy_ndarray(value) else spec.dtype - tensor = Tensor(value, dtype=dtype, requires_grad=self.is_training) if not isinstance(value, Tensor) else value + tensor = Tensor(value, dtype=dtype) if not isinstance(value, Tensor) else value if tensor.dtype is not spec.dtype: warnings.warn(f"input {name} has mismatch on dtype. Expected {spec.dtype}, received {tensor.dtype}.") for dim, (onnx_dim, user_dim_input) in enumerate(zip(spec.shape, tensor.shape, strict=True)): if isinstance(onnx_dim, str): @@ -566,10 +566,10 @@ def Identity(x:Tensor): return x def Constant(sparse_value:Tensor|None=None, value:Tensor|None=None, value_float:float|None=None, value_floats:list[float]|None=None, value_int:int|None=None, value_ints:list[int]|None=None, value_string:str|None=None, value_strings:list[str]|None=None): if value is not None: return value - if value_float is not None: return Tensor(value_float, dtype=dtypes.float32, requires_grad=False) - if value_floats is not None: return Tensor(list(value_floats), dtype=dtypes.float32, requires_grad=False) - if value_int is not None: return Tensor(value_int, dtype=dtypes.int64, requires_grad=False) - if value_ints is not None: return Tensor(list(value_ints), dtype=dtypes.int64, requires_grad=False) + if value_float is not None: return Tensor(value_float, dtype=dtypes.float32) + if value_floats is not None: return Tensor(list(value_floats), dtype=dtypes.float32) + if value_int is not None: return Tensor(value_int, dtype=dtypes.int64) + if value_ints is not None: return Tensor(list(value_ints), dtype=dtypes.int64) if value_string is not None or value_strings is not None or sparse_value is not None: raise NotImplementedError('Constant OP not implemented for value_string, value_strings and sparse_value') @@ -586,7 +586,7 @@ def ImageDecoder(encoded_stream:bytes, pixel_format="RGB"): raise ValueError(f"pixel_format={pixel_format!r} is not supported.") def EyeLike(x:Tensor, dtype:int|None=None, k:int=0): - ret = Tensor.eye(cast(int, min(x.shape)), dtype=OnnxDataType(dtype).to_dtype() if dtype is not None else x.dtype) + ret = Tensor.eye(cast(int, min(x.shape)), dtype=OnnxDataType(dtype).to_dtype() if dtype is not None else x.dtype, device=x.device) return ret if x.size(0) == x.size(1) else ret.pad(tuple(None if d == ret.size(0) else (k, d-ret.shape[0]-k) for d in x.shape)) def OptionalHasElement(x:Tensor|None=None): return Tensor(x is not None and x.numel() > 0) @@ -597,7 +597,7 @@ def ConstantOfShape(shape:list[int], value:Tensor|None=None): return value.expand(shape) def Size(data:Tensor): return data.numel() - def Shape(data:Tensor, end:int|None=None, start:int=0): return Tensor(data.shape[start:end], dtype=dtypes.int64) + def Shape(data:Tensor, end:int|None=None, start:int=0): return Tensor(data.shape[start:end], dtype=dtypes.int64, device=data.device) # ***** Unary Ops (math) ***** def Not(x:Tensor): return x.logical_not() @@ -783,11 +783,11 @@ def Resize(X:Tensor, roi:list[float]|None=None, scales:list[float]|None=None, si axes:list[int]|None=None, coordinate_transformation_mode:str='half_pixel', cubic_coeff_a:float=-0.75, exclude_outside:int=0, extrapolation_value:float=0.0, keep_aspect_ratio_policy:str='stretch', mode:str='nearest', nearest_mode:str='round_prefer_floor'): def _apply_transformation(input_sz, output_sz, scale_dim, mode): - index = Tensor.arange(output_sz, requires_grad=False, device=X.device) + index = Tensor.arange(output_sz, device=X.device) if mode == "half_pixel": return (index + 0.5) / scale_dim - 0.5 - if mode == "align_corners": return index * (input_sz - 1) / (output_sz - 1) if output_sz != 1 else Tensor.zeros_like(index) + if mode == "align_corners": return index * (input_sz - 1) / (output_sz - 1) if output_sz != 1 else index.const_like(0) if mode == "asymmetric": return index / scale_dim - if mode == "pytorch_half_pixel": return ((index + 0.5) / scale_dim - 0.5) if output_sz != 1 else Tensor.zeros_like(index) + if mode == "pytorch_half_pixel": return ((index + 0.5) / scale_dim - 0.5) if output_sz != 1 else index.const_like(0) if mode == "half_pixel_symmetric": output_dim_scaled = input_sz * scale_dim return (input_sz / 2) * (1 - (output_sz / output_dim_scaled)) + (index + 0.5) / scale_dim - 0.5 @@ -908,12 +908,13 @@ def GroupNormalization(x:Tensor, scale:Tensor, bias:Tensor, num_groups:int, epsi return x * scale.reshape(1, -1, *[1] * (x.ndim-2)) + bias.reshape(1, -1, *[1] * (x.ndim-2)) def InstanceNormalization(x:Tensor, scale:Tensor, bias:Tensor, epsilon:float=1e-05): return GroupNormalization(x, scale, bias, num_groups=cast(int, x.shape[1]), epsilon=epsilon) - def LayerNormalization(x:Tensor, scale:Tensor, bias:Tensor, axis:int=-1, epsilon:float=1e-05, stash_type:int=1): + def LayerNormalization(x:Tensor, scale:Tensor, bias:Tensor|None=None, axis:int=-1, epsilon:float=1e-05, stash_type:int=1): assert stash_type == 1, "only float32 is supported" axes = tuple(i for i in range(axis if axis >= 0 else x.ndim + axis, x.ndim)) mean = (x32:=x.cast(dtypes.float)).mean(axis=axes, keepdim=True) inv_std_dev = (x32.sub(mean)).square().mean(axis=axes, keepdim=True).add(epsilon).rsqrt() - return (x32.sub(mean)*inv_std_dev).cast(x.dtype).mul(scale).add(bias), mean, inv_std_dev + ret = (x32.sub(mean)*inv_std_dev).cast(x.dtype).mul(scale) + return (ret.add(bias) if bias is not None else ret), mean, inv_std_dev def SkipLayerNormalization(x:Tensor, skip:Tensor, gamma:Tensor, beta:Tensor|None=None, bias:Tensor|None=None, epsilon:float=1e-12): x = x + skip if bias is not None: x = x + bias @@ -933,7 +934,7 @@ def embedding(x:Tensor, vocab_size, weight:Tensor) -> Tensor: return x.unsqueeze(-1).expand(*x.shape, vocab_size)._one_hot_along_dim(vocab_size) @ weight # bert embedding layer - if position_ids is None: position_ids = Tensor.arange(seq_length, requires_grad=False).unsqueeze(0).expand(*input_shape) + if position_ids is None: position_ids = Tensor.arange(seq_length, device=input_ids.device).unsqueeze(0).expand(*input_shape) wrd_embedding_res = embedding(input_ids, vocab_size, word_embedding) pos_embedding_res = embedding(position_ids, max_position_embeddings, position_embedding) @@ -963,11 +964,11 @@ def SpaceToDepth(X:Tensor, blocksize:int): # Reimplemented here because you need legacy RNG for passing ONNX tests. def dropout_7(data:Tensor, ratio:float=0.5, training_mode:bool=False, seed:int|None=None): import numpy as np - if not training_mode: return data, data.full_like(True, dtype=dtypes.bool) + if not training_mode: return data, data.const_like(True).cast(dtypes.bool) if seed is not None: - rand = Tensor(np.random.RandomState(seed).random(cast(tuple[int,...], data.shape)), requires_grad=False, dtype=data.dtype, device=data.device) + rand = Tensor(np.random.RandomState(seed).random(cast(tuple[int,...], data.shape)), dtype=data.dtype, device=data.device) else: - rand = data.rand_like(requires_grad=False) + rand = data.rand_like() mask = rand >= ratio return data * mask / (1.0 - ratio), mask # 6 with 'is_test' needed for https://github.com/MTlab/onnx2caffe/raw/refs/heads/master/model/MobileNetV2.onnx @@ -976,7 +977,7 @@ def dropout_6(data:Tensor, ratio:float=0.5, is_test=0): return dropout_7(data, r def _window(size, output_datatype, periodic, a): size = int(_resolve_const(size)) - N, n = (size if periodic else size - 1), Tensor.arange(size, requires_grad=False) + N, n = (size if periodic else size - 1), Tensor.arange(size) w = a[0] - a[1] * (n * (2 * math.pi / N)).cos() + a[2] * (n * (4 * math.pi / N)).cos() return w.cast(OnnxDataType(output_datatype).to_dtype()) def HannWindow(size, output_datatype:int=1, periodic:int=1): return _window(size, output_datatype, periodic, (0.5, 0.5, 0)) @@ -999,7 +1000,7 @@ def generate_grid(steps): if align_corners: return Tensor.linspace(-1, 1, steps, device=theta.device) return Tensor.linspace(-1+1/steps, 1-1/steps, steps, device=theta.device) grids = Tensor.meshgrid(*(generate_grid(d) for d in spatial_dims)) - base_grid = Tensor.stack(*reversed(grids), Tensor.ones_like(grids[0], device=theta.device), dim=-1) + base_grid = Tensor.stack(*reversed(grids), grids[0].const_like(1), dim=-1) base_grid = base_grid.reshape(1, prod(spatial_dims), len(grids)+1).expand(N, -1, -1) return (base_grid @ theta.transpose(1, 2)).reshape(N, *spatial_dims, -1) @@ -1031,18 +1032,18 @@ def attention_contrib(x:Tensor, weights:Tensor, bias:Tensor|None=None, mask_inde if mask_index.ndim != 1: mask = mask_index.bool() else: if mask_index.shape[0] == batch_size: - mask = Tensor.arange(attn_scores.shape[-1], requires_grad=False, device=mask_index.device).unsqueeze(0) < mask_index.unsqueeze(1) + mask = Tensor.arange(attn_scores.shape[-1], device=mask_index.device).unsqueeze(0) < mask_index.unsqueeze(1) elif mask_index.shape[0] == 2*batch_size: end_positions = mask_index[:batch_size] start_positions = mask_index[batch_size:] - arange = Tensor.arange(seq_len).unsqueeze(0) + arange = Tensor.arange(seq_len, device=mask_index.device).unsqueeze(0) mask = (arange < end_positions.unsqueeze(1)) & (arange >= start_positions.unsqueeze(1)) else: raise NotImplementedError("mask_index with shape (3 * batch_size + 2) is not implemented") while mask.ndim < 4: mask = mask.unsqueeze(1) attn_scores = mask.where(attn_scores, mask_filter_value) if unidirectional: - causal_mask = Tensor.ones((seq_len, seq_len), dtype=dtypes.bool).tril() + causal_mask = Tensor.ones((seq_len, seq_len), dtype=dtypes.bool, device=attn_scores.device, buffer=False).tril() attn_scores = causal_mask.where(attn_scores, mask_filter_value) output = attn_scores.softmax(-1) @ v @@ -1074,7 +1075,7 @@ def attention_onnx(Q:Tensor, K:Tensor, V:Tensor, attn_mask:Tensor|None=None, pas qk_matmul_return_val = scores if is_causal: - causal_mask = Tensor.ones(Q.shape[-2], K.shape[-2], device=Q.device, dtype=dtypes.bool, requires_grad=False).tril(0) + causal_mask = Tensor.ones(Q.shape[-2], K.shape[-2], device=Q.device, dtype=dtypes.bool, buffer=False).tril(0) scores = scores.masked_fill(causal_mask.logical_not(), -float("inf")) if attn_mask is not None: @@ -1131,7 +1132,7 @@ def NonZero(x:Tensor): flat_idx = Tensor.arange(mask.numel(), dtype=dtypes.int64, device=x.device).masked_select(mask) if flat_idx.ndim == 0: flat_idx = flat_idx.reshape(1) if x.ndim == 0: - return Tensor.zeros((0, flat_idx.shape[0]), dtype=dtypes.int64, device=x.device, requires_grad=False) + return Tensor.zeros((0, flat_idx.shape[0]), dtype=dtypes.int64, device=x.device) strides = [prod(int(s) for s in x.shape[i+1:]) if i+1 < x.ndim else 1 for i in range(x.ndim)] coords = [((flat_idx // stride) % int(dim)) for stride, dim in zip(strides, x.shape)] return Tensor.stack(*coords, dim=0) @@ -1198,7 +1199,7 @@ def Compress(inp:Tensor, condition:list[bool], axis:int|None=None): inp = inp.flatten() axis = 0 axis = inp._resolve_dim(axis) - con = Tensor([i for i,cond in enumerate(condition) if cond]) # compress in python + con = Tensor([i for i,cond in enumerate(condition) if cond], device=inp.device) # compress in python return inp[tuple(con if i == axis else slice(None) for i in range(inp.ndim))] # ***** Quantization Ops ***** @@ -1228,6 +1229,9 @@ def DequantizeLinear(x:Tensor, x_scale:Tensor, x_zero_point:Tensor|int=0, axis:i def QLinearConv(x:Tensor, x_scale:Tensor, x_zero_point:Tensor, w:Tensor, w_scale:Tensor, w_zero_point:Tensor, y_scale:Tensor, y_zero_point:Tensor, B:Tensor|None=None, **opts): + # align quant params for non-scalars based on https://onnx.ai/onnx/operators/onnx__QLinearConv.html + w_zero_point = w_zero_point.reshape(-1, 1, 1, 1) + w_scale = w_scale.reshape(1, -1, 1, 1) return _qlinearop_quantized(Conv, [x,w], [x_zero_point,w_zero_point], [x_scale,w_scale], y_scale, y_zero_point, **{"B":B, **opts}) def QLinearMatMul(a:Tensor, a_scale:Tensor, a_zero_point:Tensor, b:Tensor, b_scale:Tensor, b_zero_point:Tensor, y_scale:Tensor, @@ -1277,8 +1281,8 @@ def Adam(R:Tensor, T:int, *inputs:Tensor, alpha:float=0.9, beta:float=0.999, eps if T == 0: opt.b1_t, opt.b2_t = opt.b1_t.zeros_like(), opt.b2_t.zeros_like() else: # `T-1` since it's applied again at the start of `_step` - opt.b1_t = Tensor([alpha**(T-1)], dtype=dtypes.float32, device=X.device, requires_grad=False) - opt.b2_t = Tensor([beta**(T-1)], dtype=dtypes.float32, device=X.device, requires_grad=False) + opt.b1_t = Tensor([alpha**(T-1)], dtype=dtypes.float32, device=X.device) + opt.b2_t = Tensor([beta**(T-1)], dtype=dtypes.float32, device=X.device) opt.step() X = (1 - norm_coefficient_post) * X return [X, V, H] @@ -1294,8 +1298,7 @@ def Momentum(R:Tensor, T:int, *inputs:Tensor, alpha:float, beta:float, mode:str, return [X, V] def Gradient(*inputs:Tensor, y:str, intermediate_tensors:dict[str, Tensor], **_): - intermediate_tensors[y].backward() - return tuple([t.grad for t in inputs]) + return tuple(intermediate_tensors[y].gradient(*inputs)) return { # Tensor ops diff --git a/tinygrad/nn/optim.py b/tinygrad/nn/optim.py index a256c90684e55..a5a9d43a5e1e0 100644 --- a/tinygrad/nn/optim.py +++ b/tinygrad/nn/optim.py @@ -10,21 +10,21 @@ class Optimizer: """ def __init__(self, params: list[Tensor], lr: float, device=None, fused=FUSE_OPTIM): if lr < 0: raise ValueError(f"Invalid learning rate: {lr}") - self.params: list[Tensor] = dedup([x for x in params if x.requires_grad]) + self.params: list[Tensor] = dedup([x for x in params if x.is_param]) assert len(self.params) != 0, "optimizer must have at least one param" - self.buffers: list[Tensor] = dedup([x for x in params if not x.requires_grad]) # buffers are still realized + self.buffers: list[Tensor] = dedup([x for x in params if not x.is_param]) # buffers are still realized self.device = device or self.params[0].device self.param_dtype = to_dtype(getenv("OPTIM_DTYPE", "float32")) self.fused = fused # store lr in at least float32 precision - self.lr = Tensor(lr if getenv("CONST_LR") else [lr], requires_grad=False, device=self.device, + self.lr = Tensor(lr if getenv("CONST_LR") else [lr], device=self.device, dtype=least_upper_dtype(dtypes.default_float, dtypes.float32)) if self.fused: self.pos_params = list(itertools.accumulate(self.params, lambda x,y: x+y.numel(), initial=0)) def _new_optim_param(self) -> list[Tensor]: - if self.fused: return [Tensor.zeros(self.pos_params[-1], dtype=self.param_dtype, device=self.device, requires_grad=False)] - if isinstance(self.device, tuple): return [Tensor.zeros_like(t, dtype=self.param_dtype, requires_grad=False) for t in self.params] - else: return [Tensor.zeros(t.shape, dtype=self.param_dtype, device=self.device, requires_grad=False) for t in self.params] + if self.fused: return [Tensor.zeros(self.pos_params[-1], dtype=self.param_dtype, device=self.device)] + if isinstance(self.device, tuple): return [Tensor.zeros_like(t, dtype=self.param_dtype) for t in self.params] + else: return [Tensor.zeros(t.shape, dtype=self.param_dtype, device=self.device) for t in self.params] def zero_grad(self): """ @@ -154,7 +154,7 @@ def __init__(self, params: list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-6, w if weight_decay < 0: raise ValueError(f"Invalid weight_decay value: {weight_decay}") super().__init__(params, lr, device, fused) self.b1, self.b2, self.eps, self.wd, self.adam = b1, b2, eps, weight_decay, adam - self.b1_t, self.b2_t = (Tensor.ones((1,), dtype=dtypes.float32, device=self.device, requires_grad=False) for _ in [b1, b2]) + self.b1_t, self.b2_t = (Tensor.ones((1,), dtype=dtypes.float32, device=self.device).is_param_(False) for _ in [b1, b2]) self.m = self._new_optim_param() self.v = self._new_optim_param() diff --git a/tinygrad/nn/state.py b/tinygrad/nn/state.py index ebf521b2bac16..7b340c4ec71e0 100644 --- a/tinygrad/nn/state.py +++ b/tinygrad/nn/state.py @@ -3,7 +3,7 @@ from typing import Any, Callable, BinaryIO, Iterable, cast from tinygrad.tensor import Tensor from tinygrad.dtype import dtypes -from tinygrad.helpers import prod, argsort, DEBUG, Timing, CI, GlobalCounters, tqdm, round_up, T, strides_for_shape +from tinygrad.helpers import prod, argsort, DEBUG, Timing, GlobalCounters, tqdm, round_up, T, strides_for_shape class TensorIO(io.RawIOBase, BinaryIO): def __init__(self, t: Tensor): @@ -145,7 +145,7 @@ def __init__(self): model_state_dict = get_state_dict(model) if DEBUG >= 1 and len(state_dict) > len(model_state_dict): print("WARNING: unused weights in state_dict", sorted(list(state_dict.keys() - model_state_dict.keys()))) - for k,v in (t := tqdm(model_state_dict.items(), disable=CI or not verbose)): + for k,v in (t := tqdm(model_state_dict.items(), disable=None if verbose else True)): t.desc = f"ram used: {GlobalCounters.mem_used/1e9:5.2f} GB, {k:50s}: " if k not in state_dict and not strict: if DEBUG >= 1: print(f"WARNING: not loading {k}") diff --git a/tinygrad/renderer/__init__.py b/tinygrad/renderer/__init__.py index 044d6b28f2cf6..8086a59c67616 100644 --- a/tinygrad/renderer/__init__.py +++ b/tinygrad/renderer/__init__.py @@ -1,9 +1,9 @@ from __future__ import annotations from typing import Callable, cast from dataclasses import dataclass -from tinygrad.helpers import prod, Target +from tinygrad.helpers import prod, Target, EMULATED_DTYPES from tinygrad.uop.ops import Ops, UOp, sint, ssimplify, smin, GroupOp, PatternMatcher -from tinygrad.dtype import AddrSpace, PtrDType +from tinygrad.dtype import AddrSpace, PtrDType, DType, dtypes from tinygrad.codegen.opt.tc import TensorCore from tinygrad.device import Compiler @@ -38,7 +38,7 @@ def range_gate(x): return x.op is not Ops.RANGE for u in uops: if u.op in {Ops.LOAD, Ops.STORE}: buf = u - while len(buf.src): buf = buf.src[0] + while len(buf.src) and buf.op is not Ops.PARAM: buf = buf.src[0] if buf.op is Ops.PARAM: # u.src[0] is INDEX, cap at buffer size for re-reads (e.g. matmul) accessed = mem.get((buf, u.op), 0) + u.src[0].dtype.base.itemsize * mults @@ -85,3 +85,6 @@ def __reduce__(self): return self.__class__, (self.target,) def render(self, uops:list[UOp]) -> str: raise NotImplementedError("needs a renderer") def asm(self, prg:UOp, lin:UOp) -> bytes: raise NotImplementedError("needs an assembler") def aux(self, uops:list[UOp]) -> dict: raise NotImplementedError("needs aux") + def supported_dtypes(self) -> set[DType]: + # double can't be bitcast to anything without long support + return set(dtypes.all) - {dtypes.weakint} - ({dtypes.double} if dtypes.long in EMULATED_DTYPES.tolist(dtypes) else set()) diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index da4937beea356..dc0801e7a9e7b 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -3,7 +3,7 @@ from collections import defaultdict, Counter from tinygrad.codegen.opt import tc from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, range_str, axis_letters -from tinygrad.helpers import strip_parens, getenv, prod, dedup, Target, CPU_COUNT +from tinygrad.helpers import strip_parens, getenv, prod, dedup, Target, CPU_COUNT, IMAGE, FLOAT16 from tinygrad.dtype import ImageDType, dtypes, DType, PtrDType, AddrSpace, truncate, float_to_bf16 from tinygrad.renderer import Renderer from tinygrad.codegen.late.devectorizer import no_vectorized_alu @@ -47,7 +47,7 @@ (UPat.var("buf").index(UPat.var('idx')), lambda ctx,buf,idx: f"({ctx[buf]}+{strip_parens(ctx[idx]) if idx.arg == Ops.ADD else ctx[idx]})"), (UPat(Ops.LOAD, src=(UPat.var('bidx'),)), lambda ctx,bidx: f"(*{ctx[bidx]})"), (UPat(Ops.LOAD, src=(UPat.var("bidx"), UPat.var("var"), UPat.var("gate"))), lambda ctx,bidx,var,gate: f"({ctx[gate]}?*{ctx[bidx]}:{ctx[var]})"), - (UPat(Ops.STORE, src=(UPat.var('bidx'), UPat.var("var")), allow_any_len=True), lambda ctx,bidx,var: f"*{ctx[bidx]} = {ctx[var]};"), + (UPat(Ops.STORE, src=(UPat.var('bidx'), UPat.var("var"))), lambda ctx,bidx,var: f"*{ctx[bidx]} = {ctx[var]};"), # alu/gep # TODO: look for left-associative (UPat(GroupOp.ALU, name="x"), lambda ctx,x: ctx.code_for_op[x.op]( @@ -178,8 +178,8 @@ def _render(self, uops:list[UOp]) -> tuple[str, list[str], list[tuple[str,tuple[ continue if u.op in (Ops.PARAM, Ops.DEFINE_VAR): if u.op is not Ops.PARAM: r[u] = u.arg[0] - elif isinstance(u.dtype, ImageDType): r[u] = f"data{u.arg}_{u.dtype.shape[0]}x{u.dtype.shape[1]}" - else: r[u] = f"data{u.arg}_{sz}" if (sz:=u.ptrdtype.size) > 0 else f"data{u.arg}" + elif isinstance(u.dtype, ImageDType): r[u] = f"data{u.arg.slot}_{u.dtype.shape[0]}x{u.dtype.shape[1]}" + else: r[u] = f"data{u.arg.slot}_{sz}" if (sz:=u.max_numel()) > 0 else f"data{u.arg.slot}" bufs[u] = (r[u], (u.dtype, u in writable_params)) continue @@ -198,7 +198,7 @@ def _render(self, uops:list[UOp]) -> tuple[str, list[str], list[tuple[str,tuple[ if u.op in {Ops.ENDIF, Ops.END}: depth -= 1 if (u.op is not Ops.CAST or u.dtype.vcount == 1) and (u.op in {Ops.CONST, Ops.GEP, Ops.INDEX, Ops.CUSTOMI} or \ - (u.op is Ops.LOAD and u.src[0].ptrdtype.addrspace == AddrSpace.REG) or \ + (u.op is Ops.LOAD and u.src[0].addrspace == AddrSpace.REG) or \ (u.op is Ops.CAST and isinstance(u.dtype, PtrDType)) or \ (u.op in {Ops.STACK, *(GroupOp.ALU-{Ops.WHERE}), Ops.CAST, Ops.BITCAST} and child_count[u] == 1 and not getenv("EXPAND_SSA"))): r[u] = l @@ -272,12 +272,14 @@ def render_kernel(self, function_name, kernel, bufs, uops, prefix=None) -> str: defines = '\n'.join(self._render_defines(uops)) return defines + "\n" + self._render_body(function_name, kernel, bufs, uops, prefix) + "\n" + self._render_entry(function_name, bufs) -class ClangJITRenderer(ClangRenderer): + def supported_dtypes(self): + return {d for d in super().supported_dtypes() if (d != dtypes.bfloat16 or self.target.arch.startswith(("x86", "arm"))) and d not in dtypes.fp8s} + def __init__(self, target:Target): super().__init__(target) - from tinygrad.runtime.support.compiler_cpu import ClangJITCompiler + from tinygrad.runtime.support.compiler_cpu import ClangCompiler if "AMX" in target.arch: self.tensor_cores = tc.amx - self.compiler = ClangJITCompiler([x for x in target.arch.split(",") if x != "AMX"]) + self.compiler = ClangCompiler([x for x in target.arch.split(",") if x != "AMX"]) class OpenCLRenderer(CStyleLanguage): has_aux = True @@ -306,7 +308,7 @@ class OpenCLRenderer(CStyleLanguage): (UPat(Ops.LOAD, dtype=dtypes.float.vec(4), src=(UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')),)), lambda ctx,buf,idx_y,idx_x: f"read_imagef({ctx[buf]}, smp, (int2)({ctx[idx_x]},{ctx[idx_y]}))"), (UPat(Ops.STORE, src=(UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')), - UPat.var("var", dtypes.float.vec(4))), allow_any_len=True), + UPat.var("var", dtypes.float.vec(4)))), lambda ctx,buf,idx_y,idx_x,var: f"write_imagef({ctx[buf]}, (int2)({ctx[idx_x]},{ctx[idx_y]}), {ctx[var]});"), ]) + base_rewrite @@ -317,10 +319,14 @@ def render_kernel(self, function_name, kernel, bufs, uops, prefix=None) -> str: def aux(self, uops:list[UOp]): arg_dtypes:list[list[tuple[int, DType]]] = [] for i,u in enumerate(u for u in uops if u.op is Ops.PARAM): - if len(arg_dtypes) >= u.arg: arg_dtypes.append([]) - arg_dtypes[u.arg].append((i, u.dtype)) + while len(arg_dtypes) <= u.arg.slot: arg_dtypes.append([]) + arg_dtypes[u.arg.slot].append((i, u.dtype)) return tuple(tuple(a) for a in arg_dtypes), + def supported_dtypes(self): return {d for d in super().supported_dtypes() + if (d != dtypes.half or "cl_khr_fp16" in self.target.arch) and + (d != dtypes.double or "cl_khr_fp64" in self.target.arch) and d not in dtypes.fp8s} + class IntelRenderer(OpenCLRenderer): suffix, kernel_typedef = "INTEL", "__attribute__((intel_reqd_sub_group_size(8)))\n" + "__kernel void" tensor_cores = tc.intel @@ -343,7 +349,7 @@ class MetalRenderer(CStyleLanguage): def __init__(self, target:Target): super().__init__(target) from tinygrad.runtime.ops_metal import MetalCompiler - self.compiler, self.tensor_cores = MetalCompiler(), tc.metal if target.arch == "arm64" else [] + self.compiler, self.tensor_cores = MetalCompiler(), tc.metal if target.arch.startswith("Apple") and int(target.arch[5:]) >= 7 else [] # language options kernel_typedef = "kernel void" @@ -382,6 +388,10 @@ def render_kernel(self, function_name, kernel, bufs, uops, prefix=None): simdgroup_multiply_accumulate(mat_c, mat_a, mat_b, mat_c);\n return {dstr_out}(mat_c.thread_elements()[0], mat_c.thread_elements()[1]);\n}}""") return super().render_kernel(function_name, kernel, bufs, uops, prefix) + def supported_dtypes(self): + return {d for d in super().supported_dtypes() if (d != dtypes.bfloat16 or ((arch:=self.target.arch).startswith("Apple") and int(arch[5:]) >= 6)) + and d not in dtypes.fp8s+(dtypes.double,)} + _nms = list("xyzwabcdefghijkl") + [f'v{i}' for i in range(16, 32)] class CUDARenderer(CStyleLanguage): @@ -456,6 +466,11 @@ def render_kernel(self, function_name, kernel, bufs, uops, prefix=None): return super().render_kernel(function_name, kernel, bufs, uops, prefix=prefix) + def supported_dtypes(self): + ver = int(self.target.arch[3:]) + return {d for d in super().supported_dtypes() if (d != dtypes.half or ver >= 53) and (d != dtypes.bfloat16 or ver >= 80) + and (d not in dtypes.fp8_ocp or ver >= 89) and d not in dtypes.fp8_fnuz} + class NVCCRenderer(CUDARenderer): def __init__(self, target:Target): super().__init__(target, use_nvcc=True) @@ -482,6 +497,10 @@ def __init__(self, target:Target, use_hipcc=False): # gfx942 => MI300, gfx1100 = (UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{x.arg[0]}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]}," f" {fp8_index(x.src[0].dtype)}, {fp8_index(x.src[0].dtype)}, 0, 0, 0, 0)" if x.arg[1][2] == 128 else None), (UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{x.arg[0]}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]}, 0, 0, 0)"), + (UPat(Ops.CONST, dtypes.fp8s, name="x"), lambda ctx,x: f"f32_to_fp8({ctx.nan}, {fp8_index(x.dtype)})" if math.isnan(x.arg) else None), + (UPat(Ops.CONST, dtypes.fp8s, arg=math.inf, name="x"), lambda ctx,x: f"f32_to_fp8({ctx.infinity}, {fp8_index(x.dtype)})"), + (UPat(Ops.CONST, dtypes.fp8s, arg=-math.inf, name="x"), lambda ctx,x: f"f32_to_fp8(-{ctx.infinity}, {fp8_index(x.dtype)})"), + (UPat(Ops.CONST, dtypes.fp8s, name="x"), lambda ctx,x: f"f32_to_fp8({x.arg}f, {fp8_index(x.dtype)})"), (UPat(Ops.CAST, dtypes.fp8s, (UPat(dtype=dtypes.float),), name="x",), lambda ctx,x: f"f32_to_fp8({ctx[x.src[0]]}, {fp8_index(x.dtype)})"), (UPat(Ops.CAST, dtypes.float, (UPat.var("y", dtypes.fp8s),), name="x",), @@ -535,7 +554,8 @@ def render_kernel(self, function_name, kernel, bufs, uops, prefix=None) -> str: if any(dt.scalar() == dtypes.half for dt in used_dtypes): prefix.append("#define half _Float16") if any(dt.scalar() in dtypes.fp8s for dt in used_dtypes): prefix += ["typedef unsigned char hip_bf8;", "typedef unsigned char hip_fp8;"] - if any(u.op is Ops.CAST and u.dtype in dtypes.fp8s and u.src[0].dtype == dtypes.float for u in uops): + if any((u.op is Ops.CAST and u.dtype in dtypes.fp8s and u.src[0].dtype == dtypes.float) or + (u.op is Ops.CONST and u.dtype in dtypes.fp8s) for u in uops): prefix.append("""static inline __attribute__((device)) unsigned char f32_to_fp8(float v, int is_bf8) { v = (((*(unsigned*)&v)&0x7F800000)!=0x7F800000)?__builtin_amdgcn_fmed3f(v,is_bf8?57344.0f:448.0f,is_bf8?-57344.0f:-448.0f) : v; return (unsigned char)(is_bf8?__builtin_amdgcn_cvt_pk_bf8_f32(v,v,0,false):__builtin_amdgcn_cvt_pk_fp8_f32(v,v,0,false));\n}""") @@ -559,6 +579,9 @@ def render_kernel(self, function_name, kernel, bufs, uops, prefix=None) -> str: for (int n = 0; n < 8; n++) { d[n] = c_frag[n*2]; } return d;\n}""") return super().render_kernel(function_name, kernel, bufs, uops, prefix) + def supported_dtypes(self): return {d for d in super().supported_dtypes() + if (d not in dtypes.fp8_ocp or self.target.arch == "gfx950") and d not in dtypes.fp8_fnuz} + class HIPCCRenderer(HIPRenderer): def __init__(self, target:Target): super().__init__(target, use_hipcc=True) @@ -567,3 +590,8 @@ def __init__(self, target:Target): super().__init__(target) from tinygrad.runtime.support.compiler_qcom import QCOMCompiler self.compiler = QCOMCompiler(target.arch) + + # QCOM compiler is flaky with half + def supported_dtypes(self): + return {d for d in Renderer.supported_dtypes(self) + if (d != dtypes.float16 or (bool(IMAGE) and bool(FLOAT16))) and d not in dtypes.fp8s+(dtypes.bfloat16,dtypes.double)} diff --git a/tinygrad/renderer/isa/__init__.py b/tinygrad/renderer/isa/__init__.py new file mode 100644 index 0000000000000..de214c76e3c74 --- /dev/null +++ b/tinygrad/renderer/isa/__init__.py @@ -0,0 +1,42 @@ +from __future__ import annotations +import itertools +from dataclasses import dataclass, field +from tinygrad.renderer import Renderer +from tinygrad.uop.ops import PatternMatcher, UOp, Ops, consumer_map_from_toposort + +@dataclass(frozen=True) +class Register: + name: str + index: int + _cons: tuple[Register, ...] = field(default_factory=tuple) + @property + def cons(self): return self._cons or (self,) + def __repr__(self): return self.name + +class IselContext: + def __init__(self, sink:UOp): + self.uses = consumer_map_from_toposort(sink.toposort()) + self.reg_n = itertools.count() + arg_order = {Ops.PARAM: 0, Ops.DEFINE_VAR: 1, Ops.SPECIAL: 2} + self.func_args = sorted([u for u in self.uses if u.op in arg_order], key=lambda k: (arg_order[k.op], k.arg)) + + def vreg(self, cons:tuple[Register, ...]|Register): + return Register(f"v{next(self.reg_n)}", 0, _cons=cons if isinstance(cons, tuple) else (cons,)) + +@dataclass +class PreRegAllocContext: + lock: UOp|None = None + clobbered: set[UOp] = field(default_factory=set) + +class ISARenderer(Renderer): + pre_isel_matcher: PatternMatcher + isel_matcher: PatternMatcher + pre_regalloc_matcher: PatternMatcher|None = None + post_regalloc_matcher: PatternMatcher + + def is_two_address(self, x:UOp) -> bool: return False + def stack_pointer(self) -> UOp: raise NotImplementedError("arch specific") + def copy(self, x:UOp, reg:Register) -> UOp: raise NotImplementedError("arch specific") + def spill(self, disp:UOp, x:UOp) -> UOp: raise NotImplementedError("arch specific") + def fill(self, disp:UOp, x:UOp, reg:Register) -> UOp: raise NotImplementedError("arch specific") + def asm_str(self, uops:list[UOp], function_name:str) -> str: raise NotImplementedError("arch specific") \ No newline at end of file diff --git a/tinygrad/renderer/isa/x86.py b/tinygrad/renderer/isa/x86.py new file mode 100644 index 0000000000000..5fe2ac1029156 --- /dev/null +++ b/tinygrad/renderer/isa/x86.py @@ -0,0 +1,905 @@ +# flake8: noqa: E702 +# allow semicolons to put multiple ops on one line +import sys, struct, functools +from typing import cast +from tinygrad.dtype import dtypes, PtrDType, DType, truncate, AddrSpace +from tinygrad.uop import FastEnum, auto, Ops, GroupOp +from tinygrad.uop.ops import UOp, UPat, PatternMatcher +from tinygrad.renderer.isa import ISARenderer, IselContext, Register, PreRegAllocContext +from tinygrad.helpers import getenv, CPU_COUNT, unwrap, Target + +# ***** X86 Ops ***** + +class X86Ops(FastEnum): + # NOTE: X86Ops with i suffix are variants that take an immediate, m suffix are variants that can write to memory instead of read from + # these aren't real instructions + FRAME_INDEX = auto(); LABEL = auto() + # index + LEA = auto() + # register / memory / immediate moves + MOV = auto(); MOVm = auto(); MOVi = auto(); MOVABS = auto() + VMOVSS = auto(); VMOVSD = auto(); VMOVUPS = auto() + VMOVSSm = auto(); VMOVSDm = auto(); VMOVUPSm = auto() + # casts + MOVZX = auto(); MOVSX = auto(); MOVSXD = auto() + VPMOVZXBW = auto(); VPMOVZXBD = auto(); VPMOVZXBQ = auto() + VPMOVZXWD = auto(); VPMOVZXWQ = auto(); VPMOVZXDQ = auto() + VPMOVSXBW = auto(); VPMOVSXBD = auto(); VPMOVSXBQ = auto() + VPMOVSXWD = auto(); VPMOVSXWQ = auto(); VPMOVSXDQ = auto() + VCVTDQ2PS = auto(); VCVTDQ2PD = auto(); VCVTTPS2DQ = auto(); VCVTTPD2DQ = auto() + VCVTPH2PS = auto(); VCVTPS2PH = auto(); VCVTPS2PD = auto(); VCVTPD2PS = auto() + VCVTSS2SD = auto(); VCVTSD2SS = auto(); VCVTSI2SS = auto(); VCVTSI2SD = auto() + VCVTTSS2SI = auto(); VCVTTSD2SI = auto() + # bitcasts + VMOVD = auto(); VMOVQ = auto(); VMOVDm = auto(); VMOVQm = auto() + # comparisons + VUCOMISS = auto(); VUCOMISD = auto() + VCMPSS = auto(); VCMPSD = auto(); VCMPPS = auto(); VCMPPD = auto() + VPCMPGTB = auto(); VPCMPGTW = auto(); VPCMPGTD = auto(); VPCMPGTQ = auto() + VPCMPEQB = auto(); VPCMPEQW = auto(); VPCMPEQD = auto(); VPCMPEQQ = auto() + SETNE = auto(); SETE = auto(); SETL = auto(); SETB = auto() + # where + CMOVNE = auto(); CMOVE = auto(); CMOVL = auto(); CMOVB = auto() + VPBLENDVB = auto(); VBLENDVPS = auto(); VBLENDVPD = auto() + # jumps + JNE = auto(); JE = auto(); JL = auto(); JB = auto(); JGE = auto(); JMP = auto() + # vectorize / gep + VSHUFPS = auto(); VSHUFPD = auto(); VINSERTPS = auto(); VPSRLDQ = auto() + VPEXTRB = auto(); VPEXTRW = auto(); VPEXTRD = auto(); VPEXTRQ = auto() + VPINSRB = auto(); VPINSRW = auto(); VPINSRD = auto(); VPINSRQ = auto() + VPBROADCASTB = auto(); VPBROADCASTW = auto(); VPBROADCASTD = auto(); VPBROADCASTQ = auto() + VBROADCASTSS = auto() + # int binary + IDIV = auto(); DIV = auto() + ADD = auto(); ADDi = auto(); SUB = auto(); SUBi = auto(); IMUL = auto(); IMULi = auto() + AND = auto(); ANDi = auto(); XOR = auto(); XORi = auto(); OR = auto(); ORi = auto() + SHL = auto(); SHLi = auto(); SHR = auto(); SHRi = auto(); SAR = auto(); SARi = auto(); CMP = auto(); CMPi = auto() + # float unary (sometimes not unary) + VROUNDSS = auto(); VROUNDSD = auto(); VROUNDPS = auto(); VROUNDPD = auto() + VSQRTSS = auto(); VSQRTSD = auto(); VSQRTPS = auto(); VSQRTPD = auto() + # float scalar / vector binary + VADDSS = auto(); VADDSD = auto(); VADDPS = auto(); VADDPD = auto() + VSUBSS = auto(); VSUBSD = auto(); VSUBPS = auto(); VSUBPD = auto() + VMULSS = auto(); VMULSD = auto(); VMULPS = auto(); VMULPD = auto() + VDIVSS = auto(); VDIVSD = auto(); VDIVPS = auto(); VDIVPD = auto() + VMAXSS = auto(); VMAXSD = auto(); VMAXPS = auto(); VMAXPD = auto() + VMINSS = auto(); VMINSD = auto(); VMINPS = auto(); VMINPD = auto() + # int vector binary + VPADDB = auto(); VPADDW = auto(); VPADDD = auto(); VPADDQ = auto() + VPSUBB = auto(); VPSUBW = auto(); VPSUBD = auto(); VPSUBQ = auto() + VPMULLW = auto(); VPMULLD = auto() + # packed bitwise + VPAND = auto(); VPOR = auto(); VPXOR = auto() + # packed variable shifts + VPSLLVD = auto(); VPSLLVQ = auto(); VPSRLVD = auto(); VPSRLVQ = auto(); VPSRAVD = auto() + # fused multiply add + VFMADD213SS = auto(); VFMADD213SD = auto(); VFMADD213PS = auto(); VFMADD213PD = auto() + # return + RET = auto() + +class X86GroupOp: + # X86Ops whose first src is also the destination + TwoAddress = {X86Ops.ADD, X86Ops.ADDi, X86Ops.AND, X86Ops.ANDi, X86Ops.XOR, X86Ops.XORi, X86Ops.OR, X86Ops.ORi, X86Ops.IMUL, + X86Ops.SUB, X86Ops.SUBi, X86Ops.SHL, X86Ops.SHLi, X86Ops.SHR, X86Ops.SHRi, X86Ops.SAR, X86Ops.SARi, + X86Ops.IDIV, X86Ops.DIV, X86Ops.VFMADD213SS, X86Ops.VFMADD213SD, X86Ops.VFMADD213PS, X86Ops.VFMADD213PD, + X86Ops.CMOVNE, X86Ops.CMOVE, X86Ops.CMOVL, X86Ops.CMOVB} + + # X86Ops whose first src can read from memory + ReadMem1st = {X86Ops.MOV, X86Ops.VMOVSS, X86Ops.VMOVSD, X86Ops.VMOVUPS, X86Ops.MOVZX, X86Ops.MOVSX, X86Ops.MOVSXD, X86Ops.VMOVD, X86Ops.VMOVQ, + X86Ops.VPMOVZXBW, X86Ops.VPMOVZXBD, X86Ops.VPMOVZXBQ, X86Ops.VPMOVZXWD, X86Ops.VPMOVZXWQ, X86Ops.VPMOVZXDQ, + X86Ops.VPMOVSXBW, X86Ops.VPMOVSXBD, X86Ops.VPMOVSXBQ, X86Ops.VPMOVSXWD, X86Ops.VPMOVSXWQ, X86Ops.VPMOVSXDQ, + X86Ops.VCVTDQ2PS, X86Ops.VCVTDQ2PD, X86Ops.VCVTTPS2DQ, X86Ops.VCVTTPD2DQ, X86Ops.VCVTTSS2SI, X86Ops.VCVTTSD2SI, + X86Ops.VCVTPH2PS, X86Ops.VCVTPS2PD, X86Ops.VCVTPD2PS, X86Ops.VROUNDPS, X86Ops.VROUNDPD, X86Ops.VSQRTPS, X86Ops.VSQRTPD, + X86Ops.VPBROADCASTB, X86Ops.VPBROADCASTW, X86Ops.VPBROADCASTD, X86Ops.VPBROADCASTQ, X86Ops.VBROADCASTSS, + X86Ops.CMPi, X86Ops.IMULi, X86Ops.LEA} + + # X86Ops whose second src can read from memory NOTE: some of these are TwoAddress so the second src is actually the first + ReadMem2nd = {X86Ops.ADD, X86Ops.SUB, X86Ops.AND, X86Ops.OR, X86Ops.XOR, X86Ops.SHL, X86Ops.SHR, X86Ops.SAR, X86Ops.IMUL, X86Ops.CMP, + X86Ops.VADDSS, X86Ops.VADDSD, X86Ops.VADDPS, X86Ops.VADDPD, X86Ops.VSUBSS, X86Ops.VSUBSD, X86Ops.VSUBPS, X86Ops.VSUBPD, + X86Ops.VMULSS, X86Ops.VMULSD, X86Ops.VMULPS, X86Ops.VMULPD, X86Ops.VDIVSS, X86Ops.VDIVSD, X86Ops.VDIVPS, X86Ops.VDIVPD, + X86Ops.VPADDB, X86Ops.VPADDW, X86Ops.VPADDD, X86Ops.VPADDQ, X86Ops.VPSUBB, X86Ops.VPSUBW, X86Ops.VPSUBD, X86Ops.VPSUBQ, + X86Ops.VPCMPEQB, X86Ops.VPCMPEQW, X86Ops.VPCMPEQD, X86Ops.VPCMPEQQ, X86Ops.VPBLENDVB, X86Ops.VBLENDVPS, X86Ops.VBLENDVPD, + X86Ops.VPCMPGTB, X86Ops.VPCMPGTW, X86Ops.VPCMPGTD, X86Ops.VPCMPGTQ, X86Ops.VCMPSS, X86Ops.VCMPSD, X86Ops.VCMPPS, X86Ops.VCMPPD, + X86Ops.VPMULLW, X86Ops.VPMULLD, X86Ops.VROUNDSS, X86Ops.VROUNDSD, X86Ops.VSQRTSS, X86Ops.VSQRTSD, X86Ops.VSHUFPS, X86Ops.VINSERTPS, + X86Ops.VPINSRB, X86Ops.VPINSRW, X86Ops.VPINSRD, X86Ops.VPINSRQ, X86Ops.VPAND, X86Ops.VPOR, X86Ops.VPXOR, X86Ops.VPSLLVD, + X86Ops.VPSLLVQ, X86Ops.VPSRLVD, X86Ops.VPSRLVQ, X86Ops.VPSRAVD, X86Ops.CMOVNE, X86Ops.CMOVE, X86Ops.CMOVL, X86Ops.CMOVB, + X86Ops.VMAXSS, X86Ops.VMAXSD, X86Ops.VMAXPS, X86Ops.VMAXPD, X86Ops.VMINSS, X86Ops.VMINSD, X86Ops.VMINPS, X86Ops.VMINPD, + X86Ops.VCVTSI2SS, X86Ops.VCVTSI2SD, X86Ops.VCVTSS2SD, X86Ops.VCVTSD2SS, X86Ops.VUCOMISS, X86Ops.VUCOMISD, X86Ops.IDIV, X86Ops.DIV, + X86Ops.VSHUFPD} + + # X86Ops whose third src can read from memory NOTE: these are TwoAddress so the third src is actually the second + ReadMem3rd = {X86Ops.VFMADD213SS, X86Ops.VFMADD213SD, X86Ops.VFMADD213PS, X86Ops.VFMADD213PD} + + # X86Ops that can write to memory + WriteMem = {X86Ops.MOVm, X86Ops.MOVi, X86Ops.VMOVSSm, X86Ops.VMOVSDm, X86Ops.VMOVUPSm, X86Ops.VMOVDm, X86Ops.VMOVQm, + X86Ops.ADDi, X86Ops.SUBi, X86Ops.ANDi, X86Ops.ORi, X86Ops.XORi, X86Ops.SHLi, X86Ops.SHRi, X86Ops.SARi, X86Ops.SETNE, + X86Ops.SETE, X86Ops.SETL, X86Ops.SETB, X86Ops.VCVTPS2PH, X86Ops.VPEXTRB, X86Ops.VPEXTRW, X86Ops.VPEXTRD, X86Ops.VPEXTRQ} + + # X86Ops that read flags + ReadFlags = {X86Ops.CMOVB, X86Ops.CMOVL, X86Ops.CMOVE, X86Ops.CMOVNE, X86Ops.SETB, X86Ops.SETL, X86Ops.SETE, X86Ops.SETNE, X86Ops.JB, X86Ops.JL, + X86Ops.JE, X86Ops.JNE, X86Ops.JGE} + + # X86Ops that write flags or can modify flags to undefined values + WriteFlags = {X86Ops.CMP, X86Ops.CMPi, X86Ops.ADD, X86Ops.ADDi, X86Ops.SUB, X86Ops.SUBi, X86Ops.IMUL, X86Ops.IMULi, X86Ops.IDIV, X86Ops.DIV, + X86Ops.SHL, X86Ops.SHLi, X86Ops.SHR, X86Ops.SHRi, X86Ops.SAR, X86Ops.SARi, X86Ops.AND, X86Ops.ANDi, X86Ops.XOR, X86Ops.XORi, + X86Ops.OR, X86Ops.ORi, X86Ops.VUCOMISS, X86Ops.VUCOMISD} + + # X86Ops whose first src is the rm field + Rm1st = ReadMem1st | (ReadMem2nd & TwoAddress) | {X86Ops.VPSRLDQ} + + # X86Ops whose second src is the rm field + Rm2nd = ReadMem2nd | (ReadMem3rd & TwoAddress) + + All = set(X86Ops) + +# ***** X86 legalization ***** + +extra_matcher = PatternMatcher([ + # bool CMPNE is XOR, bool CMPEQ is XOR+XOR, bool CMPLT is XOR+AND + (UPat.var('x', dtypes.bool).ne(UPat.var('y')), lambda x,y: x^y), + (UPat.var('x', dtypes.bool).alu(Ops.CMPEQ, UPat.var('y')), lambda x,y: (x^y)^True), + (UPat.var('x', dtypes.bool)> 1).cast(dtypes.int64).cast(x.dtype) * 2 + (y & 1).cast(dtypes.int64).cast(x.dtype)), + # no int8 mul or cmove, cast to int16 + (UPat.var("a", dtypes.int8s) * UPat.var("b"), lambda a,b: (a.cast(dtypes.int16) * b.cast(dtypes.int16)).cast(a.dtype)), + (UPat.var("m").where(UPat.var("a", (dtypes.bool,)+dtypes.int8s), UPat.var("b")), + lambda m,a,b: m.where(a.cast(dtypes.int16), b.cast(dtypes.int16)).cast(a.dtype) if a.dtype.count == 1 else None), + # float16 alus are done in float32 + (UPat(GroupOp.ALU, dtypes.float16, name="x"), lambda x: UOp(x.op, dtypes.float.vec(x.dtype.count), + tuple(s.cast(dtypes.float) if s.dtype != dtypes.bool else s for s in x.src)).cast(x.dtype)), + (UPat(GroupOp.Comparison, src=(UPat.var("a", dtypes.float16), UPat.var("b")), name="x"), + lambda x,a,b: UOp(x.op, x.dtype, (a.cast(dtypes.float32), b.cast(dtypes.float32))).cast(x.dtype)), + # no cmpne for packed ints, y != x => !(y==x) + (UPat(Ops.CMPNE, src=(UPat.var("y", dtypes.ints), UPat.var("x")), name="cmp"), + lambda y,x,cmp: UOp(Ops.CMPEQ, cmp.dtype, (y,x))^True if y.dtype.count > 1 else None), + # float where expects a mask + (UPat.var("m", dtypes.bool).where(UPat.var("a", dtypes.floats), UPat.var("b")), + lambda m,a,b: m.cast(a.dtype).ne(0).where(a, b) if m.src[0].dtype not in dtypes.floats else None), + # rewrite -x -> 0 - x + (UPat(Ops.NEG, name="x"), lambda x: UOp(Ops.SUB, x.dtype, (x.const_like(0),) + x.src)), + # TODO: add support for mod, requires support for accessing the 2nd+ reg of a multi output instruction + (UPat(Ops.CMOD, src=(UPat.var("x"), UPat.var("y"))), lambda x,y: x - y * x.alu(Ops.CDIV, y)), +]) + +# ***** X86 pre instruction selection ***** + +def gated_load(ctx, base:UOp, idx:UOp, cast:UOp, alt:UOp, gate:UOp, x:UOp): + local = UOp(Ops.DEFINE_LOCAL, base.dtype.base.ptr(x.dtype.count, AddrSpace.LOCAL), arg=next(ctx)) + local_idx = local.index(UOp.const(dtypes.int32, 0), ptr=True) + ptr = gate.where(base.index(idx, ptr=True), local_idx).after((local_idx if x.dtype.count == 1 else local).store(alt)) + return ptr.cast(cast.dtype).load(dtype=x.dtype) + +def gated_store(base:UOp, idx:UOp, cast:UOp, gate:UOp, val:UOp): + local = UOp(Ops.DEFINE_LOCAL, base.dtype.base.ptr(val.dtype.count, AddrSpace.LOCAL), arg=-1) + ptr = gate.where(base.index(idx, ptr=True), local.index(UOp.const(dtypes.int32, 0), ptr=True)) + return ptr.cast(cast.dtype).store(val) + +# these must be done in a separate matcher because they violate the spec +pre_isel_matcher = PatternMatcher([ + # zero extending scalar 32bit int is a noop + (UPat.var("y", dtypes.uint32).cast(dtypes.int64s, name="x"), lambda y,x: x.replace(op=Ops.NOOP) if y.dtype.count == 1 else None), + # cast between signed and unsigned int is a noop + (UPat.var("y", dtypes.ints+(dtypes.bool,)).cast(dtypes.ints, name="x"), + lambda y,x: x.replace(op=Ops.NOOP) if x.dtype.itemsize == y.dtype.itemsize else None), + # cast to < scalar int is a noop + (UPat.var("y", dtypes.ints).cast(dtypes.ints, name="x"), + lambda y,x: x.replace(op=Ops.NOOP) if x.dtype.itemsize < y.dtype.itemsize and y.dtype.count == 1 else None), + # bitcasts between scalar floats and ints are real, rest are noops + (UPat.var("y").bitcast().named("x"), lambda y,x: None if y.dtype in dtypes.floats and x.dtype in dtypes.ints or \ + y.dtype in dtypes.ints and x.dtype in dtypes.floats else x.replace(op=Ops.NOOP)), + # noop of a noop is removed + (UPat(Ops.NOOP, src=(UPat(Ops.NOOP),), name="x"), lambda x: x.replace(src=x.src[0].src)), + # moving elements of a single register to another without shuffling is a noop + (UPat(Ops.STACK, src=(UPat.var("y"),), allow_any_len=True, name="x"), + lambda y,x: UOp(Ops.NOOP, x.dtype, y.src) if all(s.op is Ops.GEP and s.src == y.src and s.arg[0] == i for i,s in enumerate(x.src)) else None), + # gated load/store become a conditional move on the index, the load/store are unconditional + (UPat.var("base").index(UPat.var("idx")).or_casted(name="cast").load(UPat.var("alt"), UPat.var("gate"), name="x"), gated_load), + (UPat.var("base").index(UPat.var("idx")).or_casted(name="cast").store(UPat.var("val"), UPat.var("gate")), gated_store), + # TODO: remove this once we allow all flag producing ops in cmove + # if gate in scalar int cmove is not a comparison need to add one to set the flag + (UPat.var("m", dtypes.bool).where(UPat.var("a"), UPat.var("b")), + lambda m,a,b: m.ne(0).where(a,b) if m.op not in GroupOp.Comparison and a.dtype.count == 1 else None), +]) + +# ***** X86 registers ***** + +RAX = Register("rax", 0) +RCX = Register("rcx", 1) +RDX = Register("rdx", 2) +RBX = Register("rbx", 3) +RSP = Register("rsp", 4) +RBP = Register("rbp", 5) +RSI = Register("rsi", 6) +RDI = Register("rdi", 7) +GPR = (RAX, RCX, RDX, RBX, RSP, RBP, RSI, RDI) + tuple(Register(f"r{i}", i) for i in range(8, 16)) +XMM = tuple(Register(f"xmm{i}", i) for i in range(16)) +# gprs you can write to +WGPR = tuple(r for r in GPR if r != RSP) + +CALLEE_SAVED = (RBX, RBP, GPR[12], GPR[13], GPR[14], GPR[15]) + ((RSI, RDI) + XMM[6:16] if sys.platform == "win32" else ()) + +reg_strs = {"rax": {4:"eax", 2:"ax", 1:"al"}, "rcx": {4:"ecx", 2:"cx", 1:"cl"}, "rdx": {4:"edx", 2:"dx", 1:"dl"}, "rbx": {4:"ebx", 2:"bx", 1:"bl"}, + "rsp": {4:"esp", 2:"sp", 1:"spl"}, "rbp": {4:"ebp", 2:"bp", 1:"bpl"}, "rsi": {4:"esi", 2:"si", 1:"sil"}, "rdi": {4:"edi", 2:"di", 1:"dil"}, + **{f"r{i}": {4:f"r{i}d", 2:f"r{i}w", 1:f"r{i}b"} for i in range(8, 16)}, **{f"xmm{i}": {64:f"zmm{i}", 32:f"ymm{i}"} for i in range(16)}} + +# ***** X86 instruction selection ***** +# if s is used multiple times we don't fold +def is_foldable(ctx:IselContext, x:UOp, s:UOp) -> bool: return len(ctx.uses[s]) == x.src.count(s) == 1 +def base(x:UOp, i:int) -> UOp: return s.src[0] if (s:=x.src[i]).op is Ops.GEP else s +def lane(x:UOp, i:int) -> int: return s.arg[0] if (s:=x.src[i]).op is Ops.GEP else 0 +def to_int(dt:DType): return {dtypes.float16: dtypes.int16, dtypes.float32: dtypes.int32, dtypes.float64: dtypes.int64}[dt] +def def_reg(dt:DType, reg:Register|None=None) -> UOp: return UOp(Ops.DEFINE_REG, dt, tag=None if reg is None else (reg,)) +def imm(dt:DType, v:int) -> UOp: return UOp.const(dt, truncate[dt](v)).rtag() +def to_imm(c:UOp) -> UOp|None: + if c.op is not Ops.CONST: return None + if c.dtype is dtypes.int64: return imm(dtypes.int32, c.arg) if not c.overflows(dtypes.int32) else None + if c.dtype is dtypes.uint64: return imm(dtypes.uint32, c.arg) if not c.overflows(dtypes.uint32) else None + if c.dtype in dtypes.ints+(dtypes.bool,): return imm(c.dtype, c.arg) + return None +def cmp(x:UOp) -> UOp: + if x.src[0].dtype is dtypes.float32: return x.ins(X86Ops.VUCOMISS, dtype=dtypes.void) + if x.src[0].dtype is dtypes.float64: return x.ins(X86Ops.VUCOMISD, dtype=dtypes.void) + return x.ins(X86Ops.CMP, dtype=dtypes.void) if (i:=to_imm(x.src[1])) is None else x.ins(X86Ops.CMPi, dtype=dtypes.void, src=(x.src[0], i)) +def vcmp(x:UOp) -> UOp: + v = imm(dtypes.uint8, {Ops.CMPLT: 1, Ops.CMPNE: 4, Ops.CMPEQ: 0}[x.op]) + if x.dtype.scalar() is dtypes.float32: return x.ins(X86Ops.VCMPSS if x.dtype.count == 1 else X86Ops.VCMPPS, src=x.src + (v,)) + return x.ins(X86Ops.VCMPSD if x.dtype.count == 1 else X86Ops.VCMPPD, src=x.src + (v,)) + +# vshufps xmm2, xmm0, xmm1, imm +# for 128 bit xmm2 selects its lower 2 32 bits from xmm0 and its upper 2 32 bits from xmm1 according to imm +# for 256 bit ymm2 repeats the shuffle for its upper 128 bits selecting from the upper 128 bits of ymm0 and ymm1 +def vshufps(x:UOp) -> UOp|None: + a, b = base(x, 0), base(x, 2) + if not (a is base(x, 1) and b is base(x, 3)) or any(lane(x, i) > 3 for i in range(4)): return None + if len(x.src) == 8: + if not (a is base(x, 4) is base(x, 5) and b is base(x, 6) is base(x, 7)) or any(lane(x, i+4) != lane(x, i)+4 for i in range(4)): return None + return x.ins(X86Ops.VSHUFPS, src=(a, b, imm(dtypes.uint8, sum(lane(x, i) << 2*i for i in range(4))))) + +# vshufpd xmm2, xmm0, xmm1, imm +# for 128 bit xmm2 selects its lower 64 bits from xmm0 and its upper 64 bits from xmm1 according to imm +# for 256 bit ymm2 also selects its upper 128 bits from the upper 128 bits of ymm0 and ymm1 following the same constraint +def vshufpd(x:UOp) -> UOp|None: + a, b = base(x, 0), base(x, 1) + if lane(x, 0) > 1 or lane(x, 1) > 1: return None + if len(x.src) == 4 and not (a is base(x, 2) and b is base(x, 3) and lane(x, 2) > 1 and lane(x, 3) > 1): return None + return x.ins(X86Ops.VSHUFPD, src=(a, b, imm(dtypes.uint8, sum(lane(x, i) << i for i in range(len(x.src)))))) + +# vinsertps xmm2, xmm0, xmm1, imm +# inserts any 32 bit element in xmm1 into any position in xmm0 according to immm, result is written to xmm2 +# this is the fallback slow case for when you can't match more a powerful shuffle +def vinsertps(x:UOp) -> UOp: + def _insert(ret:UOp, i:int) -> UOp: + s, v = base(x, i), lane(x, i) + # moving the 0th element into the 0th position does nothing + return s if i == v == 0 else x.ins(X86Ops.VINSERTPS, src=(ret, s, imm(dtypes.uint8, v << 6 | i << 4))) + return functools.reduce(_insert, range(len(x.src)), def_reg(x.dtype)) + +# vpinsq xmm2, xmm0, rax, imm +# inserts element in rax into any position in xmm0, result is written to xmm2 according to imm +def vpins(x:UOp) -> UOp: + op = {1: X86Ops.VPINSRB, 2: X86Ops.VPINSRW, 4: X86Ops.VPINSRD, 8: X86Ops.VPINSRQ}[x.dtype.scalar().itemsize] + return functools.reduce(lambda ret,i: x.ins(op, src=(ret, x.src[i], imm(dtypes.uint8, i))), range(len(x.src)), def_reg(x.dtype)) + +# vpbroadcastd xmm1, xmm0 +# inserts scalar int in xmm0 into all lanes of xmm1 +def vpbroadcast(ctx:IselContext, x:UOp, y:UOp) -> UOp: + n = x.ins({1: X86Ops.VPBROADCASTB, 2: X86Ops.VPBROADCASTW, 4: X86Ops.VPBROADCASTD, 8: X86Ops.VPBROADCASTQ}[y.dtype.itemsize], src=(y,)) + if y.op is Ops.LOAD and len(y.src) == 1 and is_foldable(ctx, n, y): return n + # if there isn't a load we can fold we need to move y from gpr to xmm + # this is hacky but required because int.vec(1) isn't supported + y = y if y.dtype.itemsize > 1 else y.cast(dtypes.int16) + return n.replace(src=(y.bitcast({2:dtypes.float16, 4:dtypes.float32, 8:dtypes.float64}[y.dtype.itemsize]),)) + +# we don't call ctx.vreg on the srcs to avoid duplicates, a rewrite will assign the tuple of valid registers to a vreg +def idiv(ctx:IselContext, x:UOp) -> UOp: + op = X86Ops.DIV if x.dtype in dtypes.uints else X86Ops.IDIV + # for >8bit need to zero/sign extend rax to rdx + if x.dtype in dtypes.int8s: ext = [] + elif x.dtype in dtypes.uints: ext = [x.ins(X86Ops.MOVi, src=(imm(min(dtypes.uint32, x.dtype), 0),), tag=(RDX,))] + else: ext = [x.ins(X86Ops.SARi, src=(x.src[0], imm(dtypes.uint8, x.dtype.itemsize * 8 - 1)), tag=(RDX,))] + # for 8bit need to zero/sign extend al to ah + if x.dtype is dtypes.uint8: dividend = UOp(Ops.INS, arg=X86Ops.MOVZX, dtype=dtypes.int16, src=(x.src[0],), tag=(RAX,)) + elif x.dtype is dtypes.int8: dividend = UOp(Ops.INS, arg=X86Ops.MOVSX, dtype=dtypes.int16, src=(x.src[0],), tag=(RAX,)) + else: dividend = x.ins(X86Ops.MOV, src=(x.src[0],), tag=(RAX,)) + # divisor can't be in rax or rdx + divisor = x.ins(X86Ops.MOV, src=(x.src[1],), tag=tuple(r for r in WGPR if r not in (RAX, RDX))) + # for >8bit both rax and rdx are written to + defs = (ctx.vreg(RAX),) if x.dtype in dtypes.int8s else (ctx.vreg(RAX), ctx.vreg(RDX)) + idiv = x.ins(op, src=(dividend, divisor) + tuple(ext), tag=defs) + # this move "cleanses" the register constraints (rax/rdx) of idiv as that only applies on definition and not on the uses of idiv + return x.ins(X86Ops.MOV, src=(idiv,)) + +def fold_address(x:UOp) -> tuple[UOp, UOp, UOp]: + def _disp(v:int) -> UOp: return imm(dtypes.int32 if abs(v) > dtypes.int8.max else dtypes.int8, v) + def _cast(v:UOp) -> UOp: return v.cast(dtypes.int64) if v.vmin < 0 else v + if x.op is not Ops.INDEX: return (x, UOp(Ops.NOOP), _disp(0)) + base, idx = x.src + disp_scale = base.dtype.itemsize if isinstance(base.dtype, PtrDType) else 1 + if idx.op is Ops.ADD and idx.src[1].op is Ops.CONST: return (base, _cast(idx.src[0]), _disp(idx.src[1].arg * disp_scale)) + if idx.op is Ops.CONST: return (base, UOp(Ops.NOOP), _disp(idx.arg * disp_scale)) + return (base, _cast(idx), _disp(0)) + +def abi(ctx:IselContext, x:UOp) -> UOp|None: + if isinstance(x.tag, tuple): return None + i = ctx.func_args.index(x) + def _stack_arg(disp:int): return (def_reg(dtypes.uint64, RSP), UOp(Ops.NOOP), UOp(Ops.INS, arg=X86Ops.FRAME_INDEX, dtype=dtypes.int32, tag=disp)) + if sys.platform == "win32": src = (x.replace(tag=((RCX, RDX, GPR[8], GPR[9])[i],)),) if i < 4 else _stack_arg((i-3)*8+32) + else: src = (x.replace(tag=((RDI, RSI, RDX, RCX, GPR[8], GPR[9])[i],)),) if i < 6 else _stack_arg((i-5)*8) + # this move "cleanses" the abi register constraint + return x.ins(X86Ops.MOV, src=src) + +def alloc_vregs(ctx:IselContext, x:UOp) -> UOp|None: + # real registers + if x.op is Ops.DEFINE_REG and x.tag is not None: return None + # this is an immediate + if x.arg is X86Ops.FRAME_INDEX: return None + # no register definition + if x.dtype is dtypes.void: return None + # already allocated vregs + if isinstance(x.tag, tuple) and x.tag[0]._cons: return None + # allocate vreg definitions + defs = [] + if isinstance(x.tag, tuple): defs = [ctx.vreg(x.tag)] + elif x.dtype in dtypes.ints+(dtypes.bool,) or isinstance(x.dtype, PtrDType): defs = [ctx.vreg(WGPR)] + elif x.dtype in dtypes.floats or x.dtype.count > 1: defs = [ctx.vreg(XMM)] + # TODO: add this once the scheduler can track register pressure + # if x.arg in X86GroupOp.WriteFlags: defs.append(ctx.vreg(RFLAGS)) + return x.replace(tag=tuple(defs)) + +dts = dtypes.ints + (dtypes.bool, dtypes.float16, dtypes.float32, dtypes.float64) +dt_16bit = tuple(dt.vec(l) for dt in dts for l in [2,1] if l*dt.itemsize == 2 and dt not in dtypes.int16s) +dt_32bit = tuple(dt.vec(l) for dt in dts for l in [4,2,1] if l*dt.itemsize == 4 and dt not in dtypes.int32s) +dt_64bit = tuple(dt.vec(l) for dt in dts for l in [8,4,2,1] if l*dt.itemsize == 8 and dt not in dtypes.int64s) +dt_128bit = tuple(dt.vec(l) for dt in dts for l in [16,8,4,2,1] if l*dt.itemsize == 16) + +isel_matcher = PatternMatcher([ + # **** Op -> Op **** + # cast to pointer is a noop + (UPat.var("y").cast(name="x"), lambda y,x: y if isinstance(x.dtype, PtrDType) or y.dtype == dtypes.void else None), + # float gep(0) is a noop as it just moves the 0th element from one xmm register to another + # this is done here to not interfere with shuffles + (UPat(dtype=dtypes.floats).gep(0, name="x"), lambda x: x.replace(op=Ops.NOOP, arg=None)), + # range is lowered to acc, cmp, jmp after regalloc + (UPat(Ops.RANGE, src=(UPat.cvar("c"),), allow_any_len=True, name="x"), lambda c,x: x.replace(src=(imm(c.dtype, c.arg),) + x.src[1:])), + (UPat(Ops.RANGE, name="x"), lambda ctx,x: x.replace(tag=(ctx.vreg(WGPR),)) if not isinstance(x.tag, tuple) else None), + # **** Op -> X86Op **** + # add callee saved registers to the RET, these will be scheduled at the top of the kernel and will be saved/restored if they are used in regalloc + # so regalloc builds the prologue/epilogue naturally + (UPat(Ops.SINK, name="x"), lambda x: + x.replace(src=(x.ins(X86Ops.RET, src=x.src + tuple(def_reg(dtypes.uint64 if r in GPR else dtypes.float64.vec(2), r) for r in CALLEE_SAVED)),)) \ + if not x.src or x.src[0].arg is not X86Ops.RET else None), + # function abi constraints + (UPat((Ops.PARAM, Ops.DEFINE_VAR, Ops.SPECIAL), name="x"), abi), + # these are treated the same for now + (UPat(Ops.DEFINE_REG, name="x"), lambda x: + x.replace(op=Ops.DEFINE_LOCAL, dtype=x.dtype.base.ptr(x.dtype.size, AddrSpace.LOCAL)) if isinstance(x.arg, int) else None), + # constants that can't be immediates, move them to registers + (UPat.cvar("x", dtypes.int64s), lambda x: x.ins(X86Ops.MOVABS, src=(imm(x.dtype, x.arg),)) if not x.tag else None), + (UPat.cvar("x", dtypes.ints+(dtypes.bool,)), lambda x: x.ins(X86Ops.MOVi, src=(imm(x.dtype, x.arg),)) if not x.tag else None), + (UPat.cvar("x", dtypes.floats), lambda x: + UOp.const(dt:=to_int(x.dtype), struct.unpack(dt.fmt, struct.pack(x.dtype.fmt, x.arg))[0]).bitcast(x.dtype) if not x.tag else None), + # TODO: these should use a.maximum(b) / a.minimum(b) + ((UPat.var("a") < UPat.var("b")).where(UPat.var("b", dtypes.float32), UPat.var("a")), lambda a,b: + a.ins(X86Ops.VMAXSS if a.dtype.count == 1 else X86Ops.VMAXPS, src=(a, b))), + ((UPat.var("a") < UPat.var("b")).where(UPat.var("b", dtypes.float64), UPat.var("a")), lambda a,b: + a.ins(X86Ops.VMAXSD if a.dtype.count == 1 else X86Ops.VMAXPD, src=(a, b))), + ((UPat.var("a") < UPat.var("b")).where(UPat.var("a", dtypes.float32), UPat.var("b")), lambda a,b: + a.ins(X86Ops.VMINSS if a.dtype.count == 1 else X86Ops.VMINPS, src=(a, b))), + ((UPat.var("a") < UPat.var("b")).where(UPat.var("a", dtypes.float64), UPat.var("b")), lambda a,b: + a.ins(X86Ops.VMINSD if a.dtype.count == 1 else X86Ops.VMINPD, src=(a, b))), + # conditional moves that use masks NOTE: these currently assume a mask producing cmp exists + (UPat.var("m").where(UPat.var("a", dtypes.ints), UPat.var("b")), lambda m,a,b: + a.ins(X86Ops.VPBLENDVB, src=(b, a, m.replace(dtype=m.src[0].dtype))) if a.dtype.count > 1 else None), + (UPat.var("m").where(UPat.var("a", dtypes.float32), UPat.var("b")), lambda m,a,b: + a.ins(X86Ops.VBLENDVPS, src=(b, a, m.replace(dtype=m.src[0].dtype)))), + (UPat.var("m").where(UPat.var("a", dtypes.float64), UPat.var("b")), lambda m,a,b: + a.ins(X86Ops.VBLENDVPD, src=(b, a, m.replace(dtype=m.src[0].dtype)))), + # in this case we have a mask producing comparison whose user expects a bool, so we convert to bool + (UPat(GroupOp.Comparison, dtypes.bool, (UPat.var("y", (dtypes.float32, dtypes.float64)), UPat()), name="x"), lambda y,x: + x.replace(dtype=y.dtype).bitcast(to_int(y.dtype)).bitwise_and(1).f(Ops.NOOP, dtype=dtypes.bool)), + # conditional moves that use flags + (UPat(Ops.CMPLT, src=(UPat(dtype=dtypes.sints), UPat()), name="m").where(UPat.var("a"), UPat.var("b")), lambda m,a,b: + a.ins(X86Ops.CMOVL, src=(b, a, cmp(m)))), + (UPat(Ops.CMPLT, name="m").where(UPat.var("a"), UPat.var("b")), lambda m,a,b: a.ins(X86Ops.CMOVB, src=(b, a, cmp(m)))), + (UPat(Ops.CMPEQ, name="m").where(UPat.var("a"), UPat.var("b")), lambda m,a,b: a.ins(X86Ops.CMOVE, src=(b, a, cmp(m)))), + (UPat(Ops.CMPNE, name="m").where(UPat.var("a"), UPat.var("b")), lambda m,a,b: a.ins(X86Ops.CMOVNE, src=(b, a, cmp(m)))), + # jumps, use flags + (UPat(Ops.IF, src=(UPat(Ops.CMPLT, src=(UPat(dtype=dtypes.uints), UPat()), name="y"),), name="x"), lambda y,x: x.ins(X86Ops.JB, src=(cmp(y),))), + (UPat(Ops.IF, src=(UPat(Ops.CMPLT, name="y"),), name="x"), lambda y,x: x.ins(X86Ops.JL, src=(cmp(y),))), + (UPat(Ops.IF, src=(UPat(Ops.CMPEQ, name="y"),), name="x"), lambda y,x: x.ins(X86Ops.JE, src=(cmp(y),))), + (UPat(Ops.IF, src=(UPat(Ops.CMPNE, name="y"),), name="x"), lambda y,x: x.ins(X86Ops.JNE, src=(cmp(y),))), + # comparisons whose user doesn't use the flag, move flag result to register + (UPat(Ops.CMPLT, dtypes.bool, (UPat(dtype=dtypes.uints), UPat()), name="x"), lambda x: x.ins(X86Ops.SETB, src=(cmp(x),))), + (UPat(Ops.CMPLT, dtypes.bool, name="x"), lambda x: x.ins(X86Ops.SETL, src=(cmp(x),))), + (UPat(Ops.CMPEQ, dtypes.bool, name="x"), lambda x: x.ins(X86Ops.SETE, src=(cmp(x),))), + (UPat(Ops.CMPNE, dtypes.bool, name="x"), lambda x: x.ins(X86Ops.SETNE, src=(cmp(x),))), + # comparisons that produce masks (these aren't bool dtype) + (UPat(GroupOp.Comparison, src=(UPat(dtype=(dtypes.float32, dtypes.float64)), UPat()), name="x"), vcmp), + (UPat(Ops.CMPEQ, src=(UPat(dtype=dtypes.int8s), UPat()), name="x"), lambda x: x.ins(X86Ops.VPCMPEQB)), + (UPat(Ops.CMPEQ, src=(UPat(dtype=dtypes.int16s), UPat()), name="x"), lambda x: x.ins(X86Ops.VPCMPEQW)), + (UPat(Ops.CMPEQ, src=(UPat(dtype=dtypes.int32s), UPat()), name="x"), lambda x: x.ins(X86Ops.VPCMPEQD)), + (UPat(Ops.CMPEQ, src=(UPat(dtype=dtypes.int64s), UPat()), name="x"), lambda x: x.ins(X86Ops.VPCMPEQQ)), + (UPat(Ops.CMPLT, src=(UPat.var("a", dtypes.int8s), UPat.var("b")), name="x"), lambda a,b,x: x.ins(X86Ops.VPCMPGTB, src=(b, a))), + (UPat(Ops.CMPLT, src=(UPat.var("a", dtypes.int16s), UPat.var("b")), name="x"), lambda a,b,x: x.ins(X86Ops.VPCMPGTW, src=(b, a))), + (UPat(Ops.CMPLT, src=(UPat.var("a", dtypes.int32s), UPat.var("b")), name="x"), lambda a,b,x: x.ins(X86Ops.VPCMPGTD, src=(b, a))), + (UPat(Ops.CMPLT, src=(UPat.var("a", dtypes.int64s), UPat.var("b")), name="x"), lambda a,b,x: x.ins(X86Ops.VPCMPGTQ, src=(b, a))), + # float unary + (UPat.var("y", dtypes.float32).sqrt().named("x"), lambda y,x: x.ins(X86Ops.VSQRTSS, src=(y, y)) if x.dtype.count == 1 else x.ins(X86Ops.VSQRTPS)), + (UPat.var("y", dtypes.float64).sqrt().named("x"), lambda y,x: x.ins(X86Ops.VSQRTSD, src=(y, y)) if x.dtype.count == 1 else x.ins(X86Ops.VSQRTPD)), + (UPat.var("y", dtypes.float32).trunc().named("x"), lambda y,x: + x.ins(X86Ops.VROUNDSS, src=(y, y, imm(dtypes.uint8, 3))) if x.dtype.count == 1 else x.ins(X86Ops.VROUNDPS, src=(y, imm(dtypes.uint8, 3)))), + (UPat.var("y", dtypes.float64).trunc().named("x"), lambda y,x: + x.ins(X86Ops.VROUNDSD, src=(y, y, imm(dtypes.uint8, 3))) if x.dtype.count == 1 else x.ins(X86Ops.VROUNDPD, src=(y, imm(dtypes.uint8, 3)))), + # shufles + (UPat.var("y", dtypes.float32).broadcast(name="x"), lambda y,x: x.ins(X86Ops.VBROADCASTSS, src=(y,))), + # for float16 we route the srcs through gprs unless we can fold them, this is suboptimal for values in xmms, in that case we want vpunpcklwd + (UPat(Ops.STACK, dtypes.float16, name="x"), lambda ctx,x: + vpins(x.replace(src=tuple(s if s.op is Ops.LOAD and is_foldable(ctx, x, s) else s.bitcast(dtypes.int16) for s in x.src)))), + (UPat(Ops.STACK, (dtypes.float32.vec(4), dtypes.float32.vec(8)), name="x"), vshufps), + (UPat(Ops.STACK, (dtypes.float64.vec(2), dtypes.float64.vec(4)), name="x"), vshufpd), + (UPat(Ops.STACK, dtypes.float32, name="x"), vinsertps), + (UPat.var("y", dtypes.ints+(dtypes.bool,)).broadcast(name="x"), vpbroadcast), + (UPat(Ops.STACK, dtypes.ints+(dtypes.bool,), name="x"), vpins), + # gep + (UPat.var("y", dtypes.int8s+(dtypes.bool,)).gep(name="x"), lambda y,x: x.ins(X86Ops.VPEXTRB, src=(y, imm(dtypes.uint8, x.arg[0])))), + (UPat.var("y", dtypes.int16s).gep(name="x"), lambda y,x: x.ins(X86Ops.VPEXTRW, src=(y, imm(dtypes.uint8, x.arg[0])))), + (UPat.var("y", dtypes.int32s).gep(name="x"), lambda y,x: x.ins(X86Ops.VPEXTRD, src=(y, imm(dtypes.uint8, x.arg[0])))), + (UPat.var("y", dtypes.int64s).gep(name="x"), lambda y,x: x.ins(X86Ops.VPEXTRQ, src=(y, imm(dtypes.uint8, x.arg[0])))), + (UPat.var("y", dtypes.floats).gep(name="x"), lambda y,x: x.ins(X86Ops.VPSRLDQ, src=(y, imm(dtypes.uint8, x.arg[0] * x.dtype.itemsize)))), + # fused multiply add + ((UPat(Ops.MUL, dtypes.float32, name="a") + UPat.var("b")).named("c"), lambda ctx,a,b,c: + a.ins(X86Ops.VFMADD213SS if a.dtype.count == 1 else X86Ops.VFMADD213PS, src=(*a.src, b)) if is_foldable(ctx, c, a) else None), + ((UPat(Ops.MUL, dtypes.float64, name="a") + UPat.var("b")).named("c"), lambda ctx,a,b,c: + a.ins(X86Ops.VFMADD213SD if a.dtype.count == 1 else X86Ops.VFMADD213PD, src=(*a.src, b)) if is_foldable(ctx, c, a) else None), + # packed bitwise + ((UPat() & UPat()).named("x"), lambda x: x.ins(X86Ops.VPAND) if x.dtype.count > 1 else None), + ((UPat() | UPat()).named("x"), lambda x: x.ins(X86Ops.VPOR) if x.dtype.count > 1 else None), + ((UPat() ^ UPat()).named("x"), lambda x: x.ins(X86Ops.VPXOR) if x.dtype.count > 1 else None), + # packed int binary + ((UPat(dtype=dtypes.int32s) << UPat()).named("x"), lambda x: x.ins(X86Ops.VPSLLVD) if x.dtype.count > 1 else None), + ((UPat(dtype=dtypes.int64s) << UPat()).named("x"), lambda x: x.ins(X86Ops.VPSLLVQ) if x.dtype.count > 1 else None), + ((UPat(dtype=dtypes.uint32) >> UPat()).named("x"), lambda x: x.ins(X86Ops.VPSRLVD) if x.dtype.count > 1 else None), + ((UPat(dtype=dtypes.uint64) >> UPat()).named("x"), lambda x: x.ins(X86Ops.VPSRLVQ) if x.dtype.count > 1 else None), + ((UPat(dtype=dtypes.int32) >> UPat()).named("x"), lambda x: x.ins(X86Ops.VPSRAVD) if x.dtype.count > 1 else None), + ((UPat(dtype=dtypes.int8s) + UPat()).named("x"), lambda x: x.ins(X86Ops.VPADDB) if x.dtype.count > 1 else None), + ((UPat(dtype=dtypes.int16s) + UPat()).named("x"), lambda x: x.ins(X86Ops.VPADDW) if x.dtype.count > 1 else None), + ((UPat(dtype=dtypes.int32s) + UPat()).named("x"), lambda x: x.ins(X86Ops.VPADDD) if x.dtype.count > 1 else None), + ((UPat(dtype=dtypes.int64s) + UPat()).named("x"), lambda x: x.ins(X86Ops.VPADDQ) if x.dtype.count > 1 else None), + (UPat(Ops.SUB, dtypes.int8s, name="x"), lambda x: x.ins(X86Ops.VPSUBB) if x.dtype.count > 1 else None), + (UPat(Ops.SUB, dtypes.int16s, name="x"), lambda x: x.ins(X86Ops.VPSUBW) if x.dtype.count > 1 else None), + (UPat(Ops.SUB, dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPSUBD) if x.dtype.count > 1 else None), + (UPat(Ops.SUB, dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPSUBQ) if x.dtype.count > 1 else None), + (UPat(Ops.MUL, dtypes.int16s, name="x"), lambda x: x.ins(X86Ops.VPMULLW) if x.dtype.count > 1 else None), + (UPat(Ops.MUL, dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPMULLD) if x.dtype.count > 1 else None), + # scalar int binary + ((UPat(dtype=dtypes.ints).alu(Ops.CDIV, UPat())).named("x"), idiv), + # scalar int binary with immediate + (UPat.var("a", dtypes.ints) << UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SHLi, src=(a, imm(dtypes.uint8, c.arg)))), + (UPat.var("a", dtypes.uints) >> UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SHRi, src=(a, imm(dtypes.uint8, c.arg)))), + (UPat.var("a", dtypes.sints) >> UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SARi, src=(a, imm(dtypes.uint8, c.arg)))), + (UPat.var("a", dtypes.ints) + UPat.cvar("c"), lambda a,c: a.ins(X86Ops.ADDi, src=(a, i)) if (i:=to_imm(c)) is not None else None), + (UPat.var("a", dtypes.ints) * UPat.cvar("c"), lambda a,c: a.ins(X86Ops.IMULi, src=(a, i)) if (i:=to_imm(c)) is not None else None), + (UPat.var("a", dtypes.ints+(dtypes.bool,)) & UPat.cvar("c"), lambda a,c: a.ins(X86Ops.ANDi, src=(a, i)) if (i:=to_imm(c)) is not None else None), + (UPat.var("a", dtypes.ints+(dtypes.bool,)) | UPat.cvar("c"), lambda a,c: a.ins(X86Ops.ORi, src=(a, i)) if (i:=to_imm(c)) is not None else None), + (UPat.var("a", dtypes.ints+(dtypes.bool,)) ^ UPat.cvar("c"), lambda a,c: a.ins(X86Ops.XORi, src=(a, i)) if (i:=to_imm(c)) is not None else None), + (UPat(Ops.SUB, dtypes.ints, (UPat.var("a"), UPat.cvar("c"))), lambda a,c: a.ins(X86Ops.SUBi, src=(a, i)) if (i:=to_imm(c)) is not None else None), + # scalar int binary with register + (UPat.var("a", dtypes.ints) << UPat.var("b"), lambda a,b: a.ins(X86Ops.SHL, src=(a, b))), + (UPat.var("a", dtypes.uints) >> UPat.var("b"), lambda a,b: a.ins(X86Ops.SHR, src=(a, b))), + (UPat.var("a", dtypes.sints) >> UPat.var("b"), lambda a,b: a.ins(X86Ops.SAR, src=(a, b))), + (UPat.var("a", dtypes.ints) + UPat.var("b"), lambda a,b: a.ins(X86Ops.ADD, src=(a, b))), + (UPat.var("a", dtypes.ints) * UPat.var("b"), lambda a,b: a.ins(X86Ops.IMUL, src=(a, b))), + (UPat.var("a", dtypes.ints+(dtypes.bool,)) & UPat.var("b"), lambda a,b: a.ins(X86Ops.AND, src=(a, b))), + (UPat.var("a", dtypes.ints+(dtypes.bool,)) | UPat.var("b"), lambda a,b: a.ins(X86Ops.OR, src=(a, b))), + (UPat.var("a", dtypes.ints+(dtypes.bool,)) ^ UPat.var("b"), lambda a,b: a.ins(X86Ops.XOR, src=(a, b))), + (UPat(Ops.SUB, dtypes.ints, (UPat.var("a"), UPat.var("b"))), lambda a,b: a.ins(X86Ops.SUB, src=(a, b))), + # float binary + ((UPat(dtype=dtypes.float32) + UPat()).named("x"), lambda x: x.ins(X86Ops.VADDSS if x.dtype.count == 1 else X86Ops.VADDPS)), + ((UPat(dtype=dtypes.float64) + UPat()).named("x"), lambda x: x.ins(X86Ops.VADDSD if x.dtype.count == 1 else X86Ops.VADDPD)), + ((UPat(dtype=dtypes.float32) * UPat()).named("x"), lambda x: x.ins(X86Ops.VMULSS if x.dtype.count == 1 else X86Ops.VMULPS)), + ((UPat(dtype=dtypes.float64) * UPat()).named("x"), lambda x: x.ins(X86Ops.VMULSD if x.dtype.count == 1 else X86Ops.VMULPD)), + (UPat(Ops.SUB, dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VSUBSS if x.dtype.count == 1 else X86Ops.VSUBPS)), + (UPat(Ops.SUB, dtypes.float64, name="x"), lambda x: x.ins(X86Ops.VSUBSD if x.dtype.count == 1 else X86Ops.VSUBPD)), + (UPat(Ops.FDIV, dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VDIVSS if x.dtype.count == 1 else X86Ops.VDIVPS)), + (UPat(Ops.FDIV, dtypes.float64, name="x"), lambda x: x.ins(X86Ops.VDIVSD if x.dtype.count == 1 else X86Ops.VDIVPD)), + # casts + (UPat(dtype=dtypes.int32).cast(dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VCVTDQ2PS) if x.dtype.count > 1 else None), + (UPat(dtype=dtypes.int32).cast(dtypes.float64, name="x"), lambda x: x.ins(X86Ops.VCVTDQ2PD) if x.dtype.count > 1 else None), + (UPat(dtype=dtypes.float32).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VCVTTPS2DQ) if x.dtype.count > 1 else None), + (UPat(dtype=dtypes.float64).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VCVTTPD2DQ) if x.dtype.count > 1 else None), + (UPat(dtype=dtypes.float32).cast(dtypes.float64, name="x"), lambda x: x.ins(X86Ops.VCVTPS2PD) if x.dtype.count > 1 else None), + (UPat(dtype=dtypes.float64).cast(dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VCVTPD2PS) if x.dtype.count > 1 else None), + (UPat(dtype=dtypes.float32).cast(dtypes.float16, name="x"), lambda x: x.ins(X86Ops.VCVTPS2PH, src=x.src + (imm(dtypes.uint8, 4),))), + (UPat(dtype=dtypes.float16).cast(dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VCVTPH2PS)), + (UPat(dtype=dtypes.float32).cast(dtypes.int32s+dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VCVTTSS2SI)), + (UPat(dtype=dtypes.float64).cast(dtypes.int32s+dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VCVTTSD2SI)), + (UPat.var("y", dtypes.float32).cast(dtypes.float64, name="x"), lambda y,x: x.ins(X86Ops.VCVTSS2SD, src=(y, y))), + (UPat.var("y", dtypes.float64).cast(dtypes.float32, name="x"), lambda y,x: x.ins(X86Ops.VCVTSD2SS, src=(y, y))), + (UPat.var("y", (dtypes.int32, dtypes.int64)).cast(dtypes.float32, name="x"), lambda y,x: x.ins(X86Ops.VCVTSI2SS, src=(def_reg(x.dtype), y))), + (UPat.var("y", (dtypes.int32, dtypes.int64)).cast(dtypes.float64, name="x"), lambda y,x: x.ins(X86Ops.VCVTSI2SD, src=(def_reg(x.dtype), y))), + (UPat(dtype=dtypes.uints+(dtypes.bool,)).cast(dtypes.ints, name="x"), lambda x: x.ins(X86Ops.MOVZX) if x.dtype.count == 1 else None), + (UPat(dtype=dtypes.int32).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.MOVSXD) if x.dtype.count == 1 else None), + (UPat(dtype=dtypes.sints).cast(dtypes.ints, name="x"), lambda x: x.ins(X86Ops.MOVSX) if x.dtype.count == 1 else None), + (UPat(dtype=(dtypes.uint8, dtypes.bool)).cast(dtypes.int16s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXBW)), + (UPat(dtype=(dtypes.uint8, dtypes.bool)).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXBD)), + (UPat(dtype=(dtypes.uint8, dtypes.bool)).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXBQ)), + (UPat(dtype=dtypes.uint16).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXWD)), + (UPat(dtype=dtypes.uint16).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXWQ)), + (UPat(dtype=dtypes.uint32).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXDQ)), + (UPat(dtype=dtypes.int8).cast(dtypes.int16s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXBW)), + (UPat(dtype=dtypes.int8).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXBD)), + (UPat(dtype=dtypes.int8).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXBQ)), + (UPat(dtype=dtypes.int16).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXWD)), + (UPat(dtype=dtypes.int16).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXWQ)), + (UPat(dtype=dtypes.int32).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXDQ)), + # bitcasts + (UPat.var("y", dtypes.float16).bitcast(dtypes.int16s).named("x"), lambda y,x: x.ins(X86Ops.VPEXTRW, src=(y, imm(dtypes.uint8, 0)))), + (UPat(dtype=dtypes.int16s).bitcast(dtypes.float16).named("x"), vpins), + (UPat(dtype=dtypes.int32s).bitcast(dtypes.float32).named("x"), lambda x: x.ins(X86Ops.VMOVD)), + (UPat(dtype=dtypes.int64s).bitcast(dtypes.float64).named("x"), lambda x: x.ins(X86Ops.VMOVQ)), + (UPat(dtype=dtypes.float32).bitcast(dtypes.int32s).named("x"), lambda x: x.ins(X86Ops.VMOVDm)), + (UPat(dtype=dtypes.float64).bitcast(dtypes.int64s).named("x"), lambda x: x.ins(X86Ops.VMOVQm)), + # index + (UPat(Ops.INDEX, name="x"), lambda x: x.ins(X86Ops.LEA, src=fold_address(x))), + # TODO: fuse stores, very few cases -- store cmp becomes setcc, store gep int becomes vpextr, store bitcast to int becomes vmovd/q + # copy, load, store + # NOTE: copy here violates the spec, it only happens post register allocation when a reg to reg move needs to be inserted + (UPat(Ops.COPY, dt_128bit, name="x"), lambda x: x.ins(X86Ops.VMOVUPS)), + (UPat(Ops.COPY, dt_64bit, name="x"), lambda x: x.ins(X86Ops.VMOVSD)), + (UPat(Ops.COPY, dt_32bit+dt_16bit, name="x"), lambda x: x.ins(X86Ops.VMOVSS)), + (UPat(Ops.COPY, dtypes.ints+(dtypes.bool,), name="x"), lambda x: x.ins(X86Ops.MOV)), + (UPat(Ops.LOAD, dt_128bit, src=(UPat(name="a"),), name="x"), lambda x,a: x.ins(X86Ops.VMOVUPS, src=fold_address(a))), + (UPat(Ops.LOAD, dt_64bit, src=(UPat(name="a"),), name="x"), lambda x,a: x.ins(X86Ops.VMOVSD, src=fold_address(a))), + (UPat(Ops.LOAD, dt_32bit, src=(UPat(name="a"),), name="x"), lambda x,a: x.ins(X86Ops.VMOVSS, src=fold_address(a))), + (UPat(Ops.LOAD, dt_16bit, src=(UPat(name="a"),), name="x"), lambda x,a: + x.ins(X86Ops.VPINSRW, src=(def_reg(x.dtype, x.tag),) + fold_address(a) + (imm(dtypes.uint8, 0),))), + (UPat(Ops.LOAD, dtypes.ints+(dtypes.bool,), src=(UPat(name="a"),), name="x"), lambda x,a: x.ins(X86Ops.MOV, src=fold_address(a))), + (UPat.var("a").store(UPat.var("b", dt_128bit), name="x"), lambda a,b,x: x.ins(X86Ops.VMOVUPSm, src=fold_address(a) + (b,))), + (UPat.var("a").store(UPat.var("b", dt_64bit), name="x"), lambda a,b,x: x.ins(X86Ops.VMOVSDm, src=fold_address(a) + (b,))), + (UPat.var("a").store(UPat.var("b", dt_32bit), name="x"), lambda a,b,x: x.ins(X86Ops.VMOVSSm, src=fold_address(a) + (b,))), + (UPat.var("a").store(UPat.var("b", dt_16bit), name="x"), lambda a,b,x: x.ins(X86Ops.VPEXTRW, src=fold_address(a) + (b, imm(dtypes.uint8, 0)))), + (UPat.var("a").store(UPat.var("b", dtypes.ints+(dtypes.bool,)), name="x"), lambda a,b,x: + x.ins(X86Ops.MOVm, src=fold_address(a) + (b,)) if (i:=to_imm(b)) is None else x.ins(X86Ops.MOVi, src=fold_address(a) + (i,))), + # **** X86Op -> X86Op **** + # fold loads into X86Ops that allow it, if beneficial + (UPat(Ops.INS, src=(UPat(Ops.LOAD, src=(UPat(name="a"),), name="y"),), allow_any_len=True, name="x"), lambda ctx,y,a,x: + x.replace(src=fold_address(a) + x.src[1:]) if x.arg in X86GroupOp.ReadMem1st and is_foldable(ctx, x, y) else None), + (UPat(Ops.INS, src=(UPat(), UPat(Ops.LOAD, src=(UPat(name="a"),), name="y")), allow_any_len=True, name="x"), lambda ctx,y,a,x: + x.replace(src=x.src[:1] + fold_address(a) + x.src[2:]) if x.arg in X86GroupOp.ReadMem2nd and is_foldable(ctx, x, y) else None), + (UPat(Ops.INS, src=(UPat(), UPat(), UPat(Ops.LOAD, src=(UPat(name="a"),), name="y")), allow_any_len=True, name="x"), lambda ctx,y,a,x: + x.replace(src=x.src[:2] + fold_address(a) + x.src[3:]) if x.arg in X86GroupOp.ReadMem3rd and is_foldable(ctx, x, y) else None), + # allocate virtual registers + (UPat((Ops.INS, Ops.DEFINE_REG, Ops.DEFINE_LOCAL), name="x"), alloc_vregs), +]) + +# ***** pre register allocation ***** +# this handles flag clobbers. Unfortunately x86 doesn't have a good way to store/restore the flag register (then regalloc would handle it) +# so we rematerialize. This is different from rematerialization you might want to do in regalloc because it is not optional, +# regalloc shouldn't rematerialize if a src of the instruction is dead, but here you need to as there's no fallback load from stack +def flag_rematerialize(ctx:PreRegAllocContext, x:UOp): + flag_def = x if x.arg in X86GroupOp.WriteFlags or x.op in (Ops.RANGE, Ops.END) else x.src[-1] if x.arg in X86GroupOp.ReadFlags else None + if flag_def is None: return None + if ctx.lock is not None and ctx.lock is not flag_def: ctx.clobbered.add(ctx.lock) + ctx.lock = flag_def + if flag_def not in ctx.clobbered: return None + ctx.clobbered.remove(flag_def) + return (x, [flag_def, x]) + +pre_regalloc_matcher = PatternMatcher([ + (UPat((Ops.INS, Ops.RANGE, Ops.END), name="x"), flag_rematerialize), +]) + +# ***** post register allocation ***** +# TODO: control flow should be overhauled so that this isn't necessary +def lower_range(ctx, x:UOp) -> tuple[UOp, list[UOp]]: + loop_label = "_".join(str(i) for i in x.arg[:-1]) + acc = x.ins(X86Ops.MOVi, src=(imm(x.dtype, 0),) + x.src[1:]) + label = UOp(Ops.INS, arg=X86Ops.LABEL, tag=f".LOOP_{loop_label}") + cmp = UOp(Ops.INS, arg=X86Ops.CMPi if x.src[0].op is Ops.CONST else X86Ops.CMP, src=(acc, x.src[0])) + jump_out = UOp(Ops.INS, arg=X86Ops.JGE, src=(cmp,), tag=f".LOOP_OUT_{loop_label}") + ctx.loop_label[acc] = loop_label + return (acc, [acc, label, cmp, jump_out]) + +# final rewrite to match the isa spec +post_regalloc_matcher = PatternMatcher([ + # rewrite FRAME_INDEX to IMM now that the stack size is known + (UPat(Ops.INS, arg=X86Ops.FRAME_INDEX, name="x"), lambda ctx,x: (nx:=x.const_like(ctx.stack_size + x.tag), [nx])), + # rewrite RANGE to ACC = 0 -> LABEL -> JUMP if ACC >= loop bound + (UPat(Ops.RANGE, name="x"), lambda ctx,x: lower_range(ctx, x)), + # rewrite END to ACC + 1 -> JUMP -> LABEL, also add the out of loop JUMP to the src so this becomes the jump target + (UPat(Ops.END, name="x"), lambda ctx,x: (jmp:=UOp(Ops.INS, arg=X86Ops.JMP, tag=f".LOOP_{ctx.loop_label[x.src[1]]}"), + [x.src[1].ins(X86Ops.ADDi, src=(imm(x.src[1].dtype, 1),)), jmp, UOp(Ops.INS, arg=X86Ops.LABEL, tag=f".LOOP_OUT_{ctx.loop_label[x.src[1]]}")])), + # rewrite two address instructions to two address form, if reused src wasn't coalesced insert a move + (UPat(Ops.INS, name="x"), lambda ctx,x: (nx:=x.replace(src=x.src[1:]), + [ctx.ren.copy(x.src[0], x.reg), nx] if x.reg != x.src[0].reg else [nx]) if x.arg in X86GroupOp.TwoAddress else None), +]) + +# ***** X86 instruction encoding ***** + +def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) -> bytes|None: + def _encode(reg_uop:UOp|None, rm_uop:UOp, idx_uop:UOp|None=None, disp_uop:UOp|None=None, vvvv_uop:UOp|None=None, imm_uop:UOp|None=None) -> bytes: + nonlocal reg, opc + # get the encoding values of the different fields + reg = cast(int, cast(Register, reg_uop.reg).index if reg_uop is not None else reg) + rm = cast(Register, rm_uop.reg).index + idx = cast(Register, idx_uop.reg).index if idx_uop is not None and idx_uop.reg is not None else 4 + rm_sz = 8 if isinstance(rm_uop.dtype, PtrDType) and disp_uop is None else rm_uop.dtype.itemsize + reg_sz = (reg_uop.dtype.itemsize if not isinstance(reg_uop.dtype, PtrDType) else 8) if reg_uop is not None else 0 + sz = reg_sz or rm_sz + + # encode instruction + inst = bytes([]) + assert 0 <= reg <= 15 and 0 <= idx <= 15 and 0 <= rm <= 15 + # r extends reg field, x extends index field, b extends rm or base field + r, _x, b = reg >> 3, idx >> 3, rm >> 3 + if sel: # VEX bytes + vvvv = cast(Register, vvvv_uop.reg).index if vvvv_uop is not None else 0 + l = (max(reg_sz, rm_sz) > 16) & 0b1 + if sel == 1 and _x == b == we == 0: inst += bytes([0xC5, (~r & 0b1) << 7 | (~vvvv & 0b1111) << 3 | l << 2 | pp]) + else: inst += bytes([0xC4, (~r & 0b1) << 7 | (~_x & 0b1) << 6 | (~b & 0b1) << 5 | sel, we << 7 | (~vvvv & 0b1111) << 3 | l << 2 | pp]) + else: # optional PREFIX and REX bytes + # PREFIX byte signaling 16 bit variant of instruction + if sz == 2: inst += bytes([0x66]) + # bit signaling 64 bit variant of instruction + w = sz == 8 + # REX byte is required when 64 bit or an extended reg is used (index 8 - 15) or lower 8 bits of (rsp, rbp, rsi, rdi) are accessed + if w | r | _x | b | (reg_sz == 1 & reg >> 2) | (rm_sz == 1 & rm >> 2): inst += bytes([0b0100 << 4 | w << 3 | r << 2 | _x << 1 | b]) + # legacy 8bit opcode is 1 less than 16-64bit variants + if (rm_sz == 1 or reg_sz == 1) and x.arg not in X86GroupOp.ReadFlags | {X86Ops.LEA}: opc -= 1 + # OPCODE byte + inst += opc.to_bytes((opc.bit_length() + 7) // 8, 'big') + # MODRM byte + # now we only care about the lower 3 bits + idx, rm, reg = idx & 0b111, rm & 0b111, reg & 0b111 + # 0b00 -- signals memory access with no displacement + # 0b01 -- signals memory access with 8bit displacement + # 0b10 -- signals memory access with 32bit displacement + # 0b11 -- signals no memory access + if disp_uop is not None: + assert disp_uop.op is Ops.CONST, "displacement must be a constant" + assert disp_uop.dtype in (dtypes.int8, dtypes.int32), "displacement can only be 1 or 4 byte signed int" + # rbp/r13 always require a displacement + if disp_uop.arg != 0 or rm == 0b101: mod = 0b01 if disp_uop.dtype.itemsize == 1 else 0b10 + else: mod = 0b00 + else: mod = 0b11 + # x 0b0 and idx 0b100 means rsp which means no index exists + # rm 0b100 (rsp/r12) signals a sib byte is required, rm then is encoded in the base field of SIB + _rm = rm if idx == 0b100 and _x == 0b0 else 0b100 + inst += bytes([mod << 6 | reg << 3 | _rm]) + # SIB byte + if _rm == 0b100 and mod != 0b11: + scale = {1: 0b00, 2: 0b01, 4: 0b10, 8: 0b11}[1 if idx == 0b100 and _x == 0b0 else rm_sz] + inst += bytes([scale << 6 | idx << 3 | rm]) + # DISP byte + if mod == 0b01 or mod == 0b10: + assert disp_uop is not None + inst += struct.pack(unwrap(disp_uop.dtype.fmt), disp_uop.arg) + # IMM byte + if imm_uop is not None: + if imm_uop.op is Ops.CONST: inst += struct.pack(unwrap(imm_uop.dtype.fmt), imm_uop.arg) + elif isinstance(imm_uop.reg, Register): inst += bytes([(imm_uop.reg.index & 0b1111) << 4 | 0b0000]) + return inst + + # get the encoding structure of the uop + # when a uop writes to memory it takes the form of a store, dtype is void, no definition + address:tuple[UOp|None, ...] + if x.arg in X86GroupOp.WriteMem: + if len(x.src) > 3: address, rest = x.src[:3], x.src[3:] + else: address, rest = (x, None, None), x.src + return _encode(rest[0], *address, *(None, *rest[1:])) if reg is None else _encode(None, *address, *(None, *rest[:1])) + + if x.arg in X86GroupOp.Rm1st: + if len(x.src) > 2: address, rest = x.src[:3], x.src[3:] + else: address, rest = (x.src[0], None, None), x.src[1:] + imm_uop = rest[:1] if rest and rest[0].op is Ops.CONST else (None,) + return _encode(x, *address, *(None, *imm_uop)) if reg is None else _encode(None, *address, *(x if sel else None, *imm_uop)) + + if x.arg in X86GroupOp.Rm2nd: + if len(x.src) > 3: address, rest = x.src[1:4], x.src[:1] + x.src[4:] + else: address, rest = (x.src[1], None, None), x.src[:1] + x.src[2:] + # cmp/vucomiss reg, rm don't define a new register + return _encode(x, *address, *rest) if x.dtype is not dtypes.void else _encode(rest[0], *address) + + return None + +# https://www.felixcloutier.com/x86/ +# legacy version -> VEX version +# prefix field: None -> 0 | 66 -> 1 | F3 -> 2 | F2 -> 3 +# opcode map select: 0F -> 1 | 0F38 -> 2 | 0F3A -> 3 +encodings = { + # moves + X86Ops.MOVABS: lambda x: + bytes([0b0100 << 4 | 0b1 << 3 | 0b00 << 2 | x.reg.index >> 3, 0xB8 + (x.reg.index & 0b111)]) + struct.pack(x.dtype.fmt, x.src[0].arg), + X86Ops.MOV: lambda x: encode(x, 0x8B), X86Ops.MOVi: lambda x: encode(x, 0xC7, reg=0), + X86Ops.MOVm: lambda x: encode(x, 0x89), X86Ops.LEA: lambda x: encode(x, 0x8D), + X86Ops.VMOVSS: lambda x: encode(x, 0x10, pp=2, sel=1), X86Ops.VMOVSSm: lambda x: encode(x, 0x11, pp=2, sel=1), + X86Ops.VMOVSD: lambda x: encode(x, 0x10, pp=3, sel=1), X86Ops.VMOVSDm: lambda x: encode(x, 0x11, pp=3, sel=1), + X86Ops.VMOVUPS: lambda x: encode(x, 0x10, pp=0, sel=1), X86Ops.VMOVUPSm: lambda x: encode(x, 0x11, pp=0, sel=1), + X86Ops.VMOVD: lambda x: encode(x, 0x6E, pp=1, sel=1), X86Ops.VMOVQ: lambda x: encode(x, 0x6E, pp=1, sel=1, we=1), + X86Ops.VMOVDm: lambda x: encode(x, 0x7E, pp=1, sel=1), X86Ops.VMOVQm: lambda x: encode(x, 0x7E, pp=1, sel=1, we=1), + # casts + X86Ops.MOVZX: lambda x: encode(x, 0x0FB7), + X86Ops.MOVSX: lambda x: encode(x, 0x0FBF), X86Ops.MOVSXD: lambda x: encode(x, 0x63), + X86Ops.VPMOVZXBW: lambda x: encode(x, 0x30, pp=1, sel=2), X86Ops.VPMOVZXBD: lambda x: encode(x, 0x31, pp=1, sel=2), + X86Ops.VPMOVZXBQ: lambda x: encode(x, 0x32, pp=1, sel=2), X86Ops.VPMOVZXWD: lambda x: encode(x, 0x33, pp=1, sel=2), + X86Ops.VPMOVZXWQ: lambda x: encode(x, 0x34, pp=1, sel=2), X86Ops.VPMOVZXDQ: lambda x: encode(x, 0x35, pp=1, sel=2), + X86Ops.VPMOVSXBW: lambda x: encode(x, 0x20, pp=1, sel=2), X86Ops.VPMOVSXBD: lambda x: encode(x, 0x21, pp=1, sel=2), + X86Ops.VPMOVSXBQ: lambda x: encode(x, 0x22, pp=1, sel=2), X86Ops.VPMOVSXWD: lambda x: encode(x, 0x23, pp=1, sel=2), + X86Ops.VPMOVSXWQ: lambda x: encode(x, 0x24, pp=1, sel=2), X86Ops.VPMOVSXDQ: lambda x: encode(x, 0x25, pp=1, sel=2), + X86Ops.VCVTSS2SD: lambda x: encode(x, 0x5A, pp=2, sel=1), X86Ops.VCVTSD2SS: lambda x: encode(x, 0x5A, pp=3, sel=1), + X86Ops.VCVTPH2PS: lambda x: encode(x, 0x13, pp=1, sel=2), X86Ops.VCVTPS2PH: lambda x: encode(x, 0x1D, pp=1, sel=3), + X86Ops.VCVTDQ2PS: lambda x: encode(x, 0x5B, pp=0, sel=1), X86Ops.VCVTDQ2PD: lambda x: encode(x, 0xE6, pp=2, sel=1), + X86Ops.VCVTPS2PD: lambda x: encode(x, 0x5A, pp=0, sel=1), X86Ops.VCVTPD2PS: lambda x: encode(x, 0x5A, pp=1, sel=1), + X86Ops.VCVTTPS2DQ: lambda x: encode(x, 0x5B, pp=2, sel=1), X86Ops.VCVTTPD2DQ: lambda x: encode(x, 0xE6, pp=1, sel=1), + X86Ops.VCVTSI2SS: lambda x: encode(x, 0x2A, pp=2, sel=1, we=x.src[1].dtype.itemsize == 8), + X86Ops.VCVTSI2SD: lambda x: encode(x, 0x2A, pp=3, sel=1, we=x.src[1].dtype.itemsize == 8), + X86Ops.VCVTTSS2SI: lambda x: encode(x, 0x2C, pp=2, sel=1, we=x.dtype.itemsize == 8), + X86Ops.VCVTTSD2SI: lambda x: encode(x, 0x2C, pp=3, sel=1, we=x.dtype.itemsize == 8), + # int division + X86Ops.IDIV: lambda x: encode(x, 0xF7, reg=7), X86Ops.DIV: lambda x: encode(x, 0xF7, reg=6), + # scalar int binary + X86Ops.SHLi: lambda x: encode(x, 0xC1, reg=4), + X86Ops.SHRi: lambda x: encode(x, 0xC1, reg=5), X86Ops.SARi: lambda x: encode(x, 0xC1, reg=7), + X86Ops.ADD: lambda x: encode(x, 0x03), X86Ops.ADDi: lambda x: encode(x, 0x81, reg=0), + X86Ops.SUB: lambda x: encode(x, 0x2B), X86Ops.SUBi: lambda x: encode(x, 0x81, reg=5), + X86Ops.AND: lambda x: encode(x, 0x23), X86Ops.ANDi: lambda x: encode(x, 0x81, reg=4), + X86Ops.XOR: lambda x: encode(x, 0x33), X86Ops.XORi: lambda x: encode(x, 0x81, reg=6), + X86Ops.OR: lambda x: encode(x, 0x0B), X86Ops.ORi: lambda x: encode(x, 0x81, reg=1), + X86Ops.CMP: lambda x: encode(x, 0x3B), X86Ops.CMPi: lambda x: encode(x, 0x81, reg=7), + X86Ops.IMUL: lambda x: encode(x, 0x0FAF), X86Ops.IMULi: lambda x: encode(x, 0x69), + X86Ops.SETB: lambda x: encode(x, 0x0F92, reg=0), X86Ops.SETL: lambda x: encode(x, 0x0F9C, reg=0), + X86Ops.SETE: lambda x: encode(x, 0x0F94, reg=0), X86Ops.SETNE: lambda x: encode(x, 0x0F95, reg=0), + # packed bitwise NOTE: only bitwise and packed + X86Ops.VPAND: lambda x: encode(x, 0xDB, pp=1, sel=1), X86Ops.VPXOR: lambda x: encode(x, 0xEF, pp=1, sel=1), + X86Ops.VPOR: lambda x: encode(x, 0xEB, pp=1, sel=1), + # unary + X86Ops.VSQRTSS: lambda x: encode(x, 0x51, pp=2, sel=1), X86Ops.VSQRTPS: lambda x: encode(x, 0x51, pp=0, sel=1), + X86Ops.VSQRTSD: lambda x: encode(x, 0x51, pp=3, sel=1), X86Ops.VSQRTPD: lambda x: encode(x, 0x51, pp=1, sel=1), + X86Ops.VROUNDSS: lambda x: encode(x, 0x0A, pp=1, sel=3), X86Ops.VROUNDPS: lambda x: encode(x, 0x08, pp=1, sel=3), + X86Ops.VROUNDSD: lambda x: encode(x, 0x0B, pp=1, sel=3), X86Ops.VROUNDPD: lambda x: encode(x, 0x09, pp=1, sel=3), + # packed int binary + X86Ops.VPSLLVD: lambda x: encode(x, 0x47, pp=1, sel=2), X86Ops.VPSLLVQ: lambda x: encode(x, 0x47, pp=1, sel=2, we=1), + X86Ops.VPSRLVD: lambda x: encode(x, 0x45, pp=1, sel=2), X86Ops.VPSRLVQ: lambda x: encode(x, 0x45, pp=1, sel=2, we=1), + X86Ops.VPCMPGTB: lambda x: encode(x, 0x64, pp=1, sel=1), X86Ops.VPCMPGTW: lambda x: encode(x, 0x65, pp=1, sel=1), + X86Ops.VPCMPGTD: lambda x: encode(x, 0x66, pp=1, sel=1), X86Ops.VPCMPGTQ: lambda x: encode(x, 0x37, pp=1, sel=2), + X86Ops.VPCMPEQB: lambda x: encode(x, 0x74, pp=1, sel=1), X86Ops.VPCMPEQW: lambda x: encode(x, 0x75, pp=1, sel=1), + X86Ops.VPCMPEQD: lambda x: encode(x, 0x76, pp=1, sel=1), X86Ops.VPCMPEQQ: lambda x: encode(x, 0x29, pp=1, sel=2), + X86Ops.VPMULLW: lambda x: encode(x, 0xD5, pp=1, sel=1), X86Ops.VPMULLD: lambda x: encode(x, 0x40, pp=1, sel=2), + X86Ops.VPADDB: lambda x: encode(x, 0xFC, pp=1, sel=1), X86Ops.VPADDW: lambda x: encode(x, 0xFD, pp=1, sel=1), + X86Ops.VPADDD: lambda x: encode(x, 0xFE, pp=1, sel=1), X86Ops.VPADDQ: lambda x: encode(x, 0xD4, pp=1, sel=1), + X86Ops.VPSUBB: lambda x: encode(x, 0xF8, pp=1, sel=1), X86Ops.VPSUBW: lambda x: encode(x, 0xF9, pp=1, sel=1), + X86Ops.VPSUBD: lambda x: encode(x, 0xFA, pp=1, sel=1), X86Ops.VPSUBQ: lambda x: encode(x, 0xFB, pp=1, sel=1), + X86Ops.VPSRAVD: lambda x: encode(x, 0x46, pp=1, sel=2), + # float cmp + X86Ops.VUCOMISS: lambda x: encode(x, 0x2E, pp=0, sel=1), X86Ops.VUCOMISD: lambda x: encode(x, 0x2E, pp=1, sel=1), + # scalar / packed float binary + X86Ops.VADDSS: lambda x: encode(x, 0x58, pp=2, sel=1), X86Ops.VADDPS: lambda x: encode(x, 0x58, pp=0, sel=1), + X86Ops.VADDSD: lambda x: encode(x, 0x58, pp=3, sel=1), X86Ops.VADDPD: lambda x: encode(x, 0x58, pp=1, sel=1), + X86Ops.VSUBSS: lambda x: encode(x, 0x5C, pp=2, sel=1), X86Ops.VSUBPS: lambda x: encode(x, 0x5C, pp=0, sel=1), + X86Ops.VSUBSD: lambda x: encode(x, 0x5C, pp=3, sel=1), X86Ops.VSUBPD: lambda x: encode(x, 0x5C, pp=1, sel=1), + X86Ops.VMULSS: lambda x: encode(x, 0x59, pp=2, sel=1), X86Ops.VMULPS: lambda x: encode(x, 0x59, pp=0, sel=1), + X86Ops.VMULSD: lambda x: encode(x, 0x59, pp=3, sel=1), X86Ops.VMULPD: lambda x: encode(x, 0x59, pp=1, sel=1), + X86Ops.VDIVSS: lambda x: encode(x, 0x5E, pp=2, sel=1), X86Ops.VDIVPS: lambda x: encode(x, 0x5E, pp=0, sel=1), + X86Ops.VDIVSD: lambda x: encode(x, 0x5E, pp=3, sel=1), X86Ops.VDIVPD: lambda x: encode(x, 0x5E, pp=1, sel=1), + X86Ops.VCMPSS: lambda x: encode(x, 0xC2, pp=2, sel=1), X86Ops.VCMPPS: lambda x: encode(x, 0xC2, pp=0, sel=1), + X86Ops.VCMPSD: lambda x: encode(x, 0xC2, pp=3, sel=1), X86Ops.VCMPPD: lambda x: encode(x, 0xC2, pp=1, sel=1), + X86Ops.VMAXSS: lambda x: encode(x, 0x5F, pp=2, sel=1), X86Ops.VMAXPS: lambda x: encode(x, 0x5F, pp=0, sel=1), + X86Ops.VMAXSD: lambda x: encode(x, 0x5F, pp=3, sel=1), X86Ops.VMAXPD: lambda x: encode(x, 0x5F, pp=1, sel=1), + X86Ops.VMINSS: lambda x: encode(x, 0x5D, pp=2, sel=1), X86Ops.VMINPS: lambda x: encode(x, 0x5D, pp=0, sel=1), + X86Ops.VMINSD: lambda x: encode(x, 0x5D, pp=3, sel=1), X86Ops.VMINPD: lambda x: encode(x, 0x5D, pp=1, sel=1), + # ternary + X86Ops.CMOVB: lambda x: encode(x, 0x0F42), X86Ops.CMOVL: lambda x: encode(x, 0x0F4C), + X86Ops.CMOVE: lambda x: encode(x, 0x0F44), X86Ops.CMOVNE: lambda x: encode(x, 0x0F45), + X86Ops.VFMADD213SS: lambda x: encode(x, 0xA9, pp=1, sel=2), X86Ops.VFMADD213SD: lambda x: encode(x, 0xA9, pp=1, sel=2, we=1), + X86Ops.VFMADD213PS: lambda x: encode(x, 0xA8, pp=1, sel=2), X86Ops.VFMADD213PD: lambda x: encode(x, 0xA8, pp=1, sel=2, we=1), + X86Ops.VBLENDVPS: lambda x: encode(x, 0x4A, pp=1, sel=3), X86Ops.VBLENDVPD: lambda x: encode(x, 0x4B, pp=1, sel=3), + X86Ops.VPBLENDVB: lambda x: encode(x, 0x4C, pp=1, sel=3), + # shuffles + X86Ops.VPBROADCASTB: lambda x: encode(x, 0x78, pp=1, sel=2), X86Ops.VPBROADCASTW: lambda x: encode(x, 0x79, pp=1, sel=2), + X86Ops.VPBROADCASTD: lambda x: encode(x, 0x58, pp=1, sel=2), X86Ops.VPBROADCASTQ: lambda x: encode(x, 0x59, pp=1, sel=2), + X86Ops.VBROADCASTSS: lambda x: encode(x, 0x18, pp=1, sel=2), X86Ops.VPSRLDQ: lambda x: encode(x, 0x73, reg=3, pp=1, sel=1), + X86Ops.VPINSRB: lambda x: encode(x, 0x20, pp=1, sel=3), X86Ops.VPINSRW: lambda x: encode(x, 0xC4, pp=1, sel=1), + X86Ops.VPINSRD: lambda x: encode(x, 0x22, pp=1, sel=3), X86Ops.VPINSRQ: lambda x: encode(x, 0x22, pp=1, sel=3, we=1), + X86Ops.VSHUFPS: lambda x: encode(x, 0xC6, pp=0, sel=1), X86Ops.VSHUFPD: lambda x: encode(x, 0xC6, pp=1, sel=1), + X86Ops.VINSERTPS: lambda x: encode(x, 0x21, pp=1, sel=3), + # extract + X86Ops.VPEXTRB: lambda x: encode(x, 0x14, pp=1, sel=3), X86Ops.VPEXTRW: lambda x: encode(x, 0x15, pp=1, sel=3), + X86Ops.VPEXTRD: lambda x: encode(x, 0x16, pp=1, sel=3), X86Ops.VPEXTRQ: lambda x: encode(x, 0x16, pp=1, sel=3, we=1), + # jumps are encoded with a placeholder which gets patched later once the real offset is known + X86Ops.JE: lambda x: bytes([0x0F, 0x84]) + int(0).to_bytes(4), + X86Ops.JNE: lambda x: bytes([0x0F, 0x85]) + int(0).to_bytes(4), + X86Ops.JL: lambda x: bytes([0x0F, 0x8C]) + int(0).to_bytes(4), + X86Ops.JB: lambda x: bytes([0x0F, 0x82]) + int(0).to_bytes(4), + X86Ops.JGE: lambda x: bytes([0x0F, 0x8D]) + int(0).to_bytes(4), + X86Ops.JMP: lambda x: bytes([0xE9]) + int(0).to_bytes(4), + X86Ops.RET: lambda x: bytes([0xC3]), +} + +class X86Renderer(ISARenderer): + device = "CPU" + has_local = False + has_threads = bool(getenv("THREADS", 1)) + global_max = (CPU_COUNT.value, 0, 0) + extra_matcher = extra_matcher + pre_isel_matcher = pre_isel_matcher + isel_matcher = isel_matcher + pre_regalloc_matcher = pre_regalloc_matcher + post_regalloc_matcher = post_regalloc_matcher + code_for_op = {x: lambda: None for x in (Ops.SQRT, Ops.AND, Ops.OR, Ops.SHL, Ops.SHR, Ops.NEG, Ops.SUB, Ops.FDIV, Ops.CMPLT, Ops.CMPEQ)} + def __init__(self, target:Target): + super().__init__(target) + from tinygrad.runtime.support.compiler_cpu import X86Compiler + self.compiler = X86Compiler() + def is_two_address(self, x:UOp) -> bool: return x.arg in X86GroupOp.TwoAddress + def stack_pointer(self) -> UOp: return def_reg(dtypes.uint64, RSP) + # nasty hacks to deal with pointers TODO: rm pointers + def copy(self, x:UOp, reg:Register): + dt = dtypes.uint64 if isinstance(x.dtype, PtrDType) else x.dtype + ret = isel_matcher.rewrite(UOp(Ops.COPY, dt, (x,), tag=reg)) + assert ret is not None + return ret.replace(dtype=x.dtype) + + def spill(self, disp:UOp, x:UOp) -> UOp: + nx = x.replace(dtype=dtypes.uint64 if isinstance(x.dtype, PtrDType) else x.dtype) + ret = isel_matcher.rewrite(self.stack_pointer().index(disp).store(nx)) + assert ret is not None + return ret.replace(src=(s if s is not nx else x for s in ret.src)) + + def fill(self, disp:UOp, x:UOp, reg:Register) -> UOp: + ndt = dtypes.uint64 if isinstance(x.dtype, PtrDType) else x.dtype + ret = isel_matcher.rewrite(self.stack_pointer().index(disp).load(dtype=ndt, tag=reg)) + assert ret is not None + return ret.replace(dtype=x.dtype) + + def asm_str(self, uops:list[UOp], function_name:str) -> str: + def _format_op(x:UOp) -> str: return f" {(o[7:-1] if (o:=str(x.arg))[-1] in ('i', 'm') else o[7:]).lower():7s}" + def _format_operands(x:UOp) -> str: + def _format(src:tuple[UOp, ...]) -> list[str]: + return [str(s.arg) if s.op is Ops.CONST else reg_strs[o].get(s.dtype.itemsize if not isinstance(s.dtype, PtrDType) else 8, o) if \ + (o:=str(s.reg)) in reg_strs else o for s in src if s.reg is not None] + def _mem_adress(base:UOp, idx:UOp, disp:UOp) -> list[str]: + return [f"[{base.reg}" + (f" + {idx.reg}*{base.dtype.itemsize}" if idx.reg else "") + (f" + {disp.arg}" if disp.arg else "") + "]"] + + if len(x.src) > 3 and x.arg in X86GroupOp.WriteMem: ret = _mem_adress(*x.src[:3]) + _format(x.src[3:]) + elif len(x.src) > 2 and x.arg in X86GroupOp.Rm1st: ret = _format((x,)) + _mem_adress(*x.src[:3]) + _format(x.src[3:]) + elif len(x.src) > 3 and x.arg in X86GroupOp.Rm2nd: ret = _format((x, x.src[0])) + _mem_adress(*x.src[1:4]) + _format(x.src[4:]) + else: ret = _format((x,) + x.src) + return ", ".join(ret) + + asm = [f".{function_name}:"] + for u in uops: + if u.op is not Ops.INS: continue + if u.arg is X86Ops.LABEL: asm.append(f"{str(u.tag)}:") + elif u.arg is X86Ops.RET: asm.append(_format_op(u)) + else: asm.append(_format_op(u) + " " + _format_operands(u)) + return "\n".join(asm) + + def render(self, uops:list[UOp]) -> str: + targets: dict[str, int] = {} + jumps: dict[UOp, int] = {} + binary = bytearray() + for u in uops: + if u.op is not Ops.INS: continue + if u.arg is X86Ops.LABEL: + targets[u.tag] = len(binary) + continue + if u.arg not in encodings or (l:=encodings[u.arg](u)) is None: + raise RuntimeError(f"failed to encode {u.arg} with {u.dtype} srcs {[x.dtype for x in u.src]}") + binary.extend(l) + if u.arg in (X86Ops.JL, X86Ops.JB, X86Ops.JE, X86Ops.JNE, X86Ops.JGE, X86Ops.JMP): jumps[u] = len(binary) + # fixup jump targets now that encoding size is known + for u in uops: + if (t:=jumps.get(u)) is not None: binary[t-4:t] = (targets[u.tag] - t).to_bytes(4, 'little', signed=True) + return binary.hex() + + def supported_dtypes(self): return {d for d in super().supported_dtypes() if d not in dtypes.fp8s+(dtypes.bfloat16,)} diff --git a/tinygrad/renderer/llvmir.py b/tinygrad/renderer/llvmir.py index e0ab52f697dbb..536564425d709 100644 --- a/tinygrad/renderer/llvmir.py +++ b/tinygrad/renderer/llvmir.py @@ -6,7 +6,7 @@ from tinygrad.uop.decompositions import xexp2, xlog2 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, GroupOp, range_str from tinygrad.dtype import dtypes, float_to_fp8, DType, PtrDType, truncate -from tinygrad.helpers import prod, Target, CPU_COUNT, getenv +from tinygrad.helpers import prod, Target, CPU_COUNT, getenv, OSX def ldt(dt:DType): if dt.vcount > 1: return f"<{dt.vcount} x {ldt(dt.scalar())}>" @@ -83,9 +83,8 @@ def render_wmma_amd(ctx, wmma: UOp, cdna=False) -> str: f" {ctx[x]}_yes = load {ldt(x.dtype)}, {ldt(idx.dtype)} {ctx[idx]}\n" f" br label {ctx[x]}_exit\n{ctx[x][1:]}_exit:\n" f" {ctx[x]} = phi {ldt(x.dtype)} [{ctx[x]}_yes, {ctx[x]}_load], [{ctx[alt]}, {ctx[x]}_entry]"), - (UPat(Ops.LOAD, src=(UPat.var('idx'),), name="x"), - lambda ctx,x,idx: f" {ctx[x]} = load {ldt(x.dtype)}, {ldt(idx.dtype)} {ctx[idx]}"), - (UPat(Ops.STORE, name="x"), lambda ctx,x: f" store {ldt(x.src[1].dtype)} {ctx[x.src[1]]}, {ldt(x.src[0].dtype)} {ctx[x.src[0]]}"), + (UPat.var('idx').load(name="x"), lambda ctx,x,idx: f" {ctx[x]} = load {ldt(x.dtype)}, {ldt(idx.dtype)} {ctx[idx]}"), + (UPat.var('idx').store(UPat.var("var")), lambda ctx,idx,var: f" store {ldt(var.dtype)} {ctx[var]}, {ldt(idx.dtype)} {ctx[idx]}"), # GEP/VECTORIZE/CAST for float4 support (UPat(Ops.GEP, name="x"), lambda ctx,x: f" {ctx[x]} = extractelement {ldt(x.src[0].dtype)} {ctx[x.src[0]]}, i32 {x.arg[0]}"), @@ -165,7 +164,7 @@ def _render_kernel(self, uops: list[UOp], prefix:list[str]|None=None) -> tuple[t if u.arg is not None: name = u.arg.function_name continue if u.op in (Ops.PARAM, Ops.DEFINE_VAR): - r[u] = f"%data{u.arg}" if u.op is Ops.PARAM else f"%{u.expr}" + r[u] = f"%data{u.arg.slot}" if u.op is Ops.PARAM else f"%{u.expr}" args.append((r[u], u.dtype)) elif u.op in (Ops.DEFINE_LOCAL, Ops.DEFINE_REG): r[u] = f"%{'local' if u.op is Ops.DEFINE_LOCAL else 'reg'}_{str(u.arg).replace('(', '').replace(')', '').replace(',', '_').replace(' ', '')}" @@ -206,6 +205,11 @@ def __init__(self, target:Target): if "AMX" in target.arch: self.tensor_cores = tc.amx self.compiler = CPULLVMCompiler([x for x in target.arch.split(",") if x != "AMX"]) + # FIXME: fp16 works on non-osx, but only if the cpu supports it + def supported_dtypes(self): + return {d for d in super().supported_dtypes() if + (d != dtypes.bfloat16 or self.target.arch.startswith(("x86", "arm"))) and (d != dtypes.half or OSX) and d not in dtypes.fp8s} + barrier = 'fence syncscope("workgroup") release\ntail call void @llvm.amdgcn.s.barrier()\nfence syncscope("workgroup") acquire\n' code_for_workitem = {"g": lambda x: f"tail call i32 @llvm.amdgcn.workgroup.id.{chr(120+int(x))}()", "l": lambda x: f"tail call i32 @llvm.amdgcn.workitem.id.{chr(120+int(x))}()"} @@ -291,3 +295,6 @@ def __init__(self, target:Target): lambda x: UOp(Ops.WMMA, dtypes.float.vec(8), (x.src[0].bitcast(dtypes.uint16.vec(8)), x.src[1].bitcast(dtypes.uint16.vec(8)), x.src[2]), (*x.arg,)) if x.src[0].dtype == dtypes.bfloat16.vec(8) else None) ]) + + def supported_dtypes(self): return {d for d in super().supported_dtypes() + if (d not in dtypes.fp8_ocp or self.target.arch == "gfx950") and d not in dtypes.fp8_fnuz} diff --git a/tinygrad/renderer/nir.py b/tinygrad/renderer/nir.py index 230e9a5f5d46d..e81713a78c136 100644 --- a/tinygrad/renderer/nir.py +++ b/tinygrad/renderer/nir.py @@ -138,7 +138,9 @@ class NIRRenderer(Renderer): # load/store use pointer arithmetic, and the cast does nothing. NOTE: this doesn't apply to image indexing cause it's 1-D (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off")), name="x"), lambda x,buf,off: x.replace( src=(buf,off.cast(dtypes.long))) if buf.dtype.addrspace != AddrSpace.REG and off.op not in (Ops.CAST, Ops.STACK) else None), - (UPat(Ops.CAST, name="x"), lambda x: x.src[0] if isinstance(x.dtype, PtrDType) or x.src[0].dtype == dtypes.void else None), + # images need index to be int for nir + (UPat.var("buf").index(UPat.var("idx_y"), UPat.var("idx_x")), + lambda buf,idx_y,idx_x: buf.index(idx_y.cast(dtypes.int), idx_x.cast(dtypes.int))), ]) def_rewrite = PatternMatcher([ @@ -146,13 +148,13 @@ class NIRRenderer(Renderer): (UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx.param(ctx.b, x, 8)), (UPat(Ops.DEFINE_VAR, name="x"), lambda ctx,x: ctx.param(ctx.b, x, 4)), (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: nchannel(ctx.b, {'g':ngid, 'l':nlid, 'i': nid}[x.arg[0]](ctx.b), int(x.arg[-1]))), - (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat.var("buf"),UPat.var("off"))), UPat.var("val")), allow_any_len=True), - lambda ctx,buf,off,val: nstore(ctx.b, buf.ptrdtype.addrspace, nidx(ctx.b, ctx.r[buf], ctx.r[off], buf.dtype), ctx.r[val], val.dtype)), - (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off"))), UPat.var("alt"), UPat.var("gate")), name="x"), + (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat.var("buf"),UPat.var("off"))).or_casted(), UPat.var("val"))), + lambda ctx,buf,off,val: nstore(ctx.b, buf.addrspace, nidx(ctx.b, ctx.r[buf], ctx.r[off], buf.dtype), ctx.r[val], val.dtype)), + (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off"))).or_casted(), UPat.var("alt"), UPat.var("gate")), name="x"), lambda ctx,x,buf,off,alt,gate: if_phi(ctx.b, ctx.r[gate], - lambda: nload(ctx.b, buf.ptrdtype.addrspace, nidx(ctx.b, ctx.r[buf], ctx.r[off], buf.dtype, ctx.r[gate]), x.dtype), lambda: ctx.r[alt])), - (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off"))),), name="x"), - lambda ctx,x,buf,off: nload(ctx.b, buf.ptrdtype.addrspace, nidx(ctx.b, ctx.r[buf], ctx.r[off], buf.dtype), x.dtype)), + lambda: nload(ctx.b, buf.addrspace, nidx(ctx.b, ctx.r[buf], ctx.r[off], buf.dtype, ctx.r[gate]), x.dtype), lambda: ctx.r[alt])), + (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off"))).or_casted(),), name="x"), + lambda ctx,x,buf,off: nload(ctx.b, buf.addrspace, nidx(ctx.b, ctx.r[buf], ctx.r[off], buf.dtype), x.dtype)), (UPat(Ops.STACK, name="x"), lambda ctx,x: nalu(ctx.b, f"vec{x.dtype.count}", *[ctx.r[src] for src in x.src])), (UPat(GroupOp.ALU, name="x"), lambda ctx,x: nalu(ctx.b, aop[x.src[0].dtype.scalar()][x.op], *[ctx.r[src] for src in x.src])), (UPat(Ops.CAST, name="x"), lambda ctx,x: ncast(ctx.b, ctx.r[x.src[0]], x.src[0].dtype, x.dtype)), @@ -188,6 +190,7 @@ def render(self, uops:list[UOp]): for u in uops: if u.op in {Ops.NOOP, Ops.GROUP, Ops.INDEX}: pass + elif u.op is Ops.CAST and isinstance(u.dtype, PtrDType): pass elif u.op is Ops.AFTER: self.r[u] = self.r[u.src[0]] elif u.op == Ops.SINK: @@ -226,11 +229,15 @@ def render(self, uops:list[UOp]): return ret + def supported_dtypes(self): return {d for d in Renderer.supported_dtypes(self) if d not in dtypes.fp8s+(dtypes.bfloat16,)} + class NAKRenderer(NIRRenderer): param = nir_instr(nc=1, num_components=1, bs=lambda sz:sz*8, also=lambda self,sz: setattr(self, "param_idx", self.param_idx + sz), intrins={"ALIGN_MUL":lambda sz:sz}, srcs=lambda self,b: [nsrc(nimm(b, 0, dtypes.int)), nsrc(nimm(b, self.param_idx, dtypes.int))])( lambda self, b, x, sz: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_ldc_nv)) + def supported_dtypes(self): return {d for d in super().supported_dtypes() if (d != dtypes.half or int(self.target.arch[3:]) >= 53)} + class LVPRenderer(NIRRenderer): has_local = False has_shared = False @@ -294,3 +301,5 @@ def postrender(self, uops:list[UOp]): self.b.shader.contents.info.num_ubos = len([u for u in bufs if not isinstance(u.dtype, ImageDType)]) self.b.shader.contents.info.num_images = texs() + imgs() + + def supported_dtypes(self): return {d for d in NIRRenderer.supported_dtypes(self) if d != dtypes.double} diff --git a/tinygrad/renderer/ptx.py b/tinygrad/renderer/ptx.py index dd9ac068b69fb..f20ae63218af7 100644 --- a/tinygrad/renderer/ptx.py +++ b/tinygrad/renderer/ptx.py @@ -54,8 +54,6 @@ def render_val(x, dtype): (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx")), name="op"), lambda buf,idx,op: UOp(Ops.INDEX, dtype=dtypes.int64, src=(buf, buf.cast(dtypes.int64)+idx.cast(dtypes.int64)*buf.dtype.itemsize)+op.src[2:]) \ if op.dtype != dtypes.int64 and buf.dtype.addrspace != AddrSpace.REG else None), - # load/store use pointer arithmetic, and the cast does nothing - (UPat(Ops.CAST, name="x"), lambda x: x.src[0] if isinstance(x.dtype, PtrDType) or x.src[0].dtype == dtypes.void else None), # ptx shr and shl instructions require y to be uint (UPat.var("x") << UPat.var("y"), lambda x,y: UOp(Ops.SHL, x.dtype, (x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None), (UPat.var("x") >> UPat.var("y"), lambda x,y: UOp(Ops.SHR, x.dtype, (x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None), @@ -65,7 +63,7 @@ def mem_type(x:UOp) -> str: match x.op: case Ops.AFTER: return mem_type(x.src[0]) case Ops.DEFINE_LOCAL: return 'shared' - case Ops.PARAM: return 'global' + case Ops.PARAM: return 'shared' if x.addrspace == AddrSpace.LOCAL else 'global' case _: raise RuntimeError(f"{x.op} needs to be memory") def render_wmma(ctx: "PTXRenderer", wmma: UOp): @@ -92,7 +90,7 @@ def modifier(a: DType, b: DType): return '.rzi' if dtypes.is_int(a) and dtypes.i (UPat.cvar("x", dtypes.bool), lambda ctx, x: f"setp.ne.s16 {ctx.r[x]}, {render_val(x.arg, x.dtype)}, 0;"), (UPat.cvar("x"), lambda ctx, x: f"mov.b{ctx.types[x.dtype][1:]} {ctx.r[x]}, {render_val(x.arg, x.dtype)};"), (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: f"mov.u32 %{x.arg}, %{'ctaid' if x.arg[0] == 'g' else 'tid'}.{chr(120+int(x.arg[-1]))};"), - (UPat(Ops.PARAM, name="x"), lambda ctx, x: f"ld.param.{ctx.types[dtypes.ulong]} {ctx.r[x]}, [data{x.arg}+0];"), + (UPat(Ops.PARAM, name="x"), lambda ctx, x: f"ld.param.{ctx.types[dtypes.ulong]} {ctx.r[x]}, [data{x.arg.slot}+0];"), (UPat((Ops.CMPLT, Ops.CMPNE, Ops.CMPEQ), name="x", allow_any_len=True, src=(UPat.var("src0"),)), lambda ctx, x, src0: ctx.code_for_op[x.op](ctx.r[x], *[ctx.r[v] for v in x.src], src0.dtype, ctx.types[src0.dtype])), (UPat(GroupOp.ALU, name="x"), lambda ctx, x: ctx.code_for_op[x.op](ctx.r[x], *[ctx.r[v] for v in x.src], x.dtype, ctx.types[x.dtype])), @@ -102,18 +100,18 @@ def modifier(a: DType, b: DType): return '.rzi' if dtypes.is_int(a) and dtypes.i (UPat(Ops.CAST, name="x", src=(UPat.var("a"),)), lambda ctx, x, a: f"cvt{modifier(x.dtype, a.dtype)}.{ctx.cast_types[x.dtype]}.{ctx.cast_types[a.dtype]} {ctx.r[x]}, {ctx.r[a]};"), # store / gated load / load - (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("loc"))), UPat.var("var")), allow_any_len=True), + (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("loc"))).or_casted(), UPat.var("var"))), lambda ctx, loc, var, buf: f"st.{mem_type(buf)}" + \ f"{f'.v{cnt}' if ((cnt:=var.dtype.count)>1) else ''}.{ctx.mem_types[var.dtype.scalar()]} " + \ f"[{ctx.r[loc]}+0], {('{' + ', '.join(ctx.r[var]) + '}') if var.dtype.count > 1 else ctx.r[var]};"), - (UPat(Ops.LOAD, name="x", src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("loc"))), UPat.var("alt"), UPat.var("gate")), allow_any_len=True), + (UPat(Ops.LOAD, name="x", src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("loc"))).or_casted(), UPat.var("alt"), UPat.var("gate"))), lambda ctx, x, loc, alt, gate, buf: flatten([ [f"mov.{ctx.mem_types[x.dtype.scalar()]} {v}, {render_val(0, x.dtype.scalar())};" for v in ctx.r[x]], [f"@{ctx.r[gate]} ld.{mem_type(buf)}.v{x.dtype.count}.{ctx.mem_types[x.dtype.scalar()]} {{{', '.join(ctx.r[x])}}}, [{ctx.r[loc]}+0];"] ]) if alt.dtype.count > 1 else [ f"@{ctx.r[gate]} ld.{mem_type(buf)}.{ctx.mem_types[x.dtype.scalar()]} {ctx.r[x]}, [{ctx.r[loc]}+0];", f"@!{ctx.r[gate]} mov.b{ctx.types[x.dtype.scalar()][1:]} {ctx.r[x]}, {ctx.r[alt]};"]), - (UPat(Ops.LOAD, name="x", src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("loc"))),), allow_any_len=True), + (UPat(Ops.LOAD, name="x", src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("loc"))).or_casted(),)), lambda ctx, x, loc, buf: f"ld.{mem_type(buf)}.v{x.dtype.count}.{ctx.mem_types[x.dtype.scalar()]} {{{', '.join(ctx.r[x])}}}, [{ctx.r[loc]}+0];" \ if x.dtype.count > 1 else f"ld.{mem_type(buf)}.{ctx.mem_types[x.dtype]} {ctx.r[x]}, [{ctx.r[loc]}+0];"), # simple @@ -187,6 +185,7 @@ def ssa(prefix:str, u:UOp|None=None, dtype:str|None=None) -> str: name = "test" for u in uops: if u.op in {Ops.NOOP, Ops.GROUP}: continue + if u.op is Ops.CAST and isinstance(u.dtype, PtrDType): continue if u.op is Ops.AFTER: self.r[u] = self.r[u.src[0]] continue @@ -203,7 +202,7 @@ def ssa(prefix:str, u:UOp|None=None, dtype:str|None=None) -> str: r[u] = r[u.src[0]] continue if u.op is Ops.DEFINE_REG: - r[u] = [ssa("reg", u, self.types[u.dtype.base.scalar()]) for _ in range(u.ptrdtype.size)] + r[u] = [ssa("reg", u, self.types[u.dtype.base.scalar()]) for _ in range(u.max_numel())] continue if u.op in {Ops.INDEX, Ops.LOAD, Ops.STORE} and isinstance(u.src[0].dtype, PtrDType) and u.src[0].dtype.addrspace == AddrSpace.REG: if u.op is Ops.INDEX: @@ -219,9 +218,8 @@ def ssa(prefix:str, u:UOp|None=None, dtype:str|None=None) -> str: if u.op is Ops.SPECIAL: r[u] = "%" + u.arg elif u.op is Ops.DEFINE_VAR: bufs.append((u.expr, u.dtype)) elif u.op is Ops.LOAD: - assert u.src[0].dtype == dtypes.int64, "load isn't int64" r[u] = [ssa('val', dtype=self.types[u.dtype.scalar()]) for _ in range(u.dtype.count)] if u.dtype.count > 1 else ssa('val', u) - elif u.op is Ops.PARAM: bufs.append((f"data{u.arg}", u.dtype)) + elif u.op is Ops.PARAM: bufs.append((f"data{u.arg.slot}", u.dtype)) elif u.op is Ops.WMMA: # registers for packing/unpacking input and acc self.wmma_r = [[ssa("wmma_in", dtype="b32") for _ in range(0, len(r[u.src[0]]), 4 // u.src[0].dtype.scalar().itemsize)], @@ -239,3 +237,6 @@ def ssa(prefix:str, u:UOp|None=None, dtype:str|None=None) -> str: if u.op is Ops.SPECIAL: kernel = [f".reg .u32 %{u.arg};"] + kernel return self.render_kernel(kernel, name, bufs, c.items(), uops) + + def supported_dtypes(self): return {d for d in super().supported_dtypes() + if (d != dtypes.half or int(self.target.arch[3:]) >= 53) and d not in dtypes.fp8s+(dtypes.bfloat16,)} diff --git a/tinygrad/renderer/wgsl.py b/tinygrad/renderer/wgsl.py index f7abf0d54f5fd..d6db9ea0f88ac 100644 --- a/tinygrad/renderer/wgsl.py +++ b/tinygrad/renderer/wgsl.py @@ -86,7 +86,7 @@ class WGSLRenderer(CStyleLanguage): (UPat.load(UPat.var("b"), UPat.cvar("v"), UPat.var("gate")), lambda ctx,b,v,gate: f"select({ctx[v]}, {ctx.render_load(ctx[b],b.src[0].dtype)}, {ctx[gate]})"), (UPat.load(UPat.var("b")), lambda ctx, b: ctx.render_load(ctx[b], b.dtype)), - (UPat.store(UPat.var("b"), UPat.var("v"), allow_any_len=True),lambda ctx,b,v:\ + (UPat.store(UPat.var("b"), UPat.var("v")), lambda ctx,b,v:\ # (load & mask) | var -> mask = v.src[0].src[1], var = v.src[1] f"atomicAnd(&{ctx[b]},{ctx[v.src[0].src[1]]});\n atomicAdd(&{ctx[b]},{ctx[v.src[1]]});" if is_packed(b.src[0].dtype) \ else f"{ctx[b]} = {ctx[v]};"), @@ -112,3 +112,6 @@ def render_kernel(self, function_name:str, kernel:list[str], bufs:list[tuple[str f"{name}:{f'array<{self.buf_map(dtype.base)}>' if isinstance(dtype,PtrDType) else self.buf_map(dtype)};" for name,(dtype,_) in bufs]) prg += f"\n@compute @workgroup_size({','.join([str(x) for x in local_size])}) fn {function_name}(@builtin(workgroup_id) gindex: vec3," return prg + "@builtin(local_invocation_id) lindex: vec3) {\n" + "\n".join(kernel) + "\n}" + + def supported_dtypes(self): return {dtypes.bool, dtypes.char, dtypes.uchar, dtypes.short, dtypes.ushort, dtypes.int32, dtypes.uint32, + dtypes.float, *((dtypes.half,) if "shader-f16" in self.target.arch else ())} diff --git a/tinygrad/runtime/autogen/__init__.py b/tinygrad/runtime/autogen/__init__.py index 6d32d34902973..eb4b61c846df8 100644 --- a/tinygrad/runtime/autogen/__init__.py +++ b/tinygrad/runtime/autogen/__init__.py @@ -54,7 +54,7 @@ def __getattr__(nm): ["/usr/include/string.h", "/usr/include/elf.h", "/usr/include/unistd.h", "/usr/include/asm-generic/mman-common.h"]), dll="'c'", errno=True) case "avcodec": return load("avcodec", ["{}/libavcodec/hevc/hevc.h", "{}/libavcodec/cbs_h265.h"], srcs=ffmpeg_src) case "opencl": return load("opencl", ["{}/CL/cl.h"], dll="'OpenCL'", args=["-I{}"], srcs=opencl_src) - case "cuda": return load("cuda", ["{}/include/cuda.h"], dll="'cuda'", args=["-D__CUDA_API_VERSION_INTERNAL"], srcs=cudart_src, macros=False) + case "cuda": return load("cuda", ["{}/include/cuda.h"], dll="'nvcuda' if WIN else 'cuda'", args=["-D__CUDA_API_VERSION_INTERNAL"], srcs=cudart_src, macros=False, prolog=["from tinygrad.helpers import WIN"]) case "nvrtc": return load("nvrtc", ["{}/include/nvrtc.h"], dll="'nvrtc'", paths=nv_lib_path, srcs=nvrtc_src, prolog=["import sysconfig"]) case "nvjitlink": load("nvjitlink", [root/"extra/nvJitLink.h"], dll="'nvJitLink'", paths=nv_lib_path, prolog=["import sysconfig"]) case "kfd": return load("kfd", [root/"extra/hip_gpu_driver/kfd_ioctl.h"]) @@ -155,7 +155,7 @@ def __getattr__(nm): *[f"python3 src/compiler/{s}_h.py > gen/{s.split('/')[-1]}.h" for s in ["nir/nir_opcodes", "nir/nir_builder_opcodes"]], *[f"python3 src/compiler/nir/nir_{s}_h.py --outdir gen" for s in ["intrinsics", "intrinsics_indices"]]]), cwd=path, shell=True, check=True), srcs="https://gitlab.freedesktop.org/mesa/mesa/-/archive/mesa-25.2.7/mesa-25.2.7.tar.gz", - dll="([] if (_cpu:=DEV.renderer == 'LVP') else ['tinymesa']) + ['tinymesa_cpu'], " \ + dll="'tinymesa_cpu' if (_cpu:=DEV.renderer == 'LVP') else 'tinymesa', " \ 'emsg="not available on this platform" if WIN or (OSX and (platform.machine() != "arm64" or (_mv:=platform.mac_ver()[0][:2]) not in {"14","15","26"})) or (platform.system() == "Linux" and platform.machine() not in {"x86_64", "aarch64"}) else ' \ 'f"run `sudo curl -fL https://github.com/sirhcm/tinymesa/releases/download/v1/libtinymesa{\'_cpu\'*_cpu}-mesa-25.2.7-{\'macos-\'+_mv if OSX else \'linux\'}-{\'amd64\' if ARCH_X86 else \'arm64\'}.{\'dylib\' if OSX else \'so\'} -o /usr/local/lib/libtinymesa{\'_cpu\'*_cpu}.{\'dylib\' if OSX else \'so\'}`"', prolog=["from tinygrad.helpers import DEV, ARCH_X86, WIN, OSX", "import gzip, base64, platform"], diff --git a/tinygrad/runtime/autogen/amd/__init__.py b/tinygrad/runtime/autogen/amd/__init__.py new file mode 100644 index 0000000000000..e69de29bb2d1d diff --git a/tinygrad/runtime/autogen/cuda.py b/tinygrad/runtime/autogen/cuda.py index 07b0abd2ebfa0..c6c805369e7ba 100644 --- a/tinygrad/runtime/autogen/cuda.py +++ b/tinygrad/runtime/autogen/cuda.py @@ -4,7 +4,8 @@ from typing import Literal, TypeAlias from tinygrad.runtime.support.c import _IO, _IOW, _IOR, _IOWR from tinygrad.runtime.support import c -dll = c.DLL('cuda', 'cuda') +from tinygrad.helpers import WIN +dll = c.DLL('cuda', 'nvcuda' if WIN else 'cuda') cuuint32_t: TypeAlias = ctypes.c_uint32 cuuint64_t: TypeAlias = ctypes.c_uint64 CUdeviceptr_v2: TypeAlias = ctypes.c_uint64 diff --git a/tinygrad/runtime/autogen/mesa.py b/tinygrad/runtime/autogen/mesa.py index 80ff1cf648331..611e3057d52e2 100644 --- a/tinygrad/runtime/autogen/mesa.py +++ b/tinygrad/runtime/autogen/mesa.py @@ -6,7 +6,7 @@ from tinygrad.runtime.support import c from tinygrad.helpers import DEV, ARCH_X86, WIN, OSX import gzip, base64, platform -dll = c.DLL('mesa', ([] if (_cpu:=DEV.renderer == 'LVP') else ['tinymesa']) + ['tinymesa_cpu'], emsg="not available on this platform" if WIN or (OSX and (platform.machine() != "arm64" or (_mv:=platform.mac_ver()[0][:2]) not in {"14","15","26"})) or (platform.system() == "Linux" and platform.machine() not in {"x86_64", "aarch64"}) else f"run `sudo curl -fL https://github.com/sirhcm/tinymesa/releases/download/v1/libtinymesa{'_cpu'*_cpu}-mesa-25.2.7-{'macos-'+_mv if OSX else 'linux'}-{'amd64' if ARCH_X86 else 'arm64'}.{'dylib' if OSX else 'so'} -o /usr/local/lib/libtinymesa{'_cpu'*_cpu}.{'dylib' if OSX else 'so'}`") +dll = c.DLL('mesa', 'tinymesa_cpu' if (_cpu:=DEV.renderer == 'LVP') else 'tinymesa', emsg="not available on this platform" if WIN or (OSX and (platform.machine() != "arm64" or (_mv:=platform.mac_ver()[0][:2]) not in {"14","15","26"})) or (platform.system() == "Linux" and platform.machine() not in {"x86_64", "aarch64"}) else f"run `sudo curl -fL https://github.com/sirhcm/tinymesa/releases/download/v1/libtinymesa{'_cpu'*_cpu}-mesa-25.2.7-{'macos-'+_mv if OSX else 'linux'}-{'amd64' if ARCH_X86 else 'arm64'}.{'dylib' if OSX else 'so'} -o /usr/local/lib/libtinymesa{'_cpu'*_cpu}.{'dylib' if OSX else 'so'}`") class struct_u_printf_info(c.Struct): pass u_printf_info: TypeAlias = struct_u_printf_info uint32_t: TypeAlias = ctypes.c_uint32 diff --git a/tinygrad/runtime/graph/metal.py b/tinygrad/runtime/graph/metal.py index 25c89257e557b..228275d787d1a 100644 --- a/tinygrad/runtime/graph/metal.py +++ b/tinygrad/runtime/graph/metal.py @@ -23,7 +23,7 @@ def __init__(self, linear, input_uops=()): self.icb = self.dev.sysdevice.newIndirectCommandBufferWithDescriptor_maxCommandCount_options(icb_descriptor, len(self.calls), metal.MTLResourceCPUCacheModeDefaultCache) if self.icb.value is None: raise GraphException("create indirect command buffer failed, does your system support this?") - self.needs_icb_fix = int(self.dev.gpu_family < 9) # ICB fix not required on M3+ (Apple9+) + self.needs_icb_fix = int(not self.dev.arch.startswith("Apple") or int(self.dev.arch[5:]) < 9) # ICB fix not required on M3+ (Apple9+) if len(self.vars): self.int_buf = self.dev.allocator.alloc(len(self.vars)*dtypes.int32.itemsize) @@ -109,5 +109,5 @@ def __del__(self): @staticmethod def supports_uop(batch_devs, new_call:UOp) -> bool: # Metal ICB replay encodes offsets as uint32; reject if any Metal buffer offset exceeds 32-bit range. - if any(b.op is Ops.BUFFER_VIEW and b.arg[1] * b.dtype.itemsize > 0xFFFFFFFF for b in new_call.src[1:]): return False + if any(b.op is Ops.SLICE and b.src[1].arg * b.src[0].dtype.itemsize > 0xFFFFFFFF for b in new_call.src[1:]): return False return GraphRunner.supports_uop(batch_devs, new_call) diff --git a/tinygrad/runtime/ops_cl.py b/tinygrad/runtime/ops_cl.py index 3ee42f9db83f3..13a4e96d54c2e 100644 --- a/tinygrad/runtime/ops_cl.py +++ b/tinygrad/runtime/ops_cl.py @@ -112,18 +112,18 @@ def __init__(self, device:str=""): self.context = checked(cl.clCreateContext(None, 1, self.device_id, CC_CB(), None, status := ctypes.c_int32()), status) self.queue = checked(cl.clCreateCommandQueue(self.context, self.device_id, cl.CL_QUEUE_PROFILING_ENABLE, status), status) self.pending_copyin: list[memoryview] = [] - self.device_exts = (cl.clGetDeviceInfo(self.device_id, cl.CL_DEVICE_EXTENSIONS, 4096, - ctypes.byref(buf := ctypes.create_string_buffer(4096)), - ctypes.byref(total := ctypes.c_size_t())), - ctypes.string_at(buf, size=total.value).decode())[1] + check(cl.clGetDeviceInfo(self.device_id, cl.CL_DEVICE_EXTENSIONS, 0, None, ctypes.byref(exts_len:=ctypes.c_size_t()))) + self.device_exts = (cl.clGetDeviceInfo(self.device_id, cl.CL_DEVICE_EXTENSIONS, exts_len.value, + ctypes.byref(buf := ctypes.create_string_buffer(exts_len.value)), None), + ctypes.string_at(buf).decode().split())[1] renderer = IntelRenderer if "cl_intel_subgroup_matrix_multiply_accumulate" in self.device_exts else OpenCLRenderer self.cl_compiler = CLCompiler(self, f"{hashlib.md5(self.device_name.encode() + self.driver_version.encode()).hexdigest()}") + arch = ",".join(self.device_exts) if "cl_khr_image2d_from_buffer" in self.device_exts: check(cl.clGetDeviceInfo(self.device_id, cl.CL_DEVICE_IMAGE_PITCH_ALIGNMENT, 4, ctypes.byref(ipa := ctypes.c_uint32()), None)) - arch = f"IMAGE_PITCH_ALIGNMENT={ipa.value}" - else: arch = "" + arch += f",IMAGE_PITCH_ALIGNMENT={ipa.value}" super().__init__(device, CLAllocator(self), [renderer], functools.partial(CLProgram, self), arch=arch) def count(self) -> int: return len(unwrap(self.device_ids)) diff --git a/tinygrad/runtime/ops_cpu.py b/tinygrad/runtime/ops_cpu.py index 1ce80d7aa882a..37c59c9338c4a 100644 --- a/tinygrad/runtime/ops_cpu.py +++ b/tinygrad/runtime/ops_cpu.py @@ -4,9 +4,10 @@ from tinygrad.device import BufferSpec from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, HWQueue, HCQArgsState, HCQSignal, HCQProgram, MMIOInterface from tinygrad.runtime.support.hcq import CLikeArgsState -from tinygrad.renderer.cstyle import ClangJITRenderer +from tinygrad.renderer.cstyle import ClangRenderer from tinygrad.renderer.llvmir import CPULLVMRenderer from tinygrad.renderer.nir import LVPRenderer +from tinygrad.renderer.isa.x86 import X86Renderer from tinygrad.runtime.support.elf import jit_loader from tinygrad.uop.ops import sint @@ -137,5 +138,5 @@ class CPUDevice(HCQCompiled): def __init__(self, device:str=""): self.tasks:queue.Queue = queue.Queue() CPUWorker(self, self.tasks, thread_id=0).start() - super().__init__(device, CPUAllocator(self), [ClangJITRenderer, CPULLVMRenderer, LVPRenderer], functools.partial(CPUProgram, self), CPUSignal, - CPUComputeQueue, arch={'amd64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine().lower(), m)+",native") + super().__init__(device, CPUAllocator(self), [ClangRenderer, CPULLVMRenderer, LVPRenderer, X86Renderer], functools.partial(CPUProgram, self), + CPUSignal, CPUComputeQueue, arch={'amd64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine().lower(), m)+",native") diff --git a/tinygrad/runtime/ops_dsp.py b/tinygrad/runtime/ops_dsp.py index 70f7b048ae71c..8bcfcb563ee9e 100644 --- a/tinygrad/runtime/ops_dsp.py +++ b/tinygrad/runtime/ops_dsp.py @@ -68,6 +68,8 @@ def _render_entry(self, function_name:str, bufs:list[tuple[str,tuple[DType,bool] msrc += ["return 0; }"] return '\n'.join(msrc) + def supported_dtypes(self): return {d for d in super().supported_dtypes() if d not in dtypes.fp8s+(dtypes.bfloat16,)} + def rpc_sc(method=0, ins=0, outs=0, fds=0): return (method << 24) | (ins << 16) | (outs << 8) | fds def rpc_prep_args(ins=None, outs=None, in_fds=None): ins, outs, in_fds = ins or list(), outs or list(), in_fds or list() diff --git a/tinygrad/runtime/ops_metal.py b/tinygrad/runtime/ops_metal.py index bc9747ba7635e..1e14bce6a2d5d 100644 --- a/tinygrad/runtime/ops_metal.py +++ b/tinygrad/runtime/ops_metal.py @@ -37,12 +37,8 @@ def __init__(self, device:str): self.timeline_signal = self.sysdevice.newSharedEvent() self.timeline_value = 0 - # probe GPU family: Apple9=M3/M4, Apple8=M2, Apple7=M1, etc. values are 1000+N. - self.gpu_family = 0 - for i in range(15, 0, -1): - if self.sysdevice.supportsFamily(1000 + i): - self.gpu_family = i - break + # https://developer.apple.com/documentation/metal/mtlgpufamily + def check_family(f): return next(filter(self.sysdevice.supportsFamily, reversed([v for v, nm in metal.enum_MTLGPUFamily.items() if f in nm])), 0) Compiled.profile_events += [ProfileDeviceEvent(device)] @@ -51,7 +47,7 @@ def __init__(self, device:str): # This can be reproduced locally with any virtualization software (like utm) that can create macOS VMs with apple's own virtualization framework. super().__init__(device, MetalAllocator(self), [MetalRenderer], functools.partial(MetalProgram, self), MetalGraph if 'virtual' not in from_ns_str(self.sysdevice.name()).lower() else None, - arch=platform.machine()) + arch=metal.enum_MTLGPUFamily[check_family("Apple") or check_family("Mac")][12:]) def synchronize(self): for cbuf in self.mtl_buffers_in_flight: diff --git a/tinygrad/runtime/ops_nv.py b/tinygrad/runtime/ops_nv.py index de70781e9d2c9..3be5a10c30fe8 100644 --- a/tinygrad/runtime/ops_nv.py +++ b/tinygrad/runtime/ops_nv.py @@ -554,7 +554,6 @@ class PCIIface(PCIIfaceBase): def __init__(self, dev, dev_id): # PCIIface's MAP_FIXED mmap will overwrite UVM allocations made by NVKIface, so don't try PCIIface if kernel driver was already used. if NVKIface.root is not None: raise RuntimeError("Cannot use PCIIface after NVKIface has been initialized (would corrupt UVM memory)") - if not OSX: System.reserve_hugepages(64) super().__init__(dev, dev_id, vendor=0x10de, devices=((0xff00, (0x2200,0x2400,0x2500,0x2600,0x2700,0x2800,0x2b00,0x2c00,0x2d00,0x2f00)),), base_class=0x03, vram_bar=1, va_start=NVMemoryManager.va_allocator.base, va_size=NVMemoryManager.va_allocator.size, dev_impl_t=NVDev) diff --git a/tinygrad/runtime/ops_python.py b/tinygrad/runtime/ops_python.py index 7c5a2c365aeaa..30ffa4625c64c 100644 --- a/tinygrad/runtime/ops_python.py +++ b/tinygrad/runtime/ops_python.py @@ -41,57 +41,68 @@ def generic_wmma_helper(inp, warp_size, WARP_THREADS, K, NUM_A, NUM_B, NUM_C, a_ class PythonProgram: def __init__(self, name:str, lib:bytes, **kwargs): - self.uops: list[tuple[Ops, DType, list[int], Any]] = pickle.loads(lib) + self.uops: list[UOp] = pickle.loads(lib) + self.uop_to_index: dict[UOp, int] = {u:i for i,u in enumerate(self.uops)} + self.loop_ends: dict[UOp, int] = {u.src[1]:i for i, u in enumerate(self.uops) if u.op == Ops.END} def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False, **kw): st = time.perf_counter() warp = list(itertools.product(*[range(x) for x in local_size[::-1]])) warp_size = len(warp) void_ops = {Ops.END, Ops.BARRIER, Ops.IF, Ops.ENDIF, Ops.SINK, Ops.NOOP, Ops.GROUP, Ops.STORE} - loop_ends: dict[int, int] = {srcs[1]:i for i, (uop, _, srcs, _) in enumerate(self.uops) if uop == Ops.END} for idxs in itertools.product(*[range(x) for x in global_size[::-1]]): - values: dict[int, Any] = {} + values: dict[UOp, Any] = {} pbufs: list[memoryview] = list(bufs) pvals: list[int] = list(vals) + exec_masks = [[True] * warp_size] i = 0 while i < len(self.uops): - uop, dtype, srcs, arg = self.uops[i] - src_values = [values[v] for v in srcs if self.uops[v][0] not in void_ops] - src_dtypes = [self.uops[v][1] for v in srcs if self.uops[v][0] not in void_ops] - if getenv("TRACE"): print(i, uop, dtype, arg, src_values, src_dtypes) - if uop is Ops.END: - i = srcs[1] + u = self.uops[i] + src_values = [values[v] for v in u.src if v.op not in void_ops] + src_dtypes = [v.dtype for v in u.src if v.op not in void_ops] + if getenv("TRACE"): print(i, u.op, u.dtype, u.arg, src_values, src_dtypes) + if u.op is Ops.END: + i = self.uop_to_index[u.src[1]] continue - if uop in (Ops.BARRIER, Ops.IF, Ops.ENDIF, Ops.SINK, Ops.NOOP, Ops.GROUP): + if u.op is Ops.IF: + exec_masks.append([x and y for x,y in zip(exec_masks[-1], src_values[0])]) + i += 1 + continue + if u.op is Ops.ENDIF: + exec_masks.pop() + i += 1 + continue + if u.op in (Ops.BARRIER, Ops.SINK, Ops.NOOP, Ops.GROUP): # in the python emulator, the warp is always in sync i += 1 continue - assert dtype is not None, f"{uop} is missing a dtype" - if uop is Ops.STORE: - store_gate = src_values[2] if len(src_values) >= 3 else [True] * warp_size + assert u.dtype is not None, f"{u.op} is missing a dtype" + if u.op is Ops.STORE: + assert len(src_values) == 2, f"STORE must be lowered to 2 srcs, got {len(src_values)}" + store_gate = exec_masks[-1] for j,val in enumerate(src_values[1] if src_dtypes[1].count > 1 else [src_values[1]]): for (m,o),v,g in zip(src_values[0], val, store_gate): if g: _store(m, o+j, v, src_dtypes[1].scalar()) i += 1 continue - if uop is Ops.AFTER: values[i] = src_values[0] - elif uop in {Ops.PARAM, Ops.DEFINE_LOCAL, Ops.DEFINE_REG}: - assert isinstance(dtype, PtrDType), dtype - storage_fmt = storage_fmt_for_dtype(dtype.base.scalar()) - if storage_fmt is None: raise RuntimeError(f"{dtype=} is not supported") + if u.op is Ops.AFTER: values[u] = src_values[0] + elif u.op in {Ops.PARAM, Ops.DEFINE_LOCAL, Ops.DEFINE_REG}: + assert isinstance(u.dtype, PtrDType), u.dtype + storage_fmt = storage_fmt_for_dtype(u.dtype.base.scalar()) + if storage_fmt is None: raise RuntimeError(f"dtype={u.dtype} is not supported") if TYPE_CHECKING or sys.version_info < (3, 12): assert storage_fmt != "e" - if uop is Ops.DEFINE_REG: + if u.op is Ops.DEFINE_REG: # REGs are per thread - values[i] = [memoryview(bytearray(dtype.size*dtype.itemsize)).cast(storage_fmt) for _ in range(warp_size)] + values[u] = [memoryview(bytearray(u.dtype.size*u.dtype.itemsize)).cast(storage_fmt) for _ in range(warp_size)] else: - buf = memoryview(bytearray(dtype.size*dtype.itemsize)) if uop is not Ops.PARAM else pbufs.pop(0) - values[i] = [buf.cast(storage_fmt)] * warp_size - elif uop is Ops.DEFINE_VAR: - values[i] = [pvals.pop(0)] * warp_size - elif uop is Ops.SPECIAL: - if arg[0] == 'g': values[i] = [idxs[2-int(arg[-1])]] * warp_size - elif arg[0] == 'l': values[i] = [x[2-int(arg[-1])] for x in warp] - elif uop is Ops.CONST: values[i] = [arg] * warp_size - elif uop is Ops.INDEX: + buf = memoryview(bytearray(u.dtype.size*u.dtype.itemsize)) if u.op is not Ops.PARAM else pbufs.pop(0) + values[u] = [buf.cast(storage_fmt)] * warp_size + elif u.op is Ops.DEFINE_VAR: + values[u] = [pvals.pop(0)] * warp_size + elif u.op is Ops.SPECIAL: + if u.arg[0] == 'g': values[u] = [idxs[2-int(u.arg[-1])]] * warp_size + elif u.arg[0] == 'l': values[u] = [x[2-int(u.arg[-1])] for x in warp] + elif u.op is Ops.CONST: values[u] = [u.arg] * warp_size + elif u.op is Ops.INDEX: ret:list = [] if isinstance(src_dtypes[0], ImageDType): assert len(src_values) == 3, "image index must be 3 srcs" @@ -101,33 +112,33 @@ def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tup else: assert len(src_values) == 2, "non-image index must be 2 srcs" for m,o in zip(*src_values): ret.append((m,o)) - values[i] = ret - elif uop is Ops.CAST and isinstance(dtype, PtrDType): - values[i] = src_values[0] - elif uop is Ops.RANGE: - if i not in values: values[i] = [0] * warp_size + values[u] = ret + elif u.op is Ops.CAST and isinstance(u.dtype, PtrDType): + values[u] = src_values[0] + elif u.op is Ops.RANGE: + if u not in values: values[u] = [0] * warp_size else: - for j in range(len(values[i])): - values[i][j] += 1 - if values[i][0] == src_values[0][0]: - del values[i] - i = loop_ends[i] + 1 + for j in range(len(values[u])): + values[u][j] += 1 + if values[u][0] == src_values[0][0]: + del values[u] + i = self.loop_ends[u] + 1 continue - elif uop is Ops.STACK: values[i] = src_values - elif uop is Ops.BITCAST: values[i] = [bitcast(x, src_dtypes[0], dtype) for x in src_values[0]] - elif uop is Ops.CAST: - values[i] = [truncate.get(dtype, lambda dt: dt)(dtype.const(x)) for x in src_values[0]] - elif uop is Ops.LOAD: - if dtype.count > 1: - values[i] = [load([src_values[i][j] if i != 0 and src_dtypes[i].count > 1 else src_values[i] \ - for i in range(len(src_values))], j, dtype.scalar()) for j in range(dtype.count)] + elif u.op is Ops.STACK: values[u] = src_values + elif u.op is Ops.BITCAST: values[u] = [bitcast(x, src_dtypes[0], u.dtype) for x in src_values[0]] + elif u.op is Ops.CAST: + values[u] = [truncate.get(u.dtype, lambda dt: dt)(u.dtype.const(x)) for x in src_values[0]] + elif u.op is Ops.LOAD: + if u.dtype.count > 1: + values[u] = [load([src_values[k][j] if k != 0 and src_dtypes[k].count > 1 else src_values[k] \ + for k in range(len(src_values))], j, u.dtype.scalar()) for j in range(u.dtype.count)] else: - values[i] = load(src_values, 0, dtype) - elif uop is Ops.GEP: values[i] = src_values[0][get_single_element(arg)] - elif uop is Ops.WMMA: - first_src_dtype = self.uops[srcs[0]][1] + values[u] = load(src_values, 0, u.dtype) + elif u.op is Ops.GEP: values[u] = src_values[0][get_single_element(u.arg)] + elif u.op is Ops.WMMA: + first_src_dtype = u.src[0].dtype assert isinstance(first_src_dtype, DType) # mypy - dims, dtype_in, device, threads = arg[1], first_src_dtype.scalar(), arg[4], arg[5] + dims, dtype_in, device, threads = u.arg[1], first_src_dtype.scalar(), u.arg[4], u.arg[5] wmma_helper = functools.partial(generic_wmma_helper, src_values, warp_size) # TODO: refactor these to a shared TensorCoreLayout if device == "METAL": @@ -135,17 +146,17 @@ def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tup def a_b_elem(x, i, j, goff): return x[(i%2)][goff+(i//2)%2+(j%4)*2+(i//4)*8+(j//4)*16] # (i, j), C, D (2 elements on 32 threads): row major same as A/B def c_map(lane, elem): return (elem + ((lane%2)*2) + ((lane//8)%2)*4, ((lane//2)%4) + (lane//16)*4) - values[i] = wmma_helper(32, 8, 2, 2, 2, a_b_elem, a_b_elem, c_map) + values[u] = wmma_helper(32, 8, 2, 2, 2, a_b_elem, a_b_elem, c_map) elif device == "AMD" and threads == 64: def a_elem(x, k, row, goff): return x[k%(dims[2]//4)][goff + (k//(dims[2]//4))*16 + row] def b_elem(x, col, k, goff): return a_elem(x, k, col, goff) # pylint: disable=arguments-out-of-order def c_map(lane, elem): return (lane%16, (lane//16)*4 + elem) - values[i] = wmma_helper(64, dims[2], len(src_values[0]), len(src_values[1]), len(src_values[2]), a_elem, b_elem, c_map) + values[u] = wmma_helper(64, dims[2], len(src_values[0]), len(src_values[1]), len(src_values[2]), a_elem, b_elem, c_map) elif device == "AMD" and len(src_values[0]) == 8: # RDNA4 def a_elem(x, k, row, goff): return x[k - [0, 4, 4, 8][k//4]][goff + row + [0, 16, 0, 16][k//4]] def b_elem(x, col, k, goff): return a_elem(x, k, col, goff) def c_map(lane, elem): return (lane%16, (lane//16)*8 + elem) - values[i] = wmma_helper(32, 16, 8, 8, 8, a_elem, b_elem, c_map) + values[u] = wmma_helper(32, 16, 8, 8, 8, a_elem, b_elem, c_map) elif device == "AMD": # A (16 elements on 32 threads): col major, lane 16-32 == lane 0-15 def a_elem(x, k, row, goff): @@ -154,7 +165,7 @@ def a_elem(x, k, row, goff): # B (16 elements on 32 threads): row major, lane 16-32 == lane 0-15 def b_elem(x, col, k, goff): return a_elem(x, k, col, goff) # pylint: disable=arguments-out-of-order def c_map(lane, elem): return (lane%16, lane//16+elem*2) # (i, j), C, D (8 elements on 32 threads): row major - values[i] = wmma_helper(32, 16, 16, 16, 8, a_elem, b_elem, c_map) + values[u] = wmma_helper(32, 16, 16, 16, 8, a_elem, b_elem, c_map) elif device == "CUDA": # (col, row) given (lane, elem) for C & D (4 elements on 32 threads); shared by all tc shapes with M=16 N=8 def c_map(lane, elem): return (elem%2 + (lane%4)*2, lane//4 + (elem//2)*8) @@ -162,24 +173,24 @@ def c_map(lane, elem): return (elem%2 + (lane%4)*2, lane//4 + (elem//2)*8) if dims == (8,16,16): def a_elem(x, k, row, goff): return x[k%2 + (row//8)*2 + (k//8)*4][goff + (k//2)%4 + (row%8)*4] def b_elem(x, col, k, goff): return x[k%2 + (k//8)*2][goff + (k//2)%4 + col*4] - values[i] = wmma_helper(32, 16, 8, 4, 4, a_elem, b_elem, c_map) + values[u] = wmma_helper(32, 16, 8, 4, 4, a_elem, b_elem, c_map) elif dims == (8,16,32): def a_elem(x, k, row, goff): return x[k%4 + (row//8)*4 + (k//16)*8][goff + (k//4)%4 + (row%8)*4] def b_elem(x, col, k, goff): return x[k%4 + (k//16)*4][goff + (k//4)%4 + col*4] - values[i] = wmma_helper(32, 32, 16, 8, 4, a_elem, b_elem, c_map) + values[u] = wmma_helper(32, 32, 16, 8, 4, a_elem, b_elem, c_map) elif dims == (8,16,8) and dtype_in == dtypes.half: def a_elem(x, k, row, goff): return x[k%2 + (row//8)*2][goff + k//2 + (row%8)*4] def b_elem(x, col, k, goff): return x[k%2][goff + k//2 + col*4] - values[i] = wmma_helper(32, 8, 4, 2, 4, a_elem, b_elem, c_map) + values[u] = wmma_helper(32, 8, 4, 2, 4, a_elem, b_elem, c_map) elif dims == (8,16,8) and dtype_in == dtypes.float: def a_elem(x, k, row, goff): return x[(k//4)*2 + row//8][goff + k%4 + (row%8)*4] def b_elem(x, col, k, goff): return x[k//4][goff + k%4 + col*4] - values[i] = wmma_helper(32, 8, 4, 2, 4, a_elem, b_elem, c_map) + values[u] = wmma_helper(32, 8, 4, 2, 4, a_elem, b_elem, c_map) - else: raise NotImplementedError(f"unimplemented tensor core {arg}") + else: raise NotImplementedError(f"unimplemented tensor core {u.arg}") elif device == "INTEL": # A (16 elements on 8 threads) def a_elem(x, k, row, goff): return x[k%2+row*2][goff+k//2] @@ -187,17 +198,17 @@ def a_elem(x, k, row, goff): return x[k%2+row*2][goff+k//2] def b_elem(x, col, k, goff): return x[k][goff+col] # C, D (8 elements on 8 threads) def c_map(lane, elem): return (lane, elem) - values[i] = wmma_helper(8, 16, 16, 16, 8, a_elem, b_elem, c_map) + values[u] = wmma_helper(8, 16, 16, 16, 8, a_elem, b_elem, c_map) elif device == "CPU": def elem(x, col, row, _): return x[col+row][0] # k is always 0 def c_map(lane, elem): return (elem%16, elem//16) - values[i] = wmma_helper(1, 1, 16, 16, 256, elem, elem, c_map) - else: raise NotImplementedError(f"unimplemented tensor core {arg}") - elif uop in GroupOp.ALU: - assert all_same([len(x) for x in src_values]), f"{[len(x) for x in src_values]} doesn't match on {uop}" - assert all_same([dtype] + src_dtypes) or uop in {*GroupOp.Comparison, Ops.WHERE}, f"dtype mismatch on {uop}" - values[i] = [exec_alu(uop, dtype, p) for p in zip(*src_values)] - assert i in values, (uop, dtype, srcs, arg) + values[u] = wmma_helper(1, 1, 16, 16, 256, elem, elem, c_map) + else: raise NotImplementedError(f"unimplemented tensor core {u.arg}") + elif u.op in GroupOp.ALU: + assert all_same([len(x) for x in src_values]), f"{[len(x) for x in src_values]} doesn't match on {u.op}" + assert all_same([u.dtype] + src_dtypes) or u.op in {*GroupOp.Comparison, Ops.WHERE}, f"dtype mismatch on {u.op}" + values[u] = [exec_alu(u.op, u.dtype, p) for p in zip(*src_values)] + assert u in values, u i += 1 return time.perf_counter() - st @@ -224,10 +235,9 @@ def __init__(self, target:Target): elif IMAGE and not target.arch: self.target = replace(target, arch="IMAGE_PITCH_ALIGNMENT=1") else: self.target = target - def render(self, uops:list[UOp]) -> str: - # the value of SPECIAL comes from local/global_size, not form its source - lops = [(u.op, u.dtype, [uops.index(v) for v in u.src if u.op is not Ops.SPECIAL], u.arg) for u in uops] - return base64.b64encode(pickle.dumps(lops)).decode() + def render(self, uops:list[UOp]) -> str: return base64.b64encode(pickle.dumps(uops)).decode() + + def supported_dtypes(self): return {d for d in super().supported_dtypes() if d != dtypes.half or sys.version_info >= (3, 12)} class PythonAllocator(Allocator['PythonDevice']): def _alloc(self, size, options): return memoryview(bytearray(size)) diff --git a/tinygrad/runtime/ops_qcom.py b/tinygrad/runtime/ops_qcom.py index 8f5fec2f41f1f..155449cc9a155 100644 --- a/tinygrad/runtime/ops_qcom.py +++ b/tinygrad/runtime/ops_qcom.py @@ -20,7 +20,7 @@ def dcache_flush(): from tinygrad.uop.ops import UOp, Ops, KernelInfo from tinygrad.codegen import to_program - buf, n = UOp(Ops.PARAM, dtypes.uint8.ptr(), arg=0), UOp(Ops.PARAM, dtypes.uint8.ptr(), arg=1) + buf, n = UOp.param(0, dtypes.uint8.ptr()), UOp.param(1, dtypes.uint8.ptr()) i = UOp.range(n.cast(dtypes.int), 0, dtype=dtypes.int) flush = UOp(Ops.CUSTOM, dtypes.void, (buf.cast(dtypes.ulong) + i.cast(dtypes.ulong) * UOp.const(dtypes.ulong, 64),), arg='__asm__ volatile("dc cvac, %0" :: "r"({0}) : "memory");') diff --git a/tinygrad/runtime/ops_webgpu.py b/tinygrad/runtime/ops_webgpu.py index 2e861a0dffc31..0157848ebdcbb 100644 --- a/tinygrad/runtime/ops_webgpu.py +++ b/tinygrad/runtime/ops_webgpu.py @@ -1,157 +1,120 @@ import functools, struct from tinygrad.device import Compiled, Allocator, BufferSpec from tinygrad.renderer.wgsl import WGSLRenderer -from tinygrad.helpers import round_up, suppress_finalizing +from tinygrad.helpers import round_up, suppress_finalizing, getenv, to_mv from tinygrad.runtime.autogen import webgpu from tinygrad.runtime.support import c -from typing import cast, List, Any, TypeAlias +from typing import Callable import ctypes -import os -WGPUDevPtr: TypeAlias = webgpu.WGPUDevice -WGPUBufPtr: TypeAlias = webgpu.WGPUBuffer - -backend_types = {v: k for k, v in webgpu.enum_WGPUBackendType.items() } - -instance = webgpu.wgpuCreateInstance(webgpu.WGPUInstanceDescriptor(features = webgpu.WGPUInstanceFeatures(timedWaitAnyEnable = True))) - -def to_c_string(_str:str) -> ctypes.Array: return ctypes.create_string_buffer(_str.encode('utf-8')) - -def from_wgpu_str(string_view:webgpu.struct_WGPUStringView) -> str: return ctypes.string_at(string_view.data, string_view.length).decode("utf-8") - -def to_wgpu_str(_str:str) -> webgpu.struct_WGPUStringView: - return webgpu.WGPUStringView(data=ctypes.cast(ctypes.pointer(to_c_string(_str)), ctypes.POINTER(ctypes.c_char)), length=len(_str)) - -def _wait(future:webgpu.struct_WGPUFuture): - assert webgpu.wgpuInstanceWaitAny(instance, 1, webgpu.WGPUFutureWaitInfo(future=future), 2**64-1) == webgpu.WGPUWaitStatus_Success, "Future failed" - -def write_buffer(device:WGPUDevPtr, buf:WGPUBufPtr, offset:int, src:memoryview|bytearray|bytes): - src = bytearray(src) - webgpu.wgpuQueueWriteBuffer(webgpu.wgpuDeviceGetQueue(device), buf, offset, (ctypes.c_uint8 * len(src)).from_buffer(src), len(src)) - -def _run(async_fun, cb_info_type, cb_type, status_enum, res_idx:int|None, msg_idx:int|None, *params): - result: List[Any] = [] - - def cb(*params): - result[:] = params - if msg_idx: result[msg_idx] = from_wgpu_str(result[msg_idx]) - - cb_info = cb_info_type(mode=webgpu.WGPUCallbackMode_WaitAnyOnly, callback=cb_type(cb)) - _wait(async_fun(*params, cb_info)) - - if result[0] != 1: raise RuntimeError(f"[{status_enum.get(result[0]) if status_enum else 'ERROR'}]{result[msg_idx] if msg_idx else ''}") - return result[res_idx] if res_idx else None - -def copy_buffer_to_buffer(dev:WGPUDevPtr, src:WGPUBufPtr, src_offset:int, dst:WGPUBufPtr, dst_offset:int, size:int): - encoder = webgpu.wgpuDeviceCreateCommandEncoder(dev, webgpu.WGPUCommandEncoderDescriptor()) - webgpu.wgpuCommandEncoderCopyBufferToBuffer(encoder, src, src_offset, dst, dst_offset, size) - cb = webgpu.wgpuCommandEncoderFinish(encoder, webgpu.WGPUCommandBufferDescriptor()) - webgpu.wgpuQueueSubmit(webgpu.wgpuDeviceGetQueue(dev), 1, (webgpu.WGPUCommandBuffer*1)(cb)) - webgpu.wgpuCommandBufferRelease(cb) - webgpu.wgpuCommandEncoderRelease(encoder) - -def read_buffer(dev:WGPUDevPtr, buf:WGPUBufPtr) -> memoryview: - size = webgpu.wgpuBufferGetSize(buf) - tmp_buffer = webgpu.wgpuDeviceCreateBuffer(dev, webgpu.WGPUBufferDescriptor(size=size, - usage=webgpu.WGPUBufferUsage_CopyDst | webgpu.WGPUBufferUsage_MapRead, mappedAtCreation=False)) - copy_buffer_to_buffer(dev, buf, 0, tmp_buffer, 0, size) - _run(webgpu.wgpuBufferMapAsync2, webgpu.WGPUBufferMapCallbackInfo2, webgpu.WGPUBufferMapCallback2, webgpu.WGPUBufferMapAsyncStatus, None, 0, - tmp_buffer, webgpu.WGPUMapMode_Read, 0, size) - void_ptr = ctypes.cast(webgpu.wgpuBufferGetConstMappedRange(tmp_buffer, 0, size), ctypes.c_void_p) - buf_copy = bytearray((ctypes.c_uint8 * size).from_address(void_ptr.value)) - webgpu.wgpuBufferUnmap(tmp_buffer) - webgpu.wgpuBufferDestroy(tmp_buffer) - return memoryview(buf_copy).cast("B") - -def pop_error(device:WGPUDevPtr) -> str: - return _run(webgpu.wgpuDevicePopErrorScopeF, webgpu.WGPUPopErrorScopeCallbackInfo, webgpu.WGPUPopErrorScopeCallback, None, 2, 2, device) - -def create_uniform(wgpu_device:WGPUDevPtr, val:int|float) -> WGPUBufPtr: - buf = webgpu.wgpuDeviceCreateBuffer(wgpu_device, - webgpu.WGPUBufferDescriptor(size=4, usage=webgpu.WGPUBufferUsage_Uniform | webgpu.WGPUBufferUsage_CopyDst)) - write_buffer(wgpu_device, buf, 0, val.to_bytes(4, "little") if isinstance(val, int) else struct.pack(' str: return ctypes.string_at(string_view.data, string_view.length).decode() +def to_wgpu_str(_str:str) -> webgpu.WGPUStringView: return webgpu.WGPUStringView(data=ctypes.create_string_buffer(_str.encode()), length=len(_str)) + +# gets a memoryview from a buffer, which is assumed to have MAP_READ (see _readable_buffer) +def buf_to_mv(buf:webgpu.WGPUBuffer) -> memoryview: + BufferMapAsync(buf, webgpu.WGPUMapMode_Read, 0, size:=webgpu.wgpuBufferGetSize(buf)) + return to_mv(webgpu.wgpuBufferGetConstMappedRange(buf, 0, size), size) + +# turns a webgpu function returning a future into python-synchronous function +# the new function handles the status code and optional error message, returning the other callback arguments +def synchronous(status_enum:dict[int, str], has_emsg:bool=False): + def wrap(fn:Callable[..., webgpu.WGPUFuture]) -> Callable: + @functools.wraps(fn) + def wrapper(*args): + status, payload, emsg = 0, [], None + + @next(ty for nm, ty, *_ in fn.argtypes[-1]._real_fields_ if nm == "callback") # type: ignore + def cb(s:int, *args): + nonlocal status, payload, emsg + # the last two arguments are "userdata1" and "userdata2", which we drop + # we must process wgpu strings in this callback, as they will be freed after we return + status, (*payload, emsg) = s, [from_wgpu_str(a) if type(a) is webgpu.WGPUStringView else a for a in args[:-2]] + ([] if has_emsg else [None]) + + future = fn(*args, fn.argtypes[-1](mode=webgpu.WGPUCallbackMode_WaitAnyOnly, callback=cb)) # type: ignore + if (future_status:=webgpu.wgpuInstanceWaitAny(instance, 1, webgpu.WGPUFutureWaitInfo(future), 2**64-1)) != webgpu.WGPUWaitStatus_Success: + raise RuntimeError(f"error while waiting for future ({fn.__name__}): {webgpu.enum_WGPUWaitStatus.get(future_status)}") + + if status != 1: raise RuntimeError(f"[{status_enum.get(status)}]{emsg or ''}") + return payload if len(payload) > 1 else payload[0] if len(payload) == 1 else None + return wrapper + return wrap + +BufferMapAsync = synchronous(webgpu.enum_WGPUBufferMapAsyncStatus, True)(webgpu.wgpuBufferMapAsync2) +DevicePopErrorScope = synchronous(webgpu.enum_WGPUPopErrorScopeStatus)(webgpu.wgpuDevicePopErrorScope2) +DeviceCreateComputePipeline = synchronous(webgpu.enum_WGPUCreatePipelineAsyncStatus, True)(webgpu.wgpuDeviceCreateComputePipelineAsync2) +InstanceRequestAdapter = synchronous(webgpu.enum_WGPURequestAdapterStatus, True)(webgpu.wgpuInstanceRequestAdapter2) +AdapterRequestDevice = synchronous(webgpu.enum_WGPURequestDeviceStatus, True)(webgpu.wgpuAdapterRequestDevice2) +QueueOnSubmittedWorkDone = synchronous(webgpu.enum_WGPUQueueWorkDoneStatus)(webgpu.wgpuQueueOnSubmittedWorkDone2) class WebGPUProgram: - def __init__(self, dev:tuple[WGPUDevPtr, bool], name:str, lib:bytes, **kwargs): - (self.dev, self.timestamp_supported) = dev + def __init__(self, dev:'WebGpuDevice', name:str, lib:bytes, **kwargs): + self.dev, self.name = dev, to_wgpu_str(name) # Creating shader module shader = webgpu.WGPUShaderModuleWGSLDescriptor(code=to_wgpu_str(lib.decode()), - chain=webgpu.WGPUChainedStruct(sType=webgpu.WGPUSType_ShaderSourceWGSL)) - module = webgpu.WGPUShaderModuleDescriptor() - module.nextInChain = ctypes.cast(ctypes.pointer(shader), c.POINTER[webgpu.struct_WGPUChainedStruct]) + chain=webgpu.WGPUChainedStruct(sType=webgpu.WGPUSType_ShaderSourceWGSL)) + module = webgpu.WGPUShaderModuleDescriptor(nextInChain=ctypes.cast(ctypes.pointer(shader), ctypes.POINTER(webgpu.struct_WGPUChainedStruct))) # Check compiler error - webgpu.wgpuDevicePushErrorScope(self.dev, webgpu.WGPUErrorFilter_Validation) - shader_module = webgpu.wgpuDeviceCreateShaderModule(self.dev, module) + webgpu.wgpuDevicePushErrorScope(self.dev.device_res, webgpu.WGPUErrorFilter_Validation) + self.prg = webgpu.wgpuDeviceCreateShaderModule(self.dev.device_res, module) + if err := self.dev.pop_error(): raise RuntimeError(f"Shader compilation failed: {err}") - if err := pop_error(self.dev): raise RuntimeError(f"Shader compilation failed: {err}") + @suppress_finalizing + def __del__(self): webgpu.wgpuShaderModuleRelease(self.prg) - self.name, self.lib, self.prg = name, lib, shader_module - def __call__(self, *bufs:WGPUBufPtr, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), + def __call__(self, *bufs:webgpu.WGPUBuffer, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False, **kw) -> float|None: - wait = wait and self.timestamp_supported - tmp_bufs = [*bufs] - buf_patch = False - - # WebGPU does not allow using the same buffer for input and output - for i in range(1, len(bufs)): - if ctypes.addressof(bufs[i]) == ctypes.addressof(bufs[0]): - tmp_bufs[0] = webgpu.wgpuDeviceCreateBuffer(self.dev, - webgpu.WGPUBufferDescriptor(size=webgpu.wgpuBufferGetSize(bufs[0]), usage=webgpu.wgpuBufferGetUsage(bufs[0]))) - buf_patch = True + wait = wait and webgpu.WGPUFeatureName_TimestampQuery in self.dev.features # Creating bind group layout - binding_layouts = [webgpu.WGPUBindGroupLayoutEntry(binding=0, visibility= webgpu.WGPUShaderStage_Compute, - buffer=webgpu.WGPUBufferBindingLayout(type=webgpu.WGPUBufferBindingType_Uniform))] - binding_layouts += [webgpu.WGPUBindGroupLayoutEntry(binding=i+1, visibility=webgpu.WGPUShaderStage_Compute, - buffer=webgpu.WGPUBufferBindingLayout(type=webgpu.WGPUBufferBindingType_Uniform if i >= len(tmp_bufs) - else webgpu.WGPUBufferBindingType_Storage)) for i in range(len(tmp_bufs)+len(vals))] + def bgl_entry(n:int, ty:str): + return webgpu.WGPUBindGroupLayoutEntry(binding=n, visibility=webgpu.WGPUShaderStage_Compute, + buffer=webgpu.WGPUBufferBindingLayout(type=getattr(webgpu, f'WGPUBufferBindingType_{ty}'))) + bind_entries = (webgpu.WGPUBindGroupLayoutEntry * (1+len(bufs)+len(vals)))( + bgl_entry(0, 'Uniform'), *(bgl_entry(i+1, 'Uniform' if i >= len(bufs) else 'Storage') for i in range(len(bufs)+len(vals)))) - bl_arr_type = webgpu.WGPUBindGroupLayoutEntry * len(binding_layouts) - webgpu.wgpuDevicePushErrorScope(self.dev, webgpu.WGPUErrorFilter_Validation) - bind_group_layouts = [webgpu.wgpuDeviceCreateBindGroupLayout(self.dev, webgpu.WGPUBindGroupLayoutDescriptor( - entryCount=len(binding_layouts), entries=ctypes.cast(bl_arr_type(*binding_layouts), ctypes.POINTER(webgpu.WGPUBindGroupLayoutEntry))))] + webgpu.wgpuDevicePushErrorScope(self.dev.device_res, webgpu.WGPUErrorFilter_Validation) + bind_layout = webgpu.wgpuDeviceCreateBindGroupLayout(self.dev.device_res, + webgpu.WGPUBindGroupLayoutDescriptor(entryCount=len(bind_entries), entries=bind_entries)) - if bg_layout_err := pop_error(self.dev): raise RuntimeError(f"Error creating bind group layout: {bg_layout_err}") + if err := self.dev.pop_error(): raise RuntimeError(f"Error creating bind group layout: {err}") # Creating pipeline layout - pipeline_layout_desc = webgpu.WGPUPipelineLayoutDescriptor(bindGroupLayoutCount=len(bind_group_layouts), - bindGroupLayouts = (webgpu.WGPUBindGroupLayout * len(bind_group_layouts))(*bind_group_layouts)) - - webgpu.wgpuDevicePushErrorScope(self.dev, webgpu.WGPUErrorFilter_Validation) - pipeline_layout = webgpu.wgpuDeviceCreatePipelineLayout(self.dev, pipeline_layout_desc) + pipeline_layout_desc = webgpu.WGPUPipelineLayoutDescriptor(bindGroupLayoutCount=1, bindGroupLayouts=(webgpu.WGPUBindGroupLayout*1)(bind_layout)) - if pipe_err := pop_error(self.dev): raise RuntimeError(f"Error creating pipeline layout: {pipe_err}") + webgpu.wgpuDevicePushErrorScope(self.dev.device_res, webgpu.WGPUErrorFilter_Validation) + pipeline_layout = webgpu.wgpuDeviceCreatePipelineLayout(self.dev.device_res, pipeline_layout_desc) + if err := self.dev.pop_error(): raise RuntimeError(f"Error creating pipeline layout: {err}") # Creating bind group - bindings = [webgpu.WGPUBindGroupEntry(binding=0, buffer=create_uniform(self.dev, float('inf')), offset=0, size=4)] - bindings += [webgpu.WGPUBindGroupEntry(binding=i+1, buffer=create_uniform(self.dev, cast(int, x)) if i >= len(tmp_bufs) else x, offset=0, - size=4 if i >= len(tmp_bufs) else webgpu.wgpuBufferGetSize(x)) for i,x in enumerate(tuple(tmp_bufs)+vals)] + def bg_entry(n:int, x:webgpu.WGPUBuffer|int|float): + buf = x if isinstance(x, webgpu.WGPUBuffer) else self.dev.create_uniform(x) + return webgpu.WGPUBindGroupEntry(binding=n, buffer=buf, offset=0, size=webgpu.wgpuBufferGetSize(buf)) + bindings = (webgpu.WGPUBindGroupEntry * (1+len(bufs)+len(vals)))(bg_entry(0, float('inf')), *(bg_entry(i+1, x) for i,x in enumerate(bufs+vals))) - bg_arr_type = webgpu.WGPUBindGroupEntry * len(bindings) - bind_group_desc = webgpu.WGPUBindGroupDescriptor(layout=bind_group_layouts[0], entryCount=len(bindings), entries=bg_arr_type(*bindings)) - webgpu.wgpuDevicePushErrorScope(self.dev, webgpu.WGPUErrorFilter_Validation) - bind_group = webgpu.wgpuDeviceCreateBindGroup(self.dev, bind_group_desc) - - if bind_err := pop_error(self.dev): raise RuntimeError(f"Error creating bind group: {bind_err}") + bind_group_desc = webgpu.WGPUBindGroupDescriptor(layout=bind_layout, entryCount=len(bindings), entries=bindings) + webgpu.wgpuDevicePushErrorScope(self.dev.device_res, webgpu.WGPUErrorFilter_Validation) + bind_group = webgpu.wgpuDeviceCreateBindGroup(self.dev.device_res, bind_group_desc) + if err := self.dev.pop_error(): raise RuntimeError(f"Error creating bind group: {err}") # Creating compute pipeline compute_desc = webgpu.WGPUComputePipelineDescriptor(layout=pipeline_layout, - compute=webgpu.WGPUComputeState(module=self.prg, entryPoint=to_wgpu_str(self.name))) - pipeline_result = _run(webgpu.wgpuDeviceCreateComputePipelineAsync2, webgpu.WGPUCreateComputePipelineAsyncCallbackInfo2, - webgpu.WGPUCreateComputePipelineAsyncCallback2, webgpu.WGPUCreatePipelineAsyncStatus, 1, None, self.dev, compute_desc) + compute=webgpu.WGPUComputeState(module=self.prg, entryPoint=self.name)) + pipeline_result = DeviceCreateComputePipeline(self.dev.device_res, compute_desc) - command_encoder = webgpu.wgpuDeviceCreateCommandEncoder(self.dev, webgpu.WGPUCommandEncoderDescriptor()) + command_encoder = webgpu.wgpuDeviceCreateCommandEncoder(self.dev.device_res, webgpu.WGPUCommandEncoderDescriptor()) comp_pass_desc = webgpu.WGPUComputePassDescriptor() if wait: - query_set = webgpu.wgpuDeviceCreateQuerySet(self.dev, webgpu.WGPUQuerySetDescriptor(type=webgpu.WGPUQueryType_Timestamp, count=2)) - query_buf = webgpu.wgpuDeviceCreateBuffer(self.dev, - webgpu.WGPUBufferDescriptor(size=16, usage=webgpu.WGPUBufferUsage_QueryResolve | webgpu.WGPUBufferUsage_CopySrc)) - comp_pass_desc.timestampWrites = c.pointer(webgpu.WGPUComputePassTimestampWrites( - querySet=query_set, beginningOfPassWriteIndex=0, endOfPassWriteIndex=1)) + query_set = webgpu.wgpuDeviceCreateQuerySet(self.dev.device_res, webgpu.WGPUQuerySetDescriptor(type=webgpu.WGPUQueryType_Timestamp, count=2)) + query_buf = webgpu.wgpuDeviceCreateBuffer( + self.dev.device_res, webgpu.WGPUBufferDescriptor(size=16, usage=webgpu.WGPUBufferUsage_QueryResolve | webgpu.WGPUBufferUsage_CopySrc)) + comp_pass_desc.timestampWrites = c.pointer(webgpu.WGPUComputePassTimestampWrites(querySet=query_set, beginningOfPassWriteIndex=0, + endOfPassWriteIndex=1)) # Begin compute pass compute_pass = webgpu.wgpuCommandEncoderBeginComputePass(command_encoder, comp_pass_desc) @@ -163,63 +126,96 @@ def __call__(self, *bufs:WGPUBufPtr, global_size:tuple[int,int,int]=(1,1,1), loc if wait: webgpu.wgpuCommandEncoderResolveQuerySet(command_encoder, query_set, 0, 2, query_buf, 0) cmd_buf = webgpu.wgpuCommandEncoderFinish(command_encoder, webgpu.WGPUCommandBufferDescriptor()) - webgpu.wgpuQueueSubmit(webgpu.wgpuDeviceGetQueue(self.dev), 1, (webgpu.WGPUCommandBuffer*1)(cmd_buf)) + webgpu.wgpuQueueSubmit(self.dev.queue, 1, (webgpu.WGPUCommandBuffer*1)(cmd_buf)) - if buf_patch: - copy_buffer_to_buffer(self.dev, tmp_bufs[0], 0, bufs[0], 0, webgpu.wgpuBufferGetSize(bufs[0])) - webgpu.wgpuBufferDestroy(tmp_bufs[0]) + # release created objects + webgpu.wgpuBindGroupLayoutRelease(bind_layout) + webgpu.wgpuPipelineLayoutRelease(pipeline_layout) + webgpu.wgpuBindGroupRelease(bind_group) + webgpu.wgpuComputePipelineRelease(pipeline_result) + webgpu.wgpuCommandEncoderRelease(command_encoder) + webgpu.wgpuComputePassEncoderRelease(compute_pass) + webgpu.wgpuCommandBufferRelease(cmd_buf) if wait: - time = ((timestamps:=read_buffer(self.dev, query_buf).cast("Q").tolist())[1] - timestamps[0]) / 1e9 - webgpu.wgpuBufferDestroy(query_buf) + time = ((timestamps:=buf_to_mv(tmp_buf:=self.dev._readable_buffer(query_buf)).cast("Q").tolist())[1] - timestamps[0]) / 1e9 + self.dev.free(query_buf) + self.dev.free(tmp_buf) webgpu.wgpuQuerySetDestroy(query_set) + webgpu.wgpuQuerySetRelease(query_set) return time return None class WebGpuAllocator(Allocator['WebGpuDevice']): - def _alloc(self, size:int, options:BufferSpec) -> WGPUBufPtr: + def _alloc(self, size:int, options:BufferSpec) -> webgpu.WGPUBuffer: # WebGPU buffers have to be 4-byte aligned return webgpu.wgpuDeviceCreateBuffer(self.dev.device_res, webgpu.WGPUBufferDescriptor(size=round_up(size, 4), usage=webgpu.WGPUBufferUsage_Storage | webgpu.WGPUBufferUsage_CopyDst | webgpu.WGPUBufferUsage_CopySrc)) - def _copyin(self, dest:WGPUBufPtr, src:memoryview): + def _copyin(self, dest:webgpu.WGPUBuffer, src:memoryview): if src.nbytes % 4: padded_src = bytearray(round_up(src.nbytes, 4)) padded_src[:src.nbytes] = src - write_buffer(self.dev.device_res, dest, 0, padded_src if src.nbytes % 4 else src) - def _copyout(self, dest:memoryview, src:WGPUBufPtr): - buffer_data = read_buffer(self.dev.device_res, src) - dest[:] = buffer_data[:dest.nbytes] if webgpu.wgpuBufferGetSize(src) > dest.nbytes else buffer_data - @suppress_finalizing - def _free(self, opaque:WGPUBufPtr, options:BufferSpec): webgpu.wgpuBufferDestroy(opaque) + self.dev.write_buffer(dest, padded_src if src.nbytes % 4 else src) + def _copyout(self, dest:memoryview, src:webgpu.WGPUBuffer): + dest[:] = buf_to_mv(tmp_buf:=self.dev._readable_buffer(src))[:dest.nbytes] + self.dev.free(tmp_buf) + + def _free(self, opaque:webgpu.WGPUBuffer, options:BufferSpec): self.dev.free(opaque) class WebGpuDevice(Compiled): def __init__(self, device:str): # Requesting an adapter - adapter_res = _run(webgpu.wgpuInstanceRequestAdapterF, webgpu.WGPURequestAdapterCallbackInfo, webgpu.WGPURequestAdapterCallback, - webgpu.WGPURequestAdapterStatus, 1, 2, instance, webgpu.WGPURequestAdapterOptions(powerPreference=webgpu.WGPUPowerPreference_HighPerformance, - backendType=backend_types.get(os.getenv("WEBGPU_BACKEND", ""), 0))) + adapter_res = InstanceRequestAdapter(instance, webgpu.WGPURequestAdapterOptions( + powerPreference=webgpu.WGPUPowerPreference_HighPerformance, backendType=backend_types.get(getenv("WEBGPU_BACKEND", ""), 0))) # Get supported features - supported_features = webgpu.WGPUSupportedFeatures() - webgpu.wgpuAdapterGetFeatures(adapter_res, supported_features) - supported = [supported_features.features[i] for i in range(supported_features.featureCount)] - features = [feat for feat in [webgpu.WGPUFeatureName_TimestampQuery, webgpu.WGPUFeatureName_ShaderF16] if feat in supported] - dev_desc = webgpu.WGPUDeviceDescriptor(requiredFeatureCount=len(features), - requiredFeatures=c.Array(webgpu.WGPUFeatureName, len(features))(*features)) # type: ignore + webgpu.wgpuAdapterGetFeatures(adapter_res, supported_features:=webgpu.WGPUSupportedFeatures()) + self.features = [feat for i in range(supported_features.featureCount) + if (feat:=supported_features.features[i]) in [webgpu.WGPUFeatureName_TimestampQuery, webgpu.WGPUFeatureName_ShaderF16]] + webgpu.wgpuSupportedFeaturesFreeMembers(supported_features) + dev_desc = webgpu.WGPUDeviceDescriptor(requiredFeatureCount=len(self.features), + requiredFeatures=(webgpu.WGPUFeatureName * len(self.features))(*self.features)) # Limits - supported_limits = webgpu.WGPUSupportedLimits() - webgpu.wgpuAdapterGetLimits(adapter_res, ctypes.cast(ctypes.pointer(supported_limits),ctypes.POINTER(webgpu.struct_WGPUSupportedLimits))) - limits = webgpu.WGPURequiredLimits(limits=supported_limits.limits) - dev_desc.requiredLimits = c.pointer(limits) + webgpu.wgpuAdapterGetLimits(adapter_res, supported_limits:=webgpu.WGPUSupportedLimits()) + dev_desc.requiredLimits = c.pointer(webgpu.WGPURequiredLimits(limits=supported_limits.limits)) # Requesting a device - self.device_res = _run(webgpu.wgpuAdapterRequestDeviceF, webgpu.WGPURequestDeviceCallbackInfo, webgpu.WGPURequestDeviceCallback, - webgpu.WGPURequestDeviceStatus, 1, 2, adapter_res, dev_desc) + self.device_res = AdapterRequestDevice(adapter_res, dev_desc) + self.queue = webgpu.wgpuDeviceGetQueue(self.device_res) - super().__init__(device, WebGpuAllocator(self), [WGSLRenderer], - functools.partial(WebGPUProgram, (self.device_res, webgpu.WGPUFeatureName_TimestampQuery in supported))) + webgpu.wgpuAdapterRelease(adapter_res) - def synchronize(self): - _run(webgpu.wgpuQueueOnSubmittedWorkDone2, webgpu.WGPUQueueWorkDoneCallbackInfo2, webgpu.WGPUQueueWorkDoneCallback2, - webgpu.WGPUQueueWorkDoneStatus, None, None, webgpu.wgpuDeviceGetQueue(self.device_res)) + super().__init__(device, WebGpuAllocator(self), [WGSLRenderer], functools.partial(WebGPUProgram, self), + arch="shader-f16" * (webgpu.WGPUFeatureName_ShaderF16 in self.features)) + + def synchronize(self): QueueOnSubmittedWorkDone(self.queue) + + @suppress_finalizing + def free(self, buf:webgpu.WGPUBuffer): + if webgpu.wgpuBufferGetMapState(buf) == webgpu.WGPUBufferMapState_Mapped: webgpu.wgpuBufferUnmap(buf) + webgpu.wgpuBufferDestroy(buf) + webgpu.wgpuBufferRelease(buf) + + def pop_error(self) -> str: return DevicePopErrorScope(self.device_res)[1] + def create_uniform(self, val:int|float) -> webgpu.WGPUBuffer: + buf = webgpu.wgpuDeviceCreateBuffer(self.device_res, + webgpu.WGPUBufferDescriptor(size=4, usage=webgpu.WGPUBufferUsage_Uniform | webgpu.WGPUBufferUsage_CopyDst)) + self.write_buffer(buf, val.to_bytes(4, "little") if isinstance(val, int) else struct.pack(' webgpu.WGPUBuffer: + size = webgpu.wgpuBufferGetSize(buf) + ret = webgpu.wgpuDeviceCreateBuffer(self.device_res, + webgpu.WGPUBufferDescriptor(size=size, usage=webgpu.WGPUBufferUsage_CopyDst | webgpu.WGPUBufferUsage_MapRead, mappedAtCreation=False)) + + # copy_buffer_to_buffer + encoder = webgpu.wgpuDeviceCreateCommandEncoder(self.device_res, webgpu.WGPUCommandEncoderDescriptor()) + webgpu.wgpuCommandEncoderCopyBufferToBuffer(encoder, buf, 0, ret, 0, size) + cmd_buf = webgpu.wgpuCommandEncoderFinish(encoder, webgpu.WGPUCommandBufferDescriptor()) + webgpu.wgpuQueueSubmit(self.queue, 1, (webgpu.WGPUCommandBuffer*1)(cmd_buf)) + webgpu.wgpuCommandBufferRelease(cmd_buf) + webgpu.wgpuCommandEncoderRelease(encoder) + + return ret + def write_buffer(self, buf:webgpu.WGPUBuffer, src:memoryview|bytearray|bytes): + webgpu.wgpuQueueWriteBuffer(self.queue, buf, 0, (ctypes.c_uint8 * len(src)).from_buffer_copy(src), len(src)) diff --git a/tinygrad/runtime/support/am/ip.py b/tinygrad/runtime/support/am/ip.py index 93688b9daca99..bbe1804230c5f 100644 --- a/tinygrad/runtime/support/am/ip.py +++ b/tinygrad/runtime/support/am/ip.py @@ -304,7 +304,7 @@ def fini_hw(self): self._dequeue_hqds() def reset_mec(self): self._dequeue_hqds() - if self.adev.ip_ver[am.GC_HWIP] < (10,0,0): # gfx10+ uses mec_pipe0_reset + if self.adev.ip_ver[am.GC_HWIP] < (12,0,0): # gfx12+ uses mec_pipe0_reset for xcc in range(self.xccs): self.adev.regGRBM_SOFT_RESET.write(soft_reset_cp=1, soft_reset_cpc=1, inst=xcc) time.sleep(0.05) for xcc in range(self.xccs): self.adev.regGRBM_SOFT_RESET.write(0x0, inst=xcc) diff --git a/tinygrad/runtime/support/amd.py b/tinygrad/runtime/support/amd.py index 27d3e1049fd48..cc7e07009881a 100644 --- a/tinygrad/runtime/support/amd.py +++ b/tinygrad/runtime/support/amd.py @@ -15,7 +15,7 @@ def fields_mask(self, *names) -> int: @dataclass class AMDIP: - name:str; version:tuple[int, ...]; bases:dict[int, tuple[int, ...]] # noqa: E702 + name:str; version:tuple[int, int, int]; bases:dict[int, tuple[int, ...]] # noqa: E702 @functools.cached_property def regs(self): return import_asic_regs(self.name, self.version, cls=functools.partial(AMDReg, bases=self.bases)) @@ -26,9 +26,11 @@ def __getattr__(self, name:str): raise AttributeError(f"{self.name.upper()} has no register {name}") # load the greatest module with matching major version that's less than or equal to the target version -# this is not universally correct, see below for an example, but appears reliable for recent gpus +# this is not universally correct, see below for an example, but appears reliable for most recent gpus # https://github.com/torvalds/linux/blob/9207d47f966be9f4d52e7e0119ac2b7a7e366f3e/drivers/gpu/drm/amd/amdgpu/amdgpu_discovery.c#L3163 -def import_module(name:str, target:tuple[int, ...], submod=""): +def import_module(name:str, target:tuple[int, int, int], submod=""): + # version overrides + target = {("smu", (13, 0, 7)): (13, 0, 0)}.get((name, target), target) mod = getattr(tinygrad.runtime.autogen.am, submod) if submod else tinygrad.runtime.autogen.am if (children:=[c for c in mod.__all__ if c.startswith(name) and (v:=tuple(map(int, c.split('_')[1:])))[0] == target[0] and v <= target]): return getattr(mod, children[-1]) @@ -41,5 +43,5 @@ def import_pmc(ip) -> dict[str, tuple[str, int]]: # NOTE: precise arch for mi300+, generic for others, since rocm headers lack some archs return {k:x for k,v in pmc.counters.items() if (x:=v.get(f"gfx{ip[0]}{ip[1]:x}{ip[2]:x}" if ip[0] == 9 else f"gfx{ip[0]}", None)) is not None} -def import_asic_regs(prefix:str, version:tuple[int, ...], cls=AMDReg) -> dict[str, AMDReg]: +def import_asic_regs(prefix:str, version:tuple[int, int, int], cls=AMDReg) -> dict[str, AMDReg]: return {reg:cls(name=reg, offset=off, segment=seg, fields=fields) for reg,(off,seg,fields) in import_module(prefix, version, submod="regs").items()} diff --git a/tinygrad/runtime/support/c.py b/tinygrad/runtime/support/c.py index d97e491fb95ee..42a3f34b0184d 100644 --- a/tinygrad/runtime/support/c.py +++ b/tinygrad/runtime/support/c.py @@ -132,6 +132,7 @@ def wrapper(*args): nonlocal cfunc if cfunc is None: (cfunc:=getattr(self, fn.__name__)).argtypes, cfunc.restype = argtypes, restype return cfunc(*args) + wrapper.restype, wrapper.argtypes = restype, argtypes # type: ignore return wrapper return wrap diff --git a/tinygrad/runtime/support/compiler_cpu.py b/tinygrad/runtime/support/compiler_cpu.py index a2db204ba17e0..2fc5ede229398 100644 --- a/tinygrad/runtime/support/compiler_cpu.py +++ b/tinygrad/runtime/support/compiler_cpu.py @@ -4,7 +4,7 @@ from tinygrad.runtime.support.elf import jit_loader from tinygrad.runtime.autogen import llvm -class ClangJITCompiler(Compiler): +class ClangCompiler(Compiler): def __init__(self, arch:list[str], cachekey="compile_clang_jit"): assert len(arch) >= 2, f"invalid arch string: {','.join(arch)!r}, expected ',,[]' (eg. 'x86_64,znver2')" self.arch, cpu, *feats = arch @@ -98,7 +98,12 @@ def __init__(self, arch:list[str], cache_key=None): if cpu == "native": cpu = ctypes.string_at(llvm.LLVMGetHostCPUName()).decode() featstr = (featstr + "," if featstr else "") + ctypes.string_at(llvm.LLVMGetHostCPUFeatures()).decode() - # +reserve-x18 here does the same thing as -ffixed-x18 in ClangJITCompiler, see comments there for why it's needed on arm osx + # +reserve-x18 here does the same thing as -ffixed-x18 in ClangCompiler, see comments there for why it's needed on arm osx super().__init__(self.arch, cpu, ('+reserve-x18,' if self.arch == "arm64" else '') + featstr, cache_key) def disassemble(self, lib:bytes): capstone_flatdump(lib, self.arch) + +class X86Compiler(Compiler): + def __init__(self): super().__init__(None) + def compile(self, src:str) -> bytes: return bytes.fromhex(src) + def disassemble(self, lib:bytes): return capstone_flatdump(lib, "x86_64") diff --git a/tinygrad/runtime/support/nv/ip.py b/tinygrad/runtime/support/nv/ip.py index 47bc610a312b2..f943ae1c08973 100644 --- a/tinygrad/runtime/support/nv/ip.py +++ b/tinygrad/runtime/support/nv/ip.py @@ -163,9 +163,9 @@ def __patch(cmd_id, cmd): patched_image[(cmd_off:=self.desc_v3.IMEMLoadSize+dmem.cmd_in_buffer_offset) : cmd_off+len(cmd)] = cmd patched_image[(sig_off:=self.desc_v3.IMEMLoadSize+self.desc_v3.PKCDataOffset) : sig_off+0x180] = signature[-0x180:] - return self.nvdev._alloc_sysmem(len(patched_image), contiguous=True, data=patched_image) + return self.nvdev._alloc_boot_mem(len(patched_image), data=patched_image, sysmem=False) - _, self.frts_image_sysmem = __patch(0x15, bytes(frts_cmd)) + _, self.frts_image_paddr, _ = __patch(0x15, bytes(frts_cmd)) def prep_booter(self): sha = {"ga102":"4497e3eff7e95c774b8a569d17b27c08c9650158d10b229d2be81cdcad9a085b", @@ -179,14 +179,14 @@ def prep_booter(self): (patched_image:=bytearray(b[h.data_offset:h.data_offset + h.data_size]))[patch_loc:patch_loc+sig_len] = sig - _, self.booter_image_sysmem = self.nvdev._alloc_sysmem(len(patched_image), contiguous=True, data=patched_image) + _, self.booter_image_paddr, _ = self.nvdev._alloc_boot_mem(len(patched_image), data=patched_image, sysmem=False) self.booter_data_off, self.booter_data_sz, self.booter_code_off, self.booter_code_sz = lh.os_data_offset, lh.os_data_size, app.offset, app.size def init_hw(self): self.falcon, self.sec2 = 0x00110000, 0x00840000 self.reset(self.falcon) - self.execute_hs(self.falcon, self.frts_image_sysmem[0], code_off=0x0, data_off=self.desc_v3.IMEMLoadSize, + self.execute_hs(self.falcon, self.frts_image_paddr, code_off=0x0, data_off=self.desc_v3.IMEMLoadSize, imemPa=self.desc_v3.IMEMPhysBase, imemVa=self.desc_v3.IMEMVirtBase, imemSz=self.desc_v3.IMEMLoadSize, dmemPa=self.desc_v3.DMEMPhysBase, dmemVa=0x0, dmemSz=self.desc_v3.DMEMLoadSize, pkc_off=self.desc_v3.PKCDataOffset, engid=self.desc_v3.EngineIdMask, ucodeid=self.desc_v3.UcodeId) @@ -195,12 +195,12 @@ def init_hw(self): self.reset(self.falcon, riscv=True) # set up the mailbox - self.nvdev.NV_PGSP_FALCON_MAILBOX0.write(lo32(self.nvdev.gsp.libos_args_sysmem[0])) - self.nvdev.NV_PGSP_FALCON_MAILBOX1.write(hi32(self.nvdev.gsp.libos_args_sysmem[0])) + self.nvdev.NV_PGSP_FALCON_MAILBOX0.write(lo32(self.nvdev.gsp.libos_args_sysmem)) + self.nvdev.NV_PGSP_FALCON_MAILBOX1.write(hi32(self.nvdev.gsp.libos_args_sysmem)) # booter self.reset(self.sec2) - mbx = self.execute_hs(self.sec2, self.booter_image_sysmem[0], code_off=self.booter_code_off, data_off=self.booter_data_off, + mbx = self.execute_hs(self.sec2, self.booter_image_paddr, code_off=self.booter_code_off, data_off=self.booter_data_off, imemPa=0x0, imemVa=self.booter_code_off, imemSz=self.booter_code_sz, dmemPa=0x0, dmemVa=0x0, dmemSz=self.booter_data_sz, pkc_off=0x10, engid=1, ucodeid=3, mailbox=self.nvdev.gsp.wpr_meta_sysmem) assert mbx[0] == 0x0, f"Booter failed to execute, mailbox is {mbx[0]:08x}, {mbx[1]:08x}" @@ -208,11 +208,11 @@ def init_hw(self): self.nvdev.NV_PFALCON_FALCON_OS.with_base(self.falcon).write(0x0) assert self.nvdev.NV_PRISCV_RISCV_CPUCTL.with_base(self.falcon).read_bitfields()['active_stat'] == 1, "GSP Core is not active" - def execute_dma(self, base:int, cmd:int, dest:int, mem_off:int, sysmem:int, size:int): + def execute_dma(self, base:int, cmd:int, dest:int, mem_off:int, src:int, size:int): wait_cond(lambda: self.nvdev.NV_PFALCON_FALCON_DMATRFCMD.with_base(base).read_bitfields()['full'], value=0, msg="DMA does not progress") - self.nvdev.NV_PFALCON_FALCON_DMATRFBASE.with_base(base).write(lo32(sysmem >> 8)) - self.nvdev.NV_PFALCON_FALCON_DMATRFBASE1.with_base(base).write(hi32(sysmem >> 8) & 0x1ff) + self.nvdev.NV_PFALCON_FALCON_DMATRFBASE.with_base(base).write(lo32(src >> 8)) + self.nvdev.NV_PFALCON_FALCON_DMATRFBASE1.with_base(base).write(hi32(src >> 8) & 0x1ff) xfered = 0 while xfered < size: @@ -232,19 +232,19 @@ def start_cpu(self, base:int): def wait_cpu_halted(self, base): wait_cond(lambda: self.nvdev.NV_PFALCON_FALCON_CPUCTL.with_base(base).read_bitfields()['halted'], msg="not halted") - def execute_hs(self, base, img_sysmem, code_off, data_off, imemPa, imemVa, imemSz, dmemPa, dmemVa, dmemSz, pkc_off, engid, ucodeid, mailbox=None): + def execute_hs(self, base, img_paddr, code_off, data_off, imemPa, imemVa, imemSz, dmemPa, dmemVa, dmemSz, pkc_off, engid, ucodeid, mailbox=None): self.disable_ctx_req(base) - self.nvdev.NV_PFALCON_FBIF_TRANSCFG.with_base(base)[ctx_dma:=0].update(target=self.nvdev.NV_PFALCON_FBIF_TRANSCFG_TARGET_COHERENT_SYSMEM, - mem_type=self.nvdev.NV_PFALCON_FBIF_TRANSCFG_MEM_TYPE_PHYSICAL) + # target=0 is FB (not in published headers) + self.nvdev.NV_PFALCON_FBIF_TRANSCFG.with_base(base)[ctx_dma:=0].update(target=0, mem_type=self.nvdev.NV_PFALCON_FBIF_TRANSCFG_MEM_TYPE_PHYSICAL) cmd = self.nvdev.NV_PFALCON_FALCON_DMATRFCMD.with_base(base).encode(write=0, size=self.nvdev.NV_PFALCON_FALCON_DMATRFCMD_SIZE_256B, ctxdma=ctx_dma, imem=1, sec=1) - self.execute_dma(base, cmd, dest=imemPa, mem_off=imemVa, sysmem=img_sysmem+code_off-imemVa, size=imemSz) + self.execute_dma(base, cmd, dest=imemPa, mem_off=imemVa, src=img_paddr+code_off-imemVa, size=imemSz) cmd = self.nvdev.NV_PFALCON_FALCON_DMATRFCMD.with_base(base).encode(write=0, size=self.nvdev.NV_PFALCON_FALCON_DMATRFCMD_SIZE_256B, ctxdma=ctx_dma, imem=0, sec=0) - self.execute_dma(base, cmd, dest=dmemPa, mem_off=dmemVa, sysmem=img_sysmem+data_off-dmemVa, size=dmemSz) + self.execute_dma(base, cmd, dest=dmemPa, mem_off=dmemVa, src=img_paddr+data_off-dmemVa, size=dmemSz) self.nvdev.NV_PFALCON2_FALCON_BROM_PARAADDR.with_base(base)[0].write(pkc_off) self.nvdev.NV_PFALCON2_FALCON_BROM_ENGIDMASK.with_base(base).write(engid) @@ -293,7 +293,9 @@ def init_sw(self): self.nvdev.include("dev_fsp_pri", "gh100") self.nvdev.include("dev_bus", "tu102") - self.fmc_boot_args_view, self.fmc_boot_args_sysmem = self.nvdev._alloc_boot_struct(nv.GSP_FMC_BOOT_PARAMS()) + self.fmc_boot_args_view, _, fmc_boot_addrs = self.nvdev._alloc_boot_mem(ctypes.sizeof(nv.GSP_FMC_BOOT_PARAMS), + data=bytes(nv.GSP_FMC_BOOT_PARAMS())) + self.fmc_boot_args_sysmem = fmc_boot_addrs[0] self.init_fmc_image() def init_fmc_image(self): @@ -302,18 +304,19 @@ def init_fmc_image(self): def _section(s): return next((sh.content for sh in sections if sh.name == s)) self.fmc_booter_image, self.fmc_booter_hash = _section("image"), memoryview(_section("hash")).cast('I') self.fmc_booter_sig, self.fmc_booter_pkey = memoryview(_section("signature")).cast('I'), memoryview(_section("publickey") + b"\x00" * 3).cast('I') - _, self.fmc_booter_sysmem = self.nvdev._alloc_sysmem(len(self.fmc_booter_image), contiguous=True, data=self.fmc_booter_image) + _, _, fmc_booter_addrs = self.nvdev._alloc_boot_mem(len(self.fmc_booter_image), data=self.fmc_booter_image) + self.fmc_booter_bar1 = fmc_booter_addrs[0] def init_hw(self): self.falcon = 0x00110000 boot_args = nv.GSP_ACR_BOOT_GSP_RM_PARAMS(gspRmDescOffset=self.nvdev.gsp.wpr_meta_sysmem, gspRmDescSize=ctypes.sizeof(nv.GspFwWprMeta), target=nv.GSP_DMA_TARGET_COHERENT_SYSTEM, bIsGspRmBoot=True) - rm_args = nv.GSP_RM_PARAMS(bootArgsOffset=self.nvdev.gsp.libos_args_sysmem[0], target=nv.GSP_DMA_TARGET_COHERENT_SYSTEM) + rm_args = nv.GSP_RM_PARAMS(bootArgsOffset=self.nvdev.gsp.libos_args_sysmem, target=nv.GSP_DMA_TARGET_COHERENT_SYSTEM) self.fmc_boot_args_view[:ctypes.sizeof(nv.GSP_FMC_BOOT_PARAMS)] = bytes(nv.GSP_FMC_BOOT_PARAMS(bootGspRmParams=boot_args, gspRmParams=rm_args)) cot_payload = nv.NVDM_PAYLOAD_COT(version=0x2, size=ctypes.sizeof(nv.NVDM_PAYLOAD_COT), frtsVidmemOffset=0x1c00000, frtsVidmemSize=0x100000, - gspBootArgsSysmemOffset=self.fmc_boot_args_sysmem, gspFmcSysmemOffset=self.fmc_booter_sysmem[0]) + gspBootArgsSysmemOffset=self.fmc_boot_args_sysmem, gspFmcSysmemOffset=self.fmc_booter_bar1) for i,x in enumerate(self.fmc_booter_hash): cot_payload.hash384[i] = x for i,x in enumerate(self.fmc_booter_sig): cot_payload.signature[i] = x for i,x in enumerate(self.fmc_booter_pkey): cot_payload.publicKey[i] = x @@ -360,7 +363,7 @@ def init_rm_args(self, queue_size=0x40000): # Alloc queues pte_cnt = ((queue_pte_cnt:=(queue_size * 2) // 0x1000)) + round_up(queue_pte_cnt * 8, 0x1000) // 0x1000 pt_size = round_up(pte_cnt * 8, 0x1000) - queues_view, queues_sysmem = self.nvdev._alloc_sysmem(pt_size + queue_size * 2, contiguous=False) + queues_view, _, queues_sysmem = self.nvdev._alloc_boot_mem(pt_size + queue_size * 2, sysmem=True) # Fill up ptes for i, sysmem in enumerate(queues_sysmem): queues_view.view(i * 0x8, 0x8, fmt='Q')[0] = sysmem @@ -368,7 +371,9 @@ def init_rm_args(self, queue_size=0x40000): # Fill up arguments queue_args = nv.MESSAGE_QUEUE_INIT_ARGUMENTS(sharedMemPhysAddr=queues_sysmem[0], pageTableEntryCount=pte_cnt, cmdQueueOffset=pt_size, statQueueOffset=pt_size + queue_size) - _, self.rm_args_sysmem = self.nvdev._alloc_boot_struct(nv.GSP_ARGUMENTS_CACHED(bDmemStack=True, messageQueueInitArguments=queue_args)) + _, _, rm_args_addrs = self.nvdev._alloc_boot_mem(ctypes.sizeof(nv.GSP_ARGUMENTS_CACHED), + data=bytes(nv.GSP_ARGUMENTS_CACHED(bDmemStack=True, messageQueueInitArguments=queue_args))) + self.rm_args_sysmem = rm_args_addrs[0] # Build command queue header # self.cmd_q_va, self.stat_q_va = queues_view.addr + pt_size, queues_view.addr + pt_size + queue_size @@ -380,11 +385,12 @@ def init_rm_args(self, queue_size=0x40000): self.cmd_q = NVRpcQueue(self, self.cmd_q_view, None) def init_libos_args(self): - _, logbuf_sysmem = self.nvdev._alloc_sysmem((2 << 20), contiguous=True) - libos_args_view, self.libos_args_sysmem = self.nvdev._alloc_sysmem(0x1000, contiguous=True) + _, _, logbuf_addrs = self.nvdev._alloc_boot_mem(2 << 20) + libos_args_view, _, libos_addrs = self.nvdev._alloc_boot_mem(0x1000) + self.libos_args_sysmem = libos_addrs[0] libos_structs = [nv.LibosMemoryRegionInitArgument(kind=nv.LIBOS_MEMORY_REGION_CONTIGUOUS, loc=nv.LIBOS_MEMORY_REGION_LOC_SYSMEM, size=0x10000, - id8=int.from_bytes(bytes(f"LOG{name}", 'utf-8'), 'big'), pa=logbuf_sysmem[0] + 0x10000 * i) + id8=int.from_bytes(bytes(f"LOG{name}", 'utf-8'), 'big'), pa=logbuf_addrs[0] + 0x10000 * i) for i, name in enumerate(["INIT", "INTR", "RM", "MNOC", "KRNL"])] libos_structs.append(nv.LibosMemoryRegionInitArgument(kind=nv.LIBOS_MEMORY_REGION_CONTIGUOUS, loc=nv.LIBOS_MEMORY_REGION_LOC_SYSMEM, size=0x1000, id8=int.from_bytes(bytes("RMARGS", 'utf-8'), 'big'), pa=self.rm_args_sysmem)) @@ -400,7 +406,7 @@ def init_gsp_image(self): for i in range(3, 0, -1): npages[i-1] = ((npages[i] - 1) >> (nv.LIBOS_MEMORY_REGION_RADIX_PAGE_LOG2 - 3)) + 1 offsets = [sum(npages[:i]) * 0x1000 for i in range(4)] - radix_view, self.gsp_radix3_sysmem = self.nvdev._alloc_sysmem(offsets[-1] + len(self.gsp_image), contiguous=False) + radix_view, _, self.gsp_radix3_addrs = self.nvdev._alloc_boot_mem(offsets[-1] + len(self.gsp_image)) # Copy image radix_view.view(offsets[-1], len(self.gsp_image))[:] = self.gsp_image @@ -408,10 +414,11 @@ def init_gsp_image(self): # Copy level and image pages. for i in range(0, 3): cur_offset = sum(npages[:i+1]) - radix_view.view(offsets[i], npages[i+1] * 8, fmt='Q')[:] = array.array('Q', self.gsp_radix3_sysmem[cur_offset:cur_offset+npages[i+1]]) + radix_view.view(offsets[i], npages[i+1] * 8, fmt='Q')[:] = array.array('Q', self.gsp_radix3_addrs[cur_offset:cur_offset+npages[i+1]]) # Copy signature - _, self.gsp_signature_sysmem = self.nvdev._alloc_sysmem(len(signature), contiguous=True, data=signature) + _, _, gsp_sig_addrs = self.nvdev._alloc_boot_mem(len(signature), data=signature) + self.gsp_signature_bar1 = gsp_sig_addrs[0] def init_boot_binary_image(self): sha = {"ga102":"82428f532240727e95bb3083fbaaba9b2cc7b937314323f2d546ce7245f27fad", @@ -419,15 +426,16 @@ def init_boot_binary_image(self): "gb202":"d40b48e431d1707dc77af3605db358ed7a32ebfc2830eb74de2eddb4d3025071"}[self.nvdev.fw_name] h = nv.struct_nvfw_bin_hdr.from_buffer_copy(b:=fetch_fw(f"nvidia/{self.nvdev.fw_name}/gsp", "bootloader-570.144.bin", sha)) self.booter_image, self.booter_desc = b[h.data_offset:h.data_offset+h.data_size], nv.RM_RISCV_UCODE_DESC.from_buffer_copy(b, h.header_offset) - _, self.booter_sysmem = self.nvdev._alloc_sysmem(len(self.booter_image), contiguous=True, data=self.booter_image) + _, _, booter_addrs = self.nvdev._alloc_boot_mem(len(self.booter_image), data=self.booter_image) + self.booter_bar1 = booter_addrs[0] def init_wpr_meta(self): self.init_gsp_image() self.init_boot_binary_image() - common = {'sizeOfBootloader':(boot_sz:=len(self.booter_image)), 'sysmemAddrOfBootloader':self.booter_sysmem[0], - 'sizeOfRadix3Elf':(radix3_sz:=len(self.gsp_image)), 'sysmemAddrOfRadix3Elf': self.gsp_radix3_sysmem[0], - 'sizeOfSignature': 0x1000, 'sysmemAddrOfSignature': self.gsp_signature_sysmem[0], + common = {'sizeOfBootloader':(boot_sz:=len(self.booter_image)), 'sysmemAddrOfBootloader':self.booter_bar1, + 'sizeOfRadix3Elf':(radix3_sz:=len(self.gsp_image)), 'sysmemAddrOfRadix3Elf': self.gsp_radix3_addrs[0], + 'sizeOfSignature': 0x1000, 'sysmemAddrOfSignature': self.gsp_signature_bar1, 'bootloaderCodeOffset': self.booter_desc.monitorCodeOffset, 'bootloaderDataOffset': self.booter_desc.monitorDataOffset, 'bootloaderManifestOffset': self.booter_desc.manifestOffset, 'revision':nv.GSP_FW_WPR_META_REVISION, 'magic':nv.GSP_FW_WPR_META_MAGIC} @@ -441,7 +449,8 @@ def init_wpr_meta(self): gspFwHeapOffset=(gsp_heap_off:=round_down(gsp_off-gsp_heap_sz, 0x100000)), gspFwWprStart=(wpr_st:=round_down(gsp_heap_off-0x1000, 0x100000)), nonWprHeapSize=(non_wpr_sz:=0x100000), nonWprHeapOffset=(non_wpr_off:=round_down(wpr_st-non_wpr_sz, 0x100000)), gspFwRsvdStart=non_wpr_off) assert self.nvdev.flcn.frts_offset == m.frtsOffset, f"FRTS mismatch: {self.nvdev.flcn.frts_offset} != {m.frtsOffset}" - self.wpr_meta, self.wpr_meta_sysmem = self.nvdev._alloc_boot_struct(m) + self.wpr_meta, _, wpr_meta_addrs = self.nvdev._alloc_boot_mem(ctypes.sizeof(type(m)), data=bytes(m)) + self.wpr_meta_sysmem = wpr_meta_addrs[0] def promote_ctx(self, client:int, subdevice:int, obj:int, ctxbufs:dict[int, GRBufDesc], bufs=None, virt=None, phys=None): res, prom = {}, nv_gpu.NV2080_CTRL_GPU_PROMOTE_CTX_PARAMS(entryCount=len(ctxbufs), engineType=0x1, hChanClient=client, hObject=obj) @@ -527,8 +536,8 @@ def rpc_rm_alloc(self, hParent:int, hClass:int, params:Any, client=None) -> int: params.ramfcMem = nv_gpu.NV_MEMORY_DESC_PARAMS(base=ramfc_alloc.paddrs[0][0], size=0x200, addressSpace=2, cacheAttrib=0) params.instanceMem = nv_gpu.NV_MEMORY_DESC_PARAMS(base=ramfc_alloc.paddrs[0][0], size=0x1000, addressSpace=2, cacheAttrib=0) - _, method_sysmem = self.nvdev._alloc_sysmem(0x5000, contiguous=True) - params.mthdbufMem = nv_gpu.NV_MEMORY_DESC_PARAMS(base=method_sysmem[0], size=0x5000, addressSpace=1, cacheAttrib=0) + _, method_paddr, _ = self.nvdev._alloc_boot_mem(0x5000, sysmem=False) + params.mthdbufMem = nv_gpu.NV_MEMORY_DESC_PARAMS(base=method_paddr, size=0x5000, addressSpace=2, cacheAttrib=0) if client is not None and client != self.priv_root and params.hObjectError != 0: params.errorNotifierMem = nv_gpu.NV_MEMORY_DESC_PARAMS(base=0, size=0xecc, addressSpace=0, cacheAttrib=0) @@ -629,8 +638,8 @@ def run_cpu_seq(self, seq_buf:bytes): elif op == 0x8: # core resume self.nvdev.flcn.reset(self.nvdev.flcn.falcon, riscv=True) - self.nvdev.NV_PGSP_FALCON_MAILBOX0.write(lo32(self.libos_args_sysmem[0])) - self.nvdev.NV_PGSP_FALCON_MAILBOX1.write(hi32(self.libos_args_sysmem[0])) + self.nvdev.NV_PGSP_FALCON_MAILBOX0.write(lo32(self.libos_args_sysmem)) + self.nvdev.NV_PGSP_FALCON_MAILBOX1.write(hi32(self.libos_args_sysmem)) self.nvdev.flcn.start_cpu(self.nvdev.flcn.sec2) wait_cond(lambda: self.nvdev.NV_PGC6_BSI_SECURE_SCRATCH_14.read_bitfields()['boot_stage_3_handoff'], msg="SEC2 didn't hand off") diff --git a/tinygrad/runtime/support/nv/nvdev.py b/tinygrad/runtime/support/nv/nvdev.py index 7f210bbba981d..cadc599cca4fd 100644 --- a/tinygrad/runtime/support/nv/nvdev.py +++ b/tinygrad/runtime/support/nv/nvdev.py @@ -1,6 +1,6 @@ from __future__ import annotations -import ctypes, time, functools, tinygrad.runtime.autogen.nv_regs -from tinygrad.helpers import getenv, DEBUG, getbits +import time, functools, tinygrad.runtime.autogen.nv_regs +from tinygrad.helpers import getenv, DEBUG, getbits, round_up from tinygrad.runtime.autogen import pci from tinygrad.runtime.support.memory import TLSFAllocator, MemoryManager, AddrSpace from tinygrad.runtime.support.nv.ip import NV_FLCN, NV_FLCN_COT, NV_GSP @@ -145,15 +145,17 @@ def _early_mmu_init(self): self.mm = NVMemoryManager(self, self.vram_size - (64 << 20), boot_size=(2 << 20), pt_t=NVPageTableEntry, va_bits=bits, va_shifts=shifts, va_base=0, palloc_ranges=[(x, x) for x in [512 << 20, 2 << 20, 4 << 10]], reserve_ptable=not self.large_bar) - def _alloc_sysmem(self, size:int, vaddr:int=0, contiguous:bool=False, data:bytes|None=None) -> tuple[MMIOInterface, list[int]]: - view, paddrs = self.pci_dev.alloc_sysmem(size, vaddr, contiguous=contiguous) + def _alloc_boot_mem(self, size:int, data:bytes|None=None, contiguous:bool=False, sysmem:bool|None=None) -> tuple[MMIOInterface,int|None,list[int]]: + sz = round_up(size, 0x1000) + if sysmem is True or (sysmem is None and not self.large_bar): + view, sysaddr = self.pci_dev.alloc_sysmem(size, 0, contiguous=contiguous) + paddr = None + else: + paddr = self.mm.palloc(sz, boot=False) + view = self.vram.view(paddr, sz) + sysaddr = [self.pci_dev.bar_info(1)[0] + paddr + i * 0x1000 for i in range(sz // 0x1000)] if data is not None: view[:size] = data - return view, paddrs - - def _alloc_boot_struct(self, struct:ctypes.Structure) -> tuple[MMIOInterface, int]: - view, paddrs = self._alloc_sysmem(sz:=ctypes.sizeof(type(struct)), contiguous=True) - view[:sz] = bytes(struct) - return view, paddrs[0] + return view, paddr, sysaddr def include(self, name:str, arch:str): for k,v in getattr(getattr(tinygrad.runtime.autogen.nv_regs, name), arch or 'regs').items(): diff --git a/tinygrad/runtime/support/system.py b/tinygrad/runtime/support/system.py index a188c7131cb69..552ffa521b15d 100644 --- a/tinygrad/runtime/support/system.py +++ b/tinygrad/runtime/support/system.py @@ -38,8 +38,6 @@ def vfio(self) -> FileIOInterface|None: return vfio_fd except OSError: return None - def reserve_hugepages(self, cnt): os.system(f"sudo sh -c 'echo {cnt} > /proc/sys/vm/nr_hugepages'") - @functools.cache def reserve_va(self, va_start, va_size): # cached, runs only once per range. used to not collide with other mappings. @@ -64,7 +62,9 @@ def read_prop(svc, key) -> int: return int.from_bytes(bytes(buf), "little") iokit.IOServiceGetMatchingServices(0, iokit.IOServiceMatching(b"IOPCIDevice"), ctypes.byref(iterator:=ctypes.c_uint())) - while svc:=iokit.IOIteratorNext(iterator): all_devs.append((v:=read_prop(svc, "vendor-id"), d:=read_prop(svc, "device-id"), f"{v:x}:{d:x}")) + while svc:=iokit.IOIteratorNext(iterator): + if base_class is not None and read_prop(svc, "class-code") >> 16 != base_class: continue + all_devs.append((v:=read_prop(svc, "vendor-id"), d:=read_prop(svc, "device-id"), f"{v:x}:{d:x}")) else: try: devs = FileIOInterface("/sys/bus/pci/devices") except FileNotFoundError: raise RuntimeError("no pcie") diff --git a/tinygrad/schedule/__init__.py b/tinygrad/schedule/__init__.py index 29f36a6b94c09..d7eb642507c86 100644 --- a/tinygrad/schedule/__init__.py +++ b/tinygrad/schedule/__init__.py @@ -77,7 +77,7 @@ def create_new_buffer(ctx:tuple[dict[UOp, UOp], tuple[UOp, ...]], b:UOp): return ret pm_post_sched_cache = PatternMatcher([ - (UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx[1][x.arg]), + (UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx[1][x.arg.slot]), # create new BUFFERs for LUNIQUE BUFFERs from rangeify (UPat(Ops.BUFFER, src=(UPat(Ops.LUNIQUE), UPat(Ops.DEVICE)), name="b"), create_new_buffer), ]) diff --git a/tinygrad/schedule/allreduce.py b/tinygrad/schedule/allreduce.py index 77bb716397d09..b2d61660823ba 100644 --- a/tinygrad/schedule/allreduce.py +++ b/tinygrad/schedule/allreduce.py @@ -55,8 +55,7 @@ def handle_allreduce(buf:UOp, red:UOp) -> UOp|None: return UOp.usum(*[c.pad(((s,numel-e),)) for (s,e),c in zip(chunks, copied_chunks)]).reshape(shape) def create_allreduce_function(buf:UOp, red:UOp, output:UOp|None=None) -> UOp|None: - # BUFFER without unique have unique added later - if output is None: output = UOp.unique_const(Invalid, red.dtype, red.device, red.shape).contiguous() + if output is None: output = UOp.const(red.dtype, Invalid, red.device, red.shape).clone() to = red.param_like(0) src = buf.param_like(1) red = src.allreduce(red.arg, red.src[1]) diff --git a/tinygrad/schedule/indexing.py b/tinygrad/schedule/indexing.py index 792bd2e309d02..be49c2eff6f7e 100644 --- a/tinygrad/schedule/indexing.py +++ b/tinygrad/schedule/indexing.py @@ -7,7 +7,7 @@ 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 -ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.AFTER, Ops.COPY, Ops.BUFFER, Ops.BUFFER_VIEW, +ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.AFTER, Ops.COPY, Ops.BUFFER, Ops.SLICE, Ops.CONST, Ops.BIND, Ops.DEVICE, Ops.MSELECT, Ops.MSTACK, Ops.PARAM, Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.LOAD, Ops.CALL, Ops.FUNCTION} @@ -18,8 +18,8 @@ def realize_srcs(ctx:dict[UOp, None], rb:UOp) -> None: if s.base.op not in ALWAYS_CONTIGUOUS: ctx[s] = None def realize_store_after_src(ctx:dict[UOp, None], dest:UOp, src:UOp): - # don't realize COPY/BUFFER_VIEW when they are the direct source of STORE+AFTER — the target buffer is the output - if src.op in {Ops.COPY, Ops.BUFFER_VIEW} and src in ctx \ + # don't realize COPY/SLICE when they are the direct source of STORE+AFTER — the target buffer is the output + if src.op in {Ops.COPY, Ops.SLICE} and src in ctx \ and not dest.op_in_backward_slice_with_self(Ops.SHRINK, Ops.PERMUTE, Ops.FLIP, Ops.PAD): del ctx[src] # you don't usually have to do this for assign unless there's a WAR hazard like TestAssign.test_assign_double_diamond_reduce @@ -58,7 +58,7 @@ def create_bufferize_and_index_based_on_ranges(ctx:IndexingContext, x:UOp): new_srcs = [] for s in x.src: new_src = s - if s.op in {Ops.PARAM, Ops.BUFFER_VIEW, Ops.MSTACK, Ops.MSELECT, Ops.AFTER}: + if s.op in {Ops.PARAM, Ops.BUFFER, Ops.SLICE, Ops.MSTACK, Ops.MSELECT, Ops.AFTER}: if x in ctx.range_map: new_src = new_src.index(*ctx.range_map[x][0]) elif s in ctx.realize_map: realized_ranges = ctx.realize_map[s] @@ -128,15 +128,15 @@ def _apply_reshape(in_shape:tuple[sint,...], out_shape:tuple[sint, ...], urngs:U @functools.cache def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UOp, ...]) -> tuple[UOp, ...]: match op: - case Ops.SHRINK: rngs = tuple(a if ss == 0 else a+ss for a,(ss,_) in zip(rngs, arg)) + case Ops.SHRINK: rngs = tuple(a if off == 0 else a+off for a,(off,_) in zip(rngs, arg)) case Ops.PERMUTE: rngs = tuple(rngs[p] for p in argsort(arg)) case Ops.FLIP: rngs = tuple(((s-1)-a) if f else a for a,s,f in zip(rngs, in_shape, arg)) case Ops.EXPAND: rngs = tuple(a if in_sh == out_sh else a.const_like(0) for a,in_sh,out_sh in zip(rngs, in_shape, arg)) case Ops.PAD: # NOTE: the .where(r-s, i) is not inside the graph_rewrite so that `convert_pad_to_where_to_keep_behavior_local` # wraps the pad with only the newly added valid - rngs = tuple(r if (s == 0 and e == 0) else graph_rewrite((r >= s) & (r < (sh+s)), - symbolic+pm_simplify_valid, name="pad").where(r-s, UOp.invalid()) for r,sh,(s,e) in zip(rngs, in_shape, arg)) + rngs = tuple(r if (sz == sh and off == 0) else graph_rewrite((r >= off) & (r < (sh+off)), + symbolic+pm_simplify_valid, name="pad").where(r-off, UOp.invalid()) for r,sh,(off,sz) in zip(rngs, in_shape, arg)) case Ops.RESHAPE: sink = UOp.sink(*rngs).simplify() # NOTE: this applies any commutative flips to the rngs early sub_array = {r:UOp.range(r.src[0], i, AxisType.PLACEHOLDER) for i,r in enumerate(sink.ranges)} diff --git a/tinygrad/schedule/memory.py b/tinygrad/schedule/memory.py index 783770e044d5f..56c16ef5132c8 100644 --- a/tinygrad/schedule/memory.py +++ b/tinygrad/schedule/memory.py @@ -52,12 +52,11 @@ def _key(b:UOp): return (b.device, 1 if b in copy_bufs else 0) peaks[_key(buf)] = (max(peaks[_key(buf)][0], offsets[buf] + buf.arg * buf.dtype.itemsize), peaks[_key(buf)][1]) arena_sizes = {key: round_up(peak, block_size) for key, (peak, _) in peaks.items()} - # build replace_map: each buffer becomes a BUFFER_VIEW into a shared per-device-lane arena + # build replace_map: each buffer becomes a SLICE into a shared per-device-lane arena arenas = {key: UOp.new_buffer(key[0], sz, dtypes.int8) for key, sz in arena_sizes.items()} replace_map:dict[UOp, UOp] = {} for buf_uop, offset in offsets.items(): - assert offset % buf_uop.dtype.itemsize == 0, f"offset {offset} not aligned to {buf_uop.dtype.itemsize}" - replace_map[buf_uop] = UOp(Ops.BUFFER_VIEW, buf_uop.dtype, (arenas[_key(buf_uop)],), (buf_uop.arg, offset // buf_uop.dtype.itemsize)) + replace_map[buf_uop] = UOp(Ops.SLICE, buf_uop.dtype, (arenas[_key(buf_uop)], UOp.const(dtypes.weakint, offset)), buf_uop.arg) if DEBUG >= 1 and (omem:=sum(nbytes.values()) / 1e6) != (nmem:=sum(arena_sizes.values()) / 1e6): print(f"memory reduced from {omem:.2f} MB -> {nmem:.2f} MB, {len(first_appearance)} -> {len(arenas)} bufs") diff --git a/tinygrad/schedule/multi.py b/tinygrad/schedule/multi.py index c3f6c68ae11c0..c1ebb68b6f239 100644 --- a/tinygrad/schedule/multi.py +++ b/tinygrad/schedule/multi.py @@ -10,7 +10,7 @@ def mstack_early_shrink(ms:UOp, shrink:UOp): def apply_shrink(s:UOp, i:int) -> UOp: new_arg = [tuple([x.substitute({dvar[0]:dvar[0].const_like(i)}) if isinstance(x, UOp) and (dvar:=[v for v in x.variables() if v.expr=='_device_num']) else x for x in ss]) for ss in shrink.marg] - return s.shrink(tuple(new_arg)) + return s._mop(Ops.SHRINK, tuple(new_arg)) for i, x in enumerate(ms.src): if x.op is Ops.COPY: ret.append(apply_shrink(x.src[0], i).copy_to_device(x.device)) @@ -43,7 +43,9 @@ def apply_shrink(s:UOp, i:int) -> UOp: def alu_multi(root:UOp): msrcs = root.src - assert all_same([x.device for x in msrcs]), f"all buffers must have the same device {[x.device for x in msrcs]}" + devices = [x.device for x in msrcs if x.device is not None] + assert all_same(devices), f"all buffers must have the same device {devices}" + dcount = len(devices[0]) axis = root.axis assert axis is not None @@ -52,7 +54,7 @@ def alu_multi(root:UOp): if mlb.axis is None: # no axis, shard it assert mlb.op is not Ops.MULTI - srcs.append(mlb._shard(axis)) + srcs.append(mlb._shard(axis, dcount)) else: assert mlb.op is Ops.MULTI if mlb.axis == axis: @@ -60,7 +62,7 @@ def alu_multi(root:UOp): srcs.append(mlb.src[0]) else: # axis mismatch, copy to all devices, and shard it correctly - srcs.append(copy_multi(mlb, mlb.device)._shard(axis)) + srcs.append(copy_multi(mlb, mlb.device)._shard(axis, dcount)) return srcs[0].alu(root.op, *srcs[1:]).multi(axis) def reduce_multi(root:UOp, multi:UOp): @@ -86,22 +88,25 @@ def expand_multi(root:UOp, multi:UOp): return multi.src[0].expand(new_shape).multi(multi.axis) def pad_multi(root:UOp, multi:UOp): - assert multi.axis is None or root.marg[multi.axis] == (0,0), f"padding not supported for {root.marg=}" - return multi.src[0].pad(root.marg).multi(multi.axis) + assert multi.axis is None or root.marg[multi.axis] == (0, multi.shape[multi.axis]), f"padding not supported for {root.marg=}" + local_pad = tuple((0, multi.src[0].shape[multi.axis]) if a == multi.axis else s for a,s in enumerate(root.marg)) + return multi.src[0]._mop(Ops.PAD, local_pad).multi(multi.axis) def permute_multi(root:UOp, multi:UOp): # all permutes supported! return multi.src[0].permute(root.marg).multi(root.axis) def shrink_multi(root:UOp, multi:UOp): - assert multi.axis is None or root.marg[multi.axis] == (0, multi.shape[multi.axis]) or root.marg[multi.axis] in multi.bounds, \ + shard_bounds = tuple((s,e-s) for s,e in multi.bounds) if multi.axis is not None else () + assert multi.axis is None or root.marg[multi.axis] == (0, multi.shape[multi.axis]) or root.marg[multi.axis] in shard_bounds, \ f"shrinking not supported for {root.marg=}" - if multi.axis is not None and root.marg[multi.axis] in multi.bounds and root.marg[multi.axis] != (0, multi.shape[multi.axis]): + if multi.axis is not None and root.marg[multi.axis] in shard_bounds and root.marg[multi.axis] != (0, multi.shape[multi.axis]): # NOTE: shrink on the shard axis is only allowed when result is a single partition, denoted by the new real # we just copy it to all the devices, no real. this will be optimized out later non_shard_shrink = tuple((0, multi.src[0].shape[i]) if i == multi.axis else s for i, s in enumerate(root.marg)) - return multi.src[0].copy_to_device(multi.device, arg=multi.bounds.index(root.marg[multi.axis])).shrink(non_shard_shrink) - return multi.src[0].shrink(tuple((0, multi.src[0].shape[multi.axis]) if a == multi.axis else s for a,s in enumerate(root.marg))).multi(multi.axis) + return multi.src[0].copy_to_device(multi.device, arg=shard_bounds.index(root.marg[multi.axis]))._mop(Ops.SHRINK, non_shard_shrink) + local_shrink = tuple((0, multi.src[0].shape[multi.axis]) if a == multi.axis else s for a,s in enumerate(root.marg)) + return multi.src[0]._mop(Ops.SHRINK, local_shrink).multi(multi.axis) def flip_multi(root:UOp, multi:UOp): assert multi.axis is None or not root.marg[multi.axis], "flipping not supported on sharded axis" @@ -109,6 +114,9 @@ def flip_multi(root:UOp, multi:UOp): def copy_multi(multi:UOp, device:str | tuple[str, ...] | UOp): assert multi.axis is not None, "all multi ops have axis" + if isinstance(device, UOp) and isinstance(device.arg, str): + pieces = [multi.src[0].mselect(i).copy_to_device(device) for i in range(len(multi.device))] + return pieces[0].cat(*pieces[1:], dim=multi.axis) return multi.src[0]._unshard(multi.axis).allreduce(Ops.ADD, device) def store_after_multi(dest:UOp, src:UOp): return dest.after(dest.store(src.src[0])).multi(src.axis) @@ -129,7 +137,7 @@ def rewrite_into_function(call:UOp): def param_to_multi(p:UOp): if p.axis is None: return None - return UOp.param(p.arg, p.dtype, p.shard_shape, p._device).multi(p.axis) + return UOp.param(p.arg.slot, p.dtype, p.shard_shape, p.device, p.arg.vmin_vmax, p.arg.name, p.arg.addrspace).multi(p.axis) # NOTE: this is the same pattern as Ops.UNROLL multi_pm = PatternMatcher([ diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index f3f375fa0a967..04a0109ea187a 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -1,7 +1,7 @@ from dataclasses import dataclass, field, replace import itertools from tinygrad.dtype import dtypes, PtrDType, AddrSpace, Invalid -from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, _substitute, KernelInfo +from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, _substitute, KernelInfo, ParamArg from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, profile_matches, identity_element from tinygrad.uop.symbolic import symbolic from tinygrad.helpers import prod, all_same, getenv, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS @@ -52,7 +52,7 @@ def found_after(ctx:dict[UOp, UOp], after:UOp, src:UOp): if x.op is Ops.PERMUTE: x, after = x.src[0], after.permute(argsort(x.marg)) elif x.op is Ops.RESHAPE: x, after = x.src[0], after.reshape(x.src[0].shape) elif x.op is Ops.WHERE and x.src[2].base.arg == Invalid and x.src[1].op is Ops.PAD: - x, after = x.src[1].src[0], after.shrink(tuple((l, s-r) for (l,r),s in zip(x.src[1].marg, x.shape))) + x, after = x.src[1].src[0], after.shrink(tuple((o, s+o) for (o,_),s in zip(x.src[1].marg, x.src[1].src[0].shape))) else: break ctx[x] = after @@ -64,10 +64,19 @@ def found_after(ctx:dict[UOp, UOp], after:UOp, src:UOp): ]) # movement op on INDEX as a PatternMatcher +def _mop_index(r:UOp, idx:UOp): + idxs = idx.src[1:] + if len(idxs) == len(r.shape): + return r.src[0].index(*apply_movement_op(r.op, r.src[0].shape, r.marg, idxs), dtype=idx.dtype, arg=idx.arg) + if r.op is Ops.RESHAPE: + src_prefix = len(r.src[0].shape) - len(r.shape[len(idxs):]) + if src_prefix >= 0 and r.src[0].shape[src_prefix:] == r.shape[len(idxs):]: + if src_prefix == 0: return r.src[0] if r.src[0].dtype == idx.dtype else None + ret = r.src[0].index(*apply_movement_op(r.op, r.src[0].shape[:src_prefix], r.shape[:len(idxs)], idxs), dtype=idx.dtype, arg=idx.arg) + return ret if ret.shape == idx.shape else None + pm_mops = PatternMatcher([ - (UPat(GroupOp.Movement, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"), - lambda r,idx: r.src[0].index(*apply_movement_op(r.op, r.src[0].shape, r.marg, idx.src[1:]), dtype=idx.dtype, arg=idx.arg) - if len(idx.src[1:]) == len(r.shape) else None), + (UPat(GroupOp.Movement, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"), _mop_index), # move movement ops and INDEX after AFTER (but not when AFTER has a raw STORE with shaped children — from replace_contig_with_store_after) (UPat(GroupOp.Movement|{Ops.INDEX}, name="r").after(name="a", allow_any_len=True), lambda r,a: UOp(r.op, r.dtype, (a.replace(src=(r.src[0],)+a.src[1:]),)+r.src[1:], r.arg)), @@ -121,15 +130,15 @@ def resolve_function(c:UOp, allow_param_mismatch=True) -> UOp|None: if c.arg.precompile: return None params: list[UOp] = [] graph_rewrite(c.src[0], pm_gather_params, bottom_up=True, ctx=params, name="gather params") - params = sorted(params, key=lambda x: x.arg) + params = sorted(params, key=lambda x: x.arg.slot) args = c.src[1:] # NOTE: this isn't really needed. it's okay if there's unused args in the function if not allow_param_mismatch: - if [x.arg for x in params] != list(range(len(params))): raise RuntimeError(f"params not in order: {[x.arg for x in params]}") + if [x.arg.slot for x in params] != list(range(len(params))): raise RuntimeError(f"params not in order: {[x.arg.slot for x in params]}") if len(params) != len(args): raise TypeError(f"expected {len(params)} args, got {len(args)}") - dict_map = {x:args[x.arg] for x in params} + dict_map = {x:args[x.arg.slot] for x in params} for i, (p, a) in enumerate(dict_map.items()): if p.axis != a.axis: raise TypeError(f"arg {i} axis mismatch: expected {p.axis}, got {a.axis}") if p.max_shape != a.max_shape: raise TypeError(f"arg {i} shape mismatch: expected {p.shape}, got {a.shape}") @@ -298,7 +307,7 @@ def buf_gate(x:UOp): return src.substitute(replaced, extra_pm=pm_gate_substitute) def remove_noop_bufferize(idx,b2): - if idx.src[1:] != b2.src[1:] or idx.src[0].op is Ops.BUFFER_VIEW: return None + if idx.src[1:] != b2.src[1:] or idx.src[0].op is Ops.SLICE: return None return idx.src[0].shrink(tuple((0, s) for s in b2.shape)) if b2.shape else idx.src[0] pm_const_buffer_folding = pm_mops+PatternMatcher([ @@ -338,11 +347,12 @@ def late_buffer_view(t:UOp, b:UOp): assert x.op not in GroupOp.Elementwise, "can't buffer view elementwise" x = x.src[0] x = next(u for u in x.src if u.op is Ops.INDEX) + assert x.op is Ops.INDEX, "must be INDEX" if len(shape) == 0: offset = x.src[1].arg else: offset = max(sum(idx.vmin for idx in x.src[1:]), 0) - return b.replace(src=(UOp(Ops.BUFFER_VIEW, t.dtype, (x.base,), (size, offset)), b.src[1])) + return b.replace(src=(UOp(Ops.SLICE, t.dtype, (x.src[0], UOp.const(dtypes.weakint, offset)), size),)) to_bufferview = PatternMatcher([ (UPat(Ops.STAGE, src=(UPat((Ops.BITCAST, Ops.CONTIGUOUS), name="t"), UPat()), name="b"), late_buffer_view), @@ -350,7 +360,7 @@ def late_buffer_view(t:UOp, b:UOp): DEVICE_MAX_BUFS = {"METAL": 31, "WEBGPU": 8} # TODO: get from device? def limit_bufs(ctx:IndexingContext, root:UOp): - if (device:=root._device) is None: return None # no device, index related calculations + if (device:=root.device) is None: return None # no device, index related calculations device = device if isinstance(device, str) else device[0].split(":")[0] if not (MAX_BUFS:=MAX_KERNEL_BUFFERS.value or DEVICE_MAX_BUFS.get(device, 0)): return None @@ -364,7 +374,7 @@ def gate_input(u:UOp): if len(bufs) > MAX_BUFS - 1: # NOTE: this -1 is for the output buffer srcs = [] for s in root.src: - if s.op in GroupOp.Elementwise and s._device is not None: + if s.op in GroupOp.Elementwise and s.device is not None: # Insert bufferize: all AxisType.REDUCE before bufferize are AxisType.LOOP orig_ranges, end_ranges = s.ranges, [x.replace(arg=(next(ctx.range_idx), AxisType.LOOP)) if x.op is Ops.RANGE else x for x in s.ranges] s = s.substitute(dict(zip(orig_ranges, end_ranges))).bufferize(*end_ranges, arg=BufferizeOpts(device=s.device)).index(*orig_ranges) @@ -381,7 +391,7 @@ def gate_input(u:UOp): # NOTE: this has been fixed up a bit def bufferize_to_store(ctx:itertools.count, x:UOp, idx:UOp, allow_locals=True): - size = prod(x.shape) + size = prod(x.shape) // x.dtype.count rngs = sorted(idx.ranges, key=lambda x: x.arg) assert size > 0 and isinstance(size, int), f"no zero sized or symbolic sized buffers {size}" @@ -404,7 +414,11 @@ def bufferize_to_store(ctx:itertools.count, x:UOp, idx:UOp, allow_locals=True): # NOTE: the DEFINE_LOCAL needs to be disambiguated here if sdtype.addrspace == AddrSpace.GLOBAL: buf = UOp(Ops.BUFFER, x.dtype, (UOp(Ops.LUNIQUE, arg=next(ctx)), UOp(Ops.DEVICE, arg=x.arg.device)), size) - do_store = buf.index(idx, dtype=sdtype).store(x.src[0]).end(*rngs) + if x.src[0].op is Ops.SLICE: + # no INDEX on SLICE, this could be cleaner + do_store = buf.store(x.src[0]).end(*rngs) + else: + do_store = buf.index(idx, dtype=sdtype).store(x.src[0]).end(*rngs) return buf.after(do_store) if allow_locals: @@ -442,6 +456,7 @@ def flatten_bufferize(x:UOp): # remove invalid writes (UPat(Ops.STORE, src=(UPat(), UPat(Ops.CONTIGUOUS, src=(UPat(Ops.CONST, arg=Invalid),)))), lambda: UOp(Ops.NOOP)), + (UPat(Ops.STORE, src=(UPat(), UPat(Ops.CONST, arg=Invalid))), lambda: UOp(Ops.NOOP)), (UPat(Ops.AFTER, src=(UPat.var("x"), UPat(Ops.NOOP, src=()))), lambda x: x), (UPat(Ops.AFTER, src=(UPat.var("x"), UPat(Ops.END, src=(UPat(Ops.NOOP, src=()),), allow_any_len=True))), lambda x: x), ]) @@ -462,7 +477,7 @@ class LocalAddBufferContext: opts:tuple|None = None def debuf(ctx:LocalAddBufferContext, buf:UOp): - ret = UOp(Ops.PARAM, buf.dtype.ptr(prod(buf.max_shape)), arg=ctx.dg).reshape(buf.max_shape) + ret = UOp(Ops.PARAM, buf.dtype.ptr(prod(buf.max_shape), buf.addrspace), arg=ParamArg(ctx.dg, addrspace=buf.addrspace)).reshape(buf.max_shape) # if the buffer has symbolic shape, shrink the max-sized view to the actual shape if buf.max_shape != buf.shape: ret = ret.shrink(tuple((0, s) for s in buf.shape)) if buf not in ctx.map: ctx.map[buf] = buf @@ -474,7 +489,7 @@ def unbind_kernel(ctx:LocalAddBufferContext, b:UOp): return b.src[0] def handle_after(ctx:LocalAddBufferContext, after:UOp): - if isinstance(after.dtype, PtrDType) and after.ptrdtype.addrspace == AddrSpace.LOCAL: return None + if isinstance(after.dtype, PtrDType) and after.addrspace == AddrSpace.LOCAL: return None buf = after.buf_uop # HACK to put the buffer in the MAP instead of MSTACK/MSELECT if buf.op in {Ops.MSTACK, Ops.MSELECT}: buf = buf.src[0] @@ -497,9 +512,11 @@ def find_bufs(x:UOp): to_define_global = PatternMatcher([ (UPat(Ops.STORE, name="x"), find_bufs), (UPat(Ops.BUFFER, name="buf"), debuf), - (UPat(Ops.PARAM, src=(UPat(), UPat(Ops.DEVICE)), name="buf"), debuf), - (UPat(Ops.PARAM, src=(UPat(), UPat(), UPat.cvar('vmin'), UPat.cvar('vmax'), UPat.var("nm")), name="v"), - lambda v, vmin, vmax, nm: UOp.variable(nm.arg, vmin.arg, vmax.arg, v.dtype)), + (UPat(Ops.PARAM, name="v"), lambda v: + UOp.variable(v.arg.name, v.arg.vmin_vmax[0], v.arg.vmin_vmax[1], v.dtype) + if v.arg.name is not None and v.arg.vmin_vmax is not None else None), + (UPat(Ops.PARAM, name="buf"), lambda ctx, buf: + None if isinstance(buf.dtype, PtrDType) or buf.arg.name is not None or buf._shape is None else debuf(ctx, buf)), (UPat(Ops.INDEX, src=(UPat(Ops.DEFINE_VAR, name="v"),)), lambda v: v), (UPat(Ops.BIND, name="b"), unbind_kernel), @@ -549,11 +566,11 @@ def split_store(x:UOp) -> UOp|None: lctx = LocalAddBufferContext() ret = graph_rewrite(x, to_define_global+pm_flatten_range+rangeify_codegen, ctx=lctx, name="kernel split", bottom_up=True) - # SINK requires all buffers on the same device, but COPY/BUFFER_VIEW are cross-device or special hardware ops + # SINK requires all buffers on the same device, but COPY/SLICE are cross-device or special hardware ops if ret.op is Ops.STORE: stored = ret.src[1] elif ret.op is Ops.END and ret.src[0].op is Ops.STORE: stored = ret.src[0].src[1] else: raise RuntimeError(f"unknown kernel type {ret.op}") - if stored.op in {Ops.COPY, Ops.BUFFER_VIEW}: ret = stored.replace(src=stored.src + ret.ended_ranges) + if stored.op in {Ops.COPY, Ops.SLICE}: ret = stored.replace(src=stored.src + ret.ended_ranges) else: ret = ret.sink(arg=KernelInfo(opts_to_apply=lctx.opts)) kernel = ret.call(*lctx.map.values(), *lctx.vars.keys()) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 2c1fb8adb3c67..e0ee4531b70ec 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -61,8 +61,8 @@ def _frompy(x:list|tuple|bytes, dtype:DType, device:str|tuple[str,...]) -> UOp: ret.buffer.allocate(memoryview(data if device != "PYTHON" else bytearray(data))) return ret -def _get_winograd_matcols(mat, dims:int, shp:tuple[sint, ...], device:str|tuple[str, ...], dtype:DType) -> list[list[Tensor]]: - return [[Tensor.cat(*[Tensor.full(shp[:dim] + (1,) + shp[dim+1:], float(m[k]), device=device, dtype=dtype) for m in mat], dim=dim) +def _get_winograd_matcols(mat, dims:int, shp:tuple[sint, ...], device:str|tuple[str, ...]|None, dtype:DType) -> list[list[Tensor]]: + return [[Tensor.cat(*[Tensor.full(shp[:dim] + (1,) + shp[dim+1:], float(m[k]), device=device, dtype=dtype, buffer=False) for m in mat], dim=dim) for k in range(len(mat[0]))] for dim in range(dims)] # winograd conv 3 kernel f(4x4,3x3) see: http://arxiv.org/abs/1509.09308 @@ -88,14 +88,14 @@ class Tensor(OpMixin): np.set_printoptions(precision=4) ``` """ - __slots__ = "uop", "requires_grad", "grad" + __slots__ = "uop", "is_param", "grad" training: ClassVar[bool] = False def __init__(self, data:ConstType|bytes|list|tuple|UOp|'numpy.ndarray'|pathlib.Path|None, - device:str|tuple|list|None=None, dtype:DTypeLike|None=None, requires_grad:bool=True, _force_unique:bool=False): + device:str|tuple|list|None=None, dtype:DTypeLike|None=None): if device is None: if isinstance(data, pathlib.Path): device = f"DISK:{data.resolve()}" # keep it on the disk if device is None - elif isinstance(data, UOp): device = data._device + elif isinstance(data, UOp): device = data.device _dtype:DType|None = to_dtype(dtype) if dtype is not None else None _device:str|tuple[str, ...] = canonicalize_device(device) del device, dtype @@ -103,18 +103,17 @@ def __init__(self, data:ConstType|bytes|list|tuple|UOp|'numpy.ndarray'|pathlib.P # tensors can have gradients if you have called .backward self.grad:Tensor|None = None - self.requires_grad:bool = requires_grad + self.is_param:bool = True # create a UOp from the different types of inputs if isinstance(data, UOp): assert _dtype is None or _dtype==data.dtype or data.dtype==dtypes.weakint, f"dtype mismatch: {_dtype} vs {data.dtype}" # if data is dtype.weakint that means that this is a symbolic int and we need to lower it to something we can make a Tensor out of - if data.dtype == dtypes.weakint: data = Tensor.from_uop(data, device=_device).uop + if data.dtype == dtypes.weakint: data = _index_to_concrete_int(data) elif data is None: - data = UOp.const(_dtype or dtypes.default_float, 0, _device) + data = UOp.const(_dtype or dtypes.default_float, 0) elif isinstance(data, get_args(ConstType)): - dt = _dtype or dtypes.from_py(data) - data = UOp.unique_const(data, dt, _device) if _force_unique or (requires_grad and dtypes.is_float(dt)) else UOp.const(dt, data, _device) + data = UOp.const(_dtype or dtypes.from_py(data), data) elif isinstance(data, bytes): data = _frompy(data, _dtype or dtypes.uint8, _device) elif isinstance(data, (list, tuple)): if _dtype is None: @@ -126,7 +125,7 @@ def __init__(self, data:ConstType|bytes|list|tuple|UOp|'numpy.ndarray'|pathlib.P import numpy as np assert isinstance(data, np.ndarray), f"expected np.ndarray, got {data}" if data.shape == (): - data = UOp.const(_dtype or _from_np_dtype(data.dtype), data.item(), _device) + data = UOp.const(_dtype or _from_np_dtype(data.dtype), data.item()) else: data = _fromnp(data.astype(npdtype) if _dtype is not None and (npdtype:=_to_np_dtype(_dtype)) is not None else data) elif isinstance(data, pathlib.Path): @@ -137,7 +136,7 @@ def __init__(self, data:ConstType|bytes|list|tuple|UOp|'numpy.ndarray'|pathlib.P if not isinstance(data, UOp): raise RuntimeError(f"can't create Tensor from {data!r} with type {type(data)}") # data might be on a different device - self.uop:UOp = data if data.device == _device else data.copy_to_device(_device) + self.uop:UOp = data if data.device is None or data.device == _device else data.copy_to_device(_device) # add to all_tensors after construction succeeds all_tensors[weakref.ref(self)] = None @@ -151,22 +150,25 @@ def _apply_uop(self, fxn:Callable[..., UOp], *x:Tensor, extra_args=(), **kwargs) if TRACEMETA >= 1 and (metadata:=_METADATA.get()) is not None: all_metadata[new_uop] = (metadata,) # directly create the Tensor ret = Tensor.__new__(Tensor) - ret.uop, ret.grad = new_uop, None - ret.requires_grad = any(t.requires_grad for t in srcs) + ret.uop, ret.grad, ret.is_param = new_uop, None, True # add to all_tensors after construction succeeds all_tensors[weakref.ref(ret)] = None return ret # alu and const_like are used by the mixins def alu(self, op: Ops, *src: Tensor) -> Tensor: return self._apply_uop(lambda *u: u[0].alu(op, *u[1:]), *src) - def const_like(self, b:ConstType) -> Tensor: return Tensor(self.uop.const_like(b), requires_grad=False) + def const_like(self, b:ConstType) -> Tensor: return Tensor(self.uop.const_like(b)) @staticmethod - def unique_const(fill_value:ConstType|UOp, **kwargs) -> Tensor: return Tensor(fill_value, _force_unique=True, **kwargs) + def const(dtype:DType, b:ConstType|UOp, device:str|tuple[str, ...]|None=None) -> Tensor: + return Tensor(UOp.const(dtype, b, device)) + @staticmethod + def unique_const(fill_value:ConstType|UOp, **kwargs) -> Tensor: + if isinstance(fill_value, UOp): return Tensor(fill_value, **kwargs) + dtype, device = kwargs.pop("dtype", None), kwargs.pop("device", None) + return Tensor(UOp.unique_const(fill_value, dtype, device), **kwargs) - def requires_grad_(self, requires_grad:bool=True) -> Tensor: - # make the UOp unique if it's a CONST to prevent gradient accumulation bugs with cached const UOps - if requires_grad and self.uop.op is Ops.CONST: self.replace(Tensor(self.uop.arg, device=self.device, dtype=self.dtype, requires_grad=True)) - self.requires_grad = requires_grad + def is_param_(self, is_param:bool=True) -> Tensor: + self.is_param = is_param return self class train(ContextDecorator): @@ -189,7 +191,7 @@ def __len__(self): return self.shape[0] @property - def device(self) -> str|tuple[str, ...]: return self.uop.device + def device(self) -> str|tuple[str, ...]|None: return self.uop.device @property def shape(self) -> tuple[sint, ...]: return self.uop.shape @@ -201,7 +203,7 @@ def dtype(self) -> DType: return self.uop.dtype def as_param(self, slot:int): if self.uop.axis is not None: - param = UOp.param(slot, self.dtype, self.uop.shard_shape, self.device).multi(self.uop.axis) + param = UOp.param(slot, self.dtype, self.uop.shard_shape, self.device, axis=self.uop.axis) else: param = UOp.param(slot, self.dtype, self.shape, self.device) return Tensor(param) @@ -238,7 +240,7 @@ def schedule_linear(self, *lst:Tensor) -> UOp: @disable_gc() def realize(self, *lst:Tensor, do_update_stats=True) -> Tensor: """Triggers the computation needed to create these Tensor(s).""" - if len(to_realize:=[x for x in (self,)+lst if not x.uop.has_buffer_identity()]): + if len(to_realize:=[x for x in (self,)+lst if x.uop.device is not None and not x.uop.has_buffer_identity()]): run_linear(*Tensor.linear_with_vars(*to_realize), update_stats=do_update_stats) return self @@ -258,7 +260,8 @@ def assign(self, x:Tensor|PyConst|list|tuple) -> Tensor: # broadcast x (shape only, dtype must match) if self.shape != x.shape: x = x._broadcast_to(self.shape) if self.shape != x.shape: raise RuntimeError(f"assign shape mismatch {self.shape} != {x.shape}") - if not is_disk and self.device != x.device: raise RuntimeError(f"assign device mismatch {self.device} != {x.device}") + if not is_disk and x.uop.device is not None and self.device is not None and self.device != x.device: + raise RuntimeError(f"assign device mismatch {self.device} != {x.device}") if not is_disk and self.dtype != x.dtype: raise RuntimeError(f"assign dtype mismatch {self.dtype} != {x.dtype}") if isinstance(self.device, tuple) and self.uop.axis != x.uop.axis: raise RuntimeError(f"multi axis mismatch {self.uop.axis} != {x.uop.axis}") @@ -279,19 +282,13 @@ def assign(self, x:Tensor|PyConst|list|tuple) -> Tensor: self.uop = assign return self - def detach(self) -> Tensor: - """ - Returns a new tensor with the same data as this tensor, but detached from the autograd graph. - """ - return Tensor(self.uop.detach(), requires_grad=False) - def _buffer(self) -> Buffer: from tinygrad.engine.realize import capturing if capturing and not getenv("UNSAFE_ALLOW_JIT_BUFFER"): from tinygrad.engine.jit import JitError raise JitError("cannot access tensor data during JIT capture, the value will be baked in") x = self.cast(self.dtype.base).contiguous() - if isinstance(self.device, tuple): x = x.to("CPU") + if self.uop.device is None or isinstance(self.device, tuple): x = x.clone("CPU") return cast(Buffer, x.realize().uop.buffer).ensure_allocated() def _data(self) -> memoryview: return self._buffer().as_memoryview() @@ -356,22 +353,24 @@ def numpy(self) -> 'numpy.ndarray': if 0 in self.shape: return np.empty(self.shape, dtype=_to_np_dtype(self.dtype.base)) return self._buffer().numpy().reshape(self.shape) - def clone(self) -> Tensor: + def clone(self, device:str|tuple[str, ...]|None=None) -> Tensor: """ Creates a clone of this tensor allocating a separate buffer for the data. + If `device` is specified, the clone is placed on that device. """ - ret = self.empty_like() - if self.grad is not None: ret.grad = self.grad.clone() - return ret.assign(self) + ret = Tensor(self.uop.clone(device=device)) + if self.grad is not None: ret.grad = self.grad.clone(device=device) + return ret.is_param_(self.is_param) def to(self, device:str|tuple[str, ...]|None) -> Tensor: """ Moves the tensor to the given device. """ + if self.uop.device is None: return self if (device:=canonicalize_device(device)) == self.device: return self - ret = Tensor(self.uop.copy_to_device(device), requires_grad=self.requires_grad) + ret = Tensor(self.uop.copy_to_device(device)) if self.grad is not None: ret.grad = self.grad.to(device) - return ret + return ret.is_param_(self.is_param) def to_(self, device:str|tuple[str, ...]|None) -> Tensor: """ @@ -390,11 +389,12 @@ def shard(self, devices:tuple[str, ...], axis:int|None=None) -> Tensor: print(t.shard((t.device, t.device), axis=1).uop) ``` """ + if self.uop.device is None: return self if not isinstance(self.device, str): raise RuntimeError("can't shard a multi-device tensor") if len(devices) == 1: return self.to(devices[0]) devices = cast(tuple[str, ...], canonicalize_device(devices)) uop = self.uop.shard(devices, self._resolve_dim(axis)) if axis is not None else self.uop.copy_to_device(devices) - return Tensor(uop, requires_grad=self.requires_grad) + return Tensor(uop).is_param_(self.is_param) def shard_(self, devices:tuple[str, ...], axis:int|None=None) -> Tensor: """ @@ -406,6 +406,7 @@ def shard_like(self, y:Tensor) -> Tensor: """ Shards the tensor the same way as `y` (same devices and axis). """ + if y.device is None: return self if isinstance(y.device, str): return self.to(y.device) return self if isinstance(self.device, tuple) and (y.device, y.uop.axis) == (self.device, self.uop.axis) else self.shard(y.device, y.uop.axis) @@ -462,20 +463,6 @@ def fs_store(self) -> Tensor: return data[:16].contiguous() - @staticmethod - def from_uop(y:UOp, **kwargs) -> Tensor: - # TODO: remove this and stay in weakint - if y.dtype == dtypes.weakint: y = _index_to_concrete_int(y) - if y.op is Ops.BIND: - var, val = y.unbind() - _device = canonicalize_device(kwargs.get("device")) - const = UOp.const(var.dtype, val, _device, ()) - return Tensor(y.replace(src=(var.replace(src=const.src), const)), **kwargs, requires_grad=False) - if y.op is Ops.CONST: return Tensor(y.arg, **kwargs, requires_grad=False) - if y.op is Ops.MUL: return Tensor.from_uop(y.src[0]) * Tensor.from_uop(y.src[1]) - if y.op is Ops.ADD: return Tensor.from_uop(y.src[0]) + Tensor.from_uop(y.src[1]) - raise RuntimeError(f"unhandled UOp {y}") - # ***** creation entrypoint ***** @staticmethod @@ -498,7 +485,7 @@ def empty_like(self, dtype:DTypeLike|None=None, device:str|tuple[str, ...]|None= Creates an empty tensor with the same shape as `self`. If `dtype` is not specified, the dtype of `self` is used. """ - return Tensor(self.uop.empty_like(dtype, device), **kwargs) + return Tensor(self.uop.empty_like(dtype, self.device if device is None else device), **kwargs) @staticmethod def from_blob(ptr:int, shape:tuple[int, ...], **kwargs) -> Tensor: @@ -552,8 +539,8 @@ def manual_seed(seed=0) -> None: def _next_counter(device:str, num:int) -> tuple[Tensor, Tensor]: if device not in Tensor._device_seeds: seed = [int.from_bytes(hashlib.sha256(len(Tensor._device_seeds).to_bytes(4, "big")).digest(), "big"), Tensor._seed] - Tensor._device_seeds[device] = Tensor(seed, device=device, dtype=dtypes.uint32, requires_grad=False) - Tensor._device_rng_counters[device] = Tensor([0, 0], device=device, dtype=dtypes.uint32, requires_grad=False) + Tensor._device_seeds[device] = Tensor(seed, device=device, dtype=dtypes.uint32) + Tensor._device_rng_counters[device] = Tensor([0, 0], device=device, dtype=dtypes.uint32) counter = Tensor._device_rng_counters[device] new_low = counter[0:1] + (num & 0xffffffff) new_high = counter[1:2] + (num >> 32) + (new_low < counter[0]) @@ -563,7 +550,7 @@ def _next_counter(device:str, num:int) -> tuple[Tensor, Tensor]: return Tensor._device_seeds[device], low.cat(high) @staticmethod - def rand(*shape, device:str|None=None, dtype:DTypeLike|None=None, requires_grad:bool=True, contiguous:bool=True) -> Tensor: + def rand(*shape, device:str|None=None, dtype:DTypeLike|None=None, contiguous:bool=True) -> Tensor: """ Creates a tensor with the given shape, filled with random values from a uniform distribution over the interval `[0, 1)`. @@ -582,41 +569,24 @@ def rand(*shape, device:str|None=None, dtype:DTypeLike|None=None, requires_grad: device = cast(str, canonicalize_device(device)) # if shape has 0, return zero tensor - if (numel := prod(shape)) == 0: return Tensor.zeros(shape, device=device, dtype=dt, requires_grad=requires_grad) + if (numel := prod(shape)) == 0: return Tensor.zeros(shape, device=device, dtype=dt) num = ceildiv(numel * dt.itemsize, 4) key, counter = Tensor._next_counter(device, num) bits = Tensor.random_bits(key, counter, num) - out = Tensor._bits_to_rand(bits, shape, dt).requires_grad_(requires_grad) + out = Tensor._bits_to_rand(bits, shape, dt) return out.contiguous() if contiguous else out # ***** creation helper functions ***** - @classmethod - def eye(cls, n:int, m:int|None=None, dtype=None, device=None, requires_grad:bool=True) -> Tensor: - """ - Returns a 2-D tensor with `n` rows and `m` columns, with ones on the diagonal and zeros elsewhere. - - You can pass in `dtype` and `device` keyword arguments to control the data type and device of the tensor. - Additionally, all other keyword arguments are passed to the constructor of the tensor. - - ```python exec="true" source="above" session="tensor" result="python" - print(Tensor.eye(3).numpy()) - ``` - - ```python exec="true" source="above" session="tensor" result="python" - print(Tensor.eye(2, 4).numpy()) - ``` - """ - return super().eye(n, m, dtype, device).requires_grad_(requires_grad) - def _multi_like(self, fxn, *args, **kwargs) -> Tensor: dtype = kwargs.pop("dtype", self.dtype) - if kwargs.get("device") is not None: raise RuntimeError("cannot specify `device` on `*_like` of a multi device tensor") + if kwargs.pop("device", None) is not None: raise RuntimeError("cannot specify `device` on `*_like` of a multi device tensor") + assert isinstance(self.device, tuple), f"_multi_like needs a multi device tensor, got {self.device}" if self.uop.axis is None: return fxn(self.shape, *args, dtype=dtype, **kwargs).shard(self.device) stacked = UOp.mstack(*[fxn(self.uop.shard_shape, *args, device=d, dtype=dtype, **kwargs).uop for d in self.device]) - return Tensor(stacked.multi(self.uop.axis), requires_grad=kwargs.get("requires_grad", True)) + return Tensor(stacked.multi(self.uop.axis)) - def full_like(self, fill_value:ConstType, dtype=None, device=None, requires_grad:bool=False) -> Tensor: + def full_like(self, fill_value:ConstType, dtype=None, device=None) -> Tensor: """ Creates a tensor with the same shape as `self`, filled with the given value. If `dtype` is not specified, the dtype of `self` is used. @@ -628,9 +598,8 @@ def full_like(self, fill_value:ConstType, dtype=None, device=None, requires_grad print(Tensor.full_like(t, 42).numpy()) ``` """ - if device is None: return super().full_like(fill_value, dtype).requires_grad_(requires_grad) - if isinstance(self.device, tuple): raise RuntimeError("cannot specify `device` on `full_like` of a multi device tensor") - return Tensor.full(self.shape, fill_value, dtype=dtype or self.dtype, device=device).requires_grad_(requires_grad) + if isinstance(self.device, tuple): return self._multi_like(Tensor.full, fill_value, dtype=dtype or self.dtype, device=device) + return Tensor.full(self.shape, fill_value, dtype=dtype or self.dtype, device=self.device if device is None else device) def rand_like(self, **kwargs) -> Tensor: """ @@ -649,7 +618,7 @@ def rand_like(self, **kwargs) -> Tensor: # ***** random functions ***** - def randn_like(self, dtype:DTypeLike|None=None, requires_grad:bool=True, **kwargs) -> Tensor: + def randn_like(self, dtype:DTypeLike|None=None, **kwargs) -> Tensor: """ Creates a tensor with the same shape and sharding as `self`, filled with random values from a normal distribution with mean 0 and variance 1. @@ -663,10 +632,10 @@ def randn_like(self, dtype:DTypeLike|None=None, requires_grad:bool=True, **kwarg """ src = self.stack(self).rand_like(**{**kwargs, "dtype": dtypes.float32}) # https://en.wikipedia.org/wiki/Box%E2%80%93Muller_transform - return (src[0].mul(2*math.pi).cos().mul((1 - src[1]).log().mul(-2).sqrt()).cast(dtype or self.dtype)).requires_grad_(requires_grad) + return src[0].mul(2*math.pi).cos().mul((1 - src[1]).log().mul(-2).sqrt()).cast(dtype or self.dtype) @staticmethod - def randn(*shape, dtype:DTypeLike|None=None, requires_grad:bool=True, **kwargs) -> Tensor: + def randn(*shape, dtype:DTypeLike|None=None, **kwargs) -> Tensor: """ Creates a tensor with the given shape, filled with random values from a normal distribution with mean `0` and standard deviation `1`. If `dtype` is not specified, the default type is used. @@ -679,7 +648,7 @@ def randn(*shape, dtype:DTypeLike|None=None, requires_grad:bool=True, **kwargs) print(Tensor.randn(2, 3).numpy()) ``` """ - return Tensor.empty(*shape, **kwargs).randn_like(dtype=dtype, requires_grad=requires_grad) + return Tensor.empty(*shape, **kwargs).randn_like(dtype=dtype) @staticmethod def randint(*shape, low=0, high=10, dtype=dtypes.int32, **kwargs) -> Tensor: @@ -701,7 +670,7 @@ def randint(*shape, low=0, high=10, dtype=dtypes.int32, **kwargs) -> Tensor: return Tensor.uniform(*shape, low=low, high=high, dtype=dtype, **kwargs) @staticmethod - def normal(*shape, mean=0.0, std=1.0, requires_grad:bool=True, **kwargs) -> Tensor: + def normal(*shape, mean=0.0, std=1.0, **kwargs) -> Tensor: """ Creates a tensor with the given shape, filled with random values from a normal distribution with the given `mean` and standard deviation `std`. Requires `std >= 0`. @@ -715,10 +684,10 @@ def normal(*shape, mean=0.0, std=1.0, requires_grad:bool=True, **kwargs) -> Tens ``` """ if std < 0: raise ValueError(f"Tensor.normal requires std >= 0, got {std=}") - return (std * Tensor.randn(*shape, **kwargs) + mean).requires_grad_(requires_grad) + return std * Tensor.randn(*shape, **kwargs) + mean @staticmethod - def uniform(*shape, low=0.0, high=1.0, dtype:DTypeLike|None=None, requires_grad:bool=True, **kwargs) -> Tensor: + def uniform(*shape, low=0.0, high=1.0, dtype:DTypeLike|None=None, **kwargs) -> Tensor: """ Creates a tensor with the given shape, filled with random values from a uniform distribution over the interval `[low, high)`. Requires `low < high`. @@ -733,7 +702,7 @@ def uniform(*shape, low=0.0, high=1.0, dtype:DTypeLike|None=None, requires_grad: """ if not all_int(shape:=argfix(*shape)) or not all(s >= 0 for s in shape): raise ValueError(f"invalid input {shape=}") if low >= high: raise ValueError(f"Tensor.uniform requires low < high, got {low=}, {high=}") - return (((high-low) * Tensor.rand(*shape, **kwargs)).cast(dtype or dtypes.default_float) + low).requires_grad_(requires_grad) + return ((high-low) * Tensor.rand(*shape, **kwargs)).cast(dtype or dtypes.default_float) + low @staticmethod def scaled_uniform(*shape, **kwargs) -> Tensor: @@ -809,7 +778,7 @@ def randperm(n:int, device=None, dtype=dtypes.int32, **kwargs) -> Tensor: print(Tensor.randperm(6).numpy()) ``` """ - return Tensor.rand(n, device=device, **kwargs).argsort().cast(dtype).requires_grad_(kwargs.get("requires_grad", True)) + return Tensor.rand(n, device=device, **kwargs).argsort().cast(dtype) def multinomial(self:Tensor, num_samples:int = 1, replacement:bool = False) -> Tensor: """ @@ -856,7 +825,7 @@ def gradient(self, *targets:Tensor, gradient:Tensor|None=None) -> list[Tensor]: """ assert gradient is not None or self.shape == tuple(), "when no gradient is provided, backward must be called on a scalar tensor" if not (self.is_floating_point() and all(t.is_floating_point() for t in targets)): raise RuntimeError("only float Tensors have gradient") - if gradient is None: gradient = Tensor(1.0, dtype=self.dtype, device=self.device, requires_grad=False) + if gradient is None: gradient = Tensor(1.0, dtype=self.dtype, device=self.device) target_uops = [x.uop for x in targets] grads = compute_gradient(self.uop, gradient.uop, set(target_uops)) ret:list[Tensor] = [] @@ -870,14 +839,15 @@ def backward(self, gradient:Tensor|None=None) -> Tensor: Propagates the gradient of a tensor backwards through the computation graph. If the 'gradient' argument is not provided, the tensor must be a scalar, and the gradient is implicitly set to 1.0. ```python exec="true" source="above" session="tensor" result="python" - t = Tensor([1.0, 2.0, 3.0, 4.0], requires_grad=True) + t = Tensor([1.0, 2.0, 3.0, 4.0]) t.sum().backward() print(t.grad.numpy()) ``` """ all_uops = self.uop.toposort() + # backward fills .grad for every in-scope non-CONST float tensor tensors_need_grad: list[Tensor] = [t for tref in all_tensors if (t:=tref()) is not None and \ - t.uop in all_uops and t.requires_grad and t.is_floating_point()] + t.uop in all_uops and t.is_floating_point() and t.uop.op is not Ops.CONST] # clear contexts for t,g in zip(tensors_need_grad, self.gradient(*tensors_need_grad, gradient=gradient)): assert g.shape == t.shape, f"grad shape must match tensor shape, {g.shape!r} != {t.shape!r}" @@ -903,12 +873,13 @@ def _getitem(self, indices, v: Tensor|None = None) -> Tensor: match index: case Tensor(): if not dtypes.is_int(index.dtype): raise IndexError(f"index dtype {index.dtype} is not supported") - if index.device != self.device: raise RuntimeError(f"expected index and self on the same device, {index.device=}, {self.device=}") + if index.device is not None and self.device is not None and index.device != self.device: + raise RuntimeError(f"expected index and self on the same device, {index.device=}, {self.device=}") assert isinstance(size, int), "size must be an int" index = (index < 0).where(index+size, index) # treat negative index values case list() | tuple(): if not dtypes.is_int((ti:=Tensor(index)).dtype): raise IndexError(f"{index=} contains non-int element") - index = Tensor([i+size if i<0 else i for i in fully_flatten(index)], self.device, requires_grad=False).reshape(ti.shape) + index = Tensor([i+size if i<0 else i for i in fully_flatten(index)], self.device).reshape(ti.shape) case _: parsed = self._parse_view_index(index, size) indices_parsed.append({**parsed, "index":index}) if index is not None: dim += 1 @@ -1023,11 +994,12 @@ def __setitem__(self, indices, v:Tensor|PyConst|list|tuple) -> None: if isinstance(v, Tensor) and v.dtype != self.dtype: raise RuntimeError(f"setitem dtype mismatch: {self.dtype=} != {v.dtype=}") # raise if mutation would diverge from eager (allow only pure views of a realized buffer; exclude +=/-= RHS via v_uop/v_bw) v_uop, v_bw = (v.uop, v.uop.backward_slice) if isinstance(v, Tensor) else (None, {}) - shared = self.uop.base if self.uop.base.is_realized else None - if any(self.uop in t.uop.backward_slice_with_self and t.uop.base is not shared for tref in all_tensors - if (t:=tref()) is not None and t is not self and t.uop is not v_uop and t.uop not in v_bw): - raise RuntimeError("can't setitem on a tensor with other uses") - if not self.uop.base.is_realized and self.is_floating_point() and (self.requires_grad or (isinstance(v, Tensor) and v.requires_grad)): + if self.uop.op_in_backward_slice_with_self(Ops.BUFFER): + shared = self.uop.base if self.uop.base.is_realized else None + if any(self.uop in t.uop.backward_slice_with_self and t.uop.base is not shared for tref in all_tensors + if (t:=tref()) is not None and t is not self and t.uop is not v_uop and t.uop not in v_bw): + raise RuntimeError("can't setitem on a tensor with other uses") + if not self.uop.base.is_realized and self.is_floating_point(): if not isinstance(v, Tensor): v = Tensor(v, device=self.device, dtype=self.dtype) # __iadd__/__isub__ creates AFTER(view, STORE(view, computed)); unwrap to get the computed value if v.uop.op is Ops.AFTER and any(s.op is Ops.STORE for s in v.uop.src[1:]): v = v._apply_uop(lambda x: x.src[1].src[1]) @@ -1076,9 +1048,9 @@ def masked_select(self, mask, size:int|None=None, fill_value:ConstType=0): x, mask = self.flatten(), mask._broadcast_to(self.shape).flatten() mask_cumsum = mask.cumsum() if size is None: - counts = Tensor.zeros(mask_cumsum[-1].item() if mask.numel() else 0, dtype=dtypes.int32, device=self.device) + counts = Tensor.zeros(mask_cumsum[-1].item() if mask.numel() else 0, dtype=dtypes.int32, device=self.device, buffer=False) return x[counts.scatter(0, mask_cumsum, 1, reduce='add').cumsum()] - counts = Tensor.zeros(size, dtype=dtypes.int32, device=self.device).scatter(0, mask_cumsum, 1, reduce='add') + counts = Tensor.zeros(size, dtype=dtypes.int32, device=self.device, buffer=False).scatter(0, mask_cumsum, 1, reduce='add') return (Tensor.arange(size, device=self.device) < mask.sum()).where(x[counts.cumsum()], fill_value).cast(self.dtype) def nonzero(self, size:int|None=None, fill_value:ConstType=0) -> Tensor: @@ -1154,7 +1126,7 @@ def ctensor(l: Sequence[PyConst], dtype: DType = dtypes.uint64): data = (data.flatten(1) ^ pad_mask).reshape(*data.shape[:2], 200).bitcast(dtypes.uint64) - state = Tensor.zeros(bs, 25, device=self.device, dtype=dtypes.uint64) + state = Tensor.zeros(bs, 25, device=self.device, dtype=dtypes.uint64, buffer=False) for k in range(int(data.shape[1])): state = state ^ data[:, k] for i in range(24): # f1600 @@ -1283,7 +1255,7 @@ def contiguous(self, *args, **kwargs) -> Tensor: def ufix(self, x) -> Tensor: # TODO: x:ConstType|UOp does not work because mixin only accepts Self | ConstType assert isinstance(x, (*get_args(ConstType), UOp)), f"{type(x)=}, {x=}" - return Tensor(x, self.device, self.dtype if self._ufix_keep_dtype(x) else None, requires_grad=False) + return Tensor(self.uop.ufix(x)) def where(self:Tensor, x:Tensor|ConstType|sint, y:Tensor|ConstType|sint) -> Tensor: """ @@ -1305,7 +1277,7 @@ def where(self:Tensor, x:Tensor|ConstType|sint, y:Tensor|ConstType|sint) -> Tens """ if isinstance(x, Tensor): x, y = x._broadcasted(y) elif isinstance(y, Tensor): y, x = y._broadcasted(x) - else: x, y = Tensor(x, self.device, requires_grad=False)._broadcasted(y) + else: x, y = self.ufix(x)._broadcasted(y) out_shape = _broadcast_shape(self.shape, x.shape) return self.cast(dtypes.bool)._broadcast_to(out_shape)._apply_uop(UOp.where, x._broadcast_to(out_shape), y._broadcast_to(out_shape)) @@ -1361,8 +1333,8 @@ def dropout(self, p=0.5) -> Tensor: """ if not 0 <= p <= 1: raise ValueError(f"{p=} is out of range [0, 1]") if not Tensor.training or p == 0: return self - if p == 1: return self.zeros_like() - return (Tensor.rand_like(self, requires_grad=False, dtype=dtypes.default_float, contiguous=False) >= p).contiguous().where(self, 0) / (1.0 - p) + if p == 1: return self.const_like(0) + return (Tensor.rand_like(self, dtype=dtypes.default_float, contiguous=False) >= p).contiguous().where(self, 0) / (1.0 - p) def scaled_dot_product_attention(self, key:Tensor, value:Tensor, attn_mask:Tensor|None=None, dropout_p:float=0.0, is_causal:bool=False, enable_gqa:bool=False) -> Tensor: @@ -1389,7 +1361,7 @@ def scaled_dot_product_attention(self, key:Tensor, value:Tensor, attn_mask:Tenso # handle attention mask if is_causal: if attn_mask is not None: raise RuntimeError("cannot set attn_mask when is_causal=True") - attn_mask = qk.ones_like(requires_grad=False, dtype=dtypes.bool).tril() + attn_mask = qk.const_like(1).cast(dtypes.bool).tril() if attn_mask is not None: if attn_mask.dtype == dtypes.bool: attn_mask = attn_mask.where(0, -float("inf")) qk = qk + attn_mask diff --git a/tinygrad/uop/__init__.py b/tinygrad/uop/__init__.py index 7e3ebdb87fcd3..98ab6490e9cc7 100644 --- a/tinygrad/uop/__init__.py +++ b/tinygrad/uop/__init__.py @@ -27,7 +27,7 @@ class Ops(FastEnum): # uops that aren't rendered NOOP = auto(); REWRITE_ERROR = auto() # FUNCTION has a TUPLE body and is gradient-able; CALL is an opaque kernel invocation - PARAM = auto(); FUNCTION = auto(); CALL = auto(); PATCH = auto() + PARAM = auto(); FUNCTION = auto(); CALL = auto() # renderer # LINEAR is a list of UOps, SOURCE has a str arg that's human readable, BINARY has bytes arg that's compiled @@ -43,6 +43,9 @@ class Ops(FastEnum): # tuple/gettuple for function with multiple returns TUPLE = auto(); GETTUPLE = auto() + # hcq specific + GETADDR = auto() + # ** 3 -- load/store ** # INDEX is a BinaryOp similar to ADD, but it operates on pointers @@ -75,8 +78,8 @@ class Ops(FastEnum): # control flow ops BARRIER = auto(); RANGE = auto(); IF = auto(); END = auto(); ENDIF = auto(); WAIT = auto() - # consts. VCONST is a vectorized const - VCONST = auto(); CONST = auto() + # const. + CONST = auto() # CUSTOM/CUSTOMI are used to output strings into codegen. the I makes the string inline CUSTOM = auto(); CUSTOMI = auto() @@ -96,7 +99,7 @@ class Ops(FastEnum): CONTIGUOUS = auto(); CONTIGUOUS_BACKWARD = auto(); DETACH = auto() # buffer ops - STAGE = auto(); COPY = auto(); BUFFER = auto(); BUFFER_VIEW = auto(); MSELECT = auto(); MSTACK = auto(); CUSTOM_FUNCTION = auto() + STAGE = auto(); COPY = auto(); BUFFER = auto(); SLICE = auto(); MSELECT = auto(); MSTACK = auto(); CUSTOM_FUNCTION = auto() # the core 6 movement ops! these only exist in the tensor graph RESHAPE = auto(); PERMUTE = auto(); EXPAND = auto(); PAD = auto(); SHRINK = auto(); FLIP = auto() @@ -114,6 +117,7 @@ class GroupOp: Ops.XOR, Ops.SHL, Ops.SHR, Ops.OR, Ops.AND, Ops.THREEFRY, Ops.SUB, Ops.FDIV, Ops.POW, Ops.FLOORDIV, Ops.FLOORMOD} Ternary = {Ops.WHERE, Ops.MULACC} ALU = set.union(Unary, Binary, Ternary) + Broadcastable = set.union(Binary, Ternary, {Ops.GROUP, Ops.STORE}) # TODO: is BITCAST always Elementwise if it's shape changing? Elementwise = set.union(ALU, {Ops.CAST, Ops.BITCAST}) diff --git a/tinygrad/uop/decompositions.py b/tinygrad/uop/decompositions.py index d468f1a467bb4..3f58cc3d61bca 100644 --- a/tinygrad/uop/decompositions.py +++ b/tinygrad/uop/decompositions.py @@ -1,10 +1,10 @@ from typing import Callable import math, functools from tinygrad.dtype import dtypes, DType, promo_lattice, truncate -from tinygrad.device import is_dtype_supported -from tinygrad.helpers import flatten, polyN, Target, EMULATED_DTYPES +from tinygrad.helpers import flatten, polyN, DEBUG, EMULATED_DTYPES from tinygrad.uop import GroupOp from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, graph_rewrite +from tinygrad.renderer import Renderer TRANSCENDENTAL_DTYPES = (dtypes.float16, dtypes.float32, dtypes.float64) @@ -279,9 +279,10 @@ def magicgu(vmax:int, d:int) -> tuple[int,int]: return m, s assert False -def fast_idiv(target: Target, x: UOp, d: int, dont_cast=False) -> UOp|None: +def fast_idiv(ren: Renderer, x: UOp, d: int, dont_cast=False) -> UOp|None: + from tinygrad.renderer.cstyle import MetalRenderer # NOTE: disable for METAL due to compiler bug. keccak with -O0 works but not with optimization - if target.device.startswith("METAL"): return None + if isinstance(ren, MetalRenderer): return None # If d is a power of two this is not valid for signed ints! is_unsigned = x.vmin>=0 or x.dtype in dtypes.uints assert d>0, "Sign should have been taken out of divisor" @@ -293,11 +294,11 @@ def fast_idiv(target: Target, x: UOp, d: int, dont_cast=False) -> UOp|None: # before we try casting to a larger dtype (slow), we see if there are powers of two in d we can shift to make x smaller # use explicit Ops.CDIV (trunc) since the recursion assumes trunc semantics throughout if (largest_factor_of_two_in_d := (d & -d)) > 1: - if (ret:=fast_idiv(target, x.alu(Ops.CDIV, x.const_like(largest_factor_of_two_in_d)), + if (ret:=fast_idiv(ren, x.alu(Ops.CDIV, x.const_like(largest_factor_of_two_in_d)), d//largest_factor_of_two_in_d, dont_cast=True)) is not None: return ret if dont_cast: return None # promo_lattice needs to return an unsigned type if the type is unsigned - if dtypes.is_int(next_dtype := promo_lattice[x.dtype.scalar()][-1]) and is_dtype_supported(next_dtype, target): + if dtypes.is_int(next_dtype := promo_lattice[x.dtype.scalar()][-1]) and next_dtype in ren.supported_dtypes(): if m*vmin >= next_dtype.min and m*vmax <= next_dtype.max: return ((x.cast(next_dtype)*m) >> s).cast(x.dtype) if is_unsigned else ((x.cast(next_dtype)*m) >> s).cast(x.dtype) + (x<0).where(x.ufix(1), 0) return None @@ -384,7 +385,7 @@ def l2i(op: Ops, dt: DType, *uops:UOp): def rne(v: UOp, s) -> UOp: return shr(v, s) + ((shr(v, s - 1) & 1) & ((v & ((1 << (s - 1)) - 1)).ne(0).cast(v.dtype) | (shr(v, s) & 1))) -def f2f(v, fr:DType, to:DType): +def f2f(v, fr:DType, to:DType, sat=True): fs, fb, (fe, fm), ts, tb, (te, tm) = fr.bitsize, exponent_bias(fr), dtypes.finfo(fr), to.bitsize, exponent_bias(to), dtypes.finfo(to) # NB: denormals are zero! if fe <= te and fm < tm: @@ -399,7 +400,7 @@ def f2f(v, fr:DType, to:DType): 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) elif fe >= te and fm > tm: - v = f2f_clamp(v.bitcast(fr), to).bitcast(f2f_dt[fr]) + v = f2f_clamp(v.bitcast(fr), to, sat).bitcast(f2f_dt[fr]) sign, nosign = shr(v, fs - ts) & shl(1, ts - 1), v & (shl(1, fs - 1) - 1) norm = (rne(nosign, fm - tm) - shl(fb - tb, tm)).cast(f2f_dt[to]) underflow = (shr(v, fm) & (shl(1, fe) - 1)) < (1 + fb - tb) @@ -410,12 +411,12 @@ def f2f(v, fr:DType, to:DType): return is_nan.where(nan, sign.cast(f2f_dt[to]) | underflow.where(0, norm)) else: raise NotImplementedError(f"unsupported decomp {fr} -> {to}") -def f2f_clamp(val:UOp, dt:DType) -> UOp: +def f2f_clamp(val:UOp, dt:DType, sat=True) -> UOp: e, m = dtypes.finfo(dt) if dt in dtypes.fp8_fnuz: max_exp, max_man = (1 << e) - 1, (1 << m) - 1 else: max_exp, max_man = ((1 << e) - 1, (1 << m) - 2) if dt == dtypes.fp8e4m3 else ((1 << e) - 2, (1 << m) - 1) mx = val.const_like(2.0**(max_exp - exponent_bias(dt)) * (1.0 + max_man / (1 << m))) - sat = mx if dt in dtypes.fp8s else val.const_like(float('inf')) + sat = mx if dt in dtypes.fp8s and sat else val.const_like(float('inf')) # FIXME: CMPLT of nan is undefined return val.ne(val).where(val, (val < -mx).where(-sat, (mx < val).where(sat, val))) @@ -476,7 +477,7 @@ def get_late_rewrite_patterns(ops:tuple[Ops, ...], disable_fast_idiv:bool) -> Pa if (v:=powers_of_two.get(c.arg, 0)) else None)] if not disable_fast_idiv: # fast_idiv handles non-pow2: only fire on non-negative inputs (signed magic-mul is unreliable for x<0) - pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.ints), UPat.cvar("d", vec=False))), + pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.ints), UPat.cvar("d"))), lambda ctx, x, d: fast_idiv(ctx, x, d.arg) if x.vmin >= 0 or x.dtype in dtypes.uints else None)] # rewrite raw CMOD -> x - d*CDIV(x,d) so fast_idiv can pick up the CDIV. only on non-negative inputs; # avoids disturbing floormod_to_mod's general-path output (which uses a trunc Ops.CMOD as an implementation detail) @@ -492,7 +493,7 @@ def get_late_rewrite_patterns(ops:tuple[Ops, ...], disable_fast_idiv:bool) -> Pa ((UPat.cvar("c", dtypes.sints) < UPat.var("x", dtypes.sints)).logical_not(), lambda x,c: x x==c ] if Ops.CMPEQ in ops: pat += [(UPat.var('x').ne(UPat.var('y')).logical_not(), lambda x,y: x.alu(Ops.CMPEQ, y))] @@ -553,13 +554,13 @@ def get_late_rewrite_patterns(ops:tuple[Ops, ...], disable_fast_idiv:bool) -> Pa f2f_store(st, idx, val, *ctx) if val.dtype.scalar() == ctx[1] and (idx:=idx.src[0] if idx.op == Ops.CAST else idx).tag == ctx[0] else None), ]) -def do_dtype_decomps(sink:UOp, ctx:tuple[set[DType], Target]) -> UOp: - def _should_emulate(dt): return dt in EMULATED_DTYPES.tolist(dtypes) or not is_dtype_supported(dt, ctx[1]) +def do_dtype_decomps(sink:UOp, ctx:tuple[set[DType], Renderer]) -> UOp: + def _should_emulate(dt): return dt in EMULATED_DTYPES.tolist(dtypes) or dt not in ctx[1].supported_dtypes() for fr in sorted(filter(_should_emulate, ctx[0])): - if fr in dtypes.floats: - to = dtypes.half if not _should_emulate(dtypes.half) and fr in dtypes.fp8s else dtypes.float - sink = graph_rewrite(sink, pm_float_decomp, name=f"decomp {fr} -> {to}", ctx=(fr, to), bottom_up=True) - else: sink = graph_rewrite(sink, pm_long_decomp, name="decomp long -> int", bottom_up=True) + to = dtypes.int if fr == dtypes.long else dtypes.half if not _should_emulate(dtypes.half) and fr in dtypes.fp8s else dtypes.float + if DEBUG >= 2: print(f"emulating {fr} as {to}") + sink = graph_rewrite(sink, pm_float_decomp if fr in dtypes.floats else pm_long_decomp, name=f"decomp {fr} -> {to}", ctx=(fr, to), bottom_up=True) + ctx[0].clear() return sink pm_dtype_decomps = PatternMatcher([ diff --git a/tinygrad/uop/divandmod.py b/tinygrad/uop/divandmod.py index 752c4dc92fd4b..17544173894e2 100644 --- a/tinygrad/uop/divandmod.py +++ b/tinygrad/uop/divandmod.py @@ -65,7 +65,7 @@ def fold_divmod_general(d: UOp) -> UOp|None: # nest_by_factor: x//c -> (x//f)//(c//f), x%c -> (x//f%(c//f))*f + b where b=x%f # FLOORDIV identity holds for any sign of x; FLOORMOD reconstruction needs x.vmin>=0 results = [] - for div in {abs(f) for u, f in zip(uops_no_const, factors) if u.op not in (Ops.CONST, Ops.VCONST) and 1 < abs(f) < c and (c%f)==0}: + for div in {abs(f) for u, f in zip(uops_no_const, factors) if u.op is not Ops.CONST and 1 < abs(f) < c and (c%f)==0}: if (newxs := fold_divmod_general(x//div)) is not None: if d.op is Ops.FLOORDIV: results.append((len(newxs.backward_slice), newxs // (c // div))) @@ -108,7 +108,7 @@ def fold_divmod_general(d: UOp) -> UOp|None: # (x//c+a)//d -> (x+a*c)//(c*d) for c>0, d>0 ((UPat.var("x")//UPat.cvar("c") + UPat.cvar("a"))//UPat.cvar("d"), lambda x,c,a,d: (x+a*c)//(c*d) if c.vmin>0 and d.vmin>0 else None), # (x+c)//d -> (x+c%d)//d + c//d for d>0 (split out the multiple of d in the constant) - ((UPat.var("x", dtypes.weakint)+UPat.cvar("c", vec=False))//UPat.cvar("d", vec=False), + ((UPat.var("x", dtypes.weakint)+UPat.cvar("c"))//UPat.cvar("d"), lambda x,c,d: (x+c.arg%d.arg)//d + c.arg//d.arg if c.arg%d.arg!=c.arg and d.arg>0 else None), # ** 2. Slow Rules ** diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 5597593fac9ff..4b9ed82f6c0eb 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -9,7 +9,7 @@ from tinygrad.device import Buffer, MultiBuffer, canonicalize_device from tinygrad.helpers import ContextVar, all_int, prod, getenv, all_same, Context, partition, temp, unwrap, T, argfix, Metadata, flatten, TRACEMETA from tinygrad.helpers import PROFILE, dedup, cdiv, cmod, floordiv, floormod, diskcache_put, to_function_name, cpu_profile, TracingKey -from tinygrad.helpers import VIZ, SPEC, CAPTURE_PROCESS_REPLAY +from tinygrad.helpers import VIZ, SPEC, CAPTURE_PROCESS_REPLAY, DISALLOW_BROADCAST from tinygrad.helpers import colored, ansilen, printable if TYPE_CHECKING: from tinygrad.renderer import Estimates @@ -18,6 +18,19 @@ class AxisType(Enum): def __repr__(self): return str(self) GLOBAL = auto(); WARP = auto(); LOCAL = auto(); LOOP = auto(); GROUP_REDUCE = auto(); REDUCE = auto(); UPCAST = auto(); UNROLL = auto() # noqa: E702 THREAD = auto(); PLACEHOLDER = auto() # noqa: E702 + +@dataclass(frozen=True, order=True) +class ParamArg: + slot: int + vmin_vmax: tuple[PyConst, PyConst]|None = None + name: str|None = None + addrspace: AddrSpace = AddrSpace.GLOBAL + axis: int|None = None + device: str|tuple[str, ...]|None = None + def __repr__(self): + fields = (("vmin_vmax", None), ("name", None), ("addrspace", AddrSpace.GLOBAL), ("axis", None), ("device", None)) + args = [str(self.slot)] + [f"{k}={v!r}" for k,default in fields if (v:=getattr(self, k)) != default] + return f"ParamArg({', '.join(args)})" axis_letters = {AxisType.GLOBAL: "g", AxisType.THREAD: "t", AxisType.LOCAL: "l", AxisType.WARP: "w", AxisType.LOOP: "L", AxisType.UPCAST: "u", AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r"} axis_colors = {AxisType.GLOBAL: "blue", AxisType.THREAD: "BLUE", AxisType.LOCAL: "cyan", AxisType.WARP: "CYAN", AxisType.LOOP: "WHITE", @@ -28,7 +41,7 @@ def __repr__(self): return str(self) AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5} range_start = {Ops.STAGE: 1, Ops.REDUCE: 1, Ops.WMMA: 3, Ops.END: 1, Ops.CALL: 1, Ops.FUNCTION: 1, - Ops.COPY: 2, Ops.BUFFER_VIEW: 1, Ops.LINEAR: 0} + Ops.COPY: 2, Ops.SLICE: 2, Ops.LINEAR: 0} # https://en.wikipedia.org/wiki/Identity_element def identity_element(op:Ops, dt:DType) -> PyConst: return dt.const({Ops.ADD:0, Ops.MUL:1, Ops.MAX:dt.min}[op]) @@ -51,7 +64,11 @@ def _align_left(*shapes:tuple[sint, ...]) -> tuple[tuple[sint, ...], ...]: max_dim = max(len(s) for s in shapes) return tuple((1,)*(max_dim-len(s))+s for s in shapes) def _broadcast_shape(*shapes:tuple[sint, ...]) -> tuple[sint, ...]: - return tuple(0 if 0 in nth_dim_sizes else smax(nth_dim_sizes) for nth_dim_sizes in zip(*_align_left(*shapes))) + shaped_aligned_left = _align_left(*shapes) + ret = tuple(0 if 0 in nth_dim_sizes else smax(nth_dim_sizes) for nth_dim_sizes in zip(*shaped_aligned_left)) + if not all(resolve(s == ns) or resolve(s == 1) for shape in shaped_aligned_left for s,ns in zip(shape, ret)): + raise IndexError(f"shape mismatch: objects cannot be broadcast to a single shape {shapes}") + return ret def ssimplify(uop:sint): return uop.ssimplify() if isinstance(uop, UOp) else uop def sym_infer(uop: UOp|int, var_vals: dict[str, int]) -> int: return uop.sym_infer(var_vals) if isinstance(uop, UOp) else uop @@ -101,7 +118,7 @@ def __call__(cls, op:Ops, dtype:DType=dtypes.void, src:tuple[UOp,...]=tuple(), a return created # some uops map to other stuff -buffers:weakref.WeakKeyDictionary[UOp, Buffer|MultiBuffer] = weakref.WeakKeyDictionary() # this maps BUFFER/BUFFER_VIEW uops to their device Buffers +buffers:weakref.WeakKeyDictionary[UOp, Buffer|MultiBuffer] = weakref.WeakKeyDictionary() # this maps BUFFER/SLICE uops to their device Buffers all_metadata:weakref.WeakKeyDictionary[UOp, tuple[Metadata, ...]] = weakref.WeakKeyDictionary() # TODO: should this be here? # recursive_property replaces functools.cached_property in recursive UOp functions to prevent RecursionError @@ -209,10 +226,14 @@ def _shape(self) -> tuple[sint, ...]|None: match self.op: # late ops don't have shape case Ops.UNIQUE | Ops.LUNIQUE | Ops.DEVICE | Ops.IF | Ops.BARRIER | Ops.CUSTOM | Ops.CUSTOMI | \ - Ops.CONTRACT | Ops.SINK | Ops.END | Ops.REWRITE_ERROR | Ops.PTRCAT | Ops.ENDIF | \ - Ops.LINEAR | Ops.PROGRAM | Ops.SOURCE | Ops.BINARY | Ops.INS | Ops.TUPLE | Ops.CALL | Ops.FUNCTION: + Ops.SINK | Ops.END | Ops.REWRITE_ERROR | Ops.PTRCAT | Ops.ENDIF | \ + Ops.LINEAR | Ops.PROGRAM | Ops.SOURCE | Ops.INS | Ops.TUPLE | Ops.CALL | Ops.FUNCTION: return None + # special (terrible) case for RESHAPE on NOOP + case Ops.RESHAPE: + if self.src[0].op is Ops.NOOP: return self.marg + # hacks for NOOP case Ops.NOOP: return self.src[0]._shape if len(self.src) >= 1 else None @@ -229,6 +250,8 @@ def _shape(self) -> tuple[sint, ...]|None: return inner_shape case Ops.CAST: + # if it has a vec dtype, set the shape + if self.dtype.count > 1: return (self.dtype.count,) # when PTX casts from ptr to non ptr, remove the shape of the buffer if isinstance(self.src[0].dtype, PtrDType) and not isinstance(self.src[0].dtype, ImageDType) and not isinstance(self.dtype, PtrDType): return () @@ -238,36 +261,52 @@ def _shape(self) -> tuple[sint, ...]|None: for s in self.src[1:]: shp.extend(list(s.shape)) return tuple(shp) + self.src[0].shape[len(self.src[1:]):] - # TODO: these should have the shape of the dtype.count - case Ops.CONST | Ops.DEFINE_VAR: return () - case Ops.GEP | Ops.STACK | Ops.VCONST | Ops.VCAT: return () + case Ops.GEP: + return (len(self.arg),) if len(self.arg) > 1 else () + case Ops.STACK: + if len(self.src) == 0: return () + if isinstance(self.dtype, PtrDType): + # TODO: this is broken + return self.src[0].shape + else: + return (len(self.src),) + self.src[0].shape + # TODO: contract and unroll should be deleted + case Ops.CONST | Ops.DEFINE_VAR | Ops.CONTRACT | Ops.UNROLL | Ops.VCAT: + return (self.dtype.count,) if self.dtype.count > 1 else () # some ops init the shape - case Ops.BIND | Ops.RANGE | Ops.SPECIAL | Ops.UNROLL: return () + case Ops.GETADDR: return () + case Ops.BIND | Ops.RANGE | Ops.SPECIAL: return () + case Ops.BINARY: return (len(self.arg),) case Ops.BUFFER: return (self.arg,) - case Ops.BUFFER_VIEW: - # HACK: BUFFER_VIEW is used inside kernels, so we set the shape to () if it's on an INDEX + case Ops.SLICE: + # HACK: SLICE is used inside kernels, so we set the shape to () if it's on an INDEX if self.src[0].op is Ops.INDEX: return () - return (self.arg[0],) + return (self.arg,) case Ops.CUSTOM_FUNCTION: return None - case Ops.STAGE: return tuple([int(r.vmax+1) for r in self.src[1:]]) - case Ops.DEFINE_LOCAL | Ops.DEFINE_REG: return (self.ptrdtype.size,) + case Ops.STAGE: + # STAGE adds the existing shape to the front, opposite of INDEX + return tuple([int(r.vmax+1) for r in self.src[1:]])+self.src[0].shape + case Ops.DEFINE_LOCAL | Ops.DEFINE_REG: + if isinstance(self.dtype, PtrDType): + return (self.ptrdtype.size, self.dtype.count) if self.dtype.count > 1 else (self.ptrdtype.size,) + return (self.dtype.count,) if self.dtype.count > 1 else () case Ops.PARAM: + if isinstance(self.dtype, ImageDType): return self.dtype.shape if isinstance(self.dtype, PtrDType): return (self.ptrdtype.size,) - # NOTE: copied from marg - if len(self.src) >= 1: return tuple(self.src[0].sgep(i) for i in range(self.src[0].dtype.count)) - return None + return tuple(self.src[0].sgep(i) for i in range(self.src[0].dtype.count)) if len(self.src) >= 1 else None # wmma output shape = accumulator shape (src[2]) case Ops.WMMA | Ops.SHAPED_WMMA: return self.src[2]._shape # passthrough ops - case Ops.MSTACK | Ops.MSELECT | Ops.DETACH | Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.AFTER | Ops.PATCH | Ops.LOAD | \ + case Ops.MSTACK | Ops.MSELECT | Ops.DETACH | Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.AFTER | Ops.LOAD | \ Ops.COPY | Ops.ALLREDUCE: return self.src[0]._shape # REDUCE with empty axis is passthrough (lowered form) case Ops.REDUCE if len(self.arg[1]) == 0: - return self.src[0]._shape + # these can mismatch if there's a horizonal reduce + return (self.dtype.count,) if self.dtype.count > 1 else () # TODO: disallow shape changing bitcast case Ops.BITCAST: @@ -277,7 +316,7 @@ def _shape(self) -> tuple[sint, ...]|None: return ps[:-1]+(ssimplify((ps[-1]*input_sz) // output_sz),) if len(ps) > 0 else ps return ps - # MULTI marker (axis info in PARAM sources) has no shape + # MULTI marker has no shape case Ops.MULTI if len(self.src) == 0: return None # movement ops change the shape @@ -299,31 +338,36 @@ def _shape(self) -> tuple[sint, ...]|None: return tuple(ps[i] for i in self.marg) case Ops.PAD: # TODO: why do i need resolve here? - if len(ps) != len(self.marg) or not all(resolve(b>=0) and resolve(e>=0) for b,e in self.marg): raise ValueError(f"invalid pad {self.marg}") - return tuple(ssimplify(s+b+e) for s,(b,e) in zip(ps, self.marg)) + if len(ps) != len(self.marg) or not all(resolve(sz>=0) and resolve(0<=o) and resolve(o+s<=sz) for s,(o,sz) in zip(ps, self.marg)): + raise ValueError(f"invalid pad {self.marg} for {ps}") + return tuple(ssimplify(sz) for _,sz in self.marg) case Ops.SHRINK: # TODO: why do i need resolve here? - if len(ps) != len(self.marg) or not all(resolve(0<=b) and resolve(b<=e) and resolve(e<=s) for s,(b,e) in zip(ps, self.marg)): + if len(ps) != len(self.marg) or not all(resolve(0<=o) and resolve(sz>=0) and resolve(o+sz<=s) for s,(o,sz) in zip(ps, self.marg)): raise ValueError(f"invalid shrink {self.marg} for {ps}") - return tuple(ssimplify(e-s) for s,e in self.marg) + return tuple(ssimplify(sz) for _,sz in self.marg) case Ops.FLIP: if len(ps) != len(self.marg) or not all(isinstance(x, bool) for x in self.marg): raise ValueError(f"bad flip on {ps}, {self.marg}") return ps case Ops.MULTI: return tuple(s*len(self.device) if a == self.axis else s for a,s in enumerate(ps)) case Ops.REDUCE: - axis_arg = self.arg[1] if self.op is Ops.REDUCE else self.arg[7] + axis_arg = self.arg[1] if not isinstance(axis_arg, tuple) or not all(isinstance(x, int) and x>=0 and x 0 and all(x is not None for x in input_shapes), f"None input shape not supported for {self.op}" - # TODO: add broadcasting here - if not all_same(input_shapes): + if DISALLOW_BROADCAST and not all_same(input_shapes): raise RuntimeError(f"shape mismatch at {self.op}: {input_shapes} {[x.op for x in self.src]}") - return input_shapes[0] + # broadcasting lives in _shape property now + return _broadcast_shape(*input_shapes) # all Ops must be explicitly handled raise NotImplementedError(f"no shape handling for {self.op} with {self.dtype}") @@ -335,6 +379,7 @@ def shape(self) -> tuple[sint, ...]: @property def max_shape(self) -> tuple[int, ...]: return to_max_shape(self.shape) + def max_numel(self) -> int: return prod(self.max_shape) @property def shard_shape(self) -> tuple[sint, ...]: @@ -378,8 +423,8 @@ def ranges(self) -> dict[UOp, None]: # *** uop evaluation *** def simplify(self, tracked=False): - if self.op in {Ops.CONST, Ops.VCONST}: return self - if self.op is Ops.SINK and all(s.op in {Ops.CONST, Ops.VCONST} or (s.op is Ops.STACK and len(s.src) == 0) for s in self.src): return self + if self.op is Ops.CONST: return self + if self.op is Ops.SINK and all(s.op is Ops.CONST or (s.op is Ops.STACK and len(s.src) == 0) for s in self.src): return self # late import! from tinygrad.uop.symbolic import symbolic with Context(TRACK_MATCH_STATS=0 if not tracked else TRACK_MATCH_STATS.value): @@ -445,7 +490,7 @@ def __getitem__(self, idx): return self.index(*[UOp.const(dtypes.weakint, x) if isinstance(x, int) else x for x in idx]) def const_like(self, b:ConstLike, dtype:DType|None=None): # constants can optionally have a DEVICE source - ret = UOp.const(dtype or self.dtype.base, b, device=self._device, shape=self.shard_shape if self.axis is not None else self._shape) + ret = UOp.const(dtype or self.dtype.base, b, device=self.device, shape=self.shard_shape if self.axis is not None else self._shape) return ret.multi(self.axis) if self.axis is not None else ret def ufix(self, x): if isinstance(x, UOp): return x @@ -465,7 +510,6 @@ def gep(self, i:tuple[int, ...]|int): if isinstance(i, int): # NOTE: these are just shortcuts to not have to create and fold later if self.op is Ops.STACK: return self.src[i] - if self.op is Ops.VCONST: return UOp.const(self.dtype.scalar(), self.arg[i]) if self.op is Ops.CONST: return UOp.const(self.dtype.scalar(), self.arg) i = (i,) return UOp(Ops.GEP, self.dtype.scalar().vec(len(i)) if len(i) > 1 else self.dtype.scalar(), (self,), i) @@ -478,6 +522,7 @@ def wait(self, src:UOp|ConstType, **kwargs): def end(self, *src:UOp): return UOp(Ops.END, src=(self,)+src) if len(src) else self def after(self, *src:UOp, **kwargs): return UOp(Ops.AFTER, self.dtype, (self,)+src, **kwargs) if len(src) else self def barrier(self, *src:UOp): return UOp(Ops.BARRIER, src=(self,)+src) + def ins(self, arg, **kwargs): return UOp(Ops.INS, kwargs.pop("dtype", self.dtype), kwargs.pop("src", self.src), arg, kwargs.pop("tag", self.tag)) def contract(self, *rngs:UOp): assert all(x.arg[-1] == AxisType.UPCAST for x in rngs), "all contract ranges must be upcast" return UOp(Ops.CONTRACT, dtype=self.dtype.vec(prod([x.vmax+1 for x in rngs])), src=(self,), arg=tuple((x.arg[0], x.vmax+1) for x in rngs)) @@ -492,14 +537,16 @@ def alu(self, op, *src:UOp, **kwargs): return UOp(op, out_dtype, all_srcs, **kwargs) @staticmethod def const(dtype:DType, b:ConstLike, device:str|tuple[str, ...]|None=None, shape:tuple[sint, ...]|None=None): - if isinstance(b, UOp): return b.unbind()[0] if b.op is Ops.BIND else b + if isinstance(b, UOp): return b.cast(dtype) if isinstance(b, tuple) and all_same(b): assert len(b) > 0, "can't create const from empty tuple" - b = b[0] # doesn't have to be a VCONST if they are all the same - ret = UOp(Ops.VCONST if isinstance(b, tuple) else Ops.CONST, dtype, - arg=dtype.const(b), - src=(UOp(Ops.DEVICE, arg=device),) if device is not None else ()) - return ret.reshape((1,)*len(shape)).expand(shape) if shape is not None and ret.shape != shape else ret + b = b[0] # doesn't have to be a STACK if they are all the same + if isinstance(b, tuple): + stk = [UOp(Ops.CONST, dtype.scalar(), arg=dtype.const(c), src=(UOp(Ops.DEVICE, arg=device),) if device is not None else ()) for c in b] + ret = UOp.vectorize(*stk) + else: + ret = UOp(Ops.CONST, dtype, arg=dtype.const(b), src=(UOp(Ops.DEVICE, arg=device),) if device is not None else ()) + return ret.reshape((1,)*len(shape)).expand(shape) if shape is not None and shape != () and ret.shape != shape else ret @staticmethod def unique_const(fill_value:ConstType, dtype:DTypeLike|None=None, device:str|tuple[str, ...]|None=None, # type: ignore[override] shape:tuple[sint, ...]|None=None, unique=True): @@ -519,7 +566,7 @@ def _rop(self, op:Ops, axis:tuple[int, ...]): @staticmethod def invalid(count=1): return UOp(Ops.CONST, dtypes.weakint.vec(count), src=(), arg=Invalid) def valid(self, cond): - return self if cond.op is Ops.WHERE and cond.arg else cond.where(self.cast(dtypes.weakint), UOp.invalid(self.dtype.count)) + return cond.where(self.cast(dtypes.weakint), UOp.invalid(self.dtype.count)) def get_idx(self) -> UOp: assert self.dtype.scalar() is dtypes.weakint, "Can only call get_idx on index dtype" return self.src[1] if self.op is Ops.WHERE and self.src[2].arg is Invalid else self @@ -533,6 +580,7 @@ def reduce(self, *src:UOp, **kwargs): def contiguous(self, *args, **kwargs): if self.op is Ops.CONTIGUOUS: return self + if self.device is None: return self if self.has_buffer_identity(): return self return UOp(Ops.CONTIGUOUS, dtype=self.dtype, src=(self,)+args, **kwargs) def bufferize(self, *args, **kwargs): return UOp(Ops.STAGE, dtype=self.dtype, src=(self,)+args, **kwargs) @@ -546,6 +594,13 @@ def split_uop(self:UOp, sep:Ops) -> Iterator[UOp]: for s in self.src: yield from s.split_uop(sep) else: yield self + @property + def reg(self:UOp): + # TODO: add a way to access the nth element in src + if self.op in (Ops.NOOP, Ops.AFTER) and self.src: return self.src[0].reg + if isinstance(self.tag, tuple): return self.tag[0] + return self.tag + # *** multi-device helpers *** def multi(self, axis:int|None): @@ -567,11 +622,7 @@ def axis(self) -> int|None: if self.op is Ops.GETTUPLE: in_tuple = self.src[0].src[0] if self.src[0].op is Ops.FUNCTION else self.src[0] return in_tuple.src[self.arg].axis if in_tuple.op is Ops.TUPLE else None - # PARAM: axis is stored as a MULTI source - if self.op is Ops.PARAM: - for s in self.src: - if s.op is Ops.MULTI: return s.arg - return None + if self.op is Ops.PARAM: return self.arg.axis # NOTE: they all have to share an axis, we always choose [-1] if self.op in GroupOp.ALU: return axes[-1] if (axes := dedup([x.axis for x in self.src if x.axis is not None])) else None if len(self.src) == 0: return None @@ -594,14 +645,13 @@ def _unshard(self, axis:int) -> UOp: dnum = UOp.variable("_device_num", 0, dcount-1) return self.pad(tuple((0,0) if a != axis else (bsz*dnum, bsz*(dcount-1) - bsz*dnum) for a in range(len(self.shape)))) - def _shard(self, axis:int) -> UOp: + def _shard(self, axis:int, dcount:int) -> UOp: if len(self.shape) == 0: return self # scalars broadcast, no sharding needed - dcount = len(self.device) dnum = UOp.variable("_device_num", 0, dcount-1) if self.shape[axis] % dcount != 0: raise RuntimeError(f"multi axis uneven: {self.shape[axis]=} {axis=} {dcount=}") sz = self.shape[axis] // dcount return self.shrink(tuple((0,s) if i != axis else (dnum*sz,dnum*sz+sz) for i,s in enumerate(self.shape))) - def shard(self, devices:tuple[str, ...], axis:int) -> UOp: return self.copy_to_device(devices)._shard(axis).multi(axis) + def shard(self, devices:tuple[str, ...], axis:int) -> UOp: return self.copy_to_device(devices)._shard(axis, len(devices)).multi(axis) def copy_to_device(self, device:str|tuple[str, ...]|UOp, arg=None): assert arg is None or isinstance(self.device, tuple) @@ -631,7 +681,6 @@ def multibase(self) -> UOp: def sgep(self, i:int) -> sint: match self.op: case Ops.CONST: return self.arg - case Ops.VCONST: return self.arg[i] case Ops.STACK: return self.src[i].sintify() case _: raise RuntimeError(f"no sgep on {self.op}") @@ -683,20 +732,40 @@ def empty_like(self, dtype:DTypeLike|None=None, device:str|tuple[str, ...]|None= device = canonicalize_device(self.device if device is None else device) axis = self.axis if isinstance(device, tuple) else None return UOp.empty(self.shard_shape if axis is not None else self.shape, self.dtype if dtype is None else dtype, device, axis) - @property - def device(self) -> str|tuple[str, ...]: return unwrap(self._device) + def clone(self, device=None) -> UOp: + device = device or self.device + ret = self.empty_like(device=device) + src = self if self.device is None or self.device == device else self.copy_to_device(device) + return ret.after(ret.store(src)) @recursive_property - def _device(self) -> str|tuple[str, ...]|None: + def device(self) -> str|tuple[str, ...]|None: + if self.op is Ops.PARAM: return self.arg.device if self.op is Ops.DEVICE: return self.arg if self.op is Ops.STAGE: return self.arg.device - if self.op is Ops.AFTER: return self.src[0]._device + if self.op is Ops.AFTER: return self.src[0].device if self.op is Ops.MSELECT: assert isinstance(self.src[0].device, tuple), f"mselect must be on tuple device, getting {self.src[0].device}" return self.src[0].device[self.arg] if self.op is Ops.MSTACK: return tuple(cast(str, x.device) for x in self.src) if self.op in {Ops.COPY, Ops.BUFFER, Ops.ALLREDUCE}: return self.src[1].device for x in self.src: - if x._device is not None: return x._device + if x.device is not None: return x.device + return None + @recursive_property + def addrspace(self) -> AddrSpace|None: + if self.op is Ops.PARAM: return self.arg.addrspace + if self.op is Ops.BUFFER: return AddrSpace.GLOBAL + if self.op is Ops.DEFINE_LOCAL: return AddrSpace.LOCAL + if self.op is Ops.DEFINE_REG: return AddrSpace.REG + # LOAD brings things into registers + if self.op is Ops.LOAD: return AddrSpace.REG + if self.op in {Ops.INDEX, Ops.CAST, Ops.AFTER, Ops.REDUCE, Ops.GEP}: + return self.src[0].addrspace + if self.op in GroupOp.Movement: return self.src[0].addrspace + if self.op is Ops.STACK or self.op in GroupOp.Elementwise: + ad = [x.addrspace for x in self.src if x.addrspace is not None] + if not len(ad) or not all_same(ad): return None + return ad[0] return None @property def buf_uop(self) -> UOp: @@ -728,7 +797,7 @@ def has_buffer_identity(self): """Check if this UOp has a concrete buffer identity in the graph (RESHAPE/MULTI -> BUFFER chain).""" if self.op in {Ops.RESHAPE, Ops.MULTI}: return self.src[0].has_buffer_identity() if self.op is Ops.GETTUPLE and self.src[0].op is Ops.TUPLE: return self.src[0].src[self.arg].has_buffer_identity() - return self.op in {Ops.BUFFER, Ops.BUFFER_VIEW, Ops.PARAM} + return self.op in {Ops.BUFFER, Ops.SLICE, Ops.PARAM} def _base_buffer_is_realized(self) -> bool: """Walk through AFTER chain to find if the underlying buffer is realized (has allocated memory).""" @@ -750,16 +819,17 @@ def buffer(self) -> Buffer|MultiBuffer: buf = self.src[0].buffer assert isinstance(buf, Buffer), "must be a Buffer for BITCAST" return buf.view(prod(self.max_shape), self.dtype, 0) - if self.op is Ops.BUFFER_VIEW: + if self.op is Ops.SLICE: if (cret:=buffers.get(self)) is not None: return cret buf = self.src[0].buffer + offset = self.src[1].arg if isinstance(buf, MultiBuffer): mbuf = MultiBuffer.__new__(MultiBuffer) - mbuf.bufs = [b.view(self.arg[0], self.dtype, self.arg[1] * self.dtype.itemsize) for b in buf.bufs] + mbuf.bufs = [b.view(self.arg, self.dtype, offset * self.src[0].dtype.itemsize) for b in buf.bufs] buffers[self] = mbuf return mbuf - assert isinstance(buf, Buffer), "must be a Buffer for BUFFER_VIEW" - buffers[self] = bv = buf.view(self.arg[0], self.dtype, self.arg[1] * self.dtype.itemsize) + assert isinstance(buf, Buffer), "must be a Buffer for SLICE" + buffers[self] = bv = buf.view(self.arg, self.dtype, offset * self.src[0].dtype.itemsize) return bv if self.op is Ops.MSELECT: ret = self.src[0].buffer @@ -827,14 +897,16 @@ def const_factor(self) -> int: """largest known int that divides self""" # TODO: for negatives it's not the largest if self.op is Ops.CONST: return self.arg - if self.op is Ops.VCONST: return math.gcd(*self.arg) + if self.op is Ops.STACK: return math.gcd(*[x.const_factor() for x in self.src]) if self.op is Ops.ADD: return math.gcd(self.src[0].const_factor(), self.src[1].const_factor()) if self.op is Ops.MUL: return self.src[0].arg if self.src[0].op is Ops.CONST else self.src[1].arg if self.src[1].op is Ops.CONST else 1 return 1 def divides(self, v:int) -> UOp|None: if v==1: return self if self.op is Ops.CONST: return self.const_like(self.arg//v) if self.arg%v == 0 else None - if self.op is Ops.VCONST: return self.const_like(tuple(x//v for x in self.arg)) if all(x%v == 0 for x in self.arg) else None + if self.op is Ops.STACK: + srcs = tuple(s.divides(v) for s in self.src) + return None if any(s is None for s in srcs) else UOp(Ops.STACK, self.dtype, cast(tuple[UOp, ...], srcs)) if self.op is Ops.ADD: return d0+d1 if (d0:=self.src[0].divides(v)) is not None and (d1:=self.src[1].divides(v)) is not None else None if self.op is Ops.MUL: if (d0:=self.src[0].divides(v)) is not None: return d0 * self.src[1] @@ -902,13 +974,12 @@ def _min_max(self) -> tuple[PyConst, PyConst]: # float has NAN issue and we use explicit NAN in transcendental if self.op is Ops.WHERE and dtypes.is_int(self.dtype): return min(self.src[1].vmin, self.src[2].vmin), max(self.src[1].vmax, self.src[2].vmax) # NOTE: returned UOp is assumed to be CONST - if self.op is Ops.PARAM and len(self.src) >= 4: return self.src[2].arg, self.src[3].arg + if self.op is Ops.PARAM and self.arg.vmin_vmax is not None: return self.arg.vmin_vmax if self.op is Ops.DEFINE_VAR and self.arg: return self.arg[1], self.arg[2] if self.op in (Ops.RANGE, Ops.SPECIAL): return 0, (self.src[0]-1).vmax if self.op is Ops.BIND: return self.src[0]._min_max # ignore the bound value if self.op in {Ops.UNROLL, Ops.STACK}: return min(x.vmin for x in self.src), max(x.vmax for x in self.src) if self.op is Ops.CONST and self.arg is not Invalid: return self.arg, self.arg - if self.op is Ops.VCONST and Invalid not in self.arg: return (min(self.arg), max(self.arg)) if self.op is Ops.GEP: return self.src[0]._min_max # TODO: CAST to bool/unsigned is not monotone, still some case can be simplified if self.op is Ops.CAST and self.dtype in dtypes.floats+dtypes.sints+(dtypes.weakint,): @@ -948,7 +1019,8 @@ def pyrender(self): @staticmethod def placeholder(shape:tuple[int, ...], dtype:DType, slot:int, addrspace=AddrSpace.GLOBAL): lookup = {AddrSpace.GLOBAL: Ops.PARAM, AddrSpace.LOCAL: Ops.DEFINE_LOCAL, AddrSpace.REG: Ops.DEFINE_REG} - ret = UOp(lookup[addrspace], dtype.ptr(prod(shape), addrspace), arg=slot) + arg = ParamArg(slot, addrspace=addrspace) if addrspace is AddrSpace.GLOBAL else slot + ret = UOp(lookup[addrspace], dtype.ptr(prod(shape), addrspace), arg=arg) if len(shape) > 1: ret = ret.reshape(shape) return ret def placeholder_like(self, slot:int): @@ -961,21 +1033,20 @@ def set(self:UOp, val:UOp|ConstType, end:UOp|tuple[UOp, ...]|list[UOp]=()) -> UO # TODO: this should replace placeholder @staticmethod - def param(slot:int, dtype:DType, shape:tuple[sint, ...]|None=None, device=None, vmin_vmax:tuple[PyConst, PyConst]|None=None, name=None): - src: tuple[UOp, ...] = (UOp(Ops.NOOP) if shape is None else shape_to_shape_arg(shape),) + \ - (UOp(Ops.NOOP) if device is None else UOp(Ops.DEVICE, arg=device),) - if vmin_vmax is not None: src += (UOp.const(dtype, vmin_vmax[0]), UOp.const(dtype.scalar(), vmin_vmax[1])) - if name is not None: src += (UOp(Ops.NOOP, arg=name),) - return UOp(Ops.PARAM, dtype, src, arg=slot) + def param(slot:int, dtype:DType, shape:tuple[sint, ...]|None=None, device=None, vmin_vmax:tuple[PyConst, PyConst]|None=None, name=None, + addrspace=AddrSpace.GLOBAL, axis:int|None=None): + if shape is not None and axis is not None and isinstance(device, tuple): + shape = tuple(s*len(device) if i == axis else s for i,s in enumerate(shape)) + src: tuple[UOp, ...] = (UOp(Ops.NOOP) if shape is None else shape_to_shape_arg(shape),) + return UOp(Ops.PARAM, dtype, src, arg=ParamArg(slot, vmin_vmax, name, addrspace, axis, device)) def param_like(self, slot:int): + addrspace = self.addrspace if isinstance(self.dtype, (PtrDType, ImageDType)) else AddrSpace.GLOBAL if self.op is Ops.BIND: - return UOp.param(slot, self.dtype, self._shape, self._device, self._min_max, self.src[0].arg[0]) - p = UOp.param(slot, self.dtype, self._shape, self._device) - if self.axis is not None: p = p.replace(src=p.src + (UOp(Ops.MULTI, arg=self.axis),)) - return p + return UOp.param(slot, self.dtype, self._shape, self.device, cast(tuple[int, int], self._min_max), self.src[0].arg[0], addrspace) + return UOp.param(slot, self.dtype, self.shard_shape if self.axis is not None else self._shape, self.device, addrspace=addrspace, axis=self.axis) # opaque bodies stay as Ops.CALL; value-producing bodies become Ops.FUNCTION (wrapped in TUPLE) - _OPAQUE_CALL_BODIES = {Ops.SINK, Ops.PROGRAM, Ops.LINEAR, Ops.COPY, Ops.BUFFER_VIEW, Ops.CUSTOM_FUNCTION} + _OPAQUE_CALL_BODIES = {Ops.SINK, Ops.PROGRAM, Ops.LINEAR, Ops.COPY, Ops.SLICE, Ops.CUSTOM_FUNCTION} def call(self, *srcs:UOp, grad_fxn:Callable|None=None, metadata:tuple[Metadata, ...]=(), name:str|None=None, precompile:bool=False, precompile_backward:bool=False) -> UOp: assert len(self.ranges) == 0, f"ranges {self.ranges} are leaking out of the call in {self.pyrender()}" @@ -1036,10 +1107,10 @@ def from_sink(sink:UOp, aux:tuple=()) -> ProgramInfo: local_size: list[int]|None = [1, 1, 1] for u in sink.toposort(): if u.op is Ops.DEFINE_VAR: _vars.append(u) - if u.op is Ops.PARAM: _globals.append(u.arg) + if u.op is Ops.PARAM: _globals.append(u.arg.slot) if u.op in (Ops.STORE, Ops.LOAD): if (idx:=u.src[0]).op is Ops.INDEX or (u.src[0].op is Ops.CAST and (idx:=u.src[0].src[0]).op is Ops.INDEX): - if (buf:=idx.src[0]).op is Ops.PARAM: (outs if u.op is Ops.STORE else ins).append(buf.arg) + if (buf:=idx.src[0]).op is Ops.PARAM: (outs if u.op is Ops.STORE else ins).append(buf.arg.slot) if u.op is Ops.SPECIAL: if u.arg[0] == 'i': local_size = None special_size = local_size if u.arg[0] == 'l' else global_size @@ -1111,13 +1182,15 @@ def get_location() -> tuple[str, int]: return frm.f_code.co_filename, frm.f_lineno class UPat(OpMixin): - __slots__ = ("op", "match_dtype", "arg", "name", "src", "is_any") + __slots__ = ("op", "match_dtype", "match_tag", "arg", "name", "src", "is_any") def __init__(self, op:Ops|tuple[Ops, ...]|set[Ops]|None=None, dtype:DType|tuple[DType, ...]|set[DType]|None=None, src:tuple[UPat, ...]|list[UPat]|UPat|None=None, arg:Any=None, - name:str|None=None, allow_any_len:bool=False, custom_early_reject:set[Ops]|None=None, location=None, is_any:bool=False): + name:str|None=None, allow_any_len:bool=False, custom_early_reject:set[Ops]|None=None, location=None, is_any:bool=False, + tag:Any=None): assert op is None or isinstance(op, (Ops, tuple, set)), f"op must be Ops or tuple of Ops, not {op!r}" self.op: tuple[Ops, ...]|None = (op,) if isinstance(op, Ops) else (tuple(op) if isinstance(op, set) else op) self.match_dtype: tuple[DType, ...]|None = (dtype,) if isinstance(dtype, DType) else (tuple(dtype) if isinstance(dtype, set) else dtype) + self.match_tag: tuple[Any, ...]|None = (tag,) if isinstance(tag, str) else (tuple(tag) if isinstance(tag, set) else tag) self.arg, self.name, self._in_src, self.custom_early_reject = arg, name, src, custom_early_reject self.src: Any = None self.is_any = is_any @@ -1147,8 +1220,10 @@ def _check_dtype(self) -> None: pass def _ensure_float(self) -> UPat: return self def __reduce__(self): - return UPat, (self.op, self.match_dtype, self._in_src, self.arg, self.name, not self.strict_length, self.custom_early_reject, self.location) - def named(self, name:str): return UPat(self.op, self.match_dtype, self._in_src, self.arg, name, not self.strict_length, self.custom_early_reject) + return UPat, (self.op, self.match_dtype, self._in_src, self.arg, self.name, not self.strict_length, self.custom_early_reject, self.location, + self.is_any, self.match_tag) + def named(self, name:str): + return UPat(self.op, self.match_dtype, self._in_src, self.arg, name, not self.strict_length, self.custom_early_reject, tag=self.match_tag) @staticmethod def any(*src): return UPat(src=src, is_any=True) @@ -1161,10 +1236,9 @@ def or_after(self, name:str|None=None): def var(name:str|None=None, dtype:DType|tuple[DType, ...]|None=None): return UPat(dtype=dtype, name=name) @staticmethod @functools.cache - def cvar(name:str|None=None, dtype:DType|tuple[DType, ...]|None=None, vec=True, arg=None): - return UPat((Ops.CONST,Ops.VCONST) if vec else Ops.CONST, dtype, name=name, arg=arg) + def cvar(name:str|None=None, dtype:DType|tuple[DType, ...]|None=None, arg=None): return UPat(Ops.CONST, dtype, name=name, arg=arg) @staticmethod - def const(dtype:DType|tuple[DType, ...]|None, b:ConstType): return UPat(Ops.CONST, dtype=dtype, arg=b) + def const(dtype:DType|tuple[DType, ...]|None, b:ConstType, device=None): return UPat(Ops.CONST, dtype=dtype, arg=b) # lil helper def f(self, op, **kwargs): return UPat(op, src=(self,), **kwargs) @@ -1209,6 +1283,7 @@ def match(self:UPat, uop:UOp, store:dict[str, UOp]) -> list[dict[str, UOp]]: (self.name is not None and store.setdefault(self.name, uop) is not uop) or \ (self.match_dtype is not None and uop.dtype not in self.match_dtype and uop.dtype.scalar() not in self.match_dtype) or \ (self.arg is not None and self.arg != uop.arg) or \ + (self.match_tag is not None and uop.tag not in self.match_tag) or \ (len(uop.src) < self.required_len) or \ (self.strict_length and len(uop.src) != self.required_len): return [] if self.src is None: return [store] @@ -1549,8 +1624,7 @@ def select_dtype(u): return (dtypes.long if u.overflows(dtypes.int32) else dtype # There are no Unary ops at this point in symbolic, those are introduced later (UPat(GroupOp.Binary, name="u", src=(UPat.var("x").cast(dtypes.weakint), UPat.var("y").cast(dtypes.weakint))), lambda u,x,y: x.cast(dt:=least_upper_dtype(select_dtype(u), x.dtype, y.dtype)).alu(u.op, y.cast(dt)).cast(u.dtype)), - (UPat((Ops.CONST, Ops.VCONST), dtype=dtypes.weakint, name="u"), - lambda u: u.replace(dtype=select_dtype(u)).cast(u.dtype) if u.arg!=Invalid else None), + (UPat(Ops.CONST, dtype=dtypes.weakint, name="u"), lambda u: u.replace(dtype=select_dtype(u)).cast(u.dtype) if u.arg!=Invalid else None), (UPat(Ops.WHERE, dtypes.weakint, src=(UPat.var("cond"), UPat.var("x").cast(dtypes.weakint), UPat.var("y").cast(dtypes.weakint))), lambda cond,x,y: cond.where(x.cast(dt:=least_upper_dtype(x.dtype, y.dtype)), y.cast(dt)).cast(dtypes.weakint)), (UPat(Ops.RANGE, src=(UPat.var("end").cast(dtypes.weakint)), name="r"), lambda r,end: r.replace(dtype=end.dtype, src=(end,)).cast(dtypes.weakint)), @@ -1566,13 +1640,21 @@ def select_dtype(u): return (dtypes.long if u.overflows(dtypes.int32) else dtype (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx", dtypes.ints).cast()),), lambda buf,idx: buf.index(idx, ptr=True)), (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("gate").where(UPat.var("idx", dtypes.ints).cast(), UPat(Ops.CONST, arg=Invalid)))), lambda buf,idx,gate: buf.index(gate.where(idx, idx.const_like(Invalid)), ptr=True)), + # remove hanging casts for images + (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx_y", dtypes.ints).cast(), UPat.var("idx_x", dtypes.ints).cast()),), + lambda buf,idx_x,idx_y: buf.index(idx_y, idx_x, ptr=True)), + (UPat(Ops.INDEX, src=(UPat.var("buf"), + UPat.var("gate").where(UPat.var("idx_y", dtypes.ints).cast(), UPat(Ops.CONST, arg=Invalid)), + UPat.var("gate").where(UPat.var("idx_x", dtypes.ints).cast(), UPat(Ops.CONST, arg=Invalid)))), + lambda buf,idx_x,idx_y,gate: buf.index(gate.where(idx_y, idx_y.const_like(Invalid)), + gate.where(idx_x, idx_x.const_like(Invalid)), ptr=True)), (UPat((Ops.SINK, Ops.NOOP, Ops.END), name="n"), lambda n: n.replace(src=tuple(s.src[0] if s.op is Ops.CAST and s.dtype == dtypes.weakint else s for s in n.src))), ]) def _index_to_concrete_int(u:UOp) -> UOp: return graph_rewrite(u.sink(), pm_lower_index_dtype).src[0] _substitute = PatternMatcher([(UPat(tuple(Ops), name="x"), lambda ctx,x: ctx.get(x,None))]) -_pm_resolve_params = PatternMatcher([(UPat(Ops.PARAM, name="p"), lambda ctx,p: ctx[p.arg])]) +_pm_resolve_params = PatternMatcher([(UPat(Ops.PARAM, name="p"), lambda ctx,p: ctx[p.arg.slot])]) remove_all_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)]) def gate_kernel_sink(x:UOp) -> bool: diff --git a/tinygrad/uop/render.py b/tinygrad/uop/render.py index bed89469fff30..8b87605299c86 100644 --- a/tinygrad/uop/render.py +++ b/tinygrad/uop/render.py @@ -34,10 +34,10 @@ def strip_binary_parens(x:UOp, left:str, right:str, code_for_op) -> str: renderer = PatternMatcher([ (UPat((Ops.DEFINE_VAR,), name="x"), lambda x: x.expr), - (UPat(Ops.PARAM, src=(UPat(), UPat(), UPat(), UPat(), UPat(Ops.NOOP, name="x"))), lambda x: x.arg), + (UPat(Ops.PARAM, name="x"), lambda x: x.arg.name if x.arg.name is not None else f"p{x.arg.slot}"), (UPat((Ops.SPECIAL), name="x"), lambda x: x.arg), (UPat(Ops.RANGE, name="x"), lambda x: f"r{range_str(x)}"), - (UPat((Ops.CONST, Ops.VCONST), name="x"), lambda x: str(x.arg)), + (UPat(Ops.CONST, name="x"), lambda x: str(x.arg)), (UPat(Ops.UNROLL, name="x"), lambda ctx,x,u: f"UNROLL({ctx[x.src[0]]}, {u.arg})"), (UPat(Ops.CAST, name="x"), lambda ctx,x: f"({str(x.dtype)[7:]})({ctx[x.src[0]]})"), (UPat(Ops.BIND, name="x"), lambda ctx,x: ctx[x.src[0]]), @@ -108,10 +108,10 @@ def render_marg(ctx,x:UOp): (UPat(Ops.CMOD, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.alu(Ops.CMOD, {ctx[x.src[1]]})"), # NOTE: only match CONSTs without UNIQUE (len(src)==1), unique_const needs explicit rendering (UPat(set(syms.keys())-{Ops.SUB, Ops.CMPNE, Ops.CDIV, Ops.CMOD}, src=(UPat(Ops.CONST, src=(UPat(Ops.DEVICE),), name="y"), UPat(name="z")), - name="x"), lambda ctx,x,y,z: strip_binary_parens(x, str(y.arg), ctx[z], lambda a,b: f"({a}{syms[x.op]}{b})")), + name="x"), lambda ctx,x,y,z: strip_binary_parens(x, str(y.arg), ctx[z], lambda a,b: f"({a}{syms[x.op]}{b})") if y.device==z.device else None), # NOTE: sub doesn't work cause it's written as add/mul (UPat(set(syms.keys())-{Ops.SUB, Ops.CDIV, Ops.CMOD}, src=(UPat(name="y"), UPat(Ops.CONST, src=(UPat(Ops.DEVICE),), name="z")), name="x"), - lambda ctx,x,y,z: strip_binary_parens(x, ctx[y], str(z.arg), lambda a,b: f"({a}{syms[x.op]}{b})")), + lambda ctx,x,y,z: strip_binary_parens(x, ctx[y], str(z.arg), lambda a,b: f"({a}{syms[x.op]}{b})") if y.device==z.device else None), (UPat(set(syms.keys())-{Ops.SUB, Ops.CDIV, Ops.CMOD}, name="x"), lambda ctx,x: strip_binary_parens(x, ctx[x.src[0]], ctx[x.src[1]], lambda a,b: f"({a}{syms[x.op]}{b})")), (UPat(sugar, src=(), name="x"), lambda x: f"UOp.{x.op.name.lower()}("+', '.join(([f'arg={repr(x.arg)}'] if x.arg is not None else []))+")"), @@ -146,7 +146,7 @@ def pyrender(ast:UOp) -> str: lst = list(ast.toposort()) cmap = consumer_map_from_toposort(lst) - not_rendered = {Ops.CONST, Ops.VCONST, Ops.DEVICE} + not_rendered = {Ops.CONST, Ops.DEVICE} always_rendered = {Ops.PARAM, Ops.LOAD, Ops.SPECIAL, Ops.RANGE, Ops.CONTIGUOUS, Ops.STACK, Ops.BUFFER, Ops.COPY, Ops.CALL, Ops.FUNCTION, Ops.WHERE, Ops.END} diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index d0a55cb9cf9f5..1836ab7a8c9fc 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -1,6 +1,6 @@ import math from typing import cast, Any -from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, AxisType, KernelInfo +from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, AxisType, KernelInfo, ParamArg from tinygrad.uop.render import print_uops, pyrender from tinygrad.dtype import DType, ImageDType, dtypes, PtrDType, AddrSpace, Invalid, ConstFloat from tinygrad.helpers import DEBUG, Context, prod, SPEC, Metadata, panic, CHECK_OOB @@ -70,20 +70,20 @@ def type_verify(ast:UOp|list[UOp], check_spec:PatternMatcher): (UPat(Ops.INDEX, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(dtypes.is_int(y.dtype) for y in x.src[1:]) or None), (UPat(Ops.END, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(u.op is Ops.RANGE for u in x.src[1:])), - # PARAM (that's really a DEFINE_GLOBAL) - (UPat(Ops.PARAM, name="x"), lambda x: isinstance(x.dtype, (PtrDType, ImageDType)) and x.dtype.addrspace == AddrSpace.GLOBAL), + # PARAM + (UPat(Ops.PARAM, name="x"), lambda x: isinstance(x.arg, ParamArg)), # GROUP of stores (or groups, or NOOPs) # TODO: remove UNROLL here, it's for SPEC=2 - (UPat(Ops.GROUP, dtypes.void, src=UPat((Ops.GROUP, Ops.STORE, Ops.NOOP, Ops.UNROLL))), lambda: True), + (UPat(Ops.GROUP, dtypes.void, src=UPat((Ops.GROUP, Ops.STORE, Ops.NOOP, Ops.UNROLL, Ops.INS))), lambda: True), # TOOD: these should be buffer with different addrspace (UPat(Ops.DEFINE_LOCAL, name="x"), lambda x: isinstance(x.dtype, PtrDType) and x.dtype.addrspace == AddrSpace.LOCAL), - (UPat(Ops.DEFINE_REG, src=(), name="x"), lambda x: isinstance(x.arg, int)), + (UPat(Ops.DEFINE_REG, src=()), lambda: True), # AFTER on Movement Op, PARAM, BUFFER, CONTIGUOUS, or another AFTER (UPat(Ops.AFTER, src=(UPat(GroupOp.Movement.union({Ops.PARAM, Ops.BUFFER, Ops.CONTIGUOUS, Ops.DEFINE_REG, Ops.DEFINE_LOCAL, Ops.AFTER, Ops.MULTI, - Ops.BITCAST})),), + Ops.BITCAST, Ops.INS})),), allow_any_len=True), lambda: True), # CUSTOM (inline and non inline) @@ -124,15 +124,12 @@ def type_verify(ast:UOp|list[UOp], check_spec:PatternMatcher): # CONST with a UNIQUE or DEVICE (UPat(Ops.CONST, src=(UPat(Ops.DEVICE),)), lambda: True), - (UPat(Ops.CONST, src=(UPat((Ops.UNIQUE, Ops.LUNIQUE)), UPat(Ops.DEVICE))), lambda: True), + (UPat(Ops.CONST, src=(UPat((Ops.UNIQUE, Ops.LUNIQUE)), UPat(Ops.DEVICE)), name="c"), lambda c: c.arg is Invalid), # BUFFER (UPat(Ops.BUFFER, src=(UPat((Ops.UNIQUE, Ops.LUNIQUE)), UPat(Ops.DEVICE)), name="buf"), lambda buf: isinstance(buf.arg, int) and isinstance(buf.dtype, DType)), - # PARAM (that's really a variable) - (UPat(Ops.PARAM, src=(UPat(), UPat(), UPat(), UPat(), UPat()), name="x"), lambda x: True), - # Tensor variable bindings (UPat(Ops.BIND, (dtypes.int, dtypes.weakint,), (UPat(Ops.DEFINE_VAR), UPat.cvar(dtype=(dtypes.int,dtypes.weakint,))), arg=None), lambda: True), @@ -147,18 +144,14 @@ def type_verify(ast:UOp|list[UOp], check_spec:PatternMatcher): (UPat(Ops.TUPLE, dtypes.void), lambda: True), (UPat(Ops.GETTUPLE, src=(UPat((Ops.FUNCTION, Ops.TUPLE)),), name="g"), lambda g: isinstance(g.arg, int)), - # PARAM - (UPat(Ops.PARAM, src=(UPat(), UPat(Ops.NOOP)), name="x"), lambda x: True), # TODO: why does this have NOOP? - (UPat(Ops.PARAM, src=(UPat(), UPat(Ops.DEVICE)), name="x"), lambda x: True), - (UPat(Ops.PARAM, src=(UPat(), UPat(Ops.DEVICE), UPat(Ops.MULTI)), name="x"), lambda x: True), - # inputs to movement ops - (UPat((Ops.STACK, Ops.VCONST)), lambda: True), + (UPat(Ops.STACK), lambda: True), (UPat({Ops.ADD, Ops.MUL, Ops.CDIV, Ops.FLOORDIV}, dtype=dtypes.weakint), lambda: True), # movement ops (UPat((Ops.RESHAPE, Ops.EXPAND), src=(UPat(), UPat(dtype=dtypes.weakint))), lambda: True), - (UPat((Ops.PAD, Ops.SHRINK), src=(UPat(), UPat(dtype=dtypes.weakint), UPat(dtype=dtypes.weakint))), lambda: True), + (UPat((Ops.PAD, Ops.SHRINK), src=(UPat(), UPat(dtype=dtypes.weakint), UPat(dtype=dtypes.weakint)), name="x"), + lambda x: x.src[1].dtype.count == x.src[2].dtype.count), (UPat((Ops.PERMUTE, Ops.FLIP), name="mv", src=(UPat(),)), lambda mv: isinstance(mv.arg, tuple)), # REDUCE has arg=(op, axis_tuple), src[1:] are ranges after lowering @@ -198,6 +191,19 @@ def type_verify(ast:UOp|list[UOp], check_spec:PatternMatcher): # these ops can exist in programs but not the tensor spec. example: LOAD spec_program = PatternMatcher([ + # weakint is not allowed in programs + (UPat(GroupOp.All, dtypes.weakint), lambda: False), + + # movement ops are not allowed in programs + (UPat(GroupOp.Movement), lambda: False), + + # Invalid is not allowed in program + (UPat(Ops.CONST, arg=Invalid), lambda: False), + + # shape of uop must match dtype.count in program + (UPat(GroupOp.All-{Ops.INS, Ops.NOOP}, name="x"), + lambda x: False if x.dtype.count > 1 and (x.dtype.count,) != x.shape else None), + # STACK/GEP in program. TODO: this should match Tensor (UPat(Ops.STACK, name="x"), lambda x: len(x.src)>1 and len(x.src) == x.dtype.vcount and all(x.dtype == y.dtype.vec(len(x.src)) for y in x.src)), (UPat(Ops.GEP, src=(UPat.var("src"),), name="gep"), lambda gep,src: gep.dtype == src.dtype.scalar()), @@ -209,12 +215,12 @@ def type_verify(ast:UOp|list[UOp], check_spec:PatternMatcher): # these are intermediate ops. everything should be deleted from here spec_full = PatternMatcher([ - # BUFFER_VIEW on BUFFER is allowed if BUFFER is - (UPat(Ops.BUFFER_VIEW, src=(UPat((Ops.BUFFER, Ops.PARAM)),)), lambda: True), + # SLICE on BUFFER is allowed if BUFFER is + (UPat(Ops.SLICE, src=(UPat(GroupOp.Movement.union({Ops.BUFFER, Ops.PARAM, Ops.STAGE, Ops.AFTER})), + UPat(Ops.CONST, dtype=dtypes.weakint)), allow_any_len=True, name="bv"), + lambda bv: isinstance(bv.arg, int)), - # TODO: BUFFER_VIEW shouldn't go on INDEX. why is this allowed? remove these both - (UPat(Ops.BUFFER_VIEW, src=(UPat((Ops.INDEX,)),), allow_any_len=True), lambda: True), - (UPat(Ops.CALL, src=(UPat((Ops.BUFFER_VIEW,)),), allow_any_len=True), lambda: True), + (UPat(Ops.CALL, src=(UPat((Ops.SLICE,)),), allow_any_len=True), lambda: True), # codegen may end ranges after gpudims has replaced RANGE with SPECIAL. (UPat(Ops.END, src=(UPat(), UPat()), allow_any_len=True), lambda: True), @@ -250,7 +256,7 @@ def type_verify(ast:UOp|list[UOp], check_spec:PatternMatcher): 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} + "ConstFloat": ConstFloat, "ParamArg": ParamArg} 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 316fee67b426f..99d2469f2edd1 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -22,6 +22,15 @@ def fold_bitcast(root:UOp, c:UOp) -> UOp|None: def convert(v:ConstType) -> ConstType: return struct.unpack(to_fmt, struct.pack(from_fmt, v))[0] return root.const_like(convert(c.arg) if root.dtype.count == 1 else tuple(map(convert, c.arg))) +def const_arg(u:UOp) -> ConstType|tuple[ConstType, ...]|None: + if u.op is Ops.CONST: return u.arg + if u.op is Ops.STACK and all(s.op is Ops.CONST for s in u.src): return tuple(s.arg for s in u.src) + return None + +def fold_const_alu(a:UOp) -> UOp|None: + vals = [const_arg(s) for s in a.src] + return None if any(v is None for v in vals) else a.const_like(exec_alu(a.op, a.dtype, vals, False)) + invalid_pat = UPat(Ops.CONST, arg=Invalid, name="i") invalid_gate = UPat.var("cond").where(UPat.var("x"), invalid_pat) @@ -71,6 +80,10 @@ def fold_add_divmod_recombine(x:UOp) -> UOp|None: (UPat.var("a").where(UPat.var("b"), invalid_gate), lambda cond,i,x,a,b: (a|cond).where(a.where(b, x), i) if b.arg != Invalid else None), (UPat(Ops.BITCAST, src=(invalid_pat,), name="bc"), lambda bc,i: i.cast(bc.dtype)), (UPat(Ops.BITCAST, src=(invalid_gate,), name="bc"), lambda bc,cond,x,i: cond.where(x.bitcast(bc.dtype), i.bitcast(bc.dtype))), + # fold gated LOAD/STORE + (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat(), invalid_pat), allow_any_len=True).or_casted(), UPat())), lambda i: UOp(Ops.NOOP)), + (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(), invalid_pat), allow_any_len=True).or_casted(),), allow_any_len=True, name="x"), + lambda x,i: x.src[1] if len(x.src) > 1 else x.const_like(0)), ]) symbolic_simple = propagate_invalid + PatternMatcher([ @@ -85,14 +98,14 @@ def fold_add_divmod_recombine(x:UOp) -> UOp|None: ((UPat.var() % UPat.var("y")).named("base") % UPat.var("y"), lambda base,y: base), # (x%y)%y = -> x%y (rewritten with base for speed) # variations of (x%c)+(x//c)*c = x (UPat(Ops.ADD, dtype=dtypes.weakint, name="x"), fold_add_divmod_recombine), - (UPat.var("x", dtype=dtypes.bool) & UPat.cvar("c", vec=False), lambda x,c: x if c.arg else c), - (UPat.var("x", dtype=dtypes.bool) | UPat.cvar("c", vec=False), lambda x,c: c if c.arg else x), + (UPat.var("x", dtype=dtypes.bool) & UPat.cvar("c"), lambda x,c: x if c.arg else c), + (UPat.var("x", dtype=dtypes.bool) | UPat.cvar("c"), lambda x,c: c if c.arg else x), (UPat(GroupOp.Idempotent, src=(UPat.var("x"), UPat.var("x"))), lambda x: x), (UPat.var("x", dtype=dtypes.bool).logical_not().logical_not(), lambda x: x), (UPat.var("x", dtype=dtypes.bool).where(UPat.const(dtypes.bool, True), UPat.const(dtypes.bool, False)), lambda x: x), (UPat.var("x", dtype=dtypes.bool).where(UPat.const(dtypes.bool, False), UPat.const(dtypes.bool, True)), lambda x: x.logical_not()), # CAST(bool -> int) != const — CAST(True)=1, CAST(False)=0, so fold based on const value - (UPat.var("x", dtype=dtypes.bool).cast(dtypes.ints+(dtypes.weakint,)) != UPat.cvar("c", vec=False), + (UPat.var("x", dtype=dtypes.bool).cast(dtypes.ints+(dtypes.weakint,)) != UPat.cvar("c"), lambda x,c: x if c.arg == 0 else x.logical_not() if c.arg == 1 else x.const_like(True)), (UPat.var("x", dtype=dtypes.ints+(dtypes.bool, dtypes.weakint)).trunc(), lambda x: x), # ** zero folding ** @@ -102,18 +115,16 @@ def fold_add_divmod_recombine(x:UOp) -> UOp|None: (UPat.var("x") & 0, lambda x: x.const_like(0)), # x&0 -> 0 # (x&mask)>>k -> x>>k when mask only clears bits below k # TODO: combine this with "# rules for threefry" below - ((UPat.var("x") & UPat.cvar("mask", vec=False)) >> UPat.cvar("k", vec=False), + ((UPat.var("x") & UPat.cvar("mask")) >> UPat.cvar("k"), lambda x,mask,k: x >> k.arg if mask.arg | ((1 << k.arg) - 1) == -1 else None), (UPat.var("x", dtype=dtypes.ints+(dtypes.bool, dtypes.weakint)) != UPat.var("x"), lambda x: x.const_like(False).cast(dtypes.bool.vec(x.dtype.count))), # x != x -> False (only ints) # ** constant folding ** - (UPat(GroupOp.Unary, src=(UPat((Ops.VCONST, Ops.CONST)),), name="a"), lambda a: a.const_like(exec_alu(a.op, a.dtype, [a.src[0].arg], False))), - (UPat(GroupOp.Binary-{Ops.THREEFRY}, src=(UPat((Ops.VCONST, Ops.CONST)),)*2, name="a"), - lambda a: a.const_like(exec_alu(a.op, a.dtype, [a.src[0].arg, a.src[1].arg], False))), + (UPat(GroupOp.Unary, src=(UPat((Ops.CONST, Ops.STACK)),), name="a"), fold_const_alu), + (UPat(GroupOp.Binary-{Ops.THREEFRY}, src=(UPat((Ops.CONST, Ops.STACK)),)*2, name="a"), fold_const_alu), (UPat(Ops.THREEFRY, src=(UPat.cvar("x"), UPat.cvar("key")), name="a"), lambda a, x, key: a.const_like(threefry2x32(x, key).simplify().arg)), - (UPat(GroupOp.Ternary, src=(UPat((Ops.VCONST, Ops.CONST)),)*3, name="a"), - lambda a: a.const_like(exec_alu(a.op, a.dtype, [a.src[0].arg, a.src[1].arg, a.src[2].arg], False))), + (UPat(GroupOp.Ternary, src=(UPat((Ops.CONST, Ops.STACK)),)*3, name="a"), fold_const_alu), # bool MUL is AND, ADD/MAX is OR. prevents other rules to rewrite bool ADD/MUL incorrectly (UPat.var('x', dtype=dtypes.bool) * UPat.var('y', dtype=dtypes.bool), lambda x,y: x&y), (UPat.var('x', dtype=dtypes.bool) + UPat.var('y', dtype=dtypes.bool), lambda x,y: x|y), @@ -137,9 +148,9 @@ def fold_add_divmod_recombine(x:UOp) -> UOp|None: (UPat.var('x').cast(name="a").cast(name="b"), lambda x,a,b: x if x.dtype == b.dtype and can_lossless_cast(b.dtype, a.dtype) else None), (UPat.var("x").cast(dtypes.bool), lambda x: x != 0), # ** pow ** - (UPat.var("x").alu(Ops.POW, UPat.cvar("c", vec=False)), simplify_pow), + (UPat.var("x").alu(Ops.POW, UPat.cvar("c")), simplify_pow), # positive const ** x - (UPat.cvar("c", vec=False).alu(Ops.POW, UPat.var("x")), lambda c,x: c if c.arg == 1 else (x*math.log2(c.arg)).exp2() if c.arg > 0 else None), + (UPat.cvar("c").alu(Ops.POW, UPat.var("x")), lambda c,x: c if c.arg == 1 else (x*math.log2(c.arg)).exp2() if c.arg > 0 else None), # rules for threefry ((UPat.var('x', dtypes.uint64)&0xFFFFFFFF).cast(dtypes.uint32), lambda x: x.cast(dtypes.uint32)), (((UPat.var(None, dtypes.uint64)*(1<<32)) | UPat.var('y', dtypes.uint32).cast(dtypes.uint64)).cast(dtypes.uint32), lambda y: y), @@ -149,7 +160,7 @@ def fold_add_divmod_recombine(x:UOp) -> UOp|None: # ** simple where folding ** # a conditional with the same results either way is a noop, also fold const conditionals (UPat.var().where(UPat.var("val"), UPat.var("val")), lambda val: val), - (UPat.cvar("gate", vec=False).where(UPat.var("c0"), UPat.var("c1")), lambda gate, c0, c1: c0 if gate.arg else c1), + (UPat.cvar("gate").where(UPat.var("c0"), UPat.var("c1")), lambda gate, c0, c1: c0 if gate.arg else c1), # a.where(b.where(c, d), d) -> (a & b).where(c, d) (UPat.var("a").where(UPat.var("b").where(UPat.var("c"), UPat.var("d")), UPat.var("d")), lambda a,b,c,d: (a&b).where(c,d)), ]) @@ -189,13 +200,12 @@ def gep_through_wmma(gep:UOp, wmma:UOp) -> UOp|None: return UOp(Ops.WMMA, gep.dtype, tuple(tsrcs), wmma.arg) gep_pushing = PatternMatcher([ - # GEP/VECTORIZE, GEP/GEP, GEP/CONST, GEP/VCONST + # GEP/VECTORIZE, GEP/GEP, GEP/CONST (UPat(Ops.GEP, name='g2').f(Ops.GEP, name='g1'), lambda g1, g2: g2.src[0].gep(tuple(g2.arg[g1.arg[i]] for i in range(len(g1.arg))))), (UPat(Ops.STACK, name='vec').f(Ops.GEP, name='gep'), lambda gep, vec: UOp(Ops.STACK, gep.dtype, tuple(vec.src[i] for i in gep.arg)) if len(gep.arg) > 1 else vec.src[gep.arg[0]]), - (UPat.cvar("c", vec=False).f(Ops.GEP, name="gep"), lambda gep, c: gep.const_like(c.arg)), - (UPat(Ops.VCONST, name="c").f(Ops.GEP, name="gep"), lambda gep, c: gep.const_like(tuple(c.arg[x] for x in gep.arg))), + (UPat.cvar("c").f(Ops.GEP, name="gep"), lambda gep, c: gep.const_like(c.arg)), # GEP on void is skipped (UPat(Ops.GEP, src=(UPat(dtype=dtypes.void, name="x"),)), lambda x: x), # GEP in order is removed @@ -216,7 +226,8 @@ def gep_through_wmma(gep:UOp, wmma:UOp) -> UOp|None: commutative = PatternMatcher([ # ** COMMUTATIVE flipping (only for index) ** # NOTE: this can break merging vector math by only flipping some of them - (UPat(GroupOp.Commutative, dtype=dtypes.weakint, name='x'), lambda x: x.replace(src=x.src[::-1]) if x.src[1].tuplize < x.src[0].tuplize else None), + (UPat(GroupOp.Commutative, dtype=dtypes.weakint, name='x'), lambda x: + x.replace(src=x.src[::-1]) if x.src[1].tuplize < x.src[0].tuplize and not x.src[0].tuplize < x.src[1].tuplize else None), ]) symbolic = symbolic_simple+commutative+PatternMatcher([ @@ -259,20 +270,20 @@ def gep_through_wmma(gep:UOp, wmma:UOp) -> UOp|None: ((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 0 and c1.arg > 0 else None), # c0*x x0 - ((UPat.var("x", dtype=dtypes.weakint)//UPat.cvar("d", vec=False)) 0 else None), # ** move add/mul consts to end (NOTE: this is still happening before constant folding) ** ((UPat.var("x") + UPat.cvar("c1")) + UPat.var("y"), lambda x,c1,y: (x+y)+c1), ((UPat.var("x") * UPat.cvar("c1")) * UPat.var("y"), lambda x,c1,y: (x*y)*c1), # *** rules from symbolic *** # generic lt folding - (UPat.var("x", dtypes.weakint) 0. NOTE: not x < 1 means x > 0 ((UPat.var("x", dtypes.weakint)<1).ne(True), lambda x: (newx<1).ne(True) if (newx:=canonicalize_simplex(x)) is not None else None), @@ -311,7 +322,6 @@ def parse_valid(v:UOp) -> tuple[UOp, bool, int]|None: if v.op is Ops.CMPLT and dtypes.is_int(v.src[0].dtype): # X < c -> X <= c-1 return v.src[0], True, int((v.src[1]).vmax)-1 - # NOTE: v.src[1].op can be Ops.VCONST return None def uop_given_valid(valid:UOp, uop:UOp, try_simplex=True) -> UOp: @@ -428,7 +438,15 @@ def gated_given_valid(cond:UOp, x:UOp, i:UOp) -> UOp|None: ]) # this is symbolic 2.0 -REMOVE_FROM_SINK_LIKE = {Ops.UNROLL, Ops.NOOP, Ops.STACK, Ops.SINK} +REMOVE_FROM_SINK_LIKE = {Ops.UNROLL, Ops.NOOP, Ops.STACK, Ops.SINK, Ops.GROUP} +pm_clean_up_group_sink = PatternMatcher([ + # clean up GROUP/SINK + (UPat(Ops.GROUP, src=(UPat.var("x"),)), lambda x: x), + (UPat((Ops.SINK, Ops.GROUP), name="root"), + lambda root: UOp(root.op, root.dtype, tuple(flatten(x.src if x.op in REMOVE_FROM_SINK_LIKE else (x,) for x in root.src)), root.arg) + if any(x.op in REMOVE_FROM_SINK_LIKE for x in root.src) else None), +]) + sym = symbolic+pm_simplify_valid+PatternMatcher([ # reorder ALU/VECTORIZE (UPat(GroupOp.ALU, src=(UPat(Ops.STACK, src=UPat(name='x')), UPat(Ops.STACK, src=UPat(name='y'))), name='alu'), @@ -446,9 +464,6 @@ def gated_given_valid(cond:UOp, x:UOp, i:UOp) -> UOp|None: UPat.load(UPat(Ops.INDEX, name="index")))), lambda index, gate, alt: UOp.store(index.src[0].index(gate.where(index.src[1], UOp.invalid())), alt)), # fold gated LOAD/STORE - (UPat(Ops.STORE, src=(UPat().index(UPat.const(dtypes.weakint, Invalid)).or_casted(), UPat())), lambda: UOp(Ops.NOOP)), - (UPat(Ops.LOAD, src=(UPat().index(UPat.const(dtypes.weakint, Invalid)).or_casted(),), allow_any_len=True, name="x"), - lambda x: x.src[1] if len(x.src) > 1 else x.const_like(0)), # invalid load produces 0, or the alt value if we have one (UPat(Ops.STORE, src=(UPat(), invalid_pat)), lambda i: UOp(Ops.NOOP)), # store of where with invalid -> gated store (UPat(Ops.STORE, src=(UPat(Ops.INDEX, name="index"), UPat.var("cond").where(UPat.var("val"), invalid_pat))), @@ -460,16 +475,11 @@ def gated_given_valid(cond:UOp, x:UOp, i:UOp) -> UOp|None: (UPat.var("x") * ((1+UPat.var("x")).reciprocal().named("d")*UPat.var("y")), lambda x,y,d: y*(1-d)), (UPat.var("x") * ((1+UPat.var("x")).reciprocal().named("d")+UPat.var("y")), lambda x,y,d: (1-d)+x*y), # move const multiply after REDUCE (NOTE: the mul chain can do this, but only if it's a same dtype reduce) - ((UPat.var("x")*UPat.cvar("c", vec=False)).reduce(arg=Ops.ADD, name="r", allow_any_len=True), lambda x,c,r: r.replace(src=(x,)+r.src[1:])*c.arg), + ((UPat.var("x")*UPat.cvar("c")).reduce(arg=Ops.ADD, name="r", allow_any_len=True), lambda x,c,r: r.replace(src=(x,)+r.src[1:])*c.arg), # reduce mul chain, move muls after the reduce (UPat(Ops.MUL).reduce(name="r", allow_any_len=True), reduce_mul_chain), - # clean up GROUP/SINK - (UPat(Ops.GROUP, src=(UPat.var("x"),)), lambda x: x), - (UPat((Ops.SINK, Ops.GROUP), name="root"), - lambda root: UOp(root.op, root.dtype, tuple(flatten(x.src if x.op in REMOVE_FROM_SINK_LIKE else (x,) for x in root.src)), root.arg) - if any(x.op in REMOVE_FROM_SINK_LIKE for x in root.src) else None), # ** combine terms (opinionated) ** (-1 * (UPat.var("x") + UPat.var("y")), lambda x,y: (-x)+(-y)), # -(x+y) -> -x + -y # (x+y)*c -> x*c+y*c. only for int, float has inf*0=nan issue ((UPat.var("x", dtypes.weakint) + UPat.var("y")) * UPat.cvar("c"), lambda x,y,c: x*c+y*c), -]) +])+pm_clean_up_group_sink diff --git a/tinygrad/uop/upat.py b/tinygrad/uop/upat.py index 3e18ef083b7a7..0849045f05970 100644 --- a/tinygrad/uop/upat.py +++ b/tinygrad/uop/upat.py @@ -26,6 +26,10 @@ def _get_clause(self:UPat, base:UOp, depth=0) -> UOp: if len(self.match_dtype) > 1: and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.BIND, arg=tuple(self.match_dtype))), arg="({0}.dtype in {1} or {0}.dtype._scalar in {1})")) else: and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.BIND, arg=self.match_dtype[0])), arg="({0}.dtype == {1} or {0}.dtype._scalar == {1})")) + if self.match_tag is not None: + if len(self.match_tag) > 1: + and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.BIND, arg=tuple(self.match_tag))), arg="{0}.tag in {1}")) + else: and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.BIND, arg=self.match_tag[0])), arg="{0}.tag == {1}")) if self.src is not None: # single match if len(self.src) == 1 and isinstance(self.src[0], tuple): diff --git a/tinygrad/viz/__init__.py b/tinygrad/viz/__init__.py new file mode 100644 index 0000000000000..e69de29bb2d1d diff --git a/tinygrad/viz/cli.py b/tinygrad/viz/cli.py index 5ac88b918a9f3..f498d9f2a9996 100755 --- a/tinygrad/viz/cli.py +++ b/tinygrad/viz/cli.py @@ -52,6 +52,10 @@ def to_str(k:str, v) -> str: return f"{k}={v}" def fmt_data(data:dict) -> str: return " ".join((p:=to_str(k, v))+" "*max(0, 14-ansilen(p)) for k,v in data.items()) +def marker_st(markers:list[dict], name:str) -> int: + try: return next(e["ts"] for e in markers if e["name"] == name) + except StopIteration: raise RuntimeError(f"marker not found: {name}") from None + def get(data:dict, key:str): for k,v in data.items(): if ansistrip(k) == key: return v @@ -140,12 +144,15 @@ def hex_colored(st:str, color:str) -> str: return f"\x1b[38;2;{int(color[1:3],16 # ** Profiler printer else: timelines = [(n,l) for n,l in profile["layout"].items() if isinstance(l, dict) and l.get("event_type") == 0] + markers = profile.get("markers", []) + interval:tuple[int, int]|None = None if not args.interval else (marker_st(markers, args.interval[0]), marker_st(markers, args.interval[1])) def produce_top_kernels() -> Iterator[dict]: tagged = ((n,e) for n,l in timelines for e in l["events"]) if not args.src else ((args.src[0],e) for e in unwrap(data)["events"]) agg:dict[tuple[str,str], tuple[float, int, int|None, dict[str, float]]] = {} # map (device, kernel name) to (total time, count, ref, est) est_keys = ("FLOPS", "B/s mem", "B/s lds") total = 0 for dev,e in tagged: + if interval and not interval[0] <= e["st"] <= interval[1]: continue et = e["dur"] * 1e-3 t, c, ref, est = agg.get((dev,e["name"]), (0.0, 0, None, {})) est.update({k:est.get(k, 0.0)+e["fmt"][k]*e["dur"]*1e-6 for k in est_keys if k in e["fmt"]}) @@ -166,8 +173,9 @@ def produce_all_kernels() -> Iterator[dict]: if not args.src: for n,l in profile["layout"].items(): if not isinstance(l, dict) or l.get("event_type") != 0: yield {"device":"SOURCE", "name":n, "st_ms":0, "ref":None, "ext":None} - marker_stream = sorted([(m["ts"], "MARKER", m) for m in profile.get("markers", [])], key=lambda t:t[0]) + marker_stream = sorted([(m["ts"], "MARKER", m) for m in markers], key=lambda t:t[0]) for ts,dev,e in heapq.merge(*event_streams, marker_stream, key=lambda t:t[0]): + if interval is not None and not interval[0] <= ts <= interval[1]: continue if dev == "MARKER": yield {"device":dev, "name":fmt_colored(e["name"]), "st_ms":ts*1e-3, "ref":None, "ext":None} continue @@ -199,7 +207,8 @@ def render_event(k:dict, ls=args.list) -> None: if DEBUG >= 3 and s["name"] == "View Base AST": print_step(s) if DEBUG >= 4 and s["name"] == "View Source": print_step(s) if DEBUG >= 5 or ls: print(emit(" "*s["depth"]+s["name"]+(f" - {s['match_count']}" if s.get('match_count', 0) else ''))) - if DEBUG >= 6 or (DEBUG >= 5 and s["name"] == "View Kernel Graph") or (s["name"] in args.src): print_step(s, print_graph=True) + if DEBUG >= 6 or (DEBUG >= 5 and s["name"] == "View Kernel Graph") or (s["name"] in args.src): + print_step(s, print_graph=True, reconstruct_matches=s["name"] in args.src) if DEBUG >= 7: print_step(s, reconstruct_matches=True) elif DEBUG >= 3 and k.get("ext"): print(emit(k["ext"])) for k in (produce_top_kernels if args.t else produce_all_kernels)(): render_event(k) @@ -208,6 +217,7 @@ def get_arg_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(prog="python -m tinygrad.viz.cli") parser.add_argument("-s", "--src", nargs="+", default=[], metavar="NAME", help="Select a data source (default: all)") parser.add_argument("--list", "--ls", dest="list", action="store_true", help="List sources") + parser.add_argument("--interval", nargs=2, metavar=("START", "END"), help="Optional start and end marker") parser.add_argument("-t", nargs="?", type=int, const=20, metavar="COUNT", help="Aggregate top kernels (optional count, default 20)") parser.add_argument("--profile-path", type=str, metavar="PATH", help="Optional path to profile.pkl (default: latest profile)", default=temp("profile.pkl", append_user=True)) diff --git a/tinygrad/viz/index.html b/tinygrad/viz/index.html index 35804472bc1e4..6163e2f8bd250 100644 --- a/tinygrad/viz/index.html +++ b/tinygrad/viz/index.html @@ -102,18 +102,18 @@ fill: #FFD700; stroke: #B8860B; } - g.tag.collapsed circle { + g.tag.collapsed circle, g.tag.collapsed rect { fill: #5CD68D; stroke: #4a4b57; } - g.tag.expanded circle { + g.tag.expanded circle, g.tag.expanded rect { fill: #9FDDE6; stroke: #4a4b57; } g.port circle { fill: #b3dcc2; } - g.tag circle, #edge-labels circle { + g.tag circle, g.tag rect, #edge-labels circle { stroke-width: 0.8; } g.tag text, #edge-labels text { diff --git a/tinygrad/viz/js/index.js b/tinygrad/viz/js/index.js index e35e718f5ac01..b47a9012e3379 100644 --- a/tinygrad/viz/js/index.js +++ b/tinygrad/viz/js/index.js @@ -57,11 +57,14 @@ function intersectRect(r1, r2) { } function addTags(root, path) { - root.selectAll("circle").data(d => [d]).join("circle").attr("r", 5); + root.selectAll("circle").data(d => d.rect ? [] : [d]).join("circle").attr("r", 5).style("fill", d => d.fill ?? null).style("stroke", d => d.stroke ?? null); + root.selectAll("rect").data(d => d.rect ? [d] : []).join("rect").attr("x", d => -d.width/2).attr("y", d => -d.height/2) + .attr("width", d => d.width).attr("height", d => d.height).style("fill", d => d.fill ?? null).style("stroke", d => d.stroke ?? null); if (path != null) root.selectAll("path").data(d => [d]).join("path").attr("d", path); - else root.selectAll("text").data(d => [d]).join("text").text(d => d).attr("dy", "0.35em"); + else root.selectAll("text").data(d => [d]).join("text").text(d => d.text).attr("dy", "0.35em"); } +anchor = null; const drawGraph = (data) => { const g = dagre.graphlib.json.read(data); // draw nodes @@ -69,11 +72,6 @@ const drawGraph = (data) => { const callCount = g.graph().callCount; const nodes = d3.select("#nodes").selectAll("g").data(g.nodes().map(id => g.node(id)), d => d).join("g").attr("class", d => d.className ?? "node") .attr("transform", d => `translate(${d.x},${d.y})`).on("click", (e,d) => { - if (d.callNode) { - if (state.callSrcMask.has(d.id)) state.callSrcMask.delete(d.id); else state.callSrcMask.add(d.id); - if (state.callSrcMask.size >= callCount) { showCallSrc.toggle.checked = !showCallSrc.toggle.checked; state.callSrcMask.clear(); } - return setState({}); - } const parents = g.predecessors(d.id); const children = g.successors(d.id); if (parents == null && children == null) return; @@ -89,7 +87,7 @@ const drawGraph = (data) => { .attr("x", d => -d.width/2).attr("y", d => -d.height/2).classed("node", true); const STROKE_WIDTH = 1.4, textSpace = g.graph().textSpace; const labels = nodes.selectAll("g.label").data(d => [d]).join("g").attr("class", "label"); - labels.attr("transform", d => `translate(-${d.labelWidth/2}, -${d.labelHeight/2+STROKE_WIDTH*2})`); + labels.attr("transform", d => `translate(${d.labelX-d.labelWidth/2}, -${d.labelHeight/2+STROKE_WIDTH*2})`); const rectGroup = labels.selectAll("g.rect-group").data(d => [d]).join("g").attr("class", "rect-group"); const tokens = labels.selectAll("g.text-group").data(d => [d]).join("g").attr("class", "text-group").selectAll("text").data(d => { if (Array.isArray(d.label)) return [d.label]; @@ -116,11 +114,26 @@ const drawGraph = (data) => { tokensBg.classed("highlight", (d, i, nodes) => !nodes[i].classList.contains("highlight") && d.keys.some(k => keys?.includes(k))); }); addTags(nodes.selectAll("g.tag").data(d => d.tag != null ? [d] : []).join("g").attr("class", "tag") - .attr("transform", d => `translate(${-d.width/2+8}, ${-d.height/2+8})`).datum(e => e.tag)); - addTags(nodes.selectAll("g.type").data(d => d.callNode ? [d] : []).join("g").attr("class", d => `tag ${d.collapsed ? 'collapsed' : 'expanded'}`) - .attr("transform", d => `translate(${-d.width/2}, ${0})`).datum(d => d.collapsed ? "+" : "−")); + .attr("transform", d => `translate(${-d.width/2+8}, ${-d.height/2+8})`).datum(e => ({ text:e.tag }))); + addTags(nodes.selectAll("g.addrspace").data(d => d.addrspace != null ? [d] : []).join("g").attr("class", "tag addrspace") + .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) => { + e.stopPropagation(); + const t = d3.zoomTransform(document.getElementById("graph-svg")); + const [x, y] = t.apply([d.x, d.y]); + anchor = {id:d.id, x, y, k:t.k}; + if (d.callNode) { + if (state.callSrcMask.has(d.id)) state.callSrcMask.delete(d.id); else state.callSrcMask.add(d.id); + if (state.callSrcMask.size >= callCount) { showCallSrc.toggle.checked = !showCallSrc.toggle.checked; state.callSrcMask.clear(); } + } else { if (state.expandedNodes.has(d.id)) state.expandedNodes.delete(d.id); else state.expandedNodes.add(d.id); } + return setState({}); + })); addTags(nodes.selectAll("g.ref").data(d => d.ref != null ? [d] : []).join("g").attr("class", "tag ref") - .attr("transform", d => `translate(${d.width/2-2}, ${-d.height/2+2})`).on("click", (e,d) => { e.stopPropagation(); switchCtx(d.ref); }), + .attr("transform", d => `translate(${d.width/2-2}, ${-d.height/2+2})`).on("click", (e,d) => { e.stopPropagation(); switchCtx(d.ref); }).datum(d => ({ref:d.ref})), "M-1.7 1.7 L1.7 -1.7 M-0.55 -1.7 H1.7 V0.55"); // draw edges const line = d3.line().x(d => d.x).y(d => d.y).curve(d3.curveBasis), edges = g.edges(); @@ -131,6 +144,7 @@ const drawGraph = (data) => { points.push(intersectRect(g.node(e.w), points[points.length-1])); return line(points); }).attr("marker-end", "url(#arrowhead)").attr("stroke", e => g.edge(e).color || "#4a4b57"); + return g; } // ** UOp graph @@ -155,7 +169,7 @@ function renderDag(layoutSpec, { recenter }) { const data = e.data.result; displaySelection("#graph"); updateProgress(Status.COMPLETE); - drawGraph(data); + const g = drawGraph(data); addTags(d3.select("#edge-labels").selectAll("g").data(data.edges).join("g").attr("transform", (e) => { // get a point near the end const [p1, p2] = e.value.points.slice(-2); @@ -170,8 +184,12 @@ function renderDag(layoutSpec, { recenter }) { const x = p2.x - ux * offset; const y = p2.y - uy * offset; return `translate(${x}, ${y})` - }).attr("class", e => e.value.label.type).attr("id", e => `${e.v}-${e.w}`).datum(e => e.value.label.text)); - if (recenter) document.getElementById("zoom-to-fit-btn").click(); + }).attr("class", e => e.value.label.type).attr("id", e => `${e.v}-${e.w}`).datum(e => ({ text:e.value.label.text }))); + if (anchor != null) { + const n = g.node(anchor.id); + d3.select("#graph-svg").call(svgZoom.transform, d3.zoomIdentity.translate(anchor.x-n.x*anchor.k, anchor.y-n.y*anchor.k).scale(anchor.k)); + } else if (recenter) document.getElementById("zoom-to-fit-btn").click(); + anchor = null; }; worker.onerror = (e) => { e.preventDefault(); @@ -859,7 +877,7 @@ const evtSources = []; // rewrite: a single UOp transformation // step: collection of rewrites // context: collection of steps -const state = {currentCtx:-1, currentStep:0, currentRewrite:0, expandSteps:false, callSrcMask:new Set()}; +const state = {currentCtx:-1, currentStep:0, currentRewrite:0, expandSteps:false, callSrcMask:new Set(), expandedNodes:new Set()}; function setState(ns) { saveToHistory(state); const { ctx:prevCtx, step:prevStep } = select(state.currentCtx, state.currentStep); @@ -1054,12 +1072,13 @@ async function main() { if (ret.length === 0) return; // ** center graph const data = ret[currentRewrite]; - const render = (opts) => renderDag({ data, opts }, { recenter:currentRewrite === 0 }); - const getOpts = () => ({ showIndexing:showIndexing.toggle.checked, showCallSrc:showCallSrc.toggle.checked, showSink:showSink.toggle.checked, callSrcMask:state.callSrcMask }); - render(getOpts()); - showIndexing.toggle.onchange = () => render(getOpts()); - showCallSrc.toggle.onchange = () => { state.callSrcMask.clear(); render(getOpts()); } - showSink.toggle.onchange = () => render(getOpts()); + const render = (layoutOpts, renderOpts) => renderDag({ data, opts:layoutOpts }, renderOpts); + const getOpts = () => ({ showIndexing:showIndexing.toggle.checked, showCallSrc:showCallSrc.toggle.checked, showSink:showSink.toggle.checked, + callSrcMask:state.callSrcMask, expandedNodes:state.expandedNodes }); + render(getOpts(), { recenter:currentRewrite === 0 }); + showIndexing.toggle.onchange = () => render(getOpts(), { recenter:true }); + showCallSrc.toggle.onchange = () => { state.callSrcMask.clear(); render(getOpts(), { recenter:true }); } + showSink.toggle.onchange = () => render(getOpts(), { recenter:true }); // ** right sidebar metadata metadata.innerHTML = ""; if (ckey.includes("rewrites")) metadata.append(showIndexing.label, showCallSrc.label, showSink.label); diff --git a/tinygrad/viz/js/worker.js b/tinygrad/viz/js/worker.js index 645cca58b1677..a809de89ec740 100644 --- a/tinygrad/viz/js/worker.js +++ b/tinygrad/viz/js/worker.js @@ -31,7 +31,7 @@ const layoutCfg = (g, { blocks, paths, pc_tokens }) => { width = Math.max(width, ctx.measureText(tokens.map((t) => t.st).join("")).width); height += lineHeight; } - g.setNode(lead, { ...rectDims(width, height), label, id:lead, color:"#1a1b26" }); + g.setNode(lead, { ...rectDims(width, height), label, labelX:0, id:lead, color:"#1a1b26", addrspace:null }); } // paths become edges between basic blocks const pathColors = {0:"#3f7564", 1:"#7a4540", 2:"#3b5f7e"}; @@ -45,9 +45,9 @@ const layoutUOp = (g, { graph, change }, opts) => { const lineHeight = 14; g.setGraph({ rankdir: "LR", font:"sans-serif", lh:lineHeight }); ctx.font = `350 ${lineHeight}px ${g.graph().font}`; - if (change?.length) g.setNode("overlay", {label:"", labelWidth:0, labelHeight:0, className:"overlay"}); + if (change?.length) g.setNode("overlay", {label:"", labelWidth:0, labelHeight:0, labelX:0, className:"overlay"}); let callCount = 0; - for (const [k, {label, src, ref, color, tag }] of Object.entries(graph)) { + for (const [k, {label, src, ref, color, tag, exclude, addrspace}] of Object.entries(graph)) { // adjust node dims by label size (excluding escape codes) + add padding let [width, height] = [0, 0]; for (line of label.replace(/\u001B\[(?:K|.*?m)/g, "").split("\n")) { @@ -56,7 +56,7 @@ const layoutUOp = (g, { graph, change }, opts) => { } const callNode = label.startsWith("CALL\n") || label.startsWith("FUNCTION\n"); if (callNode) callCount++; - g.setNode(k, {...rectDims(width, height), label, ref, id:k, color, tag, callNode}); + g.setNode(k, {...rectDims(width, height), label, labelX:0, ref, id:k, color, tag, callNode, exclude, addrspace}); // add edges const edgeCounts = {}; for (const [_, s] of src) edgeCounts[s] = (edgeCounts[s] || 0)+1; @@ -77,33 +77,38 @@ const layoutUOp = (g, { graph, change }, opts) => { if (node.label.includes("dtypes.weakint")) g.removeNode(n); } } - if (!opts.showCallSrc || opts.callSrcMask.size > 0) { - // remove edges from src[0] to CALL nodes, track affected nodes - const disconnected = new Set(); - for (const n of g.nodes()) { - const node = g.node(n); - if (node.callNode && (opts.showCallSrc ? opts.callSrcMask.has(n) : !opts.callSrcMask.has(n))) { - node.collapsed = true; - for (const pred of (g.predecessors(n) || [])) { - const edge = g.edge(pred, n); - if (edge?.label?.text === 0) { - g.removeEdge(pred, n); - disconnected.add(pred); - } - } - } + // optionally remove node srcs, track affected nodes + const disconnected = new Set(); + const CALL_TAG_WIDTH = 14; + for (const n of g.nodes()) { + const node = g.node(n); + for (const consumerId of (g.successors(n) || [])) { + 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; + 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; } + // 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; + consumer.collapsed = true; + g.removeEdge(n, consumerId); + disconnected.add(n); } - // remove nodes that are now disconnected (no successors), only from affected subtree - let changed = true; - while (changed) { - changed = false; - for (const n of disconnected) { - if (!g.hasNode(n)) continue; - if ((g.successors(n) || []).length === 0) { - for (const pred of (g.predecessors(n) || [])) disconnected.add(pred); - g.removeNode(n); - changed = true; - } + } + // remove nodes that are now disconnected (no successors), only from affected subtree + let changed = true; + while (changed) { + changed = false; + for (const n of disconnected) { + if (!g.hasNode(n)) continue; + if ((g.successors(n) || []).length === 0) { + for (const pred of (g.predecessors(n) || [])) disconnected.add(pred); + g.removeNode(n); + changed = true; } } } diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index 11366f516dea4..6ba40c0f2c583 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -43,19 +43,22 @@ def stream_json(self, source:Generator): from tinygrad.uop.ops import KernelInfo from tinygrad.uop.render import print_uops, pyrender from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry, ProfileProgramEvent -from tinygrad.dtype import dtypes +from tinygrad.dtype import dtypes, AddrSpace -uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0", Ops.VCONST: "#e0e0e0", Ops.REDUCE: "#FF5B5B", +uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0", Ops.REDUCE: "#FF5B5B", **{x:"#f2cb91" for x in {Ops.DEFINE_LOCAL, Ops.DEFINE_REG}}, Ops.SHAPED_WMMA: "#FF5B5B", Ops.RANGE: "#c8a0e0", Ops.BARRIER: "#ff8080", Ops.IF: "#c8b0c0", Ops.SPECIAL: "#c0c0ff", - Ops.INDEX: "#cef263", Ops.WMMA: "#efefc0", Ops.MULTI: "#f6ccff", Ops.INS: "#eec4ff", + Ops.INDEX: "#D8F9E4", Ops.STACK: "#D8F9E4", + Ops.WMMA: "#efefc0", Ops.MULTI: "#f6ccff", Ops.INS: "#eec4ff", **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80", - Ops.BUFFER_VIEW: "#E5EAFF", Ops.BUFFER: "#B0BDFF", Ops.COPY: "#a040a0", Ops.CUSTOM_FUNCTION: "#bf71b6", - Ops.CALL: "#00B7C8", Ops.FUNCTION: "#C07788", Ops.PARAM: "#14686F", Ops.PATCH: "#7AA5AB", Ops.SOURCE: "#c0c0c0", Ops.BINARY: "#404040", + Ops.SLICE: "#E5EAFF", Ops.BUFFER: "#B0BDFF", Ops.GETADDR: "#9DB1F0", Ops.COPY: "#a040a0", Ops.CUSTOM_FUNCTION: "#bf71b6", + Ops.CALL: "#00B7C8", Ops.FUNCTION: "#C07788", Ops.PARAM: "#14686F", Ops.SOURCE: "#c0c0c0", Ops.BINARY: "#404040", Ops.LINEAR: "#7DF4FF", Ops.ALLREDUCE: "#ff40a0", Ops.MSELECT: "#d040a0", Ops.MSTACK: "#d040a0", Ops.CONTIGUOUS: "#FFC14D", Ops.STAGE: "#AC640D", Ops.REWRITE_ERROR: "#ff2e2e", Ops.AFTER: "#8A7866", Ops.END: "#524C46"} +addrspace_colors = {AddrSpace.REG:"#e68181", AddrSpace.LOCAL:"#e7c86a", AddrSpace.GLOBAL:"#75bd7b"} + # VIZ API # A step is a lightweight descriptor for a trace entry @@ -118,15 +121,15 @@ def uop_to_json(data:VizData, x:UOp) -> dict[int, dict]: # always exclude DEVICE/CONST/UNIQUE if u.op in {Ops.DEVICE, Ops.CONST, Ops.UNIQUE, Ops.LUNIQUE} and u is not x: excluded.add(u) if u.op is Ops.CONST and len(u.src) and u.src[0].op in {Ops.UNIQUE, Ops.LUNIQUE}: excluded.remove(u) - if u.op is Ops.VCONST and u.dtype.scalar() == dtypes.weakint and u is not x: excluded.add(u) if u.op is Ops.STACK and len(u.src) == 0: excluded.add(u) # exclude RESHAPE/EXPAND that only serve to broadcast a CONST if u.op in {Ops.RESHAPE, Ops.EXPAND} and len(u.src) >= 1 and u.src[0] in excluded and u is not x: excluded.add(u) + if u.op in GroupOp.Movement: excluded.update(s for s in u.src if s.op is Ops.STACK and all(x.op is Ops.CONST for x in s.src)) for u in toposort: - if u in excluded: continue argst = codecs.decode(str(u.arg), "unicode_escape") if u.op in GroupOp.Movement: argst = (mask_to_str if u.op in {Ops.SHRINK, Ops.PAD} else shape_to_str)(u.marg) if u.op is Ops.BINARY: argst = f"<{len(u.arg)} bytes>" + if u.op is Ops.CONST and dtypes.is_float(u.dtype): argst = f"{u.arg:g}" wrap_len = 200 if u.op is Ops.SOURCE else 80 label = f"{str(u.op).split('.')[1]}{(chr(10)+word_wrap(argst.replace(':', ''), wrap=wrap_len)) if u.arg is not None else ''}" if u.dtype != dtypes.void: label += f"\n{u.dtype}" @@ -160,8 +163,9 @@ def uop_to_json(data:VizData, x:UOp) -> dict[int, dict]: # limit SOURCE labels line count if u.op is Ops.SOURCE and len(lines:=label.split("\n")) > 40: label = "\n".join(lines[:30]) + "\n..." - graph[id(u)] = {"label":label, "src":[(i,id(x)) for i,x in enumerate(u.src) if x not in excluded], "color":uops_colors.get(u.op, "#ffffff"), - "ref":ref, "tag":repr(u.tag) if u.tag is not None else None} + graph[id(u)] = {"label":label, "src":[(i,id(x)) for i,x in enumerate(u.src)], "exclude":u in excluded, "color":uops_colors.get(u.op, "#ffffff"), + "ref":ref, "tag":repr(u.tag) if u.tag is not None else None, + "addrspace":addrspace_colors.get(u.addrspace, None) if u.addrspace is not None else None} return graph def _reconstruct(data:VizData, a:int, depth:int|None=None):