123 lines
4.1 KiB
Python
123 lines
4.1 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.
|
|
|
|
import re
|
|
|
|
import pytest
|
|
|
|
import tvm
|
|
import tvm.contrib.hexagon as hexagon
|
|
import tvm.testing
|
|
from tvm.script import ir as I
|
|
from tvm.script import tirx as T
|
|
from tvm.testing import env
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def register_linker():
|
|
original_linker = hexagon.hexagon_link()
|
|
# Register a phony linker, so that we can test codegen without a Hexagon toolchain.
|
|
hexagon.register_linker(lambda: "/bin/true")
|
|
yield None
|
|
# Restore registration.
|
|
hexagon.register_linker(original_linker)
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_hexagon(), reason="need hexagon")
|
|
def test_basic():
|
|
target = tvm.target.Target("qcom/hexagon-v66")
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(
|
|
C: T.Buffer((128,), "uint8"),
|
|
A: T.Buffer((128,), "uint8"),
|
|
A_1: T.Buffer((128,), "uint8"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i in range(128):
|
|
with T.sblock("C"):
|
|
v_i = T.axis.spatial(128, i)
|
|
T.reads(A[v_i], A_1[v_i])
|
|
T.writes(C[v_i])
|
|
C[v_i] = A[v_i] + A_1[v_i]
|
|
|
|
hexm = tvm.compile(Module, target=tvm.target.Target(target, target))
|
|
asm = hexm.inspect_source("s")
|
|
vadds = re.findall(r"v[0-9]+.b = vadd\(v[0-9]+.b,v[0-9]+.b\)", asm)
|
|
assert vadds # Check that it's non-empty
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_hexagon(), reason="need hexagon")
|
|
def test_llvm_target_features():
|
|
target = tvm.target.Target("qcom/hexagon-v66")
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def add_one(C: T.Buffer((128,), "int32"), A: T.Buffer((128,), "uint8")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i in range(128):
|
|
with T.sblock("C"):
|
|
v_i = T.axis.spatial(128, i)
|
|
T.reads(A[v_i])
|
|
T.writes(C[v_i])
|
|
C[v_i] = T.Cast("int32", A[v_i]) + 1
|
|
|
|
m = tvm.compile(Module, target=tvm.target.Target(target, target))
|
|
llvm_ir = m.inspect_source("ll")
|
|
# Make sure we find +hvx-length128b in "attributes".
|
|
fs = re.findall(r"attributes.*\+hvx-length128b", llvm_ir)
|
|
assert fs # Check that it's non-empty
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_hexagon(), reason="need hexagon")
|
|
def test_llvm_options():
|
|
target = tvm.target.Target(
|
|
{
|
|
"kind": "hexagon",
|
|
"mtriple": "hexagon",
|
|
"mcpu": "hexagonv66",
|
|
"mattr": ["+hvxv66", "+hvx-length128b"],
|
|
"num-cores": 4,
|
|
"vtcm-capacity": 262144,
|
|
"llvm-options": ["-hexagon-noopt"],
|
|
}
|
|
)
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(compute: T.Buffer((10,), "int32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for _ in range(10):
|
|
with T.sblock("compute"):
|
|
v__ = T.axis.spatial(10, _)
|
|
T.reads()
|
|
T.writes(compute[v__])
|
|
compute[v__] = 0
|
|
|
|
# Check that BuildHexagon hasn't crashed because of target attribute
|
|
# type mismatch.
|
|
tvm.compile(Module, target=tvm.target.Target(target, target))
|
|
assert re.search("-hexagon-noopt", str(target))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|