std::align is not supported, so this method mimic its behavior. NOTE(aaroey): according to the TensorRT API, nvinfer1::IGpuAllocator::allocate() uses uint64_t type for size and alignment parameters, so here we use the same type to make it compatible.
| 32 | // nvinfer1::IGpuAllocator::allocate() uses uint64_t type for size and alignment |
| 33 | // parameters, so here we use the same type to make it compatible. |
| 34 | void* Align(uint64_t alignment, uint64_t size, void*& ptr, uint64_t& space) { |
| 35 | QCHECK_GT(alignment, 0ul) << "alignment must be greater than 0."; |
| 36 | QCHECK_EQ(0, alignment & (alignment - 1)) << "Alignment must be power of 2."; |
| 37 | QCHECK_GT(size, 0ul) << "size must be greater than 0."; |
| 38 | QCHECK(ptr) << "ptr must not be nullptr."; |
| 39 | QCHECK_GT(space, 0ul) << "space must be greater than 0."; |
| 40 | const uintptr_t ptr_val = reinterpret_cast<uintptr_t>(ptr); |
| 41 | QCHECK_GE(ptr_val + space, ptr_val) << "Provided space overflows."; |
| 42 | |
| 43 | if (size > space) return nullptr; |
| 44 | const uintptr_t aligned_ptr_val = ((ptr_val + alignment - 1) & -alignment); |
| 45 | if (aligned_ptr_val > ptr_val + space - size) return nullptr; |
| 46 | ptr = reinterpret_cast<void*>(aligned_ptr_val); |
| 47 | const uintptr_t diff = aligned_ptr_val - ptr_val; |
| 48 | space -= diff; |
| 49 | return ptr; |
| 50 | } |
| 51 | |
| 52 | } // namespace tensorrt |
| 53 | } // namespace tensorflow |