wehub-resource-sync
119 lines
3.4 KiB
Python
119 lines
3.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.
|
|
# ruff: noqa: F401
|
|
|
|
from typing import Optional
|
|
|
|
import pytest
|
|
from tvm_ffi.access_path import AccessPath
|
|
|
|
from tvm.script import tirx as T
|
|
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def _func():
|
|
T.evaluate(-1)
|
|
T.evaluate(1)
|
|
T.evaluate(2)
|
|
T.evaluate(3)
|
|
T.evaluate(4)
|
|
T.evaluate(5)
|
|
T.evaluate(6)
|
|
T.evaluate(7)
|
|
|
|
|
|
def test_annotation_multi_access_paths():
|
|
result = _func.with_attr("global_symbol", "main").script(
|
|
path_to_annotate={
|
|
AccessPath.root().attr("body").attr("seq").array_item(1): "annotation 1",
|
|
AccessPath.root().attr("body").attr("seq").array_item(3): "annotation 3",
|
|
AccessPath.root().attr("body").attr("seq").array_item(5): "annotation 5",
|
|
AccessPath.root().attr("body").attr("seq").array_item(7): "annotation 7",
|
|
}
|
|
)
|
|
assert (
|
|
result
|
|
== """# from tvm.script import tirx as T
|
|
# from tvm.tirx.layout import Axis
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def main():
|
|
T.evaluate(-1)
|
|
T.evaluate(1) # annotation 1
|
|
T.evaluate(2)
|
|
T.evaluate(3) # annotation 3
|
|
T.evaluate(4)
|
|
T.evaluate(5) # annotation 5
|
|
T.evaluate(6)
|
|
T.evaluate(7) # annotation 7"""
|
|
)
|
|
|
|
|
|
def test_annotate_from_multi_obj():
|
|
result = _func.with_attr("global_symbol", "main").script(
|
|
obj_to_annotate={
|
|
_func.body.seq[1]: "annotation 1",
|
|
_func.body.seq[3]: "annotation 3",
|
|
_func.body.seq[5]: "annotation 5",
|
|
_func.body.seq[7]: "annotation 7",
|
|
}
|
|
)
|
|
assert (
|
|
result
|
|
== """# from tvm.script import tirx as T
|
|
# from tvm.tirx.layout import Axis
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def main():
|
|
T.evaluate(-1)
|
|
T.evaluate(1) # annotation 1
|
|
T.evaluate(2)
|
|
T.evaluate(3) # annotation 3
|
|
T.evaluate(4)
|
|
T.evaluate(5) # annotation 5
|
|
T.evaluate(6)
|
|
T.evaluate(7) # annotation 7"""
|
|
)
|
|
|
|
|
|
def test_disable_concise_scoping_when_scope_annotated():
|
|
@T.prim_func(s_tir=True)
|
|
def _func():
|
|
x = 1
|
|
y = x + 1
|
|
T.evaluate(y - 1)
|
|
|
|
# In fork, each bare `x = expr` lowers to AllocBuffer + BufferStore (local_scalar);
|
|
# the printer fuses each pair into a single `y: T.int32 = x + 1` line. Annotate the
|
|
# AllocBuffer that originates this fused line.
|
|
result = _func.with_attr("global_symbol", "main").script(
|
|
obj_to_annotate={
|
|
_func.body.seq[2]: "annotation 1",
|
|
}
|
|
)
|
|
assert (
|
|
result
|
|
== """# from tvm.script import tirx as T
|
|
# from tvm.tirx.layout import Axis
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def main():
|
|
x: T.int32 = 1
|
|
y: T.int32 = x + 1 # annotation 1
|
|
T.evaluate(y - 1)"""
|
|
)
|