First version of MATX dense2sparse conversion (dispatch to cuSPARSE) (#868)

Author: aartbik

examples/sparse_tensor.cu

@@ -149,5 +149,22 @@ int main([[maybe_unused]] int argc, [[maybe_unused]] char **argv)
   (X = solve(Acsr, Y)).run(exec);
   print(X);
 
+  //
+  // A direct dense2sparse conversion. This is the correct way of
+  // performing an efficient sparse operation. Note, however,
+  // that assigning a right-hand-side value to a sparse tensor
+  // (viz. the lval Acoo) is an experimental operation recently
+  // added to MatX, and it is currently restricted to a direct
+  // "dense2sparse" operation at the right-hand-side.
+  //
+  auto D = make_tensor<float, 2>({4, 8});
+  D.SetVals({
+    {0, 11,  0, 12,  0,  0,  0,  0},
+    {0,  0, 13,  0,  0,  0,  0,  0},
+    {0,  0,  0,  0,  0,  0,  0, 14},
+    {0, 15,  0,  0, 16,  0, 17,  0}});
+  (Acoo = dense2sparse(D)).run(exec);
+  print(Acoo);
+
   MATX_EXIT_HANDLER();
 }

include/matx/core/sparse_tensor.h

@@ -36,8 +36,33 @@
 
 #include "matx/core/sparse_tensor_format.h"
 #include "matx/core/tensor_impl.h"
+#include "matx/operators/base_operator.h"
 
 namespace matx {
+
+namespace detail {
+
+//
+// A sparse_set operation. Assigning to a sparse tensor is very different
+// from all other MatX assignments, because the underlying storage and
+// buffers may have to be resized to accomodate the output. Therefore,
+// for now, we provide a customized set operation that passes a direct
+// reference to the executor.
+//
+template <typename T, typename Op>
+class sparse_set : public BaseOp<sparse_set<T, Op>> {
+private:
+  T &out_;
+  mutable typename detail::base_type_t<Op> op_;
+public:
+  inline sparse_set(T &out, const Op &op) : out_(out), op_(op) {}
+  template <typename Ex> __MATX_INLINE__ void run(Ex &&ex) {
+    op_.Exec(out_, std::forward<Ex>(ex));
+  }
+};
+
+} // end namespace detail
+
 namespace experimental {
 
 //
@@ -61,6 +86,7 @@ class sparse_tensor_t
   using crd_type = CRD;
   using pos_type = POS;
   using Format = TF;
+
   static constexpr int DIM = TF::DIM;
   static constexpr int LVL = TF::LVL;
 
@@ -84,13 +110,33 @@ class sparse_tensor_t
       : detail::tensor_impl_t<VAL, DIM, DimDesc,
                               detail::SparseTensorData<VAL, CRD, POS, TF>>(
             shape) {
-    // Initialize primary and secondary storage.
     values_ = std::move(vals);
     for (int l = 0; l < LVL; l++) {
       coordinates_[l] = std::move(crd[l]);
       positions_[l] = std::move(pos[l]);
     }
-    // Set the sparse data in tensor_impl.
+    SetSparseDataImpl();
+  }
+
+  // Default destructor.
+  __MATX_INLINE__ ~sparse_tensor_t() = default;
+
+  // Sets value storage.
+  __MATX_INLINE__ void SetVal(StorageV &&val) { values_ = std::move(val); }
+
+  // Sets coordinates storage.
+  __MATX_INLINE__ void SetCrd(int l, StorageC &&crd) {
+    coordinates_[l] = std::move(crd);
+  }
+
+  // Sets positions storage.
+  __MATX_INLINE__ void SetPos(int l, StorageP &&pos) {
+    positions_[l] = std::move(pos);
+  }
+
+  // Sets sparse data in tensor_impl_t. This method must be called
+  // every time changes are made to the underlying storage objects.
+  void SetSparseDataImpl() {
     VAL *v = values_.data();
     CRD *c[LVL];
     POS *p[LVL];
@@ -104,13 +150,23 @@ class sparse_tensor_t
     this->SetSparseData(v, c, p);
   }
 
-  // Default destructor.
-  __MATX_INLINE__ ~sparse_tensor_t() = default;
+  // A direct sparse tensor assignment (viz. (Acoo = ...).exec();).
+  template <typename T>
+  [[nodiscard]] __MATX_INLINE__ __MATX_HOST__ auto operator=(const T &op) {
+    [[maybe_unused]] typename T::dense2sparse_xform_op valid = true;
+    return detail::sparse_set(*this, op);
+  }
 
   // Size getters.
-  index_t Nse() const { return static_cast<index_t>(values_.size() / sizeof(VAL)); }
-  index_t crdSize(int l) const { return static_cast<index_t>(coordinates_[l].size() / sizeof(CRD)); }
-  index_t posSize(int l) const { return static_cast<index_t>(positions_[l].size() / sizeof(POS)); }
+  index_t Nse() const {
+    return static_cast<index_t>(values_.size() / sizeof(VAL));
+  }
+  index_t crdSize(int l) const {
+    return static_cast<index_t>(coordinates_[l].size() / sizeof(CRD));
+  }
+  index_t posSize(int l) const {
+    return static_cast<index_t>(positions_[l].size() / sizeof(POS));
+  }
 
 private:
   // Primary storage of sparse tensor (explicitly stored element values).

include/matx/operators/dense2sparse.h

@@ -0,0 +1,103 @@
+////////////////////////////////////////////////////////////////////////////////
+// BSD 3-Clause License
+//
+// Copyright (c) 2025, NVIDIA Corporation
+// All rights reserved.
+//
+// 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.
+//
+// 2. Redistributions in binary form must reproduce the above copyright notice,
+//    this list of conditions and the following disclaimer in the documentation
+//    and/or other materials provided with the distribution.
+//
+// 3. Neither the name of the copyright holder nor the names of its
+//    contributors may be used to endorse or promote products derived from
+//    this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+/////////////////////////////////////////////////////////////////////////////////
+
+#pragma once
+
+#include "matx/core/type_utils.h"
+#include "matx/operators/base_operator.h"
+#include "matx/transforms/convert/dense2sparse_cusparse.h"
+
+namespace matx {
+namespace detail {
+
+template <typename OpA>
+class Dense2SparseOp : public BaseOp<Dense2SparseOp<OpA>> {
+private:
+  typename detail::base_type_t<OpA> a_;
+
+public:
+  using matxop = bool;
+  using matx_transform_op = bool;
+  using dense2sparse_xform_op = bool;
+  using value_type = typename OpA::value_type;
+
+  __MATX_INLINE__ Dense2SparseOp(const OpA &a) : a_(a) {}
+
+  __MATX_INLINE__ std::string str() const {
+    return "dense2sparse(" + get_type_str(a_) + ")";
+  }
+
+  static __MATX_INLINE__ constexpr __MATX_HOST__ __MATX_DEVICE__ int32_t
+  Rank() {
+    return remove_cvref_t<OpA>::Rank();
+  }
+
+  constexpr __MATX_INLINE__ __MATX_HOST__ __MATX_DEVICE__ index_t
+  Size(int dim) const {
+    return a_.Size(dim);
+  }
+
+  template <typename Out, typename Executor>
+  void Exec([[maybe_unused]] Out &&out, [[maybe_unused]] Executor &&ex) const {
+    if constexpr (is_sparse_tensor_v<OpA>) {
+      MATX_THROW(matxNotSupported, "Cannot use dense2sparse on sparse input");
+    } else {
+      // NOTE: sparse assignment A = dense2sparse(B) takes direct reference!
+      if constexpr (is_sparse_tensor_v<Out>) {
+        dense2sparse_impl(out, a_, ex);
+      } else {
+        MATX_THROW(matxNotSupported,
+                   "Cannot use dense2sparse for dense output");
+      }
+    }
+  }
+};
+
+} // end namespace detail
+
+/**
+ * Convert a dense tensor into a sparse tensor.
+ *
+ * @tparam OpA
+ *    Data type of A tensor
+ *
+ * @param A
+ *   Dense input tensor
+ *
+ * @return
+ *   Sparse output tensor
+ */
+template <typename OpA> __MATX_INLINE__ auto dense2sparse(const OpA &A) {
+  return detail::Dense2SparseOp(A);
+}
+
+} // end namespace matx

include/matx/operators/operators.h

@@ -52,6 +52,7 @@
 #include "matx/operators/cov.h"
 #include "matx/operators/cross.h"
 #include "matx/operators/cumsum.h"
+#include "matx/operators/dense2sparse.h"
 #include "matx/operators/diag.h"
 #include "matx/operators/dct.h"
 #include "matx/operators/det.h"