TUR-306: MLIR: name the rows a constant is read over in the db IR

query/ir/codegen/DBPassPipeline.cpp

@@ -10,7 +10,7 @@ namespace {
 
 using DBPassFactory = std::unique_ptr<mlir::Pass> (*)(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<DBPassFactory, DB_PASS_COUNT> 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); },

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<mlir::db::BroadcastConstant>();
+
+    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,20 +5102,34 @@ 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;
         }
 
         return column;
     }
 
     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<mlir::db::BroadcastConstant>(_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<mlir::db::ColumnType>(column.getType())) {
             continue;
         }
 
-        _opBuilder.setInsertionPoint(&op);
-        for (mlir::OpOperand& operand : op.getOpOperands()) {
-            operand.set(_opBuilder.create<mlir::db::BroadcastConstant>(loc, noneType, operand.get(), driver).getResult());
-        }
+        operand.set(builder.create<mlir::db::BroadcastConstant>(op->getLoc(), noneType, column, driver).getResult());
     }
 }
 
@@ -5384,6 +5411,8 @@ void DBProgramGenerator::generateOutput(const Projection* projection, const Unio
     llvm::SmallVector<llvm::StringRef> 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<size_t> 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<mlir::Value> deduped;

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

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<Value, bool>& classified) {
+    const bool constantThroughOperands = op->hasTrait<OpTrait::ConstantThroughOperands>() && ::db::runsWithoutRaising(op);
+    const bool holdsOneValue = op->hasTrait<OpTrait::ConstantLike>() || op->hasTrait<OpTrait::ConstantResult>();
+    const bool laysOutAConstant = isa<BroadcastConstant>(op);
+    if (constantThroughOperands || holdsOneValue || laysOutAConstant || op->hasTrait<OpTrait::IsTerminator>()) {
+        return false;
+    }
+
+    bool readsAColumn = false;
+    for (const Value operand : op->getOperands()) {
+        if (!isa<ColumnType>(operand.getType())) {
+            continue;
+        }
+
+        if (!::db::yieldsConstantColumn(operand, classified)) {
+            return false;
+        }
+
+        readsAColumn = true;
+    }
+
+    return readsAColumn;
+}

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<Value, bool>& classified);
+
+}

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: