apache--tvm/python/tvm/relax/backend/cuda/cutlass.py

# 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: E731
"""Pattern table for CUTLASS backend"""

import operator
from collections.abc import Mapping, Sequence
from functools import reduce

import tvm
from tvm.contrib.cutlass.build import is_shape_valid_for_cutlass_matmul
from tvm.relax import (
    Call,
    ExternFunc,
    Function,
    PyExprMutator,
    Var,
    expr_functor,
    transform,
)
from tvm.relax.dpl import rewrite_call
from tvm.relax.dpl.pattern import GlobalVarPattern, TuplePattern, is_op, wildcard
from tvm.relax.transform import PatternCheckContext

from ..pattern_registry import get_patterns_with_prefix, register_patterns
from ..patterns import (
    make_attention_pattern,
    make_attention_rewrite_pattern,
    make_fused_bias_activation_pattern,
    make_layer_norm_pattern,
    make_matmul_pattern,
    make_residual_block_pattern,
    make_rms_norm_pattern,
    make_stacked_attention_pattern,
)
from ..utils import has_leaking_intermediate_variables


def _is_supported_dtype(lhs_dtype, rhs_dtype):
    """Check if dtypes in the given workload are supported by CUTLASS."""
    return (
        (lhs_dtype == "float16" and rhs_dtype == "float16")
        or (lhs_dtype == "float32" and rhs_dtype == "float32")
        or (lhs_dtype in ("int8", "uint8") and rhs_dtype in ("int8", "uint8"))
    )


def _shape_1d(shape):
    return reduce(operator.mul, shape, 1)


def _has_dependency(from_var: Var, to_var: Var, var_usages: Mapping[Var, Sequence[Var]]):
    if from_var == to_var:
        return True

    checked = set()
    vars_to_check = [to_var]
    while vars_to_check:
        current_var = vars_to_check.pop()
        for user in var_usages.get(current_var, []):
            if user == from_var:
                return True
            if user not in checked:
                checked.add(user)
                vars_to_check.append(user)

    return False


def _is_same_shape(shape1, shape2):
    analyzer = tvm.arith.Analyzer()
    return all([analyzer.can_prove_equal(s1, s2) for s1, s2 in zip(shape1, shape2)])


def _is_bias_like(shape, out_channel):
    return shape[-1] == out_channel and _shape_1d(shape) == out_channel


def _check_residual(root_call: Call, context: PatternCheckContext) -> bool:
    if "residual" in context.annotated_expr:
        residual = context.annotated_expr["residual"]
        if not isinstance(residual, Var):
            if residual not in context.value_to_bound_var:
                return False

            residual = context.value_to_bound_var[residual]

        root_var = context.value_to_bound_var[root_call]
        if _has_dependency(from_var=residual, to_var=root_var, var_usages=context.var_usages):
            # If residual depends on the result of the root call, this cannot be handled by cutlass.
            return False

        shape1 = root_var.ty.shape
        shape2 = residual.ty.shape
        out_channel = shape1[-1]

        if not _is_same_shape(shape1, shape2) and not _is_bias_like(shape2, out_channel):
            return False

    return True


def _check_conv2d(context: PatternCheckContext) -> bool:
    """Check if the given conv2d workload can be offloaded to CUTLASS."""
    if has_leaking_intermediate_variables(context):
        return False

    conv2d_call = context.annotated_expr["root"]
    data_layout = conv2d_call.attrs.data_layout
    kernel_layout = conv2d_call.attrs.kernel_layout
    data, weight, *_ = conv2d_call.args
    if (
        data_layout != "NHWC"
        or kernel_layout != "OHWI"
        or not _is_supported_dtype(data.ty.dtype, weight.ty.dtype)
    ):
        return False

    if not _check_residual(conv2d_call, context):
        return False

    # Check if any dimensions are symbolic.
    for dim in data.ty.shape.values:
        if isinstance(dim, tvm.tirx.Var):
            return False

    # pylint: disable=invalid-name
    IC = data.ty.shape.values[3]
    OC = weight.ty.shape.values[0]
    # not depthwise conv2d
    return not IC == OC == conv2d_call.attrs.groups


def _check_matmul(context: PatternCheckContext) -> bool:
    """Check if the given matmul workload can be offloaded to CUTLASS."""
    if has_leaking_intermediate_variables(context):
        return False

    lhs = context.annotated_expr["lhs"]
    rhs = context.annotated_expr["rhs"]

    lhs_dtype = lhs.ty.dtype
    rhs_dtype = rhs.ty.dtype
    if not _is_supported_dtype(lhs_dtype, rhs_dtype):
        return False

    if not _check_residual(context.annotated_expr["root"], context):
        return False

    lhs_shape = lhs.ty.shape.values
    rhs_shape = rhs.ty.shape.values
    return is_shape_valid_for_cutlass_matmul(lhs_shape, rhs_shape)


def _get_activation_from_name(pattern_name):
    if "_relu" in pattern_name:
        return "relax.nn.relu"
    elif "_gelu_tanh" in pattern_name:
        return "relax.nn.gelu_tanh"
    elif "_gelu" in pattern_name:
        return "relax.nn.gelu"
    elif "_silu" in pattern_name:
        return "relax.nn.silu"
    else:
        return None


def matmul_patterns():
    """
    Returns a list of all matmul patterns in cutlass BYOC backend.
    """

    def _matmul_pattern(pattern_name):
        transposed_rhs = "_transposed" in pattern_name
        with_bias = "_bias" in pattern_name
        activation = _get_activation_from_name(pattern_name)

        return (
            pattern_name,
            *make_matmul_pattern(
                transposed_rhs=transposed_rhs,
                with_bias=with_bias,
                activation=activation,
            ),
            _check_matmul,
        )

    return [
        _matmul_pattern("cutlass.matmul"),
        _matmul_pattern("cutlass.matmul_bias"),
        _matmul_pattern("cutlass.matmul_bias_relu"),
        _matmul_pattern("cutlass.matmul_bias_gelu"),
        _matmul_pattern("cutlass.matmul_transposed"),
        _matmul_pattern("cutlass.matmul_transposed_bias"),
        _matmul_pattern("cutlass.matmul_transposed_bias_relu"),
        _matmul_pattern("cutlass.matmul_transposed_bias_gelu"),
    ]


def _check_decode_matmul(ctx):
    """Check if the given decode -> matmul workload can be offloaded to CUTLASS."""
    if has_leaking_intermediate_variables(ctx):
        return False

    root = ctx.annotated_expr["root"]

    if not _check_residual(root, ctx):
        return False

    # out_dtype = "float32" not supported unless matmul is followed by cast to fp16.
    if root.ty.dtype == "float32":
        return False

    call_tir_decode = ctx.annotated_expr["w_decoded"]
    if "decode" not in call_tir_decode.args[0].name_hint:
        return False

    N = root.ty.shape[-1]

    if ctx.annotated_expr["lhs"].ty.dtype != "float16":
        return False

    # weight needs to be packed to int8.
    packed_weight = ctx.annotated_expr["w_encoded"]

    if packed_weight.ty.dtype != "int8":
        return False

    # The kernel expects the weight to be preprocessed by this packed function.
    if (
        isinstance(packed_weight, Call)
        and isinstance(packed_weight.args[0], ExternFunc)
        and packed_weight.args[0].global_symbol != "cutlass.ft_preprocess_weight"
    ):
        return False

    scales = ctx.annotated_expr["scales"]

    if scales.ty.dtype != "float16":
        return False

    # scale shape needs to be (N,) or (1, N) or (K // group_size, N)
    if len(scales.ty.shape) > 2 or scales.ty.shape[-1] != N:
        return False

    if "bias" in ctx.annotated_expr:
        out_shape = root.ty.shape
        bias_shape = ctx.annotated_expr["bias"].ty.shape

        # bias shape needs to be (N,), possibly with additional axes on the front.
        # It can also have the same shape as the output.
        if not _is_bias_like(bias_shape, N) and not _is_same_shape(out_shape, bias_shape):
            return False

    return True


def decode_matmul_patterns():
    """Returns a list of supported decode -> matmul patterns."""

    def _decode_matmul_pattern(name):
        scales = wildcard()
        x = wildcard()
        w_packed = wildcard()

        w = is_op("relax.call_tir")(
            GlobalVarPattern(),
            TuplePattern([w_packed, scales]),
        )
        matmul = is_op("relax.matmul")(x, w)

        if "cast" in name:
            matmul = is_op("relax.astype")(matmul)

        annotations = {
            "root": matmul,
            "lhs": x,
            "w_encoded": w_packed,
            "w_decoded": w,
            "scales": scales,
        }

        if "bias" in name:
            annotations["bias"] = bias = wildcard()
            out = is_op("relax.add")(matmul, bias)
        else:
            out = matmul

        if "gelu" in name:
            out = is_op("relax.nn.gelu")(out)

        return name, out, annotations, _check_decode_matmul

    return [
        _decode_matmul_pattern("cutlass.decode_matmul"),
        _decode_matmul_pattern("cutlass.decode_matmul_bias"),
        _decode_matmul_pattern("cutlass.decode_matmul_cast"),
        _decode_matmul_pattern("cutlass.decode_matmul_cast_bias"),
        _decode_matmul_pattern("cutlass.decode_matmul_bias_gelu"),
        _decode_matmul_pattern("cutlass.decode_matmul_cast_bias_gelu"),
    ]


def conv2d_patterns():
    """
    Returns a list of all conv2d patterns in cutlass BYOC backend.
    """

    def _conv2d_pattern(pattern_name):
        with_bias = "_bias" in pattern_name
        activation = _get_activation_from_name(pattern_name)

        return (
            pattern_name,
            *make_fused_bias_activation_pattern(
                "relax.nn.conv2d",
                with_bias=with_bias,
                activation=activation,
            ),
            _check_conv2d,
        )

    return [
        _conv2d_pattern("cutlass.conv2d"),
        _conv2d_pattern("cutlass.conv2d_bias"),
        _conv2d_pattern("cutlass.conv2d_bias_relu"),
        _conv2d_pattern("cutlass.conv2d_bias_silu"),
    ]


def residual_block_patterns():
    """
    Returns a list of all residual block patterns in cutlass BYOC backend.
    """
    patterns = []

    for activation, name_postfix in [(None, ""), ("relax.nn.relu", "_relu")]:
        for check, base_patterns in [
            (_check_conv2d, conv2d_patterns()),
            (_check_matmul, matmul_patterns()),
            (_check_decode_matmul, decode_matmul_patterns()),
        ]:
            for name, pat, arg_pat, _ in base_patterns:
                # Append residual patterns only to those base patterns with bias add,
                # since conv2d or matmul + residual add without bias is already supported
                # via conv2d or matmul + bias patterns (the residual input is treated as "bias").
                if "bias" in name:
                    for bin_op in ["relax.add", "relax.multiply"]:
                        patterns.append(
                            (
                                name + "_residual_" + bin_op.split(".")[-1] + name_postfix,
                                *make_residual_block_pattern(
                                    (pat, arg_pat), binary_op=bin_op, activation=activation
                                ),
                                check,
                            )
                        )

    return patterns


def _check_stacked_attention(context: PatternCheckContext) -> bool:
    """Check if the given stacked attention workload can be offloaded to CUTLASS."""
    if has_leaking_intermediate_variables(context):
        return False
    if not context.annotated_expr["stacked_qkv"].ty.ndim == 3:
        return False
    if "split" in context.annotated_expr:
        split_op = context.annotated_expr["split"]
        if not split_op.attrs.axis == 2:
            return False
    else:
        get_const_int_list = lambda tup: [int(e.value) for e in tup]
        last_end = 0
        for name in ["query", "key", "value"]:
            assert f"strided_slice_{name}" in context.annotated_expr
            strided_slice_op = context.annotated_expr[f"strided_slice_{name}"]
            axes = get_const_int_list(strided_slice_op.args[1])
            begins = get_const_int_list(strided_slice_op.args[2])
            ends = get_const_int_list(strided_slice_op.args[3])
            strides = get_const_int_list(strided_slice_op.args[4])

            if axes != [2]:
                return False
            if begins != [last_end]:
                return False
            if not len(ends) == 1:
                return False
            if strides != [1]:
                return False
            last_end = ends[0]
    return True


def attention_patterns():
    """
    Returns a list of all attention patterns in cutlass BYOC backend.
    """
    return [
        (
            "cutlass.attention",
            *make_attention_pattern(),
        ),
        (
            "cutlass.attention_bias",
            *make_attention_pattern(with_bias=True),
        ),
        (
            "cutlass.stacked_attention",
            *make_stacked_attention_pattern(start_op="split"),
            _check_stacked_attention,
        ),
        (
            "cutlass.stacked_attention",
            *make_stacked_attention_pattern(start_op="split", with_bias=True),
            _check_stacked_attention,
        ),
        (
            "cutlass.stacked_attention",
            *make_stacked_attention_pattern(start_op="strided_slice"),
            _check_stacked_attention,
        ),
        (
            "cutlass.stacked_attention",
            *make_stacked_attention_pattern(start_op="strided_slice", with_bias=True),
            _check_stacked_attention,
        ),
        (
            "cutlass.attention_var_len",
            *make_attention_pattern(var_len=True),
        ),
    ]


def _check_layer_norm(context: PatternCheckContext) -> bool:
    attrs = context.matched_expr.attrs

    if not attrs.center or not attrs.scale:
        return False

    if len(attrs.axes) != 1:
        # Contiguous inner-most axes can be supported, but reject it for now for simplicity.
        return False

    axis = int(attrs.axes[0])
    rank = len(context.matched_expr.ty.shape)

    if axis < 0:
        axis += rank

    return axis == rank - 1


def layer_norm_pattern():
    """Create a layer norm pattern for CUTLASS."""
    return [
        (
            "cutlass.layer_norm",
            *make_layer_norm_pattern(),
            _check_layer_norm,
        ),
    ]


def _check_rms_norm(ctx: PatternCheckContext) -> bool:
    rms_norm = ctx.annotated_expr["rms_norm"]
    if "rms_norm" not in rms_norm.args[0].name_hint:
        return False

    return True


def rms_norm_pattern():
    """Create a RMS norm pattern for CUTLASS."""
    return [
        (
            "cutlass.rms_norm",
            *make_rms_norm_pattern(),
            _check_rms_norm,
        ),
    ]


def attention_rewrite_patterns():
    """
    Returns a list of all attention rewriting patterns in cutlass BYOC backend.
    """
    patterns = []
    for qkv_layout in ["BSNH", "BSH"]:
        for out_layout in ["BSNH", "BSH"]:
            for with_bias in [True, False]:
                for with_cast in [True, False]:
                    patterns.append(
                        make_attention_rewrite_pattern(qkv_layout, out_layout, with_bias, with_cast)
                    )
    return patterns


register_patterns(
    [
        *conv2d_patterns(),
        *matmul_patterns(),
        *decode_matmul_patterns(),
        *residual_block_patterns(),
        *attention_patterns(),
        *layer_norm_pattern(),
        *rms_norm_pattern(),
    ]
)

_REWRITE_PATTERNS = [*attention_rewrite_patterns()]


@expr_functor.mutator
class WorkspaceAnnotator(PyExprMutator):
    """Annotate a workspace requirement for each CUTLASS-offloaded function."""

    def __init__(self, mod):
        super().__init__(mod)

    def visit_function_(self, f):
        if "Composite" not in f.attrs:
            body = super().visit_expr(f.body)
            new_f = Function(f.params, body, f.ret_ty, f.is_pure, f.attrs, f.span)

            if "global_symbol" in f.attrs and "cutlass" in f.attrs["global_symbol"]:
                composite_func = body.blocks[0].bindings[0].value
                if "WorkspaceSize" in composite_func.attrs:
                    return new_f.with_attr("WorkspaceSize", composite_func.attrs["WorkspaceSize"])

            return new_f

        if "attention" in f.attrs["Composite"] and "cutlass" in f.attrs["Composite"]:
            # Workspace is needed only for larger head sizes, but for simplicity we always allocate.
            out_dtype = f.ret_ty.dtype
            out_size_1d = _shape_1d(f.ret_ty.shape)
            # This needs to be in sync with the actual value that the kernel expects.
            workspace_size_bytes = out_size_1d * {"float16": 2, "float32": 4}[out_dtype]
            if not isinstance(workspace_size_bytes, int | tvm.tirx.expr.IntImm):
                # Tempororay workaround for dynamic shape workload. Will be removed when
                # workspace for dynamic shape workload is implemented.
                workspace_size_bytes = 8
            return f.with_attr("WorkspaceSize", workspace_size_bytes)

        return f


@tvm.transform.module_pass(opt_level=0)
def annotate_workspace(mod, _):
    """Pass to annotate a workspace requirement for each CUTLASS-offloaded function."""
    annotator = WorkspaceAnnotator(mod)
    for name, f in mod.functions_items():
        if isinstance(f, Function):
            new_f = annotator.visit_expr(f)
            mod.update_func(name, new_f)
    return mod


def partition_for_cutlass(mod, annotate_codegen=True, use_flash_mqa=True):
    """
    Partition the input module into CUTLASS-supported subgraphs.

    Parameters
    ----------
    mod: tvm.IRModule
        The IRModule to be partitioned.

    annotate_codegen: bool
        Whether to wrap each created composite function with another function, whose
        body consists only of a call to the composite function. See the doc of FuseOpsByPattern
        for more detail.

    use_flash_mqa: bool
        Whether to consider a rewrite pattern for multi-query attention, which is supported by
        the Flash Attention kernel.

    Returns
    -------
    mod: tvm.IRModule
        The resulting IRModule, containing partitioned subgraphs to be
        compiled by the CUTLASS backend.
    """
    for func_name, func in mod.functions_items():
        if isinstance(func, Function):
            if use_flash_mqa:
                mqa_pattern, rewriter = make_attention_rewrite_pattern(
                    "BSNH", "BSNH", with_bias=False, with_cast=True, with_kv_repeat=True
                )
                func = rewrite_call(mqa_pattern, rewriter, func)

            for pattern, rewriter in _REWRITE_PATTERNS:
                func = rewrite_call(pattern, rewriter, func)

        mod[func_name] = func

    patterns = get_patterns_with_prefix("cutlass")
    return tvm.transform.Sequential(
        [
            transform.FuseOpsByPattern(
                patterns, bind_constants=False, annotate_codegen=annotate_codegen
            ),
            annotate_workspace,
            transform.AllocateWorkspace(),
        ]
    )(mod)