convert_indices_from_csr_to_coo_cpu Class — pytorch Architecture

Source Code

aten/src/ATen/native/TensorConversions.cpp lines 1887–1933

template <typename input_t, typename output_t>
void convert_indices_from_csr_to_coo_cpu(
    const Tensor& indices,
    const Tensor& crow_indices,
    const Tensor& col_indices,
    const bool transpose = false) {
  int64_t nrows = crow_indices.size(-1) - 1;
  int64_t nnz = col_indices.size(-1);
  if (nrows == 0 || nnz == 0) {
    indices.zero_(); // is this needed as indices has a zero-valued
                     // dimension when nrows or nnz is 0?
    return;
  }
  auto crow_indices_ = crow_indices.expect_contiguous();
  int64_t total_nnz = col_indices.numel();
  int64_t batch_ndim = crow_indices.dim() - 1;
  if (batch_ndim > 0) {
    auto batch_indices = indices.narrow(0, 0, batch_ndim);
    batch_indices.copy_(
        at::sparse::full_coo_indices(
            crow_indices.sizes().slice(0, batch_ndim), crow_indices.options())
            .repeat_interleave(nnz, 1));
  }
  const input_t* crow_indices_data_in =
      crow_indices_->const_data_ptr<input_t>();
  TORCH_INTERNAL_ASSERT(indices.is_contiguous());
  auto row0 = indices.select(0, transpose ? batch_ndim + 1 : batch_ndim + 0);
  auto row1 = indices.select(0, transpose ? batch_ndim + 0 : batch_ndim + 1);
  output_t* data_out = row0.data_ptr<output_t>();
  auto col_indices_ = col_indices.expect_contiguous();
  row1.copy_(col_indices_->view({-1}));
  at::parallel_for(
      0,
      nrows * total_nnz / nnz,
      at::internal::GRAIN_SIZE,
      [&](int64_t start, int64_t end) {
        for (const auto i_ : c10::irange(start, end)) {
          auto b = i_ / nrows;
          auto i = i_ % nrows;
          std::fill(
              &data_out[b * nnz + crow_indices_data_in[b * (nrows + 1) + i]],
              &data_out
                  [b * nnz + crow_indices_data_in[b * (nrows + 1) + i + 1]],
              static_cast<output_t>(i));
        }
      });
}