Add ell scale function

ginkgo-project · Jan 12, 2024 · fba8164 · fba8164
1 parent 273153e
commit fba8164
Show file tree

Hide file tree

Showing 4 changed files with 51 additions and 0 deletions.
diff --git a/core/matrix/batch_ell.cpp b/core/matrix/batch_ell.cpp
@@ -229,6 +229,24 @@ void Ell<ValueType, IndexType>::move_to(
 GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(GKO_DECLARE_BATCH_ELL_MATRIX);
 
 
+template <typename ValueType, typename IndexType>
+void two_sided_scale(const array<ValueType>& col_scale,
+                     const array<ValueType>& row_scale,
+                     batch::matrix::Ell<ValueType, IndexType>* in_out)
+{
+    GKO_NOT_IMPLEMENTED;
+}
+
+
+#define GKO_DECLARE_TWO_SIDED_BATCH_SCALE(_vtype, _itype) \
+    void two_sided_scale(const array<_vtype>& col_scale,  \
+                         const array<_vtype>& row_scale,  \
+                         batch::matrix::Ell<_vtype, _itype>* in_out)
+
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INT32_TYPE(
+    GKO_DECLARE_TWO_SIDED_BATCH_SCALE);
+
+
 }  // namespace matrix
 }  // namespace batch
 }  // namespace gko
diff --git a/include/ginkgo/core/matrix/batch_csr.hpp b/include/ginkgo/core/matrix/batch_csr.hpp
@@ -354,6 +354,15 @@ class Csr final
 };
 
 
+/**
+ * Performs in-place row and column scaling for a given matrix.
+ *
+ * @param col_scale  the column scalars
+ * @param row_scale  the row scalars
+ * @param in_out  the matrix to be scaled
+ *
+ * @note the operation is performed in-place
+ */
 template <typename ValueType, typename IndexType>
 void two_sided_scale(const array<ValueType>& col_scale,
                      const array<ValueType>& row_scale,

diff --git a/include/ginkgo/core/matrix/batch_dense.hpp b/include/ginkgo/core/matrix/batch_dense.hpp
@@ -371,6 +371,15 @@ class Dense final : public EnableBatchLinOp<Dense<ValueType>>,
 };
 
 
+/**
+ * Performs in-place row and column scaling for a given matrix.
+ *
+ * @param col_scale  the column scalars
+ * @param row_scale  the row scalars
+ * @param in_out  the matrix to be scaled
+ *
+ * @note the operation is performed in-place
+ */
 template <typename ValueType>
 void two_sided_scale(const array<ValueType>& col_scale,
                      const array<ValueType>& row_scale,

diff --git a/include/ginkgo/core/matrix/batch_ell.hpp b/include/ginkgo/core/matrix/batch_ell.hpp
@@ -371,6 +371,21 @@ class Ell final
 };
 
 
+/**
+ * Performs in-place row and column scaling for a given matrix.
+ *
+ * @param col_scale  the column scalars
+ * @param row_scale  the row scalars
+ * @param in_out  the matrix to be scaled
+ *
+ * @note the operation is performed in-place
+ */
+template <typename ValueType, typename IndexType>
+void two_sided_scale(const array<ValueType>& col_scale,
+                     const array<ValueType>& row_scale,
+                     batch::matrix::Ell<ValueType, IndexType>* in_out);
+
+
 }  // namespace matrix
 }  // namespace batch
 }  // namespace gko