Skip to content

Commit

Permalink
Add "zero" matrix factor methods for COO,CSR,CSC (#870)
Browse files Browse the repository at this point in the history 
* Add "zero" matrix factor methods for COO,CSR,CSC

* more concise code

* added TODO on zero buffer init
  • Loading branch information
@aartbik
aartbik authored Feb 6, 2025
1 parent 4c29dee commit 9114af3
Show file tree
Hide file tree
Showing 2 changed files with 63 additions and 19 deletions.
76 changes: 61 additions & 15 deletions include/matx/core/make_sparse_tensor.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,8 @@
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice, this
// list of conditions and the following disclaimer.
// 1. Redistributions of source code must retain the above copyright notice,
// this list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright notice,
// this list of conditions and the following disclaimer in the documentation
Expand Down Expand Up @@ -37,9 +37,24 @@
namespace matx {
namespace experimental {

// Helper method to create zero storage.
template <typename T>
__MATX_INLINE__ static auto
makeDefaultNonOwningZeroStorage(index_t sz, matxMemorySpace_t space) {
T *ptr;
matxAlloc((void **)&ptr, sz * sizeof(T), space, 0);
// TODO: introduce a more efficient matxCalloc or matxMemset?
for (index_t i = 0; i < sz; i++) {
ptr[i] = 0;
}
raw_pointer_buffer<T, matx_allocator<T>> buf{ptr, sz * sizeof(T),
/*owning=*/false};
return basic_storage<decltype(buf)>{std::move(buf)};
}

// Helper method to create empty storage.
template <typename T>
__MATX_INLINE__ static auto makeDefaultNonOwningStorage() {
__MATX_INLINE__ static auto makeDefaultNonOwningEmptyStorage() {
raw_pointer_buffer<T, matx_allocator<T>> buf{nullptr, 0, /*owning=*/false};
return basic_storage<decltype(buf)>{std::move(buf)};
}
Expand Down Expand Up @@ -72,18 +87,25 @@ auto make_tensor_coo(ValTensor &val, CrdTensor &row, CrdTensor &col,
// However, under the formal DSL specifications, the top level
// compression should set up pos[0] = {0, nse}. This is done
// here, using the same memory space as the other data.
POS *ptr;
matxMemorySpace_t space = GetPointerKind(val.GetStorage().data());
matxAlloc((void **)&ptr, 2 * sizeof(POS), space, 0);
ptr[0] = 0;
ptr[1] = val.Size(0);
raw_pointer_buffer<POS, matx_allocator<POS>> topp{ptr, 2 * sizeof(POS),
/*owning=*/false};
basic_storage<decltype(topp)> tp{std::move(topp)};
auto tp = makeDefaultNonOwningZeroStorage<POS>(2, space);
tp.data()[1] = val.Size(0);
// Construct COO.
return sparse_tensor_t<VAL, CRD, POS, COO>(
shape, val.GetStorage(), {row.GetStorage(), col.GetStorage()},
{tp, makeDefaultNonOwningStorage<POS>()});
{tp, makeDefaultNonOwningEmptyStorage<POS>()});
}

// Constructs a zero sparse matrix in COO format (viz. nse=0).
template <typename VAL, typename CRD, typename POS = index_t>
auto make_zero_tensor_coo(const index_t (&shape)[2],
matxMemorySpace_t space = MATX_MANAGED_MEMORY) {
return sparse_tensor_t<VAL, CRD, POS, COO>(
shape, makeDefaultNonOwningEmptyStorage<VAL>(),
{makeDefaultNonOwningEmptyStorage<CRD>(),
makeDefaultNonOwningEmptyStorage<CRD>()},
{makeDefaultNonOwningZeroStorage<POS>(2, space),
makeDefaultNonOwningEmptyStorage<POS>()});
}

// Constructs a sparse matrix in CSR format directly from the values, the
Expand All @@ -106,8 +128,20 @@ auto make_tensor_csr(ValTensor &val, PosTensor &rowp, CrdTensor &col,
// Construct CSR.
return sparse_tensor_t<VAL, CRD, POS, CSR>(
shape, val.GetStorage(),
{makeDefaultNonOwningStorage<CRD>(), col.GetStorage()},
{makeDefaultNonOwningStorage<POS>(), rowp.GetStorage()});
{makeDefaultNonOwningEmptyStorage<CRD>(), col.GetStorage()},
{makeDefaultNonOwningEmptyStorage<POS>(), rowp.GetStorage()});
}

// Constructs a zero sparse matrix in CSR format (viz. nse=0).
template <typename VAL, typename CRD, typename POS>
auto make_zero_tensor_csr(const index_t (&shape)[2],
matxMemorySpace_t space = MATX_MANAGED_MEMORY) {
return sparse_tensor_t<VAL, CRD, POS, CSR>(
shape, makeDefaultNonOwningEmptyStorage<VAL>(),
{makeDefaultNonOwningEmptyStorage<CRD>(),
makeDefaultNonOwningEmptyStorage<CRD>()},
{makeDefaultNonOwningEmptyStorage<POS>(),
makeDefaultNonOwningZeroStorage<POS>(shape[0] + 1, space)});
}

// Constructs a sparse matrix in CSC format directly from the values, the
Expand All @@ -130,8 +164,20 @@ auto make_tensor_csc(ValTensor &val, PosTensor &colp, CrdTensor &row,
// Construct CSC.
return sparse_tensor_t<VAL, CRD, POS, CSC>(
shape, val.GetStorage(),
{makeDefaultNonOwningStorage<CRD>(), row.GetStorage()},
{makeDefaultNonOwningStorage<POS>(), colp.GetStorage()});
{makeDefaultNonOwningEmptyStorage<CRD>(), row.GetStorage()},
{makeDefaultNonOwningEmptyStorage<POS>(), colp.GetStorage()});
}

// Constructs a zero sparse matrix in CSC format (viz. nse=0).
template <typename VAL, typename CRD, typename POS>
auto make_zero_tensor_csc(const index_t (&shape)[2],
matxMemorySpace_t space = MATX_MANAGED_MEMORY) {
return sparse_tensor_t<VAL, CRD, POS, CSC>(
shape, makeDefaultNonOwningEmptyStorage<VAL>(),
{makeDefaultNonOwningEmptyStorage<CRD>(),
makeDefaultNonOwningEmptyStorage<CRD>()},
{makeDefaultNonOwningEmptyStorage<POS>(),
makeDefaultNonOwningZeroStorage<POS>(shape[1] + 1, space)});
}

} // namespace experimental
Expand Down
6 changes: 2 additions & 4 deletions include/matx/core/print.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -570,16 +570,14 @@ namespace matx {
fprintf(fp, "format = ");
Format::print();
for (int lvlIdx = 0; lvlIdx < Format::LVL; lvlIdx++) {
if (op.POSData(lvlIdx)) {
const index_t pend = op.posSize(lvlIdx);
if (const index_t pend = op.posSize(lvlIdx)) {
fprintf(fp, "pos[%d] = (", lvlIdx);
for (index_t i = 0; i < pend; i++) {
PrintVal(fp, op.POSData(lvlIdx)[i]);
}
fprintf(fp, ")\n");
}
if (op.CRDData(lvlIdx)) {
const index_t cend = op.crdSize(lvlIdx);
if (const index_t cend = op.crdSize(lvlIdx)) {
fprintf(fp, "crd[%d] = (", lvlIdx);
for (index_t i = 0; i < cend; i++) {
PrintVal(fp, op.CRDData(lvlIdx)[i]);
Expand Down

0 comments on commit 9114af3

Please sign in to comment.