From 7d94fe8c5f60b8d8343ccf431804f2efe5ace881 Mon Sep 17 00:00:00 2001 From: Karol Herbst Date: Tue, 4 Mar 2025 00:15:44 +0100 Subject: [PATCH] rusticl/queue: cache bound CSO Part-of: --- src/gallium/frontends/rusticl/core/kernel.rs | 39 +++++---- src/gallium/frontends/rusticl/core/queue.rs | 85 +++++++++++++++++++ .../frontends/rusticl/mesa/pipe/context.rs | 5 +- 3 files changed, 110 insertions(+), 19 deletions(-) diff --git a/src/gallium/frontends/rusticl/core/kernel.rs b/src/gallium/frontends/rusticl/core/kernel.rs index d5c1072bd09..497516ef7ed 100644 --- a/src/gallium/frontends/rusticl/core/kernel.rs +++ b/src/gallium/frontends/rusticl/core/kernel.rs @@ -10,6 +10,7 @@ use crate::impl_cl_type_trait; use mesa_rust::compiler::clc::*; use mesa_rust::compiler::nir::*; use mesa_rust::nir_pass; +use mesa_rust::pipe::context::PipeContext; use mesa_rust::pipe::context::RWFlags; use mesa_rust::pipe::resource::*; use mesa_rust::pipe::screen::ResourceType; @@ -343,7 +344,7 @@ pub struct KernelInfo { } /// Wraps around a compute state object which is safe to share between pipe_contexts. -struct SharedCSOWrapper { +pub struct SharedCSOWrapper { cso_ptr: *mut c_void, dev: &'static Device, } @@ -364,6 +365,16 @@ impl SharedCSOWrapper { } } + /// # Safety + /// + /// `self` needs to live until another CSOWrapper is bound to `ctx` + pub unsafe fn bind_to_ctx(&self, ctx: &PipeContext) { + // SAFETY: We make it the callers responsibility to uphold the safety requirements. + unsafe { + ctx.bind_compute_state(self.cso_ptr); + } + } + fn get_cso_info(&self) -> pipe_compute_state_object_info { self.dev.helper_ctx().compute_state_info(self.cso_ptr) } @@ -375,13 +386,13 @@ impl Drop for SharedCSOWrapper { } } -enum KernelDevStateVariant { +pub enum KernelDevStateVariant { Cso(SharedCSOWrapper), Nir(NirShader), } -#[derive(Debug, PartialEq)] -enum NirKernelVariant { +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum NirKernelVariant { /// Can be used under any circumstance. Default, @@ -436,7 +447,7 @@ impl NirKernelBuilds { } } -struct NirKernelBuild { +pub struct NirKernelBuild { nir_or_cso: KernelDevStateVariant, constant_buffer: Option>, info: pipe_compute_state_object_info, @@ -496,6 +507,10 @@ impl NirKernelBuild { None } } + + pub fn nir_or_cso(&self) -> &KernelDevStateVariant { + &self.nir_or_cso + } } pub struct Kernel { @@ -1558,18 +1573,8 @@ impl Kernel { globals.push(unsafe { input.as_mut_ptr().byte_add(offset) }.cast()); } - let temp_cso; - let cso = match &nir_kernel_build.nir_or_cso { - KernelDevStateVariant::Cso(cso) => cso, - KernelDevStateVariant::Nir(nir) => { - // SAFETY: this isn't safe at all, but we'll fix this in a later commit. - temp_cso = unsafe { SharedCSOWrapper::new(q.device, nir) }; - &temp_cso - } - }; - let sviews_len = sviews.len(); - ctx.bind_compute_state(cso.cso_ptr); + ctx.bind_kernel(&nir_kernel_builds, variant)?; ctx.bind_sampler_states(&samplers); ctx.set_sampler_views(sviews); ctx.set_shader_images(&iviews); @@ -1613,8 +1618,6 @@ impl Kernel { ctx.clear_sampler_views(sviews_len as u32); ctx.clear_sampler_states(samplers.len() as u32); - ctx.bind_compute_state(ptr::null_mut()); - ctx.memory_barrier(PIPE_BARRIER_GLOBAL_BUFFER); samplers.iter().for_each(|s| ctx.delete_sampler_state(*s)); diff --git a/src/gallium/frontends/rusticl/core/queue.rs b/src/gallium/frontends/rusticl/core/queue.rs index 6ce9ccb30fe..fb8f06fdac4 100644 --- a/src/gallium/frontends/rusticl/core/queue.rs +++ b/src/gallium/frontends/rusticl/core/queue.rs @@ -2,18 +2,24 @@ use crate::api::icd::*; use crate::core::context::*; use crate::core::device::*; use crate::core::event::*; +use crate::core::kernel::*; use crate::core::platform::*; use crate::impl_cl_type_trait; +use mesa_rust::compiler::nir::NirShader; use mesa_rust::pipe::context::PipeContext; use mesa_rust_gen::*; use mesa_rust_util::properties::*; use rusticl_opencl_gen::*; +use std::cell::RefCell; use std::cmp; +use std::ffi::c_void; use std::mem; use std::mem::ManuallyDrop; use std::ops::Deref; +use std::ptr; +use std::ptr::NonNull; use std::sync::mpsc; use std::sync::Arc; use std::sync::Mutex; @@ -21,6 +27,32 @@ use std::sync::Weak; use std::thread; use std::thread::JoinHandle; +struct CSOWrapper<'a> { + ctx: &'a PipeContext, + cso: NonNull, +} + +impl<'a> CSOWrapper<'a> { + fn new(ctx: &QueueContext<'a>, nir: &NirShader) -> Option> { + Some(Self { + ctx: ctx.ctx, + cso: NonNull::new(ctx.create_compute_state(nir, nir.shared_size()))?, + }) + } +} + +impl Drop for CSOWrapper<'_> { + fn drop(&mut self) { + self.ctx.delete_compute_state(self.cso.as_ptr()); + } +} + +struct QueueKernelState<'a> { + builds: Option>, + variant: NirKernelVariant, + cso: Option>, +} + /// State tracking wrapper for [PipeContext] /// /// Used for tracking bound GPU state to lower CPU overhead and centralize state tracking @@ -28,9 +60,50 @@ pub struct QueueContext<'a> { ctx: &'a PipeContext, pub dev: &'static Device, use_stream: bool, + kernel_state: RefCell>, } impl QueueContext<'_> { + // TODO: figure out how to make it &mut self without causing tons of borrowing issues. + pub fn bind_kernel( + &self, + builds: &Arc, + variant: NirKernelVariant, + ) -> CLResult<()> { + // this should never panic, but you never know. + let mut state = self.kernel_state.borrow_mut(); + + // If we already set the CSO then we don't have to bind again. + if let Some(stored_builds) = &state.builds { + if Arc::ptr_eq(stored_builds, builds) && state.variant == variant { + return Ok(()); + } + } + + let nir_kernel_build = &builds[variant]; + match nir_kernel_build.nir_or_cso() { + // SAFETY: We keep the cso alive until a new one is set. + KernelDevStateVariant::Cso(cso) => unsafe { + cso.bind_to_ctx(self); + }, + // TODO: We could cache the cso here. + KernelDevStateVariant::Nir(nir) => { + let cso = CSOWrapper::new(self, nir).ok_or(CL_OUT_OF_HOST_MEMORY)?; + unsafe { + self.bind_compute_state(cso.cso.as_ptr()); + } + state.cso.replace(cso); + } + }; + + // We can only store the new builds after we bound the new cso otherwise we might drop it + // too early. + state.builds = Some(Arc::clone(builds)); + state.variant = variant; + + Ok(()) + } + pub fn update_cb0(&self, data: &[u8]) -> CLResult<()> { // only update if we actually bind data if !data.is_empty() { @@ -58,6 +131,13 @@ impl Deref for QueueContext<'_> { impl Drop for QueueContext<'_> { fn drop(&mut self) { self.set_constant_buffer(0, &[]); + if self.kernel_state.get_mut().builds.is_some() { + // SAFETY: We simply unbind here. The bound cso will only be dropped at the end of this + // drop handler. + unsafe { + self.ctx.bind_compute_state(ptr::null_mut()); + } + } } } @@ -81,6 +161,11 @@ impl SendableQueueContext { fn ctx(&self) -> QueueContext { QueueContext { ctx: &self.ctx, + kernel_state: RefCell::new(QueueKernelState { + builds: None, + variant: NirKernelVariant::Default, + cso: None, + }), dev: self.dev, use_stream: self.dev.prefers_real_buffer_in_cb0(), } diff --git a/src/gallium/frontends/rusticl/mesa/pipe/context.rs b/src/gallium/frontends/rusticl/mesa/pipe/context.rs index dae76659a76..e02a26ea7fb 100644 --- a/src/gallium/frontends/rusticl/mesa/pipe/context.rs +++ b/src/gallium/frontends/rusticl/mesa/pipe/context.rs @@ -309,7 +309,10 @@ impl PipeContext { unsafe { self.pipe.as_ref().create_compute_state.unwrap()(self.pipe.as_ptr(), &state) } } - pub fn bind_compute_state(&self, state: *mut c_void) { + /// # Safety + /// + /// The state pointer needs to point to valid memory until a new one is set. + pub unsafe fn bind_compute_state(&self, state: *mut c_void) { unsafe { self.pipe.as_ref().bind_compute_state.unwrap()(self.pipe.as_ptr(), state) } }