Siyuan Feng
c09a24dcdc
This PR supports a simple meta-programming paradigm for TVMScript, which allows users to get access to var definition in the Python environment.
252 lines
9.5 KiB
Python
252 lines
9.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.
|
|
"""TVM Script Context Maintainer for TIR"""
|
|
|
|
from typing import List, Mapping, Union, Optional, Dict, Callable
|
|
import synr
|
|
|
|
|
|
import tvm
|
|
from tvm.ir import Span
|
|
from tvm.ir.expr import Range
|
|
from tvm.tir import Var, Buffer, PrimExpr, Stmt, MatchBufferRegion
|
|
from tvm.runtime import Object
|
|
from tvm.tir.expr import IterVar
|
|
from .tir.node import BufferSlice
|
|
|
|
|
|
class BlockInfo:
|
|
"""Information for block and block_realize signature
|
|
|
|
Examples
|
|
----------
|
|
.. code-block:: python
|
|
|
|
@T.prim_func
|
|
def example_func(a: T.handle, b: T.handle, c: T.handle) -> None:
|
|
A = T.match_buffer(a, (16, 16), "float32")
|
|
B = T.match_buffer(b, (16, 16), "float32")
|
|
C = T.match_buffer(a, (16, 16), "float32")
|
|
|
|
for i, j, k in T.grid(16, 16, 16):
|
|
with T.block("matmul"):
|
|
vi = T.axis.S(16, i)
|
|
vj = T.axis.S(16, j)
|
|
vk = T.axis.R(16, k) # iter_bindings = {vj: i, vj: j, vk: k}
|
|
|
|
T.where(True) # predicate of the block_realize
|
|
|
|
T.reads(A[0:16, 0:16], B[0: 16, 0: 16]) # reads region of the block
|
|
T.writes(C[0: 16, 0: 16]) # writes region of the block
|
|
T.block_attr({"attr_key": "attr_value"}) # block annotations
|
|
|
|
# alloc_buffers inside the block
|
|
CC = T.alloc_buffer((1, 1), dtype="float32")
|
|
|
|
# match_buffers of the block,
|
|
# which bind a sub-region of source buffer into a new buffer
|
|
D = T.match_buffer(C[vi, vj], ())
|
|
|
|
# init part of the block, executed when all reduce axes are the beginning value
|
|
with T.init():
|
|
C[vi, vj] = T.float32(0)
|
|
|
|
# block body
|
|
CC[0, 0] = A[vi, vk] * B[vj, vk]
|
|
D[()] += CC[0, 0] # The same as C[vi, vj] += CC[0, 0]
|
|
"""
|
|
|
|
alloc_buffers: List[Buffer] = []
|
|
"""List[Buffer]: list of T.alloc_buffer statements in the block signature"""
|
|
match_buffers: List[MatchBufferRegion] = []
|
|
"""List[MatchBufferRegion]: list of T.match_buffer statements in the block signature"""
|
|
iter_values: List[PrimExpr] = []
|
|
"""List[PrimExpr]: list of binding values for iter vars"""
|
|
iter_vars: List[IterVar] = []
|
|
"""List[PrimExpr]: list of iter vars in the block"""
|
|
reads: Optional[List[BufferSlice]] = None
|
|
"""Optional[List[BufferSlice]]:
|
|
list of T.reads statements in the block signature, None for not-visited"""
|
|
writes: Optional[List[BufferSlice]] = None
|
|
"""Optional[List[BufferSlice]]:
|
|
list of T.writes statements in the block signature, None for not-visited"""
|
|
annotations: Optional[Mapping[str, Object]] = None
|
|
"""Optional[Mapping[str, Object]]:
|
|
list of T.block_attr statements in the block signature, None for not-visited"""
|
|
predicate: Optional[PrimExpr] = None
|
|
"""Optional[PrimExpr]: block realize predicate, None for not-visited"""
|
|
init: Optional[Stmt] = None
|
|
"""Optional[Stmt]: init part of the block, None for not-visited"""
|
|
|
|
def __init__(self):
|
|
self.alloc_buffers = []
|
|
self.match_buffers = []
|
|
self.iter_values = []
|
|
self.iter_vars = []
|
|
self.reads = None
|
|
self.writes = None
|
|
self.annotations = None
|
|
self.predicate = None
|
|
self.init = None
|
|
|
|
|
|
class ContextMaintainer:
|
|
"""Maintain all the necessary context info
|
|
Parameters
|
|
----------
|
|
_report_error : Callable[[str, Union[Span, synr.ast.Span]], None]
|
|
The report error function handle
|
|
"""
|
|
|
|
# scope context
|
|
node_stack: List[List[synr.ast.Node]] = []
|
|
"""List[List[synr.ast.Node]]: The ast nodes insides the current scope"""
|
|
block_info_stack: List[BlockInfo] = []
|
|
"""List[BlockInfo]: The block info for the current block scope"""
|
|
loop_stack: Dict[Var, Range] = {}
|
|
"""Dict[Var, Range]: The dict from loop var to its domain outside the block"""
|
|
symbols: List[Dict[str, Union[Var, Buffer]]] = []
|
|
"""List[Dict[str, Union[Var, Buffer]]]: Symbol map from name to object for the current scope"""
|
|
closure_vars: Dict[str, Object] = {}
|
|
"""ClosureVars: The closure vars defined in Python interpreter"""
|
|
|
|
# function context
|
|
func_params: List[Var] = []
|
|
"""List[Var]: The function parameters"""
|
|
func_buffer_map: Mapping[Var, Buffer] = {}
|
|
"""Mapping[Var, Buffer]: The function buffer map"""
|
|
func_preflattened_buffer_map: Mapping[Var, Buffer] = {}
|
|
"""Mapping[Var, Buffer]: The function buffer map, prior to any flattening."""
|
|
func_dict_attr: Mapping[str, Object] = {}
|
|
"""Mapping[str, Object]: The function attrs"""
|
|
func_var_env_dict: Mapping[Var, str] = {}
|
|
"""Mapping[Var, str]: The map from var to env thread"""
|
|
|
|
# parser and analyzer
|
|
analyzer: tvm.arith.Analyzer = tvm.arith.Analyzer()
|
|
"""tvm.arith.Analyzer: The analyzer for simplifying"""
|
|
_report_error: Callable[[str, Union[Span, synr.ast.Span]], None]
|
|
"""Callable[[str, Union[Span, synr.ast.Span]], None]: The report error function handle"""
|
|
|
|
# root alloc_buffer
|
|
root_alloc_buffers: List[Buffer] = []
|
|
"""List[Buffer]: The buffers allocated under root block"""
|
|
|
|
def __init__(
|
|
self,
|
|
_report_error: Callable[[str, Union[Span, synr.ast.Span]], None],
|
|
closure_vars: Dict[str, Object],
|
|
):
|
|
# scope context
|
|
self.node_stack = []
|
|
self.block_info_stack = []
|
|
self.loop_stack = {}
|
|
self.symbols = []
|
|
self.closure_vars = closure_vars
|
|
# function context
|
|
self.func_params = []
|
|
self.func_buffer_map = {}
|
|
self.func_preflattened_buffer_map = {}
|
|
self.func_dict_attr = {}
|
|
self.func_var_env_dict = {}
|
|
# parser and analyzer
|
|
self._report_error = _report_error
|
|
self.analyzer = tvm.arith.Analyzer()
|
|
# root alloc_buffer
|
|
self.root_alloc_buffers = []
|
|
|
|
def enter_scope(self, nodes: Optional[List[synr.ast.Node]] = None):
|
|
"""Creates a new scope
|
|
|
|
Note
|
|
----
|
|
This function is used for normal scopes that do not involve
|
|
a `with block` scope. Use `enter_block_scope`
|
|
for block scope cases.
|
|
|
|
Parameters
|
|
----------
|
|
nodes : Optional[List[synr.ast.Node]]
|
|
The synr AST nodes in new scope
|
|
"""
|
|
if nodes is None:
|
|
nodes = []
|
|
self.node_stack.append(list(reversed(nodes)))
|
|
self.symbols.append(dict())
|
|
|
|
def enter_block_scope(self, nodes: Optional[List[synr.ast.Node]] = None):
|
|
"""Creates a new block scope, the function will call `enter_scope` implicitly
|
|
Besides the behaviors of `enter_scope`, it will update loop_stack and block_info_stack
|
|
to maintain block info.
|
|
|
|
Note
|
|
----
|
|
This function should be used to handle a block scope,
|
|
aka the blocks that involve a `with block` scope.
|
|
|
|
Parameters
|
|
----------
|
|
nodes : Optional[List[synr.ast.Node]]
|
|
The synr AST nodes in new scope
|
|
"""
|
|
self.enter_scope(nodes)
|
|
# Create a new BlockInfo for the new block
|
|
self.block_info_stack.append(BlockInfo())
|
|
|
|
def exit_scope(self):
|
|
"""Pop the inner most scope"""
|
|
self.symbols.pop()
|
|
self.node_stack.pop()
|
|
|
|
def exit_block_scope(self):
|
|
"""Pop the inner most block scope, the function will call `exit_scope` implicitly"""
|
|
self.exit_scope()
|
|
# Pop block_info
|
|
self.block_info_stack.pop()
|
|
|
|
def update_symbol(self, name: str, symbol: Union[Buffer, Var], node: synr.ast.Node):
|
|
"""Append a symbol into current scope"""
|
|
if isinstance(symbol, Buffer):
|
|
if name in self.symbols[0]:
|
|
self.report_error("Duplicate Buffer name: " + symbol.name, node.span)
|
|
self.symbols[0][name] = symbol
|
|
else:
|
|
self.symbols[-1][name] = symbol
|
|
|
|
def remove_symbol(self, name: str):
|
|
"""Remove a symbol"""
|
|
for symbols in reversed(self.symbols):
|
|
if name in symbols:
|
|
symbols.pop(name)
|
|
return
|
|
raise RuntimeError("Internal error of tvm script parser: no symbol named " + name)
|
|
|
|
def lookup_symbol(self, name: str) -> Optional[Union[Buffer, Var]]:
|
|
"""Look up symbol by name"""
|
|
for symbols in reversed(self.symbols):
|
|
if name in symbols:
|
|
return symbols[name]
|
|
return self.closure_vars.get(name)
|
|
|
|
def report_error(self, message: str, span: Union[Span, synr.ast.Span]):
|
|
self._report_error(message, span)
|
|
|
|
def current_block_scope(self) -> BlockInfo:
|
|
if self.block_info_stack:
|
|
return self.block_info_stack[-1]
|
|
return None
|