use functools.partial

Author: masahi

python/tvm/contrib/cutlass/gen_gemm.py

@@ -16,6 +16,7 @@
 # under the License.
 # pylint: disable=invalid-name
 """GEMM kernel generator and profiler for CUTLASS."""
+from functools import partial
 import os
 import re
 import tempfile
@@ -37,7 +38,7 @@
 )
 
 
-def _create_gemm_operator(
+def create_gemm_operator(
     tile_descriptions,
     data_type,
     alignment_constraints,
@@ -132,25 +133,6 @@ def _create_gemm_operator(
     return ret
 
 
-def create_gemm_operator(batched):
-    # TODO: replace with partial
-    def op_creator(
-        tile_descriptions,
-        data_type,
-        alignment_constraints,
-        swizzling_functor=SwizzlingFunctor.Identity8,
-    ):
-        return _create_gemm_operator(
-            tile_descriptions,
-            data_type,
-            alignment_constraints,
-            swizzling_functor,
-            batched=batched,
-        )
-
-    return op_creator
-
-
 GENERATOR_FUNC_TABLE = {
     75: generate_sm75_tensor_op_1688,
     80: generate_sm80_tensor_op_16816,
@@ -193,7 +175,7 @@ def get_default(self, out_dtype, batched=False):
         For now, the default kernel was picked arbitrary.
         """
         ops = GENERATOR_FUNC_TABLE[self.sm](
-            out_dtype, op_creator=create_gemm_operator(batched)
+            out_dtype, op_creator=partial(create_gemm_operator, batched=batched)
         )
         default_kernel_name = DEFAULT_KERNELS[self.sm][out_dtype]
         filtered = list(filter(lambda op: op["name"] == default_kernel_name, ops))
@@ -211,7 +193,7 @@ def profile(
             return self.cache[(M, N, K)]
 
         ops = GENERATOR_FUNC_TABLE[self.sm](
-            out_dtype, op_creator=create_gemm_operator(batched)
+            out_dtype, op_creator=partial(create_gemm_operator, batched=batched)
         )
         ops = list(filter(lambda op: self.check_align(op["name"], M), ops))
 

python/tvm/contrib/cutlass/gen_tensor_op.py

@@ -15,10 +15,9 @@
 # specific language governing permissions and limitations
 # under the License.
 # pylint: disable=invalid-name
-"""GEMM kernel generator and profiler for CUTLASS."""
+"""Common functions and classes for CUTLASS GEMM and Conv2d geneator."""
 import logging
 import os
-import re
 import tempfile
 import subprocess
 import multiprocessing
@@ -56,7 +55,7 @@ def generate_tensor_op_common(
 
 
 def generate_sm75_tensor_op_1688(out_dtype, op_creator):
-    """Generate GEMM kernels for Turing."""
+    """Generate GEMM or Conv2D kernels for Turing."""
     assert out_dtype in ["float32", "float16"]
     math_instructions = {
         "float32": [
@@ -102,7 +101,7 @@ def get_tile_descriptions(math_inst):
 
 
 def generate_sm80_tensor_op_16816(out_dtype, op_creator):
-    """Generate GEMM kernels for Ampere."""
+    """Generate GEMM or Conv2D kernels for Ampere."""
     assert out_dtype in ["float32", "float16"]
     math_instructions = {
         "float32": [