diff --git a/docs_input/basics/sparse_tensor.rst b/docs_input/basics/sparse_tensor.rst index 46d8dfd6..41de6e1c 100644 --- a/docs_input/basics/sparse_tensor.rst +++ b/docs_input/basics/sparse_tensor.rst @@ -57,13 +57,13 @@ its constituent buffers is constructed as follows:: The result of the print statement is shown below:: tensor_impl_2_f32: SparseTensor{float} Rank: 2, Sizes:[4, 8], Levels:[4, 8] - nse = 5 format = ( d0, d1 ) -> ( d0 : compressed(non-unique), d1 : singleton ) + space = CUDA managed memory + nse = 5 pos[0] = ( 0 5 ) crd[0] = ( 0 0 3 3 3 ) crd[1] = ( 0 1 2 3 5 ) values = ( 1.0000e+00 2.0000e+00 3.0000e+00 4.0000e+00 5.0000e+00 ) - space = CUDA managed memory Note that, like dense tensors, sparse tensors provide ()-operations for indexing. However, users should **never** use the ()-operator diff --git a/examples/sparse_tensor.cu b/examples/sparse_tensor.cu index b085d225..ec40385a 100644 --- a/examples/sparse_tensor.cu +++ b/examples/sparse_tensor.cu @@ -71,12 +71,12 @@ int main([[maybe_unused]] int argc, [[maybe_unused]] char **argv) // | 0, 0, 3, 4, 0, 5, 0, 0 | // // tensor_impl_2_f32: SparseTensor{float} Rank: 2, Sizes:[4, 8], Levels:[4, 8] - // nse = 5 // format = ( d0, d1 ) -> ( d0 : compressed(non-unique), d1 : singleton ) + // space = CUDA managed memory + // nse = 5 // crd[0] = ( 0 0 3 3 3 ) // crd[1] = ( 0 1 2 3 5 ) // values = ( 1.0000e+00 2.0000e+00 3.0000e+00 4.0000e+00 5.0000e+00 ) - // space = CUDA managed memory // auto vals = make_tensor({5}); auto idxi = make_tensor({5}); diff --git a/include/matx/core/print.h b/include/matx/core/print.h index b53f2408..b022d232 100644 --- a/include/matx/core/print.h +++ b/include/matx/core/print.h @@ -565,31 +565,34 @@ namespace matx { cudaDeviceSynchronize(); if constexpr (is_sparse_tensor_v) { using Format = typename Op::Format; - index_t nse = op.Nse(); - fprintf(fp, "nse = %" MATX_INDEX_T_FMT "\n", nse); fprintf(fp, "format = "); - Format::print(); - for (int lvlIdx = 0; lvlIdx < Format::LVL; 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]); + Format::print(); + const auto kind = GetPointerKind(op.Data()); + fprintf(fp, ")\nspace = %s\n", SpaceString(kind).c_str()); + const auto nse = op.Nse(); + fprintf(fp, "nse = %" MATX_INDEX_T_FMT "\n", nse); + if (HostPrintable(kind)) { + for (int lvlIdx = 0; lvlIdx < Format::LVL; 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"); } - fprintf(fp, ")\n"); - } - 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]); + 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]); + } + fprintf(fp, ")\n"); } - fprintf(fp, ")\n"); + } + fprintf(fp, "values = ("); + for (index_t i = 0; i < nse; i++) { + PrintVal(fp, op.Data()[i]); } } - fprintf(fp, "values = ("); - for (index_t i = 0; i < nse; i++) { - PrintVal(fp, op.Data()[i]); - } - fprintf(fp, ")\nspace = %s\n", SpaceString(GetPointerKind(op.Data())).c_str()); } else if constexpr (is_tensor_view_v) { // If the user is printing a tensor with a const pointer underlying the data, we need to do the lookup