chore: import upstream snapshot with attribution

2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
@@ -0,0 +1,140 @@
+# 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 and partitioning for the TensorRT BYOC backend.
+
+The composite name of each pattern is "tensorrt.<op>", matching the runtime
+converter registered under the same name (the converters are keyed by
+"tensorrt." + op_name). ``partition_for_tensorrt`` carves the matched subgraphs
+out of the module and annotates them for the ``tensorrt`` codegen.
+"""
+
+from collections.abc import Mapping
+
+from tvm.ir import IRModule
+from tvm.relax.dpl.pattern import DFPattern, is_op, wildcard
+from tvm.relax.transform import FuseOpsByPattern, MergeCompositeFunctions
+
+from ..pattern_registry import get_patterns_with_prefix, register_patterns
+
+Pattern = tuple[str, DFPattern, Mapping[str, DFPattern]]
+
+
+def _op_pattern(composite_name: str, op_name: str, num_args: int) -> Pattern:
+    """A pattern matching a single op called with ``num_args`` wildcard arguments."""
+    args = [wildcard() for _ in range(num_args)]
+    return (composite_name, is_op(op_name)(*args), {})
+
+
+def _tensorrt_patterns() -> list[Pattern]:
+    patterns: list[Pattern] = []
+
+    # Activations and unary elementwise ops (single tensor argument).
+    for composite, op in [
+        ("tensorrt.nn.relu", "relax.nn.relu"),
+        ("tensorrt.sigmoid", "relax.sigmoid"),
+        ("tensorrt.tanh", "relax.tanh"),
+        ("tensorrt.exp", "relax.exp"),
+        ("tensorrt.log", "relax.log"),
+        ("tensorrt.sqrt", "relax.sqrt"),
+        ("tensorrt.abs", "relax.abs"),
+        ("tensorrt.negative", "relax.negative"),
+        ("tensorrt.sin", "relax.sin"),
+        ("tensorrt.cos", "relax.cos"),
+        ("tensorrt.atan", "relax.atan"),
+        ("tensorrt.ceil", "relax.ceil"),
+        ("tensorrt.floor", "relax.floor"),
+        ("tensorrt.erf", "relax.erf"),
+        ("tensorrt.nn.softmax", "relax.nn.softmax"),
+        ("tensorrt.nn.batch_flatten", "relax.nn.batch_flatten"),
+        ("tensorrt.expand_dims", "relax.expand_dims"),
+        ("tensorrt.squeeze", "relax.squeeze"),
+        ("tensorrt.transpose", "relax.permute_dims"),
+        ("tensorrt.layout_transform", "relax.layout_transform"),
+        ("tensorrt.nn.max_pool2d", "relax.nn.max_pool2d"),
+        ("tensorrt.nn.avg_pool2d", "relax.nn.avg_pool2d"),
+        ("tensorrt.nn.max_pool3d", "relax.nn.max_pool3d"),
+        ("tensorrt.nn.avg_pool3d", "relax.nn.avg_pool3d"),
+        ("tensorrt.nn.adaptive_avg_pool2d", "relax.nn.adaptive_avg_pool2d"),
+        ("tensorrt.sum", "relax.sum"),
+        ("tensorrt.prod", "relax.prod"),
+        ("tensorrt.max", "relax.max"),
+        ("tensorrt.min", "relax.min"),
+        ("tensorrt.mean", "relax.mean"),
+        ("tensorrt.concatenate", "relax.concat"),
+        ("tensorrt.split", "relax.split"),
+    ]:
+        patterns.append(_op_pattern(composite, op, 1))
+
+    # Binary elementwise ops (two tensor arguments).
+    for composite, op in [
+        ("tensorrt.add", "relax.add"),
+        ("tensorrt.subtract", "relax.subtract"),
+        ("tensorrt.multiply", "relax.multiply"),
+        ("tensorrt.divide", "relax.divide"),
+        ("tensorrt.power", "relax.power"),
+        ("tensorrt.maximum", "relax.maximum"),
+        ("tensorrt.minimum", "relax.minimum"),
+    ]:
+        patterns.append(_op_pattern(composite, op, 2))
+
+    # Convolutions and matmul (data + weight).
+    for composite, op in [
+        ("tensorrt.nn.conv1d", "relax.nn.conv1d"),
+        ("tensorrt.nn.conv2d", "relax.nn.conv2d"),
+        ("tensorrt.nn.conv3d", "relax.nn.conv3d"),
+        ("tensorrt.nn.conv2d_transpose", "relax.nn.conv2d_transpose"),
+        ("tensorrt.nn.conv3d_transpose", "relax.nn.conv3d_transpose"),
+        ("tensorrt.nn.batch_matmul", "relax.matmul"),
+        ("tensorrt.reshape", "relax.reshape"),
+    ]:
+        patterns.append(_op_pattern(composite, op, 2))
+
+    # layer_norm (data, gamma, beta) and clip (data, min, max).
+    patterns.append(_op_pattern("tensorrt.nn.layer_norm", "relax.nn.layer_norm", 3))
+    patterns.append(_op_pattern("tensorrt.clip", "relax.clip", 3))
+
+    # strided_slice is called either with or without the optional strides argument.
+    patterns.append(_op_pattern("tensorrt.strided_slice", "relax.strided_slice", 5))
+    patterns.append(_op_pattern("tensorrt.strided_slice", "relax.strided_slice", 4))
+
+    return patterns
+
+
+register_patterns(_tensorrt_patterns())
+
+
+def partition_for_tensorrt(mod: IRModule) -> IRModule:
+    """Partition the module, offloading TensorRT-supported subgraphs.
+
+    Parameters
+    ----------
+    mod : tvm.ir.IRModule
+        The module to partition. Bind model parameters (e.g. via
+        ``relax.transform.BindParams``) before calling this so that weights are
+        available to TensorRT as constants.
+
+    Returns
+    -------
+    mod : tvm.ir.IRModule
+        The module with TensorRT-supported subgraphs grouped into composite
+        functions annotated for the ``tensorrt`` codegen.
+    """
+    patterns = get_patterns_with_prefix("tensorrt")
+    mod = FuseOpsByPattern(patterns, bind_constants=True, annotate_codegen=False)(mod)
+    mod = MergeCompositeFunctions()(mod)
+    return mod