Revision control

Copy as Markdown

Other Tools

/*!
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)]) {}
}