Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions doc/releases/changelog-dev.md
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,9 @@

<h3>Improvements 🛠</h3>

* A new runtime transport layer for remote/local executors is introduced.
[(#3043)](https://github.com/PennyLaneAI/catalyst/pull/3043)

* A `BufferizableOpInterface` implementation is now added for `catalyst.launch_kernel` operation and it is now bufferizable.
[(#3024)](https://github.com/PennyLaneAI/catalyst/pull/3024)

Expand Down Expand Up @@ -469,6 +472,7 @@ JiaRung Jian,
Jacob Kitchen,
Korbinian Kottmann,
Christina Lee,
Joseph Lee,
Rylan Malarchick,
Mehrdad Malekmohammadi,
River McCubbin,
Expand Down
198 changes: 198 additions & 0 deletions runtime/include/Transport.hpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,198 @@
// Copyright 2026 Xanadu Quantum Technologies Inc.

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at

// http://www.apache.org/licenses/LICENSE-2.0

// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#pragma once
#include <cstddef>
#include <cstdint>
#include <string>

namespace catalyst::transport {

/**
* @brief Data-plane strategy: which engine issues the transfer.
*/
enum class DataPath : std::uint8_t {
CpuVerbs, // Plain ibverbs on CPU.
GpuEngine, // Gpu-initiated comms.
Other,
};

/**
* @brief Memory kind: selects the allocation and registration path.
*/
enum class MemKind : std::uint8_t {
CpuRam,
GpuHbm,
Ddr,
Other,
};

/**
* @brief Out-of-band connection parameters for bringing up a session.
*/
struct ConnectInfo {
std::string peer;
std::uint16_t oob_port;
};

/**
* @brief Locally allocated and registered memory region.
*/
struct MemRegion {
void *addr = nullptr;
std::uint64_t size = 0;
std::uint32_t lkey = 0;
std::uint32_t rkey = 0;
MemKind kind = MemKind::CpuRam;
};

/**
* @brief Handle to a peer's memory region, exchanged over the out-of-band channel.
*/
struct PeerRef {
std::uint32_t rkey = 0;
std::uint64_t remote_addr = 0;
std::uint64_t size = 0;
};

/**
* @brief Configuration for the data-movement channel a session uses.
*/
struct ChannelDesc {
DataPath data_path = DataPath::CpuVerbs;
};

/**
* @brief Stateful transport session shared by the controller and coprocessor roles.
*
* Methods must be called in this order:
* 1. connect - bring up QPs + the out-of-band channel
* 2. alloc_memory - register the region (needs the connected context)
Comment thread
josephleekl marked this conversation as resolved.
* 3. exchange_keys - swap region handles over the out-of-band channel
* 4. establish_channel - program the channel from the local + peer regions
* 5. (coprocessor) set_coprocessor_fn before start()
* 6. start / collect / stop
*/
class TransportSession {
public:
virtual ~TransportSession() = default;

/**
* @brief Bring up the connection (out-of-band handshake and QP transition to RTS).
*
* @param info Peer address and out-of-band port.
*
* @return `int`
*/
virtual int connect(const ConnectInfo &info) = 0;

/**
* @brief Allocate and register a memory region on the device.
*
* @param size Size of the region in bytes.
* @param kind Memory kind selecting the allocation and registration path.
* @param access Access flags for the registration.
*
* @return `MemRegion` The allocated and registered region.
*/
virtual MemRegion alloc_memory(std::size_t size, MemKind kind, std::uint32_t access) = 0;

/**
* @brief Advertise a local region and receive the peer's region over the out-of-band channel.
*
* @param local The local region to advertise.
*
* @return `PeerRef` The peer's advertised region.
*/
virtual PeerRef exchange_keys(const MemRegion &local) = 0;

/**
* @brief Program the data movement this session will run (single channel per session).
*
* @param desc Channel configuration.
* @param local The local memory region.
* @param peer The peer's memory region.
*/
virtual void establish_channel(const ChannelDesc &desc, const MemRegion &local,
const PeerRef &peer) = 0;

/**
* @brief Launch the engine (non-blocking; runs until stop()).
*/
virtual void start() = 0;

/**
* @brief Wait for a result and write it out.
*
* @param replies Output buffer to write the result into.
* @param bytes Capacity of the output buffer, in bytes.
*
* @return `int`
*/
virtual int collect(void *replies, std::uint64_t bytes) = 0;

/**
* @brief Stop the engine and join. Idempotent.
*/
virtual void stop() = 0;

/**
* @brief Last round-trip time, in nanoseconds (for testing purposes).
*
* @return `std::uint64_t`
*/
virtual std::uint64_t last_rtt_ns() const { return 0; }
};

/**
* @brief Controller role: writes messages out and receives replies.
*/
class ControllerSession : public TransportSession {
public:
// Build the work item in slot `work_item_idx` from in_bytes and out_bytes.
virtual void commit_work_item(std::uint32_t work_item_idx, std::uint64_t in_bytes,
std::uint64_t out_bytes) = 0;

// Fire one round using work item `work_item_idx` and whatever payload is currently in
// data_slot(). Pairs with a subsequent collect(). Returns 0 on success.
virtual int kick(std::uint32_t work_item_idx) = 0;

// Current round's outbound slot in the transport-owned ring.
virtual void *data_slot() = 0;
};

/**
* @brief Opaque function to run on the coprocessor. May include a persistent kernel on the GPU.
*/
using CoprocessorFn = std::size_t (*)(const void *in, std::size_t in_len, void *out,
std::size_t out_cap, void *ctx);

/**
* @brief Coprocessor role: receives messages, process, and returns replies.
*/
class CoprocessorSession : public TransportSession {
public:
/**
* @brief Bind the coprocessor function this session runs.
*
* Call before start(). `fn` is a local function pointer; `ctx` is passed
* back to `fn` on every invocation and may be null.
*
* @param fn The coprocessor function to run per received message.
* @param ctx Opaque context passed to `fn` on each invocation; may be null.
*/
virtual void set_coprocessor_fn(CoprocessorFn fn, void *ctx) = 0;
};

} // namespace catalyst::transport
Loading