vyre-lower 0.6.5

//! Output type for the coalescence analysis.

use serde::{Deserialize, Serialize};

use crate::analyses::AccessKind;

/// Categorization of one global-memory access site's index pattern
/// across the workgroup's threads.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum AccessPattern {
    /// Index = `base + LocalInvocationId.x`. Thread `t` reads address
    /// `base + t`. Hardware coalesces into one transaction. Best case.
    CoalescedUnitStride,
    /// Index = `base + k * LocalInvocationId.x` for some constant `k > 1`.
    /// Each thread reads a separate transaction; throughput drops by
    /// roughly `k`× on most architectures.
    Strided { stride: u32 },
    /// Index depends on data we can't prove constant-stride (indirect
    /// load, conditional offset, runtime-computed). Treated as
    /// scattered for cost purposes; rewrites that target this category
    /// usually need shared-memory promotion (PERF B12) or explicit
    /// gather/scatter primitives.
    Scattered,
    /// Index is a literal constant  -  every thread reads the same
    /// address. Hardware broadcasts; this is fine for reads, wasteful
    /// for writes (last-writer-wins with race semantics).
    Broadcast,
}

impl AccessPattern {
    /// Estimated throughput multiplier vs the ideal coalesced case.
    /// Coalesced = 1.0; lower = slower.
    #[must_use]
    pub fn throughput_factor(&self) -> f32 {
        match self {
            Self::CoalescedUnitStride => 1.0,
            Self::Strided { stride } => 1.0 / (*stride as f32),
            // Scattered and broadcast costs are workload-dependent;
            // these are pessimistic order-of-magnitude defaults that
            // err on the side of "this is bad" for prioritization.
            Self::Scattered => 1.0 / 32.0,
            Self::Broadcast => 1.0,
        }
    }

    #[must_use]
    pub const fn is_problematic(&self) -> bool {
        matches!(self, Self::Strided { .. } | Self::Scattered)
    }
}

/// One global-memory access site identified during analysis.
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct AccessSite {
    /// Index of the op in the kernel body's flat `ops` Vec. Identifies
    /// which load/store this report refers to.
    pub op_index: usize,
    /// Whether this is a load or store. Both are analyzed because
    /// store coalescing is just as expensive as load coalescing on
    /// every substrate.
    pub kind: AccessKind,
    /// Binding slot the access reads/writes.
    pub binding_slot: u32,
    /// Classified access pattern.
    pub pattern: AccessPattern,
}

/// Rewrite class recommended for a non-ideal memory access site.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum CoalescenceRewrite {
    /// Constant-stride access: prefer vec2/vec4 packing, layout rewrite,
    /// or tile transposition so adjacent lanes hit adjacent words.
    VectorPackOrTile,
    /// Runtime-dependent access: prefer shared-memory staging or an
    /// explicit gather/scatter primitive with a measured cost model.
    SharedMemoryGather,
    /// Every lane writes the same address. This is normally a race or
    /// a last-writer-wins reduction smell; use an atomic/reduction op.
    BroadcastStoreReduction,
}

/// Structured warning emitted from a coalescence report.
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct CoalescenceWarning {
    /// Stable warning code for tooling.
    pub code: &'static str,
    /// Source access site.
    pub site: AccessSite,
    /// Rewrite family that should handle this site.
    pub rewrite: CoalescenceRewrite,
    /// Human-readable, fix-oriented message.
    pub message: String,
}

/// Full report for one `KernelDescriptor`.
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct CoalescenceReport {
    /// Stable kernel id (echoes `KernelDescriptor::id`).
    pub kernel_id: String,
    /// Every global access site, in body-order.
    pub sites: Vec<AccessSite>,
}

impl CoalescenceReport {
    /// Number of sites with `pattern.is_problematic()`. A perfect
    /// kernel has `problematic_count() == 0`.
    #[must_use]
    pub fn problematic_count(&self) -> usize {
        self.sites
            .iter()
            .filter(|s| s.pattern.is_problematic())
            .count()
    }

    /// Sum of `1 - throughput_factor` across all sites  -  a single
    /// score where 0.0 means perfect and higher means worse. Useful
    /// for sorting kernels by which to optimize first.
    #[must_use]
    pub fn waste_score(&self) -> f32 {
        self.sites
            .iter()
            .map(|s| 1.0 - s.pattern.throughput_factor())
            .sum()
    }

    /// Convert raw access classifications into structured warnings
    /// that emitters, benches, and CI gates can consume directly.
    ///
    /// This is the B14 warning path: analyses stay backend-neutral,
    /// while each backend decides whether to apply a rewrite or emit
    /// a native diagnostic from the returned warning.
    #[must_use]
    pub fn warnings(&self) -> Vec<CoalescenceWarning> {
        self.sites
            .iter()
            .filter_map(|site| warning_for_site(site))
            .collect()
    }
}

fn warning_for_site(site: &AccessSite) -> Option<CoalescenceWarning> {
    match (site.kind, site.pattern) {
        (AccessKind::Load | AccessKind::Store, AccessPattern::Strided { stride }) => {
            Some(CoalescenceWarning {
                code: "W-COALESCE-STRIDED",
                site: site.clone(),
                rewrite: CoalescenceRewrite::VectorPackOrTile,
                message: format!(
                    "Fix: binding {} {:?} at op {} uses stride {stride}; apply vector packing, tiling, or layout rewrite before emission.",
                    site.binding_slot, site.kind, site.op_index
                ),
            })
        }
        (AccessKind::Load | AccessKind::Store, AccessPattern::Scattered) => {
            Some(CoalescenceWarning {
                code: "W-COALESCE-SCATTERED",
                site: site.clone(),
                rewrite: CoalescenceRewrite::SharedMemoryGather,
                message: format!(
                    "Fix: binding {} {:?} at op {} is scattered; stage through shared memory or use a measured gather/scatter primitive.",
                    site.binding_slot, site.kind, site.op_index
                ),
            })
        }
        (AccessKind::Store, AccessPattern::Broadcast) => Some(CoalescenceWarning {
            code: "W-COALESCE-BROADCAST-STORE",
            site: site.clone(),
            rewrite: CoalescenceRewrite::BroadcastStoreReduction,
            message: format!(
                "Fix: binding {} Store at op {} writes one address from every lane; replace with a reduction, atomic, or single-lane guard.",
                site.binding_slot, site.op_index
            ),
        }),
        _ => None,
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn coalesced_unit_stride_factor_is_one() {
        assert!((AccessPattern::CoalescedUnitStride.throughput_factor() - 1.0).abs() < 1e-6);
    }

    #[test]
    fn strided_4_factor_is_one_quarter() {
        let f = AccessPattern::Strided { stride: 4 }.throughput_factor();
        assert!((f - 0.25).abs() < 1e-6);
    }

    #[test]
    fn scattered_factor_is_one_thirty_second() {
        let f = AccessPattern::Scattered.throughput_factor();
        assert!((f - 1.0 / 32.0).abs() < 1e-6);
    }

    #[test]
    fn coalesced_and_broadcast_not_problematic() {
        assert!(!AccessPattern::CoalescedUnitStride.is_problematic());
        assert!(!AccessPattern::Broadcast.is_problematic());
    }

    #[test]
    fn strided_and_scattered_problematic() {
        assert!(AccessPattern::Strided { stride: 2 }.is_problematic());
        assert!(AccessPattern::Scattered.is_problematic());
    }

    #[test]
    fn empty_report_has_no_problems_and_zero_waste() {
        let r = CoalescenceReport {
            kernel_id: "empty".into(),
            sites: vec![],
        };
        assert_eq!(r.problematic_count(), 0);
        assert!((r.waste_score() - 0.0).abs() < 1e-6);
    }

    #[test]
    fn waste_score_sums_per_site_costs() {
        let r = CoalescenceReport {
            kernel_id: "k".into(),
            sites: vec![
                AccessSite {
                    op_index: 0,
                    kind: AccessKind::Load,
                    binding_slot: 0,
                    pattern: AccessPattern::CoalescedUnitStride,
                },
                AccessSite {
                    op_index: 1,
                    kind: AccessKind::Load,
                    binding_slot: 0,
                    pattern: AccessPattern::Strided { stride: 4 },
                },
                AccessSite {
                    op_index: 2,
                    kind: AccessKind::Store,
                    binding_slot: 1,
                    pattern: AccessPattern::Scattered,
                },
            ],
        };
        // 0 (coalesced) + 0.75 (strided 4) + (1 - 1/32) (scattered)
        let expected = 0.0 + 0.75 + (1.0 - 1.0 / 32.0);
        assert!((r.waste_score() - expected).abs() < 1e-5);
        assert_eq!(r.problematic_count(), 2);
    }

    #[test]
    fn warnings_map_strided_and_scattered_to_rewrites() {
        let r = CoalescenceReport {
            kernel_id: "k".into(),
            sites: vec![
                AccessSite {
                    op_index: 1,
                    kind: AccessKind::Load,
                    binding_slot: 4,
                    pattern: AccessPattern::Strided { stride: 4 },
                },
                AccessSite {
                    op_index: 2,
                    kind: AccessKind::Store,
                    binding_slot: 5,
                    pattern: AccessPattern::Scattered,
                },
            ],
        };
        let warnings = r.warnings();
        assert_eq!(warnings.len(), 2);
        assert_eq!(warnings[0].code, "W-COALESCE-STRIDED");
        assert_eq!(warnings[0].rewrite, CoalescenceRewrite::VectorPackOrTile);
        assert!(warnings[0].message.contains("stride 4"));
        assert_eq!(warnings[1].code, "W-COALESCE-SCATTERED");
        assert_eq!(warnings[1].rewrite, CoalescenceRewrite::SharedMemoryGather);
    }

    #[test]
    fn warnings_include_broadcast_store_but_not_broadcast_load() {
        let r = CoalescenceReport {
            kernel_id: "k".into(),
            sites: vec![
                AccessSite {
                    op_index: 1,
                    kind: AccessKind::Load,
                    binding_slot: 0,
                    pattern: AccessPattern::Broadcast,
                },
                AccessSite {
                    op_index: 2,
                    kind: AccessKind::Store,
                    binding_slot: 1,
                    pattern: AccessPattern::Broadcast,
                },
            ],
        };
        let warnings = r.warnings();
        assert_eq!(warnings.len(), 1);
        assert_eq!(warnings[0].code, "W-COALESCE-BROADCAST-STORE");
        assert_eq!(
            warnings[0].rewrite,
            CoalescenceRewrite::BroadcastStoreReduction
        );
    }

    #[test]
    fn report_round_trips_serde_byte_stable() {
        let r = CoalescenceReport {
            kernel_id: "rt".into(),
            sites: vec![AccessSite {
                op_index: 0,
                kind: AccessKind::Load,
                binding_slot: 7,
                pattern: AccessPattern::Strided { stride: 8 },
            }],
        };
        // Use a JSON round-trip for serde stability check; we don't
        // depend on serde_json directly, so do this via the lower
        // crate's serde infra: constructed values should equal after
        // bincode-style serialize-then-deserialize. Simplest assertion
        // here is structural equality after clone  -  the strong
        // round-trip is exercised in vyre-lower's descriptor tests.
        let r2 = r.clone();
        assert_eq!(r, r2);
    }
}