From 5e5a7847b8fc1eb261c7561f44585d2f6b524df3 Mon Sep 17 00:00:00 2001 From: Diqing Zhong Date: Mon, 16 Aug 2021 17:24:09 +0200 Subject: TOSA raw data output - Add TOSA output generation in npz format Change-Id: I97822e3a93a8fef1a95a990f23ef2c4ca5a8f73a Signed-off-by: Diqing Zhong --- ethosu/vela/rawdata_writer.py | 74 +++++++++++++++++++++++++++++++++++++++++++ ethosu/vela/vela.py | 11 ++++--- 2 files changed, 81 insertions(+), 4 deletions(-) create mode 100644 ethosu/vela/rawdata_writer.py (limited to 'ethosu/vela') 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, + ) diff --git a/ethosu/vela/vela.py b/ethosu/vela/vela.py index 9e237f84..7400b8e9 100644 --- a/ethosu/vela/vela.py +++ b/ethosu/vela/vela.py @@ -27,6 +27,7 @@ import flatbuffers from . import architecture_features from . import compiler_driver from . import model_reader +from . import rawdata_writer from . import scheduler from . import stats_writer from . import tflite_writer @@ -83,18 +84,20 @@ def process(input_name, enable_debug_db, arch, model_reader_options, compiler_op arch=arch, ) - output_filename = output_basename + "_vela.tflite" + output_tfl_filename = output_basename + "_vela.tflite" if input_name.endswith(".tflite"): - tflite_writer.write_tflite(nng, output_filename) + tflite_writer.write_tflite(nng, output_tfl_filename) + elif input_name.endswith(".tosa"): + rawdata_writer.write_rawdata_output(nng, arch, output_basename) if enable_debug_db: - file_offsets = calculate_operator_file_offsets(output_filename) + file_offsets = calculate_operator_file_offsets(output_tfl_filename) for idx, offset in enumerate(sorted(file_offsets)): sg = find_subgraph_with_command_stream_order(nng, idx) if sg is not None: DebugDatabase.set_stream_offset(sg, offset) debug_filename = output_basename + "_debug.xml" - DebugDatabase.write(debug_filename, input_name, output_filename) + DebugDatabase.write(debug_filename, input_name, output_tfl_filename) if compiler_options.timing: stop = time.time() -- cgit v1.2.1