Save/restore QMI window 1 registers across calls to flash_exit_xip().

This works around the RP2350 ROM reinitialising window 1 even when it
does not actually issue an XIP exit sequence to that QSPI device. If no
exit sequence is issued, then the original configuration still applies.

This is not valid in the case where the ROM *does* issue an XIP exit
sequence to the second chip select, because the original mode is no
longer valid once the device is in the serial command state. The ROM
does this when FLASH_DEVINFO (loaded from OTP, stored in boot RAM) has a
nonzero size for chip select 1. To distinguish the two cases, this patch
adds getters/setters for FLASH_DEVINFO.

If chip select 1 has a nonzero size in FLASH_DEVINFO then the ROM's
register changes for window 1 are not reverted, and instead the
hardware_flash code makes additional changes to ensure writes also work
(as the ROM only reinitialises for reads).

This means that, in the case you have opted into ROM support for the
second chip select by setting a nonzero CS1 size in FLASH_DEVINFO, the
device on chip select 1 will continue to function correctly, but full
performance won't be restored until it is reinitialised. The best way to
handle this is initialising both chip select devices in your XIP setup
function (your boot2).

src/rp2_common/hardware_flash/flash.c

@@ -12,6 +12,7 @@
 #include "hardware/structs/ssi.h"
 #else
 #include "hardware/structs/qmi.h"
+#include "hardware/regs/otp_data.h"
 #endif
 #include "hardware/xip_cache.h"
 
@@ -70,6 +71,43 @@ static void __no_inline_not_in_flash_func(flash_enable_xip_via_boot2)(void) {
 
 #endif
 
+#if PICO_RP2350
+// This is specifically for saving/restoring the registers modified by RP2350
+// flash_exit_xip() ROM func, not the entirety of the QMI window state.
+typedef struct flash_rp2350_qmi_save_state {
+    uint32_t timing;
+    uint32_t rcmd;
+    uint32_t rfmt;
+} flash_rp2350_qmi_save_state_t;
+
+static void __no_inline_not_in_flash_func(flash_rp2350_save_qmi_cs1)(flash_rp2350_qmi_save_state_t *state) {
+    state->timing = qmi_hw->m[1].timing;
+    state->rcmd = qmi_hw->m[1].rcmd;
+    state->rfmt = qmi_hw->m[1].rfmt;
+}
+
+static void __no_inline_not_in_flash_func(flash_rp2350_restore_qmi_cs1)(const flash_rp2350_qmi_save_state_t *state) {
+    if (flash_devinfo_get_cs_size(1) == FLASH_DEVINFO_SIZE_NONE) {
+        // Case 1: The RP2350 ROM sets QMI to a clean (03h read) configuration
+        // during flash_exit_xip(), even though when CS1 is not enabled via
+        // FLASH_DEVINFO it does not issue an XIP exit sequence to CS1. In
+        // this case, restore the original register config for CS1 as it is
+        // still the correct config.
+        qmi_hw->m[1].timing = state->timing;
+        qmi_hw->m[1].rcmd = state->rcmd;
+        qmi_hw->m[1].rfmt = state->rfmt;
+    } else {
+        // Case 2: If RAM is attached to CS1, and the ROM has issued an XIP
+        // exit sequence to it, then the ROM re-initialisation of the QMI
+        // registers has actually not gone far enough. The old XIP write mode
+        // is no longer valid when the QSPI RAM is returned to a serial
+        // command state. Restore the default 02h serial write command config.
+        qmi_hw->m[1].wfmt = QMI_M1_WFMT_RESET;
+        qmi_hw->m[1].wcmd = QMI_M1_WCMD_RESET;
+    }
+}
+#endif
+
 //-----------------------------------------------------------------------------
 // Actual flash programming shims (work whether or not PICO_NO_FLASH==1)
 
@@ -87,6 +125,10 @@ void __no_inline_not_in_flash_func(flash_range_erase)(uint32_t flash_offs, size_
     flash_init_boot2_copyout();
     // Commit any pending writes to external RAM, to avoid losing them in the subsequent flush:
     xip_cache_clean_all();
+#if PICO_RP2350
+    flash_rp2350_qmi_save_state_t qmi_save;
+    flash_rp2350_save_qmi_cs1(&qmi_save);
+#endif
 
     // No flash accesses after this point
     __compiler_memory_barrier();
@@ -96,6 +138,9 @@ void __no_inline_not_in_flash_func(flash_range_erase)(uint32_t flash_offs, size_
     flash_range_erase_func(flash_offs, count, FLASH_BLOCK_SIZE, FLASH_BLOCK_ERASE_CMD);
     flash_flush_cache_func(); // Note this is needed to remove CSn IO force as well as cache flushing
     flash_enable_xip_via_boot2();
+#if PICO_RP2350
+    flash_rp2350_restore_qmi_cs1(&qmi_save);
+#endif
 }
 
 void __no_inline_not_in_flash_func(flash_flush_cache)(void) {
@@ -116,6 +161,10 @@ void __no_inline_not_in_flash_func(flash_range_program)(uint32_t flash_offs, con
     assert(connect_internal_flash_func && flash_exit_xip_func && flash_range_program_func && flash_flush_cache_func);
     flash_init_boot2_copyout();
     xip_cache_clean_all();
+#if PICO_RP2350
+    flash_rp2350_qmi_save_state_t qmi_save;
+    flash_rp2350_save_qmi_cs1(&qmi_save);
+#endif
 
     __compiler_memory_barrier();
 
@@ -124,6 +173,9 @@ void __no_inline_not_in_flash_func(flash_range_program)(uint32_t flash_offs, con
     flash_range_program_func(flash_offs, data, count);
     flash_flush_cache_func(); // Note this is needed to remove CSn IO force as well as cache flushing
     flash_enable_xip_via_boot2();
+#if PICO_RP2350
+    flash_rp2350_restore_qmi_cs1(&qmi_save);
+#endif
 }
 
 //-----------------------------------------------------------------------------
@@ -157,6 +209,10 @@ void __no_inline_not_in_flash_func(flash_do_cmd)(const uint8_t *txbuf, uint8_t *
     assert(connect_internal_flash_func && flash_exit_xip_func && flash_flush_cache_func);
     flash_init_boot2_copyout();
     xip_cache_clean_all();
+#if PICO_RP2350
+    flash_rp2350_qmi_save_state_t qmi_save;
+    flash_rp2350_save_qmi_cs1(&qmi_save);
+#endif
 
     __compiler_memory_barrier();
     connect_internal_flash_func();
@@ -204,6 +260,9 @@ void __no_inline_not_in_flash_func(flash_do_cmd)(const uint8_t *txbuf, uint8_t *
 
     flash_flush_cache_func();
     flash_enable_xip_via_boot2();
+#if PICO_RP2350
+    flash_rp2350_restore_qmi_cs1(&qmi_save);
+#endif
 }
 #endif
 
@@ -225,3 +284,84 @@ void flash_get_unique_id(uint8_t *id_out) {
         id_out[i] = rxbuf[i + 1 + FLASH_RUID_DUMMY_BYTES];
 #endif
 }
+
+#if !PICO_RP2040
+// This is a static symbol because the layout of FLASH_DEVINFO is liable to change from device to
+// device, so fields must have getters/setters.
+static io_rw_16 * flash_devinfo_ptr(void) {
+    // Note the lookup returns a pointer to a 32-bit pointer literal in the ROM
+    io_rw_16 **p = (io_rw_16 **) rom_data_lookup(ROM_DATA_FLASH_DEVINFO16_PTR);
+    assert(p);
+    return *p;
+}
+
+static void flash_devinfo_update_field(uint16_t wdata, uint16_t mask) {
+    // Boot RAM does not support exclusives, but does support RWTYPE SET/CLR/XOR (with byte
+    // strobes). Can't use hw_write_masked because it performs a 32-bit write.
+    io_rw_16 *devinfo = flash_devinfo_ptr();
+    *hw_xor_alias(devinfo) = (*devinfo ^ wdata) & mask;
+}
+
+// This is a RAM function because may be called during flash programming to enable save/restore of
+// QMI window 1 registers on RP2350:
+flash_devinfo_size_t __no_inline_not_in_flash_func(flash_devinfo_get_cs_size)(uint cs) {
+    invalid_params_if(HARDWARE_FLASH, cs > 1);
+    io_ro_16 *devinfo = (io_ro_16 *) flash_devinfo_ptr();
+    if (cs == 0u) {
+#ifdef PICO_FLASH_SIZE_BYTES
+        // A flash size explicitly specified for the build (e.g. from the board header) takes
+        // precedence over whatever was found in OTP.
+        return (flash_devinfo_size_t) (
+            __builtin_ctz(PICO_FLASH_SIZE_BYTES / 8192u) + (uint)FLASH_DEVINFO_SIZE_8K
+        );
+#else
+        return (flash_devinfo_size_t) (
+            (*devinfo & OTP_DATA_FLASH_DEVINFO_CS0_SIZE_BITS) >> OTP_DATA_FLASH_DEVINFO_CS0_SIZE_LSB
+        );
+#endif
+    } else {
+        return (flash_devinfo_size_t) (
+            (*devinfo & OTP_DATA_FLASH_DEVINFO_CS1_SIZE_BITS) >> OTP_DATA_FLASH_DEVINFO_CS1_SIZE_LSB
+        );
+    }
+}
+
+void flash_devinfo_set_cs_size(uint cs, flash_devinfo_size_t size) {
+    invalid_params_if(HARDWARE_FLASH, cs > 1);
+    invalid_params_if(HARDWARE_FLASH, (uint)size > (uint)FLASH_DEVINFO_SIZE_MAX);
+    uint cs_shift = cs == 0u ? OTP_DATA_FLASH_DEVINFO_CS0_SIZE_LSB : OTP_DATA_FLASH_DEVINFO_CS1_SIZE_LSB;
+    uint cs_mask = OTP_DATA_FLASH_DEVINFO_CS0_SIZE_BITS >> OTP_DATA_FLASH_DEVINFO_CS0_SIZE_LSB;
+    flash_devinfo_update_field(
+        (uint)size << cs_shift,
+        cs_mask << cs_shift
+    );
+}
+
+bool flash_devinfo_get_d8h_erase_supported(void) {
+    return *flash_devinfo_ptr() & OTP_DATA_FLASH_DEVINFO_D8H_ERASE_SUPPORTED_BITS;
+}
+
+void flash_devinfo_set_d8h_erase_supported(bool supported) {
+    flash_devinfo_update_field(
+        (uint)supported << OTP_DATA_FLASH_DEVINFO_D8H_ERASE_SUPPORTED_LSB,
+        OTP_DATA_FLASH_DEVINFO_D8H_ERASE_SUPPORTED_BITS
+    );
+}
+
+uint flash_devinfo_get_cs_gpio(uint cs) {
+    invalid_params_if(HARDWARE_FLASH, cs != 1);
+    (void)cs;
+    return (*flash_devinfo_ptr() & OTP_DATA_FLASH_DEVINFO_CS1_GPIO_BITS) >> OTP_DATA_FLASH_DEVINFO_CS1_GPIO_LSB;
+}
+
+void flash_devinfo_set_cs_gpio(uint cs, uint gpio) {
+    invalid_params_if(HARDWARE_FLASH, cs != 1);
+    invalid_params_if(HARDWARE_FLASH, gpio >= NUM_BANK0_GPIOS);
+    (void)cs;
+    flash_devinfo_update_field(
+        gpio << OTP_DATA_FLASH_DEVINFO_CS1_GPIO_LSB,
+        OTP_DATA_FLASH_DEVINFO_CS1_GPIO_BITS
+    );
+}
+
+#endif // !PICO_RP2040

src/rp2_common/hardware_flash/include/hardware/flash.h

@@ -14,10 +14,10 @@
  *
  * \brief Low level flash programming and erase API
  *
- * Note these functions are *unsafe* if you are using both cores, and the other 
- * is executing from flash concurrently with the operation. In this could be the 
- * case, you must perform your own synchronisation to make sure that no XIP
- * accesses take place during flash programming. One option is to use the
+ * Note these functions are *unsafe* if you are using both cores, and the other
+ * is executing from flash concurrently with the operation. In this case, you
+ * must perform your own synchronisation to make sure that no XIP accesses take
+ * place during flash programming. One option is to use the
  * \ref multicore_lockout functions.
  *
  * Likewise they are *unsafe* if you have interrupt handlers or an interrupt
@@ -121,6 +121,132 @@ void flash_do_cmd(const uint8_t *txbuf, uint8_t *rxbuf, size_t count);
 
 void flash_flush_cache(void);
 
+#if !PICO_RP2040
+typedef enum {
+	FLASH_DEVINFO_SIZE_NONE = 0x0,
+	FLASH_DEVINFO_SIZE_8K = 0x1,
+	FLASH_DEVINFO_SIZE_16K = 0x2,
+	FLASH_DEVINFO_SIZE_32K = 0x3,
+	FLASH_DEVINFO_SIZE_64K = 0x4,
+	FLASH_DEVINFO_SIZE_128K = 0x5,
+	FLASH_DEVINFO_SIZE_256K = 0x6,
+	FLASH_DEVINFO_SIZE_512K = 0x7,
+	FLASH_DEVINFO_SIZE_1M = 0x8,
+	FLASH_DEVINFO_SIZE_2M = 0x9,
+	FLASH_DEVINFO_SIZE_4M = 0xa,
+	FLASH_DEVINFO_SIZE_8M = 0xb,
+	FLASH_DEVINFO_SIZE_16M = 0xc,
+	FLASH_DEVINFO_SIZE_MAX = 0xc
+} flash_devinfo_size_t;
+
+/*! \brief Convert a flash/PSRAM size enum to an integer size in bytes
+ *  \ingroup hardware_flash
+ */
+static inline uint32_t flash_devinfo_size_to_bytes(flash_devinfo_size_t size) {
+	if (size == FLASH_DEVINFO_SIZE_NONE) {
+		return 0;
+	} else {
+		return 4096u << (uint)size;
+	}
+}
+
+/*! \brief Convert an integer flash/PSRAM size in bytes to a size enum, as
+  !  stored in OTP and used by the ROM.
+ *  \ingroup hardware_flash
+ */
+static inline flash_devinfo_size_t flash_devinfo_bytes_to_size(uint32_t bytes) {
+	// Must be zero or a power of two
+	valid_params_if(HARDWARE_FLASH, (bytes & (bytes - 1)) == 0u);
+	uint sectors = bytes / 4096u;
+	if (sectors <= 1u) {
+		return FLASH_DEVINFO_SIZE_NONE;
+	} else {
+		return (flash_devinfo_size_t) __builtin_ctz(sectors);
+	}
+}
+
+/*! \brief Get the size of the QSPI device attached to chip select cs, according to FLASH_DEVINFO
+ *  \ingroup hardware_flash
+ *
+ * \param cs Chip select index: 0 is QMI chip select 0 (QSPI CS pin), 1 is QMI chip select 1.
+ *
+ * The bootrom reads the FLASH_DEVINFO OTP data entry from OTP into boot RAM during startup. This
+ * contains basic information about the flash device which can be queried without communicating
+ * with the external device.(There are several methods to determine the size of a QSPI device over
+ * QSPI, but none are universally supported.)
+ *
+ * Since the FLASH_DEVINFO information is stored in boot RAM at runtime, it can be updated. Updates
+ * made in this way persist until the next reboot. The ROM uses this device information to control
+ * some low-level flash API behaviour, such as issuing an XIP exit sequence to CS 1 if its size is
+ * nonzero.
+ *
+ * If the macro PICO_FLASH_SIZE_BYTES is specified, this overrides the value for chip select 0. This
+ * can be specified in a board header if a board is always equipped with the same size of flash.
+ */
+flash_devinfo_size_t flash_devinfo_get_cs_size(uint cs);
+
+/*! \brief Update the size of the QSPI device attached to chip select cs in the runtime copy
+ *         of FLASH_DEVINFO.
+ *
+ *  \ingroup hardware_flash
+ *
+ * \param cs Chip select index: 0 is QMI chip select 0 (QSPI CS pin), 1 is QMI chip select 1.
+ *
+ * \param size The size of the attached device, or FLASH_DEVINFO_SIZE_NONE if there is none on this
+ *  chip select.
+ *
+ * The bootrom maintains a copy in boot RAM of the FLASH_DEVINFO information read from OTP during
+ * startup. This function updates that copy to reflect runtime information about the sizes of
+ * attached QSPI devices.
+ *
+ * This controls the behaviour of some ROM flash APIs, such as bounds checking addresses for
+ * erase/programming in the checked_flash_op() API, or issuing an XIP exit sequence to CS 1 in
+ * flash_exit_xip() if the size is nonzero.
+ */
+void flash_devinfo_set_cs_size(uint cs, flash_devinfo_size_t size);
+
+/*! \brief Check whether all attached devices support D8h block erase with 64k size, according to
+ *         FLASH_DEVINFO.
+ *
+ *  \ingroup hardware_flash
+ *
+ * This controls whether checked_flash_op() ROM API uses D8h 64k block erase where possible, for
+ * faster erase times. If not, this ROM API always uses 20h 4k sector erase.
+ *
+ * The bootrom loads this flag from the OTP FLASH_DEVINFO data entry during startup, and stores it
+ * in boot RAM. You can update the boot RAM copy based on runtime knowledge of the attached QSPI
+ * devices.
+ */
+bool flash_devinfo_get_d8h_erase_supported(void);
+
+/*! \brief Specify whether all attached devices support D8h block erase with 64k size, in the
+ *         runtime copy of FLASH_DEVINFO
+ *
+ *  \ingroup hardware_flash
+ *
+ * This function updates the boot RAM copy of OTP FLASH_DEVINFO. The flag passed here is visible to
+ * ROM APIs, and is also returned in the next call to flash_devinfo_get_d8h_erase_supported()
+ */
+void flash_devinfo_set_d8h_erase_supported(bool supported);
+
+/*! \brief Check the GPIO allocated for each chip select, according to FLASH_DEVINFO
+ *  \ingroup hardware_flash
+ *
+ * \param cs Chip select index (only the value 1 is supported on RP2350)
+ */
+uint flash_devinfo_get_cs_gpio(uint cs);
+
+/*! \brief Update the GPIO allocated for each chip select in the runtime copy of FLASH_DEVINFO
+ *  \ingroup hardware_flash
+ *
+ * \param cs Chip select index (only the value 1 is supported on RP2350)
+ *
+ * \param gpio GPIO index (must be less than NUM_BANK0_GPIOS)
+ */
+void flash_devinfo_set_cs_gpio(uint cs, uint gpio);
+
+#endif // !PICO_RP2040
+
 #ifdef __cplusplus
 }
 #endif