blob: 39a7f21f2d8c67d54eeb23a26d3bcc15bcf4687e [file] [log] [blame]
erik.andersson@arm.com460c6892021-02-24 14:38:09 +01001# Copyright (C) 2020-2021 Arm Limited or its affiliates. All rights reserved.
Tim Hall79d07d22020-04-27 18:20:16 +01002#
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# Serialises and packs an NPU subgraph into tensors.
Diego Russoea6111a2020-04-14 18:41:58 +010018import numpy as np
19
20from . import driver_actions
Tim Hall79d07d22020-04-27 18:20:16 +010021from .data_type import DataType
Diego Russoe8a10452020-04-21 17:39:10 +010022from .nn_graph import PassPlacement
Louis Verhaardaee5d752020-09-30 09:01:52 +020023from .operation import Op
Diego Russoe8a10452020-04-21 17:39:10 +010024from .operation import Operation
25from .tensor import MemArea
Patrik Gustavssoneca2e952020-05-27 09:15:11 +020026from .tensor import MemType
Diego Russoe8a10452020-04-21 17:39:10 +010027from .tensor import Tensor
28from .tensor import TensorFormat
29from .tensor import TensorPurpose
Tim Hall79d07d22020-04-27 18:20:16 +010030
31
Patrik Gustavssoneca2e952020-05-27 09:15:11 +020032def make_memory_tensor(name, mem_area, mem_type, sz, want_values, arch):
Tim Hall79d07d22020-04-27 18:20:16 +010033 tens = Tensor([sz], DataType.uint8, name)
34 tens.mem_area = mem_area
Patrik Gustavssoneca2e952020-05-27 09:15:11 +020035 tens.mem_type = mem_type
Tim Hall79d07d22020-04-27 18:20:16 +010036 tens.purpose = TensorPurpose.FeatureMap
37 tens.set_format(TensorFormat.NHWC, arch)
38 if want_values:
39 tens.values = np.zeros(tens.shape, np.uint8)
40 return tens
41
42
43def copy_compressed_values_to_memory_tensor(memory_tensor, src_tensor):
44 start_addr = src_tensor.address
Tim Halld8339a72021-05-27 18:49:40 +010045 end_addr = src_tensor.address + src_tensor.storage_size()
46 memory_tensor.values[start_addr:end_addr] = src_tensor.buffer.copy()
Tim Hall79d07d22020-04-27 18:20:16 +010047
Tim Hallc30f4952020-06-15 20:47:35 +010048
Charles Xu78792222020-05-13 10:15:26 +020049def copy_ifm_values_to_memory_tensor(memory_tensor, src_tensor):
50 start_addr = src_tensor.address
Fredrik Svedberg0f98b362020-09-29 10:00:39 +020051 values = src_tensor.quant_values.flatten() if src_tensor.quant_values is not None else src_tensor.values.flatten()
Fredrik Svedbergbb1a92a2020-08-27 15:51:50 +020052 if src_tensor.dtype.size_in_bytes() > 1:
53 values = np.frombuffer(values.tobytes(), dtype=np.uint8)
Charles Xu9a03fdf2020-07-02 15:12:40 +020054 end_addr = start_addr + values.size
55 memory_tensor.values[start_addr:end_addr] = values
Tim Hall79d07d22020-04-27 18:20:16 +010056
Tim Hallc30f4952020-06-15 20:47:35 +010057
Tim Hall03d40a22021-04-22 12:08:28 +010058def serialise_npu_subgraph_into_tensors(sg, arch, scratch_tens, scratch_fast_tens, flash_tens):
Tim Hall79d07d22020-04-27 18:20:16 +010059 if sg.placement != PassPlacement.Npu:
Patrik Gustavsson3ab94522020-06-29 17:36:55 +020060 return scratch_tens, scratch_fast_tens, flash_tens
Tim Hall79d07d22020-04-27 18:20:16 +010061
62 flash_area = arch.permanent_storage_mem_area
Patrik Gustavssoneca2e952020-05-27 09:15:11 +020063 scratch_area = arch.feature_map_storage_mem_area
Patrik Gustavsson3ab94522020-06-29 17:36:55 +020064 scratch_fast_area = arch.fast_storage_mem_area
Tim Hall79d07d22020-04-27 18:20:16 +010065
66 flash_size = sg.memory_used.get(flash_area, 0)
67 scratch_size = sg.memory_used.get(scratch_area, 0)
68
Louis Verhaard52078302020-11-18 13:35:06 +010069 payload_bytes = driver_actions.create_driver_payload(sg.register_command_stream, arch)
Tim Hall79d07d22020-04-27 18:20:16 +010070
71 command_stream_size_bytes = len(payload_bytes)
72
Diego Russoea6111a2020-04-14 18:41:58 +010073 if flash_tens == scratch_tens is None:
Tim Hall79d07d22020-04-27 18:20:16 +010074 # First Npu subgraph, create scratch and flash tensors
Patrik Gustavssoneca2e952020-05-27 09:15:11 +020075 sg.scratch_tensor = make_memory_tensor(
76 sg.name + "_scratch", scratch_area, MemType.Scratch, scratch_size, False, arch
77 )
Tim Hall79d07d22020-04-27 18:20:16 +010078 sg.scratch_tensor.purpose = TensorPurpose.Scratch
Patrik Gustavssoneca2e952020-05-27 09:15:11 +020079 sg.flash_tensor = make_memory_tensor(
80 sg.name + "_flash", flash_area, MemType.Permanent_CPU, flash_size, True, arch
81 )
Patrik Gustavsson3ab94522020-06-29 17:36:55 +020082 sg.scratch_fast_tensor = make_memory_tensor(
Jacob Bohlin268394d2020-08-13 13:24:59 +020083 sg.name + "_scratch_fast", scratch_fast_area, MemType.Scratch_fast, 0, False, arch
Patrik Gustavsson3ab94522020-06-29 17:36:55 +020084 )
85 sg.scratch_fast_tensor.purpose = TensorPurpose.Scratch
Tim Hall79d07d22020-04-27 18:20:16 +010086 else:
87 sg.scratch_tensor = scratch_tens
88 sg.scratch_tensor.shape[0] += scratch_size
89 sg.flash_tensor = flash_tens
90 sg.flash_tensor.shape[0] += flash_size
91
Patrik Gustavsson3ab94522020-06-29 17:36:55 +020092 sg.scratch_fast_tensor = scratch_fast_tens
93 sg.scratch_fast_tensor.shape[0] = 0
94
Tim Halld8339a72021-05-27 18:49:40 +010095 for sched_op in sg.sched_ops:
96 ifm_tensor, ifm2_tensor, _, _, _ = sched_op.parent_op.get_ifm_ifm2_weights_biases_ofm()
Tim Hall79d07d22020-04-27 18:20:16 +010097
Tim Halld8339a72021-05-27 18:49:40 +010098 op_info = sg.schedule.cost_map[sched_op]
99 if op_info.npu_weights_tensor:
100 copy_compressed_values_to_memory_tensor(sg.flash_tensor, op_info.npu_weights_tensor)
Charles Xu78792222020-05-13 10:15:26 +0200101
Tim Halld8339a72021-05-27 18:49:40 +0100102 if ifm_tensor and ifm_tensor.mem_type not in (MemType.Scratch, MemType.Scratch_fast):
103 copy_ifm_values_to_memory_tensor(sg.flash_tensor, ifm_tensor)
104 if ifm2_tensor and (ifm2_tensor.mem_type not in (MemType.Scratch, MemType.Scratch_fast)):
105 copy_ifm_values_to_memory_tensor(sg.flash_tensor, ifm2_tensor)
106
107 if sched_op.parent_op.activation_lut:
108 copy_ifm_values_to_memory_tensor(sg.flash_tensor, sched_op.parent_ps.lut_tensor)
109
Tim Hall79d07d22020-04-27 18:20:16 +0100110 sg.command_stream_tensor = make_memory_tensor(
Patrik Gustavssoneca2e952020-05-27 09:15:11 +0200111 sg.name + "_command_stream", flash_area, MemType.Permanent_CPU, command_stream_size_bytes, True, arch
Tim Hall79d07d22020-04-27 18:20:16 +0100112 )
113 sg.command_stream_tensor.values = np.frombuffer(payload_bytes, dtype=np.uint8)
114
Patrik Gustavsson3ab94522020-06-29 17:36:55 +0200115 return sg.scratch_tensor, sg.scratch_fast_tensor, sg.flash_tensor
Tim Hall79d07d22020-04-27 18:20:16 +0100116
117
118def add_const_tens_to_startup_cascaded_pass(startup_cps, tens):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200119 op = Operation(Op.Const, tens.name + "_const")
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100120 op.set_output_tensor(tens)
Tim Hall79d07d22020-04-27 18:20:16 +0100121 startup_cps.passes[0].ops.insert(0, op)
122 startup_cps.passes[0].outputs.insert(0, tens)
123 startup_cps.outputs.insert(0, tens)
124
125
Tim Hall03d40a22021-04-22 12:08:28 +0100126def rewrite_npu_call_ops(sg, arch):
Tim Hall79d07d22020-04-27 18:20:16 +0100127 if sg.placement != PassPlacement.Cpu:
128 return
129
130 startup_cps = sg.cascaded_passes[0]
131
132 for idx, cps in enumerate(sg.cascaded_passes):
133 for ps in cps.passes:
134 for op in ps.ops:
Louis Verhaardaee5d752020-09-30 09:01:52 +0200135 if op.type == Op.CustomNpuOp:
Tim Hall79d07d22020-04-27 18:20:16 +0100136 callee = op.attrs["subgraph"]
Tim Hall79d07d22020-04-27 18:20:16 +0100137
138 sz = 0
Patrik Gustavsson3ab94522020-06-29 17:36:55 +0200139 for tens in [
140 callee.scratch_fast_tensor,
141 callee.scratch_tensor,
142 callee.flash_tensor,
143 callee.command_stream_tensor,
144 ]:
Tim Hall79d07d22020-04-27 18:20:16 +0100145 op.inputs.insert(0, tens)
146 ps.inputs.insert(0, tens)
147 cps.inputs.insert(0, tens)
Patrik Gustavsson3ab94522020-06-29 17:36:55 +0200148 if tens != callee.scratch_tensor and tens != callee.scratch_fast_tensor:
Tim Hall79d07d22020-04-27 18:20:16 +0100149 add_const_tens_to_startup_cascaded_pass(startup_cps, tens)
150 sz += tens.storage_size()
151
152 for prev_cps in sg.cascaded_passes[: idx + 1]:
153 prev_cps.sram_used += sz
154
155 if callee.scratch_tensor is not None:
Patrik Gustavssoneca2e952020-05-27 09:15:11 +0200156 if callee.scratch_tensor.mem_area == MemArea.Sram:
157 cps.sram_used += callee.scratch_tensor.storage_size()
Patrik Gustavsson3ab94522020-06-29 17:36:55 +0200158
159 if callee.scratch_fast_tensor is not None:
160 if callee.scratch_fast_tensor.mem_area == MemArea.Sram:
161 cps.sram_used += callee.scratch_fast_tensor.storage_size()