xlog-prob 0.9.2

Probabilistic inference engines for XLOG
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

use std::sync::Arc;

use xlog_core::MemoryBudget;
use xlog_cuda::memory::TrackedCudaSlice;
use xlog_cuda::{CudaDevice, CudaKernelProvider, GpuMemoryManager};
use xlog_prob::compilation::gpu_cnf::GpuCnfVarTables;
use xlog_prob::compilation::gpu_weights::{build_evidence_by_var_gpu, build_weights_gpu};

fn try_provider() -> Option<Arc<CudaKernelProvider>> {
    let device = match CudaDevice::new(0) {
        Ok(d) => Arc::new(d),
        Err(e) => {
            eprintln!("Skipping test: CUDA runtime unavailable: {}", e);
            return None;
        }
    };
    let memory = Arc::new(GpuMemoryManager::new(
        device.clone(),
        MemoryBudget::with_limit(1 << 30),
    ));
    match CudaKernelProvider::new(device, memory) {
        Ok(p) => Some(Arc::new(p)),
        Err(e) => {
            eprintln!(
                "Skipping test: failed to create CUDA kernel provider: {}",
                e
            );
            None
        }
    }
}

fn upload_u32(provider: &Arc<CudaKernelProvider>, host: &[u32]) -> TrackedCudaSlice<u32> {
    let memory = provider.memory();
    let mut buf = memory.alloc::<u32>(host.len()).expect("alloc u32");
    provider
        .device()
        .inner()
        .htod_sync_copy_into(host, &mut buf)
        .expect("upload u32");
    buf
}

fn upload_f64(provider: &Arc<CudaKernelProvider>, host: &[f64]) -> TrackedCudaSlice<f64> {
    let memory = provider.memory();
    let mut buf = memory.alloc::<f64>(host.len()).expect("alloc f64");
    provider
        .device()
        .inner()
        .htod_sync_copy_into(host, &mut buf)
        .expect("upload f64");
    buf
}

fn upload_u8(provider: &Arc<CudaKernelProvider>, host: &[u8]) -> TrackedCudaSlice<u8> {
    let memory = provider.memory();
    let mut buf = memory.alloc::<u8>(host.len()).expect("alloc u8");
    provider
        .device()
        .inner()
        .htod_sync_copy_into(host, &mut buf)
        .expect("upload u8");
    buf
}

fn ln_prob(p: f64) -> f64 {
    if p == 0.0 {
        f64::NEG_INFINITY
    } else {
        p.ln()
    }
}

#[test]
fn gpu_weights_builds_log_tables_and_evidence() {
    let Some(provider) = try_provider() else {
        return;
    };

    // var_cap = 6 (DIMACS 1-based)
    let node_var = upload_u32(&provider, &[4, 0, 0]); // node 0 -> var 4
    let leaf_var = upload_u32(&provider, &[2, 5, 0]); // leaf 0 -> var 2, leaf 1 -> var 5
    let choice_var = upload_u32(&provider, &[3, 6]); // choice 0 -> var 3, choice 1 -> var 6

    let vars = GpuCnfVarTables {
        node_var,
        leaf_var,
        choice_var,
        max_var: 6,
    };

    let leaf_probs = upload_f64(&provider, &[0.2, 0.7]);
    let choice_true = upload_f64(&provider, &[0.1, 0.6]);
    let choice_false = upload_f64(&provider, &[0.9, 0.4]);

    let evidence_nodes = upload_u32(&provider, &[0]);
    let evidence_vals = upload_u8(&provider, &[1]); // node 0 true -> var 4 true

    let evidence_by_var = build_evidence_by_var_gpu(
        &vars.node_var,
        &evidence_nodes,
        &evidence_vals,
        vars.max_var,
        &provider,
    )
    .expect("evidence map");

    let weights = build_weights_gpu(
        &vars,
        &leaf_probs,
        &choice_true,
        &choice_false,
        &evidence_by_var,
        &provider,
    )
    .expect("weights");

    let mut host_true = vec![0.0f64; (vars.max_var as usize) + 1];
    let mut host_false = vec![0.0f64; (vars.max_var as usize) + 1];

    provider
        .device()
        .inner()
        .dtoh_sync_copy_into(&weights.log_true, &mut host_true)
        .expect("read log_true");
    provider
        .device()
        .inner()
        .dtoh_sync_copy_into(&weights.log_false, &mut host_false)
        .expect("read log_false");

    // leaf var 2
    assert!((host_true[2] - ln_prob(0.2)).abs() < 1e-9);
    assert!((host_false[2] - ln_prob(0.8)).abs() < 1e-9);
    // leaf var 5
    assert!((host_true[5] - ln_prob(0.7)).abs() < 1e-9);
    assert!((host_false[5] - ln_prob(0.3)).abs() < 1e-9);

    // choice var 3
    assert!((host_true[3] - ln_prob(0.1)).abs() < 1e-9);
    assert!((host_false[3] - ln_prob(0.9)).abs() < 1e-9);
    // choice var 6
    assert!((host_true[6] - ln_prob(0.6)).abs() < 1e-9);
    assert!((host_false[6] - ln_prob(0.4)).abs() < 1e-9);

    // evidence on var 4 (true) forces log_false to -inf
    assert!(host_false[4].is_infinite() && host_false[4].is_sign_negative());
}