PathTracer: added PBR texture mapping

data/shaders/trace.glsl

@@ -22,12 +22,16 @@ layout(binding = 1) uniform config_ { PathTracerConfig u_config; };
 
 layout(binding = 2, rgba8) uniform image2D g_raster_image;
 
-layout(binding = 4) buffer buf04_ { uint g_indices[]; };
-layout(binding = 5) buffer buf05_ { float g_vertices[]; };
+layout(binding = 3) buffer buf03_ { uint g_indices[]; };
+layout(binding = 4) buffer buf04_ { float g_vertices[]; };
+layout(binding = 5) buffer buf05_ { vec2 g_tex_coords[]; };
 layout(binding = 6) uniform accelerationStructureEXT g_accelStruct;
 
-layout(binding = 10) uniform sampler2D opaque_texture;
-layout(binding = 11) uniform sampler2D transparent_texture;
+layout(binding = 10) uniform sampler2D albedo_tex;
+layout(binding = 11) uniform sampler2D normal_tex;
+layout(binding = 12) uniform sampler2D pbr_tex;
+layout(binding = 13) uniform sampler2D env_map;
+
 
 #include "%shader_debug"
 DEBUG_SETUP(1, 12)
@@ -45,12 +49,20 @@ vec3 getVertex(uint idx) {
 	return vec3(g_vertices[idx * 3 + 0], g_vertices[idx * 3 + 1], g_vertices[idx * 3 + 2]);
 }
 
-void getTriangleVertices(uint tri_id, out vec3 tri0, out vec3 tri1, out vec3 tri2) {
-	uint idx0 = g_indices[tri_id * 3 + 0], idx1 = g_indices[tri_id * 3 + 1],
-		 idx2 = g_indices[tri_id * 3 + 2];
-	tri0 = getVertex(idx0);
-	tri1 = getVertex(idx1);
-	tri2 = getVertex(idx2);
+uvec3 getTriangleIndices(uint tri_id) {
+	return uvec3(g_indices[tri_id * 3 + 0], g_indices[tri_id * 3 + 1], g_indices[tri_id * 3 + 2]);
+}
+
+void getTriangleVertices(uvec3 tri_indices, out vec3 tri0, out vec3 tri1, out vec3 tri2) {
+	tri0 = getVertex(tri_indices[0]);
+	tri1 = getVertex(tri_indices[1]);
+	tri2 = getVertex(tri_indices[2]);
+}
+
+void getTriangleTexCoords(uvec3 tri_indices, out vec2 tri0, out vec2 tri1, out vec2 tri2) {
+	tri0 = g_tex_coords[tri_indices[0]];
+	tri1 = g_tex_coords[tri_indices[1]];
+	tri2 = g_tex_coords[tri_indices[2]];
 }
 
 void getTriangleVectors(in vec3 tri0, in vec3 tri1, in vec3 tri2, out vec3 tangent, out vec3 normal,
@@ -62,6 +74,7 @@ void getTriangleVectors(in vec3 tri0, in vec3 tri1, in vec3 tri2, out vec3 tange
 
 struct TraceResult {
 	float dist;
+	vec2 barycentric;
 	uint num_iters;
 	uint tri_id;
 };
@@ -81,7 +94,7 @@ TraceResult rayTraceAS(vec3 origin, vec3 dir) {
 	if(rayQueryGetIntersectionTypeEXT(rq, true) != 0) {
 		result.dist = rayQueryGetIntersectionTEXT(rq, true);
 		result.tri_id = rayQueryGetIntersectionPrimitiveIndexEXT(rq, true);
-		// rayQueryGetIntersectionBarycentricsEXT
+		result.barycentric = rayQueryGetIntersectionBarycentricsEXT(rq, true);
 	} else {
 		result.dist = MAX_ISECT_DIST;
 		result.tri_id = INVALID_TRI_ID;
@@ -102,11 +115,22 @@ void getScreenRay(ivec2 pixel_pos, out vec3 origin, out vec3 dir) {
 	dir = u_config.frustum.ws_dir0.xyz + float(pixel_pos.x) * u_config.frustum.ws_dirx.xyz +
 		  float(pixel_pos.y) * u_config.frustum.ws_diry.xyz;
 	dir += vec3(0.0000001); // avoiding division by 0
+	dir = normalize(dir);
+}
+
+float randomFloat(inout uint rngState)
+{
+  // Condensed version of pcg_output_rxs_m_xs_32_32, with simple conversion to floating-point [0,1].
+  rngState  = rngState * 747796405 + 1;
+  uint word = ((rngState >> ((rngState >> 28) + 4)) ^ rngState) * 277803737;
+  word      = (word >> 22) ^ word;
+  return float(word) / 4294967295.0f;
 }
 
 float computeAO(uint tri_id, vec3 hit_point) {
+	uvec3 tri_indices = getTriangleIndices(tri_id);
 	vec3 tri[3];
-	getTriangleVertices(tri_id, tri[0], tri[1], tri[2]);
+	getTriangleVertices(tri_indices, tri[0], tri[1], tri[2]);
 	vec3 tri_vecs[3];
 	getTriangleVectors(tri[0], tri[1], tri[2], tri_vecs[0], tri_vecs[1], tri_vecs[2]);
 
@@ -127,24 +151,38 @@ float computeAO(uint tri_id, vec3 hit_point) {
 	return max(0.0, (total - hits) / float(total) - 0.1) * (1.0 / 0.9);
 }
 
+vec2 longLat(vec3 normal) {
+	// convert normal to longitude and latitude
+	float latitude = acos(normal.y) / PI;
+	float longitude = (atan(normal.x, normal.z) + PI) / (2.0 * PI);
+	return vec2(longitude, latitude);
+}
+
 void traceBin() {
 	ivec2 pixel_pos = ivec2(LIX & 31, LIX >> 5) + s_bin_pos;
+	uint random_seed = pixel_pos.x + (pixel_pos.y << 16);
+
 	vec3 ray_origin, ray_dir;
 	getScreenRay(pixel_pos, ray_origin, ray_dir);
 
 	TraceResult result = rayTraceAS(ray_origin, ray_dir);
-	float ao_value = 1.0;
-	if(result.dist < MAX_ISECT_DIST) {
-		vec3 hit_point = ray_origin + ray_dir * result.dist;
-		ao_value = computeAO(result.tri_id, hit_point);
-	}
 
 	vec3 vcolor = vec3(0.0);
 	if(result.dist < MAX_ISECT_DIST) {
-		float value = sqrt(result.dist);
-		vcolor = vec3(1.0 - value * 0.02, 1.0 - value * 0.05, 1.0 - value * 0.1);
+		vec2 uvs[3];
+		uvec3 tri_indices = getTriangleIndices(result.tri_id);
+		getTriangleTexCoords(tri_indices, uvs[0], uvs[1], uvs[2]);
+		vec3 bary = vec3(1.0 - result.barycentric[0] - result.barycentric[1], result.barycentric);
+		vec2 uv = uvs[0] * bary.x + uvs[1] * bary.y + uvs[2] * bary.z;
+		vcolor = texture(albedo_tex, uv).rgb;
+
+		//float value = sqrt(result.dist);
+		//vcolor = vec3(1.0 - value * 0.02, 1.0 - value * 0.05, 1.0 - value * 0.1);
+	}
+	else {
+		vec2 tex_coord = longLat(ray_dir);
+		vcolor = texture(env_map, tex_coord).rgb;
 	}
-	vcolor *= ao_value;
 
 	outputPixel(pixel_pos, SATURATE(vec4(vcolor, 1.0)));
 }

src/path_tracer.cpp

@@ -120,12 +120,15 @@ Ex<void> PathTracer::exConstruct(VulkanDevice &device, ShaderCompiler &compiler,
 Ex<> PathTracer::updateScene(VulkanDevice &device, Scene &scene) {
 	m_scene_id = scene.id;
 
+	// TODO: multi-material support
+
 	auto blas =
 		EX_PASS(VulkanAccelStruct::buildBottom(device, scene.verts.positions, scene.tris_ib));
 	VAccelStructInstance instance{blas, Matrix4::identity()};
 	m_accel_struct = EX_PASS(VulkanAccelStruct::buildTop(device, {instance}));
 	m_indices = scene.tris_ib;
 	m_vertices = scene.verts.positions;
+	m_tex_coords = scene.verts.tex_coords;
 	// TODO: wait until AS is built?
 
 	return {};
@@ -152,12 +155,23 @@ void PathTracer::render(const Context &ctx) {
 	auto raster_image = swap_chain->acquiredImage();
 	ds.setStorageImage(2, raster_image, VImageLayout::general);
 
-	ds.set(4, m_indices, m_vertices);
+	ds.set(3, m_indices, m_vertices, m_tex_coords);
 	ds.set(6, m_accel_struct);
 
 	auto sampler = ctx.device.getSampler(ctx.config.sampler_setup);
-	ds.set(10, {{sampler, ctx.opaque_tex}});
-	ds.set(11, {{sampler, ctx.trans_tex}});
+	//ds.set(10, {{sampler, ctx.opaque_tex}});
+	//ds.set(11, {{sampler, ctx.trans_tex}});
+
+	DASSERT(ctx.scene.materials);
+	auto &material = ctx.scene.materials.front();
+	// TODO: different default textures for different map types
+	auto &albedo_map = material.maps[SceneMapType::albedo];
+	auto &normal_map = material.maps[SceneMapType::normal];
+	auto &pbr_map = material.maps[SceneMapType::pbr];
+	ds.set(10, {{sampler, albedo_map.vk_image},
+				{sampler, normal_map.vk_image},
+				{sampler, pbr_map.vk_image},
+				{sampler, ctx.lighting.env_map}});
 
 	if(m_opts & Opt::debug) {
 		ds.set(12, m_debug_buffer);

src/path_tracer.h

@@ -57,6 +57,7 @@ class PathTracer {
 
 	VBufferSpan<u32> m_indices;
 	VBufferSpan<float3> m_vertices;
+	VBufferSpan<float2> m_tex_coords;
 	PVAccelStruct m_accel_struct;
 
 	int2 m_bin_counts;