diff --git a/drivers/gpu/drm/sis/sis_drv.c b/drivers/gpu/drm/sis/sis_drv.c
index 6ad0b857ba2ed5ac149653e246a4d2e1503ca46c..64d329f7d9b6fc012b0f42e1e668af6642d4f30f 100644
--- a/drivers/gpu/drm/sis/sis_drv.c
+++ b/drivers/gpu/drm/sis/sis_drv.c
@@ -48,6 +48,7 @@ static int sis_driver_load(struct drm_device *dev, unsigned long chipset)
 
 	dev->dev_private = (void *)dev_priv;
 	dev_priv->chipset = chipset;
+	idr_init(&dev->object_name_idr);
 	ret = drm_sman_init(&dev_priv->sman, 2, 12, 8);
 	if (ret)
 		kfree(dev_priv);
@@ -60,6 +61,9 @@ static int sis_driver_unload(struct drm_device *dev)
 	drm_sis_private_t *dev_priv = dev->dev_private;
 
 	drm_sman_takedown(&dev_priv->sman);
+	idr_remove_all(&dev_priv->object_idr);
+	idr_destroy(&dev_priv->object_idr);
+
 	kfree(dev_priv);
 
 	return 0;
diff --git a/drivers/gpu/drm/sis/sis_drv.h b/drivers/gpu/drm/sis/sis_drv.h
index 194303c177ad65375af2801200e219f69894d3c6..fcdd06ad5dc35b65bd8d75793984bdce88b963f4 100644
--- a/drivers/gpu/drm/sis/sis_drv.h
+++ b/drivers/gpu/drm/sis/sis_drv.h
@@ -60,6 +60,8 @@ typedef struct drm_sis_private {
 	int agp_initialized;
 	unsigned long vram_offset;
 	unsigned long agp_offset;
+	/** Mapping of userspace keys to mm objects */
+	struct idr object_idr;
 } drm_sis_private_t;
 
 extern int sis_idle(struct drm_device *dev);
diff --git a/drivers/gpu/drm/sis/sis_mm.c b/drivers/gpu/drm/sis/sis_mm.c
index c76a118812a9ff124e6ae0a9072e19af6bbeedc9..21d36df02a886d15581c601cc04aec751b5776c6 100644
--- a/drivers/gpu/drm/sis/sis_mm.c
+++ b/drivers/gpu/drm/sis/sis_mm.c
@@ -125,7 +125,7 @@ static int sis_drm_alloc(struct drm_device *dev, struct drm_file *file,
 {
 	drm_sis_private_t *dev_priv = dev->dev_private;
 	drm_sis_mem_t *mem = data;
-	int retval = 0;
+	int retval = 0, user_key;
 	struct drm_memblock_item *item;
 	struct sis_file_private *file_priv = file->driver_priv;
 
@@ -141,23 +141,44 @@ static int sis_drm_alloc(struct drm_device *dev, struct drm_file *file,
 
 	mem->size = (mem->size + SIS_MM_ALIGN_MASK) >> SIS_MM_ALIGN_SHIFT;
 	item = drm_sman_alloc(&dev_priv->sman, pool, mem->size, 0, 0);
+	if (!item) {
+		retval = -ENOMEM;
+		goto fail_alloc;
+	}
 
-	if (item) {
-		list_add(&item->owner_list, &file_priv->obj_list);
-		mem->offset = ((pool == 0) ?
-			      dev_priv->vram_offset : dev_priv->agp_offset) +
-		    (item->mm->
-		     offset(item->mm, item->mm_info) << SIS_MM_ALIGN_SHIFT);
-		mem->free = item->user_hash.key;
-		mem->size = mem->size << SIS_MM_ALIGN_SHIFT;
-	} else {
-		mem->offset = 0;
-		mem->size = 0;
-		mem->free = 0;
+again:
+	if (idr_pre_get(&dev_priv->object_idr, GFP_KERNEL) == 0) {
 		retval = -ENOMEM;
+		goto fail_idr;
 	}
+
+	retval = idr_get_new_above(&dev_priv->object_idr, item, 1, &user_key);
+	if (retval == -EAGAIN)
+		goto again;
+	if (retval)
+		goto fail_idr;
+
+	list_add(&item->owner_list, &file_priv->obj_list);
+	mutex_unlock(&dev->struct_mutex);
+
+	mem->offset = ((pool == 0) ?
+		      dev_priv->vram_offset : dev_priv->agp_offset) +
+	    (item->mm->
+	     offset(item->mm, item->mm_info) << SIS_MM_ALIGN_SHIFT);
+	mem->free = user_key;
+	mem->size = mem->size << SIS_MM_ALIGN_SHIFT;
+
+	return 0;
+
+fail_idr:
+	drm_sman_free(item);
+fail_alloc:
 	mutex_unlock(&dev->struct_mutex);
 
+	mem->offset = 0;
+	mem->size = 0;
+	mem->free = 0;
+
 	DRM_DEBUG("alloc %d, size = %d, offset = %d\n", pool, mem->size,
 		  mem->offset);
 
@@ -168,10 +189,18 @@ static int sis_drm_free(struct drm_device *dev, void *data, struct drm_file *fil
 {
 	drm_sis_private_t *dev_priv = dev->dev_private;
 	drm_sis_mem_t *mem = data;
+	struct drm_memblock_item *obj;
 	int ret;
 
 	mutex_lock(&dev->struct_mutex);
-	ret = drm_sman_free_key(&dev_priv->sman, mem->free);
+	obj = idr_find(&dev_priv->object_idr, mem->free);
+	if (obj == NULL) {
+		mutex_unlock(&dev->struct_mutex);
+		return -EINVAL;
+	}
+
+	idr_remove(&dev_priv->object_idr, mem->free);
+	drm_sman_free(obj);
 	mutex_unlock(&dev->struct_mutex);
 	DRM_DEBUG("free = 0x%lx\n", mem->free);
 
diff --git a/drivers/gpu/drm/via/via_map.c b/drivers/gpu/drm/via/via_map.c
index b09f6596f15f308da5eada15e0ce0ab22bec4e21..fa5afbcf7e2860610c585d10cb432865744e5385 100644
--- a/drivers/gpu/drm/via/via_map.c
+++ b/drivers/gpu/drm/via/via_map.c
@@ -126,6 +126,8 @@ int via_driver_unload(struct drm_device *dev)
 	drm_via_private_t *dev_priv = dev->dev_private;
 
 	drm_sman_takedown(&dev_priv->sman);
+	idr_remove_all(&dev_priv->object_idr);
+	idr_destroy(&dev_priv->object_idr);
 
 	kfree(dev_priv);