Skip to content

Commit f4e5bff

Browse files
JoseEmilio-ARMJoseEmilio-ARM
committed
Add transparency to shadows
1 parent aa553f2 commit f4e5bff

File tree

2 files changed

+167
-4
lines changed

2 files changed

+167
-4
lines changed

‎attachments/38_ray_tracing.cpp‎

Lines changed: 114 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,8 @@ class HelloTriangleApplication {
142142
vk::raii::DeviceMemory vertexBufferMemory = nullptr;
143143
vk::raii::Buffer indexBuffer = nullptr;
144144
vk::raii::DeviceMemory indexBufferMemory = nullptr;
145+
vk::raii::Buffer uvBuffer = nullptr;
146+
vk::raii::DeviceMemory uvBufferMemory = nullptr;
145147

146148
std::vector<vk::raii::Buffer> blasBuffers;
147149
std::vector<vk::raii::DeviceMemory> blasMemories;
@@ -157,6 +159,14 @@ class HelloTriangleApplication {
157159
vk::raii::DeviceMemory tlasScratchMemory = nullptr;
158160
vk::raii::AccelerationStructureKHR tlas = nullptr;
159161

162+
struct InstanceLUT {
163+
uint32_t materialID;
164+
uint32_t indexBufferOffset;
165+
};
166+
std::vector<InstanceLUT> instanceLUTs;
167+
vk::raii::Buffer instanceLUTBuffer = nullptr;
168+
vk::raii::DeviceMemory instanceLUTBufferMemory = nullptr;
169+
160170
UniformBufferObject ubo{};
161171

162172
std::vector<vk::raii::Buffer> uniformBuffers;
@@ -232,7 +242,9 @@ class HelloTriangleApplication {
232242
createTextureSampler();
233243
createVertexBuffer();
234244
createIndexBuffer();
245+
createUVBuffer();
235246
createAccelerationStructures();
247+
createInstanceLUTBuffer();
236248
createUniformBuffers();
237249
createDescriptorPool();
238250
createDescriptorSets();
@@ -514,6 +526,9 @@ class HelloTriangleApplication {
514526
std::array global_bindings = {
515527
vk::DescriptorSetLayoutBinding( 0, vk::DescriptorType::eUniformBuffer, 1, vk::ShaderStageFlagBits::eVertex | vk::ShaderStageFlagBits::eFragment, nullptr),
516528
vk::DescriptorSetLayoutBinding( 1, vk::DescriptorType::eAccelerationStructureKHR, 1, vk::ShaderStageFlagBits::eFragment, nullptr),
529+
vk::DescriptorSetLayoutBinding( 2, vk::DescriptorType::eStorageBuffer, 1, vk::ShaderStageFlagBits::eFragment, nullptr),
530+
vk::DescriptorSetLayoutBinding( 3, vk::DescriptorType::eStorageBuffer, 1, vk::ShaderStageFlagBits::eFragment, nullptr),
531+
vk::DescriptorSetLayoutBinding( 4, vk::DescriptorType::eStorageBuffer, 1, vk::ShaderStageFlagBits::eFragment, nullptr)
517532
};
518533

519534
vk::DescriptorSetLayoutCreateInfo globalLayoutInfo{ .bindingCount = static_cast<uint32_t>(global_bindings.size()), .pBindings = global_bindings.data() };
@@ -957,6 +972,47 @@ class HelloTriangleApplication {
957972
copyBuffer(stagingBuffer, indexBuffer, bufferSize);
958973
}
959974

975+
void createUVBuffer() {
976+
// Extract all texCoords into a separate vector
977+
std::vector<glm::vec2> uvs;
978+
uvs.reserve(vertices.size());
979+
for (auto& v: vertices) {
980+
uvs.push_back(v.texCoord);
981+
}
982+
983+
vk::DeviceSize bufferSize = sizeof(uvs[0]) * uvs.size();
984+
985+
vk::raii::Buffer stagingBuffer({});
986+
vk::raii::DeviceMemory stagingBufferMemory({});
987+
createBuffer(bufferSize, vk::BufferUsageFlagBits::eTransferSrc, vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent, stagingBuffer, stagingBufferMemory);
988+
989+
void* dataStaging = stagingBufferMemory.mapMemory(0, bufferSize);
990+
memcpy(dataStaging, uvs.data(), bufferSize);
991+
stagingBufferMemory.unmapMemory();
992+
993+
createBuffer(bufferSize, vk::BufferUsageFlagBits::eTransferDst | vk::BufferUsageFlagBits::eStorageBuffer,
994+
vk::MemoryPropertyFlagBits::eDeviceLocal, uvBuffer, uvBufferMemory);
995+
996+
copyBuffer(stagingBuffer, uvBuffer, bufferSize);
997+
}
998+
999+
void createInstanceLUTBuffer() {
1000+
vk::DeviceSize bufferSize = sizeof(InstanceLUT) * instanceLUTs.size();
1001+
1002+
vk::raii::Buffer stagingBuffer({});
1003+
vk::raii::DeviceMemory stagingBufferMemory({});
1004+
createBuffer(bufferSize, vk::BufferUsageFlagBits::eTransferSrc, vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent, stagingBuffer, stagingBufferMemory);
1005+
1006+
void* dataStaging = stagingBufferMemory.mapMemory(0, bufferSize);
1007+
memcpy(dataStaging, instanceLUTs.data(), bufferSize);
1008+
stagingBufferMemory.unmapMemory();
1009+
1010+
createBuffer(bufferSize, vk::BufferUsageFlagBits::eTransferDst | vk::BufferUsageFlagBits::eStorageBuffer,
1011+
vk::MemoryPropertyFlagBits::eDeviceLocal, instanceLUTBuffer, instanceLUTBufferMemory);
1012+
1013+
copyBuffer(stagingBuffer, instanceLUTBuffer, bufferSize);
1014+
}
1015+
9601016
void createUniformBuffers() {
9611017
uniformBuffers.clear();
9621018
uniformBuffersMemory.clear();
@@ -1008,10 +1064,14 @@ class HelloTriangleApplication {
10081064
vk::AccelerationStructureGeometryDataKHR geomData(trianglesData);
10091065
vk::AccelerationStructureGeometryKHR blasGeometry{
10101066
.geometryType = vk::GeometryTypeKHR::eTriangles,
1011-
.geometry = geomData,
1012-
.flags = vk::GeometryFlagBitsKHR::eOpaque
1067+
.geometry = geomData
10131068
};
10141069

1070+
if (!submesh.alphaCut)
1071+
{
1072+
blasGeometry.flags = vk::GeometryFlagBitsKHR::eOpaque;
1073+
}
1074+
10151075
vk::AccelerationStructureBuildRangeInfoKHR blasRangeInfo{
10161076
.primitiveCount = static_cast<uint32_t>(submesh.indexCount / 3),
10171077
.primitiveOffset = 0,
@@ -1080,11 +1140,13 @@ class HelloTriangleApplication {
10801140

10811141
vk::AccelerationStructureInstanceKHR instance{};
10821142
instance.setTransform(tm)
1143+
.setInstanceCustomIndex(static_cast<uint32_t>(i)) // Used to retrieve intersection information from instance LUT
10831144
.setMask(0xFF)
10841145
.setAccelerationStructureReference(blasDeviceAddr)
10851146
.setFlags(vk::GeometryInstanceFlagBitsKHR::eTriangleFacingCullDisable);
10861147

10871148
instances.push_back(instance);
1149+
instanceLUTs.push_back({ static_cast<uint32_t>(submesh.materialID), submesh.indexOffset });
10881150
}
10891151

10901152
// Prepare instance data for the TLAS
@@ -1272,6 +1334,7 @@ class HelloTriangleApplication {
12721334
std::array poolSize {
12731335
vk::DescriptorPoolSize( vk::DescriptorType::eUniformBuffer, MAX_FRAMES_IN_FLIGHT),
12741336
vk::DescriptorPoolSize( vk::DescriptorType::eAccelerationStructureKHR, MAX_FRAMES_IN_FLIGHT),
1337+
vk::DescriptorPoolSize( vk::DescriptorType::eStorageBuffer, MAX_FRAMES_IN_FLIGHT * 3), // indices, UVs, instance LUT
12751338
vk::DescriptorPoolSize( vk::DescriptorType::eSampler, MAX_FRAMES_IN_FLIGHT),
12761339
vk::DescriptorPoolSize( vk::DescriptorType::eSampledImage, (uint32_t)materials.size())
12771340
};
@@ -1330,7 +1393,55 @@ class HelloTriangleApplication {
13301393
.descriptorType = vk::DescriptorType::eAccelerationStructureKHR
13311394
};
13321395

1333-
std::array<vk::WriteDescriptorSet, 2> descriptorWrites{bufferWrite, asWrite};
1396+
// Indices SSBO
1397+
vk::DescriptorBufferInfo indexBufferInfo{
1398+
.buffer = indexBuffer,
1399+
.offset = 0,
1400+
.range = sizeof(uint32_t) * indices.size()
1401+
};
1402+
1403+
vk::WriteDescriptorSet indexBufferWrite{
1404+
.dstSet = globalDescriptorSets[i],
1405+
.dstBinding = 2,
1406+
.dstArrayElement = 0,
1407+
.descriptorCount = 1,
1408+
.descriptorType = vk::DescriptorType::eStorageBuffer,
1409+
.pBufferInfo = &indexBufferInfo
1410+
};
1411+
1412+
// UVs SSBO
1413+
vk::DescriptorBufferInfo uvBufferInfo{
1414+
.buffer = uvBuffer,
1415+
.offset = 0,
1416+
.range = sizeof(glm::vec2) * vertices.size()
1417+
};
1418+
1419+
vk::WriteDescriptorSet uvBufferWrite{
1420+
.dstSet = globalDescriptorSets[i],
1421+
.dstBinding = 3,
1422+
.dstArrayElement = 0,
1423+
.descriptorCount = 1,
1424+
.descriptorType = vk::DescriptorType::eStorageBuffer,
1425+
.pBufferInfo = &uvBufferInfo
1426+
};
1427+
1428+
// Instance LUT SSBO
1429+
vk::DescriptorBufferInfo instanceLUTBufferInfo{
1430+
.buffer = instanceLUTBuffer,
1431+
.offset = 0,
1432+
.range = sizeof(InstanceLUT) * instanceLUTs.size()
1433+
};
1434+
1435+
vk::WriteDescriptorSet instanceLUTBufferWrite{
1436+
.dstSet = globalDescriptorSets[i],
1437+
.dstBinding = 4,
1438+
.dstArrayElement = 0,
1439+
.descriptorCount = 1,
1440+
.descriptorType = vk::DescriptorType::eStorageBuffer,
1441+
.pBufferInfo = &instanceLUTBufferInfo
1442+
};
1443+
1444+
std::array<vk::WriteDescriptorSet, 5> descriptorWrites{bufferWrite, asWrite, indexBufferWrite, uvBufferWrite, instanceLUTBufferWrite};
13341445

13351446
device.updateDescriptorSets(descriptorWrites, {});
13361447
}

‎attachments/38_ray_tracing.slang‎

Lines changed: 53 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,19 @@ ConstantBuffer<UniformBuffer> ubo;
1616
[[vk::binding(1,0)]]
1717
RaytracingAccelerationStructure accelerationStructure;
1818

19+
[[vk::binding(2,0)]]
20+
StructuredBuffer<uint> indexBuffer;
21+
22+
[[vk::binding(3,0)]]
23+
StructuredBuffer<float2> uvBuffer;
24+
25+
struct InstanceLUT {
26+
uint materialID;
27+
uint indexBufferOffset;
28+
};
29+
[[vk::binding(4,0)]]
30+
StructuredBuffer<InstanceLUT> instanceLUTBuffer;
31+
1932
struct VSOutput
2033
{
2134
float4 pos : SV_Position;
@@ -53,6 +66,24 @@ static const float3 lightDir = float3(-6.0, 0.0, 6.0);
5366
// Small epsilon to avoid self-intersection
5467
static const float EPSILON = 0.01;
5568

69+
float2 intersection_uv(uint instanceID, uint primIndex, float2 barycentrics) {
70+
uint indexOffset = instanceLUTBuffer[instanceID].indexBufferOffset;
71+
72+
uint i0 = indexBuffer[indexOffset + (primIndex * 3 + 0)];
73+
uint i1 = indexBuffer[indexOffset + (primIndex * 3 + 1)];
74+
uint i2 = indexBuffer[indexOffset + (primIndex * 3 + 2)];
75+
76+
float2 uv0 = uvBuffer[i0];
77+
float2 uv1 = uvBuffer[i1];
78+
float2 uv2 = uvBuffer[i2];
79+
80+
float w0 = 1.0 - barycentrics.x - barycentrics.y;
81+
float w1 = barycentrics.x;
82+
float w2 = barycentrics.y;
83+
84+
return w0 * uv0 + w1 * uv1 + w2 * uv2;
85+
}
86+
5687
bool in_shadow(float3 P)
5788
{
5889
// Build the shadow ray from the world position toward the light
@@ -70,7 +101,25 @@ bool in_shadow(float3 P)
70101

71102
sq.TraceRayInline(accelerationStructure, rayFlags, 0xFF, shadowRayDesc);
72103

73-
sq.Proceed();
104+
while (sq.Proceed())
105+
{
106+
uint instanceID = sq.CandidateInstanceID();
107+
uint primIndex = sq.CandidatePrimitiveIndex();
108+
109+
float2 uv = intersection_uv(instanceID, primIndex, sq.CandidateTriangleBarycentrics());
110+
111+
uint materialID = instanceLUTBuffer[instanceID].materialID;
112+
float4 intersection_color = textures[materialID].Sample(textureSampler, uv);
113+
114+
if (intersection_color.a < 0.5) {
115+
// If the triangle is transparent, we continue to trace
116+
// to find the next opaque triangle.
117+
} else {
118+
// If we hit an opaque triangle, we stop tracing.
119+
sq.CommitNonOpaqueTriangleHit();
120+
}
121+
122+
}
74123

75124
// If the shadow ray intersects an opaque triangle, we consider the pixel in shadow
76125
bool hit = (sq.CommittedStatus() == COMMITTED_TRIANGLE_HIT);
@@ -82,6 +131,9 @@ bool in_shadow(float3 P)
82131
float4 fragMain(VSOutput vertIn) : SV_TARGET {
83132
float4 baseColor = textures[pc.materialIndex].Sample(textureSampler, vertIn.fragTexCoord);
84133

134+
// Alpha test
135+
if (baseColor.a < 0.5) discard;
136+
85137
bool inShadow = in_shadow(vertIn.worldPos);
86138

87139
// Darken if in shadow

0 commit comments

Comments
 (0)