wehub-resource-sync
158 lines
5.0 KiB
Python
158 lines
5.0 KiB
Python
# 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
|