use std::ffi::{CStr, c_char, c_int, c_uint, c_void};
use thiserror::Error as ThisError;
use crate::color::Color;
pub(crate) type CUresult = c_int;
pub(crate) type CUdevice = c_int;
pub(crate) type CUcontext = *mut c_void;
pub(crate) type CUdeviceptr = u64;
type CUmodule = *mut c_void;
type CUfunction = *mut c_void;
const CUDA_SUCCESS: CUresult = 0;
const CU_MEMORYTYPE_DEVICE: c_uint = 2;
#[repr(C)]
#[derive(Clone, Copy, Default)]
struct CudaMemcpy2D {
src_x_in_bytes: usize,
src_y: usize,
src_memory_type: c_uint,
src_host: *const c_void,
src_device: CUdeviceptr,
src_array: *mut c_void,
src_pitch: usize,
dst_x_in_bytes: usize,
dst_y: usize,
dst_memory_type: c_uint,
dst_host: *mut c_void,
dst_device: CUdeviceptr,
dst_array: *mut c_void,
dst_pitch: usize,
width_in_bytes: usize,
height: usize,
}
const CONVERT_PTX: &str = r#"
.version 6.0
.target sm_50
.address_size 64
.visible .entry bgra_to_luma(
.param .u64 dst,
.param .u32 dst_pitch,
.param .u64 src,
.param .u32 src_pitch,
.param .u32 width,
.param .u32 height
)
{
.reg .pred %p<4>;
.reg .b16 %rs<8>;
.reg .b32 %r<24>;
.reg .f32 %f<20>;
.reg .b64 %rd<12>;
ld.param.u64 %rd1, [dst];
ld.param.u32 %r1, [dst_pitch];
ld.param.u64 %rd2, [src];
ld.param.u32 %r2, [src_pitch];
ld.param.u32 %r3, [width];
ld.param.u32 %r4, [height];
mov.u32 %r5, %ctaid.x;
mov.u32 %r6, %ntid.x;
mov.u32 %r7, %tid.x;
mad.lo.s32 %r8, %r5, %r6, %r7;
mov.u32 %r9, %ctaid.y;
mov.u32 %r10, %ntid.y;
mov.u32 %r11, %tid.y;
mad.lo.s32 %r12, %r9, %r10, %r11;
setp.ge.u32 %p1, %r8, %r3;
@%p1 bra LUMA_DONE;
setp.ge.u32 %p2, %r12, %r4;
@%p2 bra LUMA_DONE;
mul.lo.s32 %r13, %r12, %r2;
shl.b32 %r14, %r8, 2;
add.s32 %r15, %r13, %r14;
cvt.u64.u32 %rd3, %r15;
add.s64 %rd4, %rd2, %rd3;
ld.global.u8 %rs1, [%rd4];
ld.global.u8 %rs2, [%rd4+1];
ld.global.u8 %rs3, [%rd4+2];
cvt.u32.u16 %r16, %rs1;
cvt.rn.f32.u32 %f1, %r16;
cvt.u32.u16 %r17, %rs2;
cvt.rn.f32.u32 %f2, %r17;
cvt.u32.u16 %r18, %rs3;
cvt.rn.f32.u32 %f3, %r18;
mul.f32 %f4, %f3, 0f3E59B3D0;
mul.f32 %f5, %f2, 0f3F371759;
add.f32 %f6, %f4, %f5;
mul.f32 %f7, %f1, 0f3D93DD98;
add.f32 %f8, %f6, %f7;
mul.f32 %f9, %f8, 0f435B0000;
div.rn.f32 %f10, %f9, 0f437F0000;
add.f32 %f11, %f10, 0f41800000;
max.f32 %f12, %f11, 0f00000000;
min.f32 %f13, %f12, 0f437F0000;
add.f32 %f14, %f13, 0f3F000000;
cvt.rzi.u32.f32 %r19, %f14;
cvt.u16.u32 %rs4, %r19;
mul.lo.s32 %r20, %r12, %r1;
add.s32 %r21, %r20, %r8;
cvt.u64.u32 %rd5, %r21;
add.s64 %rd6, %rd1, %rd5;
st.global.u8 [%rd6], %rs4;
LUMA_DONE:
ret;
}
.visible .entry bgra_to_chroma(
.param .u64 dst,
.param .u32 dst_pitch,
.param .u64 src,
.param .u32 src_pitch,
.param .u32 half_width,
.param .u32 half_height
)
{
.reg .pred %p<4>;
.reg .b16 %rs<20>;
.reg .b32 %r<48>;
.reg .f32 %f<32>;
.reg .b64 %rd<16>;
ld.param.u64 %rd1, [dst];
ld.param.u32 %r1, [dst_pitch];
ld.param.u64 %rd2, [src];
ld.param.u32 %r2, [src_pitch];
ld.param.u32 %r3, [half_width];
ld.param.u32 %r4, [half_height];
mov.u32 %r5, %ctaid.x;
mov.u32 %r6, %ntid.x;
mov.u32 %r7, %tid.x;
mad.lo.s32 %r8, %r5, %r6, %r7;
mov.u32 %r9, %ctaid.y;
mov.u32 %r10, %ntid.y;
mov.u32 %r11, %tid.y;
mad.lo.s32 %r12, %r9, %r10, %r11;
setp.ge.u32 %p1, %r8, %r3;
@%p1 bra CHROMA_DONE;
setp.ge.u32 %p2, %r12, %r4;
@%p2 bra CHROMA_DONE;
shl.b32 %r13, %r12, 1;
shl.b32 %r14, %r8, 1;
mul.lo.s32 %r15, %r13, %r2;
shl.b32 %r16, %r14, 2;
add.s32 %r17, %r15, %r16;
cvt.u64.u32 %rd3, %r17;
add.s64 %rd4, %rd2, %rd3;
cvt.u64.u32 %rd5, %r2;
add.s64 %rd6, %rd4, %rd5;
ld.global.u8 %rs1, [%rd4];
ld.global.u8 %rs2, [%rd4+1];
ld.global.u8 %rs3, [%rd4+2];
ld.global.u8 %rs4, [%rd4+4];
ld.global.u8 %rs5, [%rd4+5];
ld.global.u8 %rs6, [%rd4+6];
ld.global.u8 %rs7, [%rd6];
ld.global.u8 %rs8, [%rd6+1];
ld.global.u8 %rs9, [%rd6+2];
ld.global.u8 %rs10, [%rd6+4];
ld.global.u8 %rs11, [%rd6+5];
ld.global.u8 %rs12, [%rd6+6];
cvt.u32.u16 %r18, %rs1;
cvt.u32.u16 %r19, %rs4;
add.s32 %r20, %r18, %r19;
cvt.u32.u16 %r21, %rs7;
add.s32 %r22, %r20, %r21;
cvt.u32.u16 %r23, %rs10;
add.s32 %r24, %r22, %r23;
cvt.rn.f32.u32 %f1, %r24;
mul.f32 %f2, %f1, 0f3E800000;
cvt.u32.u16 %r25, %rs2;
cvt.u32.u16 %r26, %rs5;
add.s32 %r27, %r25, %r26;
cvt.u32.u16 %r28, %rs8;
add.s32 %r29, %r27, %r28;
cvt.u32.u16 %r30, %rs11;
add.s32 %r31, %r29, %r30;
cvt.rn.f32.u32 %f3, %r31;
mul.f32 %f4, %f3, 0f3E800000;
cvt.u32.u16 %r32, %rs3;
cvt.u32.u16 %r33, %rs6;
add.s32 %r34, %r32, %r33;
cvt.u32.u16 %r35, %rs9;
add.s32 %r36, %r34, %r35;
cvt.u32.u16 %r37, %rs12;
add.s32 %r38, %r36, %r37;
cvt.rn.f32.u32 %f5, %r38;
mul.f32 %f6, %f5, 0f3E800000;
mul.f32 %f7, %f6, 0f3E59B3D0;
mul.f32 %f8, %f4, 0f3F371759;
add.f32 %f9, %f7, %f8;
mul.f32 %f10, %f2, 0f3D93DD98;
add.f32 %f11, %f9, %f10;
sub.f32 %f12, %f2, %f11;
div.rn.f32 %f13, %f12, 0f3FED844D;
mul.f32 %f14, %f13, 0f43600000;
div.rn.f32 %f15, %f14, 0f437F0000;
add.f32 %f16, %f15, 0f43000000;
max.f32 %f17, %f16, 0f00000000;
min.f32 %f18, %f17, 0f437F0000;
add.f32 %f19, %f18, 0f3F000000;
cvt.rzi.u32.f32 %r39, %f19;
cvt.u16.u32 %rs13, %r39;
sub.f32 %f20, %f6, %f11;
div.rn.f32 %f21, %f20, 0f3FC9930C;
mul.f32 %f22, %f21, 0f43600000;
div.rn.f32 %f23, %f22, 0f437F0000;
add.f32 %f24, %f23, 0f43000000;
max.f32 %f25, %f24, 0f00000000;
min.f32 %f26, %f25, 0f437F0000;
add.f32 %f27, %f26, 0f3F000000;
cvt.rzi.u32.f32 %r40, %f27;
cvt.u16.u32 %rs14, %r40;
mul.lo.s32 %r41, %r12, %r1;
shl.b32 %r42, %r8, 1;
add.s32 %r43, %r41, %r42;
cvt.u64.u32 %rd7, %r43;
add.s64 %rd8, %rd1, %rd7;
st.global.u8 [%rd8], %rs13;
st.global.u8 [%rd8+1], %rs14;
CHROMA_DONE:
ret;
}
.visible .entry extract_alpha(
.param .u64 dst,
.param .u32 dst_pitch,
.param .u64 src,
.param .u32 src_pitch,
.param .u32 width,
.param .u32 height
)
{
.reg .pred %p<4>;
.reg .b16 %rs<4>;
.reg .b32 %r<20>;
.reg .b64 %rd<12>;
ld.param.u64 %rd1, [dst];
ld.param.u32 %r1, [dst_pitch];
ld.param.u64 %rd2, [src];
ld.param.u32 %r2, [src_pitch];
ld.param.u32 %r3, [width];
ld.param.u32 %r4, [height];
mov.u32 %r5, %ctaid.x;
mov.u32 %r6, %ntid.x;
mov.u32 %r7, %tid.x;
mad.lo.s32 %r8, %r5, %r6, %r7;
mov.u32 %r9, %ctaid.y;
mov.u32 %r10, %ntid.y;
mov.u32 %r11, %tid.y;
mad.lo.s32 %r12, %r9, %r10, %r11;
setp.ge.u32 %p1, %r8, %r3;
@%p1 bra ALPHA_DONE;
setp.ge.u32 %p2, %r12, %r4;
@%p2 bra ALPHA_DONE;
mul.lo.s32 %r13, %r12, %r2;
shl.b32 %r14, %r8, 2;
add.s32 %r15, %r13, %r14;
cvt.u64.u32 %rd3, %r15;
add.s64 %rd4, %rd2, %rd3;
ld.global.u8 %rs1, [%rd4+3];
mul.lo.s32 %r16, %r12, %r1;
add.s32 %r17, %r16, %r8;
cvt.u64.u32 %rd5, %r17;
add.s64 %rd6, %rd1, %rd5;
st.global.u8 [%rd6], %rs1;
ALPHA_DONE:
ret;
}
.visible .entry extract_alpha_half(
.param .u64 dst,
.param .u32 dst_pitch,
.param .u64 src,
.param .u32 src_pitch,
.param .u32 half_width,
.param .u32 half_height
)
{
.reg .pred %p<4>;
.reg .b16 %rs<8>;
.reg .b32 %r<30>;
.reg .b64 %rd<14>;
ld.param.u64 %rd1, [dst];
ld.param.u32 %r1, [dst_pitch];
ld.param.u64 %rd2, [src];
ld.param.u32 %r2, [src_pitch];
ld.param.u32 %r3, [half_width];
ld.param.u32 %r4, [half_height];
mov.u32 %r5, %ctaid.x;
mov.u32 %r6, %ntid.x;
mov.u32 %r7, %tid.x;
mad.lo.s32 %r8, %r5, %r6, %r7;
mov.u32 %r9, %ctaid.y;
mov.u32 %r10, %ntid.y;
mov.u32 %r11, %tid.y;
mad.lo.s32 %r12, %r9, %r10, %r11;
setp.ge.u32 %p1, %r8, %r3;
@%p1 bra ALPHA_HALF_DONE;
setp.ge.u32 %p2, %r12, %r4;
@%p2 bra ALPHA_HALF_DONE;
shl.b32 %r13, %r12, 1;
shl.b32 %r14, %r8, 1;
mul.lo.s32 %r15, %r13, %r2;
shl.b32 %r16, %r14, 2;
add.s32 %r17, %r15, %r16;
cvt.u64.u32 %rd3, %r17;
add.s64 %rd4, %rd2, %rd3;
cvt.u64.u32 %rd5, %r2;
add.s64 %rd6, %rd4, %rd5;
ld.global.u8 %r18, [%rd4+3];
ld.global.u8 %r19, [%rd4+7];
ld.global.u8 %r20, [%rd6+3];
ld.global.u8 %r21, [%rd6+7];
add.s32 %r22, %r18, %r19;
add.s32 %r23, %r22, %r20;
add.s32 %r24, %r23, %r21;
add.s32 %r25, %r24, 2;
shr.u32 %r26, %r25, 2;
mul.lo.s32 %r27, %r12, %r1;
add.s32 %r28, %r27, %r8;
cvt.u64.u32 %rd7, %r28;
add.s64 %rd8, %rd1, %rd7;
cvt.u16.u32 %rs1, %r26;
st.global.u8 [%rd8], %rs1;
ALPHA_HALF_DONE:
ret;
}
"#;
#[cfg_attr(windows, link(name = "nvcuda", kind = "raw-dylib"))]
#[cfg_attr(not(windows), link(name = "cuda"))]
unsafe extern "C" {
fn cuInit(flags: c_uint) -> CUresult;
fn cuDeviceGet(device: *mut CUdevice, ordinal: c_int) -> CUresult;
fn cuDevicePrimaryCtxRetain(ctx: *mut CUcontext, device: CUdevice) -> CUresult;
fn cuDevicePrimaryCtxRelease_v2(device: CUdevice) -> CUresult;
fn cuCtxPushCurrent_v2(ctx: CUcontext) -> CUresult;
fn cuCtxPopCurrent_v2(ctx: *mut CUcontext) -> CUresult;
fn cuMemcpy2D_v2(copy: *const CudaMemcpy2D) -> CUresult;
fn cuMemAlloc_v2(ptr: *mut CUdeviceptr, size: usize) -> CUresult;
fn cuMemFree_v2(ptr: CUdeviceptr) -> CUresult;
fn cuMemcpyHtoD_v2(dst: CUdeviceptr, src: *const c_void, size: usize) -> CUresult;
fn cuMemsetD2D8_v2(
dst: CUdeviceptr,
dst_pitch: usize,
value: u8,
width: usize,
height: usize,
) -> CUresult;
fn cuMemsetD2D16_v2(
dst: CUdeviceptr,
dst_pitch: usize,
value: u16,
width: usize,
height: usize,
) -> CUresult;
fn cuGetErrorString(error: CUresult, str_: *mut *const c_char) -> CUresult;
fn cuModuleLoadData(module: *mut CUmodule, image: *const c_void) -> CUresult;
fn cuModuleUnload(module: CUmodule) -> CUresult;
fn cuModuleGetFunction(
func: *mut CUfunction,
module: CUmodule,
name: *const c_char,
) -> CUresult;
fn cuLaunchKernel(
f: CUfunction,
grid_x: c_uint,
grid_y: c_uint,
grid_z: c_uint,
block_x: c_uint,
block_y: c_uint,
block_z: c_uint,
shared_bytes: c_uint,
stream: *mut c_void,
params: *mut *mut c_void,
extra: *mut *mut c_void,
) -> CUresult;
fn cuCtxSynchronize() -> CUresult;
}
const BLEND_PTX: &str = r#"
.version 6.0
.target sm_50
.address_size 64
.visible .entry blend_plane(
.param .u64 dst,
.param .u32 dst_pitch,
.param .u64 src,
.param .u32 src_pitch,
.param .u32 width,
.param .u32 height,
.param .u32 alpha
)
{
.reg .pred %p<4>;
.reg .b16 %rs<4>;
.reg .b32 %r<32>;
.reg .b64 %rd<16>;
ld.param.u64 %rd1, [dst];
ld.param.u32 %r1, [dst_pitch];
ld.param.u64 %rd2, [src];
ld.param.u32 %r2, [src_pitch];
ld.param.u32 %r3, [width];
ld.param.u32 %r4, [height];
ld.param.u32 %r5, [alpha];
mov.u32 %r6, %ctaid.x;
mov.u32 %r7, %ntid.x;
mov.u32 %r8, %tid.x;
mad.lo.s32 %r9, %r6, %r7, %r8;
mov.u32 %r10, %ctaid.y;
mov.u32 %r11, %ntid.y;
mov.u32 %r12, %tid.y;
mad.lo.s32 %r13, %r10, %r11, %r12;
setp.ge.u32 %p1, %r9, %r3;
@%p1 bra DONE;
setp.ge.u32 %p2, %r13, %r4;
@%p2 bra DONE;
mad.lo.s32 %r14, %r13, %r1, %r9;
cvt.u64.u32 %rd3, %r14;
add.s64 %rd4, %rd1, %rd3;
mad.lo.s32 %r15, %r13, %r2, %r9;
cvt.u64.u32 %rd5, %r15;
add.s64 %rd6, %rd2, %rd5;
ld.global.u8 %r16, [%rd4];
ld.global.u8 %r17, [%rd6];
mul.lo.s32 %r18, %r17, %r5;
sub.s32 %r19, 255, %r5;
mul.lo.s32 %r20, %r16, %r19;
add.s32 %r21, %r18, %r20;
add.s32 %r22, %r21, 127;
div.u32 %r23, %r22, 255;
cvt.u16.u32 %rs1, %r23;
st.global.u8 [%rd4], %rs1;
DONE:
ret;
}
.visible .entry blend_masked(
.param .u64 dst,
.param .u32 dst_pitch,
.param .u64 mask,
.param .u32 mask_pitch,
.param .u32 width,
.param .u32 height,
.param .u32 value_even,
.param .u32 value_odd,
.param .u32 opacity,
.param .u32 mask_shift
)
{
.reg .pred %p<4>;
.reg .b16 %rs<4>;
.reg .b32 %r<40>;
.reg .b64 %rd<16>;
ld.param.u64 %rd1, [dst];
ld.param.u32 %r1, [dst_pitch];
ld.param.u64 %rd2, [mask];
ld.param.u32 %r2, [mask_pitch];
ld.param.u32 %r3, [width];
ld.param.u32 %r4, [height];
ld.param.u32 %r5, [value_even];
ld.param.u32 %r6, [value_odd];
ld.param.u32 %r7, [opacity];
ld.param.u32 %r31, [mask_shift];
mov.u32 %r8, %ctaid.x;
mov.u32 %r9, %ntid.x;
mov.u32 %r10, %tid.x;
mad.lo.s32 %r11, %r8, %r9, %r10;
mov.u32 %r12, %ctaid.y;
mov.u32 %r13, %ntid.y;
mov.u32 %r14, %tid.y;
mad.lo.s32 %r15, %r12, %r13, %r14;
setp.ge.u32 %p1, %r11, %r3;
@%p1 bra MDONE;
setp.ge.u32 %p2, %r15, %r4;
@%p2 bra MDONE;
mad.lo.s32 %r16, %r15, %r1, %r11;
cvt.u64.u32 %rd3, %r16;
add.s64 %rd4, %rd1, %rd3;
shr.u32 %r32, %r11, %r31;
mad.lo.s32 %r17, %r15, %r2, %r32;
cvt.u64.u32 %rd5, %r17;
add.s64 %rd6, %rd2, %rd5;
ld.global.u8 %r18, [%rd4];
ld.global.u8 %r19, [%rd6];
mul.lo.s32 %r20, %r19, %r7;
add.s32 %r21, %r20, 127;
div.u32 %r22, %r21, 255;
and.b32 %r23, %r11, 1;
setp.eq.u32 %p3, %r23, 0;
selp.b32 %r24, %r5, %r6, %p3;
mul.lo.s32 %r25, %r24, %r22;
sub.s32 %r26, 255, %r22;
mul.lo.s32 %r27, %r18, %r26;
add.s32 %r28, %r25, %r27;
add.s32 %r29, %r28, 127;
div.u32 %r30, %r29, 255;
cvt.u16.u32 %rs2, %r30;
st.global.u8 [%rd4], %rs2;
MDONE:
ret;
}
.visible .entry blend_plane_masked(
.param .u64 dst,
.param .u32 dst_pitch,
.param .u64 src,
.param .u32 src_pitch,
.param .u64 mask,
.param .u32 mask_pitch,
.param .u32 width,
.param .u32 height,
.param .u32 opacity,
.param .u32 mask_shift
)
{
.reg .pred %p<4>;
.reg .b16 %rs<4>;
.reg .b32 %r<40>;
.reg .b64 %rd<20>;
ld.param.u64 %rd1, [dst];
ld.param.u32 %r1, [dst_pitch];
ld.param.u64 %rd7, [src];
ld.param.u32 %r33, [src_pitch];
ld.param.u64 %rd2, [mask];
ld.param.u32 %r2, [mask_pitch];
ld.param.u32 %r3, [width];
ld.param.u32 %r4, [height];
ld.param.u32 %r7, [opacity];
ld.param.u32 %r31, [mask_shift];
mov.u32 %r8, %ctaid.x;
mov.u32 %r9, %ntid.x;
mov.u32 %r10, %tid.x;
mad.lo.s32 %r11, %r8, %r9, %r10;
mov.u32 %r12, %ctaid.y;
mov.u32 %r13, %ntid.y;
mov.u32 %r14, %tid.y;
mad.lo.s32 %r15, %r12, %r13, %r14;
setp.ge.u32 %p1, %r11, %r3;
@%p1 bra PMDONE;
setp.ge.u32 %p2, %r15, %r4;
@%p2 bra PMDONE;
mad.lo.s32 %r16, %r15, %r1, %r11;
cvt.u64.u32 %rd3, %r16;
add.s64 %rd4, %rd1, %rd3;
mad.lo.s32 %r34, %r15, %r33, %r11;
cvt.u64.u32 %rd8, %r34;
add.s64 %rd9, %rd7, %rd8;
shr.u32 %r32, %r11, %r31;
mad.lo.s32 %r17, %r15, %r2, %r32;
cvt.u64.u32 %rd5, %r17;
add.s64 %rd6, %rd2, %rd5;
ld.global.u8 %r18, [%rd4];
ld.global.u8 %r19, [%rd6];
ld.global.u8 %r24, [%rd9];
mul.lo.s32 %r20, %r19, %r7;
add.s32 %r21, %r20, 127;
div.u32 %r22, %r21, 255;
mul.lo.s32 %r25, %r24, %r22;
sub.s32 %r26, 255, %r22;
mul.lo.s32 %r27, %r18, %r26;
add.s32 %r28, %r25, %r27;
add.s32 %r29, %r28, 127;
div.u32 %r30, %r29, 255;
cvt.u16.u32 %rs2, %r30;
st.global.u8 [%rd4], %rs2;
PMDONE:
ret;
}
"#;
#[derive(Debug, ThisError)]
pub enum CudaDriverError {
#[error("{call} failed: {message}")]
Call {
call: &'static str,
message: String,
},
#[error("the CUDA driver rejected this crate's blend kernel: {0}")]
KernelRejected(String),
#[error("a coverage mask smaller than 2x2 has no chroma samples to blend into")]
EmptyMask,
}
fn check(call: &'static str, result: CUresult) -> Result<(), CudaDriverError> {
if result == CUDA_SUCCESS {
return Ok(());
}
let mut raw: *const c_char = std::ptr::null();
let message = unsafe {
if cuGetErrorString(result, &mut raw) == CUDA_SUCCESS && !raw.is_null() {
CStr::from_ptr(raw).to_string_lossy().into_owned()
} else {
format!("CUDA error {result}")
}
};
Err(CudaDriverError::Call { call, message })
}
pub(crate) struct CudaDriver {
device: CUdevice,
ctx: CUcontext,
module: CUmodule,
blend: CUfunction,
blend_masked: CUfunction,
#[allow(dead_code)]
blend_plane_masked: CUfunction,
convert_module: CUmodule,
bgra_to_luma: CUfunction,
bgra_to_chroma: CUfunction,
#[allow(dead_code)]
extract_alpha: CUfunction,
#[allow(dead_code)]
extract_alpha_half: CUfunction,
}
unsafe impl Send for CudaDriver {}
unsafe impl Sync for CudaDriver {}
impl CudaDriver {
pub(crate) fn retain_primary() -> Result<Self, CudaDriverError> {
unsafe {
check("cuInit", cuInit(0))?;
let mut device: CUdevice = 0;
check("cuDeviceGet", cuDeviceGet(&mut device, 0))?;
let mut ctx: CUcontext = std::ptr::null_mut();
check(
"cuDevicePrimaryCtxRetain",
cuDevicePrimaryCtxRetain(&mut ctx, device),
)?;
check("cuCtxPushCurrent", cuCtxPushCurrent_v2(ctx))?;
let loaded = load_module(
BLEND_PTX,
["blend_plane", "blend_masked", "blend_plane_masked"],
)
.and_then(|blend| {
match load_module(
CONVERT_PTX,
[
"bgra_to_luma",
"bgra_to_chroma",
"extract_alpha",
"extract_alpha_half",
],
) {
Ok(convert) => Ok((blend, convert)),
Err(error) => {
cuModuleUnload(blend.0);
Err(error)
}
}
});
let mut popped: CUcontext = std::ptr::null_mut();
check("cuCtxPopCurrent", cuCtxPopCurrent_v2(&mut popped))?;
let (
(module, [blend, blend_masked, blend_plane_masked]),
(
convert_module,
[
bgra_to_luma,
bgra_to_chroma,
extract_alpha,
extract_alpha_half,
],
),
) = match loaded {
Ok(modules) => modules,
Err(error) => {
cuDevicePrimaryCtxRelease_v2(device);
return Err(error);
}
};
Ok(Self {
device,
ctx,
module,
blend,
blend_masked,
blend_plane_masked,
convert_module,
bgra_to_luma,
bgra_to_chroma,
extract_alpha,
extract_alpha_half,
})
}
}
fn with_context<T>(
&self,
f: impl FnOnce() -> Result<T, CudaDriverError>,
) -> Result<T, CudaDriverError> {
unsafe { check("cuCtxPushCurrent", cuCtxPushCurrent_v2(self.ctx))? };
let value = f();
let mut popped: CUcontext = std::ptr::null_mut();
unsafe { check("cuCtxPopCurrent", cuCtxPopCurrent_v2(&mut popped))? };
value
}
pub(crate) fn fill_nv12(
&self,
surface: Nv12Surface,
width: u32,
height: u32,
color: Color,
) -> Result<(), CudaDriverError> {
let (y, u, v) = rgb_to_bt709_limited(color);
let chroma = u16::from(u) | (u16::from(v) << 8);
self.with_context(|| unsafe {
check(
"cuMemsetD2D8",
cuMemsetD2D8_v2(
surface.luma,
surface.luma_pitch,
y,
width as usize,
height as usize,
),
)?;
check(
"cuMemsetD2D16",
cuMemsetD2D16_v2(
surface.chroma,
surface.chroma_pitch,
chroma,
(width / 2) as usize,
(height / 2) as usize,
),
)
})
}
pub(crate) fn blit_nv12(
&self,
source: Nv12Surface,
destination: Nv12Surface,
region: Nv12Region,
) -> Result<(), CudaDriverError> {
let Nv12Region {
source_x,
source_y,
destination_x,
destination_y,
width,
height,
} = region;
debug_assert!(
[
source_x,
source_y,
destination_x,
destination_y,
width,
height
]
.iter()
.all(|value| value.is_multiple_of(2)),
"NV12 blits must be aligned to the 2x2 chroma grid"
);
if width == 0 || height == 0 {
return Ok(());
}
self.with_context(|| unsafe {
let luma = CudaMemcpy2D {
src_memory_type: CU_MEMORYTYPE_DEVICE,
src_device: source.luma,
src_pitch: source.luma_pitch,
src_x_in_bytes: source_x as usize,
src_y: source_y as usize,
dst_memory_type: CU_MEMORYTYPE_DEVICE,
dst_device: destination.luma,
dst_pitch: destination.luma_pitch,
dst_x_in_bytes: destination_x as usize,
dst_y: destination_y as usize,
width_in_bytes: width as usize,
height: height as usize,
..CudaMemcpy2D::default()
};
check("cuMemcpy2D", cuMemcpy2D_v2(&luma))?;
let chroma = CudaMemcpy2D {
src_memory_type: CU_MEMORYTYPE_DEVICE,
src_device: source.chroma,
src_pitch: source.chroma_pitch,
src_x_in_bytes: source_x as usize,
src_y: (source_y / 2) as usize,
dst_memory_type: CU_MEMORYTYPE_DEVICE,
dst_device: destination.chroma,
dst_pitch: destination.chroma_pitch,
dst_x_in_bytes: destination_x as usize,
dst_y: (destination_y / 2) as usize,
width_in_bytes: width as usize,
height: (height / 2) as usize,
..CudaMemcpy2D::default()
};
check("cuMemcpy2D", cuMemcpy2D_v2(&chroma))
})
}
pub(crate) fn blit_bgra(
&self,
source: BgraSurface,
destination: BgraSurface,
source_x: u32,
source_y: u32,
width: u32,
height: u32,
) -> Result<(), CudaDriverError> {
if width == 0 || height == 0 {
return Ok(());
}
self.with_context(|| unsafe {
let copy = CudaMemcpy2D {
src_memory_type: CU_MEMORYTYPE_DEVICE,
src_device: source.pixels,
src_pitch: source.pitch,
src_x_in_bytes: source_x as usize * 4,
src_y: source_y as usize,
dst_memory_type: CU_MEMORYTYPE_DEVICE,
dst_device: destination.pixels,
dst_pitch: destination.pitch,
width_in_bytes: width as usize * 4,
height: height as usize,
..CudaMemcpy2D::default()
};
check("cuMemcpy2D", cuMemcpy2D_v2(©))
})
}
pub(crate) fn blend_nv12(
&self,
source: Nv12Surface,
destination: Nv12Surface,
region: Nv12Region,
alpha: u8,
) -> Result<(), CudaDriverError> {
let Nv12Region {
source_x,
source_y,
destination_x,
destination_y,
width,
height,
} = region;
debug_assert!(
[
source_x,
source_y,
destination_x,
destination_y,
width,
height
]
.iter()
.all(|value| value.is_multiple_of(2)),
"NV12 blends must be aligned to the 2x2 chroma grid"
);
if width == 0 || height == 0 {
return Ok(());
}
self.with_context(|| {
self.launch_blend(
destination.luma
+ u64::from(destination_y) * destination.luma_pitch as u64
+ u64::from(destination_x),
destination.luma_pitch,
source.luma + u64::from(source_y) * source.luma_pitch as u64 + u64::from(source_x),
source.luma_pitch,
width,
height,
alpha,
)?;
self.launch_blend(
destination.chroma
+ u64::from(destination_y / 2) * destination.chroma_pitch as u64
+ u64::from(destination_x),
destination.chroma_pitch,
source.chroma
+ u64::from(source_y / 2) * source.chroma_pitch as u64
+ u64::from(source_x),
source.chroma_pitch,
width,
height / 2,
alpha,
)
})
}
#[allow(clippy::too_many_arguments)]
fn launch_blend(
&self,
mut dst: CUdeviceptr,
dst_pitch: usize,
mut src: CUdeviceptr,
src_pitch: usize,
width: u32,
height: u32,
alpha: u8,
) -> Result<(), CudaDriverError> {
const BLOCK: u32 = 16;
let mut dst_pitch = dst_pitch as u32;
let mut src_pitch = src_pitch as u32;
let mut width = width;
let mut height = height;
let mut alpha = u32::from(alpha);
let mut params: [*mut c_void; 7] = [
(&mut dst) as *mut _ as *mut c_void,
(&mut dst_pitch) as *mut _ as *mut c_void,
(&mut src) as *mut _ as *mut c_void,
(&mut src_pitch) as *mut _ as *mut c_void,
(&mut width) as *mut _ as *mut c_void,
(&mut height) as *mut _ as *mut c_void,
(&mut alpha) as *mut _ as *mut c_void,
];
unsafe {
check(
"cuLaunchKernel",
cuLaunchKernel(
self.blend,
width.div_ceil(BLOCK),
height.div_ceil(BLOCK),
1,
BLOCK,
BLOCK,
1,
0,
std::ptr::null_mut(),
params.as_mut_ptr(),
std::ptr::null_mut(),
),
)
}
}
pub(crate) fn bgra_to_nv12(
&self,
source: BgraSurface,
destination: Nv12Surface,
width: u32,
height: u32,
) -> Result<(), CudaDriverError> {
const BLOCK: u32 = 16;
self.with_context(|| unsafe {
let mut luma = destination.luma;
let mut luma_pitch = destination.luma_pitch as u32;
let mut pixels = source.pixels;
let mut source_pitch = source.pitch as u32;
let mut width = width;
let mut height = height;
let mut luma_params: [*mut c_void; 6] = [
(&mut luma) as *mut _ as *mut c_void,
(&mut luma_pitch) as *mut _ as *mut c_void,
(&mut pixels) as *mut _ as *mut c_void,
(&mut source_pitch) as *mut _ as *mut c_void,
(&mut width) as *mut _ as *mut c_void,
(&mut height) as *mut _ as *mut c_void,
];
check(
"cuLaunchKernel",
cuLaunchKernel(
self.bgra_to_luma,
width.div_ceil(BLOCK),
height.div_ceil(BLOCK),
1,
BLOCK,
BLOCK,
1,
0,
std::ptr::null_mut(),
luma_params.as_mut_ptr(),
std::ptr::null_mut(),
),
)?;
let mut chroma = destination.chroma;
let mut chroma_pitch = destination.chroma_pitch as u32;
let mut half_width = width / 2;
let mut half_height = height / 2;
let mut chroma_params: [*mut c_void; 6] = [
(&mut chroma) as *mut _ as *mut c_void,
(&mut chroma_pitch) as *mut _ as *mut c_void,
(&mut pixels) as *mut _ as *mut c_void,
(&mut source_pitch) as *mut _ as *mut c_void,
(&mut half_width) as *mut _ as *mut c_void,
(&mut half_height) as *mut _ as *mut c_void,
];
check(
"cuLaunchKernel",
cuLaunchKernel(
self.bgra_to_chroma,
half_width.div_ceil(BLOCK),
half_height.div_ceil(BLOCK),
1,
BLOCK,
BLOCK,
1,
0,
std::ptr::null_mut(),
chroma_params.as_mut_ptr(),
std::ptr::null_mut(),
),
)
})
}
pub(crate) fn synchronize(&self) -> Result<(), CudaDriverError> {
self.with_context(|| unsafe { check("cuCtxSynchronize", cuCtxSynchronize()) })
}
}
pub(crate) struct CudaMask {
ctx: CUcontext,
full: CUdeviceptr,
half: CUdeviceptr,
pub(crate) width: u32,
pub(crate) height: u32,
}
unsafe impl Send for CudaMask {}
unsafe impl Sync for CudaMask {}
pub(crate) struct CudaOverlayScratch {
ctx: CUcontext,
luma: CUdeviceptr,
chroma: CUdeviceptr,
alpha_full: CUdeviceptr,
alpha_half: CUdeviceptr,
pub(crate) width: u32,
pub(crate) height: u32,
}
unsafe impl Send for CudaOverlayScratch {}
impl Drop for CudaOverlayScratch {
fn drop(&mut self) {
unsafe {
if cuCtxPushCurrent_v2(self.ctx) == CUDA_SUCCESS {
cuMemFree_v2(self.luma);
cuMemFree_v2(self.chroma);
cuMemFree_v2(self.alpha_full);
cuMemFree_v2(self.alpha_half);
let mut popped: CUcontext = std::ptr::null_mut();
cuCtxPopCurrent_v2(&mut popped);
}
}
}
}
impl Drop for CudaMask {
fn drop(&mut self) {
unsafe {
if cuCtxPushCurrent_v2(self.ctx) == CUDA_SUCCESS {
cuMemFree_v2(self.full);
cuMemFree_v2(self.half);
let mut popped: CUcontext = std::ptr::null_mut();
cuCtxPopCurrent_v2(&mut popped);
}
}
}
}
impl CudaDriver {
#[allow(dead_code)]
pub(crate) fn overlay_scratch(
&self,
width: u32,
height: u32,
) -> Result<CudaOverlayScratch, CudaDriverError> {
let width = width & !1;
let height = height & !1;
if width == 0 || height == 0 {
return Err(CudaDriverError::EmptyMask);
}
let pixels = (width * height) as usize;
self.with_context(|| unsafe {
let mut taken = Vec::with_capacity(4);
let mut alloc = |bytes: usize| -> Result<CUdeviceptr, CudaDriverError> {
let mut pointer = 0;
match check("cuMemAlloc", cuMemAlloc_v2(&mut pointer, bytes)) {
Ok(()) => {
taken.push(pointer);
Ok(pointer)
}
Err(error) => {
for pointer in taken.drain(..) {
cuMemFree_v2(pointer);
}
Err(error)
}
}
};
let luma = alloc(pixels)?;
let chroma = alloc(pixels / 2)?;
let alpha_full = alloc(pixels)?;
let alpha_half = alloc(pixels / 4)?;
Ok(CudaOverlayScratch {
ctx: self.ctx,
luma,
chroma,
alpha_full,
alpha_half,
width,
height,
})
})
}
pub(crate) fn upload_mask(
&self,
coverage: &[u8],
width: u32,
height: u32,
) -> Result<CudaMask, CudaDriverError> {
let full_width = width & !1;
let full_height = height & !1;
debug_assert_eq!(coverage.len(), (width * height) as usize);
let (half_width, half_height) = (full_width / 2, full_height / 2);
if full_width == 0 || full_height == 0 {
return Err(CudaDriverError::EmptyMask);
}
let mut half = vec![0u8; (half_width * half_height) as usize];
for y in 0..half_height as usize {
for x in 0..half_width as usize {
let at = |dy: usize, dx: usize| {
u32::from(coverage[(y * 2 + dy) * width as usize + x * 2 + dx])
};
half[y * half_width as usize + x] =
((at(0, 0) + at(0, 1) + at(1, 0) + at(1, 1) + 2) / 4) as u8;
}
}
self.with_context(|| unsafe {
let mut full = 0;
check(
"cuMemAlloc",
cuMemAlloc_v2(&mut full, (full_width * full_height) as usize),
)?;
let mut half_ptr = 0;
check("cuMemAlloc", cuMemAlloc_v2(&mut half_ptr, half.len()))?;
for y in 0..full_height as usize {
let row = &coverage[y * width as usize..y * width as usize + full_width as usize];
check(
"cuMemcpyHtoD",
cuMemcpyHtoD_v2(
full + (y * full_width as usize) as u64,
row.as_ptr().cast(),
full_width as usize,
),
)?;
}
check(
"cuMemcpyHtoD",
cuMemcpyHtoD_v2(half_ptr, half.as_ptr().cast(), half.len()),
)?;
Ok(CudaMask {
ctx: self.ctx,
full,
half: half_ptr,
width: full_width,
height: full_height,
})
})
}
#[allow(dead_code, clippy::too_many_arguments)]
pub(crate) fn blend_bgra_nv12(
&self,
source: BgraSurface,
scratch: &CudaOverlayScratch,
destination: Nv12Surface,
region: Nv12Region,
opacity: u8,
) -> Result<(), CudaDriverError> {
let Nv12Region {
source_x,
source_y,
destination_x,
destination_y,
width,
height,
} = region;
if width == 0 || height == 0 {
return Ok(());
}
debug_assert!(
width <= scratch.width && height <= scratch.height,
"scratch is smaller than the region it is asked to hold"
);
let pixels =
source.pixels + u64::from(source_y) * source.pitch as u64 + u64::from(source_x) * 4;
let scratch_pitch = scratch.width;
let half_pitch = scratch.width / 2;
self.with_context(|| {
self.launch_plane(
self.bgra_to_luma,
scratch.luma,
scratch_pitch,
pixels,
source.pitch as u32,
width,
height,
)?;
self.launch_plane(
self.bgra_to_chroma,
scratch.chroma,
scratch_pitch,
pixels,
source.pitch as u32,
width / 2,
height / 2,
)?;
self.launch_plane(
self.extract_alpha,
scratch.alpha_full,
scratch_pitch,
pixels,
source.pitch as u32,
width,
height,
)?;
self.launch_plane(
self.extract_alpha_half,
scratch.alpha_half,
half_pitch,
pixels,
source.pitch as u32,
width / 2,
height / 2,
)?;
self.launch_plane_masked(
destination.luma
+ u64::from(destination_y) * destination.luma_pitch as u64
+ u64::from(destination_x),
destination.luma_pitch as u32,
scratch.luma,
scratch_pitch,
scratch.alpha_full,
scratch_pitch,
width,
height,
opacity,
0,
)?;
self.launch_plane_masked(
destination.chroma
+ u64::from(destination_y / 2) * destination.chroma_pitch as u64
+ u64::from(destination_x),
destination.chroma_pitch as u32,
scratch.chroma,
scratch_pitch,
scratch.alpha_half,
half_pitch,
width,
height / 2,
opacity,
1,
)
})
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn blend_mask_nv12(
&self,
destination: Nv12Surface,
x: u32,
y: u32,
mask: &CudaMask,
mask_x: u32,
mask_y: u32,
width: u32,
height: u32,
color: Color,
opacity: u8,
) -> Result<(), CudaDriverError> {
if width == 0 || height == 0 {
return Ok(());
}
let (luma, u, v) = rgb_to_bt709_limited(color);
self.with_context(|| {
self.launch_masked(
destination.luma + u64::from(y) * destination.luma_pitch as u64 + u64::from(x),
destination.luma_pitch as u32,
mask.full + u64::from(mask_y) * u64::from(mask.width) + u64::from(mask_x),
mask.width,
width,
height,
(u32::from(luma), u32::from(luma)),
opacity,
0,
)?;
self.launch_masked(
destination.chroma
+ u64::from(y / 2) * destination.chroma_pitch as u64
+ u64::from(x),
destination.chroma_pitch as u32,
mask.half
+ u64::from(mask_y / 2) * u64::from(mask.width / 2)
+ u64::from(mask_x / 2),
mask.width / 2,
width,
height / 2,
(u32::from(u), u32::from(v)),
opacity,
1,
)
})
}
#[allow(clippy::too_many_arguments)]
fn launch_plane(
&self,
entry: CUfunction,
mut dst: CUdeviceptr,
mut dst_pitch: u32,
mut src: CUdeviceptr,
mut src_pitch: u32,
mut width: u32,
mut height: u32,
) -> Result<(), CudaDriverError> {
const BLOCK: u32 = 16;
let (grid_x, grid_y) = (width.div_ceil(BLOCK), height.div_ceil(BLOCK));
let mut params: [*mut c_void; 6] = [
(&mut dst) as *mut _ as *mut c_void,
(&mut dst_pitch) as *mut _ as *mut c_void,
(&mut src) as *mut _ as *mut c_void,
(&mut src_pitch) as *mut _ as *mut c_void,
(&mut width) as *mut _ as *mut c_void,
(&mut height) as *mut _ as *mut c_void,
];
unsafe {
check(
"cuLaunchKernel",
cuLaunchKernel(
entry,
grid_x,
grid_y,
1,
BLOCK,
BLOCK,
1,
0,
std::ptr::null_mut(),
params.as_mut_ptr(),
std::ptr::null_mut(),
),
)
}
}
#[allow(clippy::too_many_arguments)]
fn launch_plane_masked(
&self,
mut dst: CUdeviceptr,
mut dst_pitch: u32,
mut src: CUdeviceptr,
mut src_pitch: u32,
mut mask: CUdeviceptr,
mut mask_pitch: u32,
mut width: u32,
mut height: u32,
opacity: u8,
mut shift: u32,
) -> Result<(), CudaDriverError> {
const BLOCK: u32 = 16;
let (grid_x, grid_y) = (width.div_ceil(BLOCK), height.div_ceil(BLOCK));
let mut opacity = u32::from(opacity);
let mut params: [*mut c_void; 10] = [
(&mut dst) as *mut _ as *mut c_void,
(&mut dst_pitch) as *mut _ as *mut c_void,
(&mut src) as *mut _ as *mut c_void,
(&mut src_pitch) as *mut _ as *mut c_void,
(&mut mask) as *mut _ as *mut c_void,
(&mut mask_pitch) as *mut _ as *mut c_void,
(&mut width) as *mut _ as *mut c_void,
(&mut height) as *mut _ as *mut c_void,
(&mut opacity) as *mut _ as *mut c_void,
(&mut shift) as *mut _ as *mut c_void,
];
unsafe {
check(
"cuLaunchKernel",
cuLaunchKernel(
self.blend_plane_masked,
grid_x,
grid_y,
1,
BLOCK,
BLOCK,
1,
0,
std::ptr::null_mut(),
params.as_mut_ptr(),
std::ptr::null_mut(),
),
)
}
}
#[allow(clippy::too_many_arguments)]
fn launch_masked(
&self,
mut dst: CUdeviceptr,
mut dst_pitch: u32,
mut mask: CUdeviceptr,
mut mask_pitch: u32,
width: u32,
height: u32,
values: (u32, u32),
opacity: u8,
shift: u32,
) -> Result<(), CudaDriverError> {
const BLOCK: u32 = 16;
let mut width = width;
let mut height = height;
let (mut value_even, mut value_odd) = values;
let mut opacity = u32::from(opacity);
let mut shift = shift;
let mut params: [*mut c_void; 10] = [
(&mut dst) as *mut _ as *mut c_void,
(&mut dst_pitch) as *mut _ as *mut c_void,
(&mut mask) as *mut _ as *mut c_void,
(&mut mask_pitch) as *mut _ as *mut c_void,
(&mut width) as *mut _ as *mut c_void,
(&mut height) as *mut _ as *mut c_void,
(&mut value_even) as *mut _ as *mut c_void,
(&mut value_odd) as *mut _ as *mut c_void,
(&mut opacity) as *mut _ as *mut c_void,
(&mut shift) as *mut _ as *mut c_void,
];
unsafe {
check(
"cuLaunchKernel",
cuLaunchKernel(
self.blend_masked,
width.div_ceil(BLOCK),
height.div_ceil(BLOCK),
1,
BLOCK,
BLOCK,
1,
0,
std::ptr::null_mut(),
params.as_mut_ptr(),
std::ptr::null_mut(),
),
)
}
}
}
unsafe fn load_module<const N: usize>(
ptx: &str,
entry_names: [&str; N],
) -> Result<(CUmodule, [CUfunction; N]), CudaDriverError> {
unsafe {
let image = std::ffi::CString::new(ptx)
.map_err(|error| CudaDriverError::KernelRejected(error.to_string()))?;
let mut module: CUmodule = std::ptr::null_mut();
check(
"cuModuleLoadData",
cuModuleLoadData(&mut module, image.as_ptr().cast()),
)
.map_err(|error| CudaDriverError::KernelRejected(error.to_string()))?;
let mut entries = [std::ptr::null_mut(); N];
for (entry, name) in entries.iter_mut().zip(entry_names) {
let name = std::ffi::CString::new(name).expect("a literal without a nul");
if let Err(error) = check(
"cuModuleGetFunction",
cuModuleGetFunction(entry, module, name.as_ptr()),
) {
cuModuleUnload(module);
return Err(CudaDriverError::KernelRejected(error.to_string()));
}
}
Ok((module, entries))
}
}
impl Drop for CudaDriver {
fn drop(&mut self) {
unsafe {
if cuCtxPushCurrent_v2(self.ctx) == CUDA_SUCCESS {
cuModuleUnload(self.module);
cuModuleUnload(self.convert_module);
let mut popped: CUcontext = std::ptr::null_mut();
cuCtxPopCurrent_v2(&mut popped);
}
cuDevicePrimaryCtxRelease_v2(self.device);
}
}
}
#[derive(Debug, Clone, Copy)]
pub(crate) struct Nv12Region {
pub(crate) source_x: u32,
pub(crate) source_y: u32,
pub(crate) destination_x: u32,
pub(crate) destination_y: u32,
pub(crate) width: u32,
pub(crate) height: u32,
}
#[derive(Debug, Clone, Copy)]
pub(crate) struct Nv12Surface {
pub(crate) luma: CUdeviceptr,
pub(crate) luma_pitch: usize,
pub(crate) chroma: CUdeviceptr,
pub(crate) chroma_pitch: usize,
}
impl Nv12Surface {
pub(crate) fn from_frame(frame: &ffmpeg_next::frame::Video) -> Option<Self> {
let (luma, chroma, luma_pitch, chroma_pitch) = unsafe {
let ptr = frame.as_ptr();
(
(*ptr).data[0],
(*ptr).data[1],
(*ptr).linesize[0],
(*ptr).linesize[1],
)
};
if luma.is_null() || chroma.is_null() || luma_pitch <= 0 || chroma_pitch <= 0 {
return None;
}
Some(Self {
luma: luma as CUdeviceptr,
luma_pitch: luma_pitch as usize,
chroma: chroma as CUdeviceptr,
chroma_pitch: chroma_pitch as usize,
})
}
}
#[derive(Debug, Clone, Copy)]
pub(crate) struct BgraSurface {
pub(crate) pixels: CUdeviceptr,
pub(crate) pitch: usize,
}
impl BgraSurface {
pub(crate) fn from_frame(frame: &ffmpeg_next::frame::Video) -> Option<Self> {
let (pixels, pitch) = unsafe {
let ptr = frame.as_ptr();
((*ptr).data[0], (*ptr).linesize[0])
};
(!pixels.is_null() && pitch > 0).then_some(Self {
pixels: pixels as CUdeviceptr,
pitch: pitch as usize,
})
}
}
fn rgb_to_bt709_limited(color: Color) -> (u8, u8, u8) {
bt709_limited(
f32::from(color.red),
f32::from(color.green),
f32::from(color.blue),
)
}
pub(crate) fn bt709_limited(r: f32, g: f32, b: f32) -> (u8, u8, u8) {
let y = 0.2126 * r + 0.7152 * g + 0.0722 * b;
let u = (b - y) / 1.8556;
let v = (r - y) / 1.5748;
(
(16.0 + y * 219.0 / 255.0).round().clamp(0.0, 255.0) as u8,
(128.0 + u * 224.0 / 255.0).round().clamp(0.0, 255.0) as u8,
(128.0 + v * 224.0 / 255.0).round().clamp(0.0, 255.0) as u8,
)
}
#[cfg(test)]
mod tests {
use std::sync::{Arc, Mutex};
use ffmpeg_next::{self as ffmpeg};
use super::*;
use crate::{
buffer::MediaBuffer,
control::ControlMsg,
element::{Element, ElementType, Sink, Source, element_pp_log},
elements::{CudaDownload, CudaFrameFormat, CudaUpload},
pool::UnboundObjectPool,
pp_log::PpLog,
test_support::try_cuda_device,
};
struct CapturingSink {
pp_log: PpLog,
received: Arc<Mutex<Vec<MediaBuffer>>>,
}
impl Element for CapturingSink {
fn name(&self) -> Arc<str> {
"capture".into()
}
fn element_type(&self) -> ElementType {
ElementType::Other
}
fn pp_log(&self) -> &PpLog {
&self.pp_log
}
fn pp_log_mut(&mut self) -> &mut PpLog {
&mut self.pp_log
}
}
impl Sink for CapturingSink {
fn consume(&mut self, buf: MediaBuffer) -> crate::error::Result<()> {
self.received.lock().unwrap().push(buf);
Ok(())
}
fn control(&mut self, _msg: ControlMsg) -> crate::error::Result<()> {
Ok(())
}
}
fn capture(element: &mut dyn Source) -> Arc<Mutex<Vec<MediaBuffer>>> {
let received = Arc::new(Mutex::new(Vec::new()));
element.src_pads()[0].link(Box::new(CapturingSink {
received: received.clone(),
pp_log: element_pp_log(ElementType::Other, "capture", None),
}));
received
}
fn cuda_surface(
device: &crate::elements::CudaDevice,
width: u32,
height: u32,
luma: u8,
) -> Option<MediaBuffer> {
let Ok(mut upload) =
CudaUpload::new("upload", device, CudaFrameFormat::Nv12, width, height)
else {
eprintln!("skipping: this machine has no usable CUDA frames context");
return None;
};
let uploaded = capture(&mut upload);
let mut frame = ffmpeg::frame::Video::new(ffmpeg::format::Pixel::NV12, width, height);
let y_stride = frame.stride(0);
frame.data_mut(0)[..y_stride * height as usize].fill(luma);
let uv_stride = frame.stride(1);
frame.data_mut(1)[..uv_stride * (height / 2) as usize].fill(128);
let pool = UnboundObjectPool::new(0, ffmpeg::frame::Video::empty, |_| {});
let mut slot = pool.get();
*slot = frame;
upload
.consume(MediaBuffer::Video(Arc::new(slot)))
.expect("upload");
Some(uploaded.lock().unwrap().remove(0))
}
fn download(
device: &crate::elements::CudaDevice,
frame: MediaBuffer,
width: u32,
height: u32,
) -> Arc<crate::pool::UnboundObjectPoolRef<ffmpeg::frame::Video>> {
let mut download =
CudaDownload::new("download", device, CudaFrameFormat::Nv12, width, height);
let received = capture(&mut download);
download.consume(frame).expect("download");
let buf = received.lock().unwrap().remove(0);
match buf {
MediaBuffer::Video(frame) => frame,
other => panic!("expected a Video buffer, got {}", other.kind()),
}
}
#[test]
fn fill_then_blit_writes_the_expected_rectangles() {
let Some((device, _cuda_lock)) = try_cuda_device() else {
return;
};
let driver = match CudaDriver::retain_primary() {
Ok(driver) => driver,
Err(error) => {
eprintln!("skipping: no usable CUDA driver context ({error})");
return;
}
};
let (width, height) = (64u32, 64u32);
let Some(canvas) = cuda_surface(&device, width, height, 0) else {
return;
};
let Some(layer) = cuda_surface(&device, 32, 32, 200) else {
return;
};
let (MediaBuffer::Video(canvas_frame), MediaBuffer::Video(layer_frame)) = (&canvas, &layer)
else {
panic!("expected Video buffers");
};
let canvas_surface = Nv12Surface::from_frame(canvas_frame).expect("canvas planes");
let layer_surface = Nv12Surface::from_frame(layer_frame).expect("layer planes");
driver
.fill_nv12(canvas_surface, width, height, Color::WHITE)
.expect("fill");
driver
.blit_nv12(
layer_surface,
canvas_surface,
Nv12Region {
source_x: 0,
source_y: 0,
destination_x: 16,
destination_y: 8,
width: 32,
height: 32,
},
)
.expect("blit");
let out = download(&device, canvas.clone(), width, height);
let stride = out.stride(0);
let at = |x: usize, y: usize| out.data(0)[y * stride + x];
assert_eq!(at(0, 0), 235, "the fill did not cover the top-left corner");
assert_eq!(at(63, 63), 235, "the fill did not cover the bottom-right");
assert_eq!(at(16, 8), 200, "the blit missed its top-left corner");
assert_eq!(at(47, 39), 200, "the blit missed its bottom-right corner");
assert_eq!(at(15, 8), 235, "the blit wrote left of its rectangle");
assert_eq!(at(48, 8), 235, "the blit wrote right of its rectangle");
assert_eq!(at(16, 7), 235, "the blit wrote above its rectangle");
assert_eq!(at(16, 40), 235, "the blit wrote below its rectangle");
let uv_stride = out.stride(1);
assert_eq!(
out.data(1)[uv_stride * 20 + 4],
128,
"chroma did not survive the fill"
);
}
fn cuda_bgra_surface(
device: &crate::elements::CudaDevice,
width: u32,
height: u32,
pixel: impl Fn(u32, u32) -> [u8; 4],
) -> Option<MediaBuffer> {
let Ok(mut upload) =
CudaUpload::new("upload", device, CudaFrameFormat::Bgra, width, height)
else {
eprintln!("skipping: this machine has no usable CUDA frames context");
return None;
};
let uploaded = capture(&mut upload);
let mut frame = ffmpeg::frame::Video::new(ffmpeg::format::Pixel::BGRA, width, height);
let stride = frame.stride(0);
for y in 0..height {
let row = &mut frame.data_mut(0)[y as usize * stride..];
for x in 0..width {
row[x as usize * 4..x as usize * 4 + 4].copy_from_slice(&pixel(x, y));
}
}
let pool = UnboundObjectPool::new(0, ffmpeg::frame::Video::empty, |_| {});
let mut slot = pool.get();
*slot = frame;
upload
.consume(MediaBuffer::Video(Arc::new(slot)))
.expect("upload");
Some(uploaded.lock().unwrap().remove(0))
}
#[test]
fn a_bgra_layer_blends_under_its_own_alpha_rather_than_covering() {
let Some((device, _cuda_lock)) = crate::test_support::try_cuda_device() else {
return;
};
let Ok(driver) = CudaDriver::retain_primary() else {
eprintln!("skipping: no usable CUDA driver on this machine");
return;
};
const WIDTH: u32 = 64;
const HEIGHT: u32 = 32;
const CANVAS_LUMA: u8 = 90;
let pixel = |x: u32, _y: u32| {
if x < WIDTH / 2 {
[0, 0, 255, 255]
} else {
[0, 0, 255, 0]
}
};
let Some(layer) = cuda_bgra_surface(&device, WIDTH, HEIGHT, pixel) else {
return;
};
let Some(canvas) = cuda_surface(&device, WIDTH, HEIGHT, CANVAS_LUMA) else {
return;
};
let (MediaBuffer::Video(layer_frame), MediaBuffer::Video(canvas_frame)) = (&layer, &canvas)
else {
panic!("both uploads produce Video buffers");
};
let scratch = driver
.overlay_scratch(WIDTH, HEIGHT)
.expect("scratch for a layer this size");
driver
.blend_bgra_nv12(
BgraSurface::from_frame(layer_frame).expect("a BGRA surface"),
&scratch,
Nv12Surface::from_frame(canvas_frame).expect("an NV12 surface"),
Nv12Region {
source_x: 0,
source_y: 0,
destination_x: 0,
destination_y: 0,
width: WIDTH,
height: HEIGHT,
},
255,
)
.expect("blend");
driver.synchronize().expect("synchronize");
let blended = download(&device, canvas, WIDTH, HEIGHT);
let stride = blended.stride(0);
let luma = blended.data(0);
let (expected_red, _, _) = bt709_limited(255.0, 0.0, 0.0);
for y in 0..HEIGHT as usize {
for x in 0..(WIDTH / 2) as usize {
assert_eq!(
luma[y * stride + x],
expected_red,
"the opaque half should be the layer's own colour at ({x}, {y})"
);
}
for x in (WIDTH / 2) as usize..WIDTH as usize {
assert_eq!(
luma[y * stride + x],
CANVAS_LUMA,
"the transparent half should leave the canvas alone at ({x}, {y})"
);
}
}
}
#[test]
fn bgra_converts_to_nv12_exactly_as_the_shared_definition_says() {
let Some((device, _cuda_lock)) = crate::test_support::try_cuda_device() else {
return;
};
let Ok(driver) = CudaDriver::retain_primary() else {
eprintln!("skipping: no usable CUDA driver on this machine");
return;
};
const WIDTH: u32 = 64;
const HEIGHT: u32 = 32;
let pixel = |x: u32, y: u32| {
[
(x * 4 % 256) as u8,
(y * 8 % 256) as u8,
((x + y) * 3 % 256) as u8,
255,
]
};
let Some(source) = cuda_bgra_surface(&device, WIDTH, HEIGHT, pixel) else {
return;
};
let Some(destination) = cuda_surface(&device, WIDTH, HEIGHT, 0) else {
return;
};
let (MediaBuffer::Video(source_frame), MediaBuffer::Video(destination_frame)) =
(&source, &destination)
else {
panic!("both uploads produce Video buffers");
};
driver
.bgra_to_nv12(
BgraSurface::from_frame(source_frame).expect("a BGRA surface"),
Nv12Surface::from_frame(destination_frame).expect("an NV12 surface"),
WIDTH,
HEIGHT,
)
.expect("convert");
driver.synchronize().expect("synchronize");
let converted = download(&device, destination, WIDTH, HEIGHT);
let luma_stride = converted.stride(0);
let chroma_stride = converted.stride(1);
for y in 0..HEIGHT {
for x in 0..WIDTH {
let [b, g, r, _] = pixel(x, y);
let (expected, _, _) = bt709_limited(f32::from(r), f32::from(g), f32::from(b));
assert_eq!(
converted.data(0)[y as usize * luma_stride + x as usize],
expected,
"luma at {x},{y}"
);
}
}
for cy in 0..HEIGHT / 2 {
for cx in 0..WIDTH / 2 {
let mut sums = [0u32; 3];
for (dx, dy) in [(0, 0), (1, 0), (0, 1), (1, 1)] {
let [b, g, r, _] = pixel(cx * 2 + dx, cy * 2 + dy);
sums[0] += u32::from(b);
sums[1] += u32::from(g);
sums[2] += u32::from(r);
}
let (_, expected_u, expected_v) = bt709_limited(
sums[2] as f32 / 4.0,
sums[1] as f32 / 4.0,
sums[0] as f32 / 4.0,
);
let at = cy as usize * chroma_stride + cx as usize * 2;
assert_eq!(converted.data(1)[at], expected_u, "u at {cx},{cy}");
assert_eq!(converted.data(1)[at + 1], expected_v, "v at {cx},{cy}");
}
}
}
fn cuda_surface_with(
device: &crate::elements::CudaDevice,
width: u32,
height: u32,
luma: impl Fn(u32, u32) -> u8,
chroma: u8,
) -> Option<MediaBuffer> {
let Ok(mut upload) =
CudaUpload::new("upload", device, CudaFrameFormat::Nv12, width, height)
else {
eprintln!("skipping: this machine has no usable CUDA frames context");
return None;
};
let uploaded = capture(&mut upload);
let mut frame = ffmpeg::frame::Video::new(ffmpeg::format::Pixel::NV12, width, height);
let y_stride = frame.stride(0);
let plane = frame.data_mut(0);
for y in 0..height {
for x in 0..width {
plane[y as usize * y_stride + x as usize] = luma(x, y);
}
}
let uv_stride = frame.stride(1);
frame.data_mut(1)[..uv_stride * (height / 2) as usize].fill(chroma);
let pool = UnboundObjectPool::new(0, ffmpeg::frame::Video::empty, |_| {});
let mut slot = pool.get();
*slot = frame;
upload
.consume(MediaBuffer::Video(Arc::new(slot)))
.expect("upload");
Some(uploaded.lock().unwrap().remove(0))
}
#[test]
fn the_blend_kernel_matches_a_cpu_reference_byte_for_byte() {
let Some((device, _cuda_lock)) = try_cuda_device() else {
return;
};
let driver = match CudaDriver::retain_primary() {
Ok(driver) => driver,
Err(error) => {
eprintln!("skipping: no usable CUDA driver context ({error})");
return;
}
};
let (width, height) = (64u32, 64u32);
let Some(destination) = cuda_surface_with(&device, width, height, |_, y| (y * 3) as u8, 90)
else {
return;
};
let Some(source) = cuda_surface_with(&device, width, height, |x, _| (x * 4) as u8, 200)
else {
return;
};
let (MediaBuffer::Video(dst_frame), MediaBuffer::Video(src_frame)) =
(&destination, &source)
else {
panic!("expected Video buffers");
};
let dst_surface = Nv12Surface::from_frame(dst_frame).expect("destination planes");
let src_surface = Nv12Surface::from_frame(src_frame).expect("source planes");
let alpha = 77u8;
driver
.blend_nv12(
src_surface,
dst_surface,
Nv12Region {
source_x: 0,
source_y: 0,
destination_x: 0,
destination_y: 0,
width,
height,
},
alpha,
)
.expect("blend");
driver.synchronize().expect("synchronize");
let out = download(&device, destination.clone(), width, height);
let blend = |dst: u32, src: u32| {
((src * u32::from(alpha) + dst * (255 - u32::from(alpha)) + 127) / 255) as u8
};
let stride = out.stride(0);
for y in 0..height {
for x in 0..width {
let expected = blend(u32::from((y * 3) as u8), u32::from((x * 4) as u8));
let actual = out.data(0)[y as usize * stride + x as usize];
assert_eq!(
actual, expected,
"luma mismatch at ({x}, {y}): kernel {actual} != cpu {expected}"
);
}
}
let uv_stride = out.stride(1);
let expected_chroma = blend(90, 200);
for y in 0..height / 2 {
for x in 0..width {
let actual = out.data(1)[y as usize * uv_stride + x as usize];
assert_eq!(
actual, expected_chroma,
"chroma mismatch at ({x}, {y}): kernel {actual} != cpu {expected_chroma}"
);
}
}
}
#[test]
fn alpha_endpoints_replace_and_preserve_exactly() {
let Some((device, _cuda_lock)) = try_cuda_device() else {
return;
};
let Ok(driver) = CudaDriver::retain_primary() else {
eprintln!("skipping: no usable CUDA driver context");
return;
};
let (width, height) = (32u32, 32u32);
for (alpha, expected) in [(255u8, 200u8), (0, 60)] {
let Some(destination) = cuda_surface_with(&device, width, height, |_, _| 60, 128)
else {
return;
};
let Some(source) = cuda_surface_with(&device, width, height, |_, _| 200, 128) else {
return;
};
let (MediaBuffer::Video(dst_frame), MediaBuffer::Video(src_frame)) =
(&destination, &source)
else {
panic!("expected Video buffers");
};
driver
.blend_nv12(
Nv12Surface::from_frame(src_frame).expect("source planes"),
Nv12Surface::from_frame(dst_frame).expect("destination planes"),
Nv12Region {
source_x: 0,
source_y: 0,
destination_x: 0,
destination_y: 0,
width,
height,
},
alpha,
)
.expect("blend");
driver.synchronize().expect("synchronize");
let out = download(&device, destination.clone(), width, height);
assert_eq!(
out.data(0)[out.stride(0) * 5 + 5],
expected,
"alpha {alpha} must produce {expected}"
);
}
}
#[test]
fn black_and_white_map_to_limited_range_endpoints() {
assert_eq!(rgb_to_bt709_limited(Color::BLACK), (16, 128, 128));
let (y, u, v) = rgb_to_bt709_limited(Color::WHITE);
assert_eq!(y, 235);
assert!(
u.abs_diff(128) <= 1 && v.abs_diff(128) <= 1,
"white must be chroma-neutral, got ({u}, {v})"
);
}
}