mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-08-16 17:18:27 +00:00
Compare commits
39
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
baa2f16bff | ||
|
|
ecf1477c0e | ||
|
|
e4186add83 | ||
|
|
e0694fdb8e | ||
|
|
678f83e41b | ||
|
|
a11b686c71 | ||
|
|
a0cbbc35ad | ||
|
|
fe94453d52 | ||
|
|
f793cdeb87 | ||
|
|
1bcea19846 | ||
|
|
c1cc277fc3 | ||
|
|
2551a60d97 | ||
|
|
e7aa26ed29 | ||
|
|
cf8232ec6a | ||
|
|
658c566e22 | ||
|
|
a8a9ac0e95 | ||
|
|
250f05a776 | ||
|
|
da9425c1a7 | ||
|
|
ae51bdd06a | ||
|
|
80d99d52a5 | ||
|
|
375ee2c576 | ||
|
|
1dc500426e | ||
|
|
585bd95b50 | ||
|
|
6af29b913b | ||
|
|
baab7e334d | ||
|
|
51420d1f99 | ||
|
|
43bce1f39f | ||
|
|
9f9a8b0b5b | ||
|
|
6e6059dde0 | ||
|
|
20d98b19c3 | ||
|
|
bb5671a837 | ||
|
|
be05028419 | ||
|
|
615ec6acf0 | ||
|
|
c4732a18bd | ||
|
|
5986d656a2 | ||
|
|
fc2bd53700 | ||
|
|
89ec2b3a74 | ||
|
|
84fc34b274 | ||
|
|
28edea5d67 |
@@ -52,12 +52,12 @@ jobs:
|
||||
- name: reset process replay
|
||||
run: python3.11 test/external/process_replay/reset.py
|
||||
- name: Run Stable Diffusion
|
||||
run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=1000 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
|
||||
run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=800 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
|
||||
- name: Run Stable Diffusion without fp16
|
||||
run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=1000 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
|
||||
run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=900 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
|
||||
- name: Run Stable Diffusion v2
|
||||
# TODO: very slow step time
|
||||
run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=100000 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
|
||||
run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=10000 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
|
||||
# process replay can't capture this, the graph is too large
|
||||
# TODO: too slow
|
||||
# - name: Run SDXL
|
||||
@@ -101,7 +101,7 @@ jobs:
|
||||
- name: Run GPT2
|
||||
run: |
|
||||
BENCHMARK_LOG=gpt2_nojit JIT=0 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
|
||||
BENCHMARK_LOG=gpt2 JIT=1 ASSERT_MIN_STEP_TIME=16 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
|
||||
BENCHMARK_LOG=gpt2 JIT=1 ASSERT_MIN_STEP_TIME=13 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
|
||||
- name: Run GPT2 w HALF
|
||||
run: BENCHMARK_LOG=gpt2_half HALF=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
|
||||
- name: Run GPT2 w HALF/BEAM
|
||||
@@ -246,9 +246,9 @@ jobs:
|
||||
- name: Run GPT2
|
||||
run: |
|
||||
BENCHMARK_LOG=gpt2_nojit NV=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
|
||||
BENCHMARK_LOG=gpt2 NV=1 JIT=1 ASSERT_MIN_STEP_TIME=10 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
|
||||
BENCHMARK_LOG=gpt2 NV=1 JIT=1 ASSERT_MIN_STEP_TIME=4 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
|
||||
- name: Run GPT2 w HALF
|
||||
run: BENCHMARK_LOG=gpt2_half NV=1 HALF=1 ASSERT_MIN_STEP_TIME=10 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
|
||||
run: BENCHMARK_LOG=gpt2_half NV=1 HALF=1 ASSERT_MIN_STEP_TIME=6 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
|
||||
- name: Run GPT2 w HALF/BEAM
|
||||
run: BENCHMARK_LOG=gpt2_half_beam NV=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
|
||||
- uses: actions/upload-artifact@v4
|
||||
@@ -316,11 +316,11 @@ jobs:
|
||||
- name: Train MNIST
|
||||
run: time PYTHONPATH=. NV=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
|
||||
- name: Run 10 CIFAR training steps
|
||||
run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=850 NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
|
||||
run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=270 NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
|
||||
- name: Run 10 CIFAR training steps w HALF
|
||||
run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=680 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
|
||||
run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=310 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
|
||||
- name: Run 10 CIFAR training steps w BF16
|
||||
run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=750 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
|
||||
run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=310 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
|
||||
# TODO: too slow
|
||||
# - name: Run 10 CIFAR training steps w winograd
|
||||
# run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=350 NV=1 CAPTURE_PROCESS_REPLAY=0 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
|
||||
@@ -426,7 +426,7 @@ jobs:
|
||||
- name: Test AM warm start time
|
||||
run: time AMD=1 python3 test/test_tiny.py TestTiny.test_plus
|
||||
- name: Run Stable Diffusion
|
||||
run: BENCHMARK_LOG=stable_diffusion ASSERT_MIN_STEP_TIME=900 AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
|
||||
run: BENCHMARK_LOG=stable_diffusion ASSERT_MIN_STEP_TIME=550 AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
|
||||
# TODO: too slow
|
||||
# - name: Run SDXL
|
||||
# run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=3200 CAPTURE_PROCESS_REPLAY=0 AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
|
||||
@@ -520,9 +520,9 @@ jobs:
|
||||
- name: Train MNIST
|
||||
run: time PYTHONPATH=. AMD=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
|
||||
- name: Run 10 CIFAR training steps
|
||||
run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=400 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
|
||||
run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=330 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
|
||||
- name: Run 10 CIFAR training steps w HALF
|
||||
run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=500 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
|
||||
run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=330 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
|
||||
# - name: Run 10 CIFAR training steps w BF16
|
||||
# run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
|
||||
# TODO: too slow
|
||||
@@ -625,15 +625,15 @@ jobs:
|
||||
- name: benchmark openpilot 0.9.9 dmonitoring
|
||||
run: BENCHMARK_LOG=openpilot_0_9_9_dmonitoring PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
|
||||
- name: openpilot compile3 0.9.9 driving_vision
|
||||
run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
|
||||
run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=22 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
|
||||
- name: openpilot compile3 0.9.9 driving_policy
|
||||
run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
|
||||
run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=7 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
|
||||
- name: openpilot compile3 0.9.9 dmonitoring
|
||||
run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
|
||||
run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=15 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
|
||||
- name: openpilot compile3 Space Lab policy + vision
|
||||
run: |
|
||||
PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/22aec22a10ce09384d4a4af2a0bbff08d54af7e0c888503508f356fae4ff0e29
|
||||
PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/c824f68646a3b94f117f01c70dc8316fb466e05fbd42ccdba440b8a8dc86914b
|
||||
PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/22aec22a10ce09384d4a4af2a0bbff08d54af7e0c888503508f356fae4ff0e29
|
||||
PYTHONPATH="." ASSERT_MIN_STEP_TIME=26 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/c824f68646a3b94f117f01c70dc8316fb466e05fbd42ccdba440b8a8dc86914b
|
||||
- name: benchmark MobileNetV2 on DSP
|
||||
run: |
|
||||
# generate quantized weights
|
||||
|
||||
@@ -377,7 +377,7 @@ jobs:
|
||||
llvm: 'true'
|
||||
- name: Test openpilot model kernel count and gate usage
|
||||
run: |
|
||||
ALLOWED_KERNEL_COUNT=190 ALLOWED_READ_IMAGE=2041 ALLOWED_GATED_READ_IMAGE=33 FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
|
||||
ALLOWED_KERNEL_COUNT=190 ALLOWED_READ_IMAGE=2041 ALLOWED_GATED_READ_IMAGE=543 FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
|
||||
- name: Test openpilot alt model correctness (float32)
|
||||
run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/3799fe46b3a629e491d4b8498b8ae83e4c88c304/selfdrive/modeld/models/supercombo.onnx
|
||||
- name: Test openpilot fastvits model correctness (float32)
|
||||
@@ -451,8 +451,7 @@ jobs:
|
||||
- name: Test Bert training
|
||||
run: NULL=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
|
||||
- name: Test llama 3 training
|
||||
# TODO: remove LLAMA_LAYERS once it's fast
|
||||
run: NULL=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=8 LLAMA_LAYERS=4 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
|
||||
run: NULL=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=8 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
|
||||
- name: Run process replay tests
|
||||
uses: ./.github/actions/process-replay
|
||||
|
||||
|
||||
@@ -80,7 +80,6 @@ print("******** third, the UOp ***********")
|
||||
|
||||
from tinygrad.engine.realize import run_schedule
|
||||
from tinygrad.engine.schedule import create_schedule_with_vars
|
||||
from tinygrad.helpers import RANGEIFY
|
||||
from tinygrad.schedule.rangeify import get_rangeify_map
|
||||
|
||||
# allocate some values + load in values
|
||||
|
||||
@@ -3,7 +3,7 @@ from pathlib import Path
|
||||
import multiprocessing
|
||||
|
||||
from tinygrad import Device, GlobalCounters, Tensor, TinyJit, dtypes
|
||||
from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, FUSE_CONV_BW, Profiling
|
||||
from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, Profiling
|
||||
from tinygrad.nn.state import get_parameters, get_state_dict, load_state_dict, safe_load, safe_save
|
||||
from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, Adam, AdamW
|
||||
|
||||
@@ -707,7 +707,7 @@ def train_unet3d():
|
||||
```BASEDIR=<folder_path> ./examples/mlperf/scripts/setup_kits19_dataset.sh```
|
||||
|
||||
2) To start training the model, run the following:
|
||||
```time PYTHONPATH=. WANDB=1 TRAIN_BEAM=3 FUSE_CONV_BW=1 GPUS=6 BS=6 MODEL=unet3d python3 examples/mlperf/model_train.py```
|
||||
```time PYTHONPATH=. WANDB=1 TRAIN_BEAM=3 GPUS=6 BS=6 MODEL=unet3d python3 examples/mlperf/model_train.py```
|
||||
"""
|
||||
from examples.mlperf.losses import dice_ce_loss
|
||||
from examples.mlperf.metrics import dice_score
|
||||
@@ -749,7 +749,6 @@ def train_unet3d():
|
||||
"train_beam": TRAIN_BEAM,
|
||||
"eval_beam": EVAL_BEAM,
|
||||
"wino": WINO.value,
|
||||
"fuse_conv_bw": FUSE_CONV_BW.value,
|
||||
"gpus": GPUS,
|
||||
"default_float": dtypes.default_float.name
|
||||
}
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
import os, sys, pickle, time
|
||||
import os, sys, pickle, time, re
|
||||
import numpy as np
|
||||
if "FLOAT16" not in os.environ: os.environ["FLOAT16"] = "1"
|
||||
if "IMAGE" not in os.environ: os.environ["IMAGE"] = "2"
|
||||
@@ -52,6 +52,8 @@ def compile(onnx_file):
|
||||
kernel_count += 1
|
||||
read_image_count += ei.prg.p.src.count("read_image")
|
||||
gated_read_image_count += ei.prg.p.src.count("?read_image")
|
||||
for v in [m.group(1) for m in re.finditer(r'(val\d+)\s*=\s*read_imagef\(', ei.prg.p.src)]:
|
||||
if len(re.findall(fr'[\?\:]{v}\.[xyzw]', ei.prg.p.src)) > 0: gated_read_image_count += 1
|
||||
print(f"{kernel_count=}, {read_image_count=}, {gated_read_image_count=}")
|
||||
if (allowed_kernel_count:=getenv("ALLOWED_KERNEL_COUNT", -1)) != -1:
|
||||
assert kernel_count == allowed_kernel_count, f"different kernels! {kernel_count=}, {allowed_kernel_count=}"
|
||||
@@ -77,13 +79,20 @@ def test_vs_compile(run, new_inputs, test_val=None):
|
||||
**{k:Tensor(v, device="NPY").realize() for k,v in new_inputs_numpy.items() if 'img' not in k}}
|
||||
|
||||
# run 20 times
|
||||
step_times = []
|
||||
for _ in range(20):
|
||||
st = time.perf_counter()
|
||||
out = run(**inputs)
|
||||
mt = time.perf_counter()
|
||||
val = out.numpy()
|
||||
et = time.perf_counter()
|
||||
print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {(et-st)*1e3:6.2f} ms")
|
||||
step_times.append((et-st)*1e3)
|
||||
print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {step_times[-1]:6.2f} ms")
|
||||
|
||||
if (assert_time:=getenv("ASSERT_MIN_STEP_TIME")):
|
||||
min_time = min(step_times)
|
||||
assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"
|
||||
|
||||
print(out, val.shape, val.dtype)
|
||||
if test_val is not None: np.testing.assert_equal(test_val, val)
|
||||
print("**** test done ****")
|
||||
|
||||
@@ -269,12 +269,14 @@ if __name__ == "__main__":
|
||||
|
||||
# load in weights
|
||||
with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
|
||||
load_state_dict(model, torch_load(fetch('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', 'sd-v1-4.ckpt'))['state_dict'], strict=False)
|
||||
load_state_dict(model, torch_load(fetch('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', 'sd-v1-4.ckpt'))['state_dict'], verbose=False, strict=False, realize=False)
|
||||
|
||||
if args.fp16:
|
||||
for k,v in get_state_dict(model).items():
|
||||
if k.startswith("model"):
|
||||
v.replace(v.cast(dtypes.float16).realize())
|
||||
v.replace(v.cast(dtypes.float16))
|
||||
|
||||
Tensor.realize(*get_state_dict(model).values())
|
||||
|
||||
# run through CLIP to get context
|
||||
tokenizer = Tokenizer.ClipTokenizer()
|
||||
|
||||
@@ -49,8 +49,7 @@ def rangeify_kernel3():
|
||||
b = Tensor.empty(N,N)
|
||||
c = a@b
|
||||
#c = c.reshape((32,2,16,4,32,2,16,4)).contiguous()
|
||||
with Context(RANGEIFY=1):
|
||||
sink = c.schedule()[-1].ast
|
||||
sink = c.schedule()[-1].ast
|
||||
#print(sink)
|
||||
|
||||
opts = [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.LOCAL, 0, 16), Opt(OptOps.UPCAST, 0, 2)]
|
||||
@@ -329,7 +328,7 @@ if __name__ == "__main__":
|
||||
elif HL == 1: hprg = hl_spec_kernel3()
|
||||
else: hprg = hand_spec_kernel3()
|
||||
if HL == 3:
|
||||
with Context(RANGEIFY=1, BLOCK_REORDER=0):
|
||||
with Context(BLOCK_REORDER=0):
|
||||
prg = get_program(hprg, Device.default.renderer)
|
||||
else:
|
||||
prg = get_program(hprg, Device.default.renderer)
|
||||
|
||||
@@ -155,6 +155,7 @@ class RGP:
|
||||
device_event = device_events[device]
|
||||
sqtt_events = [x for x in profile if isinstance(x, ProfileSQTTEvent) and x.device == device_event.device]
|
||||
if len(sqtt_events) == 0: raise RuntimeError(f"Device {device_event.device} doesn't contain SQTT data")
|
||||
device_props = sqtt_events[0].props
|
||||
sqtt_itrace_enabled = any([event.itrace for event in sqtt_events])
|
||||
sqtt_itrace_masked = not all_same([event.itrace for event in sqtt_events])
|
||||
sqtt_itrace_se_mask = functools.reduce(lambda a,b: a|b, [int(event.itrace) << event.se for event in sqtt_events], 0) if sqtt_itrace_masked else 0
|
||||
@@ -192,14 +193,14 @@ class RGP:
|
||||
flags=0,
|
||||
trace_shader_core_clock=0x93f05080,
|
||||
trace_memory_clock=0x4a723a40,
|
||||
device_id=0x744c,
|
||||
device_id={110000: 0x744c, 110003: 0x7480}[device_props['gfx_target_version']],
|
||||
device_revision_id=0xc8,
|
||||
vgprs_per_simd=1536,
|
||||
sgprs_per_simd=128*16,
|
||||
shader_engines=6,
|
||||
compute_unit_per_shader_engine=16,
|
||||
simd_per_compute_unit=2,
|
||||
wavefronts_per_simd=16,
|
||||
shader_engines=device_props['array_count'] // device_props['simd_arrays_per_engine'],
|
||||
compute_unit_per_shader_engine=device_props['simd_count'] // device_props['simd_per_cu'] // (device_props['array_count'] // device_props['simd_arrays_per_engine']),
|
||||
simd_per_compute_unit=device_props['simd_per_cu'],
|
||||
wavefronts_per_simd=device_props['max_waves_per_simd'],
|
||||
minimum_vgpr_alloc=4,
|
||||
vgpr_alloc_granularity=8,
|
||||
minimum_sgpr_alloc=128,
|
||||
@@ -218,7 +219,7 @@ class RGP:
|
||||
vram_bus_width=384, # 384-bit
|
||||
l2_cache_size=6 * 1024 * 1024, # 6 MB
|
||||
l1_cache_size=32 * 1024, # 32 KB per SIMD (?)
|
||||
lds_size=65536, # 64 KB per CU
|
||||
lds_size=device_props['lds_size_in_kb'] * 1024,
|
||||
gpu_name=b'NAVI31',
|
||||
alu_per_clock=0,
|
||||
texture_per_clock=0,
|
||||
|
||||
@@ -39,10 +39,6 @@ if __name__ == "__main__":
|
||||
step_times.append(t:=(time.perf_counter_ns() - st)*1e-6)
|
||||
print(f"jitted: {t:7.4f} ms")
|
||||
|
||||
if (assert_time:=getenv("ASSERT_MIN_STEP_TIME")):
|
||||
min_time = min(step_times)
|
||||
assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"
|
||||
|
||||
suffix = ""
|
||||
if IMAGE.value < 2: suffix += f"_image{IMAGE.value}" # image=2 has no suffix for compatibility
|
||||
if getenv("FLOAT16") == 1: suffix += "_float16"
|
||||
|
||||
Vendored
+18
-20
@@ -4,7 +4,7 @@ import numpy as np
|
||||
import torch
|
||||
|
||||
from tinygrad import GlobalCounters, Tensor, Device
|
||||
from tinygrad.helpers import getenv, Context, RANGEIFY
|
||||
from tinygrad.helpers import getenv
|
||||
from tinygrad.nn.state import get_parameters
|
||||
from tinygrad.engine.realize import capturing
|
||||
from tinygrad.tensor import _to_np_dtype
|
||||
@@ -164,7 +164,7 @@ class TestOpt(unittest.TestCase):
|
||||
|
||||
def test_permute_was_pushed(self):
|
||||
a = Tensor.randn(16, 16, 16)
|
||||
with CLCache(1 if RANGEIFY else 2):
|
||||
with CLCache(1):
|
||||
c = a.sum(2)
|
||||
d = c.permute(1,0).contiguous()
|
||||
d.realize()
|
||||
@@ -172,7 +172,7 @@ class TestOpt(unittest.TestCase):
|
||||
|
||||
def test_permute_was_pushed_through_contract_reshape(self):
|
||||
a = Tensor.randn(4, 4, 4, 4, 4)
|
||||
with CLCache(1 if RANGEIFY else 2):
|
||||
with CLCache(1):
|
||||
c = a.sum(-1)
|
||||
d = c.reshape(16,16).permute(1,0).contiguous()
|
||||
d.realize()
|
||||
@@ -180,7 +180,7 @@ class TestOpt(unittest.TestCase):
|
||||
|
||||
def test_permute_was_pushed_through_contractw1s_reshape(self):
|
||||
a = Tensor.randn(4, 4, 4, 4, 4)
|
||||
with CLCache(1 if RANGEIFY else 2):
|
||||
with CLCache(1):
|
||||
c = a.sum(-1)
|
||||
d = c.reshape(16,1,16).permute(2,1,0).contiguous()
|
||||
d.realize()
|
||||
@@ -188,7 +188,7 @@ class TestOpt(unittest.TestCase):
|
||||
|
||||
def test_permute_was_pushed_through_expand_reshape(self):
|
||||
a = Tensor.randn(16, 16, 16)
|
||||
with CLCache(1 if RANGEIFY else 2):
|
||||
with CLCache(1):
|
||||
c = a.sum(2)
|
||||
d = c.reshape(4,4,4,4).permute(2,3,0,1).contiguous()
|
||||
d.realize()
|
||||
@@ -217,24 +217,22 @@ class TestOpt(unittest.TestCase):
|
||||
assert cache_len == 1, "reduceop was rerun!"
|
||||
|
||||
def test_expand_reduce_is_folded_on_same_axis(self):
|
||||
with Context(FUSE_CONV_BW=1):
|
||||
for axis in [0, 1]:
|
||||
for n in [4, 8, 16]:
|
||||
b = torch.ones(n, n).sum(axis).reshape(n, 1).expand(n, n).sum(axis)
|
||||
with CLCache(allowed=3 if RANGEIFY else 2):
|
||||
a = Tensor.ones(n, n).contiguous().sum(axis).reshape(n, 1).expand(n, n).sum(axis)
|
||||
a.realize()
|
||||
np.testing.assert_allclose(a.numpy(), b.numpy(), rtol=1e-3, atol=1e-5)
|
||||
|
||||
def test_expand_reduce_is_folded_on_different_axes(self):
|
||||
with Context(FUSE_CONV_BW=1):
|
||||
axis1, axis2 = 0, 1
|
||||
for axis in [0, 1]:
|
||||
for n in [4, 8, 16]:
|
||||
b = torch.ones(n, n).sum(axis1).reshape(n, 1).expand(n, n).sum(axis2)
|
||||
with CLCache(allowed=3 if RANGEIFY else 2):
|
||||
a = Tensor.ones(n, n).contiguous().sum(axis1).reshape(n, 1).expand(n, n).sum(axis2)
|
||||
b = torch.ones(n, n).sum(axis).reshape(n, 1).expand(n, n).sum(axis)
|
||||
with CLCache(allowed=3):
|
||||
a = Tensor.ones(n, n).contiguous().sum(axis).reshape(n, 1).expand(n, n).sum(axis)
|
||||
a.realize()
|
||||
np.testing.assert_allclose(a.numpy(), b.numpy(), rtol=1e-3, atol=1e-5)
|
||||
|
||||
def test_expand_reduce_is_folded_on_different_axes(self):
|
||||
axis1, axis2 = 0, 1
|
||||
for n in [4, 8, 16]:
|
||||
b = torch.ones(n, n).sum(axis1).reshape(n, 1).expand(n, n).sum(axis2)
|
||||
with CLCache(allowed=3):
|
||||
a = Tensor.ones(n, n).contiguous().sum(axis1).reshape(n, 1).expand(n, n).sum(axis2)
|
||||
a.realize()
|
||||
np.testing.assert_allclose(a.numpy(), b.numpy(), rtol=1e-3, atol=1e-5)
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
|
||||
Vendored
+13
-4
@@ -1,7 +1,8 @@
|
||||
import gc
|
||||
from tinygrad import Tensor, UOp, Device
|
||||
from tinygrad import Tensor, UOp, Device, nn
|
||||
from tinygrad.shape.shapetracker import views_to_valid_uop
|
||||
from tinygrad.engine.realize import method_cache, get_program
|
||||
from test.test_tiny import TestTiny
|
||||
|
||||
def uops_allocated(): return sum([isinstance(x, UOp) for x in gc.get_objects()])
|
||||
def print_uops():
|
||||
@@ -46,9 +47,16 @@ def realized_gradient():
|
||||
z = y.matmul(x).sum()
|
||||
z.backward()
|
||||
Tensor.realize(x, y, z, x.grad, y.grad)
|
||||
def nn_batchnorm(): nn.BatchNorm(64)
|
||||
def nn_conv2d(): nn.Conv2d(64, 64, 3)
|
||||
def plus(): TestTiny().test_plus()
|
||||
def mnist(): TestTiny().test_mnist()
|
||||
def mnist_backward(): TestTiny().test_mnist_backward()
|
||||
|
||||
tests = [start, single_tensor, two_plus_two, two_plus_two_schedule, two_plus_two_kernel,
|
||||
two_plus_two_linearize, two_plus_two_realize, two_plus_two_item, gradient_test,
|
||||
realized_eye, realized_list, kernel_matmul, realized_matmul, realized_gradient]
|
||||
realized_eye, realized_list, kernel_matmul, realized_matmul, realized_gradient,
|
||||
nn_batchnorm, nn_conv2d, plus, mnist, mnist_backward]
|
||||
|
||||
if __name__ == "__main__":
|
||||
gc.disable()
|
||||
@@ -61,11 +69,12 @@ if __name__ == "__main__":
|
||||
# these caches will keep uops alive
|
||||
method_cache.clear()
|
||||
views_to_valid_uop.cache_clear()
|
||||
Tensor._device_seeds.clear()
|
||||
Tensor._device_rng_counters.clear()
|
||||
|
||||
new_uops = uops_allocated()
|
||||
print_uops()
|
||||
gc.collect()
|
||||
new_uops_gc = uops_allocated()
|
||||
print(f"{t.__name__:30s}: {new_uops:3d} -> {new_uops_gc:3d}")
|
||||
if new_uops != start_uops: print_uops()
|
||||
assert new_uops == start_uops
|
||||
#print_uops()
|
||||
|
||||
+2
-5
@@ -1,4 +1,4 @@
|
||||
import time, struct, unittest
|
||||
import time, struct
|
||||
from typing import Any, Callable
|
||||
import numpy as np
|
||||
from tinygrad import Tensor, dtypes, Device
|
||||
@@ -7,7 +7,7 @@ from tinygrad.tensor import _to_np_dtype
|
||||
from tinygrad.engine.realize import Runner
|
||||
from tinygrad.dtype import DType
|
||||
from tinygrad.nn.state import get_parameters
|
||||
from tinygrad.helpers import T, CI, RANGEIFY
|
||||
from tinygrad.helpers import T, CI
|
||||
from tinygrad.codegen import full_rewrite
|
||||
from tinygrad.runtime.ops_python import PythonProgram, PythonRenderer, PythonCompiler
|
||||
|
||||
@@ -62,6 +62,3 @@ def not_support_multi_device():
|
||||
|
||||
# NOTE: This will open REMOTE if it's the default device
|
||||
REAL_DEV = (Device.DEFAULT if Device.DEFAULT != "REMOTE" else Device['REMOTE'].properties.real_device)
|
||||
|
||||
def expect_rangeify_fails(fxn): return (unittest.expectedFailure if RANGEIFY else (lambda f:f))(fxn)
|
||||
def expect_nonrangeify_fails(fxn): return (unittest.expectedFailure if not RANGEIFY else (lambda f:f))(fxn)
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
import ctypes, ctypes.util, time
|
||||
import ctypes, time
|
||||
import tinygrad.runtime.autogen.nv_gpu as nv_gpu
|
||||
from enum import Enum, auto
|
||||
from test.mockgpu.gpu import VirtGPU
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
import unittest
|
||||
from tinygrad import Device, Tensor, dtypes
|
||||
from tinygrad.helpers import CI, RANGEIFY
|
||||
from tinygrad.helpers import CI
|
||||
from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
|
||||
|
||||
# TODO: write a clean version of this
|
||||
@@ -351,7 +351,6 @@ class TestKernelOpts(unittest.TestCase):
|
||||
] + [[Opt(OptOps.THREAD, 0, 4)] if Device[Device.DEFAULT].renderer.global_max[0] >= 4 else []]
|
||||
+ [[Opt(OptOps.THREAD, 0, 8)] if Device[Device.DEFAULT].renderer.global_max[0] >= 8 else []])
|
||||
|
||||
@unittest.skipUnless(RANGEIFY>=1, "Kernel only fuses with rangeify")
|
||||
def test_double_sum_group(self):
|
||||
a = Tensor.rand(4, 4, 4)
|
||||
r = a.sum((1, 2)).sum()
|
||||
|
||||
+2
-2
@@ -1,7 +1,7 @@
|
||||
import unittest
|
||||
import numpy as np
|
||||
from tinygrad import Tensor, GlobalCounters, dtypes, nn, Device, Variable
|
||||
from tinygrad.helpers import CI, Context, getenv, RANGEIFY
|
||||
from tinygrad.helpers import CI, Context, getenv
|
||||
from tinygrad.engine.realize import run_schedule
|
||||
from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
|
||||
from tinygrad.uop.ops import Ops
|
||||
@@ -95,7 +95,7 @@ class TestIndexing(unittest.TestCase):
|
||||
X = dataset[idxs]
|
||||
assert X.shape == (4,DDIM)
|
||||
sched = X.schedule()
|
||||
self.assertEqual(len(sched), 1 if RANGEIFY else 2)
|
||||
self.assertEqual(len(sched), 1)
|
||||
run_schedule(sched)
|
||||
assert GlobalCounters.global_ops < 4*DSET, f"too many ops {GlobalCounters.global_ops} != {4*DSET}"
|
||||
np.testing.assert_allclose(real_index, X.numpy())
|
||||
|
||||
+26
-36
@@ -1,10 +1,9 @@
|
||||
#!/usr/bin/env python
|
||||
import unittest
|
||||
import contextlib
|
||||
import numpy as np
|
||||
from tinygrad import dtypes, Tensor, TinyJit, GlobalCounters, Variable
|
||||
from tinygrad.device import is_dtype_supported
|
||||
from tinygrad.helpers import temp, RANGEIFY
|
||||
from tinygrad.helpers import temp
|
||||
|
||||
N = 200 # has to be bigger than the cache to fail
|
||||
|
||||
@@ -271,8 +270,6 @@ class TestAssign(unittest.TestCase):
|
||||
b.assign(a.contiguous()).realize()
|
||||
assert GlobalCounters.kernel_count - kc == 2
|
||||
|
||||
# passing in RANGEIFY=1, RANGEIFY=0 asserts permuted assigns it can't fuse
|
||||
def assert_permuted_assign(self): return self.assertRaisesRegex(RuntimeError, "contiguous") if not RANGEIFY else contextlib.nullcontext()
|
||||
def test_permuted_assignment(self):
|
||||
a = Tensor(np.arange(N*N, dtype=np.float32)).reshape(N,N)
|
||||
b = Tensor(np.arange(N*N, dtype=np.float32)).reshape(N,N)
|
||||
@@ -280,14 +277,13 @@ class TestAssign(unittest.TestCase):
|
||||
b.realize()
|
||||
ba1 = a.uop.base.realized
|
||||
bb1 = b.uop.base.realized
|
||||
with self.assert_permuted_assign():
|
||||
a = a.permute(1,0)
|
||||
a += b
|
||||
a.realize()
|
||||
ba2 = a.uop.base.realized
|
||||
np.testing.assert_allclose(a.numpy(), np.arange(N*N).reshape((N,N)) + np.arange(N*N).reshape((N,N)).transpose(1,0))
|
||||
# permute and base are the same buffer
|
||||
assert ba1 == ba2 and ba1 != bb1
|
||||
a = a.permute(1,0)
|
||||
a += b
|
||||
a.realize()
|
||||
ba2 = a.uop.base.realized
|
||||
np.testing.assert_allclose(a.numpy(), np.arange(N*N).reshape((N,N)) + np.arange(N*N).reshape((N,N)).transpose(1,0))
|
||||
# permute and base are the same buffer
|
||||
assert ba1 == ba2 and ba1 != bb1
|
||||
|
||||
def test_post_permuted_assignment(self):
|
||||
a = Tensor(np.arange(N*N, dtype=np.float32)).reshape(N,N)
|
||||
@@ -297,15 +293,13 @@ class TestAssign(unittest.TestCase):
|
||||
#GlobalCounters.cache = []
|
||||
ba1 = a.uop.base.realized # noqa: F841
|
||||
bb1 = b.uop.base.realized # noqa: F841
|
||||
with self.assert_permuted_assign():
|
||||
a.assign(a.permute(1,0) + b) # this should not work!
|
||||
a.realize()
|
||||
ba2 = a.uop.base.realized # noqa: F841
|
||||
# NOTE: don't test that it's assigned
|
||||
#assert ba1 == ba2 and ba1 != bb1
|
||||
np.testing.assert_allclose(a.numpy(), np.arange(N*N).reshape((N,N)) + np.arange(N*N).reshape((N,N)).transpose(1,0))
|
||||
a.assign(a.permute(1,0) + b) # this should not work!
|
||||
a.realize()
|
||||
ba2 = a.uop.base.realized # noqa: F841
|
||||
# NOTE: don't test that it's assigned
|
||||
#assert ba1 == ba2 and ba1 != bb1
|
||||
np.testing.assert_allclose(a.numpy(), np.arange(N*N).reshape((N,N)) + np.arange(N*N).reshape((N,N)).transpose(1,0))
|
||||
|
||||
@unittest.skipUnless(RANGEIFY, "only correct in rangeify")
|
||||
def test_post_permuted_assignment_alt(self):
|
||||
a = Tensor.arange(N*N).reshape(N,N).contiguous().realize()
|
||||
b = Tensor.arange(N*N).reshape(N,N).contiguous().realize()
|
||||
@@ -345,21 +339,18 @@ class TestAssign(unittest.TestCase):
|
||||
def test_permuted_assignment_correct(self):
|
||||
a = Tensor.arange(4 * 4).reshape(4, 4).contiguous().realize()
|
||||
b = Tensor.arange(4 * 4).reshape(4, 4).contiguous().realize()
|
||||
# TODO: swizzler.py limitation, should NOT raise AssertionError from numpy.
|
||||
with self.assert_permuted_assign():
|
||||
a = a.permute(1, 0)
|
||||
new_val = a + b
|
||||
a.assign(new_val)
|
||||
np.testing.assert_equal(a.numpy(), np.arange(4 * 4).reshape(4, 4).transpose(1, 0) + np.arange(4 * 4).reshape(4, 4))
|
||||
a = a.permute(1, 0)
|
||||
new_val = a + b
|
||||
a.assign(new_val)
|
||||
np.testing.assert_equal(a.numpy(), np.arange(4 * 4).reshape(4, 4).transpose(1, 0) + np.arange(4 * 4).reshape(4, 4))
|
||||
|
||||
def test_permuted_reduceop_child_dual_use(self):
|
||||
a = Tensor.randn(32, 32, 32).realize()
|
||||
b = Tensor.full((32, 32), 1.).contiguous().realize()
|
||||
with self.assert_permuted_assign():
|
||||
r = a.sum(axis=1)
|
||||
b.assign(r + b.permute(1, 0))
|
||||
b.realize()
|
||||
np.testing.assert_allclose(b.numpy(), a.numpy().sum(axis=1)+np.ones((32, 32)).transpose(1, 0), atol=1e-6, rtol=1e-3)
|
||||
r = a.sum(axis=1)
|
||||
b.assign(r + b.permute(1, 0))
|
||||
b.realize()
|
||||
np.testing.assert_allclose(b.numpy(), a.numpy().sum(axis=1)+np.ones((32, 32)).transpose(1, 0), atol=1e-6, rtol=1e-3)
|
||||
|
||||
@unittest.skip("multi output not supported anymore")
|
||||
def test_permuted_reduceop_multioutput_dual_use(self):
|
||||
@@ -401,11 +392,10 @@ class TestAssign(unittest.TestCase):
|
||||
|
||||
def test_permuted_assignment_masked_view_not_contiguous(self):
|
||||
a = Tensor.ones(4, 4).contiguous().realize()
|
||||
with self.assert_permuted_assign():
|
||||
b = a.shrink((None, (0, 2))).pad((None, (0, 2)), value=2).permute(1, 0)
|
||||
a.assign(a + b)
|
||||
a.realize()
|
||||
self.assertListEqual(a.tolist(), [[2.,2.,2.,2.],[2.,2.,2.,2.],[3.,3.,3.,3.], [3.,3.,3.,3.]])
|
||||
b = a.shrink((None, (0, 2))).pad((None, (0, 2)), value=2).permute(1, 0)
|
||||
a.assign(a + b)
|
||||
a.realize()
|
||||
self.assertListEqual(a.tolist(), [[2.,2.,2.,2.],[2.,2.,2.,2.],[3.,3.,3.,3.], [3.,3.,3.,3.]])
|
||||
|
||||
# TODO: is there a way to sneak in a permute such that it returns the wrong answer?
|
||||
|
||||
|
||||
@@ -3,7 +3,6 @@ from tinygrad import Tensor, Device, dtypes
|
||||
from tinygrad.dtype import DType, ConstType
|
||||
from tinygrad.uop.ops import Ops, UOp
|
||||
from tinygrad.codegen import full_rewrite_to_sink
|
||||
from tinygrad.helpers import RANGEIFY
|
||||
from tinygrad.device import is_dtype_supported
|
||||
import numpy as np
|
||||
from test.helpers import not_support_multi_device
|
||||
@@ -69,9 +68,12 @@ class TestBinaryOpsConstFolding(unittest.TestCase):
|
||||
def test_tensor_one_mul(self):
|
||||
_check_ast_count(0, Tensor.ones(4) * Tensor([1.0, 2, 3, 4]))
|
||||
|
||||
# TODO: these will be fixed with better folding
|
||||
@unittest.expectedFailure
|
||||
def test_bool_tensor_mul_bool(self):
|
||||
_check_ast_count(0, Tensor([True, False]) * True)
|
||||
_check_ast_count(0, Tensor([True, False]) * False)
|
||||
@unittest.expectedFailure
|
||||
def test_bool_mul_bool_tensor(self):
|
||||
_check_ast_count(0, True * Tensor([True, False]))
|
||||
_check_ast_count(0, False * Tensor([True, False]))
|
||||
@@ -155,8 +157,7 @@ class TestMovedConstFolding(unittest.TestCase):
|
||||
_check_ast_count(0, Tensor([1.0, 2, 3, 4]) + Tensor.zeros(6).shrink(((1, 5),)))
|
||||
|
||||
def test_add_padded_zero(self):
|
||||
# TODO: it's 1 now, this might be possible to fold
|
||||
_check_ast_count(0 if RANGEIFY else 1, Tensor([1.0, 2, 3, 4]) + Tensor.zeros(2).pad(((1, 1),)))
|
||||
_check_ast_count(0, Tensor([1.0, 2, 3, 4]) + Tensor.zeros(2).pad(((1, 1),)))
|
||||
|
||||
def test_mul_shrunk_one(self):
|
||||
_check_ast_count(0, Tensor([1.0, 2, 3, 4]) * Tensor.ones(6).shrink(((1, 5),)))
|
||||
@@ -165,16 +166,16 @@ class TestMovedConstFolding(unittest.TestCase):
|
||||
_check_ast_count(1, Tensor([1.0, 2, 3, 4]) * Tensor.ones(2).pad(((1, 1),)))
|
||||
|
||||
def test_cast_padded(self):
|
||||
# NOTE: RANGEIFY or not, it's always 1 kernel when calling .numpy, limitation of _check_ast_count
|
||||
# NOTE: it's always 1 kernel when calling .numpy, limitation of _check_ast_count
|
||||
if is_dtype_supported(dtypes.int16):
|
||||
_check_ast_count(1 if RANGEIFY else 0, Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int16))
|
||||
_check_ast_count(1, Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int16))
|
||||
np.testing.assert_equal(Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int16).numpy(), [0, 1, 1, 1, 1, 0])
|
||||
if is_dtype_supported(dtypes.uint16):
|
||||
_check_ast_count(1 if RANGEIFY else 0, Tensor.full(4, fill_value=-1).pad(((1, 1),)).cast(dtypes.uint16))
|
||||
_check_ast_count(1, Tensor.full(4, fill_value=-1).pad(((1, 1),)).cast(dtypes.uint16))
|
||||
np.testing.assert_equal(Tensor.full(4, fill_value=-1).pad(((1, 1),)).cast(dtypes.uint16).numpy(), [0, 65535, 65535, 65535, 65535, 0])
|
||||
# folded
|
||||
if is_dtype_supported(dtypes.int64):
|
||||
_check_ast_count(1 if RANGEIFY else 0, Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int64))
|
||||
_check_ast_count(1, Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int64))
|
||||
np.testing.assert_equal(Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int64).numpy(), [0, 1, 1, 1, 1, 0])
|
||||
|
||||
class TestReduceOpsConstFolding(unittest.TestCase):
|
||||
@@ -246,7 +247,7 @@ class TestReduceOpsConstFolding(unittest.TestCase):
|
||||
t = Tensor.ones(16, dtype=dt).reshape(4, 4)
|
||||
assert t.sum().dtype == t.contiguous().sum().dtype
|
||||
|
||||
@unittest.skipIf(not_support_multi_device() or RANGEIFY, "no multi, RANGEIFY doesn't support multi const folding")
|
||||
@unittest.skipIf(not_support_multi_device() or True, "no multi, RANGEIFY doesn't support multi const folding")
|
||||
class TestMultiConstFolding(unittest.TestCase):
|
||||
def test_multi_const_folding_literal(self):
|
||||
ds = tuple(f"{Device.DEFAULT}:{i}" for i in range(4))
|
||||
|
||||
@@ -4,7 +4,7 @@ from tinygrad import Device, dtypes, Tensor, Context
|
||||
from tinygrad.device import LRUAllocator, is_dtype_supported
|
||||
from tinygrad.dtype import ImageDType
|
||||
from tinygrad.engine.realize import lower_schedule
|
||||
from tinygrad.helpers import prod, unwrap, RANGEIFY
|
||||
from tinygrad.helpers import prod, unwrap
|
||||
from test.helpers import REAL_DEV
|
||||
|
||||
IMAGE_SUPPORTED_DEVICES = ("QCOM", "CL")
|
||||
@@ -139,7 +139,7 @@ class TestImageDType(unittest.TestCase):
|
||||
# NOTE: the w1 grad must realize to a seperate kernel
|
||||
assert w1.grad.uop.is_realized, f"never realized {w1.grad}"
|
||||
self.assertEqual(w1.grad.uop.base.buffer.dtype, dtypes.float32)
|
||||
self.assertEqual(len(sched), 9 if RANGEIFY else 10)
|
||||
self.assertEqual(len(sched), 9)
|
||||
|
||||
@unittest.skipUnless(REAL_DEV in IMAGE_SUPPORTED_DEVICES, "Images not supported")
|
||||
class TestImageRealization(unittest.TestCase):
|
||||
|
||||
@@ -8,7 +8,7 @@ from tinygrad.uop.ops import UOp, Ops, GroupOp
|
||||
from tinygrad.device import Device, Buffer, is_dtype_supported
|
||||
from tinygrad.tensor import Tensor, _to_np_dtype
|
||||
from tinygrad.engine.realize import run_schedule, lower_schedule, CompiledRunner, get_program
|
||||
from tinygrad.helpers import Context, flatten, dedup, TC_SELECT, TC_OPT, RANGEIFY
|
||||
from tinygrad.helpers import Context, flatten, dedup, TC_SELECT, TC_OPT
|
||||
from tinygrad.dtype import DType, dtypes, PtrDType, AddrSpace
|
||||
from tinygrad.renderer.ptx import PTXRenderer
|
||||
|
||||
@@ -314,7 +314,7 @@ class TestLinearizer(unittest.TestCase):
|
||||
a.realize()
|
||||
np.testing.assert_equal(a.flatten().numpy(), [1.,1.,1.,1.,2.,2.,2.,2.,1.,1.,1.,1.,1.,1.,1.,1.])
|
||||
|
||||
@unittest.skipIf(RANGEIFY and isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX indexes differently. might be ok?")
|
||||
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX indexes differently. might be ok?")
|
||||
def test_where_fold(self):
|
||||
a = Tensor.ones(4, 4).contiguous().realize()
|
||||
b = a.shrink(((1, 2), None)).pad(((1, 2), None))
|
||||
@@ -432,7 +432,7 @@ def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]:
|
||||
def helper_linearizer_ast(ast:UOp, inputs:list[Tensor], *args, **kwargs):
|
||||
assert isinstance(ast, UOp), "ast must be UOp"
|
||||
inbufs = [x.uop.base.buffer for x in inputs]
|
||||
outbufs = [Buffer(inbufs[-1].device if inbufs else Device.DEFAULT, out.st_arg.size, out.src[1].dtype).allocate() for out in ast.src]
|
||||
outbufs = [Buffer(inbufs[-1].device if inbufs else Device.DEFAULT, out.size, out.src[1].dtype).allocate() for out in ast.src]
|
||||
_helper_linearizer_opt_ast(ast, outbufs+inbufs, *args, **kwargs)
|
||||
|
||||
def helper_linearizer_opt(r:Tensor|list[Tensor], *args, **kwargs):
|
||||
|
||||
+3
-4
@@ -2,7 +2,7 @@ import time, math, unittest, functools, platform, warnings
|
||||
import numpy as np
|
||||
from typing import List, Callable
|
||||
import torch
|
||||
from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, TRANSCENDENTAL, CPU_LLVM, AMD_LLVM, RANGEIFY
|
||||
from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, TRANSCENDENTAL, CPU_LLVM, AMD_LLVM
|
||||
from tinygrad import Tensor, Device, dtypes
|
||||
from tinygrad.tensor import _to_np_dtype
|
||||
from tinygrad.device import is_dtype_supported
|
||||
@@ -3040,7 +3040,6 @@ class TestOps(unittest.TestCase):
|
||||
pos_weight=torch.tensor(pos_weight)),
|
||||
lambda x,y: x.binary_crossentropy_logits(y.clip(0,1),pos_weight=Tensor(pos_weight)))
|
||||
|
||||
@unittest.skipIf(RANGEIFY > 1, "broken on RANGEIFY > 1, TODO: fix")
|
||||
def test_cross_entropy_class_probabilities(self):
|
||||
helper_test_op([(32,), (32,)], lambda x,y: torch.nn.functional.cross_entropy(x, y), lambda x,y: x.cross_entropy(y))
|
||||
helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.cross_entropy(x, y), lambda x,y: x.cross_entropy(y))
|
||||
@@ -3164,8 +3163,8 @@ class TestOps(unittest.TestCase):
|
||||
helper_test_op([(32,10)], lambda x: x.masked_fill((x>0.1).detach(), -math.inf))
|
||||
helper_test_op([(32,10)], lambda x: x.masked_fill((x<0.1).detach(), -math.inf))
|
||||
|
||||
@unittest.skipIf(RANGEIFY and (getenv("MOCKGPU") or Device.DEFAULT == "PYTHON"), "very slow on MOCKGPU because reduce does not fold")
|
||||
@unittest.skipIf(RANGEIFY and Device.DEFAULT == "WEBGPU", "webgpu runtime issue")
|
||||
@unittest.skipIf((getenv("MOCKGPU") or Device.DEFAULT == "PYTHON"), "very slow on MOCKGPU because reduce does not fold")
|
||||
@unittest.skipIf(Device.DEFAULT == "WEBGPU", "webgpu runtime issue")
|
||||
def test_masked_select(self):
|
||||
helper_test_op([(32, 10)], lambda x: x.masked_select(x>0.5), lambda x: x.masked_select(x>0.5), forward_only=True)
|
||||
helper_test_op([(32, 10)], lambda x: x.masked_select(torch.tensor(True)), lambda x: x.masked_select(Tensor(True)), forward_only=True)
|
||||
|
||||
@@ -1,9 +1,8 @@
|
||||
import unittest
|
||||
from tinygrad import Tensor, nn
|
||||
from tinygrad.helpers import RANGEIFY, Context, GlobalCounters
|
||||
from tinygrad.helpers import Context, GlobalCounters
|
||||
from tinygrad.uop.ops import UOp, graph_rewrite, PatternMatcher, UPat, Ops
|
||||
|
||||
@unittest.skipIf(RANGEIFY<1, "tests only for RANGEIFY")
|
||||
class TestRangeifyAssign(unittest.TestCase):
|
||||
def test_assign_permuted(self):
|
||||
A = Tensor.empty(4, 4, dtype='int')
|
||||
@@ -55,7 +54,6 @@ class TestRangeifyOpt(unittest.TestCase):
|
||||
A = Tensor.empty(8,8,8,8).permute(1,0,3,2).flatten()
|
||||
A.sum().realize()
|
||||
|
||||
@unittest.skipIf(RANGEIFY<1, "tests only for RANGEIFY")
|
||||
class TestRangeify(unittest.TestCase):
|
||||
def test_groupnorm(self):
|
||||
# ranges 1 and 3 are merging
|
||||
@@ -230,7 +228,6 @@ class TestRangeify(unittest.TestCase):
|
||||
# contiguous + reduce can support ranges?
|
||||
|
||||
@unittest.skip("okay to disable this for now")
|
||||
@unittest.skipIf(RANGEIFY<1, "tests only for RANGEIFY")
|
||||
class TestOuterworld(unittest.TestCase):
|
||||
def test_passthrough_range(self):
|
||||
t = Tensor.rand(10, 10).realize()
|
||||
|
||||
+41
-59
@@ -17,7 +17,6 @@ from tinygrad.helpers import CI, DEBUG, SPLIT_REDUCEOP, GlobalCounters, Context,
|
||||
from tinygrad.schedule.rangeify import get_rangeify_map, Kernel
|
||||
from tinygrad.engine.schedule import create_schedule_with_vars
|
||||
from tinygrad.engine.realize import CompiledRunner, run_schedule, lower_schedule
|
||||
from test.helpers import expect_rangeify_fails, expect_nonrangeify_fails
|
||||
|
||||
class KernelCountException(Exception): pass
|
||||
def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Tensor]|None=None, filter_sink=True):
|
||||
@@ -45,7 +44,7 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te
|
||||
def _realize_weights(m):
|
||||
for p in nn.state.get_parameters(m): p.realize()
|
||||
|
||||
def _test_conv2d(allowed:int, dtype:DType=dtypes.float, **kwargs):
|
||||
def _test_conv2d(allowed:int, dtype:DType=dtypes.float):
|
||||
old_default_float, dtypes.default_float = dtypes.default_float, dtype
|
||||
dtypes.default_float = dtype
|
||||
Tensor.manual_seed(0)
|
||||
@@ -54,7 +53,7 @@ def _test_conv2d(allowed:int, dtype:DType=dtypes.float, **kwargs):
|
||||
w = Tensor.uniform(16, CIN, 3, 3, requires_grad=True).realize()
|
||||
ret = Tensor.conv2d(img, w).relu().mean().backward()
|
||||
dtypes.default_float = old_default_float
|
||||
with Context(**kwargs): s = Tensor.schedule(ret, img.grad, w.grad)
|
||||
s = Tensor.schedule(ret, img.grad, w.grad)
|
||||
run_schedule(s.copy())
|
||||
cnt = len([si for si in s if si.ast.op is Ops.SINK])
|
||||
assert cnt == allowed, f"expected {allowed} kernels, got {cnt}"
|
||||
@@ -117,7 +116,7 @@ class TestSchedule(unittest.TestCase):
|
||||
a = Tensor.empty(10)
|
||||
b = Tensor.empty((1,), device="CPU").expand(10).contiguous()
|
||||
c = a+b
|
||||
with self.assertRaisesRegex(RuntimeError, "all buffers must be on the same device"): check_schedule(c, 2 if RANGEIFY else 1)
|
||||
with self.assertRaisesRegex(RuntimeError, "all buffers must be on the same device"): check_schedule(c, 2)
|
||||
|
||||
@unittest.skipUnless(is_dtype_supported(dtypes.half) and getenv("CAST_AFTER_EXPAND"), "need half and CAST_AFTER_EXPAND=1")
|
||||
@unittest.skip("CAST_AFTER_EXPAND is not supported")
|
||||
@@ -343,7 +342,7 @@ class TestSchedule(unittest.TestCase):
|
||||
r1 = (x - r0).sum(axis=0).div(2)
|
||||
out0 = r0 + y
|
||||
out1 = r1 + y
|
||||
schedule = check_schedule([out0, out1], 2 if RANGEIFY else 4)
|
||||
schedule = check_schedule([out0, out1], 2)
|
||||
reduceops = [x for si in schedule for x in si.ast.toposort() if x.op in {Ops.REDUCE_AXIS, Ops.REDUCE}]
|
||||
assert len(reduceops) in [2,3] # why is RANGEIFY different?
|
||||
|
||||
@@ -470,15 +469,14 @@ class TestSchedule(unittest.TestCase):
|
||||
check_schedule(opt.schedule_step(), cnt)
|
||||
|
||||
def test_fold_batchnorm_backward(self):
|
||||
with Context(FUSE_CONV_BW=1):
|
||||
with Tensor.train():
|
||||
x = Tensor.empty((2, 16, 8, 8)).contiguous()
|
||||
bn = nn.BatchNorm2d(16)
|
||||
bn.weight.requires_grad = bn.bias.requires_grad = x.requires_grad = True
|
||||
fw = bn(x).contiguous_backward().relu().contiguous()
|
||||
fw.sum().backward()
|
||||
# TODO: this is too many
|
||||
check_schedule([x.grad, bn.weight.grad, bn.bias.grad, fw], 10)
|
||||
with Tensor.train():
|
||||
x = Tensor.empty((2, 16, 8, 8)).contiguous()
|
||||
bn = nn.BatchNorm2d(16)
|
||||
bn.weight.requires_grad = bn.bias.requires_grad = x.requires_grad = True
|
||||
fw = bn(x).contiguous_backward().relu().contiguous()
|
||||
fw.sum().backward()
|
||||
# TODO: this is too many
|
||||
check_schedule([x.grad, bn.weight.grad, bn.bias.grad, fw], 10)
|
||||
|
||||
def test_fold_conv_relu(self):
|
||||
c1 = nn.Conv2d(3,16,3)
|
||||
@@ -713,7 +711,7 @@ class TestSchedule(unittest.TestCase):
|
||||
check_schedule(b, 0)
|
||||
self.assertEqual(b.item(), 1)
|
||||
|
||||
@expect_rangeify_fails
|
||||
@unittest.expectedFailure
|
||||
def test_multioutput_ast(self):
|
||||
a = Tensor.zeros(1, dtype=dtypes.int).contiguous().realize().uop
|
||||
b = Tensor.zeros(1, dtype=dtypes.int).contiguous().realize().uop
|
||||
@@ -920,7 +918,7 @@ class TestSchedule(unittest.TestCase):
|
||||
out0 = a.sum() + 2
|
||||
out1 = a.sum() + 4
|
||||
out2 = out0 * out1
|
||||
run_schedule(check_schedule([out0, out1, out2], 1 if RANGEIFY else 4))
|
||||
run_schedule(check_schedule([out0, out1, out2], 1))
|
||||
np.testing.assert_allclose(out0.numpy(), out0_np:=a.numpy().sum()+2, atol=1e-4, rtol=1e-6)
|
||||
np.testing.assert_allclose(out1.numpy(), out1_np:=a.numpy().sum()+4, atol=1e-4, rtol=1e-6)
|
||||
np.testing.assert_allclose(out2.numpy(), out0_np*out1_np, atol=1e-4, rtol=1e-6)
|
||||
@@ -931,7 +929,7 @@ class TestSchedule(unittest.TestCase):
|
||||
out0 = a.sum().exp2()
|
||||
# out1 has two paths to a.sum()
|
||||
out1 = a.sum() + out0
|
||||
run_schedule(check_schedule([out0, out1], 1 if RANGEIFY else 3))
|
||||
run_schedule(check_schedule([out0, out1], 1))
|
||||
np.testing.assert_allclose(out0.numpy(), out0_np:=np.exp2(a.numpy().sum()), atol=1e-4, rtol=1e-4)
|
||||
np.testing.assert_allclose(out1.numpy(), a.numpy().sum()+out0_np, atol=1e-4, rtol=1e-6)
|
||||
|
||||
@@ -1023,7 +1021,7 @@ class TestSchedule(unittest.TestCase):
|
||||
b = Tensor.empty(10,)
|
||||
c = a.sum() + b[0]
|
||||
d = a.sum() + 2
|
||||
check_schedule([c, d], 1 if RANGEIFY else 3)
|
||||
check_schedule([c, d], 1)
|
||||
|
||||
def test_reduce_multiple_paths_midshrink(self):
|
||||
a = Tensor.empty(4, 4)
|
||||
@@ -1187,14 +1185,14 @@ class TestSchedule(unittest.TestCase):
|
||||
np.testing.assert_allclose(out.numpy(), expected, atol=1e-4, rtol=1e-4)
|
||||
|
||||
@unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
|
||||
@expect_rangeify_fails
|
||||
@unittest.expectedFailure
|
||||
def test_softmax_upcast(self):
|
||||
# input half, softmax in float
|
||||
Tensor.manual_seed(0)
|
||||
x = Tensor.randn(4, 12, 64, 64, dtype=dtypes.half).realize()
|
||||
out = x.softmax(dtype=dtypes.float)
|
||||
sched = out.schedule()
|
||||
self.assertEqual(len(sched), 2 if RANGEIFY else 3)
|
||||
self.assertEqual(len(sched), 2)
|
||||
self.assertEqual(sched[0].bufs[0].dtype, dtypes.half)
|
||||
|
||||
# input float, softmax in float
|
||||
@@ -1321,10 +1319,10 @@ class TestSchedule(unittest.TestCase):
|
||||
opt = nn.optim.SGD(nn.state.get_parameters([c1, c2, c3, c4]))
|
||||
opt.zero_grad()
|
||||
c4(c3(c2(c1(img).relu()).relu()).relu()).relu().sum().backward()
|
||||
with Context(FUSE_CONV_BW=1): check_schedule(opt.schedule_step(), 14)
|
||||
check_schedule(opt.schedule_step(), 14)
|
||||
|
||||
@unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
|
||||
@expect_rangeify_fails
|
||||
@unittest.expectedFailure
|
||||
def test_prefer_half_buffer(self):
|
||||
x = Tensor.ones(4).contiguous().realize()
|
||||
# y = Tensor.ones(4).contiguous().realize()
|
||||
@@ -1476,7 +1474,7 @@ class TestSchedule(unittest.TestCase):
|
||||
e = c * d
|
||||
f = b.sum() - e
|
||||
# run_schedule(check_schedule([c, d, e, f], 1))
|
||||
run_schedule(check_schedule([c, d, e, f], 2 if RANGEIFY else 5))
|
||||
run_schedule(check_schedule([c, d, e, f], 2))
|
||||
np.testing.assert_allclose(c.numpy(), c_np:=a.numpy().sum()+2, atol=1e-4, rtol=1e-4)
|
||||
np.testing.assert_allclose(d.numpy(), d_np:=a.numpy().sum()*2, atol=1e-4, rtol=1e-4)
|
||||
np.testing.assert_allclose(e.numpy(), e_np:=c_np*d_np, atol=1e-4, rtol=1e-4)
|
||||
@@ -1625,15 +1623,14 @@ class TestSchedule(unittest.TestCase):
|
||||
out = x.argmax(1)
|
||||
run_schedule(check_schedule(out, 2))
|
||||
|
||||
def test_conv2d(self): _test_conv2d(5 if RANGEIFY else 7)
|
||||
def test_conv2d_fused(self): _test_conv2d(5 if RANGEIFY else 5, FUSE_CONV_BW=1)
|
||||
def test_conv2d(self): _test_conv2d(5 if SPLIT_REDUCEOP else 4)
|
||||
def test_conv2d_fused(self): _test_conv2d(5 if SPLIT_REDUCEOP else 4)
|
||||
|
||||
@unittest.skipUnless(is_dtype_supported(dtypes.half) and is_dtype_supported(dtypes.ulong), "need half and ulong")
|
||||
def test_conv2d_half(self): _test_conv2d(5 if RANGEIFY else 7, dtype=dtypes.half)
|
||||
def test_conv2d_half(self): _test_conv2d(5 if SPLIT_REDUCEOP else 4, dtype=dtypes.half)
|
||||
@unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
|
||||
@unittest.skipIf(Device.DEFAULT == "WEBGPU", "Causes other tests to fail")
|
||||
@unittest.skipIf(not RANGEIFY, "passes on RANGEIFY")
|
||||
def test_conv2d_fused_half(self): _test_conv2d(5, dtype=dtypes.half)
|
||||
def test_conv2d_fused_half(self): _test_conv2d(5 if SPLIT_REDUCEOP else 4, dtype=dtypes.half)
|
||||
|
||||
def test_schedule_mem_used(self):
|
||||
base = GlobalCounters.mem_used
|
||||
@@ -1692,7 +1689,7 @@ class TestSchedule(unittest.TestCase):
|
||||
def test_late_fusion_post_expand(self):
|
||||
self._test_fusion([(32, 32)], lambda a:a-a.sum(1), 2)
|
||||
|
||||
@expect_rangeify_fails
|
||||
@unittest.expectedFailure
|
||||
def test_cast_padded_view(self):
|
||||
a = Tensor.arange(4).reshape(1, 4)
|
||||
casted_view = a.pad(((0, 1), (0, 0))).cast(dtypes.float)
|
||||
@@ -1722,7 +1719,7 @@ class TestSchedule(unittest.TestCase):
|
||||
self.assertListEqual(realized_const_view.tolist(), [[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1]])
|
||||
|
||||
@given(strat.sampled_from(dtypes.all), strat.sampled_from(dtypes.all))
|
||||
@expect_rangeify_fails
|
||||
@unittest.expectedFailure
|
||||
def test_cast_padded_const(self, dt1, dt2):
|
||||
assume(is_dtype_supported(dt1) and is_dtype_supported(dt2))
|
||||
a = Tensor(1, dtype=dt1).reshape(1, 1).pad(((1, 1), None))
|
||||
@@ -1893,9 +1890,7 @@ class TestSchedule(unittest.TestCase):
|
||||
tst = x.shrink((None, (0, 2))).assign(a).realize()
|
||||
xref[:, :2] = np.arange(8).reshape(4, 2)+y.numpy()
|
||||
np.testing.assert_equal(x.numpy(), xref)
|
||||
if RANGEIFY > 0:
|
||||
# NOTE: this is a bug on non rangeify
|
||||
np.testing.assert_equal(tst.numpy(), a.numpy())
|
||||
np.testing.assert_equal(tst.numpy(), a.numpy())
|
||||
|
||||
def test_setitem_sched(self, mop=lambda x:x, expected_kcount=1):
|
||||
a = Tensor.arange(16, device="CPU").reshape(4, 4).contiguous().realize()
|
||||
@@ -1906,7 +1901,6 @@ class TestSchedule(unittest.TestCase):
|
||||
run_schedule(sched)
|
||||
self.assertListEqual(a.tolist(), expected)
|
||||
self.assertEqual(kcount, expected_kcount)
|
||||
@unittest.skipUnless(RANGEIFY>0, "this asserts on non rangeify")
|
||||
def test_setitem_permuted_sched(self): self.test_setitem_sched(lambda x: x.T, 2)
|
||||
def test_setitem_paddded_sched(self): self.test_setitem_sched(lambda x: x.shrink_to(4, 1).pad_to(4, 4), 1)
|
||||
|
||||
@@ -1945,7 +1939,7 @@ class TestSchedule(unittest.TestCase):
|
||||
r = (X+Tensor.arange(16).reshape(4, 4)).sum()
|
||||
out0 = r+2
|
||||
out1 = r+3
|
||||
run_schedule(check_schedule([out0, out1], 1 if RANGEIFY else 3))
|
||||
run_schedule(check_schedule([out0, out1], 1))
|
||||
r_ref = (X.numpy()+np.arange(16).reshape(4, 4)).sum()
|
||||
np.testing.assert_allclose(out0.numpy(), r_ref+2, rtol=2e-7)
|
||||
np.testing.assert_allclose(out1.numpy(), r_ref+3, rtol=2e-7)
|
||||
@@ -2090,7 +2084,7 @@ class TestView(unittest.TestCase):
|
||||
run_schedule(sched)
|
||||
np.testing.assert_equal(b.numpy(), 0)
|
||||
|
||||
@expect_rangeify_fails
|
||||
@unittest.expectedFailure
|
||||
def test_mask_dim_1(self):
|
||||
# mask out dim = 1 works too
|
||||
a = Tensor.rand(10, 10).realize()
|
||||
@@ -2238,7 +2232,6 @@ class TestCopyFolding(unittest.TestCase):
|
||||
b.realize()
|
||||
self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])
|
||||
|
||||
@expect_nonrangeify_fails
|
||||
def test_permute_on_disk_contiguous(self):
|
||||
with open(temp('dt_arange_4_permute'), "wb") as f: f.write(Tensor.arange(4).realize().uop.base.buffer.as_buffer())
|
||||
a = Tensor.empty(4, dtype=dtypes.int32, device=f"disk:{temp('dt_arange_4_permute')}")
|
||||
@@ -2253,8 +2246,6 @@ class TestCopyFolding(unittest.TestCase):
|
||||
self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])
|
||||
|
||||
# NOTE: disk permute must come after COPY
|
||||
# TODO: this is wrong because of the permute
|
||||
@expect_nonrangeify_fails
|
||||
def test_permute_after_shrink_on_disk(self):
|
||||
with open(temp('dt_arange_5_permute'), "wb") as f: f.write(Tensor.arange(5).realize().uop.base.buffer.as_buffer())
|
||||
a = Tensor.empty(5, dtype=dtypes.int32, device=f"disk:{temp('dt_arange_5_permute')}")
|
||||
@@ -2291,7 +2282,7 @@ class TestBufferUOp(unittest.TestCase):
|
||||
|
||||
def test_buffer_view_not_allowed(self):
|
||||
permuted_view = Tensor.empty(1, 2, 3).permute(0, 2, 1)
|
||||
with self.assertRaisesRegex(AssertionError, "VIEW only works here if it's contiguous"):
|
||||
with self.assertRaisesRegex(AssertionError, "can only be RESHAPE"):
|
||||
permuted_view.uop.buffer # cannot access Buffer of a non contiguous VIEW
|
||||
|
||||
def test_buffer_only_after_realize(self):
|
||||
@@ -2398,12 +2389,8 @@ class TestUOpBecome(unittest.TestCase):
|
||||
a = Tensor.empty(4, 1)
|
||||
b = a.expand(4, 4).reciprocal()
|
||||
check_schedule(b, 1)
|
||||
if RANGEIFY:
|
||||
self.assertEqual(b.uop.base.buffer.size, 4)
|
||||
self.assertEqual(b.uop.shape, (4, 4))
|
||||
return
|
||||
self.assertEqual(b.uop.base.buffer.size, 16)
|
||||
self.assertEqual(b.uop.st, ShapeTracker.from_shape((4, 4)))
|
||||
self.assertEqual(b.uop.base.buffer.size, 4)
|
||||
self.assertEqual(b.uop.shape, (4, 4))
|
||||
|
||||
def test_reorder_expand_alt(self):
|
||||
x = Tensor.empty(4, 1)
|
||||
@@ -2412,7 +2399,7 @@ class TestUOpBecome(unittest.TestCase):
|
||||
z = (img*x) / y
|
||||
check_schedule(z, 1)
|
||||
|
||||
@expect_rangeify_fails
|
||||
@unittest.expectedFailure
|
||||
def test_become_existing_buffer(self):
|
||||
a = Tensor.empty(4, 4)
|
||||
b = a*1
|
||||
@@ -2446,7 +2433,7 @@ class TestUOpBecome(unittest.TestCase):
|
||||
assert UPat(Ops.CONST, arg=3).match(const_add.uop.base, {})
|
||||
|
||||
# tensors can become another realized tensor source
|
||||
@expect_rangeify_fails
|
||||
@unittest.expectedFailure
|
||||
def test_become_existing_buf_simple(self):
|
||||
a = Tensor.empty(4, 4)
|
||||
b = a+0
|
||||
@@ -2455,14 +2442,14 @@ class TestUOpBecome(unittest.TestCase):
|
||||
self.assertIs(a.uop, b.uop)
|
||||
|
||||
# they can also chain other movement ops on top of the tensor source
|
||||
@expect_rangeify_fails
|
||||
@unittest.expectedFailure
|
||||
def test_become_existing_buf_view(self):
|
||||
a = Tensor.empty(4, 4)
|
||||
b = a.permute((1, 0))+0
|
||||
check_schedule(b, 0)
|
||||
self.assertEqual(b.uop.st, a.uop.permute((1, 0)).st)
|
||||
|
||||
@expect_rangeify_fails
|
||||
@unittest.expectedFailure
|
||||
def test_become_existing_buf_view_alt(self):
|
||||
a = Tensor.empty(4, 4)
|
||||
b = a.permute((1, 0)).reshape((8, 2))+0
|
||||
@@ -2470,7 +2457,7 @@ class TestUOpBecome(unittest.TestCase):
|
||||
self.assertEqual(b.uop.st, a.uop.permute((1, 0)).reshape((8, 2)).st)
|
||||
|
||||
# they can also have other base parents that simplified, in that case we just backtrack to the chained mops
|
||||
@expect_rangeify_fails
|
||||
@unittest.expectedFailure
|
||||
def test_become_existing_buf_complex(self):
|
||||
a = Tensor.empty(4, 4)
|
||||
b = (a.permute((1, 0))+0).reshape((8, 2))+0
|
||||
@@ -2478,7 +2465,7 @@ class TestUOpBecome(unittest.TestCase):
|
||||
self.assertEqual(b.uop.st, a.uop.permute((1, 0)).reshape((8, 2)).st)
|
||||
assert b.uop.base.op is Ops.BUFFER
|
||||
|
||||
@expect_rangeify_fails
|
||||
@unittest.expectedFailure
|
||||
def test_become_multiple_choices(self):
|
||||
a = Tensor.empty(16)
|
||||
b = (a.reshape(1, 1, 4, 1, 4)+0).reshape(1, 1, 4, 4).shrink(((0, 1), (0, 1), (0, 3), (0, 3)))+0
|
||||
@@ -2496,13 +2483,8 @@ class TestUOpBecome(unittest.TestCase):
|
||||
b.realize()
|
||||
assert a.uop.is_realized
|
||||
assert a.uop.buffer._base is None
|
||||
# b is a subbuffer of a (buffer_view in non rangeify, rangeify just makes a shrink)
|
||||
if RANGEIFY:
|
||||
assert b.uop.op_in_backward_slice_with_self(Ops.SHRINK)
|
||||
assert b.uop.base is a.uop.base
|
||||
return
|
||||
assert b.uop.op is Ops.BUFFER_VIEW
|
||||
assert b.uop.src[0] is a.uop
|
||||
assert b.uop.op_in_backward_slice_with_self(Ops.SHRINK)
|
||||
assert b.uop.base is a.uop.base
|
||||
|
||||
def test_setitem_offset(self):
|
||||
a = Tensor.full((16,), 0.).contiguous().realize()
|
||||
|
||||
@@ -2,7 +2,7 @@ import unittest
|
||||
import numpy as np
|
||||
from tinygrad import Tensor, GlobalCounters, Context, Device
|
||||
from tinygrad.dtype import DTypeLike, dtypes
|
||||
from tinygrad.helpers import DEBUG, get_single_element, RANGEIFY
|
||||
from tinygrad.helpers import DEBUG, get_single_element
|
||||
from tinygrad.engine.realize import lower_schedule_item
|
||||
from tinygrad.device import is_dtype_supported
|
||||
|
||||
@@ -39,17 +39,17 @@ class TestFuse(unittest.TestCase):
|
||||
np_multi = fxn(*args, **kwargs).numpy()
|
||||
np.testing.assert_allclose(np_single, np_multi, atol=atol)
|
||||
|
||||
@unittest.skipIf(0<RANGEIFY<2, "needs RANGEIFY>1")
|
||||
@unittest.skip("needs RANGEIFY>1")
|
||||
def test_fuse_norm(self):
|
||||
a = Tensor.rand(50,50).realize()
|
||||
self._test_fuse(lambda a: a / a.mean(axis=1), a)
|
||||
|
||||
@unittest.skipIf(0<RANGEIFY<2, "needs RANGEIFY>1")
|
||||
@unittest.skip("needs RANGEIFY>1")
|
||||
def test_fuse_argmax(self):
|
||||
a = Tensor.rand(50,50).realize()
|
||||
self._test_fuse(lambda a: a.argmax(axis=-1), a)
|
||||
|
||||
@unittest.skipIf(0<RANGEIFY<2, "needs RANGEIFY>1")
|
||||
@unittest.skip("needs RANGEIFY>1")
|
||||
def test_fuse_softmax(self):
|
||||
a = Tensor.rand(50,50).realize()
|
||||
self._test_fuse(lambda a: a.softmax(axis=-1), a)
|
||||
@@ -60,7 +60,7 @@ class TestFuse(unittest.TestCase):
|
||||
self._test_fuse(lambda a,b: ((a@b).relu()+a).contiguous().softmax(axis=-1), a,b, allow_multiple=True)
|
||||
|
||||
@unittest.skipUnless(is_dtype_supported(dtypes.float16, Device.DEFAULT), f"no float16 on {Device.DEFAULT}")
|
||||
@unittest.skipIf(0<RANGEIFY<2, "needs RANGEIFY>1")
|
||||
@unittest.skip("needs RANGEIFY>1")
|
||||
def test_fuse_softmax_dtype(self):
|
||||
a = Tensor.rand(50,50).realize()
|
||||
self._test_fuse(lambda a: a.softmax(axis=-1, dtype='half'), a, atol=3e-4)
|
||||
@@ -68,7 +68,7 @@ class TestFuse(unittest.TestCase):
|
||||
def test_fuse_arange_eye(self):
|
||||
self._test_fuse(lambda: Tensor.arange(10).reshape(10,1).expand(10,10) == Tensor.arange(10).reshape(1,10).expand(10,10))
|
||||
|
||||
@unittest.skipIf(0<RANGEIFY<2, "needs RANGEIFY>1")
|
||||
@unittest.skip("needs RANGEIFY>1")
|
||||
def test_double_gemm(self):
|
||||
N = 32
|
||||
with Context(TRACK_MATCH_STATS=0, DEBUG=0):
|
||||
@@ -91,7 +91,7 @@ class TestFuse(unittest.TestCase):
|
||||
return (arange == idx).mul(vals).sum(-2, dtype=vals.dtype)
|
||||
self._test_fuse(embedding, a, atol=1e-5)
|
||||
|
||||
@unittest.skipIf(0<RANGEIFY<2, "needs RANGEIFY>1")
|
||||
@unittest.skip("needs RANGEIFY>1")
|
||||
def test_attention_kernel_count(self):
|
||||
wq = Tensor.empty(32, 32)
|
||||
wk = Tensor.empty(32, 32)
|
||||
@@ -104,7 +104,7 @@ class TestFuse(unittest.TestCase):
|
||||
s = attn.schedule()
|
||||
self.assertEqual(len(s), 4) # 3 matmul and 1 attention
|
||||
|
||||
@unittest.skipIf(0<RANGEIFY<2, "needs RANGEIFY>1")
|
||||
@unittest.skip("needs RANGEIFY>1")
|
||||
def test_flash_attention(self):
|
||||
BS = 4
|
||||
HEADS = 2
|
||||
@@ -172,7 +172,7 @@ class TestSoftmaxFusion(unittest.TestCase):
|
||||
|
||||
np.testing.assert_allclose(sout.numpy(), out.numpy(), atol=3e-7)
|
||||
|
||||
@unittest.skipIf(0<RANGEIFY<2, "needs RANGEIFY>1")
|
||||
@unittest.skip("needs RANGEIFY>1")
|
||||
def test_auto_softmax(self):
|
||||
print("*** softmax ***")
|
||||
with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
|
||||
|
||||
@@ -2,7 +2,6 @@ import unittest
|
||||
|
||||
from test.helpers import assert_jit_cache_len
|
||||
from tinygrad import Variable, Tensor, TinyJit
|
||||
from tinygrad.helpers import RANGEIFY
|
||||
import numpy as np
|
||||
|
||||
class TestSymbolicJit(unittest.TestCase):
|
||||
@@ -27,7 +26,7 @@ class TestSymbolicJit(unittest.TestCase):
|
||||
symbolic = jf(a[:, :vi]).numpy()
|
||||
expected = f(a[:, :i]).numpy()
|
||||
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
|
||||
assert_jit_cache_len(jf, 1 if RANGEIFY else 2) # one add and one pad, can be one kernel?
|
||||
assert_jit_cache_len(jf, 1)
|
||||
|
||||
def test_add(self):
|
||||
def f(a, b): return (a+b).realize()
|
||||
@@ -80,7 +79,7 @@ class TestSymbolicJit(unittest.TestCase):
|
||||
symbolic = jf(q, k[:, :vi], v[:, :vi])[:2, :4, :1, :8].numpy()
|
||||
expected = f(q, k[:, :i], v[:, :i]).numpy()
|
||||
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
|
||||
assert_jit_cache_len(jf, 4 if RANGEIFY else 5)
|
||||
assert_jit_cache_len(jf, 4)
|
||||
|
||||
def test_cat_dim0(self):
|
||||
def f(a, b): return a.cat(b, dim=0).realize()
|
||||
|
||||
+6
-13
@@ -4,7 +4,7 @@ import torch
|
||||
import unittest, copy, mmap, random, math, array
|
||||
from tinygrad import Tensor, Device, dtypes
|
||||
from tinygrad.tensor import _METADATA
|
||||
from tinygrad.helpers import getenv, temp, mv_address, RANGEIFY
|
||||
from tinygrad.helpers import getenv, temp, mv_address
|
||||
from extra.gradcheck import numerical_jacobian, jacobian, gradcheck
|
||||
from hypothesis import given, settings, strategies as strat
|
||||
from tinygrad.device import is_dtype_supported
|
||||
@@ -872,18 +872,11 @@ class TestTensorMetadata(unittest.TestCase):
|
||||
self.assertEqual(y.grad.uop.metadata[0].name, "sigmoid")
|
||||
self.assertTrue(y.grad.uop.metadata[0].backward)
|
||||
si = Tensor.schedule(out, x.grad, y.grad)[-1]
|
||||
if not RANGEIFY:
|
||||
self.assertEqual(len(si.metadata), 4, f"failed with {si.metadata}")
|
||||
self.assertSetEqual(set(m.name for m in si.metadata), {"sigmoid", "__mul__", "relu"})
|
||||
bw = [m for m in si.metadata if m.backward]
|
||||
self.assertEqual(len(bw), 2)
|
||||
self.assertEqual(bw[0].name, "sigmoid")
|
||||
else:
|
||||
self.assertEqual(len(si.metadata), 3, f"failed with {si.metadata}")
|
||||
self.assertSetEqual(set(m.name for m in si.metadata), {"sigmoid", "relu"})
|
||||
bw = [m for m in si.metadata if m.backward]
|
||||
self.assertEqual(len(bw), 1)
|
||||
self.assertEqual(bw[0].name, "sigmoid")
|
||||
self.assertEqual(len(si.metadata), 3, f"failed with {si.metadata}")
|
||||
self.assertSetEqual(set(m.name for m in si.metadata), {"sigmoid", "relu"})
|
||||
bw = [m for m in si.metadata if m.backward]
|
||||
self.assertEqual(len(bw), 1)
|
||||
self.assertEqual(bw[0].name, "sigmoid")
|
||||
|
||||
class TestIdxUpcast(unittest.TestCase):
|
||||
def _find_op(self, ast: UOp, op: Ops):
|
||||
|
||||
@@ -4,6 +4,7 @@ import unittest
|
||||
from tinygrad import Tensor, Device, dtypes
|
||||
from tinygrad.engine.realize import run_schedule
|
||||
from tinygrad.uop.ops import Ops, UOp, UPat
|
||||
from tinygrad.helpers import SPLIT_REDUCEOP
|
||||
|
||||
class TestTensorUOp(unittest.TestCase):
|
||||
def test_fromcpu_shape_tracker(self):
|
||||
@@ -94,6 +95,7 @@ class TestTensorUOp(unittest.TestCase):
|
||||
self.assertEqual(out.tolist(), Tensor.zeros(4, 8).tolist())
|
||||
|
||||
reduce_kernel = UPat(Ops.SINK, src=(UPat(Ops.STORE, allow_any_len=True, src=(UPat(), UPat((Ops.REDUCE_AXIS, Ops.REDUCE))))))
|
||||
@unittest.skipUnless(SPLIT_REDUCEOP, "only for SPLIT_REDUCEOP")
|
||||
class TestReduceOp(unittest.TestCase):
|
||||
def test_no_split_reduce_kernel(self):
|
||||
a = Tensor.rand(4, 4).realize()
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
import unittest
|
||||
from tinygrad import Tensor
|
||||
from tinygrad.helpers import getenv, GlobalCounters, EMULATE, RANGEIFY
|
||||
from tinygrad.helpers import getenv, GlobalCounters, EMULATE
|
||||
from tinygrad.engine.realize import lower_schedule_item, ProgramSpec, get_program
|
||||
from tinygrad.renderer import Estimates
|
||||
from tinygrad.codegen import full_rewrite
|
||||
@@ -51,11 +51,8 @@ class TestMemoryCount(unittest.TestCase):
|
||||
a = Tensor.empty(1024, 1, dtype=dtypes.uint8).expand(1024, 1024)
|
||||
b = Tensor.empty(1024, 1, dtype=dtypes.uint8).expand(1024, 1024)
|
||||
_, mem = get_stats(a+b)
|
||||
if RANGEIFY:
|
||||
# rangeify is smart!
|
||||
self.assertEqual(mem, 1024 + 2*1024) # 2 lil reads + 1 lil write
|
||||
else:
|
||||
self.assertEqual(mem, 1024*1024 + 2*1024) # 2 lil reads + 1 write
|
||||
# rangeify is smart!
|
||||
self.assertEqual(mem, 1024 + 2*1024) # 2 lil reads + 1 lil write
|
||||
|
||||
def test_self_add(self):
|
||||
a = Tensor.empty(1024, 1024, dtype=dtypes.uint8)
|
||||
|
||||
@@ -1,12 +1,10 @@
|
||||
import unittest
|
||||
from tinygrad import Tensor, dtypes, TinyJit, UOp
|
||||
from tinygrad.helpers import RANGEIFY
|
||||
from tinygrad.apps.llm import apply_rope
|
||||
#from tinygrad.engine.realize import run_schedule
|
||||
|
||||
# TODO: test_scheduler, but just in uint
|
||||
class TestAttention(unittest.TestCase):
|
||||
@unittest.skipIf(RANGEIFY > 0, "not half on rangeify")
|
||||
def test_half_qkv_buffers(self):
|
||||
BS, seqlen, dim = 10, 4, 100
|
||||
q = Tensor.ones(BS, seqlen, dim, dtype=dtypes.half).contiguous().realize()
|
||||
@@ -14,12 +12,11 @@ class TestAttention(unittest.TestCase):
|
||||
v = Tensor.ones(BS, seqlen, dim, dtype=dtypes.half).contiguous().realize()
|
||||
attn = q.scaled_dot_product_attention(k, v)
|
||||
sched = attn.schedule()
|
||||
#run_schedule(sched[:])
|
||||
# attention has 5 kernels now
|
||||
self.assertEqual(len(sched), 4 if RANGEIFY else 5)
|
||||
softmax_inputs = sched[1:4]
|
||||
for i,si in enumerate(softmax_inputs):
|
||||
assert all(b.dtype == dtypes.half for b in si.bufs), f"non half {si.bufs=} in kernel {i}"
|
||||
# attention has 4 kernels now
|
||||
self.assertEqual(len(sched), 4)
|
||||
# softmax_inputs = sched[1:4]
|
||||
# for i,si in enumerate(softmax_inputs):
|
||||
# assert all(b.dtype == dtypes.half for b in si.bufs), f"non half {si.bufs=} in kernel {i}"
|
||||
|
||||
def test_apply_rope(self):
|
||||
x = Tensor.randn(1, 2, 4, 8, dtype=dtypes.float32)
|
||||
|
||||
@@ -29,8 +29,11 @@ class TestKeccak(unittest.TestCase):
|
||||
out_shape = Tensor.randint(*s[i:], high=255, dtype=dtypes.uint8).keccak().shape
|
||||
self.assertTupleEqual(s[i:-1], out_shape[:-1])
|
||||
|
||||
@unittest.skipUnless(Device.DEFAULT=="METAL", "slow")
|
||||
def test_sha3_224(self): self._test_preset("sha3_224", [143, 144])
|
||||
@unittest.skipUnless(Device.DEFAULT=="METAL", "slow")
|
||||
def test_sha3_256(self): self._test_preset("sha3_256", [135, 136])
|
||||
@unittest.skipUnless(Device.DEFAULT=="METAL", "slow")
|
||||
def test_shake_128(self): self._test_preset("shake_128", [167, 168], lambda d: hashlib.shake_128(d).digest(16))
|
||||
|
||||
def _test_preset(self, name: str, special_sizes: list[int], hasher: Callable[[bytes], bytes] | None = None):
|
||||
|
||||
@@ -3,7 +3,7 @@ from tinygrad import Variable
|
||||
from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray, CI, mv_address
|
||||
from tinygrad.helpers import merge_dicts, strip_parens, prod, round_up, fetch, fully_flatten, from_mv, to_mv, polyN, time_to_str, cdiv, cmod, getbits
|
||||
from tinygrad.tensor import Tensor, get_shape
|
||||
from tinygrad.shape.view import get_contraction, get_contraction_with_reduce
|
||||
from tinygrad.shape.view import get_contraction
|
||||
import numpy as np
|
||||
|
||||
VARIABLE = ContextVar("VARIABLE", 0)
|
||||
@@ -219,20 +219,6 @@ class TestMemoryview(unittest.TestCase):
|
||||
print(f"from_mv vs mv_address: {fmv_us:8.3f} µs vs {mva_us:8.3f} µs")
|
||||
|
||||
class TestGetContraction(unittest.TestCase):
|
||||
def test_contraction_with_reduce(self):
|
||||
r = get_contraction((16, 1, 1, 1), (16, 1, 1))
|
||||
self.assertEqual(r, [[0], [], [1, 2, 3]])
|
||||
r = get_contraction_with_reduce((16, 1, 1, 1), (16, 1, 1), (1,))
|
||||
self.assertEqual(r, [[0], [1, 2], [3]])
|
||||
|
||||
r = get_contraction((16, 1, 1, 1, 1), (16, 1, 1, 1))
|
||||
self.assertEqual(r, [[0], [], [], [1, 2, 3, 4]])
|
||||
r = get_contraction_with_reduce((16, 1, 1, 1, 1), (16, 1, 1, 1), (1,))
|
||||
self.assertEqual(r, [[0], [1, 2], [3], [4]])
|
||||
|
||||
r = get_contraction_with_reduce((2, 512, 1, 1), (2, 1, 512), (1,))
|
||||
self.assertIsNone(r)
|
||||
|
||||
def test_contraction(self):
|
||||
r = get_contraction((1,2,3,4), (2,3,4))
|
||||
self.assertEqual(r, [[0, 1], [2], [3]])
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
import unittest
|
||||
from tinygrad import Tensor
|
||||
from tinygrad.uop import Ops
|
||||
from tinygrad.helpers import RANGEIFY
|
||||
|
||||
class TestKernelize(unittest.TestCase):
|
||||
def test_add_reshaped(self):
|
||||
@@ -18,8 +17,8 @@ class TestKernelize(unittest.TestCase):
|
||||
a1 = a.sum(axis=1)
|
||||
a0 = a1.sum(axis=0)
|
||||
a0.kernelize()
|
||||
self.assertEqual(len([s for s in a0.uop.toposort() if s.op is Ops.KERNEL]), 2 if RANGEIFY else 3)
|
||||
self.assertIs(a1.uop.base.op, Ops.REDUCE_AXIS if RANGEIFY else Ops.ASSIGN)
|
||||
self.assertEqual(len([s for s in a0.uop.toposort() if s.op is Ops.KERNEL]), 2)
|
||||
self.assertIs(a1.uop.base.op, Ops.REDUCE_AXIS)
|
||||
# input Tensor and user contiguous kernelize
|
||||
self.assertIs(a0.uop.base.op, Ops.ASSIGN)
|
||||
self.assertIs(a.uop.base.op, Ops.ASSIGN)
|
||||
|
||||
@@ -3,14 +3,14 @@ import unittest
|
||||
import numpy as np
|
||||
from tinygrad.dtype import dtypes, Invalid
|
||||
from tinygrad.helpers import prod
|
||||
from tinygrad.shape.shapetracker import ShapeTracker, View
|
||||
from tinygrad.shape.shapetracker import ShapeTracker, View, views_to_valid_uop
|
||||
from tinygrad import Variable
|
||||
from tinygrad.uop.ops import UOp, Ops, graph_rewrite
|
||||
from tinygrad.codegen.late.devectorizer import sym
|
||||
from itertools import product
|
||||
|
||||
def shapetracker_getitem(st:ShapeTracker, val:int):
|
||||
valid_idx = st.reshape((st.size,)).to_valid_uop([UOp.const(dtypes.int, val)])
|
||||
valid_idx = views_to_valid_uop(st.reshape((st.size,)).views, (UOp.const(dtypes.int, val),))
|
||||
idx, valid = valid_idx.get_idx(), valid_idx.get_valid()
|
||||
idx, valid = graph_rewrite(idx, sym), graph_rewrite(valid, sym)
|
||||
assert idx.op is Ops.CONST and valid.op is Ops.CONST
|
||||
@@ -175,12 +175,6 @@ class TestRealSimplifies(unittest.TestCase):
|
||||
View.create((8, 3, 3, 11, 2, 28), (924, 308, 0, 28, 0, 1), 0, None),
|
||||
View.create((8, 1, 6, 10, 28, 3, 2, 1), (5544, 0, 0, 56, 1, 1848, 672, 0), 0, None)))
|
||||
|
||||
class TestViewMinify(unittest.TestCase):
|
||||
def test_minifies(self):
|
||||
assert len(View.create((10,10)).minify().shape) == 1
|
||||
assert len(View.create((10,10)).permute((1,0)).minify().shape) == 2
|
||||
assert len(View.create((10,10,10,10)).permute((1,0,2,3)).minify().shape) == 3
|
||||
|
||||
class TestIndexExpressions2d(unittest.TestCase):
|
||||
def setUp(self):
|
||||
shapes = [(30, 5), (15, 10), (15, 1), (5, 10), (5, 1)] # Make sure dim0 is a multiple of 5, one of the tests divides this dimension by 5
|
||||
|
||||
@@ -1,11 +1,11 @@
|
||||
import unittest
|
||||
import multiprocessing.shared_memory as shared_memory
|
||||
from tinygrad.helpers import CI, WIN, RANGEIFY
|
||||
from tinygrad.helpers import CI, WIN
|
||||
from tinygrad.tensor import Tensor, Device
|
||||
import numpy as np
|
||||
|
||||
class TestRawShmBuffer(unittest.TestCase):
|
||||
@unittest.skipIf(WIN and CI and RANGEIFY, "only fails with RANGEIFY on CI windows instance")
|
||||
@unittest.skipIf(WIN and CI, "only fails on CI windows instance")
|
||||
def test_e2e(self):
|
||||
t = Tensor.randn(2, 2, 2).realize()
|
||||
|
||||
|
||||
@@ -10,25 +10,6 @@ class TestView(unittest.TestCase):
|
||||
v = View.create(shape=(4,3,2), strides=(1,4,10), mask=((0,4),(0,3),(0,2)))
|
||||
self.assertIsNone(v.mask)
|
||||
|
||||
def test_minify_zero_strided_dims(self):
|
||||
target = View.create(shape=(2,2), strides=(30,2), offset=7, mask=None)
|
||||
v = View.create(shape=(2,1,2), strides=(30,0,2), offset=7, mask=None)
|
||||
self.assertEqual(v.minify(), target)
|
||||
v = View.create(shape=(1,2,2), strides=(0,30,2), offset=7, mask=None)
|
||||
self.assertEqual(v.minify(), target)
|
||||
v = View.create(shape=(2,2,1), strides=(30,2,0), offset=7, mask=None)
|
||||
self.assertEqual(v.minify(), target)
|
||||
v = View.create(shape=(2,1,1,2), strides=(30,0,0,2), offset=7, mask=None)
|
||||
self.assertEqual(v.minify(), target)
|
||||
v = View.create(shape=(1,1,2,2), strides=(0,0,30,2), offset=7, mask=None)
|
||||
self.assertEqual(v.minify(), target)
|
||||
v = View.create(shape=(2,2,1,1), strides=(30,2,0,0), offset=7, mask=None)
|
||||
self.assertEqual(v.minify(), target)
|
||||
v = View.create(shape=(1,2,2,1), strides=(0,30,2,0), offset=7, mask=None)
|
||||
self.assertEqual(v.minify(), target)
|
||||
v = View.create(shape=(1,2,1,2), strides=(0,30,0,2), offset=7, mask=None)
|
||||
self.assertEqual(v.minify(), target)
|
||||
|
||||
def test_empty_mask_contiguous(self):
|
||||
v1 = View.create(shape=(2,2,2), strides=(4,2,1), mask=None)
|
||||
v2 = View.create(shape=(2,2,2), strides=(4,2,1), mask=((0,2),(0,2),(0,2)))
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
import unittest, sys
|
||||
import numpy as np
|
||||
from tinygrad import Tensor, GlobalCounters, dtypes, Context, nn
|
||||
from tinygrad.helpers import CI, Profiling, WINO, RANGEIFY
|
||||
from tinygrad.helpers import CI, Profiling, WINO
|
||||
|
||||
@unittest.skipIf(sys.platform.startswith("win"), "flaky on Windows")
|
||||
class TestWinogradClose(unittest.TestCase):
|
||||
@@ -35,14 +35,14 @@ class TestWinograd(unittest.TestCase):
|
||||
def test_forward_kernels(self):
|
||||
x,w = Tensor.rand(1,4,9,9).realize(), Tensor.rand(4,4,3,3).realize()
|
||||
out = Tensor.conv2d(x,w)
|
||||
self.assertEqual(len(out.schedule()), 2 if RANGEIFY else 4)
|
||||
self.assertEqual(len(out.schedule()), 2)
|
||||
|
||||
def test_backward_kernels(self):
|
||||
x,w = Tensor.empty(1,4,9,9,requires_grad=True).realize(), Tensor.empty(4,4,3,3,requires_grad=True).realize()
|
||||
out = Tensor.conv2d(x,w, padding=1)
|
||||
out.mean().backward()
|
||||
backward_schedule = Tensor.schedule(x.grad, w.grad)
|
||||
self.assertEqual(len(backward_schedule), 4 if RANGEIFY else 9)
|
||||
self.assertEqual(len(backward_schedule), 4)
|
||||
|
||||
def test_counters(self):
|
||||
IC, OC, X, Y = 4,4,9,9
|
||||
@@ -61,9 +61,9 @@ class TestWinograd(unittest.TestCase):
|
||||
print(f"ops: normal {ops_normal:9d} wino {ops_wino:9d} ratio {ops_ratio:.2f}")
|
||||
print(f"mem: normal {mem_normal:9d} wino {mem_wino:9d} ratio {mem_ratio:.2f}")
|
||||
|
||||
if not RANGEIFY:
|
||||
self.assertLess(ops_ratio, 2.6) # TODO: there's issues with factorization now
|
||||
self.assertLess(mem_ratio, 10)
|
||||
# TODO: what's optimal on this?
|
||||
self.assertLess(ops_ratio, 4.3)
|
||||
self.assertLess(mem_ratio, 3)
|
||||
|
||||
def test_dtype(self):
|
||||
IC, OC, X, Y = 4,4,9,9
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
from typing import Any, Callable
|
||||
import functools
|
||||
from dataclasses import dataclass
|
||||
from tinygrad.helpers import QUANTIZE, DEVECTORIZE, TRANSCENDENTAL, RANGEIFY
|
||||
from tinygrad.helpers import QUANTIZE, DEVECTORIZE, TRANSCENDENTAL
|
||||
from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype
|
||||
from tinygrad.uop.spec import type_verify
|
||||
from tinygrad.renderer import Renderer
|
||||
@@ -38,11 +38,10 @@ rewrites_for_linearizer = [
|
||||
|
||||
def get_rewrites_for_renderer(opts:Renderer, optimize:bool=True, linearizer:bool=True) -> list[RewriteStep]:
|
||||
# cache with the values of the context vars
|
||||
return _get_rewrites_for_renderer(opts, optimize, linearizer, QUANTIZE.value, DEVECTORIZE.value, TRANSCENDENTAL.value, RANGEIFY.value)
|
||||
return _get_rewrites_for_renderer(opts, optimize, linearizer, QUANTIZE.value, DEVECTORIZE.value, TRANSCENDENTAL.value)
|
||||
|
||||
@functools.cache
|
||||
def _get_rewrites_for_renderer(opts:Renderer, optimize:bool, linearizer:bool, _QUANTIZE, _DEVECTORIZE, _TRANSCENDENTAL,
|
||||
_RANGEIFY) -> list[RewriteStep]:
|
||||
def _get_rewrites_for_renderer(opts:Renderer, optimize:bool, linearizer:bool, _QUANTIZE, _DEVECTORIZE, _TRANSCENDENTAL) -> list[RewriteStep]:
|
||||
# ** lowerer (rewrite_shapetracker_with_index) **
|
||||
ret: list[RewriteStep] = []
|
||||
|
||||
@@ -52,8 +51,7 @@ def _get_rewrites_for_renderer(opts:Renderer, optimize:bool, linearizer:bool, _Q
|
||||
if _QUANTIZE and opts.device in {"CPU", "DSP"}: ret.append(RewriteStep(pm_quant, name="quantize"))
|
||||
|
||||
# split ranges
|
||||
if _RANGEIFY:
|
||||
ret.append(RewriteStep(pm_split_ranges+pm_flatten_range, ctx=lambda _: {}, name="split ranges"))
|
||||
ret.append(RewriteStep(pm_split_ranges+pm_flatten_range, ctx=lambda _: {}, name="split ranges"))
|
||||
|
||||
# symbolic (NOTE: this is a requirement for pm_simplify_ranges to be correct)
|
||||
ret.append(RewriteStep(sym+pm_flatten_range, name="initial symbolic"))
|
||||
|
||||
@@ -2,7 +2,7 @@ from __future__ import annotations
|
||||
import math, itertools
|
||||
from collections import defaultdict
|
||||
from typing import cast, Final
|
||||
from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, can_pad, GroupOp
|
||||
from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, GroupOp
|
||||
from tinygrad.device import Buffer
|
||||
from tinygrad.dtype import AddrSpace, dtypes, ImageDType
|
||||
from tinygrad.helpers import colored, BEAM, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, prod, merge_dicts, get_single_element
|
||||
@@ -188,7 +188,7 @@ class Scheduler:
|
||||
check(rng.arg[-1] is not AxisType.THREAD, "cannot pad thread")
|
||||
# ok to pad SUM if all parent ALU ops have f(0) = 0
|
||||
if (r:=self.reduceop) is not None and rng.arg[-1] in (AxisType.GROUP_REDUCE, AxisType.REDUCE):
|
||||
check(r.arg[0] is Ops.ADD and can_pad(r, {}), f"cannot pad {r}")
|
||||
check(r.arg[0] is Ops.ADD and not r.op_in_backward_slice_with_self(*GroupOp.UnsafePad), f"cannot pad {r}")
|
||||
new_sz = round_up(int(rng.vmax+1), cast(int, opt.arg))
|
||||
check(rng.vmax+1 > new_sz//4, "pad adds more than quadruple the work")
|
||||
replaced_rng = UOp.range(new_sz, *rng.arg)
|
||||
|
||||
@@ -129,7 +129,8 @@ def reduce_collapse(red:UOp):
|
||||
|
||||
def reduce_unparented(red:UOp):
|
||||
if red.arg not in {Ops.ADD, Ops.MAX, Ops.MUL}: return None
|
||||
reduce_parented, reduce_unparented = partition(red.src[1:], lambda x: x in red.src[0].backward_slice_with_self)
|
||||
assert all(x.op is Ops.RANGE for x in red.src[1:]), "some reduce srcs aren't ranges"
|
||||
reduce_parented, reduce_unparented = partition(red.src[1:], lambda x: x in red.src[0].ranges)
|
||||
if len(reduce_unparented) == 0: return None
|
||||
ret = red.replace(src=(red.src[0],)+tuple(reduce_parented)) if len(reduce_parented) or red.dtype != red.src[0].dtype else red.src[0]
|
||||
if red.arg is Ops.ADD:
|
||||
|
||||
+1
-1
@@ -183,7 +183,7 @@ class dtypes:
|
||||
uints = (uint8, uint16, uint32, uint64)
|
||||
sints = (int8, int16, int32, int64)
|
||||
ints = uints + sints
|
||||
all = floats + ints + (bool, index)
|
||||
all = floats + ints + (bool, index) # noqa: A003
|
||||
|
||||
if (env_default_float := getenv("DEFAULT_FLOAT", "")):
|
||||
dtypes.default_float = getattr(dtypes, env_default_float.lower())
|
||||
|
||||
@@ -133,7 +133,6 @@ JIT, JIT_BATCH_SIZE = ContextVar("JIT", 2 if OSX and ARCH_X86 else 1), ContextVa
|
||||
WINO, CAPTURING, TRACEMETA = ContextVar("WINO", 0), ContextVar("CAPTURING", 1), ContextVar("TRACEMETA", 1)
|
||||
USE_TC, TC_SELECT, TC_OPT, AMX = ContextVar("TC", 1), ContextVar("TC_SELECT", -1), ContextVar("TC_OPT", 0), ContextVar("AMX", 0)
|
||||
TRANSCENDENTAL, NOLOCALS = ContextVar("TRANSCENDENTAL", 1), ContextVar("NOLOCALS", 0)
|
||||
FUSE_CONV_BW = ContextVar("FUSE_CONV_BW", 0)
|
||||
SPLIT_REDUCEOP, NO_MEMORY_PLANNER, RING = ContextVar("SPLIT_REDUCEOP", 1), ContextVar("NO_MEMORY_PLANNER", 0), ContextVar("RING", 1)
|
||||
PICKLE_BUFFERS, LRU = ContextVar("PICKLE_BUFFERS", 1), ContextVar("LRU", 1)
|
||||
CACHELEVEL, IGNORE_BEAM_CACHE, DEVECTORIZE = ContextVar("CACHELEVEL", 2), ContextVar("IGNORE_BEAM_CACHE", 0), ContextVar("DEVECTORIZE", 1)
|
||||
|
||||
+11
-10
@@ -1,13 +1,13 @@
|
||||
from __future__ import annotations
|
||||
from typing import cast, ClassVar
|
||||
import os, ctypes, ctypes.util, struct, hashlib, functools, importlib, mmap, errno, array, contextlib, sys, weakref, itertools
|
||||
import os, ctypes, struct, hashlib, functools, importlib, mmap, errno, array, contextlib, sys, weakref, itertools
|
||||
assert sys.platform != 'win32'
|
||||
from dataclasses import dataclass
|
||||
from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, HWQueue, CLikeArgsState, HCQSignal, HCQProgram, FileIOInterface
|
||||
from tinygrad.runtime.support.hcq import MMIOInterface, BumpAllocator
|
||||
from tinygrad.uop.ops import sint
|
||||
from tinygrad.device import Compiled, DMAFdRef, BufferSpec, CompilerPairT
|
||||
from tinygrad.helpers import getenv, to_mv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, suppress_finalizing, lo32, hi32
|
||||
from tinygrad.helpers import getenv, to_mv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, suppress_finalizing, lo32, hi32, colored
|
||||
from tinygrad.renderer.cstyle import AMDRenderer
|
||||
from tinygrad.renderer.llvmir import AMDLLVMRenderer
|
||||
from tinygrad.runtime.autogen import kfd, hsa, pci, sqtt
|
||||
@@ -27,6 +27,9 @@ WAIT_REG_MEM_FUNCTION_GEQ = 5 # >=
|
||||
AQL_HDR = (1 << hsa.HSA_PACKET_HEADER_BARRIER) | (hsa.HSA_FENCE_SCOPE_SYSTEM << hsa.HSA_PACKET_HEADER_SCACQUIRE_FENCE_SCOPE) \
|
||||
| (hsa.HSA_FENCE_SCOPE_SYSTEM << hsa.HSA_PACKET_HEADER_SCRELEASE_FENCE_SCOPE)
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ProfileSQTTEvent(ProfileEvent): device:str; se:int; props:dict; blob:bytes; itrace:bool # noqa: E702
|
||||
|
||||
class AMDSignal(HCQSignal):
|
||||
def __init__(self, *args, **kwargs): super().__init__(*args, **{**kwargs, 'timestamp_divider': 100})
|
||||
|
||||
@@ -497,9 +500,6 @@ class AMDAllocator(HCQAllocator['AMDDevice']):
|
||||
|
||||
def _map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ProfileSQTTEvent(ProfileEvent): device:str; se:int; blob:bytes; itrace:bool # noqa: E702
|
||||
|
||||
@dataclass
|
||||
class AMDQueueDesc:
|
||||
ring: MMIOInterface
|
||||
@@ -803,7 +803,7 @@ class AMDDevice(HCQCompiled):
|
||||
# SQTT is disabled by default because of runtime overhead and big file sizes (~200mb to Tensor.full() two 4096x4096 tensors and matmul them)
|
||||
self.sqtt_enabled = PROFILE and bool(getenv("SQTT", 0))
|
||||
if self.sqtt_enabled:
|
||||
if self.arch != 'gfx1100': raise RuntimeError('SQ Thread Tracing is only supported on 7900XTX')
|
||||
if self.target[0] != 11: raise RuntimeError(f'SQ Thread Tracing is not supported on gc:{self.target}')
|
||||
if not self.is_am() and (ppfeaturemask:=int(FileIOInterface('/sys/module/amdgpu/parameters/ppfeaturemask', os.O_RDONLY).read(), 16))&0x8000:
|
||||
raise RuntimeError("SQTT can't be enabled because of hardware bug, to workaround either use AMD_IFACE=PCI or add "
|
||||
f"ppfeaturemask={(ppfeaturemask&~0x8000):#x} (current {ppfeaturemask=:#x} & ~PP_GFXOFF_MASK) to amdgpu module parameters\n"
|
||||
@@ -871,13 +871,14 @@ class AMDDevice(HCQCompiled):
|
||||
cast(AMDComputeQueue, self.hw_compute_queue_t()).sqtt_stop(len(self.sqtt_buffers), wptrs_buf) \
|
||||
.signal(self.timeline_signal, self.next_timeline()).submit(self)
|
||||
self.synchronize()
|
||||
if DEBUG>=2: print('Saving SQTT in profile...')
|
||||
if DEBUG >= 2: print(f'{self.device}: Saving SQTT in profile...')
|
||||
for i,buf0 in enumerate(self.sqtt_buffers):
|
||||
wptr = ((struct.unpack('<I', wptrs[i*4:i*4+4])[0] & 0x1FFFFFFF) - ((buf0.va_addr//32) & 0x1FFFFFFF)) * 32
|
||||
if DEBUG>=2: print(f'Se {i} blob size {wptr:#x}')
|
||||
if DEBUG >= 2: print(f'\t{self.device}: SE {i} blob size {wptr:#x}')
|
||||
assert wptr >= 0 and wptr <= buf0.size, f"{wptr} > {buf0.size}, should never happen"
|
||||
# When sqtt buffer overflows, wptr stops at the last dword
|
||||
if wptr >= buf0.size-32: print(f"WARNING: SQTT BUFFER IS FULL (SE {i})! INCREASE SQTT BUFFER SIZE WITH SQTT_BUFFER_SIZE=X (in MB)")
|
||||
if wptr >= buf0.size - 32:
|
||||
print(colored(f"{self.device}: Warning: SQTT buffer is full (SE {i})! Increase SQTT buffer with SQTT_BUFFER_SIZE=X (in MB)", "yellow"))
|
||||
self.allocator._copyout(sqtt_buf:=memoryview(bytearray(wptr)), buf0)
|
||||
Compiled.profile_events += [ProfileSQTTEvent(self.device, i, bytes(sqtt_buf), bool((self.sqtt_itrace_se_mask >> i) & 0b1))]
|
||||
Compiled.profile_events += [ProfileSQTTEvent(self.device, i, self.iface.props, bytes(sqtt_buf), bool((self.sqtt_itrace_se_mask >> i) & 0b1))]
|
||||
super()._at_profile_finalize()
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
from __future__ import annotations
|
||||
import ctypes, ctypes.util, functools
|
||||
import ctypes, functools
|
||||
from tinygrad.helpers import DEBUG, getenv, mv_address, init_c_var, init_c_struct_t, suppress_finalizing
|
||||
from tinygrad.device import Compiled, BufferSpec, LRUAllocator, CompilerPairT
|
||||
from tinygrad.renderer.cstyle import CUDARenderer
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
import subprocess, hashlib, tempfile, ctypes, ctypes.util, re, pathlib
|
||||
import subprocess, hashlib, tempfile, ctypes, re, pathlib
|
||||
from typing import Callable
|
||||
from tinygrad.helpers import to_char_p_p, colored, init_c_var, getenv
|
||||
import tinygrad.runtime.autogen.nvrtc as nvrtc
|
||||
|
||||
@@ -3,10 +3,11 @@ from typing import cast, Callable, Type, TypeVar, Generic, Any, Sequence
|
||||
import contextlib, decimal, statistics, time, ctypes, array, os, struct, traceback, collections
|
||||
try: import fcntl # windows misses that
|
||||
except ImportError: fcntl = None #type:ignore[assignment]
|
||||
from tinygrad.helpers import PROFILE, getenv, to_mv, round_up, ProfileRangeEvent
|
||||
from tinygrad.helpers import PROFILE, getenv, to_mv, ProfileRangeEvent
|
||||
from tinygrad.device import BufferSpec, Compiled, LRUAllocator, ProfileDeviceEvent, ProfileProgramEvent, CompilerPairT
|
||||
from tinygrad.uop.ops import sym_infer, sint, UOp
|
||||
from tinygrad.runtime.autogen import libc
|
||||
from tinygrad.runtime.support.memory import BumpAllocator
|
||||
|
||||
class MMIOInterface:
|
||||
def __init__(self, addr:int, nbytes:int, fmt='B'): self.mv, self.addr, self.nbytes, self.fmt = to_mv(addr, nbytes).cast(fmt), addr, nbytes, fmt
|
||||
@@ -62,15 +63,6 @@ ProgramType = TypeVar('ProgramType', bound='HCQProgram')
|
||||
ArgsStateType = TypeVar('ArgsStateType', bound='HCQArgsState')
|
||||
QueueType = TypeVar('QueueType', bound='HWQueue')
|
||||
|
||||
class BumpAllocator:
|
||||
def __init__(self, size:int, base:int=0, wrap:bool=True): self.size, self.ptr, self.base, self.wrap = size, 0, base, wrap
|
||||
def alloc(self, size:int, alignment:int=1) -> int:
|
||||
if round_up(self.ptr, alignment) + size > self.size:
|
||||
if not self.wrap: raise RuntimeError("Out of memory")
|
||||
self.ptr = 0
|
||||
self.ptr = (res:=round_up(self.ptr, alignment)) + size
|
||||
return res + self.base
|
||||
|
||||
class HWQueue(Generic[SignalType, HCQDeviceType, ProgramType, ArgsStateType]):
|
||||
"""
|
||||
A base class for hardware command queues in the HCQ (Hardware Command Queue) API.
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
import ctypes, ctypes.util, os, sys, subprocess
|
||||
import ctypes.util, os, sys, subprocess
|
||||
from tinygrad.helpers import DEBUG, OSX, getenv
|
||||
|
||||
if sys.platform == 'win32':
|
||||
|
||||
@@ -2,6 +2,15 @@ import collections, functools, dataclasses
|
||||
from typing import Any, ClassVar
|
||||
from tinygrad.helpers import round_up, getenv
|
||||
|
||||
class BumpAllocator:
|
||||
def __init__(self, size:int, base:int=0, wrap:bool=True): self.size, self.ptr, self.base, self.wrap = size, 0, base, wrap
|
||||
def alloc(self, size:int, alignment:int=1) -> int:
|
||||
if round_up(self.ptr, alignment) + size > self.size:
|
||||
if not self.wrap: raise RuntimeError("Out of memory")
|
||||
self.ptr = 0
|
||||
self.ptr = (res:=round_up(self.ptr, alignment)) + size
|
||||
return res + self.base
|
||||
|
||||
class TLSFAllocator:
|
||||
"""
|
||||
The allocator is based on the Two-Level Segregated Fit (TLSF) algorithm. The allocator maintains 2 level of buckets:
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
import ctypes, ctypes.util, os, subprocess, platform, sysconfig
|
||||
import ctypes.util, os, subprocess, platform, sysconfig
|
||||
from tinygrad.helpers import OSX
|
||||
|
||||
WEBGPU_PATH: str | None
|
||||
|
||||
@@ -1,11 +1,10 @@
|
||||
from typing import Iterator
|
||||
from typing import Iterator, Sequence
|
||||
import functools, operator, itertools
|
||||
from dataclasses import dataclass, field
|
||||
from tinygrad.dtype import dtypes, AddrSpace
|
||||
from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp
|
||||
from tinygrad.uop.symbolic import sym
|
||||
from tinygrad.helpers import argsort, all_same, Context
|
||||
from tinygrad.uop.ops import graph_rewrite, sint, AxisType
|
||||
from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType
|
||||
from tinygrad.uop.symbolic import sym, symbolic
|
||||
from tinygrad.helpers import argsort, all_same, cpu_profile, TracingKey
|
||||
|
||||
ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.ASSIGN, Ops.COPY, Ops.BUFFER, Ops.BUFFER_VIEW,
|
||||
Ops.CONST, Ops.BIND, Ops.DEVICE, Ops.MSELECT, Ops.MSTACK, Ops.DEFINE_GLOBAL,
|
||||
@@ -43,7 +42,6 @@ class BufferizeOpts:
|
||||
class IndexingContext:
|
||||
realize_map: dict[UOp, None] = field(default_factory=dict)
|
||||
range_map: dict[UOp, tuple[list[UOp], list[UOp]]] = field(default_factory=dict)
|
||||
pads_gate: dict[UOp, UOp] = field(default_factory=dict)
|
||||
|
||||
# create ranges
|
||||
range_idx: Iterator[int] = field(default_factory=itertools.count)
|
||||
@@ -56,10 +54,10 @@ def create_bufferize_and_index_based_on_ranges(ctx:IndexingContext, x:UOp):
|
||||
new_srcs = []
|
||||
for s in x.src:
|
||||
new_src = s
|
||||
if s.op in {Ops.BUFFER, Ops.MSTACK, Ops.MSELECT} or (s.op is Ops.ASSIGN and s.src[1].op is Ops.KERNEL):
|
||||
if s.op in {Ops.BUFFER, Ops.BUFFER_VIEW, Ops.MSTACK, Ops.MSELECT} or (s.op is Ops.ASSIGN and s.src[1].op is Ops.KERNEL):
|
||||
if x in ctx.range_map: new_src = new_src.index(*ctx.range_map[x][0])
|
||||
elif s in ctx.realize_map:
|
||||
new_src = UOp(Ops.BUFFERIZE, s.dtype, src=(s,)+tuple(ctx.range_map[s][1]), arg=BufferizeOpts(device=s.device), tag=s.tag)
|
||||
new_src = UOp(Ops.BUFFERIZE, s.dtype, src=(new_src,)+tuple(ctx.range_map[s][1]), arg=BufferizeOpts(device=s.device), tag=s.tag)
|
||||
if x in ctx.range_map: new_src = new_src.index(*ctx.range_map[x][0])
|
||||
new_srcs.append(new_src)
|
||||
# NOTE: do we need this?
|
||||
@@ -67,7 +65,8 @@ def create_bufferize_and_index_based_on_ranges(ctx:IndexingContext, x:UOp):
|
||||
|
||||
def convert_pad_to_where_to_keep_behavior_local(ctx:IndexingContext, x:UOp):
|
||||
if x not in ctx.range_map: return None
|
||||
ret = ctx.pads_gate[x].where(x.src[0], UOp.const(x.dtype, 0))
|
||||
valid: UOp = functools.reduce(operator.and_, [r.get_valid() for r in ctx.range_map[x][0]], UOp.const(dtypes.bool, True))
|
||||
ret = valid.where(x.src[0], UOp.const(x.dtype, 0))
|
||||
ctx.range_map[ret] = ctx.range_map[x]
|
||||
return ret
|
||||
|
||||
@@ -103,23 +102,52 @@ pm_apply_rangeify = PatternMatcher([
|
||||
(UPat((Ops.CONST, Ops.DEFINE_VAR), name="c"), lambda ctx,c: c.replace(src=()) if c in ctx.range_map else None),
|
||||
])
|
||||
|
||||
# this is the definition of the movement ops
|
||||
def apply_movement_op(x:UOp, rngs:Sequence[UOp]) -> list[UOp]:
|
||||
match x.op:
|
||||
case Ops.SHRINK: rngs = [a if ss == 0 else a+ss for a,(ss,_) in zip(rngs, x.arg)]
|
||||
case Ops.PERMUTE: rngs = [rngs[p] for p in argsort(x.arg)]
|
||||
case Ops.FLIP: rngs = [((s-1)-a) if f else a for a,s,f in zip(rngs, x.shape, x.arg)]
|
||||
case Ops.EXPAND: rngs = [a if in_sh == out_sh else a.const_like(0) for a,in_sh,out_sh in zip(rngs, x.src[0].shape, x.shape)]
|
||||
case Ops.PAD:
|
||||
# TODO: why is multiple graph_rewrites faster than one here?
|
||||
rngs = [r if (s == 0 and e == 0) else graph_rewrite(((r >= s) & (r < (sh-e))).where(r-s, UOp.invalid()), sym, name="pad")
|
||||
for r,sh,(s,e) in zip(rngs, x.shape, x.arg)]
|
||||
case Ops.RESHAPE:
|
||||
acc = 1
|
||||
axes_in:list[UOp] = []
|
||||
for s,src in list(zip(x.shape, rngs))[::-1]:
|
||||
axes_in.append(acc*src)
|
||||
acc *= s
|
||||
combined_axes = sum(axes_in, start=UOp.const(dtypes.index, 0))
|
||||
axes_out:list[UOp] = []
|
||||
for s in x.src[0].shape[::-1]:
|
||||
axes_out.append(combined_axes % s)
|
||||
combined_axes //= s
|
||||
# this simplify is doing a lot of heavy lifting. this is the replacement for the reshape view merging code
|
||||
rngs = list(graph_rewrite(UOp.sink(*axes_out[::-1]), symbolic, name="reshape").src)
|
||||
case _: raise RuntimeError(f"{x.op} is not a MovementOp")
|
||||
return rngs
|
||||
|
||||
@cpu_profile(TracingKey("run_rangeify"), "TINY")
|
||||
def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
|
||||
rctx = IndexingContext()
|
||||
|
||||
# get ops to realize
|
||||
graph_rewrite(tsink, pm_generate_realize_map, ctx=rctx.realize_map, name="Input Graph")
|
||||
|
||||
# get the traversal order
|
||||
with cpu_profile(TracingKey("reverse toposort"), "TINY"):
|
||||
tsink_reverse_toposort = tsink.reverse_toposort(consumer_map:=tsink.get_consumer_map())
|
||||
|
||||
# explicit rangeify
|
||||
ending_ranges: dict[UOp, bool] = {}
|
||||
for x in tsink.reverse_toposort(consumer_map:=tsink.get_consumer_map()):
|
||||
for x in tsink_reverse_toposort:
|
||||
if x.op in {Ops.DEVICE, Ops.UNIQUE}: continue
|
||||
ending_ranges[x] = any(ending_ranges[u] for u in consumer_map[x])
|
||||
|
||||
# if this element has weight and it's ending a range, we (force) realize it
|
||||
if ending_ranges[x] and x.op in GroupOp.Elementwise.union({Ops.REDUCE_AXIS}):
|
||||
if x.op_in_backward_slice_with_self(Ops.BUFFER, Ops.BUFFERIZE, Ops.CONTIGUOUS):
|
||||
if x.op_in_backward_slice_with_self(Ops.REDUCE_AXIS):
|
||||
rctx.realize_map[x] = None
|
||||
if ending_ranges[x] and x.op in GroupOp.Elementwise.union({Ops.REDUCE_AXIS}): rctx.realize_map[x] = None
|
||||
|
||||
# *** the ranges on the output are
|
||||
# 1. new if this op is realized
|
||||
@@ -159,7 +187,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
|
||||
if all_all_same:
|
||||
# the new valid is the OR of all the children valids
|
||||
minimum_valid = functools.reduce(operator.or_, valids, UOp.const(dtypes.bool, False))
|
||||
out_rngs.append(minimum_valid.where(local_rngs[0], UOp.invalid()).simplify())
|
||||
out_rngs.append(graph_rewrite(minimum_valid.where(local_rngs[0], UOp.invalid()), symbolic, name="minimum_valid"))
|
||||
else:
|
||||
out_rngs.append(rctx.new_range(x.shape[i]))
|
||||
|
||||
@@ -177,39 +205,9 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
|
||||
|
||||
rngs = out_rngs # rngs is the input ranges
|
||||
|
||||
# apply movement ops. this is the definition of them
|
||||
if x.op is Ops.SHRINK: rngs = [a+ss if resolve(ss != 0) else a for a,(ss,_) in zip(rngs, x.arg)]
|
||||
if x.op is Ops.PERMUTE: rngs = [rngs[p] for p in argsort(x.arg)]
|
||||
if x.op is Ops.FLIP: rngs = [((s-1)-a) if f else a for a,s,f in zip(rngs, x.shape, x.arg)]
|
||||
if x.op is Ops.EXPAND:
|
||||
rngs = [a if resolve(x==y, False) else a.const_like(0) for a,x,y in zip(rngs, x.src[0].shape, x.shape)]
|
||||
ending_ranges[x] = True
|
||||
if x.op is Ops.PAD:
|
||||
rngs = rngs[:]
|
||||
bigwhere = UOp.const(dtypes.bool, True)
|
||||
for i,(sh,(s,e)) in enumerate(zip(x.shape, x.arg)):
|
||||
if s == 0 and e == 0: continue
|
||||
where = UOp.const(dtypes.bool, True)
|
||||
if resolve(e > 0): where = where & (rngs[i] < (sh-e))
|
||||
if resolve(s > 0): where = where & (rngs[i] >= s)
|
||||
bigwhere = bigwhere & where
|
||||
with Context(TRACK_MATCH_STATS=0):
|
||||
rngs[i] = graph_rewrite(where.where(rngs[i]-s, UOp.invalid()), sym)
|
||||
# PAD is replaced with a WHERE in the big graph to inject the 0s at the right place
|
||||
rctx.pads_gate[x] = bigwhere.simplify()
|
||||
if x.op is Ops.RESHAPE:
|
||||
acc = 1
|
||||
to_sum = []
|
||||
for s,src in list(zip(x.shape, rngs))[::-1]:
|
||||
to_sum.append(acc*src)
|
||||
acc *= s
|
||||
mish = sum(to_sum, start=UOp.const(dtypes.index, 0))
|
||||
ret:list[UOp] = []
|
||||
for s in x.src[0].shape[::-1]:
|
||||
ret.append(mish % s) # NOTE: simplify will turn this to CONST
|
||||
mish //= s
|
||||
# this simplify is doing a lot of heavy lifting. this is the replacement for the view merger in RESHAPE
|
||||
rngs = list(UOp.sink(*ret[::-1]).simplify().src)
|
||||
# apply movement ops
|
||||
if x.op in GroupOp.Movement: rngs = apply_movement_op(x, rngs)
|
||||
if x.op is Ops.EXPAND: ending_ranges[x] = True
|
||||
|
||||
# REDUCE_AXIS creates ranges for the axes it is reducing
|
||||
if x.op is Ops.REDUCE_AXIS:
|
||||
|
||||
@@ -2,13 +2,12 @@ from typing import cast
|
||||
from dataclasses import dataclass, field
|
||||
from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace
|
||||
from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, _substitute, ssimplify, KernelInfo
|
||||
from tinygrad.uop.symbolic import sym, symbolic_simple
|
||||
from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, RANGEIFY, Context, flatten, dedup, unwrap, all_int, DEBUG, SPLIT_REDUCEOP
|
||||
from tinygrad.helpers import Metadata
|
||||
from tinygrad.uop.ops import track_rewrites, graph_rewrite, identity_element, sint, AxisType
|
||||
from tinygrad.uop.symbolic import symbolic_simple
|
||||
from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, flatten, dedup, unwrap, all_int, DEBUG, SPLIT_REDUCEOP, Metadata
|
||||
from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_unparented
|
||||
from tinygrad.codegen.opt import Opt
|
||||
from tinygrad.schedule.indexing import run_rangeify, BufferizeOpts, ALWAYS_CONTIGUOUS, IndexingContext
|
||||
from tinygrad.schedule.indexing import run_rangeify, BufferizeOpts, ALWAYS_CONTIGUOUS, IndexingContext, apply_movement_op
|
||||
|
||||
# creation can recurse a lot
|
||||
import sys
|
||||
@@ -100,69 +99,11 @@ earliest_rewrites = PatternMatcher([
|
||||
|
||||
# *****************
|
||||
# 3a. rangeify (movement)
|
||||
# NOTE: this can be deleted after the cleanup is refactored
|
||||
|
||||
def map_reshape(idx:UOp, r:UOp):
|
||||
acc = 1
|
||||
to_sum = []
|
||||
for s,src in list(zip(idx.shape, idx.src[1:]))[::-1]:
|
||||
to_sum.append(acc*src)
|
||||
acc *= s
|
||||
mish = sum(to_sum, start=UOp.const(dtypes.index, 0))
|
||||
ret:list[UOp] = []
|
||||
for s in r.src[0].shape[::-1]:
|
||||
ret.append(mish % s) # NOTE: simplify will turn this to CONST
|
||||
mish //= s
|
||||
tret = UOp.sink(*ret[::-1]).simplify().src
|
||||
return r.src[0].index(*tret, dtype=idx.dtype, arg=idx.arg)
|
||||
|
||||
def map_pad(idx:UOp, r:UOp):
|
||||
ret = list(idx.src[1:])
|
||||
bigwhere = UOp.const(dtypes.bool, True)
|
||||
for i,(sh,(s,e)) in enumerate(zip(r.shape, r.arg)):
|
||||
if s == 0 and e == 0: continue
|
||||
where = UOp.const(dtypes.bool, True)
|
||||
if resolve(e > 0): where = where & (ret[i] < (sh-e))
|
||||
if resolve(s > 0): where = where & (ret[i] >= s)
|
||||
bigwhere = bigwhere & where
|
||||
with Context(TRACK_MATCH_STATS=0):
|
||||
ret[i] = graph_rewrite(where.where(ret[i]-s, UOp.invalid()), sym)
|
||||
# PAD is with 0
|
||||
return bigwhere.simplify().where(r.src[0].index(*ret, dtype=idx.dtype, arg=idx.arg), UOp.const(r.dtype, 0))
|
||||
|
||||
def map_expand(r:UOp, idx:UOp):
|
||||
new_rngs = []
|
||||
ending_ranges = []
|
||||
non_ending_ranges = []
|
||||
for a,x,y in zip(idx.src[1:], r.src[0].shape, r.shape):
|
||||
axis_to_range = [u for u in a.toposort() if u.op is Ops.RANGE]
|
||||
if resolve(x==y, False):
|
||||
non_ending_ranges.extend(axis_to_range)
|
||||
new_rngs.append(a)
|
||||
else:
|
||||
ending_ranges.extend(axis_to_range)
|
||||
new_rngs.append(a.const_like(0))
|
||||
# if RANGEIFY >= 2, we are aggressive about not ending ranges
|
||||
if RANGEIFY >= 2: ending_ranges = [x.arg for x in ending_ranges if x not in non_ending_ranges]
|
||||
# if RANGEIFY=1, if it's ending at all we end it
|
||||
else: ending_ranges = [x.arg for x in ending_ranges]
|
||||
if idx.arg is not None: ending_ranges.append(idx.arg)
|
||||
return r.src[0].index(*new_rngs, arg=min(ending_ranges) if ending_ranges else None)
|
||||
|
||||
# movement op on INDEX as a PatternMatcher
|
||||
pm_mops = PatternMatcher([
|
||||
# this is like the definitions of these
|
||||
(UPat(Ops.SHRINK, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"),
|
||||
lambda r,idx: r.src[0].index(*[a+ss if resolve(ss != 0) else a for a,(ss,_) in zip(idx.src[1:], r.arg)], dtype=idx.dtype, arg=idx.arg)),
|
||||
(UPat(Ops.PERMUTE, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"),
|
||||
lambda r,idx: r.src[0].index(*[idx.src[1+p] for p in argsort(idx.src[0].arg)], dtype=idx.dtype, arg=idx.arg)),
|
||||
(UPat(Ops.FLIP, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"),
|
||||
lambda r,idx: r.src[0].index(*[((s-1)-a) if f else a for a,s,f in zip(idx.src[1:], r.shape, r.arg)], dtype=idx.dtype, arg=idx.arg)),
|
||||
# expand needs to end ranges
|
||||
(UPat(Ops.EXPAND, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"), map_expand),
|
||||
# reshape does a lot of symbolic stuff
|
||||
(UPat(Ops.RESHAPE, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"), map_reshape),
|
||||
# pad adds min and max
|
||||
(UPat(Ops.PAD, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"), map_pad),
|
||||
(UPat(GroupOp.Movement, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"),
|
||||
lambda r,idx: r.src[0].index(*apply_movement_op(r, idx.src[1:]), dtype=idx.dtype, arg=idx.arg)),
|
||||
])
|
||||
|
||||
# *****************
|
||||
@@ -183,7 +124,7 @@ def cleanup_dead_axes(b:UOp):
|
||||
# skip for symbolic. TODO: fix this
|
||||
if rng.op is Ops.RANGE and rng.src[0].op is not Ops.CONST: return None
|
||||
# CONSTs are already dead axes
|
||||
if rng.op is Ops.CONST or (rng.op is Ops.RANGE and rng not in b.src[0].backward_slice_with_self):
|
||||
if rng.op is Ops.CONST or (rng.op is Ops.RANGE and rng not in b.src[0].ranges):
|
||||
reshape.append(1)
|
||||
hit = True
|
||||
else:
|
||||
@@ -208,23 +149,29 @@ def remove_bufferize(src:UOp, buf:UOp, idx:UOp):
|
||||
# *** here is where we compute the cost ***
|
||||
# if we return None, the bufferize is kept
|
||||
|
||||
accessed_buffers = []
|
||||
accessed_buffers: list[UOp] = []
|
||||
reduces: list[UOp] = []
|
||||
def red_gate(x:UOp):
|
||||
if x.op is Ops.INDEX:
|
||||
accessed_buffers.append(x)
|
||||
return False
|
||||
if x.op is Ops.REDUCE: reduces.append(x)
|
||||
return True
|
||||
ran = src.toposort(gate=red_gate)
|
||||
src.toposort(gate=red_gate)
|
||||
del red_gate
|
||||
|
||||
# if this is generated from multiple buffers, don't remove this buffer
|
||||
if len(dedup([x.src[0] for x in accessed_buffers])) > 2: return None
|
||||
|
||||
# const reduce is okay
|
||||
# TODO: move the reduce folder to before this to prevent the need for this
|
||||
def okay_reduce(x:UOp): return all(y.op not in {Ops.BUFFER, Ops.BUFFERIZE, Ops.COPY} for y in x.backward_slice_with_self)
|
||||
|
||||
# always run this list of ops
|
||||
if any(x.op is Ops.REDUCE and not okay_reduce(x) for x in ran): return None
|
||||
# if any reduces access a buffer, don't remove this buffer
|
||||
buffer_in_reduce = False
|
||||
def buf_gate(x:UOp):
|
||||
nonlocal buffer_in_reduce
|
||||
if x.op in {Ops.BUFFER, Ops.BUFFERIZE}: buffer_in_reduce = True
|
||||
return not buffer_in_reduce
|
||||
UOp.sink(*[x.src[0] for x in reduces]).toposort(gate=buf_gate)
|
||||
del buf_gate
|
||||
if buffer_in_reduce: return None
|
||||
|
||||
# if it makes it here, the bufferize is removed
|
||||
# this is the ranges replaced
|
||||
@@ -524,9 +471,9 @@ def do_sub_recurse(s:UOp):
|
||||
return UOp(Ops.SUBSTITUTE, dtype=x.dtype, src=(x.src[0], sub_k, sub_v))
|
||||
# here we actually do the SUBSTITUTE
|
||||
if x in keys: return values[keys.index(x)]
|
||||
# we filter any keys that aren't in the backward slice. this keeps the algorithm O(output graph size)
|
||||
# NOTE: if k was x, it would trigger above, so self doesn't have to be included in backward_slice
|
||||
new_kv = {k:v for k,v in zip(keys,values) if k in x.backward_slice}
|
||||
# we filter any keys where the ranges don't overlap. this keeps the algorithm O(output graph size)
|
||||
x_ranges = x.ranges
|
||||
new_kv = {k:v for k,v in zip(keys,values) if any(r in x_ranges for r in k.ranges)}
|
||||
# if there's no SUBSTITUTEs left, we can just return x
|
||||
if len(new_kv) == 0: return x
|
||||
# then we add SUBSTITUTE to all parents
|
||||
|
||||
@@ -12,7 +12,6 @@ from tinygrad.uop.ops import UOp, Ops, graph_rewrite, Variable, sint, sint_to_uo
|
||||
def views_to_valid_uop(views: tuple[View, ...], _idxs:tuple[UOp, ...]|None=None) -> UOp:
|
||||
idx = views[-1].to_valid_uop(_idxs)
|
||||
for view in reversed(views[0:-1]):
|
||||
view = view.minify()
|
||||
idx = view.to_valid_uop([sint_to_uop(i) for i in unravel(view.shape, idx)])
|
||||
with Context(TRACK_MATCH_STATS=0):
|
||||
return graph_rewrite(idx, sym, name="indexing sym @ 1")
|
||||
@@ -54,11 +53,6 @@ class ShapeTracker:
|
||||
@property
|
||||
def size(self) -> int: return self.views[-1].size()
|
||||
|
||||
def reduce(self, axis:tuple[int, ...]) -> tuple[sint, ...]: return tuple(1 if i in axis else s for i,s in enumerate(self.shape))
|
||||
|
||||
def to_valid_uop(self, _idxs:list[UOp]|tuple[UOp, ...]|None=None) -> UOp:
|
||||
return views_to_valid_uop(self.views, tuple(_idxs) if _idxs is not None else None)
|
||||
|
||||
def vars(self) -> set[Variable]: return set().union(*[v.vars() for v in self.views])
|
||||
|
||||
@property
|
||||
@@ -68,11 +62,9 @@ class ShapeTracker:
|
||||
unbound_views, var_vals = zip(*[v.unbind() for v in self.views])
|
||||
if all(len(x) == 0 for x in var_vals): return self, {}
|
||||
return ShapeTracker(tuple(unbound_views)), merge_dicts(var_vals)
|
||||
def substitute(self, dvars:dict[UOp, UOp]): return ShapeTracker(tuple(x.substitute(dvars) for x in self.views))
|
||||
|
||||
def real_strides(self, ignore_valid=False) -> tuple[sint|None, ...]:
|
||||
with Context(TRACK_MATCH_STATS=0): return views_to_real_strides(self.views, ignore_valid)
|
||||
def unit_stride_axes(self, ignore_valid=False) -> list[int]: return [i for i,st in enumerate(self.real_strides(ignore_valid)) if st == 1]
|
||||
|
||||
def simplify(self) -> ShapeTracker:
|
||||
if len(self.views) >= 2 and (new_view := self.views[-2] + self.views[-1]) is not None:
|
||||
|
||||
+1
-31
@@ -4,7 +4,7 @@ from dataclasses import dataclass
|
||||
from typing import cast, Sequence
|
||||
from tinygrad.dtype import dtypes
|
||||
from tinygrad.uop.ops import resolve, UOp, Variable, sint, smax, smin, sint_to_uop, Ops, ssimplify
|
||||
from tinygrad.helpers import prod, all_int, argsort, flatten, ceildiv
|
||||
from tinygrad.helpers import prod, all_int, flatten, ceildiv
|
||||
|
||||
# returns the axes to create new_shape if new_shape can be created by combining axis from old_shape
|
||||
def get_contraction(old_shape:tuple[sint, ...], new_shape:tuple[sint, ...]) -> list[list[int]]|None:
|
||||
@@ -13,24 +13,6 @@ def get_contraction(old_shape:tuple[sint, ...], new_shape:tuple[sint, ...]) -> l
|
||||
except ValueError: return None
|
||||
return [list(range(st,ed)) for st,ed in zip([0]+split[:-1], split[:-1]+[len(old_shape)])]
|
||||
|
||||
def get_contraction_with_reduce(old_shape:tuple[sint, ...], new_shape:tuple[sint, ...], reduce_axis:tuple[int, ...]) -> list[list[int]]|None:
|
||||
if (contraction:=get_contraction(old_shape, new_shape)) is None: return None
|
||||
# contraction returns the 1s as right justified as possible
|
||||
# normally this contraction is good, but sometimes the reduce dim is empty. borrow from the next one, leaving one
|
||||
# this ensures there's always ones available in the reduce dimension. this is also a valid contraction
|
||||
for i in range(len(contraction)):
|
||||
if i in reduce_axis and len(contraction[i]) == 0:
|
||||
take_from = i+1
|
||||
while take_from < len(contraction) and len(contraction[take_from]) == 0:
|
||||
assert new_shape[take_from] == 1
|
||||
take_from += 1
|
||||
if take_from == len(contraction) or new_shape[take_from] != 1: return None # nothing to take
|
||||
for j in range(take_from, i, -1):
|
||||
assert len(contraction[j]) > 0
|
||||
contraction[j-1] = contraction[j][:-1]
|
||||
contraction[j] = contraction[j][-1:]
|
||||
return contraction
|
||||
|
||||
@functools.cache
|
||||
def canonicalize_strides(shape:tuple[sint, ...], strides:tuple[sint, ...]) -> tuple[sint, ...]:
|
||||
return tuple(0 if s == 1 else st for s, st in zip(shape, strides))
|
||||
@@ -244,18 +226,6 @@ class View:
|
||||
|
||||
return View.create(vm1.shape, tuple(strides), ssimplify(sum(o * s for o, s in zip(origin, vm2.strides)) + vm2.offset))
|
||||
|
||||
@functools.cache # pylint: disable=method-cache-max-size-none
|
||||
def invert(self, out_shape:tuple[sint, ...]) -> View|None:
|
||||
ret = View.create(self.shape)
|
||||
if self.mask: ret = ret.shrink(self.mask)
|
||||
ret = ret.flip(tuple(x < 0 for x in self.strides)).permute(argsort(tuple(-x if x > 0 else x for x in self.strides)))
|
||||
return ret if prod(ret.shape) == prod(out_shape) else None # don't support shrink, expand, or stride != (-1, 1)
|
||||
|
||||
@functools.cache # pylint: disable=method-cache-max-size-none
|
||||
def minify(self):
|
||||
min_shape = tuple(x[0] for x in merge_dims(self.shape, self.strides, self.mask))
|
||||
return nv if (nv := self.reshape(min_shape)) else self
|
||||
|
||||
def __unsafe_resize(self, arg: tuple[tuple[sint, sint], ...], mask=None) -> View:
|
||||
offset = sum([s * x[0] for s, x in zip(self.strides,arg)])
|
||||
if self.mask:
|
||||
|
||||
+2
-2
@@ -6,7 +6,7 @@ from typing import Callable, ClassVar, Sequence, cast, get_args, Literal, Suppor
|
||||
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
|
||||
from tinygrad.helpers import argfix, make_tuple, flatten, prod, all_int, round_up, merge_dicts, argsort, getenv, all_same, fully_flatten, dedup
|
||||
from tinygrad.helpers import IMAGE, WINO, Metadata, TRACEMETA, ceildiv, fetch, polyN, unwrap, DEBUG, is_numpy_ndarray, RANGEIFY, FUSE_ATTENTION
|
||||
from tinygrad.helpers import IMAGE, WINO, Metadata, TRACEMETA, ceildiv, fetch, polyN, unwrap, DEBUG, is_numpy_ndarray, FUSE_ATTENTION
|
||||
from tinygrad.helpers import suppress_finalizing
|
||||
from tinygrad.gradient import compute_gradient
|
||||
from tinygrad.uop.ops import smax, smin, resolve, UOp, Ops, sint, MathTrait, identity_element, all_metadata, _index_to_concrete_int, sint_to_uop, \
|
||||
@@ -227,7 +227,7 @@ class Tensor(MathTrait):
|
||||
# verify Tensors match the spec
|
||||
if __debug__: type_verify(list(big_sink.toposort()), tensor_uop_spec)
|
||||
|
||||
if RANGEIFY and any(isinstance(x._device, tuple) for x in big_sink.toposort()):
|
||||
if any(isinstance(x._device, tuple) for x in big_sink.toposort()):
|
||||
_apply_map_to_tensors(get_multi_map(big_sink), "Apply Multi Map")
|
||||
big_sink = UOp.sink(*flatten([x.uop.src if x.uop.op is Ops.MULTI else [x.uop] for x in (self,)+lst]))
|
||||
|
||||
|
||||
@@ -108,6 +108,4 @@ class GroupOp:
|
||||
# do not preserve f(0) = 0
|
||||
UnsafePad = {Ops.RECIP, Ops.LOG2, Ops.EXP2, Ops.IDIV, Ops.POW}
|
||||
|
||||
Meta = {Ops.COPY, Ops.BUFFER_VIEW}
|
||||
|
||||
All = set(Ops)
|
||||
|
||||
+25
-46
@@ -23,9 +23,6 @@ range_start = {Ops.BUFFERIZE: 1, Ops.REDUCE: 1, Ops.STORE: 2, Ops.WMMA: 3}
|
||||
# https://en.wikipedia.org/wiki/Identity_element
|
||||
def identity_element(op:Ops, dt:DType) -> ConstType: return dtypes.as_const({Ops.ADD:0, Ops.MUL:1, Ops.MAX:dtypes.min(dt)}[op], dt)
|
||||
|
||||
def can_pad(root:UOp, edges:dict[UOp, None]) -> bool:
|
||||
return all(u.op not in GroupOp.UnsafePad for u in root.toposort(gate=lambda x:x not in edges))
|
||||
|
||||
# With True as the default, this matches the old symbolic behavior
|
||||
def resolve(x:UOp|bool, default:bool=True):
|
||||
if isinstance(x, bool): return x
|
||||
@@ -124,11 +121,12 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
|
||||
|
||||
def f(self, op, **kwargs): return UOp(op, dtype=kwargs.pop("dtype", self.dtype), src=(self,), **kwargs)
|
||||
|
||||
@recursive_property
|
||||
@functools.cached_property
|
||||
def backward_slice(self:UOp) -> dict[UOp, None]:
|
||||
ret = {s:None for s in self.src}
|
||||
for s in self.src: ret.update(s.backward_slice)
|
||||
return ret
|
||||
res: dict[UOp, None] = self.toposort()
|
||||
res.pop(self)
|
||||
return res
|
||||
|
||||
@property
|
||||
def backward_slice_with_self(self:UOp) -> dict[UOp, None]: return {self:None, **self.backward_slice}
|
||||
def op_in_backward_slice_with_self(self, *ops:Ops): return any(x.op in ops for x in self.backward_slice_with_self)
|
||||
@@ -207,24 +205,22 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
|
||||
sz = self.ptrdtype.size
|
||||
return ShapeTracker.from_shape((sz,)) if sz > 0 else None
|
||||
|
||||
# CONTIGUOUS with RANGE
|
||||
# TODO: how are these not RANGE?
|
||||
if self.op is Ops.CONTIGUOUS and len(self.src) > 1 and all(x.op is Ops.RANGE for x in self.src[1:]):
|
||||
return ShapeTracker.from_shape((tuple([int(x.vmax+1) for x in self.src[1:]])+self.src[0].shape))
|
||||
|
||||
# hack for PTX, CASTing the ptr loses the shape
|
||||
if self.op is Ops.CAST and self.src[0].op is Ops.DEFINE_GLOBAL: return None
|
||||
|
||||
# otherwise we get the shape from sources
|
||||
if not (src_sts := [x.st for x in self.src if x.st is not None]): return None
|
||||
assert all_same([x.shape for x in src_sts]), f"UOp sources must have the same shape {self} {[x.shape for x in src_sts]}"
|
||||
shape = src_sts[0].shape
|
||||
# shape changing ops
|
||||
match self.op:
|
||||
case Ops.MULTI: shape = tuple(self.src[0].shape[a]*len(self.device) if a == self.axis else s for a,s in enumerate(self.src[0].shape))
|
||||
case Ops.MULTI: shape = tuple(s*len(self.device) if a == self.axis else s for a,s in enumerate(shape))
|
||||
case Ops.BITCAST:
|
||||
shape = src_sts[0].shape
|
||||
if self.dtype.itemsize != (input_sz:=self.src[0].dtype.itemsize): shape = shape[:-1]+((shape[-1]*input_sz) // self.dtype.itemsize,)
|
||||
case Ops.REDUCE_AXIS | Ops.WMMA: shape = src_sts[0].reduce(self.axis_arg)
|
||||
case _: shape = src_sts[0].shape
|
||||
if (output_sz:=self.dtype.itemsize) != (input_sz:=self.src[0].dtype.itemsize): shape = shape[:-1]+((shape[-1]*input_sz) // output_sz,)
|
||||
case Ops.REDUCE_AXIS | Ops.WMMA:
|
||||
axis_arg = self.arg[1] if self.op is Ops.REDUCE_AXIS else self.arg[7]
|
||||
assert isinstance(axis_arg, tuple) and all(isinstance(x, int) for x in axis_arg), f"invalid type for axis: {axis_arg}"
|
||||
shape = tuple(1 if i in axis_arg else s for i,s in enumerate(shape))
|
||||
return ShapeTracker.from_shape(shape)
|
||||
|
||||
@property
|
||||
@@ -235,7 +231,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
|
||||
def size(self) -> int: return self.arg[0] if self.op is Ops.BUFFER_VIEW else self.arg if self.op is Ops.BUFFER else unwrap(self.st).size
|
||||
|
||||
# determine what ranges this is in
|
||||
@functools.cached_property
|
||||
@recursive_property
|
||||
def _ranges(self) -> dict[UOp, None]:
|
||||
ret: dict[UOp, None] = {}
|
||||
if self.op in range_start.keys():
|
||||
@@ -286,16 +282,6 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
|
||||
|
||||
# *** uop syntactic sugar ***
|
||||
|
||||
@property
|
||||
def st_arg(self) -> ShapeTracker:
|
||||
assert self.op in GroupOp.Buffer, f"st_arg called on {self.op}"
|
||||
return unwrap(self.st)
|
||||
@property
|
||||
def axis_arg(self) -> tuple[int, ...]:
|
||||
assert self.op in {Ops.REDUCE_AXIS, Ops.WMMA}, f"axis_arg called on {self.op}"
|
||||
ret = self.arg[1] if self.op is Ops.REDUCE_AXIS else self.arg[7]
|
||||
assert isinstance(ret, tuple) and all(isinstance(x, int) for x in ret), f"axis_arg trying to return {ret}"
|
||||
return ret
|
||||
def sink(*srcs:UOp|None, **kwargs): # pylint: disable=no-self-argument
|
||||
return UOp(Ops.SINK, dtypes.void, tuple([x for x in srcs if x is not None]), **kwargs)
|
||||
def detach(self): return UOp(Ops.DETACH, self.dtype, (self,))
|
||||
@@ -501,7 +487,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
|
||||
def buffer(self) -> Buffer|MultiBuffer:
|
||||
from tinygrad.device import Buffer, MultiBuffer
|
||||
if self is not self.base:
|
||||
assert unwrap(self.st).contiguous, "VIEW only works here if it's contiguous"
|
||||
assert self.op is Ops.RESHAPE, f"can only be RESHAPE {self}"
|
||||
return self.src[0].buffer
|
||||
if self.op is Ops.MSELECT:
|
||||
ret = self.src[0].buffer
|
||||
@@ -759,8 +745,8 @@ class UPat(MathTrait):
|
||||
def var(name:str|None=None, dtype:DType|tuple[DType, ...]|None=None): return UPat(dtype=dtype, name=name)
|
||||
@staticmethod
|
||||
@functools.cache
|
||||
def cvar(name:str|None=None, dtype:DType|tuple[DType, ...]|None=None, vec=True):
|
||||
return UPat((Ops.CONST,Ops.VCONST) if vec else Ops.CONST, dtype, name=name)
|
||||
def cvar(name:str|None=None, dtype:DType|tuple[DType, ...]|None=None, vec=True, arg=None):
|
||||
return UPat((Ops.CONST,Ops.VCONST) if vec else Ops.CONST, dtype, name=name, arg=arg)
|
||||
@staticmethod
|
||||
def const(dtype:DType|tuple[DType, ...]|None, b:ConstType|InvalidType): return UPat(Ops.CONST, dtype=dtype, arg=b)
|
||||
|
||||
@@ -787,13 +773,6 @@ class UPat(MathTrait):
|
||||
asrc = (self,)+src
|
||||
return UPat(op, dtypes.bool if op in {Ops.CMPLT, Ops.CMPNE} else asrc[-1].dtype, list(asrc) if op in GroupOp.Commutative else asrc)
|
||||
|
||||
def __repr__(self):
|
||||
def rep(x):
|
||||
form = "UPat(%s, %s, name=%s, dtype=%s, allow_any_len=%s, src=%s)"
|
||||
return form % (None if x.op is None else ('(%s)'%', '.join(map(str, x.op))), x.arg, repr(x.name),
|
||||
set(x.dtype) if x.dtype else None, not x.strict_length, "[%s]" if x.src and len(x.src)>1 else ("(%s)" if x.src else "%s"))
|
||||
return pretty_print(self, rep, srcfn=lambda x:None if x.src is None else [next(x.src[0])] if isinstance(x.src[0], itertools.repeat) else x.src[0])
|
||||
|
||||
def match(self:UPat, uop:UOp, store:dict[str, UOp]) -> list[dict[str, UOp]]:
|
||||
if (self.op is not None and uop.op not in self.op) or \
|
||||
(self.name is not None and store.setdefault(self.name, uop) is not uop) or \
|
||||
@@ -879,11 +858,11 @@ match_stats:dict[UPat, list[int|float]] = dict()
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TrackedGraphRewrite:
|
||||
loc:tuple[str, int] # location that called graph_rewrite
|
||||
sink:int # the sink input to graph_rewrite
|
||||
matches:list[tuple[int, int, tuple]] # before/after UOp, UPat location
|
||||
name:str|None # optional name of the rewrite
|
||||
depth:int # depth if it's a subrewrite
|
||||
loc:tuple[str, int] # location that called graph_rewrite
|
||||
sink:int # the sink input to graph_rewrite
|
||||
matches:list[tuple[int, int, tuple, float]] # before/after UOp, UPat location and time
|
||||
name:str|None # optional name of the rewrite
|
||||
depth:int # depth if it's a subrewrite
|
||||
bottom_up:bool
|
||||
|
||||
tracked_keys:list[TracingKey] = []
|
||||
@@ -955,14 +934,14 @@ class TrackedPatternMatcher(PatternMatcher):
|
||||
try: ret = match(uop, ctx)
|
||||
except Exception:
|
||||
if TRACK_MATCH_STATS >= 2 and active_rewrites:
|
||||
active_rewrites[-1].matches.append((track_uop(uop), track_uop(UOp(Ops.REWRITE_ERROR, src=uop.src, arg=str(sys.exc_info()[1]))), p.location))
|
||||
active_rewrites[-1].matches.append((track_uop(uop), track_uop(UOp(Ops.REWRITE_ERROR,src=uop.src,arg=str(sys.exc_info()[1]))),p.location,0))
|
||||
raise
|
||||
if ret is not None and ret is not uop:
|
||||
match_stats[p][0] += 1
|
||||
match_stats[p][3] += (et:=time.perf_counter()-st)
|
||||
if TRACK_MATCH_STATS >= 3: print(f"{et*1e6:7.2f} us -- ", printable(p.location))
|
||||
if TRACK_MATCH_STATS >= 2 and isinstance(ret, UOp) and active_rewrites:
|
||||
active_rewrites[-1].matches.append((track_uop(uop), track_uop(ret), p.location))
|
||||
active_rewrites[-1].matches.append((track_uop(uop), track_uop(ret), p.location, et))
|
||||
return ret
|
||||
match_stats[p][2] += time.perf_counter()-st
|
||||
return None
|
||||
@@ -992,7 +971,7 @@ if TRACK_MATCH_STATS or PROFILE:
|
||||
if not int(os.getenv("VIZ", "0")) and not int(os.getenv("PROFILE", "0")) and not int(os.getenv("SQTT", "0")):
|
||||
args = ['--kernels', getenv("VIZ_DATA", "")] if getenv("VIZ_DATA", "") else []
|
||||
args += ['--profile', getenv("PROFILE_DATA", "")] if getenv("PROFILE_DATA", "") else []
|
||||
os.execv(sys.executable, [sys.executable] + [os.path.join(os.path.dirname(__file__), "../", "viz", "serve.py")] + args)
|
||||
os.execv(sys.executable, [sys.executable] + [pathlib.Path(__file__).resolve().parent.parent / "viz" / "serve.py"] + args)
|
||||
|
||||
# *** simple graph rewrite engine ***
|
||||
|
||||
|
||||
@@ -51,8 +51,6 @@ symbolic_simple = propagate_invalid + PatternMatcher([
|
||||
(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("x") / UPat.var("x"), lambda x: x.const_like(1)), # x/x -> 1
|
||||
((UPat.var("x") * UPat.var("x2")) / UPat.var("x2"), lambda x,x2: x), # (x*x2)/x2 -> 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)
|
||||
# 4 variations of (x%c)+(x//c)*c = x TODO: add sorting to remove some variations
|
||||
(UPat.var("x")%UPat.cvar("c")+(UPat.var("x")//UPat.cvar("c"))*UPat.cvar("c"), lambda x,c: x), # (x%c)+(x//c)*c = x
|
||||
@@ -76,10 +74,6 @@ symbolic_simple = propagate_invalid + PatternMatcher([
|
||||
(UPat.var("x") % UPat.var("x"), lambda x: x.const_like(0)), # x%x -> 0
|
||||
(UPat.var("x", dtype=dtypes.ints+(dtypes.bool, dtypes.index)) != UPat.var("x"),
|
||||
lambda x: x.const_like(False).cast(dtypes.bool.vec(x.dtype.count))), # x != x -> False (only ints)
|
||||
# x*0 -> 0 or 0*x -> 0
|
||||
# if x is nan or inf it should render the nan value.
|
||||
# NOTE: this can be wrong for loaded NaN
|
||||
(UPat.var("x") * 0, lambda x: x.const_like(float("nan") if isinstance(x.arg, float) and (math.isnan(x.arg) or math.isinf(x.arg)) else 0)),
|
||||
# ** constant folding **
|
||||
# TODO: add const folding for Ops.THREEFRY
|
||||
(UPat(GroupOp.Unary, src=(UPat((Ops.VCONST, Ops.CONST)),), name="a"), lambda a: a.const_like(exec_alu(a.op, a.dtype, [a.src[0].arg], False))),
|
||||
@@ -91,6 +85,17 @@ symbolic_simple = propagate_invalid + PatternMatcher([
|
||||
(UPat.var('x', dtype=dtypes.bool) * UPat.var('y', dtype=dtypes.bool), lambda x,y: x&y),
|
||||
(UPat.var('x', dtype=dtypes.bool) + UPat.var('y', dtype=dtypes.bool), lambda x,y: x|y),
|
||||
(UPat.var('x', dtype=dtypes.bool).maximum(UPat.var('y', dtype=dtypes.bool)), lambda x,y: x|y),
|
||||
# *** div rules ***
|
||||
(UPat.cvar('x', arg=0) / 0, lambda x: x.const_like(float('nan'))), # 0/0 -> nan
|
||||
((UPat.var("x") * 0) / 0, lambda x: x.const_like(float('nan'))), # (x*0)/0 -> nan
|
||||
# can be wrong if x or x2 is 0
|
||||
(UPat.var("x") / UPat.var("x"), lambda x: x.const_like(1)), # x/x -> 1
|
||||
((UPat.var("x") * UPat.var("x2")) / UPat.var("x2"), lambda x,x2: x), # (x*x2)/x2 -> x
|
||||
# x*0 -> 0 or 0*x -> 0
|
||||
# if x is nan or inf it should render the nan value.
|
||||
# NOTE: this can be wrong for loaded NaN
|
||||
(UPat.var("x") * 0, lambda x: x.const_like(float("nan") if x.op is Ops.CONST
|
||||
and isinstance(x.arg, float) and (math.isnan(x.arg) or math.isinf(x.arg)) else 0)),
|
||||
# *** cast/bitcast ***
|
||||
(UPat(Ops.CAST, name="root", src=(UPat.cvar("c"),)), lambda root, c: root.const_like(c.arg)),
|
||||
(UPat((Ops.CAST, Ops.BITCAST), name="root"), lambda root: root.src[0] if root.dtype == root.src[0].dtype else None),
|
||||
|
||||
+41
-17
@@ -137,17 +137,18 @@ const formatUnit = (d, unit="") => d3.format(".3~s")(d)+unit;
|
||||
|
||||
const colorScheme = {TINY:["#1b5745", "#354f52", "#354f52", "#1d2e62", "#63b0cd"],
|
||||
DEFAULT:["#2b2e39", "#2c2f3a", "#31343f", "#323544", "#2d303a", "#2e313c", "#343746", "#353847", "#3c4050", "#404459", "#444862", "#4a4e65"],
|
||||
BUFFER:["#3A57B7","#5066C1","#6277CD","#7488D8","#8A9BE3","#A3B4F2"],
|
||||
BUFFER:["#342483", "#3E2E94", "#4938A4", "#5442B4", "#5E4CC2", "#674FCA"],
|
||||
CATEGORICAL:["#ff8080", "#F4A261", "#C8F9D4", "#8D99AE", "#F4A261", "#ffffa2", "#ffffc0", "#87CEEB"],}
|
||||
const cycleColors = (lst, i) => lst[i%lst.length];
|
||||
|
||||
const rescaleTrack = (source, tid, k) => {
|
||||
for (const e of source.shapes) {
|
||||
for (let i=0; i<e.y0.length; i++) {
|
||||
e.y0[i] = e.y0[i]*k;
|
||||
e.y1[i] = e.y1[i]*k;
|
||||
for (const shapes of source.views)
|
||||
for (const e of shapes) {
|
||||
for (let i=0; i<e.y0.length; i++) {
|
||||
e.y0[i] = e.y0[i]*k;
|
||||
e.y1[i] = e.y1[i]*k;
|
||||
}
|
||||
}
|
||||
}
|
||||
const change = (source.height*k)-source.height;
|
||||
const div = document.getElementById(tid);
|
||||
div.style.height = rect(div).height+change+"px";
|
||||
@@ -221,14 +222,15 @@ async function renderProfiler() {
|
||||
const base = colorMap.get(colorKey), s = Math.min(Math.pow(1/0.7, depth), 240 / Math.max(base.r, base.g, base.b));
|
||||
const fillColor = d3.rgb(base.r*s, base.g*s, base.b*s).toString();
|
||||
const label = parseColors(e.name).map(({ color, st }) => ({ color, st, width:ctx.measureText(st).width }));
|
||||
if (e.ref != null) ref = {ctx:e.ref, step:0};
|
||||
let shapeRef = e.ref;
|
||||
if (shapeRef != null) { ref = {ctx:e.ref, step:0}; shapeRef = ref; }
|
||||
else if (ref != null) {
|
||||
const start = ref.step>0 ? ref.step+1 : 0;
|
||||
const stepIdx = ctxs[ref.ctx+1].steps.findIndex((s, i) => i >= start && s.name == e.name);
|
||||
ref = stepIdx === -1 ? null : {ctx:ref.ctx, step:stepIdx};
|
||||
if (stepIdx !== -1) { ref.step = stepIdx; shapeRef = ref; }
|
||||
}
|
||||
const htmlLabel = label.map(({color, st}) => `<span style="color:${color}">${st}</span>`).join('');
|
||||
const arg = { tooltipText:htmlLabel+"\n"+formatTime(e.dur)+(e.info != null ? "\n"+e.info : ""), ...ref };
|
||||
const arg = { tooltipText:htmlLabel+"\n"+formatTime(e.dur)+(e.info != null ? "\n"+e.info : ""), ...shapeRef };
|
||||
// offset y by depth
|
||||
shapes.push({x:e.st, y:levelHeight*depth, width:e.dur, height:levelHeight, arg, label, fillColor });
|
||||
}
|
||||
@@ -237,7 +239,7 @@ async function renderProfiler() {
|
||||
const peak = u64();
|
||||
let x = 0, y = 0;
|
||||
const buf_shapes = new Map(), temp = new Map();
|
||||
const timestamps = [];
|
||||
const timestamps = [], valueMap = new Map();
|
||||
for (let j=0; j<eventsLen; j++) {
|
||||
const alloc = u8(), ts = u32(), key = u32();
|
||||
if (alloc) {
|
||||
@@ -245,11 +247,11 @@ async function renderProfiler() {
|
||||
const shape = {x:[x], y:[y], dtype, sz, nbytes, key};
|
||||
buf_shapes.set(key, shape); temp.set(key, shape);
|
||||
timestamps.push(ts);
|
||||
x += 1; y += nbytes;
|
||||
x += 1; y += nbytes; valueMap.set(ts, y);
|
||||
} else {
|
||||
const free = buf_shapes.get(key);
|
||||
timestamps.push(ts);
|
||||
x += 1; y -= free.nbytes;
|
||||
x += 1; y -= free.nbytes; valueMap.set(ts, y);
|
||||
free.x.push(x);
|
||||
free.y.push(free.y.at(-1));
|
||||
temp.delete(key);
|
||||
@@ -273,14 +275,34 @@ async function renderProfiler() {
|
||||
const arg = {tooltipText:`${dtype} len:${formatUnit(sz)}\n${formatUnit(nbytes, "B")}\nnum:${num}\nalive for ${formatTime(dur)}`};
|
||||
shapes.push({ x, y0:y.map(yscale), y1:y.map(y0 => yscale(y0+nbytes)), arg, fillColor:cycleColors(colorScheme.BUFFER, shapes.length) });
|
||||
}
|
||||
data.tracks.set(k, { shapes, visible, offsetY, height, peak, scaleFactor:maxheight*4/height });
|
||||
// generic polygon merger
|
||||
const base0 = yscale(0);
|
||||
const allX = Array.from(new Set(shapes.flatMap(s => s.x))).sort((a,b)=>a-b);
|
||||
const idxs = new Map(allX.map((x,i) => [x, i]));
|
||||
const maxY = new Map(allX.map(x => [x, base0]));
|
||||
// for every [a,b) update the max y at x
|
||||
for (const sh of shapes) {
|
||||
for (let i=0; i<sh.x.length-1; i++) {
|
||||
const startIdx = idxs.get(sh.x[i]), endIdx = idxs.get(sh.x[i+1]);
|
||||
const shapeY = sh.y1[i];
|
||||
for (let k=startIdx; k<endIdx; k++) {
|
||||
const x = allX[k]; maxY.set(x, Math.min(maxY.get(x), shapeY));
|
||||
}
|
||||
}
|
||||
}
|
||||
const sum = {x:[], y0:[], y1:[], fillColor:"#2B1B72"};
|
||||
for (let i=0; i<allX.length-1; i++) {
|
||||
sum.x.push(allX[i], allX[i+1]);
|
||||
const y = maxY.get(allX[i]); sum.y1.push(y, y); sum.y0.push(base0, base0);
|
||||
}
|
||||
data.tracks.set(k, { shapes:[sum], visible, offsetY, height, peak, scaleFactor:maxheight*4/height, views:[[sum], shapes], valueMap });
|
||||
div.style("height", height+padding+"px").style("cursor", "pointer").on("click", (e) => {
|
||||
const newFocus = e.currentTarget.id === focusedDevice ? null : e.currentTarget.id;
|
||||
let offset = 0;
|
||||
for (const [tid, track] of data.tracks) {
|
||||
track.offsetY += offset;
|
||||
if (tid === newFocus) offset += rescaleTrack(track, tid, track.scaleFactor);
|
||||
else if (tid === focusedDevice) offset += rescaleTrack(track, tid, 1/track.scaleFactor);
|
||||
if (tid === newFocus) { track.shapes = track.views[1]; offset += rescaleTrack(track, tid, track.scaleFactor); }
|
||||
else if (tid === focusedDevice) { track.shapes = track.views[0]; offset += rescaleTrack(track, tid, 1/track.scaleFactor); }
|
||||
}
|
||||
data.axes.y = newFocus != null ? { domain:[0, (t=data.tracks.get(newFocus)).peak], range:[t.offsetY+t.height, t.offsetY], fmt:"B" } : null;
|
||||
focusedDevice = newFocus;
|
||||
@@ -301,7 +323,7 @@ async function renderProfiler() {
|
||||
const st = visibleX[0], et = visibleX[1];
|
||||
xscale.domain(visibleX);
|
||||
// draw shapes
|
||||
for (const [_, { offsetY, shapes, visible }] of data.tracks) {
|
||||
for (const [_, { offsetY, shapes, visible, valueMap }] of data.tracks) {
|
||||
visible.length = 0;
|
||||
for (const e of shapes) {
|
||||
// generic polygon
|
||||
@@ -312,7 +334,9 @@ async function renderProfiler() {
|
||||
ctx.moveTo(x[0], offsetY+e.y0[0]);
|
||||
for (let i=1; i<x.length; i++) {
|
||||
ctx.lineTo(x[i], offsetY+e.y0[i]);
|
||||
visible.push({ x0:x[i-1], x1:x[i], y0:offsetY+e.y1[i-1], y1:offsetY+e.y0[i], arg:e.arg });
|
||||
let arg = e.arg;
|
||||
if (arg == null && valueMap != null) arg = {tooltipText: `Total: ${formatUnit(valueMap.get(e.x[i-1]), 'B')}`}
|
||||
visible.push({ x0:x[i-1], x1:x[i], y0:offsetY+e.y1[i-1], y1:offsetY+e.y0[i], arg });
|
||||
}
|
||||
for (let i=x.length-1; i>=0; i--) ctx.lineTo(x[i], offsetY+e.y1[i]);
|
||||
ctx.closePath();
|
||||
|
||||
@@ -97,12 +97,13 @@ def _reconstruct(a:int, i:int):
|
||||
def get_details(ctx:TrackedGraphRewrite, i:int=0) -> Generator[GraphRewriteDetails, None, None]:
|
||||
yield {"graph":uop_to_json(next_sink:=_reconstruct(ctx.sink, i)), "uop":str(next_sink), "changed_nodes":None, "diff":None, "upat":None}
|
||||
replaces: dict[UOp, UOp] = {}
|
||||
for u0_num,u1_num,upat_loc in tqdm(ctx.matches):
|
||||
for u0_num,u1_num,upat_loc,dur in tqdm(ctx.matches):
|
||||
replaces[u0:=_reconstruct(u0_num, i)] = u1 = _reconstruct(u1_num, i)
|
||||
try: new_sink = next_sink.substitute(replaces)
|
||||
except RuntimeError as e: new_sink = UOp(Ops.NOOP, arg=str(e))
|
||||
match_repr = f"# {dur*1e6:.2f} us\n"+printable(upat_loc)
|
||||
yield {"graph":(sink_json:=uop_to_json(new_sink)), "uop":str(new_sink), "changed_nodes":[id(x) for x in u1.toposort() if id(x) in sink_json],
|
||||
"diff":list(difflib.unified_diff(str(u0).splitlines(), str(u1).splitlines())), "upat":(upat_loc, printable(upat_loc))}
|
||||
"diff":list(difflib.unified_diff(str(u0).splitlines(),str(u1).splitlines())), "upat":(upat_loc, match_repr)}
|
||||
if not ctx.bottom_up: next_sink = new_sink
|
||||
|
||||
# encoder helpers
|
||||
|
||||
Reference in New Issue
Block a user