#ifdef CONFIG_SMP
 
 static void
-__do_set_cpus_allowed(struct task_struct *p, const struct cpumask *new_mask, u32 flags);
+__do_set_cpus_allowed(struct task_struct *p, struct affinity_context *ctx);
 
 static int __set_cpus_allowed_ptr(struct task_struct *p,
-                                 const struct cpumask *new_mask,
-                                 u32 flags);
+                                 struct affinity_context *ctx);
 
 static void migrate_disable_switch(struct rq *rq, struct task_struct *p)
 {
+       struct affinity_context ac = {
+               .new_mask  = cpumask_of(rq->cpu),
+               .flags     = SCA_MIGRATE_DISABLE,
+       };
+
        if (likely(!p->migration_disabled))
                return;
 
        /*
         * Violates locking rules! see comment in __do_set_cpus_allowed().
         */
-       __do_set_cpus_allowed(p, cpumask_of(rq->cpu), SCA_MIGRATE_DISABLE);
+       __do_set_cpus_allowed(p, &ac);
 }
 
 void migrate_disable(void)
 void migrate_enable(void)
 {
        struct task_struct *p = current;
+       struct affinity_context ac = {
+               .new_mask  = &p->cpus_mask,
+               .flags     = SCA_MIGRATE_ENABLE,
+       };
 
        if (p->migration_disabled > 1) {
                p->migration_disabled--;
         */
        preempt_disable();
        if (p->cpus_ptr != &p->cpus_mask)
-               __set_cpus_allowed_ptr(p, &p->cpus_mask, SCA_MIGRATE_ENABLE);
+               __set_cpus_allowed_ptr(p, &ac);
        /*
         * Mustn't clear migration_disabled() until cpus_ptr points back at the
         * regular cpus_mask, otherwise things that race (eg.
  * sched_class::set_cpus_allowed must do the below, but is not required to
  * actually call this function.
  */
-void set_cpus_allowed_common(struct task_struct *p, const struct cpumask *new_mask, u32 flags)
+void set_cpus_allowed_common(struct task_struct *p, struct affinity_context *ctx)
 {
-       if (flags & (SCA_MIGRATE_ENABLE | SCA_MIGRATE_DISABLE)) {
-               p->cpus_ptr = new_mask;
+       if (ctx->flags & (SCA_MIGRATE_ENABLE | SCA_MIGRATE_DISABLE)) {
+               p->cpus_ptr = ctx->new_mask;
                return;
        }
 
-       cpumask_copy(&p->cpus_mask, new_mask);
-       p->nr_cpus_allowed = cpumask_weight(new_mask);
+       cpumask_copy(&p->cpus_mask, ctx->new_mask);
+       p->nr_cpus_allowed = cpumask_weight(ctx->new_mask);
 }
 
 static void
-__do_set_cpus_allowed(struct task_struct *p, const struct cpumask *new_mask, u32 flags)
+__do_set_cpus_allowed(struct task_struct *p, struct affinity_context *ctx)
 {
        struct rq *rq = task_rq(p);
        bool queued, running;
         *
         * XXX do further audits, this smells like something putrid.
         */
-       if (flags & SCA_MIGRATE_DISABLE)
+       if (ctx->flags & SCA_MIGRATE_DISABLE)
                SCHED_WARN_ON(!p->on_cpu);
        else
                lockdep_assert_held(&p->pi_lock);
        if (running)
                put_prev_task(rq, p);
 
-       p->sched_class->set_cpus_allowed(p, new_mask, flags);
+       p->sched_class->set_cpus_allowed(p, ctx);
 
        if (queued)
                enqueue_task(rq, p, ENQUEUE_RESTORE | ENQUEUE_NOCLOCK);
 
 void do_set_cpus_allowed(struct task_struct *p, const struct cpumask *new_mask)
 {
-       __do_set_cpus_allowed(p, new_mask, 0);
+       struct affinity_context ac = {
+               .new_mask  = new_mask,
+               .flags     = 0,
+       };
+
+       __do_set_cpus_allowed(p, &ac);
 }
 
 int dup_user_cpus_ptr(struct task_struct *dst, struct task_struct *src,
  * Called with both p->pi_lock and rq->lock held; drops both before returning.
  */
 static int __set_cpus_allowed_ptr_locked(struct task_struct *p,
-                                        const struct cpumask *new_mask,
-                                        u32 flags,
+                                        struct affinity_context *ctx,
                                         struct rq *rq,
                                         struct rq_flags *rf)
        __releases(rq->lock)
                cpu_valid_mask = cpu_online_mask;
        }
 
-       if (!kthread && !cpumask_subset(new_mask, cpu_allowed_mask)) {
+       if (!kthread && !cpumask_subset(ctx->new_mask, cpu_allowed_mask)) {
                ret = -EINVAL;
                goto out;
        }
         * Must re-check here, to close a race against __kthread_bind(),
         * sched_setaffinity() is not guaranteed to observe the flag.
         */
-       if ((flags & SCA_CHECK) && (p->flags & PF_NO_SETAFFINITY)) {
+       if ((ctx->flags & SCA_CHECK) && (p->flags & PF_NO_SETAFFINITY)) {
                ret = -EINVAL;
                goto out;
        }
 
-       if (!(flags & SCA_MIGRATE_ENABLE)) {
-               if (cpumask_equal(&p->cpus_mask, new_mask))
+       if (!(ctx->flags & SCA_MIGRATE_ENABLE)) {
+               if (cpumask_equal(&p->cpus_mask, ctx->new_mask))
                        goto out;
 
                if (WARN_ON_ONCE(p == current &&
                                 is_migration_disabled(p) &&
-                                !cpumask_test_cpu(task_cpu(p), new_mask))) {
+                                !cpumask_test_cpu(task_cpu(p), ctx->new_mask))) {
                        ret = -EBUSY;
                        goto out;
                }
         * for groups of tasks (ie. cpuset), so that load balancing is not
         * immediately required to distribute the tasks within their new mask.
         */
-       dest_cpu = cpumask_any_and_distribute(cpu_valid_mask, new_mask);
+       dest_cpu = cpumask_any_and_distribute(cpu_valid_mask, ctx->new_mask);
        if (dest_cpu >= nr_cpu_ids) {
                ret = -EINVAL;
                goto out;
        }
 
-       __do_set_cpus_allowed(p, new_mask, flags);
+       __do_set_cpus_allowed(p, ctx);
 
-       if (flags & SCA_USER)
+       if (ctx->flags & SCA_USER)
                user_mask = clear_user_cpus_ptr(p);
 
-       ret = affine_move_task(rq, p, rf, dest_cpu, flags);
+       ret = affine_move_task(rq, p, rf, dest_cpu, ctx->flags);
 
        kfree(user_mask);
 
  * call is not atomic; no spinlocks may be held.
  */
 static int __set_cpus_allowed_ptr(struct task_struct *p,
-                                 const struct cpumask *new_mask, u32 flags)
+                                 struct affinity_context *ctx)
 {
        struct rq_flags rf;
        struct rq *rq;
 
        rq = task_rq_lock(p, &rf);
-       return __set_cpus_allowed_ptr_locked(p, new_mask, flags, rq, &rf);
+       return __set_cpus_allowed_ptr_locked(p, ctx, rq, &rf);
 }
 
 int set_cpus_allowed_ptr(struct task_struct *p, const struct cpumask *new_mask)
 {
-       return __set_cpus_allowed_ptr(p, new_mask, 0);
+       struct affinity_context ac = {
+               .new_mask  = new_mask,
+               .flags     = 0,
+       };
+
+       return __set_cpus_allowed_ptr(p, &ac);
 }
 EXPORT_SYMBOL_GPL(set_cpus_allowed_ptr);
 
                                     const struct cpumask *subset_mask)
 {
        struct cpumask *user_mask = NULL;
+       struct affinity_context ac;
        struct rq_flags rf;
        struct rq *rq;
        int err;
                p->user_cpus_ptr = user_mask;
        }
 
-       return __set_cpus_allowed_ptr_locked(p, new_mask, 0, rq, &rf);
+       ac = (struct affinity_context){
+               .new_mask = new_mask,
+       };
+
+       return __set_cpus_allowed_ptr_locked(p, &ac, rq, &rf);
 
 err_unlock:
        task_rq_unlock(rq, p, &rf);
 }
 
 static int
-__sched_setaffinity(struct task_struct *p, const struct cpumask *mask);
+__sched_setaffinity(struct task_struct *p, struct affinity_context *ctx);
 
 /*
  * Restore the affinity of a task @p which was previously restricted by a
 void relax_compatible_cpus_allowed_ptr(struct task_struct *p)
 {
        struct cpumask *user_mask = p->user_cpus_ptr;
+       struct affinity_context ac = {
+               .new_mask  = user_mask,
+       };
        unsigned long flags;
 
        /*
         * we free the mask explicitly to avoid it being inherited across
         * a subsequent fork().
         */
-       if (!user_mask || !__sched_setaffinity(p, user_mask))
+       if (!user_mask || !__sched_setaffinity(p, &ac))
                return;
 
        raw_spin_lock_irqsave(&p->pi_lock, flags);
 #else /* CONFIG_SMP */
 
 static inline int __set_cpus_allowed_ptr(struct task_struct *p,
-                                        const struct cpumask *new_mask,
-                                        u32 flags)
+                                        struct affinity_context *ctx)
 {
-       return set_cpus_allowed_ptr(p, new_mask);
+       return set_cpus_allowed_ptr(p, ctx->new_mask);
 }
 
 static inline void migrate_disable_switch(struct rq *rq, struct task_struct *p) { }
 #endif
 
 static int
-__sched_setaffinity(struct task_struct *p, const struct cpumask *mask)
+__sched_setaffinity(struct task_struct *p, struct affinity_context *ctx)
 {
        int retval;
        cpumask_var_t cpus_allowed, new_mask;
        }
 
        cpuset_cpus_allowed(p, cpus_allowed);
-       cpumask_and(new_mask, mask, cpus_allowed);
+       cpumask_and(new_mask, ctx->new_mask, cpus_allowed);
+
+       ctx->new_mask = new_mask;
+       ctx->flags |= SCA_CHECK;
 
        retval = dl_task_check_affinity(p, new_mask);
        if (retval)
                goto out_free_new_mask;
 again:
-       retval = __set_cpus_allowed_ptr(p, new_mask, SCA_CHECK | SCA_USER);
+       retval = __set_cpus_allowed_ptr(p, ctx);
        if (retval)
                goto out_free_new_mask;
 
 
 long sched_setaffinity(pid_t pid, const struct cpumask *in_mask)
 {
+       struct affinity_context ac = {
+               .new_mask = in_mask,
+       };
        struct task_struct *p;
        int retval;
 
        if (retval)
                goto out_put_task;
 
-       retval = __sched_setaffinity(p, in_mask);
+       retval = __sched_setaffinity(p, &ac);
 out_put_task:
        put_task_struct(p);
        return retval;
  */
 void __init init_idle(struct task_struct *idle, int cpu)
 {
+#ifdef CONFIG_SMP
+       struct affinity_context ac = (struct affinity_context) {
+               .new_mask  = cpumask_of(cpu),
+               .flags     = 0,
+       };
+#endif
        struct rq *rq = cpu_rq(cpu);
        unsigned long flags;
 
         *
         * And since this is boot we can forgo the serialization.
         */
-       set_cpus_allowed_common(idle, cpumask_of(cpu), 0);
+       set_cpus_allowed_common(idle, &ac);
 #endif
        /*
         * We're having a chicken and egg problem, even though we are