1 /*------------------------------------------------------------------------------------------------*
2 * Copyright (C) by the DBCSR developers group - All rights reserved *
3 * This file is part of the DBCSR library. *
4 * *
5 * For information on the license, see the LICENSE file. *
6 * For further information please visit https://dbcsr.cp2k.org *
7 * SPDX-License-Identifier: GPL-2.0+ *
8 *------------------------------------------------------------------------------------------------*/
9
10 #include <vector>
11 #include <string>
12 #include <iostream>
13 #include <algorithm>
14 #include <cstdlib>
15 #include <cstdio>
16 #include <cstdint>
17 #include <random>
18 #include <mpi.h>
19 #include <dbcsr.h>
20 #include <tensors/dbcsr_tensor.h>
21
22 //-------------------------------------------------------------------------------------------------!
23 // Example: tensor contraction (13|2)x(54|21)=(3|45)
24 // tensor1 x tensor2 = tensor3
25 //-------------------------------------------------------------------------------------------------!
26
random_dist(int dist_size,int nbins)27 std::vector<int> random_dist(int dist_size, int nbins)
28 {
29
30 std::vector<int> dist(dist_size);
31
32 for(int i=0; i < dist_size; i++)
33 dist[i] = i % nbins;
34
35 return dist;
36 }
37
printvec(std::vector<int> & v)38 void printvec(std::vector<int>& v) {
39
40 for (auto i : v) {
41 std::cout << i << " ";
42 }
43 std::cout << '\n' << std::endl;
44
45 }
46
fill_random(void * tensor,std::vector<std::vector<int>> nzblocks)47 void fill_random(void* tensor, std::vector<std::vector<int>> nzblocks) {
48
49 int myrank, mpi_size;
50 int dim = nzblocks.size();
51
52 MPI_Comm_rank(MPI_COMM_WORLD, &myrank);
53 MPI_Comm_size(MPI_COMM_WORLD, &mpi_size);
54
55 std::random_device rd;
56 std::mt19937 gen(rd());
57 std::uniform_real_distribution<> dis(-1.0, 1.0);
58
59 if (myrank == 0) std::cout << "Filling Tensor..." << std::endl;
60 if (myrank == 0) std::cout << "Dimension: " << dim << std::endl;
61
62 int nblocks = nzblocks[0].size();
63 std::vector<std::vector<int>> mynzblocks(dim);
64 std::vector<int> idx(dim);
65
66 for (int i = 0; i != nblocks; ++i) {
67
68 // make index out of nzblocks
69 for (int j = 0; j != dim; ++j) idx[j] = nzblocks[j][i];
70
71 int proc = -1;
72
73 c_dbcsr_t_get_stored_coordinates(tensor, idx.data(), &proc);
74
75 if (proc == myrank) {
76 for (int j = 0; j != dim; ++j)
77 mynzblocks[j].push_back(idx[j]);
78 }
79
80 }
81
82 std::vector<int*> dataptr(4, nullptr);
83
84 for (int i = 0; i != dim; ++i) {
85 dataptr[i] = mynzblocks[i].size() == 0 ? nullptr : &mynzblocks[i][0];
86 }
87
88 if (myrank == 0) std::cout << "Reserving blocks..." << std::endl;
89
90 if (mynzblocks[0].size() != 0)
91 c_dbcsr_t_reserve_blocks_index(tensor, mynzblocks[0].size(), dataptr[0], dataptr[1], dataptr[2], dataptr[3]);
92
93 auto fill_rand = [&](std::vector<double>& blk) {
94 for (auto& e : blk) {
95 e = dis(gen);
96 }
97 };
98
99 void* iter = nullptr;
100
101 c_dbcsr_t_iterator_start(&iter, tensor);
102
103 std::vector<int> loc_idx(dim);
104 std::vector<int> blk_sizes(dim);
105 std::vector<double> block(1);
106
107 int blk = 0;
108 int blk_proc = 0;
109
110 while(c_dbcsr_t_iterator_blocks_left(iter)) {
111
112 c_dbcsr_t_iterator_next_block(iter, loc_idx.data(), &blk, &blk_proc, blk_sizes.data(), nullptr);
113
114 int tot = 1;
115 for (int i = 0; i != dim; ++i) {
116 tot *= blk_sizes[i];
117 }
118
119 block.resize(tot);
120
121 fill_rand(block);
122
123 c_dbcsr_t_put_block(tensor, loc_idx.data(), blk_sizes.data(), block.data(), nullptr, nullptr);
124
125 }
126
127 c_dbcsr_t_iterator_stop(&iter);
128
129 MPI_Barrier(MPI_COMM_WORLD);
130
131 }
132
133
main(int argc,char * argv[])134 int main(int argc, char* argv[])
135 {
136 MPI_Init(&argc, &argv);
137
138 int mpi_size, mpi_rank;
139 MPI_Comm_size(MPI_COMM_WORLD, &mpi_size);
140 MPI_Comm_rank(MPI_COMM_WORLD, &mpi_rank);
141
142 c_dbcsr_init_lib(MPI_COMM_WORLD, nullptr);
143
144 void* pgrid_3d = nullptr;
145 void* pgrid_4d = nullptr;
146
147 std::vector<int> dims4(4);
148 std::vector<int> dims3(3);
149
150 MPI_Fint fcomm = MPI_Comm_c2f(MPI_COMM_WORLD);
151
152 c_dbcsr_t_pgrid_create(&fcomm, dims3.data(), dims3.size(), &pgrid_3d, nullptr);
153
154 c_dbcsr_t_pgrid_create(&fcomm, dims4.data(), dims4.size(), &pgrid_4d, nullptr);
155
156 if (mpi_rank == 0) {
157
158 std::cout << "pgrid3-dimensions:" << std::endl;
159 printvec(dims3);
160
161 std::cout << "pgrid4-dimensions:" << std::endl;
162 printvec(dims4);
163 }
164
165 // block sizes
166 std::vector<int> blk1, blk2, blk3, blk4, blk5;
167 // blk indices of non-zero blocks
168 std::vector<int> nz11, nz12, nz13, nz21, nz22, nz24, nz25, nz33, nz34, nz35;
169
170 blk1 = {3, 9, 12, 1};
171 blk2 = {4, 2, 3, 1, 9, 2, 32, 10, 5, 8, 7};
172 blk3 = {7, 3, 8, 7, 9, 5, 10, 23, 2};
173 blk4 = {8, 1, 4, 13, 6};
174 blk5 = {4, 2, 22};
175
176 nz11 = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
177 0, 1, 1, 1, 1, 1, 1, 1, 2, 2,
178 2, 2, 2, 2, 2, 2, 2, 3, 3, 3,
179 3, 3};
180 nz12 = {2, 4, 4, 4, 5, 5, 6, 7, 9,10,10,
181 0, 0, 3, 6, 6, 8, 9 ,1, 1, 4, 5,
182 7, 7, 8,10,10, 1 ,3, 4, 4, 7};
183 nz13 = {6, 2, 4, 8, 5, 7, 1, 7, 2, 1, 2,
184 0, 3, 5, 1, 6, 4, 7, 2, 6, 0, 3,
185 2, 6, 7, 4, 7, 8, 5, 0, 1, 6};
186
187 nz21 = { 0, 0, 0, 0, 0, 1, 1, 1, 1, 1,
188 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
189 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
190 3, 3, 3, 3, 3, 3 };
191 nz22 = { 0, 2, 3, 5, 9, 1, 1, 3, 4, 4,
192 5, 5, 5, 6, 6, 8, 8, 8, 9, 10,
193 0, 2, 2, 3, 4, 5, 7, 8, 10, 10,
194 0, 2, 3, 5, 9, 10 };
195 nz24 = { 2, 4, 1, 2, 1, 2, 4, 0, 0, 3,
196 1, 2, 3, 0, 3, 2, 3, 3, 1, 0,
197 2, 0, 0, 2, 3, 2, 3, 1, 1, 2,
198 0, 0, 2, 1, 4, 4 };
199 nz25 = { 0, 2, 1, 0, 0, 1, 2, 0, 2, 0,
200 1, 2, 1, 0, 2, 1, 2, 1, 0, 1,
201 2, 0, 1, 2, 1, 1, 1, 2, 0, 1,
202 0, 2, 1, 0, 2, 1 };
203
204 nz33 = { 1, 3, 4, 4, 4, 5, 5, 7 };
205 nz34 = { 2, 1, 0, 0, 2, 1, 3, 4 };
206 nz35 = { 2, 1, 0, 1, 2, 1, 0, 0 };
207
208 // (13|2)x(54|21)=(3|45)
209 // distribute blocks
210 std::vector<int> dist11 = random_dist(blk1.size(), dims3[0]);
211 std::vector<int> dist12 = random_dist(blk2.size(), dims3[1]);
212 std::vector<int> dist13 = random_dist(blk3.size(), dims3[2]);
213
214 std::vector<int> dist21 = random_dist(blk1.size(), dims4[0]);
215 std::vector<int> dist22 = random_dist(blk2.size(), dims4[1]);
216 std::vector<int> dist23 = random_dist(blk4.size(), dims4[2]);
217 std::vector<int> dist24 = random_dist(blk5.size(), dims4[3]);
218
219 std::vector<int> dist31 = random_dist(blk3.size(), dims3[0]);
220 std::vector<int> dist32 = random_dist(blk4.size(), dims3[1]);
221 std::vector<int> dist33 = random_dist(blk5.size(), dims3[2]);
222
223
224 if (mpi_rank == 0) {
225
226 std::cout << "dist11:" << std::endl;
227 printvec(dist11);
228
229 std::cout << "dist12:" << std::endl;
230 printvec(dist12);
231
232 std::cout << "dist13:" << std::endl;
233 printvec(dist13);
234
235 std::cout << "dist21:" << std::endl;
236 printvec(dist21);
237
238 std::cout << "dist22:" << std::endl;
239 printvec(dist22);
240
241 std::cout << "dist23:" << std::endl;
242 printvec(dist23);
243
244 std::cout << "dist24:" << std::endl;
245 printvec(dist24);
246
247 std::cout << "dist31:" << std::endl;
248 printvec(dist31);
249
250 std::cout << "dist32:" << std::endl;
251 printvec(dist32);
252
253 std::cout << "dist33:" << std::endl;
254 printvec(dist33);
255
256 }
257
258 void* dist1 = nullptr;
259 void* dist2 = nullptr;
260 void* dist3 = nullptr;
261
262 // (13|2)x(54|21)=(3|45)
263 std::vector<int> map11, map12, map21, map22, map31, map32;
264
265 map11 = {0, 2};
266 map12 = {1};
267 map21 = {3, 2};
268 map22 = {1, 0};
269 map31 = {0};
270 map32 = {1, 2};
271
272 if (mpi_rank == 0) std::cout << "Creating dist objects..." << '\n' << std::endl;
273
274 // create distribution objects
275 c_dbcsr_t_distribution_new(&dist1, pgrid_3d, dist11.data(), dist11.size(),
276 dist12.data(), dist12.size(), dist13.data(), dist13.size(), nullptr, 0);
277
278 c_dbcsr_t_distribution_new(&dist2, pgrid_4d, dist21.data(), dist21.size(),
279 dist22.data(), dist22.size(), dist23.data(), dist23.size(),
280 dist24.data(), dist24.size());
281
282 c_dbcsr_t_distribution_new(&dist3, pgrid_3d, dist31.data(), dist31.size(),
283 dist32.data(), dist32.size(), dist33.data(), dist33.size(), nullptr, 0);
284
285 MPI_Barrier(MPI_COMM_WORLD);
286
287 // create tensors
288 // (13|2)x(54|21)=(3|45)
289
290 void* tensor1 = nullptr;
291 void* tensor2 = nullptr;
292 void* tensor3 = nullptr;
293
294 if (mpi_rank == 0) std::cout << "Creating tensors..." << std::endl;
295
296 c_dbcsr_t_create_new(&tensor1, "(13|2)", dist1, map11.data(), map11.size(), map12.data(), map12.size(), nullptr, blk1.data(),
297 blk1.size(), blk2.data(), blk2.size(), blk3.data(), blk3.size(), nullptr, 0);
298
299 c_dbcsr_t_create_new(&tensor2, "(54|21)", dist2, map21.data(), map21.size(), map22.data(), map22.size(), nullptr, blk1.data(),
300 blk1.size(), blk2.data(), blk2.size(), blk4.data(), blk4.size(), blk5.data(), blk5.size());
301
302 c_dbcsr_t_create_new(&tensor3, "(3|45)", dist3, map31.data(), map31.size(), map32.data(), map32.size(), nullptr, blk3.data(),
303 blk3.size(), blk4.data(), blk4.size(), blk5.data(), blk5.size(), nullptr, 0);
304
305 MPI_Barrier(MPI_COMM_WORLD);
306
307 // fill the tensors
308
309 if (mpi_rank == 0) std::cout << "Tensor 1" << '\n' << std::endl;
310 fill_random(tensor1, {nz11, nz12, nz13});
311 if (mpi_rank == 0) std::cout << "Tensor 2" << '\n' << std::endl;
312 fill_random(tensor2, {nz21, nz22, nz24, nz25});
313 if (mpi_rank == 0) std::cout << "Tensor 3" << '\n' << std::endl;
314 fill_random(tensor3, {nz33, nz34, nz35});
315
316 // contracting
317
318 // (13|2)x(54|21)=(3|45)
319
320 MPI_Barrier(MPI_COMM_WORLD);
321
322 if (mpi_rank == 0) std::cout << "Contracting..." << std::endl;
323
324 // cn : indices to be contracted
325 // noncn : indices not to be contracted
326 // mapn : how nonc indices map to tensor 3
327 std::vector<int> c1, nonc1, c2, nonc2, map1, map2;
328 c1 = {0,1};
329 nonc1 = {2};
330 c2 = {0,1};
331 nonc2 = {2,3};
332 map1 = {0};
333 map2 = {1,2};
334
335 int unit_nr = -1;
336 if (mpi_rank == 0) unit_nr = 6;
337 bool log_verbose = true;
338
339 // tensor_3(map_1, map_2) := 0.2 * tensor_1(notcontract_1, contract_1)
340 // * tensor_2(contract_2, notcontract_2)
341 // + 0.8 * tensor_3(map_1, map_2)
342
343 c_dbcsr_t_contract_r_dp (0.2, tensor1, tensor2, 0.8, tensor3, c1.data(), c1.size(), nonc1.data(), nonc1.size(),
344 c2.data(), c2.size(), nonc2.data(), nonc2.size(), map1.data(), map1.size(), map2.data(),
345 map2.size(), nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr,
346 nullptr, nullptr, nullptr, &unit_nr, &log_verbose);
347
348
349
350 c_dbcsr_t_destroy(&tensor1);
351 c_dbcsr_t_destroy(&tensor2);
352 c_dbcsr_t_destroy(&tensor3);
353
354 c_dbcsr_t_pgrid_destroy(&pgrid_3d, nullptr);
355 c_dbcsr_t_pgrid_destroy(&pgrid_4d, nullptr);
356
357 c_dbcsr_t_distribution_destroy(&dist1);
358 c_dbcsr_t_distribution_destroy(&dist2);
359 c_dbcsr_t_distribution_destroy(&dist3);
360
361 c_dbcsr_finalize_lib();
362
363 MPI_Finalize();
364
365 return 0;
366 }
367