From b83596accd1b474e44535d65680a87313f0bf93f Mon Sep 17 00:00:00 2001
From: Wenzel Jakob <wenzel@cs.cornell.edu>
Date: Mon, 1 Oct 2012 00:55:10 -0400
Subject: [PATCH] more MacOS-related OpenMP bugfixes

---
 include/mitsuba/core/thread.h         | 12 +++++++++++-
 src/integrators/photonmapper/bre.cpp  |  2 +-
 src/integrators/photonmapper/sppm.cpp |  4 +---
 src/libcore/thread.cpp                | 19 ++++++++++++++-----
 src/subsurface/bluenoise.cpp          |  2 +-
 5 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/include/mitsuba/core/thread.h b/include/mitsuba/core/thread.h
index 136fcfa4..f2c0b4bd 100644
--- a/include/mitsuba/core/thread.h
+++ b/include/mitsuba/core/thread.h
@@ -174,9 +174,19 @@ private:
 	boost::scoped_ptr<ThreadPrivate> d;
 };
 
-/// Variant of \c omp_get_max_threads that works on all platforms
+#if defined(MTS_OPENMP)
+#if defined(__OSX__)
+/// Variant of \c omp_get_max_threads that works on OSX
 extern MTS_EXPORT_CORE int mts_omp_get_max_threads();
 
+/// Variant of \c omp_get_thread_num that works on OSX
+extern MTS_EXPORT_CORE int mts_omp_get_thread_num();
+#else
+#define mts_omp_get_max_threads omp_get_max_threads
+#define mts_omp_get_thread_num omp_get_thread_num
+#endif
+#endif
+
 MTS_NAMESPACE_END
 
 #endif /* __MITSUBA_CORE_THREAD_H_ */
diff --git a/src/integrators/photonmapper/bre.cpp b/src/integrators/photonmapper/bre.cpp
index de9b3e5f..c17c6ea2 100644
--- a/src/integrators/photonmapper/bre.cpp
+++ b/src/integrators/photonmapper/bre.cpp
@@ -57,7 +57,7 @@ BeamRadianceEstimator::BeamRadianceEstimator(const PhotonMap *pmap, size_t looku
 	#endif
 	for (int i=0; i<(int) m_photonCount; ++i) {
 		#if defined(MTS_OPENMP)
-			int tid = omp_get_thread_num();
+			int tid = mts_omp_get_thread_num();
 		#else
 			int tid = 0;
 		#endif
diff --git a/src/integrators/photonmapper/sppm.cpp b/src/integrators/photonmapper/sppm.cpp
index b292524f..6fb2c0e0 100644
--- a/src/integrators/photonmapper/sppm.cpp
+++ b/src/integrators/photonmapper/sppm.cpp
@@ -181,8 +181,6 @@ public:
 			samplers[i] = clonedSampler.get();
 		}
 
-		Thread::initializeOpenMP(Scheduler::getInstance()->getLocalWorkerCount());
-
 		int samplerResID = sched->registerMultiResource(samplers);
 
 #ifdef MTS_DEBUG_FP
@@ -224,7 +222,7 @@ public:
 		for (int i=0; i<(int) m_gatherBlocks.size(); ++i) {
 			std::vector<GatherPoint> &gatherPoints = m_gatherBlocks[i];
 			#if defined(MTS_OPENMP)
-				Sampler *sampler = static_cast<Sampler *>(samplers[omp_get_thread_num()]);
+				Sampler *sampler = static_cast<Sampler *>(samplers[mts_omp_get_thread_num()]);
 			#else
 				Sampler *sampler = static_cast<Sampler *>(samplers[0]);
 			#endif
diff --git a/src/libcore/thread.cpp b/src/libcore/thread.cpp
index 5caa52ad..705e3cb2 100644
--- a/src/libcore/thread.cpp
+++ b/src/libcore/thread.cpp
@@ -390,18 +390,22 @@ std::string Thread::toString() const {
 	return oss.str();
 }
 
-#if defined(MTS_OPENMP)
-static int __omp_threadCount = 0;
-#endif
 static std::vector<Thread *> __unmanagedThreads;
 static boost::mutex __unmanagedMutex;
 
-#if defined(MTS_OPENMP)
+#if defined(MTS_OPENMP) && defined(__OSX__)
+static int __omp_threadCount = 0;
+static pthread_key_t __omp_key;
+
 int mts_omp_get_max_threads() {
 	/* This function exists to sidestep an annoying
 	   implementation bug that causes crashes on OSX */
 	return __omp_threadCount;
 }
+
+int mts_omp_get_thread_num() {
+	return reinterpret_cast<int>(pthread_getspecific(__omp_key));
+}
 #endif
 
 void Thread::staticInitialization() {
@@ -457,7 +461,11 @@ void Thread::initializeOpenMP(size_t threadCount) {
 	ref<Logger> logger = Thread::getThread()->getLogger();
 	ref<FileResolver> fResolver = Thread::getThread()->getFileResolver();
 
-	__omp_threadCount = threadCount;
+	#if defined(__OSX__)
+		__omp_threadCount = threadCount;
+		pthread_key_create(&__omp_key, NULL);
+	#endif
+
 	omp_set_num_threads((int) threadCount);
 	omp_set_dynamic(false);
 
@@ -465,6 +473,7 @@ void Thread::initializeOpenMP(size_t threadCount) {
 
 	#pragma omp parallel
 	{
+		pthread_setspecific(__omp_key, reinterpret_cast<void *>(counter));
 		detail::initializeLocalTLS();
 		Thread *thread = Thread::getThread();
 		if (!thread) {
diff --git a/src/subsurface/bluenoise.cpp b/src/subsurface/bluenoise.cpp
index 20aecbb3..f4c73db9 100644
--- a/src/subsurface/bluenoise.cpp
+++ b/src/subsurface/bluenoise.cpp
@@ -114,7 +114,7 @@ void blueNoisePointSet(const Scene *scene, const std::vector<Shape *> &shapes,
 	#endif
 	for (int i=0; i<nsamples; ++i) {
 		#if defined(MTS_OPENMP)
-			int tid = omp_get_thread_num();
+			int tid = mts_omp_get_thread_num();
 		#else
 			int tid = 0;
 		#endif