vyre-lower 0.6.5

Substrate-neutral lowering: vyre Program → KernelDescriptor consumed by vyre-emit-* crates.
Documentation 
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
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91

//! Shared recursive load-site counting for memory-placement analyses.
//!
//! Texture promotion, AoS-to-SoA layout rewrites, and later cache-placement
//! analyses all need the same conservative traversal over structured kernel
//! bodies. Keeping the traversal here prevents each analysis from guessing
//! which operands are child-body IDs.

use rustc_hash::FxHashMap;

use super::child_body_operands;
use crate::{KernelBody, KernelOpKind};

pub(crate) fn count_global_loads_by_slot<F>(
    body: &KernelBody,
    is_eligible_slot: &F,
    counts: &mut FxHashMap<u32, u32>,
) where
    F: Fn(u32) -> bool,
{
    for op in &body.ops {
        if matches!(op.kind, KernelOpKind::LoadGlobal) {
            if let Some(slot) = op.operands.first() {
                if is_eligible_slot(*slot) {
                    *counts.entry(*slot).or_insert(0) += 1;
                }
            }
        }
        for child_id in child_body_operands(&op.kind, &op.operands) {
            if let Some(child) = body.child_bodies.get(child_id as usize) {
                count_global_loads_by_slot(child, is_eligible_slot, counts);
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use rustc_hash::FxHashMap;

    use super::*;
    use crate::{KernelOp, LiteralValue};

    fn body_with_load(slot: u32) -> KernelBody {
        KernelBody {
            ops: vec![KernelOp {
                kind: KernelOpKind::LoadGlobal,
                operands: vec![slot, 0],
                result: Some(slot),
            }],
            child_bodies: vec![],
            literals: vec![],
        }
    }

    #[test]
    fn counts_if_else_children_and_ignores_for_loop_bound_operands() {
        let body = KernelBody {
            ops: vec![
                KernelOp {
                    kind: KernelOpKind::StructuredForLoop {
                        loop_var: "i".into(),
                    },
                    operands: vec![0, 1, 2],
                    result: None,
                },
                KernelOp {
                    kind: KernelOpKind::StructuredIfThenElse,
                    operands: vec![99, 3, 4],
                    result: None,
                },
            ],
            child_bodies: vec![
                body_with_load(7),
                body_with_load(7),
                body_with_load(7),
                body_with_load(7),
                body_with_load(7),
            ],
            literals: vec![LiteralValue::U32(0)],
        };

        let mut counts = FxHashMap::default();
        count_global_loads_by_slot(&body, &|slot| slot == 7, &mut counts);

        assert_eq!(
            counts.get(&7).copied(),
            Some(3),
            "Fix: load counting must include real structured child bodies without treating loop bound operands as child indices."
        );
    }
}