Implement sparsity support in existing operations

Tasks have to assemble the sparsity of their input nodes into a SparsityInfo
and use that to create their result node. A SparsityManager is used
to populate the sparsity information assembled on the host to the device
tensor.

Host tensors manage sparsity using a range-based encoding. TensorViews
use an array-based representation where the sparsity information is encoded
in the same memory as the tensor payload. That way we can transfer it directly
between processes.

Everything that can be sparse derives from one of the sparsity base classes
(tensor, tensorviews) and coordinate with the base class the layout of the
sparsity information. The sparsity information of Tensors, TensorViews, and
function nodes can be accessed using the `sparsity()` member call.

This makes heavy use of the conceptification of the API to pun away
the sparsity base classes. TensorView should not be used directly.
DenseTensorView is an alias for a dense TensorView. Similar aliases
exist for sparse tensors. Where possible, concepts should be used.
Also tightens the rules on const-ness in TensorView so that
we cannot assign from a `const TensorView` to a non-const TensorView.
Instead, we have to use a `TensorView<const T>` to adhere to constness.

Signed-off-by: Joseph Schuchart <joseph.schuchart@stonybrook.edu>

Author: devreal

include/mra/kernels/compress.h

@@ -70,24 +70,24 @@ namespace mra {
       T* d_sumsq,
       const concepts::TensorViewArray<NDIM+1, Key<NDIM>::num_children()> auto in_views)
     {
-      const bool is_t0 = (0 == thread_id());
       const size_type K2NDIM    = std::pow(  K,NDIM);
       const size_type TWOK2NDIM = std::pow(2*K,NDIM);
-      using tensorview_t = decltype(p_in(0));
       SHARED std::array<decltype(in_views[0](0)), Key<NDIM>::num_children()> block_in_views;
       SHARED T* workspace;
-      SHARED tensorview_t s, p, d;
+      SHARED DenseTensorView<T, NDIM> s, p, d;
       int blockId = blockIdx.x;
       T* block_tmp = &tmp[blockId*compress_tmp_size<NDIM>(K)];
 
-      if (is_t0) {
-        s = tensorview_t(&block_tmp[0], 2*K);
+      if (is_team_lead()) {
+        s = DenseTensorView<T, NDIM>(&block_tmp[0], 2*K);
         workspace = &block_tmp[TWOK2NDIM];
       }
-
+      assert(result_in.is_any_nonzero() && "why did we even get here?!");
       for (size_type fnid = blockId; fnid < N; fnid += gridDim.x) {
-        /* no need to sync threads here */
-        if (is_t0) {
+        if (result_in.is_zero(fnid)) {
+          continue; // output is zero so skip computation and leave it zero
+        }
+        if (is_team_lead()) {
           for (int i = 0; i < Key<NDIM>::num_children(); ++i) {
             block_in_views[i] = in_views[i](fnid);
           }

include/mra/kernels/derivative.h

@@ -267,7 +267,7 @@ namespace mra {
         size_type blockId = blockIdx.x;
         T* block_tmp_ptr = &tmp[blockId*derivative_tmp_size<NDIM>(K)];
         const size_type K2NDIM = std::pow(K, NDIM);
-        if(is_team_lead()){
+        if (is_team_lead()) {
           tmp_result = DenseTensorView<T, NDIM+1>(&block_tmp_ptr[       0], make_dims<NDIM+1>(2, K));
           left_tmp   = DenseTensorView<T, NDIM>(&block_tmp_ptr[2*K2NDIM], K);
           center_tmp = DenseTensorView<T, NDIM>(&block_tmp_ptr[3*K2NDIM], K);
@@ -318,8 +318,13 @@ namespace mra {
       static_assert(node_center.ndim() == NDIM+1, "node_center must be of dimension NDIM+1");
       static_assert(node_right.ndim() == NDIM+1, "node_right must be of dimension NDIM+1");
 
-      SHARED DenseTensorView<T, NDIM> node_left_view, node_center_view, node_right_view, deriv_view;
+      SHARED DenseTensorView<T, NDIM> deriv_view;
+      SHARED DenseTensorView<const T, NDIM> node_left_view, node_center_view, node_right_view;
       for (size_type blockid = blockIdx.x; blockid < N; blockid += gridDim.x) {
+        if (deriv.is_zero(blockid)) {
+          /* nothing to do */
+          continue;
+        }
         if (is_team_lead()) {
           node_left_view = node_left(blockid);
           node_center_view = node_center(blockid);

include/mra/kernels/fcoeffs.h

@@ -140,6 +140,10 @@ namespace mra {
 
       /* adjust pointers for the function of each block */
       for (size_type fnid = blockIdx.x; fnid < N; fnid += gridDim.x) {
+        if (coeffs_view.is_zero(fnid)) {
+          if (is_team_lead()) is_leaf[fnid] = false;
+          continue; // skip sparse entries
+        }
         if (is_team_lead()) {
           /* get the coefficient inputs */
           coeffs       = coeffs_view(fnid);

include/mra/kernels/gaxpy.h

@@ -33,8 +33,13 @@ namespace mra {
       const T scalarB,
       size_type N)
     {
-      SHARED DenseTensorView<T, NDIM> nodeA, nodeB, nodeR;
+      SHARED DenseTensorView<const T, NDIM> nodeA, nodeB;
+      SHARED DenseTensorView<T, NDIM> nodeR;
       for (size_type blockid = blockIdx.x; blockid < N; blockid += gridDim.x) {
+        if (nodeR_view.is_zero(blockid)) {
+          /* no work to do */
+          continue;
+        }
         if (is_team_lead()) {
           nodeA = nodeA_view(blockid);
           nodeB = nodeB_view(blockid);

include/mra/kernels/multiply.h

@@ -104,7 +104,8 @@ namespace mra {
       size_type N,
       size_type K)
     {
-      SHARED DenseTensorView<T, NDIM> nodeA, nodeB, nodeR, cnodesR, cnodesD;
+      SHARED DenseTensorView<const T, NDIM> nodeA, nodeB;
+      SHARED DenseTensorView<T, NDIM> nodeR, cnodesR, cnodesD;
       SHARED DenseTensorView<T, NDIM+1> cnodesA, cnodesB, r1;
       SHARED T* workspace;
       size_type blockId = blockIdx.x;
@@ -122,6 +123,10 @@ namespace mra {
       }
 
       for (size_type fnid = blockId; fnid < N; fnid += gridDim.x){
+        if (nodeR_view.is_zero(fnid)) {
+          /* no work to do */
+          continue;
+        }
         if (is_team_lead()) {
           nodeA = nodeA_view(fnid);
           nodeB = nodeB_view(fnid);

include/mra/kernels/norm.h

@@ -37,12 +37,16 @@ namespace mra {
       size_type K,
       const Key<NDIM>& key)
     {
-      const bool is_t0 = (0 == thread_id());
       const size_type TWOK2NDIM = std::pow(2*K, NDIM);
-      SHARED DenseTensorView<T, NDIM> n;
+      SHARED DenseTensorView<const T, NDIM> n;
       SHARED std::array<T, Key<NDIM>::num_children()> block_child_norms;
       for (size_type blockid = blockIdx.x; blockid < N; blockid += gridDim.x) {
-        if (is_t0) {
+        if (node.is_zero(blockid)) {
+          /* no work to do */
+          if (is_team_lead()) result_norms[blockid] = T(0.0);
+          continue;
+        }
+        if (is_team_lead()) {
           n = node(blockid);
           for (size_type i = 0; i < Key<NDIM>::num_children(); ++i) {
             block_child_norms[i] = (child_norms[i] != nullptr) ? child_norms[i][blockid] : T(0.0);

include/mra/kernels/reconstruct.h

@@ -61,7 +61,7 @@ namespace mra {
       Key<NDIM> key,
       size_type N,
       size_type K,
-      concepts::TensorView<NDIM+1> auto node_view,
+      const concepts::TensorView<NDIM+1> auto node_view,
       T* tmp_ptr,
       const concepts::TensorView<2> auto hg,
       const concepts::TensorView<NDIM+1> auto from_parent_view,
@@ -73,7 +73,8 @@ namespace mra {
       SHARED std::array<decltype(r_arr[0](0)), Key<NDIM>::num_children()> block_r_arr;
       SHARED DenseTensorView<T, NDIM> s, tmp_node;
       SHARED T* workspace;
-      SHARED DenseTensorView<T, NDIM> node, from_parent;
+      SHARED DenseTensorView<const T, NDIM> node;
+      SHARED DenseTensorView<const T, NDIM> from_parent;
 
       size_type blockId = blockIdx.x;
       T* block_tmp_ptr = &tmp_ptr[blockId*reconstruct_tmp_size<NDIM>(K)];
@@ -84,7 +85,13 @@ namespace mra {
         workspace   = &block_tmp_ptr[2*TWOK2NDIM];
       }
 
+      assert(node_view.is_any_nonzero() || from_parent.is_any_nonzero() && "why did we even get here?!");
+
       for (size_type fnid = blockId; fnid < N; fnid += gridDim.x){
+        if (node.is_zero(fnid) && from_parent.is_zero(fnid)) {
+          /* no work to do */
+          continue;
+        }
         if (is_t0) {
           node = node_view(fnid);
           from_parent = from_parent_view(fnid);
@@ -103,7 +110,7 @@ namespace mra {
     const Key<NDIM>& key,
     size_type N,
     size_type K,
-    concepts::TensorView<NDIM+1> auto& node,
+    const concepts::TensorView<NDIM+1> auto& node,
     const concepts::TensorView<2> auto& hg,
     const concepts::TensorView<NDIM+1> auto& from_parent,
     const concepts::TensorViewArray<NDIM+1, mra::Key<NDIM>::num_children()> auto& r_arr,
@@ -125,7 +132,7 @@ namespace mra {
     const Key<3>& key,
     size_type N,
     size_type K,
-    SparseTensorView<double, 3+1>& node,
+    const SparseTensorView<double, 3+1>& node,
     const SparseTensorView<double, 2>& hg,
     const SparseTensorView<double, 3+1>& from_parent,
     const std::array<SparseTensorView<double, 3+1>, mra::Key<3>::num_children()>& r_arr,

include/mra/tasks/compress.h

@@ -9,6 +9,7 @@
 #include "mra/misc/options.h"
 #include "mra/misc/functiondata.h"
 #include "mra/misc/functionset.h"
+#include "mra/tensor/sparsitymanager.h"
 #include "mra/tensor/tensor.h"
 #include "mra/tensor/tensorview.h"
 #include "mra/tensor/functionnode.h"
@@ -87,22 +88,29 @@ namespace mra
 
           /* some inputs are on the device so submit a kernel */
 
+          SparsityInfo sparsity(N);
+          sparsity.nonzero_if_any(in0, in1, in2, in3, in4, in5, in6, in7);
+          //std::cout << name << " " << key << " sparsity: " << sparsity << std::endl;
+
           // allocate the result
-          result = mra::FunctionsCompressedNode<T, NDIM>(key, N, K, ttg::scope::Allocate);
-          auto& d = result.coeffs();
+          result.allocate(sparsity, K, ttg::scope::Allocate);
+
           // Collect child leaf info
           mra::apply_leaf_info(result, in0, in1, in2, in3, in4, in5, in6, in7);
-          p = mra::FunctionsReconstructedNode<T, NDIM>(key, N, K, ttg::scope::Allocate);
+          p.allocate(sparsity, K, ttg::scope::Allocate);
           p.set_all_leaf(false);
           assert(p.is_all_leaf() == false);
           FunctionNorms<T, NDIM> norms(name, in0, in1, in2, in3, in4, in5, in6, in7, result);
 
+
           const std::size_t tmp_size = compress_tmp_size<NDIM>(K)*N;
           ttg::Buffer<T, DeviceAllocator<T>> tmp_scratch(tmp_size, TempScope);
           const auto& hgT = functiondata.get_hgT();
           /* stores sumsq for each child and for result at the end of the kernel */
           auto d_sumsq = ttg::Buffer<T, DeviceAllocator<T>>(N, TempScope);
 
+          auto& d = result.coeffs();
+
 #ifndef MRA_ENABLE_HOST
           auto input = ttg::device::Input(p.coeffs().buffer(), d.buffer(), hgT.buffer(),
                                           tmp_scratch, d_sumsq);
@@ -131,6 +139,9 @@ namespace mra
           auto input_views = std::array{in0.coeffs().current_view(), in1.coeffs().current_view(), in2.coeffs().current_view(), in3.coeffs().current_view(),
                                         in4.coeffs().current_view(), in5.coeffs().current_view(), in6.coeffs().current_view(), in7.coeffs().current_view()};
 
+          auto sparseman = make_sparsity_manager(d, p);
+          sparseman.populate_device_sparsity();
+
           auto coeffs_view = p.coeffs().current_view();
           auto rcoeffs_view = d.current_view();
           auto hgT_view = hgT.current_view();
@@ -176,11 +187,11 @@ namespace mra
           for (std::size_t i = 0; i < N; ++i) {
             if (std::abs(p.sum(i) - 1.0) > 1e-12) {
               all_correct = false;
-              std::cout << "At root of compressed tree " << key.batch() << " fn " << i << ": total normsq is " << p.sum(i) << std::endl;
+              std::cout << name << ": at root of compressed tree " << key.batch() << " fn " << i << ": total normsq is " << p.sum(i) << std::endl;
             }
           }
           if (all_correct) {
-            std::cout << "At root of compressed tree " << key.batch() << ": all norms are 1.0 with 1e-12 tolerance" << std::endl;
+            std::cout << name << ": at root of compressed tree " << key.batch() << ": all norms are 1.0 with 1e-12 tolerance" << std::endl;
           }
 #ifndef MRA_ENABLE_HOST
           co_await ttg::device::forward(

include/mra/tasks/derivative.h

@@ -232,7 +232,11 @@ namespace mra{
            * We can finally compute the derivative.
            */
           if ((!left.empty() || key.is_left_boundary(axis)) && (!right.empty() || key.is_right_boundary(axis))){
-            mra::FunctionsReconstructedNode<T, NDIM> result(key, N, K, ttg::scope::Allocate);
+
+            SparsityInfo sparsity(N);
+            sparsity.nonzero_if_any(left, center, right);
+
+            mra::FunctionsReconstructedNode<T, NDIM> result(key, sparsity, K, ttg::scope::Allocate);
             result.set_all_leaf(true);
             auto tmp = ttg::Buffer<T>(derivative_tmp_size<NDIM>(K)*N, TempScope);
             const DenseTensor<T, 2+1>& operators = functiondata.get_operators();
@@ -261,6 +265,9 @@ namespace mra{
             co_await ttg::device::select(input);
 #endif // MRA_ENABLE_HOST
 
+            SparsityManager sparseman(result);
+            sparseman.populate_device_sparsity();
+
             auto& D = *db.current_device_ptr();
             auto result_view = result.coeffs().current_view();
             submit_derivative_kernel(D, key, left.key(), center.key(), right.key(), left.coeffs().current_view(),

include/mra/tasks/gaxpy.h

@@ -73,20 +73,16 @@ namespace mra{
         send_out(t1);
       } else {
 
-        auto out = mra::FunctionsCompressedNode<T, NDIM>(key, N, K, ttg::scope::Allocate);
+
+        SparsityInfo sparsity(N);
+        sparsity.nonzero_if_all(t1, t2);
+
+        auto out = mra::FunctionsCompressedNode<T, NDIM>(key, sparsity, K, ttg::scope::Allocate);
 
         /* adapt the leaf information of the result: if the children of both nodes are leafs then
         * the children of the output node are leafs as well. */
         mra::apply_leaf_info(out, t1, t2);
         //std::cout << name << " " << key << " all leafs " << out.is_all_child_leaf() << std::endl;
-#if 0
-        for (size_type i = 0; i < N; ++i) {
-          for (auto child : children(key)) {
-            auto childidx = child.childindex();
-            out.is_child_leaf(i)[childidx] = t1.is_child_leaf(i)[childidx] && t2.is_child_leaf(i)[childidx];
-          }
-        }
-#endif // 0
 
         auto norms = FunctionNorms(name, out, t1, t2);
 
@@ -105,6 +101,9 @@ namespace mra{
         auto t2_view = t2.coeffs().current_view();
         auto out_view = out.coeffs().current_view();
 
+        auto sparseman = make_sparsity_manager(out);
+        sparseman.populate_device_sparsity();
+
         submit_gaxpy_kernel(key, t1_view, t2_view, out_view,
                             scalarA, scalarB, N, K, ttg::device::current_stream());