Collect kernel artifacts: device-IR node-link dump (.ir.jsonl)

[IshanAryendu]

Collect kernel artifacts: device-IR node-link dump (.ir.jsonl)

Summary

Adds a lossless dump of Helion's device IR to the autotune telemetry, as the ir_features artifact for the kernel-artifact dataset. For each autotune run, the union of all device-side torch.fx graphs is serialized as a node-link graph (one JSON record per run) to a new <autotune_log>.ir.jsonl sidecar, alongside the existing .csv and .meta.jsonl, and joined to them on run_id.

Stacked on #2737 ("Collect kernel artifacts: append-mode autotune telemetry with run_id"). PR 2737 adds the append-mode sink, .meta.jsonl, and the run_id join key this one builds on.

Motivation

The cost-model dataset needs the kernel's computation graph as features, not just source text + perf. Device IR is the structured representation Helion lowers to before config-specific codegen, so it's the right artifact to capture. It is config-independent (built once per bound kernel), so it is collected once per run_id and joins cleanly to the per-config CSV rows and the per-run identity record.

What it does

• New module helion/autotuner/ir_features.py with extract_ir_graph(device_ir, *, run_id, kernel_id, kernel_name, input_shapes) -> IrGraphRecord — a plain-JSON, networkx-node-link-compatible dump (can be parsed using: nx.node_link_graph(record, edges="links")).
• Only-true edges are captured:
  • data edges from node.args/kwargs (with arg_positions, via fx map_arg for covering all container types).
  • region edges only from live control-flow call nodes (_for_loop/_for_loop_step/_if/_while_loop), mapping their live argument lists onto the child graph's placeholders.
  • Reduction loops are rolled alternates, not called subgraphs, so they are not given fabricated dataflow edges. That relationship is recorded as typed graph-level metadata (rolled_reductions) instead.
• Per-node features (best-effort, None when absent): op_kind, target (address-free), lowering_class, dtype/shape/concrete_shape/stride/concrete_stride/device/value, operand input_dtypes/input_shapes, reduction_type/reduction_ranges, pointwise_ranges, block_ids, graph_id, region_kind, source_loc. Symbolic dims are stringified with best-effort concrete resolution and large values are length-capped.
• Typed schema via TypedDicts (IrNode composed from ValFeatures + LoweringFeatures + node-core; IrEdge, IrGraphMeta, IrGraphRecord) plus a schema_version field (IR_SCHEMA_VERSION = 1) and a documented consumer gate (record.get("schema_version", 0)).
• Wiring: AutotuneLogSink appends one .ir.jsonl record per run. BaseSearch._prepare extracts the device IR once, gated on autotune_log and best-effort (a missing/odd host_function or non-Triton backend degrades to "no IR artifact" and never breaks autotuning).
• Drift safety: unknown GraphInfo region kinds and malformed control-flow nodes log a one-time warning (never raise). A unit test asserts the known-region-kinds set matches the concrete GraphInfo subclasses, so upstream IR changes fail CI loudly.

Resulting artifact (per run_id)

<base>.ir.jsonl   {schema_version, run_id, directed, multigraph,
                   graph{run_id, kernel_id, kernel_name, input_shapes, num_graphs,
                         root_ids, graphs[], rolled_reductions[]},
                   nodes[], links[]}

Join to .meta.jsonl and .csv on run_id; group by kernel_id.

Testing

• test/test_ir_features.py: pointwise/reduction/control-flow extraction on real kernels, networkx round-trip, control-flow region-edge resolution, rolled-reduction metadata (no fake edges), and negative/edge cases (missing/scalar val, no lowering, address-free target, unconvertible dim, malformed/non-control-flow region specs warn behavior, empty device IR, bool graph-id rejection, value/source_loc capping, schema_version gate, region-kind drift guard).
• test/test_kernel_metadata.py: .ir.jsonl is appended and joins .meta.jsonl + CSV rows on run_id; metadata-without-ir-graph writes no .ir.jsonl.

Validation

• ruff clean; pyrefly 0 errors on helion/ (CI excludes test/).
• test_ir_features + test_kernel_metadata + test_autotuner all pass (no autotuner regressions).
• End-to-end on add/softmax/attention: all three sidecars produced, networkx rebuilds every record, run_id joins .ir.jsonl ↔ .meta.jsonl ↔ .csv, with correct per-kernel structure (pointwise = data edges only; softmax/layer_norm = rolled_reductions, no fake edges; bmm/attention = resolved _for_loop region edges).