1 files changed, 15 insertions, 12 deletions

diff --git a/mali_kbase/mali_kbase_ctx_sched.c b/mali_kbase/mali_kbase_ctx_sched.c
index 3e58500..beb2928 100644
--- a/mali_kbase/mali_kbase_ctx_sched.c
+++ b/mali_kbase/mali_kbase_ctx_sched.c
@@ -119,7 +119,7 @@ int kbase_ctx_sched_retain_ctx(struct kbase_context *kctx)
 	if (atomic_inc_return(&kctx->refcount) == 1) {
 		int const free_as = kbasep_ctx_sched_find_as_for_ctx(kctx);
 
-		if (free_as != KBASEP_AS_NR_INVALID) {
+		if (free_as >= 0) {
 			kbdev->as_free &= ~(1u << free_as);
 			/* Only program the MMU if the context has not been
 			 * assigned the same address space before.
@@ -173,8 +173,10 @@ void kbase_ctx_sched_retain_ctx_refcount(struct kbase_context *kctx)
          */
 	WARN_ON(!atomic_read(&kctx->refcount));
 #endif
-	WARN_ON(kctx->as_nr == KBASEP_AS_NR_INVALID);
-	WARN_ON(kbdev->as_to_kctx[kctx->as_nr] != kctx);
+	if (likely((kctx->as_nr >= 0) && (kctx->as_nr < BASE_MAX_NR_AS)))
+		WARN_ON(kbdev->as_to_kctx[kctx->as_nr] != kctx);
+	else
+		WARN(true, "Invalid as_nr(%d)", kctx->as_nr);
 
 	atomic_inc(&kctx->refcount);
 }
@@ -188,16 +190,17 @@ void kbase_ctx_sched_release_ctx(struct kbase_context *kctx)
 
 	new_ref_count = atomic_dec_return(&kctx->refcount);
 	if (new_ref_count == 0) {
-		kbdev->as_free |= (1u << kctx->as_nr);
-		if (kbase_ctx_flag(kctx, KCTX_AS_DISABLED_ON_FAULT)) {
-			KBASE_TLSTREAM_TL_KBASE_CTX_UNASSIGN_AS(
-				kbdev, kctx->id);
-			kbdev->as_to_kctx[kctx->as_nr] = NULL;
-			kctx->as_nr = KBASEP_AS_NR_INVALID;
-			kbase_ctx_flag_clear(kctx, KCTX_AS_DISABLED_ON_FAULT);
+		if (likely((kctx->as_nr >= 0) && (kctx->as_nr < BASE_MAX_NR_AS))) {
+			kbdev->as_free |= (1u << kctx->as_nr);
+			if (kbase_ctx_flag(kctx, KCTX_AS_DISABLED_ON_FAULT)) {
+				KBASE_TLSTREAM_TL_KBASE_CTX_UNASSIGN_AS(kbdev, kctx->id);
+				kbdev->as_to_kctx[kctx->as_nr] = NULL;
+				kctx->as_nr = KBASEP_AS_NR_INVALID;
+				kbase_ctx_flag_clear(kctx, KCTX_AS_DISABLED_ON_FAULT);
 #if !MALI_USE_CSF
-			kbase_backend_slot_kctx_purge_locked(kbdev, kctx);
+				kbase_backend_slot_kctx_purge_locked(kbdev, kctx);
 #endif
+			}
 		}
 	}
 
@@ -214,7 +217,7 @@ void kbase_ctx_sched_remove_ctx(struct kbase_context *kctx)
 
 	WARN_ON(atomic_read(&kctx->refcount) != 0);
 
-	if (kctx->as_nr != KBASEP_AS_NR_INVALID) {
+	if ((kctx->as_nr >= 0) && (kctx->as_nr < BASE_MAX_NR_AS)) {
 		if (kbdev->pm.backend.gpu_powered)
 			kbase_mmu_disable(kctx);