[RFC] Add third-party malloc library to improve pytorch memory performance on Windows

[xuhancn]

🚀 The feature, motivation and pitch

This doc is requesting comments for add third-party malloc library to improve pytorch memory performance on Windows.

During debug the issue: #62387 , We figure out the major performance gap between Windows to Linux is that, Windows has bad memory allocation performance.

I also write a simple malloc benchmark project, bench_malloc. Which can proof the third-party malloc (tc_malloc) can improve the memory alloction performance on Windows.

After that, I tried to evaluate some popular third-party malloc library and make a brief summay here:

From the summary, We can only select two candidate libraries:

Option 1: tc_malloc from gperftools.

1. It is the best performance one, and it improve performance from original 11.1s to 2.9s.
2. It need upstream same code changes to gperftools. PRs for pytorch gperftools/gperftools#1396. Actually I made it, but still not get response so far.

Note: I found the gperftools repo is inactivate more than one year, latest commit is gperftools/gperftools@bf8b714 on May 31, 2022

Option 2: mimalloc

1. It could improve performance from original 11.1s to 3.9s. It not as good as tc_malloc, but still a huge improvement.
2. It better compatibility to pytorch, and can integrate to pytorch directly.
3. I have a PR for this option: Enable mimalloc on pytorch Windows #102595

Alternatives

Option 3: Implement a caching memory allocator for CPU in PyTorch.

1. It is none additional depends on third party library.
2. Need a lot of effort to develop and test.
3. Its principle is similar to tc_malloc and mimalloc. Optimize existing third party library is better.

Additional context

---

My proposal

1. I'm not sure whether the PR for gperftools can be accecpted. We can't always wait on option 1.
2. We can enable option 2 (mimalloc) to optimize pytorch Windows firstly.
3. Design build option to switch malloc library, such as system malloc, mimalloc, and (further tc_malloc).
4. Maybe We can select option 2, enable the mimalloc. And then optimize mimalloc for pytorch.

cc @peterjc123 @mszhanyi @skyline75489 @nbcsm @vladimir-aubrecht @iremyux @Blackhex @cristianPanaite @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10 @ngimel

[mjp41]

@xuhancn, If you are interested in trying https://github.com/microsoft/snmalloc, I would be happy to help. It has CMake and Windows support.

[xuhancn]

@xuhancn, If you are interested in trying https://github.com/microsoft/snmalloc, I would be happy to help. It has CMake and Windows support.

I will try it soon. Thanks.

[xuhancn]

@mjp41 Hi, I wrote a simple project to study snmalloc: https://github.com/xuhancn/research_embedded_snmalloc/blob/main/src/main.cpp#L4
I have include the headers and have some question to enable it:

1. how to wrapper and call snmalloc's malloc/free/align_malloc?
2. whether it will override the system malloc/new functions? how to disable override?

[xuhancn]

Option 2 merged:
#102595
#104497

[louie-tsai]

@aice-support