From 75c2c42def014b786b110b415dbcf9fb79f8298d Mon Sep 17 00:00:00 2001 From: nimlgen <138685161+nimlgen@users.noreply.github.com> Date: Thu, 31 Jul 2025 13:57:12 +0300 Subject: [PATCH] suppress exceptions only during finalization (#11451) * suppress exceptions only during finalization * fix * fix typing * fix more warns * fix * better? * Revert "better?" This reverts commit a068aa57936f180e68f837247758c48cac8a344f. * mm? * no as e --- tinygrad/helpers.py | 7 +++++++ tinygrad/runtime/ops_amd.py | 11 +++++------ tinygrad/runtime/ops_cuda.py | 7 +++---- tinygrad/runtime/ops_gpu.py | 7 +++---- tinygrad/runtime/ops_nv.py | 9 ++++----- tinygrad/runtime/ops_webgpu.py | 7 +++---- 6 files changed, 25 insertions(+), 23 deletions(-) diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index e60840dd22..34306e22ac 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -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__ diff --git a/tinygrad/runtime/ops_amd.py b/tinygrad/runtime/ops_amd.py index 3c440768ab..abb4b603e3 100644 --- a/tinygrad/runtime/ops_amd.py +++ b/tinygrad/runtime/ops_amd.py @@ -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) diff --git a/tinygrad/runtime/ops_cuda.py b/tinygrad/runtime/ops_cuda.py index 574c42e9fc..326ae84f01 100644 --- a/tinygrad/runtime/ops_cuda.py +++ b/tinygrad/runtime/ops_cuda.py @@ -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)) diff --git a/tinygrad/runtime/ops_gpu.py b/tinygrad/runtime/ops_gpu.py index 3b9b1d6965..c9ebb338c6 100644 --- a/tinygrad/runtime/ops_gpu.py +++ b/tinygrad/runtime/ops_gpu.py @@ -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), diff --git a/tinygrad/runtime/ops_nv.py b/tinygrad/runtime/ops_nv.py index ebb459bdf2..81dbbd3715 100644 --- a/tinygrad/runtime/ops_nv.py +++ b/tinygrad/runtime/ops_nv.py @@ -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) diff --git a/tinygrad/runtime/ops_webgpu.py b/tinygrad/runtime/ops_webgpu.py index 1386c963cf..eef0a5cc30 100644 --- a/tinygrad/runtime/ops_webgpu.py +++ b/tinygrad/runtime/ops_webgpu.py @@ -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):