实用的改进
作者：Miracle
- 2025-09-30 四川
本文字数：2629 字
阅读完需：约 9 分钟
硬编码在代码里面图像的大小实在 太 low 我们需要更多的适应性以满足现实的需求，我们需要一种灵活的方式传递给 shader 参数，特别注意的是，shader 都是处理大容量的 计算密集型 的数据，所以暂时不考虑动态的容器，我们引入 Vec<T> 的 概念表示 类型为 T 的 缓冲区
这样我们的 zeta 程序 变成这样
fn yuv2rgb(y, u, v) { let r = (y + 1.402f32 * (v - 128.0)).clamp(0.0, 255.0); let g = (y - 0.344136f32 * (u - 128.0) - 0.714136f32 * (v - 128.0)).clamp(0.0, 255.0); let b = (y + 1.772f32 * (u - 128.0)).clamp(0.0, 255.0); let rgb = (r as u32) | (g as u32) << 8 | (b as u32) << 16 | 0xFF000000u32; return rgb;}
pub fn main(y_plane: Vec<u8>, u_plane: Vec<u8>, v_plane: Vec<u8>, mut rgb: Vec<u32>) { let g_id = spirv_group(); let id = spirv_local(); let x = g_id[0] * 32 + id[0]; // x 方向的坐标 分成 32 x 32 的小块 let y = g_id[1] * 32 + id[1]; // y 方向的坐标 分成 32 x 32 的小块 let offset = y * 2560 + x; let rgb_val = yuv2rgb(y_plane[offset], u_plane[offset], v_plane[offset]); rgb[offset] = rgb_val;}复制代码
使用 y u v 三个独立的平面传递数据，使用 rgb 作为输出的缓冲区
我们使用下面的代码 将 程序的参数转换为 DescriptSet
pub fn add_arg(&mut self, is_mut: bool, ty: Type, binding: u32) { let ty_id = self.get_type(SpirvType::Var(ty.clone())); let ty_pt = if is_mut { self.builder.type_pointer(None, StorageClass::StorageBuffer, ty_id) } else { self.builder.type_pointer(None, StorageClass::StorageBuffer, ty_id) }; let id = self.builder.variable(ty_pt, None, StorageClass::StorageBuffer, None); self.builder.decorate(id, spirv::Decoration::DescriptorSet, vec![rspirv::dr::Operand::LiteralBit32(0)]); //描述集 0 以后可以设置跟 loader 一致 self.builder.decorate(id, spirv::Decoration::Binding, vec![rspirv::dr::Operand::LiteralBit32(binding)]); //绑定 self.vars.push(Var{ id, ty: SpirvType::Pointer(ty, StorageClass::StorageBuffer)}); }复制代码
每一个参数有独立的绑定点
为了方便调用，我们封装了 Arg 结构 可以方便的创建 参数的数据
impl Args { pub fn new(device: Arc<Device>)-> Self { let allocator = std::sync::Arc::new(StandardMemoryAllocator::new_default(device.clone())); Self{allocator, set_writers: Vec::new()} } pub fn add_input<T: BufferContents>(&mut self, value: T) { let idx = self.set_writers.len() as u32; self.set_writers.push(WriteDescriptorSet::buffer(idx, Buffer::from_data(self.allocator.clone(), BufferCreateInfo {usage: BufferUsage::STORAGE_BUFFER | BufferUsage::TRANSFER_SRC,..Default::default() }, AllocationCreateInfo { memory_type_filter: MemoryTypeFilter::PREFER_DEVICE| MemoryTypeFilter::HOST_SEQUENTIAL_WRITE, ..Default::default() }, value).unwrap())); } pub fn add_input_vec<T: BufferContents, F: FnMut(&mut [T])>(&mut self, len: u64, mut f: F) { let idx = self.set_writers.len() as u32; let buf = Buffer::new_slice::<T>(self.allocator.clone(), BufferCreateInfo {usage: BufferUsage::STORAGE_BUFFER | BufferUsage::TRANSFER_SRC,..Default::default() }, AllocationCreateInfo { memory_type_filter: MemoryTypeFilter::PREFER_DEVICE| MemoryTypeFilter::HOST_SEQUENTIAL_WRITE, ..Default::default() }, len).unwrap(); let mut mapping = buf.write().unwrap(); f(&mut mapping.as_mut()); self.set_writers.push(WriteDescriptorSet::buffer(idx, buf.clone())); }
pub fn add_output_vec<T: BufferContents>(&mut self, len: u64)-> Subbuffer<[T]> { let idx = self.set_writers.len() as u32; let buf = Buffer::new_slice::<T>(self.allocator.clone(), BufferCreateInfo {usage: BufferUsage::STORAGE_BUFFER | BufferUsage::TRANSFER_DST,..Default::default() }, AllocationCreateInfo { memory_type_filter: MemoryTypeFilter::PREFER_DEVICE | MemoryTypeFilter::HOST_SEQUENTIAL_WRITE, ..Default::default() }, len).unwrap(); self.set_writers.push(WriteDescriptorSet::buffer(idx, buf.clone())); buf }}复制代码
这样在调用的时候，可以 简单的用下面的代码创建参数
let mut rt = spirv::Runtime::default(); let mut args = rt.get_args(); args.add_input_vec::<u8, _>(2560 * 1440, |buf| { buf.copy_from_slice(yuv.y_plane.borrow()); }); args.add_input_vec::<u8, _>(2560 * 1440, |buf| { buf.copy_from_slice(yuv.u_plane.borrow()); }); args.add_input_vec::<u8, _>(2560 * 1440, |buf| { buf.copy_from_slice(yuv.v_plane.borrow()); }); let buf = args.add_output_vec::<u32>(2560 * 1440); rt.prepare(&spirv_code, args)?; rt.run([80, 44, 1])?; let buf = buf.read()?; let image_buf: &[u8] = bytemuck::cast_slice(&buf); let img = image::RgbaImage::from_raw(2560, 1440, image_buf.to_vec()).unwrap(); img.save("output.png")?;复制代码
正好对应了 刚才 zeta 程序的 三个输入一个输出参数
这样数据的大小就可以动态确定了
我们 把之前的 run 分割为 prepare 和 run 两个部分 ，这样可以方便的进行性能测试
转换的结果是这样的
完美！！！我接下来会把 host 主机程序 源代码 全部开放，还有 Shader 的 zeta 代码 以及 shader 的 可执行模块，这样读者就可以 在任何 rust 环境下 使用强大的 Shader
、Shader 会在国庆期间逐步提供 计划的路线是这样的
、
YUV <-> RGB
Mandl;ebrot 集合的绘制
高性能并行排序算法 Bitonic
高性能 BWT 算法
高性能障碍绘制 和 寻路算法
敬请期待
划线
评论
复制
发布于: 刚刚阅读数: 3
Miracle
关注
三十年资深码农 2019-10-25 加入
还未添加个人简介
评论