#include <linux/module.h>
 #include <linux/slab.h>
 #include <linux/of.h>
+#include <linux/dma-mapping.h>
 
 #include "xhci.h"
 
        if (!res)
                return -ENODEV;
 
+       /* Initialize dma_mask and coherent_dma_mask to 32-bits */
+       ret = dma_set_coherent_mask(&pdev->dev, DMA_BIT_MASK(32));
+       if (ret)
+               return ret;
+       if (!pdev->dev.dma_mask)
+               pdev->dev.dma_mask = &pdev->dev.coherent_dma_mask;
+       else
+               dma_set_mask(&pdev->dev, DMA_BIT_MASK(32));
+
        hcd = usb_create_hcd(driver, &pdev->dev, dev_name(&pdev->dev));
        if (!hcd)
                return -ENOMEM;
 
        struct xhci_hcd         *xhci;
        struct device           *dev = hcd->self.controller;
        int                     retval;
-       u32                     temp;
 
        /* Accept arbitrarily long scatter-gather lists */
        hcd->self.sg_tablesize = ~0;
                /* xHCI private pointer was set in xhci_pci_probe for the second
                 * registered roothub.
                 */
-               xhci = hcd_to_xhci(hcd);
-               temp = xhci_readl(xhci, &xhci->cap_regs->hcc_params);
-               if (HCC_64BIT_ADDR(temp)) {
-                       xhci_dbg(xhci, "Enabling 64-bit DMA addresses.\n");
-                       dma_set_mask(hcd->self.controller, DMA_BIT_MASK(64));
-               } else {
-                       dma_set_mask(hcd->self.controller, DMA_BIT_MASK(32));
-               }
                return 0;
        }
 
                goto error;
        xhci_dbg(xhci, "Reset complete\n");
 
-       temp = xhci_readl(xhci, &xhci->cap_regs->hcc_params);
-       if (HCC_64BIT_ADDR(temp)) {
+       /* Set dma_mask and coherent_dma_mask to 64-bits,
+        * if xHC supports 64-bit addressing */
+       if (HCC_64BIT_ADDR(xhci->hcc_params) &&
+                       !dma_set_mask(dev, DMA_BIT_MASK(64))) {
                xhci_dbg(xhci, "Enabling 64-bit DMA addresses.\n");
-               dma_set_mask(hcd->self.controller, DMA_BIT_MASK(64));
-       } else {
-               dma_set_mask(hcd->self.controller, DMA_BIT_MASK(32));
+               dma_set_coherent_mask(dev, DMA_BIT_MASK(64));
        }
 
        xhci_dbg(xhci, "Calling HCD init\n");