diff --git a/docs/_source/tinygrad-features-graph-hip.rst b/docs/_source/tinygrad-features-graph-hip.rst new file mode 100644 index 000000000..f93637eb8 --- /dev/null +++ b/docs/_source/tinygrad-features-graph-hip.rst @@ -0,0 +1,17 @@ +tinygrad features.graph.hip +--------------------------- + +.. note:: You likely want the upstream tinygrad, not tinygrab. + Tinygrab contains AI generated docstrings for a tinygrad snapshot. + Upstream: https://tinygrad.org + + +Fails to build from source unless ROCm is installed: + +``OSError: /opt/rocm/lib/libhiprtc.so: cannot open shared object file: No such file or directory`` + +.. automodule:: tinygrad.features.graph.hip + :members: + :undoc-members: + :show-inheritance: + diff --git a/docs/_source/tinygrad-runtime-ops_hip.rst b/docs/_source/tinygrad-runtime-ops_hip.rst index d89e44558..f781743a0 100644 --- a/docs/_source/tinygrad-runtime-ops_hip.rst +++ b/docs/_source/tinygrad-runtime-ops_hip.rst @@ -1,15 +1,16 @@ tinygrad runtime.ops_hip ------------------------ +.. note:: You likely want the upstream tinygrad, not tinygrab. + Tinygrab contains AI generated docstrings for a tinygrad snapshot. + Upstream: https://tinygrad.org + + Fails to build from source unless ROCm is installed: ``OSError: /opt/rocm/lib/libhiprtc.so: cannot open shared object file: No such file or directory`` -.. 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_hip :members: :undoc-members: diff --git a/docs/_source/tinygrad.rst b/docs/_source/tinygrad.rst index 9528193bd..31013375b 100644 --- a/docs/_source/tinygrad.rst +++ b/docs/_source/tinygrad.rst @@ -24,6 +24,7 @@ tinygrad tinygrad-features-image tinygrad-features-search tinygrad-features-graph-cuda + tinygrad-features-graph-hip tinygrad-runtime-ops_clang tinygrad-runtime-ops_cpu tinygrad-runtime-ops_cuda @@ -31,7 +32,6 @@ tinygrad tinygrad-runtime-ops_gpu tinygrad-runtime-ops_hip tinygrad-runtime-ops_llvm - tinygrad-runtime-ops_metal tinygrad-runtime-ops_torch tinygrad-runtime-ops_webgpu :maxdepth: 3 diff --git a/tinygrad/features/graph/hip.py b/tinygrad/features/graph/hip.py index 82522f42f..379644d09 100644 --- a/tinygrad/features/graph/hip.py +++ b/tinygrad/features/graph/hip.py @@ -7,19 +7,53 @@ from tinygrad.features.graph.cuda import CUDAGraph class HIPGraph(CUDAGraph): + """ + The HIPGraph class, a subclass of CUDAGraph. + + Attributes: + graph (hip.hipGraph_t): The graph object. + instance (hip.hipGraphExec_t): The graph execution object. + """ + def __del__(self): + """ + Destructor for the HIPGraph class. + + This method is called when an instance of the HIPGraph class is about to be destroyed. It destroys the associated graph and graph execution objects. + """ check(hip.hipGraphDestroy(self.graph)) check(hip.hipGraphExecDestroy(self.instance)) def encode_args_info(self): + """ + Encode argument information for the HIPGraph object. + + Returns: + tuple: A tuple containing a hipDeviceptr_t object and a tuple of integers (1, 2, 3). + """ return (hip.hipDeviceptr_t, (1, 2, 3)) def graph_create(self): + """ + Create the HIP graph. + + Returns: + hip.hipGraph_t: The created graph object. + """ return init_c_var( hip.hipGraph_t(), lambda x: check(hip.hipGraphCreate(ctypes.byref(x), 0)) ) def graph_instantiate(self, graph): + """ + Instantiate the HIP graph. + + Args: + graph (hip.hipGraph_t): The graph object to instantiate. + + Returns: + hip.hipGraphExec_t: The instantiated graph execution object. + """ return init_c_var( hip.hipGraphExec_t(), lambda x: check( @@ -28,6 +62,17 @@ class HIPGraph(CUDAGraph): ) def graph_add_kernel_node(self, graph, c_deps, c_params): + """ + Add a kernel node to the HIP graph. + + Args: + graph (hip.hipGraph_t): The graph object. + c_deps (ctypes array): The dependencies for the kernel. + c_params (ctypes pointer): The parameters for the kernel. + + Returns: + hip.hipGraphNode_t: The created graph node object. + """ return init_c_var( hip.hipGraphNode_t(), lambda x: check( @@ -42,12 +87,46 @@ class HIPGraph(CUDAGraph): ) def graph_launch(self, *args, wait=False): + """ + Launch the HIP graph execution. + + Args: + *args: Variable length argument list. + wait (bool, optional): Whether to wait for completion or not. Defaults to False. + """ return hip_time_execution(lambda: check(hip.hipGraphLaunch(*args)), enable=wait) def graph_exec_kernel_node_set_params(self, *args): + """Set the parameters of a kernel node in the graph execution. + + This function sets the parameters for a given kernel node within a hipGraphExec. + It takes any number of arguments and passes them directly to the + `hip.hipGraphExecKernelNodeSetParams` function. + + Attributes: + *args (tuple): Variable length argument list which are passed as is to the underlying function. + + Returns: + The result returned by the hip.hipGraphExecKernelNodeSetParams function. + """ return check(hip.hipGraphExecKernelNodeSetParams(*args)) def build_kernel_node_params(self, prg, global_size, local_size, c_config): + """Build the parameters for a kernel node. + + This function constructs and returns a hip.hipKernelNodeParams object which holds + all necessary parameters to launch a HIP kernel. It takes a program object, + global and local sizes as tuples, and a c_config parameter. + + Attributes: + prg (object): The program object containing the kernel code. + global_size (tuple): A tuple representing the global size of the grid for kernel launch. + local_size (tuple): A tuple representing the local size of each thread block for kernel launch. + c_config (object): An object holding the configuration parameters for the kernel launch. + + Returns: + The constructed hipKernelNodeParams object with all necessary parameters. + """ return hip.hipKernelNodeParams( hip.dim3(*local_size), c_config, @@ -60,6 +139,19 @@ class HIPGraph(CUDAGraph): def set_kernel_node_launch_dims( self, node, global_size: Tuple[int, int, int], local_size: Tuple[int, int, int] ): + """Set the launch dimensions for a kernel node. + + This function sets the launch dimensions (blockDim and gridDim) for a given + HIP kernel node. It takes a node object, global size as a tuple, and local size as a tuple. + + Attributes: + node (object): The node object representing the kernel to be launched. + global_size (tuple): A tuple representing the global size of the grid for kernel launch. + local_size (tuple): A tuple representing the local size of each thread block for kernel launch. + + Returns: + Nothing. The launch dimensions are set directly on the node object. + """ ( node.blockDim.x, node.blockDim.y,