Apply comment

apache · Aug 6, 2021 · 479aac5 · 479aac5
1 parent 1a931a1
commit 479aac5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/python/tvm/topi/cuda/conv2d_nhwc.py b/python/tvm/topi/cuda/conv2d_nhwc.py
@@ -54,7 +54,7 @@ def schedule_conv2d_nhwc_direct(cfg, s, Conv):
     cfg.define_knob("vthread_n", [1] if dynamic_batch else [1, 2])
     cfg.define_knob("vthread_c", [1, 2])
     cfg.define_knob("step", [16, 3, 32, 64])
-    cfg.define_knob("vectorize", [4, 2, 8, 16])
+    cfg.define_knob("vectorize", [1, 2, 4, 8, 16])
 
     # fallback support
     target = tvm.target.Target.current()
@@ -97,7 +97,7 @@ def schedule_conv2d_nhwc_direct(cfg, s, Conv):
     ty, ni = s[output].split(ni, factor=tile_n)
     tyz, ty = s[output].split(ty, factor=num_thread_n)
     by, tyz = s[output].split(tyz, factor=vthread_n)
-    s[output].reorder(bz, by, bx, tyz, txz, ty, tx, ni, fi)
+    s[output].reorder(bz, by, bx, tyz, txz, ty, tx, ni, fi, vec)
     s[output].bind(bz, block_z)
     s[output].bind(by, block_y)
     s[output].bind(bx, block_x)