wehub-resource-sync
346 lines
11 KiB
Python
346 lines
11 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.
|
|
# ruff: noqa: F841
|
|
import pytest
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm.script import ir as I
|
|
from tvm.script import tirx as T
|
|
|
|
|
|
def test_annotate_entry_func_single_primfunc():
|
|
@tvm.script.ir_module
|
|
class MockModule:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def func1(A: T.Buffer((16,), "float32")):
|
|
for i in T.serial(16):
|
|
if i == 5:
|
|
if i == 5:
|
|
A[i] = 0.0
|
|
|
|
mod = MockModule
|
|
assert mod
|
|
assert "tirx.is_entry_func" not in (mod["func1"].attrs or {})
|
|
after = tvm.tirx.transform.AnnotateEntryFunc()(mod)
|
|
assert (
|
|
after["func1"].attrs
|
|
and "tirx.is_entry_func" in after["func1"].attrs
|
|
and after["func1"].attrs["tirx.is_entry_func"]
|
|
)
|
|
|
|
|
|
# Test module
|
|
@tvm.script.ir_module
|
|
class MockModule:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def func1(A: T.Buffer((16,), "float32")):
|
|
for i in T.serial(16):
|
|
if i == 5:
|
|
if i == 5:
|
|
A[i] = 0.0
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def func2(A: T.Buffer((32,), "float32")):
|
|
for i in T.serial(32):
|
|
if i == 15:
|
|
if i == 15:
|
|
A[i] = 0.0
|
|
|
|
|
|
@pytest.mark.xfail
|
|
def test_annotate_entry_func_multiple_primfunc():
|
|
mod = MockModule
|
|
assert mod
|
|
assert "target" not in (mod["func1"].attrs or {})
|
|
assert "target" not in (mod["func2"].attrs or {})
|
|
# This should fail
|
|
after = tvm.tirx.transform.AnnotateEntryFunc()(mod)
|
|
|
|
|
|
def test_bind_target():
|
|
mod = MockModule
|
|
assert mod
|
|
|
|
target = tvm.target.Target("cuda")
|
|
assert "target" not in (mod["func1"].attrs or {})
|
|
assert "target" not in (mod["func2"].attrs or {})
|
|
after = tvm.tirx.transform.BindTarget(target)(mod)
|
|
|
|
assert "target" in after["func1"].attrs
|
|
assert after["func1"].attrs["target"] == target
|
|
assert "target" in after["func2"].attrs
|
|
assert after["func2"].attrs["target"] == target
|
|
|
|
|
|
def test_bind_target_adds_attribute():
|
|
"""BindTarget adds the "target" attribute"""
|
|
|
|
@I.ir_module
|
|
class Before:
|
|
@T.prim_func(s_tir=True)
|
|
def main():
|
|
T.evaluate(0)
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@T.prim_func(s_tir=True)
|
|
def main():
|
|
T.func_attr({"target": T.target("cuda")})
|
|
T.evaluate(0)
|
|
|
|
After = tvm.tirx.transform.BindTarget(tvm.target.Target("cuda"))(Before)
|
|
tvm.ir.assert_structural_equal(After, Expected)
|
|
|
|
|
|
def test_bind_target_with_host_to_exposed_function():
|
|
"""BindTarget adds the host target to externally-exposed functions"""
|
|
|
|
@I.ir_module
|
|
class Before:
|
|
@T.prim_func(s_tir=True)
|
|
def main():
|
|
T.func_attr({"global_symbol": "main"})
|
|
T.evaluate(0)
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@T.prim_func(s_tir=True)
|
|
def main():
|
|
T.func_attr({"global_symbol": "main", "target": T.target("cuda", host="llvm")})
|
|
T.evaluate(0)
|
|
|
|
After = tvm.tirx.transform.BindTarget(tvm.target.Target("cuda", host="llvm"))(Before)
|
|
tvm.ir.assert_structural_equal(After, Expected)
|
|
|
|
|
|
def test_bind_target_with_host_to_internal_function():
|
|
"""Internal functions have a target annotation, but without the host
|
|
|
|
The host portion of the target annotation provides host
|
|
parameters, and is used to expose a function externally as part of
|
|
`MakePackedAPI` and `MakeUnpackedAPI`. For internal functions, no
|
|
external exposure is required, so the host attribute should not be
|
|
used.
|
|
"""
|
|
|
|
@I.ir_module
|
|
class Before:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def main():
|
|
T.evaluate(0)
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def main():
|
|
T.func_attr({"target": T.target("cuda")})
|
|
T.evaluate(0)
|
|
|
|
After = tvm.tirx.transform.BindTarget(tvm.target.Target("cuda", host="llvm"))(Before)
|
|
tvm.ir.assert_structural_equal(After, Expected)
|
|
|
|
|
|
def test_bind_target_ignores_existing():
|
|
"""BindTarget should not replace existing annotations"""
|
|
|
|
@I.ir_module
|
|
class Before:
|
|
@T.prim_func(s_tir=True)
|
|
def main():
|
|
T.func_attr({"target": T.target("nvptx")})
|
|
T.evaluate(0)
|
|
|
|
Expected = Before
|
|
|
|
After = tvm.tirx.transform.BindTarget(tvm.target.Target("cuda"))(Before)
|
|
tvm.ir.assert_structural_equal(After, Expected)
|
|
|
|
|
|
def test_bind_target_updates_host():
|
|
"""BindTarget should update host for existing annotations"""
|
|
|
|
@I.ir_module
|
|
class Before:
|
|
@T.prim_func(s_tir=True)
|
|
def main():
|
|
T.func_attr({"global_symbol": "func", "target": T.target("nvptx")})
|
|
T.evaluate(0)
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@T.prim_func(s_tir=True)
|
|
def main():
|
|
T.func_attr(
|
|
{
|
|
"global_symbol": "func",
|
|
"target": T.target("nvptx", host={"kind": "llvm", "opt-level": 0}),
|
|
}
|
|
)
|
|
T.evaluate(0)
|
|
|
|
After = tvm.tirx.transform.BindTarget(
|
|
tvm.target.Target("cuda", host={"kind": "llvm", "opt-level": 0})
|
|
)(Before)
|
|
tvm.ir.assert_structural_equal(After, Expected)
|
|
|
|
|
|
def test_bind_target_multiple_functions():
|
|
"""BindTarget may apply to multiple functions in a module"""
|
|
|
|
@I.ir_module
|
|
class Before:
|
|
@T.prim_func(s_tir=True)
|
|
def func1():
|
|
T.evaluate(0)
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def func2():
|
|
T.evaluate(0)
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@T.prim_func(s_tir=True)
|
|
def func1():
|
|
T.func_attr({"target": T.target("cuda")})
|
|
T.evaluate(0)
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def func2():
|
|
T.func_attr({"target": T.target("cuda")})
|
|
T.evaluate(0)
|
|
|
|
After = tvm.tirx.transform.BindTarget(tvm.target.Target("cuda"))(Before)
|
|
tvm.ir.assert_structural_equal(After, Expected)
|
|
|
|
|
|
def test_bind_target_with_device_host_call_same_func():
|
|
"""BindTarget should bind the device target to the function if it is called from device"""
|
|
|
|
@I.ir_module
|
|
class Before:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def add(a: T.int32, b: T.int32) -> T.int32:
|
|
return a + b
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def main(
|
|
A: T.Buffer((128, 128), "int32"),
|
|
B: T.Buffer((128, 128), "int32"),
|
|
C: T.Buffer((128, 128), "int32"),
|
|
):
|
|
T.func_attr({"global_symbol": "main"})
|
|
length: T.let[T.int32] = Before.add(64, 64) # Call from host
|
|
for bx in T.thread_binding(length, "blockIdx.x"):
|
|
for tx in T.thread_binding(length, "threadIdx.x"):
|
|
C[bx, tx] = Before.add(A[bx, tx], B[bx, tx]) # Call from device
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def add(a: T.int32, b: T.int32) -> T.int32:
|
|
T.func_attr({"target": T.target("cuda")})
|
|
return a + b
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def add_host(a: T.int32, b: T.int32) -> T.int32:
|
|
T.func_attr({"target": T.target({"kind": "llvm", "opt-level": 0})})
|
|
return a + b
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def main(
|
|
A: T.Buffer((128, 128), "int32"),
|
|
B: T.Buffer((128, 128), "int32"),
|
|
C: T.Buffer((128, 128), "int32"),
|
|
):
|
|
T.func_attr(
|
|
{
|
|
"global_symbol": "main",
|
|
"target": T.target("cuda", host={"kind": "llvm", "opt-level": 0}),
|
|
}
|
|
)
|
|
length: T.let[T.int32] = Expected.add_host(64, 64) # Call from host
|
|
for bx in T.thread_binding(length, "blockIdx.x"):
|
|
for tx in T.thread_binding(length, "threadIdx.x"):
|
|
C[bx, tx] = Expected.add(A[bx, tx], B[bx, tx]) # Call from device
|
|
|
|
After = tvm.tirx.transform.BindTarget(
|
|
tvm.target.Target("cuda", host={"kind": "llvm", "opt-level": 0})
|
|
)(Before)
|
|
tvm.ir.assert_structural_equal(After, Expected)
|
|
|
|
|
|
def test_filter_primfunc():
|
|
mod = MockModule
|
|
assert mod
|
|
# Annotate each function for testing
|
|
mod["func1"] = mod["func1"].with_attr("temp", "test1")
|
|
mod["func2"] = mod["func2"].with_attr("temp", "test2")
|
|
|
|
# Test condition that does not filter out anything
|
|
def checker_filter_out_none(func: tvm.tirx.PrimFunc):
|
|
return "temp" in func.attrs
|
|
|
|
after = tvm.tirx.transform.Filter(checker_filter_out_none)(mod)
|
|
assert len(after.functions) == 2
|
|
# Filtered functions should satisfy the given condition.
|
|
assert checker_filter_out_none(after["func1"])
|
|
assert checker_filter_out_none(after["func2"])
|
|
|
|
# Test condition that selectively filters out primfuncs
|
|
def checker_filter_out_one(func: tvm.tirx.PrimFunc):
|
|
return ("temp" in func.attrs) and func.attrs["temp"] == "test1"
|
|
|
|
after = tvm.tirx.transform.Filter(checker_filter_out_one)(mod)
|
|
assert len(after.functions) == 1
|
|
# Filtered functions should satisfy the given condition.
|
|
assert checker_filter_out_one(after["func1"])
|
|
|
|
# Test condition that filters out everything
|
|
def checker_filter_out_both(func: tvm.tirx.PrimFunc):
|
|
return "invalid_attr" in func.attrs
|
|
|
|
after = tvm.tirx.transform.Filter(checker_filter_out_both)(mod)
|
|
assert len(after.functions) == 0
|
|
|
|
|
|
def test_filter_removes_global_var_map():
|
|
"""Filtering out a function should be identical to never adding it
|
|
|
|
This test is to guard against hidden state in the IRModule that
|
|
remains after filtering. Previously, this was observed in the
|
|
`IRModuleNode::global_var_map_`, which retained entries of
|
|
filtered-out functions.
|
|
"""
|
|
|
|
@I.ir_module
|
|
class Before:
|
|
@T.prim_func(s_tir=True)
|
|
def func():
|
|
T.evaluate(0)
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
pass
|
|
|
|
After = tvm.tirx.transform.Filter(lambda prim_func: False)(Before)
|
|
tvm.ir.assert_structural_equal(After, Expected)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|