def test_tflite_model_u2_usecase_two_models_with_a_single_external_pool(model_urls, usmp_algo):
"""This checks for inference using a single large enough common pool"""
pytest.importorskip("tflite")
import tvm.relay.testing.tf as tf_testing
use_unpacked_api = True
interface_api = "c"
target = tvm.target.Target("c")
workspace_memory_pools = WorkspaceMemoryPools(
[PoolInfo("my_memory_pool", {target: PoolInfo.READ_WRITE_ACCESS})]
)
test_runner = AOTTestRunner(
pass_config={"tir.usmp.enable": True, "tir.usmp.algorithm": usmp_algo},
prologue=f"""
#define MAX(A, B) ((A > B) ? A : B)
__attribute__((section(".data.tvm"), aligned(16)))
static uint8_t my_memory_pool[MAX({_get_workspace_size_define_macro("my_memory_pool", "mod1")},{_get_workspace_size_define_macro("my_memory_pool", "mod2")})];
""",
)
tflite_model_file1 = tf_testing.get_workload_official(
model_urls[0][0],
model_urls[0][1],
)
mod1, inputs1, params1 = create_relay_module_and_inputs_from_tflite_file(tflite_model_file1)
output_list1 = generate_ref_data(mod1, inputs1, params1)
tflite_model_file2 = tf_testing.get_workload_official(
model_urls[1][0],
model_urls[1][1],
)
mod2, inputs2, params2 = create_relay_module_and_inputs_from_tflite_file(tflite_model_file2)
output_list2 = generate_ref_data(mod2, inputs2, params2)
compiled_test_mods = compile_models(
[
AOTTestModel(
name="mod1", module=mod1, inputs=inputs1, outputs=output_list1, params=params1
),
AOTTestModel(
name="mod2", module=mod2, inputs=inputs2, outputs=output_list2, params=params2
),
],
interface_api=interface_api,
use_unpacked_api=use_unpacked_api,
pass_config=test_runner.pass_config,
workspace_memory_pools=workspace_memory_pools,
target=target,
)
for compiled_model in compiled_test_mods:
check_for_no_tvm_backendallocworkspace_calls(compiled_model.executor_factory.lib)
run_and_check(
models=compiled_test_mods,
runner=test_runner,
interface_api=interface_api,
)
def test_quant_mobilenet_tfl():
"""Since in AOT we pass directly the output buffer from the user, in quantized networks sharing the output buffers is not possible.
This is because the output data type is int8 and the intermediate buffer are int32 or int16. We use mobilenet quantized to stress this
situation and verify that the output buffer sharing is disabled in AOT."""
pytest.importorskip("tflite")
import tvm.relay.testing.tf as tf_testing
use_unpacked_api = True
interface_api = "c"
test_runner = AOT_DEFAULT_RUNNER
tflite_model_file = tf_testing.get_workload_official(
"https://storage.googleapis.com/download.tensorflow.org/"
"models/mobilenet_v1_2018_08_02/mobilenet_v1_1.0_224_quant.tgz",
"mobilenet_v1_1.0_224_quant.tflite",
)
with open(tflite_model_file, "rb") as f:
tflite_model_buf = f.read()
data_shape = (1, 224, 224, 3)
in_min, in_max = (0, 255)
data = np.random.randint(in_min, high=in_max, size=data_shape, dtype="uint8")
mod, params = convert_to_relay(tflite_model_buf)
inputs = {"input": data}
output_list = generate_ref_data(mod, inputs, params)
compile_and_run(
AOTTestModel(module=mod, inputs=inputs, outputs=output_list, params=params),
test_runner,
interface_api,
use_unpacked_api,
)
def test_quant_mobilenet_tfl():
"""Since in AOT we pass directly the output buffer from the user, in quantized networks sharing the output buffers is not possible.
This is because the output data type is int8 and the intermediate buffer are int32 or int16. We use mobilenet quantized to stress this
situation and verify that the output buffer sharing is disabled in AOT."""
pytest.importorskip("tflite")
import tvm.relay.testing.tf as tf_testing
use_unpacked_api = True
interface_api = "c"
test_runner = AOT_DEFAULT_RUNNER
tflite_model_file = tf_testing.get_workload_official(
"https://storage.googleapis.com/download.tensorflow.org/"
"models/mobilenet_v1_2018_08_02/mobilenet_v1_1.0_224_quant.tgz",
"mobilenet_v1_1.0_224_quant.tflite",
)
mod, inputs, params = create_relay_module_and_inputs_from_tflite_file(
tflite_model_file)
output_list = generate_ref_data(mod, inputs, params)
compile_and_run(
AOTTestModel(module=mod,
inputs=inputs,
outputs=output_list,
params=params),
test_runner,
interface_api,
use_unpacked_api,
)
def test_forward_mobilenet_v1(accel_type):
"""Test the Mobilenet V1 TF Lite model."""
np.random.seed(23)
tflite_model_file = tf_testing.get_workload_official(
"https://storage.googleapis.com/download.tensorflow.org/"
"models/mobilenet_v1_2018_08_02/mobilenet_v1_1.0_224_quant.tgz",
"mobilenet_v1_1.0_224_quant.tflite",
)
with open(tflite_model_file, "rb") as f:
tflite_model_buf = f.read()
input_tensor = "input"
input_dtype = "uint8"
input_shape = (1, 224, 224, 3)
in_min, in_max = util.get_range_for_dtype_str(input_dtype)
input_data = np.random.randint(in_min, high=in_max, size=input_shape, dtype=input_dtype)
relay_mod, params = convert_to_relay(tflite_model_buf)
input_data = {input_tensor: input_data}
output_data = generate_ref_data(relay_mod, input_data)
mod = partition_for_ethosu(relay_mod, params)
compiled_models = infra.build_source(
mod, input_data, output_data, accel_type, output_tolerance=10
)
infra.verify_source(compiled_models, accel_type)
def get_tflite_model(model_url):
"""Get a TFLite model from URL."""
tflite_model_file = tf_testing.get_workload_official(
model_url[0], model_url[1])
with open(tflite_model_file, "rb") as f:
tflite_model_buf = f.read()
return tflite_model_buf
def get_model():
model_path = tf_testing.get_workload_official(
"https://storage.googleapis.com/download.tensorflow.org/models/tflite/model_zoo/upload_20180427/squeezenet_2018_04_27.tgz",
"squeezenet.tflite",
)
inputs = {"Placeholder": ((1, 224, 224, 3), "float32")}
mod, params = _get_tflite_model(model_path, inputs_dict=inputs)
return mod, params, inputs
def test_tflite_model_u3_usecase_two_external_pools(model_url, usmp_algo):
"""This checks for inference using two external pools placed in the application"""
pytest.importorskip("tflite")
import tvm.relay.testing.tf as tf_testing
use_unpacked_api = True
interface_api = "c"
target = tvm.target.Target("c")
workspace_memory_pools = WorkspaceMemoryPools([
PoolInfo("my_memory_pool_1", {target: PoolInfo.READ_WRITE_ACCESS},
size_hint_bytes=2500000),
PoolInfo("my_memory_pool_2", {target: PoolInfo.READ_WRITE_ACCESS}),
])
test_runner = AOTTestRunner(
pass_config={
"tir.usmp.enable": True,
"tir.usmp.algorithm": usmp_algo
},
prologue=f"""
__attribute__((section(".data.tvm"), aligned(16)))
static uint8_t my_memory_pool_1[{_get_workspace_size_define_macro("my_memory_pool_1")}];
__attribute__((section(".data.tvm"), aligned(16)))
static uint8_t my_memory_pool_2[{_get_workspace_size_define_macro("my_memory_pool_2")}];
""",
)
tflite_model_file = tf_testing.get_workload_official(
model_url[0],
model_url[1],
)
mod, inputs, params = _get_relay_module_and_inputs_from_tflite_file(
tflite_model_file)
output_list = generate_ref_data(mod, inputs, params)
compiled_test_mods = compile_models(
AOTTestModel(module=mod,
inputs=inputs,
outputs=output_list,
params=params),
interface_api=interface_api,
use_unpacked_api=use_unpacked_api,
pass_config=test_runner.pass_config,
workspace_memory_pools=workspace_memory_pools,
target=target,
)
for compiled_model in compiled_test_mods:
check_for_no_tvm_backendallocworkspace_calls(
compiled_model.executor_factory.lib)
run_and_check(
models=compiled_test_mods,
runner=test_runner,
interface_api=interface_api,
)
def get_model():
model_path = tf_testing.get_workload_official(
"https://storage.googleapis.com/download.tensorflow.org/"
"models/mobilenet_v1_2018_08_02/mobilenet_v1_1.0_224_quant.tgz",
"mobilenet_v1_1.0_224_quant.tflite",
)
inputs = {"input": ((1, 224, 224, 3), "uint8")}
mod, params = _get_tflite_model(model_path, inputs_dict=inputs)
return mod, params, inputs
def test_tflite_model_u1_usecase(model_url, usmp_algo, workspace_size,
constant_size):
"""
This checks for ML models and the memory used by them
when using USMP with different algorithms
"""
pytest.importorskip("tflite")
import tvm.relay.testing.tf as tf_testing # pylint: disable=import-outside-toplevel
use_unpacked_api = True
interface_api = "c"
test_runner = AOTTestRunner(pass_config={
"tir.usmp.enable": True,
"tir.usmp.algorithm": usmp_algo
})
tflite_model_file = tf_testing.get_workload_official(
model_url[0],
model_url[1],
)
mod, inputs, params = create_relay_module_and_inputs_from_tflite_file(
tflite_model_file)
output_list = generate_ref_data(mod, inputs, params)
compiled_test_mods = compile_models(
AOTTestModel(module=mod,
inputs=inputs,
outputs=output_list,
params=params),
interface_api=interface_api,
use_unpacked_api=use_unpacked_api,
pass_config=test_runner.pass_config,
)
for compiled_model in compiled_test_mods:
_check_for_no_tvm_backendallocworkspace_calls(
compiled_model.executor_factory.lib)
# Checking the workspace size reported in model library format
mlf_memory_map = mlf._build_function_memory_map(
compiled_test_mods[0].executor_factory.function_metadata)
assert mlf_memory_map["main"][0]["workspace_size_bytes"] == workspace_size
assert mlf_memory_map["main"][0]["constants_size_bytes"] == constant_size
# That should match to workspace size that will be codegen'd to the entry point.
allocated_pool_info_size = sum([
_.allocated_size for _ in list(
dict(compiled_test_mods[0].executor_factory.
executor_codegen_metadata.pool_inputs).values())
])
assert allocated_pool_info_size == workspace_size + constant_size
run_and_check(
models=compiled_test_mods,
runner=test_runner,
interface_api=interface_api,
)
def test_tflite_model(model_url, usmp_algo, workspace_size):
"""This checks for ML models and the memory used by them when using USMP with different algorithms"""
pytest.importorskip("tflite")
import tvm.relay.testing.tf as tf_testing
use_unpacked_api = True
interface_api = "c"
test_runner = AOTTestRunner(pass_config={
"tir.usmp.enable": True,
"tir.usmp.algorithm": usmp_algo
})
tflite_model_file = tf_testing.get_workload_official(
model_url[0],
model_url[1],
)
with open(tflite_model_file, "rb") as f:
tflite_model_buf = f.read()
data_shape = (1, 224, 224, 3)
in_min, in_max = (0, 255)
data = np.random.randint(in_min,
high=in_max,
size=data_shape,
dtype="uint8")
mod, params = convert_to_relay(tflite_model_buf, data, "input")
inputs = {"input": data}
output_list = generate_ref_data(mod, inputs, params)
compiled_test_mods = compile_models(
AOTTestModel(module=mod,
inputs=inputs,
outputs=output_list,
params=params),
interface_api=interface_api,
use_unpacked_api=use_unpacked_api,
pass_config=test_runner.pass_config,
)
for compiled_model in compiled_test_mods:
check_for_no_tvm_backendallocworkspace_calls(
compiled_model.executor_factory.lib)
# Checking the workspace size
assert (sum(
compiled_model.executor_factory.function_metadata["__tvm_main__"].
workspace_sizes.values()) == workspace_size)
run_and_check(
models=compiled_test_mods,
runner=test_runner,
interface_api=interface_api,
)
def test_tflite_model_u3_usecase_single_external_pool(model_url, usmp_algo):
"""This checks for inference with USMP using external pool placed in the application"""
pytest.importorskip("tflite")
import tvm.relay.testing.tf as tf_testing # pylint: disable=import-outside-toplevel
use_unpacked_api = True
interface_api = "c"
pool_name = "my_memory_pool"
target = tvm.target.Target("c")
workspace_memory_pools = WorkspaceMemoryPools(
[WorkspacePoolInfo(pool_name, [target])])
test_runner = AOTTestRunner(
pass_config={
"tir.usmp.enable": True,
"tir.usmp.algorithm": usmp_algo
},
prologue=f"""
__attribute__((section(".data.tvm"), aligned(16)))
static uint8_t {pool_name}[{_get_workspace_size_define_macro(pool_name)}];
""",
)
tflite_model_file = tf_testing.get_workload_official(
model_url[0],
model_url[1],
)
mod, inputs, params = create_relay_module_and_inputs_from_tflite_file(
tflite_model_file)
output_list = generate_ref_data(mod, inputs, params)
compiled_test_mods = compile_models(
AOTTestModel(module=mod,
inputs=inputs,
outputs=output_list,
params=params),
interface_api=interface_api,
use_unpacked_api=use_unpacked_api,
pass_config=test_runner.pass_config,
workspace_memory_pools=workspace_memory_pools,
target=target,
)
for compiled_model in compiled_test_mods:
_check_for_no_tvm_backendallocworkspace_calls(
compiled_model.executor_factory.lib)
run_and_check(
models=compiled_test_mods,
runner=test_runner,
interface_api=interface_api,
)
def get_model():
if model_url[-3:] in ("tgz", "zip"):
model_path = tf_testing.get_workload_official(
model_url,
model_sub_path,
)
else:
model_path = download.download_testdata(
model_url,
model_sub_path,
)
return _get_tflite_model(model_path, input_dict, "uint8")
def test_forward_ssd_mobilenet_v1():
"""Test the SSD Mobilenet V1 TF Lite model."""
# SSD MobilenetV1
tflite_model_file = tf_testing.get_workload_official(
"https://raw.githubusercontent.com/dmlc/web-data/master/tensorflow/models/object_detection/ssd_mobilenet_v1_coco_2018_01_28_nopp.tgz",
"ssd_mobilenet_v1_coco_2018_01_28_nopp.tflite")
with open(tflite_model_file, "rb") as f:
tflite_model_buf = f.read()
data = np.random.uniform(size=(1, 300, 300, 3)).astype('float32')
tflite_output = run_tflite_graph(tflite_model_buf, data)
tvm_output = run_tvm_graph(tflite_model_buf, data, 'normalized_input_image_tensor')
tvm.testing.assert_allclose(np.squeeze(tvm_output[0]), np.squeeze(tflite_output[0]),
rtol=1e-5, atol=1e-5)
def test_forward_mobilenet_v2():
"""Test the Mobilenet V2 TF Lite model."""
# MobilenetV2
tflite_model_file = tf_testing.get_workload_official(
"http://download.tensorflow.org/models/tflite_11_05_08/mobilenet_v2_1.0_224.tgz",
"mobilenet_v2_1.0_224.tflite")
with open(tflite_model_file, "rb") as f:
tflite_model_buf = f.read()
data = np.random.uniform(size=(1, 224, 224, 3)).astype('float32')
tflite_output = run_tflite_graph(tflite_model_buf, data)
tvm_output = run_tvm_graph(tflite_model_buf, data, 'input')
tvm.testing.assert_allclose(np.squeeze(tvm_output[0]), np.squeeze(tflite_output[0]),
rtol=1e-5, atol=1e-5)
def test_forward_inception_v4_net():
"""Test the Inception V4 TF Lite model."""
# InceptionV4
tflite_model_file = tf_testing.get_workload_official(
"https://storage.googleapis.com/download.tensorflow.org/models/tflite/model_zoo/upload_20180427/inception_v4_2018_04_27.tgz",
"inception_v4.tflite")
with open(tflite_model_file, "rb") as f:
tflite_model_buf = f.read()
data = np.random.uniform(size=(1, 299, 299, 3)).astype('float32')
tflite_output = run_tflite_graph(tflite_model_buf, data)
tvm_output = run_tvm_graph(tflite_model_buf, data, 'input')
tvm.testing.assert_allclose(np.squeeze(tvm_output[0]), np.squeeze(tflite_output[0]),
rtol=1e-5, atol=1e-5)
def test_forward_inception_v4_net():
"""Test the Inception V4 TF Lite model."""
# InceptionV4
tflite_model_file = tf_testing.get_workload_official(
"https://storage.googleapis.com/download.tensorflow.org/models/tflite/model_zoo/upload_20180427/inception_v4_2018_04_27.tgz",
"inception_v4.tflite")
with open(tflite_model_file, "rb") as f:
tflite_model_buf = f.read()
data = np.random.uniform(size=(1, 299, 299, 3)).astype('float32')
tvm_data = np.transpose(data, axes=(0, 3, 1, 2))
tflite_output = run_tflite_graph(tflite_model_buf, data)
tvm_output = run_tvm_graph(tflite_model_buf, tvm_data, 'input')
tvm.testing.assert_allclose(np.squeeze(tvm_output[0]), np.squeeze(tflite_output[0]),
rtol=1e-5, atol=1e-5)
def test_forward_mobilenet_v1():
'''test mobilenet v1 tflite model'''
# MobilenetV1
tflite_model_file = tf_testing.get_workload_official(
"http://download.tensorflow.org/models/mobilenet_v1_2018_08_02/mobilenet_v1_1.0_224.tgz",
"mobilenet_v1_1.0_224.tflite")
with open(tflite_model_file, "rb") as f:
tflite_model_buf = f.read()
data = np.random.uniform(size=(1, 224, 224, 3)).astype('float32')
tvm_data = np.transpose(data, axes=(0, 3, 1, 2))
tflite_output = run_tflite_graph(tflite_model_buf, data)
tvm_output = run_tvm_graph(tflite_model_buf, tvm_data, 'input')
tvm.testing.assert_allclose(np.squeeze(tvm_output[0]), np.squeeze(tflite_output[0]),
rtol=1e-5, atol=1e-5)
def test_forward_mobilenet_v2():
"""Test the Mobilenet V2 TF Lite model."""
# MobilenetV2
tflite_model_file = tf_testing.get_workload_official(
"http://download.tensorflow.org/models/tflite_11_05_08/mobilenet_v2_1.0_224.tgz",
"mobilenet_v2_1.0_224.tflite")
with open(tflite_model_file, "rb") as f:
tflite_model_buf = f.read()
data = np.random.uniform(size=(1, 224, 224, 3)).astype('float32')
tvm_data = np.transpose(data, axes=(0, 3, 1, 2))
tflite_output = run_tflite_graph(tflite_model_buf, data)
tvm_output = run_tvm_graph(tflite_model_buf, tvm_data, 'input')
tvm.testing.assert_allclose(np.squeeze(tvm_output[0]), np.squeeze(tflite_output[0]),
rtol=1e-5, atol=1e-5)
def test_tflite_model_u1_usecase(model_url, usmp_algo, workspace_size):
"""This checks for ML models and the memory used by them when using USMP with different algorithms"""
pytest.importorskip("tflite")
import tvm.relay.testing.tf as tf_testing
use_unpacked_api = True
interface_api = "c"
test_runner = AOTTestRunner(pass_config={
"tir.usmp.enable": True,
"tir.usmp.algorithm": usmp_algo
})
tflite_model_file = tf_testing.get_workload_official(
model_url[0],
model_url[1],
)
mod, inputs, params = _get_relay_module_and_inputs_from_tflite_file(
tflite_model_file)
output_list = generate_ref_data(mod, inputs, params)
compiled_test_mods = compile_models(
AOTTestModel(module=mod,
inputs=inputs,
outputs=output_list,
params=params),
interface_api=interface_api,
use_unpacked_api=use_unpacked_api,
pass_config=test_runner.pass_config,
)
for compiled_model in compiled_test_mods:
check_for_no_tvm_backendallocworkspace_calls(
compiled_model.executor_factory.lib)
# Checking the workspace size
assert (sum(
compiled_model.executor_factory.function_metadata["__tvm_main__"].
workspace_sizes.values()) == workspace_size)
run_and_check(
models=compiled_test_mods,
runner=test_runner,
interface_api=interface_api,
)
def test_forward_inception_v3_net():
'''test inception v3 tflite model'''
# InceptionV3
temp = util.tempdir()
tflite_model_file = tf_testing.get_workload_official(
"https://storage.googleapis.com/download.tensorflow.org/models/tflite/model_zoo/upload_20180427/inception_v3_2018_04_27.tgz",
"inception_v3.tflite", temp)
with open(tflite_model_file, "rb") as f:
tflite_model_buf = f.read()
data = np.random.uniform(size=(1, 299, 299, 3)).astype('float32')
tvm_data = np.transpose(data, axes=(0, 3, 1, 2))
tflite_output = run_tflite_graph(tflite_model_buf, data)
tvm_output = run_tvm_graph(tflite_model_buf, tvm_data, 'input')
tvm.testing.assert_allclose(np.squeeze(tvm_output[0]),
np.squeeze(tflite_output[0]),
rtol=1e-5,
atol=1e-5)
temp.remove()
def test_networks_without_usmp(accel_type, model_url, workspace_size,
tolerance):
np.random.seed(23)
tflite_model_file = tf_testing.get_workload_official(
model_url[0], model_url[1])
mod, input_data, params = create_relay_module_and_inputs_from_tflite_file(
tflite_model_file)
output_data = generate_ref_data(mod, input_data, params)
mod = partition_for_ethosu(mod, params)
compiled_models = infra.build_source(mod,
input_data,
output_data,
accel_type,
output_tolerance=tolerance,
enable_usmp=False)
mlf_memory_map = mlf._build_function_memory_map(
compiled_models[0].executor_factory.function_metadata)
assert mlf_memory_map["main"][0]["workspace_size_bytes"] == workspace_size
infra.verify_source(compiled_models, accel_type, enable_usmp=False)
def test_networks_with_usmp(accel_type, model_url, workspace_size, tolerance):
np.random.seed(23)
tflite_model_file = tf_testing.get_workload_official(
model_url[0], model_url[1])
mod, input_data, params = create_relay_module_and_inputs_from_tflite_file(
tflite_model_file)
output_data = generate_ref_data(mod, input_data, params)
mod = partition_for_ethosu(mod, params)
compiled_models = infra.build_source(mod,
input_data,
output_data,
accel_type,
output_tolerance=tolerance,
enable_usmp=True)
allocated_pool_info = list(
dict(compiled_models[0].executor_factory.executor_codegen_metadata.
pool_inputs).values())[0]
assert allocated_pool_info.allocated_size == workspace_size
infra.verify_source(compiled_models, accel_type, enable_usmp=True)
def test_forward_qnn_mobilenet_v2_net():
"""Test the Quantized TFLite Mobilenet V2 model."""
# MobilenetV2
tflite_model_file = tf_testing.get_workload_official(
"https://storage.googleapis.com/download.tensorflow.org/models/tflite_11_05_08/mobilenet_v2_1.0_224_quant.tgz",
"mobilenet_v2_1.0_224_quant.tflite")
with open(tflite_model_file, "rb") as f:
tflite_model_buf = f.read()
# Checking the labels because the requantize implementation is different between TFLite and
# Relay. This cause final output numbers to mismatch. So, testing accuracy via labels.
np.random.seed(0)
data = np.random.random_integers(low=0, high=128, size=(1, 224, 224, 3)).astype('uint8')
tflite_output = run_tflite_graph(tflite_model_buf, data)
tflite_predictions = np.squeeze(tflite_output)
tflite_sorted_labels = tflite_predictions.argsort()[-3:][::-1]
tvm_output = run_tvm_graph(tflite_model_buf, data, 'input')
tvm_predictions = np.squeeze(tvm_output)
tvm_sorted_labels = tvm_predictions.argsort()[-3:][::-1]
tvm.testing.assert_allclose(tvm_sorted_labels, tflite_sorted_labels)
def test_tflite_model_u3_usecase_var_cons_ext_pools(model_url, usmp_algo):
"""This checks for inference using one external workspace and one external constant
pools placed in the application"""
pytest.importorskip("tflite")
import tvm.relay.testing.tf as tf_testing # pylint: disable=import-outside-toplevel
use_unpacked_api = True
interface_api = "c"
target = tvm.target.Target("c")
workspace_mem_pools = WorkspaceMemoryPools([
WorkspacePoolInfo("my_memory_pool_1", [target],
PoolInfoProperties(size_hint_bytes=8500000)),
])
constant_mem_pools = ConstantMemoryPools([
ConstantPoolInfo("my_const_pool_1", [target], []),
])
test_runner = AOTTestRunner(
pass_config={
"tir.usmp.enable": True,
"tir.usmp.algorithm": usmp_algo
},
prologue=f"""
__attribute__((section(".bss.noinit"), aligned(TVM_RUNTIME_ALLOC_ALIGNMENT_BYTES)))
static uint8_t my_memory_pool_1[{_get_workspace_size_define_macro("my_memory_pool_1")}];
__attribute__((section(".rodata.tvm"), aligned(TVM_RUNTIME_CONST_ALLOC_ALIGNMENT_BYTES)))
static uint8_t my_const_pool_1[{_get_constant_size_define_macro("my_const_pool_1")}] = {{ {_get_constant_data_define_macro("my_const_pool_1")} }};
""",
)
tflite_model_file = tf_testing.get_workload_official(
model_url[0],
model_url[1],
)
mod, inputs, params = create_relay_module_and_inputs_from_tflite_file(
tflite_model_file)
output_list = generate_ref_data(mod, inputs, params)
compiled_test_mods = compile_models(
AOTTestModel(module=mod,
inputs=inputs,
outputs=output_list,
params=params),
interface_api=interface_api,
use_unpacked_api=use_unpacked_api,
pass_config=test_runner.pass_config,
workspace_memory_pools=workspace_mem_pools,
constant_memory_pools=constant_mem_pools,
target=target,
)
for compiled_model in compiled_test_mods:
_check_for_no_tvm_backendallocworkspace_calls(
compiled_model.executor_factory.lib)
run_and_check(
models=compiled_test_mods,
runner=test_runner,
interface_api=interface_api,
)
def test_tflite_model_u4_usecase_two_external_pools(model_url, usmp_algo):
"""This checks for inference with USMP using external pool placed in the application"""
pytest.importorskip("tflite")
import tvm.relay.testing.tf as tf_testing # pylint: disable=import-outside-toplevel
use_unpacked_api = True
interface_api = "c"
target = tvm.target.Target("c")
workspace_memory_pools = WorkspaceMemoryPools([
PoolInfo("my_memory_pool_1", {target: PoolInfo.READ_WRITE_ACCESS},
size_hint_bytes=2500000),
PoolInfo("my_memory_pool_2", {target: PoolInfo.READ_WRITE_ACCESS}),
])
tflite_model_file = tf_testing.get_workload_official(
model_url[0],
model_url[1],
)
mod, inputs, params = create_relay_module_and_inputs_from_tflite_file(
tflite_model_file)
output_list = generate_ref_data(mod, inputs, params)
input_name, input_data = list(inputs.items())[0]
input_size_bytes = input_data.size * input_data.itemsize
test_runner = AOTTestRunner(
pass_config={
"tir.usmp.enable": True,
"tir.usmp.algorithm": usmp_algo,
"tir.usmp.use_workspace_io": True,
},
prologue=f"""
#include <string.h>
__attribute__((section(".data.tvm"), aligned(16)))
static uint8_t my_memory_pool_1[{_get_workspace_size_define_macro("my_memory_pool_1")}];
__attribute__((section(".data.tvm"), aligned(16)))
static uint8_t my_memory_pool_2[{_get_workspace_size_define_macro("my_memory_pool_2")}];
struct {_add_module_prefix("workspace_pools")} {_add_module_prefix("workspace_pools")} = {{
.my_memory_pool_1 = my_memory_pool_1,
.my_memory_pool_2 = my_memory_pool_2,
}};
struct {_add_module_prefix("inputs")} {_add_module_prefix("inputs")} = {_add_module_prefix("map_inputs")}(&{_add_module_prefix("workspace_pools")});
memcpy({_add_module_prefix("inputs")}.{input_name}, tvmgen_default_input_data_input, {input_size_bytes});
struct {_add_module_prefix("outputs")} {_add_module_prefix("outputs")} = {_add_module_prefix("map_outputs")}(&{_add_module_prefix("workspace_pools")});
""",
)
compiled_test_mods = compile_models(
AOTTestModel(module=mod,
inputs=inputs,
outputs=output_list,
params=params),
interface_api=interface_api,
use_unpacked_api=use_unpacked_api,
pass_config=test_runner.pass_config,
workspace_memory_pools=workspace_memory_pools,
target=target,
)
for compiled_model in compiled_test_mods:
_check_for_no_tvm_backendallocworkspace_calls(
compiled_model.executor_factory.lib)
run_and_check(
models=compiled_test_mods,
runner=test_runner,
interface_api=interface_api,
use_workspace_io=True,
)
