diff options
author | Diqing Zhong <diqing.zhong@arm.com> | 2021-08-16 17:24:09 +0200 |
---|---|---|
committer | patrik.gustavsson <patrik.gustavsson@arm.com> | 2021-09-01 12:23:28 +0000 |
commit | 5e5a7847b8fc1eb261c7561f44585d2f6b524df3 (patch) | |
tree | 586d7a83b6b4da362f879620940b0146ca4428a7 /ethosu/vela/rawdata_writer.py | |
parent | e389652ccb9821f8e959d533efd553f0d5e200d9 (diff) | |
download | ethos-u-vela-5e5a7847b8fc1eb261c7561f44585d2f6b524df3.tar.gz |
TOSA raw data output
- Add TOSA output generation in npz format
Change-Id: I97822e3a93a8fef1a95a990f23ef2c4ca5a8f73a
Signed-off-by: Diqing Zhong <diqing.zhong@arm.com>
Diffstat (limited to 'ethosu/vela/rawdata_writer.py')
-rw-r--r-- | ethosu/vela/rawdata_writer.py | 74 |
1 files changed, 74 insertions, 0 deletions
diff --git a/ethosu/vela/rawdata_writer.py b/ethosu/vela/rawdata_writer.py new file mode 100644 index 00000000..76765e60 --- /dev/null +++ b/ethosu/vela/rawdata_writer.py @@ -0,0 +1,74 @@ +# Copyright (C) 2021 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: +# Functions used to write to a raw format (.npz) file. +import numpy as np + +from .high_level_command_to_npu_op import get_region +from .nn_graph import PassPlacement +from .operation import Op + + +def write_rawdata_output(nng, arch, filename): + subgraphs_to_write = [sg for sg in nng.subgraphs if sg.placement == PassPlacement.Cpu] + + for sg_idx, sg in enumerate(subgraphs_to_write): + custom_op = None + for ps in sg.passes: + for op in ps.ops: + if op.type == Op.CustomNpuOp: + custom_op = op + break + if custom_op: + break + + if custom_op: + ifm_shapes = [] + ifm_regions = [] + ifm_offsets = [] + ofm_shapes = [] + ofm_regions = [] + ofm_offsets = [] + cmd_stream_tensor, weight_tensor, scratch_tensor, scratch_fast_tensor = custom_op.inputs[:4] + weight_region = get_region(weight_tensor.mem_type, arch) + scratch_region = get_region(scratch_tensor.mem_type, arch) + scratch_fast_region = get_region(scratch_fast_tensor.mem_type, arch) + for ifm in custom_op.inputs[4:]: + ifm_shapes.append(ifm.shape) + ifm_regions.append(get_region(ifm.mem_type, arch)) + ifm_offsets.append(ifm.address) + for ofm in custom_op.outputs: + ofm_shapes.append(ofm.shape) + ofm_regions.append(get_region(ofm.mem_type, arch)) + ofm_offsets.append(ofm.address) + + filename_sg = f"{filename}_sg{sg_idx}_vela.npz" + np.savez( + filename_sg, + cmd_data=cmd_stream_tensor.values, + weight_data=weight_tensor.values, + weight_region=weight_region, + scratch_shape=scratch_tensor.shape, + scratch_region=scratch_region, + scratch_fast_shape=scratch_fast_tensor.shape, + scratch_fast_region=scratch_fast_region, + input_shape=ifm_shapes, + input_region=ifm_regions, + input_offset=ifm_offsets, + output_shape=ofm_shapes, + output_region=ofm_regions, + output_offset=ofm_offsets, + ) |