From ebf3dbee8d6b8d849ecd1f17d1a3689ed8d6a2b1 Mon Sep 17 00:00:00 2001 From: Eduard-Mihai Burtescu Date: Thu, 10 Dec 2020 21:57:15 +0200 Subject: [PATCH] Support "Unroll" Loop Control via function-scoped `#[spirv(unroll_loops)]`. (#337) * Generalize the zombie serialization system to arbitrary custom decorations. * Support "Unroll" Loop Control via function-scoped `#[spirv(unroll_loops)]`. * Pacify the merciless clippy. --- .../src/codegen_cx/declare.rs | 6 + .../rustc_codegen_spirv/src/codegen_cx/mod.rs | 52 +++--- .../{finalizing_passes.rs => decorations.rs} | 151 +++++++++++++----- crates/rustc_codegen_spirv/src/lib.rs | 2 +- .../src/linker/duplicates.rs | 3 +- crates/rustc_codegen_spirv/src/linker/mod.rs | 10 +- .../src/linker/new_structurizer.rs | 20 ++- .../src/linker/structurizer.rs | 24 ++- .../rustc_codegen_spirv/src/linker/zombies.rs | 36 +---- crates/rustc_codegen_spirv/src/symbols.rs | 2 + crates/spirv-builder/src/test/basic.rs | 62 +++++++ 11 files changed, 274 insertions(+), 94 deletions(-) rename crates/rustc_codegen_spirv/src/{finalizing_passes.rs => decorations.rs} (50%) diff --git a/crates/rustc_codegen_spirv/src/codegen_cx/declare.rs b/crates/rustc_codegen_spirv/src/codegen_cx/declare.rs index 7bfa98631b..6d0156fe29 100644 --- a/crates/rustc_codegen_spirv/src/codegen_cx/declare.rs +++ b/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 {}); + } _ => {} } } diff --git a/crates/rustc_codegen_spirv/src/codegen_cx/mod.rs b/crates/rustc_codegen_spirv/src/codegen_cx/mod.rs index 71f6220453..3158d4c433 100644 --- a/crates/rustc_codegen_spirv/src/codegen_cx/mod.rs +++ b/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>, + zombie_decorations: RefCell>, + /// 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>, pub kernel_mode: bool, /// Cache of all the builtin symbols we need pub sym: Box, @@ -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 } diff --git a/crates/rustc_codegen_spirv/src/finalizing_passes.rs b/crates/rustc_codegen_spirv/src/decorations.rs similarity index 50% rename from crates/rustc_codegen_spirv/src/finalizing_passes.rs rename to crates/rustc_codegen_spirv/src/decorations.rs index 1ed38b71dd..0646f14d1e 100644 --- a/crates/rustc_codegen_spirv/src/finalizing_passes.rs +++ b/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` +/// +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)` 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, +} + +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, + pub span: Option, +} + +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 { - // Zombies may not always have valid spans. +impl SerializedSpan { + pub fn from_rustc(span: Span, source_map: &SourceMap) -> Option { + // 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, - 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); - } -} diff --git a/crates/rustc_codegen_spirv/src/lib.rs b/crates/rustc_codegen_spirv/src/lib.rs index b4f18a54a2..4e47f36228 100644 --- a/crates/rustc_codegen_spirv/src/lib.rs +++ b/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; diff --git a/crates/rustc_codegen_spirv/src/linker/duplicates.rs b/crates/rustc_codegen_spirv/src/linker/duplicates.rs index a1771066d4..6faa3b4d8d 100644 --- a/crates/rustc_codegen_spirv/src/linker/duplicates.rs +++ b/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 = super::zombies::collect_zombies(module) + let zombies: HashSet = ZombieDecoration::decode_all(module) .map(|(z, _)| z) .collect(); diff --git a/crates/rustc_codegen_spirv/src/linker/mod.rs b/crates/rustc_codegen_spirv/src/linker/mod.rs index bc36c6a4d5..1d434ae11f 100644 --- a/crates/rustc_codegen_spirv/src/linker/mod.rs +++ b/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, opts: &Options) -> Result>(); + 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 diff --git a/crates/rustc_codegen_spirv/src/linker/new_structurizer.rs b/crates/rustc_codegen_spirv/src/linker/new_structurizer.rs index cd55630f1c..e3079ae8e5 100644 --- a/crates/rustc_codegen_spirv/src/linker/new_structurizer.rs +++ b/crates/rustc_codegen_spirv/src/linker/new_structurizer.rs @@ -1,3 +1,4 @@ +use crate::decorations::UnrollLoopsDecoration; use indexmap::{indexmap, IndexMap}; use rspirv::dr::{Block, Builder, Function, InsertPoint, Module, Operand}; use rspirv::spirv::{LoopControl, Op, SelectionControl, Word}; @@ -30,7 +31,10 @@ impl FuncBuilder<'_> { } } -pub fn structurize(module: Module) -> Module { +pub fn structurize( + module: Module, + unroll_loops_decorations: HashMap, +) -> Module { let mut builder = Builder::new_from_module(module); for func_idx in 0..builder.module_ref().functions.len() { @@ -39,6 +43,13 @@ pub fn structurize(module: Module) -> Module { builder: &mut builder, }; + let func_id = func.function().def.as_ref().unwrap().result_id.unwrap(); + + let loop_control = match unroll_loops_decorations.get(&func_id) { + Some(UnrollLoopsDecoration {}) => LoopControl::UNROLL, + None => LoopControl::NONE, + }; + let block_id_to_idx = func .blocks() .iter() @@ -49,6 +60,7 @@ pub fn structurize(module: Module) -> Module { Structurizer { func, block_id_to_idx, + loop_control, incoming_edge_count: vec![], regions: HashMap::new(), } @@ -90,6 +102,10 @@ struct Structurizer<'a> { func: FuncBuilder<'a>, block_id_to_idx: HashMap, + /// `LoopControl` to use in all loops' `OpLoopMerge` instruction. + /// Currently only affected by function-scoped `#[spirv(unroll_loops)]`. + loop_control: LoopControl, + /// Number of edges pointing to each block. /// Computed by `post_order` and updated when structuring loops /// (backedge count is subtracted to hide them from outer regions). @@ -313,7 +329,7 @@ impl Structurizer<'_> { .loop_merge( while_exit_block_id, while_body_merge_id, - LoopControl::NONE, + self.loop_control, iter::empty(), ) .unwrap(); diff --git a/crates/rustc_codegen_spirv/src/linker/structurizer.rs b/crates/rustc_codegen_spirv/src/linker/structurizer.rs index e0a64cabfb..23be88af05 100644 --- a/crates/rustc_codegen_spirv/src/linker/structurizer.rs +++ b/crates/rustc_codegen_spirv/src/linker/structurizer.rs @@ -1,6 +1,7 @@ // This pass inserts merge instructions for structured control flow with the assumption the spir-v is reducible. use super::simple_passes::outgoing_edges; +use crate::decorations::UnrollLoopsDecoration; use rspirv::spirv::{Op, SelectionControl, Word}; use rspirv::{ dr::{Block, Builder, InsertPoint, Module, Operand}, @@ -123,7 +124,11 @@ impl ControlFlowInfo { } } -pub fn structurize(sess: &Session, module: Module) -> Module { +pub fn structurize( + sess: &Session, + module: Module, + unroll_loops_decorations: HashMap, +) -> Module { let mut builder = Builder::new_from_module(module); for func_idx in 0..builder.module_ref().functions.len() { @@ -131,7 +136,19 @@ pub fn structurize(sess: &Session, module: Module) -> Module { builder.select_function(Some(func_idx)).unwrap(); - insert_loop_merge_on_conditional_branch(&mut builder, &mut cf_info); + let func_id = builder.module_ref().functions[func_idx] + .def + .as_ref() + .unwrap() + .result_id + .unwrap(); + + let loop_control = match unroll_loops_decorations.get(&func_id) { + Some(UnrollLoopsDecoration {}) => LoopControl::UNROLL, + None => LoopControl::NONE, + }; + + insert_loop_merge_on_conditional_branch(&mut builder, &mut cf_info, loop_control); retarget_loop_children_if_needed(&mut builder, &cf_info); insert_selection_merge_on_conditional_branch(sess, &mut builder, &mut cf_info); defer_loop_internals(&mut builder, &cf_info); @@ -763,6 +780,7 @@ pub fn insert_selection_merge_on_conditional_branch( pub fn insert_loop_merge_on_conditional_branch( builder: &mut Builder, cf_info: &mut ControlFlowInfo, + loop_control: LoopControl, ) { let mut branch_conditional_ops = Vec::new(); @@ -825,7 +843,7 @@ pub fn insert_loop_merge_on_conditional_branch( InsertPoint::FromEnd(1), merge_block_id, continue_block_id, - LoopControl::NONE, + loop_control, None, ) .unwrap(); diff --git a/crates/rustc_codegen_spirv/src/linker/zombies.rs b/crates/rustc_codegen_spirv/src/linker/zombies.rs index 1b4aecb033..9a5c621e06 100644 --- a/crates/rustc_codegen_spirv/src/linker/zombies.rs +++ b/crates/rustc_codegen_spirv/src/linker/zombies.rs @@ -1,38 +1,14 @@ //! See documentation on `CodegenCx::zombie` for a description of the zombie system. -use crate::finalizing_passes::ZombieDecoration; +use crate::decorations::{CustomDecoration, ZombieDecoration}; use rspirv::dr::{Instruction, Module}; -use rspirv::spirv::{Decoration, Op, Word}; +use rspirv::spirv::{Op, Word}; use rustc_session::Session; use rustc_span::{Span, DUMMY_SP}; use std::collections::HashMap; use std::env; use std::iter::once; -// HACK(eddyb) `impl FnOnce() -> ...` allows some callers to avoid deserialization. -pub fn collect_zombies( - module: &Module, -) -> impl Iterator ZombieDecoration + '_)> + '_ { - module.annotations.iter().filter_map(|inst| { - // TODO: Temp hack. We hijack UserTypeGOOGLE right now, since the compiler never emits this. - 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(); - return Some((id, move || serde_json::from_str(encoded).unwrap())); - } - None - }) -} - -fn remove_zombie_annotations(module: &mut Module) { - module.annotations.retain(|inst| { - inst.class.opcode != Op::DecorateString - || inst.operands[1].unwrap_decoration() != Decoration::UserTypeGOOGLE - }) -} - #[derive(Clone)] struct ZombieInfo<'a> { reason: &'a str, @@ -167,9 +143,9 @@ fn report_error_zombies(sess: &Session, module: &Module, zombie: &HashMap u32 { + let mut i = 0; + while i < 10 { + x = 31 * x + y; + i += 1; + } + x +} +#[allow(unused_attributes)] +#[spirv(fragment)] +pub fn main() { + java_hash_ten_times(7, 42); +} +"#, + "java_hash_ten_times", + // NOTE(eddyb) this is very verbose because of the new structurizer + // producing messier control-flow than necessary, but the important part + // being tested is `OpLoopMerge` having `Unroll` as its "Loop Control". + r#"%1 = OpFunction %2 None %3 +%4 = OpFunctionParameter %2 +%5 = OpFunctionParameter %2 +%6 = OpLabel +OpBranch %7 +%7 = OpLabel +OpBranch %8 +%8 = OpLabel +%9 = OpPhi %10 %11 %7 %12 %13 +%14 = OpPhi %2 %4 %7 %15 %13 +%16 = OpPhi %17 %18 %7 %19 %13 +OpLoopMerge %20 %13 Unroll +OpBranchConditional %16 %21 %20 +%21 = OpLabel +%22 = OpSLessThan %17 %9 %23 +OpSelectionMerge %24 None +OpBranchConditional %22 %25 %26 +%25 = OpLabel +%27 = OpIMul %2 %28 %14 +%29 = OpIAdd %2 %27 %5 +%30 = OpIAdd %10 %9 %31 +OpBranch %24 +%26 = OpLabel +OpReturnValue %14 +%24 = OpLabel +%12 = OpPhi %10 %30 %25 +%15 = OpPhi %2 %29 %25 +%19 = OpPhi %17 %32 %25 +OpBranch %13 +%13 = OpLabel +OpBranch %8 +%20 = OpLabel +OpUnreachable +OpFunctionEnd"#, + ); +}