| 1815 | |
| 1816 | template <typename T> |
| 1817 | port::Status ROCMBlas::AllocateStridedBuffer( |
| 1818 | const std::vector<typename RocBlasTypeConversionHelper<T>::mapped_type *> |
| 1819 | &raw_ptrs, |
| 1820 | int batch_count, uint64_t batch_stride, ScratchAllocator *scratch_allocator, |
| 1821 | Stream *stream, |
| 1822 | std::unique_ptr<TemporaryDeviceMemory< |
| 1823 | typename RocBlasTypeConversionHelper<T>::mapped_type>> *temp_memory, |
| 1824 | DeviceMemory<typename RocBlasTypeConversionHelper<T>::mapped_type> |
| 1825 | *device_memory) { |
| 1826 | assert(device_memory != nullptr); |
| 1827 | |
| 1828 | using MAPPED_T = typename RocBlasTypeConversionHelper<T>::mapped_type; |
| 1829 | |
| 1830 | bool needs_allocate_strided = false; |
| 1831 | for (int i = 1; i < batch_count; ++i) { |
| 1832 | uint64_t tmp_batch_stride = raw_ptrs[i] - raw_ptrs[i - 1]; |
| 1833 | if (tmp_batch_stride != batch_stride) { |
| 1834 | needs_allocate_strided = true; |
| 1835 | break; |
| 1836 | } |
| 1837 | } |
| 1838 | |
| 1839 | size_t matrix_byte_size = batch_stride * sizeof(MAPPED_T); |
| 1840 | size_t matrix_batch_byte_size = matrix_byte_size * batch_count; |
| 1841 | |
| 1842 | // No need to do re-allocation, take the short cut and return |
| 1843 | if (!needs_allocate_strided) { |
| 1844 | *device_memory = DeviceMemory<MAPPED_T>( |
| 1845 | DeviceMemoryBase(raw_ptrs[0], matrix_batch_byte_size)); |
| 1846 | return port::Status::OK(); |
| 1847 | } |
| 1848 | |
| 1849 | if (scratch_allocator != nullptr) { |
| 1850 | SE_ASSIGN_OR_RETURN( |
| 1851 | DeviceMemory<uint8> batch_matrix_bytes, |
| 1852 | scratch_allocator->AllocateBytes(matrix_batch_byte_size)); |
| 1853 | *device_memory = DeviceMemory<MAPPED_T>(batch_matrix_bytes); |
| 1854 | } else { |
| 1855 | assert(temp_memory != nullptr); |
| 1856 | SE_ASSIGN_OR_RETURN(*temp_memory, stream->AllocateTemporaryArray<MAPPED_T>( |
| 1857 | matrix_batch_byte_size)); |
| 1858 | *device_memory = |
| 1859 | DeviceMemory<MAPPED_T>(*(*temp_memory)->mutable_device_memory()); |
| 1860 | } |
| 1861 | |
| 1862 | for (int i = 0; i < batch_count; ++i) { |
| 1863 | char *device_memory_ptr = static_cast<char *>(device_memory->opaque()); |
| 1864 | DeviceMemoryBase src_mem = DeviceMemoryBase(raw_ptrs[i], matrix_byte_size); |
| 1865 | DeviceMemoryBase target_mem = DeviceMemoryBase( |
| 1866 | device_memory_ptr + i * matrix_byte_size, matrix_byte_size); |
| 1867 | bool a_status = |
| 1868 | stream->ThenMemcpy(&target_mem, src_mem, matrix_byte_size).ok(); |
| 1869 | if (!a_status) { |
| 1870 | return port::Status( |
| 1871 | port::error::INTERNAL, |
| 1872 | "failed to copy device memory in ROCMBlas::DoBlasGemmBatched"); |
| 1873 | } |
| 1874 | } |
nothing calls this directly
no test coverage detected