github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-02-19 02:44:40 -05:00
Code Issues Packages Projects Releases Wiki Activity

nv limit max local memory with envvar (#7265)

Browse Source
This commit is contained in:
nimlgen
2024-10-24 16:01:50 +03:00
committed by GitHub
parent f20b651ee0
commit 98f8d0ccf9
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tinygrad/runtime/ops_nv.py
View File

@@ -534,7 +534,7 @@ class NVDevice(HCQCompiled):
self.timeline_value += 2
def _ensure_has_local_memory(self, required):
if self.slm_per_thread >= required: return
if self.slm_per_thread >= required or ((maxlm:=getenv("NV_MAX_LOCAL_MEMORY_PER_THREAD")) > 0 and required >= maxlm): return
if self.shader_local_mem is not None: self.allocator.free(self.shader_local_mem, self.shader_local_mem.size)