csr_transpose.cc 3.28 KB
Newer Older
1
2
/*!
 *  Copyright (c) 2020 by Contributors
3
4
 * @file array/cuda/csr_transpose.cc
 * @brief CSR transpose (convert to CSC)
5
6
 */
#include <dgl/array.h>
7

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

namespace dgl {

using runtime::NDArray;

namespace aten {
namespace impl {

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

template <>
24
CSRMatrix CSRTranspose<kDGLCUDA, int32_t>(CSRMatrix csr) {
25
  auto* thr_entry = runtime::CUDAThreadEntry::ThreadLocal();
26
  cudaStream_t stream = runtime::getCurrentCUDAStream();
27
28
29
30
  // allocate cusparse handle if needed
  if (!thr_entry->cusparse_handle) {
    CUSPARSE_CALL(cusparseCreate(&(thr_entry->cusparse_handle)));
  }
31
  CUSPARSE_CALL(cusparseSetStream(thr_entry->cusparse_handle, stream));
32
33
34
35
36

  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;
37
  if (aten::IsNullArray(data)) data = aten::Range(0, nnz, bits, ctx);
38
39
40
41
  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;

42
43
  // (BarclayII) csr2csc doesn't seem to clear the content of cscColPtr if nnz
  // == 0. We need to do it ourselves.
44
  NDArray t_indptr = aten::Full(0, csr.num_cols + 1, bits, ctx);
45
46
47
48
49
50
  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;

51
#if CUDART_VERSION >= 10010
52
53
54
55
  auto device = runtime::DeviceAPI::Get(csr.indptr->ctx);
  // workspace
  size_t workspace_size;
  CUSPARSE_CALL(cusparseCsr2cscEx2_bufferSize(
56
57
58
      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,
      CUDA_R_32F, CUSPARSE_ACTION_NUMERIC, CUSPARSE_INDEX_BASE_ZERO,
59
60
61
62
      CUSPARSE_CSR2CSC_ALG1,  // see cusparse doc for reference
      &workspace_size));
  void* workspace = device->AllocWorkspace(ctx, workspace_size);
  CUSPARSE_CALL(cusparseCsr2cscEx2(
63
64
65
      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,
      CUDA_R_32F, CUSPARSE_ACTION_NUMERIC, CUSPARSE_INDEX_BASE_ZERO,
66
67
68
69
70
      CUSPARSE_CSR2CSC_ALG1,  // see cusparse doc for reference
      workspace));
  device->FreeWorkspace(ctx, workspace);
#else
  CUSPARSE_CALL(cusparseScsr2csc(
71
      thr_entry->cusparse_handle, csr.num_rows, csr.num_cols, nnz,
72
73
      static_cast<const float*>(data_ptr), indptr_ptr, indices_ptr,
      static_cast<float*>(t_data_ptr), t_indices_ptr, t_indptr_ptr,
74
      CUSPARSE_ACTION_NUMERIC, CUSPARSE_INDEX_BASE_ZERO));
75
76
#endif

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

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

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

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