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
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108

//! Vector operation tests.
//!
//! Tests vector operations: Vectorize, Gep, VConst, Cat, PtrCat.

use smallvec::smallvec;

use morok_dtype::{AddrSpace, DType};

use crate::{ConstValue, UOp};

// =========================================================================
// Vectorize Tests
// =========================================================================

#[test]
fn test_vectorize_basic() {
    // Should be Float32 vector of size 4
    assert_eq!(
        UOp::vectorize(smallvec![
            UOp::native_const(1.0f32),
            UOp::native_const(2.0f32),
            UOp::native_const(3.0f32),
            UOp::native_const(4.0f32)
        ])
        .dtype(),
        DType::Float32.vec(4)
    );
}

#[test]
fn test_vectorize_preserves_base_dtype() {
    let vec = UOp::vectorize(smallvec![UOp::native_const(1i32), UOp::native_const(2i32)]);
    assert_eq!(vec.dtype(), DType::Int32.vec(2));
}

// =========================================================================
// Gep (Get Element Pointer) Tests
// =========================================================================

#[test]
fn test_gep_basic() {
    // Create a vector
    let vec = UOp::vectorize(smallvec![
        UOp::native_const(1.0f32),
        UOp::native_const(2.0f32),
        UOp::native_const(3.0f32),
        UOp::native_const(4.0f32)
    ]);

    // GEP operation exists (actual behavior may vary based on implementation)
    let _elem = vec.gep(vec![0]);
    // Just verify it compiles and creates a UOp
}

#[test]
fn test_gep_multiple_indices() {
    let vec = UOp::vectorize(smallvec![
        UOp::native_const(10i32),
        UOp::native_const(20i32),
        UOp::native_const(30i32),
        UOp::native_const(40i32),
    ]);

    // Extract multiple elements -> produces vector of extracted elements
    let result = vec.gep(vec![0, 2]);
    assert_eq!(result.dtype(), DType::Int32.vec(2));
}

// =========================================================================
// VConst Tests
// =========================================================================

#[test]
fn test_vconst_basic() {
    let values = vec![ConstValue::Float(1.0), ConstValue::Float(2.0), ConstValue::Float(3.0), ConstValue::Float(4.0)];

    let vec = UOp::vconst(values, DType::Float64);
    assert_eq!(vec.dtype(), DType::Float64.vec(4));
}

// =========================================================================
// Cat Tests
// =========================================================================

#[test]
fn test_cat_basic() {
    let a = UOp::vectorize(smallvec![UOp::native_const(1.0f32), UOp::native_const(2.0f32),]);
    let b = UOp::vectorize(smallvec![UOp::native_const(3.0f32), UOp::native_const(4.0f32),]);

    let result = UOp::cat().sources(vec![a, b]).call();
    // Cat concatenates vectors: <2 x f32> + <2 x f32> = <4 x f32>
    assert_eq!(result.dtype(), DType::Float32.vec(4));
}

// =========================================================================
// PtrCat Tests
// =========================================================================

#[test]
fn test_ptrcat_basic() {
    let ptr_dtype = DType::Float32.ptr(None, AddrSpace::Global);
    let a = UOp::const_(ptr_dtype.clone(), ConstValue::Int(0));
    let b = UOp::const_(ptr_dtype.clone(), ConstValue::Int(0));

    let result = UOp::ptrcat().sources(vec![a, b]).call();
    // PTRCAT of 2 scalar pointers → vcount=2
    assert_eq!(result.dtype(), ptr_dtype.vec(2));
}