morok-schedule 0.1.0-alpha.2

Optimization passes and pattern engine for the Morok ML compiler
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

//! Test helpers for DCE tests.

use morok_ir::types::ConstValue;
use morok_ir::{Op, UOp};
use std::sync::Arc;

use crate::TypedPatternMatcher;
use crate::symbolic::symbolic_simple;

/// Get the symbolic_simple pattern matcher (reduces duplication).
pub fn get_matcher() -> &'static TypedPatternMatcher {
    symbolic_simple()
}

/// Assert that a UOp transforms to a specific constant value.
///
/// # Panics
/// Panics if the UOp is not a Const or doesn't match the expected value.
pub fn assert_const_value(uop: &Arc<UOp>, expected: ConstValue) {
    match uop.op() {
        Op::Const(cv) => {
            assert_eq!(cv.0, expected, "Expected Const({:?}), got Const({:?})", expected, cv.0);
        }
        other => panic!("Expected Const({:?}), got {:?}", expected, other),
    }
}

/// Assert that an END operation unwraps to its computation (no ranges).
///
/// Returns the computation UOp for further assertions.
///
/// # Panics
/// Panics if the UOp is still an END operation.
pub fn assert_end_unwrapped(uop: &Arc<UOp>) -> Arc<UOp> {
    match uop.op() {
        Op::End { .. } => {
            panic!("Expected END to be unwrapped, but got END operation: {:?}", uop.op())
        }
        _ => uop.clone(),
    }
}

/// Assert that an END operation has a specific number of ranges.
///
/// Returns the computation UOp and ranges for further assertions.
///
/// # Panics
/// Panics if the UOp is not an END operation or range count doesn't match.
pub fn assert_end_range_count(uop: &Arc<UOp>, expected_count: usize) -> (Arc<UOp>, Vec<Arc<UOp>>) {
    match uop.op() {
        Op::End { computation, ranges } => {
            assert_eq!(
                ranges.len(),
                expected_count,
                "Expected END with {} ranges, got {} ranges",
                expected_count,
                ranges.len()
            );
            (Arc::clone(computation), ranges.iter().cloned().collect())
        }
        other => panic!("Expected END operation, got {:?}", other),
    }
}