fixes

python/tvm/topi/cuda/tensorcore_alter_op.py

@@ -148,16 +148,18 @@ def _dense_legalize(attrs, inputs, arg_types):
 
     # Pad input and output channels to use tensorcore schedule.
     if dtype in ["float16", "int8", "uint8"]:
-        # The shape of (M, K, N) must be multiple of (16, 16, 16) or (32, 16, 8) or (8, 16, 32)
+        # The shape of (M, K, N) must be multiple of
+        # (16, 16, 16) or (32, 16, 8) or (8, 16, 32) or (4, 4, 4)
         if (
             (M % 8 == 0 and K % 16 == 0 and N % 32 == 0)
             or (M % 16 == 0 and K % 16 == 0 and N % 16 == 0)
             or (M % 32 == 0 and K % 16 == 0 and N % 8 == 0)
+            or (M % 4 == 0 and K % 4 == 0 and N % 4 == 0)
         ):
             # no need to pad
             return None
 
-        candidates = [(16, 16, 16), (32, 16, 8), (8, 16, 32)]
+        candidates = [(16, 16, 16), (32, 16, 8), (8, 16, 32), (4, 4, 4)]
     elif dtype in ["int4", "uint4"]:
         if M % 8 == 0 and K % 32 == 0 and N % 8 == 0:
             # no need to pad

tests/python/frontend/onnx/test_forward.py

@@ -5170,7 +5170,6 @@ def verify_eyelike(indata):
 @pytest.mark.parametrize("onnx_test", onnx_test_folders)
 @tvm.testing.parametrize_targets
 def test_onnx_nodes(target, dev, onnx_test):
-    # breakpoint()
     target = "cuda"
     dev = tvm.cuda(0)
     target_kind = tvm.target.Target(target).kind.name
@@ -5212,8 +5211,6 @@ def test_onnx_nodes(target, dev, onnx_test):
             else:
                 raise ImportError(str(tensor) + " not labeled as an import or an output")
 
-    breakpoint()
-    print(inputs)
     tvm_val = get_tvm_output_with_vm(onnx_model, inputs, target, dev)
     if len(outputs) == 1:
         tvm.testing.assert_allclose(outputs[0], tvm_val, rtol=rtol, atol=atol)