Scan the TTGIR and return (line_index, record_text) pairs. Each pair means "insert this proton.record line BEFORE lines[line_index]".
(lines: list[str])
| 75 | |
| 76 | |
| 77 | def _find_insertion_points(lines: list[str]) -> list[tuple[int, str]]: |
| 78 | """Scan the TTGIR and return (line_index, record_text) pairs. |
| 79 | |
| 80 | Each pair means "insert this proton.record line BEFORE lines[line_index]". |
| 81 | """ |
| 82 | insertions: list[tuple[int, str]] = [] |
| 83 | |
| 84 | # ── Kernel ── |
| 85 | # Full kernel scope from entry to return. |
| 86 | func_line = _find_first(lines, "tt.func public") |
| 87 | return_line = _find_first(lines, "tt.return") |
| 88 | if func_line is not None: |
| 89 | insertions.append((func_line + 1, 'start "kernel"')) |
| 90 | if return_line is not None: |
| 91 | insertions.append((return_line, 'end "kernel"')) |
| 92 | |
| 93 | # ── Setup ── |
| 94 | # From kernel entry (after tt.func) to just before the persistent scf.for. |
| 95 | # Covers: temperature load, grid dims, TMA descriptors, first tile prefetch. |
| 96 | persistent_for = _find_first(lines, "scf.for") |
| 97 | if func_line is not None: |
| 98 | insertions.append((func_line + 1, 'start "setup"')) |
| 99 | if persistent_for is not None: |
| 100 | insertions.append((persistent_for, 'end "setup"')) |
| 101 | |
| 102 | # Find the epilogue scf.if: the one whose body contains tt.reduce. |
| 103 | # This is the guard that fires when a tile's matmul is complete. |
| 104 | epilogue_if = _find_epilogue_if(lines) |
| 105 | if epilogue_if is None: |
| 106 | print("WARNING: could not find epilogue scf.if", file=sys.stderr) |
| 107 | return insertions |
| 108 | |
| 109 | # ── Mask ── |
| 110 | # V-masking and temperature scaling (first ops in the epilogue). |
| 111 | divf_line = _find_first_after(lines, "arith.divf", epilogue_if) |
| 112 | if epilogue_if is not None: |
| 113 | insertions.append((epilogue_if + 1, 'start "mask"')) |
| 114 | if divf_line is not None: |
| 115 | insertions.append((divf_line + 1, 'end "mask"')) |
| 116 | |
| 117 | # ── Tile-mgmt ── |
| 118 | # Next-tile coordinate computation (swizzle grouping). Sits between |
| 119 | # temperature scaling and Gumbel noise seed. In the DSL, h_start_c is the |
| 120 | # last tile-mgmt variable before the sample loop. The compiler preserves |
| 121 | # this name in TTGIR as "%h_start_c = arith.muli ...". |
| 122 | h_start_line = _find_first_after(lines, "%h_start_c =", epilogue_if) |
| 123 | if divf_line is not None and h_start_line is not None: |
| 124 | insertions.append((divf_line + 1, 'start "tile-mgmt"')) |
| 125 | insertions.append((h_start_line + 1, 'end "tile-mgmt"')) |
| 126 | |
| 127 | # ── Sample ── |
| 128 | # Gumbel noise generation through tt.reduce (argmax). |
| 129 | # Starts right after tile-mgmt (seed offset + Philox RNG + log transform). |
| 130 | reduce_close = _find_reduce_close(lines) |
| 131 | if h_start_line is not None: |
| 132 | insertions.append((h_start_line + 1, 'start "sample"')) |
| 133 | if reduce_close is not None: |
| 134 | insertions.append((reduce_close + 1, 'end "sample"')) |
no test coverage detected