1 files changed, 10 insertions, 3 deletions

diff --git a/drivers/infiniband/core/device.c b/drivers/infiniband/core/device.c
index 8872453e26c0..238ec42778ef 100644
--- a/drivers/infiniband/core/device.c
+++ b/drivers/infiniband/core/device.c
@@ -156,19 +156,26 @@ struct ib_device *ib_device_get_by_index(u32 index)
         down_read(&lists_rwsem);
         device = __ib_device_get_by_index(index);
         if (device) {
-                /* Do not return a device if unregistration has started. */
-                if (!refcount_inc_not_zero(&device->refcount))
+                if (!ib_device_try_get(device))
                         device = NULL;
         }
         up_read(&lists_rwsem);
         return device;
 }
 
+/**
+ * ib_device_put - Release IB device reference
+ * @device: device whose reference to be released
+ *
+ * ib_device_put() releases reference to the IB device to allow it to be
+ * unregistered and eventually free.
+ */
 void ib_device_put(struct ib_device *device)
 {
         if (refcount_dec_and_test(&device->refcount))
                 complete(&device->unreg_completion);
 }
+EXPORT_SYMBOL(ib_device_put);
 
 static struct ib_device *__ib_device_get_by_name(const char *name)
 {
@@ -303,7 +310,6 @@ struct ib_device *ib_alloc_device(size_t size)
         rwlock_init(&device->client_data_lock);
         INIT_LIST_HEAD(&device->client_data_list);
         INIT_LIST_HEAD(&device->port_list);
-        refcount_set(&device->refcount, 1);
         init_completion(&device->unreg_completion);
 
         return device;
@@ -620,6 +626,7 @@ int ib_register_device(struct ib_device *device, const char *name,
                 goto cg_cleanup;
         }
 
+        refcount_set(&device->refcount, 1);
         device->reg_state = IB_DEV_REGISTERED;
 
         list_for_each_entry(client, &client_list, list)