2 files changed, 4 insertions, 2 deletions

diff --git a/arch/arm/mm/dma-mapping.c b/arch/arm/mm/dma-mapping.c
index 1a77450e728a..11b3914660d2 100644
--- a/arch/arm/mm/dma-mapping.c
+++ b/arch/arm/mm/dma-mapping.c
@@ -1358,7 +1358,7 @@ static void *arm_iommu_alloc_attrs(struct device *dev, size_t size,
         *handle = DMA_ERROR_CODE;
         size = PAGE_ALIGN(size);
 
-        if (gfp & GFP_ATOMIC)
+        if (!(gfp & __GFP_WAIT))
                 return __iommu_alloc_atomic(dev, size, handle);
 
         /*
diff --git a/arch/x86/kernel/pci-dma.c b/arch/x86/kernel/pci-dma.c
index 872079a67e4d..f7d0672481fd 100644
--- a/arch/x86/kernel/pci-dma.c
+++ b/arch/x86/kernel/pci-dma.c
@@ -100,8 +100,10 @@ void *dma_generic_alloc_coherent(struct device *dev, size_t size,
         flag |= __GFP_ZERO;
 again:
         page = NULL;
-        if (!(flag & GFP_ATOMIC))
+        /* CMA can be used only in the context which permits sleeping */
+        if (flag & __GFP_WAIT)
                 page = dma_alloc_from_contiguous(dev, count, get_order(size));
+        /* fallback */
         if (!page)
                 page = alloc_pages_node(dev_to_node(dev), flag, get_order(size));
         if (!page)