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)