chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,914 @@
|
||||
# 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.
|
||||
"""The base parser for tirx"""
|
||||
|
||||
import ast
|
||||
import contextlib
|
||||
from copy import deepcopy
|
||||
from functools import partial
|
||||
from typing import Any
|
||||
|
||||
import tvm
|
||||
from tvm.ir import Expr, GlobalVar, PrimType
|
||||
from tvm.script.ir_builder import ir as I
|
||||
from tvm.script.ir_builder.base import IRBuilder
|
||||
from tvm.script.ir_builder.base import IRBuilderFrame as Frame
|
||||
from tvm.script.parser._core import Parser, dispatch, doc
|
||||
from tvm.script.parser.core.doc import from_doc
|
||||
from tvm.tirx import Buffer, IterVar, Layout, Var
|
||||
from tvm.tirx.script import builder as T
|
||||
from tvm.tirx.script.builder.ir import name_meta_class_value
|
||||
from tvm.tirx.stmt import BufferRegion
|
||||
|
||||
from .entry import constexpr as _constexpr_sentinel
|
||||
from .entry import inline
|
||||
|
||||
|
||||
def slice_buffer_from_region(br: BufferRegion) -> Buffer:
|
||||
"""Create a matched DeclBuffer from a BufferRegion.
|
||||
|
||||
Slices the layout (if present) or computes elem_offset for the sub-region,
|
||||
producing a DeclBuffer that views the same underlying data.
|
||||
"""
|
||||
import functools # pylint: disable=import-outside-toplevel
|
||||
|
||||
buf = br.buffer
|
||||
region = br.region
|
||||
new_shape = [r.extent for r in region]
|
||||
sliced_layout = None
|
||||
if buf.layout is not None:
|
||||
range_pairs = [(r.min, r.min + r.extent) for r in region]
|
||||
sliced_layout = buf.layout.slice(list(buf.shape), range_pairs)
|
||||
if sliced_layout is not None:
|
||||
return T.decl_buffer(
|
||||
new_shape,
|
||||
buf.dtype,
|
||||
buf.data,
|
||||
buf.strides,
|
||||
buf.elem_offset,
|
||||
None,
|
||||
buf.scope(),
|
||||
buf.data_alignment,
|
||||
buf.offset_factor,
|
||||
"",
|
||||
buf.axis_separators,
|
||||
sliced_layout,
|
||||
)
|
||||
# Fallback: compute elem_offset for default/no layout
|
||||
strides = []
|
||||
for i in range(len(buf.shape)):
|
||||
stride = functools.reduce(
|
||||
lambda x, y: x * y, buf.shape[i + 1 :], tvm.tirx.const(1, "int32")
|
||||
)
|
||||
strides.append(stride)
|
||||
offset = tvm.tirx.const(0, "int32")
|
||||
for i, r in enumerate(region):
|
||||
offset = offset + r.min * strides[i]
|
||||
new_elem_offset = buf.elem_offset + offset
|
||||
return T.decl_buffer(
|
||||
new_shape,
|
||||
buf.dtype,
|
||||
buf.data,
|
||||
buf.strides,
|
||||
new_elem_offset,
|
||||
None,
|
||||
buf.scope(),
|
||||
buf.data_alignment,
|
||||
buf.offset_factor,
|
||||
"",
|
||||
buf.axis_separators,
|
||||
buf.layout,
|
||||
)
|
||||
|
||||
|
||||
def bind_with_value(self: Parser, node: doc.expr, var_name: str, value: Any) -> Any:
|
||||
"""Value binding methods when parsing with statement.
|
||||
e.g. binding i, j, k with T.grid(128, 128, 128), when parsing
|
||||
with T.grid(128, 128, 18) as i, j, k.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The current parser.
|
||||
|
||||
node : doc.expr
|
||||
The doc AST expression node for error reporting.
|
||||
|
||||
var_name : str
|
||||
The variable name.
|
||||
|
||||
value : Any
|
||||
The value to be bound with.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The bound value.
|
||||
"""
|
||||
if isinstance(value, list | tuple):
|
||||
for i, v in enumerate(value):
|
||||
bind_with_value(self, node, f"{var_name}_{i}", v)
|
||||
return value
|
||||
elif isinstance(value, Buffer | Var):
|
||||
IRBuilder.name(var_name, value)
|
||||
return value
|
||||
else:
|
||||
self.report_error(node, f"Do not know how to bind type: {type(value)} in with statement")
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
def bind_for_value(self: Parser, node: doc.expr, var_name: str, value: Any) -> Any:
|
||||
"""Value binding methods when parsing for statement.
|
||||
e.g. binding i, j, k with T.grid(128, 128, 128), when parsing
|
||||
for i, j, k in T.grid(128, 128, 128).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The current parser.
|
||||
|
||||
node : doc.expr
|
||||
The doc AST expression node for error reporting.
|
||||
|
||||
var_name : str
|
||||
The variable name.
|
||||
|
||||
value : Any
|
||||
The value to be bound with.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The bound value.
|
||||
"""
|
||||
if isinstance(value, list | tuple | tvm.ir.Array):
|
||||
for i, v in enumerate(value):
|
||||
bind_for_value(self, node, f"{var_name}_{i}", v)
|
||||
return value
|
||||
elif isinstance(value, Var):
|
||||
IRBuilder.name(var_name, value)
|
||||
return value
|
||||
else:
|
||||
self.report_error(node, f"Do not know how to bind type: {type(value)} in for statement")
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
def bind_assign_value(self: Parser, node: doc.expr, var_name: str, value: Any) -> Any:
|
||||
"""Value binding methods when parsing assign statement.
|
||||
e.g. binding vi, vj, vk with T.axis.remap("SSR", [i, j, k]), when parsing
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k]).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The current parser.
|
||||
|
||||
node : doc.expr
|
||||
The doc AST expression node for error reporting.
|
||||
|
||||
var_name : str
|
||||
The variable name.
|
||||
|
||||
value : Any
|
||||
The value to be bound with.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The bound value.
|
||||
"""
|
||||
if isinstance(value, T.scalar_wrapper): # pylint: disable=protected-access
|
||||
# special case for scalar, name the buffer, but the var is used as BufferLoad
|
||||
assert isinstance(value.scalar, T.BufferLoad)
|
||||
IRBuilder.name(var_name, value.scalar.buffer)
|
||||
return value.scalar
|
||||
if isinstance(value, T.meta_var):
|
||||
return value.value
|
||||
elif getattr(type(value), "_is_meta_class", False):
|
||||
name_meta_class_value(var_name, value)
|
||||
return value
|
||||
elif isinstance(value, list | tuple):
|
||||
# Tuple-unpacking with a starred target (e.g. ``vi, *vs = T.axis.remap(...)``)
|
||||
# collects multiple elements into a single list bound here. Recurse so each
|
||||
# element gets a per-index name; this matches apache's behavior.
|
||||
for i, v in enumerate(value):
|
||||
bind_assign_value(self, node, f"{var_name}_{i}", v)
|
||||
return value
|
||||
elif isinstance(value, BufferRegion):
|
||||
return value
|
||||
elif isinstance(value, Frame):
|
||||
value.add_callback(partial(value.__exit__, None, None, None))
|
||||
res = value.__enter__()
|
||||
IRBuilder.name(var_name, res)
|
||||
return res
|
||||
elif isinstance(value, Buffer | IterVar | Layout) or (
|
||||
isinstance(value, Var) and not self.var_table.exist(value)
|
||||
):
|
||||
IRBuilder.name(var_name, value)
|
||||
return value
|
||||
else:
|
||||
if not tvm.ir.is_prim_expr(value):
|
||||
value = tvm.tirx.const(value)
|
||||
if not isinstance(value, tvm.tirx.StringImm):
|
||||
# x = expr -> scalar (auto-typed from value)
|
||||
scalar = T.local_scalar(dtype=str(value.ty.dtype))
|
||||
IRBuilder.name(var_name, scalar.scalar.buffer)
|
||||
T.buffer_store(scalar.scalar.buffer, value, [0])
|
||||
return scalar.scalar
|
||||
else:
|
||||
# StringImm: x = expr -> immutable Bind var
|
||||
ann_var = tvm.tirx.Var(var_name, value.ty)
|
||||
IRBuilder.name(var_name, ann_var)
|
||||
T.Bind(value, var=ann_var)
|
||||
return ann_var
|
||||
|
||||
|
||||
def find_decorator_annotation(node: doc.FunctionDef, annotation: str, default: bool = True) -> bool:
|
||||
"""
|
||||
Check the value of given annotation (argument name) in the prim_func decorator.
|
||||
Returns the value of the annotation if present, otherwise giving the default value.
|
||||
"""
|
||||
# look for the named argument in the prim_func / jit decorator
|
||||
for dec in node.decorator_list:
|
||||
if not isinstance(dec, doc.Call) or dec.func.attr not in ("prim_func", "jit"):
|
||||
continue
|
||||
for keyword in dec.keywords:
|
||||
if keyword.arg == annotation:
|
||||
return keyword.value.value
|
||||
return default
|
||||
|
||||
|
||||
@dispatch.register(token="tirx", type_name="For")
|
||||
def visit_for(self: Parser, node: doc.For) -> None:
|
||||
"""The for visiting method for tirx.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.For
|
||||
The doc AST for node.
|
||||
"""
|
||||
# Intercept range() at AST level so it works with both Python ints and PrimExprs.
|
||||
# In other contexts (e.g. list comprehensions), range remains Python's builtin.
|
||||
if (
|
||||
isinstance(node.iter, doc.Call)
|
||||
and isinstance(node.iter.func, doc.Name)
|
||||
and node.iter.func.id == "range"
|
||||
):
|
||||
args = [self.eval_expr(a) for a in node.iter.args]
|
||||
kwargs = {kw.arg: self.eval_expr(kw.value) for kw in node.iter.keywords}
|
||||
if len(args) == 1:
|
||||
for_frame = T.serial(0, args[0], **kwargs)
|
||||
elif len(args) == 2:
|
||||
for_frame = T.serial(args[0], args[1], **kwargs)
|
||||
elif len(args) == 3:
|
||||
for_frame = T.serial(args[0], args[1], step=args[2], **kwargs)
|
||||
else:
|
||||
self.report_error(node.iter, "range() takes 1 to 3 arguments")
|
||||
else:
|
||||
for_frame = self.eval_expr(node.iter)
|
||||
if not isinstance(for_frame, T.frame.ForFrame):
|
||||
self.report_error(
|
||||
node.iter,
|
||||
"Expect the for loop to be one of the following: "
|
||||
"range, T.serial, T.grid, T.parallel, T.vectorized, T.unroll, T.thread_binding",
|
||||
)
|
||||
with self.var_table.with_frame():
|
||||
with for_frame as iters:
|
||||
self.eval_assign(target=node.target, source=iters, bind_value=bind_for_value)
|
||||
self.visit_body(node.body)
|
||||
|
||||
|
||||
@dispatch.register(token="tirx", type_name="While")
|
||||
def visit_while(self: Parser, node: doc.While) -> None:
|
||||
"""The while visiting method for tirx.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.While
|
||||
The doc AST while node.
|
||||
"""
|
||||
with self.var_table.with_frame():
|
||||
cond = self.eval_expr(node.test)
|
||||
with T.While(cond):
|
||||
self.visit_body(node.body)
|
||||
|
||||
|
||||
@dispatch.register(token="tirx", type_name="Break")
|
||||
def visit_break(self: Parser, node: doc.Break) -> None:
|
||||
"""The break visiting method for tir.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.Break
|
||||
The doc AST break node.
|
||||
"""
|
||||
T.evaluate(T.break_loop())
|
||||
|
||||
|
||||
@dispatch.register(token="tirx", type_name="Continue")
|
||||
def visit_continue(self: Parser, node: doc.Continue) -> None:
|
||||
"""The continue visiting method for tir.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.Continue
|
||||
The doc AST continue node.
|
||||
"""
|
||||
T.evaluate(T.continue_loop())
|
||||
|
||||
|
||||
@dispatch.register(token="tirx", type_name="Assign")
|
||||
def visit_assign(self: Parser, node: doc.Assign) -> None:
|
||||
"""The assign visiting method for tirx.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.Assign
|
||||
The doc AST assign node.
|
||||
"""
|
||||
if len(node.targets) != 1:
|
||||
self.report_error(node, "Consequential assignments like 'a = b = c' are not supported.")
|
||||
lhs = node.targets[0]
|
||||
|
||||
if isinstance(node.value, doc.Subscript):
|
||||
check_slices = []
|
||||
if isinstance(node.value.slice, doc.Slice):
|
||||
check_slices = [node.value.slice]
|
||||
elif isinstance(node.value.slice, doc.Tuple):
|
||||
for p in node.value.slice.elts:
|
||||
if isinstance(p, doc.Slice):
|
||||
check_slices.append(p)
|
||||
for s in check_slices:
|
||||
if not s.step and s.upper and s.lower:
|
||||
s.step = doc.Constant(
|
||||
1,
|
||||
None,
|
||||
s.upper.lineno,
|
||||
s.upper.end_col_offset + 1,
|
||||
s.upper.lineno,
|
||||
s.upper.end_col_offset + 2,
|
||||
)
|
||||
|
||||
rhs = self.eval_expr(node.value)
|
||||
if isinstance(lhs, doc.Subscript):
|
||||
if isinstance(lhs.slice, doc.Tuple):
|
||||
indices = []
|
||||
for index in lhs.slice.elts:
|
||||
if isinstance(index, doc.Starred):
|
||||
# x[*y]
|
||||
indices.extend(self.eval_expr(index.value))
|
||||
else:
|
||||
indices.append(self.eval_expr(index))
|
||||
else:
|
||||
indices = self.eval_expr(lhs.slice)
|
||||
T.buffer_store(self.eval_expr(lhs.value), rhs, indices)
|
||||
else:
|
||||
# special case for scalar buffers
|
||||
# scalar = xxx <=> scalar.buffer[()] = xxx
|
||||
# or for a normal 1-dim buffer with shape (1,)
|
||||
# buffer = xxx <=> buffer[()] = xxx
|
||||
# Try to resolve lhs as a buffer/scalar variable. eval_expr may raise
|
||||
# if the name is not yet defined (i.e. this is a new variable binding),
|
||||
# which is the expected fallthrough case.
|
||||
lhs_value = None
|
||||
try:
|
||||
lhs_copy = deepcopy(lhs)
|
||||
if hasattr(lhs_copy, "ctx"):
|
||||
lhs_copy.ctx = doc.Load()
|
||||
lhs_value = self.eval_expr(lhs_copy)
|
||||
except Exception: # pylint: disable=broad-except
|
||||
pass
|
||||
# Buffer check and store are intentionally outside the try/except so
|
||||
# that genuine errors (e.g. wrong shape, bad store) are not swallowed.
|
||||
# Only TypeError from FFI type mismatch (e.g. rhs is a meta_var, not
|
||||
# a Expr or auto-convertible scalar) triggers fallthrough.
|
||||
if isinstance(lhs_value, T.scalar_wrapper | T.BufferLoad | tvm.tirx.Buffer):
|
||||
if isinstance(lhs_value, T.scalar_wrapper):
|
||||
buffer = lhs_value.scalar.buffer
|
||||
else:
|
||||
buffer = lhs_value.buffer if isinstance(lhs_value, T.BufferLoad) else lhs_value
|
||||
if len(buffer.shape) == 1 and bool(buffer.shape[0] == 1):
|
||||
# only 1-dim buffer with shape (1,) can be assigned directly
|
||||
# Note that shape can be a Expr, so we only judge by
|
||||
# bool(shape[0] == 1) rather than int(shape[0]) == 1.
|
||||
try:
|
||||
T.buffer_store(buffer, rhs, [0])
|
||||
return
|
||||
except TypeError:
|
||||
pass # rhs not compatible with buffer_store, fall through
|
||||
# otherwise
|
||||
self.eval_assign(target=lhs, source=rhs, bind_value=bind_assign_value)
|
||||
|
||||
|
||||
@dispatch.register(token="tirx", type_name="AugAssign")
|
||||
def visit_aug_assign(self: Parser, node: doc.AugAssign) -> None:
|
||||
"""The augmented assign visiting method for tirx.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.AugAssign
|
||||
The doc AST augmented assign node.
|
||||
"""
|
||||
lhs_pos = (
|
||||
node.target.lineno,
|
||||
node.target.col_offset,
|
||||
node.target.end_lineno,
|
||||
node.target.end_col_offset,
|
||||
)
|
||||
rhs_pos = (
|
||||
node.value.lineno,
|
||||
node.value.col_offset,
|
||||
node.value.end_lineno,
|
||||
node.value.end_col_offset,
|
||||
)
|
||||
node.target.ctx = doc.Load()
|
||||
with self.var_table.with_frame():
|
||||
lhs_name = "__tvm_tmp_value_aug_assign_lhs"
|
||||
rhs_name = "__tvm_tmp_value_aug_assign_rhs"
|
||||
lhs_expr = self.eval_expr(node.target)
|
||||
rhs_expr = self.eval_expr(node.value)
|
||||
self.var_table.add(lhs_name, lhs_expr)
|
||||
self.var_table.add(rhs_name, rhs_expr)
|
||||
op = doc.BinOp(
|
||||
doc.Name(lhs_name, doc.Load(), *lhs_pos),
|
||||
node.op,
|
||||
doc.Name(rhs_name, doc.Load(), *rhs_pos),
|
||||
*lhs_pos,
|
||||
)
|
||||
rhs = self.eval_expr(op)
|
||||
lhs = node.target
|
||||
lhs.ctx = doc.Store()
|
||||
if isinstance(lhs, doc.Subscript):
|
||||
if isinstance(lhs.slice, doc.Tuple):
|
||||
indices = []
|
||||
for index in lhs.slice.elts:
|
||||
if isinstance(index, doc.Starred):
|
||||
# x[*y]
|
||||
indices.extend(self.eval_expr(index.value))
|
||||
else:
|
||||
indices.append(self.eval_expr(index))
|
||||
else:
|
||||
indices = [self.eval_expr(lhs.slice)]
|
||||
T.buffer_store(self.eval_expr(lhs.value), rhs, indices)
|
||||
else:
|
||||
lhs_value = None
|
||||
try:
|
||||
lhs_copy = deepcopy(lhs)
|
||||
if hasattr(lhs_copy, "ctx"):
|
||||
lhs_copy.ctx = doc.Load()
|
||||
lhs_value = self.eval_expr(lhs_copy)
|
||||
except Exception: # pylint: disable=broad-except
|
||||
pass
|
||||
if isinstance(lhs_value, T.scalar_wrapper | T.BufferLoad | tvm.tirx.Buffer):
|
||||
if isinstance(lhs_value, T.scalar_wrapper):
|
||||
buffer = lhs_value.scalar.buffer
|
||||
else:
|
||||
buffer = lhs_value.buffer if isinstance(lhs_value, T.BufferLoad) else lhs_value
|
||||
if len(buffer.shape) == 1 and bool(buffer.shape[0] == 1):
|
||||
try:
|
||||
T.buffer_store(buffer, rhs, [0])
|
||||
return
|
||||
except TypeError:
|
||||
pass
|
||||
self.eval_assign(target=lhs, source=rhs, bind_value=bind_assign_value)
|
||||
|
||||
|
||||
@dispatch.register(token="tirx", type_name="AnnAssign")
|
||||
def visit_ann_assign(self: Parser, node: doc.AnnAssign) -> None:
|
||||
"""The annotated assign visiting method for tirx.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.AnnAssign
|
||||
The doc AST annotated assign node.
|
||||
"""
|
||||
lhs = node.target
|
||||
rhs = self.eval_expr(node.value) if node.value is not None else None
|
||||
raw_ann = self.eval_expr(node.annotation)
|
||||
|
||||
if isinstance(raw_ann, T.LocalVectorAnnotation):
|
||||
# x: T.float32[N] or x: T.f32[M, N] -> local buffer allocation
|
||||
if rhs is not None:
|
||||
self.report_error(node, "Vector annotation does not support initial value")
|
||||
buf = T.alloc_local(shape=raw_ann.shape, dtype=raw_ann.dtype)
|
||||
self.eval_assign(target=lhs, source=buf, bind_value=bind_assign_value)
|
||||
elif isinstance(raw_ann, T.LetAnnotation):
|
||||
# T.let or T.let[type] -> immutable Bind var
|
||||
if rhs is None:
|
||||
self.report_error(node, "T.let annotation requires a value")
|
||||
if not isinstance(rhs, Expr):
|
||||
if isinstance(rhs, str):
|
||||
rhs = tvm.tirx.StringImm(rhs)
|
||||
else:
|
||||
rhs = tvm.tirx.const(rhs)
|
||||
if raw_ann.type_spec is not None:
|
||||
ann_var = raw_ann.as_var()
|
||||
else:
|
||||
ann_var = raw_ann.as_var(rhs_dtype=rhs.ty)
|
||||
if not isinstance(ann_var, Var):
|
||||
self.report_error(node.annotation, "Annotation should resolve to Var")
|
||||
self.eval_assign(target=lhs, source=ann_var, bind_value=bind_assign_value)
|
||||
T.Bind(rhs, var=ann_var)
|
||||
else:
|
||||
ann_var = raw_ann() if callable(raw_ann) else raw_ann
|
||||
if not isinstance(ann_var, Var):
|
||||
self.report_error(node.annotation, "Annotation should resolve to Var")
|
||||
if not isinstance(ann_var.ty, PrimType):
|
||||
self.report_error(
|
||||
node.annotation,
|
||||
"Use T.let[...] for non-PrimType annotations (e.g. PointerType, handle)",
|
||||
)
|
||||
if str(ann_var.ty) == "handle":
|
||||
self.report_error(
|
||||
node.annotation,
|
||||
"handle type cannot be used as scalar annotation; use T.let[T.handle] instead",
|
||||
)
|
||||
# x: T.int32 = expr -> scalar (mutable scalar buffer)
|
||||
scalar = T.local_scalar(dtype=str(ann_var.ty))
|
||||
self.eval_assign(target=lhs, source=scalar, bind_value=bind_assign_value)
|
||||
if rhs is not None:
|
||||
T.buffer_store(scalar.scalar.buffer, rhs, [0])
|
||||
|
||||
|
||||
@dispatch.register(token="tirx", type_name="With")
|
||||
def visit_with(self: Parser, node: doc.With) -> None:
|
||||
"""The with visiting method for tirx.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.With
|
||||
The doc AST with node.
|
||||
"""
|
||||
with contextlib.ExitStack() as stack:
|
||||
stack.enter_context(self.var_table.with_frame())
|
||||
for item in node.items:
|
||||
frame = self.eval_expr(item.context_expr)
|
||||
if not isinstance(frame, Frame) and not (
|
||||
hasattr(frame, "__enter__") and hasattr(frame, "__exit__")
|
||||
):
|
||||
self.report_error(
|
||||
item.context_expr,
|
||||
"Invalid context expression in the with-statement.",
|
||||
)
|
||||
rhs = stack.enter_context(frame)
|
||||
if item.optional_vars is not None:
|
||||
self.eval_assign(target=item.optional_vars, source=rhs, bind_value=bind_with_value)
|
||||
self.visit_body(node.body)
|
||||
|
||||
|
||||
@dispatch.register(token="tirx", type_name="FunctionDef")
|
||||
def visit_function_def(self: Parser, node: doc.FunctionDef) -> None:
|
||||
"""The function definition visiting method for tirx.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.FunctionDef
|
||||
The doc AST function definition node.
|
||||
"""
|
||||
supplied_annotation = self.function_annotations
|
||||
func_annotation = supplied_annotation.get(node.name, {})
|
||||
privacy = find_decorator_annotation(node, "private", default=False)
|
||||
s_tir = find_decorator_annotation(node, "s_tir", default=False)
|
||||
persistent = find_decorator_annotation(node, "persistent", default=False)
|
||||
self.function_annotations = None
|
||||
with self.var_table.with_frame():
|
||||
prim_func_ctx = T.prim_func(is_private=privacy, s_tir=s_tir, persistent=persistent)
|
||||
with prim_func_ctx:
|
||||
T.func_name(node.name)
|
||||
if node.returns is not None:
|
||||
ret_type = self.eval_expr(node.returns)
|
||||
if callable(ret_type):
|
||||
ret_type = ret_type().ty
|
||||
T.func_ret(ret_type)
|
||||
with self.with_dispatch_token("tirx"):
|
||||
# TODO: handle different types of arguments:
|
||||
# - vararg: arg | None
|
||||
# - kwonlyargs: list[arg]
|
||||
# - kw_defaults: list[expr | None]
|
||||
# - kwarg: arg | None
|
||||
# - defaults: list[expr]
|
||||
# - posonlyargs: list[arg]
|
||||
for arg in node.args.args:
|
||||
if arg.annotation is None:
|
||||
self.report_error(arg, "Type annotation required for function parameters.")
|
||||
try:
|
||||
ann = self.eval_expr(arg.annotation)
|
||||
if callable(ann) and ann is not _constexpr_sentinel:
|
||||
ann = ann()
|
||||
except Exception: # pylint: disable=broad-except
|
||||
ann = func_annotation.get(arg.arg, None)
|
||||
if ann is None:
|
||||
raise
|
||||
if ann is _constexpr_sentinel:
|
||||
# T.constexpr param: value was bound in extra_vars by
|
||||
# TIRJit.specialize() and lives in an outer var_table
|
||||
# frame; do not register a runtime PrimFunc param.
|
||||
continue
|
||||
param = T.arg(arg.arg, ann)
|
||||
self.var_table.add(arg.arg, param)
|
||||
self.visit_body(node.body)
|
||||
self.function_annotations = supplied_annotation
|
||||
|
||||
|
||||
@dispatch.register(token="tir.inline", type_name="FunctionDef")
|
||||
def visit_inline_function_def(self: Parser, node: doc.FunctionDef) -> None:
|
||||
"""The function definition visiting method for inline functions in tir.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.FunctionDef
|
||||
The doc AST function definition node.
|
||||
"""
|
||||
# remove the inline decorator
|
||||
node.decorator_list.pop()
|
||||
# adjust the node location to the source code location
|
||||
node.lineno += self.diag.source.start_line - 1
|
||||
node.col_offset += self.diag.source.start_column + 1
|
||||
node.end_lineno += self.diag.source.start_line - 1
|
||||
node.end_col_offset += self.diag.source.start_column + 1
|
||||
|
||||
# Record definition depth for LEGB late binding
|
||||
definition_depth = len(self.var_table.frames)
|
||||
|
||||
def get_func():
|
||||
func_ast = from_doc(node)
|
||||
module_ast = ast.Module(body=[func_ast], type_ignores=[])
|
||||
ast.fix_missing_locations(module_ast)
|
||||
# set the filename to the source name, so that the error message can be reported correctly
|
||||
code_obj = compile(module_ast, filename=self.diag.source.source_name, mode="exec")
|
||||
namespace = self.var_table.get()
|
||||
exec(code_obj, namespace) # pylint: disable=exec-used
|
||||
func_name = func_ast.name
|
||||
func = namespace[func_name]
|
||||
return func, func_name
|
||||
|
||||
func, func_name = get_func()
|
||||
wrapper = inline(func, definition_depth=definition_depth, defining_var_table=self.var_table)
|
||||
|
||||
self.var_table.add(func_name, wrapper, allow_shadowing=False)
|
||||
return None
|
||||
|
||||
|
||||
@dispatch.register(token="tirx", type_name="tvm_annotation")
|
||||
def visit_tvm_annotation(self: Parser, node: doc.expr):
|
||||
"""The TVM annotation visiting method for tirx.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.expr
|
||||
The doc AST expr node.
|
||||
"""
|
||||
annotation = self.eval_expr(node)
|
||||
if callable(annotation):
|
||||
annotation = annotation()
|
||||
return annotation
|
||||
|
||||
|
||||
@dispatch.register(token="tirx", type_name="Expr")
|
||||
def visit_expr_stmt(self: Parser, node: doc.Expr) -> None:
|
||||
"""The expr statement visiting method for tirx.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.Expr
|
||||
The doc AST Expr node.
|
||||
"""
|
||||
|
||||
res = self.eval_expr(node.value)
|
||||
if res is None:
|
||||
pass
|
||||
elif isinstance(res, Frame):
|
||||
res.add_callback(partial(res.__exit__, None, None, None))
|
||||
res.__enter__()
|
||||
elif hasattr(res, "frames") and hasattr(res, "__enter__"):
|
||||
# _FrameScope from T.attr({...}) — enter each inner frame for concise scoping
|
||||
for f in res.frames:
|
||||
f.add_callback(partial(f.__exit__, None, None, None))
|
||||
f.__enter__()
|
||||
elif isinstance(res, Var):
|
||||
# Standalone Var expression (e.g. from T.bind(value, var=v)) --
|
||||
# the Bind statement was already emitted to the parent frame by the FFI call,
|
||||
# so just discard the returned Var.
|
||||
pass
|
||||
elif tvm.ir.is_prim_expr(res):
|
||||
T.evaluate(res)
|
||||
elif isinstance(res, int | bool):
|
||||
T.evaluate(tvm.tirx.const(res))
|
||||
elif isinstance(res, tvm.ir.Call) and not tvm.ir.is_prim_expr(res):
|
||||
if isinstance(res.op, tvm.ir.GlobalVar) and res.ty.is_missing():
|
||||
# GlobalVar calls with a missing return type are ambiguous, as each IR has a
|
||||
# different function Call representation. Convert to the TIR representation.
|
||||
T.evaluate(tvm.tirx.call_tir(res.op, *res.args))
|
||||
else:
|
||||
# Pointer-valued TIR calls are general Expr rather than PrimExpr,
|
||||
# but are still valid standalone Evaluate statements.
|
||||
T.evaluate(res)
|
||||
elif isinstance(res, str):
|
||||
# Ignore docstrings
|
||||
pass
|
||||
elif isinstance(res, tvm.tirx.stmt.BufferStore):
|
||||
T.buffer_store(res.buffer, res.value, res.indices, res.predicate)
|
||||
elif isinstance(res, tvm.tirx.Buffer):
|
||||
# ``T.match_buffer(...)`` used as a bare statement (no LHS) — the
|
||||
# buffer object is discarded; the underlying side effect (the
|
||||
# match_buffer node) has already been emitted into the frame.
|
||||
pass
|
||||
else:
|
||||
self.report_error(node, f"Parsing resulted in unexpected type {type(res)}")
|
||||
|
||||
|
||||
@dispatch.register(token="tirx", type_name="If")
|
||||
def visit_if(self: Parser, node: doc.If) -> None:
|
||||
"""The if visiting method for tirx.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.If
|
||||
The doc AST if node.
|
||||
"""
|
||||
with self.var_table.with_frame():
|
||||
predicate = self.eval_expr(node.test)
|
||||
if tvm.ir.is_prim_expr(predicate) or isinstance(predicate, tvm.tirx.expr.ExprOp):
|
||||
with T.If(self.eval_expr(node.test)):
|
||||
with T.Then():
|
||||
with self.var_table.with_frame():
|
||||
self.visit_body(node.body)
|
||||
if node.orelse:
|
||||
with T.Else():
|
||||
with self.var_table.with_frame():
|
||||
self.visit_body(node.orelse)
|
||||
elif isinstance(predicate, bool):
|
||||
if predicate:
|
||||
with self.var_table.with_frame():
|
||||
self.visit_body(node.body)
|
||||
elif node.orelse:
|
||||
with self.var_table.with_frame():
|
||||
self.visit_body(node.orelse)
|
||||
else:
|
||||
self.report_error(
|
||||
node.test,
|
||||
f"If condition must be a boolean expression, but got {predicate}",
|
||||
)
|
||||
|
||||
|
||||
@dispatch.register(token="tirx", type_name="Assert")
|
||||
def visit_assert(self: Parser, node: doc.Assert) -> None:
|
||||
"""The assert visiting method for tirx.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.Assert
|
||||
The doc AST assert node.
|
||||
|
||||
The assert message can be either:
|
||||
- A plain string: ``assert cond, "message"``
|
||||
- A tuple of (kind, [parts...]): ``assert cond, ("ValueError", ["part0", "part1"])``
|
||||
"""
|
||||
cond = self.eval_expr(node.test)
|
||||
msg = self.eval_expr(node.msg)
|
||||
|
||||
kind = "RuntimeError"
|
||||
message = msg
|
||||
|
||||
if isinstance(msg, tuple):
|
||||
if len(msg) != 2:
|
||||
self.report_error(
|
||||
node,
|
||||
f"Assert message tuple must have exactly 2 elements (kind, [parts...]), "
|
||||
f"got {len(msg)} elements",
|
||||
)
|
||||
kind_str, parts = msg
|
||||
if isinstance(kind_str, tvm.tirx.StringImm):
|
||||
kind_str = kind_str.value
|
||||
if not isinstance(kind_str, str):
|
||||
self.report_error(
|
||||
node,
|
||||
f"Assert message tuple first element must be a string (error kind like "
|
||||
f'"ValueError"), got {type(kind_str).__name__}',
|
||||
)
|
||||
kind = kind_str
|
||||
message = parts
|
||||
|
||||
if isinstance(message, list | tuple):
|
||||
message = [p.value if isinstance(p, tvm.tirx.StringImm) else str(p) for p in message]
|
||||
|
||||
frame = T.Assert(cond, message, error_kind=kind)
|
||||
frame.add_callback(partial(frame.__exit__, None, None, None))
|
||||
frame.__enter__()
|
||||
|
||||
|
||||
@dispatch.register(token="tirx", type_name="Return")
|
||||
def visit_return(self: Parser, node: doc.Return) -> None:
|
||||
"""The return visiting method for tirx.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.Return
|
||||
The doc AST return node.
|
||||
"""
|
||||
value = self.eval_expr(node.value)
|
||||
if value is None:
|
||||
self.report_error(node, "Expression to be returned must be a Expr")
|
||||
T.evaluate(tvm.tirx.ret(value))
|
||||
|
||||
|
||||
@dispatch.register(token="tirx", type_name="tvm_declare_function")
|
||||
def visit_tvm_declare_function(self: Parser, node: doc.FunctionDef) -> GlobalVar:
|
||||
"""The function declaration step for tirx
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.Return
|
||||
The doc AST return node.
|
||||
"""
|
||||
|
||||
supplied_annotation = self.function_annotations
|
||||
func_annotation = supplied_annotation.get(node.name, {})
|
||||
|
||||
ret_type = None
|
||||
with self.var_table.with_frame():
|
||||
if node.returns is not None:
|
||||
ret_type = self.eval_expr(node.returns)
|
||||
if callable(ret_type):
|
||||
ret_type = ret_type().ty
|
||||
|
||||
arg_annotations = []
|
||||
for arg in node.args.args:
|
||||
if arg.annotation is None:
|
||||
self.report_error(arg, "Type annotation required for function parameters.")
|
||||
try:
|
||||
ann = self.eval_expr(arg.annotation)
|
||||
if callable(ann):
|
||||
ann = ann()
|
||||
except Exception: # pylint: disable=broad-except
|
||||
ann = func_annotation.get(arg.arg, None)
|
||||
if ann is None:
|
||||
raise
|
||||
|
||||
IRBuilder.name(arg.arg, ann)
|
||||
arg_annotations.append(ann)
|
||||
|
||||
func_signature = tvm.tirx.PrimFunc(arg_annotations, None, ret_type=ret_type)
|
||||
return I.decl_function(node.name, func_signature)
|
||||
Reference in New Issue
Block a user