review updates

- fix and test csr column permutation for hypersparse matrices - fix inverse nonsymm scaled permutation for dense - extract permutation dimension validation - simplify permutation dispatch - re-add templated constructors to preserve interface - move more code to source - add documentation for permutation to permute_mode instead of matrix classes - improve documentation - make scaled_permutation final - remove unnecessary InvalidValueError - improve error message on dispatch_helper run(...) - simplify rectangular permutation tests - fix dense tests not calling scaled permutation - test apply in (scaled) permutation on the device - throw DimensionMismatch instead of ValueMismatch - add more round-trip and inverted permutation tests - add dpcpp subgroup size Co-authored-by: Pratik Nayak <[email protected]> Co-authored-by: Yuhsiang M. Tsai <[email protected]> Co-authored-by: Marcel Koch <[email protected]>
ginkgo-project · Oct 24, 2023 · 8c4a46d · 8c4a46d
1 parent a69266f
commit 8c4a46d
Show file tree

Hide file tree

Showing 22 changed files with 557 additions and 384 deletions.
diff --git a/common/unified/matrix/csr_kernels.cpp b/common/unified/matrix/csr_kernels.cpp
@@ -62,7 +62,7 @@ void inv_col_permute(std::shared_ptr<const DefaultExecutor> exec,
 {
     auto num_rows = orig->get_size()[0];
     auto nnz = orig->get_num_stored_elements();
-    auto size = std::max(num_rows, nnz);
+    auto size = std::max(num_rows + 1, nnz);
     run_kernel(
         exec,
         [] GKO_KERNEL(auto tid, auto num_rows, auto num_nonzeros,
@@ -95,7 +95,7 @@ void inv_col_scale_permute(std::shared_ptr<const DefaultExecutor> exec,
 {
     auto num_rows = orig->get_size()[0];
     auto nnz = orig->get_num_stored_elements();
-    auto size = std::max(num_rows, nnz);
+    auto size = std::max(num_rows + 1, nnz);
     run_kernel(
         exec,
         [] GKO_KERNEL(auto tid, auto num_rows, auto num_nonzeros, auto scale,

diff --git a/common/unified/matrix/dense_kernels.template.cpp b/common/unified/matrix/dense_kernels.template.cpp
@@ -602,7 +602,7 @@ void inv_nonsymm_scale_permute(std::shared_ptr<const DefaultExecutor> exec,
                       auto col_scale, auto col_perm, auto orig, auto permuted) {
             const auto row = row_perm[i];
             const auto col = col_perm[j];
-            permuted(row, col) = orig(i, j) / (row_scale[i] * col_scale[j]);
+            permuted(row, col) = orig(i, j) / (row_scale[row] * col_scale[col]);
         },
         orig->get_size(), row_scale, row_perm, col_scale, col_perm, orig,
         permuted);

diff --git a/core/base/dispatch_helper.hpp b/core/base/dispatch_helper.hpp
@@ -97,9 +97,9 @@ void run(T obj, Func f, Args... args)
  */
 template <template <typename> class Base, typename T, typename Func,
           typename... Args>
-void run(T, Func, Args...)
+void run(T obj, Func, Args...)
 {
-    GKO_NOT_IMPLEMENTED;
+    GKO_NOT_SUPPORTED(obj);
 }
 
 /**

diff --git a/core/matrix/csr.cpp b/core/matrix/csr.cpp
@@ -45,6 +45,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include <ginkgo/core/matrix/ell.hpp>
 #include <ginkgo/core/matrix/fbcsr.hpp>
 #include <ginkgo/core/matrix/identity.hpp>
+#include <ginkgo/core/matrix/permutation.hpp>
 #include <ginkgo/core/matrix/scaled_permutation.hpp>
 #include <ginkgo/core/matrix/sellp.hpp>
 #include <ginkgo/core/matrix/sparsity_csr.hpp>
@@ -58,6 +59,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include "core/matrix/csr_kernels.hpp"
 #include "core/matrix/ell_kernels.hpp"
 #include "core/matrix/hybrid_kernels.hpp"
+#include "core/matrix/permutation.hpp"
 #include "core/matrix/sellp_kernels.hpp"
 
 
@@ -535,6 +537,7 @@ std::unique_ptr<Csr<ValueType, IndexType>> Csr<ValueType, IndexType>::permute(
     const auto exec = this->get_executor();
     const auto size = this->get_size();
     const auto nnz = this->get_num_stored_elements();
+    validate_permute_dimensions(size, permutation->get_size(), mode);
     if ((mode & permute_mode::symmetric) == permute_mode::none) {
         return this->clone();
     }
@@ -552,7 +555,8 @@ std::unique_ptr<Csr<ValueType, IndexType>> Csr<ValueType, IndexType>::permute(
     std::unique_ptr<const Permutation<IndexType>> inv_permutation;
     const auto perm_idxs = local_permutation->get_const_permutation();
     const IndexType* inv_perm_idxs{};
-    // to permute columns, we need to know the inverse permutation
+    // Due to the sparse storage, we can only inverse-permute columns, so we
+    // need to compute the inverse for forward-permutations.
     bool needs_inverse =
         (mode & permute_mode::inverse_columns) == permute_mode::columns;
     if (needs_inverse) {
@@ -580,7 +584,7 @@ std::unique_ptr<Csr<ValueType, IndexType>> Csr<ValueType, IndexType>::permute(
         exec->run(csr::make_inv_symm_permute(perm_idxs, this, result.get()));
         break;
     default:
-        GKO_ASSERT(false);
+        GKO_INVALID_STATE("Invalid permute mode");
     }
     result->make_srow();
     if ((mode & permute_mode::columns) == permute_mode::columns) {
@@ -598,8 +602,8 @@ std::unique_ptr<Csr<ValueType, IndexType>> Csr<ValueType, IndexType>::permute(
     const auto exec = this->get_executor();
     const auto size = this->get_size();
     const auto nnz = this->get_num_stored_elements();
-    GKO_ASSERT_EQ(size[0], row_permutation->get_size()[0]);
-    GKO_ASSERT_EQ(size[1], col_permutation->get_size()[0]);
+    GKO_ASSERT_EQUAL_ROWS(this, row_permutation);
+    GKO_ASSERT_EQUAL_COLS(this, col_permutation);
     auto result = Csr::create(exec, size, nnz, this->get_strategy()->copy());
     auto local_row_permutation = make_temporary_clone(exec, row_permutation);
     auto local_col_permutation = make_temporary_clone(exec, col_permutation);
@@ -630,18 +634,10 @@ Csr<ValueType, IndexType>::scale_permute(
     const auto exec = this->get_executor();
     const auto size = this->get_size();
     const auto nnz = this->get_num_stored_elements();
+    validate_permute_dimensions(size, permutation->get_size(), mode);
     if ((mode & permute_mode::symmetric) == permute_mode::none) {
         return this->clone();
     }
-    if ((mode & permute_mode::symmetric) == permute_mode::symmetric) {
-        GKO_ASSERT_IS_SQUARE_MATRIX(this);
-    }
-    if ((mode & permute_mode::rows) == permute_mode::rows) {
-        GKO_ASSERT_EQ(size[0], permutation->get_size()[0]);
-    }
-    if ((mode & permute_mode::columns) == permute_mode::columns) {
-        GKO_ASSERT_EQ(size[1], permutation->get_size()[0]);
-    }
     auto result = Csr::create(exec, size, nnz, this->get_strategy()->copy());
     auto local_permutation = make_temporary_clone(exec, permutation);
     std::unique_ptr<const ScaledPermutation<ValueType, IndexType>>
@@ -684,7 +680,7 @@ Csr<ValueType, IndexType>::scale_permute(
                                                    this, result.get()));
         break;
     default:
-        GKO_ASSERT(false);
+        GKO_INVALID_STATE("Invalid permute mode");
     }
     result->make_srow();
     if ((mode & permute_mode::columns) == permute_mode::columns) {
@@ -704,8 +700,8 @@ Csr<ValueType, IndexType>::scale_permute(
     const auto exec = this->get_executor();
     const auto size = this->get_size();
     const auto nnz = this->get_num_stored_elements();
-    GKO_ASSERT_EQ(size[0], row_permutation->get_size()[0]);
-    GKO_ASSERT_EQ(size[1], col_permutation->get_size()[0]);
+    GKO_ASSERT_EQUAL_ROWS(this, row_permutation);
+    GKO_ASSERT_EQUAL_COLS(this, col_permutation);
     auto result = Csr::create(exec, size, nnz, this->get_strategy()->copy());
     auto local_row_permutation = make_temporary_clone(exec, row_permutation);
     auto local_col_permutation = make_temporary_clone(exec, col_permutation);

diff --git a/core/matrix/dense.cpp b/core/matrix/dense.cpp
@@ -61,6 +61,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include "core/components/prefix_sum_kernels.hpp"
 #include "core/matrix/dense_kernels.hpp"
 #include "core/matrix/hybrid_kernels.hpp"
+#include "core/matrix/permutation.hpp"
 
 
 namespace gko {
@@ -1133,19 +1134,11 @@ void Dense<ValueType>::permute_impl(const Permutation<IndexType>* permutation,
     const auto exec = this->get_executor();
     const auto size = this->get_size();
     GKO_ASSERT_EQUAL_DIMENSIONS(this, output);
+    validate_permute_dimensions(size, permutation->get_size(), mode);
     if ((mode & permute_mode::symmetric) == permute_mode::none) {
         output->copy_from(this);
         return;
     }
-    if ((mode & permute_mode::symmetric) == permute_mode::symmetric) {
-        GKO_ASSERT_IS_SQUARE_MATRIX(this);
-    }
-    if ((mode & permute_mode::rows) == permute_mode::rows) {
-        GKO_ASSERT_EQ(size[0], permutation->get_size()[0]);
-    }
-    if ((mode & permute_mode::columns) == permute_mode::columns) {
-        GKO_ASSERT_EQ(size[1], permutation->get_size()[0]);
-    }
     auto local_output = make_temporary_output_clone(exec, output);
     auto local_perm = make_temporary_clone(exec, permutation);
     switch (mode) {
@@ -1174,7 +1167,7 @@ void Dense<ValueType>::permute_impl(const Permutation<IndexType>* permutation,
             local_perm->get_const_permutation(), this, local_output.get()));
         break;
     default:
-        GKO_ASSERT(false);  // cannot happen
+        GKO_INVALID_STATE("Invalid permute mode");
     }
 }
 
@@ -1189,8 +1182,8 @@ void Dense<ValueType>::permute_impl(
     auto exec = this->get_executor();
     auto size = this->get_size();
     GKO_ASSERT_EQUAL_DIMENSIONS(this, output);
-    GKO_ASSERT_EQ(size[0], row_permutation->get_size()[0]);
-    GKO_ASSERT_EQ(size[1], col_permutation->get_size()[0]);
+    GKO_ASSERT_EQUAL_ROWS(this, row_permutation);
+    GKO_ASSERT_EQUAL_COLS(this, col_permutation);
     auto local_output = make_temporary_output_clone(exec, output);
     auto local_row_perm = make_temporary_clone(exec, row_permutation);
     auto local_col_perm = make_temporary_clone(exec, col_permutation);
@@ -1215,19 +1208,11 @@ void Dense<ValueType>::scale_permute_impl(
     const auto exec = this->get_executor();
     const auto size = this->get_size();
     GKO_ASSERT_EQUAL_DIMENSIONS(this, output);
+    validate_permute_dimensions(size, permutation->get_size(), mode);
     if ((mode & permute_mode::symmetric) == permute_mode::none) {
         output->copy_from(this);
         return;
     }
-    if ((mode & permute_mode::symmetric) == permute_mode::symmetric) {
-        GKO_ASSERT_IS_SQUARE_MATRIX(this);
-    }
-    if ((mode & permute_mode::rows) == permute_mode::rows) {
-        GKO_ASSERT_EQ(size[0], permutation->get_size()[0]);
-    }
-    if ((mode & permute_mode::columns) == permute_mode::columns) {
-        GKO_ASSERT_EQ(size[1], permutation->get_size()[0]);
-    }
     auto local_output = make_temporary_output_clone(exec, output);
     auto local_perm = make_temporary_clone(exec, permutation);
     switch (mode) {
@@ -1262,7 +1247,7 @@ void Dense<ValueType>::scale_permute_impl(
             this, local_output.get()));
         break;
     default:
-        GKO_ASSERT(false);  // cannot happen
+        GKO_INVALID_STATE("Invalid permute mode");
     }
 }
 
@@ -1277,8 +1262,8 @@ void Dense<ValueType>::scale_permute_impl(
     auto exec = this->get_executor();
     auto size = this->get_size();
     GKO_ASSERT_EQUAL_DIMENSIONS(this, output);
-    GKO_ASSERT_EQ(size[0], row_permutation->get_size()[0]);
-    GKO_ASSERT_EQ(size[1], col_permutation->get_size()[0]);
+    GKO_ASSERT_EQUAL_ROWS(this, row_permutation);
+    GKO_ASSERT_EQUAL_COLS(this, col_permutation);
     auto local_output = make_temporary_output_clone(exec, output);
     auto local_row_perm = make_temporary_clone(exec, row_permutation);
     auto local_col_perm = make_temporary_clone(exec, col_permutation);

diff --git a/core/matrix/permutation.cpp b/core/matrix/permutation.cpp
@@ -56,6 +56,76 @@ GKO_REGISTER_OPERATION(combine, permutation::combine);
 }  // namespace permutation
 
 
+void validate_permute_dimensions(dim<2> size, dim<2> permutation_size,
+                                 permute_mode mode)
+{
+    if ((mode & permute_mode::symmetric) == permute_mode::symmetric) {
+        GKO_ASSERT_IS_SQUARE_MATRIX(size);
+    }
+    if ((mode & permute_mode::rows) == permute_mode::rows) {
+        if (size[0] != permutation_size[0]) {
+            throw DimensionMismatch(
+                __FILE__, __LINE__, __func__, "matrix", size[0], size[1],
+                "permutation", permutation_size[0], permutation_size[0],
+                "expected the permutation size to match the number of rows");
+        };
+    }
+    if ((mode & permute_mode::columns) == permute_mode::columns) {
+        if (size[1] != permutation_size[0]) {
+            throw DimensionMismatch(
+                __FILE__, __LINE__, __func__, "matrix", size[0], size[1],
+                "permutation", permutation_size[0], permutation_size[0],
+                "expected the permutation size to match the number of columns");
+        };
+    }
+}
+
+
+permute_mode operator|(permute_mode a, permute_mode b)
+{
+    return static_cast<permute_mode>(static_cast<unsigned>(a) |
+                                     static_cast<unsigned>(b));
+}
+
+
+permute_mode operator&(permute_mode a, permute_mode b)
+{
+    return static_cast<permute_mode>(static_cast<unsigned>(a) &
+                                     static_cast<unsigned>(b));
+}
+
+
+permute_mode operator^(permute_mode a, permute_mode b)
+{
+    return static_cast<permute_mode>(static_cast<unsigned>(a) ^
+                                     static_cast<unsigned>(b));
+}
+
+
+std::ostream& operator<<(std::ostream& stream, permute_mode mode)
+{
+    switch (mode) {
+    case permute_mode::none:
+        return stream << "none";
+    case permute_mode::rows:
+        return stream << "rows";
+    case permute_mode::columns:
+        return stream << "columns";
+    case permute_mode::symmetric:
+        return stream << "symmetric";
+    case permute_mode::inverse:
+        return stream << "inverse";
+    case permute_mode::inverse_rows:
+        return stream << "inverse_rows";
+    case permute_mode::inverse_columns:
+        return stream << "inverse_columns";
+    case permute_mode::inverse_symmetric:
+        return stream << "inverse_symmetric";
+    }
+    return stream;
+}
+
+
 template <typename IndexType>
 std::unique_ptr<const Permutation<IndexType>>
 Permutation<IndexType>::create_const(
@@ -118,34 +188,10 @@ Permutation<IndexType>::Permutation(std::shared_ptr<const Executor> exec,
 }
 
 
-template <typename IndexType>
-Permutation<IndexType>::Permutation(std::shared_ptr<const Executor> exec,
-                                    const dim<2>& size,
-                                    array<index_type> permutation_indices)
-    : Permutation{std::move(exec), std::move(permutation_indices)}
-{
-    GKO_ASSERT_EQ(size[0], permutation_.get_num_elems());
-    GKO_ASSERT_IS_SQUARE_MATRIX(size);
-}
-
-
-template <typename IndexType>
-Permutation<IndexType>::Permutation(std::shared_ptr<const Executor> exec,
-                                    const dim<2>& size,
-                                    array<index_type> permutation_indices,
-                                    const mask_type& enabled_permute)
-    : Permutation{std::move(exec), std::move(permutation_indices)}
-{
-    GKO_ASSERT_EQ(enabled_permute, row_permute);
-    GKO_ASSERT_EQ(size[0], permutation_.get_num_elems());
-    GKO_ASSERT_IS_SQUARE_MATRIX(size);
-}
-
-
 template <typename IndexType>
 size_type Permutation<IndexType>::get_permutation_size() const noexcept
 {
-    return permutation_.get_num_elems();
+    return this->get_size()[0];
 }
 
 
@@ -212,14 +258,8 @@ void dispatch_dense(const LinOp* op, Functor fn)
 {
     using matrix::Dense;
     using std::complex;
-    if (dynamic_cast<const ConvertibleTo<Dense<double>>*>(op)) {
-        run<const Dense<double>*, const Dense<float>*>(op, fn);
-    } else if (dynamic_cast<const ConvertibleTo<Dense<complex<double>>>*>(op)) {
-        run<const Dense<complex<double>>*, const Dense<complex<float>>*>(op,
-                                                                         fn);
-    } else {
-        GKO_NOT_SUPPORTED(*op);
-    }
+    run<const Dense<double>*, const Dense<float>*,
+        const Dense<complex<double>>*, const Dense<complex<float>>*>(op, fn);
 }