Refine print order and skip device contents (#908)

NVIDIA · Mar 11, 2025 · 6b20e04 · 6b20e04
1 parent 4445f72
commit 6b20e04
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 24 deletions.
diff --git 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
@@ -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<float>({5});
   auto idxi = make_tensor<int>({5}); 

diff --git 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<Op>) {
         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<Op>) {
         // If the user is printing a tensor with a const pointer underlying the data, we need to do the lookup