lint

apache · Feb 13, 2022 · f7de3bf · f7de3bf
1 parent a786bde
commit f7de3bf
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/python/tvm/topi/x86/dense.py b/python/tvm/topi/x86/dense.py
@@ -208,6 +208,7 @@ def _callback(op):
 
 
 def dense_vnni_compute(X, packedW, bias=None):
+    """Compute for uint8 x int8 -> int32 dense"""
     m, k = X.shape
     n_o, _, n_i, _ = packedW.shape
     ak = te.reduce_axis((0, k), name="k")
@@ -231,6 +232,7 @@ def dense_vnni_compute(X, packedW, bias=None):
 
 
 def dense_vnni_schedule(s, C, O):
+    """Schedule dense compute using VNNI vpdpbusd instruction"""
     # C: The output of GEMM
     # O: The output of the fused op
     if C != O:

diff --git a/python/tvm/topi/x86/dense_alter_op.py b/python/tvm/topi/x86/dense_alter_op.py
@@ -39,8 +39,8 @@ def _alter_dense_layout(attrs, inputs, tinfos, out_type):
 
     if (
         target_has_vnni(mcpu)
-        and "uint8" == data_tensor.dtype
-        and "int8" == weight_tensor.dtype
+        and data_tensor.dtype == "uint8"
+        and weight_tensor.dtype == "int8"
         and weight_tensor.shape[0] % 16 == 0
         and weight_tensor.shape[1] % 4 == 0
     ):