Add inject-transport-session pass to emit transport session IR

doc/releases/changelog-dev.md

@@ -35,6 +35,10 @@
   transport runtime CAPI.
   [(#3048)](https://github.com/PennyLaneAI/catalyst/pull/3048)
 
+* An `inject-transport-session` pass is added, which reads the `catalyst.backline` module
+  attribute and emits the transport session lifecycle into the host entry function.
+  [(#3063)](https://github.com/PennyLaneAI/catalyst/pull/3063)
+
 * A new remote/local executor infrastructure has been added to Catalyst, enabling qnode kernels to
   be dispatched to a separate executor process.
 

mlir/include/Transport/IR/CMakeLists.txt

@@ -5,4 +5,6 @@ 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)
+mlir_tablegen(TransportAttributes.h.inc -gen-attrdef-decls -attrdefs-dialect=transport)
+mlir_tablegen(TransportAttributes.cpp.inc -gen-attrdef-defs -attrdefs-dialect=transport)
 add_public_tablegen_target(MLIRTransportEnumsIncGen)

mlir/include/Transport/IR/TransportDialect.td

@@ -37,6 +37,7 @@ def Transport_Dialect : Dialect {
     let name = "transport";
     let cppNamespace = "::catalyst::transport";
     let useDefaultTypePrinterParser = 1;
+    let useDefaultAttributePrinterParser = 1;
     let usePropertiesForAttributes = 1;
 }
 
@@ -86,6 +87,72 @@ def Transport_CoprocessorSession : Type<
           "::catalyst::transport::Role::Coprocessor">,
     "coprocessor transport session">;
 
+//===----------------------------------------------------------------------===//
+// Attributes.
+//===----------------------------------------------------------------------===//
+
+class Transport_Attr<string name, string attrMnemonic, list<Trait> traits = []>
+    : AttrDef<Transport_Dialect, name, traits> {
+    let mnemonic = attrMnemonic;
+}
+
+// A string field that reads as empty when omitted, so consumers need no null checks.
+def Transport_StrField : DefaultValuedParameter<"mlir::StringAttr",
+                                                "mlir::StringAttr::get($_ctxt, \"\")">;
+
+def Transport_NodeAttr : Transport_Attr<"Node", "node"> {
+    let summary = "A backline participant: a controller or a coprocessor.";
+
+    let parameters = (ins
+        Transport_StrField:$name,
+        Transport_StrField:$peer,
+        OptionalParameter<"mlir::IntegerAttr">:$oob_port,
+        Transport_StrField:$backend_lib,
+        Transport_StrField:$config,
+        Transport_StrField:$data_path,
+        Transport_StrField:$triple,
+        Transport_StrField:$address,
+        Transport_StrField:$symbol,
+        OptionalParameter<"mlir::BoolAttr">:$remote,
+        OptionalParameter<"mlir::IntegerAttr">:$in_bytes,
+        OptionalParameter<"mlir::IntegerAttr">:$out_bytes,
+        OptionalParameter<"mlir::IntegerAttr">:$work_item_idx
+    );
+
+    let assemblyFormat = "`<` struct(params) `>`";
+
+    let extraClassDeclaration = [{
+        /// Session registry key: `name` when set and non-empty, else `fallback`.
+        mlir::StringAttr keyOr(llvm::StringRef fallback) const;
+
+        /// `data_path` when set, else `dflt`.
+        mlir::StringAttr dataPathOr(llvm::StringRef dflt) const;
+
+        /// True only when `remote` is present and set.
+        bool isRemote() const;
+
+        /// Integer fields, with the defaults applied.
+        int64_t oobPort() const;
+        int64_t inBytes() const;
+        int64_t outBytes() const;
+        int64_t workItemIdx() const;
+    }];
+}
+
+def Transport_BacklineAttr : Transport_Attr<"Backline", "backline"> {
+    let summary = "A backline placement: one controller and the coprocessors it drives.";
+
+    let parameters = (ins
+        "mlir::StringAttr":$transport,
+        "::catalyst::transport::NodeAttr":$controller,
+        OptionalArrayRefParameter<"::catalyst::transport::NodeAttr">:$coprocessors
+    );
+
+    let assemblyFormat = [{ `<` `transport` `=` $transport `,` `controller` `=` $controller
+                            (`,` `coprocessors` `=` `[` $coprocessors^ `]`)? `>` }];
+    let genVerifyDecl = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // Operation base.
 //===----------------------------------------------------------------------===//

mlir/include/Transport/IR/TransportOps.h

@@ -22,5 +22,8 @@
 
 #include "Transport/IR/TransportDialect.h"
 
+#define GET_ATTRDEF_CLASSES
+#include "Transport/IR/TransportAttributes.h.inc"
+
 #define GET_OP_CLASSES
 #include "Transport/IR/TransportOps.h.inc"

mlir/include/Transport/Transforms/Passes.td

@@ -29,4 +29,17 @@ def ConvertTransportToLLVMPass : Pass<"convert-transport-to-llvm", "mlir::Module
     ];
 }
 
+def InjectTransportSessionPass : Pass<"inject-transport-session", "mlir::ModuleOp"> {
+    let summary = "Emit transport session bring-up/teardown from a catalyst.backline attribute.";
+    let description = [{
+        Reads the `catalyst.backline` module attribute and injects the session
+        lifecycle into the host entry function.
+    }];
+
+    let dependentDialects = [
+        "::catalyst::transport::TransportDialect",
+        "mlir::func::FuncDialect"
+    ];
+}
+
 #endif // TRANSPORT_PASSES

mlir/lib/Transport/IR/TransportDialect.cpp

@@ -37,13 +37,78 @@ using namespace catalyst::transport;
 #define GET_TYPEDEF_CLASSES
 #include "Transport/IR/TransportOpsTypes.cpp.inc"
 
+//===----------------------------------------------------------------------===//
+// Transport attribute definitions.
+//===----------------------------------------------------------------------===//
+
+#define GET_ATTRDEF_CLASSES
+#include "Transport/IR/TransportAttributes.cpp.inc"
+
+StringAttr NodeAttr::keyOr(llvm::StringRef fallback) const
+{
+    if (StringAttr n = getName(); n && !n.getValue().empty()) {
+        return n;
+    }
+    return StringAttr::get(getContext(), fallback);
+}
+
+StringAttr NodeAttr::dataPathOr(llvm::StringRef dflt) const
+{
+    if (StringAttr p = getDataPath(); p && !p.getValue().empty()) {
+        return p;
+    }
+    return StringAttr::get(getContext(), dflt);
+}
+
+bool NodeAttr::isRemote() const
+{
+    BoolAttr r = getRemote();
+    return r && r.getValue();
+}
+
+static int64_t intOr(IntegerAttr field, int64_t dflt) { return field ? field.getInt() : dflt; }
+
+// Default per-message payload width, matching the current backend's defaults.
+constexpr int64_t kDefaultMessageBytes = 8;
+
+int64_t NodeAttr::oobPort() const { return intOr(getOobPort(), 0); }
+int64_t NodeAttr::inBytes() const { return intOr(getInBytes(), kDefaultMessageBytes); }
+int64_t NodeAttr::outBytes() const { return intOr(getOutBytes(), kDefaultMessageBytes); }
+int64_t NodeAttr::workItemIdx() const { return intOr(getWorkItemIdx(), 0); }
+
+LogicalResult BacklineAttr::verify(function_ref<InFlightDiagnostic()> emitError,
+                                   StringAttr transport, NodeAttr controller,
+                                   llvm::ArrayRef<NodeAttr> coprocessors)
+{
+    if (!controller) {
+        return emitError() << "backline requires a controller";
+    }
+    for (NodeAttr c : coprocessors) {
+        if (!c) {
+            return emitError() << "null coprocessor";
+        }
+        if (!c.getPeer() || c.getPeer().getValue().empty()) {
+            return emitError() << "coprocessor requires a 'peer'";
+        }
+        if (!c.getSymbol() || c.getSymbol().getValue().empty()) {
+            return emitError() << "coprocessor requires a 'symbol'";
+        }
+    }
+    return success();
+}
+
 void TransportDialect::initialize()
 {
     addTypes<
 #define GET_TYPEDEF_LIST
 #include "Transport/IR/TransportOpsTypes.cpp.inc"
         >();
 
+    addAttributes<
+#define GET_ATTRDEF_LIST
+#include "Transport/IR/TransportAttributes.cpp.inc"
+        >();
+
     addOperations<
 #define GET_OP_LIST
 #include "Transport/IR/TransportOps.cpp.inc"

mlir/lib/Transport/Transforms/CMakeLists.txt

@@ -2,6 +2,7 @@ set(LIBRARY_NAME transport-transforms)
 
 file(GLOB SRC
     TransportToLLVM.cpp
+    InjectTransportSession.cpp
 )
 
 get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)