mm.c 9.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372
  1. /* StarPU --- Runtime system for heterogeneous multicore architectures.
  2. *
  3. * Copyright (C) 2016 Inria
  4. *
  5. * StarPU is free software; you can redistribute it and/or modify
  6. * it under the terms of the GNU Lesser General Public License as published by
  7. * the Free Software Foundation; either version 2.1 of the License, or (at
  8. * your option) any later version.
  9. *
  10. * StarPU is distributed in the hope that it will be useful, but
  11. * WITHOUT ANY WARRANTY; without even the implied warranty of
  12. * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
  13. *
  14. * See the GNU Lesser General Public License in COPYING.LGPL for more details.
  15. */
  16. /*
  17. * This example illustrates how to distribute a pre-existing data structure to
  18. * a set of computing nodes using StarPU-MPI routines.
  19. */
  20. #include <stdlib.h>
  21. #include <stdio.h>
  22. #include <assert.h>
  23. #include <math.h>
  24. #include <starpu.h>
  25. #include <starpu_mpi.h>
  26. #include "helper.h"
  27. #define VERBOSE 0
  28. static int N = 16; /* Matrix size */
  29. static int BS = 4; /* Block size */
  30. #define NB ((N)/(BS)) /* Number of blocks */
  31. /* Matrices. Will be allocated as regular, linearized C arrays */
  32. static double *A = NULL; /* A will be partitioned as BS rows x N cols blocks */
  33. static double *B = NULL; /* B will be partitioned as N rows x BS cols blocks */
  34. static double *C = NULL; /* C will be partitioned as BS rows x BS cols blocks */
  35. /* Arrays of data handles for managing matrix blocks */
  36. static starpu_data_handle_t *A_h;
  37. static starpu_data_handle_t *B_h;
  38. static starpu_data_handle_t *C_h;
  39. static int comm_rank; /* mpi rank of the process */
  40. static int comm_size; /* size of the mpi session */
  41. static void alloc_matrices(void)
  42. {
  43. /* Regular 'malloc' can also be used instead, however, starpu_malloc make sure that
  44. * the area is allocated in suitably pinned memory to improve data transfers, especially
  45. * with CUDA */
  46. starpu_malloc((void **)&A, N*N*sizeof(double));
  47. starpu_malloc((void **)&B, N*N*sizeof(double));
  48. starpu_malloc((void **)&C, N*N*sizeof(double));
  49. }
  50. static void free_matrices(void)
  51. {
  52. starpu_free(A);
  53. starpu_free(B);
  54. starpu_free(C);
  55. }
  56. static void init_matrices(void)
  57. {
  58. int row,col;
  59. for (row = 0; row < N; row++)
  60. {
  61. for (col = 0; col < N; col++)
  62. {
  63. A[row*N+col] = (row==col)?2:0;
  64. B[row*N+col] = row*N+col;
  65. C[row*N+col] = 0;
  66. }
  67. }
  68. }
  69. #if VERBOSE
  70. static void disp_matrix(double *m)
  71. {
  72. int row,col;
  73. for (row = 0; row < N; row++)
  74. {
  75. for (col = 0; col < N; col++)
  76. {
  77. printf("\t%.2lf", m[row*N+col]);
  78. }
  79. printf("\n");
  80. }
  81. }
  82. #endif
  83. static void check_result(void) {
  84. int row,col;
  85. for (row = 0; row < N; row++)
  86. {
  87. for (col = 0; col < N; col++)
  88. {
  89. if (fabs(C[row*N+col] - 2*(row*N+col)) > 1.0)
  90. {
  91. fprintf(stderr, "check failed\n");
  92. exit(1);
  93. }
  94. }
  95. }
  96. #if VERBOSE
  97. printf("success\n");
  98. #endif
  99. }
  100. /* Register the matrix blocks to StarPU and to StarPU-MPI */
  101. static void register_matrices()
  102. {
  103. A_h = calloc(NB, sizeof(starpu_data_handle_t));
  104. B_h = calloc(NB, sizeof(starpu_data_handle_t));
  105. C_h = calloc(NB*NB, sizeof(starpu_data_handle_t));
  106. /* Memory region, where the data being registered resides.
  107. * In this example, all blocks are allocated by node 0, thus
  108. * - node 0 specifies STARPU_MAIN_RAM to indicate that it owns the block in its main memory
  109. * - nodes !0 specify -1 to indicate that they don't have a copy of the block initially
  110. */
  111. int mr = (comm_rank == 0) ? STARPU_MAIN_RAM : -1;
  112. /* mpi tag used for the block */
  113. int tag = 0;
  114. int b_row,b_col;
  115. for (b_row = 0; b_row < NB; b_row++) {
  116. /* Register a block to StarPU */
  117. starpu_matrix_data_register(&A_h[b_row],
  118. mr,
  119. (comm_rank == 0)?(uintptr_t)(A+b_row*BS*N):0, N, N, BS,
  120. sizeof(double));
  121. /* Register a block to StarPU-MPI, specifying the mpi tag to use for transfering the block
  122. * and the rank of the owner node.
  123. *
  124. * Note: StarPU-MPI is an autonomous layer built on top of StarPU, hence the two separate
  125. * registration steps.
  126. */
  127. starpu_data_set_coordinates(A_h[b_row], 2, 0, b_row);
  128. starpu_mpi_data_register(A_h[b_row], tag++, 0);
  129. }
  130. for (b_col = 0; b_col < NB; b_col++) {
  131. starpu_matrix_data_register(&B_h[b_col],
  132. mr,
  133. (comm_rank == 0)?(uintptr_t)(B+b_col*BS):0, N, BS, N,
  134. sizeof(double));
  135. starpu_data_set_coordinates(B_h[b_col], 2, b_col, 0);
  136. starpu_mpi_data_register(B_h[b_col], tag++, 0);
  137. }
  138. for (b_row = 0; b_row < NB; b_row++) {
  139. for (b_col = 0; b_col < NB; b_col++) {
  140. starpu_matrix_data_register(&C_h[b_row*NB+b_col],
  141. mr,
  142. (comm_rank == 0)?(uintptr_t)(C+b_row*BS*N+b_col*BS):0, N, BS, BS,
  143. sizeof(double));
  144. starpu_data_set_coordinates(C_h[b_row*NB+b_col], 2, b_col, b_row);
  145. starpu_mpi_data_register(C_h[b_row*NB+b_col], tag++, 0);
  146. }
  147. }
  148. }
  149. /* Transfer ownership of the C matrix blocks following some user-defined distribution over the nodes.
  150. * Note: since C will be Write-accessed, it will implicitly define which node perform the task
  151. * associated to a given block. */
  152. static void distribute_matrix_C(void)
  153. {
  154. int b_row,b_col;
  155. for (b_row = 0; b_row < NB; b_row++)
  156. {
  157. for (b_col = 0; b_col < NB; b_col++)
  158. {
  159. starpu_data_handle_t h = C_h[b_row*NB+b_col];
  160. /* Select the node where the block should be computed. */
  161. int target_rank = (b_row+b_col)%comm_size;
  162. /* Move the block on to its new owner. */
  163. starpu_mpi_data_migrate(MPI_COMM_WORLD, h, target_rank);
  164. }
  165. }
  166. }
  167. /* Transfer ownership of the C matrix blocks back to node 0, for display purpose. This is not mandatory. */
  168. static void undistribute_matrix_C(void)
  169. {
  170. int b_row,b_col;
  171. for (b_row = 0; b_row < NB; b_row++) {
  172. for (b_col = 0; b_col < NB; b_col++) {
  173. starpu_data_handle_t h = C_h[b_row*NB+b_col];
  174. starpu_mpi_data_migrate(MPI_COMM_WORLD, h, 0);
  175. }
  176. }
  177. }
  178. /* Unregister matrices from the StarPU management. */
  179. static void unregister_matrices()
  180. {
  181. int b_row,b_col;
  182. for (b_row = 0; b_row < NB; b_row++) {
  183. starpu_data_unregister(A_h[b_row]);
  184. }
  185. for (b_col = 0; b_col < NB; b_col++) {
  186. starpu_data_unregister(B_h[b_col]);
  187. }
  188. for (b_row = 0; b_row < NB; b_row++) {
  189. for (b_col = 0; b_col < NB; b_col++) {
  190. starpu_data_unregister(C_h[b_row*NB+b_col]);
  191. }
  192. }
  193. free(A_h);
  194. free(B_h);
  195. free(C_h);
  196. }
  197. /* Perform the actual computation. In a real-life case, this would rather call a BLAS 'gemm' routine
  198. * instead. */
  199. static void cpu_mult(void *handles[], STARPU_ATTRIBUTE_UNUSED void *arg)
  200. {
  201. double *block_A = (double *)STARPU_MATRIX_GET_PTR(handles[0]);
  202. double *block_B = (double *)STARPU_MATRIX_GET_PTR(handles[1]);
  203. double *block_C = (double *)STARPU_MATRIX_GET_PTR(handles[2]);
  204. unsigned n_col_A = STARPU_MATRIX_GET_NX(handles[0]);
  205. unsigned n_col_B = STARPU_MATRIX_GET_NX(handles[1]);
  206. unsigned n_col_C = STARPU_MATRIX_GET_NX(handles[2]);
  207. unsigned n_row_A = STARPU_MATRIX_GET_NY(handles[0]);
  208. unsigned n_row_B = STARPU_MATRIX_GET_NY(handles[1]);
  209. unsigned n_row_C = STARPU_MATRIX_GET_NY(handles[2]);
  210. unsigned ld_A = STARPU_MATRIX_GET_LD(handles[0]);
  211. unsigned ld_B = STARPU_MATRIX_GET_LD(handles[1]);
  212. unsigned ld_C = STARPU_MATRIX_GET_LD(handles[2]);
  213. /* Sanity check, not needed in real life case */
  214. assert(n_col_C == n_col_B);
  215. assert(n_row_C == n_row_A);
  216. assert(n_col_A == n_row_B);
  217. unsigned i,j,k;
  218. for (k = 0; k < n_row_C; k++) {
  219. for (j = 0; j < n_col_C; j++) {
  220. for (i = 0; i < n_col_A; i++) {
  221. block_C[k*ld_C+j] += block_A[k*ld_A+i] * block_B[i*ld_B+j];
  222. }
  223. #if VERBOSE
  224. /* For illustration purpose, shows which node computed
  225. * the block in the decimal part of the cell */
  226. block_C[k*ld_C+j] += comm_rank / 100.0;
  227. #endif
  228. }
  229. }
  230. }
  231. /* Define a StarPU 'codelet' structure for the matrix multiply kernel above.
  232. * This structure enable specifying multiple implementations for the kernel (such as CUDA or OpenCL versions)
  233. */
  234. static struct starpu_codelet gemm_cl =
  235. {
  236. .cpu_funcs = {cpu_mult}, /* cpu implementation(s) of the routine */
  237. .nbuffers = 3, /* number of data handles referenced by this routine */
  238. .modes = {STARPU_R, STARPU_R, STARPU_RW} /* access modes for each data handle */
  239. };
  240. int main(int argc, char *argv[])
  241. {
  242. /* Initializes the StarPU core */
  243. int ret = starpu_init(NULL);
  244. STARPU_CHECK_RETURN_VALUE(ret, "starpu_init");
  245. /* Initializes the StarPU-MPI layer */
  246. ret = starpu_mpi_init(&argc, &argv, 1);
  247. STARPU_CHECK_RETURN_VALUE(ret, "starpu_mpi_init");
  248. if (starpu_cpu_worker_get_count() == 0)
  249. {
  250. FPRINTF(stderr, "We need at least 1 CPU worker.\n");
  251. starpu_mpi_shutdown();
  252. starpu_shutdown();
  253. return STARPU_TEST_SKIPPED;
  254. }
  255. /* Parse the matrix size and block size optional args */
  256. if (argc > 1) {
  257. N = atoi(argv[1]);
  258. if (N < 1) {
  259. fprintf(stderr, "invalid matrix size\n");
  260. exit(1);
  261. }
  262. if (argc > 2) {
  263. BS = atoi(argv[2]);
  264. }
  265. if (BS < 1 || N % BS != 0) {
  266. fprintf(stderr, "invalid block size\n");
  267. exit(1);
  268. }
  269. }
  270. /* Get the process rank and session size */
  271. starpu_mpi_comm_rank(MPI_COMM_WORLD, &comm_rank);
  272. starpu_mpi_comm_size(MPI_COMM_WORLD, &comm_size);
  273. if (comm_rank == 0)
  274. {
  275. #if VERBOSE
  276. printf("N = %d\n", N);
  277. printf("BS = %d\n", BS);
  278. printf("NB = %d\n", NB);
  279. printf("comm_size = %d\n", comm_size);
  280. #endif
  281. /* In this example, node rank 0 performs all the memory allocations and initializations,
  282. * and the blocks are later distributed on the other nodes.
  283. * This is not mandatory however, and blocks could be allocated on other nodes right
  284. * from the beginning, depending on the application needs (in particular for the case
  285. * where the session wide data footprint is larger than a single node available memory. */
  286. alloc_matrices();
  287. init_matrices();
  288. }
  289. /* Register matrices to StarPU and StarPU-MPI */
  290. register_matrices();
  291. /* Distribute C blocks */
  292. distribute_matrix_C();
  293. int b_row,b_col;
  294. for (b_row = 0; b_row < NB; b_row++)
  295. {
  296. for (b_col = 0; b_col < NB; b_col++)
  297. {
  298. starpu_mpi_task_insert(MPI_COMM_WORLD, &gemm_cl,
  299. STARPU_R, A_h[b_row],
  300. STARPU_R, B_h[b_col],
  301. STARPU_RW, C_h[b_row*NB+b_col],
  302. 0);
  303. }
  304. }
  305. starpu_task_wait_for_all();
  306. undistribute_matrix_C();
  307. unregister_matrices();
  308. if (comm_rank == 0) {
  309. #if VERBOSE
  310. disp_matrix(C);
  311. #endif
  312. check_result();
  313. free_matrices();
  314. }
  315. starpu_mpi_shutdown();
  316. starpu_shutdown();
  317. return 0;
  318. }