habanalabs: Avoid running restore chunks if no execute chunks

CS with no chunks for execute phase is invalid, so its
context_switch/restore phase should not be run.
Hence, move the check of the execute chunks number to the beginning of
hl_cs_ioctl().

Signed-off-by: Tomer Tayar <ttayar@habana.ai>
Reviewed-by: Oded Gabbay <oded.gabbay@gmail.com>
Signed-off-by: Oded Gabbay <oded.gabbay@gmail.com>

1 files changed, 21 insertions, 20 deletions

diff --git a/drivers/misc/habanalabs/command_submission.c b/drivers/misc/habanalabs/command_submission.c
index 7cb6910378bf..73ef0f9d758a 100644
--- a/drivers/misc/habanalabs/command_submission.c
+++ b/drivers/misc/habanalabs/command_submission.c
@@ -657,8 +657,8 @@ int hl_cs_ioctl(struct hl_fpriv *hpriv, void *data)
 	struct hl_device *hdev = hpriv->hdev;
 	union hl_cs_args *args = data;
 	struct hl_ctx *ctx = hpriv->ctx;
-	void __user *chunks;
-	u32 num_chunks;
+	void __user *chunks_execute, *chunks_restore;
+	u32 num_chunks_execute, num_chunks_restore;
 	u64 cs_seq = ULONG_MAX;
 	int rc, do_ctx_switch;
 	bool need_soft_reset = false;
@@ -671,13 +671,25 @@ int hl_cs_ioctl(struct hl_fpriv *hpriv, void *data)
 		goto out;
 	}
 
+	chunks_execute = (void __user *) (uintptr_t) args->in.chunks_execute;
+	num_chunks_execute = args->in.num_chunks_execute;
+
+	if (!num_chunks_execute) {
+		dev_err(hdev->dev,
+			"Got execute CS with 0 chunks, context %d\n",
+			ctx->asid);
+		rc = -EINVAL;
+		goto out;
+	}
+
 	do_ctx_switch = atomic_cmpxchg(&ctx->thread_ctx_switch_token, 1, 0);
 
 	if (do_ctx_switch || (args->in.cs_flags & HL_CS_FLAGS_FORCE_RESTORE)) {
 		long ret;
 
-		chunks = (void __user *)(uintptr_t)args->in.chunks_restore;
-		num_chunks = args->in.num_chunks_restore;
+		chunks_restore =
+			(void __user *) (uintptr_t) args->in.chunks_restore;
+		num_chunks_restore = args->in.num_chunks_restore;
 
 		mutex_lock(&hpriv->restore_phase_mutex);
 
@@ -705,13 +717,13 @@ int hl_cs_ioctl(struct hl_fpriv *hpriv, void *data)
 
 		hdev->asic_funcs->restore_phase_topology(hdev);
 
-		if (num_chunks == 0) {
+		if (!num_chunks_restore) {
 			dev_dbg(hdev->dev,
 			"Need to run restore phase but restore CS is empty\n");
 			rc = 0;
 		} else {
-			rc = _hl_cs_ioctl(hpriv, chunks, num_chunks,
-						&cs_seq);
+			rc = _hl_cs_ioctl(hpriv, chunks_restore,
+						num_chunks_restore, &cs_seq);
 		}
 
 		mutex_unlock(&hpriv->restore_phase_mutex);
@@ -724,7 +736,7 @@ int hl_cs_ioctl(struct hl_fpriv *hpriv, void *data)
 		}
 
 		/* Need to wait for restore completion before execution phase */
-		if (num_chunks > 0) {
+		if (num_chunks_restore) {
 			ret = _hl_cs_wait_ioctl(hdev, ctx,
 					jiffies_to_usecs(hdev->timeout_jiffies),
 					cs_seq);
@@ -752,18 +764,7 @@ int hl_cs_ioctl(struct hl_fpriv *hpriv, void *data)
 		}
 	}
 
-	chunks = (void __user *)(uintptr_t)args->in.chunks_execute;
-	num_chunks = args->in.num_chunks_execute;
-
-	if (num_chunks == 0) {
-		dev_err(hdev->dev,
-			"Got execute CS with 0 chunks, context %d\n",
-			ctx->asid);
-		rc = -EINVAL;
-		goto out;
-	}
-
-	rc = _hl_cs_ioctl(hpriv, chunks, num_chunks, &cs_seq);
+	rc = _hl_cs_ioctl(hpriv, chunks_execute, num_chunks_execute, &cs_seq);
 
 out:
 	if (rc != -EAGAIN) {