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

Merge branch 'master' into no_merge_views

Browse Source
This commit is contained in:
George Hotz
2025-08-07 14:42:07 -07:00
committed by GitHub
parent b8791e962c 82be8abfd2
commit 9ed409d0f8
90 changed files with 226 additions and 188 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/benchmark.yml
+7 -6
View File
@@ -325,7 +325,7 @@ jobs:
run: BENCHMARK_LOG=resnet_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
- name: Run 10 MLPerf Bert training steps (6 gpu)
# TODO: remove BERT_LAYERS once scheduler is fast
run: BENCHMARK_LOG=bert_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 FUSE_ARANGE=1 FUSE_ARANGE_UINT=0 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
run: BENCHMARK_LOG=bert_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
- uses: actions/upload-artifact@v4
with:
name: Speed (NVIDIA Training)
@@ -576,7 +576,7 @@ jobs:
run: BENCHMARK_LOG=resnet_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
- name: Run 10 MLPerf Bert training steps (6 gpu)
# TODO: remove BERT_LAYERS once scheduler is fast
run: BENCHMARK_LOG=bert_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 FUSE_ARANGE=1 FUSE_ARANGE_UINT=0 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
run: BENCHMARK_LOG=bert_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
- uses: actions/upload-artifact@v4
with:
name: Speed (AMD MLPerf)
@@ -692,11 +692,12 @@ jobs:
run: AMD=1 TESTFILE=/raid/downloads/llama3-8b-sfr/model-00001-of-00004.safetensors python3 test/external/external_benchmark_disk_raw.py
- name: Run full CIFAR training w 1 GPU
run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
run: BENCHMARK_LOG=resnet_10steps AMD=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee am_train_resnet_one_gpu.txt
# TODO: enable
# - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
# run: BENCHMARK_LOG=resnet_10steps AMD=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee am_train_resnet_one_gpu.txt
- name: Run 10 MLPerf Bert training steps (1 gpu)
# TODO: remove BERT_LAYERS once scheduler is fast
run: BENCHMARK_LOG=bert_10steps AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 FUSE_ARANGE=1 FUSE_ARANGE_UINT=0 MODEL=bert python3 examples/mlperf/model_train.py | tee am_train_bert_one_gpu.txt
run: BENCHMARK_LOG=bert_10steps AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee am_train_bert_one_gpu.txt
- uses: actions/upload-artifact@v4
with:
name: Speed (AM Driver)
@@ -754,7 +755,7 @@ jobs:
run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee nv_train_resnet_one_gpu.txt
- name: Run 10 MLPerf Bert training steps (1 gpu)
# TODO: remove BERT_LAYERS once scheduler is fast
run: BENCHMARK_LOG=bert_10steps NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 FUSE_ARANGE=1 FUSE_ARANGE_UINT=0 MODEL=bert python3 examples/mlperf/model_train.py | tee nv_train_bert_one_gpu.txt
run: BENCHMARK_LOG=bert_10steps NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee nv_train_bert_one_gpu.txt
- uses: actions/upload-artifact@v4
with:
name: Speed (NV Driver)
.github/workflows/test.yml
+1 -1
View File
@@ -624,7 +624,7 @@ jobs:
- name: Test LLVM=1 DEVECTORIZE=0 for model
run: PYTHONPATH="." LLVM=1 DEVECTORIZE=0 python3 test/models/test_efficientnet.py
- name: Test CPU=1 DEVECTORIZE=0
run: CPU=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"
run: CPU=1 DEVECTORIZE=0 FUSE_ARANGE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"
testwebgpu:
name: Linux (WebGPU)
docs/developer/layout.md
+2 -2
View File
@@ -18,11 +18,11 @@ Group UOps into kernels.
---
## tinygrad/opt
## tinygrad/codegen/opt
Transforms the ast into an optimized ast. This is where BEAM search and heuristics live.
::: tinygrad.opt.get_optimized_ast
::: tinygrad.codegen.opt.get_optimized_ast
options:
members: false
show_labels: false
examples/mlperf/dataloader.py
+41 -1
View File
@@ -564,10 +564,13 @@ class GPTDataset:
# check for cache
cache_hash = hashlib.sha256(f"{samples}:{seqlen}:{seed}:{shuffle}".encode()).hexdigest()
cache_path = base_path.with_name(f"{base_path.name}.{cache_hash}.index_cache")
print(f"try loading GPTDataset from {cache_path}...")
if cache_path.exists():
print("cache found, loading...")
with open(cache_path, "rb") as f:
self.doc_idx, self.sample_idx, self.shuffle_idx = pickle.load(f)
else:
print("cache not found, building index...")
self.doc_idx = self._build_doc_idx()
self.sample_idx = self._build_sample_idx()
self.shuffle_idx = self._build_shuffle_idx()
@@ -692,10 +695,47 @@ class BlendedGPTDataset:
self.datasets = [GPTDataset(path, samples_per_blend[i], seqlen, seed + i, shuffle) for i,path in enumerate(paths)]
# check for cache
cache_hash = hashlib.sha256(f"{samples}:{seqlen}:{seed}:{shuffle}".encode()).hexdigest()
cache_path = paths[0].with_name(f"{paths[0].name}.{cache_hash}.blend_cache")
print(f"try loading BlendedGPTDataset from {cache_path}...")
if cache_path.exists():
print("cache found, loading...")
with open(cache_path, "rb") as f:
self.dataset_idx, self.dataset_sample_idx = pickle.load(f)
else:
print("cache not found, building index...")
self.dataset_idx, self.dataset_sample_idx = self._build_blend_idx()
# save cache
with open(cache_path, "wb") as f:
pickle.dump((self.dataset_idx, self.dataset_sample_idx), f)
def get(self, idx:int):
tokens = self.datasets[0][idx]
tokens = self.datasets[self.dataset_idx[idx]][self.dataset_sample_idx[idx]]
return tokens
def _build_blend_idx(self):
dataset_idx = np.zeros(self.samples, dtype=np.int16)
dataset_sample_idx = np.zeros(self.samples, dtype=np.int64)
unspent_datasets = set(range(len(self.datasets)))
dataset_sample_counts = [0] * len(self.datasets)
for i in tqdm(range(self.samples)):
error_argmax, error_max = 0, 0.0
for di in unspent_datasets:
error = self.weights[di] * max(i, 1) - dataset_sample_counts[di]
if error > error_max:
error_max = error
error_argmax = di
dataset_idx[i] = error_argmax
dataset_sample_idx[i] = dataset_sample_counts[error_argmax]
dataset_sample_counts[error_argmax] += 1
return dataset_idx, dataset_sample_idx
def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True):
if val:
dataset = BlendedGPTDataset([
examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/bert/implementations/tinybox_1xMI300X/dev_beam.sh
+2
View File
@@ -4,6 +4,8 @@ export PYTHONPATH="." AMD=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" GPUS=1 BS=128 EVAL_BS=128
export IGNORE_OOB=1
export BEAM=3 BEAM_UOPS_MAX=4000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1
# export BEAM_LOG_SURPASS_MAX=1
examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/dev_beam.sh
+2
View File
@@ -5,6 +5,8 @@ export MODEL="bert"
export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
export IGNORE_OOB=1
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
export BASEDIR="/raid/datasets/wiki"
examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/dev_run.sh
+2
View File
@@ -8,6 +8,8 @@ export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
export TRAIN_STEPS=3900
export IGNORE_OOB=1
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
export BASEDIR="/raid/datasets/wiki"
examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
+2
View File
@@ -11,6 +11,8 @@ export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
export TRAIN_STEPS=3900
export IGNORE_OOB=1
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
export BASEDIR="/raid/datasets/wiki"
examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/bert/implementations/tinybox_green/dev_beam.sh
+2 -2
View File
@@ -2,9 +2,9 @@
export PYTHONPATH="." NV=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
export IGNORE_OOB=1
export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1
examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/bert/implementations/tinybox_green/dev_run.sh
+2 -2
View File
@@ -2,9 +2,9 @@
export PYTHONPATH="." NV=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
export IGNORE_OOB=1
export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1
examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
+2 -2
View File
@@ -5,9 +5,9 @@ set -o pipefail # Make pipeline fail if any command fails
export PYTHONPATH="." NV=1
export MODEL="bert"
export SUBMISSION_PLATFORM="tinybox_green"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
export IGNORE_OOB=1
export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1
examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/bert/implementations/tinybox_red/dev_beam.sh
+2 -2
View File
@@ -2,9 +2,9 @@
export PYTHONPATH="." AMD=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
export IGNORE_OOB=1
export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1
examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/bert/implementations/tinybox_red/dev_run.sh
+2 -2
View File
@@ -2,9 +2,9 @@
export PYTHONPATH="." AMD=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
export IGNORE_OOB=1
export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1
examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
+2 -2
View File
@@ -5,9 +5,9 @@ set -o pipefail # Make pipeline fail if any command fails
export PYTHONPATH="." AMD=1
export MODEL="bert"
export SUBMISSION_PLATFORM="tinybox_red"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
export IGNORE_OOB=1
export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1
extra/assembly/assembly.py
+1 -1
View File
@@ -1,5 +1,5 @@
from typing import Tuple, List, NamedTuple, Any, Dict, Optional, Union, DefaultDict, cast
from tinygrad.opt.kernel import Ops, MemOp, UOp
from tinygrad.codegen.opt.kernel import Ops, MemOp, UOp
from tinygrad.uop.ops import BinaryOps, UnaryOps
from tinygrad.dtype import DType, dtypes
from tinygrad.helpers import DEBUG
extra/assembly/assembly_arm64.py
+1 -1
View File
@@ -3,7 +3,7 @@ from platform import system
from typing import Tuple, Dict, List, Optional
from tinygrad import dtypes
from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
from tinygrad.opt.kernel import Ops, UOp
from tinygrad.codegen.opt.kernel import Ops, UOp
from tinygrad.helpers import CI
from tinygrad.codegen.assembly import uops_to_asmstyle, AssemblyLanguage
extra/assembly/assembly_ptx.py
+1 -1
View File
@@ -1,7 +1,7 @@
from typing import List
import struct
from tinygrad.codegen.assembly import uops_to_asmstyle, AssemblyLanguage
from tinygrad.opt.kernel import Ops, UOp
from tinygrad.codegen.opt.kernel import Ops, UOp
from tinygrad import dtypes
from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
from tinygrad.runtime.ops_cuda import arch
extra/assembly/assembly_rdna.py
+1 -1
View File
@@ -2,7 +2,7 @@ import yaml
from typing import Tuple, Set, Dict
from tinygrad import dtypes
from tinygrad.codegen.assembly import AssemblyCodegen, Register
from tinygrad.opt.kernel import Ops
from tinygrad.codegen.opt.kernel import Ops
from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
from tinygrad.runtime.ops_gpu import ROCM_LLVM_PATH
extra/backends/triton.py
+1 -1
View File
@@ -2,7 +2,7 @@ from typing import Dict, List, Final, Callable, DefaultDict
from collections import defaultdict
from tinygrad.uop.ops import UnaryOps, BinaryOps, TernaryOps, Op
from tinygrad.helpers import DType, PtrDType, dtypes, ImageDType, DEBUG, getenv
from tinygrad.opt.kernel import UOp, Ops
from tinygrad.codegen.opt.kernel import UOp, Ops
from triton.compiler import compile as triton_compile
import linecache
import math
extra/gemm/amd_uop_matmul.py
+1 -1
View File
@@ -6,7 +6,7 @@ from tinygrad.schedule.kernelize import merge_views, view_left
from tinygrad.helpers import getenv, colored, prod, unwrap
from tinygrad.shape.shapetracker import ShapeTracker, View
from tinygrad.shape.view import strides_for_shape
from tinygrad.opt.kernel import axis_colors
from tinygrad.codegen.opt.kernel import axis_colors
def to_colored(full_shape, axis_types): return '_'.join([colored(str(s), axis_colors[at]) for s,at in zip(full_shape, axis_types)])
extra/gemm/max_matmul.py
+2 -2
View File
@@ -5,9 +5,9 @@ from typing import Optional, List, Tuple, cast, Dict, Final, DefaultDict, Self
from tinygrad.engine.realize import get_program
# for copied uops
from tinygrad.opt.kernel import Kernel, KernelOptError
from tinygrad.codegen.opt.kernel import Kernel, KernelOptError
from tinygrad.uop.ops import UOp, Ops, BinaryOps, UnaryOps, TernaryOps, KernelInfo
from tinygrad.opt.search import Opt, OptOps
from tinygrad.codegen.opt.search import Opt, OptOps
from tinygrad import Device, dtypes, Tensor
from tinygrad.dtype import PtrDType, DType, DTYPES_DICT
from tinygrad.shape.shapetracker import ShapeTracker
extra/gemm/simple_matmul.py
+1 -1
View File
@@ -2,7 +2,7 @@ import numpy as np
from tinygrad import dtypes, Tensor
from tinygrad.helpers import getenv, get_single_element
from tinygrad.dtype import _to_np_dtype
from tinygrad.opt.kernel import OptOps
from tinygrad.codegen.opt.kernel import OptOps
from tinygrad.engine.realize import lower_schedule
dtype_in = dtypes.half if getenv("HALF") else dtypes.bfloat16 if getenv("BFLOAT16") else dtypes.float
extra/gemm/tinygrad_nv_matmul.py
+1 -1
View File
@@ -1,6 +1,6 @@
from tinygrad import Tensor, dtypes, Device
from tinygrad.helpers import getenv, DEBUG
from tinygrad.opt.kernel import Kernel, Opt, OptOps
from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps
from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
from dataclasses import replace
extra/gemm/tvm_gemm.py
+1 -1
View File
@@ -37,7 +37,7 @@ B = Tensor.rand(K, N, device="CPU")
C = (A.reshape(M, 1, K) * B.permute(1,0).reshape(1, N, K)).sum(axis=2)
sched = C.schedule()
from tinygrad.opt.kernel import Kernel
from tinygrad.codegen.opt.kernel import Kernel
from tinygrad.device import CompilerOptions
lin = Kernel(sched[-1].ast, CompilerOptions(has_local=False, supports_float4=False))
lin.to_program()
extra/hcqfuzz/tests/bert.py
-1
View File
@@ -8,7 +8,6 @@ bert_train_params = {
"BS": 96,
"EVAL_BS": 96,
"FUSE_ARANGE": 1,
"FUSE_ARANGE_UINT": 0,
"BASEDIR": "/raid/datasets/wiki",
}
extra/mcts_search.py
+2 -2
View File
@@ -4,9 +4,9 @@ import numpy as np
np.set_printoptions(suppress=True)
import math, functools, time, random, statistics
from tinygrad.helpers import DEBUG, getenv, CACHELEVEL, diskcache_get, diskcache_put, colored, Profiling
from tinygrad.opt.kernel import Kernel
from tinygrad.codegen.opt.kernel import Kernel
from tinygrad.device import Buffer, Device, CompileError
from tinygrad.opt.search import _ensure_buffer_alloc, get_kernel_actions, _time_program
from tinygrad.codegen.opt.search import _ensure_buffer_alloc, get_kernel_actions, _time_program
from tinygrad.engine.realize import get_program
class MCTSNode:
extra/optimization/extract_policynet.py
+3 -3
View File
@@ -5,9 +5,9 @@ from tinygrad.nn import Linear
from tinygrad.tensor import Tensor
from tinygrad.nn.optim import Adam
from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
from tinygrad.opt.search import actions
from tinygrad.codegen.opt.search import actions
from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats, assert_same_lin
from tinygrad.opt.kernel import Kernel
from tinygrad.codegen.opt.kernel import Kernel
from tinygrad.helpers import getenv
# stuff needed to unpack a kernel
@@ -17,7 +17,7 @@ from tinygrad.shape.shapetracker import ShapeTracker
from tinygrad.shape.view import View
from tinygrad.uop.ops import Variable
inf, nan = float('inf'), float('nan')
from tinygrad.opt.kernel import Opt, OptOps
from tinygrad.codegen.opt.kernel import Opt, OptOps
INNER = 256
class PolicyNet:
extra/optimization/extract_sa_pairs.py
+3 -3
View File
@@ -10,11 +10,11 @@ from tinygrad.shape.shapetracker import ShapeTracker
from tinygrad.shape.view import View
from tinygrad.uop.ops import Variable
inf, nan = float('inf'), float('nan')
from tinygrad.opt.kernel import Opt, OptOps
from tinygrad.codegen.opt.kernel import Opt, OptOps
# more stuff
from tinygrad.opt.kernel import Kernel
from tinygrad.opt.search import actions
from tinygrad.codegen.opt.kernel import Kernel
from tinygrad.codegen.opt.search import actions
from extra.optimization.helpers import lin_to_feats
from extra.optimization.pretrain_valuenet import ValueNet
from tinygrad.nn.optim import Adam
extra/optimization/get_action_space.py
+3 -3
View File
@@ -1,8 +1,8 @@
import random
from extra.optimization.helpers import load_worlds, ast_str_to_lin
from tinygrad.opt.search import actions
from tinygrad.opt.kernel import Kernel
from tinygrad.opt.heuristic import hand_coded_optimizations
from tinygrad.codegen.opt.search import actions
from tinygrad.codegen.opt.kernel import Kernel
from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
from tinygrad.helpers import tqdm
tactions = set()
extra/optimization/helpers.py
+3 -3
View File
@@ -1,6 +1,6 @@
# stuff needed to unpack a kernel
from tinygrad import Variable
from tinygrad.opt.kernel import Opt, OptOps
from tinygrad.codegen.opt.kernel import Opt, OptOps
from tinygrad.uop.ops import UOp, Ops, KernelInfo
from tinygrad.dtype import dtypes, PtrDType
from tinygrad.shape.shapetracker import ShapeTracker
@@ -11,7 +11,7 @@ inf, nan = float('inf'), float('nan')
UOps = Ops
# kernel unpacker
from tinygrad.opt.kernel import Kernel
from tinygrad.codegen.opt.kernel import Kernel
def ast_str_to_ast(ast_str:str) -> UOp: return eval(ast_str)
def ast_str_to_lin(ast_str:str, opts=None): return Kernel(ast_str_to_ast(ast_str), opts=opts)
def kern_str_to_lin(kern_str:str, opts=None):
@@ -103,7 +103,7 @@ def lin_to_feats(lin:Kernel, use_sts=True):
return ret
from tinygrad.device import Device, Buffer
from tinygrad.opt.search import _ensure_buffer_alloc, _time_program
from tinygrad.codegen.opt.search import _ensure_buffer_alloc, _time_program
from tinygrad.helpers import to_function_name, CACHELEVEL, diskcache_get, diskcache_put
def time_linearizer(lin:Kernel, rawbufs:list[Buffer], allow_test_size=True, max_global_size=65536, cnt=3, disable_cache=False, clear_l2=False) -> float: # noqa: E501
extra/optimization/pretrain_valuenet.py
+2 -2
View File
@@ -1,4 +1,4 @@
from tinygrad.opt.kernel import Kernel
from tinygrad.codegen.opt.kernel import Kernel
from tqdm import tqdm, trange
import math
import random
@@ -14,7 +14,7 @@ from tinygrad.shape.shapetracker import ShapeTracker
from tinygrad.shape.view import View
from tinygrad.uop.ops import Variable
inf, nan = float('inf'), float('nan')
from tinygrad.opt.kernel import Opt, OptOps
from tinygrad.codegen.opt.kernel import Opt, OptOps
from extra.optimization.helpers import lin_to_feats, MAX_DIMS
extra/optimization/rl.py
+1 -1
View File
@@ -3,7 +3,7 @@ import numpy as np
import math, random
from tinygrad.tensor import Tensor
from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
from tinygrad.opt.search import actions, bufs_from_lin, get_kernel_actions
from tinygrad.codegen.opt.search import actions, bufs_from_lin, get_kernel_actions
from tinygrad.nn.optim import Adam
from extra.optimization.extract_policynet import PolicyNet
from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats, time_linearizer
extra/optimization/run_qnet.py
+2 -2
View File
@@ -1,6 +1,6 @@
from typing import List, Tuple
from tinygrad.opt.kernel import Kernel
from tinygrad.opt.search import get_kernel_actions, actions
from tinygrad.codegen.opt.kernel import Kernel
from tinygrad.codegen.opt.search import get_kernel_actions, actions
_net = None
def beam_q_estimate(beam:List[Tuple[Kernel, float]]) -> List[Tuple[Kernel, float]]:
extra/optimization/search.py
+2 -2
View File
@@ -4,8 +4,8 @@ from extra.optimization.helpers import ast_str_to_lin, time_linearizer
from tinygrad import dtypes
from tinygrad.helpers import BEAM, getenv
from tinygrad.device import Device, Compiled
from tinygrad.opt.kernel import Kernel
from tinygrad.opt.search import beam_search, bufs_from_lin
from tinygrad.codegen.opt.kernel import Kernel
from tinygrad.codegen.opt.search import beam_search, bufs_from_lin
if __name__ == '__main__':
extra/optimization/test_net.py
+2 -2
View File
@@ -6,8 +6,8 @@ from copy import deepcopy
from tinygrad.helpers import getenv, colored
from tinygrad.tensor import Tensor
from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
from tinygrad.opt.search import bufs_from_lin, actions, get_kernel_actions
from tinygrad.opt.heuristic import hand_coded_optimizations
from tinygrad.codegen.opt.search import bufs_from_lin, actions, get_kernel_actions
from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats, time_linearizer
from extra.optimization.extract_policynet import PolicyNet
from extra.optimization.pretrain_valuenet import ValueNet
extra/optimization/test_time_linearizer.py
+1 -1
View File
@@ -1,5 +1,5 @@
from extra.optimization.helpers import load_worlds, ast_str_to_lin, time_linearizer
from tinygrad.opt.search import bufs_from_lin, get_kernel_actions
from tinygrad.codegen.opt.search import bufs_from_lin, get_kernel_actions
if __name__ == "__main__":
ast_strs = load_worlds()
extra/replay_pkl.py
+2 -2
View File
@@ -6,8 +6,8 @@ from tinygrad.helpers import getenv, BEAM
from tinygrad.engine.jit import TinyJit
from tinygrad.engine.realize import CompiledRunner, ExecItem, ScheduleItem, lower_schedule_item, get_program
from tinygrad.renderer import ProgramSpec
from tinygrad.opt.kernel import Kernel, Opt, OptOps
from tinygrad.opt.heuristic import hand_coded_optimizations
from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps
from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
import numpy as np
def move_jit_captured_to_dev(captured, device="DSP"):
setup.py
+1 -1
View File
@@ -26,7 +26,7 @@ setup(name='tinygrad',
long_description_content_type='text/markdown',
packages = ['tinygrad', 'tinygrad.runtime.autogen', 'tinygrad.runtime.autogen.am', 'tinygrad.codegen', 'tinygrad.nn',
'tinygrad.renderer', 'tinygrad.engine', 'tinygrad.viz', 'tinygrad.runtime', 'tinygrad.runtime.support', 'tinygrad.schedule',
'tinygrad.runtime.support.am', 'tinygrad.runtime.graph', 'tinygrad.shape', 'tinygrad.uop', 'tinygrad.opt',
'tinygrad.runtime.support.am', 'tinygrad.runtime.graph', 'tinygrad.shape', 'tinygrad.uop', 'tinygrad.codegen.opt',
'tinygrad.runtime.support.nv', 'tinygrad.apps'],
package_data = {'tinygrad': ['py.typed'], 'tinygrad.viz': ['index.html', 'assets/**/*', 'js/*']},
classifiers=[
test/external/external_benchmark_hcopt.py
Vendored
+2 -2
View File
@@ -1,7 +1,7 @@
import random
from tinygrad.helpers import getenv
from tinygrad.opt.search import beam_search, bufs_from_lin
from tinygrad.opt.heuristic import hand_coded_optimizations
from tinygrad.codegen.opt.search import beam_search, bufs_from_lin
from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
from extra.optimization.helpers import load_worlds, ast_str_to_lin, time_linearizer
def optimize_kernel(k):
test/external/external_benchmark_sdxl_softmax.py
Vendored
+1 -1
View File
@@ -8,7 +8,7 @@ if __name__ == "__main__":
GlobalCounters.reset()
t.softmax(-1, dtype="half", _single_kernel=True).realize()
from tinygrad.opt.kernel import Kernel, Opt, OptOps
from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps
from tinygrad.helpers import get_single_element
GlobalCounters.reset()
si = get_single_element(t.softmax(-1, dtype="half", _single_kernel=True).schedule())
test/external/external_debug_metal_sd_conv.py
Vendored
+2 -2
View File
@@ -1,8 +1,8 @@
# ruff: noqa: E501
from tinygrad.opt.kernel import Kernel, Opt, OptOps
from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps
from tinygrad.dtype import dtypes
from tinygrad.engine.realize import CompiledRunner, get_program
from tinygrad.opt.search import bufs_from_lin
from tinygrad.codegen.opt.search import bufs_from_lin
from tinygrad.uop.ops import UOp, Ops
from tinygrad.shape.shapetracker import ShapeTracker
from tinygrad.shape.view import View
test/external/external_test_hcq_fuzz_failures.py
Vendored
+2 -2
View File
@@ -4,10 +4,10 @@ os.environ["VALIDATE_HCQ"]="1"
import unittest, random
import numpy as np
from tinygrad.opt.kernel import Kernel, KernelOptError
from tinygrad.codegen.opt.kernel import Kernel, KernelOptError
from tinygrad.device import is_dtype_supported
from tinygrad.uop.ops import UOp, Ops
from tinygrad.opt.search import Opt, OptOps
from tinygrad.codegen.opt.search import Opt, OptOps
from tinygrad import Device, dtypes, Tensor
from test.external.fuzz_linearizer import compare_linearizer, compare_states, get_fuzz_rawbuf_like
test/external/external_test_hip_compile.py
Vendored
+1 -1
View File
@@ -3,7 +3,7 @@ from tinygrad.runtime.support.hip_comgr import compile_hip
from tinygrad import Tensor
from tinygrad.device import Device
from tinygrad.engine.schedule import create_schedule
from tinygrad.opt.kernel import Kernel
from tinygrad.codegen.opt.kernel import Kernel
class TestHIPCompileSpeed(unittest.TestCase):
@unittest.skipIf(Device.DEFAULT != "HIP", "only run on HIP")
test/external/external_test_nv.py
Vendored
+2 -2
View File
@@ -2,11 +2,11 @@ import unittest, struct, array, ctypes
from tinygrad import Device, dtypes, Tensor
from tinygrad.helpers import to_mv
from tinygrad.runtime.ops_nv import NVDevice, HWQueue
from tinygrad.opt.search import Opt, OptOps
from tinygrad.codegen.opt.search import Opt, OptOps
from tinygrad.engine.realize import get_runner, CompiledRunner, get_program
from test.external.fuzz_linearizer import get_fuzz_rawbufs
from tinygrad.opt.kernel import Kernel
from tinygrad.codegen.opt.kernel import Kernel
from tinygrad.uop.ops import LazyOp, Ops, ReduceOps, BufferOps, MemBuffer
from tinygrad.shape.shapetracker import ShapeTracker
from tinygrad.shape.view import View
test/external/external_test_train_gpt2.py
Vendored
+2 -2
View File
@@ -2,11 +2,11 @@
import unittest
from tinygrad.uop.ops import UOp, Ops
from tinygrad.opt.search import Opt, OptOps
from tinygrad.codegen.opt.search import Opt, OptOps
from tinygrad.dtype import dtypes
from tinygrad.shape.shapetracker import ShapeTracker
from tinygrad.shape.view import View
from tinygrad.opt.kernel import Kernel
from tinygrad.codegen.opt.kernel import Kernel
from test.external.fuzz_linearizer import run_linearizer
test/external/fuzz_linearizer.py
Vendored
+3 -3
View File
@@ -21,9 +21,9 @@ if os.getenv("VALIDATE_HCQ", 0) != 0:
from tinygrad import Tensor, Device, dtypes
from tinygrad.tensor import _to_np_dtype
from tinygrad.opt.kernel import Kernel
from tinygrad.opt.kernel import Opt, OptOps
from tinygrad.opt.search import get_kernel_actions, bufs_from_lin
from tinygrad.codegen.opt.kernel import Kernel
from tinygrad.codegen.opt.kernel import Opt, OptOps
from tinygrad.codegen.opt.search import get_kernel_actions, bufs_from_lin
from tinygrad.engine.realize import CompiledRunner
from tinygrad.helpers import getenv, from_mv, prod, colored, Context, DEBUG, Timing
from tinygrad.uop.ops import UOp, Ops
test/external/speed_beam_v_hcopt.py
Vendored
+2 -2
View File
@@ -1,7 +1,7 @@
from tinygrad import Device
from tinygrad.helpers import getenv, DEBUG, BEAM
from tinygrad.opt.search import beam_search, bufs_from_lin
from tinygrad.opt.heuristic import hand_coded_optimizations
from tinygrad.codegen.opt.search import beam_search, bufs_from_lin
from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
from extra.optimization.helpers import load_worlds, ast_str_to_lin, time_linearizer
if __name__ == "__main__":
test/external/speed_compare_amd_am.py
Vendored
+2 -2
View File
@@ -2,8 +2,8 @@ from tinygrad import Device, dtypes
from tinygrad.helpers import getenv, colorize_float, DEBUG
from extra.optimization.helpers import load_worlds, ast_str_to_lin
from test.external.fuzz_linearizer import get_fuzz_rawbufs
from tinygrad.opt.heuristic import hand_coded_optimizations
from tinygrad.opt.search import bufs_from_lin
from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
from tinygrad.codegen.opt.search import bufs_from_lin
from tinygrad.engine.realize import CompiledRunner, get_program
from tinygrad.tensor import _to_np_dtype
from tinygrad.runtime.ops_amd import AMDDevice
test/external/speed_compare_cuda_nv.py
Vendored
+2 -2
View File
@@ -2,8 +2,8 @@ from tinygrad import Device, dtypes
from tinygrad.helpers import getenv, colorize_float
from extra.optimization.helpers import load_worlds, ast_str_to_lin
from test.external.fuzz_linearizer import get_fuzz_rawbufs
from tinygrad.opt.heuristic import hand_coded_optimizations
from tinygrad.opt.search import bufs_from_lin
from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
from tinygrad.codegen.opt.search import bufs_from_lin
from tinygrad.engine.realize import CompiledRunner, get_program
from tinygrad.tensor import _to_np_dtype
import numpy as np
test/external/speed_compare_cuda_ptx.py
Vendored
+2 -2
View File
@@ -1,10 +1,10 @@
import itertools
from tinygrad import Device
from tinygrad.engine.realize import CompiledRunner, get_program
from tinygrad.opt.heuristic import hand_coded_optimizations
from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
from tinygrad.helpers import getenv, colorize_float
from extra.optimization.helpers import load_worlds, ast_str_to_lin
from tinygrad.opt.search import bufs_from_lin
from tinygrad.codegen.opt.search import bufs_from_lin
from tinygrad.runtime.ops_cuda import PTXCompiler, PTXRenderer, CUDACompiler
if __name__ == "__main__":
test/external/verify_kernel.py
Vendored
+1 -1
View File
@@ -3,7 +3,7 @@ from collections import defaultdict
from extra.optimization.helpers import kern_str_to_lin, time_linearizer
from test.external.fuzz_linearizer import compare_linearizer
from tinygrad.helpers import colored
from tinygrad.opt.kernel import Kernel
from tinygrad.codegen.opt.kernel import Kernel
# Use this with the LOGKERNS options to verify that all executed kernels are valid and evaluate to the same ground truth results
test/test_arange.py
+2 -2
View File
@@ -3,9 +3,9 @@ import numpy as np
from tinygrad import Tensor, GlobalCounters, dtypes, nn, Device, Variable
from tinygrad.helpers import CI, Context, getenv
from tinygrad.engine.realize import run_schedule
from tinygrad.opt.kernel import Opt, OptOps, Kernel, KernelOptError
from tinygrad.codegen.opt.kernel import Opt, OptOps, Kernel, KernelOptError
from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
from tinygrad.opt.search import get_kernel_actions
from tinygrad.codegen.opt.search import get_kernel_actions
from tinygrad.uop.ops import Ops
from tinygrad.codegen import apply_rewrites, rewrites_for_views
test/test_const_folding.py
+3 -4
View File
@@ -139,10 +139,9 @@ class TestBitcastConstFolding(unittest.TestCase):
class TestIndexingConstFolding(unittest.TestCase):
def test_scalar_index(self):
t = Tensor.arange(16).float().reshape(1,1,4,4).realize()
# TODO: fold these
_check_ast_count(2, t[:,:,Tensor(1),:])
_check_ast_count(2, t[:,:,Tensor(1)+2,:])
_check_ast_count(2, t[:,:,Tensor(1),Tensor(0)])
_check_ast_count(1, t[:,:,Tensor(1),:])
_check_ast_count(1, t[:,:,Tensor(1)+2,:])
_check_ast_count(1, t[:,:,Tensor(1),Tensor(0)])
@unittest.expectedFailure
def test_const_tensor_index(self):
test/test_hcq.py
+2 -2
View File
@@ -6,7 +6,7 @@ from tinygrad.runtime.support.hcq import HCQCompiled, HCQBuffer
from tinygrad.runtime.autogen import libc
from tinygrad.runtime.support.system import PCIIfaceBase
from tinygrad.engine.realize import get_runner, CompiledRunner, get_program
from tinygrad.opt.kernel import Opt, OptOps
from tinygrad.codegen.opt.kernel import Opt, OptOps
from tinygrad import Variable
MOCKGPU = getenv("MOCKGPU")
@@ -336,7 +336,7 @@ class TestHCQ(unittest.TestCase):
et = float(sig_en.timestamp - sig_st.timestamp)
print(f"exec kernel time: {et:.2f} us")
assert 0.1 <= et <= (15000 if MOCKGPU or Device.DEFAULT in {"CPU", "LLVM"} else 100)
assert 0.1 <= et <= (100000 if MOCKGPU or Device.DEFAULT in {"CPU", "LLVM"} else 100)
def test_speed_copy_bandwidth(self):
if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
test/test_linearizer.py
+3 -3
View File
@@ -2,7 +2,7 @@ import numpy as np
import unittest
from dataclasses import replace
from tinygrad.opt.kernel import Opt, OptOps, KernelOptError, Kernel, AxisType
from tinygrad.codegen.opt.kernel import Opt, OptOps, KernelOptError, Kernel, AxisType
from tinygrad.codegen.gpudims import get_grouped_dims
from tinygrad.uop.ops import UOp, Ops, GroupOp, KernelInfo
from tinygrad.device import Device, Buffer, is_dtype_supported
@@ -10,7 +10,7 @@ from tinygrad.shape.shapetracker import ShapeTracker
from tinygrad.shape.view import View
from tinygrad.tensor import Tensor, _to_np_dtype
from tinygrad.engine.realize import run_schedule, lower_schedule, CompiledRunner, get_program
from tinygrad.opt.heuristic import hand_coded_optimizations
from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
from tinygrad.helpers import prod, Context, getenv, CI, flatten, dedup, AMX, AMD_LLVM
from tinygrad.dtype import DType, dtypes, AddrSpace
from tinygrad.codegen import apply_rewrites, rewrites_for_views
@@ -410,7 +410,7 @@ class TestLinearizer(unittest.TestCase):
np.testing.assert_allclose(result, golden_result, atol=0.1, rtol=0.2)
# check that get_kernel_actions produces all 9 options
from tinygrad.opt.search import get_kernel_actions
from tinygrad.codegen.opt.search import get_kernel_actions
tc_actions = [k for i, k in get_kernel_actions(Kernel(realized_ast), False).items() if k.applied_opts[0].op == OptOps.TC]
available_tc = len([x for x in Device[Device.DEFAULT].renderer.tensor_cores if x.dtype_in == tc.dtype_in and x.dtype_out == tc.dtype_out])
test/test_linearizer_dumb.py
+2 -2
View File
@@ -8,8 +8,8 @@ from tinygrad.device import is_dtype_supported
from tinygrad.uop.ops import UOp, Ops
from tinygrad.helpers import getenv
from tinygrad.shape.shapetracker import ShapeTracker, View
from tinygrad.opt.search import Opt, OptOps
from tinygrad.opt.kernel import Kernel
from tinygrad.codegen.opt.search import Opt, OptOps
from tinygrad.codegen.opt.kernel import Kernel
from tinygrad.engine.realize import get_program
class TestLinearizerDumb(unittest.TestCase):
test/test_linearizer_overflows.py
+2 -2
View File
@@ -1,8 +1,8 @@
# ruff: noqa: E501
import unittest
from tinygrad import dtypes
from tinygrad.opt.kernel import Kernel
from tinygrad.opt.search import Opt, OptOps, bufs_from_lin
from tinygrad.codegen.opt.kernel import Kernel
from tinygrad.codegen.opt.search import Opt, OptOps, bufs_from_lin
from extra.optimization.helpers import time_linearizer
# stuff needed to unpack a kernel
test/test_nn.py
+2 -2
View File
@@ -401,7 +401,7 @@ class TestNN(unittest.TestCase):
torch_z = torch_layer(torch_x)
np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=1e-8, rtol=1e-8)
def test_embedding_one_kernel(self, ops=41410, kcount=3):
def test_embedding_one_kernel(self, ops=612000, kcount=2):
GlobalCounters.reset()
layer = Embedding(20, 30)
layer.weight = Tensor.zeros_like(layer.weight).contiguous()
@@ -409,7 +409,7 @@ class TestNN(unittest.TestCase):
[12, 19, 8, 1]])
result = layer(a)
schedule = result.schedule()
self.assertEqual(kcount, len([item for item in schedule if item.ast.op is Ops.SINK]), "first run realizes weight and embedding")
self.assertEqual(len([item for item in schedule if item.ast.op is Ops.SINK]), kcount, "first run realizes weight and embedding")
run_schedule(schedule)
b = Tensor([[1, 2, 3],
test/test_ops.py
+3 -2
View File
@@ -957,8 +957,9 @@ class TestOps(unittest.TestCase):
helper_test_op([(45,65)], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6)
helper_test_op([(45,65)], lambda t: torch.nn.functional.softplus(t, beta=3), lambda t: Tensor.softplus(t, beta=3), grad_atol=1e-6)
helper_test_op([(45,65)], lambda t: torch.nn.functional.softplus(t, beta=1/3), lambda t: Tensor.softplus(t, beta=1/3), grad_atol=1e-6)
# # TODO: support threshold and enable this
# helper_test_op([(45,65)], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6, low=300, high=400)
helper_test_op([(45,65)], lambda t: torch.nn.functional.softplus(t, beta=3, threshold=0.5),
lambda t: Tensor.softplus(t, beta=3, threshold=0.5), grad_atol=1e-6)
helper_test_op([(45,65)], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6, low=300, high=400)
helper_test_op([(45,65)], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6, low=-400, high=-300)
helper_test_op([()], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6)
test/test_opt_gemm.py
+1 -1
View File
@@ -2,7 +2,7 @@ import numpy as np
import unittest
from tinygrad import Tensor
from tinygrad.helpers import get_single_element
from tinygrad.opt.kernel import Opt, OptOps
from tinygrad.codegen.opt.kernel import Opt, OptOps
from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
class TestOptGemm(unittest.TestCase):
test/test_quantize_onnx.py
+2 -2
View File
@@ -4,9 +4,9 @@ import unittest
from dataclasses import replace
from tinygrad import Tensor, Context, Device, dtypes
from tinygrad.uop.ops import Ops, UOp # noqa: F401 # pylint: disable=unused-import
from tinygrad.opt.kernel import Kernel, Opt, OptOps
from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps
from tinygrad.engine.realize import CompiledRunner, ExecItem, lower_schedule_item, get_program
from tinygrad.opt.search import bufs_from_lin
from tinygrad.codegen.opt.search import bufs_from_lin
from tinygrad.shape.shapetracker import ShapeTracker, View # noqa: F401 # pylint: disable=unused-import
N = 512
test/test_schedule.py
+9 -9
View File
@@ -70,7 +70,7 @@ def _test_conv2d(allowed:int, dtype:DType=dtypes.float, **kwargs):
def schedule_graph_rewrite(big_sink:UOp): return get_kernelize_map(big_sink)[big_sink]
class TestSchedule(unittest.TestCase):
def test_arange_avgpool2d(self, kcount=2):
def test_arange_avgpool2d(self, kcount=1):
x = Tensor.arange(25).reshape(1,1,5,5).cast(dtypes.float32)
t = x.avg_pool2d(padding=1)
sched = t.schedule()
@@ -1028,14 +1028,14 @@ class TestSchedule(unittest.TestCase):
Tensor.manual_seed(0)
x = Tensor.randn(4, 32).realize()
out = x.argmin(-1)
run_schedule(check_schedule(out, 3))
run_schedule(check_schedule(out, 2))
np.testing.assert_equal(out.numpy(), x.numpy().argmin(axis=-1))
def test_argmax_multireduce_fusion(self):
Tensor.manual_seed(0)
x = Tensor.randn(4, 32).realize()
out = x.argmax(-1)
run_schedule(check_schedule(out, 3))
run_schedule(check_schedule(out, 2))
np.testing.assert_equal(out.numpy(), x.numpy().argmax(axis=-1))
def test_scaled_dot_product_attention_multireduce_fusion(self):
@@ -1613,7 +1613,7 @@ class TestSchedule(unittest.TestCase):
Tensor.manual_seed(0)
x = Tensor.randn(10, 20).realize()
out = x.argmax(1)
run_schedule(check_schedule(out, 3)) # TODO: push a reduceop through a reshape
run_schedule(check_schedule(out, 2))
def test_conv2d(self): _test_conv2d(7)
def test_conv2d_fused(self): _test_conv2d(5, FUSE_CONV_BW=1)
@@ -1727,7 +1727,8 @@ class TestIndexing(unittest.TestCase):
s = Tensor.schedule(*lst)
lowered = [x[1] for x in lower_schedule(s.copy())]
kernels = [ei for ei in list(lowered) if isinstance(ei.prg, CompiledRunner)]
if FUSE_ARANGE: self.assertEqual(len(kernels), cnt)
if FUSE_ARANGE and len(kernels) != cnt:
raise KernelCountException(f"{len(kernels)} != {cnt}")
for ei in lowered: ei.run(do_update_stats=True)
return s
@@ -1741,7 +1742,7 @@ class TestIndexing(unittest.TestCase):
def test_simple_indexing_alt(self):
X = Tensor.arange(16).reshape(4, 4)
xt = X[[1, 2], [1, 2]]
self.check_schedule(xt, 5)
self.check_schedule(xt, 3)
np.testing.assert_equal(xt.numpy(), (np.arange(16).reshape(4, 4))[[1, 2], [1, 2]])
def test_advanced_indexing(self):
@@ -1753,13 +1754,13 @@ class TestIndexing(unittest.TestCase):
def test_advanced_indexing_alt(self):
X = Tensor.arange(6).reshape(3, 2)+1
xt = X[[Tensor([2]), Tensor([1])]]
self.check_schedule(xt, 6)
self.check_schedule(xt, 3)
np.testing.assert_equal(xt.numpy(), 6)
def test_advanced_simple_indexing_combined(self):
X = Tensor.arange(16).reshape(4, 4)
xt = X[1:2, [1, 2]]
self.check_schedule(xt, 4)
self.check_schedule(xt, 2)
def test_push_through_reshape(self):
Tensor.manual_seed(0)
@@ -2066,7 +2067,6 @@ class TestSwizzle(unittest.TestCase):
np.testing.assert_allclose(t.numpy(), x.numpy().sum(axis=1)+y.numpy().sum(axis=1), atol=1e-6, rtol=1e-3)
# kernels can only have 1 or n in each dim
@unittest.expectedFailure
def test_dont_parallelize_different_n(self):
Tensor.manual_seed(0)
x = Tensor.randn(4, 2, 2).realize()
test/test_search.py
+5 -5
View File
@@ -1,8 +1,8 @@
import unittest
from tinygrad.opt.kernel import Opt, OptOps, Kernel
from tinygrad.codegen.opt.kernel import Opt, OptOps, Kernel
from tinygrad.uop.ops import UOp, Ops
from tinygrad.opt.search import bufs_from_lin, actions, beam_search
from tinygrad.codegen.opt.search import bufs_from_lin, actions, beam_search
from tinygrad.device import Device
from tinygrad.tensor import Tensor
from tinygrad.dtype import dtypes
@@ -36,7 +36,7 @@ class TestBEAM(unittest.TestCase):
def test_get_kernel_actions_dedup(self):
from test.test_linearizer import helper_realized_ast
from tinygrad.opt.search import get_kernel_actions
from tinygrad.codegen.opt.search import get_kernel_actions
a = Tensor.empty(4, 3)
b = Tensor.empty(3)
realized_ast, _ = helper_realized_ast(a @ b)
@@ -61,7 +61,7 @@ class TestBEAM(unittest.TestCase):
@unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
def test_search_over_shape(self):
from test.test_linearizer import helper_realized_ast
from tinygrad.opt.search import get_kernel_actions
from tinygrad.codegen.opt.search import get_kernel_actions
dtype_pairs = [(tc.dtype_in, tc.dtype_out) for tc in Device[Device.DEFAULT].renderer.tensor_cores]
multi_shape_dtype_pairs = [dts for dts in dtype_pairs if dtype_pairs.count(dts) > 1]
@@ -78,7 +78,7 @@ class TestBEAM(unittest.TestCase):
def test_get_kernel_actions_preserves_actions_state(self):
from test.test_linearizer import helper_realized_ast
from tinygrad.opt.search import get_kernel_actions
from tinygrad.codegen.opt.search import get_kernel_actions
a = Tensor.rand(16, 16)
b = Tensor.rand(16, 16)
realized_ast, _ = helper_realized_ast(a @ b)
test/test_softmax_fusion.py
-1
View File
@@ -111,7 +111,6 @@ class TestFuse(unittest.TestCase):
with Context(NOOPT=1):
self._test_fuse(Tensor.scaled_dot_product_attention, q, k, v, atol=1e-5)
@unittest.expectedFailure
def test_mismatch_reduce(self):
a = Tensor.ones(16, 10).contiguous().realize()
b = Tensor.ones(16, 20).contiguous().realize()
test/test_uops.py
+1 -1
View File
@@ -14,7 +14,7 @@ from tinygrad.engine.realize import CompiledRunner, get_program
from tinygrad.codegen import full_rewrite
from tinygrad.uop.symbolic import sym
from tinygrad.device import is_dtype_supported
from tinygrad.opt.kernel import Opt, OptOps
from tinygrad.codegen.opt.kernel import Opt, OptOps
def to_uops_list(u:list[UOp], opts=None, skip_check=False) -> list[UOp]: return full_rewrite(UOp.sink(*u), opts)
test/test_uops_stats.py
+1 -1
View File
@@ -6,7 +6,7 @@ from tinygrad.renderer import Estimates
from tinygrad.codegen import full_rewrite
from tinygrad.uop.ops import Ops, UOp
from tinygrad.dtype import dtypes
from tinygrad.opt.kernel import Opt, OptOps, KernelOptError
from tinygrad.codegen.opt.kernel import Opt, OptOps, KernelOptError
from tinygrad.device import Device
def flops_mem(uops, ignore_indexing=False):
test/unit/test_linearizer_rewrite.py
+1 -1
View File
@@ -1,7 +1,7 @@
import unittest
from tinygrad import Tensor, Context, Device
from tinygrad.engine.realize import get_program
from tinygrad.opt.kernel import Opt, OptOps
from tinygrad.codegen.opt.kernel import Opt, OptOps
from tinygrad.uop.ops import KernelInfo
class TestLinearizerRewrite(unittest.TestCase):
test/unit/test_llm_tokenizer.py
+2 -1
View File
@@ -1,7 +1,8 @@
import unittest, base64, functools
import unittest, base64, functools, sys
from tinygrad.apps.llm import SimpleTokenizer, get_llama_re
from tinygrad.helpers import fetch
@unittest.skipIf(sys.platform == 'win32', "fetch race condition on Windows")
class TestLLMTokenizer(unittest.TestCase):
@functools.cached_property
def basic_tok(self): return SimpleTokenizer(".*", { b"a": 0, b"b": 1, b"c": 2, b"ab": 3, b"bc": 4 }, { "<x>": 5, "<y>": 6, "<z>": 7 })
test/unit/test_search.py
+2 -2
View File
@@ -1,8 +1,8 @@
import unittest
from tinygrad import Tensor, Device
from tinygrad.opt.kernel import Kernel
from tinygrad.codegen.opt.kernel import Kernel
from tinygrad.device import Buffer
from tinygrad.opt.search import get_test_global_size, bufs_from_lin
from tinygrad.codegen.opt.search import get_test_global_size, bufs_from_lin
from tinygrad.helpers import GlobalCounters
from extra.optimization.helpers import time_linearizer
from test.test_linearizer import push_views
test/unit/test_uop_spec.py
-11
View File
@@ -34,17 +34,6 @@ class TestUOpSpec(unittest.TestCase):
store = UOp(Ops.STORE, dtypes.void, (buf_0.view(ShapeTracker.from_shape((32, 1))), a+b))
helper_test_verify_ast(store)
def test_exactly_one_full_shape(self):
dtype = dtypes.int
bufs = [UOp(Ops.DEFINE_GLOBAL, dtype.ptr(), (), i) for i in range(6)]
a = UOp(Ops.LOAD, dtype, (bufs[2].view(ShapeTracker.from_shape((32, 1))),))
b = UOp(Ops.LOAD, dtype, (bufs[3].view(ShapeTracker.from_shape((32, 1))),))
st0 = UOp.store(bufs[0].view(ShapeTracker.from_shape((32, 1))), a+b)
a = UOp(Ops.LOAD, dtype, (bufs[4].view(ShapeTracker.from_shape((32, 32))),))
b = UOp(Ops.LOAD, dtype, (bufs[5].view(ShapeTracker.from_shape((32, 32))),))
st1 = UOp.store(bufs[1].view(ShapeTracker.from_shape((32, 32))), a+b)
with self.assertRaises(InvalidASTException): helper_test_verify_ast(st0, st1)
def test_no_implicit_broadcasting(self):
bufs = [UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), i) for i in range(2)]
a = UOp(Ops.LOAD, dtypes.float, (bufs[1].view(ShapeTracker.from_shape((4, 32))),))
tinygrad/codegen/__init__.py
+2 -2
View File
@@ -16,8 +16,8 @@ from tinygrad.codegen.devectorizer import load_store_folding, load_store_indexin
ReduceContext, correct_load_store, pm_render
from tinygrad.codegen.optional import get_late_rewrite_patterns
from tinygrad.codegen.linearize import block_create, pm_blockend_merge, block_merge, pm_finalize, BlockContext
from tinygrad.opt import pm_optimize
from tinygrad.opt.swizzler import view_left, view_right, fix_kernel_ops
from tinygrad.codegen.opt import pm_optimize
from tinygrad.codegen.opt.swizzler import view_left, view_right, fix_kernel_ops
@dataclass
class RewriteStep:
tinygrad/opt/__init__.py → tinygrad/codegen/opt/__init__.py
+3 -3
View File
@@ -1,7 +1,7 @@
# opt opinionatedly transforms an ast into an optimized ast using either heuristics or beam search
from tinygrad.opt.kernel import Kernel
from tinygrad.opt.heuristic import hand_coded_optimizations
from tinygrad.codegen.opt.kernel import Kernel
from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops
from tinygrad.helpers import NOOPT, BEAM, USE_TC, getenv
from tinygrad.renderer import Renderer
@@ -24,7 +24,7 @@ def get_optimized_ast(ast:UOp, renderer:Renderer) -> UOp:
elif not NOOPT:
if not k.apply_tensor_cores(USE_TC.value): k.apply_opts(hand_coded_optimizations(k))
if BEAM >= 1:
from tinygrad.opt.search import beam_search, bufs_from_lin
from tinygrad.codegen.opt.search import beam_search, bufs_from_lin
kb = Kernel(ast, opts=renderer)
rawbufs = bufs_from_lin(kb, allocate=False)
k = beam_search(kb, rawbufs, BEAM.value, bool(getenv("BEAM_ESTIMATE", 1)))
tinygrad/opt/heuristic.py → tinygrad/codegen/opt/heuristic.py
+15 -12
View File
@@ -1,5 +1,5 @@
import itertools
from tinygrad.opt.kernel import Kernel, Opt, OptOps, KernelOptError, AxisType
from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps, KernelOptError, AxisType
from tinygrad.helpers import getenv, DEBUG, prod, NOLOCALS
from tinygrad.dtype import ImageDType
from tinygrad.uop.ops import Ops, resolve
@@ -80,18 +80,21 @@ def hand_coded_optimizations(k:Kernel) -> list[Opt]:
else: break
# if last reduce dim is small(ish), loop unroll the reduce
upcast_size = prod(k.full_shape[a] for a in k.axes_of(AxisType.UPCAST, AxisType.UNROLL))
if k.unrollable_dims and (upcast_size <= 4 or not k.axes_of(AxisType.UNROLL)) and (upcast_size < 64):
if (s:=k.full_shape[k.unrollable_dims[-1]]) <= 32:
k.apply_opt(Opt(OptOps.UNROLL, len(k.unrollable_dims)-1, 0))
# if it's small, upcast a second reduce dimension too
if k.unrollable_dims and s <= 3 and k.full_shape[k.unrollable_dims[-1]] <= 3:
# NOTE: this can fail on multireduce with mismatching dimensions, this is okay
try:
upcast_size = prod(k.full_shape[a] for a in k.axes_of(AxisType.UPCAST, AxisType.UNROLL))
if k.unrollable_dims and (upcast_size <= 4 or not k.axes_of(AxisType.UNROLL)) and (upcast_size < 64):
if (s:=k.full_shape[k.unrollable_dims[-1]]) <= 32:
k.apply_opt(Opt(OptOps.UNROLL, len(k.unrollable_dims)-1, 0))
else:
for splits in [4]:
if k.full_shape[axis:=k.unrollable_dims[-1]]%splits == 0:
k.apply_opt(Opt(OptOps.UNROLL, len(k.unrollable_dims)-1, splits))
break
# if it's small, upcast a second reduce dimension too
if k.unrollable_dims and s <= 3 and k.full_shape[k.unrollable_dims[-1]] <= 3:
k.apply_opt(Opt(OptOps.UNROLL, len(k.unrollable_dims)-1, 0))
else:
for splits in [4]:
if k.full_shape[axis:=k.unrollable_dims[-1]]%splits == 0:
k.apply_opt(Opt(OptOps.UNROLL, len(k.unrollable_dims)-1, splits))
break
except KernelOptError: pass
# if nothing at all is upcasted and it's easy to, do an upcast
for splits in [4]:
tinygrad/opt/kernel.py → tinygrad/codegen/opt/kernel.py
+6 -3
View File
@@ -8,13 +8,13 @@ from enum import Enum, auto
from tinygrad.uop.ops import GroupOp, KernelInfo, UOp, Ops, can_pad, resolve, Variable, sint, graph_rewrite, AxisType
from tinygrad.uop.spec import type_verify, ast_spec
from tinygrad.device import Device
from tinygrad.opt.tc import TensorCore
from tinygrad.codegen.opt.tc import TensorCore
from tinygrad.renderer import Renderer
from tinygrad.dtype import ImageDType, AddrSpace
from tinygrad.helpers import all_same, colored, ansilen, dedup, prod, round_up, to_function_name, unwrap, argfix, DEBUG, TC_SELECT, TC_OPT, AMX
from tinygrad.shape.shapetracker import ShapeTracker
from tinygrad.shape.view import strides_for_shape, get_contraction
from tinygrad.opt.swizzler import view_left, view_left_through_load
from tinygrad.codegen.opt.swizzler import view_left, view_left_through_load
class OptOps(Enum):
TC = auto(); UPCAST = auto(); UNROLL = auto(); LOCAL = auto() # noqa: E702
@@ -271,7 +271,10 @@ class Kernel:
check(isinstance(opt.arg, int), "arg should be int")
amt = arg if (arg:=cast(int, opt.arg)) != 0 else self.full_shape[axis]
check(isinstance(amt, int) and amt != 1, f"shift/padto of {amt=}, 1 or symbolic amount is meaningless")
if opt.op is not OptOps.PADTO: check(self.full_shape[axis] % amt == 0, f"no longer valid shift {self.full_shape[axis]=}, {amt=}")
if opt.op is not OptOps.PADTO:
# we check both the full_shape and each shape
check(self.full_shape[axis] % amt == 0, f"no longer valid shift {self.full_shape[axis]=}, {amt=}")
for st in self.sts: check(st.shape[axis] == 1 or st.shape[axis] % amt == 0, f"no longer valid shift {st.shape[axis]=}, {amt=}")
else: amt = -1
if self.reduceop is not None and (opt.op in {OptOps.GROUP, OptOps.GROUPTOP} or \
tinygrad/opt/search.py → tinygrad/codegen/opt/search.py
+1 -1
View File
@@ -7,7 +7,7 @@ from tinygrad.device import Device, Buffer, Compiler
from tinygrad.helpers import prod, flatten, DEBUG, CACHELEVEL, diskcache_get, diskcache_put, getenv, Context, colored, time_to_str
from tinygrad.helpers import IGNORE_BEAM_CACHE, TC_SEARCH_OVER_SHAPE
from tinygrad.dtype import ImageDType, PtrDType
from tinygrad.opt.kernel import Kernel, Opt, OptOps, KernelOptError
from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps, KernelOptError
from tinygrad.tensor import Tensor
from tinygrad.engine.realize import CompiledRunner, get_program
from tinygrad.renderer import ProgramSpec
tinygrad/opt/swizzler.py → tinygrad/codegen/opt/swizzler.py
View File
tinygrad/opt/tc.py → tinygrad/codegen/opt/tc.py
View File
tinygrad/engine/realize.py
+2 -2
View File
@@ -8,7 +8,7 @@ from tinygrad.device import Device, Buffer
from tinygrad.renderer import Renderer, ProgramSpec, Estimates
from tinygrad.engine.schedule import ScheduleItem
from tinygrad.codegen import full_rewrite
from tinygrad.opt.kernel import Opt
from tinygrad.codegen.opt.kernel import Opt
# **************** Program Creation ****************
@@ -77,7 +77,7 @@ class CompiledRunner(Runner):
global_size, local_size = self.p.launch_dims(var_vals)
if global_size is not None and local_size is None and all_int(self.p.global_size): # type: ignore[arg-type]
# TODO: this is copied from get_program
from tinygrad.opt.search import optimize_local_size
from tinygrad.codegen.opt.search import optimize_local_size
local_size = optimize_local_size(self._prg, global_size, rawbufs)
global_size = [g//l if g%l == 0 else g/l for g,l in zip(global_size, local_size)]
self.p = replace(self.p, global_size=global_size, local_size=local_size)
tinygrad/helpers.py
+2 -2
View File
@@ -131,7 +131,7 @@ JIT_BATCH_SIZE = ContextVar("JIT_BATCH_SIZE", 32)
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, TC_SEARCH_OVER_SHAPE, NOLOCALS = ContextVar("TRANSCENDENTAL", 1), ContextVar("TC_SEARCH_OVER_SHAPE", 1), ContextVar("NOLOCALS", 0)
FUSE_ARANGE, FUSE_CONV_BW = ContextVar("FUSE_ARANGE", 0), ContextVar("FUSE_CONV_BW", 0)
FUSE_ARANGE, FUSE_CONV_BW = ContextVar("FUSE_ARANGE", 1), ContextVar("FUSE_CONV_BW", 0)
SPLIT_REDUCEOP, NO_MEMORY_PLANNER, RING = ContextVar("SPLIT_REDUCEOP", 1), ContextVar("NO_MEMORY_PLANNER", 0), ContextVar("RING", 1)
PICKLE_BUFFERS, PROFILE, LRU = ContextVar("PICKLE_BUFFERS", 1), ContextVar("PROFILE", getenv("VIZ")), ContextVar("LRU", 1)
CACHELEVEL, IGNORE_BEAM_CACHE, DEVECTORIZE = ContextVar("CACHELEVEL", 2), ContextVar("IGNORE_BEAM_CACHE", 0), ContextVar("DEVECTORIZE", 1)
@@ -218,7 +218,7 @@ def cpu_profile(name:str|TracingKey, device="CPU", is_copy=False, display=True)
cache_dir: str = os.path.join(getenv("XDG_CACHE_HOME", os.path.expanduser("~/Library/Caches" if OSX else "~/.cache")), "tinygrad")
CACHEDB: str = getenv("CACHEDB", os.path.abspath(os.path.join(cache_dir, "cache.db")))
VERSION = 21
VERSION = 22
_db_connection = None
def db_connection():
global _db_connection
tinygrad/renderer/__init__.py
+2 -2
View File
@@ -6,8 +6,8 @@ from tinygrad.helpers import to_function_name, dedup, prod
from tinygrad.uop.ops import Ops, UOp, sym_infer, sint, Variable, ssimplify, GroupOp, PatternMatcher
from tinygrad.dtype import AddrSpace, PtrDType
if TYPE_CHECKING:
from tinygrad.opt.tc import TensorCore
from tinygrad.opt.kernel import Opt
from tinygrad.codegen.opt.tc import TensorCore
from tinygrad.codegen.opt.kernel import Opt
@dataclass(frozen=True)
class Estimates:
tinygrad/renderer/cstyle.py
+1 -1
View File
@@ -1,7 +1,7 @@
from typing import Literal, Callable, cast
import os, math, sys
from collections import defaultdict, Counter
from tinygrad.opt import tc
from tinygrad.codegen.opt import tc
from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat
from tinygrad.helpers import strip_parens, getenv, prod, dedup, AMX
from tinygrad.dtype import ImageDType, dtypes, DType, PtrDType, AddrSpace, truncate
tinygrad/renderer/llvmir.py
+1 -1
View File
@@ -1,6 +1,6 @@
from typing import cast
import math, struct, sys
from tinygrad.opt import tc
from tinygrad.codegen.opt import tc
from tinygrad.renderer import Renderer
from tinygrad.renderer.cstyle import AMDRenderer
from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, GroupOp
tinygrad/renderer/ptx.py
+1 -1
View File
@@ -1,7 +1,7 @@
from typing import cast, Callable
import struct
from collections import defaultdict
from tinygrad.opt import tc
from tinygrad.codegen.opt import tc
from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp
from tinygrad.dtype import dtypes, DType, PtrDType, AddrSpace
from tinygrad.renderer import Renderer
tinygrad/runtime/ops_python.py
+1 -1
View File
@@ -7,7 +7,7 @@ import pickle, base64, itertools, time, struct, sys
from tinygrad.dtype import DType, dtypes, ImageDType, PtrDType, truncate
from tinygrad.helpers import all_same, getenv, flatten, get_single_element
from tinygrad.device import Compiled, Compiler, Allocator
from tinygrad.opt import tc
from tinygrad.codegen.opt import tc
from tinygrad.uop.ops import exec_alu, Ops, UOp, GroupOp
from tinygrad.renderer import Renderer
tinygrad/schedule/grouper.py
+2 -2
View File
@@ -26,8 +26,8 @@ def realize_before_view(ctx:dict[UOp, None], view:UOp, tr:UOp) -> None:
do_realize = PatternMatcher([
# always realize SINK parents
(UPat(Ops.SINK, name="s"), lambda ctx,s: ctx.update((x.base, None) for x in s.src if x.base.op not in ALWAYS_CONTIGUOUS)),
# always realize ASSIGN/CONTIGUOUS/GroupOp.Meta
(UPat({Ops.ASSIGN, Ops.CONTIGUOUS, *GroupOp.Meta}, name="tr"), realize),
# always realize ASSIGN/CONTIGUOUS/COPY/BUFFER_VIEW
(UPat({Ops.ASSIGN, Ops.CONTIGUOUS, Ops.COPY, Ops.BUFFER_VIEW}, name="tr"), realize),
# realize before expand or unsafe pad ops
(UPat(Ops.VIEW, src=(UPat(GroupOp.All-ALWAYS_CONTIGUOUS, name="tr"),), name="view"), realize_before_view),
# realize parents of COPY, MSELECT, MSTACK
tinygrad/schedule/kernelize.py
+1 -3
View File
@@ -7,7 +7,7 @@ from tinygrad.helpers import Metadata, all_int, all_same, prod, dedup, unwrap, g
from tinygrad.dtype import ImageDType
from tinygrad.schedule.multi import multi_pm
from tinygrad.schedule.grouper import group_realizes, ALWAYS_CONTIGUOUS
from tinygrad.opt.swizzler import apply_swizzle, swizzle_reduceop
from tinygrad.codegen.opt.swizzler import apply_swizzle, swizzle_reduceop
# creation can recurse a lot
import sys
@@ -263,8 +263,6 @@ def fuse_arange(root:UOp):
for u in toposort:
for s in u.src: local_children.setdefault(s, []).append(u)
fuse_rep: dict[UOp, UOp] = {}
# skip if root depends on aranges with different ndims. This can be improved
if any(len(set(dims)) > 1 for dims in zip(*[r.src[0].shape for r in local_arange])): return
for r in local_arange:
# skip if already fused
if len(r.arg) > 2: continue
tinygrad/tensor.py
+3 -2
View File
@@ -3507,15 +3507,16 @@ class Tensor(MathTrait):
"""
return self * self.softplus().tanh()
def softplus(self, beta=1) -> Tensor:
def softplus(self, beta=1.0, threshold=20.0) -> Tensor:
"""
Applies the Softplus function element-wise.
For numerical stability, the implementation folds into identity function when `self * beta > threshold`.
```python exec="true" source="above" session="tensor" result="python"
print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).softplus().numpy())
```
"""
return (1/beta) * (1 + (self*beta).exp()).log()
return (self * beta > threshold).where(self, (1/beta) * (1 + (self*beta).exp()).log())
def softsign(self) -> Tensor:
"""
tinygrad/uop/spec.py
+1 -10
View File
@@ -1,5 +1,5 @@
from typing import cast, Callable
from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, print_uops, python_alu, graph_rewrite, resolve
from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, print_uops, python_alu, graph_rewrite
from tinygrad.dtype import DType, ImageDType, dtypes, PtrDType, AddrSpace
from tinygrad.helpers import all_same, prod, DEBUG, ContextVar, Context
from tinygrad.shape.shapetracker import ShapeTracker
@@ -207,16 +207,7 @@ spec = PatternMatcher([
# *** this is the UOp AST spec ***
def verify_sink_dims(sink:UOp):
if not all_same([s.shape for s in sink.src]): return False
for dims in zip(*[x.shape for x in sink.toposort() if x.op is Ops.VIEW]):
if len(n_dims:={s for s in dims if resolve(s!=1)}) > 1:
print(f"# INVALID KERNEL DIMS: can only have 1 or n in each dimension: {n_dims}")
return False
ast_spec = PatternMatcher([
# shapes must have either 1 or n in each dimension
(UPat(Ops.SINK, src=UPat(Ops.STORE), name="sink"), verify_sink_dims),
# VIEW can only exist in the edges
(UPat(Ops.VIEW, src=(UPat((Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL),))), lambda: True),
(UPat(Ops.VIEW, name="view"), lambda view: len(view.src) == 0),
tinygrad/viz/index.html
+5
View File
@@ -55,6 +55,11 @@
opacity: 0.4;
pointer-events: none;
}
label {
display: inline-flex;
align-items: center;
gap: 4px;
}
.graph svg {
width: 100%;
height: 100%;
tinygrad/viz/js/index.js
+1 -1
View File
@@ -109,7 +109,7 @@ function formatTime(ts, dur=ts) {
}
const formatUnit = (d, unit="") => d3.format(".3~s")(d)+unit;
const colorScheme = {TINY:["#1b5745", "#354f52", "#354f52", "#46acc2", "#1d2e62", "#63b0cd"],
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"],
CATEGORICAL:["#ff8080", "#F4A261", "#C8F9D4", "#8D99AE", "#F4A261", "#ffffa2", "#ffffc0", "#87CEEB"],}