}
        }
 
-       if (HFI1_CAP_IS_USET(TID_UNMAP)) {
-               fd->mmu_rb_insert = mmu_rb_insert;
-               fd->mmu_rb_remove = mmu_rb_remove;
-       } else {
-               fd->mmu_rb_insert = hfi1_mmu_rb_insert;
-               fd->mmu_rb_remove = hfi1_mmu_rb_remove;
-       }
-
        /*
         * PSM does not have a good way to separate, count, and
         * effectively enforce a limit on RcvArray entries used by
        node->freed = false;
        memcpy(node->pages, pages, sizeof(struct page *) * npages);
 
-       ret = fd->mmu_rb_insert(root, &node->mmu);
+       if (HFI1_CAP_IS_USET(TID_UNMAP))
+               ret = mmu_rb_insert(root, &node->mmu);
+       else
+               ret = hfi1_mmu_rb_insert(root, &node->mmu);
 
        if (ret) {
                hfi1_cdbg(TID, "Failed to insert RB node %u 0x%lx, 0x%lx %d",
        node = fd->entry_to_rb[rcventry];
        if (!node || node->rcventry != (uctxt->expected_base + rcventry))
                return -EBADF;
-       fd->mmu_rb_remove(&fd->tid_rb_root, &node->mmu);
+       if (HFI1_CAP_IS_USET(TID_UNMAP))
+               mmu_rb_remove(&fd->tid_rb_root, &node->mmu);
+       else
+               hfi1_mmu_rb_remove(&fd->tid_rb_root, &node->mmu);
 
        if (grp)
                *grp = node->grp;
                                                          uctxt->expected_base];
                                if (!node || node->rcventry != rcventry)
                                        continue;
-                               fd->mmu_rb_remove(root, &node->mmu);
+                               if (HFI1_CAP_IS_USET(TID_UNMAP))
+                                       mmu_rb_remove(&fd->tid_rb_root,
+                                                     &node->mmu);
+                               else
+                                       hfi1_mmu_rb_remove(&fd->tid_rb_root,
+                                                          &node->mmu);
                                clear_tid_node(fd, -1, node);
                        }
                }