warp-lang 0.11.0__py3-none-manylinux2014_x86_64.whl → 1.0.0__py3-none-manylinux2014_x86_64.whl

warp/jax_experimental.py

@@ -0,0 +1,339 @@
+# Copyright (c) 2024 NVIDIA CORPORATION.  All rights reserved.
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import ctypes
+import warp as wp
+from warp.types import array_t, launch_bounds_t, strides_from_shape
+from warp.context import type_str
+import jax
+import jax.numpy as jp
+
+_jax_warp_p = None
+
+# Holder for the custom callback to keep it alive.
+_cc_callback = None
+_registered_kernels = [None]
+_registered_kernel_to_id = {}
+
+
+def jax_kernel(wp_kernel):
+    """Create a Jax primitive from a Warp kernel.
+
+    NOTE: This is an experimental feature under development.
+
+    Current limitations:
+    - All kernel arguments must be arrays.
+    - Kernel launch dimensions are inferred from the shape of the first argument.
+    - Input arguments are followed by output arguments in the Warp kernel definition.
+    - There must be at least one input argument and at least one output argument.
+    - Output shapes must match the launch dimensions (i.e., output shapes must match the shape of the first argument).
+    - All arrays must be contiguous.
+    - Only the CUDA backend is supported.
+    """
+
+    if _jax_warp_p == None:
+        # Create and register the primitive
+        _create_jax_warp_primitive()
+    if not wp_kernel in _registered_kernel_to_id:
+        id = len(_registered_kernels)
+        _registered_kernels.append(wp_kernel)
+        _registered_kernel_to_id[wp_kernel] = id
+    else:
+        id = _registered_kernel_to_id[wp_kernel]
+
+    def bind(*args):
+        return _jax_warp_p.bind(*args, kernel=id)
+
+    return bind
+
+
+def _warp_custom_callback(stream, buffers, opaque, opaque_len):
+    # The descriptor is the form
+    # <kernel-id>|<launch-dims>|<arg-dims-list>
+    # Example:  42|16,32|16,32;100;16,32
+    kernel_id_str, dim_str, args_str = opaque.decode().split("|")
+
+    # Get the kernel from the registry.
+    kernel_id = int(kernel_id_str)
+    kernel = _registered_kernels[kernel_id]
+
+    # Parse launch dimensions.
+    dims = [int(d) for d in dim_str.split(",")]
+    bounds = launch_bounds_t(dims)
+
+    # Parse arguments.
+    arg_strings = args_str.split(";")
+    num_args = len(arg_strings)
+    assert num_args == len(kernel.adj.args), "Incorrect number of arguments"
+
+    # First param is the launch bounds.
+    kernel_params = (ctypes.c_void_p * (1 + num_args))()
+    kernel_params[0] = ctypes.addressof(bounds)
+
+    # Parse array descriptors.
+    args = []
+    for i in range(num_args):
+        dtype = kernel.adj.args[i].type.dtype
+        shape = [int(d) for d in arg_strings[i].split(",")]
+        strides = strides_from_shape(shape, dtype)
+
+        arr = array_t(buffers[i], 0, len(shape), shape, strides)
+        args.append(arr)  # keep a reference
+        arg_ptr = ctypes.addressof(arr)
+
+        kernel_params[i + 1] = arg_ptr
+
+    # Get current device.
+    device = wp.device_from_jax(_get_jax_device())
+
+    # Get kernel hooks.
+    # Note: module was loaded during jit lowering.
+    hooks = kernel.module.get_kernel_hooks(kernel, device)
+    assert hooks.forward, "Failed to find kernel entry point"
+
+    # Launch the kernel.
+    wp.context.runtime.core.cuda_launch_kernel(
+        device.context, hooks.forward, bounds.size, 0, kernel_params, stream
+    )
+
+
+# TODO: is there a simpler way of getting the Jax "current" device?
+def _get_jax_device():
+    # check if jax.default_device() context manager is active
+    device = jax.config.jax_default_device
+    # if default device is not set, use first device
+    if device is None:
+        device = jax.devices()[0]
+    return device
+
+
+def _create_jax_warp_primitive():
+    from functools import reduce
+    import jax
+    from jax._src.interpreters import batching
+    from jax.interpreters import mlir
+    from jax.interpreters.mlir import ir
+    from jaxlib.hlo_helpers import custom_call
+
+    global _jax_warp_p
+    global _cc_callback
+
+    # Create and register the primitive.
+    # TODO add default implementation that calls the kernel via warp.
+    _jax_warp_p = jax.core.Primitive("jax_warp")
+    _jax_warp_p.multiple_results = True
+
+    # TODO Just launch the kernel directly, but make sure the argument
+    # shapes are massaged the same way as below so that vmap works.
+    def impl(*args):
+        raise Exception("Not implemented")
+
+    _jax_warp_p.def_impl(impl)
+
+    # Auto-batching. Make sure all the arguments are fully broadcasted
+    # so that Warp is not confused about dimensions.
+    def vectorized_multi_batcher(args, dims, **params):
+        # Figure out the number of outputs.
+        wp_kernel = _registered_kernels[params["kernel"]]
+        output_count = len(wp_kernel.adj.args) - len(args)
+        shape, dim = next((a.shape, d) for a, d in zip(args, dims) if d is not None)
+        size = shape[dim]
+        args = [batching.bdim_at_front(a, d, size) if len(a.shape) else a for a, d in zip(args, dims)]
+        # Create the batched primitive.
+        return _jax_warp_p.bind(*args, **params), [dims[0]] * output_count
+
+    batching.primitive_batchers[_jax_warp_p] = vectorized_multi_batcher
+
+    def get_vecmat_shape(warp_type):
+        if hasattr(warp_type.dtype, "_shape_"):
+            return warp_type.dtype._shape_
+        return []
+
+    def strip_vecmat_dimensions(warp_arg, actual_shape):
+        shape = get_vecmat_shape(warp_arg.type)
+        for i, s in enumerate(reversed(shape)):
+            item = actual_shape[-i - 1]
+            if s != item:
+                raise Exception(f"The vector/matrix shape for argument {warp_arg.label} does not match")
+        return actual_shape[: len(actual_shape) - len(shape)]
+
+    def collapse_into_leading_dimension(warp_arg, actual_shape):
+        if len(actual_shape) < warp_arg.type.ndim:
+            raise Exception(f"Argument {warp_arg.label} has too few non-matrix/vector dimensions")
+        index_rest = len(actual_shape) - warp_arg.type.ndim + 1
+        leading_size = reduce(lambda x, y: x * y, actual_shape[:index_rest])
+        return [leading_size] + actual_shape[index_rest:]
+
+    # Infer array dimensions from input type.
+    def infer_dimensions(warp_arg, actual_shape):
+        actual_shape = strip_vecmat_dimensions(warp_arg, actual_shape)
+        return collapse_into_leading_dimension(warp_arg, actual_shape)
+
+    def base_type_to_jax(warp_dtype):
+        if hasattr(warp_dtype, "_wp_scalar_type_"):
+            return wp.jax.dtype_to_jax(warp_dtype._wp_scalar_type_)
+        return wp.jax.dtype_to_jax(warp_dtype)
+
+    def base_type_to_jax_ir(warp_dtype):
+        warp_to_jax_dict = {
+            wp.float16: ir.F16Type.get(),
+            wp.float32: ir.F32Type.get(),
+            wp.float64: ir.F64Type.get(),
+            wp.int8: ir.IntegerType.get_signless(8),
+            wp.int16: ir.IntegerType.get_signless(16),
+            wp.int32: ir.IntegerType.get_signless(32),
+            wp.int64: ir.IntegerType.get_signless(64),
+            wp.uint8: ir.IntegerType.get_unsigned(8),
+            wp.uint16: ir.IntegerType.get_unsigned(16),
+            wp.uint32: ir.IntegerType.get_unsigned(32),
+            wp.uint64: ir.IntegerType.get_unsigned(64),
+        }
+        if hasattr(warp_dtype, "_wp_scalar_type_"):
+            warp_dtype = warp_dtype._wp_scalar_type_
+        jax_dtype = warp_to_jax_dict.get(warp_dtype)
+        if jax_dtype is None:
+            raise TypeError(f"Invalid or unsupported data type: {warp_dtype}")
+        return jax_dtype
+
+    def base_type_is_compatible(warp_type, jax_ir_type):
+        jax_ir_to_warp = {
+            "f16": wp.float16,
+            "f32": wp.float32,
+            "f64": wp.float64,
+            "i8": wp.int8,
+            "i16": wp.int16,
+            "i32": wp.int32,
+            "i64": wp.int64,
+            "ui8": wp.uint8,
+            "ui16": wp.uint16,
+            "ui32": wp.uint32,
+            "ui64": wp.uint64,
+        }
+        expected_warp_type = jax_ir_to_warp.get(str(jax_ir_type))
+        if expected_warp_type is not None:
+            if hasattr(warp_type, "_wp_scalar_type_"):
+                return warp_type._wp_scalar_type_ == expected_warp_type
+            else:
+                return warp_type == expected_warp_type
+        else:
+            raise TypeError(f"Invalid or unsupported data type: {jax_ir_type}")
+
+    # Abstract evaluation.
+    def jax_warp_abstract(*args, kernel=None):
+        wp_kernel = _registered_kernels[kernel]
+        # All the extra arguments to the warp kernel are outputs.
+        warp_outputs = [o.type for o in wp_kernel.adj.args[len(args) :]]
+        # TODO. Let's just use the first input dimension to infer the output's dimensions.
+        dims = strip_vecmat_dimensions(wp_kernel.adj.args[0], list(args[0].shape))
+        jax_outputs = []
+        for o in warp_outputs:
+            shape = list(dims) + list(get_vecmat_shape(o))
+            dtype = base_type_to_jax(o.dtype)
+            jax_outputs.append(jax.core.ShapedArray(shape, dtype))
+        return jax_outputs
+
+    _jax_warp_p.def_abstract_eval(jax_warp_abstract)
+
+    # Lowering to MLIR.
+
+    # Create python-land custom call target.
+    CCALLFUNC = ctypes.CFUNCTYPE(
+        ctypes.c_voidp, ctypes.c_void_p, ctypes.POINTER(ctypes.c_void_p), ctypes.c_char_p, ctypes.c_size_t
+    )
+    _cc_callback = CCALLFUNC(_warp_custom_callback)
+    ccall_address = ctypes.cast(_cc_callback, ctypes.c_void_p)
+
+    # Put the custom call into a capsule, as required by XLA.
+    PyCapsule_Destructor = ctypes.CFUNCTYPE(None, ctypes.py_object)
+    PyCapsule_New = ctypes.pythonapi.PyCapsule_New
+    PyCapsule_New.restype = ctypes.py_object
+    PyCapsule_New.argtypes = (ctypes.c_void_p, ctypes.c_char_p, PyCapsule_Destructor)
+    capsule = PyCapsule_New(ccall_address.value, b"xla._CUSTOM_CALL_TARGET", PyCapsule_Destructor(0))
+
+    # Register the callback in XLA.
+    jax.lib.xla_client.register_custom_call_target("warp_call", capsule, platform="gpu")
+
+    def default_layout(shape):
+        return range(len(shape) - 1, -1, -1)
+
+    def warp_call_lowering(ctx, *args, kernel=None):
+        if not kernel:
+            raise Exception("Unknown kernel id " + str(kernel))
+        wp_kernel = _registered_kernels[kernel]
+
+        # TODO This may not be necessary, but it is perhaps better not to be
+        # mucking with kernel loading while already running the workload.
+        module = wp_kernel.module
+        device = wp.device_from_jax(_get_jax_device())
+        if not module.load(device):
+            raise Exception("Could not load kernel on device")
+
+        # Infer dimensions from the first input.
+        warp_arg0 = wp_kernel.adj.args[0]
+        actual_shape0 = ir.RankedTensorType(args[0].type).shape
+        dims = strip_vecmat_dimensions(warp_arg0, actual_shape0)
+        warp_dims = collapse_into_leading_dimension(warp_arg0, dims)
+
+        # Figure out the types and shapes of the input arrays.
+        arg_strings = []
+        operand_layouts = []
+        for actual, warg in zip(args, wp_kernel.adj.args):
+            wtype = warg.type
+            rtt = ir.RankedTensorType(actual.type)
+
+            if not isinstance(wtype, wp.array):
+                raise Exception("Only contiguous arrays are supported for Jax kernel arguments")
+
+            if not base_type_is_compatible(wtype.dtype, rtt.element_type):
+                raise TypeError(f"Incompatible data type for argument '{warg.label}', expected {type_str(wtype.dtype)}, got {rtt.element_type}")
+
+            # Infer array dimension (by removing the vector/matrix dimensions and
+            # collapsing the initial dimensions).
+            shape = infer_dimensions(warg, rtt.shape)
+
+            if len(shape) != wtype.ndim:
+                raise TypeError(f"Incompatible array dimensionality for argument '{warg.label}'")
+
+            arg_strings.append(",".join([str(d) for d in shape]))
+            operand_layouts.append(default_layout(rtt.shape))
+
+        # Figure out the types and shapes of the output arrays.
+        result_types = []
+        result_layouts = []
+        for warg in wp_kernel.adj.args[len(args) :]:
+            wtype = warg.type
+
+            if not isinstance(wtype, wp.array):
+                raise Exception("Only contiguous arrays are supported for Jax kernel arguments")
+
+            # Infer dimensions from the first input.
+            arg_strings.append(",".join([str(d) for d in warp_dims]))
+
+            result_shape = list(dims) + list(get_vecmat_shape(wtype))
+            result_types.append(ir.RankedTensorType.get(result_shape, base_type_to_jax_ir(wtype.dtype)))
+            result_layouts.append(default_layout(result_shape))
+
+        # Build opaque descriptor for callback.
+        shape_str = ",".join([str(d) for d in warp_dims])
+        args_str = ";".join(arg_strings)
+        descriptor = f"{kernel}|{shape_str}|{args_str}"
+
+        out = custom_call(
+            b"warp_call",
+            result_types=result_types,
+            operands=args,
+            backend_config=descriptor.encode("utf-8"),
+            operand_layouts=operand_layouts,
+            result_layouts=result_layouts,
+        ).results
+        return out
+
+    mlir.register_lowering(
+        _jax_warp_p,
+        warp_call_lowering,
+        platform="gpu",
+    )

warp/native/builtin.h

@@ -354,6 +354,12 @@ inline CUDA_CALLABLE uint32 sign(uint32 x) { return 1; }
 inline CUDA_CALLABLE uint64 sign(uint64 x) { return 1; }
 
 
+// Catch-all for non-float types
+template<typename T>
+inline bool CUDA_CALLABLE isfinite(const T&)
+{
+    return true;
+}
 
 inline bool CUDA_CALLABLE isfinite(half x)
 {
@@ -368,6 +374,12 @@ inline bool CUDA_CALLABLE isfinite(double x)
     return ::isfinite(x);
 }
 
+template<typename T>
+inline CUDA_CALLABLE void print(const T&)
+{
+    printf("<type without print implementation>\n");
+}
+
 inline CUDA_CALLABLE void print(float16 f)
 {
     printf("%g\n", half_to_float(f));

warp/native/bvh.cu

@@ -373,16 +373,16 @@ LinearBVHBuilderGPU::LinearBVHBuilderGPU()
     , total_upper(NULL)
     , total_inv_edges(NULL)
 {
-    total_lower = (vec3*)alloc_temp_device(WP_CURRENT_CONTEXT, sizeof(vec3));
-    total_upper = (vec3*)alloc_temp_device(WP_CURRENT_CONTEXT, sizeof(vec3));
-    total_inv_edges = (vec3*)alloc_temp_device(WP_CURRENT_CONTEXT, sizeof(vec3));
+    total_lower = (vec3*)alloc_device(WP_CURRENT_CONTEXT, sizeof(vec3));
+    total_upper = (vec3*)alloc_device(WP_CURRENT_CONTEXT, sizeof(vec3));
+    total_inv_edges = (vec3*)alloc_device(WP_CURRENT_CONTEXT, sizeof(vec3));
 }
 
 LinearBVHBuilderGPU::~LinearBVHBuilderGPU()
 {
-    free_temp_device(WP_CURRENT_CONTEXT, total_lower);
-    free_temp_device(WP_CURRENT_CONTEXT, total_upper);
-    free_temp_device(WP_CURRENT_CONTEXT, total_inv_edges);
+    free_device(WP_CURRENT_CONTEXT, total_lower);
+    free_device(WP_CURRENT_CONTEXT, total_upper);
+    free_device(WP_CURRENT_CONTEXT, total_inv_edges);
 }
 
 
@@ -390,12 +390,12 @@ LinearBVHBuilderGPU::~LinearBVHBuilderGPU()
 void LinearBVHBuilderGPU::build(BVH& bvh, const vec3* item_lowers, const vec3* item_uppers, int num_items, bounds3* total_bounds)
 {
     // allocate temporary memory used during  building
-    indices = (int*)alloc_temp_device(WP_CURRENT_CONTEXT, sizeof(int)*num_items*2); 	// *2 for radix sort
-    keys = (int*)alloc_temp_device(WP_CURRENT_CONTEXT, sizeof(int)*num_items*2);	    // *2 for radix sort
-    deltas = (int*)alloc_temp_device(WP_CURRENT_CONTEXT, sizeof(int)*num_items);    	// highest differenting bit between keys for item i and i+1
-    range_lefts = (int*)alloc_temp_device(WP_CURRENT_CONTEXT, sizeof(int)*bvh.max_nodes);
-    range_rights = (int*)alloc_temp_device(WP_CURRENT_CONTEXT, sizeof(int)*bvh.max_nodes);
-    num_children = (int*)alloc_temp_device(WP_CURRENT_CONTEXT, sizeof(int)*bvh.max_nodes);
+    indices = (int*)alloc_device(WP_CURRENT_CONTEXT, sizeof(int)*num_items*2); 	// *2 for radix sort
+    keys = (int*)alloc_device(WP_CURRENT_CONTEXT, sizeof(int)*num_items*2);	    // *2 for radix sort
+    deltas = (int*)alloc_device(WP_CURRENT_CONTEXT, sizeof(int)*num_items);    	// highest differenting bit between keys for item i and i+1
+    range_lefts = (int*)alloc_device(WP_CURRENT_CONTEXT, sizeof(int)*bvh.max_nodes);
+    range_rights = (int*)alloc_device(WP_CURRENT_CONTEXT, sizeof(int)*bvh.max_nodes);
+    num_children = (int*)alloc_device(WP_CURRENT_CONTEXT, sizeof(int)*bvh.max_nodes);
 
     // if total bounds supplied by the host then we just 
     // compute our edge length and upload it to the GPU directly
@@ -445,13 +445,13 @@ void LinearBVHBuilderGPU::build(BVH& bvh, const vec3* item_lowers, const vec3* i
     wp_launch_device(WP_CURRENT_CONTEXT, build_hierarchy, num_items, (num_items, bvh.root, deltas, num_children, range_lefts, range_rights, bvh.node_parents, bvh.node_lowers, bvh.node_uppers));
 
     // free temporary memory
-    free_temp_device(WP_CURRENT_CONTEXT, indices);
-    free_temp_device(WP_CURRENT_CONTEXT, keys);
-    free_temp_device(WP_CURRENT_CONTEXT, deltas);
+    free_device(WP_CURRENT_CONTEXT, indices);
+    free_device(WP_CURRENT_CONTEXT, keys);
+    free_device(WP_CURRENT_CONTEXT, deltas);
 
-    free_temp_device(WP_CURRENT_CONTEXT, range_lefts);
-    free_temp_device(WP_CURRENT_CONTEXT, range_rights);
-    free_temp_device(WP_CURRENT_CONTEXT, num_children);
+    free_device(WP_CURRENT_CONTEXT, range_lefts);
+    free_device(WP_CURRENT_CONTEXT, range_rights);
+    free_device(WP_CURRENT_CONTEXT, num_children);
 
 }
 

warp/native/clang/clang.cpp

@@ -81,7 +81,7 @@ static void initialize_llvm()
     llvm::InitializeAllAsmPrinters();
 }
 
-static std::unique_ptr<llvm::Module> cpp_to_llvm(const std::string& input_file, const char* cpp_src, const char* include_dir, bool debug, llvm::LLVMContext& context)
+static std::unique_ptr<llvm::Module> cpp_to_llvm(const std::string& input_file, const char* cpp_src, const char* include_dir, bool debug, bool verify_fp, llvm::LLVMContext& context)
 {
     // Compilation arguments
     std::vector<const char*> args;
@@ -126,6 +126,11 @@ static std::unique_ptr<llvm::Module> cpp_to_llvm(const std::string& input_file,
         compiler_instance.getPreprocessorOpts().addMacroDef("NDEBUG");
     }
 
+    if(verify_fp)
+    {
+        compiler_instance.getPreprocessorOpts().addMacroDef("WP_VERIFY_FP");
+    }
+
     compiler_instance.getLangOpts().MicrosoftExt = 1;  // __forceinline / __int64
     compiler_instance.getLangOpts().DeclSpecKeyword = 1;  // __declspec
 
@@ -201,12 +206,12 @@ static std::unique_ptr<llvm::Module> cuda_to_llvm(const std::string& input_file,
 
 extern "C" {
 
-WP_API int compile_cpp(const char* cpp_src, const char *input_file, const char* include_dir, const char* output_file, bool debug)
+WP_API int compile_cpp(const char* cpp_src, const char *input_file, const char* include_dir, const char* output_file, bool debug, bool verify_fp)
 {
     initialize_llvm();
 
     llvm::LLVMContext context;
-    std::unique_ptr<llvm::Module> module = cpp_to_llvm(input_file, cpp_src, include_dir, debug, context);
+    std::unique_ptr<llvm::Module> module = cpp_to_llvm(input_file, cpp_src, include_dir, debug, verify_fp, context);
 
     if(!module)
     {

warp/native/cuda_util.cpp

@@ -9,6 +9,7 @@
 #if WP_ENABLE_CUDA
 
 #include "cuda_util.h"
+#include "error.h"
 
 #if defined(_WIN32)
 #define WIN32_LEAN_AND_MEAN
@@ -19,6 +20,9 @@
 #include <dlfcn.h>
 #endif
 
+#include <set>
+#include <stack>
+
 // the minimum CUDA version required from the driver
 #define WP_CUDA_DRIVER_VERSION 11030
 
@@ -63,6 +67,7 @@ static PFN_cuDeviceGetUuid_v11040 pfn_cuDeviceGetUuid;
 static PFN_cuDevicePrimaryCtxRetain_v7000 pfn_cuDevicePrimaryCtxRetain;
 static PFN_cuDevicePrimaryCtxRelease_v11000 pfn_cuDevicePrimaryCtxRelease;
 static PFN_cuDeviceCanAccessPeer_v4000 pfn_cuDeviceCanAccessPeer;
+static PFN_cuMemGetInfo_v3020 pfn_cuMemGetInfo;
 static PFN_cuCtxGetCurrent_v4000 pfn_cuCtxGetCurrent;
 static PFN_cuCtxSetCurrent_v4000 pfn_cuCtxSetCurrent;
 static PFN_cuCtxPushCurrent_v4000 pfn_cuCtxPushCurrent;
@@ -72,18 +77,23 @@ static PFN_cuCtxGetDevice_v2000 pfn_cuCtxGetDevice;
 static PFN_cuCtxCreate_v3020 pfn_cuCtxCreate;
 static PFN_cuCtxDestroy_v4000 pfn_cuCtxDestroy;
 static PFN_cuCtxEnablePeerAccess_v4000 pfn_cuCtxEnablePeerAccess;
+static PFN_cuCtxDisablePeerAccess_v4000 pfn_cuCtxDisablePeerAccess;
 static PFN_cuStreamCreate_v2000 pfn_cuStreamCreate;
 static PFN_cuStreamDestroy_v4000 pfn_cuStreamDestroy;
 static PFN_cuStreamSynchronize_v2000 pfn_cuStreamSynchronize;
 static PFN_cuStreamWaitEvent_v3020 pfn_cuStreamWaitEvent;
+static PFN_cuStreamGetCaptureInfo_v11030 pfn_cuStreamGetCaptureInfo;
+static PFN_cuStreamUpdateCaptureDependencies_v11030 pfn_cuStreamUpdateCaptureDependencies;
 static PFN_cuEventCreate_v2000 pfn_cuEventCreate;
 static PFN_cuEventDestroy_v4000 pfn_cuEventDestroy;
 static PFN_cuEventRecord_v2000 pfn_cuEventRecord;
+static PFN_cuEventRecordWithFlags_v11010 pfn_cuEventRecordWithFlags;
 static PFN_cuModuleLoadDataEx_v2010 pfn_cuModuleLoadDataEx;
 static PFN_cuModuleUnload_v2000 pfn_cuModuleUnload;
 static PFN_cuModuleGetFunction_v2000 pfn_cuModuleGetFunction;
 static PFN_cuLaunchKernel_v4000 pfn_cuLaunchKernel;
 static PFN_cuMemcpyPeerAsync_v4000 pfn_cuMemcpyPeerAsync;
+static PFN_cuPointerGetAttribute_v4000 pfn_cuPointerGetAttribute;
 static PFN_cuGraphicsMapResources_v3000 pfn_cuGraphicsMapResources;
 static PFN_cuGraphicsUnmapResources_v3000 pfn_cuGraphicsUnmapResources;
 static PFN_cuGraphicsResourceGetMappedPointer_v3020 pfn_cuGraphicsResourceGetMappedPointer;
@@ -171,6 +181,7 @@ bool init_cuda_driver()
     get_driver_entry_point("cuDevicePrimaryCtxRetain", &(void*&)pfn_cuDevicePrimaryCtxRetain);
     get_driver_entry_point("cuDevicePrimaryCtxRelease", &(void*&)pfn_cuDevicePrimaryCtxRelease);
     get_driver_entry_point("cuDeviceCanAccessPeer", &(void*&)pfn_cuDeviceCanAccessPeer);
+    get_driver_entry_point("cuMemGetInfo", &(void*&)pfn_cuMemGetInfo);
     get_driver_entry_point("cuCtxSetCurrent", &(void*&)pfn_cuCtxSetCurrent);
     get_driver_entry_point("cuCtxGetCurrent", &(void*&)pfn_cuCtxGetCurrent);
     get_driver_entry_point("cuCtxPushCurrent", &(void*&)pfn_cuCtxPushCurrent);
@@ -180,18 +191,23 @@ bool init_cuda_driver()
     get_driver_entry_point("cuCtxCreate", &(void*&)pfn_cuCtxCreate);
     get_driver_entry_point("cuCtxDestroy", &(void*&)pfn_cuCtxDestroy);
     get_driver_entry_point("cuCtxEnablePeerAccess", &(void*&)pfn_cuCtxEnablePeerAccess);
+    get_driver_entry_point("cuCtxDisablePeerAccess", &(void*&)pfn_cuCtxDisablePeerAccess);
     get_driver_entry_point("cuStreamCreate", &(void*&)pfn_cuStreamCreate);
     get_driver_entry_point("cuStreamDestroy", &(void*&)pfn_cuStreamDestroy);
     get_driver_entry_point("cuStreamSynchronize", &(void*&)pfn_cuStreamSynchronize);
     get_driver_entry_point("cuStreamWaitEvent", &(void*&)pfn_cuStreamWaitEvent);
+    get_driver_entry_point("cuStreamGetCaptureInfo", &(void*&)pfn_cuStreamGetCaptureInfo);
+    get_driver_entry_point("cuStreamUpdateCaptureDependencies", &(void*&)pfn_cuStreamUpdateCaptureDependencies);
     get_driver_entry_point("cuEventCreate", &(void*&)pfn_cuEventCreate);
     get_driver_entry_point("cuEventDestroy", &(void*&)pfn_cuEventDestroy);
     get_driver_entry_point("cuEventRecord", &(void*&)pfn_cuEventRecord);
+    get_driver_entry_point("cuEventRecordWithFlags", &(void*&)pfn_cuEventRecordWithFlags);
     get_driver_entry_point("cuModuleLoadDataEx", &(void*&)pfn_cuModuleLoadDataEx);
     get_driver_entry_point("cuModuleUnload", &(void*&)pfn_cuModuleUnload);
     get_driver_entry_point("cuModuleGetFunction", &(void*&)pfn_cuModuleGetFunction);
     get_driver_entry_point("cuLaunchKernel", &(void*&)pfn_cuLaunchKernel);
     get_driver_entry_point("cuMemcpyPeerAsync", &(void*&)pfn_cuMemcpyPeerAsync);
+    get_driver_entry_point("cuPointerGetAttribute", &(void*&)pfn_cuPointerGetAttribute);
     get_driver_entry_point("cuGraphicsMapResources", &(void*&)pfn_cuGraphicsMapResources);
     get_driver_entry_point("cuGraphicsUnmapResources", &(void*&)pfn_cuGraphicsUnmapResources);
     get_driver_entry_point("cuGraphicsResourceGetMappedPointer", &(void*&)pfn_cuGraphicsResourceGetMappedPointer);
@@ -209,16 +225,16 @@ bool is_cuda_driver_initialized()
     return cuda_driver_initialized;
 }
 
-bool check_cuda_result(cudaError_t code, const char* file, int line)
+bool check_cuda_result(cudaError_t code, const char* func, const char* file, int line)
 {
     if (code == cudaSuccess)
         return true;
 
-    fprintf(stderr, "Warp CUDA error %u: %s (%s:%d)\n", unsigned(code), cudaGetErrorString(code), file, line);
+    wp::set_error_string("Warp CUDA error %u: %s (in function %s, %s:%d)", unsigned(code), cudaGetErrorString(code), func, file, line);
     return false;
 }
 
-bool check_cu_result(CUresult result, const char* file, int line)
+bool check_cu_result(CUresult result, const char* func, const char* file, int line)
 {
     if (result == CUDA_SUCCESS)
         return true;
@@ -228,13 +244,56 @@ bool check_cu_result(CUresult result, const char* file, int line)
         pfn_cuGetErrorString(result, &errString);
 
     if (errString)
-        fprintf(stderr, "Warp CUDA error %u: %s (%s:%d)\n", unsigned(result), errString, file, line);
+        wp::set_error_string("Warp CUDA error %u: %s (in function %s, %s:%d)", unsigned(result), errString, func, file, line);
     else
-        fprintf(stderr, "Warp CUDA error %u (%s:%d)\n", unsigned(result), file, line);
+        wp::set_error_string("Warp CUDA error %u (in function %s, %s:%d)", unsigned(result), func, file, line);
 
     return false;
 }
 
+bool get_capture_dependencies(CUstream stream, std::vector<CUgraphNode>& dependencies_ret)
+{
+    CUstreamCaptureStatus status;
+    size_t num_dependencies = 0;
+    const CUgraphNode* dependencies = NULL;
+    dependencies_ret.clear();
+    if (check_cu(cuStreamGetCaptureInfo_f(stream, &status, NULL, NULL, &dependencies, &num_dependencies)))
+    {
+        if (dependencies && num_dependencies > 0)
+            dependencies_ret.insert(dependencies_ret.begin(), dependencies, dependencies + num_dependencies);
+        return true;
+    }
+    return false;
+}
+
+bool get_graph_leaf_nodes(cudaGraph_t graph, std::vector<cudaGraphNode_t>& leaf_nodes_ret)
+{
+    if (!graph)
+        return false;
+
+    size_t node_count = 0;
+    if (!check_cuda(cudaGraphGetNodes(graph, NULL, &node_count)))
+        return false;
+
+    std::vector<cudaGraphNode_t> nodes(node_count);
+    if (!check_cuda(cudaGraphGetNodes(graph, nodes.data(), &node_count)))
+        return false;
+
+    leaf_nodes_ret.clear();
+
+    for (cudaGraphNode_t node : nodes)
+    {
+        size_t dependent_count;
+        if (!check_cuda(cudaGraphNodeGetDependentNodes(node, NULL, &dependent_count)))
+            return false;
+
+        if (dependent_count == 0)
+            leaf_nodes_ret.push_back(node);
+    }
+
+    return true;
+}
+
 
 #define DRIVER_ENTRY_POINT_ERROR driver_entry_point_error(__FUNCTION__)
 
@@ -311,6 +370,11 @@ CUresult cuDeviceCanAccessPeer_f(int* can_access, CUdevice dev, CUdevice peer_de
     return pfn_cuDeviceCanAccessPeer ? pfn_cuDeviceCanAccessPeer(can_access, dev, peer_dev) : DRIVER_ENTRY_POINT_ERROR;
 }
 
+CUresult cuMemGetInfo_f(size_t* free, size_t* total)
+{
+    return pfn_cuMemGetInfo ? pfn_cuMemGetInfo(free, total) : DRIVER_ENTRY_POINT_ERROR;
+}
+
 CUresult cuCtxGetCurrent_f(CUcontext* ctx)
 {
     return pfn_cuCtxGetCurrent ? pfn_cuCtxGetCurrent(ctx) : DRIVER_ENTRY_POINT_ERROR;
@@ -356,6 +420,11 @@ CUresult cuCtxEnablePeerAccess_f(CUcontext peer_ctx, unsigned int flags)
     return pfn_cuCtxEnablePeerAccess ? pfn_cuCtxEnablePeerAccess(peer_ctx, flags) : DRIVER_ENTRY_POINT_ERROR;
 }
 
+CUresult cuCtxDisablePeerAccess_f(CUcontext peer_ctx)
+{
+    return pfn_cuCtxDisablePeerAccess ? pfn_cuCtxDisablePeerAccess(peer_ctx) : DRIVER_ENTRY_POINT_ERROR;
+}
+
 CUresult cuStreamCreate_f(CUstream* stream, unsigned int flags)
 {
     return pfn_cuStreamCreate ? pfn_cuStreamCreate(stream, flags) : DRIVER_ENTRY_POINT_ERROR;
@@ -376,6 +445,16 @@ CUresult cuStreamWaitEvent_f(CUstream stream, CUevent event, unsigned int flags)
     return pfn_cuStreamWaitEvent ? pfn_cuStreamWaitEvent(stream, event, flags) : DRIVER_ENTRY_POINT_ERROR;
 }
 
+CUresult cuStreamGetCaptureInfo_f(CUstream stream, CUstreamCaptureStatus *captureStatus_out, cuuint64_t *id_out, CUgraph *graph_out, const CUgraphNode **dependencies_out, size_t *numDependencies_out)
+{
+    return pfn_cuStreamGetCaptureInfo ? pfn_cuStreamGetCaptureInfo(stream, captureStatus_out, id_out, graph_out, dependencies_out, numDependencies_out) : DRIVER_ENTRY_POINT_ERROR;
+}
+
+CUresult cuStreamUpdateCaptureDependencies_f(CUstream stream, CUgraphNode *dependencies, size_t numDependencies, unsigned int flags)
+{
+    return pfn_cuStreamUpdateCaptureDependencies ? pfn_cuStreamUpdateCaptureDependencies(stream, dependencies, numDependencies, flags) : DRIVER_ENTRY_POINT_ERROR;
+}
+
 CUresult cuEventCreate_f(CUevent* event, unsigned int flags)
 {
     return pfn_cuEventCreate ? pfn_cuEventCreate(event, flags) : DRIVER_ENTRY_POINT_ERROR;
@@ -391,6 +470,11 @@ CUresult cuEventRecord_f(CUevent event, CUstream stream)
     return pfn_cuEventRecord ? pfn_cuEventRecord(event, stream) : DRIVER_ENTRY_POINT_ERROR;
 }
 
+CUresult cuEventRecordWithFlags_f(CUevent event, CUstream stream, unsigned int flags)
+{
+    return pfn_cuEventRecordWithFlags ? pfn_cuEventRecordWithFlags(event, stream, flags) : DRIVER_ENTRY_POINT_ERROR;
+}
+
 CUresult cuModuleLoadDataEx_f(CUmodule *module, const void *image, unsigned int numOptions, CUjit_option *options, void **optionValues)
 {
     return pfn_cuModuleLoadDataEx ? pfn_cuModuleLoadDataEx(module, image, numOptions, options, optionValues) : DRIVER_ENTRY_POINT_ERROR;
@@ -416,6 +500,11 @@ CUresult cuMemcpyPeerAsync_f(CUdeviceptr dst_ptr, CUcontext dst_ctx, CUdeviceptr
     return pfn_cuMemcpyPeerAsync ? pfn_cuMemcpyPeerAsync(dst_ptr, dst_ctx, src_ptr, src_ctx, n, stream) : DRIVER_ENTRY_POINT_ERROR;
 }
 
+CUresult cuPointerGetAttribute_f(void* data, CUpointer_attribute attribute, CUdeviceptr ptr)
+{
+    return pfn_cuPointerGetAttribute ? pfn_cuPointerGetAttribute(data, attribute, ptr) : DRIVER_ENTRY_POINT_ERROR;
+}
+
 CUresult cuGraphicsMapResources_f(unsigned int count, CUgraphicsResource* resources, CUstream stream)
 {
     return pfn_cuGraphicsMapResources ? pfn_cuGraphicsMapResources(count, resources, stream) : DRIVER_ENTRY_POINT_ERROR;