Multirs baker: support for threaded baking

3 files changed, 169 insertions, 71 deletions

diff --git a/source/blender/editors/object/object_bake.c b/source/blender/editors/object/object_bake.c
index 8b40379f88d..6feb8096aff 100644
--- a/source/blender/editors/object/object_bake.c
+++ b/source/blender/editors/object/object_bake.c
@@ -100,6 +100,7 @@ typedef struct {
 	float bias;
 	int raytrace_structure;
 	int octree_resolution;
+	int threads;
 } MultiresBakeJob;
 
 static int multiresbake_check(bContext *C, wmOperator *op)
@@ -319,6 +320,7 @@ static int multiresbake_image_exec_locked(bContext *C, wmOperator *op)
 		bkr.number_of_rays = scene->r.bake_rays_number;
 		bkr.raytrace_structure = scene->r.raytrace_structure;
 		bkr.octree_resolution = scene->r.ocres;
+		bkr.threads = scene->r.mode & R_FIXED_THREADS ? scene->r.threads : 0;
 
 		/* create low-resolution DM (to bake to) and hi-resolution DM (to bake from) */
 		bkr.hires_dm = multiresbake_create_hiresdm(scene, ob, &bkr.tot_lvl, &bkr.simple);
@@ -356,6 +358,7 @@ static void init_multiresbake_job(bContext *C, MultiresBakeJob *bkj)
 	bkj->number_of_rays = scene->r.bake_rays_number;
 	bkj->raytrace_structure = scene->r.raytrace_structure;
 	bkj->octree_resolution = scene->r.ocres;
+	bkj->threads = scene->r.mode & R_FIXED_THREADS ? scene->r.threads : 0;
 
 	CTX_DATA_BEGIN (C, Base *, base, selected_editable_bases)
 	{
@@ -422,6 +425,7 @@ static void multiresbake_startjob(void *bkv, short *stop, short *do_update, floa
 		bkr.number_of_rays = bkj->number_of_rays;
 		bkr.raytrace_structure = bkj->raytrace_structure;
 		bkr.octree_resolution = bkj->octree_resolution;
+		bkr.threads = bkj->threads;
 
 		RE_multires_bake_images(&bkr);
 
diff --git a/source/blender/render/extern/include/RE_multires_bake.h b/source/blender/render/extern/include/RE_multires_bake.h
index 5e8ebdd8a18..04cfe55e3a3 100644
--- a/source/blender/render/extern/include/RE_multires_bake.h
+++ b/source/blender/render/extern/include/RE_multires_bake.h
@@ -50,7 +50,8 @@ typedef struct MultiresBakeRender {
 
 	int raytrace_structure;
 	int octree_resolution;
-
+	int threads;
+	
 	short *stop;
 	short *do_update;
 	float *progress;
diff --git a/source/blender/render/intern/source/multires_bake.c b/source/blender/render/intern/source/multires_bake.c
index 44fc6476977..50be28bdc12 100644
--- a/source/blender/render/intern/source/multires_bake.c
+++ b/source/blender/render/intern/source/multires_bake.c
@@ -39,6 +39,7 @@
 
 #include "BLI_math.h"
 #include "BLI_listbase.h"
+#include "BLI_threads.h"
 
 #include "BKE_ccg.h"
 #include "BKE_context.h"
@@ -326,104 +327,196 @@ static int multiresbake_test_break(MultiresBakeRender *bkr)
 		return 0;
 	}
 
-	return G.is_break;
+	return *bkr->stop || G.is_break;
 }
 
-static void do_multires_bake(MultiresBakeRender *bkr, Image *ima, int require_tangent, MPassKnownData passKnownData,
-                             MInitBakeData initBakeData, MApplyBakeData applyBakeData, MFreeBakeData freeBakeData)
+/* **** Threading routines **** */
+
+typedef struct MultiresBakeQueue {
+	int cur_face;
+	int tot_face;
+	SpinLock spin;
+} MultiresBakeQueue;
+
+typedef struct MultiresBakeThread {
+	/* this data is actually shared between all the threads */
+	MultiresBakeQueue *queue;
+	MultiresBakeRender *bkr;
+	Image *image;
+	void *bake_data;
+
+	/* thread-specific data */
+	MBakeRast bake_rast;
+	MResolvePixelData data;
+} MultiresBakeThread;
+
+static int multires_bake_queue_next_face(MultiresBakeQueue *queue)
 {
-	DerivedMesh *dm = bkr->lores_dm;
-	ImBuf *ibuf = BKE_image_acquire_ibuf(ima, NULL, NULL);
-	const int lvl = bkr->lvl;
-	const int tot_face = dm->getNumTessFaces(dm);
-	MVert *mvert = dm->getVertArray(dm);
-	MFace *mface = dm->getTessFaceArray(dm);
-	MTFace *mtface = dm->getTessFaceDataArray(dm, CD_MTFACE);
-	float *pvtangent = NULL;
+	int face = -1;
 
-	if (require_tangent) {
-		if (CustomData_get_layer_index(&dm->faceData, CD_TANGENT) == -1)
-			DM_add_tangent_layer(dm);
+	/* TODO: it could worth making it so thread will handle neighbor faces
+	 *       for better memory cache utilization
+	 */
 
-		pvtangent = DM_get_tessface_data_layer(dm, CD_TANGENT);
+	BLI_spin_lock(&queue->spin);
+	if (queue->cur_face < queue->tot_face) {
+		face = queue->cur_face;
+		queue->cur_face++;
 	}
+	BLI_spin_unlock(&queue->spin);
 
-	if (tot_face > 0) {  /* sanity check */
-		int f = 0;
-		MBakeRast bake_rast;
-		MResolvePixelData data = {NULL};
-
-		data.mface = mface;
-		data.mvert = mvert;
-		data.mtface = mtface;
-		data.pvtangent = pvtangent;
-		data.precomputed_normals = dm->getTessFaceDataArray(dm, CD_NORMAL);  /* don't strictly need this */
-		data.w = ibuf->x;
-		data.h = ibuf->y;
-		data.lores_dm = dm;
-		data.hires_dm = bkr->hires_dm;
-		data.lvl = lvl;
-		data.pass_data = passKnownData;
+	return face;
+}
 
-		if (initBakeData)
-			data.bake_data = initBakeData(bkr, ima);
+static void *do_multires_bake_thread(void *data_v)
+{
+	MultiresBakeThread *handle = (MultiresBakeThread *) data_v;
+	MResolvePixelData *data = &handle->data;
+	MBakeRast *bake_rast = &handle->bake_rast;
+	MultiresBakeRender *bkr = handle->bkr;
+	int f;
 
-		init_bake_rast(&bake_rast, ibuf, &data, flush_pixel);
+	while ((f = multires_bake_queue_next_face(handle->queue)) >= 0) {
+		MTFace *mtfate = &data->mtface[f];
+		int verts[3][2], nr_tris, t;
 
-		for (f = 0; f < tot_face; f++) {
-			MTFace *mtfate = &mtface[f];
-			int verts[3][2], nr_tris, t;
+		if (multiresbake_test_break(bkr))
+			break;
 
-			if (multiresbake_test_break(bkr))
-				break;
+		if (mtfate->tpage != handle->image)
+			continue;
 
-			if (mtfate->tpage != ima)
-				continue;
+		data->face_index = f;
 
-			data.face_index = f;
-			data.ibuf = ibuf;
+		/* might support other forms of diagonal splits later on such as
+		 * split by shortest diagonal.*/
+		verts[0][0] = 0;
+		verts[1][0] = 1;
+		verts[2][0] = 2;
 
-			/* might support other forms of diagonal splits later on such as
-			 * split by shortest diagonal.*/
-			verts[0][0] = 0;
-			verts[1][0] = 1;
-			verts[2][0] = 2;
+		verts[0][1] = 0;
+		verts[1][1] = 2;
+		verts[2][1] = 3;
 
-			verts[0][1] = 0;
-			verts[1][1] = 2;
-			verts[2][1] = 3;
+		nr_tris = data->mface[f].v4 != 0 ? 2 : 1;
+		for (t = 0; t < nr_tris; t++) {
+			data->i0 = verts[0][t];
+			data->i1 = verts[1][t];
+			data->i2 = verts[2][t];
 
-			nr_tris = mface[f].v4 != 0 ? 2 : 1;
-			for (t = 0; t < nr_tris; t++) {
-				data.i0 = verts[0][t];
-				data.i1 = verts[1][t];
-				data.i2 = verts[2][t];
+			bake_rasterize(bake_rast, mtfate->uv[data->i0], mtfate->uv[data->i1], mtfate->uv[data->i2]);
 
-				bake_rasterize(&bake_rast, mtfate->uv[data.i0], mtfate->uv[data.i1], mtfate->uv[data.i2]);
+			/* tag image buffer for refresh */
+			if (data->ibuf->rect_float)
+				data->ibuf->userflags |= IB_RECT_INVALID;
 
-				if (ibuf->rect_float)
-					ibuf->userflags |= IB_RECT_INVALID;
+			data->ibuf->userflags |= IB_DISPLAY_BUFFER_INVALID;
+		}
 
-				ibuf->userflags |= IB_DISPLAY_BUFFER_INVALID;
-			}
+		/* update progress */
+		BLI_spin_lock(&handle->queue->spin);
+		bkr->baked_faces++;
+
+		if (bkr->do_update)
+			*bkr->do_update = TRUE;
+
+		if (bkr->progress)
+			*bkr->progress = ((float)bkr->baked_objects + (float)bkr->baked_faces / handle->queue->tot_face) / bkr->tot_obj;
+		BLI_spin_unlock(&handle->queue->spin);
+	}
 
-			bkr->baked_faces++;
+	return NULL;
+}
+
+static void do_multires_bake(MultiresBakeRender *bkr, Image *ima, int require_tangent, MPassKnownData passKnownData,
+                             MInitBakeData initBakeData, MApplyBakeData applyBakeData, MFreeBakeData freeBakeData)
+{
+	DerivedMesh *dm = bkr->lores_dm;
+	const int lvl = bkr->lvl;
+	const int tot_face = dm->getNumTessFaces(dm);
+
+	if (tot_face > 0) {
+		MultiresBakeThread *handles;
+		MultiresBakeQueue queue;
+
+		ImBuf *ibuf = BKE_image_acquire_ibuf(ima, NULL, NULL);
+		MVert *mvert = dm->getVertArray(dm);
+		MFace *mface = dm->getTessFaceArray(dm);
+		MTFace *mtface = dm->getTessFaceDataArray(dm, CD_MTFACE);
+		float *precomputed_normals = dm->getTessFaceDataArray(dm, CD_NORMAL);
+		float *pvtangent = NULL;
 
-			if (bkr->do_update)
-				*bkr->do_update = TRUE;
+		ListBase threads;
+		int i, tot_thread = bkr->threads > 0 ? bkr->threads : BLI_system_thread_count();
 
-			if (bkr->progress)
-				*bkr->progress = ((float)bkr->baked_objects + (float)bkr->baked_faces / tot_face) / bkr->tot_obj;
+		void *bake_data = NULL;
+
+		if (require_tangent) {
+			if (CustomData_get_layer_index(&dm->faceData, CD_TANGENT) == -1)
+				DM_add_tangent_layer(dm);
+
+			pvtangent = DM_get_tessface_data_layer(dm, CD_TANGENT);
+		}
+
+		/* all threads shares the same custom bake data */
+		if (initBakeData)
+			bake_data = initBakeData(bkr, ima);
+
+		if (tot_thread > 1)
+			BLI_init_threads(&threads, do_multires_bake_thread, tot_thread);
+
+		handles = MEM_callocN(tot_thread * sizeof(MultiresBakeThread), "do_multires_bake handles");
+
+		/* faces queue */
+		queue.cur_face = 0;
+		queue.tot_face = tot_face;
+		BLI_spin_init(&queue.spin);
+
+		/* fill in threads handles */
+		for (i = 0; i < tot_thread; i++) {
+			MultiresBakeThread *handle = &handles[i];
+
+			handle->bkr = bkr;
+			handle->image = ima;
+			handle->queue = &queue;
+
+			handle->data.mface = mface;
+			handle->data.mvert = mvert;
+			handle->data.mtface = mtface;
+			handle->data.pvtangent = pvtangent;
+			handle->data.precomputed_normals = precomputed_normals;  /* don't strictly need this */
+			handle->data.w = ibuf->x;
+			handle->data.h = ibuf->y;
+			handle->data.lores_dm = dm;
+			handle->data.hires_dm = bkr->hires_dm;
+			handle->data.lvl = lvl;
+			handle->data.pass_data = passKnownData;
+			handle->data.bake_data = bake_data;
+			handle->data.ibuf = ibuf;
+
+			init_bake_rast(&handle->bake_rast, ibuf, &handle->data, flush_pixel);
+
+			if (tot_thread > 1)
+				BLI_insert_thread(&threads, handle);
 		}
 
+		/* run threads */
+		if (tot_thread > 1)
+			BLI_end_threads(&threads);
+		else
+			do_multires_bake_thread(&handles[0]);
+
+		BLI_spin_end(&queue.spin);
+
+		/* finalize baking */
 		if (applyBakeData)
-			applyBakeData(data.bake_data);
+			applyBakeData(bake_data);
 
 		if (freeBakeData)
-			freeBakeData(data.bake_data);
-	}
+			freeBakeData(bake_data);
 
-	BKE_image_release_ibuf(ima, ibuf, NULL);
+		BKE_image_release_ibuf(ima, ibuf, NULL);
+	}
 }
 
 /* mode = 0: interpolate normals,