wehub-resource-sync
261 lines
6.4 KiB
Python
261 lines
6.4 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 T.inline / T.inline with Python LEGB scoping semantics."""
|
|
|
|
from tvm.ir import assert_structural_equal
|
|
from tvm.script import tirx as T
|
|
|
|
# Module-level constant for testing global visibility
|
|
MODULE_CONST = 42
|
|
|
|
|
|
def test_local_shadows_enclosing():
|
|
"""A local parameter in the inline shadows a variable from the enclosing scope."""
|
|
|
|
@T.prim_func(private=True)
|
|
def func(A: T.Buffer((128,), "int32")) -> None:
|
|
T.int32(10)
|
|
|
|
@T.inline
|
|
def write(x):
|
|
# x here is the parameter, not the enclosing x=10
|
|
A[0] = x
|
|
|
|
write(T.int32(20))
|
|
|
|
@T.prim_func(private=True)
|
|
def expected(A: T.Buffer((128,), "int32")) -> None:
|
|
T.int32(10)
|
|
A[0] = T.int32(20)
|
|
|
|
assert_structural_equal(func, expected)
|
|
|
|
|
|
def test_enclosing_variable_capture():
|
|
"""Inline captures a variable from its enclosing scope (not a parameter)."""
|
|
val = 64
|
|
|
|
@T.inline
|
|
def write_val(A):
|
|
A[0] = val
|
|
|
|
@T.prim_func(private=True)
|
|
def func(A: T.Buffer((128,), "int32")) -> None:
|
|
write_val(A)
|
|
|
|
@T.prim_func(private=True)
|
|
def expected(A: T.Buffer((128,), "int32")) -> None:
|
|
A[0] = 64
|
|
|
|
assert_structural_equal(func, expected)
|
|
|
|
|
|
def test_nested_inline():
|
|
"""Inner inline can call outer inline (inline-in-inline)."""
|
|
|
|
@T.inline
|
|
def add_one(A):
|
|
A[0] = A[0] + 1
|
|
|
|
@T.inline
|
|
def add_two(A):
|
|
add_one(A)
|
|
add_one(A)
|
|
|
|
@T.prim_func(private=True)
|
|
def func(A: T.Buffer((128,), "int32")) -> None:
|
|
add_two(A)
|
|
|
|
@T.prim_func(private=True)
|
|
def expected(A: T.Buffer((128,), "int32")) -> None:
|
|
A[0] = A[0] + 1
|
|
A[0] = A[0] + 1
|
|
|
|
assert_structural_equal(func, expected)
|
|
|
|
|
|
def test_module_globals_visible():
|
|
"""Inline can see module-level globals."""
|
|
|
|
@T.inline
|
|
def write_const(A):
|
|
A[0] = MODULE_CONST
|
|
|
|
@T.prim_func(private=True)
|
|
def func(A: T.Buffer((128,), "int32")) -> None:
|
|
write_const(A)
|
|
|
|
@T.prim_func(private=True)
|
|
def expected(A: T.Buffer((128,), "int32")) -> None:
|
|
A[0] = 42
|
|
|
|
assert_structural_equal(func, expected)
|
|
|
|
|
|
def test_shadowing_in_inner_scope():
|
|
"""An inline defined inside a for-loop captures the loop variable."""
|
|
|
|
@T.prim_func(private=True)
|
|
def func(A: T.Buffer((10,), "int32")) -> None:
|
|
for i in T.serial(10):
|
|
|
|
@T.inline
|
|
def write_i(A):
|
|
A[i] = i
|
|
|
|
write_i(A)
|
|
|
|
@T.prim_func(private=True)
|
|
def expected(A: T.Buffer((10,), "int32")) -> None:
|
|
for i in range(10):
|
|
A[i] = i
|
|
|
|
assert_structural_equal(func, expected)
|
|
|
|
|
|
def test_lexical_not_dynamic():
|
|
"""An inline defined outside prim_func does NOT see the caller's locals.
|
|
Specifically, x_value captured at definition time (128) is used,
|
|
not the loop variable x_value from the caller."""
|
|
x_value = 128
|
|
|
|
@T.inline
|
|
def static_capture(A, B):
|
|
B[()] = A[x_value]
|
|
|
|
@T.prim_func(private=True)
|
|
def func(A: T.Buffer((1024,), "int32"), B: T.Buffer((), "int32")) -> None:
|
|
for x_value in T.serial(10):
|
|
static_capture(A, B)
|
|
|
|
@T.prim_func(private=True)
|
|
def expected(A: T.Buffer((1024,), "int32"), B: T.Buffer((), "int32")) -> None:
|
|
for x_value in range(10):
|
|
B[()] = A[128]
|
|
|
|
assert_structural_equal(func, expected)
|
|
|
|
|
|
def test_callback_pattern():
|
|
"""Inline passed as an argument to another inline."""
|
|
|
|
@T.inline
|
|
def apply_fn(fn, A):
|
|
fn(A)
|
|
|
|
@T.inline
|
|
def inc(A):
|
|
A[0] = A[0] + 1
|
|
|
|
@T.prim_func(private=True)
|
|
def func(A: T.Buffer((128,), "int32")) -> None:
|
|
apply_fn(inc, A)
|
|
|
|
@T.prim_func(private=True)
|
|
def expected(A: T.Buffer((128,), "int32")) -> None:
|
|
A[0] = A[0] + 1
|
|
|
|
assert_structural_equal(func, expected)
|
|
|
|
|
|
def test_sibling_calls():
|
|
"""Two independent inlines called in sequence."""
|
|
|
|
@T.inline
|
|
def write_a(A):
|
|
A[0] = 1
|
|
|
|
@T.inline
|
|
def write_b(A):
|
|
A[1] = 2
|
|
|
|
@T.prim_func(private=True)
|
|
def func(A: T.Buffer((128,), "int32")) -> None:
|
|
write_a(A)
|
|
write_b(A)
|
|
|
|
@T.prim_func(private=True)
|
|
def expected(A: T.Buffer((128,), "int32")) -> None:
|
|
A[0] = 1
|
|
A[1] = 2
|
|
|
|
assert_structural_equal(func, expected)
|
|
|
|
|
|
def test_recursive_inline():
|
|
"""Recursive inline (defined inside prim_func)."""
|
|
|
|
# fmt: off
|
|
@T.prim_func(private=True)
|
|
def func():
|
|
T.device_entry()
|
|
for x in T.serial(10):
|
|
|
|
@T.inline
|
|
def add(x, c):
|
|
if c > 0:
|
|
add(x, c - 1)
|
|
T.evaluate(x)
|
|
|
|
add(x, 3)
|
|
|
|
@T.prim_func(private=True)
|
|
def expected():
|
|
T.device_entry()
|
|
for x in range(10):
|
|
T.evaluate(x)
|
|
T.evaluate(x)
|
|
T.evaluate(x)
|
|
T.evaluate(x)
|
|
# fmt: on
|
|
|
|
assert_structural_equal(func, expected)
|
|
|
|
|
|
def test_late_binding():
|
|
"""Variable defined after inline but before call (inside prim_func)."""
|
|
|
|
@T.prim_func(private=True)
|
|
def func(A: T.Buffer((128,), "int32")) -> None:
|
|
@T.inline
|
|
def write(A):
|
|
A[0] = val
|
|
|
|
val = T.int32(99)
|
|
write(A)
|
|
|
|
@T.prim_func(private=True)
|
|
def expected(A: T.Buffer((128,), "int32")) -> None:
|
|
val = T.int32(99)
|
|
A[0] = val
|
|
|
|
assert_structural_equal(func, expected)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_local_shadows_enclosing()
|
|
test_enclosing_variable_capture()
|
|
test_nested_inline()
|
|
test_module_globals_visible()
|
|
test_shadowing_in_inner_scope()
|
|
test_lexical_not_dynamic()
|
|
test_callback_pattern()
|
|
test_sibling_calls()
|
|
test_recursive_inline()
|
|
test_late_binding()
|
|
print("All tests passed!")
|