vyre-libs 0.6.2

vyre Category A library ecosystem - pure-IR compositions over vyre-ops hardware primitives
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
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
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
//! GPU expression-shape graph construction builders.

#![allow(missing_docs)] // Internal VAST-builder helpers are documented at the owning module boundary.
use crate::parsing::c::lex::tokens::*;
use crate::parsing::composition::child_phase;
use crate::region::wrap_anonymous;
use vyre::ir::{BufferAccess, BufferDecl, DataType, Expr, Node, Program};

use super::helpers::*;
use super::*;

pub fn c11_build_expression_shape_nodes(
    raw_vast_nodes: &str,
    typed_vast_nodes: &str,
    num_nodes: Expr,
    out_expr_shape_nodes: &str,
) -> Program {
    c11_build_expression_shape_nodes_impl(
        raw_vast_nodes,
        typed_vast_nodes,
        num_nodes,
        out_expr_shape_nodes,
        true,
    )
}

#[must_use]
pub fn c11_build_expression_shape_nodes_no_conditional(
    raw_vast_nodes: &str,
    typed_vast_nodes: &str,
    num_nodes: Expr,
    out_expr_shape_nodes: &str,
) -> Program {
    c11_build_expression_shape_nodes_impl(
        raw_vast_nodes,
        typed_vast_nodes,
        num_nodes,
        out_expr_shape_nodes,
        false,
    )
}

fn c11_build_expression_shape_nodes_impl(
    raw_vast_nodes: &str,
    typed_vast_nodes: &str,
    num_nodes: Expr,
    out_expr_shape_nodes: &str,
    include_conditional_shapes: bool,
) -> Program {
    let t = Expr::InvocationId { axis: 0 };
    let vast_base = Expr::mul(t.clone(), Expr::u32(VAST_NODE_STRIDE_U32));
    let out_base = Expr::mul(t.clone(), Expr::u32(C_EXPR_SHAPE_STRIDE_U32));

    let mut loop_body = vec![
        Node::let_bind("raw_kind", Expr::load(raw_vast_nodes, vast_base.clone())),
        Node::let_bind(
            "typed_kind",
            Expr::load(typed_vast_nodes, vast_base.clone()),
        ),
        Node::let_bind(
            "cur_parent",
            Expr::load(raw_vast_nodes, Expr::add(vast_base.clone(), Expr::u32(1))),
        ),
        Node::let_bind(
            "shape_kind",
            c_expr_shape_kind(Expr::var("raw_kind"), Expr::var("typed_kind")),
        ),
        Node::let_bind(
            "precedence",
            c_expr_operator_precedence(Expr::var("raw_kind"), Expr::var("typed_kind")),
        ),
        Node::let_bind(
            "associativity",
            c_expr_operator_associativity(Expr::var("typed_kind")),
        ),
        Node::let_bind(
            "shape_is_expr",
            Expr::ne(Expr::var("shape_kind"), Expr::u32(C_EXPR_SHAPE_NONE)),
        ),
        Node::let_bind(
            "shape_is_conditional",
            Expr::eq(Expr::var("shape_kind"), Expr::u32(C_EXPR_SHAPE_CONDITIONAL)),
        ),
    ];

    if include_conditional_shapes {
        loop_body.extend(emit_prior_ternary_boundary_flag(
            raw_vast_nodes,
            Expr::var("cur_parent"),
            t.clone(),
            Expr::var("shape_is_expr"),
            "bin",
        ));
        loop_body.extend(emit_expr_segment_bounds(
            raw_vast_nodes,
            Expr::var("cur_parent"),
            t.clone(),
            num_nodes.clone(),
            "bin_plain",
            false,
            Expr::var("shape_is_expr"),
        ));
        loop_body.extend(emit_expr_segment_bounds(
            raw_vast_nodes,
            Expr::var("cur_parent"),
            t.clone(),
            num_nodes.clone(),
            "bin_ternary",
            true,
            Expr::var("shape_is_expr"),
        ));
    } else {
        loop_body.extend(emit_expr_segment_bounds(
            raw_vast_nodes,
            Expr::var("cur_parent"),
            t.clone(),
            num_nodes.clone(),
            "bin_plain",
            false,
            Expr::var("shape_is_expr"),
        ));
        loop_body.extend([
            Node::let_bind("bin_use_ternary_boundaries", Expr::u32(0)),
            Node::let_bind("bin_ternary_seg_start", Expr::var("bin_plain_seg_start")),
            Node::let_bind("bin_ternary_seg_end", Expr::var("bin_plain_seg_end")),
        ]);
    }
    loop_body.extend(vec![
        Node::let_bind(
            "bin_seg_start",
            Expr::select(
                Expr::eq(Expr::var("bin_use_ternary_boundaries"), Expr::u32(1)),
                Expr::var("bin_ternary_seg_start"),
                Expr::var("bin_plain_seg_start"),
            ),
        ),
        Node::let_bind(
            "bin_seg_end",
            Expr::select(
                Expr::eq(Expr::var("bin_use_ternary_boundaries"), Expr::u32(1)),
                Expr::var("bin_ternary_seg_end"),
                Expr::var("bin_plain_seg_end"),
            ),
        ),
        Node::let_bind("bin_left_bound", Expr::var("bin_seg_start")),
        Node::let_bind("bin_right_bound", Expr::var("bin_seg_end")),
        Node::let_bind("bin_left_parent_op", Expr::u32(SENTINEL)),
        Node::let_bind("bin_right_parent_op", Expr::u32(SENTINEL)),
        Node::if_then(
            Expr::var("shape_is_expr"),
            vec![Node::loop_for(
                "bin_parent_scan",
                Expr::var("bin_seg_start"),
                Expr::var("bin_seg_end"),
                vec![
                    Node::let_bind(
                        "bin_parent_base",
                        Expr::mul(
                            Expr::var("bin_parent_scan"),
                            Expr::u32(VAST_NODE_STRIDE_U32),
                        ),
                    ),
                    Node::let_bind(
                        "bin_parent_raw",
                        Expr::load(raw_vast_nodes, Expr::var("bin_parent_base")),
                    ),
                    Node::let_bind(
                        "bin_parent_typed",
                        Expr::load(typed_vast_nodes, Expr::var("bin_parent_base")),
                    ),
                    Node::let_bind(
                        "bin_parent_parent",
                        Expr::load(
                            raw_vast_nodes,
                            Expr::add(Expr::var("bin_parent_base"), Expr::u32(1)),
                        ),
                    ),
                    Node::let_bind(
                        "bin_parent_shape",
                        c_expr_shape_kind(
                            Expr::var("bin_parent_raw"),
                            Expr::var("bin_parent_typed"),
                        ),
                    ),
                    Node::let_bind(
                        "bin_parent_prec",
                        c_expr_operator_precedence(
                            Expr::var("bin_parent_raw"),
                            Expr::var("bin_parent_typed"),
                        ),
                    ),
                    Node::let_bind(
                        "bin_parent_is_operator",
                        Expr::and(
                            Expr::ne(Expr::var("bin_parent_shape"), Expr::u32(C_EXPR_SHAPE_NONE)),
                            Expr::ne(Expr::var("bin_parent_scan"), t.clone()),
                        ),
                    ),
                    Node::let_bind(
                        "bin_parent_equal_assoc",
                        Expr::and(
                            Expr::eq(Expr::var("bin_parent_prec"), Expr::var("precedence")),
                            Expr::or(
                                Expr::and(
                                    Expr::eq(
                                        Expr::var("associativity"),
                                        Expr::u32(C_EXPR_ASSOC_LEFT),
                                    ),
                                    Expr::lt(t.clone(), Expr::var("bin_parent_scan")),
                                ),
                                Expr::and(
                                    Expr::eq(
                                        Expr::var("associativity"),
                                        Expr::u32(C_EXPR_ASSOC_RIGHT),
                                    ),
                                    Expr::lt(Expr::var("bin_parent_scan"), t.clone()),
                                ),
                            ),
                        ),
                    ),
                    Node::let_bind(
                        "bin_parent_is_ancestor",
                        Expr::and(
                            Expr::and(
                                Expr::eq(Expr::var("bin_parent_parent"), Expr::var("cur_parent")),
                                Expr::var("bin_parent_is_operator"),
                            ),
                            Expr::or(
                                Expr::lt(Expr::var("bin_parent_prec"), Expr::var("precedence")),
                                Expr::var("bin_parent_equal_assoc"),
                            ),
                        ),
                    ),
                    Node::if_then(
                        Expr::and(
                            Expr::var("bin_parent_is_ancestor"),
                            Expr::lt(Expr::var("bin_parent_scan"), t.clone()),
                        ),
                        vec![Node::assign(
                            "bin_left_parent_op",
                            Expr::var("bin_parent_scan"),
                        )],
                    ),
                    Node::if_then(
                        Expr::and(
                            Expr::var("bin_parent_is_ancestor"),
                            Expr::and(
                                Expr::lt(t.clone(), Expr::var("bin_parent_scan")),
                                Expr::or(
                                    Expr::eq(Expr::var("bin_right_parent_op"), Expr::u32(SENTINEL)),
                                    Expr::lt(
                                        Expr::var("bin_parent_scan"),
                                        Expr::var("bin_right_parent_op"),
                                    ),
                                ),
                            ),
                        ),
                        vec![Node::assign(
                            "bin_right_parent_op",
                            Expr::var("bin_parent_scan"),
                        )],
                    ),
                ],
            )],
        ),
        Node::if_then(
            Expr::ne(Expr::var("bin_left_parent_op"), Expr::u32(SENTINEL)),
            vec![Node::assign(
                "bin_left_bound",
                Expr::add(Expr::var("bin_left_parent_op"), Expr::u32(1)),
            )],
        ),
        Node::if_then(
            Expr::ne(Expr::var("bin_right_parent_op"), Expr::u32(SENTINEL)),
            vec![Node::assign(
                "bin_right_bound",
                Expr::var("bin_right_parent_op"),
            )],
        ),
    ]);
    loop_body.extend(emit_expr_root_scan(
        raw_vast_nodes,
        typed_vast_nodes,
        Expr::var("bin_left_bound"),
        t.clone(),
        Expr::var("cur_parent"),
        Expr::var("shape_is_expr"),
        "bin_lhs",
    ));
    loop_body.extend(emit_expr_root_scan(
        raw_vast_nodes,
        typed_vast_nodes,
        Expr::add(t.clone(), Expr::u32(1)),
        Expr::var("bin_right_bound"),
        Expr::var("cur_parent"),
        Expr::var("shape_is_expr"),
        "bin_rhs",
    ));

    if include_conditional_shapes {
        loop_body.extend(emit_expr_segment_bounds(
            raw_vast_nodes,
            Expr::var("cur_parent"),
            t.clone(),
            num_nodes.clone(),
            "cond",
            false,
            Expr::var("shape_is_conditional"),
        ));
        loop_body.extend(emit_expr_segment_bounds(
            raw_vast_nodes,
            Expr::var("cur_parent"),
            t.clone(),
            num_nodes.clone(),
            "cond_condition",
            true,
            Expr::var("shape_is_conditional"),
        ));
        loop_body.extend(vec![
            Node::let_bind("cond_colon", Expr::u32(SENTINEL)),
            Node::let_bind("cond_depth", Expr::u32(0)),
            Node::if_then(
                Expr::var("shape_is_conditional"),
                vec![Node::loop_for(
                    "cond_colon_scan",
                    Expr::add(t.clone(), Expr::u32(1)),
                    Expr::var("cond_seg_end"),
                    vec![Node::if_then(
                        Expr::eq(Expr::var("cond_colon"), Expr::u32(SENTINEL)),
                        vec![
                            Node::let_bind(
                                "cond_colon_base",
                                Expr::mul(
                                    Expr::var("cond_colon_scan"),
                                    Expr::u32(VAST_NODE_STRIDE_U32),
                                ),
                            ),
                            Node::let_bind(
                                "cond_colon_raw",
                                Expr::load(raw_vast_nodes, Expr::var("cond_colon_base")),
                            ),
                            Node::let_bind(
                                "cond_colon_parent",
                                Expr::load(
                                    raw_vast_nodes,
                                    Expr::add(Expr::var("cond_colon_base"), Expr::u32(1)),
                                ),
                            ),
                            Node::if_then(
                                Expr::eq(Expr::var("cond_colon_parent"), Expr::var("cur_parent")),
                                vec![
                                    Node::if_then(
                                        Expr::eq(
                                            Expr::var("cond_colon_raw"),
                                            Expr::u32(TOK_QUESTION),
                                        ),
                                        vec![Node::assign(
                                            "cond_depth",
                                            Expr::add(Expr::var("cond_depth"), Expr::u32(1)),
                                        )],
                                    ),
                                    Node::if_then(
                                        Expr::eq(Expr::var("cond_colon_raw"), Expr::u32(TOK_COLON)),
                                        vec![
                                            Node::if_then(
                                                Expr::eq(Expr::var("cond_depth"), Expr::u32(0)),
                                                vec![Node::assign(
                                                    "cond_colon",
                                                    Expr::var("cond_colon_scan"),
                                                )],
                                            ),
                                            Node::if_then(
                                                Expr::gt(Expr::var("cond_depth"), Expr::u32(0)),
                                                vec![Node::assign(
                                                    "cond_depth",
                                                    Expr::sub(
                                                        Expr::var("cond_depth"),
                                                        Expr::u32(1),
                                                    ),
                                                )],
                                            ),
                                        ],
                                    ),
                                ],
                            ),
                        ],
                    )],
                )],
            ),
            Node::let_bind(
                "cond_has_colon",
                Expr::ne(Expr::var("cond_colon"), Expr::u32(SENTINEL)),
            ),
            Node::let_bind(
                "cond_then_end",
                Expr::select(
                    Expr::var("cond_has_colon"),
                    Expr::var("cond_colon"),
                    Expr::add(t.clone(), Expr::u32(1)),
                ),
            ),
            Node::let_bind(
                "cond_else_start",
                Expr::select(
                    Expr::var("cond_has_colon"),
                    Expr::add(Expr::var("cond_colon"), Expr::u32(1)),
                    Expr::var("cond_seg_end"),
                ),
            ),
            Node::let_bind(
                "cond_condition_start",
                Expr::var("cond_condition_seg_start"),
            ),
            Node::let_bind("cond_parent_op", Expr::u32(SENTINEL)),
            Node::if_then(
                Expr::var("shape_is_conditional"),
                vec![Node::loop_for(
                    "cond_parent_scan",
                    Expr::var("cond_seg_start"),
                    t.clone(),
                    vec![
                        Node::let_bind(
                            "cond_parent_base",
                            Expr::mul(
                                Expr::var("cond_parent_scan"),
                                Expr::u32(VAST_NODE_STRIDE_U32),
                            ),
                        ),
                        Node::let_bind(
                            "cond_parent_raw",
                            Expr::load(raw_vast_nodes, Expr::var("cond_parent_base")),
                        ),
                        Node::let_bind(
                            "cond_parent_typed",
                            Expr::load(typed_vast_nodes, Expr::var("cond_parent_base")),
                        ),
                        Node::let_bind(
                            "cond_parent_parent",
                            Expr::load(
                                raw_vast_nodes,
                                Expr::add(Expr::var("cond_parent_base"), Expr::u32(1)),
                            ),
                        ),
                        Node::let_bind(
                            "cond_parent_shape",
                            c_expr_shape_kind(
                                Expr::var("cond_parent_raw"),
                                Expr::var("cond_parent_typed"),
                            ),
                        ),
                        Node::let_bind(
                            "cond_parent_prec",
                            c_expr_operator_precedence(
                                Expr::var("cond_parent_raw"),
                                Expr::var("cond_parent_typed"),
                            ),
                        ),
                        Node::if_then(
                            Expr::and(
                                Expr::eq(Expr::var("cond_parent_parent"), Expr::var("cur_parent")),
                                Expr::and(
                                    Expr::ne(
                                        Expr::var("cond_parent_shape"),
                                        Expr::u32(C_EXPR_SHAPE_NONE),
                                    ),
                                    Expr::lt(
                                        Expr::var("cond_parent_prec"),
                                        Expr::var("precedence"),
                                    ),
                                ),
                            ),
                            vec![Node::assign(
                                "cond_parent_op",
                                Expr::var("cond_parent_scan"),
                            )],
                        ),
                    ],
                )],
            ),
            Node::if_then(
                Expr::ne(Expr::var("cond_parent_op"), Expr::u32(SENTINEL)),
                vec![Node::assign(
                    "cond_condition_start",
                    Expr::add(Expr::var("cond_parent_op"), Expr::u32(1)),
                )],
            ),
        ]);
        loop_body.extend(emit_expr_root_scan(
            raw_vast_nodes,
            typed_vast_nodes,
            Expr::var("cond_condition_start"),
            t.clone(),
            Expr::var("cur_parent"),
            Expr::var("shape_is_conditional"),
            "cond_condition",
        ));
        loop_body.extend(emit_expr_root_scan(
            raw_vast_nodes,
            typed_vast_nodes,
            Expr::add(t.clone(), Expr::u32(1)),
            Expr::var("cond_then_end"),
            Expr::var("cur_parent"),
            Expr::var("shape_is_conditional"),
            "cond_then",
        ));
        loop_body.extend(emit_expr_root_scan(
            raw_vast_nodes,
            typed_vast_nodes,
            Expr::var("cond_else_start"),
            Expr::var("cond_seg_end"),
            Expr::var("cur_parent"),
            Expr::var("shape_is_conditional"),
            "cond_else",
        ));
    } else {
        loop_body.extend(vec![
            Node::let_bind("cond_condition_root", Expr::u32(SENTINEL)),
            Node::let_bind("cond_then_root", Expr::u32(SENTINEL)),
            Node::let_bind("cond_else_root", Expr::u32(SENTINEL)),
        ]);
    }

    loop_body.extend(vec![
        Node::let_bind(
            "field5",
            Expr::select(
                Expr::eq(Expr::var("shape_kind"), Expr::u32(C_EXPR_SHAPE_CONDITIONAL)),
                Expr::var("cond_condition_root"),
                Expr::var("bin_lhs_root"),
            ),
        ),
        Node::let_bind(
            "field6",
            Expr::select(
                Expr::eq(Expr::var("shape_kind"), Expr::u32(C_EXPR_SHAPE_CONDITIONAL)),
                Expr::var("cond_then_root"),
                Expr::var("bin_rhs_root"),
            ),
        ),
        Node::let_bind(
            "field7",
            Expr::select(
                Expr::eq(Expr::var("shape_kind"), Expr::u32(C_EXPR_SHAPE_CONDITIONAL)),
                Expr::var("cond_else_root"),
                Expr::u32(SENTINEL),
            ),
        ),
        Node::store(
            out_expr_shape_nodes,
            out_base.clone(),
            Expr::var("shape_kind"),
        ),
        Node::store(
            out_expr_shape_nodes,
            Expr::add(out_base.clone(), Expr::u32(1)),
            Expr::select(
                Expr::eq(Expr::var("shape_kind"), Expr::u32(C_EXPR_SHAPE_NONE)),
                Expr::u32(SENTINEL),
                t.clone(),
            ),
        ),
        Node::store(
            out_expr_shape_nodes,
            Expr::add(out_base.clone(), Expr::u32(2)),
            Expr::var("raw_kind"),
        ),
        Node::store(
            out_expr_shape_nodes,
            Expr::add(out_base.clone(), Expr::u32(3)),
            Expr::var("precedence"),
        ),
        Node::store(
            out_expr_shape_nodes,
            Expr::add(out_base.clone(), Expr::u32(4)),
            Expr::var("associativity"),
        ),
        Node::store(
            out_expr_shape_nodes,
            Expr::add(out_base.clone(), Expr::u32(5)),
            Expr::select(
                Expr::eq(Expr::var("shape_kind"), Expr::u32(C_EXPR_SHAPE_NONE)),
                Expr::u32(SENTINEL),
                Expr::var("field5"),
            ),
        ),
        Node::store(
            out_expr_shape_nodes,
            Expr::add(out_base.clone(), Expr::u32(6)),
            Expr::select(
                Expr::eq(Expr::var("shape_kind"), Expr::u32(C_EXPR_SHAPE_NONE)),
                Expr::u32(SENTINEL),
                Expr::var("field6"),
            ),
        ),
        Node::store(
            out_expr_shape_nodes,
            Expr::add(out_base, Expr::u32(7)),
            Expr::select(
                Expr::eq(Expr::var("shape_kind"), Expr::u32(C_EXPR_SHAPE_NONE)),
                Expr::u32(SENTINEL),
                Expr::var("field7"),
            ),
        ),
    ]);

    let n = node_count(&num_nodes).max(1);
    Program::wrapped(
        vec![
            BufferDecl::storage(raw_vast_nodes, 0, BufferAccess::ReadOnly, DataType::U32)
                .with_count(n.saturating_mul(VAST_NODE_STRIDE_U32)),
            BufferDecl::storage(typed_vast_nodes, 1, BufferAccess::ReadOnly, DataType::U32)
                .with_count(n.saturating_mul(VAST_NODE_STRIDE_U32)),
            BufferDecl::output(out_expr_shape_nodes, 2, DataType::U32)
                .with_count(n.saturating_mul(C_EXPR_SHAPE_STRIDE_U32)),
        ],
        [256, 1, 1],
        vec![wrap_anonymous(
            EXPR_SHAPE_OP_ID,
            vec![Node::if_then(
                Expr::lt(t.clone(), num_nodes),
                vec![child_phase(
                    EXPR_SHAPE_OP_ID,
                    "vyre-libs::parsing::c11_build_expression_shape_nodes::node_shape_pass",
                    loop_body,
                )],
            )],
        )],
    )
    .with_entry_op_id(EXPR_SHAPE_OP_ID)
}

pub(super) fn u32_words_to_bytes(words: &[u32]) -> Vec<u8> {
    vyre_primitives::wire::pack_u32_slice(words)
}