Vulkan - Bringing material data to ray tracing shader and differentiate them.

Published June 26, 2023
Advertisement

A follow up post after I added material graph system in my mini engine.

To be more specific, it should be “bringing material data to a hit group shader and differentiate them”.

The result is almost the same as the previous demo. But it can differeniate the material in the ray tracing shadow now.

§Overview

After an hit object is confirmed in ray tracing workflow. We not only need to fetch the mesh data, but also the material data for most cases.

Unless the material data is coming from a traditional, fixed design that can be easily indexed with. (E.g. Provide fixed properties like diffuse, normal, specular, opacity and assume every material uses the same calculations.)

Otherwise, you will need a local solution in the hit group shader to fetch the material data. A piece of D3D12 HLSL code as an example:

// ray tracing local root signature
LocalRootSignature RootSigLocal =
{
	"RootConstants( num32BitConstants = 64, b2 )"   // material constant
};

// link local root signature to hit group
SubobjectToExportsAssociation LocalRootSignatureAssociation =
{
	"RootSigLocal",  // subobject name
	"RayHitGroup"    // export association 
};

TriangleHitGroup RayHitGroup =
{
    "RTAnyHit",      // AnyHit
    "RTClosestHit"   // ClosestHit
};

With the subobject export association, we can easily link a local material data to the hit group. (I'll simplize by calling it “local descriptor")

We just need to copy the material constants to the shader records as well. (Start after the shader identifier bytes)

In D3D12, this is possible. But in Vulkan, it has only the global descriptors and I have to figure out other ways to pass material data.

§Solution

With bindless rendering and dynamically indexing “only” are not enough. As I said, each material could have different set of constants and also the calculations.

Also, Vulkan doesn't support local desciptors. I need a global-to-local way to achieve my goal. So my solution is combining both shader table indexing (aka SBT) and bindless indexing.

The SBT is for evaluating different materials, here is one example translated in my system:

// another 2D index array for matching, since Vulkan doesn't implement local descriptor yet, I need this to fetch data
// each texture node takes one slot, and the value of Index will be updated in C++ side
// access via InstanceID() first, then the index calculated by graph system
struct MaterialData
{
	int TextureIndex;
	int SamplerIndex;
	float Cutoff;
	float Padding;
};
StructuredBuffer<MaterialData> UHMaterialDataTable[] : register(t0, space6);

// get material input, the simple version that has opacity only
UHMaterialInputs GetMaterialInputSimple(float2 UV0, float MipLevel, out float Cutoff)
{
	Cutoff = UHMaterialDataTable[InstanceID()][0].Cutoff;

	// material input code will be generated in C++ side
	float4 Result_1479 = UHTextureTable[UHMaterialDataTable[InstanceID()][2].TextureIndex].SampleLevel(UHSamplerTable[UHMaterialDataTable[InstanceID()][2].SamplerIndex], UV0, MipLevel);
	float4 Result_1472 = UHTextureTable[UHMaterialDataTable[InstanceID()][1].TextureIndex].SampleLevel(UHSamplerTable[UHMaterialDataTable[InstanceID()][1].SamplerIndex], UV0, MipLevel) * 2.0f - 1.0f;
	float3 Node_1462 = float3(0.50f, 0.50f, 0.50f);
	float Node_1459 = 1.00f;
	float4 Result_1456 = UHTextureTable[UHMaterialDataTable[InstanceID()][0].TextureIndex].SampleLevel(UHSamplerTable[UHMaterialDataTable[InstanceID()][0].SamplerIndex], UV0, MipLevel);
	float3 Node_1446 = float3(1.00f, 1.00f, 1.00f);

	UHMaterialInputs Input = (UHMaterialInputs)0;
	Input.Opacity = (Node_1459 * Result_1479.rgb).r;
	return Input;
	return (UHMaterialInputs)0;
}

I use InstanceID() to access the StructuredBuffer<MaterialData> descriptor array. Then, doing the literal indexing.

Why it must be this complicated? Before answering this, here is another piece of code of the same material, but in the base pass pixel shader:

cbuffer PassConstant : register(UHMAT_BIND)
{
		int Node_1456_Index;
	int Node_1472_Index;
	int Node_1479_Index;
	int DefaultAniso16_Index;
	float GCutoff;
	float GEnvCubeMipMapCount;

}

UHMaterialInputs GetMaterialInput(float2 UV0)
{
	// material input code will be generated in C++ side
	float4 Result_1479 = UHTextureTable[Node_1479_Index].Sample(UHSamplerTable[DefaultAniso16_Index], UV0);
	float4 Result_1472 = UHTextureTable[Node_1472_Index].Sample(UHSamplerTable[DefaultAniso16_Index], UV0) * 2.0f - 1.0f;
	float3 Node_1462 = float3(0.50f, 0.50f, 0.50f);
	float Node_1459 = 1.00f;
	float4 Result_1456 = UHTextureTable[Node_1456_Index].Sample(UHSamplerTable[DefaultAniso16_Index], UV0);
	float3 Node_1446 = float3(1.00f, 1.00f, 1.00f);

	UHMaterialInputs Input = (UHMaterialInputs)0;
	Input.Opacity = (Node_1459 * Result_1479.rgb).r;
	Input.Diffuse = (Node_1446 * Result_1456.rgb).rgb;
	Input.Occlusion = Node_1459.r;
	Input.Specular = Node_1462.rgb;
	Input.Normal = (Node_1459 * Result_1472.rgb).rgb;
	Input.Metallic = 0.0f;
	Input.Roughness = 1.0f;
	Input.FresnelFactor = 1.0f;
	Input.ReflectionFactor = 0.5f;
	Input.Emissive = float3(0,0,0);
	return Input;
}

Notice the difference right? In regular object pass, I can bind constant buffer with translated texture index.

And that's not a possible in the hit group shader. Since it's lack of a local descriptor feature.

So it translates 0,1,2 in hit group shader instead of Node_1456_Index, Node_1472_Index, Node_1479_Index. It's bascially doing a “index to an array of index” stuff.

Why not just going fully literal index? Because the texture index can change depends on the load order or something else. And the array of index lookup gives flexibility.

§Changes in Vulkan Side

Of course, the CPU side needs some changes.

Setting in acceleration structure instance

When setting up VkAccelerationStructureInstanceKHR, you need to bind the SBT offset as well.

	// set material buffer data index as SBT index, each material has an unique hitgroup shader
	InstanceKHR.instanceShaderBindingTableRecordOffset = Mat->GetBufferDataIndex();
	InstanceKHR.instanceCustomIndex = Mat->GetBufferDataIndex();

Based on the Microsoft DXR specs, the instanceShaderBindingTableRecordOffset corresponds to InstanceContributionToHitGroupIndex apparently.

You can adjust the indexing depend on your application.

The VkPipeline Creation

I'll skip a few lines as it could be too long.

	VkRayTracingPipelineCreateInfoKHR CreateInfo{};
	// ... skipped ... //

	// set RG shader
	std::string RGEntryName = InInfo.RayGenShader->GetEntryName();
	VkPipelineShaderStageCreateInfo RGStageInfo{};
	RGStageInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
	RGStageInfo.stage = VK_SHADER_STAGE_RAYGEN_BIT_KHR;
	RGStageInfo.module = InInfo.RayGenShader->GetShader();
	RGStageInfo.pName = RGEntryName.c_str();

	// set miss shader
	std::string MissEntryName = InInfo.MissShader->GetEntryName();
	VkPipelineShaderStageCreateInfo MissStageInfo{};
	MissStageInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
	MissStageInfo.stage = VK_SHADER_STAGE_MISS_BIT_KHR;
	MissStageInfo.module = InInfo.MissShader->GetShader();
	MissStageInfo.pName = MissEntryName.c_str();

	// set closest hit shader
	std::string CHGEntryName = InInfo.ClosestHitShader->GetEntryName();
	VkPipelineShaderStageCreateInfo CHGStageInfo{};
	CHGStageInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
	CHGStageInfo.stage = VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR;
	CHGStageInfo.module = InInfo.ClosestHitShader->GetShader();
	CHGStageInfo.pName = CHGEntryName.c_str();

	// set any hit shader
	std::vector<VkPipelineShaderStageCreateInfo> AHGStageInfos(InInfo.AnyHitShaders.size());
	std::vector<std::string> AHGEntryNames(InInfo.AnyHitShaders.size());

	for (size_t Idx = 0; Idx < InInfo.AnyHitShaders.size(); Idx++)
	{
		AHGEntryNames[Idx] = InInfo.AnyHitShaders[Idx]->GetEntryName();
		VkPipelineShaderStageCreateInfo AHGStageInfo{};
		if (InInfo.AnyHitShaders[Idx] != nullptr)
		{
			AHGStageInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
			AHGStageInfo.stage = VK_SHADER_STAGE_ANY_HIT_BIT_KHR;
			AHGStageInfo.module = InInfo.AnyHitShaders[Idx]->GetShader();
			AHGStageInfo.pName = AHGEntryNames[Idx].c_str();
			AHGStageInfos[Idx] = AHGStageInfo;
		}
	}

	std::vector<VkPipelineShaderStageCreateInfo> ShaderStages = { RGStageInfo, MissStageInfo, CHGStageInfo };
	ShaderStages.insert(ShaderStages.end(), AHGStageInfos.begin(), AHGStageInfos.end());
	CreateInfo.stageCount = static_cast<uint32_t>(ShaderStages.size());
	CreateInfo.pStages = ShaderStages.data();

	// setup group info for RG, MissG and HG
	VkRayTracingShaderGroupCreateInfoKHR RGGroupInfo{};
	RGGroupInfo.sType = VK_STRUCTURE_TYPE_RAY_TRACING_SHADER_GROUP_CREATE_INFO_KHR;
	RGGroupInfo.type = VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR;
	RGGroupInfo.closestHitShader = VK_SHADER_UNUSED_KHR;
	RGGroupInfo.anyHitShader = VK_SHADER_UNUSED_KHR;
	RGGroupInfo.intersectionShader = VK_SHADER_UNUSED_KHR;
	RGGroupInfo.generalShader = 0;

	VkRayTracingShaderGroupCreateInfoKHR MissGroupInfo{};
	MissGroupInfo.sType = VK_STRUCTURE_TYPE_RAY_TRACING_SHADER_GROUP_CREATE_INFO_KHR;
	MissGroupInfo.type = VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR;
	MissGroupInfo.closestHitShader = VK_SHADER_UNUSED_KHR;
	MissGroupInfo.anyHitShader = VK_SHADER_UNUSED_KHR;
	MissGroupInfo.intersectionShader = VK_SHADER_UNUSED_KHR;
	MissGroupInfo.generalShader = 1;

	// setup all hit groups
	std::vector<VkRayTracingShaderGroupCreateInfoKHR> HGGroupInfos(InInfo.AnyHitShaders.size());
	for (size_t Idx = 0; Idx < InInfo.AnyHitShaders.size(); Idx++)
	{
		VkRayTracingShaderGroupCreateInfoKHR HGGroupInfo{};
		HGGroupInfo.sType = VK_STRUCTURE_TYPE_RAY_TRACING_SHADER_GROUP_CREATE_INFO_KHR;
		HGGroupInfo.type = VK_RAY_TRACING_SHADER_GROUP_TYPE_TRIANGLES_HIT_GROUP_KHR;
		HGGroupInfo.closestHitShader = 2;
		HGGroupInfo.anyHitShader = static_cast<uint32_t>(3 + Idx);
		HGGroupInfo.intersectionShader = VK_SHADER_UNUSED_KHR;
		HGGroupInfo.generalShader = VK_SHADER_UNUSED_KHR;
		HGGroupInfos[Idx] = HGGroupInfo;
	}

	std::vector<VkRayTracingShaderGroupCreateInfoKHR> GroupInfos = { RGGroupInfo, MissGroupInfo };
	GroupInfos.insert(GroupInfos.end(), HGGroupInfos.begin(), HGGroupInfos.end());
	CreateInfo.groupCount = static_cast<uint32_t>(GroupInfos.size());
	CreateInfo.pGroups = GroupInfos.data();

	// set payload size
	VkRayTracingPipelineInterfaceCreateInfoKHR PipelineInterfaceInfo{};
	PipelineInterfaceInfo.sType = VK_STRUCTURE_TYPE_RAY_TRACING_PIPELINE_INTERFACE_CREATE_INFO_KHR;
	PipelineInterfaceInfo.maxPipelineRayPayloadSize = InInfo.PayloadSize;
	PipelineInterfaceInfo.maxPipelineRayHitAttributeSize = InInfo.AttributeSize;
	CreateInfo.pLibraryInterface = &PipelineInterfaceInfo;

	// create state for ray tracing pipeline
	PFN_vkCreateRayTracingPipelinesKHR CreateRTPipeline = 
		(PFN_vkCreateRayTracingPipelinesKHR)vkGetInstanceProcAddr(VulkanInstance, "vkCreateRayTracingPipelinesKHR");
	
	VkResult Result = CreateRTPipeline(LogicalDevice, VK_NULL_HANDLE, VK_NULL_HANDLE, 1, &CreateInfo, nullptr, &RTPipeline);
	if (Result != VK_SUCCESS)
	{
		UHE_LOG(L"Failed to create ray tracing pipeline!\n");
		return false;
	}

	return true;

The key is to set proper index for the generalShader, closestHitShader, anyHitShader variable. For now I'm reusing the closest hit shader so the value is fixed. I'll refactor that if necessary.

Upload Shader Record of Hit Group

struct UHShaderRecord
{
	BYTE Record[32];
};

void UHShaderClass::InitHitGroupTable(size_t NumMaterials)
{
	std::vector<BYTE> TempData(Gfx->GetShaderRecordSize());

	// get data for HG as well
	PFN_vkGetRayTracingShaderGroupHandlesKHR GetRTShaderGroupHandle = (PFN_vkGetRayTracingShaderGroupHandlesKHR)
		vkGetInstanceProcAddr(Gfx->GetInstance(), "vkGetRayTracingShaderGroupHandlesKHR");

	// create HG buffer
	HitGroupTable = Gfx->RequestRenderBuffer<UHShaderRecord>();
	HitGroupTable->CreateBuffer(NumMaterials, VK_BUFFER_USAGE_SHADER_BINDING_TABLE_BIT_KHR | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT);

	for (size_t Idx = 0; Idx < NumMaterials; Idx++)
	{
		if (GetRTShaderGroupHandle(Gfx->GetLogicalDevice(), RTState->GetRTPipeline(), 2 + static_cast<uint32_t>(Idx), 1, Gfx->GetShaderRecordSize(), TempData.data()) != VK_SUCCESS)
		{
			UHE_LOG(L"Failed to get hit group handle!\n");
			continue;
		}

		// copy HG records to the buffer, both closest hit and any hit will be put in the same hit group
		HitGroupTable->UploadData(TempData.data(), Idx);
	}
}

The shader record identifier is guaranteed to be 32 bytes based on specs. If you're still worrying it, try requesting the size from device won't hurt.

Note that the firstgroup parameter passed to vkGetRayTracingShaderGroupHandlesKHR function needs to be less than the VkRayTracingPipelineCreateInfoKHR::groupCount.

After uploading the shader records, it's ready for rendering use.

Using Shader Records in Rendering

	VkStridedDeviceAddressRegionKHR RayGenAddress{};
	RayGenAddress.deviceAddress = GetDeviceAddress(LogicalDevice, InRayGenTable->GetBuffer());
	RayGenAddress.size = InRayGenTable->GetBufferSize();
	RayGenAddress.stride = RayGenAddress.size;

	VkStridedDeviceAddressRegionKHR MissAddress{};
	MissAddress.deviceAddress = GetDeviceAddress(LogicalDevice, InMissTable->GetBuffer());
	MissAddress.size = InMissTable->GetBufferSize();
	MissAddress.stride = MissAddress.size;

	VkStridedDeviceAddressRegionKHR HitGroupAddress{};
	HitGroupAddress.deviceAddress = GetDeviceAddress(LogicalDevice, InHitGroupTable->GetBuffer());
	HitGroupAddress.size = InHitGroupTable->GetBufferSize();
	HitGroupAddress.stride = InHitGroupTable->GetBufferStride();

	VkStridedDeviceAddressRegionKHR NullAddress{};

	PFN_vkCmdTraceRaysKHR TraceRays = (PFN_vkCmdTraceRaysKHR)vkGetInstanceProcAddr(Gfx->GetInstance(), "vkCmdTraceRaysKHR");
	TraceRays(CmdList, &RayGenAddress, &MissAddress, &HitGroupAddress, &NullAddress, InExtent.width, InExtent.height, 1);

Most easiest step I can say.

§Summary

Without the support of local descriptor, we can still figure out a way to pass material data into ray tracing shader in Vulkan.

By combining SBT and bindless indexing, people should deal with most cases without issue.

Of course, the way to pass material data can vary based on your application ?

GitHub Link:

https://github.com/EasyJellySniper/Unheard-Engine

0 likes 0 comments

Comments

Nobody has left a comment. You can be the first!
You must log in to join the conversation.
Don't have a GameDev.net account? Sign up!
Advertisement