From fe2e63f860909fa550b44361c3ec149657cce331 Mon Sep 17 00:00:00 2001 From: Luke Hutton Date: Mon, 7 Aug 2023 17:18:53 +0000 Subject: [PATCH 1/2] [Relay][Strategy] Use x86 pool schedules for arm_cpu Similar to https://github.com/apache/tvm/pull/15470, x86 schedules are used in place of generic schedules to improve performance. Since the pooling strategy does not use `OpStrategy`, mocking is used to ensure the relevant `schedule_pool` function is called when lowing a Relay pooling operation with respect to a given target. Change-Id: I782fe00e29f9c9cf41b3405d33a82a79cd85a99b --- python/tvm/relay/op/strategy/arm_cpu.py | 3 +- .../strategy/test_select_implementation.py | 35 +++++++++++++++++++ 2 files changed, 36 insertions(+), 2 deletions(-) diff --git a/python/tvm/relay/op/strategy/arm_cpu.py b/python/tvm/relay/op/strategy/arm_cpu.py index 24966019db6d..f81354466664 100644 --- a/python/tvm/relay/op/strategy/arm_cpu.py +++ b/python/tvm/relay/op/strategy/arm_cpu.py @@ -74,8 +74,7 @@ def schedule_pool_arm_cpu(attrs, outs, target): and layout in ("NWC", "NHWC") ): return topi.arm_cpu.schedule_pool(outs, layout) - logger.warning("pool is not optimized for arm cpu.") - return topi.generic.schedule_pool(outs, layout) + return topi.x86.schedule_pool(outs, layout) def _get_padding_width(padding): diff --git a/tests/python/relay/strategy/test_select_implementation.py b/tests/python/relay/strategy/test_select_implementation.py index 20dfe9670ab3..95b808d29782 100644 --- a/tests/python/relay/strategy/test_select_implementation.py +++ b/tests/python/relay/strategy/test_select_implementation.py @@ -19,12 +19,14 @@ import pytest import numpy as np +from unittest.mock import MagicMock import tvm from tvm import relay from tvm import te from tvm.relay.testing import run_infer_type import tvm.testing +from tvm import topi @pytest.mark.parametrize( @@ -187,5 +189,38 @@ def test_dense(target, expected_valid_impl, expected_impl): assert selected_impl.name == expected_impl +@pytest.mark.parametrize( + "target,schedule_func", + [ + ("llvm -device=arm_cpu", topi.x86), + ("c -device=arm_cpu -mcpu=cortex-m55", topi.arm_cpu), + ], +) +def test_pool2d(target, schedule_func, monkeypatch): + target = tvm.target.Target(target) + + data_shape = (1, 2, 2, 4) + dtype = "float32" + + out = relay.nn.avg_pool2d(relay.var("data", shape=data_shape, dtype=dtype)) + out = relay.Function(relay.analysis.free_vars(out), out) + out = tvm.IRModule.from_expr(out) + + # Since pool does not use OpStrategy to determine the relevant schedule, + # we cannot simply check the schedule name that was selected with + # `select_implementation`. With this implementation of schedule selection, + # "pool.arm_cpu" will always be the schedule name, regardless of what schedule + # was selected. Instead, this test checks that the relevant schedule function + # is called when building the module. + mock_schedule = MagicMock() + mock_schedule.side_effect = lambda outs, layout: topi.generic.schedule_pool(outs, layout) + monkeypatch.setattr(schedule_func, "schedule_pool", mock_schedule) + + with target: + tvm.relay.build(out, target=target) + + mock_schedule.assert_called() + + if __name__ == "__main__": tvm.testing.main() From b18a0934b18a1269e28e1c25027ae8e8d41eddd0 Mon Sep 17 00:00:00 2001 From: Luke Hutton Date: Tue, 8 Aug 2023 10:21:59 +0000 Subject: [PATCH 2/2] simplify test case Change-Id: I9617ec8ee6f1ca1174c938d9cc4ba21b74216612 --- .../relay/strategy/test_select_implementation.py | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/tests/python/relay/strategy/test_select_implementation.py b/tests/python/relay/strategy/test_select_implementation.py index 95b808d29782..2bf1548d41d8 100644 --- a/tests/python/relay/strategy/test_select_implementation.py +++ b/tests/python/relay/strategy/test_select_implementation.py @@ -203,22 +203,18 @@ def test_pool2d(target, schedule_func, monkeypatch): dtype = "float32" out = relay.nn.avg_pool2d(relay.var("data", shape=data_shape, dtype=dtype)) - out = relay.Function(relay.analysis.free_vars(out), out) - out = tvm.IRModule.from_expr(out) + placeholders = [te.placeholder(data_shape, dtype)] + + mock_schedule = MagicMock() + monkeypatch.setattr(schedule_func, "schedule_pool", mock_schedule) # Since pool does not use OpStrategy to determine the relevant schedule, # we cannot simply check the schedule name that was selected with # `select_implementation`. With this implementation of schedule selection, # "pool.arm_cpu" will always be the schedule name, regardless of what schedule # was selected. Instead, this test checks that the relevant schedule function - # is called when building the module. - mock_schedule = MagicMock() - mock_schedule.side_effect = lambda outs, layout: topi.generic.schedule_pool(outs, layout) - monkeypatch.setattr(schedule_func, "schedule_pool", mock_schedule) - - with target: - tvm.relay.build(out, target=target) - + # is called when selecting the pooling from schedule from arm_cpu. + relay.op.strategy.arm_cpu.schedule_pool_arm_cpu(out.attrs, placeholders, target) mock_schedule.assert_called()