pylint

Author: masahi

python/tvm/tir/tensor_intrin/arm_cpu.py

@@ -14,16 +14,24 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+# pylint: disable=invalid-name
+"""Intrinsics for x86 tensorization."""
 from .. import TensorIntrin
 from tvm.script import tir as T
 
 
+# TODO(masahi): Parametrize the TVMScript description of dot product by
+# shape and dtype, and share the common description with x86.
+
 @T.prim_func
 def dot_product_4x4_i8i8i32_desc(
     A: T.Buffer((4,), "int8", offset_factor=1),
     B: T.Buffer((4, 4), "int8", offset_factor=1),
     C: T.Buffer((4,), "int32", offset_factor=1),
 ) -> None:
+    """
+    A description for 4x4 dot product.
+    """
     with T.block("root"):
         T.reads(C[0:4], A[0:4], B[0:4, 0:4])
         T.writes(C[0:4])
@@ -42,6 +50,9 @@ def dot_product_4x4_i8i8i32_neon(
     B: T.Buffer((4, 4), "int8", offset_factor=1),
     C: T.Buffer((4,), "int32", offset_factor=1),
 ) -> None:
+    """
+    A implementation for 4x4 dot product, applicable for any ARM CPUs supporting NEON.
+    """
     with T.block("root"):
         T.reads(C[0:4], A[0:4], B[0:4, 0:4])
         T.writes(C[0:4])
@@ -103,6 +114,9 @@ def dot_product_4x4_i8i8i32_sdot(
     B: T.Buffer((4, 4), "int8", offset_factor=1),
     C: T.Buffer((4,), "int32", offset_factor=1),
 ) -> None:
+    """
+    A implementation for 4x4 dot product, applicable for ARM CPUs supporting sdot.
+    """
     with T.block("root"):
         T.reads(C[0:4], A[0:4], B[0:4, 0:4])
         T.writes(C[0:4])

python/tvm/tir/tensor_intrin/x86.py

@@ -14,8 +14,10 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-from .. import TensorIntrin
+# pylint: disable=invalid-name
+"""Intrinsics for x86 tensorization."""
 from tvm.script import tir as T
+from .. import TensorIntrin
 
 
 # Tensorized intrinsic description and VNNI-specific implementation.
@@ -28,6 +30,9 @@ def dot_product_16x4_u8i8i32_desc(
     B: T.Buffer((16, 4), "int8", offset_factor=1),
     C: T.Buffer((16,), "int32", offset_factor=1),
 ) -> None:
+    """
+    A description for 16x4 dot product.
+    """
     with T.block("root"):
         T.reads(C[0:16], A[0:4], B[0:16, 0:4])
         T.writes(C[0:16])
@@ -46,6 +51,9 @@ def dot_product_16x4_u8i8i32_vnni(
     B: T.Buffer((16, 4), "int8", offset_factor=1),
     C: T.Buffer((16,), "int32", offset_factor=1),
 ) -> None:
+    """
+    A VNNI-specific implmementation for 16x4 dot product.
+    """
     with T.block("root"):
         T.reads(C[0:16], A[0:4], B[0:16, 0:4])
         T.writes(C[0:16])