Skip to content

Commit 3ea8358

Browse files
committed
[TOPI][X86] Pool operator parallel support.
1 parent 76c2392 commit 3ea8358

File tree

1 file changed

+6
-0
lines changed

1 file changed

+6
-0
lines changed

topi/python/topi/x86/pooling.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,12 @@ def traverse(OP):
9898
traverse(tensor.op)
9999
# schedule pool
100100
elif OP.tag.startswith('pool'):
101+
# Average pool accumulation and division happens in different for loops (#3607).
102+
# To ensure good parallel support, apply multi-threading on the second loop.
103+
output = outs[0]
104+
output_fused = s[output].fuse(output.op.axis[0], output.op.axis[1])
105+
s[output].parallel(output_fused)
106+
101107
PaddedInput = OP.input_tensors[0]
102108
Pool = OP.output(0)
103109
_schedule(PaddedInput, Pool)

0 commit comments

Comments
 (0)