github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-02-08 13:45:50 -05:00
Code Issues Packages Projects Releases Wiki Activity

add return type for HCQCompatAllocator _alloc (#5267)

Browse Source 
Co-authored-by: nimlgen <138685161+nimlgen@users.noreply.github.com>
This commit is contained in:
Vyacheslav Pachkov
2024-07-03 10:25:44 +03:00
committed by GitHub
parent 191463a919
commit d3e4e21759
3 changed files with 16 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tinygrad/runtime/ops_nv.py
View File

@@ -2,7 +2,7 @@ from __future__ import annotations
import os, ctypes, contextlib, pathlib, re, fcntl, functools, mmap, struct, tempfile, hashlib, subprocess, time, array
from typing import Tuple, List, Any
from dataclasses import dataclass
from tinygrad.device import HCQCompatCompiled, HCQCompatAllocator, Compiler, CompileError, BufferOptions
from tinygrad.device import HCQCompatCompiled, HCQCompatAllocator, HCQCompatAllocRes, Compiler, CompileError, BufferOptions
from tinygrad.helpers import getenv, from_mv, mv_address, init_c_struct_t, to_mv, round_up, to_char_p_p, DEBUG, prod, PROFILE
from tinygrad.renderer.cstyle import NVRenderer
from tinygrad.runtime.ops_cuda import check as cuda_check, _get_bytes, CUDACompiler, PTXCompiler, PTX
@@ -348,7 +348,7 @@ class NVProgram:
class NVAllocator(HCQCompatAllocator):
def __init__(self, device:NVDevice): super().__init__(device)
def _alloc(self, size:int, options:BufferOptions):
def _alloc(self, size:int, options:BufferOptions) -> HCQCompatAllocRes:
if options.host: return self.device._gpu_host_alloc(size)
return self.device._gpu_alloc(size, map_to_cpu=options.cpu_access, huge_page=(size > (16 << 20)))