From e2029030070392e103951c0bab032bf0b20663dc Mon Sep 17 00:00:00 2001
From: Karol Herbst <kherbst@redhat.com>
Date: Thu, 22 Feb 2024 13:05:15 +0100
Subject: [PATCH] rusticl/kernel: make it Send and Sync

There are a few APIs which don't have to be thread-safe, but we can
optimize it later.

Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/27747>
---
 src/gallium/frontends/rusticl/api/kernel.rs  | 11 +++---
 src/gallium/frontends/rusticl/core/kernel.rs | 36 +++++++++++++-------
 2 files changed, 28 insertions(+), 19 deletions(-)

diff --git a/src/gallium/frontends/rusticl/api/kernel.rs b/src/gallium/frontends/rusticl/api/kernel.rs
index 6fd1288ea07..02865dbad38 100644
--- a/src/gallium/frontends/rusticl/api/kernel.rs
+++ b/src/gallium/frontends/rusticl/api/kernel.rs
@@ -331,9 +331,10 @@ fn set_kernel_arg(
     arg_value: *const ::std::os::raw::c_void,
 ) -> CLResult<()> {
     let k = Kernel::ref_from_raw(kernel)?;
+    let arg_index = arg_index as usize;
 
     // CL_INVALID_ARG_INDEX if arg_index is not a valid argument index.
-    if let Some(arg) = k.kernel_info.args.get(arg_index as usize) {
+    if let Some(arg) = k.kernel_info.args.get(arg_index) {
         // CL_INVALID_ARG_SIZE if arg_size does not match the size of the data type for an argument
         // that is not a memory object or if the argument is a memory object and
         // arg_size != sizeof(cl_mem) or if arg_size is zero and the argument is declared with the
@@ -408,8 +409,7 @@ fn set_kernel_arg(
                 }
             }
         };
-        k.values.get(arg_index as usize).unwrap().replace(Some(arg));
-        Ok(())
+        k.set_kernel_arg(arg_index, arg)
     } else {
         Err(CL_INVALID_ARG_INDEX)
     }
@@ -442,8 +442,7 @@ fn set_kernel_arg_svm_pointer(
         }
 
         let arg_value = KernelArgValue::Constant(arg_value.to_ne_bytes().to_vec());
-        kernel.values[arg_index].replace(Some(arg_value));
-        Ok(())
+        kernel.set_kernel_arg(arg_index, arg_value)
     } else {
         Err(CL_INVALID_ARG_INDEX)
     }
@@ -521,7 +520,7 @@ fn enqueue_ndrange_kernel(
     }
 
     // CL_INVALID_KERNEL_ARGS if the kernel argument values have not been specified.
-    if k.values.iter().any(|v| v.borrow().is_none()) {
+    if k.arg_values().iter().any(|v| v.is_none()) {
         return Err(CL_INVALID_KERNEL_ARGS);
     }
 
diff --git a/src/gallium/frontends/rusticl/core/kernel.rs b/src/gallium/frontends/rusticl/core/kernel.rs
index 6c4b477fb18..f6b086945c4 100644
--- a/src/gallium/frontends/rusticl/core/kernel.rs
+++ b/src/gallium/frontends/rusticl/core/kernel.rs
@@ -18,7 +18,6 @@ use mesa_rust_util::math::*;
 use mesa_rust_util::serialize::*;
 use rusticl_opencl_gen::*;
 
-use std::cell::RefCell;
 use std::cmp;
 use std::collections::HashMap;
 use std::convert::TryInto;
@@ -26,6 +25,8 @@ use std::os::raw::c_void;
 use std::ptr;
 use std::slice;
 use std::sync::Arc;
+use std::sync::Mutex;
+use std::sync::MutexGuard;
 
 // ugh, we are not allowed to take refs, so...
 #[derive(Clone)]
@@ -299,7 +300,7 @@ pub struct Kernel {
     pub base: CLObjectBase<CL_INVALID_KERNEL>,
     pub prog: Arc<Program>,
     pub name: String,
-    pub values: Vec<RefCell<Option<KernelArgValue>>>,
+    values: Mutex<Vec<Option<KernelArgValue>>>,
     pub builds: HashMap<&'static Device, Arc<NirKernelBuild>>,
     pub kernel_info: KernelInfo,
 }
@@ -818,18 +819,12 @@ impl Kernel {
             .filter_map(|(&dev, b)| b.kernels.get(&name).map(|k| (dev, k.clone())))
             .collect();
 
-        // can't use vec!...
-        let values = kernel_info
-            .args
-            .iter()
-            .map(|_| RefCell::new(None))
-            .collect();
-
+        let values = vec![None; kernel_info.args.len()];
         Arc::new(Self {
             base: CLObjectBase::new(RusticlTypes::Kernel),
             prog: prog.clone(),
             name: name,
-            values: values,
+            values: Mutex::new(values),
             builds: builds,
             kernel_info: kernel_info,
         })
@@ -908,7 +903,8 @@ impl Kernel {
 
         self.optimize_local_size(q.device, &mut grid, &mut block);
 
-        for (arg, val) in self.kernel_info.args.iter().zip(&self.values) {
+        let arg_values = self.arg_values();
+        for (arg, val) in self.kernel_info.args.iter().zip(arg_values.iter()) {
             if arg.dead {
                 continue;
             }
@@ -920,7 +916,7 @@ impl Kernel {
             {
                 input.resize(arg.offset, 0);
             }
-            match val.borrow().as_ref().unwrap() {
+            match val.as_ref().unwrap() {
                 KernelArgValue::Constant(c) => input.extend_from_slice(c),
                 KernelArgValue::Buffer(buffer) => {
                     let res = buffer.get_res_of_dev(q.device)?;
@@ -1153,6 +1149,20 @@ impl Kernel {
         }))
     }
 
+    pub fn arg_values(&self) -> MutexGuard<Vec<Option<KernelArgValue>>> {
+        self.values.lock().unwrap()
+    }
+
+    pub fn set_kernel_arg(&self, idx: usize, arg: KernelArgValue) -> CLResult<()> {
+        self.values
+            .lock()
+            .unwrap()
+            .get_mut(idx)
+            .ok_or(CL_INVALID_ARG_INDEX)?
+            .replace(arg);
+        Ok(())
+    }
+
     pub fn access_qualifier(&self, idx: cl_uint) -> cl_kernel_arg_access_qualifier {
         let aq = self.kernel_info.args[idx as usize].spirv.access_qualifier;
 
@@ -1298,7 +1308,7 @@ impl Clone for Kernel {
             base: CLObjectBase::new(RusticlTypes::Kernel),
             prog: self.prog.clone(),
             name: self.name.clone(),
-            values: self.values.clone(),
+            values: Mutex::new(self.arg_values().clone()),
             builds: self.builds.clone(),
             kernel_info: self.kernel_info.clone(),
         }