2 files changed, 28 insertions, 7 deletions

diff --git a/arch/s390/include/asm/pci_io.h b/arch/s390/include/asm/pci_io.h
index f664e96f48c7..1a9a98de5bde 100644
--- a/arch/s390/include/asm/pci_io.h
+++ b/arch/s390/include/asm/pci_io.h
@@ -16,6 +16,7 @@
 struct zpci_iomap_entry {
         u32 fh;
         u8 bar;
+        u16 count;
 };
 
 extern struct zpci_iomap_entry *zpci_iomap_start;
diff --git a/arch/s390/pci/pci.c b/arch/s390/pci/pci.c
index 3290f11ae1d9..753a56731951 100644
--- a/arch/s390/pci/pci.c
+++ b/arch/s390/pci/pci.c
@@ -259,7 +259,10 @@ void __iowrite64_copy(void __iomem *to, const void *from, size_t count)
 }
 
 /* Create a virtual mapping cookie for a PCI BAR */
-void __iomem *pci_iomap(struct pci_dev *pdev, int bar, unsigned long max)
+void __iomem *pci_iomap_range(struct pci_dev *pdev,
+                              int bar,
+                              unsigned long offset,
+                              unsigned long max)
 {
         struct zpci_dev *zdev = get_zdev(pdev);
         u64 addr;
@@ -270,14 +273,27 @@ void __iomem *pci_iomap(struct pci_dev *pdev, int bar, unsigned long max)
 
         idx = zdev->bars[bar].map_idx;
         spin_lock(&zpci_iomap_lock);
-        zpci_iomap_start[idx].fh = zdev->fh;
-        zpci_iomap_start[idx].bar = bar;
+        if (zpci_iomap_start[idx].count++) {
+                BUG_ON(zpci_iomap_start[idx].fh != zdev->fh ||
+                       zpci_iomap_start[idx].bar != bar);
+        } else {
+                zpci_iomap_start[idx].fh = zdev->fh;
+                zpci_iomap_start[idx].bar = bar;
+        }
+        /* Detect overrun */
+        BUG_ON(!zpci_iomap_start[idx].count);
         spin_unlock(&zpci_iomap_lock);
 
         addr = ZPCI_IOMAP_ADDR_BASE | ((u64) idx << 48);
-        return (void __iomem *) addr;
+        return (void __iomem *) addr + offset;
 }
-EXPORT_SYMBOL_GPL(pci_iomap);
+EXPORT_SYMBOL_GPL(pci_iomap_range);
+
+void __iomem *pci_iomap(struct pci_dev *dev, int bar, unsigned long maxlen)
+{
+        return pci_iomap_range(dev, bar, 0, maxlen);
+}
+EXPORT_SYMBOL(pci_iomap);
 
 void pci_iounmap(struct pci_dev *pdev, void __iomem *addr)
 {
@@ -285,8 +301,12 @@ void pci_iounmap(struct pci_dev *pdev, void __iomem *addr)
 
         idx = (((__force u64) addr) & ~ZPCI_IOMAP_ADDR_BASE) >> 48;
         spin_lock(&zpci_iomap_lock);
-        zpci_iomap_start[idx].fh = 0;
-        zpci_iomap_start[idx].bar = 0;
+        /* Detect underrun */
+        BUG_ON(!zpci_iomap_start[idx].count);
+        if (!--zpci_iomap_start[idx].count) {
+                zpci_iomap_start[idx].fh = 0;
+                zpci_iomap_start[idx].bar = 0;
+        }
         spin_unlock(&zpci_iomap_lock);
 }
 EXPORT_SYMBOL_GPL(pci_iounmap);