use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
use crate::{KernelBody, KernelDescriptor, KernelOpKind, LiteralValue};
use vyre_foundation::ir::BinOp;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct IntRange {
pub min: i64,
pub max: i64,
}
impl IntRange {
pub fn singleton(v: i64) -> Self {
Self { min: v, max: v }
}
pub fn is_singleton(&self) -> bool {
self.min == self.max
}
pub fn contains(&self, v: i64) -> bool {
v >= self.min && v <= self.max
}
pub fn union(self, other: Self) -> Self {
Self {
min: self.min.min(other.min),
max: self.max.max(other.max),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
pub struct ValueRangeReport {
pub ranges: BTreeMap<u32, IntRange>,
}
impl ValueRangeReport {
pub fn known_count(&self) -> usize {
self.ranges.len()
}
pub fn get(&self, result_id: u32) -> Option<IntRange> {
self.ranges.get(&result_id).copied()
}
pub fn is_definitely(&self, result_id: u32, target: i64) -> Option<bool> {
self.ranges
.get(&result_id)
.map(|r| r.is_singleton() && r.min == target)
}
pub fn is_definitely_below(&self, result_id: u32, target: i64) -> Option<bool> {
self.ranges.get(&result_id).map(|r| r.max < target)
}
pub fn is_definitely_at_least(&self, result_id: u32, target: i64) -> Option<bool> {
self.ranges.get(&result_id).map(|r| r.min >= target)
}
pub fn as_constant(&self, result_id: u32) -> Option<i64> {
self.ranges
.get(&result_id)
.filter(|r| r.is_singleton())
.map(|r| r.min)
}
}
#[must_use]
pub fn analyze(desc: &KernelDescriptor) -> ValueRangeReport {
analyze_body(&desc.body)
}
#[must_use]
pub fn analyze_body(body: &KernelBody) -> ValueRangeReport {
let mut ranges: BTreeMap<u32, IntRange> = BTreeMap::new();
for op in &body.ops {
if matches!(op.kind, KernelOpKind::Literal) {
if let (Some(rid), Some(&pool_idx)) = (op.result, op.operands.first()) {
if let Some(lit) = body.literals.get(pool_idx as usize) {
let r = match lit {
LiteralValue::U32(v) => Some(IntRange::singleton(*v as i64)),
LiteralValue::I32(v) => Some(IntRange::singleton(*v as i64)),
LiteralValue::Bool(true) => Some(IntRange::singleton(1)),
LiteralValue::Bool(false) => Some(IntRange::singleton(0)),
_ => None,
};
if let Some(r) = r {
ranges.insert(rid, r);
}
}
}
}
}
for op in &body.ops {
if let KernelOpKind::BinOpKind(bin_op) = &op.kind {
if op.operands.len() < 2 {
continue;
}
let lhs = ranges.get(&op.operands[0]).copied();
let rhs = ranges.get(&op.operands[1]).copied();
let Some(rid) = op.result else { continue };
if let (Some(l), Some(r)) = (lhs, rhs) {
let derived = match bin_op {
BinOp::Min => Some(IntRange {
min: l.min.min(r.min),
max: l.max.min(r.max),
}),
BinOp::Max => Some(IntRange {
min: l.min.max(r.min),
max: l.max.max(r.max),
}),
BinOp::Add | BinOp::WrappingAdd => {
match (l.min.checked_add(r.min), l.max.checked_add(r.max)) {
(Some(min), Some(max)) => Some(IntRange { min, max }),
_ => None,
}
}
BinOp::Sub | BinOp::WrappingSub => {
match (l.min.checked_sub(r.max), l.max.checked_sub(r.min)) {
(Some(min), Some(max)) => Some(IntRange { min, max }),
_ => None,
}
}
BinOp::Mul => mul_range(l, r),
BinOp::BitAnd => {
if l.min < 0 || r.min < 0 {
None
} else {
Some(IntRange {
min: 0,
max: l.max.min(r.max),
})
}
}
BinOp::BitOr => {
if l.min < 0 || r.min < 0 {
None
} else {
Some(IntRange {
min: l.min.max(r.min),
max: l.max | r.max,
})
}
}
BinOp::Shl if r.is_singleton() && r.min >= 0 && r.min < 32 => {
let k = r.min as u32;
match (l.min.checked_shl(k), l.max.checked_shl(k)) {
(Some(min), Some(max)) => Some(IntRange { min, max }),
_ => None,
}
}
BinOp::Shr if r.is_singleton() && r.min >= 0 && r.min < 32 => {
let k = r.min as u32;
Some(IntRange {
min: l.min >> k,
max: l.max >> k,
})
}
_ => None,
};
if let Some(d) = derived {
ranges.insert(rid, d);
}
}
}
}
ValueRangeReport { ranges }
}
fn mul_range(l: IntRange, r: IntRange) -> Option<IntRange> {
let corners = [
l.min.checked_mul(r.min),
l.min.checked_mul(r.max),
l.max.checked_mul(r.min),
l.max.checked_mul(r.max),
];
if corners.iter().any(|c| c.is_none()) {
return None;
}
let vals: Vec<i64> = corners.into_iter().map(|c| c.unwrap()).collect();
let min = *vals.iter().min().unwrap();
let max = *vals.iter().max().unwrap();
Some(IntRange { min, max })
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{
BindingLayout, Dispatch, KernelBody, KernelDescriptor, KernelOp, KernelOpKind, LiteralValue,
};
fn build(ops: Vec<KernelOp>, lits: Vec<LiteralValue>) -> KernelDescriptor {
KernelDescriptor {
id: "k".into(),
bindings: BindingLayout { slots: vec![] },
dispatch: Dispatch::new(1, 1, 1),
body: KernelBody {
ops,
child_bodies: vec![],
literals: lits,
},
}
}
#[test]
fn empty_kernel_no_ranges() {
let r = analyze(&build(vec![], vec![]));
assert!(r.ranges.is_empty());
assert_eq!(r.known_count(), 0);
}
#[test]
fn lit_u32_yields_singleton() {
let desc = build(
vec![KernelOp {
kind: KernelOpKind::Literal,
operands: vec![0],
result: Some(0),
}],
vec![LiteralValue::U32(42)],
);
let r = analyze(&desc);
assert_eq!(r.ranges[&0], IntRange::singleton(42));
assert!(r.ranges[&0].is_singleton());
}
#[test]
fn lit_i32_negative_yields_correct_range() {
let desc = build(
vec![KernelOp {
kind: KernelOpKind::Literal,
operands: vec![0],
result: Some(0),
}],
vec![LiteralValue::I32(-7)],
);
let r = analyze(&desc);
assert_eq!(r.ranges[&0], IntRange::singleton(-7));
}
#[test]
fn bool_true_is_one_false_is_zero() {
let desc = build(
vec![
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![0],
result: Some(0),
},
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![1],
result: Some(1),
},
],
vec![LiteralValue::Bool(true), LiteralValue::Bool(false)],
);
let r = analyze(&desc);
assert_eq!(r.ranges[&0], IntRange::singleton(1));
assert_eq!(r.ranges[&1], IntRange::singleton(0));
}
#[test]
fn min_of_two_lits_propagates() {
let desc = build(
vec![
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![0],
result: Some(0),
},
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![1],
result: Some(1),
},
KernelOp {
kind: KernelOpKind::BinOpKind(BinOp::Min),
operands: vec![0, 1],
result: Some(2),
},
],
vec![LiteralValue::U32(3), LiteralValue::U32(5)],
);
let r = analyze(&desc);
assert_eq!(r.ranges[&2], IntRange::singleton(3));
}
#[test]
fn max_of_two_lits_propagates() {
let desc = build(
vec![
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![0],
result: Some(0),
},
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![1],
result: Some(1),
},
KernelOp {
kind: KernelOpKind::BinOpKind(BinOp::Max),
operands: vec![0, 1],
result: Some(2),
},
],
vec![LiteralValue::U32(3), LiteralValue::U32(5)],
);
let r = analyze(&desc);
assert_eq!(r.ranges[&2], IntRange::singleton(5));
}
#[test]
fn add_propagates_singleton_ranges() {
let desc = build(
vec![
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![0],
result: Some(0),
},
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![1],
result: Some(1),
},
KernelOp {
kind: KernelOpKind::BinOpKind(BinOp::Add),
operands: vec![0, 1],
result: Some(2),
},
],
vec![LiteralValue::U32(3), LiteralValue::U32(5)],
);
let r = analyze(&desc);
assert_eq!(r.ranges[&2], IntRange::singleton(8));
}
#[test]
fn sub_flips_operand_bounds() {
let desc = build(
vec![
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![0],
result: Some(0),
},
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![1],
result: Some(1),
},
KernelOp {
kind: KernelOpKind::BinOpKind(BinOp::Sub),
operands: vec![0, 1],
result: Some(2),
},
],
vec![LiteralValue::I32(10), LiteralValue::I32(3)],
);
let r = analyze(&desc);
assert_eq!(r.ranges[&2], IntRange::singleton(7));
}
#[test]
fn bitand_with_mask_bounds_to_zero_through_mask() {
let desc = build(
vec![
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![0],
result: Some(0),
},
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![1],
result: Some(1),
},
KernelOp {
kind: KernelOpKind::BinOpKind(BinOp::BitAnd),
operands: vec![0, 1],
result: Some(2),
},
],
vec![LiteralValue::U32(0x12345678), LiteralValue::U32(0xFF)],
);
let r = analyze(&desc);
assert_eq!(r.ranges[&2], IntRange { min: 0, max: 0xFF });
}
#[test]
fn shl_propagates_with_singleton_shift() {
let desc = build(
vec![
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![0],
result: Some(0),
},
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![1],
result: Some(1),
},
KernelOp {
kind: KernelOpKind::BinOpKind(BinOp::Shl),
operands: vec![0, 1],
result: Some(2),
},
],
vec![LiteralValue::U32(5), LiteralValue::U32(3)],
);
let r = analyze(&desc);
assert_eq!(r.ranges[&2], IntRange::singleton(40));
}
#[test]
fn shr_propagates_with_singleton_shift() {
let desc = build(
vec![
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![0],
result: Some(0),
},
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![1],
result: Some(1),
},
KernelOp {
kind: KernelOpKind::BinOpKind(BinOp::Shr),
operands: vec![0, 1],
result: Some(2),
},
],
vec![LiteralValue::U32(40), LiteralValue::U32(3)],
);
let r = analyze(&desc);
assert_eq!(r.ranges[&2], IntRange::singleton(5));
}
#[test]
fn shl_with_huge_shift_not_propagated() {
let desc = build(
vec![
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![0],
result: Some(0),
},
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![1],
result: Some(1),
},
KernelOp {
kind: KernelOpKind::BinOpKind(BinOp::Shl),
operands: vec![0, 1],
result: Some(2),
},
],
vec![LiteralValue::U32(1), LiteralValue::U32(64)],
);
let r = analyze(&desc);
assert!(!r.ranges.contains_key(&2));
}
#[test]
fn bitor_propagates_with_singletons() {
let desc = build(
vec![
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![0],
result: Some(0),
},
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![1],
result: Some(1),
},
KernelOp {
kind: KernelOpKind::BinOpKind(BinOp::BitOr),
operands: vec![0, 1],
result: Some(2),
},
],
vec![LiteralValue::U32(0xF0), LiteralValue::U32(0x0F)],
);
let r = analyze(&desc);
assert_eq!(
r.ranges[&2],
IntRange {
min: 0xF0,
max: 0xFF
}
);
}
#[test]
fn bitand_negative_operand_not_propagated() {
let desc = build(
vec![
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![0],
result: Some(0),
},
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![1],
result: Some(1),
},
KernelOp {
kind: KernelOpKind::BinOpKind(BinOp::BitAnd),
operands: vec![0, 1],
result: Some(2),
},
],
vec![LiteralValue::I32(-1), LiteralValue::I32(0xFF)],
);
let r = analyze(&desc);
assert!(!r.ranges.contains_key(&2));
}
#[test]
fn mul_singletons() {
let desc = build(
vec![
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![0],
result: Some(0),
},
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![1],
result: Some(1),
},
KernelOp {
kind: KernelOpKind::BinOpKind(BinOp::Mul),
operands: vec![0, 1],
result: Some(2),
},
],
vec![LiteralValue::I32(7), LiteralValue::I32(-3)],
);
let r = analyze(&desc);
assert_eq!(r.ranges[&2], IntRange::singleton(-21));
}
#[test]
fn mul_range_corner_helper() {
let r = mul_range(IntRange { min: 2, max: 5 }, IntRange { min: 3, max: 4 });
assert_eq!(r, Some(IntRange { min: 6, max: 20 }));
let r = mul_range(IntRange { min: -2, max: 3 }, IntRange { min: -1, max: 4 });
assert_eq!(r, Some(IntRange { min: -8, max: 12 }));
}
#[test]
fn add_chains_propagate() {
let desc = build(
vec![
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![0],
result: Some(0),
},
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![1],
result: Some(1),
},
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![2],
result: Some(2),
},
KernelOp {
kind: KernelOpKind::BinOpKind(BinOp::Add),
operands: vec![0, 1],
result: Some(3),
},
KernelOp {
kind: KernelOpKind::BinOpKind(BinOp::Add),
operands: vec![3, 2],
result: Some(4),
},
],
vec![
LiteralValue::U32(3),
LiteralValue::U32(5),
LiteralValue::U32(7),
],
);
let r = analyze(&desc);
assert_eq!(r.ranges[&4], IntRange::singleton(15));
}
#[test]
fn non_lit_op_no_range() {
let desc = build(
vec![KernelOp {
kind: KernelOpKind::LocalInvocationId,
operands: vec![0],
result: Some(0),
}],
vec![],
);
let r = analyze(&desc);
assert!(!r.ranges.contains_key(&0));
}
#[test]
fn as_constant_returns_value_for_singleton() {
let desc = build(
vec![KernelOp {
kind: KernelOpKind::Literal,
operands: vec![0],
result: Some(0),
}],
vec![LiteralValue::U32(42)],
);
let r = analyze(&desc);
assert_eq!(r.as_constant(0), Some(42));
assert_eq!(r.as_constant(99), None); }
#[test]
fn as_constant_returns_none_for_non_singleton() {
let mut report = ValueRangeReport::default();
report.ranges.insert(7, IntRange { min: 0, max: 10 });
assert_eq!(report.as_constant(7), None);
}
#[test]
fn report_accessors() {
let desc = build(
vec![
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![0],
result: Some(0),
},
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![1],
result: Some(1),
},
],
vec![LiteralValue::U32(0), LiteralValue::U32(42)],
);
let r = analyze(&desc);
assert_eq!(r.is_definitely(0, 0), Some(true));
assert_eq!(r.is_definitely(0, 1), Some(false));
assert_eq!(r.is_definitely(99, 0), None); assert_eq!(r.is_definitely_below(1, 100), Some(true));
assert_eq!(r.is_definitely_below(1, 42), Some(false)); assert_eq!(r.is_definitely_at_least(1, 42), Some(true));
assert_eq!(r.is_definitely_at_least(1, 43), Some(false));
assert_eq!(r.get(0), Some(IntRange::singleton(0)));
assert_eq!(r.get(99), None);
}
#[test]
fn range_helpers() {
let r = IntRange { min: 3, max: 7 };
assert!(r.contains(5));
assert!(r.contains(3));
assert!(r.contains(7));
assert!(!r.contains(2));
assert!(!r.contains(8));
assert!(!r.is_singleton());
let s = IntRange::singleton(42);
assert!(s.is_singleton());
let u = IntRange { min: 0, max: 5 }.union(IntRange { min: 3, max: 10 });
assert_eq!(u, IntRange { min: 0, max: 10 });
}
}