tinygrad features.graph.cuda

Note

You likely want the upstream tinygrad, not tinygrab. Tinygrab contains AI generated docstrings for a tinygrad snapshot. Upstream: https://tinygrad.org

class tinygrad.features.graph.cuda.CUDAGraph(jit_cache: List[JitItem], input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int])[source]

Bases: object

This class represents a CUDA Graph. It is initialized with a list of JitItems, input raw buffers, and a dictionary mapping variables to integers.

jit_cache

A list of JitItems.

Type:: List[JitItem]

input_rawbuffers

A list of input raw buffers.

Type:: List[Buffer]

var_vals

A dictionary mapping variables to integers.

Type:: Dict[Variable, int]

build_kernel_node_params(prg, global_size, local_size, c_kernel_config)[source]

Build parameters for the kernel node.

Parameters:

prg (Program) – CUDA program object.
global_size (Tuple[int, int, int]) – Global size of thread blocks.
local_size (Tuple[int, int, int]) – Local size of thread blocks.
c_kernel_config – Kernel configuration parameters.

Returns:

CUDA_KERNEL_NODE_PARAMS object with specified parameters.

encode_args_info()[source]

Encodes the arguments information for the class.

This method encodes the arguments information required by other methods in the class.

Returns:: A tuple containing the encoded arguments information.
Return type:: tuple

graph_add_kernel_node(graph, c_deps, c_node_params)[source]

Adds a kernel node to the cuGraph object.

This method adds a new kernel node to the given cuGraph object.

Parameters:

graph (cuda.CUgraph) – The cuGraph object to add the kernel node to.
c_deps (ctypes.c_void_p) – The dependencies of the kernel node.
c_node_params (ctypes.c_void_p) – The parameters for the new kernel node.

Returns:

The newly created kernel node.

Return type:

cuda.CUgraphNode

graph_create()[source]

Creates a new cuGraph object.

This method initializes and creates a new cuGraph object.

Returns:: The newly created cuGraph object.
Return type:: cuda.CUgraph

graph_exec_kernel_node_set_params(*args)[source]

Set parameters for the kernel node.

Parameters:: *args – Variable length argument list.
Returns:: Result of check function with CUDA call.

graph_instantiate(graph)[source]

Instantiates a cuGraphExec object from a cuGraph object.

This method creates an instance of the cuGraph execution from a given cuGraph object.

Parameters:: graph (cuda.CUgraph) – The cuGraph object to instantiate.
Returns:: The instantiated cuGraph execution object.
Return type:: cuda.CUgraphExec

graph_launch(*args, wait=False)[source]

Launches the cuGraph execution.

This method launches the execution of the cuGraph object.

Parameters:

*args (tuple) – The arguments for the graph execution.
wait (bool, optional) – Whether to wait for the execution to finish or not. Defaults to False.

set_kernel_node_launch_dims(node, global_size: Tuple[int, int, int], local_size: Tuple[int, int, int])[source]

Set launch dimensions for the kernel node.

Parameters:

node – Kernel node object.
global_size (Tuple[int, int, int]) – Global size of thread blocks.
local_size (Tuple[int, int, int]) – Local size of thread blocks.

Returns:

None

node.blockDimX, node.blockDimY, node.blockDimZ: Local dimensions.

node.gridDimX, node.gridDimY, node.gridDimZ: Global dimensions.