Add MPS support for getHostAllocator API #151913
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/151913
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 92eaeb9 with merge base f4ac9a1 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
| private: | ||
| bool m_has_unified_memory; | ||
| uint32_t m_usage; | ||
|
|
||
| static void Delete(void* ptr) { |
There was a problem hiding this comment.
Make Delete as public since it is used in MPSHostAllocator.
There was a problem hiding this comment.
It does not move MPS allocator to this unified one, but rather adds a new interface which is not used.
Which makes me wonder what is the point? I.e. shouldn't existing MPSCachingAllocator start inheriting from HostAllocator rather than introducing a new API
And is there a design document somewhere that explains unified interface all accelerators should support
|
|
||
| // A wrapper for MPSAllocator to be used in `getHostAllocator(at::kMPS)` API | ||
| struct MPSHostAllocator final : public at::HostAllocator { | ||
| public: |
There was a problem hiding this comment.
Why do you need to add public for struct, as it's public by default, isn't it?
Different from other backend,
No public document yet, only some comments inside this file. |
Stack from ghstack (oldest at bottom):
Motivation
In #151431, PyTorch provides a unified API
at::getHostAllocator(...)that facilitates writing device-agnostic code. This PR aims to add MPS support forat::getHostAllocatorAPI.Additional Context
The following APIs are not support on
at::getHostAllocator(at::kMPS)yet:at::getHostAllocator()->record_eventat::getHostAllocator()->get_statsat::getHostAllocator()->reset_accumulated_statsat::getHostAllocator()->reset_peak_stats