Skip to content

Commit aba9719

Browse files
devrealdevreal
committed
First draft of sparse-tile GEMM
Added a matrix wrapper for sparse tiles and the allocator wrapper. Implement the flow of the GEMM, including symbolic analysis and the subsequent GEMM itself. Not tested yet. Signed-off-by: Joseph Schuchart <joseph.schuchart@stonybrook.edu>
1 parent d18f338 commit aba9719

3 files changed

Lines changed: 535 additions & 0 deletions

File tree

‎examples/sparse_gemm/allocator.h‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
#ifndef HAVE_ALLOCATOR_H
2+
#define HAVE_ALLOCATOR_H
3+
4+
5+
template<typename T>
6+
struct is_device_allocator : std::false_type {};
7+
8+
template<typename T>
9+
struct is_device_allocator_v : is_device_allocator<T>::value {};
10+
11+
12+
#if defined(TILEDARRAY_HAS_DEVICE)
13+
template<typename T>
14+
using Allocator = ttg::pinned_allocator_t<T>;
15+
16+
template<typename T>
17+
struct is_device_allocator<TiledArray::device_pinned_allocator<T>> : std::true_type {};
18+
19+
inline void allocator_init(int argc, char **argv) {
20+
// initialize MADNESS so that TA allocators can be created
21+
#if defined(TTG_PARSEC_IMPORTED)
22+
madness::ParsecRuntime::initialize_with_existing_context(ttg::default_execution_context().impl().context());
23+
madness::initialize(argc, argv, /* nthread = */ 1, /* quiet = */ true);
24+
#endif // TTG_PARSEC_IMPORTED
25+
}
26+
27+
inline void allocator_fini() {
28+
#if defined(TTG_PARSEC_IMPORTED)
29+
madness::finalize();
30+
#endif // TTG_PARSEC_IMPORTED
31+
}
32+
#else // TILEDARRAY_HAS_DEVICE
33+
template<typename T>
34+
using Allocator = std::allocator<T>;
35+
36+
inline void allocator_init(int argc, char **argv) { }
37+
38+
inline void allocator_fini() { }
39+
40+
#endif // TILEDARRAY_HAS_DEVICE
41+
42+
43+
#endif // HAVE_ALLOCATOR_H

‎examples/sparse_gemm/matrix.h‎

Lines changed: 216 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,216 @@
1+
#ifndef HAVE_MATRIX_H
2+
#define HAVE_MATRIX_H
3+
4+
#include "allocator.h"
5+
6+
#include <ttg.h>
7+
#include <memory>
8+
#include <vector>
9+
10+
11+
template<typename ValueT, typename IndexT, typename AllocatorT = Allocator<ValueT>>
12+
class SparseTile {
13+
ttg::Buffer<IndexT, AllocatorT> m_col_indices;
14+
ttg::Buffer<IndexT, AllocatorT> m_row_indices;
15+
ttg::Buffer<ValueT, AllocatorT> m_values;
16+
17+
public:
18+
SparseTile() = default;
19+
20+
/**
21+
* Allocate a sparse tile with the given number of columns and nonzeros.
22+
* The caller is responsible for filling in the row indices, column indices, and values.
23+
*/
24+
SparseTile(size_t num_cols, size_t nnz)
25+
: m_col_indices(num_cols)
26+
, m_row_indices(nnz)
27+
, m_values(nnz)
28+
{ }
29+
30+
/**
31+
* Construct a sparse tile from the given row indices, column indices, and values. The sizes of the vectors must be consistent.
32+
* The data is copied into the tile's buffers.
33+
*/
34+
SparseTile(const std::vector<IndexT>& row_indices, const std::vector<IndexT>& col_indices, const std::vector<ValueT>& values)
35+
: m_col_indices(col_indices.size())
36+
, m_row_indices(row_indices.size())
37+
, m_values(values.size())
38+
{
39+
std::copy_n(row_indices.data(), row_indices.size(), m_row_indices.host_ptr());
40+
std::copy_n(col_indices.data(), col_indices.size(), m_col_indices.host_ptr());
41+
std::copy_n(values.data(), values.size(), m_values.host_ptr());
42+
}
43+
44+
ttg::Buffer<IndexT, AllocatorT>& row_indices() { return m_row_indices; }
45+
ttg::Buffer<IndexT, AllocatorT>& col_indices() { return m_col_indices; }
46+
ttg::Buffer<ValueT, AllocatorT>& values() { return m_values; }
47+
48+
const ttg::Buffer<IndexT, AllocatorT>& row_indices() const { return m_row_indices; }
49+
const ttg::Buffer<IndexT, AllocatorT>& col_indices() const { return m_col_indices; }
50+
const ttg::Buffer<ValueT, AllocatorT>& values() const { return m_values; }
51+
52+
bool empty() const {
53+
return m_values.empty();
54+
}
55+
56+
template<typename Archive>
57+
void serialize(Archive& ar, const unsigned int version) {
58+
serialize(ar);
59+
}
60+
61+
template<typename Archive>
62+
void serialize(Archive& ar) {
63+
ar & m_row_indices & m_col_indices & m_values;
64+
}
65+
};
66+
67+
#ifdef TTG_SERIALIZATION_SUPPORTS_MADNESS
68+
static_assert(madness::is_serializable_v<madness::archive::BufferOutputArchive, SparseTile<float, uint64_t>>);
69+
#endif // TTG_SERIALIZATION_SUPPORTS_MADNESS
70+
71+
72+
/**
73+
* A distributed sparse matrix class that uses 2D-cyclic distribution of tiles. Each tile is a SparseTile.
74+
* Supports shallow copying and moving, but not deep copying.
75+
*/
76+
template<typename ValueT, typename IndexT, typename AllocatorT = Allocator<ValueT>>
77+
class SparseTileMatrix {
78+
public:
79+
using tile_type = SparseTile<ValueT, IndexT, AllocatorT>;
80+
using value_type = ValueT;
81+
using index_type = IndexT;
82+
using allocator_type = AllocatorT;
83+
84+
private:
85+
std::shared_ptr<std::vector<tile_type>> m_tiles;
86+
size_t m_rows;
87+
size_t m_cols;
88+
size_t m_tile_rows;
89+
size_t m_tile_cols;
90+
size_t m_pr; // process grid rows
91+
size_t m_pc; // process grid cols
92+
93+
size_t tile_index(size_t i, size_t j) const {
94+
return (i % m_tile_rows) * m_tile_cols + (j % m_tile_cols);
95+
}
96+
97+
int tile_rank(size_t i, size_t j) const {
98+
return ((i % m_pr) * m_pc) + (j % m_pc);
99+
}
100+
101+
public:
102+
SparseTileMatrix(size_t rows, size_t cols, size_t tile_rows, size_t tile_cols, size_t pr, size_t pc)
103+
: m_rows(rows)
104+
, m_cols(cols)
105+
, m_tile_rows(tile_rows)
106+
, m_tile_cols(tile_cols)
107+
, m_pr(pr)
108+
, m_pc(pc)
109+
, m_tiles(std::make_shared<std::vector<tile_type>>(tile_rows * tile_cols))
110+
{ }
111+
112+
SparseTileMatrix(size_t rows, size_t cols, size_t tile_rows, size_t tile_cols)
113+
: SparseTileMatrix(rows, cols, tile_rows, tile_cols, 0, 0)
114+
{
115+
const int mpi_size = ttg::default_execution_context().size();
116+
const int mpi_rank = ttg::default_execution_context().rank();
117+
118+
// Auto-select a roughly-quadratic P×Q factorisation of mpi_size.
119+
int P = 1, Q = mpi_size;
120+
{
121+
int best = mpi_size;
122+
for (int p = 1; p <= (int)std::sqrt((double)mpi_size); p++) {
123+
if ((mpi_size % p) == 0) {
124+
int q = mpi_size / p;
125+
if (std::abs(p - q) <= best) {
126+
best = std::abs(p - q);
127+
P = p; Q = q;
128+
}
129+
}
130+
}
131+
}
132+
if (P * Q != mpi_size) {
133+
throw std::runtime_error("Unable to auto-select process grid for given MPI size");
134+
}
135+
m_pr = P;
136+
m_pc = Q;
137+
}
138+
139+
SparseTileMatrix(const SparseTileMatrix& other) = default;
140+
SparseTileMatrix(SparseTileMatrix&& other) = default;
141+
SparseTileMatrix& operator=(const SparseTileMatrix& other) = default;
142+
SparseTileMatrix& operator=(SparseTileMatrix&& other) = default;
143+
144+
// Get tile at position (i, j) using 2D-cyclic distribution
145+
tile_type& tile(size_t i, size_t j) {
146+
size_t tile_idx = tile_index(i, j);
147+
return (*m_tiles)[tile_idx];
148+
}
149+
150+
const tile_type& tile(size_t i, size_t j) const {
151+
size_t tile_idx = tile_index(i, j);
152+
return (*m_tiles)[tile_idx];
153+
}
154+
155+
tile_type operator()(size_t i, size_t j) {
156+
return tile(i, j);
157+
}
158+
159+
const tile_type operator()(size_t i, size_t j) const {
160+
return tile(i, j);
161+
}
162+
163+
size_t num_tiles() const { return m_tiles->size(); }
164+
165+
int rank_of(size_t i, size_t j) const {
166+
return tile_rank(i, j);
167+
}
168+
169+
bool is_local(size_t i, size_t j) const {
170+
return ttg::default_execution_context().rank() == rank_of(i, j);
171+
}
172+
173+
size_t rows() const { return m_rows; }
174+
size_t cols() const { return m_cols; }
175+
size_t tile_rows() const { return m_tile_rows; }
176+
size_t tile_cols() const { return m_tile_cols; }
177+
size_t proc_rows() const { return m_pr; }
178+
size_t proc_cols() const { return m_pc; }
179+
};
180+
181+
template<typename ValueT, typename IndexT, typename AllocatorT = Allocator<ValueT>>
182+
auto make_load_tt(SparseTileMatrix<ValueT, IndexT, AllocatorT>& A, ttg::Edge<void, void> ctl, std::string name) {
183+
using tile_type = typename SparseTileMatrix<ValueT, IndexT, AllocatorT>::tile_type;
184+
ttg::Edge<Key<2>, tile_type> toop;
185+
186+
auto load_tt = ttg::make_tt(
187+
[=](){
188+
for (size_t i = 0; i < A.tile_rows(); i++) {
189+
for (size_t j = 0; j < A.tile_cols(); j++) {
190+
Key<2> key{i, j};
191+
ttg::trace("Loading tile (", i, ", ", j, ")");
192+
ttg::send<0>(key, A(i, j));
193+
}
194+
}
195+
}, ttg::edges(ctl), ttg::edges(toop), "LoadMatrix " + name, {}, {"To Op"});
196+
197+
return std::make_pair(std::move(load_tt), toop);
198+
}
199+
200+
template<typename ValueT, typename IndexT, typename AllocatorT = Allocator<ValueT>>
201+
auto make_store_tt(SparseTileMatrix<ValueT, IndexT, AllocatorT>& A,
202+
ttg::Edge<Key<2>, typename SparseTileMatrix<ValueT, IndexT, AllocatorT>::tile_type> fromop,
203+
std::string name) {
204+
using tile_type = typename SparseTileMatrix<ValueT, IndexT, AllocatorT>::tile_type;
205+
auto store_tt = ttg::make_tt(
206+
[&](const Key<2>& key, tile_type&& tile){
207+
size_t i = key[0];
208+
size_t j = key[1];
209+
ttg::trace("Storing tile (", i, ", ", j, ")");
210+
A(i, j) = std::move(tile);
211+
}, ttg::edges(fromop), ttg::edges(), "StoreMatrix " + name, {"From Op"}, {});
212+
213+
return store_tt;
214+
}
215+
216+
#endif // HAVE_MATRIX_H

0 commit comments

Comments
 (0)