goto err;
        }
 
-       orig_src_addr = dma_alloc_coherent(dev, size + alignment,
-                                          &orig_src_phys_addr, GFP_KERNEL);
+       orig_src_addr = kzalloc(size + alignment, GFP_KERNEL);
        if (!orig_src_addr) {
                dev_err(dev, "Failed to allocate source buffer\n");
                ret = false;
                goto err;
        }
 
+       get_random_bytes(orig_src_addr, size + alignment);
+       orig_src_phys_addr = dma_map_single(dev, orig_src_addr,
+                                           size + alignment, DMA_TO_DEVICE);
+       if (dma_mapping_error(dev, orig_src_phys_addr)) {
+               dev_err(dev, "failed to map source buffer address\n");
+               ret = false;
+               goto err_src_phys_addr;
+       }
+
        if (alignment && !IS_ALIGNED(orig_src_phys_addr, alignment)) {
                src_phys_addr = PTR_ALIGN(orig_src_phys_addr, alignment);
                offset = src_phys_addr - orig_src_phys_addr;
        pci_endpoint_test_writel(test, PCI_ENDPOINT_TEST_UPPER_SRC_ADDR,
                                 upper_32_bits(src_phys_addr));
 
-       get_random_bytes(src_addr, size);
        src_crc32 = crc32_le(~0, src_addr, size);
 
-       orig_dst_addr = dma_alloc_coherent(dev, size + alignment,
-                                          &orig_dst_phys_addr, GFP_KERNEL);
+       orig_dst_addr = kzalloc(size + alignment, GFP_KERNEL);
        if (!orig_dst_addr) {
                dev_err(dev, "Failed to allocate destination address\n");
                ret = false;
-               goto err_orig_src_addr;
+               goto err_dst_addr;
+       }
+
+       orig_dst_phys_addr = dma_map_single(dev, orig_dst_addr,
+                                           size + alignment, DMA_FROM_DEVICE);
+       if (dma_mapping_error(dev, orig_dst_phys_addr)) {
+               dev_err(dev, "failed to map destination buffer address\n");
+               ret = false;
+               goto err_dst_phys_addr;
        }
 
        if (alignment && !IS_ALIGNED(orig_dst_phys_addr, alignment)) {
 
        wait_for_completion(&test->irq_raised);
 
+       dma_unmap_single(dev, orig_dst_phys_addr, size + alignment,
+                        DMA_FROM_DEVICE);
+
        dst_crc32 = crc32_le(~0, dst_addr, size);
        if (dst_crc32 == src_crc32)
                ret = true;
 
-       dma_free_coherent(dev, size + alignment, orig_dst_addr,
-                         orig_dst_phys_addr);
+err_dst_phys_addr:
+       kfree(orig_dst_addr);
 
-err_orig_src_addr:
-       dma_free_coherent(dev, size + alignment, orig_src_addr,
-                         orig_src_phys_addr);
+err_dst_addr:
+       dma_unmap_single(dev, orig_src_phys_addr, size + alignment,
+                        DMA_TO_DEVICE);
+
+err_src_phys_addr:
+       kfree(orig_src_addr);
 
 err:
        return ret;
                goto err;
        }
 
-       orig_addr = dma_alloc_coherent(dev, size + alignment, &orig_phys_addr,
-                                      GFP_KERNEL);
+       orig_addr = kzalloc(size + alignment, GFP_KERNEL);
        if (!orig_addr) {
                dev_err(dev, "Failed to allocate address\n");
                ret = false;
                goto err;
        }
 
+       get_random_bytes(orig_addr, size + alignment);
+
+       orig_phys_addr = dma_map_single(dev, orig_addr, size + alignment,
+                                       DMA_TO_DEVICE);
+       if (dma_mapping_error(dev, orig_phys_addr)) {
+               dev_err(dev, "failed to map source buffer address\n");
+               ret = false;
+               goto err_phys_addr;
+       }
+
        if (alignment && !IS_ALIGNED(orig_phys_addr, alignment)) {
                phys_addr =  PTR_ALIGN(orig_phys_addr, alignment);
                offset = phys_addr - orig_phys_addr;
                addr = orig_addr;
        }
 
-       get_random_bytes(addr, size);
-
        crc32 = crc32_le(~0, addr, size);
        pci_endpoint_test_writel(test, PCI_ENDPOINT_TEST_CHECKSUM,
                                 crc32);
        if (reg & STATUS_READ_SUCCESS)
                ret = true;
 
-       dma_free_coherent(dev, size + alignment, orig_addr, orig_phys_addr);
+       dma_unmap_single(dev, orig_phys_addr, size + alignment,
+                        DMA_TO_DEVICE);
+
+err_phys_addr:
+       kfree(orig_addr);
 
 err:
        return ret;
                goto err;
        }
 
-       orig_addr = dma_alloc_coherent(dev, size + alignment, &orig_phys_addr,
-                                      GFP_KERNEL);
+       orig_addr = kzalloc(size + alignment, GFP_KERNEL);
        if (!orig_addr) {
                dev_err(dev, "Failed to allocate destination address\n");
                ret = false;
                goto err;
        }
 
+       orig_phys_addr = dma_map_single(dev, orig_addr, size + alignment,
+                                       DMA_FROM_DEVICE);
+       if (dma_mapping_error(dev, orig_phys_addr)) {
+               dev_err(dev, "failed to map source buffer address\n");
+               ret = false;
+               goto err_phys_addr;
+       }
+
        if (alignment && !IS_ALIGNED(orig_phys_addr, alignment)) {
                phys_addr = PTR_ALIGN(orig_phys_addr, alignment);
                offset = phys_addr - orig_phys_addr;
 
        wait_for_completion(&test->irq_raised);
 
+       dma_unmap_single(dev, orig_phys_addr, size + alignment,
+                        DMA_FROM_DEVICE);
+
        crc32 = crc32_le(~0, addr, size);
        if (crc32 == pci_endpoint_test_readl(test, PCI_ENDPOINT_TEST_CHECKSUM))
                ret = true;
 
-       dma_free_coherent(dev, size + alignment, orig_addr, orig_phys_addr);
+err_phys_addr:
+       kfree(orig_addr);
 err:
        return ret;
 }
        init_completion(&test->irq_raised);
        mutex_init(&test->mutex);
 
+       if ((dma_set_mask_and_coherent(&pdev->dev, DMA_BIT_MASK(48)) != 0) &&
+           dma_set_mask_and_coherent(&pdev->dev, DMA_BIT_MASK(32)) != 0) {
+               dev_err(dev, "Cannot set DMA mask\n");
+               return -EINVAL;
+       }
+
        err = pci_enable_device(pdev);
        if (err) {
                dev_err(dev, "Cannot enable PCI device\n");
        pci_disable_device(pdev);
 }
 
+static const struct pci_endpoint_test_data default_data = {
+       .test_reg_bar = BAR_0,
+       .alignment = SZ_4K,
+       .irq_type = IRQ_TYPE_MSI,
+};
+
 static const struct pci_endpoint_test_data am654_data = {
        .test_reg_bar = BAR_2,
        .alignment = SZ_64K,
 };
 
 static const struct pci_device_id pci_endpoint_test_tbl[] = {
-       { PCI_DEVICE(PCI_VENDOR_ID_TI, PCI_DEVICE_ID_TI_DRA74x) },
-       { PCI_DEVICE(PCI_VENDOR_ID_TI, PCI_DEVICE_ID_TI_DRA72x) },
+       { PCI_DEVICE(PCI_VENDOR_ID_TI, PCI_DEVICE_ID_TI_DRA74x),
+         .driver_data = (kernel_ulong_t)&default_data,
+       },
+       { PCI_DEVICE(PCI_VENDOR_ID_TI, PCI_DEVICE_ID_TI_DRA72x),
+         .driver_data = (kernel_ulong_t)&default_data,
+       },
        { PCI_DEVICE(PCI_VENDOR_ID_FREESCALE, 0x81c0) },
        { PCI_DEVICE_DATA(SYNOPSYS, EDDA, NULL) },
        { PCI_DEVICE(PCI_VENDOR_ID_TI, PCI_DEVICE_ID_TI_AM654),