diff --git a/drivers/staging/android/ion/ion.c b/drivers/staging/android/ion/ion.c
index e4586eb6cae3..948b37f74f2b 100644
--- a/drivers/staging/android/ion/ion.c
+++ b/drivers/staging/android/ion/ion.c
@@ -7,6 +7,7 @@
  *
  */
 
+#include <linux/bitmap.h>
 #include <linux/debugfs.h>
 #include <linux/device.h>
 #include <linux/dma-buf.h>
@@ -27,7 +28,6 @@
 #include "ion_private.h"
 
 static struct ion_device *internal_dev;
-static int heap_id;
 
 /* Entry into ION allocator for rest of the kernel */
 struct dma_buf *ion_alloc(size_t len, unsigned int heap_id_mask,
@@ -222,6 +222,69 @@ static int debug_shrink_get(void *data, u64 *val)
 DEFINE_SIMPLE_ATTRIBUTE(debug_shrink_fops, debug_shrink_get,
 			debug_shrink_set, "%llu\n");
 
+static int ion_assign_heap_id(struct ion_heap *heap, struct ion_device *dev)
+{
+	int id_bit;
+	int start_bit, end_bit;
+
+	switch (heap->type) {
+	case ION_HEAP_TYPE_SYSTEM:
+		id_bit = ffs(ION_HEAP_SYSTEM);
+		break;
+	case ION_HEAP_TYPE_SYSTEM_CONTIG:
+		id_bit = ffs(ION_HEAP_SYSTEM_CONTIG);
+		break;
+	case ION_HEAP_TYPE_CHUNK:
+		id_bit = ffs(ION_HEAP_CHUNK);
+		break;
+	case ION_HEAP_TYPE_CARVEOUT:
+		id_bit = 0;
+		start_bit = ffs(ION_HEAP_CARVEOUT_START);
+		end_bit = ffs(ION_HEAP_CARVEOUT_END);
+		break;
+	case ION_HEAP_TYPE_DMA:
+		id_bit = 0;
+		start_bit = ffs(ION_HEAP_DMA_START);
+		end_bit = ffs(ION_HEAP_DMA_END);
+		break;
+	case ION_HEAP_TYPE_CUSTOM ... ION_HEAP_TYPE_MAX:
+		id_bit = 0;
+		start_bit = ffs(ION_HEAP_CUSTOM_START);
+		end_bit = ffs(ION_HEAP_CUSTOM_END);
+		break;
+	default:
+		return -EINVAL;
+	}
+
+	/* For carveout, dma & custom heaps, we first let the heaps choose their
+	 * own IDs. This allows the old behaviour of knowing the heap ids
+	 * of these type of heaps  in advance in user space. If a heap with
+	 * that ID already exists, it is an error.
+	 *
+	 * If the heap hasn't picked an id by itself, then we assign it
+	 * one.
+	 */
+	if (!id_bit) {
+		if (heap->id) {
+			id_bit = ffs(heap->id);
+			if (id_bit < start_bit || id_bit > end_bit)
+				return -EINVAL;
+		} else {
+			id_bit = find_next_zero_bit(dev->heap_ids, end_bit + 1,
+						    start_bit);
+			if (id_bit > end_bit)
+				return -ENOSPC;
+		}
+	}
+
+	if (test_and_set_bit(id_bit - 1, dev->heap_ids))
+		return -EEXIST;
+	heap->id = id_bit;
+	dev->heap_cnt++;
+
+	return 0;
+}
+
 int __ion_device_add_heap(struct ion_heap *heap, struct module *owner)
 {
 	struct ion_device *dev = internal_dev;
@@ -283,7 +346,14 @@ int __ion_device_add_heap(struct ion_heap *heap, struct module *owner)
 
 	heap->debugfs_dir = heap_root;
 	down_write(&dev->lock);
-	heap->id = heap_id++;
+	ret = ion_assign_heap_id(heap, dev);
+	if (ret) {
+		pr_err("%s: Failed to assign heap id for heap type %x\n",
+		       __func__, heap->type);
+		up_write(&dev->lock);
+		goto out_debugfs_cleanup;
+	}
+
 	/*
 	 * use negative heap->id to reverse the priority -- when traversing
 	 * the list later attempt higher id numbers first
@@ -291,11 +361,12 @@ int __ion_device_add_heap(struct ion_heap *heap, struct module *owner)
 	plist_node_init(&heap->node, -heap->id);
 	plist_add(&heap->node, &dev->heaps);
 
-	dev->heap_cnt++;
 	up_write(&dev->lock);
 
 	return 0;
 
+out_debugfs_cleanup:
+	debugfs_remove_recursive(heap->debugfs_dir);
 out_heap_cleanup:
 	ion_heap_cleanup(heap);
 out:
@@ -321,6 +392,8 @@ void ion_device_remove_heap(struct ion_heap *heap)
 			__func__, heap->name);
 	}
 	debugfs_remove_recursive(heap->debugfs_dir);
+	clear_bit(heap->id - 1, dev->heap_ids);
+	dev->heap_cnt--;
 	up_write(&dev->lock);
 }
 EXPORT_SYMBOL(ion_device_remove_heap);
diff --git a/drivers/staging/android/ion/ion_private.h b/drivers/staging/android/ion/ion_private.h
index 5cc09fbfbefe..2f5376e131c3 100644
--- a/drivers/staging/android/ion/ion_private.h
+++ b/drivers/staging/android/ion/ion_private.h
@@ -23,13 +23,16 @@
  * @dev:		the actual misc device
  * @buffers:		an rb tree of all the existing buffers
  * @buffer_lock:	lock protecting the tree of buffers
- * @lock:		rwsem protecting the tree of heaps and clients
+ * @lock:		rwsem protecting the tree of heaps, heap_bitmap and
+ *			clients
+ * @heap_ids:		bitmap of register heap ids
  */
 struct ion_device {
 	struct miscdevice dev;
 	struct rb_root buffers;
 	struct mutex buffer_lock;
 	struct rw_semaphore lock;
+	DECLARE_BITMAP(heap_ids, ION_NUM_MAX_HEAPS);
 	struct plist_head heaps;
 	struct dentry *debug_root;
 	int heap_cnt;