diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 578e76cc4b..77d17cc65b 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -310,9 +310,9 @@ jobs: - name: Fuzz Test fast idiv run: python test/external/fuzz_fast_idiv.py - name: Fuzz Test shapetracker - run: | - python test/external/fuzz_shapetracker.py - python test/external/fuzz_shapetracker_math.py + run: CNT=50 python test/external/fuzz_shapetracker.py + - name: Fuzz Test shapetracker math + run: CNT=200 python test/external/fuzz_shapetracker_math.py - name: Fuzz Test shape ops run: python test/external/fuzz_shape_ops.py @@ -377,7 +377,7 @@ 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=190 ALLOWED_READ_IMAGE=2081 ALLOWED_GATED_READ_IMAGE=28 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) diff --git a/.pylintrc b/.pylintrc index 2f1de51927..dc51be94d7 100644 --- a/.pylintrc +++ b/.pylintrc @@ -30,10 +30,6 @@ persistent=yes # Specify a configuration file. #rcfile= -# When enabled, pylint would attempt to guess common misconfiguration and emit -# user-friendly hints instead of false-positive error messages -suggestion-mode=yes - # Allow loading of arbitrary C extensions. Extensions are imported into the # active Python interpreter and may run arbitrary code. unsafe-load-any-extension=no diff --git a/autogen_stubs.sh b/autogen_stubs.sh index 1ea3b583db..4577bbde85 100755 --- a/autogen_stubs.sh +++ b/autogen_stubs.sh @@ -435,8 +435,8 @@ generate_sqtt() { -o extra/sqtt/rocprof/rocprof.py fixup extra/sqtt/rocprof/rocprof.py sed -i '1s/^/# pylint: skip-file\n/' extra/sqtt/rocprof/rocprof.py - sed -i "s/import ctypes/import ctypes\nfrom tinygrad.helpers import fetch/g" extra/sqtt/rocprof/rocprof.py - sed -i "s|FunctionFactoryStub()|ctypes.CDLL(str(fetch('https://github.com/ROCm/rocprof-trace-decoder/raw/5420409ad0963b2d76450add067b9058493ccbd0/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so')))|g" extra/sqtt/rocprof/rocprof.py + sed -i "s/import ctypes/import ctypes, ctypes.util/g" extra/sqtt/rocprof/rocprof.py + sed -i "s|FunctionFactoryStub()|ctypes.CDLL(ctypes.util.find_library('rocprof-trace-decoder'))|g" extra/sqtt/rocprof/rocprof.py } generate_webgpu() { diff --git a/examples/mlperf/model_train.py b/examples/mlperf/model_train.py index 6354155b14..a8189c4ee5 100644 --- a/examples/mlperf/model_train.py +++ b/examples/mlperf/model_train.py @@ -1188,7 +1188,9 @@ def train_bert(): if MLLOGGER and RUNMLPERF: MLLOGGER.start(key=mllog_constants.EVAL_START, value=None, metadata={"epoch_num": i*GBS, "step_num": i}) if getenv("RESET_STEP"): train_step_bert.reset() - elif getenv("FREE_INTERMEDIATE", 1) and train_step_bert.captured is not None: train_step_bert.captured.free_intermediates() + elif getenv("FREE_INTERMEDIATE", 0) and train_step_bert.captured is not None: + # TODO: FREE_INTERMEDIATE nan'ed after jit step 2 + train_step_bert.captured.free_intermediates() eval_lm_losses = [] eval_clsf_losses = [] eval_lm_accs = [] @@ -1222,7 +1224,7 @@ def train_bert(): return if getenv("RESET_STEP"): eval_step_bert.reset() - elif getenv("FREE_INTERMEDIATE", 1) and eval_step_bert.captured is not None: eval_step_bert.captured.free_intermediates() + elif getenv("FREE_INTERMEDIATE", 0) and eval_step_bert.captured is not None: eval_step_bert.captured.free_intermediates() del eval_data avg_lm_loss = sum(eval_lm_losses) / len(eval_lm_losses) diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_1xMI300X/dev_beam.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_1xMI300X/dev_beam.sh new file mode 100755 index 0000000000..68e5fdfcde --- /dev/null +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_1xMI300X/dev_beam.sh @@ -0,0 +1,17 @@ +#!/bin/bash + +export PYTHONPATH="." AMD=1 +export MODEL="bert" +export DEFAULT_FLOAT="HALF" GPUS=1 BS=128 EVAL_BS=128 + +export IGNORE_OOB=1 + +export BEAM=3 BEAM_UOPS_MAX=4000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 +export IGNORE_JIT_FIRST_BEAM=1 +# export BEAM_LOG_SURPASS_MAX=1 +# export BASEDIR="/raid/datasets/wiki" + +export RESET_STEP=1 +export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2 + +python3 examples/mlperf/model_train.py diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/README.md b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/README.md new file mode 100644 index 0000000000..844b90f949 --- /dev/null +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/README.md @@ -0,0 +1,69 @@ +# 1. Problem + +This problem uses BERT for NLP. + +## Requirements + +Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0. +``` +git clone https://github.com/tinygrad/tinygrad.git +python3 -m pip install -e ".[mlperf]" +``` +Also install gdown (for dataset), numpy, tqdm and tensorflow. +``` +pip install gdown numpy tqdm tensorflow +``` + +### tinybox_green +Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md) +This is the default on production tinybox green. + +# 2. Directions + +## Steps to download and verify data + +### 1. Download raw data + +``` +BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py +``` + +### 2. Preprocess train and validation data + +Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended. + +#### Training: +``` +BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all +``` + +Generating a specific topic (Between 0 and 499) +``` +BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42 +``` + +#### Validation: +``` +BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval +``` +## Running + +### tinybox_green + +#### Steps to run benchmark +``` +examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh +``` + +### tinybox_red + +#### Steps to run benchmark +``` +examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh +``` +### tinybox_8xMI300X + +#### Steps to run benchmark +``` +examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh +``` \ No newline at end of file diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/dev_beam.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/dev_beam.sh new file mode 100755 index 0000000000..278eff316d --- /dev/null +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/dev_beam.sh @@ -0,0 +1,17 @@ +#!/bin/bash + +export PYTHONPATH="." AMD=1 +export MODEL="bert" +export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024 +export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1 + +export IGNORE_OOB=1 +export REWRITE_STACK_LIMIT=500000 + +export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 +export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0 +export BASEDIR="/raid/datasets/wiki" + +export BENCHMARK=10 BERT_LAYERS=2 + +python3 examples/mlperf/model_train.py diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/dev_run.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/dev_run.sh new file mode 100755 index 0000000000..a6a42a6de0 --- /dev/null +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/dev_run.sh @@ -0,0 +1,20 @@ +#!/bin/bash + +export PYTHONPATH="." AMD=1 +export MODEL="bert" +export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024 + +# similar to https://github.com/mlcommons/training_results_v3.1/blob/d06288b2bd675a9d88e0e6181f5bb5626b71ec19/Quanta_Cloud_Technology/results/D54U-3U/bert/result_1.txt#L54 +export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1 +export TRAIN_STEPS=3900 + +export IGNORE_OOB=1 +export REWRITE_STACK_LIMIT=500000 + +export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 +export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0 +export BASEDIR="/raid/datasets/wiki" + +export WANDB=1 PARALLEL=0 + +RUNMLPERF=1 python3 examples/mlperf/model_train.py \ No newline at end of file diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh new file mode 100755 index 0000000000..1dbef0e48e --- /dev/null +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh @@ -0,0 +1,31 @@ +#!/bin/bash +set -e # Exit on any error +set -o pipefail # Make pipeline fail if any command fails + +export PYTHONPATH="." AMD=1 +export MODEL="bert" +export SUBMISSION_PLATFORM="tinybox_8xMI300X" +export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024 + +# similar to https://github.com/mlcommons/training_results_v3.1/blob/d06288b2bd675a9d88e0e6181f5bb5626b71ec19/Quanta_Cloud_Technology/results/D54U-3U/bert/result_1.txt#L54 +export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1 +export TRAIN_STEPS=3900 + +export IGNORE_OOB=1 +export REWRITE_STACK_LIMIT=500000 + +export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 +export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0 +export BASEDIR="/raid/datasets/wiki" + +# pip install -e ".[mlperf]" +export LOGMLPERF=1 + +export SEED=$RANDOM +DATETIME=$(date "+%m%d%H%M") +LOGFILE="bert_8xMI300x_${DATETIME}_${SEED}.log" + +BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE + +# run +PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_green/README.md b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_green/README.md new file mode 100644 index 0000000000..844b90f949 --- /dev/null +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_green/README.md @@ -0,0 +1,69 @@ +# 1. Problem + +This problem uses BERT for NLP. + +## Requirements + +Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0. +``` +git clone https://github.com/tinygrad/tinygrad.git +python3 -m pip install -e ".[mlperf]" +``` +Also install gdown (for dataset), numpy, tqdm and tensorflow. +``` +pip install gdown numpy tqdm tensorflow +``` + +### tinybox_green +Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md) +This is the default on production tinybox green. + +# 2. Directions + +## Steps to download and verify data + +### 1. Download raw data + +``` +BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py +``` + +### 2. Preprocess train and validation data + +Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended. + +#### Training: +``` +BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all +``` + +Generating a specific topic (Between 0 and 499) +``` +BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42 +``` + +#### Validation: +``` +BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval +``` +## Running + +### tinybox_green + +#### Steps to run benchmark +``` +examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh +``` + +### tinybox_red + +#### Steps to run benchmark +``` +examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh +``` +### tinybox_8xMI300X + +#### Steps to run benchmark +``` +examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh +``` \ No newline at end of file diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_green/dev_beam.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_green/dev_beam.sh new file mode 100755 index 0000000000..2865fbe06d --- /dev/null +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_green/dev_beam.sh @@ -0,0 +1,17 @@ +#!/bin/bash + +export PYTHONPATH="." NV=1 +export MODEL="bert" +export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90 + +export IGNORE_OOB=1 +export REWRITE_STACK_LIMIT=500000 + +export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 +export IGNORE_JIT_FIRST_BEAM=1 +export BEAM_LOG_SURPASS_MAX=1 +export BASEDIR="/raid/datasets/wiki" + +export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2 + +python3 examples/mlperf/model_train.py diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_green/dev_run.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_green/dev_run.sh new file mode 100755 index 0000000000..22573ae491 --- /dev/null +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_green/dev_run.sh @@ -0,0 +1,16 @@ +#!/bin/bash + +export PYTHONPATH="." NV=1 +export MODEL="bert" +export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90 + +export IGNORE_OOB=1 +export REWRITE_STACK_LIMIT=500000 + +export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 +export IGNORE_JIT_FIRST_BEAM=1 +export BASEDIR="/raid/datasets/wiki" + +export WANDB=1 PARALLEL=0 + +RUNMLPERF=1 python3 examples/mlperf/model_train.py \ No newline at end of file diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh new file mode 100755 index 0000000000..e533aea2a7 --- /dev/null +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh @@ -0,0 +1,28 @@ +#!/bin/bash +set -e # Exit on any error +set -o pipefail # Make pipeline fail if any command fails + +export PYTHONPATH="." NV=1 +export MODEL="bert" +export SUBMISSION_PLATFORM="tinybox_green" +export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90 + +export IGNORE_OOB=1 +export REWRITE_STACK_LIMIT=500000 + +export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 +export IGNORE_JIT_FIRST_BEAM=1 +export BASEDIR="/raid/datasets/wiki" + +# pip install -e ".[mlperf]" +export LOGMLPERF=1 + +export SEED=$RANDOM +DATETIME=$(date "+%m%d%H%M") +LOGFILE="bert_green_${DATETIME}_${SEED}.log" + +# init +BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE + +# run +PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_red/README.md b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_red/README.md new file mode 100644 index 0000000000..844b90f949 --- /dev/null +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_red/README.md @@ -0,0 +1,69 @@ +# 1. Problem + +This problem uses BERT for NLP. + +## Requirements + +Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0. +``` +git clone https://github.com/tinygrad/tinygrad.git +python3 -m pip install -e ".[mlperf]" +``` +Also install gdown (for dataset), numpy, tqdm and tensorflow. +``` +pip install gdown numpy tqdm tensorflow +``` + +### tinybox_green +Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md) +This is the default on production tinybox green. + +# 2. Directions + +## Steps to download and verify data + +### 1. Download raw data + +``` +BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py +``` + +### 2. Preprocess train and validation data + +Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended. + +#### Training: +``` +BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all +``` + +Generating a specific topic (Between 0 and 499) +``` +BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42 +``` + +#### Validation: +``` +BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval +``` +## Running + +### tinybox_green + +#### Steps to run benchmark +``` +examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh +``` + +### tinybox_red + +#### Steps to run benchmark +``` +examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh +``` +### tinybox_8xMI300X + +#### Steps to run benchmark +``` +examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh +``` \ No newline at end of file diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_red/dev_beam.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_red/dev_beam.sh new file mode 100755 index 0000000000..98f8d560d5 --- /dev/null +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_red/dev_beam.sh @@ -0,0 +1,18 @@ +#!/bin/bash + +export PYTHONPATH="." AMD=1 +export MODEL="bert" +export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90 + +export IGNORE_OOB=1 +export REWRITE_STACK_LIMIT=500000 + +export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 +export IGNORE_JIT_FIRST_BEAM=1 +export BEAM_LOG_SURPASS_MAX=1 +export BASEDIR="/raid/datasets/wiki" + +export RESET_STEP=1 +export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2 + +python3 examples/mlperf/model_train.py diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_red/dev_run.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_red/dev_run.sh new file mode 100755 index 0000000000..426e657ab9 --- /dev/null +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_red/dev_run.sh @@ -0,0 +1,16 @@ +#!/bin/bash + +export PYTHONPATH="." AMD=1 +export MODEL="bert" +export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90 + +export IGNORE_OOB=1 +export REWRITE_STACK_LIMIT=500000 + +export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 +export IGNORE_JIT_FIRST_BEAM=1 +export BASEDIR="/raid/datasets/wiki" + +export WANDB=1 PARALLEL=0 + +RUNMLPERF=1 python3 examples/mlperf/model_train.py \ No newline at end of file diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh new file mode 100755 index 0000000000..f54ba4b9d0 --- /dev/null +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh @@ -0,0 +1,31 @@ +#!/bin/bash +set -e # Exit on any error +set -o pipefail # Make pipeline fail if any command fails + +export PYTHONPATH="." AMD=1 +export MODEL="bert" +export SUBMISSION_PLATFORM="tinybox_red" +export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90 + +export IGNORE_OOB=1 +export REWRITE_STACK_LIMIT=500000 + +export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 +export IGNORE_JIT_FIRST_BEAM=1 +export BASEDIR="/raid/datasets/wiki" + +# pip install -e ".[mlperf]" +export LOGMLPERF=1 + +export SEED=$RANDOM +DATETIME=$(date "+%m%d%H%M") +LOGFILE="bert_red_${DATETIME}_${SEED}.log" + +export HCQDEV_WAIT_TIMEOUT_MS=100000 # prevents hang? + +# init +sleep 5 && sudo rmmod amdgpu || true +BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE + +# run +PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/resnet/implementations/tinybox_green/README.md b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/resnet/implementations/tinybox_green/README.md new file mode 100644 index 0000000000..d380cec5b5 --- /dev/null +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/resnet/implementations/tinybox_green/README.md @@ -0,0 +1,50 @@ +# 1. Problem + +This problem uses the ResNet-50 CNN to do image classification. + +## Requirements + +Install tinygrad and mlperf-logging from master. +``` +git clone https://github.com/tinygrad/tinygrad.git +python3 -m pip install -e ".[mlperf]" +``` + +### tinybox_green +Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md) +This is the default on production tinybox green. + +### tinybox_red +Disable cwsr +This is the default on production tinybox red. +``` +sudo vi /etc/modprobe.d/amdgpu.conf +cat < /etc/modprobe.d/amdgpu.conf +options amdgpu cwsr_enable=0 +EOF +sudo update-initramfs -u +sudo reboot + +# validate +sudo cat /sys/module/amdgpu/parameters/cwsr_enable #= 0 +``` + +# 2. Directions + +## Steps to download and verify data + +``` +IMGNET_TRAIN=1 python3 extra/datasets/imagenet_download.py +``` + +## Steps for one time setup + +### tinybox_red +``` +examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/setup.sh +``` + +## Steps to run benchmark +``` +examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh +``` diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/resnet/implementations/tinybox_green/dev_beam.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/resnet/implementations/tinybox_green/dev_beam.sh new file mode 100755 index 0000000000..2319da3fdc --- /dev/null +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/resnet/implementations/tinybox_green/dev_beam.sh @@ -0,0 +1,13 @@ +#!/bin/bash + +export PYTHONPATH="." NV=1 +export MODEL="resnet" +export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192 + +export RESET_STEP=0 + +export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0 + +export BENCHMARK=10 DEBUG=2 + +python3 examples/mlperf/model_train.py diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/resnet/implementations/tinybox_green/dev_run.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/resnet/implementations/tinybox_green/dev_run.sh new file mode 100755 index 0000000000..ebe927c373 --- /dev/null +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/resnet/implementations/tinybox_green/dev_run.sh @@ -0,0 +1,15 @@ +#!/bin/bash + +export PYTHONPATH="." NV=1 +export MODEL="resnet" +export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192 + +export RESET_STEP=0 + +export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0 + +export EVAL_START_EPOCH=3 EVAL_FREQ=4 + +export WANDB=1 PARALLEL=0 + +python3 examples/mlperf/model_train.py diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/resnet/implementations/tinybox_green/run_and_time.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/resnet/implementations/tinybox_green/run_and_time.sh new file mode 100755 index 0000000000..9c7193288a --- /dev/null +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/resnet/implementations/tinybox_green/run_and_time.sh @@ -0,0 +1,25 @@ +#!/bin/bash +set -e # Exit on any error +set -o pipefail # Make pipeline fail if any command fails + +export PYTHONPATH="." NV=1 +export MODEL="resnet" +export SUBMISSION_PLATFORM="tinybox_green" +export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192 + +export RESET_STEP=0 + +export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0 + +# pip install -e ".[mlperf]" +export LOGMLPERF=${LOGMLPERF:-1} + +export SEED=$RANDOM +DATETIME=$(date "+%m%d%H%M") +LOGFILE="resnet_green_${DATETIME}_${SEED}.log" + +# init +BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE + +# run +PARALLEL=0 RUNMLPERF=1 EVAL_START_EPOCH=3 EVAL_FREQ=4 python3 examples/mlperf/model_train.py | tee -a $LOGFILE diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/README.md b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/README.md new file mode 100644 index 0000000000..d380cec5b5 --- /dev/null +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/README.md @@ -0,0 +1,50 @@ +# 1. Problem + +This problem uses the ResNet-50 CNN to do image classification. + +## Requirements + +Install tinygrad and mlperf-logging from master. +``` +git clone https://github.com/tinygrad/tinygrad.git +python3 -m pip install -e ".[mlperf]" +``` + +### tinybox_green +Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md) +This is the default on production tinybox green. + +### tinybox_red +Disable cwsr +This is the default on production tinybox red. +``` +sudo vi /etc/modprobe.d/amdgpu.conf +cat < /etc/modprobe.d/amdgpu.conf +options amdgpu cwsr_enable=0 +EOF +sudo update-initramfs -u +sudo reboot + +# validate +sudo cat /sys/module/amdgpu/parameters/cwsr_enable #= 0 +``` + +# 2. Directions + +## Steps to download and verify data + +``` +IMGNET_TRAIN=1 python3 extra/datasets/imagenet_download.py +``` + +## Steps for one time setup + +### tinybox_red +``` +examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/setup.sh +``` + +## Steps to run benchmark +``` +examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh +``` diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/dev_beam.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/dev_beam.sh new file mode 100755 index 0000000000..7bcbec2f03 --- /dev/null +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/dev_beam.sh @@ -0,0 +1,13 @@ +#!/bin/bash + +export PYTHONPATH="." AMD=1 +export MODEL="resnet" +export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192 + +export RESET_STEP=0 + +export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0 + +export BENCHMARK=10 DEBUG=${DEBUG:-2} + +python3 examples/mlperf/model_train.py diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/dev_run.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/dev_run.sh new file mode 100755 index 0000000000..aad23e43df --- /dev/null +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/dev_run.sh @@ -0,0 +1,15 @@ +#!/bin/bash + +export PYTHONPATH="." AMD=1 +export MODEL="resnet" +export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192 + +export RESET_STEP=0 + +export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0 + +export EVAL_START_EPOCH=3 EVAL_FREQ=4 + +export WANDB=1 PARALLEL=0 + +python3 examples/mlperf/model_train.py diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh new file mode 100755 index 0000000000..7a93d435a5 --- /dev/null +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh @@ -0,0 +1,26 @@ +#!/bin/bash +set -e # Exit on any error +set -o pipefail # Make pipeline fail if any command fails + +export PYTHONPATH="." AMD=1 +export MODEL="resnet" +export SUBMISSION_PLATFORM="tinybox_red" +export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192 + +export RESET_STEP=0 + +export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0 + +# pip install -e ".[mlperf]" +export LOGMLPERF=${LOGMLPERF:-1} + +export SEED=$RANDOM +DATETIME=$(date "+%m%d%H%M") +LOGFILE="resnet_red_${DATETIME}_${SEED}.log" + +# init +sleep 5 && sudo rmmod amdgpu || true +BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE + +# run +PARALLEL=0 RUNMLPERF=1 EVAL_START_EPOCH=3 EVAL_FREQ=4 python3 examples/mlperf/model_train.py | tee -a $LOGFILE diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/setup.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/setup.sh new file mode 100755 index 0000000000..a9806164f4 --- /dev/null +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/setup.sh @@ -0,0 +1,8 @@ +#!/bin/bash + +rocm-smi --setprofile compute +rocm-smi --setmclk 3 +rocm-smi --setperflevel high + +# power cap to 350W +echo "350000000" | sudo tee /sys/class/drm/card{1..6}/device/hwmon/hwmon*/power1_cap diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/retinanet/implementations/tinybox_green/README.md b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/retinanet/implementations/tinybox_green/README.md new file mode 100644 index 0000000000..ce1ac9b9a3 --- /dev/null +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/retinanet/implementations/tinybox_green/README.md @@ -0,0 +1,38 @@ +# 1. Problem + +This problem uses RetinaNet for SSD. + +## Requirements + +Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0. +``` +git clone https://github.com/tinygrad/tinygrad.git +python3 -m pip install -e ".[mlperf]" +``` + +Also install the following dependencies: +``` +pip install tqdm numpy pycocotools boto3 pandas torch torchvision +``` + +### tinybox_green +Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md) +This is the default on production tinybox green. + +# 2. Directions + +## Steps to download data + +Run the following: +``` +BASEDIR=/raid/datasets/openimages python3 extra/datasets/openimages.py +``` + +## Running + +### tinybox_green + +#### Steps to run benchmark +``` +examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/retinanet/implementations/tinybox_green/run_and_time.sh +``` diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/retinanet/implementations/tinybox_green/dev_beam.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/retinanet/implementations/tinybox_green/dev_beam.sh new file mode 100755 index 0000000000..6e25bb9671 --- /dev/null +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/retinanet/implementations/tinybox_green/dev_beam.sh @@ -0,0 +1,14 @@ +#!/bin/bash + +export PYTHONPATH="." NV=1 +export MODEL="retinanet" +export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96 +export BASEDIR="/raid/datasets/openimages" + +# export RESET_STEP=0 + +export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0 + +export BENCHMARK=5 DEBUG=2 + +python examples/mlperf/model_train.py diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/retinanet/implementations/tinybox_green/dev_run.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/retinanet/implementations/tinybox_green/dev_run.sh new file mode 100755 index 0000000000..7a3ee0dfa2 --- /dev/null +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/retinanet/implementations/tinybox_green/dev_run.sh @@ -0,0 +1,15 @@ +#!/bin/bash + +export PYTHONPATH="." NV=1 +export MODEL="retinanet" +export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96 +export BASEDIR="/raid/datasets/openimages" + +# export RESET_STEP=0 + +export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0 + +export WANDB=1 PARALLEL=0 +export RUNMLPERF=1 + +python examples/mlperf/model_train.py diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/retinanet/implementations/tinybox_green/run_and_time.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/retinanet/implementations/tinybox_green/run_and_time.sh new file mode 100755 index 0000000000..74cdc87a1b --- /dev/null +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/retinanet/implementations/tinybox_green/run_and_time.sh @@ -0,0 +1,25 @@ +#!/bin/bash +set -e # Exit on any error +set -o pipefail # Make pipeline fail if any command fails + +export PYTHONPATH="." NV=1 +export MODEL="retinanet" +export SUBMISSION_PLATFORM="tinybox_green" +export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96 + +export TRAIN_BEAM=2 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0 +export IGNORE_JIT_FIRST_BEAM=1 +export BASEDIR="/raid/datasets/openimages" + +# pip install -e ".[mlperf]" +export LOGMLPERF=1 + +export SEED=$RANDOM +DATETIME=$(date "+%m%d%H%M") +LOGFILE="retinanet_green_${DATETIME}_${SEED}.log" + +# init +BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE + +# run +PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/retinanet/implementations/tinybox_red/dev_beam.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/retinanet/implementations/tinybox_red/dev_beam.sh new file mode 100755 index 0000000000..97aa5155eb --- /dev/null +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/retinanet/implementations/tinybox_red/dev_beam.sh @@ -0,0 +1,14 @@ +#!/bin/bash + +export PYTHONPATH="." AMD=1 +export MODEL="retinanet" +export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96 +export BASEDIR="/raid/datasets/openimages" + +# export RESET_STEP=0 + +export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0 + +export BENCHMARK=5 DEBUG=2 + +python examples/mlperf/model_train.py diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/retinanet/implementations/tinybox_red/dev_run.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/retinanet/implementations/tinybox_red/dev_run.sh new file mode 100755 index 0000000000..5fb4d109fd --- /dev/null +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/retinanet/implementations/tinybox_red/dev_run.sh @@ -0,0 +1,15 @@ +#!/bin/bash + +export PYTHONPATH="." AMD=1 +export MODEL="retinanet" +export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96 +export BASEDIR="/raid/datasets/openimages" + +# export RESET_STEP=0 + +export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0 + +export WANDB=1 PARALLEL=0 +export RUNMLPERF=1 + +python examples/mlperf/model_train.py diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/systems/tinybox_8xMI300X.json b/examples/mlperf/training_submission_v6.0/tinycorp/systems/tinybox_8xMI300X.json new file mode 100644 index 0000000000..1e0f789430 --- /dev/null +++ b/examples/mlperf/training_submission_v6.0/tinycorp/systems/tinybox_8xMI300X.json @@ -0,0 +1,38 @@ +{ + "submitter": "tinycorp", + "division": "closed", + "status": "Available on-premise", + "system_name": "tinybox 8xMI300X", + "number_of_nodes": "1", + "host_processors_per_node": "2", + "host_processor_model_name": "AMD EPYC 9354", + "host_processor_core_count": "32", + "host_processor_vcpu_count": "64", + "host_processor_frequency": "", + "host_processor_caches": "", + "host_processor_interconnect": "", + "host_memory_capacity": "2304GB", + "host_storage_type": "NVMe SSD", + "host_storage_capacity": "3x 4TB raid array", + "host_networking": "", + "host_networking_topology": "", + "host_memory_configuration": "24x 96GB DDR5", + "accelerators_per_node": "8", + "accelerator_model_name": "AMD Instinct MI300X 192GB HBM3", + "accelerator_host_interconnect": "PCIe 5.0 x16", + "accelerator_frequency": "", + "accelerator_on-chip_memories": "", + "accelerator_memory_configuration": "HBM3", + "accelerator_memory_capacity": "192GB", + "accelerator_interconnect": "", + "accelerator_interconnect_topology": "", + "cooling": "air", + "hw_notes": "", + "framework": "tinygrad, branch mlperf_training_v5.0", + "other_software_stack": { + "python": "3.10.16", + "ROCm": "3.0.0+94441cb" + }, + "operating_system": "Ubuntu 24.04.1 LTS", + "sw_notes": "" + } \ No newline at end of file diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/systems/tinybox_green.json b/examples/mlperf/training_submission_v6.0/tinycorp/systems/tinybox_green.json new file mode 100644 index 0000000000..24cbce1f1c --- /dev/null +++ b/examples/mlperf/training_submission_v6.0/tinycorp/systems/tinybox_green.json @@ -0,0 +1,38 @@ +{ + "submitter": "tinycorp", + "division": "closed", + "status": "Available on-premise", + "system_name": "tinybox green", + "number_of_nodes": "1", + "host_processors_per_node": "1", + "host_processor_model_name": "AMD EPYC 7532", + "host_processor_core_count": "32", + "host_processor_vcpu_count": "64", + "host_processor_frequency": "", + "host_processor_caches": "", + "host_processor_interconnect": "", + "host_memory_capacity": "128GB", + "host_storage_type": "NVMe SSD", + "host_storage_capacity": "4 TB raid array + 1 TB boot", + "host_networking": "", + "host_networking_topology": "", + "host_memory_configuration": "8x 16GB DDR4", + "accelerators_per_node": "6", + "accelerator_model_name": "NVIDIA GeForce RTX 4090", + "accelerator_host_interconnect": "PCIe 4.0 x16", + "accelerator_frequency": "", + "accelerator_on-chip_memories": "", + "accelerator_memory_configuration": "GDDR6X", + "accelerator_memory_capacity": "24GB", + "accelerator_interconnect": "", + "accelerator_interconnect_topology": "", + "cooling": "air", + "hw_notes": "", + "framework": "tinygrad, branch mlperf_training_v5.0", + "other_software_stack": { + "python": "3.10.12", + "CUDA": "12.4" + }, + "operating_system": "Ubuntu 22.04.4", + "sw_notes": "" +} \ No newline at end of file diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/systems/tinybox_red.json b/examples/mlperf/training_submission_v6.0/tinycorp/systems/tinybox_red.json new file mode 100644 index 0000000000..58b6efe77c --- /dev/null +++ b/examples/mlperf/training_submission_v6.0/tinycorp/systems/tinybox_red.json @@ -0,0 +1,37 @@ +{ + "submitter": "tinycorp", + "division": "closed", + "status": "Available on-premise", + "system_name": "tinybox red", + "number_of_nodes": "1", + "host_processors_per_node": "1", + "host_processor_model_name": "AMD EPYC 7532", + "host_processor_core_count": "32", + "host_processor_vcpu_count": "64", + "host_processor_frequency": "", + "host_processor_caches": "", + "host_processor_interconnect": "", + "host_memory_capacity": "128GB", + "host_storage_type": "NVMe SSD", + "host_storage_capacity": "4 TB raid array + 1 TB boot", + "host_networking": "", + "host_networking_topology": "", + "host_memory_configuration": "8x 16GB DDR4", + "accelerators_per_node": "6", + "accelerator_model_name": "AMD Radeon RX 7900 XTX", + "accelerator_host_interconnect": "PCIe 4.0 x16", + "accelerator_frequency": "", + "accelerator_on-chip_memories": "", + "accelerator_memory_configuration": "GDDR6", + "accelerator_memory_capacity": "24GB", + "accelerator_interconnect": "", + "accelerator_interconnect_topology": "", + "cooling": "air", + "hw_notes": "", + "framework": "tinygrad, branch mlperf_training_v5.0", + "other_software_stack": { + "python": "3.10.12" + }, + "operating_system": "Ubuntu 22.04.4", + "sw_notes": "" +} \ No newline at end of file diff --git a/extra/sqtt/rocprof/install.py b/extra/sqtt/rocprof/install.py new file mode 100755 index 0000000000..5243180602 --- /dev/null +++ b/extra/sqtt/rocprof/install.py @@ -0,0 +1,18 @@ +#!/usr/bin/env python3 +import os, shutil +from pathlib import Path +from tinygrad.helpers import fetch, OSX + +DEST = Path("/usr/local/lib") +DEST.mkdir(exist_ok=True) + +if __name__ == "__main__": + if OSX: + fp = fetch("https://github.com/ROCm/rocprof-trace-decoder/releases/download/0.1.4/rocprof-trace-decoder-macos-arm64-0.1.4-Darwin.sh") + lib = fp.parent/"rocprof-trace-decoder-macos-arm64-0.1.4-Darwin"/"lib"/"librocprof-trace-decoder.dylib" + os.chmod(fp, 0o755) + os.system(f"sudo {fp} --prefix={fp.parent} --include-subdir") + else: + lib = fetch("https://github.com/ROCm/rocprof-trace-decoder/raw/5420409ad0963b2d76450add067b9058493ccbd0/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so", name="librocprof-trace-decoder.so") + shutil.copy2(lib, DEST) + print(f"Installed {lib.name} to", DEST) diff --git a/extra/sqtt/rocprof/rocprof.py b/extra/sqtt/rocprof/rocprof.py index 1d0b151bb1..bded16acc2 100644 --- a/extra/sqtt/rocprof/rocprof.py +++ b/extra/sqtt/rocprof/rocprof.py @@ -7,8 +7,7 @@ # POINTER_SIZE is: 8 # LONGDOUBLE_SIZE is: 16 # -import ctypes -from tinygrad.helpers import fetch +import ctypes, ctypes.util class AsDictMixin: @@ -156,7 +155,7 @@ class FunctionFactoryStub: # You can either re-run clan2py with -l /path/to/library.so # Or manually fix this by comment the ctypes.CDLL loading _libraries = {} -_libraries['FIXME_STUB'] = ctypes.CDLL(str(fetch('https://github.com/ROCm/rocprof-trace-decoder/raw/5420409ad0963b2d76450add067b9058493ccbd0/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so'))) # ctypes.CDLL('FIXME_STUB') +_libraries['FIXME_STUB'] = ctypes.CDLL(ctypes.util.find_library('rocprof-trace-decoder')) # ctypes.CDLL('FIXME_STUB') diff --git a/ruff.toml b/ruff.toml index 22f9bf566b..0d5b7cb8f0 100644 --- a/ruff.toml +++ b/ruff.toml @@ -50,4 +50,7 @@ exclude = [ "E303", "E304", "E501", "E702", "E703", "E731", "W191", "W291", "W293", "UP039", "C416", "RET506", "RET507", "A", "FURB110", "RUF018", "F541", "F841" -] \ No newline at end of file +] + +[format] +exclude = ["*"] diff --git a/test/external/external_test_simple_tokenizer.py b/test/external/external_test_simple_tokenizer.py index 9c3ca8f420..8fc3299ee1 100644 --- a/test/external/external_test_simple_tokenizer.py +++ b/test/external/external_test_simple_tokenizer.py @@ -1,41 +1,50 @@ +import functools, multiprocessing from transformers import AutoTokenizer from datasets import load_dataset -from tinygrad.apps.llm import SimpleTokenizer, gpt2_decode_vocab, get_llama_re +from tinygrad.apps.llm import SimpleTokenizer from tinygrad.helpers import tqdm, getenv, partition +@functools.cache +def get_tokenizers(): + print("getting tokenizers") + base_tokenizer = AutoTokenizer.from_pretrained("NousResearch/Meta-Llama-3-8B-Instruct") + special_tokens, normal_tokens = partition(((t, tid) for t, tid in base_tokenizer.vocab.items()), lambda e: e[1] in base_tokenizer.all_special_ids) + simple_tokenizer = SimpleTokenizer(dict(normal_tokens), dict(special_tokens)) + return base_tokenizer, simple_tokenizer + +def test_tokenize(samp) -> bool: + base_tokenizer, simple_tokenizer = get_tokenizers() + idx, txt = samp + try: simple_tokens = tuple(simple_tokenizer.encode(txt)) + except RuntimeError: simple_tokens = () + base_tokens = tuple(base_tokenizer.encode(txt, add_special_tokens=False)) + if simple_tokens != base_tokens: + print(f"tokens mismatch at index: {idx}.\n") + color_codes = [91, 92, 94, 93, 95] + def color_tokens(tids): + return "".join(f"\033[{color_codes[i%len(color_codes)]}m{base_tokenizer.decode([t])}" for i, t in enumerate(tids)) + "\033[0m" + print("simple: ", color_tokens(simple_tokens)) + print("official:", color_tokens(base_tokens) + "\n") + return False + if simple_tokenizer.decode(simple_tokens) != txt: + print(f"decode mismatch at {idx}") + return False + return True + # use ALLOW_FAILED=-1 to go over the entire dataset without printing. if __name__ == "__main__": - base_tokenizer = AutoTokenizer.from_pretrained("NousResearch/Meta-Llama-3-8B-Instruct") - special_tokens, normal_tokens = partition(((t, tid) for t, tid in base_tokenizer.vocab.items()), - lambda e: e[1] in base_tokenizer.all_special_ids) - inv_vocab = { tid: word for word, tid in base_tokenizer.get_vocab().items() } - simple_tokenizer = SimpleTokenizer(get_llama_re(), gpt2_decode_vocab(dict(normal_tokens)), dict(special_tokens)) - - color_codes = [ 91, 92, 94, 93, 95 ] - def color_tokens(tids): - return "".join(f"\033[{color_codes[i%len(color_codes)]}m{base_tokenizer.decode([t])}" for i, t in enumerate(tids)) + "\033[0m" - + print("loading datasets") ds = load_dataset("OpenAssistant/oasst1") + loaded_ds = [(idx, el["text"]) for idx, el in enumerate(ds["train"])] + print(f"loaded {len(loaded_ds)}") + allow_failed = getenv("ALLOW_FAILED", 10) - fail_count, total = 0, 0 - - for idx, el in enumerate(tqdm(ds["train"])): - total += 1 - - try: simple_tokens = tuple(simple_tokenizer.encode(el["text"])) - except RuntimeError: simple_tokens = () - base_tokens = tuple(base_tokenizer.encode(el["text"], add_special_tokens=False)) - - if simple_tokens != base_tokens: - fail_count += 1 - allow_failed -= 1 - - if allow_failed >= 0: - print(f"tokens mismatch at index: {idx}.\n") - - print("simple: ", color_tokens(simple_tokens)) - print("official:", color_tokens(base_tokens) + "\n") - - if allow_failed == 0: break - print(f"{fail_count}/{total} samples are inconsistent with the official tokenizer.") + with multiprocessing.Pool(16) as pool: + for good in tqdm(pool.imap_unordered(test_tokenize, loaded_ds), total=len(loaded_ds)): + total += 1 + if not good: + fail_count += 1 + allow_failed -= 1 + if allow_failed == 0: break + print(f"{fail_count}/{total} samples are inconsistent with the official tokenizer.") diff --git a/test/external/external_uop_gc.py b/test/external/external_uop_gc.py index 3a39200929..1155c068bf 100644 --- a/test/external/external_uop_gc.py +++ b/test/external/external_uop_gc.py @@ -2,6 +2,7 @@ import gc from tinygrad import Tensor, UOp, Device, nn from tinygrad.shape.shapetracker import views_to_valid_uop from tinygrad.engine.realize import method_cache, get_program +from tinygrad.schedule.indexing import apply_movement_op from test.test_tiny import TestTiny def uops_allocated(): return sum([isinstance(x, UOp) for x in gc.get_objects()]) @@ -69,6 +70,7 @@ if __name__ == "__main__": # these caches will keep uops alive method_cache.clear() views_to_valid_uop.cache_clear() + apply_movement_op.cache_clear() Tensor._device_seeds.clear() Tensor._device_rng_counters.clear() diff --git a/test/external/process_replay/process_replay.py b/test/external/process_replay/process_replay.py index 3ffe5f70b7..dee3199881 100755 --- a/test/external/process_replay/process_replay.py +++ b/test/external/process_replay/process_replay.py @@ -42,13 +42,13 @@ class ProcessReplayWarning(Warning): pass # *** replay the function and convert return values to string -def replay_kernelize(ret:dict[UOp, UOp], big_sink:UOp) -> tuple[str, str, tuple[Any, ...]]: +def replay_get_rangeify_map(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)) 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]) - return to_str(new_sink), to_str(ret[big_sink]), (big_sink,) + return to_str(new_sink), to_str(big_sink.substitute(ret)), (big_sink,) def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer|None=None, opts:list[Opt]|None=None) -> tuple[str, str, tuple[Any, ...]]: # NOTE: this always uses the opts_to_apply path @@ -65,7 +65,7 @@ def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer|None=None, opts ast_repr = codecs.decode(str(input_ast), "unicode_escape") return to_str(p2), to_str(p), (ast_repr, renderer) -replayers: dict[str, Callable[..., tuple[str, str, tuple[Any, ...]]]] = {"get_kernelize_map":replay_kernelize, "get_program":replay_get_program} +replayers: dict[str, Callable[..., tuple[str, str, tuple[Any, ...]]]] = {"get_rangeify_map":replay_get_rangeify_map, "get_program":replay_get_program} # *** run replayers on captured rows and print diffs diff --git a/test/test_linearizer.py b/test/test_linearizer.py index a0d6d67f67..7af6294c83 100644 --- a/test/test_linearizer.py +++ b/test/test_linearizer.py @@ -8,9 +8,11 @@ from tinygrad.uop.ops import UOp, Ops, GroupOp from tinygrad.device import Device, Buffer, is_dtype_supported 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, getenv from tinygrad.dtype import DType, dtypes, PtrDType, AddrSpace from tinygrad.renderer.ptx import PTXRenderer +from tinygrad.renderer.cstyle import CUDARenderer +MOCKGPU = getenv("MOCKGPU") class TestLinearizer(unittest.TestCase): def test_arg_dedup(self): @@ -68,7 +70,8 @@ class TestLinearizer(unittest.TestCase): ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE] # RANGE -> ALU -> RANGE -> ALU + LOAD -> STORE assert any(x.op in GroupOp.ALU for x in uops[ranges[0]:ranges[1]]) - assert not any(x.op is Ops.LOAD for x in uops[ranges[0]:ranges[1]]) + # the index of the load doesnt depend on the second range + assert any(x.op is Ops.LOAD for x in uops[ranges[0]:ranges[1]]) assert any(x.op in {*GroupOp.ALU, Ops.LOAD} for x in uops[ranges[1]:]) def test_range_outer_op_before_phi(self): @@ -314,7 +317,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(MOCKGPU and isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, CUDARenderer)), "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)) diff --git a/test/test_multitensor.py b/test/test_multitensor.py index fdf07f3ca7..2fa3a614b8 100644 --- a/test/test_multitensor.py +++ b/test/test_multitensor.py @@ -390,7 +390,6 @@ 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") def test_data_parallel_resnet(self): from extra.models.resnet import ResNet18 diff --git a/test/test_schedule.py b/test/test_schedule.py index 4817272317..7761a0ca95 100644 --- a/test/test_schedule.py +++ b/test/test_schedule.py @@ -1526,7 +1526,7 @@ class TestSchedule(unittest.TestCase): # run_schedule(check_schedule(out, 1)) run_schedule(check_schedule(out, 4)) np.testing.assert_allclose(out.numpy(), np.pad(np.log2(np.abs(np.pad(np.log2(a.numpy()), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum() + \ - b.numpy())), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum(), atol=3e-4, rtol=1e-6) + b.numpy())), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum(), atol=3e-4, rtol=1e-5) def test_shrink_pad_safe(self): a = Tensor.ones((3, )).contiguous().realize() diff --git a/test/unit/test_llm_tokenizer.py b/test/unit/test_llm_tokenizer.py index 7b65818a6f..1e7f6cb48a 100644 --- a/test/unit/test_llm_tokenizer.py +++ b/test/unit/test_llm_tokenizer.py @@ -1,19 +1,21 @@ import unittest, base64, functools, sys -from tinygrad.apps.llm import SimpleTokenizer, get_llama_re +from tinygrad.apps.llm import SimpleTokenizer from tinygrad.helpers import fetch @unittest.skipIf(sys.platform == 'win32', "fetch race condition on Windows") class TestLLMTokenizer(unittest.TestCase): - @functools.cached_property - def basic_tok(self): return SimpleTokenizer(".*", { b"a": 0, b"b": 1, b"c": 2, b"ab": 3, b"bc": 4 }, { "": 5, "": 6, "": 7 }) - @functools.cached_property def llama_tok(self): # from https://github.com/tinygrad/tinygrad/blob/e0106b6b257ebc003eb3694144e3e198f7d8cc37/examples/llama3.py#L14 model_file = fetch("https://huggingface.co/bofenghuang/Meta-Llama-3-8B/resolve/main/original/tokenizer.model") with open(model_file, "rt") as fd: - str_vocab = [ line.split(maxsplit=1) for line in fd.read().splitlines() if line ] - normal_tokens = { base64.b64decode(stok): int(srank) for stok, srank in str_vocab } + str_vocab = [line.split(maxsplit=1) for line in fd.read().splitlines() if line] + + # https://github.com/openai/gpt-2/blob/9b63575ef42771a015060c964af2c3da4cf7c8ab/src/encoder.py#L9 + bs = [*range(33, 127), *range(161, 173), *range(174, 256)] # bytes that map to themselves + _byte_decoder = {chr(b): b for b in bs} | {chr(256+i): b for i,b in enumerate(b for b in range(256) if b not in bs)} + _byte_encoder = {v:k for k,v in _byte_decoder.items()} + normal_tokens = {''.join([_byte_encoder[x] for x in base64.b64decode(stok)]): int(srank) for stok, srank in str_vocab} special_tokens = [ "<|begin_of_text|>", @@ -27,22 +29,12 @@ class TestLLMTokenizer(unittest.TestCase): "<|reserved_special_token_4|>", "<|eot_id|>", ] + [ f"<|reserved_special_token_{i}|>" for i in range(5, 256 - 5) ] - return SimpleTokenizer(get_llama_re(), normal_tokens, { token: len(normal_tokens) + i for i, token in enumerate(special_tokens) }) + return SimpleTokenizer(normal_tokens, {token: len(normal_tokens) + i for i, token in enumerate(special_tokens)}) def _test_coding(self, tok: SimpleTokenizer, text: str, expected_tokens: list[int]): self.assertEqual(tok.encode(text), expected_tokens) self.assertEqual(tok.decode(expected_tokens), text) - def test_abc(self): self._test_coding(self.basic_tok, "abc", [ 3, 2 ]) - def test_abbc(self): self._test_coding(self.basic_tok, "abbc", [ 3, 4 ]) - def test_aabbbcc(self): self._test_coding(self.basic_tok, "aabbbcc", [ 0, 3, 1, 4, 2 ]) - def test_specials1(self): self._test_coding(self.basic_tok, "aaaa", [ 0, 5, 0, 6, 0, 7, 0 ]) - def test_specials2(self): self._test_coding(self.basic_tok, "aa", [ 5, 0, 6, 0, 7 ]) - def test_invalid_token(self): - with self.assertRaises(RuntimeError): self._test_coding(self.basic_tok, "L", []) - - def test_no_specials(self): self._test_coding(SimpleTokenizer(".*", { bytes([i]): i for i in range(256) }, {}), "abc", [97, 98, 99]) - # NOTE: the correct tokenization for this can only be found by looking up the text chunk in the vocab, not by applying merges def test_llama_early_tokenize(self): self._test_coding(self.llama_tok, " например", [ 111797 ]) diff --git a/test/unit/test_simplify_valid_idx.py b/test/unit/test_simplify_valid_idx.py index 534fd9697a..7f3790c217 100644 --- a/test/unit/test_simplify_valid_idx.py +++ b/test/unit/test_simplify_valid_idx.py @@ -5,6 +5,7 @@ from tinygrad.dtype import dtypes from tinygrad.uop.ops import UOp, Ops from tinygrad.uop.symbolic import simplify_valid from tinygrad.helpers import Context +from .test_uop_symbolic import check_uop_against_string def get_gated_load_uop(valid:UOp, idx:UOp): return UOp(Ops.LOAD, dtypes.float, ( @@ -49,8 +50,8 @@ class TestValidIdxSimplification(unittest.TestCase): with Context(NOOPT=1): load = full_rewrite_to_sink(load.sink()).src[0] idx, valid = load.src[0].src[1], load.src[0].src[2] - self.assertEqual(idx.render(simplify=False), sidx) - self.assertEqual(valid.render(simplify=False), svalid) + check_uop_against_string(self, idx, sidx) + check_uop_against_string(self, valid, svalid) def test_cumsum(self): gidx0 = Special("gidx0", 5) @@ -186,13 +187,13 @@ 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) ridx2 = Range(2, 4) ridx3 = Range(3, 4) - idx= ((ridx0+ridx1+ridx2+ridx3+28)//30) + # TODO: this should also work without the extra nesting + idx = (((ridx0+ridx1)+(ridx2+ridx3)+28)//30) valid = ((ridx0+ridx1)<1).ne(True) & ((ridx2+ridx3)<1).ne(True) load = get_gated_load_uop(valid, idx) self.check(load, @@ -218,10 +219,10 @@ class TestImageSimplification(unittest.TestCase): self.assertEqual(idx.op, Ops.VECTORIZE) self.assertEqual(len(idx.src), 2) idx0, idx1 = idx.src[0], idx.src[1] - self.assertEqual(idx0.render(simplify=False), sidx0) - self.assertEqual(idx1.render(simplify=False), sidx1) + check_uop_against_string(self, idx0, sidx0) + check_uop_against_string(self, idx1, sidx1) if svalid is not None: - self.assertEqual(load.src[0].src[2].render(simplify=False), svalid) + check_uop_against_string(self, load.src[0].src[2], svalid) else: self.assertEqual(len(load.src[0].src), 2, "svalid is None but load still has a valid") diff --git a/test/unit/test_tinyfs.py b/test/unit/test_tinyfs.py new file mode 100644 index 0000000000..9fe4fed13f --- /dev/null +++ b/test/unit/test_tinyfs.py @@ -0,0 +1,22 @@ +import unittest +from tinygrad import Tensor + +class TestLoadStore(unittest.TestCase): + def test_load_shape(self): + t = Tensor(bytes(16)).load(1024).kernelize() + assert t.shape == (1024,), t.shape + + def test_store_shape(self): + t = Tensor.zeros(1024).store().kernelize() + assert t.shape == (16,), t.shape + + def test_load_large_shape(self): + t = Tensor(bytes(16)).load(10_000_000).kernelize() + assert t.shape == (10_000_000,), t.shape + + def test_store_large_shape(self): + t = Tensor.zeros(10_000_000).store().kernelize() + assert t.shape == (16,), t.shape + +if __name__ == "__main__": + unittest.main() diff --git a/test/unit/test_uop_symbolic.py b/test/unit/test_uop_symbolic.py index 8c0bd638e5..3c1805ff3b 100644 --- a/test/unit/test_uop_symbolic.py +++ b/test/unit/test_uop_symbolic.py @@ -5,14 +5,17 @@ import z3 from tinygrad.dtype import dtypes, ConstType, DType, Invalid from tinygrad.codegen import full_rewrite from tinygrad.helpers import Context -from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer, track_rewrites -from tinygrad.uop.symbolic import sym +from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer +from tinygrad.uop.symbolic import sym, commutative from tinygrad.uop.spec import uops_to_z3 -@track_rewrites(name="simplify symbolic uop") -def render(v) -> UOp: - v_simplified = graph_rewrite(v, sym) - return v_simplified +def check_uop_against_string(self, v:UOp, s:str): + sym_vars = {v.render():v for v in v.toposort() if v.op in (Ops.DEFINE_VAR, Ops.RANGE, Ops.SPECIAL)} + s_eval = eval(s, sym_vars) + if isinstance(s_eval, int) and v.dtype==dtypes.index: s_eval = UOp.const(dtypes.index, s_eval) + elif isinstance(s_eval, (bool, int, float)): s_eval = UOp.const(dtypes.from_py(s_eval), s_eval) + s_eval = graph_rewrite(s_eval, commutative, name="cannonicalize eval") + self.assertIs(s_eval, v, f"eval did not match simplified: {s_eval} != {v} for {s}") def Variable(name: str, min_val: ConstType, max_val: ConstType, dtype: DType=dtypes.index): return UOp.variable(name,min_val,max_val,dtype) def uconst(val): return UOp.const(dtypes.index, val) @@ -33,11 +36,11 @@ class TestSymbolic(unittest.TestCase): self.assertEqual(solver.check(expr1 != expr2), z3.unsat, "simplified expression not equal to original") def helper_test_variable(self, v, n, m, s, test_z3:bool=True): - v_simplified = render(v) + v_simplified = graph_rewrite(v, sym, name="simplify symbolic uop") if test_z3: self.check_equal_z3(v, v_simplified) - rendered, nmin, nmax = v_simplified.render(simplify=False), v_simplified.vmin, v_simplified.vmax - if isinstance(s, tuple): self.assertIn(rendered, s) - else: self.assertEqual(rendered, s) + nmin, nmax = v_simplified.vmin, v_simplified.vmax + check_uop_against_string(self, v_simplified, s) + # eval the test string and see if we get the same uop self.assertEqual(nmin, n) self.assertEqual(nmax, m) @@ -76,7 +79,7 @@ class TestSymbolic(unittest.TestCase): def test_lt_factors(self): expr = (Variable("idx1", 0, 511)*4 + Variable("FLOAT4_INDEX", 0, 256)) < 512 - self.helper_test_variable(expr, 0, 1, ("(((idx1*4)+FLOAT4_INDEX)<512)", "((FLOAT4_INDEX+(idx1*4))<512)")) + self.helper_test_variable(expr, 0, 1, "(((idx1*4)+FLOAT4_INDEX)<512)") def test_div_reduction(self): self.helper_test_variable(Variable("a", 2, 3)//2, 1, 1, "1") @@ -187,7 +190,7 @@ class TestSymbolic(unittest.TestCase): self.helper_test_variable(Variable("a", 0, 8)%1, 0, 0, "0") def test_max_folds(self): - self.helper_test_variable(Variable("a", 0, 20).maximum(10).maximum(11), 11, 20, "max(a, 11)") + self.helper_test_variable(Variable("a", 0, 20).maximum(10).maximum(11), 11, 20, "a.maximum(11)") def test_add_min_max(self): self.helper_test_variable(Variable("a", 0, 8) * 2 + 12, 12, 16+12, "((a*2)+12)") @@ -216,7 +219,7 @@ class TestSymbolic(unittest.TestCase): self.helper_test_variable(usum([Variable("a", 0, 7)*4, Variable("b", 0, 3)*4]) % 2, 0, 0, "0") def test_sum_div_some_factor(self): - self.helper_test_variable(usum([Variable("a", 0, 7)*5, Variable("b", 0, 3)*4]) // 2, 0, 23, ("(((a*5)//2)+(b*2))", "((b*2)+((a*5)//2))")) + self.helper_test_variable(usum([Variable("a", 0, 7)*5, Variable("b", 0, 3)*4]) // 2, 0, 23, "(((a*5)//2)+(b*2))") def test_sum_div_trim_const(self): self.helper_test_variable((Variable("a", 0, 7)*4 + Variable("b", 0, 3)*4 + 7) // 16, 0, 2, "(((a+b)+1)//4)") @@ -279,7 +282,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)") 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): @@ -301,8 +304,7 @@ class TestSymbolic(unittest.TestCase): def test_sum_lt_fold(self): self.helper_test_variable(usum([Variable("a", 0, 7) * 4, Variable("b", 0, 3)]) < 16, 0, 1, "(a<4)") - self.helper_test_variable(usum([Variable("a", 0, 7) * 4, Variable("b", 0, 4)]) < 16, 0, 1, - ("(((a*4)+b)<16)", "((b+(a*4))<16)")) + self.helper_test_variable(usum([Variable("a", 0, 7) * 4, Variable("b", 0, 4)]) < 16, 0, 1, "(((a*4)+b)<16)") self.helper_test_variable(usum([Variable("uidx", 0, 3), Variable("a", 0, 1529) * 12]) < (4 * 67), 0, 1, "(a<23)") def test_mul_mod_large(self): @@ -364,7 +366,7 @@ class TestSymbolic(unittest.TestCase): self.helper_test_variable((1+Variable("a", 0, 3))*(-2)+12, 4, 10, "((a*-2)+10)") def test_mod_mul_sum(self): - self.helper_test_variable(usum([Variable("b", 0, 2), Variable("a", 0, 5)*10])%9, 0, 7, ("(b+a)", "(a+b)")) + self.helper_test_variable(usum([Variable("b", 0, 2), Variable("a", 0, 5)*10])%9, 0, 7, "(b+a)") def test_sum_0(self): self.helper_test_variable(usum([Variable("a", 0, 7)]), 0, 7, "a") @@ -395,11 +397,11 @@ class TestSymbolic(unittest.TestCase): def test_lt_sum_factor_rhs_partial(self): self.helper_test_variable((Variable("a", 0, 6)*6 + Variable("b", 0, 6)*4 + Variable("c", 0, 6)*8) < 4, 0, 1, - ("((((a*3)+(b*2))+(c*4))<2)", "(((b*2)+((a*3)+(c*4)))<2)")) + "((((a*3)+(b*2))+(c*4))<2)") def test_lt_sum_factor_rhs_all(self): self.helper_test_variable((Variable("a", 0, 6)*6 + Variable("b", 0, 6)*4 + Variable("c", 0, 6)*8) < 2, 0, 1, - ("((((a*3)+(b*2))+(c*4))<1)", "(((b*2)+((a*3)+(c*4)))<1)")) + "((((a*3)+(b*2))+(c*4))<1)") def test_and_fold(self): self.helper_test_variable(uand([uconst(0), Variable("a", 0, 1)]), 0, 0, "0") @@ -561,38 +563,35 @@ 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))")) + "((((((gidx0//5)+lidx2)//5)+lidx1)//2)+(((gidx2*32)+(gidx1*8))+(lidx0*16)))") def test_sum_div_complex2(self): gidx0 = Variable("gidx0", 0, 7) lidx2 = Variable("lidx2", 0, 1) lidx3 = Variable("lidx3", 0, 1) - self.helper_test_variable((gidx0*4+lidx2*2+1)//10, 0, 3, ("(((gidx0*2)+lidx2)//5)", "((lidx2+(gidx0*2))//5)")) - self.helper_test_variable((gidx0*4+lidx2*2+lidx3)//10, 0, 3, ("(((gidx0*2)+lidx2)//5)", "((lidx2+(gidx0*2))//5)")) + self.helper_test_variable((gidx0*4+lidx2*2+1)//10, 0, 3, "(((gidx0*2)+lidx2)//5)") + self.helper_test_variable((gidx0*4+lidx2*2+lidx3)//10, 0, 3, "(((gidx0*2)+lidx2)//5)") self.helper_test_variable((gidx0*2+lidx2)//10, 0, 1, "(gidx0//5)") def test_sum_div_complex3(self): gidx0 = Variable("gidx0", 0, 7) lidx2 = Variable("lidx2", 0, 12) lidx3 = Variable("lidx3", 0, 1) - self.helper_test_variable((gidx0*4+lidx2*2+lidx3)//12, 0, 4, ("(((lidx2//2)+gidx0)//3)", "((gidx0+(lidx2//2))//3)")) - self.helper_test_variable((lidx2*2+gidx0*4+lidx3)//12, 0, 4, ("(((lidx2//2)+gidx0)//3)", "((gidx0+(lidx2//2))//3)")) + self.helper_test_variable((gidx0*4+lidx2*2+lidx3)//12, 0, 4, "(((lidx2//2)+gidx0)//3)") + self.helper_test_variable((lidx2*2+gidx0*4+lidx3)//12, 0, 4, "(((lidx2//2)+gidx0)//3)") @unittest.expectedFailure # TODO: improve nest_div_by_smallest_factor def test_sum_div_complex4(self): gidx0 = Variable("gidx0", 0, 2) lidx2 = Variable("lidx2", 0, 12) lidx3 = Variable("lidx3", 0, 12) - self.helper_test_variable((gidx0*3+lidx2*19+lidx3*38)//(3*19), 0, 12, ("((lidx2+(lidx3*2))//3)")) + self.helper_test_variable((gidx0*3+lidx2*19+lidx3*38)//(3*19), 0, 12, "((lidx2+(lidx3*2))//3)") def test_sum_mul_distribute(self): gidx0 = Variable("gidx0", 0, 7) lidx2 = Variable("lidx2", 0, 12) lidx3 = Variable("lidx3", 0, 1) - self.helper_test_variable((gidx0+lidx2+lidx3)*4, 0, 80, - ("(((gidx0*4)+(lidx2*4))+(lidx3*4))","((lidx3*4)+((gidx0*4)+(lidx2*4)))")) + self.helper_test_variable((gidx0+lidx2+lidx3)*4, 0, 80, "(((gidx0*4)+(lidx2*4))+(lidx3*4))") @unittest.expectedFailure def test_variable_divmod(self): @@ -662,7 +661,7 @@ class TestSymbolic(unittest.TestCase): idx = Variable("idx", 0, 24) self.helper_test_variable(idx//4, 0, 6, "(idx//4)") # TODO: simplify the true branch - self.helper_test_variable((idx<4).where(idx//4, idx.const_like(-1)), -1, 6, "((idx//4) if (idx<4) else -1)") + self.helper_test_variable((idx<4).where(idx//4, idx.const_like(-1)), -1, 6, "(idx<4).where((idx//4), -1)") def test_idiv_lt(self): idx = Variable("idx", 0, 24) @@ -681,8 +680,8 @@ class TestSymbolic(unittest.TestCase): self.helper_test_variable((a*3+b*4<1).ne(True), 0, 1, "(((a+b)<1)!=True)") self.helper_test_variable((a*(-3)+b*4<1).ne(True), 0, 1, "((((a*-3)+(b*4))<1)!=True)") # negative coeff, should not be simplified self.helper_test_variable((a*3+d*4<1).ne(True), 0, 1, "((((a*3)+(d*4))<1)!=True)") # var can be negative, should not be simplified - self.helper_test_variable((a+b+c*2<1).ne(True), 0, 1, ("((((a+b)+c)<1)!=True)", "(((c+(a+b))<1)!=True)", '(((b+(a+c))<1)!=True)')) - self.helper_test_variable((a+b*2+c*4<1).ne(True), 0, 1, ("((((a+b)+c)<1)!=True)", "(((c+(a+b))<1)!=True)", '(((b+(a+c))<1)!=True)')) + self.helper_test_variable((a+b+c*2<1).ne(True), 0, 1, "((((a+b)+c)<1)!=True)") + self.helper_test_variable((a+b*2+c*4<1).ne(True), 0, 1, "((((a+b)+c)<1)!=True)") def test_where_removal(self): cond = Variable("a", 0, 3) < 2 @@ -700,30 +699,30 @@ class TestSymbolic(unittest.TestCase): c = Variable("c", 0, 3) aa = cond.where(a, a.ufix(0)) bb = cond.where(b, b.ufix(1)) - self.helper_test_variable(aa, 0, 3, "(a if (x<2) else 0)") - self.helper_test_variable(bb, 0, 3, "(b if (x<2) else 1)") - self.helper_test_variable(aa+bb, 0, 6, "((a+b) if (x<2) else 1)") - self.helper_test_variable(aa.maximum(bb), 0, 3, "(max(a, b) if (x<2) else 1)") - self.helper_test_variable((c+aa)+bb, 0, 9, "(c+((a+b) if (x<2) else 1))") + self.helper_test_variable(aa, 0, 3, "(x<2).where(a, 0)") + self.helper_test_variable(bb, 0, 3, "(x<2).where(b, 1)") + self.helper_test_variable(aa+bb, 0, 6, "(x<2).where((a+b), 1)") + self.helper_test_variable(aa.maximum(bb), 0, 3, "(x<2).where(a.maximum(b), 1)") + self.helper_test_variable((c+aa)+bb, 0, 9, "(c+(x<2).where((a+b), 1))") # not combining because it increased total ALU cc = cond.where(c, c+1) - self.helper_test_variable(bb+cc, 0, 7, "((b if (x<2) else 1)+(c if (x<2) else (c+1)))") + self.helper_test_variable(bb+cc, 0, 7, "((x<2).where(b, 1)+(x<2).where(c, (c+1)))") # not combining # TODO: can combine if it can further simplify? ab = cond.where(a, b) ba = cond.where(b, a) - self.helper_test_variable(ab+ba, 0, 6, "((a if (x<2) else b)+(b if (x<2) else a))") + self.helper_test_variable(ab+ba, 0, 6, "((x<2).where(a, b)+(x<2).where(b, a))") # not combining # TODO: can combine if one is identity element const - self.helper_test_variable(aa+ab, 0, 6, "((a if (x<2) else b)+(a if (x<2) else 0))") + self.helper_test_variable(aa+ab, 0, 6, "((x<2).where(a, b)+(x<2).where(a, 0))") def test_negation_in_where(self): cond = Variable("x", 0, 3) < 2 a = Variable("a", 0, 3) b = Variable("b", 0, 3) w = cond.logical_not().where(a, b) - self.helper_test_variable(w, 0, 3, "(b if (x<2) else a)") + self.helper_test_variable(w, 0, 3, "(x<2).where(b, a)") def test_neg_in_comp(self): a = Variable("a", 0, 3) @@ -750,7 +749,7 @@ class TestSymbolic(unittest.TestCase): a = Variable("a", 0, 3) b = Variable("b", 0, 3) expr = cond1.where(cond2.where(a, b), b) - self.helper_test_variable(expr, 0, 3, "(a if ((s<6)&(2 (a if (s<5) else b) - self.helper_test_variable(expr, 0, 3, "(a if (s<5) else b)") + self.helper_test_variable(expr, 0, 3, "(s<5).where(a, b)") def test_symbolic_div(self): # from symbolic arange @@ -774,7 +773,7 @@ class TestSymbolic(unittest.TestCase): a = Variable("a", 1, 10, dtypes.float) # TODO: bounds for reciprocal # TODO: should z3 work? - self.helper_test_variable(2*(2*a).reciprocal(), -math.inf, math.inf, "(1/a)", test_z3=False) + self.helper_test_variable(2*(2*a).reciprocal(), -math.inf, math.inf, "a.reciprocal()", test_z3=False) def test_trunc_noop(self): a = Variable("a", 1, 10, dtypes.int) @@ -783,8 +782,8 @@ class TestSymbolic(unittest.TestCase): def test_do_math_in_int32(self): a = Variable("a", 1, 10, dtypes.int) b = Variable("b", 1, 10, dtypes.int) - self.helper_test_variable(a.cast(dtypes.long)+b.cast(dtypes.long), 2, 20, "(long)((a+b))") - self.helper_test_variable(a.cast(dtypes.long)*b.cast(dtypes.long), 1, 100, "(long)((a*b))") + self.assertIn((a.cast(dtypes.long)+b.cast(dtypes.long)).render(), "(long)((a+b))") + self.assertIn((a.cast(dtypes.long)*b.cast(dtypes.long)).render(), "(long)((a*b))") class TestSymbolicNumeric(unittest.TestCase): def helper_test_numeric(self, f): diff --git a/test/unit/test_viz.py b/test/unit/test_viz.py index fbfc37e76f..5ba00735eb 100644 --- a/test/unit/test_viz.py +++ b/test/unit/test_viz.py @@ -30,6 +30,7 @@ class BaseTestViz(unittest.TestCase): # clear the global context for lst in [tracked_keys, tracked_ctxs, active_rewrites, _name_cnt]: lst.clear() Buffer.profile_events.clear() + cpu_events.clear() self.tms = TRACK_MATCH_STATS.value self.profile = PROFILE.value TRACK_MATCH_STATS.value = 2 @@ -462,5 +463,21 @@ class TestVizMemoryLayout(BaseTestViz): self.assertEqual(ret["peak"], 2) self.assertEqual(len(ret["events"]), 4) + def test_free_last(self): + bufs = [] + for _ in range(3): + bufs.append(_alloc(1)) + profile_marker("alloc") + device = bufs[0].device + while bufs: + b = bufs.pop() + del b + profile_marker("free") + profile = load_profile(cpu_events+Buffer.profile_events) + ret = profile["layout"][f"{device} Memory"] + self.assertEqual(ret["peak"], 3) + self.assertEqual(len(ret["events"]), 6) + self.assertEqual(len(profile["markers"]), 6) + if __name__ == "__main__": unittest.main() diff --git a/test/unit/test_winograd.py b/test/unit/test_winograd.py index 7f419b838c..d8909f7620 100644 --- a/test/unit/test_winograd.py +++ b/test/unit/test_winograd.py @@ -42,7 +42,7 @@ class TestWinograd(unittest.TestCase): 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), 5) def test_counters(self): IC, OC, X, Y = 4,4,9,9 diff --git a/tinygrad/apps/llm.py b/tinygrad/apps/llm.py index a718170259..df0d6d6db7 100644 --- a/tinygrad/apps/llm.py +++ b/tinygrad/apps/llm.py @@ -1,63 +1,61 @@ from __future__ import annotations -import sys, argparse, typing, re, itertools, unicodedata +import sys, argparse, typing, re, unicodedata from tinygrad import Tensor, nn, UOp, TinyJit, getenv, helpers -def gpt2_decode_vocab(voc: dict[str, int]): # https://github.com/openai/gpt-2/blob/9b63575ef42771a015060c964af2c3da4cf7c8ab/src/encoder.py#L9 - c2b = { chr(cp): cp for cp in itertools.chain(range(ord("!"), ord("~")+1), range(ord("¡"), ord("¬")+1), range(ord("®"), ord("ÿ")+1)) } - c2b.update({ chr(256+off): cp for off, cp in enumerate(cp for cp in range(256) if chr(cp) not in c2b) }) - return { bytes(c2b[c] for c in tok): tid for tok, tid in voc.items() } - -def get_llama_re(): - def ucat_range(pre: str): return "".join(re.escape(chr(cp)) for cp in range(sys.maxunicode + 1) if unicodedata.category(chr(cp)).startswith(pre)) - r_ws, r_p_N, r_p_L = r"\t\n\x0b\x0c\r\x85" + ucat_range("Z"), ucat_range("N"), ucat_range("L") - # https://github.com/ggml-org/llama.cpp/blob/94933c8c2eeaa9a7983e3f6c08af76bd86724094/src/llama-vocab.cpp#L286 - return "(?i:'s|'t|'re|'ve|'m|'ll|'d)|" + \ - f"[^\\r\\n{r_p_N}{r_p_L}]?[{r_p_L}]+|[{r_p_N}]{{1,3}}| ?[^{r_ws}{r_p_N}{r_p_L}]+[\\r\\n]*|[{r_ws}]*[\\r\\n]+|[{r_ws}]+(?![^{r_ws}])|[{r_ws}]+" - class SimpleTokenizer: - def __init__(self, pat: str, normal_tokens: dict[bytes, int], special_tokens: dict[str, int]): - self._normal_tokens, self._special_tokens, self._pat = normal_tokens, special_tokens, re.compile(pat) - self._tok2str = { tid: tok.encode() for tok, tid in special_tokens.items() } | { tid: tok for tok, tid in normal_tokens.items() } - self._special_re = re.compile("|".join(re.escape(tok) for tok in self._special_tokens.keys()) if special_tokens else r"(?!)") + def __init__(self, normal_tokens:dict[str, int], special_tokens:dict[str, int]): + # https://github.com/openai/gpt-2/blob/9b63575ef42771a015060c964af2c3da4cf7c8ab/src/encoder.py#L9 + bs = [*range(33, 127), *range(161, 173), *range(174, 256)] # bytes that map to themselves + self._byte_decoder = {chr(b): b for b in bs} | {chr(256+i): b for i,b in enumerate(b for b in range(256) if b not in bs)} + + # https://github.com/ggml-org/llama.cpp/blob/94933c8c2eeaa9a7983e3f6c08af76bd86724094/src/llama-vocab.cpp#L286 + def ucat_range(pre: str): return "".join(re.escape(chr(cp)) for cp in range(sys.maxunicode + 1) if unicodedata.category(chr(cp)).startswith(pre)) + r_ws, r_p_N, r_p_L = r"\t\n\x0b\x0c\r\x85" + ucat_range("Z"), ucat_range("N"), ucat_range("L") + self._split_to_word = re.compile("(?i:'s|'t|'re|'ve|'m|'ll|'d)|" + \ + f"[^\\r\\n{r_p_N}{r_p_L}]?[{r_p_L}]+|[{r_p_N}]{{1,3}}| ?[^{r_ws}{r_p_N}{r_p_L}]+[\\r\\n]*|[{r_ws}]*[\\r\\n]+|[{r_ws}]+(?![^{r_ws}])|[{r_ws}]+") + self._split_to_sentence = re.compile("|".join(re.escape(tok) for tok in special_tokens.keys()) if special_tokens else r"(?!)") + + self._normal_tokens = {bytes(self._byte_decoder[c] for c in tok): tid for tok, tid in normal_tokens.items()} + self._special_tokens = special_tokens + self._tok2bytes = {tid: tok for tok, tid in self._normal_tokens.items()} | {tid: tok.encode() for tok, tid in self._special_tokens.items()} @staticmethod - def from_gguf_kv(kv: dict): + def from_gguf_kv(kv:dict): # https://github.com/ggml-org/llama.cpp/blob/94933c8c2eeaa9a7983e3f6c08af76bd86724094/src/llama-vocab.cpp#L1818-L1820 if kv["tokenizer.ggml.pre"] not in ("llama3","llama-v3","llama-bpe"): raise ValueError(f"Invalid tokenizer preset '{kv['tokenizer.ggml.pre']}'") vocab: typing.Iterable[tuple[str, int]] = ((tok, idx) for idx, tok in enumerate(kv["tokenizer.ggml.tokens"])) normal_tokens, special_tokens = helpers.partition(vocab, lambda e: kv["tokenizer.ggml.token_type"][e[1]] == 1) - return SimpleTokenizer(get_llama_re(), gpt2_decode_vocab(dict(normal_tokens)), dict(special_tokens)) + return SimpleTokenizer(dict(normal_tokens), dict(special_tokens)) - def encode(self, text: str): + def _encode_word(self, word:bytes) -> list[int]: + if (early_token:=self._normal_tokens.get(word)) is not None: return [early_token] + parts = [bytes([b]) for b in word] + # greedily merge any parts that we can + while True: + i = min([(sys.maxsize, -1)] + [(self._normal_tokens.get(parts[j]+parts[j+1], sys.maxsize), j) for j in range(len(parts)-1)])[1] + if i == -1: break + parts[i:i+2] = [parts[i] + parts[i+1]] + try: return [self._normal_tokens[p] for p in parts] + except KeyError: raise RuntimeError("token not found") + def _encode_sentence(self, chunk:str) -> list[int]: + return [tok for word in self._split_to_word.findall(chunk) for tok in self._encode_word(word.encode())] + def encode(self, text:str) -> list[int]: tokens: list[int] = [] pos = 0 - for match in self._special_re.finditer(text): + for match in self._split_to_sentence.finditer(text): tokens.extend(self._encode_sentence(text[pos:match.start(0)]) + [self._special_tokens[text[match.start(0):match.end(0)]]]) pos = match.end(0) return tokens + self._encode_sentence(text[pos:]) - def decode(self, ids: list[int]) -> str: return b''.join(self._tok2str[tid] for tid in ids).decode() + def decode(self, ids:list[int]) -> str: return b''.join(self._tok2bytes[tid] for tid in ids).decode() def role(self, role:str): return self.encode("<|start_header_id|>" + role + "<|end_header_id|>\n\n") - def _encode_sentence(self, chunk: str): return [ tok for word in self._pat.findall(chunk) for tok in self._encode_word(word.encode()) ] - def _encode_word(self, word: bytes): - if (early_token:=self._normal_tokens.get(word)) is not None: return [early_token] - parts = [word[i:i+1] for i in range(len(word))] - while True: - min_tid, min_idx = 2**32, -1 - for idx, (p1, p2) in enumerate(zip(parts[:-1], parts[1:])): - tid = self._normal_tokens.get(p1 + p2, min_tid) - if tid < min_tid: min_tid, min_idx = tid, idx - if min_idx == -1: break - parts = parts[:min_idx] + [parts[min_idx] + parts[min_idx+1]] + parts[min_idx+2:] - try: return [ self._normal_tokens[p] for p in parts ] - except KeyError: raise RuntimeError("token not found") - def apply_rope(x:Tensor, start_pos:int|UOp, base:float = 10000.0) -> Tensor: B, H, T, Hd = x.shape - assert (Hd & 1) == 0, "RoPE requires an even head dimension" + assert isinstance(Hd, int) and (Hd & 1) == 0, "RoPE requires an even head dimension" half = Hd // 2 - angles = (Tensor.arange(T, dtype="float32") + start_pos)[:, None] * (base ** (-(Tensor.arange(half, dtype="float32") / half)))[None, :] + t_start_pos = start_pos if isinstance(start_pos, int) else Tensor(start_pos) + angles = (Tensor.arange(T, dtype="float32") + t_start_pos)[:, None] * (base ** (-(Tensor.arange(half, dtype="float32") / half)))[None, :] # contiguous here allows RoPE to be pruned in the JIT cos, sin = angles.cos().reshape(1, 1, T, half).cast(x.dtype).contiguous(), angles.sin().reshape(1, 1, T, half).cast(x.dtype).contiguous() x_pairs = x.reshape(B, H, T, half, 2) diff --git a/tinygrad/codegen/opt/heuristic.py b/tinygrad/codegen/opt/heuristic.py index cfc61dd511..b7b87c3120 100644 --- a/tinygrad/codegen/opt/heuristic.py +++ b/tinygrad/codegen/opt/heuristic.py @@ -178,7 +178,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler: if k.opts.has_threads and k.opts.global_max is not None: for threads in [32,16,12,8,6,5,4,3,2]: - # Skip is too many threads. Heuristic: use about 128K ops per thread + # Skip if too many threads. Heuristic: use about 128K ops per thread if threads > k.opts.global_max[0] or resolve(prod(k.full_shape) // (128 << 10) < threads): continue for axis in k.axes_of(AxisType.LOOP): if k.full_shape[axis] % threads == 0: diff --git a/tinygrad/codegen/opt/search.py b/tinygrad/codegen/opt/search.py index b5c3a1ebce..21cce836f3 100644 --- a/tinygrad/codegen/opt/search.py +++ b/tinygrad/codegen/opt/search.py @@ -165,15 +165,22 @@ def beam_search(lin:Scheduler, rawbufs:list[Buffer], amt:int, allow_test_size=Tr if isinstance(e, RuntimeError): continue raise timed_lins.append((acted_lins[i], min(tms))) - if BEAM_DEBUG > 1: print(f"{time.perf_counter() - st:7.2f}s: {i:5d} {len(cast(list, p.uops)):5d} uops {time_to_str(compile_et, w=12)} compile/{time_to_str(timed_lins[-1][1], w=12)} run {len(timed_lins):4d}/{len(acted_lins):4d} {timed_lins[-1][0].colored_shape()}") # noqa: E501 - elif DEBUG >= 2: print(f"\r{time.perf_counter() - st:7.2f}s: {time_to_str(timed_lins[-1][1], w=12)} {len(timed_lins):4d}/{len(acted_lins):4d} {timed_lins[-1][0].colored_shape()}\033[K", end="") # noqa: E501 + if BEAM_DEBUG > 1: + print(f"{time.perf_counter() - st:7.2f}s: {i:5d} {len(cast(list, p.uops)):5d} uops", + f"{time_to_str(compile_et, w=12)} compile/{time_to_str(timed_lins[-1][1], w=12)} run", + f" {len(timed_lins):4d}/{len(acted_lins):4d} {timed_lins[-1][0].colored_shape()}") + elif DEBUG >= 2: + print(f"\r{time.perf_counter() - st:7.2f}s: {time_to_str(timed_lins[-1][1], w=12)}", + f" {len(timed_lins):4d}/{len(acted_lins):4d} {timed_lins[-1][0].colored_shape()}\033[K", end="") # done opts = sorted(timed_lins, key=lambda x: x[1]) exiting = len(opts) == 0 or (opts[0][1] < min_progress) or (len(beam) > 0 and ((beam[0][1]-opts[0][1]) < min_progress)) if not exiting: beam = opts[:amt] elif len(opts) > 0 and opts[0][1] < beam[0][1]: beam = opts[:1] - if DEBUG >= 2: print(f"\r{time.perf_counter() - st:7.2f}s:", colored(time_to_str(beam[0][1], w=12), "green" if exiting else None), f"from {len(acted_lins):3d} -> {len(opts):3d} actions\033[K", beam[0][0].colored_shape()) # noqa: E501 + if DEBUG >= 2: + print(f"\r{time.perf_counter() - st:7.2f}s:", colored(time_to_str(beam[0][1], w=12), "green" if exiting else None), + f"from {len(acted_lins):3d} -> {len(opts):3d} actions\033[K", beam[0][0].colored_shape()) except KeyboardInterrupt as e: if beam_pool is not None: beam_pool.terminate() raise e diff --git a/tinygrad/device.py b/tinygrad/device.py index 8d93252918..3e1788c946 100644 --- a/tinygrad/device.py +++ b/tinygrad/device.py @@ -23,7 +23,7 @@ class _Device: def __getitem__(self, ix:str) -> Compiled: return self.__get_canonicalized_item(self.canonicalize(ix)) @functools.cache # this class is a singleton, pylint: disable=method-cache-max-size-none def __get_canonicalized_item(self, ix:str) -> Compiled: - assert ALLOW_DEVICE_USAGE or ix.split(":")[0] in ["DISK", "NPY", "PYTHON"], f"usage of device {ix} disallowed" + assert ALLOW_DEVICE_USAGE or ix.split(":")[0] in ["DISK", "TINYFS", "NPY", "PYTHON"], f"usage of device {ix} disallowed" base = (__package__ or __name__).split('.')[0] # tinygrad x = ix.split(":")[0].lower() ret = [cls for cname, cls in inspect.getmembers(importlib.import_module(f'{base}.runtime.ops_{x}')) \ @@ -39,7 +39,7 @@ class _Device: @functools.cached_property def DEFAULT(self) -> str: dev = [dev] if (dev:=getenv("DEV", "").upper()) else [] - from_env = dedup(dev + [d for d in self._devices if d not in ["DISK", "NPY"] and getenv(d) == 1]) + from_env = dedup(dev + [d for d in self._devices if d not in ["DISK", "TINYFS", "NPY"] and getenv(d) == 1]) assert len(from_env) < 2, f"multiple devices set in env: {from_env}" if len(from_env) == 1: return from_env[0] try: @@ -137,16 +137,14 @@ class Buffer: else: self._buf = opaque if opaque is not None else self.allocator.alloc(self.nbytes, self.options) if not self.device.startswith("DISK"): GlobalCounters.mem_used += self.nbytes - if PROFILE: - self._prof_num = num = len(Buffer.profile_events) - Buffer.profile_events.append(ProfilePointEvent(self.device, "alloc", num, {"dtype":self.dtype, "sz":self.size})) + if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "alloc", self.trace_num, {"dtype":self.dtype, "sz":self.size})) return self def deallocate(self): assert hasattr(self, '_buf'), "buffer must be allocated to deallocate" if DEBUG is not None and DEBUG >= 7: print(f"buffer: deallocate {self.nbytes} bytes on {self.device}") if self._base is None and (self.options is None or self.options.external_ptr is None): if GlobalCounters is not None and not self.device.startswith("DISK"): GlobalCounters.mem_used -= self.nbytes - if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "free", self._prof_num)) + if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "free", self.trace_num)) self.allocator.free(self._buf, self.nbytes, self.options) elif self._base is not None: self._base.allocated_views -= 1 del self._buf @@ -160,6 +158,10 @@ class Buffer: self.copyout(memoryview(buf)) return self.__class__, (self.device, self.size, self.dtype, None, self.options, buf, self.uop_refcount) @property + def trace_num(self) -> int: + if not hasattr(self, '_trace_num'): self._trace_num = len(Buffer.profile_events) + return self._trace_num + @property def nbytes(self): return self.size*self.dtype.itemsize def __del__(self): (not hasattr(self, '_buf')) or self.deallocate() def __repr__(self): diff --git a/tinygrad/engine/realize.py b/tinygrad/engine/realize.py index 50474a6284..053bcd9bf6 100644 --- a/tinygrad/engine/realize.py +++ b/tinygrad/engine/realize.py @@ -121,7 +121,7 @@ class BufferCopy(Runner): getattr(src.allocator.dev, 'fd', None) is not None and dest.allocator.supports_copy_from_disk if src.device.startswith("DISK") and hasattr(dest.allocator, 'copy_from_disk') and disk_supports_fast_copyout and src.nbytes >= 4096: dest.allocator.copy_from_disk(dest._buf, src._buf, src.nbytes) - elif src.device.startswith("DISK") and hasattr(dest.allocator, '_as_buffer'): + elif (src.device.startswith("DISK") or src.device.startswith("TINYFS")) and hasattr(dest.allocator, '_as_buffer'): # fast(ish) path, uses readinto in diskbuffers src.allocator._copyout(dest.allocator._as_buffer(dest._buf), src._buf) else: @@ -165,7 +165,9 @@ class ExecItem: def run(self, _var_vals:dict[str, int]|None=None, wait=False, jit=False, do_update_stats=True) -> float|None: var_vals = self.fixedvars if _var_vals is None else (_var_vals|self.fixedvars) bufs = [cast(Buffer, x) for x in self.bufs] if jit else [cast(Buffer, x).ensure_allocated() for x in self.bufs] - if PROFILE: cpu_events.append(ProfilePointEvent(self.prg.device, "exec", self.prg.display_name, {"metadata":self.metadata, "var_vals":var_vals})) + if PROFILE: + payload = {"metadata":self.metadata, "var_vals":var_vals, "bufs":[b.trace_num for b in bufs]} + cpu_events.append(ProfilePointEvent(self.prg.device, "exec", self.prg.display_name, payload)) et = self.prg(bufs, var_vals, wait=wait or DEBUG >= 2) if do_update_stats: GlobalCounters.kernel_count += 1 diff --git a/tinygrad/gradient.py b/tinygrad/gradient.py index c538555ad2..3d68868fdb 100644 --- a/tinygrad/gradient.py +++ b/tinygrad/gradient.py @@ -39,7 +39,7 @@ pm_gradient = PatternMatcher([ (UPat(Ops.EXPAND, name="ret"), lambda ctx, ret: (ctx.r(Ops.ADD, tuple(i for i,(si,so) in enumerate(zip(ret.src[0].shape, ret.arg)) if si!=so)),)), (UPat(Ops.MULTI, name="ret"), lambda ctx, ret: ctx.shard(ret.device, ret.axis).src), # there's no gradient for bitcast - (UPat(Ops.BITCAST), lambda ctx: (None,)), + (UPat(Ops.BITCAST), lambda: (None,)), ]) def _deepwalk(root:UOp, targets:set[UOp]) -> list[UOp]: diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index d9c21933d6..766ab4681b 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -23,10 +23,13 @@ def argfix(*x): if len(x) != 1: raise ValueError(f"bad arg {x}") return tuple(x[0]) return x -def argsort(x): return type(x)(sorted(range(len(x)), key=x.__getitem__)) # https://stackoverflow.com/questions/3382352/equivalent-of-numpy-argsort-in-basic-python +# https://stackoverflow.com/questions/3382352/equivalent-of-numpy-argsort-in-basic-python +def argsort(x): return type(x)(sorted(range(len(x)), key=x.__getitem__)) def all_same(items:tuple[T, ...]|list[T]): return all(x == items[0] for x in items) def all_int(t: Sequence[Any]) -> TypeGuard[tuple[int, ...]]: return all(isinstance(s, int) for s in t) -def colored(st, color:str|None, background=False): return f"\u001b[{10*background+60*(color.upper() == color)+30+['black', 'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white'].index(color.lower())}m{st}\u001b[0m" if color is not None else st # replace the termcolor library with one line # noqa: E501 +def colored(st, color:str|None, background=False): # replace the termcolor library + colors = ['black', 'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white'] + return f"\u001b[{10*background+60*(color.upper() == color)+30+colors.index(color.lower())}m{st}\u001b[0m" if color is not None else st def colorize_float(x: float): return colored(f"{x:7.2f}x", 'green' if x < 0.75 else 'red' if x > 1.15 else 'yellow') def time_to_str(t:float, w=8) -> str: return next((f"{t * d:{w}.2f}{pr}" for d,pr in [(1, "s "),(1e3, "ms")] if t > 10/d), f"{t * 1e6:{w}.2f}us") def ansistrip(s:str): return re.sub('\x1b\\[(K|.*?m)', '', s) @@ -96,7 +99,7 @@ def suppress_finalizing(func): if not getattr(sys, 'is_finalizing', lambda: True)(): raise # re-raise if not finalizing return wrapper -def unwrap_class_type(cls_t:T): return cls_t.func if isinstance(cls_t, functools.partial) else cls_t +def unwrap_class_type(cls_t): return cls_t.func if isinstance(cls_t, functools.partial) else cls_t def pluralize(st:str, cnt:int): return f"{cnt} {st}"+('' if cnt == 1 else 's') @@ -150,7 +153,7 @@ CORRECT_DIVMOD_FOLDING, FUSE_OPTIM = ContextVar("CORRECT_DIVMOD_FOLDING", 0), Co ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE = ContextVar("ALLOW_DEVICE_USAGE", 1), ContextVar("MAX_BUFFER_SIZE", 0) FUSE_ATTENTION = 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_COUNT = ContextVar("CPU_COUNT", max(1, len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1))) CPU_LLVM, AMD_LLVM = ContextVar("CPU_LLVM", 0), ContextVar("AMD_LLVM", 1) VIZ = PROFILE = ContextVar("VIZ", 0) SPEC = ContextVar("SPEC", 0) @@ -218,11 +221,12 @@ class TracingKey: class ProfileEvent: pass @dataclass -class ProfileRangeEvent(ProfileEvent): device:str; name:str|TracingKey; st:decimal.Decimal; en:decimal.Decimal|None=None; is_copy:bool=False # noqa: E702 +class ProfileRangeEvent(ProfileEvent): + device:str; name:str|TracingKey; st:decimal.Decimal; en:decimal.Decimal|None=None; is_copy:bool=False # noqa: E702 @dataclass(frozen=True) -class ProfilePointEvent(ProfileEvent): device:str; name:str; key:Any; arg:dict=field(default_factory=dict); \ - ts:decimal.Decimal=field(default_factory=perf_counter_us) # noqa: E702 +class ProfilePointEvent(ProfileEvent): + device:str; name:str; key:Any; arg:dict=field(default_factory=dict); ts:decimal.Decimal=field(default_factory=perf_counter_us) # noqa: E702 cpu_events:list[ProfileEvent] = [] @contextlib.contextmanager @@ -281,7 +285,8 @@ def diskcache_put(table:str, key:dict|str|int, val:Any, prepickled=False): ltypes = ', '.join(f"{k} {TYPES[type(key[k])]}" for k in key.keys()) cur.execute(f"CREATE TABLE IF NOT EXISTS '{table}_{VERSION}' ({ltypes}, val blob, PRIMARY KEY ({', '.join(key.keys())}))") _db_tables.add(table) - cur.execute(f"REPLACE INTO '{table}_{VERSION}' ({', '.join(key.keys())}, val) VALUES ({', '.join(['?']*len(key))}, ?)", tuple(key.values()) + (val if prepickled else pickle.dumps(val), )) # noqa: E501 + cur.execute(f"REPLACE INTO '{table}_{VERSION}' ({', '.join(key.keys())}, val) VALUES ({', '.join(['?']*len(key))}, ?)", + tuple(key.values()) + (val if prepickled else pickle.dumps(val),)) conn.commit() cur.close() return val @@ -347,10 +352,10 @@ def capstone_flatdump(lib: bytes): print(f"{instr.address:#08x}: {instr.mnemonic}\t{instr.op_str}") sys.stdout.flush() -def wait_cond(cb, value=True, timeout_ms=10000, msg="") -> bool: +def wait_cond(cb, *args, value=True, timeout_ms=10000, msg="") -> bool: start_time = int(time.perf_counter() * 1000) while int(time.perf_counter() * 1000) - start_time < timeout_ms: - if (val:=cb()) == value: return val + if (val:=cb(*args)) == value: return val raise TimeoutError(f"{msg}. Timed out after {timeout_ms} ms, condition not met: {val} != {value}") # *** ctypes helpers diff --git a/tinygrad/nn/__init__.py b/tinygrad/nn/__init__.py index d32a3d5e2f..b27ab036c0 100644 --- a/tinygrad/nn/__init__.py +++ b/tinygrad/nn/__init__.py @@ -223,7 +223,7 @@ class InstanceNorm: print(t.mean().item(), t.std().item()) ``` """ - def __init__(self, num_features:int, eps=1e-5, affine=True): + def __init__(self, num_features:int, eps:float=1e-5, affine:bool=True): self.num_features, self.eps = num_features, eps self.weight: Tensor|None = Tensor.ones(num_features) if affine else None self.bias: Tensor|None = Tensor.zeros(num_features) if affine else None @@ -249,16 +249,16 @@ class LayerNorm: print(t.mean().item(), t.std().item()) ``` """ - def __init__(self, normalized_shape:int|tuple[int, ...], eps=1e-5, elementwise_affine=True): + def __init__(self, normalized_shape:int|tuple[int, ...], eps:float=1e-5, elementwise_affine:bool=True): self.normalized_shape: tuple[int, ...] = make_tuple(normalized_shape, 1) - self.axis, self.eps, self.elementwise_affine = tuple(-1-i for i in range(len(self.normalized_shape))), eps, elementwise_affine + self.axis, self.eps = tuple(-1-i for i in range(len(self.normalized_shape))), eps self.weight: Tensor|None = Tensor.ones(*self.normalized_shape) if elementwise_affine else None self.bias: Tensor|None = Tensor.zeros(*self.normalized_shape) if elementwise_affine else None def __call__(self, x:Tensor) -> Tensor: assert self.normalized_shape == x.shape[-len(self.normalized_shape):], f"last dimensions of {x.shape} must match {self.normalized_shape}" x = x.layernorm(eps=self.eps, axis=self.axis) - if not self.elementwise_affine: return x + if self.weight is None or self.bias is None: return x return x * self.weight + self.bias class LayerNorm2d(LayerNorm): diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index 79b09c9b92..c3a8e1508d 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -108,7 +108,9 @@ class CStyleLanguage(Renderer): extra_matcher = extra_pm def render_kernel(self, function_name:str, kernel:list[str], bufs:list[tuple[str,tuple[DType,bool]]], uops:list[UOp], prefix=None) -> str: - tmp = "const sampler_t smp = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST;\n" if any(isinstance(dtype, ImageDType) for _,(dtype,_) in bufs) else "" # noqa: E501 + tmp = "" + if any(isinstance(dtype, ImageDType) for _,(dtype,_) in bufs): + tmp = "const sampler_t smp = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST;\n" buftypes = [(name, self.render_dtype(dtype, mutable)+self.buffer_suffix if isinstance(dtype, (ImageDType, PtrDType)) else self.arg_int_prefix if dtype == dtypes.int else None) for name,(dtype,mutable) in bufs] local_dims = [u.src[0] for u in uops if u.op is Ops.SPECIAL and u.arg[0] == "l"] @@ -229,10 +231,12 @@ class ClangRenderer(CStyleLanguage): # 'static' in C roughly means that function symbol isn't exported. LLVM puts those symbols at the end of object file which allows Clang JIT # to just jump at the start of a shellcode without having to deal with symbols or trampolines at all. This is better than having to inline # wmma function every time it is called or wasting complexity on a symbol parsing and a memory page on trampoline. - prefix += [f"""static {(out := self.render_dtype(dtype_in.vec(N*N)))} __{name}({self.render_dtype(dtype_in.vec(N))} data1, {self.render_dtype(dtype_in.vec(M))} data2, {out} data0){{ + out, dt1, dt2 = self.render_dtype(dtype_in.vec(N*N)), self.render_dtype(dtype_in.vec(N)), self.render_dtype(dtype_in.vec(M)) + prefix += [f"""static {out} __{name}({dt1} data1, {dt2} data2, {out} data0){{ AMX_SET(0);\n for(int ridx0 = 0; ridx0 < 16; ridx0++){{ AMX(4, (int *)(&data0), 0ull<<62 | (ridx0*4ull)<<56 | ridx0*64ull); }} AMX(0, (int *)(&data2), 0ull<<62); AMX(1, (int *)(&data1), 0ull<<62); AMX(12, 0, 0ull); - for(int ridx0 = 0; ridx0 < 16; ridx0++){{ AMX(5, (int *)(&data0), 0ull<<62 | (ridx0*4ull)<<56 | ridx0*64ull); }}\n AMX_SET(1);\n return data0;\n}}"""] # noqa: E501 + for(int ridx0 = 0; ridx0 < 16; ridx0++){{ AMX(5, (int *)(&data0), 0ull<<62 | (ridx0*4ull)<<56 | ridx0*64ull); }} + AMX_SET(1);\n return data0;\n}}"""] return prefix def _render_body(self, function_name, kernel, bufs, uops, pref=None) -> str: return super().render_kernel(function_name, kernel, bufs, uops, pref) def _render_entry(self, function_name:str, bufs:list[tuple[str,tuple[DType,bool]]]) -> str: return "" diff --git a/tinygrad/runtime/ops_amd.py b/tinygrad/runtime/ops_amd.py index d89eaff8b6..af239b8948 100644 --- a/tinygrad/runtime/ops_amd.py +++ b/tinygrad/runtime/ops_amd.py @@ -7,7 +7,7 @@ from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, H 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, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, suppress_finalizing, lo32, hi32, colored from tinygrad.renderer.cstyle import AMDRenderer from tinygrad.renderer.llvmir import AMDLLVMRenderer from tinygrad.runtime.autogen import kfd, hsa, pci, sqtt @@ -670,7 +670,7 @@ class PCIIface(PCIIfaceBase): gpus:ClassVar[list[str]] = [] def __init__(self, dev, dev_id): - super().__init__(dev, dev_id, vendor=0x1002, devices=[0x744c, 0x7480, 0x7550], bars=[0, 2, 5], vram_bar=0, + super().__init__(dev, dev_id, vendor=0x1002, devices=[0x744c, 0x7480, 0x7550, 0x7590], bars=[0, 2, 5], vram_bar=0, va_start=AMMemoryManager.va_allocator.base, va_size=AMMemoryManager.va_allocator.size) self._setup_adev(self.pci_dev.pcibus, self.pci_dev.map_bar(0), self.pci_dev.map_bar(2, fmt='Q'), self.pci_dev.map_bar(5, fmt='I')) self.pci_dev.write_config(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2) @@ -713,7 +713,7 @@ class PCIIface(PCIIfaceBase): def device_fini(self): self.dev_impl.fini() class USBIface(PCIIface): - def __init__(self, dev, dev_id): + def __init__(self, dev, dev_id): # pylint: disable=super-init-not-called self.dev = dev self.usb = ASM24Controller() self.bars = setup_pci_bars(self.usb, gpu_bus=4, mem_base=0x10000000, pref_mem_base=(32 << 30)) @@ -875,13 +875,12 @@ class AMDDevice(HCQCompiled): def _at_profile_finalize(self): if self.sqtt_enabled: wptrs_buf = self.allocator.alloc(round_up(len(self.sqtt_buffers), 0x1000), BufferSpec(cpu_access=True, nolru=True)) - wptrs = to_mv(wptrs_buf.va_addr, wptrs_buf.size) 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...') for i,buf0 in enumerate(self.sqtt_buffers): - wptr = ((struct.unpack('= 2: print(f'\t{self.device}: 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 diff --git a/tinygrad/runtime/ops_cl.py b/tinygrad/runtime/ops_cl.py index 8887c97f00..b89fdedcb4 100644 --- a/tinygrad/runtime/ops_cl.py +++ b/tinygrad/runtime/ops_cl.py @@ -23,10 +23,12 @@ class CLCompiler(Compiler): build_status: int = cl.clBuildProgram(program, 1, self.dev.device_id, None, cl.clBuildProgram.argtypes[4](), None) if build_status != 0: cl.clGetProgramBuildInfo(program, self.dev.device_id, cl.CL_PROGRAM_BUILD_LOG, 0, None, log_size := ctypes.c_size_t()) - cl.clGetProgramBuildInfo(program, self.dev.device_id, cl.CL_PROGRAM_BUILD_LOG, log_size.value, mstr := ctypes.create_string_buffer(log_size.value), None) # noqa: E501 + cl.clGetProgramBuildInfo(program, self.dev.device_id, cl.CL_PROGRAM_BUILD_LOG, + log_size.value, mstr := ctypes.create_string_buffer(log_size.value), None) raise CompileError(f"OpenCL Compile Error\n\n{mstr.value.decode()}") check(cl.clGetProgramInfo(program, cl.CL_PROGRAM_BINARY_SIZES, ctypes.sizeof(ctypes.c_size_t), binary_sizes := (ctypes.c_size_t * 1)(), None)) - check(cl.clGetProgramInfo(program, cl.CL_PROGRAM_BINARIES, ctypes.sizeof(ctypes.c_void_p), (ctypes.c_void_p * 1)(ctypes.addressof(binary := ctypes.create_string_buffer(binary_sizes[0]))), None)) # noqa: E501 + check(cl.clGetProgramInfo(program, cl.CL_PROGRAM_BINARIES, ctypes.sizeof(ctypes.c_void_p), + (ctypes.c_void_p * 1)(ctypes.addressof(binary := ctypes.create_string_buffer(binary_sizes[0]))), None)) check(cl.clReleaseProgram(program)) return bytes(binary) @@ -97,16 +99,22 @@ class CLDevice(Compiled): err = cl.clGetDeviceIDs(platform_ids[0], device_type, 0, None, num_devices := ctypes.c_uint32()) if err == 0 and num_devices.value != 0: break if DEBUG >= 1: print(f"CLDevice: got {num_platforms.value} platforms and {num_devices.value} devices") - CLDevice.device_ids = init_c_var((cl.cl_device_id * num_devices.value)(), lambda x: check(cl.clGetDeviceIDs(platform_ids[0], device_type, num_devices, x, None))) # noqa: E501 + CLDevice.device_ids = init_c_var((cl.cl_device_id * num_devices.value)(), + lambda x: check(cl.clGetDeviceIDs(platform_ids[0], device_type, num_devices, x, None))) self.device_id = CLDevice.device_ids[0 if ":" not in device else int(device.split(":")[1])] - self.device_name = (cl.clGetDeviceInfo(self.device_id, cl.CL_DEVICE_NAME, 256, buf := ctypes.create_string_buffer(256), None), buf.value.decode())[1] # noqa: E501 - self.driver_version = (cl.clGetDeviceInfo(self.device_id, cl.CL_DRIVER_VERSION, 256, buf := ctypes.create_string_buffer(256), None), buf.value.decode())[1] # noqa: E501 + self.device_name = (cl.clGetDeviceInfo(self.device_id, cl.CL_DEVICE_NAME, 256, + buf:=ctypes.create_string_buffer(256), None), buf.value.decode())[1] + self.driver_version = (cl.clGetDeviceInfo(self.device_id, cl.CL_DRIVER_VERSION, 256, + buf:=ctypes.create_string_buffer(256), None), buf.value.decode())[1] if DEBUG >= 1: print(f"CLDevice: opening {self.device_name} with version {self.driver_version}") self.context = checked(cl.clCreateContext(None, 1, self.device_id, cl.clCreateContext.argtypes[3](), None, status := ctypes.c_int32()), status) self.queue = checked(cl.clCreateCommandQueue(self.context, self.device_id, cl.CL_QUEUE_PROFILING_ENABLE, status), status) self.pending_copyin: list[memoryview] = [] - self.device_exts = (cl.clGetDeviceInfo(self.device_id, cl.CL_DEVICE_EXTENSIONS, 4096, ctypes.byref(buf := ctypes.create_string_buffer(4096)), ctypes.byref(total := ctypes.c_size_t())), ctypes.string_at(buf, size=total.value).decode())[1] # noqa: E501 + self.device_exts = (cl.clGetDeviceInfo(self.device_id, cl.CL_DEVICE_EXTENSIONS, 4096, + ctypes.byref(buf := ctypes.create_string_buffer(4096)), + ctypes.byref(total := ctypes.c_size_t())), + ctypes.string_at(buf, size=total.value).decode())[1] compilers = [(IntelRenderer if "cl_intel_subgroup_matrix_multiply_accumulate" in self.device_exts else OpenCLRenderer, functools.partial(CLCompiler, self, f"compile_cl_{hashlib.md5(self.device_name.encode() + self.driver_version.encode()).hexdigest()}"))] diff --git a/tinygrad/runtime/ops_cpu.py b/tinygrad/runtime/ops_cpu.py index 012a9e729e..3dc70103a8 100644 --- a/tinygrad/runtime/ops_cpu.py +++ b/tinygrad/runtime/ops_cpu.py @@ -105,8 +105,8 @@ class CPUAllocator(HCQAllocatorBase): else: addr = mv_address(buf:=mmap.mmap(-1, size, mmap.MAP_ANON | mmap.MAP_PRIVATE, mmap.PROT_READ | mmap.PROT_WRITE)) return HCQBuffer(va:=addr, sz:=size, meta=buf, view=MMIOInterface(va, sz, fmt='B'), owner=self.dev) def _as_buffer(self, src) -> memoryview: - self.dev.synchronize() - return to_mv(src.va_addr, src.size) + self.dev.synchronize() + return to_mv(src.va_addr, src.size) def _as_dmaref(self, buf): self.dev.synchronize() return DMACPURef(buf.va_addr, buf.size) diff --git a/tinygrad/runtime/ops_cuda.py b/tinygrad/runtime/ops_cuda.py index 7be380e5ef..7aa44dede8 100644 --- a/tinygrad/runtime/ops_cuda.py +++ b/tinygrad/runtime/ops_cuda.py @@ -10,7 +10,9 @@ if getenv("IOCTL"): import extra.nv_gpu_driver.nv_ioctl # noqa: F401 # pylint: if MOCKGPU:=getenv("MOCKGPU"): from test.mockgpu.cuda import cuda # type: ignore # pylint: disable=reimported def check(status): - if status != 0: raise RuntimeError(f"CUDA Error {status}, {ctypes.string_at(init_c_var(ctypes.POINTER(ctypes.c_char)(), lambda x: cuda.cuGetErrorString(status, ctypes.byref(x)))).decode()}") # noqa: E501 + if status != 0: + error = ctypes.string_at(init_c_var(ctypes.POINTER(ctypes.c_char)(), lambda x: cuda.cuGetErrorString(status, ctypes.byref(x)))).decode() + raise RuntimeError(f"CUDA Error {status}, {error}") def encode_args(args, vals) -> tuple[ctypes.Structure, ctypes.Array]: c_args = init_c_struct_t(tuple([(f'f{i}', cuda.CUdeviceptr_v2) for i in range(len(args))] + diff --git a/tinygrad/runtime/ops_remote.py b/tinygrad/runtime/ops_remote.py index 5c0c056a72..12c80cf255 100644 --- a/tinygrad/runtime/ops_remote.py +++ b/tinygrad/runtime/ops_remote.py @@ -424,7 +424,7 @@ class RemoteConnection: conns = RemoteConnection.all.keys() datas = {conn: conn.req.serialize() for conn in conns} reqs, hashes, hash_datas = sum(len(c.req._q) for c in conns), sum(len(c.req._h) for c in conns), sum(len(data) for data in datas.values()) - resps = [] + ret, resps = None, [] with Timing(f"*** send {reqs:-3d} requests {hashes:-3d} hashes with len {hash_datas/1024:.2f} kB in ", enabled=DEBUG>=3): for conn,data in datas.items(): conn.conn.request("POST", "/batch", data) for conn in datas.keys(): diff --git a/tinygrad/runtime/ops_tinyfs.py b/tinygrad/runtime/ops_tinyfs.py new file mode 100644 index 0000000000..69d5ff54e3 --- /dev/null +++ b/tinygrad/runtime/ops_tinyfs.py @@ -0,0 +1,137 @@ +import socket, uuid, json, asyncio, threading +from contextlib import asynccontextmanager +from tinygrad.device import Compiled, Allocator +from tinygrad.helpers import DEBUG, getenv +from tinygrad import Tensor + +TINYFS_ENDPOINT = getenv("TINYFS_ENDPOINT", "localhost:6767") + +class TinyFSDevice(Compiled): + def __init__(self, device:str): + self.op = device[len("tinyfs:"):].upper() + super().__init__(device, TinyFSAllocator(self), None, None, None) + + self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + self.sock.connect((TINYFS_ENDPOINT.rsplit(":", 1)[0], int(TINYFS_ENDPOINT.rsplit(":", 1)[1]))) + self.sfile = self.sock.makefile("rwb") + + # fetch node info + self.sfile.write(b"INFO\r\n") + self.sfile.flush() + info = self.sfile.readline() + self.node_info = json.loads(info) + if DEBUG >= 2: print(f"nodes: {self.node_info}") + + # spawn thread for async copyout + self.start_event = threading.Event() + self.t = threading.Thread(target=self._start_thread, daemon=True) + self.t.start() + self.start_event.wait() + + # connection pools + self.conn_pools: dict[str, asyncio.Queue] = {} + self.conn_pools_lock = asyncio.Lock() + + def finalize(self): + self.sfile.close() + + for pool in self.conn_pools.values(): + while not pool.empty(): + _, w = pool.get_nowait() + w.close() + asyncio.run_coroutine_threadsafe(w.wait_closed(), self.loop).result() + + if hasattr(self, "loop"): + self.loop.call_soon_threadsafe(self.loop.stop) + self.t.join() + + def _start_thread(self): + self.loop = asyncio.new_event_loop() + asyncio.set_event_loop(self.loop) + + self.start_event.set() + self.loop.run_forever() + self.loop.close() + + @asynccontextmanager + async def connection(self, loc): + if loc not in self.conn_pools: + await self.conn_pools_lock.acquire() + if loc not in self.conn_pools: + self.conn_pools[loc] = asyncio.Queue(nw:=getenv("ASYNC_COPY_WORKERS", 4)) + conn_tasks = [asyncio.open_connection(*self.node_info[loc][-1].rsplit(":", 1)) for _ in range(nw)] + connections = await asyncio.gather(*conn_tasks) + for reader, writer in connections: self.conn_pools[loc].put_nowait((reader, writer)) + self.conn_pools_lock.release() + + reader, writer = await self.conn_pools[loc].get() + try: + yield reader, writer + finally: + await self.conn_pools[loc].put((reader, writer)) + +class TinyFSBuffer: + def __init__(self, device:TinyFSDevice, size:int, offset=0, request_id=None, copyout_queue=None): + self.device, self.size, self.offset = device, size, offset + self.request_id: uuid.UUID|None = request_id + self.copyout_queue = copyout_queue or [] + def __repr__(self): return f"" + +class TinyFSAllocator(Allocator[TinyFSDevice]): + def _alloc(self, size, options): + return TinyFSBuffer(self.dev, size) + + def _copyin(self, dest:TinyFSBuffer, src:memoryview): + if DEBUG >= 2: print(f"Copying in {dest.size} bytes to TINYFS:{dest.device.op}") + self.dev.sfile.write(f"{dest.device.op}_IN {dest.size}\r\n".encode()) + + if dest.device.op == "STORE": + self.dev.sfile.flush() + dest.request_id = uuid.UUID(bytes=self.dev.sfile.read(16)) + if DEBUG >= 2: print(f"Request ID: {dest.request_id}") + + self.dev.sfile.write(src) + self.dev.sfile.flush() + + if dest.device.op == "LOAD": + locs = self.dev.sfile.readline() + locs = json.loads(locs) + + dest.copyout_queue = [] + for i, loc in enumerate(locs): + dest.copyout_queue.append((i, loc, src[i*16:(i+1)*16])) + + def _copyout(self, dest:memoryview, src:TinyFSBuffer): + if DEBUG >= 2: print(f"Copying out {src.size} bytes from TINYFS:{src.device.op}") + if src.device.op == "LOAD": + asyncio.run_coroutine_threadsafe(self._copyout_async(dest, src), src.device.loop).result() + else: + self.dev.sfile.write(f"{src.device.op}_OUT {src.size} {src.request_id}\r\n".encode()) + self.dev.sfile.flush() + src.request_id = uuid.UUID(bytes=self.dev.sfile.read(16)) + if DEBUG >= 2: print(f"Request ID: {src.request_id}") + self.dev.sfile.readinto(dest) + + async def _copyout_async(self, dest:memoryview, src:TinyFSBuffer): + async def _worker(item): + i, loc, h = item + async with self.dev.connection(loc) as (reader, writer): + ptr = i * Tensor.CHUNK_SIZE + size = min(len(dest[ptr:ptr+Tensor.CHUNK_SIZE]), Tensor.CHUNK_SIZE) + + writer.write(f"CHUNK_OUT {size}\r\n".encode()) + writer.write(h) + await writer.drain() + + chunk = await reader.readexactly(size) + + view = dest[ptr:ptr+len(chunk)] + view[:] = chunk + del view + + workers = [asyncio.create_task(_worker(item)) for item in src.copyout_queue] + await asyncio.gather(*workers) + src.copyout_queue.clear() + + def _offset(self, buf:TinyFSBuffer, size:int, offset:int): + return TinyFSBuffer(buf.device, size, offset, buf.request_id, buf.copyout_queue) diff --git a/tinygrad/runtime/support/am/ip.py b/tinygrad/runtime/support/am/ip.py index e6ff7a24e2..7dc47643d8 100644 --- a/tinygrad/runtime/support/am/ip.py +++ b/tinygrad/runtime/support/am/ip.py @@ -113,7 +113,7 @@ class AM_GMC(AM_IP): for eng_i in range(18): self.adev.wreg_pair(f"reg{ip}VM_INVALIDATE_ENG{eng_i}_ADDR_RANGE", "_LO32", "_HI32", 0x1fffffffff) self.hub_initted[ip] = True - @functools.cache + @functools.cache # pylint: disable=method-cache-max-size-none def get_pte_flags(self, pte_lv, is_table, frag, uncached, system, snooped, valid, extra=0): extra |= (am.AMDGPU_PTE_SYSTEM * system) | (am.AMDGPU_PTE_SNOOPED * snooped) | (am.AMDGPU_PTE_VALID * valid) | am.AMDGPU_PTE_FRAG(frag) if not is_table: extra |= (am.AMDGPU_PTE_WRITEABLE | am.AMDGPU_PTE_READABLE | am.AMDGPU_PTE_EXECUTABLE) @@ -175,7 +175,7 @@ class AM_SMU(AM_IP): def _send_msg(self, msg:int, param:int, read_back_arg=False, timeout=10000, debug=False): # default timeout is 10 seconds self._smu_cmn_send_msg(msg, param, debug=debug) - wait_cond(lambda: (self.adev.mmMP1_SMN_C2PMSG_90 if not debug else self.adev.mmMP1_SMN_C2PMSG_54).read(), value=1, timeout_ms=timeout, + wait_cond((self.adev.mmMP1_SMN_C2PMSG_90 if not debug else self.adev.mmMP1_SMN_C2PMSG_54).read, value=1, timeout_ms=timeout, msg=f"SMU msg {msg:#x} timeout") return (self.adev.mmMP1_SMN_C2PMSG_82 if not debug else self.adev.mmMP1_SMN_C2PMSG_53).read() if read_back_arg else None diff --git a/tinygrad/runtime/support/compiler_amd.py b/tinygrad/runtime/support/compiler_amd.py index 88608c71bc..8f26780d92 100644 --- a/tinygrad/runtime/support/compiler_amd.py +++ b/tinygrad/runtime/support/compiler_amd.py @@ -60,7 +60,11 @@ def compile_hip(prg:str, arch="gfx1100", asm=False) -> bytes: check(comgr.amd_comgr_set_data_name(data_src, b"")) check(comgr.amd_comgr_data_set_add(data_set_src, data_src)) # -include hiprtc_runtime.h was removed - check(set_options(action_info, f"-O3 -mcumode --hip-version=6.0.32830 -DHIP_VERSION_MAJOR=6 -DHIP_VERSION_MINOR=0 -DHIP_VERSION_PATCH=32830 -D__HIPCC_RTC__ -std=c++14 -nogpuinc -Wno-gnu-line-marker -Wno-missing-prototypes --offload-arch={arch} -I/opt/rocm/include -Xclang -disable-llvm-passes -Xclang -aux-triple -Xclang x86_64-unknown-linux-gnu".encode())) # noqa: E501 + options = [ + "-O3", "-mcumode", "--hip-version=6.0.32830", "-DHIP_VERSION_MAJOR=6", "-DHIP_VERSION_MINOR=0", "-DHIP_VERSION_PATCH=32830", + "-D__HIPCC_RTC__", "-std=c++14", "-nogpuinc", "-Wno-gnu-line-marker", "-Wno-missing-prototypes", f"--offload-arch={arch}", + "-I/opt/rocm/include", "-Xclang -disable-llvm-passes", "-Xclang -aux-triple", "-Xclang x86_64-unknown-linux-gnu"] + check(set_options(action_info, ' '.join(options).encode())) status = comgr.amd_comgr_do_action(comgr.AMD_COMGR_ACTION_COMPILE_SOURCE_WITH_DEVICE_LIBS_TO_BC, action_info, data_set_src, data_set_bc) if status != 0: print(_get_comgr_data(data_set_bc, comgr.AMD_COMGR_DATA_KIND_LOG).decode()) diff --git a/tinygrad/runtime/support/compiler_cuda.py b/tinygrad/runtime/support/compiler_cuda.py index 5c16aef2fc..3ba9945881 100644 --- a/tinygrad/runtime/support/compiler_cuda.py +++ b/tinygrad/runtime/support/compiler_cuda.py @@ -22,10 +22,12 @@ def jitlink_check(status, ctx=None): def pretty_ptx(s): # all expressions match `` and replace it with `color()` - s = re.sub(r'([!@<\[\s,\+\-;\n])((?:[_%$][\w%\$_]+(?:\.[xyz])?\:?)|(?:buf\d+))([<>\]\s,\+\-;\n\)])', lambda m:m[1]+colored(m[2], "blue")+m[3], s, flags=re.M) # identifiers # noqa: E501 + s = re.sub(r'([!@<\[\s,\+\-;\n])((?:[_%$][\w%\$_]+(?:\.[xyz])?\:?)|(?:buf\d+))([<>\]\s,\+\-;\n\)])', + lambda m:m[1]+colored(m[2], "blue")+m[3], s, flags=re.M) # identifiers s = re.sub(r'(.)((?:b|s|u|f)(?:8|16|32|64)|pred)([\.\s])', lambda m:m[1]+colored(m[2], "green")+m[3], s, flags=re.M) # types s = re.sub(r'^(\s*)([\w]+)(.*?;$)', lambda m:m[1]+colored(m[2], "yellow")+m[3], s, flags=re.M) # instructions - s = re.sub(r'([<>\[\]\s,\+\-;])((?:0[fF][0-9a-fA-F]{8})|(?:[0-9]+)|(?:0[xX][0-9a-fA-F]+))([<>\[\]\s,\+\-;])', lambda m:m[1]+colored(m[2], "yellow")+m[3], s, flags=re.M) # numbers # noqa: E501 + s = re.sub(r'([<>\[\]\s,\+\-;])((?:0[fF][0-9a-fA-F]{8})|(?:[0-9]+)|(?:0[xX][0-9a-fA-F]+))([<>\[\]\s,\+\-;])', + lambda m:m[1]+colored(m[2], "yellow")+m[3], s, flags=re.M) # numbers s = re.sub(r'(\.)(param|reg|global)', lambda m:m[1]+colored(m[2], "magenta"), s, flags=re.M) # space s = re.sub(r'(\.)(version|target|address_size|visible|entry)', lambda m:m[1]+colored(m[2], "magenta"), s, flags=re.M) # derivatives return s diff --git a/tinygrad/runtime/support/elf.py b/tinygrad/runtime/support/elf.py index 3276e6adb8..3e5f61bafd 100644 --- a/tinygrad/runtime/support/elf.py +++ b/tinygrad/runtime/support/elf.py @@ -33,7 +33,7 @@ def elf_loader(blob:bytes, force_section_align:int=1) -> tuple[memoryview, list[ for sh, trgt_sh_name, c_rels in rel + rela: target_image_off = next(tsh for tsh in sections if tsh.name == trgt_sh_name).header.sh_addr rels = [(r.r_offset, symtab[libc.ELF64_R_SYM(r.r_info)], libc.ELF64_R_TYPE(r.r_info), getattr(r, "r_addend", 0)) for r in c_rels] - for roff, sym, r_type_, r_addend in rels: + for _, sym, _, _ in rels: if sym.st_shndx == 0: raise RuntimeError(f'Attempting to relocate against an undefined symbol {repr(_strtab(sh_strtab, sym.st_name))}') relocs += [(target_image_off + roff, sections[sym.st_shndx].header.sh_addr + sym.st_value, rtype, raddend) for roff, sym, rtype, raddend in rels] diff --git a/tinygrad/runtime/support/memory.py b/tinygrad/runtime/support/memory.py index e5624515e5..1c22c1ecd9 100644 --- a/tinygrad/runtime/support/memory.py +++ b/tinygrad/runtime/support/memory.py @@ -30,10 +30,10 @@ class TLSFAllocator: self.blocks:dict[int, tuple[int, int|None, int|None, bool]] = {0: (size, None, None, True)} # size, next, prev, is_free self._insert_block(0, size) - @functools.cache + @functools.cache # pylint: disable=method-cache-max-size-none def lv1(self, size): return size.bit_length() - @functools.cache + @functools.cache # pylint: disable=method-cache-max-size-none def lv2(self, size): return (size - (1 << (size.bit_length() - 1))) // (1 << max(0, size.bit_length() - self.l2_cnt)) def _insert_block(self, start:int, size:int, prev:int|None=None): @@ -209,7 +209,7 @@ class MemoryManager: if getenv("MM_DEBUG", 0): print(f"mm {self.dev.devfmt}: unmapping {vaddr=:#x} ({size=:#x})") ctx = PageTableTraverseContext(self.dev, self.root_page_table, vaddr, free_pts=True) - for off, pt, pte_idx, pte_cnt, pte_covers in ctx.next(size): + for _, pt, pte_idx, pte_cnt, _ in ctx.next(size): for pte_id in range(pte_idx, pte_idx + pte_cnt): assert pt.valid(pte_id), f"PTE not mapped: {pt.entry(pte_id):#x}" pt.set_entry(pte_id, paddr=0x0, valid=False) diff --git a/tinygrad/runtime/support/nv/ip.py b/tinygrad/runtime/support/nv/ip.py index 2037960215..eda20117e6 100644 --- a/tinygrad/runtime/support/nv/ip.py +++ b/tinygrad/runtime/support/nv/ip.py @@ -124,6 +124,7 @@ class NV_FLCN(NV_IP): def __patch(cmd_id, cmd): patched_image = bytearray(image) + dmem_offset = 0 hdr = nv.FALCON_APPLICATION_INTERFACE_HEADER_V1.from_buffer_copy(image[(app_hdr_off:=self.desc_v3.IMEMLoadSize+self.desc_v3.InterfaceOffset):]) ents = (nv.FALCON_APPLICATION_INTERFACE_ENTRY_V1 * hdr.entryCount).from_buffer_copy(image[app_hdr_off + ctypes.sizeof(hdr):]) for i in range(hdr.entryCount): @@ -334,7 +335,7 @@ class NV_GSP(NV_IP): # Fill up arguments queue_args = nv.MESSAGE_QUEUE_INIT_ARGUMENTS(sharedMemPhysAddr=queues_sysmem[0], pageTableEntryCount=pte_cnt, cmdQueueOffset=pt_size, statQueueOffset=pt_size + queue_size) - rm_args, self.rm_args_sysmem = self.nvdev._alloc_boot_struct(nv.GSP_ARGUMENTS_CACHED(bDmemStack=True, messageQueueInitArguments=queue_args)) + _, self.rm_args_sysmem = self.nvdev._alloc_boot_struct(nv.GSP_ARGUMENTS_CACHED(bDmemStack=True, messageQueueInitArguments=queue_args)) # Build command queue header self.cmd_q_va, self.stat_q_va = queues_va + pt_size, queues_va + pt_size + queue_size @@ -481,7 +482,7 @@ class NV_GSP(NV_IP): params.ramfcMem = nv_gpu.NV_MEMORY_DESC_PARAMS(base=ramfc_alloc.paddrs[0][0], size=0x200, addressSpace=2, cacheAttrib=0) params.instanceMem = nv_gpu.NV_MEMORY_DESC_PARAMS(base=ramfc_alloc.paddrs[0][0], size=0x1000, addressSpace=2, cacheAttrib=0) - method_va, method_sysmem = System.alloc_sysmem(0x5000, contiguous=True) + _, method_sysmem = System.alloc_sysmem(0x5000, contiguous=True) params.mthdbufMem = nv_gpu.NV_MEMORY_DESC_PARAMS(base=method_sysmem[0], size=0x5000, addressSpace=1, cacheAttrib=0) if client is not None and client != self.priv_root and params.hObjectError != 0: @@ -557,7 +558,7 @@ class NV_GSP(NV_IP): self.nvdev.wreg(addr, (self.nvdev.rreg(addr) & ~mask) | (val & mask)) elif op == 0x2: # reg poll addr, mask, val, _, _ = next(cmd_iter), next(cmd_iter), next(cmd_iter), next(cmd_iter), next(cmd_iter) - wait_cond(lambda: (self.nvdev.rreg(addr) & mask), value=val, msg=f"Register {addr:#x} not equal to {val:#x} after polling") + wait_cond(lambda a, m: (self.nvdev.rreg(a) & m), addr, mask, value=val, msg=f"Register {addr:#x} not equal to {val:#x} after polling") elif op == 0x3: time.sleep(next(cmd_iter) / 1e6) # delay us elif op == 0x4: # save reg addr, index = next(cmd_iter), next(cmd_iter) diff --git a/tinygrad/runtime/support/nv/nvdev.py b/tinygrad/runtime/support/nv/nvdev.py index 6831b5e8b1..496d8ec5c8 100644 --- a/tinygrad/runtime/support/nv/nvdev.py +++ b/tinygrad/runtime/support/nv/nvdev.py @@ -152,6 +152,8 @@ class NVDev(PCIDevImplBase): return gzip.decompress(struct.pack("<4BL2B", 0x1f, 0x8b, 8, 0, 0, 0, 3) + image) if "COMPRESSION: YES" in info else image def include(self, file:str): + def _do_eval(s:str): return eval(s) # pylint: disable=eval-used + regs_off = {'NV_PFALCON_FALCON': 0x0, 'NV_PGSP_FALCON': 0x0, 'NV_PSEC_FALCON': 0x0, 'NV_PRISCV_RISCV': 0x1000, 'NV_PGC6_AON': 0x0, 'NV_PFSP': 0x0, 'NV_PGC6_BSI': 0x0, 'NV_PFALCON_FBIF': 0x600, 'NV_PFALCON2_FALCON': 0x1000, 'NV_PBUS': 0x0, 'NV_PFB': 0x0, 'NV_PMC': 0x0, 'NV_PGSP_QUEUE': 0x0, 'NV_VIRTUAL_FUNCTION':0xb80000} @@ -163,13 +165,13 @@ class NVDev(PCIDevImplBase): name, hi, lo = m.groups() reg = next((r for r in self.reg_names if name.startswith(r+"_")), None) - if reg is not None: self.__dict__[reg].add_field(name[len(reg)+1:].lower(), eval(lo), eval(hi)) - else: self.reg_offsets[name] = (eval(lo), eval(hi)) + if reg is not None: self.__dict__[reg].add_field(name[len(reg)+1:].lower(), _do_eval(lo), _do_eval(hi)) + else: self.reg_offsets[name] = (_do_eval(lo), _do_eval(hi)) continue if m:=re.match(r'#define\s+(\w+)\s*\(\s*(\w+)\s*\)\s*(.+)', raw): # reg set fn = m.groups()[2].strip().rstrip('\\').split('/*')[0].rstrip() - name, value = m.groups()[0], eval(f"lambda {m.groups()[1]}: {fn}") + name, value = m.groups()[0], _do_eval(f"lambda {m.groups()[1]}: {fn}") elif m:=re.match(r'#define\s+(\w+)\s+([0-9A-Fa-fx]+)(?![^\n]*:)', raw): name, value = m.groups()[0], int(m.groups()[1], 0) # reg value else: continue diff --git a/tinygrad/runtime/support/system.py b/tinygrad/runtime/support/system.py index 66b2f78615..df575b89fe 100644 --- a/tinygrad/runtime/support/system.py +++ b/tinygrad/runtime/support/system.py @@ -10,14 +10,14 @@ MAP_FIXED, MAP_LOCKED, MAP_POPULATE, MAP_NORESERVE = 0x10, 0 if OSX else 0x2000, class _System: def reserve_hugepages(self, cnt): os.system(f"sudo sh -c 'echo {cnt} > /proc/sys/vm/nr_hugepages'") - def memory_barrier(self): lib.atomic_thread_fence(__ATOMIC_SEQ_CST:=5) if (lib:=self.atomic_lib()) is not None else None + def memory_barrier(self): lib.atomic_thread_fence(__ATOMIC_SEQ_CST:=5) if (lib:=self.atomic_lib) is not None else None def lock_memory(self, addr:int, size:int): if libc.mlock(ctypes.c_void_p(addr), size): raise RuntimeError(f"Failed to lock memory at {addr:#x} with size {size:#x}") def system_paddrs(self, vaddr:int, size:int) -> list[int]: - self.pagemap().seek(vaddr // mmap.PAGESIZE * 8) - return [(x & ((1<<55) - 1)) * mmap.PAGESIZE for x in array.array('Q', self.pagemap().read(size//mmap.PAGESIZE*8, binary=True))] + self.pagemap.seek(vaddr // mmap.PAGESIZE * 8) + return [(x & ((1<<55) - 1)) * mmap.PAGESIZE for x in array.array('Q', self.pagemap.read(size//mmap.PAGESIZE*8, binary=True))] def alloc_sysmem(self, size:int, vaddr:int=0, contiguous:bool=False, data:bytes|None=None) -> tuple[int, list[int]]: assert not contiguous or size <= (2 << 20), "Contiguous allocation is only supported for sizes up to 2MB" @@ -36,17 +36,17 @@ class _System: if vendor == target_vendor and device in target_devices: result.append(pcibus) return sorted(result) - @functools.cache + @functools.cached_property def atomic_lib(self): return ctypes.CDLL(ctypes.util.find_library('atomic')) if sys.platform == "linux" else None - @functools.cache + @functools.cached_property def pagemap(self) -> FileIOInterface: if FileIOInterface(reloc_sysfs:="/proc/sys/vm/compact_unevictable_allowed", os.O_RDONLY).read()[0] != "0": os.system(cmd:=f"sudo sh -c 'echo 0 > {reloc_sysfs}'") assert FileIOInterface(reloc_sysfs, os.O_RDONLY).read()[0] == "0", f"Failed to disable migration of locked pages. Please run {cmd} manually." return FileIOInterface("/proc/self/pagemap", os.O_RDONLY) - @functools.cache + @functools.cached_property def vfio(self) -> FileIOInterface|None: try: if not FileIOInterface.exists("/sys/module/vfio"): os.system("sudo modprobe vfio-pci disable_idle_d3=1") @@ -90,7 +90,7 @@ class PCIDevice: " to allow python accessing device or run with sudo") from e raise RuntimeError(f"Cannot resize BAR {i}: {e}. Ensure the resizable BAR option is enabled on your system.") from e - if getenv("VFIO", 0) and (vfio_fd:=System.vfio()) is not None: + if getenv("VFIO", 0) and (vfio_fd:=System.vfio) is not None: FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/driver_override", os.O_WRONLY).write("vfio-pci") FileIOInterface("/sys/bus/pci/drivers_probe", os.O_WRONLY).write(self.pcibus) iommu_group = FileIOInterface.readlink(f"/sys/bus/pci/devices/{self.pcibus}/iommu_group").split('/')[-1] diff --git a/tinygrad/runtime/support/usb.py b/tinygrad/runtime/support/usb.py index 285e3cf287..2340c944cb 100644 --- a/tinygrad/runtime/support/usb.py +++ b/tinygrad/runtime/support/usb.py @@ -229,7 +229,7 @@ class ASM24Controller: for i in range(0, len(ops), bs:=(4 if OSX else 16)): self.exec_ops(list(itertools.chain.from_iterable(ops[i:i+bs]))) class USBMMIOInterface(MMIOInterface): - def __init__(self, usb, addr, size, fmt, pcimem=True): + def __init__(self, usb, addr, size, fmt, pcimem=True): # pylint: disable=super-init-not-called self.usb, self.addr, self.nbytes, self.fmt, self.pcimem, self.el_sz = usb, addr, size, fmt, pcimem, struct.calcsize(fmt) def __getitem__(self, index): return self._access_items(index) @@ -256,13 +256,14 @@ class USBMMIOInterface(MMIOInterface): acc, acc_size = self._acc_size(sz) return bytes(array.array(acc, [self._acc_one(off + i * acc_size, acc_size) for i in range(sz // acc_size)])) - else: # write op - data = struct.pack(self.fmt, data) if isinstance(data, int) else bytes(data) - if not self.pcimem: - # Fast path for writing into buffer 0xf000 - use_cache = 0xa800 <= self.addr <= 0xb000 - return self.usb.scsi_write(bytes(data)) if self.addr == 0xf000 else self.usb.write(self.addr + off, bytes(data), ignore_cache=not use_cache) + # write op + data = struct.pack(self.fmt, data) if isinstance(data, int) else bytes(data) - _, acc_sz = self._acc_size(len(data) * struct.calcsize(self.fmt)) - self.usb.pcie_mem_write(self.addr+off, [int.from_bytes(data[i:i+acc_sz], "little") for i in range(0, len(data), acc_sz)], acc_sz) + if not self.pcimem: + # Fast path for writing into buffer 0xf000 + use_cache = 0xa800 <= self.addr <= 0xb000 + return self.usb.scsi_write(bytes(data)) if self.addr == 0xf000 else self.usb.write(self.addr + off, bytes(data), ignore_cache=not use_cache) + + _, acc_sz = self._acc_size(len(data) * struct.calcsize(self.fmt)) + self.usb.pcie_mem_write(self.addr+off, [int.from_bytes(data[i:i+acc_sz], "little") for i in range(0, len(data), acc_sz)], acc_sz) diff --git a/tinygrad/schedule/indexing.py b/tinygrad/schedule/indexing.py index 0b4eb1ba35..6af4154a5e 100644 --- a/tinygrad/schedule/indexing.py +++ b/tinygrad/schedule/indexing.py @@ -1,9 +1,9 @@ -from typing import Iterator, Sequence +from typing import Iterator import functools, operator, itertools from dataclasses import dataclass, field from tinygrad.dtype import dtypes, AddrSpace from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType -from tinygrad.uop.symbolic import sym, symbolic +from tinygrad.uop.symbolic import symbolic, pm_simplify_valid, pm_drop_and_clauses from tinygrad.helpers import argsort, all_same, cpu_profile, TracingKey ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.ASSIGN, Ops.COPY, Ops.BUFFER, Ops.BUFFER_VIEW, @@ -41,7 +41,7 @@ class BufferizeOpts: @dataclass class IndexingContext: realize_map: dict[UOp, None] = field(default_factory=dict) - range_map: dict[UOp, tuple[list[UOp], list[UOp]]] = field(default_factory=dict) + range_map: dict[UOp, tuple[tuple[UOp, ...], tuple[UOp, ...]]] = field(default_factory=dict) # create ranges range_idx: Iterator[int] = field(default_factory=itertools.count) @@ -103,30 +103,34 @@ pm_apply_rangeify = PatternMatcher([ ]) # this is the definition of the movement ops -def apply_movement_op(x:UOp, rngs:Sequence[UOp]) -> list[UOp]: - match x.op: - case Ops.SHRINK: rngs = [a if ss == 0 else a+ss for a,(ss,_) in zip(rngs, x.arg)] - case Ops.PERMUTE: rngs = [rngs[p] for p in argsort(x.arg)] - case Ops.FLIP: rngs = [((s-1)-a) if f else a for a,s,f in zip(rngs, x.shape, x.arg)] - case Ops.EXPAND: rngs = [a if in_sh == out_sh else a.const_like(0) for a,in_sh,out_sh in zip(rngs, x.src[0].shape, x.shape)] +@functools.cache +def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UOp, ...]) -> tuple[UOp, ...]: + match op: + case Ops.SHRINK: rngs = tuple(a if ss == 0 else a+ss for a,(ss,_) in zip(rngs, arg)) + case Ops.PERMUTE: rngs = tuple(rngs[p] for p in argsort(arg)) + case Ops.FLIP: rngs = tuple(((s-1)-a) if f else a for a,s,f in zip(rngs, in_shape, arg)) + case Ops.EXPAND: rngs = tuple(a if in_sh == out_sh else a.const_like(0) for a,in_sh,out_sh in zip(rngs, in_shape, arg)) case Ops.PAD: # TODO: why is multiple graph_rewrites faster than one here? - rngs = [r if (s == 0 and e == 0) else graph_rewrite(((r >= s) & (r < (sh-e))).where(r-s, UOp.invalid()), sym, name="pad") - for r,sh,(s,e) in zip(rngs, x.shape, x.arg)] + # TODO: the .where(r-s, i) is not inside the graph_rewrite so that `convert_pad_to_where_to_keep_behavior_local` + # wraps the pad with only the newly added valid + rngs = tuple(r if (s == 0 and e == 0) else graph_rewrite(((r >= s) & (r < (sh+s))), + symbolic+pm_simplify_valid, name="pad").where(r-s, UOp.invalid()) for r,sh,(s,e) in zip(rngs, in_shape, arg)) case Ops.RESHAPE: acc = 1 axes_in:list[UOp] = [] - for s,src in list(zip(x.shape, rngs))[::-1]: + for s,src in list(zip(arg, rngs))[::-1]: axes_in.append(acc*src) acc *= s combined_axes = sum(axes_in, start=UOp.const(dtypes.index, 0)) axes_out:list[UOp] = [] - for s in x.src[0].shape[::-1]: + for s in in_shape[::-1]: axes_out.append(combined_axes % s) combined_axes //= s # this simplify is doing a lot of heavy lifting. this is the replacement for the reshape view merging code - rngs = list(graph_rewrite(UOp.sink(*axes_out[::-1]), symbolic, name="reshape").src) - case _: raise RuntimeError(f"{x.op} is not a MovementOp") + rngs = graph_rewrite(graph_rewrite(UOp.sink(*axes_out[::-1]), symbolic+pm_simplify_valid, name="reshape"), + pm_drop_and_clauses, name="reshape drop ands").src + case _: raise RuntimeError(f"{op} is not a MovementOp") return rngs @cpu_profile(TracingKey("run_rangeify"), "TINY") @@ -157,7 +161,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]: consumer_rngs = [rctx.range_map[c][0] for c in consumer_map[x] if c in rctx.range_map] if x in rctx.realize_map: # if this is in the realize_map, we create new ranges (at the output) - out_rngs = [rctx.new_range(s) if not isinstance(s, UOp) or s.op is not Ops.RANGE else s for s in x.shape] + out_rngs = tuple(rctx.new_range(s) if not isinstance(s, UOp) or s.op is not Ops.RANGE else s for s in x.shape) # all ranges are ended now ending_ranges[x] = False elif x.op in {Ops.MSTACK, Ops.MSELECT}: @@ -181,15 +185,16 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]: # TODO: in RANGEIFY > 1 all_all_same isn't required all_all_same = all(same_rngs for _,_,same_rngs in rngs_valids) - out_rngs = [] + _out_rngs = [] for i,(local_rngs,valids,same_rngs) in enumerate(rngs_valids): # we compare the ranges without their valids if all_all_same: # the new valid is the OR of all the children valids minimum_valid = functools.reduce(operator.or_, valids, UOp.const(dtypes.bool, False)) - out_rngs.append(graph_rewrite(minimum_valid.where(local_rngs[0], UOp.invalid()), symbolic, name="minimum_valid")) + _out_rngs.append(graph_rewrite(minimum_valid.where(local_rngs[0], UOp.invalid()), symbolic, name="minimum_valid")) else: - out_rngs.append(rctx.new_range(x.shape[i])) + _out_rngs.append(rctx.new_range(x.shape[i])) + out_rngs = tuple(_out_rngs) # we have to realize here if there's new ranges if not all_all_same: rctx.realize_map[x] = None @@ -203,18 +208,16 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]: # 2. newly created for REDUCE_AXIS # 3. passed through for everything else - rngs = out_rngs # rngs is the input ranges + rngs = out_rngs # rngs is the input ranges # pylint: disable=possibly-used-before-assignment # apply movement ops - if x.op in GroupOp.Movement: rngs = apply_movement_op(x, rngs) + if x.op in GroupOp.Movement: rngs = apply_movement_op(x.op, x.src[0].shape, x.arg, rngs) # if the EXPAND is used to inject a range, we don't mark it as ending_ranges. otherwise we do. if x.op is Ops.EXPAND and all(isinstance(y, int) or y.op is not Ops.RANGE for y in x.shape): ending_ranges[x] = True # REDUCE_AXIS creates ranges for the axes it is reducing if x.op is Ops.REDUCE_AXIS: - rngs = rngs[:] - for i,s in enumerate(x.src[0].shape): - if i in x.arg[1]: rngs[i] = rctx.new_range(s, axistype=AxisType.REDUCE) + rngs = tuple(rctx.new_range(s, axistype=AxisType.REDUCE) if i in x.arg[1] else r for i,(r,s) in enumerate(zip(rngs, x.src[0].shape))) if debug: print("***" if x in rctx.realize_map else " ", len(consumer_map[x]), f"{str(x.op):20s}", diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index fe968c5f74..97f0dd0076 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -103,7 +103,7 @@ earliest_rewrites = PatternMatcher([ # movement op on INDEX as a PatternMatcher pm_mops = PatternMatcher([ (UPat(GroupOp.Movement, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"), - lambda r,idx: r.src[0].index(*apply_movement_op(r, idx.src[1:]), dtype=idx.dtype, arg=idx.arg)), + lambda r,idx: r.src[0].index(*apply_movement_op(r.op, r.src[0].shape, r.arg, idx.src[1:]), dtype=idx.dtype, arg=idx.arg)), # type: ignore ]) # ***************** @@ -207,7 +207,7 @@ pm_cleanups = pm_mops+PatternMatcher([ ]) def late_buffer_view(t:UOp, b:UOp): - if isinstance(b.device, str) and b.device.startswith("DISK"): + if isinstance(b.device, str) and (b.device.startswith("DISK") or b.device.startswith("TINYFS")): rngs = b.src[1:] size = prod(shape := [int(r.vmax+1) for r in rngs]) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 0c3a11581b..ca05714038 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -9,8 +9,8 @@ from tinygrad.helpers import argfix, make_tuple, flatten, prod, all_int, round_u from tinygrad.helpers import IMAGE, WINO, Metadata, TRACEMETA, ceildiv, fetch, polyN, unwrap, DEBUG, is_numpy_ndarray, FUSE_ATTENTION from tinygrad.helpers import suppress_finalizing from tinygrad.gradient import compute_gradient -from tinygrad.uop.ops import smax, smin, resolve, UOp, Ops, sint, MathTrait, identity_element, all_metadata, _index_to_concrete_int, sint_to_uop, \ - srender +from tinygrad.uop.mathtraits import MathTrait +from tinygrad.uop.ops import smax, smin, resolve, UOp, Ops, sint, identity_element, all_metadata, _index_to_concrete_int, sint_to_uop, srender from tinygrad.uop.spec import tensor_uop_spec, type_verify from tinygrad.device import Device, Buffer from tinygrad.engine.realize import run_schedule @@ -411,6 +411,59 @@ class Tensor(MathTrait): """ return self.replace(self.shard(devices, axis)) + CHUNK_SIZE = 2**20 + def load(self, size:int) -> Tensor: + """ + Load a tensor from storage. + + self should be a tensor of the hash to load + """ + # TODO: this should work locally as well + assert self.dtype == dtypes.uint8, "hash is expected to be uint8" + h = self.contiguous().flatten() + assert h.shape[0] == 16, "expected hash" + + base_chunks = math.ceil(size / Tensor.CHUNK_SIZE) + tree_depth = math.ceil(math.log(base_chunks, Tensor.CHUNK_SIZE // 16)) + data, level_chunks = h, 0 + for i in reversed(range(tree_depth + 1)): + data = data.to("tinyfs:load") + + # if not last level, its still hashes + if i > 0 or tree_depth == 0: + level_chunks = max(1, math.ceil(base_chunks / (Tensor.CHUNK_SIZE // 16)**(i-1))) + pad_amt = 16 * level_chunks + else: pad_amt = Tensor.CHUNK_SIZE * level_chunks + if (tsize := data.shape[0]) < pad_amt: data = data.pad((0, pad_amt - tsize)) + data = data[:pad_amt].contiguous() + if i != 0: data = data.to(self.device) + + return data[:size] + + def store(self) -> Tensor: + """ + Store a tensor to storage. + """ + # TODO: this should work locally as well + data = self.contiguous().flatten().bitcast(dtypes.uint8) + + # pad to a multiple of 1mb + if (tsize := data.shape[0]) % Tensor.CHUNK_SIZE != 0: data = data.pad((0, Tensor.CHUNK_SIZE - tsize % Tensor.CHUNK_SIZE)) + size = data.shape[0] + + base_chunks = math.ceil(size / Tensor.CHUNK_SIZE) + tree_depth = math.ceil(math.log(base_chunks, Tensor.CHUNK_SIZE // 16)) + + to_device = "CPU" if isinstance(self.device, str) and self.device.startswith("DISK") else self.device + + level_chunks = base_chunks + for _ in range(tree_depth + 1): + data = data.to("tinyfs:store")[:level_chunks * 16].contiguous().to(to_device) + if (tsize := data.shape[0]) % Tensor.CHUNK_SIZE != 0: data = data.pad((0, Tensor.CHUNK_SIZE - tsize % Tensor.CHUNK_SIZE)) + level_chunks = math.ceil(data.shape[0] / Tensor.CHUNK_SIZE) + + return data[:16].contiguous() + @staticmethod def from_uop(y:UOp, **kwargs) -> Tensor: if y.op is Ops.BIND: return Tensor(y, **kwargs, requires_grad=False) @@ -1159,6 +1212,7 @@ class Tensor(MathTrait): match index: case Tensor(): if not dtypes.is_int(index.dtype): raise IndexError(f"index dtype {index.dtype} is not supported") + assert isinstance(size, int), "size must be an int" index = (index < 0).where(index+size, index).to(self.device) # treat negative index values case list() | tuple(): if not dtypes.is_int((ti:=Tensor(index)).dtype): raise IndexError(f"{index=} contains non-int element") @@ -2484,17 +2538,20 @@ class Tensor(MathTrait): if IMAGE: return self.image_conv2d(weight, bias, groups, stride, dilation, padding, dtype) (bs,cin_), (cout,cin), HW = self.shape[:2], weight.shape[:2], weight.shape[2:] padding_ = self._resolve_pool_pads(padding, len(HW)) - assert groups*cin == cin_ and len(self.shape) == len(weight.shape), f"Input Tensor shape {self.shape} does not match the shape of the weights {weight.shape}. ({groups*cin} vs. {cin_})" # noqa: E501 + assert groups*cin == cin_ and len(self.shape) == len(weight.shape),\ + f"Input Tensor shape {self.shape} does not match the shape of the weights {weight.shape}. ({groups*cin} vs. {cin_})" # conv2d is a pooling op (with padding) x = self.pad(padding_)._pool(HW, stride, dilation) # (bs, groups*cin, oy, ox, H, W) rcout, oyx = cout//groups, x.shape[2:-len(HW)] if not all(x == 3 for x in HW) or stride != 1 or dilation != 1 or not WINO: # normal conv - x = x.reshape(bs, groups, cin, 1, *oyx, *HW).expand(bs, groups, cin, rcout, *oyx, *HW).permute(0,1,3,*[4+i for i in range(len(oyx))],2,*[4+len(oyx)+i for i in range(len(HW))]) # noqa: E501 + x = x.reshape(bs, groups, cin, 1, *oyx, *HW).expand(bs, groups, cin, rcout, *oyx, *HW)\ + .permute(0,1,3,*[4+i for i in range(len(oyx))],2,*[4+len(oyx)+i for i in range(len(HW))]) # conv! broadcasted to (bs, groups, rcout, *oyx, cin, *HW) - ret = (x * weight.reshape(1, groups, rcout, *[1] * len(oyx), cin, *HW)).sum([-1-i for i in range(1+len(oyx))], keepdim=True, dtype=dtype).reshape(bs, cout, *oyx) # noqa: E501 + ret = (x * weight.reshape(1, groups, rcout, *[1] * len(oyx), cin, *HW))\ + .sum([-1-i for i in range(1+len(oyx))], keepdim=True, dtype=dtype).reshape(bs, cout, *oyx) return ret if bias is None else ret.add(bias.reshape(1, -1, *[1] * len(HW))) HWI, HWO = (6,) * len(HW), (4,) * len(HW) # F(4x4,3x3) winograd tiles @@ -2505,7 +2562,8 @@ class Tensor(MathTrait): # TODO: stride == dilation # use padding to round up to 4x4 output tiles # (bs, cin_, tyx, HWI) - d = self.pad(sum([[padding_[i*2], padding_[i*2+1] + (-(dim + sum(padding_[i * 2:(i + 1) * 2]) - 2) % 4)] for i, dim in enumerate(self.shape[-len(HW):])], []))._pool(HWI, HWO) # noqa: E501 + pads = [[padding_[i*2], padding_[i*2+1] + (-(dim + sum(padding_[i * 2:(i + 1) * 2]) - 2) % 4)] for i, dim in enumerate(self.shape[-len(HW):])] + d = self.pad(sum(pads, []))._pool(HWI, HWO) # move HW to the front: # (HWI, bs, cin_, tyx) d = d.permute(*range(len(d.shape)-len(HW),len(d.shape)), *range(len(d.shape)-len(HW))) tyx = d.shape[-len(HWI):] # dim of tiling @@ -2627,7 +2685,8 @@ class Tensor(MathTrait): base = ret[..., -1]._cumalu(-1, op, _include_initial=True) base = base.unsqueeze(-1).expand(*base.shape, ret.shape[-1]) def fix(x: Tensor) -> Tensor: return x.flatten(start_dim=-2)[..., -s:].transpose(axis,-1) - return {Ops.ADD: Tensor.__add__, Ops.MAX: Tensor.maximum, Ops.MUL: Tensor.__mul__}[op](fix(ret), fix(base)) + reduce_fxns: dict[Ops, Callable[[Tensor, Tensor], Tensor]] = {Ops.ADD: Tensor.__add__, Ops.MAX: Tensor.maximum, Ops.MUL: Tensor.__mul__} + return reduce_fxns[op](fix(ret), fix(base)) def cumsum(self, axis:int=0) -> Tensor: """ @@ -3666,7 +3725,7 @@ class Tensor(MathTrait): if self.dtype != dtypes.bool and not dtypes.is_int(self.dtype): raise RuntimeError(f"{self.dtype} is not supported") return self.logical_not() if self.dtype == dtypes.bool else self ^ -1 - def lshift(self, x:int, reverse=False) -> Tensor: + def lshift(self, x:Tensor|int, reverse=False) -> Tensor: """ Computes left arithmetic shift of `self` by `x` bits. `self` must have unsigned dtype. Equivalent to `self << x`. @@ -3678,7 +3737,7 @@ class Tensor(MathTrait): assert dtypes.is_unsigned(self.dtype) and isinstance(x, int) and x >= 0 and not reverse, f"not supported {self.dtype=} {x=}" return self.mul(2 ** x, reverse) - def rshift(self, x:int, reverse=False) -> Tensor: + def rshift(self, x:Tensor|int, reverse=False) -> Tensor: """ Computes right arithmetic shift of `self` by `x` bits. `self` must have unsigned dtype. Equivalent to `self >> x`. @@ -3794,18 +3853,20 @@ class Tensor(MathTrait): def __rpow__(self, x) -> Tensor: return self.pow(x, True) def __rmatmul__(self, x) -> Tensor: return self.matmul(x, True) - def __iadd__(self, x) -> Tensor: return self.assign(self.add(x)) - def __isub__(self, x) -> Tensor: return self.assign(self.sub(x)) - def __imul__(self, x) -> Tensor: return self.assign(self.mul(x)) - def __ipow__(self, x) -> Tensor: return self.assign(self.pow(x)) - def __itruediv__(self, x) -> Tensor: return self.assign(self.div(x)) def __ifloordiv__(self, x) -> Tensor: return self.assign(self.__floordiv__(x)) + def __ipow__(self, x) -> Tensor: return self.assign(self.pow(x)) def __imatmul__(self, x) -> Tensor: return self.assign(self.matmul(x)) - def __iand__(self, x) -> Tensor: return self.assign(self.bitwise_and(x)) - def __ior__(self, x) -> Tensor: return self.assign(self.bitwise_or(x)) - def __ixor__(self, x) -> Tensor: return self.assign(self.bitwise_xor(x)) - def __ilshift__(self, x) -> Tensor: return self.assign(self.lshift(x)) - def __irshift__(self, x) -> Tensor: return self.assign(self.rshift(x)) + + # unlike Tensors, UOps are immutable, so these don't go in MathTraits + def __iadd__(self, x) -> Tensor: return self.assign(self.add(x)) # type: ignore[misc] + def __isub__(self, x) -> Tensor: return self.assign(self.sub(x)) # type: ignore[misc] + def __imul__(self, x) -> Tensor: return self.assign(self.mul(x)) # type: ignore[misc] + def __itruediv__(self, x) -> Tensor: return self.assign(self.div(x)) # type: ignore[misc] + def __iand__(self, x) -> Tensor: return self.assign(self.bitwise_and(x)) # type: ignore[misc] + def __ior__(self, x) -> Tensor: return self.assign(self.bitwise_or(x)) # type: ignore[misc] + def __ixor__(self, x) -> Tensor: return self.assign(self.bitwise_xor(x)) # type: ignore[misc] + def __ilshift__(self, x) -> Tensor: return self.assign(self.lshift(x)) # type: ignore[misc] + def __irshift__(self, x) -> Tensor: return self.assign(self.rshift(x)) # type: ignore[misc] def __lt__(self, x) -> Tensor: return self._apply_broadcasted_uop(UOp.__lt__, x, False) def __gt__(self, x) -> Tensor: return self._apply_broadcasted_uop(UOp.__lt__, x, True) diff --git a/tinygrad/uop/mathtraits.py b/tinygrad/uop/mathtraits.py index 0de976c90b..a1f5d7eca2 100644 --- a/tinygrad/uop/mathtraits.py +++ b/tinygrad/uop/mathtraits.py @@ -1,15 +1,17 @@ +from typing import TypeVar from tinygrad.uop import Ops -from tinygrad.helpers import T -from tinygrad.dtype import dtypes +from tinygrad.dtype import dtypes, ConstType +TMT = TypeVar("TMT", bound="MathTrait") class MathTrait: # required to implement - def alu(self:T, op:Ops, *src) -> T: raise NotImplementedError - def const_like(self:T, b) -> T: raise NotImplementedError + def alu(self:TMT, op:Ops, *src:TMT) -> TMT: raise NotImplementedError + def const_like(self:TMT, b:ConstType) -> TMT: raise NotImplementedError # great functions you get! - def ufix(self, x): return self.const_like(x) if not isinstance(x, MathTrait) else x - def _binop(self, op, x, reverse): return self.ufix(x).alu(op, self) if reverse else self.alu(op, self.ufix(x)) + def ufix(self:TMT, x:TMT|ConstType) -> TMT: return self.const_like(x) if not isinstance(x, MathTrait) else x + def _binop(self:TMT, op:Ops, x:TMT|ConstType, reverse:bool) -> TMT: + return self.ufix(x).alu(op, self) if reverse else self.alu(op, self.ufix(x)) def logical_not(self): return self.ne(True) def neg(self): if (dtype:=getattr(self, 'dtype')) is None: raise TypeError(f"MathTraits __neg__ requires a dtype, {self=}") @@ -18,7 +20,7 @@ class MathTrait: if (dtype:=getattr(self, 'dtype')) is not None: if isinstance(dtype, tuple): dtype = dtype[0] if not (dtypes.is_bool(dtype) or dtypes.is_int(dtype)): raise RuntimeError(f"{dtype} is not supported") - def add(self, x, reverse=False): + def add(self:TMT, x:TMT|ConstType, reverse:bool=False): """ Adds `self` and `x`. Equivalent to `self + x`. @@ -36,7 +38,7 @@ class MathTrait: ``` """ return self._binop(Ops.ADD, x, reverse) - def mul(self, x, reverse=False): + def mul(self:TMT, x:TMT|ConstType, reverse:bool=False): """ Multiplies `self` and `x`. Equivalent to `self * x`. @@ -55,7 +57,7 @@ class MathTrait: ``` """ return self._binop(Ops.MUL, x, reverse) - def bitwise_and(self, x, reverse=False): + def bitwise_and(self:TMT, x:TMT|ConstType, reverse:bool=False): """ Computes the bitwise AND of `self` and `x`. Equivalent to `self & x`. @@ -69,7 +71,7 @@ class MathTrait: """ self._check_dtype() return self._binop(Ops.AND, x, reverse) - def bitwise_or(self, x, reverse=False): + def bitwise_or(self:TMT, x:TMT|ConstType, reverse:bool=False): """ Computes the bitwise OR of `self` and `x`. Equivalent to `self | x`. @@ -83,7 +85,7 @@ class MathTrait: """ self._check_dtype() return self._binop(Ops.OR, x, reverse) - def bitwise_xor(self, x, reverse=False): + def bitwise_xor(self:TMT, x:TMT|ConstType, reverse:bool=False): """ Computes bitwise xor of `self` and `x`. Equivalent to `self ^ x`. @@ -98,7 +100,7 @@ class MathTrait: """ self._check_dtype() return self._binop(Ops.XOR, x, reverse) - def idiv(self, x, reverse=False): + def idiv(self:TMT, x:TMT|ConstType, reverse:bool=False): """ Divides `self` by `x`. Equivalent to `self // x`. @@ -110,61 +112,61 @@ class MathTrait: ``` """ return self._binop(Ops.IDIV, x, reverse) - def mod(self, x, reverse=False): return self._binop(Ops.MOD, x, reverse) - def sub(self, x, reverse=False): return self.ufix(x).alu(Ops.ADD, -self) if reverse else self.alu(Ops.ADD, self.ufix(-x)) - def div(self, x, reverse=False): return (self.ufix(x)*self.alu(Ops.RECIP)) if reverse else (self*self.ufix(x).alu(Ops.RECIP)) + def mod(self:TMT, x:TMT|ConstType, reverse:bool=False): return self._binop(Ops.MOD, x, reverse) + def sub(self:TMT, x:TMT|ConstType, reverse:bool=False): return self.ufix(x).alu(Ops.ADD, -self) if reverse else self.alu(Ops.ADD, self.ufix(-x)) + def div(self:TMT, x:TMT|ConstType, reverse:bool=False): return (self.ufix(x)*self.alu(Ops.RECIP)) if reverse else (self*self.ufix(x).alu(Ops.RECIP)) def __neg__(self): return self.neg() - def __add__(self, x): return self.add(x) - def __sub__(self, x): return self.sub(x) - def __mul__(self, x): return self.mul(x) - def __truediv__(self, x): return self.div(x) - def __floordiv__(self, x): return self.idiv(x) # TODO: idiv is trunc div, not floordiv - def __mod__(self, x): return self.mod(x) - def __and__(self, x): return self.bitwise_and(x) - def __or__(self, x): return self.bitwise_or(x) - def __xor__(self, x): return self.bitwise_xor(x) + def __add__(self:TMT, x:TMT|ConstType): return self.add(x) + def __sub__(self:TMT, x:TMT|ConstType): return self.sub(x) + def __mul__(self:TMT, x:TMT|ConstType): return self.mul(x) + def __truediv__(self:TMT, x:TMT|ConstType): return self.div(x) + def __floordiv__(self:TMT, x:TMT|ConstType): return self.idiv(x) # TODO: idiv is trunc div, not floordiv + def __mod__(self:TMT, x:TMT|ConstType): return self.mod(x) + def __and__(self:TMT, x:TMT|ConstType): return self.bitwise_and(x) + def __or__(self:TMT, x:TMT|ConstType): return self.bitwise_or(x) + def __xor__(self:TMT, x:TMT|ConstType): return self.bitwise_xor(x) - def __radd__(self, x): return self.add(x, True) - def __rsub__(self, x): return self.sub(x, True) - def __rmul__(self, x): return self.mul(x, True) - def __rtruediv__(self, x): return self.div(x, True) - def __rfloordiv__(self, x): return self.idiv(x, True) - def __rand__(self, x): return self.bitwise_and(x, True) - def __ror__(self, x): return self.bitwise_or(x, True) - def __rxor__(self, x): return self.bitwise_xor(x, True) - def __rmod__(self, x): return self.mod(x, True) + def __radd__(self:TMT, x:TMT|ConstType): return self.add(x, True) + def __rsub__(self:TMT, x:TMT|ConstType): return self.sub(x, True) + def __rmul__(self:TMT, x:TMT|ConstType): return self.mul(x, True) + def __rtruediv__(self:TMT, x:TMT|ConstType): return self.div(x, True) + def __rfloordiv__(self:TMT, x:TMT|ConstType): return self.idiv(x, True) + def __rand__(self:TMT, x:TMT|ConstType): return self.bitwise_and(x, True) + def __ror__(self:TMT, x:TMT|ConstType): return self.bitwise_or(x, True) + def __rxor__(self:TMT, x:TMT|ConstType): return self.bitwise_xor(x, True) + def __rmod__(self:TMT, x:TMT|ConstType): return self.mod(x, True) - def __lt__(self, x): return self.alu(Ops.CMPLT, self.ufix(x)) - def __gt__(self, x): return self.ufix(x).alu(Ops.CMPLT, self) - def __ge__(self, x): return (self < x).logical_not() - def __le__(self, x): return (self > x).logical_not() + def __lt__(self:TMT, x:TMT|ConstType): return self.alu(Ops.CMPLT, self.ufix(x)) + def __gt__(self:TMT, x:TMT|ConstType): return self.ufix(x).alu(Ops.CMPLT, self) + def __ge__(self:TMT, x:TMT|ConstType): return (self < x).logical_not() + def __le__(self:TMT, x:TMT|ConstType): return (self > x).logical_not() - def ne(self, x): return self.alu(Ops.CMPNE, self.ufix(x)) - def eq(self, x): return self.ne(x).logical_not() - def __ne__(self, x): return self.ne(x) + def ne(self:TMT, x:TMT|ConstType): return self.alu(Ops.CMPNE, self.ufix(x)) + def eq(self:TMT, x:TMT|ConstType): return self.ne(x).logical_not() + def __ne__(self:TMT, x:TMT|ConstType): return self.ne(x) # type: ignore[override] # NOTE: __eq__ isn't overridden, and means the same thing as is by default - def lshift(self, x, reverse=False): return self._binop(Ops.SHL, x, reverse) - def rshift(self, x, reverse=False): return self._binop(Ops.SHR, x, reverse) - def __lshift__(self, x): return self.lshift(x) - def __rshift__(self, x): return self.rshift(x) - def __rlshift__(self, x): return self.lshift(x, True) - def __rrshift__(self, x): return self.rshift(x, True) + def lshift(self:TMT, x:TMT|int, reverse:bool=False): return self._binop(Ops.SHL, x, reverse) + def rshift(self:TMT, x:TMT|int, reverse:bool=False): return self._binop(Ops.SHR, x, reverse) + def __lshift__(self:TMT, x:TMT|int): return self.lshift(x) + def __rshift__(self:TMT, x:TMT|int): return self.rshift(x) + def __rlshift__(self:TMT, x:TMT|int): return self.lshift(x, True) + def __rrshift__(self:TMT, x:TMT|int): return self.rshift(x, True) - def maximum(self, x): return self.alu(Ops.MAX, self.ufix(x)) - def minimum(self, x): return -(-self).maximum(-x) - def where(self, x, y): - if type(self) is type(x): return self.alu(Ops.WHERE, x, x.ufix(y)) - if type(self) is type(y): return self.alu(Ops.WHERE, y.ufix(x), y) + def maximum(self:TMT, x:TMT|ConstType): return self.alu(Ops.MAX, self.ufix(x)) + def minimum(self:TMT, x:TMT|ConstType): return -(-self).maximum(-x) + def where(self:TMT, x:TMT|ConstType, y:TMT|ConstType): + if isinstance(x, type(self)): return self.alu(Ops.WHERE, x, x.ufix(y)) + if isinstance(y, type(self)): return self.alu(Ops.WHERE, y.ufix(x), y) raise RuntimeError("where needs at least one UOp arg") - def threefry(self, seed): return self.alu(Ops.THREEFRY, seed) + def threefry(self:TMT, seed:TMT): return self.alu(Ops.THREEFRY, seed) def reciprocal(self): return self.alu(Ops.RECIP) def trunc(self): return self.alu(Ops.TRUNC) def sqrt(self): return self.alu(Ops.SQRT) def sin(self): return self.alu(Ops.SIN) def log2(self): return self.alu(Ops.LOG2) def exp2(self): return self.alu(Ops.EXP2) - def pow(self, x): return self.alu(Ops.POW, self.ufix(x)) - def __pow__(self, x): return self.pow(x) + def pow(self:TMT, x:TMT|ConstType): return self.alu(Ops.POW, self.ufix(x)) + def __pow__(self:TMT, x:TMT|ConstType): return self.pow(x) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 00353dd037..786c77b498 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -35,11 +35,11 @@ def resolve(x:UOp|bool, default:bool=True): def _suop(lst, uop_fxn, python_fxn): uops, nums = partition(lst, lambda x: isinstance(x, UOp)) return ssimplify(functools.reduce(uop_fxn, uops + ([python_fxn(nums)] if nums else []))) -def smax(*lst): return _suop(argfix(*lst), UOp.maximum, max) -def smin(*lst): return _suop(argfix(*lst), UOp.minimum, min) -def srender(x) -> str: return x.render() if isinstance(x, UOp) else str(x) +def smax(*lst) -> sint: return _suop(argfix(*lst), UOp.maximum, max) +def smin(*lst) -> sint: return _suop(argfix(*lst), UOp.minimum, min) +def srender(x:sint) -> str: return x.render() if isinstance(x, UOp) else str(x) -def ssimplify(uop): return uop.ssimplify() if isinstance(uop, UOp) else uop +def ssimplify(uop:sint): return uop.ssimplify() if isinstance(uop, UOp) else uop def sym_infer(uop: UOp|int, var_vals: dict[str, int]) -> int: return uop.sym_infer(var_vals) if isinstance(uop, UOp) else uop def range_str(u:UOp) -> str: return '_'.join([str(x) if x >= 0 else "m"+str(-x) for x in u.arg[0:-1]]) @@ -251,11 +251,11 @@ class UOp(MathTrait, metaclass=UOpMetaClass): # *** uop evaluation *** - def simplify(self, tracked=False): + def simplify(self, tracked=False, full_symbolic=True): # late import! - from tinygrad.uop.symbolic import symbolic + from tinygrad.uop.symbolic import symbolic, commutative with Context(TRACK_MATCH_STATS=0 if not tracked else TRACK_MATCH_STATS.value): - return graph_rewrite(self, symbolic, name="simplify") + return graph_rewrite(self, symbolic if full_symbolic else commutative, name="simplify") def ssimplify(self) -> UOp|ConstType: return ret.arg if (ret:=self.simplify()).op is Ops.CONST else ret def _eval(self, dtype, expected_type:Type[T]) -> T: assert self.dtype in dtype, f"eval with wrong dtype {self}" diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 84580039be..3e43d13161 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -177,8 +177,8 @@ def fold_binary_numerator(d: UOp, x: UOp, y: UOp) -> UOp|None: x,const = x.pop_const() terms, factors = zip(*[(u.divides(f:=u.const_factor()),f) for u in x.split_uop(Ops.ADD)]) if len(terms)==1 and (v:=terms[0]).vmax-v.vmin == 1: - y1 = cmod(factors[0]*v.vmin+const, c) if d.op is Ops.MOD else cdiv(factors[0]*v.vmin+const, c) # type: ignore - y2 = cmod(factors[0]*v.vmax+const, c) if d.op is Ops.MOD else cdiv(factors[0]*v.vmax+const, c) # type: ignore + y1 = cmod(factors[0]*v.vmin+const, c) if d.op is Ops.MOD else cdiv(factors[0]*v.vmin+const, c) + y2 = cmod(factors[0]*v.vmax+const, c) if d.op is Ops.MOD else cdiv(factors[0]*v.vmax+const, c) return (y2-y1)*(v-v.vmin) + y1 return None @@ -397,7 +397,7 @@ def parse_valid(valid:UOp) -> tuple[UOp, bool, int]: if valid.op is Ops.CMPLT and dtypes.is_int(valid.src[0].dtype): return valid.src[0], True, int((valid.src[1]).vmax)-1 raise ValueError(f"not able to parse {valid=}") -def uop_given_valid(valid:UOp, uop:UOp) -> UOp: +def uop_given_valid(valid:UOp, uop:UOp, try_simplex=True) -> UOp: # return simplified uop (might be the same as input) # first, parse valid into {expr: (lower_bound, upper_bound)} @@ -411,26 +411,33 @@ def uop_given_valid(valid:UOp, uop:UOp) -> UOp: uop = uop.substitute((load_subs:={u: UOp(Ops.NOOP, arg=u) for u in uop.toposort() if u.op is Ops.INDEX})) # simplify uop given that valid is True - for expr,v in bounds.items(): + all_candidates = [] + for i,(expr,v) in enumerate(bounds.items()): v0, v1 = (expr.vmin if v[0] is None else v[0], expr.vmax if v[1] is None else v[1]) expr = expr.substitute(load_subs) # make sure expr appears in same form in the uop - # some expr has lower bound > upper bound -> valid is an empty set and we return None - # every candidate is a set of constrained UOp based on valid, and if every item in a set simplifies the uop into a same output, we rewrite uop - candidates = [] - if expr.op is Ops.ADD and v0 == 1 and all(u.op in GroupOp.Irreducible for u in expr.split_uop(Ops.ADD)): - # if the constraint is a simplex: X0 + X1 + ... > 0, we can check if all Xi > 0 simplify into the same output - candidates.append([(Xi, UOp.variable("fake", 1, Xi.vmax, Xi.dtype)) for Xi in expr.split_uop(Ops.ADD)]) # try checking the whole clause - candidates.append([(expr, UOp.variable("fake", v0, v1, expr.dtype))]) + all_candidates.append((expr, UOp.variable(f"fake{i}", v0, v1, expr.dtype))) - for candidate in candidates: - # if every branch in candidate gives the same simplified uop, we can rewrite the uop - newuops = [uop.substitute({X:newX}).simplify().substitute({newX:X}).simplify() for X,newX in candidate] - if uop.op is Ops.VECTORIZE and len(uop.src) == 2: - if all_same([uops.src[0] for uops in newuops]): uop = uop.replace(src=(newuops[0].src[0], uop.src[1])) - if all_same([uops.src[1] for uops in newuops]): uop = uop.replace(src=(uop.src[0], newuops[0].src[1])) - elif all_same(newuops): uop = newuops[0] + if try_simplex: + # every candidate is a set of constrained UOp based on valid, and if every item in a set simplifies the uop into a same output, we rewrite uop + candidates = [[all_candidates[-1]]] + if expr.op is Ops.ADD and v0 == 1 and all(u.op in GroupOp.Irreducible for u in expr.split_uop(Ops.ADD)): + # if the constraint is a simplex: X0 + X1 + ... > 0, we can check if all Xi > 0 simplify into the same output + candidates.append([(Xi, UOp.variable(f"fake{i}", 1, Xi.vmax, Xi.dtype)) for Xi in expr.split_uop(Ops.ADD)]) + for candidate in candidates: + # if every branch in candidate gives the same simplified uop, we can rewrite the uop + newuops = [uop.substitute({X:newX}) for X,newX in candidate] + if any(u is uop for u in newuops): continue # if any branch doesnt appear in uop, skip + newuops = [u.simplify().substitute({newX:X}).simplify(full_symbolic=False) for (X,newX),u in zip(candidate,newuops)] + if uop.op is Ops.VECTORIZE and len(uop.src) == 2: + if all_same([uops.src[0] for uops in newuops]): uop = uop.replace(src=(newuops[0].src[0], uop.src[1])) + if all_same([uops.src[1] for uops in newuops]): uop = uop.replace(src=(uop.src[0], newuops[0].src[1])) + elif all_same(newuops): uop = newuops[0] + + # try all the valids together (but only the whole expressions) + if (s_uop:=uop.substitute(sub_dict:=dict(all_candidates))) is not uop: + uop = s_uop.simplify().substitute({newX:X for X,newX in sub_dict.items()}).simplify(full_symbolic=False) # put the loads back in uop = uop.substitute({v:k for k,v in load_subs.items()}) return uop @@ -463,14 +470,21 @@ def reduce_mul_chain(r:UOp): if len(outside) == 0: return None return r.replace(src=(prod(inside) if len(inside) else r.src[0].const_like(1),)+r.src[1:])*prod(outside) +def drop_and_clauses(cond:UOp, x:UOp, i:UOp) -> UOp|None: + if not (dropped_clauses:=[c for c in cond.split_uop(Ops.AND) if not any(r in x.ranges for r in c.ranges)]): return None + return functools.reduce(operator.and_, [c for c in cond.split_uop(Ops.AND) if c not in dropped_clauses], UOp.const(dtypes.bool, True)).where(x, i) +pm_drop_and_clauses = PatternMatcher([(UPat.var("cond").where(UPat.var("x", dtype=dtypes.index), invalid_pat), drop_and_clauses)]) + +pm_simplify_valid = PatternMatcher([ + # simplify valid + (UPat(Ops.AND, name="valid"), simplify_valid), + (UPat.var("c").where(UPat.var("x", dtype=dtypes.index), invalid_pat), lambda c,x,i: c.where(uop_given_valid(c, x, try_simplex=False), i)), +]) + # this is symbolic 2.0 REMOVE_FROM_SINK = {Ops.SINK, Ops.UNROLL, Ops.PTRCAT, Ops.CAT, Ops.NOOP} REMOVE_FROM_BARRIER = {Ops.VECTORIZE, Ops.SINK, Ops.CAT, Ops.PTRCAT, Ops.NOOP} -sym = symbolic_flat+PatternMatcher([ - # simplify valid - (UPat(Ops.AND, name="valid"), simplify_valid), - (UPat.var("cond").where(UPat.var("x", dtype=dtypes.index), invalid_pat), lambda cond,x,i: cond.where(newx, i) if - (newx:=uop_given_valid(cond, x)) is not x else None), +sym = symbolic_flat+pm_simplify_valid+PatternMatcher([ # LOAD/STORE -> NOOP (UPat.var('x').store(UPat.var('x').load(), allow_any_len=True), lambda x: None if x.dtype.addrspace != AddrSpace.REG else x.src[0].src[0]), (UPat(Ops.LOAD, src=(UPat.cvar('c'))), lambda c: c), @@ -505,8 +519,8 @@ sym = symbolic_flat+PatternMatcher([ (UPat((Ops.LOAD, Ops.STORE), src=(UPat().index(UPat.const(dtypes.index, Invalid)).or_casted(),), allow_any_len=True, name="x"), lambda x: UOp(Ops.NOOP) if x.op is Ops.STORE else x.const_like(0)), # invalid store does nothing. invalid load produces 0 # # Where after gated load becomes alt value, TODO: this is sort of duplicated with rules in devectorizer - (UPat.var("c1").where(UPat(Ops.LOAD, src=(UPat().index(UPat.var("c2").where(UPat(), invalid_pat)).or_casted(),), allow_any_len=True, name="l"), 0), - lambda c1,c2,l,i: l.replace(src=(l.src[0],)+l.src[1:]) if any(c in list(c2.split_uop(Ops.AND)) for c in c1.split_uop(Ops.AND)) else None), + (UPat.var("c1").where(UPat(Ops.LOAD, src=(UPat().index(UPat.var("c2").where(UPat(), invalid_pat)).or_casted(),), name="l"), 0), + lambda c1,c2,l,i: l.replace(src=(l.src[0],)+l.src[1:]) if all(c in list(c2.split_uop(Ops.AND)) for c in c1.split_uop(Ops.AND)) else None), # remove VECTORIZE from SINK/BARRIER. TODO: SINK/BARRIER are really the same thing at GLOBAL/LOCAL levels (UPat(Ops.BARRIER, name="root"), lambda root: UOp(Ops.BARRIER, root.dtype, tuple(flatten(x.src if x.op in REMOVE_FROM_BARRIER else (x,) for x in root.src)), root.arg) diff --git a/tinygrad/viz/js/index.js b/tinygrad/viz/js/index.js index 92223b6682..8d2fd5cf2a 100644 --- a/tinygrad/viz/js/index.js +++ b/tinygrad/viz/js/index.js @@ -18,6 +18,9 @@ const ANSI_COLORS_LIGHT = ["#d9d9d9","#ff9999","#99cc99","#ffff99","#9999ff","#f const parseColors = (name, defaultColor="#ffffff") => Array.from(name.matchAll(/(?:\u001b\[(\d+)m([\s\S]*?)\u001b\[0m)|([^\u001b]+)/g), ([_, code, colored_st, st]) => ({ st: colored_st ?? st, color: code != null ? (code>=90 ? ANSI_COLORS_LIGHT : ANSI_COLORS)[(parseInt(code)-30+60)%60] : defaultColor })); +const colored = n => d3.create("span").call(s => s.selectAll("span").data(typeof n === "string" ? parseColors(n) : n).join("span") + .style("color", d => d.color).text(d => d.st)).node(); + const rect = (s) => (typeof s === "string" ? document.querySelector(s) : s).getBoundingClientRect(); let timeout = null; @@ -174,7 +177,7 @@ function tabulate(rows) { var data, focusedDevice, focusedShape, canvasZoom, zoomLevel = d3.zoomIdentity; async function renderProfiler() { displayGraph("profiler"); - d3.select(".metadata").html(""); + d3.select(".metadata").node().replaceChildren(focusedShape?.html ?? ""); // layout once! if (data != null) return updateProgress({ start:false }); const profiler = d3.select(".profiler").html(""); @@ -236,8 +239,7 @@ async function renderProfiler() { const stepIdx = ctxs[ref.ctx+1].steps.findIndex((s, i) => i >= start && s.name == e.name); if (stepIdx !== -1) { ref.step = stepIdx; shapeRef = ref; } } - const htmlLabel = label.map(({color, st}) => `${st}`).join(''); - const arg = { tooltipText:htmlLabel+"\n"+formatTime(e.dur)+(e.info != null ? "\n"+e.info : ""), ...shapeRef }; + const arg = { tooltipText:colored(e.name).outerHTML+"\n"+formatTime(e.dur)+(e.info != null ? "\n"+e.info : ""), ...shapeRef }; // offset y by depth shapes.push({x:e.st, y:levelHeight*depth, width:e.dur, height:levelHeight, arg, label, fillColor }); } @@ -257,8 +259,8 @@ async function renderProfiler() { x += 1; y += nbytes; valueMap.set(ts, y); } else { const free = buf_shapes.get(key); - timestamps.push(ts); - x += 1; y -= free.nbytes; valueMap.set(ts, y); + timestamps.push(ts); valueMap.set(ts, y); + x += 1; y -= free.nbytes; free.x.push(x); free.y.push(free.y.at(-1)); temp.delete(key); @@ -333,6 +335,7 @@ async function renderProfiler() { const st = visibleX[0], et = visibleX[1]; xscale.domain(visibleX); // draw shapes + const paths = []; for (const [_, { offsetY, shapes, visible, valueMap }] of data.tracks) { visible.length = 0; for (const e of shapes) { @@ -340,18 +343,18 @@ async function renderProfiler() { if (e.width == null) { if (e.x[0]>et || e.x.at(-1)=0; i--) ctx.lineTo(x[i], offsetY+e.y1[i]); - ctx.closePath(); - ctx.fillStyle = e.fillColor; ctx.fill(); - if (focusedShape && e.arg?.key === focusedShape) { ctx.lineWidth = 1.4; ctx.strokeStyle = "#c9a8ff"; ctx.stroke(); } + for (let i=x.length-1; i>=0; i--) p.lineTo(x[i], offsetY+e.y1[i]); + p.closePath(); + ctx.fillStyle = e.fillColor; ctx.fill(p); + if (focusedShape && e.arg?.key === focusedShape.key) { paths.push(p); } continue; } // contiguous rect @@ -400,11 +403,13 @@ async function renderProfiler() { } } // draw markers + ctx.textBaseline = "top"; for (const m of markers) { const x = xscale(m.ts); drawLine(ctx, [x, x], [0, canvas.clientHeight], { color:m.color }); ctx.fillText(m.name, x+2, 1); } + for (const p of paths) { ctx.lineWidth = 1.4; ctx.strokeStyle = "#c9a8ff"; ctx.stroke(p); } } function resize() { @@ -445,7 +450,7 @@ async function renderProfiler() { e.preventDefault(); const foundRect = findRectAtPosition(e.clientX, e.clientY); if (foundRect?.step != null) return setCtxWithHistory(foundRect.ctx, foundRect.step); - if (foundRect?.key != focusedShape) { focusedShape = foundRect?.key; render(zoomLevel); } + if (foundRect?.key != focusedShape?.key) { focusedShape = foundRect; render(zoomLevel); } return document.querySelector(".metadata").replaceChildren(foundRect?.html ?? ""); }); @@ -589,7 +594,7 @@ async function main() { const ul = ctxList.appendChild(document.createElement("ul")); ul.id = `ctx-${i}`; const p = ul.appendChild(document.createElement("p")); - p.innerHTML = parseColors(name).map(c => `${c.st}`).join(""); + p.appendChild(colored(name)); p.onclick = () => { setState(i === state.currentCtx ? { expandSteps:!state.expandSteps } : { expandSteps:true, currentCtx:i, currentStep:0, currentRewrite:0 }); } @@ -703,9 +708,7 @@ async function main() { metadata.appendChild(codeBlock(upat[1], "python", { loc:upat[0], wrap:true })); const diffCode = metadata.appendChild(document.createElement("pre")).appendChild(document.createElement("code")); for (const line of diff) { - const span = diffCode.appendChild(document.createElement("span")); - span.style.color = line.startsWith("+") ? "#3aa56d" : line.startsWith("-") ? "#d14b4b" : "#f0f0f5"; - span.innerText = line; + diffCode.appendChild(colored([{st:line, color:line.startsWith("+") ? "#3aa56d" : line.startsWith("-") ? "#d14b4b" : "#f0f0f5"}])); diffCode.appendChild(document.createElement("br")); } diffCode.className = "wrap";