Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 12 additions & 1 deletion include/af/array.h
Original file line number Diff line number Diff line change
Expand Up @@ -655,11 +655,19 @@ namespace af

/**
Copy array data to host and return host pointer

For a CSR or COO sparse array the dense representation is copied,
so the returned buffer holds dims().elements() values. CSC sparse
arrays are not supported.
*/
template<typename T> T* host() const;

/**
Copy array data to existing host pointer

For a CSR or COO sparse array the dense representation is copied,
so \p ptr must hold dims().elements() values. CSC sparse arrays
are not supported.
*/
void host(void *ptr) const;

Expand Down Expand Up @@ -1684,7 +1692,10 @@ extern "C" {
/**
Copy data from an af_array to a C pointer.

Needs to used in conjunction with the two functions above
Needs to used in conjunction with the two functions above. For a CSR
or COO sparse array the dense representation is copied, so \p data
must hold af_get_elements() values. CSC sparse arrays are not
supported.
*/
AFAPI af_err af_get_data_ptr(void *data, const af_array arr);

Expand Down
57 changes: 42 additions & 15 deletions src/api/c/array.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -36,25 +36,52 @@ using detail::uint;
using detail::uintl;
using detail::ushort;

namespace {
// Releases a temporary af_array when it goes out of scope
struct ReleaseOnExit {
af_array arr;
~ReleaseOnExit() {
if (arr) { af_release_array(arr); }
}
};
} // namespace

af_err af_get_data_ptr(void *data, const af_array arr) {
try {
af_dtype type = getInfo(arr).getType();
const ArrayInfo &info = getInfo(arr, false);
af_dtype type = info.getType();

// A sparse array reports the dense shape through dims() and
// elements(), which is what callers size their buffer from, so hand
// back the dense data rather than rejecting the array.
ReleaseOnExit dense{nullptr};
if (info.isSparse()) {
if (getSparseArrayBase(arr).getStorage() == AF_STORAGE_CSC) {
AF_ERROR(
"Copying a CSC sparse array to the host is not supported; "
"convert it to CSR or COO first",
AF_ERR_NOT_SUPPORTED);
}
AF_CHECK(af_sparse_to_dense(&dense.arr, arr));
}
const af_array src = dense.arr ? dense.arr : arr;

// clang-format off
switch (type) {
case f32: copyData(static_cast<float* >(data), arr); break;
case c32: copyData(static_cast<cfloat* >(data), arr); break;
case f64: copyData(static_cast<double* >(data), arr); break;
case c64: copyData(static_cast<cdouble* >(data), arr); break;
case b8: copyData(static_cast<char* >(data), arr); break;
case s32: copyData(static_cast<int* >(data), arr); break;
case u32: copyData(static_cast<unsigned*>(data), arr); break;
case s8: copyData(static_cast<schar* >(data), arr); break;
case u8: copyData(static_cast<uchar* >(data), arr); break;
case s64: copyData(static_cast<intl* >(data), arr); break;
case u64: copyData(static_cast<uintl* >(data), arr); break;
case s16: copyData(static_cast<short* >(data), arr); break;
case u16: copyData(static_cast<ushort* >(data), arr); break;
case f16: copyData(static_cast<half* >(data), arr); break;
case f32: copyData(static_cast<float* >(data), src); break;
case c32: copyData(static_cast<cfloat* >(data), src); break;
case f64: copyData(static_cast<double* >(data), src); break;
case c64: copyData(static_cast<cdouble* >(data), src); break;
case b8: copyData(static_cast<char* >(data), src); break;
case s32: copyData(static_cast<int* >(data), src); break;
case u32: copyData(static_cast<unsigned*>(data), src); break;
case s8: copyData(static_cast<schar* >(data), src); break;
case u8: copyData(static_cast<uchar* >(data), src); break;
case s64: copyData(static_cast<intl* >(data), src); break;
case u64: copyData(static_cast<uintl* >(data), src); break;
case s16: copyData(static_cast<short* >(data), src); break;
case u16: copyData(static_cast<ushort* >(data), src); break;
case f16: copyData(static_cast<half* >(data), src); break;
default: TYPE_ERROR(1, type);
}
// clang-format on
Expand Down
57 changes: 57 additions & 0 deletions test/sparse.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -297,6 +297,63 @@ TYPED_TEST(Sparse, EmptyDeepCopy) {
EXPECT_EQ(0, sparseGetNNZ(b));
}

TEST(Sparse, HostOfSparseReturnsDense_ISSUE_3703) {
SUPPORTED_TYPE_CHECK(double);
const int rows = 8, cols = 4;
array values = af::constant(1.0, rows, f64);
array row_ptr = af::iota(dim4(rows + 1), dim4(1), s32);
array col_idx = af::constant(0, rows, s32);
array sp = af::sparse(rows, cols, values, row_ptr, col_idx,
AF_STORAGE_CSR);
ASSERT_TRUE(sp.issparse());

// One non-zero per row, all in column 0
vector<double> gold(rows * cols, 0.0);
for (int i = 0; i < rows; i++) { gold[i] = 1.0; }

double *h = sp.host<double>();
ASSERT_NE(h, nullptr);
for (int i = 0; i < rows * cols; i++) {
ASSERT_EQ(gold[i], h[i]) << "at " << i;
}
af::freeHost(h);

vector<double> into(rows * cols, -1.0);
sp.host(into.data());
ASSERT_EQ(gold, into);
}

TEST(Sparse, HostOfSparseMatchesDenseForEveryStorage_ISSUE_3703) {
const int rows = 6, cols = 5;
array dense = randu(rows, cols, f32);
// Sparsify roughly half of it
dense(dense < 0.5f) = 0.f;
vector<float> gold(rows * cols);
dense.host(gold.data());

const af_storage storages[] = {AF_STORAGE_CSR, AF_STORAGE_COO};
for (af_storage storage : storages) {
array sp = af::sparse(dense, storage);
ASSERT_TRUE(sp.issparse());
vector<float> got(rows * cols, -1.f);
sp.host(got.data());
ASSERT_EQ(gold, got) << "storage " << static_cast<int>(storage);
}

// CSC cannot be created from dense; the CSC of A is the CSR of A^T with
// the index arrays reinterpreted. CSC to dense is not implemented, so
// host() must reject it clearly rather than copy garbage.
array csrT = af::sparse(dense.T(), AF_STORAGE_CSR);
// For CSC the rowIdx argument holds the nnz row indices and the colIdx
// argument holds the cols + 1 column pointers
array csc = af::sparse(rows, cols, sparseGetValues(csrT),
sparseGetColIdx(csrT), sparseGetRowIdx(csrT),
AF_STORAGE_CSC);
ASSERT_TRUE(csc.issparse());
vector<float> gotCsc(rows * cols, -1.f);
ASSERT_THROW(csc.host(gotCsc.data()), af::exception);
}

TEST(Sparse, CPPSparseFromHostArrays) {
//! [ex_sparse_host_arrays]

Expand Down
Loading