From b47821ab779ee75015f6c04c5b14b506ea95cf1a Mon Sep 17 00:00:00 2001 From: Animesh Jain Date: Tue, 4 Feb 2020 15:25:46 -0800 Subject: [PATCH] [TOPI][x86] Injective schedule improvement (#4786) * [TOPI][x86] Injective Schedule Improvement. * Add tiling. * Vectorize when there is an axis. --- topi/python/topi/x86/injective.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/topi/python/topi/x86/injective.py b/topi/python/topi/x86/injective.py index 8c97214ea4bb..d6bb7622d640 100644 --- 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"])