feat(encode/nvenc): SPIR-V cursor blend over Vulkan-allocated input slots — retire the PTX kernels
A vendored PTX blob is JIT'd against the driver's ISA ceiling, so the cursor-blend module silently dies on drivers older than the generating toolkit (CUDA_ERROR_UNSUPPORTED_PTX_VERSION/INVALID_PTX, 222/218 — the KWin leg's invisible composite cursor on driver 595/CUDA 13.2 vs a CUDA 13.3 blob). SPIR-V has no such coupling, and the in-tree precedent already exists twice (vulkan_video's CSC blend, VkBridge's exportable OPAQUE_FD → cuImportExternalMemory bridge). New pf_zerocopy::vkslot::VkSlotBlend: the direct-SDK NVENC encoder now allocates its input ring as exportable Vulkan buffers CUDA-imports (same contiguous InputSurface layouts, pitch = row bytes rounded to 256), and the cursor composite is a compute dispatch over the cursor's rectangle (cursor_blend.comp, vendored .spv; spec-constant selects ARGB/NV12/YUV444; BT.709 limited, matching the retired .cu). The surface SSBO is uint[] with every invocation owning whole words — no 8-bit-storage device dependency. Cursor-bearing frames force the existing CPU-synced submit path so the CUDA copy → Vulkan dispatch (fence-waited) → NVENC encode ordering is CPU-established; cursorless frames keep the stream-ordered fast path untouched. Any bring-up/alloc/registration failure falls back wholesale to plain pitched CUDA surfaces (never a mixed or short ring): sessions always encode, composite mode just loses the cursor, warned once. cursor_blend.cu / cursor_blend.ptx and the CursorBlend PTX loader are deleted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Claude Fable 5
@@ -7,8 +7,12 @@
|
||||
//! and ffmpeg's `hevc_nvenc` (encode thread) — each thread makes it current before use;
|
||||
//! * device memory: pitched allocations, the reusable `BufferPool`/`DeviceBuffer`, IPC
|
||||
//! export/import, host readback, and the plane copies;
|
||||
//! * GL / external-memory interop (`RegisteredTexture`, `ExternalDmabuf`); and
|
||||
//! * the CUDA cursor-blend kernel (`CursorBlend`).
|
||||
//! * GL / external-memory interop (`RegisteredTexture`, `ExternalDmabuf`).
|
||||
//!
|
||||
//! (The CUDA cursor-blend PTX kernel that used to live here is retired: vendored PTX is JIT'd
|
||||
//! against the driver's ISA ceiling and silently dies on older drivers. The NVENC cursor blend
|
||||
//! is now the SPIR-V compute pass in [`super::vkslot`], dispatched over Vulkan-allocated,
|
||||
//! CUDA-imported input slots.)
|
||||
//!
|
||||
//! (We use GL interop, not EGL interop: `cuGraphicsEGLRegisterImage` is Tegra-only on the desktop
|
||||
//! driver — see [`super::egl`].)
|
||||
@@ -18,7 +22,6 @@
|
||||
#![deny(clippy::undocumented_unsafe_blocks)]
|
||||
|
||||
use anyhow::{bail, Result};
|
||||
use std::ffi::CStr;
|
||||
use std::os::raw::{c_uint, c_void};
|
||||
use std::sync::{Arc, Mutex, OnceLock};
|
||||
|
||||
@@ -280,258 +283,6 @@ pub fn copy_stream_handle() -> *mut c_void {
|
||||
/// Max cursor-overlay bitmap edge (px) uploaded to the device blend buffer — matches the Vulkan path.
|
||||
pub const CURSOR_MAX: u32 = 256;
|
||||
|
||||
/// GPU cursor-overlay compositor for the NVENC path (cursor-as-metadata): loads the `cursor_blend`
|
||||
/// PTX module once and blends a straight-alpha RGBA cursor into an encoder-OWNED NVENC input surface
|
||||
/// (ARGB / NV12 / YUV444) with a small kernel launched over the cursor's rectangle — no full-frame
|
||||
/// pass, and the compositor's dmabuf is never touched. The cursor bitmap lives in a device buffer
|
||||
/// re-uploaded only when it changes. Requires `context()` to have succeeded (driver present).
|
||||
pub struct CursorBlend {
|
||||
module: CUmodule,
|
||||
f_argb: CUfunction,
|
||||
f_nv12: CUfunction,
|
||||
f_yuv444: CUfunction,
|
||||
cur_buf: CUdeviceptr, // device RGBA staging (CURSOR_MAX²·4, tight rows)
|
||||
}
|
||||
|
||||
// SAFETY: process-lifetime driver handles used only from the encode thread with the shared context
|
||||
// current — like [`DeviceBuffer`], moving the struct between threads cannot dangle or race.
|
||||
unsafe impl Send for CursorBlend {}
|
||||
|
||||
impl CursorBlend {
|
||||
/// Load the embedded PTX image and resolve the three blend kernels + a device cursor buffer.
|
||||
pub fn new(ptx: &[u8]) -> Result<CursorBlend> {
|
||||
// cuModuleLoadData reads a PTX image as a NUL-terminated string; the embedded .ptx is not,
|
||||
// so append a terminator.
|
||||
let mut image = ptx.to_vec();
|
||||
image.push(0);
|
||||
let mut module: CUmodule = std::ptr::null_mut();
|
||||
// SAFETY: `&mut module` is a live out-param the driver fills; `image` is a NUL-terminated PTX
|
||||
// byte image that outlives the synchronous load. `ck` bails on error before `module` is used.
|
||||
unsafe {
|
||||
ck(
|
||||
cuModuleLoadData(&mut module, image.as_ptr() as *const c_void),
|
||||
"cuModuleLoadData(cursor_blend)",
|
||||
)?;
|
||||
}
|
||||
let getf = |name: &CStr| -> Result<CUfunction> {
|
||||
let mut f: CUfunction = std::ptr::null_mut();
|
||||
// SAFETY: `module` loaded above; each name is a valid NUL-terminated symbol present in
|
||||
// the module (verified in the .ptx `.entry` list); `&mut f` is a live out-param.
|
||||
unsafe {
|
||||
ck(
|
||||
cuModuleGetFunction(&mut f, module, name.as_ptr()),
|
||||
"cuModuleGetFunction",
|
||||
)?;
|
||||
}
|
||||
Ok(f)
|
||||
};
|
||||
let f_argb = getf(c"blend_argb")?;
|
||||
let f_nv12 = getf(c"blend_nv12")?;
|
||||
let f_yuv444 = getf(c"blend_yuv444")?;
|
||||
let mut cur_buf: CUdeviceptr = 0;
|
||||
// SAFETY: `&mut cur_buf` is a live out-param; the size fits the CURSOR_MAX² RGBA buffer.
|
||||
unsafe {
|
||||
ck(
|
||||
cuMemAlloc_v2(&mut cur_buf, (CURSOR_MAX * CURSOR_MAX * 4) as usize),
|
||||
"cuMemAlloc(cursor)",
|
||||
)?;
|
||||
}
|
||||
Ok(CursorBlend {
|
||||
module,
|
||||
f_argb,
|
||||
f_nv12,
|
||||
f_yuv444,
|
||||
cur_buf,
|
||||
})
|
||||
}
|
||||
|
||||
/// Upload the cursor RGBA (`cw*ch*4`, tight rows) into the device blend buffer. Call only when
|
||||
/// the bitmap changes; position moves are just kernel args.
|
||||
pub fn upload(&self, rgba: &[u8], cw: u32, ch: u32) -> Result<()> {
|
||||
let cw = cw.min(CURSOR_MAX);
|
||||
let ch = ch.min(CURSOR_MAX);
|
||||
let row = cw as usize * 4;
|
||||
let copy = CUDA_MEMCPY2D {
|
||||
srcMemoryType: 1, // HOST
|
||||
srcHost: rgba.as_ptr() as *const c_void,
|
||||
srcPitch: row,
|
||||
dstMemoryType: CU_MEMORYTYPE_DEVICE,
|
||||
dstDevice: self.cur_buf,
|
||||
dstPitch: row,
|
||||
WidthInBytes: row,
|
||||
Height: ch as usize,
|
||||
..Default::default()
|
||||
};
|
||||
// SAFETY: HOST→DEVICE 2D copy of `row*ch` bytes; `rgba` covers at least that (caller passes
|
||||
// `cw*ch*4`), `cur_buf` is the CURSOR_MAX²·4 device alloc (row ≤ CURSOR_MAX·4, ch ≤ CURSOR_MAX).
|
||||
// Synchronous via `copy_blocking`. Requires the context current (caller's contract).
|
||||
unsafe { copy_blocking(©, "cursor HtoD") }
|
||||
}
|
||||
|
||||
/// Blend into a packed 4-byte (NVENC ARGB) owned surface at `(ox,oy)`.
|
||||
#[allow(clippy::too_many_arguments)] // surface geometry + cursor size + offset — a struct would just be unpacked at the call
|
||||
pub fn blend_argb(
|
||||
&self,
|
||||
surf: CUdeviceptr,
|
||||
pitch: usize,
|
||||
w: u32,
|
||||
h: u32,
|
||||
cw: u32,
|
||||
ch: u32,
|
||||
ox: i32,
|
||||
oy: i32,
|
||||
sync: bool,
|
||||
) -> Result<()> {
|
||||
let (mut a_surf, mut a_cur) = (surf, self.cur_buf);
|
||||
let (mut a_pitch, mut a_w, mut a_h) = (pitch as i32, w as i32, h as i32);
|
||||
let (mut a_cw, mut a_ch) = (cw.min(CURSOR_MAX) as i32, ch.min(CURSOR_MAX) as i32);
|
||||
let (mut a_ox, mut a_oy) = (ox, oy);
|
||||
let mut args: [*mut c_void; 9] = [
|
||||
&mut a_surf as *mut _ as *mut c_void,
|
||||
&mut a_pitch as *mut _ as *mut c_void,
|
||||
&mut a_w as *mut _ as *mut c_void,
|
||||
&mut a_h as *mut _ as *mut c_void,
|
||||
&mut a_cur as *mut _ as *mut c_void,
|
||||
&mut a_cw as *mut _ as *mut c_void,
|
||||
&mut a_ch as *mut _ as *mut c_void,
|
||||
&mut a_ox as *mut _ as *mut c_void,
|
||||
&mut a_oy as *mut _ as *mut c_void,
|
||||
];
|
||||
self.launch(self.f_argb, a_cw as u32, a_ch as u32, &mut args, sync)
|
||||
}
|
||||
|
||||
/// Blend into an owned planar YUV444 surface (3 stacked full-res planes) at `(ox,oy)`.
|
||||
#[allow(clippy::too_many_arguments)] // surface geometry + cursor size + offset — a struct would just be unpacked at the call
|
||||
pub fn blend_yuv444(
|
||||
&self,
|
||||
base: CUdeviceptr,
|
||||
pitch: usize,
|
||||
w: u32,
|
||||
h: u32,
|
||||
cw: u32,
|
||||
ch: u32,
|
||||
ox: i32,
|
||||
oy: i32,
|
||||
sync: bool,
|
||||
) -> Result<()> {
|
||||
let (mut a_base, mut a_cur) = (base, self.cur_buf);
|
||||
let (mut a_pitch, mut a_w, mut a_h) = (pitch as i32, w as i32, h as i32);
|
||||
let (mut a_cw, mut a_ch) = (cw.min(CURSOR_MAX) as i32, ch.min(CURSOR_MAX) as i32);
|
||||
let (mut a_ox, mut a_oy) = (ox, oy);
|
||||
let mut args: [*mut c_void; 9] = [
|
||||
&mut a_base as *mut _ as *mut c_void,
|
||||
&mut a_pitch as *mut _ as *mut c_void,
|
||||
&mut a_w as *mut _ as *mut c_void,
|
||||
&mut a_h as *mut _ as *mut c_void,
|
||||
&mut a_cur as *mut _ as *mut c_void,
|
||||
&mut a_cw as *mut _ as *mut c_void,
|
||||
&mut a_ch as *mut _ as *mut c_void,
|
||||
&mut a_ox as *mut _ as *mut c_void,
|
||||
&mut a_oy as *mut _ as *mut c_void,
|
||||
];
|
||||
self.launch(self.f_yuv444, a_cw as u32, a_ch as u32, &mut args, sync)
|
||||
}
|
||||
|
||||
/// Blend into an owned NV12 surface (Y plane at `base`, interleaved UV at `base + pitch*h`).
|
||||
#[allow(clippy::too_many_arguments)] // surface geometry + cursor size + offset — a struct would just be unpacked at the call
|
||||
pub fn blend_nv12(
|
||||
&self,
|
||||
base: CUdeviceptr,
|
||||
pitch: usize,
|
||||
w: u32,
|
||||
h: u32,
|
||||
cw: u32,
|
||||
ch: u32,
|
||||
ox: i32,
|
||||
oy: i32,
|
||||
sync: bool,
|
||||
) -> Result<()> {
|
||||
let (mut a_yb, mut a_uvb, mut a_cur) = (base, base + pitch as u64 * h as u64, self.cur_buf);
|
||||
let (mut a_yp, mut a_uvp) = (pitch as i32, pitch as i32);
|
||||
let (mut a_w, mut a_h) = (w as i32, h as i32);
|
||||
let (mut a_cw, mut a_ch) = (cw.min(CURSOR_MAX) as i32, ch.min(CURSOR_MAX) as i32);
|
||||
let (mut a_ox, mut a_oy) = (ox, oy);
|
||||
let mut args: [*mut c_void; 11] = [
|
||||
&mut a_yb as *mut _ as *mut c_void,
|
||||
&mut a_yp as *mut _ as *mut c_void,
|
||||
&mut a_uvb as *mut _ as *mut c_void,
|
||||
&mut a_uvp as *mut _ as *mut c_void,
|
||||
&mut a_w as *mut _ as *mut c_void,
|
||||
&mut a_h as *mut _ as *mut c_void,
|
||||
&mut a_cur as *mut _ as *mut c_void,
|
||||
&mut a_cw as *mut _ as *mut c_void,
|
||||
&mut a_ch as *mut _ as *mut c_void,
|
||||
&mut a_ox as *mut _ as *mut c_void,
|
||||
&mut a_oy as *mut _ as *mut c_void,
|
||||
];
|
||||
// One thread per 2x2 luma block → grid over ceil(cw/2) × ceil(ch/2).
|
||||
self.launch(
|
||||
self.f_nv12,
|
||||
(a_cw as u32).div_ceil(2),
|
||||
(a_ch as u32).div_ceil(2),
|
||||
&mut args,
|
||||
sync,
|
||||
)
|
||||
}
|
||||
|
||||
/// Launch `f` over a `work_w × work_h` grid (16×16 blocks) on the copy stream; `sync` waits
|
||||
/// for it, `!sync` leaves completion to the stream (stream-ordered consumers only — the
|
||||
/// kernel PARAMETERS are copied at launch time, so the arg locals need not outlive the call).
|
||||
fn launch(
|
||||
&self,
|
||||
f: CUfunction,
|
||||
work_w: u32,
|
||||
work_h: u32,
|
||||
args: &mut [*mut c_void],
|
||||
sync: bool,
|
||||
) -> Result<()> {
|
||||
if work_w == 0 || work_h == 0 {
|
||||
return Ok(());
|
||||
}
|
||||
const B: u32 = 16;
|
||||
let stream = copy_stream();
|
||||
// SAFETY: `f` is a resolved kernel from our loaded module; `args` holds pointers to live
|
||||
// locals whose types match the kernel's C parameters (per the call site above) — CUDA
|
||||
// copies the parameter values during `cuLaunchKernel` itself, so they need not outlive
|
||||
// the call. Grid/block dims are non-zero. Launched on the copy stream (ordered after the
|
||||
// input-surface copy issued on the same stream); `sync` waits, `!sync` leaves ordering to
|
||||
// the stream (the NVENC IO-stream binding). Requires the context current.
|
||||
unsafe {
|
||||
ck(
|
||||
cuLaunchKernel(
|
||||
f,
|
||||
work_w.div_ceil(B),
|
||||
work_h.div_ceil(B),
|
||||
1,
|
||||
B,
|
||||
B,
|
||||
1,
|
||||
0,
|
||||
stream,
|
||||
args.as_mut_ptr(),
|
||||
std::ptr::null_mut(),
|
||||
),
|
||||
"cuLaunchKernel(cursor)",
|
||||
)?;
|
||||
if sync {
|
||||
ck(cuStreamSynchronize(stream), "cuStreamSynchronize(cursor)")?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for CursorBlend {
|
||||
fn drop(&mut self) {
|
||||
// SAFETY: `cur_buf`/`module` are our own handles, freed exactly once here; the context is
|
||||
// current on the encode thread that drops the encoder. Errors are ignored on teardown.
|
||||
unsafe {
|
||||
let _ = cuMemFree_v2(self.cur_buf);
|
||||
let _ = cuModuleUnload(self.module);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Allocate one pitched device buffer for `width`x`height` 4-byte pixels; returns `(ptr, pitch)`.
|
||||
fn alloc_pitched(width: u32, height: u32) -> Result<(CUdeviceptr, usize)> {
|
||||
let mut ptr: CUdeviceptr = 0;
|
||||
|
||||
@@ -9,7 +9,7 @@
|
||||
#![deny(clippy::undocumented_unsafe_blocks)]
|
||||
|
||||
use anyhow::{bail, Result};
|
||||
use std::os::raw::{c_char, c_int, c_uint, c_void};
|
||||
use std::os::raw::{c_int, c_uint, c_void};
|
||||
use std::sync::OnceLock;
|
||||
|
||||
pub type CUresult = c_uint; // CUDA_SUCCESS == 0
|
||||
@@ -20,8 +20,6 @@ pub type CUdeviceptr = u64;
|
||||
pub type CUgraphicsResource = *mut c_void;
|
||||
pub type CUarray = *mut c_void;
|
||||
pub type CUexternalMemory = *mut c_void; // opaque CUextMemory_st*
|
||||
pub type CUmodule = *mut c_void; // opaque CUmod_st*
|
||||
pub type CUfunction = *mut c_void; // opaque CUfunc_st*
|
||||
|
||||
/// `CUmemorytype` (cuda.h): HOST=1, DEVICE=2, ARRAY=3, UNIFIED=4.
|
||||
pub const CU_MEMORYTYPE_DEVICE: c_uint = 2;
|
||||
@@ -144,26 +142,6 @@ pub(crate) struct CudaApi {
|
||||
cuIpcGetMemHandle: unsafe extern "C" fn(*mut CUipcMemHandle, CUdeviceptr) -> CUresult,
|
||||
cuIpcOpenMemHandle: unsafe extern "C" fn(*mut CUdeviceptr, CUipcMemHandle, c_uint) -> CUresult,
|
||||
cuIpcCloseMemHandle: unsafe extern "C" fn(CUdeviceptr) -> CUresult,
|
||||
// Cursor-overlay blend: a linear device alloc + a PTX module with the blend kernels launched
|
||||
// over the cursor's small rectangle (see [`CursorBlend`]).
|
||||
cuMemAlloc_v2: unsafe extern "C" fn(*mut CUdeviceptr, usize) -> CUresult,
|
||||
cuModuleLoadData: unsafe extern "C" fn(*mut CUmodule, *const c_void) -> CUresult,
|
||||
cuModuleUnload: unsafe extern "C" fn(CUmodule) -> CUresult,
|
||||
cuModuleGetFunction: unsafe extern "C" fn(*mut CUfunction, CUmodule, *const c_char) -> CUresult,
|
||||
#[allow(clippy::type_complexity)]
|
||||
cuLaunchKernel: unsafe extern "C" fn(
|
||||
CUfunction,
|
||||
c_uint,
|
||||
c_uint,
|
||||
c_uint,
|
||||
c_uint,
|
||||
c_uint,
|
||||
c_uint,
|
||||
c_uint,
|
||||
CUstream,
|
||||
*mut *mut c_void,
|
||||
*mut *mut c_void,
|
||||
) -> CUresult,
|
||||
}
|
||||
// SAFETY: every field is a bare `extern "C" fn` address into the leaked, process-lifetime
|
||||
// `libcuda` mapping (`cuda_api` `forget`s the `Library`, so it is never unloaded) — an immutable
|
||||
@@ -236,11 +214,6 @@ pub(crate) fn cuda_api() -> Option<&'static CudaApi> {
|
||||
.or_else(|_| lib.get(b"cuIpcOpenMemHandle\0"))
|
||||
.ok()?,
|
||||
cuIpcCloseMemHandle: *lib.get(b"cuIpcCloseMemHandle\0").ok()?,
|
||||
cuMemAlloc_v2: *lib.get(b"cuMemAlloc_v2\0").ok()?,
|
||||
cuModuleLoadData: *lib.get(b"cuModuleLoadData\0").ok()?,
|
||||
cuModuleUnload: *lib.get(b"cuModuleUnload\0").ok()?,
|
||||
cuModuleGetFunction: *lib.get(b"cuModuleGetFunction\0").ok()?,
|
||||
cuLaunchKernel: *lib.get(b"cuLaunchKernel\0").ok()?,
|
||||
};
|
||||
std::mem::forget(lib); // keep libcuda mapped for the fn pointers' lifetime (process)
|
||||
Some(api)
|
||||
@@ -304,53 +277,6 @@ pub(crate) unsafe fn cuMemFree_v2(dptr: CUdeviceptr) -> CUresult {
|
||||
None => CU_ERROR_NOT_LOADED,
|
||||
}
|
||||
}
|
||||
pub(crate) unsafe fn cuMemAlloc_v2(dptr: *mut CUdeviceptr, size: usize) -> CUresult {
|
||||
match cuda_api() {
|
||||
Some(a) => (a.cuMemAlloc_v2)(dptr, size),
|
||||
None => CU_ERROR_NOT_LOADED,
|
||||
}
|
||||
}
|
||||
pub(crate) unsafe fn cuModuleLoadData(m: *mut CUmodule, image: *const c_void) -> CUresult {
|
||||
match cuda_api() {
|
||||
Some(a) => (a.cuModuleLoadData)(m, image),
|
||||
None => CU_ERROR_NOT_LOADED,
|
||||
}
|
||||
}
|
||||
pub(crate) unsafe fn cuModuleUnload(m: CUmodule) -> CUresult {
|
||||
match cuda_api() {
|
||||
Some(a) => (a.cuModuleUnload)(m),
|
||||
None => CU_ERROR_NOT_LOADED,
|
||||
}
|
||||
}
|
||||
pub(crate) unsafe fn cuModuleGetFunction(
|
||||
f: *mut CUfunction,
|
||||
m: CUmodule,
|
||||
name: *const c_char,
|
||||
) -> CUresult {
|
||||
match cuda_api() {
|
||||
Some(a) => (a.cuModuleGetFunction)(f, m, name),
|
||||
None => CU_ERROR_NOT_LOADED,
|
||||
}
|
||||
}
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) unsafe fn cuLaunchKernel(
|
||||
f: CUfunction,
|
||||
gx: c_uint,
|
||||
gy: c_uint,
|
||||
gz: c_uint,
|
||||
bx: c_uint,
|
||||
by: c_uint,
|
||||
bz: c_uint,
|
||||
shmem: c_uint,
|
||||
stream: CUstream,
|
||||
params: *mut *mut c_void,
|
||||
extra: *mut *mut c_void,
|
||||
) -> CUresult {
|
||||
match cuda_api() {
|
||||
Some(a) => (a.cuLaunchKernel)(f, gx, gy, gz, bx, by, bz, shmem, stream, params, extra),
|
||||
None => CU_ERROR_NOT_LOADED,
|
||||
}
|
||||
}
|
||||
pub(crate) unsafe fn cuMemcpy2DAsync_v2(copy: *const CUDA_MEMCPY2D, stream: CUstream) -> CUresult {
|
||||
match cuda_api() {
|
||||
Some(a) => (a.cuMemcpy2DAsync_v2)(copy, stream),
|
||||
|
||||
@@ -0,0 +1,170 @@
|
||||
#version 450
|
||||
// Cursor-overlay blend for the direct-SDK NVENC path (cursor-as-metadata), dispatched over the
|
||||
// cursor's rectangle only — the Vulkan replacement for the retired cursor_blend.cu PTX kernels
|
||||
// (PTX is JIT'd against the driver's ISA ceiling, so a vendored blob silently dies on older
|
||||
// drivers: CUDA errors 222/218 on-glass; SPIR-V has no such coupling). The NVENC input surface is
|
||||
// Vulkan-allocated, CUDA-imported external memory (see vkslot.rs), so this shader writes the very
|
||||
// bytes NVENC encodes.
|
||||
//
|
||||
// MODE (spec constant): 0 = packed 4-byte ARGB (NVENC byte order B,G,R,A), 1 = NV12 (Y plane +
|
||||
// interleaved half-res UV at row surfH), 2 = planar YUV444 (3 full-res planes stacked at
|
||||
// pitch*surfH). BT.709 limited-range coefficients — identical to rgb2nv12_buf.comp and the
|
||||
// retired .cu, so the cursor colour matches the frame regardless of backend.
|
||||
//
|
||||
// The surface SSBO is uint[] (no 8-bit storage dependency — maximum driver reach): every
|
||||
// invocation exclusively owns the 32-bit words it read-modify-writes. ARGB: one invocation per
|
||||
// cursor pixel = one word. NV12/YUV444: one invocation per WORD-ALIGNED 4-px luma span (per two
|
||||
// rows for NV12, whose 2 chroma bytes-pairs land in one exclusive word). Spans are aligned to the
|
||||
// SURFACE, not the cursor, so neighbouring invocations never share a word even at odd `ox`.
|
||||
//
|
||||
// Rebuild: glslc cursor_blend.comp -o cursor_blend.spv (vendored beside this file)
|
||||
|
||||
layout(local_size_x = 8, local_size_y = 8, local_size_z = 1) in;
|
||||
|
||||
layout(constant_id = 0) const uint MODE = 0;
|
||||
|
||||
layout(std430, binding = 0) buffer Surf { uint surf[]; };
|
||||
layout(std430, binding = 1) readonly buffer Cur { uint cur[]; };
|
||||
|
||||
layout(push_constant) uniform Push {
|
||||
uint pitch; // surface row stride, bytes (4-aligned by construction)
|
||||
uint surfW; // content width, px
|
||||
uint surfH; // luma rows (plane stride multiplier)
|
||||
uint curW; // cursor bitmap width, px
|
||||
uint curH; // cursor bitmap height, px
|
||||
int ox; // cursor top-left on the surface, px (may be negative)
|
||||
int oy;
|
||||
} pc;
|
||||
|
||||
// Cursor texel (straight-alpha RGBA, tight rows) or (0,0,0,0) outside the bitmap.
|
||||
uvec4 cursor_px(int cx, int cy) {
|
||||
if (cx < 0 || cy < 0 || cx >= int(pc.curW) || cy >= int(pc.curH)) return uvec4(0);
|
||||
uint w = cur[uint(cy) * pc.curW + uint(cx)];
|
||||
return uvec4(w & 0xFFu, (w >> 8) & 0xFFu, (w >> 16) & 0xFFu, (w >> 24) & 0xFFu); // R,G,B,A
|
||||
}
|
||||
|
||||
uint blend8(uint dst, uint src, uint a) {
|
||||
return (src * a + dst * (255u - a)) / 255u;
|
||||
}
|
||||
|
||||
// BT.709 limited RGB→Y/U/V (matches the retired .cu / rgb2nv12_buf.comp).
|
||||
uint y_of(uvec4 s) {
|
||||
return uint(clamp(16.0 + 0.1826 * float(s.r) + 0.6142 * float(s.g) + 0.0620 * float(s.b) + 0.5, 0.0, 255.0));
|
||||
}
|
||||
float u_of(uvec4 s) { return 128.0 - 0.1006 * float(s.r) - 0.3386 * float(s.g) + 0.4392 * float(s.b); }
|
||||
float v_of(uvec4 s) { return 128.0 + 0.4392 * float(s.r) - 0.3989 * float(s.g) - 0.0403 * float(s.b); }
|
||||
|
||||
// Read-modify-write one byte lane of a word index.
|
||||
void rmw_byte(uint word_idx, uint lane, uint val8, uint a) {
|
||||
uint w = surf[word_idx];
|
||||
uint shift = lane * 8u;
|
||||
uint d = (w >> shift) & 0xFFu;
|
||||
uint b = blend8(d, val8, a);
|
||||
surf[word_idx] = (w & ~(0xFFu << shift)) | (b << shift);
|
||||
}
|
||||
|
||||
void main() {
|
||||
if (MODE == 0u) {
|
||||
// ARGB: one invocation per cursor pixel; each surface pixel is one exclusive word.
|
||||
int cx = int(gl_GlobalInvocationID.x);
|
||||
int cy = int(gl_GlobalInvocationID.y);
|
||||
if (cx >= int(pc.curW) || cy >= int(pc.curH)) return;
|
||||
int px = pc.ox + cx, py = pc.oy + cy;
|
||||
if (px < 0 || py < 0 || px >= int(pc.surfW) || py >= int(pc.surfH)) return;
|
||||
uvec4 s = cursor_px(cx, cy);
|
||||
if (s.a == 0u) return;
|
||||
uint idx = (uint(py) * pc.pitch + uint(px) * 4u) / 4u;
|
||||
uint w = surf[idx];
|
||||
uint b = blend8(w & 0xFFu, s.b, s.a); // B lane
|
||||
uint g = blend8((w >> 8) & 0xFFu, s.g, s.a); // G lane
|
||||
uint r = blend8((w >> 16) & 0xFFu, s.r, s.a); // R lane
|
||||
surf[idx] = (w & 0xFF000000u) | (r << 16) | (g << 8) | b;
|
||||
return;
|
||||
}
|
||||
|
||||
// NV12 / YUV444: one invocation per SURFACE-word-aligned 4-px luma span. Span origin:
|
||||
// x0 = floor(ox/4)*4 + span*4 (surface px), rows walk the cursor rect.
|
||||
int span = int(gl_GlobalInvocationID.x);
|
||||
int row = int(gl_GlobalInvocationID.y);
|
||||
int x0 = (pc.ox >> 2) << 2; // word-aligned start at/left-of ox (ox may be negative)
|
||||
int px0 = x0 + span * 4;
|
||||
|
||||
if (MODE == 2u) {
|
||||
// YUV444: rows walk cursor rows one at a time.
|
||||
if (row >= int(pc.curH)) return;
|
||||
int py = pc.oy + row;
|
||||
if (py < 0 || py >= int(pc.surfH)) return;
|
||||
uint plane = pc.pitch * pc.surfH;
|
||||
for (int i = 0; i < 4; i++) {
|
||||
int px = px0 + i;
|
||||
int cx = px - pc.ox;
|
||||
if (px < 0 || px >= int(pc.surfW)) continue;
|
||||
uvec4 s = cursor_px(cx, row);
|
||||
if (s.a == 0u) continue;
|
||||
uint off = uint(py) * pc.pitch + uint(px);
|
||||
uint U = uint(clamp(u_of(s) + 0.5, 0.0, 255.0));
|
||||
uint V = uint(clamp(v_of(s) + 0.5, 0.0, 255.0));
|
||||
rmw_byte(off / 4u, off % 4u, y_of(s), s.a);
|
||||
rmw_byte((plane + off) / 4u, (plane + off) % 4u, U, s.a);
|
||||
rmw_byte((2u * plane + off) / 4u, (2u * plane + off) % 4u, V, s.a);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
// NV12: rows walk 2-row luma blocks (row = block row). The span's 4 luma px × 2 rows are
|
||||
// exclusive words; its 2 chroma samples (4 bytes) are one exclusive word.
|
||||
int base_cy = row * 2;
|
||||
if (base_cy >= int(pc.curH)) return;
|
||||
// Luma: 4 px × 2 rows.
|
||||
for (int j = 0; j < 2; j++) {
|
||||
int cy = base_cy + j;
|
||||
int py = pc.oy + cy;
|
||||
if (cy >= int(pc.curH) || py < 0 || py >= int(pc.surfH)) continue;
|
||||
for (int i = 0; i < 4; i++) {
|
||||
int px = px0 + i;
|
||||
int cx = px - pc.ox;
|
||||
if (px < 0 || px >= int(pc.surfW)) continue;
|
||||
uvec4 s = cursor_px(cx, cy);
|
||||
if (s.a == 0u) continue;
|
||||
uint off = uint(py) * pc.pitch + uint(px);
|
||||
rmw_byte(off / 4u, off % 4u, y_of(s), s.a);
|
||||
}
|
||||
}
|
||||
// Chroma: two UV samples covering the span's 2x2 blocks, alpha-weighted like the .cu kernel.
|
||||
// The UV plane starts at row surfH; sample (uvx, uvy) lives at uv_base + uvy*pitch + uvx*2.
|
||||
// Guard: only spans whose px0 is 4-aligned own their chroma word (px0 is by construction).
|
||||
int py_top = pc.oy + base_cy;
|
||||
int uvy = py_top >> 1;
|
||||
if (py_top < 0 || uvy < 0 || uvy * 2 >= int(pc.surfH)) return;
|
||||
uint uv_base = pc.pitch * pc.surfH;
|
||||
for (int hf = 0; hf < 2; hf++) {
|
||||
// Each hf = one 2x2 luma block = one UV sample (2 bytes).
|
||||
int bx = px0 + hf * 2;
|
||||
if (bx < 0 || bx >= int(pc.surfW)) continue;
|
||||
int uvx = bx >> 1;
|
||||
float ua = 0.0, va = 0.0, wa = 0.0;
|
||||
int cnt = 0;
|
||||
for (int j = 0; j < 2; j++) {
|
||||
for (int i = 0; i < 2; i++) {
|
||||
int px = bx + i;
|
||||
int py = py_top + j;
|
||||
int cx = px - pc.ox;
|
||||
int cy = base_cy + j;
|
||||
if (px < 0 || py < 0 || px >= int(pc.surfW) || py >= int(pc.surfH)) continue;
|
||||
uvec4 s = cursor_px(cx, cy);
|
||||
if (s.a == 0u) continue;
|
||||
ua += u_of(s) * float(s.a);
|
||||
va += v_of(s) * float(s.a);
|
||||
wa += float(s.a);
|
||||
cnt++;
|
||||
}
|
||||
}
|
||||
if (wa <= 0.0 || cnt == 0) continue;
|
||||
uint U = uint(clamp(ua / wa + 0.5, 0.0, 255.0));
|
||||
uint V = uint(clamp(va / wa + 0.5, 0.0, 255.0));
|
||||
uint amean = uint(clamp(wa / float(cnt) + 0.5, 0.0, 255.0));
|
||||
uint off = uv_base + uint(uvy) * pc.pitch + uint(uvx) * 2u;
|
||||
rmw_byte(off / 4u, off % 4u, U, amean);
|
||||
rmw_byte((off + 1u) / 4u, (off + 1u) % 4u, V, amean);
|
||||
}
|
||||
}
|
||||
Binary file not shown.
@@ -15,6 +15,7 @@ pub mod client;
|
||||
pub mod cuda;
|
||||
pub mod egl;
|
||||
pub mod proto;
|
||||
pub mod vkslot;
|
||||
pub mod vulkan;
|
||||
pub mod worker;
|
||||
|
||||
|
||||
@@ -0,0 +1,710 @@
|
||||
//! Vulkan-allocated NVENC input slots + the Vulkan compute cursor blend — the driver-portable
|
||||
//! replacement for the retired `cursor_blend.cu` PTX kernels (design: remote-desktop-sweep §8,
|
||||
//! Phase A). A vendored PTX blob is JIT'd against the driver's ISA ceiling, so it silently dies
|
||||
//! on drivers older than the generating toolkit (CUDA errors 222/218 on-glass — the KWin leg's
|
||||
//! invisible composite cursor); SPIR-V has no such coupling.
|
||||
//!
|
||||
//! ```text
|
||||
//! exportable VkBuffer ──vkGetMemoryFdKHR(OPAQUE_FD)──▶ cuImportExternalMemory ──▶ CUdeviceptr
|
||||
//! ▲ │ NVENC registers + encodes
|
||||
//! └── cursor_blend.comp dispatch (cursor rect only) ◀───┘ CUDA copies frames in
|
||||
//! ```
|
||||
//!
|
||||
//! The direct-SDK NVENC encoder allocates its input ring through [`VkSlotBlend::alloc_slot`]
|
||||
//! instead of `cuMemAllocPitch`: same contiguous layouts (`InputSurface` docs), but the memory is
|
||||
//! Vulkan external memory both APIs address. Per cursor-bearing frame the encoder CPU-syncs its
|
||||
//! CUDA copy, then [`VkSlotBlend::blend`] dispatches the compute blend over the cursor's
|
||||
//! rectangle and fence-waits — the same coherence ceremony [`super::vulkan::VkBridge`] ships for
|
||||
//! its CSC (fence-ordered cross-API access on NVIDIA, no queue-family transfer needed). Frames
|
||||
//! without a cursor never touch Vulkan, keeping the stream-ordered fast path intact.
|
||||
//!
|
||||
//! Falls back cleanly: if bring-up fails the encoder allocates plain CUDA surfaces and composite
|
||||
//! mode degrades to no cursor (warned once) — never a failed session.
|
||||
|
||||
// Every `unsafe` block in this file carries a `// SAFETY:` proof; enforce it (unsafe-proof program).
|
||||
#![deny(clippy::undocumented_unsafe_blocks)]
|
||||
|
||||
use super::cuda::{self, CUdeviceptr};
|
||||
use anyhow::{anyhow, Context as _, Result};
|
||||
use ash::vk;
|
||||
|
||||
/// Max cursor-overlay bitmap edge (px) — matches [`cuda::CURSOR_MAX`] and the capture-side clamp.
|
||||
pub const CURSOR_MAX: u32 = cuda::CURSOR_MAX;
|
||||
|
||||
/// The vendored SPIR-V for `cursor_blend.comp` (beside this file; rebuild with
|
||||
/// `glslc cursor_blend.comp -o cursor_blend.spv`).
|
||||
const CURSOR_SPV: &[u8] = include_bytes!("cursor_blend.spv");
|
||||
|
||||
/// NVENC input-surface layout — selects the spec-constant `MODE` pipeline and the allocation
|
||||
/// arithmetic (mirroring `InputSurface`'s contiguous layouts).
|
||||
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
|
||||
pub enum SlotFormat {
|
||||
/// Packed 4-byte ARGB (NVENC byte order B,G,R,A): `pitch × height`.
|
||||
Argb,
|
||||
/// NV12: Y rows `[0, H)` + interleaved UV rows `[H, 3H/2)` under one pitch.
|
||||
Nv12,
|
||||
/// Planar YUV444: three full-res planes stacked at `pitch × height` intervals.
|
||||
Yuv444,
|
||||
}
|
||||
|
||||
impl SlotFormat {
|
||||
fn mode(self) -> u32 {
|
||||
match self {
|
||||
SlotFormat::Argb => 0,
|
||||
SlotFormat::Nv12 => 1,
|
||||
SlotFormat::Yuv444 => 2,
|
||||
}
|
||||
}
|
||||
fn row_bytes(self, width: u32) -> u64 {
|
||||
match self {
|
||||
SlotFormat::Argb => width as u64 * 4,
|
||||
SlotFormat::Nv12 | SlotFormat::Yuv444 => width as u64,
|
||||
}
|
||||
}
|
||||
fn rows(self, height: u32) -> u64 {
|
||||
match self {
|
||||
SlotFormat::Argb => height as u64,
|
||||
SlotFormat::Nv12 => height as u64 + (height as u64 / 2).max(1),
|
||||
SlotFormat::Yuv444 => height as u64 * 3,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// What the encoder holds per ring slot: the CUDA view it registers with NVENC plus the id it
|
||||
/// hands back to [`VkSlotBlend::blend`]. The backing Vulkan objects + CUDA mapping live in the
|
||||
/// [`VkSlotBlend`] (freed by [`free_slots`](VkSlotBlend::free_slots) / drop), so this is Copy —
|
||||
/// the encoder's ring keeps its existing shape.
|
||||
#[derive(Clone, Copy)]
|
||||
pub struct VkSlotRef {
|
||||
/// Device pointer NVENC registers (CUDA's mapping of the Vulkan memory).
|
||||
pub ptr: CUdeviceptr,
|
||||
/// Row stride in bytes (ours: row bytes rounded up to 256).
|
||||
pub pitch: usize,
|
||||
/// Luma rows (the plane-stride multiplier, as in `InputSurface`).
|
||||
pub height: u32,
|
||||
/// Index into the blend's slot table.
|
||||
pub id: usize,
|
||||
}
|
||||
|
||||
/// One allocated slot's backing objects, freed together in reverse order (CUDA mapping first).
|
||||
struct SlotAlloc {
|
||||
buffer: vk::Buffer,
|
||||
memory: vk::DeviceMemory,
|
||||
/// CUDA's import of the exported OPAQUE_FD — must drop BEFORE the Vulkan memory is freed.
|
||||
cuda: cuda::ExternalDmabuf,
|
||||
size: u64,
|
||||
}
|
||||
|
||||
/// 28-byte push-constant block matching `cursor_blend.comp`'s `Push`.
|
||||
#[repr(C)]
|
||||
struct Push {
|
||||
pitch: u32,
|
||||
surf_w: u32,
|
||||
surf_h: u32,
|
||||
cur_w: u32,
|
||||
cur_h: u32,
|
||||
ox: i32,
|
||||
oy: i32,
|
||||
}
|
||||
|
||||
pub struct VkSlotBlend {
|
||||
_entry: ash::Entry,
|
||||
instance: ash::Instance,
|
||||
device: ash::Device,
|
||||
ext_fd: ash::khr::external_memory_fd::Device,
|
||||
queue: vk::Queue,
|
||||
cmd_pool: vk::CommandPool,
|
||||
cmd: vk::CommandBuffer,
|
||||
fence: vk::Fence,
|
||||
mem_props: vk::PhysicalDeviceMemoryProperties,
|
||||
shader: vk::ShaderModule,
|
||||
desc_layout: vk::DescriptorSetLayout,
|
||||
pipe_layout: vk::PipelineLayout,
|
||||
desc_pool: vk::DescriptorPool,
|
||||
desc_set: vk::DescriptorSet,
|
||||
/// One pipeline per [`SlotFormat`], indexed by `mode()` (spec constant).
|
||||
pipelines: [vk::Pipeline; 3],
|
||||
/// Host-visible cursor bitmap staging (CURSOR_MAX²·4, tight rows), persistently mapped.
|
||||
cur_buf: vk::Buffer,
|
||||
cur_mem: vk::DeviceMemory,
|
||||
cur_map: *mut u8,
|
||||
slots: Vec<SlotAlloc>,
|
||||
}
|
||||
|
||||
// SAFETY: raw Vulkan handles + a persistently-mapped pointer, all uniquely owned by this struct
|
||||
// and destroyed exactly once in `Drop`; used from the encoder thread but moved with it. `Send`
|
||||
// only (not `Sync`), matching the single-thread use — transferring opaque handles cannot dangle.
|
||||
unsafe impl Send for VkSlotBlend {}
|
||||
|
||||
impl VkSlotBlend {
|
||||
/// Bring up the device + blend pipelines. Requires the CUDA shared context (the encoder's) to
|
||||
/// be established; picks the NVIDIA physical device (the NVENC path is NVIDIA by definition).
|
||||
pub fn new() -> Result<VkSlotBlend> {
|
||||
// SAFETY: standard ash bring-up, same shape as `VkBridge::new` — every call is `unsafe`
|
||||
// only because ash cannot statically verify handle/CreateInfo validity. Every
|
||||
// `*CreateInfo`/`AllocateInfo` is built from locals that live for the duration of the
|
||||
// synchronous call reading them; every handle passed was created and `?`-checked in this
|
||||
// same function. Shares nothing across threads.
|
||||
unsafe {
|
||||
let entry = ash::Entry::load().context("load libvulkan")?;
|
||||
let app = vk::ApplicationInfo::default().api_version(vk::API_VERSION_1_1);
|
||||
let instance = entry
|
||||
.create_instance(
|
||||
&vk::InstanceCreateInfo::default().application_info(&app),
|
||||
None,
|
||||
)
|
||||
.context("vkCreateInstance")?;
|
||||
let phys = match instance
|
||||
.enumerate_physical_devices()
|
||||
.context("enumerate GPUs")?
|
||||
.into_iter()
|
||||
.find(|&p| instance.get_physical_device_properties(p).vendor_id == 0x10DE)
|
||||
{
|
||||
Some(p) => p,
|
||||
None => {
|
||||
instance.destroy_instance(None);
|
||||
return Err(anyhow!("no NVIDIA Vulkan device"));
|
||||
}
|
||||
};
|
||||
let mem_props = instance.get_physical_device_memory_properties(phys);
|
||||
let qf = match instance
|
||||
.get_physical_device_queue_family_properties(phys)
|
||||
.iter()
|
||||
.position(|q| q.queue_flags.contains(vk::QueueFlags::COMPUTE))
|
||||
{
|
||||
Some(i) => i as u32,
|
||||
None => {
|
||||
instance.destroy_instance(None);
|
||||
return Err(anyhow!("no compute-capable queue family"));
|
||||
}
|
||||
};
|
||||
let prio = [1.0f32];
|
||||
let qci = [vk::DeviceQueueCreateInfo::default()
|
||||
.queue_family_index(qf)
|
||||
.queue_priorities(&prio)];
|
||||
let exts = [ash::khr::external_memory_fd::NAME.as_ptr()];
|
||||
let device = match instance.create_device(
|
||||
phys,
|
||||
&vk::DeviceCreateInfo::default()
|
||||
.queue_create_infos(&qci)
|
||||
.enabled_extension_names(&exts),
|
||||
None,
|
||||
) {
|
||||
Ok(d) => d,
|
||||
Err(e) => {
|
||||
instance.destroy_instance(None);
|
||||
return Err(e).context("vkCreateDevice (external_memory_fd supported?)");
|
||||
}
|
||||
};
|
||||
// From here teardown-on-error goes through `destroy_partial`, which tolerates null
|
||||
// handles — build everything into an incrementally-filled struct.
|
||||
let ext_fd = ash::khr::external_memory_fd::Device::new(&instance, &device);
|
||||
let queue = device.get_device_queue(qf, 0);
|
||||
let mut me = VkSlotBlend {
|
||||
_entry: entry,
|
||||
instance,
|
||||
device,
|
||||
ext_fd,
|
||||
queue,
|
||||
cmd_pool: vk::CommandPool::null(),
|
||||
cmd: vk::CommandBuffer::null(),
|
||||
fence: vk::Fence::null(),
|
||||
mem_props,
|
||||
shader: vk::ShaderModule::null(),
|
||||
desc_layout: vk::DescriptorSetLayout::null(),
|
||||
pipe_layout: vk::PipelineLayout::null(),
|
||||
desc_pool: vk::DescriptorPool::null(),
|
||||
desc_set: vk::DescriptorSet::null(),
|
||||
pipelines: [vk::Pipeline::null(); 3],
|
||||
cur_buf: vk::Buffer::null(),
|
||||
cur_mem: vk::DeviceMemory::null(),
|
||||
cur_map: std::ptr::null_mut(),
|
||||
slots: Vec::new(),
|
||||
};
|
||||
me.init_objects(qf).inspect_err(|_| {
|
||||
// `Drop` runs the same teardown and tolerates the nulls left by a partial init.
|
||||
})?;
|
||||
tracing::info!(
|
||||
"Vulkan slot blend ready (exportable NVENC inputs + SPIR-V cursor blend)"
|
||||
);
|
||||
Ok(me)
|
||||
}
|
||||
}
|
||||
|
||||
/// The non-device objects: command machinery, cursor staging, descriptor + pipelines.
|
||||
fn init_objects(&mut self, qf: u32) -> Result<()> {
|
||||
// SAFETY: same contract as `new` — ash calls on the live `self.device` with builder infos
|
||||
// from locals outliving each synchronous call; created handles are stored into `self`
|
||||
// immediately so the caller's `Drop` frees them on any later failure.
|
||||
unsafe {
|
||||
let d = &self.device;
|
||||
self.cmd_pool = d
|
||||
.create_command_pool(
|
||||
&vk::CommandPoolCreateInfo::default()
|
||||
.queue_family_index(qf)
|
||||
.flags(vk::CommandPoolCreateFlags::RESET_COMMAND_BUFFER),
|
||||
None,
|
||||
)
|
||||
.context("create command pool")?;
|
||||
self.cmd = d
|
||||
.allocate_command_buffers(
|
||||
&vk::CommandBufferAllocateInfo::default()
|
||||
.command_pool(self.cmd_pool)
|
||||
.level(vk::CommandBufferLevel::PRIMARY)
|
||||
.command_buffer_count(1),
|
||||
)
|
||||
.context("allocate command buffer")?[0];
|
||||
self.fence = d
|
||||
.create_fence(&vk::FenceCreateInfo::default(), None)
|
||||
.context("create fence")?;
|
||||
|
||||
// Cursor staging: host-visible+coherent SSBO, persistently mapped.
|
||||
let cur_size = (CURSOR_MAX * CURSOR_MAX * 4) as u64;
|
||||
self.cur_buf = d
|
||||
.create_buffer(
|
||||
&vk::BufferCreateInfo::default()
|
||||
.size(cur_size)
|
||||
.usage(vk::BufferUsageFlags::STORAGE_BUFFER),
|
||||
None,
|
||||
)
|
||||
.context("create cursor buffer")?;
|
||||
let reqs = d.get_buffer_memory_requirements(self.cur_buf);
|
||||
let mem_type = self
|
||||
.memory_type(
|
||||
reqs.memory_type_bits,
|
||||
vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT,
|
||||
)
|
||||
.context("cursor buffer memory type")?;
|
||||
self.cur_mem = d
|
||||
.allocate_memory(
|
||||
&vk::MemoryAllocateInfo::default()
|
||||
.allocation_size(reqs.size)
|
||||
.memory_type_index(mem_type),
|
||||
None,
|
||||
)
|
||||
.context("allocate cursor memory")?;
|
||||
d.bind_buffer_memory(self.cur_buf, self.cur_mem, 0)
|
||||
.context("bind cursor memory")?;
|
||||
self.cur_map = d
|
||||
.map_memory(self.cur_mem, 0, cur_size, vk::MemoryMapFlags::empty())
|
||||
.context("map cursor memory")? as *mut u8;
|
||||
|
||||
// Descriptor set: binding 0 = surface SSBO (rebound per blend), 1 = cursor SSBO.
|
||||
let bindings = [
|
||||
vk::DescriptorSetLayoutBinding::default()
|
||||
.binding(0)
|
||||
.descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
|
||||
.descriptor_count(1)
|
||||
.stage_flags(vk::ShaderStageFlags::COMPUTE),
|
||||
vk::DescriptorSetLayoutBinding::default()
|
||||
.binding(1)
|
||||
.descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
|
||||
.descriptor_count(1)
|
||||
.stage_flags(vk::ShaderStageFlags::COMPUTE),
|
||||
];
|
||||
self.desc_layout = d
|
||||
.create_descriptor_set_layout(
|
||||
&vk::DescriptorSetLayoutCreateInfo::default().bindings(&bindings),
|
||||
None,
|
||||
)
|
||||
.context("create descriptor layout")?;
|
||||
let pc = [vk::PushConstantRange::default()
|
||||
.stage_flags(vk::ShaderStageFlags::COMPUTE)
|
||||
.size(std::mem::size_of::<Push>() as u32)];
|
||||
let dl = [self.desc_layout];
|
||||
self.pipe_layout = d
|
||||
.create_pipeline_layout(
|
||||
&vk::PipelineLayoutCreateInfo::default()
|
||||
.set_layouts(&dl)
|
||||
.push_constant_ranges(&pc),
|
||||
None,
|
||||
)
|
||||
.context("create pipeline layout")?;
|
||||
let pool_sizes = [vk::DescriptorPoolSize::default()
|
||||
.ty(vk::DescriptorType::STORAGE_BUFFER)
|
||||
.descriptor_count(2)];
|
||||
self.desc_pool = d
|
||||
.create_descriptor_pool(
|
||||
&vk::DescriptorPoolCreateInfo::default()
|
||||
.max_sets(1)
|
||||
.pool_sizes(&pool_sizes),
|
||||
None,
|
||||
)
|
||||
.context("create descriptor pool")?;
|
||||
let dls = [self.desc_layout];
|
||||
self.desc_set = d
|
||||
.allocate_descriptor_sets(
|
||||
&vk::DescriptorSetAllocateInfo::default()
|
||||
.descriptor_pool(self.desc_pool)
|
||||
.set_layouts(&dls),
|
||||
)
|
||||
.context("allocate descriptor set")?[0];
|
||||
|
||||
// The shader + one pipeline per MODE (spec constant 0).
|
||||
if CURSOR_SPV.len() % 4 != 0 {
|
||||
anyhow::bail!("cursor_blend.spv is not word-aligned");
|
||||
}
|
||||
let words: Vec<u32> = CURSOR_SPV
|
||||
.chunks_exact(4)
|
||||
.map(|c| u32::from_le_bytes([c[0], c[1], c[2], c[3]]))
|
||||
.collect();
|
||||
self.shader = d
|
||||
.create_shader_module(&vk::ShaderModuleCreateInfo::default().code(&words), None)
|
||||
.context("create blend shader module")?;
|
||||
for mode in 0u32..3 {
|
||||
let entries = [vk::SpecializationMapEntry::default()
|
||||
.constant_id(0)
|
||||
.offset(0)
|
||||
.size(4)];
|
||||
let data = mode.to_le_bytes();
|
||||
let spec = vk::SpecializationInfo::default()
|
||||
.map_entries(&entries)
|
||||
.data(&data);
|
||||
let stage = vk::PipelineShaderStageCreateInfo::default()
|
||||
.stage(vk::ShaderStageFlags::COMPUTE)
|
||||
.module(self.shader)
|
||||
.name(c"main")
|
||||
.specialization_info(&spec);
|
||||
let info = [vk::ComputePipelineCreateInfo::default()
|
||||
.stage(stage)
|
||||
.layout(self.pipe_layout)];
|
||||
let p = d
|
||||
.create_compute_pipelines(vk::PipelineCache::null(), &info, None)
|
||||
.map_err(|(_, e)| e)
|
||||
.context("create blend pipeline")?[0];
|
||||
self.pipelines[mode as usize] = p;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn memory_type(&self, type_bits: u32, flags: vk::MemoryPropertyFlags) -> Result<u32> {
|
||||
(0..self.mem_props.memory_type_count)
|
||||
.find(|&i| {
|
||||
type_bits & (1 << i) != 0
|
||||
&& self.mem_props.memory_types[i as usize]
|
||||
.property_flags
|
||||
.contains(flags)
|
||||
})
|
||||
.ok_or_else(|| anyhow!("no memory type for flags {flags:?}"))
|
||||
}
|
||||
|
||||
/// Allocate one NVENC input slot as exportable Vulkan memory mapped into CUDA. Layout matches
|
||||
/// `InputSurface` (contiguous planes under one pitch); pitch = row bytes rounded to 256.
|
||||
pub fn alloc_slot(&mut self, fmt: SlotFormat, width: u32, height: u32) -> Result<VkSlotRef> {
|
||||
let pitch = (fmt.row_bytes(width) + 255) & !255;
|
||||
let size = pitch * fmt.rows(height);
|
||||
// SAFETY: exportable-buffer allocation, the exact `VkBridge::ensure_dst` incantation:
|
||||
// `ExternalMemoryBufferCreateInfo`/`ExportMemoryAllocateInfo` declare OPAQUE_FD,
|
||||
// `MemoryDedicatedAllocateInfo` ties the memory to the buffer; every info is a local
|
||||
// outliving its synchronous call and every failure path destroys the objects created so
|
||||
// far exactly once. `get_memory_fd` hands us an fd that `import_owned_fd` either adopts
|
||||
// (driver owns it) or closes on failure.
|
||||
unsafe {
|
||||
let d = &self.device;
|
||||
let mut ext_info = vk::ExternalMemoryBufferCreateInfo::default()
|
||||
.handle_types(vk::ExternalMemoryHandleTypeFlags::OPAQUE_FD);
|
||||
let buffer = d
|
||||
.create_buffer(
|
||||
&vk::BufferCreateInfo::default()
|
||||
.size(size)
|
||||
.usage(vk::BufferUsageFlags::STORAGE_BUFFER)
|
||||
.push_next(&mut ext_info),
|
||||
None,
|
||||
)
|
||||
.context("create slot buffer")?;
|
||||
let reqs = d.get_buffer_memory_requirements(buffer);
|
||||
let mem_type = match self
|
||||
.memory_type(reqs.memory_type_bits, vk::MemoryPropertyFlags::DEVICE_LOCAL)
|
||||
{
|
||||
Ok(t) => t,
|
||||
Err(e) => {
|
||||
d.destroy_buffer(buffer, None);
|
||||
return Err(e);
|
||||
}
|
||||
};
|
||||
let mut export = vk::ExportMemoryAllocateInfo::default()
|
||||
.handle_types(vk::ExternalMemoryHandleTypeFlags::OPAQUE_FD);
|
||||
let mut dedicated = vk::MemoryDedicatedAllocateInfo::default().buffer(buffer);
|
||||
let memory = match d.allocate_memory(
|
||||
&vk::MemoryAllocateInfo::default()
|
||||
.allocation_size(reqs.size)
|
||||
.memory_type_index(mem_type)
|
||||
.push_next(&mut export)
|
||||
.push_next(&mut dedicated),
|
||||
None,
|
||||
) {
|
||||
Ok(m) => m,
|
||||
Err(e) => {
|
||||
d.destroy_buffer(buffer, None);
|
||||
return Err(e).context("allocate exportable slot memory");
|
||||
}
|
||||
};
|
||||
if let Err(e) = d.bind_buffer_memory(buffer, memory, 0) {
|
||||
d.free_memory(memory, None);
|
||||
d.destroy_buffer(buffer, None);
|
||||
return Err(e).context("bind slot memory");
|
||||
}
|
||||
let fd = match self.ext_fd.get_memory_fd(
|
||||
&vk::MemoryGetFdInfoKHR::default()
|
||||
.memory(memory)
|
||||
.handle_type(vk::ExternalMemoryHandleTypeFlags::OPAQUE_FD),
|
||||
) {
|
||||
Ok(f) => f,
|
||||
Err(e) => {
|
||||
d.free_memory(memory, None);
|
||||
d.destroy_buffer(buffer, None);
|
||||
return Err(e).context("vkGetMemoryFdKHR(slot)");
|
||||
}
|
||||
};
|
||||
let ext = match cuda::ExternalDmabuf::import_owned_fd(fd, reqs.size) {
|
||||
Ok(c) => c,
|
||||
Err(e) => {
|
||||
d.free_memory(memory, None);
|
||||
d.destroy_buffer(buffer, None);
|
||||
return Err(e).context("cuImportExternalMemory(slot OPAQUE_FD)");
|
||||
}
|
||||
};
|
||||
let r = VkSlotRef {
|
||||
ptr: ext.ptr,
|
||||
pitch: pitch as usize,
|
||||
height,
|
||||
id: self.slots.len(),
|
||||
};
|
||||
self.slots.push(SlotAlloc {
|
||||
buffer,
|
||||
memory,
|
||||
cuda: ext,
|
||||
size,
|
||||
});
|
||||
Ok(r)
|
||||
}
|
||||
}
|
||||
|
||||
/// Free every allocated slot (encoder teardown, alongside its ring clear). CUDA mappings drop
|
||||
/// first (field order in [`SlotAlloc`] frees `cuda` via its own `Drop` before we free the VK
|
||||
/// objects explicitly here).
|
||||
pub fn free_slots(&mut self) {
|
||||
for s in self.slots.drain(..) {
|
||||
drop(s.cuda); // CUDA's view of the memory goes first
|
||||
// SAFETY: `buffer`/`memory` were created in `alloc_slot`, are uniquely owned by the
|
||||
// drained `SlotAlloc`, and are destroyed exactly once here. No queue work can be
|
||||
// in flight: every `blend` fence-waits before returning.
|
||||
unsafe {
|
||||
self.device.destroy_buffer(s.buffer, None);
|
||||
self.device.free_memory(s.memory, None);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Upload the cursor RGBA (`cw*ch*4`, tight rows) into the mapped staging buffer. Call only
|
||||
/// when the bitmap changes; position moves are push constants.
|
||||
pub fn upload_cursor(&mut self, rgba: &[u8], cw: u32, ch: u32) {
|
||||
let cw = cw.min(CURSOR_MAX);
|
||||
let ch = ch.min(CURSOR_MAX);
|
||||
let len = (cw * ch * 4) as usize;
|
||||
let len = len.min(rgba.len());
|
||||
// SAFETY: `cur_map` is the live persistent mapping of the CURSOR_MAX²·4 host-coherent
|
||||
// allocation (created in `init_objects`, unmapped only in `Drop`); `len` is clamped to
|
||||
// both the source slice and the buffer capacity. No blend is in flight (every `blend`
|
||||
// fence-waits before returning), so no GPU read races this host write.
|
||||
unsafe {
|
||||
std::ptr::copy_nonoverlapping(rgba.as_ptr(), self.cur_map, len);
|
||||
}
|
||||
}
|
||||
|
||||
/// Blend the uploaded cursor into `slot` at `(ox, oy)`: record, submit, fence-wait. The
|
||||
/// caller has CPU-synced its CUDA frame copy first; the fence wait makes the shader's writes
|
||||
/// visible to the subsequent NVENC encode (the `VkBridge` precedent: fence-ordered cross-API
|
||||
/// access, no queue-family transfer — NVIDIA-only path).
|
||||
#[allow(clippy::too_many_arguments)] // surface geometry + cursor rect — unpacked kernel args
|
||||
pub fn blend_ref(
|
||||
&mut self,
|
||||
slot: &VkSlotRef,
|
||||
fmt: SlotFormat,
|
||||
surf_w: u32,
|
||||
cw: u32,
|
||||
ch: u32,
|
||||
ox: i32,
|
||||
oy: i32,
|
||||
) -> Result<()> {
|
||||
let alloc = self
|
||||
.slots
|
||||
.get(slot.id)
|
||||
.ok_or_else(|| anyhow!("bad slot id {}", slot.id))?;
|
||||
let cw = cw.min(CURSOR_MAX);
|
||||
let ch = ch.min(CURSOR_MAX);
|
||||
if cw == 0 || ch == 0 {
|
||||
return Ok(());
|
||||
}
|
||||
let push = Push {
|
||||
pitch: slot.pitch as u32,
|
||||
surf_w,
|
||||
surf_h: slot.height,
|
||||
cur_w: cw,
|
||||
cur_h: ch,
|
||||
ox,
|
||||
oy,
|
||||
};
|
||||
// Dispatch geometry (must match cursor_blend.comp): ARGB = per cursor px; NV12/YUV444 =
|
||||
// per word-aligned 4-px span × (2-row blocks | rows).
|
||||
let (gx, gy) = match fmt {
|
||||
SlotFormat::Argb => (cw.div_ceil(8), ch.div_ceil(8)),
|
||||
_ => {
|
||||
let x0 = (ox >> 2) << 2;
|
||||
let spans = ((ox + cw as i32) - x0 + 3).div_euclid(4).max(1) as u32;
|
||||
let rows = match fmt {
|
||||
SlotFormat::Nv12 => ch.div_ceil(2),
|
||||
_ => ch,
|
||||
};
|
||||
(spans.div_ceil(8), rows.div_ceil(8))
|
||||
}
|
||||
};
|
||||
// SAFETY: single-threaded record/submit/wait on handles this struct owns. The descriptor
|
||||
// update is safe because no prior submission is in flight (every blend fence-waits and
|
||||
// the fence is reset before reuse). Buffer infos and barrier structs are locals outliving
|
||||
// their synchronous calls. The dispatch's shader accesses stay in-bounds by the shader's
|
||||
// own guards (surfW/surfH/curW/curH from `push`) against the slot allocation sized in
|
||||
// `alloc_slot` for exactly that geometry.
|
||||
unsafe {
|
||||
let d = &self.device;
|
||||
let surf_info = [vk::DescriptorBufferInfo::default()
|
||||
.buffer(alloc.buffer)
|
||||
.offset(0)
|
||||
.range(alloc.size)];
|
||||
let cur_info = [vk::DescriptorBufferInfo::default()
|
||||
.buffer(self.cur_buf)
|
||||
.offset(0)
|
||||
.range((CURSOR_MAX * CURSOR_MAX * 4) as u64)];
|
||||
let writes = [
|
||||
vk::WriteDescriptorSet::default()
|
||||
.dst_set(self.desc_set)
|
||||
.dst_binding(0)
|
||||
.descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
|
||||
.buffer_info(&surf_info),
|
||||
vk::WriteDescriptorSet::default()
|
||||
.dst_set(self.desc_set)
|
||||
.dst_binding(1)
|
||||
.descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
|
||||
.buffer_info(&cur_info),
|
||||
];
|
||||
d.update_descriptor_sets(&writes, &[]);
|
||||
|
||||
d.begin_command_buffer(
|
||||
self.cmd,
|
||||
&vk::CommandBufferBeginInfo::default()
|
||||
.flags(vk::CommandBufferUsageFlags::ONE_TIME_SUBMIT),
|
||||
)
|
||||
.context("begin blend cmd")?;
|
||||
// CUDA wrote the frame into this memory outside Vulkan's view — make it visible to
|
||||
// the shader (external-memory coherence ceremony; NVIDIA honors this with the fence
|
||||
// ordering alone, the barrier is the spec-shaped belt-and-braces).
|
||||
let acquire = [vk::MemoryBarrier::default()
|
||||
.src_access_mask(vk::AccessFlags::MEMORY_WRITE)
|
||||
.dst_access_mask(vk::AccessFlags::SHADER_READ | vk::AccessFlags::SHADER_WRITE)];
|
||||
d.cmd_pipeline_barrier(
|
||||
self.cmd,
|
||||
vk::PipelineStageFlags::TOP_OF_PIPE,
|
||||
vk::PipelineStageFlags::COMPUTE_SHADER,
|
||||
vk::DependencyFlags::empty(),
|
||||
&acquire,
|
||||
&[],
|
||||
&[],
|
||||
);
|
||||
d.cmd_bind_pipeline(
|
||||
self.cmd,
|
||||
vk::PipelineBindPoint::COMPUTE,
|
||||
self.pipelines[fmt.mode() as usize],
|
||||
);
|
||||
d.cmd_bind_descriptor_sets(
|
||||
self.cmd,
|
||||
vk::PipelineBindPoint::COMPUTE,
|
||||
self.pipe_layout,
|
||||
0,
|
||||
&[self.desc_set],
|
||||
&[],
|
||||
);
|
||||
let bytes = std::slice::from_raw_parts(
|
||||
(&push as *const Push) as *const u8,
|
||||
std::mem::size_of::<Push>(),
|
||||
);
|
||||
d.cmd_push_constants(
|
||||
self.cmd,
|
||||
self.pipe_layout,
|
||||
vk::ShaderStageFlags::COMPUTE,
|
||||
0,
|
||||
bytes,
|
||||
);
|
||||
d.cmd_dispatch(self.cmd, gx.max(1), gy.max(1), 1);
|
||||
// Release the shader's writes so the post-fence CUDA/NVENC reads see them.
|
||||
let release = [vk::MemoryBarrier::default()
|
||||
.src_access_mask(vk::AccessFlags::SHADER_WRITE)
|
||||
.dst_access_mask(vk::AccessFlags::MEMORY_READ)];
|
||||
d.cmd_pipeline_barrier(
|
||||
self.cmd,
|
||||
vk::PipelineStageFlags::COMPUTE_SHADER,
|
||||
vk::PipelineStageFlags::BOTTOM_OF_PIPE,
|
||||
vk::DependencyFlags::empty(),
|
||||
&release,
|
||||
&[],
|
||||
&[],
|
||||
);
|
||||
d.end_command_buffer(self.cmd).context("end blend cmd")?;
|
||||
let cmds = [self.cmd];
|
||||
let submit = [vk::SubmitInfo::default().command_buffers(&cmds)];
|
||||
d.queue_submit(self.queue, &submit, self.fence)
|
||||
.context("submit blend")?;
|
||||
let r = d.wait_for_fences(&[self.fence], true, 1_000_000_000);
|
||||
d.reset_fences(&[self.fence]).ok();
|
||||
r.context("blend fence wait")?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for VkSlotBlend {
|
||||
fn drop(&mut self) {
|
||||
self.free_slots();
|
||||
// SAFETY: every handle below was created in `new`/`init_objects` (or is null from a
|
||||
// partial init — Vulkan destroy/free calls are defined no-ops on null handles) and is
|
||||
// uniquely owned; each is destroyed exactly once here, pipelines/layouts/pools before the
|
||||
// device, the device before the instance. No work is in flight (`blend` fence-waits).
|
||||
unsafe {
|
||||
let d = &self.device;
|
||||
for p in self.pipelines {
|
||||
if p != vk::Pipeline::null() {
|
||||
d.destroy_pipeline(p, None);
|
||||
}
|
||||
}
|
||||
if self.shader != vk::ShaderModule::null() {
|
||||
d.destroy_shader_module(self.shader, None);
|
||||
}
|
||||
if self.desc_pool != vk::DescriptorPool::null() {
|
||||
d.destroy_descriptor_pool(self.desc_pool, None);
|
||||
}
|
||||
if self.pipe_layout != vk::PipelineLayout::null() {
|
||||
d.destroy_pipeline_layout(self.pipe_layout, None);
|
||||
}
|
||||
if self.desc_layout != vk::DescriptorSetLayout::null() {
|
||||
d.destroy_descriptor_set_layout(self.desc_layout, None);
|
||||
}
|
||||
if !self.cur_map.is_null() {
|
||||
d.unmap_memory(self.cur_mem);
|
||||
}
|
||||
if self.cur_buf != vk::Buffer::null() {
|
||||
d.destroy_buffer(self.cur_buf, None);
|
||||
}
|
||||
if self.cur_mem != vk::DeviceMemory::null() {
|
||||
d.free_memory(self.cur_mem, None);
|
||||
}
|
||||
if self.fence != vk::Fence::null() {
|
||||
d.destroy_fence(self.fence, None);
|
||||
}
|
||||
if self.cmd_pool != vk::CommandPool::null() {
|
||||
d.destroy_command_pool(self.cmd_pool, None);
|
||||
}
|
||||
d.destroy_device(None);
|
||||
self.instance.destroy_instance(None);
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user