bevy/crates/bevy_wgpu/src/renderer/wgpu_render_context.rs
2020-04-24 17:57:20 -07:00

566 lines
22 KiB
Rust

use super::WgpuRenderResourceContext;
use crate::{
wgpu_type_converter::{OwnedWgpuVertexBufferDescriptor, WgpuInto},
WgpuRenderPass, WgpuResourceRefs,
};
use bevy_asset::{AssetStorage, Handle};
use bevy_render::{
pass::{
PassDescriptor, RenderPass, RenderPassColorAttachmentDescriptor,
RenderPassDepthStencilAttachmentDescriptor, TextureAttachment,
},
pipeline::{BindGroupDescriptor, BindType, PipelineDescriptor},
render_resource::{
RenderResource, RenderResourceAssignments, RenderResourceSetId, ResourceInfo,
},
renderer::{RenderContext, RenderResourceContext},
shader::Shader,
texture::{Extent3d, TextureDescriptor},
};
use bevy_window::WindowId;
use std::{collections::HashMap, sync::Arc};
#[derive(Default)]
pub struct LazyCommandEncoder {
command_encoder: Option<wgpu::CommandEncoder>,
}
impl LazyCommandEncoder {
pub fn get_or_create(&mut self, device: &wgpu::Device) -> &mut wgpu::CommandEncoder {
match self.command_encoder {
Some(ref mut command_encoder) => command_encoder,
None => {
self.create(device);
self.command_encoder.as_mut().unwrap()
}
}
}
pub fn is_some(&self) -> bool {
self.command_encoder.is_some()
}
pub fn create(&mut self, device: &wgpu::Device) {
let command_encoder =
device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
self.command_encoder = Some(command_encoder);
}
pub fn take(&mut self) -> Option<wgpu::CommandEncoder> {
self.command_encoder.take()
}
pub fn set(&mut self, command_encoder: wgpu::CommandEncoder) {
self.command_encoder = Some(command_encoder);
}
}
pub struct WgpuRenderContext {
pub device: Arc<wgpu::Device>,
// TODO: remove this
pub primary_window: Option<WindowId>,
pub command_encoder: LazyCommandEncoder,
pub render_resources: WgpuRenderResourceContext,
}
impl WgpuRenderContext {
pub fn new(device: Arc<wgpu::Device>, resources: WgpuRenderResourceContext) -> Self {
WgpuRenderContext {
device,
primary_window: None,
render_resources: resources,
command_encoder: LazyCommandEncoder::default(),
}
}
/// Consume this context, finalize the current CommandEncoder (if it exists), and take the current WgpuResources.
/// This is intended to be called from a worker thread right before synchronizing with the main thread.
pub fn finish(&mut self) -> Option<wgpu::CommandBuffer> {
self.command_encoder.take().map(|encoder| encoder.finish())
}
}
impl RenderContext for WgpuRenderContext {
fn create_texture_with_data(
&mut self,
texture_descriptor: &TextureDescriptor,
bytes: &[u8],
) -> RenderResource {
self.render_resources
.wgpu_resources
.create_texture_with_data(
&self.device,
self.command_encoder.get_or_create(&self.device),
texture_descriptor,
bytes,
)
}
fn copy_buffer_to_buffer(
&mut self,
source_buffer: RenderResource,
source_offset: u64,
destination_buffer: RenderResource,
destination_offset: u64,
size: u64,
) {
let command_encoder = self.command_encoder.get_or_create(&self.device);
self.render_resources.wgpu_resources.copy_buffer_to_buffer(
command_encoder,
source_buffer,
source_offset,
destination_buffer,
destination_offset,
size,
);
}
fn resources(&self) -> &dyn RenderResourceContext {
&self.render_resources
}
fn resources_mut(&mut self) -> &mut dyn RenderResourceContext {
&mut self.render_resources
}
fn create_bind_group(
&mut self,
bind_group_descriptor: &BindGroupDescriptor,
render_resource_assignments: &RenderResourceAssignments,
) -> Option<RenderResourceSetId> {
if let Some((render_resource_set_id, _indices)) =
render_resource_assignments.get_render_resource_set_id(bind_group_descriptor.id)
{
if !self
.render_resources
.wgpu_resources
.has_bind_group(bind_group_descriptor.id, *render_resource_set_id)
{
log::trace!(
"start creating bind group for RenderResourceSet {:?}",
render_resource_set_id
);
let wgpu_bind_group = {
let textures = self
.render_resources
.wgpu_resources
.texture_views
.read()
.unwrap();
let samplers = self
.render_resources
.wgpu_resources
.samplers
.read()
.unwrap();
let buffers = self.render_resources.wgpu_resources.buffers.read().unwrap();
let bindings = bind_group_descriptor
.bindings
.iter()
.map(|binding| {
if let Some(resource) = render_resource_assignments.get(&binding.name) {
let mut wgpu_resource = None;
self.resources().get_resource_info(
resource,
&mut |resource_info| {
log::trace!(
"found binding {} ({}) resource: {:?} {:?}",
binding.index,
binding.name,
resource,
resource_info
);
wgpu_resource = match &binding.bind_type {
BindType::SampledTexture { .. } => {
if let Some(ResourceInfo::Texture) = resource_info {
let texture = textures.get(&resource).unwrap();
Some(wgpu::BindingResource::TextureView(
texture,
))
} else {
panic!("expected a Texture resource");
}
}
BindType::Sampler { .. } => {
if let Some(ResourceInfo::Sampler) = resource_info {
let sampler = samplers.get(&resource).unwrap();
Some(wgpu::BindingResource::Sampler(sampler))
} else {
panic!("expected a Sampler resource");
}
}
BindType::Uniform { .. } => {
if let Some(ResourceInfo::Buffer(buffer_info)) =
resource_info
{
let buffer = buffers.get(&resource).unwrap();
Some(wgpu::BindingResource::Buffer {
buffer,
range: 0..buffer_info.size as u64,
})
} else {
panic!("expected a Buffer resource");
}
}
_ => panic!("unsupported bind type"),
}
},
);
wgpu::Binding {
binding: binding.index,
resource: wgpu_resource.expect("No resource binding found"),
}
} else {
panic!(
"No resource assigned to uniform \"{}\" for RenderResourceAssignments {:?}",
binding.name,
render_resource_assignments.id
);
}
})
.collect::<Vec<wgpu::Binding>>();
let bind_group_layouts = self
.render_resources
.wgpu_resources
.bind_group_layouts
.read()
.unwrap();
let bind_group_layout =
bind_group_layouts.get(&bind_group_descriptor.id).unwrap();
let wgpu_bind_group_descriptor = wgpu::BindGroupDescriptor {
label: None,
layout: bind_group_layout,
bindings: bindings.as_slice(),
};
self.render_resources.wgpu_resources.create_bind_group(
&self.device,
*render_resource_set_id,
&wgpu_bind_group_descriptor,
)
};
self.render_resources.wgpu_resources.set_bind_group(
bind_group_descriptor.id,
*render_resource_set_id,
wgpu_bind_group,
);
return Some(*render_resource_set_id);
}
}
None
}
fn create_render_pipeline(
&mut self,
pipeline_handle: Handle<PipelineDescriptor>,
pipeline_descriptor: &PipelineDescriptor,
shader_storage: &AssetStorage<Shader>,
) {
if let Some(_) = self
.render_resources
.wgpu_resources
.render_pipelines
.read()
.unwrap()
.get(&pipeline_handle)
{
return;
}
let layout = pipeline_descriptor.get_layout().unwrap();
for bind_group in layout.bind_groups.iter() {
if self
.render_resources
.wgpu_resources
.bind_group_layouts
.read()
.unwrap()
.get(&bind_group.id)
.is_none()
{
let bind_group_layout_binding = bind_group
.bindings
.iter()
.map(|binding| wgpu::BindGroupLayoutEntry {
binding: binding.index,
visibility: wgpu::ShaderStage::VERTEX | wgpu::ShaderStage::FRAGMENT,
ty: (&binding.bind_type).wgpu_into(),
})
.collect::<Vec<wgpu::BindGroupLayoutEntry>>();
self.render_resources
.wgpu_resources
.create_bind_group_layout(
&self.device,
bind_group.id,
&wgpu::BindGroupLayoutDescriptor {
bindings: bind_group_layout_binding.as_slice(),
label: None,
},
);
}
}
let pipeline_layout = {
let bind_group_layouts = self
.render_resources
.wgpu_resources
.bind_group_layouts
.read()
.unwrap();
// setup and collect bind group layouts
let bind_group_layouts = layout
.bind_groups
.iter()
.map(|bind_group| bind_group_layouts.get(&bind_group.id).unwrap())
.collect::<Vec<&wgpu::BindGroupLayout>>();
self.device
.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
bind_group_layouts: bind_group_layouts.as_slice(),
})
};
let owned_vertex_buffer_descriptors = layout
.vertex_buffer_descriptors
.iter()
.map(|v| v.wgpu_into())
.collect::<Vec<OwnedWgpuVertexBufferDescriptor>>();
let color_states = pipeline_descriptor
.color_states
.iter()
.map(|c| c.wgpu_into())
.collect::<Vec<wgpu::ColorStateDescriptor>>();
if self
.render_resources
.wgpu_resources
.shader_modules
.read()
.unwrap()
.get(&pipeline_descriptor.shader_stages.vertex)
.is_none()
{
self.render_resources
.create_shader_module(pipeline_descriptor.shader_stages.vertex, shader_storage);
}
if let Some(fragment_handle) = pipeline_descriptor.shader_stages.fragment {
if self
.render_resources
.wgpu_resources
.shader_modules
.read()
.unwrap()
.get(&fragment_handle)
.is_none()
{
self.render_resources
.create_shader_module(fragment_handle, shader_storage);
}
};
let wgpu_pipeline = {
let shader_modules = self
.render_resources
.wgpu_resources
.shader_modules
.read()
.unwrap();
let vertex_shader_module = shader_modules
.get(&pipeline_descriptor.shader_stages.vertex)
.unwrap();
let fragment_shader_module = match pipeline_descriptor.shader_stages.fragment {
Some(fragment_handle) => Some(shader_modules.get(&fragment_handle).unwrap()),
None => None,
};
let render_pipeline_descriptor = wgpu::RenderPipelineDescriptor {
layout: &pipeline_layout,
vertex_stage: wgpu::ProgrammableStageDescriptor {
module: &vertex_shader_module,
entry_point: "main",
},
fragment_stage: match pipeline_descriptor.shader_stages.fragment {
Some(_) => Some(wgpu::ProgrammableStageDescriptor {
entry_point: "main",
module: fragment_shader_module.as_ref().unwrap(),
}),
None => None,
},
rasterization_state: pipeline_descriptor
.rasterization_state
.as_ref()
.map(|r| r.wgpu_into()),
primitive_topology: pipeline_descriptor.primitive_topology.wgpu_into(),
color_states: &color_states,
depth_stencil_state: pipeline_descriptor
.depth_stencil_state
.as_ref()
.map(|d| d.wgpu_into()),
vertex_state: wgpu::VertexStateDescriptor {
index_format: pipeline_descriptor.index_format.wgpu_into(),
vertex_buffers: &owned_vertex_buffer_descriptors
.iter()
.map(|v| v.into())
.collect::<Vec<wgpu::VertexBufferDescriptor>>(),
},
sample_count: pipeline_descriptor.sample_count,
sample_mask: pipeline_descriptor.sample_mask,
alpha_to_coverage_enabled: pipeline_descriptor.alpha_to_coverage_enabled,
};
self.render_resources
.wgpu_resources
.create_render_pipeline(&self.device, &render_pipeline_descriptor)
};
self.render_resources
.wgpu_resources
.set_render_pipeline(pipeline_handle, wgpu_pipeline);
}
fn begin_pass(
&mut self,
pass_descriptor: &PassDescriptor,
render_resource_assignments: &RenderResourceAssignments,
run_pass: &mut dyn Fn(&mut dyn RenderPass),
) {
if !self.command_encoder.is_some() {
self.command_encoder.create(&self.device);
}
let resource_lock = self.render_resources.wgpu_resources.read();
let refs = resource_lock.refs();
let mut encoder = self.command_encoder.take().unwrap();
{
let render_pass = create_render_pass(
pass_descriptor,
render_resource_assignments,
&refs,
&mut encoder,
);
let mut wgpu_render_pass = WgpuRenderPass {
render_context: self,
render_pass,
render_resources: refs,
bound_bind_groups: HashMap::default(),
};
run_pass(&mut wgpu_render_pass);
}
self.command_encoder.set(encoder);
}
fn copy_buffer_to_texture(
&mut self,
source_buffer: RenderResource,
source_offset: u64,
source_bytes_per_row: u32,
destination_texture: RenderResource,
destination_origin: [u32; 3],
destination_mip_level: u32,
destination_array_layer: u32,
size: Extent3d,
) {
self.render_resources.wgpu_resources.copy_buffer_to_texture(
self.command_encoder.get_or_create(&self.device),
source_buffer,
source_offset,
source_bytes_per_row,
destination_texture,
destination_origin,
destination_mip_level,
destination_array_layer,
size,
)
}
}
pub fn create_render_pass<'a, 'b>(
pass_descriptor: &PassDescriptor,
global_render_resource_assignments: &'b RenderResourceAssignments,
refs: &WgpuResourceRefs<'a>,
encoder: &'a mut wgpu::CommandEncoder,
) -> wgpu::RenderPass<'a> {
encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
color_attachments: &pass_descriptor
.color_attachments
.iter()
.map(|c| {
create_wgpu_color_attachment_descriptor(global_render_resource_assignments, refs, c)
})
.collect::<Vec<wgpu::RenderPassColorAttachmentDescriptor>>(),
depth_stencil_attachment: pass_descriptor.depth_stencil_attachment.as_ref().map(|d| {
create_wgpu_depth_stencil_attachment_descriptor(
global_render_resource_assignments,
refs,
d,
)
}),
})
}
fn get_texture_view<'a>(
global_render_resource_assignments: &RenderResourceAssignments,
refs: &WgpuResourceRefs<'a>,
attachment: &TextureAttachment,
) -> &'a wgpu::TextureView {
match attachment {
TextureAttachment::Name(name) => match global_render_resource_assignments.get(&name) {
Some(resource) => refs.textures.get(&resource).unwrap(),
None => {
panic!("Color attachment {} does not exist", name);
}
},
TextureAttachment::RenderResource(render_resource) => refs.textures.get(&render_resource).unwrap_or_else(|| &refs.swap_chain_outputs.get(&render_resource).unwrap().view),
TextureAttachment::Input(_) => panic!("Encountered unset TextureAttachment::Input. The RenderGraph executor should always set TextureAttachment::Inputs to TextureAttachment::RenderResource before running. This is a bug"),
}
}
fn create_wgpu_color_attachment_descriptor<'a>(
global_render_resource_assignments: &RenderResourceAssignments,
refs: &WgpuResourceRefs<'a>,
color_attachment_descriptor: &RenderPassColorAttachmentDescriptor,
) -> wgpu::RenderPassColorAttachmentDescriptor<'a> {
let attachment = get_texture_view(
global_render_resource_assignments,
refs,
&color_attachment_descriptor.attachment,
);
let resolve_target = color_attachment_descriptor
.resolve_target
.as_ref()
.map(|target| get_texture_view(global_render_resource_assignments, refs, &target));
wgpu::RenderPassColorAttachmentDescriptor {
store_op: color_attachment_descriptor.store_op.wgpu_into(),
load_op: color_attachment_descriptor.load_op.wgpu_into(),
clear_color: color_attachment_descriptor.clear_color.wgpu_into(),
attachment,
resolve_target,
}
}
fn create_wgpu_depth_stencil_attachment_descriptor<'a>(
global_render_resource_assignments: &RenderResourceAssignments,
refs: &WgpuResourceRefs<'a>,
depth_stencil_attachment_descriptor: &RenderPassDepthStencilAttachmentDescriptor,
) -> wgpu::RenderPassDepthStencilAttachmentDescriptor<'a> {
let attachment = get_texture_view(
global_render_resource_assignments,
refs,
&depth_stencil_attachment_descriptor.attachment,
);
wgpu::RenderPassDepthStencilAttachmentDescriptor {
attachment,
clear_depth: depth_stencil_attachment_descriptor.clear_depth,
clear_stencil: depth_stencil_attachment_descriptor.clear_stencil,
depth_load_op: depth_stencil_attachment_descriptor
.depth_load_op
.wgpu_into(),
depth_store_op: depth_stencil_attachment_descriptor
.depth_store_op
.wgpu_into(),
stencil_load_op: depth_stencil_attachment_descriptor
.stencil_load_op
.wgpu_into(),
stencil_store_op: depth_stencil_attachment_descriptor
.stencil_store_op
.wgpu_into(),
}
}