324 lines
14 KiB
Rust
324 lines
14 KiB
Rust
use pyo3::Py;
|
|
use rspirv::{
|
|
binary::Assemble,
|
|
dr::{self, Operand},
|
|
spirv,
|
|
};
|
|
use std::collections::HashMap;
|
|
|
|
use crate::pybytecode::{FunctionDef, Instruction, PossibleConsts, ShaderInput};
|
|
|
|
#[pyo3::pyfunction]
|
|
pub fn spv_gen(kernel: (u32, u32, u32), input: Py<ShaderInput>) -> Vec<u8> {
|
|
let input = input.get();
|
|
let main_name = input.main_function.as_str();
|
|
|
|
let mut func_map: HashMap<String, &FunctionDef> = HashMap::new();
|
|
for func in &input.functions {
|
|
let f = func.get();
|
|
func_map.insert(f.name.clone(), f);
|
|
}
|
|
|
|
let main_func = func_map.get(main_name).unwrap();
|
|
|
|
let mut b = rspirv::dr::Builder::new();
|
|
b.set_version(1, 3);
|
|
b.capability(spirv::Capability::Shader);
|
|
b.memory_model(spirv::AddressingModel::Logical, spirv::MemoryModel::GLSL450);
|
|
|
|
let void = b.type_void();
|
|
let voidf = b.type_function(void, vec![]);
|
|
|
|
let int32 = b.type_int(32, 1);
|
|
let float32 = b.type_float(32, None);
|
|
|
|
let mut param_vars: HashMap<String, (u32, u32, u32, bool)> = HashMap::new();
|
|
|
|
for (i, param) in main_func.inputs.iter().enumerate() {
|
|
let p = param.get();
|
|
let is_array = p.ty.starts_with("list[");
|
|
let base_type = match p.ty.as_str() {
|
|
"float" | "list[float]" => float32,
|
|
"int" | "list[int]" => int32,
|
|
_ => int32,
|
|
};
|
|
let member_type = if is_array {
|
|
let runtime_array = b.type_runtime_array(base_type);
|
|
b.decorate(
|
|
runtime_array,
|
|
spirv::Decoration::ArrayStride,
|
|
vec![Operand::LiteralBit32(4)],
|
|
);
|
|
runtime_array
|
|
} else {
|
|
base_type
|
|
};
|
|
let struct_type = b.type_struct(vec![member_type]);
|
|
b.decorate(struct_type, spirv::Decoration::Block, vec![]);
|
|
b.member_decorate(
|
|
struct_type,
|
|
0,
|
|
spirv::Decoration::Offset,
|
|
vec![Operand::LiteralBit32(0)],
|
|
);
|
|
let ptr_storage = b.type_pointer(None, spirv::StorageClass::StorageBuffer, struct_type);
|
|
let input_buf = b.variable(ptr_storage, None, spirv::StorageClass::StorageBuffer, None);
|
|
b.decorate(
|
|
input_buf,
|
|
spirv::Decoration::DescriptorSet,
|
|
vec![Operand::LiteralBit32(0)],
|
|
);
|
|
b.decorate(
|
|
input_buf,
|
|
spirv::Decoration::Binding,
|
|
vec![Operand::LiteralBit32(i as u32)],
|
|
);
|
|
param_vars.insert(p.name.clone(), (base_type, ptr_storage, input_buf, is_array));
|
|
}
|
|
|
|
let mainf = b
|
|
.begin_function(
|
|
void,
|
|
None,
|
|
spirv::FunctionControl::DONT_INLINE | spirv::FunctionControl::CONST,
|
|
voidf,
|
|
)
|
|
.unwrap();
|
|
b.begin_block(None).unwrap();
|
|
|
|
generate_body(
|
|
&mut b,
|
|
&main_func.instructions.iter().map(|i| i.get()).collect::<Vec<&Instruction>>(),
|
|
¶m_vars,
|
|
);
|
|
|
|
b.ret().unwrap();
|
|
b.end_function().unwrap();
|
|
b.entry_point(
|
|
spirv::ExecutionModel::GLCompute,
|
|
mainf,
|
|
"main",
|
|
vec![],
|
|
);
|
|
b.execution_mode(
|
|
mainf,
|
|
spirv::ExecutionMode::LocalSize,
|
|
vec![kernel.0, kernel.1, kernel.2],
|
|
);
|
|
b.module()
|
|
.assemble()
|
|
.into_iter()
|
|
.flat_map(|item| item.to_le_bytes())
|
|
.collect()
|
|
}
|
|
|
|
fn make_const(b: &mut rspirv::dr::Builder, type_id: u32, value: u32) -> u32 {
|
|
let id = b.id();
|
|
let inst = dr::Instruction::new(
|
|
spirv::Op::Constant,
|
|
Some(type_id),
|
|
Some(id),
|
|
vec![Operand::LiteralBit32(value)],
|
|
);
|
|
b.module_mut().types_global_values.push(inst);
|
|
id
|
|
}
|
|
|
|
fn generate_body(
|
|
b: &mut rspirv::dr::Builder,
|
|
instrs: &[&Instruction],
|
|
param_vars: &HashMap<String, (u32, u32, u32, bool)>,
|
|
) {
|
|
let int32 = b.type_int(32, 1);
|
|
let float32 = b.type_float(32, None);
|
|
|
|
let mut stack: Vec<(u32, u32, Option<(u32, u32)>)> = Vec::new();
|
|
let mut locals: HashMap<String, (u32, u32)> = HashMap::new();
|
|
let mut array_refs: HashMap<String, (u32, u32)> = HashMap::new();
|
|
|
|
for instr in instrs {
|
|
match instr {
|
|
Instruction::LoadConst { val } => {
|
|
let val = val.get();
|
|
match val {
|
|
PossibleConsts::Int(v) => {
|
|
let cid = make_const(b, int32, *v as u32);
|
|
stack.push((cid, int32, None));
|
|
}
|
|
PossibleConsts::Float(v) => {
|
|
let cid = make_const(b, float32, v.to_bits());
|
|
stack.push((cid, float32, None));
|
|
}
|
|
PossibleConsts::Bool(v) => {
|
|
let cid = make_const(b, int32, if *v { 1 } else { 0 });
|
|
stack.push((cid, int32, None));
|
|
}
|
|
PossibleConsts::Null() => {}
|
|
PossibleConsts::Str(_) => {}
|
|
}
|
|
}
|
|
Instruction::LoadFast { name } => {
|
|
if let Some(&(vid, vtype)) = locals.get(name) {
|
|
let aref = array_refs.get(name).copied();
|
|
stack.push((vid, vtype, aref));
|
|
} else if let Some(&(base_type, _ptr_storage, input_buf, is_array)) = param_vars.get(name) {
|
|
let zero_const = make_const(b, int32, 0);
|
|
let ptr_type = b.type_pointer(None, spirv::StorageClass::StorageBuffer, base_type);
|
|
let ptr = b.access_chain(ptr_type, None, input_buf, vec![zero_const]).unwrap();
|
|
let loaded = b.load(base_type, None, ptr, None, vec![]).unwrap();
|
|
locals.insert(name.clone(), (loaded, base_type));
|
|
if is_array {
|
|
array_refs.insert(name.clone(), (input_buf, base_type));
|
|
}
|
|
let aref = if is_array { Some((input_buf, base_type)) } else { None };
|
|
stack.push((loaded, base_type, aref));
|
|
}
|
|
}
|
|
Instruction::LoadFastLoadFast { names } => {
|
|
for name in [&names.0, &names.1] {
|
|
if let Some(&(vid, vtype)) = locals.get(name) {
|
|
let aref = array_refs.get(name).copied();
|
|
stack.push((vid, vtype, aref));
|
|
} else if let Some(&(base_type, _ptr_storage, input_buf, is_array)) = param_vars.get(name) {
|
|
let zero_const = make_const(b, int32, 0);
|
|
let ptr_type = b.type_pointer(None, spirv::StorageClass::StorageBuffer, base_type);
|
|
let ptr = b.access_chain(ptr_type, None, input_buf, vec![zero_const]).unwrap();
|
|
let loaded = b.load(base_type, None, ptr, None, vec![]).unwrap();
|
|
locals.insert(name.clone(), (loaded, base_type));
|
|
if is_array {
|
|
array_refs.insert(name.clone(), (input_buf, base_type));
|
|
}
|
|
let aref = if is_array { Some((input_buf, base_type)) } else { None };
|
|
stack.push((loaded, base_type, aref));
|
|
}
|
|
}
|
|
}
|
|
Instruction::StoreFast { name } => {
|
|
if let Some((vid, vtype, _aref)) = stack.pop() {
|
|
locals.insert(name.clone(), (vid, vtype));
|
|
if let Some(&(base_type, _ptr_storage, input_buf, _is_array)) = param_vars.get(name) {
|
|
let zero_const = make_const(b, int32, 0);
|
|
let ptr_type = b.type_pointer(None, spirv::StorageClass::StorageBuffer, base_type);
|
|
let ptr = b.access_chain(ptr_type, None, input_buf, vec![zero_const]).unwrap();
|
|
b.store(ptr, vid, None, vec![]).unwrap();
|
|
}
|
|
}
|
|
}
|
|
Instruction::BinOp { op } => {
|
|
if let (Some((rhs, rhs_type, _)), Some((lhs, lhs_type, _))) = (stack.pop(), stack.pop()) {
|
|
let (res_type, res) = if rhs_type == float32 || lhs_type == float32 {
|
|
match op.as_str() {
|
|
"+" | "+=" => (float32, b.f_add(float32, None, lhs, rhs).unwrap()),
|
|
"-" | "-=" => (float32, b.f_sub(float32, None, lhs, rhs).unwrap()),
|
|
"*" | "*=" => (float32, b.f_mul(float32, None, lhs, rhs).unwrap()),
|
|
"/" | "//" | "/=" => (float32, b.f_div(float32, None, lhs, rhs).unwrap()),
|
|
"%" | "%=" => (float32, b.f_mod(float32, None, lhs, rhs).unwrap()),
|
|
_ => (float32, lhs),
|
|
}
|
|
} else {
|
|
match op.as_str() {
|
|
"+" | "+=" => (int32, b.i_add(int32, None, lhs, rhs).unwrap()),
|
|
"-" | "-=" => (int32, b.i_sub(int32, None, lhs, rhs).unwrap()),
|
|
"*" | "*=" => (int32, b.i_mul(int32, None, lhs, rhs).unwrap()),
|
|
"/" | "//" | "/=" => (int32, b.s_div(int32, None, lhs, rhs).unwrap()),
|
|
"%" | "%=" => (int32, b.s_rem(int32, None, lhs, rhs).unwrap()),
|
|
"**" => (int32, lhs),
|
|
_ => (int32, lhs),
|
|
}
|
|
};
|
|
stack.push((res, res_type, None));
|
|
}
|
|
}
|
|
Instruction::UnaryOp { op } => {
|
|
if let Some((val, vtype, _)) = stack.pop() {
|
|
let (res_type, res) = if vtype == float32 {
|
|
match op.as_str() {
|
|
"-" => (float32, b.f_negate(float32, None, val).unwrap()),
|
|
"!" => {
|
|
let bool_type = b.type_bool();
|
|
let zero = make_const(b, float32, 0.0f32.to_bits());
|
|
(bool_type, b.f_ord_equal(bool_type, None, val, zero).unwrap())
|
|
}
|
|
_ => (vtype, val),
|
|
}
|
|
} else {
|
|
match op.as_str() {
|
|
"-" => (int32, b.s_negate(int32, None, val).unwrap()),
|
|
"!" => {
|
|
let bool_type = b.type_bool();
|
|
let zero = make_const(b, int32, 0);
|
|
(bool_type, b.i_equal(bool_type, None, val, zero).unwrap())
|
|
}
|
|
_ => (vtype, val),
|
|
}
|
|
};
|
|
stack.push((res, res_type, None));
|
|
}
|
|
}
|
|
Instruction::CompareOp { op } => {
|
|
if let (Some((rhs, rhs_type, _)), Some((lhs, lhs_type, _))) = (stack.pop(), stack.pop()) {
|
|
let bool_type = b.type_bool();
|
|
let (res, res_type) = if rhs_type == float32 || lhs_type == float32 {
|
|
match op.as_str() {
|
|
"== " | "==" => (b.f_ord_equal(bool_type, None, lhs, rhs).unwrap(), bool_type),
|
|
"!=" => (b.f_ord_not_equal(bool_type, None, lhs, rhs).unwrap(), bool_type),
|
|
"<" => (b.f_ord_less_than(bool_type, None, lhs, rhs).unwrap(), bool_type),
|
|
"<=" => (b.f_ord_less_than_equal(bool_type, None, lhs, rhs).unwrap(), bool_type),
|
|
">" => (b.f_ord_greater_than(bool_type, None, lhs, rhs).unwrap(), bool_type),
|
|
">=" => (b.f_ord_greater_than_equal(bool_type, None, lhs, rhs).unwrap(), bool_type),
|
|
_ => (b.f_ord_equal(bool_type, None, lhs, rhs).unwrap(), bool_type),
|
|
}
|
|
} else {
|
|
match op.as_str() {
|
|
"== " | "==" => (b.i_equal(bool_type, None, lhs, rhs).unwrap(), bool_type),
|
|
"!=" => (b.i_not_equal(bool_type, None, lhs, rhs).unwrap(), bool_type),
|
|
"<" => (b.s_less_than(bool_type, None, lhs, rhs).unwrap(), bool_type),
|
|
"<=" => (b.s_less_than_equal(bool_type, None, lhs, rhs).unwrap(), bool_type),
|
|
">" => (b.s_greater_than(bool_type, None, lhs, rhs).unwrap(), bool_type),
|
|
">=" => (b.s_greater_than_equal(bool_type, None, lhs, rhs).unwrap(), bool_type),
|
|
_ => (b.i_equal(bool_type, None, lhs, rhs).unwrap(), bool_type),
|
|
}
|
|
};
|
|
stack.push((res, res_type, None));
|
|
}
|
|
}
|
|
Instruction::PopTop {} => {
|
|
stack.pop();
|
|
}
|
|
Instruction::ReturnValue {} => {
|
|
stack.pop();
|
|
}
|
|
Instruction::LoadGlobal { .. } => {}
|
|
Instruction::BinarySubscr {} => {
|
|
if let (Some((idx, _, _)), Some((_, base_type, Some((input_buf, _))))) = (stack.pop(), stack.pop()) {
|
|
let zero_const = make_const(b, int32, 0);
|
|
let ptr_type = b.type_pointer(None, spirv::StorageClass::StorageBuffer, base_type);
|
|
let ptr = b.access_chain(ptr_type, None, input_buf, vec![zero_const, idx]).unwrap();
|
|
let loaded = b.load(base_type, None, ptr, None, vec![]).unwrap();
|
|
stack.push((loaded, base_type, None));
|
|
}
|
|
}
|
|
Instruction::StoreSubscr {} => {
|
|
if let (Some((idx, _, _)), Some((_, _, Some((input_buf, base_type)))), Some((val, _vtype, _))) = (stack.pop(), stack.pop(), stack.pop()) {
|
|
let zero_const = make_const(b, int32, 0);
|
|
let ptr_type = b.type_pointer(None, spirv::StorageClass::StorageBuffer, base_type);
|
|
let ptr = b.access_chain(ptr_type, None, input_buf, vec![zero_const, idx]).unwrap();
|
|
b.store(ptr, val, None, vec![]).unwrap();
|
|
}
|
|
}
|
|
Instruction::Copy { n } => {
|
|
let n = *n as usize;
|
|
let len = stack.len();
|
|
let idx = len - n;
|
|
stack.push(stack[idx].clone());
|
|
}
|
|
Instruction::Swap { n } => {
|
|
let n = *n as usize;
|
|
let len = stack.len();
|
|
stack.swap(len - 1, len - n);
|
|
}
|
|
Instruction::Call {} => {}
|
|
}
|
|
}
|
|
}
|