[microNPU] Integrate the cascader

[ekalda]

Integrate the cascader into the codegen and optionally enable it
with the enable_cascader flag. Includes placeholder MemoryRegions until
integration with the PoolInfos provided by a user.

Co-authored-by: Matthew Barrett matthew.barrett@arm.com

[ekalda]

cc: @manupa-arm@lhutton1@NicolaLancellotti

[ekalda]

This patch supersedes #10377, so that one can be closed.

[manupak]

Thanks @ekalda!

It is nice to see this working and green in CI.

I have left few comments around testing and user facing errors.

[manupak]

[Maybe for a subsequent PR] Will it be possible plumb these values using this :

tvm/src/relay/backend/build_module.cc

Lines 414 to 417 in fcdf463

	IRModule func_module = WithAttrs(IRModule::FromExpr(func),
	{{tvm::attr::kExecutor, executor_},
	{tvm::attr::kRuntime, runtime_},
	{tvm::attr::kWorkspaceMemoryPools, workspace_memory_pools_}});

An example test :

https://github.com/apache/tvm/blob/main/tests/python/relay/aot/test_crt_aot_usmp.py#L341-L383

[ekalda]

Done, with the caveat that we'll assume that there is one workspace pool in the system - is it ok to assume that for now or should we handle the case where there are several workspace pools where some of them are not accessible for the NPU?

[manupak]

For these tests, I think we need to ensure the memory usage is reduced.

We could use calculate_workspace_size TIR analysis utility for that. E.g. :

tvm/tests/python/unittest/test_tir_analysis_calculate_workspace.py

Line 106 in fcdf463

asserttvm.tir.analysis.calculate_workspace_bytes(primfunc, alignment) ==size

OR we could use the final memory calculation in a non-USMP flow as follows :

tvm/tests/python/relay/aot/test_crt_aot.py

Lines 937 to 956 in fcdf463

	deftest_workspace_calculation(workspace_byte_alignment, main_workspace_size):
	mod, params=tvm.relay.testing.synthetic.get_workload()
	target="c"
	runtime=Runtime("crt")
	executor=Executor(
	"aot",
	{
	"workspace-byte-alignment": workspace_byte_alignment,
	},
	)
	withtvm.transform.PassContext(
	opt_level=3,
	config={
	"tir.disable_vectorize": True,
	},
	):
	lib=tvm.relay.build(mod, target, executor=executor, runtime=runtime, params=params)
	mlf_memory_map=mlf._build_function_memory_map(lib.function_metadata)
	assertmlf_memory_map["main"][0]["workspace_size_bytes"] ==main_workspace_size

[ekalda]

I changed the tests to check for reduction in memory instead of bitwise accuracy with TFLite (I think it would be still good to have some FVP based tests when cascader is enabled, but it looks like that would need quite a bit of infra refactor, so I'll do it in a separate patch, if this is ok)

[lhutton1]

Looks good to me modulo @manupa-arm's comments :)

[NicolaLancellotti]

I think it is better is_cascader_enabled, otherwise, it seems you are enabling it with this function.

[ekalda]

Good point! I changed the function name to is_cascader_enabled, but kept the flag/variable as enable_cascader to align with the philosophy of enable_usmp

[NicolaLancellotti]

LGTM! Thank you @ekalda.

[manupak]

Broadly looks great!

just a nit and a question

[manupak]

out of curiosity, should we not check for exact values ?

[ekalda]

Yes good point, I changed it to check for the exact values

[manupak]

I think it would be better for this utility to be part of the cascader and not get exposed to the codegen here.

[ekalda]

Done

[manupak]

LGTM with a suggestion for a follow up.

[manupak]

For a followup : Please consider absorbing the call to extract_memory_info inside the cascader. (Sorry for not being clear before). Ideally, we'd want to remove the "MemoryRegion" construct and to get there in the current direction of travel, we should try to confine the usage of it inside the cascader. Therefore the interface of the cascader should be made to accept MemoryPool(s).

[ekalda]

I agree that this would be the right way to go about it. Let's do it in a follow up!

[manupak]

Thanks @ekalda@mbaret@lhutton1@NicolaLancellotti !

This is merged now!

[ekalda]

Integrate the cascader into the codegen and optionally enable it
with the enable_cascader flag. Includes placeholder MemoryRegions until
integration with the PoolInfos provided by a user.

Co-authored-by: Matthew Barrett matthew.barrett@arm.com

[ekalda]

cc: @manupa-arm@lhutton1@NicolaLancellotti

[ekalda]

This patch supersedes #10377, so that one can be closed.

[manupak]

Thanks @ekalda!

It is nice to see this working and green in CI.

I have left few comments around testing and user facing errors.

[manupak]

[Maybe for a subsequent PR] Will it be possible plumb these values using this :

tvm/src/relay/backend/build_module.cc

Lines 414 to 417 in fcdf463

	IRModule func_module = WithAttrs(IRModule::FromExpr(func),
	{{tvm::attr::kExecutor, executor_},
	{tvm::attr::kRuntime, runtime_},
	{tvm::attr::kWorkspaceMemoryPools, workspace_memory_pools_}});

An example test :

https://github.com/apache/tvm/blob/main/tests/python/relay/aot/test_crt_aot_usmp.py#L341-L383

[ekalda]

Done, with the caveat that we'll assume that there is one workspace pool in the system - is it ok to assume that for now or should we handle the case where there are several workspace pools where some of them are not accessible for the NPU?

[manupak]

For these tests, I think we need to ensure the memory usage is reduced.

We could use calculate_workspace_size TIR analysis utility for that. E.g. :

tvm/tests/python/unittest/test_tir_analysis_calculate_workspace.py

Line 106 in fcdf463

asserttvm.tir.analysis.calculate_workspace_bytes(primfunc, alignment) ==size

OR we could use the final memory calculation in a non-USMP flow as follows :

tvm/tests/python/relay/aot/test_crt_aot.py

Lines 937 to 956 in fcdf463

	deftest_workspace_calculation(workspace_byte_alignment, main_workspace_size):
	mod, params=tvm.relay.testing.synthetic.get_workload()
	target="c"
	runtime=Runtime("crt")
	executor=Executor(
	"aot",
	{
	"workspace-byte-alignment": workspace_byte_alignment,
	},
	)
	withtvm.transform.PassContext(
	opt_level=3,
	config={
	"tir.disable_vectorize": True,
	},
	):
	lib=tvm.relay.build(mod, target, executor=executor, runtime=runtime, params=params)
	mlf_memory_map=mlf._build_function_memory_map(lib.function_metadata)
	assertmlf_memory_map["main"][0]["workspace_size_bytes"] ==main_workspace_size

[ekalda]

I changed the tests to check for reduction in memory instead of bitwise accuracy with TFLite (I think it would be still good to have some FVP based tests when cascader is enabled, but it looks like that would need quite a bit of infra refactor, so I'll do it in a separate patch, if this is ok)

[lhutton1]

Looks good to me modulo @manupa-arm's comments :)

[NicolaLancellotti]

I think it is better is_cascader_enabled, otherwise, it seems you are enabling it with this function.

[ekalda]

Good point! I changed the function name to is_cascader_enabled, but kept the flag/variable as enable_cascader to align with the philosophy of enable_usmp

[NicolaLancellotti]

LGTM! Thank you @ekalda.

[manupak]

Broadly looks great!

just a nit and a question

[manupak]

out of curiosity, should we not check for exact values ?

[ekalda]

Yes good point, I changed it to check for the exact values

[manupak]

I think it would be better for this utility to be part of the cascader and not get exposed to the codegen here.

[ekalda]

Done

[manupak]

LGTM with a suggestion for a follow up.

[manupak]

For a followup : Please consider absorbing the call to extract_memory_info inside the cascader. (Sorry for not being clear before). Ideally, we'd want to remove the "MemoryRegion" construct and to get there in the current direction of travel, we should try to confine the usage of it inside the cascader. Therefore the interface of the cascader should be made to accept MemoryPool(s).

[ekalda]

I agree that this would be the right way to go about it. Let's do it in a follow up!

[manupak]

Thanks @ekalda@mbaret@lhutton1@NicolaLancellotti !

This is merged now!

[ekalda]

Integrate the cascader into the codegen and optionally enable it
with the enable_cascader flag. Includes placeholder MemoryRegions until
integration with the PoolInfos provided by a user.

Co-authored-by: Matthew Barrett matthew.barrett@arm.com

[ekalda]

cc: @manupa-arm@lhutton1@NicolaLancellotti

[ekalda]

This patch supersedes #10377, so that one can be closed.

[manupak]

Thanks @ekalda!

It is nice to see this working and green in CI.

I have left few comments around testing and user facing errors.

[manupak]

[Maybe for a subsequent PR] Will it be possible plumb these values using this :

tvm/src/relay/backend/build_module.cc

Lines 414 to 417 in fcdf463

	IRModule func_module = WithAttrs(IRModule::FromExpr(func),
	{{tvm::attr::kExecutor, executor_},
	{tvm::attr::kRuntime, runtime_},
	{tvm::attr::kWorkspaceMemoryPools, workspace_memory_pools_}});

An example test :

https://github.com/apache/tvm/blob/main/tests/python/relay/aot/test_crt_aot_usmp.py#L341-L383

[ekalda]

Done, with the caveat that we'll assume that there is one workspace pool in the system - is it ok to assume that for now or should we handle the case where there are several workspace pools where some of them are not accessible for the NPU?

[manupak]

For these tests, I think we need to ensure the memory usage is reduced.

We could use calculate_workspace_size TIR analysis utility for that. E.g. :

tvm/tests/python/unittest/test_tir_analysis_calculate_workspace.py

Line 106 in fcdf463

asserttvm.tir.analysis.calculate_workspace_bytes(primfunc, alignment) ==size

OR we could use the final memory calculation in a non-USMP flow as follows :

tvm/tests/python/relay/aot/test_crt_aot.py

Lines 937 to 956 in fcdf463

	deftest_workspace_calculation(workspace_byte_alignment, main_workspace_size):
	mod, params=tvm.relay.testing.synthetic.get_workload()
	target="c"
	runtime=Runtime("crt")
	executor=Executor(
	"aot",
	{
	"workspace-byte-alignment": workspace_byte_alignment,
	},
	)
	withtvm.transform.PassContext(
	opt_level=3,
	config={
	"tir.disable_vectorize": True,
	},
	):
	lib=tvm.relay.build(mod, target, executor=executor, runtime=runtime, params=params)
	mlf_memory_map=mlf._build_function_memory_map(lib.function_metadata)
	assertmlf_memory_map["main"][0]["workspace_size_bytes"] ==main_workspace_size

[ekalda]

I changed the tests to check for reduction in memory instead of bitwise accuracy with TFLite (I think it would be still good to have some FVP based tests when cascader is enabled, but it looks like that would need quite a bit of infra refactor, so I'll do it in a separate patch, if this is ok)

[lhutton1]

Looks good to me modulo @manupa-arm's comments :)

[NicolaLancellotti]

I think it is better is_cascader_enabled, otherwise, it seems you are enabling it with this function.

[ekalda]

Good point! I changed the function name to is_cascader_enabled, but kept the flag/variable as enable_cascader to align with the philosophy of enable_usmp

[NicolaLancellotti]

LGTM! Thank you @ekalda.

[manupak]

Broadly looks great!

just a nit and a question

[manupak]

out of curiosity, should we not check for exact values ?

[ekalda]

Yes good point, I changed it to check for the exact values

[manupak]

I think it would be better for this utility to be part of the cascader and not get exposed to the codegen here.

[ekalda]

Done

[manupak]

LGTM with a suggestion for a follow up.

[manupak]

For a followup : Please consider absorbing the call to extract_memory_info inside the cascader. (Sorry for not being clear before). Ideally, we'd want to remove the "MemoryRegion" construct and to get there in the current direction of travel, we should try to confine the usage of it inside the cascader. Therefore the interface of the cascader should be made to accept MemoryPool(s).

[ekalda]

I agree that this would be the right way to go about it. Let's do it in a follow up!

[manupak]

Thanks @ekalda@mbaret@lhutton1@NicolaLancellotti !

This is merged now!

[ekalda]

Integrate the cascader into the codegen and optionally enable it
with the enable_cascader flag. Includes placeholder MemoryRegions until
integration with the PoolInfos provided by a user.

Co-authored-by: Matthew Barrett matthew.barrett@arm.com

[ekalda]

cc: @manupa-arm@lhutton1@NicolaLancellotti

[ekalda]

This patch supersedes #10377, so that one can be closed.

[manupak]

Thanks @ekalda!

It is nice to see this working and green in CI.

I have left few comments around testing and user facing errors.

[manupak]

[Maybe for a subsequent PR] Will it be possible plumb these values using this :

tvm/src/relay/backend/build_module.cc

Lines 414 to 417 in fcdf463

	IRModule func_module = WithAttrs(IRModule::FromExpr(func),
	{{tvm::attr::kExecutor, executor_},
	{tvm::attr::kRuntime, runtime_},
	{tvm::attr::kWorkspaceMemoryPools, workspace_memory_pools_}});

An example test :

https://github.com/apache/tvm/blob/main/tests/python/relay/aot/test_crt_aot_usmp.py#L341-L383

[ekalda]

Done, with the caveat that we'll assume that there is one workspace pool in the system - is it ok to assume that for now or should we handle the case where there are several workspace pools where some of them are not accessible for the NPU?

[manupak]

For these tests, I think we need to ensure the memory usage is reduced.

We could use calculate_workspace_size TIR analysis utility for that. E.g. :

tvm/tests/python/unittest/test_tir_analysis_calculate_workspace.py

Line 106 in fcdf463

asserttvm.tir.analysis.calculate_workspace_bytes(primfunc, alignment) ==size

OR we could use the final memory calculation in a non-USMP flow as follows :

tvm/tests/python/relay/aot/test_crt_aot.py

Lines 937 to 956 in fcdf463

	deftest_workspace_calculation(workspace_byte_alignment, main_workspace_size):
	mod, params=tvm.relay.testing.synthetic.get_workload()
	target="c"
	runtime=Runtime("crt")
	executor=Executor(
	"aot",
	{
	"workspace-byte-alignment": workspace_byte_alignment,
	},
	)
	withtvm.transform.PassContext(
	opt_level=3,
	config={
	"tir.disable_vectorize": True,
	},
	):
	lib=tvm.relay.build(mod, target, executor=executor, runtime=runtime, params=params)
	mlf_memory_map=mlf._build_function_memory_map(lib.function_metadata)
	assertmlf_memory_map["main"][0]["workspace_size_bytes"] ==main_workspace_size

[ekalda]

I changed the tests to check for reduction in memory instead of bitwise accuracy with TFLite (I think it would be still good to have some FVP based tests when cascader is enabled, but it looks like that would need quite a bit of infra refactor, so I'll do it in a separate patch, if this is ok)

[lhutton1]

Looks good to me modulo @manupa-arm's comments :)

[NicolaLancellotti]

I think it is better is_cascader_enabled, otherwise, it seems you are enabling it with this function.

[ekalda]

Good point! I changed the function name to is_cascader_enabled, but kept the flag/variable as enable_cascader to align with the philosophy of enable_usmp

[NicolaLancellotti]

LGTM! Thank you @ekalda.

[manupak]

Broadly looks great!

just a nit and a question

[manupak]

out of curiosity, should we not check for exact values ?

[ekalda]

Yes good point, I changed it to check for the exact values

[manupak]

I think it would be better for this utility to be part of the cascader and not get exposed to the codegen here.

[ekalda]

Done

[manupak]

LGTM with a suggestion for a follow up.

[manupak]

For a followup : Please consider absorbing the call to extract_memory_info inside the cascader. (Sorry for not being clear before). Ideally, we'd want to remove the "MemoryRegion" construct and to get there in the current direction of travel, we should try to confine the usage of it inside the cascader. Therefore the interface of the cascader should be made to accept MemoryPool(s).

[ekalda]

I agree that this would be the right way to go about it. Let's do it in a follow up!

[manupak]

Thanks @ekalda@mbaret@lhutton1@NicolaLancellotti !

This is merged now!

[ekalda]

Integrate the cascader into the codegen and optionally enable it
with the enable_cascader flag. Includes placeholder MemoryRegions until
integration with the PoolInfos provided by a user.

Co-authored-by: Matthew Barrett matthew.barrett@arm.com

[ekalda]

cc: @manupa-arm@lhutton1@NicolaLancellotti

[ekalda]

This patch supersedes #10377, so that one can be closed.

[manupak]

Thanks @ekalda!

It is nice to see this working and green in CI.

I have left few comments around testing and user facing errors.

[manupak]

[Maybe for a subsequent PR] Will it be possible plumb these values using this :

tvm/src/relay/backend/build_module.cc

Lines 414 to 417 in fcdf463

	IRModule func_module = WithAttrs(IRModule::FromExpr(func),
	{{tvm::attr::kExecutor, executor_},
	{tvm::attr::kRuntime, runtime_},
	{tvm::attr::kWorkspaceMemoryPools, workspace_memory_pools_}});

An example test :

https://github.com/apache/tvm/blob/main/tests/python/relay/aot/test_crt_aot_usmp.py#L341-L383

[ekalda]

Done, with the caveat that we'll assume that there is one workspace pool in the system - is it ok to assume that for now or should we handle the case where there are several workspace pools where some of them are not accessible for the NPU?

[manupak]

For these tests, I think we need to ensure the memory usage is reduced.

We could use calculate_workspace_size TIR analysis utility for that. E.g. :

tvm/tests/python/unittest/test_tir_analysis_calculate_workspace.py

Line 106 in fcdf463

asserttvm.tir.analysis.calculate_workspace_bytes(primfunc, alignment) ==size

OR we could use the final memory calculation in a non-USMP flow as follows :

tvm/tests/python/relay/aot/test_crt_aot.py

Lines 937 to 956 in fcdf463

	deftest_workspace_calculation(workspace_byte_alignment, main_workspace_size):
	mod, params=tvm.relay.testing.synthetic.get_workload()
	target="c"
	runtime=Runtime("crt")
	executor=Executor(
	"aot",
	{
	"workspace-byte-alignment": workspace_byte_alignment,
	},
	)
	withtvm.transform.PassContext(
	opt_level=3,
	config={
	"tir.disable_vectorize": True,
	},
	):
	lib=tvm.relay.build(mod, target, executor=executor, runtime=runtime, params=params)
	mlf_memory_map=mlf._build_function_memory_map(lib.function_metadata)
	assertmlf_memory_map["main"][0]["workspace_size_bytes"] ==main_workspace_size

[ekalda]

I changed the tests to check for reduction in memory instead of bitwise accuracy with TFLite (I think it would be still good to have some FVP based tests when cascader is enabled, but it looks like that would need quite a bit of infra refactor, so I'll do it in a separate patch, if this is ok)

[lhutton1]

Looks good to me modulo @manupa-arm's comments :)

[NicolaLancellotti]

I think it is better is_cascader_enabled, otherwise, it seems you are enabling it with this function.

[ekalda]

Good point! I changed the function name to is_cascader_enabled, but kept the flag/variable as enable_cascader to align with the philosophy of enable_usmp

[NicolaLancellotti]

LGTM! Thank you @ekalda.

[manupak]

Broadly looks great!

just a nit and a question

[manupak]

out of curiosity, should we not check for exact values ?

[ekalda]

Yes good point, I changed it to check for the exact values

[manupak]

I think it would be better for this utility to be part of the cascader and not get exposed to the codegen here.

[ekalda]

Done

[manupak]

LGTM with a suggestion for a follow up.

[manupak]

For a followup : Please consider absorbing the call to extract_memory_info inside the cascader. (Sorry for not being clear before). Ideally, we'd want to remove the "MemoryRegion" construct and to get there in the current direction of travel, we should try to confine the usage of it inside the cascader. Therefore the interface of the cascader should be made to accept MemoryPool(s).

[ekalda]

I agree that this would be the right way to go about it. Let's do it in a follow up!

[manupak]

Thanks @ekalda@mbaret@lhutton1@NicolaLancellotti !

This is merged now!

[ekalda]

Integrate the cascader into the codegen and optionally enable it
with the enable_cascader flag. Includes placeholder MemoryRegions until
integration with the PoolInfos provided by a user.

Co-authored-by: Matthew Barrett matthew.barrett@arm.com

[ekalda]

cc: @manupa-arm@lhutton1@NicolaLancellotti

[ekalda]

This patch supersedes #10377, so that one can be closed.

[manupak]

Thanks @ekalda!

It is nice to see this working and green in CI.

I have left few comments around testing and user facing errors.

[manupak]

[Maybe for a subsequent PR] Will it be possible plumb these values using this :

tvm/src/relay/backend/build_module.cc

Lines 414 to 417 in fcdf463

	IRModule func_module = WithAttrs(IRModule::FromExpr(func),
	{{tvm::attr::kExecutor, executor_},
	{tvm::attr::kRuntime, runtime_},
	{tvm::attr::kWorkspaceMemoryPools, workspace_memory_pools_}});

An example test :

https://github.com/apache/tvm/blob/main/tests/python/relay/aot/test_crt_aot_usmp.py#L341-L383

[ekalda]

Done, with the caveat that we'll assume that there is one workspace pool in the system - is it ok to assume that for now or should we handle the case where there are several workspace pools where some of them are not accessible for the NPU?

[manupak]

For these tests, I think we need to ensure the memory usage is reduced.

We could use calculate_workspace_size TIR analysis utility for that. E.g. :

tvm/tests/python/unittest/test_tir_analysis_calculate_workspace.py

Line 106 in fcdf463

asserttvm.tir.analysis.calculate_workspace_bytes(primfunc, alignment) ==size

OR we could use the final memory calculation in a non-USMP flow as follows :

tvm/tests/python/relay/aot/test_crt_aot.py

Lines 937 to 956 in fcdf463

	deftest_workspace_calculation(workspace_byte_alignment, main_workspace_size):
	mod, params=tvm.relay.testing.synthetic.get_workload()
	target="c"
	runtime=Runtime("crt")
	executor=Executor(
	"aot",
	{
	"workspace-byte-alignment": workspace_byte_alignment,
	},
	)
	withtvm.transform.PassContext(
	opt_level=3,
	config={
	"tir.disable_vectorize": True,
	},
	):
	lib=tvm.relay.build(mod, target, executor=executor, runtime=runtime, params=params)
	mlf_memory_map=mlf._build_function_memory_map(lib.function_metadata)
	assertmlf_memory_map["main"][0]["workspace_size_bytes"] ==main_workspace_size

[ekalda]

I changed the tests to check for reduction in memory instead of bitwise accuracy with TFLite (I think it would be still good to have some FVP based tests when cascader is enabled, but it looks like that would need quite a bit of infra refactor, so I'll do it in a separate patch, if this is ok)

[lhutton1]

Looks good to me modulo @manupa-arm's comments :)

[NicolaLancellotti]

I think it is better is_cascader_enabled, otherwise, it seems you are enabling it with this function.

[ekalda]

Good point! I changed the function name to is_cascader_enabled, but kept the flag/variable as enable_cascader to align with the philosophy of enable_usmp

[NicolaLancellotti]

LGTM! Thank you @ekalda.

[manupak]

Broadly looks great!

just a nit and a question

[manupak]

out of curiosity, should we not check for exact values ?

[ekalda]

Yes good point, I changed it to check for the exact values

[manupak]

I think it would be better for this utility to be part of the cascader and not get exposed to the codegen here.

[ekalda]

Done

[manupak]

LGTM with a suggestion for a follow up.

[manupak]

For a followup : Please consider absorbing the call to extract_memory_info inside the cascader. (Sorry for not being clear before). Ideally, we'd want to remove the "MemoryRegion" construct and to get there in the current direction of travel, we should try to confine the usage of it inside the cascader. Therefore the interface of the cascader should be made to accept MemoryPool(s).

[ekalda]

I agree that this would be the right way to go about it. Let's do it in a follow up!

[manupak]

Thanks @ekalda@mbaret@lhutton1@NicolaLancellotti !

This is merged now!

[ekalda]

Integrate the cascader into the codegen and optionally enable it
with the enable_cascader flag. Includes placeholder MemoryRegions until
integration with the PoolInfos provided by a user.

Co-authored-by: Matthew Barrett matthew.barrett@arm.com

[ekalda]

cc: @manupa-arm@lhutton1@NicolaLancellotti

[ekalda]

This patch supersedes #10377, so that one can be closed.

[manupak]

Thanks @ekalda!

It is nice to see this working and green in CI.

I have left few comments around testing and user facing errors.

[manupak]

[Maybe for a subsequent PR] Will it be possible plumb these values using this :

tvm/src/relay/backend/build_module.cc

Lines 414 to 417 in fcdf463

	IRModule func_module = WithAttrs(IRModule::FromExpr(func),
	{{tvm::attr::kExecutor, executor_},
	{tvm::attr::kRuntime, runtime_},
	{tvm::attr::kWorkspaceMemoryPools, workspace_memory_pools_}});

An example test :

https://github.com/apache/tvm/blob/main/tests/python/relay/aot/test_crt_aot_usmp.py#L341-L383

[ekalda]

Done, with the caveat that we'll assume that there is one workspace pool in the system - is it ok to assume that for now or should we handle the case where there are several workspace pools where some of them are not accessible for the NPU?

[manupak]

For these tests, I think we need to ensure the memory usage is reduced.

We could use calculate_workspace_size TIR analysis utility for that. E.g. :

tvm/tests/python/unittest/test_tir_analysis_calculate_workspace.py

Line 106 in fcdf463

asserttvm.tir.analysis.calculate_workspace_bytes(primfunc, alignment) ==size

OR we could use the final memory calculation in a non-USMP flow as follows :

tvm/tests/python/relay/aot/test_crt_aot.py

Lines 937 to 956 in fcdf463

	deftest_workspace_calculation(workspace_byte_alignment, main_workspace_size):
	mod, params=tvm.relay.testing.synthetic.get_workload()
	target="c"
	runtime=Runtime("crt")
	executor=Executor(
	"aot",
	{
	"workspace-byte-alignment": workspace_byte_alignment,
	},
	)
	withtvm.transform.PassContext(
	opt_level=3,
	config={
	"tir.disable_vectorize": True,
	},
	):
	lib=tvm.relay.build(mod, target, executor=executor, runtime=runtime, params=params)
	mlf_memory_map=mlf._build_function_memory_map(lib.function_metadata)
	assertmlf_memory_map["main"][0]["workspace_size_bytes"] ==main_workspace_size

[ekalda]

I changed the tests to check for reduction in memory instead of bitwise accuracy with TFLite (I think it would be still good to have some FVP based tests when cascader is enabled, but it looks like that would need quite a bit of infra refactor, so I'll do it in a separate patch, if this is ok)

[lhutton1]

Looks good to me modulo @manupa-arm's comments :)

[NicolaLancellotti]

I think it is better is_cascader_enabled, otherwise, it seems you are enabling it with this function.

[ekalda]

Good point! I changed the function name to is_cascader_enabled, but kept the flag/variable as enable_cascader to align with the philosophy of enable_usmp

[NicolaLancellotti]

LGTM! Thank you @ekalda.

[manupak]

Broadly looks great!

just a nit and a question

[manupak]

out of curiosity, should we not check for exact values ?

[ekalda]

Yes good point, I changed it to check for the exact values

[manupak]

I think it would be better for this utility to be part of the cascader and not get exposed to the codegen here.

[ekalda]

Done

[manupak]

LGTM with a suggestion for a follow up.

[manupak]

For a followup : Please consider absorbing the call to extract_memory_info inside the cascader. (Sorry for not being clear before). Ideally, we'd want to remove the "MemoryRegion" construct and to get there in the current direction of travel, we should try to confine the usage of it inside the cascader. Therefore the interface of the cascader should be made to accept MemoryPool(s).

[ekalda]

I agree that this would be the right way to go about it. Let's do it in a follow up!

[manupak]

Thanks @ekalda@mbaret@lhutton1@NicolaLancellotti !

This is merged now!

[ekalda]

Integrate the cascader into the codegen and optionally enable it
with the enable_cascader flag. Includes placeholder MemoryRegions until
integration with the PoolInfos provided by a user.

Co-authored-by: Matthew Barrett matthew.barrett@arm.com

[ekalda]

cc: @manupa-arm@lhutton1@NicolaLancellotti

[ekalda]

This patch supersedes #10377, so that one can be closed.

[manupak]

Thanks @ekalda!

It is nice to see this working and green in CI.

I have left few comments around testing and user facing errors.

[manupak]

[Maybe for a subsequent PR] Will it be possible plumb these values using this :

tvm/src/relay/backend/build_module.cc

Lines 414 to 417 in fcdf463

	IRModule func_module = WithAttrs(IRModule::FromExpr(func),
	{{tvm::attr::kExecutor, executor_},
	{tvm::attr::kRuntime, runtime_},
	{tvm::attr::kWorkspaceMemoryPools, workspace_memory_pools_}});

An example test :

https://github.com/apache/tvm/blob/main/tests/python/relay/aot/test_crt_aot_usmp.py#L341-L383

[ekalda]

Done, with the caveat that we'll assume that there is one workspace pool in the system - is it ok to assume that for now or should we handle the case where there are several workspace pools where some of them are not accessible for the NPU?

[manupak]

For these tests, I think we need to ensure the memory usage is reduced.

We could use calculate_workspace_size TIR analysis utility for that. E.g. :

tvm/tests/python/unittest/test_tir_analysis_calculate_workspace.py

Line 106 in fcdf463

asserttvm.tir.analysis.calculate_workspace_bytes(primfunc, alignment) ==size

OR we could use the final memory calculation in a non-USMP flow as follows :

tvm/tests/python/relay/aot/test_crt_aot.py

Lines 937 to 956 in fcdf463

	deftest_workspace_calculation(workspace_byte_alignment, main_workspace_size):
	mod, params=tvm.relay.testing.synthetic.get_workload()
	target="c"
	runtime=Runtime("crt")
	executor=Executor(
	"aot",
	{
	"workspace-byte-alignment": workspace_byte_alignment,
	},
	)
	withtvm.transform.PassContext(
	opt_level=3,
	config={
	"tir.disable_vectorize": True,
	},
	):
	lib=tvm.relay.build(mod, target, executor=executor, runtime=runtime, params=params)
	mlf_memory_map=mlf._build_function_memory_map(lib.function_metadata)
	assertmlf_memory_map["main"][0]["workspace_size_bytes"] ==main_workspace_size

[ekalda]

I changed the tests to check for reduction in memory instead of bitwise accuracy with TFLite (I think it would be still good to have some FVP based tests when cascader is enabled, but it looks like that would need quite a bit of infra refactor, so I'll do it in a separate patch, if this is ok)

[lhutton1]

Looks good to me modulo @manupa-arm's comments :)

[NicolaLancellotti]

I think it is better is_cascader_enabled, otherwise, it seems you are enabling it with this function.

[ekalda]

Good point! I changed the function name to is_cascader_enabled, but kept the flag/variable as enable_cascader to align with the philosophy of enable_usmp

[NicolaLancellotti]

LGTM! Thank you @ekalda.

[manupak]

Broadly looks great!

just a nit and a question

[manupak]

out of curiosity, should we not check for exact values ?

[ekalda]

Yes good point, I changed it to check for the exact values

[manupak]

I think it would be better for this utility to be part of the cascader and not get exposed to the codegen here.

[ekalda]

Done

[manupak]

LGTM with a suggestion for a follow up.

[manupak]

For a followup : Please consider absorbing the call to extract_memory_info inside the cascader. (Sorry for not being clear before). Ideally, we'd want to remove the "MemoryRegion" construct and to get there in the current direction of travel, we should try to confine the usage of it inside the cascader. Therefore the interface of the cascader should be made to accept MemoryPool(s).

[ekalda]

I agree that this would be the right way to go about it. Let's do it in a follow up!

[manupak]

Thanks @ekalda@mbaret@lhutton1@NicolaLancellotti !

This is merged now!