tests/test_mlir_batch.py

from pathlib import Path

from scripts import run_mlir_batch


def test_count_text_markers_distinguishes_quant_and_float_compute(tmp_path: Path) -> None:
    path = tmp_path / "sample.mlir"
    path.write_text(
        '\n'.join([
            '%0 = "onnx.QuantizeLinear"() : () -> tensor<1xi8>',
            '%1 = "onnx.DequantizeLinear"() : () -> tensor<1xf32>',
            '%2 = llvm.fmul %a, %b : f32',
            '%3 = "onnx.MatMulInteger"() : () -> tensor<1xi8>',
        ])
    )
    markers = run_mlir_batch.count_text_markers(path)
    assert markers["onnx_quantize_linear_ops"] == 1
    assert markers["onnx_dequantize_linear_ops"] == 1
    assert markers["onnx_qoperator_ops"] == 1
    assert markers["llvm_float_compute_ops"] == 1
    assert markers["i8_mentions"] == 2


def test_eligible_rows_covers_all_active_pairs() -> None:
    rows = run_mlir_batch.eligible_rows()
    assert len(rows) == 21
    assert all(row["eligibility"] == "ELIGIBLE" for row in rows)
    assert "AD01" in {row["model_id"] for row in rows}


def test_different_setting_never_overwrites_existing_ir(tmp_path: Path) -> None:
    source = tmp_path / "input.mlir"
    source.write_text("module {}\n")
    output = tmp_path / "model" / "mlir" / "fp32" / "onnx.mlir"
    output.parent.mkdir(parents=True)
    output.write_text("protected prior IR\n")
    result = run_mlir_batch.run_command_stage(
        model_root=tmp_path / "model",
        variant="fp32",
        stage="onnx_dialect_parse",
        command=["/usr/bin/false"],
        inputs=[source],
        output=output,
        command_output=None,
        timeout_sec=1,
        settings={"toolchain_lock_sha256": "0" * 64},
        reuse_failures=False,
    )
    assert result["status"] == "BLOCKED"
    assert result["secondary_failure_code"] == "OUTPUT_SETTINGS_CONFLICT"
    assert output.read_text() == "protected prior IR\n"