| 95 | */ |
| 96 | template <int ndim, typename SplitFactor, typename OnBlockFunc> |
| 97 | void ForEachBlock(TensorShape<ndim> start, TensorShape<ndim> end, const SplitFactor& split_factor, |
| 98 | int d, int max_split_dim, OnBlockFunc&& func) { |
| 99 | assert(start.size() == end.size()); |
| 100 | if (d > max_split_dim || d == start.size()) { |
| 101 | func(start, end); |
| 102 | return; |
| 103 | } |
| 104 | |
| 105 | if (split_factor[d] == 1) { |
| 106 | ForEachBlock(start, end, split_factor, d + 1, max_split_dim, func); |
| 107 | } else { |
| 108 | int64_t start_d = start[d]; |
| 109 | int64_t extent_d = end[d] - start_d; |
| 110 | int nblocks_d = split_factor[d]; |
| 111 | int64_t prev_end = start_d; |
| 112 | for (int b = 0; b < nblocks_d; b++) { |
| 113 | start[d] = prev_end; |
| 114 | end[d] = prev_end = extent_d * (b + 1) / nblocks_d + start_d; |
| 115 | ForEachBlock(start, end, split_factor, d + 1, max_split_dim, func); |
| 116 | } |
| 117 | } |
| 118 | } |
| 119 | |
| 120 | } // namespace kernels |
| 121 | } // namespace dali |
no test coverage detected