Skip to main content

morok_ir/
lib.rs

1//! Intermediate Representation (IR) for the Morok compiler.
2//!
3//! This crate defines the core IR data structures and operations used throughout
4//! the Morok compiler pipeline.
5//!
6//! # Module Organization
7//!
8//! - [`types`] - Fundamental type definitions (ConstValue, operation types, etc.)
9//! - [`op`] - Operation enum defining all IR operations
10//! - [`uop`] - UOp (micro-operation) struct and implementation
11//! - [`uop::constructors`] - UOp constructor methods by semantic category
12//! - [`indexing`] - Multi-dimensional indexing support
13//! - [`error`] - Error types and result handling
14//! - [`shape`] - Shape inference utilities
15//! - [`sint`] - Symbolic integers
16
17// Make this crate available as `morok_ir` for proc-macro generated code
18extern crate self as morok_ir;
19
20// Module declarations
21pub mod decompositions;
22pub mod error;
23pub mod indexing;
24pub mod kernel_info;
25pub mod op;
26pub mod prelude;
27pub mod shape;
28pub mod sint;
29pub mod types;
30pub mod uop;
31
32pub mod provenance;
33
34#[macro_use]
35pub mod pattern;
36pub mod rewrite;
37
38#[cfg(any(test, feature = "proptest"))]
39pub mod test;
40
41// Re-exports for backward compatibility
42// All types remain accessible at the crate root
43pub use error::{Error, IndexTypeMismatchSnafu, Result};
44pub use indexing::IndexSpec;
45pub use op::Op;
46pub use sint::{IntoShrinkRange, SInt, ShrinkRange, sint_max, sint_min, sint_prod};
47pub use types::{
48    AddrSpace, AxisId, AxisType, BinaryOp, BufferizeOpts, ConstValue, ConstValueHash, ContiguousHint, ReduceOp,
49    TernaryOp, UnaryOp, WmmaMetadata, WmmaUpcastAxes,
50};
51pub use uop::{IntoUOp, UOp, UOpKey};
52
53// Re-export pattern matching and rewriting infrastructure
54pub use pattern::{Matcher, RewriteResult, TypedPatternMatcher};
55pub use rewrite::graph_rewrite;
56
57// Re-export external types for convenience
58pub use morok_dtype::DType;
59pub use morok_dtype::DeviceSpec;