mawei/StarPilot
1
0
Fork 0
mirror of https://github.com/firestar5683/StarPilot.git synced 2026-08-24 01:33:46 +08:00
Code Issues Packages Projects Releases Wiki Activity

The smallest promotion

Browse Source
This commit is contained in:
firestar5683
2026-08-10 22:53:46 -05:00
parent d938c9daa6
commit 6bd48c28e8
582 changed files with 64651 additions and 17100 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad_repo/tinygrad/runtime/ops_python.py
+31 -24
View File
@@ -5,11 +5,11 @@
from typing import Any, TYPE_CHECKING
import pickle, base64, itertools, time, sys, functools
from dataclasses import replace
from tinygrad.dtype import DType, dtypes, ImageDType, AddrSpace, truncate, storage_fmt_for_dtype, to_storage_scalar, from_storage_scalar
from tinygrad.helpers import all_same, getenv, flatten, Target, IMAGE
from tinygrad.device import Compiled, Compiler, Allocator
from tinygrad.dtype import bitcast, DType, dtypes, AddrSpace, truncate, storage_fmt_for_dtype, to_storage_scalar, from_storage_scalar
from tinygrad.helpers import all_same, getenv, flatten, Target, IMAGE, is_image_shape, cpu_profile
from tinygrad.device import Buffer, Compiled, Compiler, Allocator, Program, TinyELF
from tinygrad.codegen.opt import tc
from tinygrad.uop.ops import exec_alu, python_alu, Ops, UOp, GroupOp, bitcast
from tinygrad.uop.ops import exec_alu, python_alu, Ops, UOp, GroupOp
from tinygrad.renderer import Renderer
def _load(m, i, dtype: DType):
@@ -39,16 +39,15 @@ def generic_wmma_helper(inp, warp_size, WARP_THREADS, K, NUM_A, NUM_B, NUM_C, a_
out[elem_idx][goff+lane_id] += sum(a_elem(inp[0], _k, c_j, goff) * b_elem(inp[1], c_i, _k, goff) for _k in range(K))
return out
class PythonProgram:
def __init__(self, name:str, lib:bytes, **kwargs):
self.uops: list[UOp] = pickle.loads(lib)
class PythonProgram(Program['PythonDevice']):
def __init__(self, dev:'PythonDevice', obj:TinyELF):
self.uops: list[UOp] = pickle.loads(obj.lib)
self.uop_to_index: dict[UOp, int] = {u:i for i,u in enumerate(self.uops)}
self.loop_ends: dict[UOp, int] = {u.src[1]:i for i, u in enumerate(self.uops) if u.op == Ops.END}
def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False, **kw):
st = time.perf_counter()
warp = list(itertools.product(*[range(x) for x in local_size[::-1]]))
warp_size = len(warp)
void_ops = {Ops.END, Ops.BARRIER, Ops.IF, Ops.ENDIF, Ops.SINK, Ops.NOOP, Ops.GROUP, Ops.STORE}
for idxs in itertools.product(*[range(x) for x in global_size[::-1]]):
values: dict[UOp, Any] = {}
pbufs: list[memoryview] = list(bufs)
@@ -57,11 +56,15 @@ class PythonProgram:
i = 0
while i < len(self.uops):
u = self.uops[i]
src_values = [values[v] for v in u.src if v.op not in void_ops]
src_dtypes = [v.dtype for v in u.src if v.op not in void_ops]
src_values = [values[v] for v in u.src if v.dtype is not dtypes.void]
src_dtypes = [v.dtype for v in u.src if v.dtype is not dtypes.void]
if getenv("TRACE"): print(i, u.op, u.dtype, u.arg, src_values, src_dtypes)
if u.op is Ops.END:
i = self.uop_to_index[u.src[1]]
if len(u.src) == 3:
# conditional backedge on a loop: jump back while the condition is true
if values[u.src[2]][0]: i = self.uop_to_index[u.src[1]]
else: i += 1
else: i = self.uop_to_index[u.src[1]]
continue
if u.op is Ops.IF:
exec_masks.append([x and y for x,y in zip(exec_masks[-1], src_values[0])])
@@ -71,7 +74,7 @@ class PythonProgram:
exec_masks.pop()
i += 1
continue
if u.op in (Ops.BARRIER, Ops.SINK, Ops.NOOP, Ops.GROUP):
if u.op in (Ops.BARRIER, Ops.SINK, Ops.NOOP, Ops.GROUP) or (u.op is Ops.RANGE and u.dtype == dtypes.void):
# in the python emulator, the warp is always in sync
i += 1
continue
@@ -81,13 +84,13 @@ class PythonProgram:
store_gate = exec_masks[-1]
for j,val in enumerate(src_values[1] if u.max_numel() > 1 else [src_values[1]]):
for (m,o),v,g in zip(src_values[0], val, store_gate):
if g: _store(m, o+j, v, src_dtypes[1].scalar())
if g: _store(m, o+j, v, src_dtypes[1])
i += 1
continue
if u.op is Ops.AFTER: values[u] = src_values[0]
elif u.op is Ops.PARAM and u.addrspace is None: values[u] = [pvals.pop(0)] * warp_size
elif u.op is Ops.PARAM and u.addrspace is AddrSpace.ALU: values[u] = [pvals.pop(0)] * warp_size
elif u.op in {Ops.PARAM, Ops.BUFFER}:
storage_fmt = storage_fmt_for_dtype(u.dtype.base.scalar())
storage_fmt = storage_fmt_for_dtype(u.dtype)
if storage_fmt is None: raise RuntimeError(f"dtype={u.dtype} is not supported")
if TYPE_CHECKING or sys.version_info < (3, 12): assert storage_fmt != "e"
if u.addrspace == AddrSpace.REG:
@@ -99,16 +102,15 @@ class PythonProgram:
elif u.op is Ops.SPECIAL:
if u.arg[0] == 'g': values[u] = [idxs[2-int(u.arg[-1])]] * warp_size
elif u.arg[0] == 'l': values[u] = [x[2-int(u.arg[-1])] for x in warp]
elif u.op is Ops.CONST: values[u] = [u.arg] * warp_size
elif u.op is Ops.CONST: values[u] = [u.val] * warp_size
elif u.op in {Ops.INDEX, Ops.SHRINK}:
ret:list = []
if u.src[0].addrspace == AddrSpace.ALU:
ret = [src_values[0][i][t] for t,i in enumerate(src_values[1])]
elif isinstance(src_dtypes[0], ImageDType):
assert len(src_values) == 3, "image index must be 3 srcs"
elif is_image_shape(u.src[0]._shape):
for m,oy,ox in zip(*src_values):
if ox < 0 or ox >= src_dtypes[0].shape[1] or oy < 0 or oy >= src_dtypes[0].shape[0]: ret.append((m, None))
else: ret.append((m, ox*4 + oy*src_dtypes[0].shape[1]*4))
if ox < 0 or ox >= u.src[0]._shape[1] or oy < 0 or oy >= u.src[0]._shape[0]: ret.append((m, None))
else: ret.append((m, ox*4 + oy*u.src[0]._shape[1]*4))
else:
for m,o in zip(src_values[0], src_values[1]): ret.append((m,o))
values[u] = ret
@@ -129,13 +131,13 @@ class PythonProgram:
if (load_sz := u.max_numel()) > 1:
# buf and gate are not vecs
values[u] = [load([src_values[k] if k in [0,2] else src_values[k][j] \
for k in range(len(src_values))], j, u.dtype.scalar()) for j in range(load_sz)]
for k in range(len(src_values))], j, u.dtype) for j in range(load_sz)]
else:
values[u] = load(src_values, 0, u.dtype)
elif u.op is Ops.WMMA:
first_src_dtype = u.src[0].dtype
assert isinstance(first_src_dtype, DType) # mypy
dims, dtype_in, device, threads = u.arg[1], first_src_dtype.scalar(), u.arg[4], u.arg[5]
dims, dtype_in, device, threads = u.arg[0], first_src_dtype, u.arg[2], u.arg[3]
wmma_helper = functools.partial(generic_wmma_helper, src_values, warp_size)
# TODO: refactor these to a shared TensorCoreLayout
if device == "METAL":
@@ -224,8 +226,13 @@ class PythonRenderer(Renderer):
class PythonAllocator(Allocator['PythonDevice']):
def _alloc(self, size, options): return memoryview(bytearray(size))
def _copyin(self, dest, src:memoryview): dest[:] = src
def _copyout(self, dest:memoryview, src): dest[:] = src
def _as_buffer(self, src) -> memoryview: return src
def _copyin(self, dest, src:memoryview):
with cpu_profile("TINY -> PYTHON", f"{self.dev.device}:COPY"): dest[:] = src
def _copyout(self, dest:memoryview, src):
with cpu_profile("PYTHON -> TINY", f"{self.dev.device}:COPY"): dest[:] = src
def map(self, buf:Buffer): return buf.as_memoryview(force_zero_copy=True)
def _offset(self, buf:memoryview, size:int, offset:int): return buf[offset:offset+size]
class PythonDevice(Compiled):
def __init__(self, device:str):