Is the intersection shader available in unreal?

Bender2016 · December 8, 2021, 7:38am

I am currently trying to draw some procedural geometry using ray tracing’s intersection shader. After many tests it still failed. I post some of my key codes below. I hope someone can help me check how to do it.

//AABB bounding box
class FUnitBoxAABB : public FRenderResource
{
public:
	FUnitBoxAABB() {}
	virtual ~FUnitBoxAABB() {}

	// FRenderResource interface.
	virtual void InitRHI() override;

	virtual void ReleaseRHI() override;

	virtual FString GetFriendlyName() const override
	{
		return TEXT("FUnitBoxAABB");
	}

	FRayTracingGeometry RayTracingGeometry;


private:
	FVertexBufferRHIRef VertexData;
};

void FUnitBoxAABB::InitRHI()
{
	if (IsRayTracingEnabled())
	{
		FRHIResourceCreateInfo CreateInfo;

		static const FVector Vertices[2] =
		{
			FVector(-1.0f, -1.0f, -1.0f),
			FVector(1.0f, 1.0f, 1.0f)
		};

		CreateInfo.DebugName = TEXT("UnitBoxAABB");
		VertexData = RHICreateVertexBuffer(sizeof(float) * 6, BUF_Static, CreateInfo);
		void* VoidPtr = RHILockVertexBuffer(VertexData, 0, sizeof(FVector) * 2, RLM_WriteOnly);
		FMemory::Memcpy(VoidPtr, Vertices, sizeof(FVector) * 2);
		RHIUnlockVertexBuffer(VertexData);

		FRayTracingGeometryInitializer Initializer;

		Initializer.IndexBuffer = nullptr;
		Initializer.TotalPrimitiveCount = 1; // one AABB
		Initializer.GeometryType = RTGT_Procedural;
		Initializer.bFastBuild = false;

		// Need explicit segment to mark no duplicate AnyHit
		FRayTracingGeometrySegment Segment;
		Segment.bForceOpaque = false;
		Segment.bAllowDuplicateAnyHitShaderInvocation = false;
		Segment.FirstPrimitive = 0;
		Segment.NumPrimitives = 1;
		Segment.VertexBuffer = VertexData;
		Segment.VertexBufferElementType = VET_Float3;
		Segment.VertexBufferStride = sizeof(FVector) * 2;
		Segment.VertexBufferOffset = 0;

		Initializer.Segments.Add(Segment);

		RayTracingGeometry.SetInitializer(Initializer);
		RayTracingGeometry.InitResource();
	}
}

FRayTracingScene RayTracingBVHScene;
FRayTracingScene RayTracingScene;
FRayTracingPipelineState* RayTracingPipelineState;

void ProceduralGeometryRayTracingRenderer::BuildRayTracingScene(FRHICommandListImmediate & RHICmdList)
{
	FRHIRayTracingGeometry* UnitBox = UnitBoxRayTracingGeometry.RayTracingGeometry.RayTracingGeometryRHI;
	
	{
		FProceduralGeometryState& Instance = ProceduralGeometryStates[StaticMeshIndex];
		
		int32 AABBIndex = RayTracingAABBInstances.AddDefaulted(1);
		FRayTracingGeometryInstance& AABBInstance = RayTracingAABBInstances[AABBIndex];
		AABBInstance.GeometryRHI = UnitBox;
		AABBInstance.bDoubleSided = false;
		AABBInstance.bForceOpaque = false;
		AABBInstance.Mask = 0xff;
		AABBInstance.UserData.Add(StaticMeshIndex);
		AABBInstance.NumTransforms = 1;

		FMatrix Transform = FMatrix::Identity;
		FVector ScaleFactor = ProceduralGeometryStates[StaticMeshIndex].Bound.GetSize() / 2.0f;
		Transform.M[0][0] *= ScaleFactor.X;
		Transform.M[1][1] *= ScaleFactor.Y;
		Transform.M[2][2] *= ScaleFactor.Z;
		AABBInstance.Transforms.Add(Transform.ConcatTranslation(ProceduralGeometryStates[StaticMeshIndex].Bound.GetCenter()));	
		
		FRayTracingSceneInitializer AABBInitializer;
		AABBInitializer.Instances = RayTracingAABBInstances;
		AABBInitializer.ShaderSlotsPerGeometrySegment = 1;
		AABBInitializer.NumMissShaderSlots = 1;

		RayTracingBVHScene.RayTracingSceneRHI = RHICreateRayTracingScene(AABBInitializer);
		RHICmdList.BuildAccelerationStructure(RayTracingBVHScene.RayTracingSceneRHI);
	}
	
	{
		FProceduralGeometryState& Instance = ProceduralGeometryStates[StaticMeshIndex];
		
		int32 MeshIndex = RayTracingMeshInstances.AddDefaulted(1);
		FRayTracingGeometryInstance& MeshInstance = RayTracingMeshInstances[MeshIndex];
		MeshInstance.GeometryRHI = Instance.RenderData->LODResources[0].RayTracingGeometry.RayTracingGeometryRHI;
		MeshInstance.Transforms.Add(Instance.LocalToWorld);
		MeshInstance.NumTransforms = 1;
		MeshInstance.UserData.Add((uint32)StaticMeshIndex);
		MeshInstance.Mask = 0xFF;
		MeshInstance.bForceOpaque = true;
		
		FRayTracingSceneInitializer MeshInitializer;
		MeshInitializer.Instances = RayTracingMeshInstances;
		MeshInitializer.ShaderSlotsPerGeometrySegment = 1;
		MeshInitializer.NumMissShaderSlots = 1;

		RayTracingScene.RayTracingSceneRHI = RHICreateRayTracingScene(MeshInitializer);
		RHICmdList.BuildAccelerationStructure(View->RayTracingScene.RayTracingSceneRHI);
	}
	
	
	FRayTracingPipelineStateInitializer Initializer;
	Initializer.MaxPayloadSizeInBytes = 24; 
	Initializer.bAllowHitGroupIndexing = false; 

	TShaderMapRef<FProceduralGeometryRGS> RayGenShader(GetGlobalShaderMap(View->GetFeatureLevel()));
	FRHIRayTracingShader* RayGenShaderTable[] = { RayGenShader.GetRayTracingShader() };
	Initializer.SetRayGenShaderTable(RayGenShaderTable);

	// Get the ray tracing materials
	auto HitGroupShader = View->ShaderMap->GetShader<FProceduralGeometryHitGroup>();
	FRHIRayTracingShader* HitShaderTable[] = { HitGroupShader.GetRayTracingShader() };
	Initializer.SetHitGroupTable(HitShaderTable);	

	auto MissShader = View->ShaderMap->GetShader<FProceduralGeometryMS>();
	FRHIRayTracingShader* MissShaderTable[] = { MissShader.GetRayTracingShader() };
	Initializer.SetMissShaderTable(MissShaderTable);

	RayTracingPipelineState = PipelineStateCache::GetAndOrCreateRayTracingPipelineState(RHICmdList, Initializer);
	
}


void ProceduralGeometryRayTracingRenderer::RenderRayTracingProceduralGeometry(FRHICommandListImmediate& RHICmdList)
{
	FRDGBuilder GraphBuilder(RHICmdList);
	
	FProceduralGeometryRGS::FParameters* PassParameters = GraphBuilder.AllocParameters<FProceduralGeometryRGS::FParameters>();
	//AABB cannot be draw
	PassParameters->TLAS = RayTracingBVHScene.RayTracingSceneRHI->GetShaderResourceView();
	//Staticmesh can be draw
	//PassParameters->TLAS = RayTracingScene.RayTracingSceneRHI->GetShaderResourceView();
	PassParameters->ViewUniformBuffer = View->ViewUniformBuffer;
	...
	PassParameters->ColorOutput = GraphBuilder.CreateUAV(FRDGTextureUAVDesc(RayTracingOutputTexture));

	auto RayGenShader = View->ShaderMap->GetShader<FProceduralGeometryRGS>();
	ClearUnusedGraphResources(RayGenShader, PassParameters)
	
	GraphBuilder.AddPass(
		RDG_EVENT_NAME("RayTracingProceduralGeometry %dx%d", Extent.X, Extent.Y),
		PassParameters,
		ERDGPassFlags::Compute,
		[PassParameters, this, RayGenShader, Extent](FRHICommandList& RHICmdList)
	{
		FRayTracingShaderBindingsWriter GlobalResources;
		SetShaderParameters(GlobalResources, RayGenShader, *PassParameters);		

		FRHIRayTracingScene* RayTracingSceneRHI = RayTracingScene.RayTracingSceneRHI;
		RHICmdList.RayTraceDispatch(RayTracingPipelineState, RayGenShader.GetRayTracingShader(), RayTracingSceneRHI, GlobalResources, Extent.X, Extent.Y);
	});

	TRefCountPtr<IPooledRenderTarget> PooledRenderTarget;
	GraphBuilder.QueueTextureExtraction(RayTracingOutputTexture, &PooledRenderTarget);
	GraphBuilder.Execute();
	
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
struct FProceduralGeometryAttributes
{
	float2 Normal;
};

struct FProceduralGeometryPayload : FMinimalPayload
{
	float NestingLevel;
};

[shader("intersection")]
void ProceduralGeometryIS()
{
	FProceduralGeometryAttributes Attributes = (FProceduralGeometryAttributes) 0;
	ReportHit(RayTMin(), 0, Attributes);
}

[shader("anyhit")]
void ProceduralGeometryAHS(inout FProceduralGeometryPayload Payload, in FProceduralGeometryAttributes Attributes)
{
	Payload.NestingLevel += 1.0;	
	IgnoreHit(); // need all the hits along the ray
}

[shader("closesthit")]
void ProceduralGeometryCHS(inout FProceduralGeometryPayload Payload, in FProceduralGeometryAttributes Attributes)
{
	//No-op, all work occurs in the AHS
	Payload.HitT = 1.0;
}
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
[shader("raygeneration")]
void ProceduralGeometryRGS()
{
	uint2 PixelCoord = DispatchRaysIndex().xy + View.ViewRectMin;
	float2 RenderTargetUV = (float2(PixelCoord) + .5f) * View.BufferSizeAndInvSize.zw;
	RayDesc Ray = CreatePrimaryRay(RenderTargetUV);
	
	uint RayFlags = 0;
	             // | RAY_FLAG_CULL_BACK_FACING_TRIANGLES // use back face culling
	             // | RAY_FLAG_ACCEPT_FIRST_HIT_AND_END_SEARCH // terminate immediately upon detecting primitive intersection
	             // | RAY_FLAG_FORCE_OPAQUE // don't run anyhit shader
	             // | RAY_FLAG_SKIP_CLOSEST_HIT_SHADER; // don't run closesthit shader
	const uint InstanceInclusionMask = RAY_TRACING_MASK_ALL;
	
	FProceduralGeometryPayload Payload = (FProceduralGeometryPayload) 0;
	
	TraceRay(
		TLAS, // AccelerationStructure
		RayFlags,
		0xff,
		0, // RayContributionToHitGroupIndex
		1, // MultiplierForGeometryContributionToShaderIndex
		0, // MissShaderIndex
		Ray, // RayDesc
		Payload // Payload
	);
	
	ColorOutput[PixelCoord] = Payload.IsHit() ? LightColor : BaseColor;
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
[shader("miss")]
void ProceduralGeometryMS(inout FProceduralGeometryPayload Payload)
{	
	//Payload.SetMiss();
	Payload.HitT = -1.0;
}

As the above code notes, it can draw static meshes but cannot draw AABB procedural geometry. I refer to NVIDIA’s NvRTX_Caustics branch.

If I replace RayTracingBVHScene with RayTracingScene, Then I can draw static meshes.
Thanks!

Bender2016 · December 9, 2021, 3:16am

Finally, I fixed the problem. In fact, the intersection shader works, but I don’t encode other shaders correctly. Now these shaders can work as expected. It can draw AABB correctly.