Skip to content

MLIR: support CALL of registered procedures end to end - #759

Open
rjb32 wants to merge 65 commits into
mainfrom
feature/mlir/call-procedure
Open

MLIR: support CALL of registered procedures end to end#759
rjb32 wants to merge 65 commits into
mainfrom
feature/mlir/call-procedure

Conversation

@rjb32

@rjb32 rjb32 commented Jul 30, 2026

Copy link
Copy Markdown
Contributor

Runs a Cypher CALL of a registered procedure through the MLIR engine, from the frontend down to the interpreter.

A procedure that declares a row-aligned argument now has to report the input row of the rows it emits: registration refuses one that does not, and an argument read once per call rather than per row is declared as constant.

Works today:

// a call is the whole query
CALL db.labels() YIELD id, label RETURN id, label

// no YIELD: every return value the procedure declares is emitted
CALL db.labels()

// a procedure declaring no return value, driven for its effect alone
CALL test.sideEffect()

// YIELD renames what the query reads, and may take a subset of the return values
CALL db.history() YIELD commit AS commitHash, nodeCount, edgeCount, partCount
RETURN commitHash, nodeCount, edgeCount, partCount

// YIELD ... WHERE filters the rows the call emitted
CALL db.labels() YIELD label WHERE label = 'Manager' RETURN label

// correlated: the matched node is the argument, and rides through the call
MATCH (n) CALL test.nodeScore(n) YIELD score RETURN n, score

// the matched node rides through the call even where the projection drops it: every
// column in flight is carried, which the procedure's input-row report is what allows
MATCH (n) CALL test.doubleNodeID(n) YIELD doubled RETURN doubled

// uncorrelated after a MATCH: the call is crossed with it
MATCH (n) CALL db.labels() YIELD label RETURN n, label

// several yields off one call, here two lists
MATCH (n) CALL test.nodeTags(n) YIELD tags, coords RETURN n, tags, coords

// chained calls, the second's WHERE reading what the first yielded
MATCH (n) CALL test.expandNodeID(n) YIELD copy
          CALL test.fanOutNodeID(n) YIELD value WHERE copy > 100 RETURN value

// one call per end of a traversal, projecting an expression over both yields
MATCH (n)-->(m) CALL test.expandNodeID(n) YIELD copy AS x
                CALL test.fanOutNodeID(m) YIELD value AS y RETURN x + y

// chained calls and no MATCH: the first opens the dataflow, and its yield is both the
// next call's argument and carried on past it
CALL test.twoNodes() YIELD id
CALL test.expandNodeID(id) YIELD copy
CALL test.fanOutNodeID(id) YIELD value
RETURN id, copy, value

// three calls reading nothing of each other, so all three are crossed
CALL test.twoNodes() YIELD id AS a
CALL test.twoNodes() YIELD id AS b
CALL test.twoNodes() YIELD id AS c
RETURN a, b, c

Added a test for GNN-like sampling:

MATCH (n:Seed)
CALL gnn.neighbourhoodSample(n, 3, 42)     YIELD tgt AS hop1
CALL gnn.neighbourhoodSample(hop1, 2, 42)  YIELD tgt AS hop2
CALL gnn.neighbourhoodSample(hop2, 2, 42)  YIELD tgt AS hop3
RETURN n, hop1, hop2, hop3

Added a test for a call crossing and a call joining a hop chain:

CALL test.firstNodes(2) YIELD person
MATCH (person)-->(m)
MATCH (m)-->(z)
CALL test.offsetPair(17) YIELD offset
RETURN person, z, offset

CALL test.firstNodes(2) YIELD person AS a
MATCH (a)-->(m)
CALL test.firstNodes(3) YIELD person AS z
MATCH (m)-->(z)
RETURN a, m, z

@rjb32
rjb32 marked this pull request as draft July 30, 2026 16:54
Comment thread query/ir/codegen/DBProgramGenerator.cpp Outdated
Comment thread procedures/HistoryProcedure.cpp Outdated
Comment thread procedures/Procedure.h Outdated
Comment thread procedures/ProcedureData.h Outdated
Comment thread query/ir/codegen/DBProgramGenerator.h Outdated
@rjb32
rjb32 force-pushed the feature/mlir/call-procedure branch from 40cd558 to 1a12967 Compare August 3, 2026 14:52
@rjb32
rjb32 marked this pull request as ready for review August 3, 2026 15:32
@rjb32
rjb32 requested a review from cyrusknopf August 3, 2026 15:50
Comment thread procedures/Procedure.cpp Outdated
Comment on lines +29 to +32
void Procedure::setIndices(bool indices) {
_indices = indices;
}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I do this slightly differently in #754 , we should compare and see which pattern we want

Comment thread query/ir/codegen/DBProgramGenerator.h Outdated
Comment on lines +156 to +158
// Collect those columns. One bound in another block is skipped: an op here can only
// take what this block binds.
void collectLiveColumns(LiveColumns& live);

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What are "those" columns?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removed

Comment thread query/ir/codegen/DBProgramGenerator.cpp Outdated
Comment on lines +383 to +389
// A chunk holds the rows of the loop whose body binds it, so only a column bound in this
// block is row-aligned with the rows flowing past this point. One bound in an enclosing
// block - or in a loop the dataflow has already left - holds a different row set, and an
// op that consumes the whole row set (a filter, a call's carry set) would pair its rows
// with unrelated ones. Note this is stricter than dominance on purpose: an outer block's
// value does dominate here, it is just the wrong rows.
const auto isLive = [&](const mlir::Value column) {

@cyrusknopf cyrusknopf Aug 4, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If I understand correctly this comment can be simplified to: "a column is live if it is defined in the current block"

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done

Comment on lines +2179 to +2187
func.func @main() {
%call = nl.procedure("db.labels") yields ["id", "label"]
%rows = nl.procedure_init(%call, (), {}) : (!nl.procedure_state) -> !nl.iter<!nl.chunk<!storage.label_id>, !nl.chunk<!storage.string>>
nl.for %ids, %labels in %rows : !nl.iter<!nl.chunk<!storage.label_id>, !nl.chunk<!storage.string>> {
nl.output(%ids, %labels) : !nl.chunk<!storage.label_id>, !nl.chunk<!storage.string>
}
return
}
)mlir";

@cyrusknopf cyrusknopf Aug 4, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we have some example NL and DB programs as samples in samples/mlir?

@rjb32 rjb32 Aug 18, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done

Comment thread query/ir/codegen/DBProgramGenerator.cpp Outdated
Comment thread query/ir/codegen/DBProgramGenerator.cpp Outdated
// even when the projection returns none of their columns, or the query would lose
// their cardinality.
if (!live._columns.empty() && inputs.empty()) {
generateCrossedCall(procedureName, yieldedNames, yieldedVariables, live);

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we want separate logic for cross product of a MATCH and a CALL? Ideally I would want centralised cross product logic to join islands regardless of whether they are from a MATCH or a CALL. I think it would involve maintaining CALLs in the connected components (we already tracked connected components in generateTraversal

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes it would be a better shape to have one unified procedure for the cross product.
But today in this PR, the cross product is decided before CALL statements are processed, in generateTraversal. So this is where it would require integration of the CALLs with the VDG to use the same logic based on components identification in the VDG. For now we do codegen of calls afterwards to minimise changes to the VDG until we can come with a better approach.

Comment thread query/ir/codegen/DBProgramGenerator.cpp Outdated
Comment on lines +1388 to +1391
// A standalone CALL has no projection of its own: what it emits is the columns it
// yielded, in the order the procedure declares them. A call yielding none - a
// procedure declaring no return value - emits nothing at all, so the query is the
// drive and no output op is generated.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Verbose

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done

Comment thread query/ir/lowering/DBLowering.cpp Outdated
const llvm::StringRef name = mlir::cast<mlir::StringAttr>(yield).getValue();
const size_t returnIndex = procedure->getReturnValueIndex(std::string_view(name.data(), name.size()));

chunkTypes.push_back(procedureChunkType(_builder, procedure->getReturnValueType(returnIndex)));

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this could benefit from being split onto a few lines

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done

Comment thread query/ir/codegen/DBProgramGenerator.cpp Outdated
Comment thread query/ir/codegen/DBProgramGenerator.cpp Outdated
Comment thread query/ir/codegen/DBProgramGenerator.cpp Outdated
@rjb32
rjb32 force-pushed the feature/mlir/call-procedure branch from 79c2d3e to 2394f4f Compare August 7, 2026 15:04
Comment thread query/ir/codegen/DBProgramGenerator.h
Comment thread query/ir/codegen/DBProgramGenerator.h Outdated
Comment thread query/ir/codegen/DBProgramGenerator.h Outdated
Comment thread query/ir/codegen/DBProgramGenerator.h
Comment thread query/ir/dialect/db/DBOps.cpp
@rjb32
rjb32 force-pushed the feature/mlir/call-procedure branch 4 times, most recently from 7b92414 to 0090189 Compare August 17, 2026 13:24
@rjb32
rjb32 force-pushed the feature/mlir/call-procedure branch from ff441cf to 2459549 Compare August 21, 2026 10:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants