chore: import upstream snapshot with attribution

python/tvm/contrib/__init__.py

@@ -0,0 +1,22 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Contrib APIs of TVM python package.
+
+Contrib API provides many useful not core features.
+Some of these are useful utilities to interact with
+thirdparty libraries and tools.
+"""

python/tvm/contrib/cblas.py

@@ -0,0 +1,94 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""External function interface to BLAS libraries."""
+
+import tvm
+from tvm import te
+
+
+def matmul(lhs, rhs, transa=False, transb=False, **kwargs):
+    """Create an extern op that compute matrix mult of A and rhs with CrhsLAS
+    This function serves as an example on how to call external libraries.
+
+    Parameters
+    ----------
+    lhs: Tensor
+        The left matrix operand
+    rhs: Tensor
+        The right matrix operand
+    transa: bool
+        Whether transpose lhs
+    transb: bool
+        Whether transpose rhs
+
+    Returns
+    -------
+    C: Tensor
+        The result tensor.
+    """
+    n = lhs.shape[1] if transa else lhs.shape[0]
+    m = rhs.shape[0] if transb else rhs.shape[1]
+    return te.extern(
+        (n, m),
+        [lhs, rhs],
+        lambda ins, outs: tvm.tirx.call_packed(
+            "tvm.contrib.cblas.matmul", ins[0], ins[1], outs[0], transa, transb
+        ),
+        name="C",
+        **kwargs,
+    )
+
+
+def batch_matmul(lhs, rhs, transa=False, transb=False, iterative=False, **kwargs):
+    """Create an extern op that compute batched matrix mult of A and rhs with CBLAS
+    This function serves as an example on how to call external libraries.
+
+    Parameters
+    ----------
+    lhs: Tensor
+        The left matrix operand
+    rhs: Tensor
+        The right matrix operand
+    transa: bool
+        Whether transpose lhs
+    transb: bool
+        Whether transpose rhs
+
+    Returns
+    -------
+    C: Tensor
+        The result tensor.
+    """
+    b = te.max(lhs.shape[0], rhs.shape[0])
+    n = lhs.shape[2] if transa else lhs.shape[1]
+    m = rhs.shape[1] if transb else rhs.shape[2]
+    return te.extern(
+        (b, n, m),
+        [lhs, rhs],
+        lambda ins, outs: tvm.tirx.call_packed(
+            "tvm.contrib.cblas.batch_matmul"
+            if not iterative
+            else "tvm.contrib.cblas.batch_matmul_iterative",
+            ins[0],
+            ins[1],
+            outs[0],
+            transa,
+            transb,
+        ),
+        name="C",
+        **kwargs,
+    )

python/tvm/contrib/coreml_runtime.py

@@ -0,0 +1,76 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""CoreML runtime that load and run coreml models."""
+
+import tvm_ffi
+
+from ..rpc import base as rpc_base
+
+
+def create(symbol, compiled_model_path, device):
+    """Create a runtime executor module given a coreml model and context.
+
+    Parameters
+    ----------
+    symbol : str
+        The symbol that represents the Core ML model.
+    compiled_model_path : str
+        The path of the compiled model to be deployed.
+    device : Device
+        The device to deploy the module. It can be local or remote when there
+        is only one Device.
+    Returns
+    -------
+    coreml_runtime : CoreMLModule
+        Runtime coreml module that can be used to execute the coreml model.
+    """
+    device_type = device.dlpack_device_type()
+    runtime_func = "tvm.coreml_runtime.create"
+
+    if device_type >= rpc_base.RPC_SESS_MASK:
+        fcreate = device._rpc_sess.get_function(runtime_func)
+    else:
+        fcreate = tvm_ffi.get_global_func(runtime_func)
+        assert fcreate, "Cannot find `tvm.coreml_runtime.create` function."
+
+    return CoreMLModule(fcreate(symbol, compiled_model_path))
+
+
+class CoreMLModule:
+    """Wrapper runtime module.
+
+    This is a thin wrapper of the underlying TVM module.
+    you can also directly call set_input, run, and get_output
+    of underlying module functions
+
+    Parameters
+    ----------
+    module : Module
+        The internal tvm module that holds the actual coreml functions.
+
+    Attributes
+    ----------
+    module : Module
+        The internal tvm module that holds the actual coreml functions.
+    """
+
+    def __init__(self, module):
+        self.module = module
+        self.invoke = module["invoke"]
+        self.set_input = module["set_input"]
+        self.get_output = module["get_output"]
+        self.get_num_outputs = module["get_num_outputs"]

python/tvm/contrib/cublas.py

@@ -0,0 +1,87 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""External function interface to cuBLAS libraries."""
+
+import tvm
+from tvm import te
+
+
+def matmul(lhs, rhs, transa=False, transb=False, dtype=None):
+    """Create an extern op that compute matrix mult of A and rhs with cuBLAS
+
+    Parameters
+    ----------
+    lhs : Tensor
+        The left matrix operand
+    rhs : Tensor
+        The right matrix operand
+    transa : bool
+        Whether transpose lhs
+    transb : bool
+        Whether transpose rhs
+
+    Returns
+    -------
+    C : Tensor
+        The result tensor.
+    """
+    n = lhs.shape[1] if transa else lhs.shape[0]
+    m = rhs.shape[0] if transb else rhs.shape[1]
+    dtype = dtype if dtype is not None else lhs.dtype
+    return te.extern(
+        (n, m),
+        [lhs, rhs],
+        lambda ins, outs: tvm.tirx.call_packed(
+            "tvm.contrib.cublas.matmul", ins[0], ins[1], outs[0], transa, transb
+        ),
+        dtype=dtype,
+        name="matmul_cublas",
+    )
+
+
+def batch_matmul(lhs, rhs, transa=False, transb=False, dtype=None):
+    """Create an extern op that compute batch matrix mult of A and rhs with cuBLAS
+
+    Parameters
+    ----------
+    lhs : Tensor
+        The left matrix operand
+    rhs : Tensor
+        The right matrix operand
+    transa : bool
+        Whether transpose lhs
+    transb : bool
+        Whether transpose rhs
+
+    Returns
+    -------
+    C : Tensor
+        The result tensor.
+    """
+    b = lhs.shape[0]
+    n = lhs.shape[2] if transa else lhs.shape[1]
+    m = rhs.shape[1] if transb else rhs.shape[2]
+    dtype = dtype if dtype is not None else lhs.dtype
+    return te.extern(
+        (b, n, m),
+        [lhs, rhs],
+        lambda ins, outs: tvm.tirx.call_packed(
+            "tvm.contrib.cublas.batch_matmul", ins[0], ins[1], outs[0], transa, transb
+        ),
+        dtype=dtype,
+        name="batch_matmul_cublas",
+    )

python/tvm/contrib/cublaslt.py

@@ -0,0 +1,55 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""External function interface to cuBLASlt libraries."""
+
+import tvm
+from tvm import te
+
+
+def matmul(lhs, rhs, transa=False, transb=False, n=0, m=0, dtype=None):
+    """Create an extern op that compute matrix mult of A and rhs with cuBLAS
+
+    Parameters
+    ----------
+    lhs : Tensor
+        The left matrix operand
+    rhs : Tensor
+        The right matrix operand
+    transa : bool
+        Whether transpose lhs
+    transb : bool
+        Whether transpose rhs
+
+    Returns
+    -------
+    C : Tensor
+        The result tensor.
+    """
+    if n == 0:
+        n = lhs.shape[1] if transa else lhs.shape[0]
+    if m == 0:
+        m = rhs.shape[0] if transb else rhs.shape[1]
+    dtype = dtype if dtype is not None else lhs.dtype
+    return te.extern(
+        (n, m),
+        [lhs, rhs],
+        lambda ins, outs: tvm.tirx.call_packed(
+            "tvm.contrib.cublaslt.matmul", ins[0], ins[1], outs[0], transa, transb
+        ),
+        dtype=dtype,
+        name="C",
+    )

python/tvm/contrib/cudnn.py

@@ -0,0 +1,926 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F821
+"""External function interface to CuDNN v7 library."""
+
+# pylint: disable-msg=C0103
+import ctypes
+
+import numpy as np
+import tvm_ffi
+
+import tvm
+from tvm import te
+
+# algos can be read from cudnn.h
+_FWD_ALGOS = [
+    "CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_GEMM",
+    "CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM",
+    "CUDNN_CONVOLUTION_FWD_ALGO_GEMM",
+    "CUDNN_CONVOLUTION_FWD_ALGO_DIRECT",
+    "CUDNN_CONVOLUTION_FWD_ALGO_FFT",
+    "CUDNN_CONVOLUTION_FWD_ALGO_FFT_TILING",
+    "CUDNN_CONVOLUTION_FWD_ALGO_WINOGRAD",
+    "CUDNN_CONVOLUTION_FWD_ALGO_WINOGRAD_NONFUSED",
+    "CUDNN_CONVOLUTION_FWD_ALGO_COUNT",
+]
+
+
+def exists():
+    """
+    Checks whether the local machine can use CuDNN.
+
+    Returns
+    -------
+        exists: bool
+
+            True if CuDNN support is enabled and a CuDNN-capable GPU
+            exists.  Otherwise, False.
+    """
+    func = tvm.get_global_func("tvm.contrib.cudnn.exists", allow_missing=True)
+    if func is None:
+        return False
+
+    return bool(func())
+
+
+def algo_to_index(algo_type, algo_name):
+    """Return a index represents the algorithm, which can be used in
+    calling CuDNN function
+
+    Parameters
+    ----------
+    algo_type : str
+        One of ``"fwd"``, ``"bwd_filter"``, or ``"bwd_data"``.
+
+    algo_name : str
+        Algorithm name as defined in cuDNN. For example:
+
+        * fwd: ``CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_GEMM``, etc.
+        * bwd_filter: ``CUDNN_CONVOLUTION_BWD_FILTER_ALGO_0``, etc.
+        * bwd_data: ``CUDNN_CONVOLUTION_BWD_DATA_ALGO_0``, etc.
+
+    Returns
+    -------
+    algo : int
+        Algorithm index
+
+    """
+    idx = -1
+    if algo_type == "fwd":
+        idx = _FWD_ALGOS.index(algo_name)
+    elif algo_type == "bwd_filter":
+        idx = _BWD_FILTER_ALGOS.index(algo_name)
+    elif algo_type == "bwd_data":
+        idx = _BWD_DATA_ALGOS.index(algo_name)
+    assert idx >= 0
+    return idx
+
+
+def _get_np_int32_array_handle(arr):
+    """Return a void_p handle for a numpy array
+
+    Parameters
+    ----------
+    arr: numpy.Tensor
+        source numpy array
+
+    Returns
+    -------
+    ptr:  ctypes.c_void_p
+        pointer to the data
+    """
+    assert arr.dtype == np.int32
+    ptr = arr.ctypes.data_as(ctypes.POINTER(ctypes.c_int32))
+    return ctypes.cast(ptr, ctypes.c_void_p)
+
+
+def _prepare_global_func_params(dims, pad, stride, dilation, x_shape=None, w_shape=None):
+    full_dims = dims + 2
+    if x_shape:
+        assert isinstance(x_shape, list)
+        assert len(x_shape) == full_dims
+    if w_shape:
+        assert isinstance(w_shape, list)
+        assert len(w_shape) == full_dims
+
+    pad = (
+        np.full(dims, pad, dtype=np.int32)
+        if isinstance(pad, int)
+        else np.array(pad, dtype=np.int32)
+    )
+    stride = (
+        np.full(dims, stride, dtype=np.int32)
+        if isinstance(stride, int)
+        else np.array(stride, dtype=np.int32)
+    )
+    dilation = (
+        np.full(dims, dilation, dtype=np.int32)
+        if isinstance(dilation, int)
+        else np.array(dilation, dtype=np.int32)
+    )
+
+    xshape = np.array(x_shape, dtype=np.int32) if x_shape else None
+    wshape = np.array(w_shape, dtype=np.int32) if x_shape else None
+
+    return pad, stride, dilation, xshape, wshape
+
+
+def conv_output_shape(
+    tensor_format, pad, stride, dilation, x_shape, w_shape, data_dtype, conv_dtype, groups=1
+):
+    """Get output shape of 2D or 3D convolution
+
+    Paramters
+    ---------
+    tensor_format: int
+        0: CUDNN_TENSOR_NCHW
+        1: CUDNN_TENSOR_NHWC
+        2: CUDNN_TENSOR_NCHW_VECT_C
+    pad: int or list
+        padding
+    stride: int or list
+        stride
+    dilation: int or list
+        dilation
+    x_shape: list
+        input shape
+    w_shape: list
+        weight shape
+    data_dtype: str
+        data type
+    conv_dtype: str
+        convolution type
+    groups: int
+        number of groups
+
+    Returns
+    -------
+    oshape: list
+        output shape
+    """
+
+    assert len(x_shape) == len(w_shape)
+    assert len(x_shape) in (4, 5)
+
+    if tensor_format == 0:
+        n_output = x_shape[0]
+        c_output = w_shape[0]
+        x_chan = x_shape[1]
+        w_chan_input = w_shape[1]
+        x_shape = x_shape[2:]
+        w_shape = w_shape[2:]
+
+    elif tensor_format == 1:
+        n_output = x_shape[0]
+        c_output = w_shape[0]
+        x_chan = x_shape[-1]
+        w_chan_input = w_shape[-1]
+        assert len(x_shape) == 4, "CuDNN layout NHWC is only well-defined for 4d tensors"
+        x_shape = x_shape[1:-1]
+        w_shape = w_shape[1:-1]
+
+    elif tensor_format == 2:
+        n_output = x_shape[0]
+        c_output = w_shape[0]
+        x_chan = x_shape[1]
+        w_chan_input = w_shape[1]
+        w_lanes = tvm.runtime.DataType(conv_dtype).lanes
+        assert w_lanes == 1
+        x_shape = x_shape[2:]
+        w_shape = w_shape[2:]
+
+    else:
+        raise ValueError(f"Unknown CuDNN tensor format: '{tensor_format}'")
+
+    x_lanes = tvm.runtime.DataType(data_dtype).lanes
+    assert x_chan * x_lanes == w_chan_input * groups, (
+        f"Mismatched dimensions, data has {x_chan // groups} channels/group "
+        f"(dimension {x_chan} with {x_lanes} lanes/value, {groups} groups), "
+        f"but weights require {w_chan_input} input channels/group"
+    )
+
+    output_dims = []
+    for x_shape_i, w_shape_i, pad_i, stride_i, dilation_i in zip(
+        x_shape, w_shape, pad, stride, dilation
+    ):
+        output_dim = 1 + (x_shape_i + 2 * pad_i - (((w_shape_i - 1) * dilation_i) + 1)) // stride_i
+        output_dims.append(output_dim)
+
+    if tensor_format in [0, 2]:
+        output = [n_output, c_output, *output_dims]
+    elif tensor_format == 1:
+        output = [n_output, *output_dims, c_output]
+    else:
+        raise ValueError(f"Unknown CuDNN tensor format: '{tensor_format}'")
+
+    return output
+
+
+def conv_dgrad_shape(
+    tensor_format, pad, stride, dilation, dy_shape, w_shape, output_padding=(0, 0), groups=1
+):
+    """Get output shape of conv2d gradient with respect to data
+
+    Paramters
+    ---------
+    tensor_format: int
+        0: CUDNN_TENSOR_NCHW
+        1: CUDNN_TENSOR_NHWC
+    pad: int or list
+        padding
+    stride: int or list
+        stride
+    dilation: int or list
+        dilation
+    dy_shape: list
+        output gradient shape
+    w_shape: list
+        weight shape
+    data_dtype: str
+        data type
+    conv_dtype: str
+        convolution type
+    groups: int
+        number of groups
+
+    Returns
+    -------
+    oshape: list
+        output shape
+    """
+
+    assert len(dy_shape) == len(w_shape)
+    assert len(dy_shape) == 4
+
+    if tensor_format == 0:
+        N = dy_shape[0]
+        C = w_shape[1] * groups
+        dy_shape = dy_shape[2:]
+        w_shape = w_shape[2:]
+    elif tensor_format == 1:
+        N = dy_shape[0]
+        C = w_shape[-1] * groups
+        dy_shape = dy_shape[1:-1]
+        w_shape = w_shape[1:-1]
+    else:
+        raise ValueError(f"Unsupported CuDNN tensor format: '{tensor_format}'")
+
+    input_dims = []
+    for dy_shape_i, w_shape_i, pad_i, stride_i, dilation_i, out_pad in zip(
+        dy_shape, w_shape, pad, stride, dilation, output_padding
+    ):
+        input_dim = (
+            (dy_shape_i - 1) * stride_i - 2 * pad_i + (((w_shape_i - 1) * dilation_i) + 1) + out_pad
+        )
+        input_dims.append(input_dim)
+
+    if tensor_format == 0:
+        output = [N, C, *input_dims]
+    else:
+        output = [N, *input_dims, C]
+
+    return output
+
+
+def _conv_find_algo(
+    func_name,
+    tensor_format,
+    pad,
+    stride,
+    dilation,
+    x_shape,
+    w_shape,
+    y_shape,
+    data_dtype,
+    conv_dtype,
+    groups=1,
+    verbose=False,
+):
+    """
+    Common function to choose the best cudnn convolution algorithm for the given input
+    and the convolution type.
+    """
+    dims = len(x_shape)
+    assert dims in (4, 5)
+
+    pad, stride, dilation, xshape, wshape = _prepare_global_func_params(
+        dims - 2, pad, stride, dilation, x_shape, w_shape
+    )
+    yshape = np.array(y_shape, dtype=np.int32)
+    func = tvm_ffi.get_global_func(func_name)
+    return func(
+        tensor_format,
+        dims - 2,
+        _get_np_int32_array_handle(pad),
+        _get_np_int32_array_handle(stride),
+        _get_np_int32_array_handle(dilation),
+        _get_np_int32_array_handle(xshape),
+        _get_np_int32_array_handle(wshape),
+        _get_np_int32_array_handle(yshape),
+        data_dtype,
+        conv_dtype,
+        groups,
+        verbose,
+    )
+
+
+def conv_forward_find_algo(
+    tensor_format,
+    pad,
+    stride,
+    dilation,
+    x_shape,
+    w_shape,
+    y_shape,
+    data_dtype,
+    conv_dtype,
+    groups=1,
+    verbose=True,
+):
+    """Choose the best forward algorithm for the given input.
+
+    Paramters
+    ---------
+    tensor_format: int
+        0: CUDNN_TENSOR_NCHW
+        1: CUDNN_TENSOR_NHWC
+        2: CUDNN_TENSOR_NCHW_VECT_C
+    pad: int or list
+        padding
+    stride: int or list
+        stride
+    dilation: int or list
+        dilation
+    x_shape: list
+        input shape
+    w_shape: list
+        weight shape
+    y_shape: list
+        output shape
+    data_dtype: str
+        data type
+    conv_dtype: str
+        convolution type
+    groups: int
+        number of groups
+
+    Returns
+    -------
+    algo: int
+        algo chosen by CUDNN
+    """
+    return _conv_find_algo(
+        "tvm.contrib.cudnn.conv.forward_find_algo",
+        tensor_format,
+        pad,
+        stride,
+        dilation,
+        x_shape,
+        w_shape,
+        y_shape,
+        data_dtype,
+        conv_dtype,
+        groups,
+        verbose,
+    )
+
+
+def conv_backward_data_find_algo(
+    tensor_format,
+    pad,
+    stride,
+    dilation,
+    dy_shape,
+    w_shape,
+    dx_shape,
+    data_dtype,
+    conv_dtype,
+    groups=1,
+    verbose=True,
+):
+    """Choose the best backward data algorithm for the given input.
+
+    Paramters
+    ---------
+    tensor_format: int
+        0: CUDNN_TENSOR_NCHW
+        1: CUDNN_TENSOR_NHWC
+        2: CUDNN_TENSOR_NCHW_VECT_C
+    pad: int or list
+        padding
+    stride: int or list
+        stride
+    dilation: int or list
+        dilation
+    dy_shape: list
+        output gradient shape
+    w_shape: list
+        weight shape
+    dx_shape: list
+        dgrad shape
+    data_dtype: str
+        data type
+    conv_dtype: str
+        convolution type
+    groups: int
+        number of groups
+    verbose: bool
+        whether to show the selection trials
+
+    Returns
+    -------
+    algo: int
+        algo chosen by CUDNN
+    """
+    return _conv_find_algo(
+        "tvm.contrib.cudnn.conv.backward_data_find_algo",
+        tensor_format,
+        pad,
+        stride,
+        dilation,
+        dy_shape,
+        w_shape,
+        dx_shape,
+        data_dtype,
+        conv_dtype,
+        groups,
+        verbose,
+    )
+
+
+def conv_backward_filter_find_algo(
+    tensor_format,
+    pad,
+    stride,
+    dilation,
+    dy_shape,
+    x_shape,
+    dw_shape,
+    data_dtype,
+    conv_dtype,
+    groups=1,
+    verbose=True,
+):
+    """Choose the best backward filter algorithm for the given input.
+
+    Paramters
+    ---------
+    tensor_format: int
+        0: CUDNN_TENSOR_NCHW
+        1: CUDNN_TENSOR_NHWC
+        2: CUDNN_TENSOR_NCHW_VECT_C
+    pad: int or list
+        padding
+    stride: int or list
+        stride
+    dilation: int or list
+        dilation
+    dy_shape: list
+        output gradient shape
+    x_shape: list
+        weight shape
+    dw_shape: list
+        wgrad shape
+    data_dtype: str
+        data type
+    conv_dtype: str
+        convolution type
+    groups: int
+        number of groups
+    verbose: bool
+        whether to show the selection trials
+
+    Returns
+    -------
+    algo: int
+        algo chosen by CUDNN
+    """
+    return _conv_find_algo(
+        "tvm.contrib.cudnn.conv.backward_filter_find_algo",
+        tensor_format,
+        pad,
+        stride,
+        dilation,
+        dy_shape,
+        x_shape,
+        dw_shape,
+        data_dtype,
+        conv_dtype,
+        groups,
+        verbose,
+    )
+
+
+def conv_forward(
+    x, w, pad, stride, dilation, conv_mode, tensor_format, algo, conv_dtype, groups=1, verbose=True
+):
+    """Create an extern op that compute 2D or 3D convolution with CuDNN
+
+    Parameters
+    ----------
+    x: Tensor
+        input feature map
+    w: Tensor
+        convolution weight
+    pad: int or list
+        padding
+    stride: int or list
+        stride
+    dilation: int or list
+        dilation
+    conv_mode: int
+        0: CUDNN_CONVOLUTION
+        1: CUDNN_CROSS_CORRELATION
+    tensor_format: int
+        0: CUDNN_TENSOR_NCHW
+        1: CUDNN_TENSOR_NHWC
+        2: CUDNN_TENSOR_NCHW_VECT_C
+    algo: int
+        Forward algorithm, get index from ```algo_to_index``` function
+        if algo == -1, the best algo will be chosen by CUDNN
+    conv_dtype: str
+        convolution type
+    groups: int
+        the number of groups
+    verbose: bool
+        whether to show the selection trials
+
+    Returns
+    -------
+    y: Tensor
+        The result tensor
+    """
+    dims = len(x.shape)
+    assert dims in (4, 5)
+
+    conv_dtype = x.dtype if conv_dtype is None else conv_dtype
+    pad, stride, dilation, _, _ = _prepare_global_func_params(dims - 2, pad, stride, dilation)
+
+    x_shape = list(x.shape)
+
+    if isinstance(x.shape[0], tvm.tirx.expr.IntImm):
+        oshape = conv_output_shape(
+            tensor_format,
+            pad,
+            stride,
+            dilation,
+            x_shape,
+            list(w.shape),
+            x.dtype,
+            conv_dtype,
+            groups,
+        )
+        if algo == -1:
+            # For now if we try to call `cudnnFindConvolutionForwardAlgorithm` when
+            # using INT8 data type, CuDNN will crash down.
+            # On the other hand, CuDNN only support IMPLICIT_PRECOMP_GEMM at NHWC format
+            if tensor_format == 1 and conv_dtype == "int32":
+                algo = 1
+            else:
+                algo = conv_forward_find_algo(
+                    tensor_format,
+                    pad,
+                    stride,
+                    dilation,
+                    list(x.shape),
+                    list(w.shape),
+                    oshape,
+                    x.dtype,
+                    conv_dtype,
+                    groups,
+                    verbose,
+                )
+    else:
+        # The dynamic batch size case, pretend this is a single batch
+        x_shape[0] = 1
+        oshape = conv_output_shape(
+            tensor_format,
+            pad,
+            stride,
+            dilation,
+            x_shape,
+            list(w.shape),
+            x.dtype,
+            conv_dtype,
+            groups,
+        )
+        oshape[0] = x.shape[0]
+        # This picks CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM
+        # It seems this is the fastest among algorithms that are always applicable
+        algo = 1
+
+    if dims == 4:
+        return te.extern(
+            oshape,
+            [x, w],
+            lambda ins, outs: tvm.tirx.call_packed(
+                "tvm.contrib.cudnn.conv2d.forward",
+                conv_mode,
+                tensor_format,
+                algo,
+                pad[0],
+                pad[1],
+                stride[0],
+                stride[1],
+                dilation[0],
+                dilation[1],
+                ins[0],
+                ins[1],
+                outs[0],
+                conv_dtype,
+                groups,
+            ),
+            name="y",
+        )
+
+    return te.extern(
+        oshape,
+        [x, w],
+        lambda ins, outs: tvm.tirx.call_packed(
+            "tvm.contrib.cudnn.conv3d.forward",
+            conv_mode,
+            tensor_format,
+            algo,
+            pad[0],
+            pad[1],
+            pad[2],
+            stride[0],
+            stride[1],
+            stride[2],
+            dilation[0],
+            dilation[1],
+            dilation[2],
+            ins[0],
+            ins[1],
+            outs[0],
+            conv_dtype,
+            groups,
+        ),
+        name="y",
+    )
+
+
+def conv_backward_data(
+    dy,
+    w,
+    pad,
+    stride,
+    dilation,
+    conv_mode,
+    tensor_format,
+    conv_dtype,
+    groups=1,
+    output_padding=(0, 0),
+):
+    """Create a CuDNN extern op that computes the gradient of 2D convolution with respect to data.
+
+    Parameters
+    ----------
+    dy: Tensor
+        output gradient
+    w: Tensor
+        convolution weight
+    pad: int or list
+        padding
+    stride: int or list
+        stride
+    dilation: int or list
+        dilation
+    conv_mode: int
+        0: CUDNN_CONVOLUTION
+        1: CUDNN_CROSS_CORRELATION
+    tensor_format: int
+        0: CUDNN_TENSOR_NCHW
+        1: CUDNN_TENSOR_NHWC
+    conv_dtype: str
+        convolution type
+    groups: int
+        the number of groups
+
+    Returns
+    -------
+    dx: Tensor
+        dgrad tensor
+    """
+    dims = len(dy.shape)
+    assert dims == 4
+
+    conv_dtype = dy.dtype if conv_dtype is None else conv_dtype
+    pad, stride, dilation, _, _ = _prepare_global_func_params(dims - 2, pad, stride, dilation)
+
+    assert isinstance(dy.shape[0], tvm.tirx.expr.IntImm), (
+        "Dynamic batch is not supported for cudnn conv2d backwad data yet."
+    )
+
+    dx_shape = conv_dgrad_shape(
+        tensor_format, pad, stride, dilation, dy.shape, w.shape, output_padding, groups
+    )
+
+    if exists():
+        # When cudnn exists, find the backward data algo
+        algo = conv_backward_data_find_algo(
+            tensor_format,
+            pad,
+            stride,
+            dilation,
+            list(dy.shape),
+            list(w.shape),
+            dx_shape,
+            dy.dtype,
+            conv_dtype,
+            groups,
+            True,
+        )
+    else:
+        algo = 1
+
+    return te.extern(
+        dx_shape,
+        [dy, w],
+        lambda ins, outs: tvm.tirx.call_packed(
+            "tvm.contrib.cudnn.conv2d.backward_data",
+            conv_mode,
+            tensor_format,
+            algo,
+            pad[0],
+            pad[1],
+            stride[0],
+            stride[1],
+            dilation[0],
+            dilation[1],
+            ins[0],
+            ins[1],
+            outs[0],
+            conv_dtype,
+            groups,
+        ),
+        name="dx",
+    )
+
+
+def conv_backward_filter(
+    dy, x, kernel_size, pad, stride, dilation, conv_mode, tensor_format, conv_dtype, groups=1
+):
+    """Create a CuDNN extern op that computes the gradient of 2D convolution with respect to weight.
+
+    Parameters
+    ----------
+    dy: Tensor
+        output gradient
+    x: Tensor
+        input tensor
+    kernel_size: a pair of int
+        The spatial size of the corresponding forward convolution kernel
+    pad: int or list
+        padding
+    stride: int or list
+        stride
+    dilation: int or list
+        dilation
+    conv_mode: int
+        0: CUDNN_CONVOLUTION
+        1: CUDNN_CROSS_CORRELATION
+    tensor_format: int
+        0: CUDNN_TENSOR_NCHW
+        1: CUDNN_TENSOR_NHWC
+    conv_dtype: str
+        convolution type
+    groups: int
+        the number of groups
+
+    Returns
+    -------
+    dw: Tensor
+        wgrad tensor
+    """
+    dims = len(x.shape)
+    assert dims == 4
+
+    conv_dtype = x.dtype if conv_dtype is None else conv_dtype
+    pad, stride, dilation, _, _ = _prepare_global_func_params(dims - 2, pad, stride, dilation)
+    filter_h, filter_w = kernel_size
+
+    x_shape = list(x.shape)
+
+    assert isinstance(x.shape[0], tvm.tirx.expr.IntImm), (
+        "Dynamic batch is not supported for cudnn conv2d backwad filter yet."
+    )
+
+    ic_ind = 1 if tensor_format == 0 else 3
+
+    if groups > 1:
+        assert x_shape[ic_ind] == dy.shape[ic_ind] and x_shape[ic_ind] == groups, (
+            "Only depthwise wgrad supported for groups > 1."
+        )
+        ic = 1
+    else:
+        ic = x_shape[ic_ind]
+
+    if tensor_format == 0:
+        dw_shape = [dy.shape[1], ic, filter_h, filter_w]
+    else:
+        dw_shape = [dy.shape[3], filter_h, filter_w, ic]
+
+    algo = conv_backward_filter_find_algo(
+        tensor_format,
+        pad,
+        stride,
+        dilation,
+        list(dy.shape),
+        list(x.shape),
+        dw_shape,
+        x.dtype,
+        conv_dtype,
+        groups,
+        True,
+    )
+
+    return te.extern(
+        dw_shape,
+        [dy, x],
+        lambda ins, outs: tvm.tirx.call_packed(
+            "tvm.contrib.cudnn.conv2d.backward_filter",
+            conv_mode,
+            tensor_format,
+            algo,
+            pad[0],
+            pad[1],
+            stride[0],
+            stride[1],
+            dilation[0],
+            dilation[1],
+            ins[0],
+            ins[1],
+            outs[0],
+            conv_dtype,
+            groups,
+        ),
+        name="dw",
+    )
+
+
+def softmax(x, axis=-1):
+    """Compute softmax using CuDNN
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        The input tensor
+
+    axis : int
+        The axis to compute the softmax
+
+    Returns
+    -------
+    ret : tvm.te.Tensor
+        The result tensor
+    """
+    return te.extern(
+        x.shape,
+        [x],
+        lambda ins, outs: tvm.tirx.call_packed(
+            "tvm.contrib.cudnn.softmax.forward", ins[0], outs[0], axis
+        ),
+        name="y",
+    )
+
+
+def log_softmax(x, axis=-1):
+    """Compute log_softmax using CuDNN
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        The input tensor
+
+    axis : int
+        The axis to compute log softmax over
+
+    Returns
+    -------
+    ret : tvm.te.Tensor
+        The result tensor
+    """
+    return te.extern(
+        x.shape,
+        [x],
+        lambda ins, outs: tvm.tirx.call_packed(
+            "tvm.contrib.cudnn.log_softmax.forward", ins[0], outs[0], axis
+        ),
+        name="y",
+    )

python/tvm/contrib/cutlass/__init__.py

@@ -0,0 +1,20 @@
+# isort: skip_file
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""BYOC support for CUTLASS."""
+
+from .build import has_cutlass, num_cutlass_partitions, finalize_modules

python/tvm/contrib/cutlass/_ffi_api.py

@@ -0,0 +1,21 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""FFI API for CUTLASS BYOC."""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("contrib.cutlass", __name__)

python/tvm/contrib/cutlass/attention_operation.py

@@ -0,0 +1,327 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name
+# ruff: noqa: E501
+"""Generator for CUTLASS attention kernels."""
+
+from .library import substitute_template
+
+
+def instantiate_attention_template(attrs):
+    """Return CUTLASS host code for fused multi head attention
+    based on a template and the provided attribute map."""
+
+    bias_template = """
+  TVM_FFI_ICHECK(${bias}->ndim == 4); // B, N, S, S'
+
+  p.attn_bias_ptr = reinterpret_cast<T *>(${bias}->data);
+  p.bias_strideM = ${bias_strideM};
+  p.bias_strideH = ${bias_strideH};
+  p.bias_strideB = ${bias_strideB};
+"""
+
+    var_len_template = """
+  p.seqstart_q_ptr = (int32_t*)${seqstart_q}->data;
+  p.seqstart_k_ptr = (int32_t*)${seqstart_k}->data;
+  p.num_queries = ((int32_t*)${max_seqlen_q}->data)[0];
+  p.num_batches = ${seqstart_q}->shape[0] - 1;
+"""
+
+    qkv_template = {
+        "default": """
+  p.query_ptr = reinterpret_cast<T *>(${query}->data);
+  p.key_ptr = reinterpret_cast<T *>(${key}->data);
+  p.value_ptr = reinterpret_cast<T *>(${value}->data);
+  TVM_FFI_ICHECK(${query}->ndim == 4); // B, S, N, H
+  TVM_FFI_ICHECK(${key}->ndim == 4); // B, S', N, H
+  TVM_FFI_ICHECK(${value}->ndim == 4); // B, S', N, H'
+
+  // stride for N
+  p.q_strideH = p.head_dim; // H
+  p.k_strideH = p.head_dim; // H
+  p.v_strideH = p.head_dim_value; // H'
+
+  // stride for S
+  p.q_strideM = p.q_strideH * p.num_heads; // H * N
+  p.k_strideM = p.k_strideH * p.num_heads; // H * N
+  p.v_strideM = p.v_strideH * p.num_heads; // H' * N
+
+  // stride for B
+  p.q_strideB = p.q_strideM * p.num_queries; // H * N * S
+  p.k_strideB = p.k_strideM * p.num_keys; // H * N * S'
+  p.v_strideB = p.v_strideM * p.num_keys; // H'* N * S'
+""",
+        "qkv_stacked": """
+  p.query_ptr = reinterpret_cast<T *>(${qkv}->data);
+  p.key_ptr = reinterpret_cast<T *>(${qkv}->data) + p.head_dim * p.num_heads;
+  p.value_ptr = reinterpret_cast<T *>(${qkv}->data) + p.head_dim * p.num_heads * 2;
+  TVM_FFI_ICHECK(${qkv}->ndim == 3); // B, S, NH + NH + NH'
+
+  // stride for N
+  p.q_strideH = p.head_dim; // H
+  p.k_strideH = p.head_dim; // H
+  p.v_strideH = p.head_dim_value; // H'
+
+  // stride for S
+  p.q_strideM = p.k_strideM = p.v_strideM =
+    p.q_strideH * p.num_heads +
+    p.k_strideH * p.num_heads +
+    p.v_strideH * p.num_heads; // H * N + H * N + H * N'
+
+  // stride for B
+  p.q_strideB = p.k_strideB = p.v_strideB =
+    p.q_strideM * p.num_queries; // (H * N + H * N + H * N') * S
+""",
+    }
+
+    template = """
+  using T = ${data_type};
+
+  using Attention =
+      AttentionKernel<T,
+                      /*ArchTag=*/${arch},
+                      /*is_aligned=*/${kIsAligned},
+                      /*queries_per_block=*/${kQueriesPerBlock},
+                      /*keys_per_block=*/${kKeysPerBlock},
+                      /*kMaxK=*/${kMaxK},
+                      /*supports_dropout=*/${kSupportsDropout},
+                      /*supports_bias=*/${kSupportsBias}
+      >;
+
+  typename Attention::Params p;
+  p.logsumexp_ptr = nullptr;
+  p.output_ptr = reinterpret_cast<T *>(out0->data);
+
+  p.output_accum_ptr = nullptr;
+  uint64_t accumulator_buf_size = ${output_size} * sizeof(Attention::output_accum_t);
+  bool accumulator_buf_allocated = false;
+  if (Attention::kNeedsOutputAccumulatorBuffer) {
+    if (accumulator_buf_size <= ${workspace}->shape[0]) {
+        p.output_accum_ptr = static_cast<float*>(${workspace}->data);
+    } else {
+        accumulator_buf_allocated = true;
+        cudaMalloc(
+          &p.output_accum_ptr,
+          accumulator_buf_size
+        );
+    }
+  }
+
+  p.num_heads = ${num_heads}; // N
+  p.num_batches = ${num_batches}; // B
+  p.head_dim = ${head_dim}; // H
+  p.head_dim_value = ${head_dim_value}; // H'
+  p.num_queries = ${num_queries}; // S
+  p.num_keys = ${num_keys}; // S'
+  p.scale = ${scale};
+  p.custom_mask_type = ${custom_mask_type};
+
+
+  p.o_strideM = p.head_dim_value * p.num_heads; // H' * N
+  TVM_FFI_ICHECK(out0->ndim == 4); // B, S, N, H'
+
+  ${qkv_template}
+  ${bias_template}
+  ${var_len_template}
+
+  constexpr auto kernel_fn = attention_kernel_batched_impl<Attention>;
+  int smem_bytes = sizeof(typename Attention::SharedStorage);
+  if (smem_bytes > 0xc000) {
+    static bool once = [&]() {
+      cudaFuncSetAttribute(
+          kernel_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);
+      return true;
+    }();
+  }
+
+  TVM_FFI_ICHECK(Attention::check_supported(p));
+  cudaStream_t stream = static_cast<cudaStream_t>(TVMFFIEnvGetStream(kDLCUDA, ${query}->device.device_id));
+
+  kernel_fn<<<p.getBlocksGrid(), p.getThreadsGrid(), smem_bytes, stream>>>(p);
+
+  if (accumulator_buf_allocated) {
+    cudaFree(p.output_accum_ptr);
+  }
+"""
+
+    template = substitute_template(
+        template,
+        {
+            "qkv_template": qkv_template[attrs["qkv_layout"]],
+            "bias_template": bias_template if "bias" in attrs else "",
+            "var_len_template": var_len_template if "seqstart_q" in attrs else "",
+        },
+    )
+
+    return substitute_template(template, attrs)
+
+
+def instantiate_flash_attention_template(attrs):
+    """Return host code for flash attention."""
+
+    template = """
+    int q_head_stride = ${head_dim};
+    int k_head_stride = ${head_dim};
+    int v_head_stride = ${head_dim};
+    int o_head_stride = ${head_dim};
+    int q_row_stride = q_head_stride * ${num_q_heads};
+    int k_row_stride = k_head_stride * ${num_kv_heads};
+    int v_row_stride = v_head_stride * ${num_kv_heads};
+    int o_row_stride = o_head_stride * ${num_q_heads};
+    int q_batch_stride = q_row_stride * ${num_queries};
+    int k_batch_stride = k_row_stride * ${num_keys};
+    int v_batch_stride = v_row_stride * ${num_keys};
+    int o_batch_stride = o_row_stride * ${num_queries};
+
+    cudaStream_t stream = static_cast<cudaStream_t>(TVMFFIEnvGetStream(kDLCUDA, ${query}->device.device_id));
+
+    flash_attn::flash_attention_forward(
+                            static_cast<const cutlass::half_t*>(${query}->data),
+                    static_cast<const cutlass::half_t*>(${key}->data),
+                    static_cast<const cutlass::half_t*>(${value}->data),
+                    static_cast<cutlass::half_t*>(out0->data),
+                    ${num_batches},
+                    ${num_queries},
+                    ${num_keys},
+                    ${num_q_heads},
+                    ${num_kv_heads},
+                    ${head_dim},
+                    q_batch_stride,
+                    k_batch_stride,
+                    v_batch_stride,
+                    o_batch_stride,
+                    q_head_stride,
+                    k_head_stride,
+                    v_head_stride,
+                    o_head_stride,
+                    q_row_stride,
+                    k_row_stride,
+                    v_row_stride,
+                    o_row_stride,
+                    ${scale},
+                    ${is_causal},
+                            ${window_size_left},
+                            ${window_size_right},
+                    stream);
+    """
+
+    template_stacked = """
+    int q_head_stride = ${head_dim};
+    int k_head_stride = ${head_dim};
+    int v_head_stride = ${head_dim};
+    int o_head_stride = ${head_dim};
+    int row_stride = q_head_stride * ${num_q_heads} +
+                     k_head_stride * ${num_kv_heads} +
+                     v_head_stride * ${num_kv_heads};
+    int q_row_stride = row_stride;
+    int k_row_stride = row_stride;
+    int v_row_stride = row_stride;
+    int o_row_stride = o_head_stride * ${num_q_heads};
+
+    int q_batch_stride = q_row_stride * ${num_queries};
+    int k_batch_stride = k_row_stride * ${num_keys};
+    int v_batch_stride = v_row_stride * ${num_keys};
+    int o_batch_stride = o_row_stride * ${num_queries};
+
+    cudaStream_t stream = static_cast<cudaStream_t>(TVMFFIEnvGetStream(kDLCUDA, ${query}->device.device_id));
+
+    flash_attn::flash_attention_forward(
+                            static_cast<const cutlass::half_t*>(${qkv}->data),
+                    static_cast<const cutlass::half_t*>(${qkv}->data) + ${head_dim} * ${num_q_heads},
+                    static_cast<const cutlass::half_t*>(${qkv}->data) + ${head_dim} * (${num_q_heads} + ${num_kv_heads}),
+                    static_cast<cutlass::half_t*>(out0->data),
+                    ${num_batches},
+                    ${num_queries},
+                    ${num_keys},
+                    ${num_q_heads},
+                    ${num_kv_heads},
+                    ${head_dim},
+                    q_batch_stride,
+                    k_batch_stride,
+                    v_batch_stride,
+                    o_batch_stride,
+                    q_head_stride,
+                    k_head_stride,
+                    v_head_stride,
+                    o_head_stride,
+                    q_row_stride,
+                    k_row_stride,
+                    v_row_stride,
+                    o_row_stride,
+                    ${scale},
+                    ${is_causal},
+                            ${window_size_left},
+                            ${window_size_right},
+                    stream);
+    """
+
+    if "qkv" in attrs:
+        return substitute_template(template_stacked, attrs)
+
+    return substitute_template(template, attrs)
+
+
+def instantiate_flash_attention_var_len_template(attrs):
+    """Return host code for flash attention with variable sequence lengths."""
+
+    template = """
+    int _max_seqlen_q = ((int32_t*)${max_seqlen_q}->data)[0];
+    int _max_seqlen_k = ((int32_t*)${max_seqlen_k}->data)[0];
+
+    int batch_size = ${seqstart_q}->shape[0] - 1;
+
+    int q_head_stride = ${head_dim};
+    int k_head_stride = ${head_dim};
+    int v_head_stride = ${head_dim};
+    int o_head_stride = ${head_dim};
+    int q_row_stride = q_head_stride * ${num_q_heads};
+    int k_row_stride = k_head_stride * ${num_kv_heads};
+    int v_row_stride = v_head_stride * ${num_kv_heads};
+    int o_row_stride = o_head_stride * ${num_q_heads};
+
+    cudaStream_t stream = static_cast<cudaStream_t>(TVMFFIEnvGetStream(kDLCUDA, ${query}->device.device_id));
+
+    flash_attn::flash_attention_var_len_forward(
+                            static_cast<const cutlass::half_t*>(${query}->data),
+                    static_cast<const cutlass::half_t*>(${key}->data),
+                    static_cast<const cutlass::half_t*>(${value}->data),
+                            static_cast<const int*>(${seqstart_q}->data),
+                            static_cast<const int*>(${seqstart_k}->data),
+                    static_cast<cutlass::half_t*>(out0->data),
+                    batch_size,
+                    _max_seqlen_q,
+                    _max_seqlen_k,
+                    ${num_q_heads},
+                    ${num_kv_heads},
+                    ${head_dim},
+                    q_head_stride,
+                    k_head_stride,
+                    v_head_stride,
+                    o_head_stride,
+                    q_row_stride,
+                    k_row_stride,
+                    v_row_stride,
+                    o_row_stride,
+                    ${scale},
+                    ${is_causal},
+                            // For SWA, is_causal must be false.
+                            ${is_causal} ? _max_seqlen_k : ${window_size_left},
+                            ${window_size_right},
+                    stream);
+    """
+
+    return substitute_template(template, attrs)

python/tvm/contrib/cutlass/build.py

@@ -0,0 +1,907 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name, dangerous-default-value, arguments-differ
+# ruff: noqa: F821
+"""Driver for partitioning and building a Relax module for CUTLASS offload."""
+
+import itertools
+import logging
+import multiprocessing
+import operator
+import os
+from collections.abc import Sequence
+from functools import reduce
+
+from tvm_ffi import register_global_func
+
+import tvm
+from tvm import relax, runtime
+from tvm.support.nvcc import get_cuda_version
+from tvm.topi.utils import get_const_tuple
+
+from .gen_conv2d import CutlassConv2DProfiler
+from .gen_gemm import CutlassGemmProfiler
+from .library import ConvKind, LayoutType
+
+logger = logging.getLogger("cutlass")
+
+
+def has_cutlass():
+    """Returns true if the CUTLASS custom codegen is available"""
+    return tvm.get_global_func("relax.ext.cutlass", True) is not None
+
+
+def _get_cutlass_path():
+    invalid_paths = []
+    for rel in ["../../../../", "../../../", "../../"]:
+        tvm_root = os.path.join(os.path.dirname(os.path.realpath(__file__)), rel)
+        cutlass_path = os.path.join(tvm_root, "3rdparty/cutlass")
+        if os.path.exists(cutlass_path):
+            return cutlass_path
+        invalid_paths.append(cutlass_path)
+    raise AssertionError(f"The CUTLASS root directory not found in: {invalid_paths}")
+
+
+def _get_cutlass_compile_options(sm, threads, use_fast_math=False):
+    cutlass_root = _get_cutlass_path()
+    cutlass_include = os.path.join(cutlass_root, "include")
+    cutlass_util_include = os.path.join(cutlass_root, "tools/util/include")
+    cutlass_attention_include = os.path.join(cutlass_root, "examples/41_fused_multi_head_attention")
+    cutlass_fpA_intB_gemm_include = os.path.join(cutlass_root, "../cutlass_fpA_intB_gemm")
+    flash_attn_include = os.path.join(cutlass_root, "../libflash_attn/include")
+
+    kwargs = {}
+    kwargs["cc"] = "nvcc"
+    kwargs["options"] = [
+        "-c",
+        "-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1",
+        f"-gencode=arch=compute_{sm},code=[sm_{sm},compute_{sm}]",
+        "-DNDEBUG",
+        "-Xcompiler=-fPIC",
+        "-Xcompiler=-Wconversion",
+        "-Xcompiler=-fno-strict-aliasing",
+        "-Xcompiler=-fvisibility=hidden",
+        "-O3",
+        "-std=c++17",
+        f"-I{cutlass_include}",
+        f"-I{cutlass_util_include}",
+        f"-I{cutlass_attention_include}",
+        f"-I{cutlass_fpA_intB_gemm_include}",
+        f"-I{flash_attn_include}",
+    ]
+    if use_fast_math:
+        kwargs["options"].append("-DCUTLASS_USE_TANH_FOR_SIGMOID")
+    cuda_ver = get_cuda_version()
+    if cuda_ver >= (11, 2):
+        ncpu = multiprocessing.cpu_count() if threads < 0 else threads
+        kwargs["options"].append(f"-t {ncpu}")
+    return kwargs
+
+
+def select_gemm_kernel(
+    cutlass_profiler,
+    op_type,
+    MM,
+    KK,
+    NN,
+    out_dtype,
+    arg0_dtype,
+    arg1_dtype,
+    use_3xtf32,
+    batched,
+    find_first_valid,
+    use_multiprocessing,
+):
+    """Run CUTLASS profiler to select the best kernel, or return the default one for dynamic
+    workloads."""
+    if any(isinstance(s, tvm.tirx.Any) for s in [MM, KK, NN]):
+        out = cutlass_profiler.get_default(
+            op_type, out_dtype, arg0_dtype, arg1_dtype, use_3xtf32, batched=batched
+        )
+        name, cutlass_op_def = out["name"], out["opdef"]
+        logger.info("Picked the default kernel %s", name)
+    else:
+        name, cutlass_op_def, _ = cutlass_profiler.profile(
+            op_type,
+            MM,
+            NN,
+            KK,
+            out_dtype,
+            arg0_dtype,
+            arg1_dtype,
+            use_3xtf32,
+            batched=batched,
+            find_first_valid=find_first_valid,
+            use_multiprocessing=use_multiprocessing,
+        )
+        if not find_first_valid:
+            logger.info("The best kernel is %s", name)
+        else:
+            logger.info("Picked the first kernel found %s", name)
+
+    return name, cutlass_op_def
+
+
+def handle_batch_matmul(
+    cutlass_profiler,
+    op_type,
+    arg0_shape,
+    arg1_shape,
+    out_dtype,
+    arg0_dtype,
+    arg1_dtype,
+    use_3xtf32,
+    find_first_valid,
+    use_multiprocessing,
+):
+    """Profile and select a kernel for batch_matmul op workload."""
+    MM = arg0_shape[1]
+    KK = arg0_shape[2]
+    NN = arg1_shape[1]
+
+    name, cutlass_op_def = select_gemm_kernel(
+        cutlass_profiler,
+        op_type,
+        MM,
+        KK,
+        NN,
+        out_dtype,
+        arg0_dtype,
+        arg1_dtype,
+        use_3xtf32,
+        True,
+        find_first_valid,
+        use_multiprocessing,
+    )
+
+    return {
+        "batch": arg0_shape[0],
+        "batch_stride_A": arg0_shape[1] * arg0_shape[2],
+        "batch_stride_B": arg1_shape[1] * arg1_shape[2],
+        "batch_stride_C": arg0_shape[1] * arg1_shape[1],
+        "cutlass_op_def": cutlass_op_def,
+        "cutlass_op_name": name,
+        "lda": "K",
+        "ldb": "K",
+        "ldc": "N",
+    }
+
+
+def handle_dense(
+    cutlass_profiler,
+    op_type,
+    arg0_shape,
+    arg1_shape,
+    out_dtype,
+    arg0_dtype,
+    arg1_dtype,
+    use_3xtf32,
+    find_first_valid,
+    use_multiprocessing,
+):
+    """Profile and select a kernel for dense op workload."""
+    MM = arg0_shape[0]
+    KK = arg0_shape[1]
+    NN = arg1_shape[0]
+
+    name, cutlass_op_def = select_gemm_kernel(
+        cutlass_profiler,
+        op_type,
+        MM,
+        KK,
+        NN,
+        out_dtype,
+        arg0_dtype,
+        arg1_dtype,
+        use_3xtf32,
+        False,
+        find_first_valid,
+        use_multiprocessing,
+    )
+
+    assert "tn_align" in name, "Only supports (row_major, col_major) input layout for now."
+
+    return {
+        "cutlass_op_def": cutlass_op_def,
+        "cutlass_op_name": name,
+        "lda": "K",
+        "ldb": "K",
+        "ldc": "N",
+    }
+
+
+def handle_conv2d(
+    cutlass_profiler,
+    op_type,
+    d_shape,
+    w_shape,
+    padding,
+    strides,
+    dilation,
+    out_dtype,
+    data_dtype,
+    weight_dtype,
+    use_3xtf32,
+    split_k_slices,
+    profile_all_alignments,
+    find_first_valid,
+    use_multiprocessing,
+):
+    """Profile and select a kernel for conv2d op workload."""
+    if "conv2d_transpose" in op_type:
+        conv_kind = ConvKind.Dgrad
+    elif "backward_weight" in op_type:
+        conv_kind = ConvKind.Wgrad
+    else:
+        conv_kind = ConvKind.Fprop
+
+    if any(isinstance(s, tvm.tirx.Any) for s in d_shape):
+        out = cutlass_profiler.get_default(
+            op_type, out_dtype, data_dtype, weight_dtype, use_3xtf32, conv_kind, strides
+        )
+        name, cutlass_op_def = out["name"], out["opdef"]
+        logger.info("Picked the default kernel %s", name)
+    else:
+        name, cutlass_op_def, _ = cutlass_profiler.profile(
+            op_type,
+            d_shape,
+            w_shape,
+            padding,
+            strides,
+            dilation,
+            out_dtype,
+            data_dtype,
+            weight_dtype,
+            use_3xtf32,
+            conv_kind,
+            split_k_slices,
+            profile_all_alignments,
+            find_first_valid=find_first_valid,
+            use_multiprocessing=use_multiprocessing,
+        )
+        if not find_first_valid:
+            logger.info("The best kernel is %s", name)
+        else:
+            logger.info("Picked the first kernel found %s", name)
+
+    return {"cutlass_op_def": cutlass_op_def, "cutlass_op_name": name}
+
+
+def num_cutlass_partitions(mod):
+    return sum([(1 if "cutlass" in var.name_hint else 0) for var in mod.get_global_vars()])
+
+
+def tune_cutlass_kernels(
+    mod,
+    sm,
+    use_3xtf32=True,
+    split_k_slices=[1],
+    profile_all_alignments=False,
+    find_first_valid=False,
+    use_multiprocessing=False,
+    tmp_dir="./tmp",
+):
+    """Given a module partitioned for CUTLASS offloading, profile each workload to select which
+    kernels to emit.
+
+    Parameters
+    ----------
+    mod : IRModule
+        The IRModule with cutlass partitions.
+
+    sm : int
+        An integer specifying the compute capability. For example, 75 for Turing and
+        80 or 86 for Ampere.
+
+    use_3xtf32 : bool
+        Wheter or not use slower but very accurate (compared to tf32) 3xtf32 mode for
+        fp32 inputs on tensorcore.
+
+    split_k_slices : list of int
+        Split factor candidates for split-K GEMM. If split-K > 1, the GEMM K-loop is computed in
+        parallel across split-K blocks, and a separate global reduction kernel is launched to
+        accumulate partial reductions. The profiler will pick the best split-k factor from the
+        given candidate list. Note that the larger split-K factor requires a larger workspace.
+        Currently, parallel split-k has been tested only for wgrad. For GEMM and other conv2d
+        kinds, split_k_slices is ignored.
+
+    profile_all_alignments : bool
+        When True, profile all kernal variants with smaller alignments than the largest possible.
+
+    find_first_valid : bool
+        Whether or not profile all candidate kernels, or stop profiling after
+        the first applicable kernel is found.
+
+    use_multiprocessing : bool
+        Whether or not compile profiler executables for different kernels in parallel.
+
+    tmp_dir : string, optional
+        A temporary directory where intermediate compiled artifacts will be stored.
+
+    Returns
+    -------
+    mod : IRModule
+        The updated module annotated with cutlass profiling information.
+
+    num_cutlass_partition : int
+        The number of partitioned functions created for CUTLASS.
+    """
+    gemm_profiler = CutlassGemmProfiler(sm, _get_cutlass_path(), tmp_dir)
+    conv2d_profiler = CutlassConv2DProfiler(sm, _get_cutlass_path(), tmp_dir)
+    num_cutlass_partition = 0
+    for var in mod.get_global_vars():
+        fun_name = var.name_hint
+        func = mod[fun_name]
+        if "cutlass" in fun_name:
+            num_cutlass_partition += 1
+            new_func = tune_cutlass_function(
+                func,
+                use_3xtf32,
+                split_k_slices,
+                profile_all_alignments,
+                find_first_valid,
+                use_multiprocessing,
+                gemm_profiler,
+                conv2d_profiler,
+            )
+            mod.update_func(var, new_func)
+
+    return mod, num_cutlass_partition
+
+
+def _get_call_node(expr: relax.Expr, op_name: str) -> relax.Call | None:
+    node = None
+
+    def fvisit(e):
+        nonlocal node
+        if isinstance(e, relax.Call) and e.op.name == op_name:
+            node = e
+
+    relax.analysis.post_order_visit(expr, fvisit)
+    return node
+
+
+def _extract_relax_function_signature(f):
+    signature = {}
+
+    for i, arg in enumerate(f.params):
+        ty = arg.ty
+        if isinstance(ty, relax.TensorType):
+            signature[f"arg{i}_shape"] = get_const_tuple(ty.shape)
+            signature[f"arg{i}_dtype"] = ty.dtype
+        elif isinstance(ty, relax.ShapeType):
+            signature[f"arg{i}_shape"] = get_const_tuple(ty.values)
+        else:
+            raise NotImplementedError()
+
+    ret_ty = f.ret_ty
+    if ret_ty.shape is not None:
+        signature["ret_shape"] = get_const_tuple(ret_ty.shape)
+    else:
+        signature["ret_shape"] = None
+    signature["ret_dtype"] = ret_ty.dtype
+
+    return signature
+
+
+def _extract_arg_idx(pattern_name, f):
+    extract_func = tvm.get_global_func("relax.contrib.extract_arg_idx")
+    arg_indices = extract_func(pattern_name, f)
+    return {k: int(v) for k, v in arg_indices.items()}
+
+
+def is_shape_valid_for_cutlass_matmul(
+    lhs_shape: Sequence[tvm.ir.Expr],
+    rhs_shape: Sequence[tvm.ir.Expr],
+) -> bool:
+    """
+    Check whether the shape of inputs can be handled by CUTLASS GEMM.
+
+    The stride-based batch matmul in CUTLASS cannot handle cases that some of
+    the batch dimensions need to be stretched while others don't. This means
+    it can only handle ND x ND whose batch dimensions match exactly on both side,
+    as well as ND x 2D and 2D x ND. For example, it cannot handle matmul with shape
+    (2, 1, 4, 8) x (2, 3, 8, 16), because the batch stride of lhs is not constant.
+    """
+    if not isinstance(lhs_shape[-1], tvm.tirx.expr.IntImm | int):
+        # Reduction axis must be constant
+        return False
+
+    lhs_batches = reduce(operator.mul, lhs_shape[:-2], 1)
+    rhs_batches = reduce(operator.mul, rhs_shape[:-2], 1)
+    if lhs_batches == 1 or rhs_batches == 1:
+        # This could be regular matmul or batch matmul with shape ND x 2D or 2D x ND
+        return True
+
+    analyzer = tvm.arith.Analyzer()
+    # If one side has less dimensions, use 1 to fill the gap
+    batch_dim_pairs = list(
+        itertools.zip_longest(
+            list(lhs_shape)[-3::-1],  # Remove the last two dimensions and reverse
+            list(rhs_shape)[-3::-1],
+            fillvalue=1,
+        )
+    )
+    return all(analyzer.can_prove_equal(p[0], p[1]) for p in batch_dim_pairs)
+
+
+@relax.expr_functor.mutator
+class CutlassRelaxFunctionAnnotator(relax.PyExprMutator):
+    """A Relax function mutator that tunes and annotates CUTLASS composite functions
+    with shape, dtype and generated templates.
+    """
+
+    def __init__(
+        self,
+        mod,
+        conv2d_profiler: CutlassConv2DProfiler,
+        gemm_profiler: CutlassGemmProfiler,
+        options,
+    ):
+        super().__init__(mod)
+        self.options = options
+        self.conv2d_profiler = conv2d_profiler
+        self.gemm_profiler = gemm_profiler
+
+    def handle_conv2d(self, f, op_type):
+        """Tune and annotate a conv2d op."""
+        signature = _extract_relax_function_signature(f)
+        arg_idx = _extract_arg_idx(op_type, f)
+        op_attrs = _get_call_node(f.body, "relax.nn.conv2d").attrs
+
+        data_arg = f"arg{arg_idx['lhs']}"
+        weight_arg = f"arg{arg_idx['rhs']}"
+
+        d_shape = signature[f"{data_arg}_shape"]
+        w_shape = signature[f"{weight_arg}_shape"]
+        out_shape = signature["ret_shape"]
+        data_dtype = signature[f"{data_arg}_dtype"]
+        weight_dtype = signature[f"{weight_arg}_dtype"]
+        out_dtype = signature["ret_dtype"]
+        padding = op_attrs["padding"]
+        strides = op_attrs["strides"]
+        dilation = op_attrs["dilation"]
+        conv_kind = ConvKind.Fprop
+
+        use_3xtf32 = self.options.get("use_3xtf32", False)
+        profile_all_alignments = self.options.get("profile_all_alignments", False)
+        find_first_valid = self.options.get("find_first_valid", True)
+        use_multiprocessing = self.options.get("use_multiprocessing", True)
+        split_k_slices = self.options.get("split_k_slices", [1])
+
+        op_name, op_def, _ = self.conv2d_profiler.profile(
+            op_type,
+            d_shape,
+            w_shape,
+            padding,
+            strides,
+            dilation,
+            out_dtype,
+            data_dtype,
+            weight_dtype,
+            use_3xtf32,
+            conv_kind,
+            split_k_slices,
+            profile_all_alignments,
+            find_first_valid=find_first_valid,
+            use_multiprocessing=use_multiprocessing,
+        )
+
+        attrs = {
+            "op_type": op_type,
+            "data_arg_idx": arg_idx["lhs"],
+            "weight_arg_idx": arg_idx["rhs"],
+            "bias_arg_idx": arg_idx.get("bias"),
+            "residual_arg_idx": arg_idx.get("residual"),
+            "arg0_dtype": data_dtype,
+            "arg1_dtype": weight_dtype,
+            "ret_dtype": out_dtype,
+            "arg0_shape": d_shape,
+            "arg1_shape": w_shape,
+            "ret_shape": out_shape,
+            "strides": strides,
+            "padding": padding,
+            "dilation": dilation,
+            "cutlass_op_name": op_name,
+            "cutlass_op_def": op_def,
+        }
+
+        residual_arg = arg_idx.get("residual")
+
+        if residual_arg:
+            residual_shape = signature[f"arg{residual_arg}_shape"]
+            attrs["residual_shape"] = residual_shape
+        elif "residual" in op_type:
+            attrs["residual_shape"] = d_shape
+
+        return f.with_attrs(attrs)
+
+    def handle_decode_matmul(self, f, op_type):
+        """Annotate a decode -> matmul op."""
+        arg_idx = _extract_arg_idx(op_type, f)
+        signature = _extract_relax_function_signature(f)
+        lhs_arg = f"arg{arg_idx['lhs']}"
+        rhs_arg = f"arg{arg_idx['w_encoded']}"
+        lhs_shape = signature[f"{lhs_arg}_shape"]
+        rhs_shape = signature[f"{rhs_arg}_shape"]
+        ret_shape = signature["ret_shape"]
+        scale_arg = f"arg{arg_idx['scales']}"
+        scale_shape = signature[f"{scale_arg}_shape"]
+        N = ret_shape[-1]
+
+        attrs = {
+            "op_type": op_type,
+            "lhs_arg_idx": arg_idx["lhs"],
+            "rhs_arg_idx": arg_idx["w_encoded"],
+            "scales_arg_idx": arg_idx["scales"],
+            "bias_arg_idx": arg_idx.get("bias"),
+            "activation": "identity",
+        }
+        # TODO(wuwei): find a better way to get group size
+        attrs["group_size"] = 64 if len(scale_shape) == 2 and scale_shape[0] != 1 else -1
+
+        attrs["batch_rank"] = len(lhs_shape[:-1])
+        attrs["M"] = reduce(operator.mul, lhs_shape[:-1], 1)
+
+        attrs["bias_stride"] = 0
+
+        if "bias" in arg_idx:
+            bias_shape = signature[f"arg{arg_idx['bias']}_shape"]
+            bias_shape_1d = reduce(operator.mul, bias_shape, 1)
+            if bias_shape_1d != bias_shape[-1]:
+                attrs["bias_stride"] = bias_shape[-1]
+
+        if N == rhs_shape[1]:
+            attrs["weight_nbit"] = 8
+        else:
+            assert N == rhs_shape[1] * 2
+            attrs["weight_nbit"] = 4
+
+        if "residual" in op_type:
+            residual_pos = op_type.find("residual_")
+            postfix = op_type[residual_pos + len("residual_") :]
+
+            if postfix.startswith("multiply"):
+                binary_op = "multiply"
+            else:
+                binary_op = "plus"
+
+            if "relu" in postfix:
+                unary_op = "relu"
+            else:
+                unary_op = "identity"
+
+            activation = "identity"
+
+            for act in ["relu", "silu", "gelu"]:
+                if act in op_type[op_type.find("matmul_") + len("matmul_") : residual_pos]:
+                    activation = act
+                    break
+
+            attrs.update(
+                {
+                    "unary_op": unary_op,
+                    "binary_op": binary_op,
+                    "activation": activation,
+                    "residual_arg_idx": arg_idx["residual"],
+                }
+            )
+        else:
+            for act in ["relu", "silu", "gelu"]:
+                if act in op_type:
+                    attrs["activation"] = act
+                    break
+
+        return f.with_attrs(attrs)
+
+    def handle_matmul(self, f, op_type):
+        """Tune and annotate a matmul op."""
+        signature = _extract_relax_function_signature(f)
+        arg_idx = _extract_arg_idx(op_type, f)
+
+        lhs_arg = f"arg{arg_idx['lhs']}"
+        rhs_arg = f"arg{arg_idx['rhs']}"
+
+        lhs_shape = signature[f"{lhs_arg}_shape"]
+        rhs_shape = signature[f"{rhs_arg}_shape"]
+        out_shape = signature["ret_shape"]
+        lhs_dtype = signature[f"{lhs_arg}_dtype"]
+        rhs_dtype = signature[f"{rhs_arg}_dtype"]
+        out_dtype = signature["ret_dtype"]
+
+        if not is_shape_valid_for_cutlass_matmul(lhs_shape, rhs_shape):
+            raise ValueError(f"Cannot handle the input shapes, lhs: {lhs_shape}, rhs: {rhs_shape}")
+
+        MM = lhs_shape[-2]
+        KK = lhs_shape[-1]
+        if "transposed" in op_type:
+            NN = rhs_shape[-2]
+            ldb = "K"
+            layout_b = LayoutType.ColumnMajor
+        else:
+            NN = rhs_shape[-1]
+            ldb = "N"
+            layout_b = LayoutType.RowMajor
+
+        lhs_batches = reduce(operator.mul, lhs_shape[:-2], 1)
+        rhs_batches = reduce(operator.mul, rhs_shape[:-2], 1)
+        if lhs_batches == 1 and rhs_batches == 1:
+            # Regular matmul
+            is_batched = False
+            batch_attrs = {}
+        else:
+            is_batched = True
+            batch_attrs = {
+                # If both lhs_batches and rhs_batches are greater than 1,
+                # they must be equal. This is checked by is_shape_valid_for_cutlass_matmul.
+                "batch": lhs_batches if rhs_batches == 1 else rhs_batches,
+                "batch_stride_A": 0 if lhs_batches == 1 else MM * KK,
+                "batch_stride_B": 0 if rhs_batches == 1 else KK * NN,
+                "batch_stride_C": MM * NN,
+            }
+
+        use_3xtf32 = self.options.get("use_3xtf32", False)
+        find_first_valid = self.options.get("find_first_valid", True)
+        use_multiprocessing = self.options.get("use_multiprocessing", True)
+
+        op_name, op_def, _ = self.gemm_profiler.profile(
+            op_type,
+            MM,
+            NN,
+            KK,
+            out_dtype,
+            lhs_dtype,
+            rhs_dtype,
+            use_3xtf32,
+            batched=is_batched,
+            find_first_valid=find_first_valid,
+            use_multiprocessing=use_multiprocessing,
+            layout_b=layout_b,
+        )
+
+        return f.with_attrs(
+            {
+                "op_type": op_type,
+                "lhs_arg_idx": arg_idx["lhs"],
+                "rhs_arg_idx": arg_idx["rhs"],
+                "residual_arg_idx": arg_idx.get("residual"),
+                "bias_arg_idx": arg_idx.get("bias"),
+                "arg0_dtype": signature["arg0_dtype"],
+                "arg1_dtype": signature["arg1_dtype"],
+                "ret_dtype": out_dtype,
+                "arg0_shape": signature["arg0_shape"],
+                "arg1_shape": signature["arg1_shape"],
+                "ret_shape": out_shape,
+                "lda": "K",
+                "ldb": ldb,
+                "ldc": "N",
+                "cutlass_op_name": op_name,
+                "cutlass_op_def": op_def,
+                **batch_attrs,
+            }
+        )
+
+    def handle_attention(self, f, op_type):
+        """Annotate an attention op."""
+        signature = _extract_relax_function_signature(f)
+
+        if _get_call_node(f.body, "relax.nn.attention") is not None:
+            attention_node = _get_call_node(f.body, "relax.nn.attention")
+            op_attrs = attention_node.attrs
+        elif _get_call_node(f.body, "relax.nn.attention_bias") is not None:
+            attention_node = _get_call_node(f.body, "relax.nn.attention_bias")
+            op_attrs = attention_node.attrs
+        elif _get_call_node(f.body, "relax.nn.attention_var_len") is not None:
+            attention_node = _get_call_node(f.body, "relax.nn.attention_var_len")
+            op_attrs = attention_node.attrs
+        else:
+            raise ValueError("Cannot find call node for attention")
+        arg = {}
+
+        if "stacked_attention" in op_type:
+            arg["arg0_dtype"] = signature["arg0_dtype"]
+            q_shape = get_const_tuple(attention_node.args[0].ty.shape)
+            k_shape = get_const_tuple(attention_node.args[1].ty.shape)
+            v_shape = get_const_tuple(attention_node.args[2].ty.shape)
+            if len(attention_node.args) == 4:
+                arg["bias_shape"] = get_const_tuple(attention_node.args[3].ty.shape)
+                arg["bias_dtype"] = attention_node.args[3].ty.dtype
+
+            qkv_layout = "qkv_stacked"
+        else:
+            # arg0: q, arg1: k, arg2: v,  arg3: bias, arg4: workspace
+            arg["arg0_shape"] = q_shape = signature["arg0_shape"]
+            arg["arg1_shape"] = k_shape = signature["arg1_shape"]
+            arg["arg2_shape"] = v_shape = signature["arg2_shape"]
+            arg["arg0_dtype"] = signature["arg0_dtype"]
+            arg["arg1_dtype"] = signature["arg1_dtype"]
+            arg["arg2_dtype"] = signature["arg2_dtype"]
+
+            if "arg4_dtype" in signature:
+                arg["bias_dtype"] = signature["arg3_dtype"]
+            if "arg4_shape" in signature:
+                arg["bias_shape"] = signature["arg3_shape"]
+
+            qkv_layout = "default"
+
+        out_shape = signature["ret_shape"]
+        out_dtype = signature["ret_dtype"]
+        num_batches, num_queries, num_q_heads, head_dim = q_shape
+        _, num_keys, num_kv_heads, _ = k_shape
+        _, _, _, head_dim_value = v_shape
+        scale = op_attrs.scale
+
+        if op_attrs.causal_mask is None:
+            custom_mask_type = 0
+        elif op_attrs.causal_mask == "TopLeft":
+            custom_mask_type = 1
+        elif op_attrs.causal_mask == "BottomRight":
+            custom_mask_type = 2
+        else:
+            raise NotImplementedError()
+
+        attrs = {
+            "op_type": op_type,
+            "ret_dtype": out_dtype,
+            "ret_shape": out_shape,
+            "num_batches": num_batches,
+            "num_queries": num_queries,
+            "num_keys": num_keys,
+            "num_q_heads": num_q_heads,
+            "num_kv_heads": num_kv_heads,
+            "head_dim": head_dim,
+            "head_dim_value": head_dim_value,
+            "scale": scale,
+            "arch": self.options["sm"],
+            "qkv_layout": qkv_layout,
+            "custom_mask_type": custom_mask_type,
+            **arg,
+        }
+
+        if "var_len" in op_type:
+            arg_idx = _extract_arg_idx(op_type, f)
+            for arg in ["seqstart_q", "seqstart_k", "max_seqlen_q", "max_seqlen_k"]:
+                if arg in arg_idx:
+                    attrs[arg + "_idx"] = arg_idx[arg]
+
+        if op_attrs.window_size:
+            attrs["window_size"] = op_attrs.window_size
+
+        return f.with_attrs(attrs)
+
+    def handle_norm(self, f, op_type):
+        """Annotate a layer or rms norm op."""
+        signature = _extract_relax_function_signature(f)
+        attrs = {}
+        attrs["batch_rank"] = len(signature["arg0_shape"][:-1])
+        attrs["M"] = reduce(operator.mul, signature["arg0_shape"][:-1], 1)
+        attrs["N"] = signature["arg0_shape"][-1]
+        dtype = signature["arg0_dtype"]
+        attrs["data_type"] = {"float32": "float", "float16": "cutlass::half_t"}[str(dtype)]
+
+        if "rms" in op_type:
+            attrs["rms_eps"] = self.options.get("rms_eps", 1e-5)
+        else:
+            attrs["layer_norm_eps"] = self.options.get("layer_nrom_eps", 1e-5)
+
+        return f.with_attrs(attrs)
+
+    def visit_function_(self, f):
+        if "Composite" not in f.attrs:
+            body = super().visit_expr(f.body)
+            return relax.Function(f.params, body, f.ret_ty, f.is_pure, f.attrs, f.span)
+
+        op_type = f.attrs["Composite"]
+
+        if "conv2d" in op_type:
+            return self.handle_conv2d(f, op_type)
+        elif "decode" in op_type:
+            return self.handle_decode_matmul(f, op_type)
+        elif "matmul" in op_type:
+            return self.handle_matmul(f, op_type)
+        elif "attention" in op_type:
+            return self.handle_attention(f, op_type)
+        elif "layer_norm" in op_type or "rms_norm" in op_type:
+            return self.handle_norm(f, op_type)
+
+        raise ValueError(f"Unsupported composite {op_type}")
+
+    def visit_span(self, span):
+        return span
+
+
+@register_global_func("contrib.cutlass.tune_relax_function")
+def profile_relax_function(functions, options):
+    """Tune and annotate CUTLASS composite functions with shape, dtype and generated templates."""
+    tmp_dir = options.get("tmp_dir", "./tmp")
+    sm = options.get("sm", 80)
+    conv2d_profiler = CutlassConv2DProfiler(sm, _get_cutlass_path(), tmp_dir)
+    gemm_profiler = CutlassGemmProfiler(sm, _get_cutlass_path(), tmp_dir)
+
+    annotated_functions = []
+
+    for f in functions:
+        annotator = CutlassRelaxFunctionAnnotator(
+            tvm.IRModule.from_expr(f), conv2d_profiler, gemm_profiler, options
+        )
+        annotated_functions.append(annotator.visit_expr(f))
+
+    return annotated_functions
+
+
+@register_global_func("contrib.cutlass.compile")
+def compile_cutlass_module(c_source_module, options):
+    """Compile all CUTLASS kernels in the given C-source module.
+
+    Parameters
+    ----------
+    c_source_module: runtime.Module
+        A C-source module containing CUTLASS kernels.
+
+    options: dict
+        Compilation options. Currently recognizes
+          "sm": The target architecture (compute capability), for example 75 or 80 (default: 80)
+          "threads": The number of threads to use in NVCC parallel compilation (default:
+          use all logical cores)
+          "use_fast_math": Whether or not to use faster but approximate arithmetic in some
+          CUTLASS epilogues (default: False)
+
+    Returns
+    -------
+    rt_mod : runtime.Module
+        A runtime module where all cutlass kernels have been compiled.
+    """
+    tmp_dir = options.get("tmp_dir", "./tmp")
+    defaults = {"sm": 80, "threads": -1, "use_fast_math": False}
+    compile_config = {key: options.get(key, val) for key, val in defaults.items()}
+
+    function_names = c_source_module.get_function("get_func_names")()
+    compile_options = _get_cutlass_compile_options(**compile_config)
+    lib_path = os.path.join(tmp_dir, "cutlass.o")
+    logger.info("Compiling generated CUTLASS code")
+    c_source_module.export_library(lib_path, workspace_dir=tmp_dir, **compile_options)
+
+    # Recover static library
+    return tvm.runtime.load_static_library(lib_path, function_names)
+
+
+def finalize_modules(lib, lib_path="compile.so", tmp_dir="./tmp"):
+    """Returns lib with any C source, LLVM and static library modules complied and linked in ready
+    for use by the graph or AOT executors. This method is not specific to CUTLASS, however it does
+    assume nvcc will be used for final compilation and linking. It is provided here for
+    convenience.
+
+    Parameters
+    ----------
+    lib : runtime.Module
+        The output from build.
+
+    lib_path : string
+        The path to a shared library which will be generated as the result of the build process.
+
+    tmp_dir : string
+        A temporary directory where intermediate compiled artifacts will be stored.
+
+    Returns
+    -------
+    updated_lib : runtime.Module
+        The updated library with all compilation and linking completed.
+
+    """
+    lib_path = os.path.join(tmp_dir, lib_path)
+    lib.export_library(lib_path, workspace_dir=tmp_dir, cc="nvcc")
+    return runtime.load_module(lib_path)

python/tvm/contrib/cutlass/conv2d_operation.py

@@ -0,0 +1,552 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name, unused-wildcard-import, wildcard-import
+# ruff: noqa: E501, F403, F405
+"""Generator for CUTLASS Conv2D kernels."""
+
+from .library import *
+
+
+class Conv2dOperation:
+    """Describes various attributes for instantiating Conv2d kernels."""
+
+    def __init__(
+        self,
+        conv_kind,
+        iterator_algorithm,
+        arch,
+        tile_description,
+        A,
+        B,
+        C,
+        element_epilogue,
+        stride_support,
+        epilogue_functor=EpilogueFunctor.LinearCombination,
+        swizzling_functor=SwizzlingFunctor.Identity1,
+        split_k_slices=1,
+    ):
+        self.operation_kind = OperationKind.Conv2d
+        self.arch = arch
+        self.tile_description = tile_description
+        self.conv_kind = conv_kind
+        self.A = A
+        self.B = B
+        self.C = C
+        self.element_epilogue = element_epilogue
+        self.epilogue_functor = epilogue_functor
+        self.iterator_algorithm = iterator_algorithm
+        self.stride_support = stride_support
+        self.swizzling_functor = swizzling_functor
+        self.split_k_slices = split_k_slices
+
+    def accumulator_type(self):
+        return self.tile_description.math_instruction.element_accumulator
+
+    def core_name(self):
+        """The basic operation kind is prefixed with a letter indicating the accumulation type."""
+        intermediate_type = ""
+
+        if self.tile_description.math_instruction.opcode_class == OpcodeClass.TensorOp:
+            inst_shape = "{}{}{}".format(*self.tile_description.math_instruction.instruction_shape)
+            if (
+                self.tile_description.math_instruction.element_a != self.A.element
+                and self.tile_description.math_instruction.element_a != self.accumulator_type()
+            ):
+                intermediate_type = DataTypeNames[self.tile_description.math_instruction.element_a]
+        else:
+            inst_shape = ""
+
+        return f"{ShortDataTypeNames[self.accumulator_type()]}{inst_shape}{intermediate_type}{ConvKindNames[self.conv_kind]}_{IteratorAlgorithmNames[self.iterator_algorithm]}"
+
+    def extended_name(self):
+        """Append data types if they differ from compute type."""
+        if (
+            self.C.element != self.tile_description.math_instruction.element_accumulator
+            and self.A.element != self.tile_description.math_instruction.element_accumulator
+        ):
+            extended_name = "${element_c}_${core_name}_${element_a}"
+        elif (
+            self.C.element == self.tile_description.math_instruction.element_accumulator
+            and self.A.element != self.tile_description.math_instruction.element_accumulator
+        ):
+            extended_name = "${core_name}_${element_a}"
+        else:
+            extended_name = "${core_name}"
+
+        extended_name = substitute_template(
+            extended_name,
+            {
+                "element_a": DataTypeNames[self.A.element],
+                "element_c": DataTypeNames[self.C.element],
+                "core_name": self.core_name(),
+            },
+        )
+
+        return extended_name
+
+    def layout_name(self):
+        return f"{ShortLayoutTypeNames[self.A.layout]}"
+
+    def procedural_name(self):
+        """
+        The full procedural name indicates architecture, extended name, tile size, and layout.
+        """
+        opcode_class_name = OpcodeClassNames[self.tile_description.math_instruction.opcode_class]
+
+        threadblock = f"{self.tile_description.threadblock_shape[0]}x{self.tile_description.threadblock_shape[1]}_{self.tile_description.threadblock_shape[2]}x{self.tile_description.stages}"
+
+        if self.stride_support == StrideSupport.Unity:
+            configuration_name = (
+                "cutlass_${opcode_class}_${extended_name}_${threadblock}"
+                "_${layout}_align${alignment}_unity_stride"
+            )
+        else:
+            configuration_name = (
+                "cutlass_${opcode_class}_${extended_name}_${threadblock}"
+                "_${layout}_align${alignment}"
+            )
+
+        if self.split_k_slices > 1:
+            configuration_name += f"_splitk{self.split_k_slices}"
+
+        return substitute_template(
+            configuration_name,
+            {
+                "opcode_class": opcode_class_name,
+                "extended_name": self.extended_name(),
+                "threadblock": threadblock,
+                "layout": self.layout_name(),
+                "alignment": f"{self.A.alignment}",
+            },
+        )
+
+
+class EmitConv2dInstance:
+    """Responsible for emitting a CUTLASS template definition."""
+
+    def __init__(self):
+        self.epilogue_default = """
+    ${epilogue_functor}<
+      ${element_c},
+      ${epilogue_vector_length},
+      ${element_accumulator},
+      ${element_epilogue}
+    >"""
+
+        self.epilogue_no_beta_scaling = """
+    ${epilogue_functor}<
+      ${element_c},
+      ${epilogue_vector_length},
+      ${element_accumulator},
+      ${element_epilogue},
+      cutlass::epilogue::thread::ScaleType::NoBetaScaling
+    >"""
+
+        self.epilogue_residual_block = """
+    ${epilogue_functor}<
+      ${element_c},
+      ${element_accumulator},
+      ${element_epilogue},
+      ${element_c},
+      ${epilogue_vector_length},
+      ${activation},
+      ${binary_op},
+      ${unary_op}
+    >"""
+
+        self.epilogue_wgrad = """
+    ${epilogue_functor}<
+      ${element_c},
+      4,
+      float,
+      float
+    >"""
+
+        self.template = """
+  // Conv2d${conv_kind_name} ${iterator_algorithm_name} kernel instance "${operation_name}"
+  using ${operation_name} =
+  typename cutlass::conv::kernel::DefaultConv2d${conv_kind_name}${conv_kernel_postfix}<
+    ${element_a},
+    ${layout_a},
+    ${element_b},
+    ${layout_b},
+    ${element_c},
+    ${layout_c},
+    ${element_accumulator},
+    ${opcode_class},
+    ${arch},
+    cutlass::gemm::GemmShape<${threadblock_shape_m}, ${threadblock_shape_n}, ${threadblock_shape_k}>,
+    cutlass::gemm::GemmShape<${warp_shape_m}, ${warp_shape_n}, ${warp_shape_k} >,
+    cutlass::gemm::GemmShape<${instruction_shape_m}, ${instruction_shape_n}, ${instruction_shape_k}>,
+    ${epilogue},
+    ${swizzling_functor}, // cutlass::gemm::threadblock::GemmSplitKIdentityThreadblockSwizzle<>,
+    ${stages},
+    ${math_operator},
+    ${iterator_algorithm},
+    ${stride_support},
+    ${align_a},
+    ${align_b}
+  >::Kernel;
+
+  ${reduction}
+"""
+
+        self.reduction_template = """
+using EpilogueOutputOp = ${epilogue};
+using ReductionOp = cutlass::reduction::thread::ReduceAdd<
+    ${element_accumulator},
+    ${element_accumulator},
+      EpilogueOutputOp::kCount
+      >;
+
+using ReductionKernel = cutlass::reduction::kernel::ReduceSplitK<
+    cutlass::MatrixShape<4, 32 * EpilogueOutputOp::kCount>,
+      EpilogueOutputOp,
+      ReductionOp
+      >;
+
+using ReductionDevice = cutlass::reduction::device::ReduceSplitK<ReductionKernel>;
+using ReductionStrideIndex = typename ReductionDevice::StrideIndex;
+"""
+
+    def emit(
+        self, operation, no_beta_scaling=False, residual_block_info=False, emit_reduction=False
+    ):
+        """Instantiate a Conv2d kernel from given `operation`."""
+        warp_shape = [
+            int(
+                operation.tile_description.threadblock_shape[idx]
+                / operation.tile_description.warp_count[idx]
+            )
+            for idx in range(3)
+        ]
+
+        epilogue_vector_length = int(
+            min(operation.C.alignment * DataTypeSize[operation.C.element], 128)
+            / DataTypeSize[operation.C.element]
+        )
+
+        element_c = operation.C.element
+        use_split_k_wgrad = operation.conv_kind == ConvKind.Wgrad and operation.split_k_slices > 1
+        # Gemm output always fp32 in wgrad with split k
+        element_c_gemm = DataType.f32 if use_split_k_wgrad else element_c
+
+        if emit_reduction:
+            epilogue_reduction = substitute_template(
+                self.epilogue_wgrad,
+                {
+                    "epilogue_functor": EpilogueFunctorTag[operation.epilogue_functor],
+                    "element_c": DataTypeTag[element_c],
+                },
+            )
+            reduction = substitute_template(
+                self.reduction_template,
+                {
+                    "epilogue": epilogue_reduction,
+                    "operation_name": operation.procedural_name(),
+                    "element_accumulator": DataTypeTag[operation.accumulator_type()],
+                },
+            )
+            gemm_template = substitute_template(self.template, {"reduction": reduction})
+        else:
+            gemm_template = substitute_template(self.template, {"reduction": ""})
+
+        values = {
+            "operation_name": operation.procedural_name(),
+            "conv_kind": ConvKindTag[operation.conv_kind],
+            "conv_kind_name": ConvKindNames[operation.conv_kind].capitalize(),
+            "element_a": DataTypeTag[operation.A.element],
+            "layout_a": LayoutTag[operation.A.layout],
+            "element_b": DataTypeTag[operation.B.element],
+            "layout_b": LayoutTag[operation.B.layout],
+            "element_c": DataTypeTag[element_c_gemm],
+            "layout_c": LayoutTag[operation.C.layout],
+            "element_accumulator": DataTypeTag[operation.accumulator_type()],
+            "opcode_class": OpcodeClassTag[
+                operation.tile_description.math_instruction.opcode_class
+            ],
+            "arch": f"cutlass::arch::Sm{operation.arch}",
+            "threadblock_shape_m": str(operation.tile_description.threadblock_shape[0]),
+            "threadblock_shape_n": str(operation.tile_description.threadblock_shape[1]),
+            "threadblock_shape_k": str(operation.tile_description.threadblock_shape[2]),
+            "warp_shape_m": str(warp_shape[0]),
+            "warp_shape_n": str(warp_shape[1]),
+            "warp_shape_k": str(warp_shape[2]),
+            "instruction_shape_m": str(
+                operation.tile_description.math_instruction.instruction_shape[0]
+            ),
+            "instruction_shape_n": str(
+                operation.tile_description.math_instruction.instruction_shape[1]
+            ),
+            "instruction_shape_k": str(
+                operation.tile_description.math_instruction.instruction_shape[2]
+            ),
+            "epilogue_vector_length": str(epilogue_vector_length),
+            "epilogue_functor": EpilogueFunctorTag[operation.epilogue_functor],
+            "element_epilogue": str(DataTypeTag[operation.element_epilogue]),
+            "swizzling_functor": SwizzlingFunctorTag[operation.swizzling_functor],
+            "stages": str(operation.tile_description.stages),
+            "iterator_algorithm": IteratorAlgorithmTag[operation.iterator_algorithm],
+            "iterator_algorithm_name": IteratorAlgorithmNames[
+                operation.iterator_algorithm
+            ].capitalize(),
+            "stride_support": StrideSupportTag[operation.stride_support],
+            "math_operator": MathOperationTag[
+                operation.tile_description.math_instruction.math_operation
+            ],
+            "align_a": str(operation.A.alignment),
+            "align_b": str(operation.B.alignment),
+            "conv_kernel_postfix": "",
+        }
+
+        if use_split_k_wgrad:
+            # Even if the output is fp16, gemm output is always fp32 for split k wgrad.
+            epilogue_gemm = substitute_template(
+                self.epilogue_wgrad,
+                {
+                    "epilogue_functor": EpilogueFunctorTag[operation.epilogue_functor],
+                    "element_c": "float",
+                },
+            )
+            template = substitute_template(gemm_template, {"epilogue": epilogue_gemm})
+        elif residual_block_info:
+            template = substitute_template(
+                gemm_template, {"epilogue": self.epilogue_residual_block}
+            )
+            values.update(
+                {
+                    "unary_op": residual_block_info["unary_op"],
+                    "binary_op": residual_block_info["binary_op"],
+                    "activation": residual_block_info["activation"],
+                    "conv_kernel_postfix": "WithBroadcast",
+                }
+            )
+        elif no_beta_scaling:
+            template = substitute_template(
+                gemm_template, {"epilogue": self.epilogue_no_beta_scaling}
+            )
+        else:
+            template = substitute_template(gemm_template, {"epilogue": self.epilogue_default})
+
+        return substitute_template(template, values)
+
+
+def instantiate_conv2d_template(attrs):
+    """Return CUTLASS host code for conv2d based on a template and the provided attribute map."""
+    template = """
+    ${cutlass_op_def}
+
+  using Conv2d = cutlass::conv::device::ImplicitGemmConvolution<${cutlass_op_name}>;
+  using ElementInputA = Conv2d::ElementA;
+  using ElementInputB = Conv2d::ElementB;
+  using ElementComputeEpilogue = Conv2d::ElementAccumulator;
+  int N = ${N};
+  int H = ${H};
+  int W = ${W};
+  int C = ${C};
+  int K = ${K};
+  int R = ${R};
+  int S = ${S};
+  int P = ${P};
+  int Q = ${Q};
+  int pad_h = ${pad_h};
+  int pad_w = ${pad_w};
+  int stride_h = ${stride_h};
+  int stride_w = ${stride_w};
+  int dilation_h = ${dilation_h};
+  int dilation_w = ${dilation_w};
+  int split_k_slices = ${split_k_slices};
+  cutlass::conv::Conv2dProblemSize problem_size(N, H, W, C, K, R, S, P, Q, pad_h, pad_w, stride_h, stride_w, dilation_h, dilation_w, cutlass::conv::Mode::kCrossCorrelation, split_k_slices);
+  const cutlass::conv::SplitKMode split_k_mode = cutlass::conv::SplitKMode::${split_k_mode};
+
+  void* ptr_a = (void*)(${data_arg}->data);
+  void* ptr_b = (void*)(${weight_arg}->data);
+  ${bias_decl}
+  ${residual_decl}
+  void* ptr_out = (void*)(out0->data);
+
+  ElementComputeEpilogue alpha = ElementComputeEpilogue(1);
+  ElementComputeEpilogue beta = ElementComputeEpilogue(${beta});
+  using cutlass::layout::TensorNHWC;
+  auto activation_shape = TensorNHWC::packed(cutlass::make_Coord(N, H, W, C));
+  auto weight_shape = TensorNHWC::packed(cutlass::make_Coord(K, R, S, C));
+  auto output_shape = TensorNHWC::packed(cutlass::make_Coord(N, P, Q, K));
+  ${residual_shape_decl}
+
+  TensorNHWC layout_A(${A_shape});
+  TensorNHWC layout_B(${B_shape});
+  TensorNHWC layout_C(${C_shape});
+  TensorNHWC layout_D(${D_shape});
+
+  using ElementOutput = ${ElementOutput};
+  cutlass::TensorRef<ElementOutput, TensorNHWC> tensor_c{static_cast<ElementOutput*>(${tensor_c}), ${tensor_c_layout}};
+  cutlass::TensorRef<ElementOutput, TensorNHWC> tensor_d{static_cast<ElementOutput*>(ptr_out), layout_D};
+  typename Conv2d::Arguments arguments{
+   problem_size,
+   {static_cast<ElementInputA*>(ptr_a), layout_A},
+   {static_cast<ElementInputB*>(ptr_b), layout_B},
+   ${tensor_c_arg},
+   ${tensor_d_arg},
+   {${alpha_beta}},
+   split_k_mode
+   ${additional_args}
+ };
+  Conv2d conv2d_op;
+  size_t workspace_size = conv2d_op.get_workspace_size(arguments);
+  cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);
+  cutlass::Status status = conv2d_op.can_implement(arguments);
+  TVM_FFI_ICHECK(status == cutlass::Status::kSuccess);
+  ${split_k_reset}
+  status = conv2d_op.initialize(arguments, workspace.get());
+  TVM_FFI_ICHECK(status == cutlass::Status::kSuccess);
+  ${split_k_update}
+
+  cudaStream_t stream = static_cast<cudaStream_t>(TVMFFIEnvGetStream(kDLCUDA, ${data_arg}->device.device_id));
+
+  status = conv2d_op(stream);
+  TVM_FFI_ICHECK(status == cutlass::Status::kSuccess);
+  ${split_k_reduction}
+"""
+
+    split_k_reset = """
+    arguments.ref_D.reset(reinterpret_cast<ElementComputeEpilogue*>(workspace.get()), layout_D);
+"""
+
+    split_k_update = """
+  arguments.output_op = {ElementComputeEpilogue(1), ElementComputeEpilogue(0)};
+  status = conv2d_op.update(arguments, workspace.get());
+  TVM_FFI_ICHECK(status == cutlass::Status::kSuccess);
+"""
+
+    split_k_reduction = """
+  ReductionDevice reduction_op;
+  const static cutlass::conv::Operator kConvolutionalOperator = Conv2d::kConvolutionalOperator;
+  typename ReductionDevice::Arguments reduction_args(
+     cutlass::conv::implicit_gemm_problem_size(kConvolutionalOperator, problem_size).mn(),
+     problem_size.split_k_slices,
+     cutlass::conv::implicit_gemm_tensor_c_size(kConvolutionalOperator, problem_size),
+     {
+      reinterpret_cast<Conv2d::ElementAccumulator*> (workspace.get()),
+      ReductionStrideIndex(tensor_c.stride()[Conv2d::UnderlyingKernel::kTensorCStrideIdx])
+     },
+     {
+      tensor_d.data(),
+      ReductionStrideIndex(tensor_d.stride()[Conv2d::UnderlyingKernel::kTensorCStrideIdx])
+     },
+     {
+      tensor_c.data(),
+      ReductionStrideIndex(tensor_c.stride()[Conv2d::UnderlyingKernel::kTensorCStrideIdx])
+     },
+     {alpha, beta}
+  );
+  status = reduction_op.initialize(reduction_args, nullptr);
+  status = reduction_op();
+"""
+    op_type = attrs["op_type"]
+    has_bias = "bias" in op_type
+    use_split_k = "splitk" in attrs["cutlass_op_name"]
+    is_wgrad = "backward_weight" in op_type
+    is_dgrad = "conv2d_transpose" in op_type
+    has_residual_block = "residual" in op_type
+    no_bias_scaling = op_type not in [
+        "cutlass.conv2d_bias_sigmoid",
+        "cutlass.conv2d_bias_silu",
+        "cutlass.conv2d_bias_hardswish",
+    ]
+
+    aux_map = {}
+
+    if (not has_bias or no_bias_scaling) and not has_residual_block:
+        aux_map["beta"] = 0
+    else:
+        aux_map["beta"] = 1
+
+    if has_residual_block:
+        aux_map["bias_decl"] = "void* ptr_bias = (void*)(${bias_arg}->data);\n"
+        aux_map["residual_decl"] = "void* ptr_residual = (void*)(${residual_arg}->data);"
+        aux_map["tensor_c"] = "ptr_residual"
+        aux_map["tensor_c_layout"] = "layout_C"
+    elif has_bias:
+        aux_map["bias_decl"] = "void* ptr_c_bias = (void*)(${bias_arg}->data);\n"
+        aux_map["residual_decl"] = ""
+        aux_map["tensor_c"] = "ptr_c_bias"
+        aux_map["tensor_c_layout"] = "cutlass::layout::TensorNHWC::Stride(0)"
+    else:
+        aux_map["bias_decl"] = ""
+        aux_map["residual_decl"] = ""
+        aux_map["tensor_c"] = "ptr_out"
+        aux_map["tensor_c_layout"] = "layout_C"
+
+    if has_bias and no_bias_scaling and not has_residual_block:
+        aux_map["alpha_beta"] = "alpha"
+    else:
+        aux_map["alpha_beta"] = "alpha, beta"
+
+    if has_residual_block:
+        aux_map["additional_args"] = ", static_cast<ElementOutput*>(ptr_bias), nullptr, 0, K"
+    else:
+        aux_map["additional_args"] = ""
+
+    aux_map["residual_shape_decl"] = ""
+
+    if is_wgrad:
+        aux_map["A_shape"] = "output_shape"
+        aux_map["B_shape"] = "activation_shape"
+        aux_map["C_shape"] = "weight_shape"
+        aux_map["D_shape"] = "weight_shape"
+    elif is_dgrad:
+        aux_map["A_shape"] = "output_shape"
+        aux_map["B_shape"] = "weight_shape"
+        aux_map["C_shape"] = "activation_shape"
+        aux_map["D_shape"] = "activation_shape"
+    else:
+        aux_map["A_shape"] = "activation_shape"
+        aux_map["B_shape"] = "weight_shape"
+        aux_map["D_shape"] = "output_shape"
+
+        if has_residual_block:
+            res_shape = list(attrs.pop("residual_shape"))
+            shape_str = f"cutlass::make_Coord({res_shape[0]}, {res_shape[1]}, {res_shape[2]}, K)"
+            aux_map["residual_shape_decl"] = (
+                f"auto residual_shape = TensorNHWC::packed({shape_str});"
+            )
+            aux_map["C_shape"] = "residual_shape"
+
+            if res_shape == [int(attrs[c]) for c in ["N", "H", "W", "K"]]:
+                aux_map["tensor_c_layout"] = "layout_C"
+            else:
+                # bias-like residual input
+                aux_map["tensor_c_layout"] = "cutlass::layout::TensorNHWC::Stride(0)"
+        else:
+            aux_map["C_shape"] = "output_shape"
+
+    if use_split_k:
+        aux_map["ElementOutput"] = "EpilogueOutputOp::ElementOutput"
+        aux_map["tensor_c_arg"] = "{nullptr, TensorNHWC()}"
+        aux_map["tensor_d_arg"] = "{nullptr, TensorNHWC()}"
+        aux_map["split_k_reset"] = split_k_reset
+        aux_map["split_k_update"] = split_k_update
+        aux_map["split_k_reduction"] = split_k_reduction
+    else:
+        aux_map["ElementOutput"] = "Conv2d::ElementC"
+        aux_map["tensor_c_arg"] = "tensor_c"
+        aux_map["tensor_d_arg"] = "tensor_d"
+        aux_map["split_k_reset"] = aux_map["split_k_update"] = aux_map["split_k_reduction"] = ""
+
+    template = substitute_template(template, aux_map)
+
+    return substitute_template(template, attrs)

python/tvm/contrib/cutlass/conv2d_profiler.py

@@ -0,0 +1,216 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=import-outside-toplevel, invalid-name
+# ruff: noqa: E501
+"""Instantiate a C++ source for profiling CUTLASS kernels."""
+
+from .library import DataTypeTag
+
+
+class Conv2dProfilerEmitter:
+    """Emit a C++ source for profiling CUTLASS kernels."""
+
+    def __init__(self):
+        from jinja2 import Template
+
+        self.reduction = """
+      ReductionDevice reduction_op;
+      static cutlass::conv::Operator const kConvolutionalOperator = ImplicitGemm::kConvolutionalOperator;
+      typename ReductionDevice::Arguments reduction_args(
+        cutlass::conv::implicit_gemm_problem_size(kConvolutionalOperator, problem_size).mn(),
+        problem_size.split_k_slices,
+        cutlass::conv::implicit_gemm_tensor_c_size(kConvolutionalOperator, problem_size),
+        {
+           reinterpret_cast<ImplicitGemm::ElementC*> (workspace.get()),
+             ReductionStrideIndex(tensor_c.stride()[ImplicitGemm::UnderlyingKernel::kTensorCStrideIdx])
+             },
+        {
+           tensor_d.device_data(),
+             ReductionStrideIndex(tensor_d.stride()[ImplicitGemm::UnderlyingKernel::kTensorCStrideIdx])
+             },
+        {
+           tensor_c.device_data(),
+             ReductionStrideIndex(tensor_c.stride()[ImplicitGemm::UnderlyingKernel::kTensorCStrideIdx])
+             },
+        {ElementComputeEpilogue(1), ElementComputeEpilogue(0)}
+        );
+
+      reduction_op.initialize(reduction_args, nullptr);
+      reduction_op();
+"""
+
+        self.template = Template(
+            """
+#include <iostream>
+#include "cutlass/cutlass.h"
+#include "cutlass/conv/kernel/default_conv2d_fprop.h"
+#include "cutlass/conv/kernel/default_conv2d_wgrad.h"
+#include "cutlass/conv/kernel/default_conv2d_dgrad.h"
+#include "cutlass/conv/device/implicit_gemm_convolution.h"
+#include "cutlass/util/command_line.h"
+#include "cutlass/util/host_tensor.h"
+#include "cutlass/util/reference/host/tensor_fill.h"
+#include "cutlass/reduction/device/reduce_split_k.h"
+#include "cutlass/reduction/thread/reduction_operators.h"
+
+#define CUTLASS_CHECK(status)                                                                    \
+  {                                                                                              \
+    cutlass::Status error = status;                                                              \
+    if (error != cutlass::Status::kSuccess) {                                                    \
+      std::cerr << "Got cutlass error: " << cutlassGetStatusString(error) << " at: " << __LINE__ \
+                << std::endl;                                                                    \
+      exit(EXIT_FAILURE);                                                                        \
+    }                                                                                            \
+  }
+
+{{OperatorDef}}
+using ImplicitGemm = cutlass::conv::device::ImplicitGemmConvolution<{{OperatorName}}>;
+
+struct Options {
+  cutlass::Tensor4DCoord input_size;
+  cutlass::Tensor4DCoord filter_size;
+  cutlass::Tensor4DCoord padding;
+  cutlass::MatrixCoord conv_stride;
+  cutlass::MatrixCoord dilation;
+
+  void parse(int argc, char const **args) {
+    cutlass::CommandLine cmd(argc, args);
+    cmd.get_cmd_line_argument("n", input_size.n());
+    cmd.get_cmd_line_argument("h", input_size.h());
+    cmd.get_cmd_line_argument("w", input_size.w());
+    cmd.get_cmd_line_argument("c", input_size.c());
+    cmd.get_cmd_line_argument("k", filter_size.n());
+    cmd.get_cmd_line_argument("r", filter_size.h());
+    cmd.get_cmd_line_argument("s", filter_size.w());
+    int pad_h, pad_w, stride_h, stride_w, dilation_h, dilation_w;
+    cmd.get_cmd_line_argument("pad_h", pad_h);
+    cmd.get_cmd_line_argument("pad_w", pad_w);
+    cmd.get_cmd_line_argument("stride_h", stride_h);
+    cmd.get_cmd_line_argument("stride_w", stride_w);
+    cmd.get_cmd_line_argument("dilation_h", dilation_h);
+    cmd.get_cmd_line_argument("dilation_w", dilation_w);
+    filter_size.c() = input_size.c();
+    padding = {pad_h, pad_h, pad_w, pad_w};
+    conv_stride = {stride_h, stride_w};
+    dilation = {dilation_h, dilation_w};
+  }
+
+  cutlass::Tensor4DCoord output_size() const {
+    auto dilated_h = (filter_size.h() - 1) * dilation.row() + 1;
+    auto dilated_w = (filter_size.w() - 1) * dilation.column() + 1;
+    auto h = (input_size.h() + padding.n() + padding.h() - dilated_h) / conv_stride.row() + 1;
+    auto w = (input_size.w() + padding.w() + padding.c() - dilated_w) / conv_stride.column() + 1;
+    return cutlass::Tensor4DCoord(input_size.n(), h, w, filter_size.n());
+  }
+};
+
+double profile_convolution(Options const &options) {
+  using ElementOutput = {{ElementOutput}};
+  using ElementInputA = typename ImplicitGemm::ElementA;
+  using ElementInputB = typename ImplicitGemm::ElementB;
+
+  int split_k_slices = {{SplitK}};
+  cutlass::conv::Conv2dProblemSize problem_size(
+                        options.input_size,
+                        options.filter_size,
+                        options.padding,
+                        options.conv_stride,
+                        options.dilation,
+                        options.output_size(),
+                        cutlass::conv::Mode::kCrossCorrelation,
+                        split_k_slices
+                        );
+
+  auto conv_kind = ImplicitGemm::kConvolutionalOperator;
+  auto a_extent = implicit_gemm_tensor_a_extent(conv_kind, problem_size);
+  auto b_extent = implicit_gemm_tensor_b_extent(conv_kind, problem_size);
+  auto c_extent = implicit_gemm_tensor_c_extent(conv_kind, problem_size);
+
+  using LayoutC = typename ImplicitGemm::LayoutC;
+  cutlass::HostTensor<ElementInputA, typename ImplicitGemm::LayoutA> tensor_a(a_extent);
+  cutlass::HostTensor<ElementInputB, typename ImplicitGemm::LayoutB> tensor_b(b_extent);
+  cutlass::HostTensor<ElementOutput, typename ImplicitGemm::LayoutC> tensor_c(c_extent);
+  cutlass::HostTensor<ElementOutput, LayoutC> tensor_d(c_extent);
+  cutlass::HostTensor<ImplicitGemm::ElementC, LayoutC> tensor_c_gemm(c_extent);
+
+  using ElementComputeEpilogue = typename ImplicitGemm::ElementCompute;
+
+  cutlass::conv::SplitKMode const split_k_mode = split_k_slices > 1 ?
+      cutlass::conv::SplitKMode::kParallel : cutlass::conv::SplitKMode::kSerial;
+
+  typename ImplicitGemm::Arguments arguments{
+    problem_size,
+    tensor_a.device_ref(),
+    tensor_b.device_ref(),
+    tensor_c_gemm.device_ref(),
+    tensor_c_gemm.device_ref(),
+    {ElementComputeEpilogue(1), ElementComputeEpilogue(0)},
+    split_k_mode,
+  };
+
+  ImplicitGemm implicit_gemm_op;
+  size_t workspace_size = implicit_gemm_op.get_workspace_size(arguments);
+  cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);
+  auto status = implicit_gemm_op.can_implement(arguments);
+  CUTLASS_CHECK(status);
+
+  status = implicit_gemm_op.initialize(arguments, workspace.get());
+  CUTLASS_CHECK(status);
+  status = implicit_gemm_op();
+  CUTLASS_CHECK(status);
+
+  cudaEvent_t events[2];
+  for (auto & event : events) {
+    cudaEventCreate(&event);
+  }
+  cudaEventRecord(events[0]);
+
+  for (int iteration = 0; iteration < 100; ++iteration) {
+    auto status = implicit_gemm_op();
+    CUTLASS_CHECK(status);
+    {{Reduction}}
+  }
+
+  cudaEventRecord(events[1]);
+  cudaEventSynchronize(events[1]);
+  float runtime_ms = 0;
+  cudaEventElapsedTime(&runtime_ms, events[0], events[1]);
+
+  for (auto event : events) {
+    (void)cudaEventDestroy(event);
+  }
+  return double(runtime_ms) / 100.0;
+}
+
+int main(int argc, char const **args) {
+  Options options;
+  options.parse(argc, args);
+  std::cout << profile_convolution(options) << std::endl;
+  return 0;
+}
+"""
+        )
+
+    def emit(self, op_def, op_name, element_output, split_k_slices=1):
+        src = self.template.render(
+            OperatorDef=op_def,
+            OperatorName=op_name,
+            ElementOutput=DataTypeTag[element_output],
+            SplitK=split_k_slices,
+            Reduction=self.reduction if split_k_slices > 1 else "",
+        )
+        return src

python/tvm/contrib/cutlass/gemm_operation.py

@@ -0,0 +1,478 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name, unused-wildcard-import, wildcard-import, pointless-exception-statement
+# ruff: noqa: E501, F403, F405
+"""Generator for CUTLASS GEMM kernels."""
+
+from .library import *
+
+
+class GemmOperation:
+    """Describes various attributes for instantiating GEMM kernels."""
+
+    def __init__(
+        self,
+        arch,
+        tile_description,
+        A,
+        B,
+        C,
+        element_epilogue,
+        epilogue_functor=EpilogueFunctor.LinearCombination,
+        swizzling_functor=SwizzlingFunctor.Identity8,
+    ):
+        self.operation_kind = OperationKind.Gemm
+        self.arch = arch
+        self.tile_description = tile_description
+        self.A = A
+        self.B = B
+        self.C = C
+        self.element_epilogue = element_epilogue
+        self.epilogue_functor = epilogue_functor
+        self.swizzling_functor = swizzling_functor
+
+    def accumulator_type(self):
+        return self.tile_description.math_instruction.element_accumulator
+
+    def short_math_name(self):
+        return ShortDataTypeNames[self.accumulator_type()]
+
+    def core_name(self):
+        """The basic operation kind is prefixed with a letter indicating the accumulation type."""
+        inst_shape = ""
+        intermediate_type = ""
+
+        if (
+            self.tile_description.math_instruction.opcode_class == OpcodeClass.TensorOp
+            or self.tile_description.math_instruction.opcode_class == OpcodeClass.WmmaTensorOp
+        ):
+            inst_shape = "{}{}{}".format(*self.tile_description.math_instruction.instruction_shape)
+            if (
+                self.tile_description.math_instruction.element_a != self.A.element
+                and self.tile_description.math_instruction.element_a
+                != self.tile_description.math_instruction.element_accumulator
+            ):
+                intermediate_type = DataTypeNames[self.tile_description.math_instruction.element_a]
+
+        return f"{self.short_math_name()}{inst_shape}{intermediate_type}gemm"
+
+    def extended_name(self):
+        """Append data types if they differ from compute type."""
+        if (
+            self.C.element != self.tile_description.math_instruction.element_accumulator
+            and self.A.element != self.tile_description.math_instruction.element_accumulator
+        ):
+            extended_name = "${element_c}_${core_name}_${element_a}"
+        elif (
+            self.C.element == self.tile_description.math_instruction.element_accumulator
+            and self.A.element != self.tile_description.math_instruction.element_accumulator
+        ):
+            extended_name = "${core_name}_${element_a}"
+        else:
+            extended_name = "${core_name}"
+
+        extended_name = substitute_template(
+            extended_name,
+            {
+                "element_a": DataTypeNames[self.A.element],
+                "element_c": DataTypeNames[self.C.element],
+                "core_name": self.core_name(),
+            },
+        )
+
+        return extended_name
+
+    def layout_name(self):
+        return f"{ShortLayoutTypeNames[self.A.layout]}{ShortLayoutTypeNames[self.B.layout]}"
+
+    def procedural_name(self):
+        """The full procedural name indicates architecture, extended name, tile size,
+        and layout.
+        """
+        threadblock = self.tile_description.procedural_name()
+        opcode_class_name = OpcodeClassNames[self.tile_description.math_instruction.opcode_class]
+
+        return substitute_template(
+            "cutlass_${opcode_class}_${extended_name}_${threadblock}_${layout}_align${alignment}",
+            {
+                "opcode_class": opcode_class_name,
+                "extended_name": self.extended_name(),
+                "threadblock": threadblock,
+                "layout": self.layout_name(),
+                "alignment": f"{self.A.alignment}",
+            },
+        )
+
+    def leading_dim(self):
+        """lda, ldb, ldc, according to the leading dimension."""
+        if self.A.layout == LayoutType.RowMajor:
+            lda = "K"
+        elif self.A.layout == LayoutType.ColumnMajor:
+            lda = "M"
+        else:
+            ValueError("The layout of A is not implemented.")
+
+        if self.B.layout == LayoutType.RowMajor:
+            ldb = "N"
+        elif self.B.layout == LayoutType.ColumnMajor:
+            ldb = "K"
+        else:
+            ValueError("The layout of B is not implemented.")
+
+        if self.C.layout == LayoutType.RowMajor:
+            ldc = "N"
+        elif self.C.layout == LayoutType.ColumnMajor:
+            ldc = "M"
+        else:
+            ValueError("The layout of B is not implemented.")
+
+        return substitute_template(
+            "int lda = ${lda_val};\n\tint ldb = ${ldb_val};\n\tint ldc = ${ldc_val};\n",
+            {"lda_val": lda, "ldb_val": ldb, "ldc_val": ldc},
+        )
+
+
+class EmitGemmInstance:
+    """Responsible for emitting a CUTLASS template definition."""
+
+    def __init__(self):
+        self.epilogue_default = """
+    ${epilogue_functor}<
+      ${element_c},
+      ${epilogue_vector_length},
+      ${element_accumulator},
+      ${element_epilogue}
+    >"""
+
+        self.epilogue_no_beta_scaling = """
+    ${epilogue_functor}<
+      ${element_c},
+      ${epilogue_vector_length},
+      ${element_accumulator},
+      ${element_epilogue},
+      cutlass::epilogue::thread::ScaleType::NoBetaScaling
+    >"""
+
+        self.epilogue_residual_block = """
+    ${epilogue_functor}<
+      ${element_c},
+      ${element_accumulator},
+      ${element_epilogue},
+      ${element_c},
+      ${epilogue_vector_length},
+      ${activation},
+      ${binary_op},
+      ${unary_op}
+    >"""
+
+        self.gemm_template = """
+  // Gemm operator ${operation_name}
+  using Operation_${operation_name} = cutlass::gemm::device::${kernel_name}<
+    ${element_a}, ${layout_a},
+    ${element_b}, ${layout_b},
+    ${element_c}, ${layout_c},
+    ${element_accumulator},
+    ${opcode_class},
+    ${arch},
+    cutlass::gemm::GemmShape<${threadblock_shape_m}, ${threadblock_shape_n}, ${threadblock_shape_k}>,
+    cutlass::gemm::GemmShape<${warp_shape_m}, ${warp_shape_n}, ${warp_shape_k}>,
+    cutlass::gemm::GemmShape<${instruction_shape_m}, ${instruction_shape_n}, ${instruction_shape_k}>,
+    ${epilogue},
+    ${swizzling_functor},
+    ${stages},
+    ${align_a},
+    ${align_b}
+  >;
+"""
+
+    def emit(self, operation, no_beta_scaling=False, batched=False, residual_block_info=False):
+        """Instantiate a GEMM kernel from given `operation`."""
+        warp_shape = [
+            operation.tile_description.threadblock_shape[idx]
+            // operation.tile_description.warp_count[idx]
+            for idx in range(3)
+        ]
+        epilogue_vector_length = (
+            min(operation.C.alignment * DataTypeSize[operation.C.element], 128)
+            // DataTypeSize[operation.C.element]
+        )
+        values = {
+            "operation_name": operation.procedural_name(),
+            "element_a": DataTypeTag[operation.A.element],
+            "layout_a": LayoutTag[operation.A.layout],
+            "element_b": DataTypeTag[operation.B.element],
+            "layout_b": LayoutTag[operation.B.layout],
+            "element_c": DataTypeTag[operation.C.element],
+            "layout_c": LayoutTag[operation.C.layout],
+            "element_accumulator": DataTypeTag[operation.accumulator_type()],
+            "opcode_class": OpcodeClassTag[
+                operation.tile_description.math_instruction.opcode_class
+            ],
+            "arch": f"cutlass::arch::Sm{operation.arch}",
+            "threadblock_shape_m": str(operation.tile_description.threadblock_shape[0]),
+            "threadblock_shape_n": str(operation.tile_description.threadblock_shape[1]),
+            "threadblock_shape_k": str(operation.tile_description.threadblock_shape[2]),
+            "warp_shape_m": str(warp_shape[0]),
+            "warp_shape_n": str(warp_shape[1]),
+            "warp_shape_k": str(warp_shape[2]),
+            "instruction_shape_m": str(
+                operation.tile_description.math_instruction.instruction_shape[0]
+            ),
+            "instruction_shape_n": str(
+                operation.tile_description.math_instruction.instruction_shape[1]
+            ),
+            "instruction_shape_k": str(
+                operation.tile_description.math_instruction.instruction_shape[2]
+            ),
+            "epilogue_vector_length": str(epilogue_vector_length),
+            "element_epilogue": str(DataTypeTag[operation.element_epilogue]),
+            "epilogue_functor": EpilogueFunctorTag[operation.epilogue_functor],
+            "swizzling_functor": SwizzlingFunctorTag[operation.swizzling_functor],
+            "stages": str(operation.tile_description.stages),
+            "align_a": str(operation.A.alignment),
+            "align_b": str(operation.B.alignment),
+            "math_operation": MathOperationTag[
+                operation.tile_description.math_instruction.math_operation
+            ],
+        }
+
+        values["kernel_name"] = "GemmBatched" if batched else "Gemm"
+
+        if residual_block_info:
+            values["kernel_name"] = "GemmUniversalWithBroadcast"
+            template = substitute_template(
+                self.gemm_template, {"epilogue": self.epilogue_residual_block}
+            )
+            values.update(
+                {
+                    "unary_op": residual_block_info["unary_op"],
+                    "binary_op": residual_block_info["binary_op"],
+                    "activation": residual_block_info["activation"],
+                }
+            )
+        elif no_beta_scaling:
+            template = substitute_template(
+                self.gemm_template, {"epilogue": self.epilogue_no_beta_scaling}
+            )
+        else:
+            template = substitute_template(self.gemm_template, {"epilogue": self.epilogue_default})
+
+        return substitute_template(template, values)
+
+
+def instantiate_gemm_template(attrs):
+    """Return CUTLASS host code for GEMM based on a template and the provided attribute map."""
+
+    argument_template_default = """
+  typename ${kernel}::Arguments arguments{
+   problem_size,
+   {static_cast<ElementInputA*>(ptr_a), ${lda}}, ${batch_stride_A}
+   {static_cast<ElementInputB*>(ptr_b), ${ldb}}, ${batch_stride_B}
+   {static_cast<ElementOutput*>(${ptr_c}), ${c_stride}}, ${batch_stride_C}
+   {static_cast<ElementOutput*>(ptr_out), ${ldc}}, ${batch_stride_D}
+   {${alpha_beta}},
+   ${split_k_slices_or_batch}
+  };
+    """
+
+    # See cutlass/gemm/kernel/gemm_with_fused_epilogue.h
+    argument_template_residual = """
+  typename ${kernel}::Arguments arguments{
+    cutlass::gemm::GemmUniversalMode::${gemm_universal_mode},
+    problem_size,
+    ${split_k_slices_or_batch}, // batch_count
+    {${alpha_beta}},
+    static_cast<ElementInputA*>(ptr_a),
+    static_cast<ElementInputB*>(ptr_b),
+    static_cast<ElementOutput*>(ptr_residual),
+    static_cast<ElementOutput*>(ptr_out),
+    static_cast<ElementOutput*>(ptr_bias),
+    nullptr, // ptr_Tensor
+    ${batch_stride_A}
+    ${batch_stride_B}
+    ${batch_stride_C}
+    ${batch_stride_D}
+    0, // batch_stride_Vector,
+    0, // batch_stride_Tensor,
+    ${lda},
+    ${ldb},
+    ${ldc},
+    ${ldc},
+    0, // ldv, the stride for bias
+    0, // ldt
+  };
+    """
+
+    template = """
+  using ElementInputA = ${ElementInputA};
+  using ElementInputB = ${ElementInputB};
+  using ElementOutput = ${ElementOutput};
+  using ElementComputeEpilogue = ${ElementOutput};
+
+  ${cutlass_op_def}
+
+  using ${kernel} = Operation_${cutlass_op_name};
+  int M = ${M};
+  int N = ${N};
+  int K = ${K};
+  cutlass::gemm::GemmCoord problem_size(M, N, K);
+  ElementComputeEpilogue alpha = ElementComputeEpilogue(1);
+  ElementComputeEpilogue beta = ElementComputeEpilogue(${beta});
+  void* ptr_a = (void*)(${lhs_arg}->data);
+  void* ptr_b = (void*)(${rhs_arg}->data);
+  ${bias_decl}
+  ${residual_decl}
+  void* ptr_out = (void*)(out0->data);
+
+  ${argument}
+  size_t workspace_size = ${kernel}::get_workspace_size(arguments);
+  cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);
+  ${kernel} gemm_op;
+  cutlass::Status status = gemm_op.can_implement(arguments);
+  TVM_FFI_ICHECK(status == cutlass::Status::kSuccess);
+  status = gemm_op.initialize(arguments, workspace.get());
+  TVM_FFI_ICHECK(status == cutlass::Status::kSuccess);
+
+  cudaStream_t stream = static_cast<cudaStream_t>(TVMFFIEnvGetStream(kDLCUDA, ${A_arg}->device.device_id));
+
+  status = gemm_op(stream);
+  TVM_FFI_ICHECK(status == cutlass::Status::kSuccess);
+"""
+    op_type = attrs["op_type"]
+    has_bias = "bias" in op_type
+    is_gelu = "gelu" in op_type
+    batched = "batch" in attrs
+    has_residual_block = "residual" in op_type
+    aux_map = {"kernel": "Gemm"}
+
+    if has_bias:
+        aux_map.update(
+            {
+                "bias_decl": "void* ptr_bias = (void*)(${bias_arg}->data);\n",
+                "ptr_c": "ptr_bias",
+                "c_stride": (
+                    "(${bias_arg}->ndim == 1 ||"
+                    " ${bias_arg}->shape[${bias_arg}->ndim - 2] == 1) ? 0 : " + attrs["ldc"]
+                ),
+            }
+        )
+    else:
+        aux_map.update({"bias_decl": "", "ptr_c": "ptr_out", "c_stride": attrs["ldc"]})
+
+    if is_gelu or has_residual_block:
+        # GeLU epilogue does not compile with NoBetaScaling, so we explicitly specify the scale.
+        aux_map["beta"] = 1
+    else:
+        aux_map["beta"] = 0
+
+    if has_bias and not is_gelu and not has_residual_block:
+        aux_map["alpha_beta"] = "alpha"
+    else:
+        aux_map["alpha_beta"] = "alpha, beta"
+
+    for key in ["batch_stride_A", "batch_stride_B", "batch_stride_C"]:
+        if not batched and not has_residual_block:
+            aux_map[key] = ""
+        else:
+            aux_map[key] = attrs.get(key, "0") + ","
+
+    aux_map["batch_stride_D"] = aux_map["batch_stride_C"]
+    if has_bias and batched and not has_residual_block:
+        aux_map["batch_stride_C"] = "0,"
+
+    if batched:
+        attrs["split_k_slices_or_batch"] = attrs["batch"]
+    else:
+        attrs["split_k_slices_or_batch"] = 1
+
+    if has_residual_block:
+        template = substitute_template(template, {"argument": argument_template_residual})
+        aux_map["residual_decl"] = "void* ptr_residual = (void*)(${residual_arg}->data);\n"
+        aux_map["gemm_universal_mode"] = "kBatched" if batched else "kGemm"
+    else:
+        template = substitute_template(template, {"argument": argument_template_default})
+        aux_map["residual_decl"] = ""
+
+    template = substitute_template(template, aux_map)
+
+    return substitute_template(template, attrs)
+
+
+def emit_fp16A_intB_matmul(attrs):
+    """Return CUTLASS host code for fp16 A and int4 or int8 B GEMM."""
+    if attrs["group_size"] > 0:
+        attrs["quant_op"] = "cutlass::WeightOnlyQuantOp::FINEGRAINED_SCALE_ONLY"
+    else:
+        attrs["quant_op"] = "cutlass::WeightOnlyQuantOp::PER_COLUMN_SCALE_ONLY"
+        attrs["group_size"] = "k"
+
+    attrs["template_common"] = substitute_template(
+        """
+  using namespace fastertransformer;
+  constexpr auto QuantOp = ${quant_op};
+
+  int m = ${M};
+  int n = ${B_arg}->shape[1] * ${float_per_int};
+  int k = ${B_arg}->shape[0];
+
+  cudaStream_t stream = static_cast<cudaStream_t>(
+    TVMFFIEnvGetStream(kDLCUDA, ${A_arg}->device.device_id));
+    """,
+        attrs,
+    )
+
+    template = """
+  ${template_common}
+  gemm_fp16_int_bias_act<${weight_dtype}, QuantOp>(static_cast<cutlass::half_t*>(${A_arg}->data),
+                static_cast<${weight_dtype}*>(${B_arg}->data),
+                static_cast<cutlass::half_t*>(${scales_arg}->data),
+                ${bias},
+                static_cast<cutlass::half_t*>(out0->data),
+                "${activation}",
+                m, n, k, ${group_size}, ${bias_stride}, nullptr, 0, stream);
+"""
+
+    template_residual = """
+  ${template_common}
+  gemm_fp16_int_bias_act_residual<${weight_dtype}, QuantOp>(
+                static_cast<cutlass::half_t*>(${A_arg}->data),
+                static_cast<${weight_dtype}*>(${B_arg}->data),
+                static_cast<cutlass::half_t*>(${scales_arg}->data),
+                ${bias},
+                static_cast<cutlass::half_t*>(${residual_arg}->data),
+                static_cast<cutlass::half_t*>(out0->data),
+                "${activation}", "${binary_op}", "${unary_op}",
+                m, n, k, ${group_size}, nullptr, 0, stream);
+"""
+
+    if "residual_arg" in attrs:
+        if "bias_arg" in attrs:
+            bias = "static_cast<cutlass::half_t*>(${bias_arg}->data)"
+        else:
+            bias = "nullptr"
+
+        template_residual = substitute_template(template_residual, {"bias": bias})
+        return substitute_template(template_residual, attrs)
+
+    if "bias_arg" in attrs:
+        template = substitute_template(
+            template, {"bias": "static_cast<cutlass::half_t*>(${bias_arg}->data)"}
+        )
+    else:
+        template = substitute_template(template, {"bias": "nullptr"})
+
+    return substitute_template(template, attrs)

python/tvm/contrib/cutlass/gemm_profiler.py

@@ -0,0 +1,196 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=import-outside-toplevel, invalid-name
+"""Instantiate a C++ source for profiling CUTLASS kernels."""
+
+
+class GemmProfilerEmitter:
+    """Emit a C++ source for profiling CUTLASS kernels."""
+
+    def __init__(self):
+        from jinja2 import Template
+
+        self.template = Template(
+            """
+#include <iostream>
+#include <sstream>
+#include <vector>
+#include <chrono>
+
+#include "cuda_runtime.h"
+#include "cutlass/gemm/device/gemm.h"
+
+#define CUTLASS_CHECK(status)                                                                    \\
+  {                                                                                              \\
+    cutlass::Status error = status;                                                              \\
+    if (error != cutlass::Status::kSuccess) {                                                    \\
+      std::cerr << "Got cutlass error: " << cutlassGetStatusString(error) << " at: " << __LINE__ \\
+                << std::endl;                                                                    \\
+      exit(EXIT_FAILURE);                                                                        \\
+    }                                                                                            \\
+  }
+
+#define CUDA_CHECK(status)                                              \\
+  {                                                                     \\
+    cudaError_t error = status;                                         \\
+    if (error != cudaSuccess) {                                         \\
+      std::cerr << "Got bad CUDA status: " << cudaGetErrorString(error) \\
+                << " at line: " << __LINE__ << std::endl;               \\
+      exit(EXIT_FAILURE);                                               \\
+    }                                                                   \\
+  }
+
+template<typename DTypeA, typename DTypeB, typename DTypeC>
+cudaError_t CutlassGemm(
+    int M,
+    int N,
+    int K,
+    DTypeC alpha,
+    DTypeA const *A,
+    int lda,
+    DTypeB const *B,
+    int ldb,
+    DTypeC beta,
+    DTypeC *C,
+    int ldc) {
+  using namespace std::chrono;
+  {{OperatorDef}}
+  Operation_{{OperatorName}} gemm_operator;
+  Operation_{{OperatorName}}::Arguments args({M, N, K},
+                              {A, lda},
+                              {B, ldb},
+                              {C, ldc},
+                              {C, ldc},
+                              {alpha, beta});
+  cutlass::Status status = gemm_operator(args);
+  CUTLASS_CHECK(status)
+
+  high_resolution_clock::time_point t1 = high_resolution_clock::now();
+  for (int i = 0; i < 100; ++i) {
+    status = gemm_operator(args);
+  }
+  cudaDeviceSynchronize();
+  high_resolution_clock::time_point t2 = high_resolution_clock::now();
+  duration<double> time_span = duration_cast<duration<double>>(t2 - t1);
+  std::cout << time_span.count() << std::endl;
+  return cudaSuccess;
+}
+
+
+template<typename DType>
+cudaError_t AllocateMatrix(DType **matrix, int ldm, int rows, int columns, int seed = 0) {
+  cudaError_t result;
+
+  size_t sizeof_matrix = sizeof(DType) * rows * columns;
+
+  // Allocate device memory.
+  result = cudaMalloc(reinterpret_cast<void **>(matrix), sizeof_matrix);
+
+  if (result != cudaSuccess) {
+    std::cerr << "Failed to allocate matrix: "
+      << cudaGetErrorString(result) << std::endl;
+    return result;
+  }
+
+  // Clear the allocation.
+  result = cudaMemset(*matrix, 0, sizeof_matrix);
+
+  if (result != cudaSuccess) {
+    std::cerr << "Failed to clear matrix device memory: "
+      << cudaGetErrorString(result) << std::endl;
+    return result;
+  }
+
+  if (result != cudaSuccess) {
+    std::cerr << "Failed to initialize matrix: "
+      << cudaGetErrorString(result) << std::endl;
+    return result;
+  }
+
+  return result;
+}
+
+template<typename DTypeA, typename DTypeB, typename DTypeC>
+cudaError_t TestCutlassGemm(int M, int N, int K, DTypeC alpha, DTypeC beta) {
+  cudaError_t result;
+
+  {{LeadingDim}}
+  // size_t sizeof_C = sizeof(DTypeC) * ldc * N;
+  DTypeA *A;
+  DTypeB *B;
+  DTypeC *C_cutlass;
+  result = AllocateMatrix<DTypeA>(&A, lda, M, K, 0);
+  if (result !=  cudaSuccess) {
+    return result;
+  }
+  result = AllocateMatrix<DTypeB>(&B, ldb, K, N, 17);
+  if (result !=  cudaSuccess) {
+    cudaFree(A);
+    return result;
+  }
+  result = AllocateMatrix<DTypeC>(&C_cutlass, ldc, M, N, 101);
+  if (result != cudaSuccess) {
+    cudaFree(A);
+    cudaFree(B);
+    return result;
+  }
+  result = CutlassGemm<DTypeA, DTypeB, DTypeC>(M, N, K, alpha, A, lda, B, ldb,
+                                                  beta, C_cutlass, ldc);
+  if (result != cudaSuccess) {
+    std::cerr << "CUTLASS GEMM kernel failed: "
+      << cudaGetErrorString(result) << std::endl;
+    cudaFree(C_cutlass);
+    cudaFree(B);
+    cudaFree(A);
+
+    return result;
+  }
+  cudaFree(C_cutlass);
+  cudaFree(B);
+  cudaFree(A);
+  return cudaSuccess;
+}
+
+int main(int argc, const char *arg[]) {
+  int problem[3] = { 4096, 4096, 4096 };
+  for (int i = 1; i < argc && i < 4; ++i) {
+    std::stringstream ss(arg[i]);
+    ss >> problem[i - 1];
+  }
+  float scalars[2] = { 1, 0 };
+  cudaError_t result = TestCutlassGemm< {{DTypeA}}, {{DTypeB}}, {{DTypeC}}>(
+    problem[0],     // GEMM M dimension
+    problem[1],     // GEMM N dimension
+    problem[2],     // GEMM K dimension
+    static_cast<{{DTypeC}}>(scalars[0]),     // alpha
+    static_cast<{{DTypeC}}>(scalars[1])      // beta
+  );
+  return result == cudaSuccess ? 0 : -1;
+}
+"""
+        )
+
+    def emit(self, op_name, op_def, dtype_a, dtype_b, dtype_c, ld):
+        src = self.template.render(
+            OperatorName=op_name,
+            OperatorDef=op_def,
+            DTypeA=dtype_a,
+            DTypeB=dtype_b,
+            DTypeC=dtype_c,
+            LeadingDim=ld,
+        )
+        return src

python/tvm/contrib/cutlass/gen_conv2d.py

@@ -0,0 +1,392 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name, dangerous-default-value
+# ruff: noqa: E501
+"""Conv2d kernel generator and profiler for CUTLASS."""
+
+import os
+import pickle
+from functools import partial
+
+from .conv2d_operation import Conv2dOperation, EmitConv2dInstance
+from .conv2d_profiler import Conv2dProfilerEmitter
+from .gen_gemm import CutlassGemmProfiler
+from .gen_tensor_op import EPILOGUE_MAP, GENERATOR_FUNC_TABLE, ProfilerEngine
+from .library import (
+    ConvKind,
+    DataType,
+    EpilogueFunctor,
+    IteratorAlgorithm,
+    LayoutType,
+    StrideSupport,
+    SwizzlingFunctor,
+    TensorDescription,
+)
+
+
+def create_conv2d_operator_with_epilogue(
+    conv_kind,
+    stride_support,
+    op_type,
+    tile_description,
+    data_type,
+    alignment,
+    alignment_epilogue,
+    swizzling_functor,
+    split_k_slices,
+):
+    """
+    Instantiate a cutlass kernel from the given configuration,
+    along with the epilouge functor
+    """
+    if "residual" in op_type:
+        activation_map = {
+            "cutlass.conv2d_bias_hardswish": "cutlass::epilogue::thread::HardSwish",
+            "cutlass.conv2d_bias_silu": "cutlass::epilogue::thread::SiLu",
+            "cutlass.conv2d_bias_sigmoid": "cutlass::epilogue::thread::Sigmoid",
+            "cutlass.conv2d_bias_relu": "cutlass::epilogue::thread::ReLu",
+            "cutlass.conv2d_bias": "cutlass::epilogue::thread::Identity",
+        }
+        prefix = op_type[: op_type.find("_residual")]
+        activation = activation_map[prefix]
+        binary_op = "cutlass::multiplies" if "residual_multiply" in op_type else "cutlass::plus"
+        unary_op = (
+            "cutlass::epilogue::thread::ReLu"
+            if op_type.endswith("relu")
+            else "cutlass::epilogue::thread::Identity"
+        )
+        residual_block_info = {
+            "activation": activation,
+            "binary_op": binary_op,
+            "unary_op": unary_op,
+        }
+        epilogue = EpilogueFunctor.LinearCombinationResidualBlock
+        no_beta_scaling = False
+    else:
+        residual_block_info = None
+        epilogue, no_beta_scaling = EPILOGUE_MAP[op_type]
+
+    element_a, element_b, element_c, element_epilogue = data_type
+
+    A = TensorDescription(element_a, LayoutType.TensorNHWC, alignment)
+    B = TensorDescription(element_b, LayoutType.TensorNHWC, alignment)
+    C = TensorDescription(element_c, LayoutType.TensorNHWC, alignment_epilogue)
+
+    op = Conv2dOperation(
+        conv_kind,
+        IteratorAlgorithm.Optimized,
+        tile_description.minimum_compute_capability,
+        tile_description,
+        A,
+        B,
+        C,
+        element_epilogue,
+        stride_support,
+        epilogue,
+        swizzling_functor,
+        split_k_slices,
+    )
+
+    name = op.procedural_name()
+    opdef = EmitConv2dInstance().emit(
+        op,
+        no_beta_scaling=no_beta_scaling,
+        residual_block_info=residual_block_info,
+        emit_reduction=split_k_slices > 1,
+    )
+
+    return name, opdef
+
+
+def enumerate_conv2d_operators(
+    conv_kind,
+    stride_support,
+    split_k_slices,
+    alignment_c,
+    tile_descriptions,
+    data_type,
+    alignment_constraints,
+    swizzling_functor=SwizzlingFunctor.Identity4,
+):
+    """Exhaustively instantiate all kernels from a given configuration."""
+    ret = []
+
+    kernel_emitter = EmitConv2dInstance()
+    profiler_emitter = Conv2dProfilerEmitter()
+
+    element_a, element_b, element_c, element_epilogue = data_type
+
+    if conv_kind == ConvKind.Dgrad and stride_support == StrideSupport.Strided:
+        swizzling_functor = SwizzlingFunctor.StridedDgradIdentity1
+
+    for split_k_slice in split_k_slices:
+        for tile in tile_descriptions:
+            for alignmentAB in alignment_constraints:
+                for alignmentC in alignment_c:
+                    A = TensorDescription(element_a, LayoutType.TensorNHWC, alignmentAB)
+                    B = TensorDescription(element_b, LayoutType.TensorNHWC, alignmentAB)
+                    C = TensorDescription(element_c, LayoutType.TensorNHWC, alignmentC)
+
+                    if element_c == DataType.s32 and A.alignment == 1:
+                        tile.threadblock_shape[0] = min(tile.threadblock_shape[0], 128)
+                        tile.threadblock_shape[1] = min(tile.threadblock_shape[1], 128)
+
+                    op = Conv2dOperation(
+                        conv_kind,
+                        IteratorAlgorithm.Optimized,
+                        tile.minimum_compute_capability,
+                        tile,
+                        A,
+                        B,
+                        C,
+                        element_epilogue,
+                        stride_support,
+                        EpilogueFunctor.LinearCombination,
+                        swizzling_functor,
+                        split_k_slice,
+                    )
+
+                    ret.append(
+                        {
+                            "src": profiler_emitter.emit(
+                                kernel_emitter.emit(op, emit_reduction=split_k_slice > 1),
+                                op.procedural_name(),
+                                element_output=element_c,
+                                split_k_slices=split_k_slice,
+                            ),
+                            "name": op.procedural_name(),
+                            "tile_description": tile,
+                            "alignment": alignmentAB,
+                            "alignment_epilogue": alignmentC,
+                            "data_type": data_type,
+                            "swizzle_functor": swizzling_functor,
+                            "split_k_slices": split_k_slice,
+                        }
+                    )
+
+    return ret
+
+
+class CutlassConv2DProfiler:
+    """Profile all candidate kernels and select the best one."""
+
+    def __init__(self, sm, cutlass_path, binary_path):
+        self.gemm_profiler = CutlassGemmProfiler(sm, cutlass_path, binary_path)
+        self.sm = sm
+        assert sm in GENERATOR_FUNC_TABLE, f"sm{sm} not supported yet."
+        self.engine = ProfilerEngine(sm, cutlass_path, binary_path)
+        self.cache_path = os.path.join(binary_path, "cutlass_conv2d_cache.pickle")
+        if os.path.exists(self.cache_path):
+            self.cache = pickle.load(open(self.cache_path, "rb"))
+        else:
+            self.cache = {}
+
+    def get_default(
+        self,
+        op_type,
+        out_dtype,
+        arg0_dtype,
+        arg1_dtype,
+        use_3xtf32,
+        conv_kind=ConvKind.Fprop,
+        stride=(1, 1),
+    ):
+        """Return the default kernel for the requested architecture.
+        For now, the default kernel was picked arbitrary.
+        """
+        gemm_profile_result = self.gemm_profiler.get_default(
+            op_type, out_dtype, arg0_dtype, arg1_dtype, use_3xtf32
+        )
+        tile_description = gemm_profile_result["tile_description"]
+        alignment = gemm_profile_result["alignment"]
+        data_type = gemm_profile_result["data_type"]
+        stride_support = StrideSupport.Strided if stride[0] > 1 else StrideSupport.Unity
+
+        if conv_kind == ConvKind.Dgrad and stride_support == StrideSupport.Strided:
+            swizzling_functor = SwizzlingFunctor.StridedDgradIdentity1
+        else:
+            swizzling_functor = SwizzlingFunctor.Identity4
+
+        name, opdef = create_conv2d_operator_with_epilogue(
+            conv_kind,
+            stride_support,
+            op_type,
+            tile_description,
+            data_type,
+            alignment,
+            alignment,
+            swizzling_functor,
+            split_k_slices=1,
+        )
+        return {"name": name, "opdef": opdef}
+
+    def select_op(
+        self,
+        d_shape,
+        w_shape,
+        padding,
+        stride,
+        dilation,
+        out_dtype,
+        data_dtype,
+        weight_dtype,
+        use_3xtf32,
+        conv_kind,
+        stride_support,
+        split_k_slices,
+        profile_all_alignments=False,
+        find_first_valid=False,
+        use_multiprocessing=False,
+    ):
+        """
+        Profile and select the best kernel from candidate kernels.
+        See the documentation for the profile method below.
+        """
+        N, H, W, IC = d_shape
+        OC, R, S, _ = w_shape
+
+        workload = (
+            N,
+            H,
+            W,
+            IC,
+            OC,
+            R,
+            S,
+            padding[0],
+            padding[1],
+            stride[0],
+            stride[1],
+            dilation[0],
+            dilation[1],
+        )
+
+        if workload in self.cache:
+            return self.cache[workload]
+
+        def alignments(dtype):
+            if dtype in ["float16"]:
+                alignments = [8, 4, 2, 1]
+            elif dtype in ["float", "float32"]:
+                alignments = [4, 2, 1]
+            else:
+                raise ValueError(f"Unsupported data type: {dtype}")
+            return alignments
+
+        alignments_c = [align for align in alignments(out_dtype) if OC % align == 0]
+
+        if not profile_all_alignments:
+            alignments_c = [alignments_c[0]]
+
+        ops = GENERATOR_FUNC_TABLE[self.sm](
+            out_dtype,
+            data_dtype,
+            weight_dtype,
+            partial(
+                enumerate_conv2d_operators,
+                conv_kind,
+                stride_support,
+                split_k_slices,
+                alignments_c,
+            ),
+            lambda align: all([dim % align == 0 for dim in [IC]]),
+            use_3xtf32,
+            profile_all_alignments,
+            # Use fp32 accumulation for wgrad to align with cuDNN
+            accumlator_dtype="float32" if conv_kind == ConvKind.Wgrad else out_dtype,
+        )
+
+        if not find_first_valid:
+            self.engine.compile_all(ops, use_multiprocessing)
+
+        args = "--n={} --h={} --w={} --c={} --k={} --r={} --s={} --pad_h={} --pad_w={} --stride_h={} --stride_w={} --dilation_h={} --dilation_w={}".format(
+            *workload
+        )
+
+        for op in ops:
+            out = self.engine.evaluate(op, args.split(" "))
+            op["runtime"] = out
+            if out < float("inf") and find_first_valid:
+                self.cache[workload] = op
+                return op
+
+        op = min(ops, key=lambda i: i["runtime"])
+        self.cache[workload] = op
+        with open(self.cache_path, "wb") as f:
+            pickle.dump(self.cache, f)
+        return op
+
+    def profile(
+        self,
+        op_type,
+        d_shape,
+        w_shape,
+        padding,
+        stride,
+        dilation,
+        out_dtype,
+        data_dtype,
+        weight_dtype,
+        use_3xtf32=True,
+        conv_kind=ConvKind.Fprop,
+        split_k_slices=[1],
+        profile_all_alignments=False,
+        find_first_valid=False,
+        use_multiprocessing=False,
+    ):
+        """Profile and select the best kernel from candidate kernels.
+        If find_first_valid is True, return immediately after the first applicable kernel is found.
+        If use_multiprocessing is True, compile all profiler executables in parallel.
+        """
+        # Dgrad requires Unity stride when stride == (1, 1)
+        stride_support = (
+            StrideSupport.Unity
+            if stride[0] == 1 and stride[1] == 1 and conv_kind == ConvKind.Dgrad
+            else StrideSupport.Strided
+        )
+
+        op = self.select_op(
+            d_shape,
+            w_shape,
+            padding,
+            stride,
+            dilation,
+            out_dtype,
+            data_dtype,
+            weight_dtype,
+            use_3xtf32,
+            conv_kind,
+            stride_support,
+            split_k_slices,
+            profile_all_alignments,
+            find_first_valid,
+            use_multiprocessing,
+        )
+
+        name, opdef = create_conv2d_operator_with_epilogue(
+            conv_kind,
+            stride_support,
+            op_type,
+            op["tile_description"],
+            op["data_type"],
+            op["alignment"],
+            op["alignment_epilogue"],
+            op["swizzle_functor"],
+            op["split_k_slices"],
+        )
+
+        return name, opdef, op["runtime"]

python/tvm/contrib/cutlass/gen_gemm.py

@@ -0,0 +1,352 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name
+"""GEMM kernel generator and profiler for CUTLASS."""
+
+import os
+import pickle
+from functools import partial
+
+from .gemm_operation import EmitGemmInstance, GemmOperation
+from .gemm_profiler import GemmProfilerEmitter
+from .gen_tensor_op import EPILOGUE_MAP, GENERATOR_FUNC_TABLE, ProfilerEngine
+from .library import (
+    DataType,
+    DataTypeTag,
+    EpilogueFunctor,
+    LayoutType,
+    SwizzlingFunctor,
+    TensorDescription,
+)
+
+
+def create_gemm_operator_with_epilogue(
+    op_type,
+    tile_description,
+    data_type,
+    alignment,
+    swizzling_functor,
+    batched=False,
+    layout_b=LayoutType.ColumnMajor,
+):
+    """
+    Instantiate a cutlass kernel from the given configuration,
+    along with the epilouge functor
+    """
+    element_a, element_b, element_c, element_epilogue = data_type
+
+    A = TensorDescription(element_a, LayoutType.RowMajor, alignment)
+    B = TensorDescription(element_b, layout_b, alignment)
+    C = TensorDescription(element_c, LayoutType.RowMajor, alignment)
+
+    if batched:
+        swizzling_functor = SwizzlingFunctor.Batched
+
+    if "residual" in op_type:
+        if "hardswish" in op_type:
+            activation = "cutlass::epilogue::thread::HardSwish"
+        elif "silu" in op_type:
+            activation = "cutlass::epilogue::thread::SiLu"
+        elif "sigmoid" in op_type:
+            activation = "cutlass::epilogue::thread::Sigmoid"
+        elif "gelu" in op_type:
+            activation = "cutlass::epilogue::thread::GELU"
+        elif "relu" in op_type:
+            activation = "cutlass::epilogue::thread::ReLu"
+        else:
+            activation = "cutlass::epilogue::thread::Identity"
+
+        binary_op = "cutlass::multiplies" if "residual_multiply" in op_type else "cutlass::plus"
+        unary_op = (
+            "cutlass::epilogue::thread::ReLu"
+            if op_type.endswith("relu")
+            else "cutlass::epilogue::thread::Identity"
+        )
+        residual_block_info = {
+            "activation": activation,
+            "binary_op": binary_op,
+            "unary_op": unary_op,
+        }
+        epilogue = EpilogueFunctor.LinearCombinationResidualBlock
+        no_beta_scaling = False
+    else:
+        residual_block_info = None
+        epilogue, no_beta_scaling = EPILOGUE_MAP[op_type]
+
+    op = GemmOperation(
+        tile_description.minimum_compute_capability,
+        tile_description,
+        A,
+        B,
+        C,
+        element_epilogue,
+        epilogue,
+        swizzling_functor,
+    )
+
+    return (
+        op.procedural_name(),
+        EmitGemmInstance().emit(
+            op,
+            no_beta_scaling=no_beta_scaling,
+            batched=batched,
+            residual_block_info=residual_block_info,
+        ),
+    )
+
+
+def enumerate_gemm_operators(
+    tile_descriptions,
+    data_type,
+    alignment_constraints,
+    swizzling_functor=SwizzlingFunctor.Identity8,
+    layout_b=LayoutType.ColumnMajor,
+):
+    """Exhaustively instantiate all kernels from a given configuration."""
+    ret = []
+    kernel_emitter = EmitGemmInstance()
+    profiler_emitter = GemmProfilerEmitter()
+
+    element_a, element_b, element_c, element_epilogue = data_type
+
+    for tile_description in tile_descriptions:
+        for alignment in alignment_constraints:
+            A = TensorDescription(element_a, LayoutType.RowMajor, alignment)
+            B = TensorDescription(element_b, layout_b, alignment)
+            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
+                )
+
+            op = GemmOperation(
+                tile_description.minimum_compute_capability,
+                tile_description,
+                A,
+                B,
+                C,
+                element_epilogue,
+                EpilogueFunctor.LinearCombination,
+                swizzling_functor,
+            )
+
+            src = profiler_emitter.emit(
+                op.procedural_name(),
+                kernel_emitter.emit(op, batched=False),
+                DataTypeTag[element_a],
+                DataTypeTag[element_b],
+                DataTypeTag[element_c],
+                op.leading_dim(),
+            )
+
+            ret.append(
+                {
+                    "src": src,
+                    "op": op,
+                    "name": op.procedural_name(),
+                    "tile_description": tile_description,
+                    "alignment": alignment,
+                    "data_type": data_type,
+                    "swizzle_functor": swizzling_functor,
+                }
+            )
+
+    return ret
+
+
+# TODO(masahi): A sensible way to pick reasonable default kernels
+DEFAULT_KERNELS = {
+    75: {
+        ("float16", "float16"): "cutlass_tensorop_h1688gemm_128x64_32x2_tn_align1",
+        ("float16", "float32"): "cutlass_tensorop_s1688gemm_f16_64x64_32x2_tn_align1",
+    },
+    # align1 variants do not seem to be available for sm80
+    80: {
+        ("float16", "float16"): "cutlass_tensorop_h1688gemm_128x64_32x2_tn_align1",
+        ("float16", "float32"): "cutlass_tensorop_s1688gemm_f16_64x64_32x2_tn_align1",
+        # two kernels for tf32 and 3xtf32
+        ("float32", "float32"): (
+            "cutlass_tensorop_s1688gemm_128x64_32x3_tn_align1",
+            "cutlass_tensorop_s1688gemm_64x64_16x3_tn_align1",
+        ),
+    },
+}
+
+
+class CutlassGemmProfiler:
+    """Profile all candidate kernels and select the best one."""
+
+    def __init__(self, sm, cutlass_path, binary_path):
+        assert sm in GENERATOR_FUNC_TABLE and sm in DEFAULT_KERNELS, f"sm{sm} not supported yet."
+        self.engine = ProfilerEngine(sm, cutlass_path, binary_path)
+        self.sm = sm
+        self.cache_path = os.path.join(binary_path, "cutlass_gemm_cache.pickle")
+        if os.path.exists(self.cache_path):
+            self.cache = pickle.load(open(self.cache_path, "rb"))
+        else:
+            self.cache = {}
+
+    def get_default(
+        self,
+        op_type,
+        out_dtype,
+        arg0_dtype,
+        arg1_dtype,
+        use_3xtf32=True,
+        batched=False,
+        layout_b=LayoutType.ColumnMajor,
+    ):
+        """Return the default kernel for the requested architecture.
+        For now, the default kernel was picked arbitrary.
+        """
+        ops = GENERATOR_FUNC_TABLE[self.sm](
+            out_dtype,
+            arg0_dtype,
+            arg1_dtype,
+            partial(enumerate_gemm_operators, layout_b=layout_b),
+            lambda align: align == 1,  # Only request align1 kernels
+            use_3xtf32,
+            profile_all_alignments=True,  # To include all align1 kernels
+            # TODO(masahi): Invesitigate when fp32 accumulation is needed for gemm
+            accumlator_dtype=out_dtype,
+        )
+
+        default_kernel_name = DEFAULT_KERNELS[self.sm][(arg0_dtype, out_dtype)]
+
+        if arg0_dtype == "float32":
+            default_kernel_name = (
+                default_kernel_name[0] if not use_3xtf32 else default_kernel_name[1]
+            )
+
+        filtered = list(filter(lambda op: op["name"] == default_kernel_name, ops))
+        assert len(filtered) == 1
+        op = filtered[0]
+        name, opdef = create_gemm_operator_with_epilogue(
+            op_type,
+            op["tile_description"],
+            op["data_type"],
+            op["alignment"],
+            op["swizzle_functor"],
+            batched=batched,
+            layout_b=layout_b,
+        )
+        op.update({"name": name, "opdef": opdef})
+        return op
+
+    def select_op(
+        self,
+        M,
+        N,
+        K,
+        out_dtype,
+        arg0_dtype,
+        arg1_dtype,
+        use_3xtf32,
+        profile_all_alignments=False,
+        find_first_valid=False,
+        use_multiprocessing=False,
+        layout_b=LayoutType.ColumnMajor,
+    ):
+        """
+        Profile and select the best kernel from candidate kernels.
+        See the documentation for the profile method below.
+        """
+        if (M, N, K) in self.cache:
+            op = self.cache[(M, N, K)]
+            return op
+
+        # TODO(masahi): CUTLASS alignment check on gemm kernels is too restrictive.
+        # See https://github.com/NVIDIA/cutlass/issues/362.
+        # When the above issue is resolved, we can remove the alignment check on M below.
+
+        ops = GENERATOR_FUNC_TABLE[self.sm](
+            out_dtype,
+            arg0_dtype,
+            arg1_dtype,
+            partial(enumerate_gemm_operators, layout_b=layout_b),
+            lambda align: all([dim % align == 0 for dim in [M, N, K]]),
+            use_3xtf32,
+            profile_all_alignments=profile_all_alignments,
+            # TODO(masahi): Invesitigate when fp32 accumulation is needed for gemm
+            accumlator_dtype=out_dtype,
+        )
+
+        if not find_first_valid:
+            self.engine.compile_all(ops, use_multiprocessing)
+
+        for op in ops:
+            out = self.engine.evaluate(op, [M, N, K])
+            op["runtime"] = out
+            if out < float("inf") and find_first_valid:
+                self.cache[(M, N, K)] = op
+                return op
+
+        op = min(ops, key=lambda i: i["runtime"])
+        self.cache[(M, N, K)] = op
+        with open(self.cache_path, "wb") as f:
+            pickle.dump(self.cache, f)
+        return op
+
+    def profile(
+        self,
+        op_type,
+        M,
+        N,
+        K,
+        out_dtype,
+        arg0_dtype,
+        arg1_dtype,
+        use_3xtf32=True,
+        profile_all_alignments=False,
+        find_first_valid=False,
+        use_multiprocessing=False,
+        batched=False,
+        layout_b=LayoutType.ColumnMajor,
+    ):
+        """Profile and select the best kernel from candidate kernels.
+        If find_first_valid is True, return immediately after the first applicable kernel is found.
+        If use_multiprocessing is True, compile all profiler executables in parallel.
+        """
+        op = self.select_op(
+            M,
+            N,
+            K,
+            out_dtype,
+            arg0_dtype,
+            arg1_dtype,
+            use_3xtf32,
+            profile_all_alignments=profile_all_alignments,
+            find_first_valid=find_first_valid,
+            use_multiprocessing=use_multiprocessing,
+            layout_b=layout_b,
+        )
+
+        name, opdef = create_gemm_operator_with_epilogue(
+            op_type,
+            op["tile_description"],
+            op["data_type"],
+            op["alignment"],
+            op["swizzle_functor"],
+            batched=batched,
+            layout_b=layout_b,
+        )
+
+        return name, opdef, op["runtime"]

python/tvm/contrib/cutlass/gen_tensor_op.py

@@ -0,0 +1,908 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name
+# ruff: noqa: F821
+"""Common functions and classes for CUTLASS GEMM and Conv2d geneator."""
+
+import logging
+import math
+import multiprocessing
+import os
+import re
+import subprocess
+import tempfile
+
+import tvm_ffi
+
+from tvm.runtime import Object
+from tvm.tirx import IntImm
+
+from . import _ffi_api as ffi
+from .attention_operation import (
+    instantiate_attention_template,
+    instantiate_flash_attention_template,
+    instantiate_flash_attention_var_len_template,
+)
+from .conv2d_operation import instantiate_conv2d_template
+from .gemm_operation import emit_fp16A_intB_matmul, instantiate_gemm_template
+from .layer_norm_operation import instantiate_layer_norm_template
+from .library import (
+    DataType,
+    DataTypeSize,
+    DataTypeTag,
+    EpilogueFunctor,
+    MathInstruction,
+    MathOperation,
+    OpcodeClass,
+    TileDescription,
+)
+from .rms_norm_operation import instantiate_rms_norm_template
+
+logger = logging.getLogger("cutlass")
+
+
+dtype_map = {
+    "int8": DataType.s8,
+    "uint8": DataType.u8,
+    "int32": DataType.s32,
+    "float32": DataType.f32,
+    "float16": DataType.f16,
+}
+
+
+def generate_tensor_op_common(
+    math_instructions, alignment_constraints, get_tile_descriptions, op_creator
+):
+    """Common kernel generator to be used by archtecture specific generators."""
+    ops = []
+    for math_inst in math_instructions:
+        tile_descriptions = get_tile_descriptions(math_inst)
+        data_type = [
+            math_inst.element_a,
+            math_inst.element_b,
+            math_inst.element_c,
+            math_inst.element_accumulator,
+        ]
+
+        out = op_creator(tile_descriptions, data_type, alignment_constraints)
+
+        ops.extend(out)
+
+    return ops
+
+
+def generate_sm50_simt(out_dtype, arg0_dtype, arg1_dtype, op_creator, accumulator_dtype="float32"):
+    """Gemerate GEMM or Conv2D SIMT kernels"""
+    # pylint: disable=unused-argument
+    min_cc = 50
+    max_cc = 1024
+    if arg0_dtype == "float32" and arg1_dtype == "float32":
+        assert out_dtype == "float32" and accumulator_dtype == "float32"
+        math_instructions = [
+            MathInstruction(
+                [1, 1, 1],
+                DataType.f32,
+                DataType.f32,
+                DataType.f32,
+                DataType.f32,
+                OpcodeClass.Simt,
+                MathOperation.multiply_add,
+            )
+        ]
+        alignment_constraints = [1]
+        tile_descriptions = [
+            ([128, 128, 8], 2, [4, 2, 1], min_cc, max_cc),
+            ([128, 64, 8], 2, [2, 2, 1], min_cc, max_cc),
+            ([64, 128, 8], 2, [2, 2, 1], min_cc, max_cc),
+            ([64, 64, 8], 2, [2, 1, 1], min_cc, max_cc),
+            ([128, 32, 8], 2, [2, 1, 1], min_cc, max_cc),
+            ([32, 128, 8], 2, [1, 2, 1], min_cc, max_cc),
+        ]
+
+        def get_tile_descriptions(math_inst):
+            return [
+                TileDescription(threadblock_shape, stages, warp_count, math_inst, min_cc, max_cc)
+                for threadblock_shape, stages, warp_count, min_cc, max_cc in tile_descriptions
+            ]
+
+        return generate_tensor_op_common(
+            math_instructions, alignment_constraints, get_tile_descriptions, op_creator
+        )
+    else:
+        raise NotImplementedError()
+
+
+def generate_sm75_tensor_op_1688(
+    out_dtype,
+    arg0_dtype,
+    arg1_dtype,
+    op_creator,
+    check_align,
+    _,
+    profile_all_alignments=False,
+    accumlator_dtype="float32",
+):
+    """Generate GEMM or Conv2D kernels for Turing."""
+    assert out_dtype in ["float32", "float16", "int32"]
+    min_cc = 75
+    max_cc = 1024
+
+    if arg0_dtype == "float16" and arg1_dtype == "float16":
+        math_instructions = [
+            MathInstruction(
+                [16, 8, 8],
+                DataType.f16,
+                DataType.f16,
+                dtype_map[out_dtype],
+                dtype_map[accumlator_dtype],
+                OpcodeClass.TensorOp,
+                MathOperation.multiply_add,
+            )
+        ]
+        alignment_constraints = [8, 4, 2, 1]
+        tile_descriptions = [
+            ([256, 128, 32], 2, [4, 2, 1], min_cc, max_cc),
+            ([128, 256, 32], 2, [2, 4, 1], min_cc, max_cc),
+            ([128, 128, 32], 2, [2, 2, 1], min_cc, max_cc),
+            ([64, 128, 32], 2, [2, 2, 1], min_cc, max_cc),
+            ([128, 64, 32], 2, [2, 2, 1], min_cc, max_cc),
+            ([64, 64, 32], 2, [2, 2, 1], min_cc, max_cc),
+            ([64, 128, 64], 2, [1, 2, 2], min_cc, max_cc),
+        ]
+
+    elif "int8" in arg0_dtype and "int8" in arg1_dtype:
+        assert out_dtype == "int32"
+        math_instructions = [
+            MathInstruction(
+                [8, 8, 16],
+                dtype_map[arg0_dtype],
+                dtype_map[arg1_dtype],
+                DataType.s32,
+                DataType.s32,
+                OpcodeClass.TensorOp,
+                MathOperation.multiply_add_saturate,
+            )
+        ]
+        alignment_constraints = [16, 8, 4, 2, 1]
+        tile_descriptions = [
+            ([256, 128, 64], 2, [4, 2, 1], min_cc, max_cc),
+            ([128, 256, 64], 2, [2, 4, 1], min_cc, max_cc),
+            ([128, 128, 64], 2, [2, 2, 1], min_cc, max_cc),
+            ([64, 256, 64], 2, [1, 4, 1], min_cc, max_cc),
+            ([256, 64, 64], 2, [4, 1, 1], min_cc, max_cc),
+            ([64, 128, 64], 2, [2, 2, 1], min_cc, max_cc),
+            ([128, 64, 64], 2, [2, 2, 1], min_cc, max_cc),
+            ([64, 64, 64], 2, [2, 2, 1], min_cc, max_cc),
+        ]
+    elif arg0_dtype == "float32" and arg1_dtype == "float32" and out_dtype == "float32":
+        return generate_sm50_simt(out_dtype, arg0_dtype, arg1_dtype, op_creator, accumlator_dtype)
+    else:
+        raise NotImplementedError()
+
+    alignment_constraints = [align for align in alignment_constraints if check_align(align)]
+    assert len(alignment_constraints) > 0
+
+    if not profile_all_alignments:
+        alignment_constraints = [alignment_constraints[0]]
+
+    def get_tile_descriptions(math_inst):
+        return [
+            TileDescription(threadblock_shape, stages, warp_count, math_inst, min_cc, max_cc)
+            for threadblock_shape, stages, warp_count, min_cc, max_cc in tile_descriptions
+        ]
+
+    return generate_tensor_op_common(
+        math_instructions, alignment_constraints, get_tile_descriptions, op_creator
+    )
+
+
+def generate_sm80_tensor_op_16816(
+    out_dtype,
+    arg0_dtype,
+    arg1_dtype,
+    op_creator,
+    check_align,
+    use_3xtf32=True,
+    profile_all_alignments=False,
+    accumlator_dtype="float32",
+):
+    """Generate GEMM or Conv2D kernels for Ampere."""
+    min_cc = 80
+    max_cc = 1024
+    max_cc_smem_limited = 80
+
+    def get_default_tile_descriptions(block_k_factor):
+        return [
+            ([128, 256, int(32 * block_k_factor)], 3, [2, 4, 1], min_cc, max_cc),
+            ([256, 128, int(32 * block_k_factor)], 3, [4, 2, 1], min_cc, max_cc),
+            ([256, 64, int(32 * block_k_factor)], 3, [4, 1, 1], min_cc, max_cc),
+            ([256, 64, int(32 * block_k_factor)], 4, [4, 1, 1], min_cc, max_cc),
+            ([64, 256, int(32 * block_k_factor)], 4, [1, 4, 1], min_cc, max_cc),
+            ([128, 128, int(32 * block_k_factor)], 3, [2, 2, 1], min_cc, max_cc),
+            ([128, 128, int(32 * block_k_factor)], 4, [2, 2, 1], min_cc, max_cc),
+            ([128, 128, int(32 * block_k_factor)], 5, [2, 2, 1], min_cc, max_cc),
+            ([128, 64, int(32 * block_k_factor)], 6, [2, 2, 1], min_cc, max_cc),
+            ([64, 128, int(32 * block_k_factor)], 6, [2, 2, 1], min_cc, max_cc),
+            ([64, 64, int(32 * block_k_factor)], 10, [2, 2, 1], min_cc, max_cc),
+            ([256, 128, int(64 * block_k_factor)], 3, [4, 2, 1], min_cc, max_cc_smem_limited),
+            ([128, 256, int(64 * block_k_factor)], 3, [2, 4, 1], min_cc, max_cc_smem_limited),
+            ([256, 64, int(64 * block_k_factor)], 4, [4, 1, 1], min_cc, max_cc_smem_limited),
+            ([64, 256, int(64 * block_k_factor)], 4, [1, 4, 1], min_cc, max_cc_smem_limited),
+            ([128, 128, int(64 * block_k_factor)], 4, [2, 2, 1], min_cc, max_cc),
+            ([256, 64, int(64 * block_k_factor)], 3, [4, 1, 1], min_cc, max_cc),
+            ([64, 256, int(64 * block_k_factor)], 3, [1, 4, 1], min_cc, max_cc),
+            ([128, 128, int(64 * block_k_factor)], 3, [2, 2, 1], min_cc, max_cc),
+            ([128, 64, int(64 * block_k_factor)], 3, [2, 2, 1], min_cc, max_cc),
+            ([64, 128, int(64 * block_k_factor)], 3, [2, 2, 1], min_cc, max_cc),
+            ([64, 64, int(64 * block_k_factor)], 5, [2, 2, 1], min_cc, max_cc),
+        ]
+
+    if arg0_dtype == "float16" and arg1_dtype == "float16":
+        math_instructions = [
+            MathInstruction(
+                [16, 8, 16],
+                DataType.f16,
+                DataType.f16,
+                dtype_map[out_dtype],
+                dtype_map[accumlator_dtype],
+                OpcodeClass.TensorOp,
+                MathOperation.multiply_add,
+            )
+        ]
+        alignment_constraints = [8, 4, 2]
+        tile_descriptions = get_default_tile_descriptions(1)
+    elif arg0_dtype == "float32" and arg1_dtype == "float32":
+        math_instructions = [
+            MathInstruction(
+                [16, 8, 8],
+                DataType.f32,
+                DataType.f32,
+                DataType.f32,
+                DataType.f32,
+                OpcodeClass.TensorOp,
+                MathOperation.multiply_add_fast_f32 if use_3xtf32 else MathOperation.multiply_add,
+            )
+        ]
+        alignment_constraints = [4, 2, 1]
+
+        if use_3xtf32:
+            # tf32
+            tile_descriptions = [
+                ([128, 128, 16], 4, [4, 2, 1], min_cc, max_cc),
+                ([128, 128, 16], 3, [4, 2, 1], min_cc, max_cc),
+                ([256, 64, 16], 3, [4, 2, 1], min_cc, max_cc),
+                ([64, 256, 16], 3, [2, 4, 1], min_cc, max_cc),
+                ([128, 64, 16], 4, [2, 2, 1], min_cc, max_cc),
+                ([64, 128, 16], 4, [2, 2, 1], min_cc, max_cc),
+                ([64, 64, 16], 3, [2, 2, 1], min_cc, max_cc),
+                ([128, 128, 32], 3, [4, 2, 1], min_cc, max_cc),
+                ([256, 64, 32], 3, [4, 2, 1], min_cc, max_cc_smem_limited),
+                ([64, 256, 32], 3, [2, 4, 1], min_cc, max_cc_smem_limited),
+                ([128, 64, 32], 3, [2, 2, 1], min_cc, max_cc),
+                ([64, 128, 32], 3, [2, 2, 1], min_cc, max_cc),
+                ([64, 64, 32], 3, [2, 2, 1], min_cc, max_cc),
+            ]
+        else:
+            tile_descriptions = get_default_tile_descriptions(0.5)
+    else:
+        assert out_dtype == "int32"
+        math_instructions = [
+            MathInstruction(
+                [16, 8, 32],
+                dtype_map[arg0_dtype],
+                dtype_map[arg1_dtype],
+                DataType.s32,
+                DataType.s32,
+                OpcodeClass.TensorOp,
+                MathOperation.multiply_add_saturate,
+            )
+        ]
+        alignment_constraints = [16, 8, 4]
+        tile_descriptions = get_default_tile_descriptions(2)
+
+    def get_tile_descriptions(math_inst):
+        return [
+            TileDescription(threadblock_shape, stages, warp_count, math_inst, min_cc, max_cc)
+            for threadblock_shape, stages, warp_count, min_cc, max_cc in tile_descriptions
+        ]
+
+    alignment_constraints = [align for align in alignment_constraints if check_align(align)]
+
+    if len(alignment_constraints) > 0 and not profile_all_alignments:
+        alignment_constraints = [alignment_constraints[0]]
+
+    if arg0_dtype != "float32" and arg1_dtype != "float32":
+        sm75_kernels = generate_sm75_tensor_op_1688(
+            out_dtype,
+            arg0_dtype,
+            arg1_dtype,
+            op_creator,
+            check_align,
+            False,
+            profile_all_alignments,
+            accumlator_dtype=accumlator_dtype,
+        )
+    else:
+        # TF32 (float32 + float32 case) is only supported on sm80
+        sm75_kernels = []
+
+    if len(alignment_constraints) > 0:
+        sm80_kernels = generate_tensor_op_common(
+            math_instructions, alignment_constraints, get_tile_descriptions, op_creator
+        )
+    else:
+        sm80_kernels = []
+
+    # TODO(masahi): For int8 kernels, The CUTLASS generator modifies the output tensor alignment
+    # after ops are created. Revisit how important this modification is.
+    # for op in operations:
+    #     if op.tile_description.threadblock_shape[1] >= 128:
+    #         op.C.alignment = 16
+    #     else:
+    #         op.C.alignment = 8
+
+    return sm75_kernels + sm80_kernels
+
+
+GENERATOR_FUNC_TABLE = {75: generate_sm75_tensor_op_1688, 80: generate_sm80_tensor_op_16816}
+
+
+# (Epilogue functor name, no_beta_scaling)
+EPILOGUE_MAP = {
+    "cutlass.dense": (EpilogueFunctor.LinearCombination, False),
+    "cutlass.dense_bias": (EpilogueFunctor.LinearCombinationBias, True),
+    "cutlass.dense_bias_relu": (EpilogueFunctor.LinearCombinationRelu, True),
+    "cutlass.dense_bias_gelu_fp16": (EpilogueFunctor.LinearCombinationGelu, False),
+    "cutlass.dense_bias_gelu_fp32": (EpilogueFunctor.LinearCombinationGelu, False),
+    "cutlass.matmul": (EpilogueFunctor.LinearCombination, False),
+    "cutlass.matmul_bias": (EpilogueFunctor.LinearCombinationBias, True),
+    "cutlass.matmul_bias_relu": (EpilogueFunctor.LinearCombinationRelu, True),
+    "cutlass.matmul_bias_gelu": (EpilogueFunctor.LinearCombinationGelu, False),
+    "cutlass.matmul_transposed": (EpilogueFunctor.LinearCombination, False),
+    "cutlass.matmul_transposed_bias": (EpilogueFunctor.LinearCombinationBias, True),
+    "cutlass.matmul_transposed_bias_relu": (EpilogueFunctor.LinearCombinationRelu, True),
+    "cutlass.matmul_transposed_bias_gelu": (EpilogueFunctor.LinearCombinationGelu, False),
+    "cutlass.batch_matmul": (EpilogueFunctor.LinearCombination, False),
+    "cutlass.conv2d_bias_hardswish": (EpilogueFunctor.LinearCombinationHardSwish, False),
+    "cutlass.conv2d_bias_silu": (EpilogueFunctor.LinearCombinationSilu, False),
+    "cutlass.conv2d_bias_sigmoid": (EpilogueFunctor.LinearCombinationSigmoid, False),
+    "cutlass.conv2d_bias_relu": (EpilogueFunctor.LinearCombinationRelu, True),
+    "cutlass.conv2d_bias": (EpilogueFunctor.LinearCombinationBias, True),
+    "cutlass.conv2d": (EpilogueFunctor.LinearCombination, False),
+    "cutlass.conv2d_transpose": (EpilogueFunctor.LinearCombination, False),
+    "cutlass.conv2d_backward_weight": (EpilogueFunctor.LinearCombination, False),
+}
+
+
+class ProfilerEngine:
+    """Compile and run a given profiler executable."""
+
+    def __init__(self, cuda_arch, cutlass_path, binary_prefix):
+        self.cuda_arch = cuda_arch
+        self.binary_prefix = binary_prefix
+        self.cutlass = cutlass_path
+        self.cflags = f"-I{cutlass_path}/include -I{cutlass_path}/tools/util/include -O3 -std=c++17"
+        self.cflags += " -DCUTLASS_ENABLE_TENSOR_CORE_MMA=1"
+        self.cflags += (
+            f" -gencode=arch=compute_{cuda_arch},code=[sm_{cuda_arch},compute_{cuda_arch}]"
+        )
+        self.cflags += " -Xcompiler=-Wconversion -Xcompiler=-fno-strict-aliasing"
+        self.cmd = "nvcc {cflags} {src} -o {output}"
+
+    def _compile(self, op):
+        os.makedirs(self.binary_prefix, exist_ok=True)
+        opath = os.path.join(self.binary_prefix, op["name"])
+        if os.path.exists(opath):
+            return
+        fi = tempfile.NamedTemporaryFile("w", delete=False, prefix=self.binary_prefix, suffix=".cu")
+        fi.write(op["src"])
+        fi.close()
+        cmd = self.cmd.format(cflags=self.cflags, src=fi.name, output=opath)
+        logger.info("invoking compilation %s", cmd)
+        os.system(cmd)
+        os.unlink(fi.name)
+
+    def compile_all(self, ops, use_multiprocessing=False):
+        """Compile all profiler executables."""
+        if use_multiprocessing:
+            pool = multiprocessing.Pool(multiprocessing.cpu_count())
+            pool.map(self._compile, ops)
+        else:
+            for op in ops:
+                self._compile(op)
+
+    def evaluate(self, op, args):
+        """Run the profiler executable corresponding to op_name with args."""
+        op_name = op["name"]
+        opath = os.path.join(self.binary_prefix, op_name)
+        if not os.path.exists(opath):
+            self._compile(op)
+        if not os.path.exists(opath):
+            # Bail out if compilation fails for a whatever reason (e.g. static assert failure)
+            return float("inf")
+        cmd = [opath]
+        for arg in args:
+            cmd.append(str(arg))
+        try:
+            logger.info("invoking evaluation %s", cmd)
+            sp = subprocess.run(cmd, capture_output=True, check=True)
+            rt = float(sp.stdout)
+            if rt == 0.0:
+                # This seems to happen with split-k using invalid split-k-slices
+                rt = float("inf")
+            logger.info("%s, %f", op_name, rt)
+        except subprocess.CalledProcessError:
+            rt = float("inf")
+        return rt
+
+
+class CodegenResult(Object):
+    """The holder for the generated code and required headers."""
+
+    def __init__(self, code, headers):
+        self.__init_handle_by_constructor__(ffi.CodegenResult, code, headers)
+
+
+def _get_optional_int_annotation(annotations, key, default=None):
+    value = annotations.get(key, default)
+    if value is None:
+        return default
+    return int(value)
+
+
+@tvm_ffi.register_global_func("contrib.cutlass.instantiate_template")
+def instantiate_template(func_name, annotations, func_args):
+    """Return CUTLASS host code based on a template and the provided annotations.
+
+    Parameters
+    ----------
+    func_name: str
+        A string to identify the type of the kernel (dense/matmul, batched_matmul, or conv2d).
+
+    annotations: tvm_ffi.Map
+        Key and value pairs annotated during kernel selection.
+
+    func_args: list
+        Names of the function arguments.
+
+    Returns
+    -------
+    codegen_result : CodegenResult
+        Generated CUTLASS host code and required header-file names.
+    """
+    attrs = {}
+
+    for k in ["lda", "ldb", "ldc", "cutlass_op_def", "cutlass_op_name", "op_type"]:
+        if k in annotations:
+            attrs[k] = annotations[k]
+
+    headers = ["tvm/ffi/function.h", "tvm/ffi/extra/c_env_api.h"]
+
+    if "relu" in func_name:
+        headers.append("cutlass/epilogue/thread/linear_combination_bias_relu.h")
+    elif "gelu" in func_name:
+        headers.append("cutlass/epilogue/thread/linear_combination_gelu.h")
+    elif "sigmoid" in func_name:
+        headers.append("cutlass/epilogue/thread/linear_combination_sigmoid.h")
+    elif "silu" in func_name:
+        headers.append("cutlass/epilogue/thread/linear_combination_silu.h")
+    elif "hardswish" in func_name:
+        headers.append("cutlass/epilogue/thread/linear_combination_hardswish.h")
+    else:
+        headers.append("cutlass/epilogue/thread/linear_combination.h")
+
+    if "residual" in func_name:
+        headers.append("cutlass/epilogue/thread/linear_combination_residual_block.h")
+
+    def get_dim(shape_annot, var_name, axis_idx, batched_offset=0):
+        if isinstance(shape_annot, IntImm):
+            return str(int(shape_annot))
+        return f"{var_name}->shape[{batched_offset + axis_idx}]"
+
+    def get_batch_stride(stride_annot, arg0_idx, arg1_idx, arg0_axis_idx, arg1_axis_idx):
+        if isinstance(stride_annot, IntImm):
+            return str(int(stride_annot))
+        dim1 = func_args[arg0_idx] + f"->shape[{arg0_axis_idx}]"
+        dim2 = func_args[arg1_idx] + f"->shape[{arg1_axis_idx}]"
+        return dim1 + " * " + dim2
+
+    def get_flattened_batch_dim(arg_name, batch_rank):
+        return " * ".join([f"{arg_name}->shape[{i}]" for i in range(batch_rank)])
+
+    if "decode_matmul" in func_name:
+        headers.append("cutlass_kernels/fpA_intB_gemm.h")
+        lhs_arg_idx = _get_optional_int_annotation(annotations, "lhs_arg_idx", 0)
+        rhs_arg_idx = _get_optional_int_annotation(annotations, "rhs_arg_idx", 1)
+        scales_arg_idx = _get_optional_int_annotation(annotations, "scales_arg_idx", 2)
+        bias_arg_idx = _get_optional_int_annotation(annotations, "bias_arg_idx", None)
+        residual_arg_idx = _get_optional_int_annotation(annotations, "residual_arg_idx", None)
+
+        attrs["A_arg"] = func_args[lhs_arg_idx]
+        attrs["B_arg"] = func_args[rhs_arg_idx]
+        attrs["scales_arg"] = func_args[scales_arg_idx]
+        attrs["activation"] = annotations.get("activation", "identity")
+        attrs["bias_stride"] = annotations["bias_stride"]
+        attrs["M"] = annotations["M"]
+        attrs["group_size"] = annotations["group_size"]
+
+        if not isinstance(attrs["M"], tvm.tirx.IntImm):
+            attrs["M"] = get_flattened_batch_dim(
+                func_args[lhs_arg_idx], int(annotations["batch_rank"])
+            )
+
+        if bias_arg_idx is not None:
+            attrs["bias_arg"] = func_args[bias_arg_idx]
+
+        if residual_arg_idx is not None:
+            attrs["residual_arg"] = func_args[residual_arg_idx]
+            attrs["binary_op"] = annotations["binary_op"]
+            attrs["unary_op"] = annotations["unary_op"]
+
+        if annotations["weight_nbit"] == 4:
+            attrs["weight_dtype"] = "cutlass::uint4b_t"
+            attrs["float_per_int"] = 2
+        else:
+            assert annotations["weight_nbit"] == 8
+            attrs["weight_dtype"] = "uint8_t"
+            attrs["float_per_int"] = 1
+
+        code = emit_fp16A_intB_matmul(attrs)
+        return CodegenResult(code, headers)
+
+    elif "dense" in func_name or "matmul" in func_name:
+        batched = "batch" in annotations
+        # dense is equal to transposed_matmul
+        transposed = "transposed" in func_name or "dense" in func_name
+        lhs_arg_idx = _get_optional_int_annotation(annotations, "lhs_arg_idx", 0)
+        rhs_arg_idx = _get_optional_int_annotation(annotations, "rhs_arg_idx", 1)
+        if "bias" in func_name:
+            bias_arg_idx = _get_optional_int_annotation(annotations, "bias_arg_idx", 2)
+        else:
+            bias_arg_idx = _get_optional_int_annotation(annotations, "bias_arg_idx", None)
+        residual_arg_idx = _get_optional_int_annotation(annotations, "residual_arg_idx", None)
+
+        lhs_arg = func_args[lhs_arg_idx]
+        rhs_arg = func_args[rhs_arg_idx]
+        lhs_shape = annotations[f"arg{lhs_arg_idx}_shape"]
+        rhs_shape = annotations[f"arg{rhs_arg_idx}_shape"]
+        lhs_batched_offset = len(lhs_shape) - 2
+        rhs_batched_offset = len(rhs_shape) - 2
+
+        attrs["lhs_arg"] = lhs_arg
+        attrs["rhs_arg"] = rhs_arg
+
+        if bias_arg_idx is not None:
+            attrs["bias_arg"] = func_args[bias_arg_idx]
+        if residual_arg_idx is not None:
+            attrs["residual_arg"] = func_args[residual_arg_idx]
+
+        attrs["ElementInputA"] = DataTypeTag[dtype_map[annotations[f"arg{lhs_arg_idx}_dtype"]]]
+        attrs["ElementInputB"] = DataTypeTag[dtype_map[annotations[f"arg{rhs_arg_idx}_dtype"]]]
+        attrs["ElementOutput"] = DataTypeTag[dtype_map[annotations["ret_dtype"]]]
+
+        attrs["K"] = lhs_shape[lhs_batched_offset + 1]
+        attrs["M"] = get_dim(lhs_shape[lhs_batched_offset], lhs_arg, 0, lhs_batched_offset)
+
+        if transposed:
+            attrs["N"] = get_dim(rhs_shape[rhs_batched_offset], rhs_arg, 0, rhs_batched_offset)
+        else:
+            attrs["N"] = get_dim(rhs_shape[rhs_batched_offset + 1], rhs_arg, 1, rhs_batched_offset)
+
+        if batched:
+            headers.append("cutlass/gemm/device/gemm_batched.h")
+
+            def get_batch_on_arg(arg_name, arg_shape):
+                return " * ".join(f"{arg_name}->shape[{i}]" for i in range(len(arg_shape) - 2))
+
+            if isinstance(annotations["batch"], IntImm):
+                attrs["batch"] = str(int(annotations["batch"]))
+            elif annotations["batch_stride_A"] == 0:
+                # 2D x ND
+                attrs["batch"] = get_batch_on_arg(rhs_arg, rhs_shape)
+            else:
+                # ND x 2D or ND x ND
+                attrs["batch"] = get_batch_on_arg(lhs_arg, lhs_shape)
+
+            attrs["batch_stride_A"] = get_batch_stride(
+                annotations["batch_stride_A"],
+                lhs_arg_idx,
+                lhs_arg_idx,
+                lhs_batched_offset,
+                lhs_batched_offset + 1,
+            )
+            attrs["batch_stride_B"] = get_batch_stride(
+                annotations["batch_stride_B"],
+                rhs_arg_idx,
+                rhs_arg_idx,
+                rhs_batched_offset,
+                rhs_batched_offset + 1,
+            )
+
+            if transposed:
+                attrs["batch_stride_C"] = get_batch_stride(
+                    annotations["batch_stride_C"],
+                    lhs_arg_idx,
+                    rhs_arg_idx,
+                    lhs_batched_offset,
+                    rhs_batched_offset,
+                )
+            else:
+                attrs["batch_stride_C"] = get_batch_stride(
+                    annotations["batch_stride_C"],
+                    lhs_arg_idx,
+                    rhs_arg_idx,
+                    lhs_batched_offset,
+                    rhs_batched_offset + 1,
+                )
+        else:
+            headers.append("cutlass/gemm/device/gemm.h")
+
+        if "residual" in func_name:
+            headers.append("cutlass/gemm/device/gemm_universal_with_broadcast.h")
+
+        code = instantiate_gemm_template(attrs)
+        return CodegenResult(code, headers)
+
+    elif "conv2d" in func_name:
+        data_arg_idx = _get_optional_int_annotation(annotations, "data_arg_idx", 0)
+        weight_arg_idx = _get_optional_int_annotation(annotations, "weight_arg_idx", 1)
+        bias_arg_idx = _get_optional_int_annotation(annotations, "bias_arg_idx", None)
+        residual_arg_idx = _get_optional_int_annotation(annotations, "residual_arg_idx", None)
+
+        attrs["data_arg"] = func_args[data_arg_idx]
+        attrs["weight_arg"] = func_args[weight_arg_idx]
+
+        if bias_arg_idx is not None:
+            attrs["bias_arg"] = func_args[bias_arg_idx]
+        if residual_arg_idx is not None:
+            attrs["residual_arg"] = func_args[residual_arg_idx]
+
+        activation_shape = annotations[f"arg{data_arg_idx}_shape"]
+        weight_shape = annotations[f"arg{weight_arg_idx}_shape"]
+        output_shape = annotations["ret_shape"]
+
+        if "conv2d_transpose" in func_name:
+            headers.append("cutlass/conv/kernel/default_conv2d_dgrad.h")
+            activation_shape = output_shape
+            output_shape = annotations["arg0_shape"]
+        elif "backward" in func_name:
+            headers.append("cutlass/conv/kernel/default_conv2d_wgrad.h")
+            activation_shape = annotations["arg1_shape"]
+            weight_shape = output_shape
+            output_shape = annotations["arg0_shape"]
+        elif "residual" in func_name:
+            headers.append("cutlass/conv/kernel/default_conv2d_fprop_with_broadcast.h")
+        else:
+            headers.append("cutlass/conv/kernel/default_conv2d_fprop.h")
+
+        headers.append("cutlass/conv/device/implicit_gemm_convolution.h")
+
+        op_name = attrs["cutlass_op_name"]
+
+        if "splitk" in op_name:
+            headers += [
+                "cutlass/reduction/device/reduce_split_k.h",
+                "cutlass/reduction/thread/reduction_operators.h",
+            ]
+
+        data_arg = attrs["data_arg"]
+        attrs["N"] = get_dim(activation_shape[0], data_arg, 0)
+        attrs["H"] = get_dim(activation_shape[1], data_arg, 1)
+        attrs["W"] = get_dim(activation_shape[2], data_arg, 2)
+        attrs["C"] = activation_shape[3]
+        attrs["P"] = get_dim(output_shape[1], "out0", 1)
+        attrs["Q"] = get_dim(output_shape[2], "out0", 2)
+        attrs["K"] = output_shape[3]
+        attrs["R"] = weight_shape[1]
+        attrs["S"] = weight_shape[2]
+        attrs["pad_h"] = annotations["padding"][0]
+        attrs["pad_w"] = annotations["padding"][1]
+        attrs["stride_h"] = annotations["strides"][0]
+        attrs["stride_w"] = annotations["strides"][1]
+        attrs["dilation_h"] = annotations["dilation"][0]
+        attrs["dilation_w"] = annotations["dilation"][1]
+
+        if "splitk" in op_name:
+            attrs["split_k_mode"] = "kParallel"
+            attrs["split_k_slices"] = str(re.search(r"splitk(\d+)", op_name).group(1))
+        else:
+            attrs["split_k_mode"] = "kSerial"
+            attrs["split_k_slices"] = 1
+
+        if "residual_shape" in annotations:
+            attrs["residual_shape"] = annotations["residual_shape"]
+
+        code = instantiate_conv2d_template(attrs)
+        return CodegenResult(code, headers)
+
+    elif "attention" in func_name:
+        is_var_len = "var_len" in func_name
+        data_type = dtype_map[annotations["arg0_dtype"]]
+
+        attrs["qkv_layout"] = annotations["qkv_layout"]
+        if attrs["qkv_layout"] == "default":
+            attrs["query"] = func_args[0]
+            attrs["key"] = func_args[1]
+            attrs["value"] = func_args[2]
+            attrs["num_queries"] = s = get_dim(annotations["num_queries"], func_args[0], 1)
+            attrs["num_keys"] = get_dim(annotations["num_keys"], func_args[1], 1)
+            if len(func_args) > 4 and not is_var_len:  # +1 for workspace, the last arg
+                attrs["bias"] = func_args[3]
+        elif attrs["qkv_layout"] == "qkv_stacked":
+            attrs["qkv"] = func_args[0]
+            attrs["num_queries"] = s = annotations["num_queries"]
+            attrs["num_keys"] = annotations["num_keys"]
+            if len(func_args) > 2 and not is_var_len:  # +1 for workspace, the last arg
+                attrs["bias"] = func_args[1]
+        else:
+            raise NotImplementedError()
+
+        attrs["data_type"] = DataTypeTag[data_type]
+        attrs["num_batches"] = b = annotations["num_batches"]
+        attrs["head_dim"] = h = annotations["head_dim"]
+        attrs["head_dim_value"] = h_v = annotations["head_dim_value"]
+        attrs["kMaxK"] = max(int(attrs["head_dim"]), int(attrs["head_dim_value"]))
+        attrs["scale"] = (
+            float(1 / math.sqrt(h.value)) if annotations["scale"] is None else annotations["scale"]
+        )
+
+        if is_var_len:
+            attrs["seqstart_q"] = func_args[int(annotations["seqstart_q_idx"])]
+            attrs["seqstart_k"] = func_args[int(annotations["seqstart_k_idx"])]
+            attrs["max_seqlen_q"] = func_args[int(annotations["max_seqlen_q_idx"])]
+            attrs["max_seqlen_k"] = func_args[int(annotations["max_seqlen_k_idx"])]
+
+        is_mqa = annotations["num_q_heads"] != annotations["num_kv_heads"]
+
+        use_flash = (
+            annotations["ret_dtype"] == "float16"
+            and "bias" not in attrs
+            and int(attrs["head_dim"]) <= 256
+            and int(attrs["head_dim"]) % 8 == 0
+            and int(attrs["head_dim"]) == int(attrs["head_dim_value"])
+            # For the causal case (custom mask = "BottomRight"), only use flash for multi-query
+            # attention workloads. Otherwise, CUTLASS fMHA seems faster for causal attention
+            # with a single query.
+            # In addition, sliding-window attention is only supported by flash.
+            and (
+                int(annotations["custom_mask_type"]) == 0
+                or (int(annotations["custom_mask_type"]) == 2 and is_mqa)
+                or (int(annotations["custom_mask_type"]) == 2 and "window_size" in annotations)
+            )
+            # Flash v2 is currently not supported for sm < 80
+            and int(annotations["arch"]) >= 80
+        )
+
+        # See https://github.com/Dao-AILab/flash-attention/blob/
+        # 92dd5703ecdb99aa4a4aee9817f28557907403a2/csrc/flash_attn/flash_api.cpp#L111-L116
+        if "window_size" in annotations:
+            assert use_flash, "Sliding-window attention is supported only by Flash Attention."
+            assert int(annotations["custom_mask_type"]) == 2, (
+                "Sliding-window attention is only supported for causal with bottom right mask."
+            )
+            attrs["window_size_left"] = int(annotations["window_size"]) - 1
+            attrs["window_size_right"] = 0
+            attrs["is_causal"] = False
+        else:
+            if int(annotations["custom_mask_type"]) == 2:
+                attrs["window_size_left"] = attrs["num_keys"]
+                attrs["window_size_right"] = 0
+                attrs["is_causal"] = True
+            else:
+                attrs["window_size_left"] = -1
+                attrs["window_size_right"] = -1
+                attrs["is_causal"] = False
+
+        if use_flash:
+            headers.append("flash.h")
+            attrs["num_q_heads"] = annotations["num_q_heads"]
+            attrs["num_kv_heads"] = annotations["num_kv_heads"]
+
+            if is_var_len:
+                code = instantiate_flash_attention_var_len_template(attrs)
+            else:
+                code = instantiate_flash_attention_template(attrs)
+        else:
+            headers.append("kernel_forward.h")
+
+            assert not is_mqa, (
+                "The number of query and KV heads need to be the same for CUTLASS fMHA."
+            )
+
+            attrs["num_heads"] = n = annotations["num_q_heads"]
+
+            data_type_size = DataTypeSize[data_type]
+            if (data_type_size * h // 8) % 16 == 0 and (data_type_size * h_v // 8) % 16 == 0:
+                attrs["kIsAligned"] = True
+            elif (h % 4 == 0) and (h_v % 4 == 0):
+                attrs["kIsAligned"] = False
+            else:
+                raise NotImplementedError()
+            if h_v > 64:
+                attrs["kQueriesPerBlock"] = 32
+                attrs["kKeysPerBlock"] = 128
+                attrs["kSingleValueIteration"] = h_v <= 128
+            else:
+                attrs["kQueriesPerBlock"] = 64
+                attrs["kKeysPerBlock"] = 64
+                attrs["kSingleValueIteration"] = True
+
+            assert attrs["scale"] > 0 or attrs["scale"] < 0, (
+                "Cutlass may generate nan occasionally when scale == 0.0"
+            )
+            attrs["arch"] = "cutlass::arch::Sm{}".format(annotations["arch"])
+            attrs["kSupportsDropout"] = False
+
+            attrs["output_size"] = f"{b} * {s} * {n} * {h_v}"
+
+            attrs["custom_mask_type"] = annotations["custom_mask_type"]
+
+            for arg in func_args:
+                if "workspace" in arg:
+                    attrs["workspace"] = arg
+            if "bias" in attrs:
+                attrs["kSupportsBias"] = True
+                if len(annotations["bias_shape"]) == 4:
+                    strides = "p.num_keys"
+                    if annotations["bias_shape"][2] == 1:
+                        attrs["bias_strideM"] = 0
+                    else:
+                        attrs["bias_strideM"] = strides
+                        strides = f"p.num_queries * {strides}"
+                    if annotations["bias_shape"][1] == 1:
+                        attrs["bias_strideH"] = 0
+                    else:
+                        attrs["bias_strideH"] = strides
+                        strides = f"p.num_heads * {strides}"
+                    if annotations["bias_shape"][0] == 1:
+                        attrs["bias_strideB"] = 0
+                    else:
+                        attrs["bias_strideB"] = strides
+                else:
+                    raise NotImplementedError()
+            else:
+                # To support negative scale in current Cutlass implementation,
+                # kSupportsBias should be set true, or there are nan's as result.
+                attrs["kSupportsBias"] = attrs["scale"] < 0
+
+            code = instantiate_attention_template(attrs)
+
+        return CodegenResult(code, headers)
+    elif "layer_norm" in func_name:
+        headers.append("cutlass/util/device_layernorm.h")
+        headers.append("cutlass/layout/matrix.h")
+        attrs = {"input": func_args[0], "gamma": func_args[1], "beta": func_args[2]}
+        attrs.update(dict(annotations))
+
+        if not isinstance(attrs["M"], tvm.tirx.IntImm):
+            attrs["M"] = get_flattened_batch_dim(func_args[0], int(attrs["batch_rank"]))
+
+        code = instantiate_layer_norm_template(attrs)
+        return CodegenResult(code, headers)
+    elif "rms_norm" in func_name:
+        headers.append("cutlass/util/device_rmsnorm.h")
+        headers.append("cutlass/layout/matrix.h")
+        attrs = {"input": func_args[0], "weight": func_args[1]}
+        attrs.update(dict(annotations))
+
+        if not isinstance(attrs["M"], tvm.tirx.IntImm):
+            attrs["M"] = get_flattened_batch_dim(func_args[0], int(attrs["batch_rank"]))
+
+        code = instantiate_rms_norm_template(attrs)
+        return CodegenResult(code, headers)
+
+    raise ValueError(f"Do not have a template for {func_name}")

python/tvm/contrib/cutlass/layer_norm_operation.py

@@ -0,0 +1,48 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name
+# ruff: noqa: E501
+"""Generator for CUTLASS layer norm kernels."""
+
+from .library import substitute_template
+
+
+def instantiate_layer_norm_template(attrs):
+    """
+    Return CUTLASS host code for layer norm based on
+    a template and the provided attribute map.
+    """
+    template = """
+    using data_type = ${data_type};
+    using namespace cutlass::layout;
+
+    int M = ${M};
+    int N = ${N};
+    cutlass::MatrixCoord size(M, N);
+    auto layout_2D = RowMajor::packed(size);
+    auto layout_channels = RowMajor::packed({1, N});
+
+    cutlass::TensorRef<data_type, RowMajor> _input((data_type*)${input}->data, layout_2D);
+    cutlass::TensorRef<data_type, RowMajor> _gamma((data_type*)${gamma}->data, layout_channels);
+    cutlass::TensorRef<data_type, RowMajor> _beta((data_type*)${beta}->data, layout_channels);
+    cutlass::TensorRef<data_type, RowMajor> _output((data_type*)out0->data, layout_2D);
+
+    cudaStream_t stream = static_cast<cudaStream_t>(TVMFFIEnvGetStream(kDLCUDA, ${input}->device.device_id));
+
+    cutlass::layernorm(size, _output, _input, _gamma, _beta, stream);
+    """
+    return substitute_template(template, attrs)

python/tvm/contrib/cutlass/library.py

@@ -0,0 +1,301 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name,line-too-long
+# ruff: noqa: E501
+"""Various type definitions to help instantiate CUTLASS kernels."""
+
+import enum
+import re
+from enum import auto as enum_auto
+
+from tvm.tirx.expr import FloatImm, IntImm
+
+
+class GeneratorTarget(enum.Enum):
+    Library = enum_auto()
+
+
+class DataType(enum.Enum):
+    f16 = enum_auto()
+    f32 = enum_auto()
+    s8 = enum_auto()
+    u8 = enum_auto()
+    s32 = enum_auto()
+
+
+ShortDataTypeNames = {DataType.f16: "h", DataType.f32: "s", DataType.s32: "i"}
+
+
+DataTypeNames = {
+    DataType.f16: "f16",
+    DataType.f32: "f32",
+    DataType.s8: "s8",
+    DataType.u8: "u8",
+    DataType.s32: "s32",
+}
+
+DataTypeTag = {
+    DataType.f16: "cutlass::half_t",
+    DataType.f32: "float",
+    DataType.s8: "int8_t",
+    DataType.s32: "int32_t",
+    DataType.u8: "uint8_t",
+}
+
+DataTypeSize = {
+    DataType.f16: 16,
+    DataType.f32: 32,
+    DataType.u8: 8,
+    DataType.s8: 8,
+    DataType.s32: 32,
+}
+
+
+class MathOperation(enum.Enum):
+    multiply_add = enum_auto()
+    multiply_add_saturate = enum_auto()
+    multiply_add_fast_f32 = enum_auto()
+
+
+MathOperationTag = {
+    MathOperation.multiply_add: "cutlass::arch::OpMultiplyAdd",
+    MathOperation.multiply_add_saturate: "cutlass::arch::OpMultiplyAddSaturate",
+    MathOperation.multiply_add_fast_f32: "cutlass::arch::OpMultiplyAddFastF32",
+}
+
+
+class LayoutType(enum.Enum):
+    ColumnMajor = enum_auto()
+    RowMajor = enum_auto()
+    TensorNHWC = enum_auto()
+
+
+LayoutTag = {
+    LayoutType.ColumnMajor: "cutlass::layout::ColumnMajor",
+    LayoutType.RowMajor: "cutlass::layout::RowMajor",
+    LayoutType.TensorNHWC: "cutlass::layout::TensorNHWC",
+}
+
+
+TransposedLayout = {
+    LayoutType.ColumnMajor: LayoutType.RowMajor,
+    LayoutType.RowMajor: LayoutType.ColumnMajor,
+    LayoutType.TensorNHWC: LayoutType.TensorNHWC,
+}
+
+
+ShortLayoutTypeNames = {
+    LayoutType.ColumnMajor: "n",
+    LayoutType.RowMajor: "t",
+    LayoutType.TensorNHWC: "nhwc",
+}
+
+
+class OpcodeClass(enum.Enum):
+    Simt = enum_auto()
+    TensorOp = enum_auto()
+    WmmaTensorOp = enum_auto()
+
+
+OpcodeClassNames = {
+    OpcodeClass.Simt: "simt",
+    OpcodeClass.TensorOp: "tensorop",
+    OpcodeClass.WmmaTensorOp: "wmma_tensorop",
+}
+
+OpcodeClassTag = {
+    OpcodeClass.Simt: "cutlass::arch::OpClassSimt",
+    OpcodeClass.TensorOp: "cutlass::arch::OpClassTensorOp",
+    OpcodeClass.WmmaTensorOp: "cutlass::arch::OpClassWmmaTensorOp",
+}
+
+
+class OperationKind(enum.Enum):
+    Gemm = enum_auto()
+    Conv2d = enum_auto()
+
+
+OperationKindNames = {OperationKind.Gemm: "gemm", OperationKind.Conv2d: "conv2d"}
+
+
+class Target(enum.Enum):
+    library = enum_auto()
+
+
+def substitute_template(template, values):
+    """Instantiate a kernel template using `values`."""
+    text = template
+    changed = True
+    while changed:
+        changed = False
+        for key, value in values.items():
+            if isinstance(value, int | IntImm):
+                value = str(int(value))
+            if isinstance(value, float | FloatImm):
+                value = str(float(value))
+            elif isinstance(value, bool):
+                value = str(value).lower()
+            regex = f"\\$\\{{{key}\\}}"
+            newtext = re.sub(regex, value, text)
+            if newtext != text:
+                changed = True
+            text = newtext
+    return text
+
+
+class GemmKind(enum.Enum):
+    Gemm = enum_auto()
+
+
+GemmKindNames = {GemmKind.Gemm: "gemm"}
+
+
+class EpilogueFunctor(enum.Enum):
+    LinearCombination = enum_auto()
+    LinearCombinationRelu = enum_auto()
+    LinearCombinationBias = enum_auto()
+    LinearCombinationGelu = enum_auto()
+    LinearCombinationSigmoid = enum_auto()
+    LinearCombinationSilu = enum_auto()
+    LinearCombinationHardSwish = enum_auto()
+    LinearCombinationResidualBlock = enum_auto()
+
+
+EpilogueFunctorTag = {
+    EpilogueFunctor.LinearCombination: "cutlass::epilogue::thread::LinearCombination",
+    EpilogueFunctor.LinearCombinationRelu: "cutlass::epilogue::thread::LinearCombinationRelu",
+    EpilogueFunctor.LinearCombinationBias: "cutlass::epilogue::thread::LinearCombination",
+    EpilogueFunctor.LinearCombinationGelu: "cutlass::epilogue::thread::LinearCombinationGELU",
+    EpilogueFunctor.LinearCombinationSigmoid: "cutlass::epilogue::thread::LinearCombinationSigmoid",
+    EpilogueFunctor.LinearCombinationSilu: "cutlass::epilogue::thread::LinearCombinationSilu",
+    EpilogueFunctor.LinearCombinationHardSwish: "cutlass::epilogue::thread::LinearCombinationHardSwish",
+    EpilogueFunctor.LinearCombinationResidualBlock: "cutlass::epilogue::thread::LinearCombinationResidualBlock",
+}
+
+
+class SwizzlingFunctor(enum.Enum):
+    Identity1 = enum_auto()
+    Identity2 = enum_auto()
+    Identity4 = enum_auto()
+    Identity8 = enum_auto()
+    Batched = enum_auto()
+    StridedDgradIdentity1 = enum_auto()
+    StridedDgradIdentity4 = enum_auto()
+
+
+SwizzlingFunctorTag = {
+    SwizzlingFunctor.Identity1: "cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<1>",
+    SwizzlingFunctor.Identity2: "cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<2>",
+    SwizzlingFunctor.Identity4: "cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<4>",
+    SwizzlingFunctor.Identity8: "cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<8>",
+    SwizzlingFunctor.Batched: "cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle",
+    SwizzlingFunctor.StridedDgradIdentity1: "cutlass::conv::threadblock::StridedDgradIdentityThreadblockSwizzle<1>",
+    SwizzlingFunctor.StridedDgradIdentity4: "cutlass::conv::threadblock::StridedDgradIdentityThreadblockSwizzle<4>",
+}
+
+
+class ConvKind(enum.Enum):
+    Fprop = enum_auto()
+    Dgrad = enum_auto()
+    Wgrad = enum_auto()
+
+
+ConvKindTag = {
+    ConvKind.Fprop: "cutlass::conv::Operator::kFprop",
+    ConvKind.Dgrad: "cutlass::conv::Operator::kDgrad",
+    ConvKind.Wgrad: "cutlass::conv::Operator::kWgrad",
+}
+
+
+ConvKindNames = {ConvKind.Fprop: "fprop", ConvKind.Dgrad: "dgrad", ConvKind.Wgrad: "wgrad"}
+
+
+class StrideSupport(enum.Enum):
+    Strided = enum_auto()
+    Unity = enum_auto()
+
+
+StrideSupportTag = {
+    StrideSupport.Strided: "cutlass::conv::StrideSupport::kStrided",
+    StrideSupport.Unity: "cutlass::conv::StrideSupport::kUnity",
+}
+
+
+StrideSupportNames = {StrideSupport.Strided: "", StrideSupport.Unity: "unity_stride"}
+
+
+class IteratorAlgorithm(enum.Enum):
+    Analytic = enum_auto()
+    Optimized = enum_auto()
+
+
+IteratorAlgorithmTag = {
+    IteratorAlgorithm.Analytic: "cutlass::conv::IteratorAlgorithm::kAnalytic",
+    IteratorAlgorithm.Optimized: "cutlass::conv::IteratorAlgorithm::kOptimized",
+}
+
+
+IteratorAlgorithmNames = {
+    IteratorAlgorithm.Analytic: "analytic",
+    IteratorAlgorithm.Optimized: "optimized",
+}
+
+
+class MathInstruction:
+    """Describe characteristics of a math instruction."""
+
+    def __init__(
+        self,
+        instruction_shape,
+        element_a,
+        element_b,
+        element_c,
+        element_accumulator,
+        opcode_class,
+        math_operation=MathOperation.multiply_add,
+    ):
+        self.instruction_shape = instruction_shape
+        self.element_a = element_a
+        self.element_b = element_b
+        self.element_c = element_c
+        self.element_accumulator = element_accumulator
+        self.opcode_class = opcode_class
+        self.math_operation = math_operation
+
+
+class TileDescription:
+    """Describe characteristics of a GEMM tile."""
+
+    def __init__(
+        self, threadblock_shape, stages, warp_count, math_instruction, min_compute, max_compute
+    ):
+        self.threadblock_shape = threadblock_shape
+        self.stages = stages
+        self.warp_count = warp_count
+        self.math_instruction = math_instruction
+        self.minimum_compute_capability = min_compute
+        self.maximum_compute_capability = max_compute
+
+    def procedural_name(self):
+        return f"{self.threadblock_shape[0]}x{self.threadblock_shape[1]}_{self.threadblock_shape[2]}x{self.stages}"
+
+
+class TensorDescription:
+    def __init__(self, element, layout, alignment=1):
+        self.element = element
+        self.layout = layout
+        self.alignment = alignment

python/tvm/contrib/cutlass/rms_norm_operation.py

@@ -0,0 +1,47 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name
+# ruff: noqa: E501
+"""Generator for CUTLASS rms norm kernels."""
+
+from .library import substitute_template
+
+
+def instantiate_rms_norm_template(attrs):
+    """
+    Return CUTLASS host code for rms norm based on
+    a template and the provided attribute map.
+    """
+    template = """
+    using data_type = ${data_type};
+    using namespace cutlass::layout;
+
+    int M = ${M};
+    int N = ${N};
+    cutlass::MatrixCoord size(M, N);
+    auto layout_2D = RowMajor::packed(size);
+    auto layout_channels = RowMajor::packed({1, N});
+
+    cutlass::TensorRef<data_type, RowMajor> _input((data_type*)${input}->data, layout_2D);
+    cutlass::TensorRef<data_type, RowMajor> _weight((data_type*)${weight}->data, layout_channels);
+    cutlass::TensorRef<data_type, RowMajor> _output((data_type*)out0->data, layout_2D);
+
+    cudaStream_t stream = static_cast<cudaStream_t>(TVMFFIEnvGetStream(kDLCUDA, ${input}->device.device_id));
+
+    cutlass::rmsnorm(size, _output, _input, _weight, stream, ${rms_eps});
+    """
+    return substitute_template(template, attrs)

python/tvm/contrib/dlpack.py

@@ -0,0 +1,66 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Wrapping functions to bridge frameworks with DLPack support to TVM"""
+
+import tvm.runtime
+
+
+def convert_func(tvm_func, tensor_type, to_dlpack_func):
+    """Convert a tvm function into one that accepts a tensor from another
+       framework, provided the other framework supports DLPACK
+
+    Parameters
+    ----------
+    tvm_func: Function
+        Built tvm function operating on arrays
+
+    tensor_type: Type
+        Type of the tensors of the target framework
+
+    to_dlpack_func: Function
+        Function to convert the source tensors to DLPACK
+    """
+    assert callable(tvm_func)
+
+    def _wrapper(*args):
+        args = tuple(
+            tvm.runtime.from_dlpack(to_dlpack_func(arg)) if isinstance(arg, tensor_type) else arg
+            for arg in args
+        )
+        return tvm_func(*args)
+
+    return _wrapper
+
+
+def to_pytorch_func(tvm_func):
+    """Convert a tvm function into one that accepts PyTorch tensors
+
+    Parameters
+    ----------
+    tvm_func: Function
+        Built tvm function operating on arrays
+
+    Returns
+    -------
+    wrapped_func: Function
+        Wrapped tvm function that operates on PyTorch tensors
+    """
+    # pylint: disable=import-outside-toplevel
+    import torch
+    import torch.utils.dlpack
+
+    return convert_func(tvm_func, torch.Tensor, torch.utils.dlpack.to_dlpack)

python/tvm/contrib/dnnl.py

@@ -0,0 +1,164 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""External function interface to BLAS libraries."""
+
+import tvm
+from tvm import te
+
+from ..topi.nn.utils import get_pad_tuple
+
+
+def matmul(lhs, rhs, transa=False, transb=False, **kwargs):
+    """Create an extern op that compute matrix mult of A and rhs with CrhsLAS
+    This function serves as an example on how to call external libraries.
+
+    Parameters
+    ----------
+    lhs: Tensor
+        The left matrix operand
+    rhs: Tensor
+        The right matrix operand
+    transa: bool
+        Whether transpose lhs
+    transb: bool
+        Whether transpose rhs
+
+    Returns
+    -------
+    C: Tensor
+        The result tensor.
+    """
+    n = lhs.shape[1] if transa else lhs.shape[0]
+    m = rhs.shape[0] if transb else rhs.shape[1]
+    return te.extern(
+        (n, m),
+        [lhs, rhs],
+        lambda ins, outs: tvm.tirx.call_packed(
+            "tvm.contrib.dnnl.matmul", ins[0], ins[1], outs[0], transa, transb
+        ),
+        name="C",
+        **kwargs,
+    )
+
+
+def dnnl_conv2d(
+    src,
+    weights,
+    stride,
+    padding,
+    dilation,
+    groups,
+    channel_last=False,
+    out_dtype="float32",
+    **kwargs,
+):
+    """Convolution operator in NCHW layout.
+
+    Parameters
+    ----------
+    src : tvm.te.Tensor
+        4-D with shape [batch, in_channel, in_height, in_width]
+
+    weights : tvm.te.Tensor
+        4-D with shape [num_filter, in_channel, filter_height, filter_width]
+
+    stride : int or a list/tuple of two ints
+        Stride size, or [stride_height, stride_width]
+
+    padding : int or a list/tuple of 2 or 4 ints
+        padding size, or
+        [pad_height, pad_width] for 2 ints, or
+        [pad_top, pad_left, pad_bottom, pad_right] for 4 ints
+
+    dilation: int or a list/tuple of two ints
+        dilation size, or [dilation_height, dilation_width]
+
+    groups: str
+        input data layout: NCHW or NHWC
+
+    channel_last: bool
+        chose if input/output data format is in channel_last format(NHWC) or
+        in plain format(NCHW)
+
+    out_dtype: str
+        output datatype: now only support float32
+
+    Returns
+    -------
+    Output : tvm.te.Tensor
+        4-D with shape [batch, out_channel, out_height, out_width]
+    """
+
+    assert isinstance(stride, int) or len(stride) == 2
+    assert isinstance(dilation, int) or len(dilation) == 2
+    if isinstance(stride, int):
+        stride_h = stride_w = stride
+    else:
+        stride_h, stride_w = stride
+
+    if isinstance(dilation, int):
+        dilation_h = dilation_w = dilation
+    else:
+        dilation_h, dilation_w = dilation
+
+    pre_cast = src.dtype == "float32"
+    post_cast = out_dtype == "float32"
+
+    if channel_last:
+        batch, in_height, in_width, _ = src.shape
+        kernel_h, kernel_w, _, num_filter = weights.shape
+    else:
+        batch, _, in_height, in_width = src.shape
+        num_filter, _, kernel_h, kernel_w = weights.shape
+
+    dilated_kernel_h = (kernel_h - 1) * dilation_h + 1
+    dilated_kernel_w = (kernel_w - 1) * dilation_w + 1
+    pad_top, pad_left, pad_down, pad_right = get_pad_tuple(
+        padding, (dilated_kernel_h, dilated_kernel_w)
+    )
+    out_channel = num_filter
+    out_height = (in_height - dilated_kernel_h + pad_top + pad_down) // stride_h + 1
+    out_width = (in_width - dilated_kernel_w + pad_left + pad_right) // stride_w + 1
+
+    if channel_last:
+        out_shape = (batch, out_height, out_width, out_channel)
+    else:
+        out_shape = (batch, out_channel, out_height, out_width)
+
+    return te.extern(
+        out_shape,
+        [src, weights],
+        lambda ins, outs: tvm.tirx.call_packed(
+            "tvm.contrib.dnnl.conv2d",
+            ins[0],
+            ins[1],
+            outs[0],
+            pad_top,
+            pad_down,
+            pad_left,
+            pad_right,
+            stride[0],
+            stride[1],
+            groups,
+            channel_last,
+            pre_cast,
+            post_cast,
+        ),
+        name="C",
+        dtype=out_dtype,
+        **kwargs,
+    )

python/tvm/contrib/download.py

@@ -0,0 +1,175 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-timeout
+"""Helper utility for downloading"""
+
+import logging
+import os
+import shutil
+import tempfile
+import time
+from pathlib import Path
+
+LOG = logging.getLogger("download")
+
+
+def download(url, path, overwrite=False, size_compare=False, retries=3):
+    """Downloads the file from the internet.
+    Set the input options correctly to overwrite or do the size comparison
+
+    Parameters
+    ----------
+    url : str
+        Download url.
+
+    path : str
+        Local file path to save downloaded file.
+
+    overwrite : bool, optional
+        Whether to overwrite existing file, defaults to False.
+
+    size_compare : bool, optional
+        Whether to do size compare to check downloaded file, defaults
+        to False
+
+    retries: int, optional
+        Number of time to retry download, defaults to 3.
+
+    """
+    # pylint: disable=import-outside-toplevel
+    import urllib.request as urllib2
+
+    path = Path(path).resolve()
+    if path.exists() and path.is_file() and not overwrite:
+        if size_compare:
+            import requests
+
+            file_size = path.stat().st_size
+            res_head = requests.head(url)
+            res_get = requests.get(url, stream=True)
+            if "Content-Length" not in res_head.headers:
+                res_get = urllib2.urlopen(url)
+            url_file_size = int(res_get.headers["Content-Length"])
+            if url_file_size != file_size:
+                LOG.warning("Existing file %s has incorrect size, downloading fresh copy", path)
+                download(url, path, overwrite=True, size_compare=False, retries=retries)
+                return
+
+        LOG.info("File %s exists, skipping.", path)
+        return
+
+    LOG.info("Downloading from url %s to %s", url, path)
+
+    # Stateful start time
+    start_time = time.time()
+    dirpath = path.parent
+    dirpath.mkdir(parents=True, exist_ok=True)
+
+    def _download_progress(count, block_size, total_size):
+        # pylint: disable=unused-argument
+        """Show the download progress."""
+        if count == 0:
+            return
+        duration = time.time() - start_time
+        progress_bytes = int(count * block_size)
+        progress_megabytes = progress_bytes / (1024.0 * 1024)
+        speed_kbps = int(progress_bytes / (1024 * duration))
+        percent = min(int(count * block_size * 100 / total_size), 100)
+
+        # Temporarily suppress newlines on the output stream.
+        prev_terminator = logging.StreamHandler.terminator
+        logging.StreamHandler.terminator = ""
+        LOG.debug(
+            "\r...%d%%, %.2f MB, %d KB/s, %d seconds passed",
+            percent,
+            progress_megabytes,
+            speed_kbps,
+            duration,
+        )
+        logging.StreamHandler.terminator = prev_terminator
+
+    with tempfile.TemporaryDirectory() as tempdir:
+        tempdir = Path(tempdir)
+        download_loc = tempdir.joinpath(path.name)
+
+        for i_retry in range(retries):
+            # pylint: disable=broad-except
+            try:
+                urllib2.urlretrieve(url, download_loc, reporthook=_download_progress)
+                LOG.debug("")
+                try:
+                    download_loc.rename(path)
+                except OSError:
+                    # Prefer a move, but if the tempdir and final
+                    # location are in different drives, fall back to a
+                    # copy.
+                    shutil.copy2(download_loc, path)
+                return
+
+            except Exception as err:
+                if i_retry == retries - 1:
+                    raise err
+
+                LOG.warning(
+                    "%s\nDownload attempt %d/%d failed, retrying.", repr(err), i_retry, retries
+                )
+
+
+if "TEST_DATA_ROOT_PATH" in os.environ:
+    TEST_DATA_ROOT_PATH = Path(os.environ.get("TEST_DATA_ROOT_PATH"))
+else:
+    TEST_DATA_ROOT_PATH = Path(Path("~").expanduser(), ".tvm_test_data")
+TEST_DATA_ROOT_PATH.mkdir(parents=True, exist_ok=True)
+
+
+def download_testdata(url, relpath, module=None, overwrite=False):
+    """Downloads the test data from the internet.
+
+    Parameters
+    ----------
+    url : str
+        Download url.
+
+    relpath : str
+        Relative file path.
+
+    module : Union[str, list, tuple], optional
+        Subdirectory paths under test data folder.
+
+    overwrite : bool, defaults to False
+        If True, will download a fresh copy of the file regardless of
+        the cache.  If False, will only download the file if a cached
+        version is missing.
+
+    Returns
+    -------
+    abspath : str
+        Absolute file path of downloaded file
+
+    """
+    global TEST_DATA_ROOT_PATH
+    if module is None:
+        module_path = ""
+    elif isinstance(module, str):
+        module_path = module
+    elif isinstance(module, list | tuple):
+        module_path = Path(*module)
+    else:
+        raise ValueError("Unsupported module: " + module)
+    abspath = Path(TEST_DATA_ROOT_PATH, module_path, relpath)
+    download(url, abspath, overwrite=overwrite, size_compare=False)
+    return str(abspath)

python/tvm/contrib/hexagon/__init__.py

@@ -0,0 +1,19 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Hexagon APIs."""
+
+from .tools import *

python/tvm/contrib/hexagon/_ci_env_check.py

@@ -0,0 +1,55 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Hexagon environment checks for CI usage
+
+These are required by tvm.testing and live here to avoid a circular dependency.
+"""
+
+import os
+
+import tvm
+
+ANDROID_SERIAL_NUMBER = "ANDROID_SERIAL_NUMBER"
+HEXAGON_TOOLCHAIN = "HEXAGON_TOOLCHAIN"
+
+
+def _compile_time_check():
+    """Return True if compile-time support for Hexagon is present, otherwise
+    error string.
+
+    Backs :func:`tvm.testing.env.has_hexagon_toolchain`.
+    """
+    if tvm.runtime.enabled("llvm") and tvm.target.codegen.llvm_version_major() < 7:
+        return "Hexagon requires LLVM 7 or later"
+
+    if "HEXAGON_TOOLCHAIN" not in os.environ:
+        return f"Missing environment variable {HEXAGON_TOOLCHAIN}."
+
+    return True
+
+
+def _run_time_check():
+    """Return True if run-time support for Hexagon is present, otherwise
+    error string.
+
+    Backs :func:`tvm.testing.env.has_hexagon`.
+    """
+    if ANDROID_SERIAL_NUMBER not in os.environ:
+        return f"Missing environment variable {ANDROID_SERIAL_NUMBER}."
+
+    return True

python/tvm/contrib/hexagon/generate_take_op.py

@@ -0,0 +1,98 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-docstring, invalid-name, unnecessary-comprehension, unused-argument
+
+import tvm
+import tvm.testing
+from tvm import relax
+from tvm.contrib.hexagon import hexagon_unary_ops
+
+
+def op_replace(call_node, func) -> bool:
+    if not isinstance(call_node, relax.Call):
+        return False
+    call_tir_op = tvm.ir.Op.get("relax.call_tir")
+    if call_node.op != call_tir_op:
+        return False
+    ops = [
+        "qnn.tanh",
+        "qnn.sqrt",
+        "qnn.rsqrt",
+        "qnn.exp",
+        "qnn.erf",
+        "qnn.sigmoid",
+        "qnn.hardswish",
+        "qnn.log",
+        "qnn.abs",
+    ]
+    if func.attrs["op_attrs"]["op_name"] in ops:
+        return True
+    return False
+
+
+@relax.expr_functor.mutator
+class Tanh2TakeReplace(tvm.relax.PyExprMutator):
+    def __init__(self, mod: tvm.IRModule) -> None:
+        super().__init__(mod)
+        self.mod_ = mod
+
+    def transform(self) -> tvm.IRModule:
+        # Iterate over all the nodes to check for the node replaceable
+        for global_var, func in self.mod_.functions.items():
+            # Skip non-relax functions
+            if not isinstance(func, relax.Function):
+                continue
+            updated_func = self.visit_expr(func)
+            self.builder_.normalize(updated_func)
+            self.builder_.update_func(global_var, updated_func)
+        # At the end of the transformation we return the updated IRModule from the BlockBuilder.
+        return self.builder_.get()
+
+    def visit_call_(self, call_node: relax.Call) -> relax.Call:
+        call_tir_op = tvm.ir.Op.get("relax.call_tir")
+        if call_node.op != call_tir_op:
+            return call_node
+
+        var = call_node.args[0]
+        func = self.mod_[var]
+
+        if call_node.args[1][0].ty.dtype == "uint8":
+            if op_replace(call_node, func):
+                inp, inp_scale, inp_zp, out_scale, out_zp = [x for x in call_node.args[1]]
+                # LUT node creation
+                LUT = hexagon_unary_ops.LUT_generation(
+                    inp_scale, inp_zp, out_scale, out_zp, call_node.args[0].name_hint
+                )
+                # Take operation node creation
+                take_func = hexagon_unary_ops.generate_take_primfunc(inp, call_node.ty)
+                take_func = take_func.without_attr("global_symbol")
+                take_func_gv = self.builder_.add_func(take_func, "take")
+                take_node = relax.call_tir(
+                    take_func_gv,
+                    relax.expr.Tuple(
+                        [call_node.args[1][0], relax.expr.Constant(tvm.runtime.tensor(LUT))]
+                    ),
+                    call_node.ty,
+                )
+                return take_node
+        return call_node
+
+
+@tvm.ir.transform.module_pass(opt_level=2, name="replace_tanh_take")
+class PassReplaceWithTakeOpPrimFuncs:
+    def transform_module(self, mod, ctx):
+        return Tanh2TakeReplace(mod).transform()

python/tvm/contrib/hexagon/hexagon_unary_ops.py

@@ -0,0 +1,99 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-docstring, invalid-name
+import logging
+
+import numpy as np
+from scipy import special
+
+from tvm import te
+
+logger = logging.getLogger(__name__)
+
+######################################################################
+#################### PRIMFUNC FOR LUT and Take Op ####################
+######################################################################
+
+
+def saturate(x: te.Tensor, dtype: str):
+    """Saturate value for the specified data type"""
+    return te.max(te.min_value(dtype), te.min(x, te.max_value(dtype)))
+
+
+def hardswish_func(x):
+    x_2 = np.add(x, 3.0)
+    x_2 = np.clip(x_2, 0.0, 6.0)
+    return x * x_2 / 6.0
+
+
+def LUT_generation(inp_scale, inp_zp, out_scale, out_zp, op_name) -> None:
+    LUT = []
+    for i in range(256):
+        i = np.int32(i)
+        # converting the constants to the numpy value
+        if inp_zp.data.shape == ():
+            i_zp = inp_zp.data.numpy()[()]
+        if inp_scale.data.shape == ():
+            i_scale = inp_scale.data.numpy()[()]
+        if out_zp.data.shape == ():
+            o_zp = out_zp.data.numpy()[()]
+        if out_scale.data.shape == ():
+            o_scale = out_scale.data.numpy()[()]
+        # Dequantization followed by computing the op value
+        dequant = (i - i_zp) * i_scale
+        if "tanh" in op_name:
+            op_val = np.tanh(dequant)
+        elif "rsqrt" in op_name:
+            op_val = 1 / np.sqrt(dequant)
+        elif "sqrt" in op_name:
+            op_val = np.sqrt(dequant)
+        elif "exp" in op_name:
+            op_val = np.exp(dequant)
+        elif "erf" in op_name:
+            op_val = special.erf(dequant)
+        elif "sigmoid" in op_name:
+            op_val = 1 / (1 + np.exp(np.negative(dequant)))
+        elif "hardswish" in op_name:
+            op_val = hardswish_func(dequant)
+        elif "log" in op_name:
+            op_val = np.log(dequant)
+        elif "abs" in op_name:
+            op_val = np.abs(dequant)
+        else:
+            logger.error("Error op is other than unary op")
+
+        # Quantizing the value generated and appending in the Look Up Table
+        quant = np.round((op_val) / o_scale) + o_zp
+        val = np.maximum(0, np.minimum(quant, 255)).astype(np.uint8)
+        LUT.append(val)
+    return LUT
+
+
+def generate_take_primfunc(inp, ty):
+    # Generating the take op
+    N, H, W, C = inp.ty.shape
+    data = te.placeholder((N, H, W, C), dtype=ty.dtype, name="data")
+    LUT_func = te.placeholder((256,), dtype="uint8", name="LUT")
+    take = te.compute(
+        ty.shape,
+        lambda *indices: saturate((LUT_func[data[indices].astype("uint8")]), ty.dtype).astype(
+            ty.dtype
+        ),
+        name="take_op",
+    )
+    mod = te.create_prim_func([data, LUT_func, take])
+    return mod

python/tvm/contrib/hexagon/tools.py

@@ -0,0 +1,561 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name, f-string-without-interpolation, consider-using-from-import
+# ruff: noqa: E501
+"""Tools/compilers/linkers for Hexagon"""
+
+import io
+import itertools
+import os
+import pathlib
+import re
+import subprocess
+import sys
+import tarfile
+
+import numpy
+from tvm_ffi import register_global_func
+
+import tvm
+import tvm.support.cc as cc
+
+# Linking Hexagon shared libraries.
+#
+#   link_shared(name-of-shared-library, list-of-objects, kw-args)
+#
+# To use a custom linker, define a function that returns the path to the
+# linker, and pass it to 'register_linker':
+#
+#   def custom_linker_path():
+#       return '/path/to/hexagon/linker'
+#
+#   register_linker(custom_linker_path)
+#
+# Subsequent calls to 'link_shared' will use the newly registered linker.
+
+HEXAGON_TOOLCHAIN = os.environ.get("HEXAGON_TOOLCHAIN", default="")  # pylint: disable=invalid-name
+HEXAGON_SDK_ROOT = os.environ.get("HEXAGON_SDK_ROOT", default="")  # pylint: disable=invalid-name
+HEXAGON_SDK_DOCKER_IMAGE = os.environ.get("HEXAGON_SDK_DOCKER_IMAGE", default="")  # pylint: disable=invalid-name
+HEXAGON_LINK_MAIN = pathlib.Path(HEXAGON_TOOLCHAIN) / "bin" / "hexagon-link"  # pylint: disable=invalid-name
+HEXAGON_CLANG_PLUS = pathlib.Path(HEXAGON_TOOLCHAIN) / "bin" / "hexagon-clang++"  # pylint: disable=invalid-name
+HEXAGON_SDK_INCLUDE_DIRS = [  # pylint: disable=invalid-name
+    pathlib.Path(HEXAGON_SDK_ROOT) / "incs",
+    pathlib.Path(HEXAGON_SDK_ROOT) / "incs" / "stddef",
+]
+
+HEXAGON_SIMULATOR_NAME = "simulator"
+
+
+def register_linker(f):
+    """Register a function that will return the path to the Hexagon linker."""
+    return register_global_func("tvm.contrib.hexagon.hexagon_link", f, True)
+
+
+@register_global_func("tvm.contrib.hexagon.hexagon_link")
+def hexagon_link() -> str:
+    """Return path to the Hexagon linker."""
+    return str(HEXAGON_LINK_MAIN)
+
+
+def hexagon_clang_plus() -> str:
+    """Return path to the Hexagon clang++."""
+    return str(HEXAGON_CLANG_PLUS)
+
+
+def toolchain_version(toolchain=None) -> list[int]:
+    """Return the version of the Hexagon toolchain.
+
+    Parameters
+    ----------
+    toolchain: str, optional
+        Path to the Hexagon toolchain. If not provided, the environment
+        variable HEXAGON_TOOLCHAIN is used.
+
+    Returns
+    -------
+    version: List[int]
+        List of numerical components of the version number. E.g. for version
+        "8.5.06" it will be [8, 5, 6].
+    """
+
+    if toolchain is None:
+        toolchain = HEXAGON_TOOLCHAIN
+    assert toolchain is not None, "Please specify toolchain, or set HEXAGON_TOOLCHAIN variable"
+    result = subprocess.run(
+        [f"{toolchain}/bin/hexagon-clang", "-v"], capture_output=True, check=True
+    )
+    output = result.stderr.decode()
+    for line in output.splitlines():
+        m = re.match(r".* [Cc]lang version ([0-9\.]+)", line)
+        if m:
+            assert len(m.groups()) == 1
+            return [int(v) for v in m.group(1).split(".")]
+    raise RuntimeError("Cannot establish toolchain version")
+
+
+@register_global_func("tvm.contrib.hexagon.link_shared")
+def link_shared(so_name, objs, extra_args=None):
+    """Link shared library on Hexagon using the registered Hexagon linker.
+
+    Parameters
+    ----------
+    so_name : str
+        Name of the shared library file.
+    objs : list[str, tvm.tirx.StringImm]
+    extra_args : dict (str->str) or Map<String,String>
+        Additional arguments:
+            'hex_arch' - Hexagon architecture, e.g. v68
+            'verbose'  - Print additional information if the key is present
+
+    Returns
+    -------
+    ret_val : int
+        This function returns 0 at the moment.
+    """
+
+    # The list of object files can be passed as built-in Python strings,
+    # or as tvm.tirx.StringImm's.
+    def to_str(s):
+        if isinstance(s, tvm.tirx.StringImm):
+            return s.value
+        assert isinstance(s, str), 'argument "' + str(s) + '" should be a string or StrImm'
+        return s
+
+    objs = [to_str(s) for s in objs]
+
+    if not extra_args:
+        extra_args = {}
+    hex_arch = extra_args.get("hex_arch") or "v68"
+    linker = tvm.get_global_func("tvm.contrib.hexagon.hexagon_link")()
+    if extra_args.get("verbose"):
+        print("tvm.contrib.hexagon.link_shared:")
+        print("  Using linker:", linker)
+        print("  Library name:", so_name)
+        print("  Object files:", objs)
+        print("  Architecture:", hex_arch)
+    if not os.access(linker, os.X_OK):
+        message = 'The linker "' + linker + '" does not exist or is not executable.'
+        if not os.environ.get("HEXAGON_TOOLCHAIN"):
+            message += (
+                " The environment variable HEXAGON_TOOLCHAIN is unset. Please export "
+                + "HEXAGON_TOOLCHAIN in your environment, so that ${HEXAGON_TOOLCHAIN}/bin/"
+                + "hexagon-link exists."
+            )
+        else:
+            message += (
+                " Please verify the value of the HEXAGON_LINKER environment variable "
+                + '(currently set to "'
+                + HEXAGON_TOOLCHAIN
+                + '").'
+            )
+        raise Exception(message)
+
+    libpath = os.path.join(HEXAGON_TOOLCHAIN, "target", "hexagon", "lib", hex_arch, "G0")
+    cc.create_shared(
+        so_name,
+        objs,
+        # pylint: disable=bad-whitespace
+        options=[
+            "-Bdynamic",
+            "-shared",
+            "-export-dynamic",
+            os.path.join(libpath, "pic", "libgcc.so"),
+        ],
+        cc=linker,
+    )
+    return 0
+
+
+def link_shared_macos(so_name, objs, extra_args=None):
+    """Link Hexagon shared library using docker container with proper tooling.
+
+    Parameters
+    ----------
+    so_name : str
+        Name of the shared library file.
+    objs : list[str, tvm.tirx.StringImm]
+    extra_args : dict (str->str) or Map<String,String>
+        Additional arguments:
+            'hex_arch' - Hexagon architecture, e.g. v68
+
+    Returns
+    -------
+    ret_val : int
+        This function returns 0 at the moment.
+    """
+
+    # The list of object files can be passed as built-in Python strings,
+    # or as tvm.tirx.StringImm's.
+    def to_str(s):
+        if isinstance(s, tvm.tirx.StringImm):
+            return s.value
+        assert isinstance(s, str), 'argument "' + str(s) + '" should be a string or StrImm'
+        return s
+
+    objs = [to_str(s) for s in objs]
+
+    if not extra_args:
+        extra_args = {}
+    hex_arch = extra_args.get("hex_arch") or "v68"
+
+    ses = ContainerSession(HEXAGON_SDK_DOCKER_IMAGE)
+
+    hexagon_sdk_tools_path = ses.get_env("HEXAGON_TOOLCHAIN")
+    libpath = os.path.join(hexagon_sdk_tools_path, "target", "hexagon", "lib", hex_arch, "G0")
+    linker = os.path.join(hexagon_sdk_tools_path, "bin", "hexagon-link")
+
+    # Copy input data to docker container
+    docker_objs = [ses.copy_to(obj) for obj in objs]
+    docker_so_name = ses.tmp_dir + "/" + os.path.basename(so_name)
+
+    link_cmd = [linker, "-shared", "-fPIC", "-o", docker_so_name]
+    link_cmd += docker_objs
+    link_cmd += [
+        "-Bdynamic",
+        "-export-dynamic",
+        "-L" + os.path.join(libpath, "pic"),
+        "-lgcc",
+    ]
+    ses.exec(link_cmd)
+
+    # Copy result back to host
+    ses.copy_from(docker_so_name, so_name)
+    return 0
+
+
+if sys.platform == "darwin":
+
+    def __create_shared_mac(so_name, objs, **kwargs):
+        return link_shared_macos(so_name, objs, kwargs)
+
+    create_shared = __create_shared_mac
+    register_global_func("tvm.contrib.hexagon.link_shared", f=link_shared_macos, override=True)
+else:  # Linux and Win32
+    create_shared = cc.create_shared
+    register_global_func("tvm.contrib.hexagon.link_shared", f=link_shared, override=True)
+
+
+def create_aot_shared(so_name: str | pathlib.Path, files, hexagon_arch: str, options=None):
+    """Export Hexagon AOT module."""
+    options = options or []
+    if not os.access(str(HEXAGON_CLANG_PLUS), os.X_OK):
+        raise Exception(
+            'The Clang++ "' + str(HEXAGON_CLANG_PLUS) + '" does not exist or is not executable.'
+        )
+    if not HEXAGON_TOOLCHAIN:
+        raise Exception(
+            " The environment variable HEXAGON_TOOLCHAIN is unset. Please export "
+            + "HEXAGON_TOOLCHAIN in your environment."
+        )
+    if not HEXAGON_SDK_ROOT:
+        raise Exception(
+            " The environment variable HEXAGON_SDK_ROOT is unset. Please export "
+            + "HEXAGON_SDK_ROOT in your environment."
+        )
+
+    # The AOT C codegen uses TVM runtime functions
+    # (e.g. TVMBackendAllocWorkspace) directly. On Hexagon these calls
+    # should be made using functions pointers provided as __TVM*
+    # variables in the provided context.  This workaround allows the
+    # the TVM runtime symbols to be visible to the compiled shared
+    # library.
+    #
+    # This workaround can be removed when AOT codegen can be done with
+    # LLVM codegen.
+    workaround_link_flags = os.environ.get("HEXAGON_SHARED_LINK_FLAGS")
+    if workaround_link_flags:
+        options.extend(workaround_link_flags.split())
+
+    tvm_dir = pathlib.Path(os.path.dirname(os.path.realpath(__file__))) / ".." / ".." / ".." / ".."
+    compute_arch = f"compute{hexagon_arch}"
+    compile_options = [
+        "-O3",
+        f"-I{tvm_dir / 'include'}",
+        f"-I{tvm_dir / '3rdparty' / 'dlpack' / 'include'}",
+        f"-I{pathlib.Path(HEXAGON_SDK_ROOT) / 'rtos' / 'qurt' / compute_arch / 'include' / 'posix'}",
+        f"-I{pathlib.Path(HEXAGON_SDK_ROOT) / 'rtos' / 'qurt' / compute_arch / 'include' / 'qurt'}",
+        "-D_MACH_I32=int",
+    ]
+
+    # For debugging
+    for path in HEXAGON_SDK_INCLUDE_DIRS:
+        compile_options.append(f"-I{path!s}")
+
+    cross_compile = cc.cross_compiler(compile_func=hexagon_clang_plus())
+    cross_compile.output_format = "o"
+    c_files = [str(file) for file in files]
+    cross_compile(str(so_name), c_files, options=compile_options + options)
+
+
+def pack_imports(
+    module: tvm.runtime.Module,
+    is_system_lib: bool,  # pylint: disable=unused-argument
+    c_symbol_prefix: str,
+    workspace_dir: str,
+):
+    """Create an ELF object file that contains the binary data for the modules
+    imported in `module`. This is a callback function for use as `fpack_imports`
+    in `export_library`.
+
+    Parameters
+    ----------
+    module: tvm.runtime.Module
+        Module whose imported modules need to be serialized.
+    is_system_lib: bool
+        Flag whether the exported module will be used as a system library.
+    c_symbol_prefix: str
+        Prefix to prepend to the blob symbol.
+    workspace_dir: str
+        Location for created files.
+
+    Returns
+    -------
+    file_name: str
+        The name of the created object file.
+    """
+
+    path_bin = os.path.join(workspace_dir, "imports.bin")
+    pack_to_bin_f_name = "runtime.ModulePackImportsToTensor"
+    fpack_to_bin = tvm.get_global_func(pack_to_bin_f_name)
+    assert fpack_to_bin, f"Expecting {pack_to_bin_f_name} in registry"
+
+    fpack_to_bin(module).numpy().tofile(path_bin)
+
+    mblob_symbol = c_symbol_prefix + tvm.get_global_func("runtime.ModuleImportsBlobName")()
+
+    binary_size = os.path.getsize(path_bin)
+    hexagon_toolchain = os.environ.get("HEXAGON_TOOLCHAIN")
+    assert hexagon_toolchain, "Please set HEXAGON_TOOLCHAIN variable"
+    version = toolchain_version(hexagon_toolchain)
+    assert version[0] == 8 and version[1] >= 5, (
+        "Please use Hexagon toolchain version 8.5.x or later"
+    )
+    if version[1] <= 6:
+        path_o = os.path.join(workspace_dir, f"{c_symbol_prefix}devc.o")
+        subprocess.run(
+            [
+                f"{hexagon_toolchain}/bin/hexagon-clang",
+                "-x",
+                "c",
+                "-c",
+                "/dev/null",
+                "-o",
+                path_o,
+            ],
+            check=True,
+        )
+        subprocess.run(
+            [
+                f"{hexagon_toolchain}/bin/hexagon-llvm-objcopy",
+                path_o,
+                "--add-section",
+                f".rodata={path_bin}",
+                "--add-symbol",
+                f"{mblob_symbol}=.rodata:0,object",
+            ],
+            check=True,
+        )
+        return path_o
+
+    else:  # 8.6.07+
+        path_c = os.path.join(workspace_dir, f"{c_symbol_prefix}devc.c")
+        path_o = os.path.join(workspace_dir, f"{c_symbol_prefix}devc.o")
+        with open(path_c, "w") as f:
+            f.write(
+                f"const unsigned char {mblob_symbol}[{binary_size}] "
+                f'__attribute__((section(".rodata"))) = {{0x1}};'
+            )
+        subprocess.run(
+            [f"{hexagon_toolchain}/bin/hexagon-clang", "-c", path_c, "-o", path_o], check=True
+        )
+        subprocess.run(
+            [
+                f"{hexagon_toolchain}/bin/hexagon-llvm-objcopy",
+                path_o,
+                "--update-section",
+                f".rodata={path_bin}",
+            ],
+            check=True,
+        )
+        return path_o
+
+
+def export_module(module, out_dir, binary_name="test_binary.so"):
+    """Export Hexagon shared object to a file."""
+    binary_path = pathlib.Path(out_dir) / binary_name
+    module.write_to_file(str(binary_path))
+    return binary_path
+
+
+def allocate_hexagon_array(
+    dev, tensor_shape=None, dtype=None, data=None, axis_separators=None, mem_scope=None
+):
+    """
+    Allocate a hexagon array which could be a 2D array
+    on physical memory defined by axis_separators
+    """
+    if tensor_shape is None:
+        assert data is not None, "Must provide either tensor shape or numpy data array"
+        tensor_shape = data.shape
+    elif data is not None:
+        assert tensor_shape == data.shape, (
+            "Mismatch between provided tensor shape and numpy data array shape"
+        )
+
+    if dtype is None:
+        assert data is not None, "Must provide either dtype or numpy data array"
+        dtype = data.dtype.name
+    elif data is not None:
+        assert dtype == data.dtype, "Mismatch between provided dtype and numpy data array dtype"
+
+    if axis_separators is None:
+        axis_separators = []
+
+    boundaries = [0, *axis_separators, len(tensor_shape)]
+    physical_shape = [
+        numpy.prod(tensor_shape[dim_i:dim_f]) for dim_i, dim_f in itertools.pairwise(boundaries)
+    ]
+
+    arr = tvm.runtime.empty(physical_shape, dtype=dtype, device=dev, mem_scope=mem_scope)
+
+    if data is not None:
+        arr.copyfrom(data.reshape(physical_shape))
+
+    return arr._create_view(tensor_shape)
+
+
+class ContainerSession:
+    """Docker container session
+
+    Parameters
+    ----------
+    base_image_name : str
+        Docker image name to use. Empty string means to use default "tlcpack/ci-hexagon"
+        base image.
+    """
+
+    def __init__(self, base_image_name: str = ""):
+        self._client = None
+        self._container = None
+        self.tmp_dir = None
+
+        self._client = ContainerSession._get_docker_client()
+
+        if base_image_name == "":
+            base_image_name = ContainerSession._get_latest_ci_image(self._client)
+
+        self._container = ContainerSession._find_container_or_create(self._client, base_image_name)
+
+        exit_code, tmp_dir_b = self._container.exec_run("mktemp -d -t tvm-toolbox-XXXXXXXXXX")
+        assert exit_code == 0
+
+        self.tmp_dir = tmp_dir_b.decode("utf-8").rstrip()
+
+    def __del__(self):
+        self.close()
+
+    @staticmethod
+    def _get_latest_ci_image(client) -> str:
+        ci_images = client.images.list(name="tlcpack/ci-hexagon")
+        ci_images.sort(reverse=True, key=lambda img: img.tags[0])
+        return ci_images[0].tags[0]
+
+    @staticmethod
+    def _get_docker_client():
+        try:
+            from docker import from_env  # pylint: disable=import-outside-toplevel
+            from docker.errors import DockerException
+        except (ModuleNotFoundError, ImportError):
+            raise Exception("Docker SDK module is not installed. Please install it.")
+
+        try:
+            client = from_env()
+        except DockerException:
+            raise Exception(
+                "Docker server is not available. Please verify the docker is installed, "
+                "launched and available via command line ('dokcer ps' should works)."
+            )
+
+        return client
+
+    @staticmethod
+    def _find_container_or_create(client, image_name: str):
+        all_containers = client.containers.list(all=True)
+
+        filtered_containers = []
+        for container in all_containers:
+            tags: list = container.image.tags
+            img_name: str = tags[0]
+            if img_name.startswith(image_name) and container.name.startswith("tvm-hex-toolbox"):
+                filtered_containers.append(container)
+
+        if len(filtered_containers) == 0:
+            container = client.containers.run(
+                image=image_name, detach=True, tty=True, name="tvm-hex-toolbox"
+            )
+        else:
+            container = filtered_containers[0]
+
+        if container.status != "running":
+            container.start()
+
+        return container
+
+    def exec(self, cmd) -> str:
+        """Execute command inside docker container"""
+        exit_code, res = self._container.exec_run(cmd)
+        assert exit_code == 0
+        return res.decode("utf-8")
+
+    def get_env(self, key: str) -> str:
+        """Return env var value from docker container"""
+        res: str = self.exec(f"bash -c 'echo \"${key}\"'")
+        return res.rstrip(" \n")
+
+    def copy_to(self, host_file_path: str) -> str:
+        """Upload file to docker container"""
+        file_name = os.path.basename(host_file_path)
+
+        byte_stream = io.BytesIO()
+        with tarfile.open(fileobj=byte_stream, mode="w:gz") as tar:
+            tar.add(host_file_path, arcname=file_name)
+
+        self._container.put_archive(path=self.tmp_dir, data=byte_stream.getvalue())
+
+        return f"{self.tmp_dir}/{file_name}"
+
+    def copy_from(self, container_file_path: str, host_file_path: str):
+        """Download file from docker container"""
+        tar_bytes_gen, _ = self._container.get_archive(container_file_path)
+
+        # convert to bytes
+        tar_bytes = b""
+        for chunk in tar_bytes_gen:
+            tar_bytes += chunk
+
+        tar = tarfile.open(fileobj=io.BytesIO(initial_bytes=tar_bytes))
+        assert len(tar.getmembers()) == 1
+        tar_element_reader = tar.extractfile(tar.getmembers()[0])
+        with open(host_file_path, "wb") as host_file:
+            for chunk in tar_element_reader:
+                host_file.write(chunk)
+
+    def close(self):
+        """Close docker container session"""
+        if self.tmp_dir is not None:
+            exit_code, _ = self._container.exec_run(f"rm -rf {self.tmp_dir}")
+            assert exit_code == 0

python/tvm/contrib/hipblas.py

@@ -0,0 +1,87 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""External function interface to hipBLAS libraries."""
+
+import tvm
+from tvm import te
+
+
+def matmul(lhs, rhs, transa=False, transb=False, dtype=None):
+    """Create an extern op that compute matrix mult of A and rhs with cuBLAS
+
+    Parameters
+    ----------
+    lhs : Tensor
+        The left matrix operand
+    rhs : Tensor
+        The right matrix operand
+    transa : bool
+        Whether transpose lhs
+    transb : bool
+        Whether transpose rhs
+
+    Returns
+    -------
+    C : Tensor
+        The result tensor.
+    """
+    n = lhs.shape[1] if transa else lhs.shape[0]
+    m = rhs.shape[0] if transb else rhs.shape[1]
+    dtype = dtype if dtype is not None else lhs.dtype
+    return te.extern(
+        (n, m),
+        [lhs, rhs],
+        lambda ins, outs: tvm.tirx.call_packed(
+            "tvm.contrib.hipblas.matmul", ins[0], ins[1], outs[0], transa, transb
+        ),
+        dtype=dtype,
+        name="matmul_hipblas",
+    )
+
+
+def batch_matmul(lhs, rhs, transa=False, transb=False, dtype=None):
+    """Create an extern op that compute batch matrix mult of A and rhs with cuBLAS
+
+    Parameters
+    ----------
+    lhs : Tensor
+        The left matrix operand
+    rhs : Tensor
+        The right matrix operand
+    transa : bool
+        Whether transpose lhs
+    transb : bool
+        Whether transpose rhs
+
+    Returns
+    -------
+    C : Tensor
+        The result tensor.
+    """
+    b = lhs.shape[0]
+    n = lhs.shape[2] if transa else lhs.shape[1]
+    m = rhs.shape[1] if transb else rhs.shape[2]
+    dtype = dtype if dtype is not None else lhs.dtype
+    return te.extern(
+        (b, n, m),
+        [lhs, rhs],
+        lambda ins, outs: tvm.tirx.call_packed(
+            "tvm.contrib.hipblas.batch_matmul", ins[0], ins[1], outs[0], transa, transb
+        ),
+        dtype=dtype,
+        name="batch_matmul_hipblas",
+    )

python/tvm/contrib/mkl.py

@@ -0,0 +1,127 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""External function interface to BLAS libraries."""
+
+import tvm
+from tvm import te
+
+
+def matmul(lhs, rhs, transa=False, transb=False, **kwargs):
+    """Create an extern op that compute matrix mult of A and rhs with CrhsLAS
+    This function serves as an example on how to call external libraries.
+
+    Parameters
+    ----------
+    lhs: Tensor
+        The left matrix operand
+    rhs: Tensor
+        The right matrix operand
+    transa: bool
+        Whether transpose lhs
+    transb: bool
+        Whether transpose rhs
+
+    Returns
+    -------
+    C: Tensor
+        The result tensor.
+    """
+    n = lhs.shape[1] if transa else lhs.shape[0]
+    m = rhs.shape[0] if transb else rhs.shape[1]
+    return te.extern(
+        (n, m),
+        [lhs, rhs],
+        lambda ins, outs: tvm.tirx.call_packed(
+            "tvm.contrib.mkl.matmul", ins[0], ins[1], outs[0], transa, transb
+        ),
+        name="C",
+        **kwargs,
+    )
+
+
+def matmul_u8s8s32(lhs, rhs, transa=False, transb=False, **kwargs):
+    """Create an extern op that compute matrix mult of A and rhs with CrhsLAS
+    This function serves as an example on how to call external libraries.
+
+    Parameters
+    ----------
+    lhs: Tensor
+        The left matrix operand
+    rhs: Tensor
+        The right matrix operand
+    transa: bool
+        Whether transpose lhs
+    transb: bool
+        Whether transpose rhs
+
+    Returns
+    -------
+    C: Tensor
+        The result tensor.
+    """
+    n = lhs.shape[1] if transa else lhs.shape[0]
+    m = rhs.shape[0] if transb else rhs.shape[1]
+    return te.extern(
+        (n, m),
+        [lhs, rhs],
+        lambda ins, outs: tvm.tirx.call_packed(
+            "tvm.contrib.mkl.matmul_u8s8s32", ins[0], ins[1], outs[0], transa, transb
+        ),
+        name="C",
+        **kwargs,
+    )
+
+
+def batch_matmul(lhs, rhs, transa=False, transb=False, iterative=False, **kwargs):
+    """Create an extern op that compute batched matrix mult of A and rhs with mkl
+    This function serves as an example on how to call external libraries.
+
+    Parameters
+    ----------
+    lhs: Tensor
+        The left matrix operand
+    rhs: Tensor
+        The right matrix operand
+    transa: bool
+        Whether transpose lhs
+    transb: bool
+        Whether transpose rhs
+
+    Returns
+    -------
+    C: Tensor
+        The result tensor.
+    """
+    b = te.max(lhs.shape[0], rhs.shape[0])
+    n = lhs.shape[2] if transa else lhs.shape[1]
+    m = rhs.shape[1] if transb else rhs.shape[2]
+    return te.extern(
+        (b, n, m),
+        [lhs, rhs],
+        lambda ins, outs: tvm.tirx.call_packed(
+            "tvm.contrib.mkl.batch_matmul"
+            if not iterative
+            else "tvm.contrib.mkl.batch_matmul_iterative",
+            ins[0],
+            ins[1],
+            outs[0],
+            transa,
+            transb,
+        ),
+        name="C",
+        **kwargs,
+    )

python/tvm/contrib/nnpack.py

@@ -0,0 +1,238 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F821
+"""External function interface to NNPACK libraries."""
+
+import tvm_ffi
+
+import tvm
+from tvm import te
+
+
+def is_available():
+    """Check whether NNPACK is available, that is, `nnp_initialize()`
+    returns `nnp_status_success`.
+    """
+    return _initialize() == 0
+
+
+def fully_connected_inference(lhs, rhs, nthreads=1):
+    """Create an extern op that compute fully connected of 1D tensor lhs and
+    2D tensor rhs with nnpack.
+
+    Parameters
+    ----------
+    lhs : Tensor
+        lhs 1D array input[input_channels] of FP32 elements
+    rhs : Tensor
+        lhs 2D matrix kernel[output_channels][input_channels] of FP32 elements
+
+    Returns
+    -------
+    C : Tensor
+        lhs 1D array out[output_channels] of FP32 elements.
+    """
+    m = rhs.shape[0]
+    return te.extern(
+        (m,),
+        [lhs, rhs],
+        lambda ins, outs: tvm.tirx.call_packed(
+            "tvm.contrib.nnpack.fully_connected_inference", ins[0], ins[1], outs[0], nthreads
+        ),
+        name="C",
+    )
+
+
+class ConvolutionAlgorithm:
+    AUTO = 0
+    FFT_8x8 = 1
+    FFT_16x16 = 2
+    WT_8x8 = 3
+    IMPLICIT_GEMM = 4
+    DIRECT = 5
+    WT_8x8_FP16 = 6
+
+
+class ConvolutionTransformStrategy:
+    COMPUTE = 1
+    PRECOMPUTE = 2
+
+
+def convolution_inference(
+    data, kernel, bias, padding, stride, nthreads=1, algorithm=ConvolutionAlgorithm.AUTO
+):
+    """Create an extern op to do inference convolution of 4D tensor data and
+    4D tensor kernel and 1D tensor bias with nnpack.
+
+    Parameters
+    ----------
+    data : Tensor
+        data 4D tensor input[batch][input_channels][input_height][input_width] of
+        FP32 elements.
+    kernel : Tensor
+        kernel 4D tensor kernel[output_channels][input_channels][kernel_height]
+        [kernel_width] of FP32 elements.
+    bias : Tensor
+        bias 1D array bias[output_channels][input_channels][kernel_height]
+        [kernel_width] of FP32 elements.
+    padding : list
+        padding A 4-dim list of [pad_top, pad_bottom, pad_left, pad_right],
+        which indicates the padding around the feature map.
+    stride : list
+        stride A 2-dim list of [stride_height, stride_width], which indicates
+        the stride.
+
+    Returns
+    -------
+    output : Tensor
+        output 4D tensor output[batch][output_channels][output_height][output_width]
+        of FP32 elements.
+    """
+
+    assert isinstance(padding, list) and len(padding) == 4
+    assert isinstance(stride, list) and len(stride) == 2
+    batch, _, input_height, input_width = data.shape
+    output_channels, _, kernel_height, kernel_width = kernel.shape
+    idxdiv = te.indexdiv
+    output_height = idxdiv(input_height + padding[0] + padding[1] - kernel_height, stride[0]) + 1
+    output_width = idxdiv(input_width + padding[0] + padding[1] - kernel_width, stride[1]) + 1
+
+    return te.extern(
+        (batch, output_channels, output_height, output_width),
+        [data, kernel, bias] if bias is not None else [data, kernel],
+        lambda ins, outs: tvm.tirx.call_packed(
+            "tvm.contrib.nnpack.convolution_inference",
+            ins[0],
+            ins[1],
+            ins[2] if bias is not None else 0,
+            outs[0],
+            padding[0],
+            padding[1],
+            padding[2],
+            padding[3],
+            stride[0],
+            stride[1],
+            nthreads,
+            algorithm,
+        ),
+        name="C",
+    )
+
+
+def convolution_inference_without_weight_transform(
+    data, transformed_kernel, bias, padding, stride, nthreads=1, algorithm=ConvolutionAlgorithm.AUTO
+):
+    """Create an extern op to do inference convolution of 4D tensor data and
+    4D pre-transformed tensor kernel and 1D tensor bias with nnpack.
+
+    Parameters
+    ----------
+    data : Tensor
+        data 4D tensor input[batch][input_channels][input_height][input_width] of
+        FP32 elements.
+    transformed_kernel : Tensor
+        transformed_kernel 4D tensor kernel[output_channels][input_channels][tile]
+        [tile] of FP32 elements.
+    bias : Tensor
+        bias 1D array bias[output_channels][input_channels][kernel_height]
+        [kernel_width] of FP32 elements.
+    padding : list
+        padding A 4-dim list of [pad_top, pad_bottom, pad_left, pad_right],
+        which indicates the padding around the feature map.
+    stride : list
+        stride A 2-dim list of [stride_height, stride_width], which indicates
+        the stride.
+
+    Returns
+    -------
+    output : Tensor
+        output 4D tensor output[batch][output_channels][output_height][output_width]
+        of FP32 elements.
+    """
+
+    assert algorithm in (ConvolutionAlgorithm.WT_8x8, ConvolutionAlgorithm.WT_8x8_FP16)
+    assert isinstance(padding, list) and len(padding) == 4
+    assert isinstance(stride, list) and len(stride) == 2
+    batch, _, input_height, input_width = data.shape
+    output_channels, _, _, _ = transformed_kernel.shape
+    kernel_height, kernel_width = (3, 3)
+    idxdiv = te.indexdiv
+    output_height = idxdiv(input_height + padding[0] + padding[1] - kernel_height, stride[0]) + 1
+    output_width = idxdiv(input_width + padding[0] + padding[1] - kernel_width, stride[1]) + 1
+
+    return te.extern(
+        (batch, output_channels, output_height, output_width),
+        [data, transformed_kernel, bias] if bias is not None else [data, transformed_kernel],
+        lambda ins, outs: tvm.tirx.call_packed(
+            "tvm.contrib.nnpack.convolution_inference_without_weight_transform",
+            ins[0],
+            ins[1],
+            ins[2] if bias is not None else 0,
+            outs[0],
+            padding[0],
+            padding[1],
+            padding[2],
+            padding[3],
+            stride[0],
+            stride[1],
+            nthreads,
+            algorithm,
+        ),
+        name="C",
+        dtype="float32",
+    )
+
+
+def convolution_inference_weight_transform(
+    kernel, nthreads=1, algorithm=ConvolutionAlgorithm.AUTO, dtype="float32"
+):
+    """Create an extern op to do inference convolution of 3D tensor data and
+    4D tensor kernel and 1D tensor bias with nnpack.
+
+    Parameters
+    ----------
+    kernel : Tensor
+        kernel 4D tensor kernel[output_channels][input_channels][kernel_height]
+        [kernel_width] of FP32 elements.
+
+    Returns
+    -------
+    output : Tensor
+        output 4D tensor output[output_channels][input_channels][tile][tile]
+        of FP32 elements.
+    """
+    assert algorithm in (ConvolutionAlgorithm.WT_8x8, ConvolutionAlgorithm.WT_8x8_FP16)
+    output_channels, input_channels, _, _ = kernel.shape
+    transform_tile_size = 8
+    if not isinstance(dtype, str):
+        dtype = dtype.dtype
+    return te.extern(
+        (output_channels, input_channels, transform_tile_size, transform_tile_size),
+        [kernel],
+        lambda ins, outs: tvm.tirx.call_packed(
+            "tvm.contrib.nnpack.convolution_inference_weight_transform",
+            ins[0],
+            outs[0],
+            nthreads,
+            algorithm,
+        ),
+        name="transform_kernel",
+        dtype=dtype,
+    )
+
+
+tvm_ffi.init_ffi_api("tvm.contrib.nnpack")

python/tvm/contrib/pickle_memoize.py

@@ -0,0 +1,144 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: RUF005, RUF012
+"""Memoize result of function via pickle, used for cache testcases."""
+
+# pylint: disable=broad-except,superfluous-parens
+import atexit
+import functools
+import os
+import pathlib
+import sys
+
+try:
+    import cPickle as pickle
+except ImportError:
+    import pickle
+
+
+def _get_global_cache_dir() -> pathlib.Path:
+    if "XDG_CACHE_HOME" in os.environ:
+        cache_home = pathlib.Path(os.environ.get("XDG_CACHE_HOME"))
+    else:
+        cache_home = pathlib.Path.home().joinpath(".cache")
+    return cache_home.joinpath("tvm", f"pkl_memoize_py{sys.version_info[0]}")
+
+
+GLOBAL_CACHE_DIR = _get_global_cache_dir()
+
+
+class Cache:
+    """A cache object for result cache.
+
+    Parameters
+    ----------
+    key: str
+       The file key to the function
+    save_at_exit: bool
+        Whether save the cache to file when the program exits
+    """
+
+    cache_by_key = {}
+
+    def __init__(self, key, save_at_exit):
+        self._cache = None
+
+        self.path = GLOBAL_CACHE_DIR.joinpath(key)
+        self.dirty = False
+        self.save_at_exit = save_at_exit
+
+    @property
+    def cache(self):
+        """Return the cache, initializing on first use."""
+
+        if self._cache is not None:
+            return self._cache
+
+        if self.path.exists():
+            with self.path.open("rb") as cache_file:
+                try:
+                    cache = pickle.load(cache_file)
+                except pickle.UnpicklingError:
+                    cache = {}
+        else:
+            cache = {}
+
+        self._cache = cache
+        return self._cache
+
+    def save(self):
+        if self.dirty:
+            self.path.parent.mkdir(parents=True, exist_ok=True)
+
+            with self.path.open("wb") as out_file:
+                pickle.dump(self.cache, out_file, pickle.HIGHEST_PROTOCOL)
+
+
+@atexit.register
+def _atexit():
+    """Save handler."""
+    for value in Cache.cache_by_key.values():
+        if value.save_at_exit:
+            value.save()
+
+
+def memoize(key, save_at_exit=False):
+    """Memoize the result of function and reuse multiple times.
+
+    Parameters
+    ----------
+    key: str
+        The unique key to the file
+    save_at_exit: bool
+        Whether save the cache to file when the program exits
+
+    Returns
+    -------
+    fmemoize : function
+        The decorator function to perform memoization.
+    """
+
+    def _register(f):
+        """Registration function"""
+        allow_types = (str, int, float, tuple)
+        fkey = key + "." + f.__name__ + ".pkl"
+        if fkey not in Cache.cache_by_key:
+            Cache.cache_by_key[fkey] = Cache(fkey, save_at_exit)
+        cache = Cache.cache_by_key[fkey]
+        cargs = tuple(x.cell_contents for x in f.__closure__) if f.__closure__ else ()
+        cargs = (len(cargs),) + cargs
+
+        @functools.wraps(f)
+        def _memoized_f(*args, **kwargs):
+            assert not kwargs, "Only allow positional call"
+            key = cargs + args
+            for arg in key:
+                if isinstance(arg, tuple):
+                    for x in arg:
+                        assert isinstance(x, allow_types)
+                else:
+                    assert isinstance(arg, allow_types)
+            if key in cache.cache:
+                return cache.cache[key]
+            res = f(*args)
+            cache.cache[key] = res
+            cache.dirty = True
+            return res
+
+        return _memoized_f
+
+    return _register

python/tvm/contrib/random.py

@@ -0,0 +1,117 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""External function interface to random library."""
+
+import tvm_ffi
+
+import tvm
+from tvm import te
+
+
+def randint(low, high, size, dtype="int32"):
+    """Return random integers from low (inclusive) to high (exclusive).
+    Return random integers from the "discrete uniform" distribution of the
+    specified dtype in the "half-open" interval [low, high).
+
+    Parameters
+    ----------
+    low : int
+        Lowest (signed) integer to be drawn from the distribution
+    high : int
+        One above the largest (signed) integer to be drawn from the distribution
+
+    Returns
+    -------
+    out : Tensor
+        A tensor with specified size and dtype
+    """
+    assert "int" in dtype, "the type of randint output must be int or uint"
+    return te.extern(
+        size,
+        [],
+        lambda ins, outs: tvm.tirx.call_packed(
+            "tvm.contrib.random.randint", int(low), int(high), outs[0]
+        ),
+        dtype=dtype,
+    )
+
+
+def uniform(low, high, size):
+    """Draw samples from a uniform distribution.
+
+    Samples are uniformly distributed over the half-open interval [low, high)
+    (includes low, but excludes high). In other words, any value within the
+    given interval is equally likely to be drawn by uniform.
+
+    Parameters
+    ----------
+    low : float
+        Lower boundary of the output interval. All values generated will be
+        greater than or equal to low.
+    high : float
+        Upper boundary of the output interval. All values generated will be
+        less than high.
+    size : tuple of ints
+        Output shape. If the given shape is, e.g., (m, n, k), then m * n * k
+        samples are drawn.
+
+    Returns
+    -------
+    out : Tensor
+        A tensor with specified size and dtype.
+    """
+    return te.extern(
+        size,
+        [],
+        lambda ins, outs: tvm.tirx.call_packed(
+            "tvm.contrib.random.uniform", float(low), float(high), outs[0]
+        ),
+        dtype="float32",
+    )
+
+
+def normal(loc, scale, size):
+    """Draw samples from a normal distribution.
+
+    Return random samples from a normal distribution.
+
+    Parameters
+    ----------
+    loc : float
+        loc of the distribution.
+    scale : float
+        Standard deviation of the distribution.
+    size : tuple of ints
+        Output shape. If the given shape is, e.g., (m, n, k), then m * n * k
+        samples are drawn.
+
+    Returns
+    ------
+    out : Tensor
+        A tensor with specified size and dtype
+    """
+    return te.extern(
+        size,
+        [],
+        lambda ins, outs: tvm.tirx.call_packed(
+            "tvm.contrib.random.normal", float(loc), float(scale), outs[0]
+        ),
+        dtype="float32",
+    )
+
+
+tvm_ffi.init_ffi_api("tvm.contrib.random")

python/tvm/contrib/thrust.py

@@ -0,0 +1,47 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Utilities for thrust"""
+
+import logging
+
+from tvm_ffi import get_global_func
+
+
+def maybe_warn(target, func_name):
+    if (
+        "thrust" in list(target.attrs.get("libs", []))
+        and get_global_func(func_name, allow_missing=True) is None
+    ):
+        logging.warning("thrust is requested but TVM is not built with thrust.")
+
+
+def can_use_thrust(target, func_name):
+    maybe_warn(target, func_name)
+    return (
+        target.kind.name in ["cuda", "nvptx"]
+        and "thrust" in list(target.attrs.get("libs", []))
+        and get_global_func(func_name, allow_missing=True)
+    )
+
+
+def can_use_rocthrust(target, func_name):
+    maybe_warn(target, func_name)
+    return (
+        target.kind.name == "rocm"
+        and "thrust" in list(target.attrs.get("libs", []))
+        and get_global_func(func_name, allow_missing=True)
+    )

python/tvm/contrib/tvmjs.py

@@ -0,0 +1,422 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: E501, F401
+"""Namespace to store utilities for building web runtime."""
+
+import hashlib
+import json
+import math
+import os
+import shutil
+
+# pylint: disable=unused-import
+import sys
+from collections.abc import Iterator, Mapping
+from types import GeneratorType
+from typing import Any, Optional, Union
+
+import numpy as np
+
+try:
+    import ml_dtypes
+except ImportError:
+    ml_dtypes = None
+
+import tvm
+from tvm.runtime import DataType
+from tvm.support.emcc import create_tvmjs_wasm, find_wasm_lib
+
+
+def _convert_f32_to_bf16(value):
+    cap = np.finfo("float32").max
+    assert -np.finfo("float32").max == np.finfo("float32").min
+    bf16_limit = ((np.array([cap.view("uint32")]) >> 16) << 16).view("float32")[0]
+    # When the value is in [-bf16_limit, bf16_limit], round to nearest even.
+    # We can afford to do it in dumping phase to reduce overall rounding error.
+    #
+    # When the value is out of bound(usually mask values in attention), use truncation
+    # so it is equivalent to clip to the limit values
+    data = value.view("uint32")
+    rounding_bias = np.where(
+        np.logical_and(value < bf16_limit, value > -bf16_limit),
+        ((data >> 16) & 1) + 0x7FFF,
+        np.zeros_like(data),
+    )
+    return ((data + rounding_bias) >> 16).astype("uint16")
+
+
+def _convert_bf16_to_f32(value):
+    data = value.view("uint16")
+    return (data.astype("uint32") << 16).view("float32")
+
+
+def _calculate_md5(filename):
+    hash_md5 = hashlib.md5()
+    with open(filename, "rb") as file:
+        for chunk in iter(lambda: file.read(8192), b""):
+            hash_md5.update(chunk)
+    return hash_md5.hexdigest()
+
+
+class TensorCacheShardingManager:
+    """Internal helper to shard ndarrays."""
+
+    def __init__(
+        self,
+        cache_dir: str,
+        prefix: str,
+        shard_cap_nbytes: int,
+        initial_shard_records: Mapping[str, Any] | None = None,
+    ):
+        self.cache_dir = cache_dir
+        self.prefix = prefix
+        self.curr_records = []
+        self.curr_data = bytearray()
+        self.shard_records = []
+        self.shard_cap_nbytes = shard_cap_nbytes
+        self.counter = 0
+        self.name_to_record: Mapping[str, tuple[int, Mapping[str, Any]]] = {}
+        self.updated_shards: set[int] = set()
+
+        if initial_shard_records is not None:
+            self.shard_records = initial_shard_records
+            self.counter = len(initial_shard_records)
+            for idx, shard in enumerate(initial_shard_records):
+                for rec in shard["records"]:
+                    self.name_to_record[rec["name"]] = (idx, rec)
+
+    def append_or_update(self, data, name, shape, dtype, encode_format, allow_update: bool = False):
+        """Commit a record to the manager.
+
+        Parameters
+        ----------
+        data: bytes
+            Raw bytes to be appended.
+
+        name: str
+            The name of the parameter
+
+        shape: tuple
+            The shape of the array
+
+        dtype: str
+            The dtype information
+
+        encode_format:
+            The encode format of the entry
+
+        allow_update: bool
+            If the record already exists, update the record. Otherwise, raise an error.
+        """
+        rec = {
+            "name": name,
+            "shape": shape,
+            "dtype": dtype,
+            "format": encode_format,
+            "nbytes": len(data),
+        }
+        if name in self.name_to_record:
+            if not allow_update:
+                raise ValueError(f"Duplicate name {name} found in the cache.")
+            self.update_single_record(rec, data)
+            return
+
+        self.name_to_record[name] = (self.counter, rec)
+
+        if self.pending_nbytes + len(data) >= self.shard_cap_nbytes:
+            if len(data) * 2 >= self.shard_cap_nbytes:
+                # out of band data
+                rec["byteOffset"] = 0
+                self._commit_internal(data, [rec])
+                return
+            self.commit()
+        rec["byteOffset"] = self.pending_nbytes
+        self.curr_records.append(rec)
+        self.curr_data += data
+
+    def update_single_record(self, rec, data):
+        """Update a single record in a shard file."""
+        name = rec["name"]
+        idx, old_rec = self.name_to_record[name]
+        if old_rec["nbytes"] != rec["nbytes"]:
+            raise ValueError(f"Cannot update record {name}, size mismatch.")
+        data_path = self.shard_records[idx]["dataPath"]
+        full_path = os.path.join(self.cache_dir, data_path)
+        with open(full_path, "r+b") as outfile:
+            outfile.seek(old_rec["byteOffset"])
+            outfile.write(data)
+        self.name_to_record[name] = (idx, rec)
+        self.updated_shards.add(idx)
+
+    def commit(self):
+        """Commit a record"""
+        if self.pending_nbytes != 0:
+            self._commit_internal(self.curr_data, self.curr_records)
+            self.curr_data = bytearray()
+            self.curr_records = []
+
+    def finish(self):
+        """Finish building and return shard records."""
+        self.commit()
+        for idx in self.updated_shards:
+            full_path = os.path.join(self.cache_dir, self.shard_records[idx]["dataPath"])
+            self.shard_records[idx]["md5sum"] = _calculate_md5(full_path)
+        return self.shard_records
+
+    def _commit_internal(self, data, records):
+        data_path = f"{self.prefix}_{self.counter}.bin"
+        full_path = os.path.join(self.cache_dir, data_path)
+        self.counter += 1
+        with open(full_path, "wb") as outfile:
+            outfile.write(data)
+
+        shard_record = {
+            "dataPath": data_path,
+            "format": "raw-shard",
+            "nbytes": len(data),
+            "records": records,
+            "md5sum": _calculate_md5(full_path),
+        }
+        self.shard_records.append(shard_record)
+
+    @property
+    def pending_nbytes(self):
+        """Return total bytes stored so far"""
+        return len(self.curr_data)
+
+
+def dump_tensor_cache(
+    params: Mapping[str, np.ndarray | tvm.runtime.Tensor]
+    | Iterator[tuple[str, np.ndarray | tvm.runtime.Tensor]],
+    cache_dir: str,
+    encode_format="f32-to-bf16",
+    meta_data=None,
+    shard_cap_mb=32,
+    show_progress: bool = True,
+    update_if_exists: bool = False,
+):
+    """Dump parameters to Tensor cache.
+
+    Parameters
+    ----------
+    params: Union[
+        Mapping[str, Union[np.ndarray, tvm.runtime.Tensor]],
+        Iterator[Tuple[str, Union[np.ndarray, tvm.runtime.Tensor]]],
+    ]
+        The parameter dictionary or generator
+
+    cache_dir: str
+        The path to the cache
+
+    encode_format: {"f32-to-bf16", "raw"}
+        Encoding format.
+
+    meta_data: json-compatible-struct or Callable[[], Any]
+        Extra meta_data to be stored in the cache json file,
+        or a callable that returns the metadata.
+
+    shard_cap_mb: int
+        Maxinum number of MB to be kept per shard
+
+    show_progress: bool
+        A boolean indicating if to show the dump progress.
+
+    update_if_exists: bool
+        If the cache already exists, update the cache. When set to False, it will overwrite the
+        existing files.
+    """
+    if encode_format not in ("raw", "f32-to-bf16"):
+        raise ValueError(f"Invalie encode_format {encode_format}")
+
+    records = []
+    from_generator = isinstance(params, GeneratorType)
+    total_bytes = 0
+    counter = 0
+    max_out_length = 0
+
+    if not os.path.exists(cache_dir):
+        os.makedirs(cache_dir)
+
+    f32_to_bf16_triggered = False
+
+    print(f"Start storing to cache {cache_dir}")
+    shard_cap_nbytes = shard_cap_mb * (1 << 20)
+
+    nd_cache_json = os.path.join(cache_dir, "tensor-cache.json")
+    if update_if_exists and os.path.exists(nd_cache_json):
+        with open(nd_cache_json) as infile:
+            old_data = json.load(infile)
+            if meta_data is None:
+                meta_data = old_data["metadata"]
+            records = old_data["records"]
+
+    shard_manager = TensorCacheShardingManager(
+        cache_dir, "params_shard", shard_cap_nbytes, initial_shard_records=records
+    )
+
+    param_generator = params.items() if not from_generator else params
+    for k, origin_v in param_generator:
+        shape = list(origin_v.shape)
+        v = origin_v
+        if not isinstance(v, np.ndarray):
+            v = v.numpy()
+
+        # prefer to preserve original dtype, especially if the format was bfloat16
+        dtype = origin_v.dtype if isinstance(origin_v, tvm.runtime.Tensor) else v.dtype
+
+        if dtype in DataType._NUMPY_DTYPE_TO_STR:
+            dtype = DataType._NUMPY_DTYPE_TO_STR[dtype]
+        else:
+            dtype = str(dtype)
+
+        total_bytes += math.prod(v.shape) * np.dtype(v.dtype).itemsize
+
+        # convert fp32 to bf16
+        if encode_format == "f32-to-bf16" and dtype == "float32":
+            data = _convert_f32_to_bf16(v).tobytes()
+            f32_to_bf16_triggered = True
+        else:
+            data = v.tobytes()
+
+        shard_manager.append_or_update(
+            data,
+            name=k,
+            shape=shape,
+            dtype=dtype,
+            encode_format=encode_format,
+            allow_update=update_if_exists,
+        )
+
+        counter += 1
+        if show_progress:
+            last_cmd = f"[{counter:04d}] saving {k}"
+            flush = "\r" + (" " * max_out_length) + "\r"
+            max_out_length = max(len(last_cmd), max_out_length)
+            sys.stdout.write(flush + last_cmd)
+
+    records = shard_manager.finish()
+    meta_data = {} if meta_data is None else meta_data if not callable(meta_data) else meta_data()
+
+    with open(nd_cache_json, "w") as outfile:
+        json.dump({"metadata": meta_data, "records": records}, outfile, indent=4)
+    print(
+        f"\nAll finished, {shard_manager.counter} total shards committed, record saved to {nd_cache_json}"
+    )
+
+    if f32_to_bf16_triggered:
+        for shard in records:
+            for item in shard["records"]:
+                if item["dtype"] == "float32":
+                    item["format"] = "raw"
+                    item["dtype"] = "bfloat16"
+        b16_nd_cache_json = os.path.join(cache_dir, "tensor-cache-b16.json")
+        # also dump a file that contains bf16
+        with open(b16_nd_cache_json, "w") as outfile:
+            json.dump({"metadata": meta_data, "records": records}, outfile, indent=4)
+        print(f"Also saved a bf16 record to {b16_nd_cache_json}")
+
+
+def load_tensor_cache(cachepath: str, device: tvm.runtime.Device):
+    """Load the tensor cache from the directory or json.
+
+
+    Parameters
+    ----------
+    cachepath: str
+        Path to the location or json file.
+
+    device: tvm.runtime.Device
+        The device we would like to load the data from.
+    """
+    if not cachepath.endswith(".json"):
+        cachepath = os.path.join(cachepath, "tensor-cache.json")
+
+    cachedir = os.path.dirname(cachepath)
+    json_info = json.loads(open(cachepath).read())
+    result_dict = {}
+
+    for shard_rec in json_info["records"]:
+        data_path = shard_rec["dataPath"]
+        full_data_path = os.path.join(cachedir, data_path)
+        raw_data = open(full_data_path, "rb").read()
+        assert shard_rec["format"] == "raw-shard"
+        assert shard_rec["nbytes"] == len(raw_data)
+
+        for rec in shard_rec["records"]:
+            name = rec["name"]
+            shape = rec["shape"]
+            dtype = rec["dtype"]
+            encode_format = rec["format"]
+            offset = rec["byteOffset"]
+            nbytes = rec["nbytes"]
+
+            arr = tvm.runtime.empty(shape, dtype, device=device)
+            assert offset + nbytes <= len(raw_data)
+            buffer_source = raw_data[offset : offset + nbytes]
+            if dtype == "float8_e4m3fn":
+                if ml_dtypes is not None:
+                    dtype = ml_dtypes.float8_e4m3fn
+                else:
+                    raise RuntimeError(
+                        "ml_dtypes is not installed, cannot convert float8_e4m3fn array to numpy."
+                    )
+            if dtype == "float8_e5m2":
+                if ml_dtypes is not None:
+                    dtype = ml_dtypes.float8_e5m2
+                else:
+                    raise RuntimeError(
+                        "ml_dtypes is not installed, cannot convert float8_e5m2 array to numpy."
+                    )
+            if encode_format == "f32-to-bf16" and dtype == "float32":
+                data = np.frombuffer(buffer_source, dtype="uint16").reshape(shape)
+                arr.copyfrom(_convert_bf16_to_f32(data))
+            elif dtype == "bfloat16":
+                data = np.frombuffer(buffer_source, dtype="uint16").reshape(shape)
+                arr.copyfrom(data)
+            else:
+                data = np.frombuffer(buffer_source, dtype=dtype).reshape(shape)
+                arr.copyfrom(data)
+            result_dict[name] = arr
+    return result_dict, json_info["metadata"]
+
+
+def export_runtime(runtime_dir):
+    """Export TVMJS runtime to the runtime_dir
+
+    Parameters
+    ----------
+    runtime_dir: str
+        The runtime directory
+    """
+    web_hint = (
+        "make sure you setup tvm web runtime correctly."
+        + " obtain a copy of TVM source code, set TVM_HOME env variable:\n"
+        + " cd /path/to/tvm/web; make; npm run bundle"
+    )
+
+    jsbundle = find_wasm_lib("tvmjs.bundle.js", optional=True)
+    if not jsbundle:
+        raise RuntimeError("Cannot find tvmjs.bundle.js, " + web_hint)
+
+    wasi = find_wasm_lib("tvmjs_runtime.wasi.js", optional=True)
+    if not wasi:
+        raise RuntimeError("Cannot find tvmjs_runtime.wasi.js, " + web_hint)
+
+    print(f"Copy {jsbundle[0]} to {runtime_dir}")
+    shutil.copy(jsbundle[0], runtime_dir)
+    print(f"Copy {wasi[0]} to {runtime_dir}")
+    shutil.copy(wasi[0], runtime_dir)