blob: 0053e7924361baf04f167e37c534aaf7590faec5 [file] [log] [blame]
Tim Hall79d07d22020-04-27 18:20:16 +01001# Copyright (C) 2020 Arm Limited or its affiliates. All rights reserved.
2#
3# SPDX-License-Identifier: Apache-2.0
4#
5# Licensed under the Apache License, Version 2.0 (the License); you may
6# not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9# www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an AS IS BASIS, WITHOUT
13# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
Tim Hall79d07d22020-04-27 18:20:16 +010016# Description:
17# Contains classes that hold commands for the high-level command stream (one command per DMA or NPU stripe).
Diego Russoea6111a2020-04-14 18:41:58 +010018from enum import IntEnum
19
Tim Hall79d07d22020-04-27 18:20:16 +010020import numpy as np
Diego Russoea6111a2020-04-14 18:41:58 +010021
Tim Hall79d07d22020-04-27 18:20:16 +010022from .numeric_util import round_up_divide
Diego Russoe8a10452020-04-21 17:39:10 +010023from .operation import NpuBlockType
24from .range_set import AccessDirection
25from .range_set import MemoryAccessSet
Tim Hall79d07d22020-04-27 18:20:16 +010026
27
28class Box:
29 def __init__(self, start_coord, end_coord):
30 self.start_coord = list(start_coord)
31 self.end_coord = list(end_coord)
32 assert len(self.start_coord) == len(end_coord)
33 for i in range(len(self.start_coord)):
34 assert self.start_coord[i] <= self.end_coord[i]
35
36 def transform_with_strides_and_skirt(
Tim Hallc30f4952020-06-15 20:47:35 +010037 self,
38 strides,
39 skirt,
40 ifm_shape,
41 npu_block_type,
42 concat_axis=0,
43 concat_offset=0,
44 split_offset=None,
45 k_height=1,
46 upscaling_factor=1,
Tim Hall79d07d22020-04-27 18:20:16 +010047 ):
48 new_start_coord = list(self.start_coord)
49 new_end_coord = list(self.end_coord)
50
51 new_start_coord[concat_axis] -= concat_offset
52 new_end_coord[concat_axis] -= concat_offset
53
Diego Russoea6111a2020-04-14 18:41:58 +010054 if split_offset is not None:
Tim Hall79d07d22020-04-27 18:20:16 +010055 for idx in range(len(split_offset)):
56 new_start_coord[idx] += split_offset[idx]
57 new_end_coord[idx] += split_offset[idx]
58
Diego Russoea6111a2020-04-14 18:41:58 +010059 if split_offset is None and npu_block_type in set((NpuBlockType.ConvolutionMxN, NpuBlockType.VectorProduct)):
Tim Hall79d07d22020-04-27 18:20:16 +010060 # these types of operations do a "dot product" over the entire IFM
61 new_start_coord[-1] = 0
62 new_end_coord[-1] = ifm_shape[-1]
63
Louis Verhaarde0ef2732020-06-03 08:56:44 +020064 if npu_block_type == NpuBlockType.ElementWise and min(len(new_end_coord), len(ifm_shape)) >= 1:
65 new_end_coord[-1] = min(new_end_coord[-1], ifm_shape[-1])
Tim Hall79d07d22020-04-27 18:20:16 +010066 if min(len(new_end_coord), len(ifm_shape)) >= 2:
Jacob Bohlin9b64ba02020-07-07 17:15:22 +020067 new_end_coord[-2] = min(new_end_coord[-2], ifm_shape[-2] * upscaling_factor)
Tim Hall79d07d22020-04-27 18:20:16 +010068 if min(len(new_end_coord), len(ifm_shape)) >= 3:
Jacob Bohlin9b64ba02020-07-07 17:15:22 +020069 original_end_coord = list(new_end_coord)
70 new_end_coord[-3] = min(new_end_coord[-3], ifm_shape[-3] * upscaling_factor)
Tim Hall79d07d22020-04-27 18:20:16 +010071
72 pad_top = 0
73 pad_bottom = 0
74 if strides is not None and skirt is not None:
75 if len(new_start_coord) >= 2:
76 stride = strides[2]
77 new_start_coord[-2] = max(new_start_coord[-2] * stride - skirt[1], 0)
78 new_end_coord[-2] = min(new_end_coord[-2] * stride + skirt[3], ifm_shape[-2])
79
80 if len(new_start_coord) >= 3:
81 stride = strides[1]
Jacob Bohlin9b64ba02020-07-07 17:15:22 +020082 skirt_top_remainder = skirt[0] % upscaling_factor
Tim Hall79d07d22020-04-27 18:20:16 +010083
84 total_stride = stride * (new_end_coord[-3] - new_start_coord[-3] - 1)
Jacob Bohlin9b64ba02020-07-07 17:15:22 +020085 new_start_coord[-3] = new_start_coord[-3] * stride - skirt[0] + skirt_top_remainder
Tim Hall79d07d22020-04-27 18:20:16 +010086
Jacob Bohlin9b64ba02020-07-07 17:15:22 +020087 pad_top = max(0, 0 - new_start_coord[-3]) + skirt_top_remainder
Tim Hall79d07d22020-04-27 18:20:16 +010088 new_start_coord[-3] = max(new_start_coord[-3], 0)
89
90 while len(ifm_shape) < 3:
91 ifm_shape = [1] + ifm_shape
Jacob Bohlin9b64ba02020-07-07 17:15:22 +020092
93 if (new_end_coord[-3] * stride + skirt[2]) > (ifm_shape[-3] * upscaling_factor):
Tim Hall79d07d22020-04-27 18:20:16 +010094 # pad_bottom is calculated based the diff between the end position of the weight kernel,
95 # after last stride and the ifm height.
Jacob Bohlin9b64ba02020-07-07 17:15:22 +020096 if upscaling_factor != 1 and original_end_coord[-3] > ifm_shape[-3] * upscaling_factor:
97 # Special case for Transpose Convolution with VALID padding.
98 pad_bottom = original_end_coord[-3] - (ifm_shape[-3] * upscaling_factor)
99 else:
100 k_start = new_start_coord[-3] - pad_top
101 pad_bottom = max(0, k_start + total_stride + k_height - (ifm_shape[-3] * upscaling_factor))
Tim Hall79d07d22020-04-27 18:20:16 +0100102
Jacob Bohlin9b64ba02020-07-07 17:15:22 +0200103 # Adjust for upscaling
104 new_start_coord[-3] = max(new_start_coord[-3] // upscaling_factor, 0)
105 new_end_coord[-3] = new_end_coord[-3] * stride + skirt[2] + (skirt[2] % upscaling_factor)
106 new_end_coord[-3] = min(new_end_coord[-3] // upscaling_factor, ifm_shape[-3])
Tim Hall79d07d22020-04-27 18:20:16 +0100107
108 return Box(new_start_coord, new_end_coord), pad_top, pad_bottom
109
110 def make_weight_box(weight_shape, npu_block_type, oc_range_start=None, oc_range_end=None, weights_transposed=False):
111 start = [0] * len(weight_shape)
112 end = list(weight_shape)
113 if oc_range_start is not None and oc_range_end is not None:
114 if npu_block_type == NpuBlockType.ConvolutionDepthWise:
115 # input range is output range divided by channel multiplier
116 if weights_transposed:
117 start[-1] = oc_range_start // weight_shape[-2]
118 end[-1] = oc_range_end // weight_shape[-2]
119 else:
120 start[-2] = oc_range_start // weight_shape[-1]
121 end[-2] = oc_range_end // weight_shape[-1]
122 else:
123 start[-1] = oc_range_start
124 end[-1] = oc_range_end
125 for i in range(len(end)):
126 assert 0 <= start[i] < weight_shape[i]
127 assert 0 < end[i] <= weight_shape[i]
128
129 return Box(start, end)
130
131 def get_size_shape(self):
132 return [int(self.end_coord[i] - self.start_coord[i]) for i in range(len(self.end_coord))]
133
134 def get_size(self):
135 return int(np.prod(self.get_size_shape()))
136
137 def __str__(self):
138 return "<Box %s - %s>" % (self.start_coord, self.end_coord)
139
140 __repr__ = __str__
141
142
143class CommandType(IntEnum):
144 NpuStripe = 0
145 DMA = 1
146 Size = 2
147
148
149class Command:
150 def get_ofm_y_range_for_pass(self, ps_requested):
151 return None
152
153 def is_npu_pass_command(self):
154 return False
155
156 def get_memory_accesses(self):
157 return None
158
159 def get_operation_count(self):
160 # returns numpy array of (DPU blocks, dma_ops). Should line up with the CommandType enum
161 return np.array((0, 0))
162
163
164class NpuStripe(Command):
165 def __init__(
166 self,
167 ps,
168 block_config,
169 is_first,
170 is_last,
171 is_first_h_stripe,
172 is_last_h_stripe,
173 ifm_tensor,
174 ifm_box,
175 ofm_tensor,
176 ofm_box,
177 weight_tensor=None,
178 weight_box=None,
179 scale_tensor=None,
180 concat_axis=0,
181 concat_offset=0,
182 ifm2_tensor=None,
183 ifm2_box=None,
184 pad_top=0,
185 pad_bottom=0,
186 ):
187 self.cmdtype = CommandType.NpuStripe
188 self.ps = ps
189 self.block_config = block_config
190 self.is_first = is_first
191 self.is_last = is_last
192 self.is_first_h_stripe = is_first_h_stripe
193 self.is_last_h_stripe = is_last_h_stripe
194 self.ifm_tensor = ifm_tensor
195 self.ifm_box = ifm_box
196 self.ifm2_tensor = ifm2_tensor
197 self.ifm2_box = ifm2_box
198 self.ofm_tensor = ofm_tensor
199 self.ofm_box = ofm_box
200 self.weight_tensor = weight_tensor
201 self.scale_tensor = scale_tensor
202 self.weight_box = weight_box
203 self.concat_axis = concat_axis
204 self.concat_offset = concat_offset
205 self.pad_top = pad_top
206 self.pad_bottom = pad_bottom
207 for i in range(len(self.ofm_box.end_coord)):
208 assert self.ofm_box.end_coord[i] <= self.ofm_tensor.shape[i]
209
210 def get_memory_accesses(self):
211 res = MemoryAccessSet()
212 if self.ifm_tensor is not None and self.ifm_tensor.shape != []:
213 res.add(
214 self.ifm_tensor.get_address_ranges_for_coordinates(self.ifm_box.start_coord, self.ifm_box.end_coord),
215 AccessDirection.Read,
216 )
217 if self.ifm2_tensor is not None and self.ifm2_tensor.shape != []:
218 res.add(
219 self.ifm2_tensor.get_address_ranges_for_coordinates(self.ifm2_box.start_coord, self.ifm2_box.end_coord),
220 AccessDirection.Read,
221 )
222 if self.ofm_tensor is not None:
223 res.add(
224 self.ofm_tensor.get_address_ranges_for_coordinates(self.ofm_box.start_coord, self.ofm_box.end_coord),
225 AccessDirection.Write,
226 )
227 if self.weight_tensor is not None:
228 res.add(
229 self.weight_tensor.get_address_ranges_for_coordinates(
230 self.weight_box.start_coord, self.weight_box.end_coord
231 ),
232 AccessDirection.Read,
233 )
234 return res
235
236 def is_npu_pass_command(self):
237 return True
238
239 def __str__(self):
240 return "<NPUStripe: ps=%s, ifm_box=%s, ifm2_box=%s, ofm_box=%s, weight_box=%s, block_config=%s>" % (
241 self.ps.name,
242 self.ifm_box,
243 self.ifm2_box,
244 self.ofm_box,
245 self.weight_box,
246 self.block_config,
247 )
248
249 __repr__ = __str__
250
251 def get_ofm_y_range_for_pass(self, ps_requested):
252 if ps_requested != self.ps:
253 return None
254 if len(self.ofm_box.start_coord) >= 3:
255 return (self.ofm_box.start_coord[-3], self.ofm_box.end_coord[-3])
256 return None
257
258 def get_block_dimensions(self):
259 ofm_box = self.ofm_box
260 block_config = self.block_config
261
262 out_height = 1
263 out_width = 1
264 out_depth = ofm_box.end_coord[-1] - ofm_box.start_coord[-1]
265 if len(ofm_box.end_coord) >= 4:
266 out_width = ofm_box.end_coord[-2] - ofm_box.start_coord[-2]
267 out_height = ofm_box.end_coord[-3] - ofm_box.start_coord[-3]
268
269 assert out_height >= 0
270 assert out_width >= 0
271 assert out_depth >= 0
272 return (
273 round_up_divide(out_height, block_config[0]),
274 round_up_divide(out_width, block_config[1]),
275 round_up_divide(out_depth, block_config[3]),
276 )
277
278 def get_operation_count(self):
279 # returns numpy array of (DPU blocks, dma_ops)
280 return np.array((self.get_n_blocks(), 0))
281
282 def get_n_blocks(self):
283 h, w, d = self.get_block_dimensions()
284 res = h * w * d
285 assert res >= 0
286 return res
287
288 def get_single_block_command(self, block_idx):
289 block_cfg = (self.block_config[0], self.block_config[1], self.block_config[3])
290 dims = self.get_block_dimensions()
291 strides = dims[1] * dims[2], dims[2], 1
292 coord = []
293 idx_left = block_idx
294 for s in strides:
295 c = idx_left // s
296 idx_left -= c * s
297 coord.append(c)
298
299 assert idx_left == 0
300
301 # put in dummy height/widths in case we're dealing with FC layers
302 ofm_start = list(self.ofm_box.start_coord)
303 ofm_end = list(self.ofm_box.end_coord)
304
305 # cut out a nice block shape
306 for idx in (-1, -2, -3):
307 if len(ofm_start) >= -idx:
308 ofm_start[idx] += block_cfg[idx] * coord[idx]
309 ofm_end[idx] = min(ofm_end[idx], ofm_start[idx] + block_cfg[idx])
310
311 ps = self.ps
312 strides = None
313 skirt = None
314 if ps.primary_op is not None:
315 strides = ps.primary_op.attrs.get("strides", None)
316 skirt = ps.primary_op.attrs.get("skirt", None)
317 npu_block_type = ps.npu_block_type
318
319 ofm_box = Box(ofm_start, ofm_end)
320 ifm_box, _, _ = ofm_box.transform_with_strides_and_skirt(
321 strides, skirt, self.ifm_tensor.shape, npu_block_type, self.concat_axis, self.concat_offset
322 )
323
324 weight_box = None
325 if self.weight_tensor is not None:
326 weight_oc_start = ofm_start[-1]
327 weight_oc_end = ofm_end[-1]
328 if self.concat_axis - len(self.weight_tensor.shape) == -1:
329 weight_oc_start -= self.concat_offset
330 weight_oc_end -= self.concat_offset
331
332 weight_box = Box.make_weight_box(
333 self.weight_tensor.shape,
334 npu_block_type,
335 weight_oc_start,
336 weight_oc_end,
337 self.weight_tensor.weight_transpose_depthwise,
338 )
339
340 return NpuStripe(
341 self.ps,
342 self.block_config,
343 self.is_first,
344 self.is_last,
345 self.is_first_h_stripe,
346 self.is_last_h_stripe,
347 self.ifm_tensor,
348 ifm_box,
349 self.ofm_tensor,
350 ofm_box,
351 self.weight_tensor,
352 weight_box,
353 self.scale_tensor,
354 self.concat_axis,
355 self.concat_offset,
356 )
357
358
359class DMA(Command):
360 def __init__(self, in_tensor, out_tensor, box):
361 self.cmdtype = CommandType.DMA
362 self.in_tensor = in_tensor
363 self.out_tensor = out_tensor
364 self.box = box
365
366 def __str__(self):
367 return "<DMA: in=%s, out=%s, box=%s>" % (self.in_tensor.name, self.out_tensor.name, self.box)
368
369 __repr__ = __str__
370
371 def get_memory_accesses(self):
372 res = MemoryAccessSet()
373
374 res.add(
375 self.in_tensor.get_address_ranges_for_coordinates(self.box.start_coord, self.box.end_coord),
376 AccessDirection.Read,
377 )
378 res.add(
379 self.out_tensor.get_address_ranges_for_coordinates(self.box.start_coord, self.box.end_coord),
380 AccessDirection.Write,
381 )
382 return res
383
384 def get_operation_count(self):
385 # returns numpy array of (DPU blocks, dma_ops)
386 return np.array((0, 1))