o3de/Gems/Meshlets/Assets/Shaders/MeshletsDebugRenderShader.azsl

/*
* Modifications Copyright (c) Contributors to the Open 3D Engine Project. 
* For complete copyright and license terms please see the LICENSE at the root of this distribution.
* 
* SPDX-License-Identifier: Apache-2.0 OR MIT
*
*/

// To use the object Id and VertexUtility the ObjectSrg must be created 
// and passed with the Object Id here.
// This will require more than several changes also on the CPU side, utilizing the 
// transform (m_transformService->ReserveObjectId()) and setting the Id in the 
// PerObjectSrg (objectSrg->SetConstant(objectIdIndex, m_objectId.GetIndex())).
// Look at ModelDataInstance::Init(Data::Instance<RPI::Model> model) to see how
// to set the model Id - this will lead to MeshDrawPacket::DoUpdate that will demonstrate 
// how to set the Srgs.
// There is no need to set all the Srgs in the DrawPacket - only the ones that the pass
// does not add manually in the command list (check the inheritance of the commit method).
//#include <Atom/Features/PBR/DefaultObjectSrg.azsli>
#include <Atom/Features/SrgSemantics.azsli>

//------------------------------------------------------------------------------
//ShaderResourceGroup PassSrg : SRG_PerPass_WithFallback
ShaderResourceGroup PassSrg : SRG_PerPass
{   // No need for StructuredBuffer - change to buffer across all
    StructuredBuffer<int>     m_meshletsSharedBuffer;
}

// Vertex Assembly streams will NOT be used - it is shown below as reference
// Instead we use the buffers directly with index offset into the data.
struct VSInput
{
    uint m_vertexIndex : SV_VertexID;

/*  Use the following if ever passed as vertex assembly streams
    float3 m_position : POSITION;
    float3 m_normal : NORMAL;
    float4 m_tangent : TANGENT; 
    float3 m_bitangent : BITANGENT; 
    float2 m_uv : UV0;
*/
};

struct VSOutput
{
    float4 m_position : SV_Position;
    float3 m_normal: NORMAL;
    float3 m_tangent : TANGENT; 
    float3 m_bitangent : BITANGENT; 
    float3 m_worldPosition : UV0;
    float2 m_uv : UV1;
};

struct MeshletsPassOutput
{
    float4 m_color  : SV_Target0;
};

struct MeshletsPassOutputWithDepth
{
    float4 m_color  : SV_Target0;
    float m_depth : SV_Depth;
};

// Shader Resource Groups
#include <viewsrg.srgi>
#include <MeshletsPerObjectRenderSrg.azsli>

#include <Atom/RPI/Math.azsli>

//! Utility function for vertex shaders to transform vertex tangent, bitangent, and normal vectors into world space.
void ConstructTBN(float3 vertexNormal, float4 vertexTangent, float3 vertexBitangent, 
    float4x4 localToWorld, float3x3 localToWorldInverseTranspose, 
    out float3 normalWS, out float3 tangentWS, out float3 bitangentWS)
{
    normalWS = normalize(mul(localToWorldInverseTranspose, vertexNormal));
    tangentWS = normalize(mul(localToWorld, float4(vertexTangent.xyz, 0)).xyz);
    bitangentWS = normalize(mul(localToWorld, float4(vertexBitangent, 0)).xyz);
}

//! @param skipShadowCoords can be useful for example when PixelDepthOffset is enable, because the pixel shader will have to run before the final world position is known
void VertexHelper(uint vertexIndex, inout VSOutput OUT, float3 worldPosition)
{
    OUT.m_worldPosition = worldPosition;

    // This paragraph is temporary to force the shader to include the Srg.
    float dummyVariableToForceSrg = 0.000001 * (float)PassSrg::m_meshletsSharedBuffer[vertexIndex];
    OUT.m_worldPosition.z += (abs(dummyVariableToForceSrg) < 0.001) ? dummyVariableToForceSrg : 0;

    OUT.m_position = mul(ViewSrg::m_viewProjectionMatrix, float4(OUT.m_worldPosition, 1.0));

    float4x4 objectToWorld = MeshletsObjectRenderSrg::GetWorldMatrix();
    float3x3 objectToWorldIT = MeshletsObjectRenderSrg::GetWorldMatrixInverseTranspose();

    ConstructTBN( 
        MeshletsObjectRenderSrg::GetNormal(vertexIndex),  
        MeshletsObjectRenderSrg::GetTangent(vertexIndex), 
        MeshletsObjectRenderSrg::GetBiTangent(vertexIndex), 
        objectToWorld, objectToWorldIT, 
        OUT.m_normal, OUT.m_tangent, OUT.m_bitangent);
}

VSOutput MeshletsDebugRender_MainPassVS(VSInput IN)
{
    VSOutput OUT;
 
    // Idea - test vertexIndex by sending to render as a position.
    float3 localPosition = MeshletsObjectRenderSrg::GetPosition(IN.m_vertexIndex);
    float3 worldPosition = mul(MeshletsObjectRenderSrg::GetWorldMatrix(), float4(localPosition, 1.0)).xyz;
 
    VertexHelper(IN.m_vertexIndex, OUT, worldPosition);

    OUT.m_uv = MeshletsObjectRenderSrg::GetUV(IN.m_vertexIndex);

    return OUT;
}

MeshletsPassOutput MeshletsDebugRender_MainPassPS(VSOutput IN)
{
    // ------- Output -------
    MeshletsPassOutput OUT;

    // Debug purposes - uv will be colored by the Compute that 
    // will generate the meshlets indices.
    OUT.m_color = float4(IN.m_uv.rg, 0, 1);

    return OUT;
}