def _test():
target = Target("c")
global_workspace_pool = usmp_utils.PoolInfo(
pool_name="global_workspace",
target_access={target: usmp_utils.PoolInfo.READ_WRITE_ACCESS},
)
bi_a = usmp_utils.BufferInfo(name_hint="bi_a",
size_bytes=10,
pool_candidates=[global_workspace_pool])
bi_b = usmp_utils.BufferInfo(name_hint="bi_b",
size_bytes=10,
pool_candidates=[global_workspace_pool])
bi_c = usmp_utils.BufferInfo(name_hint="bi_c",
size_bytes=10,
pool_candidates=[global_workspace_pool])
bi_a.set_conflicts([bi_b, bi_c])
bi_b.set_conflicts([bi_c, bi_a])
bi_c.set_conflicts([bi_a, bi_b])
buffer_info_arr = [bi_a, bi_b, bi_c]
fusmp_algo = tvm.get_global_func(f"tir.usmp.algo.{algorithm}")
buffer_pool_allocations = fusmp_algo(buffer_info_arr, 0)
assert buffer_pool_allocations[bi_a].byte_offset == 20
assert buffer_pool_allocations[bi_b].byte_offset == 10
assert buffer_pool_allocations[bi_c].byte_offset == 0
def test_no_pool_error():
target = Target("c")
tiny_workspace_pool = usmp_utils.PoolInfo(
pool_name="tiny_workspace",
target_access={target: usmp_utils.PoolInfo.READ_WRITE_ACCESS},
size_hint_bytes=10,
)
bi_a = usmp_utils.BufferInfo(name_hint="bi_a",
size_bytes=10,
pool_candidates=[tiny_workspace_pool])
bi_b = usmp_utils.BufferInfo(name_hint="bi_b",
size_bytes=10,
pool_candidates=[tiny_workspace_pool])
bi_c = usmp_utils.BufferInfo(name_hint="bi_c",
size_bytes=10,
pool_candidates=[tiny_workspace_pool])
bi_a.set_conflicts([bi_b])
bi_b.set_conflicts([bi_c])
bi_c.set_conflicts([bi_a])
buffer_info_arr = [bi_a, bi_b, bi_c]
fusmp_algo = tvm.get_global_func(f"tir.usmp.algo.greedy_by_size")
with pytest.raises(
tvm.TVMError,
match=
"TVM USMP Error: the space available in the provided pools exceeded"
):
buffer_pool_allocations = fusmp_algo(buffer_info_arr, 0)
def test_create_pool_allocation():
pool_info = usmp_utils.PoolInfo(
pool_name="foo_workspace",
target_access={Target("c"): usmp_utils.PoolInfo.READ_WRITE_ACCESS},
)
pool_allocation = usmp_utils.PoolAllocation(pool_info=pool_info,
byte_offset=64)
assert pool_allocation.pool_info == pool_info
assert pool_allocation.byte_offset == 64
def test_linear(algorithm, workspace_size):
"""
The test case here represent BufferInfo objects
that could get generated for a linear sequence
such as :
(Op A)
|
bi_a
|
(Op B)
|
bi_b
|
.
.
.
(Op F)
|
bi_f
"""
target = Target("c")
global_workspace_pool = usmp_utils.PoolInfo(
pool_name="global_workspace",
target_access={target: usmp_utils.PoolInfo.READ_WRITE_ACCESS},
)
bi_a = usmp_utils.BufferInfo(name_hint="bi_a",
size_bytes=10,
pool_candidates=[global_workspace_pool])
bi_b = usmp_utils.BufferInfo(name_hint="bi_b",
size_bytes=20,
pool_candidates=[global_workspace_pool])
bi_c = usmp_utils.BufferInfo(name_hint="bi_c",
size_bytes=100,
pool_candidates=[global_workspace_pool])
bi_d = usmp_utils.BufferInfo(name_hint="bi_d",
size_bytes=40,
pool_candidates=[global_workspace_pool])
bi_e = usmp_utils.BufferInfo(name_hint="bi_e",
size_bytes=50,
pool_candidates=[global_workspace_pool])
bi_f = usmp_utils.BufferInfo(name_hint="bi_f",
size_bytes=50,
pool_candidates=[global_workspace_pool])
# Creating conflicts for a linear graph
bi_a.set_conflicts([bi_b])
bi_b.set_conflicts([bi_a, bi_c])
bi_c.set_conflicts([bi_b, bi_d])
bi_d.set_conflicts([bi_c, bi_e])
bi_e.set_conflicts([bi_d, bi_f])
bi_f.set_conflicts([bi_e])
buffer_info_arr = [bi_a, bi_b, bi_c, bi_d, bi_e, bi_f]
fusmp_algo = tvm.get_global_func(f"tir.usmp.algo.{algorithm}")
buffer_pool_allocations = fusmp_algo(buffer_info_arr, 0)
_check_max_workspace_size(buffer_pool_allocations, global_workspace_pool,
workspace_size)
def test_mobilenet_subgraph(algorithm, fast_memory_size, slow_memory_size):
target = Target("c")
fast_memory_pool = usmp_utils.PoolInfo(
pool_name="fast_memory",
target_access={target: usmp_utils.PoolInfo.READ_WRITE_ACCESS},
size_hint_bytes=200704,
)
slow_memory_pool = usmp_utils.PoolInfo(
pool_name="slow_memory",
target_access={target: usmp_utils.PoolInfo.READ_WRITE_ACCESS})
tir_mod = MobilenetStructure
tir_mod = _assign_targets_to_primfuncs_irmodule(tir_mod, target)
tir_mod = _assign_poolinfos_to_allocates_in_irmodule(
tir_mod, [fast_memory_pool, slow_memory_pool])
main_func = tir_mod["run_model"]
buffer_info_analysis = tvm.tir.usmp.analysis.extract_buffer_info(
main_func, tir_mod)
assert buffer_info_analysis.memory_pressure == 1117718
fcreate_array_bi = tvm.get_global_func("tir.usmp.CreateArrayBufferInfo")
buffer_info_arr = fcreate_array_bi(buffer_info_analysis.buffer_info_stmts)
fusmp_algo = tvm.get_global_func(f"tir.usmp.algo.{algorithm}")
buffer_pool_allocations = fusmp_algo(buffer_info_arr,
buffer_info_analysis.memory_pressure)
buffer_info_map_names = dict()
for buf_info in buffer_info_arr:
buffer_info_map_names[buf_info.name_hint] = buf_info
# check conflicts
_verify_conflicts("PaddedInput_7", ["sid_9", "sid_8", "Conv2dOutput_7"],
buffer_info_map_names)
_verify_conflicts("tensor_2", ["sid_8"], buffer_info_map_names)
_verify_conflicts("sid_9", ["PaddedInput_7"], buffer_info_map_names)
_verify_conflicts("sid_8", ["PaddedInput_7", "Conv2dOutput_7", "tensor_2"],
buffer_info_map_names)
_verify_conflicts("Conv2dOutput_7", ["sid_8", "PaddedInput_7"],
buffer_info_map_names)
_check_max_workspace_size(buffer_pool_allocations, slow_memory_pool,
slow_memory_size)
_check_max_workspace_size(buffer_pool_allocations, fast_memory_pool,
fast_memory_size)
def test_create_pool_info():
target = Target("c")
pool_info = usmp_utils.PoolInfo(
pool_name="foo_workspace",
target_access={target: usmp_utils.PoolInfo.READ_WRITE_ACCESS},
)
assert pool_info.pool_name == "foo_workspace"
assert dict(pool_info.target_access) == {
target: usmp_utils.PoolInfo.READ_WRITE_ACCESS
}
# default pool size constraint
assert pool_info.size_hint_bytes == -1
pool_info = usmp_utils.PoolInfo(
pool_name="bar_workspace",
target_access={target: usmp_utils.PoolInfo.READ_ONLY_ACCESS},
size_hint_bytes=1425,
)
assert pool_info.pool_name == "bar_workspace"
assert dict(pool_info.target_access) == {
target: usmp_utils.PoolInfo.READ_ONLY_ACCESS
}
assert pool_info.size_hint_bytes == 1425
def test_create_array_buffer_info():
target = Target("c")
global_ws_pool = usmp_utils.PoolInfo(
pool_name="global_workspace",
target_access={target: usmp_utils.PoolInfo.READ_WRITE_ACCESS},
)
fcreate_array_bi = tvm.get_global_func("tir.usmp.CreateArrayBufferInfo")
tir_mod = LinearStructure
tir_mod = _assign_targets_to_primfuncs_irmodule(tir_mod, target)
tir_mod = _assign_poolinfos_to_allocates_in_irmodule(
tir_mod, [global_ws_pool])
main_func = tir_mod["tvmgen_default_run_model"]
buffer_info_map = tvm.tir.usmp.analysis.extract_buffer_info(
main_func, tir_mod)
buffer_info_array = fcreate_array_bi(buffer_info_map)
for buffer_info in buffer_info_array:
assert buffer_info in buffer_info_map.keys()
def test_create_buffer_info():
global_ws_pool = usmp_utils.PoolInfo(
pool_name="global_workspace",
target_access={Target("c"): usmp_utils.PoolInfo.READ_WRITE_ACCESS},
)
buffer_info_obj = tvm.tir.usmp.BufferInfo(name_hint="buf1",
size_bytes=256,
pool_candidates=[global_ws_pool])
assert buffer_info_obj.name_hint == "buf1"
assert buffer_info_obj.size_bytes == 256
assert list(buffer_info_obj.pool_candidates) == [global_ws_pool]
# default workspace alignment
assert buffer_info_obj.alignment == 1
buffer_info_obj = tvm.tir.usmp.BufferInfo("buf2", 512, [global_ws_pool], 8)
assert buffer_info_obj.name_hint == "buf2"
assert buffer_info_obj.size_bytes == 512
assert list(buffer_info_obj.pool_candidates) == [global_ws_pool]
assert buffer_info_obj.alignment == 8
def test_resnet_subgraph(algorithm, workspace_size):
target = Target("c")
global_workspace_pool = usmp_utils.PoolInfo(
pool_name="global_workspace",
target_access={target: usmp_utils.PoolInfo.READ_WRITE_ACCESS},
)
tir_mod = ResnetStructure
tir_mod = _assign_targets_to_primfuncs_irmodule(tir_mod, target)
tir_mod = _assign_poolinfos_to_allocates_in_irmodule(
tir_mod, [global_workspace_pool])
main_func = tir_mod["tvmgen_default_run_model"]
buffer_info_analysis = tvm.tir.usmp.analysis.extract_buffer_info(
main_func, tir_mod)
assert buffer_info_analysis.memory_pressure == 7200256
fcreate_array_bi = tvm.get_global_func("tir.usmp.CreateArrayBufferInfo")
buffer_info_arr = fcreate_array_bi(buffer_info_analysis.buffer_info_stmts)
fusmp_algo = tvm.get_global_func(f"tir.usmp.algo.{algorithm}")
buffer_pool_allocations = fusmp_algo(buffer_info_arr,
buffer_info_analysis.memory_pressure)
buffer_info_map_names = dict()
for buf_info in buffer_info_arr:
buffer_info_map_names[buf_info.name_hint] = buf_info
# check conflicts
_verify_conflicts(
"sid_7",
[
"PaddedInput_1",
"sid_2",
"Conv2dOutput_1",
"PaddedInput_2",
],
buffer_info_map_names,
)
_verify_conflicts(
"Conv2dOutput_3",
[
"PaddedInput_3",
"sid_6",
],
buffer_info_map_names,
)
_verify_conflicts(
"sid_6",
[
"Conv2dOutput_2",
"PaddedInput_2",
"sid_2",
"PaddedInput_3",
"Conv2dOutput_3",
],
buffer_info_map_names,
)
_verify_conflicts(
"Conv2dOutput",
[
"sid_8",
"sid_2",
"PaddedInput",
],
buffer_info_map_names,
)
_verify_conflicts(
"PaddedInput_3",
[
"sid_6",
"sid_2",
"Conv2dOutput_3",
],
buffer_info_map_names,
)
_verify_conflicts(
"Conv2dOutput_2",
[
"PaddedInput_2",
"sid_2",
"sid_6",
],
buffer_info_map_names,
)
_verify_conflicts(
"PaddedInput_1",
[
"sid_8",
"sid_2",
"sid_7",
"Conv2dOutput_1",
],
buffer_info_map_names,
)
_verify_conflicts(
"Conv2dOutput_1",
[
"sid_7",
"PaddedInput_1",
"sid_2",
],
buffer_info_map_names,
)
_verify_conflicts(
"PaddedInput",
[
"sid_2",
"sid_8",
"Conv2dOutput",
],
buffer_info_map_names,
)
_verify_conflicts(
"sid_8",
[
"PaddedInput",
"sid_2",
"Conv2dOutput",
"PaddedInput_1",
],
buffer_info_map_names,
)
_verify_conflicts(
"sid_2",
[
"PaddedInput",
"sid_8",
"Conv2dOutput",
"PaddedInput_1",
"sid_7",
"Conv2dOutput_1",
"PaddedInput_2",
"Conv2dOutput_2",
"sid_6",
"PaddedInput_3",
],
buffer_info_map_names,
)
_verify_conflicts(
"PaddedInput_2",
[
"sid_7",
"sid_2",
"Conv2dOutput_2",
"sid_6",
],
buffer_info_map_names,
)
_check_max_workspace_size(buffer_pool_allocations, global_workspace_pool,
workspace_size)
