warp-lang 1.7.0__py3-none-manylinux_2_28_x86_64.whl

warp/native/mathdx.cpp

@@ -0,0 +1,87 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "builtin.h"
+
+// stubs for platforms where there is no CUDA
+#if !WP_ENABLE_CUDA || !WP_ENABLE_MATHDX
+
+extern "C"
+{
+
+WP_API
+bool cuda_compile_fft(
+                      const char* ltoir_output_path,
+                      const char* symbol_name, int num_include_dirs,
+                      const char** include_dirs,
+                      const char* mathdx_include_dir,
+                      int arch,
+                      int size,
+                      int elements_per_thread,
+                      int direction,
+                      int precision,
+                      int* shared_memory_size)
+{
+    printf("CUDA is disabled and/or Warp was not compiled with MathDx support.\n");
+    return false;
+}
+
+WP_API bool cuda_compile_dot(
+                             const char* fatbin_output_path,
+                             const char* ltoir_output_path,
+                             const char* symbol_name,
+                             int num_include_dirs,
+                             const char** include_dirs,
+                             const char* mathdx_include_dir,
+                             int arch,
+                             int M,
+                             int N,
+                             int K,
+                             int precision_A,
+                             int precision_B,
+                             int precision_C,
+                             int type,
+                             int a_arrangement,
+                             int b_arrangement,
+                             int c_arrangement,
+                             int num_threads)
+{
+    printf("CUDA is disabled and/or Warp was not compiled with MathDx support.\n");
+    return false;
+}
+
+WP_API bool cuda_compile_solver(
+                                const char* ltoir_output_path,
+                                const char* symbol_name,
+                                int num_include_dirs,
+                                const char** include_dirs,
+                                const char* mathdx_include_dir,
+                                int arch,
+                                int M,
+                                int N,
+                                int function,
+                                int precision,
+                                int fill_mode,
+                                int num_threads)
+{
+    printf("CUDA is disabled and/or Warp was not compiled with MathDx support.\n");
+    return false;
+}
+
+} // extern "C"
+
+#endif // !WP_ENABLE_CUDA || !WP_ENABLE_MATHDX

warp/native/matnn.h

@@ -0,0 +1,343 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+namespace wp
+{
+
+
+CUDA_CALLABLE inline int dense_index(int stride, int i, int j)
+{
+    return i*stride + j;
+}
+
+template <bool transpose>
+CUDA_CALLABLE inline int dense_index(int rows, int cols, int i, int j)
+{
+    if (transpose)
+        return j*rows + i;
+    else
+        return i*cols + j;
+}
+
+
+
+template <bool t1, bool t2, bool add>
+CUDA_CALLABLE inline void dense_gemm_impl(int m, int n, int p, const float* __restrict__ A, const float* __restrict__ B, float* __restrict__ C)
+{
+    for (int i=0; i < m; i++)
+    {
+        for (int j=0; j < n; ++j)
+        {
+            float sum = 0.0f;
+
+            for (int k=0; k < p; ++k)
+            {
+                sum += A[dense_index<t1>(m, p, i, k)]*B[dense_index<t2>(p, n, k, j)];
+            }
+            
+            if (add)
+                C[i*n + j] += sum;
+            else
+                C[i*n + j] = sum;
+        }
+    }
+}
+
+
+template <bool add=false>
+CUDA_CALLABLE inline void dense_gemm(int m, int n, int p, int t1, int t2, const array_t<float>& A, const array_t<float>& B, array_t<float>& C)
+{
+    if (t1 == 0 && t2 == 0)
+        dense_gemm_impl<false, false, add>(m, n, p, A.data, B.data, C.data);
+    else if (t1 == 1 && t2 == 0)
+        dense_gemm_impl<true, false, add>(m, n, p, A.data, B.data, C.data);
+    else if (t1 == 0 && t2 == 1)
+        dense_gemm_impl<false, true, add>(m, n, p, A.data, B.data, C.data);
+    else if (t1 == 1 && t2 == 1)
+        dense_gemm_impl<true, true, add>(m, n, p, A.data, B.data, C.data);
+}
+
+
+
+
+void  CUDA_CALLABLE inline dense_chol(int n, const array_t<float>& A, float regularization, array_t<float>& L)
+{
+    for (int j=0; j < n; ++j)
+    {
+        float s = A.data[dense_index(n, j, j)] + regularization;
+
+        for (int k=0; k < j; ++k)
+        {
+            float r = L.data[dense_index(n, j, k)];
+            s -= r*r;
+        }
+
+        s = sqrt(s);
+        const float invS = 1.0f/s;
+
+        L.data[dense_index(n, j, j)] = s;
+
+        for (int i=j+1; i < n; ++i)
+        {
+            s = A.data[dense_index(n, i, j)];
+            
+            for (int k=0; k < j; ++k)
+            {
+                s -= L.data[dense_index(n, i, k)]*L.data[dense_index(n, j, k)];
+            }
+
+            L.data[dense_index(n, i, j)] = s*invS;
+        }
+    }
+}
+
+
+
+
+// Solves (L*L^T)x = b given the Cholesky factor L 
+CUDA_CALLABLE inline void dense_subs(int n, const array_t<float>& L, const array_t<float>& b, array_t<float>& x)
+{
+    // forward substitution
+    for (int i=0; i < n; ++i)
+    {
+        float s = b.data[i];
+
+        for (int j=0; j < i; ++j)
+        {
+            s -= L.data[dense_index(n, i, j)]*x.data[j];
+        }
+
+        x.data[i] = s/L.data[dense_index(n, i, i)];
+    }
+
+    // backward substitution
+    for (int i=n-1; i >= 0; --i)
+    {
+        float s = x.data[i];
+
+        for (int j=i+1; j < n; ++j)
+        {
+            s -= L.data[dense_index(n, j, i)]*x.data[j];
+        }
+
+        x.data[i] = s/L.data[dense_index(n, i, i)];
+    }
+}
+
+CUDA_CALLABLE inline void dense_solve(int n, const array_t<float>& A, const array_t<float>& L, const array_t<float>& b, array_t<float>& x)
+{
+    dense_subs(n, L, b, x);
+}
+
+
+// CUDA_CALLABLE inline void print_matrix(const char* name, int m, int n, const float* data)
+// {
+//     printf("%s = [", name);
+
+//     for (int i=0; i < m; ++i)
+//     {
+//         for (int j=0; j < n; ++j)
+//         {
+//             printf("%f ", data[dense_index(n, i, j)]);
+//         }
+
+//         printf(";\n");
+//     }
+
+//     printf("]\n");
+// }
+
+// adjoint methods
+CUDA_CALLABLE inline void adj_dense_gemm(
+    int m, int n, int p, int t1, int t2, const array_t<float>& A, const array_t<float>& B, array_t<float>& C,
+    int adj_m, int adj_n, int adj_p, int adj_t1, int adj_t2, array_t<float>& adj_A, array_t<float>& adj_B, const array_t<float>& adj_C)
+{
+
+    // print_matrix("A", m, p, A);
+    // print_matrix("B", p, n, B);
+    // printf("t1: %d t2: %d\n", t1, t2);
+
+    if (t1)
+    {
+        dense_gemm<true>(p, m, n, 0, 1, B, adj_C, adj_A);
+        dense_gemm<true>(p, n, m, int(!t1), 0, A, adj_C, adj_B);
+    }
+    else
+    {
+        dense_gemm<true>(m, p, n, 0, int(!t2), adj_C, B, adj_A);
+        dense_gemm<true>(p, n, m, int(!t1), 0, A, adj_C, adj_B);
+    }
+}
+
+
+CUDA_CALLABLE inline void adj_dense_chol(
+    int n, const array_t<float>& A, float regularization, array_t<float>& L,
+    int adj_n, const array_t<float>& adj_A, float adj_regularization, array_t<float>& adj_L)
+{
+    // nop, use dense_solve to differentiate through (A^-1)b = x
+}
+
+CUDA_CALLABLE inline void adj_dense_subs(
+    int n, const array_t<float>& L, const array_t<float>& b, array_t<float>& x,
+    int adj_n, const array_t<float>& adj_L, const array_t<float>& adj_b, array_t<float>& adj_x)
+{
+    // nop, use dense_solve to differentiate through (A^-1)b = x
+}
+
+
+CUDA_CALLABLE inline void adj_dense_solve(int n,
+    const array_t<float>& A, const array_t<float>& L, const array_t<float>& b, const array_t<float>& x,
+    int adj_n, array_t<float>& adj_A, array_t<float>& adj_L, array_t<float>& adj_b, const array_t<float>& adj_x)
+{
+    // see https://people.maths.ox.ac.uk/gilesm/files/NA-08-01.pwp, section 2.3.1
+    dense_subs(n, L, adj_x, adj_b);
+
+    // A* = -adj_b*x^T
+    for (int i=0; i < n; ++i)
+    {
+        for (int j=0; j < n; ++j)
+        {
+            adj_A.data[dense_index(n, i, j)] += -adj_b.data[i]*x.data[j];
+        }
+    }
+}
+
+
+template <typename F>
+CUDA_CALLABLE inline void mlp(const array_t<float>& weights, const array_t<float>& bias, F activation, int index, const array_t<float>& x, array_t<float>& out)
+{
+    const int m = weights.shape[0];
+    const int n = weights.shape[1];
+    const int b = x.shape[1];
+
+    for (int i=0; i < m; ++i)
+    {
+        float tmp = bias.data[i];
+
+        for(int j=0; j < n; ++j)
+        {
+            tmp += weights.data[i*n + j]*x.data[index + b*j];
+        }
+
+        out.data[index + b*i] = activation(tmp);
+    }
+}
+
+template <typename F, typename AdjF>
+CUDA_CALLABLE inline void adj_mlp(const array_t<float>& weights, const array_t<float>& bias, F activation, int index, const array_t<float>& x, array_t<float>& out,
+                                  array_t<float>& adj_weights, array_t<float>& adj_bias, AdjF adj_activation, int adj_index, array_t<float>& adj_x, array_t<float>& adj_out)
+{
+    const int m = weights.shape[0];
+    const int n = weights.shape[1];
+    const int b = x.shape[1];
+
+    for (int i=0; i < m; ++i)
+    {
+        // recompute forward pass so we don't have to store pre-activation outputs
+        float tmp = bias.data[i];
+
+        for(int j=0; j < n; ++j)
+        {
+            tmp += weights.data[i*n + j]*x.data[index + b*j];
+        }
+
+        // adjoint w.r.t to activation
+        float adj_f = 0.0f;
+    
+        if (adj_out.data)
+            adj_activation(tmp, adj_f, adj_out.data[index + b*i]);
+
+        for (int j=0; j < n; ++j)
+        {
+            // adjoint w.r.t M_i
+            if (adj_weights.data)
+                atomic_add(&adj_weights.data[i*n + j], x.data[index + b*j]*adj_f);    // todo: reduce these atomic stores using warp/block level reductions
+
+            // adjoint w.r.t x
+            if (adj_x.data)
+                atomic_add(&adj_x.data[index + b*j], weights.data[i*n + j]*adj_f);
+        }
+
+        // adjoint w.r.t b
+        if (adj_bias.data)
+            atomic_add(&adj_bias.data[i], adj_f);
+
+    }
+}
+
+
+// template <typename F>
+// CUDA_CALLABLE inline void mlp(const array_t<float>& weights, const array_t<float>& bias, F activation, int m, int n, int b, int index, const array_t<float>& x, array_t<float>& out)
+// {
+//     x += index*n;
+//     out += index*m;
+
+
+//     for (int i=0; i < m; ++i)
+//     {
+//         float tmp = bias[i];
+
+//         for(int j=0; j < n; ++j)
+//         {
+//             tmp += weights[i*n + j]*x[j];
+//         }
+
+//         out[i] = activation(tmp);
+//     }
+// }
+
+// template <typename F, typename AdjF>
+// CUDA_CALLABLE inline void adj_mlp(const array_t<float>& weights, const array_t<float>& bias, F activation, int m, int n, int b, int index, const array_t<float>& x, const array_t<float>& out,
+//                                   array_t<float>& adj_weights, array_t<float>& adj_bias, AdjF adj_activation, int adj_m, int adj_n, int adj_b, int adj_index, array_t<float>& adj_x, array_t<float>& adj_out)
+// {
+//     x += index*n;
+//     out += index*m;
+
+//     adj_x += index*n;
+//     adj_out += index*m;
+
+//     for (int i=0; i < m; ++i)
+//     {
+//         // recompute forward pass so we don't have to store pre-activation outputs
+//         float tmp = bias[i];
+
+//         for(int j=0; j < n; ++j)
+//         {
+//             tmp += weights[i*n + j]*x[index + b*j];            
+//         }
+
+//         // adjoint w.r.t to activation
+//         float adj_f = 0.0f;
+//         adj_activation(tmp, adj_f, adj_out[index + b*i]);
+
+//         for (int j=0; j < n; ++j)
+//         {
+//             // adjoint w.r.t M_i
+//             adj_weights[i*n + j] += x[j]*adj_f;
+
+//             // adjoint w.r.t x
+//             adj_x[index + b*j] += weights[i*n + j]*adj_f;
+//         }
+
+//         // adjoint w.r.t b
+//         adj_bias[i] += adj_f;
+//     }
+// }
+
+} // namespace wp

warp/native/mesh.cpp

@@ -0,0 +1,266 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "mesh.h"
+#include "bvh.h"
+#include "warp.h"
+#include "cuda_util.h"
+
+using namespace wp;
+
+#include <map>
+
+namespace 
+{
+    // host-side copy of mesh descriptors, maps GPU mesh address (id) to a CPU desc
+    std::map<uint64_t, Mesh> g_mesh_descriptors;
+
+} // anonymous namespace
+
+
+namespace wp
+{
+
+bool mesh_get_descriptor(uint64_t id, Mesh& mesh)
+{
+    const auto& iter = g_mesh_descriptors.find(id);
+    if (iter == g_mesh_descriptors.end())
+        return false;
+    else
+        mesh = iter->second;
+        return true;
+}
+
+bool mesh_set_descriptor(uint64_t id, const Mesh& mesh)
+{
+    const auto& iter = g_mesh_descriptors.find(id);
+    if (iter == g_mesh_descriptors.end())
+        return false;
+    else
+        iter->second = mesh;
+    return true;
+}
+
+void mesh_add_descriptor(uint64_t id, const Mesh& mesh)
+{
+    g_mesh_descriptors[id] = mesh;
+}
+
+void mesh_rem_descriptor(uint64_t id)
+{
+    g_mesh_descriptors.erase(id);
+
+}
+
+} // namespace wp
+
+void bvh_refit_with_solid_angle_recursive_host(BVH& bvh, int index, Mesh& mesh)
+{
+    BVHPackedNodeHalf& lower = bvh.node_lowers[index];
+    BVHPackedNodeHalf& upper = bvh.node_uppers[index];
+
+    if (lower.b)
+    {
+        // Leaf, compute properties
+         const int start = lower.i;
+         const int end = upper.i;
+         // loops through primitives in the leaf
+         for (int primitive_counter = start; primitive_counter < end; primitive_counter++)
+         {
+             int primitive_index = mesh.bvh.primitive_indices[primitive_counter];
+             if (primitive_counter == start)
+             {
+                 precompute_triangle_solid_angle_props(mesh.points[mesh.indices[primitive_index * 3 + 0]], mesh.points[mesh.indices[primitive_index * 3 + 1]],
+                     mesh.points[mesh.indices[primitive_index * 3 + 2]], mesh.solid_angle_props[index]);
+             }
+             else
+             {
+                 SolidAngleProps triangle_solid_angle_props;
+                 precompute_triangle_solid_angle_props(mesh.points[mesh.indices[primitive_index * 3 + 0]], mesh.points[mesh.indices[primitive_index * 3 + 1]],
+                     mesh.points[mesh.indices[primitive_index * 3 + 2]], triangle_solid_angle_props);
+                 mesh.solid_angle_props[index] = combine_precomputed_solid_angle_props(&mesh.solid_angle_props[index], &triangle_solid_angle_props);
+             }
+         }
+
+         (vec3&)lower = mesh.solid_angle_props[index].box.lower;
+         (vec3&)upper = mesh.solid_angle_props[index].box.upper;
+    }
+    else
+    {
+        int left_index = lower.i;
+        int right_index = upper.i;
+
+        bvh_refit_with_solid_angle_recursive_host(bvh, left_index, mesh);
+        bvh_refit_with_solid_angle_recursive_host(bvh, right_index, mesh);
+
+        // combine
+        SolidAngleProps* left_child_data = &mesh.solid_angle_props[left_index];
+        SolidAngleProps* right_child_data = (left_index != right_index) ? &mesh.solid_angle_props[right_index] : NULL;
+        
+        combine_precomputed_solid_angle_props(mesh.solid_angle_props[index], left_child_data, right_child_data);
+
+        // compute union of children
+        const vec3& left_lower = (vec3&)bvh.node_lowers[left_index];
+        const vec3& left_upper = (vec3&)bvh.node_uppers[left_index];
+
+        const vec3& right_lower = (vec3&)bvh.node_lowers[right_index];
+        const vec3& right_upper = (vec3&)bvh.node_uppers[right_index];
+
+        // union of child bounds
+        vec3 new_lower = min(left_lower, right_lower);
+        vec3 new_upper = max(left_upper, right_upper);
+        
+        // write new BVH nodes
+        (vec3&)lower = new_lower;
+        (vec3&)upper = new_upper;        
+    }
+}
+
+void bvh_refit_with_solid_angle_host(BVH& bvh, Mesh& mesh)
+{
+    bvh_refit_with_solid_angle_recursive_host(bvh, 0, mesh);
+}
+
+uint64_t mesh_create_host(array_t<wp::vec3> points, array_t<wp::vec3> velocities, array_t<int> indices, int num_points, int num_tris, int support_winding_number, int constructor_type)
+{
+    Mesh* m = new Mesh(points, velocities, indices, num_points, num_tris);
+
+    m->lowers = new vec3[num_tris];
+    m->uppers = new vec3[num_tris];
+
+    float sum = 0.0;
+    for (int i=0; i < num_tris; ++i)
+    {
+        wp::vec3& p0 = points[indices[i*3+0]];
+        wp::vec3& p1 = points[indices[i*3+1]];
+        wp::vec3& p2 = points[indices[i*3+2]];
+        
+        // compute triangle bounds
+        bounds3 b;       
+        b.add_point(p0);
+        b.add_point(p1);
+        b.add_point(p2);
+
+        m->lowers[i] = b.lower;
+        m->uppers[i] = b.upper;
+
+        // compute edge lengths
+        sum += length(p0-p1) + length(p0-p2) + length(p2-p1);
+    }
+    m->average_edge_length = sum / (num_tris*3);
+
+    wp::bvh_create_host(m->lowers, m->uppers, num_tris, constructor_type, m->bvh);
+    
+    if (support_winding_number) 
+    {
+        // Let's first compute the sold
+        int num_bvh_nodes = 2*num_tris-1;
+        m->solid_angle_props = new SolidAngleProps[num_bvh_nodes];
+        bvh_refit_with_solid_angle_host(m->bvh, *m);
+    }
+
+    return (uint64_t)m;
+}
+
+
+void mesh_destroy_host(uint64_t id)
+{
+    Mesh* m = (Mesh*)(id);
+
+    delete[] m->lowers;
+    delete[] m->uppers;
+
+    if (m->solid_angle_props) {
+        delete [] m->solid_angle_props;
+    }
+    bvh_destroy_host(m->bvh);
+
+    delete m;
+}
+
+void mesh_refit_host(uint64_t id)
+{
+    Mesh* m = (Mesh*)(id);
+
+    float sum = 0.0;
+    for (int i=0; i < m->num_tris; ++i)
+    {
+        wp::vec3 p0 = m->points.data[m->indices.data[i*3+0]];
+        wp::vec3 p1 = m->points.data[m->indices.data[i*3+1]];
+        wp::vec3 p2 = m->points.data[m->indices.data[i*3+2]];
+
+        // compute triangle bounds
+        bounds3 b;       
+        b.add_point(p0);
+        b.add_point(p1);
+        b.add_point(p2);
+
+        m->lowers[i] = b.lower;
+        m->uppers[i] = b.upper;
+
+        sum += length(p0-p1) + length(p0-p2) + length(p2-p1);
+    }
+    m->average_edge_length = sum / (m->num_tris*3);
+
+    if (m->solid_angle_props)
+    {
+        // If solid angle were used, use refit solid angle
+        bvh_refit_with_solid_angle_host(m->bvh, *m);
+    }
+    else
+    {
+        bvh_refit_host(m->bvh);
+    }
+}
+
+void mesh_set_points_host(uint64_t id, wp::array_t<wp::vec3> points)
+{
+    Mesh* m = (Mesh*)(id);
+    if (points.ndim != 1 || points.shape[0] != m->points.shape[0])
+    {
+        fprintf(stderr, "The new points input for mesh_set_points_host does not match the shape of the original points!\n");
+        return;
+    }
+
+    m->points = points;
+
+    mesh_refit_host(id);
+}
+
+void mesh_set_velocities_host(uint64_t id, wp::array_t<wp::vec3> velocities)
+{
+    Mesh* m = (Mesh*)(id);
+    if (velocities.ndim != 1 || velocities.shape[0] != m->velocities.shape[0])
+    {
+        fprintf(stderr, "The new velocities input for mesh_set_velocities_host does not match the shape of the original velocities!\n");
+        return;
+    }
+    m->velocities = velocities;
+}
+
+// stubs for non-CUDA platforms
+#if !WP_ENABLE_CUDA
+
+
+WP_API uint64_t mesh_create_device(void* context, wp::array_t<wp::vec3> points, wp::array_t<wp::vec3> velocities, wp::array_t<int> tris, int num_points, int num_tris, int support_winding_number, int constructor_type) { return 0; }
+WP_API void mesh_destroy_device(uint64_t id) {}
+WP_API void mesh_refit_device(uint64_t id) {}
+WP_API void mesh_set_points_device(uint64_t id, wp::array_t<wp::vec3> points) {};
+WP_API void mesh_set_velocities_device(uint64_t id, wp::array_t<wp::vec3> points) {};
+
+
+#endif // !WP_ENABLE_CUDA