Line data Source code
1 : /*----------------------------------------------------------------------------*/
2 : /* CP2K: A general program to perform molecular dynamics simulations */
3 : /* Copyright 2000-2025 CP2K developers group <https://cp2k.org> */
4 : /* */
5 : /* SPDX-License-Identifier: BSD-3-Clause */
6 : /*----------------------------------------------------------------------------*/
7 :
8 : #ifndef GRID_DGEMM_TENSOR_LOCAL_H
9 : #define GRID_DGEMM_TENSOR_LOCAL_H
10 :
11 : #include <stdbool.h>
12 : #include <stdio.h>
13 : #include <stdlib.h>
14 : #include <string.h>
15 : #include <unistd.h>
16 :
17 : typedef struct tensor_ {
18 : int dim_;
19 : int size[4];
20 : size_t alloc_size_;
21 : size_t old_alloc_size_;
22 : int offsets[4];
23 : double *data;
24 : int ld_;
25 : int window_shift[4]; /* lower corner of the window. Should be between lower
26 : * corner and upper corner of the local grid */
27 : int window_size[4]; /* size of the window where computations should be
28 : * done */
29 : int full_size[4]; /* size of the global grid */
30 : int lower_corner[4]; /* coordinates of the lower corner of the local part of
31 : * the grid. It can be different from the window where
32 : * computations should be done. The upper corner can be
33 : * deduced with the sum of the grid size and the lower
34 : * corner */
35 : /* only relevant when the tensor represents a grid */
36 : double dh[3][3];
37 : double dh_inv[3][3];
38 : bool orthogonal[3];
39 : } tensor;
40 :
41 : extern void tensor_copy(tensor *const b, const tensor *const a);
42 :
43 : /* initialize a tensor structure for a tensor of dimension dim <= 4 */
44 :
45 1115471 : static inline void initialize_tensor(struct tensor_ *a, const int dim,
46 : const int *const sizes) {
47 1115471 : if (a == NULL)
48 : return;
49 :
50 1115471 : a->dim_ = dim;
51 4084749 : for (int d = 0; d < dim; d++)
52 2969278 : a->size[d] = sizes[d];
53 :
54 : // we need proper alignment here. But can be done later
55 : /* a->ld_ = (sizes[a->dim_ - 1] / 32 + 1) * 32; */
56 1115471 : a->ld_ = sizes[a->dim_ - 1];
57 1115471 : switch (a->dim_) {
58 90334 : case 4: {
59 90334 : a->offsets[0] = a->ld_ * a->size[1] * a->size[2];
60 90334 : a->offsets[1] = a->ld_ * a->size[2];
61 90334 : a->offsets[2] = a->ld_;
62 90334 : break;
63 : }
64 557668 : case 3: {
65 557668 : a->offsets[0] = a->ld_ * a->size[1];
66 557668 : a->offsets[1] = a->ld_;
67 557668 : } break;
68 467469 : case 2: { // matrix case
69 467469 : a->offsets[0] = a->ld_;
70 467469 : } break;
71 : case 1:
72 : break;
73 : }
74 :
75 1115471 : a->alloc_size_ = a->offsets[0] * a->size[0];
76 1115471 : return;
77 : }
78 :
79 : /* initialize a tensor structure for a tensor of dimension dim = 2 */
80 :
81 418368 : static inline void initialize_tensor_2(struct tensor_ *a, int n1, int n2) {
82 113215 : if (a == NULL)
83 0 : return;
84 :
85 418368 : int size_[2] = {n1, n2};
86 418368 : initialize_tensor(a, 2, size_);
87 : }
88 :
89 : /* initialize a tensor structure for a tensor of dimension dim = 2 */
90 :
91 557588 : static inline void initialize_tensor_3(struct tensor_ *a, int n1, int n2,
92 : int n3) {
93 416600 : if (a == NULL)
94 0 : return;
95 557588 : int size_[3] = {n1, n2, n3};
96 487094 : initialize_tensor(a, 3, size_);
97 : }
98 :
99 : /* initialize a tensor structure for a tensor of dimension dim = 2 */
100 :
101 90334 : static inline void initialize_tensor_4(struct tensor_ *a, int n1, int n2,
102 : int n3, int n4) {
103 90334 : if (a == NULL)
104 0 : return;
105 90334 : int size_[4] = {n1, n2, n3, n4};
106 90334 : initialize_tensor(a, 4, size_);
107 : }
108 :
109 : /* initialize a tensor structure for a tensor of dimension dim = 2 */
110 :
111 : static inline tensor *create_tensor(const int dim, const int *sizes) {
112 : tensor *a = (tensor *)malloc(sizeof(struct tensor_));
113 :
114 : if (a == NULL)
115 : abort();
116 :
117 : initialize_tensor(a, dim, sizes);
118 : a->data = (double *)malloc(sizeof(double) * a->alloc_size_);
119 : if (a->data == NULL)
120 : abort();
121 : a->old_alloc_size_ = a->alloc_size_;
122 : return a;
123 : }
124 :
125 : /* destroy a tensor created with the function above */
126 : static inline void destroy_tensor(tensor *a) {
127 : if (a->data)
128 : free(a->data);
129 : free(a);
130 : }
131 :
132 : static inline size_t tensor_return_memory_size(const struct tensor_ *const a) {
133 : if (a == NULL)
134 : abort();
135 :
136 : return a->alloc_size_;
137 : }
138 :
139 : static inline void tensor_assign_memory(struct tensor_ *a, void *data) {
140 : if (a == NULL)
141 : abort();
142 : a->data = (double *)data;
143 : }
144 :
145 : static inline int tensor_get_leading_dimension(struct tensor_ *a) {
146 : if (a == NULL)
147 : abort();
148 : return a->ld_;
149 : }
150 :
151 : static inline void tensor_set_leading_dimension(struct tensor_ *a,
152 : const int ld) {
153 : if (a == NULL)
154 : abort();
155 : a->ld_ = ld;
156 : }
157 :
158 : static inline void recompute_tensor_offsets(struct tensor_ *a) {
159 : if (a == NULL)
160 : abort();
161 :
162 : switch (a->dim_) {
163 : case 5: {
164 : a->offsets[0] = a->ld_ * a->size[1] * a->size[2] * a->size[3];
165 : a->offsets[1] = a->ld_ * a->size[1] * a->size[2];
166 : a->offsets[2] = a->ld_ * a->size[2];
167 : a->offsets[3] = a->ld_;
168 : break;
169 : }
170 : case 4: {
171 : a->offsets[0] = a->ld_ * a->size[1] * a->size[2];
172 : a->offsets[1] = a->ld_ * a->size[2];
173 : a->offsets[2] = a->ld_;
174 : break;
175 : }
176 : case 3: {
177 : a->offsets[0] = a->ld_ * a->size[1];
178 : a->offsets[1] = a->ld_;
179 : } break;
180 : case 2: { // matrix case
181 : a->offsets[0] = a->ld_;
182 : } break;
183 : case 1:
184 : break;
185 : }
186 : }
187 :
188 90334 : static inline size_t compute_memory_space_tensor_3(const int n1, const int n2,
189 : const int n3) {
190 90334 : return (n1 * n2 * n3);
191 : }
192 :
193 0 : static inline size_t compute_memory_space_tensor_4(const int n1, const int n2,
194 : const int n3, const int n4) {
195 0 : return (n1 * n2 * n3 * n4);
196 : }
197 :
198 1280 : static inline void setup_global_grid_size(tensor *const grid,
199 : const int *const full_size) {
200 1280 : switch (grid->dim_) {
201 0 : case 1:
202 0 : grid->full_size[0] = full_size[0];
203 0 : break;
204 0 : case 2: {
205 0 : grid->full_size[1] = full_size[0];
206 0 : grid->full_size[0] = full_size[1];
207 0 : } break;
208 1280 : case 3: {
209 1280 : grid->full_size[0] = full_size[2];
210 1280 : grid->full_size[1] = full_size[1];
211 1280 : grid->full_size[2] = full_size[0];
212 1280 : } break;
213 : default:
214 0 : for (int d = 0; d < grid->dim_; d++)
215 0 : grid->full_size[d] = full_size[grid->dim_ - d - 1];
216 : break;
217 : }
218 1280 : }
219 :
220 0 : static inline void setup_grid_window(tensor *const grid,
221 : const int *const shift_local,
222 : const int *const border_width,
223 : const int border_mask) {
224 0 : for (int d = 0; d < grid->dim_; d++) {
225 0 : grid->lower_corner[d] = shift_local[grid->dim_ - d - 1];
226 0 : grid->window_shift[d] = 0;
227 0 : grid->window_size[d] = grid->size[d];
228 0 : if (grid->size[d] != grid->full_size[d]) {
229 0 : grid->window_size[d]--;
230 : }
231 : }
232 :
233 0 : if (border_width) {
234 0 : if (border_mask & (1 << 0))
235 0 : grid->window_shift[2] += border_width[0];
236 0 : if (border_mask & (1 << 1))
237 0 : grid->window_size[2] -= border_width[0];
238 0 : if (border_mask & (1 << 2))
239 0 : grid->window_shift[1] += border_width[1];
240 0 : if (border_mask & (1 << 3))
241 0 : grid->window_size[1] -= border_width[1];
242 0 : if (border_mask & (1 << 4))
243 0 : grid->window_shift[0] += border_width[2];
244 0 : if (border_mask & (1 << 5))
245 0 : grid->window_size[0] -= border_width[2];
246 : }
247 0 : }
248 :
249 : extern size_t realloc_tensor(tensor *t);
250 : extern void alloc_tensor(tensor *t);
251 :
252 : #define idx5(a, i, j, k, l, m) \
253 : a.data[(i) * a.offsets[0] + (j) * a.offsets[1] + (k) * a.offsets[2] + \
254 : (l) * a.ld_ + m]
255 : #define idx4(a, i, j, k, l) \
256 : a.data[(i) * a.offsets[0] + (j) * a.offsets[1] + (k) * a.ld_ + (l)]
257 : #define idx3(a, i, j, k) a.data[(i) * a.offsets[0] + (j) * a.ld_ + (k)]
258 : #define idx2(a, i, j) a.data[(i) * a.ld_ + (j)]
259 : #endif
|
