warp-lang 1.9.0__py3-none-win_amd64.whl → 1.10.0rc2__py3-none-win_amd64.whl

warp/native/builtin.h

@@ -35,6 +35,11 @@
     #define CUDA_CALLABLE_DEVICE __device__
 #endif
 
+// Tile block dimension used while building the warp core library
+#ifndef WP_TILE_BLOCK_DIM
+#define WP_TILE_BLOCK_DIM 256
+#endif
+
 #ifdef WP_VERIFY_FP
 #define FP_CHECK 1
 #define DO_IF_FPCHECK(X) {X}
@@ -48,8 +53,28 @@
 #define RAD_TO_DEG 57.29577951308232087679
 #define DEG_TO_RAD  0.01745329251994329577
 
+#ifndef M_PI_F
+#define M_PI_F 3.14159265358979323846f
+#endif
+
+#ifndef M_2_SQRT_PI_F
+#define M_2_SQRT_PI_F 1.1283791670955125739f  // 2/sqrt(pi)
+#endif
+
+#ifndef M_2_SQRT_PI
+#define M_2_SQRT_PI 1.1283791670955125738961589031215  // 2/sqrt(pi)
+#endif
+
+#ifndef M_SQRT_PI_F_2
+#define M_SQRT_PI_F_2 0.88622692545275801364f  // sqrt(pi)/2
+#endif
+
+#ifndef M_SQRT_PI_2
+#define M_SQRT_PI_2 0.88622692545275801364908374167057  // sqrt(pi)/2
+#endif
+
 #if defined(__CUDACC__) && !defined(_MSC_VER)
-__device__ void __debugbreak() { __brkpt(); }
+__device__ inline void __debugbreak() { __brkpt(); }
 #endif
 
 #if defined(__clang__) && defined(__CUDA__) && defined(__CUDA_ARCH__)
@@ -267,6 +292,13 @@ inline CUDA_CALLABLE half operator / (half a,half b)
 
 
 
+template<typename TRet, typename T>
+inline CUDA_CALLABLE TRet cast(T a)
+{
+    static_assert(sizeof(TRet) == sizeof(T), "source and destination must have the same size");
+    return *reinterpret_cast<TRet*>(&a);
+}
+
 template <typename T>
 CUDA_CALLABLE inline float cast_float(T x) { return (float)(x); }
 
@@ -782,6 +814,7 @@ inline CUDA_CALLABLE half floordiv(half a, half b)
 #endif
     return floorf(float(a/b));
 }
+
 inline CUDA_CALLABLE float floordiv(float a, float b)
 {
 #if FP_CHECK
@@ -793,6 +826,7 @@ inline CUDA_CALLABLE float floordiv(float a, float b)
 #endif
     return floorf(a/b);
 }
+
 inline CUDA_CALLABLE double floordiv(double a, double b)
 {
 #if FP_CHECK
@@ -805,6 +839,210 @@ inline CUDA_CALLABLE double floordiv(double a, double b)
     return ::floor(a/b);
 }
 
+inline CUDA_CALLABLE half erf(half a)
+{
+    return erff(float(a));
+}
+
+inline CUDA_CALLABLE float erf(float a)
+{
+    return erff(a);
+}
+
+inline CUDA_CALLABLE double erf(double a)
+{
+    return ::erf(a);
+}
+
+inline CUDA_CALLABLE half erfc(half a)
+{
+    return erfcf(float(a));
+}
+
+inline CUDA_CALLABLE float erfc(float a)
+{
+    return erfcf(a);
+}
+
+inline CUDA_CALLABLE double erfc(double a)
+{
+    return ::erfc(a);
+}
+
+inline CUDA_CALLABLE half erfinv(half a)
+{
+#if FP_CHECK
+    if (float(a) < -1.0f || float(a) > 1.0f)
+    {
+        printf("%s:%d erfinv(%f)\n", __FILE__, __LINE__, float(a));
+        assert(0);
+    }
+#endif
+    return ::erfinvf(float(a));
+}
+
+inline CUDA_CALLABLE float erfinv(float a)
+{
+#if FP_CHECK
+    if (a < -1.0f || a > 1.0f)
+    {
+        printf("%s:%d erfinv(%f)\n", __FILE__, __LINE__, a);
+        assert(0);
+    }
+#endif
+    return ::erfinvf(a);
+}
+
+inline CUDA_CALLABLE double erfinv(double a)
+{
+#if FP_CHECK
+    if (a < -1.0 || a > 1.0)
+    {
+        printf("%s:%d erfinv(%f)\n", __FILE__, __LINE__, a);
+        assert(0);
+    }
+#endif
+    return ::erfinv(a);
+}
+
+inline CUDA_CALLABLE half erfcinv(half a)
+{
+#if FP_CHECK
+    if (float(a) < 0.0f || float(a) > 2.0f)
+    {
+        printf("%s:%d erfcinv(%f)\n", __FILE__, __LINE__, float(a));
+        assert(0);
+    }
+#endif
+    return ::erfcinvf(float(a));
+}
+
+inline CUDA_CALLABLE float erfcinv(float a)
+{
+#if FP_CHECK
+    if (a < 0.0f || a > 2.0f)
+    {
+        printf("%s:%d erfcinv(%f)\n", __FILE__, __LINE__, a);
+        assert(0);
+    }
+#endif
+    return ::erfcinvf(a);
+}
+
+inline CUDA_CALLABLE double erfcinv(double a)
+{
+#if FP_CHECK
+    if (a < 0.0 || a > 2.0)
+    {
+        printf("%s:%d erfcinv(%f)\n", __FILE__, __LINE__, a);
+        assert(0);
+    }
+#endif
+    return ::erfcinv(a);
+}
+
+inline CUDA_CALLABLE void adj_erf(half a, half& adj_a, half adj_ret)
+{
+    adj_a += half(M_2_SQRT_PI_F * ::expf(-float(a)*float(a))) * adj_ret;
+}
+
+inline CUDA_CALLABLE void adj_erf(float a, float& adj_a, float adj_ret)
+{
+    adj_a += M_2_SQRT_PI_F * ::expf(-a*a) * adj_ret;
+}
+
+inline CUDA_CALLABLE void adj_erf(double a, double& adj_a, double adj_ret)
+{
+    adj_a += M_2_SQRT_PI * ::exp(-a*a) * adj_ret;
+}
+
+inline CUDA_CALLABLE void adj_erfc(half a, half& adj_a, half adj_ret)
+{
+    adj_a -= half(M_2_SQRT_PI_F * ::expf(-float(a)*float(a))) * adj_ret;
+}
+
+inline CUDA_CALLABLE void adj_erfc(float a, float& adj_a, float adj_ret)
+{
+    adj_a -= M_2_SQRT_PI_F * ::expf(-a*a) * adj_ret;
+}
+
+inline CUDA_CALLABLE void adj_erfc(double a, double& adj_a, double adj_ret)
+{
+    adj_a -= M_2_SQRT_PI * ::exp(-a*a) * adj_ret;
+}
+
+inline CUDA_CALLABLE void adj_erfinv(half a, half ret, half& adj_a, half adj_ret)
+{
+#if FP_CHECK
+    if (float(a) < -1.0f || float(a) > 1.0f)
+    {
+        printf("%s:%d adj_erfinv(%f)\n", __FILE__, __LINE__, float(a));
+        assert(0);
+    }
+#endif
+    adj_a += half(M_SQRT_PI_F_2 * ::expf(float(ret)*float(ret))) * adj_ret;
+}
+
+inline CUDA_CALLABLE void adj_erfinv(float a, float ret, float& adj_a, float adj_ret)
+{
+#if FP_CHECK
+    if (a < -1.0f || a > 1.0f)
+    {
+        printf("%s:%d adj_erfinv(%f)\n", __FILE__, __LINE__, a);
+        assert(0);
+    }
+#endif
+    adj_a += M_SQRT_PI_F_2 * ::expf(ret*ret) * adj_ret;
+}
+
+inline CUDA_CALLABLE void adj_erfinv(double a, double ret, double& adj_a, double adj_ret)
+{
+#if FP_CHECK
+    if (a < -1.0 || a > 1.0)
+    {
+        printf("%s:%d adj_erfinv(%f)\n", __FILE__, __LINE__, a);
+        assert(0);
+    }
+#endif
+    adj_a += M_SQRT_PI_2 * ::exp(ret*ret) * adj_ret;
+}
+
+inline CUDA_CALLABLE void adj_erfcinv(half a, half ret, half& adj_a, half adj_ret)
+{
+#if FP_CHECK
+    if (float(a) < 0.0f || float(a) > 2.0f)
+    {
+        printf("%s:%d adj_erfcinv(%f)\n", __FILE__, __LINE__, float(a));
+        assert(0);
+    }
+#endif
+    adj_a -= half(M_SQRT_PI_F_2 * ::expf(float(ret)*float(ret))) * adj_ret;
+}
+
+inline CUDA_CALLABLE void adj_erfcinv(float a, float ret, float& adj_a, float adj_ret)
+{
+#if FP_CHECK
+    if (a < 0.0f || a > 2.0f)
+    {
+        printf("%s:%d adj_erfcinv(%f)\n", __FILE__, __LINE__, a);
+        assert(0);
+    }
+#endif
+    adj_a -= M_SQRT_PI_F_2 * ::expf(ret*ret) * adj_ret;
+}
+
+inline CUDA_CALLABLE void adj_erfcinv(double a, double ret, double& adj_a, double adj_ret)
+{
+#if FP_CHECK
+    if (a < 0.0 || a > 2.0)
+    {
+        printf("%s:%d adj_erfcinv(%f)\n", __FILE__, __LINE__, a);
+        assert(0);
+    }
+#endif
+    adj_a -= M_SQRT_PI_2 * ::exp(ret*ret) * adj_ret;
+}
+
 inline CUDA_CALLABLE float leaky_min(float a, float b, float r) { return min(a, b); }
 inline CUDA_CALLABLE float leaky_max(float a, float b, float r) { return max(a, b); }
 
@@ -1086,23 +1324,6 @@ DECLARE_ADJOINTS(float16)
 DECLARE_ADJOINTS(float32)
 DECLARE_ADJOINTS(float64)
 
-template <typename C, typename T>
-CUDA_CALLABLE inline T select(const C& cond, const T& a, const T& b)
-{
-    // The double NOT operator !! casts to bool without compiler warnings.
-    return (!!cond) ? b : a;
-}
-
-template <typename C, typename T>
-CUDA_CALLABLE inline void adj_select(const C& cond, const T& a, const T& b, C& adj_cond, T& adj_a, T& adj_b, const T& adj_ret)
-{
-    // The double NOT operator !! casts to bool without compiler warnings.
-    if (!!cond)
-        adj_b += adj_ret;
-    else
-        adj_a += adj_ret;
-}
-
 template <typename C, typename T>
 CUDA_CALLABLE inline T where(const C& cond, const T& a, const T& b)
 {
@@ -1110,8 +1331,8 @@ CUDA_CALLABLE inline T where(const C& cond, const T& a, const T& b)
     return (!!cond) ? a : b;
 }
 
-template <typename C, typename T>
-CUDA_CALLABLE inline void adj_where(const C& cond, const T& a, const T& b, C& adj_cond, T& adj_a, T& adj_b, const T& adj_ret)
+template <typename C, typename TA, typename TB, typename TRet>
+CUDA_CALLABLE inline void adj_where(const C& cond, const TA& a, const TB& b, C& adj_cond, TA& adj_a, TB& adj_b, const TRet& adj_ret)
 {
     // The double NOT operator !! casts to bool without compiler warnings.
     if (!!cond)
@@ -1157,12 +1378,30 @@ CUDA_CALLABLE inline T& operator += (T& a, const T& b) { a = add(a, b); return a
 template <typename T>
 CUDA_CALLABLE inline T& operator -= (T& a, const T& b) { a = sub(a, b); return a; }
 
+template <typename T>
+CUDA_CALLABLE inline T& operator &= (T& a, const T& b) { a = bit_and(a, b); return a; }
+
+template <typename T>
+CUDA_CALLABLE inline T& operator |= (T& a, const T& b) { a = bit_or(a, b); return a; }
+
+template <typename T>
+CUDA_CALLABLE inline T& operator ^= (T& a, const T& b) { a = bit_xor(a, b); return a; }
+
 template <typename T>
 CUDA_CALLABLE inline T operator+(const T& a, const T& b) { return add(a, b); }
 
 template <typename T>
 CUDA_CALLABLE inline T operator-(const T& a, const T& b) { return sub(a, b); }
 
+template <typename T>
+CUDA_CALLABLE inline T operator&(const T& a, const T& b) { return bit_and(a, b); }
+
+template <typename T>
+CUDA_CALLABLE inline T operator|(const T& a, const T& b) { return bit_or(a, b); }
+
+template <typename T>
+CUDA_CALLABLE inline T operator^(const T& a, const T& b) { return bit_xor(a, b); }
+
 template <typename T>
 CUDA_CALLABLE inline T pos(const T& x) { return x; }
 template <typename T>
@@ -1278,8 +1517,8 @@ inline CUDA_CALLABLE_DEVICE void tid(int& i, int& j, int& k, int& l, size_t inde
 }
 
 // should match types.py
-constexpr int SLICE_BEGIN = (1U << (sizeof(int) - 1)) - 1; // std::numeric_limits<int>::max()
-constexpr int SLICE_END = -(1U << (sizeof(int) - 1)); // std::numeric_limits<int>::min()
+static const int SLICE_BEGIN = (1LL << (sizeof(int) * 8 - 1)) - 1; // std::numeric_limits<int>::max()
+static const int SLICE_END = -(1LL << (sizeof(int) * 8 - 1)); // std::numeric_limits<int>::min()
 
 struct slice_t
 {
@@ -1366,26 +1605,18 @@ inline CUDA_CALLABLE T atomic_add(T* buf, T value)
 #endif
 }
 
-// emulate atomic int64 add with atomicCAS()
 template <>
-inline CUDA_CALLABLE int64 atomic_add(int64* address, int64 val)
+inline CUDA_CALLABLE int64 atomic_add(int64* buf, int64 value)
 {
-#if defined(__CUDA_ARCH__)
-    unsigned long long int *address_as_ull = (unsigned long long int*)address;
-    unsigned long long int old = *address_as_ull, assumed;
-
-    while (val < (int64)old)
-    {
-        assumed = old;
-        old = atomicCAS(address_as_ull, assumed, (int64)val);
-    }
-
-    return (int64)old;
-
-#else
-    int64 old = *address;
-    *address = min(old, val);
+#if !defined(__CUDA_ARCH__)
+    int64 old = buf[0];
+    buf[0] += value;
     return old;
+#else  // CUDA compiled by NVRTC
+    unsigned long long int *buf_as_ull = (unsigned long long int*)buf;
+    unsigned long long int unsigned_value = static_cast<unsigned long long int>(value);
+    unsigned long long int result = atomicAdd(buf_as_ull, unsigned_value);
+    return static_cast<int64>(result);
 #endif
 }
 
@@ -1727,6 +1958,52 @@ CUDA_CALLABLE inline void adj_atomic_exch(T* address, T val, T* adj_address, T&
 }
 
 
+template<typename T>
+inline CUDA_CALLABLE T atomic_and(T* buf, T value)
+{
+#if defined(__CUDA_ARCH__)
+    return atomicAnd(buf, value);
+#else
+    T old = buf[0];
+    buf[0] &= value;
+    return old;
+#endif
+}
+
+template<typename T>
+inline CUDA_CALLABLE T atomic_or(T* buf, T value)
+{
+#if defined(__CUDA_ARCH__)
+    return atomicOr(buf, value);
+#else
+    T old = buf[0];
+    buf[0] |= value;
+    return old;
+#endif
+}
+
+template<typename T>
+inline CUDA_CALLABLE T atomic_xor(T* buf, T value)
+{
+#if defined(__CUDA_ARCH__)
+    return atomicXor(buf, value);
+#else
+    T old = buf[0];
+    buf[0] ^= value;
+    return old;
+#endif
+}
+
+
+// for bitwise operations we do not accumulate gradients
+template<typename T>
+CUDA_CALLABLE inline void adj_atomic_and(T* buf, T* adj_buf, T &value, T &adj_value) { }
+template<typename T>
+CUDA_CALLABLE inline void adj_atomic_or(T* buf, T* adj_buf, T &value, T &adj_value) { }
+template<typename T>
+CUDA_CALLABLE inline void adj_atomic_xor(T* buf, T* adj_buf, T &value, T &adj_value) { }
+
+
 } // namespace wp
 
 

warp/native/bvh.cpp

@@ -39,20 +39,19 @@ class TopDownBVHBuilder
 public:
 
     void build(BVH& bvh, const vec3* lowers, const vec3* uppers, int n, int in_constructor_type);
+    void rebuild(BVH& bvh, int in_constructor_type);
 
 private:
 
     void initialize_empty(BVH& bvh);
 
     bounds3 calc_bounds(const vec3* lowers, const vec3* uppers, const int* indices, int start, int end);
-
+    int build_recursive(BVH& bvh, const vec3* lowers, const vec3* uppers, int start, int end, int depth, int parent);
     int partition_median(const vec3* lowers, const vec3* uppers, int* indices, int start, int end, bounds3 range_bounds);
     int partition_midpoint(const vec3* lowers, const vec3* uppers, int* indices, int start, int end, bounds3 range_bounds);
     float partition_sah(BVH& bvh, const vec3* lowers, const vec3* uppers,
        int start, int end, bounds3 range_bounds, int& split_axis);
 
-    int build_recursive(BVH& bvh, const vec3* lowers, const vec3* uppers, int start, int end, int depth, int parent);
-
     int constructor_type = -1;
 };
 
@@ -69,6 +68,7 @@ void TopDownBVHBuilder::initialize_empty(BVH& bvh)
     bvh.root = nullptr;
     bvh.primitive_indices = nullptr;
     bvh.num_leaf_nodes = 0;
+    bvh.num_items = 0;
 }
 
 void TopDownBVHBuilder::build(BVH& bvh, const vec3* lowers, const vec3* uppers, int n, int in_constructor_type)
@@ -112,6 +112,7 @@ void TopDownBVHBuilder::build(BVH& bvh, const vec3* lowers, const vec3* uppers,
     bvh.node_uppers = new BVHPackedNodeHalf[bvh.max_nodes];
     bvh.node_parents = new int[bvh.max_nodes];
     bvh.node_counts = nullptr;
+    bvh.num_items = n;
 
     // root is always in first slot for top down builders
     bvh.root = new int[1];
@@ -124,6 +125,27 @@ void TopDownBVHBuilder::build(BVH& bvh, const vec3* lowers, const vec3* uppers,
     build_recursive(bvh, lowers, uppers,  0, n, 0, -1);
 }
 
+void TopDownBVHBuilder::rebuild(BVH& bvh, int in_constructor_type)
+{
+    if (in_constructor_type != BVH_CONSTRUCTOR_SAH && in_constructor_type != BVH_CONSTRUCTOR_MEDIAN)
+    {
+        fprintf(stderr, "Unrecognized Constructor type: %d! For CPU constructor it should be either SAH (%d) or Median (%d)!\n",
+            in_constructor_type, BVH_CONSTRUCTOR_SAH, BVH_CONSTRUCTOR_MEDIAN);
+        return;
+    }
+    if (bvh.num_items == 0)
+        return;
+
+    constructor_type = in_constructor_type;
+    for (int i = 0; i < bvh.num_items; ++i)
+        bvh.primitive_indices[i] = i;
+
+    bvh.max_depth = 0;
+    bvh.num_nodes = 0;
+    bvh.num_leaf_nodes = 0;
+    build_recursive(bvh, bvh.item_lowers, bvh.item_uppers, 0, bvh.num_items, 0, -1);
+}
+
 
 bounds3 TopDownBVHBuilder::calc_bounds(const vec3* lowers, const vec3* uppers, const int* indices, int start, int end)
 {
@@ -323,7 +345,7 @@ int TopDownBVHBuilder::build_recursive(BVH& bvh, const vec3* lowers, const vec3*
 
     // If the depth exceeds BVH_QUERY_STACK_SIZE, an out-of-bounds access bug may occur during querying.
     // In that case, we merge the following nodes into a single large leaf node.
-    if (n <= BVH_LEAF_SIZE || depth >= BVH_QUERY_STACK_SIZE - 1)
+    if (n <= bvh.leaf_size || depth >= BVH_QUERY_STACK_SIZE - 1)
     {
         bvh.node_lowers[node_index] = make_node(b.lower, start, true);
         bvh.node_uppers[node_index] = make_node(b.upper, end, false);
@@ -421,7 +443,11 @@ void bvh_refit_host(BVH& bvh)
 {
     bvh_refit_recursive(bvh, 0);
 }
-
+void bvh_rebuild_host(BVH& bvh, int constructor_type)
+{
+    TopDownBVHBuilder builder;
+    builder.rebuild(bvh, constructor_type);
+}
 
 } // namespace wp
 
@@ -464,13 +490,14 @@ void bvh_rem_descriptor(uint64_t id)
 
 
 // create in-place given existing descriptor
-void bvh_create_host(vec3* lowers, vec3* uppers, int num_items, int constructor_type, BVH& bvh)
+void bvh_create_host(vec3* lowers, vec3* uppers, int num_items, int constructor_type, BVH& bvh, int leaf_size)
 {
     memset(&bvh, 0, sizeof(BVH));
 
     bvh.item_lowers = lowers;
     bvh.item_uppers = uppers;
     bvh.num_items = num_items;
+    bvh.leaf_size = leaf_size;
 
     TopDownBVHBuilder builder;
     builder.build(bvh, lowers, uppers, num_items, constructor_type);
@@ -496,10 +523,10 @@ void bvh_destroy_host(BVH& bvh)
 
 } // namespace wp
 
-uint64_t wp_bvh_create_host(vec3* lowers, vec3* uppers, int num_items, int constructor_type)
+uint64_t wp_bvh_create_host(vec3* lowers, vec3* uppers, int num_items, int constructor_type, int leaf_size)
 {
     BVH* bvh = new BVH();
-    wp::bvh_create_host(lowers, uppers, num_items, constructor_type, *bvh);
+    wp::bvh_create_host(lowers, uppers, num_items, constructor_type, *bvh, leaf_size);
 
     return (uint64_t)bvh;
 }
@@ -510,6 +537,12 @@ void wp_bvh_refit_host(uint64_t id)
     wp::bvh_refit_host(*bvh);
 }
 
+void wp_bvh_rebuild_host(uint64_t id, int constructor_type)
+{
+    BVH* bvh = (BVH*)(id);
+    wp::bvh_rebuild_host(*bvh, constructor_type);
+}
+
 void wp_bvh_destroy_host(uint64_t id)
 {
     BVH* bvh = (BVH*)(id);
@@ -521,8 +554,9 @@ void wp_bvh_destroy_host(uint64_t id)
 // stubs for non-CUDA platforms
 #if !WP_ENABLE_CUDA
 
-uint64_t wp_bvh_create_device(void* context, wp::vec3* lowers, wp::vec3* uppers, int num_items, int constructor_type) { return 0; }
+uint64_t wp_bvh_create_device(void* context, wp::vec3* lowers, wp::vec3* uppers, int num_items, int constructor_type, int leaf_size) { return 0; }
 void wp_bvh_refit_device(uint64_t id) {}
 void wp_bvh_destroy_device(uint64_t id) {}
+void wp_bvh_rebuild_device(uint64_t id) {}
 
 #endif // !WP_ENABLE_CUDA