add_dense_sparse_worker_non_hybrid_cpu Class — pytorch Architecture

Source Code

aten/src/ATen/native/sparse/SparseTensorMath.cpp lines 595–616

template <typename scalar_t>
static void add_dense_sparse_worker_non_hybrid_cpu(Tensor& r, const Scalar& value, const SparseTensor& sparse, const Tensor& indices, const Tensor& values) {
  auto indices_accessor = indices.accessor<int64_t, 2>();
  auto values_accessor = values.accessor<scalar_t, 1>();

  scalar_t* r_ptr = r.data_ptr<scalar_t>();
  scalar_t cast_value = value.to<scalar_t>();
  const int64_t sparse_dim = sparse.sparse_dim();
  std::vector<int64_t> result_stride(sparse_dim);
  for (const auto d: c10::irange(sparse_dim)) {
    result_stride[d] = r.stride(d);
  }
  at::parallel_for(0, sparse._nnz(), 0, [&](int64_t start, int64_t end) {
    for (const auto k: c10::irange(start, end)) {
      int64_t index = r.storage_offset();
      for (auto d: c10::irange(sparse_dim)) {
        index += result_stride[d] * indices_accessor[d][k];
      }
      r_ptr[index] += cast_value * values_accessor[k];
    }
  });
}