Skip to content

Commit

Permalink
introduced ttg::matrix::Triplet (replacement for Eigen::Triplet) to e…
Browse files Browse the repository at this point in the history 
…nable moves
  • Loading branch information
@evaleev
evaleev committed Jul 24, 2024
1 parent 4c6cc15 commit 2a8621c
Show file tree
Hide file tree
Showing 3 changed files with 79 additions and 41 deletions.
11 changes: 6 additions & 5 deletions examples/spmm/spmm.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
#endif

#include "ttg.h"
#include "../ttg_matrix.h"

using namespace ttg;

Expand Down Expand Up @@ -103,7 +104,7 @@ using blk_t = double;
template <typename T = blk_t>
using SpMatrix = Eigen::SparseMatrix<T>;
template <typename T = blk_t>
using SpMatrixTriplet = Eigen::Triplet<T>; // {row,col,value}
using SpMatrixTriplet = ttg::matrix::Triplet<T>; // {row,col,value}

#if defined(BLOCK_SPARSE_GEMM) && defined(BTAS_IS_USABLE)

Expand Down Expand Up @@ -990,7 +991,7 @@ static void initSpRmat(const std::function<int(const Key<2> &)> &keymap, const c
boost::minstd_rand gen(seed);
boost::rmat_iterator<boost::minstd_rand, boost::directed_graph<>> rmat_it(gen, N, E, a, b, c, d);

using triplet_t = Eigen::Triplet<blk_t>;
using triplet_t = ttg::matrix::Triplet<blk_t>;
std::vector<triplet_t> A_elements;
for (int i = 0; i < N; i++) {
nnz++;
Expand Down Expand Up @@ -1026,7 +1027,7 @@ static void initSpHardCoded(const std::function<int(const Key<2> &)> &keymap, Sp
C.resize(m, n);
// We initialize the same matrices on all the ranks, but we will use only the local part
// following the keymap
using triplet_t = Eigen::Triplet<blk_t>;
using triplet_t = ttg::matrix::Triplet<blk_t>;
std::vector<triplet_t> A_elements;
A_elements.emplace_back(0, 1, 12.3);
A_elements.emplace_back(0, 2, 10.7);
Expand Down Expand Up @@ -1073,7 +1074,7 @@ static void initBlSpHardCoded(const std::function<int(const Key<2> &)> &keymap,

int rank = ttg::default_execution_context().rank();

using triplet_t = Eigen::Triplet<blk_t>;
using triplet_t = ttg::matrix::Triplet<blk_t>;
std::vector<triplet_t> A_elements;
std::vector<triplet_t> Aref_elements;
#if defined(BTAS_IS_USABLE)
Expand Down Expand Up @@ -1239,7 +1240,7 @@ static void initBlSpRandom(const std::function<int(const Key<2> &)> &keymap, siz
std::mt19937 genv(seed + 1);

std::uniform_int_distribution<> dist(minTs, maxTs); // randomly pick any value in the range minTs, maxTs
using triplet_t = Eigen::Triplet<blk_t>;
using triplet_t = ttg::matrix::Triplet<blk_t>;
std::vector<triplet_t> A_elements;
std::vector<triplet_t> B_elements;
std::vector<triplet_t> Aref_elements;
Expand Down
12 changes: 7 additions & 5 deletions examples/spmm/spmm_cuda.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,8 @@ using namespace ttg;
#include <sys/time.h>
#endif

#include "../ttg_matrix.h"

#if defined(BLOCK_SPARSE_GEMM) && defined(BTAS_IS_USABLE)

template <typename _T, class _Range, class _Storage>
Expand Down Expand Up @@ -399,7 +401,7 @@ using blk_t = double;
template <typename T = blk_t>
using SpMatrix = Eigen::SparseMatrix<T>;
template <typename T = blk_t>
using SpMatrixTriplet = Eigen::Triplet<T>; // {row,col,value}
using SpMatrixTriplet = ttg::matrix::Triplet<T>; // {row,col,value}

#if defined(BLOCK_SPARSE_GEMM) && defined(BTAS_IS_USABLE)

Expand Down Expand Up @@ -1197,7 +1199,7 @@ static void initSpRmat(const std::function<int(const Key<2> &)> &keymap, const c
boost::minstd_rand gen(seed);
boost::rmat_iterator<boost::minstd_rand, boost::directed_graph<>> rmat_it(gen, N, E, a, b, c, d);

using triplet_t = Eigen::Triplet<blk_t>;
using triplet_t = ttg::matrix::Triplet<blk_t>;
std::vector<triplet_t> A_elements;
for (int i = 0; i < N; i++) {
nnz++;
Expand Down Expand Up @@ -1232,7 +1234,7 @@ static void initSpHardCoded(const std::function<int(const Key<2> &)> &keymap, Sp
C.resize(m, n);
// We initialize the same matrices on all the ranks, but we will use only the local part
// following the keymap
using triplet_t = Eigen::Triplet<blk_t>;
using triplet_t = ttg::matrix::Triplet<blk_t>;
std::vector<triplet_t> A_elements;
A_elements.emplace_back(0, 1, 12.3);
A_elements.emplace_back(0, 2, 10.7);
Expand Down Expand Up @@ -1279,7 +1281,7 @@ static void initBlSpHardCoded(const std::function<int(const Key<2> &)> &keymap,

int rank = ttg::default_execution_context().rank();

using triplet_t = Eigen::Triplet<blk_t>;
using triplet_t = ttg::matrix::Triplet<blk_t>;
std::vector<triplet_t> A_elements;
std::vector<triplet_t> Aref_elements;
#if defined(BTAS_IS_USABLE)
Expand Down Expand Up @@ -1446,7 +1448,7 @@ static void initBlSpRandom(const std::function<int(const Key<2> &)> &keymap, siz
std::mt19937 genv(seed + 1);

std::uniform_int_distribution<> dist(minTs, maxTs); // randomly pick any value in the range minTs, maxTs
using triplet_t = Eigen::Triplet<blk_t>;
using triplet_t = ttg::matrix::Triplet<blk_t>;
std::vector<triplet_t> A_elements;
std::vector<triplet_t> B_elements;
std::vector<triplet_t> Aref_elements;
Expand Down
97 changes: 66 additions & 31 deletions examples/ttg_matrix.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,11 +8,46 @@
#include <vector>

#include "ttg/serialization/std/vector.h"
#include "ttg/util/multiindex.h"

#include <Eigen/SparseCore>

namespace ttg {

namespace matrix {

/// element of a sparse matrix = {row index, col index, value}

/// movable replacement for Eigen::Triplet
template<typename Value, typename StorageIndex=typename Eigen::SparseMatrix<Value>::StorageIndex >
class Triplet {
public:
Triplet() = default;
Triplet(const Triplet&) = default;
Triplet(Triplet&&) = default;
Triplet& operator=(const Triplet&) = default;
Triplet& operator=(Triplet&&) = default;

Triplet(StorageIndex r, const StorageIndex c, const Value& v)
: m_row(r), m_col(c), m_value(v)
{}
Triplet(StorageIndex r, const StorageIndex c, Value&& v = Value{})
: m_row(r), m_col(c), m_value(std::move(v))
{}

/** \returns the row index of the element */
const StorageIndex& row() const { return m_row; }

/** \returns the column index of the element */
const StorageIndex& col() const { return m_col; }

/** \returns the value of the element */
const Value& value() const { return m_value; }
protected:
StorageIndex m_row = -1, m_col = -1;
Value m_value;
};

// matrix shape = maps {column,row} index to {row,column} indices
class Shape : public std::vector<std::vector<long>> {
using base_t = std::vector<std::vector<long>>;
Expand Down Expand Up @@ -192,22 +227,22 @@ namespace ttg {
}

// compute shape of an existing SpMatrix on rank 0
template <typename Blk = blk_t>
template <typename Blk>
class ReadShape : public TT<void, std::tuple<Out<void, Shape>>, ReadShape<Blk>, ttg::typelist<void>> {
public:
using baseT = typename ReadShape::ttT;
static constexpr const int owner = 0; // where data resides

ReadShape(const char *label, const SpMatrix<Blk> &matrix, Edge<void, void> &in, Edge<void, Shape> &out)
ReadShape(const char *label, const Eigen::SparseMatrix<Blk> &matrix, Edge<void, void> &in, Edge<void, Shape> &out)
: baseT(edges(in), edges(out), std::string("read_spmatrix_shape(") + label + ")", {"ctl"},
{std::string("shape[") + label + "]"},
/* keymap */ []() { return owner; })
, matrix_(matrix) {}

void op(std::tuple<Out<void, Shape>> &out) { ::sendv<0>(Shape(matrix_), out); }
void op(std::tuple<Out<void, Shape>> &out) { ttg::sendv<0>(Shape(matrix_), out); }

private:
const SpMatrix<Blk> &matrix_;
const Eigen::SparseMatrix<Blk> &matrix_;
};

// flow data from an existing SpMatrix on rank 0
Expand All @@ -217,67 +252,67 @@ namespace ttg {
// but will only be efficient if can do random access (slow with CSC format used by Eigen matrices)
// - this could be generalized to read efficiently from a distributed data structure
// Use Read_SpMatrix if need to read all data from a data structure localized on 1 process
template <typename Blk = blk_t>
class Read : public TT<Key<2>, std::tuple<Out<Key<2>, Blk>>, Read<Blk>, ttg::typelist<void>> {
template <typename Blk>
class Read : public TT<MultiIndex<2>, std::tuple<Out<MultiIndex<2>, Blk>>, Read<Blk>, ttg::typelist<void>> {
public:
using baseT = TT<Key<2>, std::tuple<Out<Key<2>, Blk>>, Read<Blk>, void>;
using baseT = TT<MultiIndex<2>, std::tuple<Out<MultiIndex<2>, Blk>>, Read<Blk>, void>;
static constexpr const int owner = 0; // where data resides

Read(const char *label, const SpMatrix<Blk> &matrix, Edge<Key<2>, void> &in, Edge<Key<2>, Blk> &out)
Read(const char *label, const Eigen::SparseMatrix<Blk> &matrix, Edge<MultiIndex<2>, void> &in, Edge<MultiIndex<2>, Blk> &out)
: baseT(edges(in), edges(out), std::string("read_spmatrix(") + label + ")", {"ctl[ij]"},
{std::string(label) + "[ij]"},
/* keymap */ [](auto key) { return owner; })
, matrix_(matrix) {}

void op(const Key<2> &key, std::tuple<Out<Key<2>, Blk>> &out) {
void op(const MultiIndex<2> &key, std::tuple<Out<MultiIndex<2>, Blk>> &out) {
// random access in CSC format is inefficient, this is only to demonstrate the way to go for hash-based storage
// for whatever reason coeffRef does not work on a const SpMatrix&
::send<0>(key, static_cast<const Blk &>(const_cast<SpMatrix<Blk> &>(matrix_).coeffRef(key[0], key[1])), out);
ttg::send<0>(key, static_cast<const Blk &>(const_cast<Eigen::SparseMatrix<Blk> &>(matrix_).coeffRef(key[0], key[1])), out);
}

private:
const SpMatrix<Blk> &matrix_;
const Eigen::SparseMatrix<Blk> &matrix_;
};

// WriteShape commits shape to an existing SpMatrix on rank 0 and sends it on
// since SpMatrix supports random inserts there is no need to commit the shape into the matrix, other than get the
// dimensions
template <typename Blk = blk_t>
template <typename Blk>
class WriteShape : public TT<void, std::tuple<Out<void, Shape>>, WriteShape<Blk>, ttg::typelist<Shape>> {
public:
using baseT = typename WriteShape::ttT;
static constexpr const int owner = 0; // where data resides

WriteShape(const char *label, SpMatrix<Blk> &matrix, Edge<void, Shape> &in, Edge<void, Shape> &out)
WriteShape(const char *label, Eigen::SparseMatrix<Blk> &matrix, Edge<void, Shape> &in, Edge<void, Shape> &out)
: baseT(edges(in), edges(out), std::string("write_spmatrix_shape(") + label + ")",
{std::string("shape_in[") + label + "]"}, {std::string("shape_out[") + label + "]"},
/* keymap */ []() { return owner; })
, matrix_(matrix) {}

void op(typename baseT::input_values_tuple_type &&ins, std::tuple<Out<void, Shape>> &out) {
const auto &shape = baseT::template get<0>(ins);
::ttg::trace("Resizing ", static_cast<void *>(&matrix_));
ttg::trace("Resizing ", static_cast<void *>(&matrix_));
matrix_.resize(shape.nrows(), shape.ncols());
::sendv<0>(shape, out);
ttg::sendv<0>(shape, out);
}

private:
SpMatrix<Blk> &matrix_;
Eigen::SparseMatrix<Blk> &matrix_;
};

// flow (move?) data into an existing SpMatrix on rank 0
template <typename Blk = blk_t>
class Write : public TT<Key<2>, std::tuple<>, Write<Blk>, Blk, ttg::typelist<void>> {
template <typename Blk>
class Write : public TT<MultiIndex<2>, std::tuple<>, Write<Blk>, ttg::typelist<void>> {
public:
using baseT = typename Write::ttT;

Write(const char *label, SpMatrix<Blk> &matrix, Edge<Key<2>, Blk> &data_in, Edge<Key<2>, void> &ctl_in)
Write(const char *label, Eigen::SparseMatrix<Blk> &matrix, Edge<MultiIndex<2>, Blk> &data_in, Edge<MultiIndex<2>, void> &ctl_in)
: baseT(edges(data_in, ctl_in), edges(), std::string("write_spmatrix(") + label + ")",
{std::string(label) + "[ij]", std::string("ctl[ij]")}, {},
/* keymap */ [](auto key) { return 0; })
, matrix_(matrix) {}

void op(const Key<2> &key, typename baseT::input_values_tuple_type &&elem, std::tuple<> &) {
void op(const MultiIndex<2> &key, typename baseT::input_values_tuple_type &&elem, std::tuple<> &) {
std::lock_guard<std::mutex> lock(mtx_);
ttg::trace("rank =", default_execution_context().rank(),
"/ thread_id =", reinterpret_cast<std::uintptr_t>(pthread_self()),
Expand Down Expand Up @@ -309,8 +344,8 @@ namespace ttg {

private:
std::mutex mtx_;
SpMatrix<Blk> &matrix_;
std::vector<SpMatrixTriplet<Blk>> values_;
Eigen::SparseMatrix<Blk> &matrix_;
std::vector<Triplet<Blk>> values_;
mutable std::shared_ptr<std::shared_future<void>> completion_status_;
};

Expand All @@ -325,36 +360,36 @@ namespace ttg {
/* keymap */ []() { return owner; }) {}

void op(typename baseT::input_values_tuple_type &&ins, std::tuple<Out<void, Shape>> &out) {
::sendv<0>(Shape::add(baseT::template get<0>(ins), baseT::template get<1>(ins)), out);
ttg::sendv<0>(Shape::add(baseT::template get<0>(ins), baseT::template get<1>(ins)), out);
}
};

// pushes all blocks given by the shape
class Push : public TT<void, std::tuple<Out<Key<2>, void>>, Push, ttg::typelist<Shape>> {
class Push : public TT<void, std::tuple<Out<MultiIndex<2>, void>>, Push, ttg::typelist<Shape>> {
public:
using baseT = typename Push::ttT;
static constexpr const int owner = 0; // where data resides

Push(const char *label, Edge<void, Shape> &in, Edge<Key<2>, void> &out)
Push(const char *label, Edge<void, Shape> &in, Edge<MultiIndex<2>, void> &out)
: baseT(edges(in), edges(out), std::string("push_spmatrix(") + label + ")",
{std::string("shape[") + label + "]"}, {"ctl[ij]"},
/* keymap */ []() { return owner; }) {}

void op(typename baseT::input_values_tuple_type &&ins, std::tuple<Out<Key<2>, void>> &out) {
void op(typename baseT::input_values_tuple_type &&ins, std::tuple<Out<MultiIndex<2>, void>> &out) {
const auto &shape = baseT::get<0>(ins);
if (shape.type() == Shape::Type::col2row) {
long colidx = 0;
for (const auto &col : shape) {
for (const auto rowidx : col) {
::sendk<0>(Key<2>({rowidx, colidx}), out);
ttg::sendk<0>(MultiIndex<2>({rowidx, colidx}), out);
}
++colidx;
}
} else if (shape.type() == Shape::Type::row2col) {
long rowidx = 0;
for (const auto &row : shape) {
for (const auto colidx : row) {
::sendk<0>(Key<2>({rowidx, colidx}), out);
ttg::sendk<0>(MultiIndex<2>({rowidx, colidx}), out);
}
++rowidx;
}
Expand All @@ -372,9 +407,9 @@ namespace ttg {
class Matrix {
public:
using shape_t = matrix::Shape;
using data_edge_t = Edge<Key<2>, T>;
using data_edge_t = Edge<MultiIndex<2>, T>;
using shape_edge_t = Edge<void, shape_t>;
using ctl_edge_t = Edge<Key<2>, void>;
using ctl_edge_t = Edge<MultiIndex<2>, void>;

Matrix() = default;

Expand Down Expand Up @@ -403,7 +438,7 @@ namespace ttg {
/// @return an std::future<void> object indicating the status; @c destination_matrix is ready if calling has_value()
/// on the return value of this function is true.
/// @note up to the user to ensure completion before reading destination_matrix
auto operator>>(SpMatrix<T> &destination_matrix) {
auto operator>>(Eigen::SparseMatrix<T> &destination_matrix) {
// shape writer writes shape to destination_matrix
ttg_register_ptr(world_, std::make_shared<matrix::WriteShape<T>>("Matrix.WriteShape", destination_matrix,
shape_edge_, shape_writer_push_edge_));
Expand Down

0 comments on commit 2a8621c

Please sign in to comment.