3 files changed, 24 insertions, 23 deletions

diff --git a/drivers/gpu/drm/xe/xe_device.c b/drivers/gpu/drm/xe/xe_device.c
index 6d01932f934c..89a85f193f3a 100644
--- a/drivers/gpu/drm/xe/xe_device.c
+++ b/drivers/gpu/drm/xe/xe_device.c
@@ -877,15 +877,17 @@ int xe_device_probe(struct xe_device *xe)
 
 	err = xe_pxp_init(xe);
 	if (err)
-		goto err_fini_display;
+		goto err_remove_display;
 
 	err = drm_dev_register(&xe->drm, 0);
 	if (err)
-		goto err_fini_display;
+		goto err_remove_display;
 
 	xe_display_register(xe);
 
-	xe_oa_register(xe);
+	err = xe_oa_register(xe);
+	if (err)
+		goto err_unregister_display;
 
 	xe_pmu_register(&xe->pmu);
 
@@ -902,7 +904,9 @@ int xe_device_probe(struct xe_device *xe)
 
 	return devm_add_action_or_reset(xe->drm.dev, xe_device_sanitize, xe);
 
-err_fini_display:
+err_unregister_display:
+	xe_display_unregister(xe);
+err_remove_display:
 	xe_display_driver_remove(xe);
 
 	return err;
@@ -971,8 +975,6 @@ void xe_device_remove(struct xe_device *xe)
 
 	xe_display_driver_remove(xe);
 
-	xe_oa_unregister(xe);
-
 	xe_heci_gsc_fini(xe);
 
 	xe_device_call_remove_actions(xe);
diff --git a/drivers/gpu/drm/xe/xe_oa.c b/drivers/gpu/drm/xe/xe_oa.c
index 2c640185bdec..d89e6cabf5a5 100644
--- a/drivers/gpu/drm/xe/xe_oa.c
+++ b/drivers/gpu/drm/xe/xe_oa.c
@@ -2423,36 +2423,36 @@ err_unlock:
 	return ret;
 }
 
+static void xe_oa_unregister(void *arg)
+{
+	struct xe_oa *oa = arg;
+
+	if (!oa->metrics_kobj)
+		return;
+
+	kobject_put(oa->metrics_kobj);
+	oa->metrics_kobj = NULL;
+}
+
 /**
  * xe_oa_register - Xe OA registration
  * @xe: @xe_device
  *
  * Exposes the metrics sysfs directory upon completion of module initialization
  */
-void xe_oa_register(struct xe_device *xe)
+int xe_oa_register(struct xe_device *xe)
 {
 	struct xe_oa *oa = &xe->oa;
 
 	if (!oa->xe)
-		return;
+		return 0;
 
 	oa->metrics_kobj = kobject_create_and_add("metrics",
 						  &xe->drm.primary->kdev->kobj);
-}
-
-/**
- * xe_oa_unregister - Xe OA de-registration
- * @xe: @xe_device
- */
-void xe_oa_unregister(struct xe_device *xe)
-{
-	struct xe_oa *oa = &xe->oa;
-
 	if (!oa->metrics_kobj)
-		return;
+		return -ENOMEM;
 
-	kobject_put(oa->metrics_kobj);
-	oa->metrics_kobj = NULL;
+	return devm_add_action_or_reset(xe->drm.dev, xe_oa_unregister, oa);
 }
 
 static u32 num_oa_units_per_gt(struct xe_gt *gt)
diff --git a/drivers/gpu/drm/xe/xe_oa.h b/drivers/gpu/drm/xe/xe_oa.h
index eb36ce250c61..e510826f9efc 100644
--- a/drivers/gpu/drm/xe/xe_oa.h
+++ b/drivers/gpu/drm/xe/xe_oa.h
@@ -15,8 +15,7 @@ struct xe_gt;
 struct xe_hw_engine;
 
 int xe_oa_init(struct xe_device *xe);
-void xe_oa_register(struct xe_device *xe);
-void xe_oa_unregister(struct xe_device *xe);
+int xe_oa_register(struct xe_device *xe);
 int xe_oa_stream_open_ioctl(struct drm_device *dev, u64 data, struct drm_file *file);
 int xe_oa_add_config_ioctl(struct drm_device *dev, u64 data, struct drm_file *file);
 int xe_oa_remove_config_ioctl(struct drm_device *dev, u64 data, struct drm_file *file);