[PATCH v3 4/6] arm: snapdragon: Disable MMU early before U-Boot reset vector

[Balaji Selvanathan]

When U-Boot is loaded by XBL in snagboot mode, XBL leaves the MMU
enabled with its own page tables. This causes stale TLB entries and
incorrect memory mappings when U-Boot initializes its own MMU.

Disable the MMU and invalidate TLBs at the earliest entry point,
before branching to the reset vector, handling EL1, EL2, and EL3.
This ensures a clean MMU state for U-Boot initialization.

The changes in this patch are based on:
https://lore.kernel.org/all/20260508-qcom_spl-v7-4-7d0e22aaaa8f@seznam.cz/
https://lore.kernel.org/u-boot/20260515130934.388894-1-aswin.murugan@oss.qualcomm.com/

Signed-off-by: Balaji Selvanathan <balaji.selvanathan at oss.qualcomm.com>
---
Changes in v3:
- Brought the MMU disable logic under CONFIG_QCOM_SNAGBOOT_SUPPORT
  ifdef condition
- Moved the MMU disable codes to a new seperate file

Changes in v2:
- Newly introduced in v2
---
 arch/arm/mach-snapdragon/include/mach/boot0.h      |  2 +
 .../mach-snapdragon/include/mach/snagboot_boot0.h  | 53 ++++++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/arch/arm/mach-snapdragon/include/mach/boot0.h b/arch/arm/mach-snapdragon/include/mach/boot0.h
index 032a9fbdb12..60ad15dba57 100644
--- a/arch/arm/mach-snapdragon/include/mach/boot0.h
+++ b/arch/arm/mach-snapdragon/include/mach/boot0.h
@@ -7,6 +7,8 @@
 #include "msm8916_boot0.h"
 #elif defined(CONFIG_QCOM_EL2_GUNYAH_EXIT_SUPPORT)
 #include "gunyah_exit_boot0.h"
+#elif defined(CONFIG_QCOM_SNAGBOOT_SUPPORT)
+#include "snagboot_boot0.h"
 #else
 	b	reset
 #endif
diff --git a/arch/arm/mach-snapdragon/include/mach/snagboot_boot0.h b/arch/arm/mach-snapdragon/include/mach/snagboot_boot0.h
new file mode 100644
index 00000000000..fb146384cba
--- /dev/null
+++ b/arch/arm/mach-snapdragon/include/mach/snagboot_boot0.h
@@ -0,0 +1,53 @@
+/* SPDX-License-Identifier: GPL-2.0+ */
+/*
+ * Early MMU Disable for Snagboot Mode
+ *
+ * Disable MMU at the earliest possible point. This must be done before
+ * any other U-Boot code executes. Handle all exception levels (EL1, EL2, EL3).
+ *
+ * After disabling MMU, invalidate TLB to clear any stale entries that might
+ * cause issues when MMU is re-enabled later.
+ */
+
+	mrs	x0, CurrentEL
+	cmp	x0, #(3 << 2)	/* EL3? */
+	b.eq	disable_mmu_el3
+	cmp	x0, #(2 << 2)	/* EL2? */
+	b.eq	disable_mmu_el2
+
+	/* EL1 */
+	cmp	x0, #(1 << 2)	/* Confirm EL1 */
+	b.ne	reset		/* Skip MMU disable if not EL1 (e.g., EL0) */
+	mrs	x0, sctlr_el1
+	bic	x0, x0, #1	/* Clear M bit (MMU enable) */
+	msr	sctlr_el1, x0
+	isb
+	/* Invalidate TLB for EL1 */
+	tlbi	vmalle1
+	dsb	sy
+	isb
+	b	reset
+
+disable_mmu_el3:
+	mrs	x0, sctlr_el3
+	bic	x0, x0, #1	/* Clear M bit (MMU enable) */
+	bic	x0, x0, #(1 << 19)  /* Clear WXN bit (Write XOR Execute) */
+	msr	sctlr_el3, x0
+	isb
+	/* Invalidate entire TLB for all ELs */
+	tlbi	alle3
+	dsb	sy
+	isb
+	b	reset
+
+disable_mmu_el2:
+	mrs	x0, sctlr_el2
+	bic	x0, x0, #1	/* Clear M bit (MMU enable) */
+	msr	sctlr_el2, x0
+	isb
+	/* Invalidate TLB for EL2 */
+	tlbi	alle2
+	dsb	sy
+	isb
+	b	reset
+

-- 
2.34.1