Crate morok_schedule 

Schedule module for Morok compiler.

This module implements optimization passes for the IR, including symbolic simplification and graph transformations.

Module Organization

• symbolic - Symbolic simplification patterns
• [rangeify] - RANGEIFY transformation (movement ops → kernels)
  • Phases 1-4: Movement ops to BUFFERIZE with symbolic simplification
  • Phase 5: Kernel splitting at STORE boundaries
• [linearize] - Priority-aware topological sort for GPU/NPU backends
• optimizer - Kernel optimization layer (OptOps, Scheduler, heuristics)
• expand - Pre-expansion pass for UNROLL/UPCAST range handling

Pattern Matching and Rewriting

Pattern matching infrastructure has moved to morok_ir::pattern and morok_ir::rewrite. This crate re-exports these modules for convenience.

Re-exports

pub use linearize::CFGContext;

pub use linearize::linearize;

pub use linearize::linearize_with_cfg;

pub use rangeify::RangeifyResult;

pub use rangeify::rangeify;

pub use rangeify::rangeify_with_map;

pub use rangeify::run_kernel_split_pipeline;

pub use expand::expander;

pub use expand::pm_group_for_reduce;

pub use expand::pm_pre_expander;

pub use expand::pre_expand;

pub use devectorize::devectorize;

pub use gpudims::pm_add_gpudims;

pub use passes::build_linear_index;

pub use passes::compute_row_major_strides;

pub use passes::count_divmod;

pub use passes::extract_index_dimension;

pub use optimizer::BeamConfig;

pub use optimizer::BeamResult;

pub use optimizer::HeuristicsConfig;

pub use optimizer::OptError;

pub use optimizer::OptStrategy;

pub use optimizer::OptimizerConfig;

pub use optimizer::Renderer as OptimizerRenderer;

pub use optimizer::Scheduler;

pub use optimizer::TcOptLevel;

pub use optimizer::TcSelect;

pub use optimizer::TcUsage;

pub use optimizer::apply_post_optimization;

pub use optimizer::apply_post_optimization_with_renderer;

pub use optimizer::beam_search_cached;

pub use optimizer::hand_coded_optimizations;

pub use optimizer::optimize_kernel;

pub use optimizer::optimize_kernel_with_config;

pub use optimizer::optimize_kernel_with_strategy;

pub use optimizer::prepare_scheduler;

Modules

devectorize

Devectorize pass — single-pass combined matcher matching Tinygrad’s pm_devectorize.

expand

Pre-expander pass for UNROLL/CONTRACT expansion.

gpudims

GPU dimension injection for kernel execution.

linearize

Linearization module for converting UOp DAGs to linear instruction sequences.

optimizer

Kernel optimization layer for morok-schedule.

passes

Backend-agnostic IR transformation passes.

pattern

Pattern matching infrastructure for UOp graphs.

rangeify

RANGEIFY transformation: convert movement ops to BUFFERIZE+INDEX, then split into kernels.

rewrite

Graph rewrite engine with fixed-point iteration.

symbolic

Symbolic simplification patterns.

Macros

cached_patterns

Like patterns! but wraps the matcher in LazyLock for zero-cost reuse.

patterns

Proc-macro for declarative pattern rewrite rules.

Structs

UOp

Micro-operation node in the computation graph.

Enums

RewriteResult

Result of applying a pattern rewrite.

Traits

Matcher

Trait for pattern matchers used by the rewrite engine.

Functions

graph_rewrite

Apply graph rewriting. Patterns see OPTIMIZED children (Stage 1).

Type Aliases

TypedPatternMatcher

Type alias for backwards compatibility.