rusticl/kernel: support for images

Signed-off-by: Karol Herbst <kherbst@redhat.com>
Acked-by: Alyssa Rosenzweig <alyssa.rosenzweig@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/15439>

src/gallium/frontends/rusticl/api/kernel.rs

@@ -290,6 +290,10 @@ pub fn set_kernel_arg(
                         }
                     }
                     KernelArgType::MemLocal => KernelArgValue::LocalMem(arg_size),
+                    KernelArgType::Image | KernelArgType::Texture => {
+                        let img: *const cl_mem = arg_value.cast();
+                        KernelArgValue::MemObject((*img).get_ref()?)
+                    }
                     KernelArgType::Sampler => {
                         let ptr: *const cl_sampler = arg_value.cast();
                         KernelArgValue::Sampler((*ptr).get_ref()?)

src/gallium/frontends/rusticl/core/kernel.rs

@@ -20,6 +20,7 @@ use std::cell::RefCell;
 use std::collections::HashMap;
 use std::collections::HashSet;
 use std::convert::TryInto;
+use std::os::raw::c_void;
 use std::ptr;
 use std::slice;
 use std::sync::Arc;
@@ -37,7 +38,9 @@ pub enum KernelArgValue {
 #[derive(PartialEq, Eq, Clone)]
 pub enum KernelArgType {
     Constant, // for anything passed by value
+    Image,
     Sampler,
+    Texture,
     MemGlobal,
     MemConstant,
     MemLocal,
@@ -93,7 +96,15 @@ impl KernelArg {
                     KernelArgType::MemLocal
                 }
                 clc_kernel_arg_address_qualifier::CLC_KERNEL_ARG_ADDRESS_GLOBAL => {
-                    KernelArgType::MemGlobal
+                    if unsafe { glsl_type_is_image(nir.type_) } {
+                        if nir.data.access() == gl_access_qualifier::ACCESS_NON_WRITEABLE.0 {
+                            KernelArgType::Texture
+                        } else {
+                            KernelArgType::Image
+                        }
+                    } else {
+                        KernelArgType::MemGlobal
+                    }
                 }
             };
 
@@ -229,6 +240,13 @@ fn lower_and_optimize_nir_pre_inputs(dev: &Device, nir: &mut NirShader, lib_clc:
     nir.pass0(nir_opt_deref);
 }
 
+extern "C" fn can_remove_var(var: *mut nir_variable, _: *mut c_void) -> bool {
+    unsafe {
+        let var = var.as_ref().unwrap();
+        !glsl_type_is_image(var.type_)
+    }
+}
+
 fn lower_and_optimize_nir_late(
     dev: &Device,
     nir: &mut NirShader,
@@ -242,19 +260,24 @@ fn lower_and_optimize_nir_late(
     };
     let mut lower_state = rusticl_lower_state::default();
 
+    let dv_opts = nir_remove_dead_variables_options {
+        can_remove_var: Some(can_remove_var),
+        can_remove_var_data: ptr::null_mut(),
+    };
     nir.pass2(
         nir_remove_dead_variables,
         nir_variable_mode::nir_var_uniform
+            | nir_variable_mode::nir_var_image
             | nir_variable_mode::nir_var_mem_constant
+            | nir_variable_mode::nir_var_mem_shared
             | nir_variable_mode::nir_var_function_temp,
-        ptr::null(),
+        &dv_opts,
     );
+
+    // TODO inline samplers
     nir.pass1(nir_lower_readonly_images_to_tex, false);
-    nir.pass2(
-        nir_remove_dead_variables,
-        nir_variable_mode::nir_var_mem_shared | nir_variable_mode::nir_var_function_temp,
-        ptr::null(),
-    );
+    nir.pass0(nir_lower_cl_images);
+
     nir.reset_scratch_size();
     nir.pass2(
         nir_lower_vars_to_explicit_types,
@@ -418,6 +441,9 @@ impl Kernel {
         let mut resource_info = Vec::new();
         let mut local_size: u32 = nir.shared_size();
         let printf_size = q.device.printf_buffer_size() as u32;
+        let mut samplers = Vec::new();
+        let mut iviews = Vec::new();
+        let mut sviews = Vec::new();
 
         for i in 0..3 {
             if block[i] == 0 {
@@ -431,18 +457,34 @@ impl Kernel {
             if arg.dead {
                 continue;
             }
-            input.append(&mut vec![0; arg.offset - input.len()]);
+
+            if arg.kind != KernelArgType::Image
+                && arg.kind != KernelArgType::Texture
+                && arg.kind != KernelArgType::Sampler
+            {
+                input.resize(arg.offset, 0);
+            }
             match val.borrow().as_ref().unwrap() {
                 KernelArgValue::Constant(c) => input.extend_from_slice(c),
                 KernelArgValue::MemObject(mem) => {
-                    input.extend_from_slice(&mem.offset.to_ne_bytes());
-                    resource_info.push((Some(mem.get_res_of_dev(&q.device)?.clone()), arg.offset));
+                    let res = mem.get_res_of_dev(&q.device)?;
+                    if mem.is_buffer() {
+                        input.extend_from_slice(&mem.offset.to_ne_bytes());
+                        resource_info.push((Some(res.clone()), arg.offset));
+                    } else if arg.kind == KernelArgType::Image {
+                        iviews.push(res.pipe_image_view())
+                    } else {
+                        sviews.push(res.clone())
+                    }
                 }
                 KernelArgValue::LocalMem(size) => {
                     // TODO 32 bit
                     input.extend_from_slice(&[0; 8]);
                     local_size += *size as u32;
                 }
+                KernelArgValue::Sampler(sampler) => {
+                    samplers.push(sampler.pipe());
+                }
                 KernelArgValue::None => {
                     assert!(
                         arg.kind == KernelArgType::MemGlobal
@@ -450,7 +492,6 @@ impl Kernel {
                     );
                     input.extend_from_slice(&[0; 8]);
                 }
-                _ => panic!("unhandled arg type"),
             }
         }
 
@@ -498,6 +539,13 @@ impl Kernel {
             let mut globals: Vec<*mut u32> = Vec::new();
             let printf_format = nir.printf_format();
             let printf_buf = printf_buf.clone();
+            let iviews = iviews.clone();
+
+            let mut sviews: Vec<_> = sviews.iter().map(|s| ctx.create_sampler_view(s)).collect();
+            let samplers: Vec<_> = samplers
+                .iter()
+                .map(|s| ctx.create_sampler_state(s))
+                .collect();
 
             for (res, offset) in resource_info.clone() {
                 resources.push(res);
@@ -516,12 +564,23 @@ impl Kernel {
             let cso = ctx.create_compute_state(nir, input.len() as u32, local_size);
 
             ctx.bind_compute_state(cso);
+            ctx.bind_sampler_states(&samplers);
+            ctx.set_sampler_views(&mut sviews);
+            ctx.set_shader_images(&iviews);
             ctx.set_global_binding(resources.as_slice(), &mut globals);
+
             ctx.launch_grid(work_dim, block, grid, &input);
+
             ctx.clear_global_binding(globals.len() as u32);
+            ctx.clear_shader_images(iviews.len() as u32);
+            ctx.clear_sampler_views(sviews.len() as u32);
+            ctx.clear_sampler_states(samplers.len() as u32);
             ctx.delete_compute_state(cso);
             ctx.memory_barrier(PIPE_BARRIER_GLOBAL_BUFFER);
 
+            samplers.iter().for_each(|s| ctx.delete_sampler_state(*s));
+            sviews.iter().for_each(|v| ctx.sampler_view_destroy(*v));
+
             if let Some(printf_buf) = &printf_buf {
                 let tx = ctx
                     .buffer_map(printf_buf, 0, printf_size as i32, true)

src/gallium/frontends/rusticl/core/memory.rs

@@ -877,4 +877,32 @@ impl Sampler {
             props: props,
         })
     }
+
+    pub fn pipe(&self) -> pipe_sampler_state {
+        let mut res = pipe_sampler_state::default();
+
+        let wrap = match self.addressing_mode {
+            CL_ADDRESS_CLAMP_TO_EDGE => pipe_tex_wrap::PIPE_TEX_WRAP_CLAMP_TO_EDGE,
+            CL_ADDRESS_CLAMP => pipe_tex_wrap::PIPE_TEX_WRAP_CLAMP_TO_BORDER,
+            CL_ADDRESS_REPEAT => pipe_tex_wrap::PIPE_TEX_WRAP_REPEAT,
+            CL_ADDRESS_MIRRORED_REPEAT => pipe_tex_wrap::PIPE_TEX_WRAP_MIRROR_REPEAT,
+            // TODO: what's a reasonable default?
+            _ => pipe_tex_wrap::PIPE_TEX_WRAP_CLAMP_TO_EDGE,
+        };
+
+        let img_filter = match self.filter_mode {
+            CL_FILTER_NEAREST => pipe_tex_filter::PIPE_TEX_FILTER_NEAREST,
+            CL_FILTER_LINEAR => pipe_tex_filter::PIPE_TEX_FILTER_LINEAR,
+            _ => panic!("unkown filter_mode"),
+        };
+
+        res.set_min_img_filter(img_filter);
+        res.set_mag_img_filter(img_filter);
+        res.set_normalized_coords(self.normalized_coords.into());
+        res.set_wrap_r(wrap);
+        res.set_wrap_s(wrap);
+        res.set_wrap_t(wrap);
+
+        res
+    }
 }

src/gallium/frontends/rusticl/mesa/compiler/clc/spirv.rs

@@ -189,6 +189,7 @@ impl SPIRVBin {
                 kernel: true,
                 kernel_image: true,
                 linkage: true,
+                literal_sampler: true,
                 printf: true,
                 ..Default::default()
             },

src/gallium/frontends/rusticl/mesa/pipe/context.rs

@@ -222,6 +222,39 @@ impl PipeContext {
         unsafe { self.pipe.as_ref().delete_compute_state.unwrap()(self.pipe.as_ptr(), state) }
     }
 
+    pub fn create_sampler_state(&self, state: &pipe_sampler_state) -> *mut c_void {
+        unsafe { self.pipe.as_ref().create_sampler_state.unwrap()(self.pipe.as_ptr(), state) }
+    }
+
+    pub fn bind_sampler_states(&self, samplers: &[*mut c_void]) {
+        let mut samplers = samplers.to_owned();
+        unsafe {
+            self.pipe.as_ref().bind_sampler_states.unwrap()(
+                self.pipe.as_ptr(),
+                pipe_shader_type::PIPE_SHADER_COMPUTE,
+                0,
+                samplers.len() as u32,
+                samplers.as_mut_ptr(),
+            )
+        }
+    }
+
+    pub fn clear_sampler_states(&self, count: u32) {
+        unsafe {
+            self.pipe.as_ref().bind_sampler_states.unwrap()(
+                self.pipe.as_ptr(),
+                pipe_shader_type::PIPE_SHADER_COMPUTE,
+                0,
+                count,
+                ptr::null_mut(),
+            )
+        }
+    }
+
+    pub fn delete_sampler_state(&self, ptr: *mut c_void) {
+        unsafe { self.pipe.as_ref().delete_sampler_state.unwrap()(self.pipe.as_ptr(), ptr) }
+    }
+
     pub fn launch_grid(&self, work_dim: u32, block: [u32; 3], grid: [u32; 3], input: &[u8]) {
         let info = pipe_grid_info {
             pc: 0,
@@ -253,6 +286,17 @@ impl PipeContext {
         }
     }
 
+    pub fn create_sampler_view(&self, res: &PipeResource) -> *mut pipe_sampler_view {
+        let template = res.pipe_sampler_view_template();
+        unsafe {
+            self.pipe.as_ref().create_sampler_view.unwrap()(
+                self.pipe.as_ptr(),
+                res.pipe(),
+                &template,
+            )
+        }
+    }
+
     pub fn clear_global_binding(&self, count: u32) {
         unsafe {
             self.pipe.as_ref().set_global_binding.unwrap()(
@@ -265,6 +309,64 @@ impl PipeContext {
         }
     }
 
+    pub fn set_sampler_views(&self, views: &mut [*mut pipe_sampler_view]) {
+        unsafe {
+            self.pipe.as_ref().set_sampler_views.unwrap()(
+                self.pipe.as_ptr(),
+                pipe_shader_type::PIPE_SHADER_COMPUTE,
+                0,
+                views.len() as u32,
+                0,
+                false,
+                views.as_mut_ptr(),
+            )
+        }
+    }
+
+    pub fn clear_sampler_views(&self, count: u32) {
+        unsafe {
+            self.pipe.as_ref().set_sampler_views.unwrap()(
+                self.pipe.as_ptr(),
+                pipe_shader_type::PIPE_SHADER_COMPUTE,
+                0,
+                count,
+                0,
+                false,
+                ptr::null_mut(),
+            )
+        }
+    }
+
+    pub fn sampler_view_destroy(&self, view: *mut pipe_sampler_view) {
+        unsafe { self.pipe.as_ref().sampler_view_destroy.unwrap()(self.pipe.as_ptr(), view) }
+    }
+
+    pub fn set_shader_images(&self, images: &[pipe_image_view]) {
+        unsafe {
+            self.pipe.as_ref().set_shader_images.unwrap()(
+                self.pipe.as_ptr(),
+                pipe_shader_type::PIPE_SHADER_COMPUTE,
+                0,
+                images.len() as u32,
+                0,
+                images.as_ptr(),
+            )
+        }
+    }
+
+    pub fn clear_shader_images(&self, count: u32) {
+        unsafe {
+            self.pipe.as_ref().set_shader_images.unwrap()(
+                self.pipe.as_ptr(),
+                pipe_shader_type::PIPE_SHADER_COMPUTE,
+                0,
+                count,
+                0,
+                ptr::null_mut(),
+            )
+        }
+    }
+
     pub fn memory_barrier(&self, barriers: u32) {
         unsafe { self.pipe.as_ref().memory_barrier.unwrap()(self.pipe.as_ptr(), barriers) }
     }
@@ -289,6 +391,7 @@ impl Drop for PipeContext {
 fn has_required_cbs(c: &pipe_context) -> bool {
     c.destroy.is_some()
         && c.bind_compute_state.is_some()
+        && c.bind_sampler_states.is_some()
         && c.blit.is_some()
         && c.buffer_map.is_some()
         && c.buffer_subdata.is_some()
@@ -297,11 +400,15 @@ fn has_required_cbs(c: &pipe_context) -> bool {
         && c.clear_texture.is_some()
         && c.create_compute_state.is_some()
         && c.delete_compute_state.is_some()
+        && c.delete_sampler_state.is_some()
         && c.flush.is_some()
         && c.launch_grid.is_some()
         && c.memory_barrier.is_some()
         && c.resource_copy_region.is_some()
+        && c.sampler_view_destroy.is_some()
         && c.set_global_binding.is_some()
+        && c.set_sampler_views.is_some()
+        && c.set_shader_images.is_some()
         && c.texture_map.is_some()
         && c.texture_subdata.is_some()
         && c.texture_unmap.is_some()

src/gallium/frontends/rusticl/mesa/pipe/resource.rs

@@ -20,6 +20,50 @@ impl PipeResource {
     pub(super) fn pipe(&self) -> *mut pipe_resource {
         self.pipe
     }
+
+    fn as_ref(&self) -> &pipe_resource {
+        unsafe { self.pipe.as_ref().unwrap() }
+    }
+
+    pub fn pipe_image_view(&self) -> pipe_image_view {
+        let u = if self.as_ref().target() == pipe_texture_target::PIPE_BUFFER {
+            pipe_image_view__bindgen_ty_1 {
+                buf: pipe_image_view__bindgen_ty_1__bindgen_ty_2 {
+                    offset: 0,
+                    size: self.as_ref().width0,
+                },
+            }
+        } else {
+            let mut tex = pipe_image_view__bindgen_ty_1__bindgen_ty_1::default();
+            tex.set_level(0);
+            tex.set_first_layer(0);
+            if self.as_ref().target() == pipe_texture_target::PIPE_TEXTURE_3D {
+                tex.set_last_layer((self.as_ref().depth0 - 1).into());
+            } else if self.as_ref().array_size > 0 {
+                tex.set_last_layer((self.as_ref().array_size - 1).into());
+            } else {
+                tex.set_last_layer(0);
+            }
+
+            pipe_image_view__bindgen_ty_1 { tex: tex }
+        };
+
+        pipe_image_view {
+            resource: self.pipe(),
+            format: self.as_ref().format(),
+            access: 0,
+            shader_access: PIPE_IMAGE_ACCESS_WRITE as u16,
+            u: u,
+        }
+    }
+
+    pub fn pipe_sampler_view_template(&self) -> pipe_sampler_view {
+        let mut res = pipe_sampler_view::default();
+        unsafe {
+            u_sampler_view_default_template(&mut res, self.pipe, self.as_ref().format());
+        }
+        res
+    }
 }
 
 impl Drop for PipeResource {

src/gallium/frontends/rusticl/meson.build

@@ -214,6 +214,12 @@ rusticl_mesa_bindings_rs = rust.bindgen(
     '--allowlist-var', 'stderr',
     '--allowlist-var', 'stdout',
     '--bitfield-enum', 'nir_lower_int64_options',
+    '--allowlist-type', 'pipe_tex_wrap',
+    '--constified-enum-module', 'pipe_tex_wrap',
+    '--allowlist-type', 'pipe_tex_filter',
+    '--constified-enum-module', 'pipe_tex_filter',
+    '--allowlist-type', 'gl_access_qualifier',
+    '--bitfield-enum', 'gl_access_qualifier',
   ],
 )
 

src/gallium/frontends/rusticl/rusticl_mesa_bindings.h

@@ -2,6 +2,7 @@
 
 #include "compiler/clc/clc.h"
 #include "compiler/clc/clc_helpers.h"
+#include "compiler/shader_enums.h"
 #include "nir_types.h"
 #include "spirv/nir_spirv.h"
 
@@ -12,5 +13,6 @@
 #include "pipe-loader/pipe_loader.h"
 
 #include "util/u_printf.h"
+#include "util/u_sampler.h"
 
 #include "rusticl_nir.h"