Created
December 12, 2023 05:43
-
-
Save chyyran/0df5fbf7638dc932ec128c150779e498 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
use std::borrow::Cow; | |
use std::collections::hash_map::Entry; | |
use rspirv::dr::{Builder, InsertPoint, Instruction, Module, Operand}; | |
use rustc_hash::{FxHashMap, FxHashSet}; | |
use spirv::{Decoration, Op, StorageClass, Word}; | |
pub struct Pass<'a> { | |
pub builder: &'a mut Builder, | |
} | |
#[derive(Clone, Debug)] | |
pub struct CombinedImageSampler { | |
texture_variable: spirv::Word, | |
sampler_variable: spirv::Word, | |
sampler_pointer_type: spirv::Word, | |
original_uniform_pointer_type_id: spirv::Word, | |
original_uniform_type: Instruction, | |
target_texture_type_id: spirv::Word, | |
target_texture_pointer_type_id: spirv::Word, | |
// Always OpTypeImage | |
base_type: Instruction, | |
} | |
#[derive(Debug)] | |
pub struct OpAccessChain<'a> { | |
sampled_image: &'a CombinedImageSampler, | |
index: Operand, | |
original_result_type: Word, | |
target_pointer_type: Word, | |
pub result_id: Word, | |
} | |
#[derive(Debug)] | |
pub struct OpFunctionCallParam<'a> { | |
sampled_image: &'a CombinedImageSampler, | |
type_handle: spirv::Word, | |
} | |
impl<'a> Pass<'a> { | |
pub fn new(builder: &'a mut Builder) -> Self { | |
let mut val = Self { builder }; | |
val | |
} | |
fn do_pass(&mut self) { | |
let mut combined_image_samplers = self.collect_global_sampled_images(); | |
self.retype_combined_image_sampler_uniforms(&combined_image_samplers); | |
self.put_variables_to_end(); | |
// First rewrite global loads | |
self.rewrite_loads(&combined_image_samplers); | |
let mut op_access_chains = self.collect_op_access_chain(&combined_image_samplers); | |
self.rewrite_global_op_access_chain_loads(&op_access_chains); | |
// Collect functions that reference combined image samplers indirectly (todo) (also rewrite the found OpFunctionCall) | |
while !combined_image_samplers.is_empty() { | |
let op_functions = self.rewrite_function_calls(&op_access_chains, &combined_image_samplers); | |
combined_image_samplers = self.rewrite_functions_definitions(&op_functions); | |
self.rewrite_loads(&combined_image_samplers); | |
op_access_chains = self.collect_op_access_chain(&combined_image_samplers); | |
eprintln!("{:?}", &op_access_chains); | |
self.rewrite_global_op_access_chain_loads(&op_access_chains); | |
} | |
// notes: OpFunctionCall %void %assign_s21_ %tex (i.e for scalar no OpLoad precedes, OpLoad and OpAccessChain is in function), this is true even when passing array of samplers | |
// However, if passing single sampler from array, OpAccessChain happens in outer scope. | |
// Rewrite functions that do so by adding sampler param | |
// rewrite loads + op_access_chains for functions | |
} | |
fn put_variables_to_end(&mut self) { | |
// this is easier than doing proper topo sort. | |
let mut vars = Vec::new(); | |
self.builder.module_mut() | |
.types_global_values.retain(|instr| { | |
if instr.class.opcode == spirv::Op::Variable { | |
vars.push(instr.clone()); | |
return false; | |
}; | |
true | |
}); | |
self.builder.module_mut() | |
.types_global_values.append(&mut vars); | |
} | |
fn ensure_op_type_sampler(&mut self) { | |
self.builder.type_sampler(); | |
} | |
fn find_instruction(&self, word: Word) -> Option<&Instruction> { | |
self.builder | |
.module_ref() | |
.all_inst_iter() | |
.find(|i| i.result_id == Some(word)) | |
} | |
fn find_global_instruction(&self, word: Word) -> Option<&Instruction> { | |
self.builder | |
.module_ref() | |
.global_inst_iter() | |
.find(|i| i.result_id == Some(word)) | |
} | |
fn find_global_instruction_mut(&mut self, word: Word) -> Option<&mut Instruction> { | |
self.builder | |
.module_mut() | |
.global_inst_iter_mut() | |
.find(|i| i.result_id == Some(word)) | |
} | |
fn create_sampler_name(&mut self, word: Word) -> Option<String> { | |
self.builder.module_ref().debug_names.iter().find_map(|i| { | |
if i.class.opcode != spirv::Op::Name { | |
return None; | |
} | |
let Some(&Operand::IdRef(target)) = &i.operands.get(0) else { | |
return None; | |
}; | |
if target != word { | |
return None; | |
} | |
let Some(Operand::LiteralString(string)) = &i.operands.get(1) else { | |
return None; | |
}; | |
return Some(format!("_{string}_sampler")); | |
}) | |
} | |
fn get_base_type_for_sampled_image( | |
&'a self, | |
mut inst: &'a Instruction, | |
) -> Option<&'a Instruction> { | |
if inst.class.opcode != spirv::Op::TypeSampledImage { | |
return None; | |
} | |
let Some(&Operand::IdRef(referand)) = inst.operands.get(0) else { | |
return None; | |
}; | |
self.find_global_instruction(referand) | |
} | |
// Create a sampler OpVariable at the selected block | |
fn create_sampler_uniform( | |
&mut self, | |
uniform_type: spirv::Word, | |
combined_image_sampler: spirv::Word, | |
) -> (spirv::Word, spirv::Word) { | |
let sampler_pointer_type = | |
self.builder | |
.type_pointer(None, StorageClass::UniformConstant, uniform_type); | |
let sampler_uniform = self.builder.variable( | |
sampler_pointer_type, | |
None, | |
StorageClass::UniformConstant, | |
None, | |
); | |
let decorations: Vec<Instruction> = self | |
.builder | |
.module_ref() | |
.annotations | |
.iter() | |
.filter_map(|f| { | |
if f.class.opcode == spirv::Op::Decorate | |
&& f.operands[0] == Operand::IdRef(combined_image_sampler) | |
{ | |
Some(f.clone()) | |
} else { | |
None | |
} | |
}) | |
.collect(); | |
if let Some(name) = self.create_sampler_name(combined_image_sampler) { | |
self.builder.name(sampler_uniform, name); | |
} | |
// Clone decorations to the created sampler | |
for decoration in decorations { | |
let Operand::Decoration(decoration_type) = decoration.operands[1] else { | |
continue; | |
}; | |
self.builder.decorate( | |
sampler_uniform, | |
decoration_type, | |
decoration.operands[2..].iter().map(|f| f.clone()), | |
) | |
} | |
(sampler_pointer_type, sampler_uniform) | |
} | |
fn collect_global_sampled_images(&mut self) -> FxHashMap<spirv::Word, CombinedImageSampler> { | |
let mut image_sampler_cadidates = Vec::new(); | |
let mut image_sampler_types = FxHashMap::default(); | |
for global in self.builder.module_ref().types_global_values.iter() { | |
if global.class.opcode == spirv::Op::Variable | |
&& global.operands[0] == Operand::StorageClass(StorageClass::UniformConstant) | |
{ | |
let pointer_type = global.result_type; | |
let Some(pointer_type) = pointer_type else { | |
continue; | |
}; | |
image_sampler_cadidates.push((pointer_type, global.result_id)) | |
} | |
} | |
for (pointer_type_id, global_variable) in image_sampler_cadidates { | |
let Some(pointer_type) = self.find_global_instruction(pointer_type_id).cloned() else { | |
continue; | |
}; | |
if pointer_type.class.opcode == spirv::Op::TypePointer | |
&& pointer_type.operands[0] == Operand::StorageClass(StorageClass::UniformConstant) | |
{ | |
let Some(&Operand::IdRef(sampled_image_type)) = pointer_type.operands.get(1) else { | |
continue; | |
}; | |
let Some(uniform_type) = self.find_global_instruction(sampled_image_type).cloned() | |
else { | |
continue; | |
}; | |
if uniform_type.class.opcode == spirv::Op::TypeSampledImage { | |
let Some(base_type) = | |
self.get_base_type_for_sampled_image(&uniform_type).cloned() | |
else { | |
continue; | |
}; | |
let Some(combined_image_sampler) = global_variable else { | |
continue; | |
}; | |
// insert the sampler | |
if base_type.class.opcode != spirv::Op::TypeImage { | |
continue; | |
} | |
let Some(base_type_id) = base_type.result_id else { | |
continue; | |
}; | |
let sampler_type = self.builder.type_sampler(); | |
let (sampler_pointer_type, sampler_uniform) = | |
self.create_sampler_uniform(sampler_type, combined_image_sampler); | |
image_sampler_types.insert( | |
combined_image_sampler, | |
CombinedImageSampler { | |
texture_variable: combined_image_sampler, | |
sampler_variable: sampler_uniform, | |
original_uniform_type: uniform_type, | |
target_texture_type_id: base_type_id, | |
original_uniform_pointer_type_id: pointer_type_id, | |
base_type, | |
sampler_pointer_type, | |
target_texture_pointer_type_id: pointer_type_id, | |
}, | |
); | |
continue; | |
} | |
if uniform_type.class.opcode == spirv::Op::TypeArray { | |
let Some(&Operand::IdRef(array_base_type)) = uniform_type.operands.get(0) | |
else { | |
continue; | |
}; | |
let Some(&Operand::IdRef(array_length)) = uniform_type.operands.get(1) else { | |
continue; | |
}; | |
let Some(sampled_image_type) = | |
self.find_global_instruction(array_base_type).cloned() | |
else { | |
continue; | |
}; | |
let Some(base_type) = self | |
.get_base_type_for_sampled_image(&sampled_image_type) | |
.cloned() | |
else { | |
continue; | |
}; | |
let Some(combined_image_sampler) = global_variable else { | |
continue; | |
}; | |
// insert the sampler | |
if base_type.class.opcode != spirv::Op::TypeImage { | |
continue; | |
} | |
let sampler_type = self.builder.type_sampler(); | |
let sampler_array_type = self.builder.type_array(sampler_type, array_length); | |
let (sampler_pointer_type, sampler_uniform) = | |
self.create_sampler_uniform(sampler_array_type, combined_image_sampler); | |
// insert target types | |
let Some(base_type_id) = base_type.result_id else { | |
continue; | |
}; | |
let target_texture_type_id = self.builder | |
.type_array(base_type_id, array_length); | |
let target_texture_pointer_type_id = self.builder | |
.type_pointer(None, StorageClass::UniformConstant, target_texture_type_id); | |
image_sampler_types.insert( | |
combined_image_sampler, | |
CombinedImageSampler { | |
texture_variable: combined_image_sampler, | |
sampler_variable: sampler_uniform, | |
original_uniform_type: uniform_type, | |
target_texture_type_id, | |
original_uniform_pointer_type_id: pointer_type_id, | |
base_type, | |
sampler_pointer_type, | |
target_texture_pointer_type_id, | |
}, | |
); | |
continue; | |
} | |
} | |
} | |
image_sampler_types | |
} | |
fn retype_combined_image_sampler_uniforms( | |
&mut self, | |
combined_image_samplers: &FxHashMap<spirv::Word, CombinedImageSampler>, | |
) { | |
// Need to rebuild the global instructions because we need to insert new types... | |
let mut instructions = Vec::new(); | |
for instr in self.builder.module_ref().types_global_values.clone() { | |
let Some(result_id) = instr.result_id else { | |
instructions.push(instr); | |
continue; | |
}; | |
let Some(sampled_image) = combined_image_samplers.get(&result_id) else { | |
// We need to fix.. | |
instructions.push(instr); | |
continue; | |
}; | |
let Some(base_type_id) = sampled_image.base_type.result_id else { | |
instructions.push(instr); | |
continue; | |
}; | |
// If it's a OpTypeSampledImage, we want to change the variable type to &TypeImage. | |
if sampled_image.original_uniform_type.class.opcode == spirv::Op::TypeSampledImage { | |
// keep labels in sync | |
let mut op_variable = instr; | |
op_variable.result_type = Some(sampled_image.target_texture_pointer_type_id); | |
instructions.push(op_variable); | |
continue; | |
} | |
// Re-type array globals. | |
// We don't need to worry about the pointer type of the load, as | |
// we can instantiate that later. | |
if sampled_image.original_uniform_type.class.opcode == spirv::Op::TypeArray { | |
// let Some(&Operand::IdRef(array_length)) = | |
// sampled_image.original_uniform_type.operands.get(1) | |
// else { | |
// instructions.push(instr); | |
// continue; | |
// }; | |
let mut op_variable = instr; | |
op_variable.result_type = Some(sampled_image.target_texture_pointer_type_id); | |
instructions.push(op_variable); | |
} | |
} | |
// replace | |
self.builder.module_mut().types_global_values = instructions; | |
} | |
fn rewrite_loads( | |
&mut self, | |
combined_image_samplers: &FxHashMap<spirv::Word, CombinedImageSampler>, | |
) { | |
let op_type_sampler = self.builder.type_sampler(); | |
// need to clone | |
let mut functions = self.builder.module_ref().functions.clone(); | |
for function in functions.iter_mut() { | |
for block in function.blocks.iter_mut() { | |
let mut instructions = Vec::new(); | |
for instr in block.instructions.drain(..) { | |
if instr.class.opcode != Op::Load { | |
instructions.push(instr); | |
continue; | |
} | |
// This doesn't affect array loads because array loads load the result of the OpAccessChain which can be done in a separate pass. | |
let Some(Operand::IdRef(op_variable_id)) = &instr.operands.get(0) else { | |
instructions.push(instr); | |
continue; | |
}; | |
let Some(combined_image_sampler) = combined_image_samplers.get(op_variable_id) | |
else { | |
// Ignore what we didn't process before. | |
instructions.push(instr); | |
continue; | |
}; | |
let op_load_texture_id = self.builder.id(); | |
let op_load_sampler_id = self.builder.id(); | |
let mut load_instr = instr.clone(); | |
load_instr.result_type = combined_image_sampler.base_type.result_id; | |
load_instr.result_id = Some(op_load_texture_id); | |
instructions.push(load_instr); | |
// load the sampler | |
instructions.push(Instruction { | |
class: rspirv::grammar::CoreInstructionTable::get(spirv::Op::Load), | |
result_type: Some(op_type_sampler), | |
result_id: Some(op_load_sampler_id), | |
operands: vec![Operand::IdRef(combined_image_sampler.sampler_variable)], | |
}); | |
// reuse the old id for the OpSampleImage | |
instructions.push(Instruction { | |
class: rspirv::grammar::CoreInstructionTable::get(spirv::Op::SampledImage), | |
result_type: combined_image_sampler.original_uniform_type.result_id, | |
result_id: instr.result_id, | |
operands: vec![ | |
Operand::IdRef(op_load_texture_id), | |
Operand::IdRef(op_load_sampler_id), | |
], | |
}); | |
} | |
block.instructions = instructions; | |
} | |
} | |
self.builder.module_mut().functions = functions; | |
} | |
fn collect_op_access_chain<'b>( | |
&mut self, | |
combined_image_samplers: &'b FxHashMap<spirv::Word, CombinedImageSampler>, | |
) -> FxHashMap<spirv::Word, OpAccessChain<'b>> { | |
// need to clone | |
let mut functions = self.builder.module_ref().functions.clone(); | |
let mut seen_op_access_chain = FxHashMap::default(); | |
for function in functions.iter_mut() { | |
for block in function.blocks.iter_mut() { | |
let mut instructions = Vec::new(); | |
// This needs to be done in two passes, first to find the OpAccessChains and | |
// update their type to the pointer type of the scalar value, | |
// Then to update the OpLoads. | |
for instr in block.instructions.clone() { | |
if instr.class.opcode != Op::AccessChain { | |
instructions.push(instr); | |
continue; | |
} | |
// Need to refer to it later | |
let Some(result_id) = instr.result_id else { | |
instructions.push(instr); | |
continue; | |
}; | |
// This doesn't affect array loads because array loads load the result of the OpAccessChain which can be done in a separate pass. | |
let Some(Operand::IdRef(op_variable)) = &instr.operands.get(0) else { | |
instructions.push(instr); | |
continue; | |
}; | |
// Ensure that this is an OpAccessChain of our combined image sampler. | |
let Some(sampled_image) = combined_image_samplers.get(op_variable) else { | |
// Ignore what we didn't process before. | |
instructions.push(instr); | |
continue; | |
}; | |
let Some(original_result_type) = instr | |
.result_type else { | |
// Ignore what we didn't process before. | |
instructions.push(instr); | |
continue; | |
}; | |
let Some(index) = instr | |
.operands.get(1).cloned() else { | |
// Ignore what we didn't process before. | |
instructions.push(instr); | |
continue; | |
}; | |
let Some(base_type_id) = sampled_image.base_type.result_id else { | |
// Ignore what we didn't process before. | |
instructions.push(instr); | |
continue; | |
}; | |
let op_pointer_type_id = self.builder.type_pointer( | |
None, | |
StorageClass::UniformConstant, | |
base_type_id, | |
); | |
let mut op_access_chain = instr; | |
let Some(result_id) = op_access_chain.result_id else { | |
instructions.push(op_access_chain); | |
continue; | |
}; | |
op_access_chain.result_type = Some(op_pointer_type_id); | |
instructions.push(op_access_chain); | |
seen_op_access_chain.insert( | |
result_id, | |
OpAccessChain { | |
sampled_image, | |
index, | |
original_result_type, | |
target_pointer_type: op_pointer_type_id, | |
result_id | |
}, | |
); | |
} | |
block.instructions = instructions; | |
} | |
} | |
self.builder.module_mut().functions = functions; | |
seen_op_access_chain | |
} | |
fn rewrite_global_op_access_chain_loads( | |
&mut self, | |
op_access_chains: &FxHashMap<spirv::Word, OpAccessChain>, | |
) { | |
let op_type_sampler = self.builder.type_sampler(); | |
let op_type_sampler_pointer = | |
self.builder | |
.type_pointer(None, StorageClass::UniformConstant, op_type_sampler); | |
let mut functions = self.builder.module_ref().functions.clone(); | |
for function in functions.iter_mut() { | |
for block in function.blocks.iter_mut() { | |
let mut instructions = Vec::new(); | |
// This needs to be done in two passes, first to find the OpAccessChains and | |
// update their type to the pointer type of the scalar value, | |
// Then to update the OpLoads. | |
for instr in block.instructions.clone() { | |
if instr.class.opcode != Op::Load { | |
instructions.push(instr); | |
continue; | |
} | |
let Some(Operand::IdRef(op_variable)) = &instr.operands.get(0) else { | |
instructions.push(instr); | |
continue; | |
}; | |
// Ensure that this is an OpLoad of the previous OpAccessChain | |
let Some(&OpAccessChain { | |
sampled_image, | |
ref index, | |
.. | |
}) = op_access_chains.get(op_variable) | |
else { | |
// Ignore what we didn't process before. | |
instructions.push(instr); | |
continue; | |
}; | |
let Some(base_type_id) = sampled_image.base_type.result_id else { | |
// Ignore what we didn't process before. | |
instructions.push(instr); | |
continue; | |
}; | |
let op_load_texture_id = self.builder.id(); | |
let op_load_sampler_id = self.builder.id(); | |
let op_access_chain_sampler_id = self.builder.id(); | |
let op_type_sampled_image = self.builder.type_sampled_image(base_type_id); | |
let mut load_instr = instr.clone(); | |
eprintln!("ADDING: %{op_load_texture_id} = OpLoad %{:?}", sampled_image.base_type.result_id); | |
eprintln!("{:?}", sampled_image.base_type); | |
load_instr.result_type = sampled_image.base_type.result_id; | |
load_instr.result_id = Some(op_load_texture_id); | |
instructions.push(load_instr); | |
// load the sampler | |
instructions.push(Instruction { | |
class: rspirv::grammar::CoreInstructionTable::get(spirv::Op::AccessChain), | |
result_type: Some(op_type_sampler_pointer), | |
result_id: Some(op_access_chain_sampler_id), | |
operands: vec![ | |
Operand::IdRef(sampled_image.sampler_variable), | |
index.clone(), | |
], | |
}); | |
instructions.push(Instruction { | |
class: rspirv::grammar::CoreInstructionTable::get(spirv::Op::Load), | |
result_type: Some(op_type_sampler), | |
result_id: Some(op_load_sampler_id), | |
operands: vec![Operand::IdRef(op_access_chain_sampler_id)], | |
}); | |
// reuse the old id for the OpSampleImage | |
instructions.push(Instruction { | |
class: rspirv::grammar::CoreInstructionTable::get(spirv::Op::SampledImage), | |
result_type: Some(op_type_sampled_image), | |
result_id: instr.result_id, | |
operands: vec![ | |
Operand::IdRef(op_load_texture_id), | |
Operand::IdRef(op_load_sampler_id), | |
], | |
}); | |
} | |
block.instructions = instructions; | |
} | |
} | |
self.builder.module_mut().functions = functions; | |
} | |
// Returns nested hashmap of { OpFunction: { OpType: CombinedImageSampler } }, | |
// indicating that the listed parameters should change | |
fn rewrite_function_calls<'b>( | |
&mut self, | |
op_access_chains: &'b FxHashMap<spirv::Word, OpAccessChain>, | |
combined_image_samplers: &'b FxHashMap<spirv::Word, CombinedImageSampler>, | |
) -> FxHashMap<spirv::Word, FxHashMap<spirv::Word, Cow<'b, CombinedImageSampler>>> { | |
let mut seen_functions: FxHashMap<spirv::Word, FxHashMap<spirv::Word, Cow<'b, CombinedImageSampler>>> = FxHashMap::default(); | |
// First pass, rewrite function calls involving sampled image access | |
for instr in self.builder.module_mut().all_inst_iter_mut() { | |
if instr.class.opcode != spirv::Op::FunctionCall { | |
continue; | |
} | |
let Some(&Operand::IdRef(function_id)) = instr.operands.get(0) else { | |
continue; | |
}; | |
if !instr.operands[1..].iter().any(|param| { | |
let &Operand::IdRef(function_id) = param else { | |
return false; | |
}; | |
combined_image_samplers.contains_key(&function_id) | |
}) { | |
continue; | |
}; | |
let mut function_call_operands = Vec::with_capacity(instr.operands.len()); | |
for operand in instr.operands.drain(..) { | |
let Operand::IdRef(op_ref_id) = operand else { | |
function_call_operands.push(operand); | |
continue; | |
}; | |
let Some(sampled_image) = combined_image_samplers.get(&op_ref_id) else { | |
function_call_operands.push(operand); | |
continue; | |
}; | |
// todo: assumes homogenous types, but need to check behaviour for array types and | |
// doing op_load after. | |
// Will need to do like rewrite_function_calls_with_global_op_access_chain | |
function_call_operands.push(operand); | |
function_call_operands.push(Operand::IdRef(sampled_image.sampler_variable)); | |
match seen_functions.entry(function_id) { | |
Entry::Occupied(mut vec) => { | |
vec.get_mut().insert(sampled_image.original_uniform_pointer_type_id, Cow::Borrowed(sampled_image)); | |
} | |
Entry::Vacant(vec) => { | |
let mut map = FxHashMap::default(); | |
// eprintln!("fn({}), {:#?}", function_id, sampled_image); | |
map.insert(sampled_image.original_uniform_pointer_type_id, Cow::Borrowed(sampled_image)); | |
vec.insert(map); | |
} | |
} | |
} | |
instr.operands = function_call_operands; | |
} | |
// Deal with OpAccessChains. | |
let op_type_sampler = self.builder.type_sampler(); | |
let op_type_sampler_pointer = | |
self.builder | |
.type_pointer(None, StorageClass::UniformConstant, op_type_sampler); | |
let mut functions = self.builder.module_ref().functions.clone(); | |
for function in functions.iter_mut() { | |
for block in function.blocks.iter_mut() { | |
let mut instructions = Vec::new(); | |
// This needs to be done in two passes, first to find the OpAccessChains and | |
// update their type to the pointer type of the scalar value, | |
// Then to update the OpLoads. | |
for mut instr in block.instructions.clone() { | |
if instr.class.opcode != Op::FunctionCall { | |
instructions.push(instr); | |
continue; | |
} | |
let Some(&Operand::IdRef(function_id)) = instr.operands.get(0) else { | |
instructions.push(instr); | |
continue; | |
}; | |
// Ensure that this is an OpFunctionCall involving some previous OpAccessChain | |
let relevant_operands = instr.operands[1..].iter().filter(|param| { | |
let &Operand::IdRef(op_ref_id) = param else { | |
return false; | |
}; | |
op_access_chains.contains_key(&op_ref_id) | |
}).collect::<Vec<_>>(); | |
if relevant_operands.is_empty() { | |
instructions.push(instr); | |
continue; | |
} | |
// Maps OpAccessChain of texture to OpAccessChain of sampler | |
let mut op_access_sampler_mapping = FxHashMap::default(); | |
// Insert necessary OpAccessChains | |
for operand in relevant_operands { | |
let &Operand::IdRef(op_ref_id) = operand else { | |
continue; | |
}; | |
let Some(&OpAccessChain { | |
sampled_image, | |
ref index, | |
.. | |
}) = op_access_chains.get(&op_ref_id) else { | |
continue; | |
}; | |
let op_access_chain_sampler_id = self.builder.id(); | |
// access chain the sampler | |
instructions.push(Instruction { | |
class: rspirv::grammar::CoreInstructionTable::get(spirv::Op::AccessChain), | |
result_type: Some(op_type_sampler_pointer), | |
result_id: Some(op_access_chain_sampler_id), | |
operands: vec![ | |
Operand::IdRef(sampled_image.sampler_variable), | |
index.clone(), | |
], | |
}); | |
op_access_sampler_mapping.insert(op_ref_id, op_access_chain_sampler_id); | |
} | |
let mut function_call_operands = Vec::with_capacity(instr.operands.len()); | |
for operand in instr.operands.drain(..) { | |
let Operand::IdRef(op_ref_id) = operand else { | |
function_call_operands.push(operand); | |
continue; | |
}; | |
let Some(&OpAccessChain { | |
sampled_image, | |
ref index, | |
original_result_type, | |
target_pointer_type, | |
result_id | |
}) = op_access_chains.get(&op_ref_id) else { | |
function_call_operands.push(operand); | |
continue; | |
}; | |
let Some(base_type_id) = sampled_image.base_type.result_id else { | |
function_call_operands.push(operand); | |
continue; | |
}; | |
let Some(&op_access_chain_sampler) = op_access_sampler_mapping | |
.get(&op_ref_id) else { | |
function_call_operands.push(operand); | |
continue; | |
}; | |
function_call_operands.push(operand); | |
function_call_operands.push(Operand::IdRef(op_access_chain_sampler)); | |
let original_type = self | |
.find_global_instruction(original_result_type) | |
.cloned() | |
.expect("huh"); | |
let sampled_image = CombinedImageSampler { | |
texture_variable: result_id, | |
sampler_variable: op_access_chain_sampler, | |
sampler_pointer_type: op_type_sampler_pointer, | |
original_uniform_pointer_type_id: original_result_type, | |
original_uniform_type: original_type, | |
target_texture_type_id: base_type_id, | |
target_texture_pointer_type_id: target_pointer_type, | |
base_type: sampled_image.base_type.clone(), | |
}; | |
match seen_functions.entry(function_id) { | |
Entry::Occupied(mut vec) => { | |
vec.get_mut().insert(original_result_type, Cow::Owned(sampled_image)); | |
} | |
Entry::Vacant(vec) => { | |
let mut map = FxHashMap::default(); | |
map.insert(original_result_type, Cow::Owned(sampled_image)); | |
vec.insert(map); | |
} | |
} | |
} | |
instr.operands = function_call_operands; | |
// Restore the instructions | |
instructions.push(instr); | |
} | |
block.instructions = instructions; | |
} | |
} | |
self.builder.module_mut().functions = functions; | |
seen_functions | |
} | |
fn rewrite_functions_definitions<'b>(&mut self, mappings: &FxHashMap<spirv::Word, FxHashMap<spirv::Word, Cow<'b, CombinedImageSampler>>>) | |
-> FxHashMap<spirv::Word, CombinedImageSampler> | |
{ | |
let mut sampled_refs = FxHashMap::default(); | |
let mut functions = self.builder.module_ref().functions.clone(); | |
for function in functions.iter_mut() { | |
let Some(def_id) = function.def_id() else { | |
continue; | |
}; | |
let Some(function_def) = function.def.clone() else { | |
continue; | |
}; | |
let Some(param_mappings) = mappings.get(&def_id) else { | |
continue; | |
}; | |
let &Some(function_return_type_id) = &function_def.result_type else { | |
continue; | |
}; | |
let Some(&Operand::IdRef(function_type_id)) = &function_def.operands.last() else { | |
continue; | |
}; | |
let mut parameters = Vec::new(); | |
let mut param_types = Vec::new(); | |
for param in function.parameters.drain(..) { | |
let Some(param_id) = param.result_id else { | |
parameters.push(param); | |
continue; | |
}; | |
let Some(param_type) = param.result_type else { | |
parameters.push(param); | |
continue; | |
}; | |
let Some(sampled_image) = param_mappings.get(¶m_type) else { | |
parameters.push(param); | |
continue; | |
}; | |
let op_function_parameter_sampler_id = self.builder.id(); | |
// todo: updaite param type to base type | |
// parameters.push(param); | |
parameters.push(Instruction { | |
class: rspirv::grammar::CoreInstructionTable::get(spirv::Op::FunctionParameter), | |
result_type: Some(sampled_image.target_texture_pointer_type_id), | |
result_id: param.result_id, | |
operands: vec![], | |
}); | |
parameters.push(Instruction { | |
class: rspirv::grammar::CoreInstructionTable::get(spirv::Op::FunctionParameter), | |
result_type: Some(sampled_image.sampler_pointer_type), | |
result_id: Some(op_function_parameter_sampler_id), | |
operands: vec![], | |
}); | |
param_types.push(sampled_image.target_texture_pointer_type_id); | |
param_types.push(sampled_image.sampler_pointer_type); | |
sampled_refs.insert(param_id, CombinedImageSampler { | |
texture_variable: param_id, | |
sampler_variable: op_function_parameter_sampler_id, | |
sampler_pointer_type: sampled_image.sampler_pointer_type, | |
original_uniform_pointer_type_id: param_type, | |
original_uniform_type: sampled_image.original_uniform_type.clone(), | |
target_texture_type_id: sampled_image.target_texture_type_id, | |
target_texture_pointer_type_id: sampled_image.target_texture_pointer_type_id, | |
base_type: sampled_image.base_type.clone() | |
}); | |
} | |
let new_type = self.builder | |
.type_function(function_return_type_id, param_types); | |
if let Some(function) = &mut function.def { | |
if let Some(function_type) = function.operands.last_mut() { | |
*function_type = Operand::IdRef(new_type); | |
} | |
} | |
function.parameters = parameters; | |
} | |
self.builder.module_mut().functions = functions; | |
sampled_refs | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
MIT Licensed