fixup! [USMP] Adding support for U4 usecase

Change-Id: I9676b5198ae5485815ac601eb30d496d09f981ad

Author: manupak

tests/python/relay/aot/test_c_device_api.py

@@ -20,6 +20,7 @@
 
 import numpy as np
 import pytest
+import re
 
 from tvm import relay
 from tvm.ir.module import IRModule
@@ -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 (
@@ -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])