warp-lang 0.10.1__py3-none-win_amd64.whl → 0.11.0__py3-none-win_amd64.whl

warp/native/warp.cu

@@ -73,10 +73,15 @@ struct DeviceInfo
     static constexpr int kNameLen = 128;
 
     CUdevice device = -1;
+    CUuuid uuid = {0};
     int ordinal = -1;
+    int pci_domain_id = -1;
+    int pci_bus_id = -1;
+    int pci_device_id = -1;
     char name[kNameLen] = "";
     int arch = 0;
     int is_uva = 0;
+    int is_memory_pool_supported = 0;
 };
 
 struct ContextInfo
@@ -125,7 +130,12 @@ int cuda_init()
                 g_devices[i].device = device;
                 g_devices[i].ordinal = i;
                 check_cu(cuDeviceGetName_f(g_devices[i].name, DeviceInfo::kNameLen, device));
+                check_cu(cuDeviceGetUuid_f(&g_devices[i].uuid, device));
+                check_cu(cuDeviceGetAttribute_f(&g_devices[i].pci_domain_id, CU_DEVICE_ATTRIBUTE_PCI_DOMAIN_ID, device));
+                check_cu(cuDeviceGetAttribute_f(&g_devices[i].pci_bus_id, CU_DEVICE_ATTRIBUTE_PCI_BUS_ID, device));
+                check_cu(cuDeviceGetAttribute_f(&g_devices[i].pci_device_id, CU_DEVICE_ATTRIBUTE_PCI_DEVICE_ID, device));
                 check_cu(cuDeviceGetAttribute_f(&g_devices[i].is_uva, CU_DEVICE_ATTRIBUTE_UNIFIED_ADDRESSING, device));
+                check_cu(cuDeviceGetAttribute_f(&g_devices[i].is_memory_pool_supported, CU_DEVICE_ATTRIBUTE_MEMORY_POOLS_SUPPORTED, device));
                 int major = 0;
                 int minor = 0;
                 check_cu(cuDeviceGetAttribute_f(&major, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, device));
@@ -216,6 +226,26 @@ void* alloc_device(void* context, size_t s)
     return ptr;
 }
 
+void* alloc_temp_device(void* context, size_t s)
+{
+    // "cudaMallocAsync ignores the current device/context when determining where the allocation will reside. Instead,
+    // cudaMallocAsync determines the resident device based on the specified memory pool or the supplied stream."
+    ContextGuard guard(context);
+
+    void* ptr;
+
+    if (cuda_context_is_memory_pool_supported(context))
+    {
+        check_cuda(cudaMallocAsync(&ptr, s, get_current_stream()));
+    }
+    else
+    {
+        check_cuda(cudaMalloc(&ptr, s));
+    }
+
+    return ptr;
+}
+
 void free_device(void* context, void* ptr)
 {
     ContextGuard guard(context);
@@ -223,6 +253,20 @@ void free_device(void* context, void* ptr)
     check_cuda(cudaFree(ptr));
 }
 
+void free_temp_device(void* context, void* ptr)
+{
+    ContextGuard guard(context);
+
+    if (cuda_context_is_memory_pool_supported(context))
+    {
+        check_cuda(cudaFreeAsync(ptr, get_current_stream()));
+    }
+    else
+    {
+        check_cuda(cudaFree(ptr));
+    }
+}
+
 void memcpy_h2d(void* context, void* dest, void* src, size_t n)
 {
     ContextGuard guard(context);
@@ -266,7 +310,7 @@ void memset_device(void* context, void* dest, int value, size_t n)
 {
     ContextGuard guard(context);
 
-    if ((n%4) > 0)
+    if (true)// ((n%4) > 0)
     {
         // for unaligned lengths fallback to CUDA memset
         check_cuda(cudaMemsetAsync(dest, value, n, get_current_stream()));
@@ -448,6 +492,125 @@ static __global__ void array_copy_4d_kernel(void* dst, const void* src,
 }
 
 
+static __global__ void array_copy_from_fabric_kernel(wp::fabricarray_t<void> src,
+                                                     void* dst_data, int dst_stride, const int* dst_indices,
+                                                     int elem_size)
+{
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (tid < src.size)
+    {
+        int dst_idx = dst_indices ? dst_indices[tid] : tid;
+        void* dst_ptr = (char*)dst_data + dst_idx * dst_stride;
+        const void* src_ptr = fabricarray_element_ptr(src, tid, elem_size);
+        memcpy(dst_ptr, src_ptr, elem_size);
+    }
+}
+
+static __global__ void array_copy_from_fabric_indexed_kernel(wp::indexedfabricarray_t<void> src,
+                                                             void* dst_data, int dst_stride, const int* dst_indices,
+                                                             int elem_size)
+{
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (tid < src.size)
+    {
+        int src_index = src.indices[tid];
+        int dst_idx = dst_indices ? dst_indices[tid] : tid;
+        void* dst_ptr = (char*)dst_data + dst_idx * dst_stride;
+        const void* src_ptr = fabricarray_element_ptr(src.fa, src_index, elem_size);
+        memcpy(dst_ptr, src_ptr, elem_size);
+    }
+}
+
+static __global__ void array_copy_to_fabric_kernel(wp::fabricarray_t<void> dst,
+                                                   const void* src_data, int src_stride, const int* src_indices,
+                                                   int elem_size)
+{
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (tid < dst.size)
+    {
+        int src_idx = src_indices ? src_indices[tid] : tid;
+        const void* src_ptr = (const char*)src_data + src_idx * src_stride;
+        void* dst_ptr = fabricarray_element_ptr(dst, tid, elem_size);
+        memcpy(dst_ptr, src_ptr, elem_size);
+    }
+}
+
+static __global__ void array_copy_to_fabric_indexed_kernel(wp::indexedfabricarray_t<void> dst,
+                                                           const void* src_data, int src_stride, const int* src_indices,
+                                                           int elem_size)
+{
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (tid < dst.size)
+    {
+        int src_idx = src_indices ? src_indices[tid] : tid;
+        const void* src_ptr = (const char*)src_data + src_idx * src_stride;
+        int dst_idx = dst.indices[tid];
+        void* dst_ptr = fabricarray_element_ptr(dst.fa, dst_idx, elem_size);
+        memcpy(dst_ptr, src_ptr, elem_size);
+    }
+}
+
+
+static __global__ void array_copy_fabric_to_fabric_kernel(wp::fabricarray_t<void> dst, wp::fabricarray_t<void> src, int elem_size)
+{
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (tid < dst.size)
+    {
+        const void* src_ptr = fabricarray_element_ptr(src, tid, elem_size);
+        void* dst_ptr = fabricarray_element_ptr(dst, tid, elem_size);
+        memcpy(dst_ptr, src_ptr, elem_size);
+    }
+}
+
+
+static __global__ void array_copy_fabric_to_fabric_indexed_kernel(wp::indexedfabricarray_t<void> dst, wp::fabricarray_t<void> src, int elem_size)
+{
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (tid < dst.size)
+    {
+        const void* src_ptr = fabricarray_element_ptr(src, tid, elem_size);
+        int dst_index = dst.indices[tid];
+        void* dst_ptr = fabricarray_element_ptr(dst.fa, dst_index, elem_size);
+        memcpy(dst_ptr, src_ptr, elem_size);
+    }
+}
+
+
+static __global__ void array_copy_fabric_indexed_to_fabric_kernel(wp::fabricarray_t<void> dst, wp::indexedfabricarray_t<void> src, int elem_size)
+{
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (tid < dst.size)
+    {
+        int src_index = src.indices[tid];
+        const void* src_ptr = fabricarray_element_ptr(src.fa, src_index, elem_size);
+        void* dst_ptr = fabricarray_element_ptr(dst, tid, elem_size);
+        memcpy(dst_ptr, src_ptr, elem_size);
+    }
+}
+
+
+static __global__ void array_copy_fabric_indexed_to_fabric_indexed_kernel(wp::indexedfabricarray_t<void> dst, wp::indexedfabricarray_t<void> src, int elem_size)
+{
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (tid < dst.size)
+    {
+        int src_index = src.indices[tid];
+        int dst_index = dst.indices[tid];
+        const void* src_ptr = fabricarray_element_ptr(src.fa, src_index, elem_size);
+        void* dst_ptr = fabricarray_element_ptr(dst.fa, dst_index, elem_size);
+        memcpy(dst_ptr, src_ptr, elem_size);
+    }
+}
+
+
 WP_API size_t array_copy_device(void* context, void* dst, void* src, int dst_type, int src_type, int elem_size)
 {
     if (!src || !dst)
@@ -466,6 +629,12 @@ WP_API size_t array_copy_device(void* context, void* dst, void* src, int dst_typ
     const int*const* src_indices = NULL;
     const int*const* dst_indices = NULL;
 
+    const wp::fabricarray_t<void>* src_fabricarray = NULL;
+    wp::fabricarray_t<void>* dst_fabricarray = NULL;
+
+    const wp::indexedfabricarray_t<void>* src_indexedfabricarray = NULL;
+    wp::indexedfabricarray_t<void>* dst_indexedfabricarray = NULL;
+
     const int* null_indices[wp::ARRAY_MAX_DIMS] = { NULL };
 
     if (src_type == wp::ARRAY_TYPE_REGULAR)
@@ -487,9 +656,19 @@ WP_API size_t array_copy_device(void* context, void* dst, void* src, int dst_typ
         src_strides = src_arr.arr.strides;
         src_indices = src_arr.indices;
     }
+    else if (src_type == wp::ARRAY_TYPE_FABRIC)
+    {
+        src_fabricarray = static_cast<const wp::fabricarray_t<void>*>(src);
+        src_ndim = 1;
+    }
+    else if (src_type == wp::ARRAY_TYPE_FABRIC_INDEXED)
+    {
+        src_indexedfabricarray = static_cast<const wp::indexedfabricarray_t<void>*>(src);
+        src_ndim = 1;
+    }
     else
     {
-        fprintf(stderr, "Warp error: Invalid array type (%d)\n", src_type);
+        fprintf(stderr, "Warp copy error: Invalid array type (%d)\n", src_type);
         return 0;
     }
 
@@ -512,33 +691,149 @@ WP_API size_t array_copy_device(void* context, void* dst, void* src, int dst_typ
         dst_strides = dst_arr.arr.strides;
         dst_indices = dst_arr.indices;
     }
+    else if (dst_type == wp::ARRAY_TYPE_FABRIC)
+    {
+        dst_fabricarray = static_cast<wp::fabricarray_t<void>*>(dst);
+        dst_ndim = 1;
+    }
+    else if (dst_type == wp::ARRAY_TYPE_FABRIC_INDEXED)
+    {
+        dst_indexedfabricarray = static_cast<wp::indexedfabricarray_t<void>*>(dst);
+        dst_ndim = 1;
+    }
     else
     {
-        fprintf(stderr, "Warp error: Invalid array type (%d)\n", dst_type);
+        fprintf(stderr, "Warp copy error: Invalid array type (%d)\n", dst_type);
         return 0;
     }
 
     if (src_ndim != dst_ndim)
     {
-        fprintf(stderr, "Warp error: Incompatible array dimensionalities (%d and %d)\n", src_ndim, dst_ndim);
+        fprintf(stderr, "Warp copy error: Incompatible array dimensionalities (%d and %d)\n", src_ndim, dst_ndim);
         return 0;
     }
 
-    bool has_grad = (src_grad && dst_grad);
-    size_t n = 1;
+    ContextGuard guard(context);
+
+    // handle fabric arrays
+    if (dst_fabricarray)
+    {
+        size_t n = dst_fabricarray->size;
+        if (src_fabricarray)
+        {
+            // copy from fabric to fabric
+            if (src_fabricarray->size != n)
+            {
+                fprintf(stderr, "Warp copy error: Incompatible array sizes\n");
+                return 0;
+            }
+            wp_launch_device(WP_CURRENT_CONTEXT, array_copy_fabric_to_fabric_kernel, n,
+                            (*dst_fabricarray, *src_fabricarray, elem_size));
+            return n;
+        }
+        else if (src_indexedfabricarray)
+        {
+            // copy from fabric indexed to fabric
+            if (src_indexedfabricarray->size != n)
+            {
+                fprintf(stderr, "Warp copy error: Incompatible array sizes\n");
+                return 0;
+            }
+            wp_launch_device(WP_CURRENT_CONTEXT, array_copy_fabric_indexed_to_fabric_kernel, n,
+                            (*dst_fabricarray, *src_indexedfabricarray, elem_size));
+            return n;
+        }
+        else
+        {
+            // copy to fabric
+            if (size_t(src_shape[0]) != n)
+            {
+                fprintf(stderr, "Warp copy error: Incompatible array sizes\n");
+                return 0;
+            }
+            wp_launch_device(WP_CURRENT_CONTEXT, array_copy_to_fabric_kernel, n,
+                            (*dst_fabricarray, src_data, src_strides[0], src_indices[0], elem_size));
+            return n;
+        }
+    }
+    if (dst_indexedfabricarray)
+    {
+        size_t n = dst_indexedfabricarray->size;
+        if (src_fabricarray)
+        {
+            // copy from fabric to fabric indexed
+            if (src_fabricarray->size != n)
+            {
+                fprintf(stderr, "Warp copy error: Incompatible array sizes\n");
+                return 0;
+            }
+            wp_launch_device(WP_CURRENT_CONTEXT, array_copy_fabric_to_fabric_indexed_kernel, n,
+                            (*dst_indexedfabricarray, *src_fabricarray, elem_size));
+            return n;
+        }
+        else if (src_indexedfabricarray)
+        {
+            // copy from fabric indexed to fabric indexed
+            if (src_indexedfabricarray->size != n)
+            {
+                fprintf(stderr, "Warp copy error: Incompatible array sizes\n");
+                return 0;
+            }
+            wp_launch_device(WP_CURRENT_CONTEXT, array_copy_fabric_indexed_to_fabric_indexed_kernel, n,
+                            (*dst_indexedfabricarray, *src_indexedfabricarray, elem_size));
+            return n;
+        }
+        else
+        {
+            // copy to fabric indexed
+            if (size_t(src_shape[0]) != n)
+            {
+                fprintf(stderr, "Warp copy error: Incompatible array sizes\n");
+                return 0;
+            }
+            wp_launch_device(WP_CURRENT_CONTEXT, array_copy_to_fabric_indexed_kernel, n,
+                             (*dst_indexedfabricarray, src_data, src_strides[0], src_indices[0], elem_size));
+            return n;
+        }
+    }
+    else if (src_fabricarray)
+    {
+        // copy from fabric
+        size_t n = src_fabricarray->size;
+        if (size_t(dst_shape[0]) != n)
+        {
+            fprintf(stderr, "Warp copy error: Incompatible array sizes\n");
+            return 0;
+        }
+        wp_launch_device(WP_CURRENT_CONTEXT, array_copy_from_fabric_kernel, n,
+                         (*src_fabricarray, dst_data, dst_strides[0], dst_indices[0], elem_size));
+        return n;
+    }
+    else if (src_indexedfabricarray)
+    {
+        // copy from fabric indexed
+        size_t n = src_indexedfabricarray->size;
+        if (size_t(dst_shape[0]) != n)
+        {
+            fprintf(stderr, "Warp copy error: Incompatible array sizes\n");
+            return 0;
+        }
+        wp_launch_device(WP_CURRENT_CONTEXT, array_copy_from_fabric_indexed_kernel, n,
+                         (*src_indexedfabricarray, dst_data, dst_strides[0], dst_indices[0], elem_size));
+        return n;
+    }
 
+    size_t n = 1;
     for (int i = 0; i < src_ndim; i++)
     {
         if (src_shape[i] != dst_shape[i])
         {
-            fprintf(stderr, "Warp error: Incompatible array shapes\n");
+            fprintf(stderr, "Warp copy error: Incompatible array shapes\n");
             return 0;
         }
         n *= src_shape[i];
     }
 
-    ContextGuard guard(context);
-
     switch (src_ndim)
     {
     case 1:
@@ -547,13 +842,6 @@ WP_API size_t array_copy_device(void* context, void* dst, void* src, int dst_typ
                                                                    dst_strides[0], src_strides[0],
                                                                    dst_indices[0], src_indices[0],
                                                                    src_shape[0], elem_size));
-        if (has_grad)
-        {
-            wp_launch_device(WP_CURRENT_CONTEXT, array_copy_1d_kernel, n, (dst_grad, src_grad,
-                                                                       dst_strides[0], src_strides[0],
-                                                                       dst_indices[0], src_indices[0],
-                                                                       src_shape[0], elem_size));
-        }
         break;
     }
     case 2:
@@ -568,13 +856,6 @@ WP_API size_t array_copy_device(void* context, void* dst, void* src, int dst_typ
                                                                    dst_strides_v, src_strides_v,
                                                                    dst_indices_v, src_indices_v,
                                                                    shape_v, elem_size));
-        if (has_grad)
-        {
-            wp_launch_device(WP_CURRENT_CONTEXT, array_copy_2d_kernel, n, (dst_grad, src_grad,
-                                                                       dst_strides_v, src_strides_v,
-                                                                       dst_indices_v, src_indices_v,
-                                                                       shape_v, elem_size));
-        }
         break;
     }
     case 3:
@@ -589,13 +870,6 @@ WP_API size_t array_copy_device(void* context, void* dst, void* src, int dst_typ
                                                                    dst_strides_v, src_strides_v,
                                                                    dst_indices_v, src_indices_v,
                                                                    shape_v, elem_size));
-        if (has_grad)
-        {
-            wp_launch_device(WP_CURRENT_CONTEXT, array_copy_3d_kernel, n, (dst_grad, src_grad,
-                                                                       dst_strides_v, src_strides_v,
-                                                                       dst_indices_v, src_indices_v,
-                                                                       shape_v, elem_size));
-        }
         break;
     }
     case 4:
@@ -610,17 +884,10 @@ WP_API size_t array_copy_device(void* context, void* dst, void* src, int dst_typ
                                                                    dst_strides_v, src_strides_v,
                                                                    dst_indices_v, src_indices_v,
                                                                    shape_v, elem_size));
-        if (has_grad)
-        {
-            wp_launch_device(WP_CURRENT_CONTEXT, array_copy_4d_kernel, n, (dst_grad, src_grad,
-                                                                       dst_strides_v, src_strides_v,
-                                                                       dst_indices_v, src_indices_v,
-                                                                       shape_v, elem_size));
-        }
         break;
     }
     default:
-        fprintf(stderr, "Warp error: invalid array dimensionality (%d)\n", src_ndim);
+        fprintf(stderr, "Warp copy error: invalid array dimensionality (%d)\n", src_ndim);
         return 0;
     }
 
@@ -717,6 +984,32 @@ static __global__ void array_fill_4d_kernel(void* data,
 }
 
 
+static __global__ void array_fill_fabric_kernel(wp::fabricarray_t<void> fa, const void* value, int value_size)
+{
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    if (tid < fa.size)
+    {
+        void* dst_ptr = fabricarray_element_ptr(fa, tid, value_size);
+        memcpy(dst_ptr, value, value_size);
+    }
+}
+
+
+static __global__ void array_fill_fabric_indexed_kernel(wp::indexedfabricarray_t<void> ifa, const void* value, int value_size)
+{
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    if (tid < ifa.size)
+    {
+        size_t idx = size_t(ifa.indices[tid]);
+        if (idx < ifa.fa.size)
+        {
+            void* dst_ptr = fabricarray_element_ptr(ifa.fa, idx, value_size);
+            memcpy(dst_ptr, value, value_size);
+        }
+    }
+}
+
+
 WP_API void array_fill_device(void* context, void* arr_ptr, int arr_type, const void* value_ptr, int value_size)
 {
     if (!arr_ptr || !value_ptr)
@@ -728,6 +1021,9 @@ WP_API void array_fill_device(void* context, void* arr_ptr, int arr_type, const
     const int* strides = NULL;
     const int*const* indices = NULL;
 
+    wp::fabricarray_t<void>* fa = NULL;
+    wp::indexedfabricarray_t<void>* ifa = NULL;
+
     const int* null_indices[wp::ARRAY_MAX_DIMS] = { NULL };
 
     if (arr_type == wp::ARRAY_TYPE_REGULAR)
@@ -748,9 +1044,17 @@ WP_API void array_fill_device(void* context, void* arr_ptr, int arr_type, const
         strides = ia.arr.strides;
         indices = ia.indices;
     }
+    else if (arr_type == wp::ARRAY_TYPE_FABRIC)
+    {
+        fa = static_cast<wp::fabricarray_t<void>*>(arr_ptr);
+    }
+    else if (arr_type == wp::ARRAY_TYPE_FABRIC_INDEXED)
+    {
+        ifa = static_cast<wp::indexedfabricarray_t<void>*>(arr_ptr);
+    }
     else
     {
-        fprintf(stderr, "Warp error: Invalid array type id %d\n", arr_type);
+        fprintf(stderr, "Warp fill error: Invalid array type id %d\n", arr_type);
         return;
     }
 
@@ -765,6 +1069,21 @@ WP_API void array_fill_device(void* context, void* arr_ptr, int arr_type, const
     check_cuda(cudaMalloc(&value_devptr, value_size));
     check_cuda(cudaMemcpyAsync(value_devptr, value_ptr, value_size, cudaMemcpyHostToDevice, get_current_stream()));
 
+    // handle fabric arrays
+    if (fa)
+    {
+        wp_launch_device(WP_CURRENT_CONTEXT, array_fill_fabric_kernel, n,
+                         (*fa, value_devptr, value_size));
+        return;
+    }
+    else if (ifa)
+    {
+        wp_launch_device(WP_CURRENT_CONTEXT, array_fill_fabric_indexed_kernel, n,
+                         (*ifa, value_devptr, value_size));
+        return;
+    }
+
+    // handle regular or indexed arrays
     switch (ndim)
     {
     case 1:
@@ -801,7 +1120,7 @@ WP_API void array_fill_device(void* context, void* arr_ptr, int arr_type, const
         break;
     }
     default:
-        fprintf(stderr, "Warp error: invalid array dimensionality (%d)\n", ndim);
+        fprintf(stderr, "Warp fill error: invalid array dimensionality (%d)\n", ndim);
         return;
     }
 }
@@ -830,6 +1149,11 @@ int cuda_toolkit_version()
     return CUDA_VERSION;
 }
 
+bool cuda_driver_is_initialized()
+{
+    return is_cuda_driver_initialized();
+}
+
 int nvrtc_supported_arch_count()
 {
     int count;
@@ -884,6 +1208,32 @@ int cuda_device_get_arch(int ordinal)
     return 0;
 }
 
+void cuda_device_get_uuid(int ordinal, char uuid[16])
+{
+    memcpy(uuid, g_devices[ordinal].uuid.bytes, sizeof(char)*16);
+}
+
+int cuda_device_get_pci_domain_id(int ordinal)
+{
+    if (ordinal >= 0 && ordinal < int(g_devices.size()))
+        return g_devices[ordinal].pci_domain_id;
+    return -1;
+}
+
+int cuda_device_get_pci_bus_id(int ordinal)
+{
+    if (ordinal >= 0 && ordinal < int(g_devices.size()))
+        return g_devices[ordinal].pci_bus_id;
+    return -1;
+}
+
+int cuda_device_get_pci_device_id(int ordinal)
+{
+    if (ordinal >= 0 && ordinal < int(g_devices.size()))
+        return g_devices[ordinal].pci_device_id;
+    return -1;
+}
+
 int cuda_device_is_uva(int ordinal)
 {
     if (ordinal >= 0 && ordinal < int(g_devices.size()))
@@ -891,6 +1241,13 @@ int cuda_device_is_uva(int ordinal)
     return 0;
 }
 
+int cuda_device_is_memory_pool_supported(int ordinal)
+{
+    if (ordinal >= 0 && ordinal < int(g_devices.size()))
+        return g_devices[ordinal].is_memory_pool_supported;
+    return false;
+}
+
 void* cuda_context_get_current()
 {
     return get_current_context();
@@ -999,6 +1356,16 @@ int cuda_context_is_primary(void* context)
     return 0;
 }
 
+int cuda_context_is_memory_pool_supported(void* context)
+{
+    int ordinal = cuda_context_get_device_ordinal(context);
+    if (ordinal != -1)
+    {
+        return cuda_device_is_memory_pool_supported(ordinal);
+    }
+    return 0;
+}
+
 void* cuda_context_get_stream(void* context)
 {
     ContextInfo* info = get_context_info(static_cast<CUcontext>(context));
@@ -1208,10 +1575,10 @@ void* cuda_graph_end_capture(void* context)
         //cudaGraphDebugDotPrint(graph, "graph.dot", cudaGraphDebugDotFlagsVerbose);
 
         cudaGraphExec_t graph_exec = NULL;
-        check_cuda(cudaGraphInstantiate(&graph_exec, graph, NULL, NULL, 0));
+        //check_cuda(cudaGraphInstantiate(&graph_exec, graph, NULL, NULL, 0));
         
         // can use after CUDA 11.4 to permit graphs to capture cudaMallocAsync() operations
-        //check_cuda(cudaGraphInstantiateWithFlags(&graph_exec, graph, cudaGraphInstantiateFlagAutoFreeOnLaunch));
+        check_cuda(cudaGraphInstantiateWithFlags(&graph_exec, graph, cudaGraphInstantiateFlagAutoFreeOnLaunch));
 
         // free source graph
         check_cuda(cudaGraphDestroy(graph));
@@ -1513,14 +1880,34 @@ void* cuda_get_kernel(void* context, void* module, const char* name)
     return kernel;
 }
 
-size_t cuda_launch_kernel(void* context, void* kernel, size_t dim, void** args)
+size_t cuda_launch_kernel(void* context, void* kernel, size_t dim, int max_blocks, void** args)
 {
     ContextGuard guard(context);
 
     const int block_dim = 256;
     // CUDA specs up to compute capability 9.0 says the max x-dim grid is 2**31-1, so
     // grid_dim is fine as an int for the near future
-    const int grid_dim = (dim + block_dim - 1)/block_dim;
+    int grid_dim = (dim + block_dim - 1)/block_dim;
+
+    if (max_blocks <= 0) {
+        max_blocks = 2147483647;
+    }
+
+    if (grid_dim < 0)
+    {
+#if defined(_DEBUG)
+        fprintf(stderr, "Warp warning: Overflow in grid dimensions detected for %zu total elements and 256 threads "
+                "per block.\n    Setting block count to %d.\n", dim, max_blocks);
+#endif
+        grid_dim =  max_blocks;
+    }
+    else 
+    {
+        if (grid_dim > max_blocks)
+        {
+            grid_dim = max_blocks;
+        }
+    }
 
     CUresult res = cuLaunchKernel_f(
         (CUfunction)kernel,