Skip to content

Commit

Permalink
fixup! [USMP] Adding support for U4 usecase
Browse files Browse the repository at this point in the history
Change-Id: I9676b5198ae5485815ac601eb30d496d09f981ad
  • Loading branch information
manupak committed Mar 29, 2022
1 parent 1da8837 commit b7ba416
Showing 1 changed file with 6 additions and 6 deletions.
12 changes: 6 additions & 6 deletions tests/python/relay/aot/test_c_device_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@

import numpy as np
import pytest
import re

from tvm import relay
from tvm.ir.module import IRModule
Expand Down Expand Up @@ -133,7 +134,6 @@ def compile_to_main_func(interface_api="c", use_unpacked_api=True):
def test_device_api_hooks_unpacked_api(device_api_main_func):
"""Check for Device API hooks with unpacked internal calls"""
main_func = device_api_main_func(interface_api="c", use_unpacked_api=True)
input_name = main_func.params[0].name

# Activate Device
assert (
Expand All @@ -150,12 +150,12 @@ def test_device_api_hooks_unpacked_api(device_api_main_func):
+ " device_context_ethos_u))\n"
)
# Device Call
assert (
str(main_func.body[1][0][0][1])
== "tir.tvm_check_return(0, -1, tir.call_extern("
+ '"tvmgen_default_ethos_u_main_0",'
+ f" {input_name}_buffer_var, output_buffer_var, device_context_ethos_u))\n"
# We dont need to check exact input and output var names in this test.
# Hence, using a regex to cover any legal I/O name.
regex = re.compile(
'tir\.tvm_check_return\(0, -1, tir\.call_extern\("tvmgen_default_ethos_u_main_0", \w+, \w+, device_context_ethos_u\)\)'
)
assert regex.match(str(main_func.body[1][0][0][1]))
# Close Device
assert (
str(main_func.body[1][0][0][2])
Expand Down

0 comments on commit b7ba416

Please sign in to comment.