Add vectorization to inner loop

apache · Aug 23, 2021 · e19c67a · e19c67a
1 parent d32193a
commit e19c67a
Showing 1 changed file with 1 addition and 0 deletions.
diff --git 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)