Merge "gpu: ion: Reference count protect/unprotect calls" into msm-3.4
diff --git a/drivers/gpu/ion/ion_cp_heap.c b/drivers/gpu/ion/ion_cp_heap.c
index a857988a..6edc30b 100644
--- a/drivers/gpu/ion/ion_cp_heap.c
+++ b/drivers/gpu/ion/ion_cp_heap.c
@@ -97,6 +97,7 @@
 	int iommu_map_all;
 	int iommu_2x_map_domain;
 	unsigned int has_outer_cache;
+	atomic_t protect_cnt;
 };
 
 enum {
@@ -131,7 +132,7 @@
 		container_of(heap, struct ion_cp_heap, heap);
 	int ret_value = 0;
 
-	if (cp_heap->heap_protected == HEAP_NOT_PROTECTED) {
+	if (atomic_inc_return(&cp_heap->protect_cnt) == 1) {
 		/* Make sure we are in C state when the heap is protected. */
 		if (cp_heap->reusable && !cp_heap->allocated_bytes) {
 			ret_value = fmem_set_state(FMEM_C_STATE);
@@ -150,6 +151,7 @@
 					pr_err("%s: unable to transition heap to T-state\n",
 						__func__);
 			}
+			atomic_dec(&cp_heap->protect_cnt);
 		} else {
 			cp_heap->heap_protected = HEAP_PROTECTED;
 			pr_debug("Protected heap %s @ 0x%lx\n",
@@ -157,6 +159,9 @@
 		}
 	}
 out:
+	pr_debug("%s: protect count is %d\n", __func__,
+		atomic_read(&cp_heap->protect_cnt));
+	BUG_ON(atomic_read(&cp_heap->protect_cnt) < 0);
 	return ret_value;
 }
 
@@ -170,7 +175,7 @@
 	struct ion_cp_heap *cp_heap =
 		container_of(heap, struct ion_cp_heap, heap);
 
-	if (cp_heap->heap_protected == HEAP_PROTECTED) {
+	if (atomic_dec_and_test(&cp_heap->protect_cnt)) {
 		int error_code = ion_cp_unprotect_mem(
 			cp_heap->secure_base, cp_heap->secure_size,
 			cp_heap->permission_type);
@@ -189,6 +194,9 @@
 			}
 		}
 	}
+	pr_debug("%s: protect count is %d\n", __func__,
+		atomic_read(&cp_heap->protect_cnt));
+	BUG_ON(atomic_read(&cp_heap->protect_cnt) < 0);
 }
 
 ion_phys_addr_t ion_cp_allocate(struct ion_heap *heap,
@@ -925,6 +933,7 @@
 	cp_heap->secure_base = cp_heap->base;
 	cp_heap->secure_size = heap_data->size;
 	cp_heap->has_outer_cache = heap_data->has_outer_cache;
+	atomic_set(&cp_heap->protect_cnt, 0);
 	if (heap_data->extra_data) {
 		struct ion_cp_heap_pdata *extra_data =
 				heap_data->extra_data;