chore: import upstream snapshot with attribution

python/tvm/relax/backend/cuda/cudnn.py

@@ -0,0 +1,202 @@
+# 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.
+
+"""Pattern table for cuDNN backend"""
+
+import operator
+from functools import partial, reduce
+
+import tvm
+from tvm import relax
+from tvm.relax import PyExprMutator, expr_functor, transform
+from tvm.relax.transform import PatternCheckContext
+
+from ..pattern_registry import get_patterns_with_prefix, register_patterns
+from ..patterns import make_conv2d_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 cuDNN BYOC."""
+    return (lhs_dtype == "float16" and rhs_dtype == "float16") or (
+        lhs_dtype == "float32" and rhs_dtype == "float32"
+    )
+
+
+def _is_supported_format(data_layout, kernel_layout):
+    """Check if layouts in the given workload are supported by cuDNN BYOC."""
+    return (data_layout == "NHWC" and kernel_layout == "OHWI") or (
+        data_layout == "NCHW" and kernel_layout == "OIHW"
+    )
+
+
+def _check_conv2d(context: PatternCheckContext) -> bool:
+    if has_leaking_intermediate_variables(context):
+        return False
+    # Retrieve the annotated expression from context
+    conv2d_call = context.annotated_expr["root"]
+    input_expr = context.annotated_expr["input"]
+    weight_expr = context.annotated_expr["weight"]
+
+    # Check if the data types of input and weights are supported by cuDNN BYOC
+    input_dtype = input_expr.ty.dtype
+    weight_dtype = weight_expr.ty.dtype
+    if not _is_supported_dtype(input_dtype, weight_dtype):
+        return False
+
+    input_layout = conv2d_call.attrs.data_layout
+    weight_layout = conv2d_call.attrs.kernel_layout
+    if not _is_supported_format(input_layout, weight_layout):
+        return False
+
+    return True
+
+
+def _check_stacked_attention(context: PatternCheckContext, layout: str) -> bool:
+    """Check if the given stacked attention workload can be offloaded to cuDNN."""
+    if has_leaking_intermediate_variables(context):
+        return False
+    if layout == "BS3NH":
+        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
+    elif layout == "SBN3H":
+        if not context.annotated_expr["stacked_qkv"].ty.ndim == 4:
+            return False
+        if "split" in context.annotated_expr:
+            split_op = context.annotated_expr["split"]
+            if not split_op.attrs.axis == 3:
+                return False
+    else:
+        raise NotImplementedError(f"Unsupported layout: {layout}")
+    return True
+
+
+register_patterns(
+    [
+        (
+            "cudnn.conv2d.nhwc_ohwi",
+            *make_conv2d_pattern(
+                with_bias=False,
+            ),
+            _check_conv2d,
+        ),
+        (
+            "cudnn.conv2d.nhwc_ohwi_bias",
+            *make_conv2d_pattern(
+                with_bias=True,
+            ),
+            _check_conv2d,
+        ),
+        (
+            "cudnn.conv2d.nhwc_ohwi_bias_relu",
+            *make_conv2d_pattern(
+                with_bias=True,
+                activation="relax.nn.relu",
+            ),
+            _check_conv2d,
+        ),
+        (
+            "cudnn.attention.BS3NH",
+            *make_stacked_attention_pattern(start_op="split", layout="BS3NH"),
+            partial(_check_stacked_attention, layout="BS3NH"),
+        ),
+        (
+            "cudnn.attention.SBN3H",
+            *make_stacked_attention_pattern(start_op="split", layout="SBN3H"),
+            partial(_check_stacked_attention, layout="SBN3H"),
+        ),
+    ]
+)
+
+
+def partition_for_cudnn(mod):
+    """
+    Partition the input module into cuDNN-supported subgraphs.
+
+    Parameters
+    ----------
+    mod: tvm.IRModule
+        The IRModule to be partitioned.
+
+    Returns
+    -------
+    mod: tvm.IRModule
+        The resulting IRModule, containing partitioned subgraphs to be
+        offloaded to the cuDNN backend.
+    """
+
+    patterns = get_patterns_with_prefix("cudnn")
+    return tvm.transform.Sequential(
+        [
+            transform.FuseOpsByPattern(patterns, bind_constants=False, annotate_codegen=True),
+            annotate_workspace,
+            transform.AllocateWorkspace(),
+        ]
+    )(mod)
+
+
+def _shape_1d(shape):
+    return reduce(operator.mul, shape, 1)
+
+
+@expr_functor.mutator
+class WorkspaceAnnotator(PyExprMutator):
+    """Annotate a workspace requirement for each cuDNN-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 = relax.Function(f.params, body, f.ret_ty, f.is_pure, f.attrs, f.span)
+
+            if "global_symbol" in f.attrs and "cudnn" 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 "cudnn" 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 cuDNN-offloaded function."""
+    annotator = WorkspaceAnnotator(mod)
+    for name, f in mod.functions_items():
+        if isinstance(f, relax.Function):
+            new_f = annotator.visit_expr(f)
+            mod.update_func(name, new_f)
+    return mod