From b296a4d75d83e0aed8164bd76d91c4d685c79aa5 Mon Sep 17 00:00:00 2001 From: ysun67 Date: Wed, 23 Sep 2026 20:24:27 -0400 Subject: [PATCH] feat(cuda): add opt-in exact bitmap join projections --- docs/getting_started.md | 35 ++ src/srdatalog/dsl.py | 4 + src/srdatalog/ir/codegen/cuda/api.py | 9 + src/srdatalog/ir/codegen/cuda/bitmap.py | 113 ++++ .../ir/codegen/cuda/complete_runner.py | 3 + src/srdatalog/ir/codegen/cuda/orchestrator.py | 31 +- src/srdatalog/ir/hir/index.py | 30 ++ src/srdatalog/ir/hir/lower.py | 26 +- src/srdatalog/ir/hir/plan.py | 101 +++- src/srdatalog/ir/hir/types.py | 1 + src/srdatalog/ir/mir/__init__.py | 1 + src/srdatalog/ir/mir/passes.py | 6 + src/srdatalog/ir/mir/types.py | 9 + .../gpu/runtime/jit/bitmap_join.h | 487 ++++++++++++++++++ tests/test_bitmap_plan.py | 69 +++ tests/test_compile_api.py | 27 + 16 files changed, 946 insertions(+), 6 deletions(-) create mode 100644 src/srdatalog/ir/codegen/cuda/bitmap.py create mode 100644 src/srdatalog/runtime/generalized_datalog/gpu/runtime/jit/bitmap_join.h create mode 100644 tests/test_bitmap_plan.py create mode 100644 tests/test_compile_api.py diff --git a/docs/getting_started.md b/docs/getting_started.md index 28e8d93..f24a44d 100644 --- a/docs/getting_started.md +++ b/docs/getting_started.md @@ -102,6 +102,41 @@ Relation pragmas you'll actually use: Rule builders: `atom <= body1 & body2 & ~neg & Filter((v,), "...")`. See the {py:mod}`srdatalog.dsl` API reference for every operator. +## Opt-in exact bitmap join plans + +`with_plan(dedup_bitmap=True)` selects an exact bitmap implementation for a +binary set-valued join projection. It changes the physical plan, not the rule's +logical tuple set, and is not enabled automatically: + +```python +from srdatalog import Program, Relation, Var + +join, value, destination = Var("join"), Var("value"), Var("destination") +Assign = Relation("Assign", 2, input_file="Assign.csv") +Points = Relation("Points", 2, input_file="Points.csv") +Output = Relation("Output", 2) +prog = Program([ + (Output(value, destination) <= Assign(join, destination) & Points(join, value)) + .with_plan(dedup_bitmap=True) +]) +``` + +For recursive rules, select the intended semi-naive variants explicitly with +`with_plan(delta=..., dedup_bitmap=True)`. The compiler maintains the required +dictionary indexes across producer strata and recursive iterations. + +This specialization supports unconstrained binary `NoProvenance` projections: +the shared join variable is eliminated and one distinct variable from each +source is retained. Filters, negation, constants, incompatible execution +strategies, count-only rules, unsupported index types, and plans matching no +recursive variant are rejected rather than silently approximated. Runtime +shared-memory capacity limits also fail explicitly. + +Bitmap plans require a complete runner (`build_project` or `compile_runner`), +not standalone `compile_pipeline` kernel emission. Keep the ordinary plan as a +correctness/performance control; enabling a bitmap plan is not by itself evidence +of an end-to-end speedup. + ## Translating from Nim The upstream Nim reference has a few dozen benchmark programs under diff --git a/src/srdatalog/dsl.py b/src/srdatalog/dsl.py index 62fa34f..2de4136 100644 --- a/src/srdatalog/dsl.py +++ b/src/srdatalog/dsl.py @@ -342,6 +342,7 @@ class PlanEntry: - work_stealing -> mid-level work-stealing (task queue + steal loop) - block_group -> block-group work partitioning - dedup_hash -> GPU hash table for in-kernel existential dedup + - dedup_bitmap -> exact bitmap dedup for binary join projections `balanced_root` / `balanced_sources` drive balanced partitioning for skewed joins (not yet lowered in Python). ''' @@ -355,6 +356,7 @@ class PlanEntry: dedup_hash: bool = False balanced_root: tuple[str, ...] = () balanced_sources: tuple[str, ...] = () + dedup_bitmap: bool = False @dataclass(frozen=True) @@ -406,6 +408,7 @@ def with_plan( work_stealing: bool = False, block_group: bool = False, dedup_hash: bool = False, + dedup_bitmap: bool = False, balanced_root: tuple[str, ...] | list[str] | None = None, balanced_sources: tuple[str, ...] | list[str] | None = None, ) -> Rule: @@ -420,6 +423,7 @@ def with_plan( work_stealing=work_stealing, block_group=block_group, dedup_hash=dedup_hash, + dedup_bitmap=dedup_bitmap, balanced_root=tuple(balanced_root) if balanced_root is not None else (), balanced_sources=tuple(balanced_sources) if balanced_sources is not None else (), ) diff --git a/src/srdatalog/ir/codegen/cuda/api.py b/src/srdatalog/ir/codegen/cuda/api.py index a178b7b..735b5c3 100644 --- a/src/srdatalog/ir/codegen/cuda/api.py +++ b/src/srdatalog/ir/codegen/cuda/api.py @@ -23,6 +23,10 @@ spec for count-phase shape, since the runner files contain count bodies but no isolated count-only goldens exist). +Plans requiring preprocessing or multiple coordinated kernel phases must use +`compile_runner`. Exact bitmap plans cannot be represented by standalone +kernel entry points and are rejected there. + See: - docs/stage2_emitter_audit.md — the per-milestone migration plan. - docs/ir_lowering_semantics.md — the formal lowering rules. @@ -86,6 +90,9 @@ def compile_runner( The byte-equivalence gate (`tests/test_runner_byte_equivalence.py`) anchors this entry point to the upstream goldens throughout the migration. + + Exact bitmap plans use this entry point to emit their execute-only runner, + including preprocessing and exact count/materialization. ''' from srdatalog.ir.codegen.cuda.runner import emit_runner_full @@ -134,6 +141,8 @@ def compile_kernel_body( slots advance by 2 per FULL_VER D2L source — matching legacy `compute_view_slot_offsets`. Pass {} or None for plain DSAI. ''' + if ep.bitmap_join is not None: + raise ValueError('dedup_bitmap requires a complete runner; use compile_runner') from srdatalog.ir.codegen.cuda.emit import EmitCtx, emit from srdatalog.ir.codegen.cuda.envelope import ( assign_handle_positions, diff --git a/src/srdatalog/ir/codegen/cuda/bitmap.py b/src/srdatalog/ir/codegen/cuda/bitmap.py new file mode 100644 index 0000000..4b5bbce --- /dev/null +++ b/src/srdatalog/ir/codegen/cuda/bitmap.py @@ -0,0 +1,113 @@ +'''Execute-only CUDA runner for exact binary projection bitmap plans.''' + +from __future__ import annotations + +import srdatalog.ir.mir.types as m +from srdatalog.ir.hir.types import Version + + +def _bind_relation(source: m.ColumnSource, name: str) -> str: + full = f'get_relation_by_schema<{source.rel_name}, FULL_VER>(db)' + if source.version is Version.DELTA: + relation = ( + f'(iteration == 0) ? {full} : get_relation_by_schema<{source.rel_name}, DELTA_VER>(db)' + ) + elif source.version is Version.FULL: + relation = full + else: + raise ValueError('dedup_bitmap sources must use FULL or DELTA') + cols = ', '.join(map(str, source.index)) + return ( + f' auto& {name}_relation = {relation};\n' + f' const auto& {name}_index = {name}_relation.get_index(SRDatalog::IndexSpec{{{{{cols}}}}});\n' + ) + + +def gen_bitmap_runner( + node: m.ExecutePipeline, + db_type_name: str, + rel_index_types: dict[str, str], +) -> tuple[str, str]: + '''Emit an exact set-projection runner, without hash or ordinary join kernels.''' + plan = node.bitmap_join + if plan is None or len(node.dest_specs) != 1: + raise ValueError('dedup_bitmap requires one binary destination') + if node.count or node.dedup_hash or node.work_stealing or node.block_group or node.use_fan_out: + raise ValueError('dedup_bitmap cannot be combined with other execution strategies') + dest = node.dest_specs[0] + if dest.version is not Version.NEW or sorted(dest.index) != [0, 1] or len(dest.vars) != 2: + raise ValueError('dedup_bitmap requires a binary NEW destination') + for source in (plan.assign, plan.points): + if sorted(source.index) != [0, 1] or source.prefix_vars: + raise ValueError('dedup_bitmap requires unconstrained binary source indexes') + index_type = rel_index_types.get(source.rel_name, '') + if index_type and not any( + t in index_type for t in ('DeviceSortedArrayIndex', 'Device2LevelIndex') + ): + raise ValueError(f'dedup_bitmap does not support index type {index_type!r}') + + runner = f'JitRunner_{node.rule_name}' + declaration = ( + f'struct {runner} {{\n' + f' using DB = {db_type_name};\n' + ' static void execute(DB& db, uint32_t iteration);\n' + '};\n\n' + ) + dictionary = ', '.join(map(str, reversed(plan.points.index))) + body = f'''void {runner}::execute(DB& db, uint32_t iteration) {{ + nvtxRangePushA("{node.rule_name}"); + struct RangeEnd {{ ~RangeEnd() {{ nvtxRangePop(); }} }} range_end; + namespace bitmap = SRDatalog::GPU::bitmap; + bitmap::Input input{{}}; + auto columns = [](const auto& index) -> bitmap::Columns {{ + using Index = std::remove_cvref_t; + static_assert(Index::arity == 2 && std::is_same_v); + if (index.size() == 0) return {{}}; + return {{index.size(), index.data().template column_ptr<0>(), + index.data().template column_ptr<1>()}}; + }}; + auto segments = [&](const auto& index, auto& full, auto& head) {{ + if constexpr (requires {{ index.full(); index.head(); }}) {{ + full = columns(index.full()); + head = columns(index.head()); + }} else {{ + full = columns(index); + }} + }}; + auto keys = [](const auto& index) -> bitmap::Keys {{ + if (index.size() == 0) return {{}}; + if (index.num_unique_root_values() == 0) + throw std::runtime_error("dedup_bitmap: nonempty index has no value-key cache"); + return {{index.num_unique_root_values(), index.root_unique_values().data()}}; + }}; +''' + body += _bind_relation(plan.assign, 'assign') + body += _bind_relation(plan.points, 'points') + body += f''' segments(assign_index, input.assign, input.assign_head); + segments(points_index, input.points_full, input.points_head); + const auto& value_index = points_relation.get_index(SRDatalog::IndexSpec{{{{{dictionary}}}}}); + auto collect_keys = [&](const auto& index) {{ + if constexpr (requires {{ index.full(); index.head(); }}) {{ + input.heaps_full = keys(index.full()); + input.heaps_head = keys(index.head()); + }} else {{ + input.heaps_full = keys(index); + }} + }}; + collect_keys(value_index); + auto& destination = get_relation_by_schema<{dest.rel_name}, NEW_VER>(db); + using Destination = std::remove_reference_t; + static_assert(!has_provenance_v); + bitmap::execute(input, [&](uint64_t rows) -> bitmap::Output {{ + const uint64_t old_rows = destination.size(); + constexpr uint64_t limit = std::numeric_limits::max(); + if (old_rows > limit || rows > limit - old_rows) + throw std::overflow_error("dedup_bitmap: NEW relation exceeds uint32 row limit"); + destination.resize_interned_columns(static_cast(old_rows + rows), 0); + return {{destination.template interned_column<0>() + old_rows, + destination.template interned_column<1>() + old_rows}}; + }}); +}} + +''' + return declaration, '#include "gpu/runtime/jit/bitmap_join.h"\n\n' + declaration + body diff --git a/src/srdatalog/ir/codegen/cuda/complete_runner.py b/src/srdatalog/ir/codegen/cuda/complete_runner.py index d8401a6..a690120 100644 --- a/src/srdatalog/ir/codegen/cuda/complete_runner.py +++ b/src/srdatalog/ir/codegen/cuda/complete_runner.py @@ -36,6 +36,7 @@ # NumSources and diverging from upstream. import srdatalog.ir.dialects.relation.d2l # noqa: F401 import srdatalog.ir.mir.types as m +from srdatalog.ir.codegen.cuda.bitmap import gen_bitmap_runner from srdatalog.ir.codegen.cuda.materialized import is_materialized_pipeline from srdatalog.ir.codegen.cuda.pipeline_utils import ( assign_handle_positions, @@ -910,6 +911,8 @@ def gen_complete_runner( orchestrator can call `JitRunner_X::execute()`). ''' assert isinstance(node, m.ExecutePipeline) + if node.bitmap_join is not None: + return gen_bitmap_runner(node, db_type_name, rel_index_types or {}) if rel_index_types is None: rel_index_types = {} diff --git a/src/srdatalog/ir/codegen/cuda/orchestrator.py b/src/srdatalog/ir/codegen/cuda/orchestrator.py index ec8e898..6b95965 100644 --- a/src/srdatalog/ir/codegen/cuda/orchestrator.py +++ b/src/srdatalog/ir/codegen/cuda/orchestrator.py @@ -167,6 +167,22 @@ def collect_canonical_specs( # ----------------------------------------------------------------------------- +def _required_index_sources(ep: m.ExecutePipeline) -> list[m.MirNode]: + if ep.bitmap_join is None: + return ep.source_specs + plan = ep.bitmap_join + return [ + *ep.source_specs, + plan.assign, + plan.points, + m.ColumnSource( + rel_name=plan.points.rel_name, + version=plan.points.version, + index=list(reversed(plan.points.index)), + ), + ] + + def _gen_execute_pipeline( instr: m.ExecutePipeline, indent: str, @@ -174,6 +190,10 @@ def _gen_execute_pipeline( count_only_rels: set[str], ) -> str: runner_name = f"JitRunner_{instr.rule_name}" + if instr.bitmap_join is not None: + if instr.count or is_count_only_pipeline(instr, count_only_rels): + raise ValueError('dedup_bitmap does not support count-only execution') + return indent + f"{runner_name}::execute(db, {iter_var});\n" if instr.count or is_count_only_pipeline(instr, count_only_rels): out = indent + "// Count-only query mode\n" out += indent + "{\n" @@ -226,7 +246,7 @@ def _gen_parallel_group( else: other_ops.append(op) - has_dedup = any(op.dedup_hash for op in exec_ops) + has_dedup = any(op.dedup_hash or op.bitmap_join is not None for op in exec_ops) if len(exec_ops) <= 1: out = indent + "// === ParallelGroup (single rule, sequential) ===\n" @@ -241,7 +261,10 @@ def _gen_parallel_group( return out if has_dedup: - out = indent + "// === ParallelGroup (sequential, dedup_hash present) ===\n" + strategy = ( + "dedup_bitmap" if any(op.bitmap_join is not None for op in exec_ops) else "dedup_hash" + ) + out = indent + f"// === ParallelGroup (sequential, {strategy} present) ===\n" for op in exec_ops: if op.count or is_count_only_pipeline(op, count_only_rels): out += _gen_execute_pipeline(op, indent, iter_var, count_only_rels) @@ -663,7 +686,7 @@ def gen_fixpoint_body( exec_pipelines.append(op) for ep in exec_pipelines: - for src_spec in ep.source_specs: + for src_spec in _required_index_sources(ep): if isinstance(src_spec, m.ColumnSource): ver = version_string(src_spec.version.code) spec_type = gen_index_spec_type(src_spec.rel_name, ver, list(src_spec.index)) @@ -851,7 +874,7 @@ def gen_non_recursive_block( if isinstance(op, m.ExecutePipeline): exec_pipelines.append(op) for ep in exec_pipelines: - for src_spec in ep.source_specs: + for src_spec in _required_index_sources(ep): rel_name, raw_ver, idx = extract_source_info(src_spec) if rel_name: ver = version_string(raw_ver) diff --git a/src/srdatalog/ir/hir/index.py b/src/srdatalog/ir/hir/index.py index 39024ee..93ab2c4 100644 --- a/src/srdatalog/ir/hir/index.py +++ b/src/srdatalog/ir/hir/index.py @@ -26,6 +26,7 @@ from __future__ import annotations from srdatalog.ir.hir.pass_ import IRLevel, PassInfo, PassLevel +from srdatalog.ir.hir.plan import bitmap_join_patterns from srdatalog.ir.hir.types import ( HirProgram, HirRuleVariant, @@ -101,9 +102,28 @@ def _append_unique_indices( dest.append(cidx) +def bitmap_index_requirements(hir: HirProgram) -> dict[str, set[tuple[int, ...]]]: + '''Include the value-first dictionary index even when no logical source uses it. + + These indices must also survive producer strata: DELTA uses FULL on iteration zero. + ''' + out: dict[str, set[tuple[int, ...]]] = {} + for stratum in hir.strata: + for variant in (*stratum.base_variants, *stratum.recursive_variants): + patterns = bitmap_join_patterns(variant) + if patterns is None: + continue + assign, points = patterns + for pattern in patterns: + out.setdefault(pattern.rel_name, set()).add(tuple(pattern.index_cols)) + out[points.rel_name].add(tuple(reversed(points.index_cols))) + return out + + def select_indices(hir: HirProgram) -> HirProgram: '''Pass 5 entry. Mutates and returns the HirProgram.''' decls = hir.relation_decls + bitmap_indices = bitmap_index_requirements(hir) # ----- First pass: global_index_map over all strata. for stratum in hir.strata: @@ -117,6 +137,12 @@ def select_indices(hir: HirProgram) -> HirProgram: if idx_list not in hir.global_index_map[rel_name]: hir.global_index_map[rel_name].append(idx_list) + for rel_name, indices in bitmap_indices.items(): + registered = hir.global_index_map.setdefault(rel_name, []) + for index in sorted(indices): + if list(index) not in registered: + registered.append(list(index)) + # ----- Second pass: per-stratum required + canonical. for stratum in hir.strata: if stratum.is_recursive: @@ -141,6 +167,10 @@ def select_indices(hir: HirProgram) -> HirProgram: if rel_name in all_idx: _append_unique_indices(indices, all_idx[rel_name], arity) + # Bitmap sources and their value dictionary are explicit physical + # requirements, including when a producer already has other local indices. + _append_unique_indices(indices, bitmap_indices.get(rel_name, set()), arity) + # Fallback 1: global index map (indices used by OTHER strata). if not indices and rel_name in hir.global_index_map: for idx in hir.global_index_map[rel_name]: diff --git a/src/srdatalog/ir/hir/lower.py b/src/srdatalog/ir/hir/lower.py index f6340c4..4196a1c 100644 --- a/src/srdatalog/ir/hir/lower.py +++ b/src/srdatalog/ir/hir/lower.py @@ -22,7 +22,8 @@ import srdatalog.ir.mir.types as mir from srdatalog.dsl import ArgKind, Atom, Filter, Let -from srdatalog.ir.hir.index import complete_index, get_arity +from srdatalog.ir.hir.index import bitmap_index_requirements, complete_index, get_arity +from srdatalog.ir.hir.plan import bitmap_join_patterns from srdatalog.ir.hir.types import AccessPattern, HirProgram, HirRuleVariant, HirStratum, Version @@ -599,6 +600,7 @@ def wrap_in_execute_pipeline( block_group: bool = False, count: bool = False, dedup_hash: bool = False, + bitmap_join: mir.BitmapJoin | None = None, ) -> mir.ExecutePipeline: '''Wrap a pipeline body in an ExecutePipeline node, extracting source specs (flattened through ColumnJoin/CartesianJoin) and dest specs @@ -621,6 +623,22 @@ def wrap_in_execute_pipeline( block_group=block_group, dedup_hash=dedup_hash, count=count, + bitmap_join=bitmap_join, + ) + + +def _lower_bitmap_join(variant: HirRuleVariant) -> mir.BitmapJoin | None: + patterns = bitmap_join_patterns(variant) + if patterns is None: + return None + assign, points = patterns + return mir.BitmapJoin( + assign=mir.ColumnSource( + rel_name=assign.rel_name, version=assign.version, index=list(assign.index_cols) + ), + points=mir.ColumnSource( + rel_name=points.rel_name, version=points.version, index=list(points.index_cols) + ), ) @@ -658,6 +676,7 @@ def lower_hir_to_mir_steps(hir: HirProgram) -> list[tuple[mir.MirNode, bool]]: ''' out: list[tuple[mir.MirNode, bool]] = [] decls = hir.relation_decls + bitmap_indices = bitmap_index_requirements(hir) for stratum in hir.strata: if stratum.is_recursive: @@ -727,6 +746,7 @@ def lower_hir_to_mir_steps(hir: HirProgram) -> list[tuple[mir.MirNode, bool]]: block_group=variant.block_group, count=variant.count, dedup_hash=variant.dedup_hash, + bitmap_join=_lower_bitmap_join(variant), ) ) else: @@ -741,6 +761,7 @@ def lower_hir_to_mir_steps(hir: HirProgram) -> list[tuple[mir.MirNode, bool]]: block_group=variant.block_group, count=variant.count, dedup_hash=variant.dedup_hash, + bitmap_join=_lower_bitmap_join(variant), ) ) @@ -764,6 +785,7 @@ def lower_hir_to_mir_steps(hir: HirProgram) -> list[tuple[mir.MirNode, bool]]: full_needed: set[tuple[int, ...]] = set() for raw_idx in full_map.get(rel_name, set()): full_needed.add(tuple(complete_index(list(raw_idx), arity))) + full_needed.update(bitmap_indices.get(rel_name, set())) loop_ops.extend( generate_loop_maintenance( rel_name, @@ -846,6 +868,7 @@ def lower_hir_to_mir_steps(hir: HirProgram) -> list[tuple[mir.MirNode, bool]]: block_group=variant.block_group, count=variant.count, dedup_hash=variant.dedup_hash, + bitmap_join=_lower_bitmap_join(variant), ) ) else: @@ -860,6 +883,7 @@ def lower_hir_to_mir_steps(hir: HirProgram) -> list[tuple[mir.MirNode, bool]]: block_group=variant.block_group, count=variant.count, dedup_hash=variant.dedup_hash, + bitmap_join=_lower_bitmap_join(variant), ) ) diff --git a/src/srdatalog/ir/hir/plan.py b/src/srdatalog/ir/hir/plan.py index 155b44d..0205fa4 100644 --- a/src/srdatalog/ir/hir/plan.py +++ b/src/srdatalog/ir/hir/plan.py @@ -22,7 +22,7 @@ from srdatalog.dsl import Agg, ArgKind, Atom, Filter, Let, Negation, PlanEntry, Rule, Split from srdatalog.ir.hir.pass_ import IRLevel, PassInfo, PassLevel -from srdatalog.ir.hir.types import AccessPattern, HirProgram, HirRuleVariant, Version +from srdatalog.ir.hir.types import AccessPattern, HirProgram, HirRuleVariant, RelationDecl, Version # ----------------------------------------------------------------------------- # Rule Analysis @@ -449,6 +449,81 @@ def compute_temp_vars(rule: Rule, split_at: int) -> list[str]: return result +def _validate_bitmap_plan(rule: Rule, plan: PlanEntry, decls: dict[str, RelationDecl]) -> None: + label = f"dedup_bitmap for rule {rule.name or ''!r}" + incompatible = [ + name + for name in ( + "dedup_hash", + "fanout", + "work_stealing", + "block_group", + "balanced_root", + "balanced_sources", + ) + if getattr(plan, name) + ] + if rule.count: + incompatible.append("count") + if incompatible: + raise ValueError(f"{label} is incompatible with {', '.join(incompatible)}") + if ( + len(rule.heads) != 1 + or len(rule.head.args) != 2 + or len(rule.body) != 2 + or any(not isinstance(atom, Atom) or len(atom.args) != 2 for atom in rule.body) + ): + raise ValueError( + f"{label} requires one binary head and exactly two positive binary atoms " + "(no filters, negation, aggregates, lets, or split)" + ) + atoms = (rule.head, *rule.body) + if any( + arg.kind is not ArgKind.LVAR or arg.var_name is None for atom in atoms for arg in atom.args + ): + raise ValueError(f"{label} requires only variables, without constants or expressions") + value, destination = (arg.var_name for arg in rule.head.args) + body_vars = [{arg.var_name for arg in atom.args} for atom in rule.body] + shared = body_vars[0] & body_vars[1] + if ( + value == destination + or len(shared) != 1 + or len(body_vars[0] | body_vars[1]) != 3 + or value in shared + or destination in shared + or {value, destination} != (body_vars[0] | body_vars[1]) - shared + ): + raise ValueError( + f"{label} requires C(value, destination) :- A(join, destination), B(join, value) " + "with three distinct variables (either body column/clause order is allowed)" + ) + join = next(iter(shared)) + if plan.var_order and ( + len(plan.var_order) != 3 + or set(plan.var_order) != {join, value, destination} + or plan.var_order[0] != join + ): + raise ValueError(f"{label} requires var_order to contain all three variables, join first") + if plan.clause_order and sorted(plan.clause_order) != [0, 1]: + raise ValueError(f"{label} requires clause_order to be a permutation of [0, 1]") + for atom in atoms: + decl = decls.get(atom.rel) + semiring = decl.semiring if decl is not None else getattr(atom.relation, "semiring", None) + if semiring is None or semiring.rsplit("::", 1)[-1] != "NoProvenance": + raise ValueError(f"{label} requires NoProvenance set semantics for relation {atom.rel!r}") + if (decl is not None and decl.count_only) or getattr(atom.relation, "count_only", False): + raise ValueError(f"{label} is incompatible with count-only relation {atom.rel!r}") + + +def bitmap_join_patterns(v: HirRuleVariant) -> tuple[AccessPattern, AccessPattern] | None: + '''Return validated source roles without adding a synthetic body clause.''' + if not v.dedup_bitmap: + return None + destination = v.original_rule.head.args[1].var_name + first, second = v.access_patterns + return (first, second) if destination in first.access_order else (second, first) + + def _plan_variant(v: HirRuleVariant) -> None: rule = v.original_rule analysis = analyze_rule(rule) @@ -476,8 +551,11 @@ def _plan_variant(v: HirRuleVariant) -> None: v.work_stealing = plan.work_stealing v.block_group = plan.block_group v.dedup_hash = plan.dedup_hash + v.dedup_bitmap = plan.dedup_bitmap v.balanced_root = list(plan.balanced_root) v.balanced_sources = list(plan.balanced_sources) + if plan.dedup_bitmap and plan.clause_order: + clause_order = list(plan.clause_order) v.count = rule.count v.clause_order = clause_order @@ -507,6 +585,27 @@ def _plan_variant(v: HirRuleVariant) -> None: def plan_joins(hir: HirProgram) -> HirProgram: '''HIR Pass 4 entry. Mutates and returns the HirProgram.''' + decls = {decl.rel_name: decl for decl in hir.relation_decls} + variants = [ + variant + for stratum in hir.strata + for variant in (*stratum.base_variants, *stratum.recursive_variants) + ] + for variant in variants: + rule = variant.original_rule + for plan in rule.plans: + if plan.dedup_bitmap: + _validate_bitmap_plan(rule, plan, decls) + if not any( + candidate.original_rule is rule + and candidate.delta_idx == plan.delta + and _find_plan(rule, candidate.delta_idx) is plan + for candidate in variants + ): + raise ValueError( + f"dedup_bitmap for rule {rule.name or ''!r}: " + f"plan delta={plan.delta} does not select an evaluated variant" + ) for stratum in hir.strata: for v in stratum.base_variants: _plan_variant(v) diff --git a/src/srdatalog/ir/hir/types.py b/src/srdatalog/ir/hir/types.py index c9f03fa..e571d87 100644 --- a/src/srdatalog/ir/hir/types.py +++ b/src/srdatalog/ir/hir/types.py @@ -102,6 +102,7 @@ class HirRuleVariant: block_group: bool = False dedup_hash: bool = False count: bool = False + dedup_bitmap: bool = False @dataclass diff --git a/src/srdatalog/ir/mir/__init__.py b/src/srdatalog/ir/mir/__init__.py index 902f52f..d6cae06 100644 --- a/src/srdatalog/ir/mir/__init__.py +++ b/src/srdatalog/ir/mir/__init__.py @@ -18,6 +18,7 @@ from srdatalog.ir.mir.types import ( Aggregate, BalancedScan, + BitmapJoin, Block, CartesianJoin, CheckSize, diff --git a/src/srdatalog/ir/mir/passes.py b/src/srdatalog/ir/mir/passes.py index 0c21c3f..1848939 100644 --- a/src/srdatalog/ir/mir/passes.py +++ b/src/srdatalog/ir/mir/passes.py @@ -91,6 +91,12 @@ def _collect_needed_indices(node: mir.MirNode, rel_name: str, out: set[tuple[int for s in node.sources: _collect_needed_indices(s, rel_name, out) elif isinstance(node, mir.ExecutePipeline): + if node.bitmap_join is not None: + for source in (node.bitmap_join.assign, node.bitmap_join.points): + _collect_needed_indices(source, rel_name, out) + points = node.bitmap_join.points + if points.rel_name == rel_name: + out.add(tuple(reversed(points.index))) for op in node.pipeline: _collect_needed_indices(op, rel_name, out) elif isinstance(node, mir.ParallelGroup): diff --git a/src/srdatalog/ir/mir/types.py b/src/srdatalog/ir/mir/types.py index 8fc3a2c..4e0f4cc 100644 --- a/src/srdatalog/ir/mir/types.py +++ b/src/srdatalog/ir/mir/types.py @@ -283,6 +283,14 @@ class RebuildIndexFromIndex: # ----------------------------------------------------------------------------- +@dataclass +class BitmapJoin: + '''Binary projection sources, each indexed in logical (join, other) order.''' + + assign: ColumnSource + points: ColumnSource + + @dataclass class ExecutePipeline: '''(execute-pipeline #:rule N #:sources (tuple ...) #:dests (tuple ...) )''' @@ -299,6 +307,7 @@ class ExecutePipeline: dedup_hash: bool = False count: bool = False concurrent_write: bool = False + bitmap_join: BitmapJoin | None = None @dataclass diff --git a/src/srdatalog/runtime/generalized_datalog/gpu/runtime/jit/bitmap_join.h b/src/srdatalog/runtime/generalized_datalog/gpu/runtime/jit/bitmap_join.h new file mode 100644 index 0000000..2963b98 --- /dev/null +++ b/src/srdatalog/runtime/generalized_datalog/gpu/runtime/jit/bitmap_join.h @@ -0,0 +1,487 @@ +#pragma once + +#include "gpu/device_array.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace SRDatalog::GPU::bitmap { + +struct Columns { + uint64_t rows; + const uint32_t* first; + const uint32_t* second; +}; + +struct Keys { + uint64_t size; + const uint32_t* data; +}; + +// Borrowed columns are paired (join, destination) or (join, value). +// Each points segment is sorted by join. Each heap cache is sorted and unique; +// their union must contain every consumed value, but may be a safe superset. +// Heap-cache boundaries are independent of points FULL/head boundaries. +struct Input { + Columns assign; + Columns assign_head; + Columns points_full; + Columns points_head; + Keys heaps_full; + Keys heaps_head; +}; + +struct Output { + uint32_t* value; + uint32_t* destination; +}; + +namespace detail { + +inline constexpr unsigned kThreads = 256; +inline constexpr unsigned kWarpSize = 32; +inline constexpr unsigned kWarps = kThreads / kWarpSize; +inline constexpr uint64_t kRankMapBytes = 64ull * 1024 * 1024; + +inline void check_cuda(cudaError_t status, const char* stage) { + if (status != cudaSuccess) { + throw std::runtime_error(std::string("bitmap join: ") + stage + ": " + + cudaGetErrorString(status)); + } +} + +// Declare after all scratch arrays so exceptional exits finish GPU work before +// their destructors release storage. Normal exits synchronize with error checks. +struct SynchronizeOnError { + int exceptions = std::uncaught_exceptions(); + ~SynchronizeOnError() { + if (std::uncaught_exceptions() > exceptions) cudaStreamSynchronize(nullptr); + } +}; + +template +inline void require_addressable(uint64_t size, const char* name) { + // DeviceArray rounds byte counts up to uint32 storage; Thrust differences are + // signed. Check before either multiplication or pointer arithmetic occurs. + constexpr uint64_t limit = + std::min(std::numeric_limits::max() - sizeof(uint32_t) + 1, + std::numeric_limits::max()); + if (size > limit / sizeof(T)) { + throw std::overflow_error(std::string("bitmap join: ") + name + + " exceeds addressable size"); + } +} + +inline uint64_t checked_add(uint64_t first, uint64_t second, const char* name) { + if (first > std::numeric_limits::max() - second) { + throw std::overflow_error(std::string("bitmap join: ") + name + " overflows uint64"); + } + return first + second; +} + +inline void require_columns(const Columns& columns, const char* name) { + require_addressable(columns.rows, name); + if (columns.rows != 0 && (columns.first == nullptr || columns.second == nullptr)) { + throw std::invalid_argument(std::string("bitmap join: missing ") + name + " columns"); + } +} + +inline void require_keys(const Keys& keys) { + require_addressable(keys.size, "heap cache"); + if (keys.size != 0 && keys.data == nullptr) { + throw std::invalid_argument("bitmap join: missing heap cache data"); + } +} + +inline unsigned blocks_for(uint64_t size) { + return static_cast(std::min(size / kThreads + (size % kThreads != 0), + 65535)); +} + +template +__global__ void reverse_assign(Columns full, Columns head, uint64_t edges, uint64_t* reversed) { + for (uint64_t edge = uint64_t(blockIdx.x) * blockDim.x + threadIdx.x; + edge < edges; edge += uint64_t(gridDim.x) * blockDim.x) { + const Columns segment = edge < full.rows ? full : head; + const uint64_t row = edge < full.rows ? edge : edge - full.rows; + reversed[edge] = (uint64_t(segment.second[row]) << 32) | segment.first[row]; + } +} + +struct DestinationStart { + const uint64_t* reversed; + __host__ __device__ bool operator()(uint64_t edge) const { + return edge == 0 || (reversed[edge] >> 32) != (reversed[edge - 1] >> 32); + } +}; + +template +__global__ void finish_destinations(const uint64_t* reversed, uint64_t edges, + uint32_t destinations, unsigned segments, + uint32_t* ids, uint64_t* offsets) { + for (uint64_t destination = uint64_t(blockIdx.x) * blockDim.x + threadIdx.x; + destination <= destinations; destination += uint64_t(gridDim.x) * blockDim.x) { + if (destination == destinations) { + offsets[destination] = edges * segments; + } else { + const uint64_t edge = offsets[destination]; + ids[destination] = static_cast(reversed[edge] >> 32); + offsets[destination] = edge * segments; + } + } +} + +__device__ inline uint64_t source_bound(const uint32_t* variables, uint64_t rows, + uint32_t source, bool upper) { + uint64_t low = 0; + uint64_t high = rows; + while (low < high) { + const uint64_t middle = low + (high - low) / 2; + const uint32_t value = variables[middle]; + if (value < source || (upper && value == source)) low = middle + 1; + else high = middle; + } + return low; +} + +template +__global__ void build_point_ranges(const uint64_t* reversed, uint64_t edges, + Columns full, Columns head, unsigned segments, + uint64_t* begin, uint64_t* end) { + for (uint64_t edge = uint64_t(blockIdx.x) * blockDim.x + threadIdx.x; + edge < edges; edge += uint64_t(gridDim.x) * blockDim.x) { + const uint32_t source = static_cast(reversed[edge]); + uint64_t slot = edge * segments; + if (full.rows != 0) { + begin[slot] = source_bound(full.first, full.rows, source, false); + end[slot] = source_bound(full.first, full.rows, source, true); + ++slot; + } + if (head.rows != 0) { + begin[slot] = full.rows + source_bound(head.first, head.rows, source, false); + end[slot] = full.rows + source_bound(head.first, head.rows, source, true); + } + } +} + +template +__global__ void build_heap_ranks(const uint32_t* ids, uint32_t heaps, uint32_t* ranks) { + for (uint64_t rank = uint64_t(blockIdx.x) * blockDim.x + threadIdx.x; + rank < heaps; rank += uint64_t(gridDim.x) * blockDim.x) { + // No sentinel: every consumed value belongs to the cache union. + ranks[ids[rank]] = static_cast(rank); + } +} + +struct BitmapInput { + uint32_t heaps; + const uint32_t* destination_ids; + const uint32_t* heap_ids; + const uint32_t* heap_ranks; + const uint64_t* edge_offsets; + const uint64_t* point_begin; + const uint64_t* point_end; + const uint32_t* points_full; + const uint32_t* points_head; + uint64_t full_rows; +}; + +// One block owns one destination. The count and emit passes construct exactly +// the same shared bitmap; output is unique, but not globally value-sorted. +template +__global__ void bitmap_kernel(BitmapInput input, uint32_t* counts, + const uint64_t* output_offsets, Output output) { + extern __shared__ uint32_t bitmap[]; + __shared__ uint32_t warp_totals[kWarps]; + const uint32_t destination = blockIdx.x; + const unsigned thread = threadIdx.x; + const unsigned lane = thread % kWarpSize; + const unsigned warp = thread / kWarpSize; + const uint64_t words = (uint64_t(input.heaps) + 31) / 32; + for (uint64_t word = thread; word < words; word += kThreads) bitmap[word] = 0; + __syncthreads(); + + const uint64_t edge_begin = input.edge_offsets[destination]; + const uint64_t edge_end = input.edge_offsets[uint64_t(destination) + 1]; + for (uint64_t edge = edge_begin + warp; edge < edge_end; edge += kWarps) { + const uint64_t begin = input.point_begin[edge]; + const uint64_t end = input.point_end[edge]; + for (uint64_t point = begin + lane; point < end; point += kWarpSize) { + const uint32_t original = point < input.full_rows + ? input.points_full[point] : input.points_head[point - input.full_rows]; + uint32_t heap; + if (input.heap_ranks != nullptr) { + heap = input.heap_ranks[original]; + } else { + uint32_t low = 0; + uint32_t high = input.heaps; + while (low < high) { + const uint32_t middle = low + (high - low) / 2; + if (input.heap_ids[middle] < original) low = middle + 1; + else high = middle; + } + heap = low; + } + const uint32_t mask = uint32_t(1) << (heap % 32); + cuda::atomic_ref word(bitmap[heap / 32]); + // Racing reads must also be atomic. Monotone bit setting lets an observed + // set bit bypass the shared-memory read-modify-write operation. + if ((word.load(cuda::memory_order_relaxed) & mask) == 0) { + word.fetch_or(mask, cuda::memory_order_relaxed); + } + } + } + __syncthreads(); + + uint32_t local_count = 0; + for (uint64_t word = thread; word < words; word += kThreads) { + local_count += __popc(bitmap[word]); + } + uint32_t inclusive = local_count; + for (unsigned distance = 1; distance < kWarpSize; distance *= 2) { + const uint32_t preceding = __shfl_up_sync(0xffffffffu, inclusive, distance); + if (lane >= distance) inclusive += preceding; + } + if (lane == kWarpSize - 1) warp_totals[warp] = inclusive; + __syncthreads(); + + if constexpr (Emit) { + uint32_t prefix = inclusive - local_count; + for (unsigned preceding_warp = 0; preceding_warp < warp; ++preceding_warp) { + prefix += warp_totals[preceding_warp]; + } + if (local_count == 0) return; + uint64_t offset = output_offsets[destination] + prefix; + const uint32_t destination_id = input.destination_ids[destination]; + for (uint64_t word = thread; word < words; word += kThreads) { + uint32_t bits = bitmap[word]; + while (bits != 0) { + const unsigned bit = unsigned(__ffs(static_cast(bits)) - 1); + const uint64_t heap = word * 32 + bit; + output.value[offset] = input.heap_ids[heap]; + output.destination[offset] = destination_id; + ++offset; + bits &= bits - 1; + } + } + } else if (thread == 0) { + uint32_t total = 0; + for (unsigned preceding_warp = 0; preceding_warp < kWarps; ++preceding_warp) { + total += warp_totals[preceding_warp]; + } + counts[destination] = total; + } +} + +template +inline size_t configure_bitmap(uint32_t heaps, int device) { + const size_t shared_bytes = size_t((uint64_t(heaps) + 31) / 32) * sizeof(uint32_t); + cudaFuncAttributes attributes{}; + check_cuda(cudaFuncGetAttributes(&attributes, bitmap_kernel), "inspect bitmap kernel"); + int default_shared = 0; + check_cuda(cudaDeviceGetAttribute(&default_shared, cudaDevAttrMaxSharedMemoryPerBlock, device), + "get default shared-memory limit"); + if (shared_bytes + attributes.sharedSizeBytes > size_t(default_shared)) { + int optin_shared = 0; + check_cuda(cudaDeviceGetAttribute(&optin_shared, cudaDevAttrMaxSharedMemoryPerBlockOptin, device), + "get opt-in shared-memory limit"); + if (shared_bytes + attributes.sharedSizeBytes > size_t(optin_shared) || + shared_bytes > size_t(std::numeric_limits::max())) { + throw std::runtime_error("bitmap join: exact value bitmap exceeds device shared-memory limit"); + } + check_cuda(cudaFuncSetAttribute(bitmap_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(shared_bytes)), "configure exact bitmap capacity"); + } + return shared_bytes; +} + +struct WidenCount { + __host__ __device__ uint64_t operator()(uint32_t count) const { return count; } +}; + +} // namespace detail + +// Computes the exact set C(value, destination). All uint32 IDs are valid, including +// UINT32_MAX. The dictionary must fit one block's shared memory; unsupported sizes +// fail rather than falling back to approximate deduplication. No GPU scratch is +// retained. Allocate is called exactly once for nonzero output, with a uint64 row +// count, and returns pointers to the appended NEW columns in (value, destination) +// order. It must use the default stream; borrowed inputs remain valid through emit. +template +void execute(const Input& spec, Allocate&& allocate) { + using namespace detail; + if ((spec.assign.rows == 0 && spec.assign_head.rows == 0) || + (spec.points_full.rows == 0 && spec.points_head.rows == 0)) return; + require_columns(spec.assign, "assign FULL"); + require_columns(spec.assign_head, "assign head"); + require_columns(spec.points_full, "points FULL"); + require_columns(spec.points_head, "points head"); + require_keys(spec.heaps_full); + require_keys(spec.heaps_head); + const uint64_t edges = checked_add(spec.assign.rows, spec.assign_head.rows, "assign rows"); + checked_add(spec.points_full.rows, spec.points_head.rows, "points segment offsets"); + const uint64_t heap_capacity = checked_add(spec.heaps_full.size, spec.heaps_head.size, + "heap cache union"); + if (heap_capacity == 0) { + throw std::invalid_argument("bitmap join: nonempty points have no heap caches"); + } + require_addressable(heap_capacity, "heap cache union"); + const unsigned segments = unsigned(spec.points_full.rows != 0) + unsigned(spec.points_head.rows != 0); + require_addressable(edges, "reversed assign"); + const uint64_t ranges = segments == 2 ? checked_add(edges, edges, "source ranges") : edges; + require_addressable(ranges, "source ranges"); + + const cudaStream_t stream = nullptr; + auto policy = rmm::exec_policy(stream); + DeviceArray heap_union(0, stream); + DeviceArray heap_ranks(0, stream); + DeviceArray reversed(0, stream); + DeviceArray destination_ids(0, stream); + DeviceArray edge_offsets(0, stream); + DeviceArray point_begin(0, stream); + DeviceArray point_end(0, stream); + DeviceArray counts(0, stream); + DeviceArray output_offsets(0, stream); + const SynchronizeOnError synchronize_on_error; + + // A union cannot be smaller than either unique cache. Reject an impossible + // dictionary before allocating its union; exact kernel static usage is checked + // below once the union's cardinality is known. + int device = 0; + int default_shared = 0; + int optin_shared = 0; + check_cuda(cudaGetDevice(&device), "get device"); + check_cuda(cudaDeviceGetAttribute(&default_shared, cudaDevAttrMaxSharedMemoryPerBlock, device), + "get default shared-memory limit"); + check_cuda(cudaDeviceGetAttribute(&optin_shared, cudaDevAttrMaxSharedMemoryPerBlockOptin, device), + "get opt-in shared-memory limit"); + const uint64_t minimum_heaps = std::max(spec.heaps_full.size, spec.heaps_head.size); + if ((minimum_heaps / 32 + (minimum_heaps % 32 != 0)) > + uint64_t(std::max(default_shared, optin_shared)) / sizeof(uint32_t)) { + throw std::runtime_error("bitmap join: exact value bitmap exceeds device shared-memory limit"); + } + + const uint32_t* heap_ids; + uint64_t heap_count; + if (spec.heaps_full.size == 0) { + heap_ids = spec.heaps_head.data; + heap_count = spec.heaps_head.size; + } else if (spec.heaps_head.size == 0) { + heap_ids = spec.heaps_full.data; + heap_count = spec.heaps_full.size; + } else { + heap_union = DeviceArray(heap_capacity, stream); + const auto end = thrust::set_union(policy, + spec.heaps_full.data, spec.heaps_full.data + spec.heaps_full.size, + spec.heaps_head.data, spec.heaps_head.data + spec.heaps_head.size, heap_union.data()); + check_cuda(cudaGetLastError(), "union heap caches"); + heap_ids = heap_union.data(); + heap_count = end - heap_ids; + } + if (heap_count == 0 || heap_count > std::numeric_limits::max()) { + throw std::runtime_error("bitmap join: heap cache union exceeds nonzero uint32 rank capacity"); + } + const uint32_t heaps = static_cast(heap_count); + const size_t shared_bytes = configure_bitmap(heaps, device); + configure_bitmap(heaps, device); + + uint32_t maximum_heap = 0; + check_cuda(cudaMemcpyAsync(&maximum_heap, heap_ids + heap_count - 1, sizeof(maximum_heap), + cudaMemcpyDeviceToHost, stream), "read maximum value ID"); + check_cuda(cudaStreamSynchronize(stream), "finish heap preprocessing"); + // Widen before adding: UINT32_MAX selects binary search, never wraps to zero. + const uint64_t rank_entries = uint64_t(maximum_heap) + 1; + if (rank_entries <= kRankMapBytes / sizeof(uint32_t)) { + heap_ranks = DeviceArray(rank_entries, stream); + build_heap_ranks<><<>>(heap_ids, heaps, heap_ranks.data()); + check_cuda(cudaGetLastError(), "build value rank map"); + } + + reversed = DeviceArray(edges, stream); + reverse_assign<><<>>( + spec.assign, spec.assign_head, edges, reversed.data()); + check_cuda(cudaGetLastError(), "reverse assign segments"); + thrust::sort(policy, reversed.data(), reversed.data() + edges); + check_cuda(cudaGetLastError(), "sort reversed assign"); + const auto first_edge = thrust::make_counting_iterator(0); + const DestinationStart starts{reversed.data()}; + const uint64_t destination_count = thrust::count_if(policy, first_edge, first_edge + edges, starts); + check_cuda(cudaGetLastError(), "count destinations"); + int max_grid_x = 0; + check_cuda(cudaDeviceGetAttribute(&max_grid_x, cudaDevAttrMaxGridDimX, device), "get grid limit"); + if (destination_count > uint64_t(max_grid_x)) { + throw std::runtime_error("bitmap join: destination count exceeds device grid capacity"); + } + const uint32_t destinations = static_cast(destination_count); + const uint64_t offset_count = destination_count + 1; + require_addressable(offset_count, "destination offsets"); + destination_ids = DeviceArray(destinations, stream); + edge_offsets = DeviceArray(offset_count, stream); + thrust::copy_if(policy, first_edge, first_edge + edges, edge_offsets.data(), starts); + check_cuda(cudaGetLastError(), "extract destination boundaries"); + finish_destinations<><<>>( + reversed.data(), edges, destinations, segments, destination_ids.data(), edge_offsets.data()); + check_cuda(cudaGetLastError(), "build destination edge offsets"); + point_begin = DeviceArray(ranges, stream); + point_end = DeviceArray(ranges, stream); + build_point_ranges<><<>>( + reversed.data(), edges, spec.points_full, spec.points_head, segments, + point_begin.data(), point_end.data()); + check_cuda(cudaGetLastError(), "build source ranges"); + check_cuda(cudaStreamSynchronize(stream), "finish source preprocessing"); + reversed = DeviceArray(0, stream); + + const BitmapInput input{heaps, destination_ids.data(), heap_ids, + heap_ranks.empty() ? nullptr : heap_ranks.data(), edge_offsets.data(), + point_begin.data(), point_end.data(), spec.points_full.second, + spec.points_head.second, spec.points_full.rows}; + output_offsets = DeviceArray(offset_count, stream); + counts = DeviceArray(destinations, stream); + bitmap_kernel<<>>( + input, counts.data(), nullptr, Output{nullptr, nullptr}); + check_cuda(cudaGetLastError(), "count exact destination bitmaps"); + check_cuda(cudaMemsetAsync(output_offsets.data(), 0, sizeof(uint64_t), stream), "initialize output offset"); + const auto first_count = thrust::make_transform_iterator(counts.data(), WidenCount{}); + // Each count is at most heaps, and both dimensions are uint32, so the widened + // scan cannot overflow uint64 even before the output addressability check. + thrust::inclusive_scan(policy, first_count, first_count + destinations, output_offsets.data() + 1, + thrust::plus()); + check_cuda(cudaGetLastError(), "scan exact output counts"); + uint64_t output_rows = 0; + check_cuda(cudaMemcpyAsync(&output_rows, output_offsets.data() + destinations, sizeof(output_rows), + cudaMemcpyDeviceToHost, stream), "read exact output total"); + check_cuda(cudaStreamSynchronize(stream), "finish exact output count"); + counts = DeviceArray(0, stream); + if (output_rows == 0) return; + require_addressable(output_rows, "output columns"); + const Output output = std::forward(allocate)(output_rows); + if (output.value == nullptr || output.destination == nullptr) { + throw std::runtime_error("bitmap join: output allocator returned null columns"); + } + bitmap_kernel<<>>( + input, nullptr, output_offsets.data(), output); + check_cuda(cudaGetLastError(), "emit exact destination bitmaps"); + check_cuda(cudaStreamSynchronize(stream), "finish bitmap emission"); +} + +} // namespace SRDatalog::GPU::bitmap diff --git a/tests/test_bitmap_plan.py b/tests/test_bitmap_plan.py new file mode 100644 index 0000000..9480b8d --- /dev/null +++ b/tests/test_bitmap_plan.py @@ -0,0 +1,69 @@ +'''Unsafe bitmap specializations must fail rather than discard rule semantics.''' + +import pytest + +from srdatalog.dsl import Filter, Program, Relation, Var +from srdatalog.ir.hir import compile_to_hir + + +def projection(): + join, value, destination = Var('join'), Var('value'), Var('destination') + assign = Relation('Assign', 2) + points = Relation('Points', 2) + output = Relation('Output', 2) + return join, value, destination, assign, points, output + + +@pytest.mark.parametrize('restriction', ['filter', 'negation', 'constant', 'retained_join']) +def test_bitmap_rejects_rules_that_are_not_unconstrained_projection(restriction): + join, value, destination, assign, points, output = projection() + if restriction == 'filter': + rule = output(value, destination) <= assign(join, destination) & points(join, value) & Filter( + ('value',), 'return value > 0;' + ) + elif restriction == 'negation': + excluded = Relation('Excluded', 2) + rule = output(value, destination) <= ( + assign(join, destination) & points(join, value) & ~excluded(value, destination) + ) + elif restriction == 'constant': + rule = output(value, destination) <= assign(7, destination) & points(7, value) + else: + rule = output(join, destination) <= assign(join, destination) & points(join, value) + with pytest.raises(ValueError): + compile_to_hir(Program([rule.named('Projection').with_plan(dedup_bitmap=True)])) + + +def test_bitmap_rejects_provenance_instead_of_erasing_annotations(): + join, value, destination, assign, _, output = projection() + points = Relation('AnnotatedPoints', 2, semiring='BooleanSR') + rule = output(value, destination) <= assign(join, destination) & points(join, value) + with pytest.raises(ValueError): + compile_to_hir(Program([rule.with_plan(dedup_bitmap=True)])) + + +@pytest.mark.parametrize('mode', ['hash', 'count']) +def test_bitmap_rejects_conflicting_execution_modes(mode): + join, value, destination, assign, points, output = projection() + rule = output(value, destination) <= assign(join, destination) & points(join, value) + if mode == 'hash': + rule = rule.with_plan(dedup_bitmap=True, dedup_hash=True) + else: + rule = rule.with_count().with_plan(dedup_bitmap=True) + with pytest.raises(ValueError): + compile_to_hir(Program([rule])) + + +def test_bitmap_rejects_a_plan_that_matches_no_recursive_variant(): + join, value, destination, assign, _, output = projection() + seed = Relation('Seed', 2) + rule = output(value, destination) <= assign(join, destination) & output(value, join) + with pytest.raises(ValueError): + compile_to_hir( + Program( + [ + (output(value, destination) <= seed(value, destination)).named('Base'), + rule.named('Step').with_plan(dedup_bitmap=True), + ] + ) + ) diff --git a/tests/test_compile_api.py b/tests/test_compile_api.py new file mode 100644 index 0000000..341d594 --- /dev/null +++ b/tests/test_compile_api.py @@ -0,0 +1,27 @@ +'''Public compilation errors, independent of generated C++ spelling.''' + +import pytest +from test_bitmap_plan import projection + +from srdatalog.compile import compile_pipeline +from srdatalog.dsl import Program +from srdatalog.ir.codegen.cuda.batchfile import _collect_pipelines +from srdatalog.ir.hir import compile_to_mir + + +def _projection_pipeline(*, bitmap=False): + join, value, destination, assign, points, output = projection() + rule = (output(value, destination) <= assign(join, destination) & points(join, value)).with_plan( + dedup_bitmap=bitmap + ) + return _collect_pipelines(compile_to_mir(Program([rule])))[0] + + +def test_compile_pipeline_rejects_unknown_target(): + with pytest.raises(ValueError): + compile_pipeline(_projection_pipeline(), target='cpp_tbb') # type: ignore[arg-type] + + +def test_bitmap_plan_requires_a_complete_runner(): + with pytest.raises(ValueError): + compile_pipeline(_projection_pipeline(bitmap=True))