1 files changed, 32 insertions, 22 deletions

diff --git a/drivers/misc/habanalabs/common/context.c b/drivers/misc/habanalabs/common/context.c
index 2f4620b7990c..9c8b1b37b510 100644
--- a/drivers/misc/habanalabs/common/context.c
+++ b/drivers/misc/habanalabs/common/context.c
@@ -9,38 +9,46 @@
 
 #include <linux/slab.h>
 
-void hl_encaps_handle_do_release(struct kref *ref)
+static void encaps_handle_do_release(struct hl_cs_encaps_sig_handle *handle, bool put_hw_sob,
+					bool put_ctx)
 {
-	struct hl_cs_encaps_sig_handle *handle =
-		container_of(ref, struct hl_cs_encaps_sig_handle, refcount);
 	struct hl_encaps_signals_mgr *mgr = &handle->ctx->sig_mgr;
 
+	if (put_hw_sob)
+		hw_sob_put(handle->hw_sob);
+
 	spin_lock(&mgr->lock);
 	idr_remove(&mgr->handles, handle->id);
 	spin_unlock(&mgr->lock);
 
-	hl_ctx_put(handle->ctx);
+	if (put_ctx)
+		hl_ctx_put(handle->ctx);
+
 	kfree(handle);
 }
 
-static void hl_encaps_handle_do_release_sob(struct kref *ref)
+void hl_encaps_release_handle_and_put_ctx(struct kref *ref)
 {
 	struct hl_cs_encaps_sig_handle *handle =
-		container_of(ref, struct hl_cs_encaps_sig_handle, refcount);
-	struct hl_encaps_signals_mgr *mgr = &handle->ctx->sig_mgr;
+			container_of(ref, struct hl_cs_encaps_sig_handle, refcount);
 
-	/* if we're here, then there was a signals reservation but cs with
-	 * encaps signals wasn't submitted, so need to put refcount
-	 * to hw_sob taken at the reservation.
-	 */
-	hw_sob_put(handle->hw_sob);
+	encaps_handle_do_release(handle, false, true);
+}
 
-	spin_lock(&mgr->lock);
-	idr_remove(&mgr->handles, handle->id);
-	spin_unlock(&mgr->lock);
+static void hl_encaps_release_handle_and_put_sob(struct kref *ref)
+{
+	struct hl_cs_encaps_sig_handle *handle =
+			container_of(ref, struct hl_cs_encaps_sig_handle, refcount);
 
-	hl_ctx_put(handle->ctx);
-	kfree(handle);
+	encaps_handle_do_release(handle, true, false);
+}
+
+void hl_encaps_release_handle_and_put_sob_ctx(struct kref *ref)
+{
+	struct hl_cs_encaps_sig_handle *handle =
+			container_of(ref, struct hl_cs_encaps_sig_handle, refcount);
+
+	encaps_handle_do_release(handle, true, true);
 }
 
 static void hl_encaps_sig_mgr_init(struct hl_encaps_signals_mgr *mgr)
@@ -49,8 +57,7 @@ static void hl_encaps_sig_mgr_init(struct hl_encaps_signals_mgr *mgr)
 	idr_init(&mgr->handles);
 }
 
-static void hl_encaps_sig_mgr_fini(struct hl_device *hdev,
-			struct hl_encaps_signals_mgr *mgr)
+static void hl_encaps_sig_mgr_fini(struct hl_device *hdev, struct hl_encaps_signals_mgr *mgr)
 {
 	struct hl_cs_encaps_sig_handle *handle;
 	struct idr *idp;
@@ -58,11 +65,14 @@ static void hl_encaps_sig_mgr_fini(struct hl_device *hdev,
 
 	idp = &mgr->handles;
 
+	/* The IDR is expected to be empty at this stage, because any left signal should have been
+	 * released as part of CS roll-back.
+	 */
 	if (!idr_is_empty(idp)) {
-		dev_warn(hdev->dev, "device released while some encaps signals handles are still allocated\n");
+		dev_warn(hdev->dev,
+			"device released while some encaps signals handles are still allocated\n");
 		idr_for_each_entry(idp, handle, id)
-			kref_put(&handle->refcount,
-					hl_encaps_handle_do_release_sob);
+			kref_put(&handle->refcount, hl_encaps_release_handle_and_put_sob);
 	}
 
 	idr_destroy(&mgr->handles);