github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-04-07 03:00:26 -04:00
Code Issues Packages Projects Releases Wiki Activity

fix: getenv("NOLOCALS")/NOLOCALS context var (#10927)

Browse Source 
OptOps shouldn't rely on os.environ.
This commit is contained in:
qazal
2025-06-23 11:23:59 +03:00
committed by GitHub
parent 22f9696522
commit 4e864bd304
2 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tinygrad/helpers.py
View File

@@ -114,7 +114,7 @@ DEBUG, IMAGE, BEAM, NOOPT = ContextVar("DEBUG", 0), ContextVar("IMAGE", 0), Cont
JIT = ContextVar("JIT", 2 if platform.system() == 'Darwin' and ('Intel' in platform.processor() or 'i386' in platform.processor()) else 1)
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 = ContextVar("TRANSCENDENTAL", 1), ContextVar("TC_SEARCH_OVER_SHAPE", 1)
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)
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)

6
tinygrad/opt/heuristic.py
View File

@@ -1,6 +1,6 @@
import itertools
from tinygrad.opt.kernel import Kernel, Opt, OptOps, KernelOptError
from tinygrad.helpers import getenv, DEBUG, all_int, prod
from tinygrad.helpers import getenv, DEBUG, all_int, prod, NOLOCALS
from tinygrad.dtype import ImageDType
from tinygrad.uop.ops import Ops, resolve
@@ -112,7 +112,7 @@ def hand_coded_optimizations(k:Kernel) -> list[Opt]:
# **** local groups ****
if k.opts.has_local:
if getenv("NOLOCALS") and k.local_dims == 0 and not k.group_for_reduces:
if NOLOCALS and k.local_dims == 0 and not k.group_for_reduces:
k.apply_opt(Opt(OptOps.NOLOCALS))
else:
# prioritize making expand axes local
@@ -130,4 +130,4 @@ def hand_coded_optimizations(k:Kernel) -> list[Opt]:
k.apply_opt(Opt(OptOps.LOCAL, axis, local_sz))
if will_delete_shape: deleted_shape += 1
return k.applied_opts
return k.applied_opts