gh-225: add matrix row / col extraction

include/spla.h

@@ -377,6 +377,8 @@ SPLA_API spla_Status spla_Exec_m_reduce_by_row(spla_Vector r, spla_Matrix M, spl
 SPLA_API spla_Status spla_Exec_m_reduce_by_column(spla_Vector r, spla_Matrix M, spla_OpBinary op_reduce, spla_Scalar init, spla_Descriptor desc, spla_ScheduleTask* task);
 SPLA_API spla_Status spla_Exec_m_reduce(spla_Scalar r, spla_Scalar s, spla_Matrix M, spla_OpBinary op_reduce, spla_Descriptor desc, spla_ScheduleTask* task);
 SPLA_API spla_Status spla_Exec_m_transpose(spla_Matrix R, spla_Matrix M, spla_OpUnary op_apply, spla_Descriptor desc, spla_ScheduleTask* task);
+SPLA_API spla_Status spla_Exec_m_extract_row(spla_Vector r, spla_Matrix M, spla_uint index, spla_OpUnary op_apply, spla_Descriptor desc, spla_ScheduleTask* task);
+SPLA_API spla_Status spla_Exec_m_extract_column(spla_Vector r, spla_Matrix M, spla_uint index, spla_OpUnary op_apply, spla_Descriptor desc, spla_ScheduleTask* task);
 SPLA_API spla_Status spla_Exec_v_eadd(spla_Vector r, spla_Vector u, spla_Vector v, spla_OpBinary op, spla_Descriptor desc, spla_ScheduleTask* task);
 SPLA_API spla_Status spla_Exec_v_emult(spla_Vector r, spla_Vector u, spla_Vector v, spla_OpBinary op, spla_Descriptor desc, spla_ScheduleTask* task);
 SPLA_API spla_Status spla_Exec_v_eadd_fdb(spla_Vector r, spla_Vector v, spla_Vector fdb, spla_OpBinary op, spla_Descriptor desc, spla_ScheduleTask* task);

include/spla/exec.hpp

@@ -322,6 +322,50 @@ namespace spla {
             ref_ptr<Descriptor>    desc     = ref_ptr<Descriptor>(),
             ref_ptr<ScheduleTask>* task_hnd = nullptr);
 
+    /**
+     * @brief Execute (schedule) matrix row extract
+     *
+     * @note Pass valid `task_hnd` to store as a task, rather then execute immediately.
+     *
+     * @param r Result vector
+     * @param M Source matrix
+     * @param index index of row
+     * @param op_apply Unary op to transform value
+     * @param desc Scheduled task descriptor; default is null
+     * @param task_hnd Optional task hnd; pass not-null pointer to store task
+     *
+     * @return Status on task execution or status on hnd creation
+     */
+    SPLA_API Status exec_m_extract_row(
+            ref_ptr<Vector>        r,
+            ref_ptr<Matrix>        M,
+            uint                   index,
+            ref_ptr<OpUnary>       op_apply,
+            ref_ptr<Descriptor>    desc     = ref_ptr<Descriptor>(),
+            ref_ptr<ScheduleTask>* task_hnd = nullptr);
+
+    /**
+     * @brief Execute (schedule) matrix column extract
+     *
+     * @note Pass valid `task_hnd` to store as a task, rather then execute immediately.
+     *
+     * @param r Result vector
+     * @param M Source matrix
+     * @param index Index of column
+     * @param op_apply Unary op to transform value
+     * @param desc Scheduled task descriptor; default is null
+     * @param task_hnd Optional task hnd; pass not-null pointer to store task
+     *
+     * @return Status on task execution or status on hnd creation
+     */
+    SPLA_API Status exec_m_extract_column(
+            ref_ptr<Vector>        r,
+            ref_ptr<Matrix>        M,
+            uint                   index,
+            ref_ptr<OpUnary>       op_apply,
+            ref_ptr<Descriptor>    desc     = ref_ptr<Descriptor>(),
+            ref_ptr<ScheduleTask>* task_hnd = nullptr);
+
     /**
      * @brief Execute (schedule) element-wise addition by structure of two vectors
      *

python/example.py

@@ -1,7 +1,11 @@
 from pyspla import *
 
-M = Matrix.from_lists([0, 0, 1], [0, 1, 1], [1, 2, 3], (2, 2), INT)
-print(M.eadd(INT.MULT, M.transpose()))
+M = Matrix.from_lists([0, 0, 1, 2], [1, 2, 3, 0], [-1, 1, 2, 3], (3, 4), INT)
+print(M)
+print(M.extract_row(0))
+print(M.extract_row(0, op_apply=INT.AINV))
 
-M = Matrix.from_lists([0, 0, 1], [0, 1, 1], [1, 2, 3], (2, 2), INT)
-print(M.emult(INT.MULT, M.transpose()))
+M = Matrix.from_lists([0, 1, 1, 2], [1, 0, 3, 1], [-1, 1, 2, 3], (3, 4), INT)
+print(M)
+print(M.extract_column(1))
+print(M.extract_column(1, op_apply=INT.AINV))

python/pyspla/bridge.py

@@ -558,6 +558,8 @@ def load_library(lib_path):
     _spla.spla_Exec_m_reduce_by_column.restype = _status_t
     _spla.spla_Exec_m_reduce.restype = _status_t
     _spla.spla_Exec_m_transpose.restype = _status_t
+    _spla.spla_Exec_m_extract_row.restype = _status_t
+    _spla.spla_Exec_m_extract_column.restype = _status_t
     _spla.spla_Exec_v_eadd.restype = _status_t
     _spla.spla_Exec_v_emult.restype = _status_t
     _spla.spla_Exec_v_eadd_fdb.restype = _status_t
@@ -588,6 +590,10 @@ def load_library(lib_path):
         [_object_t, _object_t, _object_t, _object_t, _object_t, _p_object_t]
     _spla.spla_Exec_m_transpose.argtypes = \
         [_object_t, _object_t, _object_t, _object_t, _p_object_t]
+    _spla.spla_Exec_m_extract_row.argtypes = \
+        [_object_t, _object_t, _uint, _object_t, _object_t, _p_object_t]
+    _spla.spla_Exec_m_extract_column.argtypes = \
+        [_object_t, _object_t, _uint, _object_t, _object_t, _p_object_t]
     _spla.spla_Exec_v_eadd.argtypes = \
         [_object_t, _object_t, _object_t, _object_t, _object_t, _p_object_t]
     _spla.spla_Exec_v_emult.argtypes = \

python/pyspla/matrix.py

@@ -1338,6 +1338,130 @@ def transpose(self, out=None, op_apply=None, desc=None):
 
         return out
 
+    def extract_row(self, index, out=None, op_apply=None, desc=None):
+        """
+        Extract matrix row.
+
+        >>> M = Matrix.from_lists([0, 0, 1, 2], [1, 2, 3, 0], [-1, 1, 2, 3], (3, 4), INT)
+        >>> print(M)
+        '
+            0 1 2 3
+         0| .-1 1 .|  0
+         1| . . . 2|  1
+         2| 3 . . .|  2
+            0 1 2 3
+        '
+
+        >>> print(M.extract_row(0))
+        '
+         0| .
+         1|-1
+         2| 1
+         3| .
+        '
+
+        >>> print(M.extract_row(0, op_apply=INT.AINV))
+        '
+         0| .
+         1| 1
+         2|-1
+         3| .
+        '
+
+        :param index: int.
+            Index of row to extract.
+
+        :param out: optional: Vector: default: none.
+            Optional vector to store result.
+
+        :param op_apply: optional: OpUnary. default: None.
+            Optional unary function to apply on extraction.
+
+        :param desc: optional: Descriptor. default: None.
+            Optional descriptor object to configure the execution.
+
+        :return: Vector.
+        """
+
+        from .vector import Vector
+
+        if out is None:
+            out = Vector(shape=self.n_cols, dtype=self.dtype)
+        if op_apply is None:
+            op_apply = self.dtype.IDENTITY
+
+        assert out
+        assert op_apply
+        assert out.dtype == self.dtype
+        assert out.n_rows == self.n_cols
+        assert 0 <= index < self.n_rows
+
+        check(backend().spla_Exec_m_extract_row(out.hnd, self.hnd, ctypes.c_uint(index), op_apply.hnd,
+                                                self._get_desc(desc), self._get_task(None)))
+
+        return out
+
+    def extract_column(self, index, out=None, op_apply=None, desc=None):
+        """
+        Extract matrix column.
+
+        >>> M = Matrix.from_lists([0, 1, 1, 2], [1, 0, 3, 1], [-1, 1, 2, 3], (3, 4), INT)
+        >>> print(M)
+        '
+            0 1 2 3
+         0| .-1 . .|  0
+         1| 1 . . 2|  1
+         2| . 3 . .|  2
+            0 1 2 3
+        '
+
+        >>> print(M.extract_column(1))
+        '
+         0|-1
+         1| .
+         2| 3
+        '
+
+        >>> print(M.extract_column(1, op_apply=INT.AINV))
+        '
+         0| 1
+         1| .
+         2|-3
+        '
+
+        :param index: int.
+            Index of column to extract.
+
+        :param out: optional: Vector: default: none.
+            Optional vector to store result.
+
+        :param op_apply: optional: OpUnary. default: None.
+            Optional unary function to apply on extraction.
+
+        :param desc: optional: Descriptor. default: None.
+            Optional descriptor object to configure the execution.
+
+        :return: Vector.
+        """
+
+        from .vector import Vector
+
+        if out is None:
+            out = Vector(shape=self.n_rows, dtype=self.dtype)
+        if op_apply is None:
+            op_apply = self.dtype.IDENTITY
+
+        assert out
+        assert op_apply
+        assert out.dtype == self.dtype
+        assert out.n_rows == self.n_rows
+        assert 0 <= index < self.n_cols
+
+        check(backend().spla_Exec_m_extract_column(out.hnd, self.hnd, ctypes.c_uint(index), op_apply.hnd,
+                                                   self._get_desc(desc), self._get_task(None)))
+
+        return out
+
     def __str__(self):
         return self.to_string()
 

src/binding/c_exec.cpp

@@ -76,6 +76,12 @@ spla_Status spla_Exec_m_reduce(spla_Scalar r, spla_Scalar s, spla_Matrix M, spla
 spla_Status spla_Exec_m_transpose(spla_Matrix R, spla_Matrix M, spla_OpUnary op_apply, spla_Descriptor desc, spla_ScheduleTask* task) {
     SPLA_WRAP_EXEC(exec_m_transpose, AS_M(R), AS_M(M), AS_OU(op_apply));
 }
+spla_Status spla_Exec_m_extract_row(spla_Vector r, spla_Matrix M, spla_uint index, spla_OpUnary op_apply, spla_Descriptor desc, spla_ScheduleTask* task) {
+    SPLA_WRAP_EXEC(exec_m_extract_row, AS_V(r), AS_M(M), index, AS_OU(op_apply));
+}
+spla_Status spla_Exec_m_extract_column(spla_Vector r, spla_Matrix M, spla_uint index, spla_OpUnary op_apply, spla_Descriptor desc, spla_ScheduleTask* task) {
+    SPLA_WRAP_EXEC(exec_m_extract_column, AS_V(r), AS_M(M), index, AS_OU(op_apply));
+}
 spla_Status spla_Exec_v_eadd(spla_Vector r, spla_Vector u, spla_Vector v, spla_OpBinary op, spla_Descriptor desc, spla_ScheduleTask* task) {
     SPLA_WRAP_EXEC(exec_v_eadd, AS_V(r), AS_V(u), AS_V(v), AS_OB(op));
 }

src/cpu/cpu_algo_registry.cpp

@@ -34,6 +34,8 @@
 #include <cpu/cpu_kron.hpp>
 #include <cpu/cpu_m_eadd.hpp>
 #include <cpu/cpu_m_emult.hpp>
+#include <cpu/cpu_m_extract_column.hpp>
+#include <cpu/cpu_m_extract_row.hpp>
 #include <cpu/cpu_m_reduce.hpp>
 #include <cpu/cpu_m_reduce_by_column.hpp>
 #include <cpu/cpu_m_reduce_by_row.hpp>
@@ -121,6 +123,16 @@ namespace spla {
         g_registry->add(MAKE_KEY_CPU_0("m_transpose", UINT), std::make_shared<Algo_m_transpose_cpu<T_UINT>>());
         g_registry->add(MAKE_KEY_CPU_0("m_transpose", FLOAT), std::make_shared<Algo_m_transpose_cpu<T_FLOAT>>());
 
+        // algorthm m_extract_row
+        g_registry->add(MAKE_KEY_CPU_0("m_extract_row", INT), std::make_shared<Algo_m_extract_row_cpu<T_INT>>());
+        g_registry->add(MAKE_KEY_CPU_0("m_extract_row", UINT), std::make_shared<Algo_m_extract_row_cpu<T_UINT>>());
+        g_registry->add(MAKE_KEY_CPU_0("m_extract_row", FLOAT), std::make_shared<Algo_m_extract_row_cpu<T_FLOAT>>());
+
+        // algorthm m_extract_column
+        g_registry->add(MAKE_KEY_CPU_0("m_extract_column", INT), std::make_shared<Algo_m_extract_column_cpu<T_INT>>());
+        g_registry->add(MAKE_KEY_CPU_0("m_extract_column", UINT), std::make_shared<Algo_m_extract_column_cpu<T_UINT>>());
+        g_registry->add(MAKE_KEY_CPU_0("m_extract_column", FLOAT), std::make_shared<Algo_m_extract_column_cpu<T_FLOAT>>());
+
         // algorthm mxv_masked
         g_registry->add(MAKE_KEY_CPU_0("mxv_masked", INT), std::make_shared<Algo_mxv_masked_cpu<T_INT>>());
         g_registry->add(MAKE_KEY_CPU_0("mxv_masked", UINT), std::make_shared<Algo_mxv_masked_cpu<T_UINT>>());

src/cpu/cpu_format_coo_vec.hpp

@@ -30,13 +30,33 @@
 
 #include <cpu/cpu_formats.hpp>
 
+#include <algorithm>
+#include <vector>
+
 namespace spla {
 
     /**
      * @addtogroup internal
      * @{
      */
 
+    template<typename T>
+    void cpu_coo_vec_sort(CpuCooVec<T>& vec) {
+        std::vector<std::pair<uint, T>> buffer;
+        buffer.reserve(vec.values);
+
+        for (uint i = 0; i < vec.values; i++) {
+            buffer.emplace_back(vec.Ai[i], vec.Ax[i]);
+        }
+
+        std::sort(buffer.begin(), buffer.end(), [](auto& a, auto& b) { return a.first < b.first; });
+
+        for (uint i = 0; i < vec.values; i++) {
+            vec.Ai[i] = buffer[i].first;
+            vec.Ax[i] = buffer[i].second;
+        }
+    }
+
     template<typename T>
     void cpu_coo_vec_resize(const uint    n_values,
                             CpuCooVec<T>& vec) {