Skip to content

Commit

Permalink
clean up byte buffer copy
Browse files Browse the repository at this point in the history
  • Loading branch information
Mike-Leo-Smith committed Feb 11, 2025
1 parent 5646420 commit c353e5d
Show file tree
Hide file tree
Showing 3 changed files with 176 additions and 104 deletions.
35 changes: 32 additions & 3 deletions include/luisa/runtime/buffer.h
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
#pragma once

#include <luisa/core/concepts.h>
#include <luisa/core/mathematics.h>
#include <luisa/runtime/rhi/command.h>
#include <luisa/runtime/rhi/resource.h>
Expand Down Expand Up @@ -36,6 +35,9 @@ struct buffer_element_impl {

}// namespace detail

class ByteBuffer;
class ByteBufferView;

template<typename T>
using is_buffer = detail::is_buffer_impl<std::remove_cvref_t<T>>;

Expand Down Expand Up @@ -135,7 +137,7 @@ class Buffer final : public Resource {
[[nodiscard]] auto view() const noexcept {
_check_is_valid();
return BufferView<T>{this->native_handle(), this->handle(), _element_stride, 0u, _size, _size};
}
}
[[nodiscard]] auto view(size_t offset, size_t count) const noexcept {
return view().subview(offset, count);
}
Expand All @@ -144,17 +146,29 @@ class Buffer final : public Resource {
[[nodiscard]] auto copy_to(void *data) const noexcept {
return this->view().copy_to(data);
}
// copy buffer's data to another buffer
[[nodiscard]] auto copy_to(BufferView<T> dst) const noexcept {
return this->view().copy_to(dst);
}
// copy buffer's data to a byte buffer
[[nodiscard]] auto copy_to(const ByteBufferView &dst) const noexcept {
return this->view().copy_to(dst);
}
// copy pointer's data to buffer
[[nodiscard]] auto copy_from(const void *data) const noexcept {
return this->view().copy_from(data);
}
[[nodiscard]] auto copy_from(const void *data, luisa::move_only_function<void(void*)>&& upload_callback) const noexcept {
[[nodiscard]] auto copy_from(const void *data, luisa::move_only_function<void(void *)> &&upload_callback) const noexcept {
return this->view().copy_from(data, std::move(upload_callback));
}
// copy source buffer's data to buffer
[[nodiscard]] auto copy_from(BufferView<T> source) const noexcept {
return this->view().copy_from(source);
}
// copy source byte buffer's data to buffer
[[nodiscard]] auto copy_from(const ByteBufferView &source) const noexcept {
return this->view().copy_from(source);
}
// DSL interface
[[nodiscard]] auto operator->() const noexcept {
_check_is_valid();
Expand Down Expand Up @@ -241,6 +255,12 @@ class BufferView {
[[nodiscard]] auto copy_to(void *data) const noexcept {
return luisa::make_unique<BufferDownloadCommand>(_handle, offset_bytes(), size_bytes(), data);
}
// copy buffer's data to another buffer
[[nodiscard]] auto copy_to(BufferView<T> dst) const noexcept {
return dst.copy_from(*this);
}
// copy buffer's data to a byte buffer
[[nodiscard]] luisa::unique_ptr<BufferCopyCommand> copy_to(const ByteBufferView &dst) const noexcept;
// copy pointer's data to buffer
[[nodiscard]] auto copy_from(const void *data) const noexcept {
return luisa::make_unique<BufferUploadCommand>(this->handle(), this->offset_bytes(), this->size_bytes(), data);
Expand All @@ -255,6 +275,8 @@ class BufferView {
source.offset_bytes(), this->offset_bytes(),
this->size_bytes());
}
// copy source byte buffer's data to buffer
[[nodiscard]] luisa::unique_ptr<BufferCopyCommand> copy_from(const ByteBufferView &source) const noexcept;
// DSL interface
[[nodiscard]] auto operator->() const noexcept {
return reinterpret_cast<const detail::BufferExprProxy<BufferView<T>> *>(this);
Expand All @@ -268,12 +290,19 @@ template<typename T>
BufferView(BufferView<T>) -> BufferView<T>;

namespace detail {

template<typename T>
struct is_buffer_impl<Buffer<T>> : std::true_type {};

template<typename T>
struct is_buffer_view_impl<BufferView<T>> : std::true_type {};

template<>
struct is_buffer_impl<ByteBuffer> : std::true_type {};

template<>
struct is_buffer_view_impl<ByteBufferView> : std::true_type {};

template<typename T>
struct buffer_element_impl<Buffer<T>> {
using type = T;
Expand Down
183 changes: 88 additions & 95 deletions include/luisa/runtime/byte_buffer.h
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ namespace luisa::compute {
namespace detail {
class ByteBufferExprProxy;
}// namespace detail

class ByteBufferView;

class LC_RUNTIME_API ByteBuffer final : public Resource {
Expand All @@ -32,46 +33,26 @@ class LC_RUNTIME_API ByteBuffer final : public Resource {
}
ByteBuffer &operator=(ByteBuffer const &) noexcept = delete;
using Resource::operator bool;
[[nodiscard]] auto copy_to(void *data) const noexcept {
_check_is_valid();
return luisa::make_unique<BufferDownloadCommand>(handle(), 0u, _size_bytes, data);
}
[[nodiscard]] auto copy_from(const void *data) noexcept {
_check_is_valid();
return luisa::make_unique<BufferUploadCommand>(handle(), 0u, _size_bytes, data);
}
[[nodiscard]] auto copy_from(const void *data, size_t buffer_offset, size_t size_bytes) noexcept {
_check_is_valid();
if (size_bytes > _size_bytes) [[unlikely]] {
detail::error_buffer_copy_sizes_mismatch(size_bytes, _size_bytes);
}
return luisa::make_unique<BufferUploadCommand>(handle(), buffer_offset, size_bytes, data);
}
template<typename T>
[[nodiscard]] auto copy_from(BufferView<T> source) noexcept {
_check_is_valid();
if (source.size_bytes() != _size_bytes) [[unlikely]] {
detail::error_buffer_copy_sizes_mismatch(source.size_bytes(), _size_bytes);
}
return luisa::make_unique<BufferCopyCommand>(
source.handle(), this->handle(),
source.offset_bytes(), 0u,
this->size_bytes());
}
[[nodiscard]] auto copy_from(const ByteBuffer &source, size_t offset, size_t size_bytes) noexcept {
_check_is_valid();
if (size_bytes > _size_bytes) [[unlikely]] {
detail::error_buffer_copy_sizes_mismatch(size_bytes, _size_bytes);
}
return luisa::make_unique<BufferCopyCommand>(
source.handle(), this->handle(),
offset, 0u,
size_bytes);
}

[[nodiscard]] ByteBufferView view() const noexcept;

[[nodiscard]] luisa::unique_ptr<BufferUploadCommand> copy_from(const void *data) const noexcept;
[[nodiscard]] luisa::unique_ptr<BufferCopyCommand> copy_from(ByteBufferView source) const noexcept;
[[nodiscard]] luisa::unique_ptr<BufferDownloadCommand> copy_to(void *data) const noexcept;
[[nodiscard]] luisa::unique_ptr<BufferCopyCommand> copy_to(ByteBufferView source) const noexcept;

template<typename T>
[[nodiscard]] luisa::unique_ptr<BufferCopyCommand> copy_from(const Buffer<T> &source) const noexcept;

template<typename T>
[[nodiscard]] luisa::unique_ptr<BufferCopyCommand> copy_from(BufferView<T> source) const noexcept;

template<typename T>
[[nodiscard]] auto copy_from(ByteBufferView source) const noexcept;
[[nodiscard]] auto copy_from(ByteBufferView source, size_t offset, size_t size_bytes) const noexcept;
[[nodiscard]] luisa::unique_ptr<BufferCopyCommand> copy_to(const Buffer<T> &dst) const noexcept;

template<typename T>
[[nodiscard]] luisa::unique_ptr<BufferCopyCommand> copy_to(BufferView<T> dst) const noexcept;

// DSL interface
[[nodiscard]] auto operator->() const noexcept {
_check_is_valid();
Expand All @@ -98,17 +79,16 @@ class ByteBufferView {
}

ByteBufferView(const ByteBuffer &buffer) noexcept : ByteBufferView{buffer.view()} {}
ByteBufferView(const ByteBufferView &) noexcept = default;

template<typename T>
requires(is_buffer_v<T> && !std::is_same_v<ByteBuffer, std::remove_cvref_t<T>>)
ByteBufferView(T const &buffer) : ByteBufferView{
buffer.native_handle(), buffer.handle(),
0, buffer.size_bytes(), buffer.size_bytes()} {}
explicit ByteBufferView(BufferView<T> buffer_view) noexcept
: ByteBufferView{buffer_view.native_handle(), buffer_view.handle(),
buffer_view.offset_bytes(), buffer_view.size_bytes(), buffer_view.total_size_bytes()} {}

template<typename T>
requires(is_buffer_view_v<T> && !std::is_same_v<ByteBufferView, std::remove_cvref_t<T>>)
ByteBufferView(T buffer_view) : ByteBufferView{
buffer_view.native_handle(), buffer_view.handle(),
buffer_view.offset_bytes(), buffer_view.size_bytes(), buffer_view.total_size_bytes()} {}
explicit ByteBufferView(const Buffer<T> &buffer) noexcept
: ByteBufferView{buffer.view()} {}

ByteBufferView() noexcept : ByteBufferView{nullptr, invalid_resource_handle, 0, 0, 0} {}
[[nodiscard]] explicit operator bool() const noexcept { return _handle != invalid_resource_handle; }
Expand All @@ -120,17 +100,8 @@ class ByteBufferView {
[[nodiscard]] auto size_bytes() const noexcept { return _size; }
[[nodiscard]] auto total_size_bytes() const noexcept { return _total_size; }

[[nodiscard]] auto original() const noexcept {
return ByteBufferView{_native_handle, _handle,
0u, _total_size, _total_size};
}
[[nodiscard]] auto subview(size_t offset_bytes, size_t size_bytes) const noexcept {
if (offset_bytes + size_bytes > _size) [[unlikely]] {
detail::error_buffer_subview_overflow(offset_bytes, size_bytes, _size);
}
return ByteBufferView{_native_handle, _handle, _offset_bytes + offset_bytes,
size_bytes, _total_size};
}
[[nodiscard]] ByteBufferView original() const noexcept;
[[nodiscard]] ByteBufferView subview(size_t offset_bytes, size_t size_bytes) const noexcept;

template<typename U>
requires(!is_custom_struct_v<U>)
Expand All @@ -141,51 +112,73 @@ class ByteBufferView {
return BufferView<U>{_native_handle, _handle, sizeof(U), _offset_bytes,
this->size_bytes() / sizeof(U), _total_size / sizeof(U)};
}
[[nodiscard]] auto copy_to(void *data) const noexcept {
return luisa::make_unique<BufferDownloadCommand>(_handle, offset_bytes(), size_bytes(), data);
}
[[nodiscard]] luisa::unique_ptr<BufferDownloadCommand> copy_to(void *data) const noexcept;
// copy pointer's data to buffer
[[nodiscard]] auto copy_from(const void *data) const noexcept {
return luisa::make_unique<BufferUploadCommand>(this->handle(), this->offset_bytes(), this->size_bytes(), data);
}
[[nodiscard]] luisa::unique_ptr<BufferUploadCommand> copy_from(const void *data) const noexcept;
// copy source buffer's data to buffer
[[nodiscard]] auto copy_from(ByteBufferView source) const noexcept {
if (source.size_bytes() != this->size_bytes()) [[unlikely]] {
detail::error_buffer_copy_sizes_mismatch(source.size_bytes(), this->size_bytes());
}
return luisa::make_unique<BufferCopyCommand>(
source.handle(), this->handle(),
source.offset_bytes(), this->offset_bytes(),
this->size_bytes());
}
[[nodiscard]] luisa::unique_ptr<BufferCopyCommand> copy_from(ByteBufferView source) const noexcept;
// copy buffer's data to source buffer
[[nodiscard]] luisa::unique_ptr<BufferCopyCommand> copy_to(ByteBufferView source) const noexcept;

template<typename T>
[[nodiscard]] luisa::unique_ptr<BufferCopyCommand> copy_to(const Buffer<T> &dst) const noexcept { return copy_to(ByteBufferView{dst}); }
template<typename T>
[[nodiscard]] luisa::unique_ptr<BufferCopyCommand> copy_to(BufferView<T> dst) const noexcept { return copy_to(ByteBufferView{dst}); }
template<typename T>
[[nodiscard]] luisa::unique_ptr<BufferCopyCommand> copy_from(const Buffer<T> &source) const noexcept { return copy_from(ByteBufferView{source}); }
template<typename T>
[[nodiscard]] luisa::unique_ptr<BufferCopyCommand> copy_from(BufferView<T> source) const noexcept { return copy_from(ByteBufferView{source}); }
};

// some type traits
template<typename T>
[[nodiscard]] inline auto ByteBuffer::copy_from(ByteBufferView source) const noexcept {
_check_is_valid();
if (source.size_bytes() != _size_bytes) [[unlikely]] {
detail::error_buffer_copy_sizes_mismatch(source.size_bytes(), _size_bytes);
}
return luisa::make_unique<BufferCopyCommand>(
source.handle(), this->handle(),
source.offset_bytes(), 0u,
this->size_bytes());
struct is_byte_buffer : std::is_same<std::remove_cvref_t<T>, ByteBuffer> {};

template<typename T>
struct is_byte_buffer_view : std::is_same<std::remove_cvref_t<T>, ByteBufferView> {};

template<typename T>
struct is_byte_buffer_or_view : std::disjunction<is_byte_buffer<T>, is_byte_buffer_view<T>> {};

template<typename T>
constexpr auto is_byte_buffer_v = is_byte_buffer<T>::value;

template<typename T>
constexpr auto is_byte_buffer_view_v = is_byte_buffer_view<T>::value;

template<typename T>
constexpr auto is_byte_buffer_or_view_v = is_byte_buffer_or_view<T>::value;

// implementations for ByteBuffer
template<typename T>
luisa::unique_ptr<BufferCopyCommand> ByteBuffer::copy_from(const Buffer<T> &source) const noexcept {
return view().copy_from(source);
}
[[nodiscard]] inline auto ByteBuffer::copy_from(ByteBufferView source, size_t offset, size_t size_bytes) const noexcept {
_check_is_valid();
if (size_bytes > _size_bytes) [[unlikely]] {
detail::error_buffer_copy_sizes_mismatch(size_bytes, _size_bytes);
}
return luisa::make_unique<BufferCopyCommand>(
source.handle(), this->handle(),
offset, 0u,
size_bytes);

template<typename T>
luisa::unique_ptr<BufferCopyCommand> ByteBuffer::copy_from(BufferView<T> source) const noexcept {
return view().copy_from(source);
}

namespace detail {
LC_RUNTIME_API void error_buffer_size_not_aligned(size_t align) noexcept;
template<>
struct is_buffer_impl<ByteBuffer> : std::true_type {};
}// namespace detail
template<typename T>
luisa::unique_ptr<BufferCopyCommand> ByteBuffer::copy_to(const Buffer<T> &dst) const noexcept {
return view().copy_to(dst);
}

template<typename T>
luisa::unique_ptr<BufferCopyCommand> ByteBuffer::copy_to(BufferView<T> dst) const noexcept {
return view().copy_to(dst);
}

// implementations for typed Buffer
template<typename T>
luisa::unique_ptr<BufferCopyCommand> BufferView<T>::copy_to(const ByteBufferView &dst) const noexcept {
return dst.copy_from(*this);
}

template<typename T>
luisa::unique_ptr<BufferCopyCommand> BufferView<T>::copy_from(const ByteBufferView &source) const noexcept {
return source.copy_to(*this);
}

}// namespace luisa::compute
Loading

0 comments on commit c353e5d

Please sign in to comment.