isl_schedule_node.c: collect_filter_prefix: allow traversal of subset of list
[isl.git] / isl_band.c
blobaaee1f1c9510c2cf7864148d89ed4ca1954aa9e5
1 /*
2 * Copyright 2011 INRIA Saclay
3 * Copyright 2012-2013 Ecole Normale Superieure
5 * Use of this software is governed by the MIT license
7 * Written by Sven Verdoolaege, INRIA Saclay - Ile-de-France,
8 * Parc Club Orsay Universite, ZAC des vignes, 4 rue Jacques Monod,
9 * 91893 Orsay, France
10 * and Ecole Normale Superieure, 45 rue d'Ulm, 75230 Paris, France
13 #include <isl_band_private.h>
14 #include <isl_schedule_private.h>
16 #undef BASE
17 #define BASE band
19 #include <isl_list_templ.c>
21 isl_ctx *isl_band_get_ctx(__isl_keep isl_band *band)
23 return band ? isl_union_pw_multi_aff_get_ctx(band->pma) : NULL;
26 __isl_give isl_band *isl_band_alloc(isl_ctx *ctx)
28 isl_band *band;
30 band = isl_calloc_type(ctx, isl_band);
31 if (!band)
32 return NULL;
34 band->ref = 1;
36 return band;
39 /* Create a duplicate of the given band. The duplicate refers
40 * to the same schedule and parent as the input, but does not
41 * increment their reference counts.
43 __isl_give isl_band *isl_band_dup(__isl_keep isl_band *band)
45 int i;
46 isl_ctx *ctx;
47 isl_band *dup;
49 if (!band)
50 return NULL;
52 ctx = isl_band_get_ctx(band);
53 dup = isl_band_alloc(ctx);
54 if (!dup)
55 return NULL;
57 dup->n = band->n;
58 dup->coincident = isl_alloc_array(ctx, int, band->n);
59 if (band->n && !dup->coincident)
60 goto error;
62 for (i = 0; i < band->n; ++i)
63 dup->coincident[i] = band->coincident[i];
65 dup->pma = isl_union_pw_multi_aff_copy(band->pma);
66 dup->schedule = band->schedule;
67 dup->parent = band->parent;
69 if (!dup->pma)
70 goto error;
72 return dup;
73 error:
74 isl_band_free(dup);
75 return NULL;
78 /* We not only increment the reference count of the band,
79 * but also that of the schedule that contains this band.
80 * This ensures that the schedule won't disappear while there
81 * is still a reference to the band outside of the schedule.
82 * There is no need to increment the reference count of the parent
83 * band as the parent band is part of the same schedule.
85 __isl_give isl_band *isl_band_copy(__isl_keep isl_band *band)
87 if (!band)
88 return NULL;
90 band->ref++;
91 band->schedule->ref++;
92 return band;
95 /* If this is not the last reference to the band (the one from within the
96 * schedule), then we also need to decrement the reference count of the
97 * containing schedule as it was incremented in isl_band_copy.
99 __isl_null isl_band *isl_band_free(__isl_take isl_band *band)
101 if (!band)
102 return NULL;
104 if (--band->ref > 0) {
105 isl_schedule_free(band->schedule);
106 return NULL;
109 isl_union_pw_multi_aff_free(band->pma);
110 isl_band_list_free(band->children);
111 free(band->coincident);
112 free(band);
114 return NULL;
117 int isl_band_has_children(__isl_keep isl_band *band)
119 if (!band)
120 return -1;
122 return band->children != NULL;
125 __isl_give isl_band_list *isl_band_get_children(
126 __isl_keep isl_band *band)
128 if (!band)
129 return NULL;
130 if (!band->children)
131 isl_die(isl_band_get_ctx(band), isl_error_invalid,
132 "band has no children", return NULL);
133 return isl_band_list_dup(band->children);
136 int isl_band_n_member(__isl_keep isl_band *band)
138 return band ? band->n : 0;
141 /* Is the given scheduling dimension coincident within the band and
142 * with respect to the coincidence constraints.
144 int isl_band_member_is_coincident(__isl_keep isl_band *band, int pos)
146 if (!band)
147 return -1;
149 if (pos < 0 || pos >= band->n)
150 isl_die(isl_band_get_ctx(band), isl_error_invalid,
151 "invalid member position", return -1);
153 return band->coincident[pos];
156 /* Return the schedule that leads up to this band.
158 __isl_give isl_union_map *isl_band_get_prefix_schedule(
159 __isl_keep isl_band *band)
161 isl_union_set *domain;
162 isl_union_pw_multi_aff *prefix;
163 isl_band *a;
165 if (!band)
166 return NULL;
168 prefix = isl_union_pw_multi_aff_copy(band->pma);
169 domain = isl_union_pw_multi_aff_domain(prefix);
170 prefix = isl_union_pw_multi_aff_from_domain(domain);
172 for (a = band->parent; a; a = a->parent) {
173 isl_union_pw_multi_aff *partial;
175 partial = isl_union_pw_multi_aff_copy(a->pma);
176 prefix = isl_union_pw_multi_aff_flat_range_product(partial,
177 prefix);
180 return isl_union_map_from_union_pw_multi_aff(prefix);
183 /* Return the schedule of the band in isolation.
185 __isl_give isl_union_pw_multi_aff *
186 isl_band_get_partial_schedule_union_pw_multi_aff(__isl_keep isl_band *band)
188 return band ? isl_union_pw_multi_aff_copy(band->pma) : NULL;
191 /* Return the schedule of the band in isolation.
193 __isl_give isl_union_map *isl_band_get_partial_schedule(
194 __isl_keep isl_band *band)
196 isl_union_pw_multi_aff *sched;
198 sched = isl_band_get_partial_schedule_union_pw_multi_aff(band);
199 return isl_union_map_from_union_pw_multi_aff(sched);
202 __isl_give isl_union_pw_multi_aff *
203 isl_band_get_suffix_schedule_union_pw_multi_aff(__isl_keep isl_band *band);
205 /* Return the schedule for the given band list.
206 * For each band in the list, the schedule is composed of the partial
207 * and suffix schedules of that band.
209 __isl_give isl_union_pw_multi_aff *
210 isl_band_list_get_suffix_schedule_union_pw_multi_aff(
211 __isl_keep isl_band_list *list)
213 isl_ctx *ctx;
214 int i, n;
215 isl_space *space;
216 isl_union_pw_multi_aff *suffix;
218 if (!list)
219 return NULL;
221 ctx = isl_band_list_get_ctx(list);
222 space = isl_space_alloc(ctx, 0, 0, 0);
223 suffix = isl_union_pw_multi_aff_empty(space);
224 n = isl_band_list_n_band(list);
225 for (i = 0; i < n; ++i) {
226 isl_band *el;
227 isl_union_pw_multi_aff *partial;
228 isl_union_pw_multi_aff *suffix_i;
230 el = isl_band_list_get_band(list, i);
231 partial = isl_band_get_partial_schedule_union_pw_multi_aff(el);
232 suffix_i = isl_band_get_suffix_schedule_union_pw_multi_aff(el);
233 suffix_i = isl_union_pw_multi_aff_flat_range_product(
234 partial, suffix_i);
235 suffix = isl_union_pw_multi_aff_union_add(suffix, suffix_i);
237 isl_band_free(el);
240 return suffix;
243 /* Return the schedule for the given band list.
244 * For each band in the list, the schedule is composed of the partial
245 * and suffix schedules of that band.
247 __isl_give isl_union_map *isl_band_list_get_suffix_schedule(
248 __isl_keep isl_band_list *list)
250 isl_union_pw_multi_aff *suffix;
252 suffix = isl_band_list_get_suffix_schedule_union_pw_multi_aff(list);
253 return isl_union_map_from_union_pw_multi_aff(suffix);
256 /* Return the schedule for the forest underneath the given band.
258 __isl_give isl_union_pw_multi_aff *
259 isl_band_get_suffix_schedule_union_pw_multi_aff(__isl_keep isl_band *band)
261 isl_union_pw_multi_aff *suffix;
263 if (!band)
264 return NULL;
266 if (!isl_band_has_children(band)) {
267 isl_union_set *domain;
269 suffix = isl_union_pw_multi_aff_copy(band->pma);
270 domain = isl_union_pw_multi_aff_domain(suffix);
271 suffix = isl_union_pw_multi_aff_from_domain(domain);
272 } else {
273 isl_band_list *list;
275 list = isl_band_get_children(band);
276 suffix =
277 isl_band_list_get_suffix_schedule_union_pw_multi_aff(list);
278 isl_band_list_free(list);
281 return suffix;
284 /* Return the schedule for the forest underneath the given band.
286 __isl_give isl_union_map *isl_band_get_suffix_schedule(
287 __isl_keep isl_band *band)
289 isl_union_pw_multi_aff *suffix;
291 suffix = isl_band_get_suffix_schedule_union_pw_multi_aff(band);
292 return isl_union_map_from_union_pw_multi_aff(suffix);
295 /* Call "fn" on each band (recursively) in the list
296 * in depth-first post-order.
298 int isl_band_list_foreach_band(__isl_keep isl_band_list *list,
299 int (*fn)(__isl_keep isl_band *band, void *user), void *user)
301 int i, n;
303 if (!list)
304 return -1;
306 n = isl_band_list_n_band(list);
307 for (i = 0; i < n; ++i) {
308 isl_band *band;
309 int r = 0;
311 band = isl_band_list_get_band(list, i);
312 if (isl_band_has_children(band)) {
313 isl_band_list *children;
315 children = isl_band_get_children(band);
316 r = isl_band_list_foreach_band(children, fn, user);
317 isl_band_list_free(children);
320 if (!band)
321 r = -1;
322 if (r == 0)
323 r = fn(band, user);
325 isl_band_free(band);
326 if (r)
327 return r;
330 return 0;
333 /* Internal data used during the construction of the schedule
334 * for the tile loops.
336 * sizes contains the tile sizes
337 * scale is set if the tile loops should be scaled
338 * tiled collects the result for a single statement
339 * res collects the result for all statements
341 struct isl_band_tile_data {
342 isl_multi_val *sizes;
343 isl_union_pw_multi_aff *res;
344 isl_pw_multi_aff *tiled;
345 int scale;
348 /* Given part of the schedule of a band, construct the corresponding
349 * schedule for the tile loops based on the tile sizes in data->sizes
350 * and add the result to data->tiled.
352 * If data->scale is set, then dimension i of the schedule will be
353 * of the form
355 * m_i * floor(s_i(x) / m_i)
357 * where s_i(x) refers to the original schedule and m_i is the tile size.
358 * If data->scale is not set, then dimension i of the schedule will be
359 * of the form
361 * floor(s_i(x) / m_i)
364 static int multi_aff_tile(__isl_take isl_set *set,
365 __isl_take isl_multi_aff *ma, void *user)
367 struct isl_band_tile_data *data = user;
368 isl_pw_multi_aff *pma;
369 int i, n;
370 isl_val *v;
372 n = isl_multi_aff_dim(ma, isl_dim_out);
374 for (i = 0; i < n; ++i) {
375 isl_aff *aff;
377 aff = isl_multi_aff_get_aff(ma, i);
378 v = isl_multi_val_get_val(data->sizes, i);
380 aff = isl_aff_scale_down_val(aff, isl_val_copy(v));
381 aff = isl_aff_floor(aff);
382 if (data->scale)
383 aff = isl_aff_scale_val(aff, isl_val_copy(v));
384 isl_val_free(v);
386 ma = isl_multi_aff_set_aff(ma, i, aff);
389 pma = isl_pw_multi_aff_alloc(set, ma);
390 data->tiled = isl_pw_multi_aff_union_add(data->tiled, pma);
392 return 0;
395 /* Given part of the schedule of a band, construct the corresponding
396 * schedule for the tile loops based on the tile sizes in data->sizes
397 * and add the result to data->res.
399 static int pw_multi_aff_tile(__isl_take isl_pw_multi_aff *pma, void *user)
401 struct isl_band_tile_data *data = user;
403 data->tiled = isl_pw_multi_aff_empty(isl_pw_multi_aff_get_space(pma));
405 if (isl_pw_multi_aff_foreach_piece(pma, &multi_aff_tile, data) < 0)
406 goto error;
408 isl_pw_multi_aff_free(pma);
409 data->res = isl_union_pw_multi_aff_add_pw_multi_aff(data->res,
410 data->tiled);
412 return 0;
413 error:
414 isl_pw_multi_aff_free(pma);
415 isl_pw_multi_aff_free(data->tiled);
416 return -1;
419 /* Given the schedule of a band, construct the corresponding
420 * schedule for the tile loops based on the given tile sizes
421 * and return the result.
423 static isl_union_pw_multi_aff *isl_union_pw_multi_aff_tile(
424 __isl_take isl_union_pw_multi_aff *sched,
425 __isl_keep isl_multi_val *sizes)
427 isl_ctx *ctx;
428 isl_space *space;
429 struct isl_band_tile_data data = { sizes };
431 ctx = isl_multi_val_get_ctx(sizes);
433 space = isl_union_pw_multi_aff_get_space(sched);
434 data.res = isl_union_pw_multi_aff_empty(space);
435 data.scale = isl_options_get_tile_scale_tile_loops(ctx);
437 if (isl_union_pw_multi_aff_foreach_pw_multi_aff(sched,
438 &pw_multi_aff_tile, &data) < 0)
439 goto error;
441 isl_union_pw_multi_aff_free(sched);
442 return data.res;
443 error:
444 isl_union_pw_multi_aff_free(sched);
445 isl_union_pw_multi_aff_free(data.res);
446 return NULL;
449 /* Extract the range space from "pma" and store it in *user.
450 * All entries are expected to have the same range space, so we can
451 * stop after extracting the range space from the first entry.
453 static int extract_range_space(__isl_take isl_pw_multi_aff *pma, void *user)
455 isl_space **space = user;
457 *space = isl_space_range(isl_pw_multi_aff_get_space(pma));
458 isl_pw_multi_aff_free(pma);
460 return -1;
463 /* Extract the range space of "band". All entries in band->pma should
464 * have the same range space. Furthermore, band->pma should have at least
465 * one entry.
467 static __isl_give isl_space *band_get_range_space(__isl_keep isl_band *band)
469 isl_space *space;
471 if (!band)
472 return NULL;
474 space = NULL;
475 isl_union_pw_multi_aff_foreach_pw_multi_aff(band->pma,
476 &extract_range_space, &space);
478 return space;
481 /* Construct and return an isl_multi_val in the given space, with as entries
482 * the first elements of "v", padded with ones if the size of "v" is smaller
483 * than the dimension of "space".
485 static __isl_give isl_multi_val *multi_val_from_vec(__isl_take isl_space *space,
486 __isl_take isl_vec *v)
488 isl_ctx *ctx;
489 isl_multi_val *mv;
490 int i, n, size;
492 if (!space || !v)
493 goto error;
495 ctx = isl_space_get_ctx(space);
496 mv = isl_multi_val_zero(space);
497 n = isl_multi_val_dim(mv, isl_dim_set);
498 size = isl_vec_size(v);
499 if (n < size)
500 size = n;
502 for (i = 0; i < size; ++i) {
503 isl_val *val = isl_vec_get_element_val(v, i);
504 mv = isl_multi_val_set_val(mv, i, val);
506 for (i = size; i < n; ++i)
507 mv = isl_multi_val_set_val(mv, i, isl_val_one(ctx));
509 isl_vec_free(v);
510 return mv;
511 error:
512 isl_space_free(space);
513 isl_vec_free(v);
514 return NULL;
517 /* Tile the given band using the specified tile sizes.
518 * The given band is modified to refer to the tile loops and
519 * a child band is created to refer to the point loops.
520 * The children of this point loop band are the children
521 * of the original band.
523 * If the scale tile loops option is set, then the tile loops
524 * are scaled by the tile sizes. If the shift point loops option is set,
525 * then the point loops are shifted to start at zero.
526 * In particular, these options affect the tile and point loop schedules
527 * as follows
529 * scale shift original tile point
531 * 0 0 i floor(i/s) i
532 * 1 0 i s * floor(i/s) i
533 * 0 1 i floor(i/s) i - s * floor(i/s)
534 * 1 1 i s * floor(i/s) i - s * floor(i/s)
536 int isl_band_tile(__isl_keep isl_band *band, __isl_take isl_vec *sizes)
538 isl_ctx *ctx;
539 isl_band *child;
540 isl_band_list *list = NULL;
541 isl_union_pw_multi_aff *sched = NULL, *child_sched = NULL;
542 isl_space *space;
543 isl_multi_val *mv_sizes;
545 if (!band || !sizes)
546 goto error;
548 ctx = isl_vec_get_ctx(sizes);
549 child = isl_band_dup(band);
550 list = isl_band_list_alloc(ctx, 1);
551 list = isl_band_list_add(list, child);
552 if (!list)
553 goto error;
555 space = band_get_range_space(band);
556 mv_sizes = multi_val_from_vec(space, isl_vec_copy(sizes));
557 sched = isl_union_pw_multi_aff_copy(band->pma);
558 sched = isl_union_pw_multi_aff_tile(sched, mv_sizes);
560 child_sched = isl_union_pw_multi_aff_copy(child->pma);
561 if (isl_options_get_tile_shift_point_loops(ctx)) {
562 isl_union_pw_multi_aff *scaled;
563 scaled = isl_union_pw_multi_aff_copy(sched);
564 if (!isl_options_get_tile_scale_tile_loops(ctx))
565 scaled = isl_union_pw_multi_aff_scale_multi_val(scaled,
566 isl_multi_val_copy(mv_sizes));
567 child_sched = isl_union_pw_multi_aff_sub(child_sched, scaled);
569 isl_multi_val_free(mv_sizes);
570 if (!sched || !child_sched)
571 goto error;
573 child->children = band->children;
574 band->children = list;
575 child->parent = band;
576 isl_union_pw_multi_aff_free(band->pma);
577 band->pma = sched;
578 isl_union_pw_multi_aff_free(child->pma);
579 child->pma = child_sched;
581 isl_vec_free(sizes);
582 return 0;
583 error:
584 isl_union_pw_multi_aff_free(sched);
585 isl_union_pw_multi_aff_free(child_sched);
586 isl_band_list_free(list);
587 isl_vec_free(sizes);
588 return -1;
591 /* Internal data structure used inside isl_union_pw_multi_aff_drop.
593 * "pos" is the position of the first dimension to drop.
594 * "n" is the number of dimensions to drop.
595 * "res" accumulates the result.
597 struct isl_union_pw_multi_aff_drop_data {
598 int pos;
599 int n;
600 isl_union_pw_multi_aff *res;
603 /* Drop the data->n output dimensions starting at data->pos from "pma"
604 * and add the result to data->res.
606 static int pw_multi_aff_drop(__isl_take isl_pw_multi_aff *pma, void *user)
608 struct isl_union_pw_multi_aff_drop_data *data = user;
610 pma = isl_pw_multi_aff_drop_dims(pma, isl_dim_out, data->pos, data->n);
612 data->res = isl_union_pw_multi_aff_add_pw_multi_aff(data->res, pma);
613 if (!data->res)
614 return -1;
616 return 0;
619 /* Drop the "n" output dimensions starting at "pos" from "sched".
621 static isl_union_pw_multi_aff *isl_union_pw_multi_aff_drop(
622 __isl_take isl_union_pw_multi_aff *sched, int pos, int n)
624 isl_space *space;
625 struct isl_union_pw_multi_aff_drop_data data = { pos, n };
627 space = isl_union_pw_multi_aff_get_space(sched);
628 data.res = isl_union_pw_multi_aff_empty(space);
630 if (isl_union_pw_multi_aff_foreach_pw_multi_aff(sched,
631 &pw_multi_aff_drop, &data) < 0)
632 data.res = isl_union_pw_multi_aff_free(data.res);
634 isl_union_pw_multi_aff_free(sched);
635 return data.res;
638 /* Drop the "n" dimensions starting at "pos" from "band".
640 static int isl_band_drop(__isl_keep isl_band *band, int pos, int n)
642 int i;
643 isl_union_pw_multi_aff *sched;
645 if (!band)
646 return -1;
647 if (n == 0)
648 return 0;
650 sched = isl_union_pw_multi_aff_copy(band->pma);
651 sched = isl_union_pw_multi_aff_drop(sched, pos, n);
652 if (!sched)
653 return -1;
655 isl_union_pw_multi_aff_free(band->pma);
656 band->pma = sched;
658 for (i = pos + n; i < band->n; ++i)
659 band->coincident[i - n] = band->coincident[i];
661 band->n -= n;
663 return 0;
666 /* Split the given band into two nested bands, one with the first "pos"
667 * dimensions of "band" and one with the remaining band->n - pos dimensions.
669 int isl_band_split(__isl_keep isl_band *band, int pos)
671 isl_ctx *ctx;
672 isl_band *child;
673 isl_band_list *list;
675 if (!band)
676 return -1;
678 ctx = isl_band_get_ctx(band);
680 if (pos < 0 || pos > band->n)
681 isl_die(ctx, isl_error_invalid, "position out of bounds",
682 return -1);
684 child = isl_band_dup(band);
685 if (isl_band_drop(child, 0, pos) < 0)
686 child = isl_band_free(child);
687 list = isl_band_list_alloc(ctx, 1);
688 list = isl_band_list_add(list, child);
689 if (!list)
690 return -1;
692 if (isl_band_drop(band, pos, band->n - pos) < 0) {
693 isl_band_list_free(list);
694 return -1;
697 child->children = band->children;
698 band->children = list;
699 child->parent = band;
701 return 0;
704 __isl_give isl_printer *isl_printer_print_band(__isl_take isl_printer *p,
705 __isl_keep isl_band *band)
707 isl_union_map *prefix, *partial, *suffix;
709 prefix = isl_band_get_prefix_schedule(band);
710 partial = isl_band_get_partial_schedule(band);
711 suffix = isl_band_get_suffix_schedule(band);
713 p = isl_printer_print_str(p, "(");
714 p = isl_printer_print_union_map(p, prefix);
715 p = isl_printer_print_str(p, ",");
716 p = isl_printer_print_union_map(p, partial);
717 p = isl_printer_print_str(p, ",");
718 p = isl_printer_print_union_map(p, suffix);
719 p = isl_printer_print_str(p, ")");
721 isl_union_map_free(prefix);
722 isl_union_map_free(partial);
723 isl_union_map_free(suffix);
725 return p;