chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,399 @@
|
||||
# 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()
|
||||
Reference in New Issue
Block a user