Skip to content

Commit fae7c7f

Browse files
mehrdad2mjosephleeklrniczhmultiphaseCFD
authored
Add inject-transport-session pass to emit transport session IR (#3063)
**Context:** Backline places a QNode's controller and its coprocessors on separate executors that talk over a transport. The placement is serialized onto the module as the catalyst.backline attribute (controller + coprocessors + per-node backend/peer/triple/data-path info). This pass turns that declarative attribute into the concrete transport-dialect session lifecycle, wired into the setup/teardown functions the runtime already calls around every execution. **Description of the Change:** Reading `catalyst.backline,` the pass injects, per session: bring-up into `@setup`: `transport.create` → `connect` → `exchange_keys` → `establish_channel` → `commit_work_item` (controller) / `set_coprocessor_fn` (coprocessor) → start tear-down into `@teardown`: `get_session` → `stop` → `destroy` Topologies handled: - Controller-only: single controller session, straight-line bring-up/teardown. - Local coprocessors: controller + coprocessor sessions created together in the host @setup; the coprocessor's `connect/exchange_keys` run async so the two ends handshake concurrently. - Remote coprocessors: each coprocessor's ops go into a generated `module_coproc` tagged with its catalyst.target triple and catalyst.dispatch address; the host launches `coproc_serve` (nonblocking) / `coproc_stop`. - Remote controller: the controller's ops go into its role-tagged target module as `setup_transport/teardown_transport` and the host `@setup/@teardown` become the single orchestration point that launches every dispatched role in dependency order (serve → controller setup; coproc stop → controller teardown), and `@setup` is marked catalyst.backline_bringup so the object-shipping pass knows where to ship from. **Benefits:** **Possible Drawbacks:** **Related GitHub Issues:** --------- Co-authored-by: Joseph Lee <joseph.lee@xanadu.ai> Co-authored-by: Joseph Lee <40768758+josephleekl@users.noreply.github.com> Co-authored-by: Hong-Sheng Zheng <mathan0203@gmail.com> Co-authored-by: Shuli Shu <31480676+multiphaseCFD@users.noreply.github.com>
1 parent f94268a commit fae7c7f

11 files changed

Lines changed: 993 additions & 3 deletions

File tree

‎doc/releases/changelog-dev.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,10 @@
3535
transport runtime CAPI.
3636
[(#3048)](https://github.com/PennyLaneAI/catalyst/pull/3048)
3737

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

‎mlir/include/Transport/IR/CMakeLists.txt‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,4 +5,6 @@ add_mlir_doc(TransportOps TransportOps Transport/ -gen-op-doc)
55
set(LLVM_TARGET_DEFINITIONS TransportOps.td)
66
mlir_tablegen(TransportEnums.h.inc -gen-enum-decls)
77
mlir_tablegen(TransportEnums.cpp.inc -gen-enum-defs)
8+
mlir_tablegen(TransportAttributes.h.inc -gen-attrdef-decls -attrdefs-dialect=transport)
9+
mlir_tablegen(TransportAttributes.cpp.inc -gen-attrdef-defs -attrdefs-dialect=transport)
810
add_public_tablegen_target(MLIRTransportEnumsIncGen)

‎mlir/include/Transport/IR/TransportDialect.td‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ def Transport_Dialect : Dialect {
3737
let name = "transport";
3838
let cppNamespace = "::catalyst::transport";
3939
let useDefaultTypePrinterParser = 1;
40+
let useDefaultAttributePrinterParser = 1;
4041
let usePropertiesForAttributes = 1;
4142
}
4243

@@ -86,6 +87,72 @@ def Transport_CoprocessorSession : Type<
8687
"::catalyst::transport::Role::Coprocessor">,
8788
"coprocessor transport session">;
8889

90+
//===----------------------------------------------------------------------===//
91+
// Attributes.
92+
//===----------------------------------------------------------------------===//
93+
94+
class Transport_Attr<string name, string attrMnemonic, list<Trait> traits = []>
95+
: AttrDef<Transport_Dialect, name, traits> {
96+
let mnemonic = attrMnemonic;
97+
}
98+
99+
// A string field that reads as empty when omitted, so consumers need no null checks.
100+
def Transport_StrField : DefaultValuedParameter<"mlir::StringAttr",
101+
"mlir::StringAttr::get($_ctxt, \"\")">;
102+
103+
def Transport_NodeAttr : Transport_Attr<"Node", "node"> {
104+
let summary = "A backline participant: a controller or a coprocessor.";
105+
106+
let parameters = (ins
107+
Transport_StrField:$name,
108+
Transport_StrField:$peer,
109+
OptionalParameter<"mlir::IntegerAttr">:$oob_port,
110+
Transport_StrField:$backend_lib,
111+
Transport_StrField:$config,
112+
Transport_StrField:$data_path,
113+
Transport_StrField:$triple,
114+
Transport_StrField:$address,
115+
Transport_StrField:$symbol,
116+
OptionalParameter<"mlir::BoolAttr">:$remote,
117+
OptionalParameter<"mlir::IntegerAttr">:$in_bytes,
118+
OptionalParameter<"mlir::IntegerAttr">:$out_bytes,
119+
OptionalParameter<"mlir::IntegerAttr">:$work_item_idx
120+
);
121+
122+
let assemblyFormat = "`<` struct(params) `>`";
123+
124+
let extraClassDeclaration = [{
125+
/// Session registry key: `name` when set and non-empty, else `fallback`.
126+
mlir::StringAttr keyOr(llvm::StringRef fallback) const;
127+
128+
/// `data_path` when set, else `dflt`.
129+
mlir::StringAttr dataPathOr(llvm::StringRef dflt) const;
130+
131+
/// True only when `remote` is present and set.
132+
bool isRemote() const;
133+
134+
/// Integer fields, with the defaults applied.
135+
int64_t oobPort() const;
136+
int64_t inBytes() const;
137+
int64_t outBytes() const;
138+
int64_t workItemIdx() const;
139+
}];
140+
}
141+
142+
def Transport_BacklineAttr : Transport_Attr<"Backline", "backline"> {
143+
let summary = "A backline placement: one controller and the coprocessors it drives.";
144+
145+
let parameters = (ins
146+
"mlir::StringAttr":$transport,
147+
"::catalyst::transport::NodeAttr":$controller,
148+
OptionalArrayRefParameter<"::catalyst::transport::NodeAttr">:$coprocessors
149+
);
150+
151+
let assemblyFormat = [{ `<` `transport` `=` $transport `,` `controller` `=` $controller
152+
(`,` `coprocessors` `=` `[` $coprocessors^ `]`)? `>` }];
153+
let genVerifyDecl = 1;
154+
}
155+
89156
//===----------------------------------------------------------------------===//
90157
// Operation base.
91158
//===----------------------------------------------------------------------===//

‎mlir/include/Transport/IR/TransportOps.h‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,5 +22,8 @@
2222

2323
#include "Transport/IR/TransportDialect.h"
2424

25+
#define GET_ATTRDEF_CLASSES
26+
#include "Transport/IR/TransportAttributes.h.inc"
27+
2528
#define GET_OP_CLASSES
2629
#include "Transport/IR/TransportOps.h.inc"

‎mlir/include/Transport/Transforms/Passes.td‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,4 +29,17 @@ def ConvertTransportToLLVMPass : Pass<"convert-transport-to-llvm", "mlir::Module
2929
];
3030
}
3131

32+
def InjectTransportSessionPass : Pass<"inject-transport-session", "mlir::ModuleOp"> {
33+
let summary = "Emit transport session bring-up/teardown from a catalyst.backline attribute.";
34+
let description = [{
35+
Reads the `catalyst.backline` module attribute and injects the session
36+
lifecycle into the host entry function.
37+
}];
38+
39+
let dependentDialects = [
40+
"::catalyst::transport::TransportDialect",
41+
"mlir::func::FuncDialect"
42+
];
43+
}
44+
3245
#endif // TRANSPORT_PASSES

‎mlir/lib/Transport/IR/TransportDialect.cpp‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,13 +37,78 @@ using namespace catalyst::transport;
3737
#define GET_TYPEDEF_CLASSES
3838
#include "Transport/IR/TransportOpsTypes.cpp.inc"
3939

40+
//===----------------------------------------------------------------------===//
41+
// Transport attribute definitions.
42+
//===----------------------------------------------------------------------===//
43+
44+
#define GET_ATTRDEF_CLASSES
45+
#include "Transport/IR/TransportAttributes.cpp.inc"
46+
47+
StringAttr NodeAttr::keyOr(llvm::StringRef fallback) const
48+
{
49+
if (StringAttr n = getName(); n && !n.getValue().empty()) {
50+
return n;
51+
}
52+
return StringAttr::get(getContext(), fallback);
53+
}
54+
55+
StringAttr NodeAttr::dataPathOr(llvm::StringRef dflt) const
56+
{
57+
if (StringAttr p = getDataPath(); p && !p.getValue().empty()) {
58+
return p;
59+
}
60+
return StringAttr::get(getContext(), dflt);
61+
}
62+
63+
bool NodeAttr::isRemote() const
64+
{
65+
BoolAttr r = getRemote();
66+
return r && r.getValue();
67+
}
68+
69+
static int64_t intOr(IntegerAttr field, int64_t dflt) { return field ? field.getInt() : dflt; }
70+
71+
// Default per-message payload width, matching the current backend's defaults.
72+
constexpr int64_t kDefaultMessageBytes = 8;
73+
74+
int64_t NodeAttr::oobPort() const { return intOr(getOobPort(), 0); }
75+
int64_t NodeAttr::inBytes() const { return intOr(getInBytes(), kDefaultMessageBytes); }
76+
int64_t NodeAttr::outBytes() const { return intOr(getOutBytes(), kDefaultMessageBytes); }
77+
int64_t NodeAttr::workItemIdx() const { return intOr(getWorkItemIdx(), 0); }
78+
79+
LogicalResult BacklineAttr::verify(function_ref<InFlightDiagnostic()> emitError,
80+
StringAttr transport, NodeAttr controller,
81+
llvm::ArrayRef<NodeAttr> coprocessors)
82+
{
83+
if (!controller) {
84+
return emitError() << "backline requires a controller";
85+
}
86+
for (NodeAttr c : coprocessors) {
87+
if (!c) {
88+
return emitError() << "null coprocessor";
89+
}
90+
if (!c.getPeer() || c.getPeer().getValue().empty()) {
91+
return emitError() << "coprocessor requires a 'peer'";
92+
}
93+
if (!c.getSymbol() || c.getSymbol().getValue().empty()) {
94+
return emitError() << "coprocessor requires a 'symbol'";
95+
}
96+
}
97+
return success();
98+
}
99+
40100
void TransportDialect::initialize()
41101
{
42102
addTypes<
43103
#define GET_TYPEDEF_LIST
44104
#include "Transport/IR/TransportOpsTypes.cpp.inc"
45105
>();
46106

107+
addAttributes<
108+
#define GET_ATTRDEF_LIST
109+
#include "Transport/IR/TransportAttributes.cpp.inc"
110+
>();
111+
47112
addOperations<
48113
#define GET_OP_LIST
49114
#include "Transport/IR/TransportOps.cpp.inc"

‎mlir/lib/Transport/Transforms/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ set(LIBRARY_NAME transport-transforms)
22

33
file(GLOB SRC
44
TransportToLLVM.cpp
5+
InjectTransportSession.cpp
56
)
67

78
get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)

0 commit comments

Comments
 (0)