xash3d-fwgs/ref/vk/ray_pass.c
2023-02-16 10:30:31 -08:00

286 lines
8.7 KiB
C

#include "ray_pass.h"
#include "shaders/ray_interop.h" // for SPEC_SBT_RECORD_SIZE_INDEX
#include "ray_resources.h"
#include "vk_pipeline.h"
#include "vk_descriptor.h"
// FIXME this is only needed for MAX_CONCURRENT_FRAMES
// TODO specify it externally as ctor arg
#include "vk_framectl.h"
#define MAX_STAGES 16
#define MAX_MISS_GROUPS 8
#define MAX_HIT_GROUPS 8
typedef enum {
RayPassType_Compute,
RayPassType_Tracing,
} ray_pass_type_t;
typedef struct ray_pass_s {
ray_pass_type_t type; // TODO remove this in favor of VkPipelineStageFlagBits
VkPipelineStageFlagBits pipeline_type;
char debug_name[32];
struct {
int write_from;
vk_descriptors_t riptors;
VkDescriptorSet sets[MAX_CONCURRENT_FRAMES];
} desc;
} ray_pass_t;
typedef struct {
ray_pass_t header;
vk_pipeline_ray_t pipeline;
} ray_pass_tracing_impl_t;
typedef struct {
ray_pass_t header;
VkPipeline pipeline;
} ray_pass_compute_impl_t;
static void initPassDescriptors( ray_pass_t *header, const ray_pass_layout_t *layout ) {
header->desc.riptors = (vk_descriptors_t) {
.bindings = layout->bindings,
.num_bindings = layout->bindings_count,
.num_sets = COUNTOF(header->desc.sets),
.desc_sets = header->desc.sets,
.push_constants = layout->push_constants,
};
VK_DescriptorsCreate(&header->desc.riptors);
header->desc.write_from = layout->write_from;
}
static void finalizePassDescriptors( ray_pass_t *header, const ray_pass_layout_t *layout ) {
const size_t bindings_size = sizeof(layout->bindings[0]) * layout->bindings_count;
VkDescriptorSetLayoutBinding *bindings = Mem_Malloc(vk_core.pool, bindings_size);
memcpy(bindings, layout->bindings, bindings_size);
header->desc.riptors.bindings = bindings;
header->desc.riptors.values = Mem_Malloc(vk_core.pool, sizeof(header->desc.riptors.values[0]) * layout->bindings_count);
}
struct ray_pass_s *RayPassCreateTracing( const ray_pass_create_tracing_t *create ) {
ray_pass_tracing_impl_t *const pass = Mem_Malloc(vk_core.pool, sizeof(*pass));
ray_pass_t *const header = &pass->header;
// TODO support external specialization
ASSERT(!create->specialization);
const struct SpecializationData {
uint32_t sbt_record_size;
} spec_data = {
.sbt_record_size = vk_core.physical_device.sbt_record_size,
};
const VkSpecializationMapEntry spec_map[] = {
{.constantID = SPEC_SBT_RECORD_SIZE_INDEX, .offset = offsetof(struct SpecializationData, sbt_record_size), .size = sizeof(uint32_t) },
};
const VkSpecializationInfo spec = {
.mapEntryCount = COUNTOF(spec_map),
.pMapEntries = spec_map,
.dataSize = sizeof(spec_data),
.pData = &spec_data,
};
initPassDescriptors(header, &create->layout);
{
int stage_index = 0;
vk_shader_stage_t stages[MAX_STAGES];
int miss_index = 0;
int misses[MAX_MISS_GROUPS];
int hit_index = 0;
vk_pipeline_ray_hit_group_t hits[MAX_HIT_GROUPS];
vk_pipeline_ray_create_info_t prci = {
.debug_name = create->debug_name,
.layout = header->desc.riptors.pipeline_layout,
.stages = stages,
.groups = {
.hit = hits,
.miss = misses,
},
};
stages[stage_index++] = (vk_shader_stage_t) {
.module = create->raygen_module,
.filename = NULL,
.stage = VK_SHADER_STAGE_RAYGEN_BIT_KHR,
.specialization_info = &spec,
};
for (int i = 0; i < create->miss_count; ++i) {
const VkShaderModule shader_module = create->miss_module ? create->miss_module[i] : VK_NULL_HANDLE;
ASSERT(stage_index < MAX_STAGES);
ASSERT(miss_index < MAX_MISS_GROUPS);
// TODO handle duplicate filenames
// TODO really, there should be a global table of shader modules as some of them are used across several passes (e.g. any hit alpha test)
misses[miss_index++] = stage_index;
stages[stage_index++] = (vk_shader_stage_t) {
.module = shader_module,
.filename = NULL,
.stage = VK_SHADER_STAGE_MISS_BIT_KHR,
.specialization_info = &spec,
};
}
for (int i = 0; i < create->hit_count; ++i) {
const ray_pass_hit_group_t *const group = create->hit + i;
ASSERT(hit_index < MAX_HIT_GROUPS);
// TODO handle duplicate filenames
if (group->any_module) {
ASSERT(stage_index < MAX_STAGES);
hits[hit_index].any = stage_index;
stages[stage_index++] = (vk_shader_stage_t) {
.module = group->any_module,
.filename = NULL,
.stage = VK_SHADER_STAGE_ANY_HIT_BIT_KHR,
.specialization_info = &spec,
};
} else {
hits[hit_index].any = -1;
}
if (group->closest_module) {
ASSERT(stage_index < MAX_STAGES);
hits[hit_index].closest = stage_index;
stages[stage_index++] = (vk_shader_stage_t) {
.module = group->closest_module,
.filename = NULL,
.stage = VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR,
.specialization_info = &spec,
};
} else {
hits[hit_index].closest = -1;
}
++hit_index;
}
prci.groups.hit_count = hit_index;
prci.groups.miss_count = miss_index;
prci.stages_count = stage_index;
pass->pipeline = VK_PipelineRayTracingCreate(&prci);
}
if (pass->pipeline.pipeline == VK_NULL_HANDLE) {
VK_DescriptorsDestroy(&header->desc.riptors);
Mem_Free(pass);
return NULL;
}
finalizePassDescriptors(header, &create->layout);
Q_strncpy(header->debug_name, create->debug_name, sizeof(header->debug_name));
header->type = RayPassType_Tracing;
header->pipeline_type = VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR;
return header;
}
struct ray_pass_s *RayPassCreateCompute( const ray_pass_create_compute_t *create ) {
ray_pass_compute_impl_t *const pass = Mem_Malloc(vk_core.pool, sizeof(*pass));
ray_pass_t *const header = &pass->header;
initPassDescriptors(header, &create->layout);
const vk_pipeline_compute_create_info_t pcci = {
.layout = header->desc.riptors.pipeline_layout,
.shader_module = create->shader_module,
.specialization_info = create->specialization,
};
pass->pipeline = VK_PipelineComputeCreate( &pcci );
if (pass->pipeline == VK_NULL_HANDLE) {
VK_DescriptorsDestroy(&header->desc.riptors);
Mem_Free(pass);
return NULL;
}
finalizePassDescriptors(header, &create->layout);
Q_strncpy(header->debug_name, create->debug_name, sizeof(header->debug_name));
header->type = RayPassType_Compute;
header->pipeline_type = VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT;
return header;
}
void RayPassDestroy( struct ray_pass_s *pass ) {
switch (pass->type) {
case RayPassType_Tracing:
{
ray_pass_tracing_impl_t *tracing = (ray_pass_tracing_impl_t*)pass;
VK_PipelineRayTracingDestroy(&tracing->pipeline);
break;
}
case RayPassType_Compute:
{
ray_pass_compute_impl_t *compute = (ray_pass_compute_impl_t*)pass;
vkDestroyPipeline(vk_core.device, compute->pipeline, NULL);
break;
}
}
VK_DescriptorsDestroy(&pass->desc.riptors);
Mem_Free(pass->desc.riptors.values);
Mem_Free((void*)pass->desc.riptors.bindings);
Mem_Free(pass);
}
static void performTracing( VkCommandBuffer cmdbuf, int set_slot, const ray_pass_tracing_impl_t *tracing, int width, int height ) {
vkCmdBindPipeline(cmdbuf, VK_PIPELINE_BIND_POINT_RAY_TRACING_KHR, tracing->pipeline.pipeline);
vkCmdBindDescriptorSets(cmdbuf, VK_PIPELINE_BIND_POINT_RAY_TRACING_KHR, tracing->header.desc.riptors.pipeline_layout, 0, 1, tracing->header.desc.riptors.desc_sets + set_slot, 0, NULL);
VK_PipelineRayTracingTrace(cmdbuf, &tracing->pipeline, width, height);
}
static void performCompute( VkCommandBuffer cmdbuf, int set_slot, const ray_pass_compute_impl_t *compute, int width, int height) {
const uint32_t WG_W = 8;
const uint32_t WG_H = 8;
vkCmdBindPipeline(cmdbuf, VK_PIPELINE_BIND_POINT_COMPUTE, compute->pipeline);
vkCmdBindDescriptorSets(cmdbuf, VK_PIPELINE_BIND_POINT_COMPUTE, compute->header.desc.riptors.pipeline_layout, 0, 1, compute->header.desc.riptors.desc_sets + set_slot, 0, NULL);
vkCmdDispatch(cmdbuf, (width + WG_W - 1) / WG_W, (height + WG_H - 1) / WG_H, 1);
}
void RayPassPerform(struct ray_pass_s *pass, VkCommandBuffer cmdbuf, ray_pass_perform_args_t args ) {
R_VkResourcesPrepareDescriptorsValues(cmdbuf,
(vk_resources_write_descriptors_args_t){
.pipeline = pass->pipeline_type,
.resources = args.resources,
.resources_map = args.resources_map,
.values = pass->desc.riptors.values,
.count = pass->desc.riptors.num_bindings,
.write_begin = pass->desc.write_from,
}
);
VK_DescriptorsWrite(&pass->desc.riptors, args.frame_set_slot);
DEBUG_BEGIN(cmdbuf, pass->debug_name);
switch (pass->type) {
case RayPassType_Tracing:
{
ray_pass_tracing_impl_t *tracing = (ray_pass_tracing_impl_t*)pass;
performTracing(cmdbuf, args.frame_set_slot, tracing, args.width, args.height);
break;
}
case RayPassType_Compute:
{
ray_pass_compute_impl_t *compute = (ray_pass_compute_impl_t*)pass;
performCompute(cmdbuf, args.frame_set_slot, compute, args.width, args.height);
break;
}
}
DEBUG_END(cmdbuf);
}