Impl quad swap for hlsl, msl and wgsl, finish spv front

gfx-rs · May 16, 2024 · aeae440 · aeae440
1 parent ae9eccd
commit aeae440
Show file tree

Hide file tree

Showing 4 changed files with 80 additions and 5 deletions.
diff --git a/naga/src/back/hlsl/writer.rs b/naga/src/back/hlsl/writer.rs
@@ -2214,7 +2214,33 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                 }
                 writeln!(self.out, ");")?;
             }
-            Statement::SubgroupQuadSwap { direction, argument, result } => {}
+            Statement::SubgroupQuadSwap { direction, argument, result } => {
+                write!(self.out, "{level}")?;
+                write!(self.out, "const ")?;
+                let name = format!("{}{}", back::BAKE_PREFIX, result.index());
+                match func_ctx.info[result].ty {
+                    proc::TypeResolution::Handle(handle) => self.write_type(module, handle)?,
+                    proc::TypeResolution::Value(ref value) => {
+                        self.write_value_type(module, value)?
+                    }
+                };
+                write!(self.out, " {name} = ")?;
+                self.named_expressions.insert(result, name);
+
+                match direction {
+                    crate::Direction::X => {
+                        write!(self.out, "QuadReadAcrossX(")?;
+                    },
+                    crate::Direction::Y => {
+                        write!(self.out, "QuadReadAcrossY(")?;
+                    },
+                    crate::Direction::Diagonal => {
+                        write!(self.out, "QuadReadAcrossDiagonal(")?;
+                    },
+                }
+                self.write_expr(module, argument, func_ctx)?;
+                writeln!(self.out, ");")?;
+            }
         }
 
         Ok(())

diff --git a/naga/src/back/msl/writer.rs b/naga/src/back/msl/writer.rs
@@ -3272,7 +3272,27 @@ impl<W: Write> Writer<W> {
                     }
                     writeln!(self.out, ");")?;
                 },
-                crate::Statement::SubgroupQuadSwap { direction, argument, result } => {}
+                crate::Statement::SubgroupQuadSwap { direction, argument, result } => {
+                    write!(self.out, "{level}")?;
+                    let name = self.namer.call("");
+                    self.start_baking_expression(result, &context.expression, &name)?;
+                    self.named_expressions.insert(result, name);
+                    write!(self.out, "{NAMESPACE}::quad_shuffle_xor(")?;
+                    self.put_expression(argument, &context.expression, true)?;
+                    write!(self.out, ", ")?;
+                    match direction {
+                        crate::Direction::X => {
+                            write!(self.out, "0x01")?;
+                        },
+                        crate::Direction::Y => {
+                            write!(self.out, "0x10")?;
+                        },
+                        crate::Direction::Diagonal => {
+                            write!(self.out, "0x11")?;
+                        },
+                    }
+                    writeln!(self.out, ");")?;
+                }
             }
         }
 

diff --git a/naga/src/back/wgsl/writer.rs b/naga/src/back/wgsl/writer.rs
@@ -1046,7 +1046,26 @@ impl<W: Write> Writer<W> {
                 }
                 writeln!(self.out, ");")?;
             }
-            Statement::SubgroupQuadSwap { direction, argument, result } => {}
+            Statement::SubgroupQuadSwap { direction, argument, result } => {
+                write!(self.out, "{level}")?;
+                let res_name = format!("{}{}", back::BAKE_PREFIX, result.index());
+                self.start_named_expr(module, result, func_ctx, &res_name)?;
+                self.named_expressions.insert(result, res_name);
+
+                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)?;
+                writeln!(self.out, ");")?;
+            }
         }
 
         Ok(())

diff --git a/naga/src/front/spv/mod.rs b/naga/src/front/spv/mod.rs
@@ -3955,7 +3955,7 @@ impl<I: Iterator<Item = u32>> Frontend<I> {
                     let result_id = self.next()?;
                     let exec_scope_id = self.next()?;
                     let argument_id = self.next()?;
-                    let direction = self.next()?;
+                    let direction_id = self.next()?;
 
                     let argument_lookup = self.lookup_expression.lookup(argument_id)?;
                     let argument_handle = get_expr_handle!(argument_id, argument_lookup);
@@ -3965,6 +3965,16 @@ impl<I: Iterator<Item = u32>> Frontend<I> {
                         .filter(|exec_scope| *exec_scope == spirv::Scope::Subgroup as u32)
                         .ok_or(Error::InvalidBarrierScope(exec_scope_id))?;
 
+                    let direction_const = self.lookup_constant.lookup(direction_id)?;
+                    let direction_const = resolve_constant(ctx.gctx(), &direction_const.inner)
+                        .ok_or(Error::InvalidOperand)?;
+                    let direction = match direction_const {
+                        0 => crate::Direction::X,
+                        1 => crate::Direction::Y,
+                        2 => crate::Direction::Diagonal,
+                        _ => unreachable!()
+                    };
+
                     let result_type = self.lookup_type.lookup(result_type_id)?;
 
                     let result_handle = ctx.expressions.append(
@@ -3984,7 +3994,7 @@ impl<I: Iterator<Item = u32>> Frontend<I> {
 
                     block.push(
                         crate::Statement::SubgroupQuadSwap {
-                            direction: crate::Direction::X,
+                            direction,
                             result: result_handle,
                             argument: argument_handle,
                         },