reshape.c 9.66 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
/*******************************************************************************
 * Copyright 2019 UChicago Argonne, LLC.
 * (c.f. AUTHORS, LICENSE)
 *
 * This file is part of the AML project.
 * For more info, see https://xgitlab.cels.anl.gov/argo/aml
 *
 * SPDX-License-Identifier: BSD-3-Clause
*******************************************************************************/

#include <stdlib.h>

#include "aml.h"
#include "aml/layout/dense.h"
#include "aml/layout/reshape.h"

17
18
19
static int aml_layout_reshape_alloc(struct aml_layout **ret,
				    const size_t ndims,
				    const size_t target_ndims)
20
{
21
22
	struct aml_layout *layout;
	struct aml_layout_data_reshape *data;
23

Nicolas Denoyelle's avatar
Nicolas Denoyelle committed
24
25
26
	layout = AML_INNER_MALLOC_ARRAY(2*ndims + target_ndims, size_t,
					struct aml_layout,
					struct aml_layout_data_reshape);
27
28
29
30
	if (layout == NULL) {
		*ret = NULL;
		return -AML_ENOMEM;
	}
31

Nicolas Denoyelle's avatar
Nicolas Denoyelle committed
32
33
34
35
	data = AML_INNER_MALLOC_GET_FIELD(layout, 2,
					  struct aml_layout,
					  struct aml_layout_data_reshape);

36
	layout->data = (struct aml_layout_data *)data;
Nicolas Denoyelle's avatar
Nicolas Denoyelle committed
37
38
39
40
41
42
43
44

	data->dims = AML_INNER_MALLOC_GET_ARRAY(layout,
						size_t,
						struct aml_layout,
						struct aml_layout_data_reshape);

	data->coffsets = data->dims + ndims;
	data->target_dims = data->coffsets + ndims;
45

46
47
48
49
50
	data->target = NULL;
	data->target_ndims = target_ndims;
	data->ndims = ndims;
	*ret = layout;
	return AML_SUCCESS;
51
52
53
54
55
56
57
58
59
60
61
62
63
}

int aml_layout_reshape_create(struct aml_layout **layout,
			      struct aml_layout *target,
			      const int order,
			      const size_t ndims,
			      const size_t *dims)
{
	struct aml_layout *output;
	struct aml_layout_data_reshape *data;
	size_t target_ndims;
	size_t prod;
	size_t target_prod;
64
	int err;
65
66
67
68
69

	if (layout == NULL || target == NULL || ndims == 0)
		return -AML_EINVAL;

	target_ndims = aml_layout_ndims(target);
70
71
72
	err = aml_layout_reshape_alloc(&output, ndims, target_ndims);
	if (err)
		return err;
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127

	data = (struct aml_layout_data_reshape *)output->data;
	data->target = target;

	switch (AML_LAYOUT_ORDER(order)) {
	case AML_LAYOUT_ORDER_ROW_MAJOR:
		output->ops = &aml_layout_reshape_row_ops;
		for (size_t i = 0; i < ndims; i++)
			data->dims[i] = dims[ndims-i-1];
		break;

	case AML_LAYOUT_ORDER_COLUMN_MAJOR:
		output->ops = &aml_layout_reshape_column_ops;
		memcpy(data->dims, dims, ndims * sizeof(size_t));
		break;
	default:
		free(output);
		return -AML_EINVAL;

	}

	size_t target_dims[target_ndims];

	switch (aml_layout_order(target)) {
	case AML_LAYOUT_ORDER_ROW_MAJOR:
		aml_layout_dims(target, target_dims);
		for (size_t i = 0; i < target_ndims; i++)
			data->target_dims[i] = target_dims[target_ndims-i-1];
		break;
	case AML_LAYOUT_ORDER_COLUMN_MAJOR:
		aml_layout_dims(target, data->target_dims);
		break;
	default:
		free(output);
		return -AML_EINVAL;
	}

	prod = 1;
	for (size_t i = 0; i < ndims; i++) {
		data->coffsets[i] = prod;
		prod *= data->dims[i];
	}
	target_prod = 1;
	for (size_t i = 0; i < data->target_ndims; i++)
		target_prod *= data->target_dims[i];

	if (target_prod != prod) {
		free(output);
		return -AML_EINVAL;
	}

	*layout = output;
	return AML_SUCCESS;
}

128
int aml_layout_reshape_duplicate(const struct aml_layout *layout,
Swann Perarnau's avatar
Swann Perarnau committed
129
                                 struct aml_layout **out)
130
131
132
133
134
135
136
137
{
	const struct aml_layout_data_reshape *data;
	struct aml_layout_data_reshape *dret;
	struct aml_layout *ret;
	int err;

	data = (const struct aml_layout_data_reshape *)layout->data;

Swann Perarnau's avatar
Swann Perarnau committed
138
	if (layout->data == NULL || out == NULL)
139
140
141
142
143
144
145
146
147
148
149
150
		return -AML_EINVAL;

	err = aml_layout_reshape_alloc(&ret, data->ndims, data->target_ndims);
	if (err)
		return err;

	ret->ops = layout->ops;
	dret = (struct aml_layout_data_reshape *)ret->data;
	aml_layout_duplicate(data->target, &dret->target);
	/* small optimization, copying all data at the end of the structure */
	memcpy(dret->dims, data->dims,
	       (2 * data->ndims + data->target_ndims) * sizeof(size_t));
Swann Perarnau's avatar
Swann Perarnau committed
151
	*out = ret;
152
153
154
	return AML_SUCCESS;
}

155
void aml_layout_reshape_destroy(struct aml_layout *l)
156
{
157
158
159
160
161
162
	assert(l != NULL);

	struct aml_layout_data_reshape *data =
	        (struct aml_layout_data_reshape *)l->data;
	aml_layout_destroy(&data->target);
	free(l);
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
}

/*******************************************************************************
 * COLUMN OPERATORS:
 ******************************************************************************/

void *aml_layout_reshape_column_deref(const struct aml_layout_data *data,
				       const size_t *coords)
{
	const struct aml_layout_data_reshape *d;

	d = (const struct aml_layout_data_reshape *)data;

	size_t offset = 0;
	size_t target_coords[d->target_ndims];

	for (size_t i = 0; i < d->ndims; i++)
		offset += coords[i] * d->coffsets[i];

	for (size_t i = 0; i < d->target_ndims; i++) {
		target_coords[i] = offset % d->target_dims[i];
		offset /= d->target_dims[i];
	}

	return d->target->ops->deref_native(d->target->data, target_coords);
}

190
191
192
193
194
195
196
197
198
void *aml_layout_reshape_rawptr(const struct aml_layout_data *data)
{
	const struct aml_layout_data_reshape *d;

	d = (const struct aml_layout_data_reshape *)data;

	return d->target->ops->rawptr(d->target->data);
}

199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
int aml_layout_reshape_column_order(const struct aml_layout_data *data)
{
	(void) data;
	return AML_LAYOUT_ORDER_COLUMN_MAJOR;
}

int aml_layout_reshape_column_dims(const struct aml_layout_data *data,
				    size_t *dims)
{
	const struct aml_layout_data_reshape *d;

	d = (const struct aml_layout_data_reshape *)data;

	memcpy((void *)dims, (void *)d->dims, sizeof(size_t)*d->ndims);

	return 0;
}

size_t aml_layout_reshape_ndims(const struct aml_layout_data *data)
{
	const struct aml_layout_data_reshape *d;

	d = (const struct aml_layout_data_reshape *)data;

	return d->ndims;
}

size_t aml_layout_reshape_element_size(const struct aml_layout_data *data)
{
	const struct aml_layout_data_reshape *d;

	d = (const struct aml_layout_data_reshape *) data;

	return aml_layout_element_size(d->target);
}

235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
int aml_layout_reshape_column_fprintf(const struct aml_layout_data *data,
				      FILE *stream, const char *prefix)
{
	const struct aml_layout_data_reshape *d;

	fprintf(stream, "%s: layout-reshape: %p: column-major\n", prefix,
		(void *)data);
	if (data == NULL)
		return AML_SUCCESS;

	d = (const struct aml_layout_data_reshape *)data;

	fprintf(stream, "%s: ndims: %zu\n", prefix, d->ndims);
	fprintf(stream, "%s: target-ndims: %zu\n", prefix, d->target_ndims);
	for (size_t i = 0; i < d->ndims; i++) {
250
251
252
253
254
255
		fprintf(stream, "%s: %16zu: %16zu %16zu\n", prefix, i,
		        d->dims[i], d->coffsets[i]);
	}
	for (size_t i = 0; i < d->target_ndims; i++) {
		fprintf(stream, "%s: %16zu: %16zu\n", prefix, i,
		        d->target_dims[i]);
256
257
258
259
260
261
262
	}
	fprintf(stream, "%s: target: begin\n", prefix);
	aml_layout_fprintf(stream, prefix, d->target);
	fprintf(stream, "%s: target: end\n", prefix);
	return AML_SUCCESS;
}

263
struct aml_layout_ops aml_layout_reshape_column_ops = {
264
265
266
267
268
269
270
271
272
273
274
275
        aml_layout_reshape_column_deref,
        aml_layout_reshape_column_deref,
        aml_layout_reshape_rawptr,
        aml_layout_reshape_column_order,
        aml_layout_reshape_column_dims,
        aml_layout_reshape_column_dims,
        aml_layout_reshape_ndims,
        aml_layout_reshape_element_size,
        NULL,
        NULL,
        NULL,
        aml_layout_reshape_column_fprintf,
276
        aml_layout_reshape_duplicate,
277
        aml_layout_reshape_destroy,
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
};

/*******************************************************************************
 * ROW OPERATORS:
 ******************************************************************************/

void *aml_layout_reshape_row_deref(const struct aml_layout_data *data,
				   const size_t *coords)
{
	const struct aml_layout_data_reshape *d;

	d = (const struct aml_layout_data_reshape *)data;

	size_t offset = 0;
	size_t target_coords[d->target_ndims];

	for (size_t i = 0; i < d->ndims; i++)
		offset += coords[d->ndims - i - 1] * d->coffsets[i];

	for (size_t i = 0; i < d->target_ndims; i++) {
		target_coords[i] = offset % d->target_dims[i];
		offset /= d->target_dims[i];
	}

	return d->target->ops->deref_native(d->target->data, target_coords);
}

int aml_layout_reshape_row_order(const struct aml_layout_data *data)
{
	(void) data;
	return AML_LAYOUT_ORDER_ROW_MAJOR;
}

int aml_layout_reshape_row_dims(const struct aml_layout_data *data,
				size_t *dims)
{
	const struct aml_layout_data_reshape *d;

	d = (const struct aml_layout_data_reshape *)data;

	for (size_t i = 0; i < d->ndims; i++)
		dims[i] = d->dims[d->ndims - i - 1];

	return 0;
}

324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
int aml_layout_reshape_row_fprintf(const struct aml_layout_data *data,
				   FILE *stream, const char *prefix)
{
	const struct aml_layout_data_reshape *d;

	fprintf(stream, "%s: layout-reshape: %p: row-major\n", prefix,
		(void *)data);
	if (data == NULL)
		return AML_SUCCESS;

	d = (const struct aml_layout_data_reshape *)data;

	fprintf(stream, "%s: ndims: %zu\n", prefix, d->ndims);
	fprintf(stream, "%s: target-ndims: %zu\n", prefix, d->target_ndims);
	for (size_t i = 0; i < d->ndims; i++) {
		size_t j = d->ndims - i - 1;

341
342
		fprintf(stream, "%s: %16zu: %16zu %16zu\n", prefix, i,
		        d->dims[j], d->coffsets[j]);
343
	}
344
345
346
347
348
	for (size_t i = 0; i < d->target_ndims; i++) {
		fprintf(stream, "%s: %16zu: %16zu\n", prefix, i,
		        d->target_dims[i]);
	}

349
350
351
352
353
354
	fprintf(stream, "%s: target: begin\n", prefix);
	aml_layout_fprintf(stream, prefix, d->target);
	fprintf(stream, "%s: target: end\n", prefix);
	return AML_SUCCESS;
}

355
struct aml_layout_ops aml_layout_reshape_row_ops = {
356
357
358
359
360
361
362
363
364
365
366
367
        aml_layout_reshape_row_deref,
        aml_layout_reshape_column_deref,
        aml_layout_reshape_rawptr,
        aml_layout_reshape_row_order,
        aml_layout_reshape_row_dims,
        aml_layout_reshape_column_dims,
        aml_layout_reshape_ndims,
        aml_layout_reshape_element_size,
        NULL,
        NULL,
        NULL,
        aml_layout_reshape_row_fprintf,
368
        aml_layout_reshape_duplicate,
369
        aml_layout_reshape_destroy,
370
};