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

warp/native/vec.h

@@ -16,9 +16,11 @@ namespace wp
 template<unsigned Length, typename Type>
 struct vec_t
 {
-    Type c[Length] = {};
+    Type c[Length];
 
-    inline vec_t() = default;
+    inline CUDA_CALLABLE vec_t()
+        : c()
+    {}
     
     inline CUDA_CALLABLE vec_t(Type s)
     {
@@ -27,6 +29,15 @@ struct vec_t
             c[i] = s;
         }
     }
+    
+    template <typename OtherType>
+    inline explicit CUDA_CALLABLE vec_t(const vec_t<Length, OtherType>& other)
+    {
+        for( unsigned i=0; i < Length; ++i )
+        {
+            c[i] = static_cast<Type>(other[i]);
+        }
+    }
 
     inline CUDA_CALLABLE vec_t(Type x, Type y)
     {
@@ -275,12 +286,41 @@ inline CUDA_CALLABLE vec_t<2, Type> div(vec_t<2, Type> a, Type s)
     return vec_t<2, Type>(a.c[0]/s,a.c[1]/s);
 }
 
+template<unsigned Length, typename Type>
+inline CUDA_CALLABLE vec_t<Length, Type> div(Type s, vec_t<Length, Type> a)
+{
+    vec_t<Length, Type> ret;
+    for (unsigned i=0; i < Length; ++i)
+    {
+        ret[i] = s / a[i];
+    }
+    return ret;
+}
+
+template<typename Type>
+inline CUDA_CALLABLE vec_t<3, Type> div(Type s, vec_t<3, Type> a)
+{
+    return vec_t<3, Type>(s/a.c[0],s/a.c[1],s/a.c[2]);
+}
+
+template<typename Type>
+inline CUDA_CALLABLE vec_t<2, Type> div(Type s, vec_t<2, Type> a)
+{
+    return vec_t<2, Type>(s/a.c[0],s/a.c[1]);
+}
+
 template<unsigned Length, typename Type>
 inline CUDA_CALLABLE vec_t<Length, Type> operator / (vec_t<Length, Type> a, Type s)
 {
     return div(a,s);
 }
 
+template<unsigned Length, typename Type>
+inline CUDA_CALLABLE vec_t<Length, Type> operator / (Type s, vec_t<Length, Type> a)
+{
+    return div(s, a);
+}
+
 // component wise division
 template<unsigned Length, typename Type>
 inline CUDA_CALLABLE vec_t<Length, Type> cw_div(vec_t<Length, Type> a, vec_t<Length, Type> b)
@@ -374,7 +414,7 @@ inline CUDA_CALLABLE Type tensordot(vec_t<Length, Type> a, vec_t<Length, Type> b
 
 
 template<unsigned Length, typename Type>
-inline CUDA_CALLABLE Type index(const vec_t<Length, Type> & a, int idx)
+inline CUDA_CALLABLE Type extract(const vec_t<Length, Type> & a, int idx)
 {
 #ifndef NDEBUG
     if (idx < 0 || idx >= Length)
@@ -388,7 +428,21 @@ inline CUDA_CALLABLE Type index(const vec_t<Length, Type> & a, int idx)
 }
 
 template<unsigned Length, typename Type>
-inline CUDA_CALLABLE void indexset(vec_t<Length, Type>& v, int idx, Type value)
+inline CUDA_CALLABLE Type* index(vec_t<Length, Type>& v, int idx)
+{
+#ifndef NDEBUG
+    if (idx < 0 || idx >= Length)
+    {
+        printf("vec index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
+        assert(0);
+    }
+#endif
+
+    return &v[idx];
+}
+
+template<unsigned Length, typename Type>
+inline CUDA_CALLABLE Type* indexref(vec_t<Length, Type>* v, int idx)
 {
 #ifndef NDEBUG
     if (idx < 0 || idx >= Length)
@@ -398,17 +452,23 @@ inline CUDA_CALLABLE void indexset(vec_t<Length, Type>& v, int idx, Type value)
     }
 #endif
 
-    v[idx] = value;
+    return &((*v)[idx]);
 }
 
 template<unsigned Length, typename Type>
-inline CUDA_CALLABLE void adj_indexset(vec_t<Length, Type>& v, int idx, const Type& value,
+inline CUDA_CALLABLE void adj_index(vec_t<Length, Type>& v, int idx,
                                        vec_t<Length, Type>& adj_v, int adj_idx, const Type& adj_value)
 {
     // nop
 }
 
 
+template<unsigned Length, typename Type>
+inline CUDA_CALLABLE void adj_indexref(vec_t<Length, Type>* v, int idx, 
+                                       vec_t<Length, Type>& adj_v, int adj_idx, const Type& adj_value)
+{
+    // nop
+}
 
 
 template<unsigned Length, typename Type>
@@ -572,7 +632,7 @@ inline CUDA_CALLABLE void expect_near(const vec_t<Length, Type>& actual, const v
     }
     if (diff > tolerance)
     {
-        printf("Error, expect_near() failed with torerance "); print(tolerance);
+        printf("Error, expect_near() failed with tolerance "); print(tolerance);
         printf("\t Expected: "); print(expected); 
         printf("\t Actual: "); print(actual);
     }
@@ -630,6 +690,15 @@ inline CUDA_CALLABLE void adj_vec_t(Type s, Type& adj_s, const vec_t<Length, Typ
     }
 }
 
+// adjoint for the casting constructor
+template<unsigned Length, typename Type, typename OtherType>
+inline CUDA_CALLABLE void adj_vec_t(const vec_t<Length, OtherType>& other, vec_t<Length, OtherType>& adj_other, const vec_t<Length, Type>& adj_ret)
+{
+    for( unsigned i=0; i < Length; ++i )
+    {
+        adj_other[i] += static_cast<OtherType>(adj_ret[i]);
+    }
+}
 
 template<typename Type>
 CUDA_CALLABLE inline void adj_vec_t(const vec_t<3,Type>& w, const vec_t<3,Type>& v, vec_t<3,Type>& adj_w, vec_t<3,Type>& adj_v, const vec_t<6,Type>& adj_ret)
@@ -697,9 +766,30 @@ inline CUDA_CALLABLE void adj_div(vec_t<Length, Type> a, Type s, vec_t<Length, T
 }
 
 template<unsigned Length, typename Type>
-inline CUDA_CALLABLE void adj_cw_div(vec_t<Length, Type> a, vec_t<Length, Type> b, vec_t<Length, Type>& adj_a, vec_t<Length, Type>& adj_b, const vec_t<Length, Type>& adj_ret) {
+inline CUDA_CALLABLE void adj_div(Type s, vec_t<Length, Type> a, Type& adj_s, vec_t<Length, Type>& adj_a, const vec_t<Length, Type>& adj_ret)
+{
+
+    adj_s -= dot(a , adj_ret)/ (s * s); // - a / s^2
+
+    for( unsigned i=0; i < Length; ++i )
+    {
+        adj_a[i] += s / adj_ret[i];
+    }
+
+#if FP_CHECK
+    if (!isfinite(a) || !isfinite(s) || !isfinite(adj_a) || !isfinite(adj_s) || !isfinite(adj_ret))
+    {
+        // \TODO: How shall we implement this error message?
+        // printf("adj_div((%f %f %f %f), %f, (%f %f %f %f), %f, (%f %f %f %f)\n", a.x, a.y, a.z, a.w, s, adj_a.x, adj_a.y, adj_a.z, adj_a.w, adj_s, adj_ret.x, adj_ret.y, adj_ret.z, adj_ret.w);
+        assert(0);
+    }
+#endif
+}
+
+template<unsigned Length, typename Type>
+inline CUDA_CALLABLE void adj_cw_div(vec_t<Length, Type> a, vec_t<Length, Type> b, vec_t<Length, Type>& ret, vec_t<Length, Type>& adj_a, vec_t<Length, Type>& adj_b, const vec_t<Length, 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));
 }
 
 template<unsigned Length, typename Type>
@@ -798,7 +888,7 @@ inline CUDA_CALLABLE void adj_dot(vec_t<3, Type> a, vec_t<3, Type> b, vec_t<3, T
 
 
 template<unsigned Length, typename Type>
-inline CUDA_CALLABLE void adj_index(const vec_t<Length, Type> & a, int idx, vec_t<Length, Type> & adj_a, int & adj_idx, Type & adj_ret)
+inline CUDA_CALLABLE void adj_extract(const vec_t<Length, Type> & a, int idx, vec_t<Length, Type> & adj_a, int & adj_idx, Type & adj_ret)
 {
 #ifndef NDEBUG
     if (idx < 0 || idx > Length)
@@ -812,9 +902,12 @@ inline CUDA_CALLABLE void adj_index(const vec_t<Length, Type> & a, int idx, vec_
 }
 
 template<unsigned Length, typename Type>
-inline CUDA_CALLABLE void adj_length(vec_t<Length, Type> a, vec_t<Length, Type>& adj_a, const Type adj_ret)
+inline CUDA_CALLABLE void adj_length(vec_t<Length, Type> a, Type ret, vec_t<Length, Type>& adj_a, const Type adj_ret)
 {
-    adj_a += normalize(a)*adj_ret;
+    if (ret > Type(kEps))
+    {
+        adj_a += div(a, ret) * adj_ret;
+    }
 
 #if FP_CHECK
     if (!isfinite(adj_a))
@@ -842,7 +935,7 @@ inline CUDA_CALLABLE void adj_length_sq(vec_t<Length, Type> a, vec_t<Length, Typ
 }
 
 template<unsigned Length, typename Type>
-inline CUDA_CALLABLE void adj_normalize(vec_t<Length, Type> a, vec_t<Length, Type>& adj_a, const vec_t<Length, Type>& adj_ret)
+inline CUDA_CALLABLE void adj_normalize(vec_t<Length, Type> a, vec_t<Length, Type>& ret, vec_t<Length, Type>& adj_a, const vec_t<Length, Type>& adj_ret)
 {
     Type d = length(a);
     
@@ -850,9 +943,7 @@ inline CUDA_CALLABLE void adj_normalize(vec_t<Length, Type> a, vec_t<Length, Typ
     {
         Type invd = Type(1.0f)/d;
 
-        vec_t<Length, Type> ahat = normalize(a);
-
-        adj_a += (adj_ret*invd - ahat*(dot(ahat, adj_ret))*invd);
+        adj_a += (adj_ret*invd - ret*(dot(ret, adj_ret))*invd);
 
 #if FP_CHECK
         if (!isfinite(adj_a))
@@ -913,8 +1004,8 @@ inline CUDA_CALLABLE void adj_max(const vec_t<Length,Type> &v, vec_t<Length,Type
 
 // Do I need to specialize these for different lengths?
 template<unsigned Length, typename Type>
-inline CUDA_CALLABLE vec_t<Length, Type> atomic_add(vec_t<Length, Type> * addr, vec_t<Length, Type> value) {
-
+inline CUDA_CALLABLE vec_t<Length, Type> atomic_add(vec_t<Length, Type> * addr, vec_t<Length, Type> value)
+{
     vec_t<Length, Type> ret;
     for( unsigned i=0; i < Length; ++i )
     {
@@ -925,8 +1016,8 @@ inline CUDA_CALLABLE vec_t<Length, Type> atomic_add(vec_t<Length, Type> * addr,
 }
 
 template<unsigned Length, typename Type>
-inline CUDA_CALLABLE vec_t<Length, Type> atomic_min(vec_t<Length, Type> * addr, vec_t<Length, Type> value) {
-
+inline CUDA_CALLABLE vec_t<Length, Type> atomic_min(vec_t<Length, Type> * addr, vec_t<Length, Type> value)
+{
     vec_t<Length, Type> ret;
     for( unsigned i=0; i < Length; ++i )
     {
@@ -937,8 +1028,8 @@ inline CUDA_CALLABLE vec_t<Length, Type> atomic_min(vec_t<Length, Type> * addr,
 }
 
 template<unsigned Length, typename Type>
-inline CUDA_CALLABLE vec_t<Length, Type> atomic_max(vec_t<Length, Type> * addr, vec_t<Length, Type> value) {
-
+inline CUDA_CALLABLE vec_t<Length, Type> atomic_max(vec_t<Length, Type> * addr, vec_t<Length, Type> value)
+{
     vec_t<Length, Type> ret;
     for( unsigned i=0; i < Length; ++i )
     {
@@ -948,6 +1039,17 @@ inline CUDA_CALLABLE vec_t<Length, Type> atomic_max(vec_t<Length, Type> * addr,
     return ret;
 }
 
+template<unsigned Length, typename Type>
+inline CUDA_CALLABLE void adj_atomic_minmax(
+    vec_t<Length,Type> *addr,
+    vec_t<Length,Type> *adj_addr,
+    const vec_t<Length,Type> &value,
+    vec_t<Length,Type> &adj_value)
+{
+    for (unsigned i=0; i < Length; ++i)
+        adj_atomic_minmax(&(addr->c[i]), &(adj_addr->c[i]), value[i], adj_value[i]);
+}
+
 // ok, the original implementation of this didn't take the absolute values.
 // I wouldn't consider this expected behavior. It looks like it's only
 // being used for bounding boxes at the moment, where this doesn't matter,
@@ -956,11 +1058,11 @@ inline CUDA_CALLABLE vec_t<Length, Type> atomic_max(vec_t<Length, Type> * addr,
 template<unsigned Length, typename Type>
 CUDA_CALLABLE inline int longest_axis(const vec_t<Length, Type>& v)
 {
-    Type lmax = fabs(v[0]);
+    Type lmax = abs(v[0]);
     int ret(0);
     for( unsigned i=1; i < Length; ++i )
     {
-        Type l = fabs(v[i]);
+        Type l = abs(v[i]);
         if( l > lmax )
         {
             ret = i;

warp/native/volume.h

@@ -232,6 +232,126 @@ CUDA_CALLABLE inline void adj_volume_sample_i(uint64_t id, vec3 uvw, uint64_t& a
     // NOP
 }
 
+// Sampling the volume at the given index-space coordinates, uvw can be fractional
+CUDA_CALLABLE inline float volume_sample_grad_f(uint64_t id, vec3 uvw, int sampling_mode, vec3& grad)
+{
+    const pnanovdb_buf_t buf = volume::id_to_buffer(id);
+    const pnanovdb_root_handle_t root = volume::get_root(buf);
+    const pnanovdb_vec3_t uvw_pnano{ uvw[0], uvw[1], uvw[2] };
+
+    if (sampling_mode == volume::CLOSEST)
+    {
+        const pnanovdb_coord_t ijk = pnanovdb_vec3_round_to_coord(uvw_pnano);
+        float val;
+        pnano_read(val, buf, root, PNANOVDB_REF(ijk));
+        grad = vec3(0.0f, 0.0f, 0.0f);
+        return val;
+    }
+    else if (sampling_mode == volume::LINEAR)
+    {
+        // NB. linear sampling is not used on int volumes
+        constexpr pnanovdb_coord_t OFFSETS[] = {
+            { 0, 0, 0 }, { 0, 0, 1 }, { 0, 1, 0 }, { 0, 1, 1 }, { 1, 0, 0 }, { 1, 0, 1 }, { 1, 1, 0 }, { 1, 1, 1 },
+        };
+
+        const pnanovdb_vec3_t ijk_base{ floorf(uvw_pnano.x), floorf(uvw_pnano.y), floorf(uvw_pnano.z) };
+        const pnanovdb_vec3_t ijk_frac{ uvw_pnano.x - ijk_base.x, uvw_pnano.y - ijk_base.y, uvw_pnano.z - ijk_base.z };
+        const pnanovdb_coord_t ijk{ (pnanovdb_int32_t)ijk_base.x, (pnanovdb_int32_t)ijk_base.y, (pnanovdb_int32_t)ijk_base.z };
+
+        pnanovdb_readaccessor_t accessor;
+        pnanovdb_readaccessor_init(PNANOVDB_REF(accessor), root);
+        float val = 0.0f;
+        const float wx[2]{ 1 - ijk_frac.x, ijk_frac.x };
+        const float wy[2]{ 1 - ijk_frac.y, ijk_frac.y };
+        const float wz[2]{ 1 - ijk_frac.z, ijk_frac.z };
+
+        const float sign_dx[8] = {-1.0f, -1.0f, -1.0f, -1.0f, 1.0f, 1.0f, 1.0f, 1.0f};
+        const float sign_dy[8] = {-1.0f, -1.0f, 1.0f, 1.0f, -1.0f, -1.0f, 1.0f, 1.0f};
+        const float sign_dz[8] = {-1.0f, 1.0f, -1.0f, 1.0f, -1.0f, 1.0f, -1.0f, 1.0f};
+        
+        float dfdx = 0.0f;
+        float dfdy = 0.0f;
+        float dfdz = 0.0f;
+        for (int idx = 0; idx < 8; ++idx)
+        {
+            const pnanovdb_coord_t& offs = OFFSETS[idx];
+            const pnanovdb_coord_t ijk_shifted = pnanovdb_coord_add(ijk, offs);
+            float v;
+            pnano_read(v, buf, PNANOVDB_REF(accessor), PNANOVDB_REF(ijk_shifted));
+            val = add(val, wx[offs.x] * wy[offs.y] * wz[offs.z] * v);
+            dfdx = add(dfdx, wy[offs.y] * wz[offs.z] * sign_dx[idx] * v);
+            dfdy = add(dfdy, wx[offs.x] * wz[offs.z] * sign_dy[idx] * v);
+            dfdz = add(dfdz, wx[offs.x] * wy[offs.y] * sign_dz[idx] * v);
+        }
+        grad = vec3(dfdx, dfdy, dfdz);
+        return val;
+    }
+    return 0.0f;
+}
+
+CUDA_CALLABLE inline void adj_volume_sample_grad_f(
+    uint64_t id, vec3 uvw, int sampling_mode, vec3& grad, uint64_t& adj_id, vec3& adj_uvw, int& adj_sampling_mode, vec3& adj_grad, const float& adj_ret)
+{
+    if (volume::get_grid_type(volume::id_to_buffer(id)) != PNANOVDB_GRID_TYPE_FLOAT) return;
+
+    if (sampling_mode != volume::LINEAR) {
+        return; // NOP
+    }
+
+    const pnanovdb_buf_t buf = volume::id_to_buffer(id);
+    const pnanovdb_root_handle_t root = volume::get_root(buf);
+    const pnanovdb_vec3_t uvw_pnano{ uvw[0], uvw[1], uvw[2] };
+
+    constexpr pnanovdb_coord_t OFFSETS[] = {
+        { 0, 0, 0 }, { 0, 0, 1 }, { 0, 1, 0 }, { 0, 1, 1 }, { 1, 0, 0 }, { 1, 0, 1 }, { 1, 1, 0 }, { 1, 1, 1 },
+    };
+
+    const pnanovdb_vec3_t ijk_base{ floorf(uvw_pnano.x), floorf(uvw_pnano.y), floorf(uvw_pnano.z) };
+    const pnanovdb_vec3_t ijk_frac{ uvw_pnano.x - ijk_base.x, uvw_pnano.y - ijk_base.y, uvw_pnano.z - ijk_base.z };
+    const pnanovdb_coord_t ijk{ (pnanovdb_int32_t)ijk_base.x, (pnanovdb_int32_t)ijk_base.y, (pnanovdb_int32_t)ijk_base.z };
+
+    pnanovdb_readaccessor_t accessor;
+    pnanovdb_readaccessor_init(PNANOVDB_REF(accessor), root);
+    const float wx[2]{ 1 - ijk_frac.x, ijk_frac.x };
+    const float wy[2]{ 1 - ijk_frac.y, ijk_frac.y };
+    const float wz[2]{ 1 - ijk_frac.z, ijk_frac.z };
+    const float sign_dx[8] = {-1.0f, -1.0f, -1.0f, -1.0f, 1.0f, 1.0f, 1.0f, 1.0f};
+    const float sign_dy[8] = {-1.0f, -1.0f, 1.0f, 1.0f, -1.0f, -1.0f, 1.0f, 1.0f};
+    const float sign_dz[8] = {-1.0f, 1.0f, -1.0f, 1.0f, -1.0f, 1.0f, -1.0f, 1.0f};
+    
+    float dfdxdy = 0.0f;
+    float dfdxdz = 0.0f;
+    float dfdydx = 0.0f;
+    float dfdydz = 0.0f;
+    float dfdzdx = 0.0f;    
+    float dfdzdy = 0.0f;    
+    vec3 dphi(0,0,0);
+    for (int idx = 0; idx < 8; ++idx)
+    {
+        const pnanovdb_coord_t& offs = OFFSETS[idx];
+        const pnanovdb_coord_t ijk_shifted = pnanovdb_coord_add(ijk, offs);
+        float v;
+        pnano_read(v, buf, PNANOVDB_REF(accessor), PNANOVDB_REF(ijk_shifted));
+        const vec3 signs(offs.x * 2 - 1, offs.y * 2 - 1, offs.z * 2 - 1);
+        const vec3 grad_w(signs[0] * wy[offs.y] * wz[offs.z], signs[1] * wx[offs.x] * wz[offs.z], signs[2] * wx[offs.x] * wy[offs.y]); 
+        dphi = add(dphi, mul(v, grad_w));
+
+        dfdxdy = add(dfdxdy, signs[1] * wz[offs.z] * sign_dx[idx] * v);
+        dfdxdz = add(dfdxdz, wy[offs.y] * signs[2] * sign_dx[idx] * v);
+
+        dfdydx = add(dfdydx, signs[0] * wz[offs.z] * sign_dy[idx] * v);
+        dfdydz = add(dfdydz, wx[offs.x] * signs[2] * sign_dy[idx] * v);
+
+        dfdzdx = add(dfdzdx, signs[0] * wy[offs.y] * sign_dz[idx] * v);        
+        dfdzdy = add(dfdzdy, wx[offs.x] * signs[1] * sign_dz[idx] * v);        
+    }
+
+    adj_uvw += mul(dphi, adj_ret);
+    adj_uvw[0] += adj_grad[1] * dfdydx + adj_grad[2] * dfdzdx;
+    adj_uvw[1] += adj_grad[0] * dfdxdy + adj_grad[2] * dfdzdy;
+    adj_uvw[2] += adj_grad[0] * dfdxdz + adj_grad[1] * dfdydz;
+}
+
 CUDA_CALLABLE inline float volume_lookup_f(uint64_t id, int32_t i, int32_t j, int32_t k)
 {
     if (volume::get_grid_type(volume::id_to_buffer(id)) != PNANOVDB_GRID_TYPE_FLOAT) return 0.f;