ruffle/render/naga-agal/src/builder.rs

759 lines
27 KiB
Rust

use std::io::Read;
use naga::{
ArraySize, BuiltIn, Constant, ConstantInner, EntryPoint, FunctionArgument, FunctionResult,
GlobalVariable, Interpolation, ScalarValue, ShaderStage, StructMember, SwizzleComponent,
};
use naga::{BinaryOperator, MathFunction};
use naga::{
Binding, Expression, Function, Handle, LocalVariable, Module, ScalarKind, Span, Statement,
Type, TypeInner, VectorSize,
};
use num_traits::FromPrimitive;
use crate::{
types::*, Error, ShaderType, VertexAttributeFormat, ENTRY_POINT, MAX_VERTEX_ATTRIBUTES,
};
const VERTEX_PROGRAM_CONTANTS: u64 = 128;
const FRAGMENT_PROGRAM_CONSTANTS: u64 = 28;
pub type Result<T> = std::result::Result<T, Error>;
pub(crate) struct NagaBuilder<'a> {
module: Module,
func: Function,
// This evaluate to a Pointer to the temporary 'main' destiation location
// (the output position for a vertex shader, or the output color for a fragment shader)
// which can be used with Expression::Load and Expression::Store
// This is needed because an AGAL shader can write to the output register
// multiple times.
dest: Handle<Expression>,
shader_config: ShaderConfig<'a>,
// Whenever we read from a varying register in a fragment shader,
// we create a new argument binding for it.
argument_expressions: Vec<Option<Handle<Expression>>>,
varying_pointers: Vec<Option<Handle<Expression>>>,
// An `Expression::GlobalVariables` for the uniform buffer
// that stores all of the program constants.
constant_registers: Handle<Expression>,
// The function return type being built up. Each time a vertex
// shader writes to a varying register, we add a new member to this
return_type: Type,
// The Naga representation of 'vec4f'
vec4f: Handle<Type>,
// The Naga representation of 'mat4x4f'
matrix4x4f: Handle<Type>,
}
impl VertexAttributeFormat {
fn to_naga_type(self, module: &mut Module) -> Handle<Type> {
if let VertexAttributeFormat::Float1 = self {
return module.types.insert(
Type {
name: None,
inner: TypeInner::Scalar {
kind: ScalarKind::Float,
width: 4,
},
},
Span::UNDEFINED,
);
}
let (size, width, kind) = match self {
VertexAttributeFormat::Float1 => unreachable!(),
VertexAttributeFormat::Float2 => (VectorSize::Bi, 4, ScalarKind::Float),
VertexAttributeFormat::Float3 => (VectorSize::Tri, 4, ScalarKind::Float),
VertexAttributeFormat::Float4 => (VectorSize::Quad, 4, ScalarKind::Float),
VertexAttributeFormat::Bytes4 => (VectorSize::Quad, 1, ScalarKind::Uint),
};
module.types.insert(
Type {
name: None,
inner: TypeInner::Vector { size, kind, width },
},
Span::UNDEFINED,
)
}
fn extend_to_float4(
&self,
base_expr: Handle<Expression>,
builder: &mut NagaBuilder,
) -> Result<Handle<Expression>> {
Ok(match self {
// This does 'vec4f(my_vec3, 1.0)'
VertexAttributeFormat::Float3 => {
let expr = base_expr;
let mut components = vec![];
for i in 0..3 {
components.push(builder.evaluate_expr(Expression::AccessIndex {
base: expr,
index: i,
}));
}
components.push(builder.func.expressions.append(
Expression::Constant(builder.module.constants.append(
Constant {
name: None,
specialization: None,
inner: ConstantInner::Scalar {
width: 4,
value: ScalarValue::Float(1.0),
},
},
Span::UNDEFINED,
)),
Span::UNDEFINED,
));
builder.evaluate_expr(Expression::Compose {
ty: builder.vec4f,
components,
})
}
VertexAttributeFormat::Float4 => base_expr,
_ => {
return Err(Error::Unimplemented(format!(
"Unsupported conversion from {self:?} to float4",
)))
}
})
}
}
#[derive(Debug)]
pub struct ShaderConfig<'a> {
pub shader_type: ShaderType,
pub vertex_attributes: &'a [Option<VertexAttributeFormat>; 8],
}
impl<'a> NagaBuilder<'a> {
pub fn process_agal(
mut agal: &[u8],
vertex_attributes: &[Option<VertexAttributeFormat>; MAX_VERTEX_ATTRIBUTES],
) -> Result<Module> {
let data = &mut agal;
let mut header = [0; 7];
data.read_exact(&mut header)?;
if header[0..6] != [0xA0, 0x01, 0x00, 0x00, 0x00, 0xA1] {
return Err(Error::InvalidHeader);
}
let shader_type = match header[6] {
0x00 => ShaderType::Vertex,
0x01 => ShaderType::Fragment,
_ => return Err(Error::InvalidShaderType(header[6])),
};
let mut builder = NagaBuilder::new(ShaderConfig {
shader_type,
vertex_attributes,
});
while !data.is_empty() {
let mut token = [0; 24];
data.read_exact(&mut token)?;
let raw_opcode = u32::from_le_bytes(token[0..4].try_into().unwrap());
// FIXME - this is a clippy false-positive
#[allow(clippy::or_fun_call)]
let opcode = Opcode::from_u32(raw_opcode).ok_or(Error::InvalidOpcode(raw_opcode))?;
let dest = DestField::parse(u32::from_le_bytes(token[4..8].try_into().unwrap()))?;
let source1 = SourceField::parse(u64::from_le_bytes(token[8..16].try_into().unwrap()))?;
let source2 = if let Opcode::Tex = opcode {
Source2::Sampler(SamplerField::parse(u64::from_le_bytes(
token[16..24].try_into().unwrap(),
))?)
} else {
Source2::SourceField(SourceField::parse(u64::from_le_bytes(
token[16..24].try_into().unwrap(),
))?)
};
builder.process_opcode(&opcode, &dest, &source1, &source2)?;
}
builder.finish()
}
fn new(shader_config: ShaderConfig<'a>) -> Self {
let mut module = Module::default();
let mut func = Function::default();
let vec4f = VertexAttributeFormat::Float4.to_naga_type(&mut module);
let matrix4x4f = module.types.insert(
Type {
name: None,
inner: TypeInner::Matrix {
columns: VectorSize::Quad,
rows: VectorSize::Quad,
width: 4,
},
},
Span::UNDEFINED,
);
// The return type always has at least one component - the vec4f that's the 'main'
// output of our shader (the position for the vertex shader, and the color for the fragment shader)
let return_type = match shader_config.shader_type {
ShaderType::Vertex => Type {
name: None,
inner: TypeInner::Struct {
members: vec![StructMember {
name: None,
ty: vec4f,
binding: Some(Binding::BuiltIn(BuiltIn::Position { invariant: false })),
offset: 0,
}],
span: 16,
},
},
ShaderType::Fragment => Type {
name: None,
inner: TypeInner::Struct {
members: vec![StructMember {
name: None,
ty: vec4f,
binding: Some(Binding::Location {
location: 0,
interpolation: None,
sampling: None,
}),
offset: 0,
}],
span: 16,
},
},
};
match shader_config.shader_type {
ShaderType::Vertex => {
func.result = Some(FunctionResult {
ty: vec4f,
binding: Some(Binding::BuiltIn(BuiltIn::Position { invariant: false })),
});
}
ShaderType::Fragment => {
func.result = Some(FunctionResult {
ty: vec4f,
binding: Some(Binding::Location {
location: 0,
interpolation: None,
sampling: None,
}),
});
}
}
// Holds the value we're going to return.
// This corresponds to RegisterType::Output
let output_temp_handle = func.local_variables.append(
LocalVariable {
name: Some("dest_temp".to_string()),
ty: vec4f,
init: None,
},
Span::UNDEFINED,
);
let dest = func.expressions.append(
Expression::LocalVariable(output_temp_handle),
Span::UNDEFINED,
);
let num_const_registers = module.constants.append(
Constant {
name: None,
specialization: None,
inner: ConstantInner::Scalar {
width: 4,
value: ScalarValue::Uint(match shader_config.shader_type {
ShaderType::Vertex => VERTEX_PROGRAM_CONTANTS,
ShaderType::Fragment => FRAGMENT_PROGRAM_CONSTANTS,
}),
},
},
Span::UNDEFINED,
);
let binding_num = match shader_config.shader_type {
ShaderType::Vertex => 0,
ShaderType::Fragment => 1,
};
let constant_registers_global = module.global_variables.append(
GlobalVariable {
name: Some("constant_registers".to_string()),
space: naga::AddressSpace::Uniform,
binding: Some(naga::ResourceBinding {
group: 0,
binding: binding_num,
}),
ty: module.types.insert(
Type {
name: None,
inner: TypeInner::Array {
base: vec4f,
size: ArraySize::Constant(num_const_registers),
stride: std::mem::size_of::<f32>() as u32 * 4,
},
},
Span::UNDEFINED,
),
init: None,
},
Span::UNDEFINED,
);
let constant_registers = func.expressions.append(
Expression::GlobalVariable(constant_registers_global),
Span::UNDEFINED,
);
NagaBuilder {
module,
func,
dest,
shader_config,
argument_expressions: vec![],
varying_pointers: vec![],
return_type,
matrix4x4f,
vec4f,
constant_registers,
}
}
fn get_vertex_input(&mut self, index: usize) -> Result<Handle<Expression>> {
if index >= self.argument_expressions.len() {
self.argument_expressions.resize(index + 1, None);
// FIXME - this is a clippy false-positive
#[allow(clippy::or_fun_call)]
let ty = self.shader_config.vertex_attributes[index]
.as_ref()
.ok_or(Error::MissingVertexAttributeData(index))?
.to_naga_type(&mut self.module);
self.func.arguments.push(FunctionArgument {
name: None,
ty,
binding: Some(Binding::Location {
location: index as u32,
interpolation: None,
sampling: None,
}),
});
// Arguments map one-to-one to vertex attributes.
let expr = self
.func
.expressions
.append(Expression::FunctionArgument(index as u32), Span::UNDEFINED);
self.argument_expressions[index] = Some(expr);
}
Ok(self.argument_expressions[index].unwrap())
}
fn get_varying_pointer(&mut self, index: usize) -> Result<Handle<Expression>> {
if index >= self.varying_pointers.len() {
self.varying_pointers.resize(index + 1, None);
}
if self.varying_pointers[index].is_none() {
match self.shader_config.shader_type {
ShaderType::Vertex => {
// We can write to varying variables in the vertex shader,
// and the fragment shader will receive them is input.
// Therefore, we create a local variable for each varying,
// and return them at the end of the function.
let local = self.func.local_variables.append(
LocalVariable {
name: Some(format!("varying_{index}")),
ty: self.vec4f,
init: None,
},
Span::UNDEFINED,
);
let expr = self
.func
.expressions
.append(Expression::LocalVariable(local), Span::UNDEFINED);
let _range = self
.func
.expressions
.range_from(self.func.expressions.len() - 1);
if let TypeInner::Struct { members, .. } = &mut self.return_type.inner {
members.push(StructMember {
name: Some(format!("varying_{index}")),
ty: self.vec4f,
binding: Some(Binding::Location {
location: index as u32,
interpolation: Some(naga::Interpolation::Perspective),
sampling: None,
}),
offset: 0,
});
} else {
unreachable!();
}
self.varying_pointers[index] = Some(expr);
}
ShaderType::Fragment => {
self.func.arguments.push(FunctionArgument {
name: None,
ty: self.vec4f,
binding: Some(Binding::Location {
location: index as u32,
interpolation: Some(Interpolation::Perspective),
sampling: None,
}),
});
let expr = self
.func
.expressions
.append(Expression::FunctionArgument(index as u32), Span::UNDEFINED);
self.varying_pointers[index] = Some(expr);
}
};
};
Ok(self.varying_pointers[index].unwrap())
}
fn emit_const_register_load(&mut self, index: usize) -> Result<Handle<Expression>> {
let index_const = self.module.constants.append(
Constant {
name: None,
specialization: None,
inner: ConstantInner::Scalar {
width: 4,
value: ScalarValue::Uint(index as u64),
},
},
Span::UNDEFINED,
);
let index_expr = self
.func
.expressions
.append(Expression::Constant(index_const), Span::UNDEFINED);
let register_pointer = self.func.expressions.append(
Expression::Access {
base: self.constant_registers,
index: index_expr,
},
Span::UNDEFINED,
);
Ok(self.evaluate_expr(Expression::Load {
pointer: register_pointer,
}))
}
fn emit_varying_load(&mut self, index: usize) -> Result<Handle<Expression>> {
// A LocalVariable evaluates to a pointer, so we need to load it
let varying_expr = self.get_varying_pointer(index)?;
Ok(match self.shader_config.shader_type {
ShaderType::Vertex => self.evaluate_expr(Expression::Load {
pointer: varying_expr,
}),
ShaderType::Fragment => varying_expr,
})
}
fn emit_source_field_load(
&mut self,
source: &SourceField,
extend_to_vec4: bool,
) -> Result<Handle<Expression>> {
let (mut base_expr, source_type) = match source.register_type {
// We can use a function argument directly - we don't need
// a separate Expression::Load
RegisterType::Attribute => (
self.get_vertex_input(source.reg_num as usize)?,
// FIXME - this is a clippy false-positive
#[allow(clippy::or_fun_call)]
self.shader_config.vertex_attributes[source.reg_num as usize]
.ok_or(Error::MissingVertexAttributeData(source.reg_num as usize))?,
),
RegisterType::Varying => (
self.emit_varying_load(source.reg_num as usize)?,
VertexAttributeFormat::Float4,
),
RegisterType::Constant => (
self.emit_const_register_load(source.reg_num as usize)?,
// Constants are always a vec4<f32>
VertexAttributeFormat::Float4,
),
_ => {
return Err(Error::Unimplemented(format!(
"Unimplemented source reg type {:?}",
source.register_type
)))
}
};
if matches!(source.direct_mode, DirectMode::Indirect) {
return Err(Error::Unimplemented(
"Indirect addressing not implemented".to_string(),
));
}
if extend_to_vec4 && source_type != VertexAttributeFormat::Float4 {
base_expr = source_type.extend_to_float4(base_expr, self)?;
}
// Swizzle is 'xyzw', which is a no-op. Just return the base expression.
if source.swizzle == 0b11100100 {
return Ok(base_expr);
}
let swizzle_flags = [
source.swizzle & 0b11,
(source.swizzle >> 2) & 0b11,
(source.swizzle >> 4) & 0b11,
(source.swizzle >> 6) & 0b11,
];
let swizzle_components: [SwizzleComponent; 4] = swizzle_flags
.into_iter()
.map(|flag| match flag {
0b00 => SwizzleComponent::X,
0b01 => SwizzleComponent::Y,
0b10 => SwizzleComponent::Z,
0b11 => SwizzleComponent::W,
_ => unreachable!(),
})
.collect::<Vec<_>>()
.try_into()
.unwrap();
Ok(self.func.expressions.append(
Expression::Swizzle {
size: VectorSize::Quad,
vector: base_expr,
pattern: swizzle_components,
},
Span::UNDEFINED,
))
}
fn emit_dest_store(&mut self, dest: &DestField, expr: Handle<Expression>) -> Result<()> {
let base_expr = match dest.register_type {
RegisterType::Output => self.dest,
RegisterType::Varying => self.get_varying_pointer(dest.reg_num as usize)?,
_ => {
return Err(Error::Unimplemented(format!(
"Unimplemented dest reg type: {dest:?}",
)))
}
};
// Optimization - use a Store instead of writing individual fields
// when we're writing to the entire output register.
if dest.write_mask.is_all() {
let store = Statement::Store {
pointer: base_expr,
value: expr,
};
self.func.body.push(store, Span::UNDEFINED);
} else {
for (i, mask) in [(0, Mask::X), (1, Mask::Y), (2, Mask::Z), (3, Mask::W)] {
if dest.write_mask.contains(mask) {
self.func.body.push(
Statement::Store {
pointer: self.func.expressions.append(
Expression::AccessIndex {
base: base_expr,
index: i,
},
Span::UNDEFINED,
),
value: expr,
},
Span::UNDEFINED,
);
}
}
}
Ok(())
}
/// Creates a `Statement::Emit` covering `expr`
fn evaluate_expr(&mut self, expr: Expression) -> Handle<Expression> {
let prev_len = self.func.expressions.len();
let expr = self.func.expressions.append(expr, Span::UNDEFINED);
let range = self.func.expressions.range_from(prev_len);
self.func.body.push(Statement::Emit(range), Span::UNDEFINED);
expr
}
fn process_opcode(
&mut self,
opcode: &Opcode,
dest: &DestField,
source1: &SourceField,
source2: &Source2,
) -> Result<()> {
match opcode {
// Copy the source register to the destination register
Opcode::Mov => {
// On the ActionScript side, the user might have specified something *other* than
// vec4f. In that case, we need to extend the source to a vec4f if we're writing to
// a vec4f register.
// FIXME - do we need to do this extension in other cases?
let do_extend = matches!(
dest.register_type,
RegisterType::Output | RegisterType::Varying
);
let source = self.emit_source_field_load(source1, do_extend)?;
self.emit_dest_store(dest, source)?;
}
// Perform 'M * v', where M is a 4x4 matrix, and 'v' is a column vector.
Opcode::M44 => {
let source2 = match source2 {
Source2::SourceField(source2) => source2,
_ => unreachable!(),
};
// Read each row of the matrix
let source2_row0 = self.emit_source_field_load(source2, false)?;
let source2_row1 = self.emit_source_field_load(
&SourceField {
reg_num: source2.reg_num + 1,
..source2.clone()
},
false,
)?;
let source2_row2 = self.emit_source_field_load(
&SourceField {
reg_num: source2.reg_num + 2,
..source2.clone()
},
false,
)?;
let source2_row3 = self.emit_source_field_load(
&SourceField {
reg_num: source2.reg_num + 3,
..source2.clone()
},
false,
)?;
// FIXME - The naga spv backend hits an 'unreachable!'
// if we don't create a Statement::Emit for each of these,
// even though validation passes. We should investigate this
// and report it upstream.
let matrix = self.evaluate_expr(Expression::Compose {
ty: self.matrix4x4f,
components: vec![source2_row0, source2_row1, source2_row2, source2_row3],
});
// Naga interprets each component of the matrix as a *column*.
// However, the matrix is stored in memory as a *row*, so we need
// to transpose it.
let matrix = self.evaluate_expr(Expression::Math {
fun: MathFunction::Transpose,
arg: matrix,
arg1: None,
arg2: None,
arg3: None,
});
let vector = self.emit_source_field_load(source1, true)?;
let multiply = self.evaluate_expr(Expression::Binary {
op: BinaryOperator::Multiply,
left: matrix,
right: vector,
});
self.emit_dest_store(dest, multiply)?;
}
_ => {
return Err(Error::Unimplemented(format!(
"Unimplemented opcode: {opcode:?}",
)))
}
}
Ok(())
}
fn finish(mut self) -> Result<Module> {
// Load the 'main' output (a position or color) from our temporary location.
let dest_load = self.evaluate_expr(Expression::Load { pointer: self.dest });
let mut components = vec![dest_load];
// If the vertex shader wrote to any varying registers, we need to
// return them as well.
if let ShaderType::Vertex = self.shader_config.shader_type {
for i in 0..self.varying_pointers.len() {
if self.varying_pointers[i].is_some() {
components.push(self.emit_varying_load(i)?);
}
}
}
// We're consuming 'self', so just store store garbage here so that we can continue
// to use methods on 'self'
let return_ty = std::mem::replace(
&mut self.return_type,
Type {
name: None,
inner: TypeInner::Scalar {
kind: ScalarKind::Float,
width: 0,
},
},
);
// Finalize the return type, and do emit the actual return
let return_ty = self.module.types.insert(return_ty, Span::UNDEFINED);
self.func.result = Some(FunctionResult {
ty: return_ty,
binding: None,
});
let return_expr = self.evaluate_expr(Expression::Compose {
ty: return_ty,
components,
});
self.func.body.push(
Statement::Return {
value: Some(return_expr),
},
Span::UNDEFINED,
);
let entry_point = EntryPoint {
name: ENTRY_POINT.to_string(),
stage: match self.shader_config.shader_type {
ShaderType::Vertex => ShaderStage::Vertex,
ShaderType::Fragment => ShaderStage::Fragment,
},
early_depth_test: None,
workgroup_size: [0; 3],
function: self.func,
};
self.module.entry_points.push(entry_point);
Ok(self.module)
}
}