chore: import upstream snapshot with attribution

python/tvm/s_tir/sblock_scope.py

@@ -0,0 +1,157 @@
+# 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.
+"""Definition of two pillar data structure for TensorIR scheduling: StmtSRef, SBlockScope."""
+
+from enum import IntEnum
+from typing import Optional
+
+from tvm_ffi import register_object
+
+from tvm.runtime import Object
+from tvm.tirx import For, SBlock
+
+from . import _ffi_api
+
+
+@register_object("s_tir.StmtSRef")
+class StmtSRef(Object):
+    """An object that refers to schedulable elements in the TensorIR, aka "sref".
+
+    Glossary
+    - SBlock sref: An StmtSref that points to a TensorIR block.
+    - Loop sref: An StmtSRef that points to a TensorIR for loop.
+    - Parent sref: The parent sref of an sref is the block/loop sref that points to its closest
+    schedulable statement of its ancestors on the TensorIR AST.
+    - Root sref: Sref to the root block. Every sref has exactly one parent sref
+    except for root sref.
+    - Sref tree: The parent-children-relationship of srefs that forms a tree,
+    uniquely determined by the TensorIR AST.
+    """
+
+    seq_index: int
+
+    @property
+    def stmt(self) -> SBlock | For | None:
+        """The block/for stmt the object refers to"""
+        return _ffi_api.StmtSRefStmt(self)  # type: ignore # pylint: disable=no-member
+
+    @property
+    def parent(self) -> Optional["StmtSRef"]:
+        """The parent sref"""
+        return _ffi_api.StmtSRefParent(self)  # type: ignore # pylint: disable=no-member
+
+    @staticmethod
+    def inline_mark() -> "StmtSRef":
+        """A special StmtSRef, which doesn't point to any stmt in the AST,
+        only serving as a "mark" to hint compute-at to do the work of compute-inline"""
+        return _ffi_api.StmtSRefInlineMark()  # type: ignore # pylint: disable=no-member
+
+    @staticmethod
+    def root_mark() -> "StmtSRef":
+        """A special StmtSRef, which doesn't point to any stmt in the AST,
+        only serving as a "mark" to hint compute-at to do nothing"""
+        return _ffi_api.StmtSRefRootMark()  # type: ignore # pylint: disable=no-member
+
+
+class DepKind(IntEnum):
+    """Type of dependency.
+
+    Attributes
+    ----------
+    RAW : int = 0
+        Read-after-write dependency
+    WAW : int = 1
+        Write-after-write dependency
+    WAR : int = 2
+        Write-after-read dependency. Not supported in TensorIR for now.
+    OPAQUE: int = 3
+        Opaque dependency
+    """
+
+    RAW = 0
+    WAW = 1
+    WAR = 2
+    OPAQUE = 3
+
+
+@register_object("s_tir.Dependency")
+class Dependency(Object):
+    """A tuple (src, dst, kind) representing certain types of dependency.
+    For example, (A, B, kRAW) means block B depends on block A, and the dependency kind is
+    read-after-write, which means block B reads the result written by block A.
+
+    Parameters
+    ----------
+    src : StmtSRef
+        The source of the dependency relation
+    dst : StmtSRef
+        The destination of the dependency relation
+    kind : DepKind
+        The dependency kind
+    """
+
+    src: StmtSRef
+    dst: StmtSRef
+    kind: DepKind
+
+
+@register_object("s_tir.SBlockScope")
+class SBlockScope(Object):
+    """An object corresponds to each block sref in the sref tree, which
+    tracks the producer-consumer dependency between blocks.
+
+    Glossary:
+
+    - SBlock scope: A contiguous subtree of the sref tree, rooted at
+      each SBlock sref, whose components are:
+
+      - scope root: a SBlock sref
+      - internal srefs: loop srefs
+      - scope leaves: SBlock srefs
+
+    - Child SBlock: The scope leaf SBlocks under the scope root or a specific internal sref
+    """
+
+    def get_deps_by_src(self, block: StmtSRef) -> list[Dependency]:
+        """Get all dependencies whose `src` is the target`block`.
+
+        Parameters
+        ----------
+        block: StmtSRef
+            The queried block
+
+        Returns
+        -------
+        blocks: List[Dependency]
+            The dependencies
+        """
+        return _ffi_api.SBlockScopeGetDepsBySrc(self, block)  # type: ignore # pylint: disable=no-member
+
+    def get_deps_by_dst(self, block: StmtSRef) -> list[Dependency]:
+        """Get all dependencies whose `dst` is the target `block`.
+
+        Parameters
+        ----------
+        block: StmtSRef
+            The queried block
+
+        Returns
+        -------
+        blocks: List[Dependency]
+            The dependencies
+        """
+        return _ffi_api.SBlockScopeGetDepsByDst(self, block)  # type: ignore # pylint: disable=no-member