aprender-gpu 0.30.0

Pure Rust PTX generation for NVIDIA CUDA - no LLVM, no nvcc
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

//! PTX Static Analysis Tests (using probar gpu_pixels)

#![cfg(feature = "gpu-pixels")]

use super::*;

/// ptx-pixel-fkr: GEMM tiled kernel has no shared memory bugs
#[test]
fn ptx_pixel_fkr_gemm_tiled_no_bugs() {
    let kernel = GemmKernel::tiled(32, 32, 128, 32);
    let ptx = kernel.emit_ptx();
    let result = validate_ptx(&ptx);

    assert!(
        !result.has_bug(&PtxBugClass::SharedMemU64Addressing),
        "GEMM tiled kernel uses u64 for shared memory (should be u32)"
    );

    assert!(
        !result.has_bug(&PtxBugClass::MissingBarrierSync),
        "GEMM tiled kernel missing barrier synchronization"
    );

    println!("ptx_pixel_fkr_gemm_tiled: PASS (no shared memory bugs)");
}

/// ptx-pixel-fkr: Tensor core GEMM valid
#[test]
fn ptx_pixel_fkr_gemm_tensor_core() {
    let kernel = GemmKernel::tensor_core(32, 32, 64);
    let ptx = kernel.emit_ptx();
    let result = validate_ptx(&ptx);

    assert!(
        !result.has_bug(&PtxBugClass::SharedMemU64Addressing),
        "Tensor core GEMM uses u64 for shared memory"
    );

    println!("ptx_pixel_fkr_gemm_tensor_core: PASS");
}

/// ptx-pixel-fkr: Attention kernel validation
#[test]
fn ptx_pixel_fkr_attention() {
    let kernel = AttentionKernel::new(64, 64);
    let ptx = kernel.emit_ptx();
    let result = validate_ptx(&ptx);

    assert!(
        !result.has_bug(&PtxBugClass::SharedMemU64Addressing),
        "Attention kernel uses u64 for shared memory"
    );

    assert!(
        ptx.contains("bar.sync"),
        "Attention kernel must have barrier synchronization"
    );

    println!("ptx_pixel_fkr_attention: PASS");
}

/// ptx-pixel-fkr: Causal attention has correct name
#[test]
fn ptx_pixel_fkr_attention_causal() {
    let kernel = AttentionKernel::new(64, 64).with_causal();
    let ptx = kernel.emit_ptx();

    assert!(
        ptx.contains("flash_attention_causal") || ptx.contains("causal"),
        "Causal attention should have _causal suffix"
    );

    println!("ptx_pixel_fkr_attention_causal: PASS");
}

/// ptx-pixel-fkr: Softmax kernel entry point
#[test]
fn ptx_pixel_fkr_softmax_entry() {
    let kernel = SoftmaxKernel::new(128);
    let ptx = kernel.emit_ptx();
    let result = validate_ptx(&ptx);

    assert!(
        !result.has_bug(&PtxBugClass::MissingEntryPoint),
        "Softmax kernel must have entry point"
    );

    println!("ptx_pixel_fkr_softmax: PASS");
}

/// ptx-pixel-fkr: LayerNorm kernel entry point
#[test]
fn ptx_pixel_fkr_layernorm_entry() {
    let kernel = LayerNormKernel::new(256);
    let ptx = kernel.emit_ptx();
    let result = validate_ptx(&ptx);

    assert!(
        !result.has_bug(&PtxBugClass::MissingEntryPoint),
        "LayerNorm kernel must have entry point"
    );

    println!("ptx_pixel_fkr_layernorm: PASS");
}

/// ptx-pixel-fkr: BiasActivation kernel entry point (all variants)
#[test]
fn ptx_pixel_fkr_bias_activation_entry() {
    for activation in [Activation::None, Activation::ReLU, Activation::GELU] {
        let kernel = BiasActivationKernel::new(1024, 64).with_activation(activation);
        let ptx = kernel.emit_ptx();
        let result = validate_ptx(&ptx);

        assert!(
            !result.has_bug(&PtxBugClass::MissingEntryPoint),
            "BiasActivation kernel ({:?}) must have entry point",
            activation
        );
    }

    println!("ptx_pixel_fkr_bias_activation: PASS (all variants)");
}

/// ptx-pixel-fkr: BiasActivation GELU uses approximation functions
#[test]
fn ptx_pixel_fkr_bias_activation_gelu_approx() {
    let kernel = BiasActivationKernel::new(1024, 64).with_gelu();
    let ptx = kernel.emit_ptx();

    assert!(
        ptx.contains("ex2.approx") || ptx.contains("ex2.f32"),
        "GELU should use ex2 for fast exp approximation"
    );

    println!("ptx_pixel_fkr_bias_activation_gelu: PASS (uses ex2 approximation)");
}

/// ptx-pixel-fkr: BiasActivation ReLU uses max instruction
#[test]
fn ptx_pixel_fkr_bias_activation_relu_max() {
    let kernel = BiasActivationKernel::new(1024, 64).with_relu();
    let ptx = kernel.emit_ptx();

    assert!(
        ptx.contains("max.f32"),
        "ReLU should use max.f32 instruction"
    );

    println!("ptx_pixel_fkr_bias_activation_relu: PASS (uses max.f32)");
}