warp-lang 1.0.2__py3-none-win_amd64.whl → 1.2.0__py3-none-win_amd64.whl

warp/native/matnn.h

@@ -1,334 +1,334 @@
-/** Copyright (c) 2022 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
-{
-
-
-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;
-//     }
-// }
-
+/** Copyright (c) 2022 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
+{
+
+
+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