NicolaLancellotti commented on a change in pull request #9384: URL: https://github.com/apache/tvm/pull/9384#discussion_r738597335
########## File path: tests/python/contrib/test_ethosu/test_replace_pooling.py ########## @@ -0,0 +1,158 @@ +# 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. +import pytest + +pytest.importorskip("ethosu.vela") + +import tvm +from tvm import relay +from tvm.relay.testing import run_opt_pass +from tvm.relay.backend.contrib.ethosu.tir import spec +from tvm.relay.backend.contrib.ethosu.tir.compiler import lower_to_tir +from .infra import make_ethosu_pooling, get_pooling_args + + [email protected]( + "ifm_shape, ofm_channels, ifm_layout, ofm_layout", + [ + ((1, 5, 9, 3), 3, "NHWC", "NHWC"), + ((1, 8, 3, 9, 16), 40, "NHCWB16", "NHCWB16"), + ((1, 8, 3, 9, 16), 40, "NHCWB16", "NHWC"), + ((1, 8, 9, 40), 40, "NHWC", "NHCWB16"), + ], +) [email protected]("pooling_type", ["AVG", "MAX"]) [email protected]("activation", ["NONE", "CLIP", "TANH", "SIGMOID"]) +def test_pooling_single( + ifm_shape, + ofm_channels, + ifm_layout, + ofm_layout, + pooling_type, + activation, +): + pool_shape = (3, 2) + strides = (1, 2) + padding = (1, 1, 1, 0) + ifm = relay.var("ifm", shape=ifm_shape, dtype="int8") + pooling = make_ethosu_pooling( + ifm, + pooling_type, + pool_shape, + ofm_channels, + strides, + padding, + activation, + ifm_layout, + ofm_layout, + ) + func = relay.Function(relay.analysis.free_vars(pooling), pooling) + func = run_opt_pass(func, relay.transform.InferType()) + mod, _ = lower_to_tir(func) + data = [] + + def _visit(stmt): + if isinstance(stmt, tvm.tir.Call): + data.append(get_pooling_args(stmt)) + + tvm.tir.stmt_functor.post_order_visit(mod["main"].body, _visit) + if ifm_layout == "NHWC": + ifm_stride_c = 1 + ifm_stride_w = ifm_shape[3] + ifm_stride_h = ifm_shape[2] * ifm_shape[3] + ofm_height = (ifm_shape[1] - pool_shape[0] + padding[0] + padding[0]) // strides[0] + 1 + ofm_width = (ifm_shape[2] - pool_shape[1] + padding[1] + padding[1]) // strides[1] + 1 + else: + ifm_stride_w = 16 + ifm_stride_c = 16 * ifm_shape[3] + ifm_stride_h = 16 * ifm_shape[2] * ifm_shape[3] + ofm_height = (ifm_shape[1] - pool_shape[0] + padding[0] + padding[0]) // strides[0] + 1 + ofm_width = (ifm_shape[3] - pool_shape[1] + padding[1] + padding[1]) // strides[1] + 1 + + if ofm_layout == "NHWC": + ofm_stride_c = 1 + ofm_stride_w = ofm_channels if ofm_width > 1 else 1 + ofm_stride_h = ofm_channels * ofm_width if ofm_height > 1 else 1 + else: + ofm_stride_w = 16 + ofm_stride_c = 16 * ofm_width + ofm_stride_h = 16 * ofm_width * ((ofm_channels - 1) // 16 + 1) Review comment: The tests for this and the other operators are quite verbose, and we have relied heavily on `pytest.mark.parametrize` to reduce the number of tests. Unfortunately, not all things can be parametrised, so some `if`s are necessary. If we had to write multiple tests there would be a lot of code duplication. What do other people think? ########## File path: python/tvm/relay/backend/contrib/ethosu/op/pooling.py ########## @@ -0,0 +1,182 @@ +# 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 operators for pooling""" +from typing import Tuple + +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 pooling_compute + + +def _extract_ethosu_pooling_params(attrs, args): + """Get the parameters necessary to construct a ethosu_pooling compute TE + from a ethosu_pooling Relay call.""" + ifm = args[0] + lut = args[1] + pooling_type = attrs.pooling_type + ifm_scale = attrs.ifm_scale + ifm_zero_point = attrs.ifm_zero_point + ofm_scale = attrs.ofm_scale + ofm_zero_point = attrs.ofm_zero_point + pool_shape = attrs.pool_shape + ofm_channels = attrs.ofm_channels + strides = attrs.strides + padding = attrs.padding + activation = attrs.activation + clip_min = attrs.clip_min + clip_max = attrs.clip_max + upscale = attrs.upscale + ifm_layout = attrs.ifm_layout + ofm_layout = attrs.ofm_layout + + return ( + ifm, + lut, + pooling_type, + ifm_scale, + ifm_zero_point, + ofm_scale, + ofm_zero_point, + pool_shape, + ofm_channels, + strides, + padding, + activation, + clip_min, + clip_max, + upscale, + ifm_layout, + ofm_layout, + ) + + [email protected]_op_attr("contrib.ethosu.pooling", "FTVMCompute") +def create_ethosu_pooling_compute(attrs, args, out_type): + """Create an ethosu_pooling compute op.""" + params = _extract_ethosu_pooling_params(attrs, args) + op = pooling_compute(*params) + return [op] + + [email protected]_op_attr("contrib.ethosu.pooling", "FTVMStrategy") +def pooling_strategy_ethosu(attrs, inputs, out_type, target): + strategy = OpStrategy() + strategy.add_implementation( + create_ethosu_pooling_compute, + _strategy.wrap_topi_schedule(schedule_injective), + name="ethosu_pooling", + ) + return strategy + + +def ethosu_pooling( + ifm: tvm.relay.Expr, + lut: tvm.relay.Expr, + pooling_type: str, + ifm_scale: float, + ifm_zero_point: int, + ofm_scale: float, + ofm_zero_point: int, + pool_shape: Tuple[int, int], + ofm_channels: int, + strides: Tuple[int, int] = (1, 1), + padding: Tuple[int, int, int, int] = (0, 0, 0, 0), + activation: str = "NONE", + clip_min: int = 0, + clip_max: int = 0, + upscale: str = "NONE", + ifm_layout: str = "NHWC", + ofm_layout: str = "NHWC", +) -> tvm.relay.Call: + """This is a quantized 2D pooling operation as supported by the + Ethos(TM)-U NPU. It accepts either NHWC or NHCWB16 format + for the input data. + + Parameters + ---------- + ifm : tvm.relay.Expr + The Input Feature Map tensor (IFM). + lut : tvm.relay.Expr + The look-up table values to use if activation = "LUT". Review comment: Ack ########## File path: python/tvm/relay/backend/contrib/ethosu/op/pooling.py ########## @@ -0,0 +1,182 @@ +# 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 operators for pooling""" +from typing import Tuple + +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 pooling_compute + + +def _extract_ethosu_pooling_params(attrs, args): + """Get the parameters necessary to construct a ethosu_pooling compute TE + from a ethosu_pooling Relay call.""" + ifm = args[0] + lut = args[1] + pooling_type = attrs.pooling_type + ifm_scale = attrs.ifm_scale + ifm_zero_point = attrs.ifm_zero_point + ofm_scale = attrs.ofm_scale + ofm_zero_point = attrs.ofm_zero_point + pool_shape = attrs.pool_shape + ofm_channels = attrs.ofm_channels + strides = attrs.strides + padding = attrs.padding + activation = attrs.activation + clip_min = attrs.clip_min + clip_max = attrs.clip_max + upscale = attrs.upscale + ifm_layout = attrs.ifm_layout + ofm_layout = attrs.ofm_layout + + return ( + ifm, + lut, + pooling_type, + ifm_scale, + ifm_zero_point, + ofm_scale, + ofm_zero_point, + pool_shape, + ofm_channels, + strides, + padding, + activation, + clip_min, + clip_max, + upscale, + ifm_layout, + ofm_layout, + ) + + [email protected]_op_attr("contrib.ethosu.pooling", "FTVMCompute") +def create_ethosu_pooling_compute(attrs, args, out_type): + """Create an ethosu_pooling compute op.""" + params = _extract_ethosu_pooling_params(attrs, args) + op = pooling_compute(*params) + return [op] + + [email protected]_op_attr("contrib.ethosu.pooling", "FTVMStrategy") +def pooling_strategy_ethosu(attrs, inputs, out_type, target): + strategy = OpStrategy() + strategy.add_implementation( + create_ethosu_pooling_compute, + _strategy.wrap_topi_schedule(schedule_injective), + name="ethosu_pooling", + ) + return strategy + + +def ethosu_pooling( + ifm: tvm.relay.Expr, + lut: tvm.relay.Expr, + pooling_type: str, + ifm_scale: float, + ifm_zero_point: int, + ofm_scale: float, + ofm_zero_point: int, + pool_shape: Tuple[int, int], + ofm_channels: int, + strides: Tuple[int, int] = (1, 1), + padding: Tuple[int, int, int, int] = (0, 0, 0, 0), + activation: str = "NONE", + clip_min: int = 0, + clip_max: int = 0, + upscale: str = "NONE", + ifm_layout: str = "NHWC", + ofm_layout: str = "NHWC", +) -> tvm.relay.Call: + """This is a quantized 2D pooling operation as supported by the + Ethos(TM)-U NPU. It accepts either NHWC or NHCWB16 format + for the input data. + + Parameters + ---------- + ifm : tvm.relay.Expr + The Input Feature Map tensor (IFM). + lut : tvm.relay.Expr + The look-up table values to use if activation = "LUT". + pooling_type: str + The type of the pooling. "AVG" - average pool, "MAX" - max pool. + ifm_scale : float + The quantization scale for the Input Feature Map tensor. + ifm_zero_point : int + The quantization zero point for the Input Feature Map tensor. + ofm_scale : float + The quantization scale for the Output Feature Map tensor. + ofm_zero_point : int + The quantization zero point for the Output Feature Map tensor. + pool_shape : tuple of int + The 2 dimensional pool shape as (pool_shape_height, pool_shape_width). + ofm_channels : int + The number of OFM channels Review comment: Ack ########## File path: python/tvm/relay/backend/contrib/ethosu/te/pooling.py ########## @@ -0,0 +1,134 @@ +# 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 +"""Tensor Expressions for poolings""" +from typing import Tuple + +from tvm import te +from .dma import dma_ofm_compute, dma_ifm_compute + + +def pooling_compute( + ifm: te.Tensor, + lut: te.Tensor, + pooling_type: str, + ifm_scale: float, + ifm_zero_point: int, + ofm_scale: float, + ofm_zero_point: int, + pool_shape: Tuple[int, int], + ofm_channels: int, + strides: Tuple[int, int], + padding: Tuple[int, int, int, int], + activation: str, + clip_min: int, + clip_max: int, + upscale: str, + ifm_layout: str, + ofm_layout: str, +) -> te.Tensor: + """A compute operator representing the capabilities of pooling for the NPU. + + Parameters + ---------- + ifm : te.Tensor + The Input Feature Map tensor (IFM). + lut : te.Tensor + The look-up table values to use if activation = "LUT". Review comment: Ack -- 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]
