if (switch_mmu_context) {
                        struct etnaviv_iommu_context *old_context = gpu->mmu_context;
 
-                       etnaviv_iommu_context_get(mmu_context);
-                       gpu->mmu_context = mmu_context;
+                       gpu->mmu_context = etnaviv_iommu_context_get(mmu_context);
                        etnaviv_iommu_context_put(old_context);
                }
 
 
                list_del(&mapping->obj_node);
        }
 
-       etnaviv_iommu_context_get(mmu_context);
-       mapping->context = mmu_context;
+       mapping->context = etnaviv_iommu_context_get(mmu_context);
        mapping->use = 1;
 
        ret = etnaviv_iommu_map_gem(mmu_context, etnaviv_obj,
 
                goto err_submit_objects;
 
        submit->ctx = file->driver_priv;
-       etnaviv_iommu_context_get(submit->ctx->mmu);
-       submit->mmu_context = submit->ctx->mmu;
+       submit->mmu_context = etnaviv_iommu_context_get(submit->ctx->mmu);
        submit->exec_state = args->exec_state;
        submit->flags = args->flags;
 
 
        }
 
        if (!gpu->mmu_context) {
-               etnaviv_iommu_context_get(submit->mmu_context);
-               gpu->mmu_context = submit->mmu_context;
+               gpu->mmu_context = etnaviv_iommu_context_get(submit->mmu_context);
                etnaviv_gpu_start_fe_idleloop(gpu);
        } else {
-               etnaviv_iommu_context_get(gpu->mmu_context);
-               submit->prev_mmu_context = gpu->mmu_context;
+               submit->prev_mmu_context = etnaviv_iommu_context_get(gpu->mmu_context);
        }
 
        if (submit->nr_pmrs) {
 
 struct etnaviv_iommu_context *
 etnaviv_iommu_context_init(struct etnaviv_iommu_global *global,
                           struct etnaviv_cmdbuf_suballoc *suballoc);
-static inline void etnaviv_iommu_context_get(struct etnaviv_iommu_context *ctx)
+static inline struct etnaviv_iommu_context *
+etnaviv_iommu_context_get(struct etnaviv_iommu_context *ctx)
 {
        kref_get(&ctx->refcount);
+       return ctx;
 }
 void etnaviv_iommu_context_put(struct etnaviv_iommu_context *ctx);
 void etnaviv_iommu_restore(struct etnaviv_gpu *gpu,