2 files changed, 24 insertions, 14 deletions

diff --git a/arch/powerpc/platforms/cell/spufs/sched.c b/arch/powerpc/platforms/cell/spufs/sched.c
index 34654743363d..2deeeba7eccf 100644
--- a/arch/powerpc/platforms/cell/spufs/sched.c
+++ b/arch/powerpc/platforms/cell/spufs/sched.c
@@ -312,11 +312,28 @@ static struct spu *aff_ref_location(struct spu_context *ctx, int mem_aff,
          */
         node = cpu_to_node(raw_smp_processor_id());
         for (n = 0; n < MAX_NUMNODES; n++, node++) {
+                int available_spus;
+
                 node = (node < MAX_NUMNODES) ? node : 0;
                 if (!node_allowed(ctx, node))
                         continue;
+
+                available_spus = 0;
                 mutex_lock(&cbe_spu_info[node].list_mutex);
                 list_for_each_entry(spu, &cbe_spu_info[node].spus, cbe_list) {
+                        if (spu->ctx && spu->ctx->gang
+                                        && spu->ctx->aff_offset == 0)
+                                available_spus -=
+                                        (spu->ctx->gang->contexts - 1);
+                        else
+                                available_spus++;
+                }
+                if (available_spus < ctx->gang->contexts) {
+                        mutex_unlock(&cbe_spu_info[node].list_mutex);
+                        continue;
+                }
+
+                list_for_each_entry(spu, &cbe_spu_info[node].spus, cbe_list) {
                         if ((!mem_aff || spu->has_mem_affinity) &&
                                                         sched_spu(spu)) {
                                 mutex_unlock(&cbe_spu_info[node].list_mutex);
@@ -389,6 +406,9 @@ static int has_affinity(struct spu_context *ctx)
         if (list_empty(&ctx->aff_list))
                 return 0;
 
+        if (atomic_read(&ctx->gang->aff_sched_count) == 0)
+                ctx->gang->aff_ref_spu = NULL;
+
         if (!gang->aff_ref_spu) {
                 if (!(gang->aff_flags & AFF_MERGED))
                         aff_merge_remaining_ctxs(gang);
@@ -416,14 +436,8 @@ static void spu_unbind_context(struct spu *spu, struct spu_context *ctx)
         if (spu->ctx->flags & SPU_CREATE_NOSCHED)
                 atomic_dec(&cbe_spu_info[spu->node].reserved_spus);
 
-        if (ctx->gang){
-                mutex_lock(&ctx->gang->aff_mutex);
-                if (has_affinity(ctx)) {
-                        if (atomic_dec_and_test(&ctx->gang->aff_sched_count))
-                                ctx->gang->aff_ref_spu = NULL;
-                }
-                mutex_unlock(&ctx->gang->aff_mutex);
-        }
+        if (ctx->gang)
+                atomic_dec_if_positive(&ctx->gang->aff_sched_count);
 
         spu_switch_notify(spu, NULL);
         spu_unmap_mappings(ctx);
@@ -562,10 +576,7 @@ static struct spu *spu_get_idle(struct spu_context *ctx)
                                 goto found;
                         mutex_unlock(&cbe_spu_info[node].list_mutex);
 
-                        mutex_lock(&ctx->gang->aff_mutex);
-                        if (atomic_dec_and_test(&ctx->gang->aff_sched_count))
-                                ctx->gang->aff_ref_spu = NULL;
-                        mutex_unlock(&ctx->gang->aff_mutex);
+                        atomic_dec(&ctx->gang->aff_sched_count);
                         goto not_found;
                 }
                 mutex_unlock(&ctx->gang->aff_mutex);
diff --git a/arch/powerpc/platforms/cell/spufs/sputrace.c b/arch/powerpc/platforms/cell/spufs/sputrace.c
index 8c0e95766a62..92d20e993ede 100644
--- a/arch/powerpc/platforms/cell/spufs/sputrace.c
+++ b/arch/powerpc/platforms/cell/spufs/sputrace.c
@@ -196,8 +196,7 @@ static int __init sputrace_init(void)
         struct proc_dir_entry *entry;
         int i, error = -ENOMEM;
 
-        sputrace_log = kcalloc(sizeof(struct sputrace),
-                                bufsize, GFP_KERNEL);
+        sputrace_log = kcalloc(bufsize, sizeof(struct sputrace), GFP_KERNEL);
         if (!sputrace_log)
                 goto out;