rp2/rp2_flash: Support flash writes from PSRAM.

Add a 256 byte (FLASH_PAGE_SIZE) SRAM copy buffer to allow copies from
PSRAM to flash.  This would otherwise hardfault since PSRAM is disabled
when doing a write to flash.

Changes are:
- ports/rp2/rp2_flash.c: Add 256 byte (flash page size) SRAM copy buffer
                         for PSRAM to flash copies.
- ports/rp2/rp2_flash.c: Invalidate the XIP cache to purge any PSRAM
                         data before critical flash operations.

Co-authored-by: Phil Howard <github@gadgetoid.com>
Co-authored-by: Angus Gratton <angus@redyak.com.au>
Signed-off-by: Phil Howard <github@gadgetoid.com>

1 files changed, 59 insertions, 5 deletions

diff --git a/ports/rp2/rp2_flash.c b/ports/rp2/rp2_flash.c
index f3f154a14..8739f17a4 100644
--- a/ports/rp2/rp2_flash.c
+++ b/ports/rp2/rp2_flash.c
@@ -33,9 +33,13 @@
 #include "modrp2.h"
 #include "hardware/flash.h"
 #include "pico/binary_info.h"
+#include "rp2_psram.h"
 
 #define BLOCK_SIZE_BYTES (FLASH_SECTOR_SIZE)
 
+// Size of buffer for flash writes from PSRAM, since they are mutually exclusive
+#define COPY_BUFFER_SIZE_BYTES (FLASH_PAGE_SIZE)
+
 static_assert(MICROPY_HW_ROMFS_BYTES % 4096 == 0, "ROMFS size must be a multiple of 4K");
 static_assert(MICROPY_HW_FLASH_STORAGE_BYTES % 4096 == 0, "Flash storage size must be a multiple of 4K");
 
@@ -96,10 +100,27 @@ static uint32_t begin_critical_flash_section(void) {
     if (use_multicore_lockout()) {
         multicore_lockout_start_blocking();
     }
-    return save_and_disable_interrupts();
+    uint32_t state = save_and_disable_interrupts();
+
+    #if MICROPY_HW_ENABLE_PSRAM
+    // We're about to invalidate the XIP cache, clean it first to commit any dirty writes to PSRAM
+    // Use the upper 16k of the maintenance space (0x1bffc000 through 0x1bffffff) to workaround
+    // incorrect behaviour of the XIP clean operation, where it also alters the tag of the associated
+    // cache line: https://forums.raspberrypi.com/viewtopic.php?t=378249#p2263677
+    volatile uint8_t *maintenance_ptr = (volatile uint8_t *)(XIP_SRAM_BASE + (XIP_MAINTENANCE_BASE - XIP_BASE));
+    for (int i = 1; i < 16 * 1024; i += 8) {
+        maintenance_ptr[i] = 0;
+    }
+    #endif
+
+    return state;
 }
 
 static void end_critical_flash_section(uint32_t state) {
+    #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);
     if (use_multicore_lockout()) {
         multicore_lockout_end_blocking();
@@ -192,10 +213,43 @@ static mp_obj_t rp2_flash_writeblocks(size_t n_args, const mp_obj_t *args) {
     } else {
         offset += mp_obj_get_int(args[3]);
     }
-    mp_uint_t atomic_state = begin_critical_flash_section();
-    flash_range_program(self->flash_base + offset, bufinfo.buf, bufinfo.len);
-    end_critical_flash_section(atomic_state);
-    mp_event_handle_nowait();
+
+    // If copying from SRAM, can write direct to flash.
+    // If copying from PSRAM/flash, use an SRAM buffer and write in chunks.
+    #if MICROPY_HW_ENABLE_PSRAM
+    bool write_direct = (uintptr_t)bufinfo.buf >= SRAM_BASE;
+    #else
+    bool write_direct = true;
+    #endif
+
+    if (write_direct) {
+        // If copying from SRAM, write direct
+        mp_uint_t atomic_state = begin_critical_flash_section();
+        flash_range_program(self->flash_base + offset, bufinfo.buf, bufinfo.len);
+        end_critical_flash_section(atomic_state);
+        mp_event_handle_nowait();
+    }
+    #if MICROPY_HW_ENABLE_PSRAM
+    else {
+        size_t bytes_left = bufinfo.len;
+        size_t bytes_offset = 0;
+        static uint8_t copy_buffer[COPY_BUFFER_SIZE_BYTES] = {0};
+
+        while (bytes_left) {
+            memcpy(copy_buffer, bufinfo.buf + bytes_offset, MIN(bytes_left, COPY_BUFFER_SIZE_BYTES));
+            mp_uint_t atomic_state = begin_critical_flash_section();
+            flash_range_program(self->flash_base + offset + bytes_offset, copy_buffer, MIN(bytes_left, COPY_BUFFER_SIZE_BYTES));
+            end_critical_flash_section(atomic_state);
+            bytes_offset += COPY_BUFFER_SIZE_BYTES;
+            if (bytes_left <= COPY_BUFFER_SIZE_BYTES) {
+                break;
+            }
+            bytes_left -= COPY_BUFFER_SIZE_BYTES;
+            mp_event_handle_nowait();
+        }
+    }
+    #endif
+
     // TODO check return value
     return mp_const_none;
 }