1 files changed, 30 insertions, 4 deletions

diff --git a/kernel/sched/membarrier.c b/kernel/sched/membarrier.c
index a92fddc22747..dd7908743dab 100644
--- a/kernel/sched/membarrier.c
+++ b/kernel/sched/membarrier.c
@@ -18,6 +18,7 @@
 #include <linux/membarrier.h>
 #include <linux/tick.h>
 #include <linux/cpumask.h>
+#include <linux/atomic.h>
 
 #include "sched.h"      /* for cpu_rq(). */
 
@@ -26,21 +27,26 @@
  * except MEMBARRIER_CMD_QUERY.
  */
 #define MEMBARRIER_CMD_BITMASK  \
-        (MEMBARRIER_CMD_SHARED | MEMBARRIER_CMD_PRIVATE_EXPEDITED)
+        (MEMBARRIER_CMD_SHARED | MEMBARRIER_CMD_PRIVATE_EXPEDITED       \
+        | MEMBARRIER_CMD_REGISTER_PRIVATE_EXPEDITED)
 
 static void ipi_mb(void *info)
 {
         smp_mb();       /* IPIs should be serializing but paranoid. */
 }
 
-static void membarrier_private_expedited(void)
+static int membarrier_private_expedited(void)
 {
         int cpu;
         bool fallback = false;
         cpumask_var_t tmpmask;
 
+        if (!(atomic_read(&current->mm->membarrier_state)
+                        & MEMBARRIER_STATE_PRIVATE_EXPEDITED_READY))
+                return -EPERM;
+
         if (num_online_cpus() == 1)
-                return;
+                return 0;
 
         /*
          * Matches memory barriers around rq->curr modification in
@@ -94,6 +100,24 @@ static void membarrier_private_expedited(void)
          * rq->curr modification in scheduler.
          */
         smp_mb();       /* exit from system call is not a mb */
+        return 0;
+}
+
+static void membarrier_register_private_expedited(void)
+{
+        struct task_struct *p = current;
+        struct mm_struct *mm = p->mm;
+
+        /*
+         * We need to consider threads belonging to different thread
+         * groups, which use the same mm. (CLONE_VM but not
+         * CLONE_THREAD).
+         */
+        if (atomic_read(&mm->membarrier_state)
+                        & MEMBARRIER_STATE_PRIVATE_EXPEDITED_READY)
+                return;
+        atomic_or(MEMBARRIER_STATE_PRIVATE_EXPEDITED_READY,
+                        &mm->membarrier_state);
 }
 
 /**
@@ -144,7 +168,9 @@ SYSCALL_DEFINE2(membarrier, int, cmd, int, flags)
                         synchronize_sched();
                 return 0;
         case MEMBARRIER_CMD_PRIVATE_EXPEDITED:
-                membarrier_private_expedited();
+                return membarrier_private_expedited();
+        case MEMBARRIER_CMD_REGISTER_PRIVATE_EXPEDITED:
+                membarrier_register_private_expedited();
                 return 0;
         default:
                 return -EINVAL;