utils.cu 993 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 "../../runtime/cuda/cuda_common.h"
8
9
#include "./dgl_cub.cuh"
#include "./utils.h"
10
11
12
13

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
21
  CUDA_CALL(cub::DeviceReduce::Min(
      nullptr, workspace_size, flags, rst, length, stream));
22
  void* workspace = device->AllocWorkspace(ctx, workspace_size);
23
24
  CUDA_CALL(cub::DeviceReduce::Min(
      workspace, workspace_size, flags, rst, length, stream));
25
  int8_t cpu_rst = GetCUDAScalar(device, ctx, rst);
26
27
28
29
30
31
32
  device->FreeWorkspace(ctx, workspace);
  device->FreeWorkspace(ctx, rst);
  return cpu_rst == 1;
}

}  // namespace cuda
}  // namespace dgl