librashader/librashader-runtime-d3d12/src/mipmap.rs

328 lines
11 KiB
Rust

use crate::descriptor_heap::ResourceWorkHeap;
use crate::resource::{ObtainResourceHandle, ResourceHandleStrategy};
use crate::util::dxc_validate_shader;
use crate::{error, util};
use bytemuck::{Pod, Zeroable};
use d3d12_descriptor_heap::{D3D12DescriptorHeap, D3D12DescriptorHeapSlot};
use librashader_common::Size;
use librashader_runtime::scaling::MipmapSize;
use std::mem::ManuallyDrop;
use windows::Win32::Graphics::Direct3D::Dxc::{
CLSID_DxcLibrary, CLSID_DxcValidator, DxcCreateInstance,
};
use windows::Win32::Graphics::Direct3D12::{
ID3D12DescriptorHeap, ID3D12Device, ID3D12GraphicsCommandList, ID3D12PipelineState,
ID3D12RootSignature, D3D12_COMPUTE_PIPELINE_STATE_DESC,
D3D12_DEFAULT_SHADER_4_COMPONENT_MAPPING, D3D12_RESOURCE_BARRIER, D3D12_RESOURCE_BARRIER_0,
D3D12_RESOURCE_BARRIER_TYPE_UAV, D3D12_RESOURCE_STATE_PIXEL_SHADER_RESOURCE,
D3D12_RESOURCE_STATE_UNORDERED_ACCESS, D3D12_RESOURCE_UAV_BARRIER, D3D12_SHADER_BYTECODE,
D3D12_SHADER_RESOURCE_VIEW_DESC, D3D12_SHADER_RESOURCE_VIEW_DESC_0,
D3D12_SRV_DIMENSION_TEXTURE2D, D3D12_TEX2D_SRV, D3D12_TEX2D_UAV, D3D12_UAV_DIMENSION_TEXTURE2D,
D3D12_UNORDERED_ACCESS_VIEW_DESC, D3D12_UNORDERED_ACCESS_VIEW_DESC_0,
};
use windows::Win32::Graphics::Dxgi::Common::DXGI_FORMAT;
const GENERATE_MIPMAPS_CS: &[u8] = include_bytes!("../shader/mipmap.dxil");
pub struct D3D12MipmapGen {
device: ID3D12Device,
root_signature: ID3D12RootSignature,
pipeline: ID3D12PipelineState,
own_heaps: bool,
}
#[derive(Copy, Clone, Zeroable, Pod)]
#[repr(C)]
struct MipConstants {
inv_out_texel_size: [f32; 2],
src_mip_index: u32,
}
pub struct MipmapGenContext<'a> {
gen: &'a D3D12MipmapGen,
cmd: &'a ID3D12GraphicsCommandList,
heap: &'a mut D3D12DescriptorHeap<ResourceWorkHeap>,
residuals: Vec<D3D12DescriptorHeapSlot<ResourceWorkHeap>>,
residual_barriers: Vec<D3D12_RESOURCE_BARRIER>,
}
impl<'a> MipmapGenContext<'a> {
fn new(
gen: &'a D3D12MipmapGen,
cmd: &'a ID3D12GraphicsCommandList,
heap: &'a mut D3D12DescriptorHeap<ResourceWorkHeap>,
) -> MipmapGenContext<'a> {
Self {
gen,
cmd,
heap,
residuals: Vec::new(),
residual_barriers: Vec::new(),
}
}
/// Generate a set of mipmaps for the resource.
/// This is a "cheap" action and only dispatches a compute shader.
pub fn generate_mipmaps<S: ResourceHandleStrategy<T>, T: ObtainResourceHandle>(
&mut self,
resource: &T,
miplevels: u16,
size: Size<u32>,
format: DXGI_FORMAT,
) -> error::Result<()> {
unsafe {
let (residuals_heap, residual_barriers) = self
.gen
.generate_mipmaps::<S, T>(self.cmd, resource, miplevels, size, format, self.heap)?;
// heap slots always need to be disposed
self.residuals.extend(residuals_heap);
// barriers need to be disposed if the handle strategy is incrementref
S::cleanup_handler(|| self.residual_barriers.extend(residual_barriers));
}
Ok(())
}
fn close(
self,
) -> (
Vec<D3D12DescriptorHeapSlot<ResourceWorkHeap>>,
Vec<D3D12_RESOURCE_BARRIER>,
) {
(self.residuals, self.residual_barriers)
}
}
impl D3D12MipmapGen {
pub fn new(device: &ID3D12Device, own_heaps: bool) -> error::Result<D3D12MipmapGen> {
unsafe {
let library = DxcCreateInstance(&CLSID_DxcLibrary)?;
let validator = DxcCreateInstance(&CLSID_DxcValidator)?;
let blob = dxc_validate_shader(&library, &validator, GENERATE_MIPMAPS_CS)?;
let blob =
std::slice::from_raw_parts(blob.GetBufferPointer().cast(), blob.GetBufferSize());
let root_signature: ID3D12RootSignature = device.CreateRootSignature(0, blob)?;
let desc = D3D12_COMPUTE_PIPELINE_STATE_DESC {
pRootSignature: ManuallyDrop::new(Some(root_signature.clone())),
CS: D3D12_SHADER_BYTECODE {
pShaderBytecode: blob.as_ptr().cast(),
BytecodeLength: blob.len(),
},
NodeMask: 0,
..Default::default()
};
let pipeline = device.CreateComputePipelineState(&desc)?;
drop(ManuallyDrop::into_inner(desc.pRootSignature));
Ok(D3D12MipmapGen {
device: device.clone(),
root_signature,
pipeline,
own_heaps,
})
}
}
/// If own_heap is false, this sets the compute root signature.
/// Otherwise, this does nothing and the root signature is set upon entering a
/// Mipmapping Context
pub fn pin_root_signature(&self, cmd: &ID3D12GraphicsCommandList) {
if !self.own_heaps {
unsafe {
cmd.SetComputeRootSignature(&self.root_signature);
}
}
}
/// Enters a mipmapping compute context.
/// This is a relatively expensive operation if set_heap is true,
/// and should only be done at most a few times per frame.
///
/// If own_heap is false, then you must ensure that the compute root signature
/// is already bound before entering the context.
///
/// The list of returned descriptors must be kept around until the command list has been
/// submitted.
#[must_use]
pub fn mipmapping_context<F, E>(
&self,
cmd: &ID3D12GraphicsCommandList,
work_heap: &mut D3D12DescriptorHeap<ResourceWorkHeap>,
mut f: F,
) -> Result<
(
Vec<D3D12DescriptorHeapSlot<ResourceWorkHeap>>,
Vec<D3D12_RESOURCE_BARRIER>,
),
E,
>
where
F: FnMut(&mut MipmapGenContext) -> Result<(), E>,
{
let heap: ID3D12DescriptorHeap = (&(*work_heap)).into();
unsafe {
cmd.SetPipelineState(&self.pipeline);
if self.own_heaps {
cmd.SetComputeRootSignature(&self.root_signature);
cmd.SetDescriptorHeaps(&[Some(heap)]);
}
}
let mut context = MipmapGenContext::new(self, cmd, work_heap);
f(&mut context)?;
Ok(context.close())
}
/// SAFETY:
/// - handle must be a CPU handle to an SRV
/// - work_heap must have enough descriptors to fit all miplevels.
unsafe fn generate_mipmaps<S: ResourceHandleStrategy<T>, T: ObtainResourceHandle>(
&self,
cmd: &ID3D12GraphicsCommandList,
resource: &T,
miplevels: u16,
size: Size<u32>,
format: DXGI_FORMAT,
work_heap: &mut D3D12DescriptorHeap<ResourceWorkHeap>,
) -> error::Result<(
Vec<D3D12DescriptorHeapSlot<ResourceWorkHeap>>,
Vec<D3D12_RESOURCE_BARRIER>,
)> {
// create views for mipmap generation
let srv = work_heap.allocate_descriptor()?;
unsafe {
let srv_desc = D3D12_SHADER_RESOURCE_VIEW_DESC {
Format: format,
ViewDimension: D3D12_SRV_DIMENSION_TEXTURE2D,
Shader4ComponentMapping: D3D12_DEFAULT_SHADER_4_COMPONENT_MAPPING,
Anonymous: D3D12_SHADER_RESOURCE_VIEW_DESC_0 {
Texture2D: D3D12_TEX2D_SRV {
MipLevels: miplevels as u32,
..Default::default()
},
},
};
self.device
.CreateShaderResourceView(resource.handle(), Some(&srv_desc), *srv.as_ref());
}
let mut heap_slots = Vec::with_capacity(miplevels as usize);
heap_slots.push(srv);
for i in 1..miplevels {
let descriptor = work_heap.allocate_descriptor()?;
let desc = D3D12_UNORDERED_ACCESS_VIEW_DESC {
Format: format,
ViewDimension: D3D12_UAV_DIMENSION_TEXTURE2D,
Anonymous: D3D12_UNORDERED_ACCESS_VIEW_DESC_0 {
Texture2D: D3D12_TEX2D_UAV {
MipSlice: i as u32,
..Default::default()
},
},
};
unsafe {
self.device.CreateUnorderedAccessView(
resource.handle(),
None,
Some(&desc),
*descriptor.as_ref(),
);
}
heap_slots.push(descriptor);
}
unsafe {
cmd.SetComputeRootDescriptorTable(0, *heap_slots[0].as_ref());
}
let mut residual_barriers = Vec::new();
for i in 1..miplevels as u32 {
let scaled = size.scale_mipmap(i);
let mipmap_params = MipConstants {
inv_out_texel_size: [1.0 / scaled.width as f32, 1.0 / scaled.height as f32],
src_mip_index: (i - 1),
};
let mipmap_params = bytemuck::bytes_of(&mipmap_params);
let barriers = [
util::d3d12_get_resource_transition_subresource::<S, _>(
resource,
D3D12_RESOURCE_STATE_PIXEL_SHADER_RESOURCE,
D3D12_RESOURCE_STATE_UNORDERED_ACCESS,
i - 1,
),
util::d3d12_get_resource_transition_subresource::<S, _>(
resource,
D3D12_RESOURCE_STATE_PIXEL_SHADER_RESOURCE,
D3D12_RESOURCE_STATE_UNORDERED_ACCESS,
i,
),
];
unsafe {
cmd.ResourceBarrier(&barriers);
S::cleanup_handler(|| residual_barriers.extend(barriers));
cmd.SetComputeRootDescriptorTable(1, *heap_slots[i as usize].as_ref());
cmd.SetComputeRoot32BitConstants(
2,
(std::mem::size_of::<MipConstants>() / std::mem::size_of::<u32>()) as u32,
mipmap_params.as_ptr().cast(),
0,
);
cmd.Dispatch(
std::cmp::max(scaled.width.div_ceil(8), 1),
std::cmp::max(scaled.height.div_ceil(8), 1),
1,
);
}
let uav_barrier = ManuallyDrop::new(D3D12_RESOURCE_UAV_BARRIER {
pResource: unsafe { S::obtain(resource) },
});
let barriers = [
D3D12_RESOURCE_BARRIER {
Type: D3D12_RESOURCE_BARRIER_TYPE_UAV,
Anonymous: D3D12_RESOURCE_BARRIER_0 { UAV: uav_barrier },
..Default::default()
},
util::d3d12_get_resource_transition_subresource::<S, _>(
resource,
D3D12_RESOURCE_STATE_UNORDERED_ACCESS,
D3D12_RESOURCE_STATE_PIXEL_SHADER_RESOURCE,
i,
),
util::d3d12_get_resource_transition_subresource::<S, _>(
resource,
D3D12_RESOURCE_STATE_UNORDERED_ACCESS,
D3D12_RESOURCE_STATE_PIXEL_SHADER_RESOURCE,
i - 1,
),
];
unsafe {
cmd.ResourceBarrier(&barriers);
}
S::cleanup_handler(|| residual_barriers.extend(barriers));
}
Ok((heap_slots, residual_barriers))
}
}