mbaret commented on a change in pull request #9457: URL: https://github.com/apache/tvm/pull/9457#discussion_r743789916
########## File path: python/tvm/relay/backend/contrib/ethosu/op/identity.py ########## @@ -0,0 +1,98 @@ +# 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=unused-argument +"""Relay identity operator""" + +import tvm +from tvm.relay.op import _make +from tvm.topi.generic import schedule_injective +from tvm.relay.op.op import OpStrategy +from tvm.relay.op import strategy as _strategy + +from ..te import identity_compute + + [email protected]_op_attr("contrib.ethosu.identity", "FTVMCompute") +def create_ethosu_identity_compute(attrs, args, out_type): + """Create an ethosu_identity compute op.""" + ifm = args[0] + lut = args[1] + ifm_scale = attrs.ifm_scale + ifm_zero_point = attrs.ifm_zero_point + ofm_scale = attrs.ofm_scale + ofm_zero_point = attrs.ofm_zero_point + activation = attrs.activation + op = identity_compute( + ifm, lut, ifm_scale, ifm_zero_point, ofm_scale, ofm_zero_point, activation + ) + return [op] + + [email protected]_op_attr("contrib.ethosu.identity", "FTVMStrategy") +def pooling_strategy_ethosu(attrs, inputs, out_type, target): Review comment: identity_strategy_ethosu ########## File path: python/tvm/relay/backend/contrib/ethosu/tir/identity.py ########## @@ -0,0 +1,87 @@ +# 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 +"""Extract information from the identity operator in TIR.""" +from typing import Dict, Tuple +import tvm +from .dma import get_read_params, get_write_params +from .spec import SerialKernel, SerialActivation, SerialPooling, SerialPadding +from .utils import get_op_attrs + + +def get_identity_params( + stmt: tvm.tir.AttrStmt, + producers: Dict[tvm.tir.Var, tvm.tir.AttrStmt], + consumers: Dict[tvm.tir.Var, tvm.tir.AttrStmt], +) -> Tuple[SerialPooling, tvm.tir.Var, tvm.tir.Var]: + """Get the parameters necessary to construct a call_extern for a pooling. + + Parameters + ---------- + stmt : tvm.tir.AttrStmt + The outermost attribute statement of a convolution loop nest. + producers : Dict[tvm.tir.Var, tvm.tir.AttrStmt] + A dictionary to associate pointers with the loop nest + that produces their values. + consumers : Dict[tvm.tir.Var, tvm.tir.AttrStmt] + A dictionary to associate pointers with the loop nest + that consumes their values. + + Returns + ------- + SerialPooling + The parameters needed to construct a 2D pooling. + output_pointer : tvm.tir.Var + The output pointer of the pooling operation. + replace_pointer : tvm.tir.Var + The output pointer of the DMA write operation, which is to replace + the pooling output pointer. + + """ + attrs, _ = get_op_attrs(stmt) + # Find the inner loop + while hasattr(stmt, "body"): + stmt = stmt.body + + input_pointer = stmt.value.buffer_var + output_pointer = stmt.buffer_var + + read = producers[input_pointer] + write = consumers[output_pointer] + + serial_ifm, _, _ = get_read_params(read) + serial_ofm, _, write_output_pointer = get_write_params(write) + + replace_pointer = write_output_pointer + + # TODO (maybe): Support stand alone RELU through clamping in identity Review comment: Clarify this ########## File path: tests/python/contrib/test_ethosu/test_type_inference.py ########## @@ -171,5 +180,26 @@ def test_ethosu_pooling_invalid_dtype(): run_opt_pass(func, relay.transform.InferType()) +def test_ethosu_identity_invalid_shape(): + invalid_shape = [1, 2, 3, 4, 5] + dtype = "int8" + ifm = relay.var("ifm", shape=invalid_shape, dtype=dtype) + + identity = make_ethosu_identity(ifm) + func = relay.Function([ifm], identity) + with pytest.raises(TVMError): + run_opt_pass(func, relay.transform.InferType()) + + +def test_ethosu_invalid_invalid_dtype(): Review comment: test_ethosu_identity_invalid_dtype ########## File path: python/tvm/relay/backend/contrib/ethosu/legalize.py ########## @@ -423,11 +526,15 @@ class LegalizeEthosU: def transform_module( self, mod: tvm.ir.IRModule, ctx: tvm.ir.transform.PassContext ) -> tvm.ir.IRModule: + """Legalize the oerators that can be offloaded to the NPU""" Review comment: operators ########## File path: tests/python/contrib/test_ethosu/test_type_inference.py ########## @@ -171,5 +180,26 @@ def test_ethosu_pooling_invalid_dtype(): run_opt_pass(func, relay.transform.InferType()) +def test_ethosu_identity_invalid_shape(): Review comment: I think we need to test the valid cases as well. ########## File path: python/tvm/relay/op/contrib/ethosu.py ########## @@ -481,6 +493,8 @@ def pattern_table() -> List[Tuple[str, tvm.relay.dataflow_pattern.DFPattern, Cal qnn_avgpool2d_pattern(), lambda pat: AvgPool2DParams(pat).is_valid(), ), + ("ethosu.strided_slice", strided_slice_pattern(), lambda pat: True), Review comment: Let's put the proper restrictions here so we don't accidentally offload, for example, floating point tensors. ########## File path: src/relay/op/contrib/ethosu/identity.cc ########## @@ -0,0 +1,128 @@ +/* + * 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/identity.cc + * \brief Property def of the Arm Ethos-U identity op. + */ +#include <tvm/relay/op.h> + +#include "common.h" + +namespace tvm { +namespace relay { +namespace op { +namespace contrib { +namespace ethosu { + +/*! \brief Attributes used by the Ethos(TM)-U NPU identity operator */ +struct EthosuIdentityAttrs : public tvm::AttrsNode<EthosuIdentityAttrs> { + double ifm_scale; + int ifm_zero_point; + double ofm_scale; + int ofm_zero_point; + String activation; + + TVM_DECLARE_ATTRS(EthosuIdentityAttrs, "relay.attrs.EthosuIdentityAttrs") { + 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(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(activation) + .describe( + "The activation function to use. " + "'NONE' - no activation function. " + "'TANH' - tanh activation function. " + "'SIGMOID' - sigmoid activation function. " + "'LUT' - use a look-up table to perform the activation function.") + .set_default("NONE"); + } +}; + +TVM_REGISTER_NODE_TYPE(EthosuIdentityAttrs); + +bool EthosuIdentityRel(const Array<Type>& types, int num_inputs, const Attrs& attrs, + const TypeReporter& reporter) { + int ifm_index = 0; + int result_index = 2; + ICHECK_EQ(types.size(), result_index + 1); + + const auto* ifm = types[ifm_index].as<TensorTypeNode>(); + if (ifm == nullptr) return false; + + const auto* param = attrs.as<EthosuIdentityAttrs>(); + + ICHECK(param != nullptr) << "EthosuIdentityAttrs cannot be nullptr."; + + if (ifm->dtype != DataType::UInt(8) && ifm->dtype != DataType::Int(8)) { + reporter->GetDiagCtx().EmitFatal( + Diagnostic::Error(reporter->GetSpan()) + << "Invalid operator: Expected type(uint8) or type(int8) for ifm but was " << ifm->dtype); + return false; + } + + if (ifm->shape.size() > 4) { + reporter->GetDiagCtx().EmitFatal( + Diagnostic::Error(reporter->GetSpan()) + << "Invalid operator: Input Feature Map should be at most 4 dimensional, but was " + << ifm->shape); + return false; + } + + // Assign ofm type + auto ofm_shape = ifm->shape; + reporter->Assign(types[result_index], TensorType(ofm_shape, ifm->dtype)); + return true; +} + +Expr MakeEthosuIdentity(Expr ifm, Expr lut, double ifm_scale, int ifm_zero_point, double ofm_scale, + int ofm_zero_point, String activation) { + auto attrs = make_object<EthosuIdentityAttrs>(); + attrs->ifm_scale = ifm_scale; + attrs->ifm_zero_point = ifm_zero_point; + attrs->ofm_scale = ofm_scale; + attrs->ofm_zero_point = ofm_zero_point; + attrs->activation = std::move(activation); + static const Op& op = Op::Get("contrib.ethosu.identity"); + return Call(op, {ifm, lut}, Attrs(attrs), {}); +} + +TVM_REGISTER_GLOBAL("relay.op._make.ethosu_identity").set_body_typed(MakeEthosuIdentity); + +RELAY_REGISTER_OP("contrib.ethosu.identity") + .describe(R"code(Identity operator for Ethos-U NPUs. + +This Relay operator performs the identity operation on Ethos(TM)-U NPU with a capability +to requantize the data. It accepts input with any shape that is less or equal to 4. Review comment: Maybe 'It accepts input tensors of 4 dimensions or less.'? ########## File path: python/tvm/relay/backend/contrib/ethosu/tir/identity.py ########## @@ -0,0 +1,87 @@ +# 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 +"""Extract information from the identity operator in TIR.""" +from typing import Dict, Tuple +import tvm +from .dma import get_read_params, get_write_params +from .spec import SerialKernel, SerialActivation, SerialPooling, SerialPadding +from .utils import get_op_attrs + + +def get_identity_params( + stmt: tvm.tir.AttrStmt, + producers: Dict[tvm.tir.Var, tvm.tir.AttrStmt], + consumers: Dict[tvm.tir.Var, tvm.tir.AttrStmt], +) -> Tuple[SerialPooling, tvm.tir.Var, tvm.tir.Var]: + """Get the parameters necessary to construct a call_extern for a pooling. + + Parameters + ---------- + stmt : tvm.tir.AttrStmt + The outermost attribute statement of a convolution loop nest. Review comment: identity pooling loop nest ########## File path: python/tvm/relay/backend/contrib/ethosu/tir/identity.py ########## @@ -0,0 +1,87 @@ +# 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 +"""Extract information from the identity operator in TIR.""" +from typing import Dict, Tuple +import tvm +from .dma import get_read_params, get_write_params +from .spec import SerialKernel, SerialActivation, SerialPooling, SerialPadding +from .utils import get_op_attrs + + +def get_identity_params( + stmt: tvm.tir.AttrStmt, + producers: Dict[tvm.tir.Var, tvm.tir.AttrStmt], + consumers: Dict[tvm.tir.Var, tvm.tir.AttrStmt], +) -> Tuple[SerialPooling, tvm.tir.Var, tvm.tir.Var]: + """Get the parameters necessary to construct a call_extern for a pooling. Review comment: for an identity pooling ########## File path: tests/python/contrib/test_ethosu/test_codegen.py ########## @@ -343,5 +343,178 @@ def representative_dataset(): infra.verify_source(compiled_models, accel_type) [email protected]("accel_type", ACCEL_TYPES) [email protected]("ifm_shape", [(3, 2), (1, 15, 11, 7), (3, 1, 12), (400,)]) [email protected]("ifm_scale, ifm_zp, ofm_scale, ofm_zp", [(1, 0, 1, 0), (0.015, 3, 0.2, 5)]) +def test_ethosu_identity_codegen(ifm_shape, ifm_scale, ifm_zp, ofm_scale, ofm_zp, accel_type): + # Create a "partitioned" Relay function + ifm = relay.var("ifm", shape=ifm_shape, dtype="int8") + ifm0 = relay.var("ifm0", shape=ifm_shape, dtype="int8") + identity = infra.make_ethosu_identity( + ifm0, ifm_scale=ifm_scale, ifm_zero_point=ifm_zp, ofm_scale=ofm_scale, ofm_zero_point=ofm_zp + ) + glb_ethosu = relay.GlobalVar("tvmgen_default_ethosu_main_0") + + func = ( + relay.Function([ifm0], identity) + .with_attr("Inline", 1) + .with_attr("Compiler", "ethosu") + .with_attr("global_symbol", "tvmgen_default_ethosu_main_0") + .with_attr("Primitive", 1) + ) + mod = tvm.IRModule() + mod[glb_ethosu] = func + mod = relay.transform.InferType()(mod) + + call = relay.Call(glb_ethosu, [ifm]) + mod["main"] = relay.Function([ifm], call) + mod = relay.transform.InferType()(mod) + + in_data = np.random.randint(-120, high=120, size=ifm_shape, dtype="int8") + requant_data = (ifm_scale * (in_data - ifm_zp)) / ofm_scale + ofm_zp + out_data = np.round(np.clip(requant_data, -128, 127)).astype("int8") + + compiled_model = infra.build_source( + mod, {"ifm": in_data}, [out_data], accel_type, output_tolerance=1 + ) + + imported_modules = compiled_model[0].executor_factory.lib.imported_modules + assert len(imported_modules) == 2 + ethosu_module = imported_modules[0] + + # Verify generated C source + get_cs = tvm._ffi.get_global_func("runtime.module.ethosu.getcs") + cmms = get_cs(ethosu_module) + cmms = bytes.fromhex(cmms) + + infra.print_payload(cmms) + infra.verify_source(compiled_model, accel_type) + + [email protected]("accel_type", ACCEL_TYPES) [email protected]( + "ifm_shape, new_shape", + [ + ((1, 4, 1, 2), (1, 1, 1, 8)), + ( + ( + 5, + 1, + 20, + ), + (1, 5, 1, 20), + ), + ((12, 20), (1, 6, 4, 10)), + ((12, 20), (6, 4, 10)), + ((20,), (4, 5)), + ], +) +def test_relay_reshape_codegen(ifm_shape, new_shape, accel_type): + # Create a "partitioned" Relay graph Review comment: It seems like some of this could be refactored out into a common function and reused between tests. ########## File path: tests/python/contrib/test_ethosu/test_codegen.py ########## @@ -343,5 +343,178 @@ def representative_dataset(): infra.verify_source(compiled_models, accel_type) [email protected]("accel_type", ACCEL_TYPES) [email protected]("ifm_shape", [(3, 2), (1, 15, 11, 7), (3, 1, 12), (400,)]) [email protected]("ifm_scale, ifm_zp, ofm_scale, ofm_zp", [(1, 0, 1, 0), (0.015, 3, 0.2, 5)]) +def test_ethosu_identity_codegen(ifm_shape, ifm_scale, ifm_zp, ofm_scale, ofm_zp, accel_type): + # Create a "partitioned" Relay function + ifm = relay.var("ifm", shape=ifm_shape, dtype="int8") + ifm0 = relay.var("ifm0", shape=ifm_shape, dtype="int8") + identity = infra.make_ethosu_identity( + ifm0, ifm_scale=ifm_scale, ifm_zero_point=ifm_zp, ofm_scale=ofm_scale, ofm_zero_point=ofm_zp + ) + glb_ethosu = relay.GlobalVar("tvmgen_default_ethosu_main_0") + + func = ( + relay.Function([ifm0], identity) + .with_attr("Inline", 1) + .with_attr("Compiler", "ethosu") + .with_attr("global_symbol", "tvmgen_default_ethosu_main_0") + .with_attr("Primitive", 1) + ) + mod = tvm.IRModule() + mod[glb_ethosu] = func + mod = relay.transform.InferType()(mod) + + call = relay.Call(glb_ethosu, [ifm]) + mod["main"] = relay.Function([ifm], call) + mod = relay.transform.InferType()(mod) + + in_data = np.random.randint(-120, high=120, size=ifm_shape, dtype="int8") + requant_data = (ifm_scale * (in_data - ifm_zp)) / ofm_scale + ofm_zp + out_data = np.round(np.clip(requant_data, -128, 127)).astype("int8") + + compiled_model = infra.build_source( + mod, {"ifm": in_data}, [out_data], accel_type, output_tolerance=1 + ) + + imported_modules = compiled_model[0].executor_factory.lib.imported_modules + assert len(imported_modules) == 2 + ethosu_module = imported_modules[0] + + # Verify generated C source + get_cs = tvm._ffi.get_global_func("runtime.module.ethosu.getcs") + cmms = get_cs(ethosu_module) + cmms = bytes.fromhex(cmms) + + infra.print_payload(cmms) + infra.verify_source(compiled_model, accel_type) + + [email protected]("accel_type", ACCEL_TYPES) [email protected]( Review comment: We should either test the 'special' indices for reshape (-1 and -2), or we should explicitly disallow these. ########## File path: python/tvm/relay/backend/contrib/ethosu/op/identity.py ########## @@ -0,0 +1,98 @@ +# 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=unused-argument +"""Relay identity operator""" + +import tvm +from tvm.relay.op import _make +from tvm.topi.generic import schedule_injective +from tvm.relay.op.op import OpStrategy +from tvm.relay.op import strategy as _strategy + +from ..te import identity_compute + + [email protected]_op_attr("contrib.ethosu.identity", "FTVMCompute") +def create_ethosu_identity_compute(attrs, args, out_type): + """Create an ethosu_identity compute op.""" + ifm = args[0] + lut = args[1] + ifm_scale = attrs.ifm_scale + ifm_zero_point = attrs.ifm_zero_point + ofm_scale = attrs.ofm_scale + ofm_zero_point = attrs.ofm_zero_point + activation = attrs.activation + op = identity_compute( + ifm, lut, ifm_scale, ifm_zero_point, ofm_scale, ofm_zero_point, activation + ) + return [op] + + [email protected]_op_attr("contrib.ethosu.identity", "FTVMStrategy") +def pooling_strategy_ethosu(attrs, inputs, out_type, target): Review comment: identity_strategy_ethosu ########## File path: python/tvm/relay/backend/contrib/ethosu/tir/identity.py ########## @@ -0,0 +1,87 @@ +# 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 +"""Extract information from the identity operator in TIR.""" +from typing import Dict, Tuple +import tvm +from .dma import get_read_params, get_write_params +from .spec import SerialKernel, SerialActivation, SerialPooling, SerialPadding +from .utils import get_op_attrs + + +def get_identity_params( + stmt: tvm.tir.AttrStmt, + producers: Dict[tvm.tir.Var, tvm.tir.AttrStmt], + consumers: Dict[tvm.tir.Var, tvm.tir.AttrStmt], +) -> Tuple[SerialPooling, tvm.tir.Var, tvm.tir.Var]: + """Get the parameters necessary to construct a call_extern for a pooling. + + Parameters + ---------- + stmt : tvm.tir.AttrStmt + The outermost attribute statement of a convolution loop nest. + producers : Dict[tvm.tir.Var, tvm.tir.AttrStmt] + A dictionary to associate pointers with the loop nest + that produces their values. + consumers : Dict[tvm.tir.Var, tvm.tir.AttrStmt] + A dictionary to associate pointers with the loop nest + that consumes their values. + + Returns + ------- + SerialPooling + The parameters needed to construct a 2D pooling. + output_pointer : tvm.tir.Var + The output pointer of the pooling operation. + replace_pointer : tvm.tir.Var + The output pointer of the DMA write operation, which is to replace + the pooling output pointer. + + """ + attrs, _ = get_op_attrs(stmt) + # Find the inner loop + while hasattr(stmt, "body"): + stmt = stmt.body + + input_pointer = stmt.value.buffer_var + output_pointer = stmt.buffer_var + + read = producers[input_pointer] + write = consumers[output_pointer] + + serial_ifm, _, _ = get_read_params(read) + serial_ofm, _, write_output_pointer = get_write_params(write) + + replace_pointer = write_output_pointer + + # TODO (maybe): Support stand alone RELU through clamping in identity Review comment: Clarify this ########## File path: tests/python/contrib/test_ethosu/test_type_inference.py ########## @@ -171,5 +180,26 @@ def test_ethosu_pooling_invalid_dtype(): run_opt_pass(func, relay.transform.InferType()) +def test_ethosu_identity_invalid_shape(): + invalid_shape = [1, 2, 3, 4, 5] + dtype = "int8" + ifm = relay.var("ifm", shape=invalid_shape, dtype=dtype) + + identity = make_ethosu_identity(ifm) + func = relay.Function([ifm], identity) + with pytest.raises(TVMError): + run_opt_pass(func, relay.transform.InferType()) + + +def test_ethosu_invalid_invalid_dtype(): Review comment: test_ethosu_identity_invalid_dtype ########## File path: python/tvm/relay/backend/contrib/ethosu/legalize.py ########## @@ -423,11 +526,15 @@ class LegalizeEthosU: def transform_module( self, mod: tvm.ir.IRModule, ctx: tvm.ir.transform.PassContext ) -> tvm.ir.IRModule: + """Legalize the oerators that can be offloaded to the NPU""" Review comment: operators ########## File path: tests/python/contrib/test_ethosu/test_type_inference.py ########## @@ -171,5 +180,26 @@ def test_ethosu_pooling_invalid_dtype(): run_opt_pass(func, relay.transform.InferType()) +def test_ethosu_identity_invalid_shape(): Review comment: I think we need to test the valid cases as well. ########## File path: python/tvm/relay/op/contrib/ethosu.py ########## @@ -481,6 +493,8 @@ def pattern_table() -> List[Tuple[str, tvm.relay.dataflow_pattern.DFPattern, Cal qnn_avgpool2d_pattern(), lambda pat: AvgPool2DParams(pat).is_valid(), ), + ("ethosu.strided_slice", strided_slice_pattern(), lambda pat: True), Review comment: Let's put the proper restrictions here so we don't accidentally offload, for example, floating point tensors. ########## File path: src/relay/op/contrib/ethosu/identity.cc ########## @@ -0,0 +1,128 @@ +/* + * 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/identity.cc + * \brief Property def of the Arm Ethos-U identity op. + */ +#include <tvm/relay/op.h> + +#include "common.h" + +namespace tvm { +namespace relay { +namespace op { +namespace contrib { +namespace ethosu { + +/*! \brief Attributes used by the Ethos(TM)-U NPU identity operator */ +struct EthosuIdentityAttrs : public tvm::AttrsNode<EthosuIdentityAttrs> { + double ifm_scale; + int ifm_zero_point; + double ofm_scale; + int ofm_zero_point; + String activation; + + TVM_DECLARE_ATTRS(EthosuIdentityAttrs, "relay.attrs.EthosuIdentityAttrs") { + 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(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(activation) + .describe( + "The activation function to use. " + "'NONE' - no activation function. " + "'TANH' - tanh activation function. " + "'SIGMOID' - sigmoid activation function. " + "'LUT' - use a look-up table to perform the activation function.") + .set_default("NONE"); + } +}; + +TVM_REGISTER_NODE_TYPE(EthosuIdentityAttrs); + +bool EthosuIdentityRel(const Array<Type>& types, int num_inputs, const Attrs& attrs, + const TypeReporter& reporter) { + int ifm_index = 0; + int result_index = 2; + ICHECK_EQ(types.size(), result_index + 1); + + const auto* ifm = types[ifm_index].as<TensorTypeNode>(); + if (ifm == nullptr) return false; + + const auto* param = attrs.as<EthosuIdentityAttrs>(); + + ICHECK(param != nullptr) << "EthosuIdentityAttrs cannot be nullptr."; + + if (ifm->dtype != DataType::UInt(8) && ifm->dtype != DataType::Int(8)) { + reporter->GetDiagCtx().EmitFatal( + Diagnostic::Error(reporter->GetSpan()) + << "Invalid operator: Expected type(uint8) or type(int8) for ifm but was " << ifm->dtype); + return false; + } + + if (ifm->shape.size() > 4) { + reporter->GetDiagCtx().EmitFatal( + Diagnostic::Error(reporter->GetSpan()) + << "Invalid operator: Input Feature Map should be at most 4 dimensional, but was " + << ifm->shape); + return false; + } + + // Assign ofm type + auto ofm_shape = ifm->shape; + reporter->Assign(types[result_index], TensorType(ofm_shape, ifm->dtype)); + return true; +} + +Expr MakeEthosuIdentity(Expr ifm, Expr lut, double ifm_scale, int ifm_zero_point, double ofm_scale, + int ofm_zero_point, String activation) { + auto attrs = make_object<EthosuIdentityAttrs>(); + attrs->ifm_scale = ifm_scale; + attrs->ifm_zero_point = ifm_zero_point; + attrs->ofm_scale = ofm_scale; + attrs->ofm_zero_point = ofm_zero_point; + attrs->activation = std::move(activation); + static const Op& op = Op::Get("contrib.ethosu.identity"); + return Call(op, {ifm, lut}, Attrs(attrs), {}); +} + +TVM_REGISTER_GLOBAL("relay.op._make.ethosu_identity").set_body_typed(MakeEthosuIdentity); + +RELAY_REGISTER_OP("contrib.ethosu.identity") + .describe(R"code(Identity operator for Ethos-U NPUs. + +This Relay operator performs the identity operation on Ethos(TM)-U NPU with a capability +to requantize the data. It accepts input with any shape that is less or equal to 4. Review comment: Maybe 'It accepts input tensors of 4 dimensions or less.'? ########## File path: python/tvm/relay/backend/contrib/ethosu/tir/identity.py ########## @@ -0,0 +1,87 @@ +# 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 +"""Extract information from the identity operator in TIR.""" +from typing import Dict, Tuple +import tvm +from .dma import get_read_params, get_write_params +from .spec import SerialKernel, SerialActivation, SerialPooling, SerialPadding +from .utils import get_op_attrs + + +def get_identity_params( + stmt: tvm.tir.AttrStmt, + producers: Dict[tvm.tir.Var, tvm.tir.AttrStmt], + consumers: Dict[tvm.tir.Var, tvm.tir.AttrStmt], +) -> Tuple[SerialPooling, tvm.tir.Var, tvm.tir.Var]: + """Get the parameters necessary to construct a call_extern for a pooling. + + Parameters + ---------- + stmt : tvm.tir.AttrStmt + The outermost attribute statement of a convolution loop nest. Review comment: identity pooling loop nest ########## File path: python/tvm/relay/backend/contrib/ethosu/tir/identity.py ########## @@ -0,0 +1,87 @@ +# 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 +"""Extract information from the identity operator in TIR.""" +from typing import Dict, Tuple +import tvm +from .dma import get_read_params, get_write_params +from .spec import SerialKernel, SerialActivation, SerialPooling, SerialPadding +from .utils import get_op_attrs + + +def get_identity_params( + stmt: tvm.tir.AttrStmt, + producers: Dict[tvm.tir.Var, tvm.tir.AttrStmt], + consumers: Dict[tvm.tir.Var, tvm.tir.AttrStmt], +) -> Tuple[SerialPooling, tvm.tir.Var, tvm.tir.Var]: + """Get the parameters necessary to construct a call_extern for a pooling. Review comment: for an identity pooling ########## File path: tests/python/contrib/test_ethosu/test_codegen.py ########## @@ -343,5 +343,178 @@ def representative_dataset(): infra.verify_source(compiled_models, accel_type) [email protected]("accel_type", ACCEL_TYPES) [email protected]("ifm_shape", [(3, 2), (1, 15, 11, 7), (3, 1, 12), (400,)]) [email protected]("ifm_scale, ifm_zp, ofm_scale, ofm_zp", [(1, 0, 1, 0), (0.015, 3, 0.2, 5)]) +def test_ethosu_identity_codegen(ifm_shape, ifm_scale, ifm_zp, ofm_scale, ofm_zp, accel_type): + # Create a "partitioned" Relay function + ifm = relay.var("ifm", shape=ifm_shape, dtype="int8") + ifm0 = relay.var("ifm0", shape=ifm_shape, dtype="int8") + identity = infra.make_ethosu_identity( + ifm0, ifm_scale=ifm_scale, ifm_zero_point=ifm_zp, ofm_scale=ofm_scale, ofm_zero_point=ofm_zp + ) + glb_ethosu = relay.GlobalVar("tvmgen_default_ethosu_main_0") + + func = ( + relay.Function([ifm0], identity) + .with_attr("Inline", 1) + .with_attr("Compiler", "ethosu") + .with_attr("global_symbol", "tvmgen_default_ethosu_main_0") + .with_attr("Primitive", 1) + ) + mod = tvm.IRModule() + mod[glb_ethosu] = func + mod = relay.transform.InferType()(mod) + + call = relay.Call(glb_ethosu, [ifm]) + mod["main"] = relay.Function([ifm], call) + mod = relay.transform.InferType()(mod) + + in_data = np.random.randint(-120, high=120, size=ifm_shape, dtype="int8") + requant_data = (ifm_scale * (in_data - ifm_zp)) / ofm_scale + ofm_zp + out_data = np.round(np.clip(requant_data, -128, 127)).astype("int8") + + compiled_model = infra.build_source( + mod, {"ifm": in_data}, [out_data], accel_type, output_tolerance=1 + ) + + imported_modules = compiled_model[0].executor_factory.lib.imported_modules + assert len(imported_modules) == 2 + ethosu_module = imported_modules[0] + + # Verify generated C source + get_cs = tvm._ffi.get_global_func("runtime.module.ethosu.getcs") + cmms = get_cs(ethosu_module) + cmms = bytes.fromhex(cmms) + + infra.print_payload(cmms) + infra.verify_source(compiled_model, accel_type) + + [email protected]("accel_type", ACCEL_TYPES) [email protected]( + "ifm_shape, new_shape", + [ + ((1, 4, 1, 2), (1, 1, 1, 8)), + ( + ( + 5, + 1, + 20, + ), + (1, 5, 1, 20), + ), + ((12, 20), (1, 6, 4, 10)), + ((12, 20), (6, 4, 10)), + ((20,), (4, 5)), + ], +) +def test_relay_reshape_codegen(ifm_shape, new_shape, accel_type): + # Create a "partitioned" Relay graph Review comment: It seems like some of this could be refactored out into a common function and reused between tests. ########## File path: tests/python/contrib/test_ethosu/test_codegen.py ########## @@ -343,5 +343,178 @@ def representative_dataset(): infra.verify_source(compiled_models, accel_type) [email protected]("accel_type", ACCEL_TYPES) [email protected]("ifm_shape", [(3, 2), (1, 15, 11, 7), (3, 1, 12), (400,)]) [email protected]("ifm_scale, ifm_zp, ofm_scale, ofm_zp", [(1, 0, 1, 0), (0.015, 3, 0.2, 5)]) +def test_ethosu_identity_codegen(ifm_shape, ifm_scale, ifm_zp, ofm_scale, ofm_zp, accel_type): + # Create a "partitioned" Relay function + ifm = relay.var("ifm", shape=ifm_shape, dtype="int8") + ifm0 = relay.var("ifm0", shape=ifm_shape, dtype="int8") + identity = infra.make_ethosu_identity( + ifm0, ifm_scale=ifm_scale, ifm_zero_point=ifm_zp, ofm_scale=ofm_scale, ofm_zero_point=ofm_zp + ) + glb_ethosu = relay.GlobalVar("tvmgen_default_ethosu_main_0") + + func = ( + relay.Function([ifm0], identity) + .with_attr("Inline", 1) + .with_attr("Compiler", "ethosu") + .with_attr("global_symbol", "tvmgen_default_ethosu_main_0") + .with_attr("Primitive", 1) + ) + mod = tvm.IRModule() + mod[glb_ethosu] = func + mod = relay.transform.InferType()(mod) + + call = relay.Call(glb_ethosu, [ifm]) + mod["main"] = relay.Function([ifm], call) + mod = relay.transform.InferType()(mod) + + in_data = np.random.randint(-120, high=120, size=ifm_shape, dtype="int8") + requant_data = (ifm_scale * (in_data - ifm_zp)) / ofm_scale + ofm_zp + out_data = np.round(np.clip(requant_data, -128, 127)).astype("int8") + + compiled_model = infra.build_source( + mod, {"ifm": in_data}, [out_data], accel_type, output_tolerance=1 + ) + + imported_modules = compiled_model[0].executor_factory.lib.imported_modules + assert len(imported_modules) == 2 + ethosu_module = imported_modules[0] + + # Verify generated C source + get_cs = tvm._ffi.get_global_func("runtime.module.ethosu.getcs") + cmms = get_cs(ethosu_module) + cmms = bytes.fromhex(cmms) + + infra.print_payload(cmms) + infra.verify_source(compiled_model, accel_type) + + [email protected]("accel_type", ACCEL_TYPES) [email protected]( Review comment: We should either test the 'special' indices for reshape (-1 and -2), or we should explicitly disallow these. -- 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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
