morok-ir 0.1.0-alpha.2

Intermediate representation 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

//! Intermediate Representation (IR) for the Morok compiler.
//!
//! This crate defines the core IR data structures and operations used throughout
//! the Morok compiler pipeline.
//!
//! # Module Organization
//!
//! - [`types`] - Fundamental type definitions (ConstValue, operation types, etc.)
//! - [`op`] - Operation enum defining all IR operations
//! - [`uop`] - UOp (micro-operation) struct and implementation
//! - [`uop::constructors`] - UOp constructor methods by semantic category
//! - [`indexing`] - Multi-dimensional indexing support
//! - [`error`] - Error types and result handling
//! - [`shape`] - Shape inference utilities
//! - [`sint`] - Symbolic integers

// Make this crate available as `morok_ir` for proc-macro generated code
extern crate self as morok_ir;

// Module declarations
pub mod decompositions;
pub mod error;
pub mod indexing;
pub mod kernel_info;
pub mod op;
pub mod prelude;
pub mod shape;
pub mod sint;
pub mod types;
pub mod uop;

pub mod provenance;

#[macro_use]
pub mod pattern;
pub mod rewrite;

#[cfg(any(test, feature = "proptest"))]
pub mod test;

// Re-exports for backward compatibility
// All types remain accessible at the crate root
pub use error::{Error, IndexTypeMismatchSnafu, Result};
pub use indexing::IndexSpec;
pub use op::Op;
pub use sint::{IntoShrinkRange, SInt, ShrinkRange, sint_max, sint_min, sint_prod};
pub use types::{
    AddrSpace, AxisId, AxisType, BinaryOp, BufferizeOpts, ConstValue, ConstValueHash, ContiguousHint, ReduceOp,
    TernaryOp, UnaryOp, WmmaMetadata, WmmaUpcastAxes,
};
pub use uop::{IntoUOp, UOp, UOpKey};

// Re-export pattern matching and rewriting infrastructure
pub use pattern::{Matcher, RewriteResult, TypedPatternMatcher};
pub use rewrite::graph_rewrite;

// Re-export external types for convenience
pub use morok_dtype::DType;
pub use morok_dtype::DeviceSpec;