tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-30 03:46:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'upstream/master' into new_x86_backend

Browse Source
This commit is contained in:
ttomsa
2026-03-21 19:16:59 +00:00
parent 292e1745b2 a9ceaf3c5f
commit 35fc12b839
77 changed files with 834 additions and 898 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/actions/setup-tinygrad/action.yml
+1
View File
@@ -287,6 +287,7 @@ runs:
CMAKE_ARGS="-Wno-dev -G Ninja -DOCELOT_BUILD_TOOLS=OFF -DCMAKE_BUILD_ALWAYS=0 -DBUILD_TESTS_CUDA=OFF -DCMAKE_POLICY_VERSION_MINIMUM=3.5"
if [[ "${{ runner.os }}" == "macOS" ]]; then
sudo xcode-select -s /Applications/Xcode_16.2.app/Contents/Developer
CMAKE_ARGS="$CMAKE_ARGS -DBoost_INCLUDE_DIR=$(brew --prefix boost)/include -DBoost_LIBRARY_DIR=$(brew --prefix boost)/lib"
fi
.github/workflows/benchmark.yml
+12 -20
View File
@@ -49,8 +49,6 @@ jobs:
source /tmp/tinygrad_pytest_ci/bin/activate
pytest -nauto --durations=20
- name: openpilot compile3 0.10.1 driving_vision
run: FLOAT16=1 CL=1 IMAGE=2 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
- name: IMAGE=1 openpilot compile3 0.10.1 driving_vision
run: FLOAT16=1 CL=1 IMAGE=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
testmacbenchmark:
@@ -187,13 +185,13 @@ jobs:
PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
PYTHONPATH=. ./extra/hcq/hcq_smi.py nv kill_pids
- name: UsbGPU boot time
run: sudo -E PYTHONPATH=. DEBUG=2 AM_RESET=1 AMD=1 AMD_IFACE=USB time python3.11 test/test_tiny.py TestTiny.test_plus
run: sudo -E PYTHONPATH=. GMMU=0 DEBUG=2 AM_RESET=1 AMD=1 AMD_IFACE=USB time python3.11 test/test_tiny.py TestTiny.test_plus
- name: UsbGPU tiny tests
run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/test_tiny.py
run: sudo -E PYTHONPATH=. GMMU=0 AMD=1 AMD_IFACE=USB python3.11 test/test_tiny.py
- name: UsbGPU copy speeds
run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
run: sudo -E PYTHONPATH=. GMMU=0 AMD=1 AMD_IFACE=USB python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
#- name: UsbGPU openpilot test
# run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
# run: sudo -E PYTHONPATH=. GMMU=0 AMD=1 AMD_IFACE=USB GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
- name: UsbGPU (USB4/TB) boot time
run: PYTHONPATH=. DEBUG=3 NV=1 NV_IFACE=PCI NV_NAK=1 time python3.11 test/test_tiny.py TestTiny.test_plus
- name: UsbGPU (USB4/TB) tiny tests
@@ -590,26 +588,20 @@ jobs:
rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
- name: reset process replay
run: test/external/process_replay/reset.py
- name: IMAGE=1 openpilot compile3 0.11.0 driving_vision
run: BENCHMARK_LOG=image_1_openpilot_0_11_0_vision PYTHONPATH="." DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
- name: openpilot compile3 0.11.0 driving_vision
run: BENCHMARK_LOG=openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
run: BENCHMARK_LOG=openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
- name: openpilot compile3 0.11.0 driving_policy
run: BENCHMARK_LOG=openpilot_0_11_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_policy.onnx
run: BENCHMARK_LOG=openpilot_0_11_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_policy.onnx
- name: openpilot compile3 0.11.0 dmonitoring
run: BENCHMARK_LOG=openpilot_0_11_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/dmonitoring_model.onnx
run: BENCHMARK_LOG=openpilot_0_11_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/dmonitoring_model.onnx
- name: DEBUG=2 openpilot compile3 0.10.1 driving_vision
run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
- name: DEBUG=2 IMAGE=1 openpilot compile3 0.10.1 driving_vision
run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
- name: IMAGE=1 openpilot compile3 0.10.1 driving_vision
run: BENCHMARK_LOG=image_1_openpilot_0_10_1_vision PYTHONPATH="." DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
- name: openpilot compile3 0.10.1 driving_vision
run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
- name: openpilot compile3 0.10.1 driving_policy
run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
- name: openpilot compile3 0.10.1 dmonitoring
run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
- name: benchmark MobileNetV2 on DSP
run: |
# generate quantized weights
@@ -638,9 +630,9 @@ jobs:
echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
- name: openpilot compile3 0.10.1 driving_vision
run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision PYTHONPATH="." DEV=AMD AMD_LLVM=1 AMD_IFACE=USB ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision PYTHONPATH="." GMMU=0 DEV=AMD AMD_LLVM=1 AMD_IFACE=USB ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
- name: openpilot load_pickle 0.10.1 driving_vision
run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_load_pickle PYTHONPATH="." DEV=AMD AMD_IFACE=USB ASSERT_MIN_LOAD_TIME=15 python3 examples/openpilot/load_pickle.py
run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_load_pickle PYTHONPATH="." GMMU=0 DEV=AMD AMD_IFACE=USB ASSERT_MIN_LOAD_TIME=15 python3 examples/openpilot/load_pickle.py
testreddriverbenchmark:
name: AM Benchmark
.github/workflows/test.yml
+18 -13
View File
@@ -167,10 +167,10 @@ jobs:
run: PYTHON=1 python3 test/backend/test_symbolic_ops.py
- name: test_renderer_failures with Python emulator
run: PYTHON=1 python3 -m pytest -rA test/backend/test_renderer_failures.py::TestRendererFailures
- name: Test IMAGE=2 support
- name: Test IMAGE support
run: |
IMAGE=2 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm
IMAGE=2 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_simple_conv2d
IMAGE=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm
IMAGE=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_simple_conv2d
- name: Test emulated METAL tensor cores
run: |
DEBUG=2 EMULATE=METAL FORWARD_ONLY=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_big_gemm
@@ -369,11 +369,11 @@ jobs:
key: gpu-image
deps: testing_unit
opencl: 'true'
- name: Test CL IMAGE=2 ops
- name: Test CL IMAGE=1 ops
run: |
CL=1 IMAGE=2 python -m pytest -n=auto test/backend/test_ops.py --durations=20
CL=1 IMAGE=1 python -m pytest -n=auto test/backend/test_ops.py --durations=20
# TODO: training is broken
# CL=1 IMAGE=2 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
# CL=1 IMAGE=1 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
- name: Run process replay tests
uses: ./.github/actions/process-replay
@@ -418,13 +418,13 @@ jobs:
llvm: 'true'
- name: Test openpilot model kernel count and gate usage
run: |
ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1397 ALLOWED_GATED_READ_IMAGE=94 FLOAT16=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1486 ALLOWED_GATED_READ_IMAGE=17 FLOAT16=1 CL=1 IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
- name: Test openpilot CL compile fp16
run: FLOAT16=1 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
run: FLOAT16=1 DEBUGCL=1 CL=1 IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
- name: Test openpilot CL compile fp32 (test correctness)
run: DEBUGCL=1 CL=1 IMAGE=2 SELFTEST=1 python examples/openpilot/compile3.py https://github.com/haraschax/filedump/raw/refs/heads/master/driving_vision_fp32.onnx
run: CL=1 IMAGE=1 SELFTEST=1 python examples/openpilot/compile3.py https://github.com/haraschax/filedump/raw/refs/heads/master/driving_vision_fp32.onnx
- name: Test openpilot LLVM compile fp16
run: FLOAT16=1 CPU=1 CPU_LLVM=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
run: IMAGE=1 FLOAT16=1 CPU=1 CPU_LLVM=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
- name: Run process replay tests
uses: ./.github/actions/process-replay
@@ -505,8 +505,13 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: apps_llm
- name: Test 1B LLM
run: echo "What's a male chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm | grep -i rooster
- name: Test 1B LLM (llama)
run: echo "What's a male chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm --model llama3.2:1b | tee /dev/stderr | grep -i rooster
- name: Test 1B LLM (llama q4)
run: echo "What's a male chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm --model llama3.2:1b-q4 | tee /dev/stderr | grep -i rooster
- name: Test 1B LLM (qwen)
# NOTE: qwen is dumb and only knows about female chickens
run: echo "What's a female chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm --model qwen3:0.6b | tee /dev/stderr | grep -i hen
# ****** Models Tests ******
@@ -699,7 +704,7 @@ jobs:
- name: Run test_tiny on MOCKAM
run: python test/test_tiny.py
- name: Run test_tiny on MOCKAM USB
run: AMD_IFACE=USB python test/test_tiny.py
run: GMMU=0 AMD_IFACE=USB python test/test_tiny.py
- name: Run test_hcq on MOCKAM
run: python -m pytest test/device/test_hcq.py
examples/mlperf/model_train.py
+26 -56
View File
@@ -13,6 +13,8 @@ from extra.bench_log import BenchEvent, WallTimeEvent
# TODO: fix benchmark logging and use tinygrad tqdm
from tqdm import tqdm
from tinygrad.uop.ops import UOp
def train_resnet():
from extra.models import resnet
from examples.mlperf.dataloader import batch_load_resnet
@@ -1282,7 +1284,7 @@ def train_bert():
previous_step = i
def train_llama3():
from examples.mlperf.models.llama import Transformer
from examples.mlperf.models.flat_llama import FlatTransformer, apply_grad
from examples.llama3 import MODEL_PARAMS
from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
from examples.mlperf.optim import GradAccClipAdamW
@@ -1343,7 +1345,7 @@ def train_llama3():
if (MP := getenv("MP", 1)) > 1: model_params['vocab_size'] = round_up(model_params['vocab_size'], 256 * MP)
vocab_mask:Tensor = Tensor.arange(model_params['vocab_size']).reshape(1, 1, -1) >= real_vocab_size
model = Transformer(**model_params, max_context=SEQLEN)
model = FlatTransformer(**model_params, max_context=SEQLEN)
params = get_parameters(model)
# weights are all bfloat16 for now
@@ -1353,49 +1355,25 @@ def train_llama3():
for v in get_parameters(model):
v = v.assign(Tensor.empty(v.shape))
if (DP := getenv("DP", 1)) > 1:
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
for v in get_parameters(model):
v.shard_(device, axis=None)
is_dp = (DP := getenv("DP", 1)) > 1
is_mp = (MP := getenv("MP", 1)) > 1
is_sharding = is_dp or is_mp
device_count = max(DP, MP)
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(device_count))
vocab_mask.shard_(device, axis=None)
model.shard(device, is_mp)
if (MP := getenv("MP", 1)) > 1:
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
for k,v in get_state_dict(model).items():
if 'scale' in k: v.shard_(device, axis=None) # from quantized
elif '.attention.wq' in k: v.shard_(device, axis=0)
elif '.attention.wk' in k: v.shard_(device, axis=0)
elif '.attention.wv' in k: v.shard_(device, axis=0)
elif '.attention.wqkv' in k: v.shard_(device, axis=0)
elif '.attention.wo' in k: v.shard_(device, axis=1)
elif '.feed_forward.w1.' in k: v.shard_(device, axis=0)
elif '.feed_forward.w2.' in k: v.shard_(device, axis=1)
elif '.feed_forward.w3.' in k: v.shard_(device, axis=0)
elif 'tok_embeddings.weight' in k: v.shard_(device, axis=0)
elif 'output.weight' in k: v.shard_(device, axis=0)
else:
# attention_norm, ffn_norm, norm
v.shard_(device, axis=None)
# prevents memory spike on device 0
v.realize()
vocab_mask.shard_(device, axis=2).realize()
if is_dp: vocab_mask.shard_(device, axis=None).realize()
if is_mp: vocab_mask.shard_(device, axis=2).realize()
is_offload_optim = bool(getenv("OFFLOAD_OPTIM"))
is_fake_offload = Device.DEFAULT == "NULL"
optim_device = ("CPU" if not is_fake_offload else "NULL:99") if is_offload_optim else None
optim = GradAccClipAdamW(get_parameters(model), lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2,
optim = GradAccClipAdamW(params, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2,
eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc, device=optim_device)
# init grads
if is_offload_optim:
for p in optim.params:
p.grad = Tensor.zeros(p.shape, dtype=p.dtype, device=optim_device, requires_grad=False).contiguous().realize()
else:
for p in optim.params:
p.grad = p.zeros_like().contiguous().realize()
grads: list[Tensor] = [p.grad for p in optim.params]
grads = [Tensor.zeros_like(p).contiguous() for p in optim.params]
scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)
@@ -1410,28 +1388,24 @@ def train_llama3():
@TinyJit
def minibatch(tokens:Tensor):
if (DP := getenv("DP", 1)) > 1:
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
tokens = tokens.to(None).shard(device, 0)
if (MP := getenv("MP", 1)) > 1:
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
tokens = tokens.shard(device)
if DP == 1 and MP == 1: tokens = tokens.to(None)
if is_dp: tokens = tokens.to(None).shard(device, 0)
if is_mp: tokens = tokens.shard(device)
if not is_sharding: tokens = tokens.to(None)
logits:Tensor = model(tokens[:, :-1])
loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
loss.backward()
assert all(p.grad is g for p,g in zip(optim.params, grads))
for i,(t,g) in enumerate(zip(optim.params, loss.gradient(*optim.params))):
grads[i].replace(Tensor(grads[i].uop.after(UOp.group(*apply_grad(grads[i].uop, g.uop))), device=t.device))
loss_cpu = loss.flatten().float().to("CPU")
Tensor.realize(loss_cpu, *grads)
return loss_cpu
return loss_cpu.realize(*grads)
@TinyJit
def optim_step():
grad_norm = optim.fstep(grads)
scheduler.step()
for g in grads:
g.assign(g.zeros_like())
for g in grads: g.assign(g.zeros_like())
lr_cpu = optim.lr.float().to("CPU")
grad_norm_cpu = grad_norm.float().to("CPU")
@@ -1442,13 +1416,9 @@ def train_llama3():
@TinyJit
@Tensor.train(False)
def eval_step(tokens:Tensor):
if (DP := getenv("DP", 1)) > 1:
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
tokens = tokens.to(None).shard(device, 0)
if (MP := getenv("MP", 1)) > 1:
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
tokens = tokens.shard(device)
if DP == 1 and MP == 1: tokens = tokens.to(None)
if is_dp: tokens = tokens.to(None).shard(device, 0)
if is_mp: tokens = tokens.shard(device)
if not is_sharding: tokens = tokens.to(None)
logits:Tensor = model(tokens[:, :-1])
loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
return loss.flatten().float().to("CPU")
examples/mlperf/models/flat_llama.py
+43 -25
View File
@@ -2,7 +2,7 @@ import math, os
if __name__ == "__main__":
os.environ["DEFAULT_FLOAT"] = "bfloat16"
os.environ["OPTIM_DTYPE"] = "bfloat16"
os.environ["DEV"] = "NULL"
if "DEV" not in os.environ: os.environ["DEV"] = "NULL"
# CDNA
os.environ["EMULATE"] = "AMD_CDNA4"
os.environ["DEVICE_IN_FUNCTION_BUG"] = "1"
@@ -13,10 +13,27 @@ if __name__ == "__main__":
if "ASM_GEMM" not in os.environ:
os.environ["ASM_GEMM"] = "1"
from tinygrad import Tensor, nn, function, getenv, dtypes, TinyJit
from tinygrad.helpers import Timing, colored, GlobalCounters
from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
from tinygrad.uop.ops import Ops, UOp
from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
FP8 = getenv("FP8", 0)
FP8_DTYPE = dtypes.fp8e4m3
FP8_MAX = 448.0
def quantize_fp8(x:Tensor):
scale = FP8_MAX / (x.abs().max().detach() + 1e-8)
x_scaled = x * scale
x_clamped = x_scaled + (x_scaled.detach().clamp(-FP8_MAX, FP8_MAX) - x_scaled.detach()) # STE
return x_clamped.cast(FP8_DTYPE), scale.float().reciprocal()
def matmul(x:Tensor, w:Tensor) -> Tensor:
if not FP8: return x @ w.T
# weights are already FP8, just quantize activations
x_fp8, x_scale = quantize_fp8(x)
return x_fp8.dot(w.T, dtype=dtypes.float) * x_scale
def rmsnorm(x_in:Tensor, eps:float):
x = x_in.float()
x = x * (x.square().mean(-1, keepdim=True) + eps).rsqrt()
@@ -53,12 +70,13 @@ class FlatTransformer:
def lin_per_layer(self, in_features:int, out_features:int):
bound = 1 / math.sqrt(in_features)
if getenv("ZEROS"): return Tensor.zeros(self.n_layers, out_features, in_features)
return Tensor.uniform(self.n_layers, out_features, in_features, low=-bound, high=bound)
dt = FP8_DTYPE if FP8 else None
if getenv("ZEROS"): return Tensor.zeros(self.n_layers, out_features, in_features, dtype=dt)
return Tensor.uniform(self.n_layers, out_features, in_features, low=-bound, high=bound, dtype=dt)
def attention(self, x:Tensor, freqs_cis:Tensor, attention_norm:Tensor, wqkv:Tensor, wo:Tensor):
x = rmsnorm(x, self.norm_eps) * attention_norm
xqkv = x @ wqkv.T
xqkv = matmul(x, wqkv)
bsz, seqlen, _ = xqkv.shape
# interleaved layout: each kv group has [n_rep q heads, 1 k head, 1 v head] for clean MP sharding
@@ -71,13 +89,13 @@ class FlatTransformer:
xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True).transpose(1, 2)
attn = attn.reshape(bsz, seqlen, -1)
return attn @ wo.T
return matmul(attn, wo)
def feed_forward(self, x:Tensor, ffn_norm:Tensor, w1:Tensor, w2:Tensor, w3:Tensor):
x = rmsnorm(x, self.norm_eps) * ffn_norm
x_w1 = (x @ w1.T).silu()
x_w3 = x.contiguous_backward() @ w3.T
return (x_w1 * x_w3) @ w2.T
x_w1 = matmul(x, w1).silu()
x_w3 = matmul(x.contiguous_backward(), w3)
return matmul(x_w1 * x_w3, w2)
@function(precompile=True, precompile_backward=True)
def run_layer(self, x:Tensor, freqs_cis:Tensor,
@@ -114,6 +132,16 @@ class FlatTransformer:
logits = self.output(self.norm(h))
return logits
# TODO: this shouldn't be needed, but it prevents a copy of the grads. CAT can help
def apply_grad(old_grad:UOp, new_grad:UOp) -> list[UOp]:
if new_grad.op == Ops.ADD:
return apply_grad(old_grad, new_grad.src[0])+apply_grad(old_grad, new_grad.src[1])
elif new_grad.op == Ops.PAD:
grad_shrink = tuple([(p[0], s+p[0]) for s,p in zip(new_grad.src[0].shape, new_grad.marg)])
return apply_grad(old_grad.shrink(grad_shrink), new_grad.src[0])
else:
return [old_grad.store(old_grad + new_grad)]
if __name__ == "__main__":
config = {}
BS = config["BS"] = getenv("BS", 16)
@@ -139,7 +167,7 @@ if __name__ == "__main__":
# print model size
sz = 0
for k,v in state.items():
print(f"{colored(k, 'green' if v in grads else 'white'):30s} {str(v.shape):30s} {v.dtype} {v.device} {v.nbytes()/1e9:.2f} GB")
print(f"{colored(k, 'green' if v in grads else 'white'):30s} {str(v.shape):30s} {str(v.dtype):20s} {v.device} {v.nbytes()/1e9:.2f} GB")
sz += v.nbytes()
print(f"total sz: {sz/1e9:.2f} GB")
@@ -149,27 +177,17 @@ if __name__ == "__main__":
if DP > 1: tokens = tokens.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(DP)), axis=0)
if MP > 1: tokens = tokens.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(MP)))
# TODO: this shouldn't be needed, but it prevents a copy of the grads. CAT can help
def apply_grad(old_grad:UOp, new_grad:UOp) -> list[UOp]:
if new_grad.op == Ops.ADD:
return apply_grad(old_grad, new_grad.src[0])+apply_grad(old_grad, new_grad.src[1])
elif new_grad.op == Ops.PAD:
grad_shrink = tuple([(p[0], s+p[0]) for s,p in zip(new_grad.src[0].shape, new_grad.marg)])
return apply_grad(old_grad.shrink(grad_shrink), new_grad.src[0])
else:
return [old_grad.store(old_grad + new_grad)]
@TinyJit
def jit_step(tokens:Tensor):
GlobalCounters.reset()
print(colored("*** step", "red"))
with Timing("python forward: "): loss = model(tokens[:, :-1]).sparse_categorical_crossentropy(tokens[:, 1:])
with Timing("python backward: "):
for t,g in zip(grads, loss.gradient(*grads)):
grads[t] = Tensor(grads[t].uop.after(UOp.group(*apply_grad(grads[t].uop, g.uop))), device=t.device)
with Timing("run step: "): loss.realize(*grads.values())
jit_step(tokens)
jit_step(tokens)
jit_step(tokens)
for i in range(6):
GlobalCounters.reset()
profile_marker(f"step {i}")
with Timing(colored(f"*** step {i}: ", "red")):
jit_step(tokens)
print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
examples/mlperf/models/test_flat_llama.py
+24 -1
View File
@@ -2,8 +2,9 @@ import os
os.environ["WQKV"] = "1"
import unittest
import numpy as np
from tinygrad import Tensor, nn
from tinygrad import Tensor, nn, dtypes
from tinygrad.nn.state import get_parameters
from tinygrad.device import is_dtype_supported
from examples.mlperf.models.llama import Transformer
from examples.mlperf.models.flat_llama import FlatTransformer
@@ -113,5 +114,27 @@ class TestFlatLlama(unittest.TestCase):
self.assertEqual(ref_logits.shape, flat_logits.shape)
np.testing.assert_allclose(flat_logits, ref_logits, atol=1e-4, rtol=1e-4)
@unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3), "fp8 not supported on this device")
def test_forward_fp8(self):
import examples.mlperf.models.flat_llama as flat_llama_mod
old_fp8 = flat_llama_mod.FP8
try:
flat_llama_mod.FP8 = 1
Tensor.manual_seed(42)
params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
ref = Transformer(**params)
flat = FlatTransformer(**params)
copy_weights(flat, ref)
Tensor.realize(*nn.state.get_state_dict(flat).values())
tokens = Tensor([[1, 50, 100, 999, 2]])
ref_logits = ref(tokens).numpy()
flat_logits = flat(tokens).numpy()
self.assertEqual(ref_logits.shape, flat_logits.shape)
# FP8 has lower precision, allow larger tolerance
np.testing.assert_allclose(flat_logits, ref_logits, atol=1.0, rtol=0.1)
finally:
flat_llama_mod.FP8 = old_fp8
if __name__ == "__main__":
unittest.main()
examples/mlperf/optim.py
+4 -2
View File
@@ -25,8 +25,10 @@ class GradAccClipAdamW(Optimizer):
return extra[-1]
def _step(self, params:list[Tensor], grads:list[Tensor]) -> tuple[list[Tensor], list[Tensor]]:
grads = list(grads)
for i in range(len(grads)):
if grads[i].device != self.m[i].device: grads[i].assign(grads[i].to(self.m[i].device))
if grads[i].device != self.m[i].device: grads[i] = grads[i].to(self.m[i].device)
if self.fused:
grads[0].assign(grads[0] / self.grad_acc)
@@ -52,6 +54,6 @@ class GradAccClipAdamW(Optimizer):
return ret, [self.b1_t, self.b2_t] + self.m + self.v + [total_norm]
def _apply_update(self, t:Tensor, up:Tensor) -> Tensor:
wd = self.wd if t.ndim >= 2 else 0.0
wd = self.wd if t.ndim >= 3 else 0.0
up = up.shard_like(t) + self.lr.to(t.device) * wd * t.detach()
return t.detach() - up.cast(t.dtype)
examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama405b/implementations/tinybox_8xMI350X/dev_beam.sh
+1
View File
@@ -5,6 +5,7 @@ export DEV=${DEV:-AMD}
export EMULATE="AMD_CDNA4"
export CHECK_OOB=0
export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
export DEVICE_IN_FUNCTION_BUG=1
export DEBUG=${DEBUG:-2}
export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
extra/export_model.py
+3 -3
View File
@@ -1,11 +1,11 @@
from typing import Tuple, Dict, List, Optional
from tinygrad.dtype import DType
from tinygrad.dtype import DType, dtypes
from tinygrad.renderer import ProgramSpec
from tinygrad.tensor import Device, Tensor
from tinygrad.tensor import Tensor
from tinygrad.device import Device
from tinygrad.engine.jit import TinyJit
from tinygrad.nn.state import get_state_dict
from tinygrad.helpers import Context, to_mv
from tinygrad.dtype import dtypes
from tinygrad.uop.ops import Ops
import json
from collections import OrderedDict
extra/gemm/amd_asm_matmul.py
+6 -5
View File
@@ -291,7 +291,7 @@ def build_kernel(N, arch='gfx1100'):
# MAIN GEMM LOOP
# ===========================================================================
NO_DS, NO_GLOBAL = getenv("NO_DS", 0), getenv("NO_GLOBAL", 0)
NO_ALU, NO_DS, NO_GLOBAL = getenv("NO_ALU", 0), getenv("NO_DS", 0), getenv("NO_GLOBAL", 0)
k.label('LOOP_INC')
k.emit(s_add_i32(s[S_LOOP_CTR], s[S_LOOP_CTR], 8))
@@ -350,10 +350,11 @@ def build_kernel(N, arch='gfx1100'):
# 64 dual FMACs
k.waitcnt(lgkm=0)
k.emit(s_clause(simm16=len(FMAC_PATTERN)-1))
for i, (vdst_x, vdst_y, ax, bx, ay, by) in enumerate(FMAC_PATTERN):
k.emit(VOPD(VOPDOp.V_DUAL_FMAC_F32, VOPDOp.V_DUAL_FMAC_F32,
vdstx=v[vdst_x], vdsty=v[vdst_y], srcx0=v[ax], vsrcx1=v[bx], srcy0=v[ay], vsrcy1=v[by]))
if not NO_ALU:
k.emit(s_clause(simm16=len(FMAC_PATTERN)-1))
for i, (vdst_x, vdst_y, ax, bx, ay, by) in enumerate(FMAC_PATTERN):
k.emit(VOPD(VOPDOp.V_DUAL_FMAC_F32, VOPDOp.V_DUAL_FMAC_F32,
vdstx=v[vdst_x], vdsty=v[vdst_y], srcx0=v[ax], vsrcx1=v[bx], srcy0=v[ay], vsrcy1=v[by]))
# wait for all global loads to finish
# then sync the warp so it's safe to store local
extra/gemm/amd_copy_matmul.py
+75 -49
View File
@@ -7,73 +7,99 @@ M = getenv("M", N)
K = getenv("K", N)
WARP_SIZE = 32
BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 8
TM, TN = 4, 4
LANES_PER_WAVE_M, LANES_PER_WAVE_N = 4, 8
BLOCK_M, BLOCK_N = 128, 128
BLOCK_K = getenv("BK", 16)
assert N % BLOCK_N == 0 and M % BLOCK_M == 0 and K % BLOCK_K == 0
is_kernel5 = getenv("K5", 0)
THREADS_PER_BLOCK = 128 if is_kernel5 else 256
WAVES_PER_BLOCK_N = 1 if is_kernel5 else 2
WAVES_PER_BLOCK_M = THREADS_PER_BLOCK // WARP_SIZE // WAVES_PER_BLOCK_N
REG_TILES_PER_WAVE_N = BLOCK_N // (WAVES_PER_BLOCK_N * LANES_PER_WAVE_N * TN)
REG_TILES_PER_WAVE_M = BLOCK_M // (WAVES_PER_BLOCK_M * LANES_PER_WAVE_M * TM)
use_wmma = getenv("WMMA")
if use_wmma:
WAVES_M, WAVES_N = 2, 2
LANES_PER_WAVE_M, LANES_PER_WAVE_N = 2, 16
UNROLL_M, UNROLL_N = 1, 1
assert WAVES_PER_BLOCK_M*REG_TILES_PER_WAVE_M*LANES_PER_WAVE_M*TM == BLOCK_M, "M reshape is wrong"
assert WAVES_PER_BLOCK_N*REG_TILES_PER_WAVE_N*LANES_PER_WAVE_N*TN == BLOCK_N, "N reshape is wrong"
# wmma params
WMMA_M, WMMA_N, WMMA_K = 16, 16, 16
WMMA_ACC = WMMA_M // LANES_PER_WAVE_M
else:
WAVES_M, WAVES_N = 4, 1
LANES_PER_WAVE_M, LANES_PER_WAVE_N = 4, 8
UNROLL_M, UNROLL_N = 4, 4
consts = {"wpb_m":WAVES_PER_BLOCK_M, "lpw_m":LANES_PER_WAVE_M, "rt_m":REG_TILES_PER_WAVE_M, "t_m": TM,
"wpb_n":WAVES_PER_BLOCK_N, "lpw_n":LANES_PER_WAVE_N, "rt_n":REG_TILES_PER_WAVE_N, "t_n": TN}
# WARP_SIZE * total waves
THREADS_PER_BLOCK = WARP_SIZE * WAVES_M * WAVES_N
# accumulator size
TM = BLOCK_M // (WAVES_M * LANES_PER_WAVE_M)
TN = BLOCK_N // (WAVES_N * LANES_PER_WAVE_N)
# 128x128 out, kx128, kx128 in
def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
tid = UOp.range(THREADS_PER_BLOCK, 2, AxisType.LOCAL)
wave_m = UOp.range(WAVES_M, 2, AxisType.LOCAL)
wave_n = UOp.range(WAVES_N, 3, AxisType.LOCAL)
lane = UOp.range(WARP_SIZE, -1, AxisType.WARP)
tid = (wave_m * WAVES_N + wave_n) * WARP_SIZE + lane
#tid = UOp.special(THREADS_PER_BLOCK, "lidx0")
warp, lane = tid // WARP_SIZE, tid % WARP_SIZE
wave_n, wave_m = warp % WAVES_PER_BLOCK_N, warp // WAVES_PER_BLOCK_N
lane_n, lane_m = lane % LANES_PER_WAVE_N, lane // LANES_PER_WAVE_N
# -- GLOBAL -> LOCAL --
# wmma: spatial outer, k inner (k contiguous for vectorized WMMA tile loads)
# gemm: k outer, spatial inner
A_local = UOp.placeholder((BLOCK_M, BLOCK_K) if use_wmma else (BLOCK_K, BLOCK_M), a.dtype.base, slot=0, addrspace=AddrSpace.LOCAL)
B_local = UOp.placeholder((BLOCK_N, BLOCK_K) if use_wmma else (BLOCK_K, BLOCK_N), b.dtype.base, slot=1, addrspace=AddrSpace.LOCAL)
# define locals
A_local = UOp.placeholder((BLOCK_K, BLOCK_M), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL)
B_local = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
# open the main reduction range and copy in GLOBAL -> LOCAL
a = a.reshape(K // BLOCK_K, BLOCK_K, BLOCK_M)
b = b.reshape(K // BLOCK_K, BLOCK_K, BLOCK_N)
k_tile_range = UOp.range(K // BLOCK_K, 3, AxisType.REDUCE)
A_store = A_local.reshape(-1, THREADS_PER_BLOCK)[:, tid].store(a[k_tile_range].reshape(-1, THREADS_PER_BLOCK)[:, tid])
B_store = B_local.reshape(-1, THREADS_PER_BLOCK)[:, tid].store(b[k_tile_range].reshape(-1, THREADS_PER_BLOCK)[:, tid])
k_tile = UOp.range(K // BLOCK_K, 100, AxisType.REDUCE)
# copy with transpose for wmma (input is k×spatial, LDS is spatial×k)
A_copy = A_local.permute((1,0)) if use_wmma else A_local
B_copy = B_local.permute((1,0)) if use_wmma else B_local
A_store = A_copy.reshape(-1, THREADS_PER_BLOCK)[:, tid].store(a[k_tile].reshape(-1, THREADS_PER_BLOCK)[:, tid])
B_store = B_copy.reshape(-1, THREADS_PER_BLOCK)[:, tid].store(b[k_tile].reshape(-1, THREADS_PER_BLOCK)[:, tid])
barrier = UOp.barrier(A_store, B_store)
A_local, B_local = A_local.after(barrier), B_local.after(barrier)
# define accumulator (128x128), but broadcast across tid
c_regs = UOp.placeholder((REG_TILES_PER_WAVE_M*TM, REG_TILES_PER_WAVE_N*TN), dtypes.float, slot=2, addrspace=AddrSpace.REG)
c_regs = c_regs.after(c_regs.store(UOp.const(dtypes.float, 0).reshape((1,)*len(c_regs.shape)).expand(c_regs.shape)))
# -- COMPUTE --
lane_m, lane_n = lane // LANES_PER_WAVE_N, lane % LANES_PER_WAVE_N
# define registers (NOTE: the thread count is the device count for this multi, it's sharded across the THREADS_PER_BLOCK)
A_col = UOp.placeholder((REG_TILES_PER_WAVE_M, TM), dtypes.float, slot=0, addrspace=AddrSpace.REG)
B_row = UOp.placeholder((REG_TILES_PER_WAVE_N, TN), dtypes.float, slot=1, addrspace=AddrSpace.REG)
# accumulator (unified: both paths use (TM, TN) with scalar dtypes.float)
acc = UOp.placeholder((TM, TN), dtypes.float, slot=2, addrspace=AddrSpace.REG)
acc = acc.after(acc.store(UOp.const(dtypes.float, 0).reshape((1,)*len(acc.shape)).expand(acc.shape)))
# LOCAL -> REGS
k = UOp.range(BLOCK_K, 4, AxisType.REDUCE)
A_col = A_col.after(A_col.store(A_local[k].reshape(WAVES_PER_BLOCK_M, REG_TILES_PER_WAVE_M, LANES_PER_WAVE_M, TM)[wave_m, :, lane_m, :]))
B_row = B_row.after(B_row.store(B_local[k].reshape(WAVES_PER_BLOCK_N, REG_TILES_PER_WAVE_N, LANES_PER_WAVE_N, TN)[wave_n, :, lane_n, :]))
if use_wmma:
k = UOp.range(BLOCK_K // WMMA_K, 101, AxisType.REDUCE)
tile_m = UOp.range(TM // WMMA_ACC, 200, AxisType.LOOP)
tile_n = UOp.range(TN, 201, AxisType.LOOP)
# do FMA
A_col = A_col.reshape(REG_TILES_PER_WAVE_M*TM, 1).expand(REG_TILES_PER_WAVE_M*TM, REG_TILES_PER_WAVE_N*TN)
B_row = B_row.reshape(1, REG_TILES_PER_WAVE_N*TN).expand(REG_TILES_PER_WAVE_M*TM, REG_TILES_PER_WAVE_N*TN)
c_regs = c_regs.after(c_regs.store(c_regs.after(k) + (A_col * B_row)).end(k).barrier().end(k_tile_range))
acc_frag = acc.reshape(TM // WMMA_ACC, WMMA_ACC, TN).permute(0,2,1)[tile_m, tile_n]
a_frag = A_local.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, BLOCK_K // WMMA_K, WMMA_K)[wave_m, tile_m, lane_n, k]
b_frag = B_local.reshape(WAVES_N, TN, WMMA_N, BLOCK_K // WMMA_K, WMMA_K)[wave_n, tile_n, lane_n, k]
# store back to c
c_store = c.rearrange("(wpb_m rt_m lpw_m t_m) (wpb_n rt_n lpw_n t_n) -> (wpb_m wpb_n lpw_m lpw_n) (rt_m t_m) (rt_n t_n)", **consts)
return c_store[tid].store(c_regs).end(tid)
wmma = UOp(Ops.SHAPED_WMMA, dtypes.float, (a_frag, b_frag, acc_frag.after(k)), arg=((16, 16, 16), 'AMD', 32))
acc_store = acc_frag.store(wmma).end(tile_m, tile_n)
else:
# registers for LOCAL -> REG
a_frag = UOp.placeholder((TM//UNROLL_M, UNROLL_M), dtypes.float, slot=0, addrspace=AddrSpace.REG)
b_frag = UOp.placeholder((TN//UNROLL_N, UNROLL_N), dtypes.float, slot=1, addrspace=AddrSpace.REG)
k = UOp.range(BLOCK_K, 101, AxisType.REDUCE)
a_frag = a_frag.after(a_frag.store(A_local[k].reshape(WAVES_M, TM//UNROLL_M, LANES_PER_WAVE_M, UNROLL_M)[wave_m, :, lane_m, :]))
b_frag = b_frag.after(b_frag.store(B_local[k].reshape(WAVES_N, TN//UNROLL_N, LANES_PER_WAVE_N, UNROLL_N)[wave_n, :, lane_n, :]))
# FMA
a_frag = a_frag.reshape(TM, 1).expand(TM, TN)
b_frag = b_frag.reshape(1, TN).expand(TM, TN)
acc_store = acc.store(acc.after(k) + (a_frag * b_frag))
# store accumulator and loop
acc = acc.after(acc_store.end(k).barrier().end(k_tile))
# store accumulator to output (unified)
c = c.reshape(WAVES_M, TM//UNROLL_M, LANES_PER_WAVE_M, UNROLL_M,
WAVES_N, TN//UNROLL_N, LANES_PER_WAVE_N, UNROLL_N)
c = c.permute((0,4,2,6, 1,3,5,7)).reshape(THREADS_PER_BLOCK, TM, TN)
return c[tid].store(acc).end(wave_m, wave_n, lane)
def amd_copy_matmul(c:UOp, a:UOp, b:UOp) -> UOp:
block_id_n = UOp.range(N // BLOCK_N, 0, AxisType.GLOBAL)
block_id_m = UOp.range(M // BLOCK_M, 1, AxisType.GLOBAL)
# index the output with the globals
block_id_m = UOp.range(M // BLOCK_M, 0, AxisType.GLOBAL)
block_id_n = UOp.range(N // BLOCK_N, 1, AxisType.GLOBAL)
c = c.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_N, BLOCK_N)[block_id_m, :, block_id_n, :]
a = a.T.reshape(K, M // BLOCK_M, BLOCK_M)[:, block_id_m, :]
b = b.reshape(K, N // BLOCK_N, BLOCK_N)[:, block_id_n, :]
@@ -81,4 +107,4 @@ def amd_copy_matmul(c:UOp, a:UOp, b:UOp) -> UOp:
if __name__ == "__main__":
from amd_uop_matmul import eval_custom_matmul
eval_custom_matmul(amd_copy_matmul)
eval_custom_matmul(amd_copy_matmul, dtypes.half if use_wmma else dtypes.float)
extra/gemm/amd_uop_matmul.py
+6 -6
View File
@@ -115,14 +115,14 @@ def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
return sink.sink(arg=KernelInfo(opts_to_apply=())).simplify()
def eval_custom_matmul(fxn):
a = Tensor.randn(M, K, dtype=dtypes.float)
b = Tensor.randn(K, N, dtype=dtypes.float)
def eval_custom_matmul(fxn, dt=dtypes.float):
a = Tensor.randn(M, K, dtype=dt)
b = Tensor.randn(K, N, dtype=dt)
c = Tensor.empty(M, N, dtype=dtypes.float)
with Context(DEBUG=0): Tensor.realize(a, b)
ets = []
with Context(DEBUG=max(2, DEBUG.value)):
with Context(DEBUG=max(2, DEBUG.value), DEVECTORIZE=2 if dt == dtypes.half else 0):
for _ in range(NUM_RUNS):
GlobalCounters.reset()
tst = Tensor.custom_kernel(c, a, b, fxn=fxn)[0].realize()
@@ -132,11 +132,11 @@ def eval_custom_matmul(fxn):
if getenv("VERIFY", 1):
GlobalCounters.reset()
with Context(DEBUG=2):
tc = (a @ b).realize()
tc = (a.float() @ b.float()).realize()
with Context(DEBUG=0):
err = (tc - tst).square().mean().item()
print(f"mean squared error {err}")
if err > 1e-06:
if err > (1e-2 if dt == dtypes.half else 1e-6):
raise RuntimeError("matmul is wrong!")
if __name__ == "__main__":
extra/gemm/cdna_asm_gemm.py
+6 -3
View File
@@ -2649,6 +2649,9 @@ def can_use_asm_gemm(a:Tensor, b:Tensor) -> bool:
else: dname = a.device
arch = getattr(Device[dname].renderer, "arch", "")
if batch not in {1, 2}: return todo(f"GEMM batch size {batch}")
# blacklist slow matmul
# TODO: why is this slow?
if (M,N,K) == (8192, 2304, 16384): return todo("blacklisted slow matmul")
if (M % TILE_M != 0 or N % TILE_N != 0 or K % TILE_K != 0) and arch == "gfx950":
return todo(f"GEMM shape ({M},{N},{K}) not a multiple of ({TILE_M},{TILE_N},{TILE_K})")
return True
@@ -2663,10 +2666,10 @@ def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
m = UOp.range(M, 1, AxisType.LOOP)
n = UOp.range(N, 2, AxisType.LOOP)
k = UOp.range(K, 0, AxisType.REDUCE)
mul = (A.flatten().index((m*UOp.const(dtypes.index, K)+k))*
B.flatten().index((k*UOp.const(dtypes.index, N)+n))).cast(dtypes.float32)
mul = (A.flatten().index((m*UOp.const(dtypes.weakint, K)+k))*
B.flatten().index((k*UOp.const(dtypes.weakint, N)+n))).cast(dtypes.float32)
red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype.base)
store = C.flatten().index((m*UOp.const(dtypes.index, N)+n), ptr=True).store(red).end(m, n)
store = C.flatten().index((m*UOp.const(dtypes.weakint, N)+n), ptr=True).store(red).end(m, n)
return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))
# ** backward gemm, might use the asm gemm
extra/setup_hipcc_osx.sh
Executable
+31
View File
@@ -0,0 +1,31 @@
#!/bin/sh
install_loc="$HOME/.local/bin"
docker build --platform=linux/amd64 -t rocm-hipcc:7.2 - <<'EOF'
FROM ubuntu:22.04
ENV DEBIAN_FRONTEND=noninteractive
ENV TZ=Etc/UTC
RUN apt-get update && apt-get install -y --no-install-recommends wget ca-certificates gnupg tzdata && \
wget https://repo.radeon.com/amdgpu-install/7.2/ubuntu/jammy/amdgpu-install_7.2.70200-1_all.deb && \
apt-get install -y ./amdgpu-install_7.2.70200-1_all.deb && \
amdgpu-install -y --usecase=rocm --no-dkms --no-32 && \
rm -rf /var/lib/apt/lists/*
ENV PATH=/opt/rocm/bin:$PATH
EOF
mkdir -p "$install_loc"
tee "$install_loc/hipccshim" >/dev/null <<'EOF'
#!/bin/sh
set -eu
cname="rocm-hipcc-persistent"
if ! docker inspect --format='{{.State.Running}}' "$cname" 2>/dev/null | grep -q true; then
docker rm -f "$cname" 2>/dev/null || true
docker run -d --platform=linux/amd64 --name "$cname" \
-v /var/folders:/var/folders -v "$HOME":"$HOME" \
rocm-hipcc:7.2 sleep 300 >/dev/null
fi
exec docker exec "$cname" "$(basename "$0")" "$@"
EOF
chmod +x "$install_loc/hipccshim"
for t in hipcc hipconfig; do
ln -sf "$install_loc/hipccshim" "$install_loc/$t"
done
test/backend/test_dtype.py
+7 -13
View File
@@ -18,11 +18,10 @@ settings.register_profile("my_profile", max_examples=200, deadline=None, derando
settings.load_profile("my_profile")
def get_available_cast_dtypes(dtype: DType) -> List[DType]:
# dont cast internal dtypes
dts = [v for k, v in DTYPES_DICT.items() if v != dtype and is_dtype_supported(v) and not k.startswith("_")]
if not is_dtype_supported(dtype) or dtypes.long in EMULATED_DTYPES.tolist(dtypes):
if dtype in (dtypes.long, dtypes.ulong): return [dt for dt in dts if dt != dtypes.double] # can't bitcast with no 64-bit support
else: return []
dts = [v for k, v in DTYPES_DICT.items() if v != dtype and is_dtype_supported(v) or v in dtypes.fp8s+(dtypes.half,dtypes.bfloat16,dtypes.long)]
if dtype in (dtypes.long, dtypes.ulong) and (not is_dtype_supported(dtype) or dtypes.long in EMULATED_DTYPES.tolist(dtypes)):
return [dt for dt in dts if dt != dtypes.double] # can't bitcast with no 64-bit support
if not is_dtype_supported(dtype) and dtype not in dtypes.fp8s+(dtypes.half,dtypes.bfloat16): return []
return dts
def _to_torch_storage_type(dtype:DType):
@@ -60,10 +59,8 @@ class TestDType(unittest.TestCase):
DATA: Any = None
@classmethod
def setUpClass(cls):
if not cls.DTYPE or not is_dtype_supported(cls.DTYPE): raise unittest.SkipTest("dtype not supported")
cls.DATA = rand_for_dtype(cls.DTYPE, 10)
def setUp(self):
if self.DTYPE is None: raise unittest.SkipTest("base class")
if cls.DTYPE is None: raise unittest.SkipTest("base class")
cls.DATA = rand_for_dtype(cls.DTYPE, 0x10, allow_subnormal=is_dtype_supported(cls.DTYPE))
def test_to_np(self):
_test_to_np(Tensor(self.DATA, dtype=self.DTYPE), _to_np_dtype(self.DTYPE), np.array(self.DATA, dtype=_to_np_dtype(self.DTYPE)))
@@ -132,7 +129,6 @@ class TestDType(unittest.TestCase):
def _test_ops(a_dtype:DType, b_dtype:DType, target_dtype=None):
target_dtype = target_dtype or least_upper_dtype(a_dtype, b_dtype)
if not is_dtype_supported(a_dtype) or not is_dtype_supported(b_dtype) or not is_dtype_supported(target_dtype): return
if a_dtype == dtypes.bool or b_dtype == dtypes.bool: return
_assert_eq(Tensor([1,2,3,4], dtype=a_dtype)+Tensor([1,2,3,4], dtype=b_dtype), target_dtype, [2,4,6,8])
_assert_eq((Tensor([1], dtype=a_dtype).cast(b_dtype)+Tensor([1], dtype=a_dtype).cast(b_dtype)).cast(a_dtype), a_dtype, [2])
@@ -195,7 +191,6 @@ class TestFp8sConversions(unittest.TestCase):
def test_fp8e5m2fnuz_to_float(self, x):
np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e5m2fnuz), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e5m2fnuz).float().item())
@unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), "bfloat16 not supported")
class TestBFloat16(unittest.TestCase):
def test_bf16_creation_numpy(self):
data = [-1, 1, 2]
@@ -215,7 +210,6 @@ class TestBFloat16(unittest.TestCase):
assert t.dtype == dtypes.bfloat16
np.testing.assert_allclose(t.numpy(), np.eye(3))
@unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), "bfloat16 not supported")
class TestBFloat16DType(unittest.TestCase):
def test_bf16_to_float(self):
_test_cast(Tensor([100000], dtype=dtypes.bfloat16), dtypes.float32)
@@ -229,7 +223,6 @@ class TestBFloat16DType(unittest.TestCase):
back = t.cast(dtypes.float32)
assert tuple(back.numpy().tolist()) == (9984., -1, -1000, -9984, 20)
@unittest.skipUnless(is_dtype_supported(dtypes.bfloat16) and is_dtype_supported(dtypes.float16), "bfloat16 or float16 not supported")
class TestBFloat16DTypeCast(unittest.TestCase):
def test_f16_to_bf16_conversion(self):
original_tensor = Tensor([1.0, 2.0, 3.0], dtype=dtypes.float16)
@@ -278,6 +271,7 @@ class TestFloatDType(TestDType):
_test_op(lambda: Tensor([-0.9, -0.3, 1.2], dtype=dtypes.float32).cast(dtypes.uint32), dtypes.uint32,
[0, 0, 1])
@unittest.skipUnless(is_dtype_supported(dtypes.double), f"no double on {Device.DEFAULT}")
class TestDoubleDType(TestDType):
DTYPE = dtypes.double
@unittest.skipIf((CI and Device.DEFAULT in {"CUDA", "NV"}) or \
test/backend/test_dtype_alu.py
+2 -2
View File
@@ -81,8 +81,8 @@ def universal_test_unary(a, dtype, op):
if op[0] == Tensor.cos and abs(a) > 30: return
if op[0] == Tensor.log and a <= 0: return
if dtype in dtypes.fp8s:
# normals are zero
if dtype in EMULATED_DTYPES.tolist(dtypes) and abs(ta.numpy().item()) < 0.015625: return
# denormals are zero
if dtype in EMULATED_DTYPES.tolist(dtypes) or not is_dtype_supported(dtype) and abs(ta.numpy().item()) < 0.015625: return
tensor_value = fp8_to_float(op[0](ta.realize()).bitcast(dtypes.uint8).item(), dtype)
numpy_value = truncate[dtype](v:=op[1](ta.numpy()).item())
# cuda cast f32 inf to f8 MAX, amd cast it to nan(E4M3)/inf(E5M2)
test/backend/test_image_dtype.py
+2 -205
View File
@@ -1,186 +1,17 @@
import unittest
import numpy as np
from tinygrad import Device, dtypes, Tensor, Context
from tinygrad.device import LRUAllocator, is_dtype_supported
from tinygrad.dtype import ImageDType
from tinygrad.helpers import prod, unwrap
from tinygrad.helpers import unwrap
IMAGE_SUPPORTED_DEVICES = ("QCOM", "CL")
@unittest.skipUnless(Device.DEFAULT in IMAGE_SUPPORTED_DEVICES, "Images not supported")
class TestImageCopy(unittest.TestCase):
def test_image_copyout_1x8(self, img_type=dtypes.imagef):
it = Tensor.arange(32).cast(img_type((1,8,4))).realize()
buf = it.uop.buffer
out = buf.as_memoryview()
np.testing.assert_equal(out.cast(it.dtype.fmt).tolist(), np.arange(32))
@unittest.skipUnless(is_dtype_supported(dtypes.half, device="PYTHON"), "need half")
def test_imageh_copyout_1x8(self): self.test_image_copyout_1x8(img_type=dtypes.imageh)
def test_image_numpy_1x8(self, img_type=dtypes.imagef):
it = Tensor.arange(32).cast(img_type((1,8,4))).realize()
np.testing.assert_equal(it.numpy(), np.arange(32))
def test_imageh_numpy_1x8(self): self.test_image_numpy_1x8(img_type=dtypes.imageh)
def test_image_copyout_2x4(self):
it = Tensor.arange(2*4*4).cast(dtypes.imagef((2,4,4))).realize()
buf = it.uop.buffer
out = buf.as_memoryview()
np.testing.assert_equal(out.cast('f').tolist(), np.arange(2*4*4))
def test_image_roundtrip(self):
sz = (4,2,4)
it = Tensor.rand(prod(sz)).cast(dtypes.imagef(sz)).realize()
buf = it.uop.buffer
out = buf.as_memoryview()
it2 = Tensor.rand(prod(sz)).cast(dtypes.imagef(sz)).realize()
buf2 = it2.uop.buffer
buf2.copyin(out)
assert (it == it2).sum().item() == prod(sz)
@unittest.skipUnless(Device.DEFAULT in IMAGE_SUPPORTED_DEVICES, "Images not supported")
class TestImageDType(unittest.TestCase):
def test_image_pitch(self):
def __validate(imgdt, expected_pitch):
assert imgdt.pitch == expected_pitch, f"Failed pitch for image: {imgdt}. Got 0x{imgdt.pitch:X}, expected 0x{expected_pitch:X}"
# Match opencl pitches for perf
__validate(dtypes.imageh((1, 201)), 0x680)
__validate(dtypes.imageh((16, 216)), 0x700)
__validate(dtypes.imageh((16, 9)), 0x80)
__validate(dtypes.imageh((48, 64)), 0x200)
__validate(dtypes.imageh((32, 128)), 0x400)
__validate(dtypes.imageh((96, 128)), 0x400)
__validate(dtypes.imageh((64, 256)), 0x840)
__validate(dtypes.imageh((64, 9)), 0x80)
__validate(dtypes.imageh((192, 256)), 0x840)
__validate(dtypes.imageh((64, 768)), 0x1840)
__validate(dtypes.imageh((256, 49)), 0x1C0)
__validate(dtypes.imageh((128, 9)), 0x80)
__validate(dtypes.imageh((16, 1024)), 0x2080)
__validate(dtypes.imageh((64, 512)), 0x1040)
__validate(dtypes.imageh((16, 512)), 0x1080)
__validate(dtypes.imageh((132, 64)), 0x200)
__validate(dtypes.imageh((4, 512)), 0x1200)
__validate(dtypes.imageh((8, 512)), 0x1100)
__validate(dtypes.imageh((128, 128)), 0x400)
__validate(dtypes.imageh((32, 512)), 0x1040)
__validate(dtypes.imageh((26, 64)), 0x200)
__validate(dtypes.imageh((32, 516)), 0x1040)
__validate(dtypes.imageh((32, 1024)), 0x2040)
__validate(dtypes.imageh((16, 2048)), 0x4080)
__validate(dtypes.imageh((8, 2048)), 0x4100)
__validate(dtypes.imageh((4, 4096)), 0x8200)
__validate(dtypes.imagef((16, 49)), 0x380)
__validate(dtypes.imagef((16, 1024)), 0x4080)
__validate(dtypes.imagef((256, 64)), 0x400)
__validate(dtypes.imagef((64, 512)), 0x2040)
__validate(dtypes.imagef((16, 512)), 0x2080)
__validate(dtypes.imagef((132, 64)), 0x400)
__validate(dtypes.imagef((4, 512)), 0x2200)
__validate(dtypes.imagef((4, 16)), 0x200)
__validate(dtypes.imagef((2, 16)), 0x400)
__validate(dtypes.imagef((8, 512)), 0x2100)
__validate(dtypes.imagef((12, 64)), 0x400)
__validate(dtypes.imagef((3, 32)), 0x400)
__validate(dtypes.imagef((128, 128)), 0x840)
__validate(dtypes.imagef((32, 512)), 0x2040)
__validate(dtypes.imagef((8, 3072)), 0xC100)
__validate(dtypes.imagef((4, 2048)), 0x8200)
__validate(dtypes.imagef((4, 1024)), 0x4200)
__validate(dtypes.imagef((4, 4096)), 0x10200)
__validate(dtypes.imagef((10, 384)), 0x1900)
__validate(dtypes.imagef((24, 64)), 0x400)
__validate(dtypes.imagef((128, 12)), 0xC0)
__validate(dtypes.imagef((10, 24)), 0x200)
__validate(dtypes.imagef((1, 129)), 0x840)
__validate(dtypes.imagef((1, 32)), 0x200)
__validate(dtypes.imagef((1, 64)), 0x400)
__validate(dtypes.imagef((1, 1239)), 0x4D80)
__validate(dtypes.imagef((1, 1)), 0x40)
def test_image_and_back(self):
data = Tensor.randn(9*32*4).realize()
tst = data.numpy()
it = data.cast(dtypes.imagef((9,32,4))).contiguous().realize()
assert isinstance(it.uop.base.realized.dtype, ImageDType)
np.testing.assert_equal(tst, it.numpy())
def test_image_cast_and_back_collapses(self):
data = Tensor.randn(9*27*4).realize()
tst = data.numpy()
it = data.cast(dtypes.imagef((9,27,4))).realize()
# the underlying UOp is identical
#self.assertIs(it.uop.base.realized, data.uop.base.realized)
np.testing.assert_equal(tst, it.numpy())
def test_image_and_back_wrong_shape(self):
data = Tensor.randn(9*27*4).realize()
tst = data.numpy()
it = data.cast(dtypes.imagef((9,12,4))).realize()
assert not isinstance(it.uop.base.realized.dtype, ImageDType)
np.testing.assert_equal(tst, it.numpy())
def test_shrink_load_float(self):
it = Tensor.randn(16).cast(dtypes.imagef((1,4,4))).realize()
imgv = it.numpy()
np.testing.assert_equal(imgv[0:2], it[0:2].numpy())
def test_mul_stays_image(self):
# NOTE: contiguous is needed otherwise this folds
it = Tensor.randn(16).cast(dtypes.imagef((1,4,4))).contiguous().realize()
out = (it*2).realize()
assert isinstance(out.uop.base.realized.dtype, ImageDType)
def test_sum(self):
it = Tensor.rand(8).cast(dtypes.imagef((1,2,4))).realize()
itn = it.numpy()
np.testing.assert_allclose(np.sum(itn), it.sum().numpy(), rtol=1e-6)
def test_shrink_max(self):
it = Tensor.randn(16).cast(dtypes.imagef((1,4,4))).realize()
imgv = it.numpy()
np.testing.assert_equal(np.maximum(imgv[0:3], 0), it[0:3].relu().numpy())
def test_shrink_to_float(self):
it = Tensor.randn(4, 4).cast(dtypes.imagef((1,4,4))).realize()
imgv = it.numpy()
np.testing.assert_equal(np.maximum(imgv[:, 0], 0), it[:, 0].relu().numpy())
@unittest.skipUnless(isinstance(Device.default.allocator, LRUAllocator), "Requires LRU")
def test_lru_alloc(self):
data = Tensor.randn(9*27*4).realize()
it = data.cast(dtypes.imagef((9,27,4))).realize()
b1 = it.uop.base.realized._buf
del it
it = data.cast(dtypes.imagef((9,27,4))).realize()
assert it.uop.base.realized._buf == b1
def test_no_lru_alloc(self):
data = Tensor.randn(9*32*4).realize()
it = data.cast(dtypes.imagef((9,32,4))).contiguous().realize()
b1 = it.uop.base.realized._buf
del it
it = data.reshape(9,32,4).pad_to(10, None, None).cast(dtypes.imagef((10,32,4))).contiguous().realize()
assert it.uop.base.realized._buf != b1
def test_no_lru_alloc_dtype(self):
data = Tensor.randn(9*32*4).realize()
it = data.cast(dtypes.imagef((9,32,4))).contiguous().realize()
b1 = it.uop.base.realized._buf
del it
it = data.cast(dtypes.imageh((9,32,4))).realize()
assert it.uop.base.realized._buf != b1
# issue caused by: don't realize image to image casts. this is part of a larger problem
#@unittest.expectedFailure
# update: passing after tensor_map
def test_lil_model(self):
with Context(IMAGE=2):
with Context(IMAGE=1):
x = Tensor.zeros(1, 1)
w1 = Tensor.zeros(1, 8, requires_grad=True)
w2 = Tensor.zeros(8, 2)
@@ -198,39 +29,5 @@ class TestImageDType(unittest.TestCase):
self.assertEqual(w1.grad.uop.base.buffer.dtype, dtypes.float32)
self.assertEqual(len(sched), 9)
@unittest.skipUnless(Device.DEFAULT in IMAGE_SUPPORTED_DEVICES, "Images not supported")
class TestImageRealization(unittest.TestCase):
def test_image_dtype_expand(self):
data = Tensor.randn(9*32*4).realize()
it = data.cast(dtypes.imagef((9,32,4))).contiguous().realize()
self.assertEqual(it.dtype, dtypes.imagef((9,32,4)))
it_expanded = it.reshape((9,32,4,1)).expand((9,32,4,4)).contiguous().realize()
self.assertEqual(it_expanded.dtype, dtypes.float32)
def test_image_dtype_expand_and_back(self):
data = Tensor.randn(9*32*4).realize()
it = data.cast(dtypes.imagef((9,32,4))).contiguous().realize()
self.assertEqual(it.dtype, dtypes.imagef((9,32,4)))
it_expanded = it.reshape((9,32,4,1)).expand((9,32,4,4))
it2 = it_expanded.sum(3).realize()
self.assertEqual(it2.dtype, dtypes.imagef((9,32,4)))
def test_image_alu_children(self):
data = Tensor.randn(9*32*4).realize()
it = data.cast(dtypes.imagef((9,32,4))).contiguous().realize()
self.assertEqual(it.dtype, dtypes.imagef((9,32,4)))
it_expanded = it.reshape((9,32,4,1)).expand((9,32,4,4)).contiguous()
alu1 = it_expanded+1
alu2 = it_expanded.sum(3)
it_expanded.realize()
# NOTE: the parent becomes float, but the alu child will stay image until its output cannot fit the image
self.assertEqual(alu1.dtype, dtypes.imagef((9,32,4)))
alu1.realize()
self.assertEqual(alu1.dtype, dtypes.float32)
# alu2 is back in image because it fits the dtype again
self.assertEqual(alu2.dtype, dtypes.imagef((9,32,4)))
alu2.realize()
self.assertEqual(alu2.dtype, dtypes.imagef((9,32,4)))
if __name__ == '__main__':
unittest.main()
test/backend/test_jit.py
+2 -2
View File
@@ -618,7 +618,7 @@ class TestJitFree(unittest.TestCase):
expected_savings = (len(inp) * inp.dtype.itemsize * 2) + dtypes.float32.itemsize # (t1 and t2) + out
self.assertEqual(savings_after_free, expected_savings)
self.assertGreaterEqual(savings_after_free, expected_savings)
out = fxn(Tensor([11,1,2,3,4]))
self.assertEqual(out.item(), 136)
@@ -628,7 +628,7 @@ class TestJitFree(unittest.TestCase):
fxn.captured.free_intermediates() # 2nd time to validate
savings_after_free = pre_free - GlobalCounters.mem_used
self.assertEqual(savings_after_free, expected_savings)
self.assertGreaterEqual(savings_after_free, expected_savings)
out = fxn(Tensor([11,1,2,3,4]))
self.assertEqual(out.item(), 136)
test/backend/test_linearizer_dumb.py
+9 -9
View File
@@ -12,18 +12,18 @@ class TestLinearizerFailure(unittest.TestCase):
@unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL")
def test_failure_beam_mnist(self):
c0 = UOp(Ops.PARAM, dtypes.uchar.ptr(4014080), arg=0, src=())
c1 = UOp.range(UOp.const(dtypes.index, 512), 0, AxisType.GLOBAL)
c2 = UOp.range(UOp.const(dtypes.index, 784), 1, AxisType.GLOBAL)
c3 = UOp.range(UOp.const(dtypes.index, 10), 3, AxisType.GLOBAL)
c1 = UOp.range(UOp.const(dtypes.weakint, 512), 0, AxisType.GLOBAL)
c2 = UOp.range(UOp.const(dtypes.weakint, 784), 1, AxisType.GLOBAL)
c3 = UOp.range(UOp.const(dtypes.weakint, 10), 3, AxisType.GLOBAL)
c4 = UOp(Ops.PARAM, dtypes.int.ptr(512), arg=1, src=())
c5 = c4.index(c1.valid(UOp.const(dtypes.bool, True)))
c6 = UOp.range(UOp.const(dtypes.index, 6000), 1004, AxisType.REDUCE)
c7 = UOp.range(UOp.const(dtypes.index, 3750), 2006, AxisType.REDUCE)
c8 = UOp.range(UOp.const(dtypes.index, 16), 2007, AxisType.GROUP_REDUCE)
c6 = UOp.range(UOp.const(dtypes.weakint, 6000), 1004, AxisType.REDUCE)
c7 = UOp.range(UOp.const(dtypes.weakint, 3750), 2006, AxisType.REDUCE)
c8 = UOp.range(UOp.const(dtypes.weakint, 16), 2007, AxisType.GROUP_REDUCE)
c9 = UOp(Ops.PARAM, dtypes.uchar.ptr(47040000), arg=2, src=())
c10 = c9.index((((c3*UOp.const(dtypes.index, 4704000))+c2)+(c6*UOp.const(dtypes.index, 784))).valid(UOp.const(dtypes.bool, True)))
c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(dtypes.index, 6000))+c6)+((c7*UOp.const(dtypes.index, 16))+c8)).alu(Ops.CMPLT, UOp.const(dtypes.index, 59999)).where(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(dtypes.int, -1))).where(UOp.const(dtypes.uchar, 0), c10).reduce(c6, arg=Ops.ADD)
c12 = c0.index((((c1*UOp.const(dtypes.index, 7840))+(c2*UOp.const(dtypes.index, 10)))+c3).valid(UOp.const(dtypes.bool, True))).store(c11).end(c1, c2, c3)
c10 = c9.index((((c3*UOp.const(dtypes.weakint, 4704000))+c2)+(c6*UOp.const(dtypes.weakint, 784))).valid(UOp.const(dtypes.bool, True)))
c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(dtypes.weakint, 6000))+c6)+((c7*UOp.const(dtypes.weakint, 16))+c8)).alu(Ops.CMPLT, UOp.const(dtypes.weakint, 59999)).where(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(dtypes.int, -1))).where(UOp.const(dtypes.uchar, 0), c10).reduce(c6, arg=Ops.ADD)
c12 = c0.index((((c1*UOp.const(dtypes.weakint, 7840))+(c2*UOp.const(dtypes.weakint, 10)))+c3).valid(UOp.const(dtypes.bool, True))).store(c11).end(c1, c2, c3)
ast = c12.sink(arg=KernelInfo(name='test', axis_types=(), dont_use_locals=False, applied_opts=(Opt(op=OptOps.GROUP, axis=1, arg=16),), opts_to_apply=None))
_ = get_program(ast, Device["METAL"].renderer)
test/backend/test_schedule.py
+7 -28
View File
@@ -9,7 +9,7 @@ from hypothesis import assume, given, settings, strategies as strat
from tinygrad import nn, dtypes, Device, Tensor, Variable
from tinygrad.device import is_dtype_supported
from tinygrad.dtype import DType, ImageDType
from tinygrad.dtype import DType
from tinygrad.uop.ops import UOp, Ops, UPat
from tinygrad.helpers import CI, DEBUG, OSX, GlobalCounters, Context, getenv, all_same, temp
from tinygrad.engine.realize import CompiledRunner, run_schedule
@@ -783,18 +783,6 @@ class TestSchedule(unittest.TestCase):
gc.collect()
self.assertEqual(GlobalCounters.mem_used-base, 1024)
@unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
def test_image_matmul(self):
with Context(IMAGE=2):
x = Tensor.randn((9, 9)).realize()
y = Tensor.randn((9, 9)).realize()
out = x@y
run_schedule(check_schedule(out, 3))
np.testing.assert_allclose(out.numpy(), x.numpy()@y.numpy(), atol=1e-4, rtol=1e-4)
self.assertIsInstance(out.dtype, ImageDType)
self.assertIsNotNone(out.uop.base.realized)
self.assertIsInstance(out.uop.base.realized.dtype, ImageDType)
@unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
def test_image_dot_f16_fusion(self):
with Context(FLOAT16=1, OPENPILOT_HACKS=1):
@@ -805,11 +793,8 @@ class TestSchedule(unittest.TestCase):
for si in sched: si.lower()
return len([si for si in sched if isinstance(si.prg, CompiledRunner)])
with Context(IMAGE=1): cnt1 = cnt()
with Context(IMAGE=2): cnt2 = cnt()
self.assertEqual(cnt1, 5)
self.assertEqual(cnt2, 5)
with Context(IMAGE=1):
self.assertEqual(cnt(), 5)
@unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
def test_image_f16_residual_fusion(self):
@@ -825,14 +810,10 @@ class TestSchedule(unittest.TestCase):
for si in sched: si.lower()
return len([si for si in sched if isinstance(si.prg, CompiledRunner)])
with Context(IMAGE=1): cnt1 = cnt()
with Context(IMAGE=2): cnt2 = cnt()
self.assertEqual(cnt1, 9)
self.assertEqual(cnt2, 9)
with Context(IMAGE=1):
self.assertEqual(cnt(), 9)
@unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
@unittest.expectedFailure
def test_image_conv_fusion(self):
with Context(OPENPILOT_HACKS=1):
def cnt():
@@ -843,10 +824,8 @@ class TestSchedule(unittest.TestCase):
for si in sched: si.lower()
return len([si for si in sched if isinstance(si.prg, CompiledRunner)])
with Context(IMAGE=1): cnt1 = cnt()
with Context(IMAGE=2): cnt2 = cnt()
self.assertEqual(cnt1, cnt2)
with Context(IMAGE=1):
self.assertEqual(cnt(), 5)
def _test_fusion(self, shapes, f, cnt):
with Context(DEBUG=0, TRACK_MATCH_STATS=0): args = [Tensor.randn(s).realize() for s in shapes]
test/external/external_benchmark_op_conv.py
Vendored
+2 -2
View File
@@ -23,7 +23,7 @@ def vision_conv_143():
c32 = ((c27<3)!=True)&(c27<67)
c34 = UOp(Ops.PARAM, dtypes.imageh((32, 1024, 4)), (), 1)
c38 = c5//2
c45 = (c32&c24).where((c27*64+c38+c17*4096+-12480), UOp.const(dtypes.index, Invalid))
c45 = (c32&c24).where((c27*64+c38+c17*4096+-12480), UOp.const(dtypes.weakint, Invalid))
c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
c49 = UOp(Ops.PARAM, dtypes.imageh((64, 49, 4)), (), 2)
c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
@@ -49,7 +49,7 @@ def vision_conv_153():
c32 = ((c27<3)!=True)&(c27<35)
c34 = UOp(Ops.PARAM, dtypes.imageh((16, 1024, 4)), (), 1)
c38 = c5//2
c45 = (c32&c24).where((c27*128+c38+c17*4096+-12672), UOp.const(dtypes.index, Invalid))
c45 = (c32&c24).where((c27*128+c38+c17*4096+-12672), UOp.const(dtypes.weakint, Invalid))
c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
c49 = UOp(Ops.PARAM, dtypes.imageh((128, 49, 4)), (), 2)
c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
test/external/external_test_opt.py
Vendored
+6 -2
View File
@@ -6,14 +6,18 @@ import torch
from tinygrad import GlobalCounters, Tensor, Device
from tinygrad.helpers import getenv
from tinygrad.nn.state import get_parameters
from tinygrad.engine.realize import capturing
from tinygrad.engine.realize import capturing, run_schedule
from tinygrad.engine.schedule import linear_to_schedule
from tinygrad.tensor import _to_np_dtype
class CLCache:
def __init__(self, allowed=None, strict=False, preclear=True, var_vals=None):
self.allowed, self.strict, self.preclear, self.var_vals = allowed, strict, preclear, var_vals if var_vals is not None else {}
self.count = 0
def add(self, ei): self.count += 1
def add_linear(self, linear, var_vals):
schedule = linear_to_schedule(linear)
self.count += len(schedule)
run_schedule(schedule, var_vals)
def __enter__(self):
if self.preclear:
gc.collect()
test/mockgpu/amd/emu.py
+28 -23
View File
@@ -337,7 +337,7 @@ def get_pcode(op) -> str:
pcode = pcode.replace('VCC = 0x0LL', 'VCC.u64[laneId] = 0').replace('VCC = 0x1LL', 'VCC.u64[laneId] = 1')
return pcode
def parse_pcode(pcode: str, srcs: dict[str, UOp] | None = None) -> tuple[dict, list[tuple[str, UOp]]]:
def parse_pcode(pcode: str, srcs: dict[str, UOp | int] | None = None) -> tuple[dict, list[tuple[str, UOp]]]:
env: dict = srcs.copy() if srcs else {}
assigns: list[tuple[str, UOp]] = []
raw_lines = [l.strip().rstrip(';') for l in pcode.split('\n') if l.strip() and not l.strip().startswith('//')]
@@ -620,7 +620,7 @@ class _Ctx:
elif dest.startswith('VCC'): stores.extend(self.wmask(_c(VCC_LO.offset), val))
return stores
def compile_sop_pcode(self, op, srcs: dict[str, UOp], sdst_reg: UOp, sdst_size: int) -> UOp:
def compile_sop_pcode(self, op, srcs: dict[str, UOp | int], sdst_reg: UOp, sdst_size: int) -> UOp:
"""Compile a scalar instruction with dynamic destination register."""
pcode = get_pcode(op)
srcs.update({'VCC': self.rmask(_c(VCC_LO.offset)), 'EXEC': self.rexec(), 'SCC': self.rsgpr_dyn(_c(SCC.offset)),
@@ -653,7 +653,7 @@ class _Ctx:
elif dest.startswith('VGPR['): stores.append(self.vgpr.index(val[0].cast(dtypes.int)).store(val[1].cast(dtypes.uint32)))
return UOp.sink(*stores, *self.inc_pc())
def compile_vop_pcode(self, op, srcs: dict[str, UOp], lane: UOp, vdst_reg: UOp, exec_mask: UOp,
def compile_vop_pcode(self, op, srcs: dict[str, UOp | int], lane: UOp, vdst_reg: UOp, exec_mask: UOp,
opsel_dst_hi: bool | UOp = False, sdst_reg: int | None = None, clmp: int = 0,
src0_off: UOp | None = None) -> UOp:
"""Compile VOP instruction. Returns sink with stores and inc_pc."""
@@ -688,6 +688,7 @@ class _Ctx:
if clmp and int_saturate is None and any(p in op.name for p in ('_SUB_U32', '_ADD_U32', '_SUB_U16', '_ADD_U16')):
s0, s1 = srcs.get('S0'), srcs.get('S1')
if s0 is not None and s1 is not None:
assert isinstance(s0, UOp) and isinstance(s1, UOp)
a, b = (s1.cast(dtypes.uint32), s0.cast(dtypes.uint32)) if 'SUBREV' in op.name else (s0.cast(dtypes.uint32), s1.cast(dtypes.uint32))
if 'SUB' in op.name:
int_saturate = (a < b).where(_c(0), a - b) # underflow -> 0
@@ -700,14 +701,13 @@ class _Ctx:
for dest, val in assigns:
# VGPR bit-slice assignment: VGPR[lane][reg][hi:lo] = (vgpr_idx, rhs_val, hi, lo[, cond]) -> read-modify-write
if dest.startswith('VGPR[') and re.search(r'\[\d+:\d+\]', dest):
vgpr_idx, rhs_val, hi_bit, lo_bit = val[:4]
branch_cond = val[4] if len(val) > 4 else None # optional condition from if/else branch
# VGPR bit-slice: (vgpr_idx, rhs_val, hi_bit, lo_bit) - hi/lo are UOp constants
hi_bit, lo_bit = int(val[2].arg), int(val[3].arg)
width = hi_bit - lo_bit + 1
old = self.vgpr.index(vgpr_idx.cast(dtypes.int), ptr=True).load()
new_val = _set_bits(old, _val_to_bits(rhs_val), width, lo_bit).cast(dtypes.uint32)
old = self.vgpr.index(val[0].cast(dtypes.int), ptr=True).load()
new_val = _set_bits(old, _val_to_bits(val[1]), width, lo_bit).cast(dtypes.uint32)
active = _lane_active(exec_mask, lane)
if branch_cond is not None: active = active & _to_u32(branch_cond).ne(_c(0))
raw_stores.append(('vgpr_direct', self.vgpr.index(vgpr_idx.cast(dtypes.int), active).store(new_val)))
raw_stores.append(('vgpr_direct', self.vgpr.index(val[0].cast(dtypes.int), active).store(new_val)))
continue
if 'D0' in dest and '[laneId]' in dest:
old_vcc = self.rmask(_c(VCC_LO.offset))
@@ -715,13 +715,13 @@ class _Ctx:
raw_stores.extend([('vcc', s) for s in self.wmask(_c(VCC_LO.offset), new_vcc)])
elif dest.startswith('D0'):
if (slice_match := re.match(r'D0\[(\d+)\s*:\s*(\d+)\]', dest)):
hi_bit, lo_bit = int(slice_match.group(1)), int(slice_match.group(2))
if hi_bit != 31 or lo_bit != 0:
width, slice_mask = hi_bit - lo_bit + 1, (1 << (hi_bit - lo_bit + 1)) - 1
d0_hi_bit, d0_lo_bit = int(slice_match.group(1)), int(slice_match.group(2))
if d0_hi_bit != 31 or d0_lo_bit != 0:
d0_width, slice_mask = d0_hi_bit - d0_lo_bit + 1, (1 << (d0_hi_bit - d0_lo_bit + 1)) - 1
val_bits = val.bitcast(dtypes.uint16).cast(dtypes.uint32) if val.dtype == dtypes.half else \
val.cast(dtypes.uint32) if val.dtype in (dtypes.uint16, dtypes.int16) else \
val.cast(dtypes.uint32) & UOp.const(dtypes.uint32, slice_mask)
raw_stores.append(('vgpr_slice', (lo_bit, width, val_bits)))
raw_stores.append(('vgpr_slice', (d0_lo_bit, d0_width, val_bits)))
continue
# For integer ops with clamp, use pre-computed saturated value; for floats, clamp to [0,1]
if int_saturate is not None: val = int_saturate
@@ -917,8 +917,10 @@ def _compile_sdwa(inst: irc.VOP1_SDWA | irc.VOP2_SDWA | irc.VOP2_SDWA_SDST | irc
is_vopc = isinstance(inst, irc.VOPC_SDWA_SDST)
exec_mask = ctx.rexec()
# sd=1 means use sdst register, sd=0 means use VCC (for VOPC_SDWA_SDST and VOP2_SDWA_SDST)
has_sdst = isinstance(inst, (irc.VOP2_SDWA_SDST, irc.VOPC_SDWA_SDST))
sdst_off = _c(inst.sdst.offset) if has_sdst and getattr(inst, 'sd', 0) else _c(VCC_LO.offset)
if isinstance(inst, (irc.VOP2_SDWA_SDST, irc.VOPC_SDWA_SDST)):
sdst_off = _c(inst.sdst.offset) if getattr(inst, 'sd', False) else _c(VCC_LO.offset)
else:
sdst_off = _c(VCC_LO.offset)
# Read SDWA fields (these are dynamic but shared across lanes)
src0_sel = ctx.inst_field(type(inst).src0_sel)
src0_sext = ctx.inst_field(type(inst).src0_sext)
@@ -937,7 +939,7 @@ def _compile_sdwa(inst: irc.VOP1_SDWA | irc.VOP2_SDWA | irc.VOP2_SDWA_SDST | irc
s0 = _sdwa_select(s0_raw, src0_sel, src0_sext)
s1_raw = ctx.rsgpr_dyn(vsrc1_reg) if inst.s1 else ctx.rvgpr_dyn(vsrc1_reg, lc)
s1 = _sdwa_select(s1_raw, src1_sel, src1_sext)
srcs: dict[str, UOp] = {'S0': s0, 'S1': s1, 'laneId': lc}
srcs = {'S0': s0, 'S1': s1, 'laneId': lc}
for dest, val in parse_pcode(pcode, srcs)[1]:
if '[laneId]' in dest and ('D0' in dest or 'EXEC' in dest): return val.cast(dtypes.uint32)
return _c(0)
@@ -947,27 +949,27 @@ def _compile_sdwa(inst: irc.VOP1_SDWA | irc.VOP2_SDWA | irc.VOP2_SDWA_SDST | irc
# Non-VOPC path: VOP1_SDWA, VOP2_SDWA, VOP2_SDWA_SDST — uses lane loop
lane = ctx.range()
vdst_reg = ctx.inst_field(type(inst).vdst)
vdst_reg = ctx.inst_field(type(inst).vdst) # type: ignore[union-attr]
s0_raw = ctx.rsgpr_dyn(vsrc0_reg) if inst.s0 else ctx.rvgpr_dyn(vsrc0_reg, lane)
s0 = _sdwa_select(s0_raw, src0_sel, src0_sext)
if isinstance(inst, (irc.VOP2_SDWA, irc.VOP2_SDWA_SDST)):
s1_raw = ctx.rsgpr_dyn(vsrc1_reg) if inst.s1 else ctx.rvgpr_dyn(vsrc1_reg, lane)
s1 = _sdwa_select(s1_raw, src1_sel, src1_sext)
srcs: dict[str, UOp] = {'S0': s0, 'S1': s1, 'D0': ctx.rvgpr_dyn(vdst_reg, lane)}
srcs:dict[str, UOp | int] = {'S0': s0, 'S1': s1, 'D0': ctx.rvgpr_dyn(vdst_reg, lane)}
else:
srcs = {'S0': s0}
# dst_sel and dst_unused
has_dst_sel = hasattr(type(inst), 'dst_sel')
if has_dst_sel:
dst_sel = ctx.inst_field(type(inst).dst_sel)
dst_unused = ctx.inst_field(type(inst).dst_unused)
dst_sel = ctx.inst_field(type(inst).dst_sel) # type: ignore[union-attr]
dst_unused = ctx.inst_field(type(inst).dst_unused) # type: ignore[union-attr]
srcs.update({'VCC': ctx.rmask(_c(VCC_LO.offset)), 'EXEC': exec_mask, 'SCC': ctx.rsgpr_dyn(_c(SCC.offset)),
'laneId': lane, 'VDST': vdst_reg, 'ROUND_MODE': _c(0), 'ROUND_TOWARD_ZERO': _c(0),
'ROUND_NEAREST_EVEN': _c(0), '_vgpr': ctx.vgpr, '_wave_size': ctx.wave_size,
'SDWA_SRC0_SEL': _c(0), 'BYTE0': _c(0), 'BYTE1': _c(1), 'BYTE2': _c(2), 'BYTE3': _c(3),
'WORD0': _c(0), 'WORD1': _c(1)})
_, assigns = parse_pcode(pcode, srcs)
stores: list[UOp] = []
stores = []
vcc_val = None
for dest, val in assigns:
if 'D0' in dest and '[laneId]' in dest:
@@ -1020,7 +1022,7 @@ def _compile_vop12(inst: ir3.VOP1 | ir3.VOP1_SDST | ir3.VOP2 | ir4.VOP1 | ir4.VO
src0_reg = src0_hi.where(src0_off - _c(384), _c(0))
s0 = src0_hi.where(_hi16(ctx.rvgpr_dyn(src0_reg, lane)), s0)
d0 = _cond_hi16(write_hi_half, ctx.rvgpr_dyn(vdst_reg, lane))
srcs = {'S0': s0, 'D0': d0}
srcs:dict[str, UOp | int] = {'S0': s0, 'D0': d0}
else:
vsrc1_reg = ctx.inst_field(type(inst).vsrc1)
vsrc1_hi = bits['s0'] == 16 and (vsrc1_reg >= _c(128))
@@ -1256,6 +1258,7 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp:
src1_is_vgpr = src1_off >= _c(256)
m = _re.search(r'(\d+)X(\d+)X(\d+)', op_name)
if m is None: raise ValueError(f"could not parse MFMA dimensions from {op_name}")
M, N, K = int(m.group(1)), int(m.group(2)), int(m.group(3))
is_bf16 = 'BF16' in op_name
@@ -1545,7 +1548,7 @@ def _compile_vop3p(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P, ctx: _Ctx) -> UOp:
is_pk_f32 = 'PK' in op_name and 'F32' in op_name and 'MOV' not in op_name # CDNA packed F32 ops
is_pk_mov_b32 = 'PK_MOV_B32' in op_name # CDNA packed MOV needs special handling
do_cast = any(x in op_name for x in ('F16', 'F32', 'BF16')) and 'IU' not in op_name and not is_pk_f32
literal = ctx.inst_field(type(inst).literal) if hasattr(type(inst), 'literal') else None
literal = ctx.inst_field(type(inst).literal) if hasattr(type(inst), 'literal') else None # type: ignore[union-attr]
src0 = ctx.rsrc_dyn(ctx.inst_field(type(inst).src0), lane, 16, literal=literal, do_cast=do_cast)
src1 = ctx.rsrc_dyn(ctx.inst_field(type(inst).src1), lane, 16, literal=literal, do_cast=do_cast)
src2 = ctx.rsrc_dyn(ctx.inst_field(type(inst).src2), lane, 16, literal=literal, do_cast=do_cast)
@@ -1571,6 +1574,7 @@ def _compile_vop3p(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P, ctx: _Ctx) -> UOp:
stores = [ctx.wvgpr_dyn(vdst_reg, lane, lo_out, exec_mask), ctx.wvgpr_dyn(vdst_reg + _c(1), lane, hi_out, exec_mask)]
return UOp.sink(UOp.group(*stores).end(lane), *ctx.inc_pc())
srcs: dict[str, UOp | int] = {}
if is_pk_f32:
# CDNA packed F32: read 32-bit sources, build 64-bit packed values using opsel.
# For VGPRs: opsel selects between v[reg] (0) and v[reg+1] (1) for each half.
@@ -1648,6 +1652,7 @@ def _compile_vopd(inst: ir3.VOPD | ir4.VOPD, ctx: _Ctx) -> UOp:
lane = ctx.range()
srcy0, srcy1 = ctx.rsrc_dyn(srcy0_off, lane, literal=literal), ctx.rvgpr_dyn(vsrcy1_reg, lane)
all_stores = []
srcs:dict[str, UOp | int] = {}
for op, src0_off, vsrc1_reg, vdst_reg, label in [(inst.opx, srcx0_off, vsrcx1_reg, vdstx_reg, 'X'),
(inst.opy, srcy0_off, vsrcy1_reg, vdsty_reg, 'Y')]:
vop = VOPD_TO_VOP2.get(op)
test/mockgpu/amd/pcode.py
+1 -1
View File
@@ -1073,7 +1073,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
ws = env.get('_wave_size', 32)
vgpr_idx = _to_u32(rg) * _u32(ws) + _to_u32(ln)
if assigns is not None:
assigns.append((f'VGPR[{_tok_str(lane_toks)}][{_tok_str(reg_toks)}][{hi_val}:{lo_val}]', (vgpr_idx, val, hi_val, lo_val)))
assigns.append((f'VGPR[{_tok_str(lane_toks)}][{_tok_str(reg_toks)}][{hi_val}:{lo_val}]', (vgpr_idx, val, _u32(hi_val), _u32(lo_val))))
i += 1
continue
if j < len(toks) and toks[j].type == 'DOT': j += 2 # skip .type suffix
test/mockgpu/cuda/cuda.py
+1 -1
View File
@@ -154,7 +154,7 @@ def cuMemHostAlloc(pp, bytesize: int, flags: int) -> int:
def cuMemFreeHost(p: ctypes.c_void_p) -> int: return cuMemFree_v2(p)
def cuMemcpyDtoDAsync_v2(dst, src, bytesize: int, stream: Any) -> int:
ctypes.memmove(dst.value, src.value, bytesize)
ctypes.memmove(dst if isinstance(dst, int) else dst.value, src if isinstance(src, int) else src.value, bytesize)
return orig_cuda.CUDA_SUCCESS
def cuFuncSetAttribute(hfunc, attrib: int, value: int) -> int:
test/null/test_gpudims.py
+1 -1
View File
@@ -24,7 +24,7 @@ class TestGroupedDims(unittest.TestCase):
total = math.prod(dims)
specials = sorted(dedup(flatten([[y for y in x.toposort() if y.op is Ops.SPECIAL] for x in idxs])), key=lambda u: u.arg)
# build flat index and primed flat (same expression with renamed SPECIALs)
flat = UOp.const(dtypes.index, 0)
flat = UOp.const(dtypes.weakint, 0)
for i, idx in enumerate(idxs):
flat = flat + idx * int(math.prod(dims[i+1:]))
flat_p = flat.substitute({s: UOp(Ops.SPECIAL, s.dtype, s.src, s.arg+"_p") for s in specials})
test/null/test_graph_rewrite.py
+4 -4
View File
@@ -99,21 +99,21 @@ class TestFoldingAndReduction(unittest.TestCase):
class TestModuloAndDivisionFolding(unittest.TestCase):
def test_full_graph_rewrite_modulo_folding_with_define_var(self):
# index dtype because div-mod rules only work on index
x_var_uop = UOp.variable('x', 0, 100).cast(dtypes.index)
x_var_uop = UOp.variable('x', 0, 100).cast(dtypes.weakint)
optimized_mod_uop = apply_rewrite(((x_var_uop * 4) + 2) % 4)
self.assertEqual(optimized_mod_uop.op, Ops.CONST)
self.assertEqual(optimized_mod_uop.arg, 2)
def test_full_graph_rewrite_division_folding_with_define_var(self):
# index dtype because div-mod rules only work on index
n_var_uop = UOp.variable('n', 1, 1000).cast(dtypes.index)
n_var_uop = UOp.variable('n', 1, 1000).cast(dtypes.weakint)
optimized_div_uop = apply_rewrite((n_var_uop * 6) // 3)
self.assertEqual(optimized_div_uop.op, Ops.MUL)
self.assertEqual(optimized_div_uop.src[1].arg, 2)
def test_full_graph_rewrite_complex_mod_div_folding(self):
# index dtype because div-mod rules only work on index
k_var_uop = UOp.variable('k', 0, 50).cast(dtypes.index)
k_var_uop = UOp.variable('k', 0, 50).cast(dtypes.weakint)
optimized_div_uop = apply_rewrite(((k_var_uop * 12 + 8) % 6) // 2)
self.assertEqual(optimized_div_uop.op, Ops.CONST)
self.assertEqual(optimized_div_uop.arg, 1)
@@ -132,7 +132,7 @@ class TestModuloAndDivisionFolding(unittest.TestCase):
def test_full_graph_rewrite_modulo_large_divisor(self):
# index dtype because div-mod rules only work on index
x_var_uop = UOp.variable('x', 1, 5)
self.assertIs(apply_rewrite(x_var_uop.cast(dtypes.index) % 10).render(simplify=False), x_var_uop.render(simplify=False))
self.assertIs(apply_rewrite(x_var_uop.cast(dtypes.weakint) % 10).render(simplify=False), x_var_uop.render(simplify=False))
def test_full_graph_rewrite_division_with_remainder(self):
x_var_uop = UOp.variable('x', 7, 9)
test/null/test_linearizer_failures.py
+5 -5
View File
@@ -8,12 +8,12 @@ from tinygrad.device import Device
class TestLinearizerFailures(unittest.TestCase):
def test_fail_1(self):
c0 = UOp(Ops.PARAM, dtypes.float.ptr(64), arg=0, src=())
c1 = UOp.range(UOp.const(dtypes.index, 2), 1, AxisType.LOOP)
c2 = UOp.range(UOp.const(dtypes.index, 32), 2, AxisType.LOOP)
c3 = ((c1*UOp.const(dtypes.index, 32))+c2)
c1 = UOp.range(UOp.const(dtypes.weakint, 2), 1, AxisType.LOOP)
c2 = UOp.range(UOp.const(dtypes.weakint, 32), 2, AxisType.LOOP)
c3 = ((c1*UOp.const(dtypes.weakint, 32))+c2)
c4 = UOp(Ops.PARAM, dtypes.float.ptr(163840), arg=1, src=())
c5 = UOp.range(UOp.const(dtypes.index, 2560), 0, AxisType.REDUCE)
c6 = c4.index(((((((c5//UOp.const(dtypes.index, 8))%UOp.const(dtypes.index, 8))*UOp.const(dtypes.index, 8))+(c5%UOp.const(dtypes.index, 8)))+(((c2*UOp.const(dtypes.index, 40))+(c5//UOp.const(dtypes.index, 64)))*UOp.const(dtypes.index, 64)))+(c1*UOp.const(dtypes.index, 81920))))
c5 = UOp.range(UOp.const(dtypes.weakint, 2560), 0, AxisType.REDUCE)
c6 = c4.index(((((((c5//UOp.const(dtypes.weakint, 8))%UOp.const(dtypes.weakint, 8))*UOp.const(dtypes.weakint, 8))+(c5%UOp.const(dtypes.weakint, 8)))+(((c2*UOp.const(dtypes.weakint, 40))+(c5//UOp.const(dtypes.weakint, 64)))*UOp.const(dtypes.weakint, 64)))+(c1*UOp.const(dtypes.weakint, 81920))))
c7 = UOp(Ops.PARAM, dtypes.float.ptr(64), arg=2, src=())
c8 = c7.index(c3)
c9 = ((((c6+(c8*UOp.const(dtypes.float, -1.0)))*(c6+(c8*UOp.const(dtypes.float, -1.0)))).reduce(c5, arg=Ops.ADD)*UOp.const(dtypes.float, 0.000390625))+UOp.const(dtypes.float, 1e-05)).sqrt().reciprocal()
test/null/test_simplify_valid_idx.py
+17 -27
View File
@@ -15,11 +15,11 @@ def get_gated_load_uop(valid:UOp, idx:UOp):
def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UOp]):
return UOp(Ops.LOAD, dtypes.float.vec(4), (
UOp(Ops.PARAM, dtypes.imagef(image_shape), arg=0).index(UOp(Ops.VECTORIZE, dtypes.index.vec(2), idx).valid(valid), ptr=True),
UOp(Ops.PARAM, dtypes.imagef(image_shape), arg=0).index(UOp(Ops.VECTORIZE, dtypes.weakint.vec(2), idx).valid(valid), ptr=True),
UOp(Ops.VECTORIZE, dtypes.float.vec(4), src=(UOp.const(dtypes.float, 0.0),) * 4)
))
def Special(expr, nmax): return UOp(Ops.SPECIAL, dtypes.index, (UOp.const(dtypes.index, nmax),), expr)
def Special(expr, nmax): return UOp(Ops.SPECIAL, dtypes.weakint, (UOp.const(dtypes.weakint, nmax),), expr)
def Variable(expr, nmin, nmax): return UOp.variable(expr, nmin, nmax)
def Range(n, nmax): return UOp.range(nmax, n)
@@ -454,23 +454,13 @@ class TestImageSimplification(unittest.TestCase):
load = get_load_image_uop((32, 1024, 4), valid, (alu0, alu1))
self.check(load, "(lidx1<7)", "((gidx0*2+lidx1*512+(lidx0*8192+r0*4096)+-11711)//4%1024)", "(lidx0*2+r0+-3)")
class TestUnfoldableImageChannelSelection(unittest.TestCase):
def _count_nans(self, load):
with Context(NOOPT=1, SPEC=0):
result = full_rewrite_to_sink(load.sink()).src[0]
return sum(1 for u in result.toposort() if u.op is Ops.CONST and u.arg != u.arg)
def test_bounded_channel_no_nan(self):
# unfoldable image load with bounded idx % 4 range [0,1] -> no NAN fallback needed
lidx = Special("lidx", 2)
load = UOp(Ops.LOAD, dtypes.float, (UOp(Ops.PARAM, dtypes.imagef((10, 10, 4)), arg=0).index(lidx, ptr=True), UOp.const(dtypes.float, 0)))
self.assertEqual(self._count_nans(load), 0)
def test_unbounded_channel_has_nan(self):
# variable with negative range -> x % 4 can be negative -> needs NAN fallback
x = Variable("x", -10, 10)
load = UOp(Ops.LOAD, dtypes.float, (UOp(Ops.PARAM, dtypes.imagef((10, 10, 4)), arg=0).index(x, ptr=True), UOp.const(dtypes.float, 0)))
self.assertEqual(self._count_nans(load), 1)
class TestUnfoldableImage(unittest.TestCase):
def test_unfoldable_becomes_buffer(self):
with Context(SPEC=0):
lidx = Special("lidx", 2)
load = UOp(Ops.LOAD, dtypes.float, (UOp(Ops.PARAM, dtypes.imagef((10, 10, 4)), arg=0).index(lidx, ptr=True), UOp.const(dtypes.float, 0)))
res = full_rewrite_to_sink(load.sink()).src[0]
self.assertEqual(res.src[0].src[0].dtype, dtypes.float.ptr(400))
class TestDropTrueGate(unittest.TestCase):
def test_drop_true_gate_on_index(self):
@@ -478,7 +468,7 @@ class TestDropTrueGate(unittest.TestCase):
from tinygrad.codegen.late.devectorizer import load_store_indexing
from tinygrad.uop.ops import graph_rewrite
buf = UOp(Ops.PARAM, dtypes.int.ptr(), arg=0)
idx = UOp.const(dtypes.index, 0)
idx = UOp.const(dtypes.weakint, 0)
true_gate = UOp.const(dtypes.bool, True)
index_with_gate = UOp(Ops.INDEX, dtypes.int.ptr(), (buf, idx, true_gate))
# apply the optimization
@@ -495,7 +485,7 @@ class TestRangeShrink(unittest.TestCase):
def test_range_shrink_single_guard(self):
# range 0..203 guarded by r < 4 everywhere -> shrink to 0..3
r = Range(0, 204)
load = get_gated_load_uop(r < UOp.const(dtypes.index, 4), r)
load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
ranges = self.get_ranges(load.sink())
self.assertEqual(len(ranges), 1)
self.assertEqual(ranges[0].src[0].arg, 4)
@@ -503,8 +493,8 @@ class TestRangeShrink(unittest.TestCase):
def test_range_shrink_picks_max_guard(self):
# two loads guard the same range with r < 4 and r < 8 -> shrink to max(4, 8) = 8
r = Range(0, 204)
load1 = get_gated_load_uop(r < UOp.const(dtypes.index, 4), r)
load2 = get_gated_load_uop(r < UOp.const(dtypes.index, 8), r)
load1 = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
load2 = get_gated_load_uop(r < UOp.const(dtypes.weakint, 8), r)
ranges = self.get_ranges(UOp.sink(load1, load2))
self.assertEqual(len(ranges), 1)
self.assertEqual(ranges[0].src[0].arg, 8)
@@ -512,7 +502,7 @@ class TestRangeShrink(unittest.TestCase):
def test_range_no_shrink_guard_ge_max(self):
# guard r < 300 with range max 204 -> no shrink (guard doesn't constrain)
r = Range(0, 204)
load = get_gated_load_uop(r < UOp.const(dtypes.index, 300), r)
load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 300), r)
ranges = self.get_ranges(load.sink())
self.assertEqual(len(ranges), 1)
self.assertEqual(ranges[0].src[0].arg, 204)
@@ -520,7 +510,7 @@ class TestRangeShrink(unittest.TestCase):
def test_range_no_shrink_when_unguarded_elsewhere(self):
# one load guards r < 4, but another load uses r without a gate -> no shrink
r = Range(0, 204)
load1 = get_gated_load_uop(r < UOp.const(dtypes.index, 4), r)
load1 = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
load2 = UOp(Ops.LOAD, dtypes.float, (UOp(Ops.PARAM, dtypes.float.ptr(), arg=1).index(r, ptr=True),))
ranges = self.get_ranges(UOp.sink(load1, load2))
self.assertEqual(len(ranges), 1)
@@ -529,7 +519,7 @@ class TestRangeShrink(unittest.TestCase):
def test_range_no_shrink_when_used_in_reduce(self):
# range used in both a gated load AND directly in the reduce expression -> no shrink
r = Range(0, 204)
gated_load = get_gated_load_uop(r < UOp.const(dtypes.index, 4), r)
gated_load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
red = UOp(Ops.REDUCE, dtypes.float, (r.cast(dtypes.float) + gated_load, r), Ops.ADD)
ranges = self.get_ranges(red.sink())
self.assertEqual(len(ranges), 1)
@@ -538,7 +528,7 @@ class TestRangeShrink(unittest.TestCase):
def test_range_shrink_to_single_iteration(self):
# guard r < 1 shrinks range to 1 -> single iteration, range eliminated entirely
r = Range(0, 204)
load = get_gated_load_uop(r < UOp.const(dtypes.index, 1), r)
load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 1), r)
ranges = self.get_ranges(load.sink())
self.assertEqual(len(ranges), 0)
test/null/test_tensor_metadata.py
+4 -2
View File
@@ -2,6 +2,7 @@ import unittest
from tinygrad import Tensor, dtypes
from tinygrad.tensor import _METADATA
from tinygrad.engine.realize import capturing
from tinygrad.engine.schedule import linear_to_schedule
from tinygrad.helpers import Context
@unittest.skip("tensor metadata is no longer supported")
@@ -94,10 +95,11 @@ class TestTensorMetadata(unittest.TestCase):
self.assertEqual(si.metadata, ())
def _has_metadata(self, h, name):
items = []
capturing.append(type("", (), {"add": lambda _, ei: items.append(ei)})())
linears = []
capturing.append(type("", (), {"add_linear": lambda _, linear, var_vals: linears.append(linear)})())
try: h.realize()
finally: capturing.clear()
items = [ei for linear in linears for ei in linear_to_schedule(linear)]
return any(m.name == name for ei in items for m in ei.metadata)
def test_metadata_survives_realize_pending_assign(self):
test/null/test_uop_graph.py
+14 -14
View File
@@ -205,7 +205,7 @@ class TestUOpGraph(unittest.TestCase):
def test_where_same_fold(self):
v = UOp.variable('tmp', 0, 1)
c0 = UOp.const(dtypes.index, 0)
c0 = UOp.const(dtypes.weakint, 0)
vc = UOp(Ops.CMPNE, dtypes.bool, (v, c0))
c1 = UOp.const(dtypes.float, 1.0)
out = UOp(Ops.WHERE, dtypes.float, (vc, c1, c1))
@@ -410,7 +410,7 @@ class TestUOpGraph(unittest.TestCase):
d0 = UOp(Ops.PARAM, dt.ptr(), arg=0)
v = d0.index(UOp.const(dtypes.int, 0))
uops = to_uops_list([v.bitcast(dt)])
self.assertEqual(len([x for x in uops if x.op is Ops.BITCAST]), 0, f"dtype = {dt}")
self.assertEqual(len([x for x in uops if x.op is Ops.BITCAST and x.dtype is dt]), 0, f"dtype = {dt}")
def test_sub_with_cast_folds(self):
a = Variable("a", 0, 5)
@@ -469,16 +469,16 @@ class TestUOpGraph(unittest.TestCase):
# mnist indexing with split reduceop
# Make sure we are not doign math on the loaded index, which would promote it to long
c0 = UOp(Ops.PARAM, dtypes.uchar.ptr(128000), arg=0, src=())
c1 = UOp.range(UOp.const(dtypes.index, 512), 1, AxisType.LOOP)
c2 = UOp.range(UOp.const(dtypes.index, 250), 2, AxisType.LOOP)
c1 = UOp.range(UOp.const(dtypes.weakint, 512), 1, AxisType.LOOP)
c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.LOOP)
c3 = UOp(Ops.PARAM, dtypes.int.ptr(512), arg=1, src=())
c4 = c3.index(c1)
c5 = UOp.range(UOp.const(dtypes.index, 240), 0, AxisType.REDUCE)
c6 = ((c2*UOp.const(dtypes.index, 240))+c5)
c5 = UOp.range(UOp.const(dtypes.weakint, 240), 0, AxisType.REDUCE)
c6 = ((c2*UOp.const(dtypes.weakint, 240))+c5)
c7 = UOp(Ops.PARAM, dtypes.uchar.ptr(60000), arg=2, src=())
c8 = c7.index(c6)
c9 = ((c4<0).where((c4+60000), c4)!=c6.cast(dtypes.int)).where(0, c8.cast(dtypes.uint).cast(dtypes.uchar)).reduce(c5, arg=Ops.ADD)
c10 = c0.index(((c1*UOp.const(dtypes.index, 250))+c2)).store(c9).end(c1, c2)
c10 = c0.index(((c1*UOp.const(dtypes.weakint, 250))+c2)).store(c9).end(c1, c2)
uops = to_uops_list([c10])
for u in uops:
self.assertNotEqual(u.dtype, dtypes.long)
@@ -486,19 +486,19 @@ class TestUOpGraph(unittest.TestCase):
def test_load_idx_no_math_on_loaded(self):
# test the (x+y)<c pattern where x has loads - we shouldn't do math on loaded indices
c0 = UOp(Ops.PARAM, dtypes.uchar.ptr(128000), arg=0, src=())
c1 = UOp.range(UOp.const(dtypes.index, 512), 1, AxisType.LOOP)
c2 = UOp.range(UOp.const(dtypes.index, 250), 2, AxisType.LOOP)
c1 = UOp.range(UOp.const(dtypes.weakint, 512), 1, AxisType.LOOP)
c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.LOOP)
c3 = UOp(Ops.PARAM, dtypes.int.ptr(512), arg=1, src=())
c4 = c3.index(c1) # c4 is a load
c5 = UOp.range(UOp.const(dtypes.index, 240), 0, AxisType.REDUCE)
c6 = ((c2*UOp.const(dtypes.index, 240))+c5)
c5 = UOp.range(UOp.const(dtypes.weakint, 240), 0, AxisType.REDUCE)
c6 = ((c2*UOp.const(dtypes.weakint, 240))+c5)
c7 = UOp(Ops.PARAM, dtypes.uchar.ptr(60000), arg=2, src=())
c8 = c7.index(c6)
# (loaded + range) < const pattern - loaded value shouldn't be promoted to long
loaded_idx = c4.cast(dtypes.index)
comparison = (loaded_idx + c5) < UOp.const(dtypes.index, 60000)
loaded_idx = c4.cast(dtypes.weakint)
comparison = (loaded_idx + c5) < UOp.const(dtypes.weakint, 60000)
c9 = comparison.where(c8.cast(dtypes.uint).cast(dtypes.uchar), 0).reduce(c5, arg=Ops.ADD)
c10 = c0.index(((c1*UOp.const(dtypes.index, 250))+c2)).store(c9).end(c1, c2)
c10 = c0.index(((c1*UOp.const(dtypes.weakint, 250))+c2)).store(c9).end(c1, c2)
uops = to_uops_list([c10])
for u in uops:
self.assertNotEqual(u.dtype, dtypes.long)
test/null/test_uop_symbolic.py
+12 -12
View File
@@ -12,13 +12,13 @@ from tinygrad.uop.validate import uops_to_z3
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)
if isinstance(s_eval, int) and v.dtype==dtypes.weakint: s_eval = UOp.const(dtypes.weakint, 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.render()} 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)
def Variable(name: str, min_val: ConstType, max_val: ConstType, dtype: DType=dtypes.weakint): return UOp.variable(name,min_val,max_val,dtype)
def uconst(val): return UOp.const(dtypes.weakint, val)
def usum(ops): return functools.reduce(lambda x,y: x+y, ops)
def uand(ops): return functools.reduce(lambda x,y: x*y, ops)
@@ -245,12 +245,12 @@ class TestSymbolic(unittest.TestCase):
self.assertEqual((Variable("x", -10, 0)%Variable("y", 1, 10))._min_max, (-9, 0))
def test_range_div_its_symbolic_bound(self):
a = Variable("a", 1, 10, dtypes.index)
a = Variable("a", 1, 10, dtypes.weakint)
ridx0 = UOp.range(a+2, 0)
self.helper_test_variable(ridx0//(a+2), 0, 0, "0")
def test_range_mod_its_symbolic_bound(self):
a = Variable("a", 1, 10, dtypes.index)
a = Variable("a", 1, 10, dtypes.weakint)
ridx = UOp.range(a+2, 0)
self.helper_test_variable(ridx%(a+2), 0, 11, "r0")
@@ -941,7 +941,7 @@ class TestSymbolic(unittest.TestCase):
self.helper_test_variable((numerator//denominator)<=0, 1, 1, "True")
def test_symbolic_range_doesnt_collapse(self):
r0 = UOp.range((Variable("a", 1, 10)<5).cast(dtypes.index), 0)
r0 = UOp.range((Variable("a", 1, 10)<5).cast(dtypes.weakint), 0)
self.helper_test_variable(r0, 0, 0, "r0")
def test_const_reciprocal(self):
@@ -1202,16 +1202,16 @@ class TestInvalidIndex(unittest.TestCase):
self.assertIs((UOp.invalid()<Variable("a",0,10)).simplify().dtype, dtypes.bool)
def test_alu_invalid_vconst(self):
c1 = UOp.const(dtypes.index.vec(4), (1, 1, Invalid, Invalid))
c2 = UOp.const(dtypes.index.vec(4), (1, Invalid, 1, 1))
self.assertIs((c1+c2).simplify(), UOp.const(dtypes.index.vec(4), (2, Invalid, Invalid, Invalid)))
c1 = UOp.const(dtypes.weakint.vec(4), (1, 1, Invalid, Invalid))
c2 = UOp.const(dtypes.weakint.vec(4), (1, Invalid, 1, 1))
self.assertIs((c1+c2).simplify(), UOp.const(dtypes.weakint.vec(4), (2, Invalid, Invalid, Invalid)))
class TestStoreLoadFolding(unittest.TestCase):
"""Tests for store(index, load(index)) -> NOOP rule. This rule matches patterns that EMERGE during simplification."""
def test_store_load_folding(self):
# store(idx, load(idx)) -> NOOP, including emergent patterns like store(idx, load(idx) + 0)
buf = UOp(Ops.PARAM, dtypes.int.ptr(), arg=0)
index = buf.index(UOp.const(dtypes.index, 0))
index = buf.index(UOp.const(dtypes.weakint, 0))
# Direct: store(idx, load(idx)) -> NOOP
self.assertEqual(graph_rewrite(index.store(index.load()), sym).op, Ops.NOOP)
# Emergent: store(idx, load(idx) + 0) -> store(idx, load(idx)) -> NOOP
@@ -1250,10 +1250,10 @@ class TestGatedUopGivenValid(unittest.TestCase):
idx0 = (r0 + uconst(-1)) // uconst(3)
idx1 = r0 % uconst(3)
idx:UOp = (r0 < 3).where(UOp(Ops.VECTORIZE, dtypes.index.vec(2), (idx0, idx1)), UOp.invalid())
idx:UOp = (r0 < 3).where(UOp(Ops.VECTORIZE, dtypes.weakint.vec(2), (idx0, idx1)), UOp.invalid())
idx = graph_rewrite(idx, pm_simplify_valid)
# NOTE: independent simplification: (r0-1)//3 -> 0, r0%3 -> r0 when r0 in [0,2]
expected_vec = UOp(Ops.VECTORIZE, dtypes.index.vec(2), (uconst(0), r0))
expected_vec = UOp(Ops.VECTORIZE, dtypes.weakint.vec(2), (uconst(0), r0))
self.assertEqual(idx, (r0 < 3).where(expected_vec, UOp.invalid()))
class TestRangeSplitting(unittest.TestCase):
test/null/test_uop_vmin_vmax.py
+1 -1
View File
@@ -144,7 +144,7 @@ class TestVminVmaxProperties(unittest.TestCase):
self.assertNotEqual(i.vmin, i.vmax)
def test_vmin_vmax_invalid_vconst(self):
x = UOp.const(dtypes.index.vec(4), (0, 4, Invalid, Invalid))
x = UOp.const(dtypes.weakint.vec(4), (0, 4, Invalid, Invalid))
self.assertLess(x.vmin, 0)
self.assertGreater(x.vmax, 4)
test/null/test_validate_oob.py
+2 -2
View File
@@ -126,7 +126,7 @@ class TestValidateOOB(unittest.TestCase):
buf0 = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
buf1 = UOp(Ops.PARAM, dtypes.int.ptr(64), (), 1)
r = UOp.range(42, 0, AxisType.GLOBAL)
ld0 = buf0.index(r.valid(r < 8), ptr=True).load(dtype=dtypes.int).cast(dtypes.index)
ld0 = buf0.index(r.valid(r < 8), ptr=True).load(dtype=dtypes.int).cast(dtypes.weakint)
to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 32)), ptr=True).load(dtype=dtypes.int)]) # valid
with self.assertRaises(RuntimeError):
to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 64)), ptr=True).load(dtype=dtypes.int)]) # oob
@@ -135,7 +135,7 @@ class TestValidateOOB(unittest.TestCase):
with Context(CHECK_OOB=1, SPEC=2):
buf_bool = UOp(Ops.PARAM, dtypes.bool.ptr(16), (), 0)
buf_int = UOp(Ops.PARAM, dtypes.int.ptr(8), (), 1)
gidx = UOp(Ops.SPECIAL, dtypes.index, (UOp.const(dtypes.index, 16),), "gidx0")
gidx = UOp(Ops.SPECIAL, dtypes.weakint, (UOp.const(dtypes.weakint, 16),), "gidx0")
ld_bool = buf_bool.index(gidx, ptr=True).load()
with self.assertRaises(RuntimeError):
to_uops_list([buf_int.index(gidx.valid(ld_bool), ptr=True).load()]) # gidx 0..15, buf_int size 8
test/test_tiny.py
+5 -5
View File
@@ -1,7 +1,7 @@
# basic self-contained tests of the external functionality of tinygrad
import unittest, random
from tinygrad import Tensor, Context, Variable, TinyJit, dtypes, Device, nn
from tinygrad.helpers import IMAGE, CI, getenv
from tinygrad.helpers import CI, getenv
class TestTiny(unittest.TestCase):
@@ -39,14 +39,14 @@ class TestTiny(unittest.TestCase):
out = Tensor.ones(N).contiguous().sum()
self.assertEqual(out.item(), N)
def test_gemm(self, N=getenv("GEMM_N", 64), out_dtype=dtypes.float):
def test_gemm(self, N=getenv("GEMM_N", 64)):
a = Tensor.ones(N,N).contiguous()
b = Tensor.eye(N).contiguous()
lst = (out:=a@b).tolist()
for y in range(N):
for x in range(N):
self.assertEqual(lst[y][x], 1.0, msg=f"mismatch at ({y},{x})")
if IMAGE < 2: self.assertEqual(out.dtype, out_dtype)
self.assertEqual(out.dtype, dtypes.float)
def test_gemv(self, N=getenv("GEMV_N", 64), out_dtype=dtypes.float):
a = Tensor.ones(1,N).contiguous()
@@ -54,7 +54,7 @@ class TestTiny(unittest.TestCase):
lst = (out:=a@b).tolist()
for x in range(N):
self.assertEqual(lst[0][x], 1.0, msg=f"mismatch at {x}")
if IMAGE < 2: self.assertEqual(out.dtype, out_dtype)
self.assertEqual(out.dtype, out_dtype)
# *** randomness ***
@@ -149,7 +149,7 @@ class TestTiny(unittest.TestCase):
@unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
def test_image(self):
with Context(IMAGE=2): self.test_gemm(N=4, out_dtype=dtypes.imagef((4, 1, 4)))
with Context(IMAGE=1): self.test_gemm(N=64)
def test_beam_image(self):
with Context(BEAM=1, IGNORE_BEAM_CACHE=1): self.test_image()
test/testextra/test_export_model.py
+2 -1
View File
@@ -1,6 +1,7 @@
import unittest
from extra.export_model import export_model, EXPORT_SUPPORTED_DEVICE
from tinygrad.tensor import Tensor, Device
from tinygrad.tensor import Tensor
from tinygrad.device import Device
from tinygrad import dtypes
import json
test/unit/test_dtype_spec.py
+11 -6
View File
@@ -1,8 +1,8 @@
import unittest, math, subprocess
from tinygrad.tensor import Tensor, dtypes, Device
from tinygrad.dtype import DType, DTYPES_DICT
from tinygrad.device import is_dtype_supported
from tinygrad.helpers import getenv, DEBUG
from tinygrad.tensor import Tensor
from tinygrad.dtype import dtypes, DType, DTYPES_DICT
from tinygrad.device import Device, is_dtype_supported
from tinygrad.helpers import getenv, DEBUG, EMULATED_DTYPES
from test.helpers import slow
from hypothesis import given, settings, strategies as strat
import numpy as np
@@ -24,8 +24,13 @@ def _assert_eq(tensor:Tensor, target_dtype:DType, target, tol_target_dtype:float
if DEBUG >= 2: print(tensor.numpy())
try:
assert tensor.dtype == target_dtype
np.testing.assert_allclose(tensor.numpy(), target, rtol={dtypes.float16:1e-3, dtypes.bfloat16:1e-2,
dtypes.fp8e4m3:1e-1, dtypes.fp8e5m2:5e-1, dtypes.fp8e4m3fnuz:1e-1, dtypes.fp8e5m2fnuz:5e-1}.get(target_dtype, tol_target_dtype))
# denormals are zero
if target_dtype in dtypes.floats and (not is_dtype_supported(target_dtype) or target_dtype in EMULATED_DTYPES.tolist(dtypes)):
fe, fm = dtypes.finfo(target_dtype)
kwargs = {"atol":2 ** (2 - (1 << (fe - 1))), "rtol": 2 ** (-fm)}
else: kwargs = {"rtol": {dtypes.float16:1e-3, dtypes.bfloat16:1e-2, dtypes.fp8e4m3:1e-1, dtypes.fp8e5m2:5e-1,
dtypes.fp8e4m3fnuz:1e-1, dtypes.fp8e5m2fnuz:5e-1}.get(target_dtype, tol_target_dtype)}
np.testing.assert_allclose(tensor.numpy(), target, **kwargs)
except AssertionError as e:
raise AssertionError(f"\ntensor {tensor.numpy()} dtype {tensor.dtype} does not match target {target} with dtype {target_dtype}") from e
test/unit/test_function.py
+14
View File
@@ -224,6 +224,20 @@ class TestFunction(unittest.TestCase):
s = State()
np.testing.assert_equal(s(Tensor([[5., 6.], [7., 8.]])).numpy(), [[5., 6.], [7., 8.]])
def test_single_after_store(self):
"""AFTER(buf, STORE(view, data)) should write data through the view into buf, same as the double-after pattern."""
@function
def f(buf:Tensor, x:Tensor, start_pos:int|UOp) -> Tensor:
slice_uop = buf[:, start_pos:start_pos+1].uop
assigned = Tensor(buf.uop.after(slice_uop.store(x.uop)))
return assigned
buf = Tensor.zeros(2, 8).contiguous().realize()
x = Tensor([[1.], [2.]]).realize()
v = UOp.variable("sp", 0, 7)
r0 = f(buf, x, v.bind(0)).numpy()
np.testing.assert_equal(r0, [[1.,0.,0.,0.,0.,0.,0.,0.], [2.,0.,0.,0.,0.,0.,0.,0.]])
@unittest.expectedFailure
def test_assign_slice(self):
@function
test/unit/test_gguf.py
+2 -1
View File
@@ -148,7 +148,8 @@ class TestGGUFGEMV(unittest.TestCase):
x = rng.standard_normal(cols).astype(np.float32)
np.testing.assert_allclose((tensors["weight"] @ Tensor(x)).numpy(), ref @ x, atol=1e-2, rtol=1e-2)
np.testing.assert_equal(tensors["weight"].numpy(), ref)
# can only expect the weights to be identical if we really support float16 (ie. not decompositions)
if is_dtype_supported(dtypes.half): np.testing.assert_equal(tensors["weight"].numpy(), ref)
assert np.isfinite(ref).all() and np.isfinite(tensors["weight"].numpy()).all(), f"{qtype.name} has NaN/Inf"
def test_gguf_gemv_q8_0(self): self._test_gguf_gemv(GGMLQuantizationType.Q8_0)
test/unit/test_gradient.py
+2 -2
View File
@@ -12,10 +12,10 @@ class TestTensorGradient(unittest.TestCase):
self.assertListEqual(dx.tolist(), [[2.0, 2.0, 2.0], [0.0, 0.0, 0.0], [-2.0, -2.0, -2.0]])
self.assertListEqual(dy.tolist(), [[1.0, 1.0, 1.0]])
def test_raises(self):
def test_zero_if_not_used(self):
x = Tensor([1.0, 2.0, 3.0])
w = Tensor.randn((3,))
with self.assertRaises(RuntimeError): x.sum().gradient(w)
self.assertListEqual(x.sum().gradient(w)[0].tolist(), [0.0, 0.0, 0.0])
def test_with_custom_gradient(self):
x = Tensor([1.0, 2.0, 3.0])
test/unit/test_shm_tensor.py
+1 -1
View File
@@ -1,7 +1,7 @@
import unittest
import multiprocessing.shared_memory as shared_memory
from tinygrad.helpers import CI, WIN
from tinygrad.tensor import Tensor, Device
from tinygrad import Tensor, Device
import numpy as np
class TestRawShmBuffer(unittest.TestCase):
tinygrad/apps/llm.py
+16 -9
View File
@@ -118,7 +118,7 @@ class TransformerBlock:
self.ffn_up = nn.Linear(dim, hidden_dim, bias=False)
self.ffn_down = nn.Linear(hidden_dim, dim, bias=False)
@function(precompile=bool(getenv("PRECOMPILE", 0)), allow_implicit=False)
@function
def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor:
x_norm = self.attn_norm(x) # (B,T,D)
q, k, v = self.attn_q(x_norm), self.attn_k(x_norm), self.attn_v(x_norm)
@@ -133,9 +133,14 @@ class TransformerBlock:
q = apply_rope(q, self.freqs_cis[start_pos:start_pos+T])
k = apply_rope(k, self.freqs_cis[start_pos:start_pos+T])
self.cache_kv[:, :, :, start_pos:start_pos+T, :].assign(Tensor.stack(k, v))
k = self.cache_kv[0, :, :, 0:start_pos+T, :]
v = self.cache_kv[1, :, :, 0:start_pos+T, :]
# NOTE: we don't want to change self.cache_kv, the function API doesn't support this well
assigned_kv = Tensor(self.cache_kv.uop.after(self.cache_kv[:, :, :, start_pos:start_pos+T, :].uop.store(Tensor.stack(k, v).uop)))
k = assigned_kv[0, :, :, 0:start_pos+T, :]
v = assigned_kv[1, :, :, 0:start_pos+T, :]
#self.cache_kv[:, :, :, start_pos:start_pos+T, :].assign(Tensor.stack(k, v))
#k = self.cache_kv[0, :, :, 0:start_pos+T, :]
#v = self.cache_kv[1, :, :, 0:start_pos+T, :]
# NOTE: this mask is causal_lower_right, not the causal_upper_left generated by is_casual = True
# TODO: this if statement should be removed and it shouldn't generate extra kernels
@@ -145,7 +150,7 @@ class TransformerBlock:
attn = self.attn_output(attn)
return x + attn
@function(precompile=bool(getenv("PRECOMPILE", 0)), allow_implicit=False)
@function
def _feed_forward(self, h: Tensor) -> Tensor:
h_norm = self.ffn_norm(h)
if hasattr(self, 'ffn_gate_exps'):
@@ -161,10 +166,12 @@ class TransformerBlock:
def __call__(self, x: Tensor, start_pos: int|UOp):
if not hasattr(self, "cache_kv"):
# TODO: how is the dtype of this determined?
# NOTE: clone is used to promise the creation of a specific buffer
self.cache_kv = Tensor.zeros(2, x.shape[0], self.n_kv_heads, self.max_context, self.head_dim, device=x.device).clone()
self.cache_kv = Tensor.empty(2, x.shape[0], self.n_kv_heads, self.max_context, self.head_dim, device=x.device)
self.freqs_cis = precompute_freqs_cis(self.head_dim, self.max_context, self.rope_theta)
return self._feed_forward(self._attention(x, start_pos)).contiguous()
# we pass in the weights implicitly so we unpack the GGUF on the fly
@function(precompile=True, allow_implicit=True)
def _run(x:Tensor, start_pos:int|UOp): return self._feed_forward(self._attention(x, start_pos)).contiguous()
return _run(x, start_pos)
class Transformer:
def __init__(self, *, num_blocks, dim, hidden_dim, n_heads, n_kv_heads, norm_eps, vocab_size, head_dim:int, rope_theta:float,
@@ -181,7 +188,7 @@ class Transformer:
self.rollout_jit = TinyJit(self.forward)
def forward(self, tokens:Tensor, start_pos:int|UOp) -> Tensor:
x = self.token_embd(tokens) # (B, T, D)
x = self.token_embd(tokens).float() # (B, T, D)
for block in self.blk: x = block(x, start_pos)
# TODO: add temperature
return self.output(self.output_norm(x))[:, -1, :].softmax(-1, dtype="float").argmax(-1, keepdim=True)
tinygrad/codegen/__init__.py
+5 -10
View File
@@ -1,19 +1,18 @@
from typing import cast
from dataclasses import replace
import itertools
from tinygrad.helpers import DISABLE_FAST_IDIV, EMULATED_DTYPES, DEVECTORIZE, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, TracingKey, Context
from tinygrad.helpers import DISABLE_FAST_IDIV, DEVECTORIZE, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, TracingKey, Context
from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, pyrender
from tinygrad.uop.spec import type_verify, program_spec, kernel_spec
from tinygrad.renderer import Renderer, ProgramSpec, Estimates
from tinygrad.dtype import dtypes, promo_lattice
from tinygrad.device import is_dtype_supported
from tinygrad.dtype import dtypes
from tinygrad.helpers import panic
from tinygrad.codegen.opt import Opt
# import all pattern matchers here
from tinygrad.codegen.gpudims import pm_add_gpudims
from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing, symbolic, pm_move_where_on_load
from tinygrad.uop.decompositions import get_late_rewrite_patterns, get_transcendental_patterns, pm_float_decomp, pm_long_decomp
from tinygrad.uop.decompositions import get_late_rewrite_patterns, get_transcendental_patterns, pm_dtype_decomps
from tinygrad.codegen.late.expander import expander, pm_pre_expander, pm_group_for_reduce
from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \
ReduceContext, correct_load_store, pm_render, pm_add_loads
@@ -42,7 +41,7 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -
sink = graph_rewrite(sink, pm_split_ranges+pm_flatten_range, ctx={}, name="split ranges")
# create image buffers
if IMAGE == 1 and ren.device in {"QCOM", "CL"}: sink = graph_rewrite(sink, pm_make_images, name="create image buffers", bottom_up=True)
if IMAGE and ren.device in {"QCOM", "CL", "PYTHON"}: sink = graph_rewrite(sink, pm_make_images, name="create image buffers", bottom_up=True)
# symbolic (NOTE: this is a requirement for pm_simplify_ranges to be correct)
sink = graph_rewrite(sink, sym+pm_flatten_range, name="initial symbolic")
@@ -92,11 +91,7 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -
pm_decomp = symbolic_simple+get_late_rewrite_patterns(supported_ops, ren.device, bool(DISABLE_FAST_IDIV))
pm_transcendental = symbolic_simple+get_transcendental_patterns(supported_ops, TRANSCENDENTAL>=2)
sink = graph_rewrite(sink, pm_decomp, ctx=ren.device, name="decompositions")
if not is_dtype_supported(dtypes.long, ren.device) or dtypes.long in EMULATED_DTYPES.tolist(dtypes):
sink = graph_rewrite(sink, pm_long_decomp, name="decomp long -> int", bottom_up=True)
for fr, to in [(fr, next((to for to in promo_lattice[fr] if is_dtype_supported(to, ren.device)), dtypes.float))
for fr in EMULATED_DTYPES.tolist(dtypes) if fr in dtypes.floats]:
sink = graph_rewrite(sink, pm_float_decomp, ctx=(fr, to), name=f"decomp {fr} -> {to}", bottom_up=True)
sink = graph_rewrite(sink, pm_dtype_decomps, ctx=(set(), ren.device, getattr(ren, "arch", "")), name="decomp dtypes")
sink = graph_rewrite(sink, pm_transcendental, ctx=ren.device, name="transcendental")
# final rules for the renderer (without sym)
tinygrad/codegen/gpudims.py
+2 -2
View File
@@ -35,7 +35,7 @@ def get_grouped_dims(prefix, dims:tuple[sint, ...], max_sizes:tuple[int, ...]|No
if len(limited) > len(max_sizes): raise RuntimeError(f"cannot limit dim {dims=}, {max_sizes=}")
# try to split up dims: (a,) -> (b, c)
if limited == dims: limited = _split_dims(dims, max_sizes)
raw_idxs = [UOp(Ops.SPECIAL, dtypes.index, (sint_to_uop(s),), (f"{prefix}{i}")) for i,s in enumerate(limited)]
raw_idxs = [UOp(Ops.SPECIAL, dtypes.weakint, (sint_to_uop(s),), (f"{prefix}{i}")) for i,s in enumerate(limited)]
if len(limited) < len(dims):
ret = []
if (contraction:=get_contraction(dims, limited)) is None: raise RuntimeError(f"get_contraction should not be None {dims=} {limited=}")
@@ -75,7 +75,7 @@ def add_gpudims(ctx:Renderer, s:UOp):
# get the idxs
ki: KernelInfo = s.arg
if ctx.has_threads: idxs = [UOp.variable("core_id", 0, int(global_shape[0])-1, dtypes.int).cast(dtypes.index)]
if ctx.has_threads: idxs = [UOp.variable("core_id", 0, int(global_shape[0])-1, dtypes.int).cast(dtypes.weakint)]
elif ki.dont_use_locals:
assert not local_dims, "can't use locals if there's no local dims"
idxs = get_grouped_dims("idx", global_shape, ctx.global_max, reverse=True)
tinygrad/codegen/late/devectorizer.py
+10 -20
View File
@@ -5,7 +5,7 @@ from dataclasses import dataclass
from tinygrad.dtype import dtypes, ImageDType, DType, AddrSpace, Invalid, PtrDType
from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, GroupOp, identity_element
from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate
from tinygrad.helpers import getenv, flatten, AMX, prod, IMAGE
from tinygrad.helpers import getenv, flatten, AMX, prod
from tinygrad.renderer import Renderer
# ***** image load valid simplification *****
@@ -63,7 +63,7 @@ load_store_indexing = PatternMatcher([
def expand_index(buf:UOp, vec:UOp):
# determine optimal image shapes
if IMAGE == 1 and isinstance(dt:=buf.dtype, ImageDType):
if isinstance(dt:=buf.dtype, ImageDType):
x, valid = vec.get_idx().gep(0), vec.get_valid().gep(0)
# search for dims that drop the most valid statements
best_drop, cands = -1, []
@@ -197,29 +197,19 @@ def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp):
return UOp(Ops.VCAT, ls.dtype, tuple(ret)) if ls.op is Ops.LOAD else UOp.group(*ret)
def get_image_idx(idx:UOp, width:int):
oidx = UOp(Ops.VECTORIZE, dtypes.index.vec(2), (((x:=idx.src[1].get_idx()) // 4) % width, (x // (4*width))))
oidx = UOp(Ops.VECTORIZE, dtypes.weakint.vec(2), (((x:=idx.src[1].get_idx()) // 4) % width, (x // (4*width))))
return idx.replace(src=(idx.src[0], oidx.valid(idx.src[1].get_valid())))
def image_fixup(ls:UOp):
# normal image load or store, with the CAST from expand_index
if ls.src[0].op is Ops.CAST and isinstance(image_dtype:=ls.src[0].src[0].dtype, ImageDType):
if isinstance(dt:=ls.src[0].src[0].dtype, ImageDType) and ls.src[0].op is Ops.CAST:
assert ls.src[0].dtype.count == 4, "image must be casted to 4"
idx = get_image_idx(ls.src[0].src[0], image_dtype.shape[1])
return ls.replace(src=(idx,)+ls.src[1:])
return ls.replace(src=(get_image_idx(ls.src[0].src[0], dt.shape[1]),)+ls.src[1:])
# this is an unprocessed image without a cast, aka unfoldable image load. this doesn't work for stores
if isinstance(image_dtype:=ls.src[0].dtype, ImageDType) and ls.src[0].src[1].get_idx().dtype != dtypes.index.vec(2):
assert ls.op is Ops.LOAD, "if an image store isn't upcasted to 4, we can't store it"
x, idx = ls.src[0].src[1].get_idx(), get_image_idx(ls.src[0], image_dtype.shape[1])
vec_load = ls.replace(dtype=ls.dtype.vec(4), src=(idx,)+ls.src[1:])
# image pixels have 4 channels (.xyzw), select channel based on x % 4
x_mod_4 = x % 4
def sel(ret, i): return x_mod_4.ne(i).where(ret, vec_load.gep(i))
# if x is non-negative, x % 4 is in [0, 3] and we can skip NAN fallback
if x_mod_4.vmin >= 0: return functools.reduce(sel, range(int(x_mod_4.vmin)+1, int(x_mod_4.vmax)+1), vec_load.gep(int(x_mod_4.vmin)))
return functools.reduce(sel, range(4), ls.const_like(float('nan')))
return None
# this is an unprocessed image without a cast, we should just make it a buffer
if isinstance(dt, ImageDType) and (off:=ls.src[0].src[1]).get_idx().dtype != dtypes.weakint.vec(2):
idx = ls.src[0].src[0].replace(dtype=(new_dt:=dtypes.half if dt.itemsize == 2 else dtypes.float).ptr(dt.size)).index(off)
return ls.replace(src=(idx,), dtype=new_dt).cast(dtypes.float) if ls.op is Ops.LOAD else ls.replace(src=(idx, ls[1].cast(dtypes.float)))
correct_load_store = PatternMatcher([
# split LOAD/STORE
@@ -263,7 +253,7 @@ def no_vectorized_index(buf:UOp, cast:UOp, idx:UOp, bcast:UOp|None=None):
# simple scalar index: one lane, all components
pairs = [(0, c) for c in range(cnt)]
idx_lanes, offsets = (tuple(x) for x in zip(*pairs))
return buf.broadcast(len(pairs)).index(idx.gep(idx_lanes)*cnt + UOp.const(dtypes.index.vec(len(pairs)), offsets), ptr=True)
return buf.broadcast(len(pairs)).index(idx.gep(idx_lanes)*cnt + UOp.const(dtypes.weakint.vec(len(pairs)), offsets), ptr=True)
devectorize_buf_and_index = PatternMatcher([
(UPat((Ops.DEFINE_LOCAL, Ops.DEFINE_REG), name="buf"), no_vectorized_buf),
tinygrad/codegen/late/expander.py
+10
View File
@@ -56,6 +56,16 @@ def do_expand(root:UOp):
# repeat the arg
new_srcs.append(src.broadcast(expand_sz))
# for non-PtrDType INDEX on REG buffers, expand into individual scalar INDEXes instead of one vectorized INDEX
# this avoids creating a VECTORIZE of REG pointers which the devectorizer can't resolve
if root.op is Ops.INDEX and not isinstance(root.dtype, PtrDType) and \
isinstance(root.src[0].dtype, PtrDType) and root.src[0].dtype.addrspace == AddrSpace.REG:
idxs = []
for j in range(expand_sz):
idx_srcs = tuple(s.gep(j) if isinstance(s.dtype, PtrDType) or s.dtype.count > 1 else s for s in new_srcs)
idxs.append(UOp(Ops.INDEX, root.dtype, idx_srcs, root.arg))
return UOp(Ops.UNROLL, root.dtype, (UOp(Ops.VECTORIZE, root.dtype.vec(expand_sz), tuple(idxs)),), expand_args)
new_arg = root.arg
if root.op is Ops.GEP:
assert root.dtype.count == 1
tinygrad/codegen/opt/postrange.py
+1 -1
View File
@@ -332,7 +332,7 @@ class Scheduler:
def bufs_from_ast(ast:UOp, dname:str) -> list[Buffer]:
glbls = sorted([x for x in ast.backward_slice if x.op is Ops.PARAM], key=lambda x: x.arg)
return [Buffer(dname, x.ptrdtype.size, x.dtype.base if not isinstance(x.dtype, ImageDType) else x.dtype) for x in glbls]
return [Buffer(dname, x.ptrdtype.size, x.dtype.base) for x in glbls]
def apply_opts(ast:UOp, ren:Renderer) -> UOp:
if ast.tag is not None: return ast
tinygrad/codegen/simplify.py
+3 -3
View File
@@ -1,6 +1,6 @@
import itertools
from typing import Callable
from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, graph_rewrite, _substitute, range_start
from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, graph_rewrite, _substitute, range_start, AxisType
from tinygrad.uop.symbolic import symbolic
from tinygrad.helpers import partition
from tinygrad.dtype import dtypes, ImageDType
@@ -57,7 +57,7 @@ pm_simplify_ranges = PatternMatcher([
])
def mark_range_mod(ctx:dict[UOp, UOp|None], r:UOp, c:UOp) -> None:
if r not in ctx and r.src[0].op is Ops.CONST and r.src[0].divides(c.arg) is not None: ctx[r] = c
if r not in ctx and r.arg[-1] is not AxisType.WARP and r.src[0].op is Ops.CONST and r.src[0].divides(c.arg) is not None: ctx[r] = c
def do_substitute(ctx:dict, x: UOp, sub_fxn:Callable[[UOp, UOp], UOp]) -> UOp|None:
ret = x.substitute({k:sub_fxn(k,v) for k,v in ctx.items() if v is not None})
@@ -157,5 +157,5 @@ def no_load(u:UOp) -> bool: return not any(x.op is Ops.INDEX for x in u.backward
pm_load_collapse = PatternMatcher([
(UPat(Ops.REDUCE, arg=Ops.ADD, src=(UPat.var("u"), UPat()), name="red"), reduce_load_collapse),
# we want to make sure we dont do math on a loaded index since that can cause overflow, this undoes the rule in pm_reduce_load_collapse
((UPat.var("x", dtypes.index)+UPat.var("y"))<UPat.var("c"), lambda x,y,c: x < c-y if no_load(y) and no_load(c) and not no_load(x) else None),
((UPat.var("x", dtypes.weakint)+UPat.var("y"))<UPat.var("c"), lambda x,y,c: x < c-y if no_load(y) and no_load(c) and not no_load(x) else None),
])
tinygrad/device.py
+20 -19
View File
@@ -3,11 +3,11 @@ from dataclasses import dataclass, replace
from collections import defaultdict
from typing import Any, Generic, TypeVar, Iterator, Generator, TYPE_CHECKING
import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re, atexit, pickle, decimal
from tinygrad.helpers import CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored
from tinygrad.helpers import BENCHMARKS, CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored
from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, dedup, ContextVar
from tinygrad.helpers import unwrap_class_type, suppress_finalizing, select_first_inited, VIZ, CPU_LLVM, CPU_LVP, CPU_X86, NV_PTX, CUDA_PTX, NV_NAK
from tinygrad.helpers import EMULATED_DTYPES, NULL_IR3, NULL_QCOMCL, TracingKey, size_to_str
from tinygrad.dtype import DType, ImageDType, PtrDType, dtypes, _to_np_dtype
from tinygrad.helpers import EMULATE, EMULATED_DTYPES, NULL_IR3, NULL_QCOMCL, IMAGE, FLOAT16, TracingKey, size_to_str
from tinygrad.dtype import DType, PtrDType, dtypes, _to_np_dtype
if TYPE_CHECKING: from tinygrad.renderer import Renderer
# **************** Device ****************
@@ -72,7 +72,6 @@ class ProfileGraphEvent(ProfileEvent): ents:list[ProfileGraphEntry]; deps:list[l
@dataclass(frozen=True, eq=True)
class BufferSpec:
# TODO: move device, size, dtype here?
image: ImageDType|None = None
uncached: bool = False
cpu_access: bool = False
host: bool = False
@@ -96,8 +95,7 @@ class Buffer:
profile_events:list[ProfileEvent] = []
def __init__(self, device:str, size:int, dtype:DType, opaque:Any=None, options:BufferSpec|None=None, initial_value:bytes|None=None,
uop_refcount=0, base:Buffer|None=None, offset:int=0, preallocate=False):
if isinstance(dtype, ImageDType): options = BufferSpec(image=dtype) # TODO: image hack shouldn't be here. where should it be?
else: assert isinstance(dtype, DType) and not isinstance(dtype, PtrDType)
assert isinstance(dtype, DType) and not isinstance(dtype, PtrDType)
self.device, self.size, self.dtype, self.options, self.offset, self.allocated_views = device, size, dtype, options, offset, 0
if base is None:
assert offset == 0, "base buffers can't have offset"
@@ -177,7 +175,7 @@ class Buffer:
(f" offset:{self.offset}" if self._base is not None else "") + (f" {self.options=}" if self.options is not None else "") + ">"
def as_memoryview(self, allow_zero_copy=False, force_zero_copy=False) -> memoryview:
# zero copy with as_memoryview (disabled by default due to use after free)
if (force_zero_copy or allow_zero_copy) and hasattr(self.allocator, '_as_buffer') and (self.options is None or self.options.image is None):
if (force_zero_copy or allow_zero_copy) and hasattr(self.allocator, '_as_buffer') and self.options is None:
return self.allocator._as_buffer(self._buf)
assert not force_zero_copy, "force zero copy was passed, but copy is required"
return self.copyout(memoryview(bytearray(self.nbytes)))
@@ -330,19 +328,22 @@ class Compiled:
# TODO: move this to each Device
# this only tracks if the dtype is natively supported, it may be supported in the frontend using decomps
def is_dtype_supported(dtype:DType, device:str|None=None) -> bool:
if dtype == dtypes.index: return False
def is_dtype_supported(dtype:DType, device:str|None=None, arch:str|None=None) -> bool:
if device is None: device = Device.DEFAULT
if dtype == dtypes.bfloat16:
if device == "METAL": return not CI
if device == "CUDA": return not CI and not CUDA_PTX
if device == "NV": return not CI and not NV_PTX and not NV_NAK
if device in {"CPU"}: return not CI and platform.machine() in {"arm", "arm64", "aarch64", "x86_64", "amd64"} and not CPU_LVP and not CPU_X86
if device == "METAL": return not CI or BENCHMARKS
if device == "CUDA": return (not CI or BENCHMARKS) and not CUDA_PTX
if device == "NV": return (not CI or BENCHMARKS) and not NV_PTX and not NV_NAK
if device in {"CPU"}:
return (not CI or BENCHMARKS) and platform.machine() in {"arm", "arm64", "aarch64", "x86_64", "amd64"} and not CPU_LVP and not CPU_X86
return device in {"AMD", "CL", "PYTHON", "NULL"}
if dtype in dtypes.fp8_ocp:
if device == "CUDA": return not CI and not CUDA_PTX
if device == "NV": return not CI and not NV_PTX and not NV_NAK
if device == "AMD": return not CI and getattr(Device["AMD"], "target") == (9,5,0)
if device == "CUDA": return (not CI or BENCHMARKS) and not CUDA_PTX
if device == "NV": return (not CI or BENCHMARKS) and not NV_PTX and not NV_NAK
if device == "AMD":
# TODO: open the device to get arch of device, will be fixed after triple is in the device string
if arch is None: arch = getattr(Device[device].renderer, "arch", "")
return (not CI or BENCHMARKS) and arch == "gfx950"
return device in {"PYTHON", "NULL"}
if dtype in dtypes.fp8_fnuz: return device in {"PYTHON", "NULL"}
if device == "WEBGPU": return dtype in [dtypes.bool, dtypes.char, dtypes.uchar, dtypes.short,
@@ -353,9 +354,9 @@ def is_dtype_supported(dtype:DType, device:str|None=None) -> bool:
# PYTHON supports half memoryview in 3.12+ https://github.com/python/cpython/issues/90751
# double can't be bitcast to anything without long support
if dtype == dtypes.half:
if device == "CL": return not CI and not OSX
if device == "QCOM": return False # QCOM compiler is flaky with half
if device in ["CUDA", "NV"]: return not CI
if device == "CL": return (not CI or BENCHMARKS) and not OSX
if device == "QCOM": return bool(IMAGE) and bool(FLOAT16) # QCOM compiler is flaky with half
if device in ["CUDA", "NV"]: return (not CI or BENCHMARKS) or "CUDA" in EMULATE.value
if device == "CPU" and CPU_LLVM: return OSX
if device == "PYTHON": return sys.version_info >= (3, 12)
if dtype == dtypes.float64: return (device not in {"METAL", "QCOM"} and not (OSX and device == "CL") and not NULL_IR3 and not NULL_QCOMCL
tinygrad/dtype.py
+9 -8
View File
@@ -150,7 +150,7 @@ class ImageDType(PtrDType):
ALIGN, MAXW, pxls = getenv("IMAGE_PITCH_ALIGN", 256 if OSX else 64), 16384, ptr.size // 4
if ptr.base not in (dtypes.half, dtypes.float) or ptr.size > 4*MAXW*MAXW: return []
# OSX has stricter requirements for height=1 images
if ptr.size % (ALIGN * 4) != 0: return [] if OSX or ptr.nbytes() % getenv("IMAGE_BASE_ALIGN", 64) != 0 else [(1, pxls)]
if ptr.size % (ALIGN * 4) != 0: return [] if OSX or ptr.nbytes() % getenv("IMAGE_BASE_ALIGN", 64) != 0 or pxls > MAXW else [(1, pxls)]
return [(pxls//ALIGN//k, ALIGN*k) for k in range(ceildiv(pxls//ALIGN, MAXW), min(pxls//ALIGN, MAXW//ALIGN)+1) if (pxls//ALIGN)%k == 0]
class dtypes:
@@ -159,7 +159,7 @@ class dtypes:
def is_float(x: DType) -> bool: return x.scalar() in dtypes.floats or isinstance(x, ImageDType)
@staticmethod # static methods on top, or bool in the type info will refer to dtypes.bool
@functools.cache
def is_int(x: DType) -> bool: return x.scalar() in (dtypes.ints + (dtypes.index,))
def is_int(x: DType) -> bool: return x.scalar() in (dtypes.ints + (dtypes.weakint,))
@staticmethod
@functools.cache
def is_unsigned(x: DType) -> bool: return x.scalar() in dtypes.uints
@@ -180,7 +180,7 @@ class dtypes:
return {dtypes.float16: (5, 10), dtypes.bfloat16: (8, 7), dtypes.float32: (8, 23), dtypes.float64: (11, 52),
dtypes.fp8e4m3: (4, 3), dtypes.fp8e5m2: (5, 2), dtypes.fp8e4m3fnuz: (4, 3), dtypes.fp8e5m2fnuz: (5, 2)}[dtype]
void: Final[DType] = DType.new(-1, 0, "void", None)
index: Final[DType] = DType.new(-1, 800, "index", None)
weakint: Final[DType] = DType.new(-1, 800, "weakint", None)
bool: Final[DType] = DType.new(0, 1, "bool", '?')
int8: Final[DType] = DType.new(1, 8, "signed char", 'b')
uint8: Final[DType] = DType.new(2, 8, "unsigned char", 'B')
@@ -227,7 +227,7 @@ class dtypes:
uints = (uint8, uint16, uint32, uint64)
sints = (int8, int16, int32, int64)
ints = uints + sints
all = floats + ints + (bool, index) # noqa: A003
all = floats + ints + (bool, weakint) # noqa: A003
if (env_default_float := getenv("DEFAULT_FLOAT", "")):
dtypes.default_float = getattr(dtypes, env_default_float.lower())
@@ -237,7 +237,8 @@ DTypeLike = str|DType
def to_dtype(dtype:DTypeLike) -> DType: return dtype if isinstance(dtype, DType) else getattr(dtypes, dtype.lower())
# https://jax.readthedocs.io/en/latest/jep/9407-type-promotion.html
# we don't support weak type and complex type
# we don't support complex type
# TODO: weakint and weakfloat in lattice
promo_lattice = { dtypes.bool: [dtypes.int8, dtypes.uint8], dtypes.int8: [dtypes.int16], dtypes.int16: [dtypes.int32], dtypes.int32: [dtypes.int64],
dtypes.int64: [dtypes.uint64], dtypes.uint8: [dtypes.int16, dtypes.uint16], dtypes.uint16: [dtypes.int32, dtypes.uint32],
dtypes.uint32: [dtypes.int64, dtypes.uint64], dtypes.uint64: [dtypes.fp8e4m3, dtypes.fp8e5m2, dtypes.fp8e4m3fnuz, dtypes.fp8e5m2fnuz],
@@ -254,8 +255,8 @@ def least_upper_dtype(*ds:DType) -> DType:
if not (images:=[d for d in ds if isinstance(d, ImageDType)]) else images[0]
def least_upper_float(dt:DType) -> DType: return dt if dtypes.is_float(dt) else least_upper_dtype(dt, dtypes.default_float)
DTYPES_DICT = {k: v for k, v in dtypes.__dict__.items() if isinstance(v, DType) and not k.startswith(("default", "void", "index", "_"))}
INVERSE_DTYPES_DICT = {**{v.name:k for k,v in DTYPES_DICT.items()}, "void": "void", "index":"index"}
DTYPES_DICT = {k: v for k, v in dtypes.__dict__.items() if isinstance(v, DType) and not k.startswith(("default", "void", "weakint", "_"))}
INVERSE_DTYPES_DICT = {**{v.name:k for k,v in DTYPES_DICT.items()}, "void": "void", "weakint":"weakint"}
@functools.cache
def can_lossless_cast(dt0:DType, dt1:DType) -> bool:
@@ -263,7 +264,7 @@ def can_lossless_cast(dt0:DType, dt1:DType) -> bool:
# similar to https://numpy.org/doc/stable/reference/generated/numpy.can_cast.html
if dt0 == dt1 or dt0 == dtypes.bool: return True
match dt1:
case dtypes.index: return dt0 in dtypes.ints
case dtypes.weakint: return dt0 in dtypes.ints
case dtypes.double: return dt0 in (dtypes.float, dtypes.half, dtypes.bfloat16, *dtypes.fp8s,
dtypes.uint32, dtypes.uint16, dtypes.uint8, dtypes.int32, dtypes.int16, dtypes.int8)
case dtypes.float: return dt0 in (dtypes.half, dtypes.bfloat16, *dtypes.fp8s, dtypes.uint16, dtypes.uint8, dtypes.int16, dtypes.int8)
tinygrad/engine/allocations.py
+2 -8
View File
@@ -1,7 +1,6 @@
from dataclasses import dataclass, field
from tinygrad.uop.ops import UOp, UPat, PatternMatcher, Ops, GroupOp, graph_rewrite, track_rewrites
from tinygrad.dtype import ImageDType
from tinygrad.helpers import prod, DEBUG, VIZ, pluralize, all_int
from tinygrad.helpers import VIZ, pluralize, all_int
@dataclass
class AllocCtx:
@@ -42,12 +41,7 @@ add_tags = PatternMatcher([
])
def _buffer_like(u:UOp) -> UOp:
dtype = u.dtype
if isinstance(dtype, ImageDType):
if prod(dtype.shape) != prod(u.max_shard_shape) or ([x for x in u.max_shard_shape if x != 1] or [1])[-1] % 4 != 0:
if DEBUG >= 1: print(f"demoting Image {dtype} with shape {u.max_shard_shape}")
dtype = dtype.base
buffer = UOp.new_buffer(u.device, u.shard_size, dtype).reshape(u.max_shard_shape).shrink_to(u.shard_shape)
buffer = UOp.new_buffer(u.device, u.shard_size, u.dtype).reshape(u.max_shard_shape).shrink_to(u.shard_shape)
if isinstance(u.device, tuple) and u.axis is not None: buffer = buffer.multi(u.axis)
return buffer
tinygrad/engine/jit.py
+13 -21
View File
@@ -7,10 +7,10 @@ from tinygrad.dtype import DType
from tinygrad.uop.ops import UOp, Variable, sym_infer, Ops
from tinygrad.engine.realize import ExecItem, capturing, ViewOp, BufferCopy, BufferXfer, EncDec, CompiledRunner, Runner, Estimates
from tinygrad.engine.memory import _internal_memory_planner
from tinygrad.engine.schedule import linear_to_schedule
from tinygrad.nn.state import get_parameters
from tinygrad.schedule.rangeify import mop_cleanup
from dataclasses import dataclass, replace
from weakref import WeakKeyDictionary
class GraphException(Exception): pass
class JitError(Exception): pass
@@ -280,17 +280,7 @@ class TinyJit(Generic[ReturnType]):
self.prune = prune
self.optimize = optimize
def add_buffer(self, b:Buffer) -> Buffer:
if found:=self._buffer_replace.get(b, None): return found
if b.is_allocated() or b.uop_refcount > 0: return b
if b._base is not None:
self._buffer_replace[b] = ret = Buffer(b.device, b.size, b.dtype, base=self.add_buffer(b._base), offset=b.offset)
else:
self._buffer_replace[b] = ret = Buffer(b.device, b.size, b.dtype, options=b.options)
return ret
def add(self, ei:ExecItem):
self._jit_cache.append(ExecItem(ei.ast, [self.add_buffer(buf) for buf in ei.bufs if buf is not None], ei.metadata, ei.fixedvars, ei.prg))
def add_linear(self, linear:UOp, var_vals:dict[str, int]): self._linears.append(linear)
def reset(self):
assert self.fxn is not None, "can't reset without function"
@@ -321,20 +311,20 @@ class TinyJit(Generic[ReturnType]):
# jit capture
assert self.fxn is not None
if capturing: raise RuntimeError(f"having TinyJit inside another TinyJit is not supported {len(capturing)=} {capturing=}")
self._jit_cache: list[ExecItem] = []
self._buffer_replace: WeakKeyDictionary[Buffer, Buffer] = WeakKeyDictionary()
# TODO: should we always disable the memory planner here? it must be off for prune
with Context(BEAM=getenv("JITBEAM", BEAM.value), NO_MEMORY_PLANNER=int(self.prune)):
self._linears: list[UOp] = []
with Context(BEAM=getenv("JITBEAM", BEAM.value)):
capturing.append(self)
try:
ret = self.fxn(*args, **kwargs)
if len(params:=get_parameters(ret)): Tensor.realize(*params)
finally: capturing.clear()
jit_cache = self._jit_cache
del self._buffer_replace, self._jit_cache
if not len(jit_cache): raise JitError("didn't JIT anything!")
if not len(self._linears): raise JitError("didn't JIT anything!")
_check_no_non_tensor_return(ret)
if DEBUG >= 1: print(f"JIT captured {len(jit_cache)} kernels with {len(input_buffers)} inputs")
if DEBUG >= 1: print(f"JIT captured {len(self._linears)} linears with {len(input_buffers)} inputs")
# combine all captured linears into one and convert to ExecItems
jit_cache = [ei.lower() for ei in linear_to_schedule(UOp(Ops.LINEAR, src=tuple(flatten([l.src for l in self._linears]))))]
del self._linears
# track inputs that are views of buffers
# TODO: eventually expected_buffers should live in ExecItem
@@ -367,7 +357,9 @@ class TinyJit(Generic[ReturnType]):
input_replace = get_input_replace(jit_cache, input_buffers)
if DEBUG >= 1 and len(set(input_replace.values())) != len(input_buffers): print("WARNING: some input tensors not found")
# set this for next run
# exec
for ei in jit_cache: ei.run(var_vals)
self.captured = CapturedJit(ret, jit_cache, input_replace, extra_view_inputs, names, expected_input_info)
if self.optimize: self.captured.replan_buffers_memory_layout()
elif self.cnt >= 2:
tinygrad/engine/memory.py
+2 -2
View File
@@ -4,7 +4,7 @@ from tinygrad.engine.realize import ExecItem
from tinygrad.device import Device, Buffer
from tinygrad.helpers import NO_MEMORY_PLANNER, dedup, DEBUG, round_up
from tinygrad.uop.ops import Ops
from tinygrad.dtype import dtypes, ImageDType
from tinygrad.dtype import dtypes
from tinygrad.runtime.support.memory import TLSFAllocator
LaneKey = tuple[str, int]
@@ -39,7 +39,7 @@ def _internal_memory_planner(buffers:list[list[Buffer]], copies:list[tuple[Buffe
global_planner:dict[LaneKey, tuple[int, TLSFAllocator]] = defaultdict(lambda: (0, TLSFAllocator(total_memory, block_size=BLK, lv2_cnt=32)))
for (_, is_open_ev), buf in buffer_requests:
# Check if suballocation is possible for the given buffer and device.
if hasattr(Device[buf.device].allocator, "_offset") and not isinstance(buf.dtype, ImageDType):
if hasattr(Device[buf.device].allocator, "_offset"):
if is_open_ev: buffer_replace[buf] = (None, global_planner[_key(buf)][1].alloc(round_up(buf.nbytes, BLK)))
else: global_planner[_key(buf)][1].free(cast(int, buffer_replace[buf][1]))
global_planner[_key(buf)] = (max(global_planner[_key(buf)][0], buffer_replace[buf][1] + buf.nbytes), global_planner[_key(buf)][1])
tinygrad/engine/realize.py
+2 -3
View File
@@ -1,7 +1,7 @@
from typing import cast, Callable
import time, pprint, random, itertools, math
from dataclasses import dataclass, replace, field
from tinygrad.helpers import all_same, colored, DEBUG, GlobalCounters, ansilen, BEAM, NOOPT, all_int, CAPTURING, Metadata, TRACEMETA, TracingKey
from tinygrad.helpers import all_same, colored, DEBUG, GlobalCounters, ansilen, BEAM, NOOPT, all_int, Metadata, TRACEMETA, TracingKey
from tinygrad.helpers import DEVECTORIZE, time_to_str, VALIDATE_WITH_CPU, cpu_profile, PROFILE, ProfilePointEvent, cpu_events, prod, Context, unwrap
from tinygrad.helpers import EMULATED_DTYPES
from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, sym_infer
@@ -187,12 +187,11 @@ class ExecItem:
# **************** main run function ****************
capturing: list = [] # put classes with an add method in here
capturing: list = [] # put classes with an add_linear method in here
def run_schedule(schedule:list[ExecItem], var_vals:dict[str, int]|None=None, do_update_stats=True):
while len(schedule):
ei = schedule.pop(0).lower()
if len(capturing) and CAPTURING: capturing[0].add(ei)
if VALIDATE_WITH_CPU and ei.ast.op is Ops.SINK:
# copy in allocated buffers from the GPU
bufs = [b for b in ei.bufs if b is not None]
tinygrad/engine/schedule.py
+7
View File
@@ -84,7 +84,9 @@ def linear_to_schedule(linear:UOp) -> list[ExecItem]:
return schedule
from tinygrad.engine.memory import memory_planner
from tinygrad.engine.realize import capturing
from tinygrad.schedule.rangeify import get_kernel_graph
from tinygrad.helpers import CAPTURING
from tinygrad.uop.ops import PatternMatcher, UPat
def create_new_buffer(ctx:tuple[dict[UOp, UOp], tuple[UOp, ...]], b:UOp):
@@ -156,6 +158,11 @@ def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[list[ExecItem], di
if var_vals.get(nm, val) != val: raise RuntimeError(f"bind mismatch on {nm}, {var_vals[nm]} != {val}")
var_vals[nm] = val
# jit captures this schedule, no need to execute.
if len(capturing) and CAPTURING:
capturing[0].add_linear(linear, var_vals)
return [], var_vals
# convert LINEAR to ExecItems
schedule: list[ExecItem] = linear_to_schedule(linear)
with cpu_profile(TracingKey("memory planner")): schedule = memory_planner(schedule)
tinygrad/helpers.py
+1 -1
View File
@@ -14,7 +14,7 @@ def prod(x:Iterable[T]) -> T|int: return functools.reduce(operator.mul, x, 1)
# NOTE: helpers is not allowed to import from anything else in tinygrad
OSX, WIN = platform.system() == "Darwin", sys.platform == "win32"
CI = os.getenv("CI", "") != ""
CI, BENCHMARKS = os.getenv("CI", "") != "", os.getenv("RUNNER_ENVIRONMENT", "") == "self-hosted"
ARCH_X86 = any(x in platform.processor() for x in ("Intel", "i386", "x86_64"))
BASEDIR = pathlib.Path(__file__).parent
tinygrad/nn/__init__.py
+2 -2
View File
@@ -347,12 +347,12 @@ def _embedding_bwd(grad_emb:UOp, call:UOp) -> tuple:
# each device owns [offset, offset+local_vocab_size) of the global vocabulary
dnum = UOp.variable("_device_num", 0, ndev-1)
offset = dnum * local_vocab_size
global_token_id = idx_flat[i].cast(dtypes.index)
global_token_id = idx_flat[i].cast(dtypes.weakint)
local_token_id = (global_token_id - offset).clip(0, grad_weight.shape[0]-1)
in_range = (global_token_id >= offset) & (global_token_id < (offset + local_vocab_size))
grad_val = in_range.where(grad_emb_flat[i, j].cast(dtypes.float), 0.0)
else:
local_token_id = idx_flat[i].clip(0, grad_weight.shape[0]-1).cast(dtypes.index)
local_token_id = idx_flat[i].clip(0, grad_weight.shape[0]-1).cast(dtypes.weakint)
grad_val = grad_emb_flat[i, j].cast(dtypes.float)
# atomic scatter-add: grad_weight[token_id, j] += grad_emb_flat[i, j]
if device in ("CPU", "NULL"): atomic_arg = "__atomic_fetch_add({0}, {1}, __ATOMIC_RELAXED);"
tinygrad/renderer/amd/sqtt.py
+6 -4
View File
@@ -8,6 +8,7 @@ from __future__ import annotations
from dataclasses import dataclass
from typing import Iterator
from enum import Enum
from tinygrad.helpers import getenv, colored
from tinygrad.renderer.amd.dsl import BitField, FixedBitField, Inst, bits
from tinygrad.runtime.autogen.amd.rdna3.ins import s_endpgm # same encoding as RDNA4
@@ -678,7 +679,6 @@ PACKET_COLORS = {
}
def format_packet(p) -> str:
from tinygrad.helpers import colored
name = type(p).__name__
if isinstance(p, (INST, INST_RDNA4)):
op_name = p.op.name if isinstance(p.op, (InstOp, InstOpRDNA4)) else f"0x{p.op:02x}"
@@ -695,7 +695,6 @@ def format_packet(p) -> str:
return f"{p._time:8}: {colored(f'{name:18}', PACKET_COLORS.get(name.replace('_RDNA4', ''), 'white'))} {fields}"
def print_packets(packets) -> None:
from tinygrad.helpers import getenv
skip = {"NOP", "TS_DELTA_SHORT", "TS_WAVE_STATE", "TS_DELTA_OR_MARK",
"TS_DELTA_S5_W2", "TS_DELTA_S5_W3", "TS_DELTA_S8_W3", "REG", "EVENT"} if not getenv("NOSKIP") else {"NOP"}
for data in packets:
@@ -710,7 +709,10 @@ if __name__ == "__main__":
prg_events = {e.tag: e for e in data if type(e).__name__ == "ProfileProgramEvent" and e.tag is not None}
sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
dev_targets = {e.device:f"gfx{e.props['gfx_target_version']//1000}" for e in data if type(e).__name__ == "ProfileDeviceEvent" and e.props}
evt_num = getenv("SQTT_EVENT", -1)
for i, event in enumerate(sqtt_events):
prg = prg_events.get(event.kern)
print(f"\n=== event {i} {prg.name if prg is not None else ''} ===")
print_packets(map_insts(event.blob, prg.lib, dev_targets[prg.device]) if prg is not None else decode(event.blob))
print(f"=== event {i} {prg.name if prg is not None else ''} ===")
if evt_num == -1 or i == evt_num:
print_packets(map_insts(event.blob, prg.lib, dev_targets[prg.device]) if prg is not None else decode(event.blob))
print("\n")
tinygrad/runtime/graph/hcq.py
+1 -1
View File
@@ -23,7 +23,7 @@ class HCQGraph(MultiGraphRunner):
for (j,i), input_idx in self.input_replace.items():
x = self.input_replace_to_var.setdefault((j,i), UOp.variable(f"input_{input_idx}", 0, 0xffffffffffffffff, dtype=dtypes.uint64))
self.hcq_bufs[j][i] = HCQBuffer(x, self.hcq_bufs[j][i].size, image=self.hcq_bufs[j][i].image) # Create fake buffer with variable
self.hcq_bufs[j][i] = HCQBuffer(x, self.hcq_bufs[j][i].size) # Create fake buffer with variable
# Allocate kernel args.
kernargs_size: dict[Compiled, int] = collections.defaultdict(int)
tinygrad/runtime/ops_amd.py
+4 -5
View File
@@ -355,7 +355,7 @@ class AMDComputeQueue(HWQueue):
self.wreg(self.gc.regCOMPUTE_RESTART_X, 0, 0, 0)
self.wreg(self.gc.regCOMPUTE_USER_DATA_0, *user_regs)
self.wreg(self.gc.regCOMPUTE_RESOURCE_LIMITS, 0)
self.wreg(self.gc.regCOMPUTE_RESOURCE_LIMITS, waves_per_sh=getenv("WAVES_PER_SH"))
self.wreg(self.gc.regCOMPUTE_START_X, 0, 0, 0, *local_size, 0, 0)
gfx10p = {'cs_w32_en': int(prg.wave32)} if prg.dev.target >= (10,0,0) else {}
@@ -885,7 +885,7 @@ class PCIIface(PCIIfaceBase):
class USBIface(PCIIface):
def __init__(self, dev, dev_id): # pylint: disable=super-init-not-called
self.dev, self.pci_dev = dev, USBPCIDevice(dev.__class__.__name__[:2], f"usb:{dev_id}")
self.dev, self.pci_dev, self.vram_bar = dev, USBPCIDevice(dev.__class__.__name__[:2], f"usb:{dev_id}"), 0
self.dev_impl = AMDev(self.pci_dev)
self._compute_props()
self.pci_dev.usb._pci_cacheable += [self.pci_dev.bar_info(2)] # doorbell region is cacheable
@@ -903,9 +903,8 @@ class USBIface(PCIIface):
self.sys_next_off += size
return self.sys_buf.offset(self.sys_next_off - size, size)
mapping = self.dev_impl.mm.valloc(size:=round_up(size, 4 << 10), uncached=uncached, contiguous=cpu_access)
barview = self.pci_dev.map_bar(bar=0, off=mapping.paddrs[0][0], size=mapping.size) if cpu_access else None
return HCQBuffer(mapping.va_addr, size, meta=PCIAllocationMeta(mapping, has_cpu_mapping=False), view=barview, owner=self.dev)
# force devmem
return super().alloc(size, host=host, uncached=uncached, cpu_access=cpu_access, contiguous=contiguous, force_devmem=True, **kwargs)
def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0,
xcc_id=0, idx=0):
tinygrad/runtime/ops_cl.py
+10 -20
View File
@@ -53,10 +53,10 @@ class CLProgram:
try: check(cl.clReleaseProgram(self.program))
except (TypeError, AttributeError): pass
def __call__(self, *bufs:tuple[cl.cl_mem, BufferSpec], global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]|None=None,
vals:tuple[int, ...]=(), wait=False, **kw) -> float|None:
def __call__(self, *bufs:cl.cl_mem, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]|None=None, vals:tuple[int, ...]=(),
wait=False, **kw) -> float|None:
i = 0
for i,(b,_) in enumerate(bufs):
for i,b in enumerate(bufs):
for real_i, dt in self.arg_dtypes[i]:
if isinstance(dt, ImageDType):
fmt = cl.cl_image_format(cl.CL_RGBA, {2:cl.CL_HALF_FLOAT, 4:cl.CL_FLOAT}[dt.itemsize])
@@ -78,26 +78,16 @@ class CLProgram:
return None
class CLAllocator(LRUAllocator['CLDevice']):
def _alloc(self, size:int, options:BufferSpec) -> tuple[cl.cl_mem, BufferSpec]:
# Recalculate real size for texture
if options.image is not None: size = options.image.pitch * options.image.shape[0]
return (checked(cl.clCreateBuffer(self.dev.context, cl.CL_MEM_READ_WRITE, size, None, status := ctypes.c_int32()), status), options)
def _alloc(self, size:int, options:BufferSpec) -> cl.cl_mem:
return checked(cl.clCreateBuffer(self.dev.context, cl.CL_MEM_READ_WRITE, size, None, status := ctypes.c_int32()), status)
@suppress_finalizing
def _free(self, opaque:tuple[cl.cl_mem, BufferSpec], options:BufferSpec): check(cl.clReleaseMemObject(opaque[0]))
def _copyin(self, dest:tuple[cl.cl_mem, BufferSpec], src:memoryview):
def _free(self, opaque:cl.cl_mem, options:BufferSpec): check(cl.clReleaseMemObject(opaque))
def _copyin(self, dest:cl.cl_mem, src:memoryview):
if mv_address(src) % 16: src = memoryview(bytearray(src))
if (img:=dest[1].image):
stride = img.shape[1]*img.itemsize*4
for i in range(img.shape[0]):
check(cl.clEnqueueWriteBuffer(self.dev.queue, dest[0], False, i*img.pitch, stride, mv_address(src)+(i*stride), 0, None, None))
else: check(cl.clEnqueueWriteBuffer(self.dev.queue, dest[0], False, 0, len(src)*src.itemsize, from_mv(src), 0, None, None))
check(cl.clEnqueueWriteBuffer(self.dev.queue, dest, False, 0, len(src)*src.itemsize, from_mv(src), 0, None, None))
self.dev.pending_copyin.append(src) # NOTE: these can't be freed until the GPU actually executes this command
def _copyout(self, dest:memoryview, src:tuple[cl.cl_mem, BufferSpec]):
if (img:=src[1].image):
stride = img.shape[1]*img.itemsize*4
for i in range(img.shape[0]):
check(cl.clEnqueueReadBuffer(self.dev.queue, src[0], False, i*img.pitch, stride, mv_address(dest)+(i*stride), 0, None, None))
else: check(cl.clEnqueueReadBuffer(self.dev.queue, src[0], False, 0, len(dest)*dest.itemsize, from_mv(dest), 0, None, None))
def _copyout(self, dest:memoryview, src:cl.cl_mem):
check(cl.clEnqueueReadBuffer(self.dev.queue, src, False, 0, len(dest)*dest.itemsize, from_mv(dest), 0, None, None))
self.dev.synchronize()
class CLDevice(Compiled):
tinygrad/runtime/ops_qcom.py
+14 -20
View File
@@ -214,10 +214,11 @@ class QCOMArgsState(HCQArgsState):
for i, v in enumerate(vals): self.bind_sints_to_buf(v, buf=self.buf, fmt='I', offset=prg.buf_offs[i+len(ubos)])
def _tex(b, ibo=False):
fmt = mesa.FMT6_32_32_32_32_FLOAT if (img:=b[1].image or b[0]).itemsize == 4 else mesa.FMT6_16_16_16_16_FLOAT
imgdt, buf = b
fmt = mesa.FMT6_32_32_32_32_FLOAT if imgdt.itemsize == 4 else mesa.FMT6_16_16_16_16_FLOAT
return [qreg.a6xx_tex_const_0(fmt=fmt) if ibo else qreg.a6xx_tex_const_0(0x8, swiz_x=0, swiz_y=1, swiz_z=2, swiz_w=3, fmt=fmt),
qreg.a6xx_tex_const_1(width=img.shape[1], height=img.shape[0]),
qreg.a6xx_tex_const_2(type=mesa.A6XX_TEX_2D, pitch=img.pitch, pitchalign=ctz(img.pitch)-6), 0, *data64_le(b[1].va_addr),
qreg.a6xx_tex_const_1(width=imgdt.shape[1], height=imgdt.shape[0]),
qreg.a6xx_tex_const_2(type=mesa.A6XX_TEX_2D, pitch=imgdt.pitch, pitchalign=ctz(imgdt.pitch)-6), 0, *data64_le(buf.va_addr),
qreg.a6xx_tex_const_6(plane_pitch=0x400000), qreg.a6xx_tex_const_7(13), 0, 0, 0, 0, 0, 0, 0, 0]
self.bind_sints_to_buf(*flatten(map(_tex, texs)), buf=self.buf, fmt='I', offset=prg.tex_off)
@@ -322,25 +323,18 @@ class QCOMTextureInfo:
class QCOMAllocator(HCQAllocatorBase):
def _alloc(self, size:int, opts:BufferSpec) -> HCQBuffer:
# Recalculate real size for texture
if opts.image is not None: size = opts.image.pitch * opts.image.shape[0]
return self.dev._gpu_map(opts.external_ptr, size, image=opts.image) if opts.external_ptr else self.dev._gpu_alloc(size, image=opts.image)
return self.dev._gpu_map(opts.external_ptr, size) if opts.external_ptr else self.dev._gpu_alloc(size)
def _do_copy(self, src_addr, dest_addr, src_size, real_size, src_stride, dest_stride, prof_text, dest_off=0, src_off=0):
with cpu_profile(prof_text, f"{self.dev.device}:COPY"):
while src_off < src_size:
ctypes.memmove(dest_addr+dest_off, src_addr+src_off, real_size)
src_off, dest_off = src_off+src_stride, dest_off+dest_stride
def _do_copy(self, src_addr, dest_addr, size, prof_text):
with cpu_profile(prof_text, f"{self.dev.device}:COPY"): ctypes.memmove(dest_addr, src_addr, size)
def _copyin(self, dest:HCQBuffer, src:memoryview):
stride, pitch = (dest.image.shape[1] * 4 * dest.image.itemsize, dest.image.pitch) if dest.image else (src.nbytes, src.nbytes)
self._do_copy(mv_address(src), dest.cpu_view().addr, src.nbytes, stride, stride, pitch, f"TINY -> {self.dev.device}")
self._do_copy(mv_address(src), dest.cpu_view().addr, src.nbytes, f"TINY -> {self.dev.device}")
def _copyout(self, dest:memoryview, src:HCQBuffer):
self.dev.synchronize()
stride, pitch = (src.image.shape[1] * 4 * src.image.itemsize, src.image.pitch) if src.image else (src.size, src.size)
self._do_copy(src.cpu_view().addr, mv_address(dest), src.size, stride, pitch, stride, f"{self.dev.device} -> TINY")
self._do_copy(src.cpu_view().addr, mv_address(dest), src.size, f"{self.dev.device} -> TINY")
def _as_buffer(self, src:HCQBuffer) -> memoryview:
self.dev.synchronize()
@@ -386,7 +380,7 @@ class QCOMDevice(HCQCompiled):
super().__init__(device, QCOMAllocator(self), compilers, functools.partial(QCOMProgram, self), QCOMSignal,
functools.partial(QCOMComputeQueue, self), None)
def _gpu_alloc(self, size:int, flags:int=0, uncached=False, fill_zeroes=False, **kwargs) -> HCQBuffer:
def _gpu_alloc(self, size:int, flags:int=0, uncached=False, fill_zeroes=False) -> HCQBuffer:
flags |= flag("KGSL_MEMALIGN", alignment_hint:=12) | kgsl.KGSL_MEMFLAGS_USE_CPU_MAP
if uncached: flags |= flag("KGSL_CACHEMODE", kgsl.KGSL_CACHEMODE_UNCACHED)
@@ -394,16 +388,16 @@ class QCOMDevice(HCQCompiled):
va_addr = self.fd.mmap(0, bosz, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED, alloc.id * 0x1000)
if fill_zeroes: ctypes.memset(va_addr, 0, size)
return HCQBuffer(va_addr=va_addr, size=size, meta=(alloc, True), view=MMIOInterface(va_addr, size, fmt='B'), owner=self, **kwargs)
return HCQBuffer(va_addr=va_addr, size=size, meta=(alloc, True), view=MMIOInterface(va_addr, size, fmt='B'), owner=self)
def _gpu_map(self, ptr:int, size:int, **kwargs) -> HCQBuffer:
def _gpu_map(self, ptr:int, size:int) -> HCQBuffer:
ptr_aligned, size_aligned = (ptr & ~0xfff), round_up(size + (ptr & 0xfff), 0x1000)
dcache_flush().fxn(ctypes.c_uint64(ptr_line_aligned:=ptr & ~63), ctypes.c_uint64(ceildiv(ptr + size - ptr_line_aligned, 64)))
try:
mi = kgsl.IOCTL_KGSL_MAP_USER_MEM(self.fd, hostptr=ptr_aligned, len=size_aligned, memtype=kgsl.KGSL_USER_MEM_TYPE_ADDR)
return HCQBuffer(mi.gpuaddr + (ptr - ptr_aligned), size=size, meta=(mi, False), view=MMIOInterface(ptr, size, fmt='B'), owner=self, **kwargs)
return HCQBuffer(mi.gpuaddr + (ptr - ptr_aligned), size=size, meta=(mi, False), view=MMIOInterface(ptr, size, fmt='B'), owner=self)
except OSError as e:
if e.errno == 14: return HCQBuffer(va_addr=ptr, size=size, meta=(None, False), view=MMIOInterface(ptr, size, fmt='B'), owner=self, **kwargs)
if e.errno == 14: return HCQBuffer(va_addr=ptr, size=size, meta=(None, False), view=MMIOInterface(ptr, size, fmt='B'), owner=self)
raise RuntimeError("Failed to map external pointer to GPU memory") from e
def _gpu_free(self, mem:HCQBuffer):
tinygrad/runtime/support/hcq.py
+3 -4
View File
@@ -9,7 +9,6 @@ from tinygrad.device import BufferSpec, Compiled, LRUAllocator, ProfileDeviceEve
from tinygrad.uop.ops import sym_infer, sint, UOp
from tinygrad.runtime.autogen import libc
from tinygrad.runtime.support.memory import BumpAllocator
from tinygrad.dtype import ImageDType
class MMIOInterface:
def __init__(self, addr:int, nbytes:int, fmt='B'): self.mv, self.addr, self.nbytes, self.fmt = to_mv(addr, nbytes).cast(fmt), addr, nbytes, fmt
@@ -470,14 +469,14 @@ class HCQCompiled(Compiled, Generic[SignalType]):
if hasattr(self, 'iface') and hasattr(self.iface, 'device_fini'): self.iface.device_fini()
class HCQBuffer:
def __init__(self, va_addr:sint, size:int, image:ImageDType|None=None, meta:Any=None, _base:HCQBuffer|None=None, view:MMIOInterface|None=None,
def __init__(self, va_addr:sint, size:int, meta:Any=None, _base:HCQBuffer|None=None, view:MMIOInterface|None=None,
owner:HCQCompiled|None=None):
self.va_addr, self.size, self.image, self.meta, self._base, self.view = va_addr, size, image, meta, _base, view
self.va_addr, self.size, self.meta, self._base, self.view = va_addr, size, meta, _base, view
self._devs, self.owner = ([owner] if owner is not None else []), owner
self._mappings:dict[HCQCompiled, HCQBuffer] = {} # mapping to the other devices
def offset(self, offset:int=0, size:int|None=None) -> HCQBuffer:
return HCQBuffer(self.va_addr+offset, size or (self.size - offset), owner=self.owner, image=self.image, meta=self.meta,
return HCQBuffer(self.va_addr+offset, size or (self.size - offset), owner=self.owner, meta=self.meta,
_base=self._base or self, view=(self.view.view(offset=offset, size=size) if self.view is not None else None))
def cpu_view(self) -> MMIOInterface:
tinygrad/runtime/support/memory.py
+11
View File
@@ -231,7 +231,16 @@ class MemoryManager:
assert cls.va_allocator is not None, "must be set"
return cls.va_allocator.alloc(size, max((1 << (size.bit_length() - 1)), align))
@functools.cache # pylint: disable=method-cache-max-size-none
def identity_va(self, uncached:bool) -> int:
self.map_range(va:=self.alloc_vaddr(self.vram_size, self.vram_size), self.vram_size, [(0, self.vram_size)], AddrSpace.PHYS, uncached=uncached)
return va
def valloc(self, size:int, align=0x1000, uncached=False, contiguous=False) -> VirtMapping:
if not getenv("GMMU", 1):
paddr = self.palloc(size:=round_up(size, 0x1000), align, zero=False)
return VirtMapping(self.identity_va(uncached) + paddr, size, [(paddr, size)], aspace=AddrSpace.PHYS, uncached=uncached)
# Alloc physical memory and map it to the virtual address
va = self.alloc_vaddr(size:=round_up(size, 0x1000), align)
@@ -255,6 +264,8 @@ class MemoryManager:
return self.map_range(va, size, paddrs, aspace=AddrSpace.PHYS, uncached=uncached)
def vfree(self, vm:VirtMapping):
if not getenv("GMMU", 1): return self.pfree(vm.paddrs[0][0])
assert self.va_allocator is not None, "must be set"
self.unmap_range(vm.va_addr, vm.size)
self.va_allocator.free(vm.va_addr)
tinygrad/schedule/indexing.py
+4 -4
View File
@@ -52,8 +52,8 @@ class IndexingContext:
range_idx: Iterator[int] = field(default_factory=itertools.count)
def new_range(self, s:sint, axistype:AxisType=AxisType.LOOP) -> UOp:
if isinstance(s, UOp) and s.op is Ops.RANGE: return s
# if a range has a 1 src, it's the same as UOp.const(dtypes.index, 0)
return UOp.range(s, next(self.range_idx), axistype) if resolve(s!=1) else UOp.const(dtypes.index, 0)
# if a range has a 1 src, it's the same as UOp.const(dtypes.weakint, 0)
return UOp.range(s, next(self.range_idx), axistype) if resolve(s!=1) else UOp.const(dtypes.weakint, 0)
def create_bufferize_and_index_based_on_ranges(ctx:IndexingContext, x:UOp):
if x.op in {Ops.BUFFERIZE, Ops.INDEX}: return None
@@ -118,7 +118,7 @@ def _apply_reshape(in_shape:tuple[sint,...], out_shape:tuple[sint, ...], urngs:U
for s,src in list(zip(out_shape, urngs.src))[::-1]:
axes_in.append(acc*src)
acc *= s
combined_axes = UOp.const(dtypes.index, 0).sum(*axes_in)
combined_axes = UOp.const(dtypes.weakint, 0).sum(*axes_in)
axes_out:list[UOp] = []
for s in in_shape[::-1]:
axes_out.append(combined_axes % s)
@@ -172,7 +172,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
# treat MSTACK/MSELECT like SINK
if x.op in {Ops.MSTACK, Ops.MSELECT}: continue
if x.dtype.scalar() == dtypes.index: continue # TODO: why do I need this?
if x.dtype.scalar() == dtypes.weakint: continue # TODO: why do I need this?
ending_ranges[x] = sum([ending_ranges.get(u, []) for u in consumer_map[x]], [])
# *** the ranges on the output are
tinygrad/schedule/rangeify.py
+33 -19
View File
@@ -1,6 +1,6 @@
from dataclasses import dataclass, field, replace
import itertools
from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace, Invalid
from tinygrad.dtype import dtypes, PtrDType, AddrSpace, Invalid
from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, _substitute, KernelInfo
from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, profile_matches, should_resolve_call, identity_element
from tinygrad.uop.symbolic import symbolic
@@ -22,6 +22,17 @@ def add_ranges_to_store(ctx, x):
idxs = [UOp.range(r, next(ctx), AxisType.LOOP) for r in x.src[0].shape]
return UOp.store(x.src[0].index(*idxs), x.src[1].index(*idxs)).end(*idxs)
def lower_shaped_wmma(ctx, x):
dims, device, threads = x.arg
dtype_in, dtype_out = x.src[0].dtype.base, x.dtype
upcasts = [(s, UOp.range(s.shape[-1], next(ctx), axis_type=AxisType.UPCAST)) for s in x.src]
tc_upcast_axes = tuple(((u.arg[0], s.shape[-1]),) for s, u in upcasts)
name = f"WMMA_{'_'.join(map(str, dims))}_{dtype_in.name}_{dtype_out.name}"
wmma_arg = (name, dims, dtype_in, dtype_out, device, threads, tc_upcast_axes, ())
wmma = UOp(Ops.WMMA, dtype_out.vec(x.src[2].shape[-1]), tuple(s[u].contract(u) for s, u in upcasts), arg=wmma_arg)
tmp = UOp.placeholder((x.src[2].shape[-1],), dtype_out, slot=next(ctx), addrspace=AddrSpace.REG)
return tmp.after(UOp.group(*[tmp[e].store(wmma.gep(e)) for e in range(x.src[2].shape[-1])]))
pm_store_ranges = PatternMatcher([
(UPat(Ops.STORE, name="x"), add_ranges_to_store),
])
@@ -35,20 +46,20 @@ pm_syntactic_sugar = PatternMatcher([
lambda idx,x: x.replace(src=tuple([s.index(*idx.src[1:]) for s in x.src]))),
])
def found_assign(ctx:dict[UOp, UOp], assign:UOp, src:UOp):
if (x:=src).op is Ops.CAST and x.dtype == dtypes.half and FLOAT16: x, assign = x.src[0], assign.cast(dtypes.float)
def found_after(ctx:dict[UOp, UOp], after:UOp, src:UOp):
if (x:=src).op is Ops.CAST and x.dtype == dtypes.half and FLOAT16: x, after = x.src[0], after.cast(dtypes.float)
while True:
if x.op is Ops.PERMUTE: x, assign = x.src[0], assign.permute(argsort(x.marg))
elif x.op is Ops.RESHAPE: x, assign = x.src[0], assign.reshape(x.src[0].shape)
if x.op is Ops.PERMUTE: x, after = x.src[0], after.permute(argsort(x.marg))
elif x.op is Ops.RESHAPE: x, after = x.src[0], after.reshape(x.src[0].shape)
elif x.op is Ops.WHERE and x.src[2].base.arg == Invalid and x.src[1].op is Ops.PAD:
x, assign = x.src[1].src[0], assign.shrink(tuple((l, s-r) for (l,r),s in zip(x.src[1].marg, x.shape)))
x, after = x.src[1].src[0], after.shrink(tuple((l, s-r) for (l,r),s in zip(x.src[1].marg, x.shape)))
else: break
ctx[x] = assign
ctx[x] = after
# *** fold moved AFTERs (hack for openpilot) ***
pm_fold_moved_assign = PatternMatcher([
(UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE, src=(UPat(), UPat((*GroupOp.Movement, Ops.CAST), name="src")))), name="assign"), found_assign),
# replace ALU sources with assign versions found above
pm_fold_moved_after = PatternMatcher([
(UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE, src=(UPat(), UPat((*GroupOp.Movement,Ops.CAST,Ops.WHERE), name="src")))), name="after"), found_after),
# replace ALU sources with AFTER versions found above
(UPat(GroupOp.ALU, name="alu"), lambda ctx,alu: alu.replace(src=new_src) if (new_src:=tuple(ctx.get(s, s) for s in alu.src)) != alu.src else None),
])
@@ -58,11 +69,13 @@ pm_mops = PatternMatcher([
(UPat(GroupOp.Movement, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"),
lambda r,idx: r.src[0].index(*apply_movement_op(r.op, r.src[0].shape, r.marg, idx.src[1:]), dtype=idx.dtype, arg=idx.arg)
if r.src[0]._shape is not None and len(idx.src[1:]) == len(r.shape) else None),
# move movement ops after AFTER (but not when AFTER has a raw STORE with shaped children — from replace_contig_with_store_after)
(UPat(GroupOp.Movement, name="r").after(name="a", allow_any_len=True),
# move movement ops and INDEX after AFTER (but not when AFTER has a raw STORE with shaped children — from replace_contig_with_store_after)
(UPat(GroupOp.Movement|{Ops.INDEX}, name="r").after(name="a", allow_any_len=True),
lambda r,a: UOp(r.op, r.dtype, (a.replace(src=(r.src[0],)+a.src[1:]),)+r.src[1:], r.arg)
if a.src[0]._shape is not None and not any(s.op is Ops.STORE and s.src[0]._shape is not None for s in a.src[1:]) else None),
(UPat(GroupOp.Movement, name="r").end(name="a", allow_any_len=True), lambda r,a: a.replace(src=(r.src[0],)+a.src[1:])),
# lower SHAPED_WMMA to WMMA with CONTRACT/UNROLL
(UPat(Ops.SHAPED_WMMA, name="x"), lower_shaped_wmma),
])
# *****************
@@ -94,7 +107,7 @@ def split_reduceop(reduce:UOp, x:UOp):
# split is moved to the end to provide maximum locality for the second phase reduce.
# get expanded by rangeifying the UOp x
indexed = x.index(*[UOp.range(s, i) if resolve(s>1) else UOp.const(dtypes.index, 0) for i,s in enumerate(x.shape)])
indexed = x.index(*[UOp.range(s, i) if resolve(s>1) else UOp.const(dtypes.weakint, 0) for i,s in enumerate(x.shape)])
range_nums = [y.arg[0] for y in indexed.substitute({x.base:UOp(Ops.NOOP)}, extra_pm=pm_mops).ranges]
is_expanded = [i not in range_nums for i in range(len(x.shape))]
@@ -175,6 +188,10 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
(UPat(Ops.AFTER, src=(UPat(Ops.BITCAST, src=(UPat(name="target"),)), UPat(Ops.STORE, src=(UPat(Ops.BITCAST), UPat(name="src"))))),
lambda target, src: target.after(target.store(src.bitcast(target.dtype)))),
# wrap STORE in inner AFTER when target is a view — gives the STORE its own ranges from the view shape
(UPat(Ops.AFTER, src=(UPat(name="buf"), UPat(Ops.STORE, src=(UPat(name="target"), UPat()))), name="after"),
lambda after, buf, target: after.replace(src=(buf, target.after(after.src[1]))) if target.shape != buf.shape else None),
# make source contiguous if it has hazardous movement ops on the dest buffer
(UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE, src=(UPat(name="target"), UPat(name="src")))), name="after"), fix_store_after_hazard),
@@ -291,9 +308,6 @@ def remove_noop_bufferize(idx,b2):
pm_const_buffer_folding = pm_mops+PatternMatcher([
(UPat(Ops.BUFFERIZE, name="b"), cleanup_dead_axes),
(UPat(GroupOp.All-{Ops.BUFFERIZE, Ops.PARAM}, name="x"), lambda x: x.replace(dtype=x.dtype.base) if isinstance(x.dtype, ImageDType) else None),
(UPat((Ops.BUFFERIZE), name="x"), lambda x: x.replace(dtype=x.dtype.base) if isinstance(x.dtype, ImageDType)
and (resolve(prod(x.dtype.shape)!=prod(x.shape)) or x.shape[-1]%4!=0) else None),
# remove noop buffers. if we look at the next index we can remove even more of these
(UPat(Ops.INDEX, name="idx").f(Ops.BUFFERIZE, allow_any_len=True, name="b2"), remove_noop_bufferize),
# no buffers for const (ranges don't matter for const - it's the same value everywhere)
@@ -515,11 +529,11 @@ rangeify_codegen = PatternMatcher([
# fix broadcast dtype
(UPat(Ops.AFTER, name="a").broadcast(name="b"), lambda a,b: a.broadcast(len(b.src))),
(UPat(Ops.DEFINE_LOCAL).f(Ops.AFTER, allow_any_len=True).broadcast(name="dg").f(Ops.INDEX, name="idx", allow_any_len=True),
lambda dg,idx: None if isinstance(idx.dtype, (PtrDType, ImageDType)) else
lambda dg,idx: None if isinstance(idx.dtype, PtrDType) else
idx.replace(dtype=dg.dtype, arg=None).load(dtype=dg.dtype.base.scalar().vec(dg.dtype.vcount))),
(UPat(Ops.AFTER, name="a").gep(name="b"), lambda a,b: a.gep(b.arg)),
(UPat(Ops.DEFINE_LOCAL).f(Ops.AFTER, allow_any_len=True).gep(name="dg").f(Ops.INDEX, name="idx", allow_any_len=True),
lambda dg,idx: None if isinstance(idx.dtype, (PtrDType, ImageDType)) else
lambda dg,idx: None if isinstance(idx.dtype, PtrDType) else
idx.replace(dtype=dg.dtype, arg=None).load(dtype=dg.dtype.base.scalar().vec(dg.dtype.vcount))),
])
@@ -556,7 +570,7 @@ split_kernels = PatternMatcher([
@profile_matches
def get_kernel_graph(sink:UOp) -> UOp:
tsink = graph_rewrite(sink, multi_pm, name="multi_pm")
if OPENPILOT_HACKS: tsink = graph_rewrite(tsink, pm_fold_moved_assign, ctx={}, name="fold moved assigns")
if OPENPILOT_HACKS: tsink = graph_rewrite(tsink, pm_fold_moved_after, ctx={}, name="fold moved afters")
tsink = graph_rewrite(tsink, pm_syntactic_sugar+pm_mops+earliest_rewrites, bottom_up=True, name="earliest rewrites")
# convert movement ops to ranges
tinygrad/tensor.py
+52 -60
View File
@@ -4,8 +4,8 @@ import time, math, itertools, functools, struct, sys, inspect, pathlib, string,
from contextlib import ContextDecorator
from typing import Any, Callable, ClassVar, Sequence, cast, get_args, Literal, SupportsIndex, ParamSpec, TypeVar, Generic, TYPE_CHECKING
if TYPE_CHECKING: import numpy
from tinygrad.dtype import DType, DTypeLike, dtypes, ImageDType, ConstType, least_upper_float, least_upper_dtype, sum_acc_dtype, to_dtype, truncate
from tinygrad.dtype import _from_np_dtype, _to_np_dtype, PyConst, Invalid, InvalidType
from tinygrad.dtype import DType, DTypeLike, dtypes, ConstType, least_upper_float, least_upper_dtype, sum_acc_dtype, to_dtype, truncate
from tinygrad.dtype import _from_np_dtype, _to_np_dtype, PyConst, Invalid
from tinygrad.helpers import argfix, make_tuple, flatten, prod, all_int, round_up, merge_dicts, argsort, getenv, all_same, fully_flatten
from tinygrad.helpers import IMAGE, FLOAT16, WINO, Metadata, TRACEMETA, ASM_GEMM, ceildiv, fetch, is_numpy_ndarray, TracingKey, cpu_profile
from tinygrad.helpers import suppress_finalizing, disable_gc
@@ -14,12 +14,12 @@ from tinygrad.mixin import OpMixin
from tinygrad.mixin.movement import _align_left
from tinygrad.uop.ops import smax, smin, resolve, UOp, Ops, sint, identity_element, all_metadata, _index_to_concrete_int, sint_to_uop, Variable
from tinygrad.engine.schedule import ExecItem, complete_create_schedule_with_vars
from tinygrad.device import Device, Buffer
from tinygrad.device import Buffer
from tinygrad.engine.realize import run_schedule
from tinygrad.engine.allocations import transform_to_call
# TODO: this should be the only usage of Device
def canonicalize_device(device:str|tuple|list|None) -> str|tuple[str, ...]:
from tinygrad.device import Device
if not isinstance(device, (tuple, list)): return Device.canonicalize(device)
return canonical[0] if len(canonical:=tuple(Device.canonicalize(d) for d in device)) == 1 else canonical
@@ -56,15 +56,15 @@ def get_shape(x) -> tuple[int, ...]:
if not all_same(subs:=[get_shape(xi) for xi in x]): raise ValueError(f"inhomogeneous shape from {x}")
return (len(subs),) + (subs[0] if subs else ())
def _frompy(x:list|tuple|bytes, dtype:DType) -> UOp:
def _frompy(x:list|tuple|bytes, dtype:DType, device:str|tuple[str,...]) -> UOp:
if isinstance(x, bytes): ret, data = UOp.new_buffer("PYTHON", len(x)//dtype.itemsize, dtype), x
else:
ret = UOp.new_buffer("PYTHON", prod(shape:=get_shape(x)), dtype).reshape(shape)
assert dtype.fmt is not None, f"{dtype=} has None fmt"
truncate_function = truncate[dtype]
data = struct.pack(f"{ret.size}{dtype.fmt}", *[truncate_function(dtype.const(xi)) for xi in fully_flatten(x)])
# fake realize
ret.buffer.allocate(memoryview(data if Device.DEFAULT != "PYTHON" else bytearray(data)))
# fake realize. if target device is PYTHON it needs bytearray to be writable
ret.buffer.allocate(memoryview(data if device != "PYTHON" else bytearray(data)))
return ret
def _get_winograd_matcols(mat, dims:int, shp:tuple[sint, ...], device:str|tuple[str, ...], dtype:DType) -> list[list[Tensor]]:
@@ -132,9 +132,10 @@ class Tensor(OpMixin):
# create a UOp from the different types of inputs
if isinstance(data, UOp):
assert _dtype is None or _dtype==data.dtype or data.dtype==dtypes.index, f"dtype mismatch: {_dtype} vs {data.dtype}"
# if data is dtype.index that means that this is a symbolic int and we need to lower it to something we can make a Tensor out of
if data.dtype==dtypes.index: data = _index_to_concrete_int(data)
assert _dtype is None or _dtype==data.dtype or data.dtype==dtypes.weakint, f"dtype mismatch: {_dtype} vs {data.dtype}"
# if data is dtype.weakint that means that this is a symbolic int and we need to lower it to something we can make a Tensor out of
# TODO: remove this and stay in weakint
if data.dtype==dtypes.weakint: data = _index_to_concrete_int(data)
if data.op is Ops.BIND:
var, val = data.unbind()
# give the bound constant a device
@@ -144,16 +145,13 @@ class Tensor(OpMixin):
data = UOp.const(_dtype or dtypes.default_float, 0, _device)
elif isinstance(data, get_args(ConstType)):
data = (UOp.unique_const if _force_unique or requires_grad else UOp.const)(_dtype or dtypes.from_py(data), data, _device)
elif isinstance(data, InvalidType):
assert _dtype is not None
data = UOp.const(_dtype, data, _device)
elif isinstance(data, bytes): data = _frompy(data, _dtype or dtypes.uint8)
elif isinstance(data, bytes): data = _frompy(data, _dtype or dtypes.uint8, _device)
elif isinstance(data, (list, tuple)):
if _dtype is None:
if (d := fully_flatten(data)) and all(isinstance(s, bool) for s in d): _dtype = dtypes.bool
else: _dtype = dtypes.default_int if d and all_int(d) else dtypes.default_float # NOTE: this works because all_int([True, False]) is True
if _dtype in [dtypes.bfloat16, *dtypes.fp8s]: data = _frompy(data, dtypes.float32).cast(_dtype)
else: data = _frompy(data, _dtype)
if _dtype in [dtypes.bfloat16, *dtypes.fp8s]: data = _frompy(data, dtypes.float32, _device).cast(_dtype)
else: data = _frompy(data, _dtype, _device)
elif is_numpy_ndarray(data):
import numpy as np
assert isinstance(data, np.ndarray), f"expected np.ndarray, got {data}"
@@ -1044,7 +1042,7 @@ class Tensor(OpMixin):
# ***** toposort and backward pass *****
def gradient(self, *targets:Tensor, gradient:Tensor|None=None, materialize_grads=False) -> list[Tensor]:
def gradient(self, *targets:Tensor, gradient:Tensor|None=None) -> list[Tensor]:
"""
Computes the gradient of the targets with respect to self.
@@ -1065,9 +1063,7 @@ class Tensor(OpMixin):
grads = compute_gradient(self.uop, gradient.uop, set(target_uops))
ret:list[Tensor] = []
for x in target_uops:
if (y:=grads.get(x)) is None:
if materialize_grads: y = x.const_like(0)
else: raise RuntimeError(f"{x}\n\nnot found in\n\n{self.uop}")
if (y:=grads.get(x)) is None: y = x.const_like(0)
ret.append(Tensor(y))
return ret
@@ -1085,7 +1081,7 @@ class Tensor(OpMixin):
tensors_need_grad: list[Tensor] = [t for tref in all_tensors if (t:=tref()) is not None and \
t.uop in all_uops and t.requires_grad]
# clear contexts
for t,g in zip(tensors_need_grad, self.gradient(*tensors_need_grad, gradient=gradient, materialize_grads=True)):
for t,g in zip(tensors_need_grad, self.gradient(*tensors_need_grad, gradient=gradient)):
assert g.shape == t.shape, f"grad shape must match tensor shape, {g.shape!r} != {t.shape!r}"
if t.grad is None: t.grad = g
else: t.grad.assign(t.grad + g.to(t.grad.device))
@@ -2957,8 +2953,7 @@ class Tensor(OpMixin):
if not isinstance(y, Tensor):
# make y a Tensor
assert isinstance(y, (*get_args(ConstType), UOp)), f"{type(y)=}, {y=}"
if y is Invalid or isinstance(x.dtype, ImageDType) or dtypes.is_float(x.dtype) or (dtypes.is_int(x.dtype) and isinstance(y, int)):
y_dtype = x.dtype
if y is Invalid or dtypes.is_float(x.dtype) or (dtypes.is_int(x.dtype) and isinstance(y, int)): y_dtype = x.dtype
elif not isinstance(y, UOp): y_dtype = dtypes.from_py(y)
if isinstance(y, UOp): y = Tensor.from_uop(y, device=x.device)
else: y = Tensor(y_dtype.const(y), x.device, y_dtype, requires_grad=False)
@@ -3136,9 +3131,7 @@ class Tensor(OpMixin):
"""
# NOTE: torch always return in float, we return based on the broadcasting rule.
other = self._broadcasted(other)[1]
# TODO: remove other.sign()*0?
# other.sign()*0 keeps other in the gradient graph (gradient=0) without affecting forward (works for inf unlike other*0)
return self.abs() * ((other < 0) | (other.reciprocal() < 0)).where(-1, 1) + other.sign()*0
return self.abs() * ((other < 0) | (other.reciprocal() < 0)).where(-1, 1)
def logaddexp(self, other) -> Tensor:
"""
@@ -3632,11 +3625,15 @@ class Tensor(OpMixin):
return cx.image_conv2d(cw, groups=groups, dtype=dtype).reshape(out_shape_t).transpose(self.ndim-1, self.ndim-2)
def image_conv2d(self, weight:Tensor, bias:Tensor|None=None, groups=1, stride=1, dilation=1, padding=0, dtype=None) -> Tensor:
base_image_type, dtsz = (dtypes.imageh, 2) if FLOAT16 else (dtypes.imagef, 4)
dtsz = 2 if FLOAT16 else 4
(bs,_,iy,ix), (cout,cin,H,W) = self.shape, weight.shape
x, w = self, weight.reshape(groups, (rcout := cout//groups), cin, H, W)
padding_neg, padding_pos = [min(0, p) for p in self._resolve_pool_pads(padding, 2)], [max(0, p) for p in self._resolve_pool_pads(padding, 2)]
x = x.pad(padding_neg)
iy, ix = x.shape[2:]
# hack for non multiples of 4 on cin
if cin % 4 != 0 and not (cin == 1 and groups%4 == 0):
x = x.reshape(bs, groups, cin, iy, ix) # do this always?
@@ -3661,34 +3658,30 @@ class Tensor(OpMixin):
else: w = w.reshape(cout//4,4,cin//4,4,H,W).permute(0,4,2,5,3,1)
# contiguous creates the image, and early realize static weights (TODO: test for the static weight)
if IMAGE == 1:
def is_pow2(v): return v > 0 and v & (v - 1) == 0
# pad dimension i to amt with invalids
def ipad(t, i, amt):
shape = (None,)*i + (amt,) + (None,)*(t.ndim-i-1)
return Tensor(True, device=t.device).expand(t.shape).pad_to(shape).where(t.pad_to(shape), Invalid) if amt != t.shape[i] else t
def is_pow2(v): return v > 0 and v & (v - 1) == 0
# pad dimension i to amt with invalids
def ipad(t, i, amt):
shape = (None,)*i + (amt,) + (None,)*(t.ndim-i-1)
return Tensor(True, device=t.device).expand(t.shape).pad_to(shape).where(t.pad_to(shape), Invalid) if amt != t.shape[i] else t
# align a dimension, use at to specify the dimension to pad in, defaults to first
def pad_align(t, dim, at=None, force=False):
# align to 64 pixels when height is real, otherwise 64 bytes is sufficient
align = (64 // dtsz) if prod(t.shape[:dim]) == 1 or prod(t.shape) < 16384 * 4 else 256
return ipad(t, at:=at or dim, round_up(t.shape[at] + int(force), align // math.gcd(prod(t.shape[dim:]) // t.shape[at], align)))
# align a dimension, use at to specify the dimension to pad in, defaults to first
def pad_align(t, dim, at=None, force=False):
# align to 64 pixels when height is real, otherwise 64 bytes is sufficient
align = (64 // dtsz) if prod(t.shape[:dim]) == 1 or prod(t.shape) < 16384 * 4 else 256
return ipad(t, at:=at or dim, round_up(t.shape[at] + int(force), align // math.gcd(prod(t.shape[dim:]) // t.shape[at], align)))
# bank conflicts
if cin >= 8 and is_pow2(cin // 4):
x, w = pad_align(x.reshape(bs, iy, ix, groups, cin // 4, 4), 2, at=4, force=True), pad_align(w, 1, at=2, force=True)
else: x, w = pad_align(x, 2), pad_align(w, 1)
# bank conflicts
if cin >= 8 and is_pow2(cin // 4):
x, w = pad_align(x.reshape(bs, iy, ix, groups, cin // 4, 4), 2, at=4, force=True), pad_align(w, 1, at=2, force=True)
else: x, w = pad_align(x, 2), pad_align(w, 1)
if FLOAT16: x, w = x.cast(dtypes.half).contiguous().cast(dtypes.float), w.cast(dtypes.half).contiguous().cast(dtypes.float)
else: x, w = x.contiguous(), w.contiguous()
# undo alignment hacks
if cin >= 8 and is_pow2(cin // 4): x, w = x[:, :, :ix, :, :cin // 4, :], w[:, :H, :cin // 4, ...]
else: x, w = x[:, :, :ix, :], w[:, :H, ...]
elif IMAGE: x, w = x.cast(base_image_type((bs*iy, ix*groups*cin//4, 4))).contiguous(), w.cast(base_image_type((cout//4, H*W*cin, 4))).contiguous()
if FLOAT16: x, w = x.cast(dtypes.half).contiguous().cast(dtypes.float), w.cast(dtypes.half).contiguous().cast(dtypes.float)
else: x, w = x.contiguous(), w.contiguous()
# undo alignment hacks
if cin >= 8 and is_pow2(cin // 4): x, w = x[:, :, :ix, :, :cin // 4, :], w[:, :H, :cin // 4, ...]
else: x, w = x[:, :, :ix, :], w[:, :H, ...]
# expand out
rcin_hi, rcin_lo = (cin//4, 4) if cin >= 4 else (1, 1)
group_shape, rcout_expand = (groups//4, 4) if cin == 1 else (groups, 1), (rcout//4, 4) if rcout >= 4 else (1, 1)
@@ -3697,24 +3690,23 @@ class Tensor(OpMixin):
else: w = w.reshape(cout//4, H, rcin_hi, W, rcin_lo, 4).permute(0,1,2,3,5,4)
# prepare input
x = x.permute(0,3,4,5,1,2).pad(self._resolve_pool_pads(padding,2))._pool((H,W), stride, dilation)# -> (bs, groups, rcin_hi, rcin_lo, oy, ox, H, W)
x = x.permute(0,3,4,5,1,2).pad(padding_pos)._pool((H,W), stride, dilation)# -> (bs, groups, rcin_hi, rcin_lo, oy, ox, H, W)
x = x.permute(0,4,5,1,2,3,6,7).reshape(bs, (oy := x.shape[4]), (ox := x.shape[5]), *group_shape, 1, 1, rcin_hi, rcin_lo, H, W)
# prepare weights
w = w.permute(0,4,2,5,1,3).reshape((1, 1, 1, *group_shape, *rcout_expand, rcin_hi, rcin_lo, H, W))
if IMAGE == 1:
added_ox = 0
assert isinstance(ox, int) and isinstance(cout, int)
if (ox * cout) % (64 // dtsz):
added_ox = round_up(ox, 64 // (dtsz * math.gcd(cout, 64 // dtsz))) - ox
ox = ox + added_ox
x = x.pad_to(None, None, ox, None, None, None, None, None, None, None, None)
added_ox = 0
assert isinstance(ox, int) and isinstance(cout, int)
if (ox * cout) % (64 // dtsz):
added_ox = round_up(ox, 64 // (dtsz * math.gcd(cout, 64 // dtsz))) - ox
ox = ox + added_ox
x = x.pad_to(None, None, ox, None, None, None, None, None, None, None, None)
# the conv!
ret = (x*w).cast(base_image_type((bs*oy, ox*cout//4, 4)) if IMAGE >= 2 else dtypes.float32).sum((-4, -3, -2, -1), dtype=dtype)
ret = (x*w).cast(dtypes.float32).sum((-4, -3, -2, -1), dtype=dtype)
if IMAGE == 1 and added_ox:
if added_ox:
ret = ret.reshape(bs, oy, ox, groups, rcout)[:, :, :-added_ox, ...]
ox = ox - added_ox
tinygrad/uop/__init__.py
+1 -1
View File
@@ -53,7 +53,7 @@ class Ops(FastEnum):
# ** 4 -- math **
# tensor core math op, not elementwise
WMMA = auto()
WMMA = auto(); SHAPED_WMMA = auto()
# UnaryOps
CAST = auto(); BITCAST = auto(); EXP2 = auto(); LOG2 = auto(); SIN = auto()
tinygrad/uop/decompositions.py
+26 -4
View File
@@ -2,9 +2,9 @@ from typing import Callable
import math, functools
from tinygrad.dtype import dtypes, DType, promo_lattice, truncate
from tinygrad.device import is_dtype_supported
from tinygrad.helpers import flatten, polyN
from tinygrad.helpers import flatten, polyN, EMULATED_DTYPES
from tinygrad.uop import GroupOp
from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher
from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, graph_rewrite
TRANSCENDENTAL_DTYPES = (dtypes.float16, dtypes.float32, dtypes.float64)
@@ -319,7 +319,7 @@ def threefry2x32(x: UOp, key: UOp):
l2i_dt = {dtypes.long: dtypes.int, dtypes.ulong: dtypes.uint}
def unpack32(v:UOp) -> tuple[UOp, UOp]: return v.bitcast(dtypes.uint) & 0xFFFF, shr(v.bitcast(dtypes.uint), 16)
def reindex(idx:UOp, off:int, mul=2) -> UOp: return idx.replace(src=(idx.src[0], idx.src[1]*mul+off))
def reindex(idx:UOp, off:int, mul=2) -> UOp: return idx.replace(src=(idx.src[0], idx.src[1]*mul+off, *idx.src[2:]))
# 4.3.1 is the relevant section in TAOCP
def l2i(op: Ops, dt: DType, *uops:UOp):
@@ -511,10 +511,15 @@ pm_float_decomp = PatternMatcher([
(UPat((*GroupOp.Defines, Ops.INDEX), name="x"), lambda ctx,x:
x.replace(dtype=f2f_dt[ctx[0]].ptr(x.dtype.size), tag=ctx[0]) if x.dtype.base == ctx[0] else None),
(UPat(Ops.LOAD, dtypes.floats, name="x"), lambda ctx,x: f2f_load(x, *ctx) if x.dtype.scalar() == ctx[0] else None),
# bitcasted load should just replace load
(UPat(Ops.BITCAST, src=(UPat(Ops.LOAD, name="ld"),), name="bc"), lambda ctx,bc,ld:
ld.replace(dtype=f2f_dt[ctx[0]]).bitcast(bc.dtype) if ld.dtype.bitsize == ctx[0].bitsize else None),
ld.replace(dtype=f2f_dt[ctx[0]]).bitcast(bc.dtype) if ld.dtype == ctx[0] else None),
# bitcast from
(UPat(Ops.BITCAST, src=(UPat.var("x", dtypes.floats),), name="bc"), lambda ctx,bc,x:
bc.replace(src=(f2f(x.bitcast(f2f_dt[ctx[1]]), ctx[1], ctx[0]),)) if x.dtype == ctx[1] and bc.dtype.bitsize == ctx[0].bitsize else None),
# bitcast to
(UPat(Ops.BITCAST, src=(UPat.var("x"),), name="bc"), lambda ctx,bc,x:
f2f(x.bitcast(f2f_dt[ctx[0]]), ctx[0], ctx[1]) if bc.dtype == ctx[0] else None),
(UPat(Ops.CAST, dtypes.floats, src=(UPat.var("val"),), name="x"), lambda ctx,x,val:
f2f_clamp(val.cast(ctx[1]), ctx[0]) if x.dtype.scalar() == ctx[0] else None),
(UPat(GroupOp.All-{Ops.BITCAST}, dtypes.floats, name="x"), lambda ctx,x:
@@ -525,3 +530,20 @@ pm_float_decomp = PatternMatcher([
(UPat(Ops.STORE, src=(UPat.var("idx"), UPat.var("val", dtypes.floats)), name='st'), lambda ctx,st,idx,val:
f2f_store(st, idx, val, *ctx) if val.dtype.scalar() == ctx[1] and (idx:=idx.src[0] if idx.op == Ops.CAST else idx).tag == ctx[0] else None),
])
def do_dtype_decomps(sink:UOp, ctx:tuple[set[DType], str, str]) -> UOp:
def _should_emulate(dt): return dt in EMULATED_DTYPES.tolist(dtypes) or not is_dtype_supported(dt, ctx[1], ctx[2])
for fr in sorted(filter(_should_emulate, ctx[0])):
if fr in dtypes.floats:
to = dtypes.half if not _should_emulate(dtypes.half) and fr in dtypes.fp8s else dtypes.float
sink = graph_rewrite(sink, pm_float_decomp, name=f"decomp {fr} -> {to}", ctx=(fr, to), bottom_up=True)
else: sink = graph_rewrite(sink, pm_long_decomp, name="decomp long -> int", bottom_up=True)
return sink
pm_dtype_decomps = PatternMatcher([
# detect dtypes to decompose
(UPat(GroupOp.All, (*dtypes.fp8s, dtypes.bfloat16, dtypes.half, dtypes.long, dtypes.ulong), name="x"), lambda x,ctx:
ctx[0].add({dtypes.ulong:dtypes.long}.get(dt:=x.dtype.base.scalar(), dt))),
# do the rewrites
(UPat(Ops.SINK, name="sink"), do_dtype_decomps),
])
tinygrad/uop/divandmod.py
+7 -7
View File
@@ -107,17 +107,17 @@ div_and_mod_symbolic = PatternMatcher([
# ** 1. Fast Inline Rules **
((UPat.var("x")//UPat.cvar("c") + UPat.cvar("a"))//UPat.cvar("d"), lambda x,c,a,d: (x+a*c)//(c*d)
if c.vmin>0 and d.vmin>0 and x.vmin>=0 and a.vmin>=0 else None), # (x//c+a)//d -> (x+a*c)//(c*d)
(UPat.var("x", dtypes.index) // UPat.var("d"), lambda x,d: -(x//(-d)) if d.vmax < 0 else None),
(UPat.var("x", dtypes.index) // UPat.var("d"), lambda x,d: -((-x)//d) if x.vmax <= 0 else None),
((UPat.var("x", dtypes.index)+UPat.cvar("c", vec=False)).named("n")//UPat.cvar("d", vec=False),
(UPat.var("x", dtypes.weakint) // UPat.var("d"), lambda x,d: -(x//(-d)) if d.vmax < 0 else None),
(UPat.var("x", dtypes.weakint) // UPat.var("d"), lambda x,d: -((-x)//d) if x.vmax <= 0 else None),
((UPat.var("x", dtypes.weakint)+UPat.cvar("c", vec=False)).named("n")//UPat.cvar("d", vec=False),
lambda x,c,n,d: ((x+c.arg%d.arg)//d + c.arg//d.arg) if c.arg%d.arg!=c.arg and x.vmin>=0 and n.vmin>=0 and d.arg>0 else None),
((UPat.var("x", dtypes.index)+UPat.cvar("c", vec=False)).named("n")//UPat.cvar("d", vec=False),
((UPat.var("x", dtypes.weakint)+UPat.cvar("c", vec=False)).named("n")//UPat.cvar("d", vec=False),
lambda x,c,n,d: (-(-(c.arg%d.arg + x - (d.arg-1))//d) + c.arg//d.arg) if x.vmax<=0 and n.vmin>=0 and d.arg>0 else None),
# ** 2. Slow Rules **
(UPat((Ops.IDIV, Ops.MOD), dtypes.index, name="d"), lambda d: fold_divmod_general(d, bool(CORRECT_DIVMOD_FOLDING))),
(UPat((Ops.IDIV, Ops.MOD), dtypes.weakint, name="d"), lambda d: fold_divmod_general(d, bool(CORRECT_DIVMOD_FOLDING))),
# NOTE: these have to go at the bottom or TestSymbolicOps.test_var loops
(UPat.var("x", dtypes.index) % UPat.var("d"), lambda x,d: -((-x)%d) if x.vmax <= 0 else None),
(UPat.var("x", dtypes.index) % UPat.var("d"), lambda x,d: (x%(-d)) if d.vmax < 0 else None),
(UPat.var("x", dtypes.weakint) % UPat.var("d"), lambda x,d: -((-x)%d) if x.vmax <= 0 else None),
(UPat.var("x", dtypes.weakint) % UPat.var("d"), lambda x,d: (x%(-d)) if d.vmax < 0 else None),
])
tinygrad/uop/ops.py
+40 -29
View File
@@ -59,9 +59,9 @@ def multirange_str(rngs:Iterable[UOp], color=False, pad=None) -> str:
return ret
def shape_to_shape_arg(arg:tuple[sint, ...]) -> UOp:
if len(arg) == 0: return UOp(Ops.VECTORIZE, dtypes.index.vec(0))
elif all_int(arg): return UOp.const(dtypes.index.vec(len(arg)), arg)
else: return UOp(Ops.VECTORIZE, dtypes.index.vec(len(arg)), tuple(UOp.const(dtypes.index, x) if isinstance(x, int) else x for x in arg))
if len(arg) == 0: return UOp(Ops.VECTORIZE, dtypes.weakint.vec(0))
elif all_int(arg): return UOp.const(dtypes.weakint.vec(len(arg)), arg)
else: return UOp(Ops.VECTORIZE, dtypes.weakint.vec(len(arg)), tuple(UOp.const(dtypes.weakint, x) if isinstance(x, int) else x for x in arg))
def consumer_map_from_toposort(lst:Iterable[UOp]):
ret: dict[UOp, dict[UOp, None]] = {}
@@ -249,6 +249,9 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
if len(self.src) >= 1: return tuple(self.src[0].sgep(i) for i in range(self.src[0].dtype.count))
return None
# SHAPED_WMMA output shape = accumulator shape (src[2])
case Ops.SHAPED_WMMA: return self.src[2]._shape
# passthrough ops
case Ops.REDUCE | Ops.MSTACK | Ops.MSELECT | Ops.DETACH | Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.AFTER | Ops.END:
return self.src[0]._shape
@@ -343,6 +346,9 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
def ended_ranges(self) -> tuple[UOp, ...]:
if self.op in range_start: return self.src[range_start[self.op]:]
if self.op is Ops.AFTER: return tuple(flatten([x.ended_ranges for x in self.src[1:]]))
if self.op is Ops.CONTRACT:
contract_rng_ids = {rng_id for rng_id, _ in self.arg}
return tuple(r for r in self.src[0].ranges if r.op is Ops.RANGE and r.arg[0] in contract_rng_ids)
return ()
# determine what ranges this is in
@@ -426,9 +432,9 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
assert len(idx) == len(self.shape), f"__getitem__ shape mismatch, indexing {self.shape} with {len(idx)} args"
if len(slice_idx:=[i for i,x in enumerate(idx) if isinstance(x, slice)]):
perm = self.permute(tuple([i for i in range(self.ndim) if i not in slice_idx] + slice_idx))
return perm.index(*[UOp.const(dtypes.index, x) if isinstance(x, int) else x for x in idx if not isinstance(x, slice)], ptr=True)
return perm.index(*[UOp.const(dtypes.weakint, x) if isinstance(x, int) else x for x in idx if not isinstance(x, slice)], ptr=True)
else:
return self.index(*[UOp.const(dtypes.index, x) if isinstance(x, int) else x for x in idx])
return self.index(*[UOp.const(dtypes.weakint, x) if isinstance(x, int) else x for x in idx])
def const_like(self, b:ConstLike):
# constants can optionally have a DEVICE source
return UOp.const(self.dtype, b, device=self._device, shape=self._shape)
@@ -483,21 +489,21 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
ret = ret.replace(src=(UOp.unique(None if unique is True else unique),) + ret.src)
return ret.reshape((1,)*len(shape)).expand(shape) if shape is not None else ret
@staticmethod
def range(end:sint, axis_id, axis_type=AxisType.LOOP, *arg, dtype=dtypes.index, src=(), **kwargs):
def range(end:sint, axis_id, axis_type=AxisType.LOOP, *arg, dtype=dtypes.weakint, src=(), **kwargs):
return UOp(Ops.RANGE, dtype=dtype, src=(sint_to_uop(end, dtype),)+src, arg=(axis_id, axis_type)+arg, **kwargs)
@staticmethod
def special(end:sint, name:str, dtype=dtypes.index): return UOp(Ops.SPECIAL, dtype=dtype, src=(sint_to_uop(end, dtype),), arg=name)
def special(end:sint, name:str, dtype=dtypes.weakint): return UOp(Ops.SPECIAL, dtype=dtype, src=(sint_to_uop(end, dtype),), arg=name)
def r(self, op:Ops, axis:tuple[int, ...]):
axis = tuple(sorted([x for x in axis if resolve(self.shape[x] != 1)]))
return UOp(Ops.REDUCE_AXIS, self.dtype, (self,), (op, axis)) if len(axis) else self
@staticmethod
def invalid(count=1): return UOp(Ops.CONST, dtypes.index.vec(count), src=(), arg=Invalid)
def invalid(count=1): return UOp(Ops.CONST, dtypes.weakint.vec(count), src=(), arg=Invalid)
def valid(self, cond): return self if cond.op is Ops.WHERE and cond.arg else cond.where(self, UOp.invalid(self.dtype.count))
def get_idx(self) -> UOp:
assert self.dtype.scalar() is dtypes.index, "Can only call get_idx on index dtype"
assert self.dtype.scalar() is dtypes.weakint, "Can only call get_idx on index dtype"
return self.src[1] if self.op is Ops.WHERE and self.src[2].arg is Invalid else self
def get_valid(self) -> UOp:
assert self.dtype.scalar() is dtypes.index, "Can only call get_valid on index dtype"
assert self.dtype.scalar() is dtypes.weakint, "Can only call get_valid on index dtype"
return self.src[0] if self.op is Ops.WHERE and self.src[2].arg is Invalid else UOp.const(dtypes.bool, self.arg is not Invalid)
def reduce(self, *src:UOp, **kwargs): return UOp(Ops.REDUCE, kwargs.pop('dtype', self.dtype), src=(self,)+src, **kwargs)
@@ -765,7 +771,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
# *** uop Variable stuff ***
@staticmethod
def variable(name:str, min_val:ConstType, max_val:ConstType, dtype:DType=dtypes.index) -> UOp:
def variable(name:str, min_val:ConstType, max_val:ConstType, dtype:DType=dtypes.weakint) -> UOp:
assert not isinstance(min_val, UOp) and not isinstance(max_val, UOp), f"can't create Variable {name} with {min_val}/{max_val}"
return UOp(Ops.DEFINE_VAR, dtype, arg=(name, min_val, max_val))
@property
@@ -874,7 +880,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
if self.op is Ops.VCONST and Invalid not in self.arg: return (min(self.arg), max(self.arg))
if self.op is Ops.GEP: return self.src[0]._min_max
# TODO: CAST to bool/unsigned is not monotone, still some case can be simplified
if self.op is Ops.CAST and self.dtype in dtypes.floats+dtypes.sints+(dtypes.index,):
if self.op is Ops.CAST and self.dtype in dtypes.floats+dtypes.sints+(dtypes.weakint,):
return max(self.dtype.min, self.src[0].vmin), min(self.src[0].vmax, self.dtype.max)
return self.dtype.min, self.dtype.max
@@ -999,7 +1005,7 @@ python_alu: dict[Ops, Callable] = {
def exec_alu(op:Ops, dtype:DType, operands, truncate_output=True):
if dtype.count > 1:
return tuple([exec_alu(op, dtype.scalar(), [x[i] if isinstance(x, tuple) else x for x in operands]) for i in range(dtype.count)])
if dtype==dtypes.index and op in GroupOp.Binary and Invalid in operands: return Invalid
if dtype==dtypes.weakint and op in GroupOp.Binary and Invalid in operands: return Invalid
alu = python_alu[op](*operands)
return truncate.get(dtype, lambda x: x)(alu) if truncate_output else alu
@@ -1432,23 +1438,26 @@ def graph_rewrite(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False, name=N
rewrite_ctx = RewriteContext(pm if not bottom_up else None, pm if bottom_up else bpm, ctx, enter_calls)
return rewrite_ctx.walk_rewrite(sink) if walk else rewrite_ctx.unified_rewrite(sink)
def sint_to_uop(x:sint, dtype=dtypes.index) -> UOp: return UOp.const(dtype, x) if isinstance(x, int) else x.cast(dtype)
def sint_to_uop(x:sint, dtype=dtypes.weakint) -> UOp: return UOp.const(dtype, x) if isinstance(x, int) else x.cast(dtype)
def select_dtype(u): return (dtypes.long if u.overflows(dtypes.int32) else dtypes.int).vec(u.dtype.count)
pm_lower_index_dtype = PatternMatcher([
# There are no Unary ops at this point in symbolic, those are introduced later
(UPat(GroupOp.Binary, name="u", src=(UPat.var("x").cast(dtypes.index), UPat.var("y").cast(dtypes.index))), lambda u,x,y:
(UPat(GroupOp.Binary, name="u", src=(UPat.var("x").cast(dtypes.weakint), UPat.var("y").cast(dtypes.weakint))), lambda u,x,y:
x.cast(dt:=least_upper_dtype(select_dtype(u), x.dtype, y.dtype)).alu(u.op, y.cast(dt)).cast(u.dtype)),
(UPat((Ops.CONST, Ops.VCONST), dtype=dtypes.index, name="u"), lambda u: u.replace(dtype=select_dtype(u)).cast(u.dtype) if u.arg!=Invalid else None),
(UPat(Ops.WHERE, dtypes.index, src=(UPat.var("cond"), UPat.var("x").cast(dtypes.index), UPat.var("y").cast(dtypes.index))), lambda cond,x,y:
cond.where(x.cast(dt:=least_upper_dtype(x.dtype, y.dtype)), y.cast(dt)).cast(dtypes.index)),
(UPat(Ops.RANGE, src=(UPat.var("end").cast(dtypes.index)), name="r"), lambda r,end: r.replace(dtype=end.dtype, src=(end,)).cast(dtypes.index)),
(UPat(Ops.VECTORIZE, src=UPat().cast(dtypes.index), name="v"),
lambda v: v.replace(dtype=(dt:=select_dtype(v)), src=tuple(s.src[0].cast(dt.scalar()) for s in v.src)).cast(dtypes.index)),
(UPat((Ops.CONST, Ops.VCONST), dtype=dtypes.weakint, name="u"),
lambda u: u.replace(dtype=select_dtype(u)).cast(u.dtype) if u.arg!=Invalid else None),
(UPat(Ops.WHERE, dtypes.weakint, src=(UPat.var("cond"), UPat.var("x").cast(dtypes.weakint), UPat.var("y").cast(dtypes.weakint))), lambda cond,x,y:
cond.where(x.cast(dt:=least_upper_dtype(x.dtype, y.dtype)), y.cast(dt)).cast(dtypes.weakint)),
(UPat(Ops.RANGE, src=(UPat.var("end").cast(dtypes.weakint)), name="r"), lambda r,end: r.replace(dtype=end.dtype, src=(end,)).cast(dtypes.weakint)),
(UPat(Ops.VECTORIZE, src=UPat().cast(dtypes.weakint), name="v"),
lambda v: v.replace(dtype=(dt:=select_dtype(v)), src=tuple(s.src[0].cast(dt.scalar()) for s in v.src)).cast(dtypes.weakint)),
# special can only be int32
(UPat(Ops.SPECIAL, src=(UPat.var("var").cast(dtypes.index),), name="u"), lambda u,var: u.replace(dtype=dtypes.int, src=(var,)).cast(dtypes.index)),
(UPat(Ops.DEFINE_VAR, dtype=dtypes.index, name="u"), lambda u: u.replace(dtype=dtypes.int).cast(dtypes.index)),
(UPat(Ops.BIND, src=(UPat.var("var").cast(dtypes.index), UPat.cvar("val").cast(dtypes.index))), lambda var,val: var.bind(val).cast(dtypes.index)),
(UPat(Ops.SPECIAL, src=(UPat.var("var").cast(dtypes.weakint),), name="u"),
lambda u,var: u.replace(dtype=dtypes.int, src=(var,)).cast(dtypes.weakint)),
(UPat(Ops.DEFINE_VAR, dtype=dtypes.weakint, name="u"), lambda u: u.replace(dtype=dtypes.int).cast(dtypes.weakint)),
(UPat(Ops.BIND, src=(UPat.var("var").cast(dtypes.weakint), UPat.cvar("val").cast(dtypes.weakint))),
lambda var,val: var.bind(val).cast(dtypes.weakint)),
# lower Invalid
(UPat.var("buf").index(UPat.var("cond").where(UPat.var("idx"), UPat(Ops.CONST, arg=Invalid))), lambda buf,idx,cond: buf.index(idx, cond, ptr=True)),
# remove hanging casts
@@ -1456,7 +1465,10 @@ pm_lower_index_dtype = PatternMatcher([
(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx", dtypes.ints).cast(), UPat.var("valid"))),
lambda buf,idx,valid: buf.index(idx, valid, ptr=True)),
(UPat((Ops.SINK, Ops.NOOP, Ops.END), name="n"),
lambda n: n.replace(src=tuple(s.src[0] if s.op is Ops.CAST and s.dtype == dtypes.index else s for s in n.src))),
lambda n: n.replace(src=tuple(s.src[0] if s.op is Ops.CAST and s.dtype == dtypes.weakint else s for s in n.src))),
# vectorized indexes (ie. images) must be int
(UPat(Ops.INDEX, src=(UPat(), UPat(Ops.VECTORIZE, dtypes.long, name="vec")), allow_any_len=True, name="idx"),
lambda idx,vec: idx.replace(src=(idx.src[0], UOp.vectorize(*(u.cast(dtypes.int) for u in vec.src)), *idx.src[2:])))
])
def _index_to_concrete_int(u:UOp) -> UOp: return graph_rewrite(u.sink(), pm_lower_index_dtype).src[0]
@@ -1543,7 +1555,7 @@ pm_pyrender_extra = PatternMatcher([
(UPat(Ops.CONST, src=(UPat(Ops.DEVICE, name="d"),), name="x"), lambda x,d: f"UOp.const({x.dtype}, {x.arg}, device={repr(d.arg)})"),
(UPat(Ops.CONST, name="x"), lambda x: f"UOp.const({x.dtype}, {x.arg})"),
(UPat(Ops.DEFINE_VAR, src=(), name="x"), lambda x:
f"UOp.variable(\"{x.arg[0]}\", {x.arg[1]}, {x.arg[2]}{', dtype='+str(x.dtype) if x.dtype is not dtypes.index else ''})"),
f"UOp.variable(\"{x.arg[0]}\", {x.arg[1]}, {x.arg[2]}{', dtype='+str(x.dtype) if x.dtype is not dtypes.weakint else ''})"),
(UPat((Ops.CAST, Ops.BITCAST), name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({x.dtype})"),
(UPat(Ops.SPECIAL, src=(UPat(Ops.CONST),), name="x"), lambda x: f"UOp.special({x.src[0].arg}, {repr(x.arg)}, dtype={x.dtype})"),
(UPat(Ops.BUFFER, src=(UPat(Ops.UNIQUE, name="u"), UPat(Ops.DEVICE, name="d")), name="x"), lambda x,u,d:
@@ -1554,7 +1566,7 @@ pm_pyrender_extra = PatternMatcher([
# NOTE: range has srcs sometimes after control flow
(UPat(Ops.RANGE, src=(UPat(Ops.CONST, name="c"),), allow_any_len=True, name="x"), lambda ctx,x,c:
"UOp.range("+', '.join([str(c.arg)] + [repr(y) for y in x.arg])+
(f', src={srcs(ctx, x.src[1:])}' if len(x.src) > 1 else '')+(', dtype='+str(x.dtype) if x.dtype is not dtypes.index else '')+")"),
(f', src={srcs(ctx, x.src[1:])}' if len(x.src) > 1 else '')+(', dtype='+str(x.dtype) if x.dtype is not dtypes.weakint else '')+")"),
# TODO: index shouldn't mismatch dtype
(UPat(Ops.INDEX, src=(UPat(), UPat()), allow_any_len=True, name="x"), lambda ctx,x:
f"{ctx[x.src[0]]}.index({ctx[x.src[1]]}, "+''.join([f"{ctx[xx]}, " for xx in x.src[2:]])+
@@ -1618,9 +1630,8 @@ def pyrender(ast:UOp) -> str:
if len(cmap[u]) == 1 and len([x for x in list(cmap[u].keys())[0].src if x is u]) == 1 and u.op not in always_rendered: continue
to_render.add(u)
kernels: dict[UOp, tuple[str, str]] = {}
ret = _render_with_splits(lst, pm_pyrender, to_render)
return ''.join([v[1] for v in kernels.values()]) + '\n'.join([f"{k} = {strip_parens(v)}" for k,v in ret.items()])
return '\n'.join([f"{k} = {strip_parens(v)}" for k,v in ret.items()])
# *** what was symbolic.py ***
tinygrad/uop/spec.py
+22 -18
View File
@@ -54,7 +54,7 @@ shared_spec = PatternMatcher([
(UPat(Ops.RANGE, src=(UPat.var("x"),), allow_any_len=True, name="rng"), lambda rng,x:
rng.dtype == x.dtype and isinstance(rng.arg, tuple) and len(rng.arg) >= 2 and \
all(isinstance(ra, int) for ra in rng.arg[0:-1]) and isinstance(rng.arg[-1], AxisType)),
(UPat(Ops.INDEX, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(y.dtype == dtypes.index for y in x.src[1:]) or None),
(UPat(Ops.INDEX, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(y.dtype == dtypes.weakint for y in x.src[1:]) or None),
# RANGE/SPECIAL define loops, END closes them
(UPat(Ops.END, src=(UPat(), UPat(Ops.RANGE))), lambda: True),
@@ -69,17 +69,17 @@ shared_spec = PatternMatcher([
# ***** UOp spec in the Tensor graph *****
movement_ops = PatternMatcher([
(UPat((Ops.RESHAPE, Ops.EXPAND), name="mv", src=(UPat.var("x"), UPat(dtype=dtypes.index))), lambda mv,x: True),
(UPat((Ops.PAD, Ops.SHRINK), name="mv", src=(UPat.var("x"), UPat(dtype=dtypes.index), UPat(dtype=dtypes.index))), lambda mv,x: True),
(UPat((Ops.RESHAPE, Ops.EXPAND), name="mv", src=(UPat.var("x"), UPat(dtype=dtypes.weakint))), lambda mv,x: True),
(UPat((Ops.PAD, Ops.SHRINK), name="mv", src=(UPat.var("x"), UPat(dtype=dtypes.weakint), UPat(dtype=dtypes.weakint))), lambda mv,x: True),
(UPat((Ops.PERMUTE, Ops.FLIP), name="mv", src=(UPat.var("x"),)), lambda mv,x: isinstance(mv.arg, tuple)),
# inputs to movement ops
(UPat((Ops.VECTORIZE, Ops.VCONST), dtype=dtypes.index), lambda: True),
(UPat({Ops.ADD, Ops.MUL, Ops.IDIV}, dtype=dtypes.index), lambda: True),
(UPat((Ops.VECTORIZE, Ops.VCONST), dtype=dtypes.weakint), lambda: True),
(UPat({Ops.ADD, Ops.MUL, Ops.IDIV}, dtype=dtypes.weakint), lambda: True),
# AFTER on Movement Op, BUFFER, COPY, or BITCAST
(UPat(Ops.AFTER, src=(UPat(GroupOp.Movement.union({Ops.MULTI, Ops.CONTIGUOUS, Ops.BUFFER, Ops.BITCAST, Ops.COPY})),), allow_any_len=True),
lambda: True),
# AFTER on Movement Op, INDEX, BUFFER, COPY, or BITCAST
(UPat(Ops.AFTER, src=(UPat(GroupOp.Movement.union({Ops.INDEX, Ops.MULTI, Ops.CONTIGUOUS, Ops.BUFFER, Ops.BITCAST, Ops.COPY})),),
allow_any_len=True), lambda: True),
])
_tensor_spec = PatternMatcher([
@@ -89,7 +89,7 @@ _tensor_spec = PatternMatcher([
(UPat(Ops.DEVICE, dtypes.void, (), name="d"), lambda d:
isinstance(d.arg, str) or (isinstance(d.arg, tuple) and all(isinstance(s, str) for s in d.arg))),
(UPat(Ops.BUFFER, src=(UPat((Ops.LUNIQUE, Ops.UNIQUE)), UPat(Ops.DEVICE)), name="buf"),
lambda buf: isinstance(buf.arg, int) and isinstance(buf.dtype, (DType, ImageDType))),
lambda buf: isinstance(buf.arg, int) and isinstance(buf.dtype, DType)),
# BUFFER_VIEW on BUFFER is allowed if BUFFER is
(UPat(Ops.BUFFER_VIEW, src=(UPat(Ops.BUFFER),)), lambda: True),
@@ -104,9 +104,9 @@ _tensor_spec = PatternMatcher([
(UPat(Ops.MSTACK, name="x"), lambda x: all(isinstance(x.device, str) for x in x.src)),
# Tensor variable bindings
(UPat(Ops.BIND, (dtypes.int,dtypes.index,), (UPat(Ops.DEFINE_VAR), UPat.cvar(dtype=(dtypes.int,dtypes.index,))), arg=None), lambda: True),
(UPat(Ops.BIND, (dtypes.int,dtypes.weakint,), (UPat(Ops.DEFINE_VAR), UPat.cvar(dtype=(dtypes.int,dtypes.weakint,))), arg=None), lambda: True),
# single-src BIND used for schedule cache key normalization
(UPat(Ops.BIND, (dtypes.int,dtypes.index,), (UPat(Ops.DEFINE_VAR),), arg=None), lambda: True),
(UPat(Ops.BIND, (dtypes.int,dtypes.weakint,), (UPat(Ops.DEFINE_VAR),), arg=None), lambda: True),
# device or unique
(UPat(Ops.CONST, src=(UPat(Ops.DEVICE),)), lambda: True),
@@ -189,7 +189,7 @@ shared_codegen_spec = PatternMatcher([
(UPat(Ops.INDEX, src=(UPat(GroupOp.Defines|{Ops.AFTER}, name="buf"), UPat.var("idx"), UPat.var("gate", dtype=dtypes.bool))), validate_index),
# SPECIAL
(UPat(Ops.SPECIAL, src=(UPat.var("x", (dtypes.index, dtypes.int32)),), name="s"), lambda s,x: s.dtype == x.dtype and isinstance(s.arg, str)),
(UPat(Ops.SPECIAL, src=(UPat.var("x", (dtypes.weakint, dtypes.int32)),), name="s"), lambda s,x: s.dtype == x.dtype and isinstance(s.arg, str)),
# BARRIER (on any length)
(UPat(Ops.BARRIER, dtypes.void), lambda: True),
@@ -199,12 +199,16 @@ shared_codegen_spec = PatternMatcher([
kernel_spec = PatternMatcher([
# index is allowed here
(UPat(GroupOp.Elementwise|{Ops.CONST, Ops.RANGE, Ops.DEFINE_VAR}, dtype=dtypes.index), lambda: True),
(UPat(GroupOp.Elementwise|{Ops.CONST, Ops.RANGE, Ops.DEFINE_VAR}, dtype=dtypes.weakint), lambda: True),
# UNROLL/CONTRACT is used here for WMMA
(UPat(Ops.CONTRACT, name="x"), lambda x: x.dtype.count == prod(y[1] for y in x.arg)),
(UPat(Ops.UNROLL, name="x"), lambda x: x.src[0].dtype.count == prod(y[1] for y in x.arg)),
# SHAPED_WMMA has <a, b, acc> with shaped inputs, arg=((M,N,K), device, threads), lowered to WMMA+CONTRACT later
(UPat(Ops.SHAPED_WMMA, src=(UPat(), UPat(), UPat()), name="x"),
lambda x: isinstance(x.arg, tuple) and len(x.arg) == 3 and isinstance(x.arg[0], tuple)),
# END can end multiple axes here
(UPat(Ops.END, src=(UPat(), UPat()), allow_any_len=True), lambda: True),
@@ -212,7 +216,7 @@ kernel_spec = PatternMatcher([
(UPat(Ops.BUFFERIZE, src=(UPat(),), allow_any_len=True, name="x"), lambda x: True),
# reduce must be on ranges
(UPat(Ops.REDUCE, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(y.dtype in (dtypes.index, dtypes.int) for y in x.src[1:])),
(UPat(Ops.REDUCE, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(y.dtype in (dtypes.weakint, dtypes.int) for y in x.src[1:])),
# COPY/BUFFER_VIEW can have ranges appended
(UPat(Ops.COPY, name="x", src=(UPat.var("s"), UPat(Ops.DEVICE)), allow_any_len=True, arg=None),
@@ -236,7 +240,7 @@ program_spec = PatternMatcher([
(UPat(Ops.END, src=(UPat(), UPat(Ops.RANGE)), dtype=dtypes.void), lambda: True),
# make sure all index dtypes have been lowered (except CONST/RANGE/DEFINE_VAR which are valid index-typed)
(UPat(GroupOp.All-{Ops.CONST, Ops.RANGE, Ops.DEFINE_VAR, Ops.VCONST, Ops.VECTORIZE}, dtype=dtypes.index), lambda: False),
(UPat(GroupOp.All-{Ops.CONST, Ops.RANGE, Ops.DEFINE_VAR, Ops.VCONST, Ops.VECTORIZE}, dtype=dtypes.weakint), lambda: False),
(UPat(Ops.CONST, arg=Invalid), lambda: False),
(UPat(Ops.VCONST, name="x"), lambda x: all(v is not Invalid for v in x.arg) and len(x.arg)==x.dtype.vcount>1 and
type(x.arg) is type(x.dtype.const(x.arg))),
@@ -273,13 +277,13 @@ full_spec = PatternMatcher([
(UPat(Ops.CALL, dtype=dtypes.void), lambda: True),
# where on index in rhs position is fine
(UPat(Ops.WHERE, dtype=dtypes.index, src=(UPat(dtype=dtypes.bool), UPat(), UPat(dtype=dtypes.index))), lambda: True),
(UPat(Ops.WHERE, dtype=dtypes.weakint, src=(UPat(dtype=dtypes.bool), UPat(), UPat(dtype=dtypes.weakint))), lambda: True),
# allow index dtype on a restricted set of UOps
(UPat((Ops.ADD, Ops.MUL, Ops.MOD, Ops.IDIV, Ops.MAX,
Ops.SPECIAL, Ops.CAST, Ops.RANGE, Ops.VCONST, Ops.VECTORIZE), dtype=dtypes.index), lambda: True),
Ops.SPECIAL, Ops.CAST, Ops.RANGE, Ops.VCONST, Ops.VECTORIZE), dtype=dtypes.weakint), lambda: True),
# while BIND is being casted
(UPat(Ops.BIND, (dtypes.int, dtypes.index), (UPat(), UPat()), arg=None), lambda: True),
(UPat(Ops.BIND, (dtypes.int, dtypes.weakint), (UPat(), UPat()), arg=None), lambda: True),
# in progress MSTACK may lose device
(UPat((Ops.MSELECT, Ops.MSTACK), name="x"), lambda x: True),
tinygrad/uop/symbolic.py
+21 -20
View File
@@ -51,15 +51,15 @@ def fold_add_divmod_recombine(x:UOp) -> UOp|None:
# this needs to be before symbolic so that 0*something_that_might_be_invalid doesnt become 0
propagate_invalid = PatternMatcher([
# propagate invalid, push it past children
(invalid_gate.cast(name="cast"), lambda i,x,cond,cast: x.cast(cast.dtype) if i.dtype is dtypes.index else None),
(invalid_gate.cast(name="cast"), lambda i,x,cond,cast: x.cast(cast.dtype) if i.dtype is dtypes.weakint else None),
(UPat(GroupOp.Unary, src=(invalid_gate,), name="alu"), lambda cond,x,alu,i: cond.where(x.alu(alu.op), i)),
(UPat(GroupOp.Binary-GroupOp.Comparison, src=(invalid_gate, UPat.var("y")), name="alu"), lambda cond,x,y,alu,i: cond.where(x.alu(alu.op,y), i)),
(UPat(GroupOp.Binary-GroupOp.Comparison, src=(UPat.var("y"), invalid_gate), name="alu"), lambda cond,x,y,alu,i: cond.where(y.alu(alu.op,x), i)),
# TODO: when can this happen? and is it always safe to just drop invalid?
(UPat(GroupOp.Comparison, src=(invalid_gate, UPat.var("y")), name="alu"), lambda cond,x,y,alu,i:
x.alu(alu.op,y) if i.dtype is dtypes.index else cond.where(x.alu(alu.op,y), i.cast(dtypes.bool))),
x.alu(alu.op,y) if i.dtype is dtypes.weakint else cond.where(x.alu(alu.op,y), i.cast(dtypes.bool))),
(UPat(GroupOp.Comparison, src=(UPat.var("y"), invalid_gate), name="alu"), lambda cond,x,y,alu,i:
y.alu(alu.op,x) if i.dtype is dtypes.index else cond.where(y.alu(alu.op,x), i.cast(dtypes.bool))),
y.alu(alu.op,x) if i.dtype is dtypes.weakint else cond.where(y.alu(alu.op,x), i.cast(dtypes.bool))),
# alu with invalid -> invalid
(UPat(GroupOp.Unary, src=(invalid_pat,)), lambda i: i),
(UPat(GroupOp.Binary-GroupOp.Comparison, src=[invalid_pat, UPat()]), lambda i: i),
@@ -77,26 +77,26 @@ symbolic_simple = propagate_invalid + PatternMatcher([
# ** self folding **
(UPat.var("x") + 0, lambda x: x), # x+0 -> x
(UPat.var("x") * 1, lambda x: x), # x*1 -> x
(UPat.var("x", dtype=dtypes.ints+(dtypes.bool, dtypes.index)) ^ 0, lambda x: x), # x^0 -> x
(UPat.var("x", dtype=dtypes.ints+(dtypes.bool, dtypes.weakint)) ^ 0, lambda x: x), # x^0 -> x
(UPat.var("x") // UPat.var("x"), lambda x: x.const_like(1)), # x//x -> 1
(UPat.var("x") // 1, lambda x: x), # x//1 -> x
(UPat.var("x") // -1, lambda x: -x), # x//-1 -> -x
((UPat.var() % UPat.var("y")).named("base") % UPat.var("y"), lambda base,y: base), # (x%y)%y = -> x%y (rewritten with base for speed)
# variations of (x%c)+(x//c)*c = x
(UPat(Ops.ADD, dtype=dtypes.index, name="x"), fold_add_divmod_recombine),
(UPat(Ops.ADD, dtype=dtypes.weakint, name="x"), fold_add_divmod_recombine),
(UPat.var("x", dtype=dtypes.bool) & UPat.cvar("c", vec=False), lambda x,c: x if c.arg else c),
(UPat.var("x", dtype=dtypes.bool) | UPat.cvar("c", vec=False), lambda x,c: c if c.arg else x),
(UPat(GroupOp.Idempotent, src=(UPat.var("x"), UPat.var("x"))), lambda x: x),
(UPat.var("x", dtype=dtypes.bool).logical_not().logical_not(), lambda x: x),
(UPat.var("x", dtype=dtypes.bool).where(UPat.const(dtypes.bool, True), UPat.const(dtypes.bool, False)), lambda x: x),
(UPat.var("x", dtype=dtypes.bool).where(UPat.const(dtypes.bool, False), UPat.const(dtypes.bool, True)), lambda x: x.logical_not()),
(UPat.var("x", dtype=dtypes.ints+(dtypes.bool, dtypes.index)).trunc(), lambda x: x),
(UPat.var("x", dtype=dtypes.ints+(dtypes.bool, dtypes.weakint)).trunc(), lambda x: x),
# ** zero folding **
(UPat.var("x") < UPat.var("x"), lambda x: x.const_like(False).cast(dtypes.bool.vec(x.dtype.count))), # x < x -> False
(UPat.var("x") % UPat.var("x"), lambda x: x.const_like(0)), # x%x -> 0
(UPat.var("x") ^ UPat.var("x"), lambda x: x.const_like(0)), # x^x -> 0
(UPat.var("x") & 0, lambda x: x.const_like(0)), # x&0 -> 0
(UPat.var("x", dtype=dtypes.ints+(dtypes.bool, dtypes.index)) != UPat.var("x"),
(UPat.var("x", dtype=dtypes.ints+(dtypes.bool, dtypes.weakint)) != UPat.var("x"),
lambda x: x.const_like(False).cast(dtypes.bool.vec(x.dtype.count))), # x != x -> False (only ints)
# ** constant folding **
# TODO: add const folding for Ops.THREEFRY
@@ -192,7 +192,7 @@ gep_pushing = PatternMatcher([
# GEP in order is removed
(UPat(Ops.GEP, name="g"), lambda g: g.src[0] if not isinstance(g.dtype, PtrDType) and g.arg == tuple(range(g.src[0].dtype.count)) else None),
# push all GEPs through ALUs for index (TODO: remove this)
(UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST), name='alu').f(Ops.GEP, dtype=dtypes.index, name='gep'),
(UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST), name='alu').f(Ops.GEP, dtype=dtypes.weakint, name='gep'),
lambda gep,alu: UOp(alu.op, alu.dtype.scalar().vec(gep.dtype.count), tuple(x.gep(gep.arg) for x in alu.src), alu.arg) \
if not isinstance(gep.dtype, PtrDType) and not isinstance(alu.dtype, PtrDType) else None),
# CAT can't be rendered. it's a VECTORIZE on vectors, we expand to a single VECTORIZEs with GEPs (TODO: move this later)
@@ -207,7 +207,7 @@ gep_pushing = PatternMatcher([
commutative = PatternMatcher([
# ** COMMUTATIVE flipping (only for index) **
# NOTE: this can break merging vector math by only flipping some of them
(UPat(GroupOp.Commutative, dtype=dtypes.index, name='x'), lambda x: x.replace(src=x.src[::-1]) if x.src[1].tuplize < x.src[0].tuplize else None),
(UPat(GroupOp.Commutative, dtype=dtypes.weakint, name='x'), lambda x: x.replace(src=x.src[::-1]) if x.src[1].tuplize < x.src[0].tuplize else None),
])
symbolic = symbolic_simple+commutative+PatternMatcher([
@@ -224,7 +224,7 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
((UPat.var("y") + UPat.var("x")) + UPat.var("x"), lambda y,x: y+x*2),
((UPat.var("x") / UPat.var("x2")) / UPat.var("x3"), lambda x,x2,x3: x/(x2*x3) if x2 is not x3 else None), # (x/x2)/x3 -> x/(x2*x3)
(-1 * (UPat.var("x") + UPat.cvar("c")), lambda x,c: (-x)+(-c)), # -(x+c) -> -x + -c
(UPat.cvar("y") * (UPat.var("x", dtype=dtypes.index) + UPat.cvar("c")), lambda x,y,c: (y*x)+(y*c)), # y*(x+c) -> y*x + y*c
(UPat.cvar("y") * (UPat.var("x", dtype=dtypes.weakint) + UPat.cvar("c")), lambda x,y,c: (y*x)+(y*c)), # y*(x+c) -> y*x + y*c
# ** where folding **
(UPat.var("cond", dtype=dtypes.bool).logical_not().where(UPat.var("t"), UPat.var("f")),
lambda cond, t, f: cond.where(f,t) if f.arg is not Invalid else None),
@@ -249,35 +249,35 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
((UPat.var("x") // UPat.cvar("c1")) // UPat.cvar("c2"), lambda x,c1,c2: x//(c1*c2)), # (x//c1)//c2 -> x//(c1*c2)
# ** lt **
# c0*x<c1 for positive int c0,c1
((UPat.cvar("c0", vec=False)*UPat.var("x", dtype=dtypes.index))<UPat.cvar("c1", vec=False),
((UPat.cvar("c0", vec=False)*UPat.var("x", dtype=dtypes.weakint))<UPat.cvar("c1", vec=False),
lambda x,c0,c1: x<math.ceil(c1.arg/c0.arg) if c0.arg > 0 and c1.arg > 0 else None),
# c0*x<c1 for negative int c0 and non-positive c1
((UPat.cvar("c0", vec=False)*UPat.var("x", dtype=dtypes.index))<UPat.cvar("c1", vec=False),
((UPat.cvar("c0", vec=False)*UPat.var("x", dtype=dtypes.weakint))<UPat.cvar("c1", vec=False),
lambda x,c0,c1: (-x)<(-(math.floor(-c1.arg/-c0.arg))) if c0.arg < 0 and c0.arg != -1 and c1.arg <= 0 else None),
# x//d<c
((UPat.var("x", dtype=dtypes.index)//UPat.cvar("d", vec=False))<UPat.cvar("c", vec=False),
((UPat.var("x", dtype=dtypes.weakint)//UPat.cvar("d", vec=False))<UPat.cvar("c", vec=False),
lambda x,d,c: (x<(c.arg*d.arg) if c.arg > 0 else x<(c.arg*d.arg-(d.arg-1))) if d.arg > 0 else None),
# ** move add/mul consts to end (NOTE: this is still happening before constant folding) **
((UPat.var("x") + UPat.cvar("c1")) + UPat.var("y"), lambda x,c1,y: (x+y)+c1),
((UPat.var("x") * UPat.cvar("c1")) * UPat.var("y"), lambda x,c1,y: (x*y)*c1),
# *** rules from symbolic ***
# generic lt folding
(UPat.var("x", dtypes.index)<UPat.cvar("c", vec=False), lambda x,c: lt_folding(x, c.arg) if 0 < c.arg else None),
(UPat.var("x", dtypes.index)*-1 < UPat.var("y")*-1, lambda x,y: y<x),
(UPat.var("x", dtypes.weakint)<UPat.cvar("c", vec=False), lambda x,c: lt_folding(x, c.arg) if 0 < c.arg else None),
(UPat.var("x", dtypes.weakint)*-1 < UPat.var("y")*-1, lambda x,y: y<x),
# canonicalize a simplex with positive coefficients > 0. NOTE: not x < 1 means x > 0
((UPat.var("x", dtypes.index)<1).ne(True), lambda x: (newx<1).ne(True) if (newx:=canonicalize_simplex(x)) is not None else None),
((UPat.var("x", dtypes.weakint)<1).ne(True), lambda x: (newx<1).ne(True) if (newx:=canonicalize_simplex(x)) is not None else None),
# a range mod its own upper bound is just the range
(UPat(Ops.RANGE, src=UPat.var("end"), name="r")%UPat.var("end"), lambda r,end: r),
(UPat(Ops.RANGE, src=UPat.var("end"), name="r")//UPat.var("end"), lambda r,end: r.const_like(0)),
# cast/long folding
# if the intermediate cast doesnt narrow we can do it in one cast
(UPat.var('x').cast(name="a").cast(name="b"), lambda x,a,b: x.cast(b.dtype) if can_lossless_cast(x.dtype, a.dtype) else None),
(UPat.var('x', dtypes.ints+(dtypes.index,)).cast(dtypes.ints+(dtypes.index,), name="a").cast(name="b"),
(UPat.var('x', dtypes.ints+(dtypes.weakint,)).cast(dtypes.ints+(dtypes.weakint,), name="a").cast(name="b"),
lambda x,a,b: x.cast(b.dtype) if a.dtype.min<=x.vmin and x.vmax<=a.dtype.max else None),
# try to do math in int instead of long
(UPat(GroupOp.Binary, src=(UPat.var("x", dtypes.long), UPat.var("y", dtypes.long)), name="u"), lambda u,x,y:
x.cast(dtypes.int).alu(u.op, y.cast(dtypes.int)).cast(u.dtype) if not any(v.overflows(dtypes.int) for v in (u,x,y)) else None),
((UPat.var("x", dtypes.index) + UPat.cvar("c")).cast(dtypes.sints, name="cast"), lambda x,c,cast:x.cast(cast.dtype)+c.cast(cast.dtype)),
((UPat.var("x", dtypes.weakint) + UPat.cvar("c")).cast(dtypes.sints, name="cast"), lambda x,c,cast:x.cast(cast.dtype)+c.cast(cast.dtype)),
# only RANGE/IF/STORE/KERNEL have side effects
(UPat(Ops.AFTER, name="x"), lambda x: x.replace(src=(x.src[0],)+
tuple(flatten([(y,) if y.op in {Ops.RANGE, Ops.STORE, Ops.CALL, Ops.BARRIER, Ops.END, Ops.UNROLL, Ops.LINEAR, Ops.BUFFERIZE}
@@ -394,6 +394,7 @@ pm_move_where_on_load = PatternMatcher([
])
def gated_given_valid(cond:UOp, x:UOp, i:UOp) -> UOp|None:
if x.dtype.scalar() is not dtypes.weakint: return None
# Skip if x contains DIV/MOD AND IMAGE mode is enabled -> image index e.g. openpilot
if IMAGE.value > 0 and x.op_in_backward_slice_with_self(Ops.IDIV, Ops.MOD): return None
return cond.where(uop_given_valid(cond, x, try_simplex=False), i)
@@ -432,7 +433,7 @@ sym = symbolic+pm_simplify_valid+PatternMatcher([
UPat.load(UPat(Ops.INDEX, name="index"))), allow_any_len=True, name="store"),
lambda index, gate, alt, store: UOp.store(index.src[0].index(gate.where(index.src[1], UOp.invalid())), alt, *store.src[2:])),
# fold gated LOAD/STORE
(UPat((Ops.LOAD, Ops.STORE), src=(UPat().index(UPat.const(dtypes.index, Invalid)).or_casted(),), allow_any_len=True, name="x"),
(UPat((Ops.LOAD, Ops.STORE), src=(UPat().index(UPat.const(dtypes.weakint, Invalid)).or_casted(),), allow_any_len=True, name="x"),
lambda x: UOp(Ops.NOOP) if x.op is Ops.STORE else x.const_like(0)), # invalid store does nothing. invalid load produces 0
(UPat(Ops.STORE, src=(UPat(), invalid_pat), allow_any_len=True), lambda i: UOp(Ops.NOOP)),
# store of where with invalid -> gated store
@@ -456,5 +457,5 @@ sym = symbolic+pm_simplify_valid+PatternMatcher([
# ** combine terms (opinionated) **
(-1 * (UPat.var("x") + UPat.var("y")), lambda x,y: (-x)+(-y)), # -(x+y) -> -x + -y
# (x+y)*c -> x*c+y*c. only for int, float has inf*0=nan issue
((UPat.var("x", dtypes.index) + UPat.var("y")) * UPat.cvar("c"), lambda x,y,c: x*c+y*c),
((UPat.var("x", dtypes.weakint) + UPat.var("y")) * UPat.cvar("c"), lambda x,y,c: x*c+y*c),
])
tinygrad/uop/validate.py
+7 -7
View File
@@ -20,33 +20,33 @@ def create_bounded(name:str, vmin:int, vmax:int, z3ctx:z3.Context) -> tuple[z3.A
return (s:=z3.Int(name, ctx=z3ctx)), (vmin <= s)&(s <= vmax)
z3_renderer = PatternMatcher([
(UPat.var("cond").where(UPat.var("x"), UPat.const(dtypes.index, Invalid)), lambda x,cond,ctx: (ctx[1][x], ctx[1][cond])),
(UPat.var("cond").where(UPat.var("x"), UPat.const(dtypes.weakint, Invalid)), lambda x,cond,ctx: (ctx[1][x], ctx[1][cond])),
# variables
(UPat(Ops.SPECIAL, name="x"), lambda x,ctx: create_bounded(x.arg, 0, ctx[1][x.src[0]]-1, ctx[0])),
(UPat(Ops.DEFINE_VAR, name="x"), lambda x,ctx: create_bounded(x.arg[0], x.arg[1], x.arg[2], ctx[0])),
(UPat(Ops.RANGE, name="x"), lambda x,ctx: create_bounded(x.render(simplify=False), 0, ctx[1][x.src[0]]-1, ctx[0])),
# loads are variables bounded by the min/max of the dtype. non-pointer INDEX is also a LOAD
(UPat((Ops.LOAD, Ops.INDEX), dtypes.ints+(dtypes.index,), name="x"), lambda x,ctx:
(UPat((Ops.LOAD, Ops.INDEX), dtypes.ints+(dtypes.weakint,), name="x"), lambda x,ctx:
create_bounded(f"load{len(ctx[1])}", x.dtype.min, x.dtype.max, ctx[0])),
(UPat((Ops.LOAD, Ops.INDEX), dtypes.bool, name="x"), lambda x,ctx: (z3.Bool(f"load{len(ctx[1])}", ctx=ctx[0]), None)),
# constants
(UPat(Ops.CONST, arg=Invalid, name="x"), lambda x,ctx: (z3.Int("Invalid", ctx=ctx[0]), None)),
(UPat(Ops.CONST, dtypes.ints+(dtypes.index,), name="x"), lambda x,ctx: (z3.IntVal(x.arg, ctx=ctx[0]), None)),
(UPat(Ops.CONST, dtypes.ints+(dtypes.weakint,), name="x"), lambda x,ctx: (z3.IntVal(x.arg, ctx=ctx[0]), None)),
(UPat(Ops.CONST, dtypes.bool, name="x"), lambda x,ctx: (z3.BoolVal(x.arg, ctx=ctx[0]), None)),
# casts from floats create new variables
(UPat(Ops.CAST, dtypes.ints+(dtypes.index,), src=(UPat(dtype=dtypes.floats),), name="x"), lambda x,ctx:
(UPat(Ops.CAST, dtypes.ints+(dtypes.weakint,), src=(UPat(dtype=dtypes.floats),), name="x"), lambda x,ctx:
create_bounded(f"cast{len(ctx[1])}", x.dtype.min, x.dtype.max, ctx[0])),
# A comparison between floats introduces a new bool variable
(UPat(GroupOp.Comparison, src=UPat(dtype=dtypes.floats), name="x"), lambda x,ctx: (z3.Bool(f"float_cmp{len(ctx[1])}", ctx=ctx[0]), None)),
# casts from bool/int to int/bool
(UPat(Ops.CAST, dtypes.ints+(dtypes.index,),src=(UPat.var("x", dtypes.bool),), name="c"), lambda x,c,ctx: (z3.If(ctx[1][x], 1, 0), None)),
(UPat(Ops.CAST, dtypes.ints+(dtypes.index,), src=(UPat.var("x", dtypes.ints+(dtypes.index,)),), name="c"), lambda x,c,ctx: (ctx[1][x], None)),
(UPat(Ops.CAST, dtypes.ints+(dtypes.weakint,),src=(UPat.var("x", dtypes.bool),), name="c"), lambda x,c,ctx: (z3.If(ctx[1][x], 1, 0), None)),
(UPat(Ops.CAST, dtypes.ints+(dtypes.weakint,), src=(UPat.var("x", dtypes.ints+(dtypes.weakint,)),), name="c"), lambda x,c,ctx: (ctx[1][x], None)),
(UPat(Ops.CAST, dtypes.bool, name="x"), lambda x,ctx: (ctx[1][x.src[0]]!=0, None)),
(UPat(GroupOp.ALU, name="x"), lambda x,ctx: (z3_alu[x.op](*(ctx[1][s] for s in x.src)), None)),
])
def uops_to_z3(solver:z3.Solver, *uops: UOp) -> list[z3.ExprRef]:
lst = list(UOp.sink(*uops).toposort(gate=lambda x: x.dtype.scalar() in dtypes.ints+(dtypes.bool, dtypes.index) or x.op is Ops.SINK))[:-1]
lst = list(UOp.sink(*uops).toposort(gate=lambda x: x.dtype.scalar() in dtypes.ints+(dtypes.bool, dtypes.weakint) or x.op is Ops.SINK))[:-1]
z3map: dict[UOp, z3.ExprRef] = {}
for u in lst:
z3_rewritten = z3_renderer.rewrite(u, ctx=(solver.ctx, z3map))
tinygrad/viz/js/index.js
+29 -6
View File
@@ -25,6 +25,13 @@ const colored = n => d3.create("span").call(s => s.selectAll("span").data(typeof
const rect = (s) => (typeof s === "string" ? document.querySelector(s) : s).getBoundingClientRect();
// dims of shapes on the canvas aren't tracked by the browser, we compute it
const canvasRect = (s, pixelScale) => {
const { e } = selectShape(s), t = data.tracks.get(s.split("-")[0]);
const x = pixelScale(e.x), w = pixelScale(e.x+e.width)-x, y = t.offsetY+e.y;
return {x0:x, x1:x+w, y0:y, y1:y+e.height};
};
let timeout = null;
const Status = {STARTED:0, COMPLETE:1, ERR:2}
const updateProgress = (st, msg) => {
@@ -193,7 +200,7 @@ function formatCycles(cycles) {
const formatUnit = (d, unit="") => d3.format(".3~s")(d)+unit;
const WAVE_COLORS = {VALU:"#ffffc0", SALU:"#cef263", LOAD:"#ffc0c0", STORE:"#4fa3cc", IMMEDIATE:"#f3b44a", BARRIER:"#d00000", JUMP:"#ffb703",
JUMP_NO:"#fb8500", MESSAGE:"#90dbf4", VMEM:"#b2b7c9", LDS:"#9fb4a6"};
JUMP_NO:"#fb8500", MESSAGE:"#90dbf4", VMEM:"#b2b7c9", LDS:"#9fb4a6", WAVERDY:"#1a2a2a"};
const waveColor = (op) => {
const cat = op.includes("VALU") || op === "VINTERP" ? "VALU" : op.includes("SALU") ? "SALU" : op.includes("VMEM") ? "VMEM"
: op.includes("LOAD") || op === "SMEM" ? "LOAD" : op.includes("STORE") ? "STORE" : op;
@@ -292,6 +299,8 @@ function setFocus(key) {
const [st, et] = xscale.range().map(zoomLevel.invertX, zoomLevel).map(xscale.invert, xscale);
if (x1 < st || x0 > et) zoomLevel = d3.zoomIdentity.translate(-xscale((x0+x1)/2-(et-st)/2)*zoomLevel.k, 0).scale(zoomLevel.k);
}
const link = e?.arg.link ?? data.links.get(key);
data.link = link == null ? null : [key, link];
focusedShape = key; d3.select("#timeline").call(canvasZoom.transform, zoomLevel);
}
const { eventType, e } = selectShape(key);
@@ -352,8 +361,10 @@ function setFocus(key) {
metadata.insertBefore(instList.parentElement, html.node());
}
d3.select(instList).selectAll("span").classed("highlight", false);
const instLine = document.getElementById(`inst-${e?.arg.pc}`); instLine?.classList.add("highlight");
let instLine = document.getElementById(`inst-${e?.arg.pc}`);
if (instLine == null && data.link != null) instLine = document.getElementById(`inst-${selectShape(data.link[1]).e.arg.pc}`);
if (instLine != null) {
instLine.classList.add("highlight");
const r = rect(instLine), c = rect(instList);
if (Math.max(c.top-r.bottom, r.top-c.bottom)>=-30) instList.scrollTop = instLine.offsetTop-instList.clientHeight/2+instLine.clientHeight/2;
}
@@ -366,7 +377,7 @@ async function renderProfiler(path, opts) {
displaySelection("#profiler");
// support non realtime x axis units
formatTime = opts.unit === "ms" ? formatMicroseconds : formatCycles;
if (data?.path !== path) { data = {tracks:new Map(), axes:{}, path, first:null}; focusedDevice = null; focusedShape = null; }
if (data?.path !== path) { data = {tracks:new Map(), axes:{}, path, first:null, links:new Map()}; focusedDevice = null; focusedShape = null; }
setFocus(focusedShape);
// layout once!
if (data.tracks.size !== 0) return updateProgress(Status.COMPLETE);
@@ -455,10 +466,11 @@ async function renderProfiler(path, opts) {
}
// tiny device events go straight to the rewrite rule
const key = k.startsWith("TINY") ? null : `${k}-${j}`;
let info = e.info != null ? "\n"+e.info : "", trace = null, pc = null
let info = e.info != null ? "\n"+e.info : "", trace = null, pc = null, link = null
if (info.startsWith("\nPC:")) { pc = parseInt(e.info.split(":")[1]); info = ""; }
if (info.startsWith("\nTB:")) { trace = info; info = ""; }
const arg = { tooltipText:" N:"+shapes.length+"\n"+formatTime(e.dur)+info, label, pc, trace, bufs:[], key, ctx:shapeRef?.ctx, step:shapeRef?.step };
if (info.startsWith("\nLINK:")) { link = info.replace("\nLINK:", ""); info = ""; data.links.set(link, key); }
const arg = { tooltipText:" N:"+shapes.length+"\n"+formatTime(e.dur)+info, label, pc, trace, link, bufs:[], key, ctx:shapeRef?.ctx, step:shapeRef?.step };
if (e.key != null) shapeMap.set(e.key, key);
// offset y by depth
shapes.push({x:e.st, y:levelHeight*depth, width:e.dur, height:levelHeight, arg, label:opts.hideLabels ? null : label, fillColor });
@@ -625,7 +637,9 @@ async function renderProfiler(path, opts) {
// add label
drawText(ctx, e.label, x+2, y+e.height/2, width);
}
if (focusedShape != null && e.arg?.key === focusedShape) { ctx.strokeStyle = pcolor; ctx.stroke(); }
if ((focusedShape != null && e.arg?.key === focusedShape) || (data.link != null && (e.arg?.key === data.link[0] || e.arg?.key === data.link[1]))) {
ctx.strokeStyle = pcolor; ctx.stroke();
}
}
// draw row line
if (rowBorderColor != null) {
@@ -633,6 +647,15 @@ async function renderProfiler(path, opts) {
drawLine(ctx, [0, canvasWidth], [y, y], { color:rowBorderColor });
}
}
// draw the link
if (data.link != null) {
const [a, b] = [canvasRect(data.link[0], xscale), canvasRect(data.link[1], xscale)];
const [left, right] = a.x0 <= b.x0 ? [a, b] : [b, a];
const startX = left.x1, endX = right.x0;
const leftY = (left.y0+left.y1)/2, rightY = (right.y0+right.y1)/2;
const dx = endX-startX, bend = Math.max(12, Math.min(40, dx/2));
ctx.beginPath(); ctx.moveTo(startX, leftY); ctx.bezierCurveTo(startX+bend, leftY, endX-bend, rightY, endX, rightY); ctx.strokeStyle = "#858b9d"; ctx.stroke();
}
// draw axes
ctx.translate(0, baseOffset);
const y = secondaryTick != null ? tickSize+padding : 0;
tinygrad/viz/js/worker.js
+1 -1
View File
@@ -74,7 +74,7 @@ const layoutUOp = (g, { graph, change }, opts) => {
if (!opts.showIndexing) {
for (const n of g.nodes()) {
const node = g.node(n);
if (node.label.includes("dtypes.index")) g.removeNode(n);
if (node.label.includes("dtypes.weakint")) g.removeNode(n);
}
}
if (!opts.showCallSrc || opts.callSrcMask.size > 0) {
tinygrad/viz/serve.py
+17 -4
View File
@@ -44,7 +44,7 @@ from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphE
from tinygrad.dtype import dtypes
uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0", Ops.VCONST: "#e0e0e0", Ops.REDUCE: "#FF5B5B",
**{x:"#f2cb91" for x in {Ops.DEFINE_LOCAL, Ops.DEFINE_REG}}, Ops.REDUCE_AXIS: "#FF6B6B",
**{x:"#f2cb91" for x in {Ops.DEFINE_LOCAL, Ops.DEFINE_REG}}, Ops.REDUCE_AXIS: "#FF6B6B", Ops.SHAPED_WMMA: "#FF5B5B",
Ops.RANGE: "#c8a0e0", Ops.BARRIER: "#ff8080", Ops.IF: "#c8b0c0", Ops.SPECIAL: "#c0c0ff",
Ops.INDEX: "#cef263", Ops.WMMA: "#efefc0", Ops.MULTI: "#f6ccff", Ops.INS: "#eec4ff",
**{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80",
@@ -55,7 +55,6 @@ uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0",
# VIZ API
# A step is a lightweight descriptor for a trace entry
# Includes a name, metadata and a URL path for fetching the full data
@@ -104,7 +103,7 @@ def uop_to_json(x:UOp) -> dict[int, dict]:
# always exclude DEVICE/CONST/UNIQUE
if u.op in {Ops.DEVICE, Ops.CONST, Ops.UNIQUE, Ops.LUNIQUE} and u is not x: excluded.add(u)
if u.op is Ops.CONST and len(u.src) and u.src[0].op in {Ops.UNIQUE, Ops.LUNIQUE}: excluded.remove(u)
if u.op is Ops.VCONST and u.dtype.scalar() == dtypes.index and u is not x: excluded.add(u)
if u.op is Ops.VCONST and u.dtype.scalar() == dtypes.weakint and u is not x: excluded.add(u)
if u.op is Ops.VECTORIZE and len(u.src) == 0: excluded.add(u)
# exclude RESHAPE/EXPAND that only serve to broadcast a CONST
if u.op in {Ops.RESHAPE, Ops.EXPAND} and len(u.src) >= 1 and u.src[0] in excluded and u is not x: excluded.add(u)
@@ -341,13 +340,17 @@ def load_amd_counters(ctxs:list[dict], profile:list[ProfileEvent]) -> None:
def sqtt_timeline(data:bytes, lib:bytes, target:str) -> list[ProfileEvent]:
from tinygrad.renderer.amd.sqtt import (map_insts, InstructionInfo, PacketType, INST, InstOp, VALUINST, IMMEDIATE, IMMEDIATE_MASK, VMEMEXEC,
ALUEXEC, INST_RDNA4, InstOpRDNA4, TS_DELTA_OR_MARK, TS_DELTA_OR_MARK_RDNA4, CDNA_INST, InstOpCDNA,
WAVEEND, CDNA_WAVEEND)
WAVEEND, CDNA_WAVEEND, WAVERDY)
ret:list[ProfileEvent] = []
row_ends:dict[str, Decimal] = {}
row_counts:dict[str, itertools.count] = {}
curr_barrier:dict[str, ProfileRangeEvent] = {}
exec_pending:dict[str, list[str]] = {}
NS_PER_TICK = 10 # 100MHz
prev_pair:tuple[int, int]|None = None # (shader, realtime)
is_cdna = target.startswith("gfx9")
dispatch_to_exec = {"WMMA":"VALU", "VALU":"VALU", "VALUINST":"VALU", "VINTERP":"VALU", "GLOBAL":"VMEM", "FLAT":"VMEM", "LDS":"LDS", "SALU":"SALU",
"SMEM":"SALU", "VMEM":"VMEM"}
def add(name:str, p:PacketType, op:str|None=None, wave:int|None=None, info:InstructionInfo|None=None) -> None:
row = f"WAVE:{wave}" if (wave:=getattr(p, "wave", wave)) is not None else f"{p.__class__.__name__}:0 {name}"
# barrier on this row extends to fill the time our wave was waiting
@@ -356,7 +359,12 @@ def sqtt_timeline(data:bytes, lib:bytes, target:str) -> list[ProfileEvent]:
# allow CDNA packets to overlap, NOT allowed on RDNA.
if (et:=row_ends.get(row)) is not None and e.st < et and not is_cdna: raise RuntimeError(f"packet {p} overlaps another packet in {row}.")
row_ends[row] = unwrap(e.en)
idx = next(row_counts.setdefault(row, itertools.count(0)))
if name == "BARRIER": curr_barrier[row] = e
# queue for exec linking
if isinstance(p, (VALUINST, INST, INST_RDNA4)) and (exec_type:=dispatch_to_exec.get(name.split("_")[0])) is not None:
exec_pending.setdefault(exec_type, []).append(f"{row}-{idx}")
if isinstance(p, (ALUEXEC, VMEMEXEC)) and "ALT" not in str(p.src): e.name = TracingKey(op or name, ret=f"LINK:{exec_pending[name].pop(0)}")
for p, info in map_insts(data, lib, target):
if len(ret) > getenv("MAX_SQTT_PKTS", 50_000): break
if isinstance(p, (TS_DELTA_OR_MARK, TS_DELTA_OR_MARK_RDNA4)) and p.is_marker:
@@ -372,6 +380,11 @@ def sqtt_timeline(data:bytes, lib:bytes, target:str) -> list[ProfileEvent]:
add(name, p, info=info)
if isinstance(p, (VALUINST, IMMEDIATE, WAVEEND, CDNA_WAVEEND)): add(p.__class__.__name__, p, info=info)
if isinstance(p, IMMEDIATE_MASK): add("IMMEDIATE", p, wave=unwrap(info).wave, info=info)
if isinstance(p, WAVERDY):
for wave in range(16):
if p.mask & (1 << wave):
row = f"WAVE:{wave}"
if row in curr_barrier: add("WAVERDY", p, wave=wave)
if isinstance(p, (VMEMEXEC, ALUEXEC)):
name = str(p.src).split('.')[1]
if name == "VALU_SALU":