[TIR][USMP] Greedy memory planning algorithm #9214

manupak · 2021-10-06T18:36:05Z

This commit implements a greedy memory planning algorithms
using the proposed USMP design.There are two greedy algorithms
introduced here which use the size and number of conflicts as the criteria.

Adds few test cases checks for fan-out and linear structures.
Added a test case for ResNet sub-structure
Added a test case for MobileNet sub-structure
This includes a slight fix for buffer info
extraction where non-linear network buffers
owned by the main function should not show
sporadic liveness.

manupak · 2021-11-01T19:34:03Z

a friendly ping @mbs-octoml ! this depends on #8468

mbs-octoml

just nits and my usual plea for more comments :-)

src/tir/usmp/algo/greedy_by_size.cc

* Implementation of greedy by size memory planning algorithm * Added a test case of linear sequence of operators with two pools * Added a test case with residual structures Change-Id: I03b41292eab85ddb43710356c23dd123beb24462

* Adding targets to the PrimFuncs in the tests Change-Id: Ic91947e23cbcc4fc0020eb62f0ed9df26cf919f9

This commit implements greedy algorithms for USMP based on size and number of liveness conflicts. * This includes a slight fix for buffer info extraction where non-linear network buffers owned by the main function should not show sporadic liveness. Change-Id: I957d543e75b3b0bcf5fc1fbc7870705c875c7a03

manupak · 2021-11-29T19:08:05Z

Thanks! @mbs-octoml .

I finally got around to address the comments and did some cleanup.

There was some remanants of unimplemented python APIs related BufferInfo. They are removed now. Change-Id: I4de1d817eb34187bc20da2ac2b1cb0da5b372833

Change-Id: Ia2d24753398bb388918aecba3b0191100d5100a6

manupak · 2021-11-30T10:46:44Z

@mbaret Could you take look at this PR when you have some time ?

mbaret · 2021-11-30T10:48:21Z

src/tir/usmp/algo/greedy.cc

+ */
+
+/*!
+ * \file tir/analysis/usmp/algo/greedy_by_size.cc


doesn't match current location

mbaret · 2021-11-30T11:05:49Z

tests/python/unittest/test_tir_usmp_algo.py

+    return allocates
+
+
+def assign_poolinfos_to_allocates_in_primfunc(primfunc, pool_infos):


should these also be 'private' i.e. _assign?

mbaret · 2021-11-30T11:10:58Z

tests/python/unittest/test_tir_usmp_algo.py

+    ["algorithm", "fast_memory_size", "slow_memory_size"],
+    [("greedy_by_size", 200704, 1418528), ("greedy_by_conflicts", 200704, 1418528)],
+)
+def test_linear(algorithm, fast_memory_size, slow_memory_size):


As a general comment on these tests, I'd normally favour a larger number of smaller cases. It's quite hard to parse these complex TIR modules and determine the expected behaviour of the test. Additionally, it's ambiguous what coverage we have other than it works in a particular linear and fan-out case.

The unit testing in this level is to make sure the component is functional and ensure the contract of the component. As long as greedy algorithm goes covering a linear and fan-out case covers most aspect conflict graph one could end up.
My personal preference would be to check coverage at the integration level of the USMP using these algorithm rather than using hand crafted TIR snippets at this level. WDYT ?

Alternatively, if you had some smaller cases in mind that you'd like tested here, please do share -- so we could consider adding them here.

I think it becomes a bit harder to argue for testing coverage in integration tests when this algorithm component is 'swappable'. That is, unless, those integration tests will be extended for every new algorithm that's added (which is plausible).

Regarding these specific tests, I think they would be a lot more readable using small synthetic tests rather than network extracts. The linear case has a large amount of complexity that doesn't contribute to the test if topology is the only impactful change.

That is, unless, those integration tests will be extended for every new algorithm that's added (which is plausible).

Yes, we will be doing this as that would make the most sense as that is the only sensible way to make sure the number workspace size decided by the algorithm, actually turns out to be something that runs.

Regarding these specific tests, I think they would be a lot more readable using small synthetic tests rather than network extracts. The linear case has a large amount of complexity that doesn't contribute to the test if topology is the only impactful change.

I think, in terms of unit testing this algorithm, I ll create BufferInfo object based tests rather than TIR ones -- truly thats the input to the algorithm.

src/tir/usmp/algo/greedy.cc

mbaret · 2021-11-30T11:12:24Z

src/tir/usmp/analysis/extract_buffer_info.cc

@@ -415,18 +454,30 @@ Map<BufferInfo, tir::Stmt> BufferInfoExtractor::operator()(const PrimFunc& main_

  // Traverse the liveness events using a open set to track what
  // is live while updating the conflicts through out the linear traversal
-  std::unordered_set<BufferInfo, ObjectPtrHash, ObjectPtrEqual> open_set;
+  //  std::unordered_set<BufferInfo, ObjectPtrHash, ObjectPtrEqual> open_set;


tests/python/unittest/test_tir_usmp_algo.py

src/tir/usmp/analysis/extract_buffer_info.cc

mbaret · 2021-11-30T11:18:56Z

tests/python/unittest/test_tir_usmp_algo.py

+    fusmp_algo_greedy_by_size = tvm.get_global_func(f"tir.usmp.algo.{algorithm}")
+    buffer_pool_allocations = fusmp_algo_greedy_by_size(buffer_info_arr)


The naming here (and fusmp_algo_greedy_by_size) is a bit confusing because 'algorithm' could be by conflict I think.

Thanks! Ack -- will change.

src/tir/usmp/algo/greedy.cc

mbaret · 2021-11-30T11:27:57Z

src/tir/usmp/algo/greedy.cc

+        return pool_info;
+      }
+    }
+    ICHECK(false) << "TVM USMP Internal Error: no candidate have been selected!";


When would this check be reached (if ever)? It seems like it may be hit if every pool is full, in which case I think the reported error should be quite explicit as it's something the user could reasonably run into. We should also probably test it as I don't think it would be a 'true assert'.

Yea, make sense, will change.

manupak

Thanks @mbaret . See my responses and I ll do the acknowledged changes.

manupak · 2021-11-30T11:36:09Z

src/tir/usmp/algo/greedy.cc

+ */
+
+/*!
+ * \file tir/analysis/usmp/algo/greedy_by_size.cc


manupak · 2021-11-30T11:36:48Z

src/tir/usmp/algo/greedy.cc

+        return pool_info;
+      }
+    }
+    ICHECK(false) << "TVM USMP Internal Error: no candidate have been selected!";


Yea, make sense, will change.

src/tir/usmp/algo/greedy.cc

src/tir/usmp/analysis/extract_buffer_info.cc

manupak · 2021-11-30T11:44:43Z

src/tir/usmp/analysis/extract_buffer_info.cc

+      } else {
+        open_set[le_event.buffer_info] -= 1;
+      }
+      //      open_set.erase(le_event.buffer_info);


manupak · 2021-11-30T11:44:59Z

tests/python/unittest/test_tir_usmp_algo.py

+    return allocates
+
+
+def assign_poolinfos_to_allocates_in_primfunc(primfunc, pool_infos):


manupak · 2021-11-30T11:49:03Z

tests/python/unittest/test_tir_usmp_algo.py

+    ["algorithm", "fast_memory_size", "slow_memory_size"],
+    [("greedy_by_size", 200704, 1418528), ("greedy_by_conflicts", 200704, 1418528)],
+)
+def test_linear(algorithm, fast_memory_size, slow_memory_size):


The unit testing in this level is to make sure the component is functional and ensure the contract of the component. As long as greedy algorithm goes covering a linear and fan-out case covers most aspect conflict graph one could end up.
My personal preference would be to check coverage at the integration level of the USMP using these algorithm rather than using hand crafted TIR snippets at this level. WDYT ?

Alternatively, if you had some smaller cases in mind that you'd like tested here, please do share -- so we could consider adding them here.

manupak · 2021-11-30T11:49:36Z

tests/python/unittest/test_tir_usmp_algo.py

+    fusmp_algo_greedy_by_size = tvm.get_global_func(f"tir.usmp.algo.{algorithm}")
+    buffer_pool_allocations = fusmp_algo_greedy_by_size(buffer_info_arr)


Thanks! Ack -- will change.

tests/python/unittest/test_tir_usmp_algo.py

manupak

@mbaret . I think I've addresssed you comments by now. PTAL when you have some time.

src/tir/usmp/algo/greedy.cc

tests/python/unittest/test_tir_usmp_algo.py

This commits removes commented out lines ,few trivial cleanups and few BufferInfo based tests to check the algorithm. Change-Id: I1a12b6a424370e9e4c4a55563dde0ad698b07ea3

Changed sorting criteria use alphabetic ordering as opposed to hashes of string as it seemed different accross different platforms. Change-Id: Ia7938d1b0d1374924c3ec7287526ccf374c54eb7

mbaret · 2021-12-01T14:30:40Z

src/tir/usmp/algo/greedy.cc

+        return pool_info;
+      }
+    }
+    CHECK(false) << "TVM USMP Error: no candidate have been selected for " << buf_info;


I think this needs to be more obvious because it'll be directly user-facing. Probably it needs to indicate that TVM has exceeded the pool memory limits and suggest to increase them if you want the compilation to pass. Let's not block the whole patch on this though, happy to take it in a follow-up.

mbaret

LGTM

Improving the error message Change-Id: Ib59efb172fe10b70f88a24f4874a7891e8a9cde7

manupak · 2021-12-01T15:19:22Z

Ah the CI decided to run itself again, thus I've changed the error message.

mbaret · 2021-12-02T09:53:04Z

This is now merged, thanks @manupa-arm, @mbs-octoml!

This commit implements a greedy memory planning algorithms using the proposed USMP design.There are two greedy algorithms introduced here which use the size and number of conflicts as the criteria. - Adds few test cases checks for fan-out and linear structures. - Added a test case for ResNet sub-structure - Added a test case for MobileNet sub-structure - This includes a slight fix for buffer info extraction where non-linear network buffers owned by the main function should not show sporadic liveness.

manupak requested review from areusch, comaniac, Hzfengsy, jroesch, junrushao, kparzysz-quic, masahi, merrymercy, tqchen, vinx13, yzhliu, zhiics and ZihengJiang as code owners October 6, 2021 18:36

manupak marked this pull request as draft October 6, 2021 18:36

manupak mentioned this pull request Oct 6, 2021

[RFC] Unified Static Memory Planning (USMP) Tracking Issue #8404

Closed

21 tasks

manupak force-pushed the usmp_greedy_by_size branch from 01b2055 to 897070d Compare October 7, 2021 16:57

manupak force-pushed the usmp_greedy_by_size branch from 897070d to ec29b89 Compare October 21, 2021 17:44

manupak force-pushed the usmp_greedy_by_size branch from ec29b89 to ceb9abe Compare October 30, 2021 06:41

manupak mentioned this pull request Nov 1, 2021

[TIR][USMP] adding the pass to convert to pool offsets #9418

Merged

areusch self-assigned this Nov 9, 2021

manupak force-pushed the usmp_greedy_by_size branch 3 times, most recently from 155f7a3 to 313e1ba Compare November 22, 2021 14:49

manupak mentioned this pull request Nov 23, 2021

[TIR][USMP] Integrating USMP to AoT Executor #9565

Merged

manupak marked this pull request as ready for review November 23, 2021 17:37

manupak requested a review from icemelon as a code owner November 23, 2021 17:37

manupak force-pushed the usmp_greedy_by_size branch from 313e1ba to 4e95aed Compare November 23, 2021 17:50

mbs-octoml approved these changes Nov 24, 2021

View reviewed changes

[TIR][USMP] greedy_by_size usmp algo

ecb3c3f

* Implementation of greedy by size memory planning algorithm * Added a test case of linear sequence of operators with two pools * Added a test case with residual structures Change-Id: I03b41292eab85ddb43710356c23dd123beb24462

manupak added 2 commits November 29, 2021 15:53

[TIR][USMP] greedy_by_size usmp algo

e0ca1d9

* Adding targets to the PrimFuncs in the tests Change-Id: Ic91947e23cbcc4fc0020eb62f0ed9df26cf919f9

manupak force-pushed the usmp_greedy_by_size branch from 4e95aed to 26fffb0 Compare November 29, 2021 19:03

manupak added 2 commits November 29, 2021 19:10

Removing unimplemented Python-APi for USMP

1d57832

There was some remanants of unimplemented python APIs related BufferInfo. They are removed now. Change-Id: I4de1d817eb34187bc20da2ac2b1cb0da5b372833

Removing commented code from extract buffer info pass

78e099b

Change-Id: Ia2d24753398bb388918aecba3b0191100d5100a6

mbaret reviewed Nov 30, 2021

View reviewed changes

manupak commented Nov 30, 2021

View reviewed changes

src/tir/usmp/algo/greedy.cc Outdated Show resolved Hide resolved

tests/python/unittest/test_tir_usmp_algo.py Show resolved Hide resolved

[TIR][USMP] Greedy algorithms for USMP

7c93c60

This commits removes commented out lines ,few trivial cleanups and few BufferInfo based tests to check the algorithm. Change-Id: I1a12b6a424370e9e4c4a55563dde0ad698b07ea3

manupak force-pushed the usmp_greedy_by_size branch from 5e75ed3 to 7c93c60 Compare November 30, 2021 20:09

[TIR][USMP] Greedy algorithms for USMP

6101e61

Changed sorting criteria use alphabetic ordering as opposed to hashes of string as it seemed different accross different platforms. Change-Id: Ia7938d1b0d1374924c3ec7287526ccf374c54eb7

manupak force-pushed the usmp_greedy_by_size branch from c675bb0 to 6101e61 Compare December 1, 2021 06:59

mbaret reviewed Dec 1, 2021

View reviewed changes

mbaret approved these changes Dec 1, 2021

View reviewed changes

[TIR][USMP] Greedy algorithms for USMP

87b0ac9

Improving the error message Change-Id: Ib59efb172fe10b70f88a24f4874a7891e8a9cde7

mbaret merged commit 756debc into apache:main Dec 2, 2021

driazati mentioned this pull request Jul 14, 2022

TVM v0.9.0.rc0 Release Candidate Notes #12102

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[TIR][USMP] Greedy memory planning algorithm #9214

[TIR][USMP] Greedy memory planning algorithm #9214

manupak commented Oct 6, 2021 •

edited

Loading

manupak commented Nov 1, 2021

mbs-octoml left a comment

manupak commented Nov 29, 2021

manupak commented Nov 30, 2021

mbaret Nov 30, 2021

manupak Nov 30, 2021

mbaret Nov 30, 2021

manupak Nov 30, 2021

mbaret Nov 30, 2021

manupak Nov 30, 2021

mbaret Nov 30, 2021

manupak Nov 30, 2021 •

edited

Loading

mbaret Nov 30, 2021

manupak Nov 30, 2021

mbaret Nov 30, 2021

manupak Nov 30, 2021

mbaret Nov 30, 2021

manupak Nov 30, 2021

manupak left a comment

manupak Nov 30, 2021

manupak Nov 30, 2021

manupak Nov 30, 2021

manupak Nov 30, 2021

manupak Nov 30, 2021

manupak Nov 30, 2021

manupak left a comment

mbaret Dec 1, 2021

mbaret left a comment

manupak commented Dec 1, 2021

mbaret commented Dec 2, 2021

		return allocates


		def assign_poolinfos_to_allocates_in_primfunc(primfunc, pool_infos):

		fusmp_algo_greedy_by_size = tvm.get_global_func(f"tir.usmp.algo.{algorithm}")
		buffer_pool_allocations = fusmp_algo_greedy_by_size(buffer_info_arr)

[TIR][USMP] Greedy memory planning algorithm #9214

[TIR][USMP] Greedy memory planning algorithm #9214

Conversation

manupak commented Oct 6, 2021 • edited Loading

manupak commented Nov 1, 2021

mbs-octoml left a comment

Choose a reason for hiding this comment

manupak commented Nov 29, 2021

manupak commented Nov 30, 2021

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

manupak Nov 30, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

manupak left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

manupak left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mbaret left a comment

Choose a reason for hiding this comment

manupak commented Dec 1, 2021

mbaret commented Dec 2, 2021

manupak commented Oct 6, 2021 •

edited

Loading

manupak Nov 30, 2021 •

edited

Loading