frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

Update round size with 1 division behavior (#162203)

Browse Source 
have round size return nearest power of 2 greater than or equal to size with 1 division

Fixes #161139

Pull Request resolved: https://github.com/pytorch/pytorch/pull/162203
Approved by: https://github.com/ezyang
This commit is contained in:
morrison-turnansky
2025-10-08 06:41:42 +00:00
committed by PyTorch MergeBot
parent 65aa62d50d
commit 12d2ef557f
2 changed files with 17 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
c10/cuda/CUDACachingAllocator.cpp
View File

@ -2502,6 +2502,8 @@ class DeviceCachingAllocator {
auto divisions = CUDAAllocatorConfig::roundup_power2_divisions(size);
if (divisions > 1 && size > (kMinBlockSize * divisions)) {
return roundup_power2_next_division(size, divisions);
} else if (divisions == 1) {
return llvm::PowerOf2Ceil(size);
} else {
return kMinBlockSize * ((size + kMinBlockSize - 1) / kMinBlockSize);
}