File size: 2,220 Bytes
4a28d4d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
#include "src/turbomind/core/buffer.h"
#include "src/turbomind/core/check.h"
#include "src/turbomind/core/context.h"
#include "src/turbomind/core/data_type.h"
#include "src/turbomind/core/stream.h"
namespace turbomind::core {
Buffer Buffer::view(DataType dtype) const
{
auto b = *this;
if (dtype == dtype_) {
return b;
}
b.dtype_ = dtype;
b.size_ = numel(dtype, byte_size());
if (base_) {
b.base_ = numel(dtype, turbomind::byte_size(dtype_, base_));
}
return b;
}
Buffer Buffer::slice(ssize_t base, ssize_t size) const
{
TM_CHECK_LE(base + size, size_);
auto b = *this;
b.base_ += base;
if (size == -1) {
b.size_ -= base;
}
else {
b.size_ = size;
}
return b;
}
std::ostream& operator<<(std::ostream& os, const Buffer& b)
{
os << b.dtype() << "[" << b.size() << "]@" << b.data_;
if (b.base_) {
os << "+" << b.base_;
}
return os;
}
void Copy(const Buffer& a, ssize_t n, Ref<Buffer> b_, const Stream& stream)
{
auto& b = b_.get();
TM_CHECK_EQ(a.dtype(), b.dtype());
TM_CHECK_LE(n, a.size());
TM_CHECK_LE(n, b.size());
if (auto size = byte_size(a.dtype(), n)) {
check_cuda_error(cudaMemcpyAsync(b.raw_data(), a.raw_data(), size, cudaMemcpyDefault, stream.handle()));
}
}
void Copy(const Buffer& a, ssize_t n, Ref<Buffer> b_)
{
Copy(a, n, b_, Context::stream());
}
void Copy(const Buffer& a, Ref<Buffer> b_, const Stream& stream)
{
TM_CHECK_EQ(a.size(), b_.get().size());
Copy(a, a.size(), b_, stream);
}
void Copy(const Buffer& a, Ref<Buffer> b_)
{
Copy(a, b_, Context::stream());
}
namespace detail {
void* Copy(const void* a, ssize_t n, void* b, const Stream& stream)
{
if (n) {
check_cuda_error(cudaMemcpyAsync(b, a, n, cudaMemcpyDefault, stream.handle()));
}
return (uint8_t*)b + n;
}
} // namespace detail
void Clear(Ref<Buffer> b_, const Stream& stream)
{
auto& b = b_.get();
if (auto size = b.byte_size()) {
check_cuda_error(cudaMemsetAsync(b.raw_data(), 0, b.byte_size(), stream.handle()));
}
}
void Clear(Ref<Buffer> b_)
{
Clear(b_, Context::stream());
}
} // namespace turbomind::core
|