Add is_pinned to host allocator #151439
Add is_pinned to host allocator #151439
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/151439
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (1 Unrelated Failure)As of commit 9553383 with merge base 94da452 ( UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
| bool is_pinned(void* ptr) override { | ||
| // First check if driver is broken/missing, in which case PyTorch CPU | ||
| // functionalities should still work, we should report `false` here. | ||
| if (!at::cuda::is_available()) { |
There was a problem hiding this comment.
can this call itself cause driver initialization?
There was a problem hiding this comment.
Yes, it will call cudaDeviceCount that will initialize driver runtime.
There was a problem hiding this comment.
Shouldn't we check for something like is_initialized? Because if driver is not initialized, ptr is obviously not pinned?
There was a problem hiding this comment.
Got it. I will try to improve it.
There was a problem hiding this comment.
@malfet The logic is already handled at
pytorch/aten/src/ATen/Context.h
Line 104 in ad11d63
There was a problem hiding this comment.
@malfet May I know if I have addressed your comments
There was a problem hiding this comment.
@malfet I introduce one pair API init&is_initialized to imporve the functionality of is_pinned like
bool is_pinned(const void* ptr) const override {
if (!is_initialized()) {
// If the host allocator is not initialized, we assume that the pointer
// isn't pinned to avoid early runtime call.
return false;
}
return impl_->is_pinned(ptr);
}ghstack-source-id: ebb95ba Pull Request resolved: pytorch/pytorch#151439
|
Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as |
Stack from ghstack (oldest at bottom):
Motivation
This PR aims to add the
is_pinnedfunctionality into theHostAllocatorclass, which enables centralized pinned memory verification through calls likeat::getHostAllocator(at::kCUDA)->is_pinned(ptr).Benefits include:
This architecture makes the system more maintainable and extensible for future device support.
Solution
This PR introduces a new pair of API
initandis_initializedinHostAllocatorto avoid early runtime call inis_pinned.Additional Context
It's difficult to deprecate
isPinnedPtrinAcceleratorHooksInterfacebecause some backends (such asmps,hpu,privateuser1) may not register their own host allocator using theREGISTER_HOST_ALLOCATORmechanism, which was introduced in #151431.