arch_dma_prep_coherent(sg_page(sg), sg->length);
        }
 
-       ret = iommu_map_sg_atomic(domain, iova, sgt->sgl, sgt->orig_nents, ioprot);
+       ret = iommu_map_sg(domain, iova, sgt->sgl, sgt->orig_nents, ioprot,
+                          GFP_ATOMIC);
        if (ret < 0 || ret < size)
                goto out_free_sg;
 
         * We'll leave any physical concatenation to the IOMMU driver's
         * implementation - it knows better than we do.
         */
-       ret = iommu_map_sg_atomic(domain, iova, sg, nents, prot);
+       ret = iommu_map_sg(domain, iova, sg, nents, prot, GFP_ATOMIC);
        if (ret < 0 || ret < iova_len)
                goto out_free_iova;
 
 
 }
 EXPORT_SYMBOL_GPL(iommu_unmap_fast);
 
-static ssize_t __iommu_map_sg(struct iommu_domain *domain, unsigned long iova,
-               struct scatterlist *sg, unsigned int nents, int prot,
-               gfp_t gfp)
+ssize_t iommu_map_sg(struct iommu_domain *domain, unsigned long iova,
+                    struct scatterlist *sg, unsigned int nents, int prot,
+                    gfp_t gfp)
 {
        const struct iommu_domain_ops *ops = domain->ops;
        size_t len = 0, mapped = 0;
        unsigned int i = 0;
        int ret;
 
+       might_sleep_if(gfpflags_allow_blocking(gfp));
+
+       /* Discourage passing strange GFP flags */
+       if (WARN_ON_ONCE(gfp & (__GFP_COMP | __GFP_DMA | __GFP_DMA32 |
+                               __GFP_HIGHMEM)))
+               return -EINVAL;
+
        while (i <= nents) {
                phys_addr_t s_phys = sg_phys(sg);
 
 
        return ret;
 }
-
-ssize_t iommu_map_sg(struct iommu_domain *domain, unsigned long iova,
-                    struct scatterlist *sg, unsigned int nents, int prot)
-{
-       might_sleep();
-       return __iommu_map_sg(domain, iova, sg, nents, prot, GFP_KERNEL);
-}
 EXPORT_SYMBOL_GPL(iommu_map_sg);
 
-ssize_t iommu_map_sg_atomic(struct iommu_domain *domain, unsigned long iova,
-                   struct scatterlist *sg, unsigned int nents, int prot)
-{
-       return __iommu_map_sg(domain, iova, sg, nents, prot, GFP_ATOMIC);
-}
-
 /**
  * report_iommu_fault() - report about an IOMMU fault to the IOMMU framework
  * @domain: the iommu domain where the fault has happened
 
                               unsigned long iova, size_t size,
                               struct iommu_iotlb_gather *iotlb_gather);
 extern ssize_t iommu_map_sg(struct iommu_domain *domain, unsigned long iova,
-               struct scatterlist *sg, unsigned int nents, int prot);
-extern ssize_t iommu_map_sg_atomic(struct iommu_domain *domain,
-                                  unsigned long iova, struct scatterlist *sg,
-                                  unsigned int nents, int prot);
+                           struct scatterlist *sg, unsigned int nents,
+                           int prot, gfp_t gfp);
 extern phys_addr_t iommu_iova_to_phys(struct iommu_domain *domain, dma_addr_t iova);
 extern void iommu_set_fault_handler(struct iommu_domain *domain,
                        iommu_fault_handler_t handler, void *token);
 
 static inline ssize_t iommu_map_sg(struct iommu_domain *domain,
                                   unsigned long iova, struct scatterlist *sg,
-                                  unsigned int nents, int prot)
-{
-       return -ENODEV;
-}
-
-static inline ssize_t iommu_map_sg_atomic(struct iommu_domain *domain,
-                                 unsigned long iova, struct scatterlist *sg,
-                                 unsigned int nents, int prot)
+                                  unsigned int nents, int prot, gfp_t gfp)
 {
        return -ENODEV;
 }
 static inline size_t iommu_map_sgtable(struct iommu_domain *domain,
                        unsigned long iova, struct sg_table *sgt, int prot)
 {
-       return iommu_map_sg(domain, iova, sgt->sgl, sgt->orig_nents, prot);
+       return iommu_map_sg(domain, iova, sgt->sgl, sgt->orig_nents, prot,
+                           GFP_KERNEL);
 }
 
 #ifdef CONFIG_IOMMU_DEBUGFS