utils.hip 1.08 KB
Newer Older
sangwzh's avatar
sangwzh committed
1
2
// !!! This is a file automatically generated by hipify!!!
#include "hip/hip_runtime.h"
3
/**
4
 *  Copyright (c) 2020 by Contributors
5
6
 * @file array/cuda/utils.cu
 * @brief Utilities for CUDA kernels.
7
8
 */

sangwzh's avatar
sangwzh committed
9
#include <hipcub/hipcub.hpp>
10

11
#include "../../runtime/cuda/cuda_common.h"
sangwzh's avatar
sangwzh committed
12
#include "utils.h"
13
14
15
16

namespace dgl {
namespace cuda {

17
bool AllTrue(int8_t* flags, int64_t length, const DGLContext& ctx) {
18
19
20
21
  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;
sangwzh's avatar
sangwzh committed
22
23
  hipStream_t stream = runtime::getCurrentHIPStreamMasqueradingAsCUDA();
  CUDA_CALL(hipcub::DeviceReduce::Min(
24
      nullptr, workspace_size, flags, rst, length, stream));
25
  void* workspace = device->AllocWorkspace(ctx, workspace_size);
sangwzh's avatar
sangwzh committed
26
  CUDA_CALL(hipcub::DeviceReduce::Min(
27
      workspace, workspace_size, flags, rst, length, stream));
28
  int8_t cpu_rst = GetCUDAScalar(device, ctx, rst);
29
30
31
32
33
34
35
  device->FreeWorkspace(ctx, workspace);
  device->FreeWorkspace(ctx, rst);
  return cpu_rst == 1;
}

}  // namespace cuda
}  // namespace dgl