cumsum.hip 866 Bytes
Newer Older
sangwzh's avatar
sangwzh committed
1
// !!! This is a file automatically generated by hipify!!!
2
3
4
5
6
7
/**
 *  Copyright (c) 2023 by Contributors
 *  Copyright (c) 2023, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
 * @file cuda/cumsum.cu
 * @brief Cumsum operators implementation on CUDA.
 */
sangwzh's avatar
sangwzh committed
8
#include <hipcub/hipcub.hpp>
9

sangwzh's avatar
sangwzh committed
10
#include "common.h"
11
12
13
14
15
16
17

namespace graphbolt {
namespace ops {

torch::Tensor ExclusiveCumSum(torch::Tensor input) {
  auto result = torch::empty_like(input);

18
19
20
21
22
23
  AT_DISPATCH_INTEGRAL_TYPES(input.scalar_type(), "ExclusiveCumSum", ([&] {
                               CUB_CALL(
                                   DeviceScan::ExclusiveSum,
                                   input.data_ptr<scalar_t>(),
                                   result.data_ptr<scalar_t>(), input.size(0));
                             }));
24
25
26
27
28
  return result;
}

}  // namespace ops
}  // namespace graphbolt