Mousius commented on a change in pull request #8795: URL: https://github.com/apache/tvm/pull/8795#discussion_r695768083
########## File path: python/tvm/relay/backend/contrib/ethosu/legalize.py ########## @@ -0,0 +1,218 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you 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 +# +# http://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. +# pylint: disable=invalid-name, unused-argument, import-outside-toplevel, no-value-for-parameter +""" A set of passes to legalize some of operations for the NPU""" Review comment: Remove leading space. ########## File path: python/tvm/relay/backend/contrib/ethosu/legalize.py ########## @@ -0,0 +1,218 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you 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 +# +# http://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. +# pylint: disable=invalid-name, unused-argument, import-outside-toplevel, no-value-for-parameter +""" A set of passes to legalize some of operations for the NPU""" +import numpy as np + +import tvm +from tvm import relay +from tvm import ir +from tvm.relay.dataflow_pattern import DFPatternCallback +from tvm.relay.dataflow_pattern import wildcard +from tvm.relay.dataflow_pattern import is_op +from tvm.relay.dataflow_pattern import rewrite +from tvm.relay.backend.contrib.ethosu import op as ethosu_ops +from tvm.relay.backend.contrib.ethosu.errors import UnsupportedLayout +from tvm.relay.backend.contrib.ethosu import vela_api +from tvm.relay.op.contrib import ethosu as ethosu_patterns + + +class SplitRewriter(DFPatternCallback): + """This rewriting converts split operations into a sequence of + strided_slice operations, because codegen is going to be based + on strided_slices that will define the slice of the tensor that + will be fed to the consumer. + """ + + def __init__(self): + super().__init__(require_type=True) + self.split_in = wildcard() + self.pattern = is_op("split")(self.split_in) + + @staticmethod + def get_section_begin_coords(split): + """Currently, the split operator takes an array of indices or an integer + indicating the number of splits. However, its an array of indices could + represent both cases, therefore this function just make it an array of + indices where each index represent the co-ordinate of beginning of each + section -- defines as section begins. + + Parameters + ---------- + split : relay.Expr + The Relay Call expression for a split operator + + Returns + ------- + section_begins : list + A list containing integers corresponding to section + begins + """ + indices_or_sections = split.attrs.indices_or_sections + input_shape = split.args[0].checked_type.shape + split_axis = split.attrs.axis + + if isinstance(indices_or_sections, tvm.ir.container.Array): + # 0 is the beginning of the first section. + return [0] + list(indices_or_sections) + split_axis_len = input_shape[split_axis].value + section_length = split_axis_len // indices_or_sections.value + section_begins = list(range(0, split_axis_len, section_length)) Review comment: Can just return this without using another variable. ########## File path: src/relay/backend/contrib/ethosu/preprocess.cc ########## @@ -0,0 +1,269 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 + * + * http://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. + */ +#include <tvm/ir/error.h> +#include <tvm/relay/analysis.h> +#include <tvm/relay/attrs/annotation.h> +#include <tvm/relay/expr.h> +#include <tvm/relay/expr_functor.h> +#include <tvm/relay/transform.h> + +#include <unordered_map> +#include <unordered_set> +#include <utility> +#include <vector> + +#include "../../../op/make_op.h" + +namespace tvm { +namespace relay { +namespace contrib { +namespace ethosu { + +/*! + * \brief This expression rewriter will traverse the graph to find calls + * to all external functions. If they have multiple inputs and/or + * multiple outputs, the following has to be done : + * 1) If multiple inputs are present, they needed to be concat before the call. + * 2) Inside the external function they need to be split again to their original inputs. + * 3) If there are multiple outputs, they need to be concat at the end of external function. + * 4) Then, the concat output again need to be split and made the original tuple output in the + * main. + */ +class ExternalFuncIOHandler : public ExprRewriter { + public: + explicit ExternalFuncIOHandler(const IRModule& module) : module_(module) {} + int count = 0; + + Function InferType(const Function& expr, const IRModule& m) { + IRModule mod(m); + mod->Update(mod->GetGlobalVar("main"), expr); + mod = transform::InferType()(mod); + return Downcast<Function>(mod->Lookup("main")); + } + + /*! + * \brief This function will take shape and compute + * the scalar size value for it to be use to create + * flat single dimensional tensors. + */ + int64_t CalcSize(const Array<Integer>& shape) { + int size = 1; + for (auto dim_sz : shape) { + size = size * Downcast<Integer>(dim_sz)->value; + } + return size; + } + + /*! + * \brief This will take a tensor and create a flattened + * tensor to be used by the concat. + */ + Expr CreateFlattenTensor(const Expr& input) { + auto ishape = Downcast<Array<Integer>>(Downcast<TensorType>(input->checked_type())->shape); + int flatten_size = CalcSize(ishape); + Array<Integer> oshape = {Integer(flatten_size)}; + return MakeReshape(input, oshape); Review comment: ```suggestion Array<Integer> output_shape = {Integer(flatten_size)}; return MakeReshape(input, output_shape); ``` ########## File path: src/relay/op/contrib/ethosu/common.cc ########## @@ -0,0 +1,65 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 + * + * http://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. + */ + +/*! + * \file src/relay/op/contrib/ethosu/op_common.cc Review comment: File name is wrong here. ########## File path: tests/python/contrib/test_ethosu/test_preprocess.py ########## @@ -0,0 +1,346 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you 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 +# +# http://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. +# pylint: disable=invalid-name, unused-argument + +import pytest + +pytest.importorskip("ethosu.vela") +import numpy as np + +import tvm +from tvm import relay +from tvm.relay.backend.contrib.ethosu import preprocess + + +def set_func_attr(func, compile_name, symbol_name): + """ + Helper function to attach attributes to the external function. + """ + func = func.with_attr("Primitive", tvm.tir.IntImm("int32", 1)) + func = func.with_attr("Inline", tvm.tir.IntImm("int32", 1)) + func = func.with_attr("Compiler", compile_name) + func = func.with_attr("global_symbol", symbol_name) + return func + + +def test_single_io(): + """ + This test will test the pass wont touch external functions that + have a single input and a single output. + """ + + def create_graph(): + def create_external_func1(mod_, compiler_name, symbol_name): + x_int = relay.var("x_int", shape=(10, 10)) + z0 = relay.nn.relu(x_int) + f1 = relay.Function([x_int], z0) + f1 = set_func_attr(f1, compiler_name, symbol_name) + glb_f1 = relay.GlobalVar(symbol_name) + mod_[glb_f1] = f1 + mod_ = relay.transform.InferType()(mod_) + return glb_f1, mod_ + + mod = tvm.IRModule() + x = relay.var("x", shape=(10, 10)) + + glb_symbol_f1, mod = create_external_func1(mod, "ethosu", "ethosu_0") + r = relay.Call(glb_symbol_f1, [x]) + main = relay.Function([x], r) + mod["main"] = main + mod = relay.transform.InferType()(mod) + return mod + + mod = create_graph() + exp = create_graph() + mod = preprocess.preprocess_ext_io()(mod) + assert tvm.ir.structural_equal(mod, exp, map_free_vars=True) + + +def test_2ins_single_out(): + """ + The test is check two inputs and a single output of external function + """ + + def create_graph(): + def create_external_func1(mod_, compiler_name, symbol_name): + x_int = relay.var("x_int", shape=(10, 10)) + w0_int = relay.var("w0_int", shape=(10, 10)) + z0 = relay.add(x_int, w0_int) + + f1 = relay.Function([x_int, w0_int], z0) + f1 = set_func_attr(f1, compiler_name, symbol_name) + glb_f1 = relay.GlobalVar(symbol_name) + mod_[glb_f1] = f1 + mod_ = relay.transform.InferType()(mod_) + return glb_f1, mod_ + + mod = tvm.IRModule() + + x = relay.var("x", shape=(10, 10)) + w0 = relay.var("w0", shape=(10, 10)) + + glb_symbol_f1, mod = create_external_func1(mod, "ethosu", "ethosu_0") + r = relay.Call(glb_symbol_f1, [x, w0]) + main = relay.Function([x, w0], r) + mod["main"] = main + mod = relay.transform.InferType()(mod) + return mod + + def expected(): + def create_external_func1(mod_, compiler_name, symbol_name): + ifms_int = relay.var("ifms_int", shape=[200]) + + # splits + (x_int_flat, w0_int_flat) = relay.split(ifms_int, [100]) + # reshapes + x_int = relay.reshape(x_int_flat, newshape=(10, 10)) + w0_int = relay.reshape(w0_int_flat, newshape=(10, 10)) + + z0 = relay.add(x_int, w0_int) + f1 = relay.Function([ifms_int], z0) + f1 = set_func_attr(f1, compiler_name, symbol_name) + glb_f1 = relay.GlobalVar(symbol_name) + mod_[glb_f1] = f1 + mod_ = relay.transform.InferType()(mod_) + return glb_f1, mod_ + + mod = tvm.IRModule() + + x = relay.var("x", shape=(10, 10)) + w0 = relay.var("w0", shape=(10, 10)) + + # reshapes + x_reshaped = relay.reshape(x, newshape=100) + w0_reshaped = relay.reshape(w0, newshape=100) + + # concat + ifms = relay.concatenate((x_reshaped, w0_reshaped), 0) + + glb_symbol_f1, mod = create_external_func1(mod, "ethosu", "ethosu_0") + r = relay.Call(glb_symbol_f1, [ifms]) + main = relay.Function([x, w0], r) + mod["main"] = main + mod = relay.transform.InferType()(mod) + return mod + + mod = create_graph() + exp = expected() + mod = preprocess.preprocess_ext_io()(mod) + assert tvm.ir.structural_equal(mod, exp, map_free_vars=True) + + +def test_single_in_2outs(): + """ + The test is to check a single input and two outputs of external function + """ + + def create_graph(): + def create_external_func1(mod_, compiler_name, symbol_name): + x_int = relay.var("x_int", shape=(10, 10)) + + p0 = relay.nn.relu(x_int) + q0 = relay.tanh(x_int) + f1_o_tuple = relay.Tuple([p0, q0]) + + f1 = relay.Function([x_int], f1_o_tuple) + f1 = set_func_attr(f1, compiler_name, symbol_name) + glb_f1 = relay.GlobalVar(symbol_name) + mod_[glb_f1] = f1 + mod_ = relay.transform.InferType()(mod_) + return glb_f1, mod_ + + mod = tvm.IRModule() + x = relay.var("x", shape=(10, 10)) + glb_symbol_f1, mod = create_external_func1(mod, "ethosu", "ethosu_0") + pq_tuple = relay.Call(glb_symbol_f1, [x]) + p0 = relay.TupleGetItem(pq_tuple, 0) + q0 = relay.TupleGetItem(pq_tuple, 1) + r = relay.concatenate((p0, q0), axis=0) + main = relay.Function([x], r) + mod["main"] = main + mod = relay.transform.InferType()(mod) + return mod + + def expected(): + def create_external_func1(mod_, compiler_name, symbol_name): + x_int = relay.var("x_int", shape=(10, 10)) + + p0 = relay.nn.relu(x_int) + q0 = relay.tanh(x_int) + + # reshapes + p0_reshaped = relay.reshape(p0, newshape=100) + q0_reshaped = relay.reshape(q0, newshape=100) + ofms = relay.concatenate((p0_reshaped, q0_reshaped), 0) + + f1 = relay.Function([x_int], ofms) + f1 = set_func_attr(f1, compiler_name, symbol_name) + glb_f1 = relay.GlobalVar(symbol_name) + mod_[glb_f1] = f1 + mod_ = relay.transform.InferType()(mod_) + return glb_f1, mod_ + + mod = tvm.IRModule() + x = relay.var("x", shape=(10, 10)) + glb_symbol_f1, mod = create_external_func1(mod, "ethosu", "ethosu_0") + ofms = relay.Call(glb_symbol_f1, [x]) + + # splits + (p0_flat, q0_flat) = relay.split(ofms, [100]) + # reshapes + p0_flat_reshaped = relay.reshape(p0_flat, newshape=(10, 10)) + q0_flat_reshaped = relay.reshape(q0_flat, newshape=(10, 10)) + # original output + tuple_out = relay.Tuple([p0_flat_reshaped, q0_flat_reshaped]) + + p0 = relay.TupleGetItem(tuple_out, 0) + q0 = relay.TupleGetItem(tuple_out, 1) + r = relay.concatenate((p0, q0), axis=0) + main = relay.Function([x], r) + mod["main"] = main + mod = relay.transform.InferType()(mod) + return mod + + mod = create_graph() + exp = expected() + mod = relay.transform.InferType()(mod) + mod = preprocess.preprocess_ext_io()(mod) + assert tvm.ir.structural_equal(mod, exp, map_free_vars=True) + + +def test_4ins_2outs(): + """ + The test is to check a 4 inputs and two outputs of external function. + This just stand as a general test for multiple ins/outs. + """ + + def create_graph(): + def create_external_func1(mod_, compiler_name, symbol_name): + x_int = relay.var("x_int", shape=(10, 10)) + w0_int = relay.var("w0_int", shape=(10, 10)) + w1_int = relay.var("w1_int", shape=(10, 10)) + w2_int = relay.var("w2_int", shape=(10, 10)) + + z0 = relay.add(x_int, w0_int) + p0 = relay.subtract(z0, w1_int) + q0 = relay.multiply(z0, w2_int) + f1_o_tuple = relay.Tuple([p0, q0]) + + f1 = relay.Function([x_int, w0_int, w1_int, w2_int], f1_o_tuple) + f1 = set_func_attr(f1, compiler_name, symbol_name) + glb_f1 = relay.GlobalVar(symbol_name) + mod_[glb_f1] = f1 + mod_ = relay.transform.InferType()(mod_) + return glb_f1, mod_ + + mod = tvm.IRModule() + + x = relay.var("x", shape=(10, 10)) + w0 = relay.var("w0", shape=(10, 10)) + w1 = relay.var("w1", shape=(10, 10)) + w2 = relay.var("w2", shape=(10, 10)) + + glb_symbol_f1, mod = create_external_func1(mod, "ethosu", "ethosu_0") + pq_tuple = relay.Call(glb_symbol_f1, [x, w0, w1, w2]) + + p0 = relay.TupleGetItem(pq_tuple, 0) + q0 = relay.TupleGetItem(pq_tuple, 1) + r = relay.concatenate((p0, q0), axis=0) + main = relay.Function([x, w0, w1, w2], r) + mod["main"] = main + mod = relay.transform.InferType()(mod) + return mod + + def expected(): + def create_external_func1(mod_, compiler_name, symbol_name): + ifms_int = relay.var("ifms_int", shape=[400]) + + # splits + (x_int_flat, w0_int_flat, w1_int_flat, w2_int_flat) = relay.split( + ifms_int, [100, 200, 300] + ) + # reshapes + x_int = relay.reshape(x_int_flat, newshape=(10, 10)) + w0_int = relay.reshape(w0_int_flat, newshape=(10, 10)) + w1_int = relay.reshape(w1_int_flat, newshape=(10, 10)) + w2_int = relay.reshape(w2_int_flat, newshape=(10, 10)) + + z0 = relay.add(x_int, w0_int) + p0 = relay.subtract(z0, w1_int) + q0 = relay.multiply(z0, w2_int) + # f1_o_tuple = relay.Tuple([p0, q0]) + + # reshapes + p0_reshaped = relay.reshape(p0, newshape=100) + q0_reshaped = relay.reshape(q0, newshape=100) + ofms = relay.concatenate((p0_reshaped, q0_reshaped), 0) + + f1 = relay.Function([ifms_int], ofms) + f1 = set_func_attr(f1, compiler_name, symbol_name) + glb_f1 = relay.GlobalVar(symbol_name) + mod_[glb_f1] = f1 + mod_ = relay.transform.InferType()(mod_) + return glb_f1, mod_ + + mod = tvm.IRModule() + + x = relay.var("x", shape=(10, 10)) + w0 = relay.var("w0", shape=(10, 10)) + w1 = relay.var("w1", shape=(10, 10)) + w2 = relay.var("w2", shape=(10, 10)) + + # reshapes + x_reshaped = relay.reshape(x, newshape=100) + w0_reshaped = relay.reshape(w0, newshape=100) + w1_reshaped = relay.reshape(w1, newshape=100) + w2_reshaped = relay.reshape(w2, newshape=100) + + # concat + ifms = relay.concatenate((x_reshaped, w0_reshaped, w1_reshaped, w2_reshaped), 0) + + # call + glb_func, mod = create_external_func1(mod, "ethosu", "ethosu_0") + ofms = relay.Call(glb_func, [ifms]) + + # splits + (p0_flat, q0_flat) = relay.split(ofms, [100]) + # reshapes + p0_flat_reshaped = relay.reshape(p0_flat, newshape=(10, 10)) + q0_flat_reshaped = relay.reshape(q0_flat, newshape=(10, 10)) + # original output + tuple_out = relay.Tuple([p0_flat_reshaped, q0_flat_reshaped]) + + p0 = relay.TupleGetItem(tuple_out, 0) + q0 = relay.TupleGetItem(tuple_out, 1) + + r = relay.concatenate((p0, q0), axis=0) + main = relay.Function([x, w0, w1, w2], r) + mod["main"] = main + mod = relay.transform.InferType()(mod) + return mod + + mod = create_graph() + exp = expected() + mod = preprocess.preprocess_ext_io()(mod) + assert tvm.ir.structural_equal(mod, exp, map_free_vars=True) + + +if __name__ == "__main__": + test_2ins_single_out() + test_single_io() + test_4ins_2outs() + test_single_in_2outs() Review comment: ```suggestion sys.exit(pytest.main([__file__] + sys.argv[1:])) ``` ########## File path: tests/python/contrib/test_ethosu/test_legalize.py ########## @@ -0,0 +1,333 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you 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 +# +# http://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. +# pylint: disable=invalid-name, unused-argument + +import pytest + +pytest.importorskip("ethosu.vela") +import numpy as np + +import tvm +from tvm import relay +from tvm.relay.backend.contrib import ethosu +from tvm.relay.backend.contrib.ethosu import legalize, preprocess +from tvm.relay.dataflow_pattern import * +from tvm.relay.op.contrib.ethosu import * +import relay_ir_builder + + +def test_split_indices_legalize(): + def create_graph(axis): + x = relay.var("x", shape=(1, 50, 50, 3)) + x_relu = relay.nn.relu(x) + split_o = relay.split(x_relu, [5, 20, 45], axis).tuple_value Review comment: `split_o` ? ########## File path: python/tvm/relay/backend/contrib/ethosu/legalize.py ########## @@ -0,0 +1,218 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you 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 +# +# http://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. +# pylint: disable=invalid-name, unused-argument, import-outside-toplevel, no-value-for-parameter +""" A set of passes to legalize some of operations for the NPU""" +import numpy as np + +import tvm +from tvm import relay +from tvm import ir +from tvm.relay.dataflow_pattern import DFPatternCallback +from tvm.relay.dataflow_pattern import wildcard +from tvm.relay.dataflow_pattern import is_op +from tvm.relay.dataflow_pattern import rewrite +from tvm.relay.backend.contrib.ethosu import op as ethosu_ops +from tvm.relay.backend.contrib.ethosu.errors import UnsupportedLayout +from tvm.relay.backend.contrib.ethosu import vela_api +from tvm.relay.op.contrib import ethosu as ethosu_patterns + + +class SplitRewriter(DFPatternCallback): + """This rewriting converts split operations into a sequence of + strided_slice operations, because codegen is going to be based + on strided_slices that will define the slice of the tensor that + will be fed to the consumer. + """ + + def __init__(self): + super().__init__(require_type=True) + self.split_in = wildcard() + self.pattern = is_op("split")(self.split_in) + + @staticmethod + def get_section_begin_coords(split): + """Currently, the split operator takes an array of indices or an integer + indicating the number of splits. However, its an array of indices could + represent both cases, therefore this function just make it an array of + indices where each index represent the co-ordinate of beginning of each + section -- defines as section begins. + + Parameters + ---------- + split : relay.Expr + The Relay Call expression for a split operator + + Returns + ------- + section_begins : list + A list containing integers corresponding to section + begins + """ + indices_or_sections = split.attrs.indices_or_sections + input_shape = split.args[0].checked_type.shape + split_axis = split.attrs.axis + + if isinstance(indices_or_sections, tvm.ir.container.Array): + # 0 is the beginning of the first section. + return [0] + list(indices_or_sections) + split_axis_len = input_shape[split_axis].value + section_length = split_axis_len // indices_or_sections.value + section_begins = list(range(0, split_axis_len, section_length)) + return section_begins + + def callback(self, pre, post, node_map): + splits_types = dict() + split_input = post.args[0] + for idx, field_type in enumerate(post.checked_type.fields): + split = relay.TupleGetItem(post, idx) + splits_types[split] = field_type + + split_begins = list() + split_ends = list() + section_begins_in_split_axis = self.get_section_begin_coords(post) + for split_cord in section_begins_in_split_axis: + # first begin is [0, 0, ... , 0] + begin_shape = [0 for i in range(len(split_input.checked_type.shape))] + begin_shape[post.attrs.axis] = split_cord + split_begins.append(begin_shape) + + end_shape = list(split_input.checked_type.shape) + # Only the split axis coordinate changes + end_shape[post.attrs.axis] = split_cord + split_ends.append(end_shape) + + # Coordinates needs to be shifted left because beginning + # of the next section is the end of the previous + split_ends = split_ends[1:] + # Last section end is the shape of the tensor itself. + split_ends.append(list(split_input.checked_type.shape)) + + strided_slices = list() + for sb, se in zip(split_begins, split_ends): + strided_slices.append(relay.strided_slice(split_input, sb, se)) + + return relay.Tuple(strided_slices) + + +@ir.transform.module_pass(opt_level=1) +class LegalizeSplit: + """This is the pass that wraps SplitRewriter""" + + def transform_module(self, mod, ctx): + for gv, func in mod.functions.items(): Review comment: `gv` ? ########## File path: src/relay/op/contrib/ethosu/common.cc ########## @@ -0,0 +1,65 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 + * + * http://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. + */ + +/*! + * \file src/relay/op/contrib/ethosu/op_common.cc + * \brief A set of utilities and common functionality for Arm(R) Ethos(TM)-U NPU QNN ops. + */ + +#include "common.h" + +#include "../../op_common.h" + +namespace tvm { +namespace relay { +namespace op { +namespace contrib { +namespace ethosu { + +Array<IndexExpr> EthosuInferKernelOutput(Array<IndexExpr> ifm_shape, String ifm_layout, + String ofm_layout, Array<IndexExpr> kernel_shape, + IndexExpr ofm_channels, Array<IndexExpr> dilation, + Array<IndexExpr> strides, Array<IndexExpr> padding) { + // In the case of NHCWB16, convert the ifm shape to NHW (C not required for this function) + if (ifm_layout == "NHCWB16") { + ifm_shape = {ifm_shape[0], ifm_shape[1], ifm_shape[3]}; + } + Array<IndexExpr> oshape({ifm_shape[0], 0, 0, ofm_channels}); + + IndexExpr dilated_ksize_y = 1 + (kernel_shape[0] - 1) * dilation[0]; + IndexExpr dilated_ksize_x = 1 + (kernel_shape[1] - 1) * dilation[1]; + IndexExpr pad_h, pad_w; + GetPaddingHeightWidth(padding, &pad_h, &pad_w); + oshape.Set(1, indexdiv(ifm_shape[1] + pad_h - dilated_ksize_y, strides[0]) + 1); Review comment: `output_shape` ########## File path: python/tvm/relay/backend/contrib/ethosu/errors.py ########## @@ -0,0 +1,38 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you 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 +# +# http://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. +# pylint: disable=super-init-not-called +"""This module defines all error types associated with the Arm(R) Ethos(TM)-U NPU code generator.""" + + +class EthosUCodegenError(Exception): + """Base class for all exceptions related to code generation""" + + def __init__(self, data): + self.message = "EthosUCodegenError:" + data + + def __str__(self): + return self.message + + +class UnsupportedLayout(EthosUCodegenError): + """Raised when unsupported layout is encountered during code generation.""" + + def __init__(self, layout): + super().__init__(f"Unsupported Layout {layout}") + + def __str__(self): Review comment: I don't think you need to redefine `__str__`. ########## File path: python/tvm/relay/backend/contrib/ethosu/util.py ########## @@ -0,0 +1,183 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you 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 +# +# http://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. +""" +Helper utility Enums and Functions used through out code generation. + +The rest of the utility functions are misc. +Refer to the description inside such functions +""" + +from enum import Enum +import numpy as np + +from tvm import relay +from tvm.relay.build_module import bind_params_by_name +from tvm.relay.backend.contrib.ethosu import preprocess + + +class QConv2DArgs(Enum): + """ + This is a helper enum to obtain the correct index + of qnn.conv2d arguments. + """ + + IFM = 0 + WEIGHTS = 1 + IFM_ZERO_POINT = 2 + WEIGHTS_ZERO_POINT = 3 + IFM_SCALE = 4 + WEIGHTS_SCALE = 5 + + +class RequantArgs(Enum): + """ + This is a helper enum to obtain the correct index + of qnn.requantize arguments. + """ + + IFM_SCALE = 1 + IFM_ZERO_POINT = 2 + OFM_SCALE = 3 + OFM_ZERO_POINT = 4 + + +class BiasAddArgs(Enum): + """ + This is a helper enums to obtain the correct index + of qnn.bias_add arguments. + """ + + BIASES = 1 + + +class ClipArgs(Enum): + """ + This is a helper enums to obtain the correct index + of clip arguments. + """ + + A_MIN = 1 + A_MAX = 2 + + +def is_composite_func(func, name): + """ + This method checks whether the call is to + a composite function of a given name. + + Parameters + ---------- + func : relay.Function + The header to be displayed along with the dump. + + name : str + The candidate name to be checked + + Returns + -------- + a boolean + """ + + if not hasattr(func, "attrs"): + return False + if "Composite" not in func.attrs.keys(): + return False + composite_name = func.attrs["Composite"] + + if composite_name != name: + return False + return True Review comment: ```suggestion return composite_name == name ``` ########## File path: python/tvm/relay/backend/contrib/ethosu/legalize.py ########## @@ -0,0 +1,218 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you 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 +# +# http://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. +# pylint: disable=invalid-name, unused-argument, import-outside-toplevel, no-value-for-parameter +""" A set of passes to legalize some of operations for the NPU""" +import numpy as np + +import tvm +from tvm import relay +from tvm import ir +from tvm.relay.dataflow_pattern import DFPatternCallback +from tvm.relay.dataflow_pattern import wildcard +from tvm.relay.dataflow_pattern import is_op +from tvm.relay.dataflow_pattern import rewrite +from tvm.relay.backend.contrib.ethosu import op as ethosu_ops +from tvm.relay.backend.contrib.ethosu.errors import UnsupportedLayout +from tvm.relay.backend.contrib.ethosu import vela_api +from tvm.relay.op.contrib import ethosu as ethosu_patterns + + +class SplitRewriter(DFPatternCallback): + """This rewriting converts split operations into a sequence of + strided_slice operations, because codegen is going to be based + on strided_slices that will define the slice of the tensor that + will be fed to the consumer. + """ + + def __init__(self): + super().__init__(require_type=True) + self.split_in = wildcard() + self.pattern = is_op("split")(self.split_in) + + @staticmethod + def get_section_begin_coords(split): + """Currently, the split operator takes an array of indices or an integer + indicating the number of splits. However, its an array of indices could + represent both cases, therefore this function just make it an array of + indices where each index represent the co-ordinate of beginning of each + section -- defines as section begins. + + Parameters + ---------- + split : relay.Expr + The Relay Call expression for a split operator + + Returns + ------- + section_begins : list + A list containing integers corresponding to section + begins + """ + indices_or_sections = split.attrs.indices_or_sections + input_shape = split.args[0].checked_type.shape + split_axis = split.attrs.axis + + if isinstance(indices_or_sections, tvm.ir.container.Array): + # 0 is the beginning of the first section. + return [0] + list(indices_or_sections) + split_axis_len = input_shape[split_axis].value + section_length = split_axis_len // indices_or_sections.value + section_begins = list(range(0, split_axis_len, section_length)) + return section_begins + + def callback(self, pre, post, node_map): + splits_types = dict() + split_input = post.args[0] + for idx, field_type in enumerate(post.checked_type.fields): + split = relay.TupleGetItem(post, idx) + splits_types[split] = field_type + + split_begins = list() + split_ends = list() + section_begins_in_split_axis = self.get_section_begin_coords(post) + for split_cord in section_begins_in_split_axis: + # first begin is [0, 0, ... , 0] + begin_shape = [0 for i in range(len(split_input.checked_type.shape))] + begin_shape[post.attrs.axis] = split_cord + split_begins.append(begin_shape) + + end_shape = list(split_input.checked_type.shape) + # Only the split axis coordinate changes + end_shape[post.attrs.axis] = split_cord + split_ends.append(end_shape) + + # Coordinates needs to be shifted left because beginning + # of the next section is the end of the previous + split_ends = split_ends[1:] + # Last section end is the shape of the tensor itself. + split_ends.append(list(split_input.checked_type.shape)) + + strided_slices = list() + for sb, se in zip(split_begins, split_ends): + strided_slices.append(relay.strided_slice(split_input, sb, se)) + + return relay.Tuple(strided_slices) + + +@ir.transform.module_pass(opt_level=1) +class LegalizeSplit: + """This is the pass that wraps SplitRewriter""" + + def transform_module(self, mod, ctx): + for gv, func in mod.functions.items(): + func = rewrite(SplitRewriter(), func) + mod.update_func(gv, func) + return mod + + def __call__(self, *args, **kwargs): + pass + + +class EthosUConv2DRewriter(DFPatternCallback): + """Convert conv2d related composite functions into ethosu_conv2d operators""" + + def __init__(self): + super().__init__(require_type=True) + self.pattern = (wildcard().has_attr({"Composite": "ethosu.qnn_conv2d"}))(wildcard()) + + def callback(self, pre, post, node_map): + params = ethosu_patterns.QnnConv2DParams(post.op.body) + params.ifm.tensor = post.args[0] + channels_map = { + "NHWC": 3, + } + if str(params.ofm.layout) not in channels_map.keys(): + raise UnsupportedLayout(str(params.ofm.layout)) + kernel_size_map = { + "HWIO": params.weights.shape[0:2], + "OHWI": params.weights.shape[1:3], + "HWOI": params.weights.shape[0:2], + } + if str(params.weights.layout) not in kernel_size_map.keys(): + raise UnsupportedLayout(str(params.weights.layout)) + activation_map = {"clip": "CLIP"} + weight_to_ohwi_transform_map = {"HWIO": [3, 0, 1, 2]} + weights_values = params.weights.values + weights_values_ohwi = np.transpose( + weights_values, weight_to_ohwi_transform_map[str(params.weights.layout)] + ) + if params.activation: + activation = activation_map[params.activation.op.name] + clip_min = int(params.activation.attrs.a_min) + clip_max = int(params.activation.attrs.a_max) + else: + activation = "NONE" + clip_min = 0 + clip_max = 0 + scale_bias = vela_api.pack_biases( + biases=params.biases.tensor.data.asnumpy(), + ifm_scale=params.ifm.q_params.scale_f32, + ifm_dtype=np.dtype(params.ifm.dtype), + weight_scales=params.weights.q_params.scale_f32, + ofm_scale=params.ofm.q_params.scale_f32, + is_activation_tanh_or_sigmoid=activation in ["TANH", "SIGMOID"], + ) + ethosu_conv2d = ethosu_ops.ethosu_conv2d( + ifm=post.args[0], + weight=relay.const(weights_values_ohwi, params.weights.values.dtype), + scale_bias=relay.const(scale_bias, "uint8"), + lut=relay.const([], dtype="int8"), + ifm_scale=float(params.ifm.q_params.scale_f32), + ifm_zero_point=int(params.ifm.q_params.zero_point), + weight_zero_point=int(params.weights.q_params.zero_point), + ofm_scale=float(params.ofm.q_params.scale_f32), + ofm_zero_point=int(params.ofm.q_params.zero_point), + kernel_shape=kernel_size_map[str(params.weights.layout)], + ofm_channels=params.ofm.shape[channels_map[str(params.ofm.layout)]], + strides=params.strides, + padding=params.padding, + dilation=params.dilation, + activation=activation, + clip_min=clip_min, + clip_max=clip_max, + upscale="NONE", + ifm_layout=str(params.ifm.layout), + ofm_layout=str(params.ofm.layout), + ) + return ethosu_conv2d + + +@ir.transform.module_pass(opt_level=1) +class LegalizeEthosUConv2D: + """This is the pass that wraps the EthosUConv2DRewriter""" + + def transform_module(self, mod, ctx): + for gv, func in mod.functions.items(): + func = rewrite(EthosUConv2DRewriter(), func) + mod.update_func(gv, func) + return mod + + def __call__(self, *args, **kwargs): Review comment: Same here? ########## File path: python/tvm/relay/op/contrib/ethosu.py ########## @@ -0,0 +1,251 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you 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 +# +# http://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. +"""Arm(R) Ethos(TM)-U NPU supported operators.""" +import numpy as np + +from tvm.relay.expr import Constant +from tvm.relay.op.contrib.register import register_pattern_table +from tvm.relay.dataflow_pattern import wildcard, is_op, is_constant +from tvm.relay.backend.contrib.ethosu.util import QConv2DArgs +from tvm.relay.backend.contrib.ethosu.util import BiasAddArgs +from tvm.relay.backend.contrib.ethosu.util import RequantArgs +from tvm.relay.backend.contrib.ethosu.util import get_dim_value +from ethosu.vela import api as vapi + + +def check_strides(strides): + """This function checks whether strides are within the limits supported by the NPU""" + stride_range = (1, 3) + smin, smax = stride_range + if not smax >= strides[0] >= smin: + return False + if not smax >= strides[1] >= smin: + return False + return True + + +def check_valid_dtypes(tensor_params): + """This function checks whether dtypes are supported by the NPU""" + supported_dtypes = (np.uint8, np.int8) + for tep in tensor_params: + # Check for dtypes + if np.dtype(tep.dtype) not in supported_dtypes: + return False + # Check for shape sizes + if any(dimlen > 65536 for dimlen in tep.shape): + return False + return True + + +def check_weights(weights, dilation): + """This function checks whether weight tensor is compatible with the NPU""" + dilated_height_range = (1, 64) + dilated_hxw_range = (1, 64 * 64) + weights_limit = 127 * 65536 + dilated_width = (weights.shape[get_dim_value(weights.layout, "W")] - 1) * dilation[0] + 1 + dilated_height = (weights.shape[get_dim_value(weights.layout, "H")] - 1) * dilation[1] + 1 + dh_min, dh_max = dilated_height_range + if not dh_min <= dilated_height <= dh_max: + return False + dilated_hxw = dilated_height * dilated_width + dhxw_min, dhxw_max = dilated_hxw_range + if not dhxw_min <= dilated_hxw <= dhxw_max: + return False + # A saturation upper bound check for accumulators + weights.values = weights.values - weights.q_params.zero_point + axis = ( + get_dim_value(weights.layout, "H"), + get_dim_value(weights.layout, "W"), + get_dim_value(weights.layout, "I"), + ) + sum_weights = np.amax(np.sum(np.absolute(weights.values), axis=axis)) + if not sum_weights <= weights_limit: + return False + return True + + +def check_bias(bias): + """This function checks whether the bias values fit in 40 bits""" + if bias and bias.dtype == np.dtype("int64"): + valid = all(len(bin(bias_value)[2:]) <= 40 for bias_value in bias.values) + return valid + return True + + +def check_batch_size(ifm): + """This function checks for the number of batches vela currently supports""" + if ifm.shape[0] != 1: + return False + return True + + +def check_dilation(dilation): + """This function checks whether dilation is within the limits supported by the NPU""" + dilation_range = (1, 2) + dmin, dmax = dilation_range + if not dmin <= dilation[0] <= dmax: + return False + if not dmin <= dilation[1] <= dmax: + return False + return True + + +def check_padding(padding, bounds): + """This function checks whether padding is within the limits supported by the NPU""" + if len(padding) != 4 or len(bounds) != 4: + return False + top, left, bottom, right = padding + topb, leftb, bottomb, rightb = bounds + if top > topb or left > leftb or bottom > bottomb or right > rightb: + return False + return True + + +class TensorParams: + """ + This class will parse a tvm Expr along with quantization scale + and zero point to populate parameters that are required + for the creation of tensors in Vela. + """ + + def __init__(self, tensor, layout=None, scale=None, zero_point=None): + self.tensor = tensor + if isinstance(tensor, Constant): + self.values = tensor.data.asnumpy() + else: + self.values = None + self.dtype = tensor.checked_type.dtype + self.shape = [int(i) for i in tensor.checked_type.shape] + self.layout = layout + + if scale is not None and zero_point is not None: + self.q_params = vapi.NpuQuantization( + scale.data.asnumpy().astype("float32"), zero_point.data.asnumpy().astype(self.dtype) + ) + else: + # put default values + self.q_params = vapi.NpuQuantization(1.0, 0) + + +class QnnConv2DParams: + """ + This class will parse a Call to a ethosu.qnn_conv2d composite function + and extract quantization information of all the associated tensors. + """ + + composite_name = "ethosu.qnn_conv2d" + # The NPU only supports padding upto the numbers as follows + padding_bounds = [31, 31, 32, 32] + activation_map = {"clip": "CLIP"} + + def __init__(self, func_body): + activation = None + if str(func_body.op) in self.activation_map.keys(): + activation = func_body + requantize_op = activation.args[0] + else: + requantize_op = func_body + bias_add = requantize_op.args[0] + qnn_conv2d = bias_add.args[0] + data_layout = qnn_conv2d.attrs.data_layout + kernel_layout = qnn_conv2d.attrs.kernel_layout + # We consider the weights & biases as params as it should be a Constant + self.weights = TensorParams( + qnn_conv2d.args[QConv2DArgs.WEIGHTS.value], + kernel_layout, + qnn_conv2d.args[QConv2DArgs.WEIGHTS_SCALE.value], + qnn_conv2d.args[QConv2DArgs.WEIGHTS_ZERO_POINT.value], + ) + + self.biases = TensorParams( + bias_add.args[BiasAddArgs.BIASES.value], + data_layout, + requantize_op.args[RequantArgs.IFM_SCALE.value], + requantize_op.args[RequantArgs.IFM_ZERO_POINT.value], + ) + self.ifm = TensorParams( + qnn_conv2d.args[QConv2DArgs.IFM.value], + data_layout, + qnn_conv2d.args[QConv2DArgs.IFM_SCALE.value], + qnn_conv2d.args[QConv2DArgs.IFM_ZERO_POINT.value], + ) + self.ofm = TensorParams( + func_body, + data_layout, + requantize_op.args[RequantArgs.OFM_SCALE.value], + requantize_op.args[RequantArgs.OFM_ZERO_POINT.value], + ) + self.padding = qnn_conv2d.attrs.padding + self.strides = qnn_conv2d.attrs.strides + self.dilation = qnn_conv2d.attrs.dilation + self.activation = activation + + # If groups are equal to channel, its a depthwise_conv2d + self.groups = qnn_conv2d.attrs.groups + self.is_depthwise = False + channels_axis = {"HWIO": 3, "HWOI": 2} + if qnn_conv2d.attrs.groups == self.weights.shape[channels_axis[kernel_layout]]: + self.is_depthwise = True + + def is_valid(self): + """ + This function checks whether QnnConv2D has compatible attributes with the NPU + """ + tensor_params = [self.weights, self.ifm, self.ofm] + if not check_valid_dtypes(tensor_params): + return False + if not check_weights(self.weights, self.dilation): + return False + if not check_bias(self.biases): + return False + if not check_strides(self.strides): + return False + if not check_batch_size(self.ifm): + return False + if not check_dilation(self.dilation): + return False + if not check_padding(self.padding, self.padding_bounds): + return False + legal_groups = [1, self.ofm.shape[3]] + if self.groups not in legal_groups: + return False + # This should be a valid QnnDepthwise2DParams, not QnnConv2DParams + if self.is_depthwise: + return False + return True Review comment: ```suggestion return not self.is_depthwise ``` ########## File path: python/tvm/relay/op/contrib/ethosu.py ########## @@ -0,0 +1,251 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you 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 +# +# http://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. +"""Arm(R) Ethos(TM)-U NPU supported operators.""" +import numpy as np + +from tvm.relay.expr import Constant +from tvm.relay.op.contrib.register import register_pattern_table +from tvm.relay.dataflow_pattern import wildcard, is_op, is_constant +from tvm.relay.backend.contrib.ethosu.util import QConv2DArgs +from tvm.relay.backend.contrib.ethosu.util import BiasAddArgs +from tvm.relay.backend.contrib.ethosu.util import RequantArgs +from tvm.relay.backend.contrib.ethosu.util import get_dim_value +from ethosu.vela import api as vapi + + +def check_strides(strides): + """This function checks whether strides are within the limits supported by the NPU""" + stride_range = (1, 3) + smin, smax = stride_range + if not smax >= strides[0] >= smin: + return False + if not smax >= strides[1] >= smin: + return False + return True + + +def check_valid_dtypes(tensor_params): + """This function checks whether dtypes are supported by the NPU""" + supported_dtypes = (np.uint8, np.int8) + for tep in tensor_params: + # Check for dtypes + if np.dtype(tep.dtype) not in supported_dtypes: + return False + # Check for shape sizes + if any(dimlen > 65536 for dimlen in tep.shape): + return False + return True + + +def check_weights(weights, dilation): + """This function checks whether weight tensor is compatible with the NPU""" + dilated_height_range = (1, 64) + dilated_hxw_range = (1, 64 * 64) + weights_limit = 127 * 65536 + dilated_width = (weights.shape[get_dim_value(weights.layout, "W")] - 1) * dilation[0] + 1 + dilated_height = (weights.shape[get_dim_value(weights.layout, "H")] - 1) * dilation[1] + 1 + dh_min, dh_max = dilated_height_range + if not dh_min <= dilated_height <= dh_max: + return False + dilated_hxw = dilated_height * dilated_width + dhxw_min, dhxw_max = dilated_hxw_range + if not dhxw_min <= dilated_hxw <= dhxw_max: + return False + # A saturation upper bound check for accumulators + weights.values = weights.values - weights.q_params.zero_point + axis = ( + get_dim_value(weights.layout, "H"), + get_dim_value(weights.layout, "W"), + get_dim_value(weights.layout, "I"), + ) + sum_weights = np.amax(np.sum(np.absolute(weights.values), axis=axis)) + if not sum_weights <= weights_limit: + return False + return True + + +def check_bias(bias): + """This function checks whether the bias values fit in 40 bits""" + if bias and bias.dtype == np.dtype("int64"): + valid = all(len(bin(bias_value)[2:]) <= 40 for bias_value in bias.values) + return valid + return True + + +def check_batch_size(ifm): + """This function checks for the number of batches vela currently supports""" + if ifm.shape[0] != 1: + return False + return True + + +def check_dilation(dilation): + """This function checks whether dilation is within the limits supported by the NPU""" + dilation_range = (1, 2) + dmin, dmax = dilation_range + if not dmin <= dilation[0] <= dmax: + return False + if not dmin <= dilation[1] <= dmax: + return False + return True + + +def check_padding(padding, bounds): + """This function checks whether padding is within the limits supported by the NPU""" + if len(padding) != 4 or len(bounds) != 4: + return False + top, left, bottom, right = padding + topb, leftb, bottomb, rightb = bounds + if top > topb or left > leftb or bottom > bottomb or right > rightb: + return False + return True Review comment: ```suggestion return not (top > topb or left > leftb or bottom > bottomb or right > rightb) ``` ########## File path: python/tvm/relay/op/contrib/ethosu.py ########## @@ -0,0 +1,251 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you 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 +# +# http://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. +"""Arm(R) Ethos(TM)-U NPU supported operators.""" +import numpy as np + +from tvm.relay.expr import Constant +from tvm.relay.op.contrib.register import register_pattern_table +from tvm.relay.dataflow_pattern import wildcard, is_op, is_constant +from tvm.relay.backend.contrib.ethosu.util import QConv2DArgs +from tvm.relay.backend.contrib.ethosu.util import BiasAddArgs +from tvm.relay.backend.contrib.ethosu.util import RequantArgs +from tvm.relay.backend.contrib.ethosu.util import get_dim_value +from ethosu.vela import api as vapi + + +def check_strides(strides): + """This function checks whether strides are within the limits supported by the NPU""" + stride_range = (1, 3) + smin, smax = stride_range + if not smax >= strides[0] >= smin: + return False + if not smax >= strides[1] >= smin: + return False + return True + + +def check_valid_dtypes(tensor_params): + """This function checks whether dtypes are supported by the NPU""" + supported_dtypes = (np.uint8, np.int8) + for tep in tensor_params: + # Check for dtypes + if np.dtype(tep.dtype) not in supported_dtypes: + return False + # Check for shape sizes + if any(dimlen > 65536 for dimlen in tep.shape): + return False + return True + + +def check_weights(weights, dilation): + """This function checks whether weight tensor is compatible with the NPU""" + dilated_height_range = (1, 64) + dilated_hxw_range = (1, 64 * 64) + weights_limit = 127 * 65536 + dilated_width = (weights.shape[get_dim_value(weights.layout, "W")] - 1) * dilation[0] + 1 + dilated_height = (weights.shape[get_dim_value(weights.layout, "H")] - 1) * dilation[1] + 1 + dh_min, dh_max = dilated_height_range + if not dh_min <= dilated_height <= dh_max: + return False + dilated_hxw = dilated_height * dilated_width + dhxw_min, dhxw_max = dilated_hxw_range + if not dhxw_min <= dilated_hxw <= dhxw_max: + return False + # A saturation upper bound check for accumulators + weights.values = weights.values - weights.q_params.zero_point + axis = ( + get_dim_value(weights.layout, "H"), + get_dim_value(weights.layout, "W"), + get_dim_value(weights.layout, "I"), + ) + sum_weights = np.amax(np.sum(np.absolute(weights.values), axis=axis)) + if not sum_weights <= weights_limit: + return False + return True + + +def check_bias(bias): + """This function checks whether the bias values fit in 40 bits""" + if bias and bias.dtype == np.dtype("int64"): + valid = all(len(bin(bias_value)[2:]) <= 40 for bias_value in bias.values) + return valid + return True + + +def check_batch_size(ifm): + """This function checks for the number of batches vela currently supports""" + if ifm.shape[0] != 1: + return False + return True Review comment: ```suggestion return ifm.shape[0] == 1 ``` ########## File path: python/tvm/relay/backend/contrib/ethosu/legalize.py ########## @@ -0,0 +1,218 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you 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 +# +# http://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. +# pylint: disable=invalid-name, unused-argument, import-outside-toplevel, no-value-for-parameter +""" A set of passes to legalize some of operations for the NPU""" +import numpy as np + +import tvm +from tvm import relay +from tvm import ir +from tvm.relay.dataflow_pattern import DFPatternCallback +from tvm.relay.dataflow_pattern import wildcard +from tvm.relay.dataflow_pattern import is_op +from tvm.relay.dataflow_pattern import rewrite +from tvm.relay.backend.contrib.ethosu import op as ethosu_ops +from tvm.relay.backend.contrib.ethosu.errors import UnsupportedLayout +from tvm.relay.backend.contrib.ethosu import vela_api +from tvm.relay.op.contrib import ethosu as ethosu_patterns + + +class SplitRewriter(DFPatternCallback): + """This rewriting converts split operations into a sequence of + strided_slice operations, because codegen is going to be based + on strided_slices that will define the slice of the tensor that + will be fed to the consumer. + """ + + def __init__(self): + super().__init__(require_type=True) + self.split_in = wildcard() + self.pattern = is_op("split")(self.split_in) + + @staticmethod + def get_section_begin_coords(split): + """Currently, the split operator takes an array of indices or an integer + indicating the number of splits. However, its an array of indices could + represent both cases, therefore this function just make it an array of + indices where each index represent the co-ordinate of beginning of each + section -- defines as section begins. + + Parameters + ---------- + split : relay.Expr + The Relay Call expression for a split operator + + Returns + ------- + section_begins : list + A list containing integers corresponding to section + begins + """ + indices_or_sections = split.attrs.indices_or_sections + input_shape = split.args[0].checked_type.shape + split_axis = split.attrs.axis + + if isinstance(indices_or_sections, tvm.ir.container.Array): + # 0 is the beginning of the first section. + return [0] + list(indices_or_sections) + split_axis_len = input_shape[split_axis].value + section_length = split_axis_len // indices_or_sections.value + section_begins = list(range(0, split_axis_len, section_length)) + return section_begins + + def callback(self, pre, post, node_map): + splits_types = dict() + split_input = post.args[0] + for idx, field_type in enumerate(post.checked_type.fields): + split = relay.TupleGetItem(post, idx) + splits_types[split] = field_type + + split_begins = list() + split_ends = list() + section_begins_in_split_axis = self.get_section_begin_coords(post) + for split_cord in section_begins_in_split_axis: + # first begin is [0, 0, ... , 0] + begin_shape = [0 for i in range(len(split_input.checked_type.shape))] + begin_shape[post.attrs.axis] = split_cord + split_begins.append(begin_shape) + + end_shape = list(split_input.checked_type.shape) + # Only the split axis coordinate changes + end_shape[post.attrs.axis] = split_cord + split_ends.append(end_shape) + + # Coordinates needs to be shifted left because beginning + # of the next section is the end of the previous + split_ends = split_ends[1:] + # Last section end is the shape of the tensor itself. + split_ends.append(list(split_input.checked_type.shape)) + + strided_slices = list() + for sb, se in zip(split_begins, split_ends): + strided_slices.append(relay.strided_slice(split_input, sb, se)) + + return relay.Tuple(strided_slices) + + +@ir.transform.module_pass(opt_level=1) +class LegalizeSplit: + """This is the pass that wraps SplitRewriter""" + + def transform_module(self, mod, ctx): + for gv, func in mod.functions.items(): + func = rewrite(SplitRewriter(), func) + mod.update_func(gv, func) + return mod + + def __call__(self, *args, **kwargs): Review comment: Do we need this empty call? ########## File path: tests/python/contrib/test_ethosu/relay_ir_builder.py ########## @@ -0,0 +1,295 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you 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 +# +# http://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. +"""Helper module to build relay operations for testing""" + +from pathlib import Path +import numpy as np +import math + +import tvm +from tvm import relay +from tvm.relay.op.contrib import get_pattern_table +from tvm.relay import qnn +from tvm.relay.backend.contrib.ethosu.util import get_range_for_dtype_str + + +class TensorType: Review comment: It'd be great to use `NamedTuple` to clean up this testing code to set the data on instantiation similar to https://github.com/apache/tvm/blob/main/tests/python/relay/aot/aot_test_utils.py#L47 ########## File path: python/tvm/relay/op/contrib/ethosu.py ########## @@ -0,0 +1,251 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you 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 +# +# http://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. +"""Arm(R) Ethos(TM)-U NPU supported operators.""" +import numpy as np + +from tvm.relay.expr import Constant +from tvm.relay.op.contrib.register import register_pattern_table +from tvm.relay.dataflow_pattern import wildcard, is_op, is_constant +from tvm.relay.backend.contrib.ethosu.util import QConv2DArgs +from tvm.relay.backend.contrib.ethosu.util import BiasAddArgs +from tvm.relay.backend.contrib.ethosu.util import RequantArgs +from tvm.relay.backend.contrib.ethosu.util import get_dim_value +from ethosu.vela import api as vapi + + +def check_strides(strides): + """This function checks whether strides are within the limits supported by the NPU""" + stride_range = (1, 3) + smin, smax = stride_range + if not smax >= strides[0] >= smin: + return False + if not smax >= strides[1] >= smin: + return False + return True + + +def check_valid_dtypes(tensor_params): + """This function checks whether dtypes are supported by the NPU""" + supported_dtypes = (np.uint8, np.int8) + for tep in tensor_params: + # Check for dtypes + if np.dtype(tep.dtype) not in supported_dtypes: + return False + # Check for shape sizes + if any(dimlen > 65536 for dimlen in tep.shape): + return False + return True + + +def check_weights(weights, dilation): + """This function checks whether weight tensor is compatible with the NPU""" + dilated_height_range = (1, 64) + dilated_hxw_range = (1, 64 * 64) + weights_limit = 127 * 65536 + dilated_width = (weights.shape[get_dim_value(weights.layout, "W")] - 1) * dilation[0] + 1 + dilated_height = (weights.shape[get_dim_value(weights.layout, "H")] - 1) * dilation[1] + 1 + dh_min, dh_max = dilated_height_range + if not dh_min <= dilated_height <= dh_max: + return False + dilated_hxw = dilated_height * dilated_width + dhxw_min, dhxw_max = dilated_hxw_range + if not dhxw_min <= dilated_hxw <= dhxw_max: + return False + # A saturation upper bound check for accumulators + weights.values = weights.values - weights.q_params.zero_point + axis = ( + get_dim_value(weights.layout, "H"), + get_dim_value(weights.layout, "W"), + get_dim_value(weights.layout, "I"), + ) + sum_weights = np.amax(np.sum(np.absolute(weights.values), axis=axis)) + if not sum_weights <= weights_limit: + return False + return True Review comment: ```suggestion return sum_weights <= weights_limit ``` ########## File path: python/tvm/relay/backend/contrib/ethosu/legalize.py ########## @@ -0,0 +1,218 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you 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 +# +# http://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. +# pylint: disable=invalid-name, unused-argument, import-outside-toplevel, no-value-for-parameter +""" A set of passes to legalize some of operations for the NPU""" +import numpy as np + +import tvm +from tvm import relay +from tvm import ir +from tvm.relay.dataflow_pattern import DFPatternCallback +from tvm.relay.dataflow_pattern import wildcard +from tvm.relay.dataflow_pattern import is_op +from tvm.relay.dataflow_pattern import rewrite +from tvm.relay.backend.contrib.ethosu import op as ethosu_ops +from tvm.relay.backend.contrib.ethosu.errors import UnsupportedLayout +from tvm.relay.backend.contrib.ethosu import vela_api +from tvm.relay.op.contrib import ethosu as ethosu_patterns + + +class SplitRewriter(DFPatternCallback): + """This rewriting converts split operations into a sequence of + strided_slice operations, because codegen is going to be based + on strided_slices that will define the slice of the tensor that + will be fed to the consumer. + """ + + def __init__(self): + super().__init__(require_type=True) + self.split_in = wildcard() + self.pattern = is_op("split")(self.split_in) + + @staticmethod + def get_section_begin_coords(split): + """Currently, the split operator takes an array of indices or an integer + indicating the number of splits. However, its an array of indices could + represent both cases, therefore this function just make it an array of + indices where each index represent the co-ordinate of beginning of each + section -- defines as section begins. + + Parameters + ---------- + split : relay.Expr + The Relay Call expression for a split operator + + Returns + ------- + section_begins : list + A list containing integers corresponding to section + begins + """ + indices_or_sections = split.attrs.indices_or_sections + input_shape = split.args[0].checked_type.shape + split_axis = split.attrs.axis + + if isinstance(indices_or_sections, tvm.ir.container.Array): + # 0 is the beginning of the first section. + return [0] + list(indices_or_sections) + split_axis_len = input_shape[split_axis].value + section_length = split_axis_len // indices_or_sections.value + section_begins = list(range(0, split_axis_len, section_length)) + return section_begins + + def callback(self, pre, post, node_map): + splits_types = dict() + split_input = post.args[0] + for idx, field_type in enumerate(post.checked_type.fields): + split = relay.TupleGetItem(post, idx) + splits_types[split] = field_type + + split_begins = list() + split_ends = list() + section_begins_in_split_axis = self.get_section_begin_coords(post) + for split_cord in section_begins_in_split_axis: + # first begin is [0, 0, ... , 0] + begin_shape = [0 for i in range(len(split_input.checked_type.shape))] + begin_shape[post.attrs.axis] = split_cord + split_begins.append(begin_shape) + + end_shape = list(split_input.checked_type.shape) + # Only the split axis coordinate changes + end_shape[post.attrs.axis] = split_cord + split_ends.append(end_shape) + + # Coordinates needs to be shifted left because beginning + # of the next section is the end of the previous + split_ends = split_ends[1:] + # Last section end is the shape of the tensor itself. + split_ends.append(list(split_input.checked_type.shape)) + + strided_slices = list() + for sb, se in zip(split_begins, split_ends): + strided_slices.append(relay.strided_slice(split_input, sb, se)) + + return relay.Tuple(strided_slices) + + +@ir.transform.module_pass(opt_level=1) +class LegalizeSplit: + """This is the pass that wraps SplitRewriter""" + + def transform_module(self, mod, ctx): + for gv, func in mod.functions.items(): + func = rewrite(SplitRewriter(), func) + mod.update_func(gv, func) + return mod + + def __call__(self, *args, **kwargs): + pass + + +class EthosUConv2DRewriter(DFPatternCallback): + """Convert conv2d related composite functions into ethosu_conv2d operators""" + + def __init__(self): + super().__init__(require_type=True) + self.pattern = (wildcard().has_attr({"Composite": "ethosu.qnn_conv2d"}))(wildcard()) + + def callback(self, pre, post, node_map): + params = ethosu_patterns.QnnConv2DParams(post.op.body) + params.ifm.tensor = post.args[0] + channels_map = { + "NHWC": 3, + } + if str(params.ofm.layout) not in channels_map.keys(): + raise UnsupportedLayout(str(params.ofm.layout)) + kernel_size_map = { + "HWIO": params.weights.shape[0:2], + "OHWI": params.weights.shape[1:3], + "HWOI": params.weights.shape[0:2], + } + if str(params.weights.layout) not in kernel_size_map.keys(): + raise UnsupportedLayout(str(params.weights.layout)) + activation_map = {"clip": "CLIP"} + weight_to_ohwi_transform_map = {"HWIO": [3, 0, 1, 2]} + weights_values = params.weights.values + weights_values_ohwi = np.transpose( + weights_values, weight_to_ohwi_transform_map[str(params.weights.layout)] + ) + if params.activation: + activation = activation_map[params.activation.op.name] + clip_min = int(params.activation.attrs.a_min) + clip_max = int(params.activation.attrs.a_max) + else: + activation = "NONE" + clip_min = 0 + clip_max = 0 + scale_bias = vela_api.pack_biases( + biases=params.biases.tensor.data.asnumpy(), + ifm_scale=params.ifm.q_params.scale_f32, + ifm_dtype=np.dtype(params.ifm.dtype), + weight_scales=params.weights.q_params.scale_f32, + ofm_scale=params.ofm.q_params.scale_f32, + is_activation_tanh_or_sigmoid=activation in ["TANH", "SIGMOID"], + ) + ethosu_conv2d = ethosu_ops.ethosu_conv2d( + ifm=post.args[0], + weight=relay.const(weights_values_ohwi, params.weights.values.dtype), + scale_bias=relay.const(scale_bias, "uint8"), + lut=relay.const([], dtype="int8"), + ifm_scale=float(params.ifm.q_params.scale_f32), + ifm_zero_point=int(params.ifm.q_params.zero_point), + weight_zero_point=int(params.weights.q_params.zero_point), + ofm_scale=float(params.ofm.q_params.scale_f32), + ofm_zero_point=int(params.ofm.q_params.zero_point), + kernel_shape=kernel_size_map[str(params.weights.layout)], + ofm_channels=params.ofm.shape[channels_map[str(params.ofm.layout)]], + strides=params.strides, + padding=params.padding, + dilation=params.dilation, + activation=activation, + clip_min=clip_min, + clip_max=clip_max, + upscale="NONE", + ifm_layout=str(params.ifm.layout), + ofm_layout=str(params.ofm.layout), + ) + return ethosu_conv2d + + +@ir.transform.module_pass(opt_level=1) +class LegalizeEthosUConv2D: + """This is the pass that wraps the EthosUConv2DRewriter""" + + def transform_module(self, mod, ctx): + for gv, func in mod.functions.items(): Review comment: `gv` again, I think it means `global_var`? ########## File path: tests/python/contrib/test_ethosu/test_legalize.py ########## @@ -0,0 +1,333 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you 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 +# +# http://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. +# pylint: disable=invalid-name, unused-argument + +import pytest + +pytest.importorskip("ethosu.vela") +import numpy as np + +import tvm +from tvm import relay +from tvm.relay.backend.contrib import ethosu +from tvm.relay.backend.contrib.ethosu import legalize, preprocess +from tvm.relay.dataflow_pattern import * +from tvm.relay.op.contrib.ethosu import * +import relay_ir_builder + + +def test_split_indices_legalize(): + def create_graph(axis): + x = relay.var("x", shape=(1, 50, 50, 3)) + x_relu = relay.nn.relu(x) + split_o = relay.split(x_relu, [5, 20, 45], axis).tuple_value + return relay.Function([x], split_o) + + def expected_mod_axis1(): + expected_ir_string = """ + #[version = "0.0.5"] + def @tvmgen_default_ethosu_main_0(%x: Tensor[(1, 50, 50, 3), float32]) -> (Tensor[(1, 5, 50, 3), float32],\ + Tensor[(1, 15, 50, 3), float32],\ + Tensor[(1, 25, 50, 3), float32],\ + Tensor[(1, 5, 50, 3), float32]) { + %0 = nn.relu(%x) /* ty=Tensor[(1, 50, 50, 3), float32] */; + %1 = strided_slice(%0, begin=[0, 0, 0, 0], end=[1, 5, 50, 3], strides=[1], axes=None)\ + /* ty=Tensor[(1, 5, 50, 3), float32] */; + %2 = strided_slice(%0, begin=[0, 5, 0, 0], end=[1, 20, 50, 3], strides=[1], axes=None)\ + /* ty=Tensor[(1, 15, 50, 3), float32] */; + %3 = strided_slice(%0, begin=[0, 20, 0, 0], end=[1, 45, 50, 3], strides=[1], axes=None)\ + /* ty=Tensor[(1, 25, 50, 3), float32] */; + %4 = strided_slice(%0, begin=[0, 45, 0, 0], end=[1, 50, 50, 3], strides=[1], axes=None)\ + /* ty=Tensor[(1, 5, 50, 3), float32] */; + (%1, %2, %3, %4) + } + """ + return tvm.parser.fromtext(expected_ir_string) + + def expected_mod_axis2(): + expected_ir_string = """ + #[version = "0.0.5"] + def @tvmgen_default_ethosu_main_0(%x: Tensor[(1, 50, 50, 3), float32]) -> (Tensor[(1, 50, 5, 3), float32],\ + Tensor[(1, 50, 15, 3), float32],\ + Tensor[(1, 50, 25, 3), float32],\ + Tensor[(1, 50, 5, 3), float32]) { + %0 = nn.relu(%x) /* ty=Tensor[(1, 50, 50, 3), float32] */; + %1 = strided_slice(%0, begin=[0, 0, 0, 0], end=[1, 50, 5, 3], strides=[1], axes=None)\ + /* ty=Tensor[(1, 50, 5, 3), float32] */; + %2 = strided_slice(%0, begin=[0, 0, 5, 0], end=[1, 50, 20, 3], strides=[1], axes=None)\ + /* ty=Tensor[(1, 50, 15, 3), float32] */; + %3 = strided_slice(%0, begin=[0, 0, 20, 0], end=[1, 50, 45, 3], strides=[1], axes=None)\ + /* ty=Tensor[(1, 50, 25, 3), float32] */; + %4 = strided_slice(%0, begin=[0, 0, 45, 0], end=[1, 50, 50, 3], strides=[1], axes=None)\ + /* ty=Tensor[(1, 50, 5, 3), float32] */; + (%1, %2, %3, %4) + } + """ + return tvm.parser.fromtext(expected_ir_string) + + mod_axis1 = tvm.IRModule() + mod_axis1["tvmgen_default_ethosu_main_0"] = create_graph(1) + mod_axis1 = legalize.LegalizeSplit()(mod_axis1) + expected_axis1 = expected_mod_axis1() + tvm.ir.assert_structural_equal(mod_axis1, expected_axis1) + + mod_axis2 = tvm.IRModule() + mod_axis2["tvmgen_default_ethosu_main_0"] = create_graph(2) + mod_axis2 = legalize.LegalizeSplit()(mod_axis2) + expected_axis2 = expected_mod_axis2() + tvm.ir.assert_structural_equal(mod_axis2, expected_axis2) + + +def test_split_sections_legalize(): + def create_graph(axis, sections): + x = relay.var("x", shape=(1, 50, 50, 3)) + x_abs = relay.abs(x) + split_o = relay.split(x_abs, sections, axis).tuple_value Review comment: `split_o` ? ########## File path: src/relay/op/contrib/ethosu/convolution.cc ########## @@ -0,0 +1,212 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 + * + * http://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. + */ + +/*! + * \file src/relay/op/contrib/ethosu/convolution.cc + * \brief Property def of the Arm(R) Ethos(TM)-U NPU convolution ops. + */ +#include "../../nn/convolution.h" + +#include <tvm/relay/base.h> +#include <tvm/relay/op.h> +#include <tvm/relay/qnn/attrs.h> +#include <tvm/tir/analysis.h> +#include <tvm/tir/data_layout.h> + +#include "../../../qnn/utils.h" +#include "common.h" + +namespace tvm { +namespace relay { +namespace op { +namespace contrib { +namespace ethosu { + +/*! \brief Attributes used by the Ethos(TM)-U NPU convolution operator */ +struct EthosuConv2DAttrs : public tvm::AttrsNode<EthosuConv2DAttrs> { + double ifm_scale; + int ifm_zero_point; + int weight_zero_point; + double ofm_scale; + int ofm_zero_point; + Array<IndexExpr> kernel_shape; + IndexExpr ofm_channels; + Array<IndexExpr> strides; + Array<IndexExpr> padding; + Array<IndexExpr> dilation; + String activation; + int clip_min; + int clip_max; + String upscale; + tvm::String ifm_layout; + tvm::String ofm_layout; + + TVM_DECLARE_ATTRS(EthosuConv2DAttrs, "relay.attrs.EthosuConv2DAttrs") { + TVM_ATTR_FIELD(ifm_scale).describe("The quantization scale for the Input Feature Map tensor."); + TVM_ATTR_FIELD(ifm_zero_point) + .describe("The quantization zero point for the Input Feature Map tensor."); + TVM_ATTR_FIELD(weight_zero_point) + .describe("The quantization zero point for the weight tensor."); + TVM_ATTR_FIELD(ofm_scale).describe("The quantization scale for the Output Feature Map tensor."); + TVM_ATTR_FIELD(ofm_zero_point) + .describe("The quantization zero point for the Output Feature Map tensor."); + TVM_ATTR_FIELD(kernel_shape) + .describe("The 2 dimensional kernel shape as (kernel_height, kernel_width).") + .set_default(NullValue<Array<IndexExpr> >()); Review comment: Weird space here. ########## File path: python/tvm/relay/backend/contrib/ethosu/legalize.py ########## @@ -0,0 +1,218 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you 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 +# +# http://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. +# pylint: disable=invalid-name, unused-argument, import-outside-toplevel, no-value-for-parameter +""" A set of passes to legalize some of operations for the NPU""" +import numpy as np + +import tvm +from tvm import relay +from tvm import ir +from tvm.relay.dataflow_pattern import DFPatternCallback +from tvm.relay.dataflow_pattern import wildcard +from tvm.relay.dataflow_pattern import is_op +from tvm.relay.dataflow_pattern import rewrite +from tvm.relay.backend.contrib.ethosu import op as ethosu_ops +from tvm.relay.backend.contrib.ethosu.errors import UnsupportedLayout +from tvm.relay.backend.contrib.ethosu import vela_api +from tvm.relay.op.contrib import ethosu as ethosu_patterns + + +class SplitRewriter(DFPatternCallback): + """This rewriting converts split operations into a sequence of + strided_slice operations, because codegen is going to be based + on strided_slices that will define the slice of the tensor that + will be fed to the consumer. + """ + + def __init__(self): + super().__init__(require_type=True) + self.split_in = wildcard() + self.pattern = is_op("split")(self.split_in) + + @staticmethod + def get_section_begin_coords(split): + """Currently, the split operator takes an array of indices or an integer + indicating the number of splits. However, its an array of indices could + represent both cases, therefore this function just make it an array of + indices where each index represent the co-ordinate of beginning of each + section -- defines as section begins. + + Parameters + ---------- + split : relay.Expr + The Relay Call expression for a split operator + + Returns + ------- + section_begins : list + A list containing integers corresponding to section + begins + """ + indices_or_sections = split.attrs.indices_or_sections + input_shape = split.args[0].checked_type.shape + split_axis = split.attrs.axis + + if isinstance(indices_or_sections, tvm.ir.container.Array): + # 0 is the beginning of the first section. + return [0] + list(indices_or_sections) + split_axis_len = input_shape[split_axis].value + section_length = split_axis_len // indices_or_sections.value + section_begins = list(range(0, split_axis_len, section_length)) + return section_begins + + def callback(self, pre, post, node_map): + splits_types = dict() + split_input = post.args[0] + for idx, field_type in enumerate(post.checked_type.fields): Review comment: ```suggestion splits_types = { relay.TupleGetItem(post, idx): field_type for idx, field_type in enumerate(post.checked_type.fields) } ``` ########## File path: src/relay/backend/contrib/ethosu/preprocess.cc ########## @@ -0,0 +1,269 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 + * + * http://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. + */ +#include <tvm/ir/error.h> +#include <tvm/relay/analysis.h> +#include <tvm/relay/attrs/annotation.h> +#include <tvm/relay/expr.h> +#include <tvm/relay/expr_functor.h> +#include <tvm/relay/transform.h> + +#include <unordered_map> +#include <unordered_set> +#include <utility> +#include <vector> + +#include "../../../op/make_op.h" + +namespace tvm { +namespace relay { +namespace contrib { +namespace ethosu { + +/*! + * \brief This expression rewriter will traverse the graph to find calls + * to all external functions. If they have multiple inputs and/or + * multiple outputs, the following has to be done : + * 1) If multiple inputs are present, they needed to be concat before the call. + * 2) Inside the external function they need to be split again to their original inputs. + * 3) If there are multiple outputs, they need to be concat at the end of external function. + * 4) Then, the concat output again need to be split and made the original tuple output in the + * main. + */ +class ExternalFuncIOHandler : public ExprRewriter { + public: + explicit ExternalFuncIOHandler(const IRModule& module) : module_(module) {} + int count = 0; + + Function InferType(const Function& expr, const IRModule& m) { + IRModule mod(m); + mod->Update(mod->GetGlobalVar("main"), expr); + mod = transform::InferType()(mod); + return Downcast<Function>(mod->Lookup("main")); + } + + /*! + * \brief This function will take shape and compute + * the scalar size value for it to be use to create + * flat single dimensional tensors. + */ + int64_t CalcSize(const Array<Integer>& shape) { + int size = 1; + for (auto dim_sz : shape) { Review comment: Can we go for the full `dim_size` here? ########## File path: python/tvm/relay/backend/contrib/ethosu/legalize.py ########## @@ -0,0 +1,200 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you 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 +# +# http://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. +# pylint: disable=invalid-name, unused-argument, import-outside-toplevel +""" A set of passes to legalize some of operations for the NPU""" +import numpy as np + +import tvm +from tvm import relay +from tvm.relay.dataflow_pattern import DFPatternCallback +from tvm.relay.dataflow_pattern import wildcard +from tvm.relay.dataflow_pattern import is_op +from tvm.relay.dataflow_pattern import rewrite +from tvm.relay.backend.contrib.ethosu import op as ethosu_ops +from tvm.relay.backend.contrib.ethosu.errors import UnsupportedLayout +from tvm.relay.backend.contrib.ethosu import vela_api +from tvm.relay.op.contrib import ethosu as ethosu_patterns + + +class SplitRewriter(DFPatternCallback): + """Convert split operations to bunch of strided_slice operations, + because codegen is going to be based on strided_slices that are + close to in/out boxes of Vela High-Level Command Stream (HLCS). + Moreover, Vela HLCS is a high-level description of the supported + hardware operator. + """ + + def __init__(self): + super().__init__(require_type=True) + self.split_in = wildcard() + self.pattern = is_op("split")(self.split_in) + + @staticmethod + def get_section_begin_coords(split): + """Currently, the split can take an array of indices or an integer + indicating the number of splits. This helper functions unifies + this by making it a array of section begins. + + Parameters + ---------- + split : relay.Expr + The relay expression for split operator + + Returns + ------- + section_begins : list Review comment: Should this be `list[int]` ? ########## File path: tests/python/contrib/test_ethosu/test_legalize.py ########## @@ -0,0 +1,333 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you 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 +# +# http://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. +# pylint: disable=invalid-name, unused-argument + +import pytest + +pytest.importorskip("ethosu.vela") +import numpy as np + +import tvm +from tvm import relay +from tvm.relay.backend.contrib import ethosu +from tvm.relay.backend.contrib.ethosu import legalize, preprocess +from tvm.relay.dataflow_pattern import * +from tvm.relay.op.contrib.ethosu import * +import relay_ir_builder + + +def test_split_indices_legalize(): + def create_graph(axis): + x = relay.var("x", shape=(1, 50, 50, 3)) + x_relu = relay.nn.relu(x) + split_o = relay.split(x_relu, [5, 20, 45], axis).tuple_value + return relay.Function([x], split_o) + + def expected_mod_axis1(): + expected_ir_string = """ + #[version = "0.0.5"] + def @tvmgen_default_ethosu_main_0(%x: Tensor[(1, 50, 50, 3), float32]) -> (Tensor[(1, 5, 50, 3), float32],\ + Tensor[(1, 15, 50, 3), float32],\ + Tensor[(1, 25, 50, 3), float32],\ + Tensor[(1, 5, 50, 3), float32]) { + %0 = nn.relu(%x) /* ty=Tensor[(1, 50, 50, 3), float32] */; + %1 = strided_slice(%0, begin=[0, 0, 0, 0], end=[1, 5, 50, 3], strides=[1], axes=None)\ + /* ty=Tensor[(1, 5, 50, 3), float32] */; + %2 = strided_slice(%0, begin=[0, 5, 0, 0], end=[1, 20, 50, 3], strides=[1], axes=None)\ + /* ty=Tensor[(1, 15, 50, 3), float32] */; + %3 = strided_slice(%0, begin=[0, 20, 0, 0], end=[1, 45, 50, 3], strides=[1], axes=None)\ + /* ty=Tensor[(1, 25, 50, 3), float32] */; + %4 = strided_slice(%0, begin=[0, 45, 0, 0], end=[1, 50, 50, 3], strides=[1], axes=None)\ + /* ty=Tensor[(1, 5, 50, 3), float32] */; + (%1, %2, %3, %4) + } + """ + return tvm.parser.fromtext(expected_ir_string) + + def expected_mod_axis2(): + expected_ir_string = """ + #[version = "0.0.5"] + def @tvmgen_default_ethosu_main_0(%x: Tensor[(1, 50, 50, 3), float32]) -> (Tensor[(1, 50, 5, 3), float32],\ + Tensor[(1, 50, 15, 3), float32],\ + Tensor[(1, 50, 25, 3), float32],\ + Tensor[(1, 50, 5, 3), float32]) { + %0 = nn.relu(%x) /* ty=Tensor[(1, 50, 50, 3), float32] */; + %1 = strided_slice(%0, begin=[0, 0, 0, 0], end=[1, 50, 5, 3], strides=[1], axes=None)\ + /* ty=Tensor[(1, 50, 5, 3), float32] */; + %2 = strided_slice(%0, begin=[0, 0, 5, 0], end=[1, 50, 20, 3], strides=[1], axes=None)\ + /* ty=Tensor[(1, 50, 15, 3), float32] */; + %3 = strided_slice(%0, begin=[0, 0, 20, 0], end=[1, 50, 45, 3], strides=[1], axes=None)\ + /* ty=Tensor[(1, 50, 25, 3), float32] */; + %4 = strided_slice(%0, begin=[0, 0, 45, 0], end=[1, 50, 50, 3], strides=[1], axes=None)\ + /* ty=Tensor[(1, 50, 5, 3), float32] */; + (%1, %2, %3, %4) + } + """ + return tvm.parser.fromtext(expected_ir_string) + + mod_axis1 = tvm.IRModule() + mod_axis1["tvmgen_default_ethosu_main_0"] = create_graph(1) + mod_axis1 = legalize.LegalizeSplit()(mod_axis1) + expected_axis1 = expected_mod_axis1() + tvm.ir.assert_structural_equal(mod_axis1, expected_axis1) + + mod_axis2 = tvm.IRModule() + mod_axis2["tvmgen_default_ethosu_main_0"] = create_graph(2) + mod_axis2 = legalize.LegalizeSplit()(mod_axis2) + expected_axis2 = expected_mod_axis2() + tvm.ir.assert_structural_equal(mod_axis2, expected_axis2) + + +def test_split_sections_legalize(): + def create_graph(axis, sections): + x = relay.var("x", shape=(1, 50, 50, 3)) + x_abs = relay.abs(x) + split_o = relay.split(x_abs, sections, axis).tuple_value + outputs = list() + for section_idx in range(sections): + split_single_out = relay.TupleGetItem(split_o, section_idx) + tanh = relay.tanh(split_single_out) + outputs.append(tanh) + tuple_out = relay.Tuple(outputs) + return relay.Function([x], tuple_out) + + def expected_mod_axis1(): + expected_ir_string = """ + #[version = "0.0.5"] + def @tvmgen_default_ethosu_main_0(%x: Tensor[(1, 50, 50, 3), float32]) -> (Tensor[(1, 10, 50, 3), float32],\ + Tensor[(1, 10, 50, 3), float32],\ + Tensor[(1, 10, 50, 3), float32],\ + Tensor[(1, 10, 50, 3), float32],\ + Tensor[(1, 10, 50, 3), float32]) { + %0 = abs(%x) /* ty=Tensor[(1, 50, 50, 3), float32] */; + %1 = strided_slice(%0, begin=[0, 0, 0, 0], end=[1, 10, 50, 3], strides=[1], axes=None)\ + /* ty=Tensor[(1, 10, 50, 3), float32] */; + %2 = strided_slice(%0, begin=[0, 10, 0, 0], end=[1, 20, 50, 3], strides=[1], axes=None)\ + /* ty=Tensor[(1, 10, 50, 3), float32] */; + %3 = strided_slice(%0, begin=[0, 20, 0, 0], end=[1, 30, 50, 3], strides=[1], axes=None)\ + /* ty=Tensor[(1, 10, 50, 3), float32] */; + %4 = strided_slice(%0, begin=[0, 30, 0, 0], end=[1, 40, 50, 3], strides=[1], axes=None)\ + /* ty=Tensor[(1, 10, 50, 3), float32] */; + %5 = strided_slice(%0, begin=[0, 40, 0, 0], end=[1, 50, 50, 3], strides=[1], axes=None)\ + /* ty=Tensor[(1, 10, 50, 3), float32] */; + %6 = (%1, %2, %3, %4, %5); + %7 = %6.0; + %8 = tanh(%7) /* ty=Tensor[(1, 10, 50, 3), float32] */; + %9 = %6.1; + %10 = tanh(%9) /* ty=Tensor[(1, 10, 50, 3), float32] */; + %11 = %6.2; + %12 = tanh(%11) /* ty=Tensor[(1, 10, 50, 3), float32] */; + %13 = %6.3; + %14 = tanh(%13) /* ty=Tensor[(1, 10, 50, 3), float32] */; + %15 = %6.4; + %16 = tanh(%15) /* ty=Tensor[(1, 10, 50, 3), float32] */; + (%8, %10, %12, %14, %16) + } + """ + return tvm.parser.fromtext(expected_ir_string) + + def expected_mod_axis2(): + expected_ir_string = """ + #[version = "0.0.5"] + def @tvmgen_default_ethosu_main_0(%x: Tensor[(1, 50, 50, 3), float32]) -> (Tensor[(1, 50, 10, 3), float32],\ + Tensor[(1, 50, 10, 3), float32],\ + Tensor[(1, 50, 10, 3), float32],\ + Tensor[(1, 50, 10, 3), float32],\ + Tensor[(1, 50, 10, 3), float32]) { + %0 = abs(%x) /* ty=Tensor[(1, 50, 50, 3), float32] */; + %1 = strided_slice(%0, begin=[0, 0, 0, 0], end=[1, 50, 10, 3], strides=[1], axes=None)\ + /* ty=Tensor[(1, 50, 10, 3), float32] */; + %2 = strided_slice(%0, begin=[0, 0, 10, 0], end=[1, 50, 20, 3], strides=[1], axes=None)\ + /* ty=Tensor[(1, 50, 10, 3), float32] */; + %3 = strided_slice(%0, begin=[0, 0, 20, 0], end=[1, 50, 30, 3], strides=[1], axes=None)\ + /* ty=Tensor[(1, 50, 10, 3), float32] */; + %4 = strided_slice(%0, begin=[0, 0, 30, 0], end=[1, 50, 40, 3], strides=[1], axes=None)\ + /* ty=Tensor[(1, 50, 10, 3), float32] */; + %5 = strided_slice(%0, begin=[0, 0, 40, 0], end=[1, 50, 50, 3], strides=[1], axes=None)\ + /* ty=Tensor[(1, 50, 10, 3), float32] */; + %6 = (%1, %2, %3, %4, %5); + %7 = %6.0; + %8 = tanh(%7) /* ty=Tensor[(1, 50, 10, 3), float32] */; + %9 = %6.1; + %10 = tanh(%9) /* ty=Tensor[(1, 50, 10, 3), float32] */; + %11 = %6.2; + %12 = tanh(%11) /* ty=Tensor[(1, 50, 10, 3), float32] */; + %13 = %6.3; + %14 = tanh(%13) /* ty=Tensor[(1, 50, 10, 3), float32] */; + %15 = %6.4; + %16 = tanh(%15) /* ty=Tensor[(1, 50, 10, 3), float32] */; + (%8, %10, %12, %14, %16) + } + """ + return tvm.parser.fromtext(expected_ir_string) + + mod_axis1 = tvm.IRModule() + mod_axis1["tvmgen_default_ethosu_main_0"] = create_graph(1, 5) + mod_axis1 = legalize.LegalizeSplit()(mod_axis1) + expected_axis1 = expected_mod_axis1() + tvm.ir.assert_structural_equal(mod_axis1, expected_axis1) + + mod_axis2 = tvm.IRModule() + mod_axis2["tvmgen_default_ethosu_main_0"] = create_graph(2, 5) + mod_axis2 = legalize.LegalizeSplit()(mod_axis2) + expected_axis2 = expected_mod_axis2() + tvm.ir.assert_structural_equal(mod_axis2, expected_axis2) + + +def infer_type_function_pass(func): + mod = tvm.IRModule() + mod["test"] = func + mod = relay.transform.InferType()(mod) + return mod["test"] + + +def get_shape_expr(in_expr, out_expr): + main_f = relay.Function([in_expr], out_expr) + main_f = infer_type_function_pass(main_f) + shape = [int(i) for i in main_f.body.checked_type.shape] + return shape + + +INVERSE_LAYOUT_TRANSFORM_OHWI_MAP = { + "HWIO": [1, 2, 3, 0], + "HWOI": [1, 2, 0, 3], + "OWHI": [0, 1, 2, 3], +} + + +def test_ethosu_conv2d_legalize(): + def create_graph_single(input_tensor_name, input_tensor_shape, input_tensor_dtype): + c1_params = relay_ir_builder.QnnConv2DParams(input_tensor_dtype) + c1_params.ifm.shape = input_tensor_shape + c1_params.kernel.shape = (3, 3, c1_params.ifm.shape[3], 32) + c1_params.strides = (1, 1) + c1_params.pad = "VALID" + c1_params.activation = "CLIP" + c1_params.clip_min = 23 + c1_params.clip_max = 180 + input0 = relay.var(input_tensor_name, shape=c1_params.ifm.shape, dtype=c1_params.ifm.dtype) + c1, new_params = relay_ir_builder.create_qnn_conv2d(c1_params, input0) + c1_params.ofm.shape = get_shape_expr(input0, c1) + + f = relay.Function([input0], c1) + mod = tvm.IRModule() + mod["main"] = f + return mod, [c1_params] + + def create_graph_double(input_tensor_name, input_tensor_shape, input_tensor_dtype): + c1_params = relay_ir_builder.QnnConv2DParams(input_tensor_dtype) + c1_params.ifm.shape = input_tensor_shape + c1_params.kernel.shape = (7, 7, c1_params.ifm.shape[3], 8) + c1_params.strides = (2, 2) + c1_params.pad = "VALID" + c1_params.activation = "CLIP" + c1_params.clip_min = 10 + c1_params.clip_max = 240 + input0 = relay.var(input_tensor_name, shape=c1_params.ifm.shape, dtype=c1_params.ifm.dtype) + c1, new_params = relay_ir_builder.create_qnn_conv2d(c1_params, input0) + c1_params.ofm.shape = get_shape_expr(input0, c1) + + c2_params = relay_ir_builder.QnnConv2DParams(input_tensor_dtype) + c2_params.ifm.shape = c1_params.ofm.shape + c2_params.kernel.shape = (5, 5, c2_params.ifm.shape[3], 16) + c2_params.strides = (1, 1) + c2_params.pad = "SAME" + c2, new_params = relay_ir_builder.create_qnn_conv2d(c2_params, c1) + c2_params.ofm.shape = get_shape_expr(input0, c2) + + f = relay.Function([input0], c2) + mod = tvm.IRModule() + mod["main"] = f + return mod, [c2_params, c1_params] + + def verify_tensor(tensor_type, expr): + assert list(tensor_type.shape) == list(expr.checked_type.shape) + assert str(tensor_type.dtype) == str(expr.checked_type.dtype) + + def verify_linear(ext_func, conv2d_params): + op = ext_func.body + for param in conv2d_params: + verify_tensor(param.ifm, op.args[0]) + verify_tensor(param.ofm, op) + + # This will be in OHWI layout + weights_ohwi = op.args[1].data.asnumpy() + weights_layout = str(param.kernel.layout) + weights = np.transpose(weights_ohwi, INVERSE_LAYOUT_TRANSFORM_OHWI_MAP[weights_layout]) + assert weights.shape == param.kernel.shape + assert weights.dtype == param.kernel.dtype + + assert list(op.args[2].checked_type.shape)[0] == weights_ohwi.shape[0] + + assert float(op.attrs.ifm_scale) == float(param.ifm.sc.data.asnumpy()) + assert int(op.attrs.ifm_zero_point) == int(param.ifm.zp.data.asnumpy()) + assert int(op.attrs.weight_zero_point) == int(param.kernel.zp.data.asnumpy()) + assert float(op.attrs.ofm_scale) == float(param.ofm.sc.data.asnumpy()) + assert int(op.attrs.ofm_zero_point) == int(param.ofm.zp.data.asnumpy()) + assert int(op.attrs.ofm_channels) == int(weights_ohwi.shape[0]) + assert list(op.attrs.padding) == list(param.pad) + assert list(op.attrs.strides) == list(param.strides) + assert list(op.attrs.dilation) == list(param.dilation) + assert str(op.attrs.activation) == str(param.activation) + assert int(op.attrs.clip_min) == int(param.clip_min) + assert int(op.attrs.clip_max) == int(param.clip_max) + op = op.args[0] + + test_cases = [ + (create_graph_single, ["input", (1, 299, 299, 3), "uint8"]), + (create_graph_double, ["input", (1, 128, 256, 4), "uint8"]), + ] + for test_case in test_cases: + mod, conv_params = test_case[0](*test_case[1]) + mod = ethosu.partition_for_ethosu(mod) + mod = legalize.LegalizeEthosUConv2D()(mod) + verify_linear(mod["tvmgen_default_ethosu_main_0"], conv_params) + + +def test_ethosu_conv2d_legalize_errors(): + def create_graph_single_unsupported_ifm_layout( + input_tensor_name, input_tensor_shape, input_tensor_dtype + ): + c1_params = relay_ir_builder.QnnConv2DParams(input_tensor_dtype) + c1_params.ifm.shape = input_tensor_shape + c1_params.ifm.layout = "NCHW" + c1_params.kernel.shape = (3, 3, c1_params.ifm.shape[1], 32) + c1_params.strides = (1, 1) + c1_params.pad = "VALID" + c1_params.activation = "CLIP" + c1_params.clip_min = 23 + c1_params.clip_max = 180 + input0 = relay.var(input_tensor_name, shape=c1_params.ifm.shape, dtype=c1_params.ifm.dtype) + c1, new_params = relay_ir_builder.create_qnn_conv2d(c1_params, input0) + c1_params.ofm.shape = get_shape_expr(input0, c1) + + f = relay.Function([input0], c1) + mod = tvm.IRModule() + mod["main"] = f + return mod, [c1_params] + + test_cases = [ + (create_graph_single_unsupported_ifm_layout, ["input", (1, 3, 299, 299), "uint8"]), + ] + + for test_case in test_cases: + mod, conv_params = test_case[0](*test_case[1]) + mod = ethosu.partition_for_ethosu(mod) + try: + mod = legalize.LegalizeEthosUConv2D()(mod) + except Exception as e: + assert "EthosUCodegenError: Unsupported Layout NCHW" in e.args[0] Review comment: This is missing the pytest main call: ``` if __name__ == "__main__": sys.exit(pytest.main([__file__] + sys.argv[1:])) ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@tvm.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org
