Improve PSRAM support (#2884)

CODEOWNERS

@@ -125,6 +125,7 @@ esphome/components/pn532_i2c/* @OttoWinter @jesserockz
 esphome/components/pn532_spi/* @OttoWinter @jesserockz
 esphome/components/power_supply/* @esphome/core
 esphome/components/preferences/* @esphome/core
+esphome/components/psram/* @esphome/core
 esphome/components/pulse_meter/* @stevebaxter
 esphome/components/pvvx_mithermometer/* @pasiz
 esphome/components/rc522/* @glmnet

esphome/components/display/display_buffer.cpp

@@ -5,6 +5,7 @@
 #include "esphome/core/color.h"
 #include "esphome/core/log.h"
 #include "esphome/core/hal.h"
+#include "esphome/core/helpers.h"
 
 namespace esphome {
 namespace display {
@@ -15,7 +16,8 @@ const Color COLOR_OFF(0, 0, 0, 0);
 const Color COLOR_ON(255, 255, 255, 255);
 
 void DisplayBuffer::init_internal_(uint32_t buffer_length) {
-  this->buffer_ = new (std::nothrow) uint8_t[buffer_length];  // NOLINT
+  ExternalRAMAllocator<uint8_t> allocator(ExternalRAMAllocator<uint8_t>::ALLOW_FAILURE);
+  this->buffer_ = allocator.allocate(buffer_length);
   if (this->buffer_ == nullptr) {
     ESP_LOGE(TAG, "Could not allocate buffer for display!");
     return;

esphome/components/esp32_camera/__init__.py

@@ -19,6 +19,8 @@ from esphome.cpp_helpers import setup_entity
 
 DEPENDENCIES = ["esp32"]
 
+AUTO_LOAD = ["psram"]
+
 esp32_camera_ns = cg.esphome_ns.namespace("esp32_camera")
 ESP32Camera = esp32_camera_ns.class_("ESP32Camera", cg.PollingComponent, cg.EntityBase)
 ESP32CameraFrameSize = esp32_camera_ns.enum("ESP32CameraFrameSize")
@@ -153,9 +155,7 @@ async def to_code(config):
     cg.add(var.set_frame_size(config[CONF_RESOLUTION]))
 
     cg.add_define("USE_ESP32_CAMERA")
-    cg.add_build_flag("-DBOARD_HAS_PSRAM")
 
     if CORE.using_esp_idf:
         cg.add_library("espressif/esp32-camera", "1.0.0")
         add_idf_sdkconfig_option("CONFIG_RTCIO_SUPPORT_RTC_GPIO_DESC", True)
-        add_idf_sdkconfig_option("CONFIG_ESP32_SPIRAM_SUPPORT", True)

esphome/components/esp32_camera/esp32_camera.cpp

@@ -49,9 +49,6 @@ void ESP32Camera::dump_config() {
   ESP_LOGCONFIG(TAG, "ESP32 Camera:");
   ESP_LOGCONFIG(TAG, "  Name: %s", this->name_.c_str());
   ESP_LOGCONFIG(TAG, "  Internal: %s", YESNO(this->internal_));
-#ifdef USE_ARDUINO
-  ESP_LOGCONFIG(TAG, "  Board Has PSRAM: %s", YESNO(psramFound()));
-#endif  // USE_ARDUINO
   ESP_LOGCONFIG(TAG, "  Data Pins: D0:%d D1:%d D2:%d D3:%d D4:%d D5:%d D6:%d D7:%d", conf.pin_d0, conf.pin_d1,
                 conf.pin_d2, conf.pin_d3, conf.pin_d4, conf.pin_d5, conf.pin_d6, conf.pin_d7);
   ESP_LOGCONFIG(TAG, "  VSYNC Pin: %d", conf.pin_vsync);

esphome/components/inkplate6/display.py

@@ -12,6 +12,7 @@ from esphome.const import (
 )
 
 DEPENDENCIES = ["i2c", "esp32"]
+AUTO_LOAD = ["psram"]
 
 CONF_DISPLAY_DATA_0_PIN = "display_data_0_pin"
 CONF_DISPLAY_DATA_1_PIN = "display_data_1_pin"
@@ -179,5 +180,3 @@ async def to_code(config):
 
     display_data_7 = await cg.gpio_pin_expression(config[CONF_DISPLAY_DATA_7_PIN])
     cg.add(var.set_display_data_7_pin(display_data_7))
-
-    cg.add_build_flag("-DBOARD_HAS_PSRAM")

esphome/components/inkplate6/inkplate.cpp

@@ -42,32 +42,32 @@ void Inkplate6::setup() {
   this->display();
 }
 void Inkplate6::initialize_() {
+  ExternalRAMAllocator<uint8_t> allocator(ExternalRAMAllocator<uint8_t>::ALLOW_FAILURE);
   uint32_t buffer_size = this->get_buffer_length_();
+  if (buffer_size == 0)
+    return;
 
-  if (this->partial_buffer_ != nullptr) {
+  if (this->partial_buffer_ != nullptr)
-    free(this->partial_buffer_);  // NOLINT
+    allocator.deallocate(this->partial_buffer_, buffer_size);
-  }
+  if (this->partial_buffer_2_ != nullptr)
-  if (this->partial_buffer_2_ != nullptr) {
+    allocator.deallocate(this->partial_buffer_2_, buffer_size * 2);
-    free(this->partial_buffer_2_);  // NOLINT
+  if (this->buffer_ != nullptr)
-  }
+    allocator.deallocate(this->buffer_, buffer_size);
-  if (this->buffer_ != nullptr) {
-    free(this->buffer_);  // NOLINT
-  }
 
-  this->buffer_ = (uint8_t *) ps_malloc(buffer_size);
+  this->buffer_ = allocator.allocate(buffer_size);
   if (this->buffer_ == nullptr) {
     ESP_LOGE(TAG, "Could not allocate buffer for display!");
     this->mark_failed();
     return;
   }
   if (!this->greyscale_) {
-    this->partial_buffer_ = (uint8_t *) ps_malloc(buffer_size);
+    this->partial_buffer_ = allocator.allocate(buffer_size);
     if (this->partial_buffer_ == nullptr) {
       ESP_LOGE(TAG, "Could not allocate partial buffer for display!");
       this->mark_failed();
       return;
     }
-    this->partial_buffer_2_ = (uint8_t *) ps_malloc(buffer_size * 2);
+    this->partial_buffer_2_ = allocator.allocate(buffer_size * 2);
     if (this->partial_buffer_2_ == nullptr) {
       ESP_LOGE(TAG, "Could not allocate partial buffer 2 for display!");
       this->mark_failed();

esphome/components/nextion/nextion_upload.cpp

@@ -8,6 +8,10 @@
 #include "esphome/core/log.h"
 #include "esphome/components/network/util.h"
 
+#ifdef USE_ESP32
+#include <esp_heap_caps.h>
+#endif
+
 namespace esphome {
 namespace nextion {
 static const char *const TAG = "nextion_upload";
@@ -158,12 +162,8 @@ void Nextion::upload_tft() {
   if (!begin_status) {
     this->is_updating_ = false;
     ESP_LOGD(TAG, "connection failed");
-#ifdef USE_ESP32
+    ExternalRAMAllocator<uint8_t> allocator(ExternalRAMAllocator<uint8_t>::ALLOW_FAILURE);
-    if (psramFound())
+    allocator.deallocate(this->transfer_buffer_, this->transfer_buffer_size_);
-      free(this->transfer_buffer_);  // NOLINT
-    else
-#endif
-      delete this->transfer_buffer_;
     return;
   } else {
     ESP_LOGD(TAG, "Connected");
@@ -252,7 +252,7 @@ void Nextion::upload_tft() {
   // Nextion wants 4096 bytes at a time. Make chunk_size a multiple of 4096
 #ifdef USE_ESP32
   uint32_t chunk_size = 8192;
-  if (psramFound()) {
+  if (heap_caps_get_free_size(MALLOC_CAP_SPIRAM) > 0) {
     chunk_size = this->content_length_;
   } else {
     if (ESP.getFreeHeap() > 40960) {  // 32K to keep on hand
@@ -269,32 +269,18 @@ void Nextion::upload_tft() {
 #endif
 
   if (this->transfer_buffer_ == nullptr) {
-#ifdef USE_ESP32
+    ExternalRAMAllocator<uint8_t> allocator(ExternalRAMAllocator<uint8_t>::ALLOW_FAILURE);
-    if (psramFound()) {
-      ESP_LOGD(TAG, "Allocating PSRAM buffer size %d, Free PSRAM size is %u", chunk_size, ESP.getFreePsram());
-      this->transfer_buffer_ = (uint8_t *) ps_malloc(chunk_size);
-      if (this->transfer_buffer_ == nullptr) {
-        ESP_LOGE(TAG, "Could not allocate buffer size %d!", chunk_size);
-        this->upload_end_();
-      }
-    } else {
-#endif
     // NOLINTNEXTLINE(readability-static-accessed-through-instance)
     ESP_LOGD(TAG, "Allocating buffer size %d, Heap size is %u", chunk_size, ESP.getFreeHeap());
-      // NOLINTNEXTLINE(cppcoreguidelines-owning-memory)
+    this->transfer_buffer_ = allocator.allocate(chunk_size);
-      this->transfer_buffer_ = new (std::nothrow) uint8_t[chunk_size];
     if (this->transfer_buffer_ == nullptr) {  // Try a smaller size
       ESP_LOGD(TAG, "Could not allocate buffer size: %d trying 4096 instead", chunk_size);
       chunk_size = 4096;
       ESP_LOGD(TAG, "Allocating %d buffer", chunk_size);
-        // NOLINTNEXTLINE(cppcoreguidelines-owning-memory)
+      this->transfer_buffer_ = allocator.allocate(chunk_size);
-        this->transfer_buffer_ = new uint8_t[chunk_size];
 
       if (!this->transfer_buffer_)
         this->upload_end_();
-#ifdef USE_ESP32
-      }
-#endif
     }
 
     this->transfer_buffer_size_ = chunk_size;

esphome/components/psram/__init__.py

@@ -0,0 +1,29 @@
+import esphome.codegen as cg
+import esphome.config_validation as cv
+from esphome.components.esp32 import add_idf_sdkconfig_option
+from esphome.core import CORE
+from esphome.const import (
+    CONF_ID,
+)
+
+CODEOWNERS = ["@esphome/core"]
+
+psram_ns = cg.esphome_ns.namespace("psram")
+PsramComponent = psram_ns.class_("PsramComponent", cg.Component)
+
+CONFIG_SCHEMA = cv.All(
+    cv.Schema({cv.GenerateID(): cv.declare_id(PsramComponent)}), cv.only_on_esp32
+)
+
+
+async def to_code(config):
+    if CORE.using_arduino:
+        cg.add_build_flag("-DBOARD_HAS_PSRAM")
+
+    if CORE.using_esp_idf:
+        add_idf_sdkconfig_option("CONFIG_ESP32_SPIRAM_SUPPORT", True)
+        add_idf_sdkconfig_option("CONFIG_SPIRAM_USE_CAPS_ALLOC", True)
+        add_idf_sdkconfig_option("CONFIG_SPIRAM_IGNORE_NOTFOUND", True)
+
+    var = cg.new_Pvariable(config[CONF_ID])
+    await cg.register_component(var, config)

esphome/components/psram/psram.cpp

@@ -0,0 +1,32 @@
+#include "psram.h"
+
+#ifdef USE_ESP32
+
+#include "esphome/core/log.h"
+
+#include <esp_heap_caps.h>
+#include <esp_idf_version.h>
+
+namespace esphome {
+namespace psram {
+
+static const char *const TAG = "psram";
+
+void PsramComponent::dump_config() {
+  // Technically this can be false if the PSRAM is full, but heap_caps_get_total_size() isn't always available, and it's
+  // very unlikely for the PSRAM to be full.
+  bool available = heap_caps_get_free_size(MALLOC_CAP_SPIRAM) > 0;
+
+  ESP_LOGCONFIG(TAG, "PSRAM:");
+  ESP_LOGCONFIG(TAG, "  Available: %s", YESNO(available));
+#if ESP_IDF_VERSION >= ESP_IDF_VERSION_VAL(4, 1, 0)
+  if (available) {
+    ESP_LOGCONFIG(TAG, "  Size: %d MB", heap_caps_get_total_size(MALLOC_CAP_SPIRAM) / 1024 / 1024);
+  }
+#endif
+}
+
+}  // namespace psram
+}  // namespace esphome
+
+#endif

esphome/components/psram/psram.h

@@ -0,0 +1,17 @@
+#pragma once
+
+#ifdef USE_ESP32
+
+#include "esphome/core/component.h"
+
+namespace esphome {
+namespace psram {
+
+class PsramComponent : public Component {
+  void dump_config() override;
+};
+
+}  // namespace psram
+}  // namespace esphome
+
+#endif

esphome/core/helpers.h

@@ -9,8 +9,8 @@
 #include <memory>
 #include <type_traits>
 
-#ifdef USE_ESP32_FRAMEWORK_ARDUINO
+#ifdef USE_ESP32
-#include "esp32-hal-psram.h"
+#include <esp_heap_caps.h>
 #endif
 
 #include "esphome/core/optional.h"
@@ -261,21 +261,6 @@ template<typename T> class Parented {
 
 uint32_t fnv1_hash(const std::string &str);
 
-template<typename T> T *new_buffer(size_t length) {
-  T *buffer;
-#ifdef USE_ESP32_FRAMEWORK_ARDUINO
-  if (psramFound()) {
-    buffer = (T *) ps_malloc(length);
-  } else {
-    buffer = new T[length];  // NOLINT(cppcoreguidelines-owning-memory)
-  }
-#else
-  buffer = new T[length];  // NOLINT(cppcoreguidelines-owning-memory)
-#endif
-
-  return buffer;
-}
-
 // ---------------------------------------------------------------------------------------------------------------------
 
 /// @name STL backports
@@ -486,6 +471,51 @@ template<typename T, typename U> T remap(U value, U min, U max, T min_out, T max
 
 ///@}
 
+/// @name Memory management
+///@{
+
+/** An STL allocator that uses SPI RAM.
+ *
+ * By setting flags, it can be configured to don't try main memory if SPI RAM is full or unavailable, and to return
+ * `nulllptr` instead of aborting when no memory is available.
+ */
+template<class T> class ExternalRAMAllocator {
+ public:
+  using value_type = T;
+
+  enum Flags {
+    NONE = 0,
+    REFUSE_INTERNAL = 1 << 0,  ///< Refuse falling back to internal memory when external RAM is full or unavailable.
+    ALLOW_FAILURE = 1 << 1,    ///< Don't abort when memory allocation fails.
+  };
+
+  ExternalRAMAllocator() = default;
+  ExternalRAMAllocator(Flags flags) : flags_{flags} {}
+  template<class U> constexpr ExternalRAMAllocator(const ExternalRAMAllocator<U> &other) : flags_{other.flags} {}
+
+  T *allocate(size_t n) {
+    size_t size = n * sizeof(T);
+    T *ptr = nullptr;
+#ifdef USE_ESP32
+    ptr = static_cast<T *>(heap_caps_malloc(size, MALLOC_CAP_SPIRAM));
+#endif
+    if (ptr == nullptr && (this->flags_ & Flags::REFUSE_INTERNAL) == 0)
+      ptr = static_cast<T *>(malloc(size));  // NOLINT(cppcoreguidelines-owning-memory,cppcoreguidelines-no-malloc)
+    if (ptr == nullptr && (this->flags_ & Flags::ALLOW_FAILURE) == 0)
+      abort();
+    return ptr;
+  }
+
+  void deallocate(T *p, size_t n) {
+    free(p);  // NOLINT(cppcoreguidelines-owning-memory,cppcoreguidelines-no-malloc)
+  }
+
+ private:
+  Flags flags_{Flags::NONE};
+};
+
+/// @}
+
 /// @name Deprecated functions
 ///@{