nordic-dev.net/rust-gpu
2
0
Fork 0
mirror of https://github.com/EmbarkStudios/rust-gpu.git synced 2025-02-16 17:04:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
* Generalize the zombie serialization system to arbitrary custom decorations.

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

* Pacify the merciless clippy.
This commit is contained in:
Eduard-Mihai Burtescu 2020-12-10 21:57:15 +02:00 committed by GitHub
parent 57b49d932a
commit ebf3dbee8d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
11 changed files with 274 additions and 94 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
crates/rustc_codegen_spirv/src/codegen_cx/declare.rs
View File

@ -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 {});
}
_ => {}
}
}

52
crates/rustc_codegen_spirv/src/codegen_cx/mod.rs
View File

@ -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
}

151
crates/rustc_codegen_spirv/src/finalizing_passes.rs → crates/rustc_codegen_spirv/src/decorations.rs
View File

@ -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);
}
}

2
crates/rustc_codegen_spirv/src/lib.rs
View File

@ -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;

3
crates/rustc_codegen_spirv/src/linker/duplicates.rs
View File

@ -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();

10
crates/rustc_codegen_spirv/src/linker/mod.rs
View File

@ -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

20
crates/rustc_codegen_spirv/src/linker/new_structurizer.rs
View File

@ -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<Word, UnrollLoopsDecoration>,
) -> 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<BlockId, BlockIdx>,
/// `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();

24
crates/rustc_codegen_spirv/src/linker/structurizer.rs
View File

@ -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<Word, UnrollLoopsDecoration>,
) -> 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();

36
crates/rustc_codegen_spirv/src/linker/zombies.rs
View File

@ -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<Item = (Word, impl FnOnce() -> 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<Word,
}
pub fn remove_zombies(sess: &Session, module: &mut Module) {
let zombies_owned = collect_zombies(module)
.map(|(id, decode)| {
let ZombieDecoration { reason, span } = decode();
let zombies_owned = ZombieDecoration::decode_all(module)
.map(|(id, zombie)| {
let ZombieDecoration { reason, span } = zombie.deserialize();
let span = span
.and_then(|span| span.to_rustc(sess.source_map()))
.unwrap_or(DUMMY_SP);
@ -180,7 +156,7 @@ pub fn remove_zombies(sess: &Session, module: &mut Module) {
.iter()
.map(|(id, (reason, span))| (*id, ZombieInfo::new(reason, *span)))
.collect();
remove_zombie_annotations(module);
ZombieDecoration::remove_all(module);
// Note: This is O(n^2).
while spread_zombie(module, &mut zombies) {}

2
crates/rustc_codegen_spirv/src/symbols.rs
View File

@ -339,6 +339,7 @@ impl Symbols {
("block", SpirvAttribute::Block),
("flat", SpirvAttribute::Flat),
("sampled_image", SpirvAttribute::SampledImage),
("unroll_loops", SpirvAttribute::UnrollLoops),
]
.iter()
.cloned();
@ -457,6 +458,7 @@ pub enum SpirvAttribute {
SampledImage,
Block,
Flat,
UnrollLoops,
}
// Note that we could mark the attr as used via cx.tcx.sess.mark_attr_used(attr), but unused

62
crates/spirv-builder/src/test/basic.rs
View File

@ -237,3 +237,65 @@ pub fn main() {
loop {}
}"#);
}
#[test]
fn unroll_loops() {
dis_fn(
// FIXME(eddyb) use `for _ in 0..10` here when that works.
r#"
#[allow(unused_attributes)]
#[spirv(unroll_loops)]
fn java_hash_ten_times(mut x: u32, y: u32) -> 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"#,
);
}