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

warp/native/cuda_util.h

@@ -17,8 +17,10 @@
 
 #include <stdio.h>
 
-#define check_cuda(code) (check_cuda_result(code, __FILE__, __LINE__))
-#define check_cu(code) (check_cu_result(code, __FILE__, __LINE__))
+#include <vector>
+
+#define check_cuda(code) (check_cuda_result(code, __FUNCTION__, __FILE__, __LINE__))
+#define check_cu(code) (check_cu_result(code, __FUNCTION__, __FILE__, __LINE__))
 
 
 #if defined(__CUDACC__)
@@ -55,6 +57,7 @@ CUresult cuDeviceGetUuid_f(CUuuid* uuid, CUdevice dev);
 CUresult cuDevicePrimaryCtxRetain_f(CUcontext* ctx, CUdevice dev);
 CUresult cuDevicePrimaryCtxRelease_f(CUdevice dev);
 CUresult cuDeviceCanAccessPeer_f(int* can_access, CUdevice dev, CUdevice peer_dev);
+CUresult cuMemGetInfo_f(size_t* free, size_t* total);
 CUresult cuCtxGetCurrent_f(CUcontext* ctx);
 CUresult cuCtxSetCurrent_f(CUcontext ctx);
 CUresult cuCtxPushCurrent_f(CUcontext ctx);
@@ -64,18 +67,23 @@ CUresult cuCtxGetDevice_f(CUdevice* dev);
 CUresult cuCtxCreate_f(CUcontext* ctx, unsigned int flags, CUdevice dev);
 CUresult cuCtxDestroy_f(CUcontext ctx);
 CUresult cuCtxEnablePeerAccess_f(CUcontext peer_ctx, unsigned int flags);
+CUresult cuCtxDisablePeerAccess_f(CUcontext peer_ctx);
 CUresult cuStreamCreate_f(CUstream* stream, unsigned int flags);
 CUresult cuStreamDestroy_f(CUstream stream);
 CUresult cuStreamSynchronize_f(CUstream stream);
 CUresult cuStreamWaitEvent_f(CUstream stream, CUevent event, unsigned int flags);
+CUresult cuStreamGetCaptureInfo_f(CUstream stream, CUstreamCaptureStatus *captureStatus_out, cuuint64_t *id_out, CUgraph *graph_out, const CUgraphNode **dependencies_out, size_t *numDependencies_out);
+CUresult cuStreamUpdateCaptureDependencies_f(CUstream stream, CUgraphNode *dependencies, size_t numDependencies, unsigned int flags);
 CUresult cuEventCreate_f(CUevent* event, unsigned int flags);
 CUresult cuEventDestroy_f(CUevent event);
 CUresult cuEventRecord_f(CUevent event, CUstream stream);
+CUresult cuEventRecordWithFlags_f(CUevent event, CUstream stream, unsigned int flags);
 CUresult cuModuleUnload_f(CUmodule hmod);
 CUresult cuModuleLoadDataEx_f(CUmodule *module, const void *image, unsigned int numOptions, CUjit_option *options, void **optionValues);
 CUresult cuModuleGetFunction_f(CUfunction *hfunc, CUmodule hmod, const char *name);
 CUresult cuLaunchKernel_f(CUfunction f, unsigned int gridDimX, unsigned int gridDimY, unsigned int gridDimZ, unsigned int blockDimX, unsigned int blockDimY, unsigned int blockDimZ, unsigned int sharedMemBytes, CUstream hStream, void **kernelParams, void **extra);
 CUresult cuMemcpyPeerAsync_f(CUdeviceptr dst_ptr, CUcontext dst_ctx, CUdeviceptr src_ptr, CUcontext src_ctx, size_t n, CUstream stream);
+CUresult cuPointerGetAttribute_f(void* data, CUpointer_attribute attribute, CUdeviceptr ptr);
 CUresult cuGraphicsMapResources_f(unsigned int count, CUgraphicsResource* resources, CUstream stream);
 CUresult cuGraphicsUnmapResources_f(unsigned int count, CUgraphicsResource* resources, CUstream hStream);
 CUresult cuGraphicsResourceGetMappedPointer_f(CUdeviceptr* pDevPtr, size_t* pSize, CUgraphicsResource resource);
@@ -86,13 +94,34 @@ CUresult cuGraphicsUnregisterResource_f(CUgraphicsResource resource);
 bool init_cuda_driver();
 bool is_cuda_driver_initialized();
 
-bool check_cuda_result(cudaError_t code, const char* file, int line);
-inline bool check_cuda_result(uint64_t code, const char* file, int line)
+bool check_cuda_result(cudaError_t code, const char* func, const char* file, int line);
+
+inline bool check_cuda_result(uint64_t code, const char* func, const char* file, int line)
+{
+    return check_cuda_result(static_cast<cudaError_t>(code), func, file, line);
+}
+
+bool check_cu_result(CUresult result, const char* func, const char* file, int line);
+
+inline uint64_t get_capture_id(CUstream stream)
 {
-    return check_cuda_result(static_cast<cudaError_t>(code), file, line);
+    CUstreamCaptureStatus status;
+    uint64_t id = 0;
+    check_cu(cuStreamGetCaptureInfo_f(stream, &status, &id, NULL, NULL, NULL));
+    return id;
 }
 
-bool check_cu_result(CUresult result, const char* file, int line);
+inline CUgraph get_capture_graph(CUstream stream)
+{
+    CUstreamCaptureStatus status;
+    CUgraph graph = NULL;
+    check_cu(cuStreamGetCaptureInfo_f(stream, &status, NULL, &graph, NULL, NULL));
+    return graph;
+}
+
+bool get_capture_dependencies(CUstream stream, std::vector<CUgraphNode>& dependencies_ret);
+
+bool get_graph_leaf_nodes(cudaGraph_t graph, std::vector<cudaGraphNode_t>& leaf_nodes_ret);
 
 
 //

warp/native/cutlass_gemm.cpp

@@ -16,7 +16,7 @@ extern "C"
 
 WP_API
 bool cutlass_gemm(
-                  int compute_capability,
+                  void* context, int compute_capability,
                   int m, int n, int k,
                   const char* datatype_str,
                   const void* a, const void* b, const void* c, void* d,

warp/native/cutlass_gemm.cu

@@ -8,6 +8,7 @@
 
 #include "builtin.h"
 #include "temp_buffer.h"
+#include "cuda_util.h"
 
 #include "cutlass/cutlass.h"
 #include "cutlass/gemm/device/gemm_universal.h"
@@ -226,7 +227,7 @@ extern "C" {
 
 WP_API
 bool cutlass_gemm(
-                  int compute_capability,
+                  void* context, int compute_capability,
                   int m, int n, int k,
                   const char* datatype_str,
                   const void* a, const void* b, const void* c, void* d,
@@ -237,6 +238,8 @@ bool cutlass_gemm(
 
     std::string datatype(datatype_str);
 
+    ContextGuard guard(context);
+
     // Specializations for using Tensor Cores and A/B RowMajor/ColumnMajor designations
     if (compute_capability == 80) {
         if (datatype == F64_STR) {

warp/native/error.cpp

@@ -0,0 +1,66 @@
+/** 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.
+ */
+
+#include <stdarg.h>
+#include <stdio.h>
+#include <string.h>
+
+namespace wp
+{
+static char g_error_buffer[4096] = "";
+static bool g_error_output_enabled = true;
+static FILE* g_error_stream = stderr;
+
+const char* get_error_string()
+{
+    return g_error_buffer;
+}
+
+void set_error_string(const char* fmt, ...)
+{
+    va_list args;
+    va_start(args, fmt);
+    vsnprintf(g_error_buffer, sizeof(g_error_buffer), fmt, args);
+    if (g_error_output_enabled)
+    {
+        vfprintf(g_error_stream, fmt, args);
+        fputc('\n', g_error_stream);
+        fflush(g_error_stream);
+    }
+    va_end(args);
+}
+
+void append_error_string(const char* fmt, ...)
+{
+    size_t offset = strlen(g_error_buffer);
+    if (offset + 2 > sizeof(g_error_buffer))
+        return;
+    g_error_buffer[offset++] = '\n';
+    va_list args;
+    va_start(args, fmt);
+    vsnprintf(g_error_buffer + offset, sizeof(g_error_buffer) - offset, fmt, args);
+    if (g_error_output_enabled)
+    {
+        vfprintf(g_error_stream, fmt, args);
+        fputc('\n', g_error_stream);
+        fflush(g_error_stream);
+    }
+    va_end(args);
+}
+
+void set_error_output_enabled(bool enable)
+{
+    g_error_output_enabled = enable;
+}
+
+bool is_error_output_enabled()
+{
+    return g_error_output_enabled;
+}
+
+} // end of namespace wp

warp/native/error.h

@@ -0,0 +1,27 @@
+/** 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.
+ */
+
+#pragma once
+
+namespace wp
+{
+// functions related to error reporting
+
+// get error string from Python
+const char* get_error_string();
+
+// set error message for Python
+// these functions also print the error message if error output is enabled
+void set_error_string(const char* fmt, ...);
+void append_error_string(const char* fmt, ...);
+
+// allow disabling printing errors, which is handy during tests that expect failure
+void set_error_output_enabled(bool enable);
+bool is_error_output_enabled();
+
+}

warp/native/mesh.cu

@@ -203,8 +203,8 @@ uint64_t mesh_create_device(void* context, wp::array_t<wp::vec3> points, wp::arr
         // bvh_destroy_host(bvh_host);
 
         // create lower upper arrays expected by GPU BVH builder
-        mesh.lowers = (wp::vec3*)alloc_temp_device(WP_CURRENT_CONTEXT, sizeof(wp::vec3)*num_tris);
-        mesh.uppers = (wp::vec3*)alloc_temp_device(WP_CURRENT_CONTEXT, sizeof(wp::vec3)*num_tris);
+        mesh.lowers = (wp::vec3*)alloc_device(WP_CURRENT_CONTEXT, sizeof(wp::vec3)*num_tris);
+        mesh.uppers = (wp::vec3*)alloc_device(WP_CURRENT_CONTEXT, sizeof(wp::vec3)*num_tris);
 
         wp_launch_device(WP_CURRENT_CONTEXT, wp::compute_triangle_bounds, num_tris, (num_tris, points.data, indices.data, mesh.lowers, mesh.uppers));
 

warp/native/reduce.cu

@@ -110,7 +110,7 @@ template <typename T> void array_sum_device(const T *ptr_a, T *ptr_out, int coun
 
     size_t buff_size = 0;
     check_cuda(cub::DeviceReduce::Sum(nullptr, buff_size, ptr_strided, ptr_out, count, stream));
-    void* temp_buffer = alloc_temp_device(WP_CURRENT_CONTEXT, buff_size);
+    void* temp_buffer = alloc_device(WP_CURRENT_CONTEXT, buff_size);
 
     for (int k = 0; k < type_length; ++k)
     {
@@ -118,7 +118,7 @@ template <typename T> void array_sum_device(const T *ptr_a, T *ptr_out, int coun
         check_cuda(cub::DeviceReduce::Sum(temp_buffer, buff_size, ptr_strided, ptr_out + k, count, stream));
     }
 
-    free_temp_device(WP_CURRENT_CONTEXT, temp_buffer);
+    free_device(WP_CURRENT_CONTEXT, temp_buffer);
 }
 
 template <typename T>
@@ -271,11 +271,11 @@ void array_inner_device(const ElemT *ptr_a, const ElemT *ptr_b, ScalarT *ptr_out
 
     size_t buff_size = 0;
     check_cuda(cub::DeviceReduce::Sum(nullptr, buff_size, inner_iterator, ptr_out, count, stream));
-    void* temp_buffer = alloc_temp_device(WP_CURRENT_CONTEXT, buff_size);
+    void* temp_buffer = alloc_device(WP_CURRENT_CONTEXT, buff_size);
 
     check_cuda(cub::DeviceReduce::Sum(temp_buffer, buff_size, inner_iterator, ptr_out, count, stream));
 
-    free_temp_device(WP_CURRENT_CONTEXT, temp_buffer);
+    free_device(WP_CURRENT_CONTEXT, temp_buffer);
 }
 
 template <typename T>

warp/native/runlength_encode.cu

@@ -21,13 +21,13 @@ void runlength_encode_device(int n,
         nullptr, buff_size, values, run_values, run_lengths, run_count,
         n, stream));
 
-    void* temp_buffer = alloc_temp_device(WP_CURRENT_CONTEXT, buff_size);
+    void* temp_buffer = alloc_device(WP_CURRENT_CONTEXT, buff_size);
 
     check_cuda(cub::DeviceRunLengthEncode::Encode(
         temp_buffer, buff_size, values, run_values, run_lengths, run_count,
         n, stream));
 
-    free_temp_device(WP_CURRENT_CONTEXT, temp_buffer);
+    free_device(WP_CURRENT_CONTEXT, temp_buffer);
 }
 
 void runlength_encode_int_device(

warp/native/scan.cu

@@ -20,7 +20,7 @@ void scan_device(const T* values_in, T* values_out, int n, bool inclusive)
         check_cuda(cub::DeviceScan::ExclusiveSum(NULL, scan_temp_size, values_in, values_out, n));
     }
 
-    void* temp_buffer = alloc_temp_device(WP_CURRENT_CONTEXT, scan_temp_size);
+    void* temp_buffer = alloc_device(WP_CURRENT_CONTEXT, scan_temp_size);
 
     // scan
     if (inclusive) {
@@ -29,7 +29,7 @@ void scan_device(const T* values_in, T* values_out, int n, bool inclusive)
         check_cuda(cub::DeviceScan::ExclusiveSum(temp_buffer, scan_temp_size, values_in, values_out, n, stream));
     }
 
-    free_temp_device(WP_CURRENT_CONTEXT, temp_buffer);
+    free_device(WP_CURRENT_CONTEXT, temp_buffer);
 }
 
 template void scan_device(const int*, int*, int, bool);

warp/native/sparse.cu

@@ -456,7 +456,6 @@ void bsr_transpose_device(int rows_per_block, int cols_per_block, int row_count,
     size_t buff_size = 0;
     check_cuda(cub::DeviceRadixSort::SortPairs(nullptr, buff_size, d_values,
                                               d_keys, nnz, 0, 64, stream));
-    void* temp_buffer = alloc_temp_device(WP_CURRENT_CONTEXT, buff_size);
     ScopedTemporary<> temp(context, buff_size);
     check_cuda(cub::DeviceRadixSort::SortPairs(
         temp.buffer(), buff_size, d_values, d_keys, nnz, 0, 64, stream));

warp/native/temp_buffer.h

@@ -10,13 +10,13 @@ template <typename T = char> struct ScopedTemporary
 {
 
     ScopedTemporary(void *context, size_t size)
-        : m_context(context), m_buffer(static_cast<T*>(alloc_temp_device(m_context, size * sizeof(T))))
+        : m_context(context), m_buffer(static_cast<T*>(alloc_device(m_context, size * sizeof(T))))
     {
     }
 
     ~ScopedTemporary()
     {
-        free_temp_device(m_context, m_buffer);
+        free_device(m_context, m_buffer);
     }
 
     T *buffer() const