Support "Unroll" Loop Control via function-scoped `#[spirv(unroll_loo…

…ps)]`. (#337)

* Generalize the zombie serialization system to arbitrary custom decorations.

* Support "Unroll" Loop Control via function-scoped `#[spirv(unroll_loops)]`.

* Pacify the merciless clippy.

crates/rustc_codegen_spirv/src/codegen_cx/declare.rs

@@ -1,6 +1,7 @@
 use super::CodegenCx;
 use crate::abi::ConvSpirvType;
 use crate::builder_spirv::{SpirvConst, SpirvValue, SpirvValueExt};
+use crate::decorations::UnrollLoopsDecoration;
 use crate::spirv_type::SpirvType;
 use crate::symbols::{parse_attrs, SpirvAttribute};
 use rspirv::spirv::{FunctionControl, LinkageType, StorageClass, Word};
@@ -121,6 +122,11 @@ impl<'tcx> CodegenCx<'tcx> {
                         .borrow_mut()
                         .insert(declared);
                 }
+                SpirvAttribute::UnrollLoops => {
+                    self.unroll_loops_decorations
+                        .borrow_mut()
+                        .insert(fn_id, UnrollLoopsDecoration {});
+                }
                 _ => {}
             }
         }

crates/rustc_codegen_spirv/src/codegen_cx/mod.rs

@@ -5,7 +5,9 @@ mod type_;
 
 use crate::builder::{ExtInst, InstructionTable};
 use crate::builder_spirv::{BuilderCursor, BuilderSpirv, SpirvValue, SpirvValueKind};
-use crate::finalizing_passes::export_zombies;
+use crate::decorations::{
+    CustomDecoration, SerializedSpan, UnrollLoopsDecoration, ZombieDecoration,
+};
 use crate::spirv_type::{SpirvType, SpirvTypePrinter, TypeCache};
 use crate::symbols::Symbols;
 use rspirv::dr::{Module, Operand};
@@ -47,7 +49,11 @@ pub struct CodegenCx<'tcx> {
     /// Invalid spir-v IDs that should be stripped from the final binary,
     /// each with its own reason and span that should be used for reporting
     /// (in the event that the value is actually needed)
-    zombie_values: RefCell<HashMap<Word, (&'static str, Span)>>,
+    zombie_decorations: RefCell<HashMap<Word, ZombieDecoration>>,
+    /// Functions that have `#[spirv(unroll_loops)]`, and therefore should
+    /// get `LoopControl::UNROLL` applied to all of their loops' `OpLoopMerge`
+    /// instructions, during structuralization.
+    unroll_loops_decorations: RefCell<HashMap<Word, UnrollLoopsDecoration>>,
     pub kernel_mode: bool,
     /// Cache of all the builtin symbols we need
     pub sym: Box<Symbols>,
@@ -105,7 +111,8 @@ impl<'tcx> CodegenCx<'tcx> {
             type_cache: Default::default(),
             vtables: Default::default(),
             ext_inst: Default::default(),
-            zombie_values: Default::default(),
+            zombie_decorations: Default::default(),
+            unroll_loops_decorations: Default::default(),
             kernel_mode,
             sym,
             instruction_table: InstructionTable::new(),
@@ -153,24 +160,24 @@ impl<'tcx> CodegenCx<'tcx> {
     ///
     /// Finally, if *user* code is marked as zombie, then this means that the user tried to do
     /// something that isn't supported, and should be an error.
-    pub fn zombie_with_span(&self, word: Word, span: Span, reason: &'static str) {
+    pub fn zombie_with_span(&self, word: Word, span: Span, reason: &str) {
         if self.is_system_crate() {
-            self.zombie_values.borrow_mut().insert(word, (reason, span));
+            self.zombie_even_in_user_code(word, span, reason);
         } else {
             self.tcx.sess.span_err(span, reason);
         }
     }
-    pub fn zombie_no_span(&self, word: Word, reason: &'static str) {
-        if self.is_system_crate() {
-            self.zombie_values
-                .borrow_mut()
-                .insert(word, (reason, DUMMY_SP));
-        } else {
-            self.tcx.sess.err(reason);
-        }
+    pub fn zombie_no_span(&self, word: Word, reason: &str) {
+        self.zombie_with_span(word, DUMMY_SP, reason)
     }
-    pub fn zombie_even_in_user_code(&self, word: Word, span: Span, reason: &'static str) {
-        self.zombie_values.borrow_mut().insert(word, (reason, span));
+    pub fn zombie_even_in_user_code(&self, word: Word, span: Span, reason: &str) {
+        self.zombie_decorations.borrow_mut().insert(
+            word,
+            ZombieDecoration {
+                reason: reason.to_string(),
+                span: SerializedSpan::from_rustc(span, self.tcx.sess.source_map()),
+            },
+        );
     }
 
     pub fn is_system_crate(&self) -> bool {
@@ -185,10 +192,17 @@ impl<'tcx> CodegenCx<'tcx> {
 
     pub fn finalize_module(self) -> Module {
         let mut result = self.builder.finalize();
-        export_zombies(
-            &mut result,
-            &self.zombie_values.borrow(),
-            self.tcx.sess.source_map(),
+        result.annotations.extend(
+            self.zombie_decorations
+                .into_inner()
+                .into_iter()
+                .map(|(id, zombie)| zombie.encode(id))
+                .chain(
+                    self.unroll_loops_decorations
+                        .into_inner()
+                        .into_iter()
+                        .map(|(id, unroll_loops)| unroll_loops.encode(id)),
+                ),
         );
         result
     }

crates/rustc_codegen_spirv/src/finalizing_passes.rs → crates/rustc_codegen_spirv/src/decorations.rs

@@ -1,23 +1,131 @@
+//! SPIR-V decorations specific to `rustc_codegen_spirv`, produced during
+//! the original codegen of a crate, and consumed by the `linker`.
+
 use rspirv::dr::{Instruction, Module, Operand};
 use rspirv::spirv::{Decoration, Op, Word};
 use rustc_span::{source_map::SourceMap, FileName, Pos, Span};
 use serde::{Deserialize, Serialize};
-use std::collections::HashMap;
+use std::marker::PhantomData;
 use std::path::PathBuf;
+use std::{iter, slice};
+
+/// Decorations not native to SPIR-V require some form of encoding into existing
+/// SPIR-V constructs, for which we use `OpDecorateString` with decoration type
+/// `UserTypeGOOGLE` and a JSON-encoded Rust value as the decoration string.
+///
+/// Each decoration type has to implement this trait, and use a different
+/// `ENCODING_PREFIX` from any other decoration type, to disambiguate them.
+///
+/// Also, all decorations have to be stripped by the linker at some point,
+/// ideally as soon as they're no longer needed, because no other tools
+/// processing the SPIR-V would understand them correctly.
+///
+/// TODO: uses `non_semantic` instead of piggybacking off of `UserTypeGOOGLE`
+/// <https://htmlpreview.github.io/?https://github.com/KhronosGroup/SPIRV-Registry/blob/master/extensions/KHR/SPV_KHR_non_semantic_info.html>
+pub trait CustomDecoration: for<'de> Deserialize<'de> + Serialize {
+    const ENCODING_PREFIX: &'static str;
+
+    fn encode(self, id: Word) -> Instruction {
+        // FIXME(eddyb) this allocates twice, because there is no functionality
+        // in `serde_json` for writing to something that impls `fmt::Write`,
+        // only for `io::Write`, which would require performing redundant UTF-8
+        // (re)validation, or relying on `unsafe` code, to use with `String`.
+        let json = serde_json::to_string(&self).unwrap();
+        let encoded = [Self::ENCODING_PREFIX, &json].concat();
+
+        Instruction::new(
+            Op::DecorateString,
+            None,
+            None,
+            vec![
+                Operand::IdRef(id),
+                Operand::Decoration(Decoration::UserTypeGOOGLE),
+                Operand::LiteralString(encoded),
+            ],
+        )
+    }
+
+    fn try_decode(inst: &Instruction) -> Option<(Word, LazilyDeserialized<'_, Self>)> {
+        if inst.class.opcode == Op::DecorateString
+            && inst.operands[1].unwrap_decoration() == Decoration::UserTypeGOOGLE
+        {
+            let id = inst.operands[0].unwrap_id_ref();
+            let encoded = inst.operands[2].unwrap_literal_string();
+            let json = encoded.strip_prefix(Self::ENCODING_PREFIX)?;
+
+            Some((
+                id,
+                LazilyDeserialized {
+                    json,
+                    _marker: PhantomData,
+                },
+            ))
+        } else {
+            None
+        }
+    }
+
+    fn decode_all(module: &Module) -> DecodeAllIter<'_, Self> {
+        module
+            .annotations
+            .iter()
+            .filter_map(Self::try_decode as fn(_) -> _)
+    }
+
+    fn remove_all(module: &mut Module) {
+        module
+            .annotations
+            .retain(|inst| Self::try_decode(inst).is_none())
+    }
+}
+
+// HACK(eddyb) return type of `CustomDecoration::decode_all`, in lieu of
+// `-> impl Iterator<Item = (Word, LazilyDeserialized<'_, Self>)` in the trait.
+type DecodeAllIter<'a, D> = iter::FilterMap<
+    slice::Iter<'a, Instruction>,
+    fn(&'a Instruction) -> Option<(Word, LazilyDeserialized<'a, D>)>,
+>;
+
+/// Helper allowing full deserialization to be avoided where possible.
+#[derive(Copy, Clone)]
+pub struct LazilyDeserialized<'a, D> {
+    json: &'a str,
+    _marker: PhantomData<D>,
+}
+
+impl<'a, D: Deserialize<'a>> LazilyDeserialized<'a, D> {
+    pub fn deserialize(self) -> D {
+        serde_json::from_str(self.json).unwrap()
+    }
+}
+
+/// An `OpFunction` with `#[spirv(unroll_loops)]` on the Rust `fn` definition,
+/// which should get `LoopControl::UNROLL` applied to all of its loops'
+/// `OpLoopMerge` instructions, during structuralization.
+#[derive(Deserialize, Serialize)]
+pub struct UnrollLoopsDecoration {}
+
+impl CustomDecoration for UnrollLoopsDecoration {
+    const ENCODING_PREFIX: &'static str = "U";
+}
 
 #[derive(Deserialize, Serialize)]
 pub struct ZombieDecoration {
     pub reason: String,
 
     #[serde(flatten)]
-    pub span: Option<ZombieSpan>,
+    pub span: Option<SerializedSpan>,
+}
+
+impl CustomDecoration for ZombieDecoration {
+    const ENCODING_PREFIX: &'static str = "Z";
 }
 
 /// Representation of a `rustc` `Span` that can be turned into a `Span` again
 /// in another compilation, by reloading the file. However, note that this will
 /// fail if the file changed since, which is detected using the serialized `hash`.
 #[derive(Deserialize, Serialize)]
-pub struct ZombieSpan {
+pub struct SerializedSpan {
     file: PathBuf,
     hash: serde_adapters::SourceFileHash,
     lo: u32,
@@ -65,9 +173,9 @@ mod serde_adapters {
     }
 }
 
-impl ZombieSpan {
-    fn from_rustc(span: Span, source_map: &SourceMap) -> Option<Self> {
-        // Zombies may not always have valid spans.
+impl SerializedSpan {
+    pub fn from_rustc(span: Span, source_map: &SourceMap) -> Option<Self> {
+        // Decorations may not always have valid spans.
         // FIXME(eddyb) reduce the sources of this as much as possible.
         if span.is_dummy() {
             return None;
@@ -108,7 +216,7 @@ impl ZombieSpan {
             return None;
         }
 
-        // Sanity check - assuming `ZombieSpan` isn't corrupted, this assert
+        // Sanity check - assuming `SerializedSpan` isn't corrupted, this assert
         // could only ever fail because of a hash collision.
         assert!(self.lo <= self.hi && self.hi <= file.byte_length());
 
@@ -118,32 +226,3 @@ impl ZombieSpan {
         ))
     }
 }
-
-pub fn export_zombies(
-    module: &mut Module,
-    zombies: &HashMap<Word, (&'static str, Span)>,
-    source_map: &SourceMap,
-) {
-    for (&id, &(reason, span)) in zombies {
-        let encoded = serde_json::to_string(&ZombieDecoration {
-            reason: reason.to_string(),
-            span: ZombieSpan::from_rustc(span, source_map),
-        })
-        .unwrap();
-
-        // TODO: Right now we just piggyback off UserTypeGOOGLE since we never use it elsewhere. We should, uh, fix this
-        // to use non_semantic or something.
-        // https://htmlpreview.github.io/?https://github.com/KhronosGroup/SPIRV-Registry/blob/master/extensions/KHR/SPV_KHR_non_semantic_info.html
-        let inst = Instruction::new(
-            Op::DecorateString,
-            None,
-            None,
-            vec![
-                Operand::IdRef(id),
-                Operand::Decoration(Decoration::UserTypeGOOGLE),
-                Operand::LiteralString(encoded),
-            ],
-        );
-        module.annotations.push(inst);
-    }
-}

crates/rustc_codegen_spirv/src/lib.rs

@@ -84,7 +84,7 @@ mod abi;
 mod builder;
 mod builder_spirv;
 mod codegen_cx;
-mod finalizing_passes;
+mod decorations;
 mod link;
 mod linker;
 mod spirv_type;

crates/rustc_codegen_spirv/src/linker/duplicates.rs

@@ -1,3 +1,4 @@
+use crate::decorations::{CustomDecoration, ZombieDecoration};
 use rspirv::binary::Assemble;
 use rspirv::dr::{Instruction, Module, Operand};
 use rspirv::spirv::{Op, Word};
@@ -154,7 +155,7 @@ pub fn remove_duplicate_types(module: &mut Module) {
     // Keep in mind, this algorithm requires forward type references to not exist - i.e. it's a valid spir-v module.
 
     // Include zombies in the key to not merge zombies with non-zombies
-    let zombies: HashSet<Word> = super::zombies::collect_zombies(module)
+    let zombies: HashSet<Word> = ZombieDecoration::decode_all(module)
         .map(|(z, _)| z)
         .collect();
 

crates/rustc_codegen_spirv/src/linker/mod.rs

@@ -12,6 +12,7 @@ mod simple_passes;
 mod structurizer;
 mod zombies;
 
+use crate::decorations::{CustomDecoration, UnrollLoopsDecoration};
 use rspirv::binary::Consumer;
 use rspirv::dr::{Block, Instruction, Loader, Module, ModuleHeader};
 use rspirv::spirv::{Op, Word};
@@ -136,12 +137,17 @@ pub fn link(sess: &Session, mut inputs: Vec<Module>, opts: &Options) -> Result<M
         dce::dce(&mut output);
     }
 
+    let unroll_loops_decorations = UnrollLoopsDecoration::decode_all(&output)
+        .map(|(id, unroll_loops)| (id, unroll_loops.deserialize()))
+        .collect::<HashMap<_, _>>();
+    UnrollLoopsDecoration::remove_all(&mut output);
+
     let mut output = if opts.structurize {
         let _timer = sess.timer("link_structurize");
         if opts.use_new_structurizer {
-            new_structurizer::structurize(output)
+            new_structurizer::structurize(output, unroll_loops_decorations)
         } else {
-            structurizer::structurize(sess, output)
+            structurizer::structurize(sess, output, unroll_loops_decorations)
         }
     } else {
         output