Improve reduction schedule on arm CPUs (#6110)

* Improve reduction schedule on arm CPUs

Change-Id: I9cd85deac6a57666b82ff7250d827652a4000d82

* Retrigger CI

Change-Id: I5efd99e34268e6bb990904a4b98e1edf2174b26b

diff --git a/python/tvm/relay/op/strategy/arm_cpu.py b/python/tvm/relay/op/strategy/arm_cpu.py
index e639e22..8457965 100644 (file)
--- a/python/tvm/relay/op/strategy/arm_cpu.py
+++ b/python/tvm/relay/op/strategy/arm_cpu.py
@@ -26,6 +26,12 @@ from .. import op as _op
 
 logger = logging.getLogger('strategy')
 
+@schedule_reduce.register("arm_cpu")
+def schedule_reduce_cpu(attrs, outs, target):
+    """schedule reduction ops for arm_cpu"""
+    with target:
+        return topi.x86.schedule_reduce(outs)
+
 @schedule_injective.register(["arm_cpu", "micro_dev"])
 def schedule_injective_arm_cpu(_, outs, target):
     """schedule injective ops for arm cpu"""