optimize matrix generation

core/test/utils/matrix_generator.hpp

@@ -37,7 +37,9 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include <algorithm>
 #include <iterator>
 #include <numeric>
+#include <random>
 #include <type_traits>
+#include <unordered_set>
 #include <vector>
 
 
@@ -82,22 +84,48 @@ matrix_data<ValueType, IndexType> generate_random_matrix_data(
 
     matrix_data<ValueType, IndexType> data{gko::dim<2>{num_rows, num_cols}, {}};
 
-    std::vector<size_type> col_idx(num_cols);
-    std::iota(begin(col_idx), end(col_idx), size_type(0));
+    std::vector<bool> present_cols(num_cols);
 
-    for (size_type row = 0; row < num_rows; ++row) {
+    for (IndexType row = 0; row < num_rows; ++row) {
         // randomly generate number of nonzeros in this row
-        auto nnz_in_row = static_cast<size_type>(nonzero_dist(engine));
-        nnz_in_row = std::max(size_type(0), std::min(nnz_in_row, num_cols));
-        // select a subset of `nnz_in_row` column indexes, and fill these
-        // locations with random values
-        std::shuffle(begin(col_idx), end(col_idx), engine);
-        std::for_each(
-            begin(col_idx), begin(col_idx) + nnz_in_row, [&](size_type col) {
-                data.nonzeros.emplace_back(
-                    row, col,
-                    detail::get_rand_value<ValueType>(value_dist, engine));
-            });
+        const auto nnz_in_row = std::max(
+            size_type(0),
+            std::min(static_cast<size_type>(nonzero_dist(engine)), num_cols));
+        std::uniform_int_distribution<IndexType> col_dist{
+            0, static_cast<IndexType>(num_cols) - 1};
+        if (nnz_in_row > num_cols / 2) {
+            present_cols.assign(num_cols, true);
+            // remove num_cols - nnz_in_row entries from present_cols
+            size_type count = num_cols;
+            while (count > nnz_in_row) {
+                const auto new_col = col_dist(engine);
+                if (present_cols[new_col]) {
+                    present_cols[new_col] = false;
+                    count--;
+                }
+            }
+            for (IndexType col = 0; col < num_cols; col++) {
+                if (present_cols[col]) {
+                    data.nonzeros.emplace_back(
+                        row, col,
+                        detail::get_rand_value<ValueType>(value_dist, engine));
+                }
+            }
+        } else {
+            // add nnz_in_row entries to present_cols
+            present_cols.assign(num_cols, false);
+            size_type count = 0;
+            while (count < nnz_in_row) {
+                const auto col = col_dist(engine);
+                if (!present_cols[col]) {
+                    present_cols[col] = true;
+                    count++;
+                    data.nonzeros.emplace_back(
+                        row, col,
+                        detail::get_rand_value<ValueType>(value_dist, engine));
+                }
+            }
+        }
     }
 
     data.ensure_row_major_order();
@@ -163,52 +191,72 @@ std::unique_ptr<MatrixType> generate_random_matrix(
 template <typename ValueType, typename IndexType, typename NonzeroDistribution,
           typename ValueDistribution, typename Engine>
 matrix_data<ValueType, IndexType> generate_random_triangular_matrix_data(
-    size_type num_rows, size_type num_cols, bool ones_on_diagonal,
-    bool lower_triangular, NonzeroDistribution&& nonzero_dist,
-    ValueDistribution&& value_dist, Engine&& engine)
+    size_type size, bool ones_on_diagonal, bool lower_triangular,
+    NonzeroDistribution&& nonzero_dist, ValueDistribution&& value_dist,
+    Engine&& engine)
 {
     using std::begin;
     using std::end;
 
-    matrix_data<ValueType, IndexType> data{gko::dim<2>{num_rows, num_cols}, {}};
-    ValueType one = 1.0;
-    std::vector<size_type> col_idx(num_cols);
-    std::iota(begin(col_idx), end(col_idx), size_type(0));
+    matrix_data<ValueType, IndexType> data{gko::dim<2>{size, size}, {}};
 
-    for (size_type row = 0; row < num_rows; ++row) {
+    std::vector<bool> present_cols(size);
+
+    for (IndexType row = 0; row < size; ++row) {
         // randomly generate number of nonzeros in this row
-        auto nnz_in_row = static_cast<size_type>(nonzero_dist(engine));
-        nnz_in_row = std::max(size_type(0), std::min(nnz_in_row, num_cols));
-        // select a subset of `nnz_in_row` column indexes, and fill these
-        // locations with random values
-        std::shuffle(begin(col_idx), end(col_idx), engine);
-        // add non-zeros
-        bool has_diagonal{};
-        for (size_type nz = 0; nz < nnz_in_row; ++nz) {
-            auto col = col_idx[nz];
-            // skip non-zeros outside triangle
-            if ((col > row && lower_triangular) ||
-                (col < row && !lower_triangular)) {
-                continue;
+        const auto min_col = lower_triangular ? 0 : row;
+        const auto max_col =
+            lower_triangular ? row : static_cast<IndexType>(size) - 1;
+        const auto max_row_nnz = max_col - min_col + 1;
+        const auto nnz_in_row = std::max(
+            size_type(0), std::min(static_cast<size_type>(nonzero_dist(engine)),
+                                   static_cast<size_type>(max_row_nnz)));
+        std::uniform_int_distribution<IndexType> col_dist{min_col, max_col};
+        if (nnz_in_row > max_row_nnz / 2) {
+            present_cols.assign(size, true);
+            // remove max_row_nnz - nnz_in_row entries from present_cols
+            size_type count = max_row_nnz;
+            while (nnz_in_row > count) {
+                const auto new_col = col_dist(engine);
+                if (present_cols[new_col]) {
+                    present_cols[new_col] = false;
+                    count--;
+                }
             }
-
-            // generate and store non-zero
-            auto val = detail::get_rand_value<ValueType>(value_dist, engine);
-            if (col == row) {
-                has_diagonal = true;
-                if (ones_on_diagonal) {
-                    val = one;
+            for (auto col = min_col; col <= max_col; col++) {
+                if (present_cols[col] || col == row) {
+                    data.nonzeros.emplace_back(
+                        row, col,
+                        row == col && ones_on_diagonal
+                            ? one<ValueType>()
+                            : detail::get_rand_value<ValueType>(value_dist,
+                                                                engine));
                 }
             }
-            data.nonzeros.emplace_back(row, col, val);
-        }
-
-        // add diagonal if it hasn't been added yet
-        if (!has_diagonal && row < num_cols) {
-            auto val = ones_on_diagonal ? one
-                                        : detail::get_rand_value<ValueType>(
-                                              value_dist, engine);
-            data.nonzeros.emplace_back(row, row, val);
+        } else {
+            // add nnz_in_row entries to present_cols
+            present_cols.assign(size, false);
+            size_type count = 0;
+            while (count < nnz_in_row) {
+                const auto col = col_dist(engine);
+                if (!present_cols[col]) {
+                    present_cols[col] = true;
+                    count++;
+                    data.nonzeros.emplace_back(
+                        row, col,
+                        row == col && ones_on_diagonal
+                            ? one<ValueType>()
+                            : detail::get_rand_value<ValueType>(value_dist,
+                                                                engine));
+                }
+            }
+            if (!present_cols[row]) {
+                data.nonzeros.emplace_back(
+                    row, row,
+                    ones_on_diagonal ? one<ValueType>()
+                                     : detail::get_rand_value<ValueType>(
+                                           value_dist, engine));
+            }
         }
     }
 
@@ -234,16 +282,15 @@ matrix_data<ValueType, IndexType> generate_random_triangular_matrix_data(
 template <typename MatrixType = matrix::Dense<>, typename NonzeroDistribution,
           typename ValueDistribution, typename Engine, typename... MatrixArgs>
 std::unique_ptr<MatrixType> generate_random_triangular_matrix(
-    size_type num_rows, size_type num_cols, bool ones_on_diagonal,
-    bool lower_triangular, NonzeroDistribution&& nonzero_dist,
-    ValueDistribution&& value_dist, Engine&& engine,
-    std::shared_ptr<const Executor> exec, MatrixArgs&&... args)
+    size_type size, bool ones_on_diagonal, bool lower_triangular,
+    NonzeroDistribution&& nonzero_dist, ValueDistribution&& value_dist,
+    Engine&& engine, std::shared_ptr<const Executor> exec, MatrixArgs&&... args)
 {
     using value_type = typename MatrixType::value_type;
     using index_type = typename MatrixType::index_type;
     auto result = MatrixType::create(exec, std::forward<MatrixArgs>(args)...);
     result->read(generate_random_triangular_matrix_data<value_type, index_type>(
-        num_rows, num_cols, ones_on_diagonal, lower_triangular,
+        size, ones_on_diagonal, lower_triangular,
         std::forward<NonzeroDistribution>(nonzero_dist),
         std::forward<ValueDistribution>(value_dist),
         std::forward<Engine>(engine)));
@@ -277,13 +324,13 @@ std::unique_ptr<MatrixType> generate_random_triangular_matrix(
 template <typename MatrixType = matrix::Dense<>, typename NonzeroDistribution,
           typename ValueDistribution, typename Engine, typename... MatrixArgs>
 std::unique_ptr<MatrixType> generate_random_lower_triangular_matrix(
-    size_type num_rows, size_type num_cols, bool ones_on_diagonal,
-    NonzeroDistribution&& nonzero_dist, ValueDistribution&& value_dist,
-    Engine&& engine, std::shared_ptr<const Executor> exec, MatrixArgs&&... args)
+    size_type size, bool ones_on_diagonal, NonzeroDistribution&& nonzero_dist,
+    ValueDistribution&& value_dist, Engine&& engine,
+    std::shared_ptr<const Executor> exec, MatrixArgs&&... args)
 {
     return generate_random_triangular_matrix<MatrixType>(
-        num_rows, num_cols, ones_on_diagonal, true, nonzero_dist, value_dist,
-        engine, std::move(exec), std::forward<MatrixArgs>(args)...);
+        size, ones_on_diagonal, true, nonzero_dist, value_dist, engine,
+        std::move(exec), std::forward<MatrixArgs>(args)...);
 }
 
 
@@ -313,13 +360,13 @@ std::unique_ptr<MatrixType> generate_random_lower_triangular_matrix(
 template <typename MatrixType = matrix::Dense<>, typename NonzeroDistribution,
           typename ValueDistribution, typename Engine, typename... MatrixArgs>
 std::unique_ptr<MatrixType> generate_random_upper_triangular_matrix(
-    size_type num_rows, size_type num_cols, bool ones_on_diagonal,
-    NonzeroDistribution&& nonzero_dist, ValueDistribution&& value_dist,
-    Engine&& engine, std::shared_ptr<const Executor> exec, MatrixArgs&&... args)
+    size_type size, bool ones_on_diagonal, NonzeroDistribution&& nonzero_dist,
+    ValueDistribution&& value_dist, Engine&& engine,
+    std::shared_ptr<const Executor> exec, MatrixArgs&&... args)
 {
     return generate_random_triangular_matrix<MatrixType>(
-        num_rows, num_cols, ones_on_diagonal, false, nonzero_dist, value_dist,
-        engine, std::move(exec), std::forward<MatrixArgs>(args)...);
+        size, ones_on_diagonal, false, nonzero_dist, value_dist, engine,
+        std::move(exec), std::forward<MatrixArgs>(args)...);
 }
 
 

core/test/utils/matrix_generator_test.cpp

@@ -60,11 +60,11 @@ class MatrixGenerator : public ::testing::Test {
               std::normal_distribution<real_type>(20.0, 5.0),
               std::default_random_engine(42), exec)),
           l_mtx(gko::test::generate_random_lower_triangular_matrix<mtx_type>(
-              4, 3, true, std::normal_distribution<real_type>(50, 5),
+              4, true, std::normal_distribution<real_type>(50, 5),
               std::normal_distribution<real_type>(20.0, 5.0),
               std::default_random_engine(42), exec)),
           u_mtx(gko::test::generate_random_upper_triangular_matrix<mtx_type>(
-              3, 4, true, std::normal_distribution<real_type>(50, 5),
+              4, true, std::normal_distribution<real_type>(50, 5),
               std::normal_distribution<real_type>(20.0, 5.0),
               std::default_random_engine(42), exec)),
           lower_bandwidth(2),

cuda/test/factorization/par_ic_kernels.cpp

@@ -74,7 +74,7 @@ class ParIc : public ::testing::Test {
           cuda(gko::CudaExecutor::create(0, gko::ReferenceExecutor::create()))
     {
         mtx_l = gko::test::generate_random_lower_triangular_matrix<Csr>(
-            mtx_size[0], mtx_size[0], false,
+            mtx_size[0], false,
             std::uniform_int_distribution<index_type>(10, mtx_size[0]),
             std::normal_distribution<gko::remove_complex<value_type>>(0, 10.0),
             rand_engine, ref);

cuda/test/factorization/par_ict_kernels.cpp

@@ -78,8 +78,7 @@ class ParIct : public ::testing::Test {
             std::uniform_int_distribution<>(10, mtx_size[1]),
             std::normal_distribution<>(-1.0, 1.0), rand_engine, ref);
         mtx_l = gko::test::generate_random_lower_triangular_matrix<Csr>(
-            mtx_size[0], mtx_size[0], false,
-            std::uniform_int_distribution<>(1, mtx_size[0]),
+            mtx_size[0], false, std::uniform_int_distribution<>(1, mtx_size[0]),
             std::normal_distribution<>(-1.0, 1.0), rand_engine, ref);
 
         dmtx_ani = Csr::create(cuda);

cuda/test/factorization/par_ilut_kernels.cpp

@@ -93,25 +93,24 @@ class ParIlut : public ::testing::Test {
             std::uniform_int_distribution<>(1, mtx_size[0]),
             std::normal_distribution<>(-1.0, 1.0), rand_engine, ref);
         mtx_l = gko::test::generate_random_lower_triangular_matrix<Csr>(
-            mtx_size[0], mtx_size[0], false,
+            mtx_size[0], false,
             std::uniform_int_distribution<>(10, mtx_size[0]),
             std::normal_distribution<>(-1.0, 1.0), rand_engine, ref);
         mtx_l2 = gko::test::generate_random_lower_triangular_matrix<Csr>(
-            mtx_size[0], mtx_size[0], true,
-            std::uniform_int_distribution<>(1, mtx_size[0]),
+            mtx_size[0], true, std::uniform_int_distribution<>(1, mtx_size[0]),
             std::normal_distribution<>(-1.0, 1.0), rand_engine, ref);
         mtx_l_complex =
             gko::test::generate_random_lower_triangular_matrix<ComplexCsr>(
-                mtx_size[0], mtx_size[0], false,
+                mtx_size[0], false,
                 std::uniform_int_distribution<>(10, mtx_size[0]),
                 std::normal_distribution<>(-1.0, 1.0), rand_engine, ref);
         mtx_u = gko::test::generate_random_upper_triangular_matrix<Csr>(
-            mtx_size[0], mtx_size[0], false,
+            mtx_size[0], false,
             std::uniform_int_distribution<>(10, mtx_size[0]),
             std::normal_distribution<>(-1.0, 1.0), rand_engine, ref);
         mtx_u_complex =
             gko::test::generate_random_upper_triangular_matrix<ComplexCsr>(
-                mtx_size[0], mtx_size[0], false,
+                mtx_size[0], false,
                 std::uniform_int_distribution<>(10, mtx_size[0]),
                 std::normal_distribution<>(-1.0, 1.0), rand_engine, ref);
 

cuda/test/preconditioner/isai_kernels.cpp

@@ -110,7 +110,7 @@ class Isai : public ::testing::Test {
             mtx->copy_from(spd_mtx.get());
         } else {
             mtx = gko::test::generate_random_triangular_matrix<Csr>(
-                n, n, true, for_lower_tm, nz_dist, val_dist, rand_engine, ref,
+                n, true, for_lower_tm, nz_dist, val_dist, rand_engine, ref,
                 gko::dim<2>{n, n});
         }
         inverse = clone_allocations(mtx.get());

cuda/test/solver/lower_trs_kernels.cpp

@@ -85,17 +85,16 @@ class LowerTrs : public ::testing::Test {
             std::normal_distribution<>(-1.0, 1.0), rand_engine, ref);
     }
 
-    std::unique_ptr<Mtx> gen_l_mtx(int num_rows, int num_cols)
+    std::unique_ptr<Mtx> gen_l_mtx(int size)
     {
         return gko::test::generate_random_lower_triangular_matrix<Mtx>(
-            num_rows, num_cols, false,
-            std::uniform_int_distribution<>(num_cols, num_cols),
+            size, false, std::uniform_int_distribution<>(size),
             std::normal_distribution<>(-1.0, 1.0), rand_engine, ref);
     }
 
     void initialize_data(int m, int n)
     {
-        mtx = gen_l_mtx(m, m);
+        mtx = gen_l_mtx(m);
         b = gen_mtx(m, n);
         x = gen_mtx(m, n);
         csr_mtx = CsrMtx::create(ref);

cuda/test/solver/upper_trs_kernels.cpp

@@ -85,17 +85,16 @@ class UpperTrs : public ::testing::Test {
             std::normal_distribution<>(-1.0, 1.0), rand_engine, ref);
     }
 
-    std::unique_ptr<Mtx> gen_u_mtx(int num_rows, int num_cols)
+    std::unique_ptr<Mtx> gen_u_mtx(int size)
     {
         return gko::test::generate_random_upper_triangular_matrix<Mtx>(
-            num_rows, num_cols, false,
-            std::uniform_int_distribution<>(num_cols, num_cols),
+            size, false, std::uniform_int_distribution<>(size, size),
             std::normal_distribution<>(-1.0, 1.0), rand_engine, ref);
     }
 
     void initialize_data(int m, int n)
     {
-        mtx = gen_u_mtx(m, m);
+        mtx = gen_u_mtx(m);
         b = gen_mtx(m, n);
         x = gen_mtx(m, n);
         csr_mtx = CsrMtx::create(ref);

dpcpp/test/preconditioner/isai_kernels.cpp

@@ -115,7 +115,7 @@ class Isai : public ::testing::Test {
             mtx->copy_from(spd_mtx.get());
         } else {
             mtx = gko::test::generate_random_triangular_matrix<Csr>(
-                n, n, true, for_lower_tm, nz_dist, val_dist, rand_engine, ref,
+                n, true, for_lower_tm, nz_dist, val_dist, rand_engine, ref,
                 gko::dim<2>{n, n});
         }
         inverse = clone_allocations(mtx.get());

hip/test/factorization/par_ic_kernels.hip.cpp

@@ -74,7 +74,7 @@ class ParIc : public ::testing::Test {
           hip(gko::HipExecutor::create(0, gko::ReferenceExecutor::create()))
     {
         mtx_l = gko::test::generate_random_lower_triangular_matrix<Csr>(
-            mtx_size[0], mtx_size[0], false,
+            mtx_size[0], false,
             std::uniform_int_distribution<index_type>(10, mtx_size[0]),
             std::normal_distribution<gko::remove_complex<value_type>>(0, 10.0),
             rand_engine, ref);

hip/test/factorization/par_ict_kernels.hip.cpp

@@ -78,8 +78,7 @@ class ParIct : public ::testing::Test {
             std::uniform_int_distribution<>(10, mtx_size[1]),
             std::normal_distribution<>(-1.0, 1.0), rand_engine, ref);
         mtx_l = gko::test::generate_random_lower_triangular_matrix<Csr>(
-            mtx_size[0], mtx_size[0], false,
-            std::uniform_int_distribution<>(1, mtx_size[0]),
+            mtx_size[0], false, std::uniform_int_distribution<>(1, mtx_size[0]),
             std::normal_distribution<>(-1.0, 1.0), rand_engine, ref);
 
         dmtx_ani = Csr::create(hip);