chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,216 @@
|
||||
# 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.
|
||||
"""Analysis utilities for Schedulable TensorIR (S-TIR)."""
|
||||
|
||||
# pylint: disable=invalid-name
|
||||
from typing import Optional, Union
|
||||
|
||||
import tvm
|
||||
from tvm.ir import IRModule
|
||||
from tvm.tirx.expr import Var
|
||||
from tvm.tirx.stmt import SBlock, BufferRegion
|
||||
|
||||
from tvm.tirx import Buffer, Stmt
|
||||
from tvm.tirx.function import PrimFunc
|
||||
from . import _ffi_api
|
||||
|
||||
|
||||
def get_sblock_access_region(
|
||||
block: SBlock, buffer_var_map: dict[Var, Buffer]
|
||||
) -> list[list[BufferRegion]]:
|
||||
"""Detect which regions of tensors in this block are read or written to.
|
||||
Regions are sorted by order of appearance in the AST.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
block: tvm.tirx.SBlock
|
||||
The block in which we are detecting read/write regions.
|
||||
|
||||
buffer_var_map : Dict[Var, Buffer]
|
||||
The outside buffers which may access the block. Mapping from buffer var to the buffer
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : List[List[BufferRegion]]
|
||||
Array of access regions. There are three arrays of BufferRegion:
|
||||
- first: read regions
|
||||
- second: write regions
|
||||
- third: opaque regions
|
||||
"""
|
||||
return _ffi_api.GetSBlockAccessRegion(block, buffer_var_map) # type: ignore
|
||||
|
||||
|
||||
def get_sblock_read_write_region(
|
||||
block: SBlock, buffer_var_map: dict[Var, Buffer]
|
||||
) -> list[list[BufferRegion]]:
|
||||
"""Auto detect the block read/write region according to its body stmt.
|
||||
An opaque access will be counted as both a read and a write access
|
||||
|
||||
Parameters
|
||||
----------
|
||||
block: tvm.tirx.SBlock
|
||||
The block in which we are detecting read/write regions.
|
||||
|
||||
buffer_var_map : Dict[Var, Buffer]
|
||||
The outside buffers which may access the block. Mapping from buffer var to the buffer
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : List[List[BufferRegion]]
|
||||
An array only consisting of the read regions and write regions of the input block
|
||||
"""
|
||||
return _ffi_api.GetSBlockReadWriteRegion(block, buffer_var_map) # type: ignore
|
||||
|
||||
|
||||
def detect_buffer_access_lca(func: PrimFunc) -> dict[Buffer, Stmt]:
|
||||
"""Detect the lowest common ancestor(LCA) of buffer access, including both high-level
|
||||
access (BufferLoad, BufferStore) and low-level access (BufferLoad, BufferStore and opaque
|
||||
access).
|
||||
The LCA may be a For loop or a Block.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
func: tvm.tirx.PrimFunc
|
||||
The function to be detected.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : Dict[Buffer, Stmt]
|
||||
Map from buffer to the LCA of all access to it.
|
||||
"""
|
||||
return _ffi_api.detect_buffer_access_lca(func) # type: ignore # pylint: disable=no-member
|
||||
|
||||
|
||||
def find_anchor_sblock(mod: IRModule) -> SBlock | None:
|
||||
"""Find the "anchor block" of the given module.
|
||||
|
||||
We define the anchor block to be the block with (1) an init statement and (2) having
|
||||
the biggest flops count. The latter condition is only used when there are multiple blocks
|
||||
with an init statement.
|
||||
|
||||
For example, if the input module is conv2d + fused spatial blocks, conv2d is the anchor block.
|
||||
The input module may not contain more than one such block. For example, a module having
|
||||
two conv2d is not allowed as an input.
|
||||
|
||||
However, a module created from winograd convolution has multiple blocks with an init statement
|
||||
(input transform, batched GEMM, and output transform). We use the second condition, the flops
|
||||
count, to determine that the batched GEMM block is the anchor block.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
mod: tvm.ir.IRModule
|
||||
The input TIR module.
|
||||
Returns
|
||||
-------
|
||||
anchor_block: Optional[SBlock]
|
||||
The anchor block if found, None otherwise.
|
||||
"""
|
||||
return _ffi_api.find_anchor_sblock(mod) # type: ignore # pylint: disable=no-member
|
||||
|
||||
|
||||
def verify_gpu_code(func: PrimFunc, constraints: dict[str, int]) -> bool:
|
||||
"""Verify if module contains illegal host side direct memory access.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
func: tvm.tirx.PrimFunc
|
||||
The module to be verified.
|
||||
|
||||
constraints : Dict[str, int]
|
||||
The attribute constraints.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : bool
|
||||
The result of verification.
|
||||
"""
|
||||
return _ffi_api.verify_gpu_code(func, constraints) # type: ignore
|
||||
|
||||
|
||||
def calculate_allocated_bytes(
|
||||
func_or_mod: PrimFunc | IRModule,
|
||||
) -> dict[str, dict[str, int]]:
|
||||
"""Calculate allocated memory per memory scope required by TIR PrimFuncs.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
func_or_mod: Union[PrimFunc, IRModule]
|
||||
The function or module to be detected. If a module is passed, allocated
|
||||
memory is calculated for all PrimFuncs inside the module
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : Dict[str, Dict[str, int]]
|
||||
Allocated memory size per scope in bytes for each function in the IRModule returned as a
|
||||
dict with function names as keys and a dict of allocated sizes as values. If a single
|
||||
PrimFunc is passed, the function name is returned as "main"
|
||||
"""
|
||||
if not isinstance(func_or_mod, PrimFunc | IRModule):
|
||||
raise TypeError(
|
||||
f"Expected argument to be PrimFunc or IRModule, but received {type(func_or_mod)}"
|
||||
)
|
||||
return _ffi_api.calculate_allocated_bytes(func_or_mod) # type: ignore
|
||||
|
||||
|
||||
def estimate_tir_flops(stmt_or_mod: Stmt | IRModule) -> float:
|
||||
"""Estimate the FLOPs of a TIR fragment.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
stmt_or_mod: Union[Stmt, IRModule]
|
||||
The TIR fragment or IRModule to be estimated.
|
||||
|
||||
Returns
|
||||
-------
|
||||
flops: float
|
||||
The estimated FLOPs.
|
||||
"""
|
||||
return _ffi_api.EstimateTIRFlops(stmt_or_mod) # type: ignore # pylint: disable=no-member
|
||||
|
||||
|
||||
def OOBChecker():
|
||||
"""Detect out of bounds memory access in arrays.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.OOBChecker() # type: ignore
|
||||
|
||||
|
||||
def get_vtcm_compaction_passes() -> list[tvm.transform.Pass]:
|
||||
"""Utility function to get the list of lowering passes to be applied to calculate the compacted
|
||||
VTCM allocation size
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : List[tvm.transform.Pass]
|
||||
returns list of passes
|
||||
"""
|
||||
return _ffi_api.get_vtcm_compaction_passes() # type: ignore # pylint: disable=no-member
|
||||
|
||||
|
||||
def is_pure_function(func: PrimFunc) -> bool:
|
||||
"""Checks if the function is a pure function"""
|
||||
return _ffi_api.is_pure_function(func, False) # type: ignore # pylint: disable=no-member
|
||||
|
||||
|
||||
def assert_pure_function(func: PrimFunc) -> bool:
|
||||
"""Asserts that the function is a pure function"""
|
||||
return _ffi_api.is_pure_function(func, True) # type: ignore # pylint: disable=no-member
|
||||
@@ -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.
|
||||
"""FFI APIs for tvm.s_tir.analysis"""
|
||||
|
||||
import tvm_ffi
|
||||
|
||||
tvm_ffi.init_ffi_api("s_tir.analysis", __name__)
|
||||
Reference in New Issue
Block a user