docstrings, requirements webgpu

docs/_source/tinygrad-runtime-ops_webgpu.rst

@@ -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:
+

docs/_source/tinygrad.rst

@@ -19,6 +19,7 @@ tinygrad
    tinygrad-ops
    tinygrad-realize
    tinygrad-tensor
+   tinygrad-codegen-kernel
    tinygrad-runtime-ops_clang
    tinygrad-runtime-ops_cpu
    tinygrad-runtime-ops_cuda
@@ -28,6 +29,7 @@ tinygrad
    tinygrad-runtime-ops_llvm
    tinygrad-runtime-ops_metal
    tinygrad-runtime-ops_torch
+   tinygrad-runtime-ops_webgpu
    :maxdepth: 3
    :caption: Contents:
 

requirements-docs.txt

@@ -11,3 +11,4 @@ sphinx-notfound-page
 black
 llvmlite
 torch
+wgpu

tinygrad/runtime/ops_webgpu.py

@@ -8,7 +8,26 @@ wgpu_device = get_default_device()
 
 
 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):
+        """
+        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 = (
             name,
             lib,
@@ -16,6 +35,16 @@ class WebGPUProgram:
         )  # NOTE: this is the compiler
 
     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"
         binding_layouts = [
             {
@@ -52,7 +81,23 @@ class WebGPUProgram:
 
 
 class WebGpuAllocator(Allocator):
+    """
+    WebGpuAllocator class.
+
+    Attributes:
+        Allocator (parent class): Parent class for this class.
+    """
+
     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(
             size=size,
             usage=wgpu.BufferUsage.STORAGE
@@ -61,14 +106,47 @@ class WebGpuAllocator(Allocator):
         )
 
     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)
 
     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
 
 
 class WebGpuDevice(Compiled):
+    """
+    WebGpuDevice class.
+
+    Attributes:
+        Compiled (parent class): Parent class for this class.
+    """
+
     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__(
             WebGpuAllocator(),
             LinearizerOptions(