spin_lock_init(&iovad->iova_rbtree_lock);
        iovad->rbroot = RB_ROOT;
+       iovad->cached_node = NULL;
        iovad->cached32_node = NULL;
        iovad->granule = granule;
        iovad->start_pfn = start_pfn;
 static struct rb_node *
 __get_cached_rbnode(struct iova_domain *iovad, unsigned long *limit_pfn)
 {
-       if ((*limit_pfn > iovad->dma_32bit_pfn) ||
-               (iovad->cached32_node == NULL))
+       struct rb_node *cached_node = NULL;
+       struct iova *curr_iova;
+
+       if (*limit_pfn <= iovad->dma_32bit_pfn)
+               cached_node = iovad->cached32_node;
+       if (!cached_node)
+               cached_node = iovad->cached_node;
+       if (!cached_node)
                return rb_last(&iovad->rbroot);
-       else {
-               struct rb_node *prev_node = rb_prev(iovad->cached32_node);
-               struct iova *curr_iova =
-                       rb_entry(iovad->cached32_node, struct iova, node);
-               *limit_pfn = curr_iova->pfn_lo;
-               return prev_node;
-       }
+
+       curr_iova = rb_entry(cached_node, struct iova, node);
+       *limit_pfn = min(*limit_pfn, curr_iova->pfn_lo);
+
+       return rb_prev(cached_node);
 }
 
 static void
-__cached_rbnode_insert_update(struct iova_domain *iovad,
-       unsigned long limit_pfn, struct iova *new)
+__cached_rbnode_insert_update(struct iova_domain *iovad, struct iova *new)
 {
-       if (limit_pfn != iovad->dma_32bit_pfn)
-               return;
-       iovad->cached32_node = &new->node;
+       if (new->pfn_hi < iovad->dma_32bit_pfn)
+               iovad->cached32_node = &new->node;
+       else
+               iovad->cached_node = &new->node;
 }
 
 static void
 __cached_rbnode_delete_update(struct iova_domain *iovad, struct iova *free)
 {
        struct iova *cached_iova;
-       struct rb_node *curr;
 
-       if (!iovad->cached32_node)
-               return;
-       curr = iovad->cached32_node;
-       cached_iova = rb_entry(curr, struct iova, node);
+       cached_iova = rb_entry(iovad->cached32_node, struct iova, node);
+       if (free->pfn_hi < iovad->dma_32bit_pfn &&
+           iovad->cached32_node && free->pfn_lo >= cached_iova->pfn_lo)
+               iovad->cached32_node = rb_next(&free->node);
 
-       if (free->pfn_lo >= cached_iova->pfn_lo) {
-               struct rb_node *node = rb_next(&free->node);
-               struct iova *iova = rb_entry(node, struct iova, node);
-
-               /* only cache if it's below 32bit pfn */
-               if (node && iova->pfn_lo < iovad->dma_32bit_pfn)
-                       iovad->cached32_node = node;
-               else
-                       iovad->cached32_node = NULL;
-       }
+       cached_iova = rb_entry(iovad->cached_node, struct iova, node);
+       if (iovad->cached_node && free->pfn_lo >= cached_iova->pfn_lo)
+               iovad->cached_node = rb_next(&free->node);
 }
 
 /* Insert the iova into domain rbtree by holding writer lock */
 {
        struct rb_node *prev, *curr = NULL;
        unsigned long flags;
-       unsigned long saved_pfn, new_pfn;
+       unsigned long new_pfn;
        unsigned long align_mask = ~0UL;
 
        if (size_aligned)
 
        /* Walk the tree backwards */
        spin_lock_irqsave(&iovad->iova_rbtree_lock, flags);
-       saved_pfn = limit_pfn;
        curr = __get_cached_rbnode(iovad, &limit_pfn);
        prev = curr;
        while (curr) {
 
        /* If we have 'prev', it's a valid place to start the insertion. */
        iova_insert_rbtree(&iovad->rbroot, new, prev);
-       __cached_rbnode_insert_update(iovad, saved_pfn, new);
+       __cached_rbnode_insert_update(iovad, new);
 
        spin_unlock_irqrestore(&iovad->iova_rbtree_lock, flags);