File size: 1,844 Bytes
a2ffabc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | from arm_gym.multi_agent import (
OpportunityToken,
analyzer_credit,
build_optimizer_prompt,
parse_analyzer_output,
per_role_reward,
)
def test_parses_strict_json_array():
out = '[{"kind":"vectorize_neon","target_region":"loop","rationale":"wide"}]'
toks = parse_analyzer_output(out)
assert len(toks) == 1
assert toks[0].kind == "vectorize_neon"
def test_parses_json_embedded_in_prose():
out = "thinking... [{\"kind\":\"use_ldp_stp\"}] done"
toks = parse_analyzer_output(out)
assert toks and toks[0].kind == "use_ldp_stp"
def test_ignores_garbage():
assert parse_analyzer_output("no json here") == []
assert parse_analyzer_output("[not, valid, json]") == []
def test_optimizer_prompt_contains_all_sections():
prompt = build_optimizer_prompt(
"int k(){return 0;}",
".text\nret",
[OpportunityToken("vectorize_neon")],
)
assert "C Code:" in prompt
assert "Baseline Assembly" in prompt
assert "Opportunity Tokens:" in prompt
assert "vectorize_neon" in prompt
assert "<assembly>" in prompt
def test_analyzer_credit_when_suggestion_applied():
toks = [OpportunityToken("use_ldp_stp")]
asm_with = "ldp x0, x1, [sp]\nret"
asm_without = "ldr x0, [sp]\nret"
assert analyzer_credit(toks, asm_with) == 1.0
assert analyzer_credit(toks, asm_without) == 0.0
def test_per_role_reward_splits():
toks = [OpportunityToken("fuse_multiply_add")]
asm = "fmadd s0, s1, s2, s3\nret"
split = per_role_reward(1.0, toks, asm)
assert split["optimizer"] == 1.0
assert split["analyzer"] > 0.0
def test_per_role_reward_zero_when_suggestions_unused():
toks = [OpportunityToken("fuse_multiply_add")]
asm = "ldr x0, [sp]\nret"
split = per_role_reward(1.0, toks, asm)
assert split["analyzer"] == 0.0
|