ops_clang, ops_cpu docstrings

docs/_source/conf.py

@@ -3,6 +3,12 @@
 import sys, os
 import sphinx.util.logging
 sys.path.append('../../tinygrad')
+sys.path.append('../../tinygrad/codegen')
+sys.path.append('../../tinygrad/features')
+sys.path.append('../../tinygrad/nn')
+sys.path.append('../../tinygrad/renderer')
+sys.path.append('../../tinygrad/runtime')
+sys.path.append('../../tinygrad/shape')
 
 logger = sphinx.util.logging.getLogger(__name__)
 

docs/_source/tinygrad-runtime-ops_clang.rst

@@ -0,0 +1,8 @@
+tinygrad runtime.ops_clang
+--------------------------
+
+.. automodule:: tinygrad.runtime.ops_clang
+   :members:
+   :undoc-members:
+   :show-inheritance:
+

docs/_source/tinygrad-runtime-ops_cpu.rst

@@ -0,0 +1,8 @@
+tinygrad runtime.ops_cpu
+------------------------
+
+.. automodule:: tinygrad.runtime.ops_cpu
+   :members:
+   :undoc-members:
+   :show-inheritance:
+

docs/_source/tinygrad.rst

@@ -14,6 +14,8 @@ tinygrad
    tinygrad-ops
    tinygrad-realize
    tinygrad-tensor
+   tinygrad-runtime-ops_clang
+   tinygrad-runtime-ops_cpu
    :maxdepth: 3
    :caption: Contents:
 

tinygrad/runtime/ops_clang.py

@@ -10,7 +10,19 @@ CLANG_PROGRAM_HEADER = "#include <math.h>\n#define max(x,y) ((x>y)?x:y)\n#define
 
 @diskcache
 def compile_clang(prg: str, header: str = CLANG_PROGRAM_HEADER) -> bytes:
-    # TODO: remove file write. sadly clang doesn't like the use of /dev/stdout here
+    """
+    Compile a given C program using clang.
+
+    Parameters:
+        prg (str): The C program to be compiled.
+        header (str): The header for the C program. Default is CLANG_PROGRAM_HEADER.
+
+    Returns:
+        bytes: The compiled C program in bytes format.
+
+    Note:
+        Currently, there's a TODO to remove file write due to clang not supporting /dev/stdout usage.
+    """
     with tempfile.NamedTemporaryFile(delete=True) as output_file:
         subprocess.check_output(
             args=(
@@ -23,6 +35,18 @@ def compile_clang(prg: str, header: str = CLANG_PROGRAM_HEADER) -> bytes:
 
 
 class ClangProgram:
+    """
+    A class representing a compiled Clang program.
+
+    Attributes:
+        name (str): The name of the Clang program.
+        lib (bytes): The compiled Clang program in bytes format.
+        fxn (Any): The function object of the compiled Clang program.
+
+    Note:
+        Writes to disk for loading the compiled program.
+    """
+
     def __init__(self, name: str, lib: bytes):
         self.name, self.lib = name, lib
         # write to disk so we can load it
@@ -31,6 +55,17 @@ class ClangProgram:
             self.fxn: Any = ctypes.CDLL(str(cached_file_path.name))[name]
 
     def __call__(self, *bufs, vals=(), wait=False):
+        """
+        Call the Clang program with given buffers and values.
+
+        Parameters:
+            bufs (*): Variable length buffer arguments.
+            vals (tuple): Tuple of constant integer values. Default is an empty tuple.
+            wait (bool): If True, enables waiting for the call to complete. Default is False.
+
+        Returns:
+            Any: The result of the Clang program execution.
+        """
         return cpu_time_execution(lambda: self.fxn(*bufs, *vals), enable=wait)
 
 

tinygrad/runtime/ops_cpu.py

@@ -16,12 +16,35 @@ from tinygrad.device import Interpreted, Allocator
 def shape_to_axis(
     old_shape: Tuple[int, ...], new_shape: Tuple[int, ...]
 ) -> Tuple[int, ...]:
+    """
+    Compare two shapes and return a tuple containing the indices of axes that differ between the two shapes.
+
+    Parameters:
+        old_shape (Tuple[int, ...]): The original shape to be compared.
+        new_shape (Tuple[int, ...]): The new shape to compare against the original shape.
+
+    Returns:
+        Tuple[int, ...]: A tuple of indices where the axes differ between the two shapes.
+
+    Raises:
+        AssertionError: If the dimensions of old_shape and new_shape are not equal.
+    """
     assert len(old_shape) == len(new_shape), "reduce shapes must have same dimensions"
     return tuple(i for i, (a, b) in enumerate(zip(old_shape, new_shape)) if a != b)
 
 
 # TODO: this should be global infrastructure
 def output_type(x, y):
+    """
+    Determine the datatype with higher priority between two numpy arrays.
+
+    Parameters:
+        x (numpy.ndarray): The first numpy array to compare.
+        y (numpy.ndarray): The second numpy array to compare.
+
+    Returns:
+        numpy.dtype: The datatype with higher priority between x and y.
+    """
     return (
         x.dtype
         if dtypes.from_np(x.dtype).priority > dtypes.from_np(y.dtype).priority
@@ -30,20 +53,69 @@ def output_type(x, y):
 
 
 def match_types(x, y):
+    """
+    Cast two numpy arrays to a common datatype determined by output_type() function and return the casted arrays.
+
+    Parameters:
+        x (numpy.ndarray): The first numpy array for casting.
+        y (numpy.ndarray): The second numpy array for casting.
+
+    Returns:
+        Tuple[numpy.ndarray, numpy.ndarray]: A tuple of the casted x and y arrays.
+    """
     up = output_type(x, y)
     return x.astype(up, copy=False), y.astype(up, copy=False)
 
 
 def einsum_mulacc(einsum, get_strides, expand):
+    """
+    This function returns a higher-order function named `mulacc`. The returned function performs multiplication and accumulation using the numpy.einsum function. It takes two arrays as input along with their shapes and strides.
+
+    Attributes:
+        einsum (function): A function that performs numpy einsum operation.
+        get_strides (function): Function to calculate the strides of an array.
+        expand (function): Function to perform broadcasting/expansion of arrays.
+
+    """
+
     def einscripts(x):
+        """
+        This function converts a list/array of integers into a string where each integer is mapped to a lowercase English alphabet in order. For example, [0, 1, 2] would be converted to 'abc'.
+
+        Args:
+            x (list or array): A list or array of integers.
+
+        Returns:
+            str: A string created by mapping each integer in `x` to a lowercase English alphabet.
+        """
         return "".join(["abcdefghijklmnopqrstuvwxyz"[i] for i in x])
 
     def axes_slice(strides):
+        """
+        This function returns two lists: one containing the indices of non-zero strides and another containing slices for those indices. The input is a list/array of strides.
+
+        Args:
+            strides (list or array): A list or array of integers representing strides.
+
+        Returns:
+            list, tuple: Two lists: first one contains the indices of non-zero strides and second one contains slices for those indices.
+        """
         return [i for i, s in enumerate(strides) if s != 0], tuple(
             [slice(None) if s != 0 else 0 for i, s in enumerate(strides)]
         )
 
     def mulacc(a, b, new_shape):
+        """
+        This function performs multiplication and accumulation using the numpy.einsum function on two arrays `a` and `b`. It also takes their strides and a new shape as input.
+
+        Args:
+            a (array): The first array.
+            b (array): The second array.
+            new_shape (tuple): The desired output shape.
+
+        Returns:
+            array: The result of multiplication and accumulation operation on `a` and `b`.
+        """
         (a_axes, a_slices), (b_axes, b_slices) = axes_slice(get_strides(a)), axes_slice(
             get_strides(b)
         )
@@ -119,16 +191,58 @@ numpy_fxn_for_op: Dict[Op, Callable] = {
 
 
 class NumpyAllocator(Allocator):
+    """
+    Allocator class for numpy arrays.
+
+    Attributes:
+        _alloc (method): Allocates memory for a given size and dtype.
+        as_buffer (method): Converts an np.ndarray to a memoryview object.
+        copyin (method): Copies data from a memoryview object to an np.ndarray.
+        copyout (method): Copies data from an np.ndarray to a memoryview object.
+    """
+
     def _alloc(self, size: int):
+        """
+        Allocates memory for a given size and dtype.
+
+        Parameters:
+            size (int): Size of the array.
+
+        Returns:
+            np.ndarray: Empty numpy array with specified size and dtype=np.uint8.
+        """
         return np.empty(size, dtype=np.uint8)
 
     def as_buffer(self, src: np.ndarray) -> memoryview:
+        """
+        Converts an np.ndarray to a memoryview object.
+
+        Parameters:
+            src (np.ndarray): The numpy array to be converted.
+
+        Returns:
+            memoryview: A view into the original numpy array's data.
+        """
         return flat_mv(np.require(src, requirements="C").data)
 
     def copyin(self, dest: np.ndarray, src: memoryview):
+        """
+        Copies data from a memoryview object to an np.ndarray.
+
+        Parameters:
+            dest (np.ndarray): The destination numpy array.
+            src (memoryview): The source memoryview object.
+        """
         np.copyto(dest, np.frombuffer(src, dest.dtype).reshape(dest.shape))
 
     def copyout(self, dest: memoryview, src: np.ndarray):
+        """
+        Copies data from an np.ndarray to a memoryview object.
+
+        Parameters:
+            dest (memoryview): The destination memoryview object.
+            src (np.ndarray): The source numpy array.
+        """
         np.copyto(np.frombuffer(dest, src.dtype).reshape(src.shape), src)