From de5bd772792d9f6ea2282dde65c4b1e7544b9405 Mon Sep 17 00:00:00 2001 From: Dzmitry Malyshau Date: Sat, 13 Feb 2021 18:52:00 -0500 Subject: [PATCH] ImageStore statement in IR --- src/back/glsl/mod.rs | 78 +++++++--- src/back/msl/writer.rs | 17 +++ src/back/spv/instructions.rs | 24 +++- src/back/spv/writer.rs | 271 ++++++++++++++++++++++------------- src/front/spv/mod.rs | 7 +- src/front/wgsl/mod.rs | 26 ++++ src/lib.rs | 7 + src/proc/analyzer.rs | 15 ++ src/proc/interface.rs | 13 ++ src/proc/terminator.rs | 1 + 10 files changed, 337 insertions(+), 122 deletions(-) diff --git a/src/back/glsl/mod.rs b/src/back/glsl/mod.rs index 528deed9f9..14959bbb35 100644 --- a/src/back/glsl/mod.rs +++ b/src/back/glsl/mod.rs @@ -1129,6 +1129,27 @@ impl<'a, W: Write> Writer<'a, W> { self.write_expr(value, ctx)?; writeln!(self.out, ";")? } + // Stores a value into an image. + Statement::ImageStore { + image, + coordinate, + array_index, + value, + } => { + // This will only panic if the module is invalid + let dim = match *ctx.typifier.get(image, &self.module.types) { + TypeInner::Image { dim, .. } => dim, + _ => unreachable!(), + }; + + write!(self.out, "imageStore(")?; + self.write_expr(image, ctx)?; + write!(self.out, ", ")?; + self.write_texture_coordinates(coordinate, array_index, dim, ctx)?; + write!(self.out, ", ")?; + self.write_expr(value, ctx)?; + writeln!(self.out, ");")?; + } // A `Call` is written `name(arguments)` where `arguments` is a comma separated expressions list Statement::Call { function, @@ -1136,7 +1157,7 @@ impl<'a, W: Write> Writer<'a, W> { } => { write!(self.out, "{}(", &self.names[&NameKey::Function(function)])?; self.write_slice(arguments, |this, _, arg| this.write_expr(*arg, ctx))?; - write!(self.out, ");")? + writeln!(self.out, ");")? } } @@ -1162,7 +1183,7 @@ impl<'a, W: Write> Writer<'a, W> { Expression::AccessIndex { base, index } => { self.write_expr(base, ctx)?; - match ctx.typifier.get(base, &self.module.types) { + match *ctx.typifier.get(base, &self.module.types) { TypeInner::Vector { .. } | TypeInner::Matrix { .. } | TypeInner::Array { .. } => write!(self.out, "[{}]", index)?, @@ -1306,7 +1327,7 @@ impl<'a, W: Write> Writer<'a, W> { index, } => { // This will only panic if the module is invalid - let (dim, class) = match ctx.typifier.get(image, &self.module.types) { + let (dim, class) = match *ctx.typifier.get(image, &self.module.types) { TypeInner::Image { dim, arrayed: _, @@ -1324,25 +1345,8 @@ impl<'a, W: Write> Writer<'a, W> { write!(self.out, "{}(", fun_name)?; self.write_expr(image, ctx)?; - match array_index { - Some(layer_expr) => { - let tex_coord_type = match dim { - crate::ImageDimension::D1 => "ivec2", - crate::ImageDimension::D2 => "ivec3", - crate::ImageDimension::D3 => "ivec4", - crate::ImageDimension::Cube => "ivec4", - }; - write!(self.out, ", {}(", tex_coord_type)?; - self.write_expr(coordinate, ctx)?; - write!(self.out, ", ")?; - self.write_expr(layer_expr, ctx)?; - write!(self.out, ")")?; - } - None => { - write!(self.out, ", ")?; - self.write_expr(coordinate, ctx)?; - } - } + write!(self.out, ", ")?; + self.write_texture_coordinates(coordinate, array_index, dim, ctx)?; if let Some(index_expr) = index { write!(self.out, ", ")?; @@ -1356,7 +1360,7 @@ impl<'a, W: Write> Writer<'a, W> { // - textureSamples/imageSamples Expression::ImageQuery { image, query } => { // This will only panic if the module is invalid - let (dim, class) = match ctx.typifier.get(image, &self.module.types) { + let (dim, class) = match *ctx.typifier.get(image, &self.module.types) { TypeInner::Image { dim, arrayed: _, @@ -1710,6 +1714,34 @@ impl<'a, W: Write> Writer<'a, W> { Ok(()) } + fn write_texture_coordinates( + &mut self, + coordinate: Handle, + array_index: Option>, + dim: crate::ImageDimension, + ctx: &FunctionCtx, + ) -> Result<(), Error> { + match array_index { + Some(layer_expr) => { + let tex_coord_type = match dim { + crate::ImageDimension::D1 => "ivec2", + crate::ImageDimension::D2 => "ivec3", + crate::ImageDimension::D3 => "ivec4", + crate::ImageDimension::Cube => "ivec4", + }; + write!(self.out, "{}(", tex_coord_type)?; + self.write_expr(coordinate, ctx)?; + write!(self.out, ", ")?; + self.write_expr(layer_expr, ctx)?; + write!(self.out, ")")?; + } + None => { + self.write_expr(coordinate, ctx)?; + } + } + Ok(()) + } + /// Helper method used to produce the images mapping that's returned to the user /// /// It takes an iterator of [`Function`](crate::Function) references instead of diff --git a/src/back/msl/writer.rs b/src/back/msl/writer.rs index c88f71b024..d51a9e5d05 100644 --- a/src/back/msl/writer.rs +++ b/src/back/msl/writer.rs @@ -909,6 +909,23 @@ impl Writer { self.put_expression(value, &context.expression)?; writeln!(self.out, ";")?; } + crate::Statement::ImageStore { + image, + coordinate, + array_index, + value, + } => { + self.put_expression(image, &context.expression)?; + write!(self.out, ".write(")?; + self.put_expression(value, &context.expression)?; + write!(self.out, ", ")?; + self.put_expression(coordinate, &context.expression)?; + if let Some(expr) = array_index { + write!(self.out, ", ")?; + self.put_expression(expr, &context.expression)?; + } + write!(self.out, ");")?; + } crate::Statement::Call { function, ref arguments, diff --git a/src/back/spv/instructions.rs b/src/back/spv/instructions.rs index a4fe76706e..82af2d35f1 100644 --- a/src/back/spv/instructions.rs +++ b/src/back/spv/instructions.rs @@ -233,7 +233,7 @@ impl super::Instruction { instruction.add_operand(depth as u32); instruction.add_operand(arrayed as u32); instruction.add_operand(multi as u32); - instruction.add_operand(sampled as u32); + instruction.add_operand(if sampled { 1 } else { 2 }); let format = match image_class { crate::ImageClass::Storage(format) => match format { @@ -552,6 +552,28 @@ impl super::Instruction { instruction } + pub(super) fn image_read( + result_type_id: Word, + id: Word, + image: Word, + coordinates: Word, + ) -> Self { + let mut instruction = Self::new(Op::ImageRead); + instruction.set_type(result_type_id); + instruction.set_result(id); + instruction.add_operand(image); + instruction.add_operand(coordinates); + instruction + } + + pub(super) fn image_write(image: Word, coordinates: Word, value: Word) -> Self { + let mut instruction = Self::new(Op::ImageWrite); + instruction.add_operand(image); + instruction.add_operand(coordinates); + instruction.add_operand(value); + instruction + } + // // Conversion Instructions // diff --git a/src/back/spv/writer.rs b/src/back/spv/writer.rs index 709dd8b922..f5d7fa1e67 100644 --- a/src/back/spv/writer.rs +++ b/src/back/spv/writer.rs @@ -1,6 +1,9 @@ /*! Standard Portable Intermediate Representation (SPIR-V) backend !*/ use super::{Instruction, LogicalLayout, PhysicalLayout, WriterFlags}; -use crate::proc::{Layouter, ResolveContext, ResolveError, Typifier}; +use crate::{ + arena::{Arena, Handle}, + proc::{Layouter, ResolveContext, ResolveError, Typifier}, +}; use spirv::Word; use std::collections::hash_map::Entry; use thiserror::Error; @@ -47,7 +50,7 @@ enum RawExpression { struct Function { signature: Option, parameters: Vec, - variables: crate::FastHashMap, LocalVariable>, + variables: crate::FastHashMap, LocalVariable>, blocks: Vec, } @@ -94,11 +97,11 @@ enum LocalType { width: crate::Bytes, }, Pointer { - base: crate::Handle, + base: Handle, class: crate::StorageClass, }, SampledImage { - image_type: crate::Handle, + image_type: Handle, }, } @@ -126,7 +129,7 @@ impl LocalType { #[derive(Debug, PartialEq, Hash, Eq, Copy, Clone)] enum LookupType { - Handle(crate::Handle), + Handle(Handle), Local(LocalType), } @@ -183,14 +186,14 @@ pub struct Writer { flags: WriterFlags, void_type: u32, lookup_type: crate::FastHashMap, - lookup_function: crate::FastHashMap, Word>, + lookup_function: crate::FastHashMap, Word>, lookup_function_type: crate::FastHashMap, - lookup_constant: crate::FastHashMap, Word>, + lookup_constant: crate::FastHashMap, Word>, lookup_global_variable: - crate::FastHashMap, (Word, spirv::StorageClass)>, + crate::FastHashMap, (Word, spirv::StorageClass)>, // TODO: this is a type property that depends on the global variable that uses it // so it may require us to duplicate the type! - struct_type_handles: crate::FastHashMap, crate::StorageAccess>, + struct_type_handles: crate::FastHashMap, crate::StorageAccess>, gl450_ext_inst_id: Word, layouter: Layouter, typifier: Typifier, @@ -248,7 +251,7 @@ impl Writer { fn get_type_id( &mut self, - arena: &crate::Arena, + arena: &Arena, lookup_ty: LookupType, ) -> Result { if let Entry::Occupied(e) = self.lookup_type.entry(lookup_ty) { @@ -284,7 +287,7 @@ impl Writer { fn get_global_variable_id( &mut self, ir_module: &crate::Module, - handle: crate::Handle, + handle: Handle, ) -> Result<(Word, spirv::StorageClass), Error> { Ok(match self.lookup_global_variable.entry(handle) { Entry::Occupied(e) => *e.get(), @@ -299,8 +302,8 @@ impl Writer { fn get_pointer_id( &mut self, - arena: &crate::Arena, - handle: crate::Handle, + arena: &Arena, + handle: Handle, class: crate::StorageClass, ) -> Result { let ty_id = self.get_type_id(arena, LookupType::Handle(handle))?; @@ -539,7 +542,7 @@ impl Writer { fn write_type_declaration_local( &mut self, - arena: &crate::Arena, + arena: &Arena, local_ty: LocalType, ) -> Result { let id = self.generate_id(); @@ -581,8 +584,8 @@ impl Writer { fn write_type_declaration_arena( &mut self, - arena: &crate::Arena, - handle: crate::Handle, + arena: &Arena, + handle: Handle, ) -> Result { let ty = &arena[handle]; let id = self.generate_id(); @@ -768,7 +771,7 @@ impl Writer { &mut self, id: Word, inner: &crate::ConstantInner, - types: &crate::Arena, + types: &Arena, ) -> Result<(), Error> { let instruction = match *inner { crate::ConstantInner::Scalar { width, ref value } => { @@ -847,7 +850,7 @@ impl Writer { fn write_global_variable( &mut self, ir_module: &crate::Module, - handle: crate::Handle, + handle: Handle, ) -> Result<(Instruction, Word, spirv::StorageClass), Error> { let global_variable = &ir_module.global_variables[handle]; let id = self.generate_id(); @@ -988,12 +991,93 @@ impl Writer { id } + fn write_texture_coordinates( + &mut self, + ir_module: &crate::Module, + ir_function: &crate::Function, + coordinates: Handle, + array_index: Option>, + block: &mut Block, + function: &mut Function, + ) -> Result { + let coordinate_id = + self.write_expression(ir_module, ir_function, coordinates, block, function)?; + + Ok(if let Some(array_index) = array_index { + let coordinate_scalar_type_id = self.get_type_id( + &ir_module.types, + LookupType::Local(LocalType::Scalar { + kind: crate::ScalarKind::Float, + width: 4, + }), + )?; + + let mut constituent_ids = [0u32; 4]; + let size = match *self.typifier.get(coordinates, &ir_module.types) { + crate::TypeInner::Scalar { .. } => { + constituent_ids[0] = coordinate_id; + crate::VectorSize::Bi + } + crate::TypeInner::Vector { size, .. } => { + for i in 0..size as u32 { + let id = self.generate_id(); + constituent_ids[i as usize] = id; + block.body.push(Instruction::composite_extract( + coordinate_scalar_type_id, + id, + coordinate_id, + &[i], + )); + } + match size { + crate::VectorSize::Bi => crate::VectorSize::Tri, + crate::VectorSize::Tri => crate::VectorSize::Quad, + crate::VectorSize::Quad => { + unimplemented!("Unable to extend the vec4 coordinate") + } + } + } + ref other => unimplemented!("wrong coordinate type {:?}", other), + }; + + let array_index_f32_id = self.generate_id(); + constituent_ids[size as usize - 1] = array_index_f32_id; + + let array_index_u32_id = + self.write_expression(ir_module, ir_function, array_index, block, function)?; + let cast_instruction = Instruction::unary( + spirv::Op::ConvertUToF, + coordinate_scalar_type_id, + array_index_f32_id, + array_index_u32_id, + ); + block.body.push(cast_instruction); + + let extended_coordinate_type_id = self.get_type_id( + &ir_module.types, + LookupType::Local(LocalType::Vector { + size, + kind: crate::ScalarKind::Float, + width: 4, + }), + )?; + + self.write_composite_construct( + extended_coordinate_type_id, + &constituent_ids[..size as usize], + block, + ) + } else { + coordinate_id + }) + } + /// Write an expression and return a value ID. fn write_expression<'a>( &mut self, ir_module: &'a crate::Module, ir_function: &crate::Function, - handle: crate::Handle, + handle: Handle, block: &mut Block, function: &mut Function, ) -> Result { @@ -1016,7 +1100,7 @@ impl Writer { &mut self, ir_module: &'a crate::Module, ir_function: &crate::Function, - handle: crate::Handle, + handle: Handle, block: &mut Block, function: &mut Function, ) -> Result { @@ -1036,7 +1120,7 @@ impl Writer { &mut self, ir_module: &'a crate::Module, ir_function: &crate::Function, - expr_handle: crate::Handle, + expr_handle: Handle, block: &mut Block, function: &mut Function, ) -> Result<(RawExpression, ExpressionId), Error> { @@ -1510,6 +1594,44 @@ impl Writer { RawExpression::Value(id) } + crate::Expression::ImageLoad { + image, + coordinate, + array_index, + index, + } => { + let image_id = + self.write_expression(ir_module, ir_function, image, block, function)?; + let coordinate_id = self.write_texture_coordinates( + ir_module, + ir_function, + coordinate, + array_index, + block, + function, + )?; + + let id = self.generate_id(); + let mut instruction = + Instruction::image_read(result_type_id, id, image_id, coordinate_id); + + if let Some(index) = index { + let index_id = + self.write_expression(ir_module, ir_function, index, block, function)?; + let image_ops = match *self.typifier.get(image, &ir_module.types) { + crate::TypeInner::Image { + class: crate::ImageClass::Sampled { multi: true, .. }, + .. + } => spirv::ImageOperands::SAMPLE, + _ => spirv::ImageOperands::LOD, + }; + instruction.add_operand(image_ops.bits()); + instruction.add_operand(index_id); + } + + block.body.push(instruction); + RawExpression::Value(id) + } crate::Expression::ImageSample { image, sampler, @@ -1531,84 +1653,16 @@ impl Writer { LookupType::Local(LocalType::SampledImage { image_type }), )?; - // sampler let sampler_id = self.write_expression(ir_module, ir_function, sampler, block, function)?; - - // coordinate - let mut coordinate_id = - self.write_expression(ir_module, ir_function, coordinate, block, function)?; - - if let Some(array_index) = array_index { - let coordinate_scalar_type_id = self.get_type_id( - &ir_module.types, - LookupType::Local(LocalType::Scalar { - kind: crate::ScalarKind::Float, - width: 4, - }), - )?; - - let mut constituent_ids = [0u32; 4]; - let size = match *self.typifier.get(coordinate, &ir_module.types) { - crate::TypeInner::Scalar { .. } => { - constituent_ids[0] = coordinate_id; - crate::VectorSize::Bi - } - crate::TypeInner::Vector { size, .. } => { - for i in 0..size as u32 { - let id = self.generate_id(); - constituent_ids[i as usize] = id; - block.body.push(Instruction::composite_extract( - coordinate_scalar_type_id, - id, - coordinate_id, - &[i], - )); - } - match size { - crate::VectorSize::Bi => crate::VectorSize::Tri, - crate::VectorSize::Tri => crate::VectorSize::Quad, - crate::VectorSize::Quad => { - unimplemented!("Unable to extend the vec4 coordinate") - } - } - } - ref other => unimplemented!("wrong coordinate type {:?}", other), - }; - - let array_index_f32_id = self.generate_id(); - constituent_ids[size as usize - 1] = array_index_f32_id; - - let array_index_u32_id = self.write_expression( - ir_module, - ir_function, - array_index, - block, - function, - )?; - let cast_instruction = Instruction::unary( - spirv::Op::ConvertUToF, - coordinate_scalar_type_id, - array_index_f32_id, - array_index_u32_id, - ); - block.body.push(cast_instruction); - - let extended_coordinate_type_id = self.get_type_id( - &ir_module.types, - LookupType::Local(LocalType::Vector { - size, - kind: crate::ScalarKind::Float, - width: 4, - }), - )?; - - coordinate_id = self.write_composite_construct( - extended_coordinate_type_id, - &constituent_ids[..size as usize], - block, - ); - } + let coordinate_id = self.write_texture_coordinates( + ir_module, + ir_function, + coordinate, + array_index, + block, + function, + )?; let sampled_image_id = self.generate_id(); block.body.push(Instruction::sampled_image( @@ -1995,6 +2049,29 @@ impl Writer { .body .push(Instruction::store(pointer_id, value_id, None)); } + crate::Statement::ImageStore { + image, + coordinate, + array_index, + value, + } => { + let image_id = + self.write_expression(ir_module, ir_function, image, &mut block, function)?; + let coordinate_id = self.write_texture_coordinates( + ir_module, + ir_function, + coordinate, + array_index, + &mut block, + function, + )?; + let value_id = + self.write_expression(ir_module, ir_function, value, &mut block, function)?; + + block + .body + .push(Instruction::image_write(image_id, coordinate_id, value_id)); + } crate::Statement::Call { function: local_function, ref arguments, diff --git a/src/front/spv/mod.rs b/src/front/spv/mod.rs index a57d031a58..3de5bebf4f 100644 --- a/src/front/spv/mod.rs +++ b/src/front/spv/mod.rs @@ -1582,7 +1582,12 @@ impl> Parser { self.patch_function_call_statements(body)?; self.patch_function_call_statements(continuing)?; } - S::Break | S::Continue | S::Return { .. } | S::Kill | S::Store { .. } => {} + S::Break + | S::Continue + | S::Return { .. } + | S::Kill + | S::Store { .. } + | S::ImageStore { .. } => {} S::Call { ref mut function, .. } => { diff --git a/src/front/wgsl/mod.rs b/src/front/wgsl/mod.rs index 551c9da6cf..e4febf8a08 100644 --- a/src/front/wgsl/mod.rs +++ b/src/front/wgsl/mod.rs @@ -1798,6 +1798,32 @@ impl Parser { "break" => Some(crate::Statement::Break), "continue" => Some(crate::Statement::Continue), "discard" => Some(crate::Statement::Kill), + "textureStore" => { + lexer.expect(Token::Paren('('))?; + let image_name = lexer.next_ident()?; + let image = context.lookup_ident.lookup(image_name)?; + lexer.expect(Token::Separator(','))?; + let coordinate = self.parse_general_expression(lexer, context.as_expression())?; + let arrayed = match *context.as_expression().resolve_type(image)? { + crate::TypeInner::Image { arrayed, .. } => arrayed, + _ => return Err(Error::BadTexture(image_name)), + }; + let array_index = if arrayed { + lexer.expect(Token::Separator(','))?; + Some(self.parse_general_expression(lexer, context.as_expression())?) + } else { + None + }; + lexer.expect(Token::Separator(','))?; + let value = self.parse_general_expression(lexer, context.as_expression())?; + lexer.expect(Token::Paren(')'))?; + Some(crate::Statement::ImageStore { + image, + coordinate, + array_index, + value, + }) + } // assignment or a function call ident => { if let Some(&var_expr) = context.lookup_ident.get(ident) { diff --git a/src/lib.rs b/src/lib.rs index c87af3c533..38d2157cc7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -805,6 +805,13 @@ pub enum Statement { pointer: Handle, value: Handle, }, + /// Stores a value to an image. + ImageStore { + image: Handle, + coordinate: Handle, + array_index: Option>, + value: Handle, + }, /// Calls a function with no return value. Call { function: Handle, diff --git a/src/proc/analyzer.rs b/src/proc/analyzer.rs index 8c03f0623e..b35f9c6806 100644 --- a/src/proc/analyzer.rs +++ b/src/proc/analyzer.rs @@ -335,6 +335,21 @@ impl FunctionInfo { ControlFlags::MAY_EXIT | flags } S::Store { pointer, value } => self.add_ref(pointer) | self.add_ref(value), + S::ImageStore { + image, + coordinate, + array_index, + value, + } => { + let array_flags = match array_index { + Some(expr) => self.add_ref(expr), + None => ControlFlags::empty(), + }; + array_flags + | self.add_ref(image) + | self.add_ref(coordinate) + | self.add_ref(value) + } S::Call { function, ref arguments, diff --git a/src/proc/interface.rs b/src/proc/interface.rs index ca6804388c..051d42e1be 100644 --- a/src/proc/interface.rs +++ b/src/proc/interface.rs @@ -203,6 +203,19 @@ where self.visitor.visit_lhs_expr(left, &self.expressions[left]); self.traverse_expr(value); } + S::ImageStore { + image, + coordinate, + array_index, + value, + } => { + self.visitor.visit_lhs_expr(image, &self.expressions[image]); + self.traverse_expr(coordinate); + if let Some(expr) = array_index { + self.traverse_expr(expr); + } + self.traverse_expr(value); + } S::Call { function, ref arguments, diff --git a/src/proc/terminator.rs b/src/proc/terminator.rs index 181c815323..af5eaf54dd 100644 --- a/src/proc/terminator.rs +++ b/src/proc/terminator.rs @@ -35,6 +35,7 @@ pub fn ensure_block_returns(block: &mut crate::Block) { | Some(&mut crate::Statement::Kill) => (), Some(&mut crate::Statement::Loop { .. }) | Some(&mut crate::Statement::Store { .. }) + | Some(&mut crate::Statement::ImageStore { .. }) | Some(&mut crate::Statement::Call { .. }) | None => block.push(crate::Statement::Return { value: None }), }