diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 94747f38e8..4fa90bb270 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -183,7 +183,7 @@ jobs: - name: Run Clip tests for SD MLPerf on NULL backend run: DEV=NULL python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20 - name: Run AMD emulated BERT training on NULL backend - run: DEV=NULL::gfx1201 NULL_ALLOW_COPYOUT=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 + run: DEV=NULL::gfx1201 NULL_ALLOW_COPYOUT=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py # TODO: support fake weights #- name: Run LLaMA 7B on 4 fake devices # run: DEV=NULL python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing @@ -348,9 +348,9 @@ jobs: - name: DEV=NULL beautiful_mnist_multigpu run: DEV=NULL NULL_ALLOW_COPYOUT=1 python examples/beautiful_mnist_multigpu.py - name: Test Bert training - run: DEV=NULL NULL_ALLOW_COPYOUT=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py + run: DEV=NULL NULL_ALLOW_COPYOUT=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 - run: DEV=NULL NULL_ALLOW_COPYOUT=1 CAPTURE_PROCESS_REPLAY=0 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py + run: DEV=NULL NULL_ALLOW_COPYOUT=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=1 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 diff --git a/examples/beautiful_cifar.py b/examples/beautiful_cifar.py index ff66811e78..962523939c 100644 --- a/examples/beautiful_cifar.py +++ b/examples/beautiful_cifar.py @@ -9,8 +9,7 @@ from extra.lr_scheduler import OneCycleLR GPUS = [f'{Device.DEFAULT}:{i}' for i in range(getenv("GPUS", 1))] # override tinygrad defaults -dtypes.default_float = dtypes.half -Context(FUSE_OPTIM=1).__enter__() +Context(DEFAULT_FLOAT=dtypes.half, FUSE_OPTIM=1).__enter__() # from https://github.com/tysam-code/hlb-CIFAR10/blob/main/main.py batchsize = getenv("BS", 1024) diff --git a/test/backend/test_optim.py b/test/backend/test_optim.py index 1d526dc596..299ec7a44d 100644 --- a/test/backend/test_optim.py +++ b/test/backend/test_optim.py @@ -143,12 +143,11 @@ class TestOptim(unittest.TestCase): @unittest.skipUnless(dtypes.half in Device[Device.DEFAULT].renderer.supported_dtypes(), "need half") def test_mixed_precision(self): - old_default_float, dtypes.default_float = dtypes.default_float, dtypes.half + self.enterContext(Context(DEFAULT_FLOAT=dtypes.half)) # weight update would overflow without upcasting self._test_sgd(10, {'lr': 1e10}, 1e-6, 3e-4) self._test_adam(1, {'lr': 1e10}, 1e-4, 1e-4) self._test_adamw(1, {'lr': 1e10}, 1e-4, 1e-4) - dtypes.default_float = old_default_float def test_assert_tensor_train(self): t = Tensor.ones((1,1)) diff --git a/test/backend/test_randomness.py b/test/backend/test_randomness.py index 97e4294360..ad5ee25842 100644 --- a/test/backend/test_randomness.py +++ b/test/backend/test_randomness.py @@ -232,17 +232,14 @@ class TestRandomness(unittest.TestCase): @given(strat.sampled_from([dtypes.float, dtypes.float16, dtypes.bfloat16])) def test_randn_finite(self, default_float): if default_float not in Device[Device.DEFAULT].renderer.supported_dtypes(): return - self.enterContext(Context(CAPTURE_PROCESS_REPLAY=0)) # TODO: make default dtype ContextVar - old_default_float = dtypes.default_float # low precision can result in inf from randn - dtypes.default_float = default_float + self.enterContext(Context(DEFAULT_FLOAT=default_float)) t = Tensor.randn(64, 64) mx = t.max().numpy().item() mn = t.min().numpy().item() print(f"testing with {default_float=}") assert math.isfinite(mx), mx assert math.isfinite(mn), mn - dtypes.default_float = old_default_float def test_random_counter_overflow(self): device = Device.DEFAULT diff --git a/test/external/external_benchmark_bert_matmuls.py b/test/external/external_benchmark_bert_matmuls.py index fa2afd0388..fc630cf980 100644 --- a/test/external/external_benchmark_bert_matmuls.py +++ b/test/external/external_benchmark_bert_matmuls.py @@ -1,7 +1,8 @@ from tinygrad import Tensor, dtypes -dtypes.default_float = dtypes.float16 from tinygrad.dtype import to_dtype -from tinygrad.helpers import getenv +from tinygrad.helpers import getenv, Context + +Context(DEFAULT_FLOAT=dtypes.float16).__enter__() if __name__ == "__main__": # matmuls in bert layers diff --git a/test/external/external_benchmark_bert_softmax.py b/test/external/external_benchmark_bert_softmax.py index 35b180fa94..5ed47d9492 100644 --- a/test/external/external_benchmark_bert_softmax.py +++ b/test/external/external_benchmark_bert_softmax.py @@ -1,9 +1,10 @@ from tinygrad import Tensor, dtypes, GlobalCounters -dtypes.default_float = dtypes.float16 from tinygrad.dtype import to_dtype -from tinygrad.helpers import getenv +from tinygrad.helpers import getenv, Context from test.backend.test_softmax_fusion import single_kernel_softmax +Context(DEFAULT_FLOAT=dtypes.float16).__enter__() + if __name__ == "__main__": # softmax in bert layers BS = getenv("BS", 96//6) diff --git a/test/external/external_test_jit_on_models.py b/test/external/external_test_jit_on_models.py index e2d14d9768..6af6a706cc 100644 --- a/test/external/external_test_jit_on_models.py +++ b/test/external/external_test_jit_on_models.py @@ -3,6 +3,7 @@ import unittest import numpy as np from tinygrad import Tensor, dtypes from tinygrad.engine.jit import TinyJit +from tinygrad.helpers import Context from test.helpers import derandomize_model from examples.llama import Transformer @@ -14,8 +15,7 @@ def helper_test_jitted_correctness(gen, train, train_jit): class TestJittedModels(unittest.TestCase): def test_jitted_tiny_llama(self): - old_float = dtypes.default_float - dtypes.default_float = dtypes.float16 + self.enterContext(Context(DEFAULT_FLOAT=dtypes.float16)) args_tiny = {"dim": 1024, "hidden_dim": 1024, "n_heads": 8, "n_layers": 8, "norm_eps": 1e-05, "vocab_size": 1000} model = Transformer(**args_tiny) @@ -25,7 +25,6 @@ class TestJittedModels(unittest.TestCase): @TinyJit def test_jit(t): return model(t, 0).realize() helper_test_jitted_correctness(lambda: (Tensor([[1,]]),), test, test_jit) - dtypes.default_float = old_float def test_jitted_stable_diffusion(self): from examples.stable_diffusion import UNetModel, unet_params diff --git a/test/null/test_dtype_spec.py b/test/null/test_dtype_spec.py index ac04ae5ca6..2aaebefe05 100644 --- a/test/null/test_dtype_spec.py +++ b/test/null/test_dtype_spec.py @@ -2,7 +2,7 @@ import unittest, math, struct, operator from tinygrad import Tensor, Device from tinygrad.dtype import DTYPES_DICT, dtypes, Invalid, truncate, float_to_fp16, float_to_bf16, _to_np_dtype, least_upper_dtype, least_upper_float -from tinygrad.helpers import getenv +from tinygrad.helpers import getenv, Context from hypothesis import given, settings, strategies as strat import numpy as np import torch @@ -245,19 +245,14 @@ class TestTypePromotion(unittest.TestCase): assert least_upper_dtype(dtypes.weakfloat, dtypes.float64) == dtypes.float64 class TestTypeSpec(unittest.TestCase): - def setUp(self): - self.old_default_int, self.old_default_float = dtypes.default_int, dtypes.default_float - def tearDown(self): - dtypes.default_int, dtypes.default_float = self.old_default_int, self.old_default_float - def test_set_dtype_default(self): for default_int in [dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64]: - dtypes.default_int = default_int - assert dtypes.default_int == default_int + with Context(DEFAULT_INT=default_int): + assert dtypes.default_int == default_int for default_float in [*dtypes.fp8s, dtypes.float16, dtypes.bfloat16, dtypes.float32, dtypes.float64]: - dtypes.default_float = default_float - assert dtypes.default_float == default_float + with Context(DEFAULT_FLOAT=default_float): + assert dtypes.default_float == default_float @given(strat.sampled_from(core_dtypes), strat.sampled_from([operator.gt, operator.ge, operator.le, operator.lt, operator.eq, operator.ne])) def test_bool_ops(self, dtype, op): @@ -265,7 +260,7 @@ class TestTypeSpec(unittest.TestCase): @given(strat.sampled_from(core_dtypes), strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats)) def test_functions_return_index(self, dtype, default_int, default_float): - dtypes.default_int, dtypes.default_float = default_int, default_float + self.enterContext(Context(DEFAULT_INT=default_int, DEFAULT_FLOAT=default_float)) assert Tensor([0, 1], dtype=dtype).argmax().dtype == dtypes.int32 assert Tensor([0, 1], dtype=dtype).argmin().dtype == dtypes.int32 assert Tensor([0, 1], dtype=dtype).multinomial().dtype == dtypes.int32 @@ -285,7 +280,7 @@ class TestTypeSpec(unittest.TestCase): @given(strat.sampled_from(dtype_floats), strat.sampled_from(dtype_floats)) def test_attention_returns_same_dtype(self, data_dtype, default_float): - dtypes.default_float = default_float + self.enterContext(Context(DEFAULT_FLOAT=default_float)) query = Tensor.rand(32, 8, 128, 64, dtype=data_dtype) key = Tensor.rand(32, 8, 128, 64, dtype=data_dtype) value = Tensor.rand(32, 8, 128, 64, dtype=data_dtype) @@ -296,19 +291,14 @@ class TestTypeSpec(unittest.TestCase): assert query.scaled_dot_product_attention(key, value, attn_mask=mask).dtype == data_dtype class TestAutoCastType(unittest.TestCase): - def setUp(self): - self.old_default_int, self.old_default_float = dtypes.default_int, dtypes.default_float - def tearDown(self): - dtypes.default_int, dtypes.default_float = self.old_default_int, self.old_default_float - @given(strat.sampled_from(dtype_floats), strat.sampled_from(dtype_floats)) def test_least_upper_float_input_is_float(self, input_dtype, default_float): - dtypes.default_float = default_float + self.enterContext(Context(DEFAULT_FLOAT=default_float)) self.assertEqual(least_upper_float(input_dtype), input_dtype) @given(strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats)) def test_least_upper_float_input_is_int(self, input_dtype, default_float): - dtypes.default_float = default_float + self.enterContext(Context(DEFAULT_FLOAT=default_float)) self.assertEqual(least_upper_float(input_dtype), default_float) @given(strat.sampled_from(core_dtypes)) @@ -332,7 +322,7 @@ class TestAutoCastType(unittest.TestCase): @given(strat.sampled_from(dtype_floats)) def test_int_div_int(self, default_float): - dtypes.default_float = default_float + self.enterContext(Context(DEFAULT_FLOAT=default_float)) self.assertEqual(Tensor([1]).div(Tensor([2])).dtype, default_float) def test_sum(self): diff --git a/test/null/test_real_world.py b/test/null/test_real_world.py index 0c058106ec..378e8b924b 100644 --- a/test/null/test_real_world.py +++ b/test/null/test_real_world.py @@ -47,13 +47,8 @@ class TestRealWorld(unittest.TestCase): gc.collect() global global_mem_used global_mem_used = GlobalCounters.mem_used - self.old_float = dtypes.default_float - self.enterContext(Context(CAPTURE_PROCESS_REPLAY=0)) # TODO: make default dtype ContextVar np.random.seed(2002) - def tearDown(self): - dtypes.default_float = self.old_float - @slow @unittest.skipUnless(dtypes.float16 in supported_dtypes, "need dtypes.float16") def test_stable_diffusion(self): @@ -82,7 +77,7 @@ class TestRealWorld(unittest.TestCase): @unittest.skipUnless(dtypes.float16 in supported_dtypes, "need dtypes.float16") def test_llama(self): - dtypes.default_float = dtypes.float16 + self.enterContext(Context(DEFAULT_FLOAT=dtypes.float16)) args_tiny = {"dim": 1024, "hidden_dim": 2048, "n_heads": 8, "n_layers": 8, "norm_eps": 1e-05, "vocab_size": 1000} model = LLaMaTransformer(**args_tiny) @@ -94,7 +89,7 @@ class TestRealWorld(unittest.TestCase): @unittest.skipUnless(dtypes.float16 in supported_dtypes, "need dtypes.float16") def test_gpt2(self): - dtypes.default_float = dtypes.float16 + self.enterContext(Context(DEFAULT_FLOAT=dtypes.float16)) args_tiny = {"dim": 1024, "n_heads": 8, "n_layers": 8, "norm_eps": 1e-5, "vocab_size": 1000} model = GPT2Transformer(**args_tiny) @@ -151,7 +146,7 @@ class TestRealWorld(unittest.TestCase): @unittest.skipUnless(dtypes.float16 in supported_dtypes, "need dtypes.float16") def test_train_cifar_hyp(self): - dtypes.default_float = dtypes.float16 + self.enterContext(Context(DEFAULT_FLOAT=dtypes.float16)) with Context(TRAINING=1): model = SpeedyResNet(Tensor.ones((12,3,2,2))) optimizer = optim.SGD(get_parameters(model), lr=0.01, momentum=hyp['opt']['momentum'], nesterov=True, weight_decay=hyp['opt']['bias_decay']) diff --git a/test/null/test_schedule.py b/test/null/test_schedule.py index 27f12864a5..7ca7d900ff 100644 --- a/test/null/test_schedule.py +++ b/test/null/test_schedule.py @@ -603,14 +603,12 @@ class TestSchedule(unittest.TestCase): check_schedule(p, 4) def test_conv2d(self, allowed=4, dtype=dtypes.float): - old_default_float, dtypes.default_float = dtypes.default_float, dtype - dtypes.default_float = dtype + self.enterContext(Context(DEFAULT_FLOAT=dtype)) Tensor.manual_seed(0) BS, CIN = 2, 3 img = Tensor.randn(BS, CIN, 64, 64).realize() w = Tensor.uniform(16, CIN, 3, 3).realize() ret = Tensor.conv2d(img, w).relu().mean().backward() - dtypes.default_float = old_default_float linear, var_vals = Tensor.linear_with_vars(ret, img.grad, w.grad) cnt = len([call for call in linear.src if call.src[0].op is Ops.SINK]) assert cnt == allowed, f"expected {allowed} kernels, got {cnt}" diff --git a/test/unit/test_dtype_spec.py b/test/unit/test_dtype_spec.py index 2d60d2223c..cbb3bc5ba6 100644 --- a/test/unit/test_dtype_spec.py +++ b/test/unit/test_dtype_spec.py @@ -39,11 +39,13 @@ def _assert_eq(tensor:Tensor, target_dtype:DType, target, tol_target_dtype:float raise AssertionError(f"\ntensor {tensor.numpy()} dtype {tensor.dtype} does not match target {target} with dtype {target_dtype}") from e class TestTypeSpec(unittest.TestCase): - def setUp(self): - self.old_default_int, self.old_default_float = dtypes.default_int, dtypes.default_float - self.enterContext(Context(CAPTURE_PROCESS_REPLAY=0)) # TODO: make default dtype ContextVar - def tearDown(self): - dtypes.default_int, dtypes.default_float = self.old_default_int, self.old_default_float + def test_default_dtype_context(self): + default_float, default_int = dtypes.default_float, dtypes.default_int + with Context(DEFAULT_FLOAT=dtypes.half, DEFAULT_INT=dtypes.int16): + assert dtypes.default_float is dtypes.half + assert dtypes.default_int is dtypes.int16 + assert dtypes.default_float is default_float + assert dtypes.default_int is default_int @unittest.skip("this test is slow and spawning whole pythons") def test_env_set_default_float(self): @@ -82,7 +84,7 @@ class TestTypeSpec(unittest.TestCase): @given(strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats)) def test_creation(self, default_int, default_float): - dtypes.default_int, dtypes.default_float = default_int, default_float + self.enterContext(Context(DEFAULT_INT=default_int, DEFAULT_FLOAT=default_float)) _assert_eq(Tensor(True), dtypes.bool, True) _assert_eq(Tensor(None), dtypes.weakfloat, []) _assert_eq(Tensor(2), dtypes.weakint, 2) @@ -101,7 +103,7 @@ class TestTypeSpec(unittest.TestCase): @given(strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats)) def test_full(self, default_int, default_float): - dtypes.default_int, dtypes.default_float = default_int, default_float + self.enterContext(Context(DEFAULT_INT=default_int, DEFAULT_FLOAT=default_float)) _assert_eq(Tensor.zeros((2, 3)), dtypes.default_float, np.zeros((2, 3))) _assert_eq(Tensor.zeros((2, 3), dtype=dtypes.int64), dtypes.int64, np.zeros((2, 3))) @@ -124,7 +126,7 @@ class TestTypeSpec(unittest.TestCase): @given(strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats)) def test_reduce_0d_default(self, default_int, default_float): - dtypes.default_int, dtypes.default_float = default_int, default_float + self.enterContext(Context(DEFAULT_INT=default_int, DEFAULT_FLOAT=default_float)) _assert_eq(Tensor.ones((2,3,0)).sum(2), dtypes.default_float, np.zeros((2, 3))) # TODO: what should this one be? # _assert_eq(Tensor.ones((2,3,0), dtype=dtypes.default_int).sum(2), dtypes.default_int, np.zeros((2, 3))) @@ -132,7 +134,7 @@ class TestTypeSpec(unittest.TestCase): @given(strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats)) def test_arange(self, default_int, default_float): - dtypes.default_int, dtypes.default_float = default_int, default_float + self.enterContext(Context(DEFAULT_INT=default_int, DEFAULT_FLOAT=default_float)) _assert_eq(Tensor.arange(5), dtypes.default_int, np.arange(5)) _assert_eq(Tensor.arange(120), dtypes.default_int, np.arange(120)) @@ -149,12 +151,6 @@ class TestTypeSpec(unittest.TestCase): _assert_eq(Tensor.arange(5.0, 3.0), dtypes.default_float, np.arange(5.0, 3.0)) class TestAutoCastType(unittest.TestCase): - def setUp(self): - self.old_default_int, self.old_default_float = dtypes.default_int, dtypes.default_float - self.enterContext(Context(CAPTURE_PROCESS_REPLAY=0)) # TODO: make default dtype ContextVar - def tearDown(self): - dtypes.default_int, dtypes.default_float = self.old_default_int, self.old_default_float - def test_int_sqrt(self): _assert_eq(Tensor([1, 4, 9, 16]).sqrt(), dtypes.default_float, [1, 2, 3, 4]) @@ -194,22 +190,18 @@ class TestAutoCastType(unittest.TestCase): np.testing.assert_allclose(t.prod(dtype=dtypes.float32).numpy(), 20000) def test_gradient_dtype(self): - old_default_float = dtypes.default_float - for default_dtype in dtypes.floats: if default_dtype not in supported_dtypes: continue - dtypes.default_float = default_dtype - for dtype in dtypes.floats: - if dtype not in supported_dtypes: continue - if DEBUG >= 2: - print(f"testing {default_dtype=}, {dtype=}") - a = Tensor([1, 2, 3], dtype=dtype) - b = (a * 5).sum() - b.backward() # if there is dtype mismatch, lazy should assert - assert a.grad.dtype == a.dtype - np.testing.assert_allclose(a.grad.numpy(), [5, 5, 5]) - - dtypes.default_float = old_default_float + with Context(DEFAULT_FLOAT=default_dtype): + for dtype in dtypes.floats: + if dtype not in supported_dtypes: continue + if DEBUG >= 2: + print(f"testing {default_dtype=}, {dtype=}") + a = Tensor([1, 2, 3], dtype=dtype) + b = (a * 5).sum() + b.backward() # if there is dtype mismatch, lazy should assert + assert a.grad.dtype == a.dtype + np.testing.assert_allclose(a.grad.numpy(), [5, 5, 5]) @unittest.skipIf(Device.DEFAULT == "PYTHON", "very slow") @slow diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index f690245bba..17693ed9d4 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -1,7 +1,7 @@ from dataclasses import replace, dataclass import itertools, functools from tinygrad.helpers import DISABLE_FAST_IDIV, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC -from tinygrad.helpers import ALLOW_TF32, TracingKey, Context, panic +from tinygrad.helpers import ALLOW_TF32, DEFAULT_FLOAT, DEFAULT_INT, TracingKey, Context, panic from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, ProgramInfo, GroupOp from tinygrad.uop.ops import AxisType from tinygrad.uop.render import pyrender @@ -450,7 +450,7 @@ def do_to_program(ast:UOp, renderer:Renderer) -> UOp: to_program_cache: dict[tuple, UOp] = {} def to_program(ast:UOp, renderer:Renderer) -> UOp: - config = (NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC, IMAGE, DISABLE_FAST_IDIV, TRANSCENDENTAL, ALLOW_TF32) + config = (NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC, IMAGE, DISABLE_FAST_IDIV, TRANSCENDENTAL, ALLOW_TF32, DEFAULT_FLOAT, DEFAULT_INT) key = (ast.key, type(renderer), renderer.target, *[x.value for x in config]) if (prg:=to_program_cache.get(key)) is None: to_program_cache[key] = prg = do_to_program(ast, renderer) return prg diff --git a/tinygrad/dtype.py b/tinygrad/dtype.py index eddc38a4f2..2e97e9a8de 100644 --- a/tinygrad/dtype.py +++ b/tinygrad/dtype.py @@ -2,7 +2,7 @@ from __future__ import annotations from typing import Final, ClassVar, Callable, Literal import math, struct, ctypes, functools from dataclasses import dataclass, fields -from tinygrad.helpers import getenv +from tinygrad.helpers import getenv, DEFAULT_FLOAT, DEFAULT_INT from enum import IntEnum, auto class ConstFloat(float): @@ -83,7 +83,7 @@ class DType(metaclass=DTypeMetaClass): return ConstFloat(float(val)) if dtypes.is_float(self) else bool(val) if dtypes.is_bool(self) else int(val) -class dtypes: +class DTypes: @staticmethod @functools.cache def is_float(x: DType) -> bool: return x in (dtypes.floats + (dtypes.weakfloat,)) @@ -138,8 +138,10 @@ class dtypes: uchar = uint8; ushort = uint16; uint = uint32; ulong = uint64 # noqa: E702 char = int8; short = int16; int = int32; long = int64 # noqa: E702 - default_float: ClassVar[DType] = float32 - default_int: ClassVar[DType] = int32 + @property + def default_float(self) -> DType: return to_dtype(DEFAULT_FLOAT.value) + @property + def default_int(self) -> DType: return to_dtype(DEFAULT_INT.value) fp8_ocp = (fp8e4m3, fp8e5m2) fp8_fnuz = (fp8e4m3fnuz, fp8e5m2fnuz) @@ -155,12 +157,12 @@ class dtypes: weaks = (weakint, weakfloat) all = floats + ints + (bool,) # noqa: A003 -if (env_default_float := getenv("DEFAULT_FLOAT", "")): - dtypes.default_float = getattr(dtypes, env_default_float.lower()) - assert dtypes.is_float(dtypes.default_float), f"{env_default_float} is not a float dtype" +dtypes = DTypes() DTypeLike = str|DType def to_dtype(dtype:DTypeLike) -> DType: return dtype if isinstance(dtype, DType) else getattr(dtypes, dtype.lower()) +assert dtypes.is_float(dtypes.default_float), f"{DEFAULT_FLOAT.value} is not a float dtype" +assert dtypes.is_int(dtypes.default_int), f"{DEFAULT_INT.value} is not an int dtype" def strong_dtype(dtype:DType) -> DType: return {dtypes.weakint: dtypes.default_int, dtypes.weakfloat: dtypes.default_float}.get(dtype, dtype) @@ -184,7 +186,7 @@ def least_upper_dtype(*ds:DType) -> DType: def least_upper_float(dt:DType) -> DType: return dtypes.weakfloat if dt is dtypes.weakint else dt if dtypes.is_float(dt) else least_upper_dtype(dt, dtypes.default_float) -DTYPES_DICT = {k: v for k, v in dtypes.__dict__.items() if isinstance(v, DType) and not k.startswith(("default", "void", "weak", "_"))} +DTYPES_DICT = {k: v for k, v in DTypes.__dict__.items() if isinstance(v, DType) and not k.startswith(("default", "void", "weak", "_"))} INVERSE_DTYPES_DICT = {**{v.name:k for k,v in DTYPES_DICT.items()}, "void": "void", "weakint":"weakint", "weakfloat":"weakfloat"} @functools.cache diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index 5b7e9891a8..86d0136b18 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -247,6 +247,7 @@ FUSE_OPTIM = ContextVar("FUSE_OPTIM", 0) ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE = ContextVar("ALLOW_DEVICE_USAGE", 1), ContextVar("MAX_BUFFER_SIZE", 0) MAX_KERNEL_BUFFERS = ContextVar("MAX_KERNEL_BUFFERS", 0) EMULATED_DTYPES = ContextVar("EMULATED_DTYPES", "") +DEFAULT_FLOAT, DEFAULT_INT = ContextVar("DEFAULT_FLOAT", "float32"), ContextVar("DEFAULT_INT", "int32") CAPTURE_PROCESS_REPLAY = ContextVar("CAPTURE_PROCESS_REPLAY", 0) def _get_cpu_count() -> int: # os.process_cpu_count (3.13+) respects cgroup limits