vyre-emit-ptx 0.6.1

PTX text emitter for vyre KernelDescriptor. Produces NVRTC-compatible CUDA assembly.
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

use vyre_lower::{KernelOp, KernelOpKind};

pub(super) fn is_latency_load(op: &KernelOp) -> bool {
    matches!(
        op.kind,
        KernelOpKind::LoadGlobal | KernelOpKind::LoadShared | KernelOpKind::LoadConstant
    ) && op.result.is_some()
}

pub(super) fn is_scheduling_fence(op: &KernelOp) -> bool {
    matches!(
        op.kind,
        KernelOpKind::StoreGlobal
            | KernelOpKind::StoreShared
            | KernelOpKind::Atomic { .. }
            | KernelOpKind::Barrier { .. }
            | KernelOpKind::Return
            | KernelOpKind::Region { .. }
            | KernelOpKind::StructuredBlock
            | KernelOpKind::StructuredIfThen
            | KernelOpKind::StructuredIfThenElse
            | KernelOpKind::StructuredForLoop { .. }
            | KernelOpKind::AsyncLoad { .. }
            | KernelOpKind::AsyncStore { .. }
            | KernelOpKind::AsyncWait { .. }
            | KernelOpKind::Trap { .. }
    )
}

pub(super) fn is_schedulable_pure_op(op: &KernelOp) -> bool {
    matches!(
        op.kind,
        KernelOpKind::Literal
            | KernelOpKind::LocalInvocationId
            | KernelOpKind::GlobalInvocationId
            | KernelOpKind::WorkgroupId
            | KernelOpKind::BinOpKind(_)
            | KernelOpKind::UnOpKind(_)
            | KernelOpKind::Fma
            | KernelOpKind::MatrixMma { .. }
            | KernelOpKind::Cast { .. }
            | KernelOpKind::Select
            | KernelOpKind::BufferLength
            | KernelOpKind::SubgroupLocalId
            | KernelOpKind::SubgroupSize
            | KernelOpKind::SubgroupBallot
            | KernelOpKind::SubgroupShuffle
            | KernelOpKind::SubgroupAdd
    ) && op.result.is_some()
}

pub(super) fn operand_is_immediate(op: &KernelOp, _operand: u32) -> bool {
    matches!(
        op.kind,
        KernelOpKind::Literal
            | KernelOpKind::LocalInvocationId
            | KernelOpKind::GlobalInvocationId
            | KernelOpKind::WorkgroupId
            | KernelOpKind::BufferLength
    )
}

pub(super) fn op_reads_operand(op: &KernelOp, operand: u32) -> bool {
    op.operands
        .iter()
        .any(|candidate| *candidate == operand && !operand_is_immediate(op, *candidate))
}

#[cfg(test)]
mod tests {
    use super::*;
    use vyre_lower::{MatrixMmaElement, MatrixMmaLayout, MatrixMmaShape};

    fn op(kind: KernelOpKind, operands: Vec<u32>, result: Option<u32>) -> KernelOp {
        KernelOp {
            kind,
            operands,
            result,
        }
    }

    #[test]
    fn fma_and_mma_are_schedulable_compute_fillers() {
        assert!(is_schedulable_pure_op(&op(
            KernelOpKind::Fma,
            vec![1, 2, 3],
            Some(4)
        )));
        assert!(is_schedulable_pure_op(&op(
            KernelOpKind::MatrixMma {
                shape: MatrixMmaShape::M16N8K16,
                a_layout: MatrixMmaLayout::RowMajor,
                b_layout: MatrixMmaLayout::ColMajor,
                a_type: MatrixMmaElement::F16,
                b_type: MatrixMmaElement::F16,
                accum_type: MatrixMmaElement::F32,
            },
            vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
            Some(11)
        )));
    }

    #[test]
    fn buffer_length_slot_operand_is_immediate_for_scheduling() {
        let length = op(KernelOpKind::BufferLength, vec![7], Some(12));
        assert!(is_schedulable_pure_op(&length));
        assert!(operand_is_immediate(&length, 7));
        assert!(!op_reads_operand(&length, 7));
    }

    #[test]
    fn unsupported_and_resultless_ops_are_not_latency_fillers() {
        assert!(!is_schedulable_pure_op(&op(
            KernelOpKind::Fma,
            vec![1, 2, 3],
            None
        )));
        assert!(!is_schedulable_pure_op(&op(
            KernelOpKind::Copy,
            vec![1],
            Some(2)
        )));
        assert!(!is_schedulable_pure_op(&op(
            KernelOpKind::StoreGlobal,
            vec![0, 1, 2],
            None
        )));
    }
}