add pragma annotation

Author: MooMooHorse

runtime/approx_runtime/cpp_annotation.py

@@ -34,7 +34,13 @@ class CppAnnotation:
 
 
 def parse_cpp_annotations(source: str) -> List[CppAnnotation]:
-    """Parse all @approx annotations from C/C++ source text."""
+    """Parse all @approx annotations from C/C++ source text.
+
+    Supports two annotation formats:
+      1. Comment-based: ``// @approx:decision_tree { ... }``
+      2. Pragma-based:  ``#pragma approx decision_tree key=val ...``
+         (with optional backslash line continuations)
+    """
     lines = source.splitlines()
     annotations: List[CppAnnotation] = []
     i = 0
@@ -53,6 +59,14 @@ def parse_cpp_annotations(source: str) -> List[CppAnnotation]:
                 data = _parse_single_line(raw, i + 1)
                 i += 1
 
+            func_name, arg_count = _find_next_function(lines, i)
+            ann = _build_annotation(data, func_name, arg_count, i + 1)
+            annotations.append(ann)
+        elif _is_pragma_approx(line):
+            raw, end_idx = _collect_pragma_line(lines, i)
+            data = _parse_single_line(raw, i + 1)
+            i = end_idx
+
             func_name, arg_count = _find_next_function(lines, i)
             ann = _build_annotation(data, func_name, arg_count, i + 1)
             annotations.append(ann)
@@ -126,6 +140,34 @@ def parse_and_generate(source: str, module_name: Optional[str] = None) -> str:
     return generate_cpp_annotation_mlir(annotations, module_name)
 
 
+def _is_pragma_approx(line: str) -> bool:
+    stripped = line.strip()
+    return stripped.startswith("#pragma") and "approx" in stripped
+
+
+def _collect_pragma_line(lines: List[str], start: int) -> tuple[str, int]:
+    """Collect a #pragma approx line, joining backslash continuations."""
+    parts: List[str] = []
+    i = start
+    while i < len(lines):
+        line = lines[i].rstrip()
+        if i == start:
+            # Strip the "#pragma approx" prefix
+            idx = line.index("approx") + len("approx")
+            line = line[idx:].strip()
+            # Strip optional "decision_tree" keyword
+            if line.startswith("decision_tree"):
+                line = line[len("decision_tree"):].strip()
+        if line.endswith("\\"):
+            parts.append(line[:-1].strip())
+            i += 1
+        else:
+            parts.append(line.strip())
+            i += 1
+            break
+    return " ".join(parts), i
+
+
 def _collect_block(lines: List[str], start: int) -> tuple[list[str], int]:
     block_lines = []
     i = start
@@ -173,8 +215,12 @@ def _parse_single_line(raw: str, line_number: int) -> dict:
     tokens = raw.split()
     if not tokens:
         raise AnnotationSyntaxError(f"Empty @approx annotation at line {line_number}")
-    data = {"transform_type": tokens[0]}
-    for tok in tokens[1:]:
+    data = {}
+    start = 0
+    if "=" not in tokens[0]:
+        data["transform_type"] = tokens[0]
+        start = 1
+    for tok in tokens[start:]:
         if "=" not in tok:
             continue
         key, value = tok.split("=", 1)

runtime/examples/benchmark/example_pagerank_tuning.py

@@ -110,7 +110,7 @@ def main() -> None:
     argv_base = [
         "-p", "-n", "2000000"
     ]
-    conf_min = int(os.environ.get("PAGERANK_CONF_MIN", "0"))
+    conf_min = int(os.environ.get("PAGERANK_CONF_MIN", "1"))
     conf_max = int(os.environ.get("PAGERANK_CONF_MAX", "5"))
     conf_runs = int(os.environ.get("PAGERANK_CONF_RUNS", "3"))
     conf_seed = os.environ.get("PAGERANK_CONF_SEED")

runtime/tests/test_cpp_annotation.py

@@ -263,3 +263,124 @@ def test_parse_real_lavamd_neighbor_box_accumulate():
     ann = anns[0]
     assert ann.func_name == "neighbor_box_accumulate"
     assert ann.state_indices == [7]
+
+
+# ---- Pragma-based annotation tests ----
+
+
+def test_pragma_single_line():
+    source = """
+#pragma approx decision_tree transform_type=func_substitute state_indices=[-1] thresholds=[4] decisions=[0,1] state_fn=getState
+int foo(int x, int y, int state) { return x + y; }
+"""
+    anns = parse_cpp_annotations(source)
+    assert len(anns) == 1
+    ann = anns[0]
+    assert ann.func_name == "foo"
+    assert ann.transform_type == "func_substitute"
+    assert ann.state_indices == [2]
+    assert ann.state_function == "getState"
+    assert ann.thresholds == [4]
+    assert ann.decisions == [0, 1]
+
+
+def test_pragma_multiline_backslash():
+    source = """
+#pragma approx decision_tree \\
+    transform_type=func_substitute \\
+    state_indices=[7] \\
+    state_function=approx_state_identity \\
+    thresholds=[2000] \\
+    thresholds_lower=[1] \\
+    thresholds_upper=[40] \\
+    decisions=[0,0] \\
+    decision_values=[0,1,2]
+void score_term_over_docs(
+    const char *lower_term,
+    char **lower_corpus,
+    const double *doc_lengths,
+    double avg_doc_len,
+    double idf,
+    int *scores,
+    int num_docs,
+    int state
+){ }
+"""
+    anns = parse_cpp_annotations(source)
+    assert len(anns) == 1
+    ann = anns[0]
+    assert ann.func_name == "score_term_over_docs"
+    assert ann.transform_type == "func_substitute"
+    assert ann.state_indices == [7]
+    assert ann.state_function == "approx_state_identity"
+    assert ann.thresholds == [2000]
+    assert ann.thresholds_lower == [1]
+    assert ann.thresholds_upper == [40]
+    assert ann.decisions == [0, 0]
+    assert ann.decision_values == [0, 1, 2]
+
+
+def test_pragma_loop_perforate():
+    source = """
+#pragma approx decision_tree transform_type=loop_perforate state_indices=[5] thresholds=[8] decisions=[0,0] decision_values=[0,1,2]
+int choose_cluster(const double *point, double **centroids, int k, int dim, int dist_state, int state) {
+    return 0;
+}
+"""
+    anns = parse_cpp_annotations(source)
+    ann = anns[0]
+    assert ann.func_name == "choose_cluster"
+    assert ann.transform_type == "loop_perforate"
+    assert ann.state_indices == [5]
+
+
+def test_pragma_task_skipping():
+    source = """
+#pragma approx decision_tree transform_type=task_skipping state_indices=[2] thresholds=[2] decisions=[1,2] decision_values=[0,1,2]
+void model_choose(int input, int* output, int state) { }
+"""
+    anns = parse_cpp_annotations(source)
+    ann = anns[0]
+    assert ann.func_name == "model_choose"
+    assert ann.transform_type == "task_skipping"
+    assert ann.decisions == [1, 2]
+
+
+def test_pragma_and_comment_mixed():
+    source = """
+#pragma approx decision_tree transform_type=loop_perforate state_indices=[-1] thresholds=[1] decisions=[0,1]
+int a(int x, int state) { return x; }
+
+// @approx:decision_tree {
+//   transform_type: task_skipping
+//   thresholds: [2]
+//   decisions: [0, 1]
+// }
+void b(int x, int state) { }
+"""
+    anns = parse_cpp_annotations(source)
+    assert len(anns) == 2
+    assert anns[0].func_name == "a"
+    assert anns[0].transform_type == "loop_perforate"
+    assert anns[1].func_name == "b"
+    assert anns[1].transform_type == "task_skipping"
+
+
+def test_pragma_generates_same_mlir_as_comment():
+    pragma_source = """
+#pragma approx decision_tree transform_type=func_substitute thresholds=[1] decisions=[0,1]
+int kernel(int x, int state) { return x; }
+"""
+    comment_source = """
+// @approx:decision_tree {
+//   transform_type: func_substitute
+//   thresholds: [1]
+//   decisions: [0, 1]
+// }
+int kernel(int x, int state) { return x; }
+"""
+    pragma_anns = parse_cpp_annotations(pragma_source)
+    comment_anns = parse_cpp_annotations(comment_source)
+    pragma_mlir = generate_cpp_annotation_mlir(pragma_anns)
+    comment_mlir = generate_cpp_annotation_mlir(comment_anns)
+    assert pragma_mlir == comment_mlir