wehub-resource-sync
1911 lines
54 KiB
Python
1911 lines
54 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=no-else-return, unidiomatic-typecheck, invalid-name, arguments-differ
|
|
# ruff: noqa: RUF012
|
|
"""The expression and statement functor of TIR."""
|
|
|
|
from collections.abc import Callable
|
|
|
|
import tvm_ffi
|
|
|
|
from tvm.ir import Call, Expr
|
|
from tvm.ir.utils import derived_object
|
|
|
|
from . import _ffi_api
|
|
from .expr import (
|
|
EQ,
|
|
GE,
|
|
GT,
|
|
LE,
|
|
LT,
|
|
NE,
|
|
Add,
|
|
And,
|
|
Broadcast,
|
|
BufferLoad,
|
|
Cast,
|
|
Div,
|
|
FloatImm,
|
|
FloorDiv,
|
|
FloorMod,
|
|
IntImm,
|
|
Let,
|
|
Max,
|
|
Min,
|
|
Mod,
|
|
Mul,
|
|
Not,
|
|
Or,
|
|
ProducerLoad,
|
|
Ramp,
|
|
Reduce,
|
|
Select,
|
|
Shuffle,
|
|
StringImm,
|
|
Sub,
|
|
Var,
|
|
)
|
|
from .stmt import (
|
|
AllocBuffer,
|
|
AssertStmt,
|
|
AttrStmt,
|
|
Bind,
|
|
BufferStore,
|
|
DeclBuffer,
|
|
Evaluate,
|
|
For,
|
|
IfThenElse,
|
|
SBlock,
|
|
SBlockRealize,
|
|
SeqStmt,
|
|
Stmt,
|
|
While,
|
|
)
|
|
|
|
# visitor and mutator are aliases for derived_object
|
|
visitor = derived_object
|
|
mutator = derived_object
|
|
|
|
"""
|
|
A decorator to wrap user-customized PyStmtExprMutator as TVM object _PyStmtExprMutator.
|
|
|
|
Parameters
|
|
----------
|
|
mutator_cls : PyStmtExprMutator
|
|
The user-customized PyStmtExprMutator.
|
|
|
|
Returns
|
|
-------
|
|
cls : _PyStmtExprMutator
|
|
The decorated TVM object _PyStmtExprMutator(StmtExprMutator on the C++ side).
|
|
|
|
Example
|
|
-------
|
|
.. code-block:: python
|
|
|
|
@tirx.functor.stmt_expr_mutator
|
|
class MyStmtExprMutator(PyStmtExprMutator):
|
|
# customize rewrite function
|
|
def visit_add_(self, op: Add) -> Expr:
|
|
# just for demo purposes
|
|
...
|
|
|
|
# mymutator is now a special mutator that rewrite every Add with
|
|
# user-customized visit_add_
|
|
mymutator = MyStmtExprMutator()
|
|
# apply mymutator to Expr and Stmt
|
|
mymutator.visit_expr(expr)
|
|
mymutator.visit_stmt(stmt)
|
|
"""
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.PyStmtExprVisitor")
|
|
class _PyStmtExprVisitor(tvm_ffi.core.Object):
|
|
"""
|
|
An internal wrapper to interface between C++ and Python StmtExprVisitor.
|
|
This is the TVM object that wraps PyStmtExprVisitor.
|
|
|
|
Do not use this class directly. Use PyStmtExprVisitor instead.
|
|
|
|
See also: PyStmtExprVisitor, stmt_expr_visitor
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
f_visit_stmt: Callable | None = None,
|
|
f_visit_expr: Callable | None = None,
|
|
# Stmt
|
|
f_visit_bind: Callable | None = None,
|
|
f_visit_attr_stmt: Callable | None = None,
|
|
f_visit_if_then_else: Callable | None = None,
|
|
f_visit_for: Callable | None = None,
|
|
f_visit_while: Callable | None = None,
|
|
f_visit_alloc_buffer: Callable | None = None,
|
|
f_visit_decl_buffer: Callable | None = None,
|
|
f_visit_buffer_store: Callable | None = None,
|
|
f_visit_assert_stmt: Callable | None = None,
|
|
f_visit_seq_stmt: Callable | None = None,
|
|
f_visit_evaluate: Callable | None = None,
|
|
f_visit_block: Callable | None = None,
|
|
f_visit_sblock_realize: Callable | None = None,
|
|
# Expr
|
|
f_visit_var: Callable | None = None,
|
|
f_visit_buffer_load: Callable | None = None,
|
|
f_visit_producer_load: Callable | None = None,
|
|
f_visit_let: Callable | None = None,
|
|
f_visit_call: Callable | None = None,
|
|
f_visit_add: Callable | None = None,
|
|
f_visit_sub: Callable | None = None,
|
|
f_visit_mul: Callable | None = None,
|
|
f_visit_div: Callable | None = None,
|
|
f_visit_mod: Callable | None = None,
|
|
f_visit_floor_div: Callable | None = None,
|
|
f_visit_floor_mod: Callable | None = None,
|
|
f_visit_min: Callable | None = None,
|
|
f_visit_max: Callable | None = None,
|
|
f_visit_eq: Callable | None = None,
|
|
f_visit_ne: Callable | None = None,
|
|
f_visit_lt: Callable | None = None,
|
|
f_visit_le: Callable | None = None,
|
|
f_visit_gt: Callable | None = None,
|
|
f_visit_ge: Callable | None = None,
|
|
f_visit_and: Callable | None = None,
|
|
f_visit_or: Callable | None = None,
|
|
f_visit_reduce: Callable | None = None,
|
|
f_visit_cast: Callable | None = None,
|
|
f_visit_not: Callable | None = None,
|
|
f_visit_select: Callable | None = None,
|
|
f_visit_ramp: Callable | None = None,
|
|
f_visit_broadcast: Callable | None = None,
|
|
f_visit_shuffle: Callable | None = None,
|
|
f_visit_int_imm: Callable | None = None,
|
|
f_visit_float_imm: Callable | None = None,
|
|
f_visit_string_imm: Callable | None = None,
|
|
) -> None:
|
|
"""Constructor."""
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.MakePyStmtExprVisitor, # type: ignore
|
|
f_visit_stmt,
|
|
f_visit_expr,
|
|
# Stmt
|
|
f_visit_bind,
|
|
f_visit_attr_stmt,
|
|
f_visit_if_then_else,
|
|
f_visit_for,
|
|
f_visit_while,
|
|
f_visit_alloc_buffer,
|
|
f_visit_decl_buffer,
|
|
f_visit_buffer_store,
|
|
f_visit_assert_stmt,
|
|
f_visit_seq_stmt,
|
|
f_visit_evaluate,
|
|
f_visit_block,
|
|
f_visit_sblock_realize,
|
|
# Expr
|
|
f_visit_var,
|
|
f_visit_buffer_load,
|
|
f_visit_producer_load,
|
|
f_visit_let,
|
|
f_visit_call,
|
|
f_visit_add,
|
|
f_visit_sub,
|
|
f_visit_mul,
|
|
f_visit_div,
|
|
f_visit_mod,
|
|
f_visit_floor_div,
|
|
f_visit_floor_mod,
|
|
f_visit_min,
|
|
f_visit_max,
|
|
f_visit_eq,
|
|
f_visit_ne,
|
|
f_visit_lt,
|
|
f_visit_le,
|
|
f_visit_gt,
|
|
f_visit_ge,
|
|
f_visit_and,
|
|
f_visit_or,
|
|
f_visit_reduce,
|
|
f_visit_cast,
|
|
f_visit_not,
|
|
f_visit_select,
|
|
f_visit_ramp,
|
|
f_visit_broadcast,
|
|
f_visit_shuffle,
|
|
f_visit_int_imm,
|
|
f_visit_float_imm,
|
|
f_visit_string_imm,
|
|
)
|
|
|
|
|
|
class PyStmtExprVisitor:
|
|
"""
|
|
A Python StmtExprVisitor to define custom visitor for both Stmt and Expr.
|
|
|
|
Users can customize any of the visit function.
|
|
"""
|
|
|
|
_tvm_metadata = {
|
|
"cls": _PyStmtExprVisitor,
|
|
"methods": [
|
|
"visit_stmt",
|
|
"visit_expr",
|
|
# Stmt
|
|
"visit_bind_",
|
|
"visit_attr_stmt_",
|
|
"visit_if_then_else_",
|
|
"visit_for_",
|
|
"visit_while_",
|
|
"visit_alloc_buffer_",
|
|
"visit_decl_buffer_",
|
|
"visit_buffer_store_",
|
|
"visit_assert_stmt_",
|
|
"visit_seq_stmt_",
|
|
"visit_evaluate_",
|
|
"visit_sblock_",
|
|
"visit_sblock_realize_",
|
|
# Expr
|
|
"visit_var_",
|
|
"visit_buffer_load_",
|
|
"visit_producer_load_",
|
|
"visit_let_",
|
|
"visit_call_",
|
|
"visit_add_",
|
|
"visit_sub_",
|
|
"visit_mul_",
|
|
"visit_div_",
|
|
"visit_mod_",
|
|
"visit_floor_div_",
|
|
"visit_floor_mod_",
|
|
"visit_min_",
|
|
"visit_max_",
|
|
"visit_eq_",
|
|
"visit_ne_",
|
|
"visit_lt_",
|
|
"visit_le_",
|
|
"visit_gt_",
|
|
"visit_ge_",
|
|
"visit_and_",
|
|
"visit_or_",
|
|
"visit_reduce_",
|
|
"visit_cast_",
|
|
"visit_not_",
|
|
"visit_select_",
|
|
"visit_ramp_",
|
|
"visit_broadcast_",
|
|
"visit_shuffle_",
|
|
"visit_int_imm_",
|
|
"visit_float_imm_",
|
|
"visit_string_imm_",
|
|
],
|
|
}
|
|
|
|
def visit_stmt(self, stmt: Stmt) -> None:
|
|
"""Visit a Stmt.
|
|
|
|
Parameters
|
|
----------
|
|
stmt : Stmt
|
|
The Stmt to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorVisitStmt(self._outer(), stmt) # type: ignore
|
|
|
|
def visit_expr(self, expr: Expr) -> None:
|
|
"""Visit a Expr.
|
|
|
|
Parameters
|
|
----------
|
|
expr : Expr
|
|
The Expr to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorVisitExpr(self._outer(), expr) # type: ignore
|
|
|
|
def visit_attr_stmt_(self, op: AttrStmt) -> None:
|
|
"""Visit AttrStmt.
|
|
Users can customize this function to overwrite VisitStmt_(const AttrStmtNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : AttrStmt
|
|
The AttrStmt to be visited.
|
|
"""
|
|
print("visit_attr_stmt_", op)
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitStmt(self._outer(), op) # type: ignore
|
|
|
|
def visit_if_then_else_(self, op: IfThenElse) -> None:
|
|
"""Visit IfThenElse.
|
|
Users can customize this function to overwrite VisitStmt_(const IfThenElseNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : IfThenElse
|
|
The IfThenElse to be visited.
|
|
"""
|
|
print("visit_if_then_else_", op)
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitStmt(self._outer(), op) # type: ignore
|
|
|
|
def visit_bind_(self, op: Bind) -> None:
|
|
"""Visit Bind.
|
|
Users can customize this function to overwrite VisitStmt_(const BindNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Bind
|
|
The Bind node to be visited.
|
|
"""
|
|
print("visit_bind_", op)
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitStmt(self._outer(), op) # type: ignore
|
|
|
|
def visit_for_(self, op: For) -> None:
|
|
"""Visit For.
|
|
Users can customize this function to overwrite VisitStmt_(const ForNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : For
|
|
The For to be visited.
|
|
"""
|
|
print("visit_for_", op)
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitStmt(self._outer(), op) # type: ignore
|
|
|
|
def visit_while_(self, op: While) -> None:
|
|
"""Visit While.
|
|
Users can customize this function to overwrite VisitStmt_(const WhileNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : While
|
|
The While to be visited.
|
|
"""
|
|
print("visit_while_", op)
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitStmt(self._outer(), op) # type: ignore
|
|
|
|
def visit_alloc_buffer_(self, op: AllocBuffer) -> None:
|
|
"""Visit AllocBuffer.
|
|
Users can customize this function to overwrite VisitStmt_(const AllocBufferNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : AllocBuffer
|
|
The AllocBuffer to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitStmt(self._outer(), op) # type: ignore
|
|
|
|
def visit_decl_buffer_(self, op: DeclBuffer) -> None:
|
|
"""Visit DeclBuffer.
|
|
Users can customize this function to overwrite VisitStmt_(const DeclBufferNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : DeclBuffer
|
|
The DeclBuffer to be visited.
|
|
"""
|
|
print("visit_decl_buffer_", op)
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitStmt(self._outer(), op) # type: ignore
|
|
|
|
def visit_buffer_store_(self, op: BufferStore) -> None:
|
|
"""Visit BufferStore.
|
|
Users can customize this function to overwrite VisitStmt_(const BufferStoreNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : BufferStore
|
|
The BufferStore to be visited.
|
|
"""
|
|
print("visit_buffer_store_", op)
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitStmt(self._outer(), op) # type: ignore
|
|
|
|
def visit_assert_stmt_(self, op: AssertStmt) -> None:
|
|
"""Visit AssertStmt.
|
|
Users can customize this function to overwrite VisitStmt_(const AssertStmtNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : AssertStmt
|
|
The AssertStmt to be visited.
|
|
"""
|
|
print("visit_assert_stmt_", op)
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitStmt(self._outer(), op) # type: ignore
|
|
|
|
def visit_seq_stmt_(self, op: SeqStmt) -> None:
|
|
"""Visit SeqStmt.
|
|
Users can customize this function to overwrite VisitStmt_(const SeqStmtNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : SeqStmt
|
|
The SeqStmt to be visited.
|
|
"""
|
|
print("visit_seq_stmt_", op)
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitStmt(self._outer(), op) # type: ignore
|
|
|
|
def visit_evaluate_(self, op: Evaluate) -> None:
|
|
"""Visit Evaluate.
|
|
Users can customize this function to overwrite VisitStmt_(const EvaluateNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Evaluate
|
|
The Evaluate to be visited.
|
|
"""
|
|
print("visit_evaluate_", op)
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitStmt(self._outer(), op) # type: ignore
|
|
|
|
def visit_sblock_(self, op: SBlock) -> None:
|
|
"""Visit SBlock.
|
|
Users can customize this function to overwrite VisitStmt_(const SBlockNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : SBlock
|
|
The SBlock to be visited.
|
|
"""
|
|
print("visit_sblock_", op)
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitStmt(self._outer(), op) # type: ignore
|
|
|
|
def visit_sblock_realize_(self, op: SBlockRealize) -> None:
|
|
"""Visit BlockRealize.
|
|
Users can customize this function to overwrite VisitStmt_(const SBlockRealizeNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : SBlockRealize
|
|
The BlockRealize to be visited.
|
|
"""
|
|
print("visit_sblock_realize_", op)
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitStmt(self._outer(), op) # type: ignore
|
|
|
|
def visit_var_(self, op: Var) -> None:
|
|
"""Visit Var.
|
|
|
|
Users can customize this function to overwrite VisitVar_(const VarNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Var
|
|
The Var to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_buffer_load_(self, op: BufferLoad) -> None:
|
|
"""Visit BufferLoad.
|
|
|
|
Users can customize this function to overwrite VisitBufferLoad_(const BufferLoadNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : BufferLoad
|
|
The BufferLoad to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_producer_load_(self, op: ProducerLoad) -> None:
|
|
"""Visit ProducerLoad.
|
|
|
|
Users can customize this function to overwrite
|
|
VisitProducerLoad_(const ProducerLoadNode* op) on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : ProducerLoad
|
|
The ProducerLoad to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_let_(self, op: Let) -> None:
|
|
"""Visit Let.
|
|
|
|
Users can customize this function to overwrite VisitLet_(const LetNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Let
|
|
The Let to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_call_(self, op: Call) -> None:
|
|
"""Visit Call.
|
|
|
|
Users can customize this function to overwrite VisitCall_(const CallNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Call
|
|
The Call to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_add_(self, op: Add) -> None:
|
|
"""Visit Add.
|
|
|
|
Users can customize this function to overwrite VisitAdd_(const AddNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Add
|
|
The Add to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_sub_(self, op: Sub) -> None:
|
|
"""Visit Sub.
|
|
|
|
Users can customize this function to overwrite VisitSub_(const SubNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Sub
|
|
The Sub to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_mul_(self, op: Mul) -> None:
|
|
"""Visit Mul.
|
|
|
|
Users can customize this function to overwrite VisitMul_(const MulNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Mul
|
|
The Mul to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_div_(self, op: Div) -> None:
|
|
"""Visit Div.
|
|
|
|
Users can customize this function to overwrite VisitDiv_(const DivNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Div
|
|
The Div to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_mod_(self, op: Mod) -> None:
|
|
"""Visit Mod.
|
|
|
|
Users can customize this function to overwrite VisitMod_(const ModNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Mod
|
|
The Mod to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_floor_div_(self, op: FloorDiv) -> None:
|
|
"""Visit FloorDiv.
|
|
|
|
Users can customize this function to overwrite VisitFloorDiv_(const FloorDivNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : FloorDiv
|
|
The FloorDiv to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_floor_mod_(self, op: FloorMod) -> None:
|
|
"""Visit FloorMod.
|
|
|
|
Users can customize this function to overwrite VisitFloorMod_(const FloorModNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : FloorMod
|
|
The FloorMod to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_min_(self, op: Min) -> None:
|
|
"""Visit Min.
|
|
|
|
Users can customize this function to overwrite VisitMin_(const MinNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Min
|
|
The Min to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_max_(self, op: Max) -> None:
|
|
"""Visit Max.
|
|
|
|
Users can customize this function to overwrite VisitMax_(const MaxNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Max
|
|
The Max to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_eq_(self, op: EQ) -> None:
|
|
"""Visit EQ.
|
|
|
|
Users can customize this function to overwrite VisitEQ_(const EQNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : EQ
|
|
The EQ to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_ne_(self, op: NE) -> None:
|
|
"""Visit NE.
|
|
|
|
Users can customize this function to overwrite VisitNE_(const NENode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : NE
|
|
The NE to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_lt_(self, op: LT) -> None:
|
|
"""Visit LT.
|
|
|
|
Users can customize this function to overwrite VisitLT_(const LTNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : LT
|
|
The LT to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_le_(self, op: LE) -> None:
|
|
"""Visit LE.
|
|
|
|
Users can customize this function to overwrite VisitLE_(const LENode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : LE
|
|
The LE to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_gt_(self, op: GT) -> None:
|
|
"""Visit GT.
|
|
|
|
Users can customize this function to overwrite VisitGT_(const GTNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : GT
|
|
The GT to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_ge_(self, op: GE) -> None:
|
|
"""Visit GE.
|
|
|
|
Users can customize this function to overwrite VisitGE_(const GENode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : GE
|
|
The GE to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_and_(self, op: And) -> None:
|
|
"""Visit And.
|
|
|
|
Users can customize this function to overwrite VisitAnd_(const AndNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : And
|
|
The And to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_or_(self, op: Or) -> None:
|
|
"""Visit Or.
|
|
|
|
Users can customize this function to overwrite VisitOr_(const OrNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Or
|
|
The Or to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_reduce_(self, op: Reduce) -> None:
|
|
"""Visit Reduce.
|
|
|
|
Users can customize this function to overwrite VisitReduce_(const ReduceNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Reduce
|
|
The Reduce to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_cast_(self, op: Cast) -> None:
|
|
"""Visit Cast.
|
|
|
|
Users can customize this function to overwrite VisitCast_(const CastNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Cast
|
|
The Cast to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_not_(self, op: Not) -> None:
|
|
"""Visit Not.
|
|
|
|
Users can customize this function to overwrite VisitNot_(const NotNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Not
|
|
The Not to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_select_(self, op: Select) -> None:
|
|
"""Visit Select.
|
|
|
|
Users can customize this function to overwrite VisitSelect_(const SelectNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Select
|
|
The Select to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_ramp_(self, op: Ramp) -> None:
|
|
"""Visit Ramp.
|
|
|
|
Users can customize this function to overwrite VisitRamp_(const RampNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Ramp
|
|
The Ramp to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_broadcast_(self, op: Broadcast) -> None:
|
|
"""Visit Broadcast.
|
|
|
|
Users can customize this function to overwrite VisitBroadcast_(const BroadcastNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Broadcast
|
|
The Broadcast to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_shuffle_(self, op: Shuffle) -> None:
|
|
"""Visit Shuffle.
|
|
|
|
Users can customize this function to overwrite VisitShuffle_(const ShuffleNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Shuffle
|
|
The Shuffle to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_int_imm_(self, op: IntImm) -> None:
|
|
"""Visit IntImm.
|
|
|
|
Users can customize this function to overwrite VisitIntImm_(const IntImmNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : IntImm
|
|
The IntImm to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_float_imm_(self, op: FloatImm) -> None:
|
|
"""Visit FloatImm.
|
|
|
|
Users can customize this function to overwrite VisitFloatImm_(const FloatImmNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : FloatImm
|
|
The FloatImm to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_string_imm_(self, op: StringImm) -> None:
|
|
"""Visit StringImm.
|
|
|
|
Users can customize this function to overwrite VisitStringImm_(const StringImmNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : StringImm
|
|
The StringImm to be visited.
|
|
"""
|
|
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.PyStmtExprMutator")
|
|
class _PyStmtExprMutator(tvm_ffi.core.Object):
|
|
"""
|
|
A TVM object to support customization of StmtExprMutator on the python side.
|
|
This is the decorated result returned from stmt_expr_mutator decorator.
|
|
|
|
WARNING: This is NOT the user facing class for method overwriting inheritance.
|
|
|
|
See also: stmt_expr_mutator, PyStmtExprMutator
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
f_visit_stmt: Callable | None = None,
|
|
f_visit_expr: Callable | None = None,
|
|
# Stmt
|
|
f_visit_bind: Callable | None = None,
|
|
f_visit_attr_stmt: Callable | None = None,
|
|
f_visit_if_then_else: Callable | None = None,
|
|
f_visit_for: Callable | None = None,
|
|
f_visit_while: Callable | None = None,
|
|
f_visit_alloc_buffer: Callable | None = None,
|
|
f_visit_decl_buffer: Callable | None = None,
|
|
f_visit_buffer_store: Callable | None = None,
|
|
f_visit_assert_stmt: Callable | None = None,
|
|
f_visit_seq_stmt: Callable | None = None,
|
|
f_visit_evaluate: Callable | None = None,
|
|
f_visit_block: Callable | None = None,
|
|
f_visit_sblock_realize: Callable | None = None,
|
|
# Expr
|
|
f_visit_var: Callable | None = None,
|
|
f_visit_buffer_load: Callable | None = None,
|
|
f_visit_producer_load: Callable | None = None,
|
|
f_visit_let: Callable | None = None,
|
|
f_visit_call: Callable | None = None,
|
|
f_visit_add: Callable | None = None,
|
|
f_visit_sub: Callable | None = None,
|
|
f_visit_mul: Callable | None = None,
|
|
f_visit_div: Callable | None = None,
|
|
f_visit_mod: Callable | None = None,
|
|
f_visit_floor_div: Callable | None = None,
|
|
f_visit_floor_mod: Callable | None = None,
|
|
f_visit_min: Callable | None = None,
|
|
f_visit_max: Callable | None = None,
|
|
f_visit_eq: Callable | None = None,
|
|
f_visit_ne: Callable | None = None,
|
|
f_visit_lt: Callable | None = None,
|
|
f_visit_le: Callable | None = None,
|
|
f_visit_gt: Callable | None = None,
|
|
f_visit_ge: Callable | None = None,
|
|
f_visit_and: Callable | None = None,
|
|
f_visit_or: Callable | None = None,
|
|
f_visit_reduce: Callable | None = None,
|
|
f_visit_cast: Callable | None = None,
|
|
f_visit_not: Callable | None = None,
|
|
f_visit_select: Callable | None = None,
|
|
f_visit_ramp: Callable | None = None,
|
|
f_visit_broadcast: Callable | None = None,
|
|
f_visit_shuffle: Callable | None = None,
|
|
f_visit_int_imm: Callable | None = None,
|
|
f_visit_float_imm: Callable | None = None,
|
|
f_visit_string_imm: Callable | None = None,
|
|
) -> None:
|
|
"""Constructor."""
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.MakePyStmtExprMutator, # type: ignore
|
|
f_visit_stmt,
|
|
f_visit_expr,
|
|
# Stmt
|
|
f_visit_bind,
|
|
f_visit_attr_stmt,
|
|
f_visit_if_then_else,
|
|
f_visit_for,
|
|
f_visit_while,
|
|
f_visit_alloc_buffer,
|
|
f_visit_decl_buffer,
|
|
f_visit_buffer_store,
|
|
f_visit_assert_stmt,
|
|
f_visit_seq_stmt,
|
|
f_visit_evaluate,
|
|
f_visit_block,
|
|
f_visit_sblock_realize,
|
|
# Expr
|
|
f_visit_var,
|
|
f_visit_buffer_load,
|
|
f_visit_producer_load,
|
|
f_visit_let,
|
|
f_visit_call,
|
|
f_visit_add,
|
|
f_visit_sub,
|
|
f_visit_mul,
|
|
f_visit_div,
|
|
f_visit_mod,
|
|
f_visit_floor_div,
|
|
f_visit_floor_mod,
|
|
f_visit_min,
|
|
f_visit_max,
|
|
f_visit_eq,
|
|
f_visit_ne,
|
|
f_visit_lt,
|
|
f_visit_le,
|
|
f_visit_gt,
|
|
f_visit_ge,
|
|
f_visit_and,
|
|
f_visit_or,
|
|
f_visit_reduce,
|
|
f_visit_cast,
|
|
f_visit_not,
|
|
f_visit_select,
|
|
f_visit_ramp,
|
|
f_visit_broadcast,
|
|
f_visit_shuffle,
|
|
f_visit_int_imm,
|
|
f_visit_float_imm,
|
|
f_visit_string_imm,
|
|
)
|
|
|
|
|
|
class PyStmtExprMutator:
|
|
"""
|
|
A Python StmtExprMutator to define custom mutator for both Stmt and Expr.
|
|
|
|
Users can customize any of the visit function.
|
|
"""
|
|
|
|
_tvm_metadata = {
|
|
"cls": _PyStmtExprMutator,
|
|
"methods": [
|
|
"visit_stmt",
|
|
"visit_expr",
|
|
# Stmt
|
|
"visit_bind_",
|
|
"visit_attr_stmt_",
|
|
"visit_if_then_else_",
|
|
"visit_for_",
|
|
"visit_while_",
|
|
"visit_alloc_buffer_",
|
|
"visit_decl_buffer_",
|
|
"visit_buffer_store_",
|
|
"visit_assert_stmt_",
|
|
"visit_seq_stmt_",
|
|
"visit_evaluate_",
|
|
"visit_sblock_",
|
|
"visit_sblock_realize_",
|
|
# Expr
|
|
"visit_var_",
|
|
"visit_buffer_load_",
|
|
"visit_producer_load_",
|
|
"visit_let_",
|
|
"visit_call_",
|
|
"visit_add_",
|
|
"visit_sub_",
|
|
"visit_mul_",
|
|
"visit_div_",
|
|
"visit_mod_",
|
|
"visit_floor_div_",
|
|
"visit_floor_mod_",
|
|
"visit_min_",
|
|
"visit_max_",
|
|
"visit_eq_",
|
|
"visit_ne_",
|
|
"visit_lt_",
|
|
"visit_le_",
|
|
"visit_gt_",
|
|
"visit_ge_",
|
|
"visit_and_",
|
|
"visit_or_",
|
|
"visit_reduce_",
|
|
"visit_cast_",
|
|
"visit_not_",
|
|
"visit_select_",
|
|
"visit_ramp_",
|
|
"visit_broadcast_",
|
|
"visit_shuffle_",
|
|
"visit_int_imm_",
|
|
"visit_float_imm_",
|
|
"visit_string_imm_",
|
|
],
|
|
}
|
|
|
|
def visit_expr(self, expr: Expr) -> Expr:
|
|
"""Visit Expr.
|
|
Users can customize this function to overwrite VisitExpr(const Expr& expr)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
expr : Expr
|
|
The Expr to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorVisitExpr(self._outer(), expr) # type: ignore
|
|
|
|
def visit_stmt(self, stmt: Stmt) -> Stmt:
|
|
"""Visit Stmt.
|
|
Users can customize this function to overwrite VisitStmt(const Stmt& stmt)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
stmt : Stmt
|
|
The Stmt to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Stmt
|
|
The mutated Stmt.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorVisitStmt(self._outer(), stmt) # type: ignore
|
|
|
|
def visit_attr_stmt_(self, op: AttrStmt) -> Stmt:
|
|
"""Visit AttrStmt.
|
|
Users can customize this function to overwrite VisitStmt_(const AttrStmtNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : AttrStmt
|
|
The AttrStmt to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Stmt
|
|
The mutated Stmt.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitStmt(self._outer(), op) # type: ignore
|
|
|
|
def visit_if_then_else_(self, op: IfThenElse) -> Stmt:
|
|
"""Visit IfThenElse.
|
|
Users can customize this function to overwrite VisitStmt_(const IfThenElseNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : IfThenElse
|
|
The IfThenElse to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Stmt
|
|
The mutated Stmt.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitStmt(self._outer(), op) # type: ignore
|
|
|
|
def visit_bind_(self, op: Bind) -> Stmt:
|
|
"""Visit Bind.
|
|
Users can customize this function to overwrite VisitStmt_(const BindNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Bind
|
|
The Bind node to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Stmt
|
|
The mutated Stmt.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitStmt(self._outer(), op) # type: ignore
|
|
|
|
def visit_for_(self, op: For) -> Stmt:
|
|
"""Visit For.
|
|
Users can customize this function to overwrite VisitStmt_(const ForNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : For
|
|
The For to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Stmt
|
|
The mutated Stmt.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitStmt(self._outer(), op) # type: ignore
|
|
|
|
def visit_while_(self, op: While) -> Stmt:
|
|
"""Visit While.
|
|
Users can customize this function to overwrite VisitStmt_(const WhileNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : While
|
|
The While to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Stmt
|
|
The mutated Stmt.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitStmt(self._outer(), op) # type: ignore
|
|
|
|
def visit_alloc_buffer_(self, op: AllocBuffer) -> Stmt:
|
|
"""Visit AllocBuffer.
|
|
Users can customize this function to overwrite VisitStmt_(const AllocBufferNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : AllocBuffer
|
|
The AllocBuffer to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Stmt
|
|
The mutated Stmt.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitStmt(self._outer(), op) # type: ignore
|
|
|
|
def visit_decl_buffer_(self, op: DeclBuffer) -> Stmt:
|
|
"""Visit DeclBuffer.
|
|
Users can customize this function to overwrite VisitStmt_(const DeclBufferNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : DeclBuffer
|
|
The DeclBuffer to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Stmt
|
|
The mutated Stmt.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitStmt(self._outer(), op) # type: ignore
|
|
|
|
def visit_buffer_store_(self, op: BufferStore) -> Stmt:
|
|
"""Visit BufferStore.
|
|
Users can customize this function to overwrite VisitStmt_(const BufferStoreNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : BufferStore
|
|
The BufferStore to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Stmt
|
|
The mutated Stmt.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitStmt(self._outer(), op) # type: ignore
|
|
|
|
def visit_assert_stmt_(self, op: AssertStmt) -> Stmt:
|
|
"""Visit AssertStmt.
|
|
Users can customize this function to overwrite VisitStmt_(const AssertStmtNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : AssertStmt
|
|
The AssertStmt to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Stmt
|
|
The mutated Stmt.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitStmt(self._outer(), op) # type: ignore
|
|
|
|
def visit_seq_stmt_(self, op: SeqStmt) -> Stmt:
|
|
"""Visit SeqStmt.
|
|
Users can customize this function to overwrite VisitStmt_(const SeqStmtNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : SeqStmt
|
|
The SeqStmt to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Stmt
|
|
The mutated Stmt.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitStmt(self._outer(), op) # type: ignore
|
|
|
|
def visit_evaluate_(self, op: Evaluate) -> Stmt:
|
|
"""Visit Evaluate.
|
|
Users can customize this function to overwrite VisitStmt_(const EvaluateNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Evaluate
|
|
The Evaluate to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Stmt
|
|
The mutated Stmt.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitStmt(self._outer(), op) # type: ignore
|
|
|
|
def visit_sblock_(self, op: SBlock) -> Stmt:
|
|
"""Visit SBlock.
|
|
Users can customize this function to overwrite VisitStmt_(const SBlockNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : SBlock
|
|
The SBlock to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Stmt
|
|
The mutated Stmt.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitStmt(self._outer(), op) # type: ignore
|
|
|
|
def visit_sblock_realize_(self, op: SBlockRealize) -> Stmt:
|
|
"""Visit BlockRealize.
|
|
Users can customize this function to overwrite VisitStmt_(const SBlockRealizeNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : SBlockRealize
|
|
The SBlockRealize to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Stmt
|
|
The mutated Stmt.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitStmt(self._outer(), op) # type: ignore
|
|
|
|
def visit_var_(self, op: Var) -> Expr:
|
|
"""Visit Var.
|
|
|
|
Users can customize this function to overwrite VisitVar_(const VarNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Var
|
|
The Var to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_buffer_load_(self, op: BufferLoad) -> Expr:
|
|
"""Visit BufferLoad.
|
|
|
|
Users can customize this function to overwrite VisitBufferLoad_(const BufferLoadNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : BufferLoad
|
|
The BufferLoad to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_producer_load_(self, op: ProducerLoad) -> Expr:
|
|
"""Visit ProducerLoad.
|
|
|
|
Users can customize this function to overwrite
|
|
VisitProducerLoad_(const ProducerLoadNode* op) on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : ProducerLoad
|
|
The ProducerLoad to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_let_(self, op: Let) -> Expr:
|
|
"""Visit Let.
|
|
|
|
Users can customize this function to overwrite VisitLet_(const LetNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Let
|
|
The Let to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_call_(self, op: Call) -> Expr:
|
|
"""Visit Call.
|
|
|
|
Users can customize this function to overwrite VisitCall_(const CallNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Call
|
|
The Call to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_add_(self, op: Add) -> Expr:
|
|
"""Visit Add.
|
|
|
|
Users can customize this function to overwrite VisitAdd_(const AddNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Add
|
|
The Add to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_sub_(self, op: Sub) -> Expr:
|
|
"""Visit Sub.
|
|
|
|
Users can customize this function to overwrite VisitSub_(const SubNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Sub
|
|
The Sub to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_mul_(self, op: Mul) -> Expr:
|
|
"""Visit Mul.
|
|
|
|
Users can customize this function to overwrite VisitMul_(const MulNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Mul
|
|
The Mul to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_div_(self, op: Div) -> Expr:
|
|
"""Visit Div.
|
|
|
|
Users can customize this function to overwrite VisitDiv_(const DivNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Div
|
|
The Div to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_mod_(self, op: Mod) -> Expr:
|
|
"""Visit Mod.
|
|
|
|
Users can customize this function to overwrite VisitMod_(const ModNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Mod
|
|
The Mod to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_floor_div_(self, op: FloorDiv) -> Expr:
|
|
"""Visit FloorDiv.
|
|
|
|
Users can customize this function to overwrite VisitFloorDiv_(const FloorDivNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : FloorDiv
|
|
The FloorDiv to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_floor_mod_(self, op: FloorMod) -> Expr:
|
|
"""Visit FloorMod.
|
|
|
|
Users can customize this function to overwrite VisitFloorMod_(const FloorModNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : FloorMod
|
|
The FloorMod to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_min_(self, op: Min) -> Expr:
|
|
"""Visit Min.
|
|
|
|
Users can customize this function to overwrite VisitMin_(const MinNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Min
|
|
The Min to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_max_(self, op: Max) -> Expr:
|
|
"""Visit Max.
|
|
|
|
Users can customize this function to overwrite VisitMax_(const MaxNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Max
|
|
The Max to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_eq_(self, op: EQ) -> Expr:
|
|
"""Visit EQ.
|
|
|
|
Users can customize this function to overwrite VisitEQ_(const EQNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : EQ
|
|
The EQ to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_ne_(self, op: NE) -> Expr:
|
|
"""Visit NE.
|
|
|
|
Users can customize this function to overwrite VisitNE_(const NENode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : NE
|
|
The NE to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_lt_(self, op: LT) -> Expr:
|
|
"""Visit LT.
|
|
|
|
Users can customize this function to overwrite VisitLT_(const LTNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : LT
|
|
The LT to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_le_(self, op: LE) -> Expr:
|
|
"""Visit LE.
|
|
|
|
Users can customize this function to overwrite VisitLE_(const LENode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : LE
|
|
The LE to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_gt_(self, op: GT) -> Expr:
|
|
"""Visit GT.
|
|
|
|
Users can customize this function to overwrite VisitGT_(const GTNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : GT
|
|
The GT to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_ge_(self, op: GE) -> Expr:
|
|
"""Visit GE.
|
|
|
|
Users can customize this function to overwrite VisitGE_(const GENode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : GE
|
|
The GE to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_and_(self, op: And) -> Expr:
|
|
"""Visit And.
|
|
|
|
Users can customize this function to overwrite VisitAnd_(const AndNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : And
|
|
The And to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_or_(self, op: Or) -> Expr:
|
|
"""Visit Or.
|
|
|
|
Users can customize this function to overwrite VisitOr_(const OrNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Or
|
|
The Or to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_reduce_(self, op: Reduce) -> Expr:
|
|
"""Visit Reduce.
|
|
|
|
Users can customize this function to overwrite VisitReduce_(const ReduceNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Reduce
|
|
The Reduce to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_cast_(self, op: Cast) -> Expr:
|
|
"""Visit Cast.
|
|
|
|
Users can customize this function to overwrite VisitCast_(const CastNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Cast
|
|
The Cast to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_not_(self, op: Not) -> Expr:
|
|
"""Visit Not.
|
|
|
|
Users can customize this function to overwrite VisitNot_(const NotNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Not
|
|
The Not to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_select_(self, op: Select) -> Expr:
|
|
"""Visit Select.
|
|
|
|
Users can customize this function to overwrite VisitSelect_(const SelectNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Select
|
|
The Select to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_ramp_(self, op: Ramp) -> Expr:
|
|
"""Visit Ramp.
|
|
|
|
Users can customize this function to overwrite VisitRamp_(const RampNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Ramp
|
|
The Ramp to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_broadcast_(self, op: Broadcast) -> Expr:
|
|
"""Visit Broadcast.
|
|
|
|
Users can customize this function to overwrite VisitBroadcast_(const BroadcastNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Broadcast
|
|
The Broadcast to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_shuffle_(self, op: Shuffle) -> Expr:
|
|
"""Visit Shuffle.
|
|
|
|
Users can customize this function to overwrite VisitShuffle_(const ShuffleNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : Shuffle
|
|
The Shuffle to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_int_imm_(self, op: IntImm) -> Expr:
|
|
"""Visit IntImm.
|
|
|
|
Users can customize this function to overwrite VisitIntImm_(const IntImmNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : IntImm
|
|
The IntImm to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_float_imm_(self, op: FloatImm) -> Expr:
|
|
"""Visit FloatImm.
|
|
|
|
Users can customize this function to overwrite VisitFloatImm_(const FloatImmNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : FloatImm
|
|
The FloatImm to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op) # type: ignore
|
|
|
|
def visit_string_imm_(self, op: StringImm) -> Expr:
|
|
"""Visit StringImm.
|
|
|
|
Users can customize this function to overwrite VisitStringImm_(const StringImmNode* op)
|
|
on the C++ side.
|
|
|
|
Parameters
|
|
----------
|
|
op : StringImm
|
|
The StringImm to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated Expr.
|
|
"""
|
|
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op) # type: ignore
|