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

warp/native/quat.h

@@ -19,6 +19,15 @@ struct quat_t
     // zero constructor for adjoint variable initialization
     inline CUDA_CALLABLE quat_t(Type x=Type(0), Type y=Type(0), Type z=Type(0), Type w=Type(0)) : x(x), y(y), z(z), w(w) {}    
     explicit inline CUDA_CALLABLE quat_t(const vec_t<3,Type>& v, Type w=Type(0)) : x(v[0]), y(v[1]), z(v[2]), w(w) {}
+    
+    template<typename OtherType>
+    explicit inline CUDA_CALLABLE quat_t(const quat_t<OtherType>& other)
+    {
+        x = static_cast<Type>(other.x);
+        y = static_cast<Type>(other.y);
+        z = static_cast<Type>(other.z);
+        w = static_cast<Type>(other.w);
+    }
 
     // imaginary part
     Type x;
@@ -73,7 +82,17 @@ inline CUDA_CALLABLE void adj_quat_t(const vec_t<3,Type>& v, Type w, vec_t<3,Typ
     adj_v[0] += adj_ret.x;
     adj_v[1] += adj_ret.y;
     adj_v[2] += adj_ret.z;
-    adj_w   += adj_ret.w;
+    adj_w    += adj_ret.w;
+}
+
+// casting constructor adjoint
+template<typename Type, typename OtherType>
+inline CUDA_CALLABLE void adj_quat_t(const quat_t<OtherType>& other, quat_t<OtherType>& adj_other, const quat_t<Type>& adj_ret)
+{
+    adj_other.x += static_cast<OtherType>(adj_ret.x);
+    adj_other.y += static_cast<OtherType>(adj_ret.y);
+    adj_other.z += static_cast<OtherType>(adj_ret.z);
+    adj_other.w += static_cast<OtherType>(adj_ret.w);
 }
 
 // forward methods
@@ -206,12 +225,24 @@ inline CUDA_CALLABLE quat_t<Type> div(quat_t<Type> q, Type s)
     return quat_t<Type>(q.x/s, q.y/s, q.z/s, q.w/s);
 }
 
+template<typename Type>
+inline CUDA_CALLABLE quat_t<Type> div(Type s, quat_t<Type> q)
+{
+    return quat_t<Type>(s/q.x, s/q.y, s/q.z, s/q.w);
+}
+
 template<typename Type>
 inline CUDA_CALLABLE quat_t<Type> operator / (quat_t<Type> a, Type s)
 {
     return div(a,s);
 }
 
+template<typename Type>
+inline CUDA_CALLABLE quat_t<Type> operator / (Type s, quat_t<Type> a)
+{
+    return div(s,a);
+}
+
 template<typename Type>
 inline CUDA_CALLABLE quat_t<Type> operator*(Type s, const quat_t<Type>& a)
 {
@@ -321,7 +352,7 @@ inline CUDA_CALLABLE quat_t<Type> quat_from_matrix(const mat_t<3,3,Type>& m)
 }
 
 template<typename Type>
-inline CUDA_CALLABLE Type index(const quat_t<Type>& a, int idx)
+inline CUDA_CALLABLE Type extract(const quat_t<Type>& a, int idx)
 {
 #if FP_CHECK
     if (idx < 0 || idx > 3)
@@ -357,7 +388,7 @@ CUDA_CALLABLE inline void adj_lerp(const quat_t<Type>& a, const quat_t<Type>& b,
 }
 
 template<typename Type>
-inline CUDA_CALLABLE void adj_index(const quat_t<Type>& a, int idx, quat_t<Type>& adj_a, int & adj_idx, Type & adj_ret)
+inline CUDA_CALLABLE void adj_extract(const quat_t<Type>& a, int idx, quat_t<Type>& adj_a, int & adj_idx, Type & adj_ret)
 {
 #if FP_CHECK
     if (idx < 0 || idx > 3)
@@ -367,7 +398,7 @@ inline CUDA_CALLABLE void adj_index(const quat_t<Type>& a, int idx, quat_t<Type>
     }
 #endif
 
-    // See wp::index(const quat_t<Type>& a, int idx) note
+    // See wp::extract(const quat_t<Type>& a, int idx) note
     if (idx == 0)       {adj_a.x += adj_ret;}
     else if (idx == 1)  {adj_a.y += adj_ret;}
     else if (idx == 2)  {adj_a.z += adj_ret;}
@@ -504,9 +535,14 @@ inline CUDA_CALLABLE void tensordot(const quat_t<Type>& a, const quat_t<Type>& b
 }
 
 template<typename Type>
-inline CUDA_CALLABLE void adj_length(const quat_t<Type>& a, quat_t<Type>& adj_a, const Type adj_ret)
+inline CUDA_CALLABLE void adj_length(const quat_t<Type>& a, Type ret, quat_t<Type>& adj_a, const Type adj_ret)
 {
-    adj_a += normalize(a)*adj_ret;
+    if (ret > Type(kEps))
+    {
+        Type inv_l = Type(1)/ret;
+
+        adj_a += quat_t<Type>(a.x*inv_l, a.y*inv_l, a.z*inv_l, a.w*inv_l) * adj_ret;
+    }
 }
 
 template<typename Type>
@@ -589,6 +625,13 @@ inline CUDA_CALLABLE void adj_div(quat_t<Type> a, Type s, quat_t<Type>& adj_a, T
     adj_a += adj_ret / s;
 }
 
+template<typename Type>
+inline CUDA_CALLABLE void adj_div(Type s, quat_t<Type> a, Type& adj_s, quat_t<Type>& adj_a, const quat_t<Type>& adj_ret)
+{
+    adj_s -= dot(a, adj_ret)/ (s * s); // - a / s^2
+    adj_a += s / adj_ret;
+}
+
 template<typename Type>
 inline CUDA_CALLABLE void adj_quat_rotate(const quat_t<Type>& q, const vec_t<3,Type>& p, quat_t<Type>& adj_q, vec_t<3,Type>& adj_p, const vec_t<3,Type>& adj_ret)
 {
@@ -658,7 +701,7 @@ inline CUDA_CALLABLE void adj_quat_rotate_inv(const quat_t<Type>& q, const vec_t
 }
 
 template<typename Type>
-inline CUDA_CALLABLE void adj_quat_slerp(const quat_t<Type>& q0, const quat_t<Type>& q1, Type t, quat_t<Type>& adj_q0, quat_t<Type>& adj_q1, Type& adj_t, const quat_t<Type>& adj_ret)
+inline CUDA_CALLABLE void adj_quat_slerp(const quat_t<Type>& q0, const quat_t<Type>& q1, Type t, quat_t<Type>& ret, quat_t<Type>& adj_q0, quat_t<Type>& adj_q1, Type& adj_t, const quat_t<Type>& adj_ret)
 {
     vec_t<3,Type> axis;
     Type angle;
@@ -669,7 +712,7 @@ inline CUDA_CALLABLE void adj_quat_slerp(const quat_t<Type>& q0, const quat_t<Ty
     angle = angle * 0.5;
     
     // adj_t
-    adj_t += dot(mul(quat_slerp(q0, q1, t), quat_t<Type>(angle*axis[0], angle*axis[1], angle*axis[2], Type(0))), adj_ret);
+    adj_t += dot(mul(ret, quat_t<Type>(angle*axis[0], angle*axis[1], angle*axis[2], Type(0))), adj_ret);
 
     // adj_q0
     quat_t<Type> q_inc_x_q0;

warp/native/rand.h

@@ -9,8 +9,8 @@
 # pragma once
 #include "array.h"
 
-#ifndef M_PI
-#define M_PI 3.14159265358979323846f
+#ifndef M_PI_F
+#define M_PI_F 3.14159265358979323846f
 #endif
 
 namespace wp
@@ -33,7 +33,7 @@ inline CUDA_CALLABLE float randf(uint32& state) { state = rand_pcg(state); retur
 inline CUDA_CALLABLE float randf(uint32& state, float min, float max) { return (max - min) * randf(state) + min; }
 
 // Box-Muller method
-inline CUDA_CALLABLE float randn(uint32& state) { return sqrt(-2.f * log(randf(state))) * cos(2.f * M_PI * randf(state)); }
+inline CUDA_CALLABLE float randn(uint32& state) { return sqrt(-2.f * log(randf(state))) * cos(2.f * M_PI_F * randf(state)); }
 
 inline CUDA_CALLABLE void adj_rand_init(int seed, int& adj_seed, float adj_ret) {}
 inline CUDA_CALLABLE void adj_rand_init(int seed, int offset, int& adj_seed, int& adj_offset, float adj_ret) {}
@@ -55,14 +55,14 @@ inline CUDA_CALLABLE int sample_cdf(uint32& state, const array_t<float>& cdf)
 inline CUDA_CALLABLE vec2 sample_triangle(uint32& state)
 {
     float r = sqrt(randf(state));
-    float u = 1.0 - r;
+    float u = 1.f - r;
     float v = randf(state) * r;
     return vec2(u, v);
 }
 
 inline CUDA_CALLABLE vec2 sample_unit_ring(uint32& state)
 {
-    float theta = randf(state, 0.f, 2.f*M_PI);
+    float theta = randf(state, 0.f, 2.f*M_PI_F);
     float x = cos(theta);
     float y = sin(theta);
     return vec2(x, y);
@@ -71,7 +71,7 @@ inline CUDA_CALLABLE vec2 sample_unit_ring(uint32& state)
 inline CUDA_CALLABLE vec2 sample_unit_disk(uint32& state)
 {
     float r = sqrt(randf(state));
-    float theta = randf(state, 0.f, 2.f*M_PI);
+    float theta = randf(state, 0.f, 2.f*M_PI_F);
     float x = r * cos(theta);
     float y = r * sin(theta);
     return vec2(x, y);
@@ -80,7 +80,7 @@ inline CUDA_CALLABLE vec2 sample_unit_disk(uint32& state)
 inline CUDA_CALLABLE vec3 sample_unit_sphere_surface(uint32& state)
 {
     float phi = acos(1.f - 2.f * randf(state));
-    float theta = randf(state, 0.f, 2.f*M_PI);
+    float theta = randf(state, 0.f, 2.f*M_PI_F);
     float x = cos(theta) * sin(phi);
     float y = sin(theta) * sin(phi);
     float z = cos(phi);
@@ -90,7 +90,7 @@ inline CUDA_CALLABLE vec3 sample_unit_sphere_surface(uint32& state)
 inline CUDA_CALLABLE vec3 sample_unit_sphere(uint32& state)
 {
     float phi = acos(1.f  - 2.f * randf(state));
-    float theta = randf(state, 0.f, 2.f*M_PI);
+    float theta = randf(state, 0.f, 2.f*M_PI_F);
     float r = pow(randf(state), 1.f/3.f);
     float x = r * cos(theta) * sin(phi);
     float y = r * sin(theta) * sin(phi);
@@ -101,7 +101,7 @@ inline CUDA_CALLABLE vec3 sample_unit_sphere(uint32& state)
 inline CUDA_CALLABLE vec3 sample_unit_hemisphere_surface(uint32& state)
 {
     float phi = acos(1.f - randf(state));
-    float theta = randf(state, 0.f, 2.f*M_PI);
+    float theta = randf(state, 0.f, 2.f*M_PI_F);
     float x = cos(theta) * sin(phi);
     float y = sin(theta) * sin(phi);
     float z = cos(phi);
@@ -111,7 +111,7 @@ inline CUDA_CALLABLE vec3 sample_unit_hemisphere_surface(uint32& state)
 inline CUDA_CALLABLE vec3 sample_unit_hemisphere(uint32& state)
 {
     float phi = acos(1.f - randf(state));
-    float theta = randf(state, 0.f, 2.f*M_PI);
+    float theta = randf(state, 0.f, 2.f*M_PI_F);
     float r = pow(randf(state), 1.f/3.f);
     float x = r * cos(theta) * sin(phi);
     float y = r * sin(theta) * sin(phi);
@@ -134,6 +134,15 @@ inline CUDA_CALLABLE vec3 sample_unit_cube(uint32& state)
     return vec3(x, y, z);
 }
 
+inline CUDA_CALLABLE vec4 sample_unit_hypercube(uint32& state)
+{
+    float a = randf(state) - 0.5f;
+    float b = randf(state) - 0.5f;
+    float c = randf(state) - 0.5f;
+    float d = randf(state) - 0.5f;
+    return vec4(a, b, c, d);
+}
+
 inline CUDA_CALLABLE void adj_sample_cdf(uint32& state, const array_t<float>& cdf, uint32& adj_state, array_t<float>& adj_cdf, const int& adj_ret) {}
 inline CUDA_CALLABLE void adj_sample_triangle(uint32& state, uint32& adj_state, const vec2& adj_ret) {}
 inline CUDA_CALLABLE void adj_sample_unit_ring(uint32& state, uint32& adj_state, const vec2& adj_ret) {}
@@ -144,6 +153,7 @@ inline CUDA_CALLABLE void adj_sample_unit_hemisphere_surface(uint32& state, uint
 inline CUDA_CALLABLE void adj_sample_unit_hemisphere(uint32& state, uint32& adj_state, const vec3& adj_ret) {}
 inline CUDA_CALLABLE void adj_sample_unit_square(uint32& state, uint32& adj_state, const vec2& adj_ret) {}
 inline CUDA_CALLABLE void adj_sample_unit_cube(uint32& state, uint32& adj_state, const vec3& adj_ret) {}
+inline CUDA_CALLABLE void adj_sample_unit_hypercube(uint32& state, uint32& adj_state, const vec3& adj_ret) {}
 
 /*
  * log-gamma function to support some of these distributions. The
@@ -158,17 +168,17 @@ inline CUDA_CALLABLE float random_loggam(float x)
     float x0, x2, lg2pi, gl, gl0;
     uint32 n;
 
-    const float a[10] = {8.333333333333333e-02, -2.777777777777778e-03,
-                        7.936507936507937e-04, -5.952380952380952e-04,
-                        8.417508417508418e-04, -1.917526917526918e-03,
-                        6.410256410256410e-03, -2.955065359477124e-02,
-                        1.796443723688307e-01, -1.39243221690590e+00};
+    const float a[10] = {8.333333333333333e-02f, -2.777777777777778e-03f,
+                        7.936507936507937e-04f, -5.952380952380952e-04f,
+                        8.417508417508418e-04f, -1.917526917526918e-03f,
+                        6.410256410256410e-03f, -2.955065359477124e-02f,
+                        1.796443723688307e-01f, -1.39243221690590e+00f};
 
-    if ((x == 1.0) || (x == 2.0))
+    if ((x == 1.f) || (x == 2.f))
     {
-        return 0.0;
+        return 0.f;
     }
-    else if (x < 7.0)
+    else if (x < 7.f)
     {
         n = uint32((7 - x));
     }
@@ -178,8 +188,8 @@ inline CUDA_CALLABLE float random_loggam(float x)
     }
 
     x0 = x + float(n);
-    x2 = (1.0 / x0) * (1.0 / x0);
-    // log(2 * M_PI)
+    x2 = (1.f / x0) * (1.f / x0);
+    // log(2 * M_PI_F)
     lg2pi = 1.8378770664093453f;
     gl0 = a[9];
     for (int i = 8; i >= 0; i--)
@@ -187,13 +197,13 @@ inline CUDA_CALLABLE float random_loggam(float x)
         gl0 *= x2;
         gl0 += a[i];
     }
-    gl = gl0 / x0 + 0.5 * lg2pi + (x0 - 0.5) * log(x0) - x0;
-    if (x < 7.0)
+    gl = gl0 / x0 + 0.5f * lg2pi + (x0 - 0.5f) * log(x0) - x0;
+    if (x < 7.f)
     {
         for (uint32 k = 1; k <= n; k++)
         {
-            gl -= log(x0 - 1.0);
-            x0 -= 1.0;
+            gl -= log(x0 - 1.f);
+            x0 -= 1.f;
         }
     }
     return gl;
@@ -205,7 +215,7 @@ inline CUDA_CALLABLE uint32 random_poisson_mult(uint32& state, float lam) {
 
     enlam = exp(-lam);
     X = 0;
-    prod = 1.0;
+    prod = 1.f;
 
     while (1)
     {
@@ -234,22 +244,22 @@ inline CUDA_CALLABLE uint32 random_poisson(uint32& state, float lam)
 
     slam = sqrt(lam);
     loglam = log(lam);
-    b = 0.931 + 2.53 * slam;
-    a = -0.059 + 0.02483 * b;
-    invalpha = 1.1239 + 1.1328 / (b - 3.4);
-    vr = 0.9277 - 3.6224 / (b - 2.0);
+    b = 0.931f + 2.53f * slam;
+    a = -0.059f + 0.02483f * b;
+    invalpha = 1.1239f + 1.1328f / (b - 3.4f);
+    vr = 0.9277f - 3.6224f / (b - 2.f);
 
     while (1)
     {
-        U = randf(state) - 0.5;
+        U = randf(state) - 0.5f;
         V = randf(state);
-        us = 0.5 - abs(U);
-        k = uint32(floor((2 * a / us + b) * U + lam + 0.43));
-        if ((us >= 0.07) && (V <= vr))
+        us = 0.5f - abs(U);
+        k = uint32(floor((2.f * a / us + b) * U + lam + 0.43f));
+        if ((us >= 0.07f) && (V <= vr))
         {
             return k;
         }
-        if ((us < 0.013) && (V > us))
+        if ((us < 0.013f) && (V > us))
         {
             continue;
         }
@@ -261,7 +271,7 @@ inline CUDA_CALLABLE uint32 random_poisson(uint32& state, float lam)
 }
 
 /*
-* Adpated from NumPy's implementation
+* Adapted from NumPy's implementation
 * Warp's state variable is half the precision of NumPy's so
 * poisson implementation uses half the precision used in NumPy's implementation
 * both precisions appear to converge in the statistical limit

warp/native/range.h

@@ -15,8 +15,12 @@ namespace wp
 // represents a built-in Python range() loop
 struct range_t
 {
-    CUDA_CALLABLE range_t() {}
-    CUDA_CALLABLE range_t(int) {} // for backward pass
+    CUDA_CALLABLE range_t()
+        : start(0),
+          end(0),
+          step(0),
+          i(0)
+    {}
 
     int start;
     int end;

warp/native/reduce.cpp

@@ -97,7 +97,7 @@ template <typename T> void array_sum_host(const T *ptr_a, T *ptr_out, int count,
         accumulate_func = dyn_len_sum<T>;
     }
 
-    *ptr_out = 0.0f;
+    memset(ptr_out, 0, sizeof(T)*type_length);
     for (int i = 0; i < count; ++i)
         accumulate_func(ptr_a + i * stride, ptr_out, type_length);
 }

warp/native/reduce.cu

@@ -103,23 +103,22 @@ template <typename T> void array_sum_device(const T *ptr_a, T *ptr_out, int coun
     assert((byte_stride % sizeof(T)) == 0);
     const int stride = byte_stride / sizeof(T);
 
-    void *context = cuda_context_get_current();
-    TemporaryBuffer &cub_temp = g_temp_buffer_map[context];
-
-    ContextGuard guard(context);
+    ContextGuard guard(cuda_context_get_current());
     cudaStream_t stream = static_cast<cudaStream_t>(cuda_stream_get_current());
 
     cub_strided_iterator<const T> ptr_strided{ptr_a, stride};
 
     size_t buff_size = 0;
     check_cuda(cub::DeviceReduce::Sum(nullptr, buff_size, ptr_strided, ptr_out, count, stream));
-    cub_temp.ensure_fits(buff_size);
+    void* temp_buffer = alloc_temp_device(WP_CURRENT_CONTEXT, buff_size);
 
     for (int k = 0; k < type_length; ++k)
     {
         cub_strided_iterator<const T> ptr_strided{ptr_a + k, stride};
-        check_cuda(cub::DeviceReduce::Sum(cub_temp.buffer, buff_size, ptr_strided, ptr_out + k, count, stream));
+        check_cuda(cub::DeviceReduce::Sum(temp_buffer, buff_size, ptr_strided, ptr_out + k, count, stream));
     }
+
+    free_temp_device(WP_CURRENT_CONTEXT, temp_buffer);
 }
 
 template <typename T>
@@ -265,19 +264,18 @@ void array_inner_device(const ElemT *ptr_a, const ElemT *ptr_b, ScalarT *ptr_out
     const int stride_a = byte_stride_a / sizeof(ElemT);
     const int stride_b = byte_stride_b / sizeof(ElemT);
 
-    void *context = cuda_context_get_current();
-    TemporaryBuffer &cub_temp = g_temp_buffer_map[context];
-
-    ContextGuard guard(context);
+    ContextGuard guard(cuda_context_get_current());
     cudaStream_t stream = static_cast<cudaStream_t>(cuda_stream_get_current());
 
     cub_inner_product_iterator<ElemT, ScalarT> inner_iterator{ptr_a, ptr_b, stride_a, stride_b, type_length};
 
     size_t buff_size = 0;
     check_cuda(cub::DeviceReduce::Sum(nullptr, buff_size, inner_iterator, ptr_out, count, stream));
-    cub_temp.ensure_fits(buff_size);
+    void* temp_buffer = alloc_temp_device(WP_CURRENT_CONTEXT, buff_size);
+
+    check_cuda(cub::DeviceReduce::Sum(temp_buffer, buff_size, inner_iterator, ptr_out, count, stream));
 
-    check_cuda(cub::DeviceReduce::Sum(cub_temp.buffer, buff_size, inner_iterator, ptr_out, count, stream));
+    free_temp_device(WP_CURRENT_CONTEXT, temp_buffer);
 }
 
 template <typename T>

warp/native/runlength_encode.cu

@@ -3,8 +3,6 @@
 #include "warp.h"
 #include "cuda_util.h"
 
-#include "temp_buffer.h"
-
 #define THRUST_IGNORE_CUB_VERSION_CHECK
 #include <cub/device/device_run_length_encode.cuh>
 
@@ -15,11 +13,7 @@ void runlength_encode_device(int n,
                              int *run_lengths,
                              int *run_count)
 {
-    void *context = cuda_context_get_current();
-
-    TemporaryBuffer &cub_temp = g_temp_buffer_map[context];
-
-    ContextGuard guard(context);
+    ContextGuard guard(cuda_context_get_current());
     cudaStream_t stream = static_cast<cudaStream_t>(cuda_stream_get_current());
 
     size_t buff_size = 0;
@@ -27,11 +21,13 @@ void runlength_encode_device(int n,
         nullptr, buff_size, values, run_values, run_lengths, run_count,
         n, stream));
 
-    cub_temp.ensure_fits(buff_size);
+    void* temp_buffer = alloc_temp_device(WP_CURRENT_CONTEXT, buff_size);
 
     check_cuda(cub::DeviceRunLengthEncode::Encode(
-        cub_temp.buffer, buff_size, values, run_values, run_lengths, run_count,
+        temp_buffer, buff_size, values, run_values, run_lengths, run_count,
         n, stream));
+
+    free_temp_device(WP_CURRENT_CONTEXT, temp_buffer);
 }
 
 void runlength_encode_int_device(
@@ -47,4 +43,4 @@ void runlength_encode_int_device(
         reinterpret_cast<int *>(run_values),
         reinterpret_cast<int *>(run_lengths),
         reinterpret_cast<int *>(run_count));
-}
+}

warp/native/scan.cu

@@ -1,8 +1,6 @@
 #include "warp.h"
 #include "scan.h"
 
-#include "temp_buffer.h"
-
 #define THRUST_IGNORE_CUB_VERSION_CHECK
 
 #include <cub/device/device_scan.cuh>
@@ -10,29 +8,28 @@
 template<typename T>
 void scan_device(const T* values_in, T* values_out, int n, bool inclusive)
 {
-    void *context = cuda_context_get_current();
-    TemporaryBuffer &cub_temp = g_temp_buffer_map[context];
-
-    ContextGuard guard(context);
+    ContextGuard guard(cuda_context_get_current());
 
     cudaStream_t stream = static_cast<cudaStream_t>(cuda_stream_get_current());
 
     // compute temporary memory required
 	size_t scan_temp_size;
     if (inclusive) {
-        cub::DeviceScan::InclusiveSum(NULL, scan_temp_size, values_in, values_out, n);
+        check_cuda(cub::DeviceScan::InclusiveSum(NULL, scan_temp_size, values_in, values_out, n));
     } else {
-        cub::DeviceScan::ExclusiveSum(NULL, scan_temp_size, values_in, values_out, n);
+        check_cuda(cub::DeviceScan::ExclusiveSum(NULL, scan_temp_size, values_in, values_out, n));
     }
 
-    cub_temp.ensure_fits(scan_temp_size);
+    void* temp_buffer = alloc_temp_device(WP_CURRENT_CONTEXT, scan_temp_size);
 
     // scan
     if (inclusive) {
-        cub::DeviceScan::InclusiveSum(cub_temp.buffer, scan_temp_size, values_in, values_out, n, (cudaStream_t)cuda_stream_get_current());
+        check_cuda(cub::DeviceScan::InclusiveSum(temp_buffer, scan_temp_size, values_in, values_out, n, stream));
     } else {
-        cub::DeviceScan::ExclusiveSum(cub_temp.buffer, scan_temp_size, values_in, values_out, n, (cudaStream_t)cuda_stream_get_current());
+        check_cuda(cub::DeviceScan::ExclusiveSum(temp_buffer, scan_temp_size, values_in, values_out, n, stream));
     }
+
+    free_temp_device(WP_CURRENT_CONTEXT, temp_buffer);
 }
 
 template void scan_device(const int*, int*, int, bool);

warp/native/sparse.cpp

@@ -179,10 +179,10 @@ void bsr_transpose_host(int rows_per_block, int cols_per_block, int row_count, i
     const int block_size = rows_per_block * cols_per_block;
 
     void (*block_transpose_func)(const T *, T *, int, int) = bsr_dyn_block_transpose<T>;
-    switch (row_count)
+    switch (rows_per_block)
     {
     case 1:
-        switch (col_count)
+        switch (cols_per_block)
         {
         case 1:
             block_transpose_func = bsr_fixed_block_transpose<1, 1, T>;
@@ -196,7 +196,7 @@ void bsr_transpose_host(int rows_per_block, int cols_per_block, int row_count, i
         }
         break;
     case 2:
-        switch (col_count)
+        switch (cols_per_block)
         {
         case 1:
             block_transpose_func = bsr_fixed_block_transpose<2, 1, T>;
@@ -210,7 +210,7 @@ void bsr_transpose_host(int rows_per_block, int cols_per_block, int row_count, i
         }
         break;
     case 3:
-        switch (col_count)
+        switch (cols_per_block)
         {
         case 1:
             block_transpose_func = bsr_fixed_block_transpose<3, 1, T>;