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

warp/native/marching.cu

@@ -1,8 +1,6 @@
 #include "warp.h"
 #include "cuda_util.h"
-
-#include "thrust/device_ptr.h"
-#include "thrust/sort.h"
+#include "scan.h"
 
 namespace wp {
 
@@ -162,13 +160,17 @@ namespace wp {
     };
 
 
-	// ---------------------------------------------------------------------------------------
-	struct MarchingCubes
-	{
-		MarchingCubes() 
+    // ---------------------------------------------------------------------------------------
+    struct MarchingCubes
+    {
+        MarchingCubes() 
         {
-			memset(this, 0, sizeof(MarchingCubes));
-		}
+            memset(this, 0, sizeof(MarchingCubes));
+            first_cell_vert = nullptr;
+            first_cell_tri = nullptr;
+            cell_verts = nullptr;
+            context = nullptr;
+        }
 
         __device__ __host__ int cell_index(int xi, int yi, int zi) const
         {
@@ -181,169 +183,169 @@ namespace wp {
             xi = cell_index / ny;
         }
 
-		// grid
+        // grid
         int nx;
         int ny;
         int nz;
 
-		int* first_cell_vert;
-		int* first_cell_tri;
-		int* cell_verts;
+        int* first_cell_vert;
+        int* first_cell_tri;
+        int* cell_verts;
 
         int num_cells;   
         int max_cells;
 
         void* context;
-	};
-
-
-	// -----------------------------------------------------------------------------------
-	__global__ void count_cell_verts(MarchingCubes mc, const float* density, float threshold)
-	{
-		int cell_index = blockIdx.x * blockDim.x + threadIdx.x;
-		if (cell_index >= mc.num_cells)
-			return;
-
-		int xi, yi, zi;
-		mc.cell_coord(cell_index, xi, yi, zi);
-
-		mc.first_cell_vert[cell_index] = 0;
-		if (xi >= mc.nx - 1 || yi >= mc.ny - 1 || zi >= mc.nz - 1)
-			return;
-
-		float d0 = density[cell_index];
-		float dx = density[mc.cell_index(xi + 1, yi, zi)];
-		float dy = density[mc.cell_index(xi, yi + 1, zi)];
-		float dz = density[mc.cell_index(xi, yi, zi + 1)];
-
-		int num = 0;
-		if ((d0 <= threshold && dx >= threshold) || (dx <= threshold && d0 >= threshold))
-			num++;
-		if ((d0 <= threshold && dy >= threshold) || (dy <= threshold && d0 >= threshold))
-			num++;
-		if ((d0 <= threshold && dz >= threshold) || (dz <= threshold && d0 >= threshold))
-			num++;
-
-		mc.first_cell_vert[cell_index] = num;
-	}
-
-	// -----------------------------------------------------------------------------------
-	__global__ void create_cell_verts(MarchingCubes mc, vec3* __restrict__ vertices, vec3* normals, const float* __restrict__  density, float threshold)
-	{
-		int cell_index = blockIdx.x * blockDim.x + threadIdx.x;
-		if (cell_index >= mc.num_cells)
-			return;
-
-		int xi, yi, zi;
-		mc.cell_coord(cell_index, xi, yi, zi);
-		if (xi >= mc.nx - 1 || yi >= mc.ny - 1 || zi >= mc.nz - 1)
-			return;
-
-		vec3 p = vec3(xi + 0.5f, yi + 0.5f, zi + 0.5f);
-
-		float d0 = density[cell_index];
-		float ds[3];
-		ds[0] = density[mc.cell_index(xi + 1, yi, zi)];
-		ds[1] = density[mc.cell_index(xi, yi + 1, zi)];
-		ds[2] = density[mc.cell_index(xi, yi, zi + 1)];
-
-		// vec3 n0 = densityNormal[cell_index];
-		// vec3 ns[3];
-		// ns[0] = densityNormal[mc.cell_index(xi + 1, yi, zi)];
-		// ns[1] = densityNormal[mc.cell_index(xi, yi + 1, zi)];
-		// ns[2] = densityNormal[mc.cell_index(xi, yi, zi + 1)];
-
-		int first = mc.first_cell_vert[cell_index];
-
-		for (int dim = 0; dim < 3; dim++) 
+    };
+
+
+    // -----------------------------------------------------------------------------------
+    __global__ void count_cell_verts(MarchingCubes mc, const float* density, float threshold)
+    {
+        int cell_index = blockIdx.x * blockDim.x + threadIdx.x;
+        if (cell_index >= mc.num_cells)
+            return;
+
+        int xi, yi, zi;
+        mc.cell_coord(cell_index, xi, yi, zi);
+
+        mc.first_cell_vert[cell_index] = 0;
+        if (xi >= mc.nx - 1 || yi >= mc.ny - 1 || zi >= mc.nz - 1)
+            return;
+
+        float d0 = density[cell_index];
+        float dx = density[mc.cell_index(xi + 1, yi, zi)];
+        float dy = density[mc.cell_index(xi, yi + 1, zi)];
+        float dz = density[mc.cell_index(xi, yi, zi + 1)];
+
+        int num = 0;
+        if ((d0 <= threshold && dx >= threshold) || (dx <= threshold && d0 >= threshold))
+            num++;
+        if ((d0 <= threshold && dy >= threshold) || (dy <= threshold && d0 >= threshold))
+            num++;
+        if ((d0 <= threshold && dz >= threshold) || (dz <= threshold && d0 >= threshold))
+            num++;
+
+        mc.first_cell_vert[cell_index] = num;
+    }
+
+    // -----------------------------------------------------------------------------------
+    __global__ void create_cell_verts(MarchingCubes mc, vec3* __restrict__ vertices, vec3* normals, const float* __restrict__  density, float threshold)
+    {
+        int cell_index = blockIdx.x * blockDim.x + threadIdx.x;
+        if (cell_index >= mc.num_cells)
+            return;
+
+        int xi, yi, zi;
+        mc.cell_coord(cell_index, xi, yi, zi);
+        if (xi >= mc.nx - 1 || yi >= mc.ny - 1 || zi >= mc.nz - 1)
+            return;
+
+        vec3 p = vec3(xi + 0.5f, yi + 0.5f, zi + 0.5f);
+
+        float d0 = density[cell_index];
+        float ds[3];
+        ds[0] = density[mc.cell_index(xi + 1, yi, zi)];
+        ds[1] = density[mc.cell_index(xi, yi + 1, zi)];
+        ds[2] = density[mc.cell_index(xi, yi, zi + 1)];
+
+        // vec3 n0 = densityNormal[cell_index];
+        // vec3 ns[3];
+        // ns[0] = densityNormal[mc.cell_index(xi + 1, yi, zi)];
+        // ns[1] = densityNormal[mc.cell_index(xi, yi + 1, zi)];
+        // ns[2] = densityNormal[mc.cell_index(xi, yi, zi + 1)];
+
+        int first = mc.first_cell_vert[cell_index];
+
+        for (int dim = 0; dim < 3; dim++) 
         {
-			float d = ds[dim];
-			mc.cell_verts[3 * cell_index + dim] = 0;
+            float d = ds[dim];
+            mc.cell_verts[3 * cell_index + dim] = 0;
 
-			if ((d0 <= threshold && d >= threshold) || (d <= threshold && d0 >= threshold)) 
+            if ((d0 <= threshold && d >= threshold) || (d <= threshold && d0 >= threshold)) 
             {
-				float t = (d != d0) ? clamp((threshold - d0) / (d - d0), 0.0f, 1.0f) : 0.5f;
-				int id = first++;
-
-				vec3 off;
-				off[dim] = t;
-				vertices[id] = p + off;
-
-				// vec3 n = normalize(n0 + t * (ns[dim] - n0));
-				// normals[id] = -n;
-
-				mc.cell_verts[3 * cell_index + dim] = id;
-			}
-		}
-	}
-
-	// -----------------------------------------------------------------------------------
-	__global__ void count_cell_tris(MarchingCubes mc, const float* __restrict__ density, float threshold)
-	{
-		int cell_index = blockIdx.x * blockDim.x + threadIdx.x;
-		if (cell_index >= mc.num_cells)
-			return;
-
-		int xi, yi, zi;
-		mc.cell_coord(cell_index, xi, yi, zi);
-
-		mc.first_cell_tri[cell_index] = 0;
-		if (xi >= mc.nx - 2 || yi >= mc.ny - 2 || zi >= mc.nz - 2)
-			return;
-
-		int code = 0;
-		for (int i = 0; i < 8; i++) {
-			int cxi = xi + marchingCubeCorners[i][0];
-			int cyi = yi + marchingCubeCorners[i][1];
-			int czi = zi + marchingCubeCorners[i][2];
-
-			if (density[mc.cell_index(cxi, cyi, czi)] >= threshold)
-				code |= (1 << i);
-		}
-
-		mc.first_cell_tri[cell_index] = firstMarchingCubesId[code + 1] - firstMarchingCubesId[code];
-	}
-
-	// -----------------------------------------------------------------------------------
-	__global__ void create_cell_tris(MarchingCubes mc, const float* __restrict__ density, int* __restrict__ triangles, float threshold)
-	{
-		int cell_index = blockIdx.x * blockDim.x + threadIdx.x;
-		if (cell_index >= mc.num_cells)
-			return;
-
-		int xi, yi, zi;
-		mc.cell_coord(cell_index, xi, yi, zi);
-		if (xi >= mc.nx - 2 || yi >= mc.ny - 2 || zi >= mc.nz - 2)
-			return;
-
-		int code = 0;
-		for (int i = 0; i < 8; i++)
+                float t = (d != d0) ? clamp((threshold - d0) / (d - d0), 0.0f, 1.0f) : 0.5f;
+                int id = first++;
+
+                vec3 off;
+                off[dim] = t;
+                vertices[id] = p + off;
+
+                // vec3 n = normalize(n0 + t * (ns[dim] - n0));
+                // normals[id] = -n;
+
+                mc.cell_verts[3 * cell_index + dim] = id;
+            }
+        }
+    }
+
+    // -----------------------------------------------------------------------------------
+    __global__ void count_cell_tris(MarchingCubes mc, const float* __restrict__ density, float threshold)
+    {
+        int cell_index = blockIdx.x * blockDim.x + threadIdx.x;
+        if (cell_index >= mc.num_cells)
+            return;
+
+        int xi, yi, zi;
+        mc.cell_coord(cell_index, xi, yi, zi);
+
+        mc.first_cell_tri[cell_index] = 0;
+        if (xi >= mc.nx - 2 || yi >= mc.ny - 2 || zi >= mc.nz - 2)
+            return;
+
+        int code = 0;
+        for (int i = 0; i < 8; i++) {
+            int cxi = xi + marchingCubeCorners[i][0];
+            int cyi = yi + marchingCubeCorners[i][1];
+            int czi = zi + marchingCubeCorners[i][2];
+
+            if (density[mc.cell_index(cxi, cyi, czi)] >= threshold)
+                code |= (1 << i);
+        }
+
+        mc.first_cell_tri[cell_index] = firstMarchingCubesId[code + 1] - firstMarchingCubesId[code];
+    }
+
+    // -----------------------------------------------------------------------------------
+    __global__ void create_cell_tris(MarchingCubes mc, const float* __restrict__ density, int* __restrict__ triangles, float threshold)
+    {
+        int cell_index = blockIdx.x * blockDim.x + threadIdx.x;
+        if (cell_index >= mc.num_cells)
+            return;
+
+        int xi, yi, zi;
+        mc.cell_coord(cell_index, xi, yi, zi);
+        if (xi >= mc.nx - 2 || yi >= mc.ny - 2 || zi >= mc.nz - 2)
+            return;
+
+        int code = 0;
+        for (int i = 0; i < 8; i++)
         {
-			int cxi = xi + marchingCubeCorners[i][0];
-			int cyi = yi + marchingCubeCorners[i][1];
-			int czi = zi + marchingCubeCorners[i][2];
+            int cxi = xi + marchingCubeCorners[i][0];
+            int cyi = yi + marchingCubeCorners[i][1];
+            int czi = zi + marchingCubeCorners[i][2];
 
-			if (density[mc.cell_index(cxi, cyi, czi)] >= threshold)
-				code |= (1 << i);
-		}
+            if (density[mc.cell_index(cxi, cyi, czi)] >= threshold)
+                code |= (1 << i);
+        }
 
-		int firstIn = firstMarchingCubesId[code];
-		int num = firstMarchingCubesId[code + 1] - firstIn;
-		int firstOut = mc.first_cell_tri[cell_index];
+        int firstIn = firstMarchingCubesId[code];
+        int num = firstMarchingCubesId[code + 1] - firstIn;
+        int firstOut = mc.first_cell_tri[cell_index];
 
-		for (int i = 0; i < num; i++)
+        for (int i = 0; i < num; i++)
         {
-			int eid = marchingCubesIds[firstIn + i];
+            int eid = marchingCubesIds[firstIn + i];
 
-			int exi = xi + marchingCubesEdgeLocations[eid][0];
-			int eyi = yi + marchingCubesEdgeLocations[eid][1];
-			int ezi = zi + marchingCubesEdgeLocations[eid][2];
-			int edgeNr = marchingCubesEdgeLocations[eid][3];
+            int exi = xi + marchingCubesEdgeLocations[eid][0];
+            int eyi = yi + marchingCubesEdgeLocations[eid][1];
+            int ezi = zi + marchingCubesEdgeLocations[eid][2];
+            int edgeNr = marchingCubesEdgeLocations[eid][3];
 
-			int id = mc.cell_verts[3 * mc.cell_index(exi, eyi, ezi) + edgeNr];
-			triangles[firstOut + i] = id;
-		}
-	}
+            int id = mc.cell_verts[3 * mc.cell_index(exi, eyi, ezi) + edgeNr];
+            triangles[firstOut + i] = id;
+        }
+    }
 
     // -------------------------
     void marching_cubes_resize(MarchingCubes& mc, int nx, int ny, int nz)
@@ -444,10 +446,7 @@ WP_API int marching_cubes_surface_device(
     int num_last;		
     memcpy_d2h(WP_CURRENT_CONTEXT, &num_last, &mc.first_cell_vert[mc.num_cells - 1], sizeof(int));
 
-    thrust::exclusive_scan(
-        thrust::device_ptr<int>(mc.first_cell_vert),
-        thrust::device_ptr<int>(mc.first_cell_vert + mc.num_cells),
-        thrust::device_ptr<int>(mc.first_cell_vert));
+    scan_device(mc.first_cell_vert, mc.first_cell_vert, mc.num_cells, false);
 
     int num_verts;
     memcpy_d2h(WP_CURRENT_CONTEXT, &num_verts, &mc.first_cell_vert[mc.num_cells - 1], sizeof(int));
@@ -472,10 +471,7 @@ WP_API int marching_cubes_surface_device(
     
     memcpy_d2h(WP_CURRENT_CONTEXT, &num_last, &mc.first_cell_tri[mc.num_cells - 1], sizeof(int));
 
-    thrust::exclusive_scan(
-        thrust::device_ptr<int>(mc.first_cell_tri),
-        thrust::device_ptr<int>(mc.first_cell_tri + mc.num_cells),
-        thrust::device_ptr<int>(mc.first_cell_tri));
+    scan_device(mc.first_cell_tri, mc.first_cell_tri, mc.num_cells, false);
 
 
     int num_indices;

warp/native/mat.h

@@ -21,7 +21,9 @@ struct quat_t;
 template<unsigned Rows, unsigned Cols, typename Type>
 struct mat_t
 {
-    inline mat_t() = default;
+    inline CUDA_CALLABLE mat_t()
+        : data()
+    {}
 
     inline CUDA_CALLABLE mat_t(Type s)
     {
@@ -30,6 +32,14 @@ struct mat_t
                 data[i][j] = s;
     }
     
+    template <typename OtherType>
+    inline explicit CUDA_CALLABLE mat_t(const mat_t<Rows, Cols, OtherType>& other)
+    {
+        for (unsigned i=0; i < Rows; ++i)
+            for (unsigned j=0; j < Cols; ++j)
+                data[i][j] = other.data[i][j];
+    }
+    
     inline CUDA_CALLABLE mat_t(vec_t<2,Type> c0, vec_t<2,Type> c1)
     {
         data[0][0] = c0[0];
@@ -185,7 +195,7 @@ struct mat_t
     }
 
     // row major storage assumed to be compatible with PyTorch
-    Type data[Rows][Cols] = {};
+    Type data[Rows][Cols];
 };
 
 
@@ -290,7 +300,19 @@ inline CUDA_CALLABLE mat_t<Rows,Cols,Type> atomic_max(mat_t<Rows,Cols,Type> * ad
 }
 
 template<unsigned Rows, unsigned Cols, typename Type>
-inline CUDA_CALLABLE vec_t<Cols,Type> index(const mat_t<Rows,Cols,Type>& m, int row)
+inline CUDA_CALLABLE void adj_atomic_minmax(
+    mat_t<Rows,Cols,Type> *addr,
+    mat_t<Rows,Cols,Type> *adj_addr,
+    const mat_t<Rows,Cols,Type> &value,
+    mat_t<Rows,Cols,Type> &adj_value)
+{
+    for (unsigned i=0; i < Rows; ++i)
+        for (unsigned j=0; j < Cols; ++j)
+            adj_atomic_minmax(&addr->data[i][j], &adj_addr->data[i][j], value.data[i][j], adj_value.data[i][j]);
+}
+
+template<unsigned Rows, unsigned Cols, typename Type>
+inline CUDA_CALLABLE vec_t<Cols,Type> extract(const mat_t<Rows,Cols,Type>& m, int row)
 {
     vec_t<Cols,Type> ret;
     for(unsigned i=0; i < Cols; ++i)
@@ -301,7 +323,7 @@ inline CUDA_CALLABLE vec_t<Cols,Type> index(const mat_t<Rows,Cols,Type>& m, int
 }
 
 template<unsigned Rows, unsigned Cols, typename Type>
-inline CUDA_CALLABLE Type index(const mat_t<Rows,Cols,Type>& m, int row, int col)
+inline CUDA_CALLABLE Type extract(const mat_t<Rows,Cols,Type>& m, int row, int col)
 {
 #ifndef NDEBUG
     if (row < 0 || row >= Rows)
@@ -319,7 +341,7 @@ inline CUDA_CALLABLE Type index(const mat_t<Rows,Cols,Type>& m, int row, int col
 }
 
 template<unsigned Rows, unsigned Cols, typename Type>
-inline CUDA_CALLABLE void indexset(mat_t<Rows,Cols,Type>& m, int row, vec_t<Cols, Type> value)
+inline CUDA_CALLABLE vec_t<Cols, Type>* index(mat_t<Rows,Cols,Type>& m, int row)
 {
 #ifndef NDEBUG
     if (row < 0 || row >= Rows)
@@ -329,12 +351,11 @@ inline CUDA_CALLABLE void indexset(mat_t<Rows,Cols,Type>& m, int row, vec_t<Cols
     }
 #endif
 
-    for(unsigned i=0; i < Cols; ++i)
-        m.data[row][i] = value[i];
+    return reinterpret_cast<vec_t<Cols, Type>*>(&m.data[row]);
 }
 
 template<unsigned Rows, unsigned Cols, typename Type>
-inline CUDA_CALLABLE void indexset(mat_t<Rows,Cols,Type>& m, int row, int col, Type value)
+inline CUDA_CALLABLE Type* index(mat_t<Rows,Cols,Type>& m, int row, int col)
 {
 #ifndef NDEBUG
     if (row < 0 || row >= Rows)
@@ -348,18 +369,19 @@ inline CUDA_CALLABLE void indexset(mat_t<Rows,Cols,Type>& m, int row, int col, T
         assert(0);
     }
 #endif
-    m.data[row][col] = value;
+
+    return &m.data[row][col];
 }
 
 template<unsigned Rows, unsigned Cols, typename Type>
-inline CUDA_CALLABLE void adj_indexset(const mat_t<Rows,Cols,Type>& m, int row, const vec_t<Cols, Type>& value,
+inline CUDA_CALLABLE void adj_index(const mat_t<Rows,Cols,Type>& m, int row,
                                        const mat_t<Rows,Cols,Type>& adj_m, int adj_row, const vec_t<Cols, Type>& adj_value)
 {
     // nop
 }
 
 template<unsigned Rows, unsigned Cols, typename Type>
-inline CUDA_CALLABLE void adj_indexset(const mat_t<Rows,Cols,Type>& m, int row, int col, Type value,
+inline CUDA_CALLABLE void adj_index(const mat_t<Rows,Cols,Type>& m, int row, int col,
                                        const mat_t<Rows,Cols,Type>& adj_m, int adj_row, int adj_col, Type adj_value)
 {
     // nop
@@ -417,7 +439,22 @@ inline CUDA_CALLABLE mat_t<Rows,Cols,Type> div(const mat_t<Rows,Cols,Type>& a, T
         }
     }
 
-    return t;   
+    return t;
+}
+
+template<unsigned Rows, unsigned Cols, typename Type>
+inline CUDA_CALLABLE mat_t<Rows,Cols,Type> div(Type b, const mat_t<Rows,Cols,Type>& a)
+{
+    mat_t<Rows,Cols,Type> t;
+    for (unsigned i=0; i < Rows; ++i)
+    {
+        for (unsigned j=0; j < Cols; ++j)
+        {
+            t.data[i][j] = b / a.data[i][j];
+        }
+    }
+
+    return t;
 }
 
 template<unsigned Rows, unsigned Cols, typename Type>
@@ -432,7 +469,7 @@ inline CUDA_CALLABLE mat_t<Rows,Cols,Type> mul(const mat_t<Rows,Cols,Type>& a, T
         }
     }
 
-    return t;   
+    return t;
 }
 
 template<unsigned Rows, unsigned Cols, typename Type>
@@ -465,6 +502,17 @@ inline CUDA_CALLABLE vec_t<Rows,Type> mul(const mat_t<Rows,Cols,Type>& a, const
     return r;
 }
 
+template<unsigned Rows, unsigned Cols, typename Type>
+inline CUDA_CALLABLE vec_t<Cols,Type> mul(const vec_t<Rows,Type>& b, const mat_t<Rows,Cols,Type>& a)
+{
+    vec_t<Cols,Type> r = a.get_row(0)*b[0];
+    for( unsigned i=1; i < Rows; ++i )
+    {
+        r += a.get_row(i)*b[i];
+    }
+    return r;
+}
+
 template<unsigned Rows, unsigned Cols, unsigned ColsOut, typename Type>
 inline CUDA_CALLABLE mat_t<Rows,ColsOut,Type> mul(const mat_t<Rows,Cols,Type>& a, const mat_t<Cols,ColsOut,Type>& b)
 {
@@ -608,6 +656,17 @@ inline CUDA_CALLABLE Type trace(const mat_t<Rows,Rows,Type>& m)
     return ret;
 }
 
+template<unsigned Rows, typename Type>
+inline CUDA_CALLABLE vec_t<Rows, Type> get_diag(const mat_t<Rows,Rows,Type>& m)
+{
+    vec_t<Rows, Type> ret;
+    for( unsigned i=0; i < Rows; ++i )
+    {
+        ret[i] = m.data[i][i];
+    }
+    return ret;
+}
+
 // Only implementing inverses for 2x2, 3x3 and 4x4 matrices for now...
 template<typename Type>
 inline CUDA_CALLABLE mat_t<2,2,Type> inverse(const mat_t<2,2,Type>& m)
@@ -842,14 +901,14 @@ inline CUDA_CALLABLE vec_t<3,Type> transform_vector(const mat_t<4,4,Type>& m, co
 }
 
 template<unsigned Rows, unsigned Cols, typename Type>
-inline CUDA_CALLABLE void adj_index(const mat_t<Rows,Cols,Type>& m, int row, mat_t<Rows,Cols,Type>& adj_m, int& adj_row, const vec_t<Cols,Type>& adj_ret)
+inline CUDA_CALLABLE void adj_extract(const mat_t<Rows,Cols,Type>& m, int row, mat_t<Rows,Cols,Type>& adj_m, int& adj_row, const vec_t<Cols,Type>& adj_ret)
 {
     for( unsigned col=0; col < Cols; ++col )
         adj_m.data[row][col] += adj_ret[col];
 }
 
 template<unsigned Rows, unsigned Cols, typename Type>
-inline void CUDA_CALLABLE adj_index(const mat_t<Rows,Cols,Type>& m, int row, int col, mat_t<Rows,Cols,Type>& adj_m, int& adj_row, int& adj_col, Type adj_ret)
+inline void CUDA_CALLABLE adj_extract(const mat_t<Rows,Cols,Type>& m, int row, int col, mat_t<Rows,Cols,Type>& adj_m, int& adj_row, int& adj_col, Type adj_ret)
 {
 #ifndef NDEBUG
     if (row < 0 || row > Rows)
@@ -913,6 +972,20 @@ inline CUDA_CALLABLE void adj_div(const mat_t<Rows,Cols,Type>& a, Type s, mat_t<
     }
 }
 
+template<unsigned Rows, unsigned Cols, typename Type>
+inline CUDA_CALLABLE void adj_div(Type s, const mat_t<Rows,Cols,Type>& a, Type& adj_s, mat_t<Rows,Cols,Type>& adj_a, const mat_t<Rows,Cols,Type>& adj_ret)
+{
+    adj_s -= tensordot(a , adj_ret)/ (s * s); // - a / s^2
+
+    for (unsigned i=0; i < Rows; ++i)
+    {
+        for (unsigned j=0; j < Cols; ++j)
+        {
+            adj_a.data[i][j] += s / adj_ret.data[i][j];
+        }
+    }
+}
+
 template<unsigned Rows, unsigned Cols, typename Type>
 inline CUDA_CALLABLE void adj_mul(const mat_t<Rows,Cols,Type>& a, Type b, mat_t<Rows,Cols,Type>& adj_a, Type& adj_b, const mat_t<Rows,Cols,Type>& adj_ret)
 {
@@ -946,6 +1019,13 @@ inline CUDA_CALLABLE void adj_mul(const mat_t<Rows,Cols,Type>& a, const vec_t<Co
     adj_b += mul(transpose(a), adj_ret);
 }
 
+template<unsigned Rows, unsigned Cols, typename Type>
+inline CUDA_CALLABLE void adj_mul(const vec_t<Rows,Type>& b, const mat_t<Rows,Cols,Type>& a, vec_t<Rows,Type>& adj_b, mat_t<Rows,Cols,Type>& adj_a, const vec_t<Cols,Type>& adj_ret)
+{
+    adj_a += outer(b, adj_ret);
+    adj_b += mul(adj_ret, transpose(a));
+}
+
 template<unsigned Rows, unsigned Cols, unsigned ColsOut, typename Type>
 inline CUDA_CALLABLE void adj_mul(const mat_t<Rows,Cols,Type>& a, const mat_t<Cols,ColsOut,Type>& b, mat_t<Rows,Cols,Type>& adj_a, mat_t<Cols,ColsOut,Type>& adj_b, const mat_t<Rows,ColsOut,Type>& adj_ret)
 {
@@ -973,6 +1053,13 @@ inline CUDA_CALLABLE void adj_diag(const vec_t<Rows,Type>& d, vec_t<Rows,Type>&
         adj_d[i] += adj_ret.data[i][i];
 }
 
+template<unsigned Rows, typename Type>
+inline CUDA_CALLABLE void adj_get_diag(const mat_t<Rows,Rows,Type>& m, mat_t<Rows,Rows,Type>& adj_m, const vec_t<Rows,Type>& adj_ret)
+{
+    for (unsigned i=0; i < Rows; ++i)
+        adj_m.data[i][i] += adj_ret[i];
+}
+
 template<typename Type>
 inline CUDA_CALLABLE void adj_determinant(const mat_t<2,2,Type>& m, mat_t<2,2,Type>& adj_m, Type adj_ret)
 {
@@ -1079,10 +1166,10 @@ inline CUDA_CALLABLE void adj_determinant(const mat_t<4,4,Type>& m, mat_t<4,4,Ty
 }
 
 template<unsigned Rows, typename Type>
-inline CUDA_CALLABLE void adj_inverse(const mat_t<Rows,Rows,Type>& m, mat_t<Rows,Rows,Type>& adj_m, const mat_t<Rows,Rows,Type>& adj_ret)
+inline CUDA_CALLABLE void adj_inverse(const mat_t<Rows,Rows,Type>& m, mat_t<Rows,Rows,Type>& ret, mat_t<Rows,Rows,Type>& adj_m, const mat_t<Rows,Rows,Type>& adj_ret)
 {
     // todo: how to cache this from the forward pass?
-    mat_t<Rows,Rows,Type> invt = transpose(inverse(m));
+    mat_t<Rows,Rows,Type> invt = transpose(ret);
 
     // see https://people.maths.ox.ac.uk/gilesm/files/NA-08-01.pdf 2.2.3
     adj_m -= mul(mul(invt, adj_ret), invt);
@@ -1124,10 +1211,10 @@ inline CUDA_CALLABLE void adj_cw_mul(const mat_t<Rows,Cols,Type>& a, const mat_t
 }
 
 template<unsigned Rows, unsigned Cols, typename Type>
-inline CUDA_CALLABLE void adj_cw_div(const mat_t<Rows,Cols,Type>& a, const mat_t<Rows,Cols,Type>& b, mat_t<Rows,Cols,Type>& adj_a, mat_t<Rows,Cols,Type>& adj_b, const mat_t<Rows,Cols,Type>& adj_ret)
+inline CUDA_CALLABLE void adj_cw_div(const mat_t<Rows,Cols,Type>& a, const mat_t<Rows,Cols,Type>& b, mat_t<Rows,Cols,Type>& ret, mat_t<Rows,Cols,Type>& adj_a, mat_t<Rows,Cols,Type>& adj_b, const mat_t<Rows,Cols,Type>& adj_ret)
 {
   adj_a += cw_div(adj_ret, b);
-  adj_b -= cw_mul(adj_ret, cw_div(cw_div(a, b), b));
+  adj_b -= cw_mul(adj_ret, cw_div(ret, b));
 }
 
 // adjoint for the constant constructor:
@@ -1143,6 +1230,19 @@ inline CUDA_CALLABLE void adj_mat_t(Type s, Type& adj_s, const mat_t<Rows, Cols,
     }
 }
 
+// adjoint for the casting constructor:
+template<unsigned Rows, unsigned Cols, typename Type, typename OtherType>
+inline CUDA_CALLABLE void adj_mat_t(const mat_t<Rows, Cols, OtherType>& other, mat_t<Rows, Cols, OtherType>& adj_other, const mat_t<Rows, Cols, Type>& adj_ret)
+{
+    for (unsigned i=0; i < Rows; ++i)
+    {
+        for (unsigned j=0; j < Cols; ++j)
+        {
+            adj_other.data[i][j] += adj_ret.data[i][j];
+        }
+    }
+}
+
 // adjoint for the initializer_array scalar constructor:
 template<unsigned Rows, unsigned Cols, typename Type>
 inline CUDA_CALLABLE void adj_mat_t(const initializer_array<Rows * Cols, Type> &cmps, const initializer_array<Rows * Cols, Type*> &adj_cmps, const mat_t<Rows, Cols, Type>& adj_ret)

warp/native/matnn.h

@@ -248,7 +248,7 @@ CUDA_CALLABLE inline void adj_mlp(const array_t<float>& weights, const array_t<f
             tmp += weights.data[i*n + j]*x.data[index + b*j];
         }
 
-        // adjoint w.r.t to acivation
+        // adjoint w.r.t to activation
         float adj_f = 0.0f;
     
         if (adj_out.data)
@@ -313,7 +313,7 @@ CUDA_CALLABLE inline void adj_mlp(const array_t<float>& weights, const array_t<f
 //             tmp += weights[i*n + j]*x[index + b*j];            
 //         }
 
-//         // adjoint w.r.t to acivation
+//         // adjoint w.r.t to activation
 //         float adj_f = 0.0f;
 //         adj_activation(tmp, adj_f, adj_out[index + b*i]);