chore: import upstream snapshot with attribution

python/tvm/relax/dpl/__init__.py

@@ -0,0 +1,29 @@
+# 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.
+
+"""The Relax Dataflow Pattern Language."""
+
+from .pattern import *
+from .context import *
+from .rewrite import (
+    rewrite_call,
+    rewrite_bindings,
+    PatternMatchingRewriter,
+    ExprPatternRewriter,
+    OrRewriter,
+)

python/tvm/relax/dpl/_ffi.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.
+"""DataFlow Pattern Language FFI bindings."""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("relax.dpl", __name__)

python/tvm/relax/dpl/context.py

@@ -0,0 +1,79 @@
+# 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.
+
+"""The Graph Matching Context Manager for Dataflow Pattern Language."""
+
+import tvm
+
+from ..expr import DataflowBlock, Var
+from . import _ffi as ffi
+from .pattern import DFPattern
+
+
+class PatternContext(tvm.runtime.Object):
+    """A context object for doing graph (topogical) pattern matching."""
+
+    def __init__(self, incremental=False):
+        """
+        Initialize the PatternContext
+
+        Parameters
+        ----------
+        incremental : bool, optional
+            perform incremental matching based on the recent context, by default False
+        """
+        self.__init_handle_by_constructor__(ffi.PatternContext, incremental)  # type: ignore
+
+    def __enter__(self):
+        """Enter the context"""
+        ffi.enter_context(self)  # type: ignore
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        """Exit the context"""
+        ffi.exit_context(self)  # type: ignore
+
+    @staticmethod
+    def current() -> "PatternContext":
+        """
+        Get the current context
+
+        Returns
+        -------
+        PatternContext
+            The current context
+        """
+        return ffi.current_context()  # type: ignore
+
+    def match_dfb(
+        self,
+        dfb: DataflowBlock,
+    ) -> dict[DFPattern, Var]:
+        """
+        Match a DataflowBlock via a graph of DFPattern and corresponding constraints
+
+        Parameters
+        ----------
+        dfb : DataflowBlock
+            The DataflowBlock to match
+
+        Returns
+        -------
+        Dict[DFPattern, Var]
+            The mapping from DFPattern to matched expression
+        """
+        return ffi.match_dfb(self, dfb)  # type: ignore

python/tvm/relax/dpl/rewrite.py

@@ -0,0 +1,299 @@
+# 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.
+"""APIs for pattern-based rewriting."""
+
+from collections.abc import Callable
+
+from tvm_ffi import register_object
+
+from tvm.ir import IRModule
+from tvm.runtime import Object
+
+from ..expr import Expr, Function, Var
+from . import _ffi as ffi
+from .context import PatternContext
+from .pattern import DFPattern
+
+
+@register_object("relax.dpl.PatternMatchingRewriter")
+class PatternMatchingRewriter(Object):
+    """A pattern-matching rewriter for Relax"""
+
+    @staticmethod
+    def from_pattern(
+        pattern: DFPattern,
+        func: Callable[[Expr, dict[DFPattern, Expr]], Expr],
+    ) -> "PatternMatchingRewriter":
+        """Construct from a pattern and rewriter-function
+
+        The replacements performed by the rewriter will be equivalent
+        to using the `pattern` and `func` as arguments to
+        `rewrite_call`.
+
+        Parameters
+        ----------
+        pattern: DFPattern
+
+            The pattern to be matched against.
+
+        func: Callable[[Expr, Dict[DFPattern, Expr]], Expr]
+
+            A function that returns the rewritten expression.  See
+            `rewrite_call` for details and examples.
+
+
+        Returns
+        -------
+        rewriter_obj: PatternMatchingRewriter
+
+            The rewriter object
+
+        """
+        return ffi.PatternMatchingRewriterFromPattern(
+            pattern,
+            func,
+        )  # type: ignore
+
+    @staticmethod
+    def from_module(mod: IRModule) -> "PatternMatchingRewriter":
+        """Construct a rewriter from an IRModule
+
+        The IRModule must have two publicly-exposed functions,
+        `pattern` and `replacement`, where `pattern` and `replacement`
+        have the same function signature, as shown in the example
+        below.
+
+        .. code-block:: python
+
+            @I.ir_module
+            class RewriteAddIntoMultiply:
+                @R.function
+                def pattern(A: R.Tensor):
+                    B = A + A
+                    return B
+
+                @R.function
+                def replacement(A: R.Tensor):
+                    B = A * 2
+                    return B
+
+            rewriter = PatternMatchingRewriter.from_module(RewriteAddIntoMultiply)
+            rewritten_ir_module = rewriter(ir_module)
+
+        To support the common case of defining an IRModule with
+        TVMScript, then immediately turning it into a rewriter, the
+        `@R.rewriter` annotation can be used.
+
+        .. code-block:: python
+
+            @R.rewriter
+            class RewriteAddIntoMultiply:
+                @R.function
+                def pattern(A: R.Tensor):
+                    B = A + A
+                    return B
+
+                @R.function
+                def replacement(A: R.Tensor):
+                    B = A * 2
+                    return B
+
+            rewritten_ir_module = RewriteAddIntoMultiply(ir_module)
+
+        Parameters
+        ----------
+        mod: IRModule
+
+            A module with `pattern` and `replacement` functions,
+            defining a rewrite rule.
+
+
+        Returns
+        -------
+        rewriter_obj: PatternMatchingRewriter
+
+            The rewriter object
+
+        """
+        return ffi.PatternMatchingRewriterFromModule(mod)  # type: ignore
+
+    def __call__(self, obj: Expr | IRModule) -> Expr | IRModule:
+        """Apply the rewriter
+
+        Parameters
+        ----------
+        obj: Union[Expr, IRModule])
+
+            The object to be rewritten.  May be applied to either a
+            relax expression, or an IRModule.
+
+        Returns
+        -------
+        updated: Union[Expr, IRModule]
+
+            The rewritten object
+
+        """
+        return ffi.PatternMatchingRewriterApply(self, obj)
+
+    def __or__(self, other: "PatternMatchingRewriter") -> "PatternMatchingRewriter":
+        """Compose two rewriters
+
+        Composing two rewrite rules together allows them to be applied
+        in a single Relax-level transformation.
+
+        Parameters
+        ----------
+        other: PatternMatchingRewriter
+
+            Another rewrite rule
+
+        Returns
+        -------
+        PatternMatchingRewriter
+
+            A rewriter that will apply either rewrite pattern
+
+        """
+        return OrRewriter(self, other)
+
+
+@register_object("relax.dpl.ExprPatternRewriter")
+class ExprPatternRewriter(PatternMatchingRewriter):
+    def __init__(self, pattern, func):
+        self.__init_handle_by_constructor__(
+            ffi.PatternRewriter,
+            pattern,
+            func,
+        )  # type: ignore
+
+
+@register_object("relax.dpl.OrRewriter")
+class OrRewriter(PatternMatchingRewriter):
+    def __init__(self, lhs, rhs):
+        self.__init_handle_by_constructor__(
+            ffi.OrRewriter,
+            lhs,
+            rhs,
+        )  # type: ignore
+
+
+@register_object("relax.dpl.TupleRewriter")
+class TupleRewriter(PatternMatchingRewriter):
+    def __init__(self, patterns, func):
+        self.__init_handle_by_constructor__(
+            ffi.TupleRewriter,
+            patterns,
+            func,
+        )  # type: ignore
+
+
+def rewrite_call(
+    pattern: DFPattern,
+    rewriter: Callable[[Expr, dict[DFPattern, Expr]], Expr],
+    func: Function,
+) -> Function:
+    """
+    Rewrite a function with the given pattern and the rewriter function.
+
+    Parameters
+    ----------
+    pattern: DFPattern
+        The pattern to match.
+
+    rewriter: Callable[[Expr, Dict[DFPattern, Expr]], Expr]
+        The function to be called on a successful matching for rewriting. Given the matched
+        call node and the map of patterns and matched expressions, it should return a new call node
+        to replace the original one or the original matched call node as is.
+
+        For example, to replace x + x with 2 * x, we can write the rewriter as follows:
+        ```
+        x = wildcard()
+        pattern = is_op("relax.add")(x, x)
+
+        def rewriter(orig, matchings):
+            return R.multiply(matchings[x], R.const(2, "float32"))
+        ```
+
+    func: Function
+        The function to rewrite.
+
+    Returns
+    -------
+    rewritten_func: Function
+        The rewritten or the input function, depending on the pattern matching result.
+    """
+    return ffi.rewrite_call(pattern, rewriter, func)
+
+
+def rewrite_bindings(
+    ctx: PatternContext,
+    rewriter: Callable[[dict[DFPattern, Var], dict[Var, Expr]], dict[Var, Expr]],
+    func: Function,
+) -> Function:
+    """
+    Rewrite a function with the given pattern and the rewriter function.
+
+    Parameters
+    ----------
+    ctx: PatternContext
+        The pattern constraint context under which rewriting takes place.
+
+    rewriter: Callable[[Dict[DFPattern, Var], Dict[Var, Expr]], Dict[Var, Expr]]
+        The function to be called on a successful matching for rewriting. Given the map of patterns
+        and corresponding variables (bound variables or parameters), it should return a map that
+        specifies new values for matched bound variables. It can refer to the passed bindings to
+        create the replacement expressions.
+
+        For example, to rewrite three matmuls for QKV projection in transformer models into one
+        matmul followed by slicing, one can use the follwoing rewriter:
+        ```
+        inp_pat = wildcard()
+        Q_weight_pat, K_weight_pat, V_weight_pat = wildcard(), wildcard(), wildcard()
+
+        matmul1 = is_op("relax.matmul")(inp_pat, Q_weight_pat)
+        matmul2 = is_op("relax.matmul")(inp_pat, K_weight_pat)
+        matmul3 = is_op("relax.matmul")(inp_pat, V_weight_pat)
+
+        def rewriter(matchings):
+            inp = matchings[inp_pat]
+            Q_weight = matchings[Q_weight_pat]
+            K_weight = matchings[K_weight_pat]
+            V_weight = matchings[V_weight_pat]
+            width = Q_weight.ty.shape[1]
+
+            concat = R.concat([Q_weight, K_weight, V_weight], axis=1)
+            matmul = R.matmul(inp, concat)
+            Q = R.strided_slice(matmul, axes=[2], begin=[0], end=[width])
+            K = R.strided_slice(matmul, axes=[2], begin=[width], end=[width * 2])
+            V = R.strided_slice(matmul, axes=[2], begin=[width * 2], end=[width * 3])
+
+            # matchings[matmul1] gives the bound variable in the binding whose RHS matches with
+            # the matmul1 pattern. For example, lv0 in lv0 = R.matmul(x1, w0).
+            # We want to replace the RHS of this binding with Q.
+            return {matchings[matmul1]: Q, matchings[matmul2]: K, matchings[matmul3]: V}
+        ```
+
+    func: Function
+        The function to rewrite.
+
+    Returns
+    -------
+    rewritten_func: Function
+        The rewritten or the input function, depending on the pattern matching result.
+    """
+    return ffi.rewrite_bindings(ctx, rewriter, func)