nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

suppress exceptions only during finalization (#11451)

Browse Source 
* suppress exceptions only during finalization

* fix

* fix typing

* fix more warns

* fix

* better?

* Revert "better?"

This reverts commit a068aa5793.

* mm?

* no as e
This commit is contained in:
nimlgen
2025-07-31 13:57:12 +03:00
committed by GitHub
parent 24dd0d52ed
commit 75c2c42def
6 changed files with 25 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/helpers.py
+7
View File
@@ -81,6 +81,13 @@ def word_wrap(x, wrap=80):
while len(ansistrip(x[:i])) < wrap and i < len(x): i += 1
return x[:i] + "\n" + word_wrap(x[i:], wrap)
def suppress_finalizing(func):
def wrapper(*args, **kwargs):
try: return func(*args, **kwargs)
except (AttributeError, TypeError, ImportError):
if not getattr(sys, 'is_finalizing', lambda: True)(): raise # re-raise if not finalizing
return wrapper
def pluralize(st:str, cnt:int): return f"{cnt} {st}"+('' if cnt == 1 else 's')
class LazySeq(Generic[T]): # NOTE: Mapping requires __iter__ and __len__, Sequence requires supporting __len__ and slicing in __getitem__
tinygrad/runtime/ops_amd.py
+5 -6
View File
@@ -7,7 +7,7 @@ from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, H
from tinygrad.runtime.support.hcq import MMIOInterface
from tinygrad.uop.ops import sint
from tinygrad.device import Compiled, DMAFdRef, BufferSpec
from tinygrad.helpers import getenv, to_mv, round_up, data64_le, all_same, flatten, DEBUG, AMD_LLVM, PROFILE, ProfileEvent
from tinygrad.helpers import getenv, to_mv, round_up, data64_le, all_same, flatten, DEBUG, AMD_LLVM, PROFILE, ProfileEvent, suppress_finalizing
from tinygrad.renderer.cstyle import AMDRenderer
from tinygrad.renderer.llvmir import AMDLLVMRenderer
from tinygrad.runtime.autogen import kfd, hsa, pci, sqtt
@@ -473,11 +473,10 @@ class AMDAllocator(HCQAllocator['AMDDevice']):
def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
return self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access)
@suppress_finalizing
def _free(self, opaque, options:BufferSpec):
try:
self.dev.synchronize()
self.dev.iface.free(opaque)
except AttributeError: pass
self.dev.synchronize()
self.dev.iface.free(opaque)
def _map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
@@ -593,7 +592,7 @@ class KFDIface:
def free(self, mem):
if len(mem.mapped_devs) > 0:
gpus = (ctypes.c_int32 * len(mem.mapped_devs))(*[x.gpu_id for x in mem.mapped_devs])
gpus = (ctypes.c_int32 * len(mem.mapped_devs))(*[x.iface.gpu_id for x in mem.mapped_devs])
stm = kfd.AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(gpus), n_devices=len(gpus))
assert stm.n_success == len(gpus)
if mem.va_addr: FileIOInterface.munmap(mem.va_addr, mem.size)
tinygrad/runtime/ops_cuda.py
+3 -4
View File
@@ -1,6 +1,6 @@
from __future__ import annotations
import ctypes, ctypes.util, functools
from tinygrad.helpers import DEBUG, getenv, mv_address, init_c_var, init_c_struct_t
from tinygrad.helpers import DEBUG, getenv, mv_address, init_c_var, init_c_struct_t, suppress_finalizing
from tinygrad.device import Compiled, BufferSpec, LRUAllocator
from tinygrad.renderer.cstyle import CUDARenderer
from tinygrad.renderer.ptx import PTXRenderer
@@ -45,9 +45,8 @@ class CUDAProgram:
self.prg = prg
if self.smem > 0: check(cuda.cuFuncSetAttribute(self.prg, cuda.CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, self.smem))
def __del__(self):
try: check(cuda.cuModuleUnload(self.module))
except AttributeError: pass
@suppress_finalizing
def __del__(self): check(cuda.cuModuleUnload(self.module))
def __call__(self, *args, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
check(cuda.cuCtxSetCurrent(self.dev.context))
tinygrad/runtime/ops_gpu.py
+3 -4
View File
@@ -2,7 +2,7 @@ from __future__ import annotations
from typing import cast
import ctypes, functools, hashlib
from tinygrad.runtime.autogen import opencl as cl
from tinygrad.helpers import init_c_var, to_char_p_p, from_mv, OSX, DEBUG, getenv, mv_address
from tinygrad.helpers import init_c_var, to_char_p_p, from_mv, OSX, DEBUG, getenv, mv_address, suppress_finalizing
from tinygrad.renderer.cstyle import OpenCLRenderer, IntelRenderer
from tinygrad.device import BufferSpec, LRUAllocator, Compiled, Compiler, CompileError
@@ -69,9 +69,8 @@ class CLAllocator(LRUAllocator['CLDevice']):
cl.cl_image_format(cl.CL_RGBA, {2: cl.CL_HALF_FLOAT, 4: cl.CL_FLOAT}[options.image.itemsize]),
options.image.shape[1], options.image.shape[0], 0, None, status := ctypes.c_int32()), status), options)
return (checked(cl.clCreateBuffer(self.dev.context, cl.CL_MEM_READ_WRITE, size, None, status := ctypes.c_int32()), status), options)
def _free(self, opaque:tuple[ctypes._CData, BufferSpec], options:BufferSpec):
try: check(cl.clReleaseMemObject(opaque[0]))
except AttributeError: pass
@suppress_finalizing
def _free(self, opaque:tuple[ctypes._CData, BufferSpec], options:BufferSpec): check(cl.clReleaseMemObject(opaque[0]))
def _copyin(self, dest:tuple[ctypes._CData, BufferSpec], src:memoryview):
if dest[1].image is not None:
check(cl.clEnqueueWriteImage(self.dev.queue, dest[0], False, (ctypes.c_size_t * 3)(0,0,0),
tinygrad/runtime/ops_nv.py
+4 -5
View File
@@ -7,7 +7,7 @@ from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, H
from tinygrad.runtime.support.hcq import MMIOInterface, FileIOInterface, MOCKGPU
from tinygrad.uop.ops import sint
from tinygrad.device import BufferSpec
from tinygrad.helpers import getenv, mv_address, round_up, data64, data64_le, prod, OSX, to_mv, hi32, lo32
from tinygrad.helpers import getenv, mv_address, round_up, data64, data64_le, prod, OSX, to_mv, hi32, lo32, suppress_finalizing
from tinygrad.renderer.ptx import PTXRenderer
from tinygrad.renderer.cstyle import NVRenderer
from tinygrad.runtime.support.compiler_cuda import CUDACompiler, PTXCompiler, PTX, NVPTXCompiler, NVCompiler
@@ -276,11 +276,10 @@ class NVAllocator(HCQAllocator['NVDevice']):
def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
return self.dev.iface.alloc(size, cpu_access=options.cpu_access, host=options.host)
@suppress_finalizing
def _free(self, opaque:HCQBuffer, options:BufferSpec):
try:
self.dev.synchronize()
self.dev.iface.free(opaque)
except AttributeError: pass
self.dev.synchronize()
self.dev.iface.free(opaque)
def _map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
tinygrad/runtime/ops_webgpu.py
+3 -4
View File
@@ -1,7 +1,7 @@
import functools, struct
from tinygrad.device import Compiled, Allocator, Compiler, BufferSpec
from tinygrad.renderer.wgsl import WGSLRenderer
from tinygrad.helpers import round_up
from tinygrad.helpers import round_up, suppress_finalizing
from tinygrad.runtime.autogen import webgpu
from typing import List, Any, TypeAlias
import ctypes
@@ -188,9 +188,8 @@ class WebGpuAllocator(Allocator['WGPUDevPtr']):
def _copyout(self, dest:memoryview, src:WGPUBufPtr):
buffer_data = read_buffer(self.dev, src)
dest[:] = buffer_data[:dest.nbytes] if webgpu.wgpuBufferGetSize(src) > dest.nbytes else buffer_data
def _free(self, opaque:WGPUBufPtr, options:BufferSpec):
try: webgpu.wgpuBufferDestroy(opaque)
except AttributeError: pass
@suppress_finalizing
def _free(self, opaque:WGPUBufPtr, options:BufferSpec): webgpu.wgpuBufferDestroy(opaque)
class WebGpuDevice(Compiled):
def __init__(self, device:str):