This is an automated email from the ASF dual-hosted git repository.
lukhut pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 25e98dd5e4 [ACL] Enable int8 data type in pooling operators (#13488)
25e98dd5e4 is described below
commit 25e98dd5e4bd989ade14e36df1fc2e1f81be7885
Author: Leandro Nunes <[email protected]>
AuthorDate: Mon Nov 28 13:22:24 2022 +0000
[ACL] Enable int8 data type in pooling operators (#13488)
This enables int8 data type to be used in Compute Library
for the Arm(r) Architecture (ACL) BYOC integration.
This PR covers:
- nn.max_pool2d
- nn.avg_pool2d
- nn.l2_pool2d
- nn.global_avg_pool2d
- nn.global_max_pool2d
Co-authored-by: Luke Hutton <[email protected]>
---
python/tvm/relay/op/contrib/arm_compute_lib.py | 6 +-
.../contrib/test_arm_compute_lib/test_pooling.py | 414 +++++++++++----------
2 files changed, 228 insertions(+), 192 deletions(-)
diff --git a/python/tvm/relay/op/contrib/arm_compute_lib.py
b/python/tvm/relay/op/contrib/arm_compute_lib.py
index 97f695baca..1b9abb0948 100644
--- a/python/tvm/relay/op/contrib/arm_compute_lib.py
+++ b/python/tvm/relay/op/contrib/arm_compute_lib.py
@@ -267,7 +267,7 @@ def
arm_compute_lib_pattern_table(disabled_ops=["concatenate"]):
def check_avg_pool2d(extract):
"""Check average pool2d pattern is supported by ACL."""
- if extract.attrs.dtype != "uint8":
+ if extract.attrs.dtype not in ("uint8", "int8"):
return False
pool = extract.args[0]
if pool.args[0].attrs.dtype != "int32":
@@ -440,7 +440,7 @@ def max_pool2d(expr):
if attrs.layout != "NHWC":
return False
typ = args[0].checked_type
- if typ.dtype not in ["float32", "uint8"]:
+ if typ.dtype not in ["float32", "uint8", "int8"]:
return False
return check_dilation(attrs)
@@ -468,7 +468,7 @@ def global_max_pool2d(expr):
"""Check if the external ACL codegen for gloval_maxpool2d should be
used."""
attrs, args = expr.attrs, expr.args
typ = args[0].checked_type
- if typ.dtype not in ["float32", "uint8"]:
+ if typ.dtype not in ["float32", "uint8", "int8"]:
return False
if attrs.layout != "NHWC":
return False
diff --git a/tests/python/contrib/test_arm_compute_lib/test_pooling.py
b/tests/python/contrib/test_arm_compute_lib/test_pooling.py
index b174f9a788..f08fa0059d 100644
--- a/tests/python/contrib/test_arm_compute_lib/test_pooling.py
+++ b/tests/python/contrib/test_arm_compute_lib/test_pooling.py
@@ -15,8 +15,9 @@
# specific language governing permissions and limitations
# under the License.
"""Arm Compute Library integration pooling tests."""
-
import numpy as np
+import pytest
+
import tvm
from tvm import relay, testing
@@ -46,6 +47,7 @@ def _get_pooling_model(
if len(padding) == 2:
padding = (padding[0], padding[1], padding[0], padding[1])
out = relay.var(next(var_names), shape=shape, dtype=dtype)
+ qnn_dtypes = ("uint8", "int8")
if typef == "nn.max_pool2d":
out = relay.nn.max_pool2d(
@@ -58,7 +60,7 @@ def _get_pooling_model(
layout="NHWC",
)
elif typef == "nn.avg_pool2d":
- if dtype == "uint8":
+ if dtype in qnn_dtypes:
out = relay.cast(out, "int32")
out = relay.nn.avg_pool2d(
out,
@@ -70,8 +72,8 @@ def _get_pooling_model(
count_include_pad=count_include_pad,
layout="NHWC",
)
- if dtype == "uint8":
- out = relay.cast(out, "uint8")
+ if dtype in qnn_dtypes:
+ out = relay.cast(out, dtype)
elif typef == "nn.l2_pool2d":
out = relay.power(out, relay.const(2.0))
out = relay.nn.avg_pool2d(
@@ -93,15 +95,16 @@ def _get_pooling_model(
def _get_global_pooling_model(shape, dtype, typef, var_names):
"""Return a model and any parameters it may have."""
out = relay.var(next(var_names), shape=shape, dtype=dtype)
+ qnn_dtypes = ("uint8", "int8")
if typef == "nn.global_max_pool2d":
out = relay.nn.global_max_pool2d(out, layout="NHWC")
elif typef == "nn.global_avg_pool2d":
- if dtype == "uint8":
+ if dtype in qnn_dtypes:
out = relay.cast(out, "int32")
out = relay.nn.global_avg_pool2d(out, layout="NHWC")
- if dtype == "uint8":
- out = relay.cast(out, "uint8")
+ if dtype in qnn_dtypes:
+ out = relay.cast(out, dtype)
else:
raise ValueError("Function not supported")
@@ -160,7 +163,59 @@ def _get_expected_global_pooling_codegen(shape, dtype,
typef):
return [input, node]
-def test_pooling():
+def _get_low_high_atol_rtol(dtype):
+ if dtype == "float32":
+ low, high, atol, rtol = (-127, 128, 0.001, 0.001)
+ elif dtype == "uint8":
+ low, high, atol, rtol = (0, 255, 1, 0)
+ elif dtype == "int8":
+ low, high, atol, rtol = (-127, 128, 1, 0)
+ else:
+ pytest.fail(f"dtype not expected: {dtype}")
+
+ return low, high, atol, rtol
+
+
+# fmt: off
[email protected](
+
"typef,dtype,size,stride,dilation,pad,ceil_mode,count_include_pad,input_shape,expected_ops",
+ [
+ ("nn.max_pool2d", "float32", (3, 3), (2, 2), (1, 1), (0, 0), False,
False, (27, 27, 512), (0, 1),),
+ ("nn.max_pool2d", "float32", (2, 2), (2, 2), (1, 1), (0, 0), False,
True, (16, 16, 16), (0, 1),),
+ ("nn.max_pool2d", "float32", (3, 3), (2, 2), (1, 1), (1, 1), True,
True, (15, 15, 16), (0, 1),),
+ ("nn.max_pool2d", "float32", (2, 2), (2, 2), (1, 1), (0, 1), False,
False, (16, 16, 16), (0, 1),),
+ ("nn.max_pool2d", "uint8", (3, 3), (2, 2), (1, 1), (0, 1), False,
False, (16, 16, 16), (0, 1),),
+ ("nn.max_pool2d", "uint8", (2, 2), (2, 2), (1, 1), (1, 1), True,
True, (15, 15, 16), (0, 1),),
+ ("nn.max_pool2d", "uint8", (2, 2), (2, 2), (3, 2), (1, 1), True,
True, (15, 15, 16), (1, 0),),
+ ("nn.max_pool2d", "int8", (3, 3), (2, 2), (1, 1), (0, 1), False,
False, (16, 16, 16), (0, 1),),
+ ("nn.max_pool2d", "int8", (2, 2), (2, 2), (1, 1), (1, 1), True, True,
(15, 15, 16), (0, 1),),
+ ("nn.max_pool2d", "int8", (2, 2), (2, 2), (3, 2), (1, 1), True, True,
(15, 15, 16), (1, 0),),
+ ("nn.avg_pool2d", "float32", (2, 2), (2, 2), (1, 1), (1, 1), False,
False, (16, 16, 16), (0, 1),),
+ ("nn.avg_pool2d", "float32", (2, 2), (2, 2), (1, 1), (0, 0), False,
True, (16, 16, 16), (0, 1),),
+ ("nn.avg_pool2d", "float32", (3, 3), (2, 2), (3, 2), (0, 1), True,
False, (15, 15, 16), (1, 0),),
+ # 20.05: "exclude_padding equal false is not supported for AVG Pooling
with padding on quantized types"
+ # ["nn.avg_pool2d", uint8_dtype, (2, 2), (2, 2), (1, 1), False, True,
(16, 16, 16)],
+ ("nn.avg_pool2d", "uint8", (3, 3), (2, 2), (1, 1), (0, 1), False,
False, (16, 16, 16), (0, 1),),
+ ("nn.avg_pool2d", "int8", (3, 3), (2, 2), (1, 1), (0, 1), False,
False, (16, 16, 16), (0, 1),),
+ ("nn.l2_pool2d", "float32", (2, 2), (2, 2), (1, 1), (0, 1), True,
False, (16, 16, 16), (0, 1),),
+ ("nn.l2_pool2d", "float32", (3, 3), (2, 2), (1, 1), (0, 0), False,
False, (16, 16, 16), (0, 1),),
+ ("nn.l2_pool2d", "float32", (2, 2), (2, 2), (1, 1), (1, 1), False,
True, (15, 15, 16), (0, 1),),
+
+ ],
+)
+# fmt: on
+def test_pooling(
+ typef,
+ dtype,
+ size,
+ stride,
+ dilation,
+ pad,
+ ceil_mode,
+ count_include_pad,
+ input_shape,
+ expected_ops,
+):
Device.load("test_config.json")
if skip_runtime_test():
@@ -169,91 +224,77 @@ def test_pooling():
device = Device()
np.random.seed(0)
- fp32_dtype = ("float32", -127, 128, 0.001, 0.001)
- uint8_dtype = ("uint8", 0, 255, 1, 0)
- # fmt: off
- trials = [
- ["nn.max_pool2d", fp32_dtype, (3, 3), (2, 2), (1, 1), (0, 0), False,
False, (27, 27, 512), (0, 1),],
- ["nn.max_pool2d", fp32_dtype, (2, 2), (2, 2), (1, 1), (0, 0), False,
True, (16, 16, 16), (0, 1),],
- ["nn.max_pool2d", fp32_dtype, (3, 3), (2, 2), (1, 1), (1, 1), True,
True, (15, 15, 16), (0, 1),],
- ["nn.max_pool2d", fp32_dtype, (2, 2), (2, 2), (1, 1), (0, 1), False,
False, (16, 16, 16), (0, 1),],
- ["nn.max_pool2d", uint8_dtype, (3, 3), (2, 2), (1, 1), (0, 1), False,
False, (16, 16, 16), (0, 1),],
- ["nn.max_pool2d", uint8_dtype, (2, 2), (2, 2), (1, 1), (1, 1), True,
True, (15, 15, 16), (0, 1),],
- ["nn.max_pool2d", uint8_dtype, (2, 2), (2, 2), (3, 2), (1, 1), True,
True, (15, 15, 16), (1, 0),],
- ["nn.avg_pool2d", fp32_dtype, (2, 2), (2, 2), (1, 1), (1, 1), False,
False, (16, 16, 16), (0, 1),],
- ["nn.avg_pool2d", fp32_dtype, (2, 2), (2, 2), (1, 1), (0, 0), False,
True, (16, 16, 16), (0, 1),],
- ["nn.avg_pool2d", fp32_dtype, (3, 3), (2, 2), (3, 2), (0, 1), True,
False, (15, 15, 16), (1, 0),],
- # 20.05: "exclude_padding equal false is not supported for AVG Pooling
with padding on quantized types"
- # ["nn.avg_pool2d", uint8_dtype, (2, 2), (2, 2), (1, 1), False, True,
(16, 16, 16)],
- ["nn.avg_pool2d", uint8_dtype, (3, 3), (2, 2), (1, 1), (0, 1), False,
False, (16, 16, 16), (0, 1),],
- ["nn.l2_pool2d", fp32_dtype, (2, 2), (2, 2), (1, 1), (0, 1), True,
False, (16, 16, 16), (0, 1),],
- ["nn.l2_pool2d", fp32_dtype, (3, 3), (2, 2), (1, 1), (0, 0), False,
False, (16, 16, 16), (0, 1),],
- ["nn.l2_pool2d", fp32_dtype, (2, 2), (2, 2), (1, 1), (1, 1), False,
True, (15, 15, 16), (0, 1),],
- ]
- # fmt: on
- for (
+ low, high, atol, rtol = _get_low_high_atol_rtol(dtype)
+ tvm_ops, acl_partitions = expected_ops
+
+ shape = (1, *input_shape)
+ outputs = []
+ inputs = {
+ "a": tvm.nd.array(np.random.uniform(low, high, shape).astype(dtype)),
+ }
+
+ func = _get_pooling_model(
+ shape,
+ dtype,
typef,
- (dtype, low, high, atol, rtol),
size,
stride,
dilation,
pad,
ceil_mode,
count_include_pad,
- input_shape,
- (tvm_ops, acl_partitions),
- ) in trials:
- shape = (1, *input_shape)
- outputs = []
- inputs = {
- "a": tvm.nd.array(np.random.uniform(low, high,
shape).astype(dtype)),
- }
-
- func = _get_pooling_model(
- shape,
- dtype,
- typef,
- size,
- stride,
- dilation,
- pad,
- ceil_mode,
- count_include_pad,
- iter(inputs),
+ iter(inputs),
+ )
+
+ config = {
+ "size": size,
+ "stride": stride,
+ "shape": shape,
+ "pooling type": typef,
+ "dtype": dtype,
+ "padding": pad,
+ "dilation": dilation,
+ "ceil_mode": ceil_mode,
+ "count_include_pad": count_include_pad,
+ "inputs": inputs,
+ }
+ verify_saturation = True if dtype == "uint8" else False
+ for acl in [False, True]:
+ outputs.append(
+ build_and_run(
+ func,
+ inputs,
+ 1,
+ None,
+ device,
+ enable_acl=acl,
+ tvm_ops=tvm_ops,
+ acl_partitions=acl_partitions,
+ config=config,
+ )[0]
)
- config = {
- "size": size,
- "stride": stride,
- "shape": shape,
- "pooling type": typef,
- "dtype": dtype,
- "padding": pad,
- "dilation": dilation,
- "ceil_mode": ceil_mode,
- "count_include_pad": count_include_pad,
- "inputs": inputs,
- }
- verify_saturation = True if dtype == "uint8" else False
- for acl in [False, True]:
- outputs.append(
- build_and_run(
- func,
- inputs,
- 1,
- None,
- device,
- enable_acl=acl,
- tvm_ops=tvm_ops,
- acl_partitions=acl_partitions,
- config=config,
- )[0]
- )
-
- verify(outputs, atol=atol, rtol=rtol, config=config,
verify_saturation=verify_saturation)
-
-
-def test_global_pooling():
+ verify(outputs, atol=atol, rtol=rtol, config=config,
verify_saturation=verify_saturation)
+
+
[email protected](
+ "typef,dtype,input_shape",
+ [
+ ["nn.global_max_pool2d", "float32", (8, 8, 16)],
+ ["nn.global_max_pool2d", "float32", (9, 9, 16)],
+ ["nn.global_max_pool2d", "uint8", (8, 8, 16)],
+ ["nn.global_max_pool2d", "uint8", (9, 9, 16)],
+ ["nn.global_max_pool2d", "int8", (8, 8, 16)],
+ ["nn.global_max_pool2d", "int8", (9, 9, 16)],
+ ["nn.global_avg_pool2d", "float32", (8, 8, 16)],
+ ["nn.global_avg_pool2d", "float32", (9, 9, 16)],
+ ["nn.global_avg_pool2d", "uint8", (8, 8, 16)],
+ ["nn.global_avg_pool2d", "uint8", (9, 9, 16)],
+ ["nn.global_avg_pool2d", "int8", (8, 8, 16)],
+ ["nn.global_avg_pool2d", "int8", (9, 9, 16)],
+ ],
+)
+def test_global_pooling(typef, dtype, input_shape):
Device.load("test_config.json")
if skip_runtime_test():
@@ -262,118 +303,113 @@ def test_global_pooling():
device = Device()
np.random.seed(0)
- fp32_dtype = ("float32", -127, 128, 0.001, 0.001)
- uint8_dtype = ("uint8", 0, 255, 1, 0)
-
- trials = [
- ["nn.global_max_pool2d", fp32_dtype, (8, 8, 16)],
- ["nn.global_max_pool2d", fp32_dtype, (9, 9, 16)],
- ["nn.global_max_pool2d", fp32_dtype, (8, 8, 16)],
- ["nn.global_max_pool2d", uint8_dtype, (8, 8, 16)],
- ["nn.global_max_pool2d", uint8_dtype, (9, 9, 16)],
- ["nn.global_avg_pool2d", fp32_dtype, (8, 8, 16)],
- ["nn.global_avg_pool2d", fp32_dtype, (8, 8, 16)],
- ["nn.global_avg_pool2d", fp32_dtype, (9, 9, 16)],
- ["nn.global_avg_pool2d", uint8_dtype, (8, 8, 16)],
- ["nn.global_avg_pool2d", uint8_dtype, (8, 8, 16)],
- ]
-
- for typef, (dtype, low, high, atol, rtol), input_shape in trials:
- shape = (1, *input_shape)
- outputs = []
- inputs = {
- "a": tvm.nd.array(np.random.uniform(low, high,
shape).astype(dtype)),
- }
-
- func = _get_global_pooling_model(shape, dtype, typef, iter(inputs))
- config = {
- "shape": shape,
- "pooling type": typef,
- "dtype": dtype,
- }
- verify_saturation = True if dtype == "uint8" else False
-
- for acl in [False, True]:
- outputs.append(
- build_and_run(func, inputs, 1, None, device, enable_acl=acl,
config=config)[0]
- )
-
- verify(outputs, atol=atol, rtol=rtol, config=config,
verify_saturation=verify_saturation)
-
-
-def test_codegen_pooling():
- if skip_codegen_test():
- return
+ low, high, rtol, atol = _get_low_high_atol_rtol(dtype)
- fp32_dtype = ("float32", -127, 128)
- uint8_dtype = ("uint8", 0, 255)
- # fmt: off
- trials = [
- ["nn.max_pool2d", fp32_dtype, (2, 2), (2, 2), (1, 1), (0, 0), False,
True, (16, 16, 16), (0, 1),],
- ["nn.max_pool2d", fp32_dtype, (3, 3), (2, 2), (1, 1), (1, 1), True,
True, (15, 15, 16), (0, 1),],
- ["nn.max_pool2d", fp32_dtype, (2, 2), (2, 2), (1, 1), (0, 1), False,
False, (16, 16, 16), (0, 1),],
- ["nn.max_pool2d", uint8_dtype, (3, 3), (2, 2), (1, 1), (0, 1), False,
False, (16, 16, 16), (0, 1),],
- ["nn.max_pool2d", uint8_dtype, (2, 2), (2, 2), (1, 1), (1, 1), True,
True, (15, 15, 16), (0, 1),],
- ["nn.max_pool2d", uint8_dtype, (2, 2), (2, 2), (3, 2), (1, 1), True,
True, (15, 15, 16), (1, 0),],
- ["nn.avg_pool2d", fp32_dtype, (2, 2), (2, 2), (1, 1), (1, 1), False,
False, (16, 16, 16), (0, 1),],
- ["nn.avg_pool2d", fp32_dtype, (2, 2), (2, 2), (1, 1), (1, 1), False,
False, (16, 16, 16), (0, 1),],
- ["nn.avg_pool2d", fp32_dtype, (2, 2), (2, 2), (1, 1), (0, 0), False,
True, (16, 16, 16), (0, 1),],
- ["nn.avg_pool2d", fp32_dtype, (3, 3), (2, 2), (3, 2), (0, 1), True,
False, (15, 15, 16), (1, 0),],
- ["nn.avg_pool2d", uint8_dtype, (2, 2), (2, 2), (1, 1), (1, 1), False,
True, (16, 16, 16), (0, 1),],
- ["nn.avg_pool2d", uint8_dtype, (3, 3), (2, 2), (1, 1), (0, 1), False,
False, (16, 16, 16), (0, 1),],
- ["nn.l2_pool2d", fp32_dtype, (2, 2), (2, 2), (1, 1), (0, 1), True,
False, (15, 15, 16), (0, 1),],
- ["nn.l2_pool2d", fp32_dtype, (3, 3), (2, 2), (1, 1), (0, 0), False,
False, (16, 16, 16), (0, 1),],
- ["nn.l2_pool2d", fp32_dtype, (2, 2), (2, 2), (1, 1), (1, 1), False,
True, (15, 15, 16), (0, 1),],
- ]
- # fmt: on
- for (
- typef,
- (dtype, low, high),
- size,
- stride,
- dilation,
- pad,
- ceil_mode,
- count_include_pad,
- input_shape,
- (tvm_ops, acl_partitions),
- ) in trials:
- shape = (1, *input_shape)
- inputs = {"a"}
- args = (shape, dtype, typef, size, stride, dilation, pad, False, False)
- func = _get_pooling_model(*args, iter(inputs))
- exp_codegen = _get_expected_pooling_codegen(*args)
+ shape = (1, *input_shape)
+ outputs = []
+ inputs = {
+ "a": tvm.nd.array(np.random.uniform(low, high, shape).astype(dtype)),
+ }
- verify_codegen(func, exp_codegen, acl_partitions, tvm_ops)
+ func = _get_global_pooling_model(shape, dtype, typef, iter(inputs))
+ config = {
+ "shape": shape,
+ "pooling type": typef,
+ "dtype": dtype,
+ }
+ verify_saturation = True if dtype in ("uint8", "int8") else False
+ for acl in [False, True]:
+ outputs.append(
+ build_and_run(func, inputs, 1, None, device, enable_acl=acl,
config=config)[0]
+ )
+
+ verify(outputs, atol=atol, rtol=rtol, config=config,
verify_saturation=verify_saturation)
+
+
+# fmt: off
[email protected](
+
"typef,dtype,size,stride,dilation,pad,ceil_mode,count_include_pad,input_shape,expected_ops",
+ [
+ ("nn.max_pool2d", "float32", (2, 2), (2, 2), (1, 1), (0, 0), False,
True, (16, 16, 16), (0, 1),),
+ ("nn.max_pool2d", "float32", (3, 3), (2, 2), (1, 1), (1, 1), True,
True, (15, 15, 16), (0, 1),),
+ ("nn.max_pool2d", "float32", (2, 2), (2, 2), (1, 1), (0, 1), False,
False, (16, 16, 16), (0, 1),),
+ ("nn.max_pool2d", "uint8", (3, 3), (2, 2), (1, 1), (0, 1), False,
False, (16, 16, 16), (0, 1),),
+ ("nn.max_pool2d", "uint8", (2, 2), (2, 2), (1, 1), (1, 1), True,
True, (15, 15, 16), (0, 1),),
+ ("nn.max_pool2d", "uint8", (2, 2), (2, 2), (3, 2), (1, 1), True,
True, (15, 15, 16), (1, 0),),
+ ("nn.max_pool2d", "int8", (3, 3), (2, 2), (1, 1), (0, 1), False,
False, (16, 16, 16), (0, 1),),
+ ("nn.max_pool2d", "int8", (2, 2), (2, 2), (1, 1), (1, 1), True,
True, (15, 15, 16), (0, 1),),
+ ("nn.max_pool2d", "int8", (2, 2), (2, 2), (3, 2), (1, 1), True,
True, (15, 15, 16), (1, 0),),
+ ("nn.avg_pool2d", "float32", (2, 2), (2, 2), (1, 1), (1, 1), False,
False, (16, 16, 16), (0, 1),),
+ ("nn.avg_pool2d", "float32", (2, 2), (2, 2), (1, 1), (1, 1), False,
False, (16, 16, 16), (0, 1),),
+ ("nn.avg_pool2d", "float32", (2, 2), (2, 2), (1, 1), (0, 0), False,
True, (16, 16, 16), (0, 1),),
+ ("nn.avg_pool2d", "float32", (3, 3), (2, 2), (3, 2), (0, 1), True,
False, (15, 15, 16), (1, 0),),
+ ("nn.avg_pool2d", "uint8", (2, 2), (2, 2), (1, 1), (1, 1), False,
True, (16, 16, 16), (0, 1),),
+ ("nn.avg_pool2d", "uint8", (3, 3), (2, 2), (1, 1), (0, 1), False,
False, (16, 16, 16), (0, 1),),
+ ("nn.avg_pool2d", "int8", (2, 2), (2, 2), (1, 1), (1, 1), False,
True, (16, 16, 16), (0, 1),),
+ ("nn.avg_pool2d", "int8", (3, 3), (2, 2), (1, 1), (0, 1), False,
False, (16, 16, 16), (0, 1),),
+ ("nn.l2_pool2d", "float32", (2, 2), (2, 2), (1, 1), (0, 1), True,
False, (15, 15, 16), (0, 1),),
+ ("nn.l2_pool2d", "float32", (3, 3), (2, 2), (1, 1), (0, 0), False,
False, (16, 16, 16), (0, 1),),
+ ("nn.l2_pool2d", "float32", (2, 2), (2, 2), (1, 1), (1, 1), False,
True, (15, 15, 16), (0, 1),),
+ ],
+)
+# fmt: on
+def test_codegen_pooling(
+ typef,
+ dtype,
+ size,
+ stride,
+ dilation,
+ pad,
+ ceil_mode,
+ count_include_pad,
+ input_shape,
+ expected_ops,
+):
+ if skip_codegen_test():
+ return
-def test_codegen_global_pooling():
+ low, high, _, _ = _get_low_high_atol_rtol(dtype)
+ tvm_ops, acl_partitions = expected_ops
+
+ shape = (1, *input_shape)
+ inputs = {"a"}
+ args = (shape, dtype, typef, size, stride, dilation, pad, False, False)
+ func = _get_pooling_model(*args, iter(inputs))
+ exp_codegen = _get_expected_pooling_codegen(*args)
+
+ verify_codegen(func, exp_codegen, acl_partitions, tvm_ops)
+
+
[email protected](
+ "typef,dtype,input_shape",
+ [
+ ("nn.global_max_pool2d", "float32", (8, 8, 16)),
+ ("nn.global_max_pool2d", "float32", (9, 9, 16)),
+ ("nn.global_max_pool2d", "uint8", (8, 8, 16)),
+ ("nn.global_max_pool2d", "uint8", (9, 9, 16)),
+ ("nn.global_max_pool2d", "int8", (8, 8, 16)),
+ ("nn.global_max_pool2d", "int8", (9, 9, 16)),
+ ("nn.global_avg_pool2d", "float32", (8, 8, 16)),
+ ("nn.global_avg_pool2d", "float32", (9, 9, 16)),
+ ("nn.global_avg_pool2d", "uint8", (8, 8, 16)),
+ ("nn.global_avg_pool2d", "uint8", (9, 9, 16)),
+ ("nn.global_avg_pool2d", "int8", (8, 8, 16)),
+ ("nn.global_avg_pool2d", "int8", (9, 9, 16)),
+ ],
+)
+def test_codegen_global_pooling(typef, dtype, input_shape):
if skip_codegen_test():
return
- fp32_dtype = ("float32", -127, 128)
- uint8_dtype = ("uint8", 0, 255)
-
- trials = [
- ["nn.global_max_pool2d", fp32_dtype, (8, 8, 16)],
- ["nn.global_max_pool2d", fp32_dtype, (9, 9, 16)],
- ["nn.global_max_pool2d", fp32_dtype, (8, 8, 16)],
- ["nn.global_max_pool2d", uint8_dtype, (8, 8, 16)],
- ["nn.global_max_pool2d", uint8_dtype, (9, 9, 16)],
- ["nn.global_avg_pool2d", fp32_dtype, (8, 8, 16)],
- ["nn.global_avg_pool2d", fp32_dtype, (8, 8, 16)],
- ["nn.global_avg_pool2d", fp32_dtype, (9, 9, 16)],
- ["nn.global_avg_pool2d", uint8_dtype, (8, 8, 16)],
- ["nn.global_avg_pool2d", uint8_dtype, (8, 8, 16)],
- ]
-
- for typef, (dtype, low, high), input_shape in trials:
- shape = (1, *input_shape)
- inputs = {"a"}
- args = (shape, dtype, typef)
- func = _get_global_pooling_model(*args, iter(inputs))
- exp_codegen = _get_expected_global_pooling_codegen(*args)
- verify_codegen(func, exp_codegen, 1)
+ low, high, _, _ = _get_low_high_atol_rtol(dtype)
+
+ shape = (1, *input_shape)
+ inputs = {"a"}
+ args = (shape, dtype, typef)
+ func = _get_global_pooling_model(*args, iter(inputs))
+ exp_codegen = _get_expected_global_pooling_codegen(*args)
+ verify_codegen(func, exp_codegen, 1)
if __name__ == "__main__":
