Gather operations

gfx-rs · Dec 16, 2021 · cedcc2c · cedcc2c
1 parent 8ffd6ba
commit cedcc2c
Show file tree

Hide file tree

Showing 19 changed files with 401 additions and 43 deletions.
diff --git a/src/back/dot/mod.rs b/src/back/dot/mod.rs
@@ -231,6 +231,7 @@ fn write_fun(
             E::ImageSample {
                 image,
                 sampler,
+                gather,
                 coordinate,
                 array_index,
                 offset: _,
@@ -260,7 +261,11 @@ fn write_fun(
                 if let Some(expr) = depth_ref {
                     edges.insert("depth_ref", expr);
                 }
-                ("ImageSample".into(), 5)
+                let string = match gather {
+                    Some(component) => Cow::Owned(format!("ImageGather{:?}", component)),
+                    _ => Cow::Borrowed("ImageSample"),
+                };
+                (string, 5)
             }
             E::ImageLoad {
                 image,

diff --git a/src/back/glsl/mod.rs b/src/back/glsl/mod.rs
@@ -1928,6 +1928,7 @@ impl<'a, W: Write> Writer<'a, W> {
             Expression::ImageSample {
                 image,
                 sampler: _, //TODO?
+                gather,
                 coordinate,
                 array_index,
                 offset,
@@ -1962,13 +1963,15 @@ impl<'a, W: Write> Writer<'a, W> {
                 let workaround_lod_array_shadow_as_grad = (array_index.is_some()
                     || dim == crate::ImageDimension::Cube)
                     && depth_ref.is_some()
+                    && gather.is_none()
                     && !self
                         .options
                         .writer_flags
                         .contains(WriterFlags::TEXTURE_SHADOW_LOD);
 
                 //Write the function to be used depending on the sample level
                 let fun_name = match level {
+                    crate::SampleLevel::Zero if gather.is_some() => "textureGather",
                     crate::SampleLevel::Auto | crate::SampleLevel::Bias(_) => "texture",
                     crate::SampleLevel::Zero | crate::SampleLevel::Exact(_) => {
                         if workaround_lod_array_shadow_as_grad {
@@ -2045,7 +2048,7 @@ impl<'a, W: Write> Writer<'a, W> {
                     crate::SampleLevel::Zero => {
                         if workaround_lod_array_shadow_as_grad {
                             write!(self.out, ", vec2(0,0), vec2(0,0)")?;
-                        } else {
+                        } else if gather.is_none() {
                             write!(self.out, ", 0.0")?;
                         }
                     }
@@ -2082,6 +2085,10 @@ impl<'a, W: Write> Writer<'a, W> {
                     }
                 }
 
+                if let Some(component) = gather {
+                    write!(self.out, ", {}", component as usize)?;
+                }
+
                 // End the function
                 write!(self.out, ")")?
             }

diff --git a/src/back/hlsl/writer.rs b/src/back/hlsl/writer.rs
@@ -1594,30 +1594,38 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
             Expression::ImageSample {
                 image,
                 sampler,
+                gather,
                 coordinate,
                 array_index,
                 offset,
                 level,
                 depth_ref,
             } => {
                 use crate::SampleLevel as Sl;
+                const COMPONENTS: [&str; 4] = ["", "Green", "Blue", "Alpha"];
 
-                let texture_func = match level {
-                    Sl::Auto => {
-                        if depth_ref.is_some() {
-                            "SampleCmp"
-                        } else {
-                            "Sample"
-                        }
-                    }
-                    Sl::Zero => "SampleCmpLevelZero",
-                    Sl::Exact(_) => "SampleLevel",
-                    Sl::Bias(_) => "SampleBias",
-                    Sl::Gradient { .. } => "SampleGrad",
+                let (base_str, component_str) = match gather {
+                    Some(component) => ("Gather", COMPONENTS[component as usize]),
+                    None => ("Sample", ""),
+                };
+                let cmp_str = match depth_ref {
+                    Some(_) => "Cmp",
+                    None => "",
+                };
+                let level_str = match level {
+                    Sl::Auto => "",
+                    Sl::Zero => "LevelZero",
+                    Sl::Exact(_) => "Level",
+                    Sl::Bias(_) => "Bias",
+                    Sl::Gradient { .. } => "Grad",
                 };
 
                 self.write_expr(module, image, func_ctx)?;
-                write!(self.out, ".{}(", texture_func)?;
+                write!(
+                    self.out,
+                    ".{}{}{}{}(",
+                    base_str, cmp_str, component_str, level_str
+                )?;
                 self.write_expr(module, sampler, func_ctx)?;
                 write!(self.out, ", ")?;
                 self.write_texture_coordinates(

diff --git a/src/back/msl/writer.rs b/src/back/msl/writer.rs
@@ -896,18 +896,23 @@ impl<W: Write> Writer<W> {
             crate::Expression::ImageSample {
                 image,
                 sampler,
+                gather,
                 coordinate,
                 array_index,
                 offset,
                 level,
                 depth_ref,
             } => {
-                let op = match depth_ref {
-                    Some(_) => "sample_compare",
+                let main_op = match gather {
+                    Some(_) => "gather",
                     None => "sample",
                 };
+                let comparison_op = match depth_ref {
+                    Some(_) => "_compare",
+                    None => "",
+                };
                 self.put_expression(image, context, false)?;
-                write!(self.out, ".{}(", op)?;
+                write!(self.out, ".{}{}(", main_op, comparison_op)?;
                 self.put_expression(sampler, context, true)?;
                 write!(self.out, ", ")?;
                 self.put_expression(coordinate, context, true)?;
@@ -931,6 +936,25 @@ impl<W: Write> Writer<W> {
                     };
                     write!(self.out, ", {}", coco)?;
                 }
+                match gather {
+                    None | Some(crate::SwizzleComponent::X) => {}
+                    Some(component) => {
+                        let is_cube_map = match *context.resolve_type(image) {
+                            crate::TypeInner::Image {
+                                dim: crate::ImageDimension::Cube,
+                                ..
+                            } => true,
+                            _ => false,
+                        };
+                        // Offset always comes before the gather, except
+                        // in cube mapes where it's not applicable
+                        if offset.is_none() && !is_cube_map {
+                            write!(self.out, ", int2(0)")?;
+                        }
+                        let letter = ['x', 'y', 'z', 'w'][component as usize];
+                        write!(self.out, ", {}::component::{}", NAMESPACE, letter)?;
+                    }
+                }
                 write!(self.out, ")")?;
             }
             crate::Expression::ImageLoad {

diff --git a/src/back/spv/block.rs b/src/back/spv/block.rs
@@ -885,6 +885,7 @@ impl<'w> BlockContext<'w> {
             crate::Expression::ImageSample {
                 image,
                 sampler,
+                gather,
                 coordinate,
                 array_index,
                 offset,
@@ -894,6 +895,7 @@ impl<'w> BlockContext<'w> {
                 result_type_id,
                 image,
                 sampler,
+                gather,
                 coordinate,
                 array_index,
                 offset,

diff --git a/src/back/spv/image.rs b/src/back/spv/image.rs
@@ -799,6 +799,7 @@ impl<'w> BlockContext<'w> {
         result_type_id: Word,
         image: Handle<crate::Expression>,
         sampler: Handle<crate::Expression>,
+        gather: Option<crate::SwizzleComponent>,
         coordinate: Handle<crate::Expression>,
         array_index: Option<Handle<crate::Expression>>,
         offset: Option<Handle<crate::Constant>>,
@@ -816,7 +817,7 @@ impl<'w> BlockContext<'w> {
             crate::TypeInner::Image {
                 class: crate::ImageClass::Depth { .. },
                 ..
-            } => depth_ref.is_none(),
+            } => depth_ref.is_none() && gather.is_none(),
             _ => false,
         };
         let sample_result_type_id = if needs_sub_access {
@@ -853,8 +854,23 @@ impl<'w> BlockContext<'w> {
         let mut mask = spirv::ImageOperands::empty();
         mask.set(spirv::ImageOperands::CONST_OFFSET, offset.is_some());
 
-        let mut main_instruction = match level {
-            crate::SampleLevel::Zero => {
+        let mut main_instruction = match (level, gather) {
+            (_, Some(component)) => {
+                let component_id = self.get_index_constant(component as u32);
+                let mut inst = Instruction::image_gather(
+                    sample_result_type_id,
+                    id,
+                    sampled_image_id,
+                    coordinates_id,
+                    component_id,
+                    depth_id,
+                );
+                if !mask.is_empty() {
+                    inst.add_operand(mask.bits());
+                }
+                inst
+            }
+            (crate::SampleLevel::Zero, None) => {
                 let mut inst = Instruction::image_sample(
                     sample_result_type_id,
                     id,
@@ -874,7 +890,7 @@ impl<'w> BlockContext<'w> {
 
                 inst
             }
-            crate::SampleLevel::Auto => {
+            (crate::SampleLevel::Auto, None) => {
                 let mut inst = Instruction::image_sample(
                     sample_result_type_id,
                     id,
@@ -888,7 +904,7 @@ impl<'w> BlockContext<'w> {
                 }
                 inst
             }
-            crate::SampleLevel::Exact(lod_handle) => {
+            (crate::SampleLevel::Exact(lod_handle), None) => {
                 let mut inst = Instruction::image_sample(
                     sample_result_type_id,
                     id,
@@ -905,7 +921,7 @@ impl<'w> BlockContext<'w> {
 
                 inst
             }
-            crate::SampleLevel::Bias(bias_handle) => {
+            (crate::SampleLevel::Bias(bias_handle), None) => {
                 let mut inst = Instruction::image_sample(
                     sample_result_type_id,
                     id,
@@ -922,7 +938,7 @@ impl<'w> BlockContext<'w> {
 
                 inst
             }
-            crate::SampleLevel::Gradient { x, y } => {
+            (crate::SampleLevel::Gradient { x, y }, None) => {
                 let mut inst = Instruction::image_sample(
                     sample_result_type_id,
                     id,

diff --git a/src/back/spv/instructions.rs b/src/back/spv/instructions.rs
@@ -569,6 +569,33 @@ impl super::Instruction {
         instruction
     }
 
+    pub(super) fn image_gather(
+        result_type_id: Word,
+        id: Word,
+        sampled_image: Word,
+        coordinates: Word,
+        component_id: Word,
+        depth_ref: Option<Word>,
+    ) -> Self {
+        let op = match depth_ref {
+            None => Op::ImageGather,
+            Some(_) => Op::ImageDrefGather,
+        };
+
+        let mut instruction = Self::new(op);
+        instruction.set_type(result_type_id);
+        instruction.set_result(id);
+        instruction.add_operand(sampled_image);
+        instruction.add_operand(coordinates);
+        if let Some(dref) = depth_ref {
+            instruction.add_operand(dref);
+        } else {
+            instruction.add_operand(component_id);
+        }
+
+        instruction
+    }
+
     pub(super) fn image_fetch_or_read(
         op: Op,
         result_type_id: Word,

diff --git a/src/back/wgsl/writer.rs b/src/back/wgsl/writer.rs
@@ -1239,6 +1239,7 @@ impl<W: Write> Writer<W> {
             Expression::ImageSample {
                 image,
                 sampler,
+                gather: None,
                 coordinate,
                 array_index,
                 offset,
@@ -1306,6 +1307,54 @@ impl<W: Write> Writer<W> {
 
                 write!(self.out, ")")?;
             }
+            Expression::ImageSample {
+                image,
+                sampler,
+                gather: Some(component),
+                coordinate,
+                array_index,
+                offset,
+                level: _,
+                depth_ref,
+            } => {
+                let suffix_cmp = match depth_ref {
+                    Some(_) => "Compare",
+                    None => "",
+                };
+
+                write!(self.out, "textureGather{}(", suffix_cmp)?;
+                match *func_ctx.info[image].ty.inner_with(&module.types) {
+                    TypeInner::Image {
+                        class: crate::ImageClass::Depth { multi: _ },
+                        ..
+                    } => {}
+                    _ => {
+                        write!(self.out, "{}, ", component as u8)?;
+                    }
+                }
+                self.write_expr(module, image, func_ctx)?;
+                write!(self.out, ", ")?;
+                self.write_expr(module, sampler, func_ctx)?;
+                write!(self.out, ", ")?;
+                self.write_expr(module, coordinate, func_ctx)?;
+
+                if let Some(array_index) = array_index {
+                    write!(self.out, ", ")?;
+                    self.write_expr(module, array_index, func_ctx)?;
+                }
+
+                if let Some(depth_ref) = depth_ref {
+                    write!(self.out, ", ")?;
+                    self.write_expr(module, depth_ref, func_ctx)?;
+                }
+
+                if let Some(offset) = offset {
+                    write!(self.out, ", ")?;
+                    self.write_constant(module, offset)?;
+                }
+
+                write!(self.out, ")")?;
+            }
             Expression::ImageQuery { image, query } => {
                 use crate::ImageQuery as Iq;
 

diff --git a/src/front/glsl/builtins.rs b/src/front/glsl/builtins.rs
@@ -1930,6 +1930,7 @@ fn texture_call(
             Expression::ImageSample {
                 image,
                 sampler,
+                gather: None, //TODO
                 coordinate: comps.coordinate,
                 array_index,
                 offset,

diff --git a/src/front/spv/image.rs b/src/front/spv/image.rs
@@ -606,6 +606,7 @@ impl<I: Iterator<Item = u32>> super::Parser<I> {
         let expr = crate::Expression::ImageSample {
             image: si_lexp.image,
             sampler: si_lexp.sampler,
+            gather: None, //TODO
             coordinate,
             array_index,
             offset,