1 files changed, 37 insertions, 42 deletions

diff --git a/drivers/dma/ioat/pci.c b/drivers/dma/ioat/pci.c
index 982e38fd177c..55414d88ac1b 100644
--- a/drivers/dma/ioat/pci.c
+++ b/drivers/dma/ioat/pci.c
@@ -60,14 +60,8 @@ static struct pci_device_id ioat_pci_tbl[] = {
         { 0, }
 };
 
-struct ioat_device {
-        struct pci_dev          *pdev;
-        struct ioatdma_device   *dma;
-        struct dca_provider     *dca;
-};
-
-static int __devinit ioat_probe(struct pci_dev *pdev,
-                                const struct pci_device_id *id);
+static int __devinit ioat_pci_probe(struct pci_dev *pdev,
+                                    const struct pci_device_id *id);
 static void __devexit ioat_remove(struct pci_dev *pdev);
 
 static int ioat_dca_enabled = 1;
@@ -79,17 +73,28 @@ MODULE_PARM_DESC(ioat_dca_enabled, "control support of dca service (default: 1)"
 static struct pci_driver ioat_pci_driver = {
         .name           = DRV_NAME,
         .id_table       = ioat_pci_tbl,
-        .probe          = ioat_probe,
+        .probe          = ioat_pci_probe,
         .remove         = __devexit_p(ioat_remove),
 };
 
-static int __devinit ioat_probe(struct pci_dev *pdev,
-                                const struct pci_device_id *id)
+static struct ioatdma_device *
+alloc_ioatdma(struct pci_dev *pdev, void __iomem *iobase)
+{
+        struct device *dev = &pdev->dev;
+        struct ioatdma_device *d = devm_kzalloc(dev, sizeof(*d), GFP_KERNEL);
+
+        if (!d)
+                return NULL;
+        d->pdev = pdev;
+        d->reg_base = iobase;
+        return d;
+}
+
+static int __devinit ioat_pci_probe(struct pci_dev *pdev, const struct pci_device_id *id)
 {
         void __iomem * const *iomap;
-        void __iomem *iobase;
         struct device *dev = &pdev->dev;
-        struct ioat_device *device;
+        struct ioatdma_device *device;
         int err;
 
         err = pcim_enable_device(pdev);
@@ -119,33 +124,24 @@ static int __devinit ioat_probe(struct pci_dev *pdev,
         if (!device)
                 return -ENOMEM;
 
-        device->pdev = pdev;
-        pci_set_drvdata(pdev, device);
-        iobase = iomap[IOAT_MMIO_BAR];
-
         pci_set_master(pdev);
 
-        switch (readb(iobase + IOAT_VER_OFFSET)) {
-        case IOAT_VER_1_2:
-                device->dma = ioat_dma_probe(pdev, iobase);
-                if (device->dma && ioat_dca_enabled)
-                        device->dca = ioat_dca_init(pdev, iobase);
-                break;
-        case IOAT_VER_2_0:
-                device->dma = ioat_dma_probe(pdev, iobase);
-                if (device->dma && ioat_dca_enabled)
-                        device->dca = ioat2_dca_init(pdev, iobase);
-                break;
-        case IOAT_VER_3_0:
-                device->dma = ioat_dma_probe(pdev, iobase);
-                if (device->dma && ioat_dca_enabled)
-                        device->dca = ioat3_dca_init(pdev, iobase);
-                break;
-        default:
+        device = alloc_ioatdma(pdev, iomap[IOAT_MMIO_BAR]);
+        if (!device)
+                return -ENOMEM;
+        pci_set_drvdata(pdev, device);
+
+        device->version = readb(device->reg_base + IOAT_VER_OFFSET);
+        if (device->version == IOAT_VER_1_2)
+                err = ioat1_dma_probe(device, ioat_dca_enabled);
+        else if (device->version == IOAT_VER_2_0)
+                err = ioat2_dma_probe(device, ioat_dca_enabled);
+        else if (device->version >= IOAT_VER_3_0)
+                err = ioat3_dma_probe(device, ioat_dca_enabled);
+        else
                 return -ENODEV;
-        }
 
-        if (!device->dma) {
+        if (err) {
                 dev_err(dev, "Intel(R) I/OAT DMA Engine init failed\n");
                 return -ENODEV;
         }
@@ -155,7 +151,10 @@ static int __devinit ioat_probe(struct pci_dev *pdev,
 
 static void __devexit ioat_remove(struct pci_dev *pdev)
 {
-        struct ioat_device *device = pci_get_drvdata(pdev);
+        struct ioatdma_device *device = pci_get_drvdata(pdev);
+
+        if (!device)
+                return;
 
         dev_err(&pdev->dev, "Removing dma and dca services\n");
         if (device->dca) {
@@ -163,11 +162,7 @@ static void __devexit ioat_remove(struct pci_dev *pdev)
                 free_dca_provider(device->dca);
                 device->dca = NULL;
         }
-
-        if (device->dma) {
-                ioat_dma_remove(device->dma);
-                device->dma = NULL;
-        }
+        ioat_dma_remove(device);
 }
 
 static int __init ioat_init_module(void)