From f39708abb23b7c7794c04de184213433de481a33 Mon Sep 17 00:00:00 2001
From: Scott Shawcroft <scott@tannewt.org>
Date: Fri, 21 Aug 2020 16:20:58 -0700
Subject: [PATCH] Add PSRAM support to ESP32S2

When configured the CircuitPython heap will be on the external RAM.
When not available, the heap will be 48k inside the IDF heap.
---
 main.c                                        |  5 +--
 .../boards/espressif_saola_1_wrover/sdkconfig | 33 +++++++++++++++++
 .../unexpectedmaker_feathers2/sdkconfig       | 35 +++++++++++++++++++
 ports/esp32s2/supervisor/port.c               | 23 +++++++++---
 4 files changed, 90 insertions(+), 6 deletions(-)

diff --git a/main.c b/main.c
index fb73b09b1f..3a14df0fa6 100755
--- a/main.c
+++ b/main.c
@@ -436,11 +436,12 @@ int run_repl(void) {
 }
 
 int __attribute__((used)) main(void) {
-    memory_init();
-
     // initialise the cpu and peripherals
     safe_mode_t safe_mode = port_init();
 
+    // Init memory after the port in case the port needs to set aside memory.
+    memory_init();
+
     // Turn on LEDs
     init_status_leds();
     rgb_led_status_init();
diff --git a/ports/esp32s2/boards/espressif_saola_1_wrover/sdkconfig b/ports/esp32s2/boards/espressif_saola_1_wrover/sdkconfig
index e69de29bb2..9d8bbde967 100644
--- a/ports/esp32s2/boards/espressif_saola_1_wrover/sdkconfig
+++ b/ports/esp32s2/boards/espressif_saola_1_wrover/sdkconfig
@@ -0,0 +1,33 @@
+CONFIG_ESP32S2_SPIRAM_SUPPORT=y
+
+#
+# SPI RAM config
+#
+# CONFIG_SPIRAM_TYPE_AUTO is not set
+CONFIG_SPIRAM_TYPE_ESPPSRAM16=y
+# CONFIG_SPIRAM_TYPE_ESPPSRAM32 is not set
+# CONFIG_SPIRAM_TYPE_ESPPSRAM64 is not set
+CONFIG_SPIRAM_SIZE=2097152
+
+#
+# PSRAM clock and cs IO for ESP32S2
+#
+CONFIG_DEFAULT_PSRAM_CLK_IO=30
+CONFIG_DEFAULT_PSRAM_CS_IO=26
+# end of PSRAM clock and cs IO for ESP32S2
+
+# CONFIG_SPIRAM_FETCH_INSTRUCTIONS is not set
+# CONFIG_SPIRAM_RODATA is not set
+# CONFIG_SPIRAM_SPEED_80M is not set
+CONFIG_SPIRAM_SPEED_40M=y
+# CONFIG_SPIRAM_SPEED_26M is not set
+# CONFIG_SPIRAM_SPEED_20M is not set
+CONFIG_SPIRAM=y
+CONFIG_SPIRAM_BOOT_INIT=y
+# CONFIG_SPIRAM_IGNORE_NOTFOUND is not set
+CONFIG_SPIRAM_USE_MEMMAP=y
+# CONFIG_SPIRAM_USE_CAPS_ALLOC is not set
+# CONFIG_SPIRAM_USE_MALLOC is not set
+CONFIG_SPIRAM_MEMTEST=y
+# CONFIG_SPIRAM_ALLOW_BSS_SEG_EXTERNAL_MEMORY is not set
+# end of SPI RAM config
diff --git a/ports/esp32s2/boards/unexpectedmaker_feathers2/sdkconfig b/ports/esp32s2/boards/unexpectedmaker_feathers2/sdkconfig
index e69de29bb2..b73c4a8c20 100644
--- a/ports/esp32s2/boards/unexpectedmaker_feathers2/sdkconfig
+++ b/ports/esp32s2/boards/unexpectedmaker_feathers2/sdkconfig
@@ -0,0 +1,35 @@
+CONFIG_ESP32S2_SPIRAM_SUPPORT=y
+
+#
+# SPI RAM config
+#
+# CONFIG_SPIRAM_TYPE_AUTO is not set
+# CONFIG_SPIRAM_TYPE_ESPPSRAM16 is not set
+# CONFIG_SPIRAM_TYPE_ESPPSRAM32 is not set
+CONFIG_SPIRAM_TYPE_ESPPSRAM64=y
+CONFIG_SPIRAM_SIZE=8388608
+
+#
+# PSRAM clock and cs IO for ESP32S2
+#
+CONFIG_DEFAULT_PSRAM_CLK_IO=30
+CONFIG_DEFAULT_PSRAM_CS_IO=26
+# end of PSRAM clock and cs IO for ESP32S2
+
+CONFIG_SPIRAM_SPIWP_SD3_PIN=28
+# CONFIG_SPIRAM_FETCH_INSTRUCTIONS is not set
+# CONFIG_SPIRAM_RODATA is not set
+# CONFIG_SPIRAM_USE_AHB_DBUS3 is not set
+# CONFIG_SPIRAM_SPEED_80M is not set
+CONFIG_SPIRAM_SPEED_40M=y
+# CONFIG_SPIRAM_SPEED_26M is not set
+# CONFIG_SPIRAM_SPEED_20M is not set
+CONFIG_SPIRAM=y
+CONFIG_SPIRAM_BOOT_INIT=y
+# CONFIG_SPIRAM_IGNORE_NOTFOUND is not set
+CONFIG_SPIRAM_USE_MEMMAP=y
+# CONFIG_SPIRAM_USE_CAPS_ALLOC is not set
+# CONFIG_SPIRAM_USE_MALLOC is not set
+CONFIG_SPIRAM_MEMTEST=y
+# CONFIG_SPIRAM_ALLOW_BSS_SEG_EXTERNAL_MEMORY is not set
+# end of SPI RAM config
diff --git a/ports/esp32s2/supervisor/port.c b/ports/esp32s2/supervisor/port.c
index 46590b5513..08710fe910 100644
--- a/ports/esp32s2/supervisor/port.c
+++ b/ports/esp32s2/supervisor/port.c
@@ -43,7 +43,14 @@
 #include "supervisor/memory.h"
 #include "supervisor/shared/tick.h"
 
-#include "rmt.h"
+#include "peripherals/rmt.h"
+#include "esp-idf/components/heap/include/esp_heap_caps.h"
+#include "esp-idf/components/soc/soc/esp32s2/include/soc/cache_memory.h"
+
+#define HEAP_SIZE (48 * 1024)
+
+uint32_t* heap;
+uint32_t heap_size;
 
 STATIC esp_timer_handle_t _tick_timer;
 
@@ -59,6 +66,16 @@ safe_mode_t port_init(void) {
     args.name = "CircuitPython Tick";
     esp_timer_create(&args, &_tick_timer);
     never_reset_module_internal_pins();
+
+    #ifdef CONFIG_SPIRAM
+        heap = (uint32_t*) (DRAM0_CACHE_ADDRESS_HIGH - CONFIG_SPIRAM_SIZE);
+        heap_size = CONFIG_SPIRAM_SIZE / sizeof(uint32_t);
+    #endif
+
+    if (heap == NULL) {
+        heap = malloc(HEAP_SIZE);
+        heap_size = HEAP_SIZE / sizeof(uint32_t);
+    }
     return NO_SAFE_MODE;
 }
 
@@ -86,14 +103,12 @@ void reset_to_bootloader(void) {
 void reset_cpu(void) {
 }
 
-uint32_t heap[64 / sizeof(uint32_t) * 1024];
-
 uint32_t *port_heap_get_bottom(void) {
     return heap;
 }
 
 uint32_t *port_heap_get_top(void) {
-    return heap + sizeof(heap) / sizeof(heap[0]);
+    return heap + heap_size;
 }
 
 uint32_t *port_stack_get_limit(void) {