1 files changed, 8 insertions, 4 deletions

diff --git a/drivers/dax/pmem.c b/drivers/dax/pmem.c
index fd49b24fd6af..99e2aace8078 100644
--- a/drivers/dax/pmem.c
+++ b/drivers/dax/pmem.c
@@ -105,15 +105,19 @@ static int dax_pmem_probe(struct device *dev)
         if (rc)
                 return rc;
 
-        rc = devm_add_action_or_reset(dev, dax_pmem_percpu_exit,
-                                                        &dax_pmem->ref);
-        if (rc)
+        rc = devm_add_action(dev, dax_pmem_percpu_exit, &dax_pmem->ref);
+        if (rc) {
+                percpu_ref_exit(&dax_pmem->ref);
                 return rc;
+        }
 
         dax_pmem->pgmap.ref = &dax_pmem->ref;
         addr = devm_memremap_pages(dev, &dax_pmem->pgmap);
-        if (IS_ERR(addr))
+        if (IS_ERR(addr)) {
+                devm_remove_action(dev, dax_pmem_percpu_exit, &dax_pmem->ref);
+                percpu_ref_exit(&dax_pmem->ref);
                 return PTR_ERR(addr);
+        }
 
         rc = devm_add_action_or_reset(dev, dax_pmem_percpu_kill,
                                                         &dax_pmem->ref);