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