1 /* 2 * Copyright (c) 2019-2021 Arm Limited. 3 * 4 * SPDX-License-Identifier: MIT 5 * 6 * Permission is hereby granted, free of charge, to any person obtaining a copy 7 * of this software and associated documentation files (the "Software"), to 8 * deal in the Software without restriction, including without limitation the 9 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 10 * sell copies of the Software, and to permit persons to whom the Software is 11 * furnished to do so, subject to the following conditions: 12 * 13 * The above copyright notice and this permission notice shall be included in all 14 * copies or substantial portions of the Software. 15 * 16 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 17 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 18 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 19 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 20 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 21 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 22 * SOFTWARE. 23 */ 24 #pragma once 25 26 #include <algorithm> 27 #include <array> 28 #include <cassert> 29 #include <initializer_list> 30 31 namespace arm_gemm 32 { 33 template <unsigned int D> 34 class NDRange 35 { 36 private: 37 std::array<unsigned int, D> m_sizes{}; 38 std::array<unsigned int, D> m_totalsizes{}; 39 40 class NDRangeIterator 41 { 42 private: 43 const NDRange &m_parent; 44 unsigned int m_pos = 0; 45 unsigned int m_end = 0; 46 47 public: NDRangeIterator(const NDRange & p,unsigned int s,unsigned int e)48 NDRangeIterator(const NDRange &p, unsigned int s, unsigned int e) 49 : m_parent(p), m_pos(s), m_end(e) 50 { 51 } 52 done() const53 bool done() const 54 { 55 return (m_pos >= m_end); 56 } 57 dim(unsigned int d) const58 unsigned int dim(unsigned int d) const 59 { 60 unsigned int r = m_pos; 61 62 if(d < (D - 1)) 63 { 64 r %= m_parent.m_totalsizes[d]; 65 } 66 67 if(d > 0) 68 { 69 r /= m_parent.m_totalsizes[d - 1]; 70 } 71 72 return r; 73 } 74 next_dim0()75 bool next_dim0() 76 { 77 m_pos++; 78 79 return !done(); 80 } 81 next_dim1()82 bool next_dim1() 83 { 84 m_pos += m_parent.m_sizes[0] - dim(0); 85 86 return !done(); 87 } 88 dim0_max() const89 unsigned int dim0_max() const 90 { 91 unsigned int offset = std::min(m_end - m_pos, m_parent.m_sizes[0] - dim(0)); 92 93 return dim(0) + offset; 94 } 95 }; 96 set_totalsizes()97 void set_totalsizes() 98 { 99 unsigned int t = 1; 100 101 for(unsigned int i = 0; i < D; i++) 102 { 103 if(m_sizes[i] == 0) 104 { 105 m_sizes[i] = 1; 106 } 107 108 t *= m_sizes[i]; 109 110 m_totalsizes[i] = t; 111 } 112 } 113 114 public: 115 NDRange &operator=(const NDRange &rhs) = default; 116 NDRange(const NDRange &rhs) = default; 117 118 template <typename... T> NDRange(T...ts)119 NDRange(T... ts) 120 : m_sizes{ ts... } 121 { 122 set_totalsizes(); 123 } 124 NDRange(const std::array<unsigned int,D> & n)125 NDRange(const std::array<unsigned int, D> &n) 126 : m_sizes(n) 127 { 128 set_totalsizes(); 129 } 130 iterator(unsigned int start,unsigned int end) const131 NDRangeIterator iterator(unsigned int start, unsigned int end) const 132 { 133 return NDRangeIterator(*this, start, end); 134 } 135 total_size() const136 unsigned int total_size() const 137 { 138 return m_totalsizes[D - 1]; 139 } 140 get_size(unsigned int v) const141 unsigned int get_size(unsigned int v) const 142 { 143 return m_sizes[v]; 144 } 145 }; 146 147 /** NDCoordinate builds upon a range, but specifies a starting position 148 * in addition to a size which it inherits from NDRange 149 */ 150 template <unsigned int N> 151 class NDCoordinate : public NDRange<N> 152 { 153 using int_t = unsigned int; 154 using ndrange_t = NDRange<N>; 155 156 std::array<int_t, N> m_positions{}; 157 158 public: 159 NDCoordinate &operator=(const NDCoordinate &rhs) = default; 160 NDCoordinate(const NDCoordinate &rhs) = default; NDCoordinate(const std::initializer_list<std::pair<int_t,int_t>> & list)161 NDCoordinate(const std::initializer_list<std::pair<int_t, int_t>> &list) 162 { 163 std::array<int_t, N> sizes{}; 164 165 std::size_t i = 0; 166 for(auto &p : list) 167 { 168 m_positions[i] = p.first; 169 sizes[i++] = p.second; 170 } 171 172 //update the parents sizes 173 static_cast<ndrange_t &>(*this) = ndrange_t(sizes); 174 } 175 get_position(int_t d) const176 int_t get_position(int_t d) const 177 { 178 assert(d < N); 179 180 return m_positions[d]; 181 } 182 set_position(int_t d,int_t v)183 void set_position(int_t d, int_t v) 184 { 185 assert(d < N); 186 187 m_positions[d] = v; 188 } 189 get_position_end(int_t d) const190 int_t get_position_end(int_t d) const 191 { 192 return get_position(d) + ndrange_t::get_size(d); 193 } 194 }; //class NDCoordinate 195 196 using ndrange_t = NDRange<6>; 197 using ndcoord_t = NDCoordinate<6>; 198 199 } // namespace arm_gemm 200