     // Cast validity buffer to 4 byte type
     valid32_t const * masks32 = reinterpret_cast<valid32_t const *>(masks);
 
-    CUDA_TRY(cudaMalloc(&d_count, sizeof(int)));
-    CUDA_TRY(cudaMemsetAsync(d_count, 0, sizeof(int),count_stream));
+    RMM_TRY(rmmAlloc((void**)&d_count, sizeof(int), count_stream));
+    CUDA_TRY(cudaMemsetAsync(d_count, 0, sizeof(int), count_stream));
 
     const int grid_size = (num_masks32 + block_size - 1)/block_size;
 


    CUDA_TRY( cudaGetLastError() );

    CUDA_TRY(cudaMemcpyAsync(&h_count, d_count, sizeof(int), cudaMemcpyDeviceToHost,count_stream));
    CUDA_TRY(cudaStreamSynchronize(count_stream));
    CUDA_TRY(cudaStreamDestroy(count_stream));
  }

#include <gdf/errorutils.h>
#include <gdf/utils.h>
#include <vector>
#include <cassert>
#include <cub/cub.cuh>
#include "rmm.h"


using valid32_t = uint32_t;
