[FEA] Implement Multi-output AST JIT & IR CSE - #23621
Conversation
lamarrr
added
feature request
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
💤 Files with no reviewable changes (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughSummary by CodeRabbit
WalkthroughThe PR adds ChangesJIT table transform
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related issues
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (4)
cpp/src/jit/row_ir.hpp (1)
119-120: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument and enforce the node-address invariant for
cse_nodes_andalias_.
cse_nodes_andnode::alias_store rawnode const*.nodekeeps a public defaulted move constructor and move assignment. If any code moves anodeafterinstantiate()registers it, both the map entry and everyalias_that points to it dangle, andemit_codethen dereferences freed memory.The current call paths appear safe because nodes are moved only before
instantiate(), andoutput_irs_holds them throughstd::unique_ptr. The invariant is implicit. Add a comment on the two members that states the requirement, so a later refactor does not break it silently.♻️ Suggested documentation of the invariant
std::unordered_multimap<size_t, node const*> - cse_nodes_; ///< Nodes from completed outputs, indexed by structural hash + cse_nodes_; ///< Nodes from completed outputs, indexed by structural hash. + ///< Non-owning. Registered nodes must not be moved or destroyed + ///< while this context is alive.node const* alias_ = nullptr; ///< The equivalent IR node that this IR aliases, if any. This is ///< used to avoid emitting duplicate code for equivalent IR nodes. + ///< Non-owning. The aliased node must outlive this node and must + ///< not be moved after `instantiate()`.Run the following script to confirm no
nodeis moved after instantiation:#!/bin/bash # Find moves of row_ir::node objects that could invalidate cse_nodes_/alias_ pointers. fd -e cpp -e hpp -e cu -e cuh . cpp | xargs rg -n -C4 'std::move\([^)]*\bnode\b' rg -nP -C4 '\bnode\s*&&|std::vector<\s*node\s*>' cpp/src cpp/testsAlso applies to: 284-287
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/src/jit/row_ir.hpp` around lines 119 - 120, Document on both cse_nodes_ and node::alias_ that registered nodes must not be moved or relocated after instantiate() because these raw pointers must remain valid; preserve the existing ownership and call paths without changing behavior.cpp/tests/jit/row_ir.cpp (1)
477-479: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAssert the nullability values, not only the count.
The test checks
nullability.size(). It does not check the per-output policy. The inputs here are non-nullable and both outputs usePROPAGATEoperators, so both entries must beALL_VALID. Asserting the values protects the per-output nullability logic ingenerate_code.💚 Suggested assertion
EXPECT_EQ(code, expected_code); EXPECT_EQ(null_aware, cudf::null_aware::NO); - EXPECT_EQ(nullability.size(), 2); + ASSERT_EQ(nullability.size(), 2); + EXPECT_EQ(nullability[0], cudf::output_nullability::ALL_VALID); + EXPECT_EQ(nullability[1], cudf::output_nullability::ALL_VALID); }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tests/jit/row_ir.cpp` around lines 477 - 479, Update the nullability assertions in the test around generate_code to verify both entries are ALL_VALID, not just that nullability has size two. Preserve the existing size check and assert the expected value for each output in the nullability collection.cpp/tests/ast/transform_tests.cpp (1)
1567-1590: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd a case with two identical expressions.
CommonSubexpressionsharessumbetween different root expressions. It does not cover the case where the same root expression is passed twice. That case exercisesnode::is_equivalenton twoSET_OUTPUTnodes whose subtrees are identical but whoseoutput_referenceindices differ. The newoutput_reference::operator==is what prevents the second output from aliasing the first and losing its store.💚 Suggested additional test
+TEST_F(ComputeTableJitTest, DuplicateExpressions) +{ + auto c0 = column_wrapper<int32_t>{1, 2, 3, 4}; + auto c1 = column_wrapper<int32_t>{10, 20, 30, 40}; + auto table = cudf::table_view{{c0, c1}}; + + auto ref0 = cudf::ast::column_reference{0}; + auto ref1 = cudf::ast::column_reference{1}; + auto sum = cudf::ast::operation{cudf::ast::ast_operator::ADD, ref0, ref1}; + + std::array<std::reference_wrapper<cudf::ast::expression const>, 2> expressions{sum, sum}; + auto result = cudf::compute_table_jit(table, expressions); + + auto expected_sum = column_wrapper<int32_t>{11, 22, 33, 44}; + auto expected = cudf::table_view{{expected_sum, expected_sum}}; + + CUDF_TEST_EXPECT_TABLES_EQUAL(expected, result->view()); +}🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tests/ast/transform_tests.cpp` around lines 1567 - 1590, Add a duplicate-root-expression case to the CommonSubexpression test by passing the same expression twice in the expressions array and expecting two distinct, identical output columns. Ensure the assertions verify both outputs are stored independently, exercising node::is_equivalent and output_reference::operator== behavior.cpp/src/jit/row_ir.cpp (1)
891-906: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winTrack nullable inputs per output to avoid redundant masks.
has_nullable_inputsmarks an output that reads only valid inputs asPRESERVE, somake_outputsallocates and updates an unnecessary null mask. The null-awareALL_VALIDpath is safe because generated assignments engage thecuda::std::optional<T>output, and null-mask writes are skipped when no mask exists.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/src/jit/row_ir.cpp` around lines 891 - 906, Update the nullability handling around null_policies and generate_null_aware_udf so each output’s nullable-input usage is tracked independently rather than applying has_nullable_inputs to every output. Mark outputs that only read valid inputs as ALL_VALID, and ensure make_outputs skips allocating or updating redundant null masks while preserving optional-based null propagation for null-aware assignments.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@cpp/src/jit/row_ir.cpp`:
- Around line 891-906: Update the nullability handling around null_policies and
generate_null_aware_udf so each output’s nullable-input usage is tracked
independently rather than applying has_nullable_inputs to every output. Mark
outputs that only read valid inputs as ALL_VALID, and ensure make_outputs skips
allocating or updating redundant null masks while preserving optional-based null
propagation for null-aware assignments.
In `@cpp/src/jit/row_ir.hpp`:
- Around line 119-120: Document on both cse_nodes_ and node::alias_ that
registered nodes must not be moved or relocated after instantiate() because
these raw pointers must remain valid; preserve the existing ownership and call
paths without changing behavior.
In `@cpp/tests/ast/transform_tests.cpp`:
- Around line 1567-1590: Add a duplicate-root-expression case to the
CommonSubexpression test by passing the same expression twice in the expressions
array and expecting two distinct, identical output columns. Ensure the
assertions verify both outputs are stored independently, exercising
node::is_equivalent and output_reference::operator== behavior.
In `@cpp/tests/jit/row_ir.cpp`:
- Around line 477-479: Update the nullability assertions in the test around
generate_code to verify both entries are ALL_VALID, not just that nullability
has size two. Preserve the existing size check and assert the expected value for
each output in the nullability collection.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 70e84426-a33b-4404-9b79-1a59052e160b
📒 Files selected for processing (6)
cpp/include/cudf/transform.hppcpp/src/jit/row_ir.cppcpp/src/jit/row_ir.hppcpp/src/transform/transform.cucpp/tests/ast/transform_tests.cppcpp/tests/jit/row_ir.cpp
bdice
left a comment
bdice
left a comment
There was a problem hiding this comment.
Really nice work. The CSE implementation is very clear. I would be interested in seeing some benchmarks for this, perhaps comparing to some baseline like a series of column transforms.
| { | ||
| if (auto* column = std::get_if<column_input>(&in); | ||
| column != nullptr && column->table_source.has_value() && column->column_index.has_value()) { | ||
| for (size_t i = 0; i < inputs_.size(); ++i) { |
There was a problem hiding this comment.
Can this use a find algorithm?
2 participants
Description
Adds
cudf::compute_table_jit, which evaluates multiple AST expressions in a single JIT-compiled transform and returns one output column per expression, in the supplied order.The row-IR changes:
Checklist