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) 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..d2b9081dc0 --- /dev/null +++ b/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) diff --git a/mlir/include/Transport/IR/TransportDialect.h b/mlir/include/Transport/IR/TransportDialect.h new file mode 100644 index 0000000000..675578a74c --- /dev/null +++ b/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" diff --git a/mlir/include/Transport/IR/TransportDialect.td b/mlir/include/Transport/IR/TransportDialect.td new file mode 100644 index 0000000000..118aae1734 --- /dev/null +++ b/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 traits = []> + : TypeDef { + 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:$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 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..cc5406f8be --- /dev/null +++ b/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`) 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 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..1dbeed05e2 --- /dev/null +++ b/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 +) diff --git a/mlir/lib/Transport/IR/TransportDialect.cpp b/mlir/lib/Transport/IR/TransportDialect.cpp new file mode 100644 index 0000000000..c2c41d2800 --- /dev/null +++ b/mlir/lib/Transport/IR/TransportDialect.cpp @@ -0,0 +1,51 @@ +// 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/TransportEnums.cpp.inc" +#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-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(); 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();