Skip to content

Implement subgroup quad ops #7683

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 7 commits into
base: trunk
Choose a base branch
from
Open

Implement subgroup quad ops #7683

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
54a3ae0
Rudimentary impl of quad ops, impl quad ops for spirv
valaphee May 9, 2024
8690ad4
Impl quad swap for hlsl, msl and wgsl, finish spv front
valaphee May 16, 2024
e1e92da
Cargo clippy & cargo fmt, impl valid for quad ops
valaphee May 16, 2024
f1492f0
Enable quad feature
valaphee May 16, 2024
b525a60
Add missing feature to glsl
valaphee May 16, 2024
63ee3ca
Simplifying code by making `SubgroupQuadSwap` an instance of `Subgrou…
dzamkov May 11, 2025
409aba6
Add `GroupNonUniformQuad` spv capability to Vulkan
dzamkov May 11, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 9 additions & 1 deletion naga/src/back/dot/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -379,9 +379,11 @@ impl StatementGraph {
| crate::GatherMode::Shuffle(index)
| crate::GatherMode::ShuffleDown(index)
| crate::GatherMode::ShuffleUp(index)
| crate::GatherMode::ShuffleXor(index) => {
| crate::GatherMode::ShuffleXor(index)
| crate::GatherMode::QuadBroadcast(index) => {
self.dependencies.push((id, index, "index"))
}
crate::GatherMode::QuadSwap(_) => {}
}
self.dependencies.push((id, argument, "arg"));
self.emits.push((id, result));
Expand All @@ -392,6 +394,12 @@ impl StatementGraph {
crate::GatherMode::ShuffleDown(_) => "SubgroupShuffleDown",
crate::GatherMode::ShuffleUp(_) => "SubgroupShuffleUp",
crate::GatherMode::ShuffleXor(_) => "SubgroupShuffleXor",
crate::GatherMode::QuadBroadcast(_) => "SubgroupQuadBroadcast",
crate::GatherMode::QuadSwap(direction) => match direction {
crate::Direction::X => "SubgroupQuadSwapX",
crate::Direction::Y => "SubgroupQuadSwapY",
crate::Direction::Diagonal => "SubgroupQuadSwapDiagonal",
},
}
}
};
Expand Down
1 change: 1 addition & 0 deletions naga/src/back/glsl/features.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -280,6 +280,7 @@ impl FeaturesManager {
out,
"#extension GL_KHR_shader_subgroup_shuffle_relative : require"
)?;
writeln!(out, "#extension GL_KHR_shader_subgroup_quad : require")?;
}

if self.0.contains(Features::TEXTURE_ATOMICS) {
Expand Down
18 changes: 17 additions & 1 deletion naga/src/back/glsl/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2716,6 +2716,20 @@ impl<'a, W: Write> Writer<'a, W> {
crate::GatherMode::ShuffleXor(_) => {
write!(self.out, "subgroupShuffleXor(")?;
}
crate::GatherMode::QuadBroadcast(_) => {
write!(self.out, "subgroupQuadBroadcast(")?;
}
crate::GatherMode::QuadSwap(direction) => match direction {
crate::Direction::X => {
write!(self.out, "subgroupQuadSwapHorizontal(")?;
}
crate::Direction::Y => {
write!(self.out, "subgroupQuadSwapVertical(")?;
}
crate::Direction::Diagonal => {
write!(self.out, "subgroupQuadSwapDiagonal(")?;
}
},
}
self.write_expr(argument, ctx)?;
match mode {
Expand All @@ -2724,10 +2738,12 @@ impl<'a, W: Write> Writer<'a, W> {
| crate::GatherMode::Shuffle(index)
| crate::GatherMode::ShuffleDown(index)
| crate::GatherMode::ShuffleUp(index)
| crate::GatherMode::ShuffleXor(index) => {
| crate::GatherMode::ShuffleXor(index)
| crate::GatherMode::QuadBroadcast(index) => {
write!(self.out, ", ")?;
self.write_expr(index, ctx)?;
}
crate::GatherMode::QuadSwap(_) => {}
}
writeln!(self.out, ");")?;
}
Expand Down
71 changes: 48 additions & 23 deletions naga/src/back/hlsl/writer.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2610,30 +2610,55 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
};
write!(self.out, " {name} = ")?;
self.named_expressions.insert(result, name);

if matches!(mode, crate::GatherMode::BroadcastFirst) {
write!(self.out, "WaveReadLaneFirst(")?;
self.write_expr(module, argument, func_ctx)?;
} else {
write!(self.out, "WaveReadLaneAt(")?;
self.write_expr(module, argument, func_ctx)?;
write!(self.out, ", ")?;
match mode {
crate::GatherMode::BroadcastFirst => unreachable!(),
crate::GatherMode::Broadcast(index) | crate::GatherMode::Shuffle(index) => {
self.write_expr(module, index, func_ctx)?;
}
crate::GatherMode::ShuffleDown(index) => {
write!(self.out, "WaveGetLaneIndex() + ")?;
self.write_expr(module, index, func_ctx)?;
}
crate::GatherMode::ShuffleUp(index) => {
write!(self.out, "WaveGetLaneIndex() - ")?;
self.write_expr(module, index, func_ctx)?;
match mode {
crate::GatherMode::BroadcastFirst => {
write!(self.out, "WaveReadLaneFirst(")?;
self.write_expr(module, argument, func_ctx)?;
}
crate::GatherMode::QuadBroadcast(index) => {
write!(self.out, "QuadReadLaneAt(")?;
self.write_expr(module, argument, func_ctx)?;
write!(self.out, ", ")?;
self.write_expr(module, index, func_ctx)?;
}
crate::GatherMode::QuadSwap(direction) => {
match direction {
crate::Direction::X => {
write!(self.out, "QuadReadAcrossX(")?;
}
crate::Direction::Y => {
write!(self.out, "QuadReadAcrossY(")?;
}
crate::Direction::Diagonal => {
write!(self.out, "QuadReadAcrossDiagonal(")?;
}
}
crate::GatherMode::ShuffleXor(index) => {
write!(self.out, "WaveGetLaneIndex() ^ ")?;
self.write_expr(module, index, func_ctx)?;
self.write_expr(module, argument, func_ctx)?;
}
_ => {
write!(self.out, "WaveReadLaneAt(")?;
self.write_expr(module, argument, func_ctx)?;
write!(self.out, ", ")?;
match mode {
crate::GatherMode::BroadcastFirst => unreachable!(),
crate::GatherMode::Broadcast(index)
| crate::GatherMode::Shuffle(index) => {
self.write_expr(module, index, func_ctx)?;
}
crate::GatherMode::ShuffleDown(index) => {
write!(self.out, "WaveGetLaneIndex() + ")?;
self.write_expr(module, index, func_ctx)?;
}
crate::GatherMode::ShuffleUp(index) => {
write!(self.out, "WaveGetLaneIndex() - ")?;
self.write_expr(module, index, func_ctx)?;
}
crate::GatherMode::ShuffleXor(index) => {
write!(self.out, "WaveGetLaneIndex() ^ ")?;
self.write_expr(module, index, func_ctx)?;
}
crate::GatherMode::QuadBroadcast(_) => unreachable!(),
crate::GatherMode::QuadSwap(_) => unreachable!(),
}
}
}
Expand Down
23 changes: 22 additions & 1 deletion naga/src/back/msl/writer.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3931,6 +3931,12 @@ impl<W: Write> Writer<W> {
crate::GatherMode::ShuffleXor(_) => {
write!(self.out, "{NAMESPACE}::simd_shuffle_xor(")?;
}
crate::GatherMode::QuadBroadcast(_) => {
write!(self.out, "{NAMESPACE}::quad_broadcast(")?;
}
crate::GatherMode::QuadSwap(_) => {
write!(self.out, "{NAMESPACE}::quad_shuffle_xor(")?;
}
}
self.put_expression(argument, &context.expression, true)?;
match mode {
Expand All @@ -3939,10 +3945,25 @@ impl<W: Write> Writer<W> {
| crate::GatherMode::Shuffle(index)
| crate::GatherMode::ShuffleDown(index)
| crate::GatherMode::ShuffleUp(index)
| crate::GatherMode::ShuffleXor(index) => {
| crate::GatherMode::ShuffleXor(index)
| crate::GatherMode::QuadBroadcast(index) => {
write!(self.out, ", ")?;
self.put_expression(index, &context.expression, true)?;
}
crate::GatherMode::QuadSwap(direction) => {
write!(self.out, ", ")?;
match direction {
crate::Direction::X => {
write!(self.out, "1u")?;
}
crate::Direction::Y => {
write!(self.out, "2u")?;
}
crate::Direction::Diagonal => {
write!(self.out, "3u")?;
}
}
}
}
writeln!(self.out, ");")?;
}
Expand Down
4 changes: 3 additions & 1 deletion naga/src/back/pipeline_constants.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -759,9 +759,11 @@ fn adjust_stmt(new_pos: &HandleVec<Expression, Handle<Expression>>, stmt: &mut S
| crate::GatherMode::Shuffle(ref mut index)
| crate::GatherMode::ShuffleDown(ref mut index)
| crate::GatherMode::ShuffleUp(ref mut index)
| crate::GatherMode::ShuffleXor(ref mut index) => {
| crate::GatherMode::ShuffleXor(ref mut index)
| crate::GatherMode::QuadBroadcast(ref mut index) => {
adjust(index);
}
crate::GatherMode::QuadSwap(_) => {}
}
adjust(argument);
adjust(result)
Expand Down
16 changes: 16 additions & 0 deletions naga/src/back/spv/instructions.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1203,6 +1203,22 @@ impl super::Instruction {
}
instruction.add_operand(value);

instruction
}
pub(super) fn group_non_uniform_quad_swap(
result_type_id: Word,
id: Word,
exec_scope_id: Word,
value: Word,
direction: Word,
) -> Self {
let mut instruction = Self::new(Op::GroupNonUniformQuadSwap);
instruction.set_type(result_type_id);
instruction.set_result(id);
instruction.add_operand(exec_scope_id);
instruction.add_operand(value);
instruction.add_operand(direction);

instruction
}
}
Expand Down
29 changes: 24 additions & 5 deletions naga/src/back/spv/subgroup.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -125,10 +125,6 @@ impl BlockContext<'_> {
result: Handle<crate::Expression>,
block: &mut Block,
) -> Result<(), Error> {
self.writer.require_any(
"GroupNonUniformBallot",
&[spirv::Capability::GroupNonUniformBallot],
)?;
match *mode {
crate::GatherMode::BroadcastFirst => {
self.writer.require_any(
Expand All @@ -150,6 +146,12 @@ impl BlockContext<'_> {
&[spirv::Capability::GroupNonUniformShuffleRelative],
)?;
}
crate::GatherMode::QuadBroadcast(_) | crate::GatherMode::QuadSwap(_) => {
self.writer.require_any(
"GroupNonUniformQuad",
&[spirv::Capability::GroupNonUniformQuad],
)?;
}
}

let id = self.gen_id();
Expand All @@ -174,7 +176,8 @@ impl BlockContext<'_> {
| crate::GatherMode::Shuffle(index)
| crate::GatherMode::ShuffleDown(index)
| crate::GatherMode::ShuffleUp(index)
| crate::GatherMode::ShuffleXor(index) => {
| crate::GatherMode::ShuffleXor(index)
| crate::GatherMode::QuadBroadcast(index) => {
let index_id = self.cached[index];
let op = match *mode {
crate::GatherMode::BroadcastFirst => unreachable!(),
Expand All @@ -187,6 +190,8 @@ impl BlockContext<'_> {
crate::GatherMode::ShuffleDown(_) => spirv::Op::GroupNonUniformShuffleDown,
crate::GatherMode::ShuffleUp(_) => spirv::Op::GroupNonUniformShuffleUp,
crate::GatherMode::ShuffleXor(_) => spirv::Op::GroupNonUniformShuffleXor,
crate::GatherMode::QuadBroadcast(_) => spirv::Op::GroupNonUniformQuadBroadcast,
crate::GatherMode::QuadSwap(_) => unreachable!(),
};
block.body.push(Instruction::group_non_uniform_gather(
op,
Expand All @@ -197,6 +202,20 @@ impl BlockContext<'_> {
index_id,
));
}
crate::GatherMode::QuadSwap(direction) => {
let direction = self.get_index_constant(match direction {
crate::Direction::X => 0,
crate::Direction::Y => 1,
crate::Direction::Diagonal => 2,
});
block.body.push(Instruction::group_non_uniform_quad_swap(
result_type_id,
id,
exec_scope_id,
arg_id,
direction,
));
}
}
self.cached[result] = id;
Ok(())
Expand Down
18 changes: 17 additions & 1 deletion naga/src/back/wgsl/writer.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -945,6 +945,20 @@ impl<W: Write> Writer<W> {
crate::GatherMode::ShuffleXor(_) => {
write!(self.out, "subgroupShuffleXor(")?;
}
crate::GatherMode::QuadBroadcast(_) => {
write!(self.out, "quadBroadcast(")?;
}
crate::GatherMode::QuadSwap(direction) => match direction {
crate::Direction::X => {
write!(self.out, "quadSwapX(")?;
}
crate::Direction::Y => {
write!(self.out, "quadSwapY(")?;
}
crate::Direction::Diagonal => {
write!(self.out, "quadSwapDiagonal(")?;
}
},
}
self.write_expr(module, argument, func_ctx)?;
match mode {
Expand All @@ -953,10 +967,12 @@ impl<W: Write> Writer<W> {
| crate::GatherMode::Shuffle(index)
| crate::GatherMode::ShuffleDown(index)
| crate::GatherMode::ShuffleUp(index)
| crate::GatherMode::ShuffleXor(index) => {
| crate::GatherMode::ShuffleXor(index)
| crate::GatherMode::QuadBroadcast(index) => {
write!(self.out, ", ")?;
self.write_expr(module, index, func_ctx)?;
}
crate::GatherMode::QuadSwap(_) => {}
}
writeln!(self.out, ");")?;
}
Expand Down
8 changes: 6 additions & 2 deletions naga/src/compact/statements.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -141,9 +141,11 @@ impl FunctionTracer<'_> {
| crate::GatherMode::Shuffle(index)
| crate::GatherMode::ShuffleDown(index)
| crate::GatherMode::ShuffleUp(index)
| crate::GatherMode::ShuffleXor(index) => {
| crate::GatherMode::ShuffleXor(index)
| crate::GatherMode::QuadBroadcast(index) => {
self.expressions_used.insert(index);
}
crate::GatherMode::QuadSwap(_) => {}
}
self.expressions_used.insert(argument);
self.expressions_used.insert(result);
Expand Down Expand Up @@ -350,7 +352,9 @@ impl FunctionMap {
| crate::GatherMode::Shuffle(ref mut index)
| crate::GatherMode::ShuffleDown(ref mut index)
| crate::GatherMode::ShuffleUp(ref mut index)
| crate::GatherMode::ShuffleXor(ref mut index) => adjust(index),
| crate::GatherMode::ShuffleXor(ref mut index)
| crate::GatherMode::QuadBroadcast(ref mut index) => adjust(index),
crate::GatherMode::QuadSwap(_) => {}
}
adjust(argument);
adjust(result);
Expand Down
Loading
Loading