docstrings, requirements webgpu
parent
96fb6d334b
commit
fe8d9753f0
|
@ -0,0 +1,12 @@
|
||||||
|
tinygrad runtime.ops_webgpu
|
||||||
|
---------------------------
|
||||||
|
|
||||||
|
.. note:: You likely want the upstream tinygrad, not tinygrab.
|
||||||
|
Tinygrab contains AI generated docstrings for a tinygrad snapshot.
|
||||||
|
Upstream: https://tinygrad.org
|
||||||
|
|
||||||
|
.. automodule:: tinygrad.runtime.ops_webgpu
|
||||||
|
:members:
|
||||||
|
:undoc-members:
|
||||||
|
:show-inheritance:
|
||||||
|
|
|
@ -19,6 +19,7 @@ tinygrad
|
||||||
tinygrad-ops
|
tinygrad-ops
|
||||||
tinygrad-realize
|
tinygrad-realize
|
||||||
tinygrad-tensor
|
tinygrad-tensor
|
||||||
|
tinygrad-codegen-kernel
|
||||||
tinygrad-runtime-ops_clang
|
tinygrad-runtime-ops_clang
|
||||||
tinygrad-runtime-ops_cpu
|
tinygrad-runtime-ops_cpu
|
||||||
tinygrad-runtime-ops_cuda
|
tinygrad-runtime-ops_cuda
|
||||||
|
@ -28,6 +29,7 @@ tinygrad
|
||||||
tinygrad-runtime-ops_llvm
|
tinygrad-runtime-ops_llvm
|
||||||
tinygrad-runtime-ops_metal
|
tinygrad-runtime-ops_metal
|
||||||
tinygrad-runtime-ops_torch
|
tinygrad-runtime-ops_torch
|
||||||
|
tinygrad-runtime-ops_webgpu
|
||||||
:maxdepth: 3
|
:maxdepth: 3
|
||||||
:caption: Contents:
|
:caption: Contents:
|
||||||
|
|
||||||
|
|
|
@ -11,3 +11,4 @@ sphinx-notfound-page
|
||||||
black
|
black
|
||||||
llvmlite
|
llvmlite
|
||||||
torch
|
torch
|
||||||
|
wgpu
|
||||||
|
|
|
@ -8,7 +8,26 @@ wgpu_device = get_default_device()
|
||||||
|
|
||||||
|
|
||||||
class WebGPUProgram:
|
class WebGPUProgram:
|
||||||
|
"""
|
||||||
|
This class represents a WebGPU program. It stores the name, library, and compiled shader module of a GPU program.
|
||||||
|
|
||||||
|
Attributes:
|
||||||
|
name (str): The name of the GPU program.
|
||||||
|
lib (bytes): The library containing the code for the GPU program.
|
||||||
|
prg (wgpu_device.create_shader_module): The compiled shader module of the GPU program.
|
||||||
|
"""
|
||||||
|
|
||||||
def __init__(self, name: str, lib: bytes):
|
def __init__(self, name: str, lib: bytes):
|
||||||
|
"""
|
||||||
|
Constructs a WebGPUProgram object.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
name (str): The name of the GPU program.
|
||||||
|
lib (bytes): The library containing the code for the GPU program.
|
||||||
|
|
||||||
|
Notes:
|
||||||
|
This is the compiler for the GPU program.
|
||||||
|
"""
|
||||||
self.name, self.lib, self.prg = (
|
self.name, self.lib, self.prg = (
|
||||||
name,
|
name,
|
||||||
lib,
|
lib,
|
||||||
|
@ -16,6 +35,16 @@ class WebGPUProgram:
|
||||||
) # NOTE: this is the compiler
|
) # NOTE: this is the compiler
|
||||||
|
|
||||||
def __call__(self, *bufs, global_size, local_size, vals=(), wait=False):
|
def __call__(self, *bufs, global_size, local_size, vals=(), wait=False):
|
||||||
|
"""
|
||||||
|
Executes the GPU program.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
*bufs (tuple of buffers): The input buffers for the GPU program.
|
||||||
|
global_size (tuple): The size of the global workgroup.
|
||||||
|
local_size (tuple): The size of the local workgroup.
|
||||||
|
vals (tuple, optional): Additional values to pass to the GPU program. Defaults to empty tuple.
|
||||||
|
wait (bool, optional): Whether or not to wait for the execution to finish before returning. Defaults to False.
|
||||||
|
"""
|
||||||
assert len(bufs) <= 8, "WEBGPU only supports 8 buffers"
|
assert len(bufs) <= 8, "WEBGPU only supports 8 buffers"
|
||||||
binding_layouts = [
|
binding_layouts = [
|
||||||
{
|
{
|
||||||
|
@ -52,7 +81,23 @@ class WebGPUProgram:
|
||||||
|
|
||||||
|
|
||||||
class WebGpuAllocator(Allocator):
|
class WebGpuAllocator(Allocator):
|
||||||
|
"""
|
||||||
|
WebGpuAllocator class.
|
||||||
|
|
||||||
|
Attributes:
|
||||||
|
Allocator (parent class): Parent class for this class.
|
||||||
|
"""
|
||||||
|
|
||||||
def _alloc(self, size: int):
|
def _alloc(self, size: int):
|
||||||
|
"""
|
||||||
|
Allocate memory on the device.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
size (int): Size of memory to be allocated.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Memory buffer created by wgpu_device.create_buffer().
|
||||||
|
"""
|
||||||
return wgpu_device.create_buffer(
|
return wgpu_device.create_buffer(
|
||||||
size=size,
|
size=size,
|
||||||
usage=wgpu.BufferUsage.STORAGE
|
usage=wgpu.BufferUsage.STORAGE
|
||||||
|
@ -61,14 +106,47 @@ class WebGpuAllocator(Allocator):
|
||||||
)
|
)
|
||||||
|
|
||||||
def copyin(self, dest, src: memoryview):
|
def copyin(self, dest, src: memoryview):
|
||||||
|
"""
|
||||||
|
Copy data from source to destination.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
dest: Destination of the data.
|
||||||
|
src (memoryview): Source of the data.
|
||||||
|
"""
|
||||||
wgpu_device.queue.write_buffer(dest, 0, src)
|
wgpu_device.queue.write_buffer(dest, 0, src)
|
||||||
|
|
||||||
def copyout(self, dest, src: memoryview):
|
def copyout(self, dest, src: memoryview):
|
||||||
|
"""
|
||||||
|
Copy data from source to destination.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
dest: Destination of the data.
|
||||||
|
src (memoryview): Source of the data.
|
||||||
|
|
||||||
|
Note:
|
||||||
|
This is a temporary solution and should be removed in the future.
|
||||||
|
"""
|
||||||
dest[:] = wgpu_device.queue.read_buffer(src, 0) # TODO: remove this copy
|
dest[:] = wgpu_device.queue.read_buffer(src, 0) # TODO: remove this copy
|
||||||
|
|
||||||
|
|
||||||
class WebGpuDevice(Compiled):
|
class WebGpuDevice(Compiled):
|
||||||
|
"""
|
||||||
|
WebGpuDevice class.
|
||||||
|
|
||||||
|
Attributes:
|
||||||
|
Compiled (parent class): Parent class for this class.
|
||||||
|
"""
|
||||||
|
|
||||||
def __init__(self, device: str):
|
def __init__(self, device: str):
|
||||||
|
"""
|
||||||
|
Initialize an instance of the WebGpuDevice class.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
device (str): Device identifier.
|
||||||
|
|
||||||
|
Note:
|
||||||
|
The WebGpuAllocator and LinearizerOptions classes are also initialized here.
|
||||||
|
"""
|
||||||
super().__init__(
|
super().__init__(
|
||||||
WebGpuAllocator(),
|
WebGpuAllocator(),
|
||||||
LinearizerOptions(
|
LinearizerOptions(
|
||||||
|
|
Loading…
Reference in New Issue