github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-02-16 01:26:29 -05:00
Code Issues Packages Projects Releases Wiki Activity

a little more typing [pr] (#8346)

Browse Source 
* a little more typing [pr]

* few more
This commit is contained in:
George Hotz
2024-12-19 22:09:52 -08:00
committed by GitHub
parent 62e5d96446
commit 82833f1b3c
11 changed files with 38 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tinygrad/runtime/graph/hcq.py
View File

@@ -1,5 +1,5 @@
import collections, time
from typing import List, Any, cast, Optional
from typing import Any, cast, Optional
from tinygrad.helpers import round_up, PROFILE
from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQSignal, HCQBuffer, HWQueue, HCQArgsState, BumpAllocator
from tinygrad.device import Buffer, BufferSpec, Compiled, Device, ProfileGraphEntry, ProfileGraphEvent
@@ -42,7 +42,7 @@ class HCQGraph(MultiGraphRunner):
# graph-related tasks. This synchronization uses a global timeline signal per device. Within the graph, the compute queue coordinates with
# global operations and sets a kickoff signal. Any queue accessing a buffer from another device waits for this signal from the device’s
# compute queue to ensure exclusive access. The compute queue signals the completion of the graph, synchronizing with the device's copy queue.
self.ji_schedule: dict[int, tuple[HCQCompiled, HWQueue, List, List, HCQSignal, Optional[int]]] = {}
self.ji_schedule: dict[int, tuple[HCQCompiled, HWQueue, list, list, HCQSignal, Optional[int]]] = {}
self.comp_queues: dict[HCQCompiled, HWQueue] = {dev: dev.hw_compute_queue_t() for dev in self.devices}
self.copy_queues: dict[HCQCompiled, HWQueue] = {} # lazy allocation

4
tinygrad/runtime/ops_nv.py
View File

@@ -1,7 +1,7 @@
from __future__ import annotations
import os, ctypes, contextlib, re, fcntl, functools, mmap, struct, array, sys
assert sys.platform != 'win32'
from typing import List, Any, cast, Union, Type, Optional
from typing import Any, cast, Union, Type, Optional
from dataclasses import dataclass
from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, HWQueue, CLikeArgsState, HCQProgram, HCQSignal, BumpAllocator
from tinygrad.ops import sint
@@ -285,7 +285,7 @@ class NVDevice(HCQCompiled[NVSignal]):
root = None
fd_ctl: int = -1
fd_uvm: int = -1
gpus_info: Union[List, ctypes.Array] = []
gpus_info: Union[list, ctypes.Array] = []
signals_page: Any = None
signals_pool: list[int] = []