blob: c6698297118eb92edf3c177010d7880e35ce71fb [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
Fredrik Svedberga0c36242020-06-03 15:43:31 +020059 if split_offset is None and npu_block_type in set(
60 (NpuBlockType.ConvolutionMxN, NpuBlockType.VectorProduct, NpuBlockType.ReduceSum)
61 ):
62 # these types of operations do a "dot product" or sum over the entire IFM
Tim Hall79d07d22020-04-27 18:20:16 +010063 new_start_coord[-1] = 0
64 new_end_coord[-1] = ifm_shape[-1]
65
Louis Verhaarde0ef2732020-06-03 08:56:44 +020066 if npu_block_type == NpuBlockType.ElementWise and min(len(new_end_coord), len(ifm_shape)) >= 1:
67 new_end_coord[-1] = min(new_end_coord[-1], ifm_shape[-1])
Tim Hall79d07d22020-04-27 18:20:16 +010068 if min(len(new_end_coord), len(ifm_shape)) >= 2:
Jacob Bohlin9b64ba02020-07-07 17:15:22 +020069 new_end_coord[-2] = min(new_end_coord[-2], ifm_shape[-2] * upscaling_factor)
Tim Hall79d07d22020-04-27 18:20:16 +010070 if min(len(new_end_coord), len(ifm_shape)) >= 3:
Jacob Bohlin9b64ba02020-07-07 17:15:22 +020071 original_end_coord = list(new_end_coord)
72 new_end_coord[-3] = min(new_end_coord[-3], ifm_shape[-3] * upscaling_factor)
Tim Hall79d07d22020-04-27 18:20:16 +010073
74 pad_top = 0
75 pad_bottom = 0
76 if strides is not None and skirt is not None:
77 if len(new_start_coord) >= 2:
78 stride = strides[2]
79 new_start_coord[-2] = max(new_start_coord[-2] * stride - skirt[1], 0)
80 new_end_coord[-2] = min(new_end_coord[-2] * stride + skirt[3], ifm_shape[-2])
81
82 if len(new_start_coord) >= 3:
83 stride = strides[1]
Jacob Bohlin9b64ba02020-07-07 17:15:22 +020084 skirt_top_remainder = skirt[0] % upscaling_factor
Tim Hall79d07d22020-04-27 18:20:16 +010085
86 total_stride = stride * (new_end_coord[-3] - new_start_coord[-3] - 1)
Jacob Bohlin9b64ba02020-07-07 17:15:22 +020087 new_start_coord[-3] = new_start_coord[-3] * stride - skirt[0] + skirt_top_remainder
Tim Hall79d07d22020-04-27 18:20:16 +010088
Jacob Bohlin9b64ba02020-07-07 17:15:22 +020089 pad_top = max(0, 0 - new_start_coord[-3]) + skirt_top_remainder
Tim Hall79d07d22020-04-27 18:20:16 +010090 new_start_coord[-3] = max(new_start_coord[-3], 0)
91
92 while len(ifm_shape) < 3:
93 ifm_shape = [1] + ifm_shape
Jacob Bohlin9b64ba02020-07-07 17:15:22 +020094
95 if (new_end_coord[-3] * stride + skirt[2]) > (ifm_shape[-3] * upscaling_factor):
Tim Hall79d07d22020-04-27 18:20:16 +010096 # pad_bottom is calculated based the diff between the end position of the weight kernel,
97 # after last stride and the ifm height.
Jacob Bohlin9b64ba02020-07-07 17:15:22 +020098 if upscaling_factor != 1 and original_end_coord[-3] > ifm_shape[-3] * upscaling_factor:
99 # Special case for Transpose Convolution with VALID padding.
100 pad_bottom = original_end_coord[-3] - (ifm_shape[-3] * upscaling_factor)
101 else:
102 k_start = new_start_coord[-3] - pad_top
103 pad_bottom = max(0, k_start + total_stride + k_height - (ifm_shape[-3] * upscaling_factor))
Tim Hall79d07d22020-04-27 18:20:16 +0100104
Jacob Bohlin9b64ba02020-07-07 17:15:22 +0200105 # Adjust for upscaling
106 new_start_coord[-3] = max(new_start_coord[-3] // upscaling_factor, 0)
107 new_end_coord[-3] = new_end_coord[-3] * stride + skirt[2] + (skirt[2] % upscaling_factor)
108 new_end_coord[-3] = min(new_end_coord[-3] // upscaling_factor, ifm_shape[-3])
Tim Hall79d07d22020-04-27 18:20:16 +0100109
110 return Box(new_start_coord, new_end_coord), pad_top, pad_bottom
111
112 def make_weight_box(weight_shape, npu_block_type, oc_range_start=None, oc_range_end=None, weights_transposed=False):
113 start = [0] * len(weight_shape)
114 end = list(weight_shape)
115 if oc_range_start is not None and oc_range_end is not None:
116 if npu_block_type == NpuBlockType.ConvolutionDepthWise:
117 # input range is output range divided by channel multiplier
118 if weights_transposed:
119 start[-1] = oc_range_start // weight_shape[-2]
120 end[-1] = oc_range_end // weight_shape[-2]
121 else:
122 start[-2] = oc_range_start // weight_shape[-1]
123 end[-2] = oc_range_end // weight_shape[-1]
124 else:
125 start[-1] = oc_range_start
126 end[-1] = oc_range_end
127 for i in range(len(end)):
128 assert 0 <= start[i] < weight_shape[i]
129 assert 0 < end[i] <= weight_shape[i]
130
131 return Box(start, end)
132
133 def get_size_shape(self):
134 return [int(self.end_coord[i] - self.start_coord[i]) for i in range(len(self.end_coord))]
135
136 def get_size(self):
137 return int(np.prod(self.get_size_shape()))
138
139 def __str__(self):
140 return "<Box %s - %s>" % (self.start_coord, self.end_coord)
141
142 __repr__ = __str__
143
144
145class CommandType(IntEnum):
146 NpuStripe = 0
147 DMA = 1
148 Size = 2
149
150
151class Command:
152 def get_ofm_y_range_for_pass(self, ps_requested):
153 return None
154
155 def is_npu_pass_command(self):
156 return False
157
158 def get_memory_accesses(self):
159 return None
160
161 def get_operation_count(self):
162 # returns numpy array of (DPU blocks, dma_ops). Should line up with the CommandType enum
163 return np.array((0, 0))
164
165
166class NpuStripe(Command):
167 def __init__(
168 self,
169 ps,
170 block_config,
171 is_first,
172 is_last,
173 is_first_h_stripe,
174 is_last_h_stripe,
175 ifm_tensor,
176 ifm_box,
177 ofm_tensor,
178 ofm_box,
179 weight_tensor=None,
180 weight_box=None,
181 scale_tensor=None,
182 concat_axis=0,
183 concat_offset=0,
184 ifm2_tensor=None,
185 ifm2_box=None,
186 pad_top=0,
187 pad_bottom=0,
188 ):
189 self.cmdtype = CommandType.NpuStripe
190 self.ps = ps
191 self.block_config = block_config
192 self.is_first = is_first
193 self.is_last = is_last
194 self.is_first_h_stripe = is_first_h_stripe
195 self.is_last_h_stripe = is_last_h_stripe
196 self.ifm_tensor = ifm_tensor
197 self.ifm_box = ifm_box
198 self.ifm2_tensor = ifm2_tensor
199 self.ifm2_box = ifm2_box
200 self.ofm_tensor = ofm_tensor
201 self.ofm_box = ofm_box
202 self.weight_tensor = weight_tensor
203 self.scale_tensor = scale_tensor
204 self.weight_box = weight_box
205 self.concat_axis = concat_axis
206 self.concat_offset = concat_offset
207 self.pad_top = pad_top
208 self.pad_bottom = pad_bottom
209 for i in range(len(self.ofm_box.end_coord)):
210 assert self.ofm_box.end_coord[i] <= self.ofm_tensor.shape[i]
211
212 def get_memory_accesses(self):
213 res = MemoryAccessSet()
214 if self.ifm_tensor is not None and self.ifm_tensor.shape != []:
215 res.add(
216 self.ifm_tensor.get_address_ranges_for_coordinates(self.ifm_box.start_coord, self.ifm_box.end_coord),
217 AccessDirection.Read,
218 )
219 if self.ifm2_tensor is not None and self.ifm2_tensor.shape != []:
220 res.add(
221 self.ifm2_tensor.get_address_ranges_for_coordinates(self.ifm2_box.start_coord, self.ifm2_box.end_coord),
222 AccessDirection.Read,
223 )
224 if self.ofm_tensor is not None:
225 res.add(
226 self.ofm_tensor.get_address_ranges_for_coordinates(self.ofm_box.start_coord, self.ofm_box.end_coord),
227 AccessDirection.Write,
228 )
229 if self.weight_tensor is not None:
230 res.add(
231 self.weight_tensor.get_address_ranges_for_coordinates(
232 self.weight_box.start_coord, self.weight_box.end_coord
233 ),
234 AccessDirection.Read,
235 )
236 return res
237
238 def is_npu_pass_command(self):
239 return True
240
241 def __str__(self):
242 return "<NPUStripe: ps=%s, ifm_box=%s, ifm2_box=%s, ofm_box=%s, weight_box=%s, block_config=%s>" % (
243 self.ps.name,
244 self.ifm_box,
245 self.ifm2_box,
246 self.ofm_box,
247 self.weight_box,
248 self.block_config,
249 )
250
251 __repr__ = __str__
252
253 def get_ofm_y_range_for_pass(self, ps_requested):
254 if ps_requested != self.ps:
255 return None
256 if len(self.ofm_box.start_coord) >= 3:
257 return (self.ofm_box.start_coord[-3], self.ofm_box.end_coord[-3])
258 return None
259
260 def get_block_dimensions(self):
261 ofm_box = self.ofm_box
262 block_config = self.block_config
263
264 out_height = 1
265 out_width = 1
266 out_depth = ofm_box.end_coord[-1] - ofm_box.start_coord[-1]
267 if len(ofm_box.end_coord) >= 4:
268 out_width = ofm_box.end_coord[-2] - ofm_box.start_coord[-2]
269 out_height = ofm_box.end_coord[-3] - ofm_box.start_coord[-3]
270
271 assert out_height >= 0
272 assert out_width >= 0
273 assert out_depth >= 0
274 return (
275 round_up_divide(out_height, block_config[0]),
276 round_up_divide(out_width, block_config[1]),
277 round_up_divide(out_depth, block_config[3]),
278 )
279
280 def get_operation_count(self):
281 # returns numpy array of (DPU blocks, dma_ops)
282 return np.array((self.get_n_blocks(), 0))
283
284 def get_n_blocks(self):
285 h, w, d = self.get_block_dimensions()
286 res = h * w * d
287 assert res >= 0
288 return res
289
290 def get_single_block_command(self, block_idx):
291 block_cfg = (self.block_config[0], self.block_config[1], self.block_config[3])
292 dims = self.get_block_dimensions()
293 strides = dims[1] * dims[2], dims[2], 1
294 coord = []
295 idx_left = block_idx
296 for s in strides:
297 c = idx_left // s
298 idx_left -= c * s
299 coord.append(c)
300
301 assert idx_left == 0
302
303 # put in dummy height/widths in case we're dealing with FC layers
304 ofm_start = list(self.ofm_box.start_coord)
305 ofm_end = list(self.ofm_box.end_coord)
306
307 # cut out a nice block shape
308 for idx in (-1, -2, -3):
309 if len(ofm_start) >= -idx:
310 ofm_start[idx] += block_cfg[idx] * coord[idx]
311 ofm_end[idx] = min(ofm_end[idx], ofm_start[idx] + block_cfg[idx])
312
313 ps = self.ps
314 strides = None
315 skirt = None
316 if ps.primary_op is not None:
317 strides = ps.primary_op.attrs.get("strides", None)
318 skirt = ps.primary_op.attrs.get("skirt", None)
319 npu_block_type = ps.npu_block_type
320
321 ofm_box = Box(ofm_start, ofm_end)
322 ifm_box, _, _ = ofm_box.transform_with_strides_and_skirt(
323 strides, skirt, self.ifm_tensor.shape, npu_block_type, self.concat_axis, self.concat_offset
324 )
325
326 weight_box = None
327 if self.weight_tensor is not None:
328 weight_oc_start = ofm_start[-1]
329 weight_oc_end = ofm_end[-1]
330 if self.concat_axis - len(self.weight_tensor.shape) == -1:
331 weight_oc_start -= self.concat_offset
332 weight_oc_end -= self.concat_offset
333
334 weight_box = Box.make_weight_box(
335 self.weight_tensor.shape,
336 npu_block_type,
337 weight_oc_start,
338 weight_oc_end,
339 self.weight_tensor.weight_transpose_depthwise,
340 )
341
342 return NpuStripe(
343 self.ps,
344 self.block_config,
345 self.is_first,
346 self.is_last,
347 self.is_first_h_stripe,
348 self.is_last_h_stripe,
349 self.ifm_tensor,
350 ifm_box,
351 self.ofm_tensor,
352 ofm_box,
353 self.weight_tensor,
354 weight_box,
355 self.scale_tensor,
356 self.concat_axis,
357 self.concat_offset,
358 )
359
360
361class DMA(Command):
362 def __init__(self, in_tensor, out_tensor, box):
363 self.cmdtype = CommandType.DMA
364 self.in_tensor = in_tensor
365 self.out_tensor = out_tensor
366 self.box = box
367
368 def __str__(self):
369 return "<DMA: in=%s, out=%s, box=%s>" % (self.in_tensor.name, self.out_tensor.name, self.box)
370
371 __repr__ = __str__
372
373 def get_memory_accesses(self):
374 res = MemoryAccessSet()
375
376 res.add(
377 self.in_tensor.get_address_ranges_for_coordinates(self.box.start_coord, self.box.end_coord),
378 AccessDirection.Read,
379 )
380 res.add(
381 self.out_tensor.get_address_ranges_for_coordinates(self.box.start_coord, self.box.end_coord),
382 AccessDirection.Write,
383 )
384 return res
385
386 def get_operation_count(self):
387 # returns numpy array of (DPU blocks, dma_ops)
388 return np.array((0, 1))