// Copyright Epic Games, Inc. All Rights Reserved.

/*=============================================================================
	VulkanCommon.ush: Common Vulkan shader code
=============================================================================*/

#pragma once

// Update this GUID to invalidate shader recompilation for only Vulkan shaders
#pragma message("UESHADERMETADATA_VERSION 3A87C7A8-C46C-467E-ACA3-B808491C3B15")

#define COMPILER_SUPPORTS_ULONG_TYPES 1
#define COMPILER_SUPPORTS_UINT64_IMAGE_ATOMICS (1 && COMPILER_SUPPORTS_ULONG_TYPES)

#if COMPILER_SUPPORTS_ULONG_TYPES

#define UlongType uint64_t

UlongType PackUlongType(uint2 Value)
{
	return ((UlongType)Value.y << 32) | Value.x;
}

uint2 UnpackUlongType(UlongType Value)
{
	return uint2(uint(Value), uint(Value >> 32));
}

#endif

#if COMPILER_SUPPORTS_UINT64_IMAGE_ATOMICS
void ImageInterlockedMaxUInt64(RWTexture2D<UlongType> Dest, uint2 Coord, UlongType Value)
{
	InterlockedMax(Dest[Coord], Value);
}

void ImageInterlockedAddUInt64(RWTexture2D<UlongType> Dest, uint2 Coord, UlongType Value)
{
	InterlockedAdd(Dest[Coord], Value);
}
#endif

#if UE_BINDLESS_ENABLED
	// While we use descriptor indexing, declare unbounded resource arrays
	#define VULKAN_HEAP_NAME_INNER(Prefix, Type) Prefix##Type
	#define VULKAN_HEAP_NAME(Prefix, Type) VULKAN_HEAP_NAME_INNER(Prefix, Type)
	#define VULKAN_SAMPLER_HEAP(Type) VULKAN_HEAP_NAME(BINDLESS_SAMPLER_ARRAY_PREFIX,Type)
	#define VULKAN_SRV_HEAP(Type) VULKAN_HEAP_NAME(BINDLESS_SRV_ARRAY_PREFIX,Type)
	#define VULKAN_UAV_HEAP(Type) VULKAN_HEAP_NAME(BINDLESS_UAV_ARRAY_PREFIX,Type)

	#if (RAYCALLABLESHADER || RAYHITGROUPSHADER || RAYGENSHADER || RAYMISSSHADER)
		#define VULKAN_HEAP_ACCESS(BindlessIndex) NonUniformResourceIndex(BindlessIndex)
	#else
		#define VULKAN_HEAP_ACCESS(BindlessIndex) BindlessIndex
	#endif

	#define UE_BINDLESS_SRV_HEAP_DECL(Type) Type VULKAN_SRV_HEAP(Type)[];
	#define UE_BINDLESS_UAV_HEAP_DECL(Type) Type VULKAN_UAV_HEAP(Type)[];
	#define UE_BINDLESS_SAMPLER_HEAP_DECL(Type) Type VULKAN_SAMPLER_HEAP(Type)[];

	#define GetSRVFromHeap(Type, Index) VULKAN_SRV_HEAP(Type)[VULKAN_HEAP_ACCESS(Index)]
	#define GetUAVFromHeap(Type, Index) VULKAN_UAV_HEAP(Type)[VULKAN_HEAP_ACCESS(Index)]
	#define GetSamplerFromHeap(Type, Index) VULKAN_SAMPLER_HEAP(Type)[VULKAN_HEAP_ACCESS(Index)]
#endif

#if VULKAN_SUPPORTS_SUBGROUP_SIZE_CONTROL
	#define COMPILER_SUPPORTS_WAVE_SIZE 1

	// NOTE: Wave size in Vulkan is supported at shader stage in the pipeline creation.
	// Create something unique we can read back from the shader, but that goes away on the rewrite
	COMPILER_DEFINE VULKAN_WAVESIZE(N)
	#define WAVESIZE(N) VULKAN_WAVESIZE(N)
#endif

#if UE_BINDLESS_ENABLED && (RAYCALLABLESHADER || RAYHITGROUPSHADER || RAYGENSHADER || RAYMISSSHADER)
	#define OVERRIDE_RAY_TRACING_HIT_GROUP_SYSTEM_RESOURCES 1
	#include "/Engine/Shared/RayTracingBuiltInResources.h"

	struct VulkanShaderRecordLooseParameterData
	{
		uint VulkanShaderRecordDummyGlobals;
	};

	// Combine the common FHitGroupSystemRootConstants with bindless parameters
	// :todo-jn: must match with VulkanRayTracing.h decl until we create a common header like RayTracingBuiltInResources.h
	struct FVulkanHitGroupSystemParameters
	{
		FHitGroupSystemRootConstants RootConstants;

		uint BindlessHitGroupSystemIndexBuffer;
		uint BindlessHitGroupSystemVertexBuffer;

		uint BindlessUniformBuffers[32];

		// Globals must be kept at the end of FVulkanHitGroupSystemParameters
		VulkanShaderRecordLooseParameterData Globals;
	};

	[[vk::shader_record_ext]] ConstantBuffer<FVulkanHitGroupSystemParameters> VulkanHitGroupSystemParameters;

	static const FHitGroupSystemRootConstants HitGroupSystemRootConstants = VulkanHitGroupSystemParameters.RootConstants;

	// Will get filled in with the bindless index in VulkanHitGroupSystemParameters.BindlessHitGroupSystemIndexBuffer
	ByteAddressBuffer HitGroupSystemIndexBuffer;
	// Will get filled in with the bindless index in VulkanHitGroupSystemParameters.BindlessHitGroupSystemVertexBuffer
	ByteAddressBuffer HitGroupSystemVertexBuffer;
#endif

#if PLATFORM_SUPPORTS_MESH_SHADERS_TIER0
	#define MESH_SHADER_TRIANGLE_ATTRIBUTES(InThreadsX) [numthreads(InThreadsX, 1, 1)][outputtopology("triangle")]
	#define MESH_SHADER_LINE_ATTRIBUTES(InThreadsX) [numthreads(InThreadsX, 1, 1)][outputtopology("line")]

	#define AMPLIFICATION_SHADER_TRIANGLE_ATTRIBUTES(InThreadsX) [numthreads(InThreadsX, 1, 1)][outputtopology("triangle")]
	#define AMPLIFICATION_SHADER_LINE_ATTRIBUTES(InThreadsX) [numthreads(InThreadsX, 1, 1)][outputtopology("line")]

	#define MESH_SHADER_VERTEX_EXPORT(VertexType, InMaxVertexCount) out vertices VertexType OutVertices[InMaxVertexCount]
	#define MESH_SHADER_TRIANGLE_EXPORT(InMaxTriangleCount) out indices uint3 OutTriangles[InMaxTriangleCount]
	#define MESH_SHADER_LINE_EXPORT(InMaxLineCount) out indices uint2 OutLines[InMaxLineCount]

	#define MESH_SHADER_WRITE_VERTEX(InVertexIndex, Value) OutVertices[InVertexIndex] = Value;
	#define MESH_SHADER_WRITE_TRIANGLE(InTriangleIndex, Value) OutTriangles[InTriangleIndex] = Value;
	#define MESH_SHADER_WRITE_LINE(InLineIndex, Value) OutLines[InLineIndex] = Value;
#endif

#if PLATFORM_SUPPORTS_MESH_SHADERS_TIER1
	#define MESH_SHADER_PRIMITIVE_EXPORT(PrimitiveType, InMaxPrimitiveCount) out primitives PrimitiveType OutPrimitives[InMaxPrimitiveCount]

	#define MESH_SHADER_WRITE_PRIMITIVE(InPrimitiveIndex, Value) OutPrimitives[InPrimitiveIndex] = Value;
#endif


// Types of Global Samplers (see Common.ush for types)
// Must match GetGlobalSamplerType() in SpirVShaderCompiler.inl
// and EGlobalSamplerType in VulkanShaderResources.h

#if UE_BINDLESS_ENABLED
	typedef SamplerState SafeTypeGlobalSamplers;
	SafeTypeGlobalSamplers VULKAN_SAMPLER_HEAP(GlobalSamplers)[];
	static const SamplerState VulkanGlobalPointClampedSampler     = VULKAN_SAMPLER_HEAP(GlobalSamplers)[1];
	static const SamplerState VulkanGlobalPointWrappedSampler     = VULKAN_SAMPLER_HEAP(GlobalSamplers)[2];
	static const SamplerState VulkanGlobalBilinearClampedSampler  = VULKAN_SAMPLER_HEAP(GlobalSamplers)[3];
	static const SamplerState VulkanGlobalBilinearWrappedSampler  = VULKAN_SAMPLER_HEAP(GlobalSamplers)[4];
	static const SamplerState VulkanGlobalTrilinearClampedSampler = VULKAN_SAMPLER_HEAP(GlobalSamplers)[5];
	static const SamplerState VulkanGlobalTrilinearWrappedSampler = VULKAN_SAMPLER_HEAP(GlobalSamplers)[6];
#else
	SamplerState VulkanGlobalPointWrappedSampler;
	SamplerState VulkanGlobalPointClampedSampler;
	SamplerState VulkanGlobalBilinearWrappedSampler;
	SamplerState VulkanGlobalBilinearClampedSampler;
	SamplerState VulkanGlobalTrilinearWrappedSampler;
	SamplerState VulkanGlobalTrilinearClampedSampler;
#endif

#define GetGlobalSampler(Filter, WrapMode) VulkanGlobal##Filter##WrapMode##Sampler