black

Author: masahi

python/tvm/contrib/cutlass/gen_gemm.py

@@ -89,8 +89,12 @@ def enumerate_gemm_operators(
             C = TensorDescription(element_c, LayoutType.RowMajor, alignment)
 
             if element_c == DataType.s32 and A.alignment == 1:
-                tile_description.threadblock_shape[0] = min(tile_description.threadblock_shape[0], 128)
-                tile_description.threadblock_shape[1] = min(tile_description.threadblock_shape[1], 128)
+                tile_description.threadblock_shape[0] = min(
+                    tile_description.threadblock_shape[0], 128
+                )
+                tile_description.threadblock_shape[1] = min(
+                    tile_description.threadblock_shape[1], 128
+                )
 
             op = GemmOperation(
                 tile_description.minimum_compute_capability,

python/tvm/relay/op/strategy/cuda.py

@@ -872,18 +872,8 @@ def dense_strategy_cuda(attrs, inputs, out_type, target):
                         or (i % 16 == 0 and b % 32 == 0 and o % 8 == 0)
                     )
                 )
-                or (
-                    data.dtype in ["int4", "uint4"]
-                    and i % 32 == 0
-                    and b % 8 == 0
-                    and o % 8 == 0
-                )
-                or (
-                    data.dtype in ["int1", "uint1"]
-                    and i % 128 == 0
-                    and b % 8 == 0
-                    and o % 8 == 0
-                )
+                or (data.dtype in ["int4", "uint4"] and i % 32 == 0 and b % 8 == 0 and o % 8 == 0)
+                or (data.dtype in ["int1", "uint1"] and i % 128 == 0 and b % 8 == 0 and o % 8 == 0)
             ):
                 strategy.add_implementation(
                     wrap_compute_dense(topi.cuda.dense_tensorcore),

tests/python/contrib/test_cublas.py

@@ -169,6 +169,7 @@ def test_batch_matmul():
 
     verify_batch_matmul((16, 1024, 128), (16, 128, 236), (16, 1024, 236), "int8", "int32")
 
+
 if __name__ == "__main__":
     test_matmul_add()
     test_batch_matmul()

tests/python/contrib/test_cutlass.py

@@ -742,7 +742,7 @@ def test_conv2d():
         rtol=1e-5,
         ref_target="llvm",
         data_dtype="uint8",
-        weight_dtype="int8"
+        weight_dtype="int8",
     )