| 448 | } |
| 449 | |
| 450 | void* GPUTensorPoolAllocator::BigAllocate(size_t alignment, |
| 451 | size_t num_bytes) { |
| 452 | auto id = Index(num_bytes, alignment_, alignment_offset_); |
| 453 | if (unlikely(id < 0)) { |
| 454 | return sub_allocator_->Alloc(alignment, num_bytes); |
| 455 | } |
| 456 | |
| 457 | auto b = GetBin(id); |
| 458 | if (unlikely(b == nullptr)) { |
| 459 | return sub_allocator_->Alloc(alignment, num_bytes); |
| 460 | } |
| 461 | |
| 462 | auto ptr = b->Allocate(); |
| 463 | if (likely(ptr != nullptr)) { |
| 464 | return ptr; |
| 465 | } |
| 466 | |
| 467 | return sub_allocator_->Alloc(alignment, num_bytes); |
| 468 | } |
| 469 | |
| 470 | // unlikely execute this path which do some atomic operations |
| 471 | void* GPUTensorPoolAllocator::BigAllocateStatistic(size_t alignment, |