`939f7e4`

Test WHERE with constant-scalar store value vectorizes at O3

Authored by mfwolffe <wolffemf@dukes.jmu.edu> 1 week ago

SHA: 939f7e443c3fb7b5a310852a2898212a4ac344bf
Parents: 131e2f1
Tree: 66aaafa

2 changed files

Status	File	+	-
A	`test_programs/do_loop_vectorize_where_const.f90`	26	0
A	`tests/vectorize_where_const.rs`	73	0

test_programs/do_loop_vectorize_where_const.f90added

 +! WHERE with a literal-constant store value (`b = K`). The matcher
 +! must accept a store value that is loop-invariant constant — no load,
 +! no unary, no binop in the then_block other than the Store itself.
 +! Lifts to vselect(mask, BroadcastK_vec, dest_vec) → store.
 +!
 +! a(i) = i - 16 (range -15..16); b(i) = i. Where a > 0 (lanes 17..32),
 +! set b = -1.0; otherwise leave b = i.
 +!   b(1)  = 1.0   (a(1)=-15, mask=false, b unchanged)
 +!   b(16) = 16.0  (a(16)=0, NOT > 0, mask=false)
 +!   b(17) = -1.0  (a(17)=1, mask=true, b ← -1.0)
 +!   b(32) = -1.0  (a(32)=16, mask=true, b ← -1.0)
 +!
 +! CHECK: 1.0000000E0     1.6000000E1    -1.0000000E0    -1.0000000E0
 +program test_do_loop_vectorize_where_const
 +  implicit none
 +  integer :: i
 +  real(4) :: a(32), b(32)
 +  do i = 1, 32
 +    a(i) = real(i - 16, 4)
 +    b(i) = real(i, 4)
 +  end do
 +  where (a > 0.0)
 +    b = -1.0
 +  end where
 +  print *, b(1), b(16), b(17), b(32)
 +end program test_do_loop_vectorize_where_const

tests/vectorize_where_const.rsadded

 +use std::collections::BTreeSet;
 +use std::path::PathBuf;
++
 +use armfortas::driver::OptLevel;
 +use armfortas::testing::{capture_from_path, CaptureRequest, CapturedStage, Stage};
++
 +fn fixture(name: &str) -> PathBuf {
 +    let path = PathBuf::from("test_programs").join(name);
 +    assert!(path.exists(), "missing test fixture {}", path.display());
 +    path
 +}
++
 +fn capture_text(request: CaptureRequest, stage: Stage) -> String {
 +    let result = capture_from_path(&request).expect("capture should succeed");
 +    match result.get(stage) {
 +        Some(CapturedStage::Text(text)) => text.clone(),
 +        Some(CapturedStage::Run(_)) => panic!("expected text stage for {}", stage.as_str()),
 +        None => panic!("missing requested stage {}", stage.as_str()),
 +    }
 +}
++
 +fn capture_run_stdout(request: CaptureRequest) -> String {
 +    let result = capture_from_path(&request).expect("capture should succeed");
 +    match result.get(Stage::Run) {
 +        Some(CapturedStage::Run(run)) => run.stdout.clone(),
 +        _ => panic!("missing run stage"),
 +    }
 +}
++
 +#[test]
 +fn o3_vectorizes_where_with_constant_scalar_store() {
 +    let source = fixture("do_loop_vectorize_where_const.f90");
++
 +    let o3_ir = capture_text(
 +        CaptureRequest {
 +            input: source.clone(),
 +            requested: BTreeSet::from([Stage::OptIr]),
 +            opt_level: OptLevel::O3,
 +        },
 +        Stage::OptIr,
 +    );
 +    // The constant-scalar WHERE store should lift to a VBroadcast
 +    // feeding the vselect's true arm.
 +    assert_eq!(
 +        o3_ir.matches("vselect").count(),
 +        1,
 +        "expected exactly one vselect:\n{}",
 +        o3_ir
 +    );
 +    assert!(
 +        o3_ir.contains("vbroadcast") || o3_ir.contains("VBroadcast"),
 +        "expected vbroadcast for const scalar in WHERE body:\n{}",
 +        o3_ir
 +    );
++
 +    let stdout = capture_run_stdout(CaptureRequest {
 +        input: source,
 +        requested: BTreeSet::from([Stage::Run]),
 +        opt_level: OptLevel::O3,
 +    });
 +    let trimmed: Vec<&str> = stdout
 +        .lines()
 +        .map(|l| l.trim())
 +        .filter(|l| !l.is_empty())
 +        .collect();
 +    assert_eq!(trimmed.len(), 1, "expected one output line:\n{}", stdout);
 +    assert_eq!(
 +        trimmed[0],
 +        "1.0000000E0     1.6000000E1    -1.0000000E0    -1.0000000E0",
 +        "const WHERE wrong: {:?}",
 +        trimmed[0]
 +    );
 +}