csr_transpose.cc 3.38 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
#if CUDART_VERSION < 12000
26
  auto* thr_entry = runtime::CUDAThreadEntry::ThreadLocal();
27
  cudaStream_t stream = runtime::getCurrentCUDAStream();
28
29
30
31
  // allocate cusparse handle if needed
  if (!thr_entry->cusparse_handle) {
    CUSPARSE_CALL(cusparseCreate(&(thr_entry->cusparse_handle)));
  }
32
  CUSPARSE_CALL(cusparseSetStream(thr_entry->cusparse_handle, stream));
33
34
35
36
37

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

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

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

78
79
  return CSRMatrix(
      csr.num_cols, csr.num_rows, t_indptr, t_indices, t_data, false);
80
81
82
#else
  return COOToCSR(COOTranspose(CSRToCOO(csr, false)));
#endif
83
84
}

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

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

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