From ad9a613821ab2ce620e039170d80eb0051df261e Mon Sep 17 00:00:00 2001 From: Luke Hutton Date: Mon, 10 Jun 2024 09:27:59 +0000 Subject: [PATCH 1/2] [AOT] Correctly calculate workspace for vector types When calculating the size of the workspace for a given prim func, the lanes of the data type was not being considered, meaning sizes calculated for dtypes such as "float32x4" were smaller than what they should be. This commit also considers lanes in the calculation. Change-Id: I23a1329ad3c7910784a046e7007a104676ad3664 --- src/tir/usmp/utils.cc | 6 +++++- .../test_tir_analysis_calculate_workspace.py | 17 +++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/src/tir/usmp/utils.cc b/src/tir/usmp/utils.cc index 88a6496859a5..d640e9fa073e 100644 --- a/src/tir/usmp/utils.cc +++ b/src/tir/usmp/utils.cc @@ -181,7 +181,11 @@ Map GetIOPoolAllocations( } static Integer CalculateExtentsSize(const DataType& dtype, const Array& extents) { - size_t element_size_bytes = dtype.bytes(); + if (dtype.is_scalable_vector()) { + // We cannot statically calculate workspace for scalable types + return Integer(); + } + size_t element_size_bytes = dtype.bytes() * dtype.lanes(); size_t num_elements = 1; for (const auto& ext : extents) { if (ext->IsInstance()) { diff --git a/tests/python/tir-analysis/test_tir_analysis_calculate_workspace.py b/tests/python/tir-analysis/test_tir_analysis_calculate_workspace.py index 12c892a04b07..00e9df9d8b34 100644 --- a/tests/python/tir-analysis/test_tir_analysis_calculate_workspace.py +++ b/tests/python/tir-analysis/test_tir_analysis_calculate_workspace.py @@ -91,6 +91,18 @@ def primfunc_local_allocates(placeholder_162: T.handle, placeholder_163: T.handl # fmt: on +@T.prim_func +def prim_func_decl_vector_type(a: T.handle, b: T.handle): + T.func_attr({"tir.noalias": True}) + A = T.match_buffer(a, (4,), "float32x4") + B = T.match_buffer(b, (4,), "float32x4") + C = T.decl_buffer((4,), "float32x4") + for i in range(3): + with T.block("block"): + vi = T.axis.remap("S", [i]) + B[vi] = A[vi] + C[vi] + + @pytest.mark.parametrize("alignment,size,consts", [(1, 663552, 0), (10, 663560, 0)]) def test_global_allocates(alignment, size, consts): primfunc = primfunc_global_allocates @@ -105,6 +117,11 @@ def test_local_allocates(alignment, size, consts): assert tvm.tir.analysis.calculate_workspace_bytes(primfunc, alignment) == size +def test_vector_type(): + primfunc = prim_func_decl_vector_type + assert tvm.tir.analysis.calculate_workspace_bytes(primfunc, 1) == 64 + + if __name__ == "__main__": test_global_allocates() test_local_allocates() From d0a90311ba349e54b1ef96fdd744100ad15e8d41 Mon Sep 17 00:00:00 2001 From: Luke Hutton Date: Mon, 10 Jun 2024 09:57:33 +0000 Subject: [PATCH 2/2] Use tvm.testing.main() Change-Id: Iecc9d3b762eae5c65e5a12c9a803178b976286a6 --- .../tir-analysis/test_tir_analysis_calculate_workspace.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/python/tir-analysis/test_tir_analysis_calculate_workspace.py b/tests/python/tir-analysis/test_tir_analysis_calculate_workspace.py index 00e9df9d8b34..29bfc5845870 100644 --- a/tests/python/tir-analysis/test_tir_analysis_calculate_workspace.py +++ b/tests/python/tir-analysis/test_tir_analysis_calculate_workspace.py @@ -123,5 +123,4 @@ def test_vector_type(): if __name__ == "__main__": - test_global_allocates() - test_local_allocates() + tvm.testing.main()