[GitHub] [tvm] lhutton1 commented on a diff in pull request #15506: [Relay][Strategy] Use x86 pool schedules for arm_cpu

["lhutton1 (via GitHub)" <gi...@apache.org>]

lhutton1 commented on code in PR #15506:
URL: https://github.com/apache/tvm/pull/15506#discussion_r1286859213


##########
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)

Review Comment:
   It's a good point, although I think we could do this as and when needed



##########
tests/python/relay/strategy/test_select_implementation.py:
##########
@@ -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)

Review Comment:
   Yes essentially it was setting a return value for the mocked function.
   
   The return type of `schedule_pool` is a schedule. I took a look at the object but it doesn't really offer anything indicative of an x86 schedule having been used. Also, I don't believe `schedule` can be equality checked very easily. Therefore, I think the simplest way is to check whether the expected function was called.
   
   You did give me an idea that the test could be simplified a little. Rather than check from `relay.build` we can simply check from `schedule_pool_arm_cpu`. This will remove the side-effect nastiness. 



-- 
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-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org