rusticl/program: fix building kernels

We ended up with duplicates, but also rebuilt the same kernel over and
over again for multi dev builds.

Cc: mesa-stable
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/33892>

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

@@ -102,16 +102,27 @@ impl ProgramBuild {
         self.dev_build(d).spirv.as_ref()?.kernel_info(kernel)
     }
 
-    fn args(&self, dev: &Device, kernel: &str) -> Vec<spirv::SPIRVKernelArg> {
-        self.dev_build(dev).spirv.as_ref().unwrap().args(kernel)
+    fn args(&self, dev: &Device, kernel: &str) -> Option<Vec<spirv::SPIRVKernelArg>> {
+        self.dev_build(dev).spirv.as_ref().map(|s| s.args(kernel))
     }
 
-    fn build_nirs(&mut self, devs: &[&'static Device], is_src: bool) {
-        for kernel_name in &self.kernels.clone() {
-            let kernel_args: HashSet<_> = self
-                .devs_with_build()
+    fn rebuild_kernels(&mut self, devs: &[&'static Device], is_src: bool) {
+        let mut kernels: Vec<_> = self
+            .builds
+            .values()
+            .filter_map(|b| b.spirv.as_ref())
+            .flat_map(|s| s.kernels())
+            .collect();
+
+        kernels.sort();
+        kernels.dedup();
+
+        self.kernels = kernels;
+
+        for kernel_name in &self.kernels {
+            let kernel_args: HashSet<_> = devs
                 .iter()
-                .map(|d| self.args(d, kernel_name))
+                .filter_map(|d| self.args(d, kernel_name))
                 .collect();
 
             let args = kernel_args.into_iter().next().unwrap();
@@ -119,6 +130,14 @@ impl ProgramBuild {
 
             // TODO: we could run this in parallel?
             for dev in devs {
+                let Some(build) = self.builds.get_mut(dev) else {
+                    continue;
+                };
+
+                if !build.is_success() {
+                    continue;
+                }
+
                 let build_result = convert_spirv_to_nir(self, kernel_name, &args, dev);
                 kernel_info_set.insert(build_result.kernel_info);
 
@@ -151,14 +170,6 @@ impl ProgramBuild {
         self.builds.get_mut(dev).unwrap()
     }
 
-    fn devs_with_build(&self) -> Vec<&'static Device> {
-        self.builds
-            .iter()
-            .filter(|(_, build)| build.status == CL_BUILD_SUCCESS as cl_build_status)
-            .map(|(&d, _)| d)
-            .collect()
-    }
-
     pub fn hash_key(&self, dev: &Device, name: &str) -> Option<cache_key> {
         if let Some(cache) = dev.screen().shader_cache() {
             let info = self.dev_build(dev);
@@ -388,22 +399,14 @@ impl Program {
         bins: &[&[u8]],
     ) -> Result<Arc<Program>, Vec<cl_int>> {
         let mut builds = HashMap::new();
-        let mut kernels = HashSet::new();
-
         let mut errors = vec![CL_SUCCESS as cl_int; devs.len()];
         for (idx, (&d, b)) in devs.iter().zip(bins).enumerate() {
             let build = match Self::spirv_from_bin_for_dev(d, b) {
-                Ok((spirv, bin_type)) => {
-                    for k in spirv.kernels() {
-                        kernels.insert(k);
-                    }
-
-                    ProgramDevBuild {
-                        spirv: Some(spirv),
-                        bin_type: bin_type,
-                        ..Default::default()
-                    }
-                }
+                Ok((spirv, bin_type)) => ProgramDevBuild {
+                    spirv: Some(spirv),
+                    bin_type: bin_type,
+                    ..Default::default()
+                },
                 Err(err) => {
                     errors[idx] = err;
                     ProgramDevBuild {
@@ -423,10 +426,10 @@ impl Program {
         let mut build = ProgramBuild {
             builds: builds,
             spec_constants: HashMap::new(),
-            kernels: kernels.into_iter().collect(),
+            kernels: Vec::new(),
             kernel_info: HashMap::new(),
         };
-        build.build_nirs(&devs, false);
+        build.rebuild_kernels(&devs, false);
 
         Ok(Arc::new(Self {
             base: CLObjectBase::new(RusticlTypes::Program),
@@ -579,16 +582,13 @@ impl Program {
             d.log.push_str(&log);
             d.spirv = spirv;
 
-            if let Some(spirv) = &d.spirv {
+            if d.spirv.is_some() {
                 d.bin_type = if lib {
                     CL_PROGRAM_BINARY_TYPE_LIBRARY
                 } else {
                     CL_PROGRAM_BINARY_TYPE_EXECUTABLE
                 };
                 d.status = CL_BUILD_SUCCESS as cl_build_status;
-                let mut kernels = spirv.kernels();
-                info.kernels.append(&mut kernels);
-                info.build_nirs(devs, self.is_src());
             } else {
                 d.status = CL_BUILD_ERROR;
                 d.bin_type = CL_PROGRAM_BINARY_TYPE_NONE;
@@ -596,6 +596,8 @@ impl Program {
             }
         }
 
+        info.rebuild_kernels(devs, self.is_src());
+
         res
     }
 
@@ -682,7 +684,6 @@ impl Program {
         options: String,
     ) -> Arc<Program> {
         let mut builds = HashMap::new();
-        let mut kernels = HashSet::new();
         let mut locks: Vec<_> = progs.iter().map(|p| p.build_info()).collect();
         let lib = options.contains("-create-library");
 
@@ -707,10 +708,7 @@ impl Program {
 
             let status;
             let bin_type;
-            if let Some(spirv) = &spirv {
-                for k in spirv.kernels() {
-                    kernels.insert(k);
-                }
+            if spirv.is_some() {
                 status = CL_BUILD_SUCCESS as cl_build_status;
                 bin_type = if lib {
                     CL_PROGRAM_BINARY_TYPE_LIBRARY
@@ -737,12 +735,12 @@ impl Program {
         let mut build = ProgramBuild {
             builds: builds,
             spec_constants: HashMap::new(),
-            kernels: kernels.into_iter().collect(),
+            kernels: Vec::new(),
             kernel_info: HashMap::new(),
         };
 
         // Pre build nir kernels
-        build.build_nirs(devs, false);
+        build.rebuild_kernels(devs, false);
 
         Arc::new(Self {
             base: CLObjectBase::new(RusticlTypes::Program),