Support rtio.rpc op for OQD's AWG end to end pipeline by rniczh · Pull Request #2577 · PennyLaneAI/catalyst

rniczh · 2026-03-12T15:48:32Z

Context:

A new rtio.rpc operation is added to the RTIO dialect. It represents a host RPC call triggered by the kernel, optionally carrying runtime arguments and supporting both synchronous and async modes. The op is lowered to rpc_send / rpc_recv LLVM calls (the ARTIQ RPC wire protocol). It is required by both AWG control (program_awg, awg_close) and measurement result collection (set_dataset, transfer_data).

The RTIO dialect has no representation for RPC calls. Our use cases:

Use case	Callee	Args from kernel	Mode
AWG setup	`program_awg`	none	synchronous
AWG teardown	`awg_close`	none	synchronous
Dataset init	`set_dataset`	key (str), value (array ptr)	synchronous
Measurement result	`transfer_data`	key (str), index (i64), value (f64)	async

The ARTIQ runtime exposes rpc_send(int32_t service_id, const char *tag, ...) and rpc_recv(void *result) as extern C functions linked into the ELF.

Description of the Change:

Added rtio.rpc in RTIO dialect.
This op contains the following attributes:

- callee:      The symbolic name of the host-side callable to invoke.
- tag:         type-tag string encoding the return and argument types,
               which is compatiable with the ARTIQ RPC wire protocol.
               (e.g. "n:" for void no-arg, "n:si" for void(str,i32)).
- is_async:    If present the call is async.
- rpc_id:      Populated by the RPC-ID assignment sub-pass.

Benefits:

Possible Drawbacks:

Related GitHub Issues:
[sc-113163]

github-actions · 2026-03-16T16:04:50Z

Hello. You may have forgotten to update the changelog!
Please edit doc/releases/changelog-dev.md on your branch with:

A one-to-two sentence description of the change. You may include a small working example for new features.
A link back to this PR.
Your name (or GitHub username) in the contributors section.

codecov · 2026-03-16T16:32:40Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 97.00%. Comparing base (d54fdea) to head (936a580).
⚠️ Report is 2 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2577   +/-   ##
=======================================
  Coverage   97.00%   97.00%           
=======================================
  Files         157      157           
  Lines       17470    17470           
  Branches     1686     1686           
=======================================
  Hits        16946    16946           
  Misses        387      387           
  Partials      137      137

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

mehrdad2m

Thanks @rniczh, Looks great. minor comments

mlir/lib/Driver/CompilerDriver.cpp

mehrdad2m · 2026-03-17T21:13:57Z

mlir/include/RTIO/IR/RTIOOps.td

+        - callee:    The symbolic name of the host-side callable to invoke.
+        - tag:       type-tag string encoding the return and argument types,
+                     which is compatiable with the ARTIQ RPC wire protocol.
+                     (e.g. "n:" for void no-arg, "n:si" for void(str,i32)).


This ARTIQ related part seems out of context. Ideally we should try to not pollute this dialect with ARTIQ details.

now that I think more, this could be a general enough scheme that is not ARTIQ dependant, but still you may want lower to any arbitrary tagging system, so not tying up it up to ARTIQ seams reasonable?

Yeah, that's right, for ARTIQ that tag can be derived from this operand types I think

I removed the tag adb3b83. which can be derived from the operand/result types during the conversion from rtio to artiq

mehrdad2m · 2026-03-17T21:18:00Z

mlir/include/RTIO/IR/RTIOOps.td

+                     which is compatiable with the ARTIQ RPC wire protocol.
+                     (e.g. "n:" for void no-arg, "n:si" for void(str,i32)).
+        - is_async:  If present the call is async.
+        - rpc_id:    Populated by the RPC-ID assignment sub-pass.


I think the rpc_id explanation is a bit too specific. We should explain what the attribute is or why it is useful. not what we do with it in the lowering.

Thanks! I updated the explanation here. adb3b83
which is the only way that host device to communicate with the RPC server (through the rpc id to determine which function should be used)

mehrdad2m · 2026-03-18T18:14:40Z

mlir/lib/RTIO/IR/RTIOOps.cpp

+//===----------------------------------------------------------------------===//
+
+namespace {
+/// ARTIQ RPC tag format: <return_type>:<arg_types>


Hmm again I am wondering if this would be better suited in a lower level place, maybe in the the lowering pass. What do you think?

I think this is a bit nitpicky given that currently, we only use this dialect in the context of lowering to ARTIQ. So if the work needed to move it down is substantial, I think it would be fine to have it here with a TODO for now since the dialect is not used anywhere else, but keeping it independent of the ARTIQ internal details would be nice.

Yes, adb3b83. I removed the tag, it will only be constructed when converting to the artiq llvm ir

mlir/lib/RTIO/IR/RTIOOps.cpp

mehrdad2m · 2026-03-18T19:11:55Z

mlir/lib/RTIO/Transforms/RTIOEventToARTIQPatterns.cpp

+            size_t colon = tag.find(':');
+            StringRef returnPart = colon != StringRef::npos ? tag.take_front(colon) : "n";
+
+            int slotSize = (returnPart == "n") ? 8 : getSlotSizeForReturnCode(returnPart[0]);


Why do we single out "n"? couldn't we include it in the switch case within getSlotSizeForReturnCode?

originally I was want to distinguish the n with other types (since it's void). But yeah, I think it's fine to move it into the getSlotSzieForReturnCode function. It doesn't need to be separated out. d36e3b5

mehrdad2m · 2026-03-18T19:14:37Z

mlir/lib/RTIO/Transforms/RTIOEventToARTIQPatterns.cpp

+            StringRef returnPart = colon != StringRef::npos ? tag.take_front(colon) : "n";
+
+            int slotSize = (returnPart == "n") ? 8 : getSlotSizeForReturnCode(returnPart[0]);
+            auto slotTy = LLVM::LLVMArrayType::get(i8Ty, slotSize);


Can you explain the logic here? I don't understand why we allocate an array of bytes instead of using a concrete resultTy?

This is the wire protocol of the ARTIQ, they only allow you to pass raw bytes buffer with llvm ptr, except for the the rpc id (the first arg).

Co-authored-by: Mehrdad Malek <[email protected]>

Support rtio.rpc op

804b596

rniczh changed the title ~~Support rtio.rpc op~~ [WIP] Support rtio.rpc op Mar 12, 2026

rniczh marked this pull request as draft March 12, 2026 15:49

rniczh added 3 commits March 16, 2026 09:58

Merge branch 'main' into rniczh/support-rtio-rpc-op

db71d1e

formatting

d2fbb76

Remove target dependent description from the RTIO dialect

ed66f7e

rniczh changed the title ~~[WIP] Support rtio.rpc op~~ Support rtio.rpc op Mar 16, 2026

rniczh marked this pull request as ready for review March 16, 2026 16:04

fix type

e4e13e5

rniczh added 4 commits March 16, 2026 13:51

support return values

fbd9b44

Merge branch 'main' into rniczh/support-rtio-rpc-op

62ee8bf

typo

e97835d

revert

f2a7a60

rniczh changed the title ~~Support rtio.rpc op~~ Support rtio.rpc op for OQD's AWG end to end pipeline Mar 16, 2026

mehrdad2m reviewed Mar 18, 2026

View reviewed changes

rniczh and others added 7 commits March 18, 2026 16:46

Update mlir/lib/RTIO/IR/RTIOOps.cpp

6ab8758

Co-authored-by: Mehrdad Malek <[email protected]>

Update mlir/lib/Driver/CompilerDriver.cpp

3555860

Co-authored-by: Mehrdad Malek <[email protected]>

Merge branch 'main' into rniczh/support-rtio-rpc-op

1c6760e

Merge branch 'main' into rniczh/support-rtio-rpc-op

18884c1

refactor rpc function

adb3b83

unify the return code processing

d36e3b5

refactor

936a580

Conversation

rniczh commented Mar 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Mar 16, 2026

Uh oh!

codecov bot commented Mar 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

mehrdad2m left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

rniczh commented Mar 12, 2026 •

edited

Loading

codecov bot commented Mar 16, 2026 •

edited

Loading