feat: rework pinned memory management (PROOF-928) (#259)

* rework pinned memory management

* rework concepts

* drop print statement

Author: rnburn

sxt/base/device/pinned_buffer.cc

@@ -1,6 +1,6 @@
 /** Proofs GPU - Space and Time's cryptographic proof algorithms on the CPU and GPU.
  *
- * Copyright 2024-present Space and Time Labs, Inc.
+ * Copyright 2025-present Space and Time Labs, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,17 +17,20 @@
 #include "sxt/base/device/pinned_buffer.h"
 
 #include "sxt/base/device/pinned_buffer_pool.h"
+#include "sxt/base/error/assert.h"
 
 namespace sxt::basdv {
 //--------------------------------------------------------------------------------------------------
-// consructor
+// constructor
 //--------------------------------------------------------------------------------------------------
-pinned_buffer::pinned_buffer() noexcept : handle_{get_pinned_buffer_pool()->acquire_handle()} {}
-
-pinned_buffer::pinned_buffer(pinned_buffer&& ptr) noexcept : handle_{ptr.handle_} {
-  ptr.handle_ = nullptr;
+pinned_buffer::pinned_buffer(size_t size) noexcept
+    : handle_{get_pinned_buffer_pool()->acquire_handle()}, size_{size} {
+  SXT_RELEASE_ASSERT(size_ <= this->capacity());
 }
 
+pinned_buffer::pinned_buffer(pinned_buffer&& other) noexcept
+    : handle_{std::exchange(other.handle_, nullptr)}, size_{std::exchange(other.size_, 0)} {}
+
 //--------------------------------------------------------------------------------------------------
 // destructor
 //--------------------------------------------------------------------------------------------------
@@ -40,17 +43,54 @@ pinned_buffer::~pinned_buffer() noexcept {
 //--------------------------------------------------------------------------------------------------
 // operator=
 //--------------------------------------------------------------------------------------------------
-pinned_buffer& pinned_buffer::operator=(pinned_buffer&& ptr) noexcept {
-  if (handle_ != nullptr) {
-    get_pinned_buffer_pool()->release_handle(handle_);
-  }
-  handle_ = ptr.handle_;
-  ptr.handle_ = nullptr;
+pinned_buffer& pinned_buffer::operator=(pinned_buffer&& other) noexcept {
+  this->reset();
+  handle_ = std::exchange(other.handle_, nullptr);
+  size_ = std::exchange(other.size_, 0);
   return *this;
 }
 
 //--------------------------------------------------------------------------------------------------
-// size
+// capacity
 //--------------------------------------------------------------------------------------------------
-size_t pinned_buffer::size() noexcept { return pinned_buffer_size; }
+size_t pinned_buffer::capacity() noexcept { return pinned_buffer_size; }
+
+//--------------------------------------------------------------------------------------------------
+// resize
+//--------------------------------------------------------------------------------------------------
+void pinned_buffer::resize(size_t size) noexcept {
+  SXT_RELEASE_ASSERT(size <= this->capacity());
+  if (handle_ == nullptr) {
+    handle_ = get_pinned_buffer_pool()->acquire_handle();
+  }
+  size_ = size;
+}
+
+//--------------------------------------------------------------------------------------------------
+// fill
+//--------------------------------------------------------------------------------------------------
+basct::cspan<std::byte> pinned_buffer::fill_from_host(basct::cspan<std::byte> src) noexcept {
+  if (src.empty()) {
+    return src;
+  }
+  if (handle_ == nullptr) {
+    handle_ = get_pinned_buffer_pool()->acquire_handle();
+  }
+  auto n = std::min(src.size(), this->capacity() - size_);
+  std::copy_n(src.data(), n, static_cast<std::byte*>(handle_->ptr) + size_);
+  size_ += n;
+  return src.subspan(n);
+}
+
+//--------------------------------------------------------------------------------------------------
+// reset
+//--------------------------------------------------------------------------------------------------
+void pinned_buffer::reset() noexcept {
+  if (handle_ == nullptr) {
+    return;
+  }
+  get_pinned_buffer_pool()->release_handle(handle_);
+  handle_ = nullptr;
+  size_ = 0;
+}
 } // namespace sxt::basdv

sxt/base/device/pinned_buffer.h

@@ -1,6 +1,6 @@
 /** Proofs GPU - Space and Time's cryptographic proof algorithms on the CPU and GPU.
  *
- * Copyright 2024-present Space and Time Labs, Inc.
+ * Copyright 2025-present Space and Time Labs, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,6 +16,9 @@
  */
 #pragma once
 
+#include <cstddef>
+
+#include "sxt/base/container/span.h"
 #include "sxt/base/device/pinned_buffer_handle.h"
 
 namespace sxt::basdv {
@@ -24,26 +27,48 @@ namespace sxt::basdv {
 //--------------------------------------------------------------------------------------------------
 class pinned_buffer {
 public:
-  pinned_buffer() noexcept;
-  pinned_buffer(pinned_buffer&& ptr) noexcept;
+  pinned_buffer() noexcept = default;
+
+  explicit pinned_buffer(size_t size) noexcept;
+
   pinned_buffer(const pinned_buffer&) noexcept = delete;
+  pinned_buffer(pinned_buffer&& other) noexcept;
 
   ~pinned_buffer() noexcept;
 
-  pinned_buffer& operator=(pinned_buffer&& ptr) noexcept;
-  pinned_buffer& operator=(const pinned_buffer& ptr) noexcept = delete;
+  pinned_buffer& operator=(const pinned_buffer&) noexcept = delete;
+  pinned_buffer& operator=(pinned_buffer&& other) noexcept;
+
+  bool empty() const noexcept { return size_ == 0; }
+
+  bool full() const noexcept { return size_ == this->capacity(); }
+
+  size_t size() const noexcept { return size_; }
+
+  static size_t capacity() noexcept;
 
-  static size_t size() noexcept;
+  void* data() noexcept {
+    if (handle_ == nullptr) {
+      return nullptr;
+    }
+    return handle_->ptr;
+  }
 
-  void* data() noexcept { return handle_->ptr; }
+  const void* data() const noexcept {
+    if (handle_ == nullptr) {
+      return nullptr;
+    }
+    return handle_->ptr;
+  }
 
-  const void* data() const noexcept { return handle_->ptr; }
+  void resize(size_t size) noexcept;
 
-  operator void*() noexcept { return handle_->ptr; }
+  basct::cspan<std::byte> fill_from_host(basct::cspan<std::byte> src) noexcept;
 
-  operator const void*() const noexcept { return handle_->ptr; }
+  void reset() noexcept;
 
 private:
-  pinned_buffer_handle* handle_;
+  pinned_buffer_handle* handle_ = nullptr;
+  size_t size_ = 0;
 };
 } // namespace sxt::basdv

sxt/base/device/pinned_buffer.t.cc

@@ -1,6 +1,6 @@
 /** Proofs GPU - Space and Time's cryptographic proof algorithms on the CPU and GPU.
  *
- * Copyright 2024-present Space and Time Labs, Inc.
+ * Copyright 2025-present Space and Time Labs, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,40 +16,64 @@
  */
 #include "sxt/base/device/pinned_buffer.h"
 
-#include "sxt/base/device/pinned_buffer_pool.h"
 #include "sxt/base/test/unit_test.h"
 
 using namespace sxt;
 using namespace sxt::basdv;
 
-TEST_CASE("we can manage pinned buffers") {
-  auto num_buffers = 5u;
-  auto pool = get_pinned_buffer_pool(num_buffers);
-
-  SECTION("we can acquire and release a pinned buffer") {
-    {
-      pinned_buffer buf;
-      REQUIRE(pool->size() == num_buffers - 1);
-      *reinterpret_cast<char*>(buf.data()) = 1u;
-      *(reinterpret_cast<char*>(buf.data()) + buf.size() - 1) = 2u;
-    }
-    REQUIRE(pool->size() == num_buffers);
+TEST_CASE("we can manage a buffer of pinned memory") {
+  SECTION("we can construct and deconstruct a buffer") {
+    pinned_buffer buf;
+    REQUIRE(buf.size() == 0);
+    REQUIRE(buf.empty());
+  }
+
+  SECTION("we can add a single byte to a buffer") {
+    pinned_buffer buf;
+    std::vector<std::byte> data = {std::byte{123}};
+    auto rest = buf.fill_from_host(data);
+    REQUIRE(rest.empty());
+    REQUIRE(buf.size() == 1);
+    REQUIRE(*static_cast<std::byte*>(buf.data()) == data[0]);
+  }
+
+  SECTION("we can reset a buffer") {
+    pinned_buffer buf;
+    std::vector<std::byte> data = {std::byte{123}};
+    buf.fill_from_host(data);
+    buf.reset();
+    REQUIRE(buf.empty());
   }
 
   SECTION("we can move construct a buffer") {
-    pinned_buffer buf1;
-    auto ptr = buf1.data();
-    pinned_buffer buf{std::move(buf1)};
-    REQUIRE(buf.data() == ptr);
-    REQUIRE(pool->size() == num_buffers - 1);
+    pinned_buffer buf;
+    std::vector<std::byte> data = {static_cast<std::byte>(123)};
+    buf.fill_from_host(data);
+    pinned_buffer buf_p{std::move(buf)};
+    REQUIRE(buf.empty());
+    REQUIRE(buf_p.size() == 1);
+    REQUIRE(*static_cast<std::byte*>(buf_p.data()) == data[0]);
+  }
+
+  SECTION("we can move assign a buffer") {
+    pinned_buffer buf;
+    std::vector<std::byte> data = {std::byte{123}};
+    buf.fill_from_host(data);
+
+    pinned_buffer buf_p;
+    data[0] = std::byte{3};
+    buf_p.fill_from_host(data);
+    buf_p = std::move(buf);
+    REQUIRE(buf.empty());
+    REQUIRE(*static_cast<std::byte*>(buf_p.data()) == std::byte{123});
   }
 
-  SECTION("we can move-assign a buffer") {
-    pinned_buffer buf1;
-    auto ptr = buf1.data();
+  SECTION("we can fill a buffer") {
     pinned_buffer buf;
-    buf = std::move(buf1);
-    REQUIRE(buf.data() == ptr);
-    REQUIRE(pool->size() == num_buffers - 1);
+    std::vector<std::byte> data(buf.capacity() + 1, std::byte{123});
+    auto rest = buf.fill_from_host(data);
+    REQUIRE(rest.size() == 1);
+    REQUIRE(buf.size() == buf.capacity());
+    REQUIRE(buf.full());
   }
 }

sxt/execution/device/BUILD

@@ -11,6 +11,31 @@ sxt_cc_component(
     ],
 )
 
+sxt_cc_component(
+    name = "to_device_copier",
+    impl_deps = [
+        ":synchronization",
+        "//sxt/base/device:memory_utility",
+        "//sxt/base/functional:function_ref",
+        "//sxt/execution/async:coroutine",
+    ],
+    test_deps = [
+        "//sxt/base/device:pinned_buffer_pool",
+        "//sxt/base/device:synchronization",
+        "//sxt/base/test:unit_test",
+        "//sxt/execution/schedule:scheduler",
+        "//sxt/memory/management:managed_array",
+        "//sxt/memory/resource:managed_device_resource",
+    ],
+    deps = [
+        "//sxt/base/container:span",
+        "//sxt/base/device:pinned_buffer",
+        "//sxt/base/device:stream",
+        "//sxt/base/error:assert",
+        "//sxt/execution/async:future",
+    ],
+)
+
 sxt_cc_component(
     name = "device_viewable",
     test_deps = [
@@ -192,10 +217,7 @@ sxt_cc_component(
 
 sxt_cc_component(
     name = "generate",
-    impl_deps = [
-    ],
     test_deps = [
-        "//sxt/base/device:pinned_buffer",
         "//sxt/base/device:stream",
         "//sxt/base/device:synchronization",
         "//sxt/base/test:unit_test",
@@ -217,7 +239,7 @@ sxt_cc_component(
 sxt_cc_component(
     name = "copy",
     impl_deps = [
-        ":generate",
+        ":to_device_copier",
         "//sxt/base/device:memory_utility",
         "//sxt/base/device:stream",
         "//sxt/execution/async:coroutine",