4 files changed, 18 insertions, 7 deletions

diff --git a/drivers/misc/habanalabs/command_submission.c b/drivers/misc/habanalabs/command_submission.c
index 3525236ed8d9..19c84214a7ea 100644
--- a/drivers/misc/habanalabs/command_submission.c
+++ b/drivers/misc/habanalabs/command_submission.c
@@ -179,6 +179,12 @@ static void cs_do_release(struct kref *ref)
 
         /* We also need to update CI for internal queues */
         if (cs->submitted) {
+                int cs_cnt = atomic_dec_return(&hdev->cs_active_cnt);
+
+                WARN_ONCE((cs_cnt < 0),
+                        "hl%d: error in CS active cnt %d\n",
+                        hdev->id, cs_cnt);
+
                 hl_int_hw_queue_update_ci(cs);
 
                 spin_lock(&hdev->hw_queues_mirror_lock);
diff --git a/drivers/misc/habanalabs/device.c b/drivers/misc/habanalabs/device.c
index 93d67983ddba..470d8005b50e 100644
--- a/drivers/misc/habanalabs/device.c
+++ b/drivers/misc/habanalabs/device.c
@@ -218,6 +218,7 @@ static int device_early_init(struct hl_device *hdev)
         spin_lock_init(&hdev->hw_queues_mirror_lock);
         atomic_set(&hdev->in_reset, 0);
         atomic_set(&hdev->fd_open_cnt, 0);
+        atomic_set(&hdev->cs_active_cnt, 0);
 
         return 0;
 
diff --git a/drivers/misc/habanalabs/habanalabs.h b/drivers/misc/habanalabs/habanalabs.h
index 806f0a5ee4d8..a8ee52c880cd 100644
--- a/drivers/misc/habanalabs/habanalabs.h
+++ b/drivers/misc/habanalabs/habanalabs.h
@@ -1056,13 +1056,15 @@ struct hl_device_reset_work {
  * @cb_pool_lock: protects the CB pool.
  * @user_ctx: current user context executing.
  * @dram_used_mem: current DRAM memory consumption.
- * @in_reset: is device in reset flow.
- * @curr_pll_profile: current PLL profile.
- * @fd_open_cnt: number of open user processes.
  * @timeout_jiffies: device CS timeout value.
  * @max_power: the max power of the device, as configured by the sysadmin. This
  *             value is saved so in case of hard-reset, KMD will restore this
  *             value and update the F/W after the re-initialization
+ * @in_reset: is device in reset flow.
+ * @curr_pll_profile: current PLL profile.
+ * @fd_open_cnt: number of open user processes.
+ * @cs_active_cnt: number of active command submissions on this device (active
+ *                 means already in H/W queues)
  * @major: habanalabs KMD major.
  * @high_pll: high PLL profile frequency.
  * @soft_reset_cnt: number of soft reset since KMD loading.
@@ -1128,11 +1130,12 @@ struct hl_device {
         struct hl_ctx                   *user_ctx;
 
         atomic64_t                      dram_used_mem;
+        u64                             timeout_jiffies;
+        u64                             max_power;
         atomic_t                        in_reset;
         atomic_t                        curr_pll_profile;
         atomic_t                        fd_open_cnt;
-        u64                             timeout_jiffies;
-        u64                             max_power;
+        atomic_t                        cs_active_cnt;
         u32                             major;
         u32                             high_pll;
         u32                             soft_reset_cnt;
diff --git a/drivers/misc/habanalabs/hw_queue.c b/drivers/misc/habanalabs/hw_queue.c
index 67bece26417c..ef3bb6951360 100644
--- a/drivers/misc/habanalabs/hw_queue.c
+++ b/drivers/misc/habanalabs/hw_queue.c
@@ -370,12 +370,13 @@ int hl_hw_queue_schedule_cs(struct hl_cs *cs)
                 spin_unlock(&hdev->hw_queues_mirror_lock);
         }
 
-        list_for_each_entry_safe(job, tmp, &cs->job_list, cs_node) {
+        atomic_inc(&hdev->cs_active_cnt);
+
+        list_for_each_entry_safe(job, tmp, &cs->job_list, cs_node)
                 if (job->ext_queue)
                         ext_hw_queue_schedule_job(job);
                 else
                         int_hw_queue_schedule_job(job);
-        }
 
         cs->submitted = true;