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)