This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-tvm.git
The following commit(s) were added to refs/heads/master by this push:
new ea321df [topi, x86] for 1d loop, make outer loop parallel after split
(#6455)
ea321df is described below
commit ea321df95e8428f265e2d718679bb2b3d6070809
Author: masahi <[email protected]>
AuthorDate: Sat Sep 12 03:32:05 2020 +0900
[topi, x86] for 1d loop, make outer loop parallel after split (#6455)
Co-authored-by: masa <[email protected]>
---
python/tvm/topi/x86/injective.py | 8 +++++++-
1 file changed, 7 insertions(+), 1 deletion(-)
diff --git a/python/tvm/topi/x86/injective.py b/python/tvm/topi/x86/injective.py
index eaf39db..a5e521b 100644
--- a/python/tvm/topi/x86/injective.py
+++ b/python/tvm/topi/x86/injective.py
@@ -47,8 +47,14 @@ def schedule_injective_from_existing(sch, out):
# Vectorize the inner most for loop. Tiling first to get a const extent
if len(sch[out].op.axis) >= 1:
l = sch[out].op.axis[-1]
- _, li = sch[out].split(l, factor=16)
+ lo, li = sch[out].split(l, factor=16)
sch[out].vectorize(li)
+
+ # for 1D loop, the above split will break the parallel axis
+ # Need to make the outer loop parallel again
+ if len(sch[out].op.axis) == 1:
+ sch[out].parallel(lo)
+
return sch
