Project import generated by Copybara. GitOrigin-RevId: 98383fa31a2fd6c1b306cf78cb369c0be4694ece Change-Id: I7621b0443f2dacc15c44c856dbd2cf275a3ce020

diff --git a/bifrost/r47p0/kernel/drivers/gpu/arm/midgard/mali_kbase_ctx_sched.c b/bifrost/r47p0/kernel/drivers/gpu/arm/midgard/mali_kbase_ctx_sched.c
index 41f8c9c..c5fa78b 100644
--- a/bifrost/r47p0/kernel/drivers/gpu/arm/midgard/mali_kbase_ctx_sched.c
+++ b/bifrost/r47p0/kernel/drivers/gpu/arm/midgard/mali_kbase_ctx_sched.c

@@ -209,7 +209,17 @@
 
 	mutex_lock(&kbdev->mmu_hw_mutex);
 	spin_lock_irqsave(&kbdev->hwaccess_lock, flags);
+	kbase_ctx_sched_remove_ctx_nolock(kctx);
+	spin_unlock_irqrestore(&kbdev->hwaccess_lock, flags);
+	mutex_unlock(&kbdev->mmu_hw_mutex);
+}
 
+void kbase_ctx_sched_remove_ctx_nolock(struct kbase_context *kctx)
+{
+	struct kbase_device *const kbdev = kctx->kbdev;
+
+	lockdep_assert_held(&kbdev->mmu_hw_mutex);
+	lockdep_assert_held(&kbdev->hwaccess_lock);
 	WARN_ON(atomic_read(&kctx->refcount) != 0);
 
 	if ((kctx->as_nr >= 0) && (kctx->as_nr < BASE_MAX_NR_AS)) {
@@ -220,9 +230,6 @@
 		kbdev->as_to_kctx[kctx->as_nr] = NULL;
 		kctx->as_nr = KBASEP_AS_NR_INVALID;
 	}
-
-	spin_unlock_irqrestore(&kbdev->hwaccess_lock, flags);
-	mutex_unlock(&kbdev->mmu_hw_mutex);
 }
 
 void kbase_ctx_sched_restore_all_as(struct kbase_device *kbdev)

diff --git a/bifrost/r47p0/kernel/drivers/gpu/arm/midgard/mali_kbase_ctx_sched.h b/bifrost/r47p0/kernel/drivers/gpu/arm/midgard/mali_kbase_ctx_sched.h
index b66850f..d888888 100644
--- a/bifrost/r47p0/kernel/drivers/gpu/arm/midgard/mali_kbase_ctx_sched.h
+++ b/bifrost/r47p0/kernel/drivers/gpu/arm/midgard/mali_kbase_ctx_sched.h

@@ -131,6 +131,20 @@
 void kbase_ctx_sched_remove_ctx(struct kbase_context *kctx);
 
 /**
+ * kbase_ctx_sched_remove_ctx_nolock - Unassign previously assigned address space
+ * @kctx: The context to be removed
+ *
+ * The following lock must be held by the caller:
+ * kbase_device::mmu_hw_mutex
+ * kbase_device::hwaccess_lock
+ *
+ * This function should be called when a context is being destroyed. The
+ * context must no longer have any reference. If it has been assigned an
+ * address space before then the AS will be unprogrammed.
+ */
+void kbase_ctx_sched_remove_ctx_nolock(struct kbase_context *kctx);
+
+/**
  * kbase_ctx_sched_restore_all_as - Reprogram all address spaces
  * @kbdev: The device for which address spaces to be reprogrammed
  *

diff --git a/bifrost/r47p0/kernel/drivers/gpu/arm/midgard/mali_kbase_js.c b/bifrost/r47p0/kernel/drivers/gpu/arm/midgard/mali_kbase_js.c
index 9b1bd71..3887270 100644
--- a/bifrost/r47p0/kernel/drivers/gpu/arm/midgard/mali_kbase_js.c
+++ b/bifrost/r47p0/kernel/drivers/gpu/arm/midgard/mali_kbase_js.c

@@ -796,7 +796,55 @@
 		mutex_unlock(&kbdev->js_data.runpool_mutex);
 	}
 
-	kbase_ctx_sched_remove_ctx(kctx);
+	/* A work item to handle page_fault/bus_fault/gpu_fault could be
+	 * pending for the outgoing context which we no longer care about.
+	 * Ensure that the context won't be accessed anymore by the fault
+	 * workers.
+	 */
+	while (true) {
+		unsigned long flags;
+		int refcount;
+
+		mutex_lock(&kbdev->mmu_hw_mutex);
+		spin_lock_irqsave(&kbdev->hwaccess_lock, flags);
+		refcount = atomic_read(&kctx->refcount);
+		if ((refcount != 0) && !WARN_ON_ONCE(kctx->as_nr == KBASEP_AS_NR_INVALID)) {
+			struct kbase_as *as = &kctx->kbdev->as[kctx->as_nr];
+			int new_refcount;
+
+			dev_dbg(kbdev->dev,
+				"Waiting for pending fault worker to complete when terminating context (%d_%d)",
+				kctx->tgid, kctx->id);
+			spin_unlock_irqrestore(&kbdev->hwaccess_lock, flags);
+			mutex_unlock(&kbdev->mmu_hw_mutex);
+			flush_workqueue(as->pf_wq);
+			new_refcount = atomic_read(&kctx->refcount);
+			if (refcount != new_refcount) {
+				/* Fault workers executed and released some references, re-check */
+				continue;
+			} else {
+				/* Waiting for pending fault workers to execute was not effective,
+				 * we're going to forcefully de-assign the AS from this context
+				 * because nothing else should still be accessing the context at
+				 * this point.
+				 *
+				 * This should never happen and a WARN_ON() would be printed by
+				 * kbase_ctx_sched_remove_ctx() if the refcount is non-zero.
+				 */
+				dev_warn(
+					kbdev->dev,
+					"No fault workers executed, %d refs remain for terminating context (%d_%d)",
+					new_refcount, kctx->tgid, kctx->id);
+				kbase_ctx_sched_remove_ctx(kctx);
+			}
+		} else {
+			kbase_ctx_sched_remove_ctx_nolock(kctx);
+			spin_unlock_irqrestore(&kbdev->hwaccess_lock, flags);
+			mutex_unlock(&kbdev->mmu_hw_mutex);
+		}
+		break;
+	}
+
 #if IS_ENABLED(CONFIG_MALI_TRACE_POWER_GPU_WORK_PERIOD)
 	gpu_metrics_ctx_term(kctx);
 #endif