| # Copyright (C) 2020 Arm Limited or its affiliates. All rights reserved. |
| # |
| # SPDX-License-Identifier: Apache-2.0 |
| # |
| # Licensed under the Apache License, Version 2.0 (the License); you may |
| # not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an AS IS BASIS, WITHOUT |
| # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| # Description: |
| # Shared buffer allocation works out how to allocate the Ethos-U55 shared buffer for a given pass. |
| from typing import List |
| from typing import Tuple |
| |
| import numpy as np |
| |
| from .api import NpuActivationOp |
| from .api import NpuBlockOperation |
| from .architecture_features import ArchitectureFeatures |
| from .architecture_features import Block |
| from .architecture_features import SharedBufferArea |
| from .architecture_features import SHRAMElements |
| from .errors import VelaError |
| from .ethos_u55_regs.ethos_u55_regs import resampling_mode |
| from .high_level_command_to_npu_op import to_kernel |
| from .operation import Kernel |
| from .operation import NpuBlockType |
| from .range_set import MemoryRangeSet |
| from .tensor import MemArea |
| |
| |
| class SharedBufferAllocation: |
| def __init__( |
| self, |
| arch, |
| kernel, |
| uses_lut, |
| npu_block_type, |
| all_fms_have_quant, |
| ifm_resampling_mode, |
| ifm_bits, |
| ifm_depth, |
| ifm_count, |
| ofm_shape, |
| ): |
| self.arch = arch |
| |
| self.bank_locations = np.zeros(SharedBufferArea.Size) |
| self.banks_required = np.zeros(SharedBufferArea.Size) |
| |
| self.kernel = Kernel(1, 1) if kernel is None else kernel |
| self.is_elementwise = npu_block_type == NpuBlockType.ElementWise |
| self.uses_lut = uses_lut |
| self.ifm_count = ifm_count |
| |
| self.is_equal_depth_op = self.is_elementwise or npu_block_type in ( |
| NpuBlockType.ConvolutionDepthWise, |
| NpuBlockType.Pooling, |
| ) |
| |
| self.use_accumulator_element = SHRAMElements.Acc32 |
| if self.is_elementwise: |
| self.use_ifm_element = SHRAMElements.IFM8_Elementwise |
| else: |
| self.use_ifm_element = SHRAMElements.IFM8 |
| |
| self.ifm_resampling_mode = ifm_resampling_mode |
| self.ifm_bits = ifm_bits |
| self.ifm_depth = ifm_depth |
| self.ifm_count = ifm_count |
| |
| if self.ifm_bits == 16: |
| if npu_block_type != NpuBlockType.Pooling and all_fms_have_quant: |
| self.use_accumulator_element = SHRAMElements.Acc40 |
| self.use_ifm_element = self.use_ifm_element + 1 |
| assert (self.use_ifm_element == SHRAMElements.IFM16) or ( |
| self.use_ifm_element == SHRAMElements.IFM16_Elementwise |
| ) |
| elif self.ifm_bits == 32: |
| assert self.is_elementwise or npu_block_type == NpuBlockType.ReduceSum, "Unsupported 32-bit IFM operation" |
| self.use_ifm_element = SHRAMElements.IFM32 |
| else: |
| assert self.ifm_bits == 8, "Unexpected IFM bitdepth" |
| |
| self.ifm_block_depth = arch.calc_ifm_block_depth(self.ifm_depth, self.ifm_bits) |
| self.ofm_shape = ofm_shape |
| |
| self.banks_required[SharedBufferArea.Weights] = arch.shram_reserved_weight_banks |
| self.banks_required[SharedBufferArea.OFM] = arch.shram_reserved_output_banks |
| |
| def is_valid(self): |
| # Assign zero-based bank starts (first element remains zero) |
| self.bank_locations[1:] = np.cumsum(self.banks_required)[:-1] |
| |
| # Accumulator area is measured from the end of the buffer |
| self.bank_locations[SharedBufferArea.Accumulators] = ( |
| self.arch.available_shram_banks(self.uses_lut) - self.banks_required[SharedBufferArea.Accumulators] |
| ) |
| ifm_end = self.bank_locations[SharedBufferArea.IFM] + self.banks_required[SharedBufferArea.IFM] |
| return ifm_end <= self.bank_locations[SharedBufferArea.Accumulators] |
| |
| def try_block(self, ofm_block: Block): |
| # Get IFM block configuration |
| ifm_block_depth = ofm_block.depth if self.is_equal_depth_op else self.ifm_block_depth |
| ifm_block = self.arch.get_ifm_block_size( |
| ifm_block_depth, ofm_block, self.kernel, ifm_resampling_mode=self.ifm_resampling_mode |
| ) |
| ifm_config = self.arch.get_block_config(ifm_block.width, ifm_block.height, ifm_block.depth) |
| if ifm_config is None: |
| return None |
| |
| # Get OFM block configuration |
| ofm_config = self.arch.get_block_config(ofm_block.width, ofm_block.height, ofm_block.depth) |
| if ofm_config is None: |
| return None |
| |
| acc_banks = ofm_config.banks[self.use_accumulator_element] |
| |
| # Update bank counts for IFM and Accumulator |
| self.banks_required[SharedBufferArea.IFM] = ifm_config.banks[self.use_ifm_element] * self.ifm_count |
| self.banks_required[SharedBufferArea.Accumulators] = 0 if self.is_elementwise else acc_banks |
| |
| # Validating calculates bank layout and returns validity |
| if not self.is_valid(): |
| return None |
| |
| return (ofm_block.height, ofm_block.width, ifm_block.depth, ofm_block.depth) |
| |
| def generate_used_mask(self, active_set): |
| res = np.zeros(self.arch.shram_total_banks, dtype=np.int64) |
| for kind in active_set: |
| start = int(self.bank_locations[kind]) |
| end = start + int(self.banks_required[kind]) |
| res[start:end] = 1 |
| return res |
| |
| def is_compatible(first, second): |
| """See if the bank allocations of two convolutions are compatible, |
| so that they can run back-to-back without a fence in between""" |
| |
| first_set = set((SharedBufferArea.OFM, SharedBufferArea.Accumulators)) |
| second_set = set((SharedBufferArea.IFM, SharedBufferArea.Weights)) |
| |
| first_mask = first.generate_used_mask(first_set) |
| second_mask = second.generate_used_mask(second_set) |
| |
| if np.sum(first_mask & second_mask): |
| # overlap |
| return False |
| |
| return True |
| |
| def get_shram_memory_access_range(self): |
| # Returns the SHRAM memory access range used by this shared buffer, |
| # excluding access to LUT |
| return MemoryRangeSet( |
| MemArea.Shram, 0, self.arch.available_shram_banks(self.uses_lut) * self.arch.shram_bank_size |
| ) |
| |
| |
| def _all_fms_have_quant(ifm_tensor, ofm_tensor, ifm2_tensor=None) -> bool: |
| tensors = [t for t in (ifm_tensor, ifm2_tensor, ofm_tensor) if t is not None] |
| scales = [t.quantization.scale_f32 for t in tensors if t.quantization is not None] |
| return len(tensors) == len(scales) and None not in scales |
| |
| |
| def is_acc_40bits_used(npu_block_type, ifm_tensor, ofm_tensor, ifm2_tensor=None): |
| return npu_block_type != NpuBlockType.Pooling and _all_fms_have_quant(ifm_tensor, ofm_tensor, ifm2_tensor) |
| |
| |
| def shared_buffer_allocation_for_pass(arch, ps) -> SharedBufferAllocation: |
| ifm_tensor, ifm2_tensor, _, ofm_tensor = ps.get_primary_op_ifm_ifm2_weights_ofm() |
| all_fms_have_quant = _all_fms_have_quant(ifm_tensor, ifm2_tensor, ofm_tensor) |
| |
| kernel = Kernel(1, 1) |
| is_elementwise = ps.npu_block_type == NpuBlockType.ElementWise |
| uses_lut = False |
| ifm_count = 1 |
| |
| if ps.primary_op: |
| kernel = ps.primary_op.kernel |
| uses_lut = ps.primary_op.activation_lut is not None |
| |
| ifm_resampling_mode = resampling_mode.NONE |
| ifm_bits = 0 |
| ifm_depth = 0 |
| if ifm_tensor: |
| ifm_resampling_mode = ifm_tensor.resampling_mode |
| ifm_bits = ifm_tensor.dtype.size_in_bits() |
| |
| if ifm_tensor.shape != []: |
| ifm_depth = ifm_tensor.shape[-1] |
| |
| if is_elementwise: |
| ifm_count = 2 |
| if ifm_tensor.shape == []: # Scalar in ifm1 |
| assert ifm2_tensor |
| ifm_depth = ifm2_tensor.shape[-1] |
| ifm_count = 1 |
| elif not ifm2_tensor or ifm2_tensor.shape == []: # Scalar in ifm2 |
| ifm_count = 1 |
| return SharedBufferAllocation( |
| arch, |
| kernel, |
| uses_lut, |
| npu_block_type=ps.npu_block_type, |
| all_fms_have_quant=all_fms_have_quant, |
| ifm_resampling_mode=ifm_resampling_mode, |
| ifm_bits=ifm_bits, |
| ifm_depth=ifm_depth, |
| ifm_count=ifm_count, |
| ofm_shape=ofm_tensor.shape, |
| ) |
| |
| |
| def shared_buffer_allocation_for_pass_and_block_config(arch, ps, block_config) -> SharedBufferAllocation: |
| alloc = shared_buffer_allocation_for_pass(arch, ps) |
| assert (alloc.ifm_block_depth == block_config[2]) or alloc.is_equal_depth_op |
| if alloc.try_block(Block(block_config[1], block_config[0], block_config[3])): |
| return alloc |
| |
| return None |
| |
| |
| def shared_buffer_allocation_for_npu_op( |
| arch, npu_op: NpuBlockOperation, npu_block_type: NpuBlockType, ifm_resampling_mode |
| ) -> SharedBufferAllocation: |
| uses_lut = npu_op.activation is not None and npu_op.activation.op_type == NpuActivationOp.TABLE_LOOKUP |
| fms = [npu_op.ifm, npu_op.ofm] |
| if npu_op.ifm2 is not None: |
| fms.append(npu_op.ifm2) |
| all_fms_have_quant = not any(fm.quantization is None or fm.quantization.scale_f32 is None for fm in fms) |
| ifm_bits = npu_op.ifm.data_type.size_in_bits() |
| ifm_depth = npu_op.ifm.shape.depth |
| ifm_count = 2 if npu_op.ifm2 is not None and npu_op.ifm2_scalar is None else 1 |
| ofm_shape = [1, npu_op.ofm.shape.height, npu_op.ofm.shape.width, npu_op.ofm.shape.depth] |
| return SharedBufferAllocation( |
| arch, |
| to_kernel(npu_op.kernel), |
| uses_lut, |
| npu_block_type=npu_block_type, |
| all_fms_have_quant=all_fms_have_quant, |
| ifm_resampling_mode=ifm_resampling_mode, |
| ifm_bits=ifm_bits, |
| ifm_depth=ifm_depth, |
| ifm_count=ifm_count, |
| ofm_shape=ofm_shape, |
| ) |
| |
| |
| def find_suitable_block_configs(arch, alloc: SharedBufferAllocation) -> List[Tuple]: |
| """Returns list of block configs that would fit with the given shared buffer allocation""" |
| if arch.override_block_config: |
| config = alloc.try_block(arch.override_block_config) |
| if config is None: |
| raise VelaError("Block config override '{0}' cannot be allocated".format(arch.override_block_config)) |
| return [config] |
| |
| # Constrain the search space if the OFM is smaller than the max block size |
| # - Add other block search constraints here if required |
| if len(alloc.ofm_shape) <= 2: |
| max_block_height = max_block_width = alloc.ofm_shape[0] |
| else: |
| max_block_width = alloc.ofm_shape[-2] |
| max_block_height = alloc.ofm_shape[-3] |
| |
| # Common block depth |
| max_block_depth = alloc.ofm_shape[-1] |
| |
| # Constrain to valid ranges before search |
| max_block_width = min(arch.ofm_block_max.width, max_block_width) |
| max_block_height = min(arch.ofm_block_max.height, max_block_height) |
| max_block_depth = min(arch.ofm_block_max.depth, max_block_depth) |
| |
| valid_block_configs = [] |
| # Try a range of block shapes against this pass |
| for w in range(arch.ofm_ublock.width, max_block_width + arch.ofm_ublock.width, arch.ofm_ublock.width): |
| for h in range(arch.ofm_ublock.height, max_block_height + arch.ofm_ublock.height, arch.ofm_ublock.height): |
| # Try valid OFM block depths |
| for c in range(arch.ofm_ublock.depth, max_block_depth + arch.ofm_ublock.depth, arch.ofm_ublock.depth): |
| # OFM block depth has the constraint that if it causes the OFM to be |
| # split, it must be a multiple of the OFM split size |
| if (c >= max_block_depth) or (c < max_block_depth and (c % ArchitectureFeatures.OFMSplitDepth) == 0): |
| config = alloc.try_block(Block(w, h, c)) |
| if config: |
| valid_block_configs.append(config) |
| |
| assert len(valid_block_configs) > 0 |
| return valid_block_configs |
| |
| |
| def find_block_configs_suitable_for_pass_and_shared_buffer(arch, ps) -> List[Tuple]: |
| alloc = shared_buffer_allocation_for_pass(arch, ps) |
| return find_suitable_block_configs(arch, alloc) |