Spaces:
Sleeping
Sleeping
File size: 2,294 Bytes
66c9c8a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 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 | #pragma once
namespace wp
{
// All iterable types should implement 3 methods:
//
// T iter_next(iter) - returns the current value and moves iterator to next state
// int iter_cmp(iter) - returns 0 if finished
// iter iter_reverse(iter) - return an iterator of the same type representing the reverse order
//
// iter_next() should also be registered as a built-in hidden function so that code-gen
// can call it and generate the appropriate variable storage
// represents a built-in Python range() loop
struct range_t
{
CUDA_CALLABLE range_t() {}
CUDA_CALLABLE range_t(int) {} // for backward pass
int start;
int end;
int step;
int i;
};
CUDA_CALLABLE inline range_t range(int end)
{
range_t r;
r.start = 0;
r.end = end;
r.step = 1;
r.i = r.start;
return r;
}
CUDA_CALLABLE inline range_t range(int start, int end)
{
range_t r;
r.start = start;
r.end = end;
r.step = 1;
r.i = r.start;
return r;
}
CUDA_CALLABLE inline range_t range(int start, int end, int step)
{
range_t r;
r.start = start;
r.end = end;
r.step = step;
r.i = r.start;
return r;
}
CUDA_CALLABLE inline void adj_range(int end, int adj_end, range_t& adj_ret) {}
CUDA_CALLABLE inline void adj_range(int start, int end, int adj_start, int adj_end, range_t& adj_ret) {}
CUDA_CALLABLE inline void adj_range(int start, int end, int step, int adj_start, int adj_end, int adj_step, range_t& adj_ret) {}
CUDA_CALLABLE inline int iter_next(range_t& r)
{
int iter = r.i;
r.i += r.step;
return iter;
}
CUDA_CALLABLE inline bool iter_cmp(const range_t& r)
{
// implements for-loop comparison to emulate Python range() loops with negative arguments
if (r.step == 0)
// degenerate case where step == 0
return false;
if (r.step > 0)
// normal case where step > 0
return r.i < r.end;
else
// reverse case where step < 0
return r.i > r.end;
}
CUDA_CALLABLE inline range_t iter_reverse(const range_t& r)
{
// generates a reverse range, equivalent to reversed(range())
range_t rev;
rev.start = r.end-1;
rev.end = r.start-1;
rev.step = -r.step;
rev.i = rev.start;
return rev;
}
} // namespace wp |