1 files changed, 19 insertions, 9 deletions

diff --git a/drivers/misc/habanalabs/command_buffer.c b/drivers/misc/habanalabs/command_buffer.c
index 53fddbd8e693..02d13f71b1df 100644
--- a/drivers/misc/habanalabs/command_buffer.c
+++ b/drivers/misc/habanalabs/command_buffer.c
@@ -105,10 +105,9 @@ int hl_cb_create(struct hl_device *hdev, struct hl_cb_mgr *mgr,
 		goto out_err;
 	}
 
-	if (cb_size > HL_MAX_CB_SIZE) {
-		dev_err(hdev->dev,
-			"CB size %d must be less then %d\n",
-			cb_size, HL_MAX_CB_SIZE);
+	if (cb_size > SZ_2M) {
+		dev_err(hdev->dev, "CB size %d must be less than %d\n",
+			cb_size, SZ_2M);
 		rc = -EINVAL;
 		goto out_err;
 	}
@@ -211,7 +210,7 @@ int hl_cb_ioctl(struct hl_fpriv *hpriv, void *data)
 {
 	union hl_cb_args *args = data;
 	struct hl_device *hdev = hpriv->hdev;
-	u64 handle;
+	u64 handle = 0;
 	int rc;
 
 	if (hl_device_disabled_or_in_reset(hdev)) {
@@ -223,15 +222,26 @@ int hl_cb_ioctl(struct hl_fpriv *hpriv, void *data)
 
 	switch (args->in.op) {
 	case HL_CB_OP_CREATE:
-		rc = hl_cb_create(hdev, &hpriv->cb_mgr, args->in.cb_size,
-					&handle, hpriv->ctx->asid);
+		if (args->in.cb_size > HL_MAX_CB_SIZE) {
+			dev_err(hdev->dev,
+				"User requested CB size %d must be less than %d\n",
+				args->in.cb_size, HL_MAX_CB_SIZE);
+			rc = -EINVAL;
+		} else {
+			rc = hl_cb_create(hdev, &hpriv->cb_mgr,
+						args->in.cb_size, &handle,
+						hpriv->ctx->asid);
+		}
+
 		memset(args, 0, sizeof(*args));
 		args->out.cb_handle = handle;
 		break;
+
 	case HL_CB_OP_DESTROY:
 		rc = hl_cb_destroy(hdev, &hpriv->cb_mgr,
 					args->in.cb_handle);
 		break;
+
 	default:
 		rc = -ENOTTY;
 		break;
@@ -278,7 +288,7 @@ int hl_cb_mmap(struct hl_fpriv *hpriv, struct vm_area_struct *vma)
 	cb = hl_cb_get(hdev, &hpriv->cb_mgr, handle);
 	if (!cb) {
 		dev_err(hdev->dev,
-			"CB mmap failed, no match to handle %d\n", handle);
+			"CB mmap failed, no match to handle 0x%x\n", handle);
 		return -EINVAL;
 	}
 
@@ -347,7 +357,7 @@ struct hl_cb *hl_cb_get(struct hl_device *hdev, struct hl_cb_mgr *mgr,
 	if (!cb) {
 		spin_unlock(&mgr->cb_lock);
 		dev_warn(hdev->dev,
-			"CB get failed, no match to handle %d\n", handle);
+			"CB get failed, no match to handle 0x%x\n", handle);
 		return NULL;
 	}