wehub-resource-sync
400 lines
12 KiB
Python
400 lines
12 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=missing-function-docstring,missing-module-docstring
|
|
# mypy: ignore-errors
|
|
# ruff: noqa: E501, F401
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm import s_tir, tirx
|
|
from tvm.s_tir.schedule import Instruction, InstructionKind, LoopRV, SBlockRV, Trace
|
|
from tvm.s_tir.schedule.testing import assert_structural_equal_ignore_global_symbol
|
|
from tvm.script import tirx as T
|
|
|
|
# pylint: disable=no-member,invalid-name,unused-variable
|
|
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def elementwise(a: T.handle, c: T.handle) -> None:
|
|
A = T.match_buffer(a, (128, 128))
|
|
B = T.sblock_alloc_buffer((128, 128))
|
|
C = T.match_buffer(c, (128, 128))
|
|
for i, j in T.grid(128, 128):
|
|
with T.sblock("B"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
B[vi, vj] = A[vi, vj] * 2.0
|
|
for i, j in T.grid(128, 128):
|
|
with T.sblock("C"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
C[vi, vj] = B[vi, vj] + 1.0
|
|
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def elementwise_inlined(a: T.handle, c: T.handle) -> None:
|
|
A = T.match_buffer(a, (128, 128))
|
|
C = T.match_buffer(c, (128, 128))
|
|
for i, j in T.grid(128, 128):
|
|
with T.sblock("C"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
C[vi, vj] = A[vi, vj] * 2.0 + 1.0
|
|
|
|
|
|
# pylint: enable=no-member,invalid-name,unused-variable
|
|
|
|
|
|
def _make_get_sblock(name, output):
|
|
return Instruction(
|
|
kind=InstructionKind.get("GetSBlock"),
|
|
inputs=[],
|
|
attrs=[name, "main"],
|
|
outputs=[output],
|
|
)
|
|
|
|
|
|
def _make_get_loops(input, outputs): # pylint: disable=redefined-builtin
|
|
return Instruction(
|
|
kind=InstructionKind.get("GetLoops"),
|
|
inputs=[input],
|
|
attrs=[],
|
|
outputs=outputs,
|
|
)
|
|
|
|
|
|
def _make_compute_inline(input): # pylint: disable=redefined-builtin
|
|
return Instruction(
|
|
kind=InstructionKind.get("ComputeInline"),
|
|
inputs=[input],
|
|
attrs=[],
|
|
outputs=[],
|
|
)
|
|
|
|
|
|
def _make_split(inputs, outputs): # pylint: disable=redefined-builtin
|
|
return Instruction(
|
|
kind=InstructionKind.get("Split"),
|
|
inputs=inputs,
|
|
attrs=[T.bool(True), T.bool(False)],
|
|
outputs=outputs,
|
|
)
|
|
|
|
|
|
def _make_enter_postproc():
|
|
return Instruction(
|
|
kind=InstructionKind.get("EnterPostproc"),
|
|
inputs=[],
|
|
attrs=[],
|
|
outputs=[],
|
|
)
|
|
|
|
|
|
def _make_annotate(block: SBlockRV, annotation: str):
|
|
return Instruction(
|
|
kind=InstructionKind.get("Annotate"),
|
|
inputs=[block, annotation],
|
|
attrs=["meta_schedule.auto_tensorize"],
|
|
outputs=[],
|
|
)
|
|
|
|
|
|
def _make_trace_1(b0, l1, l2): # pylint: disable=invalid-name
|
|
return Trace(
|
|
insts=[
|
|
_make_get_sblock(name="block", output=b0),
|
|
_make_get_loops(input=b0, outputs=[l1, l2]),
|
|
],
|
|
decisions={},
|
|
)
|
|
|
|
|
|
def _make_trace_2(b0): # pylint: disable=invalid-name
|
|
return Trace(
|
|
insts=[
|
|
_make_get_sblock(name="B", output=b0),
|
|
_make_compute_inline(input=b0),
|
|
],
|
|
decisions={},
|
|
)
|
|
|
|
|
|
def _make_trace_3(b0, b1, add_postproc): # pylint: disable=invalid-name
|
|
if add_postproc:
|
|
insts = [
|
|
_make_get_sblock(name="B", output=b0),
|
|
_make_compute_inline(input=b0),
|
|
_make_get_sblock(name="C", output=b1),
|
|
_make_enter_postproc(),
|
|
_make_compute_inline(input=b1),
|
|
]
|
|
else:
|
|
insts = [
|
|
_make_get_sblock(name="B", output=b0),
|
|
_make_compute_inline(input=b0),
|
|
_make_get_sblock(name="C", output=b1),
|
|
]
|
|
return Trace(insts=insts, decisions={})
|
|
|
|
|
|
def _make_trace_4(b0, l1, l2, l3): # pylint: disable=invalid-name
|
|
return Trace(
|
|
insts=[
|
|
_make_get_sblock(name="B", output=b0),
|
|
_make_get_loops(input=b0, outputs=[l1]),
|
|
_make_split([l1, None, T.int32(32)], [l2, l3]),
|
|
],
|
|
decisions={},
|
|
)
|
|
|
|
|
|
def test_trace_construct_1():
|
|
trace = _make_trace_1(SBlockRV(), LoopRV(), LoopRV())
|
|
assert str(trace) == "\n".join(
|
|
(
|
|
"# from tvm import s_tir",
|
|
"def apply_trace(sch: s_tir.Schedule) -> None:",
|
|
' b0 = sch.get_sblock(name="block", func_name="main")',
|
|
" l1, l2 = sch.get_loops(block=b0)",
|
|
)
|
|
)
|
|
assert len(trace.insts) == 2
|
|
assert len(trace.decisions) == 0
|
|
|
|
|
|
def test_trace_construct_get_decision_1():
|
|
trace = _make_trace_1(SBlockRV(), LoopRV(), LoopRV())
|
|
assert trace.get_decision(trace.insts[0]) is None
|
|
assert trace.get_decision(trace.insts[1]) is None
|
|
|
|
|
|
def test_trace_construct_append_1():
|
|
trace = _make_trace_1(SBlockRV(), LoopRV(), LoopRV())
|
|
trace.append(inst=_make_get_sblock("block2", SBlockRV()))
|
|
assert str(trace) == "\n".join(
|
|
(
|
|
"# from tvm import s_tir",
|
|
"def apply_trace(sch: s_tir.Schedule) -> None:",
|
|
' b0 = sch.get_sblock(name="block", func_name="main")',
|
|
" l1, l2 = sch.get_loops(block=b0)",
|
|
' b3 = sch.get_sblock(name="block2", func_name="main")',
|
|
)
|
|
)
|
|
|
|
|
|
def test_trace_construct_pop_1():
|
|
trace = _make_trace_1(SBlockRV(), LoopRV(), LoopRV())
|
|
last_inst = trace.insts[-1]
|
|
assert trace.pop().same_as(last_inst)
|
|
assert str(trace) == "\n".join(
|
|
(
|
|
"# from tvm import s_tir",
|
|
"def apply_trace(sch: s_tir.Schedule) -> None:",
|
|
' b0 = sch.get_sblock(name="block", func_name="main")',
|
|
)
|
|
)
|
|
|
|
|
|
def test_trace_construct_pop_2():
|
|
trace = Trace([], {})
|
|
assert str(trace) == "\n".join(
|
|
(
|
|
"# from tvm import s_tir",
|
|
"def apply_trace(sch: s_tir.Schedule) -> None:",
|
|
" pass",
|
|
)
|
|
)
|
|
assert trace.pop() is None
|
|
assert str(trace) == "\n".join(
|
|
(
|
|
"# from tvm import s_tir",
|
|
"def apply_trace(sch: s_tir.Schedule) -> None:",
|
|
" pass",
|
|
)
|
|
)
|
|
|
|
|
|
def test_trace_apply_to_schedule():
|
|
trace = _make_trace_2(SBlockRV())
|
|
sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
|
|
trace.apply_to_schedule(sch, remove_postproc=False, decision_provider=None)
|
|
assert_structural_equal_ignore_global_symbol(elementwise_inlined, sch.mod["main"])
|
|
|
|
|
|
def test_trace_as_json_1():
|
|
trace = _make_trace_1(SBlockRV(), LoopRV(), LoopRV())
|
|
obj = trace.as_json()
|
|
assert obj == [
|
|
[
|
|
["GetSBlock", [], ["block", "main"], ["b0"]],
|
|
["GetLoops", ["b0"], [], ["l1", "l2"]],
|
|
],
|
|
[],
|
|
]
|
|
|
|
|
|
def test_trace_simplified_1():
|
|
trace = _make_trace_3(SBlockRV(), SBlockRV(), add_postproc=True)
|
|
assert str(trace) == "\n".join(
|
|
(
|
|
"# from tvm import s_tir",
|
|
"def apply_trace(sch: s_tir.Schedule) -> None:",
|
|
' b0 = sch.get_sblock(name="B", func_name="main")',
|
|
" sch.compute_inline(block=b0)",
|
|
' b1 = sch.get_sblock(name="C", func_name="main")',
|
|
" sch.enter_postproc()",
|
|
" sch.compute_inline(block=b1)",
|
|
)
|
|
)
|
|
trace = trace.simplified(remove_postproc=True)
|
|
assert str(trace) == "\n".join(
|
|
(
|
|
"# from tvm import s_tir",
|
|
"def apply_trace(sch: s_tir.Schedule) -> None:",
|
|
' b0 = sch.get_sblock(name="B", func_name="main")',
|
|
" sch.compute_inline(block=b0)",
|
|
)
|
|
)
|
|
|
|
|
|
def test_trace_simplified_2():
|
|
trace = _make_trace_3(SBlockRV(), SBlockRV(), add_postproc=True)
|
|
assert str(trace) == "\n".join(
|
|
(
|
|
"# from tvm import s_tir",
|
|
"def apply_trace(sch: s_tir.Schedule) -> None:",
|
|
' b0 = sch.get_sblock(name="B", func_name="main")',
|
|
" sch.compute_inline(block=b0)",
|
|
' b1 = sch.get_sblock(name="C", func_name="main")',
|
|
" sch.enter_postproc()",
|
|
" sch.compute_inline(block=b1)",
|
|
)
|
|
)
|
|
trace = trace.simplified(remove_postproc=False)
|
|
assert str(trace) == "\n".join(
|
|
(
|
|
"# from tvm import s_tir",
|
|
"def apply_trace(sch: s_tir.Schedule) -> None:",
|
|
' b0 = sch.get_sblock(name="B", func_name="main")',
|
|
" sch.compute_inline(block=b0)",
|
|
' b1 = sch.get_sblock(name="C", func_name="main")',
|
|
" sch.enter_postproc()",
|
|
" sch.compute_inline(block=b1)",
|
|
)
|
|
)
|
|
|
|
|
|
def test_trace_simplified_3():
|
|
trace = _make_trace_4(SBlockRV(), LoopRV(), LoopRV(), LoopRV()).simplified(
|
|
remove_postproc=False
|
|
)
|
|
assert str(trace) == "\n".join(
|
|
(
|
|
"# from tvm import s_tir",
|
|
"def apply_trace(sch: s_tir.Schedule) -> None:",
|
|
' b0 = sch.get_sblock(name="B", func_name="main")',
|
|
" l1, = sch.get_loops(block=b0)",
|
|
" l2, l3 = sch.split(loop=l1, factors=[None, 32], preserve_unit_iters=True, disable_predication=False)",
|
|
)
|
|
)
|
|
|
|
|
|
def test_apply_json_to_schedule_1():
|
|
trace = _make_trace_2(SBlockRV())
|
|
json_obj = trace.as_json()
|
|
sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
|
|
Trace.apply_json_to_schedule(json_obj, sch)
|
|
assert_structural_equal_ignore_global_symbol(elementwise_inlined, sch.mod["main"])
|
|
|
|
|
|
def test_apply_json_to_schedule_sample_categorical():
|
|
var = tirx.Var("v", "int32")
|
|
trace1 = Trace(
|
|
insts=[
|
|
Instruction(
|
|
kind=InstructionKind.get("SampleCategorical"),
|
|
inputs=[],
|
|
attrs=[[tvm.tirx.IntImm("int32", 3)], [tvm.tirx.FloatImm("float32", 1.0)]],
|
|
outputs=[var],
|
|
)
|
|
],
|
|
decisions={},
|
|
)
|
|
json = trace1.as_json()
|
|
assert str(json) == "[[['SampleCategorical', [], [[3], [T.float32(1.0)]], ['v0']]], []]"
|
|
|
|
sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
|
|
# As long as the application does not fail, it is fine.
|
|
Trace.apply_json_to_schedule(json, sch)
|
|
python_str = sch.trace.as_python()
|
|
assert len(python_str) == 1
|
|
assert python_str[0] == "v0 = sch.sample_categorical(candidates=[3], probs=[1], decision=0)"
|
|
|
|
|
|
def _test_apply_annotation_trace_from_json(annotation: str):
|
|
"""Test applying an annotation works without crashing.
|
|
|
|
Designed to handle some previously failing edge cases like the
|
|
empty string.
|
|
"""
|
|
b0 = SBlockRV()
|
|
trace = Trace(
|
|
insts=[
|
|
_make_get_sblock(name="B", output=b0),
|
|
_make_annotate(block=b0, annotation=annotation),
|
|
],
|
|
decisions={},
|
|
)
|
|
json_obj = trace.as_json()
|
|
sch = tvm.s_tir.Schedule(elementwise, debug_mask="all")
|
|
Trace.apply_json_to_schedule(json_obj, sch)
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def elementwise_expected(a: T.handle, c: T.handle) -> None:
|
|
A = T.match_buffer(a, (128, 128))
|
|
B = T.sblock_alloc_buffer((128, 128))
|
|
C = T.match_buffer(c, (128, 128))
|
|
for i, j in T.grid(128, 128):
|
|
with T.sblock("B"):
|
|
T.sblock_attr({"meta_schedule.auto_tensorize": annotation})
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
B[vi, vj] = A[vi, vj] * 2.0
|
|
for i, j in T.grid(128, 128):
|
|
with T.sblock("C"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
C[vi, vj] = B[vi, vj] + 1.0
|
|
|
|
assert_structural_equal_ignore_global_symbol(elementwise_expected, sch.mod["main"])
|
|
|
|
|
|
def test_apply_annotation_from_json():
|
|
# Something reasonable
|
|
_test_apply_annotation_trace_from_json("SSRSSR")
|
|
|
|
# The empty string
|
|
_test_apply_annotation_trace_from_json("")
|
|
|
|
# A string of two quotation marks
|
|
_test_apply_annotation_trace_from_json('""')
|
|
|
|
# A string of one quotation mark
|
|
_test_apply_annotation_trace_from_json('"')
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|