[TOPI][X86] Pool operator parallel support. (#4090)

apache · Oct 9, 2019 · 3a32729 · 3a32729
1 parent 1c56c72
commit 3a32729
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/topi/python/topi/x86/pooling.py b/topi/python/topi/x86/pooling.py
@@ -98,6 +98,13 @@ def traverse(OP):
                     traverse(tensor.op)
         # schedule pool
         elif OP.tag.startswith('pool'):
+            # Average pool accumulation and division happens in different for loops (#3607).
+            # To ensure good parallel support, apply multi-threading on the second loop.
+            if OP != outs[0].op:
+                output = outs[0]
+                output_fused = s[output].fuse(output.op.axis[0], output.op.axis[1])
+                s[output].parallel(output_fused)
+
             PaddedInput = OP.input_tensors[0]
             Pool = OP.output(0)
             _schedule(PaddedInput, Pool)