Add transport dialect definitions

doc/releases/changelog-dev.md

@@ -15,6 +15,10 @@
 * A new runtime transport layer for remote/local executors is introduced.
   [(#3043)](https://github.com/PennyLaneAI/catalyst/pull/3043)
 
+* A new `Transport` MLIR dialect is added, providing typed ops for driving a transport session's
+  lifecycle at the IR level.
+  [(#3047)](https://github.com/PennyLaneAI/catalyst/pull/3047)
+
 * A `BufferizableOpInterface` implementation is now added for `catalyst.launch_kernel` operation and it is now bufferizable.
   [(#3024)](https://github.com/PennyLaneAI/catalyst/pull/3024)
 

mlir/include/CMakeLists.txt

@@ -11,4 +11,5 @@ add_subdirectory(QecPhysical)
 add_subdirectory(Quantum)
 add_subdirectory(QRef)
 add_subdirectory(RTIO)
+add_subdirectory(Transport)
 add_subdirectory(Test)

mlir/include/Transport/CMakeLists.txt

@@ -0,0 +1 @@
+add_subdirectory(IR)

mlir/include/Transport/IR/CMakeLists.txt

@@ -0,0 +1,8 @@
+add_mlir_dialect(TransportOps transport)
+add_mlir_doc(TransportDialect TransportDialect Transport/ -gen-dialect-doc)
+add_mlir_doc(TransportOps TransportOps Transport/ -gen-op-doc)
+
+set(LLVM_TARGET_DEFINITIONS TransportOps.td)
+mlir_tablegen(TransportEnums.h.inc -gen-enum-decls)
+mlir_tablegen(TransportEnums.cpp.inc -gen-enum-defs)
+add_public_tablegen_target(MLIRTransportEnumsIncGen)

mlir/include/Transport/IR/TransportDialect.h

@@ -0,0 +1,25 @@
+// 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 "mlir/Bytecode/BytecodeOpInterface.h"
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/OpDefinition.h"
+
+#include "Transport/IR/TransportEnums.h.inc"
+#include "Transport/IR/TransportOpsDialect.h.inc"
+
+#define GET_TYPEDEF_CLASSES
+#include "Transport/IR/TransportOpsTypes.h.inc"

mlir/include/Transport/IR/TransportDialect.td

@@ -0,0 +1,104 @@
+// 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.
+
+#ifndef TRANSPORT_DIALECT
+#define TRANSPORT_DIALECT
+
+include "mlir/IR/OpBase.td"
+include "mlir/IR/DialectBase.td"
+include "mlir/IR/AttrTypeBase.td"
+include "mlir/IR/EnumAttr.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+
+//===----------------------------------------------------------------------===//
+// Transport dialect definition.
+//===----------------------------------------------------------------------===//
+
+def Transport_Dialect : Dialect {
+    let summary = "Typed ops for setting up and driving a transport session.";
+    let description = [{
+        The transport dialect models a connection-oriented data-movement session
+        between two endpoints: creating a session, bringing up the connection,
+        exchanging memory handles, establishing a data path, and running rounds
+        of request/reply traffic until teardown.
+    }];
+
+    let name = "transport";
+    let cppNamespace = "::catalyst::transport";
+    let useDefaultTypePrinterParser = 1;
+    let usePropertiesForAttributes = 1;
+}
+
+//===----------------------------------------------------------------------===//
+// Enums.
+//===----------------------------------------------------------------------===//
+
+def Transport_Role : I32EnumAttr<"Role", "transport session role", [
+        I32EnumAttrCase<"Controller", 0, "controller">,
+        I32EnumAttrCase<"Coprocessor", 1, "coprocessor">
+    ]> {
+    let cppNamespace = "::catalyst::transport";
+}
+
+def Transport_DataPath : I32EnumAttr<"DataPath", "transport data-movement path", [
+        I32EnumAttrCase<"CpuVerbs", 0, "cpu_verbs">,
+        I32EnumAttrCase<"GpuEngine", 1, "gpu_engine">,
+        I32EnumAttrCase<"Other", 2, "other">
+    ]> {
+    let cppNamespace = "::catalyst::transport";
+}
+
+//===----------------------------------------------------------------------===//
+// Types.
+//===----------------------------------------------------------------------===//
+
+class Transport_Type<string name, string typeMnemonic, list<Trait> traits = []>
+    : TypeDef<Transport_Dialect, name, traits> {
+    let mnemonic = typeMnemonic;
+}
+
+// Opaque session handle, parameterized by role. Lowers to !llvm.ptr; the role is
+// compile-time only and drives op verification + the create factory selection.
+def Transport_SessionType : Transport_Type<"Session", "session"> {
+    let summary = "Opaque transport session handle (CatalystTransportSession*), tagged with its role.";
+    let parameters = (ins EnumParameter<Transport_Role>:$role);
+    let assemblyFormat = "`<` $role `>`";
+}
+
+def Transport_TokenType : Transport_Type<"Token", "token"> {
+    let summary = "Handle to an in-flight async transport step, awaited with transport.barrier.";
+}
+
+// Role-constrained session-type constraints for role-specific ops.
+def Transport_ControllerSession : Type<
+    CPred<"::llvm::isa<::catalyst::transport::SessionType>($_self) && "
+          "::llvm::cast<::catalyst::transport::SessionType>($_self).getRole() == "
+          "::catalyst::transport::Role::Controller">,
+    "controller transport session">;
+
+def Transport_CoprocessorSession : Type<
+    CPred<"::llvm::isa<::catalyst::transport::SessionType>($_self) && "
+          "::llvm::cast<::catalyst::transport::SessionType>($_self).getRole() == "
+          "::catalyst::transport::Role::Coprocessor">,
+    "coprocessor transport session">;
+
+//===----------------------------------------------------------------------===//
+// Operation base.
+//===----------------------------------------------------------------------===//
+
+// All transport ops perform side-effecting I/O
+class Transport_Op<string mnemonic, list<Trait> traits = []> :
+    Op<Transport_Dialect, mnemonic, !listconcat(traits, [MemoryEffects<[MemWrite]>])>;
+
+#endif // TRANSPORT_DIALECT

mlir/include/Transport/IR/TransportOps.h

@@ -0,0 +1,26 @@
+// 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 "mlir/Bytecode/BytecodeOpInterface.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+
+#include "Transport/IR/TransportDialect.h"
+
+#define GET_OP_CLASSES
+#include "Transport/IR/TransportOps.h.inc"

mlir/include/Transport/IR/TransportOps.td

@@ -0,0 +1,156 @@
+// 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.
+
+#ifndef TRANSPORT_OPS
+#define TRANSPORT_OPS
+
+include "mlir/IR/OpBase.td"
+include "mlir/IR/BuiltinAttributes.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+include "Transport/IR/TransportDialect.td"
+
+//===----------------------------------------------------------------------===//
+// Session creation
+//===----------------------------------------------------------------------===//
+
+def Transport_CreateOp : Transport_Op<"create"> {
+    let summary = "Create a session from a backend plugin .so with a certain role";
+    let description = [{
+        Loads the backend `.so` and builds a session. The result type's role
+        (`!transport.session<controller|coprocessor>`) selects which factory the
+        runtime looks up and constrains the role-specific ops downstream.
+    }];
+    let arguments = (ins StrAttr:$backend_lib, StrAttr:$config);
+    let results = (outs Transport_SessionType:$session);
+    let assemblyFormat = "attr-dict `->` type($session)";
+}
+
+//===----------------------------------------------------------------------===//
+// Bring-up ops
+//===----------------------------------------------------------------------===//
+
+def Transport_ConnectOp : Transport_Op<"connect"> {
+    let summary = "Bring up the connection to the peer (blocking).";
+    let arguments = (ins Transport_SessionType:$session, StrAttr:$peer, I16Attr:$oob_port);
+    let assemblyFormat = "$session attr-dict `:` type($session)";
+}
+
+def Transport_ConnectAsyncOp : Transport_Op<"connect_async"> {
+    let summary = "connect() on a worker; await with transport.barrier.";
+    let arguments = (ins Transport_SessionType:$session, StrAttr:$peer, I16Attr:$oob_port);
+    let results = (outs Transport_TokenType:$token);
+    let assemblyFormat = "$session attr-dict `:` type($session) `->` type($token)";
+}
+
+def Transport_ExchangeKeysOp : Transport_Op<"exchange_keys"> {
+    let summary = "Exchange region handles with the peer (blocking); result kept in the session.";
+    let arguments = (ins Transport_SessionType:$session);
+    let assemblyFormat = "$session attr-dict `:` type($session)";
+}
+
+def Transport_ExchangeKeysAsyncOp : Transport_Op<"exchange_keys_async"> {
+    let summary = "exchange_keys() on a worker; await with transport.barrier.";
+    let arguments = (ins Transport_SessionType:$session);
+    let results = (outs Transport_TokenType:$token);
+    let assemblyFormat = "$session attr-dict `:` type($session) `->` type($token)";
+}
+
+def Transport_BarrierOp : Transport_Op<"barrier"> {
+    let summary = "Await an async step (connect_async / exchange_keys_async).";
+    let arguments = (ins Transport_TokenType:$token);
+    let assemblyFormat = "$token attr-dict `:` type($token)";
+}
+
+def Transport_EstablishChannelOp : Transport_Op<"establish_channel"> {
+    let summary = "Set up the data channel used to transfer payloads each round.";
+    let description = [{
+        Arms the data channel for the given `data_path`, using this side's
+        registered memory region together with the peer's region that
+        `exchange_keys` learned earlier (both stored in the session). After this
+        the channel is ready for `kick`/`collect`.
+    }];
+    let arguments = (ins Transport_SessionType:$session, Transport_DataPath:$data_path);
+    let assemblyFormat = "$session $data_path attr-dict `:` type($session)";
+}
+
+//===----------------------------------------------------------------------===//
+// Controller-only ops
+//===----------------------------------------------------------------------===//
+
+def Transport_CommitWorkItemOp : Transport_Op<"commit_work_item"> {
+    let summary = "Build a work item (I/O sizes) in a slot before kicking rounds.";
+    let arguments = (ins Transport_ControllerSession:$session, I32Attr:$work_item_idx,
+                         I64Attr:$in_bytes, I64Attr:$out_bytes);
+    let assemblyFormat = "$session attr-dict `:` type($session)";
+}
+
+def Transport_KickOp : Transport_Op<"kick"> {
+    let summary = "Write the payload into the outbound slot and fire one round.";
+    let arguments = (ins Transport_ControllerSession:$session, I64:$payload, I32Attr:$work_item_idx);
+    let assemblyFormat = "$session `,` $payload attr-dict `:` type($session) `,` type($payload)";
+}
+
+//===----------------------------------------------------------------------===//
+// Coprocessor-only: bind the coprocessor function (requires a coprocessor session)
+//===----------------------------------------------------------------------===//
+
+def Transport_SetCoprocessorFnOp : Transport_Op<"set_coprocessor_fn"> {
+    let summary = "Bind the built-in coprocessor function (echo / on-device kernel).";
+    let arguments = (ins Transport_CoprocessorSession:$session);
+    let assemblyFormat = "$session attr-dict `:` type($session)";
+}
+
+//===----------------------------------------------------------------------===//
+// Run / collect / teardown
+//===----------------------------------------------------------------------===//
+
+def Transport_StartOp : Transport_Op<"start"> {
+    let summary = "Start the session (non-blocking; runs until stop()).";
+    let arguments = (ins Transport_SessionType:$session);
+    let assemblyFormat = "$session attr-dict `:` type($session)";
+}
+
+def Transport_CollectOp : Transport_Op<"collect"> {
+    let summary = "Wait for this round's reply and return it as a value.";
+    let arguments = (ins Transport_SessionType:$session, I64Attr:$bytes);
+    let results = (outs I64:$result);
+    let assemblyFormat = "$session attr-dict `:` type($session) `->` type($result)";
+}
+
+def Transport_LastRttNsOp : Transport_Op<"last_rtt_ns"> {
+    let summary = "Last round-trip time in nanoseconds.";
+    let arguments = (ins Transport_SessionType:$session);
+    let results = (outs I64:$rtt_ns);
+    let assemblyFormat = "$session attr-dict `:` type($session) `->` type($rtt_ns)";
+}
+
+def Transport_StopOp : Transport_Op<"stop"> {
+    let summary = "Stop the session. Idempotent.";
+    let arguments = (ins Transport_SessionType:$session);
+    let assemblyFormat = "$session attr-dict `:` type($session)";
+}
+
+def Transport_CloseOp : Transport_Op<"close"> {
+    let summary = "Close the transport (releases the channel).";
+    let arguments = (ins Transport_SessionType:$session);
+    let assemblyFormat = "$session attr-dict `:` type($session)";
+}
+
+def Transport_DestroyOp : Transport_Op<"destroy"> {
+    let summary = "Destroy the session and free it.";
+    let arguments = (ins Transport_SessionType:$session);
+    let assemblyFormat = "$session attr-dict `:` type($session)";
+}
+
+#endif // TRANSPORT_OPS

mlir/lib/CMakeLists.txt

@@ -13,4 +13,5 @@ add_subdirectory(QecPhysical)
 add_subdirectory(QRef)
 add_subdirectory(Quantum)
 add_subdirectory(RTIO)
+add_subdirectory(Transport)
 add_subdirectory(Test)

mlir/lib/Driver/CMakeLists.txt

@@ -76,6 +76,7 @@ set(LIBS
     ion-transforms
     MLIRRTIO
     rtio-transforms
+    MLIRTransport
     MLIRCatalystTest
     ${ENZYME_LIB}
     fmt::fmt

mlir/lib/Driver/CompilerDriver.cpp

@@ -89,6 +89,8 @@
 
 #include "RegisterAllPasses.h"
 
+#include "Transport/IR/TransportDialect.h"
+
 using namespace mlir;
 using namespace catalyst;
 using namespace catalyst::driver;
@@ -181,6 +183,7 @@ void registerAllCatalystDialects(DialectRegistry &registry)
     registry.insert<mbqc::MBQCDialect>();
     registry.insert<ion::IonDialect>();
     registry.insert<rtio::RTIODialect>();
+    registry.insert<transport::TransportDialect>();
     registry.insert<gradient::GradientDialect>();
     registry.insert<mitigation::MitigationDialect>();
     registry.insert<pauli_frame::PauliFrameDialect>();

mlir/lib/Transport/CMakeLists.txt

@@ -0,0 +1 @@
+add_subdirectory(IR)

mlir/lib/Transport/IR/CMakeLists.txt

@@ -0,0 +1,14 @@
+add_mlir_library(MLIRTransport
+    TransportDialect.cpp
+    TransportOps.cpp
+
+    ADDITIONAL_HEADER_DIRS
+    ${PROJECT_SOURCE_DIR}/include/Transport
+
+    DEPENDS
+    MLIRTransportOpsIncGen
+    MLIRTransportEnumsIncGen
+
+    LINK_LIBS PRIVATE
+    MLIRLLVMDialect
+)