knok-build 0.4.0

Build-script graph tracing and code generation for knok
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
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157

use knok_core::{CallOp, Expr};

use super::{
    context::{tuple_id, TraceVars},
    tensor::{ScalarLiteral, TraceTensor, T0},
};

/// Floating-point scalar types supported by random uniform tracing helpers.
pub trait RandomFloatElement: ScalarLiteral {
    /// Additive identity used by [`rand`].
    fn zero() -> Self;
    /// Exclusive upper bound used by [`rand`].
    fn one() -> Self;
}

impl RandomFloatElement for f32 {
    fn zero() -> Self {
        0.0
    }

    fn one() -> Self {
        1.0
    }
}

impl RandomFloatElement for f64 {
    fn zero() -> Self {
        0.0
    }

    fn one() -> Self {
        1.0
    }
}

/// Integer scalar types supported by random integer tracing helpers.
pub trait RandomIntElement: ScalarLiteral {}

impl RandomIntElement for i32 {}
impl RandomIntElement for i64 {}

/// Records a compile-time deterministic uniform random tensor.
///
/// Values are sampled in `[low, high)` from `seed` while the graph is compiled,
/// then embedded as a constant in the generated artifact.
pub fn uniform_static<Output>(seed: u64, low: Output::Elem, high: Output::Elem) -> Output
where
    Output: TraceTensor,
    Output::Elem: RandomFloatElement,
{
    trace_static_random_uniform(seed, low, high)
}

fn trace_static_random_uniform<Output>(seed: u64, low: Output::Elem, high: Output::Elem) -> Output
where
    Output: TraceTensor,
    Output::Elem: RandomFloatElement,
{
    Output::from_expr(Expr::Call {
        op: CallOp::StaticRandomUniform {
            target: Output::tensor_type(),
            seed,
        },
        args: vec![low.const_expr(), high.const_expr()],
    })
}

/// Records a compile-time deterministic integer random tensor.
///
/// Values are sampled in `[low, high)` from `seed` while the graph is compiled,
/// then embedded as a constant in the generated artifact.
pub fn randint_static<Output>(seed: u64, low: Output::Elem, high: Output::Elem) -> Output
where
    Output: TraceTensor,
    Output::Elem: RandomIntElement,
{
    trace_static_random_int(seed, low, high)
}

fn trace_static_random_int<Output>(seed: u64, low: Output::Elem, high: Output::Elem) -> Output
where
    Output: TraceTensor,
    Output::Elem: RandomIntElement,
{
    Output::from_expr(Expr::Call {
        op: CallOp::StaticRandomInt {
            target: Output::tensor_type(),
            seed,
        },
        args: vec![low.const_expr(), high.const_expr()],
    })
}

/// Records a runtime deterministic uniform random tensor and the next key.
///
/// Values are sampled in `[low, high)` from `key`. Passing the same key produces
/// the same values; pass the returned key to advance the random stream.
pub fn uniform<Output>(key: T0<i64>, low: Output::Elem, high: Output::Elem) -> (Output, T0<i64>)
where
    Output: TraceTensor,
    Output::Elem: RandomFloatElement,
{
    trace_random_uniform(key, low, high)
}

fn trace_random_uniform<Output>(
    key: T0<i64>,
    low: Output::Elem,
    high: Output::Elem,
) -> (Output, T0<i64>)
where
    Output: TraceTensor,
    Output::Elem: RandomFloatElement,
{
    let value = Expr::Call {
        op: CallOp::RandomUniform(Output::tensor_type()),
        args: vec![key.into_expr(), low.const_expr(), high.const_expr()],
    };
    <(Output, T0<i64>) as TraceVars>::from_tuple_expr(tuple_id(), value)
}

/// Records a runtime deterministic uniform random tensor in `[0, 1)`.
pub fn rand<Output>(key: T0<i64>) -> (Output, T0<i64>)
where
    Output: TraceTensor,
    Output::Elem: RandomFloatElement,
{
    uniform(key, Output::Elem::zero(), Output::Elem::one())
}

/// Records a runtime deterministic integer random tensor and the next key.
///
/// Values are sampled in `[low, high)` from `key`. Passing the same key produces
/// the same values; pass the returned key to advance the random stream.
pub fn randint<Output>(key: T0<i64>, low: Output::Elem, high: Output::Elem) -> (Output, T0<i64>)
where
    Output: TraceTensor,
    Output::Elem: RandomIntElement,
{
    trace_random_int(key, low, high)
}

fn trace_random_int<Output>(
    key: T0<i64>,
    low: Output::Elem,
    high: Output::Elem,
) -> (Output, T0<i64>)
where
    Output: TraceTensor,
    Output::Elem: RandomIntElement,
{
    let value = Expr::Call {
        op: CallOp::RandomInt(Output::tensor_type()),
        args: vec![key.into_expr(), low.const_expr(), high.const_expr()],
    };
    <(Output, T0<i64>) as TraceVars>::from_tuple_expr(tuple_id(), value)
}