diff --git a/samples/general/mobile_nerf/mobile_nerf.cpp b/samples/general/mobile_nerf/mobile_nerf.cpp index d614e79f2..1f996056e 100644 --- a/samples/general/mobile_nerf/mobile_nerf.cpp +++ b/samples/general/mobile_nerf/mobile_nerf.cpp @@ -24,7 +24,6 @@ #include "scene_graph/components/material.h" #include "scene_graph/components/mesh.h" #include "scene_graph/components/perspective_camera.h" -#include "stb_image.h" namespace { @@ -100,6 +99,8 @@ MobileNerf::MobileNerf() add_device_extension(VK_KHR_SPIRV_1_4_EXTENSION_NAME); // Required by VK_KHR_spirv_1_4 add_device_extension(VK_KHR_SHADER_FLOAT_CONTROLS_EXTENSION_NAME); + // For choosing different sets of weights + add_device_extension(VK_EXT_DESCRIPTOR_INDEXING_EXTENSION_NAME); } MobileNerf::~MobileNerf() @@ -157,6 +158,7 @@ MobileNerf::~MobileNerf() attachment.feature_0.destroy(); attachment.feature_1.destroy(); attachment.feature_2.destroy(); + attachment.weights_idx.destroy(); } } } @@ -371,13 +373,17 @@ void MobileNerf::load_shaders() // Loading first pass shaders shader_stages_first_pass[0] = load_shader("mobile_nerf/raster.vert", VK_SHADER_STAGE_VERTEX_BIT); shader_stages_first_pass[1] = load_shader( - using_original_nerf_models[0] ? "mobile_nerf/raster.frag" : "mobile_nerf/raster_morpheus.frag", + combo_mode ? + (using_original_nerf_models[0] ? "mobile_nerf/raster_combo.frag" : "mobile_nerf/raster_morpheus_combo.frag") : + (using_original_nerf_models[0] ? "mobile_nerf/raster.frag" : "mobile_nerf/raster_morpheus.frag"), VK_SHADER_STAGE_FRAGMENT_BIT); // Loading second pass shaders shader_stages_second_pass[0] = load_shader("mobile_nerf/quad.vert", VK_SHADER_STAGE_VERTEX_BIT); shader_stages_second_pass[1] = load_shader( - using_original_nerf_models[0] ? "mobile_nerf/mlp.frag" : "mobile_nerf/mlp_morpheus.frag", + combo_mode ? + (using_original_nerf_models[0] ? "mobile_nerf/mlp_combo.frag" : "mobile_nerf/mlp_morpheus_combo.frag") : + (using_original_nerf_models[0] ? "mobile_nerf/mlp.frag" : "mobile_nerf/mlp_morpheus.frag"), VK_SHADER_STAGE_FRAGMENT_BIT); } else @@ -487,6 +493,10 @@ bool MobileNerf::resize(const uint32_t width, const uint32_t height) void MobileNerf::request_gpu_features(vkb::PhysicalDevice &gpu) { + RequestFeature(gpu) + .request(VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_DESCRIPTOR_INDEXING_FEATURES_EXT, &VkPhysicalDeviceDescriptorIndexingFeaturesEXT::shaderUniformBufferArrayNonUniformIndexing) + .request(VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_DESCRIPTOR_INDEXING_FEATURES_EXT, &VkPhysicalDeviceDescriptorIndexingFeaturesEXT::runtimeDescriptorArray) + .request(VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_DESCRIPTOR_INDEXING_FEATURES_EXT, &VkPhysicalDeviceDescriptorIndexingFeaturesEXT::descriptorBindingVariableDescriptorCount); } void MobileNerf::render(float delta_time) @@ -499,11 +509,6 @@ void MobileNerf::render(float delta_time) update_uniform_buffers(); } -inline uint32_t aligned_size(uint32_t value, uint32_t alignment) -{ - return (value + alignment - 1) & ~(alignment - 1); -} - void MobileNerf::FrameBufferAttachment::destroy() { if (!image) @@ -512,8 +517,6 @@ void MobileNerf::FrameBufferAttachment::destroy() } auto &device = image->get_device(); - vkDestroySampler(device.get_handle(), sampler, nullptr); - sampler = VK_NULL_HANDLE; vkDestroyImageView(device.get_handle(), view, nullptr); image.reset(); } @@ -543,6 +546,10 @@ void MobileNerf::setup_attachment(VkFormat format, VkImageUsageFlags usage, Fram with_command_buffer([&](VkCommandBuffer command_buffer) { vkb::image_layout_transition(command_buffer, attachment.image->get_handle(), + VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, + VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, + {}, + {}, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_GENERAL, {aspectMask, 0, 1, 0, 1}); @@ -558,24 +565,6 @@ void MobileNerf::setup_attachment(VkFormat format, VkImageUsageFlags usage, Fram color_image_view.subresourceRange.layerCount = 1; color_image_view.image = attachment.image->get_handle(); VK_CHECK(vkCreateImageView(get_device().get_handle(), &color_image_view, nullptr, &attachment.view)); - - // Calculate valid filter - VkFilter filter = VK_FILTER_LINEAR; - vkb::make_filters_valid(get_device().get_gpu().get_handle(), format, &filter); - - VkSamplerCreateInfo samplerCreateInfo = {}; - - samplerCreateInfo.sType = VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO; - samplerCreateInfo.magFilter = filter; - samplerCreateInfo.minFilter = filter; - samplerCreateInfo.mipmapMode = VK_SAMPLER_MIPMAP_MODE_NEAREST; - samplerCreateInfo.addressModeU = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE; - samplerCreateInfo.addressModeV = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE; - samplerCreateInfo.addressModeW = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE; - samplerCreateInfo.minLod = 0; - samplerCreateInfo.maxLod = 16.f; - - VK_CHECK(vkCreateSampler(get_device().get_handle(), &samplerCreateInfo, 0, &attachment.sampler)); } void MobileNerf::setup_nerf_framebuffer_baseline() @@ -589,6 +578,8 @@ void MobileNerf::setup_nerf_framebuffer_baseline() setup_attachment(feature_map_format, VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_INPUT_ATTACHMENT_BIT, frameAttachments[i].feature_0); setup_attachment(feature_map_format, VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_INPUT_ATTACHMENT_BIT, frameAttachments[i].feature_1); setup_attachment(VK_FORMAT_R16G16B16A16_SFLOAT, VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_INPUT_ATTACHMENT_BIT, frameAttachments[i].feature_2); + if (combo_mode) + setup_attachment(VK_FORMAT_R8_UINT, VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_INPUT_ATTACHMENT_BIT, frameAttachments[i].weights_idx); } } @@ -608,8 +599,8 @@ void MobileNerf::setup_nerf_framebuffer_baseline() if (use_deferred) { - views.resize(5); - views[3] = depth_stencil.view; + views.resize(combo_mode ? 6 : 5); + views[depth_attach_idx] = depth_stencil.view; } else { @@ -635,10 +626,12 @@ void MobileNerf::setup_nerf_framebuffer_baseline() { if (use_deferred) { - views[0] = frameAttachments[i].feature_0.view; - views[1] = frameAttachments[i].feature_1.view; - views[2] = frameAttachments[i].feature_2.view; - views[4] = swapchain_buffers[i].view; + views[color_attach_0_idx] = frameAttachments[i].feature_0.view; + views[color_attach_1_idx] = frameAttachments[i].feature_1.view; + views[color_attach_2_idx] = frameAttachments[i].feature_2.view; + if (combo_mode) + views[color_attach_3_idx] = frameAttachments[i].weights_idx.view; + views[swapchain_attach_idx] = swapchain_buffers[i].view; } else { @@ -653,7 +646,8 @@ void MobileNerf::update_descriptor_sets_baseline() { for (int i = 0; i < nerf_framebuffers.size(); i++) { - std::array attachment_input_descriptors; + std::vector attachment_input_descriptors; + attachment_input_descriptors.resize(combo_mode ? 4 : 3); attachment_input_descriptors[0].sampler = VK_NULL_HANDLE; attachment_input_descriptors[0].imageView = frameAttachments[i].feature_0.view; @@ -667,16 +661,38 @@ void MobileNerf::update_descriptor_sets_baseline() attachment_input_descriptors[2].imageView = frameAttachments[i].feature_2.view; attachment_input_descriptors[2].imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL; + if (combo_mode) + { + attachment_input_descriptors[3].sampler = VK_NULL_HANDLE; + attachment_input_descriptors[3].imageView = frameAttachments[i].weights_idx.view; + attachment_input_descriptors[3].imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL; + } + VkWriteDescriptorSet texture_input_write_0 = vkb::initializers::write_descriptor_set(descriptor_set_baseline[i], VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, 0, &attachment_input_descriptors[0]); VkWriteDescriptorSet texture_input_write_1 = vkb::initializers::write_descriptor_set(descriptor_set_baseline[i], VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, 1, &attachment_input_descriptors[1]); VkWriteDescriptorSet texture_input_write_2 = vkb::initializers::write_descriptor_set(descriptor_set_baseline[i], VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, 2, &attachment_input_descriptors[2]); - std::vector write_descriptor_sets = { - texture_input_write_0, - texture_input_write_1, - texture_input_write_2}; + if (combo_mode) + { + VkWriteDescriptorSet texture_input_write_3 = vkb::initializers::write_descriptor_set(descriptor_set_baseline[i], VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, 3, &attachment_input_descriptors[3]); - vkUpdateDescriptorSets(get_device().get_handle(), static_cast(write_descriptor_sets.size()), write_descriptor_sets.data(), 0, VK_NULL_HANDLE); + std::vector write_descriptor_sets = { + texture_input_write_0, + texture_input_write_1, + texture_input_write_2, + texture_input_write_3}; + + vkUpdateDescriptorSets(get_device().get_handle(), static_cast(write_descriptor_sets.size()), write_descriptor_sets.data(), 0, VK_NULL_HANDLE); + } + else + { + std::vector write_descriptor_sets = { + texture_input_write_0, + texture_input_write_1, + texture_input_write_2}; + + vkUpdateDescriptorSets(get_device().get_handle(), static_cast(write_descriptor_sets.size()), write_descriptor_sets.data(), 0, VK_NULL_HANDLE); + } } } @@ -709,12 +725,25 @@ void MobileNerf::build_command_buffers_baseline() if (use_deferred) { - clear_values.resize(5); - clear_values[0].color = {{0.025f, 0.025f, 0.025f, 0.5f}}; // default_clear_color; - clear_values[1].color = {{0.025f, 0.025f, 0.025f, 0.5f}}; // default_clear_color; - clear_values[2].color = {{0.025f, 0.025f, 0.025f, 0.5f}}; // default_clear_color; - clear_values[3].depthStencil = {1.0f, 0}; - clear_values[4].color = {{1.0f, 1.0f, 1.0f, 0.5f}}; // default_clear_color; + if (combo_mode) + { + clear_values.resize(6); + clear_values[0].color = {{0.025f, 0.025f, 0.025f, 0.5f}}; // default_clear_color; + clear_values[1].color = {{0.025f, 0.025f, 0.025f, 0.5f}}; // default_clear_color; + clear_values[2].color = {{0.025f, 0.025f, 0.025f, 0.5f}}; // default_clear_color; + clear_values[3].color = {{0, 0, 0, 0}}; // default clear index for weights; + clear_values[4].depthStencil = {1.0f, 0}; + clear_values[5].color = {{1.0f, 1.0f, 1.0f, 0.5f}}; // default_clear_color; + } + else + { + clear_values.resize(5); + clear_values[0].color = {{0.025f, 0.025f, 0.025f, 0.5f}}; // default_clear_color; + clear_values[1].color = {{0.025f, 0.025f, 0.025f, 0.5f}}; // default_clear_color; + clear_values[2].color = {{0.025f, 0.025f, 0.025f, 0.5f}}; // default_clear_color; + clear_values[3].depthStencil = {1.0f, 0}; + clear_values[4].color = {{1.0f, 1.0f, 1.0f, 0.5f}}; // default_clear_color; + } } else { @@ -776,6 +805,17 @@ void MobileNerf::build_command_buffers_baseline() vkCmdBindVertexBuffers(draw_cmd_buffers[i], 0, 1, model.vertex_buffer->get(), offsets); vkCmdBindVertexBuffers(draw_cmd_buffers[i], 1, 1, instance_buffer->get(), offsets); vkCmdBindIndexBuffer(draw_cmd_buffers[i], model.index_buffer->get_handle(), 0, VK_INDEX_TYPE_UINT32); + if (use_deferred && combo_mode) + { + PushConstants constants = {static_cast(model.model_index)}; + vkCmdPushConstants( + draw_cmd_buffers[i], + pipeline_first_pass_layout, + VK_SHADER_STAGE_FRAGMENT_BIT, + 0, + sizeof(PushConstants), + &constants); + } vkCmdDrawIndexed(draw_cmd_buffers[i], static_cast(model.indices.size()) * 3, ii.dim.x * ii.dim.y * ii.dim.z, 0, 0, 0); } @@ -885,9 +925,18 @@ void MobileNerf::create_descriptor_pool() // First Pass {VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 2 * static_cast(models.size())}, {VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 1 * static_cast(models.size())}, - // Second Pass - {VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, 3 * static_cast(framebuffers.size())}, - {VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 1 * static_cast(framebuffers.size())}}; + }; + // Second Pass + if (combo_mode) + { + pool_sizes.push_back({VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, 4 * static_cast(framebuffers.size())}); + pool_sizes.push_back({VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 1 * static_cast(framebuffers.size()) * static_cast(model_path.size())}); + } + else + { + pool_sizes.push_back({VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, 3 * static_cast(framebuffers.size())}); + pool_sizes.push_back({VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 1 * static_cast(framebuffers.size())}); + } VkDescriptorPoolCreateInfo descriptor_pool_create_info = vkb::initializers::descriptor_pool_create_info(pool_sizes, static_cast(models.size() + framebuffers.size())); VK_CHECK(vkCreateDescriptorPool(get_device().get_handle(), &descriptor_pool_create_info, nullptr, &descriptor_pool)); @@ -929,7 +978,18 @@ void MobileNerf::create_pipeline_layout_fist_pass() &descriptor_set_first_pass_layout, 1); - VK_CHECK(vkCreatePipelineLayout(get_device().get_handle(), &pipeline_layout_create_info, nullptr, &pipeline_first_pass_layout)); + if (use_deferred && combo_mode) + { + VkPushConstantRange pushConstantRange = vkb::initializers::push_constant_range(VK_SHADER_STAGE_FRAGMENT_BIT, sizeof(PushConstants), 0); + + pipeline_layout_create_info.pushConstantRangeCount = 1; + pipeline_layout_create_info.pPushConstantRanges = &pushConstantRange; + VK_CHECK(vkCreatePipelineLayout(get_device().get_handle(), &pipeline_layout_create_info, nullptr, &pipeline_first_pass_layout)); + } + else + { + VK_CHECK(vkCreatePipelineLayout(get_device().get_handle(), &pipeline_layout_create_info, nullptr, &pipeline_first_pass_layout)); + } } void MobileNerf::create_pipeline_layout_baseline() @@ -941,12 +1001,34 @@ void MobileNerf::create_pipeline_layout_baseline() vkb::initializers::descriptor_set_layout_binding(VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, VK_SHADER_STAGE_FRAGMENT_BIT, 0), vkb::initializers::descriptor_set_layout_binding(VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, VK_SHADER_STAGE_FRAGMENT_BIT, 1), vkb::initializers::descriptor_set_layout_binding(VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, VK_SHADER_STAGE_FRAGMENT_BIT, 2), - // MLP weights - vkb::initializers::descriptor_set_layout_binding(VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, VK_SHADER_STAGE_FRAGMENT_BIT, 3) // SSBO }; + if (combo_mode) + { + set_layout_bindings.push_back(vkb::initializers::descriptor_set_layout_binding(VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, VK_SHADER_STAGE_FRAGMENT_BIT, 3)); + // MLP weights array, using descriptor indexing + set_layout_bindings.push_back(vkb::initializers::descriptor_set_layout_binding(VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, VK_SHADER_STAGE_FRAGMENT_BIT, 4, static_cast(model_path.size()))); + } + else + { + // MLP weights + set_layout_bindings.push_back(vkb::initializers::descriptor_set_layout_binding(VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, VK_SHADER_STAGE_FRAGMENT_BIT, 3)); + } VkDescriptorSetLayoutCreateInfo descriptor_layout = vkb::initializers::descriptor_set_layout_create_info(set_layout_bindings.data(), static_cast(set_layout_bindings.size())); - VK_CHECK(vkCreateDescriptorSetLayout(get_device().get_handle(), &descriptor_layout, nullptr, &descriptor_set_layout_baseline)); + if (combo_mode) + { + VkDescriptorBindingFlagsEXT flags[5] = {0, 0, 0, 0, VK_DESCRIPTOR_BINDING_VARIABLE_DESCRIPTOR_COUNT_BIT_EXT}; + VkDescriptorSetLayoutBindingFlagsCreateInfoEXT setLayoutBindingFlags{}; + setLayoutBindingFlags.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_BINDING_FLAGS_CREATE_INFO_EXT; + setLayoutBindingFlags.bindingCount = 5; + setLayoutBindingFlags.pBindingFlags = flags; + descriptor_layout.pNext = &setLayoutBindingFlags; + VK_CHECK(vkCreateDescriptorSetLayout(get_device().get_handle(), &descriptor_layout, nullptr, &descriptor_set_layout_baseline)); + } + else + { + VK_CHECK(vkCreateDescriptorSetLayout(get_device().get_handle(), &descriptor_layout, nullptr, &descriptor_set_layout_baseline)); + } VkPipelineLayoutCreateInfo pipeline_layout_create_info = vkb::initializers::pipeline_layout_create_info( @@ -977,13 +1059,14 @@ void MobileNerf::create_descriptor_sets_first_pass(Model &model) texture_input_descriptors[1].imageView = model.texture_input_1.image->get_vk_image_view().get_handle(); texture_input_descriptors[1].imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL; - VkDescriptorBufferInfo buffer_descriptor = create_descriptor(**model.uniform_buffer_ref); + VkDescriptorBufferInfo buffer_descriptor = create_descriptor(*uniform_buffers[model.model_index]); VkWriteDescriptorSet texture_input_write_0 = vkb::initializers::write_descriptor_set(model.descriptor_set_first_pass[i], VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 0, &texture_input_descriptors[0]); VkWriteDescriptorSet texture_input_write_1 = vkb::initializers::write_descriptor_set(model.descriptor_set_first_pass[i], VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 1, &texture_input_descriptors[1]); - VkWriteDescriptorSet uniform_buffer_write = vkb::initializers::write_descriptor_set(model.descriptor_set_first_pass[i], VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 2, &buffer_descriptor); + VkWriteDescriptorSet uniform_buffer_write = vkb::initializers::write_descriptor_set(model.descriptor_set_first_pass[i], + VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 2, &buffer_descriptor); std::vector write_descriptor_sets = { texture_input_write_0, @@ -995,7 +1078,7 @@ void MobileNerf::create_descriptor_sets_first_pass(Model &model) if (!use_deferred) { // Add in descriptor sets for MLP weights - weights_buffer_descriptor = create_descriptor(**model.weights_buffer_ref); + weights_buffer_descriptor = create_descriptor(*weights_buffers[model.model_index]); // Add in descriptor sets for MLP weights VkWriteDescriptorSet weights_buffer_write = vkb::initializers::write_descriptor_set(model.descriptor_set_first_pass[i], VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 3, &weights_buffer_descriptor); @@ -1013,9 +1096,27 @@ void MobileNerf::create_descriptor_sets_baseline() for (int i = 0; i < nerf_framebuffers.size(); i++) { VkDescriptorSetAllocateInfo descriptor_set_allocate_info = vkb::initializers::descriptor_set_allocate_info(descriptor_pool, &descriptor_set_layout_baseline, 1); - VK_CHECK(vkAllocateDescriptorSets(get_device().get_handle(), &descriptor_set_allocate_info, &descriptor_set_baseline[i])); + if (combo_mode) + { + uint32_t counts[1]; + counts[0] = static_cast(model_path.size()); + + VkDescriptorSetVariableDescriptorCountAllocateInfo set_counts = {}; + set_counts.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_VARIABLE_DESCRIPTOR_COUNT_ALLOCATE_INFO; + set_counts.descriptorSetCount = 1; + set_counts.pDescriptorCounts = counts; + + descriptor_set_allocate_info.pNext = &set_counts; + + VK_CHECK(vkAllocateDescriptorSets(get_device().get_handle(), &descriptor_set_allocate_info, &descriptor_set_baseline[i])); + } + else + { + VK_CHECK(vkAllocateDescriptorSets(get_device().get_handle(), &descriptor_set_allocate_info, &descriptor_set_baseline[i])); + } - std::array attachment_input_descriptors; + std::vector attachment_input_descriptors; + attachment_input_descriptors.resize(combo_mode ? 4 : 3); attachment_input_descriptors[0].sampler = VK_NULL_HANDLE; attachment_input_descriptors[0].imageView = frameAttachments[i].feature_0.view; @@ -1029,22 +1130,48 @@ void MobileNerf::create_descriptor_sets_baseline() attachment_input_descriptors[2].imageView = frameAttachments[i].feature_2.view; attachment_input_descriptors[2].imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL; - // TODO(tomatkinson): add in descriptor sets for MLP weights - VkDescriptorBufferInfo weights_buffer_descriptor = create_descriptor(**models[0].weights_buffer_ref); - VkWriteDescriptorSet texture_input_write_0 = vkb::initializers::write_descriptor_set(descriptor_set_baseline[i], VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, 0, &attachment_input_descriptors[0]); - VkWriteDescriptorSet texture_input_write_1 = vkb::initializers::write_descriptor_set(descriptor_set_baseline[i], VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, 1, &attachment_input_descriptors[1]); - VkWriteDescriptorSet texture_input_write_2 = vkb::initializers::write_descriptor_set(descriptor_set_baseline[i], VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, 2, &attachment_input_descriptors[2]); + VkWriteDescriptorSet texture_input_write_0 = vkb::initializers::write_descriptor_set(descriptor_set_baseline[i], VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, 0, &attachment_input_descriptors[0]); + VkWriteDescriptorSet texture_input_write_1 = vkb::initializers::write_descriptor_set(descriptor_set_baseline[i], VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, 1, &attachment_input_descriptors[1]); + VkWriteDescriptorSet texture_input_write_2 = vkb::initializers::write_descriptor_set(descriptor_set_baseline[i], VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, 2, &attachment_input_descriptors[2]); - // TODO(tomatkinson): add in descriptor sets for MLP weights - VkWriteDescriptorSet weights_buffer_write = vkb::initializers::write_descriptor_set(descriptor_set_baseline[i], VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 3, &weights_buffer_descriptor); // UBO + if (combo_mode) + { + attachment_input_descriptors[3].sampler = VK_NULL_HANDLE; + attachment_input_descriptors[3].imageView = frameAttachments[i].weights_idx.view; + attachment_input_descriptors[3].imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL; + VkWriteDescriptorSet texture_input_write_3 = vkb::initializers::write_descriptor_set(descriptor_set_baseline[i], VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, 3, &attachment_input_descriptors[3]); + + std::vector weights_buffer_descriptors; + weights_buffer_descriptors.reserve(mlp_weight_vector.size()); + for (auto &weight_buffer : weights_buffers) + { + weights_buffer_descriptors.emplace_back(create_descriptor(*weight_buffer)); + } + VkWriteDescriptorSet weights_buffer_write = vkb::initializers::write_descriptor_set( + descriptor_set_baseline[i], VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 4, weights_buffer_descriptors.data(), static_cast(weights_buffer_descriptors.size())); - std::vector write_descriptor_sets = { - texture_input_write_0, - texture_input_write_1, - texture_input_write_2, - weights_buffer_write}; + std::vector write_descriptor_sets = { + texture_input_write_0, + texture_input_write_1, + texture_input_write_2, + texture_input_write_3, + weights_buffer_write}; - vkUpdateDescriptorSets(get_device().get_handle(), static_cast(write_descriptor_sets.size()), write_descriptor_sets.data(), 0, VK_NULL_HANDLE); + vkUpdateDescriptorSets(get_device().get_handle(), static_cast(write_descriptor_sets.size()), write_descriptor_sets.data(), 0, VK_NULL_HANDLE); + } + else + { + VkDescriptorBufferInfo weights_buffer_descriptor = create_descriptor(*weights_buffers[models[0].model_index]); + VkWriteDescriptorSet weights_buffer_write = vkb::initializers::write_descriptor_set(descriptor_set_baseline[i], VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 3, &weights_buffer_descriptor); // UBO + + std::vector write_descriptor_sets = { + texture_input_write_0, + texture_input_write_1, + texture_input_write_2, + weights_buffer_write}; + + vkUpdateDescriptorSets(get_device().get_handle(), static_cast(write_descriptor_sets.size()), write_descriptor_sets.data(), 0, VK_NULL_HANDLE); + } } } @@ -1061,6 +1188,8 @@ void MobileNerf::prepare_pipelines() { blend_attachment_states.push_back(vkb::initializers::pipeline_color_blend_attachment_state(0xf, VK_FALSE)); blend_attachment_states.push_back(vkb::initializers::pipeline_color_blend_attachment_state(0xf, VK_FALSE)); + if (combo_mode) + blend_attachment_states.push_back(vkb::initializers::pipeline_color_blend_attachment_state(0xf, VK_FALSE)); } VkPipelineColorBlendStateCreateInfo color_blend_state = vkb::initializers::pipeline_color_blend_state_create_info(static_cast(blend_attachment_states.size()), blend_attachment_states.data()); @@ -1207,14 +1336,8 @@ void MobileNerf::create_uniforms() VMA_MEMORY_USAGE_CPU_TO_GPU); } - // Record a referce to vulkan buffer for each one of the models - for (Model &model : models) - { - model.uniform_buffer_ref = &uniform_buffers[model.model_index]; - model.weights_buffer_ref = &weights_buffers[model.model_index]; - } - update_uniform_buffers(); + update_weights_buffers(); } void MobileNerf::initialize_mlp_uniform_buffers(int model_index) @@ -1313,16 +1436,16 @@ void MobileNerf::initialize_mlp_uniform_buffers(int model_index) } // Each sub model will share the same mlp weights data - MLP_Weights *model_mlp = &mlp_weight_vector[model_index]; + MLP_Weights &model_mlp = mlp_weight_vector[model_index]; for (int ii = 0; ii < WEIGHTS_0_COUNT; ii++) { - model_mlp->data[ii] = weights_0_array[ii]; + model_mlp.data[ii] = weights_0_array[ii]; } for (int ii = 0; ii < WEIGHTS_1_COUNT; ii++) { - model_mlp->data[WEIGHTS_0_COUNT + ii] = weights_1_array[ii]; + model_mlp.data[WEIGHTS_0_COUNT + ii] = weights_1_array[ii]; } // We need to pad the layer 2's weights with zeros for every 3 weights to make it 16 bytes aligned @@ -1331,23 +1454,23 @@ void MobileNerf::initialize_mlp_uniform_buffers(int model_index) { if ((ii + 1) % 4 == 0) { - model_mlp->data[WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + ii] = 0.0f; + model_mlp.data[WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + ii] = 0.0f; } else { - model_mlp->data[WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + ii] = weights_2_array[raw_weight_cnt++]; + model_mlp.data[WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + ii] = weights_2_array[raw_weight_cnt++]; } } for (int ii = 0; ii < BIAS_0_COUNT; ii++) { - model_mlp->data[WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + ii] = bias_0_array[ii]; + model_mlp.data[WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + ii] = bias_0_array[ii]; } for (int ii = 0; ii < BIAS_1_COUNT; ii++) { - model_mlp->data[WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + - BIAS_0_COUNT + ii] = bias_1_array[ii]; + model_mlp.data[WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + + BIAS_0_COUNT + ii] = bias_1_array[ii]; } // We need to pad the layer 2's bias with zeros for every 3 weights to make it 16 bytes aligned @@ -1355,13 +1478,13 @@ void MobileNerf::initialize_mlp_uniform_buffers(int model_index) { if ((ii + 1) % 4 == 0) { - model_mlp->data[WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + - BIAS_0_COUNT + BIAS_1_COUNT + ii] = 0.0f; + model_mlp.data[WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + + BIAS_0_COUNT + BIAS_1_COUNT + ii] = 0.0f; } else { - model_mlp->data[WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + - BIAS_0_COUNT + BIAS_1_COUNT + ii] = bias_2_array[ii]; + model_mlp.data[WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + + BIAS_0_COUNT + BIAS_1_COUNT + ii] = bias_2_array[ii]; } } @@ -1387,14 +1510,19 @@ void MobileNerf::update_uniform_buffers() global_uniform.img_dim = glm::vec2(width, height); global_uniform.tan_half_fov = tan_half_fov; - // Note that this is a hard-coded scene setting for the lego_combo - glm::mat4x4 model_translation[4] = {glm::translate(glm::vec3(0.5, 0.75, 0)), glm::translate(glm::vec3(0.5, 0.25, 0)), - glm::translate(glm::vec3(0, -0.25, 0.5)), glm::translate(glm::vec3(0, -0.75, -0.5))}; - for (int i = 0; i < model_path.size(); i++) { - global_uniform.model = combo_mode ? model_translation[i] : glm::translate(glm::vec3(0.0f)); + // Note that this is a hard-coded scene setting for the lego_combo + global_uniform.model = combo_mode ? combo_model_transform[i] : glm::translate(glm::vec3(0.0f)); uniform_buffers[i]->update(&global_uniform, sizeof(global_uniform)); + } +} + +void MobileNerf::update_weights_buffers() +{ + // No need to be updated for every frames + for (int i = 0; i < model_path.size(); i++) + { weights_buffers[i]->update(&(mlp_weight_vector[i].data[0]), sizeof(MLP_Weights)); } } @@ -1478,7 +1606,8 @@ void MobileNerf::create_texture(int model_index, int sub_model_index, int models void MobileNerf::create_texture_helper(std::string const &texturePath, Texture &texture_input) { - texture_input = load_texture(texturePath, vkb::sg::Image::Color); + // Feature textures are in linear space instead of sRGB space + texture_input = load_texture(texturePath, vkb::sg::Image::Other); vkDestroySampler(get_device().get_handle(), texture_input.sampler, nullptr); // Calculate valid filter @@ -1555,76 +1684,121 @@ void MobileNerf::update_render_pass_nerf_forward() void MobileNerf::update_render_pass_nerf_baseline() { - std::array attachments = {}; - // Color attachment 1 - attachments[0].format = feature_map_format; - attachments[0].samples = VK_SAMPLE_COUNT_1_BIT; - attachments[0].loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR; - attachments[0].storeOp = VK_ATTACHMENT_STORE_OP_DONT_CARE; - attachments[0].stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE; - attachments[0].stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE; - attachments[0].initialLayout = VK_IMAGE_LAYOUT_UNDEFINED; - attachments[0].finalLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL; - // Color attachment 2 - attachments[1].format = feature_map_format; - attachments[1].samples = VK_SAMPLE_COUNT_1_BIT; - attachments[1].loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR; - attachments[1].storeOp = VK_ATTACHMENT_STORE_OP_DONT_CARE; - attachments[1].stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE; - attachments[1].stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE; - attachments[1].initialLayout = VK_IMAGE_LAYOUT_UNDEFINED; - attachments[1].finalLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL; - // Color attachment 3 - attachments[2].format = VK_FORMAT_R16G16B16A16_SFLOAT; - attachments[2].samples = VK_SAMPLE_COUNT_1_BIT; - attachments[2].loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR; - attachments[2].storeOp = VK_ATTACHMENT_STORE_OP_DONT_CARE; - attachments[2].stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE; - attachments[2].stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE; - attachments[2].initialLayout = VK_IMAGE_LAYOUT_UNDEFINED; - attachments[2].finalLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL; + unsigned int attachment_idx = 0; + + // Color attachment 0 - feature 0 G-buffer + color_attach_0_idx = attachment_idx++; + VkAttachmentDescription color_description_0 = {}; + color_description_0.format = feature_map_format; + color_description_0.samples = VK_SAMPLE_COUNT_1_BIT; + color_description_0.loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR; + color_description_0.storeOp = VK_ATTACHMENT_STORE_OP_DONT_CARE; + color_description_0.stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE; + color_description_0.stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE; + color_description_0.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED; + color_description_0.finalLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL; + // Color attachment 1 - feature 1 G-buffer + color_attach_1_idx = attachment_idx++; + VkAttachmentDescription color_description_1 = {}; + color_description_1.format = feature_map_format; + color_description_1.samples = VK_SAMPLE_COUNT_1_BIT; + color_description_1.loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR; + color_description_1.storeOp = VK_ATTACHMENT_STORE_OP_DONT_CARE; + color_description_1.stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE; + color_description_1.stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE; + color_description_1.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED; + color_description_1.finalLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL; + // Color attachment 2 - ray direction G-buffer + color_attach_2_idx = attachment_idx++; + VkAttachmentDescription color_description_2 = {}; + color_description_2.format = VK_FORMAT_R16G16B16A16_SFLOAT; + color_description_2.samples = VK_SAMPLE_COUNT_1_BIT; + color_description_2.loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR; + color_description_2.storeOp = VK_ATTACHMENT_STORE_OP_DONT_CARE; + color_description_2.stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE; + color_description_2.stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE; + color_description_2.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED; + color_description_2.finalLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL; + // Color attachment 3 - weight index G-buffer + VkAttachmentDescription color_description_3 = {}; + color_attach_3_idx = 3; + if (combo_mode) + { + color_attach_3_idx = attachment_idx++; + color_description_3.format = VK_FORMAT_R8_UINT; + color_description_3.samples = VK_SAMPLE_COUNT_1_BIT; + color_description_3.loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR; + color_description_3.storeOp = VK_ATTACHMENT_STORE_OP_DONT_CARE; + color_description_3.stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE; + color_description_3.stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE; + color_description_3.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED; + color_description_3.finalLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL; + } // Depth attachment - attachments[3].format = depth_format; - attachments[3].samples = VK_SAMPLE_COUNT_1_BIT; - attachments[3].loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR; - attachments[3].storeOp = VK_ATTACHMENT_STORE_OP_DONT_CARE; - attachments[3].stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE; - attachments[3].stencilStoreOp = VK_ATTACHMENT_STORE_OP_STORE; - attachments[3].initialLayout = VK_IMAGE_LAYOUT_UNDEFINED; - attachments[3].finalLayout = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL; + depth_attach_idx = attachment_idx++; + VkAttachmentDescription depth_description = {}; + depth_description.format = depth_format; + depth_description.samples = VK_SAMPLE_COUNT_1_BIT; + depth_description.loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR; + depth_description.storeOp = VK_ATTACHMENT_STORE_OP_DONT_CARE; + depth_description.stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE; + depth_description.stencilStoreOp = VK_ATTACHMENT_STORE_OP_STORE; + depth_description.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED; + depth_description.finalLayout = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL; // Swapchain attachment - attachments[4].format = get_render_context().get_swapchain().get_format(); - attachments[4].samples = VK_SAMPLE_COUNT_1_BIT; - attachments[4].loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR; - attachments[4].storeOp = VK_ATTACHMENT_STORE_OP_STORE; - attachments[4].stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE; - attachments[4].stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE; - attachments[4].initialLayout = VK_IMAGE_LAYOUT_UNDEFINED; - attachments[4].finalLayout = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR; + swapchain_attach_idx = attachment_idx++; + VkAttachmentDescription swapchain_description = {}; + swapchain_description.format = get_render_context().get_swapchain().get_format(); + swapchain_description.samples = VK_SAMPLE_COUNT_1_BIT; + swapchain_description.loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR; + swapchain_description.storeOp = VK_ATTACHMENT_STORE_OP_STORE; + swapchain_description.stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE; + swapchain_description.stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE; + swapchain_description.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED; + swapchain_description.finalLayout = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR; + + std::vector attachments; + attachments.push_back(color_description_0); + attachments.push_back(color_description_1); + attachments.push_back(color_description_2); + if (combo_mode) + attachments.push_back(color_description_3); + attachments.push_back(depth_description); + attachments.push_back(swapchain_description); VkAttachmentReference color_reference_0 = {}; - color_reference_0.attachment = 0; + color_reference_0.attachment = color_attach_0_idx; color_reference_0.layout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL; VkAttachmentReference color_reference_1 = {}; - color_reference_1.attachment = 1; + color_reference_1.attachment = color_attach_1_idx; color_reference_1.layout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL; VkAttachmentReference color_reference_2 = {}; - color_reference_2.attachment = 2; + color_reference_2.attachment = color_attach_2_idx; color_reference_2.layout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL; + VkAttachmentReference color_reference_3 = {}; + color_reference_3.attachment = color_attach_3_idx; + color_reference_3.layout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL; + VkAttachmentReference depth_reference = {}; - depth_reference.attachment = 3; + depth_reference.attachment = depth_attach_idx; depth_reference.layout = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL; + VkAttachmentReference swapchain_reference = {}; + swapchain_reference.attachment = swapchain_attach_idx; + swapchain_reference.layout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL; + std::array subpassDescriptions{}; - VkAttachmentReference color_references_feature_maps[3] = {color_reference_0, color_reference_1, color_reference_2}; + std::vector color_references_feature_maps = {color_reference_0, color_reference_1, color_reference_2}; + if (combo_mode) + color_references_feature_maps.push_back(color_reference_3); subpassDescriptions[0].pipelineBindPoint = VK_PIPELINE_BIND_POINT_GRAPHICS; - subpassDescriptions[0].colorAttachmentCount = 3; - subpassDescriptions[0].pColorAttachments = color_references_feature_maps; + subpassDescriptions[0].colorAttachmentCount = static_cast(color_references_feature_maps.size()); + subpassDescriptions[0].pColorAttachments = color_references_feature_maps.data(); subpassDescriptions[0].pDepthStencilAttachment = &depth_reference; subpassDescriptions[0].inputAttachmentCount = 0; subpassDescriptions[0].pInputAttachments = nullptr; @@ -1632,28 +1806,27 @@ void MobileNerf::update_render_pass_nerf_baseline() subpassDescriptions[0].pPreserveAttachments = nullptr; subpassDescriptions[0].pResolveAttachments = nullptr; - VkAttachmentReference swapchain_reference = {}; - swapchain_reference.attachment = 4; - swapchain_reference.layout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL; - // Color attachments written to in first sub pass will be used as input attachments to be read in the fragment shader - VkAttachmentReference inputReferences[3]; - inputReferences[0] = {0, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL}; - inputReferences[1] = {1, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL}; - inputReferences[2] = {2, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL}; + std::vector inputReferences = { + {color_attach_0_idx, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL}, // Color attachment 0 - feature 0 G-buffer + {color_attach_1_idx, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL}, // Color attachment 1 - feature 1 G-buffer + {color_attach_2_idx, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL} // Color attachment 2 - ray direction G-buffer + }; + if (combo_mode) + inputReferences.push_back({color_attach_3_idx, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL}); // Color attachment 3 - weight index G-buffer subpassDescriptions[1].pipelineBindPoint = VK_PIPELINE_BIND_POINT_GRAPHICS; subpassDescriptions[1].colorAttachmentCount = 1; subpassDescriptions[1].pColorAttachments = &swapchain_reference; subpassDescriptions[1].pDepthStencilAttachment = nullptr; - subpassDescriptions[1].inputAttachmentCount = 3; - subpassDescriptions[1].pInputAttachments = inputReferences; + subpassDescriptions[1].inputAttachmentCount = static_cast(inputReferences.size()); + subpassDescriptions[1].pInputAttachments = inputReferences.data(); subpassDescriptions[1].preserveAttachmentCount = 0; subpassDescriptions[1].pPreserveAttachments = nullptr; subpassDescriptions[1].pResolveAttachments = nullptr; // Subpass dependencies for layout transitions - std::array dependencies; + std::array dependencies{}; dependencies[0].srcSubpass = VK_SUBPASS_EXTERNAL; dependencies[0].dstSubpass = 0; @@ -1683,7 +1856,7 @@ void MobileNerf::update_render_pass_nerf_baseline() render_pass_create_info.sType = VK_STRUCTURE_TYPE_RENDER_PASS_CREATE_INFO; render_pass_create_info.attachmentCount = static_cast(attachments.size()); render_pass_create_info.pAttachments = attachments.data(); - render_pass_create_info.subpassCount = subpassDescriptions.size(); + render_pass_create_info.subpassCount = static_cast(subpassDescriptions.size()); render_pass_create_info.pSubpasses = subpassDescriptions.data(); render_pass_create_info.dependencyCount = static_cast(dependencies.size()); render_pass_create_info.pDependencies = dependencies.data(); diff --git a/samples/general/mobile_nerf/mobile_nerf.h b/samples/general/mobile_nerf/mobile_nerf.h index 63c9c01dd..6a4062443 100644 --- a/samples/general/mobile_nerf/mobile_nerf.h +++ b/samples/general/mobile_nerf/mobile_nerf.h @@ -61,6 +61,11 @@ class MobileNerf : public ApiVulkanSample alignas(4) float tan_half_fov; } global_uniform; + struct PushConstants + { + unsigned int weight_idx; + } push_constants; + #define WEIGHTS_0_COUNT (176) #define WEIGHTS_1_COUNT (256) // The third layer weights' size is changed from 48 to 64 to make sure a 16 bytes alignement @@ -78,13 +83,8 @@ class MobileNerf : public ApiVulkanSample struct Vertex { - alignas(4) glm::vec3 position; - alignas(4) glm::vec2 tex_coord; - - bool operator==(const Vertex &other) const - { - return position == other.position && tex_coord == other.tex_coord; - } + glm::vec3 position; + glm::vec2 tex_coord; }; struct InstancingInfo @@ -130,10 +130,6 @@ class MobileNerf : public ApiVulkanSample // Deferred mode will only have one set of descriptor per model std::vector descriptor_set_first_pass{VK_NULL_HANDLE}; - // Stores references to each models mlp weights and uniform buffers - std::unique_ptr *weights_buffer_ref; - std::unique_ptr *uniform_buffer_ref; - int sub_model_num; int model_index; }; @@ -147,7 +143,7 @@ class MobileNerf : public ApiVulkanSample // Uniform buffer for each model std::vector> uniform_buffers; - // buffer to store instance data + // Buffer to store instance data std::unique_ptr instance_buffer{nullptr}; // Common @@ -160,6 +156,7 @@ class MobileNerf : public ApiVulkanSample void load_scene(int model_index, int sub_model_index, int models_entry); void initialize_mlp_uniform_buffers(int model_index); void update_uniform_buffers(); + void update_weights_buffers(); void create_texture(int model_index, int sub_model_index, int models_entry); void create_texture_helper(std::string const &texturePath, Texture &texture); @@ -181,9 +178,16 @@ class MobileNerf : public ApiVulkanSample void create_descriptor_sets_first_pass(Model &model); void prepare_pipelines(); + unsigned int color_attach_0_idx; + unsigned int color_attach_1_idx; + unsigned int color_attach_2_idx; + unsigned int color_attach_3_idx; + unsigned int depth_attach_idx; + unsigned int swapchain_attach_idx; + struct Attachments_baseline { - FrameBufferAttachment feature_0, feature_1, feature_2; + FrameBufferAttachment feature_0, feature_1, feature_2, weights_idx; }; std::vector frameAttachments; @@ -209,20 +213,25 @@ class MobileNerf : public ApiVulkanSample // For loading nerf assets map json asset_map; std::vector model_path; - bool combo_mode; std::vector using_original_nerf_models; - bool use_deferred; - bool do_rotation; + bool combo_mode = false; + bool use_deferred = false; + bool do_rotation = false; glm::vec3 camera_pos = glm::vec3(-2.2f, 2.2f, 2.2f); + // Currently combo mode translation are hard-coded + glm::mat4x4 combo_model_transform[4] = { + glm::translate(glm::vec3(0.5, 0.75, 0)), glm::translate(glm::vec3(0.5, 0.25, 0)), + glm::translate(glm::vec3(0, -0.25, 0.5)), glm::translate(glm::vec3(0, -0.75, -0.5))}; + // For instancing InstancingInfo instancing_info; // Viewport Setting - float fov = 60.0f; - uint32_t view_port_width; - uint32_t view_port_height; + float fov = 60.0f; + uint32_t view_port_width = width; + uint32_t view_port_height = height; bool use_native_screen_size = false; }; diff --git a/shaders/mobile_nerf/merged.frag b/shaders/mobile_nerf/merged.frag index f914b945a..6a9ffe877 100644 --- a/shaders/mobile_nerf/merged.frag +++ b/shaders/mobile_nerf/merged.frag @@ -17,7 +17,7 @@ * ------------------------------------------------------------------------ * * THIS IS A MODIFIED VERSION OF THE ORIGINAL FILE - * + * * The original file, along with the original Apache-2.0 LICENSE can be found at: * https://github.com/google-research/jax3d/tree/main/jax3d/projects/mobilenerf * @@ -49,19 +49,19 @@ precision highp float; #define BIAS_2_COUNT (4) layout(binding = 3) uniform mlp_weights { - vec4 data[(WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + + vec4 data[(WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + BIAS_0_COUNT + BIAS_1_COUNT + BIAS_2_COUNT)/4]; // Array of floats } weights; -vec3 evaluateNetwork( vec4 f0, vec4 f1, vec4 viewdir) +vec3 evaluateNetwork( vec4 f0, vec4 f1, vec4 viewdir) { - + vec3 res; int bias_0_ind = WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT; vec4 intermediate_one[4] = vec4[]( - weights.data[bias_0_ind/4], + weights.data[bias_0_ind/4], weights.data[bias_0_ind/4 + 1], weights.data[bias_0_ind/4 + 2], weights.data[bias_0_ind/4 + 3] @@ -72,7 +72,7 @@ vec3 evaluateNetwork( vec4 f0, vec4 f1, vec4 viewdir) intermediate_one[ 0] += (multiplier) * weights.data[ weightFirstInd/4]; \ intermediate_one[ 1] += (multiplier) * weights.data[ weightFirstInd/4 + 1]; \ intermediate_one[ 2] += (multiplier) * weights.data[ weightFirstInd/4 + 2]; \ - intermediate_one[ 3] += (multiplier) * weights.data[ weightFirstInd/4 + 3]; + intermediate_one[ 3] += (multiplier) * weights.data[ weightFirstInd/4 + 3]; APPLY_WEIGHTS_0( f0.r, 0) APPLY_WEIGHTS_0( f0.g, 16) @@ -87,10 +87,10 @@ vec3 evaluateNetwork( vec4 f0, vec4 f1, vec4 viewdir) APPLY_WEIGHTS_0( -viewdir.b, 144) APPLY_WEIGHTS_0( viewdir.g, 160) - int bias_1_ind = WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + + int bias_1_ind = WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + BIAS_0_COUNT; vec4 intermediate_two[4] = vec4[]( - weights.data[bias_1_ind/4], + weights.data[bias_1_ind/4], weights.data[bias_1_ind/4 + 1], weights.data[bias_1_ind/4 + 2], weights.data[bias_1_ind/4 + 3] @@ -104,7 +104,7 @@ vec3 evaluateNetwork( vec4 f0, vec4 f1, vec4 viewdir) intermediate_two[ 2] += intermediate * weights.data[ WEIGHTS_0_COUNT/4 + oneInd * 4 + 2]; \ intermediate_two[ 3] += intermediate * weights.data[ WEIGHTS_0_COUNT/4 + oneInd * 4 + 3]; \ } - + APPLY_WEIGHTS_1( intermediate_one[0].r, 0) APPLY_WEIGHTS_1( intermediate_one[0].g, 1) APPLY_WEIGHTS_1( intermediate_one[0].b, 2) @@ -122,7 +122,7 @@ vec3 evaluateNetwork( vec4 f0, vec4 f1, vec4 viewdir) APPLY_WEIGHTS_1( intermediate_one[3].b, 14) APPLY_WEIGHTS_1( intermediate_one[3].a, 15) - int bias_2_ind = WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + + int bias_2_ind = WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + BIAS_0_COUNT + BIAS_1_COUNT; vec4 result = weights.data[bias_2_ind/4]; @@ -130,7 +130,7 @@ vec3 evaluateNetwork( vec4 f0, vec4 f1, vec4 viewdir) if(intermediate > 0.0f){ \ result += intermediate * weights.data[ WEIGHTS_0_COUNT/4 + WEIGHTS_1_COUNT/4 + oneInd]; \ } - + APPLY_WEIGHTS_2(intermediate_two[0].r, 0) APPLY_WEIGHTS_2(intermediate_two[0].g, 1) APPLY_WEIGHTS_2(intermediate_two[0].b, 2) @@ -147,9 +147,9 @@ vec3 evaluateNetwork( vec4 f0, vec4 f1, vec4 viewdir) APPLY_WEIGHTS_2(intermediate_two[3].g,13) APPLY_WEIGHTS_2(intermediate_two[3].b,14) APPLY_WEIGHTS_2(intermediate_two[3].a,15) - + result = 1.0 / (1.0 + exp(-result)); - return vec3(result * viewdir.a+(1.0-viewdir.a)); + return vec3(result * viewdir.a+(1.0-viewdir.a)); } ////////////////////////////////////////////////////////////// @@ -164,7 +164,7 @@ float Convert_sRGB_ToLinear(float value) : pow((value + 0.055) / 1.055, 2.4); } -vec3 Convert_sRGB_ToLinear(vec3 value) +vec3 Convert_sRGB_ToLinear(vec3 value) { return vec3(Convert_sRGB_ToLinear(value.x), Convert_sRGB_ToLinear(value.y), Convert_sRGB_ToLinear(value.z)); } @@ -181,20 +181,15 @@ void main(void) vec4 pixel_1 = texture(textureInput_1, flipped); vec4 feature_0 = pixel_0; vec4 feature_1 = pixel_1; - + vec4 rayDirection = vec4(normalize(rayDirectionIn), 1.0f); - - // For debugging only - // o_color_0 = vec4( texCoord_frag.x, 0.0f, 0.0f, 1.0f); // deal with iphone feature_0.a = feature_0.a*2.0-1.0; feature_1.a = feature_1.a*2.0-1.0; rayDirection.a = rayDirection.a*2.0-1.0; - - // Original + // Original o_color.rgb = Convert_sRGB_ToLinear(evaluateNetwork(feature_0,feature_1,rayDirection)); - //o_color.rgb = feature_0.rgb; o_color.a = 1.0; } diff --git a/shaders/mobile_nerf/merged_morpheus.frag b/shaders/mobile_nerf/merged_morpheus.frag index 2e93b15ea..bd3c668a6 100644 --- a/shaders/mobile_nerf/merged_morpheus.frag +++ b/shaders/mobile_nerf/merged_morpheus.frag @@ -17,7 +17,7 @@ * ------------------------------------------------------------------------ * * THIS IS A MODIFIED VERSION OF THE ORIGINAL FILE - * + * * The original file, along with the original Apache-2.0 LICENSE can be found at: * https://github.com/google-research/jax3d/tree/main/jax3d/projects/mobilenerf * @@ -49,19 +49,19 @@ precision highp float; #define BIAS_2_COUNT (4) layout(binding = 3) uniform mlp_weights { - vec4 data[(WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + + vec4 data[(WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + BIAS_0_COUNT + BIAS_1_COUNT + BIAS_2_COUNT)/4]; // Array of floats } weights; -vec3 evaluateNetwork( vec4 f0, vec4 f1, vec4 viewdir) +vec3 evaluateNetwork( vec4 f0, vec4 f1, vec4 viewdir) { - + vec3 res; int bias_0_ind = WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT; vec4 intermediate_one[4] = vec4[]( - weights.data[bias_0_ind/4], + weights.data[bias_0_ind/4], weights.data[bias_0_ind/4 + 1], weights.data[bias_0_ind/4 + 2], weights.data[bias_0_ind/4 + 3] @@ -72,7 +72,7 @@ vec3 evaluateNetwork( vec4 f0, vec4 f1, vec4 viewdir) intermediate_one[ 0] += (multiplier) * weights.data[ weightFirstInd/4]; \ intermediate_one[ 1] += (multiplier) * weights.data[ weightFirstInd/4 + 1]; \ intermediate_one[ 2] += (multiplier) * weights.data[ weightFirstInd/4 + 2]; \ - intermediate_one[ 3] += (multiplier) * weights.data[ weightFirstInd/4 + 3]; + intermediate_one[ 3] += (multiplier) * weights.data[ weightFirstInd/4 + 3]; APPLY_WEIGHTS_0( f0.r, 0) APPLY_WEIGHTS_0( f0.g, 16) @@ -87,10 +87,10 @@ vec3 evaluateNetwork( vec4 f0, vec4 f1, vec4 viewdir) APPLY_WEIGHTS_0( (-viewdir.b + 1.0 )/2, 144) APPLY_WEIGHTS_0( (viewdir.g + 1.0 )/2, 160) - int bias_1_ind = WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + + int bias_1_ind = WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + BIAS_0_COUNT; vec4 intermediate_two[4] = vec4[]( - weights.data[bias_1_ind/4], + weights.data[bias_1_ind/4], weights.data[bias_1_ind/4 + 1], weights.data[bias_1_ind/4 + 2], weights.data[bias_1_ind/4 + 3] @@ -104,7 +104,7 @@ vec3 evaluateNetwork( vec4 f0, vec4 f1, vec4 viewdir) intermediate_two[ 2] += intermediate * weights.data[ WEIGHTS_0_COUNT/4 + oneInd * 4 + 2]; \ intermediate_two[ 3] += intermediate * weights.data[ WEIGHTS_0_COUNT/4 + oneInd * 4 + 3]; \ } - + APPLY_WEIGHTS_1( intermediate_one[0].r, 0) APPLY_WEIGHTS_1( intermediate_one[0].g, 1) APPLY_WEIGHTS_1( intermediate_one[0].b, 2) @@ -122,7 +122,7 @@ vec3 evaluateNetwork( vec4 f0, vec4 f1, vec4 viewdir) APPLY_WEIGHTS_1( intermediate_one[3].b, 14) APPLY_WEIGHTS_1( intermediate_one[3].a, 15) - int bias_2_ind = WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + + int bias_2_ind = WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + BIAS_0_COUNT + BIAS_1_COUNT; vec4 result = weights.data[bias_2_ind/4]; @@ -130,7 +130,7 @@ vec3 evaluateNetwork( vec4 f0, vec4 f1, vec4 viewdir) if(intermediate > 0.0f){ \ result += intermediate * weights.data[ WEIGHTS_0_COUNT/4 + WEIGHTS_1_COUNT/4 + oneInd]; \ } - + APPLY_WEIGHTS_2(intermediate_two[0].r, 0) APPLY_WEIGHTS_2(intermediate_two[0].g, 1) APPLY_WEIGHTS_2(intermediate_two[0].b, 2) @@ -147,9 +147,9 @@ vec3 evaluateNetwork( vec4 f0, vec4 f1, vec4 viewdir) APPLY_WEIGHTS_2(intermediate_two[3].g,13) APPLY_WEIGHTS_2(intermediate_two[3].b,14) APPLY_WEIGHTS_2(intermediate_two[3].a,15) - + result = 1.0 / (1.0 + exp(-result)); - return vec3(result * viewdir.a+(1.0-viewdir.a)); + return vec3(result * viewdir.a+(1.0-viewdir.a)); } ////////////////////////////////////////////////////////////// @@ -164,7 +164,7 @@ float Convert_sRGB_ToLinear(float value) : pow((value + 0.055) / 1.055, 2.4); } -vec3 Convert_sRGB_ToLinear(vec3 value) +vec3 Convert_sRGB_ToLinear(vec3 value) { return vec3(Convert_sRGB_ToLinear(value.x), Convert_sRGB_ToLinear(value.y), Convert_sRGB_ToLinear(value.z)); } @@ -181,19 +181,15 @@ void main(void) vec4 pixel_1 = texture(textureInput_1, flipped); vec4 feature_0 = pixel_0; vec4 feature_1 = pixel_1; - + vec4 rayDirection = vec4(normalize(rayDirectionIn), 1.0f); - - // For debugging only - // o_color_0 = vec4( texCoord_frag.x, 0.0f, 0.0f, 1.0f); // deal with iphone feature_0.a = feature_0.a*2.0-1.0; feature_1.a = feature_1.a*2.0-1.0; rayDirection.a = rayDirection.a*2.0-1.0; - + // Original o_color.rgb = Convert_sRGB_ToLinear(evaluateNetwork(feature_0,feature_1,rayDirection)); - // o_color.rgb = feature_0.rgb; o_color.a = 1.0; } diff --git a/shaders/mobile_nerf/mlp.frag b/shaders/mobile_nerf/mlp.frag index 801154e66..23a88705a 100644 --- a/shaders/mobile_nerf/mlp.frag +++ b/shaders/mobile_nerf/mlp.frag @@ -17,7 +17,7 @@ * ------------------------------------------------------------------------ * * THIS IS A MODIFIED VERSION OF THE ORIGINAL FILE - * + * * The original file, along with the original Apache-2.0 LICENSE can be found at: * https://github.com/google-research/jax3d/tree/main/jax3d/projects/mobilenerf * @@ -47,52 +47,27 @@ precision highp float; #define BIAS_2_COUNT (4) layout(binding = 3) uniform mlp_weights { - vec4 data[(WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + + vec4 data[(WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + BIAS_0_COUNT + BIAS_1_COUNT + BIAS_2_COUNT)/4]; // Array of floats } weights; vec3 evaluateNetwork( vec4 f0, vec4 f1, vec4 viewdir) { - + vec3 res; int bias_0_ind = WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT; vec4 intermediate_one[4] = vec4[]( - weights.data[bias_0_ind/4], + weights.data[bias_0_ind/4], weights.data[bias_0_ind/4 + 1], weights.data[bias_0_ind/4 + 2], weights.data[bias_0_ind/4 + 3] ); - #if 0 - for (int j = 0; j < 11; ++j) { - float input_value = 0.0; - if (j < 4) { - input_value = - (j == 0) ? f0.r : ( - (j == 1) ? f0.g : ( - (j == 2) ? f0.b : f0.a)); - } else if (j < 8) { - input_value = - (j == 4) ? f1.r : ( - (j == 5) ? f1.g : ( - (j == 6) ? f1.b : f1.a)); - } else { - input_value = - (j == 8) ? viewdir.r : ( - (j == 9) ? -viewdir.b : viewdir.g); //switch y-z axes - } - for (int i = 0; i < 16; ++i) { - intermediate_one[i] += input_value * weights0[ j*16 + i]; - //texelFetch(weightsZero, ivec2(j, i), 0).x; - } - } - #endif - #define APPLY_WEIGHTS_0(multiplier, weightFirstInd) \ intermediate_one[ 0] += (multiplier) * weights.data[ weightFirstInd/4]; \ intermediate_one[ 1] += (multiplier) * weights.data[ weightFirstInd/4 + 1]; \ intermediate_one[ 2] += (multiplier) * weights.data[ weightFirstInd/4 + 2]; \ - intermediate_one[ 3] += (multiplier) * weights.data[ weightFirstInd/4 + 3]; + intermediate_one[ 3] += (multiplier) * weights.data[ weightFirstInd/4 + 3]; APPLY_WEIGHTS_0( f0.r, 0) APPLY_WEIGHTS_0( f0.g, 16) @@ -107,27 +82,15 @@ vec3 evaluateNetwork( vec4 f0, vec4 f1, vec4 viewdir) { APPLY_WEIGHTS_0( -viewdir.b, 144) APPLY_WEIGHTS_0( viewdir.g, 160) - int bias_1_ind = WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + + int bias_1_ind = WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + BIAS_0_COUNT; vec4 intermediate_two[4] = vec4[]( - weights.data[bias_1_ind/4], + weights.data[bias_1_ind/4], weights.data[bias_1_ind/4 + 1], weights.data[bias_1_ind/4 + 2], weights.data[bias_1_ind/4 + 3] ); - #if 0 - for (int j = 0; j < 16; ++j) { - if (intermediate_one[j] <= 0.0) { - continue; - } - for (int i = 0; i < 16; ++i) { - intermediate_two[i] += intermediate_one[j] * weights1[j*16 + i]; - //texelFetch(weightsOne, ivec2(j, i), 0).x; - } - } - #endif - #define APPLY_WEIGHTS_1(intermediate, oneInd) \ if(intermediate > 0.0f){ \ intermediate_two[ 0] += intermediate * weights.data[ WEIGHTS_0_COUNT/4 + oneInd * 4 + 0]; \ @@ -135,7 +98,7 @@ vec3 evaluateNetwork( vec4 f0, vec4 f1, vec4 viewdir) { intermediate_two[ 2] += intermediate * weights.data[ WEIGHTS_0_COUNT/4 + oneInd * 4 + 2]; \ intermediate_two[ 3] += intermediate * weights.data[ WEIGHTS_0_COUNT/4 + oneInd * 4 + 3]; \ } - + APPLY_WEIGHTS_1( intermediate_one[0].r, 0) APPLY_WEIGHTS_1( intermediate_one[0].g, 1) APPLY_WEIGHTS_1( intermediate_one[0].b, 2) @@ -153,27 +116,15 @@ vec3 evaluateNetwork( vec4 f0, vec4 f1, vec4 viewdir) { APPLY_WEIGHTS_1( intermediate_one[3].b, 14) APPLY_WEIGHTS_1( intermediate_one[3].a, 15) - int bias_2_ind = WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + + int bias_2_ind = WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + BIAS_0_COUNT + BIAS_1_COUNT; vec4 result = weights.data[bias_2_ind/4]; - #if 0 - for (int j = 0; j < 16; ++j) { - if (intermediate_two[j] <= 0.0) { - continue; - } - for (int i = 0; i < 3; ++i) { - result[i] += intermediate_two[j] * weights2[j*3 + i]; - //texelFetch(weightsTwo, ivec2(j, i), 0).x; - } - } - #endif - #define APPLY_WEIGHTS_2(intermediate, oneInd) \ if(intermediate > 0.0f){ \ result += intermediate * weights.data[ WEIGHTS_0_COUNT/4 + WEIGHTS_1_COUNT/4 + oneInd]; \ } - + APPLY_WEIGHTS_2(intermediate_two[0].r, 0) APPLY_WEIGHTS_2(intermediate_two[0].g, 1) APPLY_WEIGHTS_2(intermediate_two[0].b, 2) @@ -190,9 +141,9 @@ vec3 evaluateNetwork( vec4 f0, vec4 f1, vec4 viewdir) { APPLY_WEIGHTS_2(intermediate_two[3].g,13) APPLY_WEIGHTS_2(intermediate_two[3].b,14) APPLY_WEIGHTS_2(intermediate_two[3].a,15) - + result = 1.0 / (1.0 + exp(-result)); - return vec3(result * viewdir.a+(1.0-viewdir.a)); + return vec3(result * viewdir.a+(1.0-viewdir.a)); } @@ -208,7 +159,7 @@ float Convert_sRGB_ToLinear(float value) : pow((value + 0.055) / 1.055, 2.4); } -vec3 Convert_sRGB_ToLinear(vec3 value) +vec3 Convert_sRGB_ToLinear(vec3 value) { return vec3(Convert_sRGB_ToLinear(value.x), Convert_sRGB_ToLinear(value.y), Convert_sRGB_ToLinear(value.z)); } @@ -224,27 +175,13 @@ void main(void) vec4 rayDirection = subpassLoad(rayDirectionIn).rgba; if (rayDirection.a < 0.6) discard; - + //deal with iphone feature_0.a = feature_0.a*2.0-1.0; feature_1.a = feature_1.a*2.0-1.0; rayDirection.a = rayDirection.a*2.0-1.0; - - // Original + // Original o_color.rgb = Convert_sRGB_ToLinear(evaluateNetwork(feature_0,feature_1,rayDirection)); - //o_color.rgb = feature_0.rgb; o_color.a = 1.0; - - - // --------------------------------------------------------------------------- - // Debugging - - //o_color.rgb = (rayDirection.rgb)*0.5 + vec3(0.5,0.5,0.5); - //o_color.rgb = rayDirection.rgb; - - //o_color.rgba = feature_1; - //o_color.a = 1.0f; - //o_color.rgba = feature_1; - } diff --git a/shaders/mobile_nerf/mlp_combo.frag b/shaders/mobile_nerf/mlp_combo.frag new file mode 100644 index 000000000..101358bd3 --- /dev/null +++ b/shaders/mobile_nerf/mlp_combo.frag @@ -0,0 +1,189 @@ +/* Copyright (c) 2024, Qualcomm Innovation Center, Inc. All rights reserved. + * + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 the "License"; + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + * ------------------------------------------------------------------------ + * + * THIS IS A MODIFIED VERSION OF THE ORIGINAL FILE + * + * The original file, along with the original Apache-2.0 LICENSE can be found at: + * https://github.com/google-research/jax3d/tree/main/jax3d/projects/mobilenerf + * + * Modification details: Shader code was updated to work on Vulkan (originally + * built for WebGL) + * Contributor: (Qualcomm) Rodrigo Holztrattner - quic_rholztra@quicinc.com + */ +#version 460 +#extension GL_EXT_nonuniform_qualifier : enable + +layout (input_attachment_index = 0, binding = 0) uniform subpassInput inputFeature_0; +layout (input_attachment_index = 1, binding = 1) uniform subpassInput inputFeature_1; +layout (input_attachment_index = 2, binding = 2) uniform subpassInput rayDirectionIn; +layout (input_attachment_index = 3, binding = 3) uniform usubpassInput weightsIndex; + +layout(location = 0) out vec4 o_color; + +// Try defining constants in the shader itself +precision highp float; + +#define WEIGHTS_0_COUNT (176) +#define WEIGHTS_1_COUNT (256) +// The third layer's size is changed from 48 to 64 to make sure a 16 bytes alignement +//#define WEIGHTS_2_COUNT (48) +#define WEIGHTS_2_COUNT (64) +#define BIAS_0_COUNT (16) +#define BIAS_1_COUNT (16) +// The third layer bias' size is changed from 3 to 4 to make sure a 16 bytes alignement +#define BIAS_2_COUNT (4) +layout(binding = 4) uniform mlp_weights +{ + vec4 data[(WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + + BIAS_0_COUNT + BIAS_1_COUNT + BIAS_2_COUNT)/4]; // Array of floats +} weights_arr[]; + +vec3 evaluateNetwork( vec4 f0, vec4 f1, vec4 viewdir, uint idx) { + + vec3 res; + + int bias_0_ind = WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT; + vec4 intermediate_one[4] = vec4[]( + weights_arr[nonuniformEXT(idx)].data[bias_0_ind/4], + weights_arr[nonuniformEXT(idx)].data[bias_0_ind/4 + 1], + weights_arr[nonuniformEXT(idx)].data[bias_0_ind/4 + 2], + weights_arr[nonuniformEXT(idx)].data[bias_0_ind/4 + 3] + ); + +#define APPLY_WEIGHTS_0(multiplier, weightFirstInd) \ + intermediate_one[ 0] += (multiplier) * weights_arr[nonuniformEXT(idx)].data[ weightFirstInd/4]; \ + intermediate_one[ 1] += (multiplier) * weights_arr[nonuniformEXT(idx)].data[ weightFirstInd/4 + 1]; \ + intermediate_one[ 2] += (multiplier) * weights_arr[nonuniformEXT(idx)].data[ weightFirstInd/4 + 2]; \ + intermediate_one[ 3] += (multiplier) * weights_arr[nonuniformEXT(idx)].data[ weightFirstInd/4 + 3]; + + APPLY_WEIGHTS_0( f0.r, 0) + APPLY_WEIGHTS_0( f0.g, 16) + APPLY_WEIGHTS_0( f0.b, 32) + APPLY_WEIGHTS_0( f0.a, 48) + APPLY_WEIGHTS_0( f1.r, 64) + APPLY_WEIGHTS_0( f1.g, 80) + APPLY_WEIGHTS_0( f1.b, 96) + APPLY_WEIGHTS_0( f1.a, 112) + // For models form original mobile nerf, use the original code + APPLY_WEIGHTS_0( viewdir.r, 128) + APPLY_WEIGHTS_0( -viewdir.b, 144) + APPLY_WEIGHTS_0( viewdir.g, 160) + + int bias_1_ind = WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + + BIAS_0_COUNT; + vec4 intermediate_two[4] = vec4[]( + weights_arr[nonuniformEXT(idx)].data[bias_1_ind/4], + weights_arr[nonuniformEXT(idx)].data[bias_1_ind/4 + 1], + weights_arr[nonuniformEXT(idx)].data[bias_1_ind/4 + 2], + weights_arr[nonuniformEXT(idx)].data[bias_1_ind/4 + 3] + ); + +#define APPLY_WEIGHTS_1(intermediate, oneInd) \ + if(intermediate > 0.0f){ \ + intermediate_two[ 0] += intermediate * weights_arr[nonuniformEXT(idx)].data[ WEIGHTS_0_COUNT/4 + oneInd * 4 + 0]; \ + intermediate_two[ 1] += intermediate * weights_arr[nonuniformEXT(idx)].data[ WEIGHTS_0_COUNT/4 + oneInd * 4 + 1]; \ + intermediate_two[ 2] += intermediate * weights_arr[nonuniformEXT(idx)].data[ WEIGHTS_0_COUNT/4 + oneInd * 4 + 2]; \ + intermediate_two[ 3] += intermediate * weights_arr[nonuniformEXT(idx)].data[ WEIGHTS_0_COUNT/4 + oneInd * 4 + 3]; \ + } + + APPLY_WEIGHTS_1( intermediate_one[0].r, 0) + APPLY_WEIGHTS_1( intermediate_one[0].g, 1) + APPLY_WEIGHTS_1( intermediate_one[0].b, 2) + APPLY_WEIGHTS_1( intermediate_one[0].a, 3) + APPLY_WEIGHTS_1( intermediate_one[1].r, 4) + APPLY_WEIGHTS_1( intermediate_one[1].g, 5) + APPLY_WEIGHTS_1( intermediate_one[1].b, 6) + APPLY_WEIGHTS_1( intermediate_one[1].a, 7) + APPLY_WEIGHTS_1( intermediate_one[2].r, 8) + APPLY_WEIGHTS_1( intermediate_one[2].g, 9) + APPLY_WEIGHTS_1( intermediate_one[2].b, 10) + APPLY_WEIGHTS_1( intermediate_one[2].a, 11) + APPLY_WEIGHTS_1( intermediate_one[3].r, 12) + APPLY_WEIGHTS_1( intermediate_one[3].g, 13) + APPLY_WEIGHTS_1( intermediate_one[3].b, 14) + APPLY_WEIGHTS_1( intermediate_one[3].a, 15) + + int bias_2_ind = WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + + BIAS_0_COUNT + BIAS_1_COUNT; + vec4 result = weights_arr[nonuniformEXT(idx)].data[bias_2_ind/4]; + +#define APPLY_WEIGHTS_2(intermediate, oneInd) \ + if(intermediate > 0.0f){ \ + result += intermediate * weights_arr[nonuniformEXT(idx)].data[ WEIGHTS_0_COUNT/4 + WEIGHTS_1_COUNT/4 + oneInd]; \ + } + + APPLY_WEIGHTS_2(intermediate_two[0].r, 0) + APPLY_WEIGHTS_2(intermediate_two[0].g, 1) + APPLY_WEIGHTS_2(intermediate_two[0].b, 2) + APPLY_WEIGHTS_2(intermediate_two[0].a, 3) + APPLY_WEIGHTS_2(intermediate_two[1].r, 4) + APPLY_WEIGHTS_2(intermediate_two[1].g, 5) + APPLY_WEIGHTS_2(intermediate_two[1].b, 6) + APPLY_WEIGHTS_2(intermediate_two[1].a, 7) + APPLY_WEIGHTS_2(intermediate_two[2].r, 8) + APPLY_WEIGHTS_2(intermediate_two[2].g, 9) + APPLY_WEIGHTS_2(intermediate_two[2].b,10) + APPLY_WEIGHTS_2(intermediate_two[2].a,11) + APPLY_WEIGHTS_2(intermediate_two[3].r,12) + APPLY_WEIGHTS_2(intermediate_two[3].g,13) + APPLY_WEIGHTS_2(intermediate_two[3].b,14) + APPLY_WEIGHTS_2(intermediate_two[3].a,15) + + result = 1.0 / (1.0 + exp(-result)); + return vec3(result * viewdir.a+(1.0-viewdir.a)); + } + +////////////////////////////////////////////////////////////// +// MLP was trained with gamma-corrected values // +// convert to linear so sRGB conversion isn't applied twice // +////////////////////////////////////////////////////////////// + +float Convert_sRGB_ToLinear(float value) +{ + return value <= 0.04045 + ? value / 12.92 + : pow((value + 0.055) / 1.055, 2.4); +} + +vec3 Convert_sRGB_ToLinear(vec3 value) +{ + return vec3(Convert_sRGB_ToLinear(value.x), Convert_sRGB_ToLinear(value.y), Convert_sRGB_ToLinear(value.z)); +} + +////////////////////////////////////////////////////////////// +////////////////////////////////////////////////////////////// +////////////////////////////////////////////////////////////// + +void main(void) +{ + vec4 feature_0 = subpassLoad(inputFeature_0).rgba; + vec4 feature_1 = subpassLoad(inputFeature_1).rgba; + vec4 rayDirection = subpassLoad(rayDirectionIn).rgba; + uint idx = subpassLoad(weightsIndex).r; + + if (rayDirection.a < 0.6) discard; + + //deal with iphone + feature_0.a = feature_0.a*2.0-1.0; + feature_1.a = feature_1.a*2.0-1.0; + rayDirection.a = rayDirection.a*2.0-1.0; + + // Original + o_color.rgb = Convert_sRGB_ToLinear(evaluateNetwork(feature_0,feature_1,rayDirection,idx)); + o_color.a = 1.0; +} diff --git a/shaders/mobile_nerf/mlp_morpheus.frag b/shaders/mobile_nerf/mlp_morpheus.frag index 96d62c36e..e36e4397e 100644 --- a/shaders/mobile_nerf/mlp_morpheus.frag +++ b/shaders/mobile_nerf/mlp_morpheus.frag @@ -17,7 +17,7 @@ * ------------------------------------------------------------------------ * * THIS IS A MODIFIED VERSION OF THE ORIGINAL FILE - * + * * The original file, along with the original Apache-2.0 LICENSE can be found at: * https://github.com/google-research/jax3d/tree/main/jax3d/projects/mobilenerf * @@ -47,52 +47,27 @@ precision highp float; #define BIAS_2_COUNT (4) layout(binding = 3) uniform mlp_weights { - vec4 data[(WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + + vec4 data[(WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + BIAS_0_COUNT + BIAS_1_COUNT + BIAS_2_COUNT)/4]; // Array of floats } weights; vec3 evaluateNetwork( vec4 f0, vec4 f1, vec4 viewdir) { - + vec3 res; int bias_0_ind = WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT; vec4 intermediate_one[4] = vec4[]( - weights.data[bias_0_ind/4], + weights.data[bias_0_ind/4], weights.data[bias_0_ind/4 + 1], weights.data[bias_0_ind/4 + 2], weights.data[bias_0_ind/4 + 3] ); - #if 0 - for (int j = 0; j < 11; ++j) { - float input_value = 0.0; - if (j < 4) { - input_value = - (j == 0) ? f0.r : ( - (j == 1) ? f0.g : ( - (j == 2) ? f0.b : f0.a)); - } else if (j < 8) { - input_value = - (j == 4) ? f1.r : ( - (j == 5) ? f1.g : ( - (j == 6) ? f1.b : f1.a)); - } else { - input_value = - (j == 8) ? viewdir.r : ( - (j == 9) ? -viewdir.b : viewdir.g); //switch y-z axes - } - for (int i = 0; i < 16; ++i) { - intermediate_one[i] += input_value * weights0[ j*16 + i]; - //texelFetch(weightsZero, ivec2(j, i), 0).x; - } - } - #endif - #define APPLY_WEIGHTS_0(multiplier, weightFirstInd) \ intermediate_one[ 0] += (multiplier) * weights.data[ weightFirstInd/4]; \ intermediate_one[ 1] += (multiplier) * weights.data[ weightFirstInd/4 + 1]; \ intermediate_one[ 2] += (multiplier) * weights.data[ weightFirstInd/4 + 2]; \ - intermediate_one[ 3] += (multiplier) * weights.data[ weightFirstInd/4 + 3]; + intermediate_one[ 3] += (multiplier) * weights.data[ weightFirstInd/4 + 3]; APPLY_WEIGHTS_0( f0.r, 0) APPLY_WEIGHTS_0( f0.g, 16) @@ -107,27 +82,15 @@ vec3 evaluateNetwork( vec4 f0, vec4 f1, vec4 viewdir) { APPLY_WEIGHTS_0( (-viewdir.b + 1.0 )/2, 144) APPLY_WEIGHTS_0( (viewdir.g + 1.0 )/2, 160) - int bias_1_ind = WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + + int bias_1_ind = WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + BIAS_0_COUNT; vec4 intermediate_two[4] = vec4[]( - weights.data[bias_1_ind/4], + weights.data[bias_1_ind/4], weights.data[bias_1_ind/4 + 1], weights.data[bias_1_ind/4 + 2], weights.data[bias_1_ind/4 + 3] ); - #if 0 - for (int j = 0; j < 16; ++j) { - if (intermediate_one[j] <= 0.0) { - continue; - } - for (int i = 0; i < 16; ++i) { - intermediate_two[i] += intermediate_one[j] * weights1[j*16 + i]; - //texelFetch(weightsOne, ivec2(j, i), 0).x; - } - } - #endif - #define APPLY_WEIGHTS_1(intermediate, oneInd) \ if(intermediate > 0.0f){ \ intermediate_two[ 0] += intermediate * weights.data[ WEIGHTS_0_COUNT/4 + oneInd * 4 + 0]; \ @@ -135,7 +98,7 @@ vec3 evaluateNetwork( vec4 f0, vec4 f1, vec4 viewdir) { intermediate_two[ 2] += intermediate * weights.data[ WEIGHTS_0_COUNT/4 + oneInd * 4 + 2]; \ intermediate_two[ 3] += intermediate * weights.data[ WEIGHTS_0_COUNT/4 + oneInd * 4 + 3]; \ } - + APPLY_WEIGHTS_1( intermediate_one[0].r, 0) APPLY_WEIGHTS_1( intermediate_one[0].g, 1) APPLY_WEIGHTS_1( intermediate_one[0].b, 2) @@ -153,27 +116,15 @@ vec3 evaluateNetwork( vec4 f0, vec4 f1, vec4 viewdir) { APPLY_WEIGHTS_1( intermediate_one[3].b, 14) APPLY_WEIGHTS_1( intermediate_one[3].a, 15) - int bias_2_ind = WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + + int bias_2_ind = WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + BIAS_0_COUNT + BIAS_1_COUNT; vec4 result = weights.data[bias_2_ind/4]; - #if 0 - for (int j = 0; j < 16; ++j) { - if (intermediate_two[j] <= 0.0) { - continue; - } - for (int i = 0; i < 3; ++i) { - result[i] += intermediate_two[j] * weights2[j*3 + i]; - //texelFetch(weightsTwo, ivec2(j, i), 0).x; - } - } - #endif - #define APPLY_WEIGHTS_2(intermediate, oneInd) \ if(intermediate > 0.0f){ \ result += intermediate * weights.data[ WEIGHTS_0_COUNT/4 + WEIGHTS_1_COUNT/4 + oneInd]; \ } - + APPLY_WEIGHTS_2(intermediate_two[0].r, 0) APPLY_WEIGHTS_2(intermediate_two[0].g, 1) APPLY_WEIGHTS_2(intermediate_two[0].b, 2) @@ -190,9 +141,9 @@ vec3 evaluateNetwork( vec4 f0, vec4 f1, vec4 viewdir) { APPLY_WEIGHTS_2(intermediate_two[3].g,13) APPLY_WEIGHTS_2(intermediate_two[3].b,14) APPLY_WEIGHTS_2(intermediate_two[3].a,15) - + result = 1.0 / (1.0 + exp(-result)); - return vec3(result * viewdir.a+(1.0-viewdir.a)); + return vec3(result * viewdir.a+(1.0-viewdir.a)); } ////////////////////////////////////////////////////////////// @@ -207,7 +158,7 @@ float Convert_sRGB_ToLinear(float value) : pow((value + 0.055) / 1.055, 2.4); } -vec3 Convert_sRGB_ToLinear(vec3 value) +vec3 Convert_sRGB_ToLinear(vec3 value) { return vec3(Convert_sRGB_ToLinear(value.x), Convert_sRGB_ToLinear(value.y), Convert_sRGB_ToLinear(value.z)); } @@ -223,27 +174,14 @@ void main(void) vec4 rayDirection = subpassLoad(rayDirectionIn).rgba; if (rayDirection.a < 0.6) discard; - + //deal with iphone feature_0.a = feature_0.a*2.0-1.0; feature_1.a = feature_1.a*2.0-1.0; rayDirection.a = rayDirection.a*2.0-1.0; - - // Original + // Original o_color.rgb = Convert_sRGB_ToLinear(evaluateNetwork(feature_0,feature_1,rayDirection)); - //o_color.rgb = feature_0.rgb; o_color.a = 1.0; - - - // --------------------------------------------------------------------------- - // Debugging - - //o_color.rgb = (rayDirection.rgb)*0.5 + vec3(0.5,0.5,0.5); - //o_color.rgb = rayDirection.rgb; - - //o_color.rgba = feature_1; - //o_color.a = 1.0f; - //o_color.rgba = feature_1; - + } diff --git a/shaders/mobile_nerf/mlp_morpheus_combo.frag b/shaders/mobile_nerf/mlp_morpheus_combo.frag new file mode 100644 index 000000000..330950dda --- /dev/null +++ b/shaders/mobile_nerf/mlp_morpheus_combo.frag @@ -0,0 +1,189 @@ +/* Copyright (c) 2024, Qualcomm Innovation Center, Inc. All rights reserved. + * + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 the "License"; + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + * ------------------------------------------------------------------------ + * + * THIS IS A MODIFIED VERSION OF THE ORIGINAL FILE + * + * The original file, along with the original Apache-2.0 LICENSE can be found at: + * https://github.com/google-research/jax3d/tree/main/jax3d/projects/mobilenerf + * + * Modification details: Shader code was updated to work on Vulkan (originally + * built for WebGL) + * Contributor: (Qualcomm) Rodrigo Holztrattner - quic_rholztra@quicinc.com + */ +#version 460 +#extension GL_EXT_nonuniform_qualifier : enable + +layout (input_attachment_index = 0, binding = 0) uniform subpassInput inputFeature_0; +layout (input_attachment_index = 1, binding = 1) uniform subpassInput inputFeature_1; +layout (input_attachment_index = 2, binding = 2) uniform subpassInput rayDirectionIn; +layout (input_attachment_index = 3, binding = 3) uniform usubpassInput weightsIndex; + +layout(location = 0) out vec4 o_color; + +// Try defining constants in the shader itself +precision highp float; + +#define WEIGHTS_0_COUNT (176) +#define WEIGHTS_1_COUNT (256) +// The third layer's size is changed from 48 to 64 to make sure a 16 bytes alignement +//#define WEIGHTS_2_COUNT (48) +#define WEIGHTS_2_COUNT (64) +#define BIAS_0_COUNT (16) +#define BIAS_1_COUNT (16) +// The third layer bias' size is changed from 3 to 4 to make sure a 16 bytes alignement +#define BIAS_2_COUNT (4) +layout(binding = 4) uniform mlp_weights +{ + vec4 data[(WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + + BIAS_0_COUNT + BIAS_1_COUNT + BIAS_2_COUNT)/4]; // Array of floats +} weights_arr[]; + +vec3 evaluateNetwork( vec4 f0, vec4 f1, vec4 viewdir, uint idx) { + + vec3 res; + + int bias_0_ind = WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT; + vec4 intermediate_one[4] = vec4[]( + weights_arr[nonuniformEXT(idx)].data[bias_0_ind/4], + weights_arr[nonuniformEXT(idx)].data[bias_0_ind/4 + 1], + weights_arr[nonuniformEXT(idx)].data[bias_0_ind/4 + 2], + weights_arr[nonuniformEXT(idx)].data[bias_0_ind/4 + 3] + ); + +#define APPLY_WEIGHTS_0(multiplier, weightFirstInd) \ + intermediate_one[ 0] += (multiplier) * weights_arr[nonuniformEXT(idx)].data[ weightFirstInd/4]; \ + intermediate_one[ 1] += (multiplier) * weights_arr[nonuniformEXT(idx)].data[ weightFirstInd/4 + 1]; \ + intermediate_one[ 2] += (multiplier) * weights_arr[nonuniformEXT(idx)].data[ weightFirstInd/4 + 2]; \ + intermediate_one[ 3] += (multiplier) * weights_arr[nonuniformEXT(idx)].data[ weightFirstInd/4 + 3]; + + APPLY_WEIGHTS_0( f0.r, 0) + APPLY_WEIGHTS_0( f0.g, 16) + APPLY_WEIGHTS_0( f0.b, 32) + APPLY_WEIGHTS_0( f0.a, 48) + APPLY_WEIGHTS_0( f1.r, 64) + APPLY_WEIGHTS_0( f1.g, 80) + APPLY_WEIGHTS_0( f1.b, 96) + APPLY_WEIGHTS_0( f1.a, 112) + // For models form the Morpheus team, the view direction need to be handled differently + APPLY_WEIGHTS_0( (viewdir.r + 1.0 )/2, 128) + APPLY_WEIGHTS_0( (-viewdir.b + 1.0 )/2, 144) + APPLY_WEIGHTS_0( (viewdir.g + 1.0 )/2, 160) + + int bias_1_ind = WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + + BIAS_0_COUNT; + vec4 intermediate_two[4] = vec4[]( + weights_arr[nonuniformEXT(idx)].data[bias_1_ind/4], + weights_arr[nonuniformEXT(idx)].data[bias_1_ind/4 + 1], + weights_arr[nonuniformEXT(idx)].data[bias_1_ind/4 + 2], + weights_arr[nonuniformEXT(idx)].data[bias_1_ind/4 + 3] + ); + +#define APPLY_WEIGHTS_1(intermediate, oneInd) \ + if(intermediate > 0.0f){ \ + intermediate_two[ 0] += intermediate * weights_arr[nonuniformEXT(idx)].data[ WEIGHTS_0_COUNT/4 + oneInd * 4 + 0]; \ + intermediate_two[ 1] += intermediate * weights_arr[nonuniformEXT(idx)].data[ WEIGHTS_0_COUNT/4 + oneInd * 4 + 1]; \ + intermediate_two[ 2] += intermediate * weights_arr[nonuniformEXT(idx)].data[ WEIGHTS_0_COUNT/4 + oneInd * 4 + 2]; \ + intermediate_two[ 3] += intermediate * weights_arr[nonuniformEXT(idx)].data[ WEIGHTS_0_COUNT/4 + oneInd * 4 + 3]; \ + } + + APPLY_WEIGHTS_1( intermediate_one[0].r, 0) + APPLY_WEIGHTS_1( intermediate_one[0].g, 1) + APPLY_WEIGHTS_1( intermediate_one[0].b, 2) + APPLY_WEIGHTS_1( intermediate_one[0].a, 3) + APPLY_WEIGHTS_1( intermediate_one[1].r, 4) + APPLY_WEIGHTS_1( intermediate_one[1].g, 5) + APPLY_WEIGHTS_1( intermediate_one[1].b, 6) + APPLY_WEIGHTS_1( intermediate_one[1].a, 7) + APPLY_WEIGHTS_1( intermediate_one[2].r, 8) + APPLY_WEIGHTS_1( intermediate_one[2].g, 9) + APPLY_WEIGHTS_1( intermediate_one[2].b, 10) + APPLY_WEIGHTS_1( intermediate_one[2].a, 11) + APPLY_WEIGHTS_1( intermediate_one[3].r, 12) + APPLY_WEIGHTS_1( intermediate_one[3].g, 13) + APPLY_WEIGHTS_1( intermediate_one[3].b, 14) + APPLY_WEIGHTS_1( intermediate_one[3].a, 15) + + int bias_2_ind = WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + + BIAS_0_COUNT + BIAS_1_COUNT; + vec4 result = weights_arr[nonuniformEXT(idx)].data[bias_2_ind/4]; + +#define APPLY_WEIGHTS_2(intermediate, oneInd) \ + if(intermediate > 0.0f){ \ + result += intermediate * weights_arr[nonuniformEXT(idx)].data[ WEIGHTS_0_COUNT/4 + WEIGHTS_1_COUNT/4 + oneInd]; \ + } + + APPLY_WEIGHTS_2(intermediate_two[0].r, 0) + APPLY_WEIGHTS_2(intermediate_two[0].g, 1) + APPLY_WEIGHTS_2(intermediate_two[0].b, 2) + APPLY_WEIGHTS_2(intermediate_two[0].a, 3) + APPLY_WEIGHTS_2(intermediate_two[1].r, 4) + APPLY_WEIGHTS_2(intermediate_two[1].g, 5) + APPLY_WEIGHTS_2(intermediate_two[1].b, 6) + APPLY_WEIGHTS_2(intermediate_two[1].a, 7) + APPLY_WEIGHTS_2(intermediate_two[2].r, 8) + APPLY_WEIGHTS_2(intermediate_two[2].g, 9) + APPLY_WEIGHTS_2(intermediate_two[2].b,10) + APPLY_WEIGHTS_2(intermediate_two[2].a,11) + APPLY_WEIGHTS_2(intermediate_two[3].r,12) + APPLY_WEIGHTS_2(intermediate_two[3].g,13) + APPLY_WEIGHTS_2(intermediate_two[3].b,14) + APPLY_WEIGHTS_2(intermediate_two[3].a,15) + + result = 1.0 / (1.0 + exp(-result)); + return vec3(result * viewdir.a+(1.0-viewdir.a)); + } + +////////////////////////////////////////////////////////////// +// MLP was trained with gamma-corrected values // +// convert to linear so sRGB conversion isn't applied twice // +////////////////////////////////////////////////////////////// + +float Convert_sRGB_ToLinear(float value) +{ + return value <= 0.04045 + ? value / 12.92 + : pow((value + 0.055) / 1.055, 2.4); +} + +vec3 Convert_sRGB_ToLinear(vec3 value) +{ + return vec3(Convert_sRGB_ToLinear(value.x), Convert_sRGB_ToLinear(value.y), Convert_sRGB_ToLinear(value.z)); +} + +////////////////////////////////////////////////////////////// +////////////////////////////////////////////////////////////// +////////////////////////////////////////////////////////////// + +void main(void) +{ + vec4 feature_0 = subpassLoad(inputFeature_0).rgba; + vec4 feature_1 = subpassLoad(inputFeature_1).rgba; + vec4 rayDirection = subpassLoad(rayDirectionIn).rgba; + uint idx = subpassLoad(weightsIndex).r; + + if (rayDirection.a < 0.6) discard; + + //deal with iphone + feature_0.a = feature_0.a*2.0-1.0; + feature_1.a = feature_1.a*2.0-1.0; + rayDirection.a = rayDirection.a*2.0-1.0; + + // Original + o_color.rgb = Convert_sRGB_ToLinear(evaluateNetwork(feature_0,feature_1,rayDirection,idx)); + o_color.a = 1.0; +} diff --git a/shaders/mobile_nerf/quad.vert b/shaders/mobile_nerf/quad.vert index bd34fb92e..3221e88d0 100644 --- a/shaders/mobile_nerf/quad.vert +++ b/shaders/mobile_nerf/quad.vert @@ -1,4 +1,4 @@ -/* Copyright (c) 2023, Qualcomm Innovation Center, Inc. All rights reserved. +/* Copyright (c) 2024, Qualcomm Innovation Center, Inc. All rights reserved. * * SPDX-License-Identifier: Apache-2.0 * @@ -17,7 +17,7 @@ * ------------------------------------------------------------------------ * * THIS IS A MODIFIED VERSION OF THE ORIGINAL FILE - * + * * The original file, along with the original Apache-2.0 LICENSE can be found at: * https://github.com/google-research/jax3d/tree/main/jax3d/projects/mobilenerf * @@ -31,7 +31,7 @@ out gl_PerVertex { vec4 gl_Position; }; -void main() +void main() { vec2 outUV = vec2((gl_VertexIndex << 1) & 2, gl_VertexIndex & 2); gl_Position = vec4(outUV * 2.0f - 1.0f, 0.0f, 1.0f); diff --git a/shaders/mobile_nerf/raster.frag b/shaders/mobile_nerf/raster.frag index 1ffcbbbaa..f3603d9e8 100644 --- a/shaders/mobile_nerf/raster.frag +++ b/shaders/mobile_nerf/raster.frag @@ -17,7 +17,7 @@ * ------------------------------------------------------------------------ * * THIS IS A MODIFIED VERSION OF THE ORIGINAL FILE - * + * * The original file, along with the original Apache-2.0 LICENSE can be found at: * https://github.com/google-research/jax3d/tree/main/jax3d/projects/mobilenerf * @@ -38,36 +38,15 @@ layout(location = 2) out vec4 rayDirectionOut; layout(binding = 0) uniform sampler2D textureInput_0; layout(binding = 1) uniform sampler2D textureInput_1; -////////////////////////////////////////////////////////////// -// MLP was trained with gamma-corrected values // -// convert to linear so sRGB conversion isn't applied twice // -////////////////////////////////////////////////////////////// - -float Convert_sRGB_ToLinear(float value) -{ - return value <= 0.04045 - ? value / 12.92 - : pow((value + 0.055) / 1.055, 2.4); -} - -vec3 Convert_sRGB_ToLinear(vec3 value) -{ - return vec3(Convert_sRGB_ToLinear(value.x), Convert_sRGB_ToLinear(value.y), Convert_sRGB_ToLinear(value.z)); -} - -////////////////////////////////////////////////////////////// -////////////////////////////////////////////////////////////// -////////////////////////////////////////////////////////////// - void main(void) { vec2 flipped = vec2( texCoord_frag.x, 1.0 - texCoord_frag.y ); vec4 pixel_0 = texture(textureInput_0, flipped); if (pixel_0.r == 0.0) discard; vec4 pixel_1 = texture(textureInput_1, flipped); - o_color_0 = vec4(Convert_sRGB_ToLinear(pixel_0.xyz), pixel_0.w); - o_color_1 = vec4(Convert_sRGB_ToLinear(pixel_1.xyz), pixel_1.w); - + o_color_0 = vec4(pixel_0.xyz, pixel_0.w); + o_color_1 = vec4(pixel_1.xyz, pixel_1.w); + rayDirectionOut.rgb = normalize(rayDirectionIn); rayDirectionOut.a = 1.0f; } diff --git a/shaders/mobile_nerf/raster.vert b/shaders/mobile_nerf/raster.vert index e293cdc2f..2f62235e1 100644 --- a/shaders/mobile_nerf/raster.vert +++ b/shaders/mobile_nerf/raster.vert @@ -1,4 +1,4 @@ -/* Copyright (c) 2023, Qualcomm Innovation Center, Inc. All rights reserved. +/* Copyright (c) 2024, Qualcomm Innovation Center, Inc. All rights reserved. * * SPDX-License-Identifier: Apache-2.0 * @@ -17,7 +17,7 @@ * ------------------------------------------------------------------------ * * THIS IS A MODIFIED VERSION OF THE ORIGINAL FILE - * + * * The original file, along with the original Apache-2.0 LICENSE can be found at: * https://github.com/google-research/jax3d/tree/main/jax3d/projects/mobilenerf * @@ -50,9 +50,9 @@ layout(location = 1) out vec3 rayDirection; void main(void) { texCoord_frag = texCoord; - + vec3 pos = position + posOffset; gl_Position = global_uniform.proj * global_uniform.view * global_uniform.model * vec4(pos.x, -pos.y, pos.z, 1.0); rayDirection = pos - vec3(global_uniform.camera_position.x, -global_uniform.camera_position.y, global_uniform.camera_position.z); -} \ No newline at end of file +} diff --git a/shaders/mobile_nerf/mlp.comp b/shaders/mobile_nerf/raster_combo.frag similarity index 54% rename from shaders/mobile_nerf/mlp.comp rename to shaders/mobile_nerf/raster_combo.frag index 8decb4352..67507d517 100644 --- a/shaders/mobile_nerf/mlp.comp +++ b/shaders/mobile_nerf/raster_combo.frag @@ -1,4 +1,4 @@ -/* Copyright (c) 2023, Qualcomm Innovation Center, Inc. All rights reserved. +/* Copyright (c) 2024, Qualcomm Innovation Center, Inc. All rights reserved. * * SPDX-License-Identifier: Apache-2.0 * @@ -17,7 +17,7 @@ * ------------------------------------------------------------------------ * * THIS IS A MODIFIED VERSION OF THE ORIGINAL FILE - * + * * The original file, along with the original Apache-2.0 LICENSE can be found at: * https://github.com/google-research/jax3d/tree/main/jax3d/projects/mobilenerf * @@ -27,30 +27,32 @@ */ #version 460 -#extension GL_EXT_ray_query : enable - -layout(binding = 0, set = 0, rgba8) uniform image2D image; - -layout(binding = 1, set = 0) uniform sampler2D inputFeature_0; +layout(location = 0) in vec2 texCoord_frag; +layout(location = 1) in vec3 rayDirectionIn; -layout(binding = 2, set = 0) uniform sampler2D inputFeature_1; +layout(location = 0) out vec4 o_color_0; +layout(location = 1) out vec4 o_color_1; +layout(location = 2) out vec4 rayDirectionOut; +layout(location = 3) out uint weights_idx_out; -// TODO should add mlp weight in here -//layout(set = 0, binding = 2) uniform mlp_weights.... +layout(binding = 0) uniform sampler2D textureInput_0; +layout(binding = 1) uniform sampler2D textureInput_1; -layout(local_size_x = 16, local_size_y = 16, local_size_z = 1) in; +layout(push_constant) uniform PushConstants { + uint weights_idx; +} pc; -void main() +void main(void) { - const uvec3 globalSize = gl_NumWorkGroups * gl_WorkGroupSize; - const vec2 pixelCenter = vec2(gl_GlobalInvocationID.xy) + vec2(0.5); - const vec2 inUV = pixelCenter/vec2(globalSize); + vec2 flipped = vec2( texCoord_frag.x, 1.0 - texCoord_frag.y ); + vec4 pixel_0 = texture(textureInput_0, flipped); + if (pixel_0.r == 0.0) discard; + vec4 pixel_1 = texture(textureInput_1, flipped); + o_color_0 = vec4(pixel_0.xyz, pixel_0.w); + o_color_1 = vec4(pixel_1.xyz, pixel_1.w); - vec4 feature_0 = texture(inputFeature_0, inUV); - vec4 feature_1 = texture(inputFeature_1, inUV); + rayDirectionOut.rgb = normalize(rayDirectionIn); + rayDirectionOut.a = 1.0f; - vec4 o_color = vec4(((feature_0 + feature_1) / 2.0f).rgb, 1.0f); - - imageStore(image, ivec2(gl_GlobalInvocationID.xy), o_color); + weights_idx_out = pc.weights_idx; } - diff --git a/shaders/mobile_nerf/raster_morpheus.frag b/shaders/mobile_nerf/raster_morpheus.frag index a04c5e08c..90bfbaa10 100644 --- a/shaders/mobile_nerf/raster_morpheus.frag +++ b/shaders/mobile_nerf/raster_morpheus.frag @@ -17,7 +17,7 @@ * ------------------------------------------------------------------------ * * THIS IS A MODIFIED VERSION OF THE ORIGINAL FILE - * + * * The original file, along with the original Apache-2.0 LICENSE can be found at: * https://github.com/google-research/jax3d/tree/main/jax3d/projects/mobilenerf * @@ -38,39 +38,15 @@ layout(location = 2) out vec4 rayDirectionOut; layout(binding = 0) uniform sampler2D textureInput_0; layout(binding = 1) uniform sampler2D textureInput_1; -////////////////////////////////////////////////////////////// -// MLP was trained with gamma-corrected values // -// convert to linear so sRGB conversion isn't applied twice // -////////////////////////////////////////////////////////////// - -float Convert_sRGB_ToLinear(float value) -{ - return value <= 0.04045 - ? value / 12.92 - : pow((value + 0.055) / 1.055, 2.4); -} - -vec3 Convert_sRGB_ToLinear(vec3 value) -{ - return vec3(Convert_sRGB_ToLinear(value.x), Convert_sRGB_ToLinear(value.y), Convert_sRGB_ToLinear(value.z)); -} - -////////////////////////////////////////////////////////////// -////////////////////////////////////////////////////////////// -////////////////////////////////////////////////////////////// - void main(void) { vec2 flipped = vec2( texCoord_frag.x, 1.0 - texCoord_frag.y ); vec4 pixel_0 = texture(textureInput_0, flipped); // if (pixel_0.r == 0.0) discard; vec4 pixel_1 = texture(textureInput_1, flipped); - o_color_0 = vec4(Convert_sRGB_ToLinear(pixel_0.xyz), pixel_0.w); - o_color_1 = vec4(Convert_sRGB_ToLinear(pixel_1.xyz), pixel_1.w); - + o_color_0 = vec4(pixel_0.xyz, pixel_0.w); + o_color_1 = vec4(pixel_1.xyz, pixel_1.w); + rayDirectionOut.rgb = normalize(rayDirectionIn); rayDirectionOut.a = 1.0f; - - // For debugging only - // o_color_0 = vec4( texCoord_frag.x, 0.0f, 0.0f, 1.0f); } diff --git a/shaders/mobile_nerf/raster_morpheus_combo.frag b/shaders/mobile_nerf/raster_morpheus_combo.frag new file mode 100644 index 000000000..da3777165 --- /dev/null +++ b/shaders/mobile_nerf/raster_morpheus_combo.frag @@ -0,0 +1,58 @@ +/* Copyright (c) 2024, Qualcomm Innovation Center, Inc. All rights reserved. + * + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 the "License"; + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + * ------------------------------------------------------------------------ + * + * THIS IS A MODIFIED VERSION OF THE ORIGINAL FILE + * + * The original file, along with the original Apache-2.0 LICENSE can be found at: + * https://github.com/google-research/jax3d/tree/main/jax3d/projects/mobilenerf + * + * Modification details: Shader code was updated to work on Vulkan (originally + * built for WebGL) + * Contributor: (Qualcomm) Rodrigo Holztrattner - quic_rholztra@quicinc.com + */ +#version 460 + +layout(location = 0) in vec2 texCoord_frag; +layout(location = 1) in vec3 rayDirectionIn; + +layout(location = 0) out vec4 o_color_0; +layout(location = 1) out vec4 o_color_1; +layout(location = 2) out vec4 rayDirectionOut; +layout(location = 3) out uint weights_idx_out; + +layout(binding = 0) uniform sampler2D textureInput_0; +layout(binding = 1) uniform sampler2D textureInput_1; + +layout(push_constant) uniform PushConstants { + uint weights_idx; +} pc; + +void main(void) +{ + vec2 flipped = vec2( texCoord_frag.x, 1.0 - texCoord_frag.y ); + vec4 pixel_0 = texture(textureInput_0, flipped); + // if (pixel_0.r == 0.0) discard; + vec4 pixel_1 = texture(textureInput_1, flipped); + o_color_0 = vec4(pixel_0.xyz, pixel_0.w); + o_color_1 = vec4(pixel_1.xyz, pixel_1.w); + + rayDirectionOut.rgb = normalize(rayDirectionIn); + rayDirectionOut.a = 1.0f; + + weights_idx_out = pc.weights_idx; +}