1/*
2 * Licensed to the Apache Software Foundation (ASF) under one
3 * or more contributor license agreements. See the NOTICE file
4 * distributed with this work for additional information
5 * regarding copyright ownership. The ASF licenses this file
6 * to you under the Apache License, Version 2.0 (the
7 * "License"); you may not use this file except in compliance
8 * with the License. You may obtain a copy of the License at
9 *
10 * http://www.apache.org/licenses/LICENSE-2.0
11 *
12 * Unless required by applicable law or agreed to in writing,
13 * software distributed under the License is distributed on an
14 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15 * KIND, either express or implied. See the License for the
16 * specific language governing permissions and limitations
17 * under the License.
18 */
19#include "block_config.h"
20
21#include <tvm/runtime/container/array.h>
22#include <tvm/runtime/object.h>
23#include <tvm/runtime/registry.h>
24
25#include <utility>
26#include <vector>
27
28#include "common.h"
29
30namespace tvm {
31namespace contrib {
32namespace ethosu {
33namespace cascader {
34
35void BlockConfigNode::VisitAttrs(AttrVisitor* v) {
36 Array<Integer> tmp_arr = make_array(input_shape_);
37 v->Visit("_input_shape", &tmp_arr);
38 tmp_arr = make_array(output_shape_);
39 v->Visit("_output_shape", &tmp_arr);
40 v->Visit("_compute_cycles", &compute_cycles_);
41 v->Visit("_output_cycles", &output_cycles_);
42}
43
44BlockConfig::BlockConfig(const std::vector<int>& input_shape, const std::vector<int>& output_shape,
45 int compute_cycles, int output_cycles) {
46 auto n = make_object<BlockConfigNode>();
47 n->input_shape_ = std::move(input_shape);
48 n->output_shape_ = std::move(output_shape);
49 n->compute_cycles_ = compute_cycles;
50 n->output_cycles_ = output_cycles;
51 data_ = std::move(n);
52}
53
54TVM_REGISTER_GLOBAL("contrib.ethosu.cascader.BlockConfig")
55 .set_body_typed([](Array<Integer> input_shape, Array<Integer> output_shape, int compute_cycles,
56 int output_cycles) {
57 std::vector<int> vinput_shape = make_vector<int, Integer>(input_shape);
58 std::vector<int> voutput_shape = make_vector<int, Integer>(output_shape);
59 return BlockConfig(vinput_shape, voutput_shape, compute_cycles, output_cycles);
60 });
61
62TVM_REGISTER_NODE_TYPE(BlockConfigNode);
63
64} // namespace cascader
65} // namespace ethosu
66} // namespace contrib
67} // namespace tvm
68