"tests/vscode:/vscode.git/clone" did not exist on "9731e023255559b1c7537aad60b0d534738add85"
csr_transpose.cc 3.39 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/csr_transpose.cc
 * @brief CSR transpose (convert to CSC)
7
8
 */
#include <dgl/array.h>
9

10
11
12
13
14
15
16
17
18
#include "../../runtime/cuda/cuda_common.h"

namespace dgl {

using runtime::NDArray;

namespace aten {
namespace impl {

19
template <DGLDeviceType XPU, typename IdType>
20
CSRMatrix CSRTranspose(CSRMatrix csr) {
21
22
23
24
25
  LOG(FATAL) << "Unreachable codes";
  return {};
}

template <>
26
CSRMatrix CSRTranspose<kDGLCUDA, int32_t>(CSRMatrix csr) {
27
  auto* thr_entry = runtime::CUDAThreadEntry::ThreadLocal();
sangwzh's avatar
sangwzh committed
28
  hipStream_t stream = runtime::getCurrentHIPStreamMasqueradingAsCUDA();
29
30
  // allocate cusparse handle if needed
  if (!thr_entry->cusparse_handle) {
sangwzh's avatar
sangwzh committed
31
    CUSPARSE_CALL(hipsparseCreate(&(thr_entry->cusparse_handle)));
32
  }
sangwzh's avatar
sangwzh committed
33
  CUSPARSE_CALL(hipsparseSetStream(thr_entry->cusparse_handle, stream));
34
35
36
37
38

  NDArray indptr = csr.indptr, indices = csr.indices, data = csr.data;
  const int64_t nnz = indices->shape[0];
  const auto& ctx = indptr->ctx;
  const auto bits = indptr->dtype.bits;
39
  if (aten::IsNullArray(data)) data = aten::Range(0, nnz, bits, ctx);
40
41
42
43
  const int32_t* indptr_ptr = static_cast<int32_t*>(indptr->data);
  const int32_t* indices_ptr = static_cast<int32_t*>(indices->data);
  const void* data_ptr = data->data;

44
45
  // (BarclayII) csr2csc doesn't seem to clear the content of cscColPtr if nnz
  // == 0. We need to do it ourselves.
46
  NDArray t_indptr = aten::Full(0, csr.num_cols + 1, bits, ctx);
47
48
49
50
51
52
  NDArray t_indices = aten::NewIdArray(nnz, ctx, bits);
  NDArray t_data = aten::NewIdArray(nnz, ctx, bits);
  int32_t* t_indptr_ptr = static_cast<int32_t*>(t_indptr->data);
  int32_t* t_indices_ptr = static_cast<int32_t*>(t_indices->data);
  void* t_data_ptr = t_data->data;

sangwzh's avatar
sangwzh committed
53
#if DTKRT_VERSION >= 10010
54
55
56
  auto device = runtime::DeviceAPI::Get(csr.indptr->ctx);
  // workspace
  size_t workspace_size;
sangwzh's avatar
sangwzh committed
57
  CUSPARSE_CALL(hipsparseCsr2cscEx2_bufferSize(
58
59
      thr_entry->cusparse_handle, csr.num_rows, csr.num_cols, nnz, data_ptr,
      indptr_ptr, indices_ptr, t_data_ptr, t_indptr_ptr, t_indices_ptr,
sangwzh's avatar
sangwzh committed
60
61
      HIP_R_32F, HIPSPARSE_ACTION_NUMERIC, HIPSPARSE_INDEX_BASE_ZERO,
      HIPSPARSE_CSR2CSC_ALG1,  // see cusparse doc for reference
62
63
      &workspace_size));
  void* workspace = device->AllocWorkspace(ctx, workspace_size);
sangwzh's avatar
sangwzh committed
64
  CUSPARSE_CALL(hipsparseCsr2cscEx2(
65
66
      thr_entry->cusparse_handle, csr.num_rows, csr.num_cols, nnz, data_ptr,
      indptr_ptr, indices_ptr, t_data_ptr, t_indptr_ptr, t_indices_ptr,
sangwzh's avatar
sangwzh committed
67
68
      HIP_R_32F, HIPSPARSE_ACTION_NUMERIC, HIPSPARSE_INDEX_BASE_ZERO,
      HIPSPARSE_CSR2CSC_ALG1,  // see cusparse doc for reference
69
70
71
      workspace));
  device->FreeWorkspace(ctx, workspace);
#else
sangwzh's avatar
sangwzh committed
72
  CUSPARSE_CALL(hipsparseScsr2csc(
73
      thr_entry->cusparse_handle, csr.num_rows, csr.num_cols, nnz,
74
75
      static_cast<const float*>(data_ptr), indptr_ptr, indices_ptr,
      static_cast<float*>(t_data_ptr), t_indices_ptr, t_indptr_ptr,
sangwzh's avatar
sangwzh committed
76
      HIPSPARSE_ACTION_NUMERIC, HIPSPARSE_INDEX_BASE_ZERO));
77
78
#endif

79
80
  return CSRMatrix(
      csr.num_cols, csr.num_rows, t_indptr, t_indices, t_data, false);
81
82
}

83
template <>
84
CSRMatrix CSRTranspose<kDGLCUDA, int64_t>(CSRMatrix csr) {
85
86
87
  return COOToCSR(COOTranspose(CSRToCOO(csr, false)));
}

88
89
template CSRMatrix CSRTranspose<kDGLCUDA, int32_t>(CSRMatrix csr);
template CSRMatrix CSRTranspose<kDGLCUDA, int64_t>(CSRMatrix csr);
90
91
92
93

}  // namespace impl
}  // namespace aten
}  // namespace dgl