From e19c67af6850bc1564a701ca28b7de5491d55259 Mon Sep 17 00:00:00 2001 From: Egor Churaev Date: Mon, 23 Aug 2021 14:54:26 +0300 Subject: [PATCH] Add vectorization to inner loop --- python/tvm/topi/gpu/conv2d_nhwc.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/tvm/topi/gpu/conv2d_nhwc.py b/python/tvm/topi/gpu/conv2d_nhwc.py index a3cf124747167..0aa66845e8978 100644 --- a/python/tvm/topi/gpu/conv2d_nhwc.py +++ b/python/tvm/topi/gpu/conv2d_nhwc.py @@ -110,6 +110,7 @@ def schedule_conv2d_nhwc_direct(cfg, s, Conv): ni, yi, xi, fi = s[OL].op.axis ry, rx, rc = s[OL].op.reduce_axis rco, rci = s[OL].split(rc, factor=step) + s[OL].vectorize(fi) s[OL].reorder(rco, ry, rx, rci, ni, fi) s[AA].compute_at(s[OL], rx)