warp-lang 1.10.0__py3-none-macosx_11_0_arm64.whl

warp/native/mathdx.cpp

@@ -0,0 +1,93 @@
+/*
+ * 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 wp_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 wp_cuda_compile_dot(
+                                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 arrangement_A,
+                                int arrangement_B,
+                                int arrangement_C,
+                                int num_threads)
+{
+    printf("CUDA is disabled and/or Warp was not compiled with MathDx support.\n");
+    return false;
+}
+
+WP_API bool wp_cuda_compile_solver(
+                                   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 NRHS,
+                                   int function,
+                                   int side,
+                                   int diag,
+                                   int precision,
+                                   int arrangement_A,
+                                   int arrangement_B,
+                                   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,221 @@
+/*
+ * 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];
+        }
+    }
+}
+
+} // 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);
+             }
+         }
+
+         reinterpret_cast<vec3&>(lower) = mesh.solid_angle_props[index].box.lower;
+         reinterpret_cast<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 = reinterpret_cast<const vec3&>(bvh.node_lowers[left_index]);
+        const vec3& left_upper = reinterpret_cast<const vec3&>(bvh.node_uppers[left_index]);
+
+        const vec3& right_lower = reinterpret_cast<const vec3&>(bvh.node_lowers[right_index]);
+        const vec3& right_upper = reinterpret_cast<const 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
+        reinterpret_cast<vec3&>(lower) = new_lower;
+        reinterpret_cast<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 wp_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, int bvh_leaf_size)
+{
+    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, bvh_leaf_size);
+    
+    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 wp_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;
+    }
+    wp::bvh_destroy_host(m->bvh);
+
+    delete m;
+}
+
+void wp_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
+    {
+        wp::bvh_refit_host(m->bvh);
+    }
+}
+
+void wp_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 wp_mesh_set_points_host does not match the shape of the original points!\n");
+        return;
+    }
+
+    m->points = points;
+
+    wp_mesh_refit_host(id);
+}
+
+void wp_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 wp_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 wp_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, int bvh_leaf_size) { return 0; }
+WP_API void wp_mesh_destroy_device(uint64_t id) {}
+WP_API void wp_mesh_refit_device(uint64_t id) {}
+WP_API void wp_mesh_set_points_device(uint64_t id, wp::array_t<wp::vec3> points) {};
+WP_API void wp_mesh_set_velocities_device(uint64_t id, wp::array_t<wp::vec3> points) {};
+
+
+#endif // !WP_ENABLE_CUDA