[TOPI][X86] Pool operator parallel support.

apache · Oct 8, 2019 · 3ea8358 · 3ea8358
1 parent 76c2392
commit 3ea8358
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/topi/python/topi/x86/pooling.py b/topi/python/topi/x86/pooling.py
@@ -98,6 +98,12 @@ 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.
+            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)