[TOPI][x86] Injective schedule improvement (apache#4786)

* [TOPI][x86] Injective Schedule Improvement. * Add tiling. * Vectorize when there is an axis.
neo-ai · Mar 2, 2020 · b47821a · b47821a
1 parent cf6739c
commit b47821a
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/topi/python/topi/x86/injective.py b/topi/python/topi/x86/injective.py
@@ -45,6 +45,12 @@ def schedule_injective_from_existing(sch, out):
         sch[out].parallel(fused)
     elif len(sch[out].op.axis) >= 1:
         sch[out].parallel(sch[out].op.axis[0])
+
+    # 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)
+        sch[out].vectorize(li)
     return sch
 
 @generic.schedule_injective.register(["cpu"])