Active Message APIs: support scatter-gather I/O and user-defined header data (Part 1: C++ APIs)

cpp/include/ucxx/endpoint.h

@@ -376,6 +376,52 @@ class Endpoint : public Component {
     RequestCallbackUserFunction callbackFunction                     = nullptr,
     RequestCallbackUserData callbackData                             = nullptr);
 
+  /**
+   * @brief Enqueue an active message send operation with explicit policy parameters.
+   *
+   * This overload extends `amSend()` with explicit UCX datatype/flags controls and receive
+   * allocation policy metadata while keeping callback behavior identical to the legacy API.
+   *
+   * @param[in] buffer              a raw pointer to the data to be sent.
+   * @param[in] length              the size in bytes of the message to be sent.
+   * @param[in] params              active message send parameters.
+   * @param[in] enablePythonFuture  whether a python future should be created and
+   *                                subsequently notified.
+   * @param[in] callbackFunction    user-defined callback function to call upon completion.
+   * @param[in] callbackData        user-defined data to pass to the `callbackFunction`.
+   *
+   * @returns Request to be subsequently checked for the completion and its state.
+   */
+  [[nodiscard]] std::shared_ptr<Request> amSend(
+    const void* const buffer,
+    const size_t length,
+    const AmSendParams& params,
+    const bool enablePythonFuture                = false,
+    RequestCallbackUserFunction callbackFunction = nullptr,
+    RequestCallbackUserData callbackData         = nullptr);
+
+  /**
+   * @brief Enqueue an active message send operation with IOV datatype.
+   *
+   * This overload submits `UCP_DATATYPE_IOV` active message sends.
+   *
+   * @param[in] iov                 vector of IOV segments to be sent.
+   * @param[in] params              active message send parameters. Datatype must be
+   *                                `UCP_DATATYPE_IOV`.
+   * @param[in] enablePythonFuture  whether a python future should be created and
+   *                                subsequently notified.
+   * @param[in] callbackFunction    user-defined callback function to call upon completion.
+   * @param[in] callbackData        user-defined data to pass to the `callbackFunction`.
+   *
+   * @returns Request to be subsequently checked for the completion and its state.
+   */
+  [[nodiscard]] std::shared_ptr<Request> amSend(
+    std::vector<ucp_dt_iov_t> iov,
+    const AmSendParams& params,
+    const bool enablePythonFuture                = false,
+    RequestCallbackUserFunction callbackFunction = nullptr,
+    RequestCallbackUserData callbackData         = nullptr);
+
   /**
    * @brief Enqueue an active message receive operation.
    *

cpp/include/ucxx/internal/request_am.h

@@ -62,7 +62,8 @@ class RecvAmMessage {
                 ucp_ep_h ep,
                 std::shared_ptr<RequestAm> request,
                 std::shared_ptr<Buffer> buffer,
-                AmReceiverCallbackType receiverCallback = AmReceiverCallbackType());
+                AmReceiverCallbackType receiverCallback = AmReceiverCallbackType(),
+                std::string userHeader = {});
 
   /**
    * @brief Set the UCP request.

cpp/include/ucxx/request.h

@@ -224,6 +224,17 @@ class Request : public Component {
    * @return The received buffer (if applicable) or `nullptr`.
    */
   [[nodiscard]] virtual std::shared_ptr<Buffer> getRecvBuffer();
+
+  /**
+   * @brief Get the received user header.
+   *
+   * This method is used to get the user-defined header bytes for applicable derived classes
+   * (e.g., `RequestAm` receive operations), in all other cases this will return an empty
+   * string.
+   *
+   * @return The received user header (if applicable) or an empty string.
+   */
+  [[nodiscard]] virtual std::string getRecvHeader();
 };
 
 }  // namespace ucxx

cpp/include/ucxx/request_am.h

@@ -161,6 +161,8 @@ class RequestAm : public Request {
                                                  const ucp_am_recv_param_t* param);
 
   [[nodiscard]] std::shared_ptr<Buffer> getRecvBuffer() override;
+
+  [[nodiscard]] std::string getRecvHeader() override;
 };
 
 }  // namespace ucxx

cpp/include/ucxx/request_data.h

@@ -29,10 +29,18 @@ namespace data {
 class AmSend {
  public:
   const void* const _buffer{nullptr};  ///< The raw pointer where data to be sent is stored.
-  const size_t _length{0};             ///< The length of the message.
+  const size_t _length{0};             ///< Message length in bytes (contiguous datatype only).
+  const std::vector<ucp_dt_iov_t> _iov{};  ///< Segments for IOV datatype.
+  const size_t _count{0};                   ///< Count passed to `ucp_am_send_nbx`: byte count
+                                            ///< for contiguous, number of IOV segments for IOV.
+  const uint32_t _flags{UCP_AM_SEND_FLAG_REPLY};  ///< UCP AM send flags.
+  const ucp_datatype_t _datatype{ucp_dt_make_contig(1)};  ///< UCP datatype.
   const ucs_memory_type_t _memoryType{UCS_MEMORY_TYPE_HOST};  ///< Memory type used on the operation
+  const AmSendMemoryTypePolicy _memoryTypePolicy{
+    AmSendMemoryTypePolicy::FallbackToHost};  ///< Receiver allocation policy.
   const std::optional<AmReceiverCallbackInfo> _receiverCallbackInfo{
     std::nullopt};  ///< Owner name and unique identifier of the receiver callback.
+  const std::string _userHeader{};  ///< Opaque user-defined header (arbitrary bytes, not necessarily text).
 
   /**
    * @brief Constructor for Active Message-specific send data.
@@ -41,14 +49,21 @@ class AmSend {
    *
    * @param[in] buffer                  a raw pointer to the data to be sent.
    * @param[in] length                  the size in bytes of the message to be sent.
-   * @param[in] memoryType              the memory type of the buffer.
-   * @param[in] receiverCallbackInfo    the owner name and unique identifier of the receiver
-                                        callback.
+   * @param[in] params                  send parameters controlling datatype/flags/policies.
    */
   explicit AmSend(const decltype(_buffer) buffer,
                   const decltype(_length) length,
-                  const decltype(_memoryType) memoryType                     = UCS_MEMORY_TYPE_HOST,
-                  const decltype(_receiverCallbackInfo) receiverCallbackInfo = std::nullopt);
+                  const AmSendParams& params = AmSendParams{});
+
+  /**
+   * @brief Constructor for Active Message-specific send data using IOV datatype.
+   *
+   * Construct an object containing Active Message-specific send data for `UCP_DATATYPE_IOV`.
+   *
+   * @param[in] iov     vector of IOV segments to send.
+   * @param[in] params  send parameters controlling datatype/flags/policies.
+   */
+  explicit AmSend(decltype(_iov) iov, const AmSendParams& params = AmSendParams{});
 
   AmSend() = delete;
 };
@@ -62,6 +77,7 @@ class AmSend {
 class AmReceive {
  public:
   std::shared_ptr<::ucxx::Buffer> _buffer{nullptr};  ///< The AM received message buffer
+  std::string _userHeader{};  ///< User-defined header from the sender (arbitrary bytes, not necessarily text).
 
   /**
    * @brief Constructor for Active Message-specific receive data.

cpp/include/ucxx/typedefs.h

@@ -154,8 +154,8 @@ typedef uint64_t AmReceiverCallbackIdType;
  */
 class AmReceiverCallbackInfo {
  public:
-  const AmReceiverCallbackOwnerType owner;  ///< The owner name of the callback
-  const AmReceiverCallbackIdType id;        ///< The unique identifier of the callback
+  AmReceiverCallbackOwnerType owner;  ///< The owner name of the callback
+  AmReceiverCallbackIdType id;        ///< The unique identifier of the callback
 
   AmReceiverCallbackInfo() = delete;
 
@@ -168,6 +168,41 @@ class AmReceiverCallbackInfo {
   AmReceiverCallbackInfo(const AmReceiverCallbackOwnerType owner, AmReceiverCallbackIdType id);
 };
 
+/**
+ * @brief Policy used to allocate receive buffers for Active Messages.
+ *
+ * Active Message receive allocations can be strict (error if no allocator is registered for
+ * sender-provided memory type) or permissive (fallback to host allocation).
+ */
+enum class AmSendMemoryTypePolicy {
+  FallbackToHost = 0,  ///< If no allocator exists for memory type, fallback to host memory.
+  ErrorOnUnsupported,  ///< If no allocator exists for memory type, fail with unsupported error.
+};
+
+/**
+ * @brief Parameters controlling Active Message send behavior.
+ *
+ * This object is used by the extended Active Message API to expose UCX send knobs without
+ * breaking existing callers.
+ */
+struct AmSendParams {
+  uint32_t flags{UCP_AM_SEND_FLAG_REPLY};  ///< UCP AM send flags.
+  ucp_datatype_t datatype{ucp_dt_make_contig(1)};  ///< Datatype used by `ucp_am_send_nbx`.
+  ucs_memory_type_t memoryType{UCS_MEMORY_TYPE_HOST};  ///< Sender memory type hint.
+  AmSendMemoryTypePolicy memoryTypePolicy{
+    AmSendMemoryTypePolicy::FallbackToHost};  ///< Receiver allocation policy.
+  std::optional<AmReceiverCallbackInfo> receiverCallbackInfo{
+    std::nullopt};           ///< Optional receiver callback metadata.
+  std::string userHeader{};  ///< Opaque user-defined header (arbitrary bytes, not necessarily text).
+                             ///< This is serialized into the AM header parameter of
+                             ///< `ucp_am_send_nbx`, which is subject to transport-level size
+                             ///< limits. For TCP, the default segment size is ~8 KiB
+                             ///< (`UCX_TCP_TX_SEG_SIZE` / `UCX_TCP_RX_SEG_SIZE`). Headers that
+                             ///< exceed the transport limit will cause a fatal UCX error. Keep
+                             ///< user headers small (recommended < 4 KiB) or increase the
+                             ///< segment size environment variables as needed.
+};
+
 /**
  * @brief Serialized form of a remote key.
  *

cpp/src/endpoint.cpp

@@ -449,11 +449,45 @@ std::shared_ptr<Request> Endpoint::amSend(
   const bool enablePythonFuture,
   RequestCallbackUserFunction callbackFunction,
   RequestCallbackUserData callbackData)
+{
+  auto params                    = AmSendParams{};
+  params.memoryType              = memoryType;
+  params.receiverCallbackInfo    = receiverCallbackInfo;
+
+  return amSend(buffer,
+                length,
+                params,
+                enablePythonFuture,
+                callbackFunction,
+                callbackData);
+}
+
+std::shared_ptr<Request> Endpoint::amSend(const void* const buffer,
+                                          const size_t length,
+                                          const AmSendParams& params,
+                                          const bool enablePythonFuture,
+                                          RequestCallbackUserFunction callbackFunction,
+                                          RequestCallbackUserData callbackData)
+{
+  auto endpoint = std::dynamic_pointer_cast<Endpoint>(shared_from_this());
+  return registerInflightRequest(
+    createRequestAm(endpoint,
+                    data::AmSend(buffer, length, params),
+                    enablePythonFuture,
+                    callbackFunction,
+                    callbackData));
+}
+
+std::shared_ptr<Request> Endpoint::amSend(std::vector<ucp_dt_iov_t> iov,
+                                          const AmSendParams& params,
+                                          const bool enablePythonFuture,
+                                          RequestCallbackUserFunction callbackFunction,
+                                          RequestCallbackUserData callbackData)
 {
   auto endpoint = std::dynamic_pointer_cast<Endpoint>(shared_from_this());
   return registerInflightRequest(
     createRequestAm(endpoint,
-                    data::AmSend(buffer, length, memoryType, receiverCallbackInfo),
+                    data::AmSend(std::move(iov), params),
                     enablePythonFuture,
                     callbackFunction,
                     callbackData));

cpp/src/internal/request_am.cpp

@@ -18,11 +18,15 @@ RecvAmMessage::RecvAmMessage(internal::AmData* amData,
                              ucp_ep_h ep,
                              std::shared_ptr<RequestAm> request,
                              std::shared_ptr<Buffer> buffer,
-                             AmReceiverCallbackType receiverCallback)
+                             AmReceiverCallbackType receiverCallback,
+                             std::string userHeader)
   : _amData(amData), _ep(ep), _request(request)
 {
   std::visit(data::dispatch{
-               [this, buffer](data::AmReceive& amReceive) { amReceive._buffer = buffer; },
+               [this, buffer, &userHeader](data::AmReceive& amReceive) {
+                 amReceive._buffer     = buffer;
+                 amReceive._userHeader = std::move(userHeader);
+               },
                [](auto) { throw std::runtime_error("Unreachable"); },
              },
              _request->_requestData);

cpp/src/request.cpp

@@ -242,4 +242,6 @@ const std::string& Request::getOwnerString() const { return _ownerString; }
 
 std::shared_ptr<Buffer> Request::getRecvBuffer() { return nullptr; }
 
+std::string Request::getRecvHeader() { return {}; }
+
 }  // namespace ucxx