From 3ea835858cdcff572abf7619588ccd3a98a5dfbe Mon Sep 17 00:00:00 2001 From: Animesh Jain Date: Tue, 8 Oct 2019 23:53:06 +0000 Subject: [PATCH] [TOPI][X86] Pool operator parallel support. --- topi/python/topi/x86/pooling.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/topi/python/topi/x86/pooling.py b/topi/python/topi/x86/pooling.py index 9adece2c7bcba..3fc2e48917790 100644 --- 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)