IQ.Pilot Prebuilt Release @ 27f668a

artifacts/package_runtime/tinygrad/runtime/graph/cuda.py

@@ -0,0 +1,74 @@
+import ctypes
+from typing import Any, cast
+import tinygrad.runtime.autogen.cuda as cuda
+from tinygrad.runtime.support.c import init_c_var
+from tinygrad.device import Device, MultiBuffer
+from tinygrad.uop.ops import UOp, Ops
+from tinygrad.runtime.ops_cuda import CUDADevice, check, encode_args, cu_time_execution
+from tinygrad.engine.jit import MultiGraphRunner
+
+class CUDAGraph(MultiGraphRunner):
+  def __init__(self, linear, input_uops=()):
+    super().__init__(linear, input_uops)
+
+    self.nodes: list[tuple[Any, ...]] = [] # list of tuple(graph node, node params, c_args/context, is memcpy)
+    self.graph = init_c_var(cuda.CUgraph, lambda x: check(cuda.cuGraphCreate(ctypes.byref(x), 0)))
+
+    for (dev_idx, ast, bufs, device_vars), runtime in zip(self.calls, self.runtimes):
+      if ast.op is Ops.PROGRAM:
+        assert runtime is not None
+        global_size, local_size = ast.arg.launch_dims({v: 0 for v in self.vars})
+
+        c_deps, new_node = self.new_node([b.base for b in bufs], ast.arg.outs)
+        c_args, vargs = encode_args([b._buf for b in bufs], [device_vars.get(x.expr, 0) for x in ast.arg.vars], runtime.signature)
+        kern_params = cuda.CUDA_KERNEL_NODE_PARAMS_v1(runtime.prg, *global_size, *local_size, runtime.smem,
+                                                      ctypes.cast(0, ctypes.POINTER(ctypes.c_void_p)), vargs)
+        check(cuda.cuGraphAddKernelNode(ctypes.byref(new_node), self.graph, c_deps, len(c_deps or []), ctypes.byref(kern_params)))
+
+        self.nodes.append((new_node, kern_params, c_args, False))
+      elif ast.op is Ops.COPY:
+        dest, src = bufs[0], bufs[1]
+        src_dev = cast(CUDADevice, Device[src.device])
+        c_deps, new_node = self.new_node([dest.base, src.base], [0])
+        cp_params = cuda.CUDA_MEMCPY3D_v2(srcMemoryType=cuda.CU_MEMORYTYPE_DEVICE, srcDevice=src._buf, srcPitch=src.nbytes, srcHeight=1,
+                                          dstMemoryType=cuda.CU_MEMORYTYPE_DEVICE, dstDevice=dest._buf, dstPitch=dest.nbytes, dstHeight=1,
+                                          WidthInBytes=dest.nbytes, Height=1, Depth=1)
+        check(cuda.cuGraphAddMemcpyNode(ctypes.byref(new_node), self.graph, c_deps, len(c_deps or []), ctypes.byref(cp_params), src_dev.context))
+
+        self.nodes.append((new_node, cp_params, src_dev.context, True))
+
+    self.instance = init_c_var(cuda.CUgraphExec, lambda x: check(cuda.cuGraphInstantiate_v2(ctypes.byref(x), self.graph, None, None, 0)))
+    self.updatable = sorted({j for j,r in enumerate(self.uop_replace) if r} | self.var_vals_replace.keys() | self.launch_dims_replace.keys())
+
+  def new_node(self, bufs, write):
+    deps = self._access_resources(bufs, write, new_dependency=(node:=cuda.CUgraphNode()))
+    return (cuda.CUgraphNode*len(deps))(*deps) if deps else None, node
+
+  def __call__(self, input_uops:tuple[UOp, ...], var_vals:dict[str, int], wait=False):
+    # Update buffers in the c_args struct.
+    for j in self.updatable:
+      (_, params, c_args, is_copy), dev_idx = self.nodes[j], self.calls[j][0]
+      for pos, iidx in self.uop_replace[j]:
+        buf = b.bufs[dev_idx] if isinstance(b:=input_uops[iidx].buffer, MultiBuffer) else b
+        if not is_copy: setattr(c_args, f'f{pos}', buf._buf)
+        else: setattr(params, 'srcDevice' if pos == 1 else 'dstDevice', buf._buf)
+
+    # Update var_vals in the c_args struct.
+    for j, i, v in self.updated_vars(var_vals): setattr(self.nodes[j][2], f'v{i}', v)
+
+    # Update launch dims in the kern_params struct.
+    for j, global_dims, local_dims in self.updated_launch_dims(var_vals):
+      node = self.nodes[j][1]
+      node.blockDimX, node.blockDimY, node.blockDimZ, node.gridDimX, node.gridDimY, node.gridDimZ = *local_dims, *global_dims # type: ignore[misc]
+
+    # Update graph nodes with the updated structs.
+    for j in self.updatable:
+      node, c_node_params, c_args, is_copy = self.nodes[j]
+      if not is_copy: check(cuda.cuGraphExecKernelNodeSetParams(self.instance, node, ctypes.byref(c_node_params)))
+      else: check(cuda.cuGraphExecMemcpyNodeSetParams(self.instance, node, ctypes.byref(c_node_params), c_args))
+
+    return cu_time_execution(lambda: check(cuda.cuGraphLaunch(self.instance, None)), enable=wait)
+
+  def __del__(self):
+    if hasattr(self, 'graph'): check(cuda.cuGraphDestroy(self.graph))
+    if hasattr(self, 'instance'): check(cuda.cuGraphExecDestroy(self.instance))