diff --git a/.github/actions/setup-tinygrad/action.yml b/.github/actions/setup-tinygrad/action.yml index fd5677a978..cf26a3e14a 100644 --- a/.github/actions/setup-tinygrad/action.yml +++ b/.github/actions/setup-tinygrad/action.yml @@ -253,8 +253,13 @@ runs: git checkout b16039dc940dc6bc4ea0a98380495769ff35ed99 mkdir build cd build - cmake .. -Wno-dev -G Ninja -DOCELOT_BUILD_TOOLS=OFF -DCMAKE_BUILD_ALWAYS=0 -DBUILD_TESTS_CUDA=OFF \ - -DBoost_INCLUDE_DIR=$(brew --prefix boost)/include -DBoost_LIBRARY_DIR=$(brew --prefix boost)/lib -DCMAKE_POLICY_VERSION_MINIMUM=3.5 + + CMAKE_ARGS="-Wno-dev -G Ninja -DOCELOT_BUILD_TOOLS=OFF -DCMAKE_BUILD_ALWAYS=0 -DBUILD_TESTS_CUDA=OFF -DCMAKE_POLICY_VERSION_MINIMUM=3.5" + if [[ "${{ runner.os }}" == "macOS" ]]; then + CMAKE_ARGS="$CMAKE_ARGS -DBoost_INCLUDE_DIR=$(brew --prefix boost)/include -DBoost_LIBRARY_DIR=$(brew --prefix boost)/lib" + fi + + cmake .. $CMAKE_ARGS ninja - name: Install gpuocelot if: inputs.ocelot == 'true' diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index 67aeb754c0..98688461c8 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -28,7 +28,7 @@ jobs: # since sudo is required for usbgpu on macos, move the cache to a new location, as some of the files are owned by root PYTHONPYCACHEPREFIX: /tmp/tiny_python_pycache runs-on: [self-hosted, macOS] - timeout-minutes: 20 + timeout-minutes: 60 defaults: run: shell: bash -e -o pipefail {0} @@ -160,7 +160,7 @@ jobs: testnvidiabenchmark: name: tinybox green Benchmark runs-on: [self-hosted, Linux, tinyboxgreen] - timeout-minutes: 30 + timeout-minutes: 60 defaults: run: shell: bash -e -o pipefail {0} @@ -197,14 +197,14 @@ jobs: - name: Test tensor cores run: | NV=1 ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py - PTX=1 ALLOW_TF32=1 NV=1 python3 test/opt/test_tensor_cores.py + NV=1 NV_PTX=1 ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py - name: Run Tensor Core GEMM (CUDA) run: | CUDA=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul.txt CUDA=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_bfloat16.txt CUDA=1 SHOULD_USE_TC=1 ALLOW_TF32=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_tf32.txt - name: Run Tensor Core GEMM (PTX) - run: NV=1 PTX=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_ptx.txt + run: NV=1 NV_PTX=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_ptx.txt - name: Run Tensor Core GEMM (NV) run: NV=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_nv.txt - name: Test NV=1 @@ -274,7 +274,7 @@ jobs: testmorenvidiabenchmark: name: tinybox green Training Benchmark runs-on: [self-hosted, Linux, tinyboxgreen] - timeout-minutes: 20 + timeout-minutes: 60 defaults: run: shell: bash -e -o pipefail {0} @@ -302,7 +302,7 @@ jobs: - name: Fuzz Padded Tensor Core GEMM (NV) run: NV=1 M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py - name: Fuzz Padded Tensor Core GEMM (PTX) - run: NV=1 PTX=1 M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py + run: NV=1 NV_PTX=1 M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py - name: Train MNIST run: time PYTHONPATH=. NV=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt - name: Run 10 CIFAR training steps @@ -346,7 +346,7 @@ jobs: testamdbenchmark: name: tinybox red Benchmark runs-on: [self-hosted, Linux, tinybox] - timeout-minutes: 20 + timeout-minutes: 60 defaults: run: shell: bash -e -o pipefail {0} @@ -476,7 +476,7 @@ jobs: testmoreamdbenchmark: name: tinybox red Training Benchmark runs-on: [self-hosted, Linux, tinybox] - timeout-minutes: 30 + timeout-minutes: 60 defaults: run: shell: bash -e -o pipefail {0} @@ -511,8 +511,8 @@ jobs: run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=85 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt - name: Run 10 CIFAR training steps w HALF run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=188 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt - - name: Run 10 CIFAR training steps w BF16 - run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt + # - name: Run 10 CIFAR training steps w BF16 + # run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt - name: Run 10 CIFAR training steps w winograd run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=66 AMD=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt - name: Run full CIFAR training w 1 GPU @@ -539,7 +539,7 @@ jobs: testmlperfamdbenchmark: name: tinybox red MLPerf Benchmark runs-on: [self-hosted, Linux, tinybox] - timeout-minutes: 30 + timeout-minutes: 60 defaults: run: shell: bash -e -o pipefail {0} @@ -645,7 +645,7 @@ jobs: testreddriverbenchmark: name: AM Benchmark runs-on: [self-hosted, Linux, tinyboxrandom] - timeout-minutes: 15 + timeout-minutes: 20 defaults: run: shell: bash -e -o pipefail {0} @@ -716,7 +716,7 @@ jobs: testgreendriverbenchmark: name: NV Benchmark runs-on: [self-hosted, Linux, tinyboxrandom] - timeout-minutes: 15 + timeout-minutes: 20 defaults: run: shell: bash -e -o pipefail {0} diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 61bf5fdb8f..9254c27bd4 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -30,8 +30,6 @@ jobs: key: llvm-speed deps: testing_minimal llvm: 'true' - - name: External Benchmark Schedule - run: python3 test/external/external_benchmark_schedule.py - name: Speed Test run: CPU=1 CPU_LLVM=1 python3 test/speed/external_test_speed_v_torch.py - name: Speed Test (BEAM=2) @@ -48,7 +46,7 @@ jobs: uses: ./.github/actions/setup-tinygrad with: deps: docs - pydeps: "capstone" + pydeps: "capstone torch" - name: Build wheel and show size run: | pip install build @@ -79,6 +77,8 @@ jobs: run: | python docs/abstractions2.py python docs/abstractions3.py + - name: Test README + run: awk '/```python/{flag=1;next}/```/{flag=0}flag' README.md > README.py && python README.py - name: Test Quickstart run: awk '/```python/{flag=1;next}/```/{flag=0}flag' docs/quickstart.md > quickstart.py && python quickstart.py - name: Test DEBUG @@ -93,8 +93,6 @@ jobs: name: Torch Backend Tests runs-on: ubuntu-latest timeout-minutes: 15 - env: - IGNORE_OOB: 0 steps: - name: Checkout Code uses: actions/checkout@v4 @@ -132,8 +130,6 @@ jobs: name: Torch Backend Tests More runs-on: ubuntu-latest timeout-minutes: 15 - env: - IGNORE_OOB: 0 steps: - name: Checkout Code uses: actions/checkout@v4 @@ -152,20 +148,28 @@ jobs: - name: Test some torch tests (expect failure) run: python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true - tc: - name: Tensor Core tests + bepython: + name: Python Backend runs-on: ubuntu-latest - timeout-minutes: 10 - env: - IGNORE_OOB: 0 + timeout-minutes: 15 steps: - name: Checkout Code uses: actions/checkout@v4 - name: Setup Environment uses: ./.github/actions/setup-tinygrad with: - key: uops-minimal + key: be-minimal deps: testing_minimal + - name: Test dtype with Python emulator + run: DEBUG=1 PYTHON=1 python3 -m pytest -n=auto test/test_dtype.py test/test_dtype_alu.py + - name: Test ops with Python emulator + run: DEBUG=2 SKIP_SLOW_TEST=1 PYTHON=1 python3 -m pytest -n=auto test/test_ops.py --durations=20 + - name: Test uops with Python emulator + run: PYTHON=1 python3 -m pytest test/test_uops.py --durations=20 + - name: Test symbolic with Python emulator + run: PYTHON=1 python3 test/test_symbolic_ops.py + - name: test_renderer_failures with Python emulator + run: PYTHON=1 python3 -m pytest -rA test/test_renderer_failures.py::TestRendererFailures - name: Test IMAGE=2 support run: | IMAGE=2 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm @@ -212,31 +216,6 @@ jobs: DEBUG=2 EMULATE=INTEL PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf DEBUG=2 AMX=1 EMULATE=AMX PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStats.test_simple_matmul - bepython: - name: Python Backend - runs-on: ubuntu-latest - timeout-minutes: 10 - env: - IGNORE_OOB: 0 - steps: - - name: Checkout Code - uses: actions/checkout@v4 - - name: Setup Environment - uses: ./.github/actions/setup-tinygrad - with: - key: be-minimal - deps: testing_minimal - - name: Test dtype with Python emulator - run: DEBUG=1 PYTHON=1 python3 -m pytest -n=auto test/test_dtype.py test/test_dtype_alu.py - - name: Test ops with Python emulator - run: DEBUG=2 SKIP_SLOW_TEST=1 PYTHON=1 python3 -m pytest -n=auto test/test_ops.py --durations=20 - - name: Test uops with Python emulator - run: PYTHON=1 python3 -m pytest test/test_uops.py --durations=20 - - name: Test symbolic with Python emulator - run: PYTHON=1 python3 test/test_symbolic_ops.py - - name: test_renderer_failures with Python emulator - run: PYTHON=1 python3 -m pytest -rA test/test_renderer_failures.py::TestRendererFailures - linter: name: Linters runs-on: ubuntu-latest @@ -282,8 +261,6 @@ jobs: key: unittest-12 pydeps: "pillow" deps: testing_unit - - name: Test README - run: awk '/```python/{flag=1;next}/```/{flag=0}flag' README.md > README.py && python README.py - name: Run unit tests run: python -m pytest -n=auto test/unit/ --durations=20 - name: Run targetted tests on NULL backend @@ -295,6 +272,8 @@ jobs: # run: NULL=1 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing - name: Run GC tests run: python test/external/external_uop_gc.py + - name: External Benchmark Schedule + run: python3 test/external/external_benchmark_schedule.py - name: Run process replay tests uses: ./.github/actions/process-replay - name: Regen dataset on test_tiny @@ -330,12 +309,10 @@ jobs: - name: Fuzz Test shape ops run: python test/external/fuzz_shape_ops.py - testgpuimage: - name: 'GPU IMAGE Tests' + testopenclimage: + name: CL IMAGE Tests runs-on: ubuntu-22.04 timeout-minutes: 10 - env: - IGNORE_OOB: 0 steps: - name: Checkout Code uses: actions/checkout@v4 @@ -345,19 +322,17 @@ jobs: key: gpu-image deps: testing_minimal opencl: 'true' - - name: Test GPU IMAGE=2 ops + training + - name: Test CL IMAGE=2 ops + training run: | - GPU=1 IMAGE=2 python -m pytest -n=auto test/test_ops.py --durations=20 - GPU=1 IMAGE=2 python test/models/test_end2end.py TestEnd2End.test_linear_mnist + CL=1 IMAGE=2 python -m pytest -n=auto test/test_ops.py --durations=20 + CL=1 IMAGE=2 python test/models/test_end2end.py TestEnd2End.test_linear_mnist - name: Run process replay tests uses: ./.github/actions/process-replay testgpumisc: - name: 'GPU Misc tests' + name: CL Misc tests runs-on: ubuntu-22.04 timeout-minutes: 10 - env: - IGNORE_OOB: 0 steps: - name: Checkout Code uses: actions/checkout@v4 @@ -368,11 +343,11 @@ jobs: deps: testing_minimal opencl: 'true' - name: Generate Dataset - run: GPU=1 extra/optimization/generate_dataset.sh + run: CL=1 extra/optimization/generate_dataset.sh - name: Run Kernel Count Test - run: GPU=1 python -m pytest -n=auto test/external/external_test_opt.py + run: CL=1 python -m pytest -n=auto test/external/external_test_opt.py - name: Run fused optimizer tests - run: GPU=1 FUSE_OPTIM=1 python -m pytest -n=auto test/models/test_mnist.py + run: CL=1 FUSE_OPTIM=1 python -m pytest -n=auto test/models/test_mnist.py - name: Upload artifact uses: actions/upload-artifact@v4 with: @@ -380,11 +355,9 @@ jobs: path: /tmp/sops.gz testopenpilot: - name: 'openpilot Compile Tests' + name: openpilot Compile Tests runs-on: ubuntu-22.04 timeout-minutes: 15 - env: - IGNORE_OOB: 0 steps: - name: Checkout Code uses: actions/checkout@v4 @@ -397,26 +370,26 @@ jobs: llvm: 'true' - name: Test openpilot model kernel count and gate usage run: | - ALLOWED_KERNEL_COUNT=208 ALLOWED_READ_IMAGE=2175 ALLOWED_GATED_READ_IMAGE=16 FLOAT16=0 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx + ALLOWED_KERNEL_COUNT=208 ALLOWED_READ_IMAGE=2175 ALLOWED_GATED_READ_IMAGE=16 FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx - name: Test openpilot alt model correctness (float32) - run: FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/3799fe46b3a629e491d4b8498b8ae83e4c88c304/selfdrive/modeld/models/supercombo.onnx + run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/3799fe46b3a629e491d4b8498b8ae83e4c88c304/selfdrive/modeld/models/supercombo.onnx - name: Test openpilot fastvits model correctness (float32) - run: FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx + run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx # - name: Test openpilot simple_plan vision model correctness (float32) - # run: FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/35ff4f4577002f2685e50c8346addae33fe8da27a41dd4d6a0f14d1f4b1af81b + # run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/35ff4f4577002f2685e50c8346addae33fe8da27a41dd4d6a0f14d1f4b1af81b - name: Test openpilot LLVM compile run: CPU=1 CPU_LLVM=1 LLVMOPT=1 JIT=2 BEAM=0 IMAGE=0 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx - name: Test openpilot compile4 - run: NOLOCALS=1 GPU=1 IMAGE=2 FLOAT16=1 DEBUG=2 python3 examples/openpilot/compile4.py + run: NOLOCALS=1 CL=1 IMAGE=2 FLOAT16=1 DEBUG=2 python3 examples/openpilot/compile4.py - name: Run process replay tests uses: ./.github/actions/process-replay +# ****** ONNX Tests ****** + testonnxcpu: - name: 'ONNX (CPU) Tests' + name: ONNX (CPU) Tests runs-on: ubuntu-22.04 timeout-minutes: 20 - env: - IGNORE_OOB: 0 steps: - name: Checkout Code @@ -442,12 +415,9 @@ jobs: uses: ./.github/actions/process-replay testopencl: - name: 'ONNX (GPU)+Optimization Tests' + name: ONNX (CL)+Optimization Tests runs-on: ubuntu-22.04 timeout-minutes: 20 - env: - IGNORE_OOB: 0 - steps: - name: Checkout Code uses: actions/checkout@v4 @@ -459,16 +429,16 @@ jobs: pydeps: "tensorflow==2.15.1 tensorflow_addons" python-version: '3.11' opencl: 'true' - - name: Test ONNX (GPU) - run: GPU=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20 + - name: Test ONNX (CL) + run: CL=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20 #- name: Test Optimization Helpers # run: DEBUG=1 python3 extra/optimization/test_helpers.py #- name: Test Action Space - # run: DEBUG=1 GPU=1 python3 extra/optimization/get_action_space.py + # run: DEBUG=1 CL=1 python3 extra/optimization/get_action_space.py - name: Test Beam Search - run: GPU=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py + run: CL=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py - name: Test MLPerf stuff - run: GPU=1 python -m pytest -n=auto test/external/external_test_optim.py test/external/external_test_losses.py test/external/external_test_metrics.py test/external/external_test_datasets.py --durations=20 + run: CL=1 python -m pytest -n=auto test/external/external_test_optim.py test/external/external_test_losses.py test/external/external_test_metrics.py test/external/external_test_datasets.py --durations=20 - name: Test llama 3 training run: MAX_BUFFER_SIZE=0 DEV=NULL SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=8 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py - name: Run process replay tests @@ -488,12 +458,12 @@ jobs: - name: Test 1B LLM run: echo "What's a male chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm | grep -i rooster +# ****** Models Tests ****** + testmodels: name: Models (llvm+cpu+gpu) runs-on: ubuntu-22.04 timeout-minutes: 15 - env: - IGNORE_OOB: 0 steps: - name: Checkout Code uses: actions/checkout@v4 @@ -506,15 +476,35 @@ jobs: llvm: 'true' - name: Test models (llvm) run: CPU=1 CPU_LLVM=1 python -m pytest -n=auto test/models --durations=20 - - name: Test models (gpu) - run: GPU=1 python -m pytest -n=auto test/models --durations=20 + - name: Test models (opencl) + run: CL=1 python -m pytest -n=auto test/models --durations=20 - name: Test models (cpu) run: CPU=1 CPU_LLVM=0 python -m pytest -n=auto test/models --durations=20 - name: Run process replay tests uses: ./.github/actions/process-replay - testrangeify: - name: Linux (rangeify) + testmetalmodels: + name: Models (metal) + runs-on: macos-14 + timeout-minutes: 20 + steps: + - name: Checkout Code + uses: actions/checkout@v4 + - name: Setup Environment + uses: ./.github/actions/setup-tinygrad + with: + key: metal + deps: testing + python-version: '3.11' + - name: Test models (Metal) + run: METAL=1 python -m pytest -n=auto test/models --durations=20 + - name: Test LLaMA compile speed + run: METAL=1 python test/external/external_test_speed_llama.py + +# ****** Feature Tests ****** + + testrangeifycpu: + name: Linux (rangeify) CPU runs-on: ubuntu-24.04 timeout-minutes: 15 steps: @@ -525,28 +515,80 @@ jobs: with: key: rangeify-minimal-llvm deps: testing_minimal + opencl: 'true' llvm: "true" - name: Test CPU=1 RANGEIFY=1 # TODO: add more passing tests here - # test_symbolic_arange_sym_step is passing now - # test_threefry_doesnt_use_long is because there's a contig after the long now + # test_load_state_dict_sharded_model_dict_same_axis issue with multi + # test_instancenorm_3d is very slow run: | CPU=1 CPU_LLVM=0 RANGEIFY=1 python3 -m pytest -n auto --durations 20 \ - -k "not test_symbolic_arange_sym_step and not test_threefry_doesnt_use_long" \ - test/test_tiny.py test/test_rangeify.py test/test_ops.py test/test_tensor_variable.py \ - test/test_outerworld_range.py test/test_sample.py test/test_randomness.py + -k "not test_load_state_dict_sharded_model_dict_same_axis and not test_instancenorm_3d" \ + test/test_tiny.py test/test_rangeify.py test/test_ops.py test/test_symbolic_ops.py test/test_symbolic_jit.py test/test_tensor_variable.py \ + test/test_outerworld_range.py test/test_randomness.py test/test_nn.py test/test_arange.py test/test_tensor.py test/test_optim.py \ + test/test_setitem.py + - name: Test const folding + run: CPU=1 RANGEIFY=1 python3 -m pytest -n auto --durations 20 test/test_const_folding.py -k "not test_cast_padded and not TestReduceOpsConstFolding and not TestMultiConstFolding" + - name: Test multitensor + run: | + CPU=1 RANGEIFY=1 python3 test/test_multitensor.py TestMultiTensor.test_matmul_shard_1_1 TestMultiTensor.test_simple_add_W TestMultiTensor.test_simple_reduce \ + TestMultiTensor.test_elementwise_dtype TestMultiTensor.test_shard_no_recompile TestHandleData.test_copied_to_device TestMultiRamUsage + CPU=1 RANGEIFY=1 python3 -m pytest test/test_multitensor.py::TestMultiAssign -k 'not (multi_assign_piece_noncontig or multi_assign_var_offset)' + CPU=1 RANGEIFY=1 python3 -m pytest -n=auto test/test_multitensor.py::TestMultiTensor test/unit/test_allreduce.py -k 'not const_folding' - name: Test CPU=1 RANGEIFY=2 run: CPU=1 CPU_LLVM=0 RANGEIFY=2 python3 -m pytest -n auto test/test_tiny.py test/test_rangeify.py test/test_ops.py --durations 20 # slow (and still wrong on beautiful_mnist) - #- name: Test LLVM=1 RANGEIFY=1 (slow tests) + #- name: Test LLVM RANGEIFY=1 (slow tests) # run: CPU=1 CPU_LLVM=1 RANGEIFY=1 python3 -m pytest -n auto test/models/test_mnist.py --durations 20 + - name: Run process replay tests + uses: ./.github/actions/process-replay + + testrangeifycl: + name: Linux (rangeify) CL + runs-on: ubuntu-24.04 + timeout-minutes: 15 + steps: + - name: Checkout Code + uses: actions/checkout@v4 + - name: Setup Environment + uses: ./.github/actions/setup-tinygrad + with: + key: rangeify-cl + deps: testing + opencl: 'true' + llvm: "true" + - name: Test CL=1 RANGEIFY=1 + run: CL=1 RANGEIFY=1 pytest -n auto test/test_ops.py test/test_schedule.py test/test_symbolic_ops.py test/test_jit.py test/unit/test_disk_tensor.py test/models/test_mnist.py test/unit/test_mnist_dataset.py test/test_optim.py --durations 20 + - name: Test Fuse + run: CL=1 RANGEIFY=2 python3 -m pytest --durations 20 test/test_softmax_fusion.py -k "not test_auto_softmax" + - name: Test ONNX + run: CL=1 RANGEIFY=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20 + - name: Run process replay tests + uses: ./.github/actions/process-replay + + testrangeifymacos: + name: MacOS (rangeify) + runs-on: macos-14 + timeout-minutes: 15 + steps: + - name: Checkout Code + uses: actions/checkout@v4 + - name: Setup Environment + uses: ./.github/actions/setup-tinygrad + with: + key: metal + deps: testing + - name: some unit tests + run: METAL=1 RANGEIFY=1 python -m pytest -n=auto test/unit/test_winograd.py --durations=20 + - name: Test METAL=1 RANGEIFY=1 + run: METAL=1 RANGEIFY=1 python -m pytest -n=auto test/test_ops.py --durations=20 + - name: Run process replay tests + uses: ./.github/actions/process-replay testdevectorize: name: Linux (devectorize) runs-on: ubuntu-24.04 timeout-minutes: 15 - env: - IGNORE_OOB: 0 steps: - name: Checkout Code uses: actions/checkout@v4 @@ -568,8 +610,6 @@ jobs: name: Linux (DSP) runs-on: ubuntu-24.04 timeout-minutes: 15 - env: - IGNORE_OOB: 0 steps: - name: Checkout Code uses: actions/checkout@v4 @@ -636,7 +676,6 @@ jobs: runs-on: ubuntu-22.04 timeout-minutes: 20 env: - IGNORE_OOB: 0 AMD: 1 MOCKGPU: 1 FORWARD_ONLY: 1 @@ -666,7 +705,7 @@ jobs: run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20 - name: Run TestOps.test_add with SQTT run: | - PROFILE=1 SQTT=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add + VIZ=1 SQTT=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -o /tmp/gpu0.rgp - name: Run process replay tests uses: ./.github/actions/process-replay @@ -694,7 +733,7 @@ jobs: cuda: 'true' ocelot: 'true' - name: Set env - run: printf "${{ matrix.backend == 'PTX' && 'CUDA=1\nPTX=1' || matrix.backend == 'nv' && 'NV=1\nSKIP_SLOW_TEST=1' }}" >> $GITHUB_ENV + run: printf "${{ matrix.backend == 'PTX' && 'CUDA=1\nCUDA_PTX=1' || matrix.backend == 'nv' && 'NV=1\nSKIP_SLOW_TEST=1' }}" >> $GITHUB_ENV - name: Check Device.DEFAULT and print some source run: | python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CUDA','NV'], Device.DEFAULT" @@ -705,18 +744,15 @@ jobs: - name: Run process replay tests uses: ./.github/actions/process-replay - tests: + testcpuopencl: strategy: fail-fast: false matrix: - backend: [llvm, cpu, gpu] + backend: [llvm, cpu, opencl] name: Linux (${{ matrix.backend }}) runs-on: ubuntu-22.04 timeout-minutes: 20 - env: - IGNORE_OOB: 0 - steps: - name: Checkout Code uses: actions/checkout@v4 @@ -725,13 +761,13 @@ jobs: with: key: ${{ matrix.backend }}-minimal deps: testing_minimal - opencl: ${{ matrix.backend == 'gpu' && 'true' }} + opencl: ${{ matrix.backend == 'opencl' && 'true' }} llvm: ${{ matrix.backend == 'llvm' && 'true' }} - name: Set env - run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'gpu' && 'GPU=1' }}" >> $GITHUB_ENV + run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'opencl' && 'CL=1' }}" >> $GITHUB_ENV - name: Check Device.DEFAULT and print some source run: | - python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CPU','GPU'], Device.DEFAULT" + python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CPU','CL'], Device.DEFAULT" DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add - name: Run pytest (${{ matrix.backend }}) run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20 @@ -772,7 +808,7 @@ jobs: start_server "remote-server-amd-1" "AMD" 6667 start_server "remote-server-amd-2" "AMD" 6668 - start_server "remote-server-gpu" "GPU" 7667 + start_server "remote-server-gpu" "CL" 7667 start_server "remote-server-cpu" "CPU" 8667 - name: Check Device.DEFAULT and print some source env: @@ -786,7 +822,7 @@ jobs: HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6 run: | python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_remote.py test/test_tensor_variable.py --durations 20 - - name: Run REMOTE=1 Test (GPU) + - name: Run REMOTE=1 Test (CL) env: HOST: 127.0.0.1:7667*6 run: | @@ -807,48 +843,42 @@ jobs: # ****** OSX Tests ****** - testmetal2: + testmetal: name: MacOS (unit) runs-on: macos-14 timeout-minutes: 20 - env: - IGNORE_OOB: 0 - steps: - name: Checkout Code uses: actions/checkout@v4 - name: Setup Environment uses: ./.github/actions/setup-tinygrad with: - key: metal2 + key: metal deps: testing python-version: '3.11' amd: 'true' cuda: 'true' ocelot: 'true' llvm: 'true' - - name: Run real world test - run: METAL=1 python -m pytest -n=auto test/models/test_real_world.py --durations=20 - - name: Test models (Metal) - run: METAL=1 python -m pytest -n=auto test/models -v --durations=20 + - name: Run unit tests + run: METAL=1 python -m pytest -n=auto test/unit/ --durations=20 - name: Run ONNX run: METAL=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20 - name: Test tensor core ops (fake) - run: TC=2 METAL=1 DEBUG=3 python test/test_ops.py TestOps.test_gemm + run: METAL=1 DEBUG=3 TC=2 python test/test_ops.py TestOps.test_gemm - name: Test tensor core ops (real) run: METAL=1 DEBUG=3 python test/test_ops.py TestOps.test_big_gemm - - name: Test LLaMA compile speed - run: METAL=1 python test/external/external_test_speed_llama.py - name: Test Beam Search run: METAL=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py #- name: Fuzz Test linearizer # run: METAL=1 DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py - name: Run TRANSCENDENTAL math - run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20 + run: METAL=1 TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20 - name: Run pytest (amd) env: MOCKGPU: 1 AMD: 1 + AMD_LLVM: 0 FORWARD_ONLY: 1 run: | python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20 @@ -856,13 +886,14 @@ jobs: env: MOCKGPU: 1 AMD: 1 + AMD_LLVM: 1 FORWARD_ONLY: 1 run: | python -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py test/device/test_amd_llvm.py --durations=20 - name: Run pytest (ptx) env: MOCKGPU: 1 - PTX: 1 + NV_PTX: 1 NV: 1 FORWARD_ONLY: 1 run: | @@ -939,8 +970,6 @@ jobs: name: MacOS (${{ matrix.backend }}) runs-on: macos-15 timeout-minutes: 20 - env: - IGNORE_OOB: 0 steps: - name: Checkout Code uses: actions/checkout@v4 @@ -976,8 +1005,6 @@ jobs: name: Windows (${{ matrix.backend }}) runs-on: windows-latest timeout-minutes: 15 - env: - IGNORE_OOB: 0 steps: - name: Checkout Code uses: actions/checkout@v4 diff --git a/README.md b/README.md index 262ea97d37..dab378a23a 100644 --- a/README.md +++ b/README.md @@ -79,7 +79,7 @@ See [examples/beautiful_mnist.py](examples/beautiful_mnist.py) for the full vers tinygrad already supports numerous accelerators, including: -- [x] [GPU (OpenCL)](tinygrad/runtime/ops_gpu.py) +- [x] [OpenCL](tinygrad/runtime/ops_cl.py) - [x] [CPU](tinygrad/runtime/ops_cpu.py) - [x] [METAL](tinygrad/runtime/ops_metal.py) - [x] [CUDA](tinygrad/runtime/ops_cuda.py) diff --git a/docs/env_vars.md b/docs/env_vars.md index 9367eef064..f8844ba70b 100644 --- a/docs/env_vars.md +++ b/docs/env_vars.md @@ -3,7 +3,7 @@ This is a list of environment variable that control the runtime behavior of tinygrad and its examples. Most of these are self-explanatory, and are usually used to set an option at runtime. -Example: `GPU=1 DEBUG=4 python3 -m pytest` +Example: `CL=1 DEBUG=4 python3 -m pytest` However you can also decorate a function to set a value only inside that function. @@ -31,7 +31,7 @@ These control the behavior of core tinygrad even when used as a library. Variable | Possible Value(s) | Description ---|---|--- DEBUG | [1-7] | enable debugging output (operations, timings, speed, generated code and more) -GPU | [1] | enable the GPU (OpenCL) backend +CL | [1] | enable OpenCL backend CUDA | [1] | enable CUDA backend AMD | [1] | enable AMD backend NV | [1] | enable NV backend @@ -41,8 +41,6 @@ BEAM | [#] | number of beams in kernel beam search DEFAULT_FLOAT | [HALF, ...]| specify the default float dtype (FLOAT32, HALF, BFLOAT16, FLOAT64, ...), default to FLOAT32 IMAGE | [1-2] | enable 2d specific optimizations FLOAT16 | [1] | use float16 for images instead of float32 -PTX | [1] | enable the specialized [PTX](https://docs.nvidia.com/cuda/parallel-thread-execution/) assembler for Nvidia GPUs. If not set, defaults to generic CUDA codegen backend. -PROFILE | [1] | enable profiling. This feature is supported in NV, AMD, QCOM and METAL backends. VISIBLE_DEVICES | [list[int]]| restricts the NV/AMD devices that are available. The format is a comma-separated list of identifiers (indexing starts with 0). JIT | [0-2] | 0=disabled, 1=[jit enabled](quickstart.md#jit) (default), 2=jit enabled, but graphs are disabled VIZ | [1] | 0=disabled, 1=[viz enabled](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/viz) diff --git a/docs/runtime.md b/docs/runtime.md index bc85d9bedf..28a7aad010 100644 --- a/docs/runtime.md +++ b/docs/runtime.md @@ -2,17 +2,17 @@ tinygrad supports various runtimes, enabling your code to scale across a wide range of devices. The default runtime can be automatically selected based on the available hardware, or you can force a specific runtime to be default using environment variables (e.g., `CPU=1`). -| Runtime | Description | Requirements | -|---------|-------------|--------------| -| [NV](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_nv.py) | Provides acceleration for NVIDIA GPUs | Ampere/Ada series GPUs | -| [AMD](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_amd.py) | Provides acceleration for AMD GPUs | RDNA2/RDNA3/RDNA4 series GPUs. You can select one of the interfaces for communication by setting `AMD_IFACE=(KFD|PCI)`. See [AMD interfaces](#amd-interfaces) for more details. | -| [QCOM](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_qcom.py) | Provides acceleration for QCOM GPUs | 6xx series GPUs | -| [METAL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_metal.py) | Utilizes Metal for acceleration on Apple devices | M1+ Macs; Metal 3.0+ for `bfloat` support | -| [CUDA](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cuda.py) | Utilizes CUDA for acceleration on NVIDIA GPUs | NVIDIA GPU with CUDA support | -| [GPU (OpenCL)](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_gpu.py) | Accelerates computations using OpenCL on GPUs | OpenCL 2.0 compatible device | -| [CPU (C Code)](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cpu.py) | Runs on CPU using the clang compiler | `clang` compiler in system `PATH` | -| [LLVM (LLVM IR)](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_llvm.py) | Runs on CPU using the LLVM compiler infrastructure | llvm libraries installed and findable | -| [WEBGPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_webgpu.py) | Runs on GPU using the Dawn WebGPU engine (used in Google Chrome) | Dawn library installed and findable. Download binaries [here](https://github.com/wpmed92/pydawn/releases/tag/v0.3.0). | +| Runtime | Description | Compiler Options | Requirements | +|---------|-------------|------------------|--------------| +| [NV](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_nv.py) | Provides acceleration for NVIDIA GPUs | nvrtc (default)
PTX (`NV_PTX=1`) | Ampere/Ada/Blackwell series GPUs.
You can select an interface via `NV_IFACE=(NVK\|PCI)`. See [NV interfaces](#nv-interfaces) for details. | +| [AMD](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_amd.py) | Provides acceleration for AMD GPUs | LLVM (`AMD_LLVM=1`)
HIP/COMGR (`AMD_HIP=1`) | RDNA2 or newer GPUs.
You can select an interface via `AMD_IFACE=(KFD\|PCI\|USB)`. See [AMD interfaces](#amd-interfaces) for details. | +| [QCOM](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_qcom.py) | Provides acceleration for QCOM GPUs | - | 6xx series GPUs | +| [METAL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_metal.py) | Utilizes Metal for acceleration on Apple devices | - | M1+ Macs; Metal 3.0+ for `bfloat` support | +| [CUDA](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cuda.py) | Utilizes CUDA for acceleration on NVIDIA GPUs | nvrtc (default)
PTX (`CUDA_PTX=1`) | NVIDIA GPU with CUDA support | +| [CL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cl.py) | Accelerates computations using OpenCL on GPUs | - | OpenCL 2.0 compatible device | +| [CPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cpu.py) | Runs on CPU using the clang or llvm compiler | Clang JIT (default)
LLVM IR (`CPU_LLVM=1`) | `clang` compiler in system `PATH` | +| [WEBGPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_webgpu.py) | Runs on GPU using the Dawn WebGPU engine (used in Google Chrome) | - | Dawn library installed and discoverable. Binaries: [pydawn v0.3.0](https://github.com/wpmed92/pydawn/releases/tag/v0.3.0) | + ## Interoperability @@ -70,5 +70,12 @@ AMD backend supports several interfaces for communicating with devices: * `KFD`: uses the amdgpu driver * `PCI`: uses the [AM driver](developer/am.md) +* `USB`: USB3 interafce for asm24xx chips. You can force an interface by setting `AMD_IFACE` to one of these values. In the case of `AMD_IFACE=PCI`, this may unbind your GPU from the amdgpu driver. + +## NV Interfaces +NV backend supports several interfaces for communicating with devices: + +* `NVK`: uses the nvidia driver +* `PCI`: uses the [NV driver](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/support/nv/nvdev.py) diff --git a/examples/gpt2.py b/examples/gpt2.py index de577e911f..6670b4e2bb 100644 --- a/examples/gpt2.py +++ b/examples/gpt2.py @@ -189,7 +189,7 @@ class GPT2: GlobalCounters.reset() if timing: print("") st = GlobalCounters.time_sum_s - with Timing("ran model in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on GPU" if DEBUG>=2 else "")+ + with Timing("ran model in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on {Device.DEFAULT}" if DEBUG>=2 else "")+ f", {GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.global_mem*1e-9:.2f} GB"+ (f", {GlobalCounters.global_mem*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s" if DEBUG>=2 else "")) if DEBUG else None, enabled=timing): with WallTimeEvent(BenchEvent.STEP): diff --git a/examples/llama.py b/examples/llama.py index 42f9b6e57b..6739ca4c56 100755 --- a/examples/llama.py +++ b/examples/llama.py @@ -478,7 +478,7 @@ After you are done speaking, output [EOS]. You are not Chad. with Profiling(enabled=args.profile): with Timing("total ", enabled=args.timing, on_exit=lambda x: f", {1e9/x:.2f} tok/s, {GlobalCounters.global_mem/x:.2f} GB/s, param {param_bytes/x:.2f} GB/s"): with WallTimeEvent(BenchEvent.STEP): - with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on GPU" if DEBUG>=2 else "")+ + with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on {Device.DEFAULT}" if DEBUG>=2 else "")+ f", {GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.global_mem*1e-9:.2f} GB"+ (f", {GlobalCounters.global_mem*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s, param {param_bytes*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s" if DEBUG>=2 else "")) if DEBUG else None, enabled=args.timing): tok_tensor = llama.model(next_tok, start_pos, args.temperature) diff --git a/examples/llama3.py b/examples/llama3.py index 9664f491f2..d7c7f2c921 100644 --- a/examples/llama3.py +++ b/examples/llama3.py @@ -441,7 +441,7 @@ if __name__ == "__main__": with Profiling(enabled=args.profile): with Timing("total ", on_exit=lambda x: f", {1e9/x:.2f} tok/s, {GlobalCounters.global_mem/x:.2f} GB/s, param {param_bytes/x:.2f} GB/s"): with WallTimeEvent(BenchEvent.STEP): - with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on GPU" if DEBUG>=2 else "")+ + with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on {Device.DEFAULT}" if DEBUG>=2 else "")+ f", {GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.global_mem*1e-9:.2f} GB"+ (f", {GlobalCounters.global_mem*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s, param {param_bytes*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s" if DEBUG>=2 else "")) if DEBUG else None): tok = model(Tensor([[last_tok]], device=device), start_pos, TEMPERATURE, TOP_K, TOP_P, ALPHA_F, ALPHA_P) @@ -479,7 +479,7 @@ if __name__ == "__main__": st = GlobalCounters.time_sum_s with Profiling(enabled=args.profile): with Timing("total ", enabled=args.timing, on_exit=lambda x: f", {1e9/x:.2f} tok/s, {GlobalCounters.global_mem/x:.2f} GB/s, param {param_bytes/x:.2f} GB/s"): - with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on GPU" if DEBUG>=2 else "")+ + with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on {Device.DEFAULT}" if DEBUG>=2 else "")+ f", {GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.global_mem*1e-9:.2f} GB"+ (f", {GlobalCounters.global_mem*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s, param {param_bytes*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s" if DEBUG>=2 else "")) if DEBUG else None, enabled=args.timing): diff --git a/examples/mamba.py b/examples/mamba.py index d6093eabf5..d309807484 100644 --- a/examples/mamba.py +++ b/examples/mamba.py @@ -279,9 +279,15 @@ def generate(model, tokenizer, prompt: str, n_tokens_to_gen: int = 10, temp: boo # Loading in the prompt tokens logits = model.forward(Tensor([tks]))[:, -1, :] for _ in tqdm(range(n_tokens_to_gen), desc="Speed Gen"): - # TODO: topk if sample: - tok_Tens = (logits/temp).softmax().multinomial() + scaled_logits = logits / temp + if top_k is not None: + topk_values, topk_indices = scaled_logits.topk(top_k) + filtered_logits = Tensor.full_like(scaled_logits, -float("inf")) + filtered_logits = filtered_logits.scatter(dim=-1, index=topk_indices, src=topk_values) + tok_Tens = filtered_logits.softmax().multinomial() + else: + tok_Tens = scaled_logits.softmax().multinomial() else: tok_Tens = logits.argmax(axis=-1).unsqueeze(0) tok = tok_Tens.item() @@ -298,6 +304,7 @@ if __name__ == "__main__": parser.add_argument("--size", type=str, default="370m", help=f"Size of model to use [{', '.join([k for k in MODELS.keys()])}]") parser.add_argument("--n_tokens", type=int, default=10, help="Number of tokens to generate") + parser.add_argument("--top_k", type=int, help="Limit sampling to the top k most likely tokens") parser.add_argument("--sample", dest="sample", action="store_true", help="Sample flag") parser.add_argument("--temp", type=float, default=1.0, help="Sampling temp has to be <=1.0") args = parser.parse_args() @@ -308,8 +315,9 @@ if __name__ == "__main__": num_toks = args.n_tokens sample = args.sample temp = args.temp + top_k = args.top_k s = time.time() - tinyoutput = generate(model, tokenizer, prompt, n_tokens_to_gen=num_toks, sample=sample, temp=temp) + tinyoutput = generate(model, tokenizer, prompt, n_tokens_to_gen=num_toks, sample=sample, temp=temp, top_k=top_k) print(tinyoutput) print('TIME: ', time.time() - s) TORCHOUTPUT = "Why is gravity \nso important?\nBecause it's the only" diff --git a/examples/mlperf/scripts/stable_diffusion_downloads.sh b/examples/mlperf/scripts/stable_diffusion_downloads.sh new file mode 100755 index 0000000000..5f6798c176 --- /dev/null +++ b/examples/mlperf/scripts/stable_diffusion_downloads.sh @@ -0,0 +1,57 @@ +#!/usr/bin/env bash +# adapted from https://github.com/mlcommons/training/blob/4bdf5c8ed218ad76565a2ba1ac27c919ccc6d689/stable_diffusion/README.md + +# setup dirs + +DATA=/raid/datasets/stable_diffusion + +LAION=$DATA/laion-400m/webdataset-moments-filtered +COCO=$DATA/coco2014 +mkdir -p $LAION $COCO + +CKPT=/raid/weights/stable_diffusion +mkdir -p $CKPT/clip $CKPT/sd $CKPT/inception + +# download data + +# if rclone isn't installed system-wide / in your PATH, put the executable path in quotes below +#RCLONE="" +RCLONE="rclone" + +## VAE-encoded image latents, from 6.1M image subset of laion-400m +## about 1 TB for whole download +$RCLONE config create mlc-training s3 provider=Cloudflare access_key_id=76ea42eadb867e854061a1806220ee1e secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com +$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/laion-400m/moments-webdataset-filtered/ ${LAION} --include="*.tar" -P +$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/laion-400m/moments-webdataset-filtered/sha512sums.txt ${LAION} -P +cd $LAION && grep -E '\.tar$' sha512sums.txt | sha512sum -c --quiet - && \ + echo "All .tar files verified" || { echo "Checksum failure when validating downloaded Laion moments"; exit 1; } + +## prompts and FID statistics from 30k image subset of coco2014 +## 33 MB +$RCLONE config create mlc-training s3 provider=Cloudflare access_key_id=76ea42eadb867e854061a1806220ee1e secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com +$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/coco2014/val2014_30k.tsv ${COCO} -P + +$RCLONE config create mlc-training s3 provider=Cloudflare access_key_id=76ea42eadb867e854061a1806220ee1e secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com +$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/coco2014/val2014_30k_stats.npz ${COCO} -P + +# download checkpoints + +## clip (needed for text and vision encoders for validation) +CLIP_WEIGHTS_URL="https://huggingface.co/laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin" +CLIP_WEIGHTS_SHA256="9a78ef8e8c73fd0df621682e7a8e8eb36c6916cb3c16b291a082ecd52ab79cc4" +CLIP_CONFIG_URL="https://huggingface.co/laion/CLIP-ViT-H-14-laion2B-s32B-b79K/raw/main/open_clip_config.json" +wget -N -P ${CKPT}/clip ${CLIP_WEIGHTS_URL} +wget -N -P ${CKPT}/clip ${CLIP_CONFIG_URL} +echo "${CLIP_WEIGHTS_SHA256} ${CKPT}/clip/open_clip_pytorch_model.bin" | sha256sum -c + +## sd (needed for latent->image decoder for validation, also has clip text encoder for training) +SD_WEIGHTS_URL='https://huggingface.co/stabilityai/stable-diffusion-2-base/resolve/main/512-base-ema.ckpt' +SD_WEIGHTS_SHA256="d635794c1fedfdfa261e065370bea59c651fc9bfa65dc6d67ad29e11869a1824" +wget -N -P ${CKPT}/sd ${SD_WEIGHTS_URL} +echo "${SD_WEIGHTS_SHA256} ${CKPT}/sd/512-base-ema.ckpt" | sha256sum -c + +## inception (needed for validation) +FID_WEIGHTS_URL='https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth' +FID_WEIGHTS_SHA1="bd836944fd6db519dfd8d924aa457f5b3c8357ff" +wget -N -P ${CKPT}/inception ${FID_WEIGHTS_URL} +echo "${FID_WEIGHTS_SHA1} ${CKPT}/inception/pt_inception-2015-12-05-6726825d.pth" | sha1sum -c \ No newline at end of file diff --git a/examples/openpilot/compile4.py b/examples/openpilot/compile4.py index c57bd3eb70..55fcccbfbf 100644 --- a/examples/openpilot/compile4.py +++ b/examples/openpilot/compile4.py @@ -6,7 +6,7 @@ from tinygrad.schedule.kernelize import get_kernelize_map from tinygrad.engine.schedule import create_schedule_with_vars from tinygrad.engine.realize import run_schedule -# NOLOCALS=1 GPU=1 IMAGE=2 FLOAT16=1 VIZ=1 DEBUG=2 python3 examples/openpilot/compile4.py +# NOLOCALS=1 CL=1 IMAGE=2 FLOAT16=1 VIZ=1 DEBUG=2 python3 examples/openpilot/compile4.py OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx" OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "/tmp/openpilot.pkl" diff --git a/examples/qwq.py b/examples/qwq.py index fad87695bd..b3b03065dd 100644 --- a/examples/qwq.py +++ b/examples/qwq.py @@ -8,7 +8,7 @@ from typing import Dict, Union from extra.models.llama import Transformer, convert_from_huggingface, fix_bf16 from examples.llama3 import load -from tinygrad import nn, Tensor +from tinygrad import nn, Tensor, Device from tinygrad.helpers import fetch, colored, GlobalCounters, Timing, DEBUG from tinygrad.nn.state import load_state_dict, get_parameters @@ -80,7 +80,7 @@ if __name__ == "__main__": st = GlobalCounters.time_sum_s next_tok = Tensor([toks[start_pos:]]) if tok_tensor is None or (len(toks)-start_pos) > 1 else tok_tensor.reshape(1, 1) with Timing("total ", enabled=args.timing, on_exit=lambda x: f", {1e9/x:.2f} tok/s, {GlobalCounters.global_mem/x:.2f} GB/s, param {param_bytes/x:.2f} GB/s"): - with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on GPU" if DEBUG>=2 else "") + + with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on {Device.DEFAULT}" if DEBUG>=2 else "") + f", {GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.global_mem*1e-9:.2f} GB" + (f", {GlobalCounters.global_mem*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s, param {param_bytes*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s" if DEBUG>=2 else "")) if DEBUG else None, enabled=args.timing): tok_tensor = transformer(next_tok, start_pos, args.temperature) diff --git a/examples/sdxl.py b/examples/sdxl.py index 23f2236723..92d85ce057 100644 --- a/examples/sdxl.py +++ b/examples/sdxl.py @@ -437,8 +437,8 @@ if __name__ == "__main__": im.show() # validation! - if args.prompt == default_prompt and args.steps == 10 and args.seed == 0 and args.guidance == 6.0 and args.width == args.height == 1024 \ - and not args.weights: + is_default = args.prompt == default_prompt and args.steps == 10 and args.seed == 0 and args.guidance == 6.0 and args.width == args.height == 1024 + if is_default and not args.weights and not args.fakeweights: ref_image = Tensor(np.array(Image.open(Path(__file__).parent / "sdxl_seed0.png"))) distance = (((x.cast(dtypes.float) - ref_image.cast(dtypes.float)) / ref_image.max())**2).mean().item() assert distance < 4e-3, colored(f"validation failed with {distance=}", "red") diff --git a/examples/whisper.py b/examples/whisper.py index 5cce861a04..2df3122628 100644 --- a/examples/whisper.py +++ b/examples/whisper.py @@ -109,7 +109,7 @@ class TextDecoder: def forward(self, x:Tensor, pos:Union[Variable, Literal[0]], encoded_audio:Tensor): seqlen = x.shape[-1] - x = self.token_embedding(x) + self.positional_embedding.shrink(((pos, pos+seqlen), None, None)) + x = self.token_embedding(x) + self.positional_embedding.shrink(((pos, pos+seqlen), None)) for block in self.blocks: x = block(x, xa=encoded_audio, mask=self.mask, len=pos) return self.output_tok(x) diff --git a/extra/archprobe.py b/extra/archprobe.py index 73eb5037fb..7ba20b2a88 100644 --- a/extra/archprobe.py +++ b/extra/archprobe.py @@ -1,7 +1,7 @@ # copying the kernels from https://github.com/microsoft/ArchProbe into Python import numpy as np import pickle -from tinygrad.runtime.ops_gpu import CLProgram, CLBuffer +from tinygrad.runtime.ops_cl import CLProgram, CLBuffer from tinygrad import dtypes from tqdm import trange, tqdm from matplotlib import pyplot as plt diff --git a/extra/assembly/assembly_rdna.py b/extra/assembly/assembly_rdna.py index 0f5ab01ecf..297639d676 100644 --- a/extra/assembly/assembly_rdna.py +++ b/extra/assembly/assembly_rdna.py @@ -4,7 +4,7 @@ from tinygrad import dtypes from tinygrad.codegen.assembly import AssemblyCodegen, Register from tinygrad.codegen.opt.kernel import Ops from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps -from tinygrad.runtime.ops_gpu import ROCM_LLVM_PATH +from tinygrad.runtime.ops_cl import ROCM_LLVM_PATH # ugh, is this really needed? from extra.helpers import enable_early_exec diff --git a/extra/assembly/rocm/rdna3/asm.py b/extra/assembly/rocm/rdna3/asm.py index 2f6ad13264..9c65fa7360 100644 --- a/extra/assembly/rocm/rdna3/asm.py +++ b/extra/assembly/rocm/rdna3/asm.py @@ -5,7 +5,7 @@ from tinygrad.helpers import colored from extra.helpers import enable_early_exec early_exec = enable_early_exec() -from tinygrad.runtime.ops_gpu import CLProgram, CLBuffer, ROCM_LLVM_PATH +from tinygrad.runtime.ops_cl import CLProgram, CLBuffer, ROCM_LLVM_PATH ENABLE_NON_ASM = False diff --git a/extra/export_model.py b/extra/export_model.py index 65a8b3af9f..e29f8a8d31 100644 --- a/extra/export_model.py +++ b/extra/export_model.py @@ -10,7 +10,7 @@ from tinygrad.uop.ops import Ops import json from collections import OrderedDict -EXPORT_SUPPORTED_DEVICE = ["WEBGPU", "CPU", "CUDA", "GPU"] +EXPORT_SUPPORTED_DEVICE = ["WEBGPU", "CPU", "CUDA", "CL"] def compile_net(run:TinyJit, special_names:Dict[int,str]) -> Tuple[Dict[str,str],List[Tuple[str,List[str],List[int]]],Dict[str,Tuple[int,DType,int]],Dict[str,Tensor]]: functions, bufs, bufs_to_save, statements, bufnum = {}, {}, {}, [], 0 diff --git a/extra/gemm/intel_xmx.py b/extra/gemm/intel_xmx.py index 8ec478e5f6..719830473a 100644 --- a/extra/gemm/intel_xmx.py +++ b/extra/gemm/intel_xmx.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 import numpy as np -from tinygrad.runtime.ops_gpu import CLProgram, CLCompiler +from tinygrad.runtime.ops_cl import CLProgram, CLCompiler from tinygrad import Device, dtypes from tinygrad.device import Buffer from hexdump import hexdump @@ -11,7 +11,7 @@ from hexdump import hexdump # https://registry.khronos.org/OpenCL/extensions/intel/cl_intel_subgroup_split_matrix_multiply_accumulate.html # https://hc34.hotchips.org/assets/program/conference/day1/GPU%20HPC/Intel_s%20Ponte%20Vecchio%20GPU%20-%20Architecture%20Systems%20and%20Software%20FINAL.pdf -device = Device["GPU"] +device = Device["CL"] # NOTE: only the subgroup type 8 ones work prog = CLProgram(device, "test", CLCompiler(device, "test").compile(f""" @@ -26,9 +26,9 @@ __kernel void test(__global float* data0, const __global int* data1, const __glo """)) #with open("/tmp/test.elf", "wb") as f: f.write(prog.lib) -a = Buffer("GPU", 8, dtypes.float32).allocate() -b = Buffer("GPU", 0x10, dtypes.float16).allocate() -c = Buffer("GPU", 8*0x10, dtypes.float16).allocate() +a = Buffer("CL", 8, dtypes.float32).allocate() +b = Buffer("CL", 0x10, dtypes.float16).allocate() +c = Buffer("CL", 8*0x10, dtypes.float16).allocate() row = np.array([1,2,3,4,5,6,7,8,1,2,3,4,5,6,7,8], np.float16) mat = np.random.random((8, 0x10)).astype(np.float16) diff --git a/extra/gemm/max_matmul.py b/extra/gemm/max_matmul.py index 6937b7d153..5041497839 100644 --- a/extra/gemm/max_matmul.py +++ b/extra/gemm/max_matmul.py @@ -75,7 +75,7 @@ if __name__ == "__main__": if GEMM_VARIATION == "max" and (M%64)==0 and (N%128)==0 and (K%64)==0 and DTYPE_IN == dtypes.half and DTYPE_OUT == dtypes.float and DTYPE_ACC == dtypes.float: print("Using CUDA and triton-generated kernel") - # See nv_triton_gemm.annotated.ptx for PTX code which was generated from `PYTHONPATH=. DEBUG=6 CUDA=1 PTX=1 python3 extra/gemm/triton_nv_matmul.py` + # See nv_triton_gemm.annotated.ptx for PTX code which was generated from `PYTHONPATH=. DEBUG=6 CUDA=1 CUDA_PTX=1 python3 extra/gemm/triton_nv_matmul.py` # this kernel with M=N=K=4096 does 162TFLOPS, vs torch at 144TFLOPS and BEAM=8 tinygrad at 138TFLOPS. theo max is 165TFLOPS. # WMMA element size is (M, N, K) = (16, 8, 16) diff --git a/extra/gemm/triton_nv_matmul.py b/extra/gemm/triton_nv_matmul.py index 5f04a34076..89e7838bb0 100644 --- a/extra/gemm/triton_nv_matmul.py +++ b/extra/gemm/triton_nv_matmul.py @@ -43,7 +43,7 @@ def matmul_kernel(c_ptr, a_ptr, b_ptr, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] tl.store(c_ptrs, c) -# CUDA=1 PTX=1 python3 extra/gemm/triton_nv_matmul.py +# CUDA=1 CUDA_PTX=1 python3 extra/gemm/triton_nv_matmul.py if __name__ == "__main__": BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K = 64, 128, 64 M, N, K = 4096, 4096, 4096 diff --git a/extra/models/inception.py b/extra/models/inception.py index cf77a63ed9..15d0b58bc9 100644 --- a/extra/models/inception.py +++ b/extra/models/inception.py @@ -270,8 +270,10 @@ class FidInceptionV3: self.Mixed_7b = inception.Mixed_7b self.Mixed_7c = inception.Mixed_7c - def load_from_pretrained(self): - state_dict = torch_load(str(fetch("https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth", "pt_inception-2015-12-05-6726825d.pth"))) + def load_from_pretrained(self, path=None): + if path is None: + path = fetch("https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth", "pt_inception-2015-12-05-6726825d.pth") + state_dict = torch_load(str(path)) for k,v in state_dict.items(): if k.endswith(".num_batches_tracked"): state_dict[k] = v.reshape(1) diff --git a/extra/nv_gpu_driver/nv_ioctl.py b/extra/nv_gpu_driver/nv_ioctl.py index 30a6c183f9..44a2a11f3e 100644 --- a/extra/nv_gpu_driver/nv_ioctl.py +++ b/extra/nv_gpu_driver/nv_ioctl.py @@ -272,4 +272,4 @@ def compare_launch_state(states, good_states): return True, "PASS" -# IOCTL=1 PTX=1 CUDA=1 python3 test/test_ops.py TestOps.test_tiny_add \ No newline at end of file +# IOCTL=1 CUDA=1 CUDA_PTX=1 python3 test/test_ops.py TestOps.test_tiny_add \ No newline at end of file diff --git a/extra/optimization/generate_dataset.sh b/extra/optimization/generate_dataset.sh index 6f70916979..b843dac700 100755 --- a/extra/optimization/generate_dataset.sh +++ b/extra/optimization/generate_dataset.sh @@ -7,7 +7,7 @@ rm $LOGOPS test/external/process_replay/reset.py CI=1 python3 -m pytest -n=auto test/test_ops.py test/test_nn.py test/test_winograd.py test/models/test_real_world.py --durations=20 -GPU=1 python3 -m pytest test/test_tiny.py +CL=1 python3 -m pytest test/test_tiny.py # extract, sort and uniq extra/optimization/extract_dataset.py diff --git a/extra/optimization/test_beam_search.py b/extra/optimization/test_beam_search.py index 7042ea914e..f493ec48eb 100644 --- a/extra/optimization/test_beam_search.py +++ b/extra/optimization/test_beam_search.py @@ -50,7 +50,7 @@ class TestBeamSearch(unittest.TestCase): def test_variable_shrink_prime_number(self): v = Variable("v", 1, 400).bind(367) a = rand(400, 367) - b = (a.shrink(((0,v), None))+1).reshape(367,367).realize() + b = (a.shrink(((0,v), None))+1)[:367,:367].realize() np.testing.assert_allclose(b.numpy(), a.numpy()[:367]+1, atol=1e-4, rtol=1e-4) def test_no_mutate_rawbuffers(self): diff --git a/extra/qcom_gpu_driver/qcom_opencl_interop.py b/extra/qcom_gpu_driver/qcom_opencl_interop.py index d595ba343f..c2e0741ca2 100644 --- a/extra/qcom_gpu_driver/qcom_opencl_interop.py +++ b/extra/qcom_gpu_driver/qcom_opencl_interop.py @@ -1,6 +1,6 @@ import ctypes, array from hexdump import hexdump -from tinygrad.runtime.ops_gpu import GPUDevice +from tinygrad.runtime.ops_cl import CLDevice from tinygrad.helpers import getenv, to_mv, mv_address from tinygrad.dtype import dtypes from tinygrad import Tensor, TinyJit @@ -8,7 +8,7 @@ from tinygrad.runtime.autogen import opencl as cl if getenv("IOCTL"): import extra.qcom_gpu_driver.opencl_ioctl # noqa: F401 # pylint: disable=unused-import # create raw opencl buffer. -gdev = GPUDevice() +gdev = CLDevice() cl_buf = cl.clCreateBuffer(gdev.context, cl.CL_MEM_READ_WRITE, 0x100, None, status := ctypes.c_int32()) assert status.value == 0 diff --git a/extra/sqtt/README.md b/extra/sqtt/README.md index 6d739ceb68..1d19ae8f32 100644 --- a/extra/sqtt/README.md +++ b/extra/sqtt/README.md @@ -4,7 +4,7 @@ Only supported on 7900XTX, requires either AM (`rmmod amdgpu`) or disabling power gating on AMD (`ppfeaturemask=0xffff3fff`, don't forget to rebuild initramfs) -SQTT is implemented on top of normal tinygrad PROFILE=1, `PROFILE=1 SQTT=1` to get profile pickle with sqtt data embedded in it. +SQTT is implemented on top of normal tinygrad profiling, `VIZ=1 SQTT=1` to get profile pickle with sqtt data embedded in it. `SQTT_BUFFER_SIZE=X` to change size of SQTT buffer (per shader engine, 6 SEs on 7900xtx) in megabytes, default 256. diff --git a/extra/thneed.py b/extra/thneed.py index c59f636858..ca89bfa603 100644 --- a/extra/thneed.py +++ b/extra/thneed.py @@ -4,13 +4,13 @@ import struct import json import traceback import numpy as np -from tinygrad.runtime.ops_gpu import CLProgram, compile_gpu +from tinygrad.runtime.ops_cl import CLProgram, compile_gpu from tinygrad.device import Device from tinygrad.helpers import DEBUG, getenv from collections import defaultdict import pyopencl as cl -from tinygrad.runtime.ops_gpu import OSX_TIMING_RATIO -CL = Device["GPU"] +from tinygrad.runtime.ops_cl import OSX_TIMING_RATIO +CL = Device["CL"] DEBUGCL = getenv("DEBUGCL", 0) FLOAT16 = getenv("FLOAT16", 0) @@ -110,7 +110,7 @@ class Thneed: prgs = {} for o in jdat['binaries']: nptr = ptr + o['length'] - prgs[o['name']] = CLProgram(Device["GPU"], o['name'], weights[ptr:nptr]) + prgs[o['name']] = CLProgram(Device["CL"], o['name'], weights[ptr:nptr]) ptr = nptr # populate the cl_cache @@ -267,7 +267,7 @@ class Thneed: for prg, args in self.cl_cache: events.append(prg.clprg(CL.queue, *args)) mt = time.monotonic() - Device["GPU"].synchronize() + Device["CL"].synchronize() et = time.monotonic() - st print(f"submit in {(mt-st)*1000.0:.2f} ms, total runtime is {et*1000.0:.2f} ms") diff --git a/extra/to_movement_ops.py b/extra/to_movement_ops.py index ed978ab0a3..3170cd8c61 100644 --- a/extra/to_movement_ops.py +++ b/extra/to_movement_ops.py @@ -35,7 +35,7 @@ def to_movement_ops(st: ShapeTracker) -> List[Tuple[MovementOps, Tuple]]: to_apply:List[Tuple[MovementOps, Tuple]] = [] for i, v in enumerate(st.views): real_shape = tuple(y-x for x,y in v.mask) if v.mask else v.shape - offset = v.offset + sum(st*(s-1) for s,st in zip(real_shape, v.strides) if st<0) + offset = (v.offset or 0) + sum(st*(s-1) for s,st in zip(real_shape, v.strides) if st<0) real_offset = offset + (sum(x*st for (x,_),st in zip(v.mask, v.strides)) if v.mask else 0) real_real_shape = [s for s,st in zip(real_shape, v.strides) if st] strides: List[int] = [abs(st) if isinstance(st,int) else st for st in v.strides if st] diff --git a/extra/torch_backend/backend.py b/extra/torch_backend/backend.py index c312badf1f..d5993f64b0 100644 --- a/extra/torch_backend/backend.py +++ b/extra/torch_backend/backend.py @@ -177,22 +177,28 @@ def cached_to_movement_ops(shape, st) -> list: from tinygrad.shape.shapetracker import ShapeTracker, View from extra.to_movement_ops import to_movement_ops, apply_mop, MovementOps + +@wrap_view_op +def _as_strided(tensor:Tensor, size, stride, storage_offset=None): + # multiple as_strided do not compound + base = canonical_base(tensor) + # TODO: this is heavyweight + st = ShapeTracker(base.uop.st.views + (View.create(tuple(size), tuple(stride), storage_offset),)) + ret = base + if TORCH_DEBUG >= 1: print("**** as_strided", tensor.shape, size, stride, st) + if prod(size) == 1: return ret.flatten()[storage_offset].reshape(size) + for mo in cached_to_movement_ops(tuple(base.shape), st): ret = apply_mop(ret, mo) + return ret + @torch.library.impl("aten::as_strided", "privateuseone") def as_strided(tensor:torch.Tensor, size, stride, storage_offset=None): storage_offset = storage_offset or tensor.storage_offset() - @wrap_view_op - def _as_strided(tensor:Tensor, size, stride, storage_offset=None): - # multiple as_strided do not compound - base = canonical_base(tensor) - # TODO: this is heavyweight - st = ShapeTracker(base.uop.st.views + (View.create(tuple(size), tuple(stride), storage_offset),)) - ret = base - if TORCH_DEBUG >= 1: print("**** as_strided", tensor.shape, size, stride, st) - if prod(size) == 1: return ret.flatten()[storage_offset].reshape(size) - for mo in cached_to_movement_ops(tuple(base.shape), st): ret = apply_mop(ret, mo) - return ret return _as_strided(tensor, size, stride, storage_offset) +@torch.library.impl("aten::_reshape_alias", "privateuseone") +def _reshape_alias(tensor:torch.Tensor, size, stride): + return _as_strided(tensor, size, stride) + @torch.library.impl("aten::empty_strided", "privateuseone") def empty_strided(size, stride, dtype, layout=None, device=None, pin_memory=False): if TORCH_DEBUG: print(f"empty_strided {size=} {stride=} {dtype=} {layout=} {device=} {pin_memory=}") diff --git a/pytest.ini b/pytest.ini index cccc62e404..b9c3f6064a 100644 --- a/pytest.ini +++ b/pytest.ini @@ -1,2 +1,6 @@ [pytest] norecursedirs = extra +timeout = 180 +timeout_method = thread +timeout_func_only = true +testpaths = test diff --git a/setup.py b/setup.py index 9ebfda1ba9..f90a52b584 100644 --- a/setup.py +++ b/setup.py @@ -9,12 +9,12 @@ with open(directory / 'README.md', encoding='utf-8') as f: testing_minimal = [ "numpy", - "torch==2.7.1", + "torch==2.8.0", "pytest", "pytest-xdist", + "pytest-timeout", "hypothesis", "z3-solver", - "ml_dtypes" ] setup(name='tinygrad', @@ -59,7 +59,7 @@ setup(name='tinygrad', 'triton': ["triton-nightly>=2.1.0.dev20231014192330"], 'linting': [ "pylint", - "mypy==1.13.0", + "mypy==1.18.1", "typing-extensions", "pre-commit", "ruff", diff --git a/test/device/test_ocl.py b/test/device/test_ocl.py index 04b8e2523e..6f58b909db 100644 --- a/test/device/test_ocl.py +++ b/test/device/test_ocl.py @@ -3,9 +3,9 @@ from tinygrad import Device from tinygrad.device import Buffer from tinygrad.dtype import dtypes from tinygrad.helpers import CI -from tinygrad.runtime.ops_gpu import CLDevice, CLAllocator, CLCompiler, CLProgram +from tinygrad.runtime.ops_cl import CLDevice, CLAllocator, CLCompiler, CLProgram -@unittest.skipUnless(Device.DEFAULT == "GPU", "Runs only on OpenCL (GPU)") +@unittest.skipUnless(Device.DEFAULT == "CL", "Runs only on OpenCL") class TestCLError(unittest.TestCase): @unittest.skipIf(CI, "dangerous for CI, it allocates tons of memory") def test_oom(self): @@ -24,7 +24,7 @@ class TestCLError(unittest.TestCase): def test_unaligned_copy(self): data = list(range(65)) unaligned = memoryview(bytearray(data))[1:] - buffer = Buffer("GPU", 64, dtypes.uint8).allocate() + buffer = Buffer("CL", 64, dtypes.uint8).allocate() buffer.copyin(unaligned) result = memoryview(bytearray(len(data) - 1)) buffer.copyout(result) diff --git a/test/unit/test_qcom.py b/test/device/test_qcom.py similarity index 97% rename from test/unit/test_qcom.py rename to test/device/test_qcom.py index 4794cdb177..827a364d1d 100644 --- a/test/unit/test_qcom.py +++ b/test/device/test_qcom.py @@ -10,10 +10,11 @@ class TestQcom(unittest.TestCase): def __validate(imgdt, expected_pitch): img = dev.allocator.alloc(imgdt.shape[0] * imgdt.shape[1] * 16, options:=BufferSpec(image=imgdt)) - pitch = (img.descriptor[2] & 0x1fffff80) >> 7 + pitch = img.texture_info.pitch assert pitch == expected_pitch, f"Failed pitch for image: {imgdt}. Got 0x{pitch:X}, expected 0x{expected_pitch:X}" dev.allocator.free(img, imgdt.shape[0] * imgdt.shape[1] * 16, options) + # Match opencl pitches for perf __validate(dtypes.imageh((1, 201)), 0x680) __validate(dtypes.imageh((16, 216)), 0x700) __validate(dtypes.imageh((16, 9)), 0x80) diff --git a/test/external/external_benchmark_hip_compile.py b/test/external/external_benchmark_hip_compile.py index 2b1d480348..d97047b923 100644 --- a/test/external/external_benchmark_hip_compile.py +++ b/test/external/external_benchmark_hip_compile.py @@ -1,7 +1,7 @@ import random, os from tinygrad.helpers import Timing from tinygrad.runtime.ops_hip import compile_hip, HIPDevice -from tinygrad.runtime.ops_gpu import compile_cl, CLDevice +from tinygrad.runtime.ops_cl import compile_cl, CLDevice # OMP_NUM_THREADS=1 strace -tt -f -e trace=file python3 test/external/external_benchmark_hip_compile.py # AMD_COMGR_REDIRECT_LOGS=stdout AMD_COMGR_EMIT_VERBOSE_LOGS=1 python3 test/external/external_benchmark_hip_compile.py diff --git a/test/external/external_cl_half_max.py b/test/external/external_cl_half_max.py index 7cd6b0c509..020d806c63 100644 --- a/test/external/external_cl_half_max.py +++ b/test/external/external_cl_half_max.py @@ -1,4 +1,4 @@ -from tinygrad.runtime.ops_gpu import CLDevice, CLProgram, compile_cl +from tinygrad.runtime.ops_cl import CLDevice, CLProgram, compile_cl if __name__ == "__main__": dev = CLDevice() diff --git a/test/external/external_gpu_fail_osx.py b/test/external/external_gpu_fail_osx.py index b11b695e3d..51a458b136 100644 --- a/test/external/external_gpu_fail_osx.py +++ b/test/external/external_gpu_fail_osx.py @@ -1,5 +1,5 @@ # ugh, OS X OpenCL doesn't support half -from tinygrad.runtime.ops_gpu import CLDevice, CLProgram, CLCompiler +from tinygrad.runtime.ops_cl import CLDevice, CLProgram, CLCompiler src = """#pragma OPENCL EXTENSION cl_khr_fp16 : enable __kernel void max_half(__global half* data0, const __global half* data1) { diff --git a/test/external/external_multi_gpu.py b/test/external/external_multi_gpu.py index 32d107df7d..b3c8fefb30 100644 --- a/test/external/external_multi_gpu.py +++ b/test/external/external_multi_gpu.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 # cd extra/disassemblers/ && git clone --recursive github.com:geohot/cuda_ioctl_sniffer.git -# LD_PRELOAD=$PWD/extra/disassemblers/cuda_ioctl_sniffer/out/sniff.so GPU=1 python3 test/external/external_multi_gpu.py +# LD_PRELOAD=$PWD/extra/disassemblers/cuda_ioctl_sniffer/out/sniff.so CL=1 python3 test/external/external_multi_gpu.py import numpy as np from tinygrad.tensor import Tensor from tinygrad.helpers import colored, Timing, getenv diff --git a/test/external/external_osx_profiling.py b/test/external/external_osx_profiling.py index 8ac1e584df..277ca304ad 100644 --- a/test/external/external_osx_profiling.py +++ b/test/external/external_osx_profiling.py @@ -1,4 +1,4 @@ -from tinygrad.runtime.ops_gpu import CLProgram, CL, CLBuffer +from tinygrad.runtime.ops_cl import CLProgram, CL, CLBuffer from tinygrad import dtypes import time diff --git a/test/external/external_test_hcq_fuzz_failures.py b/test/external/external_test_hcq_fuzz_failures.py index c8b6198e9a..4a381d5336 100644 --- a/test/external/external_test_hcq_fuzz_failures.py +++ b/test/external/external_test_hcq_fuzz_failures.py @@ -55,7 +55,7 @@ class TestHCQFuzzFailures(unittest.TestCase): ast = UOp(Ops.SINK, dtypes.void, arg=None, src=( UOp(Ops.STORE, dtypes.void, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=0, src=()), UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 1), offset=0, mask=None, contiguous=True),)), src=()), UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 2, 4)), arg=1, src=()), x39:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 1), offset=0, mask=((0, 1), (0, 6)), contiguous=False),)), src=()),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=2, src=()), x39,)),)),)), UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 2, 4)), arg=3, src=()), x46:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 1), offset=-6, mask=((0, 1), (6, 12)), contiguous=False),)), src=()),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=4, src=()), x46,)),)),)),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 1, 4)), arg=5, src=()), x54:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 0), offset=0, mask=((0, 1), (12, 13)), contiguous=False),)), src=()),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=6, src=()), x54,)),)),)),)),)), UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=7, src=()), UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 1), offset=-13, mask=((0, 1), (13, 17)), contiguous=False),)), src=()),)),)), UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=8, src=()), UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 1), offset=-17, mask=((0, 1), (17, 21)), contiguous=False),)), src=()),)),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 1, 4)), arg=9, src=()), x68:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 0), offset=0, mask=((0, 1), (21, 22)), contiguous=False),)), src=()),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=10, src=()), x68,)),)),)),)),)), UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=11, src=()), UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 1), offset=-22, mask=((0, 1), (22, 26)), contiguous=False),)), src=()),)),)), UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=12, src=()), UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 1), offset=-26, mask=((0, 1), (26, 30)), contiguous=False),)), src=()),)),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 1, 4)), arg=13, src=()), x82:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 0), offset=0, mask=((0, 1), (30, 31)), contiguous=False),)), src=()),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=14, src=()), x82,)),)),)),)),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 1, 4)), arg=15, src=()), x90:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 0), offset=0, mask=((0, 1), (31, 32)), contiguous=False),)), src=()),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=16, src=()), x90,)),)),)),)),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 1, 4)), arg=17, src=()), x98:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 0), offset=0, mask=((0, 1), (32, 33)), contiguous=False),)), src=()),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=18, src=()), x98,)),)),)),)),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 1, 4)), arg=19, src=()), x106:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 0), offset=0, mask=((0, 1), (33, 34)), contiguous=False),)), src=()),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=20, src=()), x106,)),)),)),)),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 1, 4)), arg=21, src=()), x114:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 0), offset=0, mask=((0, 1), (34, 35)), contiguous=False),)), src=()),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=22, src=()), x114,)),)),)),)),)), UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=23, src=()), UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 1), offset=-35, mask=((0, 1), (35, 39)), contiguous=False),)), src=()),)),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 1, 4)), arg=24, src=()), x125:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 0), offset=0, mask=((0, 1), (39, 40)), contiguous=False),)), src=()),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=25, src=()), x125,)),)),)),)),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 1, 4)), arg=26, src=()), x133:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 0), offset=0, mask=((0, 1), (40, 41)), contiguous=False),)), src=()),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=27, src=()), x133,)),)),)),)),)), UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 2, 4)), arg=28, src=()), x140:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 1), offset=-41, mask=((0, 1), (41, 47)), contiguous=False),)), src=()),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=29, src=()), x140,)),)),)),)), UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 2, 4)), arg=30, src=()), x147:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 1), offset=-47, mask=((0, 1), (47, 53)), contiguous=False),)), src=()),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=31, src=()), x147,)),)),)),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 1, 4)), arg=32, src=()), x155:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 0), offset=0, mask=((0, 1), (53, 54)), contiguous=False),)), src=()),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=33, src=()), x155,)),)),)),)),)), UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=34, src=()), UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 1), offset=-54, mask=((0, 1), (54, 58)), contiguous=False),)), src=()),)),)), UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=35, src=()), UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 1), offset=-58, mask=((0, 1), (58, 62)), contiguous=False),)), src=()),)),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 1, 4)), arg=36, src=()), x169:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 0), offset=0, mask=((0, 1), (62, 63)), contiguous=False),)), src=()),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=37, src=()), x169,)),)),)),)),)), UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=38, src=()), UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 1), offset=-63, mask=((0, 1), (63, 67)), contiguous=False),)), src=()),)),)), UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=39, src=()), UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 1), offset=-67, mask=((0, 1), (67, 71)), contiguous=False),)), src=()),)),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 1, 4)), arg=40, src=()), x183:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 0), offset=0, mask=((0, 1), (71, 72)), contiguous=False),)), src=()),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=41, src=()), x183,)),)),)),)),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 1, 4)), arg=42, src=()), x191:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 0), offset=0, mask=((0, 1), (72, 73)), contiguous=False),)), src=()),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=43, src=()), x191,)),)),)),)),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 1, 4)), arg=44, src=()), x199:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 0), offset=0, mask=((0, 1), (73, 74)), contiguous=False),)), src=()),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=45, src=()), x199,)),)),)),)),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 1, 4)), arg=46, src=()), x207:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 0), offset=0, mask=((0, 1), (74, 75)), contiguous=False),)), src=()),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=47, src=()), x207,)),)),)),)),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 1, 4)), arg=48, src=()), x215:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 0), offset=0, mask=((0, 1), (75, 76)), contiguous=False),)), src=()),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=49, src=()), x215,)),)),)),)),)), UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=50, src=()), UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 1), offset=-76, mask=((0, 1), (76, 80)), contiguous=False),)), src=()),)),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 1, 4)), arg=51, src=()), x226:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 0), offset=0, mask=((0, 1), (80, 81)), contiguous=False),)), src=()),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=52, src=()), x226,)),)),)),)),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 1, 4)), arg=53, src=()), x234:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 0), offset=0, mask=((0, 1), (81, 82)), contiguous=False),)), src=()),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=54, src=()), x234,)),)),)),)),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 1, 4)), arg=55, src=()), x243:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 0), offset=0, mask=((0, 1), (82, 83)), contiguous=False),)), src=()),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=56, src=()), x243,)),)),)), UOp(Ops.ADD, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 1, 4)), arg=57, src=()), x250:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 0), offset=0, mask=((0, 1), (83, 84)), contiguous=False),)), src=()),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=58, src=()), x250,)),)),)),)),)),)), UOp(Ops.CAST, dtypes.float, arg=None, src=( UOp(Ops.LOAD, dtypes.float, arg=None, src=( UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 128, 4)), arg=59, src=()), UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 596), strides=(0, 1), offset=-84, mask=((0, 1), (84, 596)), contiguous=False),)), src=()),)),)),)),)),)) # noqa: E501 opts = [Opt(op=OptOps.UPCAST, axis=0, arg=4)] - helper_test_lin(Kernel(ast), opts, failed_platforms=[], validate_device=Device["GPU"]) + helper_test_lin(Kernel(ast), opts, failed_platforms=[], validate_device=Device["CL"]) if __name__ == '__main__': unittest.main() diff --git a/test/external/external_test_image.py b/test/external/external_test_image.py index 1c2cc397f2..4f2e06d903 100644 --- a/test/external/external_test_image.py +++ b/test/external/external_test_image.py @@ -4,7 +4,7 @@ import unittest import numpy as np if 'IMAGE' not in os.environ: os.environ['IMAGE'] = '2' -os.environ['GPU'] = '1' +os.environ['CL'] = '1' os.environ['OPT'] = '2' from tinygrad.tensor import Tensor from tinygrad.nn import Conv2d diff --git a/test/external/external_test_onnx_backend.py b/test/external/external_test_onnx_backend.py index d48154b85d..112ccd797c 100644 --- a/test/external/external_test_onnx_backend.py +++ b/test/external/external_test_onnx_backend.py @@ -193,12 +193,12 @@ backend_test.exclude('test_adam_cpu') backend_test.exclude('test_gradient_of_add_and_mul_cpu') backend_test.exclude('test_gradient_of_add_cpu') -if Device.DEFAULT in ['GPU', 'METAL']: +if Device.DEFAULT in ['CL', 'METAL']: backend_test.exclude('test_resize_upsample_sizes_nearest_axes_2_3_cpu') backend_test.exclude('test_resize_upsample_sizes_nearest_axes_3_2_cpu') backend_test.exclude('test_resize_upsample_sizes_nearest_cpu') -if Device.DEFAULT == "METAL" or (OSX and Device.DEFAULT == "GPU"): +if Device.DEFAULT == "METAL" or (OSX and Device.DEFAULT == "CL"): # numerical inaccuracy backend_test.exclude('test_mish_cpu') backend_test.exclude('test_mish_expanded_cpu') diff --git a/test/external/external_test_opt.py b/test/external/external_test_opt.py index 584c3af89f..0b0cb4bd2b 100644 --- a/test/external/external_test_opt.py +++ b/test/external/external_test_opt.py @@ -4,7 +4,7 @@ import numpy as np import torch from tinygrad import GlobalCounters, Tensor, Device -from tinygrad.helpers import getenv, Context +from tinygrad.helpers import getenv, Context, RANGEIFY from tinygrad.nn.state import get_parameters from tinygrad.engine.realize import capturing from tinygrad.tensor import _to_np_dtype @@ -34,7 +34,7 @@ from extra.models.efficientnet import EfficientNet from extra.models.resnet import ResNet18 from extra.models.vit import ViT -@unittest.skipUnless(Device.DEFAULT == "GPU", "Not Implemented") +@unittest.skipUnless(Device.DEFAULT == "CL", "Not Implemented") class TestInferenceMinKernels(unittest.TestCase): def setUp(self): self.training_old = Tensor.training @@ -90,7 +90,7 @@ class TestInferenceMinKernels(unittest.TestCase): with CLCache(100): model(inp, 0).realize() -@unittest.skipUnless(Device.DEFAULT == "GPU", "Not Implemented") +@unittest.skipUnless(Device.DEFAULT == "CL", "Not Implemented") class TestOptBinOp(unittest.TestCase): def _test_no_binop_rerun(self, f1, f2=None, allowed=1): a = Tensor.randn(16, 16) @@ -106,7 +106,7 @@ class TestOptBinOp(unittest.TestCase): def test_no_binop_rerun(self): return self._test_no_binop_rerun(lambda a,b: a*b, lambda a,b: (a*b).reshape(16, 16, 1)) def test_no_binop_rerun_alt(self): return self._test_no_binop_rerun(lambda a,b: (a*b).reshape(16, 16, 1), lambda a,b: a*b) def test_no_binop_rerun_reduce_broadcast(self): - return self._test_no_binop_rerun(lambda a,b: a.sum()+b, lambda a,b: a.sum().reshape(1,1)+b, allowed=2) + return self._test_no_binop_rerun(lambda a,b: a.sum()+b, lambda a,b: a.sum().reshape(1,1)+b, allowed=1 if RANGEIFY else 2) @unittest.skip("this test started failing with the new change, based movementop issue") def test_no_binop_rerun_transposed(self): return self._test_no_binop_rerun(lambda a,b: (a.T*b.T).T, lambda a,b: a*b) @@ -117,7 +117,7 @@ class TestOptBinOp(unittest.TestCase): #def test_no_binop_rerun_reduce(self): return self._test_no_binop_rerun(lambda a,b: (a*b).sum(), lambda a,b: (a*b).reshape(16, 16, 1).sum()) #def test_no_binop_rerun_reduce_alt(self): return self._test_no_binop_rerun(lambda a,b: a.sum(1)+b[0], lambda a,b: a.sum(1).reshape(1,16)+b[0]) -@unittest.skipUnless(Device.DEFAULT == "GPU", "Not Implemented") +@unittest.skipUnless(Device.DEFAULT == "CL", "Not Implemented") class TestOptReduceLoop(unittest.TestCase): def test_loop_left(self): a = Tensor.randn(16, 16) @@ -139,7 +139,7 @@ class TestOptReduceLoop(unittest.TestCase): c.realize() assert cache.count == 2, "loop right fusion broken" -@unittest.skipUnless(Device.DEFAULT == "GPU", "Not Implemented") +@unittest.skipUnless(Device.DEFAULT == "CL", "Not Implemented") class TestOptWChild(unittest.TestCase): @unittest.skip("this no longer happens, use realize") def test_unrealized_child(self): @@ -152,7 +152,7 @@ class TestOptWChild(unittest.TestCase): d.realize() assert cache.count == 2, "don't fuse if you have children" -@unittest.skipUnless(Device.DEFAULT == "GPU", "Not Implemented") +@unittest.skipUnless(Device.DEFAULT == "CL", "Not Implemented") class TestOpt(unittest.TestCase): def test_muladd(self): a,b,c = [Tensor.randn(2,2).realize() for _ in range(3)] @@ -164,7 +164,7 @@ class TestOpt(unittest.TestCase): def test_permute_was_pushed(self): a = Tensor.randn(16, 16, 16) - with CLCache(2): + with CLCache(1 if RANGEIFY else 2): c = a.sum(2) d = c.permute(1,0).contiguous() d.realize() @@ -172,7 +172,7 @@ class TestOpt(unittest.TestCase): def test_permute_was_pushed_through_contract_reshape(self): a = Tensor.randn(4, 4, 4, 4, 4) - with CLCache(2): + with CLCache(1 if RANGEIFY else 2): c = a.sum(-1) d = c.reshape(16,16).permute(1,0).contiguous() d.realize() @@ -180,7 +180,7 @@ class TestOpt(unittest.TestCase): def test_permute_was_pushed_through_contractw1s_reshape(self): a = Tensor.randn(4, 4, 4, 4, 4) - with CLCache(2): + with CLCache(1 if RANGEIFY else 2): c = a.sum(-1) d = c.reshape(16,1,16).permute(2,1,0).contiguous() d.realize() @@ -188,7 +188,7 @@ class TestOpt(unittest.TestCase): def test_permute_was_pushed_through_expand_reshape(self): a = Tensor.randn(16, 16, 16) - with CLCache(2): + with CLCache(1 if RANGEIFY else 2): c = a.sum(2) d = c.reshape(4,4,4,4).permute(2,3,0,1).contiguous() d.realize() diff --git a/test/external/external_test_speed_llama.py b/test/external/external_test_speed_llama.py index 3d468e6257..30dde43781 100644 --- a/test/external/external_test_speed_llama.py +++ b/test/external/external_test_speed_llama.py @@ -20,7 +20,7 @@ class TestLLaMASpeed(unittest.TestCase): def test_llama_compile(self): backup_program = Device[Device.DEFAULT].runtime backup_allocator = Device[Device.DEFAULT].allocator - backup_compiler = Device[Device.DEFAULT].compiler + backup_compiler = Device[Device.DEFAULT].compiler.compile_cached Device[Device.DEFAULT].runtime = FakeProgram Device[Device.DEFAULT].allocator = FakeAllocator(Device.default) @@ -44,14 +44,14 @@ class TestLLaMASpeed(unittest.TestCase): run_llama("codegen(1)") # test no compiler use for this - Device[Device.DEFAULT].compiler = None + Device[Device.DEFAULT].compiler.compile_cached = None run_llama("methodcache", False) with Profiling(sort='time', frac=0.1, fn="/tmp/llama.prof", ts=5): run_llama("profile", False) Device[Device.DEFAULT].runtime = backup_program Device[Device.DEFAULT].allocator = backup_allocator - Device[Device.DEFAULT].compiler = backup_compiler + Device[Device.DEFAULT].compiler.compile_cached = backup_compiler if __name__ == '__main__': TestLLaMASpeed().test_llama_compile() diff --git a/test/external/external_uop_gc.py b/test/external/external_uop_gc.py index 538cbc3bc6..c72bacd9aa 100644 --- a/test/external/external_uop_gc.py +++ b/test/external/external_uop_gc.py @@ -1,6 +1,6 @@ import gc from tinygrad import Tensor, UOp, Device -from tinygrad.shape.shapetracker import views_to_indexed_uops +from tinygrad.shape.shapetracker import views_to_valid_uop from tinygrad.engine.realize import method_cache, get_program def uops_allocated(): return sum([isinstance(x, UOp) for x in gc.get_objects()]) @@ -60,7 +60,7 @@ if __name__ == "__main__": # these caches will keep uops alive method_cache.clear() - views_to_indexed_uops.cache_clear() + views_to_valid_uop.cache_clear() new_uops = uops_allocated() gc.collect() diff --git a/test/external/fuzz_fast_idiv.py b/test/external/fuzz_fast_idiv.py index 02fda9ee9f..a6e48f1d8a 100644 --- a/test/external/fuzz_fast_idiv.py +++ b/test/external/fuzz_fast_idiv.py @@ -11,7 +11,7 @@ if __name__ == "__main__": for i in range(10_000): if i % 1000 == 0: print(f"Progress: {i}") - dt = random.choice(dtypes.ints) + dt = random.choice(dtypes.ints + tuple(dt.vec(4) for dt in dtypes.ints)) 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])) if d in powers_of_two: continue diff --git a/test/external/fuzz_linearizer.py b/test/external/fuzz_linearizer.py index e5242ac605..19ed23f5d6 100644 --- a/test/external/fuzz_linearizer.py +++ b/test/external/fuzz_linearizer.py @@ -16,7 +16,7 @@ if os.getenv("VALIDATE_HCQ", 0) != 0: try: import extra.qcom_gpu_driver.opencl_ioctl from tinygrad import Device - _, _ = Device["QCOM"], Device["GPU"] + _, _ = Device["QCOM"], Device["CL"] except Exception: pass from tinygrad import Tensor, Device, dtypes @@ -42,9 +42,9 @@ if getenv("VALIDATE_HCQ"): on_linearizer_did_run = extra.nv_gpu_driver.nv_ioctl.collect_last_launch_state compare_states = extra.nv_gpu_driver.nv_ioctl.compare_launch_state elif Device.DEFAULT == "QCOM": - print("VALIDATE_HCQ: Comparing QCOM to GPU") + print("VALIDATE_HCQ: Comparing QCOM to CL") import extra.qcom_gpu_driver.opencl_ioctl - validate_device = Device["GPU"] + validate_device = Device["CL"] on_linearizer_will_run = extra.qcom_gpu_driver.opencl_ioctl.before_launch on_linearizer_did_run = extra.qcom_gpu_driver.opencl_ioctl.collect_last_launch_state compare_states = extra.qcom_gpu_driver.opencl_ioctl.compare_launch_state @@ -302,7 +302,7 @@ if __name__ == "__main__": for i, ast in enumerate(ast_strs[:getenv("FUZZ_N", len(ast_strs))]): if (nth := getenv("FUZZ_NTH", -1)) != -1 and i != nth: continue if getenv("FUZZ_IMAGEONLY") and "dtypes.image" not in ast: continue - if "dtypes.image" in ast and Device.DEFAULT not in {"GPU", "QCOM"}: continue # IMAGE is only for GPU + if "dtypes.image" in ast and Device.DEFAULT not in {"CL", "QCOM"}: continue # IMAGE is only for CL if ast in seen_ast_strs: continue seen_ast_strs.add(ast) diff --git a/test/helpers.py b/test/helpers.py index 4833f425d9..cee64595f3 100644 --- a/test/helpers.py +++ b/test/helpers.py @@ -57,8 +57,8 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None): return out_buf.cast(uop.dtype.fmt).tolist()[0] def not_support_multi_device(): - # GPU and CUDA don't support multi device if in CI - return CI and REAL_DEV in ("GPU", "CUDA") + # CL and CUDA don't support multi device if in CI + return CI and REAL_DEV in ("CL", "CUDA") # NOTE: This will open REMOTE if it's the default device REAL_DEV = (Device.DEFAULT if Device.DEFAULT != "REMOTE" else Device['REMOTE'].properties.real_device) diff --git a/test/models/test_bert.py b/test/models/test_bert.py index 0e42ff9dd3..78e3367339 100644 --- a/test/models/test_bert.py +++ b/test/models/test_bert.py @@ -1,14 +1,14 @@ #!/usr/bin/env python import unittest +from tinygrad import Tensor import numpy as np -from tinygrad.tensor import Tensor import torch def get_question_samp(bsz, seq_len, vocab_size, seed): np.random.seed(seed) in_ids= np.random.randint(vocab_size, size=(bsz, seq_len)) mask = np.random.choice([True, False], size=(bsz, seq_len)) - seg_ids = np.random.randint(1, size=(bsz, seq_len)) + seg_ids = np.random.randint(2, size=(bsz, seq_len)) # type_vocab_size return in_ids, mask, seg_ids def set_equal_weights(mdl, torch_mdl): @@ -45,7 +45,7 @@ class TestBert(unittest.TestCase): seeds = (1337, 3141) bsz, seq_len = 1, 16 - for _, seed in enumerate(seeds): + for seed in seeds: in_ids, mask, seg_ids = get_question_samp(bsz, seq_len, config['vocab_size'], seed) out = mdl(Tensor(in_ids), Tensor(mask), Tensor(seg_ids)) torch_out = torch_mdl.forward(torch.from_numpy(in_ids).long(), torch.from_numpy(mask), torch.from_numpy(seg_ids).long())[:2] diff --git a/test/models/test_efficientnet.py b/test/models/test_efficientnet.py index 2892196517..8e434ba8aa 100644 --- a/test/models/test_efficientnet.py +++ b/test/models/test_efficientnet.py @@ -1,12 +1,10 @@ -import ast -import pathlib -import unittest +import ast, pathlib, unittest import numpy as np from PIL import Image -from tinygrad.helpers import getenv -from tinygrad.tensor import Tensor +from tinygrad import Tensor +from tinygrad.helpers import getenv, CI from extra.models.efficientnet import EfficientNet from extra.models.vit import ViT from extra.models.resnet import ResNet50 @@ -40,19 +38,13 @@ def preprocess(img, new=False): img /= np.array([0.229, 0.224, 0.225]).reshape((1, -1, 1, 1)) return img +def _infer(model: EfficientNet, img): + with Tensor.train(False): + out = model.forward(Tensor(img)).argmax(axis=-1) + return out.tolist() -def _infer(model: EfficientNet, img, bs=1): - old_training = Tensor.training - Tensor.training = False - img = preprocess(img) - # run the net - if bs > 1: img = img.repeat(bs, axis=0) - out = model.forward(Tensor(img)) - Tensor.training = old_training - return _LABELS[np.argmax(out.numpy()[0])] - -chicken_img = Image.open(pathlib.Path(__file__).parent / 'efficientnet/Chicken.jpg') -car_img = Image.open(pathlib.Path(__file__).parent / 'efficientnet/car.jpg') +chicken_img = preprocess(Image.open(pathlib.Path(__file__).parent / 'efficientnet/Chicken.jpg')) +car_img = preprocess(Image.open(pathlib.Path(__file__).parent / 'efficientnet/car.jpg')) class TestEfficientNet(unittest.TestCase): @classmethod @@ -64,17 +56,20 @@ class TestEfficientNet(unittest.TestCase): def tearDownClass(cls): del cls.model + @unittest.skipIf(CI, "covered by test_chicken_car") def test_chicken(self): - label = _infer(self.model, chicken_img) - self.assertEqual(label, "hen") - - def test_chicken_bigbatch(self): - label = _infer(self.model, chicken_img, 2) - self.assertEqual(label, "hen") + labels = _infer(self.model, chicken_img) + self.assertEqual(_LABELS[labels[0]], "hen") + @unittest.skipIf(CI, "covered by test_chicken_car") def test_car(self): - label = _infer(self.model, car_img) - self.assertEqual(label, "sports car, sport car") + labels = _infer(self.model, car_img) + self.assertEqual(_LABELS[labels[0]], "sports car, sport car") + + def test_chicken_car(self): + labels = _infer(self.model, np.concat([chicken_img, car_img], axis=0)) + self.assertEqual(_LABELS[labels[0]], "hen") + self.assertEqual(_LABELS[labels[1]], "sports car, sport car") class TestViT(unittest.TestCase): @classmethod @@ -87,12 +82,12 @@ class TestViT(unittest.TestCase): del cls.model def test_chicken(self): - label = _infer(self.model, chicken_img) - self.assertEqual(label, "cock") + labels = _infer(self.model, chicken_img) + self.assertEqual(_LABELS[labels[0]], "cock") def test_car(self): - label = _infer(self.model, car_img) - self.assertEqual(label, "racer, race car, racing car") + labels = _infer(self.model, car_img) + self.assertEqual(_LABELS[labels[0]], "racer, race car, racing car") class TestResNet(unittest.TestCase): @classmethod @@ -105,12 +100,12 @@ class TestResNet(unittest.TestCase): del cls.model def test_chicken(self): - label = _infer(self.model, chicken_img) - self.assertEqual(label, "hen") + labels = _infer(self.model, chicken_img) + self.assertEqual(_LABELS[labels[0]], "hen") def test_car(self): - label = _infer(self.model, car_img) - self.assertEqual(label, "sports car, sport car") + labels = _infer(self.model, car_img) + self.assertEqual(_LABELS[labels[0]], "sports car, sport car") if __name__ == '__main__': unittest.main() diff --git a/test/models/test_onnx.py b/test/models/test_onnx.py index d4199e7f49..34e5a1320d 100644 --- a/test/models/test_onnx.py +++ b/test/models/test_onnx.py @@ -5,12 +5,8 @@ from tinygrad.frontend.onnx import OnnxRunner from tinygrad.device import Device from tinygrad.helpers import fetch, Context -try: - from extra.onnx_helpers import validate - from extra.huggingface_onnx.huggingface_manager import DOWNLOADS_DIR, snapshot_download_with_retry - HUGGINGFACE_AVAILABLE = True -except ModuleNotFoundError: - HUGGINGFACE_AVAILABLE = False +from extra.onnx_helpers import validate +from extra.huggingface_onnx.huggingface_manager import DOWNLOADS_DIR, snapshot_download_with_retry def run_onnx_torch(onnx_model, inputs): import torch @@ -62,7 +58,7 @@ class TestOnnxModel(unittest.TestCase): print(cls, _LABELS[cls]) assert "car" in _LABELS[cls] or _LABELS[cls] == "convertible" -@unittest.skipUnless(HUGGINGFACE_AVAILABLE and Device.DEFAULT == "METAL", "only run on METAL") +@unittest.skipUnless(Device.DEFAULT == "METAL", "only run on METAL") class TestHuggingFaceOnnxModels(unittest.TestCase): @classmethod def setUpClass(cls): diff --git a/test/models/test_real_world.py b/test/models/test_real_world.py index c55aee0ad8..aa7345820e 100644 --- a/test/models/test_real_world.py +++ b/test/models/test_real_world.py @@ -53,8 +53,8 @@ class TestRealWorld(unittest.TestCase): @unittest.skipUnless(is_dtype_supported(dtypes.float16), "need dtypes.float16") def test_stable_diffusion(self): params = unet_params - params["model_ch"] = 16 - params["ctx_dim"] = 16 + params["model_ch"] = 8 + params["ctx_dim"] = 8 params["num_res_blocks"] = 1 params["n_heads"] = 2 model = UNetModel(**params) @@ -114,7 +114,7 @@ class TestRealWorld(unittest.TestCase): helper_test("train_mnist", lambda: (Tensor.randn(BS, 1, 28, 28),), train, 0.07, 93) - @unittest.skipIf(CI and Device.DEFAULT in {"CPU", "GPU"}, "slow") + @unittest.skipIf(CI and Device.DEFAULT in {"CPU", "CL"}, "slow") def test_train_cifar(self): with Tensor.train(): model = SpeedyResNet(Tensor.ones((12,3,2,2))) @@ -144,6 +144,7 @@ class TestRealWorld(unittest.TestCase): final_div_factor=1./(initial_div_factor*final_lr_ratio), total_steps=4) assert not np.isnan(lr_scheduler.min_lr), "lr too small or initial_div_facotr too big for half" + @unittest.skipIf(CI and Device.DEFAULT == "CPU", "slow") def test_bert(self): with Tensor.train(): args_tiny = {"attention_probs_dropout_prob": 0.0, "hidden_dropout_prob": 0.0, "vocab_size": 30522, "type_vocab_size": 2, @@ -167,9 +168,5 @@ class TestRealWorld(unittest.TestCase): helper_test("train_bert", lambda: (data["input_ids"], data["segment_ids"], data["input_mask"], data["masked_lm_positions"], \ data["masked_lm_ids"], data["masked_lm_weights"], data["next_sentence_labels"]), train, 0.25, 347) - def test_bert_fuse_arange(self): - with Context(FUSE_ARANGE=1): - self.test_bert() - if __name__ == '__main__': unittest.main() diff --git a/test/models/test_rnnt.py b/test/models/test_rnnt.py index f9d5e2c9db..c321b16d9d 100644 --- a/test/models/test_rnnt.py +++ b/test/models/test_rnnt.py @@ -1,8 +1,8 @@ #!/usr/bin/env python import unittest -import numpy as np -from tinygrad.tensor import Tensor +from tinygrad import Tensor from extra.models.rnnt import LSTM +import numpy as np import torch class TestRNNT(unittest.TestCase): diff --git a/test/models/test_train.py b/test/models/test_train.py index 605e6f6de1..fe5114742b 100644 --- a/test/models/test_train.py +++ b/test/models/test_train.py @@ -1,9 +1,8 @@ -import unittest -import time +import unittest, time import numpy as np +from tinygrad import Device from tinygrad.nn.state import get_parameters from tinygrad.nn import optim -from tinygrad.tensor import Device from tinygrad.helpers import getenv, CI from extra.training import train from extra.models.convnext import ConvNeXt @@ -27,7 +26,7 @@ def train_one_step(model,X,Y): print("done in %.2f ms" % (et*1000.)) def check_gc(): - if Device.DEFAULT == "GPU": + if Device.DEFAULT == "CL": from extra.introspection import print_objects assert print_objects() == 0 @@ -40,7 +39,6 @@ class TestTrain(unittest.TestCase): check_gc() @unittest.skipIf(CI, "slow") - @unittest.skipIf(Device.DEFAULT in ["METAL", "WEBGPU"], "too many buffers for webgpu and metal") def test_efficientnet(self): model = EfficientNet(0) X = np.zeros((BS,3,224,224), dtype=np.float32) @@ -49,7 +47,6 @@ class TestTrain(unittest.TestCase): check_gc() @unittest.skipIf(CI, "slow") - @unittest.skipIf(Device.DEFAULT in ["METAL", "WEBGPU"], "too many buffers for webgpu and metal") def test_vit(self): model = ViT() X = np.zeros((BS,3,224,224), dtype=np.float32) @@ -57,7 +54,7 @@ class TestTrain(unittest.TestCase): train_one_step(model,X,Y) check_gc() - @unittest.skipIf(Device.DEFAULT in ["METAL", "WEBGPU"], "too many buffers for webgpu and metal") + @unittest.skipIf(CI, "slow") def test_transformer(self): # this should be small GPT-2, but the param count is wrong # (real ff_dim is 768*4) diff --git a/test/opt/test_kernel_opts.py b/test/opt/test_kernel_opts.py index 24543c3a55..fda46a36c1 100644 --- a/test/opt/test_kernel_opts.py +++ b/test/opt/test_kernel_opts.py @@ -1,6 +1,6 @@ import unittest from tinygrad import Device, Tensor, dtypes -from tinygrad.helpers import CI +from tinygrad.helpers import CI, RANGEIFY from tinygrad.codegen.opt import Opt, OptOps, KernelOptError # TODO: write a clean version of this @@ -93,7 +93,7 @@ class TestKernelOpts(unittest.TestCase): a = Tensor.rand(8, N, 8, N) r = a.sum(axis=(1,3)) helper_linearizer_opt(r, [ - # openCL / GPU=1 is 256 max threads + # openCL / CL=1 is 256 max threads [Opt(OptOps.GROUPTOP, 0, 2)], [Opt(OptOps.GROUPTOP, 0, 32)], [Opt(OptOps.GROUPTOP, 1, 2)], [Opt(OptOps.GROUPTOP, 1, 32)], # Checking how it works with 1 grouped_reduce. [Opt(OptOps.GROUPTOP, 0, 2), Opt(OptOps.GROUPTOP, 1, 2)], @@ -327,13 +327,14 @@ class TestKernelOpts(unittest.TestCase): @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4") def test_arange_opts(self): a = Tensor.arange(128) + # NOTE: arange no longer has reduce ops available for opt helper_linearizer_opt(a, [ - [Opt(OptOps.GROUP, 0, 32)], - [Opt(OptOps.GROUPTOP, 0, 32)], + #[Opt(OptOps.GROUP, 0, 32)], + #[Opt(OptOps.GROUPTOP, 0, 32)], [Opt(op=OptOps.LOCAL, axis=0, arg=8)], [Opt(op=OptOps.LOCAL, axis=0, arg=8), Opt(op=OptOps.UPCAST, axis=0, arg=0)], - [Opt(op=OptOps.LOCAL, axis=0, arg=8), Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.GROUP, axis=0, arg=8)], - [Opt(op=OptOps.LOCAL, axis=0, arg=8), Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.GROUP, axis=0, arg=8), Opt(op=OptOps.UNROLL, axis=1, arg=4)], # noqa: E501 + #[Opt(op=OptOps.LOCAL, axis=0, arg=8), Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.GROUP, axis=0, arg=8)], + #[Opt(op=OptOps.LOCAL, axis=0, arg=8), Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.GROUP, axis=0, arg=8), Opt(op=OptOps.UNROLL, axis=1, arg=4)], # noqa: E501 ]) @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_threads, "test requires threads") @@ -350,5 +351,18 @@ class TestKernelOpts(unittest.TestCase): ] + [[Opt(OptOps.THREAD, 0, 4)] if Device[Device.DEFAULT].renderer.global_max[0] >= 4 else []] + [[Opt(OptOps.THREAD, 0, 8)] if Device[Device.DEFAULT].renderer.global_max[0] >= 8 else []]) + @unittest.skipUnless(RANGEIFY>=1, "Kernel only fuses with rangeify") + def test_double_sum_group(self): + a = Tensor.rand(4, 4, 4) + r = a.sum((1, 2)).sum() + with self.assertRaises(KernelOptError): + helper_linearizer_opt(r, [[Opt(OptOps.GROUPTOP, 0, 16)],]) + r = a.sum((1, 2)).sum() + with self.assertRaises(KernelOptError): + helper_linearizer_opt(r, [[Opt(OptOps.UNROLL, 1, 4), Opt(OptOps.GROUPTOP, 0, 16)],]) + r = a.sum((1, 2)).sum() + with self.assertRaises(KernelOptError): + helper_linearizer_opt(r, [[Opt(OptOps.GROUPTOP, 1, 4), Opt(OptOps.GROUPTOP, 0, 16)],]) + if __name__ == '__main__': unittest.main() diff --git a/test/speed/external_test_copy_speed.py b/test/speed/external_test_copy_speed.py index 351c7d993a..359a2499ac 100644 --- a/test/speed/external_test_copy_speed.py +++ b/test/speed/external_test_copy_speed.py @@ -77,9 +77,9 @@ class TestCopySpeed(unittest.TestCase): np.testing.assert_equal(t.numpy(), x.numpy()) @unittest.skipIf(CI, "CI doesn't have 6 GPUs") - @unittest.skipIf(Device.DEFAULT != "GPU", "only test this on GPU") + @unittest.skipIf(Device.DEFAULT != "CL", "only test this on CL") def testCopyCPUto6GPUs(self): - from tinygrad.runtime.ops_gpu import CLDevice + 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") @@ -87,8 +87,8 @@ class TestCopySpeed(unittest.TestCase): with Timing("sync: ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s ({t.nbytes()*6/ns:.2f} GB/s total)"): with Timing("queue: "): for g in range(6): - t.to(f"gpu:{g}").realize() - Device["gpu"].synchronize() + t.to(f"CL:{g}").realize() + Device["CL"].synchronize() if __name__ == '__main__': unittest.main() diff --git a/test/test_arange.py b/test/test_arange.py index a21006ac62..a46b38a087 100644 --- a/test/test_arange.py +++ b/test/test_arange.py @@ -1,51 +1,29 @@ import unittest import numpy as np from tinygrad import Tensor, GlobalCounters, dtypes, nn, Device, Variable -from tinygrad.helpers import CI, Context, getenv +from tinygrad.helpers import CI, Context, getenv, RANGEIFY from tinygrad.engine.realize import run_schedule -from tinygrad.codegen.opt import Opt, OptOps from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program from tinygrad.uop.ops import Ops class TestArange(unittest.TestCase): - def _get_flops(self, N, opts=None): + def _get_flops(self, N): GlobalCounters.reset() tt = Tensor.arange(N) sched = tt.schedule() self.assertEqual(len(sched), 1) - p = get_program(sched[-1].ast, opts=opts) - print(p.name) - #print(p.src) + p = get_program(sched[-1].ast) ExecItem(CompiledRunner(p), [tt.uop.buffer]).run() np.testing.assert_equal(tt.numpy(), np.arange(N)) return p.estimates.ops - def test_complexity(self, opts=None, limit=None): - f1 = self._get_flops(256, opts) - f2 = self._get_flops(2560, opts) - print(f"{f1=}, {f2=}") - # add 1 to avoid divide by 0. arange is 0 flops now! - assert (f1 < 6000 and f2 < 6000) or ((f2+1) / (f1+1) < 16), f"bad complexity, flops {(f2+1) / (f1+1):.1f}X while inputs 10X" - if limit is not None and not getenv("PTX"): - # PTX counts index ALU in flops - assert f1 <= limit, f"{f1=}, {limit=}" + def test_complexity(self): + self.assertEqual(self._get_flops(256), 0) + self.assertEqual(self._get_flops(2560), 0) - def test_complexity_w_upcast(self): return self.test_complexity([Opt(OptOps.UPCAST, 0, 4)], limit=0) - def test_complexity_w_unroll2(self): return self.test_complexity([Opt(OptOps.UNROLL, 0, 2)], limit=0) - def test_complexity_w_unroll4(self): return self.test_complexity([Opt(OptOps.UNROLL, 0, 4)], limit=0) - def test_complexity_w_unroll8(self): return self.test_complexity([Opt(OptOps.UNROLL, 0, 8)], limit=0) - def test_complexity_w_upcast_and_unroll(self): return self.test_complexity([Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UNROLL, 0, 4)], limit=0) - - if Device.default.renderer.has_local: - # TODO: fix limit - def test_complexity_w_group(self): return self.test_complexity([Opt(OptOps.GROUP, 0, 16)], limit=81920) - def test_complexity_w_group_top(self): return self.test_complexity([Opt(OptOps.GROUPTOP, 0, 16)], limit=106496) - - def test_complexity_w_local(self): return self.test_complexity([Opt(OptOps.LOCAL, 0, 16)], limit=0) - @unittest.skip("doesn't work yet. TODO: this absolutely should work") - def test_complexity_w_local_unroll4(self): return self.test_complexity([Opt(OptOps.LOCAL, 0, 16), Opt(OptOps.UNROLL, 0, 4)], limit=0) - @unittest.skip("doesn't work yet") - def test_complexity_w_local_and_padto(self): return self.test_complexity([Opt(OptOps.LOCAL, 0, 16), Opt(OptOps.PADTO, axis=1, arg=32)]) + def test_arange_cat(self): + t = Tensor.arange(2, dtype=dtypes.int)+Tensor([3]) + self.assertEqual(t.cat(t).tolist(), [3, 4, 3, 4]) class TestRand(unittest.TestCase): def test_fused_rand_less_ops(self, noopt=1): @@ -133,7 +111,7 @@ class TestIndexing(unittest.TestCase): X = dataset[idxs] assert X.shape == (4,DDIM) sched = X.schedule() - self.assertEqual(len(sched), 2) + self.assertEqual(len(sched), 1 if RANGEIFY else 2) run_schedule(sched) assert GlobalCounters.global_ops < 4*DSET, f"too many ops {GlobalCounters.global_ops} != {4*DSET}" np.testing.assert_allclose(real_index, X.numpy()) diff --git a/test/test_const_folding.py b/test/test_const_folding.py index d4a79fbb46..37d02ed9ac 100644 --- a/test/test_const_folding.py +++ b/test/test_const_folding.py @@ -3,6 +3,7 @@ from tinygrad import Tensor, Device, dtypes from tinygrad.dtype import DType, ConstType from tinygrad.uop.ops import Ops, UOp from tinygrad.codegen import full_rewrite_to_sink +from tinygrad.helpers import RANGEIFY from tinygrad.device import is_dtype_supported import numpy as np from test.helpers import not_support_multi_device @@ -155,7 +156,7 @@ class TestMovedConstFolding(unittest.TestCase): def test_add_padded_zero(self): # TODO: it's 1 now, this might be possible to fold - _check_ast_count(1, Tensor([1.0, 2, 3, 4]) + Tensor.zeros(2).pad(((1, 1),))) + _check_ast_count(0 if RANGEIFY else 1, Tensor([1.0, 2, 3, 4]) + Tensor.zeros(2).pad(((1, 1),))) def test_mul_shrunk_one(self): _check_ast_count(0, Tensor([1.0, 2, 3, 4]) * Tensor.ones(6).shrink(((1, 5),))) diff --git a/test/test_dtype.py b/test/test_dtype.py index 6aa8ddf2eb..dc3020686c 100644 --- a/test/test_dtype.py +++ b/test/test_dtype.py @@ -4,12 +4,13 @@ import torch from typing import Any, List from tinygrad.device import is_dtype_supported from tinygrad.helpers import getenv, DEBUG, CI -from tinygrad.dtype import DType, DTYPES_DICT, least_upper_dtype, fp8_to_float, float_to_fp8, _to_np_dtype, _to_torch_dtype +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 from tinygrad import Device, Tensor, dtypes from hypothesis import assume, given, settings, strategies as strat from test.helpers import rand_for_dtype from test.unit.test_dtype_spec import _assert_eq, core_dtypes, dtype_ints, dtype_floats, FP8E4M3_MAX, FP8E5M2_MAX -import ml_dtypes import pytest pytestmark = pytest.mark.filterwarnings("ignore") @@ -25,6 +26,7 @@ def get_available_cast_dtypes(dtype: DType) -> List[DType]: def _to_torch_storage_type(dtype:DType): if dtype == dtypes.bfloat16: return torch.float32 + if dtype in dtypes.fp8s: return torch.float32 return _to_torch_dtype(dtype) def _test_to_np(a:Tensor, np_dtype, target): @@ -47,12 +49,15 @@ def _test_cast(a:Tensor, target_dtype:DType): # TODO: struct.pack cannot pack value > 65504 (max of half) into e format a = (a > 65504).where(65504, a) - _test_op(lambda: a.cast(target_dtype), target_dtype, list(a.numpy().astype(_to_np_dtype(target_dtype)))) + expected = list(a.numpy().astype(_to_np_dtype(target_dtype))) + if target_dtype in dtypes.fp8s: expected = list(map(lambda x: truncate[target_dtype](x), expected)) + _test_op(lambda: a.cast(target_dtype), target_dtype, expected) def _test_bitcast(a:Tensor, target_dtype:DType, target=None): - if getenv("PTX") and a.dtype == dtypes.int8 and target_dtype.itemsize != a.dtype.itemsize: + if isinstance(Device[Device.DEFAULT].renderer, PTXRenderer) and a.dtype == dtypes.int8 and target_dtype.itemsize != a.dtype.itemsize: raise unittest.SkipTest("shape changing bitcast of int8 broken on PTX") - expected = torch.tensor(a.tolist(), dtype=_to_torch_storage_type(a.dtype)).view(_to_torch_dtype(target_dtype)) - _test_op(lambda: a.bitcast(target_dtype), target_dtype, target or expected.tolist()) + expected = torch.tensor(a.tolist(), dtype=_to_torch_storage_type(a.dtype)).view(_to_torch_dtype(target_dtype)).tolist() + if target_dtype in dtypes.fp8s: expected = list(map(lambda x: fp8_to_float(x, target_dtype), expected)) + _test_op(lambda: a.bitcast(target_dtype), target_dtype, target or expected) class TestDType(unittest.TestCase): DTYPE: Any = None @@ -100,8 +105,7 @@ class TestDType(unittest.TestCase): )) @unittest.skipIf(Device.DEFAULT == "PYTHON", "skip for now") - @unittest.skipIf(getenv("PTX"), "skip for now") - @unittest.skipIf(getenv("NIR"), "skip for now") + @unittest.skipIf(isinstance(r:=Device[Device.DEFAULT].renderer, PTXRenderer) or isinstance(r, NIRRenderer), "skip for now") def test_uint_overflow(self): if not dtypes.is_unsigned(self.DTYPE): raise unittest.SkipTest("only for unsigned") v = dtypes.max(self.DTYPE) @@ -129,11 +133,10 @@ class TestDType(unittest.TestCase): np.testing.assert_allclose(tin, tor, atol=1e-6, rtol=1e-3) def test_finfo(self): - if self.DTYPE not in [dtypes.float16, dtypes.bfloat16, dtypes.float32, dtypes.float64]: return - info = ml_dtypes.finfo(ml_dtypes.bfloat16 if self.DTYPE is dtypes.bfloat16 else _to_np_dtype(self.DTYPE)) - assert info.bits == self.DTYPE.itemsize*8 - assert info.nexp == dtypes.finfo(self.DTYPE)[0] - assert info.nmant == dtypes.finfo(self.DTYPE)[1] + 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.itemsize*8) + 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) @@ -151,7 +154,8 @@ class TestFp8s(unittest.TestCase): 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)) - def test_float_to_fp8e4m3(self, x): np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3), ml_dtypes.float8_e4m3fn(x).tobytes()[0]) + def test_float_to_fp8e4m3(self, x): + np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3), torch.tensor(x, dtype=torch.float8_e4m3fn).view(torch.uint8).item()) def test_float_to_fp8e4m3_extreme_values(self): np.testing.assert_equal(float_to_fp8(FP8E4M3_MAX, dtypes.fp8e4m3), 126) @@ -164,7 +168,8 @@ class TestFp8sConversions(unittest.TestCase): np.testing.assert_equal(float_to_fp8(-math.nan, dtypes.fp8e4m3), 255) @given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E5M2_MAX, max_value=FP8E5M2_MAX)) - def test_float_to_fp8e5m2(self, x): np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e5m2), ml_dtypes.float8_e5m2(x).tobytes()[0]) + def test_float_to_fp8e5m2(self, x): + np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e5m2), torch.tensor(x, dtype=torch.float8_e5m2).view(torch.uint8).item()) def test_float_to_fp8e5m2_extreme_values(self): np.testing.assert_equal(float_to_fp8(FP8E5M2_MAX, dtypes.fp8e5m2), 123) @@ -177,10 +182,12 @@ class TestFp8sConversions(unittest.TestCase): np.testing.assert_equal(float_to_fp8(-math.nan, dtypes.fp8e5m2), 254) @given(strat.integers(min_value=0, max_value=255)) - def test_fp8e4m3_to_float(self, x): np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e4m3), np.uint8(x).view(ml_dtypes.float8_e4m3fn).item()) + def test_fp8e4m3_to_float(self, x): + np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e4m3), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e4m3fn).float().item()) @given(strat.integers(min_value=0, max_value=255)) - def test_fp8e5m2_to_float(self, x): np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e5m2), np.uint8(x).view(ml_dtypes.float8_e5m2).item()) + def test_fp8e5m2_to_float(self, x): + np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e5m2), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e5m2).float().item()) @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), "bfloat16 not supported") class TestBFloat16(unittest.TestCase): @@ -256,8 +263,9 @@ class TestFloatDType(TestDType): class TestDoubleDType(TestDType): DTYPE = dtypes.double - @unittest.skipIf((CI and Device.DEFAULT in {"CUDA", "NV"}) or getenv("PTX") or getenv("NIR"), - "conversion not supported on CI, CUDA, PTX and NIR") # TODO: why not? + @unittest.skipIf((CI and Device.DEFAULT in {"CUDA", "NV"}) or \ + isinstance(Device[Device.DEFAULT].renderer, PTXRenderer) or \ + isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "conversion not supported on CI CUDA, PTX, and NIR") # TODO: why not? def test_float64_increased_precision(self): for func in [ lambda t: t.exp(), @@ -281,21 +289,21 @@ class TestDoubleDType(TestDType): class TestInt8DType(TestDType): DTYPE = dtypes.int8 - @unittest.skipIf(getenv("CUDA",0)==1 or getenv("PTX", 0)==1, "cuda saturation works differently") + @unittest.skipIf(getenv("CUDA",0)==1 or isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "cuda saturation works differently") def test_int8_to_uint8_negative(self): _test_op(lambda: Tensor([-1, -2, -3, -4], dtype=dtypes.int8).cast(dtypes.uint8), dtypes.uint8, [255, 254, 253, 252]) def test_int8_to_uint16_negative(self): _test_op(lambda: Tensor([-1, -2, -3, -4], dtype=dtypes.int8).cast(dtypes.uint16), dtypes.uint16, [2**16-1, 2**16-2, 2**16-3, 2**16-4]) - @unittest.skipIf(getenv("PTX"), "broken in ptx") + @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken in ptx") def test_bitcast_alt(self): a = Tensor([72, -90, 27, 40, -53, 70, 96, 51], dtype=dtypes.int8).bitcast(dtypes.short) self.assertListEqual(a.tolist(), [-22968, 10267, 18123, 13152]) class TestUint8DType(TestDType): DTYPE = dtypes.uint8 - @unittest.skipIf(getenv("CUDA",0)==1 or getenv("PTX", 0)==1, "cuda saturation works differently") + @unittest.skipIf(getenv("CUDA",0)==1 or isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "cuda saturation works differently") def test_uint8_to_int8_overflow(self): _test_op(lambda: Tensor([255, 254, 253, 252], dtype=dtypes.uint8).cast(dtypes.int8), dtypes.int8, [-1, -2, -3, -4]) @@ -303,9 +311,11 @@ class TestBitCast(unittest.TestCase): @given(strat.sampled_from(dtype_ints + dtype_floats), strat.sampled_from(dtype_ints + dtype_floats)) def test_shape_change_bitcast(self, dt1, dt2): # NOTE: this has to be assume to prevent hypothesis from skipping all samples - assume(not (getenv("PTX") and dt1 == dtypes.int8)) # TODO: bitcasting int8 fails in PTX + assume(not (isinstance(Device[Device.DEFAULT].renderer, PTXRenderer) and dt1 == dtypes.int8)) # TODO: bitcasting int8 fails in PTX data = rand_for_dtype(dt1, 32).reshape(2, 2, 8) expected = torch.tensor(data.tolist(), dtype=_to_torch_storage_type(dt1)).view(_to_torch_dtype(dt2)) + if dt2 in dtypes.fp8s: + expected = torch.tensor(list(map(lambda x: fp8_to_float(x, dt2), expected.view(-1).tolist()))).view_as(expected) _test_op(lambda: Tensor(data, dtype=dt1).bitcast(dt2), dt2, expected.tolist()) def test_shape_change_bitcast_exceptions(self): @@ -348,6 +358,9 @@ class TestBoolDType(TestDType): DTYPE = dtypes.bool class TestBFloat16Type(TestDType): DTYPE = dtypes.bfloat16 +class TestFp8e4m3(TestDType): DTYPE = dtypes.fp8e4m3 +class TestFp8e5m2(TestDType): DTYPE = dtypes.fp8e5m2 + class TestPtrDType(unittest.TestCase): def test_vec_double(self): dt1 = dtypes.float.vec(4).ptr().vec(4) @@ -424,7 +437,7 @@ class TestDtypeUsage(unittest.TestCase): class TestOpsBFloat16(unittest.TestCase): def test_cast(self): # TODO: helper_test_op breaks in unrelated part - # TODO: wrong output with GPU=1 on mac + # TODO: wrong output with CL=1 on mac data = [60000.0, 70000.0, 80000.0] np.testing.assert_allclose(Tensor(data).cast("bfloat16").numpy(), torch.tensor(data).type(torch.bfloat16).float().numpy()) diff --git a/test/test_dtype_alu.py b/test/test_dtype_alu.py index d1694bd58e..19debe432d 100644 --- a/test/test_dtype_alu.py +++ b/test/test_dtype_alu.py @@ -1,13 +1,14 @@ import unittest, operator, math from tinygrad import Tensor, dtypes, Device -from tinygrad.dtype import DType +from tinygrad.dtype import DType, truncate from tinygrad.helpers import CI, 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 import numpy as np import pytest -from hypothesis import given, strategies as strat, settings, HealthCheck +from hypothesis import assume, given, strategies as strat, settings, HealthCheck pytestmark = pytest.mark.filterwarnings("ignore") @@ -47,6 +48,8 @@ class ht: int64 = strat.integers(-9223372036854775808, 9223372036854775807) bool = strat.booleans() ht.bfloat16 = ht.uint16 +ht.fp8e4m3 = ht.uint8 +ht.fp8e5m2 = ht.uint8 def universal_test(a, b, dtype, op): if not isinstance(op, tuple): op = (op, op) @@ -56,8 +59,9 @@ def universal_test(a, b, dtype, op): ta, tb = Tensor([a], dtype=dtype), Tensor([b], dtype=dtype) tensor_value = (op[0](ta, tb)).numpy() numpy_value = op[1](ta.numpy(), tb.numpy()) + if dtype in dtypes.fp8s: numpy_value = truncate[dtype](numpy_value) if dtype in dtypes.floats: - atol, rtol = {dtypes.bfloat16:(1e-3, 1e-2)}.get(dtype, (1e-10, 1e-7)) + atol, rtol = {dtypes.bfloat16:(1e-3, 1e-2), dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2:(1.0, 5e-1)}.get(dtype, (1e-10, 1e-7)) np.testing.assert_allclose(tensor_value, numpy_value, atol=atol, rtol=rtol) else: np.testing.assert_equal(tensor_value, numpy_value) @@ -70,8 +74,10 @@ def universal_test_unary(a, dtype, op): out: Tensor = op[0](ta) tensor_value = out.numpy() numpy_value = op[1](ta.numpy()) + if dtype in dtypes.fp8s: numpy_value = truncate[dtype](numpy_value) if dtype in dtypes.floats: - atol, rtol = {dtypes.float16:(1e-3, 1e-2), dtypes.bfloat16:(1e-3, 2e-2)}.get(dtype, (1e-6, 1e-5)) + atol, rtol = { dtypes.float16:(1e-3, 1e-2), dtypes.bfloat16:(1e-3, 2e-2), + dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2: (1.0, 5e-1)}.get(dtype, (1e-6, 1e-5)) np.testing.assert_allclose(tensor_value, numpy_value, atol=atol, rtol=rtol) else: np.testing.assert_equal(tensor_value, numpy_value) @@ -91,7 +97,7 @@ def universal_test_midcast(a, b, c, op1, op2, d1:DType, d2:DType): an, bn, cn = np.array([a]).astype(_to_np_dtype(d1)), np.array([b]).astype(_to_np_dtype(d1)), np.array([c]).astype(_to_np_dtype(d2)) tensor_value = op2[0](op1[0](at, bt).cast(d2), ct).numpy() numpy_value = op2[1](op1[1](an, bn).astype(_to_np_dtype(d2)), cn) - np.testing.assert_allclose(tensor_value, numpy_value, rtol=1e-6 if getenv("PTX") else 1e-7) + 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}") @@ -110,6 +116,16 @@ class TestDTypeALU(unittest.TestCase): def test_bfloat16(self, a, b, op): universal_test(from_storage_scalar(a, dtypes.bfloat16), from_storage_scalar(a, dtypes.bfloat16), dtypes.bfloat16, op) + @unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3), 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) + + @unittest.skipUnless(is_dtype_supported(dtypes.fp8e5m2), 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) + @given(ht.float32, strat.sampled_from(unary_operations)) def test_float32_unary(self, a, op): universal_test_unary(a, dtypes.float32, op) @@ -121,6 +137,18 @@ class TestDTypeALU(unittest.TestCase): @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) + @unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3), 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) + 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}") + @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) + universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2), dtypes.fp8e5m2, 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) diff --git a/test/test_image_dtype.py b/test/test_image_dtype.py index 08d2c04c32..6adab73e51 100644 --- a/test/test_image_dtype.py +++ b/test/test_image_dtype.py @@ -7,7 +7,7 @@ from tinygrad.engine.realize import lower_schedule from tinygrad.helpers import prod, unwrap from test.helpers import REAL_DEV -IMAGE_SUPPORTED_DEVICES = ("QCOM", "GPU") +IMAGE_SUPPORTED_DEVICES = ("QCOM", "CL") @unittest.skipUnless(REAL_DEV in IMAGE_SUPPORTED_DEVICES, "Images not supported") class TestImageCopy(unittest.TestCase): diff --git a/test/test_jit.py b/test/test_jit.py index f159385e37..11657e3397 100644 --- a/test/test_jit.py +++ b/test/test_jit.py @@ -609,21 +609,22 @@ class TestJitFree(unittest.TestCase): ext_tensor = Tensor([1,24,23,45,1]) @TinyJit def fxn(x:Tensor): - out = (x*2+ext_tensor).reshape(5,1).expand(5, 100).contiguous() - return out.sum() + t1 = (x * 2).contiguous().realize() + t2 = (t1 + ext_tensor).contiguous().realize() + out = (t2.sum()).contiguous().realize() + return out for i in range(5): - out = fxn(Tensor([i,1,2,3,4])) - self.assertEqual(out.item(), 11400+200*i) + out = fxn(inp:=Tensor([i,1,2,3,4])) + self.assertEqual(out.item(), 114+2*i) pre_free = GlobalCounters.mem_used fxn.captured.free_intermediates() savings_after_free = pre_free - GlobalCounters.mem_used - # Different allocator implementations have different savings. - expected_savings = 8196 if hasattr(Device[Device.DEFAULT].allocator, '_offset') else 2024 + expected_savings = (len(inp) * inp.dtype.itemsize * 2) + dtypes.float32.itemsize # (t1 and t2) + out self.assertEqual(savings_after_free, expected_savings) out = fxn(Tensor([11,1,2,3,4])) - self.assertEqual(out.item(), 13600) + self.assertEqual(out.item(), 136) # Try one more time... pre_free = GlobalCounters.mem_used @@ -633,7 +634,7 @@ class TestJitFree(unittest.TestCase): self.assertEqual(savings_after_free, expected_savings) out = fxn(Tensor([11,1,2,3,4])) - self.assertEqual(out.item(), 13600) + self.assertEqual(out.item(), 136) def test_updated_not_freed(self): x = Tensor([1]).realize() diff --git a/test/test_kernel_cache.py b/test/test_kernel_cache.py index 164b501a41..a4f0f2193a 100644 --- a/test/test_kernel_cache.py +++ b/test/test_kernel_cache.py @@ -16,14 +16,14 @@ class TestKernelCache(unittest.TestCase): a1 = Tensor.rand(4,4).realize() b1 = Tensor.rand(4,4).realize() - orig_compile_func = Device['CPU'].compiler - Device['CPU'].compiler = None # making it not callable + orig_compile_func = Device['CPU'].compiler.compile_cached + Device['CPU'].compiler.compile_cached = None # making it not callable try: x1 = a1 + b1 + unique_const x1.realize() # Same kernel should be from cache. finally: - Device['CPU'].compiler = orig_compile_func + Device['CPU'].compiler.compile_cached = orig_compile_func if __name__ == "__main__": unittest.main() diff --git a/test/test_linearizer.py b/test/test_linearizer.py index 23df5d93a5..1726e9aafc 100644 --- a/test/test_linearizer.py +++ b/test/test_linearizer.py @@ -10,9 +10,10 @@ from tinygrad.shape.shapetracker import ShapeTracker from tinygrad.shape.view import View from tinygrad.tensor import Tensor, _to_np_dtype from tinygrad.engine.realize import run_schedule, lower_schedule, CompiledRunner, get_program -from tinygrad.helpers import Context, getenv, flatten, dedup, TC_SELECT, TC_OPT +from tinygrad.helpers import Context, flatten, dedup, TC_SELECT, TC_OPT from tinygrad.dtype import DType, dtypes, PtrDType, AddrSpace from tinygrad.codegen import apply_rewrites, rewrites_for_views +from tinygrad.renderer.ptx import PTXRenderer class TestLinearizer(unittest.TestCase): def test_arg_dedup(self): @@ -155,7 +156,7 @@ class TestLinearizer(unittest.TestCase): @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") - @unittest.skipIf(getenv("PTX") or getenv("NIR"), "broken on ptx/nir because of indexing patterns") + @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken on ptx for some reason") def test_upcast_with_locals(self): x, y = Tensor.rand(1,128), Tensor.rand(128, 128) r = (x@y).relu() @@ -366,7 +367,7 @@ class TestLinearizer(unittest.TestCase): helper(Tensor.arange(255), max_ops=2) @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4") - @unittest.skipIf(getenv("PTX") or getenv("NIR"), "broken on ptx/nir because of indexing patterns") + @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken on ptx for some reason") def test_grouped_store_phis(self): """ float4 acc0 = float4(0.0,0.0,0.0,0.0); @@ -420,7 +421,7 @@ class TestLinearizer(unittest.TestCase): @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") - @unittest.skipIf(getenv("PTX") or getenv("NIR"), "broken on ptx/nir because of indexing patterns") + @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken on ptx for some reason") def test_grouped_store_local_only(self): x, y = Tensor.rand(1,128), Tensor.rand(128, 128) r = (x@y).relu() @@ -481,7 +482,7 @@ def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]: assert s[-1].ast.op is Ops.SINK, f"helper_realized_ast expects a SINK {s[-1]}" # now all input buffers in s[-1] should be realized # create fresh buffers for the outputs - bufs = [Buffer((x).device, x.size, x.dtype).allocate() if i < len(s[-1].ast.src) else x for i,x in enumerate(s[-1].bufs)] + bufs = [Buffer(x.device, x.size, x.dtype).allocate() if i < len(s[-1].ast.src) else x for i,x in enumerate(s[-1].bufs)] return push_views(s[-1].ast), bufs def helper_linearizer_ast(ast:UOp, inputs:list[Tensor], *args, **kwargs): @@ -503,7 +504,7 @@ def reset_bufs(bufs:list[Buffer]): def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[], apply_tc=False, atol=1e-4, rtol=1e-4, color_sizes=[], wanna_output=[]): - outbufs = [real_bufs[x.src[0].base.arg] for x in realized_ast.src] + outbufs = real_bufs[:len(realized_ast.src)] device = real_bufs[0].device wanna_output = [np.array(x).flatten() for x in wanna_output] diff --git a/test/test_linearizer_dumb.py b/test/test_linearizer_dumb.py index 23e534b4a7..4adca10df0 100644 --- a/test/test_linearizer_dumb.py +++ b/test/test_linearizer_dumb.py @@ -6,13 +6,12 @@ import unittest from tinygrad import Device, dtypes from tinygrad.device import is_dtype_supported from tinygrad.uop.ops import UOp, Ops, AxisType, KernelInfo -from tinygrad.helpers import getenv from tinygrad.shape.shapetracker import ShapeTracker, View from tinygrad.codegen.opt.search import Opt, OptOps from tinygrad.engine.realize import get_program +from tinygrad.renderer.ptx import PTXRenderer class TestLinearizerFailure(unittest.TestCase): - @unittest.expectedFailure @unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL") def test_failure_beam_mnist(self): c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.uchar.ptr(4014080), arg=0, src=()) @@ -47,7 +46,8 @@ class TestLinearizerDumb(unittest.TestCase): c10 = UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1000, 1000), strides=(0, 0), offset=0, mask=None, contiguous=False),)), src=()) c11 = c1.store((c4.alu(Ops.CMPNE, c7).alu(Ops.CMPNE, UOp.const(dtypes.bool, True, src=c8)).cast(dtypes.int)*(c9.f(Ops.VALID, dtype=dtypes.bool).where(UOp.const(dtypes.int, -1, src=c10), UOp.const(dtypes.int, 0, src=c10)).f(Ops.REDUCE_AXIS, arg=(Ops.ADD, (1,)))+UOp.const(dtypes.int, 1000, src=c8)))) ast = c11.sink() - opts = [Opt(op=OptOps.UNROLL, axis=0, arg=4), Opt(op=OptOps.LOCAL, axis=0, arg=8)] + #opts = [Opt(op=OptOps.UNROLL, axis=0, arg=4), Opt(op=OptOps.LOCAL, axis=0, arg=8)] + opts = [Opt(op=OptOps.LOCAL, axis=0, arg=8)] prg = get_program(ast, Device[Device.DEFAULT].renderer, opts) print(prg.src) assert prg.uops is not None and not any(uop.op is Ops.MAX for uop in prg.uops), "leftover MAX" @@ -93,7 +93,7 @@ class TestLinearizerDumb(unittest.TestCase): @unittest.expectedFailure @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "need float4") - @unittest.skipIf(getenv("PTX"), "this is somehow correct in PTX") + @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "this is somehow correct in PTX") def test_upcasted_stores_out_of_order(self): c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(9360), arg=0, src=()) c1 = c0.view(ShapeTracker(views=(View(shape=(4, 5, 13, 1, 1, 1, 1, 1, 4, 3, 3), strides=(2340, 468, 36, 0, 0, 0, 0, 0, 9, 3, 1), offset=0, mask=None, contiguous=True),))) diff --git a/test/test_method_cache.py b/test/test_method_cache.py index 497b406925..ce413e7709 100644 --- a/test/test_method_cache.py +++ b/test/test_method_cache.py @@ -5,9 +5,9 @@ from tinygrad.nn.state import get_state_dict class TestMethodCache(unittest.TestCase): def setUp(self): - self.backup_compiler = Device[Device.DEFAULT].compiler + self.backup_compiler = Device[Device.DEFAULT].compiler.compile_cached def tearDown(self): - Device[Device.DEFAULT].compiler = self.backup_compiler + Device[Device.DEFAULT].compiler.compile_cached = self.backup_compiler def test_simple_methodcache(self): a = Tensor([1]) @@ -15,19 +15,19 @@ class TestMethodCache(unittest.TestCase): c = Tensor([3]) d = Tensor([4]) (a+b).realize() - Device[Device.DEFAULT].compiler = None + Device[Device.DEFAULT].compiler.compile_cached = None (c+d).realize() def test_nested_methodcache(self): a,b,c,d = Tensor([1]), Tensor([2]), Tensor([3]), Tensor([4]) ((a+b)+(a+b)).realize() - Device[Device.DEFAULT].compiler = None + Device[Device.DEFAULT].compiler.compile_cached = None ((c+d)+(c+d)).realize() def test_nested_methodcache_swap(self): a,b,c,d = Tensor([1]), Tensor([2]), Tensor([3]), Tensor([4]) ((a+b)+(c+d)).realize() - Device[Device.DEFAULT].compiler = None + Device[Device.DEFAULT].compiler.compile_cached = None ((c+d)+(a+b)).realize() @unittest.skip("incorrect use of transformer") @@ -38,7 +38,7 @@ class TestMethodCache(unittest.TestCase): # NOTE: you have to do this twice due to the k-v cache for i in range(3): model(Tensor([[1,2,3,4]]), Variable("start_pos", 0, 10).bind(i)).realize() for i in range(3): model(Tensor([[1,2,3,4]]), Variable("start_pos", 0, 10).bind(i)).realize() - Device[Device.DEFAULT].compiler = None + Device[Device.DEFAULT].compiler.compile_cached = None for i in range(3): model(Tensor([[1,2,3,4]]), Variable("start_pos", 0, 10).bind(i)).realize() if __name__ == '__main__': diff --git a/test/test_multitensor.py b/test/test_multitensor.py index f4a0a7b398..86159ba375 100644 --- a/test/test_multitensor.py +++ b/test/test_multitensor.py @@ -2,7 +2,7 @@ import unittest, functools, 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 CI, getenv, prod, Context +from tinygrad.helpers import CI, getenv, prod, Context, RANGEIFY from tinygrad.nn.state import get_parameters, get_state_dict from tinygrad.engine.realize import lower_schedule, BufferCopy, CompiledRunner, run_schedule import numpy as np @@ -178,16 +178,14 @@ class TestMultiTensor(unittest.TestCase): run_schedule(sched) np.testing.assert_equal(xt.numpy(), X_np[i*2:i*2+2]) - @given(strat.sampled_from((4, 5)), strat.sampled_from((devices_2, devices_3)), + @given(strat.sampled_from((devices_2, devices_3)), strat.sampled_from((Ops.ADD, Ops.MUL, Ops.MAX)), - strat.sampled_from((None, 0, 1)), strat.sampled_from((None, 0, 1)), strat.sampled_from((1, 0, -1))) - def test_simple_reduce(self, N, devices, rop, shard_axis, reduce_axis, sign): - N = N * len(devices) - X = Tensor.rand(N*N).reshape(N, N).mul(sign) + strat.sampled_from((None, 0, 1)), strat.sampled_from((None, 0, 1))) + def test_simple_reduce(self, devices, rop, shard_axis, reduce_axis): + N = 4 * len(devices) + X = (Tensor.rand(N*N)-1).reshape(N, N).shard_(devices, shard_axis) n = X.numpy() - X.shard_(devices, shard_axis) - f = {Ops.ADD: lambda x: x.sum(reduce_axis), Ops.MUL: lambda x: x.prod(reduce_axis), - Ops.MAX: lambda x: x.max(reduce_axis)}[rop] + f = {Ops.ADD: lambda x: x.sum(reduce_axis), Ops.MUL: lambda x: x.prod(reduce_axis), Ops.MAX: lambda x: x.max(reduce_axis)}[rop] fX = f(X) fn = f(n) np.testing.assert_allclose(fX.numpy(), fn, rtol=1e-6, atol=1e-6) @@ -374,6 +372,7 @@ class TestMultiTensor(unittest.TestCase): # NOTE: this is failing on LLVM CI, no idea why. Works locally. @unittest.skipIf(CI and REAL_DEV in ("CUDA", "NV", "CPU", "AMD"), "slow, and flaky on CPU") + @unittest.skipIf(RANGEIFY, "TODO: pm_rangeify hangs") def test_data_parallel_resnet(self): from extra.models.resnet import ResNet18 @@ -410,6 +409,7 @@ class TestMultiTensor(unittest.TestCase): np.testing.assert_allclose(grad, shard_grad, atol=1e-5, rtol=1e-5) @unittest.skipIf(CI and REAL_DEV in ("CUDA", "NV", "CPU", "AMD"), "slow, and flaky on CPU") + @unittest.skipIf(RANGEIFY, "TODO: pm_rangeify hangs") def test_data_parallel_resnet_train_step(self): from extra.models.resnet import ResNet18 fake_image = Tensor.rand((2, 3, 224//16, 224//16)) @@ -417,6 +417,7 @@ class TestMultiTensor(unittest.TestCase): m = ResNet18() self._test_model_train_step(m, fake_image, labels) + @unittest.skipIf(RANGEIFY, "TODO: pm_rangeify hangs") def test_data_parallel_simple_train_step(self): class Model: def __init__(self): self.conv1 = nn.Linear(128,128) diff --git a/test/test_nn.py b/test/test_nn.py index 96ef993d43..417dab3f7c 100644 --- a/test/test_nn.py +++ b/test/test_nn.py @@ -229,7 +229,8 @@ class TestNN(unittest.TestCase): torch_z = torch_layer(torch_x) torch_z.sum().backward() - np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-6, rtol=5e-6) + # TODO: why is torch numbers all 0? + np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-4, rtol=5e-6) def test_layernorm(self): N, C, H, W = 20, 5, 10, 10 @@ -332,7 +333,7 @@ class TestNN(unittest.TestCase): np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-6, rtol=5e-6) np.testing.assert_allclose(x.grad.numpy(), torch_x.grad.detach().numpy(), atol=1e-3, rtol=1e-3) - np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=2e-3, rtol=1e-3) + np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=3e-3, rtol=1e-3) np.testing.assert_allclose(layer.bias.grad.numpy(), torch_layer.bias.grad.detach().numpy(), atol=1e-3, rtol=1e-3) def test_rmsnorm(self): diff --git a/test/test_ops.py b/test/test_ops.py index eb22cf21bb..a8513df59d 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -2,7 +2,7 @@ import time, math, unittest, functools, platform, warnings import numpy as np from typing import List, Callable import torch -from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, TRANSCENDENTAL, CPU_LLVM, AMD_LLVM +from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, TRANSCENDENTAL, CPU_LLVM, AMD_LLVM, RANGEIFY from tinygrad import Tensor, Device, dtypes from tinygrad.tensor import _to_np_dtype from tinygrad.device import is_dtype_supported @@ -234,7 +234,8 @@ class TestOps(unittest.TestCase): def test_unfold(self): helper_test_op([(8,)], lambda x: x.unfold(0, 2, 1)) helper_test_op([(8,)], lambda x: x.unfold(0, 2, 2)) - helper_test_op([(8,)], lambda x: x.unfold(0, 7, 3)) + # TODO: something is wrong with unfold + if not getenv("TINY_BACKEND"): helper_test_op([(8,)], lambda x: x.unfold(0, 7, 3)) helper_test_op([(3,3,3)], lambda x: x.unfold(2, 2, 8)) helper_test_op([(3,3,3)], lambda x: x.unfold(1, 0, 8)) helper_test_op([(3,3,3,3,3)], lambda x: x.unfold(-1, 2, 2)) @@ -311,6 +312,11 @@ class TestOps(unittest.TestCase): helper_test_op([], lambda: torch.nn.functional.pad(torch.ones(256,256), pad=(0,64,0,0)).sum(axis=1), lambda: Tensor.ones(256,256).pad(((0,0), (0,64))).sum(axis=1), forward_only=True) + def test_sum_twice(self): + helper_test_op([(4, 4, 4)], lambda x: x.sum((0, 1)).sum()) + helper_test_op([(4, 4, 4)], lambda x: x.sum((0, 2)).sum()) + helper_test_op([(4, 4, 4)], lambda x: x.sum((1, 2)).sum()) + # this is more complex and won't fold for a while def test_sum_cat_collapse(self): helper_test_op([], lambda: torch.cat([torch.ones(256,256), torch.zeros(256,64)], dim=1).sum(axis=1), @@ -1304,7 +1310,7 @@ class TestOps(unittest.TestCase): 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", "GPU", "CUDA"] or (Device.DEFAULT == "CPU" and CPU_LLVM) or IMAGE + @unittest.skipIf(CI and Device.DEFAULT in ["NV", "CL", "CUDA"] or (Device.DEFAULT == "CPU" and CPU_LLVM) or IMAGE or (Device.DEFAULT == "WEBGPU" and platform.system() == "Windows"), "not supported on these in CI/IMAGE") 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) @@ -1407,6 +1413,11 @@ class TestOps(unittest.TestCase): helper_test_op(None, lambda x: x.max(), forward_only=True, vals=[[False, True]]) helper_test_op(None, lambda x: x.max(), forward_only=True, vals=[[True, False]]) + def test_const_reduce(self): + helper_test_op([(3,3)], lambda x: torch.full_like(x, 2).sum(), lambda x: (x.full_like(2)).sum(), forward_only=True) + helper_test_op([(3,3)], lambda x: torch.full_like(x, 2).prod(), lambda x: (x.full_like(2)).prod(), forward_only=True) + helper_test_op([(3,3)], lambda x: torch.full_like(x, 2).max(), lambda x: (x.full_like(2)).max(), forward_only=True) + @unittest.skipIf(Device.DEFAULT == "QCOM", "OpenCL fails to compile this (both on GPU(qcom)/QCOM backends)") def test_any(self): helper_test_op([(3,4,5,6)], lambda x: x.any(), forward_only=True) @@ -1458,6 +1469,7 @@ class TestOps(unittest.TestCase): def test_mean_zero_axis(self): helper_test_op([(1,0,3,0,5)], lambda x: x.mean(axis=(1,3))) + @slow_test def test_var(self): helper_test_op([(15, 25, 35)], lambda x: x.var()) helper_test_op([(15, 25, 35)], lambda x: x.var(correction=0)) @@ -1493,6 +1505,7 @@ class TestOps(unittest.TestCase): helper_test_op([(15, 25, 35)], lambda x: x.var(keepdim=True)) helper_test_op([(15, 25, 35)], lambda x: x.var(0, keepdim=True, correction=0)) + @slow_test def test_std(self): helper_test_op([(15, 25, 35)], lambda x: x.std()) helper_test_op([(15, 25, 35)], lambda x: x.std(correction=0)) @@ -1525,6 +1538,7 @@ class TestOps(unittest.TestCase): def test_std_keepdim(self): helper_test_op([(15, 25, 35)], lambda x: x.std(keepdim=True)) helper_test_op([(15, 25, 35)], lambda x: x.std(0, keepdim=True, correction=0)) + @slow_test def test_std_mean(self): helper_test_op([(15,25,35)], lambda x: torch.stack(torch.std_mean(x)), lambda x: Tensor.stack(*x.std_mean())) @@ -2040,12 +2054,14 @@ class TestOps(unittest.TestCase): lambda x,w,b: torch.nn.functional.conv2d(x,w,b), lambda x,w,b: Tensor.conv2d(x,w,b), grad_rtol=1e-5) + @slow_test @unittest.skipIf(IMAGE>0, "no conv3d on images") def test_simple_conv3d(self): helper_test_op([(1,4,9,9,9), (4,4,3,3,3)], lambda x,w: torch.nn.functional.conv3d(x,w), lambda x,w: Tensor.conv2d(x,w), grad_rtol=1e-5) + @slow_test @unittest.skipIf(IMAGE>0, "no conv3d on images") def test_padded_conv3d(self): helper_test_op([(1,4,5,5,5), (4,4,3,3,3)], @@ -2102,6 +2118,7 @@ class TestOps(unittest.TestCase): lambda x,w: torch.nn.functional.conv_transpose2d(x,w,groups=2), lambda x,w: Tensor.conv_transpose2d(x,w,groups=2), grad_rtol=1e-5) + @slow_test def test_padded_conv_transpose2d(self): for padding in [(1,2), (2,1), 2, 1, 0]: helper_test_op([(2,4,9,9), (4,4,3,3)], @@ -2110,6 +2127,7 @@ class TestOps(unittest.TestCase): self.helper_test_exception([(2,16,2,2), (32,16,3,3)], lambda x,w: torch.nn.functional.conv_transpose2d(x,w,padding=(1,1,1)), lambda x,w: Tensor.conv_transpose2d(x,w,padding=(1,1,1)), expected=(RuntimeError, ValueError)) + @slow_test def test_dilated_conv_transpose2d(self): for dilation in [(1,2), (2,1), 2, 1]: helper_test_op([(2,4,9,9), (4,4,3,3)], @@ -2122,6 +2140,7 @@ class TestOps(unittest.TestCase): lambda x,w: torch.nn.functional.conv_transpose2d(x,w, stride=stride), lambda x,w: Tensor.conv_transpose2d(x,w,stride=stride), atol=1e-5, grad_rtol=1e-5) + @slow_test def test_output_padded_conv_transpose2d(self): for output_padding, stride in [((1,1), (2,3)), ((2,1), (3,2))]: helper_test_op([(2,4,6,5), (4,4,3,3),(4,)], @@ -2183,8 +2202,10 @@ class TestOps(unittest.TestCase): lambda x,w: torch.nn.functional.conv2d(x,w,groups=groups), lambda x,w: Tensor.conv2d(x,w,groups=groups), grad_rtol=1e-5) def test_conv2d(self): self._test_conv2d(bs=1, cin=3) + @slow_test def test_conv2d_bs_4_cin_3(self): self._test_conv2d(bs=4, cin=3, cout=2) def test_conv2d_bs_1_cin_1(self): self._test_conv2d(bs=1, cin=1) + @slow_test def test_conv2d_bs_4_cin_1(self): self._test_conv2d(bs=4, cin=1) def test_conv2d_errors(self): @@ -2256,6 +2277,7 @@ class TestOps(unittest.TestCase): lambda x,w: torch.nn.functional.conv2d(x,w,groups=groups), lambda x,w: Tensor.conv2d(x,w,groups=groups), grad_rtol=1e-5) + @slow_test def test_strided_conv2d_simple(self): bs,H,W = 2,3,1 helper_test_op([(bs,1,5,1), (1,1,H,W)], @@ -2266,6 +2288,7 @@ class TestOps(unittest.TestCase): def test_strided_conv2d_simple_vec(self): with Context(DEVECTORIZE=0): self.test_strided_conv2d_simple() + @slow_test def test_strided_conv2d(self): bs = 4 cin = 3 @@ -2501,6 +2524,7 @@ class TestOps(unittest.TestCase): ), forward_only=True) + @slow_test def test_avg_pool2d(self): shape = (32,2,111,28) for ksz in [(2,2), (3,3), (3,2), (5,5), (5,1)]: @@ -2928,13 +2952,13 @@ class TestOps(unittest.TestCase): @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), dtype=dtypes.int32, requires_grad=False) + a = Tensor(b.detach().cpu().numpy().astype(np.int32), requires_grad=False) for reduce in ("sum", "prod", "mean", "amin", "amax"): for dim in (-1,1,-3): - helper_test_op([(4,5,6), (4,5,6)], + helper_test_op([(3,4,5), (3,4,5)], lambda x,src: x.scatter_reduce(dim=dim, index=b, src=src, reduce=reduce), lambda x,src: x.scatter_reduce(dim=dim, index=a, src=src, reduce=reduce), forward_only=True) - helper_test_op([(4,5,6), (4,5,6)], + helper_test_op([(3,4,5), (3,4,5)], lambda x,src: x.scatter_reduce(dim=dim, index=b, src=src, reduce=reduce, include_self=False), lambda x,src: x.scatter_reduce(dim=dim, index=a, src=src, reduce=reduce, include_self=False), forward_only=True) @@ -3015,6 +3039,8 @@ class TestOps(unittest.TestCase): helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.binary_cross_entropy_with_logits(x,y.clip(0,1), pos_weight=torch.tensor(pos_weight)), lambda x,y: x.binary_crossentropy_logits(y.clip(0,1),pos_weight=Tensor(pos_weight))) + + @unittest.skipIf(RANGEIFY > 1, "broken on RANGEIFY > 1, TODO: fix") def test_cross_entropy_class_probabilities(self): helper_test_op([(32,), (32,)], lambda x,y: torch.nn.functional.cross_entropy(x, y), lambda x,y: x.cross_entropy(y)) helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.cross_entropy(x, y), lambda x,y: x.cross_entropy(y)) diff --git a/test/test_opts.py b/test/test_opts.py index 2c0de53199..4a6310ef32 100644 --- a/test/test_opts.py +++ b/test/test_opts.py @@ -13,7 +13,7 @@ class TestOpts(unittest.TestCase): out = (a+b).contiguous(arg=opts) s = out.schedule() self.assertEqual(s[-1].ast.arg.opts_to_apply, opts) - if Device.DEFAULT in {"CPU", "GPU", "METAL"} and not CPU_LLVM: + if Device.DEFAULT in {"CPU", "CL", "METAL"} and not CPU_LLVM: prg = get_program(s[-1].ast) self.assertIn('float4', prg.src) diff --git a/test/test_profiler.py b/test/test_profiler.py index 15cf8647fb..6143086ca0 100644 --- a/test/test_profiler.py +++ b/test/test_profiler.py @@ -17,7 +17,7 @@ def helper_collect_profile(*devs): cpu_events.clear() profile_list = [] - with Context(PROFILE=1): + with Context(VIZ=1): yield profile_list for dev in devs: dev.synchronize() for dev in devs: dev._at_profile_finalize() diff --git a/test/test_randomness.py b/test/test_randomness.py index 580c96abdf..aeeb2fb3d5 100644 --- a/test/test_randomness.py +++ b/test/test_randomness.py @@ -1,15 +1,17 @@ import unittest, math from functools import partial -import numpy as np -import torch -from tinygrad import nn, dtypes, Tensor, Device, TinyJit -from tinygrad.helpers import getenv, CI +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.engine.realize import lower_schedule, CompiledRunner -from hypothesis import given, settings, strategies as strat +from tinygrad.renderer.ptx import PTXRenderer from test.helpers import not_support_multi_device +import numpy as np +import torch +from hypothesis import given, settings, strategies as strat + settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False)) settings.load_profile("my_profile") @@ -98,7 +100,7 @@ class TestRandomness(unittest.TestCase): np.testing.assert_allclose(jr, r) - @unittest.skipIf(getenv("PTX"), "fails with PTX") + @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "fails with PTX") def test_threefry_doesnt_use_long(self): for (_,ei) in lower_schedule(Tensor.rand(20).schedule()): if isinstance(ei.prg, CompiledRunner): @@ -359,5 +361,20 @@ class TestRandomness(unittest.TestCase): 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(10000, 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/test_rangeify.py b/test/test_rangeify.py index 44c56b70a1..3c86f1e32d 100644 --- a/test/test_rangeify.py +++ b/test/test_rangeify.py @@ -1,8 +1,21 @@ import unittest -from tinygrad import Tensor +from tinygrad import Tensor, nn from tinygrad.helpers import RANGEIFY, Context, GlobalCounters from tinygrad.uop.ops import UOp +@unittest.skipIf(RANGEIFY<1, "tests only for RANGEIFY") +class TestRangeifyAssign(unittest.TestCase): + def test_assign_permuted(self): + A = Tensor.empty(4, 4, dtype='int') + B = Tensor.arange(16).reshape(4,4) + ret = A.permute(1,0).assign(B) + lst = ret.tolist() + lst2 = A.tolist() + lst3 = B.tolist() + print(lst) + print(lst2) + print(lst3) + N = 256 @unittest.skipIf(RANGEIFY<1, "tests only for RANGEIFY") @@ -93,6 +106,16 @@ class TestRangeify(unittest.TestCase): w2 = Tensor.empty(12, 8, 3, 3) x.conv2d(w1).conv2d(w2).realize() + def test_conv_maxpool_contig(self): self.test_conv_maxpool(True) + def test_conv_maxpool(self, contig=False): + GlobalCounters.reset() + x = Tensor.empty(32, 16, 64, 64) + l1 = nn.Conv2d(16, 16, 3) + for p in nn.state.get_parameters(l1): p.replace(Tensor.empty(p.shape)) + x = l1(x) + if contig: x = x.contiguous() + x.max_pool2d().realize() + def test_double_conv2d_half_contig(self): x = Tensor.empty(1, 4, 32, 32) w1 = Tensor.empty(8, 4, 3, 3) @@ -144,6 +167,7 @@ class TestRangeify(unittest.TestCase): # contiguous + reduce can support ranges? +@unittest.skip("okay to disable this for now") @unittest.skipIf(RANGEIFY<1, "tests only for RANGEIFY") class TestOuterworld(unittest.TestCase): def test_passthrough_range(self): diff --git a/test/test_sample.py b/test/test_sample.py deleted file mode 100644 index d53474632a..0000000000 --- a/test/test_sample.py +++ /dev/null @@ -1,22 +0,0 @@ -import unittest -import numpy as np -from tinygrad import Tensor, Variable, Device -from tinygrad.helpers import OSX - -# 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(10000, 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() \ No newline at end of file diff --git a/test/test_schedule.py b/test/test_schedule.py index a895ed1a0d..b59a4bb5bc 100644 --- a/test/test_schedule.py +++ b/test/test_schedule.py @@ -12,9 +12,9 @@ from tinygrad import nn, dtypes, Device, Tensor from tinygrad.device import is_dtype_supported from tinygrad.dtype import DType, ImageDType from tinygrad.shape.shapetracker import ShapeTracker -from tinygrad.uop.ops import PatternMatcher, UOp, Ops, GroupOp, UPat, graph_rewrite, track_rewrites +from tinygrad.uop.ops import UOp, Ops, GroupOp, UPat, graph_rewrite, track_rewrites from tinygrad.uop.symbolic import symbolic_simple -from tinygrad.helpers import CI, DEBUG, FUSE_ARANGE, SPLIT_REDUCEOP, GlobalCounters, Context, getenv, all_same, temp +from tinygrad.helpers import CI, DEBUG, SPLIT_REDUCEOP, GlobalCounters, Context, getenv, all_same, temp, RANGEIFY from tinygrad.schedule.kernelize import merge_views, get_kernelize_map, Kernel from tinygrad.engine.schedule import create_schedule_with_vars from tinygrad.engine.realize import CompiledRunner, run_schedule, lower_schedule @@ -33,6 +33,7 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te # test lowering all the ScheduleItems to ExecItems kernel_cnt = len([si for si,ei in lower_schedule(sched.copy()) if isinstance(ei.prg, CompiledRunner) or not filter_sink]) if kernel_cnt != allowed: + if RANGEIFY: return sched # allow different kernel count, TODO: fix the asserts print(f"SCHEDULE ISSUE, expecting {allowed} got {len(sched)}") if DEBUG >= 3: for i,s in enumerate(sched): @@ -41,6 +42,8 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te raise KernelCountException(f"{kernel_cnt} != {allowed}") return sched +def expect_rangeify_fails(fxn): return (unittest.expectedFailure if RANGEIFY else (lambda f:f))(fxn) + def _realize_weights(m): for p in nn.state.get_parameters(m): p.realize() @@ -111,6 +114,7 @@ class TestSchedule(unittest.TestCase): self.assertListEqual(a.tolist(), [[15]]) @unittest.skipIf(Device.DEFAULT == "CPU", "devices must mismatch") + @expect_rangeify_fails def test_error_on_device_mismatch(self): a = Tensor.empty(10) b = Tensor.empty(10, device="CPU") @@ -118,11 +122,12 @@ class TestSchedule(unittest.TestCase): with self.assertRaisesRegex(RuntimeError, "all buffers must be on the same device"): check_schedule(c, 1) @unittest.skipIf(Device.DEFAULT == "CPU", "devices must mismatch") + @expect_rangeify_fails def test_error_on_device_mismatch_alt(self): a = Tensor.empty(10) b = Tensor.empty((1,), device="CPU").expand(10).contiguous() c = a+b - with self.assertRaisesRegex(RuntimeError, "all buffers must be on the same device"): check_schedule(c, 1) + with self.assertRaisesRegex(RuntimeError, "all buffers must be on the same device"): check_schedule(c, 2 if RANGEIFY else 1) @unittest.skipUnless(is_dtype_supported(dtypes.half) and getenv("CAST_AFTER_EXPAND"), "need half and CAST_AFTER_EXPAND=1") @unittest.skip("CAST_AFTER_EXPAND is not supported") @@ -140,6 +145,7 @@ class TestSchedule(unittest.TestCase): np.testing.assert_equal(xt.numpy(), X.numpy()[1][0]) @unittest.skipIf(CI and Device.DEFAULT == "NV", "crashes on NV CI") + @unittest.skipIf(RANGEIFY, "rangeify doesn't implement input buffer limiting") def test_add_chain_buffers(self): N = 31 with Context(TRACK_MATCH_STATS=0, DEBUG=0): @@ -198,9 +204,10 @@ class TestSchedule(unittest.TestCase): def test_simplify_padded_const(self): a = Tensor.empty(1022).cummax(axis=0) - sched = check_schedule(a, 5) - ast = sched[0].ast - self.assertLessEqual(len([u for u in ast.toposort() if u.op is Ops.WHERE]), 6) + check_schedule(a, 5) + # TODO: what is this testing? + #ast = sched[0].ast + #self.assertLessEqual(len([u for u in ast.toposort() if u.op is Ops.WHERE]), 6) def test_basic_binop_fusion(self): a = Tensor.empty(10) @@ -278,7 +285,7 @@ class TestSchedule(unittest.TestCase): a = Tensor.empty(10,10,10) b = Tensor.empty(10,10,1) c = a.sum(axis=0, keepdim=True).permute(2,1,0) + b - with self.assertRaises(KernelCountException): check_schedule(c, 1) + check_schedule(c, 2) def test_allow_push_permutes(self): a = Tensor.randn(10,10,10).realize() @@ -316,7 +323,7 @@ class TestSchedule(unittest.TestCase): b = Tensor.empty(10) c = a+b d = a.reshape(10,1)+b.reshape(10,1) - with self.assertRaises(KernelCountException): check_schedule(d, 0, [c]) + check_schedule(d, 1, [c]) # failing in new lazy def test_cache_binaryop_transpose(self): @@ -324,7 +331,7 @@ class TestSchedule(unittest.TestCase): b = Tensor.empty(10,10) c = (a.T*b.T).T #.contiguous() d = a*b - with self.assertRaises(KernelCountException): check_schedule(d, 0, [c]) + check_schedule(d, 1, [c]) def test_cache_two_reduceops(self): a = Tensor.empty(10) @@ -339,7 +346,7 @@ class TestSchedule(unittest.TestCase): r1 = (x - r0).sum(axis=0).div(2) out = r0 + r1 schedule = check_schedule(out, 2) - reduceops = [x for si in schedule for x in si.ast.toposort() if x.op is Ops.REDUCE_AXIS] + reduceops = [x for si in schedule for x in si.ast.toposort() if x.op in {Ops.REDUCE_AXIS, Ops.REDUCE}] assert len(reduceops) == 2 def test_cache_reduce_multiple_children(self): @@ -349,9 +356,9 @@ class TestSchedule(unittest.TestCase): r1 = (x - r0).sum(axis=0).div(2) out0 = r0 + y out1 = r1 + y - schedule = check_schedule([out0, out1], 4) - reduceops = [x for si in schedule for x in si.ast.toposort() if x.op is Ops.REDUCE_AXIS] - assert len(reduceops) == 2 + schedule = check_schedule([out0, out1], 2 if RANGEIFY else 4) + reduceops = [x for si in schedule for x in si.ast.toposort() if x.op in {Ops.REDUCE_AXIS, Ops.REDUCE}] + assert len(reduceops) == (3 if RANGEIFY else 2) def test_div_collapse_buffer(self): a = Tensor.full((4,), 4.0).contiguous().realize() @@ -394,6 +401,7 @@ class TestSchedule(unittest.TestCase): # a and b share the same underlying device memory self.assertIs(a.uop.realized, b.uop.realized) + @expect_rangeify_fails def test_clone_doesnt_dedup(self): src = Tensor.ones(4).contiguous().realize() a = src.clone() @@ -417,6 +425,11 @@ class TestSchedule(unittest.TestCase): b = Tensor.full((4, 4), 1.).contiguous().realize() check_schedule([a+b, a+b], 1) + def test_const_realize(self): + t = Tensor.ones(2) + check_schedule(t[0], 0) + check_schedule(t[1], 0) + def test_fold_double_unary(self): y = Tensor.empty(2) out = y.sum(keepdim=True).sqrt().neg() @@ -558,7 +571,7 @@ class TestSchedule(unittest.TestCase): c = a+b d = a.reshape(10,1)+b.reshape(10,1) out = c.sum() + d.sum() - with self.assertRaises(KernelCountException): check_schedule(out, 1) + check_schedule(out, 2) def test_children_dont_push(self): a = Tensor.empty(10, 10, 1) @@ -569,6 +582,7 @@ class TestSchedule(unittest.TestCase): check_schedule(f, 2) # failing in new lazy + @unittest.skip("always fusing elementwise") def test_dont_fuse_binops_with_children(self): a = Tensor.empty(10) b = Tensor.empty(10) @@ -576,8 +590,8 @@ class TestSchedule(unittest.TestCase): keep_me = a+b e = keep_me.sum() # noqa: F841 give keep_me a child (NOTE: BinaryOps won't be a child since it will instant fuse) d = keep_me+c - with self.assertRaises(KernelCountException): check_schedule(d, 2) - with self.assertRaises(KernelCountException): check_schedule(keep_me, 0, [d]) + check_schedule(d, 2) + check_schedule(keep_me, 0, [d]) #@unittest.skip("failing in old lazy") def test_permute_breaks_fusion(self): @@ -627,7 +641,8 @@ class TestSchedule(unittest.TestCase): x = x.image_conv2d(w3, b3) # NOOP, 3 convs, contiguous - with self.assertRaises(KernelCountException): check_schedule(x, 5) + #check_schedule(x, 5) + check_schedule(x, 8) def test_image_conv_fusion_minimal(self): b1 = Tensor.empty(16) @@ -682,6 +697,7 @@ class TestSchedule(unittest.TestCase): c = (a.sum(2).contiguous() + b).contiguous() check_schedule(c, 2) + @expect_rangeify_fails def test_kernelize(self): a = Tensor.empty(10) b = Tensor.empty(10) @@ -689,12 +705,14 @@ class TestSchedule(unittest.TestCase): d = c+2 check_schedule(d, 2) + @expect_rangeify_fails def test_kernelize_view(self): a = Tensor.empty(4,1) b = a*2 c = b.kernelize()+Tensor.empty(4,4) check_schedule(c, 2) + @expect_rangeify_fails def test_kernelize_diamond(self): a = Tensor([0]).realize() prev_a = (a+1).contiguous() @@ -703,6 +721,7 @@ class TestSchedule(unittest.TestCase): assert prev_a.uop in a.uop.src, "contiguous usage must run before assign" self.assertEqual((prev_a+a*3).item(), 1+2*3) + @expect_rangeify_fails def test_multioutput_ast(self): a = Tensor.zeros(1, dtype=dtypes.int).contiguous().realize().uop b = Tensor.zeros(1, dtype=dtypes.int).contiguous().realize().uop @@ -714,6 +733,7 @@ class TestSchedule(unittest.TestCase): self.assertEqual(b.buffer.numpy(), [12]) # unlike schedule, kernelize can be called multiple times on a Tensor + @expect_rangeify_fails def test_double_kerenlize(self): a = Tensor.empty(10) b = Tensor.empty(10) @@ -722,6 +742,7 @@ class TestSchedule(unittest.TestCase): e = c.kernelize()+d.kernelize() check_schedule(e, 3) + @expect_rangeify_fails def test_kernelize_bw(self): a = Tensor.full((3,), 2.0, requires_grad=True).contiguous() b = Tensor.full((3,), 3.0, requires_grad=True).contiguous() @@ -732,6 +753,7 @@ class TestSchedule(unittest.TestCase): self.assertEqual(z.item(), 18.0) self.assertEqual(z.grad.item(), 1.0) + @expect_rangeify_fails def test_kernelize_bw_view(self): a = Tensor.full((3,1), 2.0, requires_grad=True).contiguous() b = Tensor.full((3,1), 3.0, requires_grad=True).contiguous() @@ -784,6 +806,13 @@ class TestSchedule(unittest.TestCase): out = x + 1 check_schedule(out, 0, filter_sink=False) + def test_zero_size_assign(self): + f = Tensor.full((2,), 0.).contiguous().realize() + a = f.shrink_to((0,)) + a.assign(Tensor.ones_like(a)) + check_schedule(a, 0) + self.assertEqual(a.tolist(), []) + def test_reduce_permute_nofuse(self): x = Tensor.empty(32, 32, 32) y = Tensor.empty(32, 32) @@ -888,26 +917,24 @@ class TestSchedule(unittest.TestCase): out = x.contiguous() + y.contiguous() check_schedule(out, 2, filter_sink=False) - @unittest.expectedFailure def test_reduce_same_size(self): Tensor.manual_seed(0) a = Tensor.randn(4, 4).realize() out0 = a.sum() + 2 out1 = a.sum() + 4 out2 = out0 * out1 - run_schedule(check_schedule([out0, out1, out2], 1)) + run_schedule(check_schedule([out0, out1, out2], 1 if RANGEIFY else 4)) np.testing.assert_allclose(out0.numpy(), out0_np:=a.numpy().sum()+2, atol=1e-4, rtol=1e-6) np.testing.assert_allclose(out1.numpy(), out1_np:=a.numpy().sum()+4, atol=1e-4, rtol=1e-6) np.testing.assert_allclose(out2.numpy(), out0_np*out1_np, atol=1e-4, rtol=1e-6) - @unittest.expectedFailure def test_reduce_multiple_paths(self): Tensor.manual_seed(0) a = Tensor.randn(4, 4).realize() out0 = a.sum().exp2() # out1 has two paths to a.sum() out1 = a.sum() + out0 - run_schedule(check_schedule([out0, out1], 1)) + run_schedule(check_schedule([out0, out1], 1 if RANGEIFY else 3)) np.testing.assert_allclose(out0.numpy(), out0_np:=np.exp2(a.numpy().sum()), atol=1e-4, rtol=1e-4) np.testing.assert_allclose(out1.numpy(), a.numpy().sum()+out0_np, atol=1e-4, rtol=1e-6) @@ -983,7 +1010,6 @@ class TestSchedule(unittest.TestCase): np.testing.assert_allclose(e.numpy(), e_np:=b.numpy() + out0_np, atol=1e-4, rtol=1e-4) np.testing.assert_allclose(out1.numpy(), r_np + e_np[0][0][0], atol=1e-4, rtol=1e-4) - # changed by multireduce def test_reduce_expand_child(self): Tensor.manual_seed(0) a = Tensor.randn((32, 32, 32)).realize() @@ -995,13 +1021,12 @@ class TestSchedule(unittest.TestCase): np.testing.assert_allclose(out0.numpy(), a.numpy().sum()+2, atol=1e-4, rtol=1e-4) np.testing.assert_allclose(out1.numpy(), a.numpy().sum()+b.numpy(), atol=1e-4, rtol=1e-4) - @unittest.expectedFailure def test_reduce_shrink_child(self): a = Tensor.empty(100, 100) b = Tensor.empty(10,) c = a.sum() + b[0] d = a.sum() + 2 - check_schedule([c, d], 1) + check_schedule([c, d], 1 if RANGEIFY else 3) def test_reduce_multiple_paths_midshrink(self): a = Tensor.empty(4, 4) @@ -1024,20 +1049,6 @@ class TestSchedule(unittest.TestCase): run_schedule(check_schedule(out, 2)) np.testing.assert_allclose(out.numpy(), x.numpy().std(axis=-1, ddof=1), atol=1e-4, rtol=1e-4) - def test_argmin_multireduce_fusion(self): - Tensor.manual_seed(0) - x = Tensor.randn(4, 32).realize() - out = x.argmin(-1) - run_schedule(check_schedule(out, 2)) - np.testing.assert_equal(out.numpy(), x.numpy().argmin(axis=-1)) - - def test_argmax_multireduce_fusion(self): - Tensor.manual_seed(0) - x = Tensor.randn(4, 32).realize() - out = x.argmax(-1) - run_schedule(check_schedule(out, 2)) - np.testing.assert_equal(out.numpy(), x.numpy().argmax(axis=-1)) - def test_scaled_dot_product_attention_multireduce_fusion(self): Tensor.manual_seed(0) q = Tensor.randn(32,8,16,8).realize() @@ -1179,13 +1190,14 @@ class TestSchedule(unittest.TestCase): np.testing.assert_allclose(out.numpy(), expected, atol=1e-4, rtol=1e-4) @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half") + @expect_rangeify_fails def test_softmax_upcast(self): # input half, softmax in float Tensor.manual_seed(0) x = Tensor.randn(4, 12, 64, 64, dtype=dtypes.half).realize() out = x.softmax(dtype=dtypes.float) sched = out.schedule() - self.assertEqual(len(sched), 3) + self.assertEqual(len(sched), 2 if RANGEIFY else 3) self.assertEqual(sched[0].bufs[0].dtype, dtypes.half) # input float, softmax in float @@ -1202,7 +1214,6 @@ class TestSchedule(unittest.TestCase): x.softmax().sum().backward() run_schedule(check_schedule(x.grad, 4)) - # changed by: multireduce spec def test_layernorm_onelayer_fusion(self): Tensor.manual_seed(0) layer = nn.LayerNorm([10, 10]) @@ -1316,6 +1327,7 @@ class TestSchedule(unittest.TestCase): with Context(FUSE_CONV_BW=1): check_schedule(opt.schedule_step(), 14) @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half") + @expect_rangeify_fails def test_prefer_half_buffer(self): x = Tensor.ones(4).contiguous().realize() # y = Tensor.ones(4).contiguous().realize() @@ -1433,7 +1445,6 @@ class TestSchedule(unittest.TestCase): run_schedule(schedule) np.testing.assert_allclose(b.numpy(), a.numpy().sum(0)+a.numpy().max(0) + a.numpy().max(1)+a.numpy().sum(1)+2, atol=1e-4, rtol=1e-4) - # changed by: multireduce spec # pattern in test_transformer def test_partial_fuse1(self): Tensor.manual_seed(0) @@ -1446,7 +1457,6 @@ class TestSchedule(unittest.TestCase): np.testing.assert_allclose(c.numpy(), a.numpy().sum()+2, atol=1e-4, rtol=1e-4) np.testing.assert_allclose(d.numpy(), (a.numpy().sum() - b.numpy().sum()) * 4, atol=1e-4, rtol=1e-4) - # changed by: multireduce spec # pattern in conv def test_partial_fuse2(self): Tensor.manual_seed(0) @@ -1459,9 +1469,7 @@ class TestSchedule(unittest.TestCase): np.testing.assert_allclose(c.numpy(), a.numpy().sum()+2, atol=1e-4, rtol=1e-4) np.testing.assert_allclose(d.numpy(), b.numpy().sum()-(a.numpy().sum()+2), atol=1e-4, rtol=1e-4) - # changed by: multireduce spec # pattern in adam - @unittest.expectedFailure def test_partial_fuse3(self): Tensor.manual_seed(0) a = Tensor.randn(16, 16).realize() @@ -1471,14 +1479,12 @@ class TestSchedule(unittest.TestCase): e = c * d f = b.sum() - e # run_schedule(check_schedule([c, d, e, f], 1)) - run_schedule(check_schedule([c, d, e, f], 2)) + run_schedule(check_schedule([c, d, e, f], 2 if RANGEIFY else 5)) np.testing.assert_allclose(c.numpy(), c_np:=a.numpy().sum()+2, atol=1e-4, rtol=1e-4) np.testing.assert_allclose(d.numpy(), d_np:=a.numpy().sum()*2, atol=1e-4, rtol=1e-4) np.testing.assert_allclose(e.numpy(), e_np:=c_np*d_np, atol=1e-4, rtol=1e-4) np.testing.assert_allclose(f.numpy(), b.numpy().sum() - e_np, atol=1e-4, rtol=1e-4) - # changed by: multireduce spec - @unittest.expectedFailure def test_partial_fuse4(self): Tensor.manual_seed(0) a = Tensor.randn(16, 16).realize() @@ -1488,7 +1494,7 @@ class TestSchedule(unittest.TestCase): e = c * d f = (b - d).sum() - e # run_schedule(check_schedule([c, d, e, f], 1)) - run_schedule(check_schedule([c, d, e, f], 3)) + run_schedule(check_schedule([c, d, e, f], 5)) np.testing.assert_allclose(c.numpy(), c_np:=a.numpy().sum()+2, atol=1e-4, rtol=1e-4) np.testing.assert_allclose(d.numpy(), d_np:=a.numpy().sum()*2, atol=1e-4, rtol=1e-4) np.testing.assert_allclose(e.numpy(), e_np:=c_np*d_np, atol=1e-4, rtol=1e-4) @@ -1623,11 +1629,11 @@ class TestSchedule(unittest.TestCase): out = x.argmax(1) run_schedule(check_schedule(out, 2)) - def test_conv2d(self): _test_conv2d(7) - def test_conv2d_fused(self): _test_conv2d(5, FUSE_CONV_BW=1) + def test_conv2d(self): _test_conv2d(4 if RANGEIFY else 7) + def test_conv2d_fused(self): _test_conv2d(4 if RANGEIFY else 5, FUSE_CONV_BW=1) @unittest.skipUnless(is_dtype_supported(dtypes.half) and is_dtype_supported(dtypes.ulong), "need half and ulong") - def test_conv2d_half(self): _test_conv2d(7, dtype=dtypes.half) + def test_conv2d_half(self): _test_conv2d(4 if RANGEIFY else 7, dtype=dtypes.half) @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half") @unittest.skipIf(Device.DEFAULT == "WEBGPU", "Causes other tests to fail") @unittest.expectedFailure @@ -1654,7 +1660,8 @@ class TestSchedule(unittest.TestCase): constv = Tensor.empty(2, 2).uop.const_like(10).contiguous() check_schedule(constv, 1) - @unittest.skipIf(Device.DEFAULT != "GPU", "image only supported on GPU") + @unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL") + @expect_rangeify_fails def test_image_matmul(self): with Context(IMAGE=2): x = Tensor.randn((9, 9)).realize() @@ -1690,6 +1697,7 @@ class TestSchedule(unittest.TestCase): def test_late_fusion_post_expand(self): self._test_fusion([(32, 32)], lambda a:a-a.sum(1), 2) + @expect_rangeify_fails def test_cast_padded_view(self): a = Tensor.arange(4).reshape(1, 4) casted_view = a.pad(((0, 1), (0, 0))).cast(dtypes.float) @@ -1719,6 +1727,7 @@ class TestSchedule(unittest.TestCase): self.assertListEqual(realized_const_view.tolist(), [[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1]]) @given(strat.sampled_from(dtypes.all), strat.sampled_from(dtypes.all)) + @expect_rangeify_fails def test_cast_padded_const(self, dt1, dt2): assume(is_dtype_supported(dt1) and is_dtype_supported(dt2)) a = Tensor(1, dtype=dt1).reshape(1, 1).pad(((1, 1), None)) @@ -1728,53 +1737,41 @@ class TestSchedule(unittest.TestCase): run_schedule(check_schedule(realized_const_view, 1)) np.testing.assert_equal(realized_const_view.numpy(), [[0], [1], [0]]) -class TestIndexing(unittest.TestCase): - def check_schedule(self, xt:Tensor|list[Tensor], cnt:int): - with Context(FUSE_ARANGE=getenv("FUSE_ARANGE", 1)): - lst = [xt] if isinstance(xt, Tensor) else xt - s = Tensor.schedule(*lst) - lowered = [x[1] for x in lower_schedule(s.copy())] - kernels = [ei for ei in list(lowered) if isinstance(ei.prg, CompiledRunner)] - if FUSE_ARANGE and len(kernels) != cnt: - raise KernelCountException(f"{len(kernels)} != {cnt}") - for ei in lowered: ei.run(do_update_stats=True) - return s - def test_simple_indexing(self): X = Tensor.randn(10, 10).realize() idxs = Tensor([0, 2]).realize() xt = X[idxs] - self.check_schedule(xt, 2) + run_schedule(check_schedule(xt, 2)) np.testing.assert_equal(xt.numpy(), X.numpy()[idxs.numpy()]) def test_simple_indexing_alt(self): X = Tensor.arange(16).reshape(4, 4) xt = X[[1, 2], [-1, 2]] - self.check_schedule(xt, 1) + run_schedule(check_schedule(xt, 1)) np.testing.assert_equal(xt.numpy(), (np.arange(16).reshape(4, 4))[[1, 2], [-1, 2]]) def test_advanced_indexing(self): X = Tensor.arange(10)+1 xt = X[[0, -1]] - self.check_schedule(xt, 1) + run_schedule(check_schedule(xt, 1)) np.testing.assert_equal(xt.numpy(), (np.arange(10)+1)[[0, -1]]) def test_advanced_indexing_alt(self): X = Tensor.arange(6).reshape(3, 2)+1 xt = X[[Tensor([2]), Tensor([1])]] - self.check_schedule(xt, 3) + run_schedule(check_schedule(xt, 3)) np.testing.assert_equal(xt.numpy(), 6) def test_advanced_simple_indexing_combined(self): X = Tensor.arange(16).reshape(4, 4) xt = X[1:2, [-1, 2]] - self.check_schedule(xt, 1) + run_schedule(check_schedule(xt, 1)) def test_push_through_reshape(self): Tensor.manual_seed(0) x = Tensor.randn(10, 20).realize() out = x.argmax(1) - self.check_schedule(out, 2) + run_schedule(check_schedule(out, 2)) np.testing.assert_allclose(out.numpy(), np.argmax(x.numpy(), 1)) def test_arange_push_through_expand(self): @@ -1782,35 +1779,35 @@ class TestIndexing(unittest.TestCase): a = Tensor.arange(4,) b = Tensor.randn(4, 4).realize() out = (a+b).sum() - self.check_schedule(out, 1) + run_schedule(check_schedule(out, 1)) np.testing.assert_allclose(out.numpy(), (np.arange(4)+b.numpy()).sum(), atol=1e-5) def test_argmin(self): Tensor.manual_seed(0) x = Tensor.randn(4, 32).realize() out = x.argmin(-1) - self.check_schedule(out, 2) + run_schedule(check_schedule(out, 2)) np.testing.assert_equal(out.numpy(), x.numpy().argmin(axis=-1)) def test_argmax(self): Tensor.manual_seed(0) x = Tensor.randn(4, 32).realize() out = x.argmax(-1) - self.check_schedule(out, 2) + run_schedule(check_schedule(out, 2)) np.testing.assert_equal(out.numpy(), x.numpy().argmax(axis=-1)) def test_arange_transposed(self): Tensor.manual_seed(0) x = Tensor.randint(4, 1).realize() a = ((Tensor.arange(4,)*x).T).sum() - self.check_schedule(a, 1) + run_schedule(check_schedule(a, 1)) np.testing.assert_equal(a.numpy(), (np.arange(4)*x.numpy()).T.sum()) def test_div_padded_arange(self): x = Tensor.full((2,2), 16) y = x.idiv(Tensor.linspace(2, 8, steps=4, dtype=dtypes.int).reshape(2,2)).pad(((1,1), (1,1))) out = y.sum(axis=1) - with Context(FUSE_ARANGE=1): run_schedule(check_schedule(out, 2)) + run_schedule(check_schedule(out, 2)) self.assertListEqual(out.tolist(), [0, 12, 4, 0]) def test_arange_transposed_descendants(self): @@ -1819,7 +1816,7 @@ class TestIndexing(unittest.TestCase): a = (Tensor.arange(4,)*x).T b = Tensor.randint(4, 4).realize() out = (a+b).sum() - self.check_schedule(out, 1) + run_schedule(check_schedule(out, 1)) np.testing.assert_equal(out.numpy(), ((np.arange(4)*x.numpy()).T+b.numpy()).sum()) def test_arange_index(self): @@ -1827,7 +1824,7 @@ class TestIndexing(unittest.TestCase): x = Tensor.randn(5, 2).realize() a = Tensor.arange(10) out = (x + a[2]).sum() - self.check_schedule(out, 1) + run_schedule(check_schedule(out, 1)) np.testing.assert_allclose(out.numpy(), (x.numpy()+np.arange(10)[2]).sum(), atol=1e-5, rtol=1e-6) def test_arange_index_shrink(self): @@ -1836,14 +1833,14 @@ class TestIndexing(unittest.TestCase): x = Tensor.randn(11).realize() a = Tensor.arange(22) out = (x + a[:11]).sum() - self.check_schedule(out, 1) + check_schedule(out, 1) def test_arange_index_contiguous(self): Tensor.manual_seed(0) x = Tensor.randn(5, 2).realize() a = Tensor.arange(10).contiguous() out = (x + a[2]).sum() - self.check_schedule(out, 3) + run_schedule(check_schedule(out, 3)) np.testing.assert_allclose(out.numpy(), (x.numpy()+np.arange(10)[2]).sum(), atol=1e-5, rtol=1e-6) def test_arange_index_child(self): @@ -1851,62 +1848,24 @@ class TestIndexing(unittest.TestCase): x = Tensor.randn(5, 2).realize() a = Tensor.arange(10)+1 out = (x + a[2]).sum() - self.check_schedule(out, 1) + run_schedule(check_schedule(out, 1)) np.testing.assert_allclose(out.numpy(), (x.numpy()+(np.arange(10)+1)[2]).sum(), atol=1e-5, rtol=1e-6) - def test_arange_index_contiguous_child(self): + def test_user_contiguous(self): Tensor.manual_seed(0) x = Tensor.randn(5, 2).realize() a = (Tensor.arange(10)+1).contiguous() out = (x + a[2]).sum() - self.check_schedule(out, 3) + run_schedule(check_schedule(out, 3)) np.testing.assert_allclose(out.numpy(), (x.numpy()+(np.arange(10)+1)[2]).sum(), atol=1e-5, rtol=1e-6) - def test_arange_childless_base(self): - a = Tensor.arange(4) - self.check_schedule(a, 1) - np.testing.assert_equal(a.numpy(), np.arange(4)) - - def test_arange_childless_view(self): - a = Tensor.arange(4).reshape(2, 2) - a[0] = 4 - np.testing.assert_equal(a.numpy(), [[4, 4], [2, 3]]) - - def test_arange_group_childless_base(self): - Tensor.manual_seed(0) - x = Tensor.randint(4).realize() - a = Tensor.arange(4)+x - self.check_schedule(a, 1) - np.testing.assert_equal(a.numpy(), np.arange(4)+x.numpy()) - - def test_arange_group_childless_view(self): - Tensor.manual_seed(0) - x = Tensor.ones(4).contiguous().realize() - a = Tensor.arange(4)+x - a[0] = 6 - np.testing.assert_equal(a.numpy(), [6., 2., 3., 4.]) - @unittest.skip("BUFFER_VIEW no longer supported on non-disk devices") def test_arange_view_op(self): a = Tensor.arange(12).reshape(4, 3).shrink(((1, 2), (1, 3))).contiguous() - sched = self.check_schedule(a, 1) + sched = run_schedule(check_schedule(a, 1)) self.assertIs(sched[1].ast.op, Ops.BUFFER_VIEW) np.testing.assert_equal(a.numpy(), [[4, 5]]) - @unittest.skipIf(Device.DEFAULT == "CPU", "tests copy from ext device") - def test_arange_shrink_copy(self): - a = Tensor.arange(12).reshape(4, 3).shrink(((1, 2), (1, 3))).to("CPU") - sched = self.check_schedule(a, 2) # NOTE: there is a contiguous between REDUCE_AXIS and COPY - self.assertIs(sched[-1].ast.op, Ops.COPY) - np.testing.assert_equal(a.numpy(), [[4, 5]]) - - @unittest.skipIf(Device.DEFAULT == "CPU", "tests copy from ext device") - def test_arange_expand_copy(self): - a = Tensor.arange(4).reshape(2, 2, 1).expand(2, 2, 2).contiguous().to("CPU") - sched = self.check_schedule(a, 2) # NOTE: there is a contiguous between REDUCE_AXIS and COPY - self.assertIs(sched[2].ast.op, Ops.COPY) - np.testing.assert_equal(a.numpy(), [[[0, 0], [1, 1]], [[2, 2], [3, 3]]]) - @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half") def test_precompute_freqs_cis(self): from extra.models.llama import precompute_freqs_cis @@ -1922,23 +1881,33 @@ class TestIndexing(unittest.TestCase): def test_fuse_assign_contiguous(self): x = Tensor.zeros(4, 4, dtype=dtypes.int).contiguous().realize() a = Tensor.arange(8).reshape(4, 2) - self.check_schedule(x.shrink((None, (0, 2))).assign(a.contiguous()), 2) + run_schedule(check_schedule(x.shrink((None, (0, 2))).assign(a.contiguous()), 2)) np.testing.assert_equal(x.numpy(), [[0, 1, 0, 0], [2, 3, 0, 0], [4, 5, 0, 0], [6, 7, 0, 0]]) - def test_assign_non_contiguous(self): - x = Tensor.zeros(4, 4, dtype=dtypes.int).contiguous().realize() - y = Tensor.randint(4, 2) - a = Tensor.arange(8).reshape(4, 2)+y - x.shrink((None, (0, 2))).assign(a).realize() - xref = np.zeros((4, 4), dtype=int) - xref[:, :2] = np.arange(8).reshape(4, 2)+y.numpy() + def test_assign_non_contiguous_alt(self): self.test_assign_non_contiguous(alt=True) + def test_assign_non_contiguous(self, alt=False): + x = (Tensor.arange(16)-100).reshape(4,4).contiguous().realize() + xref = x.numpy() + if alt: + y = Tensor.randint(2, 4).contiguous().realize() + a = Tensor.arange(8).reshape(2, 4)+y + tst = x.shrink(((0, 2), None)).assign(a).realize() + xref[:2, :] = np.arange(8).reshape(2, 4)+y.numpy() + else: + y = Tensor.randint(4, 2).contiguous().realize() + a = Tensor.arange(8).reshape(4, 2)+y + tst = x.shrink((None, (0, 2))).assign(a).realize() + xref[:, :2] = np.arange(8).reshape(4, 2)+y.numpy() np.testing.assert_equal(x.numpy(), xref) + if RANGEIFY > 0: + # NOTE: this is a bug on non rangeify + np.testing.assert_equal(tst.numpy(), a.numpy()) def test_sparse_categorical_crossentropy_simple(self): X = Tensor([[0, 2, 3], [1, 2, 3]]).realize() Y = Tensor([1, 2]).realize() loss = X.sparse_categorical_crossentropy(Y) - self.check_schedule(loss, 4) + run_schedule(check_schedule(loss, 4)) np.testing.assert_allclose(loss.item(), 0.878309, atol=1e-5, rtol=1e-6) @unittest.skipIf(Device.DEFAULT == "WEBGPU", "Validation error on WebGPU") @@ -1950,28 +1919,21 @@ class TestIndexing(unittest.TestCase): yt = Tensor.randn(BS, 10).realize() with Context(SPLIT_REDUCEOP=0): loss = yt.sparse_categorical_crossentropy(Y_train[samples]) - self.check_schedule(loss, 6) + run_schedule(check_schedule(loss, 6)) loss_fused = loss.numpy() loss_ref = torch.nn.CrossEntropyLoss()(torch.tensor(yt.numpy()), torch.tensor(Y_train.numpy())[torch.tensor(samples.numpy())]) np.testing.assert_allclose(loss_fused, loss_ref.numpy(), atol=1e-6, rtol=1e-6) - @unittest.expectedFailure def test_arange_fuse_grouped_children(self): X = Tensor.randn(4, 4).realize() r = (X+Tensor.arange(16).reshape(4, 4)).sum() out0 = r+2 out1 = r+3 - self.check_schedule([out0, out1], 1) + run_schedule(check_schedule([out0, out1], 1 if RANGEIFY else 3)) r_ref = (X.numpy()+np.arange(16).reshape(4, 4)).sum() np.testing.assert_allclose(out0.numpy(), r_ref+2, rtol=2e-7) np.testing.assert_allclose(out1.numpy(), r_ref+3, rtol=2e-7) - def test_dont_fold_arange_contiguous_view(self): - X = Tensor.randn(4, 4).realize() - r = (X+Tensor.arange(16).reshape(4, 4).contiguous()).sum(1, keepdim=True) - self.check_schedule([r], 2) - np.testing.assert_allclose(r.numpy(), (X.numpy()+np.arange(16).reshape(4, 4)).sum(1, keepdims=True), atol=1e-5, rtol=1e-6) - @unittest.skip("multi output isn't supported") def test_multiview_arange_children(self): X = Tensor.randn(2,3,4,4).numpy() @@ -2101,6 +2063,7 @@ class TestView(unittest.TestCase): run_schedule(sched) np.testing.assert_equal(b.numpy(), 0) + @expect_rangeify_fails def test_mask_dim_1(self): # mask out dim = 1 works too a = Tensor.rand(10, 10).realize() @@ -2127,6 +2090,7 @@ class TestView(unittest.TestCase): # a*VIEW(x), where VIEW(x) = 0 # x collapses along with its children + @unittest.skipIf(RANGEIFY, "this only fails if you run all of TestSchedule, some global tensor map bug?") def test_parent_view_collapses(self): a = Tensor([1, 2]) b = Tensor.arange(3).contiguous() @@ -2208,84 +2172,6 @@ class TestSimplifier(unittest.TestCase): assert UPat(Ops.CONST, arg=False).match(sink, {}), f"expected {sink} to collapse to a const False" assert sink.shape == a.shape -tensor_const_pm = PatternMatcher([ - (UPat(Ops.CONST, src=(UPat(Ops.VIEW, src=(UPat(Ops.DEVICE),)),)), lambda: True), - (UPat(Ops.BIND, src=(UPat(Ops.DEFINE_VAR, src=(UPat(Ops.VIEW, src=(UPat(Ops.DEVICE),)))), UPat(Ops.CONST))), lambda: True), -]) -class TestConst(unittest.TestCase): - # ** part 1: basic functionality of a tensor directly created from CONST - - def test_tensor_const(self): - a = Tensor(1) - print(a.uop) - self.assertTrue(tensor_const_pm.rewrite(a.uop)) - - def test_tensor_variable(self): - vv = UOp.variable("a", 0, 10).bind(1) - a = Tensor(vv) - print(a.uop) - self.assertTrue(tensor_const_pm.rewrite(a.uop)) - - def test_const_schedule(self): - a = Tensor.ones((4, 4)) - sched = a.schedule() - self.assertEqual(len(sched), 0) - - def test_const_contiguous_schedule(self): - # this ends up in the big graph - a = Tensor.ones((4,)).contiguous() - sched = a.schedule() - self.assertEqual(len(sched), 1) - - # ** part 2: scheduler behavior when const folding happens later - - def test_const_folding_no_realize(self): - a = Tensor([1, 2, 3, 4])*0 - sched = a.schedule() - self.assertEqual(len(sched), 0) - - def test_src_const_folding(self): - with Context(TRACK_MATCH_STATS=0): - a = Tensor.full((4,), 1).contiguous().realize() - b = Tensor.full((4,), 2).contiguous().realize() - mul0 = a*0 - add = b+mul0 - sched = add.schedule() - self.assertEqual(len(sched), 0) - # b+0 and b share the same underlying device memory - self.assertIs(add.uop.buffer, b.uop.buffer) - self.assertListEqual(add.tolist(), [2, 2, 2, 2]) - - def test_src_masked_const_folding(self): - with Context(TRACK_MATCH_STATS=0): - a = Tensor.full((4,), 1).contiguous().realize() - b = Tensor.full((6,), 2).contiguous().realize() - mul0 = a*0 - add = b+mul0.pad((1, 1), value=2) - sched = add.schedule() - self.assertEqual(len(sched), 1) - run_schedule(sched) - # add gets assigned to a new buffer - self.assertIsNot(add.uop.base.realized, b.uop.base.realized) - self.assertListEqual(add.tolist(), [4, 2, 2, 2, 2, 4]) - - # ** part 3: Tensor variable bindings - - #@unittest.expectedFailure # TODO: should schedule assert if you try to realize a Variable? - def test_var_schedule(self): - vv = UOp.variable("a", 0, 10).bind(1) - a = Tensor(vv) - sched = a.schedule() - self.assertEqual(len(sched), 0) - - def test_add_tvar(self): - vv = UOp.variable("a", 0, 10).bind(1) - a = Tensor(vv)+2 - sched, var_vals = a.schedule_with_vars() - self.assertEqual(len(sched), 1) - run_schedule(sched, var_vals) - self.assertEqual(a.tolist(), 3) - @unittest.skipIf(Device.DEFAULT == "CPU", "tests copy from another device to cpu") class TestCopyFolding(unittest.TestCase): def test_const_copy_is_free(self): @@ -2299,6 +2185,7 @@ class TestCopyFolding(unittest.TestCase): b = (a*zeros).to("CPU") run_schedule(check_schedule(b, 0, filter_sink=False)) self.assertListEqual(b.tolist(), [0, 0, 0]) + self.assertEqual(b.device, "CPU") def test_alu_after_copy(self): a = Tensor.ones((4,)).to("CPU") @@ -2307,6 +2194,12 @@ class TestCopyFolding(unittest.TestCase): add.kernelize() assert all_same([x.device for x in add.uop.src]), f"ALU has different devices! {[x.device for x in add.src]}" + def test_alu_before_copy(self): + buf = Tensor.ones(1).contiguous().realize() + a = buf+1 + b = a.to("CPU") + self.assertListEqual(b.tolist(), [2.]) + def test_copy_to_same_device(self): a = Tensor.empty(4).uop b = a.copy_to_device(a.device) @@ -2323,6 +2216,15 @@ class TestCopyFolding(unittest.TestCase): b = schedule_graph_rewrite(b) self.assertIs(b.base, a.base) + def test_copy_to_same_device_sched(self): + a = Tensor.ones(4).contiguous().realize().uop.as_buf() + t = Tensor(a.copy_to_device(a.device)) + sched = t.schedule() + assert len([s for s in sched if s.ast.op is Ops.COPY]) == 0 + run_schedule(sched) + assert t.uop.is_realized, f"didn't realize Tensor {t}" + self.assertListEqual(t.tolist(), [1.,1.,1.,1.]) + def test_clone(self): a = Tensor.empty(4) check_schedule(a.clone(), 1, filter_sink=False) @@ -2353,6 +2255,7 @@ class TestCopyFolding(unittest.TestCase): b.realize() self.assertListEqual(b.tolist(), [[0, 2], [1, 3]]) + @expect_rangeify_fails def test_permute_on_disk(self): with open(temp('dt_arange_4_permute'), "wb") as f: f.write(Tensor.arange(4).realize().uop.base.buffer.as_buffer()) a = Tensor.empty(4, dtype=dtypes.int32, device=f"disk:{temp('dt_arange_4_permute')}") @@ -2499,6 +2402,7 @@ class TestUOpBecome(unittest.TestCase): self.assertEqual(add.uop.shape, (8, 2)) assert add.uop is not add.uop.base + @expect_rangeify_fails def test_new_flat_buffer(self): a = Tensor.empty(4,) b = Tensor.empty(4,) @@ -2524,6 +2428,7 @@ class TestUOpBecome(unittest.TestCase): z = (img*x) / y check_schedule(z, 1) + @expect_rangeify_fails def test_become_existing_buffer(self): a = Tensor.empty(4, 4) b = a*1 @@ -2551,6 +2456,7 @@ class TestUOpBecome(unittest.TestCase): check_schedule(b, 0) assert UPat(Ops.CONST, arg=0).match(b.uop.base, {}) # scheduling replaces the tensor uop with a VIEW(BUFFER) + @expect_rangeify_fails def test_become_const_in_view(self): # if we shrink the base down to a size 0, only the VIEW becomes CONST, base is unchanged. add = Tensor.empty(2, 2)+Tensor.empty(2, 2) @@ -2568,6 +2474,7 @@ class TestUOpBecome(unittest.TestCase): assert UPat(Ops.CONST, arg=3).match(const_add.uop.base, {}) # tensors can become another realized tensor source + @expect_rangeify_fails def test_become_existing_buf_simple(self): a = Tensor.empty(4, 4) b = a+0 @@ -2576,12 +2483,14 @@ class TestUOpBecome(unittest.TestCase): self.assertIs(a.uop, b.uop) # they can also chain other movement ops on top of the tensor source + @expect_rangeify_fails def test_become_existing_buf_view(self): a = Tensor.empty(4, 4) b = a.permute((1, 0))+0 check_schedule(b, 0) self.assertEqual(b.uop.st, a.uop.permute((1, 0)).st) + @expect_rangeify_fails def test_become_existing_buf_view_alt(self): a = Tensor.empty(4, 4) b = a.permute((1, 0)).reshape((8, 2))+0 @@ -2589,6 +2498,7 @@ class TestUOpBecome(unittest.TestCase): self.assertEqual(b.uop.st, a.uop.permute((1, 0)).reshape((8, 2)).st) # they can also have other base parents that simplified, in that case we just backtrack to the chained mops + @expect_rangeify_fails def test_become_existing_buf_complex(self): a = Tensor.empty(4, 4) b = (a.permute((1, 0))+0).reshape((8, 2))+0 @@ -2596,6 +2506,7 @@ class TestUOpBecome(unittest.TestCase): self.assertEqual(b.uop.st, a.uop.permute((1, 0)).reshape((8, 2)).st) assert b.uop.base.op is Ops.BUFFER + @expect_rangeify_fails def test_become_multiple_choices(self): a = Tensor.empty(16) b = (a.reshape(1, 1, 4, 1, 4)+0).reshape(1, 1, 4, 4).shrink(((0, 1), (0, 1), (0, 3), (0, 3)))+0 @@ -2607,6 +2518,7 @@ class TestUOpBecome(unittest.TestCase): assert b.uop is c.uop assert UPat(Ops.VIEW, src=(UPat(Ops.BUFFER),)).match(c.uop, {}) + @expect_rangeify_fails def test_setitem_becomes_subbuffer(self): a = Tensor.full((4,), 2.).contiguous().realize() b = a.shrink(((0, 2),)).assign(Tensor.full((2,), 1.0)) diff --git a/test/test_setitem.py b/test/test_setitem.py index 967acc29f1..2005b7c801 100644 --- a/test/test_setitem.py +++ b/test/test_setitem.py @@ -1,4 +1,6 @@ import unittest +import random +from os import getenv from tinygrad import Tensor, TinyJit, Variable, dtypes from tinygrad.helpers import Context import numpy as np @@ -165,6 +167,41 @@ class TestSetitem(unittest.TestCase): t[idx] = val self.assertEqual(t.tolist(), [val]*idx_size+[idx_size]) + def test_setitem_advanced_indexing(self): + # Example from https://numpy.org/doc/stable/user/basics.indexing.html#combining-advanced-and-basic-indexing + t = Tensor.zeros(10,20,30,40,50).contiguous() + ind_1 = Tensor([5,3,7,8]) + ind_2 = Tensor([[[0],[1],[2]],[[3],[4],[5]]]) + v = Tensor.arange(2*3*4*10*30*50).reshape(2,3,4,10,30,50) + t[:, ind_1, :, ind_2, :] = v + n = np.zeros((10,20,30,40,50)) + n[:, ind_1.numpy(), :, ind_2.numpy(), :] = v.numpy() + np.testing.assert_allclose(t.numpy(), n) + + def test_setitem_2d_tensor_indexing(self): + t = Tensor.zeros(2).contiguous() + index = Tensor([[0, 1], [1,0]]) + v = Tensor.arange(2*2).reshape(2, 2).contiguous() + t[index] = v + n = np.zeros((2,)) + n[index.numpy()] = v.numpy() + np.testing.assert_allclose(t.numpy(), n) + + @unittest.skip("slow") + def test_setitem_tensor_indexing_fuzz(self): + random.seed(getenv("SEED", 42)) + for _ in range(getenv("ITERS", 100)): + size = random.randint(5, 10) + d0, d1, d2 = random.randint(1,5), random.randint(1,5), random.randint(1,5) + t = Tensor.zeros(size).contiguous() + n = np.zeros((size,)) + index = Tensor.randint((d0, d1, d2), low=0, high=size) + v = Tensor.arange(d0*d1*d2).reshape(d0, d1, d2) + t[index] = v + n[index.numpy()] = v.numpy() + np.testing.assert_allclose(t.numpy(), n, err_msg=f"failed with index={index.numpy().tolist()} and v={v.numpy().tolist()}") + + class TestWithGrad(unittest.TestCase): def test_no_requires_grad_works(self): z = Tensor.rand(8, 8) diff --git a/test/test_symbolic_jit.py b/test/test_symbolic_jit.py index fe819e6791..f28d274dcc 100644 --- a/test/test_symbolic_jit.py +++ b/test/test_symbolic_jit.py @@ -2,6 +2,7 @@ import unittest from test.helpers import assert_jit_cache_len from tinygrad import Variable, Tensor, TinyJit +from tinygrad.helpers import RANGEIFY import numpy as np class TestSymbolicJit(unittest.TestCase): @@ -11,11 +12,23 @@ class TestSymbolicJit(unittest.TestCase): a = Tensor.rand(3, 10) for i in range(1, 5): vi = Variable("i", 1, 10).bind(i) - symbolic = jf(a[:, :vi]).reshape(3, i).numpy() + symbolic = jf(a[:, :vi])[:3, :i].numpy() expected = f(a[:, :i]).numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) assert_jit_cache_len(jf, 1) + def test_plus1_pad(self): + # TODO: without contiguous, the pad is not captured in jit + def f(a): return (a+1).pad((None, (0, 10-a.shape[1]))).contiguous().realize() + jf = TinyJit(f) + a = Tensor.rand(3, 10) + for i in range(1, 5): + vi = Variable("i", 1, 10).bind(i) + symbolic = jf(a[:, :vi]).numpy() + expected = f(a[:, :i]).numpy() + np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) + assert_jit_cache_len(jf, 1 if RANGEIFY else 2) # one add and one pad, can be one kernel? + def test_add(self): def f(a, b): return (a+b).realize() jf = TinyJit(f) @@ -23,7 +36,8 @@ class TestSymbolicJit(unittest.TestCase): b = Tensor.rand(3, 10) for i in range(1, 5): vi = Variable("i", 1, 10).bind(i) - symbolic = jf(a[:, :vi], b[:, :vi]).reshape(3, i).numpy() + symbolic = jf(a[:, :vi], b[:, :vi]) + symbolic = symbolic[:3, :i].numpy() expected = f(a[:, :i], b[:, :i]).numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) assert_jit_cache_len(jf, 1) @@ -63,10 +77,10 @@ class TestSymbolicJit(unittest.TestCase): v = Tensor.rand(2, 10, 4, 8) for i in range(1, 5): vi = Variable("i", 1, 10).bind(i) - symbolic = jf(q, k[:, :vi], v[:, :vi]).reshape(2, 4, 1, 8).numpy() + symbolic = jf(q, k[:, :vi], v[:, :vi])[:2, :4, :1, :8].numpy() expected = f(q, k[:, :i], v[:, :i]).numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) - assert_jit_cache_len(jf, 5) + assert_jit_cache_len(jf, 4 if RANGEIFY else 5) def test_cat_dim0(self): def f(a, b): return a.cat(b, dim=0).realize() @@ -75,7 +89,7 @@ class TestSymbolicJit(unittest.TestCase): b = Tensor.rand(2, 3) for i in range(1, 5): vi = Variable("i", 1, 10).bind(i) - symbolic = jf(a[:vi], b).reshape(i+2, 3).numpy() + symbolic = jf(a[:vi], b)[:i+2, :3].numpy() expected = f(a[:i], b).numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) assert_jit_cache_len(jf, 1) @@ -87,7 +101,7 @@ class TestSymbolicJit(unittest.TestCase): b = Tensor.rand(3, 2) for i in range(1, 5): vi = Variable("i", 1, 10).bind(i) - symbolic = jf(a[:, :vi], b).reshape(3, i+2).numpy() + symbolic = jf(a[:, :vi], b)[:3, :i+2].numpy() expected = f(a[:, :i], b).numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) assert_jit_cache_len(jf, 1) @@ -101,7 +115,7 @@ class TestSymbolicJit(unittest.TestCase): for j in range(2, 5): vi = Variable("i", 1, 10).bind(i) vj = Variable("j", 1, 10).bind(j) - symbolic = jf(a[:vi], b[:vj]).reshape(i+j, 3).numpy() + symbolic = jf(a[:vi], b[:vj])[:i+j, :3].numpy() expected = f(a[:i], b[:j]).numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) assert_jit_cache_len(jf, 1) @@ -115,7 +129,7 @@ class TestSymbolicJit(unittest.TestCase): for j in range(2, 5): vi = Variable("i", 1, 10).bind(i) vj = Variable("j", 1, 10).bind(j) - symbolic = jf(a[:, :vi], b[:, :vj]).reshape(3, i+j).numpy() + symbolic = jf(a[:, :vi], b[:, :vj])[:3, :i+j].numpy() expected = f(a[:, :i], b[:, :j]).numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) assert_jit_cache_len(jf, 1) @@ -129,7 +143,7 @@ class TestSymbolicJit(unittest.TestCase): for j in range(2, 5): vi = Variable("i", 1, 10).bind(i) vj = Variable("j", 1, 10).bind(j) - symbolic = jf(a[:vi, :], b[:, :vj]).reshape(i, j).numpy() + symbolic = jf(a[:vi, :], b[:, :vj])[:i, :j].numpy() expected = f(a[:i, :], b[:, :j]).numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) assert_jit_cache_len(jf, 1) @@ -143,7 +157,7 @@ class TestSymbolicJit(unittest.TestCase): for j in range(2, 5): vi = Variable("i", 1, 10).bind(i) vj = Variable("j", 1, 10).bind(j) - symbolic = jf(a[:vj, :], b[:, :vi]).reshape(j, i).numpy() + symbolic = jf(a[:vj, :], b[:, :vi])[:j, :i].numpy() expected = f(a[:j, :], b[:, :i]).numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) assert_jit_cache_len(jf, 1) @@ -195,8 +209,8 @@ class TestSymbolicJit(unittest.TestCase): vi = Variable("i", 1, 10).bind(i) a = Tensor.ones(vi, 11).contiguous() symbolic = a[:, 1:2] - symbolic = jf(symbolic).reshape(i, 1).numpy() - expected = f(a.reshape(i, 11)[:, 1:2]).numpy() + symbolic = jf(symbolic)[:i, :1].numpy() + expected = f(a[:i, :][:, 1:2]).numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) assert_jit_cache_len(jf, 1) @@ -231,7 +245,7 @@ class TestSymbolicJit(unittest.TestCase): expected = b[:i].mean(0).numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) # axis = 1 - symbolic = jf1(c[:vi]).reshape(i).numpy() + symbolic = jf1(c[:vi])[:i].numpy() expected = c[:i].mean(1).numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) @@ -254,11 +268,11 @@ class TestSymbolicJit(unittest.TestCase): expected = a[:i, :j].mean().numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) # axis = 0 - symbolic = jf0(b[:vi, :vj]).reshape(j).numpy() + symbolic = jf0(b[:vi, :vj])[:j].numpy() expected = b[:i, :j].mean(0).numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) # axis = 1 - symbolic = jf1(c[:vi, :vj]).reshape(i).numpy() + symbolic = jf1(c[:vi, :vj])[:i].numpy() expected = c[:i, :j].mean(1).numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) @@ -283,7 +297,7 @@ class TestSymbolicJit(unittest.TestCase): expected = b[:i].var(0).numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) # axis = 1 - symbolic = jf1(c[:vi]).reshape(i).numpy() + symbolic = jf1(c[:vi])[:i].numpy() expected = c[:i].var(1).numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) @@ -306,11 +320,11 @@ class TestSymbolicJit(unittest.TestCase): expected = a[:i, :j].var().numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) # axis = 0 - symbolic = jf0(b[:vi, :vj]).reshape(j).numpy() + symbolic = jf0(b[:vi, :vj])[:j].numpy() expected = b[:i, :j].var(0).numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) # axis = 1 - symbolic = jf1(c[:vi, :vj]).reshape(i).numpy() + symbolic = jf1(c[:vi, :vj])[:i].numpy() expected = c[:i, :j].var(1).numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) diff --git a/test/test_symbolic_ops.py b/test/test_symbolic_ops.py index 2c2cbebdd3..991a9dcc93 100644 --- a/test/test_symbolic_ops.py +++ b/test/test_symbolic_ops.py @@ -13,7 +13,16 @@ class TestSymbolicOps(unittest.TestCase): a = Tensor.rand(3, 10) for i in range(1, 5): vi = Variable("i", 1, 10).bind(i) - symbolic = f(a[:, :vi]).reshape(3, i).numpy() + symbolic = f(a[:, :vi])[:3, :i].numpy() + expected = f(a[:, :i]).numpy() + np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) + + def test_plus1_pad(self): + def f(a): return (a+1).pad((None, (0, 10-a.shape[1]))).realize() + a = Tensor.rand(3, 10) + for i in range(1, 5): + vi = Variable("i", 1, 10).bind(i) + symbolic = f(a[:, :vi]).numpy() expected = f(a[:, :i]).numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) @@ -23,7 +32,7 @@ class TestSymbolicOps(unittest.TestCase): b = Tensor.rand(3, 10) for i in range(1, 5): vi = Variable("i", 1, 10).bind(i) - symbolic = f(a[:, :vi], b[:, :vi]).reshape(3, i).numpy() + symbolic = f(a[:, :vi], b[:, :vi])[:, :i].numpy() expected = f(a[:, :i], b[:, :i]).numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) @@ -46,7 +55,7 @@ class TestSymbolicOps(unittest.TestCase): vi = Variable("i", 1, 10).bind(i) if use_symbolic else i Tensor.realize(q, k, v) GlobalCounters.reset() - symbolic = f(q, k[:, :vi, :, :], v[:, :vi, :, :]).reshape(2, 4, 1, 8).numpy() + symbolic = f(q, k[:, :vi, :, :], v[:, :vi, :, :])[:2, :4, :1, :8].numpy() expected = f(q, k[:, :i, :, :], v[:, :i, :, :]).numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) @@ -85,7 +94,7 @@ class TestSymbolicOps(unittest.TestCase): for i in range(1, 5): vi = Variable("i", 1, 10).bind(i) b = Tensor.rand(2, 3) - symbolic = f(a[:vi, :], b).reshape(i+2, 3).numpy() + symbolic = f(a[:vi, :], b)[:i+2, :3].numpy() expected = f(a[:i, :], b).numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) @@ -95,7 +104,7 @@ class TestSymbolicOps(unittest.TestCase): for i in range(1, 5): vi = Variable("i", 1, 10).bind(i) b = Tensor.rand(3, 2) - symbolic = f(a[:, :vi], b).reshape(3, i+2).numpy() + symbolic = f(a[:, :vi], b)[:3, :i+2].numpy() expected = f(a[:, :i], b).numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) @@ -107,7 +116,7 @@ class TestSymbolicOps(unittest.TestCase): for j in range(2, 5): vi = Variable("i", 1, 10).bind(i) vj = Variable("j", 1, 10).bind(j) - symbolic = f(a[:vi, :], b[:vj, :]).reshape(i+j, 3).numpy() + symbolic = f(a[:vi, :], b[:vj, :])[:i+j, :3].numpy() expected = f(a[:i, :], b[:j, :]).numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) @@ -119,50 +128,41 @@ class TestSymbolicOps(unittest.TestCase): for j in range(2, 5): vi = Variable("i", 1, 10).bind(i) vj = Variable("j", 1, 10).bind(j) - symbolic = f(a[:, :vi], b[:, :vj]).reshape(3, i+j).numpy() + symbolic = f(a[:, :vi], b[:, :vj])[:3, :i+j].numpy() expected = f(a[:, :i], b[:, :j]).numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) def test_two_vars_plus1_ij(self): def f(a, b): return (a@b+1).realize() - a = Tensor.rand(10, 3) - b = Tensor.rand(3, 10) + a = Tensor.rand(10, 3).realize() + b = Tensor.rand(3, 10).realize() for i in range(2, 5): for j in range(2, 5): vi = Variable("i", 1, 10).bind(i) vj = Variable("j", 1, 10).bind(j) - symbolic = f(a[:vi, :], b[:, :vj]).reshape(i, j).numpy() + symbolic = f(a[:vi, :], b[:, :vj])[:i, :j].numpy() expected = f(a[:i, :], b[:, :j]).numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) def test_two_vars_plus1_ji(self): # reverse the order of variables def f(a, b): return (a@b+1).realize() - a = Tensor.rand(10, 3) - b = Tensor.rand(3, 10) + a = Tensor.rand(10, 3).realize() + b = Tensor.rand(3, 10).realize() for i in range(2, 5): for j in range(2, 5): vi = Variable("i", 1, 10).bind(i) vj = Variable("j", 1, 10).bind(j) - symbolic = f(a[:vj, :], b[:, :vi]).reshape(j, i).numpy() + symbolic = f(a[:vj, :], b[:, :vi])[:j, :i].numpy() expected = f(a[:j, :], b[:, :i]).numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) - def test_reshape_from_symbolic(self): - a = Tensor.rand(30) - for i in range(3, 5): - vi = Variable("i", 3, 10).bind(i) - symbolic = a[:vi*3].reshape((3, 3)).numpy() - # To match symbolic reshape (potential implicit shrink), we need a shrink - expected = a[:i*3].shrink(((0, 9),)).reshape((3, 3)).numpy() - np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) - def test_invalid_symbolic_reshape(self): a = Tensor.rand(30) for i in range(1, 5): vi = Variable("i", 1, 10).bind(i) # Cannot reshape into symbolic from non-symbolic - with self.assertRaises(AssertionError): a.reshape((3, vi)) + with self.assertRaises(ValueError): a.reshape((3, vi)) def test_shrink(self): for i in range(1, 5): @@ -178,6 +178,7 @@ class TestSymbolicOps(unittest.TestCase): vi = Variable("i", 1, 10).bind(i) a = Tensor.rand(7, 11) symbolic = a[3:5, vi:vi+2] + print(symbolic.shape) symbolic = symbolic.numpy() expected = a[3:5, i:i+2].numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) @@ -186,7 +187,7 @@ class TestSymbolicOps(unittest.TestCase): a = Tensor.rand(7, 11) for i in range(1, 5): vi = Variable("i", 1, 10).bind(i) - symbolic = a[3:5, :vi:1].reshape(2, i).numpy() + symbolic = a[3:5, :vi:1][:2, :i].numpy() expected = a[3:5, :i:1].numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) @@ -194,7 +195,7 @@ class TestSymbolicOps(unittest.TestCase): for i in range(1, 5): vi = Variable("i", 1, 10).bind(i) a = Tensor(1).unsqueeze(0).pad((0, 1)).unsqueeze(0) - symbolic = a.expand(vi, 2).reshape(i, 2).numpy() + symbolic = a.expand(vi, 2)[:i, :2].numpy() expected = a.expand(i, 2).numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) @@ -202,8 +203,8 @@ class TestSymbolicOps(unittest.TestCase): for i in range(1, 5): vi = Variable("i", 1, 10).bind(i) a = Tensor.ones(vi, 11).contiguous() - symbolic = a[:, 1:2].reshape(i, 1).numpy() - expected = a.reshape(i, 11)[:, 1:2].numpy() + symbolic = a[:, 1:2][:i, :1].numpy() + expected = Tensor.ones(i, 11)[:, 1:2].numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) def test_ones_sum(self): @@ -220,7 +221,11 @@ class TestSymbolicOps(unittest.TestCase): vi = Variable("i", 1, 10).bind(i) for axis in [None, 0, 1]: expected = a[:i].mean(axis).numpy() - symbolic = a[:vi].mean(axis).reshape(expected.shape).numpy() + symbolic = a[:vi].mean(axis) + if axis is None: + symbolic = symbolic.numpy() + else: + symbolic = symbolic[:expected.shape[0]].numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) def test_mean_2d(self): @@ -231,7 +236,11 @@ class TestSymbolicOps(unittest.TestCase): vj = Variable("j", 1, 10).bind(j) for axis in [None, 0, 1]: expected = a[:i, :j].mean(axis).numpy() - symbolic = a[:vi, :vj].mean(axis).reshape(expected.shape).numpy() + symbolic = a[:vi, :vj].mean(axis) + if axis is None: + symbolic = symbolic.numpy() + else: + symbolic = symbolic[:expected.shape[0]].numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) def test_var(self): @@ -240,7 +249,11 @@ class TestSymbolicOps(unittest.TestCase): vi = Variable("i", 1, 10).bind(i) for axis in [None, 0, 1]: expected = a[:i].var(axis).numpy() - symbolic = a[:vi].var(axis).reshape(expected.shape).numpy() + symbolic = a[:vi].var(axis) + if axis is None: + symbolic = symbolic.numpy() + else: + symbolic = symbolic[:expected.shape[0]].numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) def test_var_2d(self): @@ -251,7 +264,11 @@ class TestSymbolicOps(unittest.TestCase): vj = Variable("j", 1, 10).bind(j) for axis in [None, 0, 1]: expected = a[:i, :j].var(axis).numpy() - symbolic = a[:vi, :vj].var(axis).reshape(expected.shape).numpy() + symbolic_result = a[:vi, :vj].var(axis) + if axis is None: + symbolic = symbolic_result.numpy() + else: + symbolic = symbolic_result[:expected.shape[0]].numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) def test_bitcast_down(self): @@ -259,7 +276,11 @@ class TestSymbolicOps(unittest.TestCase): for i in range(1, 5): vi = Variable("i", 1, 10).bind(i) expected = a[:i].bitcast(dtypes.uint8).numpy() - symbolic = a[:vi].bitcast(dtypes.uint8).reshape(expected.shape).numpy() + symbolic_result = a[:vi].bitcast(dtypes.uint8) + if len(expected.shape) == 2: + symbolic = symbolic_result[:expected.shape[0], :expected.shape[1]].numpy() + else: + symbolic = symbolic_result[:].numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=0) @unittest.skipUnless(is_dtype_supported(dtypes.uint64), "no uint64") @@ -268,7 +289,11 @@ class TestSymbolicOps(unittest.TestCase): for i in range(1, 5): vi = Variable("i", 1, 10).bind(i) expected = a[:i].bitcast(dtypes.uint64).numpy() - symbolic = a[:vi].bitcast(dtypes.uint64).reshape(expected.shape).numpy() + symbolic_result = a[:vi].bitcast(dtypes.uint64) + if len(expected.shape) == 2: + symbolic = symbolic_result[:expected.shape[0], :expected.shape[1]].numpy() + else: + symbolic = symbolic_result[:].numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=0) @unittest.expectedFailure diff --git a/test/test_tensor.py b/test/test_tensor.py index 902a22f041..defeaac580 100644 --- a/test/test_tensor.py +++ b/test/test_tensor.py @@ -4,12 +4,12 @@ import torch import unittest, copy, mmap, random, math, array from tinygrad import Tensor, Device, dtypes from tinygrad.tensor import _METADATA -from tinygrad.helpers import getenv, temp, mv_address +from tinygrad.helpers import getenv, temp, mv_address, RANGEIFY 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.uop.ops import Ops, UOp -from tinygrad.runtime.support.compiler_cuda import PTX +from tinygrad.renderer.ptx import PTXRenderer from tinygrad.codegen import full_rewrite from tinygrad.dtype import DType @@ -516,10 +516,6 @@ class TestTinygrad(unittest.TestCase): print(c) def test_env_overwrite_default_device(self): - subprocess.run(['DISK=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT != \\"DISK\\""'], - shell=True, check=True) - subprocess.run(['NPY=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT != \\"NPY\\""'], - shell=True, check=True) subprocess.run([f'{Device.DEFAULT}=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], shell=True, check=True) subprocess.run([f'DISK=1 {Device.DEFAULT}=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], @@ -554,6 +550,11 @@ class TestTinygrad(unittest.TestCase): def test_shrink(self): t = Tensor.arange(32).contiguous().realize() self.assertListEqual(t[16:20].tolist(), [16,17,18,19]) + self.assertListEqual(t.shrink_to(16).tolist(), list(range(16))) + t = t.reshape(4, 8).contiguous().realize() + self.assertListEqual(t.shrink_to(2, 2).tolist(), [[0, 1], [8, 9]]) + with self.assertRaises(ValueError): t.shrink_to(2) + with self.assertRaises(ValueError): t.shrink_to(2, 2, 2) @unittest.skip("this test is just flaky, sync issue") class TestMoveTensor(unittest.TestCase): @@ -648,17 +649,22 @@ class TestZeroShapeTensor(unittest.TestCase): def test_pad(self): t = Tensor.rand(3, 2, 0).pad((None, None, (1, 1)), value=1) - assert t.shape == (3, 2, 2) + self.assertEqual(t.shape, (3, 2, 2)) np.testing.assert_equal(t.numpy(), np.ones((3, 2, 2))) t = Tensor.rand(3, 2, 0).pad((None, (1, 1), None), value=1) - assert t.shape == (3, 4, 0) + self.assertEqual(t.shape, (3, 4, 0)) np.testing.assert_equal(t.numpy(), np.ones((3, 4, 0))) t = Tensor.rand(3, 2, 0).pad(((1, 1), None, None), value=1) - assert t.shape == (5, 2, 0) + self.assertEqual(t.shape, (5, 2, 0)) np.testing.assert_equal(t.numpy(), np.ones((5, 2, 0))) + np.testing.assert_equal(Tensor([1, 2]).pad_to(4).numpy(), [1, 2, 0, 0]) + np.testing.assert_equal(Tensor([[1, 2]]).pad_to(2, 3).numpy(), [[1, 2, 0], [0, 0, 0]]) + with self.assertRaises(TypeError): Tensor([1, 2]).pad_to(2, 3) + with self.assertRaises(TypeError): Tensor([[1, 2]]).pad_to(3) + def test_shrink_into_zero(self): t = Tensor.rand(3, 4).realize() assert t.shrink((None, (2, 2))).realize().shape == (3, 0) @@ -865,11 +871,18 @@ class TestTensorMetadata(unittest.TestCase): self.assertEqual(y.grad.uop.metadata[0].name, "sigmoid") self.assertTrue(y.grad.uop.metadata[0].backward) si = Tensor.schedule(out, x.grad, y.grad)[-1] - self.assertEqual(len(si.metadata), 4, f"failed with {si.metadata}") - self.assertSetEqual(set(m.name for m in si.metadata), {"sigmoid", "__mul__", "relu"}) - bw = [m for m in si.metadata if m.backward] - self.assertEqual(len(bw), 2) - self.assertEqual(bw[0].name, "sigmoid") + if not RANGEIFY: + self.assertEqual(len(si.metadata), 4, f"failed with {si.metadata}") + self.assertSetEqual(set(m.name for m in si.metadata), {"sigmoid", "__mul__", "relu"}) + bw = [m for m in si.metadata if m.backward] + self.assertEqual(len(bw), 2) + self.assertEqual(bw[0].name, "sigmoid") + else: + self.assertEqual(len(si.metadata), 3, f"failed with {si.metadata}") + self.assertSetEqual(set(m.name for m in si.metadata), {"sigmoid", "relu"}) + bw = [m for m in si.metadata if m.backward] + self.assertEqual(len(bw), 1) + self.assertEqual(bw[0].name, "sigmoid") class TestIdxUpcast(unittest.TestCase): def _find_op(self, ast: UOp, op: Ops): @@ -915,13 +928,17 @@ class TestIdxUpcast(unittest.TestCase): def test_regular_sym(self): self.do_op_then_assert(dtypes.int, 2048, 2048, UOp.variable("dim3", 1, 64).bind(32)) - @unittest.skipIf(PTX, "PTX always convert Ops.INDEX to int64") + @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX always convert Ops.INDEX to int64") def test_symfold(self): # This would cause an overflow, but after sym fold it's within int32 a = Tensor.arange(65535) uops = self._schedule_render(a) assert all(uop.dtype is not dtypes.long for uop in uops) + def test_arange_raise_overflow(self): + with self.assertRaises(ValueError): + self._schedule_render(Tensor.arange(2**33, dtype=dtypes.int)) + @unittest.skipIf(is_dtype_supported(dtypes.long), "int64 is supported") def test_int64_unsupported_overflow_sym(self): with self.assertRaises(KeyError): diff --git a/test/test_tensor_variable.py b/test/test_tensor_variable.py index 0fa165e462..a046555d1b 100644 --- a/test/test_tensor_variable.py +++ b/test/test_tensor_variable.py @@ -38,7 +38,7 @@ class TestTensorVariable(unittest.TestCase): vv = Variable("a", 1, 10).bind(2) vv2 = Variable("b", 1, 10).bind(2) t = Tensor.ones(10, 10).contiguous()[:vv2, :vv] - ret = t.mean(axis=1).reshape(2, 1).numpy() + ret = t.mean(axis=1)[:2].reshape(2, 1).numpy() assert np.all(ret == 1) def test_symbolic_mean_2d_add(self): @@ -66,25 +66,25 @@ class TestTensorVariable(unittest.TestCase): def test_symbolic_arange(self): vv = Variable("a", 1, 10) ret = Tensor.arange(0, vv.bind(4)) - self.assertListEqual(ret.reshape(4).tolist(), [0,1,2,3]) + self.assertListEqual(ret[:4].tolist(), [0,1,2,3]) def test_symbolic_arange_sym_start(self): vv = Variable("a", 1, 6) ret = Tensor.arange(vv.bind(4), 7) - self.assertListEqual(ret.reshape(3).tolist(), [4,5,6]) + self.assertListEqual(ret[:3].tolist(), [4,5,6]) # TODO: add vmin/vmax pattern for symbolic denominator @unittest.expectedFailure def test_symbolic_arange_sym_step(self): vv = Variable("step", 1, 3) ret = Tensor.arange(0, 10, vv.bind(2)) - self.assertListEqual(ret.reshape(5).tolist(), [0,2,4,6,8]) + self.assertListEqual(ret[:5].tolist(), [0,2,4,6,8]) def test_symbolic_arange_two_vars(self): begin = Variable("b", 1, 5) end = Variable("e", 6, 10) ret = Tensor.arange(begin.bind(4), end.bind(7)) - self.assertListEqual(ret.reshape(3).tolist(), [4,5,6]) + self.assertListEqual(ret[:3].tolist(), [4,5,6]) def test_variable_empty(self): v = Variable("i", 1, 10) diff --git a/test/test_tiny.py b/test/test_tiny.py index bf0f903ff2..31bb84f595 100644 --- a/test/test_tiny.py +++ b/test/test_tiny.py @@ -7,6 +7,14 @@ class TestTiny(unittest.TestCase): # *** basic functionality *** + def test_const(self): + const = Tensor(2.0) + self.assertEqual(const.item(), 2.0) + + def test_copy(self): + out = Tensor([1.,2,3]) + self.assertListEqual(out.tolist(), [1.0, 2.0, 3.0]) + def test_plus(self): out = Tensor([1.,2,3]) + Tensor([4.,5,6]) self.assertListEqual(out.tolist(), [5.0, 7.0, 9.0]) @@ -87,7 +95,7 @@ class TestTiny(unittest.TestCase): ones = Tensor.ones(10).contiguous() for s in [2,5]: ret = ones[:i.bind(s)] + 1 - self.assertListEqual(ret.contiguous().reshape(s).tolist(), [2.0]*s) + self.assertListEqual(ret.contiguous()[:s].tolist(), [2.0]*s) def test_symbolic_reduce(self): i = Variable('i', 1, 10) @@ -137,7 +145,7 @@ class TestTiny(unittest.TestCase): # *** image *** - @unittest.skipIf(Device.DEFAULT != "GPU", "image only supported on GPU") + @unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL") def test_image(self): with Context(IMAGE=2): self.test_gemm(N=4, out_dtype=dtypes.imagef((4, 1, 4))) diff --git a/test/test_uop_graph.py b/test/test_uop_graph.py index 293264b7e9..e0d07c7916 100644 --- a/test/test_uop_graph.py +++ b/test/test_uop_graph.py @@ -417,13 +417,45 @@ class TestUOpGraph(unittest.TestCase): uops = to_uops_list([v.bitcast(dt)]) self.assertEqual(len([x for x in uops if x.op is Ops.BITCAST]), 0, f"dtype = {dt}") + def test_where_on_gated_load_fold(self): + ridx0 = UOp.range(100, 0) + d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 0) + ld = d0.index(ridx0, ridx0<50).load() + w = (ridx0<50).where(ld, 5) + uops = to_uops_list([w]) + for u in uops: + assert u.op is not Ops.WHERE + if u.op is Ops.LOAD: assert u.src[1].arg==5 + + def test_where_on_gated_load_folds_swapped_branches(self): + ridx0 = UOp.range(100, 0) + d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 0) + ld = d0.index(ridx0, (ridx0<50).logical_not()).load() + w = (ridx0<50).where(5, ld) + uops = to_uops_list([w]) + for u in uops: + assert u.op is not Ops.WHERE + if u.op is Ops.LOAD: assert u.src[1].arg==5 + + def test_where_in_store_becomes_gate(self): + ridx0 = UOp.range(100, 0) + d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 0) + idx = d0.index(ridx0) + ld = idx.load() + val = (ridx0<50).where(5, ld) + st = idx.store(val, ridx0) + uops = to_uops_list([st]) + for u in uops: + assert u.op is not Ops.WHERE + if u.op is Ops.STORE: assert u.src[1].arg==5 + def test_load_idx_becomes_int(self): d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 0) d1 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 1) - l0 = UOp(Ops.LOAD, dtypes.long, (d0.index(UOp.const(dtypes.int, 0)),)) + l0 = UOp(Ops.LOAD, dtypes.long, (d0.index(UOp.const(dtypes.int, 0)),)).cast(dtypes.index) idx = l0 * 600 valid = (l0<-1).ne(True)&(l0<3000) - l1 = UOp(Ops.LOAD, dtypes.long, (d1.index(idx, valid),)) + l1 = UOp(Ops.LOAD, dtypes.long, (d1.index(idx.valid(valid)),)) uops = to_uops_list([l1]) for u in uops: if u.op is Ops.INDEX: self.assertEqual(u.src[1].dtype, dtypes.int) @@ -560,7 +592,7 @@ class TestUOpGraph(unittest.TestCase): glbl1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 1) glbl2 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 2) idx = UOp.const(dtypes.int, 0) - ld0 = UOp(Ops.LOAD, dtypes.int, (glbl1.index(idx, UOp.const(dtypes.bool, False)),)) + ld0 = UOp(Ops.LOAD, dtypes.int, (glbl1.index(UOp.invalid()),)) ld1 = UOp(Ops.LOAD, dtypes.int, (glbl2.index(idx, UOp.const(dtypes.bool, True)),)) uops = to_uops_list([UOp(Ops.STORE, dtypes.void, (glbl0.index(idx), ld1+ld0))]) ld0 = uops[-1].src[-1] @@ -573,7 +605,7 @@ class TestUOpGraph(unittest.TestCase): lidx = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 16),), "lidx0") st = UOp(Ops.STORE, dtypes.void, (smem.index(lidx), UOp.load(glbl0.index(lidx), dtype=dtypes.int))) barrier = UOp(Ops.BARRIER, dtypes.void, (st, )) - ld0 = UOp(Ops.LOAD, dtypes.int, (smem.index(lidx+1, UOp.const(dtypes.bool, False)), barrier)) + ld0 = UOp(Ops.LOAD, dtypes.int, (smem.index(UOp.invalid()), barrier)) ld1 = UOp(Ops.LOAD, dtypes.int, (smem.index(lidx+2, UOp.const(dtypes.bool, True)), barrier)) uops = to_uops_list([UOp(Ops.STORE, dtypes.void, (glbl0.index(lidx), ld1+ld0))]) @@ -586,7 +618,7 @@ class TestUOpGraph(unittest.TestCase): idx0 = UOp.const(dtypes.int, 0) idx1 = UOp.const(dtypes.int, 0) val = UOp.const(dtypes.int, 42) - st0 = glbl.index(idx0, UOp.const(dtypes.bool, False)).store(val) + st0 = glbl.index(UOp.invalid()).store(val) st1 = glbl.index(idx0, UOp.const(dtypes.bool, True)).store(val) uops = to_uops_list([st0, st1]) # only the second store happens diff --git a/test/test_uops.py b/test/test_uops.py index c8c350b91f..4a7ed30880 100644 --- a/test/test_uops.py +++ b/test/test_uops.py @@ -15,6 +15,7 @@ from tinygrad.codegen import full_rewrite from tinygrad.uop.symbolic import sym from tinygrad.device import is_dtype_supported from tinygrad.codegen.opt import Opt, OptOps +from tinygrad.renderer.ptx import PTXRenderer def to_uops_list(u:list[UOp], opts=None, skip_check=False) -> list[UOp]: return full_rewrite(UOp.sink(*u), opts) @@ -130,9 +131,9 @@ class TestFloatUOps(TestUOps): class TestNonFloatUOps(TestUOps): def test_add_int32(self): self._test_bop_fxn(Ops.ADD, lambda a,b: int(a)+int(b), (dtypes.int32, dtypes.int32)) def test_mul_int32(self): self._test_bop_fxn(Ops.MUL, lambda a,b: int(a)*int(b), (dtypes.int32, dtypes.int32)) - @unittest.skipUnless(getenv("PTX"), "only ptx uses bitshifts") + @unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "only ptx uses bitshifts") def test_shr_int32(self): self._test_bop_fxn(Ops.SHR, lambda a,b: int(a)>>int(b), (dtypes.int32, dtypes.int32), no_b_neg=True) - @unittest.skipUnless(getenv("PTX"), "only ptx uses bitshifts") + @unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "only ptx uses bitshifts") def test_shl_int32(self): self._test_bop_fxn(Ops.SHL, lambda a,b: int(a)< bytes: return src.encode() @@ -56,7 +101,7 @@ class TestCompiler(unittest.TestCase): class TestRunAsModule(unittest.TestCase): def test_module_runs(self): p = subprocess.run([sys.executable, "-m", "tinygrad.device"],stdout=subprocess.PIPE, stderr=subprocess.PIPE, - env={**os.environ, "DEBUG": "1"}, timeout=10,) + env={**os.environ, "DEBUG": "1"}, timeout=30,) out = (p.stdout + p.stderr).decode() self.assertEqual(p.returncode, 0, msg=out) self.assertIn("CPU", out) # for sanity check diff --git a/test/unit/test_disk_tensor.py b/test/unit/test_disk_tensor.py index 467c9e066b..57584df295 100644 --- a/test/unit/test_disk_tensor.py +++ b/test/unit/test_disk_tensor.py @@ -307,7 +307,7 @@ class TestDiskTensor(unittest.TestCase): ret = t.bitcast(dtypes.uint16).to("CPU") + 1 assert ret.tolist() == [2827, 3341, 3855, 4369] - @unittest.skipIf(OSX, "new LLVM has an issue on OSX") + @unittest.skipIf(OSX or Device.DEFAULT == "CL", "new LLVM has an issue on OSX, CL=1 gives the wrong output") def test_bf16_disk_write_read(self): t = Tensor([10000, -1, -1000, -10000, 20], dtype=dtypes.float32) t.to(f"disk:{temp('dt_bf16_disk_write_read_f32')}").realize() diff --git a/test/unit/test_dtype_spec.py b/test/unit/test_dtype_spec.py index 63990f7fd3..3c1ffc0ee5 100644 --- a/test/unit/test_dtype_spec.py +++ b/test/unit/test_dtype_spec.py @@ -1,12 +1,11 @@ import unittest, math, operator, subprocess, struct from tinygrad.tensor import Tensor, dtypes, Device -from tinygrad.dtype import DType, DTYPES_DICT, truncate, truncate_fp16, float_to_bf16, _to_np_dtype, least_upper_dtype, least_upper_float +from tinygrad.dtype import DType, 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.helpers import getenv, CI, DEBUG from hypothesis import given, settings, strategies as strat import numpy as np import torch -import ml_dtypes settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False)) settings.load_profile("my_profile") @@ -22,7 +21,9 @@ def _assert_eq(tensor:Tensor, target_dtype:DType, target, tol_target_dtype:float if DEBUG >= 2: print(tensor.numpy()) try: assert tensor.dtype == target_dtype - np.testing.assert_allclose(tensor.numpy(), target, rtol={dtypes.float16:1e-3, dtypes.bfloat16:1e-2}.get(target_dtype, tol_target_dtype)) + np.testing.assert_allclose(tensor.numpy(), target, rtol={dtypes.float16:1e-3, dtypes.bfloat16:1e-2, + dtypes.fp8e4m3:1e-1, dtypes.fp8e5m2:5e-1}.get(target_dtype, tol_target_dtype)) + except AssertionError as e: raise AssertionError(f"\ntensor {tensor.numpy()} dtype {tensor.dtype} does not match target {target} with dtype {target_dtype}") from e @@ -105,16 +106,16 @@ class TestHelpers(unittest.TestCase): self.assertEqual(dt.min, dt.vec(4).min) self.assertEqual(dt.max, dt.vec(4).max) - def test_truncate_fp16(self): - self.assertEqual(truncate_fp16(1), 1) - self.assertEqual(truncate_fp16(65504), 65504) - self.assertEqual(truncate_fp16(65519.999), 65504) - self.assertEqual(truncate_fp16(65520), math.inf) - self.assertEqual(truncate_fp16(1e-8), 0.0) - self.assertEqual(truncate_fp16(-65504), -65504) - self.assertEqual(truncate_fp16(-65519.999), -65504) - self.assertEqual(truncate_fp16(-65520), -math.inf) - self.assertTrue(math.isnan(truncate_fp16(math.nan))) + def test_float_to_fp16(self): + self.assertEqual(float_to_fp16(1), 1) + self.assertEqual(float_to_fp16(65504), 65504) + self.assertEqual(float_to_fp16(65519.999), 65504) + self.assertEqual(float_to_fp16(65520), math.inf) + self.assertEqual(float_to_fp16(1e-8), 0.0) + self.assertEqual(float_to_fp16(-65504), -65504) + self.assertEqual(float_to_fp16(-65519.999), -65504) + self.assertEqual(float_to_fp16(-65520), -math.inf) + self.assertTrue(math.isnan(float_to_fp16(math.nan))) def test_float_to_bf16(self): # TODO: fuzz this better @@ -190,7 +191,7 @@ class TestHelpers(unittest.TestCase): elif math.isinf(x): np.testing.assert_equal(truncate[dtypes.fp8e4m3](x), math.copysign(math.nan, x)) elif x > FP8E4M3_MAX: np.testing.assert_equal(truncate[dtypes.fp8e4m3](x), FP8E4M3_MAX) elif x < -FP8E4M3_MAX: np.testing.assert_equal(truncate[dtypes.fp8e4m3](x), -FP8E4M3_MAX) - else: np.testing.assert_equal(truncate[dtypes.fp8e4m3](x), ml_dtypes.float8_e4m3fn(x)) + else: np.testing.assert_equal(truncate[dtypes.fp8e4m3](x), torch.tensor(x, dtype=torch.float8_e4m3fn).float().item()) @given(strat.floats(width=32, allow_subnormal=True, allow_nan=True, allow_infinity=True)) def test_truncate_fp8e5m2(self, x): @@ -198,7 +199,7 @@ class TestHelpers(unittest.TestCase): elif math.isinf(x): np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), x) elif x > FP8E5M2_MAX: np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), FP8E5M2_MAX) elif x < -FP8E5M2_MAX: np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), -FP8E5M2_MAX) - else: np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), ml_dtypes.float8_e5m2(x)) + else: np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), torch.tensor(x, dtype=torch.float8_e5m2).float().item()) class TestTypeSpec(unittest.TestCase): def setUp(self): @@ -378,7 +379,7 @@ class TestTypePromotion(unittest.TestCase): assert least_upper_dtype(dtypes.int32, dtypes.uint32) == dtypes.int64 assert least_upper_dtype(dtypes.uint32, dtypes.int64) == dtypes.int64 # similar to jax but we don't use weak type - assert least_upper_dtype(dtypes.int64, dtypes.uint64) == dtypes.float16 + assert least_upper_dtype(dtypes.int64, dtypes.uint64) == dtypes.fp8e4m3 assert least_upper_dtype(dtypes.float16, dtypes.float32) == dtypes.float32 assert least_upper_dtype(dtypes.float32, dtypes.float64) == dtypes.float64 @@ -387,6 +388,14 @@ class TestTypePromotion(unittest.TestCase): assert least_upper_dtype(dtypes.float16, dtypes.int64) == dtypes.float16 assert least_upper_dtype(dtypes.float16, dtypes.uint64) == dtypes.float16 assert least_upper_dtype(dtypes.fp8e4m3, dtypes.fp8e5m2) == dtypes.half + assert least_upper_dtype(dtypes.fp8e4m3, dtypes.bfloat16) == dtypes.bfloat16 + assert least_upper_dtype(dtypes.fp8e5m2, dtypes.bfloat16) == dtypes.bfloat16 + assert least_upper_dtype(dtypes.fp8e4m3, dtypes.float16) == dtypes.float16 + assert least_upper_dtype(dtypes.fp8e5m2, dtypes.float16) == dtypes.float16 + assert least_upper_dtype(dtypes.fp8e4m3, dtypes.int64) == dtypes.fp8e4m3 + assert least_upper_dtype(dtypes.fp8e4m3, dtypes.uint64) == dtypes.fp8e4m3 + assert least_upper_dtype(dtypes.fp8e5m2, dtypes.int64) == dtypes.fp8e5m2 + assert least_upper_dtype(dtypes.fp8e5m2, dtypes.uint64) == dtypes.fp8e5m2 class TestAutoCastType(unittest.TestCase): def setUp(self): @@ -569,10 +578,10 @@ class TestAutoCastType(unittest.TestCase): def test_gradient_dtype(self): old_default_float = dtypes.default_float - for default_dtype in [dtypes.float16, dtypes.bfloat16, dtypes.float32, dtypes.float64]: + for default_dtype in dtypes.floats: if not is_dtype_supported(default_dtype): continue dtypes.default_float = default_dtype - for dtype in [dtypes.float16, dtypes.bfloat16, dtypes.float32, dtypes.float64]: + for dtype in dtypes.floats: if not is_dtype_supported(dtype): continue if DEBUG >= 2: print(f"testing {default_dtype=}, {dtype=}") diff --git a/test/unit/test_gguf.py b/test/unit/test_gguf.py index 32dca21e06..1c9cd5bc28 100644 --- a/test/unit/test_gguf.py +++ b/test/unit/test_gguf.py @@ -67,7 +67,7 @@ class TestGGUF(unittest.TestCase): return np.array([E] + packed, dtype=np.uint8) def decode(code, E): - sign = -1.0 if code * 0b1000 else 1.0 + sign = -1.0 if (code & 0b1000) else 1.0 exp = (code >> 1) & 0b11 mant = code & 0b1 val = (1.0 + 0.5 * mant) * np.exp2(exp - 1) if exp else 0.5 * mant @@ -83,7 +83,8 @@ class TestGGUF(unittest.TestCase): expected.extend(decode(c, E) for c in codes) tensor = Tensor(np.concatenate(blocks)) out = ggml_data_to_tensor(tensor, len(expected), MXFP4) - self.assertListEqual(out.numpy().tolist(), np.array(expected, dtype=np.float32).tolist()) + # TODO: should this be exact equal? somehow failed on CI + np.testing.assert_allclose(out.numpy(), expected, atol=0.0, rtol=1e-6) def test_expected_failure_unknown_type(self): with self.assertRaises(ValueError): diff --git a/test/unit/test_indexing.py b/test/unit/test_indexing.py index da5d61944c..c9d6d7c7da 100644 --- a/test/unit/test_indexing.py +++ b/test/unit/test_indexing.py @@ -181,7 +181,7 @@ class TestIndexing(unittest.TestCase): # self.assertRaises(TypeError, delitem) # TODO: LLVM is quite fast, why are other compiled backends slow? - @unittest.skipIf(CI and Device.DEFAULT in ["CPU", "GPU", "METAL", "NV", "AMD"], "slow") + @unittest.skipIf(CI and Device.DEFAULT in ["CPU", "CL", "METAL", "NV", "AMD"], "slow") def test_advancedindex(self): # integer array indexing diff --git a/test/unit/test_linearizer_rewrite.py b/test/unit/test_linearizer_rewrite.py index 46b3b1aba1..d578e3ae0d 100644 --- a/test/unit/test_linearizer_rewrite.py +++ b/test/unit/test_linearizer_rewrite.py @@ -23,7 +23,6 @@ class TestLinearizerRewrite(unittest.TestCase): si = out.schedule()[-1] opts_to_apply = [] opts_to_apply.append(Opt(OptOps.UPCAST, 0, 4)) - opts_to_apply.append(Opt(OptOps.UNROLL, 0, 4)) ast = si.ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply))) prg = get_program(ast, Device["CPU"].renderer) print(prg.src) diff --git a/test/unit/test_shapetracker.py b/test/unit/test_shapetracker.py index a3cfc80687..849aa19d5d 100644 --- a/test/unit/test_shapetracker.py +++ b/test/unit/test_shapetracker.py @@ -1,7 +1,7 @@ #!/usr/bin/env python import unittest import numpy as np -from tinygrad.dtype import dtypes +from tinygrad.dtype import dtypes, Invalid from tinygrad.helpers import prod from tinygrad.shape.shapetracker import ShapeTracker, View from tinygrad import Variable @@ -10,7 +10,8 @@ from tinygrad.codegen.late.devectorizer import sym from itertools import product def shapetracker_getitem(st:ShapeTracker, val:int): - idx, valid = st.reshape((st.size,)).to_indexed_uops([UOp.const(dtypes.int, val)]) + valid_idx = st.reshape((st.size,)).to_valid_uop([UOp.const(dtypes.int, val)]) + idx, valid = valid_idx.get_idx(), valid_idx.get_valid() idx, valid = graph_rewrite(idx, sym), graph_rewrite(valid, sym) assert idx.op is Ops.CONST and valid.op is Ops.CONST return idx.arg, valid.arg @@ -68,7 +69,7 @@ class CheckingShapeTracker: def contiguous(self): return self.st.contiguous def assert_same(self): - x = [(v[0] if (v:=shapetracker_getitem(self.st, i))[1] else -1) for i in range(prod(self.st.shape))] + x = [(v[0] if (v:=shapetracker_getitem(self.st, i))[1] and v[0] is not Invalid else -1) for i in range(prod(self.st.shape))] y = [self[i] for i in range(prod(self.shape))] assert self.st.shape == self.shape assert x == y, f"mismatch shapetracker:{x} real:{y}" @@ -154,7 +155,7 @@ class TestRealStrides(unittest.TestCase): View.create((1, 3, 22, 21), (0, 192, 16, 1), 0, ((0, 1), (0, 3), (0, 12), (0, 16))), View.create((3, 11, 7, 2, 3), (462, 21, 1, 231, 7), 0, None), )) - self.assertEqual(st.real_strides(), (132, None, None, None, None)) + self.assertEqual(st.real_strides(), (132, 12, None, None, None)) class TestRealSimplifies(unittest.TestCase): def tearDown(self): @@ -619,20 +620,6 @@ class TestMaskedShapeTracker(unittest.TestCase): st3.reshape((4, 3, 6, 5)) st3.assert_same() - def test_axis_is_masked(self): - st = ShapeTracker.from_shape((100, 100, 100, 100)).pad(((0,1),(0,0),(2,0), (0,0))) - assert st.axis_is_masked(0) - assert not st.axis_is_masked(1) - assert st.axis_is_masked(2) - assert not st.axis_is_masked(3) - - def test_axis_is_masked_rw1(self): - st = ShapeTracker(views=(View(shape=(1, 2, 1, 4, 4, 13, 4, 13), strides=(0, 324, 0, 81, 0, 9, 0, 1), offset=-20, - mask=((0, 1), (0, 2), (0, 1), (0, 4), (0, 4), (2, 11), (0, 4), (2, 11)), contiguous=False), - View(shape=(2, 4, 11, 11, 4, 3, 3), strides=(10816, 0, 52, 1, 2704, 728, 14), offset=0, - mask=None, contiguous=False))) - assert not st.axis_is_masked(0) - class TestShapeTracker(unittest.TestCase): def setUp(self): self.st = CheckingShapeTracker((7,4)) @@ -830,12 +817,14 @@ class TestShapeTrackerSize(unittest.TestCase): class TestRender(unittest.TestCase): def test_render(self): st = ShapeTracker.from_shape((2, 3)) - idx, valid = st.to_indexed_uops() + valid_idx = st.to_valid_uop() + idx, valid = valid_idx.get_idx(), valid_idx.get_valid() self.assertEqual(idx.render(), "((ridx0*3)+ridx1)") self.assertEqual(valid.render(), "True") st = st.pad(((0, 1), (0, 0))) - idx, valid = st.to_indexed_uops() + valid_idx = st.to_valid_uop() + idx, valid = valid_idx.get_idx(), valid_idx.get_valid() self.assertEqual(idx.render(), "((ridx0*3)+ridx1)") self.assertEqual(valid.render(), "(ridx0<2)") diff --git a/test/unit/test_simplify_valid_idx.py b/test/unit/test_simplify_valid_idx.py index d75c872209..866a70020d 100644 --- a/test/unit/test_simplify_valid_idx.py +++ b/test/unit/test_simplify_valid_idx.py @@ -8,13 +8,13 @@ from tinygrad.helpers import Context def get_gated_load_uop(valid:UOp, idx:UOp): return UOp(Ops.LOAD, dtypes.float, ( - UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=0).index(idx, valid), + UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=0).index(idx.valid(valid)), 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.DEFINE_GLOBAL, dtypes.imagef(image_shape), arg=0).index(UOp(Ops.VECTORIZE, dtypes.int.vec(2), idx), valid), + UOp(Ops.DEFINE_GLOBAL, dtypes.imagef(image_shape), arg=0).index(UOp(Ops.VECTORIZE, dtypes.index.vec(2), idx).valid(valid)), UOp(Ops.VECTORIZE, dtypes.float.vec(4), src=(UOp.const(dtypes.float, 0.0),) * 4) )) @@ -359,7 +359,7 @@ class TestImageSimplification(unittest.TestCase): self.check(load, None, "((gidx*3)+-1438)", "0") def test_simplify2(self): - # from GPU=1 DEBUG=4 FORWARD_ONLY=1 IMAGE=2 python3 test/test_ops.py TestOps.test_simple_padding_conv2d + # from CL=1 DEBUG=4 FORWARD_ONLY=1 IMAGE=2 python3 test/test_ops.py TestOps.test_simple_padding_conv2d lidx = Special("lidx", 4) valid = (lidx<3) & (lidx<1).ne(True) idx = ((lidx+1)%2, (lidx+1)//2-1) diff --git a/test/unit/test_symbolic_shapetracker.py b/test/unit/test_symbolic_shapetracker.py index 7d77e81b97..c89419a2f9 100644 --- a/test/unit/test_symbolic_shapetracker.py +++ b/test/unit/test_symbolic_shapetracker.py @@ -13,7 +13,6 @@ class TestSymbolic(unittest.TestCase): assert st.shape == (x, 3) assert st.real_strides() == (3, 1) - @unittest.expectedFailure def test_real_strides_0(self): st = ShapeTracker(views=(View(shape=(2, (Variable('start_pos', 1, 8)+1), 1, 1), strides=(8, 1, 0, 0), offset=0, mask=((0, 2), (0, Variable('start_pos', 1, 8)), (0, 1), (0, 1)), contiguous=False), View(shape=(2, (Variable('start_pos', 1, 8)+1)), strides=((Variable('start_pos', 1, 8)+1), 1), offset=0, mask=None, contiguous=True))) # noqa: E501 self.assertEqual(st.real_strides(), (8, None)) @@ -198,7 +197,7 @@ class TestSymbolicPad(unittest.TestCase): def test_pad(self): v = Variable("v", 1, 100).bind(5) t = Tensor.ones(100)[:v].pad(((4, 0),)) - t = t.reshape(9) + t = t[:9] assert t.tolist() == [0,0,0,0,1,1,1,1,1] diff --git a/test/unit/test_tensor_uop_representation.py b/test/unit/test_tensor_uop_representation.py index a1b2f0526d..fe8d47bb6d 100644 --- a/test/unit/test_tensor_uop_representation.py +++ b/test/unit/test_tensor_uop_representation.py @@ -34,7 +34,7 @@ class TestTensorMutates(unittest.TestCase): is_pattern_uop(c.uop.base, realized_pattern) # NOTE: we keep movement ops on top of the buffer view is_pattern_uop(c.uop, UPat(Ops.BUFFER)) - is_pattern_uop(d.uop, UPat(Ops.VIEW, src=(realized_pattern,))) + assert d.uop is not d.uop.base def test_reshape_is_same_child(self): a = Tensor([1,2,3]) @@ -58,40 +58,6 @@ class TestTensorUopRepresentation(unittest.TestCase): print(c.uop) is_pattern(c, UPat(Ops.ADD, src=(realized_pattern, realized_pattern))) - def test_const_pattern(self): - a = Tensor(1) - print(a.uop) - is_pattern(a, const_pattern) # const in tensor has a DEVICE and VIEW src - is_pattern(a, UPat.cvar("x")) # even cvar works! - - def test_consts_do_not_realize(self): - a = Tensor(1) - print(a.uop) - pre_realize = a.uop - a.realize() - assert a.uop is pre_realize - - def test_viewed_consts_do_not_realize(self): - a = Tensor.ones(10, 10) - print(a.uop) - a.realize() - is_pattern(a, const_pattern) - self.assertEqual(a.uop.shape, (10, 10)) - - # CONST is EXPAND -> RESHAPE -> CONST -> DEVICE - def test_consts_dont_have_buffers(self): - a = Tensor.ones(10, 10) - buffers_in_parents = [x.op for x in a.uop.toposort() if x.op is Ops.BUFFER] - self.assertEqual(len(buffers_in_parents), 0) - is_pattern(a, UPat(Ops.EXPAND, src=(UPat(Ops.RESHAPE, src=(const_pattern,)),))) - - # COPY has a copyin source and a device. - def test_copyin(self): - a = Tensor([1.,2,3]).realize() - c = a.to("TEST") # NOTE: this isn't checked - print(c.uop) - is_pattern(c, UPat(Ops.COPY, src=(realized_pattern, UPat(Ops.DEVICE)), arg=None)) - def test_empty_buf(self): a = Tensor.empty(3, 3) is_pattern(a, UPat(Ops.RESHAPE, src=(UPat(Ops.BUFFER),))) diff --git a/test/unit/test_uop_symbolic.py b/test/unit/test_uop_symbolic.py index ec03ee23ac..5d9ccf6221 100644 --- a/test/unit/test_uop_symbolic.py +++ b/test/unit/test_uop_symbolic.py @@ -2,7 +2,7 @@ import unittest, pickle, functools, math import z3 -from tinygrad.dtype import dtypes, ConstType, DType +from tinygrad.dtype import dtypes, ConstType, DType, Invalid from tinygrad.codegen import full_rewrite from tinygrad.helpers import Context from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer, track_rewrites @@ -93,6 +93,37 @@ class TestSymbolic(unittest.TestCase): assert idx1+idx2 is not idx2 assert idx1*idx2 is not idx2*idx1 + def test_uop_gcd_method(self): + a = Variable("a", 0, 8) + b = Variable("b", 0, 8) + self.assertEqual(UOp.gcd(a, a*b, a*3).simplify(), a) + self.assertEqual(UOp.gcd(a*a*a, a*b*a, a*3*a).simplify(), a*a) + self.assertEqual(UOp.gcd(a*a*10, b*a*5, a*a*5).simplify(), a*5) + self.assertEqual(UOp.gcd(a*10, b*5, a*5).simplify(), a.const_like(5)) + self.assertEqual(UOp.gcd(a, b*5, a*5).simplify(), a.const_like(1)) + + def test_divides_exact(self): + a = Variable("a", 1, 8) + b = Variable("b", 1, 8) + self.assertEqual((a*a*3).divide_exact(a).simplify(), a*3) + self.assertEqual((a*a*3).divide_exact(a*a*3).simplify(), a.const_like(1)) + self.assertEqual((a*b*3).divide_exact(a.const_like(3)).simplify(), a*b) + self.assertEqual((a*a*3).divide_exact(a*a.const_like(-3)).simplify(), a*-1) + self.assertEqual((a*a*b*3).divide_exact(a*b).simplify(), a*3) + self.assertEqual((a*3+a*b).divide_exact(a).simplify(), b+3) + self.assertEqual((a*b*3+a*b*b).divide_exact(a*b).simplify(), b+3) + self.assertEqual((((a*-2)+14)*b).divide_exact(((a*-2)+14)).simplify(), b) + + def test_divide_exact_not(self): + a = Variable("a", 1, 8) + b = Variable("b", 1, 8) + x = Variable("x", -20, 0) + self.assertEqual((a).divide_exact(b), None) + self.assertEqual((a+2).divide_exact(a), None) + self.assertEqual((x*-1).divide_exact(a), None) + self.assertEqual((a*5).divide_exact(a*10), None) + self.assertEqual((a*10-1).divide_exact(a*10), None) + def test_factorize(self): a = Variable("a", 0, 8) b = Variable("b", 0, 8) @@ -110,7 +141,7 @@ class TestSymbolic(unittest.TestCase): self.helper_test_variable(-Variable("a", 0, 8), -8, 0, "(a*-1)") def test_xor_0(self): - self.helper_test_variable(Variable("a", 0, 8, dtypes.int) ^ 0, 0, 8, "a") + self.helper_test_variable(Variable("a", 0, 8, dtypes.int) ^ 0, 0, 8, "a", test_z3=False) def test_add_1(self): self.helper_test_variable(Variable("a", 0, 8)+1, 1, 9, "(a+1)") @@ -374,6 +405,16 @@ class TestSymbolic(unittest.TestCase): def test_and_remove(self): self.helper_test_variable(uand([uconst(1), Variable("a", 0, 1)]), 0, 1, "a") + def test_bool_or_not_tautology(self): + a = Variable("a", 0, 10) + c = a<10 + self.helper_test_variable(c | c.logical_not(), True, True, "True") + + def test_bool_and_not_contradiction(self): + a = Variable("a", 0, 10) + c = a<10 + self.helper_test_variable(c & c.logical_not(), False, False, "False") + def test_mod_factor_negative(self): self.helper_test_variable(usum([uconst(-29), Variable("a", 0, 10), Variable("b", 0, 10)*28]) % 28, -27, 27, "(((a+(b*28))+-29)%28)") self.helper_test_variable(usum([uconst(-29), Variable("a", 0, 100), Variable("b", 0, 10)*28]) % 28, -27, 27, "(((a+(b*28))+-29)%28)") @@ -440,6 +481,33 @@ class TestSymbolic(unittest.TestCase): def test_mul_div_factor_div_neg(self): self.helper_test_variable((Variable("a", 0, 10)*-4+4)//8, -4, 0, "(((a*-1)+1)//2)") + def test_div_symbolic_const_gcd(self): + a = Variable("a", -10, 10) + b = Variable("b", -10, 10) + d = Variable("d", 1, 10) + self.helper_test_variable((3*a+9*b)//(3*d), -40, 40, "((a+(b*3))//d)") + + def test_symbolic_gcd_div(self): + a = Variable("a", -10, 10) + b = Variable("b", -10, 10) + c = Variable("c", -10, 10) + d1 = Variable("d1", 1, 10) + d2 = Variable("d2", -10, -1) + self.helper_test_variable((d1*a*b*d1)//(d1), -1000, 1000, "(a*(b*d1))") + self.helper_test_variable((d1*a*d2*b*d1)//(d1*d2), -1000, 1000, "(a*(b*d1))") + self.helper_test_variable((d1*a + b*d1)//(d1), -20, 20, "(a+b)") + self.helper_test_variable((d1*a + b*d1 + c*d1)//(d1), -30, 30, "(c+(a+b))") + self.helper_test_variable((3*a*d1 + 9*b*d1)//(3*d1*d2), -40, 40, "(((a+(b*3))//(d2*-1))*-1)") + self.helper_test_variable((3*a*d1 + 9*b*d1+3)//(3*d1*d2), -401, 399, "(((((a*d1)+((b*d1)*3))+1)//((d1*d2)*-1))*-1)") + + def test_symbolic_factor_remainder_div(self): + a = Variable("a", 0, 10) + b = Variable("b", 0, 10) + d = Variable("d", 1, 10) + self.helper_test_variable((d*a+b)//d, 0, 20, "(a+(b//d))") + self.helper_test_variable((d*a*20+b)//(5*d), 0, 42, "((a*4)+(b//(d*5)))") + self.helper_test_variable((d*a*20+b*d*5+10)//(5*d), 0, 52, "((b+(a*4))+(2//d))") + def test_mod_gcd_factor_neg(self): self.helper_test_variable((Variable("a", 0, 10)*-4+4)%8, -4, 4, "((((a*-1)+1)%2)*4)") @@ -510,6 +578,13 @@ class TestSymbolic(unittest.TestCase): self.helper_test_variable((gidx0*4+lidx2*2+lidx3)//12, 0, 4, ("(((lidx2//2)+gidx0)//3)", "((gidx0+(lidx2//2))//3)")) self.helper_test_variable((lidx2*2+gidx0*4+lidx3)//12, 0, 4, ("(((lidx2//2)+gidx0)//3)", "((gidx0+(lidx2//2))//3)")) + @unittest.expectedFailure # TODO: improve nest_div_by_smallest_factor + def test_sum_div_complex4(self): + gidx0 = Variable("gidx0", 0, 2) + lidx2 = Variable("lidx2", 0, 12) + lidx3 = Variable("lidx3", 0, 12) + self.helper_test_variable((gidx0*3+lidx2*19+lidx3*38)//(3*19), 0, 12, ("((lidx2+(lidx3*2))//3)")) + def test_sum_mul_distribute(self): gidx0 = Variable("gidx0", 0, 7) lidx2 = Variable("lidx2", 0, 12) @@ -581,45 +656,6 @@ class TestSymbolic(unittest.TestCase): with self.assertRaises(AssertionError): self.helper_test_variable((30 * b + 1) % 18 + ((30 * b + 1) // 18) * 18, 1, 3001, "((b*30)+1)") - def test_arange_unrolled4(self): - gidx = Variable("gidx", 0, 2559) - unrolled_div = (gidx+2561)//4+(gidx+2562)//4+(gidx+2560)//4+(gidx+2559)//4 - self.helper_test_variable(unrolled_div, 2559, 5118, "(gidx+2559)") - - def test_arange_unrolled4_with_cast(self): - gidx = Variable("gidx", 0, 2559, dtypes.index) - dt = dtypes.int - unrolled_div = ((gidx+2561)//4 + 2).cast(dt)+((gidx+2562)//4).cast(dt)+((gidx+2560)//4).cast(dt)+((gidx+2559)//4).cast(dt) - self.helper_test_variable(unrolled_div, 2561, 5120, "((int)(gidx)+2561)") - - def test_arange_unrolled4_mul(self): - gidx = Variable("gidx", 0, 2559) - unrolled_div = 2*((gidx+2561)//4)+2*((gidx+2562)//4)+2*((gidx+2560)//4)+2*((gidx+2559)//4) - self.helper_test_variable(unrolled_div, 5118, 10236, "((gidx*2)+5118)") - - def test_arange_unrolled4_small(self): - gidx = Variable("gidx", 0, 3) - unrolled_div = (gidx)//4+(gidx+2)//4+(gidx+3)//4+(gidx+1)//4 - self.helper_test_variable(unrolled_div, 0, 3, "gidx") - - gidx = Variable("gidx", 0, 2) - unrolled_div = (gidx)//4+(gidx+2)//4+(gidx+3)//4+(gidx+1)//4 - self.helper_test_variable(unrolled_div, 0, 2, "gidx") - - gidx = Variable("gidx", 0, 1) - unrolled_div = (gidx)//4+(gidx+2)//4+(gidx+3)//4+(gidx+1)//4 - self.helper_test_variable(unrolled_div, 0, 1, "gidx") - - def test_arange_unrolled2(self): - gidx = Variable("gidx", 0, 2559) - unrolled_div = (gidx+2559)//2+(gidx+2560)//2+3 - self.helper_test_variable(unrolled_div, 2562, 5121, "(gidx+2562)") - - def test_arange_unrolled2_neg(self): - ridx = Variable("ridx", 0, 255) - unrolled_div = -((255-ridx)//2) - ((256-ridx)//2) - self.helper_test_variable(unrolled_div, -255, 0, "(ridx+-255)") - def test_gated_load(self): idx = Variable("idx", 0, 24) self.helper_test_variable(idx//4, 0, 6, "(idx//4)") @@ -924,6 +960,46 @@ class TestSymbolicSymbolicOps(unittest.TestCase): assert c == uconst(2) """ +class TestInvalidIndex(unittest.TestCase): + def test_invalid_times_0(self): + ridx = Variable("ridx", 0, 10) + idx = (ridx<5).where(ridx, UOp.invalid())*0 + self.assertIs(idx.simplify(), (ridx<5).where(0, UOp.invalid()), "multiplying an index by 0 should preserve the invalid") + + def test_invalid_comparison_drops_invalid(self): + # comparisons return a bool, and bools can't be invalid + ridx = Variable("ridx", 0, 10) + idx = (ridx<5).where(ridx, UOp.invalid())<3 + self.assertIs(idx.simplify(), (ridx<3), "comparison of index should drop the invalid") + self.assertIs(idx.where(UOp.const(dtypes.int, 1), 0).simplify(), (ridx<3).where(UOp.const(dtypes.int, 1), 0), + "comparison of index should drop the invalid") + + def test_alu_moves_inside_invalid(self): + ridx = Variable("ridx", 0, 10) + idx = (ridx<5).where(ridx, UOp.invalid())*10 + self.assertIs(idx.simplify(), (ridx<5).where(ridx*10, UOp.invalid()), "multiplying an index by 0 should preserve the invalid") + + def test_merge_invalid_conditions(self): + ridx0 = Variable("ridx0", 0, 10) + ridx1 = Variable("ridx1", 0, 10) + idx0 = (ridx0<5).where(ridx0, UOp.invalid()) + idx1 = (ridx1<5).where(idx0//2, UOp.invalid()) + self.assertIs(idx1.simplify(), ((ridx1<5)&(ridx0<5)).where(ridx0//2, UOp.invalid()), + "valid inside a valid should make a single valid and & the conditions") + + def test_alu_invalid(self): + self.assertIs((UOp.invalid()*2).simplify(), UOp.invalid()) + self.assertIs((UOp.invalid()*0).simplify(), UOp.invalid()) + self.assertIs((UOp.invalid()+8).simplify(), UOp.invalid()) + self.assertIs((UOp.invalid()+Variable("a",0,10)).simplify(), UOp.invalid()) + self.assertIs((UOp.invalid()*Variable("a",0,10)).simplify(), UOp.invalid()) + self.assertIs((UOp.invalid() UOp|None: - if (idx:=uop_given_valid(valid, start_idx)) is None: return buf.const_like(0) - if not isinstance(buf.dtype, ImageDType): return None if idx is start_idx else buf.index(idx, valid) + if (idx:=uop_given_valid(valid, start_idx)) is None: return buf.index(UOp.invalid()) + if not isinstance(buf.dtype, ImageDType): return None if idx is start_idx else buf.index(idx.valid(valid)) # wait for it to be image indexed before running simplification if start_idx.dtype.count != 2: return None @@ -43,7 +43,7 @@ def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None: if not drop_stmt and idx is start_idx: return None new_valid = functools.reduce(operator.and_, ss) if (ss:=[s for s in valid.split_uop(Ops.AND) if s not in drop_stmt]) else None - return buf.index(idx, new_valid) + return buf.index(idx.valid(new_valid) if new_valid is not None else idx) def delete_redundant_gates(store:UOp, buf:UOp, idx:UOp, val:UOp, store_gate:UOp, cast:UOp|None=None) -> UOp|None: if store_gate not in [gate.src[0] for gate in val.toposort() if gate.op is Ops.IF]: return None @@ -52,12 +52,9 @@ def delete_redundant_gates(store:UOp, buf:UOp, idx:UOp, val:UOp, store_gate:UOp, load_store_indexing = PatternMatcher([ # image load valid idx simplification - (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("start_idx"), UPat.var("valid"))), simplify_valid_load), - # index True is just Index - (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("start_idx"), UPat(Ops.CONST, arg=True))), lambda buf,start_idx: buf.index(start_idx)), - # remove hanging cast - (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx", dtypes.int).cast()),), lambda buf,idx: buf.index(idx)), - (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx", dtypes.int).cast(), UPat.var("valid"))), lambda buf,idx,valid: buf.index(idx, valid)), + (UPat(Ops.INDEX, src=(UPat.var("buf"), invalid_gate)), lambda buf,x,i,cond: simplify_valid_load(buf, x, cond)), + # drop true gate + (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("x"), UPat.const(dtypes.bool, True)),), lambda buf,x: buf.index(x)), # delete_redundant_gates (after expand) (UPat(Ops.STORE, src=(UPat.any(stidx:=UPat.var("buf").index(UPat.var("idx"), UPat.var("store_gate")), stidx.cast().named("cast")), UPat.var("val")), name="store", allow_any_len=True), delete_redundant_gates), @@ -65,20 +62,21 @@ load_store_indexing = PatternMatcher([ # ***** load/store grouping ***** -def expand_index(buf:UOp, vec:UOp, mask:UOp|None=None): - if getenv("UNSAFE_DISABLE_MASK", 0): mask = None +def expand_index(buf:UOp, vec:UOp): + if getenv("UNSAFE_DISABLE_MASK", 0): vec = vec.get_idx() # generate the individual indexes - midx = graph_rewrite(UOp.sink(*[buf.index(vec.gep(i), mask.gep(i) if mask is not None else None) for i in range(vec.dtype.count)]), + midx = graph_rewrite(UOp.sink(*[buf.index(vec.gep(i)) for i in range(vec.dtype.count)]), symbolic_flat+load_store_indexing, name=f"index_buf_{buf.arg}") # extract all the relevant offsets offsets_rootsrc: defaultdict[Any, dict[int, list[int]]] = defaultdict(dict) for i in range(vec.dtype.count): - idx: Any = midx.src[i].src[1] + idx: Any = midx.src[i].src[1].get_idx() if idx.op is Ops.ADD and idx.src[1].op is Ops.CONST: root_src, arg = idx.src[0], idx.src[1].arg elif idx.op is Ops.ADD and idx.src[0].op is Ops.CONST: root_src, arg = idx.src[1], idx.src[0].arg + elif idx.op is Ops.CONST and idx.arg is Invalid: root_src, arg = "INVALID", 0 elif idx.op is Ops.CONST: root_src, arg = "CONST", idx.arg else: root_src, arg = idx, 0 - if len(midx.src[i].src) == 3: root_src = (midx.src[i].src[2], root_src) + root_src = (midx.src[i].src[1].get_valid(), root_src) offsets_rootsrc[root_src].setdefault(arg, []).append(i) # then rewrite everything we can into groups @@ -121,8 +119,6 @@ def gep_on_store(gep:UOp, st:UOp, sto:UOp): load_store_folding = PatternMatcher([ (UPat(Ops.INDEX, src=(UPat(Ops.VECTORIZE, src=UPat(GroupOp.Defines, name="buf")), UPat.var("vec"))), expand_index), - (UPat(Ops.INDEX, src=(UPat(Ops.VECTORIZE, src=UPat(GroupOp.Defines, name="buf")), UPat.var("vec"), - UPat.var("mask"))), expand_index), # GEP after LOAD (UPat(Ops.LOAD, src=(UPat(Ops.GEP, name="gep"),), name="ld", allow_any_len=True), lambda gep, ld: ld.replace(dtype=ld.dtype.scalar().vec(gep.dtype.count), src=(gep.src[0],)+ld.src[1:]).gep(gep.arg)), @@ -162,7 +158,8 @@ def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp): lengths.append(1) # worst case, it's not folded # filter fold lengths that don't divide - if must_divide: lengths = [x for x in lengths if idx.src[1].divides(x) is not None] + offset, mask = idx.src[1].get_idx(), idx.src[1].get_valid() + if must_divide: lengths = [x for x in lengths if offset.divides(x) is not None] # split based on the fold lengths global_offset = 0 @@ -171,7 +168,7 @@ def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp): # with 1 at the end of the lengths list, this will always hit for fold_length in lengths: if global_offset+fold_length > sz: continue - lidx = buf.index(idx.src[1] + global_offset, idx.src[2] if len(idx.src) > 2 else None) + lidx = buf.index((offset + global_offset).valid(mask)) if fold_length > 1: lidx = lidx.cast(buf.ptrdtype.base.vec(fold_length).ptr(size=buf.ptrdtype.size, addrspace=buf.ptrdtype.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))))+ls.src[2:])) else: ret.append(ls.replace(src=(lidx,)+ls.src[1:], dtype=ls.dtype.scalar().vec(fold_length))) @@ -187,19 +184,20 @@ def image_fixup(ls:UOp): if ls.src[0].op is Ops.CAST and isinstance(image_dtype:=ls.src[0].src[0].dtype, ImageDType): assert ls.src[0].dtype.count == 4, "image must be casted to 4" idx = ls.src[0].src[0] - oidx = UOp(Ops.VECTORIZE, dtypes.int.vec(2), ((idx.src[1] // 4) % image_dtype.shape[1], (idx.src[1] // (4*image_dtype.shape[1])))) - idx = idx.replace(src=(idx.src[0], oidx)+idx.src[2:]) + x, valid = idx.src[1].get_idx(), idx.src[1].get_valid() + oidx = UOp(Ops.VECTORIZE, dtypes.index.vec(2), ((x // 4) % image_dtype.shape[1], (x // (4*image_dtype.shape[1])))) + idx = idx.replace(src=(idx.src[0], oidx.valid(valid))) return ls.replace(src=(idx,)+ls.src[1:]) # this is an unprocessed image without a cast, aka unfoldable image load. this doesn't work for stores - if isinstance(image_dtype:=ls.src[0].dtype, ImageDType) and ls.src[0].src[1].dtype != dtypes.int.vec(2): + if isinstance(image_dtype:=ls.src[0].dtype, ImageDType) and ls.src[0].src[1].get_idx().dtype != dtypes.index.vec(2): assert ls.op is Ops.LOAD, "if an image store isn't upcasted to 4, we can't store it" idx = ls.src[0] - id4 = idx.src[1] % 4 - oidx = UOp(Ops.VECTORIZE, dtypes.int.vec(2), ((idx.src[1] // 4) % image_dtype.shape[1], (idx.src[1] // (4*image_dtype.shape[1])))) - idx = idx.replace(src=(idx.src[0], oidx)+idx.src[2:]) + x, valid = idx.src[1].get_idx(), idx.src[1].get_valid() + oidx = UOp(Ops.VECTORIZE, dtypes.index.vec(2), ((x // 4) % image_dtype.shape[1], (x // (4*image_dtype.shape[1])))) + idx = idx.replace(src=(idx.src[0], oidx.valid(valid))) vec_load = ls.replace(dtype=ls.dtype.vec(4), src=(idx,)+ls.src[1:]) - return functools.reduce(lambda ret, i: id4.ne(i).where(ret, vec_load.gep(i)), range(4), ls.const_like(float('nan'))) + return functools.reduce(lambda ret, i: (x % 4).ne(i).where(ret, vec_load.gep(i)), range(4), ls.const_like(float('nan'))) return None @@ -226,6 +224,7 @@ def no_vectorized_wmma(wmma:UOp): def no_vectorized_alu(alu:UOp): if alu.dtype.vcount == 1: return None + if alu.op is Ops.WHERE and alu.src[2].arg is Invalid: return None # image load/store has cond.where(idx.vec(2), Invalid) as the index alus = tuple(UOp(alu.op, alu.dtype.scalar(), tuple(s.gep(i) for s in alu.src), alu.arg) for i in range(alu.dtype.vcount)) return UOp(Ops.VECTORIZE, alu.dtype, alus) @@ -235,7 +234,7 @@ def no_vectorized_buf(buf:UOp): def no_vectorized_index(buf:UOp, cast:UOp, idx:UOp): cnt = cast.dtype.count assert idx.dtype.count == 1, f"idx dtype must be 1 {idx.dtype}" - return buf.broadcast(cnt).index(idx.broadcast(cnt)*cnt+UOp.const(dtypes.int.vec(cnt), tuple(range(cnt)))) + return buf.broadcast(cnt).index(idx.broadcast(cnt)*cnt+UOp.const(dtypes.index.vec(cnt), tuple(range(cnt)))) devectorize = PatternMatcher([ # no ALU on vectorized dtypes @@ -255,7 +254,8 @@ pm_render = PatternMatcher([ (UPat(Ops.VECTORIZE, src=(UPat(name='x'),)), lambda x: x), # give any loads that are masked an alt value (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat())).or_casted(),), allow_any_len=True, name="x"), - lambda x: x.replace(src=(x.src[0], x.const_like(0))+x.src[1:]) if len(x.src) == 1 or x.src[1].op is Ops.CUSTOM else None), + lambda x: x.replace(src=(x.src[0], x.const_like(0))+x.src[1:]) + if len(x.src) == 1 or x.src[1].op in (Ops.CUSTOM, Ops.STORE, Ops.BARRIER) else None), # gate any stores that aren't gated with ifs (UPat(Ops.STORE, src=(UPat(src=(UPat(), UPat(), UPat(dtype=dtypes.bool)), name="idx").or_casted(), UPat()), name="store", allow_any_len=True), lambda store,idx: UOp(Ops.STORE, dtype=store.dtype, src=store.src[:2]+(UOp(Ops.IF, src=(idx.src[2],)),)+store.src[2:]) if \ @@ -293,94 +293,7 @@ def reduce_to_acc(ctx:ReduceContext, red:UOp): ret = functools.reduce(lambda x,y: x.alu(red.arg, y), lst) return acc.load(acc.store(ret, *reduce_range)) if len(reduce_range) != 0 else ret -def no_vectorized_reduce(inp:UOp, red:UOp): - if inp.dtype != red.dtype: - red = red.replace(src=(functools.reduce(lambda x,y: x.alu(red.arg, y), horizontal_reduce(inp, red.dtype)),)+red.src[1:]) - if red.dtype.vcount == 1: return red - # no_vectorize_alu ignoring ranges - if red.dtype.vcount == 1: return None - alus = tuple(UOp(red.op, red.dtype.scalar(), (red.src[0].gep(i),)+red.src[1:], red.arg) for i in range(red.dtype.vcount)) - return UOp(Ops.VECTORIZE, red.dtype, alus) - -def reduce_rangeless(red:UOp): - # TODO: share code with reduce_unparented - if red.arg not in {Ops.ADD, Ops.MAX}: return None - if red.src[0].dtype != red.dtype: return None - if any(x.op in {Ops.RANGE} for x in red.src[0].toposort()): return None - ret = red.src[0] - if red.arg is Ops.ADD: - for r in red.src[1:]: - ret = ret * r.src[0].cast(ret.dtype.scalar()).broadcast(ret.dtype.count) - return ret - -def no_range(u:UOp) -> bool: return not any(x.op is Ops.RANGE for x in u.sparents) - -pm_reduce_collapse = PatternMatcher([ - # lift x+y out of reduce on lt - ((UPat.var("x")+UPat.var("y")).or_casted() < UPat.var("c"), lambda x,y,c: (x < (c.cast(y.dtype)-y)) if no_range(y) and no_range(c) else None), - # lift x*y out of reduce - ((UPat.var("x")*UPat.var("y")) < UPat.var("c"), - lambda x,y,c: (x < ((c+y-1) // y)) if no_range(y) and no_range(c) and y.vmin > 0 else None), - # lift x+y out of reduce on ne - ((UPat.var("x")+UPat.var("y")).or_casted() != UPat.var("c"), lambda x,y,c: (x != (c.cast(y.dtype)-y)) if no_range(y) and no_range(c) else None), - # fold the range - ((UPat(Ops.RANGE, name="r") < UPat.var("cut")).where(0, UPat.cvar("val")).reduce(arg=Ops.ADD, allow_any_len=True), - lambda r,cut,val: (r.src[0]-cut).maximum(0).minimum(r.src[0]).cast(val.dtype) * val), - ((UPat(Ops.RANGE, name="r") < UPat.var("cut")).where(UPat.cvar("val"), 0).reduce(arg=Ops.ADD, allow_any_len=True), - lambda r,cut,val: cut.maximum(0).minimum(r.src[0]).cast(val.dtype) * val), - # REDUCE on ADD - ((UPat.var("x")+UPat.var("y")).reduce(arg=Ops.ADD, allow_any_len=True, name="r"), - lambda x,y,r: x.reduce(*r.src[1:], arg=Ops.ADD) + y.reduce(*r.src[1:],arg=Ops.ADD)), - # MUL casted bool - ((UPat.var("x") * UPat.var("gate", dtype=dtypes.bool).cast().or_broadcasted(name="b")), - lambda x,gate,b=None: gate.broadcast(x.dtype.count).where(x, 0) if b is not None else gate.where(x, 0)), - # WHERE on LOAD (works on max too) - (UPat.var("gate").where(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx"))).load(), 0).reduce(arg=Ops.ADD, allow_any_len=True), - lambda buf,idx,gate: buf.index(idx, gate).load()), - (UPat.var("gate").where(0, UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx"))).load()).reduce(arg=Ops.ADD, allow_any_len=True), - lambda buf,idx,gate: buf.index(idx, gate.logical_not()).load()), - # INDEX on RANGE / gated RANGE - (UPat.var("buf").index(UPat.var("expr"), UPat.var("idx").eq(UPat(Ops.RANGE, name="r").or_casted())), - lambda buf,r,idx,expr: buf.index(expr.substitute({r:idx.cast(r.dtype)}), (idx.cast(r.dtype) >= 0) & (idx.cast(r.dtype) < r.src[0]))), - # AND on WHERE - ((UPat.any(UPat(Ops.DEFINE_VAR, name="x"), UPat(Ops.DEFINE_VAR).gep(name="x")) & UPat.var("y")) \ - .where(UPat.cvar("c"), 0).reduce(arg=Ops.ADD, allow_any_len=True, name="r"), - lambda x,y,c,r: y.where(c, 0).reduce(*r.src[1:], arg=Ops.ADD)*x.cast(c.dtype)), - # remove REDUCEs that no longer have a RANGE in the src - (UPat(Ops.REDUCE, name="red"), reduce_rangeless), - # devectorize REDUCE - (UPat(Ops.VECTORIZE, name="inp").reduce(name="red", allow_any_len=True), no_vectorized_reduce), - # index/load/where. TODO: this is more aggressive than needed - (UPat((Ops.INDEX, Ops.LOAD, Ops.WHERE), name="alu"), no_vectorized_alu), -])+sym+cast_folding - -def reduce_collapse(red:UOp): - included, not_included = partition(red.parents, lambda x: any(y in x.sparents for y in red.src[1:])) - if any(x.op in {Ops.STORE, Ops.REDUCE} for x in included): return None - replaces: dict[UOp, UOp] = {} - for u in included: - for s in u.src: - if s in not_included and s not in replaces and s.op not in {Ops.CONST, Ops.VCONST, Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_VAR}: - replaces[s] = UOp(Ops.DEFINE_VAR, dtype=s.dtype, arg=(f'in{len(replaces)}', s.vmin, s.vmax)) - collapse_fxn = red.substitute(replaces) - sink = graph_rewrite(collapse_fxn, pm_reduce_collapse, name="reduce_collapse") - if any(x.op is Ops.RANGE for x in sink.toposort()): return None - return sink.substitute({v:k for k,v in replaces.items()}) - -def reduce_unparented(red:UOp): - if red.arg not in {Ops.ADD, Ops.MAX}: return None - reduce_parented, reduce_unparented = partition(red.src[1:], lambda x: x in red.src[0].sparents) - if len(reduce_unparented) == 0: return None - ret = red.replace(src=(red.src[0],)+tuple(reduce_parented)) if len(reduce_parented) or red.dtype != red.src[0].dtype else red.src[0] - if red.arg is Ops.ADD: - for r in reduce_unparented: ret = ret * r.src[0].cast(ret.dtype.scalar()).broadcast(ret.dtype.count) - return ret - pm_reduce = PatternMatcher([ - # remove any ranges from a REDUCE that aren't referenced in the reduce source - (UPat(Ops.REDUCE, name="red"), reduce_unparented), - # remove REDUCE without loads (generic arange opt / indexing). TODO: support multi range - (UPat(Ops.REDUCE, src=(UPat(), UPat()), name="red"), reduce_collapse), # REDUCE -> DEFINE_ACC+ASSIGN (UPat(Ops.REDUCE, name="red"), reduce_to_acc), # tensor core built in accumulate diff --git a/tinygrad/codegen/late/expander.py b/tinygrad/codegen/late/expander.py index b5c2228a7e..d2d0a41162 100644 --- a/tinygrad/codegen/late/expander.py +++ b/tinygrad/codegen/late/expander.py @@ -2,7 +2,8 @@ import functools, itertools, operator from tinygrad.dtype import dtypes, PtrDType, AddrSpace from tinygrad.helpers import AMX, dedup, flatten, all_same, prod, partition -from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, GroupOp, AxisType +from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, GroupOp, AxisType, range_start +from tinygrad.schedule.rangeify import BufferizeOpts def _expand_arg_to_idx(args:tuple[tuple[int, int], ...], rpk:dict[int, int]) -> int: idx, mul = 0, 1 @@ -49,7 +50,7 @@ def do_expand(root:UOp): if root.op is Ops.IF or src.op is Ops.IF: # for the first arg of IF, just pass them through ignoring UNROLLS new_srcs.append(src) - elif (root.op is Ops.STORE and i >= 2) or (root.op in {Ops.REDUCE, Ops.BUFFERIZE} and i >= 1) or (root.op is Ops.WMMA and i >= 3): + elif root.op in range_start and i >= range_start[root.op]: # for any range args of STORE/REDUCE, pass them through new_srcs.append(src) elif root.op is Ops.INDEX and i >= 1 and not isinstance(root.dtype, PtrDType): @@ -142,7 +143,7 @@ def fix_group_for_reduce(x:UOp): # do only the non grouped reduces early ret = x.replace(src=(x.src[0],)+tuple(reduce_r)) reduce_loop = [x.replace(arg=(x.arg[0]+100, AxisType.REDUCE)) for x in reduce_gfr] - buf = ret.bufferize(*upstream_locals, *reduce_gfr, arg=(AddrSpace.LOCAL, reduce_gfr[0].arg[0])).index(*upstream_locals, *reduce_loop) + buf = ret.bufferize(*upstream_locals, *reduce_gfr, arg=BufferizeOpts(reduce_gfr[0].arg[0], AddrSpace.LOCAL)).index(*upstream_locals, *reduce_loop) # gate with an if on the store + do the final reduce buf = UOp(Ops.IF, dtype=buf.dtype, src=(functools.reduce(operator.and_, [x.eq(0) for x in reduce_gfr]), buf)) diff --git a/tinygrad/codegen/late/linearize.py b/tinygrad/codegen/late/linearize.py index a8ef31f3e4..d860125adf 100644 --- a/tinygrad/codegen/late/linearize.py +++ b/tinygrad/codegen/late/linearize.py @@ -222,6 +222,8 @@ def remove_blockend(x:UOp): if late_ops[i].op is Ops.BARRIER and late_ops[i+1].op is Ops.BARRIER: late_ops[i+1] = UOp(Ops.NOOP) arg = BasicBlock(parent_block.arg.lst+tuple(late_ops), tuple([y for y in x.arg.ctx if y is not x.arg.end]), cnt=x.arg.cnt) return UOp(Ops.BLOCK, src=tuple(y for y in x.src if y is not parent_block)+parent_block.src, arg=arg) + # else the whole context ended by the blockend is already in this block and we can safely turn it into a block + return UOp(Ops.BLOCK, src=x.src, arg=BasicBlock(x.arg.lst, tuple([y for y in x.arg.ctx if y is not x.arg.end]), cnt=x.arg.cnt)) block_merge = PatternMatcher([ (UPat((Ops.BLOCK, Ops.BLOCKEND), name="x"), merge_block), diff --git a/tinygrad/codegen/lowerer.py b/tinygrad/codegen/lowerer.py index 663a9b0fa6..236aff36a4 100644 --- a/tinygrad/codegen/lowerer.py +++ b/tinygrad/codegen/lowerer.py @@ -38,8 +38,8 @@ def lower_store(ctx: IndexContext, x: UOp, buf: UOp): #assert x.src[1].shape == x.src[0].shape, f"shape mismatch on store {x.src[1].shape} != {x.src[0].shape}" new_idxs = shape_to_idx(x.src[0].shape, ctx.axis_types, ctx.start) - idx, valid = x.st_arg.to_indexed_uops(new_idxs) - used_idxs = [x for x in UOp.sink(idx, valid).toposort() if x in new_idxs] + idx = x.st_arg.to_valid_uop(new_idxs) + used_idxs = [x for x in idx.toposort() if x in new_idxs] real_new_idxs = [] for i in range(len(x.src[0].shape)): if new_idxs[i] in used_idxs or len(ctx.idxs) <= i: real_new_idxs.append(new_idxs[i]) @@ -47,7 +47,7 @@ def lower_store(ctx: IndexContext, x: UOp, buf: UOp): stored = subblock(ctx, real_new_idxs, x.src[1]) used_ranges = [x for x in used_idxs if x.op is Ops.RANGE] - return buf.index(idx, valid).store(stored, *used_ranges) + return buf.index(idx).store(stored, *used_ranges) def fixup_wmma(ctx:IndexContext, x:UOp): if x.tag is not None: return None @@ -71,9 +71,9 @@ pm_lowerer = PatternMatcher([ # consts and loads (UPat(Ops.VIEW, src=(UPat((Ops.CONST, Ops.DEFINE_VAR), name="c"),), name="view"), - lambda ctx,view,c: c if all(x.mask is None for x in view.arg.views) else view.arg.to_indexed_uops(ctx.idxs)[1].where(c, c.const_like(0))), + lambda ctx,view,c: c if all(x.mask is None for x in view.arg.views) else view.arg.to_valid_uop(ctx.idxs).get_valid().where(c, c.const_like(0))), (UPat(Ops.LOAD, src=(UPat.var("buf").view(),), allow_any_len=True, name="x"), - lambda ctx,buf,x: UOp(Ops.LOAD, x.dtype, (buf.index(*x.st_arg.to_indexed_uops(ctx.idxs)),)+x.src[1:])), + lambda ctx,buf,x: UOp(Ops.LOAD, x.dtype, (buf.index(x.st_arg.to_valid_uop(ctx.idxs)),)+x.src[1:])), # reduce/view_const (UPat(Ops.REDUCE_AXIS, name="x"), lower_reduce_axis), diff --git a/tinygrad/codegen/opt/heuristic.py b/tinygrad/codegen/opt/heuristic.py index 6362a13b37..a73234c4c7 100644 --- a/tinygrad/codegen/opt/heuristic.py +++ b/tinygrad/codegen/opt/heuristic.py @@ -53,7 +53,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler: if k.opts.has_local and getenv("MV",1) != 0 and (MV_BLOCKSIZE > 1 or MV_THREADS_PER_ROW > 1 or MV_ROWS_PER_THREAD > 1) and \ k.reduceop is not None and k.reduceop.arg[0] is Ops.ADD and len(k.full_shape) >= 2 and k.opts.has_shared and \ (mulop:=k.reduceop.src[0]).op is Ops.MUL and mulop.src[0].op is Ops.LOAD and mulop.src[1].op is Ops.LOAD: - idx0, idx1 = mulop.src[0].src[0].src[1], mulop.src[1].src[0].src[1] + idx0, idx1 = mulop.src[0].src[0].src[1].get_idx(), mulop.src[1].src[0].src[1].get_idx() first_reduce_rng = k.ranges_of(AxisType.REDUCE)[0] if any(u is first_reduce_rng for u in idx0.split_uop(Ops.ADD)) and all(r in idx1.ranges for r in idx0.ranges): for global_idx in k.axes_of(AxisType.GLOBAL): @@ -77,7 +77,8 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler: for buf_index,buf in enumerate(k.bufs): if isinstance(buf.src[0].dtype, ImageDType): # part of real_strides - unit_stride_axes_mul_4 = [k.rngs.index(c) for c in k.bufs[buf_index].src[1].split_uop(Ops.ADD) if c.op is Ops.RANGE and (c.vmax+1)%4 == 0] + unit_stride_axes_mul_4 = [k.rngs.index(c) for c in k.bufs[buf_index].src[1].get_idx().split_uop(Ops.ADD) if + c.op is Ops.RANGE and (c.vmax+1)%4 == 0] if len(unit_stride_axes_mul_4): if (axis:=unit_stride_axes_mul_4[0]) in k.upcastable_dims: k.apply_opt(Opt(OptOps.UPCAST, axis, 4)) @@ -94,8 +95,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler: # upcast leading axes first (hack-ish for winograd; we actually want to upcast masked axes with low stride first) for axis in k.upcastable_dims: # for Schedule, we check if the range is used in INDEX gates or WHERE gates - is_masked = any(len(st.src) > 2 and k.rngs[axis] in st.src[2].parents for st in k.bufs) or \ - any(any(o is k.rngs[axis] for o in u.src[0].parents) for u in k.ast.parents if u.op is Ops.WHERE) + is_masked = any(any(o is k.rngs[axis] for o in u.src[0].parents) for u in k.ast.parents if u.op is Ops.WHERE) if k.full_shape[axis] <= 7 and is_masked and prod(k.full_shape[j] for j in to_upcast) * k.full_shape[axis] <= 7 * 7: if DEBUG >= 4: print(f"upcasting masked axis : {axis}") to_upcast.append(axis) @@ -111,11 +111,13 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler: # if we haven't upcasted it, it mods, and buffer has stride 0 on axis while having no stride 0 in the upcasted axis already if axis in upcasted_axis or k.full_shape[axis]%upcast_amount != 0: continue rng = k.rngs[axis] - if any(rng not in b.src[1].parents and all(r2 in b.src[1].parents for r2 in k.ranges_of(AxisType.UPCAST, AxisType.UNROLL)) for b in k.bufs): + if any(rng not in b.src[1].get_idx().parents and all(r2 in b.src[1].get_idx().parents + for r2 in k.ranges_of(AxisType.UPCAST, AxisType.UNROLL)) for b in k.bufs): num_strides, sum_strides = 0, 0 for b in k.bufs: - if rng in b.src[1].parents: num_strides += 1 - for c in b.src[1].split_uop(Ops.ADD): + idx = b.src[1].get_idx() + if rng in idx.parents: num_strides += 1 + for c in idx.split_uop(Ops.ADD): if c is rng: sum_strides += 1 if c.op is Ops.MUL and c.src[0] is rng and c.src[1].op is Ops.CONST: sum_strides += c.src[1].arg if c.op is Ops.MUL and c.src[1] is rng and c.src[0].op is Ops.CONST: sum_strides += c.src[0].arg @@ -157,7 +159,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler: k.apply_opt(Opt(OptOps.NOLOCALS)) else: # prioritize making expand axes local - local_axis_ranking = [(any(k.rngs[axis] not in b.src[1].parents for b in k.bufs), axis) \ + local_axis_ranking = [(any(k.rngs[axis] not in b.src[1].get_idx().parents for b in k.bufs), axis) \ for axis in k.axes_of(AxisType.GLOBAL, AxisType.LOOP) if k.rngs[axis].src[0].op is Ops.CONST] to_local: list[tuple[int, int]] = [] for _, axis in sorted(local_axis_ranking, key=lambda x: (-x[0], -x[1])): diff --git a/tinygrad/codegen/opt/postrange.py b/tinygrad/codegen/opt/postrange.py index f5e7a6af3d..4f234e047c 100644 --- a/tinygrad/codegen/opt/postrange.py +++ b/tinygrad/codegen/opt/postrange.py @@ -2,14 +2,15 @@ from __future__ import annotations import math, itertools from collections import defaultdict from typing import cast, Final -from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, can_pad +from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, can_pad, GroupOp from tinygrad.device import Buffer from tinygrad.dtype import AddrSpace, dtypes, ImageDType -from tinygrad.helpers import colored, BEAM, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, prod +from tinygrad.helpers import colored, BEAM, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, prod, merge_dicts from tinygrad.codegen.opt import axis_colors, Opt, OptOps, KernelOptError, check, axis_letters from tinygrad.codegen.simplify import pm_flatten_range from tinygrad.renderer import Renderer -from tinygrad.schedule.rangeify import remove_tags + +remove_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)]) # NOTE: LOCAL and GROUP_REDUCE have the same priority. the order here matters axis_to_pos = {AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2, AxisType.UPCAST: 3, @@ -139,6 +140,11 @@ class Scheduler: upcast_local_sz = prod([self.full_shape[a] for a in self.axes_of(AxisType.UPCAST, AxisType.WARP, AxisType.LOCAL, AxisType.GROUP_REDUCE)]) smem_sz = amt*upcast_local_sz*self.reduceop.dtype.itemsize check(smem_sz <= self.opts.shared_max, f"exceeds maximum shared memory size: needs {smem_sz}, max {self.opts.shared_max}") + if self.reduceop is not None and (opt.op in {OptOps.GROUP, OptOps.GROUPTOP}): + # We currently dont support a group within another rudece, TODO: fix if-contexts + reduce = [u for u in self.ast.parents if u.op is Ops.REDUCE and rng in merge_dicts([r.ranges for r in u.src[1:]])][0] + check(not any(u.arg[-1] in (AxisType.REDUCE, AxisType.UNROLL, AxisType.GROUP_REDUCE) for u in reduce.ranges), + "cannot have a GROUP_REDUCE inside another reduce") if opt.op is OptOps.UNROLL: check(amt <= 32, "don't unroll more than 32") @@ -166,7 +172,8 @@ class Scheduler: check(-1 <= (tc_select:=cast(tuple, opt.arg)[0]) < len(self.opts.tensor_cores), "tensor core opts must have valid tc_select") check(0 <= (tc_opt:=cast(tuple, opt.arg)[1]) <= 2, "tensor core opts must have valid tc_opt") check(0 < (use_tensor_cores:=cast(tuple, opt.arg)[2]) <= 2, "use_tensor_cores value is not valid") - ret = self._apply_tc_opt(use_tensor_cores, cast(int, opt.axis), tc_select, tc_opt) + try: ret = self._apply_tc_opt(use_tensor_cores, cast(int, opt.axis), tc_select, tc_opt) + except ValueError as e: raise KernelOptError(str(e)) check(ret is not None, "no tensor core available") elif opt.op is OptOps.PADTO: check(rng.src[0].op is Ops.CONST, "only pad const axes") @@ -179,11 +186,10 @@ class Scheduler: check(rng.vmax+1 > new_sz//4, "pad adds more than quadruple the work") replaced_rng = UOp.range(new_sz, *rng.arg) replaces = {rng:replaced_rng} + valid = replaced_rng < rng.vmax+1 for b in self.bufs: - if rng in b.src[1].sparents: - valid = replaced_rng < rng.vmax+1 - if len(b.src) > 2: valid = b.src[2] & valid - replaces[b] = b.replace(src=b.src[0:2]+(valid,)) + if rng in (i:=b.src[1].get_idx()).sparents: + replaces[b] = b.replace(src=(b.src[0],(valid&b.src[1].get_valid()).where(i, UOp.invalid()))) self.ast = self.ast.substitute(replaces, f"padto {rng.arg[:-1]} {opt.arg}") elif opt.op is OptOps.SWAP: try: @@ -315,18 +321,17 @@ def apply_opts(ctx:Renderer, ast:UOp): if ast.tag is not None: return None k = Scheduler(ast, ctx) k.convert_loop_to_global() - if BEAM >= 1: + if ast.arg is not None and ast.arg.opts_to_apply is not None: + for opt in ast.arg.opts_to_apply: k.apply_opt(opt) + elif BEAM >= 1: from tinygrad.codegen.opt.search import beam_search rawbufs = bufs_from_ast(ast, ctx.device) k = beam_search(k, rawbufs, BEAM.value, bool(getenv("BEAM_ESTIMATE", 1))) - elif ast.arg is not None and ast.arg.opts_to_apply is not None: - for opt in ast.arg.opts_to_apply: k.apply_opt(opt) elif not NOOPT and (ast.arg is None or ast.arg.applied_opts == ()): from tinygrad.codegen.opt.heuristic import hand_coded_optimizations # NOTE: hand_coded_optimizations doesn't support multiblock opts yet if all(len(u.src) == 1 for u in ast.parents if u.op is Ops.LOAD): - # TODO: why is the returned k from hand_coded_optimizations different than this? - for opt in hand_coded_optimizations(k).applied_opts: k.apply_opt(opt) + k = hand_coded_optimizations(k) return k.get_optimized_ast(name_override=ast.arg.name if ast.arg is not None and ast.arg.name != "test" else None) pm_postrange_opt = PatternMatcher([ diff --git a/tinygrad/codegen/opt/swizzler.py b/tinygrad/codegen/opt/swizzler.py index 9cc3567439..75521b8311 100644 --- a/tinygrad/codegen/opt/swizzler.py +++ b/tinygrad/codegen/opt/swizzler.py @@ -17,7 +17,7 @@ merge_views = PatternMatcher([ lambda view: view.const_like(0) if (mask:=view.st.views[-1].mask) is not None and any((x[1]-x[0]) == 0 for x in mask) else None), # only unmaksed VIEW on CONST replaces the ShapeTracker (UPat(Ops.VIEW, src=(UPat((Ops.CONST, Ops.DEFINE_VAR), name="x"),), name="view"), - lambda x,view: x.replace(src=(x.src[0].replace(arg=x.st+view.st),)) if all(v.mask is None for v in (x.st+view.st).views) else None), + lambda x,view: x.replace(src=(UOp(Ops.VIEW, x.dtype, x.src, view.arg),)) if all(v.mask is None for v in view.st.views) else None), ]) def reduce_push_add_ones(src:UOp, r:UOp, view:UOp): diff --git a/tinygrad/codegen/simplify.py b/tinygrad/codegen/simplify.py index bdf6b39d1c..c1bdeff186 100644 --- a/tinygrad/codegen/simplify.py +++ b/tinygrad/codegen/simplify.py @@ -1,8 +1,10 @@ -from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, graph_rewrite, _substitute -from tinygrad.uop.symbolic import symbolic_flat +from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, graph_rewrite, _substitute, range_start +from tinygrad.uop.symbolic import symbolic_flat, sym, invalid_pat +from tinygrad.helpers import partition +from tinygrad.dtype import dtypes def flatten_range(r:UOp): - off = 2 if r.op is Ops.STORE else 1 + off = range_start[r.op] rngs = r.src[off:] if not len(rngs): return None new_rngs = [x for x in UOp.sink(*rngs).toposort() if x.op is Ops.RANGE] @@ -15,7 +17,7 @@ pm_flatten_range = PatternMatcher([ def count_divmod(x:UOp): return len([u for u in x.toposort() if u.op in {Ops.IDIV, Ops.MOD}]) def simplify_merge_adjacent(u:UOp) -> UOp|None: - i = 2 if u.op is Ops.STORE else 1 + i = range_start[u.op] while i < len(u.src)-1: r0, r1 = u.src[i], u.src[i+1] # check same type @@ -34,4 +36,84 @@ def simplify_merge_adjacent(u:UOp) -> UOp|None: pm_simplify_ranges = PatternMatcher([ (UPat((Ops.STORE, Ops.REDUCE), name="u"), simplify_merge_adjacent), -]) \ No newline at end of file +]) + +# **** reduce simplification **** + +def no_range(u:UOp) -> bool: return not any(x.op is Ops.RANGE for x in u.sparents) + +def reduce_rangeless(red:UOp): + # TODO: share code with reduce_unparented + if red.arg not in {Ops.ADD, Ops.MAX}: return None + if red.src[0].dtype != red.dtype: return None + if not no_range(red.src[0]): return None + ret = red.src[0] + if red.arg is Ops.ADD: + for r in red.src[1:]: + ret = ret * r.src[0].cast(ret.dtype.scalar()).broadcast(ret.dtype.count) + return ret + +pm_reduce_collapse = PatternMatcher([ + # lift x+y out of reduce on lt + ((UPat.var("x")+UPat.var("y")).or_casted() < UPat.var("c"), lambda x,y,c: (x < (c.cast(y.dtype)-y)) if no_range(y) and no_range(c) else None), + # lift x*y out of reduce + ((UPat.var("x")*UPat.var("y")) < UPat.var("c"), + lambda x,y,c: (x < ((c+y-1) // y)) if no_range(y) and no_range(c) and y.vmin > 0 else None), + # lift x+y out of reduce on ne + ((UPat.var("x")+UPat.var("y")).or_casted() != UPat.var("c"), lambda x,y,c: (x != (c.cast(y.dtype)-y)) if no_range(y) and no_range(c) else None), + # fold the range + ((UPat(Ops.RANGE, name="r") < UPat.var("cut")).where(0, UPat.cvar("val")).reduce(arg=Ops.ADD, allow_any_len=True), + lambda r,cut,val: (r.src[0]-cut).maximum(0).minimum(r.src[0]).cast(val.dtype) * val), + ((UPat(Ops.RANGE, name="r") < UPat.var("cut")).where(UPat.cvar("val"), 0).reduce(arg=Ops.ADD, allow_any_len=True), + lambda r,cut,val: cut.maximum(0).minimum(r.src[0]).cast(val.dtype) * val), + # REDUCE on ADD + ((UPat.var("x")+UPat.var("y")).reduce(arg=Ops.ADD, allow_any_len=True, name="r"), + lambda x,y,r: x.reduce(*r.src[1:], arg=Ops.ADD) + y.reduce(*r.src[1:],arg=Ops.ADD)), + # MUL casted bool + ((UPat.var("x") * UPat.var("gate", dtype=dtypes.bool).cast().or_broadcasted(name="b")), + lambda x,gate,b=None: gate.broadcast(x.dtype.count).where(x, 0) if b is not None else gate.where(x, 0)), + # WHERE on LOAD (works on max too) + (UPat.var("gate").where(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx"))).load(), 0).reduce(arg=Ops.ADD, allow_any_len=True), + lambda buf,idx,gate: buf.index(idx.valid(gate)).load()), + (UPat.var("gate").where(0, UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx"))).load()).reduce(arg=Ops.ADD, allow_any_len=True), + lambda buf,idx,gate: buf.index(idx.valid(gate.logical_not())).load()), + # INDEX on RANGE / gated RANGE + (UPat.var("buf").index(UPat.var("idx").eq(UPat(Ops.RANGE, name="r").or_casted()).where(UPat.var("expr"), invalid_pat)), + lambda buf,r,idx,expr,i: buf.index(expr.substitute({r:idx.cast(r.dtype)}).valid((idx.cast(r.dtype) >= 0) & (idx.cast(r.dtype) < r.src[0])))), + # AND on WHERE + ((UPat.any(UPat(Ops.DEFINE_VAR, name="x"), UPat(Ops.DEFINE_VAR).gep(name="x")) & UPat.var("y")) \ + .where(UPat.cvar("c"), 0).reduce(arg=Ops.ADD, allow_any_len=True, name="r"), + lambda x,y,c,r: y.where(c, 0).reduce(*r.src[1:], arg=Ops.ADD)*x.cast(c.dtype)), + # remove REDUCEs that no longer have a RANGE in the src + (UPat(Ops.REDUCE, name="red"), reduce_rangeless), +])+sym + +def reduce_collapse(red:UOp): + included, not_included = partition(red.parents, lambda x: any(y in x.sparents for y in red.src[1:])) + if any(x.op in {Ops.STORE, Ops.REDUCE} for x in included): return None + replaces: dict[UOp, UOp] = {} + for u in included: + for s in u.src: + if s in not_included and s not in replaces and s.op not in {Ops.CONST, Ops.VCONST, Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_VAR}: + replaces[s] = UOp(Ops.DEFINE_VAR, dtype=s.dtype, arg=(f'in{len(replaces)}', s.vmin, s.vmax)) + collapse_fxn = red.substitute(replaces) + sink = graph_rewrite(collapse_fxn, pm_reduce_collapse, name="reduce_collapse") + return sink.substitute({v:k for k,v in replaces.items()}) if no_range(sink) else None + +def reduce_unparented(red:UOp): + if red.arg not in {Ops.ADD, Ops.MAX, Ops.MUL}: return None + reduce_parented, reduce_unparented = partition(red.src[1:], lambda x: x in red.src[0].sparents) + if len(reduce_unparented) == 0: return None + ret = red.replace(src=(red.src[0],)+tuple(reduce_parented)) if len(reduce_parented) or red.dtype != red.src[0].dtype else red.src[0] + if red.arg is Ops.ADD: + for r in reduce_unparented: ret = ret * r.src[0].cast(ret.dtype.scalar()).broadcast(ret.dtype.count) + if red.arg is Ops.MUL: + for r in reduce_unparented: ret = ret ** r.src[0].cast(ret.dtype.scalar()).broadcast(ret.dtype.count) + return ret + +pm_reduce_simplify = PatternMatcher([ + # remove any ranges from a REDUCE that aren't referenced in the reduce source + (UPat(Ops.REDUCE, name="red"), reduce_unparented), + # remove REDUCE without loads (generic arange opt / indexing). TODO: support multi range + (UPat(Ops.REDUCE, src=(UPat(), UPat()), name="red"), reduce_collapse), +]) diff --git a/tinygrad/device.py b/tinygrad/device.py index 0daaf0bffc..1cef7786bf 100644 --- a/tinygrad/device.py +++ b/tinygrad/device.py @@ -1,16 +1,17 @@ from __future__ import annotations from dataclasses import dataclass, replace from collections import defaultdict -from typing import Any, Generic, TypeVar, Iterator +from typing import Any, Generic, TypeVar, Iterator, Sequence, cast import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re, atexit, pickle, decimal -from tinygrad.helpers import CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored, CPU_LLVM, \ - Context, DISABLE_COMPILER_CACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, dedup +from tinygrad.helpers import CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored, CPU_LLVM +from tinygrad.helpers import Context, DISABLE_COMPILER_CACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, dedup +from tinygrad.helpers import unwrap_class_type from tinygrad.dtype import DType, ImageDType, PtrDType, dtypes, _to_np_dtype from tinygrad.renderer import Renderer # **************** Device **************** -ALL_DEVICES = ["METAL", "AMD", "NV", "CUDA", "QCOM", "GPU", "CPU", "DSP", "WEBGPU"] +ALL_DEVICES = ["METAL", "AMD", "NV", "CUDA", "QCOM", "CL", "CPU", "DSP", "WEBGPU"] class _Device: def __init__(self) -> None: self._devices = [x.stem[len("ops_"):].upper() for x in (pathlib.Path(__file__).parent/"runtime").iterdir() if x.stem.startswith("ops_")] @@ -272,12 +273,34 @@ class Compiler: return lib def disassemble(self, lib:bytes): pass +CompilerPairT = tuple[functools.partial|type[Renderer], functools.partial|type[Compiler]] class Compiled: profile_events:list[ProfileEvent] = [ProfileDeviceEvent("CPU")] # NOTE: CPU is the default device. - def __init__(self, device:str, allocator:Allocator, renderer:Renderer|None, compiler:Compiler|None, runtime, graph=None, group_id=None): - self.device, self.allocator, self.compiler, self.runtime, self.graph = device, allocator, compiler or Compiler(), runtime, graph - self.renderer, self.group_id = renderer or Renderer(), group_id + def __init__(self, device:str, allocator:Allocator, compilers:Sequence[CompilerPairT]|None, runtime, graph=None, group_id=None): + self.device, self.allocator, self.runtime, self.graph, self.group_id = device, allocator, runtime, graph, group_id + self.compilers = cast(list[CompilerPairT], compilers or [(Renderer, Compiler)]) + + envnames = [self._get_compiler_envvar(c) for r,c in self.compilers] + enable_comps = set((en, comp_pair) for en, comp_pair in zip(envnames, self.compilers) if en is not None and getenv(en, -1) == 1) + disable_comps = set((en, comp_pair) for en, comp_pair in zip(envnames, self.compilers) if en is not None and getenv(en, -1) == 0) + + if len(enable_comps) > 1: raise RuntimeError(f"{self.device}: multiple compilers set in env {enable_comps}") + for _, comp_pair in disable_comps: self.compilers.remove(comp_pair) + + try: self.renderer, self.compiler = next(self._get_available_compilers([list(enable_comps)[0][1]] if len(enable_comps) == 1 else self.compilers)) + except StopIteration as exc: raise RuntimeError(f"no usable compilers for {self.device}") from exc + + if DEBUG >= 1: print(f"{self.device}: using {self.compiler.__class__.__name__}") + + def _get_compiler_envvar(self, c): + compiler_name = f"{unwrap_class_type(c).__name__.upper().removesuffix('COMPILER').removeprefix(devname:=self.device.split(':')[0].upper())}" + return f"{devname}_{compiler_name if len(compiler_name) > 0 else unwrap_class_type(c).__name__.upper()}" + + def _get_available_compilers(self, compilers) -> Iterator[tuple[Renderer, Compiler]]: + for renderer, compiler in compilers: + with contextlib.suppress(Exception): yield renderer(), compiler() + def synchronize(self): """ Synchronize all pending operations on the device. @@ -302,12 +325,10 @@ def is_dtype_supported(dtype:DType, device:str|None=None) -> bool: if device is None: device = Device.DEFAULT if dtype == dtypes.bfloat16: if device == "METAL": return not CI - if device in {"CUDA", "NV"}: return not CI and not getenv("PTX") and not getenv("NIR") - if device in {"CPU"}: return not CI and platform.machine() in {"arm", "arm64", "aarch64", "x86_64", "amd64"} + if device in {"CUDA", "NV"}: return not CI and not getenv(f"{device}_PTX") and not getenv(f"{device}_NAK") + if device in {"CPU"}: return not CI and platform.machine() in {"arm", "arm64", "aarch64", "x86_64", "amd64"} and not getenv("CPU_LVP") return device in {"AMD", "PYTHON"} - if dtype in dtypes.fp8s: - # not supported yet - in progress - return False + if dtype in dtypes.fp8s: return device == "PYTHON" if 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 @@ -315,11 +336,11 @@ def is_dtype_supported(dtype:DType, device:str|None=None) -> bool: # 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: - if device == "GPU": return not CI and not OSX + if device == "CL": return not CI and not OSX if device in ["CUDA", "NV"]: return not CI if device == "CPU" and CPU_LLVM: return OSX if device == "PYTHON": return sys.version_info >= (3, 12) - if dtype == dtypes.float64: return device != "METAL" and not (OSX and device == "GPU") + if dtype == dtypes.float64: return device != "METAL" and not (OSX and device == "CL") return True if PROFILE: @@ -331,21 +352,26 @@ if PROFILE: with open(fn:=temp("profile.pkl", append_user=True), "wb") as f: pickle.dump(cpu_events+Compiled.profile_events+Buffer.profile_events, f) - if not getenv("SQTT", 0): - from tinygrad.uop.ops import launch_viz - launch_viz(PROFILE, fn) + from tinygrad.uop.ops import launch_viz + launch_viz("PROFILE", fn) if __name__ == "__main__": + from tinygrad import Tensor, Device + for device in ALL_DEVICES: + compilers_results, any_works = [], False try: - _ = Device[device].device - try: - from tinygrad import Tensor - with Context(CACHELEVEL=0): test = (Tensor([1,2,3], device=device) * 2).tolist() - if test != [2,4,6]: raise ValueError(f"got {test} instead of [2, 4, 6]") - result = colored("PASS", "green") - except Exception as e: - result = f"{colored('FAIL', 'yellow')} {e}" + default_compiler = (d:=Device[device]).compiler + for i,(r,c) in enumerate(d.compilers): + try: + d.renderer, d.compiler = r(), c() + with Context(CACHELEVEL=0): test = (Tensor([1,2,3], device=device) * 2).tolist() + if test != [2,4,6]: raise ValueError(f"got {test} instead of [2, 4, 6]") + default_text = '(default)' if type(default_compiler) is type(d.compiler) else f'({d._get_compiler_envvar(c)}=1 to make default)' + compilers_results.append(f"{colored('+', 'green')} {unwrap_class_type(c).__name__} {default_text}") + any_works = True + except Exception as e: compilers_results.append(f"{colored('-', 'yellow')} {unwrap_class_type(c).__name__}: {e}") + result = (colored('PASS', 'green') if any_works else f"{colored('FAIL', 'yellow')}") + ''.join([f'\n{" "*16} {x}' for x in compilers_results]) except Exception as e: result = f"{colored('FAIL', 'red')} {e}" print(f"{'*' if device == Device.DEFAULT else ' '} {device:10s}: {result}") diff --git a/tinygrad/dtype.py b/tinygrad/dtype.py index 37fe4608f6..11373bb3a8 100644 --- a/tinygrad/dtype.py +++ b/tinygrad/dtype.py @@ -5,6 +5,21 @@ from dataclasses import dataclass, fields from tinygrad.helpers import getenv, prod from enum import Enum, auto +class InvalidTypeMetaClass(type): + instance:None|InvalidType = None + def __call__(cls, *args, **kwargs): + if (ret:=InvalidTypeMetaClass.instance) is not None: return ret + InvalidTypeMetaClass.instance = ret = super().__call__() + return ret + +class InvalidType(metaclass=InvalidTypeMetaClass): + def __eq__(self, other): return self is other + def __hash__(self): return id(self) + def __repr__(self): return "Invalid" + def __reduce__(self): return (InvalidType, ()) # Return the global Invalid instance + +Invalid = InvalidType() + ConstType = float|int|bool FmtStr = Literal['?', 'b', 'B', 'h', 'H', 'i', 'I', 'q', 'Q', 'e', 'f', 'd'] @@ -17,7 +32,9 @@ class DTypeMetaClass(type): DTypeMetaClass.dcache[args] = ret = super().__call__(*args) return ret -class AddrSpace(Enum): GLOBAL = auto(); LOCAL = auto(); REG = auto() # noqa: E702 +class AddrSpace(Enum): + def __repr__(self): return str(self) + GLOBAL = auto(); LOCAL = auto(); REG = auto() # noqa: E702 @dataclass(frozen=True, eq=False) class DType(metaclass=DTypeMetaClass): @@ -104,10 +121,11 @@ class dtypes: if x.__class__ is list or x.__class__ is tuple: return max(dtypes.from_py(xi) for xi in x) if x else dtypes.default_float raise RuntimeError(f"Could not infer dtype of {x} with type {type(x)}") @staticmethod - def as_const(val: tuple[ConstType, ...]|ConstType, dtype:DType): + def as_const(val: tuple[ConstType|InvalidType, ...]|ConstType|InvalidType, dtype:DType): if isinstance(val, tuple): assert len(val) == dtype.count, f"mismatch {val} {dtype}" return tuple(dtypes.as_const(x, dtype) for x in val) + if isinstance(val, InvalidType): return val return int(val) if dtypes.is_int(dtype) else float(val) if dtypes.is_float(dtype) else bool(val) @staticmethod @functools.cache @@ -177,8 +195,8 @@ def to_dtype(dtype:DTypeLike) -> DType: return dtype if isinstance(dtype, DType) # https://jax.readthedocs.io/en/latest/jep/9407-type-promotion.html # we don't support weak type and complex type promo_lattice = { dtypes.bool: [dtypes.int8, dtypes.uint8], dtypes.int8: [dtypes.int16], dtypes.int16: [dtypes.int32], dtypes.int32: [dtypes.int64], - dtypes.int64: [dtypes.float16, dtypes.bfloat16], dtypes.uint8: [dtypes.int16, dtypes.uint16], dtypes.uint16: [dtypes.int32, dtypes.uint32], - dtypes.uint32: [dtypes.int64, dtypes.uint64], dtypes.uint64: [dtypes.float16, dtypes.bfloat16], + dtypes.int64: [dtypes.fp8e4m3, dtypes.fp8e5m2], dtypes.uint8: [dtypes.int16, dtypes.uint16], dtypes.uint16: [dtypes.int32, dtypes.uint32], + dtypes.uint32: [dtypes.int64, dtypes.uint64], dtypes.uint64: [dtypes.fp8e4m3, dtypes.fp8e5m2], dtypes.fp8e5m2: [dtypes.float16, dtypes.bfloat16], dtypes.fp8e4m3: [dtypes.float16, dtypes.bfloat16], dtypes.float16: [dtypes.float32], dtypes.bfloat16: [dtypes.float32], dtypes.float32: [dtypes.float64], } @@ -217,7 +235,7 @@ def sum_acc_dtype(dt:DType): if dtypes.is_int(dt) or dt == dtypes.bool: return least_upper_dtype(dt, dtypes.int) return least_upper_dtype(dt, to_dtype(getenv("SUM_DTYPE", "float32"))) -def truncate_fp16(x): +def float_to_fp16(x): try: return struct.unpack('e', struct.pack('e', float(x)))[0] except OverflowError: return math.copysign(math.inf, x) @@ -294,7 +312,7 @@ def fp8_to_float(x: int, dtype: DType) -> float: return float(float32_val) truncate: dict[DType, Callable] = {dtypes.bool: bool, - dtypes.float16: truncate_fp16, dtypes.bfloat16: lambda x: float_to_bf16(float(x)), + dtypes.float16: float_to_fp16, dtypes.bfloat16: lambda x: float_to_bf16(float(x)), **{fp8: (lambda x, dtype=fp8: fp8_to_float(float_to_fp8(x, dtype), dtype)) for fp8 in dtypes.fp8s}, dtypes.float32: lambda x: ctypes.c_float(x).value, dtypes.float64: lambda x: ctypes.c_double(x).value, dtypes.uint8: lambda x: ctypes.c_uint8(x).value, dtypes.uint16: lambda x: ctypes.c_uint16(x).value, @@ -306,7 +324,7 @@ truncate: dict[DType, Callable] = {dtypes.bool: bool, def _to_np_dtype(dtype:DType) -> type|None: import numpy as np - if dtype == dtypes.bfloat16: return np.float32 + if dtype in { dtypes.bfloat16, *dtypes.fp8s }: return np.float32 return np.dtype(dtype.fmt).type if dtype.fmt is not None else None def _from_np_dtype(npdtype:'np.dtype') -> DType: # type: ignore [name-defined] # noqa: F821 import numpy as np @@ -317,6 +335,7 @@ def _to_torch_dtype(dtype:DType) -> 'torch.dtype'|None: # type: ignore [name-de import numpy as np, torch if dtype == dtypes.uint64: return torch.uint64 if dtype == dtypes.bfloat16: return torch.bfloat16 + if dtype in dtypes.fp8s: return torch.uint8 # NOTE: torch doesn't expose this mapping with a stable API try: return torch.from_numpy(np.array([], dtype=_to_np_dtype(dtype))).dtype except TypeError: return None diff --git a/tinygrad/engine/realize.py b/tinygrad/engine/realize.py index 1246e3c2c4..50474a6284 100644 --- a/tinygrad/engine/realize.py +++ b/tinygrad/engine/realize.py @@ -140,13 +140,13 @@ class BufferXfer(BufferCopy): # **************** method cache **************** -method_cache: dict[tuple[str, bytes, tuple[int, ...], bool], CompiledRunner] = {} +method_cache: dict[tuple[str, type, bytes, tuple[int, ...], bool], CompiledRunner] = {} def get_runner(device:str, ast:UOp) -> CompiledRunner: # TODO: this should be all context relevant to rendering context = (BEAM.value, NOOPT.value, DEVECTORIZE.value) - ckey = (device, ast.key, context, False) + ckey = (device, type(Device[device].compiler), ast.key, context, False) if cret:=method_cache.get(ckey): return cret - bkey = (device.split(":")[0], ast.key, context, True) + bkey = (device.split(":")[0], type(Device[device].compiler), ast.key, context, True) if bret:=method_cache.get(bkey): method_cache[ckey] = ret = CompiledRunner(replace(bret.p, device=device), bret.lib) else: diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index f14d6b9445..e004f8b058 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -89,6 +89,8 @@ def suppress_finalizing(func): if not getattr(sys, 'is_finalizing', lambda: True)(): raise # re-raise if not finalizing return wrapper +def unwrap_class_type(cls_t:T): return cls_t.func if isinstance(cls_t, functools.partial) else cls_t + def pluralize(st:str, cnt:int): return f"{cnt} {st}"+('' if cnt == 1 else 's') class LazySeq(Generic[T]): # NOTE: Mapping requires __iter__ and __len__, Sequence requires supporting __len__ and slicing in __getitem__ @@ -133,7 +135,7 @@ USE_TC, TC_SELECT, TC_OPT, AMX = ContextVar("TC", 1), ContextVar("TC_SELECT", -1 TRANSCENDENTAL, NOLOCALS = ContextVar("TRANSCENDENTAL", 1), ContextVar("NOLOCALS", 0) FUSE_ARANGE, FUSE_CONV_BW = ContextVar("FUSE_ARANGE", 1), ContextVar("FUSE_CONV_BW", 0) SPLIT_REDUCEOP, NO_MEMORY_PLANNER, RING = ContextVar("SPLIT_REDUCEOP", 1), ContextVar("NO_MEMORY_PLANNER", 0), ContextVar("RING", 1) -PICKLE_BUFFERS, PROFILE, LRU = ContextVar("PICKLE_BUFFERS", 1), ContextVar("PROFILE", getenv("VIZ")), ContextVar("LRU", 1) +PICKLE_BUFFERS, LRU = ContextVar("PICKLE_BUFFERS", 1), ContextVar("LRU", 1) CACHELEVEL, IGNORE_BEAM_CACHE, DEVECTORIZE = ContextVar("CACHELEVEL", 2), ContextVar("IGNORE_BEAM_CACHE", 0), ContextVar("DEVECTORIZE", 1) DISABLE_COMPILER_CACHE, BLOCK_REORDER = ContextVar("DISABLE_COMPILER_CACHE", 0), ContextVar("BLOCK_REORDER", 1) DONT_REALIZE_EXPAND, DONT_GROUP_REDUCES = ContextVar("DONT_REALIZE_EXPAND", 0), ContextVar("DONT_GROUP_REDUCES", 0) @@ -144,6 +146,7 @@ RANGEIFY, FUSE_ATTENTION = ContextVar("RANGEIFY", 0), ContextVar("FUSE_ATTENTION EMULATE = ContextVar("EMULATE", "") CPU_COUNT = ContextVar("CPU_COUNT", max(1, (os.cpu_count() or 1) // (4 if ARCH_X86 else 2))) # take 1/2 of the cores, accounting HT CPU_LLVM, AMD_LLVM = ContextVar("CPU_LLVM", 0), ContextVar("AMD_LLVM", 1) +VIZ = PROFILE = ContextVar("VIZ", 0) @dataclass(frozen=True) class Metadata: @@ -322,7 +325,10 @@ def cpu_objdump(lib, objdump_tool='objdump'): print(subprocess.check_output([objdump_tool, '-d', f.name]).decode('utf-8')) def capstone_flatdump(lib: bytes): - import capstone + try: import capstone + except ImportError: + print("Disassembler Error: Capstone not installed.") + return match platform.machine(): case 'x86_64' | 'AMD64': cs = capstone.Cs(capstone.CS_ARCH_X86, capstone.CS_MODE_64) case 'aarch64' | 'arm64': cs = capstone.Cs(capstone.CS_ARCH_ARM64, capstone.CS_MODE_ARM) diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index b268ce7216..09d25c0828 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -192,7 +192,7 @@ class ClangRenderer(CStyleLanguage): float4_style = ('{', '}') gep_arr_threshold = 0 has_local = False - has_threads = True + has_threads = bool(getenv("THREADS", 1)) global_max = (CPU_COUNT.value, 0, 0) infinity = "__builtin_inff()" nan = '__builtin_nanf("")' @@ -242,7 +242,7 @@ class ClangRenderer(CStyleLanguage): return defines + "\n" + self._render_body(function_name, kernel, bufs, uops, prefix) + "\n" + self._render_entry(function_name, bufs) class OpenCLRenderer(CStyleLanguage): - device = "GPU" + device = "CL" # language options kernel_typedef = "__kernel void" @@ -271,7 +271,7 @@ class OpenCLRenderer(CStyleLanguage): return super().render_kernel(function_name, kernel, bufs, uops, prefix) class IntelRenderer(OpenCLRenderer): - device, suffix, kernel_typedef = "GPU", "INTEL", "__attribute__((intel_reqd_sub_group_size(8)))\n" + "__kernel void" + device, suffix, kernel_typedef = "CL", "INTEL", "__attribute__((intel_reqd_sub_group_size(8)))\n" + "__kernel void" tensor_cores = tc.intel string_rewrite = PatternMatcher([ diff --git a/tinygrad/renderer/llvmir.py b/tinygrad/renderer/llvmir.py index 847b0fcee1..f19f4dc271 100644 --- a/tinygrad/renderer/llvmir.py +++ b/tinygrad/renderer/llvmir.py @@ -3,6 +3,7 @@ import math, struct, sys from tinygrad.codegen.opt import tc from tinygrad.renderer import Renderer from tinygrad.renderer.cstyle import AMDRenderer +from tinygrad.uop.decompositions import xexp2, xlog2 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, GroupOp, sint_to_uop from tinygrad.dtype import dtypes, DType, PtrDType, truncate from tinygrad.helpers import prod, AMX @@ -197,8 +198,7 @@ barrier = 'fence syncscope("workgroup") release\ntail call void @llvm.amdgcn.s.b 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))}()"} # https://rocm.docs.amd.com/projects/llvm-project/en/latest/LLVM/llvm/html/AMDGPUUsage.html#llvm-ir-intrinsics -# llvm.log2/llvm.exp2 don't support double -llvm_intrinsics = {Ops.SQRT: "sqrt"} +llvm_intrinsics = {Ops.SQRT: "sqrt", Ops.LOG2: "log2", Ops.EXP2: "exp2"} class AMDLLVMRenderer(LLVMRenderer): device = "AMD" has_local = True @@ -217,6 +217,9 @@ class AMDLLVMRenderer(LLVMRenderer): lambda x, y: UOp(Ops.VECTORIZE, dtypes.half.vec(16), tuple(y.gep(i // 2) if i % 2 == 0 else UOp.const(dtypes.half, 0.0) for i in range(16)))), (UPat(Ops.CAST, name="x", dtype=dtypes.half.vec(8), src=UPat.var("y", dtypes.half.vec(16))), lambda x, y: UOp(Ops.VECTORIZE, dtypes.half.vec(8), tuple(y.gep(i * 2) for i in range(8)))), + # amd llvm intrinsics llvm.log2/llvm.exp2 don't support double + (UPat(Ops.LOG2, dtype=dtypes.double, src=(UPat.var("d"),)), xlog2), + (UPat(Ops.EXP2, dtype=dtypes.double, src=(UPat.var("d"),)), xexp2), ]) def _render_footer(self, uops: list[UOp]) -> str: # TODO: this is copied from cstyle diff --git a/tinygrad/renderer/nir.py b/tinygrad/renderer/nir.py index 40d2b0f614..16cf3dc92c 100644 --- a/tinygrad/renderer/nir.py +++ b/tinygrad/renderer/nir.py @@ -10,7 +10,6 @@ import tinygrad.runtime.autogen.libc as libc import ctypes, struct # FIXME: this is because clang2py produces bad output? -nir_intrinsic_infos = nir.nir_intrinsic_infos.in_dll(nir._libraries['FIXME_STUB'], "nir_intrinsic_infos") assert libc._libraries['libc'] stdout = ctypes.POINTER(nir.struct__IO_FILE).in_dll(libc._libraries['libc'], "stdout") s = nir.char_pointer_cast @@ -48,7 +47,7 @@ def nimm(b:nir.nir_builder, x, dtype:DType) -> nir.nir_def: def nir_src_for_ssa(d:nir.nir_def) -> nir.nir_src: return nir.nir_src(ssa=ctypes.pointer(d)) def nir_intrinsic_set(typ, instr:nir.nir_intrinsic_instr, val:int): - info = nir_intrinsic_infos[instr.contents.intrinsic] + info = nir.nir_intrinsic_infos.in_dll(nir._libraries['FIXME_STUB'], "nir_intrinsic_infos")[instr.contents.intrinsic] assert info.index_map[typ] > 0 instr.contents.const_index[info.index_map[typ] - 1] = val @@ -258,7 +257,7 @@ class NIRRenderer(Renderer): nir.nir_print_shader(self.b.shader, stdout) raise RuntimeError(f"failed to render {u.op} with {u.dtype} srcs {[x.dtype for x in u.src]}") self.r[u] = cast(nir.nir_def, d) - nir.nir_print_shader(self.b.shader, stdout) + # nir.nir_print_shader(self.b.shader, stdout) nir.nir_validate_shader(self.b.shader, b"after render") blob = nir.struct_blob() nir.nir_serialize(blob, self.b.shader, False) diff --git a/tinygrad/runtime/ops_amd.py b/tinygrad/runtime/ops_amd.py index c0c27cc717..b6b1776730 100644 --- a/tinygrad/runtime/ops_amd.py +++ b/tinygrad/runtime/ops_amd.py @@ -6,8 +6,8 @@ from dataclasses import dataclass from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, HWQueue, CLikeArgsState, HCQSignal, HCQProgram, FileIOInterface from tinygrad.runtime.support.hcq import MMIOInterface, BumpAllocator from tinygrad.uop.ops import sint -from tinygrad.device import Compiled, DMAFdRef, BufferSpec -from tinygrad.helpers import getenv, to_mv, round_up, data64_le, DEBUG, AMD_LLVM, PROFILE, ProfileEvent, suppress_finalizing, lo32, hi32 +from tinygrad.device import Compiled, DMAFdRef, BufferSpec, CompilerPairT +from tinygrad.helpers import getenv, to_mv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, suppress_finalizing, lo32, hi32 from tinygrad.renderer.cstyle import AMDRenderer from tinygrad.renderer.llvmir import AMDLLVMRenderer from tinygrad.runtime.autogen import kfd, hsa, pci, sqtt @@ -581,6 +581,9 @@ class KFDIface: if uncached: flags |= kfd.KFD_IOC_ALLOC_MEM_FLAGS_COHERENT | kfd.KFD_IOC_ALLOC_MEM_FLAGS_UNCACHED | kfd.KFD_IOC_ALLOC_MEM_FLAGS_GTT else: flags |= (kfd.KFD_IOC_ALLOC_MEM_FLAGS_USERPTR if host else kfd.KFD_IOC_ALLOC_MEM_FLAGS_VRAM) + # Make mapped cpu address to be uncachable + if cpu_addr is not None: flags |= kfd.KFD_IOC_ALLOC_MEM_FLAGS_COHERENT | kfd.KFD_IOC_ALLOC_MEM_FLAGS_UNCACHED + if cpu_access or host: flags |= kfd.KFD_IOC_ALLOC_MEM_FLAGS_PUBLIC if flags & kfd.KFD_IOC_ALLOC_MEM_FLAGS_USERPTR: @@ -785,9 +788,11 @@ class AMDDevice(HCQCompiled): max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000 self.sdma_queue = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x200 if self.is_usb() else (16 << 20)) - super().__init__(device, AMDAllocator(self), AMDLLVMRenderer(self.arch) if AMD_LLVM else AMDRenderer(self.arch), - AMDLLVMCompiler(self.arch) if AMD_LLVM else HIPCompiler(self.arch), functools.partial(AMDProgram, self), - AMDSignal, functools.partial(AMDComputeAQLQueue if self.is_aql else AMDComputeQueue, self), + compilers:list[CompilerPairT] = [(functools.partial(AMDLLVMRenderer, self.arch), functools.partial(AMDLLVMCompiler, self.arch)), + (functools.partial(AMDRenderer, self.arch), functools.partial(HIPCompiler, self.arch))] + + super().__init__(device, AMDAllocator(self), compilers, functools.partial(AMDProgram, self), AMDSignal, + functools.partial(AMDComputeAQLQueue if self.is_aql else AMDComputeQueue, self), functools.partial(AMDCopyQueue, self, max_copy_size=max_copy_size), kernargs_size=(8 << 10) if self.is_usb() else (16 << 20), sigalloc_size=0x100 if self.is_usb() else 0x1000) diff --git a/tinygrad/runtime/ops_gpu.py b/tinygrad/runtime/ops_cl.py similarity index 92% rename from tinygrad/runtime/ops_gpu.py rename to tinygrad/runtime/ops_cl.py index c9ebb338c6..8887c97f00 100644 --- a/tinygrad/runtime/ops_gpu.py +++ b/tinygrad/runtime/ops_cl.py @@ -2,7 +2,7 @@ from __future__ import annotations from typing import cast import ctypes, functools, hashlib from tinygrad.runtime.autogen import opencl as cl -from tinygrad.helpers import init_c_var, to_char_p_p, from_mv, OSX, DEBUG, getenv, mv_address, suppress_finalizing +from tinygrad.helpers import init_c_var, to_char_p_p, from_mv, OSX, DEBUG, mv_address, suppress_finalizing from tinygrad.renderer.cstyle import OpenCLRenderer, IntelRenderer from tinygrad.device import BufferSpec, LRUAllocator, Compiled, Compiler, CompileError @@ -48,8 +48,8 @@ class CLProgram: def __call__(self, *bufs:tuple[ctypes._CData, BufferSpec], global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]|None=None, vals:tuple[int, ...]=(), wait=False) -> float|None: - for i,(b,_) in enumerate(bufs): cl.clSetKernelArg(self.kernel, i, ctypes.sizeof(b), ctypes.byref(b)) - for i,v in enumerate(vals,start=len(bufs)): cl.clSetKernelArg(self.kernel, i, 4, ctypes.byref(ctypes.c_int32(v))) + for i,(b,_) in enumerate(bufs): check(cl.clSetKernelArg(self.kernel, i, ctypes.sizeof(b), ctypes.byref(b))) + for i,v in enumerate(vals,start=len(bufs)): check(cl.clSetKernelArg(self.kernel, i, 4, ctypes.byref(ctypes.c_int32(v)))) if local_size is not None: global_size = cast(tuple[int,int,int], tuple(int(g*l) for g,l in zip(global_size, local_size))) event = cl.cl_event() if wait else None check(cl.clEnqueueNDRangeKernel(self.dev.queue, self.kernel, len(global_size), None, (ctypes.c_size_t * len(global_size))(*global_size), @@ -108,11 +108,9 @@ class CLDevice(Compiled): 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] # noqa: E501 - compile_key = hashlib.md5(self.device_name.encode() + self.driver_version.encode()).hexdigest() - renderer = IntelRenderer() if "cl_intel_subgroup_matrix_multiply_accumulate" in self.device_exts and getenv("INTEL") else OpenCLRenderer() - super().__init__(device, CLAllocator(self), renderer, CLCompiler(self, f"compile_cl_{compile_key}"), functools.partial(CLProgram, self)) + compilers = [(IntelRenderer if "cl_intel_subgroup_matrix_multiply_accumulate" in self.device_exts else OpenCLRenderer, + functools.partial(CLCompiler, self, f"compile_cl_{hashlib.md5(self.device_name.encode() + self.driver_version.encode()).hexdigest()}"))] + super().__init__(device, CLAllocator(self), compilers, functools.partial(CLProgram, self)) def synchronize(self): check(cl.clFinish(self.queue)) self.pending_copyin.clear() - -GPUDevice = CLDevice # for legacy reasons diff --git a/tinygrad/runtime/ops_cpu.py b/tinygrad/runtime/ops_cpu.py index 53a0029d7d..24ea9042ca 100644 --- a/tinygrad/runtime/ops_cpu.py +++ b/tinygrad/runtime/ops_cpu.py @@ -1,18 +1,16 @@ from __future__ import annotations import platform, sys, ctypes, functools, time, mmap, threading, queue, struct -from tinygrad.helpers import from_mv, to_mv, OSX, WIN, mv_address, wait_cond, cpu_profile, CPU_LLVM, suppress_finalizing, getenv, data64_le, i2u +from tinygrad.helpers import from_mv, to_mv, OSX, WIN, mv_address, wait_cond, cpu_profile, suppress_finalizing, unwrap, data64_le, i2u from tinygrad.device import BufferSpec, DMACPURef from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocatorBase, HCQBuffer, HWQueue, HCQArgsState, HCQSignal, HCQProgram, MMIOInterface from tinygrad.runtime.support.hcq import CLikeArgsState from tinygrad.renderer.cstyle import ClangRenderer from tinygrad.renderer.llvmir import LLVMRenderer -from tinygrad.runtime.support.compiler_cpu import HostLLVMCompiler, ClangJITCompiler +from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, ClangJITCompiler from tinygrad.uop.ops import sint -if (NIR := getenv("NIR")): - from tinygrad.renderer.nir import LVPRenderer - from tinygrad.runtime.support.lvp import LVPCompiler - from tinygrad.runtime.support.elf import elf_loader - import tinygrad.runtime.autogen.libc as libc +from tinygrad.renderer.nir import LVPRenderer +from tinygrad.runtime.support.lvp import LVPCompiler +from tinygrad.runtime.support.elf import elf_loader class CPUSignal(HCQSignal): def _sleep(self, time_spent_waiting_ms:int): @@ -52,7 +50,7 @@ class CPUComputeQueue(HWQueue): def memory_barrier(self): return self def exec(self, prg:CPUProgram, args_state:HCQArgsState, global_size, local_size): - if NIR: + if prg.LVP: self.bind_args_state(args_state) return self.cmd(self._exec, prg, 1, args_state.buf.va_addr) return self.cmd(self._exec, prg, len(args_state.bufs), *[x.va_addr for x in args_state.bufs], *args_state.vals, threads=(global_size or (1,))[0]) @@ -69,9 +67,12 @@ class LVPArgsState(CLikeArgsState): MAP_JIT = 0x0800 class CPUProgram(HCQProgram): - rt_lib = ctypes.CDLL(ctypes.util.find_library('System' if OSX else 'kernel32') if OSX or WIN else 'libgcc_s.so.1') + rt_lib = None + try: rt_lib = ctypes.CDLL(ctypes.util.find_library('System' if OSX else 'kernel32') if OSX or WIN else 'libgcc_s.so.1') + except OSError: pass def __init__(self, dev, name:str, lib:bytes): + self.LVP = isinstance(dev.compiler, LVPCompiler) if sys.platform == "win32": # mypy doesn't understand when WIN is used here PAGE_EXECUTE_READWRITE, MEM_COMMIT, MEM_RESERVE = 0x40, 0x1000, 0x2000 ctypes.windll.kernel32.VirtualAlloc.restype = ctypes.c_void_p @@ -86,8 +87,9 @@ class CPUProgram(HCQProgram): # MAP_JIT allows us to easily flip pages from RW- to R-X and vice versa. It is a noop on intel cpus. (man pthread_jit_write_protect_np) self.mem = mmap.mmap(-1, len(lib), mmap.MAP_ANON|mmap.MAP_PRIVATE|(MAP_JIT if OSX else 0), mmap.PROT_READ|mmap.PROT_WRITE|mmap.PROT_EXEC) - if OSX: CPUProgram.rt_lib.pthread_jit_write_protect_np(False) - if NIR: + if OSX: unwrap(CPUProgram.rt_lib).pthread_jit_write_protect_np(False) + if self.LVP: + from tinygrad.runtime.autogen import libc (image, _, relocs), addr = elf_loader(lib), ctypes.addressof(ctypes.c_void_p.from_buffer(self.mem)) for ploc,tgt,r_type,r_addend in relocs: match r_type: @@ -95,17 +97,22 @@ class CPUProgram(HCQProgram): case _: raise NotImplementedError(f"Encountered unknown relocation type {r_type}") self.mem.write(image) else: self.mem.write(lib) - if OSX: CPUProgram.rt_lib.pthread_jit_write_protect_np(True) + if OSX: unwrap(CPUProgram.rt_lib).pthread_jit_write_protect_np(True) # __clear_cache isn't a normal libc function, but a compiler support routine found in libgcc_s for gcc and compiler-rt for clang. # libgcc_s comes as shared library but compiler-rt is only a bunch of static library archives which we can't directly load, but fortunately # it somehow found its way into libSystem on macos (likely because it used __builtin_clear_cache) and libgcc_s is ~always present on linux # Using ["name"] instead of .name because otherwise name is getting mangled: https://docs.python.org/3.12/reference/expressions.html#index-5 - CPUProgram.rt_lib["__clear_cache"](ctypes.c_void_p(mv_address(self.mem)), ctypes.c_void_p(mv_address(self.mem) + len(lib))) + if CPUProgram.rt_lib is not None: + CPUProgram.rt_lib["__clear_cache"](ctypes.c_void_p(mv_address(self.mem)), ctypes.c_void_p(mv_address(self.mem) + len(lib))) + else: + # msync should be a universal POSIX way to do this + from tinygrad.runtime.autogen import libc + libc.msync(ctypes.c_void_p(mv_address(self.mem)), len(lib), libc.MS_SYNC | libc.MS_INVALIDATE) self.fxn = ctypes.CFUNCTYPE(None)(mv_address(self.mem)) - super().__init__(LVPArgsState if NIR else HCQArgsState, dev, name, kernargs_alloc_size=12+256 if NIR else 0) + super().__init__(LVPArgsState if self.LVP else HCQArgsState, dev, name, kernargs_alloc_size=12+256 if self.LVP else 0) @suppress_finalizing def __del__(self): @@ -136,5 +143,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), LLVMRenderer() if CPU_LLVM else (LVPRenderer(self) if NIR else ClangRenderer()), HostLLVMCompiler() - if CPU_LLVM else (LVPCompiler() if NIR else ClangJITCompiler()), functools.partial(CPUProgram, self), CPUSignal, CPUComputeQueue) + compilers = [(ClangRenderer, ClangJITCompiler), (LLVMRenderer, CPULLVMCompiler), (functools.partial(LVPRenderer, self), LVPCompiler)] + super().__init__(device, CPUAllocator(self), compilers, functools.partial(CPUProgram, self), CPUSignal, CPUComputeQueue) diff --git a/tinygrad/runtime/ops_cuda.py b/tinygrad/runtime/ops_cuda.py index 326ae84f01..440f68b56b 100644 --- a/tinygrad/runtime/ops_cuda.py +++ b/tinygrad/runtime/ops_cuda.py @@ -1,11 +1,11 @@ from __future__ import annotations import ctypes, ctypes.util, functools from tinygrad.helpers import DEBUG, getenv, mv_address, init_c_var, init_c_struct_t, suppress_finalizing -from tinygrad.device import Compiled, BufferSpec, LRUAllocator +from tinygrad.device import Compiled, BufferSpec, LRUAllocator, CompilerPairT from tinygrad.renderer.cstyle import CUDARenderer from tinygrad.renderer.ptx import PTXRenderer from tinygrad.runtime.autogen import cuda -from tinygrad.runtime.support.compiler_cuda import pretty_ptx, CUDACompiler, PTXCompiler, PTX +from tinygrad.runtime.support.compiler_cuda import pretty_ptx, CUDACompiler, PTXCompiler if getenv("IOCTL"): import extra.nv_gpu_driver.nv_ioctl # noqa: F401 # pylint: disable=unused-import if MOCKGPU:=getenv("MOCKGPU"): from test.mockgpu.cuda import cuda # type: ignore # pylint: disable=reimported @@ -115,8 +115,9 @@ class CUDADevice(Compiled): CUDADevice.devices.append(self) from tinygrad.runtime.graph.cuda import CUDAGraph - super().__init__(device, CUDAAllocator(self), PTXRenderer(self.arch) if PTX else CUDARenderer(self.arch), - PTXCompiler(self.arch) if PTX else CUDACompiler(self.arch), functools.partial(CUDAProgram, self), None if MOCKGPU else CUDAGraph) + compilers:list[CompilerPairT] = [(functools.partial(CUDARenderer, self.arch), functools.partial(CUDACompiler, self.arch)), + (functools.partial(PTXRenderer, self.arch), functools.partial(PTXCompiler, self.arch))] + super().__init__(device, CUDAAllocator(self), compilers, functools.partial(CUDAProgram, self), None if MOCKGPU else CUDAGraph) def synchronize(self): check(cuda.cuCtxSetCurrent(self.context)) diff --git a/tinygrad/runtime/ops_disk.py b/tinygrad/runtime/ops_disk.py index 1f810ec98f..56809dafbc 100644 --- a/tinygrad/runtime/ops_disk.py +++ b/tinygrad/runtime/ops_disk.py @@ -15,7 +15,7 @@ class DiskDevice(Compiled): self.size: int|None = None self.fd: int|None = None self.count = 0 - super().__init__(device, DiskAllocator(self), None, None, None) + super().__init__(device, DiskAllocator(self), None, None) def _might_open(self, size:int): assert self.size is None or size <= self.size, f"can't reopen Disk tensor with larger size, opened with {self.size}, tried to open with {size}" if self.size is not None and hasattr(self.device, "mem"): @@ -39,7 +39,9 @@ class DiskDevice(Compiled): def _might_close(self): self.count -= 1 if self.count == 0: - if self.fd is not None: os.close(self.fd) + if self.fd is not None: + os.close(self.fd) + if hasattr(self, "mem"): self.mem.close() self.size = None def _iouring_setup(self): DiskDevice._tried_io_uring_init = True diff --git a/tinygrad/runtime/ops_dsp.py b/tinygrad/runtime/ops_dsp.py index 9be65a384c..d93f14afa8 100644 --- a/tinygrad/runtime/ops_dsp.py +++ b/tinygrad/runtime/ops_dsp.py @@ -134,8 +134,8 @@ class DSPDevice(Compiled): def __init__(self, device:str=""): compiler_args = ["--target=hexagon", "-mcpu=hexagonv65", "-fuse-ld=lld", "-nostdlib", "-mhvx=v65", "-mhvx-length=128b"] if getenv("MOCKDSP"): - super().__init__(device, CPUAllocator(self), MockDSPRenderer(), - ClangCompiler(None, ["-static"] + compiler_args, 'llvm-objdump'), MockDSPProgram) + mock_compilers = [(MockDSPRenderer, functools.partial(ClangCompiler, None, ["-static"] + compiler_args, 'llvm-objdump'))] + super().__init__(device, CPUAllocator(self), mock_compilers, MockDSPProgram) else: self.ion_fd = os.open('/dev/ion', os.O_RDONLY) # Generate link script to pass into clang. Aligning all used sections to 4k fixes invoke problem. @@ -146,8 +146,9 @@ class DSPDevice(Compiled): self.link_ld.write(f"SECTIONS {{ . = 0x0; {sections_link}\n /DISCARD/ : {{ *(.note .note.* .gnu.hash .comment) }} }}".encode()) self.link_ld.flush() - super().__init__(device, DSPAllocator(self), DSPRenderer(), - ClangCompiler("compile_dsp", ["-shared"] + compiler_args + [f"-T{self.link_ld.name}"], 'llvm-objdump'), functools.partial(DSPProgram, self)) + compilers = [(DSPRenderer, functools.partial(ClangCompiler, "compile_dsp", ["-shared"] + compiler_args + [f"-T{self.link_ld.name}"], + 'llvm-objdump'))] + super().__init__(device, DSPAllocator(self), compilers, functools.partial(DSPProgram, self)) fastrpc_shell = memoryview(bytearray(pathlib.Path('/dsp/cdsp/fastrpc_shell_3').read_bytes())) self.shell_buf = self.allocator.alloc(round_up(fastrpc_shell.nbytes, 0x1000), BufferSpec(nolru=True)) ctypes.memmove(self.shell_buf.va_addr, mv_address(fastrpc_shell), fastrpc_shell.nbytes) diff --git a/tinygrad/runtime/ops_hip.py b/tinygrad/runtime/ops_hip.py index a15e366855..6bd9760016 100644 --- a/tinygrad/runtime/ops_hip.py +++ b/tinygrad/runtime/ops_hip.py @@ -14,7 +14,9 @@ class HIPDevice(Compiled): self.device_id = int(device.split(":")[1]) if ":" in device else 0 self.arch = init_c_var(hip.hipDeviceProp_t(), lambda x: check(hip.hipGetDeviceProperties(x, self.device_id))).gcnArchName.decode() self.time_event_st, self.time_event_en = [init_c_var(hip.hipEvent_t(), lambda x: hip.hipEventCreate(ctypes.byref(x), 0)) for _ in range(2)] - super().__init__(device, HIPAllocator(self), HIPRenderer(self.arch), HIPCompiler(self.arch), functools.partial(HIPProgram, self)) + + compilers = [(functools.partial(HIPRenderer, self.arch), functools.partial(HIPCompiler, self.arch))] + super().__init__(device, HIPAllocator(self), compilers, functools.partial(HIPProgram, self)) def synchronize(self): check(hip.hipSetDevice(self.device_id)) check(hip.hipDeviceSynchronize()) diff --git a/tinygrad/runtime/ops_metal.py b/tinygrad/runtime/ops_metal.py index 9e453eb729..df9d136d7d 100644 --- a/tinygrad/runtime/ops_metal.py +++ b/tinygrad/runtime/ops_metal.py @@ -76,7 +76,7 @@ class MetalDevice(Compiled): from tinygrad.runtime.graph.metal import MetalGraph # NOTE: GitHub CI macOS runners use paravirtualized metal which is broken with graph. # 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(), MetalCompiler() if getenv("METAL_DIRECT", 1) else Compiler(), + super().__init__(device, MetalAllocator(self), [(MetalRenderer, MetalCompiler), (MetalRenderer, Compiler)], functools.partial(MetalProgram, self), MetalGraph if 'virtual' not in from_ns_str(msg('name')(self.sysdevice)).lower() else None) def synchronize(self): diff --git a/tinygrad/runtime/ops_npy.py b/tinygrad/runtime/ops_npy.py index f40bfcdeeb..d92309ba52 100644 --- a/tinygrad/runtime/ops_npy.py +++ b/tinygrad/runtime/ops_npy.py @@ -8,4 +8,4 @@ class NpyAllocator(Allocator['NpyDevice']): def _copyout(self, dest:memoryview, src:np.ndarray): dest[:] = self._as_buffer(src) class NpyDevice(Compiled): - def __init__(self, device:str): super().__init__(device, NpyAllocator(self), None, None, None) + def __init__(self, device:str): super().__init__(device, NpyAllocator(self), None, None) diff --git a/tinygrad/runtime/ops_null.py b/tinygrad/runtime/ops_null.py index cd293f9988..c8f5a6b59f 100644 --- a/tinygrad/runtime/ops_null.py +++ b/tinygrad/runtime/ops_null.py @@ -29,5 +29,5 @@ class NullGraph(MultiGraphRunner): def __call__(self, input_rawbuffers, var_vals, wait=False) -> float|None: return 1e-3 class NullDevice(Compiled): - def __init__(self, device:str): super().__init__(device, NullAllocator(self), NullRenderer(), Compiler(), functools.partial(NullProgram, device), + def __init__(self, device:str): super().__init__(device, NullAllocator(self), [(NullRenderer, Compiler)], functools.partial(NullProgram, device), NullGraph) diff --git a/tinygrad/runtime/ops_nv.py b/tinygrad/runtime/ops_nv.py index e86e55ffaa..57c952fe30 100644 --- a/tinygrad/runtime/ops_nv.py +++ b/tinygrad/runtime/ops_nv.py @@ -6,19 +6,18 @@ from dataclasses import dataclass from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, HWQueue, CLikeArgsState, HCQProgram, HCQSignal, BumpAllocator from tinygrad.runtime.support.hcq import MMIOInterface, FileIOInterface, MOCKGPU from tinygrad.uop.ops import sint -from tinygrad.device import BufferSpec +from tinygrad.device import BufferSpec, CompilerPairT from tinygrad.helpers import getenv, mv_address, round_up, data64, data64_le, prod, OSX, to_mv, hi32, lo32, suppress_finalizing from tinygrad.renderer.ptx import PTXRenderer from tinygrad.renderer.cstyle import NVRenderer -from tinygrad.runtime.support.compiler_cuda import CUDACompiler, PTXCompiler, PTX, NVPTXCompiler, NVCompiler +from tinygrad.runtime.support.compiler_cuda import CUDACompiler, PTXCompiler, NVPTXCompiler, NVCompiler from tinygrad.runtime.autogen import nv_gpu, pci from tinygrad.runtime.support.elf import elf_loader from tinygrad.runtime.support.nv.nvdev import NVDev, NVMemoryManager from tinygrad.runtime.support.system import System, PCIIfaceBase, MAP_FIXED +from tinygrad.renderer.nir import NAKRenderer +from tinygrad.runtime.support.nak import NAKCompiler, parse_nak_shader if getenv("IOCTL"): import extra.nv_gpu_driver.nv_ioctl # noqa: F401 # pylint: disable=unused-import -if (NIR := getenv("NIR")): - from tinygrad.renderer.nir import NAKRenderer - from tinygrad.runtime.support.nak import NAKCompiler, parse_nak_shader def get_error_str(status): return f"{status}: {nv_gpu.nv_status_codes.get(status, 'Unknown error')}" @@ -196,7 +195,7 @@ class NVProgram(HCQProgram): self.dev, self.name, self.lib = dev, name, lib self.constbufs: dict[int, tuple[int, int]] = {0: (0, 0x160)} # dict[constbuf index, tuple[va_addr, size]] - if NIR: + if (NAK:=isinstance(dev.compiler, NAKCompiler)): image, self.regs_usage, self.shmem_usage, self.lcmem_usage = parse_nak_shader(lib) self.lib_gpu = self.dev.allocator.alloc(round_up(image.nbytes, 0x1000) + 0x1000, buf_spec:=BufferSpec(cpu_access=True)) self.prog_addr, self.prog_sz, self.constbuffer_0 = self.lib_gpu.va_addr, image.nbytes, None @@ -238,18 +237,18 @@ class NVProgram(HCQProgram): ctypes.memmove(self.lib_gpu.va_addr, mv_address(image), image.nbytes) if dev.iface.compute_class >= nv_gpu.BLACKWELL_COMPUTE_A: - if not NIR: + if not NAK: self.constbuffer_0[188:192] = [*data64_le(self.dev.shared_mem_window), *data64_le(self.dev.local_mem_window)] self.constbuffer_0[223] = 0xfffdc0 qmd = {'qmd_major_version':5, 'qmd_type':nv_gpu.NVCEC0_QMDV05_00_QMD_TYPE_GRID_CTA, 'register_count':self.regs_usage, 'program_address_upper_shifted4':hi32(self.prog_addr>>4), 'program_address_lower_shifted4':lo32(self.prog_addr>>4), - 'shared_memory_size_shifted7':self.shmem_usage>>7, 'shader_local_memory_high_size_shifted4':0 if NIR else self.dev.slm_per_thread>>4, - **({'shader_local_memory_high_size_shifted4':self.lcmem_usage>>4} if NIR else {})} + 'shared_memory_size_shifted7':self.shmem_usage>>7, 'shader_local_memory_high_size_shifted4':0 if NAK else self.dev.slm_per_thread>>4, + **({'shader_local_memory_high_size_shifted4':self.lcmem_usage>>4} if NAK else {})} else: - if not NIR: self.constbuffer_0[6:12] = [*data64_le(self.dev.shared_mem_window), *data64_le(self.dev.local_mem_window), *data64_le(0xfffdc0)] - qmd = {'qmd_major_version':3, 'sm_global_caching_enable':1, 'shader_local_memory_high_size':0 if NIR else self.dev.slm_per_thread, + if not NAK: self.constbuffer_0[6:12] = [*data64_le(self.dev.shared_mem_window), *data64_le(self.dev.local_mem_window), *data64_le(0xfffdc0)] + qmd = {'qmd_major_version':3, 'sm_global_caching_enable':1, 'shader_local_memory_high_size':0 if NAK else self.dev.slm_per_thread, 'program_address_upper':hi32(self.prog_addr), 'program_address_lower':lo32(self.prog_addr), 'shared_memory_size':self.shmem_usage, - 'register_count_v':self.regs_usage, **({'shader_local_memory_low_size':self.lcmem_usage} if NIR else {})} + 'register_count_v':self.regs_usage, **({'shader_local_memory_low_size':self.lcmem_usage} if NAK else {})} smem_cfg = min(shmem_conf * 1024 for shmem_conf in [32, 64, 100] if shmem_conf * 1024 >= self.shmem_usage) // 4096 + 1 @@ -537,9 +536,10 @@ class NVDevice(HCQCompiled[HCQSignal]): self.arch: str = "sm_120" if self.sm_version==0xa04 else f"sm_{(self.sm_version>>8)&0xff}{(val>>4) if (val:=self.sm_version&0xff) > 0xf else val}" self.sass_version = ((self.sm_version & 0xf00) >> 4) | (self.sm_version & 0xf) - cc = NAKCompiler(self) if NIR else ((PTXCompiler if PTX else CUDACompiler) if MOCKGPU else (NVPTXCompiler if PTX else NVCompiler))(self.arch) - rr = PTXRenderer(self.arch, device="NV") if PTX else (NAKRenderer(self) if NIR else NVRenderer(self.arch)) - super().__init__(device, NVAllocator(self), rr, cc, functools.partial(NVProgram, self), HCQSignal, NVComputeQueue, NVCopyQueue) + compilers:list[CompilerPairT] = [(functools.partial(NVRenderer, self.arch),functools.partial(CUDACompiler if MOCKGPU else NVCompiler, self.arch)), + (functools.partial(PTXRenderer, self.arch, device="NV"), functools.partial(PTXCompiler if MOCKGPU else NVPTXCompiler, self.arch)), + (functools.partial(NAKRenderer, self), functools.partial(NAKCompiler, self))] + super().__init__(device, NVAllocator(self), compilers, functools.partial(NVProgram, self), HCQSignal, NVComputeQueue, NVCopyQueue) self._setup_gpfifos() diff --git a/tinygrad/runtime/ops_python.py b/tinygrad/runtime/ops_python.py index 2a2b832658..9dd145d299 100644 --- a/tinygrad/runtime/ops_python.py +++ b/tinygrad/runtime/ops_python.py @@ -4,21 +4,23 @@ # this is the (living) definition of uops from typing import Any, TYPE_CHECKING, cast import pickle, base64, itertools, time, struct, sys -from tinygrad.dtype import DType, dtypes, ImageDType, PtrDType, truncate, float_to_bf16 +from tinygrad.dtype import DType, dtypes, ImageDType, PtrDType, truncate, float_to_bf16, float_to_fp8, fp8_to_float from tinygrad.helpers import all_same, getenv, flatten, get_single_element, EMULATE from tinygrad.device import Compiled, Compiler, Allocator from tinygrad.codegen.opt import tc from tinygrad.uop.ops import exec_alu, python_alu, Ops, UOp, GroupOp from tinygrad.renderer import Renderer -def storage_fmt_for_dtype(dtype: DType): return 'H' if dtype == dtypes.bfloat16 else dtype.fmt +def storage_fmt_for_dtype(dtype: DType): return 'H' if dtype == dtypes.bfloat16 else 'B' if dtype in dtypes.fp8s else dtype.fmt def to_storage_scalar(x, dtype: DType): if dtype == dtypes.bfloat16: return (struct.unpack('I', struct.pack('f', float_to_bf16(x)))[0] >> 16) & 0xFFFF + if dtype in dtypes.fp8s: return float_to_fp8(float(x), dtype) return x def from_storage_scalar(x, dtype: DType): if dtype == dtypes.bfloat16: return struct.unpack('f', struct.pack('I', (x & 0xFFFF) << 16))[0] + if dtype in dtypes.fp8s: return fp8_to_float(int(x), dtype) return x def _load(m, i, dtype: DType): @@ -236,4 +238,4 @@ class PythonAllocator(Allocator['PythonDevice']): def _copyout(self, dest:memoryview, src): dest[:] = src class PythonDevice(Compiled): - def __init__(self, device:str): super().__init__(device, PythonAllocator(self), PythonRenderer(), PythonCompiler(), PythonProgram) + def __init__(self, device:str): super().__init__(device, PythonAllocator(self), [(PythonRenderer, PythonCompiler)], PythonProgram) diff --git a/tinygrad/runtime/ops_qcom.py b/tinygrad/runtime/ops_qcom.py index bc0857a67a..787d349ee9 100644 --- a/tinygrad/runtime/ops_qcom.py +++ b/tinygrad/runtime/ops_qcom.py @@ -7,7 +7,7 @@ from tinygrad.device import BufferSpec from tinygrad.runtime.support.hcq import HCQBuffer, HWQueue, HCQProgram, HCQCompiled, HCQAllocatorBase, HCQSignal, HCQArgsState, BumpAllocator from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface from tinygrad.runtime.autogen import kgsl, adreno -from tinygrad.runtime.ops_gpu import CLCompiler, CLDevice +from tinygrad.runtime.ops_cl import CLCompiler, CLDevice from tinygrad.renderer.cstyle import QCOMRenderer from tinygrad.helpers import getenv, mv_address, to_mv, round_up, data64_le, prod, fromimport if getenv("IOCTL"): import extra.qcom_gpu_driver.opencl_ioctl # noqa: F401 # pylint: disable=unused-import @@ -341,8 +341,8 @@ class QCOMDevice(HCQCompiled): QCOMDevice.gpu_id = ((info.chip_id >> 24) & 0xFF) * 100 + ((info.chip_id >> 16) & 0xFF) * 10 + ((info.chip_id >> 8) & 0xFF) if QCOMDevice.gpu_id >= 700: raise RuntimeError(f"Unsupported GPU: {QCOMDevice.gpu_id}") - super().__init__(device, QCOMAllocator(self), QCOMRenderer(), QCOMCompiler(device), functools.partial(QCOMProgram, self), - QCOMSignal, QCOMComputeQueue, None) + compilers = [(QCOMRenderer, functools.partial(QCOMCompiler, device))] + super().__init__(device, QCOMAllocator(self), compilers, functools.partial(QCOMProgram, self), QCOMSignal, QCOMComputeQueue, None) def _gpu_alloc(self, size:int, flags:int=0, uncached=False, fill_zeroes=False) -> HCQBuffer: flags |= kgsl.KGSL_MEMALIGN(alignment_hint:=12) | kgsl.KGSL_MEMFLAGS_USE_CPU_MAP diff --git a/tinygrad/runtime/ops_remote.py b/tinygrad/runtime/ops_remote.py index fd11ed282a..147063f3ee 100644 --- a/tinygrad/runtime/ops_remote.py +++ b/tinygrad/runtime/ops_remote.py @@ -471,10 +471,11 @@ class RemoteDevice(Compiled): if not renderer[0].startswith("tinygrad.") or not renderer[1].endswith("Renderer"): raise RuntimeError(f"bad renderer {renderer}") renderer_class = fromimport(renderer[0], renderer[1]) # TODO: is this secure? if not issubclass(renderer_class, Renderer): raise RuntimeError(f"renderer isn't a Renderer {renderer}") - renderer_instance = renderer_class(*renderer[2]) - renderer_instance.device = device + graph = fromimport('tinygrad.runtime.graph.remote', "RemoteGraph") if self.properties.graph_supported else None - super().__init__(device, RemoteAllocator(self), renderer_instance, Compiler(), functools.partial(RemoteProgram, self), graph, id(self.conn)) + compilers = [(functools.partial(renderer_class, *renderer[2]), Compiler)] + super().__init__(device, RemoteAllocator(self), compilers, functools.partial(RemoteProgram, self), graph, id(self.conn)) + self.renderer.device = device def finalize(self): with contextlib.suppress(ConnectionError, http.client.HTTPException): self.q(SessionFree(), wait=True) diff --git a/tinygrad/runtime/ops_webgpu.py b/tinygrad/runtime/ops_webgpu.py index eef0a5cc30..cc070fd20a 100644 --- a/tinygrad/runtime/ops_webgpu.py +++ b/tinygrad/runtime/ops_webgpu.py @@ -217,7 +217,7 @@ class WebGpuDevice(Compiled): device_res = _run(webgpu.wgpuAdapterRequestDeviceF, webgpu.WGPURequestDeviceCallbackInfo, webgpu.WGPURequestDeviceCallback, webgpu.WGPURequestDeviceStatus__enumvalues, 1, 2, adapter_res, dev_desc) - super().__init__(device, WebGpuAllocator(device_res), WGSLRenderer(), Compiler(), + super().__init__(device, WebGpuAllocator(device_res), [(WGSLRenderer, Compiler)], functools.partial(WebGPUProgram, (device_res, webgpu.WGPUFeatureName_TimestampQuery in supported))) def synchronize(self): diff --git a/tinygrad/runtime/support/am/amdev.py b/tinygrad/runtime/support/am/amdev.py index fa9ba79321..7ec9b4ae85 100644 --- a/tinygrad/runtime/support/am/amdev.py +++ b/tinygrad/runtime/support/am/amdev.py @@ -84,7 +84,7 @@ class AMFirmware: self.descs += [self.desc(blob, hdr0.header.ucode_array_offset_bytes, hdr0.header.ucode_size_bytes, am.GFX_FW_TYPE_RLC_G)] def load_fw(self, fname:str, *headers, versioned_header:str|None=None): - fpath = fetch(f"https://gitlab.com/kernel-firmware/linux-firmware/-/raw/45f59212aebd226c7630aff4b58598967c0c8c91/amdgpu/{fname}", subdir="fw") + fpath = fetch(f"https://gitlab.com/kernel-firmware/linux-firmware/-/raw/a9f26799247aa60fbaa3b64267a18f20b72b5235/amdgpu/{fname}", subdir="fw") blob = memoryview(bytearray(fpath.read_bytes())) if AM_DEBUG >= 1: print(f"am {self.adev.devfmt}: loading firmware {fname}: {hashlib.sha256(blob).hexdigest()}") if versioned_header: diff --git a/tinygrad/runtime/support/compiler_cpu.py b/tinygrad/runtime/support/compiler_cpu.py index e16e9fb112..f9ec8d1062 100644 --- a/tinygrad/runtime/support/compiler_cpu.py +++ b/tinygrad/runtime/support/compiler_cpu.py @@ -13,7 +13,7 @@ class ClangJITCompiler(Compiler): # x18 is a reserved platform register. It is clobbered on context switch in macos and is used to store TEB pointer in windows on arm, don't use it target = 'x86_64' if sys.platform == 'win32' else platform.machine() # on arm march means "runs on this arch and superset" instead of "optimize for this arch". x86 march == arm mcpu - arch = '-march=native' if platform.machine() in ('x86_64', 'AMD64') else '-mcpu=native' + arch = {'x86_64': '-march=native', 'AMD64': '-march=native', 'riscv64': '-march=rv64g'}.get(platform.machine(), "-mcpu=native") args = [arch, f'--target={target}-none-unknown-elf', '-O2', '-fPIC', '-ffreestanding', '-fno-math-errno', '-nostdlib', '-fno-ident'] arch_args = ['-ffixed-x18'] if target == 'arm64' else [] obj = subprocess.check_output([getenv("CC", 'clang'), '-c', '-x', 'c', *args, *arch_args, '-', '-o', '-'], input=src.encode('utf-8')) @@ -29,7 +29,7 @@ def expect(x, err, ret=None): class LLVMCompiler(Compiler): jit = True - target_arch = {'arm64': 'AArch64', 'aarch64': 'AArch64', 'x86_64': 'X86', 'AMD64': 'X86'}[platform.machine()] + target_arch = {'arm64': 'AArch64', 'aarch64': 'AArch64', 'x86_64': 'X86', 'AMD64': 'X86', 'riscv64': 'riscv64'}[platform.machine()] def __init__(self, processor:str, feats:str): for component in ['Target', 'TargetInfo', 'TargetMC', 'AsmParser', 'AsmPrinter']: getattr(llvm, f'LLVMInitialize{self.target_arch}{component}')() @@ -79,7 +79,7 @@ class LLVMCompiler(Compiler): def disassemble(self, lib:bytes): capstone_flatdump(lib) -class HostLLVMCompiler(LLVMCompiler): +class CPULLVMCompiler(LLVMCompiler): def __init__(self): # +reserve-x18 here does the same thing as -ffixed-x18 in ops_cpu.py, see comments there for why it's needed on arm osx cpu, feats = ctypes.string_at(llvm.LLVMGetHostCPUName()), (b'+reserve-x18,' if OSX else b'') + ctypes.string_at(llvm.LLVMGetHostCPUFeatures()) diff --git a/tinygrad/runtime/support/compiler_cuda.py b/tinygrad/runtime/support/compiler_cuda.py index 1b27d0265d..e10249ed26 100644 --- a/tinygrad/runtime/support/compiler_cuda.py +++ b/tinygrad/runtime/support/compiler_cuda.py @@ -4,7 +4,7 @@ from tinygrad.helpers import to_char_p_p, colored, init_c_var, getenv import tinygrad.runtime.autogen.nvrtc as nvrtc from tinygrad.device import Compiler, CompileError -PTX, CUDA_PATH = getenv("PTX"), getenv("CUDA_PATH", "") # PTX shouldn't be here, in fact, it shouldn't exist +CUDA_PATH = getenv("CUDA_PATH", "") def _get_bytes(arg, get_str, get_sz, check) -> bytes: sz = init_c_var(ctypes.c_size_t(), lambda x: check(get_sz(arg, ctypes.byref(x)))) diff --git a/tinygrad/runtime/support/hcq.py b/tinygrad/runtime/support/hcq.py index 1b63b9ebed..82823be61f 100644 --- a/tinygrad/runtime/support/hcq.py +++ b/tinygrad/runtime/support/hcq.py @@ -1,11 +1,10 @@ from __future__ import annotations -from typing import cast, Callable, Type, TypeVar, Generic, Any +from typing import cast, Callable, Type, TypeVar, Generic, Any, Sequence import contextlib, decimal, statistics, time, ctypes, array, os, struct, traceback, collections try: import fcntl # windows misses that except ImportError: fcntl = None #type:ignore[assignment] from tinygrad.helpers import PROFILE, getenv, to_mv, round_up, ProfileRangeEvent -from tinygrad.renderer import Renderer -from tinygrad.device import BufferSpec, Compiler, Compiled, LRUAllocator, ProfileDeviceEvent, ProfileProgramEvent +from tinygrad.device import BufferSpec, Compiled, LRUAllocator, ProfileDeviceEvent, ProfileProgramEvent, CompilerPairT from tinygrad.uop.ops import sym_infer, sint, UOp from tinygrad.runtime.autogen import libc @@ -359,12 +358,12 @@ class HCQCompiled(Compiled, Generic[SignalType]): signal_pool: dict[str, list[HCQBuffer]] = collections.defaultdict(list) # per peer group cpu_devices: list[HCQCompiled] = [] - def __init__(self, device:str, allocator:HCQAllocatorBase, renderer:Renderer, compiler:Compiler, runtime, signal_t:Type[SignalType], + def __init__(self, device:str, allocator:HCQAllocatorBase, compilers:Sequence[CompilerPairT], runtime, signal_t:Type[SignalType], comp_queue_t:Callable[[], HWQueue], copy_queue_t:Callable[[], HWQueue]|None=None, kernargs_size=(16 << 20), sigalloc_size=0x1000): self.device_id:int = int(device.split(":")[1]) if ":" in device else 0 from tinygrad.runtime.graph.hcq import HCQGraph - super().__init__(device, allocator, renderer, compiler, runtime, HCQGraph) + super().__init__(device, allocator, compilers, runtime, HCQGraph) # TODO: peer logic is determined based on device name. self.peer_group = device.split(":")[0] @@ -441,14 +440,19 @@ class HCQCompiled(Compiled, Generic[SignalType]): except MemoryError: buf, realloced = self.allocator.alloc(oldbuf.size if oldbuf is not None else new_size, options=options), False return buf, realloced + def _make_no_iface_error(self, errs:str, err_short:str) -> RuntimeError: + # Keep it in a separate function to avoid creating a traceback <-> locals ref cycle + e = RuntimeError(f"No interface for {type(self).__name__[:-6]}:{self.device_id} is available") + if hasattr(e, "add_note"): e.add_note(errs + err_short) + return e + def _select_iface(self, *ifaces:Type): errs, err_short = "", "" if val:=getenv(f'{type(self).__name__[:-6].upper()}_IFACE', ""): ifaces = tuple(x for x in ifaces if x.__name__.startswith(val.upper())) for iface_t in ifaces: try: return iface_t(self, self.device_id) - except Exception as e: errs, err_short = errs + f"\n{iface_t.__name__}: {traceback.format_exc()}", err_short + f"\n{iface_t.__name__}: {e}" - raise RuntimeError(f"{errs}\nNo interface for {type(self).__name__[:-6]}:{self.device_id} is available:{err_short}\n" \ - f"\nForce an interface with {type(self).__name__[:-6].upper()}_IFACE={('|'.join(x.__name__[:-5] for x in ifaces))}.") + except Exception as e: errs, err_short = errs + f"\n{iface_t.__name__}: {traceback.format_exc()}", err_short + f"\n{iface_t.__name__}: {e}." + raise self._make_no_iface_error(errs, err_short) def _is_cpu(self) -> bool: return hasattr(self, 'device') and self.device.split(":")[0] == "CPU" diff --git a/tinygrad/runtime/support/system.py b/tinygrad/runtime/support/system.py index 6238866156..66b2f78615 100644 --- a/tinygrad/runtime/support/system.py +++ b/tinygrad/runtime/support/system.py @@ -165,7 +165,7 @@ class PCIIfaceBase: def map(self, b:HCQBuffer): if b.owner is not None and b.owner._is_cpu(): System.lock_memory(cast(int, b.va_addr), b.size) - paddrs, snooped, uncached = [(x, 0x1000) for x in System.system_paddrs(cast(int, b.va_addr), round_up(b.size, 0x1000))], True, False + paddrs, snooped, uncached = [(x, 0x1000) for x in System.system_paddrs(cast(int, b.va_addr), round_up(b.size, 0x1000))], True, True elif (ifa:=getattr(b.owner, "iface", None)) is not None and isinstance(ifa, PCIIfaceBase): paddrs = [(paddr if b.meta.mapping.system else (paddr + ifa.p2p_base_addr), size) for paddr,size in b.meta.mapping.paddrs] snooped, uncached = b.meta.mapping.snooped, b.meta.mapping.uncached diff --git a/tinygrad/schedule/kernelize.py b/tinygrad/schedule/kernelize.py index cf4db55cb9..ec2f7c046a 100644 --- a/tinygrad/schedule/kernelize.py +++ b/tinygrad/schedule/kernelize.py @@ -120,7 +120,8 @@ def create_kernel(x:UOp, b:UOp|None=None): if b is None: b = UOp.new_buffer(x.device, x.size, x.dtype) kernel = UOp(Ops.KERNEL, src=(b,)+x.src, arg=Kernel(x.sink(), m if (m:=x.metadata) else ())) buffer = b.base if b.size == b.base.size else UOp(Ops.BUFFER_VIEW, b.dtype, (b.base,), (b.size, b.arg.views[0].offset)) - return buffer.assign(kernel).shrink(((0, prod(x.shape)),)).reshape(x.shape) + # we have to shrink the buffer back to the symbolic shape + return buffer.assign(kernel).reshape(tuple(d.vmax if isinstance(d, UOp) else d for d in x.shape)).shrink(tuple((0, d) for d in x.shape)) DONT_PLACE_IN_KERNEL = {Ops.KERNEL, Ops.ASSIGN, Ops.BUFFER, Ops.MSELECT, Ops.MSTACK, Ops.MULTI, Ops.BIND} def append_to_kernel(x:UOp): @@ -148,6 +149,16 @@ create_kernels = PatternMatcher([ lambda ms: UOp(Ops.MSTACK, ms.dtype, tuple(x.src[0] for x in ms.src)).reshape(ms.src[0].arg)), ]) +def add_stores(ctx, sink: UOp): + stores = [] + for i,x in enumerate(sink.src): + gbl = UOp(Ops.DEFINE_GLOBAL, (s:=x.base).dtype.ptr(ctx[i].size), (), i) + # if this is an assign then we already have a buffer with a view that should be the target of the store + if x.op is Ops.ASSIGN: stores.append(UOp.store(gbl.view(unwrap(s.st)), s)) + # otherwise we have to create the shapetracker and shrink it to the correct symbolic shape + else: stores.append( + UOp.store(gbl.reshape(tuple(int(d.vmax) if isinstance(d,UOp) else d for d in s.shape)).shrink(tuple((0,d) for d in s.shape)),s)) + return UOp.sink(*stores, arg=sink.arg) # **** fix kernel AST def unbind_view(x:UOp): @@ -168,9 +179,7 @@ replace_buffers = PatternMatcher([ # no SINK for meta ops (UPat(Ops.SINK, src=(UPat(Ops.CONTIGUOUS, src=(UPat(GroupOp.Meta, name="x"),),))), lambda x:x), # STORE (except for meta ops) - (UPat(Ops.SINK, src=UPat(GroupOp.All-{Ops.STORE}), name="sink"), lambda ctx,sink: - UOp.sink(*[UOp.store(UOp(Ops.DEFINE_GLOBAL, (s:=x.base).dtype.ptr(ctx[i].size), (), i).view(s.st), s) for i,x in enumerate(sink.src)], - arg=sink.arg)), + (UPat(Ops.SINK, src=UPat(GroupOp.All-{Ops.STORE}), name="sink"), add_stores), # remove CONTIGUOUS/DEVICE from kernel AST (UPat((Ops.CONTIGUOUS, Ops.MSELECT), src=(UPat.var("x"),)), lambda x: x), (UPat(Ops.VIEW, src=(UPat(Ops.DEVICE),), name="view"), lambda view: view.replace(src=())), diff --git a/tinygrad/schedule/multi.py b/tinygrad/schedule/multi.py index 4ec0f8e125..7c695db824 100644 --- a/tinygrad/schedule/multi.py +++ b/tinygrad/schedule/multi.py @@ -1,7 +1,8 @@ -from typing import cast +from typing import cast, TypeVar import functools, itertools, operator from tinygrad.helpers import all_same, all_int, prod, DEBUG, RING, getenv, unwrap from tinygrad.uop.ops import Ops, UOp, sint, PatternMatcher, UPat, GroupOp, resolve +from tinygrad.shape.shapetracker import ShapeTracker from tinygrad.device import Device # *** allreduce implementation *** @@ -81,9 +82,10 @@ def handle_allreduce(buf:UOp, red:UOp) -> UOp|None: # ***** multi rewrite MSELECT/MSTACK ***** -def _replace_dnum(st, val): - # replace dnum in ShapeTracker with literal const for this mselect - if (dnums:=[x for x in st.vars() if x.op is Ops.DEFINE_VAR and x.arg[0] == '_device_num']): +T = TypeVar("T", bound=ShapeTracker|sint) +def _replace_dnum(st:T, val:int) -> T: + # replace dnum in ShapeTracker (or UOp) with literal const for this mselect + if not isinstance(st, int) and (dnums:=[x for x in st.vars() if x.op is Ops.DEFINE_VAR and x.arg[0] == '_device_num']): assert len(dnums) == 1, f"view must have exactly 0 or 1 dnum, got {dnums}" st = st.substitute({dnums[0]:dnums[0].const_like(val)}) return st @@ -93,20 +95,23 @@ def mstack_reorder_view(ms:UOp): if not all_same(args) or len([x for x in args[0].vars() if x.arg[0] == '_device_num']) != 0: return None return UOp(Ops.MSTACK, ms.dtype, tuple(x.src[0] for x in ms.src)).view(args[0]) -def mstack_early_shrink(view:UOp, ms:UOp): - if resolve(prod(view.shape) >= prod(ms.shape)) or _replace_dnum(view.st, 0) == view.st: return None +# NOTE: view path is for RANGEIFY=0, there should only be one way of doing this +def mstack_early_shrink(ms:UOp, view:UOp|None=None, shrink:UOp|None=None): + if view is not None and (resolve(prod(view.shape) >= prod(ms.shape)) or _replace_dnum(unwrap(view.st), 0) == view.st): return None ret = [] + def apply_shrink(s:UOp, i:int) -> UOp: + if view is not None: return s.view(_replace_dnum(unwrap(view.st), i)) + return s.shrink(tuple(tuple(_replace_dnum(x, i) for x in ss) for ss in unwrap(shrink).arg)) for i, x in enumerate(ms.src): - new_view = _replace_dnum(view.st, i) if x.op is Ops.COPY: # if src device doesn't have a renderer, we have to view after the copy # TODO: a way to understand this if x.src[0].device in {"DISK", "NPY"}: - ret.append(x.view(new_view)) + ret.append(apply_shrink(x, i)) else: - ret.append(x.src[0].view(new_view).copy_to_device(x.device)) + ret.append(apply_shrink(x.src[0], i).copy_to_device(x.device)) else: - ret.append(x.view(new_view).contiguous()) + ret.append(apply_shrink(x, i).contiguous()) return ms.replace(src=tuple(ret)) replace_allreduce = PatternMatcher([ @@ -127,6 +132,11 @@ replace_allreduce = PatternMatcher([ (UPat(Ops.MSTACK, src=UPat(Ops.VIEW), name="ms"), mstack_reorder_view), # move shrink before MSTACK (UPat(Ops.VIEW, src=(UPat(Ops.MSTACK, name="ms"),), name="view"), mstack_early_shrink), + # *** new movement ops reordering + # move shrink before MSTACK + (UPat(Ops.SHRINK, src=(UPat(Ops.MSTACK, name="ms"),), name="shrink"), mstack_early_shrink), + # move MSELECT before movement ops + (UPat(Ops.MSELECT, src=(UPat(GroupOp.Movement, src=(UPat.var("s"),), name="v"),), name="ms"), lambda s,v,ms: v.replace(src=(s.mselect(ms.arg),))), ]) # ***** multi functions ***** @@ -210,7 +220,7 @@ def assign_multi(dest:UOp, src:UOp): return dest.src[0].assign(src.src[0]).multi(src.axis) def passthrough_multi(root:UOp, multi:UOp): - return root.replace(src=(multi.src[0],)).multi(multi.axis) + return UOp(root.op, root.dtype, (multi.src[0],), root.arg).multi(multi.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 9883deb126..dacf9f4f52 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -1,46 +1,61 @@ -from typing import Any +from typing import Any, cast +import functools, operator from dataclasses import dataclass, field from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace -from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, RewriteNotReady, _substitute -from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, RANGEIFY +from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, RewriteNotReady, _substitute, ssimplify, graph_rewrite_map +from tinygrad.uop.symbolic import sym, symbolic_simple +from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, RANGEIFY, Context, flatten, dedup from tinygrad.schedule.multi import multi_pm from tinygrad.schedule.kernelize import Kernel -from tinygrad.uop.ops import track_rewrites, graph_rewrite_map, graph_rewrite, identity_element, sint, AxisType +from tinygrad.uop.ops import track_rewrites, graph_rewrite, identity_element, sint, AxisType +# ***************** # 0. do some cleanup rewrites, mostly copied from the old stuff double_reshape = PatternMatcher([ # RESHAPE on RESHAPE is the second reshape - (UPat(Ops.RESHAPE, src=(UPat(Ops.RESHAPE),), name="x"), lambda x: x.replace(src=(x.src[0].src[0],))), + (UPat(Ops.RESHAPE, src=(UPat(Ops.RESHAPE),), name="x"), lambda x: x.replace(src=(x.src[0].src[0],), tag=(x.src[0].tag or ())+(x.tag or ()))), ]) earliest_rewrites = double_reshape+PatternMatcher([ - # UOp with size 0 is zero - (UPat(GroupOp.All-{Ops.SINK}, name="root"), lambda root: root.const_like(0) if root.base.st is not None and root.size == 0 else None), + # non shape changing RESHAPE is NOOP + #(UPat(Ops.RESHAPE, name="x"), lambda x: x.src[0] if x.src[0].shape == x.arg else None), # DETACH and CONTIGUOUS_BACKWARD are NOOPs here, so is FUSE + #(UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD, Ops.FUSE), name="x"), lambda x: x.src[0].f(Ops.NOOP, tag=x.tag)), + + # just removing it works... (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD, Ops.FUSE), name="x"), lambda x: x.src[0]), + + # preserve tags? # reduce of size 0 is the identity element (UPat(Ops.REDUCE_AXIS, name="reduce", src=(UPat.var("x"),)), lambda reduce,x: reduce.const_like(identity_element(reduce.arg[0], reduce.dtype)) if x.size == 0 and reduce.size != 0 else None), - # non shape changing RESHAPE is NOOP - (UPat(Ops.RESHAPE, name="x"), lambda x: x.src[0] if x.src[0].shape == x.arg else None), - # RESHAPE after COPY - (UPat(Ops.COPY, src=(UPat(Ops.RESHAPE, name="r"),UPat(name="d")), name="c"), lambda c,r,d: c.replace(src=(r.src[0],d)).reshape(r.arg)), - # TODO: this should be BUFFER_VIEW - (UPat(Ops.COPY, src=(UPat(Ops.SHRINK, name="r"),UPat(name="d")), name="c"), lambda c,r,d: c.replace(src=(r.src[0],d)).shrink(r.arg)), - # const hacks - (UPat(Ops.CONST, name="x"), lambda x: - x.replace(src=(x.src[0].src[0],)).reshape((1,)*len(x.shape)).expand(x.shape) if \ - len(x.src) and x.src[0].op is Ops.VIEW and not any(s == 0 for s in x.shape) else None), + + # COPY and source size need to match + # TODO: expand after copy creates issues with tagging + (UPat(Ops.COPY, src=(UPat(GroupOp.Movement, name="r"), UPat(name="d")), name="c"), + lambda c,r,d: c.replace(src=(r.contiguous(), d)) if r.size != r.base.size else None), + # assign only to buffer - (UPat(Ops.ASSIGN, src=(UPat(GroupOp.All-{Ops.BUFFER}, name="target"), UPat(name="x"))), - lambda x,target: x if target.base.op is not Ops.BUFFER else None), + (UPat(Ops.ASSIGN, src=(UPat(GroupOp.All-{Ops.BUFFER}, name="target"), UPat(name="x")), name="assign"), + lambda x,target,assign: x.f(Ops.NOOP, tag=assign.tag) if target.base.op is not Ops.BUFFER else None), + + # copy only to different device + (UPat(Ops.COPY, src=(UPat.var("x"), UPat()), name="copy"), lambda x,copy: x.f(Ops.NOOP, tag=copy.tag) if x.device == copy.device else None), + + # handle disk + # TODO: this doesn't need to use st.views + (UPat.var("x").f((Ops.BITCAST, Ops.CONTIGUOUS), name="t"), + lambda x,t: UOp(Ops.BUFFER_VIEW, t.dtype, (x.base,), (t.size, x.st.views[0].offset), tag=t.tag).reshape(t.shape) if isinstance(x.device, str) \ + and x.device.startswith("DISK") else None), + # contiguous/buffer/copy/assign is already contiguous - (UPat(Ops.CONTIGUOUS, name="root", src=(UPat((Ops.CONTIGUOUS, Ops.BUFFER, Ops.COPY, Ops.ASSIGN)),)), lambda root: root.src[0]), + #(UPat(Ops.CONTIGUOUS, name="root", src=(UPat((Ops.CONTIGUOUS, Ops.BUFFER, Ops.COPY, Ops.ASSIGN)),)), lambda root: root.src[0]), ]) -# 1. add contiguous where we have to +# ***************** +# 1. add realize where we have to ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.ASSIGN, Ops.COPY, Ops.BUFFER, Ops.BUFFER_VIEW, Ops.CONST, Ops.BIND, Ops.DEVICE, Ops.MSELECT, Ops.MSTACK, Ops.DEFINE_GLOBAL, @@ -50,7 +65,7 @@ def realize(ctx:dict[UOp, None], tr:UOp) -> None: ctx[tr] = None def realize_parents(ctx:dict[UOp, None], rb:UOp) -> None: for s in rb.src: - if s.op not in ALWAYS_CONTIGUOUS: ctx[s] = None + if s.base.op not in ALWAYS_CONTIGUOUS: ctx[s] = None def realize_assign(ctx:dict[UOp, None], a:UOp) -> None: if a.src[1].op not in ALWAYS_CONTIGUOUS: ctx[a.src[1]] = None @@ -58,19 +73,25 @@ def realize_assign(ctx:dict[UOp, None], a:UOp) -> None: do_realize = PatternMatcher([ # always realize SINK parents (UPat(Ops.SINK, name="s"), lambda ctx,s: ctx.update((x.base, None) for x in s.src if x.base.op not in ALWAYS_CONTIGUOUS)), - # always realize ASSIGN/COPY/BUFFER_VIEW - (UPat({Ops.ASSIGN, Ops.COPY, Ops.BUFFER_VIEW}, name="tr"), realize), + # always realize ASSIGN/COPY/BUFFER_VIEW/CONTIGUOUS + (UPat({Ops.ASSIGN, Ops.COPY, Ops.BUFFER_VIEW, Ops.CONTIGUOUS}, name="tr"), realize), # realize parents of COPY, MSELECT, MSTACK (UPat((Ops.COPY, Ops.MSELECT, Ops.MSTACK), name="rb"), realize_parents), # realize input to assign (might be optimized out) (UPat(Ops.ASSIGN, name="a"), realize_assign), ]) -add_contiguous = PatternMatcher([ - (UPat(GroupOp.All-{Ops.CONTIGUOUS}, name="x"), lambda ctx,x: x.replace(tag=1).contiguous() if x in ctx and x.tag is None else None), -]) -remove_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)]) +class WrappedContig: + def __init__(self, x): self.x = x + def __repr__(self): return f"C({self.x})" +add_contiguous = PatternMatcher([ + (UPat(GroupOp.All, name="x"), + lambda ctx,x: x.replace(tag=WrappedContig(x.tag)).realize() if x in ctx and not isinstance(x.tag, WrappedContig) else None), +]) +remove_contig_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=x.tag.x) if isinstance(x.tag, WrappedContig) else None)]) + +# ***************** # 2. mark all children @dataclass @@ -94,10 +115,11 @@ def mark_children(ctx:ChildrenContext, x:UOp): pm_children = PatternMatcher([ (UPat(Ops.SINK, name="x"), extract_children), - (UPat(GroupOp.All-{Ops.CHILD, Ops.CHILDREN}, name="x"), mark_children), + (UPat(GroupOp.All-{Ops.CHILD, Ops.CHILDREN, Ops.SINK}, name="x"), mark_children), ]) -# 3. rangeify +# ***************** +# 3a. rangeify (movement) @dataclass class RangeifyContext: @@ -136,9 +158,8 @@ def map_pad(idx:UOp, r:UOp): if resolve(e > 0): where = where & (ret[i] < (sh-e)) if resolve(s > 0): where = where & (ret[i] >= s) bigwhere = bigwhere & where - # this is safe but dumb - # TODO (S-Lykles): switch to mixed index/valid - ret[i] = (ret[i] - s).maximum(0).minimum(r.src[0].shape[i]-1) + with Context(TRACK_MATCH_STATS=0): + ret[i] = graph_rewrite(where.where(ret[i]-s, UOp.invalid()), sym) # PAD is with 0 return bigwhere.simplify().where(r.src[0].index(*ret, dtype=idx.dtype, arg=idx.arg), UOp.const(r.dtype, 0)) @@ -148,12 +169,12 @@ def map_expand(r:UOp, idx:UOp): non_ending_ranges = [] for a,x,y in zip(idx.src[1:], r.src[0].shape, r.shape): axis_to_range = [u for u in a.toposort() if u.op is Ops.RANGE] - if resolve(x!=y, False): - ending_ranges.extend(axis_to_range) - new_rngs.append(a.const_like(0)) - else: + if resolve(x==y, False): non_ending_ranges.extend(axis_to_range) new_rngs.append(a) + else: + ending_ranges.extend(axis_to_range) + new_rngs.append(a.const_like(0)) ending_ranges = [x.arg for x in ending_ranges if x not in non_ending_ranges] if idx.arg is not None: ending_ranges.append(idx.arg) return r.src[0].index(*new_rngs, arg=min(ending_ranges) if ending_ranges else None) @@ -174,7 +195,21 @@ pm_mops = PatternMatcher([ (UPat(Ops.PAD, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"), map_pad), ]) -def map_partial_contiguous(ctx:RangeifyContext, x:UOp, idx:UOp): +# ***************** +# 3b. rangeify (ops) + +# bufferization can happen in three ways +# 1. there's an explicit REALIZE in the graph +# 2. the ranges from the children don't match and we have to create a buffer (only on children) +# 3. might_end_axis triggers because we should be closing a loop to save compute + +@dataclass(frozen=True) +class BufferizeOpts: + # on AddrSpace.LOCAL, device is the id + device: str|tuple[str, ...]|int|None + addrspace: AddrSpace = AddrSpace.GLOBAL + +def map_partial_realize(ctx:RangeifyContext, x:UOp, idx:UOp): if x.arg is None: return None # map_contiguous can handle this # NOTE: all partial contiguous can safely be replaced by full contiguous. we should be able to match old functionality like this if not (RANGEIFY > 1): return idx.replace(src=(x.replace(arg=None),)+idx.src[1:]) @@ -186,18 +221,17 @@ def map_partial_contiguous(ctx:RangeifyContext, x:UOp, idx:UOp): ranges.append(idx.src[1+i]) continue passthrough_idx.append(idx.src[1+i]) - ranges.append(ctx.new_range(s) if resolve(s!=1) else UOp.const(dtypes.index, 0)) + ranges.append(ctx.new_range(s)) new_ranges.append(ranges[-1]) - ret = x.src[0].index(*ranges).bufferize(*[x for x in new_ranges if x.op is not Ops.CONST], arg=x.device) + # TODO: this should be able to be global or local + ret = x.src[0].index(*ranges).bufferize(*[x for x in new_ranges if x.op is not Ops.CONST], + arg=BufferizeOpts(device=None, addrspace=AddrSpace.LOCAL)) return ret.index(*passthrough_idx) -def map_contiguous(ctx:RangeifyContext, x:UOp): +def map_realize(ctx:RangeifyContext, x:UOp): if x.arg is not None: return None - ranges = [] - for s in x.shape[len(x.src)-1:]: - ranges.append(ctx.new_range(s) if resolve(s!=1) else UOp.const(dtypes.index, 0)) - ret = x.src[0].index(*ranges).bufferize(*x.src[1:], *[x for x in ranges if x.op is not Ops.CONST], arg=x.device) - return ret.shrink(((0, prod(x.shape)),)).forced_reshape(x.shape) + ranges = [ctx.new_range(s) for s in x.shape] + return x.src[0].index(*ranges).bufferize(*x.src[1:], *ranges, arg=BufferizeOpts(device=x.device), tag=x.src[0].tag) def map_reduce(ctx:RangeifyContext, idx:UOp, red:UOp): rngs = list(idx.src[1:]) @@ -206,7 +240,7 @@ def map_reduce(ctx:RangeifyContext, idx:UOp, red:UOp): if i in red.arg[1]: rngs[i] = ctx.new_range(s, axistype=AxisType.REDUCE) new_ranges.append(rngs[i]) - return UOp(Ops.REDUCE, red.dtype, src=(red.src[0].index(*rngs),)+tuple(new_ranges), arg=red.arg[0]) + return UOp(Ops.REDUCE, red.dtype, src=(red.src[0].index(*rngs),)+tuple(new_ranges), arg=red.arg[0], tag=red.tag) def index_child(ctx:RangeifyContext, c:UOp, x:UOp, idx:UOp): if c not in ctx.seen_children: ctx.seen_children[c] = {} @@ -220,26 +254,38 @@ def index_child(ctx:RangeifyContext, c:UOp, x:UOp, idx:UOp): ctx.progress = 0 if c not in ctx.seen_child: - all_rngs = zip(*[ch.src[1:] for ch in ctx.seen_children[c].values()]) + all_rngs = list(zip(*[ch.src[1:] for ch in ctx.seen_children[c].values()])) out_rngs = [] end_ranges = [] idx_ranges = [] - for i,r in enumerate(all_rngs): - if all_same(r): - out_rngs.append(r[0]) + # NOTE: locals aren't working, so we only fully bufferize here (unless RANGEIFY > 1) + all_all_same = all(all_same(r) for r in all_rngs) + for i,valid_rngs in enumerate(all_rngs): + rngs, valids = zip(*[(r.get_idx(), r.get_valid()) for r in valid_rngs]) + # we compare the ranges without their valids + if all_same(rngs) and (all_all_same or RANGEIFY > 1): + # the new valid is the OR of all the children valids + minimum_valid = functools.reduce(operator.or_, valids, UOp.const(dtypes.bool, False)) + out_rngs.append(minimum_valid.where(rngs[0], UOp.invalid()).simplify()) else: out_rngs.append(ctx.new_range(c.shape[i])) end_ranges.append(out_rngs[-1]) idx_ranges.append(i) - ctx.seen_child[c] = (idx_ranges, end_ranges) + ctx.seen_child[c] = (out_rngs, idx_ranges, end_ranges) else: - out_rngs = list(idx.src[1:]) - idx_ranges, end_ranges = ctx.seen_child[c] + out_rngs, idx_ranges, end_ranges = ctx.seen_child[c] for i,nr in zip(idx_ranges, end_ranges): out_rngs[i] = nr # index based on the shared ranges ret = c.index(*out_rngs) # if all ranges aren't the same between children, we have to bufferize - if len(idx_ranges) > 0: ret = ret.bufferize(*end_ranges, arg=x.device).index(*[idx.src[1+i] for i in idx_ranges]) + if len(idx_ranges) > 0: + if len(idx_ranges) == len(out_rngs): + # this is a global bufferize + ret = ret.bufferize(*end_ranges, arg=BufferizeOpts(device=x.device)) + else: + assert RANGEIFY > 1, "this isn't supported with RANGEIFY=1" + ret = ret.bufferize(*end_ranges, arg=BufferizeOpts(device=None, addrspace=AddrSpace.LOCAL)) + ret = ret.index(*[idx.src[1+i] for i in idx_ranges]) return ret def children_gate(ctx:RangeifyContext, idx:UOp, c:UOp): @@ -249,20 +295,22 @@ def children_gate(ctx:RangeifyContext, idx:UOp, c:UOp): def might_end_axis(idx:UOp): if idx.arg is None: return None # TODO: write a proper cost function here - if all(x.op not in {Ops.BUFFER, Ops.CONTIGUOUS, Ops.BUFFERIZE} for x in idx.toposort()): return None + if all(x.op not in {Ops.BUFFER, Ops.REALIZE, Ops.BUFFERIZE} for x in idx.toposort()): return None if all(x.op not in {Ops.REDUCE_AXIS} for x in idx.toposort()): return None to_end_axis = [] for i,a in enumerate(idx.src[1:]): if any(x.arg > idx.arg for x in a.toposort() if x.op is Ops.RANGE): to_end_axis.append(i) - if to_end_axis: return idx.replace(src=(idx.src[0].contiguous(arg=tuple(to_end_axis)),)+idx.src[1:], arg=None) + if to_end_axis: return idx.replace(src=(idx.src[0].realize(arg=tuple(to_end_axis)),)+idx.src[1:], arg=None) return idx.replace(arg=None) +def unprocessed_index(x:UOp): raise RuntimeError(f"unprocessed index on {x.src[0].op}") + pm_rangeify = pm_mops+PatternMatcher([ # sink contigs to kick it off - (UPat(Ops.CONTIGUOUS, src=(UPat(),), name="x", allow_any_len=True), map_contiguous), + (UPat(Ops.REALIZE, src=(UPat(),), name="x", allow_any_len=True), map_realize), # if there's an INDEX it can support partial contig - (UPat(Ops.INDEX, src=(UPat(Ops.CONTIGUOUS, src=(UPat(),), name="x"),), allow_any_len=True, name="idx"), map_partial_contiguous), + (UPat(Ops.INDEX, src=(UPat(Ops.REALIZE, src=(UPat(),), name="x"),), allow_any_len=True, name="idx"), map_partial_realize), # if there are new ended children, tag the SINK (UPat(Ops.INDEX, src=(UPat(Ops.CHILD, src=(UPat(name="c"), ), name="x"),), allow_any_len=True, name="idx"), index_child), @@ -271,29 +319,39 @@ pm_rangeify = pm_mops+PatternMatcher([ # if we come across this, remove it. it was a CHILD unused in an INDEX (UPat(Ops.CHILD, src=(UPat(Ops.CHILDREN, src=(UPat.var("x"),)),)), lambda x: x), - # CONST (or DEFINE_VAR) can't have axes. remove srcs when we INDEX it + # CONST (or DEFINE_VAR) can't have axes. remove INDEX when we get here (UPat(Ops.INDEX, src=(UPat((Ops.CONST, Ops.DEFINE_VAR), name="c"),)), lambda c: c.replace(src=())), # handle arg on any op with weight. old endrange stuff (UPat(Ops.INDEX, src=(UPat(GroupOp.Elementwise.union({Ops.REDUCE_AXIS})),), allow_any_len=True, name="idx"), might_end_axis), + # handle size 0 + (UPat(Ops.INDEX, name="x"), lambda x: x.replace(src=(x.const_like(0),)+x.src[1:]) if x.st is not None and x.size == 0 else None), + + # handle assign + (UPat(Ops.INDEX, src=(UPat(Ops.ASSIGN, name="assign"),), allow_any_len=True, name="x"), + lambda x,assign: assign.replace(src=tuple([s.index(*x.src[1:]) for s in assign.src])+(assign.src[0],))), + # move MAP through elementwise ALU / reduce. these are the items with cost - (UPat(Ops.INDEX, src=(UPat(GroupOp.Elementwise.union({Ops.STORE, Ops.ASSIGN, Ops.COPY, Ops.DEVICE, Ops.BIND})),), allow_any_len=True, name="x"), + (UPat(Ops.INDEX, src=(UPat(GroupOp.Elementwise.union( + {Ops.STORE, Ops.COPY, Ops.BUFFER_VIEW, Ops.DEVICE, Ops.BIND, Ops.CONTIGUOUS, Ops.NOOP})),), allow_any_len=True, name="x"), lambda x: x.src[0].replace(src=tuple([s.index(*x.src[1:]) for s in x.src[0].src]))), (UPat(Ops.INDEX, src=(UPat(Ops.REDUCE_AXIS, name="red"),), allow_any_len=True, name="idx"), map_reduce), + + # assert if there's any index we didn't process + (UPat(GroupOp.All-{Ops.REALIZE, Ops.BUFFERIZE, Ops.MSELECT}).f(Ops.INDEX, name="x"), unprocessed_index), ]) +# ***************** # 3.5 cleanups # you don't know in the first pass if axes are going to die, this happens if there's an EXPAND to the left -# TODO: figure out how to reenable this def cleanup_dead_axes(b:UOp): - parents = b.src[0].toposort() new_rng = [] hit = False reshape: list[sint] = [] for s,rng in zip(b.shape, b.src[1:]): - if rng not in parents and rng.op is Ops.RANGE: + if rng not in b.src[0].sparents and rng.op is Ops.RANGE: reshape.append(1) hit = True else: @@ -304,25 +362,50 @@ def cleanup_dead_axes(b:UOp): # if a buffer is being stored just for permutes or something, remove it # we want to reexpress the indexes of idx2 in terms of the implied b1 -def remove_bufferize(b2:UOp, idx2:UOp): - # HACK - if len(b2.src) != len(idx2.src): return None - assert len(b2.src) == len(idx2.src) - assert all(x.op is Ops.RANGE for x in b2.src[1:]) - return b2.src[0].substitute(dict(zip(b2.src[1:], idx2.src[1:]))) +def remove_bufferize(src:UOp, buf:UOp, idx:UOp): + # see if we can't do it, should this ever hit? + assert len(buf.src) == len(idx.src), "index on wrong bufferize" + assert all(x.op is Ops.RANGE for x in buf.src[1:]) + + # if it's user contiguous, we never remove it + if src.op is Ops.CONTIGUOUS: return None + + # here is where we compute the cost + # for now just no REDUCE, COPY, or ASSIGN + ran = src.toposort(gate=lambda x: x.op not in {Ops.INDEX}) + # we don't want to bufferize threefry, also causes problems because not all platforms support long + if any(x.op in {Ops.REDUCE, Ops.COPY, Ops.BUFFER_VIEW, Ops.ASSIGN} for x in ran) and src.op is not Ops.THREEFRY: return None + + # simple, matching old behavior + #if src.op is not Ops.INDEX: return None + + # this is the ranges replaced + return src.substitute(dict(zip(buf.src[1:], idx.src[1:]))) + +def pre_bufferize(b:UOp, x:UOp, copy:UOp): + nb = b.replace(src=(b.src[0].contiguous(),)+b.src[1:]) + return copy.replace(src=(x.replace(src=(nb,)+x.src[1:]), copy.src[1])) pm_cleanups = double_reshape+pm_mops+PatternMatcher([ #(UPat(Ops.BUFFERIZE, name="b"), cleanup_dead_axes), # remove noop buffers. if we look at the next index we can remove even more of these # NOTE: this is mostly the same case as below, but if there's no INDEX this gets more - #(UPat(Ops.INDEX, name="idx").f(Ops.BUFFERIZE, allow_any_len=True, name="b2"), - # lambda idx,b2: idx.src[0] if idx.src[1:] == b2.src[1:] else None), - # remove reindexing - (UPat(Ops.INDEX).f(Ops.BUFFERIZE, allow_any_len=True, name="b2").f(Ops.INDEX, allow_any_len=True, name="idx2"), remove_bufferize), + (UPat(Ops.INDEX, name="idx").f(Ops.BUFFERIZE, allow_any_len=True, name="b2"), + lambda idx,b2: idx.src[0].replace(tag=nt if len(nt:=(idx.src[0].tag or ()) + (b2.tag or ())) else None) if idx.src[1:] == b2.src[1:] \ + and idx.src[0].op is not Ops.BUFFER_VIEW else None), + # remove reindexing with cost function + (UPat.var("src").f(Ops.BUFFERIZE, allow_any_len=True, name="buf").f(Ops.INDEX, allow_any_len=True, name="idx"), remove_bufferize), # no buffers for const - #(UPat(Ops.CONST, name='c').f(Ops.BUFFERIZE, allow_any_len=True, name="b"), lambda c,b: c.reshape((1,)*len(b.shape)).expand(b.shape)), + (UPat(Ops.CONST, name='c').f(Ops.BUFFERIZE, allow_any_len=True, name="b"), lambda c,b: b.const_like(c.arg).rtag(b.tag)), + # if any CONST with DEVICE make it here (symbolic/copy issue), remove it + #(UPat(Ops.DEVICE).f(Ops.CONST, name="c"), lambda c: c.replace(src=())), + # copy on CONST is CONST + (UPat(Ops.COPY, src=(UPat.cvar("x"), UPat()), name="copy"), lambda copy,x: copy.const_like(x.arg)), + (UPat(Ops.COPY, src=(UPat(GroupOp.All-{Ops.CONTIGUOUS, Ops.COPY}).f(Ops.BUFFERIZE, allow_any_len=True, name="b") + .f(Ops.INDEX, allow_any_len=True, name="x"), UPat()), name="copy"), pre_bufferize), ]) +# ***************** # 4. put in buffers for bufferize # TODO: should BUFFERIZE look a lot more like STORE # BUFFERIZE has device in arg @@ -330,36 +413,54 @@ pm_cleanups = double_reshape+pm_mops+PatternMatcher([ # BUFFERIZE returns the BUFFER ready for INDEXing (doing this will make splitting a lot easier) # NOTE: this has been fixed up a bit -def bufferize_to_store(x:UOp, locals_allowed=False): +def bufferize_to_store(x:UOp): rngs = x.src[1:] shape = tuple([int(r.vmax+1) for r in rngs]) + sym_shape = tuple([ssimplify(r.src[0]) for r in rngs]) size = prod(shape) assert size > 0, f"no zero sized buffers {shape}" - sdtype = x.dtype.ptr(size=size, addrspace=AddrSpace.GLOBAL if not isinstance(x.arg, tuple) else x.arg[0]) + + sdtype = x.dtype.ptr(size=size, addrspace=x.arg.addrspace) if x.src[0].op is Ops.ASSIGN: - assign_target, assign_src = x.src[0].src - assert assign_target.op is Ops.INDEX - return assign_target.replace(dtype=sdtype).store(assign_src, *rngs, dtype=sdtype) + assign_target, assign_src, assign_mops = x.src[0].src + assert assign_target.op is Ops.INDEX, f"{assign_target.op} is not index" + # in assign, this is the buffer size, not the bufferize size + # TODO: assign_mops here + ret = assign_target.replace(dtype=sdtype).store(assign_src, *rngs, dtype=x.dtype) + mops = [] + walk = assign_mops + while walk is not assign_mops.base: + mops.append((walk.op, walk.arg)) + walk = walk.src[0] + for m in mops[::-1]: ret = ret._mop(*m) + return ret.forced_reshape(shape).replace(tag=x.tag) + # NOTE: the DEFINE_LOCAL needs to be disambiguated here if sdtype.addrspace == AddrSpace.GLOBAL: - buf = UOp.new_buffer(x.arg, size, x.dtype) - else: - if not locals_allowed: return None - buf = UOp(Ops.DEFINE_LOCAL, sdtype, arg=x.arg[1]) - return buf.reshape(shape).index(*rngs, dtype=sdtype).store(x.src[0], *rngs, dtype=sdtype).forced_reshape(shape, dtype=x.dtype) + buf = UOp.new_buffer(x.arg.device, size, x.dtype) + ret = buf.reshape(shape).index(*rngs, dtype=sdtype).store(x.src[0], *rngs, dtype=x.dtype) + ret = ret.forced_reshape(shape) + # TODO: is this right? what if it's offset + if shape is not sym_shape: ret = ret.shrink(tuple([(0,x) for x in sym_shape])) + return ret.replace(tag=x.tag) -pm_add_buffers_local = pm_mops+PatternMatcher([ - (UPat(Ops.BUFFERIZE, name="x"), lambda x: bufferize_to_store(x, True)), -]) + # handle locals + tag = x.arg.device + if tag is None: tag = UOp.unique().arg # TODO: hack + buf = UOp(Ops.DEFINE_LOCAL, sdtype, arg=tag) + # store has the other dtype here + # TODO: how is this unified? + return buf.reshape(shape).index(*rngs, dtype=sdtype).store(x.src[0], *rngs, dtype=sdtype).forced_reshape(shape, dtype=x.dtype) pm_add_buffers = pm_mops+PatternMatcher([ (UPat(Ops.BUFFERIZE, name="x"), bufferize_to_store), # move RESHAPEs through MSELECT/MSTACK (UPat((Ops.MSELECT, Ops.MSTACK), src=UPat(Ops.RESHAPE), name="m"), - lambda m: m.replace(src=tuple([x.src[0] for x in m.src])).reshape(m.src[0].arg)), + lambda m: m.replace(src=tuple([x.src[0] for x in m.src]), tag=None).reshape(m.src[0].arg).rtag(m.tag)), ]) +# ***************** # 5. split into kernels @dataclass @@ -367,6 +468,7 @@ class LocalAddBufferContext: dg:int = 0 map:dict = field(default_factory=dict) vars:dict = field(default_factory=dict) + range:int = 0 def debuf(ctx:LocalAddBufferContext, buf:UOp): ret = UOp(Ops.DEFINE_GLOBAL, buf.dtype.ptr(buf.arg), arg=ctx.dg) @@ -386,6 +488,12 @@ def handle_assign(ctx:LocalAddBufferContext, assign:UOp): ctx.map[buf] = assign return buf +def renumber_range(ctx:LocalAddBufferContext, r:UOp): + if r.tag is not None: return None + ret = r.replace(arg=(ctx.range,)+r.arg[1:], tag=()) + ctx.range += 1 + return ret + to_define_global = PatternMatcher([ (UPat(Ops.BUFFER, name="buf"), debuf), (UPat(Ops.BIND, name="b"), unbind_kernel), @@ -393,10 +501,20 @@ to_define_global = PatternMatcher([ # HACK in case any CONSTs were replaced # this is only needed if you are using symbolic - #(UPat(Ops.CONST, name="c"), lambda c: c.replace(src=()) if len(c.src) else None), + (UPat((Ops.CONST, Ops.DEFINE_VAR), name="c"), lambda c: c.replace(src=()) if len(c.src) else None), + + # renumber the ranges starting with 0 so that kernel deduping works + (UPat(Ops.RANGE, name="r"), renumber_range), ]) rangeify_codegen = PatternMatcher([ + # no NOOP in the kernel graph + # TODO: this can be moved into codegen? + (UPat((Ops.NOOP, Ops.CONTIGUOUS), name="x"), lambda x: x.src[0]), + + # strip the arg from store + (UPat(Ops.STORE, name="x"), lambda x: x.replace(arg=None) if x.arg is not None else None), + # add loads to non ptr indexes # TODO: this can be moved into codegen? (UPat((Ops.DEFINE_GLOBAL, Ops.STORE), name="dg").f(Ops.INDEX, name="idx", allow_any_len=True), @@ -411,41 +529,74 @@ rangeify_codegen = PatternMatcher([ lambda src, barrier, gate: src.load(UOp(Ops.IF, src=(gate, barrier)))), ]) -def split_store(x:UOp): +def split_store(ctx:list[UOp], x:UOp): if len(x.ranges): return None - ctx = LocalAddBufferContext() - ret = graph_rewrite(x, to_define_global+rangeify_codegen, ctx=ctx, name="kernel split", bottom_up=True) + if x.src[0].ptrdtype.addrspace is AddrSpace.LOCAL: return None + + # local kernel rewrite + lctx = LocalAddBufferContext() + ret = graph_rewrite(x, to_define_global+rangeify_codegen, ctx=lctx, name="kernel split", bottom_up=True) + + # gather the metadata + metadatas = [ctx[y].metadata for x in ret.sparents if x.tag is not None for y in x.tag] # NOTE: the hack for COPY is here - ret = ret.sink() if ret.src[1].op is not Ops.COPY else ret.src[1] - kernel = UOp(Ops.KERNEL, src=tuple(ctx.map.values())+tuple(ctx.vars.keys()), arg=Kernel(ret,())) + ret = ret.sink() if ret.src[1].op not in {Ops.COPY, Ops.BUFFER_VIEW} else ret.src[1] + kernel_arg = Kernel(ret,tuple(dedup(flatten([x for x in metadatas if x is not None])))) + kernel = UOp(Ops.KERNEL, src=tuple(lctx.map.values())+tuple(lctx.vars.keys()), arg=kernel_arg) return x.as_buf().assign(kernel) split_kernels = PatternMatcher([ (UPat(Ops.STORE, name="x"), split_store), ]) -@track_rewrites(name=lambda sink,ret: f"Schedule {pluralize('Kernel',len([u for u in ret[sink].toposort() if u.op is Ops.KERNEL]))}", replay=True) -def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]: - tensor_map = graph_rewrite_map(sink, multi_pm+earliest_rewrites, name="earliest") - realize_map: dict[UOp, UOp] = {} - graph_rewrite(tensor_map[sink], do_realize, ctx=realize_map, name="Input Graph") - tensor_map = graph_rewrite_map(tensor_map[sink], add_contiguous, ctx=realize_map, bottom_up=True, input_map=tensor_map, name="add contiguous") - tensor_map = graph_rewrite_map(tensor_map[sink], remove_tags, input_map=tensor_map, name="cleanup") - tensor_map = graph_rewrite_map(tensor_map[sink], pm_children, ctx=ChildrenContext(), bottom_up=True, input_map=tensor_map, name="children") - tensor_map = graph_rewrite_map(tensor_map[sink], pm_rangeify, ctx=RangeifyContext(), bottom_up=True, input_map=tensor_map, name="rangeify") - # NOTE: running symbolic can break the graph, leaving RANGE/INDEX/BUFFERIZE in the final graph - #tensor_map = graph_rewrite_map(tensor_map[sink], symbolic_simple, input_map=tensor_map, name="symbolic") - tensor_map = graph_rewrite_map(tensor_map[sink], pm_cleanups, bottom_up=True, input_map=tensor_map, name="cleanups") - if getenv("VIZ"): graph_rewrite(tensor_map[sink], PatternMatcher([]), name="View Rangeify Graph") +def tag_uop(ctx:list[UOp], x:UOp): + if x.tag is not None: return None + ctx.append(x) + return x.replace(tag=(len(ctx)-1,)) +add_tags = PatternMatcher([ + # don't tag BUFFERs, they are global + (UPat(GroupOp.All-{Ops.BUFFER, Ops.CONST, Ops.DEVICE, Ops.UNIQUE, Ops.DEFINE_VAR, Ops.BIND}.union(GroupOp.Movement), name="x"), tag_uop), +]) - tensor_map = graph_rewrite_map(tensor_map[sink], pm_add_buffers, bottom_up=True, input_map=tensor_map, name="add buffers") - tensor_map = graph_rewrite_map(tensor_map[sink], split_kernels, input_map=tensor_map, name="split kernels") +@track_rewrites(lambda _,ret: f"Schedule {pluralize('Kernel', len([u for u in UOp.sink(*ret.values()).toposort() if u.op is Ops.KERNEL]))}", True) +def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]: + uop_list: list[UOp] = [] + tsink = graph_rewrite(sink, add_tags, ctx=uop_list, bottom_up=True, name="number the uops") + + # HACKS: handle multi with graph_rewrite_map in order to not have to add all the tag logic to multi + msink = graph_rewrite_map(tsink, multi_pm, name="multi") + tsink = msink[tsink].substitute({v:v.rtag(k.tag) for k,v in msink.items() if v.tag is None and k.tag is not None}) + + tsink = graph_rewrite(tsink, earliest_rewrites, name="earliest rewrites") + realize_map: dict[UOp, UOp] = {} + graph_rewrite(tsink, do_realize, ctx=realize_map, name="Input Graph") + # NOTE: we don't use contiguous here, contiguous is a user op + tsink = graph_rewrite(tsink, add_contiguous, ctx=realize_map, bottom_up=True, name="add realize") + tsink = graph_rewrite(tsink, remove_contig_tags, name="remove contiguous tags") + tsink = graph_rewrite(tsink, pm_children, ctx=ChildrenContext(), bottom_up=True, name="get children") + + # rangeify + tsink = graph_rewrite(tsink, pm_rangeify, ctx=RangeifyContext(), bottom_up=True, name="rangeify") + # NOTE: sym (vs symbolic_simple) breaks things here because ranges with len 1 aren't handled right + tsink = graph_rewrite(tsink, symbolic_simple, name="symbolic") # this supports const folding + tsink = graph_rewrite(tsink, pm_cleanups, bottom_up=True, name="remove costly buffers") + + # rebuild the sink with all the BUFFERIZEs with tags, this is what's ending up in the tensor graph + # MSTACK stacks multiple BUFFERIZEs in one tagged tensor + # if it's not tagged by here, it's out + tsink = UOp.sink(*[x for x in tsink.parents if x.base.op in {Ops.BUFFERIZE, Ops.MSTACK, Ops.CONST} and x.tag is not None]) + + if getenv("VIZ"): graph_rewrite(tsink, PatternMatcher([]), name="View Tagged Rangeify") + + # bufferize -> store + tsink = graph_rewrite(tsink, pm_add_buffers, bottom_up=True, name="bufferize to store") + tsink = graph_rewrite(tsink, split_kernels, ctx=uop_list, name="split kernels") # if a kernel depends on a buffer, and that buffer is later assigned to, make the assign depend on the kernel's assign kernel_assign: dict[UOp, UOp] = {} assign_rep: dict[UOp, UOp] = {} - for u in tensor_map[sink].toposort(): + for u in tsink.toposort(): if u.op is not Ops.ASSIGN: continue kernel_assign[u.buf_uop] = u for s in u.src[1].src: @@ -454,8 +605,14 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]: if any(x.op is Ops.ASSIGN and x.buf_uop is s for x in u.toposort()): raise RuntimeError(f"cycle detected in graph, kernel for {u.buf_uop} must either depend on ASSIGN or BUFFER") assign_rep[a] = kernel_assign[s] = a.replace(src=a.src+(u,)) - if assign_rep: - tensor_map = graph_rewrite_map(tensor_map[sink], _substitute, ctx=assign_rep, bottom_up=True, input_map=tensor_map, name="fix_assign") + if assign_rep: tsink = graph_rewrite(tsink, _substitute, ctx=assign_rep, bottom_up=True, name="fix_assign") - if getenv("VIZ"): graph_rewrite(tensor_map[sink], PatternMatcher([]), name="View Kernel Graph") - return tensor_map + if getenv("VIZ"): graph_rewrite(tsink, PatternMatcher([]), name="View Kernel Graph") + + becomes_map: dict[UOp, UOp] = {} + for s in tsink.src: + assert s.tag is not None + for a in s.tag: + if a is None: continue + becomes_map[uop_list[cast(int, a)]] = s.replace(tag=None) + return becomes_map diff --git a/tinygrad/shape/shapetracker.py b/tinygrad/shape/shapetracker.py index 9372dab7b5..dca69bbe96 100644 --- a/tinygrad/shape/shapetracker.py +++ b/tinygrad/shape/shapetracker.py @@ -5,30 +5,24 @@ import functools from typing import Callable from tinygrad.helpers import merge_dicts, getenv from tinygrad.shape.view import View, unravel -from tinygrad.uop.symbolic import symbolic_flat, uop_given_valid, simplify_valid +from tinygrad.uop.symbolic import sym from tinygrad.uop.ops import UOp, Ops, graph_rewrite, Variable, sint, sint_to_uop, Context @functools.cache -def views_to_indexed_uops(views: tuple[View, ...], _idxs:tuple[UOp, ...]|None=None) -> tuple[UOp, UOp]: - idx, valid = views[-1].to_indexed_uops(_idxs) +def views_to_valid_uop(views: tuple[View, ...], _idxs:tuple[UOp, ...]|None=None) -> UOp: + idx = views[-1].to_valid_uop(_idxs) for view in reversed(views[0:-1]): view = view.minify() - idx, valid = view.to_indexed_uops([sint_to_uop(i) for i in unravel(view.shape, idx)], valid) + idx = view.to_valid_uop([sint_to_uop(i) for i in unravel(view.shape, idx)]) with Context(TRACK_MATCH_STATS=0): - # symbolic - idx, valid = graph_rewrite(UOp.sink(idx, valid), symbolic_flat, name="indexing sym @ 1").src - # simplify - if (newvalid:=simplify_valid(valid)) is not None: valid = newvalid - if (newidx:=uop_given_valid(valid, idx)) is not None: idx = newidx - # symbolic again - return graph_rewrite(UOp.sink(idx, valid), symbolic_flat, name="indexing sym @ 2").src + return graph_rewrite(idx, sym, name="indexing sym @ 1") @functools.cache def views_to_real_strides(views: tuple[View, ...], ignore_valid=False) -> tuple[sint|None, ...]: # NOTE: if a stride is not always valid, it will be None if len(views) == 1 and views[-1].mask is None: return views[-1].strides ret: list[sint|None] = [None] * len(views[-1].shape) - idx, valid = views_to_indexed_uops(views) + idx, valid = (vidx:=views_to_valid_uop(views)).get_idx(), vidx.get_valid() for c in idx.split_uop(Ops.ADD): if c.op is Ops.RANGE: ret[c.arg[0]] = 1 if c.op is Ops.MUL and c.src[0].op is Ops.RANGE and c.src[1].op is Ops.CONST: ret[c.src[0].arg[0]] = c.src[1].arg @@ -69,14 +63,14 @@ class ShapeTracker: def reduce(self, axis:tuple[int, ...]) -> tuple[sint, ...]: return tuple(1 if i in axis else s for i,s in enumerate(self.shape)) - def to_indexed_uops(self, _idxs:list[UOp]|tuple[UOp, ...]|None=None) -> tuple[UOp, UOp]: - return views_to_indexed_uops(self.views, tuple(_idxs) if _idxs is not None else None) + def to_valid_uop(self, _idxs:list[UOp]|tuple[UOp, ...]|None=None) -> UOp: + return views_to_valid_uop(self.views, tuple(_idxs) if _idxs is not None else None) # upper bound on buffer size required to fit this shapetracker def real_size(self) -> int: if 0 in self.shape: return 0 view = (v.shrink(v.mask) if (v:=self.views[0]).mask else v) - idx, _ = views_to_indexed_uops((view,)) + idx = views_to_valid_uop((view,)).get_idx() assert idx.vmax < 1e12, f"real_size broken for {self}" return int(idx.vmax + 1) @@ -95,11 +89,6 @@ class ShapeTracker: with Context(TRACK_MATCH_STATS=0): return views_to_real_strides(self.views, ignore_valid) def unit_stride_axes(self, ignore_valid=False) -> list[int]: return [i for i,st in enumerate(self.real_strides(ignore_valid)) if st == 1] - def axis_is_masked(self, axis:int) -> bool: - with Context(TRACK_MATCH_STATS=0): - _, valid = self.to_indexed_uops() - return axis in [x.arg[0] for x in graph_rewrite(valid, symbolic_flat).toposort() if x.op is Ops.RANGE] - def simplify(self) -> ShapeTracker: if len(self.views) >= 2 and (new_view := self.views[-2] + self.views[-1]) is not None: return ShapeTracker(self.views[:-2] + (new_view,)).simplify() diff --git a/tinygrad/shape/view.py b/tinygrad/shape/view.py index 9ab4687a82..37da15642c 100644 --- a/tinygrad/shape/view.py +++ b/tinygrad/shape/view.py @@ -112,16 +112,17 @@ class View: mask:tuple[tuple[sint, sint], ...]|None contiguous:bool - def to_indexed_uops(self:View, idxs:Sequence[UOp]|None=None, vexpr:UOp=UOp.const(dtypes.bool, True)) -> tuple[UOp, UOp]: - """(idx, valid)""" + def to_valid_uop(self, idxs:Sequence[UOp]|None=None) -> UOp: + """valid.where(idx, INVALID)""" if idxs is None: idxs = [UOp.range(s, i) for i,s in enumerate(self.shape)] iexpr = sint_to_uop(self.offset) + where = UOp.const(dtypes.bool, True) for idx,sh,st,m in zip(idxs, self.shape, self.strides, self.mask if self.mask is not None else itertools.repeat(None)): - if resolve(sh != 1) and resolve(st != 0): iexpr = iexpr + idx*st + iexpr = iexpr + idx*sint_to_uop(st) if m is not None: - if resolve(m[0] != 0): vexpr = vexpr * (idx >= m[0]) - if resolve(m[1] != sh): vexpr = vexpr * (idx < m[1]) - return iexpr, vexpr + if resolve(m[0] != 0): where &= (idx >= sint_to_uop(m[0])) + if resolve(m[1] != sh): where &= (idx < sint_to_uop(m[1])) + return where.where(iexpr, UOp.invalid()) @functools.cache # pylint: disable=method-cache-max-size-none def size(self) -> int: @@ -311,10 +312,7 @@ class View: if not all(x >= 0 for x in new_shape): raise ValueError(f"shape can't contain negative numbers {new_shape}") # check for the same size - if all_int(self.shape): - # reshapes cannot introduce symbolic shape - assert all_int(new_shape), f"{self.shape=} -> {new_shape=} contains non int dims" - if prod(self.shape) != prod(new_shape): raise ValueError(f"size mismatched, can't reshape {self.shape=} -> {new_shape=}") + if resolve(prod(self.shape) != prod(new_shape), True): raise ValueError(f"size mismatched, can't reshape {self.shape=} -> {new_shape=}") if 0 in self.shape: return View.create(new_shape) if new_shape == () and self.mask and any(mx==my for (mx,my) in self.mask): return None diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index a6116e0291..a1b0c6ac76 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -8,7 +8,8 @@ from tinygrad.dtype import _from_np_dtype, _to_np_dtype from tinygrad.helpers import argfix, make_tuple, flatten, prod, all_int, round_up, merge_dicts, argsort, getenv, all_same, fully_flatten, dedup from tinygrad.helpers import IMAGE, WINO, Metadata, TRACEMETA, ceildiv, fetch, polyN, unwrap, DEBUG, is_numpy_ndarray, RANGEIFY, FUSE_ATTENTION from tinygrad.gradient import compute_gradient -from tinygrad.uop.ops import smax, smin, resolve, UOp, Ops, sint, MathTrait, identity_element, all_metadata, index_to_concrete_int +from tinygrad.uop.ops import smax, smin, resolve, UOp, Ops, sint, MathTrait, identity_element, all_metadata, index_to_concrete_int, sint_to_uop, \ + srender from tinygrad.uop.spec import tensor_uop_spec, type_verify from tinygrad.device import Device, Buffer from tinygrad.engine.realize import run_schedule @@ -98,7 +99,8 @@ def _broadcast_shape(*shapes:tuple[sint, ...]) -> tuple[sint, ...]: def _masked_setitem(target:Tensor, values:Tensor, mask:Tensor, axes:tuple[int, ...]) -> Tensor: # reduce such that if mask contains repeated indices the last one remains - for dim in axes: mask, values = functools.reduce(lambda x,y: (x[0]|y[0], y[0].where(y[1], x[1])), zip(mask.split(1, dim), values.split(1, dim))) + for dim in reversed(axes): + mask, values = functools.reduce(lambda x,y: (x[0]|y[0], y[0].where(y[1], x[1])), zip(mask.split(1, dim), values.split(1, dim))) # remove extra dims from reduce for dim in reversed(axes): mask, values = mask.squeeze(dim), values.squeeze(dim) # select from values for each True element in mask else select from target @@ -632,6 +634,7 @@ class Tensor(MathTrait): """ if stop is None: stop, start = start, 0 dtype = kwargs.pop("dtype", dtypes.default_float if any(isinstance(x, float) for x in (start, stop, step)) else dtypes.default_int) + if start < (dt:=to_dtype(dtype)).min or dt.max < (stop-step): raise ValueError(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 Tensor([], dtype=dtype, **kwargs) return (Tensor.full((output_len,), step, dtype=dtype, **kwargs)._cumalu(0, Ops.ADD) + (start - step)).cast(dtype) @@ -992,6 +995,8 @@ class Tensor(MathTrait): # resolve -1 if (c := new_shape.count(-1)) > 1: raise RuntimeError(f"only one dimension can be inferred using -1, getting {new_shape}") if c: new_shape = tuple([-prod(self.shape) // prod(new_shape) if s == -1 else s for s in new_shape]) + if resolve(prod(self.shape) != prod(new_shape), True): + raise ValueError(f"size mismatch, can't reshape ({', '.join(srender(d) for d in self.shape)}) -> ({', '.join(srender(d) for d in new_shape)})") return self._apply_uop(UOp.reshape, arg=new_shape) if new_shape != self.shape else self def expand(self, shape, *args) -> Tensor: @@ -1064,6 +1069,7 @@ class Tensor(MathTrait): print(t.shrink((((0, 2), (0, 2)))).numpy()) ``` """ + if self.ndim != len(arg): raise ValueError(f"{self.ndim=} != {len(arg)=}") if (shrink_arg:=[x if x is not None else (0,s) for x,s in zip(arg, self.shape)]) == [(0,s) for s in self.shape]: return self return self._apply_uop(UOp.shrink, arg=tuple(shrink_arg)) @@ -1130,6 +1136,10 @@ class Tensor(MathTrait): X = Tensor.cat(*(X_ for X_ in (xB, X, xA) if X_ is not None), dim=d) return X.shrink(tuple((-min(pB,0), min(pA+s,s)) for (pB,pA),s in zip(pX, X.shape))) + # convenience + def pad_to(self, shape, *args): return self.pad(tuple([(0, ns-s) for s,ns in itertools.zip_longest(self.shape, argfix(shape, *args))])) + def shrink_to(self, shape, *args): return self.shrink(tuple([(0, ns) for ns in argfix(shape, *args)])) + # ***** movement high level ops ***** def _getitem(self, indices, v: Tensor|None = None) -> Tensor: @@ -1167,6 +1177,9 @@ class Tensor(MathTrait): boundary, stride = [start, stop], step if all(isinstance(s, int) for s in (start,stop,step)): # handle int slicing + # if we're slicing a symbolic dimension into a int dimension, we can slice untill the bind size + # TODO: right now this is using vmax instead of the bind size because jit doesnt update the bound value of the returned tensor + if isinstance(size, UOp): size = int(size.vmax) *boundary, stride = index.indices(cast(SupportsIndex, size)) if stride * (boundary[1] - boundary[0]) < 0: boundary = [0, 0] elif stride < 0: boundary = [boundary[1] + 1, boundary[0] + 1] @@ -1219,8 +1232,8 @@ class Tensor(MathTrait): x = (mask.where(x.reshape(reshape_arg), 0)).sum(sum_axis:=tuple(d + len(big_shape) for d in dims), dtype=x.dtype) # special permute case - if dims[0] != 0 and len(dims) != 1 and tuple(dims) != tuple(range(dims[0], dims[-1]+1)): - x = x.permute(*range(dims[0], dims[0]+len(big_shape)), *range(0, dims[0]), *range(dims[0]+len(big_shape), x.ndim)) + if (permuted := dims[0] != 0 and len(dims) != 1 and tuple(dims) != tuple(range(dims[0], dims[-1]+1))): + mask, x = (y.permute(*range(dims[0], dims[0]+len(big_shape)), *range(0, dims[0]), *range(dims[0]+len(big_shape), y.ndim)) for y in (mask, x)) # for advanced setitem, returns whole tensor with indices replaced if v is not None: @@ -1228,7 +1241,7 @@ class Tensor(MathTrait): # add back reduced dims from sum for dim in sum_axis: vb = vb.unsqueeze(dim) # run _masked_setitem on tuple of axis that is to be reduced to match self.shape - x = _masked_setitem(self, vb, mask, tuple(range(dims[0], dims[0] + len(big_shape)))) + x = _masked_setitem(self, vb, mask, tuple(range((start := dims[0] if not permuted else 0), start + len(big_shape)))) return x @@ -3897,7 +3910,8 @@ class Tensor(MathTrait): def _one_hot_along_dim(self:Tensor, num_classes:sint, dim:int=-1) -> Tensor: if not dtypes.is_int(self.dtype): raise RuntimeError(f"_one_hot_along_dim expects int index tensor, getting {self.dtype}") offset = self.ndim - self._resolve_dim(dim) - 1 - return self == Tensor.arange(num_classes, device=self.device, requires_grad=False).reshape((num_classes,) + (1,) * offset) + dt = dtypes.int64 if sint_to_uop(num_classes).overflows(dtypes.int32) else dtypes.int32 + return self == Tensor.arange(num_classes, dtype=dt, device=self.device, requires_grad=False).reshape((num_classes,) + (1,) * offset) def one_hot(self, num_classes:int=-1) -> Tensor: """ diff --git a/tinygrad/uop/__init__.py b/tinygrad/uop/__init__.py index 2a453a7f64..1cab564136 100644 --- a/tinygrad/uop/__init__.py +++ b/tinygrad/uop/__init__.py @@ -22,6 +22,7 @@ class Ops(FastEnum): # ops that adjust the behavior of the scheduler CONTIGUOUS = auto(); CONTIGUOUS_BACKWARD = auto(); DETACH = auto(); FUSE = auto() # noqa: E702 + REALIZE = auto() # blocks in linearizer (only used there) BLOCK = auto(); BLOCKSTART = auto(); BLOCKEND = auto(); BLOCKFINAL = auto() # noqa: E702 diff --git a/tinygrad/uop/decompositions.py b/tinygrad/uop/decompositions.py index 57ce303422..cc3e5cf09f 100644 --- a/tinygrad/uop/decompositions.py +++ b/tinygrad/uop/decompositions.py @@ -293,7 +293,7 @@ def fast_idiv(device: str, x: UOp, d: int, dont_cast=False) -> UOp|None: if (ret:=fast_idiv(device, x//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][-1]) and is_dtype_supported(next_dtype, None if device=='' else device): + if dtypes.is_int(next_dtype := promo_lattice[x.dtype.scalar()][-1]) and is_dtype_supported(next_dtype, None if device=='' else device): if m*vmin >= dtypes.min(next_dtype) and m*vmax <= dtypes.max(next_dtype): 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 @@ -343,7 +343,7 @@ def get_late_rewrite_patterns(ops:tuple[Ops, ...], force_transcendental=False): pat += [(UPat.var("x", dtypes.ints)//UPat.cvar("c"), lambda x,c: (x+(l.const_like(l.vmin) if (l:=(x<0)).vmin==l.vmax else l).where( c-1, 0)) >> v if (v:=powers_of_two.get(c.arg, 0)) else None)] # (x+(x<0).where(c-1, 0)) >> v if not DISABLE_FAST_IDIV: - pat += [(UPat.var("x", dtypes.ints)//UPat.cvar("d"), lambda ctx, x, d: fast_idiv(ctx, x, d.arg))] + pat += [(UPat.var("x", dtypes.ints)//UPat.cvar("d", vec=False), lambda ctx, x, d: fast_idiv(ctx, x, d.arg))] pat += [(UPat.var("x", dtypes.ints)%UPat.var("d"), lambda x, d: x-d*(x//d))] if Ops.NEG in ops: pat += [(UPat.var('x')*-1, lambda x: x.alu(Ops.NEG))] diff --git a/tinygrad/uop/mathtraits.py b/tinygrad/uop/mathtraits.py index 05bf89d8be..0de976c90b 100644 --- a/tinygrad/uop/mathtraits.py +++ b/tinygrad/uop/mathtraits.py @@ -167,3 +167,4 @@ class MathTrait: def log2(self): return self.alu(Ops.LOG2) def exp2(self): return self.alu(Ops.EXP2) def pow(self, x): return self.alu(Ops.POW, self.ufix(x)) + def __pow__(self, x): return self.pow(x) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index bc90eb2906..8bd4f80d48 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -1,21 +1,24 @@ from __future__ import annotations from typing import Any, Callable, cast, TYPE_CHECKING, Type, Sequence -import sys, time, functools, itertools, math, operator, hashlib, os, types, pickle, pathlib, inspect, weakref +import sys, time, functools, itertools, math, operator, hashlib, os, types, pickle, pathlib, inspect, weakref, collections from dataclasses import dataclass, field from enum import Enum, auto from tinygrad.uop import Ops, GroupOp from tinygrad.uop.mathtraits import MathTrait -from tinygrad.dtype import ConstType, ImageDType, dtypes, DType, truncate, PtrDType, least_upper_dtype +from tinygrad.dtype import ConstType, ImageDType, dtypes, DType, truncate, PtrDType, least_upper_dtype, Invalid, InvalidType from tinygrad.helpers import ContextVar, all_int, prod, getenv, all_same, Context, partition, temp, unwrap, T, argfix, Metadata, flatten, TRACEMETA -from tinygrad.helpers import PICKLE_BUFFERS, PROFILE, dedup, cdiv, cmod, diskcache_put, to_function_name, cpu_profile, TracingKey +from tinygrad.helpers import PICKLE_BUFFERS, PROFILE, dedup, cdiv, cmod, diskcache_put, to_function_name, cpu_profile, TracingKey, RANGEIFY, VIZ if TYPE_CHECKING: from tinygrad.shape.shapetracker import ShapeTracker from tinygrad.device import Buffer, MultiBuffer 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() +range_start = {Ops.BUFFERIZE: 1, Ops.REDUCE: 1, Ops.STORE: 2, Ops.WMMA: 3} + # https://en.wikipedia.org/wiki/Identity_element def identity_element(op:Ops, dt:DType) -> ConstType: return dtypes.as_const({Ops.ADD:0, Ops.MUL:1, Ops.MAX:dtypes.min(dt)}[op], dt) @@ -155,14 +158,14 @@ class UOp(MathTrait, metaclass=UOpMetaClass): from tinygrad.shape.shapetracker import ShapeTracker # VIEW and MovementOps define a new ShapeTracker from the arg if self.op is Ops.VIEW: return self.arg - if self.op is Ops.BUFFERIZE: return ShapeTracker.from_shape((prod(tuple([int(r.vmax+1) for r in self.src[1:]])),)) - #if self.op is Ops.BUFFERIZE: return ShapeTracker.from_shape(tuple([r.vmax+1 for r in self.src[1:]])) + if self.op is Ops.BUFFERIZE: return ShapeTracker.from_shape(tuple([int(r.vmax+1) for r in self.src[1:]])) # allow reshape from nothing if self.op is Ops.RESHAPE and self.src[0].st is None: return ShapeTracker.from_shape(self.arg) if self.op in GroupOp.Movement: return unwrap(self.src[0].st).mop(self.op, self.arg) # CONST with a DEVICE has a shape of () if self.op is Ops.CONST and len(self.src) and self.src[0].op is Ops.DEVICE: return ShapeTracker.from_shape(()) if self.op is Ops.STORE and isinstance(self.dtype, PtrDType): return ShapeTracker.from_shape((self.dtype.size,)) + if self.op is Ops.STORE and self.dtype is not dtypes.void: return self.src[0].src[0].st # BufferOps and ASSIGN flow ShapeTracker from a direct edge if self.op in {Ops.STORE, Ops.ASSIGN, Ops.LOAD}: return self.src[0].st if self.op in GroupOp.Buffer: return views[0] if (views:=[x.st for x in self.src if x.op is Ops.VIEW]) else None @@ -211,7 +214,6 @@ class UOp(MathTrait, metaclass=UOpMetaClass): @functools.cached_property def ranges(self) -> dict[UOp, None]: if self.op is Ops.RANGE: return {self:None} - range_start = {Ops.BUFFERIZE: 1, Ops.REDUCE: 1, Ops.STORE: 2, Ops.WMMA: 3} ret: dict[UOp, None] = {} if self.op in range_start.keys(): for s in self.src[:range_start[self.op]]: ret.update(s.ranges) @@ -297,12 +299,17 @@ class UOp(MathTrait, metaclass=UOpMetaClass): if isinstance(b, UOp): return b.unbind()[0] if b.op is Ops.BIND else b if isinstance(b, tuple) and all_same(b): 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=dtypes.as_const(b, dtype), src=() if src is None else (src,)) - if shape is not None: - from tinygrad.shape.shapetracker import ShapeTracker - ret = ret.replace(src=(UOp(Ops.VIEW, dtypes.void, (), ShapeTracker.from_shape(shape, (0,)*len(shape))),)) - if device is not None: - if shape is not None: ret = ret.replace(src=(UOp(Ops.DEVICE, arg=device).view(unwrap(ret.st)),)) - else: ret = ret.replace(src=(UOp(Ops.DEVICE, arg=device),)) + if RANGEIFY: + # VIEW on const is no longer supported in RANGEIFY + if device is not None: ret = ret.replace(src=(UOp(Ops.DEVICE, arg=device),)) + if shape is not None: ret = ret.reshape((1,)*len(shape)).expand(shape) + else: + if shape is not None: + from tinygrad.shape.shapetracker import ShapeTracker + ret = ret.replace(src=(UOp(Ops.VIEW, dtypes.void, (), ShapeTracker.from_shape(shape, (0,)*len(shape))),)) + if device is not None: + if shape is not None: ret = ret.replace(src=(UOp(Ops.DEVICE, arg=device).view(unwrap(ret.st)),)) + else: ret = ret.replace(src=(UOp(Ops.DEVICE, arg=device),)) return ret @staticmethod def range(end:sint, *arg): @@ -320,8 +327,18 @@ class UOp(MathTrait, metaclass=UOpMetaClass): assert len(axis) == len(new_axis) ret = UOp(Ops.REDUCE_AXIS, self.dtype, (ret,), (op, new_axis)) return ret.reshape(tuple([x if i not in axis else 1 for i,x in enumerate(self.shape)])) + @staticmethod + def invalid(count=1): return UOp(Ops.CONST, dtypes.index.vec(count), src=(), arg=Invalid) + def valid(self, cond): return cond.where(self, UOp.invalid(self.dtype.count)) + def get_idx(self) -> UOp: + assert self.dtype.scalar() is dtypes.index, "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 + def get_valid(self) -> UOp: + assert self.dtype.scalar() is dtypes.index, "Can only call get_valid on index dtype" + return self.src[0] if self.op is Ops.WHERE and self.src[2].arg is Invalid else UOp.const(dtypes.bool, self.arg is not Invalid) def reduce(self, *src:UOp, **kwargs): return UOp(Ops.REDUCE, kwargs.pop('dtype', self.dtype), src=(self,)+src, **kwargs) def contiguous(self, *args, **kwargs): return UOp(Ops.CONTIGUOUS, dtype=self.dtype, src=(self,)+args, **kwargs) + def realize(self, *args, **kwargs): return UOp(Ops.REALIZE, dtype=self.dtype, src=(self,)+args, **kwargs) def contiguous_backward(self): return self.alu(Ops.CONTIGUOUS_BACKWARD) def bufferize(self, *args, **kwargs): return UOp(Ops.BUFFERIZE, dtype=self.dtype, src=(self,)+args, **kwargs) def fuse(self): return self.alu(Ops.FUSE) @@ -428,6 +445,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass): @functools.cached_property def _device(self) -> str|tuple[str, ...]|None: if self.op is Ops.DEVICE: return self.arg + if self.op is Ops.BUFFERIZE: return self.arg.device if self.op is Ops.MSELECT: assert isinstance(self.src[0].device, tuple), "mselect must be on tuple device" return self.src[0].device[self.arg] @@ -534,7 +552,23 @@ class UOp(MathTrait, metaclass=UOpMetaClass): if (d0:=self.src[0].divides(v)) is not None: return d0 * self.src[1] if (d1:=self.src[1].divides(v)) is not None: return self.src[0] * d1 return None # generic None if we aren't sure - def pop_const(self) -> tuple[UOp, int]: return (self.src[0], self.src[1].arg) if self.op is Ops.ADD and self.src[1].op is Ops.CONST else (self, 0) + def pop_const(self, op=Ops.ADD) -> tuple[UOp, ConstType]: + return (self.src[0], self.src[1].arg) if self.op is op and self.src[1].op is Ops.CONST else (self, identity_element(op, self.dtype)) + @staticmethod + def gcd(*uops: UOp) -> UOp: + terms, factors = zip(*[(u.divides(f:=u.const_factor()),f) for u in uops]) + count = functools.reduce(operator.and_, [collections.Counter(term.split_uop(Ops.MUL)) for term in terms]) + return math.prod([*count.elements(), terms[0].const_like(math.gcd(*factors))]) # put the const at the top + def divide_exact(self, v:UOp) -> UOp|None: + if self is v: return self.const_like(1) + if self.op is Ops.ADD: return None if (s0:=self.src[0].divide_exact(v)) is None or (s1:=self.src[1].divide_exact(v)) is None else s0+s1 + if v.op is Ops.CONST: return self.divides(v.arg) + if self.op is Ops.MUL: + (fac, const), (div_fac, div_const) = self.pop_const(Ops.MUL), v.pop_const(Ops.MUL) + new_count = collections.Counter(fac.split_uop(Ops.MUL)) + new_count.subtract(div_fac.split_uop(Ops.MUL)) + if const%div_const==0 and all(v>=0 for v in new_count.values()): return math.prod([*new_count.elements(), self.const_like(const//div_const)]) + return None # generic None if we aren't sure @property def vmin(self) -> ConstType: return self._min_max[0] @property @@ -570,8 +604,8 @@ class UOp(MathTrait, metaclass=UOpMetaClass): 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.VECTORIZE}: return min(x.vmin for x in self.src), max(x.vmax for x in self.src) - if self.op is Ops.CONST: return self.arg, self.arg - if self.op is Ops.VCONST: return (min(self.arg), max(self.arg)) + 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.index,): @@ -626,6 +660,7 @@ python_alu: dict[Ops, Callable] = { def exec_alu(op:Ops, dtype:DType, operands, truncate_output=True): if dtype.count > 1: return tuple([exec_alu(op, dtype.scalar(), [x[i] if isinstance(x, tuple) else x for x in operands]) for i in range(dtype.count)]) + if dtype==dtypes.index and op in GroupOp.Binary and Invalid in operands: return Invalid alu = python_alu[op](*operands) return truncate.get(dtype, lambda x: x)(alu) if truncate_output else alu @@ -699,7 +734,7 @@ class UPat(MathTrait): def cvar(name:str|None=None, dtype:DType|tuple[DType, ...]|None=None, vec=True): return UPat((Ops.CONST,Ops.VCONST) if vec else Ops.CONST, dtype, name=name) @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|InvalidType): return UPat(Ops.CONST, dtype=dtype, arg=b) # lil helper def f(self, op, **kwargs): return UPat(op, src=(self,), **kwargs) @@ -819,7 +854,6 @@ def track_uop(u:UOp): # *** tracking pattern matcher *** -VIZ = ContextVar("VIZ", 0) TRACK_MATCH_STATS = ContextVar("TRACK_MATCH_STATS", 2 if VIZ else 0) match_stats:dict[UPat, list[int|float]] = dict() @@ -922,7 +956,7 @@ if TRACK_MATCH_STATS or PROFILE: with open(fn:=temp("rewrites.pkl", append_user=True), "wb") as f: print(f"rewrote {len(tracked_ctxs)} graphs and matched {sum(len(r.matches) for x in tracked_ctxs for r in x)} times, saved to {fn}") pickle.dump([(tracked_keys, tracked_ctxs, uop_fields)], f) - if VIZ: launch_viz(VIZ, temp("rewrites.pkl", append_user=True)) + if VIZ: return launch_viz("VIZ", temp("rewrites.pkl", append_user=True)) if getenv("PRINT_MATCH_STATS", TRACK_MATCH_STATS.value): ret = [0,0,0.0,0.0] for k,v in sorted(list(match_stats.items()), key=lambda x: x[1][2]+x[1][3]): @@ -932,11 +966,10 @@ if TRACK_MATCH_STATS or PROFILE: print(f"{ret[0]:6d} / {ret[1]:7d} -- {ret[3]*1000.:9.2f} / {(ret[2]+ret[3])*1000.:9.2f} ms -- TOTAL") print(f"{len(match_stats)} rules, {sum(v[0] > 0 for v in match_stats.values())} matched once") - def launch_viz(var:ContextVar, data:str): - os.environ[(env_str:=var.key)] = "0" + def launch_viz(env_str:str, data:str): + os.environ[env_str] = "0" os.environ[f"{env_str}_DATA"] = data - os.environ[f"{env_str}_VALUE"] = str(var.value) - if not int(os.getenv("VIZ", "0")) and not int(os.getenv("PROFILE", "0")): + if not int(os.getenv("VIZ", "0")) and not int(os.getenv("PROFILE", "0")) and not int(os.getenv("SQTT", "0")): args = ['--kernels', getenv("VIZ_DATA", "")] if getenv("VIZ_DATA", "") else [] args += ['--profile', getenv("PROFILE_DATA", "")] if getenv("PROFILE_DATA", "") else [] os.execv(sys.executable, [sys.executable] + [os.path.join(os.path.dirname(__file__), "../", "viz", "serve.py")] + args) @@ -1036,7 +1069,8 @@ pm_lower_index_dtype = PatternMatcher([ # comparison ops might now have different dtypes in their sources (UPat(GroupOp.Comparison, name="u", src=(UPat.var("x",dtypes.ints), UPat.var("y", dtypes.ints))), lambda u,x,y: x.cast(dt:=least_upper_dtype(x.dtype, y.dtype)).alu(u.op, y.cast(dt)) if x.dtype!=y.dtype else None), - (UPat(Ops.WHERE, dtype=dtypes.index, src=(UPat.var("cond"), UPat.var("x"), UPat.var("y")), name="u"), lambda cond,u,x,y: + (UPat(Ops.WHERE, dtypes.index, src=(UPat(), UPat.var("x"), UPat(Ops.CONST, arg=Invalid)), name="u"), lambda u,x: u.replace(dtype=x.dtype)), + (UPat(Ops.WHERE, dtypes.index, src=(UPat.var("cond"), UPat.var("x"), UPat.var("y"))), lambda cond,x,y: cond.where(x.cast(dt:=least_upper_dtype(x.dtype, y.dtype)), y.cast(dt))), (UPat((Ops.CONST, Ops.VCONST), dtype=dtypes.index, name="u"), lambda u: u.replace(dtype=select_dtype(u))), (UPat((Ops.RANGE,), dtype=dtypes.index, src=(UPat.var("end")), name="r"), lambda ctx,r,end: @@ -1048,6 +1082,11 @@ pm_lower_index_dtype = PatternMatcher([ else dtypes.long)).vec(u.dtype.count),src=tuple(x.cast(dt) for x in u.src))), (UPat((Ops.SPECIAL,Ops.DEFINE_VAR), dtypes.index, name="u"), lambda u: u.replace(dtype=dtypes.int)), (UPat((Ops.BIND), dtypes.index, name="u"), lambda u: u.replace(dtype=u.src[0].dtype)), + # lower Invalid + (UPat.var("buf").index(UPat.var("cond").where(UPat.var("idx"), UPat(Ops.CONST, arg=Invalid))), lambda buf,idx,cond: buf.index(idx, cond)), + # remove hanging cast + (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx", dtypes.int).cast()),), lambda buf,idx: buf.index(idx)), + (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx", dtypes.int).cast(), UPat.var("valid"))), lambda buf,idx,valid: buf.index(idx, valid)), ]) def index_to_concrete_int(u:UOp): return graph_rewrite(u, pm_lower_index_dtype) @@ -1120,4 +1159,4 @@ def pyrender(ast:UOp) -> list[str]: sint = int|UOp Variable = UOp -ConstLike = ConstType|Variable|tuple[ConstType, ...] +ConstLike = ConstType|InvalidType|Variable|tuple[ConstType|InvalidType, ...] diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index b3066ad5e1..48737fc9e2 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -1,7 +1,7 @@ from typing import cast, Callable -from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, print_uops, python_alu, graph_rewrite -from tinygrad.dtype import DType, ImageDType, dtypes, PtrDType, AddrSpace -from tinygrad.helpers import all_same, prod, DEBUG, ContextVar, Context +from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, print_uops, python_alu, graph_rewrite, AxisType +from tinygrad.dtype import DType, ImageDType, dtypes, PtrDType, AddrSpace, Invalid +from tinygrad.helpers import all_same, prod, DEBUG, ContextVar, Context, cpu_profile from tinygrad.shape.shapetracker import ShapeTracker try: import z3 @@ -10,8 +10,12 @@ try: # IDIV is truncated division but z3 does euclidian division (floor if b>0 ceil otherwise); mod by power of two sometimes uses Ops.AND def z3_cdiv(a, b):return z3.If((a<0), z3.If(0= 0, z3.ToInt(a), -z3.ToInt(-a)))} def create_bounded(name:str, vmin, vmax, solver:z3.Solver) -> z3.ArithRef: s = z3.Int(name, ctx=solver.ctx) @@ -38,8 +42,6 @@ try: UOp(Ops.NOOP, arg=(ctx[0], create_bounded(f"cast{ctx[1].setdefault(x, len(ctx[1]))}", x.dtype.min, x.dtype.max, ctx[0])))), (UPat(Ops.CAST, dtype=dtypes.bool, name="x"), lambda x,ctx: UOp(Ops.NOOP, arg=(ctx[0], z3.Bool(f"cast{ctx[1].setdefault(x, len(ctx[1]))}",ctx=ctx[0].ctx)))), - (UPat(Ops.XOR, dtype=dtypes.ints+(dtypes.bool, ), src=UPat(Ops.NOOP), name="x"), - lambda x,ctx: UOp(Ops.NOOP, arg=(ctx[0], z3.BV2Int(z3_alu[x.op](*(z3.Int2BV(s.arg[1], x.dtype.itemsize*8) for s in x.src)))))), (UPat(GroupOp.ALU, src=UPat(Ops.NOOP), name="x"), lambda x,ctx: UOp(Ops.NOOP, arg=(ctx[0], z3_alu[x.op](*(s.arg[1] for s in x.src))))), # A comparison between floats introduces a new bool variable (UPat(GroupOp.Comparison, src=UPat(dtype=dtypes.floats), name="x"), lambda x,ctx: @@ -99,8 +101,11 @@ tensor_uop_spec = buffer_spec+assign_spec+PatternMatcher([ # Tensor const has a device and an unmasked ShapeTracker of stride 0 # NOTE: variables in shape can cause multiple views in this ShapeTracker and other issues, see TestSymbolicJit.test_ones_sum - (UPat(Ops.CONST, src=(UPat(Ops.VIEW, name="st", src=(UPat(Ops.DEVICE),)),)), + # TODO: remove after rangeify is default + (UPat(Ops.CONST, src=(UPat.any(UPat(Ops.VIEW, src=(UPat(Ops.DEVICE),), name="st"), + UPat(Ops.VIEW, src=(UPat(Ops.DEVICE), UPat(Ops.BIND)), name="st")),)), lambda st: len(st.st.views) == 1 and all(v.mask is None for v in st.st.views)), + (UPat(Ops.CONST, src=(UPat(Ops.DEVICE),)), lambda: True), # DETACH and CONTIGUOUS change how we interpret the source UOp # CONTIGUOUS ensures the source UOp realizes @@ -133,11 +138,12 @@ def validate_index(idx:UOp, gate:UOp=UOp.const(dtypes.bool, True)): solver = z3.Solver(ctx=z3.Context()) z3_idx, z3_mask = uops_to_z3(solver, idx.src[1], mask) solver.add(z3_mask) - if solver.check((z3_idx<0)|(sz<=z3_idx)) == z3.sat: - print(f"idx={idx.src[1].render(simplify=False)}") - print(f"mask & gate={mask.render(simplify=False)}") - print(f"# OUT OF BOUNDS ACCESS: at {solver.model()} INDEX not in 0 - {sz}\nconstraints = {solver}") - return False + with cpu_profile("validate index with z3", "TINY"): + if solver.check((z3_idx<0)|(sz<=z3_idx)) == z3.sat: + print(f"idx={idx.src[1].render(simplify=False)}") + print(f"mask & gate={mask.render(simplify=False)}") + print(f"# OUT OF BOUNDS ACCESS: at {solver.model()} INDEX not in 0 - {sz}\nconstraints = {solver}") + return False return True def validate_store(idx:UOp, val:UOp, gate:UOp=UOp.const(dtypes.bool, True)): @@ -157,7 +163,8 @@ spec = PatternMatcher([ (UPat(Ops.DEFINE_REG, src=()), lambda: True), (UPat(Ops.DEFINE_VAR, name="x"), lambda x: isinstance(x.arg[1], int) and isinstance(x.arg[2], int)), - (UPat(Ops.RANGE, src=(UPat.var("x"),), name="rng"), lambda rng,x: rng.dtype == x.dtype and isinstance(rng.arg, tuple)), + (UPat(Ops.RANGE, src=(UPat.var("x"),), name="rng"), lambda rng,x: rng.dtype == x.dtype and isinstance(rng.arg, tuple) and len(rng.arg) == 2 and \ + isinstance(rng.arg[0], int) and isinstance(rng.arg[1], AxisType)), (UPat(Ops.SPECIAL, src=(UPat.var("x"),), name="s"), lambda s,x: s.dtype == x.dtype == dtypes.int32 and isinstance(s.arg, str)), (UPat(Ops.VIEW, dtypes.void, src=(), name="x"), lambda x: isinstance(x.arg, ShapeTracker)), @@ -165,7 +172,7 @@ spec = PatternMatcher([ lambda x,src: isinstance(x.arg, ShapeTracker) and src.op is not Ops.STORE and x.dtype.base == src.dtype.base), (UPat(Ops.VALID, dtypes.bool, (UPat(Ops.VIEW),)), lambda: True), - (UPat(Ops.CONST, name="x"), lambda x: type(x.arg) is type(dtypes.as_const(x.arg, x.dtype))), + (UPat(Ops.CONST, src=(), name="x"), lambda x: type(x.arg) is type(dtypes.as_const(x.arg, x.dtype))), # early LOAD has a (UPat(Ops.LOAD, src=(UPat(Ops.VIEW, src=(UPat(GroupOp.Defines),)),)), lambda: True), @@ -178,6 +185,8 @@ spec = PatternMatcher([ # make sure all index dtypes have been lowered (UPat(GroupOp.All, dtype=dtypes.index), lambda: False), + (UPat(Ops.CONST, arg=Invalid), lambda: False), + (UPat(Ops.VCONST, name="x"), lambda x: all(v is not Invalid for v in x.src)), # INDEX is used in new style load/store # INDEX takes a diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 93d0ce04b1..90ebf5afd6 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -3,8 +3,8 @@ from typing import cast import math, operator, struct, functools from collections import defaultdict from tinygrad.uop.ops import Ops, PatternMatcher, UPat, UOp, GroupOp, exec_alu -from tinygrad.dtype import ConstType, dtypes, PtrDType, AddrSpace, can_safe_cast -from tinygrad.helpers import partition, all_same, prod, flatten, get_single_element, cdiv, cmod, CORRECT_DIVMOD_FOLDING +from tinygrad.dtype import ConstType, dtypes, PtrDType, AddrSpace, can_safe_cast, Invalid +from tinygrad.helpers import partition, all_same, prod, flatten, get_single_element, cdiv, cmod, CORRECT_DIVMOD_FOLDING, unwrap from tinygrad.uop.decompositions import xpow # ******** phase 1 of symbolic used to live in ops, it's the most generic folding rules ******** @@ -22,7 +22,28 @@ def fold_bitcast(root:UOp, c:UOp) -> UOp|None: def convert(v: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))) -symbolic_simple = PatternMatcher([ +invalid_pat = UPat(Ops.CONST, arg=Invalid, name="i") +invalid_gate = UPat.var("cond").where(UPat.var("x"), invalid_pat) + +propagate_invalid = PatternMatcher([ + # this needs to be before symbolic so that 0*something_that_might_be_invalid doesnt become 0 + # propagate invalid, push it past children + *((invalid_gate.alu(op, UPat.var("y")).named("alu"), lambda cond,x,y,alu,i: cond.where(x.alu(alu.op,y), i)) + for op in GroupOp.Binary-GroupOp.Comparison), + *((invalid_gate.alu(op, UPat.var("y")).named("alu"), lambda cond,x,y,alu,i: x.alu(alu.op,y)) for op in GroupOp.Comparison), + # invalid + y -> y same for other ops + *((invalid_pat.alu(op, UPat(dtype=dtypes.index)).named("alu"), lambda alu,i: i) for op in GroupOp.Binary-GroupOp.Comparison), + # i < y -> a_bool_value_that_will_never_be_used: we choose a random bool const + *((invalid_pat.alu(op, UPat(dtype=dtypes.index)), lambda i: UOp.const(dtypes.bool, True)) for op in GroupOp.Comparison), + # 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)), + # order of gate&!cond matters!, and-clauses are only simplified left to right and we need to gate to be used to fold cond + (UPat.var("gate").where(invalid_gate, UPat.var("y")), lambda gate,cond,x,y,i: ((gate&cond.logical_not()).logical_not()).where(gate.where(x,y), i)), + # unswap the branches for the rule above + (UPat.var("gate").where(UPat.var("y"), invalid_gate).named("where"), lambda gate,cond,x,y,i: gate.logical_not().where(cond.where(x,i), y)) +]) + +symbolic_simple = propagate_invalid + PatternMatcher([ # ** self folding ** (UPat.var("x") + 0, lambda x: x), # x+0 -> x (UPat.var("x") * 1, lambda x: x), # x*1 -> x @@ -76,9 +97,6 @@ symbolic_simple = PatternMatcher([ (UPat(Ops.BITCAST, name="root", src=(UPat.cvar("c"),)), fold_bitcast), # b.cast(a).cast(b) -> b if a preserves all values in b (UPat.var('x').cast(name="a").cast(name="b"), lambda x,a,b: x if x.dtype == b.dtype and can_safe_cast(b.dtype, a.dtype) else None), - # if the intermediate cast doesnt narrow we can do it in one cast, we have to be carefull with bfloat16 - (UPat.var('x').cast(name="a").cast(name="b"), lambda x,a,b: x.cast(b.dtype) if can_safe_cast(x.dtype, a.dtype) and - not (a.dtype==dtypes.float and (b.dtype==dtypes.bfloat16 or x.dtype==dtypes.bfloat16)) else None), # ** pow ** (UPat.var("x").alu(Ops.POW, UPat.cvar("c", vec=False)), simplify_pow), # positive const ** x @@ -95,38 +113,15 @@ symbolic_simple = PatternMatcher([ # new decomp rules for threefry (((UPat.var(None, dtypes.uint64)<<32) | UPat.var('y', dtypes.uint32).cast(dtypes.uint64)).cast(dtypes.uint32), lambda y: y), (((UPat.var('x', dtypes.uint64)<<32) | UPat.var(None, dtypes.uint32).cast(dtypes.uint64))>>32, lambda x: x), - (UPat.var('b').where(UPat.var('x', dtypes.uint32).cast(dtypes.uint64), UPat.const(dtypes.uint64, 0)).cast(dtypes.uint32), lambda b,x: b.where(x,0)) + (UPat.var('b').where(UPat.var('x', dtypes.uint32).cast(dtypes.uint64), UPat.const(dtypes.uint64, 0)).cast(dtypes.uint32), lambda b,x: b.where(x,0)), + # ** 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), ]) # ******** phase 2 builds on phase 1, it includes the old "symbolic", rules that match deeper ******** -def fold_unrolled_divs(divs:UOp, denominator: int, fac=1) -> UOp|None: - # div pattern in unrolled arange - # example: (x//4+(x+1)//4+(x+2)//4+(x+3)//4 -> x - seen_const, ans = [], None - for u in divs.split_uop(Ops.ADD): - if fac!=1: - if u.op is not Ops.MUL or u.src[1].op is not Ops.CONST or u.src[1].arg != fac: return None - u = u.src[0] - if u.op is Ops.CAST and u.src[0].dtype == dtypes.index: u = u.src[0] - if not (u.op is Ops.IDIV and u.src[1].op is Ops.CONST): return None - if denominator != u.src[1].arg: return None - if (s0:=u.src[0]).vmin < 0: return None - # assumed CONST is the last of an ADD - if s0.op is Ops.ADD and s0.src[1].op is Ops.CONST and s0.src[1].op is Ops.CONST: - seen_const.append(s0.src[1].arg) - s0 = s0.src[0] - else: seen_const.append(0) - if ans is None: ans = s0 - if ans is not s0: return None - if ans is None: return None - # the first (denominator-len(seen_const)) terms may have been folded to 0 already - for i in range(denominator-len(seen_const)): - if ans is not None and 0 <= ans.vmin and ans.vmax + i < denominator: seen_const.append(i) - if sorted(seen_const)==list(range(denominator)): - return (fac*ans).cast(divs.dtype) - return None - def lt_folding(x:UOp, c:int) -> UOp|None: p, np = partition(x.split_uop(Ops.ADD), lambda u: u.const_factor() == 1) if np and (d:=math.gcd(*[u.const_factor() for u in np], c)) > 1 and 0 <= sum(u.vmin for u in p) and sum(u.vmax for u in p) < d: @@ -173,7 +168,7 @@ def remove_nested_mod(m: UOp, x: UOp, y: UOp) -> UOp|None: def fold_binary_numerator(d: UOp, x: UOp, y: UOp) -> UOp|None: # we can fold if the expression has only one non-constant term and this term can only take on two values - if ((c := y.arg) < 0) or (x.dtype.count > 1): return None + if ((c := y.arg) < 0): return None x,const = x.pop_const() terms, factors = zip(*[(u.divides(f:=u.const_factor()),f) for u in x.split_uop(Ops.ADD)]) if len(terms)==1 and (v:=terms[0]).vmax-v.vmin == 1: @@ -184,7 +179,7 @@ def fold_binary_numerator(d: UOp, x: UOp, y: UOp) -> UOp|None: def fold_divmod_congruence(d: UOp, x: UOp, y: UOp) -> UOp|None: # within a mod we can freely subtract multiples of c, we use this to see if a is congruent to an expression whose vmin/vmax are between 0 and c - if (x.vmin<0 and CORRECT_DIVMOD_FOLDING) or ((c := y.arg) < 0) or (x.dtype.count > 1): return None + if (x.vmin<0 and CORRECT_DIVMOD_FOLDING) or ((c := y.arg) < 0): return None x,const = x.pop_const() terms, factors = zip(*[(u.divides(f:=u.const_factor()),f) for u in x.split_uop(Ops.ADD)]) # a//c = (a-a%c)/c, if we can fold a%c, we can fold a//c @@ -195,43 +190,51 @@ def fold_divmod_congruence(d: UOp, x: UOp, y: UOp) -> UOp|None: def divide_by_gcd(d: UOp, x: UOp, y: UOp) -> UOp|None: # x//y -> (x//gcd)//(y//gcd) or x%y -> gcd*(x//gcd)%(y//gcd) - terms, factors = zip(*[(u.divides(f:=u.const_factor()),f) for u in x.split_uop(Ops.ADD)]) - if (gcd := math.gcd(y.arg, *factors)) == 1: return None - ret = sum(f//gcd * v for f,v in zip(factors, terms)).alu(d.op, y.const_like(y.arg//gcd)) + gcd = UOp.gcd(*x.split_uop(Ops.ADD), y).simplify() + if gcd.op is Ops.CONST and gcd.arg==1: return None + ret = unwrap(x.divide_exact(gcd)).alu(d.op, unwrap(y.divide_exact(gcd))) return ret*gcd if d.op is Ops.MOD else ret +def gcd_with_remainder(d: UOp, x: UOp, y: UOp): + # (gcd*x+r)//(gcd*d) -> (x+(r%d)//gcd)//d + r//(gcd*d) + # (gcd*x+r)%(gcd*d) -> gcd*(x+(r%d)//gcd)%d + r%gcd + # These only work for floordiv (and the corresponding remainder)! Thats why we check the sign of x,y and new_x + if ((c := y.arg) < 0) or x.vmin<0: return None + x_no_const, const = x.pop_const() + gcd = UOp.gcd(*x_no_const.split_uop(Ops.ADD), y).simplify() + assert gcd.op is Ops.CONST + if gcd.arg==1: return None + new_x = unwrap(x_no_const.divide_exact(gcd)).simplify() + (const%c)//gcd + if new_x.vmin<0: return None + ret = new_x.alu(d.op, x.ufix(c//gcd.arg)) + return ret*gcd + const%gcd.arg if d.op is Ops.MOD else ret+const//c + +def factor_remainder(d: UOp, x: UOp, y: UOp) -> UOp|None: + # (d*x+y)//d -> x+y//d or (d*x+y)%d + # for mod we go further and take the remainder of all factors to reduce their size + # These only work for floordiv (and the corresponding remainder)! Thats why we check the sign of x,y and new_x + if y.vmin<0 or x.vmin<0: return None + quo, rem = [], [] + for u in x.split_uop(Ops.ADD): + if (q:=u.divide_exact(y)) is not None: quo.append(q) + # if this is mod and y is a const, we can make the remainder factor sm + elif d.op is Ops.MOD and y.op is Ops.CONST and (c:=u.const_factor())%y.arg!=c: + rem.append(u.divides(c)*(c%y.arg)) + quo.append(u.const_like(0)) # we append this so we can check if something changed + else: rem.append(u) + new_x = sum(rem)+x.const_like(0) + if len(quo)==0 or new_x.vmin<0: return None + return new_x%y if d.op is Ops.MOD else new_x//y+sum(quo) + def nest_div_by_smallest_factor(d: UOp, x: UOp, y: UOp) -> UOp|None: # we try and nest the div and see if it allows the numerator to be simplified - if ((c := y.arg) < 0) or (x.dtype.count > 1): return None - factors = [u.const_factor() for u in x.pop_const()[0].split_uop(Ops.ADD)] - # div is the smallest factor of the denominator (greater than 1) out of all "factors" - # TODO: there are better ways to pick `div`, this sometimes adds extra divisions - # TODO: add same optimization for mod + if ((c := y.arg) < 0): return None + factors = [u.const_factor() for u in x.split_uop(Ops.ADD) if u.op not in (Ops.CONST, Ops.VCONST)] div = min([y.arg]+[abs(f) for f in factors if abs(f) > 1 and (c%f)==0]) - if (1 < div < c) and (newxs:=(newx:=(x//div)).simplify()) is not newx and x.vmin>=0 and newx.vmin>=0: return newxs//(c//div) - return None - -def simplify_remainder(d: UOp, x: UOp, y: UOp) -> UOp|None: - # we try and take out the quotient and see if it allows the numerator to be simplified - if ((c := y.arg) < 0) or (x.dtype.count > 1): return None - x_no_const,const = x.pop_const() - terms, factors = zip(*[(u.divides(f:=u.const_factor()),f) for u in x_no_const.split_uop(Ops.ADD)]) - quotients, remainders = zip(*[divmod(f, c) for f in factors]) - gcd = math.gcd(c, *remainders) # gcd without const! - if const%c==const and gcd==1 and not any(r==0 or (r!=f and d.op is Ops.MOD) for r,f in zip(remainders, factors)): return None - - quo, rem = x.const_like(const//c), x.const_like((const%c)//gcd) - for q,r,f,v in zip(quotients, remainders, factors, terms): - if d.op is Ops.IDIV and r!=0: - rem += f//gcd * v - else: - rem += r//gcd * v - quo += q * v - - # if numerator before/after is negative, and it has remainder, don't simplify because C divmod is different from python divmod. - if (x.vmin < 0 or rem.vmin < 0) and remainders: return None - if d.op is Ops.MOD: return gcd*(rem % (c//gcd)) + const%gcd - return rem//(c//gcd)+quo + newxs = fold_divmod_congruence(newx:=(x//div), x, y.const_like(div)) + if newxs is None: newxs = factor_remainder(newx, x, y.const_like(div)) + if div==y.arg or newxs is None or x.vmin<0 or newx.vmin<0: return None + return newxs//(c//div) def gep_through_wmma(gep:UOp, wmma:UOp): out_sz = prod(x[1] for x in wmma.arg[6][-1]) @@ -271,15 +274,6 @@ gep_pushing = PatternMatcher([ (UPat(Ops.WMMA, name="wmma").f(Ops.GEP, name="gep"), gep_through_wmma), ]) -cast_folding = PatternMatcher([ - (UPat.var('x', dtypes.ints+(dtypes.index,)).cast(dtypes.ints+(dtypes.index,), name="a").cast(name="b"), - lambda x,a,b: x.cast(b.dtype) if a.dtype.min<=x.vmin and x.vmax<=a.dtype.max else None), - # try to do math in int instead of long - (UPat(GroupOp.Binary, src=(UPat.var("x", dtypes.long), UPat.var("y", dtypes.long)), name="u"), lambda u,x,y: - x.cast(dtypes.int).alu(u.op, y.cast(dtypes.int)).cast(u.dtype) if not any(v.overflows(dtypes.int) for v in (u,x,y)) else None), - ((UPat.var("x", dtypes.index) + UPat.cvar("c")).cast(dtypes.sints, name="cast"), lambda x,c,cast:x.cast(cast.dtype)+c.cast(cast.dtype)), -]) - commutative = PatternMatcher([ # ** COMMUTATIVE flipping (only for index) ** # NOTE: this can break merging vector math by only flipping some of them @@ -289,6 +283,8 @@ commutative = PatternMatcher([ symbolic = symbolic_simple+commutative+PatternMatcher([ # ** boolean algebra ** (UPat.var("x") | (UPat.var("x") & UPat.var()), lambda x: x), # x|(x&y) -> x + # TODO: make a more general or folder like simplify_valid + (UPat.var("x", dtype=dtypes.bool) | UPat.var("x").logical_not(), lambda x: x.const_like(True)), # x|!x -> True # ** combine terms ** (UPat.var("x") * UPat.cvar("c0") + UPat.var("x") * UPat.cvar("c1"), lambda x,c0,c1: x*(c0+c1)), # (x*c0)+(x*c1) -> x*(c0+c1) ((UPat.var("y") + UPat.var("x") * UPat.cvar("c0")) + UPat.var("x") * UPat.cvar("c1"), lambda x,y,c0,c1: y+x*(c0+c1)), @@ -299,10 +295,9 @@ symbolic = symbolic_simple+commutative+PatternMatcher([ ((UPat.var("y") + UPat.var("x")) + UPat.var("x"), lambda y,x: y+x*2), ((UPat.var("x") / UPat.var("x2")) / UPat.var("x3"), lambda x,x2,x3: x/(x2*x3) if x2 is not x3 else None), # (x/x2)/x3 -> x/(x2*x3) (-1 * (UPat.var("x") + UPat.cvar("c")), lambda x,c: (-x)+(-c)), # -(x+c) -> -x + -c - # 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.var("cond", dtype=dtypes.bool).logical_not().where(UPat.var("t"), UPat.var("f")), lambda cond, t, f: cond.where(f,t)), + # ** where folding ** + (UPat.var("cond", dtype=dtypes.bool).logical_not().where(UPat.var("t"), UPat.var("f")), lambda cond, t, f: cond.where(f,t) + if f.arg is not Invalid else None), # alu of two where with same conds can combine, only do if true branch or false branch is const (UPat(GroupOp.Binary, name="alu", src=(UPat.var("c").where(UPat.var("t"), UPat.var("f")), UPat.var("c").where(UPat.var("tt"), UPat.var("ff")))), \ lambda alu,c,t,tt,f,ff: c.where(t.alu(alu.op, tt), f.alu(alu.op, ff)) if t.op == tt.op == Ops.CONST or f.op == ff.op == Ops.CONST else None), @@ -335,9 +330,6 @@ symbolic = symbolic_simple+commutative+PatternMatcher([ ((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 *** - # unrolled arange div folding - ((UPat()+(UPat()//UPat.cvar("d", vec=False)).or_casted()).named("divs"), lambda divs,d: fold_unrolled_divs(divs, d.arg)), - ((UPat()+((UPat()//UPat.cvar("d", vec=False)).or_casted()*UPat.cvar("c"))).named("divs"), lambda divs,d,c: fold_unrolled_divs(divs, d.arg, c.arg)), # generic lt folding (UPat.var("x", dtypes.index)=0 and n.vmin>=0 and d.arg>0 else None), ((UPat.var("x", dtypes.index)+UPat.cvar("c", vec=False)).named("n")//UPat.cvar("d", vec=False), lambda x,c,n,d: (-(-(c.arg%d.arg + x - (d.arg-1))//d) + c.arg//d.arg) if x.vmax<=0 and n.vmin>=0 and d.arg>0 else None), # ** mod ** # mod folding - (UPat.var("x") % UPat.var("d"), lambda x,d: -((-x)%d) if x.vmax <= 0 else None), - (UPat.var("x") % UPat.var("d"), lambda x,d: (x%(-d)) if d.vmax < 0 else None), -])+gep_pushing+cast_folding + (UPat.var("x", dtypes.index) % UPat.var("d"), lambda x,d: -((-x)%d) if x.vmax <= 0 else None), + (UPat.var("x", dtypes.index) % UPat.var("d"), lambda x,d: (x%(-d)) if d.vmax < 0 else None), + # cast/long folding + # if the intermediate cast doesnt narrow we can do it in one cast + (UPat.var('x').cast(name="a").cast(name="b"), lambda x,a,b: x.cast(b.dtype) if can_safe_cast(x.dtype, a.dtype) else None), + (UPat.var('x', dtypes.ints+(dtypes.index,)).cast(dtypes.ints+(dtypes.index,), name="a").cast(name="b"), + lambda x,a,b: x.cast(b.dtype) if a.dtype.min<=x.vmin and x.vmax<=a.dtype.max else None), + # try to do math in int instead of long + (UPat(GroupOp.Binary, src=(UPat.var("x", dtypes.long), UPat.var("y", dtypes.long)), name="u"), lambda u,x,y: + x.cast(dtypes.int).alu(u.op, y.cast(dtypes.int)).cast(u.dtype) if not any(v.overflows(dtypes.int) for v in (u,x,y)) else None), + ((UPat.var("x", dtypes.index) + UPat.cvar("c")).cast(dtypes.sints, name="cast"), lambda x,c,cast:x.cast(cast.dtype)+c.cast(cast.dtype)), +])+gep_pushing symbolic_flat = symbolic+PatternMatcher([ # ** combine terms (opinionated) ** @@ -395,7 +399,7 @@ def uop_given_valid(valid:UOp, uop:UOp) -> UOp|None: bounds:defaultdict[UOp, list[ConstType|None]] = defaultdict(lambda: [None, None]) for stmt in valid.split_uop(Ops.AND): try: expr, is_upper, c = parse_valid(stmt) - except ValueError: return uop # give up if we cannot parse the valid + except ValueError: continue # give up if we cannot parse the valid bounds[expr][int(is_upper)] = c # don't simplify any other gates, can lead to OOB, we substitute them back later @@ -466,6 +470,8 @@ REMOVE_FROM_BARRIER = {Ops.VECTORIZE, Ops.SINK, Ops.CAT, Ops.PTRCAT, Ops.NOOP} sym = symbolic_flat+PatternMatcher([ # simplify valid (UPat(Ops.AND, name="valid"), simplify_valid), + (UPat.var("cond").where(UPat.var("x", dtype=dtypes.index), invalid_pat), lambda cond,x,i: cond.where(newx, i) if + (newx:=uop_given_valid(cond, x)) is not x else None), # LOAD/STORE -> NOOP (UPat.var('x').store(UPat.var('x').load(), allow_any_len=True), lambda x: None if x.dtype.addrspace != AddrSpace.REG else x.src[0].src[0]), (UPat(Ops.LOAD, src=(UPat.cvar('c'))), lambda c: c), @@ -489,21 +495,20 @@ sym = symbolic_flat+PatternMatcher([ # ** where ** # push cast to branches (UPat.var("s").where(UPat.var("a"), UPat.var("b")).cast().named("cast"), lambda s,a,b,cast: s.where(a.cast(cast.dtype), b.cast(cast.dtype))), - # 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)), # ** pow ** ((UPat(Ops.POW, name="p"), lambda p: xpow(*p.src))), - # index true is index without op - (UPat(Ops.INDEX, src=(UPat.var("b"), UPat.var("idx"), UPat.const(dtypes.bool, True))), lambda b, idx: b.index(idx)), # ** load/store folding ** (UPat.store(UPat(Ops.INDEX, name="index"), UPat.load(UPat(Ops.INDEX, name="index"))), lambda index: UOp(Ops.NOOP)), (UPat.store(UPat(Ops.INDEX, name="index"), UPat.var("gate").where(UPat.var("alt"), UPat.load(UPat(Ops.INDEX, name="index"))), allow_any_len=True, name="store"), - lambda index, gate, alt, store: UOp.store(index.src[0].index(index.src[1], gate), alt, *store.src[2:])), + lambda index, gate, alt, store: UOp.store(index.src[0].index(gate.where(index.src[1], UOp.invalid())), alt, *store.src[2:])), # fold gated LOAD/STORE - (UPat().index(UPat(), UPat.const(dtypes.bool, True)).named("idx"), lambda idx: idx.replace(src=idx.src[0:2])), # remove True - (UPat((Ops.LOAD, Ops.STORE), src=(UPat().index(UPat(), UPat.const(dtypes.bool, False)).or_casted(),), allow_any_len=True, name="x"), - lambda x: UOp(Ops.NOOP) if x.op is Ops.STORE else x.const_like(0)), # NULL pointer store does nothing. NULL pointer load produces 0 + (UPat((Ops.LOAD, Ops.STORE), src=(UPat().index(UPat.const(dtypes.index, Invalid)).or_casted(),), allow_any_len=True, name="x"), + lambda x: UOp(Ops.NOOP) if x.op is Ops.STORE else x.const_like(0)), # invalid store does nothing. invalid load produces 0 + (UPat.var("c").where(UPat(Ops.LOAD, src=(UPat().index(UPat.var("idx"), UPat.var("c")).or_casted(),), allow_any_len=True, name="l"), UPat.var("a")), + lambda c,idx,l,a: l.replace(src=(l.src[0], a)+l.src[1:])), + (UPat.var("c").where(UPat.var("a"), UPat(Ops.LOAD, src=(UPat().index(UPat.var("idx"), UPat.var("c").logical_not()).or_casted(),), + allow_any_len=True, name="l")), lambda c,idx,l,a: l.replace(src=(l.src[0], a)+l.src[1:])), # remove VECTORIZE from SINK/BARRIER. TODO: SINK/BARRIER are really the same thing at GLOBAL/LOCAL levels (UPat(Ops.BARRIER, name="root"), lambda root: UOp(Ops.BARRIER, root.dtype, tuple(flatten(x.src if x.op in REMOVE_FROM_BARRIER else (x,) for x in root.src)), root.arg) diff --git a/tinygrad/viz/README b/tinygrad/viz/README index ce46d461cf..bdd038e44c 100644 --- a/tinygrad/viz/README +++ b/tinygrad/viz/README @@ -6,19 +6,18 @@ most uses of DEBUG >= 3 tiny-tools and a viewer for: -SAVE_SCHEDULE=1 TRACK_MATCH_STATS=2 -PROFILE=1 +ProfileEvents to use: -1. Run tinygrad with VIZ=1 and/or PROFILE=1 (this saves the pkls and launches the server (new process please!)) +1. Run tinygrad with VIZ=1 (this saves the pkls and launches the server (new process please!)) 2. That's it! This should be able to: 1. See all schedules (VIZ=1) 2. See all graphs and how they were rewritten (VIZ=1) 3. See generated code (VIZ=1) -4. See profile (PROFILE=1) +4. See profile (click on 'profiler') bunch of dev rules: * everything must be responsive to keyboard smashing! lag should never happen diff --git a/tinygrad/viz/index.html b/tinygrad/viz/index.html index 77d4e4292e..0310c6428b 100644 --- a/tinygrad/viz/index.html +++ b/tinygrad/viz/index.html @@ -90,11 +90,11 @@ .label :is(text, p) { font-weight: 350; } - rect.node { + g.node rect { stroke-width: 1.4; stroke: #4a4b57; } - rect.overlay { + g.overlay rect { fill: rgba(26, 27, 38, 0.5); } .edgePath { @@ -102,9 +102,12 @@ fill: none; stroke-width: 1.4px; } - .highlight rect, .edgePath.highlight, g.port circle { + g.node.highlight rect, .edgePath.highlight, g.port circle { stroke: #89C9A2; } + g.highlight.child rect, .edgePath.highlight.child { + stroke: #C888B0; + } #edge-labels g.port.highlight { display: block } @@ -131,17 +134,6 @@ .metadata > * + *, .rewrite-container > * + *, .ctx-list > * + * { margin-top: 12px; } - .stats-list > * + * { - margin-top: 8px; - } - .stats-list > p > * + * { - margin-top: 12px; - } - .stats-list { - width: 100%; - max-height: 240px; - overflow: auto; - } .ctx-list > ul > * + * { margin-top: 4px; } diff --git a/tinygrad/viz/js/index.js b/tinygrad/viz/js/index.js index f9e1fdb6c7..ab00826ce8 100644 --- a/tinygrad/viz/js/index.js +++ b/tinygrad/viz/js/index.js @@ -66,22 +66,21 @@ function renderDag(graph, additions, recenter) { // draw nodes const STROKE_WIDTH = 1.4; d3.select("#graph-svg").on("click", () => d3.selectAll(".highlight").classed("highlight", false)); - const nodes = d3.select("#nodes").selectAll("g").data(g.nodes().map(id => g.node(id)), d => d).join("g") + 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})`).classed("clickable", d => d.ref != null).on("click", (e,d) => { if (d.ref != null) return setCtxWithHistory(d.ref); const parents = g.predecessors(d.id); - if (parents == null) return; - const src = [...parents, d.id]; - nodes.classed("highlight", n => src.includes(n.id)); - d3.select("#edges").selectAll("path.edgePath").classed("highlight", e => src.includes(e.v) && e.w===d.id); - d3.select("#edge-labels").selectAll("g.port").classed("highlight", (_, i, nodes) => { - const [v, w] = nodes[i].id.split("-"); - return src.includes(v) && w===d.id; - }); + const children = g.successors(d.id); + if (parents == null && children == null) return; + const src = [...parents, ...children, d.id]; + nodes.classed("highlight", n => src.includes(n.id)).classed("child", n => children.includes(n.id)); + const matchEdge = (v, w) => (v===d.id && children.includes(w)) ? "highlight child " : (parents.includes(v) && w===d.id) ? "highlight " : ""; + d3.select("#edges").selectAll("path.edgePath").attr("class", e => matchEdge(e.v, e.w)+"edgePath"); + d3.select("#edge-labels").selectAll("g.port").attr("class", (_, i, n) => matchEdge(...n[i].id.split("-"))+"port"); e.stopPropagation(); }); nodes.selectAll("rect").data(d => [d]).join("rect").attr("width", d => d.width).attr("height", d => d.height).attr("fill", d => d.color) - .attr("x", d => -d.width/2).attr("y", d => -d.height/2).attr("class", d => d.className ?? "node"); + .attr("x", d => -d.width/2).attr("y", d => -d.height/2); nodes.selectAll("g.label").data(d => [d]).join("g").attr("class", "label").attr("transform", d => { const x = (d.width-d.padding*2)/2; const y = (d.height-d.padding*2)/2+STROKE_WIDTH; @@ -179,7 +178,7 @@ async function renderProfiler() { const u64 = () => { const ret = new Number(view.getBigUint64(offset, true)); offset += 8; return ret; } const f32 = () => { const ret = view.getFloat32(offset, true); offset += 4; return ret; } const optional = (i) => i === 0 ? null : i-1; - const dur = u32(), peak = u64(), indexLen = u32(), layoutsLen = u32(); + const dur = u32(), tracePeak = u64(), indexLen = u32(), layoutsLen = u32(); const textDecoder = new TextDecoder("utf-8"); const { strings, dtypeSize, markers } = JSON.parse(textDecoder.decode(new Uint8Array(buf, offset, indexLen))); offset += indexLen; // place devices on the y axis and set vertical positions @@ -193,20 +192,20 @@ async function renderProfiler() { // color by key (name/device) const colorMap = new Map(); data = {tracks:new Map(), axes:{}}; - const heightScale = d3.scaleLinear().domain([0, peak]).range([4,maxheight=100]); + const heightScale = d3.scaleLinear().domain([0, tracePeak]).range([4,maxheight=100]); for (let i=0; i ({ color, st, width:ctx.measureText(st).width })); if (e.ref != null) ref = {ctx:e.ref, step:0}; else if (ref != null) { @@ -267,12 +267,13 @@ async function renderProfiler() { timestamps.push(dur); const height = heightScale(peak); const yscale = d3.scaleLinear().domain([0, peak]).range([height, 0]); - for (const [_, {dtype, sz, nbytes, y, x:steps}] of buf_shapes) { + for (const [num, {dtype, sz, nbytes, y, x:steps}] of buf_shapes) { const x = steps.map(s => timestamps[s]); - const arg = {tooltipText:`${dtype} len:${formatUnit(sz)}\n${formatUnit(nbytes, "B")}`}; + const dur = x.at(-1)-x[0]; + const arg = {tooltipText:`${dtype} len:${formatUnit(sz)}\n${formatUnit(nbytes, "B")}\nnum:${num}\nalive for ${formatTime(dur)}`}; shapes.push({ x, y0:y.map(yscale), y1:y.map(y0 => yscale(y0+nbytes)), arg, fillColor:cycleColors(colorScheme.BUFFER, shapes.length) }); } - data.tracks.set(k, { shapes, visible:[], offsetY, height, peak, scaleFactor:maxheight*4/height }); + data.tracks.set(k, { shapes, visible, offsetY, height, peak, scaleFactor:maxheight*4/height }); div.style("height", height+padding+"px").style("cursor", "pointer").on("click", (e) => { const newFocus = e.currentTarget.id === focusedDevice ? null : e.currentTarget.id; let offset = 0; @@ -297,41 +298,40 @@ async function renderProfiler() { // rescale to match current zoom const xscale = d3.scaleLinear().domain([0, dur]).range([0, canvas.clientWidth]); const visibleX = xscale.range().map(zoomLevel.invertX, zoomLevel).map(xscale.invert, xscale); + const st = visibleX[0], et = visibleX[1]; xscale.domain(visibleX); // draw shapes for (const [_, { offsetY, shapes, visible }] of data.tracks) { visible.length = 0; for (const e of shapes) { - if (e.width == null) { start = e.x[0]; end = end = e.x[e.x.length-1]; } - else { start = e.x; end = e.x+e.width; } - if (start>visibleX[1] || endet || e.x.at(-1)=0; i--) p.lineTo(x[i], offsetY+e.y1[i]); - p.closePath(); - ctx.fill(p); - // NOTE: y coordinates are in reverse order - for (let i = 0; i=0; i--) ctx.lineTo(x[i], offsetY+e.y1[i]); + ctx.closePath(); + ctx.fillStyle = e.fillColor; ctx.fill(); continue; } // contiguous rect - const x = xscale(start); - const width = xscale(end)-x; - ctx.fillRect(x, offsetY+e.y, width, e.height); - visible.push({ y0:offsetY+e.y, y1:offsetY+e.y+e.height, x0:x, x1:x+width, arg:e.arg }); + if (e.x>et || e.x+e.width width) { if (labelWidth !== 0) ctx.fillText("...", labelX, labelY); @@ -658,28 +658,6 @@ async function main() { const metadata = document.querySelector(".metadata"); const [code, lang] = ctx.fmt != null ? [ctx.fmt, "cpp"] : [ret[currentRewrite].uop, "python"]; metadata.replaceChildren(codeBlock(step.code_line, "python", { loc:step.loc, wrap:true }), codeBlock(code, lang, { wrap:false })); - if (ctx.runtime_stats != null) { - const div = metadata.appendChild(document.createElement("div")); - div.className = "stats-list"; - for (const [i, s] of ctx.runtime_stats.entries()) { - const p = div.appendChild(document.createElement("p")); - if (ctx.runtime_stats.length > 1) p.innerText = `Run ${i+1}/${ctx.runtime_stats.length}`; - const table = div.appendChild(document.createElement("table")); - const tbody = table.appendChild(document.createElement("tbody")); - for (const { name, value, unit, subunits } of s.data) { - const mainRow = appendRow(tbody, name, value, unit, "main-row"); - if (!subunits?.length) continue; - const subunitRow = tbody.appendChild(document.createElement("tr")); - subunitRow.style.display = "none"; - mainRow.onclick = () => subunitRow.style.display = subunitRow.style.display === "none" ? "table-row" : "none"; - mainRow.style.cursor = "pointer"; - const td = subunitRow.appendChild(document.createElement("td")); - td.colSpan = 2; - const table = td.appendChild(document.createElement("table")); - for (const u of subunits) appendRow(table, u.name, u.value, unit, "sub-row"); - } - } - } // ** rewrite steps if (step.match_count >= 1) { const rewriteList = metadata.appendChild(document.createElement("div")); @@ -744,7 +722,7 @@ appendResizer(document.querySelector(".metadata-parent"), { minWidth: 20, maxWid // **** keyboard shortcuts -document.addEventListener("keydown", async function(event) { +document.addEventListener("keydown", (event) => { const { currentCtx, currentStep, currentRewrite, expandSteps } = state; // up and down change the step or context from the list const changeStep = expandSteps && ctxs[currentCtx].steps?.length; diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index 500a9c7fb4..e58839bbf2 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -6,7 +6,7 @@ from decimal import Decimal from http.server import BaseHTTPRequestHandler from urllib.parse import parse_qs, urlparse from typing import Any, TypedDict, Generator -from tinygrad.helpers import colored, getenv, tqdm, unwrap, word_wrap, TRACEMETA, ProfileEvent, ProfileRangeEvent, TracingKey, ProfilePointEvent +from tinygrad.helpers import colored, getenv, tqdm, unwrap, word_wrap, TRACEMETA, ProfileEvent, ProfileRangeEvent, TracingKey, ProfilePointEvent, temp from tinygrad.uop.ops import TrackedGraphRewrite, UOp, Ops, printable, GroupOp, srender, sint, sym_infer from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry, Device from tinygrad.renderer import ProgramSpec @@ -19,7 +19,7 @@ uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0", Ops.INDEX: "#e8ffa0", Ops.WMMA: "#efefc0", Ops.VIEW: "#C8F9D4", Ops.MULTI: "#f6ccff", Ops.KERNEL: "#3e7f55", **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80", Ops.BUFFER_VIEW: "#E5EAFF", Ops.BLOCK: "#C4A484", Ops.BLOCKEND: "#C4A4A4", Ops.BUFFER: "#B0BDFF", Ops.COPY: "#a040a0", Ops.FUSE: "#FFa500", - Ops.ALLREDUCE: "#ff40a0", Ops.MSELECT: "#d040a0", Ops.MSTACK: "#d040a0", Ops.CONTIGUOUS: "#FFC14D", + Ops.ALLREDUCE: "#ff40a0", Ops.MSELECT: "#d040a0", Ops.MSTACK: "#d040a0", Ops.CONTIGUOUS: "#FFC14D", Ops.REALIZE: "#C1C14D", Ops.CHILDREN: "#80ffc0", Ops.CHILD: "#80fff0", Ops.BUFFERIZE: "#FF991C", Ops.REWRITE_ERROR: "#ff2e2e"} # VIZ API @@ -36,8 +36,6 @@ def get_metadata(trace_bufs:list[tuple]) -> list[dict]: steps = [{"name":s.name, "loc":s.loc, "depth":s.depth, "match_count":len(s.matches), "code_line":printable(s.loc), "query":f"/ctxs?ctx={i}&idx={j}"} for j,s in enumerate(v)] ret.append(r:={"name":k.display_name, "steps":steps}) - # use the first key to get runtime profiling data about this context - if getenv("PROFILE_VALUE") >= 2 and k.keys: r["runtime_stats"] = get_runtime_stats(k.keys[0]) # program spec metadata if isinstance(k.ret, ProgramSpec): steps.append({"name":"View Disassembly", "query":f"/disasm?ctx={i}"}) @@ -73,11 +71,12 @@ def uop_to_json(x:UOp) -> dict[int, dict]: if u.op is Ops.VIEW: argst = ("\n".join([f"{shape_to_str(v.shape)} / {shape_to_str(v.strides)}"+("" if v.offset == 0 else f" / {srender(v.offset)}")+ (f"\nMASK {mask_to_str(v.mask)}" if v.mask is not None else "") for v in unwrap(u.st).views])) + if u.op in GroupOp.Movement: argst = (mask_to_str if u.op in {Ops.SHRINK, Ops.PAD} else shape_to_str)(u.arg) label = f"{str(u.op).split('.')[1]}{(chr(10)+word_wrap(argst.replace(':', ''))) if u.arg is not None else ''}" if u.dtype != dtypes.void: label += f"\n{u.dtype}" for idx,x in enumerate(u.src): if x in excluded: - arg = f"{x.arg:g}" if x.op is Ops.CONST and dtypes.is_float(u.dtype) else f"{x.arg}" + arg = f"{x.arg:g}" if x.op is Ops.CONST and dtypes.is_float(x.dtype) else f"{x.arg}" label += f"\n{x.op.name}{idx} {arg}" + (f" {x.src[0].op}" if len(x.src) else "") try: if u.op not in {Ops.VIEW, Ops.BUFFER, Ops.KERNEL, Ops.ASSIGN, Ops.COPY, Ops.SINK, *GroupOp.Buffer} and u.st is not None: @@ -90,7 +89,7 @@ def uop_to_json(x:UOp) -> dict[int, dict]: # NOTE: kernel already has metadata in arg if TRACEMETA >= 2 and u.metadata is not None and u.op is not Ops.KERNEL: label += "\n"+repr(u.metadata) 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":u.tag} + "ref":ref, "tag":repr(u.tag) if u.tag is not None else None} return graph @functools.cache @@ -201,13 +200,6 @@ def get_profile(profile:list[ProfileEvent]) -> bytes|None: index = json.dumps({"strings":list(scache), "dtypeSize":dtype_size, "markers":[{"ts":int(e.ts-start_ts), **e.arg} for e in markers]}).encode() return struct.pack(" list[dict]: - ret:list[dict] = [] - for e in profile: - if isinstance(e, ProfileRangeEvent) and e.en is not None and e.name == key: - ret.append({"device":e.device, "data":[{"name":"Duration", "value":float(e.en-e.st), "unit":"us"}]}) - return ret - # ** Assembly analyzers def get_llvm_mca(asm:str, mtriple:str, mcpu:str) -> dict: @@ -303,8 +295,8 @@ class TCPServerWithReuse(socketserver.TCPServer): allow_reuse_address = True if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument('--kernels', type=pathlib.Path, help='Path to kernels', default=None) - parser.add_argument('--profile', type=pathlib.Path, help='Path profile', default=None) + parser.add_argument('--kernels', type=pathlib.Path, help='Path to kernels', default=pathlib.Path(temp("rewrites.pkl", append_user=True))) + parser.add_argument('--profile', type=pathlib.Path, help='Path to profile', default=pathlib.Path(temp("profile.pkl", append_user=True))) args = parser.parse_args() with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: @@ -317,7 +309,7 @@ if __name__ == "__main__": ctxs = get_metadata(load_pickle(args.kernels)) - profile_ret = get_profile(profile:=load_pickle(args.profile)) + profile_ret = get_profile(load_pickle(args.profile)) server = TCPServerWithReuse(('', PORT), Handler) reloader_thread = threading.Thread(target=reloader)