Tim Hall | d8339a7 | 2021-05-27 18:49:40 +0100 | [diff] [blame] | 1 | # Copyright (C) 2021 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. |
| 16 | # |
| 17 | # Description: |
| 18 | # Groups Operators in a schedule together to form Cascades. |
Johan Alfvén | fba0a7d | 2022-10-11 20:41:41 +0200 | [diff] [blame] | 19 | from .live_range import ofm_can_reuse_ifm |
Tim Hall | d8339a7 | 2021-05-27 18:49:40 +0100 | [diff] [blame] | 20 | from .numeric_util import round_up |
| 21 | from .operation import NpuBlockType |
erik.andersson@arm.com | 6b2a0b4 | 2022-03-22 15:35:30 +0100 | [diff] [blame] | 22 | from .operation import Op |
Rickard Bolin | 9ae3455 | 2022-06-09 13:07:17 +0000 | [diff] [blame] | 23 | from .operation import Padding |
Tim Hall | d8339a7 | 2021-05-27 18:49:40 +0100 | [diff] [blame] | 24 | from .shape4d import Shape4D |
| 25 | |
| 26 | non_cascadable_blocks = ( |
| 27 | NpuBlockType.Default, |
| 28 | NpuBlockType.VectorProduct, |
Tim Hall | d8339a7 | 2021-05-27 18:49:40 +0100 | [diff] [blame] | 29 | NpuBlockType.ReduceSum, |
| 30 | ) |
| 31 | |
| 32 | |
| 33 | class CascadeInfo: |
| 34 | """Contains metadata about a cascade""" |
| 35 | |
| 36 | def __init__(self, start, end, buffers, mem_usage: int): |
| 37 | self.start = start |
| 38 | self.end = end |
| 39 | self.buffers = buffers |
| 40 | self.mem_usage = mem_usage |
| 41 | |
| 42 | |
| 43 | class BufferMap: |
| 44 | """Caches the buffers seen""" |
| 45 | |
| 46 | def __init__(self): |
| 47 | self.buffer_map = {} |
| 48 | |
| 49 | def get_buffer(self, producer, consumer, cost): |
| 50 | assert producer or consumer |
| 51 | key = (producer, consumer) |
| 52 | if key not in self.buffer_map: |
| 53 | # No cached buffer between these two SchedulerOperations |
| 54 | if consumer is None: |
| 55 | # There are either no consumers or multiple consumers - FeatureMap needs to be stored in full |
| 56 | buffer_shape = producer.ofm.shape |
| 57 | buffer_size = producer.ofm_size_in_bytes() |
| 58 | elif producer is None: |
| 59 | # First Op in subgraph or cascade - FeatureMap needs to be stored in full |
| 60 | buffer_shape = consumer.ifm.shape |
| 61 | buffer_size = consumer.ifm_size_in_bytes() |
| 62 | elif producer.requires_full_ofm or consumer.requires_full_ifm: |
| 63 | # FeatureMap needs to be stored in full |
| 64 | buffer_shape = max(producer.ofm.shape, consumer.ifm.shape) |
| 65 | buffer_size = max(producer.ofm_size_in_bytes(), consumer.ifm_size_in_bytes()) |
| 66 | else: |
| 67 | # Use a rolling buffer |
| 68 | buffer_shape = rolling_buffer_shape(cost[producer].stripe, cost[consumer].stripe_input) |
| 69 | buffer_size = buffer_shape.elements() * producer.ofm.dtype.size_in_bytes() |
| 70 | |
| 71 | self.buffer_map[key] = (buffer_shape, buffer_size) |
| 72 | |
| 73 | return self.buffer_map[key] |
| 74 | |
| 75 | |
| 76 | def rolling_buffer_shape(producer_stripe: Shape4D, consumer_stripe_input: Shape4D) -> Shape4D: |
| 77 | """Calculates the storage shape of the rolling buffer between two SchedulerOperations in a Cascade""" |
| 78 | buffer_height = round_up(producer_stripe.height + consumer_stripe_input.height, consumer_stripe_input.height) |
| 79 | # Rolling buffers have to conform to NHCWB16 format |
| 80 | return consumer_stripe_input.with_height(buffer_height).with_depth(round_up(producer_stripe.depth, 16)) |
| 81 | |
| 82 | |
| 83 | class CascadeBuilder: |
| 84 | """Class for grouping SchedulerOperations into cascades""" |
| 85 | |
| 86 | def __init__(self, sched_ops, spilling, non_local_mem_usage=None): |
| 87 | self.sched_ops = sched_ops |
| 88 | self.no_cascade = 0 |
| 89 | self.non_local_mem_usage = non_local_mem_usage if non_local_mem_usage else {} |
| 90 | self.spilling = spilling |
| 91 | |
| 92 | def _is_cascadable(self, sched_op, cost) -> bool: |
| 93 | """Checks if 'sched_op' can be cascaded""" |
erik.andersson@arm.com | 6b2a0b4 | 2022-03-22 15:35:30 +0100 | [diff] [blame] | 94 | |
Tim Hall | d8339a7 | 2021-05-27 18:49:40 +0100 | [diff] [blame] | 95 | return ( |
| 96 | sched_op.op_type.npu_block_type not in non_cascadable_blocks |
| 97 | and cost.stripe.height < sched_op.ofm.shape.height |
Johan Alfvén | ab677b3 | 2022-05-09 13:02:24 +0200 | [diff] [blame] | 98 | and sched_op.parent_op.read_offsets[0] is None |
| 99 | and sched_op.parent_op.read_offsets[1] is None |
Johan Alfvén | 0f2e59f | 2022-10-21 11:21:38 +0200 | [diff] [blame] | 100 | and self.elementwise_cascadable(sched_op) |
Johan Alfvén | dc7414a | 2022-08-18 11:12:40 +0200 | [diff] [blame] | 101 | and not sched_op.parent_op.type.is_resize_op() |
Fredrik Svedberg | 3e3faa9 | 2022-10-11 16:15:47 +0200 | [diff] [blame] | 102 | and not sched_op.parent_op.type == Op.Conv2DBackpropInputSwitchedBias |
Rickard Bolin | 9ae3455 | 2022-06-09 13:07:17 +0000 | [diff] [blame] | 103 | and sched_op.parent_op.attrs.get("padding", None) != Padding.TILE |
Tim Hall | d8339a7 | 2021-05-27 18:49:40 +0100 | [diff] [blame] | 104 | ) |
| 105 | |
| 106 | def _estimate_sram_usage(self, sched_op, cost) -> int: |
| 107 | """Estimate the SRAM required for the Op if all FeatureMaps are in SRAM""" |
| 108 | ifm2_size = sched_op.ifm2_size_in_bytes() |
| 109 | if sched_op.requires_full_ifm: |
| 110 | ifm_size = sched_op.ifm_size_in_bytes() |
| 111 | else: |
| 112 | ifm_size = ( |
| 113 | cost.stripe_input.with_depth(round_up(cost.stripe_input.depth, 16)).elements() |
| 114 | * sched_op.ifm.dtype.size_in_bytes() |
| 115 | ) |
Johan Alfvén | fba0a7d | 2022-10-11 20:41:41 +0200 | [diff] [blame] | 116 | if ofm_can_reuse_ifm(sched_op): |
| 117 | # ofm will use the ifm buffer to reduce SRAM usage, hence ofm_size = 0 |
| 118 | ofm_size = 0 |
| 119 | elif sched_op.requires_full_ofm: |
Tim Hall | d8339a7 | 2021-05-27 18:49:40 +0100 | [diff] [blame] | 120 | ofm_size = sched_op.ofm_size_in_bytes() |
| 121 | else: |
| 122 | ofm_size = ( |
| 123 | cost.stripe.with_depth(round_up(cost.stripe.depth, 16)).elements() * sched_op.ofm.dtype.size_in_bytes() |
| 124 | ) |
| 125 | |
| 126 | return ifm_size + ifm2_size + ofm_size + self.non_local_mem_usage.get(sched_op, 0) |
| 127 | |
erik.andersson@arm.com | 6b2a0b4 | 2022-03-22 15:35:30 +0100 | [diff] [blame] | 128 | @staticmethod |
Johan Alfvén | 0f2e59f | 2022-10-21 11:21:38 +0200 | [diff] [blame] | 129 | def elementwise_cascadable(sched_op): |
| 130 | """Check if the elementwise can be cascaded.""" |
erik.andersson@arm.com | 6b2a0b4 | 2022-03-22 15:35:30 +0100 | [diff] [blame] | 131 | |
Johan Alfvén | 56a71b0 | 2022-10-19 11:20:12 +0200 | [diff] [blame] | 132 | if sched_op.parent_op.type.is_binary_elementwise_op(): |
Johan Alfvén | 56a71b0 | 2022-10-19 11:20:12 +0200 | [diff] [blame] | 133 | ifm = sched_op.parent_op.ifm |
| 134 | ifm2 = sched_op.parent_op.ifm2 |
Johan Alfvén | 0f2e59f | 2022-10-21 11:21:38 +0200 | [diff] [blame] | 135 | ofm = sched_op.parent_op.ofm |
erik.andersson@arm.com | 6b2a0b4 | 2022-03-22 15:35:30 +0100 | [diff] [blame] | 136 | |
Johan Alfvén | 0f2e59f | 2022-10-21 11:21:38 +0200 | [diff] [blame] | 137 | # IFM must be non-constant/non-scalar/non-broadcast |
| 138 | ifm_cascadable = not (ifm.is_const or ifm.is_scalar or ifm.is_broadcast(ofm)) |
Johan Alfvén | 56a71b0 | 2022-10-19 11:20:12 +0200 | [diff] [blame] | 139 | |
Johan Alfvén | 0f2e59f | 2022-10-21 11:21:38 +0200 | [diff] [blame] | 140 | # IFM2 must be constant or scalar |
| 141 | ifm2_cascadable = ifm2.is_const or ifm2.is_scalar |
Johan Alfvén | 56a71b0 | 2022-10-19 11:20:12 +0200 | [diff] [blame] | 142 | |
Johan Alfvén | 0f2e59f | 2022-10-21 11:21:38 +0200 | [diff] [blame] | 143 | return ifm_cascadable and ifm2_cascadable |
Johan Alfvén | 56a71b0 | 2022-10-19 11:20:12 +0200 | [diff] [blame] | 144 | else: |
| 145 | return True |
| 146 | |
Tim Hall | d8339a7 | 2021-05-27 18:49:40 +0100 | [diff] [blame] | 147 | def build_cascades(self, ref_schedule, fallback_schedule, guiding_mem_limit): |
| 148 | ref_cost = ref_schedule.cost_map |
| 149 | fallback_cost = fallback_schedule.cost_map |
| 150 | cost = {} |
| 151 | cascade_map = {} |
| 152 | buffers = BufferMap() |
| 153 | |
| 154 | # Peak memory usage so far - updated continously, unless dedicated SRAM where this is a hard limit |
| 155 | peak_sram_usage = guiding_mem_limit |
| 156 | |
| 157 | idx = 0 |
| 158 | while idx < len(self.sched_ops): |
| 159 | op = self.sched_ops[idx] |
| 160 | if op in cost: |
| 161 | # Already processed this Op |
| 162 | idx += 1 |
| 163 | continue |
| 164 | |
| 165 | if not self._is_cascadable(op, ref_cost[op]): |
| 166 | # Op is not a candidate for cascading - assign fallback cost |
| 167 | cost[op] = fallback_cost[op] |
| 168 | if not self.spilling: |
| 169 | peak_sram_usage = max(self._estimate_sram_usage(op, fallback_cost[op]), peak_sram_usage) |
| 170 | idx += 1 |
| 171 | continue |
| 172 | |
| 173 | # Propose a cascade starting with this Op |
| 174 | cascade_start = op.index |
| 175 | # Keep track of which Ops are in the proposed cascade as well as the best cascade so far |
| 176 | ops_in_cascade = [op] |
| 177 | ops_in_best_cascade = [op] |
Rickard Bolin | fd8b500 | 2022-05-16 09:11:06 +0000 | [diff] [blame] | 178 | # Get the size of the weight buffer(s) |
| 179 | weight_buffer = sum(tens.storage_size() for tens in ref_cost[op].buffered_weight_tensors) |
Tim Hall | d8339a7 | 2021-05-27 18:49:40 +0100 | [diff] [blame] | 180 | |
| 181 | # The first IFM needs to be stored in full |
| 182 | cascade_ifm_size = op.ifm_size_in_bytes() if not self.spilling else 0 |
| 183 | |
| 184 | # Add non-local memory usage |
| 185 | cascade_ifm_size += self.non_local_mem_usage.get(op, 0) |
| 186 | |
| 187 | # Sum of all intermediate cascade buffers (including weight buffers) |
| 188 | cascade_buffers = weight_buffer |
| 189 | # Best cascade size - Initially it's the fallback cost of the first Op in the cascade |
| 190 | best_cascade_size = self._estimate_sram_usage(op, fallback_cost[op]) |
| 191 | |
| 192 | # Op is the producer of the OFM consumed by the next Op to consider |
| 193 | producer = op |
| 194 | while True: |
| 195 | dependants = producer.get_dependants() |
| 196 | if len(dependants) != 1: |
| 197 | # producer is either the last Op in the schedule or the start of a branch |
| 198 | break |
| 199 | |
| 200 | current_op = dependants[0] |
| 201 | if ( |
| 202 | current_op in cost |
| 203 | or current_op not in ref_cost |
| 204 | or not self._is_cascadable(current_op, ref_cost[current_op]) |
| 205 | or producer.ofm.shape != current_op.ifm.shape |
Louis Verhaard | 04bd3e9 | 2021-08-19 16:36:32 +0200 | [diff] [blame] | 206 | or current_op.requires_full_ifm |
| 207 | or producer.requires_full_ofm |
Tim Hall | d8339a7 | 2021-05-27 18:49:40 +0100 | [diff] [blame] | 208 | ): |
| 209 | # Current op has already been processed or cannot be cascaded |
| 210 | break |
| 211 | |
Louis Verhaard | 37ba98c | 2022-03-16 09:56:45 +0100 | [diff] [blame] | 212 | if producer.index + 1 != current_op.index: |
| 213 | # Cascading is possible, but requires reordering of operations in the schedule, |
| 214 | # this is currently not supported |
| 215 | break |
| 216 | |
Tim Hall | d8339a7 | 2021-05-27 18:49:40 +0100 | [diff] [blame] | 217 | # Get the size of the FeatureMap buffers between current and neighbouring Ops |
| 218 | op_full_ifm = current_op.ifm_size_in_bytes() |
| 219 | op_full_ofm = current_op.ofm_size_in_bytes() |
| 220 | _, op_ifm_buffer = buffers.get_buffer(producer, current_op, ref_cost) |
| 221 | |
Rickard Bolin | fd8b500 | 2022-05-16 09:11:06 +0000 | [diff] [blame] | 222 | # Get the size of the weight buffer(s) |
| 223 | op_weight_buffer = sum(tens.storage_size() for tens in ref_cost[current_op].buffered_weight_tensors) |
Tim Hall | d8339a7 | 2021-05-27 18:49:40 +0100 | [diff] [blame] | 224 | |
| 225 | # Calculate the uncascaded memory requirement for current Op |
| 226 | uncascaded_sram_usage = op_full_ifm + op_full_ofm + self.non_local_mem_usage.get(current_op, 0) |
| 227 | |
| 228 | # Add current Op to cascade |
| 229 | ops_in_cascade.append(current_op) |
| 230 | |
| 231 | # Increase the accumulated intermediate buffers in the cascade |
| 232 | cascade_buffers += op_ifm_buffer + op_weight_buffer |
| 233 | |
| 234 | if self.spilling: |
| 235 | # For Dedicated SRAM only the intermediate buffers are in SRAM |
| 236 | if uncascaded_sram_usage < peak_sram_usage or cascade_buffers > peak_sram_usage: |
| 237 | # Cascade until an Op fits in its entirety or the accumulated buffers no longer fit |
| 238 | break |
| 239 | else: |
| 240 | # Any addition to the cascade that fits is the new best cascade for Dedicated SRAM |
| 241 | ops_in_best_cascade = [op for op in ops_in_cascade] |
| 242 | best_cascade_size = cascade_buffers |
| 243 | |
| 244 | else: |
| 245 | # Calculate the total size of the current cascade |
| 246 | cascade_size = cascade_ifm_size + cascade_buffers + op_full_ofm |
| 247 | |
| 248 | # Determine if cascading search should stop |
| 249 | if ( |
| 250 | uncascaded_sram_usage < peak_sram_usage |
| 251 | and best_cascade_size < peak_sram_usage |
| 252 | or (cascade_ifm_size + cascade_buffers) > best_cascade_size |
| 253 | ): |
| 254 | # Both the existing cascade and current Op fits |
| 255 | break |
| 256 | |
Johan Alfvén | 255dad7 | 2022-07-16 18:27:05 +0200 | [diff] [blame] | 257 | """ |
| 258 | One of two conditions will update the best cascade: |
| 259 | |
| 260 | - cascade_size < best_cascade_size or |
| 261 | - cascade_size < uncascaded_sram_usage |
| 262 | |
| 263 | The last condition is illustrated below, showing an example where it is |
| 264 | better to choose a larger cascade_size (with more OPs) because it will |
| 265 | use less total SRAM usage. |
| 266 | |
| 267 | For simplicity, all featuremaps have same size. |
| 268 | |
| 269 | Cascade OP1-OP2, OP3 is standalone |
| 270 | |
| 271 | -> |OP1| -> roll buffer -> |OP2| -> FM -> |OP3| -> FM |
| 272 | / |
| 273 | |OP0| -> FM |
| 274 | \ |
| 275 | -> .... |
| 276 | |
| 277 | |
| 278 | best_cascade_size : FM + roll buffer + FM |
| 279 | uncascaded_sram_usage: FM + FM + FM |
| 280 | |
| 281 | compared with: |
| 282 | |
| 283 | Cascade OP1-OP3 |
| 284 | |
| 285 | -> |OP1| -> roll buffer -> |OP2| -> roll buffer -> |OP3| -> FM |
| 286 | / |
| 287 | |OP0| -> FM |
| 288 | \ |
| 289 | -> .... |
| 290 | |
| 291 | |
| 292 | cascade_size : FM + roll buffer + roll buffer + FM |
| 293 | |
| 294 | |
| 295 | So, for this use case the comparison will be |
| 296 | |
| 297 | (FM + roll buffer + roll buffer + FM) < (FM + roll buffer + FM) or |
| 298 | (FM + roll buffer + roll buffer + FM) < (FM + FM + FM) |
| 299 | |
| 300 | hence, better to choose Cascade OP1-OP3 in this case. |
| 301 | """ |
| 302 | if cascade_size < best_cascade_size or cascade_size < uncascaded_sram_usage: |
Tim Hall | d8339a7 | 2021-05-27 18:49:40 +0100 | [diff] [blame] | 303 | best_cascade_size = cascade_ifm_size + cascade_buffers + op_full_ofm |
| 304 | ops_in_best_cascade = [op for op in ops_in_cascade] |
| 305 | |
| 306 | producer = current_op |
| 307 | |
| 308 | if len(ops_in_best_cascade) > 1: |
| 309 | # A cascade was created - assign cascade and ref_cost to all of the Ops |
| 310 | cascade_end = cascade_start + (len(ops_in_best_cascade) - 1) |
| 311 | buffers_in_cascade = {} |
| 312 | prev_op = None |
| 313 | for cascaded_op in ops_in_best_cascade: |
Louis Verhaard | 37ba98c | 2022-03-16 09:56:45 +0100 | [diff] [blame] | 314 | assert cascade_start <= cascaded_op.index <= cascade_end |
Tim Hall | d8339a7 | 2021-05-27 18:49:40 +0100 | [diff] [blame] | 315 | cost[cascaded_op] = ref_cost[cascaded_op] |
| 316 | cost[cascaded_op].cascade = cascade_end |
| 317 | if prev_op: |
| 318 | rolling_buffer_shape, _ = buffers.get_buffer(prev_op, cascaded_op, ref_cost) |
| 319 | buffers_in_cascade[cascaded_op] = rolling_buffer_shape |
| 320 | |
| 321 | prev_op = cascaded_op |
| 322 | |
| 323 | # Create a CascadeInfo for the cascade |
| 324 | cascade_map[cascade_end] = CascadeInfo( |
| 325 | cascade_start, cascade_end, buffers_in_cascade, best_cascade_size |
| 326 | ) |
| 327 | if not self.spilling: |
| 328 | # Update peak memory usage |
| 329 | peak_sram_usage = max(best_cascade_size, peak_sram_usage) |
| 330 | else: |
| 331 | # Assign fallback cost to the initial Op |
| 332 | cost[op] = fallback_cost[op] |
| 333 | if not self.spilling: |
| 334 | peak_sram_usage = max(self._estimate_sram_usage(op, fallback_cost[op]), peak_sram_usage) |
| 335 | |
erik.andersson@arm.com | 6b2a0b4 | 2022-03-22 15:35:30 +0100 | [diff] [blame] | 336 | # Update costing and cascade information for the ref_schedule |
Tim Hall | d8339a7 | 2021-05-27 18:49:40 +0100 | [diff] [blame] | 337 | ref_schedule.cost_map = cost |
| 338 | ref_schedule.cascades = cascade_map |
| 339 | return ref_schedule |