Skip to content

Commit b9beed1

Browse files
devrealdevreal
committed
Add stubs and basic test harness for GETRF
Signed-off-by: Joseph Schuchart <joseph.schuchart@stonybrook.edu>
1 parent d07e51e commit b9beed1

2 files changed

Lines changed: 359 additions & 0 deletions

File tree

‎examples/getrf/getrf.h‎

Lines changed: 204 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,204 @@
1+
#ifndef GETRF_H
2+
#define GETRF_H
3+
4+
#include "ttg.h"
5+
#include "potrf/pmw.h"
6+
7+
8+
template <typename MatrixT>
9+
auto make_getrf(MatrixT& A,
10+
ttg::Edge<Key1, MatrixTile<typename MatrixT::element_type>>& from_input,
11+
ttg::Edge<Key1, MatrixTile<typename MatrixT::element_type>>& from_gemm,
12+
ttg::Edge<Key2, MatrixTile<typename MatrixT::element_type>>& to_result,
13+
ttg::Edge<Key2, MatrixTile<typename MatrixT::element_type>>& to_trsm) {
14+
ttg::Edge<Key2, MatrixTile<typename MatrixT::element_type>> output("output");
15+
using T = typename MatrixT::element_type;
16+
assert(A.cols() == A.rows());
17+
18+
auto f = [=](const Key1& key, MatrixTile<T>&& tile) TASKRET {
19+
const int K = key[0];
20+
21+
22+
/* from here we have a device selected */
23+
24+
/* do the LU factorization */
25+
26+
/* send the tile to the output */
27+
std::vector<Key2> trsm;
28+
for (int i = K+1; i < A.rows(); i++) {
29+
trsm.push_back(Key2{K, i});
30+
trsm.push_back(Key2{i, K});
31+
}
32+
33+
ttg::broadcast<0, 1>(std::make_tuple(Key2(K, K), std::move(trsm)), std::move(tile));
34+
};
35+
36+
return ttg::make_tt(std::move(f), // task function
37+
ttg::edges(ttg::fuse(from_input, from_gemm)), // input edges
38+
ttg::edges(to_result, to_trsm), // output edges
39+
"LU");
40+
}
41+
42+
43+
template <typename MatrixT>
44+
auto make_trsm(MatrixT& A,
45+
ttg::Edge<Key2, MatrixTile<typename MatrixT::element_type>>& from_getrf, // A
46+
ttg::Edge<Key2, MatrixTile<typename MatrixT::element_type>>& from_input, // B from memory
47+
ttg::Edge<Key2, MatrixTile<typename MatrixT::element_type>>& from_gemm, // B from previous gemm
48+
ttg::Edge<Key2, MatrixTile<typename MatrixT::element_type>>& to_result, // to memory
49+
ttg::Edge<Key3, MatrixTile<typename MatrixT::element_type>>& to_gemmA, // to gemm A input
50+
ttg::Edge<Key3, MatrixTile<typename MatrixT::element_type>>& to_gemmB // to gemm B input
51+
) {
52+
53+
54+
using T = typename MatrixT::element_type;
55+
auto f = [=](const Key2& key, const MatrixTile<T>& tile_A, MatrixTile<T>&& tile_B) {
56+
const int M = key[0];
57+
const int N = key[1];
58+
59+
bool is_lower = M > N;
60+
61+
std::vector<Key3> gemm;
62+
63+
/* TODO: do the trsm */
64+
65+
if (is_lower) {
66+
/* TODO: call kernel */
67+
for (int n = N+1; n < A.cols(); n++) {
68+
gemm.push_back(Key3{M, n, N});
69+
}
70+
ttg::broadcast<0, 1>(std::make_tuple(key, std::move(gemm)), std::move(tile_B));
71+
} else {
72+
/* TODO: call kernel */
73+
for (int m = M+1; m < A.rows(); m++) {
74+
gemm.push_back(Key3{m, N, M});
75+
}
76+
ttg::broadcast<0, 2>(std::make_tuple(key, std::move(gemm)), std::move(tile_B));
77+
}
78+
}
79+
80+
return ttg::make_tt(std::move(f), // task function
81+
ttg::edges(from_getrf, ttg::fuse(from_input, from_gemm)), // input edges
82+
ttg::edges(to_result, to_gemmA, to_gemmB), // output edges
83+
"TRSM");
84+
}
85+
86+
87+
template <typename MatrixT>
88+
auto make_gemm(MatrixT& A,
89+
ttg::Edge<Key3, MatrixTile<typename MatrixT::element_type>>& from_trsmA, // A from trsm
90+
ttg::Edge<Key3, MatrixTile<typename MatrixT::element_type>>& from_trsmB, // B from trsm
91+
ttg::Edge<Key3, MatrixTile<typename MatrixT::element_type>>& from_input, // C from memory
92+
ttg::Edge<Key2, MatrixTile<typename MatrixT::element_type>>& to_getrf, // A of LU
93+
ttg::Edge<Key2, MatrixTile<typename MatrixT::element_type>>& to_trsm, // to trsm
94+
) {
95+
using T = typename MatrixT::element_type;
96+
ttg::Edge<Key2, MatrixTile<typename MatrixT::element_type>> gemm_recurse;
97+
auto f = [=](const Key3& key, const MatrixTile<T>& tile_A, const MatrixTile<T>& tile_B, MatrixTile<T>&& tile_C) {
98+
const int M = key[0];
99+
const int N = key[1];
100+
const int K = key[2];
101+
102+
/* TODO: call GEMM kernel */
103+
104+
105+
std::vector<Key2> lu;
106+
std::vector<Key2> trsm;
107+
std::vector<Key2> gemm;
108+
if (M == K+1 && N == K+1) {
109+
lu.push_back(Key2{M, N});
110+
} else if (M > K+1 && N == K+1) {
111+
trsm.push_back(Key2{M, N});
112+
} else if (M == K+1 && N > K+1) {
113+
trsm.push_back(Key2{M, N});
114+
} else if (M > K+1 && N > K+1) {
115+
gemm.push_back(Key2{M, N, K+1});
116+
}
117+
ttg::broadcast<0, 1, 2>(std::make_tuple(std::move(lu), std::move(trsm), std::move(gemm)), std::move(tile_C));
118+
};
119+
120+
return ttg::make_tt(std::move(f), // task function
121+
ttg::edges(from_trsmA, from_trsmB, ttg::fuse(from_input, gemm_recurse)), // input edges
122+
ttg::edges(to_getrf, to_trsm, gemm_recurse), // output edges
123+
"GEMM");
124+
};
125+
126+
127+
auto make_getrf_ttg(MatrixT& A,
128+
ttg::Edge<Key2, MatrixTile<typename MatrixT::element_type>>& from_input,
129+
ttg::Edge<Key2, MatrixTile<typename MatrixT::element_type>>& to_result,
130+
bool enable_device_map) {
131+
132+
133+
auto keymap1 = [=](const Key1& key) { return A.rank_of(key[0], key[0]); };
134+
auto keymap2 = [=](const Key2& key) { return A.rank_of(key[0], key[1]); };
135+
auto keymap3 = [=](const Key3& key) { return A.rank_of(key[0], key[1]); };
136+
137+
138+
/**
139+
* Set a device map, 2d block-cyclic
140+
*/
141+
int num_devices = ttg::device::num_devices();
142+
int gp = std::sqrt(num_devices);
143+
int gq = (num_devices > 0) ? (num_devices / gp) : 1;
144+
auto mapper = [&A, gp,gq,num_devices](int i){
145+
auto device = (((i/A.P())%gp)*gq) + (i/A.Q())%gq;
146+
return device;
147+
};
148+
149+
auto devmap1 = [=](const Key1& key) { return mapper(key[0]); };
150+
auto devmap2 = [=](const Key2& key) { return mapper(key[0]); };
151+
auto devmap3 = [=](const Key3& key) { return mapper(key[0]); };
152+
153+
154+
ttg::Edge<Key1, MatrixTile<typename MatrixT::element_type>> dispatch_to_getrf("to_getrf");
155+
ttg::Edge<Key2, MatrixTile<typename MatrixT::element_type>> dispatch_to_trsm("to_trsm");
156+
ttg::Edge<Key3, MatrixTile<typename MatrixT::element_type>> dispatch_to_gemm("to_gemm");
157+
ttg::Edge<Key3, MatrixTile<typename MatrixT::element_type>> gemm_to_getrf("gemm_to_getrf");
158+
auto dispatcher = [=](const Key2& key, MatrixTile<typename MatrixT::element_type>& tile) {
159+
if (key[0] == 0 && key[1] == 0) {
160+
// diagonal goes to getrf
161+
ttg::send<0>(Key1{0}, tile);
162+
} else if (key[0] == 0 || key[1] == 0) {
163+
// lower goes to trsm
164+
ttg::send<1>(key, tile);
165+
} else {
166+
// rest goes to gemm
167+
ttg::send<2>(Key3{key[0], key[1], 0}, tile);
168+
}
169+
};
170+
auto dispatch_tt = ttg::make_tt(dispatcher, ttg::edges(from_input), ttg::edges(dispatch_to_getrf, dispatch_to_trsm, dispatch_to_gemm), "Dispatcher");
171+
dispatch_tt->set_keymap(keymap2);
172+
ttg::Edge<Key2, MatrixTile<typename MatrixT::element_type>> getrf_to_trsm("to_trsm");
173+
ttg::Edge<Key2, MatrixTile<typename MatrixT::element_type>> gemm_to_trsm("to_trsm");
174+
ttg::Edge<Key2, MatrixTile<typename MatrixT::element_type>> trsm_to_gemmA("to_gemmA");
175+
ttg::Edge<Key2, MatrixTile<typename MatrixT::element_type>> trsm_to_gemmB("to_gemmB");
176+
ttg::Edge<Key2, MatrixTile<typename MatrixT::element_type>> gemm_to_getrf("to_getrf");
177+
auto getrf_tt = make_getrf(A, dispatch_to_getrf, gemm_to_getrf, to_result, getrf_to_trsm);
178+
getrf_tt->set_keymap(keymap1);
179+
if (enable_device_map) {
180+
getrf_tt->set_devmap(devmap1);
181+
}
182+
auto trsm_tt = make_trsm(A, getrf_to_trsm, dispatch_to_trsm, gemm_to_trsm, to_result, trsm_to_gemmA, trsm_to_gemmB);
183+
trsm_tt->set_keymap(keymap2);
184+
if (enable_device_map) {
185+
trsm_tt->set_devmap(devmap2);
186+
}
187+
auto gemm_tt = make_gemm(A, trsm_to_gemmA, trsm_to_gemmB, dispatch_to_gemm, gemm_to_getrf, gemm_to_trsm);
188+
gemm_tt->set_keymap(keymap3);
189+
if (enable_device_map) {
190+
gemm_tt->set_devmap(devmap3);
191+
}
192+
193+
std::vector<std::unique_ptr<ttg::TTBase>> ops(4);
194+
auto ins = std::make_tuple(tt_dispatch->template in<0>());
195+
auto outs = std::make_tuple(tt_getrf->template out<0>(), tt_trsm->template out<0>());
196+
ops[0] = std::move(dispatch_tt);
197+
ops[1] = std::move(getrf_tt);
198+
ops[2] = std::move(trsm_tt);
199+
ops[3] = std::move(gemm_tt);
200+
201+
return make_ttg(std::move(ops), ins, outs, "GETRF TTG");
202+
}
203+
204+
#endif // GETRF_H

‎examples/getrf/testing_getrf.cc‎

Lines changed: 155 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,155 @@
1+
#include "getrf.h"
2+
3+
4+
5+
6+
int main(int argc, char **argv) {
7+
8+
9+
std::chrono::time_point<std::chrono::high_resolution_clock> beg, end;
10+
int NB = 32;
11+
int N = 5*NB;
12+
int M = N;
13+
int nthreads = -1;
14+
const char* prof_filename = nullptr;
15+
char *opt = nullptr;
16+
int ret = EXIT_SUCCESS;
17+
int niter = 3;
18+
bool print_dot = false;
19+
20+
if( (opt = getCmdOption(argv+1, argv+argc, "-N")) != nullptr ) {
21+
N = M = atoi(opt);
22+
}
23+
24+
if( (opt = getCmdOption(argv+1, argv+argc, "-t")) != nullptr ) {
25+
NB = atoi(opt);
26+
}
27+
28+
if( (opt = getCmdOption(argv+1, argv+argc, "-c")) != nullptr ) {
29+
nthreads = atoi(opt);
30+
}
31+
32+
if( (opt = getCmdOption(argv+1, argv+argc, "-dag")) != nullptr ) {
33+
prof_filename = opt;
34+
}
35+
36+
if( (opt = getCmdOption(argv+1, argv+argc, "-n")) != nullptr) {
37+
niter = atoi(opt);
38+
}
39+
40+
/* whether to print the TTG dot */
41+
print_dot = cmdOptionExists(argv+1, argv+argc, "-dot");
42+
43+
44+
bool check = !cmdOptionExists(argv+1, argv+argc, "-x");
45+
46+
/* whether we set a device mapping */
47+
bool enable_device_map = !cmdOptionExists(argv, argv+argc, "--default-device-map");
48+
49+
// TODO: need to filter out our arguments to make parsec happy
50+
ttg::initialize(1, argv, nthreads);
51+
52+
/* set up TA to get the allocator */
53+
allocator_init(argc, argv);
54+
55+
auto world = ttg::default_execution_context();
56+
if(nullptr != prof_filename) {
57+
world.profile_on();
58+
world.dag_on(prof_filename);
59+
}
60+
61+
int P = std::sqrt(world.size());
62+
int Q = (world.size() + P - 1)/P;
63+
64+
if ( (opt = getCmdOption(argv+1, argv + argc, "-P")) != nullptr) {
65+
P = atoi(opt);
66+
}
67+
if ( (opt = getCmdOption(argv+1, argv + argc, "-Q")) != nullptr) {
68+
Q = atoi(opt);
69+
}
70+
71+
72+
if(check && (P>1 || Q>1)) {
73+
std::cerr << "Check is disabled for distributed runs at this time" << std::endl;
74+
check = false;
75+
}
76+
77+
if (world.rank() == 0) {
78+
std::cout << "Creating 2D block cyclic matrix with NB " << NB << " N " << N << " M " << M << " P " << P << " Q " << Q << std::endl;
79+
}
80+
81+
parsec_matrix_block_cyclic_t dcA;
82+
parsec_matrix_block_cyclic_init(&dcA, parsec_matrix_type_t::PARSEC_MATRIX_DOUBLE,
83+
world.rank(), NB, NB, N, M,
84+
0, 0, N, M, P, Q, PARSEC_MATRIX_LOWER);
85+
dcA.mat = parsec_data_allocate((size_t)dcA.super.nb_local_tiles *
86+
(size_t)dcA.super.bsiz *
87+
(size_t)parsec_datadist_getsizeoftype(dcA.super.mtype));
88+
89+
/* would be nice to have proper abstractions for this */
90+
parsec_data_collection_t *o = &(dcA.super.super);
91+
for (int devid = 1; devid < parsec_nb_devices; ++devid) {
92+
auto* device = parsec_mca_device_get(devid);
93+
if (device->memory_register) {
94+
o->register_memory(o, device); // TODO: check device IDs
95+
}
96+
}
97+
98+
parsec_data_collection_set_key((parsec_data_collection_t*)&dcA, (char*)"Matrix A");
99+
100+
101+
for (int i = 0; i < niter; ++i) {
102+
parsec_devices_release_memory();
103+
104+
//Matrix<double>* A = new Matrix<double>(n_rows, n_cols, NB, NB);
105+
MatrixT<double> A{&dcA};
106+
/* TODO: initialize the matrix */
107+
/* This works only with the parsec backend! */
108+
int random_seed = 3872;
109+
110+
init_matrix(A, random_seed);
111+
ttg::Edge<Key2, MatrixTile<double>> startup("startup");
112+
ttg::Edge<Key2, MatrixTile<double>> result("To result");
113+
114+
auto getrf_init_tt = make_load_tt(A, startup);
115+
auto getrf_ttg = make_getrf_ttg(A, startup, result, enable_device_map);
116+
auto getrf_result_ttg = make_result_ttg(A, result);
117+
118+
auto connected = make_graph_executable(getrf_init_tt.get());
119+
assert(connected);
120+
TTGUNUSED(connected);
121+
122+
if (world.rank() == 0) {
123+
if (print_dot) {
124+
std::cout << "==== begin dot ====\n";
125+
std::cout << ttg::Dot()(getrf_init_tt.get()) << std::endl;
126+
std::cout << "==== end dot ====\n";
127+
}
128+
beg = std::chrono::high_resolution_clock::now();
129+
}
130+
131+
if (world.rank() == 0) {
132+
beg = std::chrono::high_resolution_clock::now();
133+
}
134+
135+
getrf_init_tt->invoke();
136+
ttg::execute(world);
137+
ttg::fence(world);
138+
139+
if (world.rank() == 0) {
140+
end = std::chrono::high_resolution_clock::now();
141+
auto elapsed = (std::chrono::duration_cast<std::chrono::microseconds>(end - beg).count());
142+
end = std::chrono::high_resolution_clock::now();
143+
std::cout << "TTG Execution Time (milliseconds) : "
144+
<< elapsed / 1E3 << " : Flops " << (potrf::FLOPS_DPOTRF(N)) << " " << (potrf::FLOPS_DPOTRF(N)/1e9)/(elapsed/1e6) << " GF/s" << std::endl;
145+
}
146+
#if defined(TTG_PARSEC_IMPORTED)
147+
parsec_devices_reset_load(ttg::default_execution_context().impl().context());
148+
#endif // TTG_PARSEC_IMPORTED
149+
}
150+
151+
world.dag_off();
152+
153+
ttg::finalize();
154+
155+
}

0 commit comments

Comments
 (0)