utils.cu 979 Bytes
Newer Older
1
2
/*!
 *  Copyright (c) 2020 by Contributors
3
4
 * @file array/cuda/utils.cu
 * @brief Utilities for CUDA kernels.
5
6
7
 */

#include "./utils.h"
8
#include "./dgl_cub.cuh"
9
10
11
12
13
#include "../../runtime/cuda/cuda_common.h"

namespace dgl {
namespace cuda {

14
bool AllTrue(int8_t* flags, int64_t length, const DGLContext& ctx) {
15
16
17
18
  auto device = runtime::DeviceAPI::Get(ctx);
  int8_t* rst = static_cast<int8_t*>(device->AllocWorkspace(ctx, 1));
  // Call CUB's reduction
  size_t workspace_size = 0;
19
  cudaStream_t stream = runtime::getCurrentCUDAStream();
20
  CUDA_CALL(cub::DeviceReduce::Min(nullptr, workspace_size, flags, rst, length, stream));
21
  void* workspace = device->AllocWorkspace(ctx, workspace_size);
22
23
  CUDA_CALL(cub::DeviceReduce::Min(workspace, workspace_size, flags, rst, length, stream));
  int8_t cpu_rst = GetCUDAScalar(device, ctx, rst);
24
25
26
27
28
29
30
  device->FreeWorkspace(ctx, workspace);
  device->FreeWorkspace(ctx, rst);
  return cpu_rst == 1;
}

}  // namespace cuda
}  // namespace dgl