diff --git a/drivers/gpu/ion/ion.c b/drivers/gpu/ion/ion.c
index cd15d71..82403d2 100644
--- a/drivers/gpu/ion/ion.c
+++ b/drivers/gpu/ion/ion.c
@@ -440,7 +440,7 @@
 			continue;
 		/* Do not allow un-secure heap if secure is specified */
 		if (secure_allocation &&
-			(heap->type != (enum ion_heap_type) ION_HEAP_TYPE_CP))
+		    !ion_heap_allow_secure_allocation(heap->type))
 			continue;
 		trace_ion_alloc_buffer_start(client->name, heap->name, len,
 					     heap_mask, flags);
@@ -1713,7 +1713,7 @@
 	buffer = handle->buffer;
 	heap = buffer->heap;
 
-	if (heap->type != (enum ion_heap_type) ION_HEAP_TYPE_CP) {
+	if (!ion_heap_allow_handle_secure(heap->type)) {
 		pr_err("%s: cannot secure buffer from non secure heap\n",
 			__func__);
 		goto out_unlock;
@@ -1746,7 +1746,7 @@
 	buffer = handle->buffer;
 	heap = buffer->heap;
 
-	if (heap->type != (enum ion_heap_type) ION_HEAP_TYPE_CP) {
+	if (!ion_heap_allow_handle_secure(heap->type)) {
 		pr_err("%s: cannot secure buffer from non secure heap\n",
 			__func__);
 		goto out_unlock;
@@ -1777,7 +1777,7 @@
 	mutex_lock(&dev->lock);
 	for (n = rb_first(&dev->heaps); n != NULL; n = rb_next(n)) {
 		struct ion_heap *heap = rb_entry(n, struct ion_heap, node);
-		if (heap->type != (enum ion_heap_type) ION_HEAP_TYPE_CP)
+		if (!ion_heap_allow_heap_secure(heap->type))
 			continue;
 		if (ION_HEAP(heap->id) != heap_id)
 			continue;
@@ -1805,7 +1805,7 @@
 	mutex_lock(&dev->lock);
 	for (n = rb_first(&dev->heaps); n != NULL; n = rb_next(n)) {
 		struct ion_heap *heap = rb_entry(n, struct ion_heap, node);
-		if (heap->type != (enum ion_heap_type) ION_HEAP_TYPE_CP)
+		if (!ion_heap_allow_heap_secure(heap->type))
 			continue;
 		if (ION_HEAP(heap->id) != heap_id)
 			continue;