av-denoise-core 0.4.0-alpha4

Core kernels and types for av-denoise (Do not use directly)
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

use cubecl::prelude::*;

/// The plane-local index of the first lane in the calling lane's
/// 8-lane group.
///
/// Every shuffle in this module addresses lanes relative to this, so a
/// group never reads a lane belonging to another group. Cubes are
/// launched 1-D, which makes the thread-to-plane mapping linear, so
/// `UNIT_POS_PLANE % 8` and `UNIT_POS_X % 8` agree at both wave32 and
/// wave64.
#[cube]
pub(crate) fn group_base() -> u32 {
    UNIT_POS_PLANE - UNIT_POS_PLANE % 8u32
}

/// The sum of `partial` across the calling lane's 8-lane group,
/// returned to every lane in it.
///
/// Three XOR shuffles fold 8 values into 8 copies of their sum. The
/// masks are 1, 2 and 4, all below 8, so a lane index XORed with any of
/// them stays inside its own group whatever the plane's width.
///
/// Each lane holds the squared differences for one patch column, so
/// this is what completes a candidate's distance.
#[cube]
pub(crate) fn plane_ssd_reduce8(partial: f32) -> f32 {
    let mut acc = partial;
    acc += plane_shuffle_xor(acc, 1u32);
    acc += plane_shuffle_xor(acc, 2u32);
    acc += plane_shuffle_xor(acc, 4u32);
    acc
}

/// Inserts one candidate into the group's sorted top-8.
///
/// The eight best candidates seen so far live one per lane, ascending,
/// slot 0 in the group's first lane. A new candidate shifts every slot
/// it beats one lane along and drops out the eighth. Nothing is stored
/// in shared memory and no second pass is needed.
///
/// A candidate that ties an incumbent does not displace it, so the
/// first candidate seen at a given distance keeps its slot. Exactly
/// equal distances are common on flat content, so this is what fixes
/// which member a group keeps rather than leaving it to scheduling.
///
/// `plane_shuffle_up` at the group's first lane returns a value from
/// the previous group. The `sub == 0` term discards it before it can be
/// used.
///
/// Exactly two shuffles run, one per carried value. Whether the previous
/// lane also beat `d` is not shuffled, because it follows from `prev_d`.
/// The slots ascend, so the previous lane holds `prev_d` and beat `d`
/// exactly when `d < prev_d`. Each shuffle is an LDS crossbar operation
/// on every candidate, so deriving this rather than shuffling a flag is
/// worth the line of algebra.
#[cube]
pub(crate) fn shift_insert8(best_d: &mut f32, best_pos: &mut u32, d: f32, packed: u32, sub: u32) {
    let prev_d = plane_shuffle_up(*best_d, 1u32);
    let prev_pos = plane_shuffle_up(*best_pos, 1u32);

    if d < *best_d {
        let first = sub == 0u32 || d >= prev_d;
        if first {
            *best_d = d;
            *best_pos = packed;
        } else {
            *best_d = prev_d;
            *best_pos = prev_pos;
        }
    }
}

/// [`shift_insert8`] with the shuffles skipped when the candidate
/// cannot place.
///
/// The group's eighth-best distance sits in its last lane. A candidate
/// that does not beat it changes nothing, so the two shuffles the insert
/// costs are skipped. Every lane holds the same `d` and reads the same
/// broadcast, so the branch is uniform across the group and no lane sits
/// out a shuffle another lane takes.
///
/// The broadcast costs nothing. It fuses into the compare that sets the
/// execution mask, one `v_cmpx_gt_f32` with a `dpp8` modifier, while
/// each shuffle it skips is an LDS crossbar operation. That is why this
/// is the form the kernel uses rather than a later optimisation.
///
/// This is a compute saving and never an admission decision. The eight
/// slots it produces are the same ones [`shift_insert8`] produces.
#[cube]
pub(crate) fn shift_insert8_gated(
    best_d: &mut f32,
    best_pos: &mut u32,
    d: f32,
    packed: u32,
    sub: u32,
    base: u32,
) {
    let worst = plane_shuffle(*best_d, base + 7u32);
    if d < worst {
        shift_insert8(best_d, best_pos, d, packed, sub);
    }
}

/// Transposes one 8x8 block held one column per lane into one row per
/// lane, through `buf`.
///
/// `v` holds the calling lane's 8 values on entry and its transposed 8
/// on return. `slot` is the calling lane's group index, which picks the
/// group's own 65-float region of `buf`, and the stride is padded by
/// one past 64 so that eight lanes writing eight consecutive rows never
/// collide on a bank.
///
/// The spatial row pass needs a row and a lane owns a column, so this
/// runs once before it and once after its inverse.
#[cube]
pub(crate) fn transpose8(buf: &mut SharedMemory<f32>, v: &mut Array<f32>, sub: u32, slot: u32) {
    let base = slot * 65u32;
    #[unroll]
    for i in 0..8u32 {
        buf[(base + i * 8u32 + sub) as usize] = v[i as usize];
    }
    sync_cube();
    #[unroll]
    for i in 0..8u32 {
        v[i as usize] = buf[(base + sub * 8u32 + i) as usize];
    }
    sync_cube();
}