From a901365656b7172bbf850d6883eb4896769306e9 Mon Sep 17 00:00:00 2001 From: modeco80 Date: Sun, 6 Oct 2024 04:16:13 -0400 Subject: [PATCH] begin adding provision for hardware encoding functionality --- server/Cargo.lock | 10 ++++ server/Cargo.toml | 3 +- server/src/video/encoder_thread.rs | 32 +++++++++++ server/src/video/h264_encoder.rs | 71 +++++++++++++++++++++--- server/src/video/hwdevice.rs | 78 ++++++++++++++++++++++++++ server/src/video/hwframe.rs | 88 ++++++++++++++++++++++++++++++ server/src/video/mod.rs | 11 ++++ 7 files changed, 283 insertions(+), 10 deletions(-) create mode 100644 server/src/video/hwdevice.rs create mode 100644 server/src/video/hwframe.rs diff --git a/server/Cargo.lock b/server/Cargo.lock index 3ce6baf..4c5bd03 100644 --- a/server/Cargo.lock +++ b/server/Cargo.lock @@ -273,6 +273,15 @@ dependencies = [ "typenum", ] +[[package]] +name = "cudarc" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38cd60a9a42ec83a2ed7effb0b1f073270264ea99da7acfc44f7e8d74dee0384" +dependencies = [ + "libloading", +] + [[package]] name = "data-encoding" version = "2.6.0" @@ -1317,6 +1326,7 @@ version = "0.1.0" dependencies = [ "anyhow", "axum", + "cudarc", "ffmpeg-next", "futures", "futures-util", diff --git a/server/Cargo.toml b/server/Cargo.toml index f0384b0..03f3095 100644 --- a/server/Cargo.toml +++ b/server/Cargo.toml @@ -23,8 +23,9 @@ ffmpeg = { version = "7.0.0", package = "ffmpeg-next" } rand = "0.8.5" serde = "1.0.209" serde_json = "1.0.128" +cudarc = "0.12.1" [patch.crates-io] ffmpeg = { version = "7.0.0", package = "ffmpeg-next", git = "https://github.com/hgaiser/rust-ffmpeg", branch = "codec-context-settable" } -ffmpeg-sys-next = { version = "7.0.0", git = "https://github.com/hgaiser/rust-ffmpeg-sys", branch = "cuda" } \ No newline at end of file +ffmpeg-sys-next = { version = "7.0.0", git = "https://github.com/hgaiser/rust-ffmpeg-sys", branch = "cuda" } diff --git a/server/src/video/encoder_thread.rs b/server/src/video/encoder_thread.rs index b1ef301..cffc082 100644 --- a/server/src/video/encoder_thread.rs +++ b/server/src/video/encoder_thread.rs @@ -7,6 +7,8 @@ use tokio::sync::mpsc::{self, error::TryRecvError}; use super::ffmpeg; use super::h264_encoder::H264Encoder; +use super::hwframe::HwFrameContext; + pub enum EncodeThreadInput { Init { size: crate::types::Size }, ForceKeyframe, @@ -33,6 +35,36 @@ fn set_frame_flags(frame: &mut ffmpeg::Frame, force_keyframe: bool) { } } +fn create_frame( + width: u32, + height: u32, + pixel_format: ffmpeg::format::Pixel, + context: &mut HwFrameContext, +) -> anyhow::Result { + unsafe { + let mut frame = ffmpeg::Frame::empty(); + (*frame.as_mut_ptr()).format = pixel_format as i32; + (*frame.as_mut_ptr()).width = width as i32; + (*frame.as_mut_ptr()).height = height as i32; + (*frame.as_mut_ptr()).hw_frames_ctx = context.as_raw_mut(); + + super::check_ret(ffmpeg::sys::av_hwframe_get_buffer( + context.as_raw_mut(), + frame.as_mut_ptr(), + 0, + ))?; + super::check_ret(ffmpeg::sys::av_hwframe_get_buffer( + context.as_raw_mut(), + frame.as_mut_ptr(), + 0, + ))?; + + (*frame.as_mut_ptr()).linesize[0] = (*frame.as_ptr()).width * 4; + + Ok(frame) + } +} + fn encoder_thread_main( mut rx: mpsc::Receiver, tx: mpsc::Sender, diff --git a/server/src/video/h264_encoder.rs b/server/src/video/h264_encoder.rs index db963fc..4c4e7d9 100644 --- a/server/src/video/h264_encoder.rs +++ b/server/src/video/h264_encoder.rs @@ -1,8 +1,9 @@ use super::ffmpeg; +use super::hwframe::HwFrameContext; use anyhow::Context; use ffmpeg::error::EAGAIN; -use ffmpeg::codec as lavc; // lavc +use ffmpeg::{codec as lavc, packet}; // lavc use crate::types::Size; @@ -30,8 +31,6 @@ fn create_context_and_set_common_parameters( let mut video_encoder_context = create_context_from_codec(encoder)?.encoder().video()?; - - video_encoder_context.set_width(size.width); video_encoder_context.set_height(size.height); video_encoder_context.set_frame_rate(Some(ffmpeg::Rational(1, max_framerate as i32))); @@ -63,8 +62,15 @@ fn create_context_and_set_common_parameters( /// A simple H.264 encoder. Currently software only, however /// pieces are being put in place to eventually allow HW encoding. -pub struct H264Encoder { - encoder: ffmpeg::encoder::video::Encoder, +pub enum H264Encoder { + Software { + encoder: ffmpeg::encoder::video::Encoder, + }, + + Nvenc { + encoder: ffmpeg::encoder::video::Encoder, + hw_context: HwFrameContext, + }, } impl H264Encoder { @@ -107,21 +113,68 @@ impl H264Encoder { .open_as_with(encoder, dict) .with_context(|| "While opening x264 video codec")?; - Ok(Self { encoder: encoder }) + Ok(Self::Software { encoder: encoder }) + } + + // FIXME: It's a bit pointless to have this have a mut borrow, + // but you'll probably have a mutable borrow on this already.. + pub fn is_hardware(&mut self) -> bool { + match self { + Self::Software { .. } => false, + Self::Nvenc { .. } => true, + } } pub fn send_frame(&mut self, frame: &ffmpeg::Frame) { - self.encoder.send_frame(frame).unwrap(); + match self { + Self::Software { encoder } => { + encoder.send_frame(frame).unwrap(); + } + + Self::Nvenc { + encoder, + hw_context, + } => { + // Realistically this should be the same right? + todo!("Requires support."); + } + } } pub fn send_eof(&mut self) { - self.encoder.send_eof().unwrap(); + match self { + Self::Software { encoder } => { + encoder.send_eof().unwrap(); + } + + Self::Nvenc { + encoder, + hw_context, + } => { + // Realistically this should be the same right? + todo!("Requires support."); + } + } + } + + fn receive_packet_impl(&mut self, packet: &mut ffmpeg::Packet) -> Result<(), ffmpeg::Error> { + return match self { + Self::Software { encoder } => encoder.receive_packet(packet), + + Self::Nvenc { + encoder, + hw_context, + } => { + // Realistically this should be the same right? + todo!("Requires support."); + } + }; } // Shuold this return a Result so we can make it easier to know when to continue? pub fn receive_packet(&mut self, packet: &mut ffmpeg::Packet) -> anyhow::Result<()> { loop { - match self.encoder.receive_packet(packet) { + match self.receive_packet_impl(packet) { Ok(_) => break, Err(ffmpeg::Error::Other { errno }) => { if errno != EAGAIN { diff --git a/server/src/video/hwdevice.rs b/server/src/video/hwdevice.rs new file mode 100644 index 0000000..ce9a208 --- /dev/null +++ b/server/src/video/hwdevice.rs @@ -0,0 +1,78 @@ +use std::ptr::null_mut; + +use super::check_ret; + +use super::ffmpeg; + +pub struct CudaDeviceContext { + buffer: *mut ffmpeg::sys::AVBufferRef, +} + +impl CudaDeviceContext { + fn new(buffer: *mut ffmpeg::sys::AVBufferRef) -> Self { + Self { buffer } + } + + // pub fn as_device_mut(&mut self) -> &mut ffmpeg::sys::AVHWDeviceContext { + // unsafe { &mut *((*self.buffer).data as *mut ffmpeg::sys::AVHWDeviceContext) } + // } + + // pub fn as_device(&self) -> &ffmpeg::sys::AVHWDeviceContext { + // unsafe { &*((*self.buffer).data as *const ffmpeg::sys::AVHWDeviceContext) } + // } + + pub fn as_raw_mut(&mut self) -> &mut ffmpeg::sys::AVBufferRef { + unsafe { &mut *self.buffer } + } + + // pub fn as_raw(&self) -> &ffmpeg::sys::AVBufferRef { + // unsafe { &*self.buffer } + // } +} + +pub struct CudaDeviceContextBuilder { + buffer: *mut ffmpeg::sys::AVBufferRef, +} + +impl CudaDeviceContextBuilder { + pub fn new() -> Result { + let buffer = unsafe { ffmpeg::sys::av_hwdevice_ctx_alloc(ffmpeg::sys::AVHWDeviceType::AV_HWDEVICE_TYPE_CUDA) }; + if buffer.is_null() { + return Err("could not allocate a hwdevice".to_string()); + } + + Ok(Self { buffer }) + } + + pub fn build(mut self) -> Result { + check_ret(unsafe { ffmpeg::sys::av_hwdevice_ctx_init(self.buffer) })?; + let result = Ok(CudaDeviceContext::new(self.buffer)); + self.buffer = null_mut(); + + result + } + + pub fn set_cuda_context(mut self, context: ffmpeg::sys::CUcontext) -> Self { + unsafe { + (*(self.as_device_mut().hwctx as *mut ffmpeg::sys::AVCUDADeviceContext)).cuda_ctx = context; + } + + self + } + + pub fn as_device_mut(&mut self) -> &mut ffmpeg::sys::AVHWDeviceContext { + unsafe { &mut *((*self.buffer).data as *mut ffmpeg::sys::AVHWDeviceContext) } + } + + // pub fn as_device(&self) -> &ffmpeg::sys::AVHWDeviceContext { + // unsafe { &*((*self.buffer).data as *const ffmpeg::sys::AVHWDeviceContext) } + // } + + // pub fn as_raw_mut(&mut self) -> &mut ffmpeg::sys::AVBufferRef { + // unsafe { &mut *self.buffer } + // } + + // pub fn as_raw(&self) -> &ffmpeg::sys::AVBufferRef { + // unsafe { &*self.buffer } + // } +} \ No newline at end of file diff --git a/server/src/video/hwframe.rs b/server/src/video/hwframe.rs new file mode 100644 index 0000000..e040b6c --- /dev/null +++ b/server/src/video/hwframe.rs @@ -0,0 +1,88 @@ +use std::ptr::null_mut; + +use super::ffmpeg; + +use ffmpeg::format::Pixel; + +use super::{check_ret, hwdevice::CudaDeviceContext}; + +pub struct HwFrameContext { + _cuda_device_context: CudaDeviceContext, + buffer: *mut ffmpeg::sys::AVBufferRef, +} + +impl HwFrameContext { + fn new(cuda_device_context: CudaDeviceContext, buffer: *mut ffmpeg::sys::AVBufferRef) -> Self { + Self { _cuda_device_context: cuda_device_context, buffer } + } + + // pub fn as_context_mut(&mut self) -> &mut ffmpeg::sys::AVHWFramesContext { + // unsafe { &mut *((*self.buffer).data as *mut ffmpeg::sys::AVHWFramesContext) } + // } + + // pub fn as_context(&self) -> &ffmpeg::sys::AVHWFramesContext { + // unsafe { &*((*self.buffer).data as *const ffmpeg::sys::AVHWFramesContext) } + // } + + pub fn as_raw_mut(&mut self) -> &mut ffmpeg::sys::AVBufferRef { + unsafe { &mut *self.buffer } + } + + // pub fn as_raw(&self) -> &ffmpeg::sys::AVBufferRef { + // unsafe { &*self.buffer } + // } +} + +unsafe impl Send for HwFrameContext { } + +pub struct HwFrameContextBuilder { + cuda_device_context: CudaDeviceContext, + buffer: *mut ffmpeg::sys::AVBufferRef, +} + +impl HwFrameContextBuilder { + pub fn new(mut cuda_device_context: CudaDeviceContext) -> Result { + let buffer = unsafe { ffmpeg::sys::av_hwframe_ctx_alloc(cuda_device_context.as_raw_mut()) }; + if buffer.is_null() { + return Err("could not allocate a hwframe".to_string()); + } + + Ok(Self { cuda_device_context, buffer }) + } + + pub fn build(mut self) -> Result { + check_ret(unsafe { ffmpeg::sys::av_hwframe_ctx_init(self.buffer) })?; + let result = Ok(HwFrameContext::new(self.cuda_device_context, self.buffer)); + self.buffer = null_mut(); + + result + } + + pub fn set_width(mut self, width: u32) -> Self { + self.as_frame_mut().width = width as i32; + self + } + + pub fn set_height(mut self, height: u32) -> Self { + self.as_frame_mut().height = height as i32; + self + } + + pub fn set_sw_format(mut self, sw_format: Pixel) -> Self { + self.as_frame_mut().sw_format = sw_format.into(); + self + } + + pub fn set_format(mut self, format: Pixel) -> Self { + self.as_frame_mut().format = format.into(); + self + } + + pub fn as_frame_mut(&mut self) -> &mut ffmpeg::sys::AVHWFramesContext { + unsafe { &mut *((*self.buffer).data as *mut ffmpeg::sys::AVHWFramesContext) } + } + + // pub fn as_frame(&self) -> &ffmpeg::sys::AVHWFramesContext { + // unsafe { &*((*self.buffer).data as *const ffmpeg::sys::AVHWFramesContext) } + // } +} \ No newline at end of file diff --git a/server/src/video/mod.rs b/server/src/video/mod.rs index dc337b9..e5b8689 100644 --- a/server/src/video/mod.rs +++ b/server/src/video/mod.rs @@ -5,5 +5,16 @@ pub mod h264_encoder; /// Re-export of `ffmpeg_the_third` crate in an infinitely less obtuse name. pub use ffmpeg as ffmpeg; +pub mod hwdevice; +pub mod hwframe; pub use encoder_thread::*; + +// from hgaiser/moonshine +pub fn check_ret(error_code: i32) -> Result<(), ffmpeg::Error> { + if error_code != 0 { + return Err(ffmpeg::Error::from(error_code)); + } + + Ok(()) +} \ No newline at end of file