wehub-resource-sync
1066 lines
29 KiB
Python
1066 lines
29 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.
|
|
"""
|
|
Tests for StmtVisitor and StmtMutator functionality in TVM TIR.
|
|
"""
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm import tirx as tir
|
|
from tvm.ir import Range
|
|
from tvm.script import tirx as T
|
|
from tvm.script.tirx import tile as Tx
|
|
from tvm.tirx.expr import EQ, GT, LT, Add, IntImm, Mul, Sub, Var
|
|
from tvm.tirx.stmt_functor import StmtExprMutator, StmtExprVisitor, StmtMutator, StmtVisitor
|
|
|
|
|
|
class ASTLog:
|
|
"""Helper class to log AST traversal"""
|
|
|
|
def __init__(self) -> None:
|
|
self.log = []
|
|
self.indent = "\t"
|
|
self.level = 0
|
|
|
|
def push_scope(self):
|
|
self.level += 1
|
|
|
|
def pop_scope(self):
|
|
self.level -= 1
|
|
|
|
def add(self, s: str):
|
|
self.log.append(self.indent * self.level + s)
|
|
|
|
def __str__(self) -> str:
|
|
return "\n".join(self.log)
|
|
|
|
|
|
class BasicStmtVisitor(StmtVisitor):
|
|
"""Default StmtVisitor - doesn't override any methods"""
|
|
|
|
pass
|
|
|
|
|
|
class ASTPrinter(StmtVisitor):
|
|
"""Print TIR AST in structured format."""
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.log = ASTLog()
|
|
|
|
def visit_bind_(self, op):
|
|
self.log.add("Bind")
|
|
self.log.push_scope()
|
|
self.visit_expr(op.value)
|
|
self.log.pop_scope()
|
|
|
|
def visit_attr_(self, op):
|
|
self.log.add("AttrStmt")
|
|
self.log.push_scope()
|
|
self.visit_expr(op.value)
|
|
self.visit_stmt(op.body)
|
|
self.log.pop_scope()
|
|
|
|
def visit_assert_(self, op):
|
|
self.log.add("AssertStmt")
|
|
self.log.push_scope()
|
|
self.visit_expr(op.condition)
|
|
self.visit_expr(op.message)
|
|
self.visit_stmt(op.body)
|
|
self.log.pop_scope()
|
|
|
|
def visit_for_(self, op):
|
|
self.log.add("For")
|
|
self.log.push_scope()
|
|
self.visit_expr(op.min)
|
|
self.visit_expr(op.extent)
|
|
self.visit_stmt(op.body)
|
|
self.log.pop_scope()
|
|
|
|
def visit_while_(self, op):
|
|
self.log.add("While")
|
|
self.log.push_scope()
|
|
self.visit_expr(op.condition)
|
|
self.visit_stmt(op.body)
|
|
self.log.pop_scope()
|
|
|
|
def visit_buffer_store_(self, op):
|
|
self.log.add("BufferStore")
|
|
self.log.push_scope()
|
|
self.visit_expr(op.value)
|
|
for index in op.indices:
|
|
self.visit_expr(index)
|
|
self.log.pop_scope()
|
|
|
|
def visit_seqstmt_(self, op):
|
|
self.log.add("SeqStmt")
|
|
self.log.push_scope()
|
|
for stmt in op.seq:
|
|
self.visit_stmt(stmt)
|
|
self.log.pop_scope()
|
|
|
|
def visit_evaluate_(self, op):
|
|
self.log.add("Evaluate")
|
|
self.log.push_scope()
|
|
self.visit_expr(op.value)
|
|
self.log.pop_scope()
|
|
|
|
def visit_block_(self, op):
|
|
self.log.add("Block")
|
|
self.log.push_scope()
|
|
if op.init is not None:
|
|
self.visit_stmt(op.init)
|
|
self.visit_stmt(op.body)
|
|
self.log.pop_scope()
|
|
|
|
def visit_block_realize_(self, op):
|
|
self.log.add("BlockRealize")
|
|
self.log.push_scope()
|
|
for val in op.iter_values:
|
|
self.visit_expr(val)
|
|
self.visit_expr(op.predicate)
|
|
self.visit_stmt(op.block)
|
|
self.log.pop_scope()
|
|
|
|
def visit_if_then_else_(self, op):
|
|
self.log.add("IfThenElse")
|
|
self.log.push_scope()
|
|
self.visit_expr(op.condition)
|
|
self.visit_stmt(op.then_case)
|
|
if op.else_case:
|
|
self.visit_stmt(op.else_case)
|
|
self.log.pop_scope()
|
|
|
|
def visit_decl_buffer_(self, op):
|
|
self.log.add("DeclBuffer")
|
|
self.log.push_scope()
|
|
self.visit_stmt(op.body)
|
|
self.log.pop_scope()
|
|
|
|
def visit_break_(self, op):
|
|
self.log.add("Break")
|
|
|
|
def visit_continue_(self, op):
|
|
self.log.add("Continue")
|
|
|
|
def visit_op_call_(self, op):
|
|
self.log.add("OpCall")
|
|
self.log.push_scope()
|
|
for arg in op.args:
|
|
if isinstance(arg, tir.BufferRegion):
|
|
self.visit_buffer_region_(arg)
|
|
else:
|
|
self.visit_expr(arg)
|
|
self.log.pop_scope()
|
|
|
|
def visit_buffer_region_(self, op):
|
|
self.log.add("BufferRegion")
|
|
self.log.push_scope()
|
|
for r in op.region:
|
|
self.visit_expr(r.min)
|
|
self.visit_expr(r.extent)
|
|
self.log.pop_scope()
|
|
|
|
def visit_expr(self, expr):
|
|
"""Simple expression visitor that logs expression types."""
|
|
if expr is None:
|
|
return
|
|
|
|
if isinstance(expr, Var):
|
|
self.log.add("Var")
|
|
elif isinstance(expr, IntImm):
|
|
self.log.add("IntImm")
|
|
elif isinstance(expr, Add):
|
|
self.log.add("Add")
|
|
self.log.push_scope()
|
|
self.visit_expr(expr.a)
|
|
self.visit_expr(expr.b)
|
|
self.log.pop_scope()
|
|
elif isinstance(expr, Sub):
|
|
self.log.add("Sub")
|
|
self.log.push_scope()
|
|
self.visit_expr(expr.a)
|
|
self.visit_expr(expr.b)
|
|
self.log.pop_scope()
|
|
elif isinstance(expr, Mul):
|
|
self.log.add("Mul")
|
|
self.log.push_scope()
|
|
self.visit_expr(expr.a)
|
|
self.visit_expr(expr.b)
|
|
self.log.pop_scope()
|
|
elif isinstance(expr, EQ):
|
|
self.log.add("EQ")
|
|
self.log.push_scope()
|
|
self.visit_expr(expr.a)
|
|
self.visit_expr(expr.b)
|
|
self.log.pop_scope()
|
|
elif isinstance(expr, LT):
|
|
self.log.add("LT")
|
|
self.log.push_scope()
|
|
self.visit_expr(expr.a)
|
|
self.visit_expr(expr.b)
|
|
self.log.pop_scope()
|
|
elif isinstance(expr, GT):
|
|
self.log.add("GT")
|
|
self.log.push_scope()
|
|
self.visit_expr(expr.a)
|
|
self.visit_expr(expr.b)
|
|
self.log.pop_scope()
|
|
else:
|
|
self.log.add(f"Expr::{type(expr).__name__}")
|
|
|
|
|
|
class ASTPrinterMutator(StmtMutator):
|
|
"""Print TIR AST in post-order while mutating."""
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.log = ASTLog()
|
|
|
|
def visit_bind_(self, op):
|
|
result = super().visit_bind_(op)
|
|
self.log.add("Bind")
|
|
return result
|
|
|
|
def visit_attr_(self, op):
|
|
result = super().visit_attr_(op)
|
|
self.log.add("AttrStmt")
|
|
return result
|
|
|
|
def visit_assert_(self, op):
|
|
result = super().visit_assert_(op)
|
|
self.log.add("AssertStmt")
|
|
return result
|
|
|
|
def visit_for_(self, op):
|
|
result = super().visit_for_(op)
|
|
self.log.add("For")
|
|
return result
|
|
|
|
def visit_while_(self, op):
|
|
result = super().visit_while_(op)
|
|
self.log.add("While")
|
|
return result
|
|
|
|
def visit_buffer_store_(self, op):
|
|
result = super().visit_buffer_store_(op)
|
|
self.log.add("BufferStore")
|
|
return result
|
|
|
|
def visit_seqstmt_(self, op):
|
|
result = super().visit_seqstmt_(op)
|
|
self.log.add("SeqStmt")
|
|
return result
|
|
|
|
def visit_evaluate_(self, op):
|
|
result = super().visit_evaluate_(op)
|
|
self.log.add("Evaluate")
|
|
return result
|
|
|
|
def visit_block_(self, op):
|
|
result = super().visit_block_(op)
|
|
self.log.add("Block")
|
|
return result
|
|
|
|
def visit_block_realize_(self, op):
|
|
result = super().visit_block_realize_(op)
|
|
self.log.add("BlockRealize")
|
|
return result
|
|
|
|
def visit_if_then_else_(self, op):
|
|
result = super().visit_if_then_else_(op)
|
|
self.log.add("IfThenElse")
|
|
return result
|
|
|
|
def visit_decl_buffer_(self, op):
|
|
result = super().visit_decl_buffer_(op)
|
|
self.log.add("DeclBuffer")
|
|
return result
|
|
|
|
def visit_break_(self, op):
|
|
result = super().visit_break_(op)
|
|
self.log.add("Break")
|
|
return result
|
|
|
|
def visit_continue_(self, op):
|
|
result = super().visit_continue_(op)
|
|
self.log.add("Continue")
|
|
return result
|
|
|
|
def visit_op_call_(self, op):
|
|
result = super().visit_op_call_(op)
|
|
self.log.add("OpCall")
|
|
return result
|
|
|
|
def visit_buffer_region_(self, op):
|
|
result = super().visit_buffer_region_(op)
|
|
self.log.add("BufferRegion")
|
|
return result
|
|
|
|
def visit_expr(self, expr):
|
|
"""Simple expression visitor that logs expression types."""
|
|
if expr is None:
|
|
return expr
|
|
|
|
if isinstance(expr, Var):
|
|
self.log.add("Var")
|
|
return expr
|
|
elif isinstance(expr, IntImm):
|
|
self.log.add("IntImm")
|
|
return expr
|
|
elif isinstance(expr, Add):
|
|
a = self.visit_expr(expr.a)
|
|
b = self.visit_expr(expr.b)
|
|
self.log.add("Add")
|
|
if a is expr.a and b is expr.b:
|
|
return expr
|
|
return tir.Add(a, b)
|
|
elif isinstance(expr, Sub):
|
|
a = self.visit_expr(expr.a)
|
|
b = self.visit_expr(expr.b)
|
|
self.log.add("Sub")
|
|
if a is expr.a and b is expr.b:
|
|
return expr
|
|
return tir.Sub(a, b)
|
|
elif isinstance(expr, Mul):
|
|
a = self.visit_expr(expr.a)
|
|
b = self.visit_expr(expr.b)
|
|
self.log.add("Mul")
|
|
if a is expr.a and b is expr.b:
|
|
return expr
|
|
return tir.Mul(a, b)
|
|
elif isinstance(expr, EQ):
|
|
a = self.visit_expr(expr.a)
|
|
b = self.visit_expr(expr.b)
|
|
self.log.add("EQ")
|
|
if a is expr.a and b is expr.b:
|
|
return expr
|
|
return tir.EQ(a, b)
|
|
elif isinstance(expr, LT):
|
|
a = self.visit_expr(expr.a)
|
|
b = self.visit_expr(expr.b)
|
|
self.log.add("LT")
|
|
if a is expr.a and b is expr.b:
|
|
return expr
|
|
return tir.LT(a, b)
|
|
elif isinstance(expr, GT):
|
|
a = self.visit_expr(expr.a)
|
|
b = self.visit_expr(expr.b)
|
|
self.log.add("GT")
|
|
if a is expr.a and b is expr.b:
|
|
return expr
|
|
return tir.GT(a, b)
|
|
else:
|
|
self.log.add(f"Expr::{type(expr).__name__}")
|
|
return expr
|
|
|
|
|
|
class StmtExprASTPrinter(StmtExprVisitor):
|
|
"""AST printer using StmtExprVisitor."""
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.log = ASTLog()
|
|
|
|
def visit_bind_(self, op):
|
|
self.log.add("Bind")
|
|
self.log.push_scope()
|
|
super().visit_bind_(op)
|
|
self.log.pop_scope()
|
|
|
|
def visit_attr_(self, op):
|
|
self.log.add("AttrStmt")
|
|
self.log.push_scope()
|
|
super().visit_attr_(op)
|
|
self.log.pop_scope()
|
|
|
|
def visit_assert_(self, op):
|
|
self.log.add("AssertStmt")
|
|
self.log.push_scope()
|
|
super().visit_assert_(op)
|
|
self.log.pop_scope()
|
|
|
|
def visit_for_(self, op):
|
|
self.log.add("For")
|
|
self.log.push_scope()
|
|
super().visit_for_(op)
|
|
self.log.pop_scope()
|
|
|
|
def visit_while_(self, op):
|
|
self.log.add("While")
|
|
self.log.push_scope()
|
|
super().visit_while_(op)
|
|
self.log.pop_scope()
|
|
|
|
def visit_buffer_store_(self, op):
|
|
self.log.add("BufferStore")
|
|
self.log.push_scope()
|
|
super().visit_buffer_store_(op)
|
|
self.log.pop_scope()
|
|
|
|
def visit_seqstmt_(self, op):
|
|
self.log.add("SeqStmt")
|
|
self.log.push_scope()
|
|
super().visit_seqstmt_(op)
|
|
self.log.pop_scope()
|
|
|
|
def visit_evaluate_(self, op):
|
|
self.log.add("Evaluate")
|
|
self.log.push_scope()
|
|
super().visit_evaluate_(op)
|
|
self.log.pop_scope()
|
|
|
|
def visit_block_(self, op):
|
|
self.log.add("Block")
|
|
self.log.push_scope()
|
|
super().visit_block_(op)
|
|
self.log.pop_scope()
|
|
|
|
def visit_block_realize_(self, op):
|
|
self.log.add("BlockRealize")
|
|
self.log.push_scope()
|
|
super().visit_block_realize_(op)
|
|
self.log.pop_scope()
|
|
|
|
def visit_if_then_else_(self, op):
|
|
self.log.add("IfThenElse")
|
|
self.log.push_scope()
|
|
super().visit_if_then_else_(op)
|
|
self.log.pop_scope()
|
|
|
|
def visit_decl_buffer_(self, op):
|
|
self.log.add("DeclBuffer")
|
|
self.log.push_scope()
|
|
super().visit_decl_buffer_(op)
|
|
self.log.pop_scope()
|
|
|
|
def visit_break_(self, op):
|
|
self.log.add("Break")
|
|
super().visit_break_(op)
|
|
|
|
def visit_continue_(self, op):
|
|
self.log.add("Continue")
|
|
super().visit_continue_(op)
|
|
|
|
# ExprVisitor methods
|
|
def visit_var_(self, op):
|
|
self.log.add("Var")
|
|
|
|
def visit_int_imm_(self, op):
|
|
self.log.add("IntImm")
|
|
|
|
def visit_add_(self, op):
|
|
self.log.add("Add")
|
|
self.log.push_scope()
|
|
super().visit_add_(op)
|
|
self.log.pop_scope()
|
|
|
|
def visit_sub_(self, op):
|
|
self.log.add("Sub")
|
|
self.log.push_scope()
|
|
super().visit_sub_(op)
|
|
self.log.pop_scope()
|
|
|
|
def visit_mul_(self, op):
|
|
self.log.add("Mul")
|
|
self.log.push_scope()
|
|
super().visit_mul_(op)
|
|
self.log.pop_scope()
|
|
|
|
def visit_eq_(self, op):
|
|
self.log.add("EQ")
|
|
self.log.push_scope()
|
|
super().visit_eq_(op)
|
|
self.log.pop_scope()
|
|
|
|
def visit_lt_(self, op):
|
|
self.log.add("LT")
|
|
self.log.push_scope()
|
|
super().visit_lt_(op)
|
|
self.log.pop_scope()
|
|
|
|
def visit_gt_(self, op):
|
|
self.log.add("GT")
|
|
self.log.push_scope()
|
|
super().visit_gt_(op)
|
|
self.log.pop_scope()
|
|
|
|
|
|
class StmtExprMutatorPrinter(StmtExprMutator):
|
|
"""AST mutator printer using StmtExprMutator."""
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.log = ASTLog()
|
|
|
|
def visit_bind_(self, op):
|
|
result = super().visit_bind_(op)
|
|
self.log.add("Bind")
|
|
return result
|
|
|
|
def visit_attr_(self, op):
|
|
result = super().visit_attr_(op)
|
|
self.log.add("AttrStmt")
|
|
return result
|
|
|
|
def visit_assert_(self, op):
|
|
result = super().visit_assert_(op)
|
|
self.log.add("AssertStmt")
|
|
return result
|
|
|
|
def visit_for_(self, op):
|
|
result = super().visit_for_(op)
|
|
self.log.add("For")
|
|
return result
|
|
|
|
def visit_while_(self, op):
|
|
result = super().visit_while_(op)
|
|
self.log.add("While")
|
|
return result
|
|
|
|
def visit_buffer_store_(self, op):
|
|
result = super().visit_buffer_store_(op)
|
|
self.log.add("BufferStore")
|
|
return result
|
|
|
|
def visit_seqstmt_(self, op):
|
|
result = super().visit_seqstmt_(op)
|
|
self.log.add("SeqStmt")
|
|
return result
|
|
|
|
def visit_evaluate_(self, op):
|
|
result = super().visit_evaluate_(op)
|
|
self.log.add("Evaluate")
|
|
return result
|
|
|
|
def visit_block_(self, op):
|
|
result = super().visit_block_(op)
|
|
self.log.add("Block")
|
|
return result
|
|
|
|
def visit_block_realize_(self, op):
|
|
result = super().visit_block_realize_(op)
|
|
self.log.add("BlockRealize")
|
|
return result
|
|
|
|
# ExprMutator methods
|
|
def visit_var_(self, op):
|
|
result = super().visit_var_(op)
|
|
self.log.add("Var")
|
|
return result
|
|
|
|
def visit_int_imm_(self, op):
|
|
result = super().visit_int_imm_(op)
|
|
self.log.add("IntImm")
|
|
return result
|
|
|
|
def visit_add_(self, op):
|
|
result = super().visit_add_(op)
|
|
self.log.add("Add")
|
|
return result
|
|
|
|
def visit_sub_(self, op):
|
|
result = super().visit_sub_(op)
|
|
self.log.add("Sub")
|
|
return result
|
|
|
|
def visit_mul_(self, op):
|
|
result = super().visit_mul_(op)
|
|
self.log.add("Mul")
|
|
return result
|
|
|
|
def visit_eq_(self, op):
|
|
result = super().visit_eq_(op)
|
|
self.log.add("EQ")
|
|
return result
|
|
|
|
def visit_lt_(self, op):
|
|
result = super().visit_lt_(op)
|
|
self.log.add("LT")
|
|
return result
|
|
|
|
def visit_gt_(self, op):
|
|
result = super().visit_gt_(op)
|
|
self.log.add("GT")
|
|
return result
|
|
|
|
|
|
def basic_check(stmt, visitor_str, mutator_str):
|
|
"""Check visitor and mutator behavior on the given statement."""
|
|
# Check basic visitor
|
|
basic_visitor = BasicStmtVisitor()
|
|
basic_visitor.visit_stmt(stmt)
|
|
|
|
# Check AST printer visitor
|
|
log_visitor = ASTPrinter()
|
|
log_visitor.visit_stmt(stmt)
|
|
assert str(log_visitor.log) == visitor_str
|
|
|
|
# Check AST printer mutator
|
|
log_mutator = ASTPrinterMutator()
|
|
result = log_mutator.visit_stmt(stmt)
|
|
# Check we get back structurally equivalent statement
|
|
tvm.ir.assert_structural_equal(result, stmt)
|
|
assert str(log_mutator.log) == mutator_str
|
|
|
|
|
|
def create_test_statements():
|
|
"""Create test statements for various TIR constructs."""
|
|
x = tir.Var("x", "int32")
|
|
|
|
# IntImm
|
|
int_imm = tir.IntImm("int32", 10)
|
|
|
|
# Simple expression
|
|
add_expr = tir.Add(x, int_imm)
|
|
|
|
# Evaluate
|
|
evaluate_stmt = tir.Evaluate(add_expr)
|
|
|
|
# Bind + SeqStmt (was LetStmt)
|
|
let_stmt = tir.SeqStmt([tir.Bind(x, int_imm), evaluate_stmt])
|
|
|
|
# For loop
|
|
for_loop = tir.For(x, 0, 10, tir.ForKind.SERIAL, evaluate_stmt)
|
|
|
|
# While loop
|
|
while_loop = tir.While(tir.LT(x, int_imm), evaluate_stmt)
|
|
|
|
# Buffer operations
|
|
buffer_var = tir.Var("buf", "handle")
|
|
buffer = tir.decl_buffer((10,), "int32", buffer_var.name)
|
|
buffer_store = tir.BufferStore(buffer, add_expr, [int_imm])
|
|
|
|
# Sequence of statements
|
|
seq_stmt = tir.SeqStmt([evaluate_stmt, for_loop])
|
|
|
|
# Block with iteration variables
|
|
iter_var = tir.IterVar(Range(0, 10), x, 0)
|
|
block = tir.SBlock([iter_var], [], [], "block", evaluate_stmt)
|
|
block_realize = tir.SBlockRealize([int_imm], tir.IntImm("bool", 1), block)
|
|
|
|
# IfThenElse statement
|
|
if_then_else = tir.IfThenElse(tir.LT(x, int_imm), evaluate_stmt, evaluate_stmt)
|
|
|
|
# Break and continue statements inside a for loop
|
|
@T.prim_func(s_tir=True)
|
|
def func(A: T.Buffer((10,), "int32")):
|
|
for x in range(10):
|
|
A[x] = x + 1
|
|
if x == 5:
|
|
break
|
|
continue
|
|
|
|
# DeclBuffer
|
|
buffer_decl = tir.DeclBuffer(T.buffer((10,), "int32"), evaluate_stmt)
|
|
|
|
# OpCall
|
|
@T.prim_func(s_tir=True)
|
|
def op_call(A: T.Buffer((10,), "int32"), B: T.Buffer((10,), "int32")):
|
|
Tx.add(A, B, 1.0)
|
|
|
|
return {
|
|
"evaluate": evaluate_stmt,
|
|
"let": let_stmt,
|
|
"for": for_loop,
|
|
"while": while_loop,
|
|
"buffer_store": buffer_store,
|
|
"seq_stmt": seq_stmt,
|
|
"block_realize": block_realize,
|
|
"if_then_else": if_then_else,
|
|
"for_with_break": func.body,
|
|
"decl_buffer": buffer_decl,
|
|
"op_call": op_call.body,
|
|
}
|
|
|
|
|
|
def test_evaluate():
|
|
"""Test evaluate statement."""
|
|
evaluate_stmt = create_test_statements()["evaluate"]
|
|
basic_check(
|
|
evaluate_stmt,
|
|
"\n".join(["Evaluate", "\tAdd", "\t\tVar", "\t\tIntImm"]),
|
|
"\n".join(["Var", "IntImm", "Add", "Evaluate"]),
|
|
)
|
|
|
|
|
|
def test_let():
|
|
"""Test let statement (Bind + SeqStmt)."""
|
|
let_stmt = create_test_statements()["let"]
|
|
basic_check(
|
|
let_stmt,
|
|
"\n".join(
|
|
[
|
|
"SeqStmt",
|
|
"\tBind",
|
|
"\t\tIntImm",
|
|
"\tEvaluate",
|
|
"\t\tAdd",
|
|
"\t\t\tVar",
|
|
"\t\t\tIntImm",
|
|
]
|
|
),
|
|
"\n".join(["IntImm", "Bind", "Var", "IntImm", "Add", "Evaluate", "SeqStmt"]),
|
|
)
|
|
|
|
|
|
def test_for():
|
|
"""Test for loop statement."""
|
|
for_loop = create_test_statements()["for"]
|
|
basic_check(
|
|
for_loop,
|
|
"\n".join(
|
|
["For", "\tIntImm", "\tIntImm", "\tEvaluate", "\t\tAdd", "\t\t\tVar", "\t\t\tIntImm"]
|
|
),
|
|
"\n".join(["IntImm", "IntImm", "Var", "IntImm", "Add", "Evaluate", "For"]),
|
|
)
|
|
|
|
|
|
def test_while():
|
|
"""Test while loop statement."""
|
|
while_loop = create_test_statements()["while"]
|
|
basic_check(
|
|
while_loop,
|
|
"\n".join(
|
|
[
|
|
"While",
|
|
"\tLT",
|
|
"\t\tVar",
|
|
"\t\tIntImm",
|
|
"\tEvaluate",
|
|
"\t\tAdd",
|
|
"\t\t\tVar",
|
|
"\t\t\tIntImm",
|
|
]
|
|
),
|
|
"\n".join(["Var", "IntImm", "LT", "Var", "IntImm", "Add", "Evaluate", "While"]),
|
|
)
|
|
|
|
|
|
def test_buffer_store():
|
|
"""Test buffer store statement."""
|
|
buffer_store = create_test_statements()["buffer_store"]
|
|
basic_check(
|
|
buffer_store,
|
|
"\n".join(["BufferStore", "\tAdd", "\t\tVar", "\t\tIntImm", "\tIntImm"]),
|
|
"\n".join(["Var", "IntImm", "Add", "IntImm", "BufferStore"]),
|
|
)
|
|
|
|
|
|
def test_seq_stmt():
|
|
"""Test sequence statement."""
|
|
seq_stmt = create_test_statements()["seq_stmt"]
|
|
basic_check(
|
|
seq_stmt,
|
|
"\n".join(
|
|
[
|
|
"SeqStmt",
|
|
"\tEvaluate",
|
|
"\t\tAdd",
|
|
"\t\t\tVar",
|
|
"\t\t\tIntImm",
|
|
"\tFor",
|
|
"\t\tIntImm",
|
|
"\t\tIntImm",
|
|
"\t\tEvaluate",
|
|
"\t\t\tAdd",
|
|
"\t\t\t\tVar",
|
|
"\t\t\t\tIntImm",
|
|
]
|
|
),
|
|
"\n".join(
|
|
[
|
|
"Var",
|
|
"IntImm",
|
|
"Add",
|
|
"Evaluate",
|
|
"IntImm",
|
|
"IntImm",
|
|
"Var",
|
|
"IntImm",
|
|
"Add",
|
|
"Evaluate",
|
|
"For",
|
|
"SeqStmt",
|
|
]
|
|
),
|
|
)
|
|
|
|
|
|
def test_block_realize():
|
|
"""Test block realize statement."""
|
|
block_realize = create_test_statements()["block_realize"]
|
|
basic_check(
|
|
block_realize,
|
|
"\n".join(
|
|
[
|
|
"BlockRealize",
|
|
"\tIntImm",
|
|
"\tIntImm",
|
|
"\tBlock",
|
|
"\t\tEvaluate",
|
|
"\t\t\tAdd",
|
|
"\t\t\t\tVar",
|
|
"\t\t\t\tIntImm",
|
|
]
|
|
),
|
|
"\n".join(
|
|
[
|
|
"IntImm",
|
|
"IntImm",
|
|
"IntImm",
|
|
"IntImm",
|
|
"Var",
|
|
"IntImm",
|
|
"Add",
|
|
"Evaluate",
|
|
"Block",
|
|
"BlockRealize",
|
|
]
|
|
),
|
|
)
|
|
|
|
|
|
def test_if_then_else():
|
|
"""Test if-then-else statement."""
|
|
if_then_else = create_test_statements()["if_then_else"]
|
|
basic_check(
|
|
if_then_else,
|
|
"\n".join(
|
|
[
|
|
"IfThenElse",
|
|
"\tLT",
|
|
"\t\tVar",
|
|
"\t\tIntImm",
|
|
"\tEvaluate",
|
|
"\t\tAdd",
|
|
"\t\t\tVar",
|
|
"\t\t\tIntImm",
|
|
"\tEvaluate",
|
|
"\t\tAdd",
|
|
"\t\t\tVar",
|
|
"\t\t\tIntImm",
|
|
]
|
|
),
|
|
"\n".join(
|
|
[
|
|
"Var",
|
|
"IntImm",
|
|
"LT",
|
|
"Var",
|
|
"IntImm",
|
|
"Add",
|
|
"Evaluate",
|
|
"Var",
|
|
"IntImm",
|
|
"Add",
|
|
"Evaluate",
|
|
"IfThenElse",
|
|
]
|
|
),
|
|
)
|
|
|
|
|
|
def test_for_with_break_continue():
|
|
"""Test for loop with break and continue statements."""
|
|
for_with_break = create_test_statements()["for_with_break"]
|
|
basic_check(
|
|
for_with_break,
|
|
"\n".join(
|
|
[
|
|
"For",
|
|
"\tIntImm",
|
|
"\tIntImm",
|
|
"\tSeqStmt",
|
|
"\t\tBufferStore",
|
|
"\t\t\tAdd",
|
|
"\t\t\t\tVar",
|
|
"\t\t\t\tIntImm",
|
|
"\t\t\tVar",
|
|
"\t\tIfThenElse",
|
|
"\t\t\tEQ",
|
|
"\t\t\t\tVar",
|
|
"\t\t\t\tIntImm",
|
|
"\t\t\tEvaluate",
|
|
"\t\t\t\tExpr::Call",
|
|
"\t\tEvaluate",
|
|
"\t\t\tExpr::Call",
|
|
]
|
|
),
|
|
"\n".join(
|
|
[
|
|
"IntImm",
|
|
"IntImm",
|
|
"Var",
|
|
"IntImm",
|
|
"Add",
|
|
"Var",
|
|
"BufferStore",
|
|
"Var",
|
|
"IntImm",
|
|
"EQ",
|
|
"Expr::Call",
|
|
"Evaluate",
|
|
"IfThenElse",
|
|
"Expr::Call",
|
|
"Evaluate",
|
|
"SeqStmt",
|
|
"For",
|
|
]
|
|
),
|
|
)
|
|
|
|
|
|
def test_decl_buffer():
|
|
"""Test buffer declaration statement."""
|
|
buffer_decl = create_test_statements()["decl_buffer"]
|
|
basic_check(
|
|
buffer_decl,
|
|
"\n".join(["DeclBuffer", "\tEvaluate", "\t\tAdd", "\t\t\tVar", "\t\t\tIntImm"]),
|
|
"\n".join(["Var", "IntImm", "Add", "Evaluate", "DeclBuffer"]),
|
|
)
|
|
|
|
|
|
def test_op_call():
|
|
"""Test op call statement"""
|
|
op_call = create_test_statements()["op_call"]
|
|
basic_check(
|
|
op_call,
|
|
"\n".join(
|
|
[
|
|
"OpCall",
|
|
"\tBufferRegion",
|
|
"\t\tIntImm",
|
|
"\t\tIntImm",
|
|
"\tBufferRegion",
|
|
"\t\tIntImm",
|
|
"\t\tIntImm",
|
|
"\tExpr::FloatImm",
|
|
]
|
|
),
|
|
"\n".join(
|
|
[
|
|
"IntImm",
|
|
"IntImm",
|
|
"BufferRegion",
|
|
"IntImm",
|
|
"IntImm",
|
|
"BufferRegion",
|
|
"Expr::FloatImm",
|
|
"OpCall",
|
|
]
|
|
),
|
|
)
|
|
|
|
|
|
def test_stmt_expr_mutator():
|
|
"""Test StmtExprMutator."""
|
|
evaluate_stmt = create_test_statements()["evaluate"]
|
|
mutator = StmtExprMutatorPrinter()
|
|
result = mutator.visit_stmt(evaluate_stmt)
|
|
tvm.ir.assert_structural_equal(result, evaluate_stmt)
|
|
|
|
expected = "\n".join(["Var", "IntImm", "Add", "Evaluate"])
|
|
assert str(mutator.log) == expected
|
|
|
|
|
|
def test_stmt_expr_visitor():
|
|
"""Test StmtExprVisitor."""
|
|
evaluate_stmt = create_test_statements()["evaluate"]
|
|
visitor = StmtExprASTPrinter()
|
|
visitor.visit_stmt(evaluate_stmt)
|
|
expected = "\n".join(["Evaluate", "\tAdd", "\t\tVar", "\t\tIntImm"])
|
|
assert str(visitor.log) == expected
|
|
|
|
|
|
class NegateIntImmMutator(StmtExprMutator):
|
|
"""Mutator that negates all integer immediates."""
|
|
|
|
def visit_int_imm_(self, op):
|
|
# Create a new IntImm with negated value
|
|
return tir.IntImm(op.ty.dtype, -op.value)
|
|
|
|
|
|
def test_mutator_transformation():
|
|
"""Test that mutator actually transforms the AST."""
|
|
evaluate_stmt = create_test_statements()["evaluate"]
|
|
mutator = NegateIntImmMutator()
|
|
result = mutator.visit_stmt(evaluate_stmt)
|
|
|
|
# The original has value 10, the transformed should have -10
|
|
assert isinstance(evaluate_stmt.value, tir.Add)
|
|
assert isinstance(evaluate_stmt.value.b, tir.IntImm)
|
|
assert evaluate_stmt.value.b.value == 10
|
|
|
|
assert isinstance(result.value, tir.Add)
|
|
assert isinstance(result.value.b, tir.IntImm)
|
|
assert result.value.b.value == -10
|
|
|
|
|
|
class InheritVsMixin:
|
|
"""Test inheriting vs mixing in with StmtVisitor/StmtMutator."""
|
|
|
|
class InheritedVisitor(StmtVisitor):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.log = ASTLog()
|
|
|
|
def visit_for_(self, op):
|
|
self.log.add("InheritedVisitor::For")
|
|
super().visit_for_(op)
|
|
|
|
class DerivedVisitor(InheritedVisitor):
|
|
def visit_for_(self, op):
|
|
self.log.add("DerivedVisitor::For")
|
|
super().visit_for_(op)
|
|
|
|
class BaseMutator(StmtMutator):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.log = ASTLog()
|
|
|
|
def visit_for_(self, op):
|
|
self.log.add("BaseMutator::For")
|
|
return super().visit_for_(op)
|
|
|
|
class DerivedMutator(BaseMutator):
|
|
def visit_for_(self, op):
|
|
self.log.add("DerivedMutator::For")
|
|
return super().visit_for_(op)
|
|
|
|
|
|
def test_inheritance():
|
|
"""Test inheritance with visitor and mutator classes."""
|
|
for_loop = create_test_statements()["for"]
|
|
|
|
# Test inherited visitor
|
|
visitor = InheritVsMixin.DerivedVisitor()
|
|
visitor.visit_stmt(for_loop)
|
|
expected = "\n".join(["DerivedVisitor::For", "InheritedVisitor::For"])
|
|
assert str(visitor.log) == expected
|
|
|
|
# Test derived mutator
|
|
mutator = InheritVsMixin.DerivedMutator()
|
|
result = mutator.visit_stmt(for_loop)
|
|
tvm.ir.assert_structural_equal(result, for_loop)
|
|
expected = "\n".join(["DerivedMutator::For", "BaseMutator::For"])
|
|
assert str(mutator.log) == expected
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|