wehub-resource-sync
239 lines
8.5 KiB
Python
239 lines
8.5 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.
|
|
# pylint: disable=invalid-name
|
|
"""This file defines ScheduleState, the core data structure of TensorIR scheduling."""
|
|
|
|
from collections import namedtuple
|
|
from enum import IntEnum
|
|
|
|
from tvm_ffi import register_object
|
|
|
|
from tvm.ir import IRModule
|
|
from tvm.runtime import Object
|
|
from tvm.tirx import For, PrimFunc, SBlock, SBlockRealize
|
|
|
|
from ..sblock_scope import SBlockScope, StmtSRef
|
|
from . import _ffi_api
|
|
|
|
CachedFlags = namedtuple("CachedFlags", ["affine_binding", "region_cover", "stage_pipeline"])
|
|
|
|
|
|
class ScheduleDebugMask(IntEnum):
|
|
"""The bitmask of the `debug_mask` flag in the ScheduleState class.
|
|
|
|
If the `debug_mask` flag has a certain bit on, then the correpsonding
|
|
verification pass will be conducted. For example, if `(debug_mask & VERIFY_SREF_TREE) != 0`,
|
|
then the correctness of the sref tree will be verified after each schedule instruction.
|
|
|
|
Attributes
|
|
----------
|
|
VERIFY_SREF_TREE : int = 1
|
|
Verify the correctness of the sref tree
|
|
VERIFY_CACHED_FLAGS : int = 2
|
|
Verify the correctness of affine_binding, region_cover and stage_pipeline
|
|
"""
|
|
|
|
VERIFY_SREF_TREE = 1
|
|
VERIFY_CACHED_FLAGS = 2
|
|
|
|
|
|
def _parse_mod(mod: PrimFunc | IRModule) -> IRModule:
|
|
if isinstance(mod, PrimFunc):
|
|
mod = IRModule({"main": mod})
|
|
if not isinstance(mod, IRModule):
|
|
raise TypeError(f"Expected `mod` to be PrimFunc or IRModule, but gets: {mod}")
|
|
return mod
|
|
|
|
|
|
def _parse_debug_mask(debug_mask: str | int) -> int:
|
|
if isinstance(debug_mask, str):
|
|
if debug_mask == "all":
|
|
debug_mask = ScheduleDebugMask.VERIFY_SREF_TREE | ScheduleDebugMask.VERIFY_CACHED_FLAGS
|
|
elif debug_mask == "none":
|
|
debug_mask = 0
|
|
else:
|
|
raise ValueError(f"Unrecognizable `debug_mask`: {debug_mask}")
|
|
if not isinstance(debug_mask, bool) and not isinstance(debug_mask, int):
|
|
raise TypeError(f"`debug_mask` should be integer or boolean, but gets: {debug_mask}")
|
|
return debug_mask
|
|
|
|
|
|
def _parse_enable_checks(enable_checks: bool) -> bool:
|
|
if not isinstance(enable_checks, bool):
|
|
raise TypeError(f"enable_checks only accepts bool value, got {type(enable_checks)} instead")
|
|
return enable_checks
|
|
|
|
|
|
@register_object("s_tir.ScheduleState")
|
|
class ScheduleState(Object):
|
|
"""The state of scheduling, which exposes a `Replace` method as
|
|
the primary resort for all the scheduling primitives to manipulate the TensorIR.
|
|
|
|
The data structure contains the following information
|
|
1) The AST being scheduled (mod)
|
|
2) The sref tree of schedulable statements (indicated by the srefs)
|
|
3) The dependency information of each block scope (block_info)
|
|
4) A reverse mapping from the AST nodes to that in the sref tree (get_sref)
|
|
5) A debug flag, if set, extra checking is enabled (debug_mask)
|
|
6) A enable check flag, if False, some prerequisite checks are disabled.
|
|
|
|
Parameters
|
|
----------
|
|
mod : IRModule
|
|
The AST of the module being scheduled
|
|
debug_mask : int
|
|
Do extra correctness checking after the object construction
|
|
and each time after calling the Replace method.
|
|
enable_check : bool
|
|
Indicates whether we enable prerequisite checks for some schedule primitives or not,
|
|
defaults to `True`.
|
|
"""
|
|
|
|
mod: IRModule
|
|
debug_mask: int
|
|
|
|
def __init__(
|
|
self,
|
|
mod: PrimFunc | IRModule,
|
|
*,
|
|
debug_mask: str | int = "none",
|
|
enable_check: bool = True,
|
|
) -> None:
|
|
"""Construct a schedule state from an IRModule or a PrimFunc
|
|
|
|
Parameters
|
|
----------
|
|
mod : Union[PrimFunc, IRModule]
|
|
The IRModule or PrimFunc to be scheduled
|
|
debug_mask : Union[str, int]
|
|
Do extra correctness checking after the class creation and each time
|
|
after calling the Replace method.
|
|
Possible choices of `debug_mask`:
|
|
1) "all" - Turn on all the checks
|
|
2) "none" - Turn off all the checks
|
|
3) An integer - Turn on checks according to the bitmasks provided in ScheduleDebugMask
|
|
"""
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.ScheduleState, # type: ignore # pylint: disable=no-member
|
|
_parse_mod(mod),
|
|
_parse_debug_mask(debug_mask),
|
|
_parse_enable_checks(enable_check),
|
|
)
|
|
|
|
def get_sref(self, stmt: SBlock | For) -> StmtSRef | None:
|
|
"""Return the corresponding sref that points to the stmt
|
|
|
|
Parameters
|
|
----------
|
|
stmt : Union[Block, For]
|
|
The schedulable statement in the TensorIR to be retrieved for its sref
|
|
|
|
Returns
|
|
-------
|
|
sref : StmtSRef
|
|
The corresponding sref
|
|
"""
|
|
return _ffi_api.ScheduleStateGetSRef(self, stmt) # type: ignore # pylint: disable=no-member
|
|
|
|
def get_sblock_scope(self, block_sref: StmtSRef) -> SBlockScope:
|
|
"""Get the SBlockScope correpsonding to the block sref
|
|
|
|
Parameters
|
|
----------
|
|
block_sref : StmtSRef
|
|
The block sref to be retrieved
|
|
|
|
Returns
|
|
-------
|
|
sref : StmtSRef
|
|
The corresponding sref
|
|
"""
|
|
return _ffi_api.ScheduleStateGetSBlockScope( # type: ignore # pylint: disable=no-member
|
|
self, block_sref
|
|
)
|
|
|
|
def _get_cached_flags(self, block_sref: StmtSRef) -> CachedFlags:
|
|
"""Get the cached flags of the corresponding block
|
|
|
|
Parameters
|
|
----------
|
|
block_sref : StmtSRef
|
|
The block sref to be retrieved
|
|
|
|
Returns
|
|
-------
|
|
flags : CachedFlags
|
|
Three flags: affine_binding, region_cover, stage_pipeline
|
|
|
|
Note
|
|
----
|
|
It is an API intended for internal testing use.
|
|
"""
|
|
(
|
|
affine_binding,
|
|
region_cover,
|
|
stage_pipeline,
|
|
) = _ffi_api.ScheduleStateGetCachedFlags( # type: ignore # pylint: disable=no-member
|
|
self, block_sref
|
|
)
|
|
return CachedFlags(
|
|
affine_binding=bool(affine_binding.value),
|
|
region_cover=bool(region_cover.value),
|
|
stage_pipeline=bool(stage_pipeline.value),
|
|
)
|
|
|
|
def replace(
|
|
self,
|
|
src_sref: StmtSRef,
|
|
tgt_stmt: SBlock | For | SBlockRealize,
|
|
block_sref_reuse: dict[SBlock, SBlock] | None = None,
|
|
) -> None:
|
|
"""
|
|
Replace the part of the AST, as being pointed to by `src_sref`,
|
|
with a specific statement `tgt_stmt`, and maintain the sref tree accordingly.
|
|
Replace will try to perform copy on write as much as possible when the ScheduleState holds
|
|
the only copy to the IRModule and IR nodes.
|
|
|
|
Only 3 types of replacements are allowed: from `src_sref->stmt` to `tgt_stmt`.
|
|
1) SBlock -> SBlock
|
|
2) Loop -> Loop
|
|
3) Loop -> BlockRealize
|
|
|
|
Parameters
|
|
----------
|
|
src_sref : StmtSRef
|
|
The sref to the statement to be replaced in the TensorIR AST
|
|
|
|
tgt_stmt : Union[Block, For, BlockRealize]
|
|
The statement to be replaced to
|
|
|
|
block_sref_reuse : Optional[Dict[Block, Block]] = None
|
|
Maps an old block (to be replaced in the subtree under `src_sref->stmt`)
|
|
to a new block (replaced to, in the subtree under `tgt_stmt`), and enforces
|
|
reuse of srefs between them (rather than create new srefs) i.e. after being replaced,
|
|
the sref that points to the old block will point to the new one
|
|
|
|
Note
|
|
----
|
|
The reuse of loop srefs are detected automatically according to the reuse of loop vars.
|
|
"""
|
|
if block_sref_reuse is None:
|
|
block_sref_reuse = {}
|
|
_ffi_api.ScheduleStateReplace( # type: ignore # pylint: disable=no-member
|
|
self, src_sref, tgt_stmt, block_sref_reuse
|
|
)
|