Update KMD to r47p0

Provenance: ipdelivery@ad01e50d640910a99224382bb227e6d4de627657
Change-Id: I19ac9bce34a5c5a319c1b4a388e8b037b3dfe6e7

1 files changed, 478 insertions, 458 deletions

diff --git a/mali_kbase/mmu/mali_kbase_mmu.c b/mali_kbase/mmu/mali_kbase_mmu.c
index 9dd5116..cc4de07 100644
--- a/mali_kbase/mmu/mali_kbase_mmu.c
+++ b/mali_kbase/mmu/mali_kbase_mmu.c
@@ -28,7 +28,7 @@
 #include <linux/migrate.h>
 #include <mali_kbase.h>
 #include <gpu/mali_kbase_gpu_fault.h>
-#include <gpu/mali_kbase_gpu_regmap.h>
+#include <hw_access/mali_kbase_hw_access_regmap.h>
 #include <tl/mali_kbase_tracepoints.h>
 #include <backend/gpu/mali_kbase_instr_defs.h>
 #include <mali_kbase_ctx_sched.h>
@@ -56,7 +56,7 @@
 #define MGM_DEFAULT_PTE_GROUP (0)
 
 /* Macro to convert updated PDGs to flags indicating levels skip in flush */
-#define pgd_level_to_skip_flush(dirty_pgds) (~(dirty_pgds) & 0xF)
+#define pgd_level_to_skip_flush(dirty_pgds) (~(dirty_pgds)&0xF)
 
 static int mmu_insert_pages_no_flush(struct kbase_device *kbdev, struct kbase_mmu_table *mmut,
 				     const u64 start_vpfn, struct tagged_addr *phys, size_t nr,
@@ -64,8 +64,7 @@ static int mmu_insert_pages_no_flush(struct kbase_device *kbdev, struct kbase_mm
 				     struct kbase_va_region *reg, bool ignore_page_migration);
 
 /* Small wrapper function to factor out GPU-dependent context releasing */
-static void release_ctx(struct kbase_device *kbdev,
-		struct kbase_context *kctx)
+static void release_ctx(struct kbase_device *kbdev, struct kbase_context *kctx)
 {
 #if MALI_USE_CSF
 	CSTD_UNUSED(kbdev);
@@ -75,46 +74,6 @@ static void release_ctx(struct kbase_device *kbdev,
 #endif /* MALI_USE_CSF */
 }
 
-static void mmu_hw_operation_begin(struct kbase_device *kbdev)
-{
-#if !IS_ENABLED(CONFIG_MALI_NO_MALI)
-#if MALI_USE_CSF
-	if (kbase_hw_has_issue(kbdev, BASE_HW_ISSUE_GPU2019_3878)) {
-		unsigned long flags;
-
-		lockdep_assert_held(&kbdev->mmu_hw_mutex);
-
-		spin_lock_irqsave(&kbdev->hwaccess_lock, flags);
-		WARN_ON_ONCE(kbdev->mmu_hw_operation_in_progress);
-		kbdev->mmu_hw_operation_in_progress = true;
-		spin_unlock_irqrestore(&kbdev->hwaccess_lock, flags);
-	}
-#endif /* MALI_USE_CSF */
-#endif /* !CONFIG_MALI_NO_MALI */
-}
-
-static void mmu_hw_operation_end(struct kbase_device *kbdev)
-{
-#if !IS_ENABLED(CONFIG_MALI_NO_MALI)
-#if MALI_USE_CSF
-	if (kbase_hw_has_issue(kbdev, BASE_HW_ISSUE_GPU2019_3878)) {
-		unsigned long flags;
-
-		lockdep_assert_held(&kbdev->mmu_hw_mutex);
-
-		spin_lock_irqsave(&kbdev->hwaccess_lock, flags);
-		WARN_ON_ONCE(!kbdev->mmu_hw_operation_in_progress);
-		kbdev->mmu_hw_operation_in_progress = false;
-		/* Invoke the PM state machine, the L2 power off may have been
-		 * skipped due to the MMU command.
-		 */
-		kbase_pm_update_state(kbdev);
-		spin_unlock_irqrestore(&kbdev->hwaccess_lock, flags);
-	}
-#endif /* MALI_USE_CSF */
-#endif /* !CONFIG_MALI_NO_MALI */
-}
-
 /**
  * mmu_flush_cache_on_gpu_ctrl() - Check if cache flush needs to be done
  * through GPU_CONTROL interface.
@@ -128,11 +87,7 @@ static void mmu_hw_operation_end(struct kbase_device *kbdev)
  */
 static bool mmu_flush_cache_on_gpu_ctrl(struct kbase_device *kbdev)
 {
-	uint32_t const arch_maj_cur = (kbdev->gpu_props.props.raw_props.gpu_id &
-				       GPU_ID2_ARCH_MAJOR) >>
-				      GPU_ID2_ARCH_MAJOR_SHIFT;
-
-	return arch_maj_cur > 11;
+	return kbdev->gpu_props.gpu_id.arch_major > 11;
 }
 
 /**
@@ -207,7 +162,7 @@ static void mmu_flush_invalidate_as(struct kbase_device *kbdev, struct kbase_as
 	mutex_lock(&kbdev->mmu_hw_mutex);
 	spin_lock_irqsave(&kbdev->hwaccess_lock, flags);
 
-	if (kbdev->pm.backend.gpu_ready && (kbase_mmu_hw_do_flush_locked(kbdev, as, op_param)))
+	if (kbdev->pm.backend.gpu_ready && kbase_mmu_hw_do_flush(kbdev, as, op_param))
 		dev_err(kbdev->dev, "Flush for GPU page table update did not complete");
 
 	spin_unlock_irqrestore(&kbdev->hwaccess_lock, flags);
@@ -281,7 +236,8 @@ static void mmu_flush_invalidate(struct kbase_device *kbdev, struct kbase_contex
  * interface.
  */
 static void mmu_flush_invalidate_on_gpu_ctrl(struct kbase_device *kbdev, struct kbase_context *kctx,
-					int as_nr, const struct kbase_mmu_hw_op_param *op_param)
+					     int as_nr,
+					     const struct kbase_mmu_hw_op_param *op_param)
 {
 	unsigned long flags;
 
@@ -300,8 +256,7 @@ static void mmu_flush_invalidate_on_gpu_ctrl(struct kbase_device *kbdev, struct
 }
 
 static void kbase_mmu_sync_pgd_gpu(struct kbase_device *kbdev, struct kbase_context *kctx,
-				   phys_addr_t phys, size_t size,
-				   enum kbase_mmu_op_type flush_op)
+				   phys_addr_t phys, size_t size, enum kbase_mmu_op_type flush_op)
 {
 	kbase_mmu_flush_pa_range(kbdev, kctx, phys, size, flush_op);
 }
@@ -312,8 +267,7 @@ static void kbase_mmu_sync_pgd_cpu(struct kbase_device *kbdev, dma_addr_t handle
 	 * the pages from memory
 	 */
 	if (kbdev->system_coherency == COHERENCY_NONE)
-		dma_sync_single_for_device(kbdev->dev, handle, size,
-				DMA_TO_DEVICE);
+		dma_sync_single_for_device(kbdev->dev, handle, size, DMA_TO_DEVICE);
 }
 
 /**
@@ -339,7 +293,6 @@ static void kbase_mmu_sync_pgd(struct kbase_device *kbdev, struct kbase_context
 			       phys_addr_t phys, dma_addr_t handle, size_t size,
 			       enum kbase_mmu_op_type flush_op)
 {
-
 	kbase_mmu_sync_pgd_cpu(kbdev, handle, size);
 	kbase_mmu_sync_pgd_gpu(kbdev, kctx, phys, size, flush_op);
 }
@@ -390,8 +343,7 @@ static void kbase_mmu_account_freed_pgd(struct kbase_device *kbdev, struct kbase
 }
 
 static bool kbase_mmu_handle_isolated_pgd_page(struct kbase_device *kbdev,
-					       struct kbase_mmu_table *mmut,
-					       struct page *p)
+					       struct kbase_mmu_table *mmut, struct page *p)
 {
 	struct kbase_page_metadata *page_md = kbase_page_private(p);
 	bool page_is_isolated = false;
@@ -405,13 +357,12 @@ static bool kbase_mmu_handle_isolated_pgd_page(struct kbase_device *kbdev,
 	if (PAGE_STATUS_GET(page_md->status) == PT_MAPPED) {
 		WARN_ON_ONCE(!mmut->kctx);
 		if (IS_PAGE_ISOLATED(page_md->status)) {
-			page_md->status = PAGE_STATUS_SET(page_md->status,
-							  FREE_PT_ISOLATED_IN_PROGRESS);
+			page_md->status =
+				PAGE_STATUS_SET(page_md->status, FREE_PT_ISOLATED_IN_PROGRESS);
 			page_md->data.free_pt_isolated.kbdev = kbdev;
 			page_is_isolated = true;
 		} else {
-			page_md->status =
-				PAGE_STATUS_SET(page_md->status, FREE_IN_PROGRESS);
+			page_md->status = PAGE_STATUS_SET(page_md->status, FREE_IN_PROGRESS);
 		}
 	} else if ((PAGE_STATUS_GET(page_md->status) == FREE_IN_PROGRESS) ||
 		   (PAGE_STATUS_GET(page_md->status) == ALLOCATE_IN_PROGRESS)) {
@@ -520,8 +471,8 @@ static inline void kbase_mmu_reset_free_pgds_list(struct kbase_mmu_table *mmut)
  *
  * Return: the number of backed pages to increase by
  */
-static size_t reg_grow_calc_extra_pages(struct kbase_device *kbdev,
-		struct kbase_va_region *reg, size_t fault_rel_pfn)
+static size_t reg_grow_calc_extra_pages(struct kbase_device *kbdev, struct kbase_va_region *reg,
+					size_t fault_rel_pfn)
 {
 	size_t multiple = reg->extension;
 	size_t reg_current_size = kbase_reg_current_backed_size(reg);
@@ -569,8 +520,8 @@ static size_t reg_grow_calc_extra_pages(struct kbase_device *kbdev,
 			/* same as calculating
 			 * (fault_rel_pfn - initial_commit + 1)
 			 */
-			size_t pages_after_initial = minimum_extra +
-				reg_current_size - initial_commit;
+			size_t pages_after_initial =
+				minimum_extra + reg_current_size - initial_commit;
 
 			remainder = pages_after_initial % multiple;
 		}
@@ -585,21 +536,18 @@ static size_t reg_grow_calc_extra_pages(struct kbase_device *kbdev,
 
 #ifdef CONFIG_MALI_CINSTR_GWT
 static void kbase_gpu_mmu_handle_write_faulting_as(struct kbase_device *kbdev,
-						   struct kbase_as *faulting_as,
-						   u64 start_pfn, size_t nr,
-						   u32 kctx_id, u64 dirty_pgds)
+						   struct kbase_as *faulting_as, u64 start_pfn,
+						   size_t nr, u32 kctx_id, u64 dirty_pgds)
 {
 	/* Calls to this function are inherently synchronous, with respect to
 	 * MMU operations.
 	 */
 	const enum kbase_caller_mmu_sync_info mmu_sync_info = CALLER_MMU_SYNC;
 	struct kbase_mmu_hw_op_param op_param;
+	unsigned long irq_flags;
 	int ret = 0;
 
-	mutex_lock(&kbdev->mmu_hw_mutex);
-
-	kbase_mmu_hw_clear_fault(kbdev, faulting_as,
-			KBASE_MMU_FAULT_TYPE_PAGE);
+	kbase_mmu_hw_clear_fault(kbdev, faulting_as, KBASE_MMU_FAULT_TYPE_PAGE);
 
 	/* flush L2 and unlock the VA (resumes the MMU) */
 	op_param.vpfn = start_pfn;
@@ -607,43 +555,34 @@ static void kbase_gpu_mmu_handle_write_faulting_as(struct kbase_device *kbdev,
 	op_param.op = KBASE_MMU_OP_FLUSH_PT;
 	op_param.kctx_id = kctx_id;
 	op_param.mmu_sync_info = mmu_sync_info;
+	spin_lock_irqsave(&kbdev->hwaccess_lock, irq_flags);
 	if (mmu_flush_cache_on_gpu_ctrl(kbdev)) {
-		unsigned long irq_flags;
-
-		spin_lock_irqsave(&kbdev->hwaccess_lock, irq_flags);
-		op_param.flush_skip_levels =
-				pgd_level_to_skip_flush(dirty_pgds);
+		op_param.flush_skip_levels = pgd_level_to_skip_flush(dirty_pgds);
 		ret = kbase_mmu_hw_do_flush_on_gpu_ctrl(kbdev, faulting_as, &op_param);
-		spin_unlock_irqrestore(&kbdev->hwaccess_lock, irq_flags);
 	} else {
-		mmu_hw_operation_begin(kbdev);
 		ret = kbase_mmu_hw_do_flush(kbdev, faulting_as, &op_param);
-		mmu_hw_operation_end(kbdev);
 	}
-
-	mutex_unlock(&kbdev->mmu_hw_mutex);
+	spin_unlock_irqrestore(&kbdev->hwaccess_lock, irq_flags);
 
 	if (ret)
 		dev_err(kbdev->dev,
 			"Flush for GPU page fault due to write access did not complete");
 
-	kbase_mmu_hw_enable_fault(kbdev, faulting_as,
-			KBASE_MMU_FAULT_TYPE_PAGE);
+	kbase_mmu_hw_enable_fault(kbdev, faulting_as, KBASE_MMU_FAULT_TYPE_PAGE);
 }
 
-static void set_gwt_element_page_addr_and_size(
-		struct kbasep_gwt_list_element *element,
-		u64 fault_page_addr, struct tagged_addr fault_phys)
+static void set_gwt_element_page_addr_and_size(struct kbasep_gwt_list_element *element,
+					       u64 fault_page_addr, struct tagged_addr fault_phys)
 {
 	u64 fault_pfn = fault_page_addr >> PAGE_SHIFT;
-	unsigned int vindex = fault_pfn & (NUM_4K_PAGES_IN_2MB_PAGE - 1);
+	unsigned int vindex = fault_pfn & (NUM_PAGES_IN_2MB_LARGE_PAGE - 1);
 
 	/* If the fault address lies within a 2MB page, then consider
 	 * the whole 2MB page for dumping to avoid incomplete dumps.
 	 */
 	if (is_huge(fault_phys) && (vindex == index_in_large_page(fault_phys))) {
-		element->page_addr = fault_page_addr & ~(SZ_2M - 1);
-		element->num_pages = NUM_4K_PAGES_IN_2MB_PAGE;
+		element->page_addr = fault_page_addr & ~(SZ_2M - 1UL);
+		element->num_pages = NUM_PAGES_IN_2MB_LARGE_PAGE;
 	} else {
 		element->page_addr = fault_page_addr;
 		element->num_pages = 1;
@@ -651,7 +590,7 @@ static void set_gwt_element_page_addr_and_size(
 }
 
 static void kbase_gpu_mmu_handle_write_fault(struct kbase_context *kctx,
-			struct kbase_as *faulting_as)
+					     struct kbase_as *faulting_as)
 {
 	struct kbasep_gwt_list_element *pos;
 	struct kbase_va_region *region;
@@ -659,7 +598,7 @@ static void kbase_gpu_mmu_handle_write_fault(struct kbase_context *kctx,
 	struct tagged_addr *fault_phys_addr;
 	struct kbase_fault *fault;
 	u64 fault_pfn, pfn_offset;
-	int as_no;
+	unsigned int as_no;
 	u64 dirty_pgds = 0;
 
 	as_no = faulting_as->number;
@@ -670,21 +609,20 @@ static void kbase_gpu_mmu_handle_write_fault(struct kbase_context *kctx,
 	kbase_gpu_vm_lock(kctx);
 
 	/* Find region and check if it should be writable. */
-	region = kbase_region_tracker_find_region_enclosing_address(kctx,
-			fault->addr);
+	region = kbase_region_tracker_find_region_enclosing_address(kctx, fault->addr);
 	if (kbase_is_region_invalid_or_free(region)) {
 		kbase_gpu_vm_unlock(kctx);
 		kbase_mmu_report_fault_and_kill(kctx, faulting_as,
-				"Memory is not mapped on the GPU",
-				&faulting_as->pf_data);
+						"Memory is not mapped on the GPU",
+						&faulting_as->pf_data);
 		return;
 	}
 
 	if (!(region->flags & KBASE_REG_GPU_WR)) {
 		kbase_gpu_vm_unlock(kctx);
 		kbase_mmu_report_fault_and_kill(kctx, faulting_as,
-				"Region does not have write permissions",
-				&faulting_as->pf_data);
+						"Region does not have write permissions",
+						&faulting_as->pf_data);
 		return;
 	}
 
@@ -709,8 +647,8 @@ static void kbase_gpu_mmu_handle_write_fault(struct kbase_context *kctx,
 			pos = kmalloc(sizeof(*pos), GFP_KERNEL);
 			if (pos) {
 				pos->region = region;
-				set_gwt_element_page_addr_and_size(pos,
-					fault_page_addr, *fault_phys_addr);
+				set_gwt_element_page_addr_and_size(pos, fault_page_addr,
+								   *fault_phys_addr);
 				list_add(&pos->link, &kctx->gwt_current_list);
 			} else {
 				dev_warn(kbdev->dev, "kmalloc failure");
@@ -722,14 +660,14 @@ static void kbase_gpu_mmu_handle_write_fault(struct kbase_context *kctx,
 	kbase_mmu_update_pages_no_flush(kbdev, &kctx->mmu, fault_pfn, fault_phys_addr, 1,
 					region->flags, region->gpu_alloc->group_id, &dirty_pgds);
 
-	kbase_gpu_mmu_handle_write_faulting_as(kbdev, faulting_as, fault_pfn, 1,
-					       kctx->id, dirty_pgds);
+	kbase_gpu_mmu_handle_write_faulting_as(kbdev, faulting_as, fault_pfn, 1, kctx->id,
+					       dirty_pgds);
 
 	kbase_gpu_vm_unlock(kctx);
 }
 
 static void kbase_gpu_mmu_handle_permission_fault(struct kbase_context *kctx,
-			struct kbase_as	*faulting_as)
+						  struct kbase_as *faulting_as)
 {
 	struct kbase_fault *fault = &faulting_as->pf_data;
 
@@ -739,16 +677,15 @@ static void kbase_gpu_mmu_handle_permission_fault(struct kbase_context *kctx,
 		kbase_gpu_mmu_handle_write_fault(kctx, faulting_as);
 		break;
 	case AS_FAULTSTATUS_ACCESS_TYPE_EXECUTE:
-		kbase_mmu_report_fault_and_kill(kctx, faulting_as,
-				"Execute Permission fault", fault);
+		kbase_mmu_report_fault_and_kill(kctx, faulting_as, "Execute Permission fault",
+						fault);
 		break;
 	case AS_FAULTSTATUS_ACCESS_TYPE_READ:
-		kbase_mmu_report_fault_and_kill(kctx, faulting_as,
-				"Read Permission fault", fault);
+		kbase_mmu_report_fault_and_kill(kctx, faulting_as, "Read Permission fault", fault);
 		break;
 	default:
-		kbase_mmu_report_fault_and_kill(kctx, faulting_as,
-				"Unknown Permission fault", fault);
+		kbase_mmu_report_fault_and_kill(kctx, faulting_as, "Unknown Permission fault",
+						fault);
 		break;
 	}
 }
@@ -758,7 +695,7 @@ static void kbase_gpu_mmu_handle_permission_fault(struct kbase_context *kctx,
  * estimate_pool_space_required - Determine how much a pool should be grown by to support a future
  * allocation
  * @pool:           The memory pool to check, including its linked pools
- * @pages_required: Number of 4KiB pages require for the pool to support a future allocation
+ * @pages_required: Number of small pages require for the pool to support a future allocation
  *
  * The value returned is accounting for the size of @pool and the size of each memory pool linked to
  * @pool. Hence, the caller should use @pool and (if not already satisfied) all its linked pools to
@@ -770,7 +707,7 @@ static void kbase_gpu_mmu_handle_permission_fault(struct kbase_context *kctx,
  * should keep attempting an allocation and then re-growing with a new value queried form this
  * function until the allocation succeeds.
  *
- * Return: an estimate of the amount of extra 4KiB pages in @pool that are required to satisfy an
+ * Return: an estimate of the amount of extra small pages in @pool that are required to satisfy an
  * allocation, or 0 if @pool (including its linked pools) is likely to already satisfy the
  * allocation.
  */
@@ -780,15 +717,15 @@ static size_t estimate_pool_space_required(struct kbase_mem_pool *pool, const si
 
 	for (pages_still_required = pages_required; pool != NULL && pages_still_required;
 	     pool = pool->next_pool) {
-		size_t pool_size_4k;
+		size_t pool_size_small;
 
 		kbase_mem_pool_lock(pool);
 
-		pool_size_4k = kbase_mem_pool_size(pool) << pool->order;
-		if (pool_size_4k >= pages_still_required)
+		pool_size_small = kbase_mem_pool_size(pool) << pool->order;
+		if (pool_size_small >= pages_still_required)
 			pages_still_required = 0;
 		else
-			pages_still_required -= pool_size_4k;
+			pages_still_required -= pool_size_small;
 
 		kbase_mem_pool_unlock(pool);
 	}
@@ -799,11 +736,11 @@ static size_t estimate_pool_space_required(struct kbase_mem_pool *pool, const si
  * page_fault_try_alloc - Try to allocate memory from a context pool
  * @kctx:          Context pointer
  * @region:        Region to grow
- * @new_pages:     Number of 4 KiB pages to allocate
+ * @new_pages:     Number of small pages to allocate
  * @pages_to_grow: Pointer to variable to store number of outstanding pages on failure. This can be
- *                 either 4 KiB or 2 MiB pages, depending on the number of pages requested.
+ *                 either small or 2 MiB pages, depending on the number of pages requested.
  * @grow_2mb_pool: Pointer to variable to store which pool needs to grow - true for 2 MiB, false for
- *                 4 KiB.
+ *                 pool of small pages.
  * @prealloc_sas:  Pointer to kbase_sub_alloc structures
  *
  * This function will try to allocate as many pages as possible from the context pool, then if
@@ -815,33 +752,31 @@ static size_t estimate_pool_space_required(struct kbase_mem_pool *pool, const si
  * held could invoke the OoM killer and cause an effective deadlock with kbase_cpu_vm_close().
  *
  * If 2 MiB pages are enabled and new_pages is >= 2 MiB then pages_to_grow will be a count of 2 MiB
- * pages, otherwise it will be a count of 4 KiB pages.
+ * pages, otherwise it will be a count of small pages.
  *
  * Return: true if successful, false on failure
  */
-static bool page_fault_try_alloc(struct kbase_context *kctx,
-		struct kbase_va_region *region, size_t new_pages,
-		int *pages_to_grow, bool *grow_2mb_pool,
-		struct kbase_sub_alloc **prealloc_sas)
+static bool page_fault_try_alloc(struct kbase_context *kctx, struct kbase_va_region *region,
+				 size_t new_pages, size_t *pages_to_grow, bool *grow_2mb_pool,
+				 struct kbase_sub_alloc **prealloc_sas)
 {
 	size_t total_gpu_pages_alloced = 0;
 	size_t total_cpu_pages_alloced = 0;
 	struct kbase_mem_pool *pool, *root_pool;
 	bool alloc_failed = false;
 	size_t pages_still_required;
-	size_t total_mempools_free_4k = 0;
+	size_t total_mempools_free_small = 0;
 
 	lockdep_assert_held(&kctx->reg_lock);
 	lockdep_assert_held(&kctx->mem_partials_lock);
 
-	if (WARN_ON(region->gpu_alloc->group_id >=
-		MEMORY_GROUP_MANAGER_NR_GROUPS)) {
+	if (WARN_ON(region->gpu_alloc->group_id >= MEMORY_GROUP_MANAGER_NR_GROUPS)) {
 		/* Do not try to grow the memory pool */
 		*pages_to_grow = 0;
 		return false;
 	}
 
-	if (kctx->kbdev->pagesize_2mb && new_pages >= (SZ_2M / SZ_4K)) {
+	if (kctx->kbdev->pagesize_2mb && new_pages >= NUM_PAGES_IN_2MB_LARGE_PAGE) {
 		root_pool = &kctx->mem_pools.large[region->gpu_alloc->group_id];
 		*grow_2mb_pool = true;
 	} else {
@@ -871,41 +806,40 @@ static bool page_fault_try_alloc(struct kbase_context *kctx,
 	 */
 	pages_still_required = new_pages;
 	for (pool = root_pool; pool != NULL && pages_still_required; pool = pool->next_pool) {
-		size_t pool_size_4k;
-		size_t pages_to_alloc_4k;
-		size_t pages_to_alloc_4k_per_alloc;
+		size_t pool_size_small;
+		size_t pages_to_alloc_small;
+		size_t pages_to_alloc_small_per_alloc;
 
 		kbase_mem_pool_lock(pool);
 
 		/* Allocate as much as possible from this pool*/
-		pool_size_4k = kbase_mem_pool_size(pool) << pool->order;
-		total_mempools_free_4k += pool_size_4k;
-		pages_to_alloc_4k = MIN(pages_still_required, pool_size_4k);
+		pool_size_small = kbase_mem_pool_size(pool) << pool->order;
+		total_mempools_free_small += pool_size_small;
+		pages_to_alloc_small = MIN(pages_still_required, pool_size_small);
 		if (region->gpu_alloc == region->cpu_alloc)
-			pages_to_alloc_4k_per_alloc = pages_to_alloc_4k;
+			pages_to_alloc_small_per_alloc = pages_to_alloc_small;
 		else
-			pages_to_alloc_4k_per_alloc = pages_to_alloc_4k >> 1;
+			pages_to_alloc_small_per_alloc = pages_to_alloc_small >> 1;
 
-		if (pages_to_alloc_4k) {
-			struct tagged_addr *gpu_pages =
-				kbase_alloc_phy_pages_helper_locked(region->gpu_alloc, pool,
-								    pages_to_alloc_4k_per_alloc,
-								    &prealloc_sas[0]);
+		if (pages_to_alloc_small) {
+			struct tagged_addr *gpu_pages = kbase_alloc_phy_pages_helper_locked(
+				region->gpu_alloc, pool, pages_to_alloc_small_per_alloc,
+				&prealloc_sas[0]);
 
 			if (!gpu_pages)
 				alloc_failed = true;
 			else
-				total_gpu_pages_alloced += pages_to_alloc_4k_per_alloc;
+				total_gpu_pages_alloced += pages_to_alloc_small_per_alloc;
 
 			if (!alloc_failed && region->gpu_alloc != region->cpu_alloc) {
 				struct tagged_addr *cpu_pages = kbase_alloc_phy_pages_helper_locked(
-					region->cpu_alloc, pool, pages_to_alloc_4k_per_alloc,
+					region->cpu_alloc, pool, pages_to_alloc_small_per_alloc,
 					&prealloc_sas[1]);
 
 				if (!cpu_pages)
 					alloc_failed = true;
 				else
-					total_cpu_pages_alloced += pages_to_alloc_4k_per_alloc;
+					total_cpu_pages_alloced += pages_to_alloc_small_per_alloc;
 			}
 		}
 
@@ -913,12 +847,12 @@ static bool page_fault_try_alloc(struct kbase_context *kctx,
 
 		if (alloc_failed) {
 			WARN_ON(!pages_still_required);
-			WARN_ON(pages_to_alloc_4k >= pages_still_required);
-			WARN_ON(pages_to_alloc_4k_per_alloc >= pages_still_required);
+			WARN_ON(pages_to_alloc_small >= pages_still_required);
+			WARN_ON(pages_to_alloc_small_per_alloc >= pages_still_required);
 			break;
 		}
 
-		pages_still_required -= pages_to_alloc_4k;
+		pages_still_required -= pages_to_alloc_small;
 	}
 
 	if (pages_still_required) {
@@ -942,7 +876,7 @@ static bool page_fault_try_alloc(struct kbase_context *kctx,
 				kctx->kbdev->dev,
 				"Page allocation failure of %zu pages: managed %zu pages, mempool (inc linked pools) had %zu pages available",
 				new_pages, total_gpu_pages_alloced + total_cpu_pages_alloced,
-				total_mempools_free_4k);
+				total_mempools_free_small);
 			*pages_to_grow = 0;
 		} else {
 			/* Tell the caller to try to grow the memory pool
@@ -981,14 +915,14 @@ void kbase_mmu_page_fault_worker(struct work_struct *data)
 	size_t new_pages;
 	size_t fault_rel_pfn;
 	struct kbase_as *faulting_as;
-	int as_no;
+	unsigned int as_no;
 	struct kbase_context *kctx;
 	struct kbase_device *kbdev;
 	struct kbase_va_region *region;
 	struct kbase_fault *fault;
 	int err;
 	bool grown = false;
-	int pages_to_grow;
+	size_t pages_to_grow;
 	bool grow_2mb_pool;
 	struct kbase_sub_alloc *prealloc_sas[2] = { NULL, NULL };
 	int i;
@@ -996,6 +930,7 @@ void kbase_mmu_page_fault_worker(struct work_struct *data)
 #if MALI_JIT_PRESSURE_LIMIT_BASE
 	size_t pages_trimmed = 0;
 #endif
+	unsigned long hwaccess_flags;
 
 	/* Calls to this function are inherently synchronous, with respect to
 	 * MMU operations.
@@ -1008,7 +943,7 @@ void kbase_mmu_page_fault_worker(struct work_struct *data)
 	as_no = faulting_as->number;
 
 	kbdev = container_of(faulting_as, struct kbase_device, as[as_no]);
-	dev_dbg(kbdev->dev, "Entering %s %pK, fault_pfn %lld, as_no %d", __func__, (void *)data,
+	dev_dbg(kbdev->dev, "Entering %s %pK, fault_pfn %lld, as_no %u", __func__, (void *)data,
 		fault_pfn, as_no);
 
 	/* Grab the context that was already refcounted in kbase_mmu_interrupt()
@@ -1038,60 +973,122 @@ void kbase_mmu_page_fault_worker(struct work_struct *data)
 #endif
 
 	if (unlikely(fault->protected_mode)) {
-		kbase_mmu_report_fault_and_kill(kctx, faulting_as,
-				"Protected mode fault", fault);
-		kbase_mmu_hw_clear_fault(kbdev, faulting_as,
-				KBASE_MMU_FAULT_TYPE_PAGE);
+		kbase_mmu_report_fault_and_kill(kctx, faulting_as, "Protected mode fault", fault);
+		kbase_mmu_hw_clear_fault(kbdev, faulting_as, KBASE_MMU_FAULT_TYPE_PAGE);
 
 		goto fault_done;
 	}
 
 	fault_status = fault->status;
-	switch (fault_status & AS_FAULTSTATUS_EXCEPTION_CODE_MASK) {
-
-	case AS_FAULTSTATUS_EXCEPTION_CODE_TRANSLATION_FAULT:
+	switch (AS_FAULTSTATUS_EXCEPTION_TYPE_GET(fault_status)) {
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_TRANSLATION_FAULT_0:
+		fallthrough;
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_TRANSLATION_FAULT_1:
+		fallthrough;
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_TRANSLATION_FAULT_2:
+		fallthrough;
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_TRANSLATION_FAULT_3:
+		fallthrough;
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_TRANSLATION_FAULT_4:
+#if !MALI_USE_CSF
+		fallthrough;
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_TRANSLATION_FAULT_IDENTITY:
+#endif
 		/* need to check against the region to handle this one */
 		break;
 
-	case AS_FAULTSTATUS_EXCEPTION_CODE_PERMISSION_FAULT:
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_PERMISSION_FAULT_0:
+		fallthrough;
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_PERMISSION_FAULT_1:
+		fallthrough;
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_PERMISSION_FAULT_2:
+		fallthrough;
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_PERMISSION_FAULT_3:
 #ifdef CONFIG_MALI_CINSTR_GWT
 		/* If GWT was ever enabled then we need to handle
 		 * write fault pages even if the feature was disabled later.
 		 */
 		if (kctx->gwt_was_enabled) {
-			kbase_gpu_mmu_handle_permission_fault(kctx,
-							faulting_as);
+			kbase_gpu_mmu_handle_permission_fault(kctx, faulting_as);
 			goto fault_done;
 		}
 #endif
 
-		kbase_mmu_report_fault_and_kill(kctx, faulting_as,
-				"Permission failure", fault);
+		kbase_mmu_report_fault_and_kill(kctx, faulting_as, "Permission failure", fault);
 		goto fault_done;
 
-	case AS_FAULTSTATUS_EXCEPTION_CODE_TRANSTAB_BUS_FAULT:
-		kbase_mmu_report_fault_and_kill(kctx, faulting_as,
-				"Translation table bus fault", fault);
+#if !MALI_USE_CSF
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_TRANSTAB_BUS_FAULT_0:
+		fallthrough;
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_TRANSTAB_BUS_FAULT_1:
+		fallthrough;
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_TRANSTAB_BUS_FAULT_2:
+		fallthrough;
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_TRANSTAB_BUS_FAULT_3:
+		kbase_mmu_report_fault_and_kill(kctx, faulting_as, "Translation table bus fault",
+						fault);
 		goto fault_done;
+#endif
 
-	case AS_FAULTSTATUS_EXCEPTION_CODE_ACCESS_FLAG:
+#if !MALI_USE_CSF
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_ACCESS_FLAG_0:
+		fallthrough;
+#endif
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_ACCESS_FLAG_1:
+		fallthrough;
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_ACCESS_FLAG_2:
+		fallthrough;
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_ACCESS_FLAG_3:
 		/* nothing to do, but we don't expect this fault currently */
 		dev_warn(kbdev->dev, "Access flag unexpectedly set");
 		goto fault_done;
 
-	case AS_FAULTSTATUS_EXCEPTION_CODE_ADDRESS_SIZE_FAULT:
-		kbase_mmu_report_fault_and_kill(kctx, faulting_as,
-				"Address size fault", fault);
+#if MALI_USE_CSF
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_ADDRESS_SIZE_FAULT_IN:
+		fallthrough;
+#else
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_ADDRESS_SIZE_FAULT_IN0:
+		fallthrough;
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_ADDRESS_SIZE_FAULT_IN1:
+		fallthrough;
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_ADDRESS_SIZE_FAULT_IN2:
+		fallthrough;
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_ADDRESS_SIZE_FAULT_IN3:
+		fallthrough;
+#endif
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_ADDRESS_SIZE_FAULT_OUT0:
+		fallthrough;
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_ADDRESS_SIZE_FAULT_OUT1:
+		fallthrough;
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_ADDRESS_SIZE_FAULT_OUT2:
+		fallthrough;
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_ADDRESS_SIZE_FAULT_OUT3:
+		kbase_mmu_report_fault_and_kill(kctx, faulting_as, "Address size fault", fault);
 		goto fault_done;
 
-	case AS_FAULTSTATUS_EXCEPTION_CODE_MEMORY_ATTRIBUTES_FAULT:
-		kbase_mmu_report_fault_and_kill(kctx, faulting_as,
-				"Memory attributes fault", fault);
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_MEMORY_ATTRIBUTE_FAULT_0:
+		fallthrough;
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_MEMORY_ATTRIBUTE_FAULT_1:
+		fallthrough;
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_MEMORY_ATTRIBUTE_FAULT_2:
+		fallthrough;
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_MEMORY_ATTRIBUTE_FAULT_3:
+#if !MALI_USE_CSF
+		fallthrough;
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_MEMORY_ATTRIBUTE_NONCACHEABLE_0:
+		fallthrough;
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_MEMORY_ATTRIBUTE_NONCACHEABLE_1:
+		fallthrough;
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_MEMORY_ATTRIBUTE_NONCACHEABLE_2:
+		fallthrough;
+	case AS_FAULTSTATUS_EXCEPTION_TYPE_MEMORY_ATTRIBUTE_NONCACHEABLE_3:
+#endif
+		kbase_mmu_report_fault_and_kill(kctx, faulting_as, "Memory attributes fault",
+						fault);
 		goto fault_done;
 
 	default:
-		kbase_mmu_report_fault_and_kill(kctx, faulting_as,
-				"Unknown fault code", fault);
+		kbase_mmu_report_fault_and_kill(kctx, faulting_as, "Unknown fault code", fault);
 		goto fault_done;
 	}
 
@@ -1118,46 +1115,42 @@ page_fault_retry:
 	 */
 	kbase_gpu_vm_lock(kctx);
 
-	region = kbase_region_tracker_find_region_enclosing_address(kctx,
-			fault->addr);
+	region = kbase_region_tracker_find_region_enclosing_address(kctx, fault->addr);
 	if (kbase_is_region_invalid_or_free(region)) {
 		kbase_gpu_vm_unlock(kctx);
 		kbase_mmu_report_fault_and_kill(kctx, faulting_as,
-				"Memory is not mapped on the GPU", fault);
+						"Memory is not mapped on the GPU", fault);
 		goto fault_done;
 	}
 
 	if (region->gpu_alloc->type == KBASE_MEM_TYPE_IMPORTED_UMM) {
 		kbase_gpu_vm_unlock(kctx);
 		kbase_mmu_report_fault_and_kill(kctx, faulting_as,
-				"DMA-BUF is not mapped on the GPU", fault);
+						"DMA-BUF is not mapped on the GPU", fault);
 		goto fault_done;
 	}
 
 	if (region->gpu_alloc->group_id >= MEMORY_GROUP_MANAGER_NR_GROUPS) {
 		kbase_gpu_vm_unlock(kctx);
-		kbase_mmu_report_fault_and_kill(kctx, faulting_as,
-				"Bad physical memory group ID", fault);
+		kbase_mmu_report_fault_and_kill(kctx, faulting_as, "Bad physical memory group ID",
+						fault);
 		goto fault_done;
 	}
 
-	if ((region->flags & GROWABLE_FLAGS_REQUIRED)
-			!= GROWABLE_FLAGS_REQUIRED) {
+	if ((region->flags & GROWABLE_FLAGS_REQUIRED) != GROWABLE_FLAGS_REQUIRED) {
 		kbase_gpu_vm_unlock(kctx);
-		kbase_mmu_report_fault_and_kill(kctx, faulting_as,
-				"Memory is not growable", fault);
+		kbase_mmu_report_fault_and_kill(kctx, faulting_as, "Memory is not growable", fault);
 		goto fault_done;
 	}
 
 	if ((region->flags & KBASE_REG_DONT_NEED)) {
 		kbase_gpu_vm_unlock(kctx);
 		kbase_mmu_report_fault_and_kill(kctx, faulting_as,
-				"Don't need memory can't be grown", fault);
+						"Don't need memory can't be grown", fault);
 		goto fault_done;
 	}
 
-	if (AS_FAULTSTATUS_ACCESS_TYPE_GET(fault_status) ==
-		AS_FAULTSTATUS_ACCESS_TYPE_READ)
+	if (AS_FAULTSTATUS_ACCESS_TYPE_GET(fault_status) == AS_FAULTSTATUS_ACCESS_TYPE_READ)
 		dev_warn(kbdev->dev, "Grow on pagefault while reading");
 
 	/* find the size we need to grow it by
@@ -1173,15 +1166,10 @@ page_fault_retry:
 		struct kbase_mmu_hw_op_param op_param;
 
 		dev_dbg(kbdev->dev,
-			"Page fault @ 0x%llx in allocated region 0x%llx-0x%llx of growable TMEM: Ignoring",
-				fault->addr, region->start_pfn,
-				region->start_pfn +
-				current_backed_size);
+			"Page fault @ VA 0x%llx in allocated region 0x%llx-0x%llx of growable TMEM: Ignoring",
+			fault->addr, region->start_pfn, region->start_pfn + current_backed_size);
 
-		mutex_lock(&kbdev->mmu_hw_mutex);
-
-		kbase_mmu_hw_clear_fault(kbdev, faulting_as,
-				KBASE_MMU_FAULT_TYPE_PAGE);
+		kbase_mmu_hw_clear_fault(kbdev, faulting_as, KBASE_MMU_FAULT_TYPE_PAGE);
 		/* [1] in case another page fault occurred while we were
 		 * handling the (duplicate) page fault we need to ensure we
 		 * don't loose the other page fault as result of us clearing
@@ -1192,32 +1180,23 @@ page_fault_retry:
 		 */
 		op_param.mmu_sync_info = mmu_sync_info;
 		op_param.kctx_id = kctx->id;
-		if (!mmu_flush_cache_on_gpu_ctrl(kbdev)) {
-			mmu_hw_operation_begin(kbdev);
-			err = kbase_mmu_hw_do_unlock_no_addr(kbdev, faulting_as,
-							     &op_param);
-			mmu_hw_operation_end(kbdev);
-		} else {
-			/* Can safely skip the invalidate for all levels in case
-			 * of duplicate page faults.
-			 */
-			op_param.flush_skip_levels = 0xF;
-			op_param.vpfn = fault_pfn;
-			op_param.nr = 1;
-			err = kbase_mmu_hw_do_unlock(kbdev, faulting_as,
-						     &op_param);
-		}
+		/* Can safely skip the invalidate for all levels in case
+		 * of duplicate page faults.
+		 */
+		op_param.flush_skip_levels = 0xF;
+		op_param.vpfn = fault_pfn;
+		op_param.nr = 1;
+		spin_lock_irqsave(&kbdev->hwaccess_lock, hwaccess_flags);
+		err = kbase_mmu_hw_do_unlock(kbdev, faulting_as, &op_param);
+		spin_unlock_irqrestore(&kbdev->hwaccess_lock, hwaccess_flags);
 
 		if (err) {
 			dev_err(kbdev->dev,
-				"Invalidation for MMU did not complete on handling page fault @ 0x%llx",
+				"Invalidation for MMU did not complete on handling page fault @ VA 0x%llx",
 				fault->addr);
 		}
 
-		mutex_unlock(&kbdev->mmu_hw_mutex);
-
-		kbase_mmu_hw_enable_fault(kbdev, faulting_as,
-				KBASE_MMU_FAULT_TYPE_PAGE);
+		kbase_mmu_hw_enable_fault(kbdev, faulting_as, KBASE_MMU_FAULT_TYPE_PAGE);
 		kbase_gpu_vm_unlock(kctx);
 
 		goto fault_done;
@@ -1232,41 +1211,29 @@ page_fault_retry:
 	if (new_pages == 0) {
 		struct kbase_mmu_hw_op_param op_param;
 
-		mutex_lock(&kbdev->mmu_hw_mutex);
-
 		/* Duplicate of a fault we've already handled, nothing to do */
-		kbase_mmu_hw_clear_fault(kbdev, faulting_as,
-				KBASE_MMU_FAULT_TYPE_PAGE);
+		kbase_mmu_hw_clear_fault(kbdev, faulting_as, KBASE_MMU_FAULT_TYPE_PAGE);
 
 		/* See comment [1] about UNLOCK usage */
 		op_param.mmu_sync_info = mmu_sync_info;
 		op_param.kctx_id = kctx->id;
-		if (!mmu_flush_cache_on_gpu_ctrl(kbdev)) {
-			mmu_hw_operation_begin(kbdev);
-			err = kbase_mmu_hw_do_unlock_no_addr(kbdev, faulting_as,
-							     &op_param);
-			mmu_hw_operation_end(kbdev);
-		} else {
-			/* Can safely skip the invalidate for all levels in case
-			 * of duplicate page faults.
-			 */
-			op_param.flush_skip_levels = 0xF;
-			op_param.vpfn = fault_pfn;
-			op_param.nr = 1;
-			err = kbase_mmu_hw_do_unlock(kbdev, faulting_as,
-						     &op_param);
-		}
+		/* Can safely skip the invalidate for all levels in case
+		 * of duplicate page faults.
+		 */
+		op_param.flush_skip_levels = 0xF;
+		op_param.vpfn = fault_pfn;
+		op_param.nr = 1;
+		spin_lock_irqsave(&kbdev->hwaccess_lock, hwaccess_flags);
+		err = kbase_mmu_hw_do_unlock(kbdev, faulting_as, &op_param);
+		spin_unlock_irqrestore(&kbdev->hwaccess_lock, hwaccess_flags);
 
 		if (err) {
 			dev_err(kbdev->dev,
-				"Invalidation for MMU did not complete on handling page fault @ 0x%llx",
+				"Invalidation for MMU did not complete on handling page fault @ VA 0x%llx",
 				fault->addr);
 		}
 
-		mutex_unlock(&kbdev->mmu_hw_mutex);
-
-		kbase_mmu_hw_enable_fault(kbdev, faulting_as,
-				KBASE_MMU_FAULT_TYPE_PAGE);
+		kbase_mmu_hw_enable_fault(kbdev, faulting_as, KBASE_MMU_FAULT_TYPE_PAGE);
 		kbase_gpu_vm_unlock(kctx);
 		goto fault_done;
 	}
@@ -1281,8 +1248,8 @@ page_fault_retry:
 #endif
 
 	spin_lock(&kctx->mem_partials_lock);
-	grown = page_fault_try_alloc(kctx, region, new_pages, &pages_to_grow,
-			&grow_2mb_pool, prealloc_sas);
+	grown = page_fault_try_alloc(kctx, region, new_pages, &pages_to_grow, &grow_2mb_pool,
+				     prealloc_sas);
 	spin_unlock(&kctx->mem_partials_lock);
 
 	if (grown) {
@@ -1291,8 +1258,7 @@ page_fault_retry:
 		struct kbase_mmu_hw_op_param op_param;
 
 		/* alloc success */
-		WARN_ON(kbase_reg_current_backed_size(region) >
-			region->nr_pages);
+		WARN_ON(kbase_reg_current_backed_size(region) > region->nr_pages);
 
 		/* set up the new pages */
 		pfn_offset = kbase_reg_current_backed_size(region) - new_pages;
@@ -1310,30 +1276,29 @@ page_fault_retry:
 						region->gpu_alloc->group_id, &dirty_pgds, region,
 						false);
 		if (err) {
-			kbase_free_phy_pages_helper(region->gpu_alloc,
-					new_pages);
+			kbase_free_phy_pages_helper(region->gpu_alloc, new_pages);
 			if (region->gpu_alloc != region->cpu_alloc)
-				kbase_free_phy_pages_helper(region->cpu_alloc,
-						new_pages);
+				kbase_free_phy_pages_helper(region->cpu_alloc, new_pages);
 			kbase_gpu_vm_unlock(kctx);
 			/* The locked VA region will be unlocked and the cache
 			 * invalidated in here
 			 */
 			kbase_mmu_report_fault_and_kill(kctx, faulting_as,
-					"Page table update failure", fault);
+							"Page table update failure", fault);
 			goto fault_done;
 		}
-		KBASE_TLSTREAM_AUX_PAGEFAULT(kbdev, kctx->id, as_no,
-				(u64)new_pages);
-		trace_mali_mmu_page_fault_grow(region, fault, new_pages);
+		KBASE_TLSTREAM_AUX_PAGEFAULT(kbdev, kctx->id, as_no, (u64)new_pages);
+		if (kbase_reg_is_valid(kbdev, MMU_AS_OFFSET(as_no, FAULTEXTRA)))
+			trace_mali_mmu_page_fault_extra_grow(region, fault, new_pages);
+		else
+			trace_mali_mmu_page_fault_grow(region, fault, new_pages);
 
 #if MALI_INCREMENTAL_RENDERING_JM
 		/* Switch to incremental rendering if we have nearly run out of
 		 * memory in a JIT memory allocation.
 		 */
 		if (region->threshold_pages &&
-			kbase_reg_current_backed_size(region) >
-				region->threshold_pages) {
+		    kbase_reg_current_backed_size(region) > region->threshold_pages) {
 			dev_dbg(kctx->kbdev->dev, "%zu pages exceeded IR threshold %zu",
 				new_pages + current_backed_size, region->threshold_pages);
 
@@ -1345,7 +1310,6 @@ page_fault_retry:
 #endif
 
 		/* AS transaction begin */
-		mutex_lock(&kbdev->mmu_hw_mutex);
 
 		/* clear MMU interrupt - this needs to be done after updating
 		 * the page tables but before issuing a FLUSH command. The
@@ -1355,40 +1319,34 @@ page_fault_retry:
 		 * this stage a new IRQ might not be raised when the GPU finds
 		 * a MMU IRQ is already pending.
 		 */
-		kbase_mmu_hw_clear_fault(kbdev, faulting_as,
-					 KBASE_MMU_FAULT_TYPE_PAGE);
+		kbase_mmu_hw_clear_fault(kbdev, faulting_as, KBASE_MMU_FAULT_TYPE_PAGE);
 
 		op_param.vpfn = region->start_pfn + pfn_offset;
 		op_param.nr = new_pages;
 		op_param.op = KBASE_MMU_OP_FLUSH_PT;
 		op_param.kctx_id = kctx->id;
 		op_param.mmu_sync_info = mmu_sync_info;
+		spin_lock_irqsave(&kbdev->hwaccess_lock, hwaccess_flags);
 		if (mmu_flush_cache_on_gpu_ctrl(kbdev)) {
 			/* Unlock to invalidate the TLB (and resume the MMU) */
-			op_param.flush_skip_levels =
-				pgd_level_to_skip_flush(dirty_pgds);
-			err = kbase_mmu_hw_do_unlock(kbdev, faulting_as,
-						     &op_param);
+			op_param.flush_skip_levels = pgd_level_to_skip_flush(dirty_pgds);
+			err = kbase_mmu_hw_do_unlock(kbdev, faulting_as, &op_param);
 		} else {
 			/* flush L2 and unlock the VA (resumes the MMU) */
-			mmu_hw_operation_begin(kbdev);
-			err = kbase_mmu_hw_do_flush(kbdev, faulting_as,
-						    &op_param);
-			mmu_hw_operation_end(kbdev);
+			err = kbase_mmu_hw_do_flush(kbdev, faulting_as, &op_param);
 		}
+		spin_unlock_irqrestore(&kbdev->hwaccess_lock, hwaccess_flags);
 
 		if (err) {
 			dev_err(kbdev->dev,
-				"Flush for GPU page table update did not complete on handling page fault @ 0x%llx",
+				"Flush for GPU page table update did not complete on handling page fault @ VA 0x%llx",
 				fault->addr);
 		}
 
-		mutex_unlock(&kbdev->mmu_hw_mutex);
 		/* AS transaction end */
 
 		/* reenable this in the mask */
-		kbase_mmu_hw_enable_fault(kbdev, faulting_as,
-					 KBASE_MMU_FAULT_TYPE_PAGE);
+		kbase_mmu_hw_enable_fault(kbdev, faulting_as, KBASE_MMU_FAULT_TYPE_PAGE);
 
 #ifdef CONFIG_MALI_CINSTR_GWT
 		if (kctx->gwt_enabled) {
@@ -1398,12 +1356,9 @@ page_fault_retry:
 			pos = kmalloc(sizeof(*pos), GFP_KERNEL);
 			if (pos) {
 				pos->region = region;
-				pos->page_addr = (region->start_pfn +
-							pfn_offset) <<
-							 PAGE_SHIFT;
+				pos->page_addr = (region->start_pfn + pfn_offset) << PAGE_SHIFT;
 				pos->num_pages = new_pages;
-				list_add(&pos->link,
-					&kctx->gwt_current_list);
+				list_add(&pos->link, &kctx->gwt_current_list);
 			} else {
 				dev_warn(kbdev->dev, "kmalloc failure");
 			}
@@ -1432,24 +1387,24 @@ page_fault_retry:
 				struct kbase_mem_pool *const lp_mem_pool =
 					&kctx->mem_pools.large[group_id];
 
-				pages_to_grow = (pages_to_grow +
-					((1 << lp_mem_pool->order) - 1))
-						>> lp_mem_pool->order;
+				pages_to_grow =
+					(pages_to_grow + ((1u << lp_mem_pool->order) - 1u)) >>
+					lp_mem_pool->order;
 
-				ret = kbase_mem_pool_grow(lp_mem_pool,
-					pages_to_grow, kctx->task);
+				ret = kbase_mem_pool_grow(lp_mem_pool, pages_to_grow, kctx->task);
 			} else {
 				struct kbase_mem_pool *const mem_pool =
 					&kctx->mem_pools.small[group_id];
 
-				ret = kbase_mem_pool_grow(mem_pool,
-					pages_to_grow, kctx->task);
+				ret = kbase_mem_pool_grow(mem_pool, pages_to_grow, kctx->task);
 			}
 		}
 		if (ret < 0) {
 			/* failed to extend, handle as a normal PF */
+			if (unlikely(ret == -EPERM))
+				kbase_ctx_flag_set(kctx, KCTX_PAGE_FAULT_REPORT_SKIP);
 			kbase_mmu_report_fault_and_kill(kctx, faulting_as,
-					"Page allocation failure", fault);
+							"Page allocation failure", fault);
 		} else {
 			dev_dbg(kbdev->dev, "Try again after pool_grow");
 			goto page_fault_retry;
@@ -1481,8 +1436,7 @@ fault_done:
 	dev_dbg(kbdev->dev, "Leaving page_fault_worker %pK", (void *)data);
 }
 
-static phys_addr_t kbase_mmu_alloc_pgd(struct kbase_device *kbdev,
-		struct kbase_mmu_table *mmut)
+static phys_addr_t kbase_mmu_alloc_pgd(struct kbase_device *kbdev, struct kbase_mmu_table *mmut)
 {
 	u64 *page;
 	struct page *p;
@@ -1490,7 +1444,7 @@ static phys_addr_t kbase_mmu_alloc_pgd(struct kbase_device *kbdev,
 
 	p = kbase_mem_pool_alloc(&kbdev->mem_pools.small[mmut->group_id]);
 	if (!p)
-		return KBASE_MMU_INVALID_PGD_ADDRESS;
+		return KBASE_INVALID_PHYSICAL_ADDRESS;
 
 	page = kbase_kmap(p);
 
@@ -1506,12 +1460,8 @@ static phys_addr_t kbase_mmu_alloc_pgd(struct kbase_device *kbdev,
 	if (mmut->kctx) {
 		int new_page_count;
 
-		new_page_count = atomic_add_return(1,
-			&mmut->kctx->used_pages);
-		KBASE_TLSTREAM_AUX_PAGESALLOC(
-			kbdev,
-			mmut->kctx->id,
-			(u64)new_page_count);
+		new_page_count = atomic_add_return(1, &mmut->kctx->used_pages);
+		KBASE_TLSTREAM_AUX_PAGESALLOC(kbdev, mmut->kctx->id, (u64)new_page_count);
 		kbase_process_page_usage_inc(mmut->kctx, 1);
 	}
 
@@ -1532,7 +1482,7 @@ static phys_addr_t kbase_mmu_alloc_pgd(struct kbase_device *kbdev,
 alloc_free:
 	kbase_mem_pool_free(&kbdev->mem_pools.small[mmut->group_id], p, false);
 
-	return KBASE_MMU_INVALID_PGD_ADDRESS;
+	return KBASE_INVALID_PHYSICAL_ADDRESS;
 }
 
 /**
@@ -1541,7 +1491,7 @@ alloc_free:
  * @kbdev:    Device pointer.
  * @mmut:     GPU MMU page table.
  * @pgd:      Physical addresse of level N page directory.
- * @vpfn:     The virtual page frame number.
+ * @vpfn:     The virtual page frame number, in GPU_PAGE_SIZE units.
  * @level:    The level of MMU page table (N).
  *
  * Return:
@@ -1594,7 +1544,7 @@ static int mmu_get_next_pgd(struct kbase_device *kbdev, struct kbase_mmu_table *
  *
  * @kbdev:    Device pointer.
  * @mmut:     GPU MMU page table.
- * @vpfn:     The virtual page frame number.
+ * @vpfn:     The virtual page frame number, in GPU_PAGE_SIZE units.
  * @in_level:     The level of MMU page table (N).
  * @out_level:    Set to the level of the lowest valid PGD found on success.
  *                Invalid on error.
@@ -1683,8 +1633,10 @@ static void mmu_insert_pages_failure_recovery(struct kbase_device *kbdev,
 	u64 vpfn = from_vpfn;
 	struct kbase_mmu_mode const *mmu_mode;
 
+	/* Both from_vpfn and to_vpfn are in GPU_PAGE_SIZE units */
+
 	/* 64-bit address range is the max */
-	KBASE_DEBUG_ASSERT(vpfn <= (U64_MAX / PAGE_SIZE));
+	KBASE_DEBUG_ASSERT(vpfn <= (U64_MAX / GPU_PAGE_SIZE));
 	KBASE_DEBUG_ASSERT(from_vpfn <= to_vpfn);
 
 	lockdep_assert_held(&mmut->mmu_lock);
@@ -1708,9 +1660,8 @@ static void mmu_insert_pages_failure_recovery(struct kbase_device *kbdev,
 		if (count > left)
 			count = left;
 
-		/* need to check if this is a 2MB page or a 4kB */
-		for (level = MIDGARD_MMU_TOPLEVEL;
-				level <= MIDGARD_MMU_BOTTOMLEVEL; level++) {
+		/* need to check if this is a 2MB page or a small page */
+		for (level = MIDGARD_MMU_TOPLEVEL; level <= MIDGARD_MMU_BOTTOMLEVEL; level++) {
 			idx = (vpfn >> ((3 - level) * 9)) & 0x1FF;
 			pgds[level] = pgd;
 			page = kbase_kmap(p);
@@ -1778,7 +1729,7 @@ next:
 	 * movable once they are returned to a memory pool.
 	 */
 	if (kbase_is_page_migration_enabled() && !ignore_page_migration && phys) {
-		const u64 num_pages = to_vpfn - from_vpfn + 1;
+		const u64 num_pages = (to_vpfn - from_vpfn) / GPU_PAGES_PER_CPU_PAGE;
 		u64 i;
 
 		for (i = 0; i < num_pages; i++) {
@@ -1844,7 +1795,7 @@ static void mmu_flush_invalidate_insert_pages(struct kbase_device *kbdev,
  *                The bottom PGD level.
  * @insert_level: The level of MMU page table where the chain of newly allocated
  *                PGDs needs to be linked-in/inserted.
- * @insert_vpfn:  The virtual page frame number for the ATE.
+ * @insert_vpfn:  The virtual page frame number, in GPU_PAGE_SIZE units, for the ATE.
  * @pgds_to_insert: Ptr to an array (size MIDGARD_MMU_BOTTOMLEVEL+1) that contains
  *                  the physical addresses of newly allocated PGDs from index
  *                  insert_level+1 to cur_level, and an existing PGD at index
@@ -1880,7 +1831,7 @@ static int update_parent_pgds(struct kbase_device *kbdev, struct kbase_mmu_table
 		struct page *parent_page = pfn_to_page(PFN_DOWN(parent_pgd));
 		u64 *parent_page_va;
 
-		if (WARN_ON_ONCE(target_pgd == KBASE_MMU_INVALID_PGD_ADDRESS)) {
+		if (WARN_ON_ONCE(target_pgd == KBASE_INVALID_PHYSICAL_ADDRESS)) {
 			err = -EFAULT;
 			goto failure_recovery;
 		}
@@ -1987,12 +1938,12 @@ static int mmu_insert_alloc_pgds(struct kbase_device *kbdev, struct kbase_mmu_ta
 	for (i = level_low; i <= level_high; i++) {
 		do {
 			new_pgds[i] = kbase_mmu_alloc_pgd(kbdev, mmut);
-			if (new_pgds[i] != KBASE_MMU_INVALID_PGD_ADDRESS)
+			if (new_pgds[i] != KBASE_INVALID_PHYSICAL_ADDRESS)
 				break;
 
 			mutex_unlock(&mmut->mmu_lock);
 			err = kbase_mem_pool_grow(&kbdev->mem_pools.small[mmut->group_id],
-						  level_high, NULL);
+						  (size_t)level_high, NULL);
 			mutex_lock(&mmut->mmu_lock);
 			if (err) {
 				dev_err(kbdev->dev, "%s: kbase_mem_pool_grow() returned error %d",
@@ -2002,7 +1953,7 @@ static int mmu_insert_alloc_pgds(struct kbase_device *kbdev, struct kbase_mmu_ta
 				 * from (i-1) to level_low
 				 */
 				for (i = (i - 1); i >= level_low; i--) {
-					if (new_pgds[i] != KBASE_MMU_INVALID_PGD_ADDRESS)
+					if (new_pgds[i] != KBASE_INVALID_PHYSICAL_ADDRESS)
 						kbase_mmu_free_pgd(kbdev, mmut, new_pgds[i]);
 				}
 
@@ -2032,6 +1983,7 @@ static int kbase_mmu_insert_single_page(struct kbase_context *kctx, u64 start_vp
 	enum kbase_mmu_op_type flush_op;
 	struct kbase_mmu_table *mmut = &kctx->mmu;
 	int l, cur_level, insert_level;
+	const phys_addr_t base_phys_address = as_phys_addr_t(phys);
 
 	if (WARN_ON(kctx == NULL))
 		return -EINVAL;
@@ -2045,6 +1997,10 @@ static int kbase_mmu_insert_single_page(struct kbase_context *kctx, u64 start_vp
 	if (nr == 0)
 		return 0;
 
+	/* Convert to GPU_PAGE_SIZE units. */
+	insert_vpfn *= GPU_PAGES_PER_CPU_PAGE;
+	remain *= GPU_PAGES_PER_CPU_PAGE;
+
 	/* If page migration is enabled, pages involved in multiple GPU mappings
 	 * are always treated as not movable.
 	 */
@@ -2122,21 +2078,24 @@ static int kbase_mmu_insert_single_page(struct kbase_context *kctx, u64 start_vp
 			goto fail_unlock_free_pgds;
 		}
 
-		num_of_valid_entries =
-			kbdev->mmu_mode->get_num_valid_entries(pgd_page);
+		num_of_valid_entries = kbdev->mmu_mode->get_num_valid_entries(pgd_page);
 
-		for (i = 0; i < count; i++) {
-			unsigned int ofs = vindex + i;
+		for (i = 0; i < count; i += GPU_PAGES_PER_CPU_PAGE) {
+			unsigned int j;
 
-			/* Fail if the current page is a valid ATE entry */
-			KBASE_DEBUG_ASSERT(0 == (pgd_page[ofs] & 1UL));
+			for (j = 0; j < GPU_PAGES_PER_CPU_PAGE; j++) {
+				unsigned int ofs = vindex + i + j;
+				phys_addr_t page_address = base_phys_address + (j * GPU_PAGE_SIZE);
 
-			pgd_page[ofs] = kbase_mmu_create_ate(kbdev,
-				phys, flags, MIDGARD_MMU_BOTTOMLEVEL, group_id);
+				/* Fail if the current page is a valid ATE entry */
+				WARN_ON_ONCE((pgd_page[ofs] & 1UL));
+				pgd_page[ofs] = kbase_mmu_create_ate(kbdev, as_tagged(page_address),
+								     flags, MIDGARD_MMU_BOTTOMLEVEL,
+								     group_id);
+			}
 		}
 
-		kbdev->mmu_mode->set_num_valid_entries(
-			pgd_page, num_of_valid_entries + count);
+		kbdev->mmu_mode->set_num_valid_entries(pgd_page, num_of_valid_entries + count);
 
 		dirty_pgds |= 1ULL << (newly_created_pgd ? insert_level : MIDGARD_MMU_BOTTOMLEVEL);
 
@@ -2185,10 +2144,10 @@ fail_unlock_free_pgds:
 		kbase_mmu_free_pgd(kbdev, mmut, new_pgds[l]);
 
 fail_unlock:
-	if (insert_vpfn != start_vpfn) {
+	if (insert_vpfn != (start_vpfn * GPU_PAGES_PER_CPU_PAGE)) {
 		/* Invalidate the pages we have partially completed */
-		mmu_insert_pages_failure_recovery(kbdev, mmut, start_vpfn, insert_vpfn, &dirty_pgds,
-						  NULL, true);
+		mmu_insert_pages_failure_recovery(kbdev, mmut, start_vpfn * GPU_PAGES_PER_CPU_PAGE,
+						  insert_vpfn, &dirty_pgds, NULL, true);
 	}
 
 	mmu_flush_invalidate_insert_pages(kbdev, mmut, start_vpfn, nr, dirty_pgds, mmu_sync_info,
@@ -2270,7 +2229,7 @@ static void kbase_mmu_progress_migration_on_teardown(struct kbase_device *kbdev,
 		struct page *phys_page = as_page(phys[i]);
 		struct kbase_page_metadata *page_md = kbase_page_private(phys_page);
 
-		/* Skip the 4KB page that is part of a large page, as the large page is
+		/* Skip the small page that is part of a large page, as the large page is
 		 * excluded from the migration process.
 		 */
 		if (is_huge(phys[i]) || is_partial(phys[i]))
@@ -2293,7 +2252,7 @@ static void kbase_mmu_progress_migration_on_teardown(struct kbase_device *kbdev,
 					 * status will subsequently be freed in either
 					 * kbase_page_migrate() or kbase_page_putback()
 					 */
-					phys[i] = as_tagged(0);
+					phys[i] = as_tagged(KBASE_INVALID_PHYSICAL_ADDRESS);
 				} else
 					page_md->status = PAGE_STATUS_SET(page_md->status,
 									  (u8)FREE_IN_PROGRESS);
@@ -2304,19 +2263,18 @@ static void kbase_mmu_progress_migration_on_teardown(struct kbase_device *kbdev,
 	}
 }
 
-u64 kbase_mmu_create_ate(struct kbase_device *const kbdev,
-	struct tagged_addr const phy, unsigned long const flags,
-	int const level, int const group_id)
+u64 kbase_mmu_create_ate(struct kbase_device *const kbdev, struct tagged_addr const phy,
+			 unsigned long const flags, int const level, int const group_id)
 {
 	u64 entry;
 
 	kbdev->mmu_mode->entry_set_ate(&entry, phy, flags, level);
-	return kbdev->mgm_dev->ops.mgm_update_gpu_pte(kbdev->mgm_dev,
-		group_id, level, entry);
+	return kbdev->mgm_dev->ops.mgm_update_gpu_pte(kbdev->mgm_dev, (unsigned int)group_id, level,
+						      entry);
 }
 
 static int mmu_insert_pages_no_flush(struct kbase_device *kbdev, struct kbase_mmu_table *mmut,
-				     const u64 start_vpfn, struct tagged_addr *phys, size_t nr,
+				     u64 start_vpfn, struct tagged_addr *phys, size_t nr,
 				     unsigned long flags, int const group_id, u64 *dirty_pgds,
 				     struct kbase_va_region *reg, bool ignore_page_migration)
 {
@@ -2329,6 +2287,7 @@ static int mmu_insert_pages_no_flush(struct kbase_device *kbdev, struct kbase_mm
 	unsigned int i;
 	phys_addr_t new_pgds[MIDGARD_MMU_BOTTOMLEVEL + 1];
 	int l, cur_level, insert_level;
+	struct tagged_addr *start_phys = phys;
 
 	/* Note that 0 is a valid start_vpfn */
 	/* 64-bit address range is the max */
@@ -2340,6 +2299,9 @@ static int mmu_insert_pages_no_flush(struct kbase_device *kbdev, struct kbase_mm
 	if (nr == 0)
 		return 0;
 
+	/* Convert to GPU_PAGE_SIZE units. */
+	insert_vpfn *= GPU_PAGES_PER_CPU_PAGE;
+	remain *= GPU_PAGES_PER_CPU_PAGE;
 	mutex_lock(&mmut->mmu_lock);
 
 	while (remain) {
@@ -2408,8 +2370,7 @@ static int mmu_insert_pages_no_flush(struct kbase_device *kbdev, struct kbase_mm
 			goto fail_unlock_free_pgds;
 		}
 
-		num_of_valid_entries =
-			mmu_mode->get_num_valid_entries(pgd_page);
+		num_of_valid_entries = mmu_mode->get_num_valid_entries(pgd_page);
 
 		if (cur_level == MIDGARD_MMU_LEVEL(2)) {
 			int level_index = (insert_vpfn >> 9) & 0x1FF;
@@ -2418,29 +2379,39 @@ static int mmu_insert_pages_no_flush(struct kbase_device *kbdev, struct kbase_mm
 
 			num_of_valid_entries++;
 		} else {
-			for (i = 0; i < count; i++) {
-				unsigned int ofs = vindex + i;
-				u64 *target = &pgd_page[ofs];
-
-				/* Warn if the current page is a valid ATE
-				 * entry. The page table shouldn't have anything
-				 * in the place where we are trying to put a
-				 * new entry. Modification to page table entries
-				 * should be performed with
-				 * kbase_mmu_update_pages()
-				 */
-				WARN_ON((*target & 1UL) != 0);
+			for (i = 0; i < count; i += GPU_PAGES_PER_CPU_PAGE) {
+				struct tagged_addr base_tagged_addr =
+					phys[i / GPU_PAGES_PER_CPU_PAGE];
+				phys_addr_t base_phys_address = as_phys_addr_t(base_tagged_addr);
+				unsigned int j;
+
+				for (j = 0; j < GPU_PAGES_PER_CPU_PAGE; j++) {
+					unsigned int ofs = vindex + i + j;
+					u64 *target = &pgd_page[ofs];
+					phys_addr_t page_address =
+						base_phys_address + (j * GPU_PAGE_SIZE);
+
+					/* Warn if the current page is a valid ATE
+					 * entry. The page table shouldn't have anything
+					 * in the place where we are trying to put a
+					 * new entry. Modification to page table entries
+					 * should be performed with
+					 * kbase_mmu_update_pages()
+					 */
+					WARN_ON_ONCE((*target & 1UL) != 0);
 
-				*target = kbase_mmu_create_ate(kbdev,
-					phys[i], flags, cur_level, group_id);
+					*target = kbase_mmu_create_ate(kbdev,
+								       as_tagged(page_address),
+								       flags, cur_level, group_id);
+				}
 
 				/* If page migration is enabled, this is the right time
 				 * to update the status of the page.
 				 */
 				if (kbase_is_page_migration_enabled() && !ignore_page_migration &&
-				    !is_huge(phys[i]) && !is_partial(phys[i]))
-					kbase_mmu_progress_migration_on_insert(phys[i], reg, mmut,
-									       insert_vpfn + i);
+				    !is_huge(base_tagged_addr) && !is_partial(base_tagged_addr))
+					kbase_mmu_progress_migration_on_insert(
+						base_tagged_addr, reg, mmut, insert_vpfn + i);
 			}
 			num_of_valid_entries += count;
 		}
@@ -2477,7 +2448,7 @@ static int mmu_insert_pages_no_flush(struct kbase_device *kbdev, struct kbase_mm
 			}
 		}
 
-		phys += count;
+		phys += (count / GPU_PAGES_PER_CPU_PAGE);
 		insert_vpfn += count;
 		remain -= count;
 		kbase_kunmap(p, pgd_page);
@@ -2493,10 +2464,11 @@ fail_unlock_free_pgds:
 		kbase_mmu_free_pgd(kbdev, mmut, new_pgds[l]);
 
 fail_unlock:
-	if (insert_vpfn != start_vpfn) {
+	if (insert_vpfn != (start_vpfn * GPU_PAGES_PER_CPU_PAGE)) {
 		/* Invalidate the pages we have partially completed */
-		mmu_insert_pages_failure_recovery(kbdev, mmut, start_vpfn, insert_vpfn, dirty_pgds,
-						  phys, ignore_page_migration);
+		mmu_insert_pages_failure_recovery(kbdev, mmut, start_vpfn * GPU_PAGES_PER_CPU_PAGE,
+						  insert_vpfn, dirty_pgds, start_phys,
+						  ignore_page_migration);
 	}
 
 	mmu_flush_invalidate_insert_pages(kbdev, mmut, start_vpfn, nr,
@@ -2536,6 +2508,8 @@ int kbase_mmu_insert_pages(struct kbase_device *kbdev, struct kbase_mmu_table *m
 	int err;
 	u64 dirty_pgds = 0;
 
+	CSTD_UNUSED(as_nr);
+
 	/* Early out if there is nothing to do */
 	if (nr == 0)
 		return 0;
@@ -2562,6 +2536,8 @@ int kbase_mmu_insert_pages_skip_status_update(struct kbase_device *kbdev,
 	int err;
 	u64 dirty_pgds = 0;
 
+	CSTD_UNUSED(as_nr);
+
 	/* Early out if there is nothing to do */
 	if (nr == 0)
 		return 0;
@@ -2588,6 +2564,8 @@ int kbase_mmu_insert_aliased_pages(struct kbase_device *kbdev, struct kbase_mmu_
 	int err;
 	u64 dirty_pgds = 0;
 
+	CSTD_UNUSED(as_nr);
+
 	/* Early out if there is nothing to do */
 	if (nr == 0)
 		return 0;
@@ -2605,9 +2583,7 @@ int kbase_mmu_insert_aliased_pages(struct kbase_device *kbdev, struct kbase_mmu_
 	return 0;
 }
 
-void kbase_mmu_update(struct kbase_device *kbdev,
-		struct kbase_mmu_table *mmut,
-		int as_nr)
+void kbase_mmu_update(struct kbase_device *kbdev, struct kbase_mmu_table *mmut, int as_nr)
 {
 	lockdep_assert_held(&kbdev->hwaccess_lock);
 	lockdep_assert_held(&kbdev->mmu_hw_mutex);
@@ -2620,7 +2596,9 @@ KBASE_EXPORT_TEST_API(kbase_mmu_update);
 void kbase_mmu_disable_as(struct kbase_device *kbdev, int as_nr)
 {
 	lockdep_assert_held(&kbdev->hwaccess_lock);
+#if !MALI_USE_CSF
 	lockdep_assert_held(&kbdev->mmu_hw_mutex);
+#endif
 
 	kbdev->mmu_mode->disable_as(kbdev, as_nr);
 }
@@ -2643,10 +2621,9 @@ void kbase_mmu_disable(struct kbase_context *kctx)
 	KBASE_DEBUG_ASSERT(kctx->as_nr != KBASEP_AS_NR_INVALID);
 
 	lockdep_assert_held(&kctx->kbdev->hwaccess_lock);
-	lockdep_assert_held(&kctx->kbdev->mmu_hw_mutex);
 
 	op_param.vpfn = 0;
-	op_param.nr = ~0;
+	op_param.nr = ~0U;
 	op_param.op = KBASE_MMU_OP_FLUSH_MEM;
 	op_param.kctx_id = kctx->id;
 	op_param.mmu_sync_info = mmu_sync_info;
@@ -2655,7 +2632,6 @@ void kbase_mmu_disable(struct kbase_context *kctx)
 	/* 0xF value used to prevent skipping of any levels when flushing */
 	if (mmu_flush_cache_on_gpu_ctrl(kbdev))
 		op_param.flush_skip_levels = pgd_level_to_skip_flush(0xF);
-#endif
 
 	/* lock MMU to prevent existing jobs on GPU from executing while the AS is
 	 * not yet disabled
@@ -2687,8 +2663,28 @@ void kbase_mmu_disable(struct kbase_context *kctx)
 			dev_err(kbdev->dev, "Failed to unlock AS %d for ctx %d_%d", kctx->as_nr,
 				kctx->tgid, kctx->id);
 	}
+#else
+	lockdep_assert_held(&kctx->kbdev->mmu_hw_mutex);
+
+	CSTD_UNUSED(lock_err);
+
+	/*
+	 * The address space is being disabled, drain all knowledge of it out
+	 * from the caches as pages and page tables might be freed after this.
+	 *
+	 * The job scheduler code will already be holding the locks and context
+	 * so just do the flush.
+	 */
+	flush_err = kbase_mmu_hw_do_flush(kbdev, &kbdev->as[kctx->as_nr], &op_param);
+	if (flush_err) {
+		dev_err(kbdev->dev,
+			"Flush for GPU page table update did not complete to disable AS %d for ctx %d_%d",
+			kctx->as_nr, kctx->tgid, kctx->id);
+		/* GPU reset would have been triggered by the flush function */
+	}
+
+	kbdev->mmu_mode->disable_as(kbdev, kctx->as_nr);
 
-#if !MALI_USE_CSF
 	/*
 	 * JM GPUs has some L1 read only caches that need to be invalidated
 	 * with START_FLUSH configuration. Purge the MMU disabled kctx from
@@ -2709,23 +2705,21 @@ static void kbase_mmu_update_and_free_parent_pgds(struct kbase_device *kbdev,
 
 	lockdep_assert_held(&mmut->mmu_lock);
 
-	for (current_level = level - 1; current_level >= MIDGARD_MMU_LEVEL(0);
-	     current_level--) {
+	for (current_level = level - 1; current_level >= MIDGARD_MMU_LEVEL(0); current_level--) {
 		phys_addr_t current_pgd = pgds[current_level];
 		struct page *p = phys_to_page(current_pgd);
 
 		u64 *current_page = kbase_kmap(p);
 		unsigned int current_valid_entries =
 			kbdev->mmu_mode->get_num_valid_entries(current_page);
-		int index = (vpfn >> ((3 - current_level) * 9)) & 0x1FF;
+		unsigned int index = (vpfn >> ((3 - current_level) * 9)) & 0x1FFU;
 
 		/* We need to track every level that needs updating */
 		if (dirty_pgds)
 			*dirty_pgds |= 1ULL << current_level;
 
 		kbdev->mmu_mode->entries_invalidate(&current_page[index], 1);
-		if (current_valid_entries == 1 &&
-		    current_level != MIDGARD_MMU_LEVEL(0)) {
+		if (current_valid_entries == 1 && current_level != MIDGARD_MMU_LEVEL(0)) {
 			kbase_kunmap(p, current_page);
 
 			/* Ensure the cacheline containing the last valid entry
@@ -2733,15 +2727,14 @@ static void kbase_mmu_update_and_free_parent_pgds(struct kbase_device *kbdev,
 			 * PGD page is freed.
 			 */
 			kbase_mmu_sync_pgd_gpu(kbdev, mmut->kctx,
-				current_pgd + (index * sizeof(u64)),
-				sizeof(u64), flush_op);
+					       current_pgd + (index * sizeof(u64)), sizeof(u64),
+					       flush_op);
 
 			kbase_mmu_add_to_free_pgds_list(mmut, p);
 		} else {
 			current_valid_entries--;
 
-			kbdev->mmu_mode->set_num_valid_entries(
-				current_page, current_valid_entries);
+			kbdev->mmu_mode->set_num_valid_entries(current_page, current_valid_entries);
 
 			kbase_kunmap(p, current_page);
 
@@ -2807,6 +2800,9 @@ static void mmu_flush_invalidate_teardown_pages(struct kbase_device *kbdev,
 			spin_unlock_irqrestore(&kbdev->hwaccess_lock, irq_flags);
 		}
 	}
+#else
+	CSTD_UNUSED(phys);
+	CSTD_UNUSED(phys_page_nr);
 #endif
 }
 
@@ -2817,8 +2813,13 @@ static int kbase_mmu_teardown_pgd_pages(struct kbase_device *kbdev, struct kbase
 {
 	struct kbase_mmu_mode const *mmu_mode = kbdev->mmu_mode;
 
+	CSTD_UNUSED(free_pgds_list);
+
 	lockdep_assert_held(&mmut->mmu_lock);
 	kbase_mmu_reset_free_pgds_list(mmut);
+	/* Convert to GPU_PAGE_SIZE units. */
+	vpfn *= GPU_PAGES_PER_CPU_PAGE;
+	nr *= GPU_PAGES_PER_CPU_PAGE;
 
 	while (nr) {
 		unsigned int index = vpfn & 0x1FF;
@@ -2834,9 +2835,8 @@ static int kbase_mmu_teardown_pgd_pages(struct kbase_device *kbdev, struct kbase
 		if (count > nr)
 			count = nr;
 
-		/* need to check if this is a 2MB page or a 4kB */
-		for (level = MIDGARD_MMU_TOPLEVEL;
-				level <= MIDGARD_MMU_BOTTOMLEVEL; level++) {
+		/* need to check if this is a 2MB page or a small page */
+		for (level = MIDGARD_MMU_TOPLEVEL; level <= MIDGARD_MMU_BOTTOMLEVEL; level++) {
 			phys_addr_t next_pgd;
 
 			index = (vpfn >> ((3 - level) * 9)) & 0x1FF;
@@ -2921,9 +2921,8 @@ static int kbase_mmu_teardown_pgd_pages(struct kbase_device *kbdev, struct kbase
 			 * of PGD is invalidated from the GPU cache, before the
 			 * PGD page is freed.
 			 */
-			kbase_mmu_sync_pgd_gpu(kbdev, mmut->kctx,
-				pgd + (index * sizeof(u64)),
-				pcount * sizeof(u64), flush_op);
+			kbase_mmu_sync_pgd_gpu(kbdev, mmut->kctx, pgd + (index * sizeof(u64)),
+					       pcount * sizeof(u64), flush_op);
 
 			kbase_mmu_add_to_free_pgds_list(mmut, p);
 
@@ -2941,9 +2940,9 @@ static int kbase_mmu_teardown_pgd_pages(struct kbase_device *kbdev, struct kbase
 				   kbase_dma_addr(p) + (index * sizeof(u64)), pcount * sizeof(u64),
 				   flush_op);
 next:
-	kbase_kunmap(p, page);
-	vpfn += count;
-	nr -= count;
+		kbase_kunmap(p, page);
+		vpfn += count;
+		nr -= count;
 	}
 out:
 	return 0;
@@ -2954,12 +2953,12 @@ out:
  *
  * @kbdev:    Pointer to kbase device.
  * @mmut:     Pointer to GPU MMU page table.
- * @vpfn:     Start page frame number of the GPU virtual pages to unmap.
+ * @vpfn:     Start page frame number (in PAGE_SIZE units) of the GPU virtual pages to unmap.
  * @phys:     Array of physical pages currently mapped to the virtual
  *            pages to unmap, or NULL. This is used for GPU cache maintenance
  *            and page migration support.
- * @nr_phys_pages: Number of physical pages to flush.
- * @nr_virt_pages: Number of virtual pages whose PTEs should be destroyed.
+ * @nr_phys_pages: Number of physical pages (in PAGE_SIZE units) to flush.
+ * @nr_virt_pages: Number of virtual pages (in PAGE_SIZE units) whose PTEs should be destroyed.
  * @as_nr:    Address space number, for GPU cache maintenance operations
  *            that happen outside a specific kbase context.
  * @ignore_page_migration: Whether page migration metadata should be ignored.
@@ -3042,7 +3041,7 @@ static int mmu_teardown_pages(struct kbase_device *kbdev, struct kbase_mmu_table
 		.mmu_sync_info = mmu_sync_info,
 		.kctx_id = mmut->kctx ? mmut->kctx->id : 0xFFFFFFFF,
 		.op = (flush_op == KBASE_MMU_OP_FLUSH_PT) ? KBASE_MMU_OP_FLUSH_PT :
-							    KBASE_MMU_OP_FLUSH_MEM,
+								  KBASE_MMU_OP_FLUSH_MEM,
 		.flush_skip_levels = pgd_level_to_skip_flush(dirty_pgds),
 	};
 	mmu_flush_invalidate_teardown_pages(kbdev, mmut->kctx, as_nr, phys, nr_phys_pages,
@@ -3085,12 +3084,12 @@ int kbase_mmu_teardown_imported_pages(struct kbase_device *kbdev, struct kbase_m
  *
  * @kbdev: Pointer to kbase device.
  * @mmut:  The involved MMU table
- * @vpfn:  Virtual PFN (Page Frame Number) of the first page to update
+ * @vpfn:  Virtual PFN (Page Frame Number), in PAGE_SIZE units, of the first page to update
  * @phys:  Pointer to the array of tagged physical addresses of the physical
  *         pages that are pointed to by the page table entries (that need to
  *         be updated). The pointer should be within the reg->gpu_alloc->pages
  *         array.
- * @nr:    Number of pages to update
+ * @nr:    Number of pages (in PAGE_SIZE units) to update
  * @flags: Flags
  * @group_id: The physical memory group in which the page was allocated.
  *            Valid range is 0..(MEMORY_GROUP_MANAGER_NR_GROUPS-1).
@@ -3120,6 +3119,9 @@ static int kbase_mmu_update_pages_no_flush(struct kbase_device *kbdev, struct kb
 	if (nr == 0)
 		return 0;
 
+	/* Convert to GPU_PAGE_SIZE units. */
+	vpfn *= GPU_PAGES_PER_CPU_PAGE;
+	nr *= GPU_PAGES_PER_CPU_PAGE;
 	mutex_lock(&mmut->mmu_lock);
 
 	while (nr) {
@@ -3133,7 +3135,8 @@ static int kbase_mmu_update_pages_no_flush(struct kbase_device *kbdev, struct kb
 		if (count > nr)
 			count = nr;
 
-		if (is_huge(*phys) && (index == index_in_large_page(*phys)))
+		if (is_huge(*phys) &&
+		    (index == (index_in_large_page(*phys) * GPU_PAGES_PER_CPU_PAGE)))
 			cur_level = MIDGARD_MMU_LEVEL(2);
 
 		err = mmu_get_pgd_at_level(kbdev, mmut, vpfn, cur_level, &pgd);
@@ -3148,34 +3151,38 @@ static int kbase_mmu_update_pages_no_flush(struct kbase_device *kbdev, struct kb
 			goto fail_unlock;
 		}
 
-		num_of_valid_entries =
-			kbdev->mmu_mode->get_num_valid_entries(pgd_page);
+		num_of_valid_entries = kbdev->mmu_mode->get_num_valid_entries(pgd_page);
 
 		if (cur_level == MIDGARD_MMU_LEVEL(2)) {
-			int level_index = (vpfn >> 9) & 0x1FF;
-			struct tagged_addr *target_phys =
-				phys - index_in_large_page(*phys);
+			unsigned int level_index = (vpfn >> 9) & 0x1FFU;
+			struct tagged_addr *target_phys = phys - index_in_large_page(*phys);
 
 #ifdef CONFIG_MALI_DEBUG
-			WARN_ON_ONCE(!kbdev->mmu_mode->ate_is_valid(
-					pgd_page[level_index], MIDGARD_MMU_LEVEL(2)));
+			WARN_ON_ONCE(!kbdev->mmu_mode->ate_is_valid(pgd_page[level_index],
+								    MIDGARD_MMU_LEVEL(2)));
 #endif
-			pgd_page[level_index] = kbase_mmu_create_ate(kbdev,
-					*target_phys, flags, MIDGARD_MMU_LEVEL(2),
-					group_id);
+			pgd_page[level_index] = kbase_mmu_create_ate(
+				kbdev, *target_phys, flags, MIDGARD_MMU_LEVEL(2), group_id);
 			kbase_mmu_sync_pgd(kbdev, mmut->kctx, pgd + (level_index * sizeof(u64)),
 					   kbase_dma_addr(p) + (level_index * sizeof(u64)),
 					   sizeof(u64), KBASE_MMU_OP_NONE);
 		} else {
-			for (i = 0; i < count; i++) {
+			for (i = 0; i < count; i += GPU_PAGES_PER_CPU_PAGE) {
+				phys_addr_t base_phys_address =
+					as_phys_addr_t(phys[i / GPU_PAGES_PER_CPU_PAGE]);
+				unsigned int j;
+
+				for (j = 0; j < GPU_PAGES_PER_CPU_PAGE; j++) {
+					phys_addr_t page_address =
+						base_phys_address + (j * GPU_PAGE_SIZE);
 #ifdef CONFIG_MALI_DEBUG
-				WARN_ON_ONCE(!kbdev->mmu_mode->ate_is_valid(
-						pgd_page[index + i],
-						MIDGARD_MMU_BOTTOMLEVEL));
+					WARN_ON_ONCE(!kbdev->mmu_mode->ate_is_valid(
+						pgd_page[index + i + j], MIDGARD_MMU_BOTTOMLEVEL));
 #endif
-				pgd_page[index + i] = kbase_mmu_create_ate(kbdev,
-					phys[i], flags, MIDGARD_MMU_BOTTOMLEVEL,
-					group_id);
+					pgd_page[index + i + j] = kbase_mmu_create_ate(
+						kbdev, as_tagged(page_address), flags,
+						MIDGARD_MMU_BOTTOMLEVEL, group_id);
+				}
 			}
 
 			/* MMU cache flush strategy is NONE because GPU cache maintenance
@@ -3186,13 +3193,12 @@ static int kbase_mmu_update_pages_no_flush(struct kbase_device *kbdev, struct kb
 					   count * sizeof(u64), KBASE_MMU_OP_NONE);
 		}
 
-		kbdev->mmu_mode->set_num_valid_entries(pgd_page,
-					num_of_valid_entries);
+		kbdev->mmu_mode->set_num_valid_entries(pgd_page, num_of_valid_entries);
 
 		if (dirty_pgds && count > 0)
 			*dirty_pgds |= 1ULL << cur_level;
 
-		phys += count;
+		phys += (count / GPU_PAGES_PER_CPU_PAGE);
 		vpfn += count;
 		nr -= count;
 
@@ -3299,15 +3305,17 @@ int kbase_mmu_migrate_page(struct tagged_addr old_phys, struct tagged_addr new_p
 	struct kbase_page_metadata *page_md = kbase_page_private(as_page(old_phys));
 	struct kbase_mmu_hw_op_param op_param;
 	struct kbase_mmu_table *mmut = (level == MIDGARD_MMU_BOTTOMLEVEL) ?
-					       page_md->data.mapped.mmut :
-					       page_md->data.pt_mapped.mmut;
+						     page_md->data.mapped.mmut :
+						     page_md->data.pt_mapped.mmut;
 	struct kbase_device *kbdev;
 	phys_addr_t pgd;
 	u64 *old_page, *new_page, *pgd_page, *target, vpfn;
-	int index, check_state, ret = 0;
+	unsigned int index;
+	int check_state, ret = 0;
 	unsigned long hwaccess_flags = 0;
 	unsigned int num_of_valid_entries;
 	u8 vmap_count = 0;
+	u8 pgd_entries_to_sync = (level == MIDGARD_MMU_BOTTOMLEVEL) ? GPU_PAGES_PER_CPU_PAGE : 1;
 
 	/* If page migration support is not compiled in, return with fault */
 	if (!IS_ENABLED(CONFIG_PAGE_MIGRATION_SUPPORT))
@@ -3326,7 +3334,7 @@ int kbase_mmu_migrate_page(struct tagged_addr old_phys, struct tagged_addr new_p
 		vpfn = PGD_VPFN_LEVEL_GET_VPFN(page_md->data.pt_mapped.pgd_vpfn_level);
 
 	kbdev = mmut->kctx->kbdev;
-	index = (vpfn >> ((3 - level) * 9)) & 0x1FF;
+	index = (vpfn >> ((3 - level) * 9)) & 0x1FFU;
 
 	/* Create all mappings before copying content.
 	 * This is done as early as possible because it is the only operation that may
@@ -3376,13 +3384,13 @@ int kbase_mmu_migrate_page(struct tagged_addr old_phys, struct tagged_addr new_p
 #define PGD_VPFN_MASK(level) (~((((u64)1) << ((3 - level) * 9)) - 1))
 	op_param.mmu_sync_info = CALLER_MMU_ASYNC;
 	op_param.kctx_id = mmut->kctx->id;
-	op_param.vpfn = vpfn & PGD_VPFN_MASK(level);
-	op_param.nr = 1 << ((3 - level) * 9);
+	op_param.vpfn = (vpfn / GPU_PAGES_PER_CPU_PAGE) & PGD_VPFN_MASK(level);
+	op_param.nr = 1U << ((3 - level) * 9);
 	op_param.op = KBASE_MMU_OP_FLUSH_PT;
 	/* When level is not MIDGARD_MMU_BOTTOMLEVEL, it is assumed PGD page migration */
 	op_param.flush_skip_levels = (level == MIDGARD_MMU_BOTTOMLEVEL) ?
-					     pgd_level_to_skip_flush(1ULL << level) :
-					     pgd_level_to_skip_flush(3ULL << level);
+						   pgd_level_to_skip_flush(1ULL << level) :
+						   pgd_level_to_skip_flush(3ULL << level);
 
 	mutex_lock(&mmut->mmu_lock);
 
@@ -3459,6 +3467,13 @@ int kbase_mmu_migrate_page(struct tagged_addr old_phys, struct tagged_addr new_p
 
 		ret = kbase_mmu_hw_do_lock(kbdev, as, &op_param);
 		if (!ret) {
+#if MALI_USE_CSF
+			if (mmu_flush_cache_on_gpu_ctrl(kbdev))
+				ret = kbase_gpu_cache_flush_pa_range_and_busy_wait(
+					kbdev, as_phys_addr_t(old_phys), PAGE_SIZE,
+					GPU_COMMAND_FLUSH_PA_RANGE_CLN_INV_L2_LSC);
+			else
+#endif
 				ret = kbase_gpu_cache_flush_and_busy_wait(
 					kbdev, GPU_COMMAND_CACHE_CLN_INV_L2_LSC);
 		}
@@ -3489,7 +3504,7 @@ int kbase_mmu_migrate_page(struct tagged_addr old_phys, struct tagged_addr new_p
 	/* Remap GPU virtual page.
 	 *
 	 * This code rests on the assumption that page migration is only enabled
-	 * for 4 kB pages, that necessarily live in the bottom level of the MMU
+	 * for small pages, that necessarily live in the bottom level of the MMU
 	 * page table. For this reason, the PGD level tells us inequivocably
 	 * whether the page being migrated is a "content page" or another PGD
 	 * of the page table:
@@ -3509,17 +3524,26 @@ int kbase_mmu_migrate_page(struct tagged_addr old_phys, struct tagged_addr new_p
 	num_of_valid_entries = kbdev->mmu_mode->get_num_valid_entries(pgd_page);
 
 	if (level == MIDGARD_MMU_BOTTOMLEVEL) {
-		WARN_ON_ONCE((*target & 1UL) == 0);
-		*target =
-			kbase_mmu_create_ate(kbdev, new_phys, page_md->data.mapped.reg->flags,
-					     level, page_md->data.mapped.reg->gpu_alloc->group_id);
+		phys_addr_t base_phys_address = as_phys_addr_t(new_phys);
+		unsigned int i;
+
+		for (i = 0; i < GPU_PAGES_PER_CPU_PAGE; i++) {
+			phys_addr_t page_address = base_phys_address + (i * GPU_PAGE_SIZE);
+
+			WARN_ON_ONCE((*target & 1UL) == 0);
+			*target = kbase_mmu_create_ate(
+				kbdev, as_tagged(page_address), page_md->data.mapped.reg->flags,
+				level, page_md->data.mapped.reg->gpu_alloc->group_id);
+			target++;
+		}
 	} else {
 		u64 managed_pte;
 
 #ifdef CONFIG_MALI_DEBUG
 		/* The PTE should be pointing to the page being migrated */
-		WARN_ON_ONCE(as_phys_addr_t(old_phys) != kbdev->mmu_mode->pte_to_phy_addr(
-			kbdev->mgm_dev->ops.mgm_pte_to_original_pte(
+		WARN_ON_ONCE(
+			as_phys_addr_t(old_phys) !=
+			kbdev->mmu_mode->pte_to_phy_addr(kbdev->mgm_dev->ops.mgm_pte_to_original_pte(
 				kbdev->mgm_dev, MGM_DEFAULT_PTE_GROUP, level, pgd_page[index])));
 #endif
 		kbdev->mmu_mode->entry_set_pte(&managed_pte, as_phys_addr_t(new_phys));
@@ -3529,12 +3553,14 @@ int kbase_mmu_migrate_page(struct tagged_addr old_phys, struct tagged_addr new_p
 
 	kbdev->mmu_mode->set_num_valid_entries(pgd_page, num_of_valid_entries);
 
-	/* This function always updates a single entry inside an existing PGD,
-	 * therefore cache maintenance is necessary and affects a single entry.
+	/* This function always updates a single entry inside an existing PGD when
+	 * level != MIDGARD_MMU_BOTTOMLEVEL, and would update more than one entry for
+	 * MIDGARD_MMU_BOTTOMLEVEL PGD when PAGE_SIZE is not 4K, therefore cache
+	 * maintenance is necessary.
 	 */
 	kbase_mmu_sync_pgd(kbdev, mmut->kctx, pgd + (index * sizeof(u64)),
-			   kbase_dma_addr(phys_to_page(pgd)) + (index * sizeof(u64)), sizeof(u64),
-			   KBASE_MMU_OP_FLUSH_PT);
+			   kbase_dma_addr(phys_to_page(pgd)) + (index * sizeof(u64)),
+			   pgd_entries_to_sync * sizeof(u64), KBASE_MMU_OP_FLUSH_PT);
 
 	/* Unlock MMU region.
 	 *
@@ -3575,8 +3601,8 @@ int kbase_mmu_migrate_page(struct tagged_addr old_phys, struct tagged_addr new_p
 	/* Undertaking metadata transfer, while we are holding the mmu_lock */
 	spin_lock(&page_md->migrate_lock);
 	if (level == MIDGARD_MMU_BOTTOMLEVEL) {
-		size_t page_array_index =
-			page_md->data.mapped.vpfn - page_md->data.mapped.reg->start_pfn;
+		size_t page_array_index = (page_md->data.mapped.vpfn / GPU_PAGES_PER_CPU_PAGE) -
+					  page_md->data.mapped.reg->start_pfn;
 
 		WARN_ON(PAGE_STATUS_GET(page_md->status) != ALLOCATED_MAPPED);
 
@@ -3615,7 +3641,7 @@ undo_mappings:
 }
 
 static void mmu_teardown_level(struct kbase_device *kbdev, struct kbase_mmu_table *mmut,
-			       phys_addr_t pgd, unsigned int level)
+			       phys_addr_t pgd, int level)
 {
 	u64 *pgd_page;
 	int i;
@@ -3683,12 +3709,10 @@ static void kbase_mmu_mark_non_movable(struct page *page)
 	spin_unlock(&page_md->migrate_lock);
 }
 
-int kbase_mmu_init(struct kbase_device *const kbdev,
-	struct kbase_mmu_table *const mmut, struct kbase_context *const kctx,
-	int const group_id)
+int kbase_mmu_init(struct kbase_device *const kbdev, struct kbase_mmu_table *const mmut,
+		   struct kbase_context *const kctx, int const group_id)
 {
-	if (WARN_ON(group_id >= MEMORY_GROUP_MANAGER_NR_GROUPS) ||
-	    WARN_ON(group_id < 0))
+	if (WARN_ON(group_id >= MEMORY_GROUP_MANAGER_NR_GROUPS) || WARN_ON(group_id < 0))
 		return -EINVAL;
 
 	compiletime_assert(KBASE_MEM_ALLOC_MAX_SIZE <= (((8ull << 30) >> PAGE_SHIFT)),
@@ -3699,18 +3723,17 @@ int kbase_mmu_init(struct kbase_device *const kbdev,
 	mmut->group_id = group_id;
 	mutex_init(&mmut->mmu_lock);
 	mmut->kctx = kctx;
-	mmut->pgd = KBASE_MMU_INVALID_PGD_ADDRESS;
+	mmut->pgd = KBASE_INVALID_PHYSICAL_ADDRESS;
 
 	/* We allocate pages into the kbdev memory pool, then
 	 * kbase_mmu_alloc_pgd will allocate out of that pool. This is done to
 	 * avoid allocations from the kernel happening with the lock held.
 	 */
-	while (mmut->pgd == KBASE_MMU_INVALID_PGD_ADDRESS) {
+	while (mmut->pgd == KBASE_INVALID_PHYSICAL_ADDRESS) {
 		int err;
 
-		err = kbase_mem_pool_grow(
-			&kbdev->mem_pools.small[mmut->group_id],
-			MIDGARD_MMU_BOTTOMLEVEL, kctx ? kctx->task : NULL);
+		err = kbase_mem_pool_grow(&kbdev->mem_pools.small[mmut->group_id],
+					  MIDGARD_MMU_BOTTOMLEVEL, kctx ? kctx->task : NULL);
 		if (err) {
 			kbase_mmu_term(kbdev, mmut);
 			return -ENOMEM;
@@ -3729,7 +3752,7 @@ void kbase_mmu_term(struct kbase_device *kbdev, struct kbase_mmu_table *mmut)
 	     "kctx-%d_%d must first be scheduled out to flush GPU caches+tlbs before tearing down MMU tables",
 	     mmut->kctx->tgid, mmut->kctx->id);
 
-	if (mmut->pgd != KBASE_MMU_INVALID_PGD_ADDRESS) {
+	if (mmut->pgd != KBASE_INVALID_PHYSICAL_ADDRESS) {
 		mutex_lock(&mmut->mmu_lock);
 		mmu_teardown_level(kbdev, mmut, mmut->pgd, MIDGARD_MMU_TOPLEVEL);
 		mutex_unlock(&mmut->mmu_lock);
@@ -3747,8 +3770,7 @@ void kbase_mmu_as_term(struct kbase_device *kbdev, unsigned int i)
 }
 
 void kbase_mmu_flush_pa_range(struct kbase_device *kbdev, struct kbase_context *kctx,
-			      phys_addr_t phys, size_t size,
-			      enum kbase_mmu_op_type flush_op)
+			      phys_addr_t phys, size_t size, enum kbase_mmu_op_type flush_op)
 {
 #if MALI_USE_CSF
 	unsigned long irq_flags;
@@ -3758,12 +3780,18 @@ void kbase_mmu_flush_pa_range(struct kbase_device *kbdev, struct kbase_context *
 	    kbdev->pm.backend.gpu_ready && (!kctx || kctx->as_nr >= 0))
 		mmu_flush_pa_range(kbdev, phys, size, KBASE_MMU_OP_FLUSH_PT);
 	spin_unlock_irqrestore(&kbdev->hwaccess_lock, irq_flags);
+#else
+	CSTD_UNUSED(kbdev);
+	CSTD_UNUSED(kctx);
+	CSTD_UNUSED(phys);
+	CSTD_UNUSED(size);
+	CSTD_UNUSED(flush_op);
 #endif
 }
 
 #ifdef CONFIG_MALI_VECTOR_DUMP
-static size_t kbasep_mmu_dump_level(struct kbase_context *kctx, phys_addr_t pgd,
-		int level, char ** const buffer, size_t *size_left)
+static size_t kbasep_mmu_dump_level(struct kbase_context *kctx, phys_addr_t pgd, int level,
+				    char **const buffer, size_t *size_left)
 {
 	phys_addr_t target_pgd;
 	u64 *pgd_page;
@@ -3790,7 +3818,7 @@ static size_t kbasep_mmu_dump_level(struct kbase_context *kctx, phys_addr_t pgd,
 		/* A modified physical address that contains
 		 * the page table level
 		 */
-		u64 m_pgd = pgd | level;
+		u64 m_pgd = pgd | (u64)level;
 
 		/* Put the modified physical address in the output buffer */
 		memcpy(*buffer, &m_pgd, sizeof(m_pgd));
@@ -3808,12 +3836,11 @@ static size_t kbasep_mmu_dump_level(struct kbase_context *kctx, phys_addr_t pgd,
 			if (mmu_mode->pte_is_valid(pgd_page[i], level)) {
 				target_pgd = mmu_mode->pte_to_phy_addr(
 					kbdev->mgm_dev->ops.mgm_pte_to_original_pte(
-						kbdev->mgm_dev, MGM_DEFAULT_PTE_GROUP,
-						level, pgd_page[i]));
+						kbdev->mgm_dev, MGM_DEFAULT_PTE_GROUP, level,
+						pgd_page[i]));
 
-				dump_size = kbasep_mmu_dump_level(kctx,
-						target_pgd, level + 1,
-						buffer, size_left);
+				dump_size = kbasep_mmu_dump_level(kctx, target_pgd, level + 1,
+								  buffer, size_left);
 				if (!dump_size) {
 					kbase_kunmap(pfn_to_page(PFN_DOWN(pgd)), pgd_page);
 					return 0;
@@ -3828,7 +3855,7 @@ static size_t kbasep_mmu_dump_level(struct kbase_context *kctx, phys_addr_t pgd,
 	return size;
 }
 
-void *kbase_mmu_dump(struct kbase_context *kctx, int nr_pages)
+void *kbase_mmu_dump(struct kbase_context *kctx, size_t nr_pages)
 {
 	void *kaddr;
 	size_t size_left;
@@ -3859,8 +3886,7 @@ void *kbase_mmu_dump(struct kbase_context *kctx, int nr_pages)
 		buffer = (char *)kaddr;
 		mmu_dump_buffer = buffer;
 
-		kctx->kbdev->mmu_mode->get_as_setup(&kctx->mmu,
-				&as_setup);
+		kctx->kbdev->mmu_mode->get_as_setup(&kctx->mmu, &as_setup);
 		config[0] = as_setup.transtab;
 		config[1] = as_setup.memattr;
 		config[2] = as_setup.transcfg;
@@ -3869,11 +3895,8 @@ void *kbase_mmu_dump(struct kbase_context *kctx, int nr_pages)
 		size_left -= sizeof(config);
 		size += sizeof(config);
 
-		dump_size = kbasep_mmu_dump_level(kctx,
-				kctx->mmu.pgd,
-				MIDGARD_MMU_TOPLEVEL,
-				&mmu_dump_buffer,
-				&size_left);
+		dump_size = kbasep_mmu_dump_level(kctx, kctx->mmu.pgd, MIDGARD_MMU_TOPLEVEL,
+						  &mmu_dump_buffer, &size_left);
 
 		if (!dump_size)
 			goto fail_free;
@@ -3908,7 +3931,7 @@ KBASE_EXPORT_TEST_API(kbase_mmu_dump);
 void kbase_mmu_bus_fault_worker(struct work_struct *data)
 {
 	struct kbase_as *faulting_as;
-	int as_no;
+	unsigned int as_no;
 	struct kbase_context *kctx;
 	struct kbase_device *kbdev;
 	struct kbase_fault *fault;
@@ -3944,14 +3967,11 @@ void kbase_mmu_bus_fault_worker(struct work_struct *data)
 #endif
 
 	if (unlikely(fault->protected_mode)) {
-		kbase_mmu_report_fault_and_kill(kctx, faulting_as,
-				"Permission failure", fault);
-		kbase_mmu_hw_clear_fault(kbdev, faulting_as,
-				KBASE_MMU_FAULT_TYPE_BUS_UNEXPECTED);
+		kbase_mmu_report_fault_and_kill(kctx, faulting_as, "Permission failure", fault);
+		kbase_mmu_hw_clear_fault(kbdev, faulting_as, KBASE_MMU_FAULT_TYPE_BUS_UNEXPECTED);
 		release_ctx(kbdev, kctx);
 		atomic_dec(&kbdev->faults_pending);
 		return;
-
 	}
 
 #if MALI_USE_CSF
@@ -3965,7 +3985,7 @@ void kbase_mmu_bus_fault_worker(struct work_struct *data)
 	 * we don't need to switch to unmapped
 	 */
 	if (!kbase_pm_context_active_handle_suspend(kbdev,
-				KBASE_PM_SUSPEND_HANDLER_DONT_REACTIVATE)) {
+						    KBASE_PM_SUSPEND_HANDLER_DONT_REACTIVATE)) {
 		kbase_gpu_report_bus_fault_and_kill(kctx, faulting_as, fault);
 		kbase_pm_context_idle(kbdev);
 	}