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
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153

#![allow(
    clippy::doc_lazy_continuation,
    clippy::double_must_use,
    clippy::manual_div_ceil,
    clippy::needless_range_loop,
    clippy::collapsible_if,
    clippy::match_like_matches_macro,
    clippy::redundant_closure,
    clippy::too_many_arguments,
    clippy::nonminimal_bool,
    clippy::derivable_impls
)]
//! PTX text emitter for vyre `KernelDescriptor`.
//!
//! Consumes a substrate-neutral `vyre_lower::KernelDescriptor` and
//! produces NVRTC-compatible PTX assembly text. The emitter owns only
//! PTX construction; descriptor shaping and substrate-neutral
//! analyses stay in `vyre-lower`.
//!
//! ## Op coverage
//!
//! Mirrors `vyre-emit-naga` for parity:
//! - `Literal` (U32, I32, F32, Bool)
//! - `LocalInvocationId` / `GlobalInvocationId` / `WorkgroupId` (axis 0/1/2)
//! - `LoadGlobal` / `StoreGlobal` (scalar U32/I32/F32/Bool, plus packed
//!   `v2`/`v4` U32/I32/F32 chains when the descriptor presents unit-stride
//!   adjacent accesses)
//! - `BinOpKind` for the common arithmetic/logic set
//! - `UnOpKind` for Negate / LogicalNot / BitNot
//! - `Cast` between scalar types
//! - `Select`, `Fma`
//! - `StructuredIfThen`, `StructuredIfThenElse`, `StructuredBlock`,
//!   `Region`, `Return`, workgroup-scope `Barrier`
//!
//! Out of scope (returns `EmitError::UnsupportedOp` or
//! `EmitError::InvalidDescriptor`): indirect-dispatch (host concern),
//! `MemoryOrdering::GridSync` until a native cooperative-grid lowering is
//! wired, and descriptor forms without a PTX-safe lowering.
//!
//! ## PTX output shape
//!
//! ```text
//! //
//! // Generated by vyre-emit-ptx (target sm_70)
//! //
//! .version 7.0
//! .target sm_70
//! .address_size 64
//!
//! .visible .entry main(
//!     .param .u64 _arg_<binding_name>
//! )
//! {
//!     .reg .pred  %p<N>;
//!     .reg .u32   %r<N>;
//!     .reg .s32   %s<N>;
//!     .reg .f32   %f<N>;
//!     .reg .u64   %rd<N>;
//!
//!     <body>
//!
//!     ret;
//! }
//! ```

mod emitter;
mod error;
mod index_facts;
pub mod patterns;
mod reg;
mod target;

use vyre_lower::KernelDescriptor;

pub use error::EmitError;
pub use target::{ComputeCapability, PtxEmitOptions};

pub fn emit(desc: &KernelDescriptor) -> Result<String, EmitError> {
    emit_with_target(desc, ComputeCapability::default())
}

pub fn emit_with_target(
    desc: &KernelDescriptor,
    target: ComputeCapability,
) -> Result<String, EmitError> {
    emit_with_options(desc, PtxEmitOptions::for_target(target))
}

pub fn emit_with_options(
    desc: &KernelDescriptor,
    options: PtxEmitOptions,
) -> Result<String, EmitError> {
    if options.subgroup_size == 0
        || options.subgroup_size > 32
        || !options.subgroup_size.is_power_of_two()
    {
        return Err(EmitError::InvalidDescriptor(format!(
            "invalid CUDA subgroup size {}. Fix: pass the probed CUDA warp size.",
            options.subgroup_size
        )));
    }
    emitter::emit_text(desc, options)
}

/// Emit PTX text from a `KernelDescriptor` after running the full
/// `vyre_lower::rewrites::run_all` optimization pipeline. Recommended
/// over [`emit`] for production use  -  fewer dead instructions, fewer
/// redundant loads, lower register pressure.
pub fn emit_optimized(desc: &KernelDescriptor) -> Result<String, EmitError> {
    emit_optimized_with_stats(desc).map(|(s, _)| s)
}

/// Like [`emit_optimized`] but also returns
/// [`vyre_lower::rewrites::OptimizationStats`].
pub fn emit_optimized_with_stats(
    desc: &KernelDescriptor,
) -> Result<(String, vyre_lower::rewrites::OptimizationStats), EmitError> {
    let (optimized, stats) = vyre_lower::rewrites::run_all_with_stats(desc);
    debug_assert!(
        vyre_lower::verify::verify(&optimized).is_ok(),
        "rewrite pipeline produced an invalid descriptor  -  see vyre_lower::verify for the contract"
    );
    let ptx = emit(&optimized)?;
    Ok((ptx, stats))
}

/// Same as [`emit_with_target`] but runs the optimization pipeline
/// first.
pub fn emit_optimized_with_target(
    desc: &KernelDescriptor,
    target: ComputeCapability,
) -> Result<String, EmitError> {
    emit_optimized_with_target_with_stats(desc, target).map(|(s, _)| s)
}

/// The full-power variant: optimize first AND target a specific
/// compute capability AND surface OptimizationStats. Combines
/// [`emit_optimized_with_target`] and [`emit_optimized_with_stats`].
pub fn emit_optimized_with_target_with_stats(
    desc: &KernelDescriptor,
    target: ComputeCapability,
) -> Result<(String, vyre_lower::rewrites::OptimizationStats), EmitError> {
    let (optimized, stats) = vyre_lower::rewrites::run_all_with_stats(desc);
    debug_assert!(
        vyre_lower::verify::verify(&optimized).is_ok(),
        "rewrite pipeline produced an invalid descriptor  -  see vyre_lower::verify for the contract"
    );
    let ptx = emit_with_target(&optimized, target)?;
    Ok((ptx, stats))
}

#[cfg(test)]
mod tests;