github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-02-15 09:05:40 -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_amd.py
View File

@@ -2,7 +2,7 @@ from __future__ import annotations
from typing import Tuple, List, Any
import os, fcntl, ctypes, ctypes.util, functools, re, pathlib, mmap, struct, errno, subprocess, time, array
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, init_c_struct_t, to_mv, round_up, DEBUG, PROFILE, mv_address
from tinygrad.renderer.cstyle import AMDRenderer
from tinygrad.runtime.driver.hip_comgr import compile_hip
@@ -385,7 +385,7 @@ class AMDProgram:
class AMDAllocator(HCQCompatAllocator):
def __init__(self, device:AMDDevice): super().__init__(device, batch_size=SDMA_MAX_COPY_SIZE)
def _alloc(self, size:int, options:BufferOptions):
def _alloc(self, size:int, options:BufferOptions) -> HCQCompatAllocRes:
try:
if options.host: return self.device._gpu_alloc(size, kfd.KFD_IOC_ALLOC_MEM_FLAGS_USERPTR, public=True)
return self.device._gpu_alloc(size, kfd.KFD_IOC_ALLOC_MEM_FLAGS_VRAM, public=options.cpu_access)