request_data.h

#pragma once
 
#include <memory>
#include <optional>
#include <string>
#include <variant>
#include <vector>
 
#include <ucp/api/ucp.h>
 
#include <ucxx/tag_probe.h>
#include <ucxx/typedefs.h>
 
namespace ucxx {
 
class Buffer;
 
namespace data {

class AmSend {
 public:
  const void* const _buffer{nullptr};      
  const size_t _length{0};                 
  const std::vector<ucp_dt_iov_t> _iov{};  
  const size_t _count{0};                  
  const uint32_t _flags{UCP_AM_SEND_FLAG_REPLY};              
  const ucp_datatype_t _datatype{ucp_dt_make_contig(1)};      
  const ucs_memory_type_t _memoryType{UCS_MEMORY_TYPE_HOST};  
  const AmSendMemoryTypePolicy _memoryTypePolicy{
    AmSendMemoryTypePolicy::FallbackToHost};  
  const std::optional<AmReceiverCallbackInfo> _receiverCallbackInfo{
    std::nullopt};  
  const std::vector<std::byte> _userHeader{};  

  explicit AmSend(const decltype(_buffer) buffer,
                  const decltype(_length) length,
                  const AmSendParams& params = AmSendParams{});

  explicit AmSend(decltype(_iov) iov, const AmSendParams& params = AmSendParams{});
 
  AmSend() = delete;
};

class AmReceive {
 public:
  std::shared_ptr<::ucxx::Buffer> _buffer{nullptr};  
  std::vector<std::byte> _userHeader{};              

  AmReceive();
};

class EndpointClose {
 public:
  const bool _force{false};  
  explicit EndpointClose(const decltype(_force) force);
 
  EndpointClose() = delete;
};

class Flush {
 public:
  Flush();
};

class MemPut {
 public:
  const void* const _buffer{nullptr};  
  const size_t _length{0};             
  const uint64_t _remoteAddr{0};       
  const ucp_rkey_h _rkey{};  

  explicit MemPut(const decltype(_buffer) buffer,
                  const decltype(_length) length,
                  const decltype(_remoteAddr) remoteAddr,
                  const decltype(_rkey) rkey);
 
  MemPut() = delete;
};

class MemGet {
 public:
  void* _buffer{nullptr};         
  const size_t _length{0};        
  const uint64_t _remoteAddr{0};  
  const ucp_rkey_h _rkey{};       

  explicit MemGet(decltype(_buffer) buffer,
                  const decltype(_length) length,
                  const decltype(_remoteAddr) remoteAddr,
                  const decltype(_rkey) rkey);
 
  MemGet() = delete;
};

class StreamSend {
 public:
  const void* const _buffer{nullptr};  
  const size_t _length{0};             

  explicit StreamSend(const decltype(_buffer) buffer, const decltype(_length) length);
 
  StreamSend() = delete;
};

class StreamReceive {
 public:
  void* _buffer{nullptr};     
  const size_t _length{0};    
  size_t _lengthReceived{0};  

  explicit StreamReceive(decltype(_buffer) buffer, const decltype(_length) length);
 
  StreamReceive() = delete;
};

class TagSend {
 public:
  const void* const _buffer{nullptr};  
  const size_t _length{0};             
  const ::ucxx::Tag _tag{0};           

  explicit TagSend(const decltype(_buffer) buffer,
                   const decltype(_length) length,
                   const decltype(_tag) tag);
 
  TagSend() = delete;
};

class TagReceive {
 public:
  void* _buffer{nullptr};             
  const size_t _length{0};            
  const ::ucxx::Tag _tag{0};          
  const ::ucxx::TagMask _tagMask{0};  

  explicit TagReceive(decltype(_buffer) buffer,
                      const decltype(_length) length,
                      const decltype(_tag) tag,
                      const decltype(_tagMask) tagMask);
 
  TagReceive() = delete;
};

class TagReceiveWithHandle {
 public:
  void* _buffer{nullptr};   
  const size_t _length{0};  
  const std::shared_ptr<TagProbeInfo> _probeInfo{
    nullptr};  

  explicit TagReceiveWithHandle(decltype(_buffer) buffer,
                                decltype(_length) length,
                                std::shared_ptr<TagProbeInfo> probeInfo);
 
  TagReceiveWithHandle() = delete;
};

class TagMultiSend {
 public:
  const std::vector<const void*> _buffer{};  
  const std::vector<size_t> _length{};       
  const std::vector<int> _isCUDA{};  
  const ::ucxx::Tag _tag{0};         

  explicit TagMultiSend(const decltype(_buffer)& buffer,
                        const decltype(_length)& length,
                        const decltype(_isCUDA)& isCUDA,
                        const decltype(_tag) tag);
 
  TagMultiSend() = delete;
};

class TagMultiReceive {
 public:
  const ::ucxx::Tag _tag{0};          
  const ::ucxx::TagMask _tagMask{0};  

  explicit TagMultiReceive(const decltype(_tag) tag, const decltype(_tagMask) tagMask);
 
  TagMultiReceive() = delete;
};
 
using RequestData = std::variant<std::monostate,
                                 AmSend,
                                 AmReceive,
                                 EndpointClose,
                                 Flush,
                                 MemPut,
                                 MemGet,
                                 StreamSend,
                                 StreamReceive,
                                 TagSend,
                                 TagReceive,
                                 TagReceiveWithHandle,
                                 TagMultiSend,
                                 TagMultiReceive>;
 
template <class... Ts>
struct dispatch : Ts... {
  using Ts::operator()...;
};
template <class... Ts>
dispatch(Ts...) -> dispatch<Ts...>;
 
template <class T>
RequestData getRequestData(T t)
{
  return std::visit([](auto arg) -> RequestData { return arg; }, t);
}
 
}  // namespace data
 
}  // namespace ucxx