From 439d9903f88f6cdb8ed2dd2a603f944cbc063186 Mon Sep 17 00:00:00 2001 From: Remy Boutonnet Date: Wed, 21 Oct 2026 18:38:26 +0200 Subject: [PATCH 1/6] MLIR: name the rows a constant is read over in the db IR --- query/ir/codegen/DBProgramGenerator.cpp | 84 +++++-- query/ir/codegen/DBProgramGenerator.h | 25 +- query/ir/dialect/db/DBOps.cpp | 24 ++ query/ir/dialect/db/DBOps.h | 8 + query/ir/dialect/db/DBOps.td | 6 + query/ir/dialect/db/DBPasses.cpp | 87 ++++++- query/ir/lowering/DBLowering.cpp | 237 +++++++++--------- query/ir/lowering/DBLowering.h | 8 +- test/query/ir/call/MLIRCallProcedureTest.cpp | 9 +- test/query/ir/engine/DBLoweringTest.cpp | 30 ++- test/query/ir/projection/CMakeLists.txt | 2 + test/query/ir/projection/ConstantRowsTest.cpp | 89 +++++++ 12 files changed, 436 insertions(+), 173 deletions(-) create mode 100644 test/query/ir/projection/ConstantRowsTest.cpp diff --git a/query/ir/codegen/DBProgramGenerator.cpp b/query/ir/codegen/DBProgramGenerator.cpp index 82085b026e..b6d7cf2bc6 100644 --- a/query/ir/codegen/DBProgramGenerator.cpp +++ b/query/ir/codegen/DBProgramGenerator.cpp @@ -753,6 +753,13 @@ EdgeMetadata::EdgeType reverseEdge(EdgeMetadata::EdgeType type) { } +// A constant, or one laid out over a relation, holds one value in every row +bool holdsOneValue(mlir::Value column) { + mlir::db::BroadcastConstant broadcast = column.getDefiningOp(); + + return yieldsConstantColumn(broadcast ? broadcast.getValue() : column); +} + // A constant column holds the same value in every row, so a step shaped by rows - a // dedup, a sort, a cut - is not given one: reading it would anchor that step where the // constant is bound, above the loop the other columns are read in @@ -998,7 +1005,8 @@ DBProgramGenerator::DBProgramGenerator(mlir::ModuleOp* mainModule, const mlir::db::DBPassContext& passContext) : _module(mainModule), _mlirCtxt(_module->getContext()), - _opBuilder(_module->getBodyRegion()), + _constantRowsListener(this), + _opBuilder(_module->getBodyRegion(), &_constantRowsListener), _explain(explain), _passContext(passContext) { @@ -1007,6 +1015,15 @@ DBProgramGenerator::DBProgramGenerator(mlir::ModuleOp* mainModule, DBProgramGenerator::~DBProgramGenerator() { } +DBProgramGenerator::ConstantRowsListener::ConstantRowsListener(DBProgramGenerator* generator) + : _generator(generator) +{ +} + +void DBProgramGenerator::ConstantRowsListener::notifyOperationInserted(mlir::Operation* op, mlir::OpBuilder::InsertPoint previous) { + _generator->layConstantOperandsOverRows(op); +} + mlir::db::ColumnType DBProgramGenerator::allocColumnType(mlir::Type type) { return mlir::db::ColumnType::get(_mlirCtxt, type); } @@ -5085,7 +5102,7 @@ mlir::Value DBProgramGenerator::resolveColumnInScope(ColumnPredicate accept) con } const mlir::Value column = findIt->second.back(); - if (!isRowAlignedHere(column) || !accept(column)) { + if (!holdsTheRowsInFlight(column) || !accept(column)) { continue; } @@ -5093,12 +5110,26 @@ mlir::Value DBProgramGenerator::resolveColumnInScope(ColumnPredicate accept) con } for (const YieldedColumn& yielded : _part._yieldedColumns) { - if (isRowAlignedHere(yielded._column)) { + if (holdsTheRowsInFlight(yielded._column)) { return yielded._column; } } - return mlir::Value {}; + // The entities one write created hold the same rows, so the first one bound stands for them + mlir::Value firstCreated; + for (const auto& [decl, created] : _part._createdEntities) { + const mlir::Value column = created._column; + if (!column || !holdsTheRowsInFlight(column) || !accept(column)) { + continue; + } + + const bool boundEarlier = !firstCreated || computedBefore(column, firstCreated); + if (boundEarlier) { + firstCreated = column; + } + } + + return firstCreated; } mlir::Value DBProgramGenerator::readWalkEntities(const Expr* argExpr, mlir::Value column) { @@ -5130,15 +5161,8 @@ mlir::Value DBProgramGenerator::translateAggregateInput(const Expr* argExpr, return getOrTranslateBoundColumn(argExpr); } else if (argType != EvaluatedType::Wildcard) { - mlir::Block* const block = _opBuilder.getInsertionBlock(); - const mlir::Block::iterator insertionPoint = _opBuilder.getInsertionPoint(); - mlir::Operation* const lastOpBefore = insertionPoint == block->begin() ? nullptr : &*std::prev(insertionPoint); - translateExpr(argExpr); - const mlir::Block::iterator firstTranslated = lastOpBefore ? std::next(lastOpBefore->getIterator()) : block->begin(); - alignConstantOperandsToDriver(firstTranslated, _opBuilder.getInsertionPoint()); - return alignConstantToDriver(_part._exprMap.at(argExpr)); } @@ -5159,7 +5183,13 @@ mlir::Value DBProgramGenerator::translateAggregateInput(const Expr* argExpr, } mlir::Value DBProgramGenerator::resolveRowCarryingColumn() { - return resolveColumnInScope([this](mlir::Value column) { return !yieldsConstantColumn(column, _constantColumns); }); + const mlir::Value column = resolveColumnInScope([this](mlir::Value column) { return !yieldsConstantColumn(column, _constantColumns); }); + if (column || !_part._aggregateOp) { + return column; + } + + // Under grouping keys no variable names, the groups are held by the aggregate's results alone + return _part._aggregateOp->getResult(0); } mlir::Value DBProgramGenerator::alignConstantToDriver(mlir::Value column) { @@ -5177,29 +5207,26 @@ mlir::Value DBProgramGenerator::alignConstantToDriver(mlir::Value column) { return _opBuilder.create(_opBuilder.getUnknownLoc(), noneType, column, driver).getResult(); } -void DBProgramGenerator::alignConstantOperandsToDriver(mlir::Block::iterator first, mlir::Block::iterator last) { +void DBProgramGenerator::layConstantOperandsOverRows(mlir::Operation* op) { + if (!mlir::db::readsConstantColumnsAlone(op, _constantColumns)) { + return; + } + const mlir::Value driver = resolveRowCarryingColumn(); if (!driver) { return; } - const auto isConstant = [this](mlir::Value value) { return yieldsConstantColumn(value, _constantColumns); }; - - const mlir::Location loc = _opBuilder.getUnknownLoc(); const mlir::db::ColumnType noneType = allocColumnType(mlir::NoneType::get(_mlirCtxt)); - const mlir::OpBuilder::InsertionGuard insertionGuard(_opBuilder); - for (mlir::Operation& op : llvm::make_range(first, last)) { - const bool readsConstantsAlone = op.getNumOperands() > 0 && llvm::all_of(op.getOperands(), isConstant); - const bool runsPerRow = !llvm::all_of(op.getResults(), isConstant); - if (!readsConstantsAlone || !runsPerRow) { + mlir::OpBuilder builder(op); + for (mlir::OpOperand& operand : op->getOpOperands()) { + const mlir::Value column = operand.get(); + if (!mlir::isa(column.getType())) { continue; } - _opBuilder.setInsertionPoint(&op); - for (mlir::OpOperand& operand : op.getOpOperands()) { - operand.set(_opBuilder.create(loc, noneType, operand.get(), driver).getResult()); - } + operand.set(builder.create(op->getLoc(), noneType, column, driver).getResult()); } } @@ -5384,6 +5411,8 @@ void DBProgramGenerator::generateOutput(const Projection* projection, const Unio llvm::SmallVector outputNames; translateProjection(projection, variableColumns, outputted, outputNames); + broadcastConstantProjection(outputted); + buildNamedPathItems(projection, outputted); expandPathItems(projection, outputted); @@ -6888,10 +6917,11 @@ void DBProgramGenerator::translateDistinct(const Projection* projection, llvm::SmallVector dedupedCarriedColumns; collectCarriedRowColumns(carriedColumns, dedupedCarriedColumns, dedupedColumns); - // Every column is constant, so the projection is one row repeated as many times as + // Every column holds one value, so the projection is one row repeated as many times as // the query matched: no column tells two of those rows apart, and the one the dedup // keeps is the first of them - the projection capped at a single row - if (dedupedColumns.empty()) { + const bool oneValuePerColumn = std::ranges::all_of(projected, holdsOneValue) && dedupedCarriedColumns.empty(); + if (dedupedColumns.empty() || oneValuePerColumn) { translateDistinctOverConstants(projection, projected); } else { llvm::SmallVector deduped; diff --git a/query/ir/codegen/DBProgramGenerator.h b/query/ir/codegen/DBProgramGenerator.h index 749c83558c..2e734c620e 100644 --- a/query/ir/codegen/DBProgramGenerator.h +++ b/query/ir/codegen/DBProgramGenerator.h @@ -152,8 +152,21 @@ class DBProgramGenerator { void generateExplainResult(const ExplainReport& report); private: + // An op reading constants alone runs over a single row, so every op inserted while rows + // are in flight has the constants it reads alone laid over them + class ConstantRowsListener : public mlir::OpBuilder::Listener { + public: + explicit ConstantRowsListener(DBProgramGenerator* generator); + + void notifyOperationInserted(mlir::Operation* op, mlir::OpBuilder::InsertPoint previous) override; + + private: + DBProgramGenerator* _generator {nullptr}; + }; + mlir::ModuleOp* _module {nullptr}; mlir::MLIRContext* _mlirCtxt {nullptr}; + ConstantRowsListener _constantRowsListener; mlir::OpBuilder _opBuilder; ExplainReport* _explain {nullptr}; @@ -1047,18 +1060,12 @@ class DBProgramGenerator { mlir::Value resolveColumnInScope(ColumnPredicate accept) const; mlir::Value resolveRowCarryingColumn(); + void layConstantOperandsOverRows(mlir::Operation* op); - // Pins a constant column to the relation whose rows it stands for. Left unpinned, a - // constant is aligned during lowering against whichever loop is innermost by then, - // which a collect in the same projection has already repointed at its own emit loop + // Pins a constant column to the relation whose rows it stands for, ahead of the op + // reading it, whose result type follows the column's mlir::Value alignConstantToDriver(mlir::Value column); - // An op that runs once per row, such as db.range, reads no row of its own when every - // operand is a constant, so nothing ties it to the relation it runs over: a pass erasing - // that relation's scan would leave it one row. Pins the constant operands of each such - // op in [@param first, @param last) to the driver instead. - void alignConstantOperandsToDriver(mlir::Block::iterator first, mlir::Block::iterator last); - // The column count(*) counts: the first variable bound to a column holding the rows // flowing past the insertion point. A scope of constants alone is the single row those // constants are, which is what count(*) counts there diff --git a/query/ir/dialect/db/DBOps.cpp b/query/ir/dialect/db/DBOps.cpp index c5bf4340a9..7660f77d7f 100644 --- a/query/ir/dialect/db/DBOps.cpp +++ b/query/ir/dialect/db/DBOps.cpp @@ -2382,3 +2382,27 @@ Speculation::Speculatability DivOp::getSpeculatability() { Speculation::Speculatability ModOp::getSpeculatability() { return speculatableWhen(isNumberColumn(getLhs()) && isSafeLiteralDivisor(getRhs(), false)); } + +bool mlir::db::readsConstantColumnsAlone(Operation* op, llvm::DenseMap& classified) { + const bool constantThroughOperands = op->hasTrait() && ::db::runsWithoutRaising(op); + const bool holdsOneValue = op->hasTrait() || op->hasTrait(); + const bool laysOutAConstant = isa(op); + if (constantThroughOperands || holdsOneValue || laysOutAConstant || op->hasTrait()) { + return false; + } + + bool readsAColumn = false; + for (const Value operand : op->getOperands()) { + if (!isa(operand.getType())) { + continue; + } + + if (!::db::yieldsConstantColumn(operand, classified)) { + return false; + } + + readsAColumn = true; + } + + return readsAColumn; +} diff --git a/query/ir/dialect/db/DBOps.h b/query/ir/dialect/db/DBOps.h index f484c4e245..b78ddc7d73 100644 --- a/query/ir/dialect/db/DBOps.h +++ b/query/ir/dialect/db/DBOps.h @@ -16,3 +16,11 @@ #define GET_OP_CLASSES #include "DBOps.h.inc" + +namespace mlir::db { + +// Whether @param op takes rows from its columns and every column it reads is a constant, +// against @param classified as yieldsConstantColumn reads it: such an op runs over one row +bool readsConstantColumnsAlone(Operation* op, llvm::DenseMap& classified); + +} diff --git a/query/ir/dialect/db/DBOps.td b/query/ir/dialect/db/DBOps.td index 506e892341..6a639fa62a 100644 --- a/query/ir/dialect/db/DBOps.td +++ b/query/ir/dialect/db/DBOps.td @@ -2909,6 +2909,12 @@ def BroadcastConstant : TuringOp<"broadcast_constant", [Pure]> { after the barrier - a WHERE cutting those rows, a count(*) tallying them, a pattern crossed with them - read a column that has them. + An op that reads constants alone runs over a single row, or over the rows of the region + it sits in where the op owning that region defines them. So a constant read where rows + are in flight comes through this op, and the relation it is read over is a use of the + driver like any other: `MATCH (n) RETURN 5` outputs a broadcast over the scan, and a + pass asking who reads the scan's rows finds it. + The driver is absent when no relation drives the scope, which is then the single row the constants are - `WITH 1 AS x WHERE x > 100`, whose one row the filter drops - so the value is laid out over that row: diff --git a/query/ir/dialect/db/DBPasses.cpp b/query/ir/dialect/db/DBPasses.cpp index 3b8c1dc813..a58ddb2755 100644 --- a/query/ir/dialect/db/DBPasses.cpp +++ b/query/ir/dialect/db/DBPasses.cpp @@ -5924,6 +5924,26 @@ void keepRowCarryingColumn(ValueRange columns, size_t firstIndex, llvm::SmallBit } } +// Whether the op's rows are those of its carry set: its results are its carried columns, or +// the operands it reads beside them are constants, which hold no rows of their own +bool takesItsRowsFromItsCarrySet(Operation* op, const CarrySetLayout& layout) { + if (layout._resultOffset == 0) { + return true; + } + + const size_t carriedEnd = op->getNumOperands() - layout._trailingOperandCount; + for (size_t operandIndex = 0; operandIndex < op->getNumOperands(); operandIndex++) { + const bool carried = operandIndex >= layout._operandOffset && operandIndex < carriedEnd; + const Value operand = op->getOperand(operandIndex); + + if (!carried && isa(operand.getType()) && !::db::yieldsConstantColumn(operand)) { + return false; + } + } + + return true; +} + void selectKeptCarriedColumns(Operation* op, const CarrySetLayout& layout, llvm::SmallVectorImpl& kept) { const size_t count = carriedCount(op, layout); @@ -5936,9 +5956,8 @@ void selectKeptCarriedColumns(Operation* op, const CarrySetLayout& layout, llvm: keepRequiredColumns(op, keep); - const bool standsForItsRows = layout._resultOffset == 0; - if (standsForItsRows) { - keepRowCarryingColumn(op->getOperands().drop_front(layout._operandOffset), 0, keep); + if (takesItsRowsFromItsCarrySet(op, layout)) { + keepRowCarryingColumn(op->getOperands().slice(layout._operandOffset, count), 0, keep); } for (size_t carriedIndex = 0; carriedIndex < count; carriedIndex++) { @@ -7559,6 +7578,53 @@ Value handedOnFrom(OpResult result) { return {}; } +bool drivesABroadcast(OpOperand& use) { + BroadcastConstant broadcast = dyn_cast(use.getOwner()); + + return broadcast && broadcast.getDriver() == use.get(); +} + +// Whether @param column is read once, setting aside the broadcasts it drives +bool readOnceBesidesBroadcasts(Value column) { + const size_t reads = llvm::count_if(column.getUses(), [](OpOperand& use) { + return !drivesABroadcast(use); + }); + + return reads == 1; +} + +// A broadcast reads only the rows of its driver, so one driven by a column a rewrite drops +// is driven by @param kept instead: a column of the same rows the rewrite keeps, or none +// where the dropped column held the rows alone +void redriveBroadcasts(Value dropped, Value kept, mlir::RewriterBase& rewriter) { + for (OpOperand& use : llvm::make_early_inc_range(dropped.getUses())) { + if (!drivesABroadcast(use)) { + continue; + } + + BroadcastConstant broadcast = cast(use.getOwner()); + rewriter.modifyOpInPlace(broadcast, [&broadcast, kept]() { + MutableOperandRange driver = broadcast.getDriverMutable(); + if (kept) { + driver.assign(kept); + } else { + driver.clear(); + } + }); + } +} + +// The column standing for the same rows as @param link once it is dropped from its op +Value siblingOf(OpResult link) { + for (const OpResult result : link.getOwner()->getResults()) { + if (result != link) { + return result; + } + } + + return {}; +} + bool isFactorOfTheScanAlone(Block* block) { const bool inAProduct = isa_and_nonnull(block->getParentOp()); @@ -7597,7 +7663,7 @@ bool matchFetchNodesChain(Value nodes, Value ids, FetchNodesChain& chain) { const bool readByTheEquality = user == chain._equality.getOperation(); const bool carriedByTheFilter = user == chain._filter.getOperation() && use.getOperandNumber() > 0; - if (!readByTheEquality && !carriedByTheFilter) { + if (!readByTheEquality && !carriedByTheFilter && !drivesABroadcast(use)) { return false; } } @@ -7616,7 +7682,7 @@ bool matchFetchNodesChain(Value nodes, Value ids, FetchNodesChain& chain) { return false; } - const bool readByTheLinkAlone = handedOn.hasOneUse(); + const bool readByTheLinkAlone = readOnceBesidesBroadcasts(handedOn); if (!readByTheLinkAlone || !keepsAColumnOnceDropped(link, handedOn)) { return false; } @@ -7751,6 +7817,17 @@ void fuseFetchNodes(FetchNodesChain& chain, mlir::RewriterBase& rewriter) { } } + for (const OpResult link : chain._links) { + redriveBroadcasts(link, siblingOf(link), rewriter); + } + + Block* const scanBlock = chain._scan->getBlock(); + const auto rowArgument = llvm::find_if(scanBlock->getArguments(), [](BlockArgument argument) { + return isa(argument.getType()); + }); + const Value rowsBeforeTheScan = rowArgument == scanBlock->args_end() ? Value() : Value(*rowArgument); + redriveBroadcasts(chain._scan->getResult(0), rowsBeforeTheScan, rewriter); + rewriter.setInsertionPoint(filter); FetchNodes fetch = rewriter.create(loc, fetchTypes, chain._ids, fetchCarried); diff --git a/query/ir/lowering/DBLowering.cpp b/query/ir/lowering/DBLowering.cpp index c5861c771d..d1bf09566e 100644 --- a/query/ir/lowering/DBLowering.cpp +++ b/query/ir/lowering/DBLowering.cpp @@ -899,6 +899,26 @@ nl::ChunkType groupAggregateResultChunkType(mlir::OpBuilder& builder, throw IRException("Unhandled group aggregate kind"); } +// Whether every column is a constant laid over the relation @param driver carries, one +// driver for them all or none +bool laysConstantsOverOneRelation(mlir::OperandRange columns, mlir::Value& driver) { + if (columns.empty()) { + return false; + } + + mlir::db::BroadcastConstant first = columns.front().getDefiningOp(); + if (!first) { + return false; + } + + driver = first.getDriver(); + + return llvm::all_of(columns, [driver](mlir::Value column) { + mlir::db::BroadcastConstant broadcast = column.getDefiningOp(); + return broadcast && broadcast.getDriver() == driver && yieldsConstantColumn(broadcast.getValue()); + }); +} + // The single nl.output that solely consumes every result in the range, or a null // op if any result has more than one use, a non-nl.output user, or a different // output than its siblings. Read the direction as "one shared output user => the @@ -1118,7 +1138,7 @@ mlir::func::FuncOp DBLowering::lower(mlir::func::FuncOp dbFunction, mlir::Module _constantColumns.clear(); _rootBlock = _entryBlock; _innermostLoopBody = nullptr; - _innermostCardinality = mlir::Value(); + _regionRows = mlir::Value(); _limitHandles.clear(); _loopLimitHandle.clear(); _optionalMatchLimits.clear(); @@ -1239,7 +1259,7 @@ void DBLowering::lowerOverReachingRows(mlir::Operation& operation) { for (const mlir::Value operand : operation.getOperands()) { const mlir::Value chunk = mapValue(operand); constantChunks.emplace_back(operand, chunk); - _valueMap[operand] = rowAlignedChunk(chunk, _innermostCardinality); + _valueMap[operand] = rowAlignedChunk(chunk, _regionRows); } lowerOperationByKind(operation); @@ -1250,12 +1270,12 @@ void DBLowering::lowerOverReachingRows(mlir::Operation& operation) { } void DBLowering::lowerOnceOverConstants(mlir::Operation& operation) { - const mlir::Value previousInnermostCardinality = _innermostCardinality; - _innermostCardinality = {}; + const mlir::Value previousRegionRows = _regionRows; + _regionRows = {}; lowerOperationByKind(operation); - _innermostCardinality = previousInnermostCardinality; + _regionRows = previousRegionRows; for (const mlir::Value result : operation.getResults()) { const mlir::Value row = mapValue(result); @@ -1845,9 +1865,9 @@ mlir::Operation* DBLowering::lowerElementBody(mlir::Block& dbBody, // The rows the body computes over are the elements of one row's list, so a constant it // reads is laid out over those rather than over the rows the op is read on mlir::Block* const previousInnermostLoopBody = _innermostLoopBody; - const mlir::Value previousInnermostCardinality = _innermostCardinality; + const mlir::Value previousRegionRows = _regionRows; _innermostLoopBody = nlBody; - _innermostCardinality = elementChunk; + _regionRows = elementChunk; for (mlir::Operation& operation : dbBody) { mlir::db::ComprehensionYield yield = mlir::dyn_cast(operation); @@ -1861,7 +1881,7 @@ mlir::Operation* DBLowering::lowerElementBody(mlir::Block& dbBody, } _innermostLoopBody = previousInnermostLoopBody; - _innermostCardinality = previousInnermostCardinality; + _regionRows = previousRegionRows; // ComprehensionYield::verify guarantees both, so an empty one here means unverified // IR - the defensive backstop lowerOptionalMatch keeps too. @@ -2028,9 +2048,9 @@ mlir::Operation* DBLowering::lowerReduceBody(mlir::Block& dbBody, mlir::Block* n } mlir::Block* const previousInnermostLoopBody = _innermostLoopBody; - const mlir::Value previousInnermostCardinality = _innermostCardinality; + const mlir::Value previousRegionRows = _regionRows; _innermostLoopBody = nlBody; - _innermostCardinality = elementChunk; + _regionRows = elementChunk; for (mlir::Operation& operation : dbBody) { mlir::db::ReduceYield yield = mlir::dyn_cast(operation); @@ -2042,7 +2062,7 @@ mlir::Operation* DBLowering::lowerReduceBody(mlir::Block& dbBody, mlir::Block* n } _innermostLoopBody = previousInnermostLoopBody; - _innermostCardinality = previousInnermostCardinality; + _regionRows = previousRegionRows; valueChunk = rowAlignedChunk(valueChunk, elementChunk); @@ -2130,10 +2150,10 @@ void DBLowering::lowerPatternComprehension(mlir::db::PatternComprehension compre // this step; the caller's root and innermost loop are restored once it is lowered. mlir::Block* const previousRoot = _rootBlock; mlir::Block* const previousInnermostLoopBody = _innermostLoopBody; - const mlir::Value previousInnermostCardinality = _innermostCardinality; + const mlir::Value previousRegionRows = _regionRows; _rootBlock = stepBlock; _innermostLoopBody = nullptr; - _innermostCardinality = mlir::Value(); + _regionRows = mlir::Value(); mlir::Value rowTagChunk; mlir::Value valueChunk; @@ -2160,7 +2180,7 @@ void DBLowering::lowerPatternComprehension(mlir::db::PatternComprehension compre // The matches the value is read against are the ones the tag counts, since a WHERE has // cut both together - valueChunk = rowListsValueChunk(valueChunk, rowTagChunk ? rowTagChunk : _innermostCardinality); + valueChunk = rowListsValueChunk(valueChunk, rowTagChunk ? rowTagChunk : _regionRows); llvm::SmallVector stagedChunks {valueChunk}; if (rowTagChunk) { @@ -2172,7 +2192,7 @@ void DBLowering::lowerPatternComprehension(mlir::db::PatternComprehension compre _rootBlock = previousRoot; _innermostLoopBody = previousInnermostLoopBody; - _innermostCardinality = previousInnermostCardinality; + _regionRows = previousRegionRows; // Every row gets a list, the empty one where the pattern matched nothing, so the // lists are a container chunk rather than the nullable one a list comprehension builds @@ -2997,10 +3017,10 @@ void DBLowering::lowerOptionalMatch(mlir::db::OptionalMatch optionalMatch) { // this step; the caller's root and innermost loop are restored once it is lowered. mlir::Block* const previousRoot = _rootBlock; mlir::Block* const previousInnermostLoopBody = _innermostLoopBody; - const mlir::Value previousInnermostCardinality = _innermostCardinality; + const mlir::Value previousRegionRows = _regionRows; _rootBlock = stepBlock; _innermostLoopBody = nullptr; - _innermostCardinality = mlir::Value(); + _regionRows = mlir::Value(); if (patternLimit) { mlir::db::OptionalYield patternYield = mlir::cast(patternBlock.getTerminator()); @@ -3059,7 +3079,7 @@ void DBLowering::lowerOptionalMatch(mlir::db::OptionalMatch optionalMatch) { _rootBlock = previousRoot; _innermostLoopBody = previousInnermostLoopBody; - _innermostCardinality = previousInnermostCardinality; + _regionRows = previousRegionRows; llvm::SmallVector chunkTypes; for (const mlir::Value chunk : matchedChunks) { @@ -3093,7 +3113,7 @@ void DBLowering::lowerCallSubquery(mlir::db::CallSubquery call) { } mlir::Block* const previousInnermostLoopBody = _innermostLoopBody; - const mlir::Value previousInnermostCardinality = _innermostCardinality; + const mlir::Value previousRegionRows = _regionRows; llvm::SmallVector yieldedChunks; mlir::Value yieldedTag; @@ -3107,7 +3127,7 @@ void DBLowering::lowerCallSubquery(mlir::db::CallSubquery call) { const bool bodyOpenedALoop = _innermostLoopBody != nullptr; if (!returning || !bodyOpenedALoop) { _innermostLoopBody = previousInnermostLoopBody; - _innermostCardinality = previousInnermostCardinality; + _regionRows = previousRegionRows; } if (!returning) { @@ -3137,7 +3157,7 @@ void DBLowering::lowerSubqueryBody(mlir::db::CallSubquery call, mlir::Block* const previousRoot = _rootBlock; _rootBlock = stepBlock; _innermostLoopBody = nullptr; - _innermostCardinality = mlir::Value(); + _regionRows = mlir::Value(); for (mlir::Operation& operation : bodyBlock) { mlir::db::SubqueryYield yield = mlir::dyn_cast(operation); @@ -3165,7 +3185,7 @@ void DBLowering::lowerSubqueryPerRow(mlir::db::CallSubquery call, const bool unit = call.getUnit(); mlir::Block* const previousInnermostLoopBody = _innermostLoopBody; - const mlir::Value previousInnermostCardinality = _innermostCardinality; + const mlir::Value previousRegionRows = _regionRows; llvm::SmallVector yieldedChunks; @@ -3179,7 +3199,7 @@ void DBLowering::lowerSubqueryPerRow(mlir::db::CallSubquery call, if (unit || !_innermostLoopBody) { _innermostLoopBody = previousInnermostLoopBody; - _innermostCardinality = previousInnermostCardinality; + _regionRows = previousRegionRows; setInsertionInto(stepBlock); } @@ -3217,7 +3237,7 @@ void DBLowering::lowerSubqueryPerRow(mlir::db::CallSubquery call, // relation for what follows to walk: the step's own rows are in flight still if (unit) { _innermostLoopBody = previousInnermostLoopBody; - _innermostCardinality = previousInnermostCardinality; + _regionRows = previousRegionRows; setInsertionInto(stepBlock); return; @@ -3432,10 +3452,10 @@ void DBLowering::lowerSubqueryExpression(SubqueryOp subquery) { // this step; the caller's root and innermost loop are restored once it is lowered. mlir::Block* const previousRoot = _rootBlock; mlir::Block* const previousInnermostLoopBody = _innermostLoopBody; - const mlir::Value previousInnermostCardinality = _innermostCardinality; + const mlir::Value previousRegionRows = _regionRows; _rootBlock = bodyRoot; _innermostLoopBody = nullptr; - _innermostCardinality = mlir::Value(); + _regionRows = mlir::Value(); for (mlir::Operation& operation : bodyBlock.without_terminator()) { lowerOperation(operation); @@ -3491,7 +3511,7 @@ void DBLowering::lowerSubqueryExpression(SubqueryOp subquery) { _rootBlock = previousRoot; _innermostLoopBody = previousInnermostLoopBody; - _innermostCardinality = previousInnermostCardinality; + _regionRows = previousRegionRows; setInsertionInto(bodyRoot); @@ -3531,7 +3551,7 @@ void DBLowering::lowerSubqueryExpression(SubqueryOp subquery) { } _innermostLoopBody = bodyRoot; - _innermostCardinality = stepChunks.front(); + _regionRows = stepChunks.front(); } void DBLowering::collectReadPastRowLoop(mlir::Operation* subquery, @@ -3673,10 +3693,10 @@ mlir::Block* DBLowering::lowerFactor(mlir::Region& factor, // save and restore the caller's so nested or sibling products are unaffected. mlir::Block* const previousRoot = _rootBlock; mlir::Block* const previousInnermostLoopBody = _innermostLoopBody; - const mlir::Value previousInnermostCardinality = _innermostCardinality; + const mlir::Value previousRegionRows = _regionRows; _rootBlock = rootBlock; _innermostLoopBody = nullptr; - _innermostCardinality = mlir::Value(); + _regionRows = mlir::Value(); // A factor is one self-contained block ending in a db.yield. Lower each op // as at top level; the yield names the columns this factor contributes, so @@ -3707,7 +3727,7 @@ mlir::Block* DBLowering::lowerFactor(mlir::Region& factor, _rootBlock = previousRoot; _innermostLoopBody = previousInnermostLoopBody; - _innermostCardinality = previousInnermostCardinality; + _regionRows = previousRegionRows; return factorBody; } @@ -3829,7 +3849,6 @@ void DBLowering::lowerLimit(mlir::db::Limit limit) { _valueMap[dbResults[resultIndex]] = truncatedChunks[resultIndex]; } - followCardinalityThrough(chunks, truncatedChunks); } void DBLowering::lowerSkip(mlir::db::Skip skip) { @@ -3883,7 +3902,6 @@ void DBLowering::lowerSkip(mlir::db::Skip skip) { _valueMap[dbResults[resultIndex]] = truncatedChunks[resultIndex]; } - followCardinalityThrough(chunks, truncatedChunks); } void DBLowering::detectTopKFusion(mlir::func::FuncOp dbFunction) { @@ -4006,14 +4024,14 @@ void DBLowering::lowerUnion(mlir::db::Union unionOp) { mlir::Block* const previousRoot = _rootBlock; mlir::Block* const previousInnermostLoopBody = _innermostLoopBody; - const mlir::Value previousInnermostCardinality = _innermostCardinality; + const mlir::Value previousRegionRows = _regionRows; llvm::SmallVector branchOutputs; for (mlir::Region& branch : branches) { _rootBlock = _entryBlock; _innermostLoopBody = nullptr; - _innermostCardinality = mlir::Value(); + _regionRows = mlir::Value(); mlir::db::Output branchOutput = mlir::cast(branch.front().back()); @@ -4045,7 +4063,7 @@ void DBLowering::lowerUnion(mlir::db::Union unionOp) { _rootBlock = previousRoot; _innermostLoopBody = previousInnermostLoopBody; - _innermostCardinality = previousInnermostCardinality; + _regionRows = previousRegionRows; } // The branches of a union inside a CALL body feed the rest of the body, which is lowered @@ -4075,7 +4093,7 @@ void DBLowering::lowerConditional(mlir::db::Conditional conditional) { mlir::Block* const root = _rootBlock; mlir::Block* const previousInnermostLoopBody = _innermostLoopBody; - const mlir::Value previousInnermostCardinality = _innermostCardinality; + const mlir::Value previousRegionRows = _regionRows; mlir::Block* const previousConstantComputationBlock = _constantComputationBlock; const bool returning = conditional.getNumResults() > 0; @@ -4146,7 +4164,7 @@ void DBLowering::lowerConditional(mlir::db::Conditional conditional) { if (!returning) { _innermostLoopBody = previousInnermostLoopBody; - _innermostCardinality = previousInnermostCardinality; + _regionRows = previousRegionRows; setInsertionInto(root); return; @@ -4189,7 +4207,7 @@ void DBLowering::lowerCollectedBranch(mlir::Region& branch, const mlir::Location loc = _builder.getUnknownLoc(); _innermostLoopBody = nullptr; - _innermostCardinality = mlir::Value(); + _regionRows = mlir::Value(); mlir::db::Yield branchYield = mlir::cast(branch.front().back()); @@ -4471,9 +4489,9 @@ void DBLowering::lowerRemoveDuplicates(mlir::db::RemoveDuplicates distinct) { void DBLowering::lowerCount(mlir::db::Count count) { // The nl chunk the counted column lowered to; the update reads its per-step - // non-null row count. A constant column counts the rows it stands for, not the - // one row it is, so it is laid out over the driving relation first. - const mlir::Value inputChunk = rowAlignedChunk(mapValue(count.getInput()), _innermostCardinality); + // non-null row count. A constant column holds no rows, so it is laid out over the + // rows of the region first. + const mlir::Value inputChunk = rowAlignedChunk(mapValue(count.getInput()), _regionRows); const mlir::Location loc = _builder.getUnknownLoc(); @@ -4524,10 +4542,6 @@ void DBLowering::lowerCount(mlir::db::Count count) { // into a function-scope nl.output reading it - the block that holds the chunk is // the entry block, so lowerOutput places nl.output there. _valueMap[count.getResult()] = result.getResult(); - - // The count collapses the rows in flight to its one row, so from here on that row is - // what sizes a projection of constants alone - _innermostCardinality = result.getResult(); } void DBLowering::lowerCountScanRows(mlir::db::CountScanRows countScanRows) { @@ -4542,14 +4556,13 @@ void DBLowering::lowerCountScanRows(mlir::db::CountScanRows countScanRows) { countScanRows.getPropertyScanAttr()); _valueMap[countScanRows.getResult()] = rows.getResult(); - _innermostCardinality = rows.getResult(); } void DBLowering::lowerAggregate(mlir::Value input, mlir::Value result, storage::AggregateKind kind, bool distinct) { // The nl chunk the aggregated column lowered to; the update folds its per-step - // non-null values into the accumulator. A constant column is reduced over the - // rows it stands for, so it is laid out over the driving relation first. - const mlir::Value alignedChunk = rowAlignedChunk(mapValue(input), _innermostCardinality); + // non-null values into the accumulator. A constant column holds no rows, so it is + // laid out over the rows of the region first. + const mlir::Value alignedChunk = rowAlignedChunk(mapValue(input), _regionRows); const mlir::Value inputChunk = reducedValueChunk(alignedChunk); @@ -4611,7 +4624,6 @@ void DBLowering::lowerAggregate(mlir::Value input, mlir::Value result, storage:: // The db aggregate's result maps to that chunk, so the db.output that follows // lowers into a function-scope nl.output reading it, exactly as db.count does. _valueMap[result] = aggregateResult.getResult(); - _innermostCardinality = aggregateResult.getResult(); } void DBLowering::lowerGroupAggregate(mlir::db::GroupAggregate groupAggregate) { @@ -4637,7 +4649,7 @@ void DBLowering::lowerGroupAggregate(mlir::db::GroupAggregate groupAggregate) { // The group assignment reads one row of every key per row of the step, which a constant // key - one value standing for every row - has none of. The keys are laid out first, so // that the aggregate inputs behind them are laid out over a key chunk holding rows. - rowAlignBufferedChunks(llvm::MutableArrayRef(chunks).take_front(keyCount)); + rowAlignGroupingKeys(chunks, keyCount); // A constant aggregate input is reduced over the rows of the group it falls in, // not over the single row it is, so it is laid out over the chunk the grouping @@ -4738,13 +4750,13 @@ void DBLowering::lowerCollect(mlir::db::Collect collect) { // The grouping keys are laid out first, as lowerGroupAggregate lays its own: the group // assignment reads a row of every key per row of the step, which a constant key has none of. - rowAlignBufferedChunks(llvm::MutableArrayRef(chunks).take_front(keyCount)); + rowAlignGroupingKeys(chunks, keyCount); // A constant collected column is folded over the rows of the group it falls in, not // over the single row it is, so it is laid out over the chunk the grouping keys are - // read from - the driving relation when the collect has no keys. The aggregate + // read from - the rows of the region when the collect has no keys. The aggregate // inputs beside it are laid out the same way, as lowerGroupAggregate lays its own. - const mlir::Value cardinality = keyCount > 0 ? chunks.front() : _innermostCardinality; + const mlir::Value cardinality = keyCount > 0 ? chunks.front() : _regionRows; for (size_t inputIndex = keyCount; inputIndex < chunks.size(); inputIndex++) { chunks[inputIndex] = rowAlignedChunk(chunks[inputIndex], cardinality); } @@ -5678,10 +5690,6 @@ void DBLowering::lowerMerge(mlir::db::Merge merge) { for (size_t index = 0; index < dbResults.size(); index++) { _valueMap[dbResults[index]] = nlResults[index]; } - - // The merge grows or shrinks the rows in flight, so from here on what sizes a - // projection of constants alone is its own node chunk - _innermostCardinality = nlResults.front(); } void DBLowering::lowerSetNodeProperty(mlir::db::SetNodeProperty setNodeProperty) { @@ -6256,8 +6264,8 @@ mlir::Value DBLowering::lowerCaseRegion(mlir::Region& region, } mlir::Value DBLowering::lowerBlockOverRows(mlir::Block& block, mlir::Value cardinality) { - const mlir::Value previousInnermostCardinality = _innermostCardinality; - _innermostCardinality = cardinality; + const mlir::Value previousRegionRows = _regionRows; + _regionRows = cardinality; mlir::Value value; for (mlir::Operation& operation : block) { @@ -6268,7 +6276,7 @@ mlir::Value DBLowering::lowerBlockOverRows(mlir::Block& block, mlir::Value cardi } } - _innermostCardinality = previousInnermostCardinality; + _regionRows = previousRegionRows; return rowAlignedChunk(value, cardinality); } @@ -6464,7 +6472,6 @@ void DBLowering::lowerFilter(mlir::db::FilterOp filter) { _valueMap[filteredCol] = outputChunk; } - followCardinalityThrough(columnChunks, nlFilter.getResults()); } void DBLowering::lowerFetchNodes(mlir::db::FetchNodes fetch) { @@ -6496,7 +6503,6 @@ void DBLowering::lowerFetchNodes(mlir::db::FetchNodes fetch) { _valueMap[results[resultIndex]] = nlFetch->getResult(resultIndex); } - followCardinalityThrough(inputChunks, nlFetch->getResults()); } mlir::Value DBLowering::deepestBoundChunk(llvm::ArrayRef chunks) { @@ -6537,7 +6543,41 @@ mlir::Block* DBLowering::deeperOfBlocks(mlir::Block* first, mlir::Block* second) throw IRException("db operands read together must be bound in the same loop"); } +// The constants go out against the driver's row count rather than as columns laid out +// over its rows first +nl::Output DBLowering::lowerConstantOutput(mlir::db::Output output, mlir::Value driver) { + llvm::SmallVector constants; + for (const mlir::Value column : output.getColumns()) { + mlir::db::BroadcastConstant broadcast = column.getDefiningOp(); + constants.push_back(mapValue(broadcast.getValue())); + + mlir::Operation* const layout = mapValue(broadcast.getResult()).getDefiningOp(); + if (broadcast->hasOneUse() && layout && layout->use_empty()) { + layout->erase(); + } + } + + const mlir::Value cardinality = driver ? mapValue(driver) : mlir::Value(); + setInsertionInto(cardinality ? ownerBlock(cardinality) : ownerBlock(constants.front())); + + return _builder.create(_builder.getUnknownLoc(), + constants, + mlir::Value(), + mlir::Value(), + cardinality, + output.getColumnNamesAttr()); +} + nl::Output DBLowering::lowerOutput(mlir::db::Output output) { + // A union branch's columns meet the other branches' and take one type with them, so + // they stay laid out + const bool feedsAUnion = mlir::isa(output->getParentOp()); + + mlir::Value driver; + if (!feedsAUnion && laysConstantsOverOneRelation(output.getColumns(), driver)) { + return lowerConstantOutput(output, driver); + } + llvm::SmallVector columns; for (const mlir::Value column : output.getColumns()) { columns.push_back(mapValue(column)); @@ -6555,9 +6595,9 @@ nl::Output DBLowering::lowerOutput(mlir::db::Output output) { // dropping the copy. // If there is a column which is produced by a block which is not the entry block - // (e.g. a loop block), then set the anchor to be that block. Otherwise, we have no - // loops, i.e. we are in a MATCH (n) RETURN 5 case, where the output can just be in - // the entry block since it is independent of any loop (over n in this case). + // (e.g. a loop block), then set the anchor to be that block. Otherwise the columns are + // constants alone, which stand for the single row of the top level, and the output goes + // in the entry block. // Getting the owner block of the returned column is sufficient because in nested // loops, each Cypher variable is redefined each op due to the carry set implictly // filtering. Otherwise we would need to check for the *deepest* block of all @@ -6571,29 +6611,13 @@ nl::Output DBLowering::lowerOutput(mlir::db::Output output) { } } - // Custom cardinality calculation for e.g. MATCH (n) RETURN 5 - mlir::Value cardinality; - // Output would not normally be in a loop body, but toplevel - const bool returningNonLooped = anchorBlock == _entryBlock; - if (returningNonLooped && _innermostLoopBody) { - // but if we have all constants, then it need be moved to the inner most loop to - // match cardinality. An expression over constants alone is one of them: it is - // bound where its operands are, above the loop whose rows it is projected over - const bool allConstants = llvm::all_of(columns, [this](mlir::Value column) { return yieldsConstantColumn(column, _constantColumns); }); - - if (allConstants) { - anchorBlock = _innermostLoopBody; - cardinality = _innermostCardinality; - } - } - setInsertionInto(anchorBlock); return _builder.create(_builder.getUnknownLoc(), columns, mlir::Value(), mlir::Value(), - cardinality, + mlir::Value(), output.getColumnNamesAttr()); } @@ -6815,11 +6839,6 @@ void DBLowering::buildLoopForSource(mlir::Value iterator, mlir::Operation* dbOp) // nest in dataflow order), so a cross product nests at this body. _innermostLoopBody = loopBody; - // A loop over a call that binds no return value has no variable, so there is no chunk - // here for a constant projection to be sized against. - _innermostCardinality = loopBody->getNumArguments() > 0 ? loopBody->getArgument(0) - : mlir::Value(); - const mlir::ResultRange dbResults = dbOp->getResults(); for (size_t resultIndex = 0; resultIndex < dbResults.size(); resultIndex++) { _valueMap[dbResults[resultIndex]] = loopBody->getArgument(static_cast(resultIndex)); @@ -6918,23 +6937,6 @@ mlir::Value DBLowering::mapValue(mlir::Value dbValue) const { return slotIt->second; } -void DBLowering::followCardinalityThrough(mlir::ValueRange inputChunks, mlir::ValueRange resultChunks) { - if (!_innermostCardinality) { - return; - } - - // Any chunk of the driver's relation stands for it: a cut trimmed down to a column other - // than the driver itself still narrows the relation the driver counts. - for (size_t chunkIndex = 0; chunkIndex < inputChunks.size(); chunkIndex++) { - if (!rowAlignedWith(inputChunks[chunkIndex], _innermostCardinality)) { - continue; - } - - _innermostCardinality = resultChunks[chunkIndex]; - return; - } -} - void DBLowering::rowAlignFactorChunks(llvm::SmallVectorImpl& chunks) { mlir::Value cardinality; for (const mlir::Value chunk : chunks) { @@ -6954,7 +6956,7 @@ void DBLowering::rowAlignFactorChunks(llvm::SmallVectorImpl& chunks // constants are laid out first: over the relation driving the step, or - where none drives // it - over the single row the projection is. void DBLowering::rowAlignBufferedChunks(llvm::MutableArrayRef chunks) { - mlir::Value cardinality = _innermostCardinality; + mlir::Value cardinality = _regionRows; for (const mlir::Value chunk : chunks) { if (!yieldsConstantColumn(chunk, _constantColumns)) { cardinality = chunk; @@ -6967,15 +6969,26 @@ void DBLowering::rowAlignBufferedChunks(llvm::MutableArrayRef chunk } } +// The grouping keys are laid out over the rows the op reads: those of a key that carries +// rows, else those of a column it aggregates +void DBLowering::rowAlignGroupingKeys(llvm::MutableArrayRef chunks, size_t keyCount) { + const mlir::Value rows = cardinalityDriver(chunks); + + for (mlir::Value& key : chunks.take_front(keyCount)) { + key = rowAlignedChunk(key, rows); + } +} + void DBLowering::rowAlignCutChunks(llvm::SmallVectorImpl& chunks) { - // With no relation driving the projection a constant is the single row that - // projection is, and the cut charges it as it stands. - if (!_innermostCardinality) { + // Constants alone at the top level are the single row they stand for, and the cut + // charges it as it stands + const mlir::Value cardinality = cardinalityDriver(chunks); + if (!cardinality) { return; } for (mlir::Value& chunk : chunks) { - chunk = rowAlignedChunk(chunk, _innermostCardinality); + chunk = rowAlignedChunk(chunk, cardinality); } } @@ -7131,7 +7144,7 @@ mlir::Value DBLowering::cardinalityDriver(llvm::ArrayRef chunks) { } } - return _innermostCardinality; + return _regionRows; } mlir::Value DBLowering::rowAlignedChunk(mlir::Value chunk, mlir::Value cardinality) { diff --git a/query/ir/lowering/DBLowering.h b/query/ir/lowering/DBLowering.h index f7d01c5434..ff5bb18b54 100644 --- a/query/ir/lowering/DBLowering.h +++ b/query/ir/lowering/DBLowering.h @@ -181,8 +181,9 @@ class DBLowering { // not left in the loop body - just a record of which block is innermost. mlir::Block* _innermostLoopBody {nullptr}; - // Defines the cardinality of the innermost loop, used for projection - mlir::Value _innermostCardinality; + // The rows of the region being lowered, as the op owning it defines them: what an op + // reading constants alone runs over. Null at the top level, where that is a single row. + mlir::Value _regionRows; // Each db.limit gets its own hoisted nl.limit handle - N independent budgets. // A pre-scan in lower() finds every db.limit before any loop is built, so the @@ -696,6 +697,7 @@ class DBLowering { // Emits the nl.output the sink reads this result off, and hands it back: a union // reaches for it again once every branch has said what its columns hold mlir::nl::Output lowerOutput(mlir::db::Output output); + mlir::nl::Output lowerConstantOutput(mlir::db::Output output, mlir::Value driver); // Lower one factor region of a db.cross_product or a db.hash_join into a loop // nest rooted at rootBlock, collecting its db.yield columns (mapped to their nl @@ -813,6 +815,7 @@ class DBLowering { // holds one value standing for every row of the step - so the loop's rows are the // count a cut over constants alone charges, and the count every cut chained after it // reads in turn. + void rowAlignGroupingKeys(llvm::MutableArrayRef chunks, size_t keyCount); void rowAlignCutChunks(llvm::SmallVectorImpl& chunks); // Lays the constants among a hash join factor's yielded chunks over the rows that @@ -831,7 +834,6 @@ class DBLowering { // @param resultChunks. The driver stands for the relation a projection of constants // alone is a projection of, so leaving it on the wider chunk would size that // emission by the rows the step dropped. - void followCardinalityThrough(mlir::ValueRange inputChunks, mlir::ValueRange resultChunks); // The chunk a step that walks rows is handed for @param chunk: itself when it // carries rows, and its value laid out over the rows of @param cardinality - the diff --git a/test/query/ir/call/MLIRCallProcedureTest.cpp b/test/query/ir/call/MLIRCallProcedureTest.cpp index 5d97f6aa80..0bcfdeaac1 100644 --- a/test/query/ir/call/MLIRCallProcedureTest.cpp +++ b/test/query/ir/call/MLIRCallProcedureTest.cpp @@ -1092,7 +1092,8 @@ constexpr const char* labelsConstantProgram = R"mlir( func.func @main() { %label = db.call_procedure("db.labels", {}, {}) yields ["label"] : () -> !db.column %c = db.constant(5 : i64) - db.output(%c) : !db.column + %x = db.broadcast_constant(%c, %label) : (!db.column, !db.column) -> !db.column + db.output(%x) : !db.column return } )mlir"; @@ -1696,9 +1697,9 @@ TEST_F(MLIRCallProcedureTest, constantProjectionOverDriveLoopKeepsCardinality) { const FrozenCommitTx transaction = graph->openTransaction(); const GraphReader reader = transaction.readGraph(); - // Projection preserves cardinality, so the five labels give five rows of 5. The - // drive loop binds no node-ID chunk, so this only lowers because nl.output's - // cardinality driver takes a chunk of any element type. + // The constant is laid over the five labels, so five rows of 5. The drive loop + // binds no node-ID chunk, so this only lowers because nl.output's cardinality + // driver takes a chunk of any element type. ConstSink sink; runLoweredProgram(labelsConstantProgram, reader.getView(), sink, /*chunkSize=*/2); diff --git a/test/query/ir/engine/DBLoweringTest.cpp b/test/query/ir/engine/DBLoweringTest.cpp index fa8a5cf76b..1cf2389747 100644 --- a/test/query/ir/engine/DBLoweringTest.cpp +++ b/test/query/ir/engine/DBLoweringTest.cpp @@ -678,13 +678,14 @@ func.func @main() { )mlir"; // MATCH (n) RETURN 5: a constant projected over a scan, with no per-row column in -// the projection. Projection preserves cardinality, so the result is one row of 5 -// per node - not a single row. +// the projection. The constant is laid over the scan's rows, so the result is one row +// of 5 per node - not a single row. constexpr const char* constantOverMatchProgram = R"mlir( func.func @main() { %n = db.scan_nodes() : !db.column %c = db.constant(5 : i64) - db.output(%c) : !db.column + %x = db.broadcast_constant(%c, %n) : (!db.column, !db.column) -> !db.column + db.output(%x) : !db.column return } )mlir"; @@ -696,7 +697,8 @@ func.func @main() { %x = db.constant(40 : i64) %y = db.constant(2 : i64) %s = db.add %x, %y : (!db.column, !db.column) -> !db.column - db.output(%s) : !db.column + %r = db.broadcast_constant(%s, %n) : (!db.column, !db.column) -> !db.column + db.output(%r) : !db.column return } )mlir"; @@ -706,8 +708,9 @@ constexpr const char* constantLimitOverMatchProgram = R"mlir( func.func @main() { %n = db.scan_nodes() : !db.column %c = db.constant(5 : i64) - %lc = db.limit(%c) count 3 : (!db.column) -> !db.column - db.output(%lc) : !db.column + %x = db.broadcast_constant(%c, %n) : (!db.column, !db.column) -> !db.column + %lc = db.limit(%x) count 3 : (!db.column) -> !db.column + db.output(%lc) : !db.column return } )mlir"; @@ -757,7 +760,8 @@ func.func @main() { db.yield %b : !db.column } %c = db.constant(5 : i64) - db.output(%c) : !db.column + %x = db.broadcast_constant(%c, %0#0) : (!db.column, !db.column) -> !db.column + db.output(%x) : !db.column return } )mlir"; @@ -779,7 +783,8 @@ func.func @main() { db.yield %c : !db.column } %d = db.constant(5 : i64) - db.output(%d) : !db.column + %x = db.broadcast_constant(%d, %0#0) : (!db.column, !db.column) -> !db.column + db.output(%x) : !db.column return } )mlir"; @@ -2875,9 +2880,8 @@ TEST_F(DBLoweringTest, constantProjectionOverMatchHasRelationCardinality) { CollectingConstSink sink; runLoweredProgram(constantOverMatchProgram, reader.getView(), sink); - // RETURN 5 over the diamond's four nodes: projection preserves cardinality, so - // four rows of 5 - not one. (Before the fix the output floated to function scope - // and emitted a single row.) + // RETURN 5 over the diamond's four nodes: the constant is laid over the scan's + // rows, so four rows of 5 - not one. const std::vector> expected {{5}, {5}, {5}, {5}}; EXPECT_EQ(sink.rows(), expected); } @@ -2890,8 +2894,8 @@ TEST_F(DBLoweringTest, constantExpressionProjectionOverMatchHasRelationCardinali CollectingConstSink sink; runLoweredProgram(constantExpressionOverMatchProgram, reader.getView(), sink); - // An expression over constants alone is a constant column too, so it is projected - // over the relation's rows as a literal is: four rows of 42, not the single row the + // An expression over constants alone is a constant column too, so it is laid over + // the relation's rows as a literal is: four rows of 42, not the single row the // expression is bound as above the loop. const std::vector> expected {{42}, {42}, {42}, {42}}; EXPECT_EQ(sink.rows(), expected); diff --git a/test/query/ir/projection/CMakeLists.txt b/test/query/ir/projection/CMakeLists.txt index d13615729b..5820e5d8be 100644 --- a/test/query/ir/projection/CMakeLists.txt +++ b/test/query/ir/projection/CMakeLists.txt @@ -60,3 +60,5 @@ target_include_directories(test_query_ir_bool_ordering PRIVATE ${IR_TEST_DIR}) add_call_v3_test(test_query_ir_nan_ordering NaNOrderingTest.cpp) add_call_v3_test(test_query_ir_constant_projection_rows ConstantProjectionRowsTest.cpp) + +add_call_v3_test(test_query_ir_constant_rows ConstantRowsTest.cpp) diff --git a/test/query/ir/projection/ConstantRowsTest.cpp b/test/query/ir/projection/ConstantRowsTest.cpp new file mode 100644 index 0000000000..f693b0e6eb --- /dev/null +++ b/test/query/ir/projection/ConstantRowsTest.cpp @@ -0,0 +1,89 @@ +#include + +#include +#include +#include + +#include "CallV3Test.h" +#include "StringRowSink.h" + +using namespace turing::test; + +namespace { + +using Rows = std::vector; + +bool contains(std::string_view text, std::string_view needle) { + return text.find(needle) != std::string_view::npos; +} + +} + +// An op reading constants alone runs over a single row, so codegen lays every constant an +// op reads alone over the rows in flight with db.broadcast_constant: the rows a query emits +// follow the use-def chains a pass reads +class ConstantRowsTest : public CallV3Test { +protected: + void expectRowCount(std::string_view query, size_t expected) { + StringRowSink sink; + runQuery(query, sink); + + EXPECT_EQ(sink.getRows().size(), expected) << query; + } + + void expectRows(std::string_view query, const Rows& expected) { + StringRowSink sink; + runQuery(query, sink); + + Rows rows; + sink.sortedRows(rows); + + EXPECT_EQ(rows, expected) << query; + } + + void programOf(std::string_view query, std::string& program) { + StringRowSink sink; + runQuery(std::string("EXPLAIN (db) ") + std::string(query), sink); + + ASSERT_EQ(sink.getRows().size(), 1u); + program = sink.getRows().front().back(); + } +}; + +TEST_F(ConstantRowsTest, returnsOneRowPerPath) { + expectRowCount("MATCH (a:Person)-[e*1..2]->(b) RETURN 1", 25); + expectRowCount("MATCH (a:Person)-[*1..2]->(b) RETURN null", 25); + expectRowCount("MATCH (a:Person {name: 'Remy'})-[*2..2]->(b) RETURN 1", 4); + expectRowCount("MATCH p = (a:Person {name: 'Remy'})-[*2..2]->(b) " + "RETURN toStringOrNull(relationships(p)), toStringOrNull(nodes(p))", + 4); +} + +TEST_F(ConstantRowsTest, returnsOneRowPerPathPastAHopOrACut) { + expectRowCount("MATCH (a:Person)-[e*1..2]->(b)-->(c) RETURN 1", 20); + expectRowCount("MATCH (a:Person)-[e*1..2]->(b) WHERE b:Person RETURN 1", 5); + expectRowCount("MATCH (a:Person)-[e*1..2]->(b) RETURN 1 LIMIT 100", 25); + expectRowCount("MATCH (a:Person)-[e*1..2]->(b) RETURN 1 SKIP 1", 24); + expectRowCount("MATCH (a:Person)-[e*1..2]->(b) RETURN DISTINCT 1", 1); +} + +TEST_F(ConstantRowsTest, readsTheExplorationThroughTheBroadcast) { + std::string program; + programOf("MATCH (a:Person {name: 'Remy'})-[*2..2]->(b) RETURN 1", program); + + EXPECT_TRUE(contains(program, "db.broadcast_constant")) << program; + EXPECT_FALSE(contains(program, " distinct")) << program; +} + +TEST_F(ConstantRowsTest, laysOutTheConstantsAfterAnAggregateOverItsRow) { + expectRows("MATCH (n:Person) RETURN count(*) * (size('ab') - 1) AS c", {{"8"}}); +} + +TEST_F(ConstantRowsTest, unwindsAConstantListForEveryMatchedRow) { + expectRows("MATCH (n:Person) WHERE n.age = 32 UNWIND [1, 2] AS x RETURN x", {{"1"}, {"1"}, {"2"}, {"2"}}); +} + +TEST_F(ConstantRowsTest, runsAConstantsAloneProjectionOverASingleRow) { + expectRows("RETURN 1", {{"1"}}); + expectRows("MATCH (n:Nothing) RETURN 1 / 0", {}); +} From 1ca60f4256fba3753476d10973fc0f9a7141d72d Mon Sep 17 00:00:00 2001 From: Remy Boutonnet Date: Wed, 21 Oct 2026 19:06:44 +0200 Subject: [PATCH 2/6] MLIR: drop the no-user guard of fuse_explore_distinct_ends --- query/ir/dialect/db/DBPasses.cpp | 6 +----- .../query/ir/path/ExploreDistinctEndsTest.cpp | 19 +++++++++++-------- 2 files changed, 12 insertions(+), 13 deletions(-) diff --git a/query/ir/dialect/db/DBPasses.cpp b/query/ir/dialect/db/DBPasses.cpp index a58ddb2755..a7022772f2 100644 --- a/query/ir/dialect/db/DBPasses.cpp +++ b/query/ir/dialect/db/DBPasses.cpp @@ -5722,16 +5722,12 @@ bool readsRowsAsASet(Operation* op, llvm::SmallVectorImpl& passedOn) // Whether every row the op ever emits is only ever read as a member of a set: a dedup or a // distinct aggregate settles it, a row-wise op, a filter or a further hop passes the question -// on to its own users, and anything counting, cutting or outputting rows refuses. So does an -// op with no user: a constant read after it, as in RETURN 1, is laid out over its rows. +// on to its own users, and anything counting, cutting or outputting rows refuses. bool usersReadRowsAsASet(Operation* op) { llvm::SmallPtrSet visited; llvm::SmallVector passedOn {op}; while (!passedOn.empty()) { Operation* const producer = passedOn.pop_back_val(); - if (producer->use_empty()) { - return false; - } for (Operation* const user : producer->getUsers()) { const bool firstVisit = visited.insert(user).second; diff --git a/test/query/ir/path/ExploreDistinctEndsTest.cpp b/test/query/ir/path/ExploreDistinctEndsTest.cpp index 53b45bfd27..55f9973763 100644 --- a/test/query/ir/path/ExploreDistinctEndsTest.cpp +++ b/test/query/ir/path/ExploreDistinctEndsTest.cpp @@ -125,32 +125,34 @@ func.func @main() { } )mlir"; -// MATCH (n)-[e]->{1,3}(m) RETURN 1: no op reads the exploration, and the constant goes out -// once per path +// MATCH (n)-[e]->{1,3}(m) RETURN 1: the constant is laid over the paths, one row per path const char* const constantOutputProgram = R"mlir( func.func @main() { %n = db.scan_nodes() : !db.column %0:3 = db.explore_paths(%n, {}) forward hops 1 to 3 : (!db.column) -> (!db.column, !db.column, !db.column) %c = db.constant(1 : i64) - db.output(%c) : !db.column + %x = db.broadcast_constant(%c, %0#0) : (!db.column, !db.column) -> !db.column + db.output(%x) : !db.column return } )mlir"; -// MATCH (n)-[e]->{1,3}(m)-->(p) RETURN 1: the hop passes the rows on, and no op reads it +// MATCH (n)-[e]->{1,3}(m)-->(p) RETURN 1: the hop passes the rows on to the constant laid +// over them const char* const hopThenConstantOutputProgram = R"mlir( func.func @main() { %n = db.scan_nodes() : !db.column %0:3 = db.explore_paths(%n, {}) forward hops 1 to 3 : (!db.column) -> (!db.column, !db.column, !db.column) %h:4 = db.get_out_edges(%0#1, {}) : (!db.column) -> (!db.column, !db.column, !db.column, !db.column) %c = db.constant(1 : i64) - db.output(%c) : !db.column + %x = db.broadcast_constant(%c, %h#3) : (!db.column, !db.column) -> !db.column + db.output(%x) : !db.column return } )mlir"; -// MATCH (n)-[e]->{1,3}(m) WHERE m.age > 30 RETURN 1: the filter passes the rows on, and no -// op reads it +// MATCH (n)-[e]->{1,3}(m) WHERE m.age > 30 RETURN 1: the filter passes the rows on to the +// constant laid over them const char* const filterThenConstantOutputProgram = R"mlir( func.func @main() { %n = db.scan_nodes() : !db.column @@ -160,7 +162,8 @@ func.func @main() { %mask = db.gt %age, %c : (!db.column, !db.column) -> !db.column %1 = db.filter(%mask, {%0#1}) : (!db.column, !db.column) -> !db.column %one = db.constant(1 : i64) - db.output(%one) : !db.column + %x = db.broadcast_constant(%one, %1) : (!db.column, !db.column) -> !db.column + db.output(%x) : !db.column return } )mlir"; From 5d47f48cf9ee030113bfc1e915bca2d1e6a429d8 Mon Sep 17 00:00:00 2001 From: Remy Boutonnet Date: Wed, 21 Oct 2026 19:20:38 +0200 Subject: [PATCH 3/6] MLIR: lay a hop predicate's constants over the hop's edges --- query/ir/lowering/DBLowering.cpp | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/query/ir/lowering/DBLowering.cpp b/query/ir/lowering/DBLowering.cpp index d1bf09566e..e5bf1a9c91 100644 --- a/query/ir/lowering/DBLowering.cpp +++ b/query/ir/lowering/DBLowering.cpp @@ -6740,6 +6740,11 @@ void DBLowering::lowerHopRegion(mlir::Block& dbHop, mlir::Region& nlHop, size_t // before the mask that the real terminator yields exists nl::Yield placeholder = _builder.create(loc); + // The hop's rows are its edges, one per row of the step + const mlir::Value edges = nlBlock->getArgument(mlir::hopEdgeArgument(step)); + const mlir::Value previousRegionRows = _regionRows; + _regionRows = edges; + mlir::Value mask; for (mlir::Operation& operation : dbHop) { if (mlir::db::Yield yield = mlir::dyn_cast(operation)) { @@ -6749,8 +6754,10 @@ void DBLowering::lowerHopRegion(mlir::Block& dbHop, mlir::Region& nlHop, size_t } } + _regionRows = previousRegionRows; + _builder.setInsertionPoint(placeholder); - const mlir::Value maskChunk = rowAlignedChunk(mask, nlBlock->getArgument(mlir::hopEdgeArgument(step))); + const mlir::Value maskChunk = rowAlignedChunk(mask, edges); _builder.create(loc, mlir::ValueRange {maskChunk}); placeholder.erase(); } From cbe5d90083775a1b3fd2a70452d0a8bd1894cafc Mon Sep 17 00:00:00 2001 From: Remy Boutonnet Date: Wed, 21 Oct 2026 19:20:38 +0200 Subject: [PATCH 4/6] MLIR: pass the rows on through a broadcast in fuse_explore_distinct_ends --- query/ir/dialect/db/DBPasses.cpp | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/query/ir/dialect/db/DBPasses.cpp b/query/ir/dialect/db/DBPasses.cpp index a7022772f2..39a8df1e7b 100644 --- a/query/ir/dialect/db/DBPasses.cpp +++ b/query/ir/dialect/db/DBPasses.cpp @@ -5712,7 +5712,7 @@ bool readsRowsAsASet(Operation* op, llvm::SmallVectorImpl& passedOn) return aggregatesDistinctly(groupAggregate.getKinds()); } else if (Collect collect = dyn_cast(op)) { return collectsDistinctly(collect); - } else if (isRowWiseOp(op) || isa(op) || isEdgeHop(op)) { + } else if (isRowWiseOp(op) || isa(op) || isEdgeHop(op)) { passedOn.push_back(op); return true; } @@ -5721,8 +5721,9 @@ bool readsRowsAsASet(Operation* op, llvm::SmallVectorImpl& passedOn) } // Whether every row the op ever emits is only ever read as a member of a set: a dedup or a -// distinct aggregate settles it, a row-wise op, a filter or a further hop passes the question -// on to its own users, and anything counting, cutting or outputting rows refuses. +// distinct aggregate settles it, a row-wise op, a constant laid over the rows, a filter or a +// further hop passes the question on to its own users, and anything counting, cutting or +// outputting rows refuses. bool usersReadRowsAsASet(Operation* op) { llvm::SmallPtrSet visited; llvm::SmallVector passedOn {op}; From 64c7edb73de30aa55081ceb66a7038f890d99ecf Mon Sep 17 00:00:00 2001 From: Remy Boutonnet Date: Wed, 21 Oct 2026 19:20:38 +0200 Subject: [PATCH 5/6] MLIR: test the passes over constants laid out over their rows --- test/query/ir/passes/CMakeLists.txt | 2 + .../ir/passes/ConstantRowsPassesTest.cpp | 212 ++++++++++++++++++ 2 files changed, 214 insertions(+) create mode 100644 test/query/ir/passes/ConstantRowsPassesTest.cpp diff --git a/test/query/ir/passes/CMakeLists.txt b/test/query/ir/passes/CMakeLists.txt index 5b5cc30a57..d9a46dc668 100644 --- a/test/query/ir/passes/CMakeLists.txt +++ b/test/query/ir/passes/CMakeLists.txt @@ -123,3 +123,5 @@ target_sources(test_query_ir_reroot_pattern_at_seed PRIVATE ${IR_TEST_DIR}/Strin target_include_directories(test_query_ir_reroot_pattern_at_seed PRIVATE ${IR_TEST_DIR}) add_call_v3_test(test_query_ir_push_down_filters_past_writes PushDownFiltersPastWritesTest.cpp) + +add_call_v3_test(test_query_ir_constant_rows_passes ConstantRowsPassesTest.cpp) diff --git a/test/query/ir/passes/ConstantRowsPassesTest.cpp b/test/query/ir/passes/ConstantRowsPassesTest.cpp new file mode 100644 index 0000000000..a3a3b99068 --- /dev/null +++ b/test/query/ir/passes/ConstantRowsPassesTest.cpp @@ -0,0 +1,212 @@ +#include + +#include +#include +#include + +#include "CallV3Test.h" +#include "StringRowSink.h" + +using namespace turing::test; + +namespace { + +using Rows = std::vector; + +bool contains(std::string_view text, std::string_view needle) { + return text.find(needle) != std::string_view::npos; +} + +std::string withConstant(std::string_view query, std::string_view constant) { + std::string result(query); + const size_t at = result.find("$X"); + result.replace(at, 2, constant); + + return result; +} + +// A shape a pass rewrites, the op the rewrite leaves in the program, and the rows it binds +struct PassCase { + std::string_view _match; + std::string_view _fusedOp; + size_t _rowCount {0}; +}; + +} + +// A constant read where rows are in flight is laid over them by db.broadcast_constant, so +// every pass that drops, replaces or moves a row column meets a broadcast it drives. Each +// pass must still fire, and the constant must still go out once per row. +class ConstantRowsPassesTest : public CallV3Test { +protected: + void programOf(std::string_view query, std::string& program) { + StringRowSink sink; + runQuery(std::string("EXPLAIN (db) ") + std::string(query), sink); + + ASSERT_EQ(sink.getRows().size(), 1u); + program = sink.getRows().front().back(); + } + + void rowsOf(std::string_view query, Rows& rows) { + StringRowSink sink; + runQuery(query, sink); + + sink.sortedRows(rows); + } + + // MATCH ... RETURN 1 emits the rows count(*) tallies, through the op the pass fuses + void expectFusedConstantRows(const PassCase& passCase) { + const std::string match(passCase._match); + + StringRowSink counted; + runQuery(match + " RETURN count(*)", counted); + ASSERT_EQ(counted.getRows().size(), 1u); + EXPECT_EQ(counted.getRows().front().front(), std::to_string(passCase._rowCount)) << match; + + const std::string query = match + " RETURN 1"; + + StringRowSink sink; + runQuery(query, sink); + EXPECT_EQ(sink.getRows().size(), passCase._rowCount) << query; + + std::string program; + programOf(query, program); + EXPECT_TRUE(contains(program, "db.broadcast_constant")) << program; + EXPECT_TRUE(contains(program, passCase._fusedOp)) << passCase._fusedOp << "\n" << program; + } + + // The query with a constant that may raise in place of $X emits the rows it does with the + // literal 1, which is computed once + void expectRowsOfTheLiteral(std::string_view query) { + const std::string literal = withConstant(query, "1"); + const std::string raising = withConstant(query, "(size('ab') - 1)"); + + Rows expected; + rowsOf(literal, expected); + + Rows rows; + rowsOf(raising, rows); + + EXPECT_FALSE(expected.empty()) << literal; + EXPECT_EQ(rows, expected) << raising; + } +}; + +TEST_F(ConstantRowsPassesTest, fusesNodeScansUnderAConstantProjection) { + const PassCase cases[] = { + {"MATCH (n:Person)", "db.scan_nodes_by_label", 8}, + {"MATCH (n) WHERE n:Person AND n:Founder", "db.scan_nodes_by_label", 2}, + {"MATCH (n:Person {name: 'Remy'})", "db.scan_nodes_by_property_value", 1}, + {"MATCH (n:Person) WHERE n.name = 'Remy'", "db.scan_nodes_by_property_value", 1}, + {"MATCH (n) WHERE id(n) = 1", "db.const_scan_nodes", 1}, + }; + + for (const PassCase& passCase : cases) { + expectFusedConstantRows(passCase); + } +} + +TEST_F(ConstantRowsPassesTest, fusesEdgeScansUnderAConstantProjection) { + const PassCase cases[] = { + {"MATCH (a)-[e]->(b)", "db.scan_edges", 18}, + {"MATCH (a)-[e:KNOWS_WELL]->(b)", "db.scan_edges_by_type", 3}, + {"MATCH (a:Person)-[e:KNOWS_WELL]->(b)", "db.get_out_edges_by_type", 2}, + {"MATCH (a:Person)-[e]->(b)", "db.scan_out_edges_by_label_src", 17}, + {"MATCH (a)-[e]->(b:Interest)", "db.scan_out_edges_by_label_tgt", 15}, + {"MATCH (a {name: 'Remy'})-[e]->(b:Interest)", "db.get_out_edges_by_label", 3}, + {"MATCH (a:Person)-[e:KNOWS_WELL]->(b:Person)", "db.get_out_edges_by_type_and_label", 2}, + }; + + for (const PassCase& passCase : cases) { + expectFusedConstantRows(passCase); + } +} + +TEST_F(ConstantRowsPassesTest, placesFactorsAndReRootsUnderAConstantProjection) { + const PassCase cases[] = { + {"MATCH (a:Person), (b:Interest) WHERE b.name = 'Bio'", "db.scan_nodes_by_property_value", 8}, + {"MATCH (a)-->(b:Person {name: 'Remy'})", "db.scan_out_edges_by_label_tgt", 2}, + {"MATCH (a:Interest)<--(b:Person {name: 'Adam'})", "db.scan_in_edges_by_label_tgt", 2}, + }; + + for (const PassCase& passCase : cases) { + expectFusedConstantRows(passCase); + } +} + +TEST_F(ConstantRowsPassesTest, fusesUnwoundIDsUnderAConstantProjection) { + const PassCase cases[] = { + {"UNWIND [0, 1] AS i MATCH (n) WHERE id(n) = i", "db.const_scan_nodes", 2}, + {"UNWIND range(1, 5) AS i", "db.unwind_range", 5}, + {"UNWIND [[0, 1], [2, 3]] AS r MATCH (a), (b) WHERE a = r[0] AND b = r[1]", "db.fetch_nodes", 2}, + {"WITH [0, 2] AS ss, [1, 3] AS ds UNWIND range(0, size(ss) - 1) AS i " + "MATCH (a), (b) WHERE a = ss[i] AND b = ds[i]", + "db.fetch_nodes", + 2}, + {"MATCH (n:Person) WHERE n.age = 32 UNWIND [1, 2] AS x", "db.unwind_const", 4}, + }; + + for (const PassCase& passCase : cases) { + expectFusedConstantRows(passCase); + } +} + +TEST_F(ConstantRowsPassesTest, fusesExplorationsUnderAConstantProjection) { + const PassCase cases[] = { + {"MATCH (a:Person)-[e*1..3]->(b:Interest)", "end_labels", 27}, + {"MATCH (a:Person)((x)-[r]->(y:Person)){1,3}(b)", "hop_labels", 4}, + {"MATCH (a:Person)-[e*1..3]->(b), (a)-->(b)", "end_column", 26}, + {"MATCH (a:Person)-[e*1..3]->(b), (b:Interest {name: 'Bio'})", "end_nodes", 4}, + {"MATCH (a:Person)-[e*1..3]->(b {name: 'Bio'})", "end_nodes", 4}, + {"MATCH (a:Person)-[e*1..3]->(b)-->(c {name: 'Bio'})", "end_nodes", 3}, + }; + + for (const PassCase& passCase : cases) { + expectFusedConstantRows(passCase); + } +} + +TEST_F(ConstantRowsPassesTest, fusesAnOptionalHopUnderAConstantProjection) { + expectFusedConstantRows({"MATCH (a:Person) OPTIONAL MATCH (a)-[:KNOWS_WELL]->(b)", "db.get_out_edges_by_type", 8}); +} + +TEST_F(ConstantRowsPassesTest, readsARaisingConstantOverTheRowsOfAPattern) { + expectRowsOfTheLiteral("MATCH (n:Person) WHERE n.age > $X RETURN n.name"); + expectRowsOfTheLiteral("MATCH (n:Person {name: 'Remy'}) WHERE n.age > $X RETURN n.name"); + expectRowsOfTheLiteral("MATCH (a:Person)-[e]->(b:Interest) WHERE a.age > $X RETURN a.name, b.name"); + expectRowsOfTheLiteral("MATCH (a:Person), (b:Interest) WHERE a.age > $X AND b.name = 'Bio' RETURN a.name, b.name"); + expectRowsOfTheLiteral("MATCH (a)-->(b:Person {name: 'Remy'}) WHERE a.age > $X RETURN a.name"); + expectRowsOfTheLiteral("UNWIND [0, 1] AS i MATCH (n) WHERE id(n) = i AND n.age > $X RETURN n.name"); + expectRowsOfTheLiteral("UNWIND [[0, 1], [2, 3]] AS r MATCH (a), (b) WHERE a = r[0] AND b = r[1] AND a.age > $X " + "RETURN a.name, b.name"); + expectRowsOfTheLiteral("MATCH (a:Person)-[e*1..3]->(b {name: 'Bio'}) WHERE a.age > $X RETURN a.name"); + expectRowsOfTheLiteral("MATCH (a:Person)-[e*1..3]->(b) WHERE a.age > $X RETURN DISTINCT a.name, b.name"); +} + +// The constant laid over the walk's rows reads them one for one, so the rows the DISTINCT +// dedups are still a set of the walk's ends +TEST_F(ConstantRowsPassesTest, exploresDistinctlyPastARaisingConstant) { + const std::string_view query = "MATCH (a:Person)-[e*1..3]->(b) WHERE a.age > (size('ab') - 1) RETURN DISTINCT a.name, b.name"; + + std::string program; + programOf(query, program); + + const size_t exploration = program.find("db.explore_paths"); + ASSERT_NE(exploration, std::string::npos) << program; + + const std::string_view explorationLine = std::string_view(program).substr(exploration, program.find('\n', exploration) - exploration); + EXPECT_TRUE(contains(explorationLine, " distinct")) << program; + EXPECT_TRUE(contains(program, "db.broadcast_constant")) << program; +} + +TEST_F(ConstantRowsPassesTest, readsARaisingConstantOverTheRowsOfItsRegion) { + expectRowsOfTheLiteral("MATCH (a:Person)((x)-[r]->(y) WHERE y.age > $X){1,3}(b) RETURN a.name, b.name"); + expectRowsOfTheLiteral("MATCH (a:Person)((x)-[r]->(y) WHERE $X > 0){1,3}(b) RETURN a.name, b.name"); + expectRowsOfTheLiteral("MATCH (a:Person) RETURN a.name, size([(a)-->(b) WHERE $X > 0 | b.name]) AS s"); + expectRowsOfTheLiteral("MATCH (a:Person) OPTIONAL MATCH (a)-->(b) WHERE $X > 0 RETURN a.name, b.name"); + expectRowsOfTheLiteral("MATCH (a:Person) RETURN a.name, COUNT { MATCH (a)-->(b) WHERE $X > 0 } AS c"); + expectRowsOfTheLiteral("MATCH (a:Person) CALL (a) { MATCH (a)-->(b) WHERE $X > 0 RETURN b } RETURN a.name, b.name"); + expectRowsOfTheLiteral("MATCH (a:Person) RETURN a.name, [x IN [1, 2] WHERE x > $X | x] AS l"); + expectRowsOfTheLiteral("MATCH (a:Person) RETURN a.name, CASE WHEN a.age > 30 THEN $X ELSE 0 END AS c"); + expectRowsOfTheLiteral("MATCH (a:Person) RETURN a.name AS n UNION ALL MATCH (b:Interest) WHERE $X > 0 RETURN b.name AS n"); +} From 5c07a303523b91dce57571cf5c923e735df3459a Mon Sep 17 00:00:00 2001 From: Remy Boutonnet Date: Thu, 22 Oct 2026 11:34:56 +0200 Subject: [PATCH 6/6] MLIR: drive each broadcast from the column its rows come from --- query/ir/codegen/DBPassPipeline.cpp | 3 +- query/ir/dialect/db/DBPasses.cpp | 77 +++++++++ query/ir/dialect/db/DBPasses.td | 5 + test/query/ir/passes/CMakeLists.txt | 2 + .../ir/passes/ConstantRowsPassesTest.cpp | 13 ++ .../DriveBroadcastsFromRowSourcesTest.cpp | 155 ++++++++++++++++++ 6 files changed, 254 insertions(+), 1 deletion(-) create mode 100644 test/query/ir/passes/DriveBroadcastsFromRowSourcesTest.cpp diff --git a/query/ir/codegen/DBPassPipeline.cpp b/query/ir/codegen/DBPassPipeline.cpp index 85e7d92e27..9541653cc8 100644 --- a/query/ir/codegen/DBPassPipeline.cpp +++ b/query/ir/codegen/DBPassPipeline.cpp @@ -10,7 +10,7 @@ namespace { using DBPassFactory = std::unique_ptr (*)(const mlir::db::DBPassContext*); -constexpr size_t DB_PASS_COUNT = 39; +constexpr size_t DB_PASS_COUNT = 40; // The optimisation pipeline every query runs through, in order. An EXPLAIN prefix // reporting on a pass walks the same table one pass at a time, which is what keeps the @@ -53,6 +53,7 @@ const std::array dbPassPipeline = { [](const mlir::db::DBPassContext*) { return mlir::db::createCountPathRows(); }, [](const mlir::db::DBPassContext*) { return mlir::db::createFuseFetchNodes(); }, [](const mlir::db::DBPassContext* context) { return mlir::db::createFuseHashJoin(context); }, + [](const mlir::db::DBPassContext*) { return mlir::db::createDriveBroadcastsFromRowSources(); }, [](const mlir::db::DBPassContext*) { return mlir::db::createTrimUnreadColumns(); }, [](const mlir::db::DBPassContext*) { return mlir::db::createFuseExploreDistinctEnds(); }, [](const mlir::db::DBPassContext* context) { return mlir::db::createCountFromMetadata(context); }, diff --git a/query/ir/dialect/db/DBPasses.cpp b/query/ir/dialect/db/DBPasses.cpp index 39a8df1e7b..e5389ef7f3 100644 --- a/query/ir/dialect/db/DBPasses.cpp +++ b/query/ir/dialect/db/DBPasses.cpp @@ -71,6 +71,7 @@ namespace mlir::db { #define GEN_PASS_DEF_FUSEEXPLOREENDSET #define GEN_PASS_DEF_FUSEEXPLORELISTPREDICATE #define GEN_PASS_DEF_FUSEEXPLOREDISTINCTENDS +#define GEN_PASS_DEF_DRIVEBROADCASTSFROMROWSOURCES #define GEN_PASS_DEF_COUNTPATHROWS #define GEN_PASS_DEF_TRIMUNREADCOLUMNS #define GEN_PASS_DEF_COUNTFROMMETADATA @@ -5767,6 +5768,82 @@ struct FuseExploreDistinctEnds : public impl::FuseExploreDistinctEndsBasegetOperands(), [op](Value operand) { + const bool sameBlock = operand.getParentBlock() == op->getBlock(); + return sameBlock && isa(operand.getType()) && !::db::yieldsConstantColumn(operand); + }); + + if (rowOperand == op->getOperands().end()) { + break; + } + + column = *rowOperand; + } + + return column; +} + +// Erases the row-wise ops computing @param column that nothing reads any more +void eraseUnreadComputation(Value column, mlir::RewriterBase& rewriter) { + llvm::SmallPtrSet erased; + llvm::SmallVector pending; + if (Operation* const op = column.getDefiningOp()) { + pending.push_back(op); + } + + while (!pending.empty()) { + Operation* const op = pending.pop_back_val(); + if (erased.contains(op) || !isRowWiseOp(op) || !mlir::isOpTriviallyDead(op)) { + continue; + } + + for (const Value operand : op->getOperands()) { + if (Operation* const operandOp = operand.getDefiningOp()) { + pending.push_back(operandOp); + } + } + + erased.insert(op); + rewriter.eraseOp(op); + } +} + +struct DriveBroadcastsFromRowSources : public impl::DriveBroadcastsFromRowSourcesBase { + void runOnOperation() override { + mlir::IRRewriter rewriter(&getContext()); + + llvm::SmallVector broadcasts; + getOperation()->walk([&broadcasts](BroadcastConstant broadcast) { + if (broadcast.getDriver()) { + broadcasts.push_back(broadcast); + } + }); + + for (BroadcastConstant broadcast : broadcasts) { + const Value driver = broadcast.getDriver(); + const Value source = rowSourceOf(driver); + if (source == driver) { + continue; + } + + rewriter.modifyOpInPlace(broadcast, [&broadcast, source]() { + broadcast.getDriverMutable().assign(source); + }); + + eraseUnreadComputation(driver, rewriter); + } + } +}; + // A count reads no value off a path, and a path is never null, so count(p) tallies the // rows its build ran over. The tally moves onto the node the path opens on, which the // build already read: row aligned with the path and never null in its own right. diff --git a/query/ir/dialect/db/DBPasses.td b/query/ir/dialect/db/DBPasses.td index 4a1c04f555..e6b3297dbc 100644 --- a/query/ir/dialect/db/DBPasses.td +++ b/query/ir/dialect/db/DBPasses.td @@ -163,6 +163,11 @@ def FuseExploreListPredicate : Pass<"fuse_explore_list_predicate"> { let dependentDialects = ["mlir::db::DB"]; } +def DriveBroadcastsFromRowSources : Pass<"drive_broadcasts_from_row_sources"> { + let summary = "Drive each db.broadcast_constant from the column its driver's rows come from, and drop the row-wise ops that only computed the driver"; + let dependentDialects = ["mlir::db::DB"]; +} + def FuseExploreDistinctEnds : Pass<"fuse_explore_distinct_ends"> { let summary = "Mark a path exploration whose path is unread and whose rows are only ever deduplicated as distinct, so it emits each end once"; let dependentDialects = ["mlir::db::DB"]; diff --git a/test/query/ir/passes/CMakeLists.txt b/test/query/ir/passes/CMakeLists.txt index d9a46dc668..928d8e688d 100644 --- a/test/query/ir/passes/CMakeLists.txt +++ b/test/query/ir/passes/CMakeLists.txt @@ -125,3 +125,5 @@ target_include_directories(test_query_ir_reroot_pattern_at_seed PRIVATE ${IR_TES add_call_v3_test(test_query_ir_push_down_filters_past_writes PushDownFiltersPastWritesTest.cpp) add_call_v3_test(test_query_ir_constant_rows_passes ConstantRowsPassesTest.cpp) + +add_ir_tests(test_query_ir_drive_broadcasts_from_row_sources DriveBroadcastsFromRowSourcesTest.cpp) diff --git a/test/query/ir/passes/ConstantRowsPassesTest.cpp b/test/query/ir/passes/ConstantRowsPassesTest.cpp index a3a3b99068..92b70db5d4 100644 --- a/test/query/ir/passes/ConstantRowsPassesTest.cpp +++ b/test/query/ir/passes/ConstantRowsPassesTest.cpp @@ -151,6 +151,19 @@ TEST_F(ConstantRowsPassesTest, fusesUnwoundIDsUnderAConstantProjection) { } } +// The broadcast reads only the rows of i, so nothing recomputes i from the fused scan +TEST_F(ConstantRowsPassesTest, drivesTheConstantFromTheFusedScan) { + const std::string_view query = "UNWIND [0, 1] AS i MATCH (n) WHERE id(n) = i RETURN 1"; + + std::string program; + programOf(query, program); + + EXPECT_TRUE(contains(program, "db.const_scan_nodes")) << program; + EXPECT_FALSE(contains(program, "db.element_id")) << program; + + expectFusedConstantRows({"UNWIND [0, 1] AS i MATCH (n) WHERE id(n) = i", "db.const_scan_nodes", 2}); +} + TEST_F(ConstantRowsPassesTest, fusesExplorationsUnderAConstantProjection) { const PassCase cases[] = { {"MATCH (a:Person)-[e*1..3]->(b:Interest)", "end_labels", 27}, diff --git a/test/query/ir/passes/DriveBroadcastsFromRowSourcesTest.cpp b/test/query/ir/passes/DriveBroadcastsFromRowSourcesTest.cpp new file mode 100644 index 0000000000..dba79171cf --- /dev/null +++ b/test/query/ir/passes/DriveBroadcastsFromRowSourcesTest.cpp @@ -0,0 +1,155 @@ +#include + +#include + +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/IR/OwningOpRef.h" +#include "mlir/IR/Verifier.h" +#include "mlir/Parser/Parser.h" +#include "mlir/Pass/PassManager.h" + +#include "DBDialect.h" +#include "DBOps.h" +#include "DBPasses.h" +#include "StorageDialect.h" + +namespace { + +// UNWIND [0, 1] AS i MATCH (n) WHERE id(n) = i RETURN 1 once the scan is fused: the driver is +// i, recomputed from the scanned nodes for the broadcast alone +const char* const elementIDDriverProgram = R"mlir( +func.func @main() { + %n = db.scan_nodes() : !db.column + %id = db.element_id(%n) : (!db.column) -> !db.column + %c = db.constant(1 : i64) + %x = db.broadcast_constant(%c, %id) : (!db.column, !db.column) -> !db.column + db.output(%x) : !db.column + return +} +)mlir"; + +// A chain of row-wise ops, one of them reading a constant beside the rows +const char* const chainedDriverProgram = R"mlir( +func.func @main() { + %n = db.scan_nodes() : !db.column + %age = db.get_node_properties(%n, "age") : (!db.column) -> !db.column> + %thirty = db.constant(30 : i64) + %older = db.gt %age, %thirty : (!db.column>, !db.column) -> !db.column + %c = db.constant(1 : i64) + %x = db.broadcast_constant(%c, %older) : (!db.column, !db.column) -> !db.column + db.output(%x) : !db.column + return +} +)mlir"; + +// The driver is read beside the broadcast, so only the broadcast moves +const char* const sharedDriverProgram = R"mlir( +func.func @main() { + %n = db.scan_nodes() : !db.column + %id = db.element_id(%n) : (!db.column) -> !db.column + %c = db.constant(1 : i64) + %x = db.broadcast_constant(%c, %id) : (!db.column, !db.column) -> !db.column + db.output(%x, %id) : !db.column, !db.column + return +} +)mlir"; + +// A division by a column may raise, so it is kept and the broadcast stays on it +const char* const raisingDriverProgram = R"mlir( +func.func @main() { + %n = db.scan_nodes() : !db.column + %age = db.get_node_properties(%n, "age") : (!db.column) -> !db.column> + %ratio = db.div %age, %age : (!db.column>, !db.column>) -> !db.column + %c = db.constant(1 : i64) + %x = db.broadcast_constant(%c, %ratio) : (!db.column, !db.column) -> !db.column + db.output(%x) : !db.column + return +} +)mlir"; + +} + +// A broadcast reads only the rows of its driver, so it is driven from the column those rows +// come from, and the row-wise ops that only computed the driver go +class DriveBroadcastsFromRowSourcesTest : public ::testing::Test { +protected: + DriveBroadcastsFromRowSourcesTest() { + _context.getOrLoadDialect(); + _context.getOrLoadDialect(); + _context.getOrLoadDialect(); + } + + void runPass(const char* programText, mlir::OwningOpRef& module) { + module = mlir::parseSourceString(programText, mlir::ParserConfig(&_context)); + ASSERT_TRUE(module); + + mlir::PassManager passManager(&_context); + passManager.addPass(mlir::db::createDriveBroadcastsFromRowSources()); + ASSERT_TRUE(mlir::succeeded(passManager.run(*module))); + ASSERT_TRUE(mlir::succeeded(mlir::verify(*module))); + } + + static mlir::db::BroadcastConstant findBroadcast(mlir::ModuleOp module) { + mlir::db::BroadcastConstant found; + module.walk([&found](mlir::db::BroadcastConstant broadcast) { + found = broadcast; + }); + + return found; + } + + template + static size_t countOps(mlir::ModuleOp module) { + size_t count = 0; + module.walk([&count](OpType) { + count++; + }); + + return count; + } + + mlir::MLIRContext _context; +}; + +TEST_F(DriveBroadcastsFromRowSourcesTest, drivesFromTheScanPastAnElementID) { + mlir::OwningOpRef module; + runPass(elementIDDriverProgram, module); + + mlir::db::BroadcastConstant broadcast = findBroadcast(*module); + ASSERT_TRUE(broadcast); + EXPECT_TRUE(broadcast.getDriver().getDefiningOp()); + EXPECT_EQ(countOps(*module), 0u); +} + +TEST_F(DriveBroadcastsFromRowSourcesTest, drivesFromTheScanPastAChainOfRowWiseOps) { + mlir::OwningOpRef module; + runPass(chainedDriverProgram, module); + + mlir::db::BroadcastConstant broadcast = findBroadcast(*module); + ASSERT_TRUE(broadcast); + EXPECT_TRUE(broadcast.getDriver().getDefiningOp()); + EXPECT_EQ(countOps(*module), 0u); + EXPECT_EQ(countOps(*module), 0u); +} + +TEST_F(DriveBroadcastsFromRowSourcesTest, keepsADriverReadElsewhere) { + mlir::OwningOpRef module; + runPass(sharedDriverProgram, module); + + mlir::db::BroadcastConstant broadcast = findBroadcast(*module); + ASSERT_TRUE(broadcast); + EXPECT_TRUE(broadcast.getDriver().getDefiningOp()); + EXPECT_EQ(countOps(*module), 1u); +} + +TEST_F(DriveBroadcastsFromRowSourcesTest, staysOnADriverThatMayRaise) { + mlir::OwningOpRef module; + runPass(raisingDriverProgram, module); + + mlir::db::BroadcastConstant broadcast = findBroadcast(*module); + ASSERT_TRUE(broadcast); + EXPECT_TRUE(broadcast.getDriver().getDefiningOp()); + EXPECT_EQ(countOps(*module), 1u); +}