diff options
Diffstat (limited to 'third_party/rust/naga/src/back/spv/selection.rs')
| -rw-r--r-- | third_party/rust/naga/src/back/spv/selection.rs | 257 |
1 files changed, 257 insertions, 0 deletions
diff --git a/third_party/rust/naga/src/back/spv/selection.rs b/third_party/rust/naga/src/back/spv/selection.rs new file mode 100644 index 0000000000..788b1f10ab --- /dev/null +++ b/third_party/rust/naga/src/back/spv/selection.rs @@ -0,0 +1,257 @@ +/*! +Generate SPIR-V conditional structures. + +Builders for `if` structures with `and`s. + +The types in this module track the information needed to emit SPIR-V code +for complex conditional structures, like those whose conditions involve +short-circuiting 'and' and 'or' structures. These track labels and can emit +`OpPhi` instructions to merge values produced along different paths. + +This currently only supports exactly the forms Naga uses, so it doesn't +support `or` or `else`, and only supports zero or one merged values. + +Naga needs to emit code roughly like this: + +```ignore + + value = DEFAULT; + if COND1 && COND2 { + value = THEN_VALUE; + } + // use value + +``` + +Assuming `ctx` and `block` are a mutable references to a [`BlockContext`] +and the current [`Block`], and `merge_type` is the SPIR-V type for the +merged value `value`, we can build SPIR-V for the code above like so: + +```ignore + + let cond = Selection::start(block, merge_type); + // ... compute `cond1` ... + cond.if_true(ctx, cond1, DEFAULT); + // ... compute `cond2` ... + cond.if_true(ctx, cond2, DEFAULT); + // ... compute THEN_VALUE + let merged_value = cond.finish(ctx, THEN_VALUE); + +``` + +After this, `merged_value` is either `DEFAULT` or `THEN_VALUE`, depending on +the path by which the merged block was reached. + +This takes care of writing all branch instructions, including an +`OpSelectionMerge` annotation in the header block; starting new blocks and +assigning them labels; and emitting the `OpPhi` that gathers together the +right sources for the merged values, for every path through the selection +construct. + +When there is no merged value to produce, you can pass `()` for `merge_type` +and the merge values. In this case no `OpPhi` instructions are produced, and +the `finish` method returns `()`. + +To enforce proper nesting, a `Selection` takes ownership of the `&mut Block` +pointer for the duration of its lifetime. To obtain the block for generating +code in the selection's body, call the `Selection::block` method. +*/ + +use super::{Block, BlockContext, Instruction}; +use spirv::Word; + +/// A private struct recording what we know about the selection construct so far. +pub(super) struct Selection<'b, M: MergeTuple> { + /// The block pointer we're emitting code into. + block: &'b mut Block, + + /// The label of the selection construct's merge block, or `None` if we + /// haven't yet written the `OpSelectionMerge` merge instruction. + merge_label: Option<Word>, + + /// A set of `(VALUES, PARENT)` pairs, used to build `OpPhi` instructions in + /// the merge block. Each `PARENT` is the label of a predecessor block of + /// the merge block. The corresponding `VALUES` holds the ids of the values + /// that `PARENT` contributes to the merged values. + /// + /// We emit all branches to the merge block, so we know all its + /// predecessors. And we refuse to emit a branch unless we're given the + /// values the branching block contributes to the merge, so we always have + /// everything we need to emit the correct phis, by construction. + values: Vec<(M, Word)>, + + /// The types of the values in each element of `values`. + merge_types: M, +} + +impl<'b, M: MergeTuple> Selection<'b, M> { + /// Start a new selection construct. + /// + /// The `block` argument indicates the selection's header block. + /// + /// The `merge_types` argument should be a `Word` or tuple of `Word`s, each + /// value being the SPIR-V result type id of an `OpPhi` instruction that + /// will be written to the selection's merge block when this selection's + /// [`finish`] method is called. This argument may also be `()`, for + /// selections that produce no values. + /// + /// (This function writes no code to `block` itself; it simply constructs a + /// fresh `Selection`.) + /// + /// [`finish`]: Selection::finish + pub(super) fn start(block: &'b mut Block, merge_types: M) -> Self { + Selection { + block, + merge_label: None, + values: vec![], + merge_types, + } + } + + pub(super) fn block(&mut self) -> &mut Block { + self.block + } + + /// Branch to a successor block if `cond` is true, otherwise merge. + /// + /// If `cond` is false, branch to the merge block, using `values` as the + /// merged values. Otherwise, proceed to a new block. + /// + /// The `values` argument must be the same shape as the `merge_types` + /// argument passed to `Selection::start`. + pub(super) fn if_true(&mut self, ctx: &mut BlockContext, cond: Word, values: M) { + self.values.push((values, self.block.label_id)); + + let merge_label = self.make_merge_label(ctx); + let next_label = ctx.gen_id(); + ctx.function.consume( + std::mem::replace(self.block, Block::new(next_label)), + Instruction::branch_conditional(cond, next_label, merge_label), + ); + } + + /// Emit an unconditional branch to the merge block, and compute merged + /// values. + /// + /// Use `final_values` as the merged values contributed by the current + /// block, and transition to the merge block, emitting `OpPhi` instructions + /// to produce the merged values. This must be the same shape as the + /// `merge_types` argument passed to [`Selection::start`]. + /// + /// Return the SPIR-V ids of the merged values. This value has the same + /// shape as the `merge_types` argument passed to `Selection::start`. + pub(super) fn finish(self, ctx: &mut BlockContext, final_values: M) -> M { + match self { + Selection { + merge_label: None, .. + } => { + // We didn't actually emit any branches, so `self.values` must + // be empty, and `final_values` are the only sources we have for + // the merged values. Easy peasy. + final_values + } + + Selection { + block, + merge_label: Some(merge_label), + mut values, + merge_types, + } => { + // Emit the final branch and transition to the merge block. + values.push((final_values, block.label_id)); + ctx.function.consume( + std::mem::replace(block, Block::new(merge_label)), + Instruction::branch(merge_label), + ); + + // Now that we're in the merge block, build the phi instructions. + merge_types.write_phis(ctx, block, &values) + } + } + } + + /// Return the id of the merge block, writing a merge instruction if needed. + fn make_merge_label(&mut self, ctx: &mut BlockContext) -> Word { + match self.merge_label { + None => { + let merge_label = ctx.gen_id(); + self.block.body.push(Instruction::selection_merge( + merge_label, + spirv::SelectionControl::NONE, + )); + self.merge_label = Some(merge_label); + merge_label + } + Some(merge_label) => merge_label, + } + } +} + +/// A trait to help `Selection` manage any number of merged values. +/// +/// Some selection constructs, like a `ReadZeroSkipWrite` bounds check on a +/// [`Load`] expression, produce a single merged value. Others produce no merged +/// value, like a bounds check on a [`Store`] statement. +/// +/// To let `Selection` work nicely with both cases, we let the merge type +/// argument passed to [`Selection::start`] be any type that implements this +/// `MergeTuple` trait. `MergeTuple` is then implemented for `()`, `Word`, +/// `(Word, Word)`, and so on. +/// +/// A `MergeTuple` type can represent either a bunch of SPIR-V types or values; +/// the `merge_types` argument to `Selection::start` are type ids, whereas the +/// `values` arguments to the [`if_true`] and [`finish`] methods are value ids. +/// The set of merged value returned by `finish` is a tuple of value ids. +/// +/// In fact, since Naga only uses zero- and single-valued selection constructs +/// at present, we only implement `MergeTuple` for `()` and `Word`. But if you +/// add more cases, feel free to add more implementations. Once const generics +/// are available, we could have a single implementation of `MergeTuple` for all +/// lengths of arrays, and be done with it. +/// +/// [`Load`]: crate::Expression::Load +/// [`Store`]: crate::Statement::Store +/// [`if_true`]: Selection::if_true +/// [`finish`]: Selection::finish +pub(super) trait MergeTuple: Sized { + /// Write OpPhi instructions for the given set of predecessors. + /// + /// The `predecessors` vector should be a vector of `(LABEL, VALUES)` pairs, + /// where each `VALUES` holds the values contributed by the branch from + /// `LABEL`, which should be one of the current block's predecessors. + fn write_phis( + self, + ctx: &mut BlockContext, + block: &mut Block, + predecessors: &[(Self, Word)], + ) -> Self; +} + +/// Selections that produce a single merged value. +/// +/// For example, `ImageLoad` with `BoundsCheckPolicy::ReadZeroSkipWrite` either +/// returns a texel value or zeros. +impl MergeTuple for Word { + fn write_phis( + self, + ctx: &mut BlockContext, + block: &mut Block, + predecessors: &[(Word, Word)], + ) -> Word { + let merged_value = ctx.gen_id(); + block + .body + .push(Instruction::phi(self, merged_value, predecessors)); + merged_value + } +} + +/// Selections that produce no merged values. +/// +/// For example, `ImageStore` under `BoundsCheckPolicy::ReadZeroSkipWrite` +/// either does the store or skips it, but in neither case does it produce a +/// value. +impl MergeTuple for () { + /// No phis need to be generated. + fn write_phis(self, _: &mut BlockContext, _: &mut Block, _: &[((), Word)]) {} +} |