commaai
/
tinygrad
mirror of https://github.com/commaai/tinygrad.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

device.py cleanups for -1 line

This commit is contained in:
George Hotz 2023-12-01 14:59:33 -08:00
parent f9b1de598f
commit eb595588bb
1 changed files with 11 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
tinygrad/device.py
View File

@ -2,7 +2,7 @@ from __future__ import annotations
import numpy as np
from collections import defaultdict
from typing import TYPE_CHECKING, Union, Any, List, Optional, Dict, Callable, Tuple
import importlib, inspect, functools, pathlib, time, re
import importlib, inspect, functools, pathlib, time, re, ctypes
from tinygrad.helpers import ansilen, DEBUG, getenv, GlobalCounters, colored, BEAM, NOOPT, all_int, to_function_name, DType, from_mv, dtypes, flat_mv
from tinygrad.shape.symbolic import Variable, sym_infer, sint
from tinygrad.ops import LazyOp, TernaryOps, get_lazyop_info, ReduceOps, BufferOps, BinaryOps, Op
@ -33,6 +33,8 @@ class _Device:
return "CPU"
Device = _Device()
# **************** Buffer / Allocator ****************
class Buffer:
def __init__(self, device:str, size:int, dtype:DType, opaque:Any=None):
assert isinstance(dtype, DType)
@ -103,7 +105,14 @@ class LRUAllocator(Allocator): # pylint: disable=abstract-method
opaques.clear()
def free(self, opaque:Any, size:int, dtype:DType): self.cache[(size, dtype)].append(opaque)
# **************** shared device helpers ****************
class _MallocAllocator(LRUAllocator):
def _alloc(self, size:int, dtype:DType): return (ctypes.c_uint8 * (size*dtype.itemsize))()
def as_buffer(self, src) -> memoryview: return flat_mv(memoryview(src))
def copyin(self, dest, src:memoryview): ctypes.memmove(dest, from_mv(src), len(src))
def copyout(self, dest:memoryview, src): ctypes.memmove(from_mv(dest), src, len(dest))
MallocAllocator = _MallocAllocator()
# **************** base Runner + helpers ****************
class JITRunner:
def __init__(self):
@ -271,11 +280,3 @@ def _get_optimized_linearizer(linearizer_opts:LinearizerOptions, ast:LazyOp) ->
if DEBUG >= 1: print(" < ".join(f"{nm:6s} : {lin.colored_shape(30, dense=True)} : {tm*1e6:8.2f} us" for nm, lin, tm in timed))
k = timed[0][1]
return k
import ctypes
class _MallocAllocator(LRUAllocator):
def _alloc(self, size:int, dtype:DType): return (ctypes.c_uint8 * (size*dtype.itemsize))()
def as_buffer(self, src) -> memoryview: return flat_mv(memoryview(src))
def copyin(self, dest, src:memoryview): ctypes.memmove(dest, from_mv(src), len(src))
def copyout(self, dest:memoryview, src): ctypes.memmove(from_mv(dest), src, len(dest))
MallocAllocator = _MallocAllocator()