From 18c26840ba23d43fdcf02a107e3f7a033b8ab55a Mon Sep 17 00:00:00 2001 From: Mehrdad Malekmohammadi Date: Tue, 21 Jul 2026 07:10:23 -0400 Subject: [PATCH 1/5] Add transport dialect Assisted-by: Claude Opus 4.8 --- mlir/include/CMakeLists.txt | 1 + mlir/include/Transport/CMakeLists.txt | 1 + mlir/include/Transport/IR/CMakeLists.txt | 3 + mlir/include/Transport/IR/TransportDialect.h | 24 ++++ mlir/include/Transport/IR/TransportDialect.td | 67 ++++++++++ mlir/include/Transport/IR/TransportOps.h | 26 ++++ mlir/include/Transport/IR/TransportOps.td | 117 ++++++++++++++++++ mlir/lib/CMakeLists.txt | 1 + mlir/lib/Driver/CMakeLists.txt | 1 + mlir/lib/Driver/CompilerDriver.cpp | 3 + mlir/lib/Transport/CMakeLists.txt | 1 + mlir/lib/Transport/IR/CMakeLists.txt | 13 ++ mlir/lib/Transport/IR/TransportDialect.cpp | 50 ++++++++ mlir/lib/Transport/IR/TransportOps.cpp | 24 ++++ mlir/tools/quantum-opt/CMakeLists.txt | 1 + mlir/tools/quantum-opt/quantum-opt.cpp | 3 + 16 files changed, 336 insertions(+) create mode 100644 mlir/include/Transport/CMakeLists.txt create mode 100644 mlir/include/Transport/IR/CMakeLists.txt create mode 100644 mlir/include/Transport/IR/TransportDialect.h create mode 100644 mlir/include/Transport/IR/TransportDialect.td create mode 100644 mlir/include/Transport/IR/TransportOps.h create mode 100644 mlir/include/Transport/IR/TransportOps.td create mode 100644 mlir/lib/Transport/CMakeLists.txt create mode 100644 mlir/lib/Transport/IR/CMakeLists.txt create mode 100644 mlir/lib/Transport/IR/TransportDialect.cpp create mode 100644 mlir/lib/Transport/IR/TransportOps.cpp diff --git a/mlir/include/CMakeLists.txt b/mlir/include/CMakeLists.txt index a94e4af0f9..09bb4d088e 100644 --- a/mlir/include/CMakeLists.txt +++ b/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) diff --git a/mlir/include/Transport/CMakeLists.txt b/mlir/include/Transport/CMakeLists.txt new file mode 100644 index 0000000000..f33061b2d8 --- /dev/null +++ b/mlir/include/Transport/CMakeLists.txt @@ -0,0 +1 @@ +add_subdirectory(IR) diff --git a/mlir/include/Transport/IR/CMakeLists.txt b/mlir/include/Transport/IR/CMakeLists.txt new file mode 100644 index 0000000000..d1bc88dbed --- /dev/null +++ b/mlir/include/Transport/IR/CMakeLists.txt @@ -0,0 +1,3 @@ +add_mlir_dialect(TransportOps transport) +add_mlir_doc(TransportDialect TransportDialect Transport/ -gen-dialect-doc) +add_mlir_doc(TransportOps TransportOps Transport/ -gen-op-doc) diff --git a/mlir/include/Transport/IR/TransportDialect.h b/mlir/include/Transport/IR/TransportDialect.h new file mode 100644 index 0000000000..ebfdd9ba47 --- /dev/null +++ b/mlir/include/Transport/IR/TransportDialect.h @@ -0,0 +1,24 @@ +// 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/TransportOpsDialect.h.inc" + +#define GET_TYPEDEF_CLASSES +#include "Transport/IR/TransportOpsTypes.h.inc" diff --git a/mlir/include/Transport/IR/TransportDialect.td b/mlir/include/Transport/IR/TransportDialect.td new file mode 100644 index 0000000000..78ed32fb70 --- /dev/null +++ b/mlir/include/Transport/IR/TransportDialect.td @@ -0,0 +1,67 @@ +// 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" + +//===----------------------------------------------------------------------===// +// Transport dialect definition. +//===----------------------------------------------------------------------===// + +def Transport_Dialect : Dialect { + let summary = "Runtime transport-layer ops that lower to the __catalyst__transport__* CAPI."; + let description = [{ + A thin, typed MLIR representation of the Catalyst runtime transport CAPI + (runtime/include/TransportCAPI.h). One op per controller-side CAPI entry + point; `convert-transport-to-llvm` lowers each to an `llvm.call` on the + matching `__catalyst__transport__*` symbol. The coprocessor role is an + out-of-tree backend loaded by the runtime, so only the controller side is + emitted here. + }]; + + let name = "transport"; + let cppNamespace = "::catalyst::transport"; + let useDefaultTypePrinterParser = 1; + let usePropertiesForAttributes = 1; +} + +//===----------------------------------------------------------------------===// +// Transport dialect types. +//===----------------------------------------------------------------------===// + +class Transport_Type traits = []> + : TypeDef { + let mnemonic = typeMnemonic; +} + +def Transport_SessionType : Transport_Type<"Session", "session"> { + let summary = "Opaque transport controller session handle (CatalystTransportSession*)."; +} + +def Transport_PeerType : Transport_Type<"Peer", "peer"> { + let summary = "Opaque peer-region descriptor handle (CatalystTransportPeerRef*)."; +} + +//===----------------------------------------------------------------------===// +// Transport operation base. +//===----------------------------------------------------------------------===// + +class Transport_Op traits = []> : + Op; + +#endif // TRANSPORT_DIALECT diff --git a/mlir/include/Transport/IR/TransportOps.h b/mlir/include/Transport/IR/TransportOps.h new file mode 100644 index 0000000000..0cc3624b7b --- /dev/null +++ b/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" diff --git a/mlir/include/Transport/IR/TransportOps.td b/mlir/include/Transport/IR/TransportOps.td new file mode 100644 index 0000000000..ccd7b2308e --- /dev/null +++ b/mlir/include/Transport/IR/TransportOps.td @@ -0,0 +1,117 @@ +// 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" + +//===----------------------------------------------------------------------===// +// Bring-up ops +//===----------------------------------------------------------------------===// + +def Transport_ControllerCreateOp : Transport_Op<"controller_create"> { + let summary = "Create a controller session from a named backend plugin .so."; + let arguments = (ins StrAttr:$backend_lib, StrAttr:$config); + let results = (outs Transport_SessionType:$session); + let assemblyFormat = "attr-dict `->` type($session)"; +} + +def Transport_ConnectOp : Transport_Op<"connect"> { + let summary = "Bring up the connection to the peer."; + let arguments = (ins Transport_SessionType:$session, StrAttr:$peer, I16Attr:$oob_port); + let results = (outs I32:$status); + let assemblyFormat = "$session attr-dict `:` functional-type($session, $status)"; +} + +def Transport_ExchangeKeysOp : Transport_Op<"exchange_keys"> { + let summary = "Exchange local and peer region handles; yields the peer descriptor."; + let arguments = (ins Transport_SessionType:$session); + let results = (outs I32:$status, Transport_PeerType:$peer); + let assemblyFormat = "$session attr-dict `:` type($session) `->` type($peer)"; +} + +def Transport_EstablishChannelOp : Transport_Op<"establish_channel"> { + let summary = "Program the data-movement channel from the local + peer regions."; + let arguments = (ins Transport_SessionType:$session, Transport_PeerType:$peer, + I32Attr:$data_path); + let results = (outs I32:$status); + let assemblyFormat = "$session `,` $peer attr-dict `:` type($session) `,` type($peer)"; +} + +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_SessionType:$session, I32Attr:$work_item_idx, + I64Attr:$in_bytes, I64Attr:$out_bytes); + let results = (outs I32:$status); + let assemblyFormat = "$session attr-dict `:` type($session)"; +} + +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)"; +} + +//===----------------------------------------------------------------------===// +// Per-round ops +//===----------------------------------------------------------------------===// + +def Transport_KickOp : Transport_Op<"kick"> { + let summary = "Write the payload into the outbound slot and fire one round."; + let arguments = (ins Transport_SessionType:$session, I64:$payload, I32Attr:$work_item_idx); + let results = (outs I32:$status); + let assemblyFormat = "$session `,` $payload attr-dict `:` type($session) `,` type($payload)"; +} + +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)"; +} + +//===----------------------------------------------------------------------===// +// Teardown ops +//===----------------------------------------------------------------------===// + +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 diff --git a/mlir/lib/CMakeLists.txt b/mlir/lib/CMakeLists.txt index b4ff29b1a5..190316d1c9 100644 --- a/mlir/lib/CMakeLists.txt +++ b/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) diff --git a/mlir/lib/Driver/CMakeLists.txt b/mlir/lib/Driver/CMakeLists.txt index c371e2e8e0..cd55702b17 100644 --- a/mlir/lib/Driver/CMakeLists.txt +++ b/mlir/lib/Driver/CMakeLists.txt @@ -76,6 +76,7 @@ set(LIBS ion-transforms MLIRRTIO rtio-transforms + MLIRTransport MLIRCatalystTest ${ENZYME_LIB} fmt::fmt diff --git a/mlir/lib/Driver/CompilerDriver.cpp b/mlir/lib/Driver/CompilerDriver.cpp index f3aac8efc9..8b768e590d 100644 --- a/mlir/lib/Driver/CompilerDriver.cpp +++ b/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 ®istry) registry.insert(); registry.insert(); registry.insert(); + registry.insert(); registry.insert(); registry.insert(); registry.insert(); diff --git a/mlir/lib/Transport/CMakeLists.txt b/mlir/lib/Transport/CMakeLists.txt new file mode 100644 index 0000000000..f33061b2d8 --- /dev/null +++ b/mlir/lib/Transport/CMakeLists.txt @@ -0,0 +1 @@ +add_subdirectory(IR) diff --git a/mlir/lib/Transport/IR/CMakeLists.txt b/mlir/lib/Transport/IR/CMakeLists.txt new file mode 100644 index 0000000000..9554c24fe3 --- /dev/null +++ b/mlir/lib/Transport/IR/CMakeLists.txt @@ -0,0 +1,13 @@ +add_mlir_library(MLIRTransport + TransportDialect.cpp + TransportOps.cpp + + ADDITIONAL_HEADER_DIRS + ${PROJECT_SOURCE_DIR}/include/Transport + + DEPENDS + MLIRTransportOpsIncGen + + LINK_LIBS PRIVATE + MLIRLLVMDialect +) diff --git a/mlir/lib/Transport/IR/TransportDialect.cpp b/mlir/lib/Transport/IR/TransportDialect.cpp new file mode 100644 index 0000000000..96bb105661 --- /dev/null +++ b/mlir/lib/Transport/IR/TransportDialect.cpp @@ -0,0 +1,50 @@ +// 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. + +#include "Transport/IR/TransportDialect.h" + +#include "llvm/ADT/TypeSwitch.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/DialectImplementation.h" + +#include "Transport/IR/TransportOps.h" + +using namespace mlir; +using namespace catalyst::transport; + +//===----------------------------------------------------------------------===// +// Transport dialect definitions. +//===----------------------------------------------------------------------===// + +#include "Transport/IR/TransportOpsDialect.cpp.inc" + +//===----------------------------------------------------------------------===// +// Transport type definitions. +//===----------------------------------------------------------------------===// + +#define GET_TYPEDEF_CLASSES +#include "Transport/IR/TransportOpsTypes.cpp.inc" + +void TransportDialect::initialize() +{ + addTypes< +#define GET_TYPEDEF_LIST +#include "Transport/IR/TransportOpsTypes.cpp.inc" + >(); + + addOperations< +#define GET_OP_LIST +#include "Transport/IR/TransportOps.cpp.inc" + >(); +} diff --git a/mlir/lib/Transport/IR/TransportOps.cpp b/mlir/lib/Transport/IR/TransportOps.cpp new file mode 100644 index 0000000000..8d894ef69d --- /dev/null +++ b/mlir/lib/Transport/IR/TransportOps.cpp @@ -0,0 +1,24 @@ +// 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. + +#include "Transport/IR/TransportOps.h" + +#include "mlir/IR/Builders.h" +#include "mlir/IR/OpImplementation.h" + +using namespace mlir; +using namespace catalyst::transport; + +#define GET_OP_CLASSES +#include "Transport/IR/TransportOps.cpp.inc" diff --git a/mlir/tools/quantum-opt/CMakeLists.txt b/mlir/tools/quantum-opt/CMakeLists.txt index ef16ce934d..cf7a3b2053 100644 --- a/mlir/tools/quantum-opt/CMakeLists.txt +++ b/mlir/tools/quantum-opt/CMakeLists.txt @@ -33,6 +33,7 @@ set(LIBS ion-transforms MLIRRTIO rtio-transforms + MLIRTransport MLIRQecLogical MLIRQecPhysical MLIRCatalystTest diff --git a/mlir/tools/quantum-opt/quantum-opt.cpp b/mlir/tools/quantum-opt/quantum-opt.cpp index 5b9a279fd8..8e94976461 100644 --- a/mlir/tools/quantum-opt/quantum-opt.cpp +++ b/mlir/tools/quantum-opt/quantum-opt.cpp @@ -50,6 +50,8 @@ #include "RegisterAllPasses.h" +#include "Transport/IR/TransportDialect.h" + namespace test { void registerTestDialect(mlir::DialectRegistry &); } // namespace test @@ -77,6 +79,7 @@ int main(int argc, char **argv) registry.insert(); registry.insert(); registry.insert(); + registry.insert(); registry.insert(); registry.insert(); registry.insert(); From fd1790679a015257a3f3aabebb9b52f7609408b3 Mon Sep 17 00:00:00 2001 From: Mehrdad Malekmohammadi Date: Wed, 22 Jul 2026 11:15:03 -0400 Subject: [PATCH 2/5] update --- mlir/include/Transport/IR/CMakeLists.txt | 5 + mlir/include/Transport/IR/TransportDialect.h | 1 + mlir/include/Transport/IR/TransportDialect.td | 63 +++++++++--- mlir/include/Transport/IR/TransportOps.td | 97 +++++++++++++------ mlir/lib/Transport/IR/CMakeLists.txt | 1 + mlir/lib/Transport/IR/TransportDialect.cpp | 1 + 6 files changed, 126 insertions(+), 42 deletions(-) diff --git a/mlir/include/Transport/IR/CMakeLists.txt b/mlir/include/Transport/IR/CMakeLists.txt index d1bc88dbed..d2b9081dc0 100644 --- a/mlir/include/Transport/IR/CMakeLists.txt +++ b/mlir/include/Transport/IR/CMakeLists.txt @@ -1,3 +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) diff --git a/mlir/include/Transport/IR/TransportDialect.h b/mlir/include/Transport/IR/TransportDialect.h index ebfdd9ba47..57ab9b3dcb 100644 --- a/mlir/include/Transport/IR/TransportDialect.h +++ b/mlir/include/Transport/IR/TransportDialect.h @@ -19,6 +19,7 @@ #include "mlir/IR/OpDefinition.h" #include "Transport/IR/TransportOpsDialect.h.inc" +#include "Transport/IR/TransportEnums.h.inc" #define GET_TYPEDEF_CLASSES #include "Transport/IR/TransportOpsTypes.h.inc" diff --git a/mlir/include/Transport/IR/TransportDialect.td b/mlir/include/Transport/IR/TransportDialect.td index 78ed32fb70..118aae1734 100644 --- a/mlir/include/Transport/IR/TransportDialect.td +++ b/mlir/include/Transport/IR/TransportDialect.td @@ -18,20 +18,20 @@ 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 = "Runtime transport-layer ops that lower to the __catalyst__transport__* CAPI."; + let summary = "Typed ops for setting up and driving a transport session."; let description = [{ - A thin, typed MLIR representation of the Catalyst runtime transport CAPI - (runtime/include/TransportCAPI.h). One op per controller-side CAPI entry - point; `convert-transport-to-llvm` lowers each to an `llvm.call` on the - matching `__catalyst__transport__*` symbol. The coprocessor role is an - out-of-tree backend loaded by the runtime, so only the controller side is - emitted here. + 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"; @@ -41,7 +41,26 @@ def Transport_Dialect : Dialect { } //===----------------------------------------------------------------------===// -// Transport dialect types. +// 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 traits = []> @@ -49,19 +68,37 @@ class Transport_Type 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 controller session handle (CatalystTransportSession*)."; + let summary = "Opaque transport session handle (CatalystTransportSession*), tagged with its role."; + let parameters = (ins EnumParameter:$role); + let assemblyFormat = "`<` $role `>`"; } -def Transport_PeerType : Transport_Type<"Peer", "peer"> { - let summary = "Opaque peer-region descriptor handle (CatalystTransportPeerRef*)."; +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">; + //===----------------------------------------------------------------------===// -// Transport operation base. +// Operation base. //===----------------------------------------------------------------------===// +// All transport ops perform side-effecting I/O class Transport_Op traits = []> : - Op; + Op])>; #endif // TRANSPORT_DIALECT diff --git a/mlir/include/Transport/IR/TransportOps.td b/mlir/include/Transport/IR/TransportOps.td index ccd7b2308e..cc5406f8be 100644 --- a/mlir/include/Transport/IR/TransportOps.td +++ b/mlir/include/Transport/IR/TransportOps.td @@ -21,61 +21,104 @@ include "mlir/Interfaces/SideEffectInterfaces.td" include "Transport/IR/TransportDialect.td" //===----------------------------------------------------------------------===// -// Bring-up ops +// Session creation //===----------------------------------------------------------------------===// -def Transport_ControllerCreateOp : Transport_Op<"controller_create"> { - let summary = "Create a controller session from a named backend plugin .so."; +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`) 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."; + let summary = "Bring up the connection to the peer (blocking)."; let arguments = (ins Transport_SessionType:$session, StrAttr:$peer, I16Attr:$oob_port); - let results = (outs I32:$status); - let assemblyFormat = "$session attr-dict `:` functional-type($session, $status)"; + 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 local and peer region handles; yields the peer descriptor."; + 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 I32:$status, Transport_PeerType:$peer); - let assemblyFormat = "$session attr-dict `:` type($session) `->` type($peer)"; + 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 = "Program the data-movement channel from the local + peer regions."; - let arguments = (ins Transport_SessionType:$session, Transport_PeerType:$peer, - I32Attr:$data_path); - let results = (outs I32:$status); - let assemblyFormat = "$session `,` $peer attr-dict `:` type($session) `,` type($peer)"; + 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_SessionType:$session, I32Attr:$work_item_idx, + let arguments = (ins Transport_ControllerSession:$session, I32Attr:$work_item_idx, I64Attr:$in_bytes, I64Attr:$out_bytes); - let results = (outs I32:$status); let assemblyFormat = "$session attr-dict `:` type($session)"; } -def Transport_StartOp : Transport_Op<"start"> { - let summary = "Start the session (non-blocking; runs until stop())."; - let arguments = (ins Transport_SessionType:$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)"; } //===----------------------------------------------------------------------===// -// Per-round ops +// Run / collect / teardown //===----------------------------------------------------------------------===// -def Transport_KickOp : Transport_Op<"kick"> { - let summary = "Write the payload into the outbound slot and fire one round."; - let arguments = (ins Transport_SessionType:$session, I64:$payload, I32Attr:$work_item_idx); - let results = (outs I32:$status); - let assemblyFormat = "$session `,` $payload attr-dict `:` type($session) `,` type($payload)"; +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"> { @@ -92,10 +135,6 @@ def Transport_LastRttNsOp : Transport_Op<"last_rtt_ns"> { let assemblyFormat = "$session attr-dict `:` type($session) `->` type($rtt_ns)"; } -//===----------------------------------------------------------------------===// -// Teardown ops -//===----------------------------------------------------------------------===// - def Transport_StopOp : Transport_Op<"stop"> { let summary = "Stop the session. Idempotent."; let arguments = (ins Transport_SessionType:$session); diff --git a/mlir/lib/Transport/IR/CMakeLists.txt b/mlir/lib/Transport/IR/CMakeLists.txt index 9554c24fe3..1dbeed05e2 100644 --- a/mlir/lib/Transport/IR/CMakeLists.txt +++ b/mlir/lib/Transport/IR/CMakeLists.txt @@ -7,6 +7,7 @@ add_mlir_library(MLIRTransport DEPENDS MLIRTransportOpsIncGen + MLIRTransportEnumsIncGen LINK_LIBS PRIVATE MLIRLLVMDialect diff --git a/mlir/lib/Transport/IR/TransportDialect.cpp b/mlir/lib/Transport/IR/TransportDialect.cpp index 96bb105661..aea2c90196 100644 --- a/mlir/lib/Transport/IR/TransportDialect.cpp +++ b/mlir/lib/Transport/IR/TransportDialect.cpp @@ -28,6 +28,7 @@ using namespace catalyst::transport; //===----------------------------------------------------------------------===// #include "Transport/IR/TransportOpsDialect.cpp.inc" +#include "Transport/IR/TransportEnums.cpp.inc" //===----------------------------------------------------------------------===// // Transport type definitions. From bdfab4c57a348b25d7c8be21045e56a3bd2f04db Mon Sep 17 00:00:00 2001 From: Mehrdad Malekmohammadi Date: Wed, 22 Jul 2026 11:59:34 -0400 Subject: [PATCH 3/5] format --- mlir/include/Transport/IR/TransportDialect.h | 2 +- mlir/lib/Transport/IR/TransportDialect.cpp | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/mlir/include/Transport/IR/TransportDialect.h b/mlir/include/Transport/IR/TransportDialect.h index 57ab9b3dcb..675578a74c 100644 --- a/mlir/include/Transport/IR/TransportDialect.h +++ b/mlir/include/Transport/IR/TransportDialect.h @@ -18,8 +18,8 @@ #include "mlir/IR/Dialect.h" #include "mlir/IR/OpDefinition.h" -#include "Transport/IR/TransportOpsDialect.h.inc" #include "Transport/IR/TransportEnums.h.inc" +#include "Transport/IR/TransportOpsDialect.h.inc" #define GET_TYPEDEF_CLASSES #include "Transport/IR/TransportOpsTypes.h.inc" diff --git a/mlir/lib/Transport/IR/TransportDialect.cpp b/mlir/lib/Transport/IR/TransportDialect.cpp index aea2c90196..c2c41d2800 100644 --- a/mlir/lib/Transport/IR/TransportDialect.cpp +++ b/mlir/lib/Transport/IR/TransportDialect.cpp @@ -27,8 +27,8 @@ using namespace catalyst::transport; // Transport dialect definitions. //===----------------------------------------------------------------------===// -#include "Transport/IR/TransportOpsDialect.cpp.inc" #include "Transport/IR/TransportEnums.cpp.inc" +#include "Transport/IR/TransportOpsDialect.cpp.inc" //===----------------------------------------------------------------------===// // Transport type definitions. From e4a779b1ed1253a3402fb0f193000e6a78ab92e4 Mon Sep 17 00:00:00 2001 From: Mehrdad Malekmohammadi Date: Thu, 23 Jul 2026 10:58:32 -0400 Subject: [PATCH 4/5] add include to lsp --- mlir/tools/quantum-lsp-server/quantum-lsp-server.cpp | 3 +++ 1 file changed, 3 insertions(+) diff --git a/mlir/tools/quantum-lsp-server/quantum-lsp-server.cpp b/mlir/tools/quantum-lsp-server/quantum-lsp-server.cpp index 9cc18cbcc2..0330c721b0 100644 --- a/mlir/tools/quantum-lsp-server/quantum-lsp-server.cpp +++ b/mlir/tools/quantum-lsp-server/quantum-lsp-server.cpp @@ -30,6 +30,8 @@ #include "Quantum/IR/QuantumDialect.h" #include "RTIO/IR/RTIODialect.h" +#include "Transport/IR/TransportDialect.h" + int main(int argc, char **argv) { mlir::DialectRegistry registry; @@ -44,6 +46,7 @@ int main(int argc, char **argv) registry.insert(); registry.insert(); registry.insert(); + registry.insert(); registry.insert(); registry.insert(); From 44e0d6ace199f217a50f26c5aa18f3d2696eab0e Mon Sep 17 00:00:00 2001 From: Mehrdad Malekmohammadi Date: Thu, 23 Jul 2026 11:33:29 -0400 Subject: [PATCH 5/5] add changelog --- doc/releases/changelog-dev.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 35bb08012e..1b271f2e16 100644 --- a/doc/releases/changelog-dev.md +++ b/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)