feat(flasher_stub): Increase CPU frequency and write/read speeds over USB-JTAG/Serial

Closes https://github.com/espressif/esptool/pull/781/

Closes https://github.com/espressif/esptool/issues/779/

flasher_stub/include/soc_support.h

@@ -45,6 +45,8 @@
 #define WITH_USB_OTG 1
 #endif // ESP32S3
 
+#define USE_MAX_CPU_FREQ WITH_USB_JTAG_SERIAL // Increase CPU freq only when USB-JTAG/Serial is used
+
 /**********************************************************
  * Per-SOC based peripheral register base addresses
  */
@@ -87,6 +89,7 @@
 #define GPIO_BASE_REG       0x60004000
 #define RTCCNTL_BASE_REG    0x60008000
 #define USB_DEVICE_BASE_REG 0x60043000
+#define SYSTEM_BASE_REG     0x600C0000
 #endif
 
 #ifdef ESP32C6BETA
@@ -293,12 +296,28 @@
 #define SYSTEM_CPUPERIOD_SEL_M        ((SYSTEM_CPUPERIOD_SEL_V)<<(SYSTEM_CPUPERIOD_SEL_S))
 #define SYSTEM_CPUPERIOD_SEL_V        0x3
 #define SYSTEM_CPUPERIOD_SEL_S        0
+#define SYSTEM_CPUPERIOD_MAX          2  // CPU_CLK frequency is 240 MHz
 
 #define SYSTEM_SYSCLK_CONF_REG        (SYSTEM_BASE_REG + 0x060)
 #define SYSTEM_SOC_CLK_SEL_M          ((SYSTEM_SOC_CLK_SEL_V)<<(SYSTEM_SOC_CLK_SEL_S))
 #define SYSTEM_SOC_CLK_SEL_V          0x3
 #define SYSTEM_SOC_CLK_SEL_S          10
-#endif
+#define SYSTEM_SOC_CLK_MAX            1
+#endif // ESP32S3
+
+#ifdef ESP32C3
+#define SYSTEM_CPU_PER_CONF_REG       (SYSTEM_BASE_REG + 0x008)
+#define SYSTEM_CPUPERIOD_SEL_M        ((SYSTEM_CPUPERIOD_SEL_V)<<(SYSTEM_CPUPERIOD_SEL_S))
+#define SYSTEM_CPUPERIOD_SEL_V        0x3
+#define SYSTEM_CPUPERIOD_SEL_S        0
+#define SYSTEM_CPUPERIOD_MAX          1  // CPU_CLK frequency is 160 MHz
+
+#define SYSTEM_SYSCLK_CONF_REG        (SYSTEM_BASE_REG + 0x058)
+#define SYSTEM_SOC_CLK_SEL_M          ((SYSTEM_SOC_CLK_SEL_V)<<(SYSTEM_SOC_CLK_SEL_S))
+#define SYSTEM_SOC_CLK_SEL_V          0x3
+#define SYSTEM_SOC_CLK_SEL_S          10
+#define SYSTEM_SOC_CLK_MAX            1
+#endif // ESP32C3
 
 /**********************************************************
  * Per-SOC security info buffer size

flasher_stub/include/stub_io.h

@@ -49,4 +49,8 @@ void stub_io_set_baudrate(uint32_t current_baud, uint32_t new_baud);
  */
 void stub_io_idle_hook(void);
 
-
+/* Checks if USB-Serial/JTAG is being currently used.
+ */
+#if WITH_USB_JTAG_SERIAL
+bool stub_uses_usb_jtag_serial(void);
+#endif

flasher_stub/stub_flasher.c

@@ -48,6 +48,34 @@ static uint8_t calculate_checksum(uint8_t *buf, int length)
   return res;
 }
 
+#if USE_MAX_CPU_FREQ
+static uint32_t cpu_per_conf_reg = 0;
+static uint32_t sysclk_conf_reg = 0;
+
+static void set_max_cpu_freq()
+{
+  if (stub_uses_usb_jtag_serial())
+  {
+    /* Set CPU frequency to max. This also increases SPI speed. */
+    cpu_per_conf_reg = READ_REG(SYSTEM_CPU_PER_CONF_REG);
+    sysclk_conf_reg = READ_REG(SYSTEM_SYSCLK_CONF_REG);
+    WRITE_REG(SYSTEM_CPU_PER_CONF_REG, (cpu_per_conf_reg & ~SYSTEM_CPUPERIOD_SEL_M) | (SYSTEM_CPUPERIOD_MAX << SYSTEM_CPUPERIOD_SEL_S));
+    WRITE_REG(SYSTEM_SYSCLK_CONF_REG, (sysclk_conf_reg & ~SYSTEM_SOC_CLK_SEL_M) | (SYSTEM_SOC_CLK_MAX << SYSTEM_SOC_CLK_SEL_S));
+  }
+}
+
+static void reset_cpu_freq()
+{
+  /* Restore saved sysclk_conf and cpu_per_conf registers.
+     Use only if set_max_cpu_freq() has been called. */
+  if (stub_uses_usb_jtag_serial() && sysclk_conf_reg != 0 && cpu_per_conf_reg != 0)
+  {
+    WRITE_REG(SYSTEM_CPU_PER_CONF_REG, (READ_REG(SYSTEM_CPU_PER_CONF_REG) & ~SYSTEM_CPUPERIOD_SEL_M) | (cpu_per_conf_reg & SYSTEM_CPUPERIOD_SEL_M));
+    WRITE_REG(SYSTEM_SYSCLK_CONF_REG, (READ_REG(SYSTEM_SYSCLK_CONF_REG) & ~SYSTEM_SOC_CLK_SEL_M) | (sysclk_conf_reg & SYSTEM_SOC_CLK_SEL_M));
+  }
+}
+#endif // USE_MAX_CPU_FREQ
+
 static void stub_handle_rx_byte(char byte)
 {
   int16_t r = SLIP_recv_byte(byte, (slip_state_t *)&ub.state);
@@ -309,6 +337,9 @@ void cmd_loop() {
           /* Flush the FLASH_END response before rebooting */
           stub_tx_flush();
           ets_delay_us(10000);
+          #if USE_MAX_CPU_FREQ
+            reset_cpu_freq();
+          #endif // USE_MAX_CPU_FREQ
           software_reset();
         }
         break;
@@ -324,6 +355,9 @@ void cmd_loop() {
                  function. But for our purposes so far, having a bit of
                  extra stuff on the stack doesn't really matter.
               */
+              #if USE_MAX_CPU_FREQ
+                reset_cpu_freq();
+              #endif // USE_MAX_CPU_FREQ
               entrypoint_fn();
           }
           break;
@@ -368,11 +402,9 @@ void stub_main()
   /* this points to stub_main now, clear for next boot */
   ets_set_user_start(0);
 
-#ifdef ESP32S3
-  /* Set CPU frequency to 240 MHz. This also increases SPI speed to 20 MHz. */
-  WRITE_REG(SYSTEM_CPU_PER_CONF_REG, (READ_REG(SYSTEM_CPU_PER_CONF_REG) & ~SYSTEM_CPUPERIOD_SEL_M) | (2 << SYSTEM_CPUPERIOD_SEL_S));
-  WRITE_REG(SYSTEM_SYSCLK_CONF_REG, (READ_REG(SYSTEM_SYSCLK_CONF_REG) & ~SYSTEM_SOC_CLK_SEL_M) | (1 << SYSTEM_SOC_CLK_SEL_S));
-#endif
+  #if USE_MAX_CPU_FREQ
+    set_max_cpu_freq();
+  #endif // USE_MAX_CPU_FREQ
 
   /* zero bss */
   for(uint32_t *p = &_bss_start; p < &_bss_end; p++) {
@@ -412,6 +444,8 @@ void stub_main()
   cmd_loop();
 
   /* if cmd_loop returns, it's due to ESP_RUN_USER_CODE command. */
-
+  #if USE_MAX_CPU_FREQ
+    reset_cpu_freq();
+  #endif // USE_MAX_CPU_FREQ
   return;
 }

flasher_stub/stub_io.c

@@ -40,7 +40,7 @@ void uart_isr(void *arg) {
 }
 
 #if WITH_USB_JTAG_SERIAL
-static bool stub_uses_usb_jtag_serial(void)
+bool stub_uses_usb_jtag_serial(void)
 {
   UartDevice *uart = GetUartDevice();