rp2/rp2_flash: Configure optimal flash timings.

Configure flash timings dynamically to match the system clock.  Reconfigure
timings after flash writes.

Changes are:
- ports/rp2/main.c: Set default flash timings.
- ports/rp2/modmachine.c: Configure optimal flash timings on freq change.
- ports/rp2/rp2_flash.c: Reconfigure flash when leaving critical section.

Signed-off-by: Phil Howard <github@gadgetoid.com>

4 files changed, 120 insertions, 1 deletions

diff --git a/ports/rp2/main.c b/ports/rp2/main.c
index 764af6b74..0f10f63c6 100644
--- a/ports/rp2/main.c
+++ b/ports/rp2/main.c
@@ -26,6 +26,7 @@
 
 #include <stdio.h>
 
+#include "rp2_flash.h"
 #include "py/compile.h"
 #include "py/cstack.h"
 #include "py/runtime.h"
@@ -94,6 +95,9 @@ int main(int argc, char **argv) {
     // Hook for setting up anything that needs to be super early in the boot-up process.
     MICROPY_BOARD_STARTUP();
 
+    // Set the flash divisor to an appropriate value
+    rp2_flash_set_timing();
+
     #if MICROPY_HW_ENABLE_PSRAM
     size_t psram_size = psram_init(MICROPY_HW_PSRAM_CS_PIN);
     #endif
diff --git a/ports/rp2/modmachine.c b/ports/rp2/modmachine.c
index bf4e43a61..58a3a8ae4 100644
--- a/ports/rp2/modmachine.c
+++ b/ports/rp2/modmachine.c
@@ -32,6 +32,7 @@
 #include "modmachine.h"
 #include "uart.h"
 #include "rp2_psram.h"
+#include "rp2_flash.h"
 #include "clocks_extra.h"
 #include "hardware/pll.h"
 #include "hardware/structs/rosc.h"
@@ -95,6 +96,11 @@ static mp_obj_t mp_machine_get_freq(void) {
 
 static void mp_machine_set_freq(size_t n_args, const mp_obj_t *args) {
     mp_int_t freq = mp_obj_get_int(args[0]);
+
+    // If necessary, increase the flash divider before increasing the clock speed
+    const int old_freq = clock_get_hz(clk_sys);
+    rp2_flash_set_timing_for_freq(MAX(freq, old_freq));
+
     if (!set_sys_clock_khz(freq / 1000, false)) {
         mp_raise_ValueError(MP_ERROR_TEXT("cannot change frequency"));
     }
@@ -112,6 +118,12 @@ static void mp_machine_set_freq(size_t n_args, const mp_obj_t *args) {
             }
         }
     }
+
+    // If clock speed was reduced, maybe we can reduce the flash divider
+    if (freq < old_freq) {
+        rp2_flash_set_timing_for_freq(freq);
+    }
+
     #if MICROPY_HW_ENABLE_UART_REPL
     setup_default_uart();
     mp_uart_init();
diff --git a/ports/rp2/rp2_flash.c b/ports/rp2/rp2_flash.c
index 8739f17a4..a9beabf05 100644
--- a/ports/rp2/rp2_flash.c
+++ b/ports/rp2/rp2_flash.c
@@ -34,6 +34,12 @@
 #include "hardware/flash.h"
 #include "pico/binary_info.h"
 #include "rp2_psram.h"
+#ifdef PICO_RP2350
+#include "hardware/structs/ioqspi.h"
+#include "hardware/structs/qmi.h"
+#else
+#include "hardware/structs/ssi.h"
+#endif
 
 #define BLOCK_SIZE_BYTES (FLASH_SECTOR_SIZE)
 
@@ -94,6 +100,48 @@ static bool use_multicore_lockout(void) {
     ;
 }
 
+// Function to set the flash divisor to the correct divisor, assumes interrupts disabled
+// and core1 locked out if relevant.
+static void __no_inline_not_in_flash_func(rp2_flash_set_timing_internal)(int clock_hz) {
+
+    // Use the minimum divisor assuming a 133MHz flash.
+    const int max_flash_freq = 133000000;
+    int divisor = (clock_hz + max_flash_freq - 1) / max_flash_freq;
+
+    #if PICO_RP2350
+    // Make sure flash is deselected - QMI doesn't appear to have a busy flag(!)
+    while ((ioqspi_hw->io[1].status & IO_QSPI_GPIO_QSPI_SS_STATUS_OUTTOPAD_BITS) != IO_QSPI_GPIO_QSPI_SS_STATUS_OUTTOPAD_BITS) {
+        ;
+    }
+
+    // RX delay equal to the divisor means sampling at the same time as the next falling edge of SCK after the
+    // falling edge that generated the data.  This is pretty tight at 133MHz but seems to work with the Winbond flash chips.
+    const int rxdelay = divisor;
+    qmi_hw->m[0].timing = (1 << QMI_M0_TIMING_COOLDOWN_LSB) |
+        rxdelay << QMI_M1_TIMING_RXDELAY_LSB |
+        divisor << QMI_M1_TIMING_CLKDIV_LSB;
+
+    // Force a read through XIP to ensure the timing is applied
+    volatile uint32_t *ptr = (volatile uint32_t *)0x14000000;
+    (void)*ptr;
+    #else
+    // RP2040 SSI hardware only supports even divisors
+    if (divisor & 1) {
+        divisor += 1;
+    }
+
+    // Wait for SSI not busy
+    while (ssi_hw->sr & SSI_SR_BUSY_BITS) {
+        ;
+    }
+
+    // Disable, set the new divisor, and re-enable
+    hw_clear_bits(&ssi_hw->ssienr, SSI_SSIENR_SSI_EN_BITS);
+    ssi_hw->baudr = divisor;
+    hw_set_bits(&ssi_hw->ssienr, SSI_SSIENR_SSI_EN_BITS);
+    #endif
+}
+
 // Flash erase and write must run with interrupts disabled and the other core suspended,
 // because the XIP bit gets disabled.
 static uint32_t begin_critical_flash_section(void) {
@@ -117,8 +165,9 @@ static uint32_t begin_critical_flash_section(void) {
 }
 
 static void end_critical_flash_section(uint32_t state) {
+    // The ROM function to program flash will have reset flash and PSRAM timings to defaults
+    rp2_flash_set_timing_internal(clock_get_hz(clk_sys));
     #if MICROPY_HW_ENABLE_PSRAM
-    // The ROM function to program flash will reset PSRAM timings to defaults
     psram_init(MICROPY_HW_PSRAM_CS_PIN);
     #endif
     restore_interrupts(state);
@@ -313,3 +362,23 @@ mp_obj_t mp_vfs_rom_ioctl(size_t n_args, const mp_obj_t *args) {
     }
 }
 #endif
+
+// Modify the flash timing.  Ensure flash access is suspended while
+// the timings are altered.
+void rp2_flash_set_timing_for_freq(int clock_hz) {
+    if (multicore_lockout_victim_is_initialized(1 - get_core_num())) {
+        multicore_lockout_start_blocking();
+    }
+    uint32_t state = save_and_disable_interrupts();
+
+    rp2_flash_set_timing_internal(clock_hz);
+
+    restore_interrupts(state);
+    if (multicore_lockout_victim_is_initialized(1 - get_core_num())) {
+        multicore_lockout_end_blocking();
+    }
+}
+
+void rp2_flash_set_timing(void) {
+    rp2_flash_set_timing_for_freq(clock_get_hz(clk_sys));
+}
diff --git a/ports/rp2/rp2_flash.h b/ports/rp2/rp2_flash.h
new file mode 100644
index 000000000..3c1fbbff2
--- /dev/null
+++ b/ports/rp2/rp2_flash.h
@@ -0,0 +1,34 @@
+/*
+ * This file is part of the MicroPython project, http://micropython.org/
+ *
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2025 Mike Bell
+ *                    Phil Howard
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+#ifndef MICROPY_INCLUDED_RP2_RP2_FLASH_H
+#define MICROPY_INCLUDED_RP2_RP2_FLASH_H
+
+extern void rp2_flash_set_timing_for_freq(int clock_hz);
+extern void rp2_flash_set_timing(void);
+
+#endif