int iommu_take_ownership(struct iommu_table *tbl)
 {
-       unsigned long sz = (tbl->it_size + 7) >> 3;
+       unsigned long flags, i, sz = (tbl->it_size + 7) >> 3;
+       int ret = 0;
+
+       spin_lock_irqsave(&tbl->large_pool.lock, flags);
+       for (i = 0; i < tbl->nr_pools; i++)
+               spin_lock(&tbl->pools[i].lock);
 
        if (tbl->it_offset == 0)
                clear_bit(0, tbl->it_map);
 
        if (!bitmap_empty(tbl->it_map, tbl->it_size)) {
                pr_err("iommu_tce: it_map is not empty");
-               return -EBUSY;
+               ret = -EBUSY;
+               /* Restore bit#0 set by iommu_init_table() */
+               if (tbl->it_offset == 0)
+                       set_bit(0, tbl->it_map);
+       } else {
+               memset(tbl->it_map, 0xff, sz);
        }
 
-       memset(tbl->it_map, 0xff, sz);
+       for (i = 0; i < tbl->nr_pools; i++)
+               spin_unlock(&tbl->pools[i].lock);
+       spin_unlock_irqrestore(&tbl->large_pool.lock, flags);
 
-       return 0;
+       return ret;
 }
 EXPORT_SYMBOL_GPL(iommu_take_ownership);
 
 void iommu_release_ownership(struct iommu_table *tbl)
 {
-       unsigned long sz = (tbl->it_size + 7) >> 3;
+       unsigned long flags, i, sz = (tbl->it_size + 7) >> 3;
+
+       spin_lock_irqsave(&tbl->large_pool.lock, flags);
+       for (i = 0; i < tbl->nr_pools; i++)
+               spin_lock(&tbl->pools[i].lock);
 
        memset(tbl->it_map, 0, sz);
 
        /* Restore bit#0 set by iommu_init_table() */
        if (tbl->it_offset == 0)
                set_bit(0, tbl->it_map);
+
+       for (i = 0; i < tbl->nr_pools; i++)
+               spin_unlock(&tbl->pools[i].lock);
+       spin_unlock_irqrestore(&tbl->large_pool.lock, flags);
 }
 EXPORT_SYMBOL_GPL(iommu_release_ownership);