diff --git a/include/linux/idr.h b/include/linux/idr.h
index bf70b3e..7eb4432 100644
--- a/include/linux/idr.h
+++ b/include/linux/idr.h
@@ -22,8 +22,8 @@ struct idr {
 };
 
 /*
- * The IDR API does not expose the tagging functionality of the radix tree
- * to users.  Use tag 0 to track whether a node has free space below it.
+ * Reserve one of the radix tree tags to track whether a node has free
+ * space below it.
  */
 #define IDR_FREE	0
 
@@ -93,6 +93,59 @@ static inline void *idr_remove(struct idr *idr, int id)
 	return radix_tree_delete_item(&idr->idr_rt, id, NULL);
 }
 
+/**
+ * idr_tag_set - Set a tag on an entry
+ * @idr: IDR pointer
+ * @id: ID of entry to tag
+ * @tag: Tag index to set
+ *
+ * If there is an entry at @id in this IDR, set a tag on it and return
+ * the address of the entry.  If @id is outside the range of the IDR,
+ * return NULL.  This API does not allow you to set IDR_FREE on an entry;
+ * use idr_remove() for that.  The implementation does not currently check
+ * that IDR_FREE is clear, so it is possible to set a tag on a free entry.
+ * This is not recommended and may change in the future.
+ */
+static inline void *idr_tag_set(struct idr *idr, int id, unsigned int tag)
+{
+	BUG_ON(tag == IDR_FREE);
+	return radix_tree_tag_set(&idr->idr_rt, id, tag);
+}
+
+/**
+ * idr_tag_clear - Clear a tag on an entry
+ * @idr: IDR pointer
+ * @id: ID of entry to tag
+ * @tag: Tag index to clear
+ *
+ * If there is an entry at @id in this IDR, clear its tag and return
+ * the address of the entry.  If @id is outside the range of the IDR,
+ * return NULL.  This API does not allow you to clear IDR_FREE on an entry;
+ * use idr_alloc() for that.  The implementation does not currently check
+ * that IDR_FREE is clear, so it is possible to clear a tag on a free entry.
+ * This is not recommended and may change in the future.
+ */
+static inline void *idr_tag_clear(struct idr *idr, int id, unsigned int tag)
+{
+	BUG_ON(tag == IDR_FREE);
+	return radix_tree_tag_clear(&idr->idr_rt, id, tag);
+}
+
+/**
+ * idr_tag_get - Return whether a particular entry has a tag set
+ * @idr: IDR pointer
+ * @id: ID of entry to check
+ * @tag: Tag index to check
+ *
+ * Returns true/false depending whether @tag is set on this ID.  Unlike
+ * idr_tag_set() or idr_tag_clear(), you can use the IDR_FREE tag value,
+ * as it can be useful to know whether a particular ID has been allocated.
+ */
+static inline bool idr_tag_get(const struct idr *idr, int id, unsigned int tag)
+{
+	return radix_tree_tag_get(&idr->idr_rt, id, tag);
+}
+
 static inline void idr_init(struct idr *idr)
 {
 	INIT_RADIX_TREE(&idr->idr_rt, IDR_RT_MARKER);
diff --git a/lib/radix-tree.c b/lib/radix-tree.c
index 691a9ad..6723384 100644
--- a/lib/radix-tree.c
+++ b/lib/radix-tree.c
@@ -638,18 +638,17 @@ static int radix_tree_extend(struct radix_tree_root *root, gfp_t gfp,
 		if (!node)
 			return -ENOMEM;
 
+		/* Propagate the aggregated tag info to the new child */
+		for (tag = 0; tag < RADIX_TREE_MAX_TAGS; tag++) {
+			if (root_tag_get(root, tag))
+				tag_set(node, tag, 0);
+		}
 		if (is_idr(root)) {
 			all_tag_set(node, IDR_FREE);
 			if (!root_tag_get(root, IDR_FREE)) {
 				tag_clear(node, IDR_FREE, 0);
 				root_tag_set(root, IDR_FREE);
 			}
-		} else {
-			/* Propagate the aggregated tag info to the new child */
-			for (tag = 0; tag < RADIX_TREE_MAX_TAGS; tag++) {
-				if (root_tag_get(root, tag))
-					tag_set(node, tag, 0);
-			}
 		}
 
 		BUG_ON(shift > BITS_PER_LONG);
@@ -1434,8 +1433,6 @@ void *radix_tree_tag_set(struct radix_tree_root *root,
 
 		parent = entry_to_node(node);
 		offset = radix_tree_descend(parent, &node, index);
-		BUG_ON(!node);
-
 		if (!tag_get(parent, tag, offset))
 			tag_set(parent, tag, offset);
 	}
@@ -1489,7 +1486,7 @@ static void node_tag_clear(struct radix_tree_root *root,
  *	Clear the search tag (which must be < RADIX_TREE_MAX_TAGS)
  *	corresponding to @index in the radix tree.  If this causes
  *	the leaf node to have no tags set then clear the tag in the
- *	next-to-leaf node, etc.
+ *	parent node, etc.
  *
  *	Returns the address of the tagged item on success, else NULL.  ie:
  *	has the same return value and semantics as radix_tree_lookup().
@@ -1512,8 +1509,7 @@ void *radix_tree_tag_clear(struct radix_tree_root *root,
 		offset = radix_tree_descend(parent, &node, index);
 	}
 
-	if (node)
-		node_tag_clear(root, parent, tag, offset);
+	node_tag_clear(root, parent, tag, offset);
 
 	return node;
 }
@@ -1552,11 +1548,11 @@ int radix_tree_tag_get(const struct radix_tree_root *root,
 	struct radix_tree_node *node, *parent;
 	unsigned long maxindex;
 
-	if (!root_tag_get(root, tag))
-		return 0;
-
 	radix_tree_load_root(root, &node, &maxindex);
 	if (index > maxindex)
+		return is_idr(root) && (tag == IDR_FREE);
+
+	if (!root_tag_get(root, tag))
 		return 0;
 
 	while (radix_tree_is_internal_node(node)) {
@@ -1783,7 +1779,7 @@ void __rcu **radix_tree_next_chunk(const struct radix_tree_root *root,
 			child = rcu_dereference_raw(node->slots[offset]);
 		}
 
-		if (!child)
+		if (!child && !is_idr(root))
 			goto restart;
 		if (child == RADIX_TREE_RETRY)
 			break;
@@ -1994,11 +1990,10 @@ static bool __radix_tree_delete(struct radix_tree_root *root,
 	unsigned offset = get_slot_offset(node, slot);
 	int tag;
 
+	for (tag = 0; tag < RADIX_TREE_MAX_TAGS; tag++)
+		node_tag_clear(root, node, tag, offset);
 	if (is_idr(root))
 		node_tag_set(root, node, IDR_FREE, offset);
-	else
-		for (tag = 0; tag < RADIX_TREE_MAX_TAGS; tag++)
-			node_tag_clear(root, node, tag, offset);
 
 	replace_slot(slot, NULL, node, -1, exceptional);
 	return node && delete_node(root, node, NULL, NULL);
diff --git a/tools/testing/radix-tree/idr-test.c b/tools/testing/radix-tree/idr-test.c
index 30cd0b2..fd94bee 100644
--- a/tools/testing/radix-tree/idr-test.c
+++ b/tools/testing/radix-tree/idr-test.c
@@ -135,6 +135,32 @@ void idr_null_test(void)
 	assert(idr_is_empty(&idr));
 }
 
+#define IDR_TEST	1
+
+void idr_tag_test(void)
+{
+	unsigned int i;
+	DEFINE_IDR(idr);
+
+	for (i = 0; i < 100; i++) {
+		assert(idr_alloc(&idr, NULL, 0, 0, GFP_KERNEL) == i);
+		if (i % 7 == 0)
+			idr_tag_set(&idr, i, IDR_TEST);
+	}
+
+	for (i = 0; i < 100; i += 14) {
+		assert(idr_tag_get(&idr, i, IDR_TEST));
+		idr_tag_clear(&idr, i, IDR_TEST);
+	}
+
+	for (i = 0; i < 100; i++) {
+		assert(idr_tag_get(&idr, i, IDR_TEST) == (i % 14 == 7));
+	}
+
+	idr_for_each(&idr, item_idr_free, &idr);
+	idr_destroy(&idr);
+}
+
 void idr_nowait_test(void)
 {
 	unsigned int i;
@@ -225,6 +251,7 @@ void idr_checks(void)
 	idr_replace_test();
 	idr_alloc_test();
 	idr_null_test();
+	idr_tag_test();
 	idr_nowait_test();
 	idr_get_next_test();
 }
