warp-lang 1.7.2__py3-none-manylinux_2_34_aarch64.whl → 1.8.0__py3-none-manylinux_2_34_aarch64.whl

warp/native/sparse.cu

@@ -17,6 +17,8 @@
 
 #include "cuda_util.h"
 #include "warp.h"
+#include "stdint.h"
+#include <cstdint>
 
 #define THRUST_IGNORE_CUB_VERSION_CHECK
 
@@ -45,29 +47,55 @@ template <typename T> struct BsrBlockIsNotZero
 {
     int block_size;
     const T* values;
+    T zero_mask;
 
-    CUDA_CALLABLE_DEVICE bool operator()(int i) const
+    BsrBlockIsNotZero(int block_size, const void* values, const uint64_t zero_mask)
+        : block_size(block_size), values(static_cast<const T*>(values)), zero_mask(static_cast<const T>(zero_mask)) 
+        {}
+
+    CUDA_CALLABLE_DEVICE bool operator()(int block) const
     {
         if (!values)
             return true;
 
-        const T* val = values + i * block_size;
+        const T* val = values + block * block_size;
         for (int i = 0; i < block_size; ++i, ++val)
         {
-            if (*val != T(0))
+            if ((*val & zero_mask) != 0)
                 return true;
         }
         return false;
     }
 };
 
+template <> struct BsrBlockIsNotZero<void>
+{
+    BsrBlockIsNotZero(int block_size, const void* values, const uint64_t zero_mask)
+    {}
+
+    CUDA_CALLABLE_DEVICE bool operator()(int block) const
+    {
+        return true;
+    }
+};
+
 struct BsrBlockInMask
 {
+    const int nrow;
+    const int ncol;
     const int* bsr_offsets;
     const int* bsr_columns;
+    const int* device_nnz;
 
-    CUDA_CALLABLE_DEVICE bool operator()(int row, int col) const
+    CUDA_CALLABLE_DEVICE bool operator()(int index, int row, int col) const
     {
+        if (device_nnz != nullptr && index >= *device_nnz)
+            return false;
+    
+        if (row < 0 || row >= nrow || col < 0 || col >= ncol){
+            return false;
+        }
+
         if (bsr_offsets == nullptr)
             return true;
 
@@ -93,9 +121,9 @@ struct BsrBlockInMask
 };
 
 template <typename T>
-__global__ void bsr_fill_triplet_key_values(const int nnz, const int nrow, const int* tpl_rows, const int* tpl_columns,
+__global__ void bsr_fill_triplet_key_values(const int nnz, const int* tpl_rows, const int* tpl_columns,
                                             const BsrBlockIsNotZero<T> nonZero, const BsrBlockInMask mask,
-                                            uint32_t* block_indices, BsrRowCol* tpl_row_col)
+                                            int* block_indices, BsrRowCol* tpl_row_col)
 {
     int block = blockIdx.x * blockDim.x + threadIdx.x;
     if (block >= nnz)
@@ -103,10 +131,10 @@ __global__ void bsr_fill_triplet_key_values(const int nnz, const int nrow, const
 
     const int row = tpl_rows[block];
     const int col = tpl_columns[block];
-    const bool is_valid = row >= 0 && row < nrow;
 
     const BsrRowCol row_col =
-        is_valid && nonZero(block) && mask(row, col) ? bsr_combine_row_col(row, col) : PRUNED_ROWCOL;
+        mask(block, row, col) && nonZero(block) ? bsr_combine_row_col(row, col) : PRUNED_ROWCOL;
+
     tpl_row_col[block] = row_col;
     block_indices[block] = block;
 }
@@ -153,126 +181,34 @@ __global__ void bsr_find_row_offsets(uint32_t row_count, const T* d_nnz, const B
     row_offsets[row] = lower;
 }
 
-template <typename T>
-__global__ void bsr_merge_blocks(const int* d_nnz, int block_size, const uint32_t* block_offsets,
-                                 const uint32_t* sorted_block_indices, const BsrRowCol* unique_row_cols,
-                                 const T* tpl_values, int* bsr_cols, T* bsr_values)
-
+__global__ void bsr_set_column(const int* d_nnz, const BsrRowCol* unique_row_cols, int* bsr_cols)
 {
     const uint32_t i = blockIdx.x * blockDim.x + threadIdx.x;
-
     if (i >= *d_nnz)
         return;
-
     const BsrRowCol row_col = unique_row_cols[i];
     bsr_cols[i] = bsr_get_col(row_col);
-
-    // Accumulate merged block values
-    if (row_col == PRUNED_ROWCOL || bsr_values == nullptr)
-        return;
-
-    const uint32_t beg = i ? block_offsets[i - 1] : 0;
-    const uint32_t end = block_offsets[i];
-
-    T* bsr_val = bsr_values + i * block_size;
-    const T* tpl_val = tpl_values + sorted_block_indices[beg] * block_size;
-
-    for (int k = 0; k < block_size; ++k)
-    {
-        bsr_val[k] = tpl_val[k];
-    }
-
-    for (uint32_t cur = beg + 1; cur != end; ++cur)
-    {
-        const T* tpl_val = tpl_values + sorted_block_indices[cur] * block_size;
-        for (int k = 0; k < block_size; ++k)
-        {
-            bsr_val[k] += tpl_val[k];
-        }
-    }
 }
 
 template <typename T>
-void bsr_matrix_from_triplets_device(const int rows_per_block, const int cols_per_block, const int row_count,
-                                     const int nnz, const int* tpl_rows, const int* tpl_columns, const T* tpl_values,
-                                     const bool prune_numerical_zeros, const bool masked, int* bsr_offsets,
-                                     int* bsr_columns, T* bsr_values, int* bsr_nnz, void* bsr_nnz_event)
+void launch_bsr_fill_triplet_key_values(
+    const int block_size, 
+    const int nnz, 
+    const BsrBlockInMask& mask,
+    const int* tpl_rows, 
+    const int* tpl_columns, 
+    const void* tpl_values,
+    const uint64_t scalar_zero_mask,
+    int* block_indices,
+    BsrRowCol* row_col
+     )
 {
-    const int block_size = rows_per_block * cols_per_block;
-
-    void* context = cuda_context_get_current();
-    ContextGuard guard(context);
-
-    // Per-context cached temporary buffers
-    // BsrFromTripletsTemp& bsr_temp = g_bsr_from_triplets_temp_map[context];
-
-    cudaStream_t stream = static_cast<cudaStream_t>(cuda_stream_get_current());
-
-    ScopedTemporary<uint32_t> block_indices(context, 2 * nnz + 1);
-    ScopedTemporary<BsrRowCol> combined_row_col(context, 2 * nnz);
-
-    cub::DoubleBuffer<uint32_t> d_keys(block_indices.buffer(), block_indices.buffer() + nnz);
-    cub::DoubleBuffer<BsrRowCol> d_values(combined_row_col.buffer(), combined_row_col.buffer() + nnz);
-
-    uint32_t* unique_triplet_count = block_indices.buffer() + 2 * nnz;
-
-    // Combine rows and columns so we can sort on them both
-    BsrBlockIsNotZero<T> isNotZero{block_size, prune_numerical_zeros ? tpl_values : nullptr};
-    BsrBlockInMask mask{masked ? bsr_offsets : nullptr, bsr_columns};
+    BsrBlockIsNotZero<T> isNotZero{block_size, tpl_values, scalar_zero_mask};
     wp_launch_device(WP_CURRENT_CONTEXT, bsr_fill_triplet_key_values, nnz,
-                     (nnz, row_count, tpl_rows, tpl_columns, isNotZero, mask, d_keys.Current(), d_values.Current()));
-
-    // Sort
-    {
-        size_t buff_size = 0;
-        check_cuda(cub::DeviceRadixSort::SortPairs(nullptr, buff_size, d_values, d_keys, nnz, 0, 64, stream));
-        ScopedTemporary<> temp(context, buff_size);
-        check_cuda(cub::DeviceRadixSort::SortPairs(temp.buffer(), buff_size, d_values, d_keys, nnz, 0, 64, stream));
-    }
-
-    // Runlength encode row-col sequences
-    {
-        size_t buff_size = 0;
-        check_cuda(cub::DeviceRunLengthEncode::Encode(nullptr, buff_size, d_values.Current(), d_values.Alternate(),
-                                                      d_keys.Alternate(), unique_triplet_count, nnz, stream));
-        ScopedTemporary<> temp(context, buff_size);
-        check_cuda(cub::DeviceRunLengthEncode::Encode(temp.buffer(), buff_size, d_values.Current(),
-                                                      d_values.Alternate(), d_keys.Alternate(), unique_triplet_count,
-                                                      nnz, stream));
-    }
-
-    // Compute row offsets from sorted unique blocks
-    wp_launch_device(WP_CURRENT_CONTEXT, bsr_find_row_offsets, row_count + 1,
-                     (row_count, unique_triplet_count, d_values.Alternate(), bsr_offsets));
-
-    if (bsr_nnz)
-    {
-        // Copy nnz to host, and record an event for the completed transfer if desired
-
-        memcpy_d2h(WP_CURRENT_CONTEXT, bsr_nnz, bsr_offsets + row_count, sizeof(int), stream);
-
-        if (bsr_nnz_event)
-        {
-            cuda_event_record(bsr_nnz_event, stream);
-        }
-    }
-
-    // Scan repeated block counts
-    {
-        size_t buff_size = 0;
-        check_cuda(
-            cub::DeviceScan::InclusiveSum(nullptr, buff_size, d_keys.Alternate(), d_keys.Alternate(), nnz, stream));
-        ScopedTemporary<> temp(context, buff_size);
-        check_cuda(cub::DeviceScan::InclusiveSum(temp.buffer(), buff_size, d_keys.Alternate(), d_keys.Alternate(), nnz,
-                                                 stream));
-    }
-
-    // Accumulate repeated blocks and set column indices
-    wp_launch_device(WP_CURRENT_CONTEXT, bsr_merge_blocks, nnz,
-                     (bsr_offsets + row_count, block_size, d_keys.Alternate(), d_keys.Current(), d_values.Alternate(),
-                      tpl_values, bsr_columns, bsr_values));
+                     (nnz, tpl_rows, tpl_columns, isNotZero, mask, block_indices, row_col   ));
 }
 
+
 __global__ void bsr_transpose_fill_row_col(const int nnz_upper_bound, const int row_count, const int* bsr_offsets,
                                            const int* bsr_columns, int* block_indices, BsrRowCol* transposed_row_col)
 {
@@ -283,6 +219,8 @@ __global__ void bsr_transpose_fill_row_col(const int nnz_upper_bound, const int
         // Outside of allocated bounds, do nothing
         return;
     }
+    
+    block_indices[i] = i;
 
     if (i >= bsr_offsets[row_count])
     {
@@ -291,8 +229,6 @@ __global__ void bsr_transpose_fill_row_col(const int nnz_upper_bound, const int
         return;
     }
 
-    block_indices[i] = i;
-
     // Binary search for row
     int lower = 0;
     int upper = row_count - 1;
@@ -317,144 +253,153 @@ __global__ void bsr_transpose_fill_row_col(const int nnz_upper_bound, const int
     transposed_row_col[i] = row_col;
 }
 
-template <int Rows, int Cols, typename T> struct BsrBlockTransposer
+} // namespace
+
+
+WP_API void bsr_matrix_from_triplets_device(
+    const int block_size, 
+    int scalar_size,
+    const int row_count,
+    const int col_count,
+    const int nnz, 
+    const int* tpl_nnz,
+    const int* tpl_rows, 
+    const int* tpl_columns, 
+    const void* tpl_values,
+    const uint64_t scalar_zero_mask,
+    const bool masked_topology, 
+    int* tpl_block_offsets,
+    int* tpl_block_indices,
+    int* bsr_offsets,
+    int* bsr_columns,
+    int* bsr_nnz, void* bsr_nnz_event)
 {
-    void CUDA_CALLABLE_DEVICE operator()(const T* src, T* dest) const
+    void* context = cuda_context_get_current();
+    ContextGuard guard(context);
+
+    // Per-context cached temporary buffers
+    // BsrFromTripletsTemp& bsr_temp = g_bsr_from_triplets_temp_map[context];
+
+    cudaStream_t stream = static_cast<cudaStream_t>(cuda_stream_get_current());
+
+    ScopedTemporary<BsrRowCol> combined_row_col(context, 2 * size_t(nnz));
+    ScopedTemporary<int> unique_triplet_count(context, 1);
+
+    bool return_summed_blocks = tpl_block_offsets != nullptr && tpl_block_indices != nullptr;
+    if(!return_summed_blocks)
     {
-        for (int r = 0; r < Rows; ++r)
-        {
-            for (int c = 0; c < Cols; ++c)
-            {
-                dest[c * Rows + r] = src[r * Cols + c];
-            }
-        }
+        // if not provided, allocate temporary offset and indices buffers
+        tpl_block_offsets = static_cast<int*>(alloc_device(context, size_t(nnz) * sizeof(int)));
+        tpl_block_indices = static_cast<int*>(alloc_device(context,  size_t(nnz) * sizeof(int)));
     }
-};
 
-template <typename T> struct BsrBlockTransposer<-1, -1, T>
-{
 
-    int row_count;
-    int col_count;
+    cub::DoubleBuffer<int> d_keys(tpl_block_indices, tpl_block_offsets);
+    cub::DoubleBuffer<BsrRowCol> d_values(combined_row_col.buffer(), combined_row_col.buffer() + nnz);
 
-    void CUDA_CALLABLE_DEVICE operator()(const T* src, T* dest) const
+    // Combine rows and columns so we can sort on them both,
+    // ensuring that blocks that should be pruned are moved to the end 
+    BsrBlockInMask mask{row_count, col_count, masked_topology ? bsr_offsets : nullptr, bsr_columns, tpl_nnz};
+    if (scalar_zero_mask == 0 || tpl_values == nullptr)
+        scalar_size = 0;
+    switch(scalar_size)
     {
-        for (int r = 0; r < row_count; ++r)
-        {
-            for (int c = 0; c < col_count; ++c)
-            {
-                dest[c * row_count + r] = src[r * col_count + c];
-            }
-        }
+        case sizeof(uint8_t):
+            launch_bsr_fill_triplet_key_values<uint8_t>(block_size, nnz, mask, tpl_rows, tpl_columns, tpl_values, scalar_zero_mask, d_keys.Current(), d_values.Current());
+            break;
+        case sizeof(uint16_t):
+            launch_bsr_fill_triplet_key_values<uint16_t>(block_size, nnz, mask, tpl_rows, tpl_columns, tpl_values, scalar_zero_mask, d_keys.Current(), d_values.Current());
+            break;
+        case sizeof(uint32_t):
+            launch_bsr_fill_triplet_key_values<uint32_t>(block_size, nnz, mask, tpl_rows, tpl_columns, tpl_values, scalar_zero_mask, d_keys.Current(), d_values.Current());
+            break;
+        case sizeof(uint64_t):
+            launch_bsr_fill_triplet_key_values<uint64_t>(block_size, nnz, mask, tpl_rows, tpl_columns, tpl_values, scalar_zero_mask, d_keys.Current(), d_values.Current());
+            break;
+        default:
+            // no scalar-level pruning
+            launch_bsr_fill_triplet_key_values<void>(block_size, nnz, mask, tpl_rows, tpl_columns, tpl_values, scalar_zero_mask, d_keys.Current(), d_values.Current());
+            break;
     }
-};
 
-template <int Rows, int Cols, typename T>
-__global__ void bsr_transpose_blocks(const int* nnz, const int block_size, BsrBlockTransposer<Rows, Cols, T> transposer,
-                                     const int* block_indices, const BsrRowCol* transposed_indices, const T* bsr_values,
-                                     int* transposed_bsr_columns, T* transposed_bsr_values)
-{
-    int i = blockIdx.x * blockDim.x + threadIdx.x;
-    if (i >= *nnz)
-        return;
 
-    const int src_idx = block_indices[i];
+    // Sort
+    {
+        size_t buff_size = 0;
+        check_cuda(cub::DeviceRadixSort::SortPairs(nullptr, buff_size, d_values, d_keys, nnz, 0, 64, stream));
+        ScopedTemporary<> temp(context, buff_size);
+        check_cuda(cub::DeviceRadixSort::SortPairs(temp.buffer(), buff_size, d_values, d_keys, nnz, 0, 64, stream));
 
-    transposer(bsr_values + src_idx * block_size, transposed_bsr_values + i * block_size);
+        // Depending on data size and GPU architecture buffers may have been swapped or not
+        // Ensures the sorted keys are available in summed_block_indices if needed
+        if(return_summed_blocks && d_keys.Current() != tpl_block_indices)
+        {
+            check_cuda(cudaMemcpy(tpl_block_indices, d_keys.Current(), nnz * sizeof(int), cudaMemcpyDeviceToDevice));
+        }
+    }
 
-    transposed_bsr_columns[i] = bsr_get_col(transposed_indices[i]);
-}
+    // Runlength encode row-col sequences
+    {
+        size_t buff_size = 0;
+        check_cuda(cub::DeviceRunLengthEncode::Encode(nullptr, buff_size, d_values.Current(), d_values.Alternate(),
+                                                      tpl_block_offsets, unique_triplet_count.buffer(), nnz, stream));
+        ScopedTemporary<> temp(context, buff_size);
+        check_cuda(cub::DeviceRunLengthEncode::Encode(temp.buffer(), buff_size, d_values.Current(),
+                                                      d_values.Alternate(), tpl_block_offsets, unique_triplet_count.buffer(),
+                                                      nnz, stream));
+    }
 
-template <typename T>
-void launch_bsr_transpose_blocks(int nnz, const int* d_nnz, const int block_size, const int rows_per_block,
-                                 const int cols_per_block, const int* block_indices,
-                                 const BsrRowCol* transposed_indices, const T* bsr_values, int* transposed_bsr_columns,
-                                 T* transposed_bsr_values)
-{
+    // Compute row offsets from sorted unique blocks
+    wp_launch_device(WP_CURRENT_CONTEXT, bsr_find_row_offsets, row_count + 1,
+                     (row_count, unique_triplet_count.buffer(), d_values.Alternate(), bsr_offsets));
 
-    switch (rows_per_block)
+    if (bsr_nnz)
     {
-    case 1:
-        switch (cols_per_block)
-        {
-        case 1:
-            wp_launch_device(WP_CURRENT_CONTEXT, bsr_transpose_blocks, nnz,
-                             (d_nnz, block_size, BsrBlockTransposer<1, 1, T>{}, block_indices, transposed_indices,
-                              bsr_values, transposed_bsr_columns, transposed_bsr_values));
-            return;
-        case 2:
-            wp_launch_device(WP_CURRENT_CONTEXT, bsr_transpose_blocks, nnz,
-                             (d_nnz, block_size, BsrBlockTransposer<1, 2, T>{}, block_indices, transposed_indices,
-                              bsr_values, transposed_bsr_columns, transposed_bsr_values));
-            return;
-        case 3:
-            wp_launch_device(WP_CURRENT_CONTEXT, bsr_transpose_blocks, nnz,
-                             (d_nnz, block_size, BsrBlockTransposer<1, 3, T>{}, block_indices, transposed_indices,
-                              bsr_values, transposed_bsr_columns, transposed_bsr_values));
-            return;
-        }
-    case 2:
-        switch (cols_per_block)
-        {
-        case 1:
-            wp_launch_device(WP_CURRENT_CONTEXT, bsr_transpose_blocks, nnz,
-                             (d_nnz, block_size, BsrBlockTransposer<2, 1, T>{}, block_indices, transposed_indices,
-                              bsr_values, transposed_bsr_columns, transposed_bsr_values));
-            return;
-        case 2:
-            wp_launch_device(WP_CURRENT_CONTEXT, bsr_transpose_blocks, nnz,
-                             (d_nnz, block_size, BsrBlockTransposer<2, 2, T>{}, block_indices, transposed_indices,
-                              bsr_values, transposed_bsr_columns, transposed_bsr_values));
-            return;
-        case 3:
-            wp_launch_device(WP_CURRENT_CONTEXT, bsr_transpose_blocks, nnz,
-                             (d_nnz, block_size, BsrBlockTransposer<2, 3, T>{}, block_indices, transposed_indices,
-                              bsr_values, transposed_bsr_columns, transposed_bsr_values));
-            return;
-        }
-    case 3:
-        switch (cols_per_block)
+        // Copy nnz to host, and record an event for the completed transfer if desired
+
+        memcpy_d2h(WP_CURRENT_CONTEXT, bsr_nnz, bsr_offsets + row_count, sizeof(int), stream);
+
+        if (bsr_nnz_event)
         {
-        case 1:
-            wp_launch_device(WP_CURRENT_CONTEXT, bsr_transpose_blocks, nnz,
-                             (d_nnz, block_size, BsrBlockTransposer<3, 1, T>{}, block_indices, transposed_indices,
-                              bsr_values, transposed_bsr_columns, transposed_bsr_values));
-            return;
-        case 2:
-            wp_launch_device(WP_CURRENT_CONTEXT, bsr_transpose_blocks, nnz,
-                             (d_nnz, block_size, BsrBlockTransposer<3, 2, T>{}, block_indices, transposed_indices,
-                              bsr_values, transposed_bsr_columns, transposed_bsr_values));
-            return;
-        case 3:
-            wp_launch_device(WP_CURRENT_CONTEXT, bsr_transpose_blocks, nnz,
-                             (d_nnz, block_size, BsrBlockTransposer<3, 3, T>{}, block_indices, transposed_indices,
-                              bsr_values, transposed_bsr_columns, transposed_bsr_values));
-            return;
+            cuda_event_record(bsr_nnz_event, stream);
         }
     }
 
-    wp_launch_device(WP_CURRENT_CONTEXT, bsr_transpose_blocks, nnz,
-                     (d_nnz, block_size, BsrBlockTransposer<-1, -1, T>{rows_per_block, cols_per_block}, block_indices,
-                      transposed_indices, bsr_values, transposed_bsr_columns, transposed_bsr_values));
-}
+    // Set column indices
+    wp_launch_device(WP_CURRENT_CONTEXT, bsr_set_column, nnz,
+                     (bsr_offsets + row_count, d_values.Alternate(),
+                      bsr_columns));
 
-template <typename T>
-void bsr_transpose_device(int rows_per_block, int cols_per_block, int row_count, int col_count, int nnz,
-                          const int* bsr_offsets, const int* bsr_columns, const T* bsr_values,
-                          int* transposed_bsr_offsets, int* transposed_bsr_columns, T* transposed_bsr_values)
-{
+    // Scan repeated block counts
+    if(return_summed_blocks)
+    {
+        size_t buff_size = 0;
+        check_cuda(
+            cub::DeviceScan::InclusiveSum(nullptr, buff_size, tpl_block_offsets, tpl_block_offsets, nnz, stream));
+        ScopedTemporary<> temp(context, buff_size);
+        check_cuda(cub::DeviceScan::InclusiveSum(temp.buffer(), buff_size, tpl_block_offsets, tpl_block_offsets, nnz,
+                                                 stream));
+    } else {
+        // free our temporary buffers
+        free_device(context, tpl_block_offsets);
+        free_device(context, tpl_block_indices);
+     }
+}
 
-    const int block_size = rows_per_block * cols_per_block;
 
+WP_API void bsr_transpose_device(int row_count, int col_count, int nnz,
+                          const int* bsr_offsets, const int* bsr_columns, 
+                          int* transposed_bsr_offsets, int* transposed_bsr_columns,
+                          int* src_block_indices)
+{
     void* context = cuda_context_get_current();
     ContextGuard guard(context);
 
     cudaStream_t stream = static_cast<cudaStream_t>(cuda_stream_get_current());
 
-    ScopedTemporary<int> block_indices(context, 2 * nnz);
     ScopedTemporary<BsrRowCol> combined_row_col(context, 2 * nnz);
 
-    cub::DoubleBuffer<int> d_keys(block_indices.buffer(), block_indices.buffer() + nnz);
+    cub::DoubleBuffer<int> d_keys(src_block_indices + nnz, src_block_indices);
     cub::DoubleBuffer<BsrRowCol> d_values(combined_row_col.buffer(), combined_row_col.buffer() + nnz);
 
     wp_launch_device(WP_CURRENT_CONTEXT, bsr_transpose_fill_row_col, nnz,
@@ -466,59 +411,21 @@ void bsr_transpose_device(int rows_per_block, int cols_per_block, int row_count,
         check_cuda(cub::DeviceRadixSort::SortPairs(nullptr, buff_size, d_values, d_keys, nnz, 0, 64, stream));
         ScopedTemporary<> temp(context, buff_size);
         check_cuda(cub::DeviceRadixSort::SortPairs(temp.buffer(), buff_size, d_values, d_keys, nnz, 0, 64, stream));
+
+        // Depending on data size and GPU architecture buffers may have been swapped or not
+        // Ensures the sorted keys are available in summed_block_indices if needed
+        if(d_keys.Current() != src_block_indices)
+        {
+            check_cuda(cudaMemcpy(src_block_indices, src_block_indices+nnz, size_t(nnz) * sizeof(int), cudaMemcpyDeviceToDevice));
+        }
     }
 
     // Compute row offsets from sorted unique blocks
     wp_launch_device(WP_CURRENT_CONTEXT, bsr_find_row_offsets, col_count + 1,
                      (col_count, bsr_offsets + row_count, d_values.Current(), transposed_bsr_offsets));
 
-    // Move and transpose individual blocks
-    if (transposed_bsr_values != nullptr)
-    {
-        launch_bsr_transpose_blocks(nnz, bsr_offsets + row_count, block_size, rows_per_block, cols_per_block,
-                                    d_keys.Current(), d_values.Current(), bsr_values, transposed_bsr_columns,
-                                    transposed_bsr_values);
-    }
-}
-
-} // namespace
 
-void bsr_matrix_from_triplets_float_device(int rows_per_block, int cols_per_block, int row_count, int nnz,
-                                           int* tpl_rows, int* tpl_columns, void* tpl_values,
-                                           bool prune_numerical_zeros, bool masked, int* bsr_offsets, int* bsr_columns,
-                                           void* bsr_values, int* bsr_nnz, void* bsr_nnz_event)
-{
-    return bsr_matrix_from_triplets_device<float>(rows_per_block, cols_per_block, row_count, nnz, tpl_rows, tpl_columns,
-                                                  static_cast<const float*>(tpl_values), prune_numerical_zeros, masked,
-                                                  bsr_offsets, bsr_columns, static_cast<float*>(bsr_values), bsr_nnz,
-                                                  bsr_nnz_event);
-}
-
-void bsr_matrix_from_triplets_double_device(int rows_per_block, int cols_per_block, int row_count, int nnz,
-                                            int* tpl_rows, int* tpl_columns, void* tpl_values,
-                                            bool prune_numerical_zeros, bool masked, int* bsr_offsets, int* bsr_columns,
-                                            void* bsr_values, int* bsr_nnz, void* bsr_nnz_event)
-{
-    return bsr_matrix_from_triplets_device<double>(rows_per_block, cols_per_block, row_count, nnz, tpl_rows,
-                                                   tpl_columns, static_cast<const double*>(tpl_values),
-                                                   prune_numerical_zeros, masked, bsr_offsets, bsr_columns,
-                                                   static_cast<double*>(bsr_values), bsr_nnz, bsr_nnz_event);
-}
-
-void bsr_transpose_float_device(int rows_per_block, int cols_per_block, int row_count, int col_count, int nnz,
-                                int* bsr_offsets, int* bsr_columns, void* bsr_values, int* transposed_bsr_offsets,
-                                int* transposed_bsr_columns, void* transposed_bsr_values)
-{
-    bsr_transpose_device(rows_per_block, cols_per_block, row_count, col_count, nnz, bsr_offsets, bsr_columns,
-                         static_cast<const float*>(bsr_values), transposed_bsr_offsets, transposed_bsr_columns,
-                         static_cast<float*>(transposed_bsr_values));
-}
-
-void bsr_transpose_double_device(int rows_per_block, int cols_per_block, int row_count, int col_count, int nnz,
-                                 int* bsr_offsets, int* bsr_columns, void* bsr_values, int* transposed_bsr_offsets,
-                                 int* transposed_bsr_columns, void* transposed_bsr_values)
-{
-    bsr_transpose_device(rows_per_block, cols_per_block, row_count, col_count, nnz, bsr_offsets, bsr_columns,
-                         static_cast<const double*>(bsr_values), transposed_bsr_offsets, transposed_bsr_columns,
-                         static_cast<double*>(transposed_bsr_values));
+    wp_launch_device(WP_CURRENT_CONTEXT, bsr_set_column, nnz,
+                     (bsr_offsets + row_count, d_values.Current(),
+                      transposed_bsr_columns));
 }