diff --git a/python/tvm/target/target.py b/python/tvm/target/target.py index fb7c2827efca..b027b99b17eb 100644 --- a/python/tvm/target/target.py +++ b/python/tvm/target/target.py @@ -699,6 +699,17 @@ def get_arch_version(cpu_ver): msg = "{} is not a valid Hexagon version\nvalid versions include {}" raise ValueError(msg.format(cpu_ver, valid_hex)) from None + def get_vtcm_capacity(cpu_ver): + one_mb = 2**20 + default_vtcm_sizes = { + "v65": one_mb // 4, + "v66": one_mb // 4, + "v68": 4 * one_mb, + "v69": 8 * one_mb, + "v73": 8 * one_mb, + } + return default_vtcm_sizes.get(cpu_ver, 0) + # Target configuration: arch_version = get_arch_version(cpu_ver) config = { @@ -706,7 +717,7 @@ def get_arch_version(cpu_ver): "llvm_options": None, "use_qfloat": arch_version >= 68, "use_ieee_fp": False, - "vtcm_capacity": 0, + "vtcm_capacity": get_vtcm_capacity(cpu_ver), } config.update(kwargs) diff --git a/tests/python/contrib/test_hexagon/test_vtcm.py b/tests/python/contrib/test_hexagon/test_vtcm.py index a549588e5768..524bcb859990 100644 --- a/tests/python/contrib/test_hexagon/test_vtcm.py +++ b/tests/python/contrib/test_hexagon/test_vtcm.py @@ -79,7 +79,9 @@ def _raises_exception(f): with tvm.transform.PassContext(config={"tir.vtcm_capacity": vtcm_capacity}): assert ( - _raises_exception(lambda: tvm.build(sch.mod, target=get_hexagon_target("v68"))) + _raises_exception( + lambda: tvm.build(sch.mod, target=get_hexagon_target("v68", vtcm_capacity=0)) + ) == limited ), "Case 3 - context. VTCM memory allocation limiter does not work correctly "