morok_schedule/
lib.rs

1//! Schedule module for Morok compiler.
2//!
3//! This module implements optimization passes for the IR,
4//! including symbolic simplification and graph transformations.
5//!
6//! # Module Organization
7//!
8//! - [`symbolic`] - Symbolic simplification patterns
9//! - [`rangeify`] - RANGEIFY transformation (movement ops → kernels)
10//!   - Phases 1-4: Movement ops to BUFFERIZE with symbolic simplification
11//!   - Phase 5: Kernel splitting at STORE boundaries
12//! - [`linearize`] - Priority-aware topological sort for GPU/NPU backends
13//! - [`optimizer`] - Kernel optimization layer (OptOps, Scheduler, heuristics)
14//! - [`expand`] - Pre-expansion pass for UNROLL/UPCAST range handling
15//!
16//! # Pattern Matching and Rewriting
17//!
18//! Pattern matching infrastructure has moved to `morok_ir::pattern` and `morok_ir::rewrite`.
19//! This crate re-exports these modules for convenience.
20
21pub mod devectorize;
22pub mod expand;
23pub mod gpudims;
24pub mod linearize;
25pub mod optimizer;
26pub mod passes;
27pub mod rangeify;
28pub mod symbolic;
29#[cfg(feature = "testing")]
30pub mod testing;
31
32#[cfg(feature = "z3")]
33pub mod z3;
34
35#[cfg(test)]
36pub mod test;
37
38// Re-export pattern matching and rewriting from morok_ir
39// This maintains backward compatibility while the infrastructure lives in morok_ir
40pub use morok_ir::pattern;
41pub use morok_ir::rewrite;
42
43// Re-export main types
44pub use linearize::{CFGContext, linearize, linearize_with_cfg};
45pub use morok_ir::pattern::{Matcher, RewriteResult, TypedPatternMatcher};
46pub use morok_ir::rewrite::graph_rewrite;
47pub use rangeify::{RangeifyResult, rangeify, rangeify_with_map, run_kernel_split_pipeline};
48
49// Re-export expand pass
50pub use expand::{expander, pm_group_for_reduce, pm_pre_expander, pre_expand};
51
52// Re-export devectorize pass
53pub use devectorize::devectorize;
54
55// Re-export gpudims pass
56pub use gpudims::pm_add_gpudims;
57
58// Re-export backend-agnostic passes (pm_linearize_multi_index removed: Tinygrad keeps multi-index INDEX)
59pub use passes::{build_linear_index, compute_row_major_strides, count_divmod, extract_index_dimension};
60
61// Re-export optimizer entry points
62pub use optimizer::{
63    BeamConfig, BeamResult, HeuristicsConfig, OptError, OptStrategy, OptimizerConfig, Renderer as OptimizerRenderer,
64    Scheduler, TcOptLevel, TcSelect, TcUsage, apply_post_optimization, apply_post_optimization_with_renderer,
65    beam_search_cached, hand_coded_optimizations, optimize_kernel, optimize_kernel_with_config,
66    optimize_kernel_with_strategy, prepare_scheduler,
67};
68
69// Re-export UOp for macro usage
70pub use morok_ir::UOp;
71
72// Re-export the patterns! proc-macros
73pub use morok_macros::{cached_patterns, patterns};
74
75/// Compute inverse permutation (argsort).
76pub(crate) fn argsort(perm: &[usize]) -> Vec<usize> {
77    let mut inv = vec![0; perm.len()];
78    for (i, &p) in perm.iter().enumerate() {
79        inv[p] = i;
80    }
81    inv
82}
morok_schedule/lib.rs

morok_schedule/
lib.rs
