Compare commits

..
Author SHA1 Message Date
geohot 2650f15c4b fix test_matmul_relu_cat 2025-10-07 14:15:31 +08:00
113 changed files with 3162 additions and 2438 deletions
+47 -60
View File
@@ -52,16 +52,14 @@ jobs:
- name: reset process replay
run: python3.11 test/external/process_replay/reset.py
- name: Run Stable Diffusion
run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=800 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=500 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
- name: Run Stable Diffusion without fp16
run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=900 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=700 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
- name: Run Stable Diffusion v2
# TODO: very slow step time
run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=10000 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=1600 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
# process replay can't capture this, the graph is too large
# TODO: too slow
# - name: Run SDXL
# run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=5000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
- name: Run SDXL
run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=3000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
- name: Run model inference benchmark
run: METAL=1 python3.11 test/external/external_model_benchmark.py
- name: Test speed vs torch
@@ -101,7 +99,7 @@ jobs:
- name: Run GPT2
run: |
BENCHMARK_LOG=gpt2_nojit JIT=0 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
BENCHMARK_LOG=gpt2 JIT=1 ASSERT_MIN_STEP_TIME=13 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
BENCHMARK_LOG=gpt2 JIT=1 ASSERT_MIN_STEP_TIME=8 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
- name: Run GPT2 w HALF
run: BENCHMARK_LOG=gpt2_half HALF=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
- name: Run GPT2 w HALF/BEAM
@@ -110,19 +108,14 @@ jobs:
run: BENCHMARK_LOG=olmoe python3.11 examples/olmoe.py
- name: Train MNIST
run: time PYTHONPATH=. TARGET_EVAL_ACC_PCT=96.0 python3.11 examples/beautiful_mnist.py | tee beautiful_mnist.txt
# NOTE: this is failing in CI. it is not failing on my machine and I don't really have a way to debug it
# the error is "RuntimeError: Internal Error (0000000e:Internal Error)"
#- name: Run 10 CIFAR training steps
# run: BENCHMARK_LOG=cifar_10steps JIT=1 ASSERT_MIN_STEP_TIME=3000 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar.txt
#- name: Run 10 CIFAR training steps w HALF
# run: BENCHMARK_LOG=cifar_10steps_half JIT=2 ASSERT_MIN_STEP_TIME=3000 STEPS=10 DEFAULT_FLOAT=HALF python3.11 examples/hlb_cifar10.py | tee train_cifar_half.txt
- name: Run 10 CIFAR training steps
run: BENCHMARK_LOG=cifar_10steps JIT=1 ASSERT_MIN_STEP_TIME=330 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar.txt
- name: Run 10 CIFAR training steps w HALF
run: BENCHMARK_LOG=cifar_10steps_half JIT=2 ASSERT_MIN_STEP_TIME=385 STEPS=10 DEFAULT_FLOAT=HALF python3.11 examples/hlb_cifar10.py | tee train_cifar_half.txt
#- name: Run 10 CIFAR training steps w BF16
# run: STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3.11 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
# TODO: too slow
# - name: Run 10 CIFAR training steps w winograd
# run: BENCHMARK_LOG=cifar_10steps_wino JIT=1 ASSERT_MIN_STEP_TIME=150 WINO=1 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar_wino.txt
- name: Run 10 CIFAR training steps w winograd
run: BENCHMARK_LOG=cifar_10steps_wino JIT=1 ASSERT_MIN_STEP_TIME=150 WINO=1 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar_wino.txt
- name: UsbGPU boot time
run: sudo -E PYTHONPATH=. DEBUG=2 AM_RESET=1 AMD=1 AMD_IFACE=USB time python3.11 test/test_tiny.py TestTiny.test_plus
- name: UsbGPU tiny tests
@@ -220,9 +213,8 @@ jobs:
run: DEBUG=2 CUDA=1 python -m pytest -rA test/test_tiny.py
- name: Run Stable Diffusion
run: BENCHMARK_LOG=stable_diffusion NV=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
# TODO: too slow
# - name: Run SDXL
# run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=2000 CAPTURE_PROCESS_REPLAY=0 NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
- name: Run SDXL
run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=2000 CAPTURE_PROCESS_REPLAY=0 NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
- name: Run LLaMA
run: |
BENCHMARK_LOG=llama_nojit NV=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
@@ -246,9 +238,9 @@ jobs:
- name: Run GPT2
run: |
BENCHMARK_LOG=gpt2_nojit NV=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
BENCHMARK_LOG=gpt2 NV=1 JIT=1 ASSERT_MIN_STEP_TIME=4 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
BENCHMARK_LOG=gpt2 NV=1 JIT=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
- name: Run GPT2 w HALF
run: BENCHMARK_LOG=gpt2_half NV=1 HALF=1 ASSERT_MIN_STEP_TIME=6 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
run: BENCHMARK_LOG=gpt2_half NV=1 HALF=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
- name: Run GPT2 w HALF/BEAM
run: BENCHMARK_LOG=gpt2_half_beam NV=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
- uses: actions/upload-artifact@v4
@@ -307,27 +299,24 @@ jobs:
rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
- name: reset process replay
run: test/external/process_replay/reset.py
# TODO: too slow
# - 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
# TODO: too slow
# - name: Fuzz Padded Tensor Core GEMM (PTX)
# 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: 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 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
run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=270 NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=85 NV=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=310 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=68 NV=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=310 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
# TODO: too slow
# - name: Run 10 CIFAR training steps w winograd
# run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=350 NV=1 CAPTURE_PROCESS_REPLAY=0 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=75 NV=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=35 NV=1 CAPTURE_PROCESS_REPLAY=0 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
run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
- name: Run full CIFAR training steps w 6 GPUS
run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 NV=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 NV=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
- name: Run MLPerf resnet eval on training data
run: time BENCHMARK_LOG=resnet_eval NV=1 MODEL=resnet python3 examples/mlperf/model_eval.py
#- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
@@ -426,10 +415,9 @@ jobs:
- name: Test AM warm start time
run: time AMD=1 python3 test/test_tiny.py TestTiny.test_plus
- name: Run Stable Diffusion
run: BENCHMARK_LOG=stable_diffusion ASSERT_MIN_STEP_TIME=550 AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
# TODO: too slow
# - name: Run SDXL
# run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=3200 CAPTURE_PROCESS_REPLAY=0 AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
run: BENCHMARK_LOG=stable_diffusion ASSERT_MIN_STEP_TIME=450 AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
- name: Run SDXL
run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=1400 CAPTURE_PROCESS_REPLAY=0 AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
- name: Run LLaMA 7B
run: |
BENCHMARK_LOG=llama_nojit AMD=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
@@ -520,20 +508,19 @@ jobs:
- name: Train MNIST
run: time PYTHONPATH=. AMD=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
- name: Run 10 CIFAR training steps
run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=330 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
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=330 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
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
# TODO: too slow
# - 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 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
run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
#- name: Run full CIFAR training steps w 6 GPUS
# run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
# run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
#- name: Run full CIFAR training steps w 6 GPUS (REMOTE)
# run: time BENCHMARK_LOG=cifar_6gpu_remote REMOTE=1 REMOTEDEV=AMD DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu_remote.txt
# run: time BENCHMARK_LOG=cifar_6gpu_remote REMOTE=1 REMOTEDEV=AMD DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu_remote.txt
- uses: actions/upload-artifact@v4
with:
name: Speed (AMD Training)
@@ -619,21 +606,21 @@ jobs:
- name: reset process replay
run: test/external/process_replay/reset.py
- name: benchmark openpilot 0.9.9 driving_vision
run: BENCHMARK_LOG=openpilot_0_9_9_vision PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
run: BENCHMARK_LOG=openpilot_0_9_9_vision ASSERT_MIN_STEP_TIME=30 PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
- name: benchmark openpilot 0.9.9 driving_policy
run: BENCHMARK_LOG=openpilot_0_9_9_policy PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
run: BENCHMARK_LOG=openpilot_0_9_9_policy ASSERT_MIN_STEP_TIME=45 PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
- name: benchmark openpilot 0.9.9 dmonitoring
run: BENCHMARK_LOG=openpilot_0_9_9_dmonitoring PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
run: BENCHMARK_LOG=openpilot_0_9_9_dmonitoring ASSERT_MIN_STEP_TIME=70 PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
- name: openpilot compile3 0.9.9 driving_vision
run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=22 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
- name: openpilot compile3 0.9.9 driving_policy
run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=7 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
- name: openpilot compile3 0.9.9 dmonitoring
run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=15 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
- name: openpilot compile3 Space Lab policy + vision
run: |
PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/22aec22a10ce09384d4a4af2a0bbff08d54af7e0c888503508f356fae4ff0e29
PYTHONPATH="." ASSERT_MIN_STEP_TIME=26 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/c824f68646a3b94f117f01c70dc8316fb466e05fbd42ccdba440b8a8dc86914b
PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/22aec22a10ce09384d4a4af2a0bbff08d54af7e0c888503508f356fae4ff0e29
PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/c824f68646a3b94f117f01c70dc8316fb466e05fbd42ccdba440b8a8dc86914b
- name: benchmark MobileNetV2 on DSP
run: |
# generate quantized weights
@@ -708,7 +695,7 @@ jobs:
AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
- name: Run full CIFAR training w 1 GPU
run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
# TODO: enable
# - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
# run: BENCHMARK_LOG=resnet_10steps AMD=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee am_train_resnet_one_gpu.txt
@@ -771,7 +758,7 @@ jobs:
- name: Test LLAMA-3
run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --benchmark --temperature 0 | tee nv_llama3_beam.txt
- name: Run full CIFAR training w 1 GPU
run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee nv_train_cifar_one_gpu.txt
run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee nv_train_cifar_one_gpu.txt
#- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
# run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee nv_train_resnet_one_gpu.txt
- name: Run 10 MLPerf Bert training steps (1 gpu)
+106 -10
View File
@@ -144,7 +144,7 @@ jobs:
sudo apt update || true
sudo apt install -y --no-install-recommends ninja-build
- name: Test beautiful_mnist in torch with TINY_BACKEND
run: CPU=1 CPU_LLVM=1 TARGET_EVAL_ACC_PCT=96.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
run: SPLIT_REDUCEOP=0 FUSE_ARANGE=1 CPU=1 CPU_LLVM=1 TARGET_EVAL_ACC_PCT=96.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
- name: Test some torch tests (expect failure)
run: python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true
@@ -160,8 +160,10 @@ jobs:
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 dtype with Python emulator (with RANGEIFY)
run: |
RANGEIFY=0 DEBUG=1 PYTHON=1 python3 -m pytest -n=auto test/test_dtype.py test/test_dtype_alu.py
RANGEIFY=1 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
@@ -265,13 +267,10 @@ jobs:
run: python -c "from tinygrad import Device; assert Device.DEFAULT == 'CPU', Device.DEFAULT"
- name: Run unit tests
run: CPU=1 python -m pytest -n=auto test/unit/ --durations=20
- name: Check SPEC=1
run: SPEC=1 python3 test/test_tiny.py
- name: Run targetted tests on NULL backend
run: NULL=1 python3 -m unittest test.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step test/device/test_null.py
# TODO: too slow
# - name: Run SDXL on NULL backend
# run: NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
- name: Run SDXL on NULL backend
run: NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
- name: Run Clip tests for SD MLPerf on NULL backend
run: NULL=1 python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
# TODO: support fake weights
@@ -333,6 +332,10 @@ jobs:
run: |
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: Test CL IMAGE=2 ops + training (rangeify)
run: |
RANGEIFY=1 CL=1 IMAGE=2 python -m pytest -n=auto test/test_ops.py --durations=20
RANGEIFY=1 CL=1 IMAGE=2 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
- name: Run process replay tests
uses: ./.github/actions/process-replay
@@ -377,7 +380,10 @@ jobs:
llvm: 'true'
- name: Test openpilot model kernel count and gate usage
run: |
ALLOWED_KERNEL_COUNT=190 ALLOWED_READ_IMAGE=2041 ALLOWED_GATED_READ_IMAGE=41 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
ALLOWED_KERNEL_COUNT=208 ALLOWED_READ_IMAGE=2160 ALLOWED_GATED_READ_IMAGE=16 RANGEIFY=0 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 model with rangeify
run: |
ALLOWED_KERNEL_COUNT=190 ALLOWED_READ_IMAGE=2041 ALLOWED_GATED_READ_IMAGE=33 RANGEIFY=1 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 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)
@@ -514,6 +520,88 @@ jobs:
# ****** Feature Tests ******
testrangeifycpu:
name: Linux (rangeify) CPU
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-minimal-llvm
deps: testing_minimal
opencl: 'true'
llvm: "true"
- name: Test CPU=1 RANGEIFY=1
# TODO: add more passing tests here
run: |
CPU=1 CPU_LLVM=0 RANGEIFY=1 python3 -m pytest -n auto --durations 20 \
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 test/test_assign.py test/test_multitensor.py test/test_const_folding.py
- name: Test CPU=1 DEVECTORIZE=0 (RANGEIFY=1)
run: CPU=1 CPU_LLVM=0 RANGEIFY=1 DEVECTORIZE=0 FUSE_ARANGE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"
- name: Test CPU=1 CPU_LLVM=1 RANGEIFY=1
run: |
CPU=1 CPU_LLVM=1 RANGEIFY=1 python3 -m pytest -n auto --durations 20 test/test_edgecases.py
- name: Test Docs RANGEIFY=1
run: |
RANGEIFY=1 python docs/abstractions2.py
# RANGEIFY=2 isn't supported
#- 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 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 test/unit/test_linalg.py --durations=20
- name: Test METAL=1 RANGEIFY=1
run: |
METAL=1 RANGEIFY=1 python -m pytest -n=auto test/test_ops.py test/test_multitensor.py --durations=20
METAL=1 MAX_KERNEL_BUFFERS=6 RANGEIFY=1 PYTHONPATH=. python test/test_multitensor.py TestBatchNorm.test_batchnorm
- name: Run process replay tests
uses: ./.github/actions/process-replay
testdevectorize:
name: Linux (devectorize)
runs-on: ubuntu-24.04
@@ -533,7 +621,7 @@ jobs:
- name: Test LLVM=1 DEVECTORIZE=0 for model
run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 test/models/test_efficientnet.py
- name: Test CPU=1 DEVECTORIZE=0
run: CPU=1 CPU_LLVM=0 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"
run: CPU=1 CPU_LLVM=0 DEVECTORIZE=0 FUSE_ARANGE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"
testdsp:
name: Linux (DSP)
@@ -636,6 +724,8 @@ jobs:
run: |
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 pytest (amd) with RANGEIFY
run: RANGEIFY=1 python -m pytest test/test_linearizer.py::TestLinearizer::test_where_fold
- name: Run process replay tests
uses: ./.github/actions/process-replay
@@ -955,3 +1045,9 @@ jobs:
run: |
python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
python -m pytest -n=auto test/test_tiny.py test/test_ops.py --durations=20
- name: Run pytest (${{ matrix.backend }}) with RANGEIFY
if: matrix.backend=='webgpu'
env:
RANGEIFY: 1
shell: bash
run: python -m pytest -n=auto test/test_tiny.py test/test_ops.py --durations=20
+7 -5
View File
@@ -42,6 +42,7 @@ import struct
from tinygrad.dtype import dtypes
from tinygrad.device import Buffer, Device
from tinygrad.uop.ops import UOp, Ops
from tinygrad.shape.shapetracker import ShapeTracker
# allocate some buffers + load in values
out = Buffer(DEVICE, 1, dtypes.int32).allocate()
@@ -50,14 +51,13 @@ b = Buffer(DEVICE, 1, dtypes.int32).allocate().copyin(memoryview(bytearray(struc
# NOTE: a._buf is the same as the return from cpu.allocator.alloc
# describe the computation
idx = UOp.const(dtypes.index, 0)
buf_1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 1)
buf_2 = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 2)
ld_1 = UOp(Ops.LOAD, dtypes.int32, (buf_1.index(idx),))
ld_2 = UOp(Ops.LOAD, dtypes.int32, (buf_2.index(idx),))
ld_1 = UOp(Ops.LOAD, dtypes.int32, (buf_1.view(ShapeTracker.from_shape((1,))),))
ld_2 = UOp(Ops.LOAD, dtypes.int32, (buf_2.view(ShapeTracker.from_shape((1,))),))
alu = ld_1 + ld_2
output_buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 0)
st_0 = UOp(Ops.STORE, dtypes.void, (output_buf.index(idx), alu))
st_0 = UOp(Ops.STORE, dtypes.void, (output_buf.view(ShapeTracker.from_shape((1,))), alu))
s = UOp(Ops.SINK, dtypes.void, (st_0,))
# convert the computation to a "linearized" format (print the format)
@@ -80,6 +80,8 @@ print("******** third, the UOp ***********")
from tinygrad.engine.realize import run_schedule
from tinygrad.engine.schedule import create_schedule_with_vars
from tinygrad.helpers import RANGEIFY
from tinygrad.schedule.kernelize import get_kernelize_map
from tinygrad.schedule.rangeify import get_rangeify_map
# allocate some values + load in values
@@ -93,7 +95,7 @@ out = a + b
s = UOp(Ops.SINK, dtypes.void, (out,))
# group the computation into kernels
becomes_map = get_rangeify_map(s)
becomes_map = get_rangeify_map(s) if RANGEIFY else get_kernelize_map(s)
# the compute maps to an assign
assign = becomes_map[a+b].base
+1 -1
View File
@@ -10,7 +10,7 @@ Directories are listed in order of how they are processed.
Group UOps into kernels.
::: tinygrad.schedule.rangeify.get_rangeify_map
::: tinygrad.schedule.kernelize.get_kernelize_map
options:
members: false
show_labels: false
+1 -1
View File
@@ -10,7 +10,7 @@ GPUS = [f'{Device.DEFAULT}:{i}' for i in range(getenv("GPUS", 1))]
# override tinygrad defaults
dtypes.default_float = dtypes.half
Context(FUSE_OPTIM=1).__enter__()
Context(FUSE_ARANGE=1, FUSE_OPTIM=1).__enter__()
# from https://github.com/tysam-code/hlb-CIFAR10/blob/main/main.py
batchsize = getenv("BS", 1024)
+1
View File
@@ -145,6 +145,7 @@ hyp = {
},
}
@Context(FUSE_ARANGE=getenv("FUSE_ARANGE", 1))
def train_cifar():
def set_seed(seed):
+4 -3
View File
@@ -3,7 +3,7 @@ from pathlib import Path
import multiprocessing
from tinygrad import Device, GlobalCounters, Tensor, TinyJit, dtypes
from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, Profiling
from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, FUSE_CONV_BW, Profiling
from tinygrad.nn.state import get_parameters, get_state_dict, load_state_dict, safe_load, safe_save
from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, Adam, AdamW
@@ -707,7 +707,7 @@ def train_unet3d():
```BASEDIR=<folder_path> ./examples/mlperf/scripts/setup_kits19_dataset.sh```
2) To start training the model, run the following:
```time PYTHONPATH=. WANDB=1 TRAIN_BEAM=3 GPUS=6 BS=6 MODEL=unet3d python3 examples/mlperf/model_train.py```
```time PYTHONPATH=. WANDB=1 TRAIN_BEAM=3 FUSE_CONV_BW=1 GPUS=6 BS=6 MODEL=unet3d python3 examples/mlperf/model_train.py```
"""
from examples.mlperf.losses import dice_ce_loss
from examples.mlperf.metrics import dice_score
@@ -749,6 +749,7 @@ def train_unet3d():
"train_beam": TRAIN_BEAM,
"eval_beam": EVAL_BEAM,
"wino": WINO.value,
"fuse_conv_bw": FUSE_CONV_BW.value,
"gpus": GPUS,
"default_float": dtypes.default_float.name
}
@@ -1308,7 +1309,7 @@ def train_llama3():
EVAL_BS = config["EVAL_BS"] = getenv("EVAL_BS", 16)
EVAL_TARGET = config["EVAL_TARGET"] = getenv("EVAL_TARGET", 5.6)
# LR=1e-4 TRAIN_ON_VAL=1 DEFAULT_FLOAT=bfloat16 JITBEAM=2 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B WARMUP_STEPS=36 DECAY_STEPS=360 SEQLEN=512 PYTHONPATH=. AMD=1 AMD_LLVM=0 MODEL=llama3 python3 examples/mlperf/model_train.py
# LR=1e-4 TRAIN_ON_VAL=1 DEFAULT_FLOAT=bfloat16 FUSE_ARANGE=1 JITBEAM=2 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B WARMUP_STEPS=36 DECAY_STEPS=360 SEQLEN=512 PYTHONPATH=. AMD=1 AMD_LLVM=0 MODEL=llama3 python3 examples/mlperf/model_train.py
# trains to 7
opt_adamw_beta_1 = 0.9
+2 -11
View File
@@ -1,4 +1,4 @@
import os, sys, pickle, time, re
import os, sys, pickle, time
import numpy as np
if "FLOAT16" not in os.environ: os.environ["FLOAT16"] = "1"
if "IMAGE" not in os.environ: os.environ["IMAGE"] = "2"
@@ -52,8 +52,6 @@ def compile(onnx_file):
kernel_count += 1
read_image_count += ei.prg.p.src.count("read_image")
gated_read_image_count += ei.prg.p.src.count("?read_image")
for v in [m.group(1) for m in re.finditer(r'(val\d+)\s*=\s*read_imagef\(', ei.prg.p.src)]:
if len(re.findall(fr'[\?\:]{v}\.[xyzw]', ei.prg.p.src)) > 0: gated_read_image_count += 1
print(f"{kernel_count=}, {read_image_count=}, {gated_read_image_count=}")
if (allowed_kernel_count:=getenv("ALLOWED_KERNEL_COUNT", -1)) != -1:
assert kernel_count == allowed_kernel_count, f"different kernels! {kernel_count=}, {allowed_kernel_count=}"
@@ -79,20 +77,13 @@ def test_vs_compile(run, new_inputs, test_val=None):
**{k:Tensor(v, device="NPY").realize() for k,v in new_inputs_numpy.items() if 'img' not in k}}
# run 20 times
step_times = []
for _ in range(20):
st = time.perf_counter()
out = run(**inputs)
mt = time.perf_counter()
val = out.numpy()
et = time.perf_counter()
step_times.append((et-st)*1e3)
print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {step_times[-1]:6.2f} ms")
if (assert_time:=getenv("ASSERT_MIN_STEP_TIME")):
min_time = min(step_times)
assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"
print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {(et-st)*1e3:6.2f} ms")
print(out, val.shape, val.dtype)
if test_val is not None: np.testing.assert_equal(test_val, val)
print("**** test done ****")
+3 -1
View File
@@ -2,7 +2,9 @@ import sys
from tinygrad import Tensor, fetch, GlobalCounters, dtypes
from tinygrad.uop.ops import UOp
from tinygrad.nn.onnx import OnnxRunner
from tinygrad.schedule.kernelize import get_kernelize_map
from tinygrad.schedule.rangeify import get_rangeify_map
from tinygrad.helpers import RANGEIFY
from tinygrad.engine.schedule import create_schedule_with_vars
from tinygrad.engine.realize import run_schedule
@@ -33,7 +35,7 @@ if __name__ == "__main__":
if not in_target_path[s]:
independent_set[s] = None
independent = UOp.sink(*independent_set.keys())
kernelized = get_rangeify_map(independent)
kernelized = (get_rangeify_map if RANGEIFY else get_kernelize_map)(independent)
independent = independent.substitute(kernelized)
schedule, var_vals = create_schedule_with_vars(independent)
run_schedule(schedule)
+2 -4
View File
@@ -269,14 +269,12 @@ if __name__ == "__main__":
# load in weights
with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
load_state_dict(model, torch_load(fetch('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', 'sd-v1-4.ckpt'))['state_dict'], verbose=False, strict=False, realize=False)
load_state_dict(model, torch_load(fetch('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', 'sd-v1-4.ckpt'))['state_dict'], strict=False)
if args.fp16:
for k,v in get_state_dict(model).items():
if k.startswith("model"):
v.replace(v.cast(dtypes.float16))
Tensor.realize(*get_state_dict(model).values())
v.replace(v.cast(dtypes.float16).realize())
# run through CLIP to get context
tokenizer = Tokenizer.ClipTokenizer()
+1 -1
View File
@@ -32,7 +32,7 @@ if __name__ == "__main__":
lr = 5e-3
transform = ComposeTransforms([
lambda x: [Image.fromarray(xx).resize((64, 64)) for xx in x],
lambda x: [Image.fromarray(xx, mode='L').resize((64, 64)) for xx in x],
lambda x: np.stack([np.asarray(xx) for xx in x], 0),
lambda x: x / 255.0,
lambda x: np.tile(np.expand_dims(x, 1), (1, 3, 1, 1)).astype(np.float32),
+3 -2
View File
@@ -49,7 +49,8 @@ def rangeify_kernel3():
b = Tensor.empty(N,N)
c = a@b
#c = c.reshape((32,2,16,4,32,2,16,4)).contiguous()
sink = c.schedule()[-1].ast
with Context(RANGEIFY=1):
sink = c.schedule()[-1].ast
#print(sink)
opts = [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.LOCAL, 0, 16), Opt(OptOps.UPCAST, 0, 2)]
@@ -328,7 +329,7 @@ if __name__ == "__main__":
elif HL == 1: hprg = hl_spec_kernel3()
else: hprg = hand_spec_kernel3()
if HL == 3:
with Context(BLOCK_REORDER=0):
with Context(RANGEIFY=1, BLOCK_REORDER=0):
prg = get_program(hprg, Device.default.renderer)
else:
prg = get_program(hprg, Device.default.renderer)
+1
View File
@@ -7,6 +7,7 @@ bert_train_params = {
"GPUS": 6,
"BS": 96,
"EVAL_BS": 96,
"FUSE_ARANGE": 1,
"BASEDIR": "/raid/datasets/wiki",
}
+1 -1
View File
@@ -50,7 +50,7 @@ def ioctls_from_header():
hdr = (pathlib.Path(__file__).parent / "kfd_ioctl.h").read_text().replace("\\\n", "")
pattern = r'#define\s+(AMDKFD_IOC_[A-Z0-9_]+)\s+AMDKFD_IOW?R?\((0x[0-9a-fA-F]+),\s+struct\s([A-Za-z0-9_]+)\)'
matches = re.findall(pattern, hdr, re.MULTILINE)
return {int(nr, 0x10):(name, getattr(kfd_ioctl, "struct_"+sname, None)) for name, nr, sname in matches}
return {int(nr, 0x10):(name, getattr(kfd_ioctl, "struct_"+sname)) for name, nr, sname in matches}
nrs = ioctls_from_header()
@ctypes.CFUNCTYPE(ctypes.c_int, ctypes.c_int, ctypes.c_ulong, ctypes.c_void_p)
File diff suppressed because it is too large Load Diff
+1 -1
View File
@@ -81,7 +81,7 @@ def lin_to_feats(lin:Kernel, use_sts=True):
ret = [float(x) for x in ret]
if use_sts:
my_sts = dedup([(x.shape == lin.full_shape, x.is_expanded(), any(v.mask is not None for v in x.views), len(x.views)) for x in lin.sts])
my_sts = dedup([(x.shape == lin.full_shape, x.real_strides(), any(v.mask is not None for v in x.views), len(x.views)) for x in lin.sts])
assert len(my_sts) < MAX_BUFS
sts_len = 3 + 5*MAX_DIMS
for s in my_sts:
+6 -7
View File
@@ -155,7 +155,6 @@ class RGP:
device_event = device_events[device]
sqtt_events = [x for x in profile if isinstance(x, ProfileSQTTEvent) and x.device == device_event.device]
if len(sqtt_events) == 0: raise RuntimeError(f"Device {device_event.device} doesn't contain SQTT data")
device_props = sqtt_events[0].props
sqtt_itrace_enabled = any([event.itrace for event in sqtt_events])
sqtt_itrace_masked = not all_same([event.itrace for event in sqtt_events])
sqtt_itrace_se_mask = functools.reduce(lambda a,b: a|b, [int(event.itrace) << event.se for event in sqtt_events], 0) if sqtt_itrace_masked else 0
@@ -193,14 +192,14 @@ class RGP:
flags=0,
trace_shader_core_clock=0x93f05080,
trace_memory_clock=0x4a723a40,
device_id={110000: 0x744c, 110003: 0x7480}[device_props['gfx_target_version']],
device_id=0x744c,
device_revision_id=0xc8,
vgprs_per_simd=1536,
sgprs_per_simd=128*16,
shader_engines=device_props['array_count'] // device_props['simd_arrays_per_engine'],
compute_unit_per_shader_engine=device_props['simd_count'] // device_props['simd_per_cu'] // (device_props['array_count'] // device_props['simd_arrays_per_engine']),
simd_per_compute_unit=device_props['simd_per_cu'],
wavefronts_per_simd=device_props['max_waves_per_simd'],
shader_engines=6,
compute_unit_per_shader_engine=16,
simd_per_compute_unit=2,
wavefronts_per_simd=16,
minimum_vgpr_alloc=4,
vgpr_alloc_granularity=8,
minimum_sgpr_alloc=128,
@@ -219,7 +218,7 @@ class RGP:
vram_bus_width=384, # 384-bit
l2_cache_size=6 * 1024 * 1024, # 6 MB
l1_cache_size=32 * 1024, # 32 KB per SIMD (?)
lds_size=device_props['lds_size_in_kb'] * 1024,
lds_size=65536, # 64 KB per CU
gpu_name=b'NAVI31',
alu_per_clock=0,
texture_per_clock=0,
+40
View File
@@ -0,0 +1,40 @@
import time
from extra.optimization.helpers import load_worlds, ast_str_to_ast
from tinygrad import Device
from tinygrad.codegen.lowerer import pm_lowerer, get_index
from tinygrad.uop.ops import graph_rewrite
from tinygrad.codegen.opt.kernel import Kernel
from tinygrad.codegen.opt.postrange import Scheduler
from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
from tinygrad.helpers import getenv
if __name__ == "__main__":
renderer = Device.default.renderer
ast_strs = load_worlds()
if (n:=getenv("N", -1)) != -1: ast_strs = ast_strs[n:n+1]
good = 0
for i, ast_str in enumerate(ast_strs):
ast = ast_str_to_ast(ast_str)
st = time.perf_counter()
lin = Kernel(ast, renderer)
opt1 = hand_coded_optimizations(lin)
et_lin = time.perf_counter() - st
lowered = graph_rewrite(ast, pm_lowerer, ctx=get_index(ast), bottom_up=True)
st = time.perf_counter()
sch = Scheduler(lowered, renderer)
sch.convert_loop_to_global()
sch.simplify_merge_adjacent()
opt2 = hand_coded_optimizations(sch)
et_sch = time.perf_counter() - st
if opt1 != opt2:
print(f"******* {i:6d}")
print("Kernel: ", lin.colored_shape(), "->", lin.apply_opts(opt1).colored_shape())
print("Scheduler: ", sch.colored_shape(), "->", sch.apply_opts(opt2).colored_shape())
print(opt1)
print(opt2)
else:
good += 1
print(f"******* {i:6d} MATCH {good/(i+1)*100:.2f}% -- {et_lin/et_sch:4.2f}x speedup")
+10 -9
View File
@@ -227,15 +227,16 @@ class TestTorchBackend(unittest.TestCase):
np.testing.assert_equal(result.cpu().numpy(), [3., 3., 2.])
def test_mnist_index(self):
GlobalCounters.reset()
from tinygrad.nn.datasets import mnist
X_train, Y_train, _, _ = mnist()
X_train = torch.tensor(X_train.float().numpy(), device=device)
Y_train = torch.tensor(Y_train.cast('int64').numpy(), device=device)
samples = torch.randint(0, X_train.shape[0], (32,))
X,Y = X_train[samples], Y_train[samples]
X.cpu(), Y.cpu()
self.assertLessEqual(GlobalCounters.global_ops, 10_000_000)
with Context(FUSE_ARANGE=1, SPLIT_REDUCEOP=0):
GlobalCounters.reset()
from tinygrad.nn.datasets import mnist
X_train, Y_train, _, _ = mnist()
X_train = torch.tensor(X_train.float().numpy(), device=device)
Y_train = torch.tensor(Y_train.cast('int64').numpy(), device=device)
samples = torch.randint(0, X_train.shape[0], (32,))
X,Y = X_train[samples], Y_train[samples]
X.cpu(), Y.cpu()
self.assertLessEqual(GlobalCounters.global_ops, 10_000_000)
def _test_diagonal(self, *shape):
a = torch.randn(*shape, dtype=torch.float32, device=device)
-43
View File
@@ -1,43 +0,0 @@
# extra/weekly_commits_table.py
import os, subprocess, datetime as dt
NAMES = ["chenyu","George Hotz","nimlgen","qazal","Sieds Lykles","wozeparrot"]
REPO = os.environ.get("REPO_PATH",".")
today = dt.date.today()
days = [(today - dt.timedelta(i)).strftime("%Y-%m-%d") for i in range(6,-1,-1)]
seen = {d:{n:False for n in NAMES} for d in days}
cmd = ["git","-C",REPO,"log","--use-mailmap","--since=7 days ago","--no-merges",
"--date=short","--pretty=%ad%x09%aN%x09%ae"]
out = subprocess.run(cmd, capture_output=True, text=True).stdout.splitlines()
for line in out:
try: d, name, email = line.split("\t")
except: continue
if d in seen:
low = (name+" "+email).lower()
for n in NAMES:
if n.lower() in low: seen[d][n] = True
# --- width-aware padding so emoji align ---
try:
from wcwidth import wcswidth as _wcswidth
vlen = lambda s: _wcswidth(s)
except Exception:
vlen = lambda s: sum(2 if ch in "✅❌" else 1 for ch in s)
pad = lambda s,w: s + " " * max(0, w - vlen(s))
w_date = 10
w_cols = [max(3, vlen(n)) for n in NAMES]
header = " | ".join([pad("date", w_date)] + [pad(n, w_cols[i]) for i,n in enumerate(NAMES)])
rule = "-+-".join(["-"*w_date] + ["-"*w for w in w_cols])
rows=[]
for d in days:
cells = ["" if seen[d][n] else "" for n in NAMES]
rows.append(" | ".join([pad(d, w_date)] + [pad(c, w_cols[i]) for i,c in enumerate(cells)]))
print("** Commits by day (last 7) **")
print("```")
print("\n".join([header, rule] + rows))
print("```")
+1 -1
View File
@@ -1,6 +1,6 @@
[pytest]
norecursedirs = extra
timeout = 300
timeout = 240
timeout_method = thread
timeout_func_only = true
testpaths = test
+1 -1
View File
@@ -13,6 +13,6 @@ if __name__ == "__main__":
(Tensor.empty(BS, 16, 512, 512), Tensor.empty(BS, 512, 16, 64).permute(0,2,1,3)), # qk@v
]
for t0, t1 in tensors:
print(f"{t0.shape=}, {t0.uop.st.is_expanded()=}, {t1.shape=}, {t1.uop.st.is_expanded()=}")
print(f"{t0.shape=}, {t0.uop.st.real_strides()=}, {t1.shape=}, {t1.uop.st.real_strides()=}")
for _ in range(5):
t0.dot(t1, dtype=acc_dtype).realize()
+4
View File
@@ -39,6 +39,10 @@ if __name__ == "__main__":
step_times.append(t:=(time.perf_counter_ns() - st)*1e-6)
print(f"jitted: {t:7.4f} ms")
if (assert_time:=getenv("ASSERT_MIN_STEP_TIME")):
min_time = min(step_times)
assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"
suffix = ""
if IMAGE.value < 2: suffix += f"_image{IMAGE.value}" # image=2 has no suffix for compatibility
if getenv("FLOAT16") == 1: suffix += "_float16"
+7 -8
View File
@@ -1,9 +1,8 @@
from extra.models.resnet import ResNet50
from tinygrad import Tensor, nn, Device
from tinygrad.helpers import Profiling, Timing, getenv
from tinygrad.uop.ops import Ops, UOp
from tinygrad.codegen import get_rewrites_for_renderer, apply_rewrites
from tinygrad.codegen.late.control_flow import schedule
from tinygrad.uop.ops import Ops
from tinygrad.codegen import get_rewrites_for_renderer, apply_rewrites, rewrites_for_linearizer
from tinygrad.uop.spec import type_verify
if __name__ == "__main__":
@@ -29,10 +28,10 @@ if __name__ == "__main__":
asts = list({x.ast.key:x.ast for x in sched if x.ast.op is Ops.SINK}.values())
if (restrict_kernel := getenv("RESTRICT_KERNEL", -1)) != -1: asts = asts[restrict_kernel:restrict_kernel+1]
rewrites = get_rewrites_for_renderer(Device.default.renderer, linearizer=LINEARIZE)
rewrites = get_rewrites_for_renderer(Device.default.renderer, linearizer=False)
with Profiling(PROFILE, fn="/tmp/rewrite.prof"):
with Timing("***** model rewrite in "):
rewritten_uops: list[UOp] = []
rewritten_uops = []
for u in asts:
rewritten_uops.append(apply_rewrites(u, rewrites))
@@ -40,7 +39,7 @@ if __name__ == "__main__":
with Timing("***** model linearize in "):
uops_line = []
for u in rewritten_uops:
uops_line.append(schedule(list(u.toposort())))
uops_line.append(apply_rewrites(u, rewrites_for_linearizer))
with Timing("***** model verify in "):
for u in uops_line: type_verify(u)
print(sum(len(u) for u in uops_line))
for u in uops_line: type_verify(u.arg.lst)
print(sum(len(u.arg.lst) for u in uops_line))
+20 -18
View File
@@ -4,7 +4,7 @@ import numpy as np
import torch
from tinygrad import GlobalCounters, Tensor, Device
from tinygrad.helpers import getenv
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
@@ -164,7 +164,7 @@ class TestOpt(unittest.TestCase):
def test_permute_was_pushed(self):
a = Tensor.randn(16, 16, 16)
with CLCache(1):
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(1):
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(1):
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(1):
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()
@@ -217,22 +217,24 @@ class TestOpt(unittest.TestCase):
assert cache_len == 1, "reduceop was rerun!"
def test_expand_reduce_is_folded_on_same_axis(self):
for axis in [0, 1]:
with Context(FUSE_CONV_BW=1):
for axis in [0, 1]:
for n in [4, 8, 16]:
b = torch.ones(n, n).sum(axis).reshape(n, 1).expand(n, n).sum(axis)
with CLCache(allowed=3 if RANGEIFY else 2):
a = Tensor.ones(n, n).contiguous().sum(axis).reshape(n, 1).expand(n, n).sum(axis)
a.realize()
np.testing.assert_allclose(a.numpy(), b.numpy(), rtol=1e-3, atol=1e-5)
def test_expand_reduce_is_folded_on_different_axes(self):
with Context(FUSE_CONV_BW=1):
axis1, axis2 = 0, 1
for n in [4, 8, 16]:
b = torch.ones(n, n).sum(axis).reshape(n, 1).expand(n, n).sum(axis)
with CLCache(allowed=3):
a = Tensor.ones(n, n).contiguous().sum(axis).reshape(n, 1).expand(n, n).sum(axis)
b = torch.ones(n, n).sum(axis1).reshape(n, 1).expand(n, n).sum(axis2)
with CLCache(allowed=3 if RANGEIFY else 2):
a = Tensor.ones(n, n).contiguous().sum(axis1).reshape(n, 1).expand(n, n).sum(axis2)
a.realize()
np.testing.assert_allclose(a.numpy(), b.numpy(), rtol=1e-3, atol=1e-5)
def test_expand_reduce_is_folded_on_different_axes(self):
axis1, axis2 = 0, 1
for n in [4, 8, 16]:
b = torch.ones(n, n).sum(axis1).reshape(n, 1).expand(n, n).sum(axis2)
with CLCache(allowed=3):
a = Tensor.ones(n, n).contiguous().sum(axis1).reshape(n, 1).expand(n, n).sum(axis2)
a.realize()
np.testing.assert_allclose(a.numpy(), b.numpy(), rtol=1e-3, atol=1e-5)
if __name__ == '__main__':
unittest.main()
+4 -13
View File
@@ -1,8 +1,7 @@
import gc
from tinygrad import Tensor, UOp, Device, nn
from tinygrad import Tensor, UOp, Device
from tinygrad.shape.shapetracker import views_to_valid_uop
from tinygrad.engine.realize import method_cache, get_program
from test.test_tiny import TestTiny
def uops_allocated(): return sum([isinstance(x, UOp) for x in gc.get_objects()])
def print_uops():
@@ -47,16 +46,9 @@ def realized_gradient():
z = y.matmul(x).sum()
z.backward()
Tensor.realize(x, y, z, x.grad, y.grad)
def nn_batchnorm(): nn.BatchNorm(64)
def nn_conv2d(): nn.Conv2d(64, 64, 3)
def plus(): TestTiny().test_plus()
def mnist(): TestTiny().test_mnist()
def mnist_backward(): TestTiny().test_mnist_backward()
tests = [start, single_tensor, two_plus_two, two_plus_two_schedule, two_plus_two_kernel,
two_plus_two_linearize, two_plus_two_realize, two_plus_two_item, gradient_test,
realized_eye, realized_list, kernel_matmul, realized_matmul, realized_gradient,
nn_batchnorm, nn_conv2d, plus, mnist, mnist_backward]
realized_eye, realized_list, kernel_matmul, realized_matmul, realized_gradient]
if __name__ == "__main__":
gc.disable()
@@ -69,12 +61,11 @@ if __name__ == "__main__":
# these caches will keep uops alive
method_cache.clear()
views_to_valid_uop.cache_clear()
Tensor._device_seeds.clear()
Tensor._device_rng_counters.clear()
new_uops = uops_allocated()
print_uops()
gc.collect()
new_uops_gc = uops_allocated()
print(f"{t.__name__:30s}: {new_uops:3d} -> {new_uops_gc:3d}")
if new_uops != start_uops: print_uops()
assert new_uops == start_uops
#print_uops()
+13 -1
View File
@@ -2,6 +2,7 @@ import random
from tinygrad.helpers import getenv, DEBUG, colored, trange
from tinygrad.shape.shapetracker import ShapeTracker
from test.external.fuzz_shapetracker import shapetracker_ops
from test.external.fuzz_shapetracker import do_permute, do_reshape_split_one, do_reshape_combine_two, do_flip, do_pad
from test.unit.test_shapetracker_math import st_equal, MultiShapeTracker
def fuzz_plus() -> tuple[ShapeTracker, ShapeTracker]:
@@ -13,10 +14,21 @@ def fuzz_plus() -> tuple[ShapeTracker, ShapeTracker]:
st_sum = backup + m.sts[1]
return m.sts[0], st_sum
# shrink and expand aren't invertible, and stride is only invertible in the flip case
invertible_shapetracker_ops = [do_permute, do_reshape_split_one, do_reshape_combine_two, do_flip, do_pad]
def fuzz_invert() -> tuple[ShapeTracker, ShapeTracker]:
start = ShapeTracker.from_shape((random.randint(1, 10), random.randint(1, 10), random.randint(1, 10)))
m = MultiShapeTracker([start])
for _ in range(8): random.choice(invertible_shapetracker_ops)(m)
inv = m.sts[0].invert(start.shape)
st_sum = (m.sts[0] + inv) if inv else None
return start, st_sum
if __name__ == "__main__":
if seed:=getenv("SEED"): random.seed(seed)
total = getenv("CNT", 1000)
for fuzz in [globals()[f'fuzz_{x}'] for x in getenv("FUZZ", "plus").split(",")]:
for fuzz in [globals()[f'fuzz_{x}'] for x in getenv("FUZZ", "invert,plus").split(",")]:
same_but_neq = 0
for _ in trange(total, desc=f"{fuzz}"):
st1, st2 = fuzz()
+13
View File
@@ -0,0 +1,13 @@
from tinygrad.shape.shapetracker import ShapeTracker
from test.external.fuzz_shapetracker import shapetracker_ops as st_ops
from test.unit.test_shapetracker_math import MultiShapeTracker
from tinygrad.helpers import getenv
import random
random.seed(getenv("SEED", 42))
for i in range(getenv("CNT", 2000)):
if getenv("DEBUG", 0) >= 1: print()
N = random.randint(1, 10000)
mst = MultiShapeTracker([ShapeTracker.from_shape((N,))]) # st_ops don't mutate regular shapetrackers for some reason
for j in range(20): random.choice(st_ops)(mst)
assert mst.sts[0].real_size() <= N, f"{N=}, real_size={mst.sts[0].real_size()}, st={mst.sts[0]}"
+2 -2
View File
@@ -8,7 +8,7 @@ ASSERT_DIFF = int((flag:="[pr]") in os.getenv("COMMIT_MESSAGE", flag) or flag in
if not int(os.getenv("ASSERT_PROCESS_REPLAY", "1")): ASSERT_DIFF = 0
try:
from tinygrad.schedule.rangeify import get_rangeify_map
from tinygrad.schedule.kernelize import get_kernelize_map
from tinygrad.renderer import Renderer, ProgramSpec
from tinygrad.engine.realize import get_program
from tinygrad.uop.ops import UOp, Ops, KernelInfo
@@ -44,7 +44,7 @@ class ProcessReplayWarning(Warning): pass
def replay_kernelize(ret:dict[UOp, UOp], big_sink:UOp) -> tuple[str, str, tuple[Any, ...]]:
UOp.unique_num = itertools.count(max([u.arg for u in big_sink.toposort() if u.op is Ops.UNIQUE], default=0)+1)
new_sink = big_sink.substitute(get_rangeify_map(big_sink))
new_sink = big_sink.substitute(get_kernelize_map(big_sink))
def to_str(ret:UOp) -> str:
asts = [repr(u.arg.ast) for u in ret.toposort() if u.op is Ops.KERNEL]
return "\n".join([f"{len(asts)} kernels", *asts])
+5 -2
View File
@@ -1,4 +1,4 @@
import time, struct
import time, struct, unittest
from typing import Any, Callable
import numpy as np
from tinygrad import Tensor, dtypes, Device
@@ -7,7 +7,7 @@ from tinygrad.tensor import _to_np_dtype
from tinygrad.engine.realize import Runner
from tinygrad.dtype import DType
from tinygrad.nn.state import get_parameters
from tinygrad.helpers import T, CI
from tinygrad.helpers import T, CI, RANGEIFY
from tinygrad.codegen import full_rewrite
from tinygrad.runtime.ops_python import PythonProgram, PythonRenderer, PythonCompiler
@@ -62,3 +62,6 @@ def not_support_multi_device():
# 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)
def expect_rangeify_fails(fxn): return (unittest.expectedFailure if RANGEIFY else (lambda f:f))(fxn)
def expect_nonrangeify_fails(fxn): return (unittest.expectedFailure if not RANGEIFY else (lambda f:f))(fxn)
+1 -5
View File
@@ -17,7 +17,7 @@ def ioctls_from_header():
pattern = r'#define\s+(AMDKFD_IOC_[A-Z0-9_]+)\s+AMDKFD_(IOW?R?)\((0x[0-9a-fA-F]+),\s+struct\s([A-Za-z0-9_]+)\)'
matches = re.findall(pattern, hdr, re.MULTILINE)
return type("KFD_IOCTLS", (object, ), {name: int(nr, 0x10) for name, _, nr, _ in matches}), \
{int(nr, 0x10): getattr(kfd, "struct_"+sname, None) for name, idir, nr, sname in matches}
{int(nr, 0x10): getattr(kfd, "struct_"+sname) for name, idir, nr, sname in matches}
kfd_ioctls, kfd_headers = ioctls_from_header()
class KFDFileDesc(VirtFileDesc):
@@ -115,10 +115,6 @@ class AMDDriver(VirtDriver):
struct = kfd_headers[nr].from_address(argp)
if nr == kfd_ioctls.AMDKFD_IOC_ACQUIRE_VM: pass
elif nr == kfd_ioctls.AMDKFD_IOC_RUNTIME_ENABLE: pass
elif nr == kfd_ioctls.AMDKFD_IOC_GET_VERSION:
struct.major_version = 1
struct.minor_version = 14
elif nr == kfd_ioctls.AMDKFD_IOC_ALLOC_MEMORY_OF_GPU:
if struct.gpu_id not in self.gpus: return -1
struct.handle = self._alloc_handle()
+1 -1
View File
@@ -1,4 +1,4 @@
import ctypes, time
import ctypes, ctypes.util, time
import tinygrad.runtime.autogen.nv_gpu as nv_gpu
from enum import Enum, auto
from test.mockgpu.gpu import VirtGPU
+2 -2
View File
@@ -94,7 +94,7 @@ class TestRealWorld(unittest.TestCase):
@TinyJit
def test(t, v):
with Context(JIT=0): return model(t, v).realize()
helper_test("test_gpt2", lambda: (Tensor([[1,]]),Variable("pos", 1, 100).bind(1)), test, 0.23 if CI else 0.9, 160 if CI else 468, all_jitted=True)
helper_test("test_gpt2", lambda: (Tensor([[1,]]),Variable("pos", 1, 100).bind(1)), test, 0.23 if CI else 0.9, 160 if CI else 396, all_jitted=True)
@unittest.skipIf(CI and Device.DEFAULT == "CPU", "slow")
def test_train_mnist(self):
@@ -176,7 +176,7 @@ class TestRealWorld(unittest.TestCase):
for v in data.values(): v.to_(Device.DEFAULT)
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.31, 358)
data["masked_lm_ids"], data["masked_lm_weights"], data["next_sentence_labels"]), train, 0.28, 357)
if __name__ == '__main__':
unittest.main()
-1
View File
@@ -149,7 +149,6 @@ class TestFloat4(unittest.TestCase):
assert TestFloat4.count_float4(uops) == (1, 1)
@unittest.skip("Ops.VIEW no longer exists")
def test_half4_load_unrolled(self):
# from llama 7B shard 4 gpus
ast = UOp(Ops.SINK, dtypes.void, arg=None, src=(
+2 -1
View File
@@ -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
@@ -351,6 +351,7 @@ 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()
+25 -9
View File
@@ -1,7 +1,7 @@
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.engine.realize import CompiledRunner, ExecItem, get_program
from tinygrad.uop.ops import Ops
@@ -25,6 +25,22 @@ class TestArange(unittest.TestCase):
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):
GlobalCounters.reset()
with Context(FUSE_ARANGE=0, NOOPT=noopt):
out = Tensor.rand(16384)
out.realize()
unfused_ops = GlobalCounters.global_ops
GlobalCounters.reset()
with Context(FUSE_ARANGE=1, NOOPT=noopt):
out = Tensor.rand(16384)
out.realize()
print(f"fused {GlobalCounters.global_ops} unfused {unfused_ops}")
self.assertLessEqual(GlobalCounters.global_ops, unfused_ops*2)
def test_fused_rand_less_ops_opt(self): self.test_fused_rand_less_ops(0)
DSET, DDIM = 2048, 32
class TestIndexing(unittest.TestCase):
@@ -32,7 +48,7 @@ class TestIndexing(unittest.TestCase):
needle = Tensor.zeros(16384, dtype=dtypes.int).contiguous()
needle[1337] = 1
needle.realize()
with Context(NOOPT=1):
with Context(NOOPT=1, FUSE_ARANGE=1):
GlobalCounters.reset()
out = ((Tensor.arange(1,16385)-1)*needle).sum()
sched = out.schedule()
@@ -45,7 +61,7 @@ class TestIndexing(unittest.TestCase):
idxs = Tensor([0,3,5,6]).realize()
real_index = dataset.numpy()[idxs.numpy()]
print("*** indexing ***")
with Context(NOOPT=1):
with Context(NOOPT=1, FUSE_ARANGE=1):
GlobalCounters.reset()
rng = Tensor.ones(4, DDIM, DSET, dtype=dtypes.int)._cumalu(axis=-1, op=Ops.ADD, _include_initial=True).reshape(4, DDIM, DSET, 1)
idxs = idxs.reshape(4,1,1,1).expand(4, DDIM, DSET, 1)
@@ -61,7 +77,7 @@ class TestIndexing(unittest.TestCase):
def test_index_variable(self):
dataset = Tensor.rand(DSET, DDIM).realize()
v = Variable("v", 0, DDIM-1)
with Context(NOOPT=1):
with Context(NOOPT=1, FUSE_ARANGE=1, SPLIT_REDUCEOP=0):
GlobalCounters.reset()
vb = Tensor(v.bind(12))
comp = dataset[vb].numpy()
@@ -90,12 +106,12 @@ class TestIndexing(unittest.TestCase):
idxs = Tensor([0,3,5,6]).realize()
real_index = dataset.numpy()[idxs.numpy()]
print("*** indexing ***")
with Context(NOOPT=noopt):
with Context(NOOPT=noopt, FUSE_ARANGE=1):
GlobalCounters.reset()
X = dataset[idxs]
assert X.shape == (4,DDIM)
sched = X.schedule()
self.assertEqual(len(sched), 1)
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())
@@ -105,7 +121,7 @@ class TestIndexing(unittest.TestCase):
def test_index_fused_out_of_bounds(self):
dataset = Tensor.rand(256, 256).realize()
idxs = Tensor([-19238, -257, 256, 495, 10982377]).realize()
with Context(NOOPT=1):
with Context(NOOPT=1, FUSE_ARANGE=1):
X = dataset[idxs]
np.testing.assert_equal(X.numpy(), 0)
@@ -114,7 +130,7 @@ class TestIndexing(unittest.TestCase):
if Device.DEFAULT == "WEBGPU": op_limit *= 15
from tinygrad.nn.datasets import mnist
X_train, Y_train, _, _ = mnist()
with Context(NOOPT=noopt, SPLIT_REDUCEOP=split_reduceop):
with Context(NOOPT=noopt, FUSE_ARANGE=1, SPLIT_REDUCEOP=split_reduceop):
samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0]).realize()
GlobalCounters.reset()
x = X_train[samples].numpy()
@@ -134,7 +150,7 @@ class TestIndexing(unittest.TestCase):
# TODO: why is a new realize needed here
emb_w = emb.weight.realize().numpy()
x = Tensor([1,2,3,4])
with Context(NOOPT=noopt):
with Context(NOOPT=noopt, FUSE_ARANGE=1):
GlobalCounters.reset()
z = emb(x).realize()
self.assertLessEqual(GlobalCounters.global_ops, op_limit)
+35 -26
View File
@@ -1,9 +1,10 @@
#!/usr/bin/env python
import unittest
import contextlib
import numpy as np
from tinygrad import dtypes, Tensor, TinyJit, GlobalCounters, Variable
from tinygrad.device import is_dtype_supported
from tinygrad.helpers import temp
from tinygrad.helpers import temp, RANGEIFY
N = 200 # has to be bigger than the cache to fail
@@ -270,6 +271,8 @@ class TestAssign(unittest.TestCase):
b.assign(a.contiguous()).realize()
assert GlobalCounters.kernel_count - kc == 2
# passing in RANGEIFY=1, RANGEIFY=0 asserts permuted assigns it can't fuse
def assert_permuted_assign(self): return self.assertRaisesRegex(RuntimeError, "contiguous") if not RANGEIFY else contextlib.nullcontext()
def test_permuted_assignment(self):
a = Tensor(np.arange(N*N, dtype=np.float32)).reshape(N,N)
b = Tensor(np.arange(N*N, dtype=np.float32)).reshape(N,N)
@@ -277,13 +280,13 @@ class TestAssign(unittest.TestCase):
b.realize()
ba1 = a.uop.base.realized
bb1 = b.uop.base.realized
a = a.permute(1,0)
a += b
a.realize()
ba2 = a.uop.base.realized
np.testing.assert_allclose(a.numpy(), np.arange(N*N).reshape((N,N)) + np.arange(N*N).reshape((N,N)).transpose(1,0))
# permute and base are the same buffer
assert ba1 == ba2 and ba1 != bb1
with self.assertRaises((RuntimeError, AssertionError)):
a = a.permute(1,0)
a += b
a.realize()
ba2 = a.uop.base.realized
assert ba1 != ba2 and ba1 != bb1
np.testing.assert_allclose(a.numpy(), np.arange(N*N).reshape((N,N)) + np.arange(N*N).reshape((N,N)).transpose(1,0))
def test_post_permuted_assignment(self):
a = Tensor(np.arange(N*N, dtype=np.float32)).reshape(N,N)
@@ -293,13 +296,15 @@ class TestAssign(unittest.TestCase):
#GlobalCounters.cache = []
ba1 = a.uop.base.realized # noqa: F841
bb1 = b.uop.base.realized # noqa: F841
a.assign(a.permute(1,0) + b) # this should not work!
a.realize()
ba2 = a.uop.base.realized # noqa: F841
# NOTE: don't test that it's assigned
#assert ba1 == ba2 and ba1 != bb1
np.testing.assert_allclose(a.numpy(), np.arange(N*N).reshape((N,N)) + np.arange(N*N).reshape((N,N)).transpose(1,0))
with self.assert_permuted_assign():
a.assign(a.permute(1,0) + b) # this should not work!
a.realize()
ba2 = a.uop.base.realized # noqa: F841
# NOTE: don't test that it's assigned
#assert ba1 == ba2 and ba1 != bb1
np.testing.assert_allclose(a.numpy(), np.arange(N*N).reshape((N,N)) + np.arange(N*N).reshape((N,N)).transpose(1,0))
@unittest.skipUnless(RANGEIFY, "only correct in rangeify")
def test_post_permuted_assignment_alt(self):
a = Tensor.arange(N*N).reshape(N,N).contiguous().realize()
b = Tensor.arange(N*N).reshape(N,N).contiguous().realize()
@@ -339,18 +344,21 @@ class TestAssign(unittest.TestCase):
def test_permuted_assignment_correct(self):
a = Tensor.arange(4 * 4).reshape(4, 4).contiguous().realize()
b = Tensor.arange(4 * 4).reshape(4, 4).contiguous().realize()
a = a.permute(1, 0)
new_val = a + b
a.assign(new_val)
np.testing.assert_equal(a.numpy(), np.arange(4 * 4).reshape(4, 4).transpose(1, 0) + np.arange(4 * 4).reshape(4, 4))
# TODO: swizzler.py limitation, should NOT raise AssertionError from numpy.
with self.assert_permuted_assign():
a = a.permute(1, 0)
new_val = a + b
a.assign(new_val)
np.testing.assert_equal(a.numpy(), np.arange(4 * 4).reshape(4, 4).transpose(1, 0) + np.arange(4 * 4).reshape(4, 4))
def test_permuted_reduceop_child_dual_use(self):
a = Tensor.randn(32, 32, 32).realize()
b = Tensor.full((32, 32), 1.).contiguous().realize()
r = a.sum(axis=1)
b.assign(r + b.permute(1, 0))
b.realize()
np.testing.assert_allclose(b.numpy(), a.numpy().sum(axis=1)+np.ones((32, 32)).transpose(1, 0), atol=1e-6, rtol=1e-3)
with self.assert_permuted_assign():
r = a.sum(axis=1)
b.assign(r + b.permute(1, 0))
b.realize()
np.testing.assert_allclose(b.numpy(), a.numpy().sum(axis=1)+np.ones((32, 32)).transpose(1, 0), atol=1e-6, rtol=1e-3)
@unittest.skip("multi output not supported anymore")
def test_permuted_reduceop_multioutput_dual_use(self):
@@ -392,10 +400,11 @@ class TestAssign(unittest.TestCase):
def test_permuted_assignment_masked_view_not_contiguous(self):
a = Tensor.ones(4, 4).contiguous().realize()
b = a.shrink((None, (0, 2))).pad((None, (0, 2)), value=2).permute(1, 0)
a.assign(a + b)
a.realize()
self.assertListEqual(a.tolist(), [[2.,2.,2.,2.],[2.,2.,2.,2.],[3.,3.,3.,3.], [3.,3.,3.,3.]])
with self.assert_permuted_assign():
b = a.shrink((None, (0, 2))).pad((None, (0, 2)), value=2).permute(1, 0)
a.assign(a + b)
a.realize()
self.assertListEqual(a.tolist(), [[2.,2.,2.,2.],[2.,2.,2.,2.],[3.,3.,3.,3.], [3.,3.,3.,3.]])
# TODO: is there a way to sneak in a permute such that it returns the wrong answer?
+8 -9
View File
@@ -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
@@ -68,12 +69,9 @@ class TestBinaryOpsConstFolding(unittest.TestCase):
def test_tensor_one_mul(self):
_check_ast_count(0, Tensor.ones(4) * Tensor([1.0, 2, 3, 4]))
# TODO: these will be fixed with better folding
@unittest.expectedFailure
def test_bool_tensor_mul_bool(self):
_check_ast_count(0, Tensor([True, False]) * True)
_check_ast_count(0, Tensor([True, False]) * False)
@unittest.expectedFailure
def test_bool_mul_bool_tensor(self):
_check_ast_count(0, True * Tensor([True, False]))
_check_ast_count(0, False * Tensor([True, False]))
@@ -157,7 +155,8 @@ class TestMovedConstFolding(unittest.TestCase):
_check_ast_count(0, Tensor([1.0, 2, 3, 4]) + Tensor.zeros(6).shrink(((1, 5),)))
def test_add_padded_zero(self):
_check_ast_count(0, Tensor([1.0, 2, 3, 4]) + Tensor.zeros(2).pad(((1, 1),)))
# TODO: it's 1 now, this might be possible to fold
_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),)))
@@ -166,16 +165,16 @@ class TestMovedConstFolding(unittest.TestCase):
_check_ast_count(1, Tensor([1.0, 2, 3, 4]) * Tensor.ones(2).pad(((1, 1),)))
def test_cast_padded(self):
# NOTE: it's always 1 kernel when calling .numpy, limitation of _check_ast_count
# NOTE: RANGEIFY or not, it's always 1 kernel when calling .numpy, limitation of _check_ast_count
if is_dtype_supported(dtypes.int16):
_check_ast_count(1, Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int16))
_check_ast_count(1 if RANGEIFY else 0, Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int16))
np.testing.assert_equal(Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int16).numpy(), [0, 1, 1, 1, 1, 0])
if is_dtype_supported(dtypes.uint16):
_check_ast_count(1, Tensor.full(4, fill_value=-1).pad(((1, 1),)).cast(dtypes.uint16))
_check_ast_count(1 if RANGEIFY else 0, Tensor.full(4, fill_value=-1).pad(((1, 1),)).cast(dtypes.uint16))
np.testing.assert_equal(Tensor.full(4, fill_value=-1).pad(((1, 1),)).cast(dtypes.uint16).numpy(), [0, 65535, 65535, 65535, 65535, 0])
# folded
if is_dtype_supported(dtypes.int64):
_check_ast_count(1, Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int64))
_check_ast_count(1 if RANGEIFY else 0, Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int64))
np.testing.assert_equal(Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int64).numpy(), [0, 1, 1, 1, 1, 0])
class TestReduceOpsConstFolding(unittest.TestCase):
@@ -247,7 +246,7 @@ class TestReduceOpsConstFolding(unittest.TestCase):
t = Tensor.ones(16, dtype=dt).reshape(4, 4)
assert t.sum().dtype == t.contiguous().sum().dtype
@unittest.skipIf(not_support_multi_device() or True, "no multi, RANGEIFY doesn't support multi const folding")
@unittest.skipIf(not_support_multi_device() or RANGEIFY, "no multi, RANGEIFY doesn't support multi const folding")
class TestMultiConstFolding(unittest.TestCase):
def test_multi_const_folding_literal(self):
ds = tuple(f"{Device.DEFAULT}:{i}" for i in range(4))
+2 -2
View File
@@ -4,7 +4,7 @@ from tinygrad import Device, dtypes, Tensor, Context
from tinygrad.device import LRUAllocator, is_dtype_supported
from tinygrad.dtype import ImageDType
from tinygrad.engine.realize import lower_schedule
from tinygrad.helpers import prod, unwrap
from tinygrad.helpers import prod, unwrap, RANGEIFY
from test.helpers import REAL_DEV
IMAGE_SUPPORTED_DEVICES = ("QCOM", "CL")
@@ -139,7 +139,7 @@ class TestImageDType(unittest.TestCase):
# NOTE: the w1 grad must realize to a seperate kernel
assert w1.grad.uop.is_realized, f"never realized {w1.grad}"
self.assertEqual(w1.grad.uop.base.buffer.dtype, dtypes.float32)
self.assertEqual(len(sched), 9)
self.assertEqual(len(sched), 9 if RANGEIFY else 10)
@unittest.skipUnless(REAL_DEV in IMAGE_SUPPORTED_DEVICES, "Images not supported")
class TestImageRealization(unittest.TestCase):
+1 -1
View File
@@ -838,7 +838,7 @@ class TestJitRandom(unittest.TestCase):
tst = {0:[], 1:[]}
for r in [0,1]:
Tensor.manual_seed(1337)
with Context(JIT=r):
with Context(RANGEIFY=r):
_ = Tensor.randint(4, high=3)
# this second one makes the behavior different
_ = Tensor.randint(4, high=3)
+64 -6
View File
@@ -6,10 +6,13 @@ from tinygrad.codegen.opt import Opt, OptOps
from tinygrad.codegen.gpudims import get_grouped_dims
from tinygrad.uop.ops import UOp, Ops, GroupOp
from tinygrad.device import Device, Buffer, is_dtype_supported
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, flatten, dedup, TC_SELECT, TC_OPT
from tinygrad.helpers import Context, flatten, dedup, TC_SELECT, TC_OPT, RANGEIFY
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):
@@ -36,6 +39,24 @@ class TestLinearizer(unittest.TestCase):
np.testing.assert_equal(a.numpy(), ta)
np.testing.assert_equal(b.numpy(), tb)
def test_multioutput(self):
dtype, st = dtypes.int, ShapeTracker.from_shape((8,))
g0, g1, g2, g3 = [UOp(Ops.DEFINE_GLOBAL, dtype.ptr(), arg=i) for i in range(4)]
a = UOp(Ops.LOAD, dtype, src=(g2.view(st),))
b = UOp(Ops.LOAD, dtype, src=(g3.view(st),))
out0 = UOp(Ops.STORE, dtypes.void, src=(g0.view(st), a + b))
out1 = UOp(Ops.STORE, dtypes.void, src=(g1.view(st), a * b))
sink = UOp(Ops.SINK, src=(out0, out1))
a_t = Tensor.full(st.shape, 2).contiguous().realize()
b_t = Tensor.full(st.shape, 3).contiguous().realize()
helper_linearizer_ast(sink, [a_t, b_t], wanna_output=[a_t.numpy()+b_t.numpy(), a_t.numpy()*b_t.numpy()])
uops = get_program(sink, opts=[]).uops
stores = [u for u in uops if u.op is Ops.STORE]
mutable_bufs = dedup(flatten([[x for x in u.src[0].toposort() if x.op is Ops.DEFINE_GLOBAL] for u in stores]))
assert len(mutable_bufs) == len(stores) == 2
self.assertSetEqual(set([u.arg for u in mutable_bufs]), set([0,1]))
def _test_no_nested_ranges(self, lins, skip=None):
for l in lins:
range_in_acc = flatten([[x for x in u.src if x.op is Ops.RANGE] for u in l.uops if u.op is Ops.DEFINE_REG])
@@ -212,7 +233,7 @@ class TestLinearizer(unittest.TestCase):
assert u.src[1].op in GroupOp.ALU
assert begin_range < uops.index(u) < end_range
# children of STORE are placed after ENDRANGE
if any(x.op is Ops.STORE and x.src[1].op in GroupOp.ALU for x in u.src) and u.op is not Ops.ENDRANGE:
if any(x.op is Ops.STORE and x.src[1].op in GroupOp.ALU for x in u.src):
assert end_range < uops.index(u)
def test_grouped_dims(self):
@@ -314,7 +335,7 @@ class TestLinearizer(unittest.TestCase):
a.realize()
np.testing.assert_equal(a.flatten().numpy(), [1.,1.,1.,1.,2.,2.,2.,2.,1.,1.,1.,1.,1.,1.,1.,1.])
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX indexes differently. might be ok?")
@unittest.skipIf(RANGEIFY and isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX indexes differently. might be ok?")
def test_where_fold(self):
a = Tensor.ones(4, 4).contiguous().realize()
b = a.shrink(((1, 2), None)).pad(((1, 2), None))
@@ -397,7 +418,7 @@ class TestLinearizer(unittest.TestCase):
# # check the children's vins
# TODO: src ALU are not the same, should it?
# assert barrier.src == tuple(local_stores)
assert len([u for u in uops if u.op is Ops.IF and u.src[1] == barrier]) == 1
assert len([u for u in uops if u.op is Ops.IF and u.src[-1] == barrier]) == 1
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
@@ -417,8 +438,45 @@ class TestLinearizer(unittest.TestCase):
# the global store doesn't change
assert stores[1].src[1].dtype == dtypes.float
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
def test_skip_unmatching_upcasts(self):
Tensor.manual_seed(0)
c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(9600), arg=0, src=())
c1 = c0.view(ShapeTracker(views=(View(shape=(240, 40, 1, 1), strides=(40, 1, 0, 0), offset=0, mask=None, contiguous=True),)))
c2 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(9600), arg=1, src=())
c3 = c2.view(ShapeTracker(views=(View(shape=(240, 40, 1, 1), strides=(1, 240, 0, 0), offset=0, mask=None, contiguous=False),)))
c4 = c3.load()
c5 = c1.store(c4)
ast = c5.sink()
opt = [Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.LOCAL, axis=0, arg=16),
Opt(op=OptOps.LOCAL, axis=1, arg=2), Opt(op=OptOps.UPCAST, axis=3, arg=2)]
helper_linearizer_ast(ast, [Tensor.randn(240*40).realize()], opts=[opt])
out = [u for u in get_program(ast, opts=opt).uops if u.op is Ops.STORE][0]
assert out.src[1].op is Ops.VECTORIZE and out.src[1].dtype == dtypes.float.vec(4)
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
def test_skip_unmatching_upcasts_with_gep(self):
Tensor.manual_seed(0)
c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(256), arg=0, src=())
c1 = c0.view(ShapeTracker(views=(View(shape=(8, 32, 1, 1), strides=(32, 1, 0, 0), offset=0, mask=None, contiguous=True),)))
c2 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(256), arg=1, src=())
c3 = c2.view(ShapeTracker(views=(View(shape=(8, 32, 1, 1), strides=(1, 8, 0, 0), offset=0, mask=None, contiguous=False),)))
c4 = c3.load()
c5 = c1.store(c4)
ast = c5.sink()
opt = [Opt(op=OptOps.LOCAL, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=2, arg=2), Opt(op=OptOps.LOCAL, axis=1, arg=8),
Opt(op=OptOps.UPCAST, axis=1, arg=0), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.LOCAL, axis=0, arg=8),
Opt(op=OptOps.UPCAST, axis=1, arg=0), Opt(op=OptOps.UPCAST, axis=0, arg=2)]
helper_linearizer_ast(ast, [Tensor.randn(8*32).realize()], opts=[opt])
out = [u for u in get_program(ast).uops if u.op is Ops.STORE][0]
assert out.src[1].op is Ops.VECTORIZE and out.src[1].dtype.count != 1
# *** helpers ***
def push_views(ast): return apply_rewrites(ast, rewrites_for_views)
def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]:
if isinstance(r, Tensor): r = [r]
s = Tensor.schedule(*r)
@@ -427,12 +485,12 @@ def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]:
# 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)]
return s[-1].ast, bufs
return push_views(s[-1].ast), bufs
def helper_linearizer_ast(ast:UOp, inputs:list[Tensor], *args, **kwargs):
assert isinstance(ast, UOp), "ast must be UOp"
inbufs = [x.uop.base.buffer for x in inputs]
outbufs = [Buffer(inbufs[-1].device if inbufs else Device.DEFAULT, out.size, out.src[1].dtype).allocate() for out in ast.src]
outbufs = [Buffer(inbufs[-1].device if inbufs else Device.DEFAULT, out.st_arg.size, out.src[1].dtype).allocate() for out in ast.src]
_helper_linearizer_opt_ast(ast, outbufs+inbufs, *args, **kwargs)
def helper_linearizer_opt(r:Tensor|list[Tensor], *args, **kwargs):
+1 -4
View File
@@ -28,12 +28,10 @@ class TestLinearizerFailure(unittest.TestCase):
c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(dtypes.index, 6000))+c6)+((c7*UOp.const(dtypes.index, 16))+c8)).alu(Ops.CMPLT, UOp.const(dtypes.index, 59999)).where(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(dtypes.int, -1))).where(UOp.const(dtypes.uchar, 0), c10).reduce(c6, arg=Ops.ADD)
c12 = c0.index((((c1*UOp.const(dtypes.index, 7840))+(c2*UOp.const(dtypes.index, 10)))+c3).valid(UOp.const(dtypes.bool, True))).store(c11, c1, c2, c3)
ast = c12.sink(arg=KernelInfo(name='test', axis_types=(), dont_use_locals=False, applied_opts=(Opt(op=OptOps.GROUP, axis=1, arg=16),), opts_to_apply=None))
print(ast)
#_ = get_program(ast, Device["METAL"].renderer)
_ = get_program(ast, Device["METAL"].renderer)
class TestLinearizerDumb(unittest.TestCase):
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "need local")
@unittest.skip("Ops.VALID no longer exists")
def test_max_simplify_and_cancel(self):
c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1000), arg=0, src=())
c1 = c0.view(ShapeTracker(views=(View(shape=(1000, 1), strides=(1, 0), offset=0, mask=None, contiguous=True),)))
@@ -56,7 +54,6 @@ class TestLinearizerDumb(unittest.TestCase):
# this was a bug in embedding, someday we should fold this anyway
@unittest.skipUnless(is_dtype_supported(dtypes.half), f"half dtype not supported on {Device.DEFAULT}")
@unittest.skip("UOp.view is no longer supported")
def test_llama_embedding(self):
c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(4096), arg=0, src=())
c1 = c0.view(ShapeTracker(views=(View(shape=(4096, 1, 1), strides=(1, 0, 0), offset=0, mask=None, contiguous=True),)))
+5 -6
View File
@@ -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
@@ -390,7 +390,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.skip("TODO: pm_rangeify hangs")
@unittest.skipIf(RANGEIFY, "TODO: pm_rangeify hangs")
def test_data_parallel_resnet(self):
from extra.models.resnet import ResNet18
@@ -427,7 +427,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.skip("TODO: pm_rangeify hangs")
@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))
@@ -435,7 +435,7 @@ class TestMultiTensor(unittest.TestCase):
m = ResNet18()
self._test_model_train_step(m, fake_image, labels)
@unittest.skip("TODO: pm_rangeify hangs")
@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)
@@ -800,7 +800,7 @@ class TestMultiTensor(unittest.TestCase):
t = Tensor.rand(16, 16).shard(devices_2, axis=0)
np.testing.assert_allclose(t.numpy(), t.clone().numpy())
@unittest.skip("RANGEIFY doesn't support multi const folding")
@unittest.skipIf(RANGEIFY, "RANGEIFY doesn't support multi const folding")
def test_multi_const_folding(self):
with Context(TRACK_MATCH_STATS=0):
a = Tensor.arange(3).realize()
@@ -1137,7 +1137,6 @@ class TestMultiRamUsage(unittest.TestCase):
del _
self.assertUsed(0)
@unittest.skip("flaky")
def test_zeros_copy(self):
_ = Tensor.zeros(self.N, self.N).contiguous().to(devices_2).realize()
# NOTE: the first one on the DEFAULT device should be freed
+6 -5
View File
@@ -447,11 +447,11 @@ class TestNN(unittest.TestCase):
# TODO: fused with opts uses more ops
def test_embedding_one_kernel_fused(self):
with Context(NOOPT=0):
with Context(FUSE_ARANGE=1, NOOPT=0):
self.test_embedding_one_kernel(ops=612_000, kcount=2)
def test_embedding_one_kernel_fused_noopt(self):
with Context(NOOPT=1):
with Context(FUSE_ARANGE=1, NOOPT=1):
self.test_embedding_one_kernel(ops=0, kcount=2)
def test_embedding_shape(self):
@@ -465,9 +465,10 @@ class TestNN(unittest.TestCase):
def test_embedding_regression(self):
# used to fail bounds check
embedding = Embedding(100, 1024)
input_ids = Tensor.empty(16, 16, dtype=dtypes.int)
embedding(input_ids).realize()
with Context(FUSE_ARANGE=1):
embedding = Embedding(100, 1024)
input_ids = Tensor.empty(16, 16, dtype=dtypes.int)
embedding(input_ids).realize()
def test_load_state_dict(self):
layer = Conv2d(3, 5, kernel_size=3)
+4 -3
View File
@@ -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
@@ -3040,6 +3040,7 @@ class TestOps(unittest.TestCase):
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))
@@ -3163,8 +3164,8 @@ class TestOps(unittest.TestCase):
helper_test_op([(32,10)], lambda x: x.masked_fill((x>0.1).detach(), -math.inf))
helper_test_op([(32,10)], lambda x: x.masked_fill((x<0.1).detach(), -math.inf))
@unittest.skipIf((getenv("MOCKGPU") or Device.DEFAULT == "PYTHON"), "very slow on MOCKGPU because reduce does not fold")
@unittest.skipIf(Device.DEFAULT == "WEBGPU", "webgpu runtime issue")
@unittest.skipIf(RANGEIFY and (getenv("MOCKGPU") or Device.DEFAULT == "PYTHON"), "very slow on MOCKGPU because reduce does not fold")
@unittest.skipIf(RANGEIFY and Device.DEFAULT == "WEBGPU", "webgpu runtime issue")
def test_masked_select(self):
helper_test_op([(32, 10)], lambda x: x.masked_select(x>0.5), lambda x: x.masked_select(x>0.5), forward_only=True)
helper_test_op([(32, 10)], lambda x: x.masked_select(torch.tensor(True)), lambda x: x.masked_select(Tensor(True)), forward_only=True)
+1 -1
View File
@@ -90,7 +90,7 @@ class TestOptim(unittest.TestCase):
def test_muon(self): self._test_muon(1, {'lr': 0.001}, 1e-6, 0)
def test_muon_high_lr(self): self._test_muon(1, {'lr': 10}, 1e-6, 3e-4)
def test_muon_wd(self): self._test_muon(1, {'lr': 0.001, 'weight_decay': 0.01}, 1e-6, 0)
def test_muon_high_lr_wd(self): self._test_muon(1, {'lr': 10, 'weight_decay': 0.01}, 1e-6, 5e-4)
def test_muon_high_lr_wd(self): self._test_muon(1, {'lr': 10, 'weight_decay': 0.01}, 1e-6, 3e-4)
# NOTE: momentum set to 0.95 by default, nesterov set to True by default
def test_multistep_muon_momentum_wd(self): self._test_muon(10, {'lr': 0.001, 'weight_decay': 0.01}, 1e-5, 0)
-75
View File
@@ -1,75 +0,0 @@
import unittest
from tinygrad import Tensor, UOp, GlobalCounters, Context
class TestOuterworld(unittest.TestCase):
def test_range_plus_1(self):
t = Tensor.arange(100).reshape(10,10).realize()
# passthrough ranges
a = UOp.range(10, -1)
sel = t[a] + 1
assert sel.shape == (10,)
cpy = sel.reshape(1, 10).expand(a, 10).contiguous().realize()
self.assertTrue((t+1==cpy).all().item())
def test_flip_range(self):
t = Tensor.rand(10, 10).realize()
# passthrough ranges
a = UOp.range(10, -1)
sel = t[9-a]
cpy = sel.reshape(1, 10).expand(a, 10).contiguous().realize()
self.assertTrue((t.flip(0)==cpy).all().item())
def test_vmap(self):
def f(x): return x.sum(axis=0)*2
x = Tensor.ones(3, 10, 2).contiguous()
# vmap across axis 0
a = UOp.range(3, -1)
out = f(x[a])
out = out.reshape(1, 2).expand(a, 2).contiguous()
# 3x2 grid of 20
out.realize()
self.assertTrue((out==20).all().item())
@unittest.skip("opts don't work")
def test_triple_gemm(self):
x = Tensor.rand(1, 16).realize()
W = Tensor.rand(3, 16, 16).realize()
manual = (x @ W[0] @ W[1] @ W[2]).contiguous().realize()
a = UOp.range(3, -1)
x = x.assign(x @ W[a])
out = x.contiguous(a)[-1].contiguous().realize()
self.assertTrue((manual==out).all().item())
def test_setitem_pyrange(self):
with Context(DEBUG=0):
t = Tensor.rand(10).realize()
o = Tensor.empty(10)
GlobalCounters.reset()
for i in range(10):
o[i] = t[i]
o.realize()
self.assertTrue((t==o).all().item())
@unittest.skip("TODO: fix this")
def test_setitem(self):
with Context(DEBUG=0):
t = Tensor.rand(10).realize()
o = Tensor.empty(10)
GlobalCounters.reset()
i = UOp.range(10, -1)
o[i] = t[i]
o.contiguous(i).realize()
self.assertTrue((t==o).all().item())
if __name__ == '__main__':
unittest.main()
+67 -47
View File
@@ -1,8 +1,9 @@
import unittest
from tinygrad import Tensor, nn
from tinygrad.helpers import Context, GlobalCounters
from tinygrad.uop.ops import graph_rewrite, PatternMatcher, UPat, Ops
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')
@@ -54,6 +55,7 @@ class TestRangeifyOpt(unittest.TestCase):
A = Tensor.empty(8,8,8,8).permute(1,0,3,2).flatten()
A.sum().realize()
@unittest.skipIf(RANGEIFY<1, "tests only for RANGEIFY")
class TestRangeify(unittest.TestCase):
def test_groupnorm(self):
# ranges 1 and 3 are merging
@@ -70,7 +72,6 @@ class TestRangeify(unittest.TestCase):
ret = A.sum(axis=2).contiguous(arg=(1,)).sum(axis=1)
ret.realize()
@unittest.skip("RANGEIFY=0 does nothing")
def test_double_gemm_real(self):
def go():
with Context(DEBUG=0):
@@ -200,7 +201,6 @@ class TestRangeify(unittest.TestCase):
out = blk._feed_forward(x)
out.realize()
@unittest.skip("RANGEIFY=0 does nothing")
def test_flash_attention(self):
BS, HEADS, SEQLEN, EMB = 4, 2, 16, 8
@@ -229,56 +229,76 @@ class TestRangeify(unittest.TestCase):
# contiguous + reduce can support ranges?
@unittest.skip("pm_rangeify no longer exists. test this in a different way")
class TestRangeifyPM(unittest.TestCase):
def setUp(self): self.base = Tensor.empty(10*10).reshape(10, 10).contiguous()
def assert_same(self, a, b):
def run_pm_rangeify(t:Tensor):
from tinygrad.schedule.rangeify import pm_rangeify, RangeifyContext
sink = t.uop.sink()
pm_realize = PatternMatcher([(UPat(Ops.CONTIGUOUS, name="x"), lambda x: x.replace(op=Ops.REALIZE))])
sink = graph_rewrite(sink, pm_realize)
return graph_rewrite(sink, pm_rangeify, ctx=RangeifyContext())
self.assertIs(run_pm_rangeify(a.contiguous()), run_pm_rangeify(b.contiguous()))
@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):
t = Tensor.rand(10, 10).realize()
def test_nothing_match(self):
a = self.base.pad(((0,0),(0,1)))
b = self.base.pad(((0,0),(0,1)))
self.assert_same(a, b)
# passthrough ranges
a = UOp.range(10, -1)
sel = t[a]
cpy = sel.contiguous(a).realize()
def test_reshape_match(self):
a = self.base
b = self.base.reshape(100).reshape(10, 10)
self.assert_same(a, b)
self.assertTrue((t==cpy).all().item())
def test_permute_reshape_match(self):
a = self.base
b = self.base.permute(1,0).reshape(100).reshape(10, 10).permute(1,0)
self.assert_same(a, b)
def test_flip_range(self):
t = Tensor.rand(10, 10).realize()
def test_padded_permute_match(self):
a = self.base.pad(((0,0),(0,1)))
b = self.base.permute(1,0).pad(((0,1),(0,0))).permute(1,0)
self.assert_same(a, b)
# passthrough ranges
a = UOp.range(10, -1)
sel = t[9-a]
cpy = sel.contiguous(a).realize()
@unittest.expectedFailure
def test_padded_reshape_match(self):
a = self.base.pad(((0,0),(0,1)))
b = self.base.reshape(100).reshape(10, 10).pad(((0,0),(0,1)))
self.assert_same(a, b)
self.assertTrue((t.flip(0)==cpy).all().item())
@unittest.expectedFailure
def test_padded_permute_reshape_match(self):
a = self.base.pad(((0,0),(0,1)))
b = self.base.permute(1,0).reshape(100).reshape(10, 10).pad(((0,1),(0,0))).permute(1,0)
self.assert_same(a, b)
def test_vmap(self):
def f(x): return x.sum(axis=0)*2
# why is this failing?
@unittest.expectedFailure
def test_cross_pad_match(self):
a = self.base.pad(((0,0),(0,1))).pad(((0,1),(0,0)))
b = self.base.pad(((0,1),(0,0))).pad(((0,0),(0,1)))
self.assert_same(a, b)
x = Tensor.ones(3, 10, 2).contiguous()
# vmap across axis 0
a = UOp.range(3, -1)
out = f(x[a])
out = out.contiguous(a)
# 3x2 grid of 20
out.realize()
print(out.numpy())
@unittest.skip("opts don't work")
def test_triple_gemm(self):
x = Tensor.rand(1, 16).realize()
W = Tensor.rand(3, 16, 16).realize()
manual = (x @ W[0] @ W[1] @ W[2]).contiguous().realize()
a = UOp.range(3, -1)
x = x.assign(x @ W[a])
out = x.contiguous(a)[-1].contiguous().realize()
self.assertTrue((manual==out).all().item())
def test_setitem_pyrange(self):
with Context(DEBUG=0):
t = Tensor.rand(10).realize()
o = Tensor.empty(10)
GlobalCounters.reset()
for i in range(10):
o[i] = t[i]
o.realize()
self.assertTrue((t==o).all().item())
@unittest.skip("TODO: fix this")
def test_setitem(self):
with Context(DEBUG=0):
t = Tensor.rand(10).realize()
o = Tensor.empty(10)
GlobalCounters.reset()
i = UOp.range(10, -1)
o[i] = t[i]
o.contiguous(i).realize()
self.assertTrue((t==o).all().item())
class TestRangeifyEdgeCase(unittest.TestCase):
def test_matmul_relu_cat(self):
+181 -65
View File
@@ -12,11 +12,13 @@ from tinygrad import nn, dtypes, Device, Tensor, Variable
from tinygrad.device import is_dtype_supported
from tinygrad.dtype import DType, ImageDType
from tinygrad.shape.shapetracker import ShapeTracker
from tinygrad.uop.ops import UOp, Ops, GroupOp, UPat
from tinygrad.helpers import CI, DEBUG, SPLIT_REDUCEOP, GlobalCounters, Context, getenv, all_same, temp
from tinygrad.schedule.rangeify import get_rangeify_map, Kernel
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, 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
from test.helpers import expect_rangeify_fails, expect_nonrangeify_fails
class KernelCountException(Exception): pass
def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Tensor]|None=None, filter_sink=True):
@@ -27,12 +29,12 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te
else:
assert isinstance(t, UOp), f"can't schedule {t}"
sink = UOp.sink(t) if t.op is not Ops.SINK else t
becomes_map = get_rangeify_map(sink)
becomes_map = get_kernelize_map(sink)
sched, _ = create_schedule_with_vars(sink.substitute(becomes_map))
# 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:
return sched # allow different kernel count, TODO: fix the asserts
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):
@@ -44,7 +46,7 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te
def _realize_weights(m):
for p in nn.state.get_parameters(m): p.realize()
def _test_conv2d(allowed:int, dtype:DType=dtypes.float):
def _test_conv2d(allowed:int, dtype:DType=dtypes.float, **kwargs):
old_default_float, dtypes.default_float = dtypes.default_float, dtype
dtypes.default_float = dtype
Tensor.manual_seed(0)
@@ -53,7 +55,7 @@ def _test_conv2d(allowed:int, dtype:DType=dtypes.float):
w = Tensor.uniform(16, CIN, 3, 3, requires_grad=True).realize()
ret = Tensor.conv2d(img, w).relu().mean().backward()
dtypes.default_float = old_default_float
s = Tensor.schedule(ret, img.grad, w.grad)
with Context(**kwargs): s = Tensor.schedule(ret, img.grad, w.grad)
run_schedule(s.copy())
cnt = len([si for si in s if si.ast.op is Ops.SINK])
assert cnt == allowed, f"expected {allowed} kernels, got {cnt}"
@@ -66,6 +68,9 @@ def _test_conv2d(allowed:int, dtype:DType=dtypes.float):
np.testing.assert_allclose(img.grad.numpy(), ref_img.grad.detach().numpy(), atol=1e-6 if dtype == dtypes.float else 1e-2)
np.testing.assert_allclose(w.grad.numpy(), ref_w.grad.detach().numpy(), atol=1e-6 if dtype == dtypes.float else 1e-2)
@track_rewrites(name=True)
def schedule_graph_rewrite(big_sink:UOp): return get_kernelize_map(big_sink)[big_sink]
class TestSchedule(unittest.TestCase):
def test_arange_avgpool2d(self, kcount=1):
x = Tensor.arange(25).reshape(1,1,5,5).cast(dtypes.float32)
@@ -78,33 +83,37 @@ class TestSchedule(unittest.TestCase):
np.testing.assert_allclose(t.numpy(), torch_out)
def test_arange_avgpool2d_fused_noopt(self):
with Context(NOOPT=1): self.test_arange_avgpool2d(kcount=1)
with Context(FUSE_ARANGE=1, NOOPT=1): self.test_arange_avgpool2d(kcount=1)
# linearizer error
@unittest.skip("recursion error no longer raised")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "needs supports_float4 to fail")
def test_arange_avgpool2d_fused(self):
with self.assertRaises(RecursionError):
with Context(NOOPT=0): self.test_arange_avgpool2d(kcount=1)
with Context(FUSE_ARANGE=1, NOOPT=0): self.test_arange_avgpool2d(kcount=1)
# when we're fusing a reduce, all ReduceOps must have the same N in the dimensions
# all permutes, reshapes, expands and shrinks push through the reduce
def test_arange_sum(self):
a = Tensor.arange(6).reshape(3, 2).sum(axis=1)
run_schedule(check_schedule(a, 1))
with Context(FUSE_ARANGE=1):
run_schedule(check_schedule(a, 1))
self.assertListEqual(a.tolist(), [1, 5, 9])
def test_arange_sum_alt(self):
a = (Tensor.arange(5).reshape(1,5).expand(6,5)*Tensor(2)).reshape(1,6,5).sum(axis=2)
run_schedule(check_schedule(a, 1))
with Context(FUSE_ARANGE=1):
run_schedule(check_schedule(a, 1))
np.testing.assert_equal(a.numpy(), 20)
def test_permute_arange(self):
a = Tensor.arange(6).reshape(6, 1, 1).permute(2, 0, 1).sum(axis=1)
run_schedule(check_schedule(a, 1))
with Context(FUSE_ARANGE=1):
run_schedule(check_schedule(a, 1))
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")
@@ -112,11 +121,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, 2)
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")
@@ -129,7 +139,8 @@ class TestSchedule(unittest.TestCase):
def test_indexing_scalars_simple(self):
X = Tensor.randn(2, 2).realize()
xt = X[Tensor(1)][Tensor(0)]
run_schedule(check_schedule(xt, 2))
with Context(FUSE_ARANGE=1):
run_schedule(check_schedule(xt, 2))
np.testing.assert_equal(xt.numpy(), X.numpy()[1][0])
@unittest.skipIf(CI and Device.DEFAULT == "NV", "crashes on NV CI")
@@ -149,7 +160,8 @@ class TestSchedule(unittest.TestCase):
assume(a<x and b<y)
X = Tensor.randn(x, y).realize()
xt = X[Tensor(a)][Tensor(b)]
run_schedule(check_schedule(xt, 2))
with Context(FUSE_ARANGE=1):
run_schedule(check_schedule(xt, 2))
np.testing.assert_equal(xt.numpy(), X.numpy()[a][b])
def test_push_pads_elementwise(self):
@@ -342,9 +354,9 @@ class TestSchedule(unittest.TestCase):
r1 = (x - r0).sum(axis=0).div(2)
out0 = r0 + y
out1 = r1 + y
schedule = check_schedule([out0, out1], 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}]
self.assertEqual(len(reduceops), 2) # why is RANGEIFY different?
assert len(reduceops) == (3 if RANGEIFY else 2)
def test_div_collapse_buffer(self):
a = Tensor.full((4,), 4.0).contiguous().realize()
@@ -387,6 +399,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()
@@ -394,7 +407,7 @@ class TestSchedule(unittest.TestCase):
sched = check_schedule([a, b], 2, filter_sink=False)
run_schedule(sched)
# a and b are assigned to the same device Buffer
self.assertIsNot(a.uop.base.realized, b.uop.base.realized)
self.assertIsNot(a.uop.realized, b.uop.realized)
# EMPTY is assigned to a unique device Buffer
@@ -469,14 +482,15 @@ class TestSchedule(unittest.TestCase):
check_schedule(opt.schedule_step(), cnt)
def test_fold_batchnorm_backward(self):
with Tensor.train():
x = Tensor.empty((2, 16, 8, 8)).contiguous()
bn = nn.BatchNorm2d(16)
bn.weight.requires_grad = bn.bias.requires_grad = x.requires_grad = True
fw = bn(x).contiguous_backward().relu().contiguous()
fw.sum().backward()
# TODO: this is too many
check_schedule([x.grad, bn.weight.grad, bn.bias.grad, fw], 10)
with Context(FUSE_CONV_BW=1):
with Tensor.train():
x = Tensor.empty((2, 16, 8, 8)).contiguous()
bn = nn.BatchNorm2d(16)
bn.weight.requires_grad = bn.bias.requires_grad = x.requires_grad = True
fw = bn(x).contiguous_backward().relu().contiguous()
fw.sum().backward()
# TODO: this is too many
check_schedule([x.grad, bn.weight.grad, bn.bias.grad, fw], 10)
def test_fold_conv_relu(self):
c1 = nn.Conv2d(3,16,3)
@@ -699,6 +713,9 @@ class TestSchedule(unittest.TestCase):
prev_a = (a+1).contiguous()
a.assign(Tensor([2]))
a.kernelize(prev_a)
# RANGEIFY doesn't apply the post diamond graph, it's fine since we can always apply the fixup on each kernelize call
if not RANGEIFY:
assert prev_a.uop in a.uop.src, "contiguous usage must run before assign"
self.assertEqual((prev_a+a*3).item(), 1+2*3)
def test_kernelize_sym(self):
@@ -708,13 +725,13 @@ class TestSchedule(unittest.TestCase):
check_schedule(b, 0)
self.assertEqual(b.item(), 1)
# TODO: this requires supporting multiple stores in the AST
@unittest.expectedFailure
@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
c = Tensor.arange(4).realize().uop
kernel = UOp(Ops.KERNEL, src=(a.base, b.base, c.base), arg=Kernel(UOp.sink(c.r(Ops.ADD, (0,))+1, c.r(Ops.ADD, (0,))*2)))
kernel = UOp(Ops.KERNEL, src=(a, b, c.base), arg=Kernel(UOp.sink(c.r(Ops.ADD, (0,))+1, c.r(Ops.ADD, (0,))*2)))
assert all(s.op is Ops.BUFFER for s in kernel.src), f"views are not allowed here {kernel}"
run_schedule(check_schedule(UOp.sink(a.assign(kernel), b.assign(kernel)), 1))
self.assertEqual(a.buffer.numpy(), [7])
self.assertEqual(b.buffer.numpy(), [12])
@@ -915,7 +932,7 @@ class TestSchedule(unittest.TestCase):
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)
@@ -926,7 +943,7 @@ class TestSchedule(unittest.TestCase):
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)
@@ -1018,7 +1035,7 @@ class TestSchedule(unittest.TestCase):
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)
@@ -1181,16 +1198,15 @@ class TestSchedule(unittest.TestCase):
expected = (x_exp:=np.exp(x.numpy()-x.numpy().max(-1, keepdims=True)))/x_exp.sum(-1, keepdims=True)
np.testing.assert_allclose(out.numpy(), expected, atol=1e-4, rtol=1e-4)
# TODO: rangeify stores the output in float32
@unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
@unittest.expectedFailure
@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), 2)
self.assertEqual(len(sched), 2 if RANGEIFY else 3)
self.assertEqual(sched[0].bufs[0].dtype, dtypes.half)
# input float, softmax in float
@@ -1317,10 +1333,10 @@ class TestSchedule(unittest.TestCase):
opt = nn.optim.SGD(nn.state.get_parameters([c1, c2, c3, c4]))
opt.zero_grad()
c4(c3(c2(c1(img).relu()).relu()).relu()).relu().sum().backward()
check_schedule(opt.schedule_step(), 14)
with Context(FUSE_CONV_BW=1): check_schedule(opt.schedule_step(), 14)
@unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
@unittest.expectedFailure
@expect_rangeify_fails
def test_prefer_half_buffer(self):
x = Tensor.ones(4).contiguous().realize()
# y = Tensor.ones(4).contiguous().realize()
@@ -1472,7 +1488,7 @@ 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)
@@ -1561,7 +1577,8 @@ class TestSchedule(unittest.TestCase):
x = Tensor.empty(3,3,3,3, requires_grad=True)
y = x.pad((-1,2,2,-1), mode="replicate")
dx = y.sum().gradient(x)[0]
sched = check_schedule(dx, 3)
with Context(FUSE_ARANGE=1):
sched = check_schedule(dx, 3)
run_schedule(sched)
np.testing.assert_allclose(dx.numpy(), [[[[0.,3.,9.],[0,1.,3.],[0.,0.,0.]]]*3]*3)
@@ -1621,14 +1638,15 @@ class TestSchedule(unittest.TestCase):
out = x.argmax(1)
run_schedule(check_schedule(out, 2))
def test_conv2d(self): _test_conv2d(5 if SPLIT_REDUCEOP else 4)
def test_conv2d_fused(self): _test_conv2d(5 if SPLIT_REDUCEOP else 4)
def test_conv2d(self): _test_conv2d(5 if RANGEIFY else 7)
def test_conv2d_fused(self): _test_conv2d(5 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(5 if SPLIT_REDUCEOP else 4, dtype=dtypes.half)
def test_conv2d_half(self): _test_conv2d(5 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")
def test_conv2d_fused_half(self): _test_conv2d(5 if SPLIT_REDUCEOP else 4, dtype=dtypes.half)
@unittest.skipIf(not RANGEIFY, "passes on RANGEIFY")
def test_conv2d_fused_half(self): _test_conv2d(5, dtype=dtypes.half)
def test_schedule_mem_used(self):
base = GlobalCounters.mem_used
@@ -1687,14 +1705,15 @@ 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)
casted_view.realize()
self.assertEqual(casted_view.uop.base.realized.size, 8)
contig = casted_view.contiguous().realize()
self.assertEqual(contig.uop.base.realized.size, 8)
self.assertListEqual(contig.tolist(), [[0.0, 1.0, 2.0, 3.0], [0.0, 0.0, 0.0, 0.0]])
self.assertEqual(casted_view.uop.base.realized.size, 4)
realized_view = casted_view.contiguous().realize()
self.assertEqual(realized_view.uop.base.realized.size, 8)
self.assertListEqual(realized_view.tolist(), [[0.0, 1.0, 2.0, 3.0], [0.0, 0.0, 0.0, 0.0]])
# NOTE: we only reorder CAST if it's an EXPAND
def test_cast_after_shrink(self):
@@ -1716,6 +1735,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))
@@ -1859,7 +1879,8 @@ class TestSchedule(unittest.TestCase):
from extra.models.llama import precompute_freqs_cis
args = {"dim":32 if CI else 128, "end":2048 if CI else 8192, "theta":10000}
fused = precompute_freqs_cis(**args)
run_schedule(check_schedule(fused, 3))
with Context(FUSE_ARANGE=1):
run_schedule(check_schedule(fused, 3))
if getenv("CHECK", 1):
ref = precompute_freqs_cis(**args)
run_schedule(check_schedule(ref, 3))
@@ -1886,7 +1907,9 @@ class TestSchedule(unittest.TestCase):
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)
np.testing.assert_equal(tst.numpy(), a.numpy())
if RANGEIFY > 0:
# NOTE: this is a bug on non rangeify
np.testing.assert_equal(tst.numpy(), a.numpy())
def test_setitem_sched(self, mop=lambda x:x, expected_kcount=1):
a = Tensor.arange(16, device="CPU").reshape(4, 4).contiguous().realize()
@@ -1897,6 +1920,7 @@ class TestSchedule(unittest.TestCase):
run_schedule(sched)
self.assertListEqual(a.tolist(), expected)
self.assertEqual(kcount, expected_kcount)
@unittest.skipUnless(RANGEIFY>0, "this asserts on non rangeify")
def test_setitem_permuted_sched(self): self.test_setitem_sched(lambda x: x.T, 2)
def test_setitem_paddded_sched(self): self.test_setitem_sched(lambda x: x.shrink_to(4, 1).pad_to(4, 4), 1)
@@ -1935,16 +1959,26 @@ class TestSchedule(unittest.TestCase):
r = (X+Tensor.arange(16).reshape(4, 4)).sum()
out0 = r+2
out1 = r+3
run_schedule(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)
@unittest.skip("multi output isn't supported")
def test_multiview_arange_children(self):
X = Tensor.randn(2,3,4,4).numpy()
with Context(FUSE_ARANGE=1):
compare = Tensor(X).interpolate(size=(2, 2), mode="linear").numpy()
with Context(FUSE_ARANGE=0, TRACK_MATCH_STATS=0):
ref = Tensor(X).interpolate(size=(2, 2), mode="linear").numpy()
np.testing.assert_allclose(ref, compare, atol=1e-5, rtol=1e-6)
def test_recursive_swizzle(self):
a = Tensor([1,2,3,4]).realize()
for _ in range(24): a = a + a
new_uop = a.reshape(4,1).realize().uop
assert new_uop.base.op is Ops.BUFFER
self.assertEqual(new_uop.st, ShapeTracker.from_shape((4,)).reshape((4, 1)))
self.assertEqual(swizzle_cnt(new_uop), 0)
@unittest.skipIf(CI and Device.DEFAULT == "NV", "crashes on NV CI")
def test_limit_bufs_with_var(self):
@@ -1970,6 +2004,9 @@ class TestSchedule(unittest.TestCase):
sched = z.schedule()
self.assertEqual(len(sched), kcount+1)
def swizzle_cnt(u:UOp) -> int:
return len([x for x in u.toposort() if x.op is Ops.VIEW and len(x.src) != 0 and x.src[0].op not in {Ops.BUFFER, Ops.DEFINE_GLOBAL, Ops.ASSIGN}])
class TestSwizzle(unittest.TestCase):
def test_swizzle_simple(self):
Tensor.manual_seed(0)
@@ -2080,12 +2117,14 @@ 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()
b = a.pad((None, (0, 10)))[:, 10:]
assert b.shape == (10, 10)
sched = check_schedule(b.contiguous(), 1)
self.assertEqual(sched[-1].ast.full_shape, (10, 10))
run_schedule(sched)
np.testing.assert_equal(b.numpy(), 0)
@@ -2105,6 +2144,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()
@@ -2136,6 +2176,56 @@ class TestView(unittest.TestCase):
run_schedule(s)
self.assertEqual(other_child.tolist(), [2, 3, 4])
def tensor_rewrite(t) -> UOp: return graph_rewrite(t.uop.base, merge_views+symbolic_simple)
class TestSimplifier(unittest.TestCase):
def test_sink_childless_const(self):
x = Tensor(0)
check_schedule(x, 0)
def test_sink_childless_const_alt_expanded(self):
x = Tensor.zeros(4, 4).contiguous()
check_schedule(x, 1)
def test_all_const_uops(self):
a = Tensor(4)*Tensor(2)
sink = tensor_rewrite(a)
assert UPat.cvar().match(sink, {})
def test_masked_const_elementwise(self):
a = Tensor.eye(10)@Tensor.eye(10)
sink = tensor_rewrite(a)
assert UPat(Ops.REDUCE_AXIS, src=(UPat.cvar().view()*UPat.cvar().view(),)).match(sink, {})
def test_elementwise_ops(self):
a = Tensor.empty(4, 4, dtype=dtypes.int)
sink = tensor_rewrite(a*0)
assert UPat(Ops.CONST, arg=0).match(sink, {})
self.assertIs(tensor_rewrite(a*1).base, a.uop.base)
self.assertIs(tensor_rewrite(a+0).base, a.uop.base)
def test_cast_folding(self):
a = Tensor(1.0).cast(dtypes.int)
sink = tensor_rewrite(a)
assert UPat.cvar(dtype=dtypes.int).match(sink, {})
def test_const_folding_mul(self):
a = Tensor([1])
sink = tensor_rewrite(a*0)
assert UPat(Ops.CONST, arg=0).match(sink, {}), f"expected {sink} to collapse to a const 0"
assert sink.shape == a.shape
def test_const_folding_ne(self):
a = Tensor([1])
sink = tensor_rewrite(a != a)
assert UPat(Ops.CONST, arg=False).match(sink, {}), f"expected {sink} to collapse to a const False"
assert sink.shape == a.shape
def test_const_folding_lt(self):
a = Tensor([1])
sink = tensor_rewrite(a < a)
assert UPat(Ops.CONST, arg=False).match(sink, {}), f"expected {sink} to collapse to a const False"
assert sink.shape == a.shape
@unittest.skipIf(Device.DEFAULT == "CPU", "tests copy from another device to cpu")
class TestCopyFolding(unittest.TestCase):
def test_const_copy_is_free(self):
@@ -2173,11 +2263,17 @@ class TestCopyFolding(unittest.TestCase):
a = Tensor.empty(4).uop
b = a.copy_to_device(a.device)
check_schedule(b, 0, filter_sink=False)
b = schedule_graph_rewrite(b)
# NOTE: Tensor.empty(4) always creates a VIEW(BUFFER) with ShapeTracker((4,)), we simplify this to jsut a BUFFER
# in the scheduler because buffer already has shape (4,)
self.assertIs(b, a.base)
def test_copy_to_same_device_alt(self):
a = Tensor.empty(4, 4).uop
b = a.copy_to_device(a.device)
check_schedule(b, 0, filter_sink=False)
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()
@@ -2225,6 +2321,7 @@ class TestCopyFolding(unittest.TestCase):
b.realize()
self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])
@expect_nonrangeify_fails
def test_permute_on_disk_contiguous(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')}")
@@ -2239,6 +2336,8 @@ class TestCopyFolding(unittest.TestCase):
self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])
# NOTE: disk permute must come after COPY
# TODO: this is wrong because of the permute
@expect_nonrangeify_fails
def test_permute_after_shrink_on_disk(self):
with open(temp('dt_arange_5_permute'), "wb") as f: f.write(Tensor.arange(5).realize().uop.base.buffer.as_buffer())
a = Tensor.empty(5, dtype=dtypes.int32, device=f"disk:{temp('dt_arange_5_permute')}")
@@ -2275,8 +2374,9 @@ class TestBufferUOp(unittest.TestCase):
def test_buffer_view_not_allowed(self):
permuted_view = Tensor.empty(1, 2, 3).permute(0, 2, 1)
with self.assertRaisesRegex(AssertionError, "can only be RESHAPE"):
permuted_view.uop.buffer # cannot access Buffer of a non contiguous VIEW
merged = graph_rewrite(permuted_view.uop, merge_views)
with self.assertRaisesRegex(AssertionError, "VIEW only works here if it's contiguous"):
merged.buffer # cannot access Buffer of a non contiguous VIEW
def test_buffer_only_after_realize(self):
a = Tensor([1])+Tensor([2])
@@ -2368,22 +2468,25 @@ 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,)
add = a+b
check_schedule(add, 1)
# BUFFER already has a shape (4,), this tensor just becomes a contiguous BUFFER
assert UPat(Ops.BUFFER).match(add.uop.base, {})
assert UPat(Ops.BUFFER).match(add.uop, {})
# sometimes we prefer to perform an op before movement ops, in this case we should stack the mops on top of the new buffer
# NOTE: this expand is not reordered because there's before it to fuse
@expect_rangeify_fails
def test_reorder_expand(self):
a = Tensor.empty(4, 1)
b = a.expand(4, 4).reciprocal()
check_schedule(b, 1)
self.assertEqual(b.uop.base.buffer.size, 4)
self.assertEqual(b.uop.shape, (4, 4))
self.assertEqual(b.uop.base.buffer.size, 16)
self.assertEqual(b.uop.st, ShapeTracker.from_shape((4, 4)))
def test_reorder_expand_alt(self):
x = Tensor.empty(4, 1)
@@ -2392,13 +2495,13 @@ class TestUOpBecome(unittest.TestCase):
z = (img*x) / y
check_schedule(z, 1)
# TODO: rangeify doesn't yet cleanup this kind of re-indexing
@unittest.expectedFailure
@expect_rangeify_fails
def test_become_existing_buffer(self):
a = Tensor.empty(4, 4)
b = a*1
assert UPat(Ops.MUL).match(b.uop, {}) # before scheduling it's a mul
check_schedule(b, 0)
assert UPat(Ops.VIEW, src=(UPat(Ops.BUFFER))).match(b.uop, {}) # scheduling merges all MovementOps into a single VIEW
self.assertIs(a.uop.base.buffer, b.uop.base.buffer)
def test_become_buf_with_mops(self):
@@ -2420,6 +2523,17 @@ 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)
b = add.shrink(((0, 1), (0, 0)))
check_schedule(b, 0)
assert UPat(Ops.CONST, arg=0).match(b.uop, {})
self.assertEqual(b.shape, (1, 0))
# the base is untouched.
assert UPat(Ops.ADD).match(add.uop, {})
def test_become_const_from_const(self):
const_add = Tensor(1)+Tensor(2)
assert UPat(Ops.ADD).match(const_add.uop, {})
@@ -2427,7 +2541,7 @@ class TestUOpBecome(unittest.TestCase):
assert UPat(Ops.CONST, arg=3).match(const_add.uop.base, {})
# tensors can become another realized tensor source
@unittest.expectedFailure
@expect_rangeify_fails
def test_become_existing_buf_simple(self):
a = Tensor.empty(4, 4)
b = a+0
@@ -2436,14 +2550,14 @@ class TestUOpBecome(unittest.TestCase):
self.assertIs(a.uop, b.uop)
# they can also chain other movement ops on top of the tensor source
@unittest.expectedFailure
@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)
@unittest.expectedFailure
@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
@@ -2451,7 +2565,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
@unittest.expectedFailure
@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
@@ -2459,7 +2573,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
@unittest.expectedFailure
@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
@@ -2471,14 +2585,16 @@ 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))
b.realize()
assert a.uop.is_realized
assert a.uop.buffer._base is None
assert b.uop.op_in_backward_slice_with_self(Ops.SHRINK)
assert b.uop.base is a.uop.base
# b is a subbuffer of a
assert b.uop.op is Ops.BUFFER_VIEW
assert b.uop.src[0] is a.uop
def test_setitem_offset(self):
a = Tensor.full((16,), 0.).contiguous().realize()
+9 -9
View File
@@ -2,7 +2,7 @@ import unittest
import numpy as np
from tinygrad import Tensor, GlobalCounters, Context, Device
from tinygrad.dtype import DTypeLike, dtypes
from tinygrad.helpers import DEBUG, get_single_element
from tinygrad.helpers import DEBUG, get_single_element, RANGEIFY
from tinygrad.engine.realize import lower_schedule_item
from tinygrad.device import is_dtype_supported
@@ -39,17 +39,17 @@ class TestFuse(unittest.TestCase):
np_multi = fxn(*args, **kwargs).numpy()
np.testing.assert_allclose(np_single, np_multi, atol=atol)
@unittest.skip("needs RANGEIFY>1")
@unittest.skipIf(0<RANGEIFY<2, "needs RANGEIFY>1")
def test_fuse_norm(self):
a = Tensor.rand(50,50).realize()
self._test_fuse(lambda a: a / a.mean(axis=1), a)
@unittest.skip("needs RANGEIFY>1")
@unittest.skipIf(0<RANGEIFY<2, "needs RANGEIFY>1")
def test_fuse_argmax(self):
a = Tensor.rand(50,50).realize()
self._test_fuse(lambda a: a.argmax(axis=-1), a)
@unittest.skip("needs RANGEIFY>1")
@unittest.skipIf(0<RANGEIFY<2, "needs RANGEIFY>1")
def test_fuse_softmax(self):
a = Tensor.rand(50,50).realize()
self._test_fuse(lambda a: a.softmax(axis=-1), a)
@@ -60,7 +60,7 @@ class TestFuse(unittest.TestCase):
self._test_fuse(lambda a,b: ((a@b).relu()+a).contiguous().softmax(axis=-1), a,b, allow_multiple=True)
@unittest.skipUnless(is_dtype_supported(dtypes.float16, Device.DEFAULT), f"no float16 on {Device.DEFAULT}")
@unittest.skip("needs RANGEIFY>1")
@unittest.skipIf(0<RANGEIFY<2, "needs RANGEIFY>1")
def test_fuse_softmax_dtype(self):
a = Tensor.rand(50,50).realize()
self._test_fuse(lambda a: a.softmax(axis=-1, dtype='half'), a, atol=3e-4)
@@ -68,7 +68,7 @@ class TestFuse(unittest.TestCase):
def test_fuse_arange_eye(self):
self._test_fuse(lambda: Tensor.arange(10).reshape(10,1).expand(10,10) == Tensor.arange(10).reshape(1,10).expand(10,10))
@unittest.skip("needs RANGEIFY>1")
@unittest.skipIf(0<RANGEIFY<2, "needs RANGEIFY>1")
def test_double_gemm(self):
N = 32
with Context(TRACK_MATCH_STATS=0, DEBUG=0):
@@ -91,7 +91,7 @@ class TestFuse(unittest.TestCase):
return (arange == idx).mul(vals).sum(-2, dtype=vals.dtype)
self._test_fuse(embedding, a, atol=1e-5)
@unittest.skip("needs RANGEIFY>1")
@unittest.skipIf(0<RANGEIFY<2, "needs RANGEIFY>1")
def test_attention_kernel_count(self):
wq = Tensor.empty(32, 32)
wk = Tensor.empty(32, 32)
@@ -104,7 +104,7 @@ class TestFuse(unittest.TestCase):
s = attn.schedule()
self.assertEqual(len(s), 4) # 3 matmul and 1 attention
@unittest.skip("needs RANGEIFY>1")
@unittest.skipIf(0<RANGEIFY<2, "needs RANGEIFY>1")
def test_flash_attention(self):
BS = 4
HEADS = 2
@@ -172,7 +172,7 @@ class TestSoftmaxFusion(unittest.TestCase):
np.testing.assert_allclose(sout.numpy(), out.numpy(), atol=3e-7)
@unittest.skip("needs RANGEIFY>1")
@unittest.skipIf(0<RANGEIFY<2, "needs RANGEIFY>1")
def test_auto_softmax(self):
print("*** softmax ***")
with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
+1 -1
View File
@@ -40,7 +40,7 @@ class TestStunning(unittest.TestCase):
Y_train = Y_train.one_hot(10)
X_samp, Y_samp = X_train[samples], Y_train[samples]
vi = Variable('i', 0, samples.shape[0]-1)
with Context(SPLIT_REDUCEOP=0):
with Context(FUSE_ARANGE=1, SPLIT_REDUCEOP=0):
with Tensor.train():
losses = []
for i in range(samples.shape[0]):
+3 -2
View File
@@ -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):
@@ -26,7 +27,7 @@ class TestSymbolicJit(unittest.TestCase):
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)
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()
@@ -79,7 +80,7 @@ class TestSymbolicJit(unittest.TestCase):
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, 4)
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()
+17 -11
View File
@@ -4,7 +4,7 @@ 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
@@ -595,21 +595,21 @@ class TestMoveTensor(unittest.TestCase):
np.testing.assert_equal(x.grad.numpy(), [[2,2,2],[0,0,0],[-2,-2,-2]])
class TestZeroShapeTensor(unittest.TestCase):
def test_shape_is_expanded(self):
def test_shape_stride(self):
t = Tensor.empty(3, 2, 0)
assert t.shape == (3, 2, 0)
# numpy has stride 0, 0, 0; torch has stride 2, 1, 1
assert t.uop.st.is_expanded() == (True, True, True)
assert t.uop.st.real_strides() == (0, 0, 0)
t = Tensor.empty(3, 0, 2)
assert t.shape == (3, 0, 2)
# numpy has stride 0, 0, 0; torch has stride 2, 2, 1
assert t.uop.st.is_expanded() == (True, True, True)
assert t.uop.st.real_strides() == (0, 0, 0)
t = Tensor.empty(0, 0, 0)
assert t.shape == (0, 0, 0)
# numpy has stride 0, 0, 0; torch has stride 1, 1, 1
assert t.uop.st.is_expanded() == (True, True, True)
assert t.uop.st.real_strides() == (0, 0, 0)
def test_rand(self):
t = Tensor.rand(3, 2, 0)
@@ -860,7 +860,6 @@ class TestTensorMetadata(unittest.TestCase):
self.assertEqual(len(si.metadata), 3)
self.assertEqual(set(m.name for m in si.metadata), {"relu", "sigmoid", "__mul__"})
@unittest.skip("not accurate")
def test_complex_backward(self):
x = Tensor.rand(3, requires_grad=True).realize()
y = Tensor.rand(3, requires_grad=True).realize()
@@ -872,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), 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")
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):
-2
View File
@@ -4,7 +4,6 @@ import unittest
from tinygrad import Tensor, Device, dtypes
from tinygrad.engine.realize import run_schedule
from tinygrad.uop.ops import Ops, UOp, UPat
from tinygrad.helpers import SPLIT_REDUCEOP
class TestTensorUOp(unittest.TestCase):
def test_fromcpu_shape_tracker(self):
@@ -95,7 +94,6 @@ class TestTensorUOp(unittest.TestCase):
self.assertEqual(out.tolist(), Tensor.zeros(4, 8).tolist())
reduce_kernel = UPat(Ops.SINK, src=(UPat(Ops.STORE, allow_any_len=True, src=(UPat(), UPat((Ops.REDUCE_AXIS, Ops.REDUCE))))))
@unittest.skipUnless(SPLIT_REDUCEOP, "only for SPLIT_REDUCEOP")
class TestReduceOp(unittest.TestCase):
def test_no_split_reduce_kernel(self):
a = Tensor.rand(4, 4).realize()
+13
View File
@@ -665,6 +665,19 @@ class TestUOpGraph(unittest.TestCase):
bad_gate = UOp.const(dtypes.int, 1)
with self.assertRaises(AssertionError): to_uops_list([UOp(Ops.STORE, dtypes.void, (glbl0, idx, UOp.const(dtypes.int, 42), bad_gate))])
def test_switched_range_order(self):
glbl = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
cf = UOp.const(dtypes.float, 0.0)
r1 = UOp.range(2, 0)
r2 = UOp.range(2, 1)
alu = UOp(Ops.MUL, dtypes.int, (r2, r1))
store = UOp(Ops.STORE, dtypes.void, (glbl.index(alu), cf))
uops = to_uops_list([store])
ranges = [x for x in uops if x.op is Ops.RANGE]
endranges = [x for x in uops if x.op is Ops.ENDRANGE]
# ranges are closed in the right order
self.assertEqual(endranges[-1].src[0], ranges[0])
@track_rewrites()
def expander_rewrite(sink): return graph_rewrite(sink, sym + expander)
+33 -11
View File
@@ -1,6 +1,8 @@
from typing import Optional, Any
import unittest, math
import numpy as np
from tinygrad.shape.shapetracker import ShapeTracker
from tinygrad.shape.view import View # noqa F401
from tinygrad.tensor import Tensor, _to_np_dtype
from tinygrad.helpers import CI, DEBUG, getenv, Timing
from tinygrad.dtype import dtypes, DType, AddrSpace
@@ -490,6 +492,15 @@ class TestUOpMethod(unittest.TestCase):
self.assertIs(x.replace(arg=None).arg, None)
with self.assertRaises(AssertionError): x.replace(field="a")
def test_device(self):
x = UOp(Ops.VIEW, dtypes.int, (UOp.new_buffer(Device.DEFAULT, 1, dtypes.int), UOp.const(dtypes.int, 1)), ShapeTracker.from_shape(()))
self.assertEqual(x.device, Device.DEFAULT)
# NOTE: CONST doesn't have device
buffer, const = x.src
self.assertEqual(buffer.device, Device.DEFAULT)
self.assertEqual(const._device, None)
with self.assertRaises(AssertionError): const.device
class TestUOpStr(unittest.TestCase):
def test_uop_str(self):
a = UOp(Ops.CONST, dtypes.float, (), 2.0) + UOp(Ops.CONST, dtypes.float, (), 3.0)
@@ -533,18 +544,29 @@ class TestUopsObject(unittest.TestCase):
with Timing("create 10k uops:"): ret = [UOp(Ops.CONST, dtypes.int, arg=10000000+i) for i in range(10000)]
assert len(ret) == 10000
def test_nested(self):
a = UOp.new_buffer(Device.DEFAULT, 1, dtypes.char)
for _ in range(10_000): a = a+a
self.assertEqual(a.device, Device.DEFAULT)
class TestUOpChildren(unittest.TestCase):
def test_children_exist(self):
a = UOp.variable("weird_name_234", 0, 10)
b = a*a
self.assertEqual(len(a.children), 1)
self.assertIs(list(a.children)[0](), b)
class TestUOpRender(unittest.TestCase):
def test_render_vectorize_same(self):
u = UOp(Ops.VECTORIZE, src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0)))
self.assertEqual(u.render(), "{0, ...}")
def test_render_vectorize_different(self):
u = UOp(Ops.VECTORIZE, src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1), UOp.const(dtypes.int, 2)))
self.assertEqual(u.render(), "{0,1,2}")
def test_children_cleaned_up(self):
a = UOp.variable("weird_name_235", 0, 10)
b = a*a
self.assertEqual(len(a.children), 1)
del b
self.assertEqual(len(a.children), 0)
def test_children_cleaned_up_two(self):
a = UOp.variable("weird_name_236", 0, 10)
b = a*a
c = a*2
self.assertEqual(len(a.children), 2)
del b
self.assertEqual(len(a.children), 1)
del c
self.assertEqual(len(a.children), 0)
if __name__ == '__main__':
unittest.main(verbosity=2)
+6 -4
View File
@@ -1,6 +1,6 @@
import unittest
from tinygrad import Tensor
from tinygrad.helpers import getenv, GlobalCounters, EMULATE
from tinygrad.helpers import getenv, GlobalCounters, EMULATE, RANGEIFY
from tinygrad.engine.realize import lower_schedule_item, ProgramSpec, get_program
from tinygrad.renderer import Estimates
from tinygrad.codegen import full_rewrite
@@ -51,8 +51,11 @@ class TestMemoryCount(unittest.TestCase):
a = Tensor.empty(1024, 1, dtype=dtypes.uint8).expand(1024, 1024)
b = Tensor.empty(1024, 1, dtype=dtypes.uint8).expand(1024, 1024)
_, mem = get_stats(a+b)
# rangeify is smart!
self.assertEqual(mem, 1024 + 2*1024) # 2 lil reads + 1 lil write
if RANGEIFY:
# rangeify is smart!
self.assertEqual(mem, 1024 + 2*1024) # 2 lil reads + 1 lil write
else:
self.assertEqual(mem, 1024*1024 + 2*1024) # 2 lil reads + 1 write
def test_self_add(self):
a = Tensor.empty(1024, 1024, dtype=dtypes.uint8)
@@ -208,7 +211,6 @@ class TestStatsOptimized(unittest.TestCase):
self.check_gemm(p)
self.assertEqual(p.estimates.lds, 2*N*N*N*4//4 + 4*N*N)
@unittest.skip("wrong estimate")
def test_gemm_group(self):
try:
p = get_program(self.ast_gemm, opts=[Opt(OptOps.GROUP, 0, 4)])
+8 -5
View File
@@ -1,10 +1,12 @@
import unittest
from tinygrad import Tensor, dtypes, TinyJit, UOp
from tinygrad.helpers import RANGEIFY
from tinygrad.apps.llm import apply_rope
#from tinygrad.engine.realize import run_schedule
# TODO: test_scheduler, but just in uint
class TestAttention(unittest.TestCase):
@unittest.skipIf(RANGEIFY > 0, "not half on rangeify")
def test_half_qkv_buffers(self):
BS, seqlen, dim = 10, 4, 100
q = Tensor.ones(BS, seqlen, dim, dtype=dtypes.half).contiguous().realize()
@@ -12,11 +14,12 @@ class TestAttention(unittest.TestCase):
v = Tensor.ones(BS, seqlen, dim, dtype=dtypes.half).contiguous().realize()
attn = q.scaled_dot_product_attention(k, v)
sched = attn.schedule()
# attention has 4 kernels now
self.assertEqual(len(sched), 4)
# softmax_inputs = sched[1:4]
# for i,si in enumerate(softmax_inputs):
# assert all(b.dtype == dtypes.half for b in si.bufs), f"non half {si.bufs=} in kernel {i}"
#run_schedule(sched[:])
# attention has 5 kernels now
self.assertEqual(len(sched), 4 if RANGEIFY else 5)
softmax_inputs = sched[1:4]
for i,si in enumerate(softmax_inputs):
assert all(b.dtype == dtypes.half for b in si.bufs), f"non half {si.bufs=} in kernel {i}"
def test_apply_rope(self):
x = Tensor.randn(1, 2, 4, 8, dtype=dtypes.float32)
@@ -1,7 +1,7 @@
import unittest, random
from tinygrad.dtype import dtypes
from tinygrad.uop.ops import print_uops, UOp, Ops
from tinygrad.codegen.late.control_flow import schedule
from tinygrad.codegen.late.linearize import block_reorder
from tinygrad.renderer.cstyle import OpenCLRenderer
def is_toposorted(lst:list[UOp]):
@@ -11,7 +11,7 @@ def is_toposorted(lst:list[UOp]):
seen.add(u)
return True
class TestCFGSchedule(unittest.TestCase):
class TestBlockReorder(unittest.TestCase):
def _test_randomize(self, golden:list[UOp]):
# test random order is always same
for _ in range(50):
@@ -24,7 +24,7 @@ class TestCFGSchedule(unittest.TestCase):
if p not in topolst: topolst.append(p)
assert is_toposorted(topolst)
for x,y in zip(golden, this_order:=schedule(topolst)):
for x,y in zip(golden, this_order:=block_reorder(topolst)):
if x is not y:
print_uops(golden)
print_uops(this_order)
@@ -54,7 +54,7 @@ class TestCFGSchedule(unittest.TestCase):
sink = c.store(sum(loads)).sink()
# determine golden order
golden = schedule(list(sink.toposort()))
golden = block_reorder(list(sink.toposort()))
# render for test
print(self._test_render(golden))
@@ -62,8 +62,8 @@ class TestCFGSchedule(unittest.TestCase):
# assert the loads are in this order
self.assertListEqual([g.src[0].src[1].render() for g in golden if g.op is Ops.LOAD],
['(gidx0*27)', '((gidx0*27)+1)', '((gidx0*27)+2)', '((gidx0*27)+3)',
'(gidx1*4)', '((gidx1*4)+1)', '((gidx1*4)+2)', '((gidx1*4)+3)'])
['(gidx1*4)', '((gidx1*4)+1)', '((gidx1*4)+2)', '((gidx1*4)+3)',
'(gidx0*27)', '((gidx0*27)+1)', '((gidx0*27)+2)', '((gidx0*27)+3)'])
# assert math is after loads
first_math = [i for i,g in enumerate(golden) if g.op is Ops.ADD and g.dtype == dtypes.float][0]
-3
View File
@@ -29,11 +29,8 @@ class TestKeccak(unittest.TestCase):
out_shape = Tensor.randint(*s[i:], high=255, dtype=dtypes.uint8).keccak().shape
self.assertTupleEqual(s[i:-1], out_shape[:-1])
@unittest.skipUnless(Device.DEFAULT=="METAL", "slow")
def test_sha3_224(self): self._test_preset("sha3_224", [143, 144])
@unittest.skipUnless(Device.DEFAULT=="METAL", "slow")
def test_sha3_256(self): self._test_preset("sha3_256", [135, 136])
@unittest.skipUnless(Device.DEFAULT=="METAL", "slow")
def test_shake_128(self): self._test_preset("shake_128", [167, 168], lambda d: hashlib.shake_128(d).digest(16))
def _test_preset(self, name: str, special_sizes: list[int], hasher: Callable[[bytes], bytes] | None = None):
+15 -1
View File
@@ -3,7 +3,7 @@ from tinygrad import Variable
from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray, CI, mv_address
from tinygrad.helpers import merge_dicts, strip_parens, prod, round_up, fetch, fully_flatten, from_mv, to_mv, polyN, time_to_str, cdiv, cmod, getbits
from tinygrad.tensor import Tensor, get_shape
from tinygrad.shape.view import get_contraction
from tinygrad.shape.view import get_contraction, get_contraction_with_reduce
import numpy as np
VARIABLE = ContextVar("VARIABLE", 0)
@@ -219,6 +219,20 @@ class TestMemoryview(unittest.TestCase):
print(f"from_mv vs mv_address: {fmv_us:8.3f} µs vs {mva_us:8.3f} µs")
class TestGetContraction(unittest.TestCase):
def test_contraction_with_reduce(self):
r = get_contraction((16, 1, 1, 1), (16, 1, 1))
self.assertEqual(r, [[0], [], [1, 2, 3]])
r = get_contraction_with_reduce((16, 1, 1, 1), (16, 1, 1), (1,))
self.assertEqual(r, [[0], [1, 2], [3]])
r = get_contraction((16, 1, 1, 1, 1), (16, 1, 1, 1))
self.assertEqual(r, [[0], [], [], [1, 2, 3, 4]])
r = get_contraction_with_reduce((16, 1, 1, 1, 1), (16, 1, 1, 1), (1,))
self.assertEqual(r, [[0], [1, 2], [3], [4]])
r = get_contraction_with_reduce((2, 512, 1, 1), (2, 1, 512), (1,))
self.assertIsNone(r)
def test_contraction(self):
r = get_contraction((1,2,3,4), (2,3,4))
self.assertEqual(r, [[0, 1], [2], [3]])
+3 -2
View File
@@ -971,8 +971,9 @@ class TestIndexing(unittest.TestCase):
numpy_testing_assert_equal_helper((2, 0, 4), z.shape)
# this isn't technically necessary, but matches NumPy stride calculations.
# NOTE: this is empty and shouldn't have strides
numpy_testing_assert_equal_helper((True, True, True), z.uop.st.is_expanded())
self.assertTrue(z.uop.st.contiguous)
#numpy_testing_assert_equal_helper((60, 20, 5), z.uop.st.real_strides())
# NOTE tinygrad's int slicing implementation makes this not contiguous
# self.assertTrue(z.uop.st.contiguous)
@unittest.skip("bool indexing not supported")
def test_index_getitem_copy_bools_slices(self):
+3 -2
View File
@@ -1,6 +1,7 @@
import unittest
from tinygrad import Tensor
from tinygrad.uop import Ops
from tinygrad.helpers import RANGEIFY
class TestKernelize(unittest.TestCase):
def test_add_reshaped(self):
@@ -17,8 +18,8 @@ class TestKernelize(unittest.TestCase):
a1 = a.sum(axis=1)
a0 = a1.sum(axis=0)
a0.kernelize()
self.assertEqual(len([s for s in a0.uop.toposort() if s.op is Ops.KERNEL]), 2)
self.assertIs(a1.uop.base.op, Ops.REDUCE_AXIS)
self.assertEqual(len([s for s in a0.uop.toposort() if s.op is Ops.KERNEL]), 2 if RANGEIFY else 3)
self.assertIs(a1.uop.base.op, Ops.REDUCE_AXIS if RANGEIFY else Ops.ASSIGN)
# input Tensor and user contiguous kernelize
self.assertIs(a0.uop.base.op, Ops.ASSIGN)
self.assertIs(a.uop.base.op, Ops.ASSIGN)
-1
View File
@@ -12,7 +12,6 @@ def reconstruction_helper(A:list[Tensor],B:Tensor, tolerance=1e-5):
np.testing.assert_allclose(reconstructed_tensor.numpy(),B.numpy(),atol=tolerance,rtol=tolerance)
class TestLinAlg(unittest.TestCase):
@unittest.skip("TODO: reenable this")
def test_svd_general(self):
sizes = [(2,2),(5,3),(3,5),(3,4,4),(2,2,2,2,3)]
for size in sizes:
+14 -29
View File
@@ -28,22 +28,6 @@ class TestRewriteMap(unittest.TestCase):
self.assertIs(sub_map[a+b], e)
self.assertIs(sub_map[(a+b)*c], f)
def test_multistage_substitute(self):
a = UOp.variable('a', 0, 10)
b = UOp.variable('b', 0, 10)
c = UOp.variable('c', 0, 10)
d = UOp.variable('d', 0, 10)
sub1 = {a+b:c}
start = (a+b)*c
# stage 1: (a+b)*c -> c*c
sub_map1 = graph_rewrite_map(start, _substitute, sub1, bottom_up=True)
self.assertIs(sub_map1[(a+b)*c], c*c)
# stage 2: c*c -> d
sub2 = {c*c:d}
sub_map2 = graph_rewrite_map(sub_map1[start], _substitute, sub2, input_map=sub_map1, bottom_up=True)
# (a+b)*c -> c*c -> d
self.assertIs(sub_map2[(a+b)*c], d)
def test_add_zero(self):
# Build a small graph: add(0, add(const=0, const=5))
zero_node = UOp.const(dtypes.index, 0)
@@ -144,11 +128,11 @@ class TestRewriteMap(unittest.TestCase):
yz_sum_zero = yz_sum + zero_node -> rewrites to yz_sum
yz_neg = -yz_sum_zero -> -(y+z)
yz_dneg = -yz_neg -> y+z (double neg gone)
x_plus_yz = x_var + yz_dneg -> x + (y+z)
double_neg_x = -(-x_plus_yz) -> x + (y+z)
final_expr = double_neg_x * one_node -> x + (y+z)
x_plus_yz = x_var + yz_dneg -> (x+y)+z (add nodes get sorted)
double_neg_x = -(-x_plus_yz) -> (x+y)+z
final_expr = double_neg_x * one_node -> (x+y)+z
We expect the final result to be (x + (y+z)).
We expect the final result to be ((x+y)+z).
Each original node should map to the final node that replaces it,
which might be structurally equivalent but not the same reference.
"""
@@ -163,9 +147,9 @@ class TestRewriteMap(unittest.TestCase):
yz_sum_zero = yz_sum + zero_node # (y + z) + 0
yz_neg = -yz_sum_zero # -(y+z)
yz_dneg = -yz_neg # -(-(y+z)) -> (y+z)
x_plus_yz = x_var + yz_dneg # x + (y+z)
double_neg_x = -(-x_plus_yz) # neg(neg(x+(y+z))) -> x+(y+z)
final_expr = double_neg_x * one_node # (x+(y+z)) * 1 -> x+(y+z)
x_plus_yz = x_var + yz_dneg # x + (y+z) -> (x+y)+z
double_neg_x = -(-x_plus_yz) # neg(neg(x+(y+z))) -> (x+y)+z
final_expr = double_neg_x * one_node # ((x+y)+z) * 1 -> (x+y)+z
node_map = graph_rewrite_map(final_expr, symbolic)
@@ -182,14 +166,15 @@ class TestRewriteMap(unittest.TestCase):
# -(-(y+z)) => (y+z)
self.assertEqual(node_map[yz_dneg], yz_sum)
# x + (y+z) => might get recreated if yz_dneg was changed, so compare to x + yz_sum
self.assertEqual(node_map[x_plus_yz], x_var + yz_sum)
# x + (y+z) => (x+y)+z
expected_xyz = (x_var + y_var) + z_var
self.assertEqual(node_map[x_plus_yz], expected_xyz)
# -(-(x+(y+z))) => x + (y+z)
self.assertEqual(node_map[double_neg_x], x_var + yz_sum)
# -(-(x+(y+z))) => (x+y)+z
self.assertEqual(node_map[double_neg_x], expected_xyz)
# (x+(y+z)) * 1 => x+(y+z)
self.assertEqual(node_map[final_expr], x_var + yz_sum)
# ((x+y)+z) * 1 => (x+y)+z
self.assertEqual(node_map[final_expr], expected_xyz)
# Unchanged atomic nodes map to themselves
self.assertEqual(node_map[x_var], x_var)
+110
View File
@@ -0,0 +1,110 @@
import unittest
from dataclasses import dataclass, field
from tinygrad.uop.ops import PatternMatcher, UOp, graph_rewrite, Ops, UPat, GroupOp, RewriteNotReady
# we could insert CHILDREN node
@dataclass
class ChildrenContext:
children: dict[UOp, list[UOp]]|None = None
# this is a generic child labeller
def extract_children(ctx:ChildrenContext, x:UOp):
if ctx.children is not None: return
ctx.children = {k:list(v.keys()) for k,v in x.get_children_map().items() if len(v) > 1}
def mark_children(ctx:ChildrenContext, x:UOp):
new_srcs = [(UOp(Ops.CHILD, s.dtype, src=(s,), arg=(ctx.children[s].index(x), len(ctx.children[s]))) if s in ctx.children else s) for s in x.src]
return x.replace(src=tuple(new_srcs))
pm_children = PatternMatcher([
(UPat(Ops.SINK, name="x"), extract_children),
(UPat(GroupOp.All-{Ops.CHILD}, name="x"), mark_children),
])
@dataclass
class TestContext:
seen_children: dict[UOp, set[int]] = field(default_factory=dict)
ready_children: dict[UOp, set[int]] = field(default_factory=dict)
seen_consts:int = 0
saved_seen_consts:int = 0
exp2_visit_count:int = 0
# this is a generic pattern
def visit_child(ctx:ChildrenContext, x:UOp):
if x.src[0] not in ctx.seen_children:
ctx.seen_children[x.src[0]] = set()
ctx.ready_children[x.src[0]] = set()
ctx.seen_children[x.src[0]].add(x.arg[0])
if len(ctx.seen_children[x.src[0]]) != x.arg[1]:
print(f"visit CHILD {x.arg} bottom up -- not ready {ctx.seen_children[x.src[0]]}")
raise RewriteNotReady
print(f"visit CHILD {x.arg} bottom up -- READY {ctx.seen_children[x.src[0]]}")
ctx.ready_children[x.src[0]].add(x.arg[0])
pm_child_visitor = PatternMatcher([
(UPat(Ops.CHILD, name="x"), visit_child),
])
# this is for the test
def see_const(ctx:ChildrenContext, c:UOp): ctx.seen_consts += c.arg
def see_exp2(ctx:ChildrenContext): ctx.exp2_visit_count += 1
def save_seen_consts(ctx:ChildrenContext, x:UOp): ctx.saved_seen_consts = ctx.seen_consts
pm_consts = PatternMatcher([
(UPat(Ops.DEFINE_VAR, name="x"), save_seen_consts),
(UPat()+UPat.cvar("c"), see_const),
(UPat(Ops.EXP2), see_exp2),
])
class TestChildrenRewrite(unittest.TestCase):
def test_not_ready_double_simple(self):
global_a = UOp.variable("a", 0, 10).exp2()
inter = (global_a+global_a).exp2()
global_sink = (inter+inter).sink()
sink = graph_rewrite(global_sink, pm_children, ctx=ChildrenContext(), bottom_up=True)
ctx = TestContext()
graph_rewrite(sink, pm_consts, ctx=ctx, bottom_up=True)
self.assertEqual(ctx.exp2_visit_count, 2)
def test_not_ready_double(self):
global_a = UOp.variable("a", 0, 10).exp2()
inter = ((global_a+1000)+(global_a+100)).exp2()
global_sink = ((inter+10)+(inter+1)).sink()
sink = graph_rewrite(global_sink, pm_children, ctx=ChildrenContext(), bottom_up=True)
print("test_not_ready_double")
ctx = TestContext()
graph_rewrite(sink, pm_child_visitor+pm_consts, ctx=ctx, bottom_up=True)
self.assertEqual(ctx.exp2_visit_count, 2)
self.assertEqual(ctx.seen_consts, ctx.saved_seen_consts)
self.assertEqual(ctx.seen_consts, 1111)
def test_in_srcs_twice(self):
global_a = UOp.variable("a", 0, 10).exp2()
global_sink = (global_a+global_a).sink()
ctx = TestContext()
graph_rewrite(global_sink, pm_consts, ctx=ctx, bottom_up=True)
self.assertEqual(ctx.exp2_visit_count, 1)
def test_not_ready(self):
global_a = UOp.variable("a", 0, 10).exp2()
global_sink = ((global_a+2)+(global_a+3)).sink()
# without children and not ready, we don't see both adds before the DEFINE_VAR
ctx = TestContext()
graph_rewrite(global_sink, pm_consts, ctx=ctx, bottom_up=True)
self.assertNotEqual(ctx.seen_consts, ctx.saved_seen_consts)
self.assertEqual(ctx.exp2_visit_count, 1)
# with children and not ready we do
sink = graph_rewrite(global_sink, pm_children, ctx=ChildrenContext(), bottom_up=True)
ctx = TestContext()
graph_rewrite(sink, pm_child_visitor+pm_consts, ctx=ctx, bottom_up=True)
self.assertEqual(ctx.seen_consts, ctx.saved_seen_consts)
self.assertEqual(ctx.exp2_visit_count, 1)
self.assertSetEqual(list(ctx.ready_children.values())[0], {0,1})
if __name__ == '__main__':
unittest.main()
+63
View File
@@ -0,0 +1,63 @@
import unittest
from tinygrad import Tensor
from tinygrad.uop.ops import PatternMatcher, Ops, UPat, graph_rewrite, RewriteContext, UOp
from tinygrad.schedule.kernelize import kernelize_sym, merge_views
class TestRewriteTrackedChildren(unittest.TestCase):
@unittest.skip("track_children no longer supported")
def test_children_in_context(self):
def print_children(ctx:RewriteContext, sink:UOp):
view_w_child = sink.src[0].src[0].src[0]
assert view_w_child.op is Ops.VIEW
assert set([x.arg for x in ctx.children[view_w_child]]) == set((2,3))
ctx.update_children()
assert set([x.arg for x in ctx.children[view_w_child]]) == set((3,4))
# this is the 3
assert len(ctx.children[sink.src[0].src[1]]) == 1
assert next(iter(ctx.children[sink.src[0].src[1]])).op is Ops.ADD
# this is the 4
assert len(ctx.children[sink.src[0].src[0]]) == 1
assert next(iter(ctx.children[sink.src[0].src[0]])).op is Ops.ADD
rewrite = PatternMatcher([
(UPat(Ops.CONST, arg=2, name="x"), lambda x: x.replace(arg=4)),
(UPat(Ops.SINK, name="sink"), print_children)
])
a = Tensor(2)
b = Tensor(3)
c = a + b
sink = c.uop.sink()
sink = graph_rewrite(sink, rewrite, track_children=True)
def test_simple_child(self):
rewrite = PatternMatcher([
(UPat(Ops.CONST, arg=2, name="x"), lambda x: x.replace(arg=4)),
])
a = Tensor(2)
b = Tensor(3)
c = a + b
sink = c.uop
view_w_child = a.uop.src[0]
print([x().arg for x in view_w_child.children])
print([x.arg for x in sink.get_children_map()[view_w_child]])
self.assertSetEqual(set([x.arg for x in sink.get_children_map()[view_w_child]]), set((2,3)))
# children can either be added to or removed from the map with graph_rewrite
# added to is easy to detect, just hook the UOp constructor
# when are children removed?
# * if a rewrite rule returns a UOp, the matched node is removed from the graph
sink = graph_rewrite(sink, rewrite)
print([x().arg for x in view_w_child.children])
print([x.arg for x in sink.get_children_map()[view_w_child]])
self.assertSetEqual(set([x.arg for x in sink.get_children_map()[view_w_child]]), set((3,4)))
@unittest.skip("track_children no longer supported")
def test_child_after_parent_update(self):
def print_children(ctx, r):
ctx.update_children()
print(ctx.children[r])
extra = PatternMatcher([(UPat(Ops.REDUCE_AXIS, name="r"), print_children)])
a = Tensor.empty(3, 3)
r = (a+0).sum()
graph_rewrite(r.uop, merge_views+kernelize_sym+extra, track_children=True)
if __name__ == '__main__':
unittest.main()
+78 -9
View File
@@ -3,14 +3,14 @@ import unittest
import numpy as np
from tinygrad.dtype import dtypes, Invalid
from tinygrad.helpers import prod
from tinygrad.shape.shapetracker import ShapeTracker, View, views_to_valid_uop
from tinygrad.shape.shapetracker import ShapeTracker, View
from tinygrad import Variable
from tinygrad.uop.ops import UOp, Ops, graph_rewrite
from tinygrad.codegen.late.devectorizer import sym
from itertools import product
def shapetracker_getitem(st:ShapeTracker, val:int):
valid_idx = views_to_valid_uop(st.reshape((st.size,)).views, (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
@@ -95,20 +95,23 @@ class TestRealIssues(unittest.TestCase):
class TestRealDoesntSimplify(unittest.TestCase):
def tearDown(self):
st = self.st.real_strides()
print(st)
self.st = self.st.simplify()
assert len(self.st.views) != 1
assert None in st
def test_1(self):
self.st = ShapeTracker((
View.create((8, 3, 1, 2, 11, 1), (33, 11, 0, 0, 1, 0), 0, None),
View.create((8, 6, 11), (66, 11, 1), 0, None)))
self.assertEqual(self.st.is_expanded(), (False, False, False))
self.assertEqual(self.st.real_strides(), (33, None, 1))
def test_2(self):
self.st = ShapeTracker((
View.create((2, 2, 4, 3, 3), (72, 9, 18, -3, -1), 8, None),
View.create((4, 4, 3, 3), (36, 9, 3, 1), 0, None)))
self.assertEqual(self.st.is_expanded(), (False, False, False, False))
self.assertEqual(self.st.real_strides(), (None, 18, -3, -1))
class TestRealStrides(unittest.TestCase):
def test_1(self):
@@ -116,7 +119,7 @@ class TestRealStrides(unittest.TestCase):
View.create((2048,), (1,), 0, ((0, 512),)),
View.create((16, 32, 4), (128, 4, 1), 0, None),
))
self.assertEqual(st.is_expanded(), (False, False, False))
self.assertEqual(st.real_strides(), (None, 4, 1))
def test_2(self):
# test/test_ops.py::TestOps::test_simple_padding_conv1d
@@ -125,7 +128,7 @@ class TestRealStrides(unittest.TestCase):
View.create((6, 2, 78), (140, 70, 1), 0, ((0, 6), (0, 2), (0, 70))),
View.create((6, 2, 13, 6), (156, 78, 1, 13), 0, None),
))
self.assertEqual(st.is_expanded(), (False, False, False, False))
self.assertEqual(st.real_strides(), (90, 45, None, None))
def test_3(self):
# test/test_ops.py::TestOps::test_simple_cumsum
@@ -134,7 +137,7 @@ class TestRealStrides(unittest.TestCase):
View.create((4, 131327), (131072, 1), 0, ((0, 4), (0, 131072))),
View.create((4, 511, 257), (131327, 1, 511), 0, None),
))
self.assertEqual(st.is_expanded(), (False, False, False))
self.assertEqual(st.real_strides(), (256, None, None))
def test_4(self):
# test/test_nn.py::TestNN::test_conv_transpose1d
@@ -143,7 +146,7 @@ class TestRealStrides(unittest.TestCase):
View.create((1, 4, 1, 16, 8, 121), (0, 1792, 0, 112, 0, 1), -5, ((0, 1), (0, 4), (0, 1), (0, 16), (0, 8), (5, 116))),
View.create((4, 64, 115, 16, 7), (15488, 0, 1, 968, 122), 0, None),
))
self.assertEqual(st.is_expanded(), (False, True, False, False, False))
self.assertEqual(st.real_strides(), (896, 0, None, 56, None))
def test_5(self):
# test/test_ops.py::TestOps::test_conv2d
@@ -152,12 +155,15 @@ 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.is_expanded(), (False, False, False, True, False))
self.assertEqual(st.real_strides(), (132, 12, None, None, None))
class TestRealSimplifies(unittest.TestCase):
def tearDown(self):
st = self.st.real_strides()
self.st = self.st.simplify()
assert len(self.st.views) == 1
print(self.st.views[-1].strides, st)
self.assertEqual(self.st.views[-1].strides, st)
def test_1(self):
self.st = ShapeTracker((
@@ -169,6 +175,12 @@ class TestRealSimplifies(unittest.TestCase):
View.create((8, 3, 3, 11, 2, 28), (924, 308, 0, 28, 0, 1), 0, None),
View.create((8, 1, 6, 10, 28, 3, 2, 1), (5544, 0, 0, 56, 1, 1848, 672, 0), 0, None)))
class TestViewMinify(unittest.TestCase):
def test_minifies(self):
assert len(View.create((10,10)).minify().shape) == 1
assert len(View.create((10,10)).permute((1,0)).minify().shape) == 2
assert len(View.create((10,10,10,10)).permute((1,0,2,3)).minify().shape) == 3
class TestIndexExpressions2d(unittest.TestCase):
def setUp(self):
shapes = [(30, 5), (15, 10), (15, 1), (5, 10), (5, 1)] # Make sure dim0 is a multiple of 5, one of the tests divides this dimension by 5
@@ -745,6 +757,63 @@ class TestShapeTracker(unittest.TestCase):
self.test_expand()
self.test_permute()
class TestShapeTrackerSize(unittest.TestCase):
def test_simple_size(self):
st = ShapeTracker.from_shape((100, 100))
self.assertEqual(st.real_size(), 100*100)
def test_0_in_shape_size(self):
st = ShapeTracker.from_shape((0, 100))
self.assertEqual(st.real_size(), 0)
st = ShapeTracker.from_shape((100, 0))
self.assertEqual(st.real_size(), 0)
def test_expand_size(self):
st = ShapeTracker.from_shape((100, 100))
st = st.reshape((100, 100, 1))
st = st.expand((100, 100, 100))
self.assertEqual(st.real_size(), 100*100)
def test_expand_size_flatten(self):
st = ShapeTracker.from_shape((100, 100))
st = st.reshape((100, 100, 1))
st = st.expand((100, 100, 100))
st = st.reshape((100*100*100,))
self.assertEqual(st.real_size(), 100*100)
def test_shrink_size_axis_0(self):
st = ShapeTracker.from_shape((100, 100))
st = st.shrink(((0, 50), (0, 100)))
self.assertEqual(st.real_size(), 50*100)
def test_shrink_size_axis_0_variable(self):
st = ShapeTracker.from_shape((100, 100))
st = st.shrink(((0, Variable("a", 0, 50)), (0, 100)))
self.assertEqual(st.real_size(), 50*100)
def test_shrink_size_axis_1(self):
st = ShapeTracker.from_shape((100, 100))
st = st.shrink(((0, 100), (0, 50)))
self.assertEqual(st.real_size(), 9950) # careful here
def test_size_variable(self):
st = ShapeTracker(views=(View(shape=(1, 1, 1, (Variable('start_pos', 0, 8192)+1), 1, 8, 4, 128), strides=(0, 0, 0, 1024, 0, 128, 0, 1),
offset=0, mask=None, contiguous=False), View(shape=(1, 32, 1, (Variable('start_pos', 0, 8192)+1), 128),
strides=(0, 128, 0, 4096, 1), offset=0, mask=None, contiguous=False)))
self.assertEqual(st.real_size(), 8389632)
def test_pad_size_simple(self):
st = ShapeTracker.from_shape((10,)).pad(((2,4),))
self.assertEqual(st.real_size(), 10)
def test_pad_size_multiview(self):
st = ShapeTracker.from_shape((10,10)).pad(((2,4), (3,1))).reshape((16*14,))
self.assertEqual(st.real_size(), 100)
def test_flip_size(self):
st = ShapeTracker.from_shape((10,10)).pad(((2,4), (3,1))).flip((True, True))
self.assertEqual(st.real_size(), 100)
class TestVariableShrink(unittest.TestCase):
def test_shrink(self):
st = ShapeTracker.from_shape((10,))
+57
View File
@@ -103,5 +103,62 @@ class TestShapeTrackerAddVariable(unittest.TestCase):
ret_2 = ShapeTracker((vm1,)) + ShapeTracker((vm2,)).reshape((var_i, var_j, 1))
assert ret == ret_2
class TestShapeTrackerInvert(unittest.TestCase):
def test_invert_reshape(self):
a = ShapeTracker.from_shape((10, 10))
x = a.reshape((5, 20))
ap = ShapeTracker.from_shape(x.shape) + x.invert(a.shape)
assert ap == a, f"{ap} != {a}"
def test_invert_permute(self):
a = ShapeTracker.from_shape((5, 20))
x = a.permute((1,0))
ap = x + x.invert(a.shape)
assert ap == a, f"{ap} != {a}"
def test_invert_permute_3(self):
a = ShapeTracker.from_shape((8, 4, 5))
x = a.permute((1,2,0))
ap = x + x.invert(a.shape)
assert ap == a, f"{ap} != {a}"
def test_invert_real1(self):
a = ShapeTracker.from_shape((3, 6, 10))
x = a.reshape( (3, 3, 2, 10) )
x = x.permute( (2, 1, 3, 0) )
ap = x + x.invert(a.shape)
assert ap == a, f"{ap} != {a}"
def test_cant_invert_expand(self):
a = ShapeTracker.from_shape((10, 1))
x = a.expand((10,10))
assert x.invert(a.shape) is None
def test_cant_invert_shrink(self):
a = ShapeTracker.from_shape((10, 10))
x = a.shrink(((0,10),(2,8)))
assert x.invert(a.shape) is None
def test_can_invert_flip(self):
a = ShapeTracker.from_shape((20, 10))
x = a.flip((True,False))
ap = x + x.invert(a.shape)
assert st_equal(ap, a)
def test_can_invert_flip_permute(self):
a = ShapeTracker.from_shape((20, 10))
x = a.permute((1,0))
x = x.flip((True,False))
ap = x + x.invert(a.shape)
assert st_equal(ap, a)
def test_invert_failure(self):
a = ShapeTracker.from_shape((2, 5))
x = a.pad( ((2, 0), (0, 0)) )
x = x.reshape( (2, 2, 5) )
x = x.reshape( (4, 5) )
ap = x + x.invert(a.shape)
assert st_equal(ap, a)
if __name__ == '__main__':
unittest.main()
+2 -2
View File
@@ -1,11 +1,11 @@
import unittest
import multiprocessing.shared_memory as shared_memory
from tinygrad.helpers import CI, WIN
from tinygrad.helpers import CI, WIN, RANGEIFY
from tinygrad.tensor import Tensor, Device
import numpy as np
class TestRawShmBuffer(unittest.TestCase):
@unittest.skipIf(WIN and CI, "only fails on CI windows instance")
@unittest.skipIf(WIN and CI and RANGEIFY, "only fails with RANGEIFY on CI windows instance")
def test_e2e(self):
t = Tensor.randn(2, 2, 2).realize()
+9 -10
View File
@@ -60,7 +60,7 @@ class TestValidIdxSimplification(unittest.TestCase):
load = get_gated_load_uop(gate, idx)
self.check(load,
"0",
"(((lidx0+(gidx0*4))<19)!=True)")
"((((gidx0*4)+lidx0)<19)!=True)")
def test_simplify_within_valid1(self):
ridx0 = Range(0, 4)
@@ -186,7 +186,6 @@ class TestValidIdxSimplification(unittest.TestCase):
print("The expressions are not equivalent.")
print(s.model())
@unittest.expectedFailure # TODO: improve uop_given_valid
def test_valid_becomes_const2(self):
ridx0 = Range(0, 4)
ridx1 = Range(1, 4)
@@ -306,7 +305,7 @@ class TestImageSimplification(unittest.TestCase):
idx = ((alu4+1530)%1536, alu1+((idx1+((ridx2+7)//8)+31)//32)+(-2))
load = get_load_image_uop(shape, valid, idx)
self.check(load, None, "((((idx1*48)+(r2*6))+r0)+-6)", "(((idx2*2)+r1)+-1)")
self.check(load, None, "((((idx1*48)+r0)+(r2*6))+-6)", "(((idx2*2)+r1)+-1)")
def test_openpilot_conv2(self):
# conv in test/external/external_test_valid_remove.py
@@ -327,7 +326,7 @@ class TestImageSimplification(unittest.TestCase):
idx = ((alu3+765)%768, alu1+((idx1+((ridx2+7)//8)+31)//32)+(-2))
load = get_load_image_uop(shape, valid, idx)
self.check(load, None, "((((idx1*24)+(r2*3))+r0)+-3)", "(((idx2*2)+r1)+-1)")
self.check(load, None, "((((idx1*24)+r0)+(r2*3))+-3)", "(((idx2*2)+r1)+-1)")
def test_openpilot_conv3(self):
# in openpilot 0.9.7
@@ -349,8 +348,8 @@ class TestImageSimplification(unittest.TestCase):
self.check(load,
"((((idx2*2)+r0)<11)&((((idx1*8)+r1)<3)!=True))",
"(((idx0+((idx1*512)+(r1*64)))+832)%1024)",
"((((idx2*2)+r0)+(((idx1+((r1+5)//8))+1)//2))+-4)")
"(((idx0+(idx1*512))+(r1*64))+-192)",
"((((idx2*2)+(((idx1+((r1+5)//8))+1)//2))+r0)+-4)")
def test_simplify1(self):
# idx has the form (A % m, A // m + k) and valid has (c0 < A) and (A < c1)
@@ -388,16 +387,16 @@ class TestImageSimplification(unittest.TestCase):
# TODO: can this be simplified further?
load = get_load_image_uop(shape, alu9, (((alu8+(alu2*8))%64),(alu2//8)))
self.check(load, "(idx0<256)", "(((((idx0%8)*32)+(idx0//32))+8)%64)", "((idx0%8)//2)")
self.check(load, "(idx0<256)", "((((idx0//32)+((idx0%8)*32))+8)%64)", "((idx0%8)//2)")
load = get_load_image_uop(shape, alu9, (((alu8+(alu3*8))%64),(alu3//8)))
self.check(load, "(idx0<256)", "(((((idx0%8)*32)+(idx0//32))+16)%64)", "((idx0%8)//2)")
self.check(load, "(idx0<256)", "((((idx0//32)+((idx0%8)*32))+16)%64)", "((idx0%8)//2)")
load = get_load_image_uop(shape, alu9, (((alu8+(alu4*8))%64),(alu4//8)))
self.check(load, "(idx0<256)", "(((((idx0%8)*32)+(idx0//32))+24)%64)", "((idx0%8)//2)")
self.check(load, "(idx0<256)", "((((idx0//32)+((idx0%8)*32))+24)%64)", "((idx0%8)//2)")
load = get_load_image_uop(shape, alu9, (((alu8+(alu5*8))%64),(alu5//8)))
self.check(load, "(idx0<256)", "((((idx0%8)*32)+(idx0//32))%64)", "((idx0%8)//2)")
self.check(load, "(idx0<256)", "(((idx0//32)+((idx0%8)*32))%64)", "((idx0%8)//2)")
def test_simplify5(self):
# openpilot 0.9.7, chunk replacement to simplify
+14 -12
View File
@@ -10,20 +10,22 @@ class TestSymbolic(unittest.TestCase):
def test_symbolic_st(self):
x = Variable("x", 1, 100)
st = ShapeTracker.from_shape((x, 3))
self.assert_tuple_equal(st.shape, (x, 3))
self.assert_tuple_equal(st.is_expanded(), (False, False))
assert st.shape == (x, 3)
assert st.real_strides() == (3, 1)
def test_is_expanded_0(self):
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.assert_tuple_equal(st.is_expanded(), (False, False))
self.assertEqual(st.real_strides(), (8, None))
def test_is_expanded_1(self):
@unittest.expectedFailure
def test_real_strides_1(self):
st = ShapeTracker(views=(View(shape=(3, (Variable('i', 1, 10)+2)), strides=(Variable('i', 1, 10), 1), offset=0, mask=((0, 3), (0, Variable('i', 1, 10))), contiguous=False),)) # noqa: E501
self.assert_tuple_equal(st.is_expanded(), (False, False))
self.assertEqual(st.real_strides(), (Variable('i', 1, 10), None))
def test_is_expanded_2(self):
@unittest.expectedFailure
def test_real_strides_2(self):
st = ShapeTracker(views=(View(shape=(3, (Variable('i', 1, 10)+Variable('j', 1, 10))), strides=(Variable('i', 1, 10), 1), offset=0, mask=((0, 3), (0, Variable('i', 1, 10))), contiguous=False),)) # noqa: E501
self.assert_tuple_equal(st.is_expanded(), (False, False))
self.assertEqual(st.real_strides(), (Variable('i', 1, 10), None))
def test_merge_view_recursion_err(self):
vm2 = View(shape=(Variable('j', 1, 10),), strides=(0,), offset=0, mask=None, contiguous=False)
@@ -41,18 +43,18 @@ class TestSymbolic(unittest.TestCase):
self.assertEqual(vm3.strides, vm1.strides)
self.assertEqual(vm2+vm3, vm2)
def test_cat_dim0_is_expanded(self):
def test_cat_dim0_strides(self):
i = Variable("i", 1, 5).bind(3)
j = Variable("j", 1, 5).bind(3)
k = Variable("k", 1, 5).bind(3)
t = Tensor.rand(5, 4)[:i].cat(Tensor.rand(5, 4)[:j], dim=0).cat(Tensor.rand(5, 4)[:k], dim=0)
st = t.uop.st
self.assert_tuple_equal(st.shape, (i+j+k, 4))
self.assert_tuple_equal(st.is_expanded(), (False, False))
assert st.real_strides() == (4, 1)
t = Tensor.rand(5, 3)[:i].cat(Tensor.rand(5, 3)[:i], dim=0).cat(Tensor.rand(3, 3), dim=0)
st = t.uop.st
self.assert_tuple_equal(st.shape, (2*i+3, 3))
self.assert_tuple_equal(st.is_expanded(), (False, False))
assert st.real_strides() == (3, 1)
def test_cat_dim1_strides(self):
i = Variable("i", 1, 5).bind(4)
@@ -61,7 +63,7 @@ class TestSymbolic(unittest.TestCase):
t = Tensor.rand(3, 5)[:, :i].cat(Tensor.rand(3, 5)[:, :j], dim=1).cat(Tensor.rand(3, 5)[:, :k], dim=1)
st = t.uop.st
self.assert_tuple_equal(st.shape, (3, i+j+k))
self.assert_tuple_equal(st.is_expanded(), (False, False))
self.assert_tuple_equal(st.real_strides(), (i+j+k, 1))
class TestSymbolicVarVals(unittest.TestCase):
def assert_equal(self, x, y): self.assertFalse(x != y)
@@ -6,6 +6,7 @@ from tinygrad.uop.ops import UPat, Ops, UOp
realized_pattern = UPat(Ops.BUFFER)
# after realization, base tensor uops become RESHAPE(BUFFER)
buffer_view_pattern = UPat(Ops.RESHAPE, src=(UPat(Ops.BUFFER),))
const_pattern = UPat(Ops.CONST, src=(UPat(Ops.VIEW, src=(UPat(Ops.DEVICE),),)))
def is_pattern_uop(u:UOp, pat:UPat): assert pat.match(u, {}), f"{u}\nis not\n{pat}"
def is_pattern(ten:Tensor, pat:UPat): is_pattern_uop(ten.uop, pat)
+97
View File
@@ -0,0 +1,97 @@
from __future__ import annotations
import unittest
from tinygrad import Tensor
from tinygrad.helpers import DEBUG, RANGEIFY
from tinygrad.uop.ops import UOp, Ops, print_uops
from tinygrad.uop.spec import type_verify, ast_spec, tensor_uop_spec
from tinygrad.shape.shapetracker import ShapeTracker
from tinygrad import dtypes
from tinygrad.shape.view import View
from tinygrad.engine.realize import get_program
from tinygrad.device import Device
class InvalidASTException(Exception): pass
def helper_test_verify_ast(*stores:UOp):
sink = UOp(Ops.SINK, dtypes.void, stores)
if DEBUG >= 3:
for op in stores: print(op)
try: type_verify(list(sink.toposort()), ast_spec)
except RuntimeError as e: raise InvalidASTException(e.args)
program = get_program(sink, Device[Device.DEFAULT].renderer)
if DEBUG >= 6: print_uops(program.uops)
if DEBUG >= 4: print(program.src)
class TestUOpSpec(unittest.TestCase):
def test_tiny_add(self):
dtype = dtypes.int
buf_0 = UOp(Ops.DEFINE_GLOBAL, dtype.ptr(), (), 0)
buf_1 = UOp(Ops.DEFINE_GLOBAL, dtype.ptr(), (), 1)
buf_2 = UOp(Ops.DEFINE_GLOBAL, dtype.ptr(), (), 2)
a = UOp(Ops.LOAD, dtype, (buf_1.view(ShapeTracker.from_shape((32, 1))),))
b = UOp(Ops.LOAD, dtype, (buf_2.view(ShapeTracker.from_shape((32, 1))),))
store = UOp(Ops.STORE, dtypes.void, (buf_0.view(ShapeTracker.from_shape((32, 1))), a+b))
helper_test_verify_ast(store)
def test_no_implicit_broadcasting(self):
bufs = [UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), i) for i in range(2)]
a = UOp(Ops.LOAD, dtypes.float, (bufs[1].view(ShapeTracker.from_shape((4, 32))),))
b = a + UOp(Ops.REDUCE_AXIS, dtypes.float, (a,), (Ops.MAX, (1,)))
st = UOp(Ops.STORE, dtypes.void, (bufs[0].view(ShapeTracker.from_shape((4, 32))), b))
with self.assertRaises(InvalidASTException): helper_test_verify_ast(st)
def test_shrink_ok(self):
bufs = [UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), i) for i in range(2)]
a = UOp(Ops.LOAD, dtypes.float, (bufs[1].view(ShapeTracker((View((32, 32), strides=(32, 1), offset=0, mask=None, contiguous=True),))),))
b = UOp(Ops.LOAD, dtypes.float, (bufs[1].view(ShapeTracker((View((32, 32), strides=(0, 1), offset=0, mask=None, contiguous=False),))),))
st = UOp.store(bufs[0].view(ShapeTracker.from_shape((32, 32))), a+b)
helper_test_verify_ast(st)
def test_reduce_store(self):
bufs = [UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), i) for i in range(2)]
a = UOp(Ops.LOAD, dtypes.float, (bufs[1].view(ShapeTracker.from_shape((32, 1))),))
r = UOp(Ops.REDUCE_AXIS, dtypes.float, (a,), (Ops.ADD, (0,)))
st = UOp.store(bufs[0].view(ShapeTracker.from_shape((32, 1))), r)
with self.assertRaises(InvalidASTException): helper_test_verify_ast(st)
def test_reduce_add_store(self):
bufs = [UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), i) for i in range(2)]
a = UOp(Ops.LOAD, dtypes.float, (bufs[1].view(ShapeTracker.from_shape((32, 1))),))
r = UOp(Ops.REDUCE_AXIS, dtypes.float, (a,), (Ops.ADD, (0,)))
st = UOp.store(bufs[0].view(ShapeTracker.from_shape((32, 1))), r+a)
with self.assertRaises(InvalidASTException): helper_test_verify_ast(st)
def test_assert_swizzle(self):
buf = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
a = UOp(Ops.LOAD, dtypes.float, (buf.view(ShapeTracker.from_shape((32, 1))),))
r = UOp(Ops.REDUCE_AXIS, dtypes.float, (a,), (Ops.ADD, (0,)))
st = UOp.store(buf.view(ShapeTracker.from_shape((32, 1))), r.view(r.st.expand((32, 1)))+a)
with self.assertRaisesRegex(InvalidASTException, "UOp verification failed"): helper_test_verify_ast(st)
def test_const_view_always_valid(self):
buf = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
a = UOp.const(dtypes.int, 0).replace(src=(UOp(Ops.VIEW, dtypes.void, (), ShapeTracker.from_shape(())),))
st = UOp.store(buf.view(ShapeTracker.from_shape(())), a.cast(dtypes.float))
helper_test_verify_ast(st)
@unittest.skipIf(RANGEIFY, "RANGEIFY does not push views")
def test_assert_masked_view_in_const(self):
t = Tensor(6).uop
a = t.replace(src=(t.src[0].replace(arg=t.st.reshape((1,)).pad(((0, 1),))),))
with self.assertRaisesRegex(RuntimeError, "UOp verification failed"):
type_verify([a], tensor_uop_spec)
class TestUOpSink(unittest.TestCase):
def test_0(self):
s = UOp.sink()
self.assertEqual(len(s.src), 0)
def test_1(self):
a = UOp.const(dtypes.int, 0)
s1 = UOp.sink(a)
s2 = a.sink()
self.assertIs(s1, s2)
if __name__ == '__main__':
unittest.main()
+65 -13
View File
@@ -116,6 +116,39 @@ class TestSymbolic(unittest.TestCase):
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 helper_test_factor(self, expr, *factors):
factored = expr.factor(*factors)
self.check_equal_z3(expr, factored)
for fac in factors: self.assertIn(fac, factored.toposort())
def test_uop_factor(self):
a = Variable("a", 0, 8)
b = Variable("b", 0, 8)
c = Variable("c", 0, 8)
self.helper_test_factor((1400*a+2800*b), (a+2*b))
self.helper_test_factor((1400*a+2800*b)%9000, (a+2*b))
self.helper_test_factor((a+2*b), (a+2*b))
self.helper_test_factor((a+c+2*b), (a+2*b))
self.helper_test_factor((1400*a+c+2800*b)%9000, (a+2*b))
self.helper_test_factor((1399*a+c+2800*b)%9000+1400*a+2800*b, (a+2*b))
self.helper_test_factor((1400*a+c+2800*b)%9000+1400*a+2800*b, (a+2*b))
# self.assertIsNone((a+c+3*b).factor(a+2*b))
# self.assertIsNone((1399*a+c+2800*b).factor(a+2*b))
def test_uop_multiple_factors(self):
a = Variable("a", 0, 8)
b = Variable("b", 0, 8)
c = Variable("c", 0, 8)
d = Variable("d", 0, 8)
self.helper_test_factor((1400*a+2800*b+2*c+d), (a+2*b), (2*c+d))
self.helper_test_factor((100*a+200*b+5*c), (a+2*b), (5*c))
self.helper_test_factor((3*a+6*b+2*c+4*d), (a+2*b), (c+2*d))
self.helper_test_factor((7*a+14*b+3*c+6*d), (a+2*b), (3*c+6*d))
self.helper_test_factor((10*a+20*b+10*c+30*d), (a+2*b), (c+3*d))
self.helper_test_factor((10*c+(10*a+20*b)//3+30*d), (a+2*b), (c+3*d))
self.helper_test_factor((10*c+(10*a+20*b)//3+30*d), (a+2*b), (c+3*d))
# self.assertIsNone((7*a+14*b+3*c+6*d).factor((a+8*b), (2*c+6*d)))
def test_divide_exact_not(self):
a = Variable("a", 1, 8)
b = Variable("b", 1, 8)
@@ -130,13 +163,13 @@ class TestSymbolic(unittest.TestCase):
a = Variable("a", 0, 8)
b = Variable("b", 0, 8)
self.helper_test_variable(a*2+a*3, 0, 8*5, "(a*5)")
self.helper_test_variable(b+a*2+a*3, 0, 8*6, "(b+(a*5))")
self.helper_test_variable(b+a*2+a*3, 0, 8*6, "((a*5)+b)")
def test_factorize_no_mul(self):
a = Variable("a", 0, 8)
b = Variable("b", 0, 8)
self.helper_test_variable(a+a*3, 0, 8*4, "(a*4)")
self.helper_test_variable((a+b)+a*3, 0, 8*5, "(b+(a*4))")
self.helper_test_variable((a+b)+a*3, 0, 8*5, "((a*4)+b)")
self.helper_test_variable((a*3+b)+b*3, 0, 8*7, "((a*3)+(b*4))")
def test_neg(self):
@@ -159,8 +192,15 @@ class TestSymbolic(unittest.TestCase):
b = Variable("b", 0, 8)
self.helper_test_variable(a+a, 0, 16, "(a*2)")
self.helper_test_variable((a+b)+b, 0, 24, "(a+(b*2))")
self.helper_test_variable((a*3+b)+a, 0, 40, "(b+(a*4))")
self.helper_test_variable((a+b)+a*3, 0, 40, "(b+(a*4))")
self.helper_test_variable((a*3+b)+a, 0, 40, "((a*4)+b)")
self.helper_test_variable((a+b)+a*3, 0, 40, "((a*4)+b)")
def test_add_self_seperated(self):
a = Variable("a", 0, 8)
b = Variable("b", 0, 8)
c = Variable("c", 0, 8)
self.helper_test_variable((a+b)+c+a, 0, 32, "(((a*2)+b)+c)")
self.helper_test_variable((a*3+b*2)+c*2+a*5, 0, 96, "(((a*8)+(b*2))+(c*2))")
def test_sub_self(self):
a = Variable("a", 0, 8)
@@ -279,7 +319,7 @@ class TestSymbolic(unittest.TestCase):
def test_mod_congruence_multiple_vars(self):
self.helper_test_variable((9+9*Variable("x",0,3)+9*Variable("y",0,3))%10, 3, 9, "(((x*-1)+(y*-1))+9)")
self.helper_test_variable((7+9*Variable("x",0,2)+9*Variable("y",0,2)+Variable("z",0,2))%10, 3, 9,
("(((z+(x*-1))+(y*-1))+7)", "(((y*-1)+(z+(x*-1)))+7)"))
("(((z+(x*-1))+(y*-1))+7)", "(((y*-1)+(z+(x*-1)))+7)", "((((x*-1)+(y*-1))+z)+7)"))
self.helper_test_variable((10+12*Variable("x",0,2)+Variable("y", 0, 4)%3)%13, 8, 12, "(((x*-1)+(y%3))+10)")
def test_div_congruence(self):
@@ -455,7 +495,7 @@ class TestSymbolic(unittest.TestCase):
ridx1005 = UOp.variable("ridx1005", 0, 2)
ridx1006 = UOp.variable("ridx1006", 0, 2)
self.helper_test_variable((lidx1+((gidx1*18)+(ridx1005*18)+(lidx0*162))+(gidx0*2)+(ridx1006*2)+-40)//18, -2, 20,
"(((((lidx1+(((gidx1*18)+(ridx1005*18))+(lidx0*162)))+(gidx0*2))+(ridx1006*2))+-40)//18)")
"((((((((gidx0*2)+(gidx1*18))+(lidx0*162))+lidx1)+(ridx1005*18))+(ridx1006*2))+-40)//18)")
def test_add_div(self):
# careful about the lower bounds and upper bounds
@@ -498,7 +538,7 @@ class TestSymbolic(unittest.TestCase):
self.helper_test_variable((d1*a*b*d1)//(d1), -1000, 1000, "(a*(b*d1))", test_z3=False)
self.helper_test_variable((d1*a*d2*b*d1)//(d1*d2), -1000, 1000, "(a*(b*d1))", test_z3=False)
self.helper_test_variable((d1*a + b*d1)//(d1), -20, 20, "(a+b)", test_z3=False)
self.helper_test_variable((d1*a + b*d1 + c*d1)//(d1), -30, 30, "(c+(a+b))", test_z3=False)
self.helper_test_variable((d1*a + b*d1 + c*d1)//(d1), -30, 30, "((a+b)+c)", test_z3=False)
self.helper_test_variable((3*a*d1 + 9*b*d1)//(3*d1*d2), -40, 40, "(((a+(b*3))//(d2*-1))*-1)", test_z3=False)
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)", test_z3=False)
@@ -508,7 +548,7 @@ class TestSymbolic(unittest.TestCase):
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))")
self.helper_test_variable((d*a*20+b*d*5+10)//(5*d), 0, 52, "(((a*4)+b)+(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)")
@@ -561,9 +601,7 @@ class TestSymbolic(unittest.TestCase):
lidx2 = Variable("lidx2", 0, 3)
alu0 = gidx2*640+gidx1*160+(gidx0//5)*2+lidx0*320+lidx1*10
self.helper_test_variable((alu0+lidx2*2+1)//20, 0, 8192,
("((((((gidx0//5)+lidx2)//5)+lidx1)//2)+(((gidx2*32)+(gidx1*8))+(lidx0*16)))",
"(((lidx1+((lidx2+(gidx0//5))//5))//2)+((gidx2*32)+((gidx1*8)+(lidx0*16))))",
"((((gidx1*8)+(gidx2*32))+(lidx0*16))+((lidx1+((lidx2+(gidx0//5))//5))//2))"))
("((((gidx1*8)+(gidx2*32))+(lidx0*16))+((lidx1+((lidx2+(gidx0//5))//5))//2))",))
def test_sum_div_complex2(self):
gidx0 = Variable("gidx0", 0, 7)
@@ -641,8 +679,21 @@ class TestSymbolic(unittest.TestCase):
self.helper_test_variable((gidx//4)*4+gidx%4, 0, 124, "gidx")
self.helper_test_variable(lidx+gidx%4+(gidx//4)*4, 0, 248, "(gidx+lidx)")
self.helper_test_variable(lidx+(gidx//4)*4+gidx%4, 0, 248, "(gidx+lidx)")
self.helper_test_variable(lidx+(gidx//4)*8+2*(gidx%4), 0, 372, "(lidx+(gidx*2))")
self.helper_test_variable(lidx+2*(gidx%4)+(gidx//4)*8, 0, 372, "(lidx+(gidx*2))")
self.helper_test_variable(lidx+(gidx//4)*8+2*(gidx%4), 0, 372, "((gidx*2)+lidx)")
self.helper_test_variable(lidx+2*(gidx%4)+(gidx//4)*8, 0, 372, "((gidx*2)+lidx)")
def test_div_mod_recombine_seperated(self):
gidx = Variable("gidx", 0, 124)
lidx = Variable("lidx", 0, 124)
a = Variable("a", 0, 3)
b = Variable("b", 0, 3)
c = Variable("c", 0, 3)
self.helper_test_variable(gidx%4+a+b+c+(gidx//4)*4, 0, 133, "(((a+b)+c)+gidx)")
self.helper_test_variable((gidx//4)*4+a+b*10+gidx%4, 0, 157, "((a+(b*10))+gidx)")
self.helper_test_variable(lidx+gidx%4+a+b+c//2+(gidx//4)*4, 0, 255, "((((a+b)+gidx)+lidx)+(c//2))")
self.helper_test_variable(lidx+(gidx//4)*8+b+c+a*8+2*(gidx%4), 0, 402, "(((((a*8)+b)+c)+(gidx*2))+lidx)")
# TODO: need better sorting for this one
# self.helper_test_variable(lidx+(gidx//4)*4+a*3+b*3+(c*10)%3+gidx%4, , , "")
def test_div_mod_recombine_folded_mod(self):
a = Variable("a", 0, 2)
@@ -1019,6 +1070,7 @@ class TestSymbolicRealWorld(unittest.TestCase):
("((((((((((lidx5+1)//16)*802816)+(((lidx5+1)%16)*49))+(gidx0*3211264))+(gidx1*784))+(gidx2*8))+(lidx4*100352))+lidx3)+2207744)",
'((lidx3+((((((((lidx5+1)//16)*802816)+(((lidx5+1)%16)*49))+(gidx0*3211264))+(gidx1*784))+(gidx2*8))+(lidx4*100352)))+2207744)',
'((lidx3+((lidx4*100352)+((gidx2*8)+((gidx1*784)+((gidx0*3211264)+((((lidx5+1)//16)*802816)+(((lidx5+1)%16)*49)))))))+2207744)',
'((((((((gidx0*3211264)+(gidx1*784))+(gidx2*8))+lidx3)+(lidx4*100352))+(((lidx5+1)//16)*802816))+(((lidx5+1)%16)*49))+2207744)',
))
class TestBounds(unittest.TestCase):
+19
View File
@@ -10,6 +10,25 @@ class TestView(unittest.TestCase):
v = View.create(shape=(4,3,2), strides=(1,4,10), mask=((0,4),(0,3),(0,2)))
self.assertIsNone(v.mask)
def test_minify_zero_strided_dims(self):
target = View.create(shape=(2,2), strides=(30,2), offset=7, mask=None)
v = View.create(shape=(2,1,2), strides=(30,0,2), offset=7, mask=None)
self.assertEqual(v.minify(), target)
v = View.create(shape=(1,2,2), strides=(0,30,2), offset=7, mask=None)
self.assertEqual(v.minify(), target)
v = View.create(shape=(2,2,1), strides=(30,2,0), offset=7, mask=None)
self.assertEqual(v.minify(), target)
v = View.create(shape=(2,1,1,2), strides=(30,0,0,2), offset=7, mask=None)
self.assertEqual(v.minify(), target)
v = View.create(shape=(1,1,2,2), strides=(0,0,30,2), offset=7, mask=None)
self.assertEqual(v.minify(), target)
v = View.create(shape=(2,2,1,1), strides=(30,2,0,0), offset=7, mask=None)
self.assertEqual(v.minify(), target)
v = View.create(shape=(1,2,2,1), strides=(0,30,2,0), offset=7, mask=None)
self.assertEqual(v.minify(), target)
v = View.create(shape=(1,2,1,2), strides=(0,30,0,2), offset=7, mask=None)
self.assertEqual(v.minify(), target)
def test_empty_mask_contiguous(self):
v1 = View.create(shape=(2,2,2), strides=(4,2,1), mask=None)
v2 = View.create(shape=(2,2,2), strides=(4,2,1), mask=((0,2),(0,2),(0,2)))
+3 -3
View File
@@ -379,9 +379,9 @@ class TestVizProfiler(unittest.TestCase):
j = load_profile(prof)
tracks = list(j['layout'])
self.assertEqual(tracks[0], 'NV')
self.assertEqual(tracks[1], 'NV:1')
self.assertEqual(tracks[2], 'NV Graph')
self.assertEqual(tracks[0], 'NV Graph')
self.assertEqual(tracks[1], 'NV')
self.assertEqual(tracks[2], 'NV:1')
nv_events = j['layout']['NV']['events']
self.assertEqual(nv_events[0]['name'], 'E_25_4n2')
+6 -6
View File
@@ -1,7 +1,7 @@
import unittest, sys
import numpy as np
from tinygrad import Tensor, GlobalCounters, dtypes, Context, nn
from tinygrad.helpers import CI, Profiling, WINO
from tinygrad.helpers import CI, Profiling, WINO, RANGEIFY
@unittest.skipIf(sys.platform.startswith("win"), "flaky on Windows")
class TestWinogradClose(unittest.TestCase):
@@ -35,14 +35,14 @@ class TestWinograd(unittest.TestCase):
def test_forward_kernels(self):
x,w = Tensor.rand(1,4,9,9).realize(), Tensor.rand(4,4,3,3).realize()
out = Tensor.conv2d(x,w)
self.assertEqual(len(out.schedule()), 2)
self.assertEqual(len(out.schedule()), 2 if RANGEIFY else 4)
def test_backward_kernels(self):
x,w = Tensor.empty(1,4,9,9,requires_grad=True).realize(), Tensor.empty(4,4,3,3,requires_grad=True).realize()
out = Tensor.conv2d(x,w, padding=1)
out.mean().backward()
backward_schedule = Tensor.schedule(x.grad, w.grad)
self.assertEqual(len(backward_schedule), 4)
self.assertEqual(len(backward_schedule), 4 if RANGEIFY else 9)
def test_counters(self):
IC, OC, X, Y = 4,4,9,9
@@ -61,9 +61,9 @@ class TestWinograd(unittest.TestCase):
print(f"ops: normal {ops_normal:9d} wino {ops_wino:9d} ratio {ops_ratio:.2f}")
print(f"mem: normal {mem_normal:9d} wino {mem_wino:9d} ratio {mem_ratio:.2f}")
# TODO: what's optimal on this?
self.assertLess(ops_ratio, 4.3)
self.assertLess(mem_ratio, 3)
if not RANGEIFY:
self.assertLess(ops_ratio, 2.6) # TODO: there's issues with factorization now
self.assertLess(mem_ratio, 10)
def test_dtype(self):
IC, OC, X, Y = 4,4,9,9
+30 -14
View File
@@ -1,12 +1,13 @@
from typing import Any, Callable
import functools
from dataclasses import dataclass
from tinygrad.helpers import QUANTIZE, DEVECTORIZE, TRANSCENDENTAL
from tinygrad.helpers import QUANTIZE, DEVECTORIZE, TRANSCENDENTAL, RANGEIFY
from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype
from tinygrad.uop.spec import type_verify
from tinygrad.renderer import Renderer
# import all pattern matchers here
from tinygrad.codegen.lowerer import pm_lowerer, get_index
from tinygrad.codegen.quantize import pm_quant
from tinygrad.codegen.gpudims import pm_add_gpudims
from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing, symbolic
@@ -14,7 +15,8 @@ from tinygrad.uop.decompositions import get_late_rewrite_patterns
from tinygrad.codegen.late.expander import migrate_indexing, expander, pm_pre_expander, pm_group_for_reduce
from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \
ReduceContext, correct_load_store, pm_render
from tinygrad.codegen.late.control_flow import pm_control_flow_ends, pm_control_flow_starts, CFGContext, linearize
from tinygrad.codegen.late.linearize import block_create, pm_blockend_merge, block_merge, pm_finalize, BlockContext
from tinygrad.codegen.opt.swizzler import view_left, view_right, fix_kernel_ops
from tinygrad.codegen.opt.postrange import pm_postrange_opt
from tinygrad.codegen.simplify import pm_simplify_ranges, pm_reduce_simplify, pm_flatten_range, pm_split_ranges
from tinygrad.schedule.rangeify import pm_add_buffers, rangeify_codegen
@@ -30,22 +32,39 @@ class RewriteStep:
def apply_rewrites(sink:UOp, rewrites:list[RewriteStep]): return functools.reduce(lambda x,f: f(x), rewrites, sink)
rewrites_for_views = [
RewriteStep(view_left, name="Main View Left"),
RewriteStep(view_right, name="Main View Right"),
RewriteStep(view_left+fix_kernel_ops, bottom_up=True, name="Finalize Kernel"),
]
rewrites_for_linearizer = [
RewriteStep(block_create, ctx=BlockContext.from_sink, name="Linearizer: Create Blocks", bottom_up=True),
RewriteStep(pm_blockend_merge, name="Linearizer: Merge Blockends"),
RewriteStep(block_merge, name="Linearizer: Merge Blocks"),
RewriteStep(pm_finalize, name="Linearizer: Finalize")]
def get_rewrites_for_renderer(opts:Renderer, optimize:bool=True, linearizer:bool=True) -> list[RewriteStep]:
# cache with the values of the context vars
return _get_rewrites_for_renderer(opts, optimize, linearizer, QUANTIZE.value, DEVECTORIZE.value, TRANSCENDENTAL.value)
return _get_rewrites_for_renderer(opts, optimize, linearizer, QUANTIZE.value, DEVECTORIZE.value, TRANSCENDENTAL.value, RANGEIFY.value)
@functools.cache
def _get_rewrites_for_renderer(opts:Renderer, optimize:bool, linearizer:bool, _QUANTIZE, _DEVECTORIZE, _TRANSCENDENTAL) -> list[RewriteStep]:
def _get_rewrites_for_renderer(opts:Renderer, optimize:bool, linearizer:bool, _QUANTIZE, _DEVECTORIZE, _TRANSCENDENTAL,
_RANGEIFY) -> list[RewriteStep]:
# ** lowerer (rewrite_shapetracker_with_index) **
ret: list[RewriteStep] = []
if optimize:
# view pushing
if not _RANGEIFY: ret.extend(rewrites_for_views)
# lowerer first
if _QUANTIZE and opts.device in {"CPU", "DSP"}: ret.append(RewriteStep(pm_quant, name="quantize"))
ret.append(RewriteStep(pm_lowerer, get_index, name="lowerer", bottom_up=True))
# split ranges
ret.append(RewriteStep(pm_split_ranges+pm_flatten_range, ctx=lambda _: {}, name="split ranges"))
if _RANGEIFY:
ret.append(RewriteStep(pm_split_ranges+pm_flatten_range, ctx=lambda _: {}, name="split ranges"))
# symbolic (NOTE: this is a requirement for pm_simplify_ranges to be correct)
ret.append(RewriteStep(sym+pm_flatten_range, name="initial symbolic"))
@@ -95,15 +114,11 @@ def _get_rewrites_for_renderer(opts:Renderer, optimize:bool, linearizer:bool, _Q
pm_final_rewrite = pm_decomp+pm_render+extra_matcher
ret.append(RewriteStep(pm_final_rewrite, lambda _: opts.device, name="final rewrite"))
# add control flow to the graph
ret.append(RewriteStep(pm_control_flow_ends, name="add control flow ends"))
ret.append(RewriteStep(pm_control_flow_starts, CFGContext, name="add control flow starts", bottom_up=True))
# return the list (with optional linearizer)
return ret + (rewrites_for_linearizer if linearizer else [])
# return the list
return ret
def full_rewrite_to_sink(sink:UOp, opts:Renderer|None=None, optimize:bool=True) -> UOp:
return apply_rewrites(sink, get_rewrites_for_renderer(opts if opts is not None else Renderer(), optimize))
def full_rewrite_to_sink(sink:UOp, opts:Renderer|None=None, optimize:bool=True, linearizer:bool=False) -> UOp:
return apply_rewrites(sink, get_rewrites_for_renderer(opts if opts is not None else Renderer(), optimize, linearizer))
def full_rewrite(sink:UOp, opts:Renderer|None=None) -> list[UOp]:
"""
@@ -116,6 +131,7 @@ def full_rewrite(sink:UOp, opts:Renderer|None=None) -> list[UOp]:
Returns:
Linear program in UOps.
"""
lst = linearize(full_rewrite_to_sink(sink, opts, optimize=sink.tag is None))
lst = list(full_rewrite_to_sink(sink, opts, optimize=sink.tag is None, linearizer=True).arg.lst)
if __debug__: type_verify(lst)
return lst
-105
View File
@@ -1,105 +0,0 @@
from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat
from tinygrad.helpers import dedup
from collections import defaultdict
from itertools import groupby
from functools import reduce
import heapq
def linearize(u:UOp) -> list[UOp]:
lst = list(u.toposort())
in_this_block = set(lst)
local_children: defaultdict[UOp, list[UOp]] = defaultdict(list)
in_degree:dict[UOp, int] = {}
priorities:dict[UOp, int] = {}
# get local children and assign priorities
# NOTE: this requires the lst be locally toposorted
for u in reversed(lst):
in_degree[u] = 0
for s in u.src:
if s in in_this_block:
local_children[s].append(u)
in_degree[u] += 1
# put loads in the beginning of the block and prevent priority inversion. hack for BARRIER grouping too
priority = [0] + [priorities[x] for x in local_children[u]]
if u.op is Ops.LOAD: priority.append(-1000)
if u.op is Ops.BARRIER: priority.append(-1500)
# ranges are scheduled as late as possible so anything that can be outside is
#if u.op is Ops.RANGE: priority = [2000]
# move defines and consts to the top
if u.op in {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.DEFINE_VAR, Ops.SPECIAL, Ops.CONST}: priority.append(-2000)
priorities[u] = min(priority)
# number the uops in "ideal" order
nkey = {u:i for i,u in enumerate(sorted(lst, key=lambda x: (priorities[x],)+x.tuplize))}
# then force then to be toposorted in as close to the ideal order as possible
heapq.heapify(heap:=[(nkey[u],u) for u in lst if in_degree[u] == 0])
newlst = []
while heap:
newlst.append(u:=heapq.heappop(heap)[1])
for v in local_children[u]:
in_degree[v] -= 1
if in_degree[v] == 0: heapq.heappush(heap, (nkey[v],v))
assert len(newlst) == len(lst), f"len mismatch {len(newlst)} != {len(lst)}"
return newlst
def add_endrange(x:UOp):
if not ((x.op is Ops.LOAD and x.src[-1].op is Ops.STORE) or all(s.op is Ops.STORE and any(n.op is Ops.RANGE for n in s.src) for s in x.src)):
return None
src: list[UOp] = []
for k,g in groupby(x.src, key=lambda k: tuple(dedup(s for s in k.src if s.op is Ops.RANGE))):
if not k: src.extend(g)
else: src.extend(reduce(lambda acc,rng: (UOp(Ops.ENDRANGE, src=(rng,) + acc),), reversed(k), tuple(g))) # type: ignore
return x.replace(src=tuple(src))
def add_endif(x:UOp):
groups = {k: tuple(g) for k,g in groupby(x.src, key=lambda k: k.src[2] if len(k.src) >= 3 and k.src[2].op is Ops.IF else k)}
if not any(k.op is Ops.IF for k in groups): return None
return x.replace(src=tuple(UOp(Ops.ENDIF, src=(k,) + g) if k.op is Ops.IF else k for k,g in groups.items()))
# some Ops.IF aren't closed by an Ops.STORE, in that case the Ops.SINK closes it
def close_ifs(x:UOp):
consumers = x.get_consumer_map()
if (y:=next((s for s in consumers if s.op is Ops.IF and all(n.op is not Ops.ENDIF for n in consumers[s])), None)) is not None:
return x.replace(src=(UOp(Ops.ENDIF, src=(y,) + x.src),))
return None
pm_control_flow_ends = PatternMatcher([
(UPat((Ops.SINK, Ops.NOOP, Ops.LOAD), name="x"), add_endrange),
(UPat((Ops.SINK, Ops.ENDRANGE, Ops.BARRIER), name="x"), add_endif),
(UPat(Ops.SINK, name="x"), close_ifs),
])
class CFGContext:
def __init__(self, sink:UOp):
# there are 3 relationships between ranges:
# nested, meaning endrange y is a dependency of endrange x and range x is a dependency of endrange y
# dependent, meaning endrange y is a dependency of endrange x and range x is not a dependency of endrange y
# independent, endrange y is not a dependency of endrange x
deps: dict[UOp, set[UOp]] = {}
nesting: dict[UOp, UOp] = {}
for u in sink.toposort():
deps[u] = set().union(*(deps[s] for s in u.src))
if u.op in (Ops.ENDRANGE, Ops.ENDIF):
for n in [x for x in deps[u] if x.op in (Ops.ENDRANGE, Ops.ENDIF) and u.src[0] in deps[x] and x not in nesting]: nesting[n] = u
if u.op is Ops.SINK:
for n in [x for x in deps[u] if x.op in (Ops.ENDRANGE, Ops.ENDIF) and x not in nesting]: nesting[n] = u
if u.op in (Ops.RANGE, Ops.ENDRANGE, Ops.IF, Ops.ENDIF): deps[u] |= {u}
self.edges: dict[UOp, UOp] = {}
siblings: dict[UOp, list[UOp]] = {}
for k,vv in nesting.items(): siblings.setdefault(vv, []).append(k)
for k,v in siblings.items():
# range/if that have dependencies on other siblings need to run after them
order = sorted(v, key=lambda x: len([y for y in v if y in deps[x]]))
zipped = zip(order, order[1:]) if k.op is Ops.SINK else zip([k.src[0]] + order, order)
for x,y in zipped: self.edges[y.src[0]] = x
pm_control_flow_starts = PatternMatcher([
(UPat((Ops.RANGE, Ops.IF), src=(UPat(),), name="x"), lambda ctx,x: x.replace(src=x.src+(y,)) if (y:=ctx.edges.get(x)) is not None else None),
(UPat(Ops.IF, src=(UPat(), UPat(Ops.BARRIER)), name="x"), lambda ctx,x: x.replace(src=x.src+(y,)) if (y:=ctx.edges.get(x)) is not None else None),
# optional: remove ranges from STORE. keep NOOP since they determine ordering
#(UPat(Ops.STORE, name="s"), lambda s: s.replace(src=s.src[0:2]+tuple([x for x in s.src[2:] if x.op not in {Ops.RANGE, Ops.CONST}]))),
])
+1 -1
View File
@@ -50,7 +50,7 @@ def delete_redundant_gates(store:UOp, buf:UOp, idx:UOp, val:UOp, store_gate:UOp,
# remove the gate from the index
return UOp.store(buf.index(idx).cast(cast.dtype) if cast is not None else buf.index(idx), val, *store.src[2:])
def no_load(u:UOp) -> bool: return not any(x.op is Ops.LOAD for x in u.backward_slice_with_self)
def no_load(u:UOp) -> bool: return not any(x.op is Ops.LOAD for x in u.sparents)
load_store_indexing = PatternMatcher([
# image load valid idx simplification
(UPat(Ops.INDEX, src=(UPat.var("buf"), invalid_gate)), lambda buf,x,i,cond: simplify_valid_load(buf, x, cond)),
+243
View File
@@ -0,0 +1,243 @@
from __future__ import annotations
import heapq
from collections import defaultdict
from dataclasses import dataclass, replace
from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat, GroupOp, BottomUpGate
from tinygrad.helpers import dedup, all_same, flatten, BLOCK_REORDER
# NOTE: any toposort should be valid here, unlike last time this isn't required, it's just for speed
def block_reorder(lst:list[UOp]) -> list[UOp]:
in_this_block = set(lst)
local_children: defaultdict[UOp, list[UOp]] = defaultdict(list)
in_degree:dict[UOp, int] = {}
priorities:dict[UOp, int] = {}
# get local children and assign priorities
# NOTE: this requires the lst be locally toposorted
for u in reversed(lst):
in_degree[u] = 0
for s in u.src:
if s in in_this_block:
local_children[s].append(u)
in_degree[u] += 1
# put loads in the beginning of the block and prevent priority inversion. hack for BARRIER grouping too
priority = [0] + [priorities[x] for x in local_children[u]]
if u.op is Ops.LOAD: priority.append(-1000)
if u.op is Ops.BARRIER: priority.append(-1500)
priorities[u] = min(priority)
# number the uops in "ideal" order
nkey = {u:i for i,u in enumerate(sorted(lst, key=lambda x: (priorities[x],)+x.tuplize))}
# then force then to be toposorted in as close to the ideal order as possible
heapq.heapify(heap:=[(nkey[u],u) for u in lst if in_degree[u] == 0])
newlst = []
while heap:
newlst.append(u:=heapq.heappop(heap)[1])
for v in local_children[u]:
in_degree[v] -= 1
if in_degree[v] == 0: heapq.heappush(heap, (nkey[v],v))
assert len(newlst) == len(lst), f"len mismatch {len(newlst)} != {len(lst)}"
return newlst
# ***** basic block *****
def disp(y:UOp) -> str:
if y.op is Ops.IF: return f'IF{id(y)}'
if y.op is Ops.RANGE: return str(y.arg)
return "<NONE>"
@dataclass(frozen=True, eq=False)
class BasicBlock:
lst: tuple[UOp, ...]
ctx: tuple[UOp, ...] = ()
end: UOp|None = None
cnt: int = 0
child_ctx: tuple[UOp, ...]|None = None
def __lt__(self, _:BasicBlock): raise RuntimeError("no comparing basic blocks")
def __repr__(self):
return f"{(str(disp(self.end))+' ') if self.end is not None else ''}"+f'f{self.cnt} '+\
f"{[disp(y) for y in self.ctx]} {[disp(y) for y in self.child_ctx] if self.child_ctx is not None else '-'} "+\
f"{len(self.lst)}" + "\n" + '\n'.join([str(x.op) for x in self.lst])
def last_ctx(self): return self.child_ctx if self.child_ctx is not None else self.ctx
def _sort_ctx(inp): return tuple(sorted(dedup(inp), key=lambda x: x.tuplize))
# ***** block context *****
@dataclass
class BlockContext:
child_count: dict[UOp, int]
block_ctxs: dict[UOp, tuple[UOp, ...]]
child_ctxs: dict[UOp, tuple[UOp, ...]]
def last_ctx(self, u): return self.child_ctxs.get(u, self.block_ctxs[u])
@staticmethod
def from_sink(sink:UOp) -> BlockContext:
# get children and all block contexts
ctx = BlockContext({}, {}, {})
for u in sink.toposort(gate=lambda u:u.op is not Ops.SPECIAL):
this_block_ctx: list[UOp] = []
ctx.child_count[u] = 0
# get children and accumulate the last_ctx
for s in u.src:
if s.op is Ops.SPECIAL: continue
# NOTE: if a parent appears multiple times in the src, it counts multiple times as a child
ctx.child_count[s] += 1
this_block_ctx += ctx.last_ctx(s)
# save the block ctx. SINK never has anything
ctx.block_ctxs[u] = _sort_ctx(this_block_ctx) if u.op is not Ops.SINK else ()
# RANGE/IF add to the next ctx
# STORE/ASSIGN subtract from the next ctx
if u.op in {Ops.RANGE, Ops.IF}: ctx.child_ctxs[u] = _sort_ctx(ctx.block_ctxs[u] + (u,))
elif u.op is Ops.STORE: ctx.child_ctxs[u] = tuple([y for y in ctx.block_ctxs[u] if y not in u.src])
return ctx
# ***** make blocks *****
DONT_PLACE_IN_BLOCK = {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.DEFINE_VAR, Ops.SPECIAL, Ops.CONST}
def add_blockends(base_block:UOp, new_ctx:tuple[UOp, ...], current_ctx:tuple[UOp, ...], cnt:int=1) -> UOp:
ends_to_add = [z for z in new_ctx if z not in current_ctx]
while len(ends_to_add):
r:UOp = ends_to_add.pop(-1)
new_ctx = tuple([z for z in new_ctx if z is not r])
end_uop = UOp(Ops.ENDIF if r.op is Ops.IF else Ops.ENDRANGE, src=(r,))
base_block = UOp(Ops.BLOCKEND, src=(base_block,)*cnt, arg=BasicBlock((end_uop,), tuple(new_ctx), end=r, cnt=cnt))
return base_block
def make_block_bottom_up(ctx:BlockContext, x:UOp):
if x.op is Ops.BLOCKSTART:
current_ctx, child_ctx = x.arg
lst = list(x.src)
child_count = 1
else:
current_ctx, child_count, child_ctx = ctx.block_ctxs[x], ctx.child_count[x], ctx.child_ctxs.get(x, None)
lst = [x]
# count of times we've seen this block, or a seed for a new block if we can't merge it
unmergable: defaultdict[UOp, int] = defaultdict(int)
blockseeds = defaultdict(list)
# add the srcs of this to the frontier
# NOTE: things may be in here multiple times, that's okay
frontier_nodes = list(flatten(y.src[::-1] for y in lst))
while len(frontier_nodes):
u = frontier_nodes.pop(0)
if u.op not in DONT_PLACE_IN_BLOCK and ctx.child_count[u] == unmergable[u]+1:
# count is correct
if (newctx:=ctx.block_ctxs[u]) == current_ctx:
# block has same context, merge it, and put the srcs on the frontier
lst.append(u)
frontier_nodes.extend(u.src[::-1])
else:
# block has different context, add it to blockseeds
blockseeds[(newctx, ctx.child_ctxs.get(u, None))].append(u)
del unmergable[u]
else:
# count is incorrect (or it's DONT_PLACE_IN_BLOCK), add it to unmergable
unmergable[u] += 1
# add unmergables to sources
srcs = []
for u,cnt in unmergable.items(): srcs += [add_blockends(u, ctx.block_ctxs.get(u,()), current_ctx, cnt=cnt)]*cnt
# add blockseeds, with blockends as needed
for (new_ctx, new_child_ctx), v in blockseeds.items():
base_block = UOp(Ops.BLOCKSTART, src=tuple(v), arg=(new_ctx, new_child_ctx))
srcs.append(add_blockends(base_block, new_ctx, current_ctx))
lst = lst[::-1]
if BLOCK_REORDER: lst = block_reorder(lst)
bb = BasicBlock(tuple(lst), ctx=current_ctx, cnt=child_count, child_ctx=child_ctx)
return UOp(Ops.BLOCK, src=tuple(srcs), arg=bb)
# we prevent the source of the SPECIAL from being linearized since its not part of the kernel
def raise_bottom_up_gate(): raise BottomUpGate()
block_create = PatternMatcher([
(UPat(GroupOp.All-DONT_PLACE_IN_BLOCK.union({Ops.BLOCK, Ops.BLOCKEND}), name="x"), make_block_bottom_up),
(UPat(Ops.SPECIAL), raise_bottom_up_gate)
])
# ***** blockend merging ****
def merge_blockends(sink:UOp) -> UOp|None:
# only run on the final BLOCK with the SINK in it
if sink.arg.lst[-1].op is not Ops.SINK: return None
# combine matching BLOCKENDS, the keys of this dictionary are the RANGE UOps, values are the BLOCKENDs
blockends_to_arg: dict[UOp, list[UOp]] = {}
for be in sink.toposort():
if be.op is Ops.BLOCKEND: blockends_to_arg.setdefault(be.arg.end, []).append(be)
new_forks = {}
for k,v in blockends_to_arg.items():
# NOTE: if any BLOCKEND is the parent of any other with the same arg, this algo fails
if len(v) > 1:
bb = BasicBlock(v[0].arg.lst, _sort_ctx(flatten([y.arg.ctx for y in v])), k, cnt=sum(y.arg.cnt for y in v))
out = UOp(Ops.BLOCKEND, src=tuple(flatten([x.src for x in v])), arg=bb)
# NOTE: bb.ctx != u.arg.ctx can cause problems here
for u in v: new_forks[u] = out
if len(new_forks) == 0: return None
return sink.substitute(new_forks)
pm_blockend_merge = PatternMatcher([(UPat(Ops.BLOCK, name="sink"), merge_blockends)])
# ***** block merging ****
def merge_block(x:UOp):
unmergable_blocks, mergable_blocks = [], []
mergable_dict: defaultdict[UOp, int] = defaultdict(int)
for y in x.src:
if y.op is Ops.BLOCK and x.op is Ops.BLOCK and x.arg.ctx == y.arg.ctx: mergable_dict[y] += 1
elif y.op is Ops.BLOCK and x.op is Ops.BLOCKEND and x.arg.end in y.arg.ctx: mergable_dict[y] += 1
else: unmergable_blocks.append(y)
for k,v in mergable_dict.items():
if v == k.arg.cnt: mergable_blocks.append(k)
else: unmergable_blocks.extend([k]*v)
if len(mergable_blocks) == 0: return None
del mergable_dict
# create the block
arg = replace(x.arg, lst=tuple(flatten([y.arg.lst for y in mergable_blocks]))+x.arg.lst)
return UOp(x.op, src=tuple(flatten([y.src for y in mergable_blocks])+unmergable_blocks), arg=arg)
def remove_blockend(x:UOp):
# if there's any remaining blocks that need to go in this BLOCKEND, we don't remove it
if any(x.arg.end in y.arg.ctx for y in x.src if y.op in {Ops.BLOCK, Ops.BLOCKEND}): return None
if (parent_blocks := [y for y in x.src if y.op is Ops.BLOCK and y.arg.child_ctx is not None and x.arg.end in y.arg.child_ctx]):
assert all_same(parent_blocks), f"should never have two parent blocks (has {len(parent_blocks)})"
parent_block = parent_blocks[0]
assert len(parent_blocks) == parent_block.arg.cnt
# NOTE: DEFINE_ACC doesn't have to be handled in any special way
late_ops = list(x.arg.lst)
# NOTE: we have to add a barrier at the start if barrier is used in the range
if x.op is Ops.BLOCKEND and any(y.op is Ops.BARRIER for y in late_ops) and late_ops[-1].op is Ops.ENDRANGE:
late_ops = [UOp(Ops.BARRIER)] + late_ops
# peephole opt, remove any BARRIERs next to each other
for i in range(len(late_ops)-1):
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),
(UPat(Ops.BLOCKEND, name="x"), remove_blockend),
])
# ****** finalize ******
def finalize(sink:UOp) -> UOp:
if sink.op is not Ops.BLOCK or not all(x.op in DONT_PLACE_IN_BLOCK for x in sink.src):
raise RuntimeError(f"linearize failure {sink.op} {[x.op for x in sink.src if x.op not in DONT_PLACE_IN_BLOCK]}")
# place the early things
lst = sorted(dedup(sink.src), key=lambda x: x.tuplize) + list(sink.arg.lst)
return UOp(Ops.BLOCKFINAL, arg=BasicBlock(tuple(lst)))
pm_finalize = PatternMatcher([(UPat(Ops.BLOCK, name="sink"), finalize)])
+86
View File
@@ -0,0 +1,86 @@
# the job of the lowerer is to do indexing
from dataclasses import dataclass
from tinygrad.uop.ops import KernelInfo, UOp, Ops, PatternMatcher, UPat, sint_to_uop, AxisType, graph_rewrite, resolve
# ***** indexing *****
@dataclass
class IndexContext:
axis_types: tuple[AxisType, ...]
idxs: list[UOp]
start: int = 0
def shape_to_idx(s, axis_types, start=0):
return [UOp.range(sint_to_uop(s), start+i, at) for i, (s, at) in enumerate(zip(s, axis_types))]
def get_index(ast:UOp) -> IndexContext:
axis_types = ast.arg.axis_types if isinstance(ast.arg, KernelInfo) else ()
if len(ast.full_shape) != len(axis_types) and ast.st is not None:
axis_types = tuple([AxisType.REDUCE if resolve(s != fs) else AxisType.LOOP for s,fs in zip(ast.shape, ast.full_shape)])
return IndexContext(axis_types, [], 0)
# ***** lowering (given index) *****
def subblock(ctx: IndexContext, full_new_idx: list[UOp], src: UOp):
lc = IndexContext(ctx.axis_types, full_new_idx, ctx.start+1000)
ctx.start = lc.start
return graph_rewrite(src, pm_lowerer, lc, name="subblock", bottom_up=True)
def lower_reduce_axis(ctx: IndexContext, x: UOp):
new_idxs = shape_to_idx(x.src[0].shape, ctx.axis_types, ctx.start)
full_new_idx = list(ctx.idxs)
for a in x.axis_arg: full_new_idx[a] = new_idxs[a]
ret = subblock(ctx, full_new_idx, x.src[0])
return UOp(Ops.REDUCE, x.dtype, (ret,)+tuple([full_new_idx[i] for i in x.axis_arg]), x.arg[0])
def lower_store(ctx: IndexContext, x: UOp, buf: UOp):
# TODO: reenable after REDUCE_AXIS is fixed
#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 = 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])
else: real_new_idxs.append(ctx.idxs[i])
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).store(stored, *used_ranges)
def fixup_wmma(ctx:IndexContext, x:UOp):
if x.tag is not None: return None
new_idxs = shape_to_idx(x.src[0].shape, ctx.axis_types, ctx.start)
full_new_idx = list(ctx.idxs)
for a in x.arg[-1]: full_new_idx[a] = new_idxs[a]
srcs = subblock(ctx, full_new_idx, UOp.sink(*x.src)).src
# NOTE: this assumes these are expanded. which now shouldn't change anything
new_x_arg_m2 = tuple([tuple([(full_new_idx[a].arg[0], sz) for a,sz in v]) for v in x.arg[-2]])
new_x_arg_m1 = tuple([full_new_idx[a].arg[0] for a in x.arg[-1]])
return x.replace(src=srcs, arg=x.arg[:-2]+(new_x_arg_m2, new_x_arg_m1), tag=1)
pm_lowerer = PatternMatcher([
# TODO: remove these hacks
# hack for old style CONST(VIEW) (now it's just VIEW(CONST))
(UPat((Ops.DEFINE_VAR, Ops.CONST), src=(UPat(Ops.VIEW, name="v"),), name="c"), lambda c,v: c.replace(src=()).view(v.arg)),
# hack for old style VALID (now it's just VIEW(CONST))
(UPat(Ops.VALID, src=(UPat(Ops.VIEW, name="v"),)).where(UPat.cvar("c"), UPat(Ops.CONST, arg=0)), lambda c,v: c.replace(src=()).view(v.arg)),
# 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_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_valid_uop(ctx.idxs)),)+x.src[1:])),
# reduce/view_const
(UPat(Ops.REDUCE_AXIS, name="x"), lower_reduce_axis),
(UPat(Ops.STORE, src=(UPat.var("buf").view(),), allow_any_len=True, name="x"), lower_store),
(UPat(Ops.WMMA, name="x"), fixup_wmma),
# axis fixups for WMMA
(UPat((Ops.CONTRACT, Ops.UNROLL), name="x"),
lambda ctx,x: x.replace(tag=1, arg=tuple([(ctx.idxs[a].arg[0], sz) for a,sz in x.arg])) if x.tag is None else None),
])
+5 -5
View File
@@ -51,7 +51,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
# upcast float4 images, this must be early so we don't accidentally add locals before the upcast
for buf_index,buf in enumerate(k.bufs):
if isinstance(buf.src[0].dtype, ImageDType):
# part of is_expanded
# part of real_strides
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):
@@ -96,7 +96,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(any(o is k.rngs[axis] for o in u.src[0].backward_slice) for u in k.ast.backward_slice 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)
@@ -112,12 +112,12 @@ 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].get_idx().backward_slice and all(r2 in b.src[1].get_idx().backward_slice
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:
idx = b.src[1].get_idx()
if rng in idx.backward_slice: num_strides += 1
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
@@ -160,7 +160,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].get_idx().backward_slice 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])):
+8 -8
View File
@@ -2,7 +2,7 @@ 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, GroupOp
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, merge_dicts, get_single_element
@@ -25,7 +25,7 @@ class Scheduler:
@property
def rngs(self):
# always in order by axistype
return sorted([u for u in self.ast.backward_slice if u.op is Ops.RANGE and u.vmax > 0], key=lambda x: (axis_to_pos[x.arg[-1]],) + x.arg[0:-1])
return sorted([u for u in self.ast.parents if u.op is Ops.RANGE and u.vmax > 0], key=lambda x: (axis_to_pos[x.arg[-1]],) + x.arg[0:-1])
@property
def shape_len(self): return len(self.rngs)
@property
@@ -149,7 +149,7 @@ class Scheduler:
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.backward_slice if u.op is Ops.REDUCE and rng in merge_dicts([r.ranges for r in u.src[1:]])][0]
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")
@@ -188,14 +188,14 @@ class Scheduler:
check(rng.arg[-1] is not AxisType.THREAD, "cannot pad thread")
# ok to pad SUM if all parent ALU ops have f(0) = 0
if (r:=self.reduceop) is not None and rng.arg[-1] in (AxisType.GROUP_REDUCE, AxisType.REDUCE):
check(r.arg[0] is Ops.ADD and not r.op_in_backward_slice_with_self(*GroupOp.UnsafePad), f"cannot pad {r}")
check(r.arg[0] is Ops.ADD and can_pad(r, {}), f"cannot pad {r}")
new_sz = round_up(int(rng.vmax+1), cast(int, opt.arg))
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 (i:=b.src[1].get_idx()).backward_slice_with_self:
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:
@@ -310,7 +310,7 @@ class Scheduler:
# helpers for hand_coded_optimizations
@property
def reduceop(self) -> UOp|None:
red = [x for x in self.ast.backward_slice if x.op is Ops.REDUCE]
red = [x for x in self.ast.parents if x.op is Ops.REDUCE]
if not len(red): return None
return UOp(Ops.REDUCE_AXIS, red[0].dtype, red[0].src, (red[0].arg, ()))
@property
@@ -324,7 +324,7 @@ class Scheduler:
def group_for_reduces(self) -> int: return len(self.axes_of(AxisType.GROUP_REDUCE))
def bufs_from_ast(ast:UOp, dname:str) -> list[Buffer]:
glbls = sorted([x for x in ast.backward_slice if x.op is Ops.DEFINE_GLOBAL], key=lambda x: x.arg)
glbls = sorted([x for x in ast.parents if x.op is Ops.DEFINE_GLOBAL], key=lambda x: x.arg)
return [Buffer(dname, x.ptrdtype.size, x.dtype.base if not isinstance(x.dtype, ImageDType) else x.dtype) for x in glbls]
def apply_opts(ctx:Renderer, ast:UOp):
@@ -340,7 +340,7 @@ def apply_opts(ctx:Renderer, ast:UOp):
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.backward_slice if u.op is Ops.LOAD):
if all(len(u.src) == 1 for u in ast.parents if u.op is Ops.LOAD):
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)
+135
View File
@@ -0,0 +1,135 @@
from tinygrad.uop.ops import UOp, Ops, GroupOp, PatternMatcher, UPat, graph_rewrite, resolve, sint
from tinygrad.helpers import all_same, prod, unwrap, colored
from tinygrad.shape.shapetracker import ShapeTracker
from tinygrad.shape.view import View, strides_for_shape, get_contraction_with_reduce
from tinygrad.schedule.grouper import ALWAYS_CONTIGUOUS
from tinygrad.dtype import ImageDType, dtypes
merge_views = PatternMatcher([
# merge adjacent views
(UPat(Ops.VIEW, src=(UPat(Ops.VIEW, name="v1"),), name="v2"), lambda v1,v2: v1.replace(arg=v1.arg+v2.arg)),
# replace MovementOps with VIEW
(UPat(GroupOp.Movement, src=(UPat.var("x"),), name="mop"), lambda mop,x: x.base.view(mop.st)),
# remove NOOP views
(UPat.var("x").view(name="view"),
lambda x,view: x if x.st is not None and x.op not in GroupOp.Defines and view.st.contiguous and view.shape == x.shape else None),
(UPat(GroupOp.All-{Ops.DEFINE_GLOBAL}).view(name="view"),
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=(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):
# contiguous, expand, and the same with ones removed
if unwrap(view.st).contiguous and len(r.shape) < len(view.shape) and \
tuple(x for x in r.shape if resolve(x != 1)) == tuple(x for x in view.shape if resolve(x != 1)):
new_shape: list[sint] = []
new_reduce_axis = []
if (contraction:=get_contraction_with_reduce(view.shape, r.shape, r.arg[1])) is None: return None
for i,pairs in enumerate(contraction):
new_shape_chunk = [view.shape[p] for p in pairs]
if i in r.arg[1]:
# if this is a reduce axis, we need a 1 in the view here to put it
assert len(new_shape_chunk) > 0
new_shape += [1]*(len(pairs)-1) + [src.shape[i]]
new_reduce_axis.append(len(new_shape)-1)
else:
# otherwise, pass through the new_shape_chunk
new_shape += new_shape_chunk
ret = r.replace(src=(src.reshape(tuple(new_shape)),), arg=(r.arg[0], tuple(new_reduce_axis))+r.arg[2:])
assert ret.shape == view.shape, f"shape mismatch on reduce_push_add_ones, {ret.shape} != {view.shape}"
return ret
return None
view_left = merge_views+PatternMatcher([
# view before elementwise and buffer ops
(UPat(Ops.VIEW, src=(UPat({*GroupOp.ALU, Ops.CAST, Ops.BITCAST, Ops.BIND, Ops.STORE, Ops.VALID, Ops.SINK}, name="e"),), name="view"),
lambda e,view: e.replace(src=tuple(s.view(view.st) for s in e.src))),
# if there's ones added after reduce, put this before the reduce
(UPat(Ops.VIEW, src=(UPat(Ops.REDUCE_AXIS, src=(UPat.var("src"),), name="r"),), name="view"), reduce_push_add_ones),
])
view_left_through_load = PatternMatcher([
# view before load
(UPat(Ops.VIEW, src=(UPat(Ops.LOAD, name="e"),), name="view"),
lambda e,view: e.replace(src=tuple(s.view(view.st) for s in e.src))),
])
def apply_swizzle(u:UOp) -> UOp: return graph_rewrite(u, view_left, name="Sub View Left")
# change reduceop axes and input ShapeTrackers, view gets replaced with a reshape.
def swizzle_reduceop(r:UOp, src:UOp, view:UOp, fuse=False):
# contiguous and same size can push to children
# if there's a reduce child, shapes match with ones removed
if unwrap(view.st).contiguous and view.size == r.size and \
(not (len(r.arg) == 3 and r.arg[2]) or # arg[2] = True is fuse marker
tuple((i,x) for i,x in enumerate(r.shape) if resolve(x != 1)) == tuple((i,x) for i,x in enumerate(view.shape) if resolve(x != 1))):
return None
# swizzle the input
input_st = ShapeTracker.from_shape(src.shape)
tmp = input_st.permute(tuple(i for i in range(len(input_st.shape)) if i not in r.axis_arg)+r.axis_arg)
prshape = prod(rshape:=tmp.shape[-len(r.axis_arg):])
strides = strides_for_shape(rshape)
nv = [View.create(v.shape+rshape, tuple(x*prshape for x in v.strides)+strides,
v.offset*prshape, v.mask+tuple((0,s) for s in rshape) if v.mask is not None else None) for v in unwrap(view.st).views]
new_view = tmp + ShapeTracker(tuple(nv))
swizzled_input = apply_swizzle(src.view(new_view))
# create a new reduceop
new_axis = tuple(range(len(view.shape), len(view.shape) + len(r.axis_arg)))
if fuse: red = UOp(Ops.REDUCE_AXIS, r.dtype, (swizzled_input.fuse(),), (r.arg[0], new_axis, True))
else: red = UOp(Ops.REDUCE_AXIS, r.dtype, (swizzled_input,), (r.arg[0], new_axis))
return red.reshape(view.shape)
def reduceop_view_right(src:UOp, v:UOp, r:UOp):
assert unwrap(v.st).contiguous and v.size == src.size, f"can't compute new axis for {src.shape} -> {r.shape}"
new_axis = [i for i,(s,u) in enumerate(zip(src.shape, r.shape)) if s != u]
return src.r(r.arg[0], tuple(new_axis)).reshape(r.shape)
def elementwise_view_right(root:UOp):
if not (swizzles:=[x for x in root.src if x.op is Ops.VIEW and x.base.op not in ALWAYS_CONTIGUOUS]): return None
assert all_same([x.base.size for x in swizzles]), f"swizzle inputs must have the same size {swizzles}"
# place view after applying the elementwise op
new_st = ShapeTracker.from_shape(swizzles[0].base.shape)
new_src = [x.base if x.base.shape==new_st.shape else apply_swizzle(x.view(new_st)) for x in root.src]
# reshape to match downstream shapes
return root.replace(src=tuple(new_src)).reshape(root.shape)
# push VIEW to children
view_right = merge_views+PatternMatcher([
# push a non contiguous ShapeTracker through reduceop
(UPat(Ops.VIEW, src=(UPat(Ops.REDUCE_AXIS, src=(UPat.var("src"),), name="r"),), name="view"), swizzle_reduceop),
# apply view after reduceops
(UPat(Ops.REDUCE_AXIS, src=(UPat(Ops.VIEW, src=(UPat(GroupOp.All-ALWAYS_CONTIGUOUS, name="src"),), name="v"),), name="r"), reduceop_view_right),
# apply view after elementwise ops
(UPat(GroupOp.All-{Ops.SINK, Ops.REDUCE_AXIS}, name="root"), elementwise_view_right),
# merge axes for double reduce (invert of SPLIT_REDUCEOP=1)
(UPat(Ops.REDUCE_AXIS, src=(UPat(Ops.REDUCE_AXIS, name="r1"),), name="r2"),
lambda r1,r2: r1.replace(arg=(r1.arg[0], r2.arg[1]+r1.arg[1])) if r1.arg[0] is r2.arg[0] else None),
# remove view from sink
(UPat(Ops.VIEW, name="v").sink(name="sink"), lambda v,sink: v.src[0].sink(arg=sink.arg)),
])
def check_load_st(glbl:UOp, view:UOp):
if glbl.arg != 0 or (st:=unwrap(view.st)).contiguous: return
# if it has a single view and it becomes contiguous when you shrink expanded axes, it's fine
if len(st.views) == 1 and st.shrink(tuple((0,1) if st == 0 else (0,s) for s,st in zip(st.shape, st.views[0].strides))).contiguous: return
# if it has a single view and it's equal when you shrink a contig, it's fine
if len(st.views) == 1 and (mask:=st.views[0].mask) is not None and ShapeTracker.from_shape(st.shape).shrink(mask) == st.shrink(mask): return
# otherwise, it's not fine
raise RuntimeError("self operand of augmented assign must be contiguous.\nhelp: consider using .contiguous():\n"
+colored(" - a += a.T\n", "red")+colored(" + a += a.T.contiguous()", "green"))
fix_kernel_ops = view_left_through_load+PatternMatcher([
# add view to LOAD and STORE
(UPat(Ops.DEFINE_GLOBAL, name="g").load(), lambda g: g.view(g.st).load()),
(UPat(Ops.DEFINE_GLOBAL, name="g").store(UPat.var('x')), lambda g,x: g.view(g.st).store(x)),
# VALID
(UPat(Ops.VIEW, src=(UPat.cvar(),), name="self"),
lambda self: UOp.where(UOp(Ops.VALID, dtypes.bool, (UOp(Ops.VIEW, arg=self.st),)), self.const_like(self.base.arg), 0)),
# no ImageDType after index
(UPat(GroupOp.All-{Ops.DEFINE_GLOBAL, Ops.VIEW, Ops.INDEX}, name="x"),
lambda x: x.replace(dtype=x.dtype.base) if isinstance(x.dtype, ImageDType) else None),
# if this kernel also assigns to the loaded buffer, ensure we can index it correctly
(UPat(Ops.LOAD, src=(UPat.var("glbl").view(name="view"),)), check_load_st),
])
+7 -7
View File
@@ -27,13 +27,13 @@ pm_quant = symbolic+PatternMatcher([
(UPat.var("x")*UPat.cvar("c1", dtype=dtypes.floats) + UPat.var("y")*UPat.cvar("c2", dtype=dtypes.floats),
lambda x,y,c1,c2: (x+y)*c1 if abs(c1.arg-c2.arg) < 1e-9 else None),
# mul 0 * c1 is 0
#(UPat(Ops.VALID, src=(UPat(Ops.VIEW, name="v"),)).where(UPat.cvar("c1"), UPat(Ops.CONST, arg=0)) *
# UPat(Ops.LOAD, src=(UPat().view(name="v"),)).cast(dtypes.int).cast(dtypes.float).named("ld"), lambda ld,v,c1: ld*c1),
(UPat(Ops.VALID, src=(UPat(Ops.VIEW, name="v"),)).where(UPat.cvar("c1"), UPat(Ops.CONST, arg=0)) *
UPat(Ops.LOAD, src=(UPat().view(name="v"),)).cast(dtypes.int).cast(dtypes.float).named("ld"), lambda ld,v,c1: ld*c1),
# mul (with plus) 0 * c1 is 0
#(UPat(Ops.VALID, src=(UPat(Ops.VIEW, name="v"),)).where(UPat.cvar("c1"), UPat(Ops.CONST, arg=0)) *
# (UPat(Ops.LOAD, src=(UPat().view(name="v"),)).cast(dtypes.int) + \
# UPat(Ops.VALID, src=(UPat(Ops.VIEW, name="v"),)).where(UPat.cvar(), UPat(Ops.CONST, arg=0))).cast(dtypes.float).named("ld"),
# lambda ld,v,c1: ld*c1),
(UPat(Ops.VALID, src=(UPat(Ops.VIEW, name="v"),)).where(UPat.cvar("c1"), UPat(Ops.CONST, arg=0)) *
(UPat(Ops.LOAD, src=(UPat().view(name="v"),)).cast(dtypes.int) + \
UPat(Ops.VALID, src=(UPat(Ops.VIEW, name="v"),)).where(UPat.cvar(), UPat(Ops.CONST, arg=0))).cast(dtypes.float).named("ld"),
lambda ld,v,c1: ld*c1),
# const push through add
((UPat.var("x")*UPat.cvar("c1") + UPat.var("y")*UPat.cvar("c2")) * UPat.cvar("c3"), lambda x,y,c1,c2,c3: (x*c1*c3) + (y*c2*c3)),
@@ -64,4 +64,4 @@ pm_quant = symbolic+PatternMatcher([
lambda v1,v2,c1,r: r.replace(src=(v1*v2,)) + r.replace(src=(c1*v2,))),
(UPat(Ops.REDUCE_AXIS, src=((UPat(Ops.CAST, name="v1")+UPat.var("c1")) * (UPat(Ops.CAST, name="v2",)+UPat.var("c2")),), name="r"),
lambda v1,v2,c1,c2,r: r.replace(src=(v1*v2,)) + r.replace(src=(c2*v1,)) + r.replace(src=(c1*v2,)) + r.replace(src=(c1*c2,))),
])
])
+4 -5
View File
@@ -17,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:
reduce_ranges = [x.ranges for x in u.backward_slice_with_self if x.op is Ops.REDUCE]
reduce_ranges = [x.ranges for x in u.sparents if x.op is Ops.REDUCE]
i = range_start[u.op]
while i < len(u.src)-1:
r0, r1 = u.src[i], u.src[i+1]
@@ -67,7 +67,7 @@ pm_split_ranges = PatternMatcher([
# **** reduce simplification ****
def no_range(u:UOp) -> bool: return not any(x.op is Ops.RANGE for x in u.backward_slice_with_self)
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
@@ -116,7 +116,7 @@ pm_reduce_collapse = PatternMatcher([
])+sym
def reduce_collapse(red:UOp):
included, not_included = partition(red.backward_slice, lambda x: any(y in x.backward_slice_with_self for y in red.src[1:]))
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:
@@ -129,8 +129,7 @@ def reduce_collapse(red:UOp):
def reduce_unparented(red:UOp):
if red.arg not in {Ops.ADD, Ops.MAX, Ops.MUL}: return None
assert all(x.op is Ops.RANGE for x in red.src[1:]), "some reduce srcs aren't ranges"
reduce_parented, reduce_unparented = partition(red.src[1:], lambda x: x in red.src[0].ranges)
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:
+1 -1
View File
@@ -183,7 +183,7 @@ class dtypes:
uints = (uint8, uint16, uint32, uint64)
sints = (int8, int16, int32, int64)
ints = uints + sints
all = floats + ints + (bool, index) # noqa: A003
all = floats + ints + (bool, index)
if (env_default_float := getenv("DEFAULT_FLOAT", "")):
dtypes.default_float = getattr(dtypes, env_default_float.lower())
+3 -2
View File
@@ -85,7 +85,7 @@ def word_wrap(x, wrap=80):
def suppress_finalizing(func):
def wrapper(*args, **kwargs):
try: return func(*args, **kwargs)
except (RuntimeError, AttributeError, TypeError, ImportError):
except (AttributeError, TypeError, ImportError):
if not getattr(sys, 'is_finalizing', lambda: True)(): raise # re-raise if not finalizing
return wrapper
@@ -133,6 +133,7 @@ JIT, JIT_BATCH_SIZE = ContextVar("JIT", 2 if OSX and ARCH_X86 else 1), ContextVa
WINO, CAPTURING, TRACEMETA = ContextVar("WINO", 0), ContextVar("CAPTURING", 1), ContextVar("TRACEMETA", 1)
USE_TC, TC_SELECT, TC_OPT, AMX = ContextVar("TC", 1), ContextVar("TC_SELECT", -1), ContextVar("TC_OPT", 0), ContextVar("AMX", 0)
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, LRU = ContextVar("PICKLE_BUFFERS", 1), ContextVar("LRU", 1)
CACHELEVEL, IGNORE_BEAM_CACHE, DEVECTORIZE = ContextVar("CACHELEVEL", 2), ContextVar("IGNORE_BEAM_CACHE", 0), ContextVar("DEVECTORIZE", 1)
@@ -141,7 +142,7 @@ DONT_REALIZE_EXPAND, DONT_GROUP_REDUCES = ContextVar("DONT_REALIZE_EXPAND", 0),
QUANTIZE, VALIDATE_WITH_CPU, DISABLE_FAST_IDIV = ContextVar("QUANTIZE", 0), ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("DISABLE_FAST_IDIV", 0)
CORRECT_DIVMOD_FOLDING, FUSE_OPTIM = ContextVar("CORRECT_DIVMOD_FOLDING", 0), ContextVar("FUSE_OPTIM", 0)
ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE = ContextVar("ALLOW_DEVICE_USAGE", 1), ContextVar("MAX_BUFFER_SIZE", 0)
FUSE_ATTENTION = ContextVar("FUSE_ATTENTION", 0)
RANGEIFY, FUSE_ATTENTION = ContextVar("RANGEIFY", 0), ContextVar("FUSE_ATTENTION", 0)
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)
+1 -1
View File
@@ -115,7 +115,7 @@ string_rewrite = PatternMatcher([
if x.dtype.count > 1 else f"ld.{mem_type(x)}.{ctx.mem_types[x.dtype]} {ctx.r[x]}, [{ctx.r[loc]}+0];"),
(UPat(Ops.DEFINE_REG, src=()), lambda ctx: []),
(UPat(Ops.RANGE, name="x"), lambda ctx, x: [f"mov.u32 {ctx.r[x]}, 0;", "LOOP_" + f"{ctx.r[x][1:]}:"]),
(UPat(Ops.ENDRANGE, name="x", src=(UPat.var("src0"),), allow_any_len=True), lambda ctx, x, src0: [
(UPat(Ops.ENDRANGE, name="x", src=(UPat.var("src0"),)), lambda ctx, x, src0: [
ctx.code_for_op[Ops.ADD](ctx.r[src0], ctx.r[src0], "1", dtypes.int, ctx.types[dtypes.int]),
ctx.code_for_op[Ops.CMPLT](ctx.r[x], ctx.r[x.src[0]], ctx.r[src0.src[0]], dtypes.int, ctx.types[dtypes.int]),
f"@{ctx.r[x]} bra LOOP_{ctx.r[src0][1:]};"]),
+26 -26
View File
@@ -1,13 +1,13 @@
from __future__ import annotations
from typing import cast, ClassVar
import os, ctypes, struct, hashlib, functools, importlib, mmap, errno, array, contextlib, sys, weakref, itertools
import os, ctypes, ctypes.util, struct, hashlib, functools, importlib, mmap, errno, array, contextlib, sys, weakref
assert sys.platform != 'win32'
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, CompilerPairT
from tinygrad.helpers import getenv, to_mv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, suppress_finalizing, lo32, hi32, colored
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
@@ -27,9 +27,6 @@ WAIT_REG_MEM_FUNCTION_GEQ = 5 # >=
AQL_HDR = (1 << hsa.HSA_PACKET_HEADER_BARRIER) | (hsa.HSA_FENCE_SCOPE_SYSTEM << hsa.HSA_PACKET_HEADER_SCACQUIRE_FENCE_SCOPE) \
| (hsa.HSA_FENCE_SCOPE_SYSTEM << hsa.HSA_PACKET_HEADER_SCRELEASE_FENCE_SCOPE)
@dataclass(frozen=True)
class ProfileSQTTEvent(ProfileEvent): device:str; se:int; props:dict; blob:bytes; itrace:bool # noqa: E702
class AMDSignal(HCQSignal):
def __init__(self, *args, **kwargs): super().__init__(*args, **{**kwargs, 'timestamp_divider': 100})
@@ -150,14 +147,14 @@ class AMDComputeQueue(HWQueue):
# be dispatched on something else and not be seen in instruction tracing tab. You can force the wavefronts of a kernel to be dispatched on the
# CUs you want to by disabling other CUs via bits in regCOMPUTE_STATIC_THREAD_MGMT_SE<x> and trace even kernels that only have one wavefront.
self.wreg(self.gc.regSQ_THREAD_TRACE_MASK, wtype_include=self.soc.SQ_TT_WTYPE_INCLUDE_CS_BIT, simd_sel=0, wgp_sel=0, sa_sel=0)
reg_include = self.soc.SQ_TT_TOKEN_MASK_SQDEC_BIT | self.soc.SQ_TT_TOKEN_MASK_SHDEC_BIT | self.soc.SQ_TT_TOKEN_MASK_GFXUDEC_BIT | \
REG_INCLUDE = self.soc.SQ_TT_TOKEN_MASK_SQDEC_BIT | self.soc.SQ_TT_TOKEN_MASK_SHDEC_BIT | self.soc.SQ_TT_TOKEN_MASK_GFXUDEC_BIT | \
self.soc.SQ_TT_TOKEN_MASK_COMP_BIT | self.soc.SQ_TT_TOKEN_MASK_CONTEXT_BIT | self.soc.SQ_TT_TOKEN_MASK_CONTEXT_BIT
token_exclude = 1 << self.soc.SQ_TT_TOKEN_EXCLUDE_PERF_SHIFT
TOKEN_EXCLUDE = 1 << self.soc.SQ_TT_TOKEN_EXCLUDE_PERF_SHIFT
if not (se_mask >> se) & 0b1:
token_exclude |= 1 << self.soc.SQ_TT_TOKEN_EXCLUDE_VMEMEXEC_SHIFT | 1 << self.soc.SQ_TT_TOKEN_EXCLUDE_ALUEXEC_SHIFT | \
TOKEN_EXCLUDE |= 1 << self.soc.SQ_TT_TOKEN_EXCLUDE_VMEMEXEC_SHIFT | 1 << self.soc.SQ_TT_TOKEN_EXCLUDE_ALUEXEC_SHIFT | \
1 << self.soc.SQ_TT_TOKEN_EXCLUDE_VALUINST_SHIFT | 1 << self.soc.SQ_TT_TOKEN_EXCLUDE_IMMEDIATE_SHIFT | \
1 << self.soc.SQ_TT_TOKEN_EXCLUDE_INST_SHIFT
self.wreg(self.gc.regSQ_THREAD_TRACE_TOKEN_MASK, reg_include=reg_include, token_exclude=token_exclude, bop_events_token_include=1)
self.wreg(self.gc.regSQ_THREAD_TRACE_TOKEN_MASK, reg_include=REG_INCLUDE, token_exclude=TOKEN_EXCLUDE, bop_events_token_include=1)
# Enable SQTT
self.sqtt_config(tracing=True)
# Restore global broadcasting
@@ -204,7 +201,9 @@ class AMDComputeQueue(HWQueue):
self.sqtt_userdata(sqtt.struct_rgp_sqtt_marker_event(
_0=sqtt.union_rgp_sqtt_marker_event_0(_0=sqtt.struct_rgp_sqtt_marker_event_0_0(has_thread_dims=1)),
_2=sqtt.union_rgp_sqtt_marker_event_2(cmd_id=next(prg.dev.sqtt_next_cmd_id))), *global_size)
_2=sqtt.union_rgp_sqtt_marker_event_2(cmd_id=prg.dev.cmd_id)), *global_size)
prg.dev.cmd_id += 1
def exec(self, prg:AMDProgram, args_state:CLikeArgsState, global_size:tuple[sint, ...], local_size:tuple[sint, ...]):
self.bind_args_state(args_state)
@@ -213,7 +212,7 @@ class AMDComputeQueue(HWQueue):
user_regs = []
if prg.enable_private_segment_sgpr:
assert self.dev.xccs == 1, "Only architected flat scratch is supported on multi-xcc"
assert self.dev.xccs == 1, "Only architected flat scratch is suppored on multi-xcc"
scratch_hilo = data64_le(prg.dev.scratch.va_addr)
# sgpr word1 bit31 enables swizzle
# sgpr word3 = 0x14 << 12 | 2 << 28 | 2 << 21 | 1 << 23
@@ -255,8 +254,8 @@ class AMDComputeQueue(HWQueue):
self.wreg(self.gc.regCOMPUTE_START_X, 0, 0, 0, *local_size, 0, 0)
gfx10p = {'cs_w32_en': int(prg.wave32)} if prg.dev.target >= (10,0,0) else {}
self.pkt3(self.pm4.PACKET3_DISPATCH_DIRECT, *global_size,
self.gc.regCOMPUTE_DISPATCH_INITIATOR.encode(**gfx10p, force_start_at_000=1, compute_shader_en=1))
DISPATCH_INITIATOR = self.gc.regCOMPUTE_DISPATCH_INITIATOR.encode(**gfx10p, force_start_at_000=1, compute_shader_en=1)
self.pkt3(self.pm4.PACKET3_DISPATCH_DIRECT, *global_size, DISPATCH_INITIATOR)
if prg.dev.sqtt_enabled: self.pkt3(self.pm4.PACKET3_EVENT_WRITE, self.pm4.EVENT_TYPE(self.soc.THREAD_TRACE_MARKER) | self.pm4.EVENT_INDEX(0))
self.pkt3(self.pm4.PACKET3_EVENT_WRITE, self.pm4.EVENT_TYPE(self.soc.CS_PARTIAL_FLUSH) | self.pm4.EVENT_INDEX(EVENT_INDEX_PARTIAL_FLUSH))
@@ -500,6 +499,9 @@ class AMDAllocator(HCQAllocator['AMDDevice']):
def _map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
@dataclass(frozen=True)
class ProfileSQTTEvent(ProfileEvent): device:str; se:int; blob:bytes; itrace:bool # noqa: E702
@dataclass
class AMDQueueDesc:
ring: MMIOInterface
@@ -555,9 +557,7 @@ class KFDIface:
for i in FileIOInterface(f'{ip_base}/{hw}').listdir()} for ip,hw in ip_hw }
self.drm_fd = FileIOInterface(f"/dev/dri/renderD{self.props['drm_render_minor']}", os.O_RDWR)
self.kfd_ver = ((ver_st:=kfd.AMDKFD_IOC_GET_VERSION(KFDIface.kfd)).major_version, ver_st.minor_version)
kfd.AMDKFD_IOC_ACQUIRE_VM(KFDIface.kfd, drm_fd=self.drm_fd.fd, gpu_id=self.gpu_id)
if self.kfd_ver >= (1,14): kfd.AMDKFD_IOC_RUNTIME_ENABLE(KFDIface.kfd, mode_mask=0)
# Set these for our device.
if KFDIface.event_page is None:
@@ -749,10 +749,9 @@ class AMDDevice(HCQCompiled):
if self.target < (9,4,2) or self.target >= (13,0,0): raise RuntimeError(f"Unsupported arch: {self.arch}")
if DEBUG >= 1: print(f"AMDDevice: opening {self.device_id} with target {self.target} arch {self.arch}")
self.se_cnt = self.iface.props['array_count'] // self.iface.props['simd_arrays_per_engine']
self.max_cu_id = self.iface.props['simd_count'] // self.iface.props['simd_per_cu'] // self.iface.props.get('num_xcc', 1) - 1
self.max_wave_id = (self.iface.props['max_waves_per_simd'] * self.iface.props['simd_per_cu'] - 1) if self.target >= (10,1,0) else \
(min((self.max_cu_id+1)*40, self.se_cnt * 512) - 1)
(min((self.max_cu_id+1)*40, self.iface.props['array_count'] // self.iface.props['simd_arrays_per_engine'] * 512) - 1)
self.xccs = self.iface.props.get('num_xcc', 1) if getenv("XCCS", 1) else 1
# this is what llvm refers to as "architected flat scratch"
self.has_scratch_base_registers = self.target >= (11,0,0) or self.target in {(9,4,2), (9,5,0)}
@@ -804,15 +803,16 @@ class AMDDevice(HCQCompiled):
# SQTT is disabled by default because of runtime overhead and big file sizes (~200mb to Tensor.full() two 4096x4096 tensors and matmul them)
self.sqtt_enabled = PROFILE and bool(getenv("SQTT", 0))
if self.sqtt_enabled:
if self.target[0] != 11: raise RuntimeError(f'SQ Thread Tracing is not supported on gc:{self.target}')
if self.arch != 'gfx1100': raise RuntimeError('SQ Thread Tracing is only supported on 7900XTX')
if not self.is_am() and (ppfeaturemask:=int(FileIOInterface('/sys/module/amdgpu/parameters/ppfeaturemask', os.O_RDONLY).read(), 16))&0x8000:
raise RuntimeError("SQTT can't be enabled because of hardware bug, to workaround either use AMD_IFACE=PCI or add "
f"ppfeaturemask={(ppfeaturemask&~0x8000):#x} (current {ppfeaturemask=:#x} & ~PP_GFXOFF_MASK) to amdgpu module parameters\n"
"For more information read https://github.com/tinygrad/tinygrad/blob/master/extra/sqtt/README.md")
SQTT_BUFFER_SIZE = getenv("SQTT_BUFFER_SIZE", 256) # in mb, per shader engine
self.sqtt_buffers = [self.allocator.alloc(SQTT_BUFFER_SIZE*1024*1024, BufferSpec(cpu_access=True, nolru=True)) for _ in range(self.se_cnt)]
SQTT_NUM = self.iface.props['array_count'] // self.iface.props['simd_arrays_per_engine']
self.sqtt_buffers = [self.allocator.alloc(SQTT_BUFFER_SIZE*1024*1024, BufferSpec(cpu_access=True, nolru=True)) for _ in range(SQTT_NUM)]
self.sqtt_itrace_se_mask = getenv("SQTT_ITRACE_SE_MASK", 2) # -1 enable all, 0 disable all, >0 bitmask for where to enable instruction tracing
self.sqtt_next_cmd_id = itertools.count(0)
self.cmd_id = 0
cast(AMDComputeQueue, self.hw_compute_queue_t()).sqtt_start(self.sqtt_buffers, self.sqtt_itrace_se_mask).submit(self)
def create_queue(self, queue_type, ring_size, ctx_save_restore_size=0, eop_buffer_size=0, ctl_stack_size=0, debug_memory_size=0):
@@ -843,9 +843,10 @@ class AMDDevice(HCQCompiled):
scratch_size = (self.max_cu_id+1)*self.iface.props['max_slots_scratch_cu']*wave_scratch_len # per xcc
self.scratch, ok = self._realloc(getattr(self, 'scratch', None), scratch_size*self.xccs)
if ok:
engines = self.iface.props['array_count'] // self.iface.props['simd_arrays_per_engine']
waves = wave_scratch_len // (256 if self.target >= (11,0,0) else 1024)
# >=gfx11 wavesize is per SE
wavesize = scratch_size // ((wave_scratch_len * self.se_cnt) if self.target >= (11,0,0) else wave_scratch_len)
wavesize = scratch_size // ((wave_scratch_len * engines) if self.target >= (11,0,0) else wave_scratch_len)
self.tmpring_size = waves << 12 | wavesize
self.max_private_segment_size = required
@@ -870,14 +871,13 @@ class AMDDevice(HCQCompiled):
cast(AMDComputeQueue, self.hw_compute_queue_t()).sqtt_stop(len(self.sqtt_buffers), wptrs_buf) \
.signal(self.timeline_signal, self.next_timeline()).submit(self)
self.synchronize()
if DEBUG >= 2: print(f'{self.device}: Saving SQTT in profile...')
if DEBUG>=2: print('Saving SQTT in profile...')
for i,buf0 in enumerate(self.sqtt_buffers):
wptr = ((struct.unpack('<I', wptrs[i*4:i*4+4])[0] & 0x1FFFFFFF) - ((buf0.va_addr//32) & 0x1FFFFFFF)) * 32
if DEBUG >= 2: print(f'\t{self.device}: SE {i} blob size {wptr:#x}')
if DEBUG>=2: print(f'Se {i} blob size {wptr:#x}')
assert wptr >= 0 and wptr <= buf0.size, f"{wptr} > {buf0.size}, should never happen"
# When sqtt buffer overflows, wptr stops at the last dword
if wptr >= buf0.size - 32:
print(colored(f"{self.device}: Warning: SQTT buffer is full (SE {i})! Increase SQTT buffer with SQTT_BUFFER_SIZE=X (in MB)", "yellow"))
if wptr >= buf0.size-32: print(f"WARNING: SQTT BUFFER IS FULL (SE {i})! INCREASE SQTT BUFFER SIZE WITH SQTT_BUFFER_SIZE=X (in MB)")
self.allocator._copyout(sqtt_buf:=memoryview(bytearray(wptr)), buf0)
Compiled.profile_events += [ProfileSQTTEvent(self.device, i, self.iface.props, bytes(sqtt_buf), bool((self.sqtt_itrace_se_mask >> i) & 0b1))]
Compiled.profile_events += [ProfileSQTTEvent(self.device, i, bytes(sqtt_buf), bool((self.sqtt_itrace_se_mask >> i) & 0b1))]
super()._at_profile_finalize()
+1 -1
View File
@@ -1,5 +1,5 @@
from __future__ import annotations
import ctypes, functools
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, CompilerPairT
from tinygrad.renderer.cstyle import CUDARenderer
+5 -8
View File
@@ -43,12 +43,12 @@ def fixup_ip_version(ip:str, version:tuple[int, ...]) -> list[tuple[int, ...]]:
return [version, version[:2], version[:2]+(0,), version[:1]+(0, 0)]
def header_download(file, name=None, subdir="defines", url=None) -> str:
url = url or "https://gitlab.com/linux-kernel/linux-next/-/raw/cf6d949a409e09539477d32dbe7c954e4852e744/drivers/gpu/drm/amd"
def header_download(file, name=None, subdir="defines") -> str:
url = "https://gitlab.com/linux-kernel/linux-next/-/raw/cf6d949a409e09539477d32dbe7c954e4852e744/drivers/gpu/drm/amd"
return fetch(f"{url}/{file}", name=name, subdir=subdir).read_text()
def import_header(path:str, url=None):
t = re.sub(r'//.*|/\*.*?\*/','', header_download(path, subdir="defines", url=url), flags=re.S)
def import_header(path:str):
t = re.sub(r'//.*|/\*.*?\*/','', header_download(path, subdir="defines"), flags=re.S)
return {k:int(v,0) for k,v in re.findall(r'\b([A-Za-z_]\w*)\s*=\s*(0x[0-9A-Fa-f]+|\d+)', t)}
def import_module(name:str, version:tuple[int, ...], version_prefix:str=""):
@@ -57,10 +57,7 @@ def import_module(name:str, version:tuple[int, ...], version_prefix:str=""):
except ImportError: pass
raise ImportError(f"Failed to load autogen module for {name.upper()} {'.'.join(map(str, version))}")
def import_soc(ip):
# rocm soc headers have more profiling enums than upstream linux
url = "https://raw.githubusercontent.com/ROCm/rocm-systems/cccc350dc620e61ae2554978b62ab3532dc10bd9/projects"
return type("SOC", (object,), import_header(f"aqlprofile/linux/{({9: 'vega10', 10: 'navi10', 11: 'soc21', 12: 'soc24'}[ip[0]])}_enum.h", url=url))
def import_soc(ip): return type("SOC", (object,), import_header(f"include/{({9: 'vega10', 10: 'navi10', 11: 'soc21', 12: 'soc24'}[ip[0]])}_enum.h"))
def import_asic_regs(prefix:str, version:tuple[int, ...], cls=AMDReg) -> dict[str, AMDReg]:
def _split_name(name): return name[:(pos:=next((i for i,c in enumerate(name) if c.isupper()), len(name)))], name[pos:]
+1 -1
View File
@@ -1,4 +1,4 @@
import subprocess, hashlib, tempfile, ctypes, re, pathlib
import subprocess, hashlib, tempfile, ctypes, ctypes.util, re, pathlib
from typing import Callable
from tinygrad.helpers import to_char_p_p, colored, init_c_var, getenv
import tinygrad.runtime.autogen.nvrtc as nvrtc
+10 -2
View File
@@ -3,11 +3,10 @@ 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, ProfileRangeEvent
from tinygrad.helpers import PROFILE, getenv, to_mv, round_up, ProfileRangeEvent
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
from tinygrad.runtime.support.memory import BumpAllocator
class MMIOInterface:
def __init__(self, addr:int, nbytes:int, fmt='B'): self.mv, self.addr, self.nbytes, self.fmt = to_mv(addr, nbytes).cast(fmt), addr, nbytes, fmt
@@ -63,6 +62,15 @@ ProgramType = TypeVar('ProgramType', bound='HCQProgram')
ArgsStateType = TypeVar('ArgsStateType', bound='HCQArgsState')
QueueType = TypeVar('QueueType', bound='HWQueue')
class BumpAllocator:
def __init__(self, size:int, base:int=0, wrap:bool=True): self.size, self.ptr, self.base, self.wrap = size, 0, base, wrap
def alloc(self, size:int, alignment:int=1) -> int:
if round_up(self.ptr, alignment) + size > self.size:
if not self.wrap: raise RuntimeError("Out of memory")
self.ptr = 0
self.ptr = (res:=round_up(self.ptr, alignment)) + size
return res + self.base
class HWQueue(Generic[SignalType, HCQDeviceType, ProgramType, ArgsStateType]):
"""
A base class for hardware command queues in the HCQ (Hardware Command Queue) API.
+1 -1
View File
@@ -1,4 +1,4 @@
import ctypes.util, os, sys, subprocess
import ctypes, ctypes.util, os, sys, subprocess
from tinygrad.helpers import DEBUG, OSX, getenv
if sys.platform == 'win32':
-9
View File
@@ -2,15 +2,6 @@ import collections, functools, dataclasses
from typing import Any, ClassVar
from tinygrad.helpers import round_up, getenv
class BumpAllocator:
def __init__(self, size:int, base:int=0, wrap:bool=True): self.size, self.ptr, self.base, self.wrap = size, 0, base, wrap
def alloc(self, size:int, alignment:int=1) -> int:
if round_up(self.ptr, alignment) + size > self.size:
if not self.wrap: raise RuntimeError("Out of memory")
self.ptr = 0
self.ptr = (res:=round_up(self.ptr, alignment)) + size
return res + self.base
class TLSFAllocator:
"""
The allocator is based on the Two-Level Segregated Fit (TLSF) algorithm. The allocator maintains 2 level of buckets:
+1 -1
View File
@@ -1,4 +1,4 @@
import ctypes.util, os, subprocess, platform, sysconfig
import ctypes, ctypes.util, os, subprocess, platform, sysconfig
from tinygrad.helpers import OSX
WEBGPU_PATH: str | None
+119
View File
@@ -0,0 +1,119 @@
from tinygrad.uop.ops import Ops, UOp, resolve, can_pad, GroupOp, UPat, PatternMatcher, graph_rewrite
from tinygrad.helpers import all_int, prod, unwrap, dedup, DONT_REALIZE_EXPAND, DONT_GROUP_REDUCES, FUSE_CONV_BW
from tinygrad.shape.shapetracker import ShapeTracker
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,
Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.LOAD}
# **** Grouper decides which of the UOps realize
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
def realize_before_view(ctx:dict[UOp, None], view:UOp, tr:UOp) -> None:
st = unwrap(view.st)
# always realize unsafe pad ops before masked view
if any(v.mask is not None for v in st.views) and not can_pad(tr, ctx): return realize(ctx, tr)
# fold simple pads
if len(st.views) == 1 and (m:=st.views[-1].mask) is not None and all_int(tr.shape) and resolve(prod(tr.shape) >= prod([y-x for x,y in m])): return
# realize before expand
if resolve(prod(tr.shape) < prod(st.shape)) and not DONT_REALIZE_EXPAND: return realize(ctx, tr)
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/CONTIGUOUS/COPY/BUFFER_VIEW
(UPat({Ops.ASSIGN, Ops.CONTIGUOUS, Ops.COPY, Ops.BUFFER_VIEW}, name="tr"), realize),
# realize before expand or unsafe pad ops
(UPat(Ops.VIEW, src=(UPat(GroupOp.All-ALWAYS_CONTIGUOUS, name="tr"),), name="view"), realize_before_view),
# realize parents of COPY, MSELECT, MSTACK
(UPat((Ops.COPY, Ops.MSELECT, Ops.MSTACK), name="rb"), realize_parents),
])
def recursive_group(tr:UOp, st:ShapeTracker, r:UOp, children:dict[UOp, dict[UOp, None]], realizes:dict[UOp, None],
reduce_for_op:dict[UOp, UOp], group:dict[UOp, None], cache:dict[tuple[UOp, ShapeTracker], None]) -> None:
if (tr, st) in cache: return
cache.setdefault((tr, st))
rsize = unwrap(r.st).size
if tr in realizes and tr is not r:
# can only fuse contiguous
# max one reduceop per kernel
if not st.contiguous or st.size != rsize or tr in reduce_for_op: group.setdefault(r)
return group.setdefault(tr)
for tr_next in children.get(tr, {}):
# max one reduceop per kernel
if tr_next.op is Ops.REDUCE_AXIS: return group.setdefault(r)
# can only fuse contiguous
if len(st_childs:=dedup(unwrap(x.st) for x in tr_next.src if x.base == tr)) > 1: return group.setdefault(r)
recursive_group(tr_next, st+st_childs[0], r, children, realizes, reduce_for_op, group, cache)
def group_realizes(sink:UOp) -> dict[UOp, None]:
# start by adding uops that always realize
realizes: dict[UOp, None] = {}
sink = graph_rewrite(sink, do_realize, ctx=realizes, name="do_realize")
if DONT_GROUP_REDUCES: return realizes
# construct children graph (only for bases)
children: dict[UOp, dict[UOp, None]] = {}
assigns: dict[UOp, None] = {}
for u in (toposort:=sink.toposort()):
if u.op in {Ops.VIEW, Ops.SINK}: continue
if u.op is Ops.ASSIGN: assigns[u.buf_uop] = None
for s in u.src: children.setdefault(s.base, {})[u] = None
# find all reduces, and pair them to a elementwise op. if they can't be cleanly paired, force realize the reduce (or a contig child)
reduce_for_op: dict[UOp, UOp] = {}
double_reduces: list[UOp] = []
for r in toposort:
if r.op is not Ops.REDUCE_AXIS: continue
if len(r.arg) == 3 and r.arg[2] is True: continue
if FUSE_CONV_BW and r.src[0].base.op is Ops.REDUCE_AXIS and r.src[0] is not r.src[0].base: double_reduces.append(r)
if r in realizes: continue
group: dict[UOp, None] = {}
recursive_group(r, unwrap(r.st), r, children, realizes, reduce_for_op, group, cache={})
# max one reduceop per kernel
can_chase = all(tr not in reduce_for_op for tr in group)
for u in r.toposort(gate=lambda u: u not in realizes):
if u.op is Ops.REDUCE_AXIS and u.src[0].base.op is Ops.CONST:
can_chase = False
break
# TODO: forced_realize exists because the scheduler is incapable of checking for self-contained DAGs
forced_realize = r in group
# can only have one output
if not forced_realize and len(group) > 1: forced_realize = True
# can only fuse assign if no other assign_target is used in the kernel
if not forced_realize and (assign_targets:={x.buf_uop for x in group if x.op is Ops.ASSIGN}):
parents = [r, *group]
while parents and not forced_realize:
p = parents.pop().base
if p.op is Ops.BUFFER and p in assigns and p not in assign_targets: forced_realize, can_chase = True, False
if p in realizes: continue
parents.extend(p.src)
if forced_realize or not group:
tr = r
if can_chase:
# can chase this down to contiguous children
st = unwrap(tr.st)
while len(lst:=children.get(tr, {})) == 1:
tr_next = next(iter(lst))
st_childs = dedup(unwrap(s.st) for s in tr_next.src if s.base is tr)
if len(st_childs) > 1: break
if st.size != st_childs[0].size: break
st = st + st_childs[0]
if not st.contiguous or tr_next.op is Ops.REDUCE_AXIS: break
tr = tr_next
# don't cast to higher size before store (tr cannot be realized if forced_realize)
if tr.op is Ops.CAST and tr.dtype.itemsize > tr.src[0].dtype.itemsize:
tr = tr.src[0].base
group = {tr: None}
realizes[tr] = None
reduce_for_op.update((tr, r) for tr in group)
# fuse double reduces with no other child
for reduceop in double_reduces:
top_reduce = reduceop.src[0].base
if len(children.get(top_reduce, {})) == 1: del realizes[top_reduce]
return realizes

Some files were not shown because too many files have changed in this diff Show More