wehub-resource-sync
1062 lines
34 KiB
Python
1062 lines
34 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.
|
|
"""Statement functor utilities for IR transformations"""
|
|
|
|
from typing import TypeVar
|
|
|
|
import tvm
|
|
from tvm.ir import Range
|
|
|
|
from . import _ffi_api
|
|
from .expr_functor import ExprMutator, ExprVisitor, _visit_array
|
|
from .function import PrimFunc
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
class StmtFunctor:
|
|
"""An abstract visitor over Statement, with visiting functions defined for each Stmt type."""
|
|
|
|
def __init__(self):
|
|
self._dispatch_map = {
|
|
"tirx.Bind": self.visit_bind_,
|
|
"tirx.AttrStmt": self.visit_attr_,
|
|
"tirx.IfThenElse": self.visit_if_then_else_,
|
|
"tirx.For": self.visit_for_,
|
|
"tirx.While": self.visit_while_,
|
|
"tirx.Break": self.visit_break_,
|
|
"tirx.Continue": self.visit_continue_,
|
|
"tirx.Allocate": self.visit_allocate_,
|
|
"tirx.AllocateConst": self.visit_allocate_const_,
|
|
"tirx.DeclBuffer": self.visit_decl_buffer_,
|
|
"tirx.BufferStore": self.visit_buffer_store_,
|
|
"tirx.BufferRealize": self.visit_buffer_realize_,
|
|
"tirx.AssertStmt": self.visit_assert_,
|
|
"tirx.ProducerStore": self.visit_producer_store_,
|
|
"tirx.ProducerRealize": self.visit_producer_realize_,
|
|
"tirx.Prefetch": self.visit_prefetch_,
|
|
"tirx.SeqStmt": self.visit_seqstmt_,
|
|
"tirx.Evaluate": self.visit_evaluate_,
|
|
"tirx.SBlock": self.visit_block_,
|
|
"tirx.SBlockRealize": self.visit_block_realize_,
|
|
"tirx.ScopeIdDefStmt": self.visit_scope_id_def_stmt_,
|
|
"tirx.TilePrimitiveCall": self.visit_op_call_,
|
|
"tirx.AllocBuffer": self.visit_alloc_buffer_,
|
|
}
|
|
|
|
def visit_stmt(self, stmt):
|
|
"""Apply the visitor to a statement.
|
|
|
|
Parameters
|
|
----------
|
|
stmt : tvm.tirx.Stmt
|
|
The statement to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Any
|
|
The result of the visit.
|
|
"""
|
|
if stmt is None:
|
|
return None
|
|
if isinstance(stmt, tvm.tirx.TilePrimitiveCall):
|
|
# subclass of TilePrimitiveCall only exists in python side
|
|
# and are not handled by dispatch map
|
|
key = "TilePrimitiveCall"
|
|
else:
|
|
key = stmt.__class__.__name__
|
|
if key.endswith("Node"):
|
|
key = key[:-4] # Remove the "Node" suffix
|
|
|
|
key = "tirx." + key
|
|
if key in self._dispatch_map:
|
|
return self._dispatch_map[key](stmt)
|
|
|
|
return self.visit_stmt_default_(stmt)
|
|
|
|
def visit_stmt_default_(self, op):
|
|
"""Default visitor implementation for statements."""
|
|
raise NotImplementedError(f"Do not have a default for {op.__class__.__name__}")
|
|
|
|
def visit_bind_(self, op):
|
|
"""Visitor for Bind nodes."""
|
|
return self.visit_stmt_default_(op)
|
|
|
|
def visit_attr_(self, op):
|
|
"""Visitor for AttrStmt nodes."""
|
|
return self.visit_stmt_default_(op)
|
|
|
|
def visit_if_then_else_(self, op):
|
|
"""Visitor for IfThenElse nodes."""
|
|
return self.visit_stmt_default_(op)
|
|
|
|
def visit_for_(self, op):
|
|
"""Visitor for For nodes."""
|
|
return self.visit_stmt_default_(op)
|
|
|
|
def visit_while_(self, op):
|
|
"""Visitor for While nodes."""
|
|
return self.visit_stmt_default_(op)
|
|
|
|
def visit_break_(self, op):
|
|
"""Visitor for Break nodes."""
|
|
return self.visit_stmt_default_(op)
|
|
|
|
def visit_continue_(self, op):
|
|
"""Visitor for Continue nodes."""
|
|
return self.visit_stmt_default_(op)
|
|
|
|
def visit_allocate_(self, op):
|
|
"""Visitor for Allocate nodes."""
|
|
return self.visit_stmt_default_(op)
|
|
|
|
def visit_allocate_const_(self, op):
|
|
"""Visitor for AllocateConst nodes."""
|
|
return self.visit_stmt_default_(op)
|
|
|
|
def visit_decl_buffer_(self, op):
|
|
"""Visitor for DeclBuffer nodes."""
|
|
return self.visit_stmt_default_(op)
|
|
|
|
def visit_buffer_store_(self, op):
|
|
"""Visitor for BufferStore nodes."""
|
|
return self.visit_stmt_default_(op)
|
|
|
|
def visit_buffer_realize_(self, op):
|
|
"""Visitor for BufferRealize nodes."""
|
|
raise ValueError("BufferRealize is not allowed")
|
|
|
|
def visit_assert_(self, op):
|
|
"""Visitor for AssertStmt nodes."""
|
|
return self.visit_stmt_default_(op)
|
|
|
|
def visit_producer_store_(self, op):
|
|
"""Visitor for ProducerStore nodes."""
|
|
raise ValueError("ProducerStore is not allowed")
|
|
|
|
def visit_producer_realize_(self, op):
|
|
"""Visitor for ProducerRealize nodes."""
|
|
raise ValueError("ProducerRealize is not allowed")
|
|
|
|
def visit_prefetch_(self, op):
|
|
"""Visitor for Prefetch nodes."""
|
|
raise ValueError("Prefetch is not allowed")
|
|
|
|
def visit_seqstmt_(self, op):
|
|
"""Visitor for SeqStmt nodes."""
|
|
return self.visit_stmt_default_(op)
|
|
|
|
def visit_evaluate_(self, op):
|
|
"""Visitor for Evaluate nodes."""
|
|
return self.visit_stmt_default_(op)
|
|
|
|
def visit_block_(self, op):
|
|
"""Visitor for Block nodes."""
|
|
return self.visit_stmt_default_(op)
|
|
|
|
def visit_block_realize_(self, op):
|
|
"""Visitor for BlockRealize nodes."""
|
|
return self.visit_stmt_default_(op)
|
|
|
|
def visit_scope_id_def_stmt_(self, op):
|
|
"""Visitor for ScopeIdDefStmt nodes."""
|
|
return self.visit_stmt_default_(op)
|
|
|
|
def visit_op_call_(self, op):
|
|
"""Visitor for TilePrimitiveCall nodes."""
|
|
return self.visit_stmt_default_(op)
|
|
|
|
def visit_buffer_region_(self, op):
|
|
"""Visitor for BufferRegion nodes."""
|
|
return self.visit_stmt_default_(op)
|
|
|
|
def visit_alloc_buffer_(self, op):
|
|
"""Visitor for AllocBuffer nodes."""
|
|
return self.visit_stmt_default_(op)
|
|
|
|
def __call__(self, stmt):
|
|
"""Call visitor on statement.
|
|
|
|
Parameters
|
|
----------
|
|
stmt : tvm.tirx.Stmt
|
|
The statement.
|
|
|
|
Returns
|
|
-------
|
|
result : Any
|
|
The result of visiting.
|
|
"""
|
|
return self.visit_stmt(stmt)
|
|
|
|
|
|
class StmtVisitor(StmtFunctor):
|
|
"""A visitor over Stmt.
|
|
|
|
This is a visitor that recursively traverses a statement. Subclasses can
|
|
override the visit methods to customize the behavior.
|
|
"""
|
|
|
|
def visit_expr(self, expr):
|
|
"""Visit expressions that occur in a statement.
|
|
|
|
This method can be overridden to implement expression
|
|
traversal in a statement visitor.
|
|
|
|
Parameters
|
|
----------
|
|
expr : Expr
|
|
The expression to be visited.
|
|
"""
|
|
pass
|
|
|
|
def visit_bind_(self, op):
|
|
"""Visitor implementation for Bind."""
|
|
self.visit_expr(op.value)
|
|
|
|
def visit_attr_(self, op):
|
|
"""Visitor implementation for AttrStmt."""
|
|
self.visit_expr(op.value)
|
|
self.visit_stmt(op.body)
|
|
|
|
def visit_if_then_else_(self, op):
|
|
"""Visitor implementation for IfThenElse."""
|
|
self.visit_expr(op.condition)
|
|
self.visit_stmt(op.then_case)
|
|
if op.else_case:
|
|
self.visit_stmt(op.else_case)
|
|
|
|
def visit_for_(self, op):
|
|
"""Visitor implementation for For."""
|
|
self.visit_expr(op.min)
|
|
self.visit_expr(op.extent)
|
|
if op.step is not None:
|
|
self.visit_expr(op.step)
|
|
self.visit_stmt(op.body)
|
|
|
|
def visit_while_(self, op):
|
|
"""Visitor implementation for While."""
|
|
self.visit_expr(op.condition)
|
|
self.visit_stmt(op.body)
|
|
|
|
def visit_break_(self, op):
|
|
"""Visitor implementation for Break."""
|
|
pass
|
|
|
|
def visit_continue_(self, op):
|
|
"""Visitor implementation for Continue."""
|
|
pass
|
|
|
|
def visit_allocate_(self, op):
|
|
"""Visitor implementation for Allocate."""
|
|
_visit_array(op.extents, lambda x: self.visit_expr(x))
|
|
self.visit_stmt(op.body)
|
|
self.visit_expr(op.condition)
|
|
|
|
def visit_allocate_const_(self, op):
|
|
"""Visitor implementation for AllocateConst."""
|
|
_visit_array(op.extents, lambda x: self.visit_expr(x))
|
|
self.visit_stmt(op.body)
|
|
|
|
def visit_decl_buffer_(self, op):
|
|
"""Visitor implementation for DeclBuffer."""
|
|
if hasattr(op, "body"):
|
|
self.visit_stmt(op.body)
|
|
return
|
|
return
|
|
|
|
def visit_buffer_store_(self, op):
|
|
"""Visitor implementation for BufferStore."""
|
|
self.visit_expr(op.value)
|
|
_visit_array(op.indices, lambda x: self.visit_expr(x))
|
|
if op.predicate is not None:
|
|
self.visit_expr(op.predicate)
|
|
|
|
def visit_assert_(self, op):
|
|
"""Visitor implementation for AssertStmt."""
|
|
self.visit_expr(op.condition)
|
|
for message_part in op.message_parts:
|
|
if tvm.ir.is_prim_expr(message_part):
|
|
self.visit_expr(message_part)
|
|
|
|
def visit_seqstmt_(self, op):
|
|
"""Visitor implementation for SeqStmt."""
|
|
_visit_array(op.seq, lambda s: self.visit_stmt(s))
|
|
|
|
def visit_evaluate_(self, op):
|
|
"""Visitor implementation for Evaluate."""
|
|
self.visit_expr(op.value)
|
|
|
|
def visit_block_(self, op):
|
|
"""Visitor implementation for Block."""
|
|
# Visit IterVars
|
|
for iter_var in op.iter_vars:
|
|
self.visit_expr(iter_var.dom.min)
|
|
self.visit_expr(iter_var.dom.extent)
|
|
|
|
# Visit buffer regions (reads and writes)
|
|
def _visit_buffer_region(buffer_region):
|
|
for r in buffer_region.region:
|
|
self.visit_expr(r.min)
|
|
self.visit_expr(r.extent)
|
|
|
|
_visit_array(op.reads, _visit_buffer_region)
|
|
_visit_array(op.writes, _visit_buffer_region)
|
|
|
|
# Visit match buffers
|
|
for match_buffer in op.match_buffers:
|
|
_visit_buffer_region(match_buffer.source)
|
|
|
|
# Visit init statement
|
|
if op.init is not None:
|
|
self.visit_stmt(op.init)
|
|
|
|
# Visit body
|
|
self.visit_stmt(op.body)
|
|
|
|
def visit_block_realize_(self, op):
|
|
"""Visitor implementation for BlockRealize."""
|
|
_visit_array(op.iter_values, lambda x: self.visit_expr(x))
|
|
self.visit_expr(op.predicate)
|
|
self.visit_stmt(op.block)
|
|
|
|
def visit_scope_id_def_stmt_(self, op):
|
|
"""Visitor implementation for ScopeIdDefStmt.
|
|
|
|
Mirrors the C++ visitor: walk extents and preferred_extents via
|
|
``visit_expr``; there is no body to recurse into (the def vars
|
|
themselves are leaves the visitor doesn't otherwise inspect).
|
|
"""
|
|
# The C++ field is named ``def``, which is a Python keyword,
|
|
# so it's accessed via ``getattr``.
|
|
sid = getattr(op, "def")
|
|
if sid.extents is not None:
|
|
for e in sid.extents:
|
|
self.visit_expr(e)
|
|
if sid.preferred_extents is not None:
|
|
for e in sid.preferred_extents:
|
|
self.visit_expr(e)
|
|
|
|
def visit_op_call_(self, op):
|
|
"""Visitor implementation for TilePrimitiveCall."""
|
|
for arg in op.args:
|
|
if tvm.ir.is_prim_expr(arg):
|
|
self.visit_expr(arg)
|
|
elif isinstance(arg, tvm.tirx.Stmt):
|
|
self.visit_stmt(arg)
|
|
elif isinstance(arg, tvm.tirx.BufferRegion):
|
|
self.visit_buffer_region_(arg)
|
|
for value in op.config.values():
|
|
if tvm.ir.is_prim_expr(value):
|
|
self.visit_expr(value)
|
|
elif isinstance(value, tvm.tirx.Stmt):
|
|
self.visit_stmt(value)
|
|
|
|
def visit_buffer_region_(self, op):
|
|
"""Visitor implementation for BufferRegion."""
|
|
|
|
def _visit_range(range):
|
|
self.visit_expr(range.min)
|
|
self.visit_expr(range.extent)
|
|
|
|
_visit_array(op.region, _visit_range)
|
|
|
|
def visit_alloc_buffer_(self, op):
|
|
"""Visitor implementation for AllocBuffer."""
|
|
if hasattr(op, "body"):
|
|
self.visit_stmt(op.body)
|
|
return
|
|
return
|
|
|
|
|
|
class StmtMutator(StmtFunctor):
|
|
"""A mutator over Stmt.
|
|
|
|
This is a mutator that recursively transforms a statement. Subclasses can
|
|
override the visit methods to customize the behavior.
|
|
"""
|
|
|
|
def visit_expr(self, expr):
|
|
"""Visit and mutate expressions that occur in a statement.
|
|
|
|
This method can be overridden to implement expression
|
|
mutation in a statement mutator.
|
|
|
|
Parameters
|
|
----------
|
|
expr : Expr
|
|
The expression to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated expression.
|
|
"""
|
|
return expr
|
|
|
|
def visit_bind_(self, op):
|
|
"""Mutator implementation for Bind."""
|
|
value = self.visit_expr(op.value)
|
|
|
|
if value is op.value:
|
|
return op
|
|
|
|
return tvm.tirx.Bind(op.var, value, op.span)
|
|
|
|
def visit_attr_(self, op):
|
|
"""Mutator implementation for AttrStmt."""
|
|
value = self.visit_expr(op.value)
|
|
body = self.visit_stmt(op.body)
|
|
|
|
if value is op.value and body is op.body:
|
|
return op
|
|
|
|
return tvm.tirx.AttrStmt(op.node, op.attr_key, value, body, op.span)
|
|
|
|
def visit_if_then_else_(self, op):
|
|
"""Mutator implementation for IfThenElse."""
|
|
condition = self.visit_expr(op.condition)
|
|
then_case = self.visit_stmt(op.then_case)
|
|
else_case = self.visit_stmt(op.else_case) if op.else_case else None
|
|
|
|
if condition is op.condition and then_case is op.then_case and else_case is op.else_case:
|
|
return op
|
|
|
|
return tvm.tirx.IfThenElse(condition, then_case, else_case, op.span)
|
|
|
|
def visit_for_(self, op):
|
|
"""Mutator implementation for For."""
|
|
min_val = self.visit_expr(op.min)
|
|
extent = self.visit_expr(op.extent)
|
|
step = self.visit_expr(op.step) if op.step is not None else None
|
|
body = self.visit_stmt(op.body)
|
|
|
|
if min_val is op.min and extent is op.extent and step is op.step and body is op.body:
|
|
return op
|
|
|
|
return tvm.tirx.For(
|
|
op.loop_var,
|
|
min_val,
|
|
extent,
|
|
op.kind,
|
|
body,
|
|
op.thread_binding,
|
|
op.annotations,
|
|
step,
|
|
op.span,
|
|
)
|
|
|
|
def visit_while_(self, op):
|
|
"""Mutator implementation for While."""
|
|
condition = self.visit_expr(op.condition)
|
|
body = self.visit_stmt(op.body)
|
|
|
|
if condition is op.condition and body is op.body:
|
|
return op
|
|
|
|
return tvm.tirx.While(condition, body, op.span)
|
|
|
|
def visit_break_(self, op):
|
|
"""Mutator implementation for Break."""
|
|
return op
|
|
|
|
def visit_continue_(self, op):
|
|
"""Mutator implementation for Continue."""
|
|
return op
|
|
|
|
def visit_allocate_(self, op):
|
|
"""Mutator implementation for Allocate."""
|
|
extents = [self.visit_expr(extent) for extent in op.extents]
|
|
body = self.visit_stmt(op.body)
|
|
condition = self.visit_expr(op.condition)
|
|
|
|
extents_changed = any(old is not new for old, new in zip(op.extents, extents))
|
|
|
|
if not extents_changed and body is op.body and condition is op.condition:
|
|
return op
|
|
|
|
return tvm.tirx.Allocate(
|
|
op.buffer_var, op.dtype, extents, condition, body, op.annotations, op.span
|
|
)
|
|
|
|
def visit_allocate_const_(self, op):
|
|
"""Mutator implementation for AllocateConst."""
|
|
extents = [self.visit_expr(extent) for extent in op.extents]
|
|
body = self.visit_stmt(op.body)
|
|
|
|
extents_changed = any(old is not new for old, new in zip(op.extents, extents))
|
|
|
|
if not extents_changed and body is op.body:
|
|
return op
|
|
|
|
# Create the data_or_idx parameter based on what's available
|
|
if op.data is not None:
|
|
data_or_idx = op.data
|
|
elif op.irmod_storage_idx is not None:
|
|
data_or_idx = op.irmod_storage_idx
|
|
else:
|
|
data_or_idx = None
|
|
|
|
return tvm.tirx.AllocateConst(
|
|
op.buffer_var, op.dtype, extents, data_or_idx, body, op.annotations, op.span
|
|
)
|
|
|
|
def visit_decl_buffer_(self, op):
|
|
"""Mutator implementation for DeclBuffer."""
|
|
if hasattr(op, "body"):
|
|
body = self.visit_stmt(op.body)
|
|
if body is op.body:
|
|
return op
|
|
return tvm.tirx.DeclBuffer(op.buffer, body, op.span)
|
|
return op
|
|
|
|
def visit_buffer_store_(self, op):
|
|
"""Mutator implementation for BufferStore."""
|
|
value = self.visit_expr(op.value)
|
|
indices = [self.visit_expr(idx) for idx in op.indices]
|
|
predicate = self.visit_expr(op.predicate) if op.predicate is not None else None
|
|
|
|
indices_changed = any(old is not new for old, new in zip(op.indices, indices))
|
|
|
|
if value is op.value and not indices_changed and predicate is op.predicate:
|
|
return op
|
|
|
|
return tvm.tirx.BufferStore(op.buffer, value, indices, predicate, op.span)
|
|
|
|
def visit_buffer_realize_(self, op):
|
|
"""Mutator implementation for BufferRealize."""
|
|
bounds = []
|
|
bounds_changed = False
|
|
|
|
for r in op.bounds:
|
|
new_min = self.visit_expr(r.min)
|
|
new_extent = self.visit_expr(r.extent)
|
|
|
|
if new_min is not r.min or new_extent is not r.extent:
|
|
bounds_changed = True
|
|
bounds.append(tvm.ir.Range(new_min, new_extent))
|
|
else:
|
|
bounds.append(r)
|
|
|
|
condition = self.visit_expr(op.condition)
|
|
body = self.visit_stmt(op.body)
|
|
|
|
if not bounds_changed and condition is op.condition and body is op.body:
|
|
return op
|
|
|
|
return tvm.tirx.BufferRealize(op.buffer, bounds, condition, body, op.span)
|
|
|
|
def visit_assert_(self, op):
|
|
"""Mutator implementation for AssertStmt."""
|
|
condition = self.visit_expr(op.condition)
|
|
message_parts = []
|
|
message_parts_changed = False
|
|
for message_part in op.message_parts:
|
|
if tvm.ir.is_prim_expr(message_part):
|
|
new_message_part = self.visit_expr(message_part)
|
|
if new_message_part is not message_part:
|
|
message_parts_changed = True
|
|
message_parts.append(new_message_part)
|
|
else:
|
|
message_parts.append(message_part)
|
|
|
|
if condition is op.condition and not message_parts_changed:
|
|
return op
|
|
|
|
return tvm.tirx.AssertStmt(op.kind, condition, message_parts, op.span)
|
|
|
|
def visit_producer_store_(self, op):
|
|
"""Mutator implementation for ProducerStore."""
|
|
value = self.visit_expr(op.value)
|
|
indices = [self.visit_expr(idx) for idx in op.indices]
|
|
|
|
indices_changed = any(old is not new for old, new in zip(op.indices, indices))
|
|
|
|
if value is op.value and not indices_changed:
|
|
return op
|
|
|
|
return tvm.tirx.ProducerStore(op.producer, value, indices, op.span)
|
|
|
|
def visit_producer_realize_(self, op):
|
|
"""Mutator implementation for ProducerRealize."""
|
|
bounds = []
|
|
bounds_changed = False
|
|
|
|
for r in op.bounds:
|
|
new_min = self.visit_expr(r.min)
|
|
new_extent = self.visit_expr(r.extent)
|
|
|
|
if new_min is not r.min or new_extent is not r.extent:
|
|
bounds_changed = True
|
|
bounds.append(tvm.ir.Range(new_min, new_extent))
|
|
else:
|
|
bounds.append(r)
|
|
|
|
condition = self.visit_expr(op.condition)
|
|
body = self.visit_stmt(op.body)
|
|
|
|
if not bounds_changed and condition is op.condition and body is op.body:
|
|
return op
|
|
|
|
return tvm.tirx.ProducerRealize(
|
|
op.producer, bounds, condition, body, op.storage_scope, op.span
|
|
)
|
|
|
|
def visit_prefetch_(self, op):
|
|
"""Mutator implementation for Prefetch."""
|
|
bounds = []
|
|
bounds_changed = False
|
|
|
|
for r in op.bounds:
|
|
new_min = self.visit_expr(r.min)
|
|
new_extent = self.visit_expr(r.extent)
|
|
|
|
if new_min is not r.min or new_extent is not r.extent:
|
|
bounds_changed = True
|
|
bounds.append(tvm.ir.Range(new_min, new_extent))
|
|
else:
|
|
bounds.append(r)
|
|
|
|
if not bounds_changed:
|
|
return op
|
|
|
|
return tvm.tirx.Prefetch(op.buffer, bounds, op.span)
|
|
|
|
def visit_seqstmt_(self, op):
|
|
"""Mutator implementation for SeqStmt."""
|
|
new_seq = []
|
|
changed = False
|
|
|
|
for stmt in op.seq:
|
|
new_stmt = self.visit_stmt(stmt)
|
|
if new_stmt is not stmt:
|
|
changed = True
|
|
if isinstance(new_stmt, tvm.tirx.SeqStmt):
|
|
# Flatten nested SeqStmt
|
|
new_seq.extend(new_stmt.seq)
|
|
changed = True
|
|
else:
|
|
new_seq.append(new_stmt)
|
|
|
|
if not changed:
|
|
return op
|
|
|
|
if len(new_seq) == 1:
|
|
return new_seq[0]
|
|
|
|
return tvm.tirx.SeqStmt(new_seq, op.span)
|
|
|
|
def visit_evaluate_(self, op):
|
|
"""Mutator implementation for Evaluate."""
|
|
value = self.visit_expr(op.value)
|
|
|
|
if value is op.value:
|
|
return op
|
|
|
|
return tvm.tirx.Evaluate(value, op.span)
|
|
|
|
def visit_block_(self, op):
|
|
"""Mutator implementation for Block."""
|
|
# Process iter_vars
|
|
iter_vars = []
|
|
iter_vars_changed = False
|
|
|
|
for iv in op.iter_vars:
|
|
old_dom = iv.dom
|
|
new_min = self.visit_expr(old_dom.min)
|
|
new_extent = self.visit_expr(old_dom.extent)
|
|
|
|
if new_min is not old_dom.min or new_extent is not old_dom.extent:
|
|
iter_vars_changed = True
|
|
new_dom = tvm.ir.Range(new_min, new_extent)
|
|
iter_vars.append(tvm.tirx.IterVar(new_dom, iv.var, iv.iter_type, iv.thread_tag))
|
|
else:
|
|
iter_vars.append(iv)
|
|
|
|
# Process reads/writes buffer regions
|
|
def _mutate_buffer_regions(regions):
|
|
new_regions = []
|
|
regions_changed = False
|
|
|
|
for region in regions:
|
|
new_ranges = []
|
|
ranges_changed = False
|
|
|
|
for r in region.region:
|
|
new_min = self.visit_expr(r.min)
|
|
new_extent = self.visit_expr(r.extent)
|
|
|
|
if new_min is not r.min or new_extent is not r.extent:
|
|
ranges_changed = True
|
|
new_ranges.append(tvm.ir.Range(new_min, new_extent))
|
|
else:
|
|
new_ranges.append(r)
|
|
|
|
if ranges_changed:
|
|
regions_changed = True
|
|
new_regions.append(tvm.tirx.BufferRegion(region.buffer, new_ranges))
|
|
else:
|
|
new_regions.append(region)
|
|
|
|
return new_regions, regions_changed
|
|
|
|
reads, reads_changed = _mutate_buffer_regions(op.reads)
|
|
writes, writes_changed = _mutate_buffer_regions(op.writes)
|
|
|
|
# Process match buffers
|
|
match_buffers = []
|
|
match_buffers_changed = False
|
|
|
|
for match_buffer in op.match_buffers:
|
|
source_region = match_buffer.source
|
|
new_ranges = []
|
|
ranges_changed = False
|
|
|
|
for r in source_region.region:
|
|
new_min = self.visit_expr(r.min)
|
|
new_extent = self.visit_expr(r.extent)
|
|
|
|
if new_min is not r.min or new_extent is not r.extent:
|
|
ranges_changed = True
|
|
new_ranges.append(tvm.ir.Range(new_min, new_extent))
|
|
else:
|
|
new_ranges.append(r)
|
|
|
|
if ranges_changed:
|
|
match_buffers_changed = True
|
|
new_source = tvm.tirx.BufferRegion(source_region.buffer, new_ranges)
|
|
match_buffers.append(tvm.tirx.MatchBufferRegion(match_buffer.buffer, new_source))
|
|
else:
|
|
match_buffers.append(match_buffer)
|
|
|
|
# Process init and body
|
|
init = self.visit_stmt(op.init) if op.init is not None else None
|
|
body = self.visit_stmt(op.body)
|
|
|
|
# Check if anything changed
|
|
if (
|
|
not iter_vars_changed
|
|
and not reads_changed
|
|
and not writes_changed
|
|
and not match_buffers_changed
|
|
and (init is op.init or (init is None and op.init is None))
|
|
and body is op.body
|
|
):
|
|
return op
|
|
return tvm.tirx.SBlock(
|
|
iter_vars,
|
|
reads,
|
|
writes,
|
|
op.name_hint,
|
|
body,
|
|
init,
|
|
op.alloc_buffers,
|
|
match_buffers,
|
|
op.annotations,
|
|
)
|
|
|
|
def visit_block_realize_(self, op):
|
|
"""Mutator implementation for BlockRealize."""
|
|
iter_values = [self.visit_expr(val) for val in op.iter_values]
|
|
predicate = self.visit_expr(op.predicate)
|
|
block = self.visit_stmt(op.block)
|
|
|
|
iter_values_changed = any(old is not new for old, new in zip(op.iter_values, iter_values))
|
|
|
|
if not iter_values_changed and predicate is op.predicate and block is op.block:
|
|
return op
|
|
|
|
if not isinstance(block, tvm.tirx.SBlock):
|
|
raise TypeError(f"Expected SBlock, but got {type(block)}")
|
|
|
|
return tvm.tirx.SBlockRealize(iter_values, predicate, block)
|
|
|
|
def visit_scope_id_def_stmt_(self, op):
|
|
"""Mutator implementation for ScopeIdDefStmt.
|
|
|
|
Mirrors the C++ mutator: rewrite ``extents`` and
|
|
``preferred_extents`` via ``visit_expr``. Deferred-extent defs
|
|
(extents is None) and unchanged extents pass through.
|
|
"""
|
|
from .exec_scope import _SCOPE_BINDING_TO_PARENT_CUR, ScopeIdDef
|
|
|
|
# ``def`` is a Python keyword; access the C++ field via ``getattr``.
|
|
sid = getattr(op, "def")
|
|
changed = False
|
|
|
|
def _walk(arr):
|
|
nonlocal changed
|
|
if arr is None:
|
|
return None
|
|
out = []
|
|
for e in arr:
|
|
ne = self.visit_expr(e)
|
|
if ne is not e:
|
|
changed = True
|
|
out.append(ne)
|
|
return out
|
|
|
|
new_extents = _walk(sid.extents)
|
|
new_pref = _walk(sid.preferred_extents)
|
|
if not changed:
|
|
return op
|
|
parent, cur = _SCOPE_BINDING_TO_PARENT_CUR[sid.scope]
|
|
new_def = ScopeIdDef(sid.def_ids, new_extents, parent, cur, new_pref)
|
|
return tvm.tirx.ScopeIdDefStmt(new_def, op.span)
|
|
|
|
def visit_op_call_(self, op):
|
|
"""Mutator implementation for TilePrimitiveCall."""
|
|
new_args = []
|
|
args_changed = False
|
|
|
|
for arg in op.args:
|
|
if tvm.ir.is_prim_expr(arg):
|
|
new_arg = self.visit_expr(arg)
|
|
elif isinstance(arg, tvm.tirx.Stmt):
|
|
new_arg = self.visit_stmt(arg)
|
|
elif isinstance(arg, tvm.tirx.BufferRegion):
|
|
new_arg = self.visit_buffer_region_(arg)
|
|
else:
|
|
new_arg = arg
|
|
|
|
if new_arg is not arg:
|
|
args_changed = True
|
|
new_args.append(new_arg)
|
|
|
|
# Also mutate Expr values in the config map
|
|
new_config = {}
|
|
config_changed = False
|
|
for key, value in op.config.items():
|
|
if tvm.ir.is_prim_expr(value):
|
|
new_value = self.visit_expr(value)
|
|
elif isinstance(value, tvm.tirx.Stmt):
|
|
new_value = self.visit_stmt(value)
|
|
else:
|
|
new_value = value
|
|
if new_value is not value:
|
|
config_changed = True
|
|
new_config[key] = new_value
|
|
|
|
if not args_changed and not config_changed:
|
|
return op
|
|
|
|
return tvm.tirx.TilePrimitiveCall(
|
|
*new_args,
|
|
op=op.op,
|
|
workspace=op.workspace,
|
|
config=new_config,
|
|
dispatch=op.dispatch,
|
|
scope=op.scope,
|
|
)
|
|
|
|
def visit_buffer_region_(self, op):
|
|
"""Mutator implementation for BufferRegion."""
|
|
|
|
def _mutate_range(range):
|
|
new_min = self.visit_expr(range.min)
|
|
new_extent = self.visit_expr(range.extent)
|
|
|
|
if new_min is range.min and new_extent is range.extent:
|
|
return range
|
|
else:
|
|
return Range.from_min_extent(new_min, new_extent)
|
|
|
|
region = [_mutate_range(r) for r in op.region]
|
|
|
|
if all(old_r is new_r for old_r, new_r in zip(op.region, region)):
|
|
return op
|
|
else:
|
|
return tvm.tirx.BufferRegion(op.buffer, region)
|
|
|
|
def visit_alloc_buffer_(self, op):
|
|
"""Mutator implementation for AllocBuffer."""
|
|
if hasattr(op, "body"):
|
|
body = self.visit_stmt(op.body)
|
|
if body is op.body:
|
|
return op
|
|
return tvm.tirx.AllocBuffer(op.buffer, body, op.annotations, op.span)
|
|
return op
|
|
|
|
def __call__(self, stmt):
|
|
"""Call mutator on statement.
|
|
|
|
Parameters
|
|
----------
|
|
stmt : tvm.tirx.Stmt
|
|
The statement to be mutated.
|
|
|
|
Returns
|
|
-------
|
|
result : tvm.tirx.Stmt
|
|
The mutated statement
|
|
"""
|
|
return self.visit_stmt(stmt)
|
|
|
|
|
|
class StmtExprVisitor(StmtVisitor, ExprVisitor):
|
|
"""A visitor over both statements and expressions.
|
|
|
|
This class inherits from both StmtVisitor and ExprVisitor to recursively visit
|
|
both statements and expressions.
|
|
"""
|
|
|
|
def __init__(self):
|
|
StmtVisitor.__init__(self)
|
|
self._stmt_dispatch_map = self._dispatch_map.copy()
|
|
ExprVisitor.__init__(self)
|
|
self._expr_dispatch_map = self._dispatch_map.copy()
|
|
self._dispatch_map = {}
|
|
self._dispatch_map.update(self._stmt_dispatch_map)
|
|
self._dispatch_map.update(self._expr_dispatch_map)
|
|
|
|
def visit_expr(self, expr):
|
|
"""Visit an expression used in a statement.
|
|
|
|
Parameters
|
|
----------
|
|
expr : Expr
|
|
The expression to be visited.
|
|
"""
|
|
return ExprVisitor.visit_expr(self, expr)
|
|
|
|
|
|
class StmtExprMutator(StmtMutator, ExprMutator):
|
|
"""A mutator over both statements and expressions.
|
|
|
|
This class inherits from both StmtMutator and ExprMutator to recursively transform
|
|
both statements and expressions.
|
|
"""
|
|
|
|
def __init__(self):
|
|
StmtMutator.__init__(self)
|
|
self._stmt_dispatch_map = self._dispatch_map.copy()
|
|
ExprMutator.__init__(self)
|
|
self._expr_dispatch_map = self._dispatch_map.copy()
|
|
self._dispatch_map = {}
|
|
self._dispatch_map.update(self._stmt_dispatch_map)
|
|
self._dispatch_map.update(self._expr_dispatch_map)
|
|
|
|
def visit_expr(self, expr):
|
|
"""Mutate an expression used in a statement.
|
|
|
|
Parameters
|
|
----------
|
|
expr : Expr
|
|
The expression to be mutated.
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The mutated expression.
|
|
"""
|
|
return ExprMutator.visit_expr(self, expr)
|
|
|
|
|
|
def ir_transform(stmt, preorder, postorder, only_enable=None):
|
|
"""Recursively visit and transform ir nodes in post DFS order.
|
|
|
|
Parameters
|
|
----------
|
|
stmt : tvm.tirx.Stmt
|
|
The input to be transformed.
|
|
|
|
preorder: function
|
|
The function called in before recursive mutation
|
|
If preorder returns None, then the transform will proceed to recursive call.
|
|
If preorder returns a not None tvm.tirx.Stmt/Expr, the transformer will simply return it and
|
|
won't do further recursion.
|
|
|
|
postorder : function
|
|
The function called after recursive mutation.
|
|
|
|
only_enable : Optional[List[str]]
|
|
List of types that we only enable.
|
|
|
|
Returns
|
|
-------
|
|
result : tvm.tirx.Stmt
|
|
The result.
|
|
"""
|
|
return _ffi_api.IRTransform(stmt, preorder, postorder, only_enable) # type: ignore
|
|
|
|
|
|
def post_order_visit(node, fvisit):
|
|
"""Recursively visit a statement or expression in post DFS order, applying fvisit.
|
|
Each node is guaranteed to be visited only once.
|
|
|
|
Parameters
|
|
----------
|
|
node : tvm.tirx.Stmt or tvm.ir.Expr
|
|
The statement or expression to visit.
|
|
|
|
fvisit: function
|
|
The visitor function.
|
|
"""
|
|
return _ffi_api.PostOrderVisit(node, fvisit) # type: ignore
|
|
|
|
|
|
def pre_order_visit(node, fvisit):
|
|
"""Recursively visit a statement or expression in pre-order, applying fvisit.
|
|
If fvisit returns False, it won't visit the children of the node.
|
|
|
|
Parameters
|
|
----------
|
|
node : tvm.tirx.Stmt or tvm.ir.Expr
|
|
The statement or expression to visit.
|
|
|
|
fvisit: function of the signature Object -> bool
|
|
The visitor function.
|
|
"""
|
|
return _ffi_api.PreOrderVisit(node, fvisit) # type: ignore
|
|
|
|
|
|
def substitute(node, vmap):
|
|
"""Substitute the var specified by vmap.
|
|
|
|
Parameters
|
|
----------
|
|
node: ObjectRef
|
|
The input.
|
|
|
|
vmap : Dict[Var, Expr]
|
|
The variable mapping.
|
|
|
|
Returns
|
|
-------
|
|
result : tvm.tirx.Stmt
|
|
The result.
|
|
"""
|
|
return _ffi_api.Substitute(node, vmap) # type: ignore
|
|
|
|
|
|
def renew_defs(func: PrimFunc):
|
|
"""Re-generate the definition nodes for a TIR, including VarDef, BufferDef.
|
|
This pass works as a simple DeepCopy to duplicate a function with different Vars and
|
|
Buffers but the same behavior
|
|
|
|
Parameters
|
|
----------
|
|
func: PrimFunc
|
|
The input function
|
|
|
|
Returns
|
|
-------
|
|
result : PrimFunc
|
|
The new generated func.
|
|
"""
|
|
return _ffi_api.RenewDefs(func) # type: ignore
|