update sched_op transactions enter/leave in get_multi_worker_task

src/drivers/driver_common/driver_common.c

@@ -430,6 +430,7 @@ int _starpu_get_multi_worker_task(struct _starpu_worker *workers, struct starpu_
 	/* This assumes only 1 worker */
 	STARPU_ASSERT_MSG(nworkers == 1, "Multiple workers is not yet possible in blocking drivers mode\n");
 	STARPU_PTHREAD_MUTEX_LOCK_SCHED(&workers[0].sched_mutex);
+	_starpu_worker_enter_sched_op(&workers[0]);
 #endif
 	for (i = 0; i < nworkers; i++)
 	{
@@ -457,15 +458,18 @@ int _starpu_get_multi_worker_task(struct _starpu_worker *workers, struct starpu_
 			STARPU_PTHREAD_MUTEX_LOCK_SCHED(&workers[i].sched_mutex);
 #endif
 			_starpu_worker_set_status_scheduling(workers[i].workerid);
-			_starpu_set_local_worker_key(&workers[i]);
+#ifdef STARPU_NON_BLOCKING_DRIVERS
 			_starpu_worker_enter_sched_op(&workers[i]);
+#endif
+			_starpu_set_local_worker_key(&workers[i]);
 			tasks[i] = _starpu_pop_task(&workers[i]);
-			_starpu_worker_leave_sched_op(&workers[i]);
+
 			if(tasks[i] != NULL)
 			{
 				_starpu_worker_set_status_scheduling_done(workers[i].workerid);
 				_starpu_worker_set_status_wakeup(workers[i].workerid);
 #ifdef STARPU_NON_BLOCKING_DRIVERS
+				_starpu_worker_leave_sched_op(&workers[i]);
 				STARPU_PTHREAD_MUTEX_UNLOCK_SCHED(&workers[i].sched_mutex);
 #endif
 
@@ -504,6 +508,7 @@ int _starpu_get_multi_worker_task(struct _starpu_worker *workers, struct starpu_
 			{
 				_starpu_worker_set_status_sleeping(workers[i].workerid);
 #ifdef STARPU_NON_BLOCKING_DRIVERS
+				_starpu_worker_leave_sched_op(&workers[i]);
 				STARPU_PTHREAD_MUTEX_UNLOCK_SCHED(&workers[i].sched_mutex);
 #endif
 			}
@@ -529,6 +534,7 @@ int _starpu_get_multi_worker_task(struct _starpu_worker *workers, struct starpu_
 		if (_starpu_worker_can_block(memnode, worker)
 				&& !_starpu_sched_ctx_last_worker_awake(worker))
 		{
+			_starpu_worker_leave_sched_op(worker);
 			do
 			{
 				STARPU_PTHREAD_COND_WAIT(&worker->sched_cond, &worker->sched_mutex);
@@ -538,6 +544,7 @@ int _starpu_get_multi_worker_task(struct _starpu_worker *workers, struct starpu_
 		}
 		else
 		{
+			_starpu_worker_leave_sched_op(worker);
 			STARPU_PTHREAD_MUTEX_UNLOCK_SCHED(&worker->sched_mutex);
 			if (_starpu_machine_is_running())
 				_starpu_exponential_backoff(worker);
@@ -549,6 +556,7 @@ int _starpu_get_multi_worker_task(struct _starpu_worker *workers, struct starpu_
 	worker->spinning_backoff = BACKOFF_MIN;
 #endif /* !STARPU_SIMGRID */
 
+	_starpu_worker_leave_sched_op(&workers[0]);
 	STARPU_PTHREAD_MUTEX_UNLOCK_SCHED(&workers[0].sched_mutex);
 #endif /* !STARPU_NON_BLOCKING_DRIVERS */