chore: import upstream snapshot with attribution

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