Use PSRAM for internal buffers if available (#1)

* add support for external psram

* don't cast malloc pointer

* fall back to internal memory

* bump version

Author: kahrendt

library.json

@@ -1,6 +1,6 @@
 {
     "name": "ESPMicroSpeechFeatures",
-    "version": "1.0.0",
+    "version": "1.1.0",
     "description": "Generates TensorFlow micro spectrogram features from audio samples",
     "keywords": "tensorflow, spectrogram, audio",
     "repository":

src/fft_util.c

@@ -17,51 +17,84 @@ limitations under the License.
 
 #include <stdio.h>
 
+#ifdef USE_ESP32
+#include <esp_heap_caps.h>
+#endif
+
 #include "kiss_fftr.h"
 
-int FftPopulateState(struct FftState *state, size_t input_size) {
+int FftPopulateState(struct FftState *state, size_t input_size)
+{
   state->input_size = input_size;
   state->fft_size = 1;
-  while (state->fft_size < state->input_size) {
+  while (state->fft_size < state->input_size)
+  {
     state->fft_size <<= 1;
   }
 
-  state->input = (int16_t *) (malloc(state->fft_size * sizeof(*state->input)));
-  if (state->input == NULL) {
+#ifdef USE_ESP32
+  state->input =
+      heap_caps_malloc(state->fft_size * sizeof(*state->input), MALLOC_CAP_SPIRAM | MALLOC_CAP_8BIT);
+  if (state->input == NULL)
+#endif
+  {
+    state->input = malloc(state->fft_size * sizeof(*state->input));
+  }
+  if (state->input == NULL)
+  {
     fprintf(stderr, "Failed to alloc fft input buffer\n");
     return 0;
   }
 
-  // Cast to int32_t instead of the original complex_int16_t, easy way to fix issues converting the original file into C (may not be the best/proper way!)
-  state->output = (int32_t *) (malloc((state->fft_size / 2 + 1) * sizeof(*state->output) * 2));
-  if (state->output == NULL) {
+#ifdef USE_ESP32
+  state->output = heap_caps_malloc((state->fft_size / 2 + 1) * sizeof(*state->output) * 2,
+                                   MALLOC_CAP_SPIRAM | MALLOC_CAP_8BIT);
+  if (state->output == NULL)
+#endif
+  {
+    state->output = malloc((state->fft_size / 2 + 1) * sizeof(*state->output) * 2);
+  }
+  if (state->output == NULL)
+  {
     fprintf(stderr, "Failed to alloc fft output buffer\n");
     return 0;
   }
 
   // Ask kissfft how much memory it wants.
   size_t scratch_size = 0;
   kiss_fftr_cfg kfft_cfg = kiss_fftr_alloc(state->fft_size, 0, NULL, &scratch_size);
-  if (kfft_cfg != NULL) {
+  if (kfft_cfg != NULL)
+  {
     fprintf(stderr, "Kiss memory sizing failed.\n");
     return 0;
   }
-  state->scratch = malloc(scratch_size);
-  if (state->scratch == NULL) {
+
+#ifdef USE_ESP32
+  state->scratch = heap_caps_malloc(scratch_size, MALLOC_CAP_SPIRAM | MALLOC_CAP_8BIT);
+  if (state->scratch == NULL)
+#endif
+  {
+    state->scratch = malloc(scratch_size);
+  }
+
+  if (state->scratch == NULL)
+  {
     fprintf(stderr, "Failed to alloc fft scratch buffer\n");
     return 0;
   }
   state->scratch_size = scratch_size;
   // Let kissfft configure the scratch space we just allocated
   kfft_cfg = kiss_fftr_alloc(state->fft_size, 0, state->scratch, &scratch_size);
-  if (kfft_cfg != state->scratch) {
+  if (kfft_cfg != state->scratch)
+  {
     fprintf(stderr, "Kiss memory preallocation strategy failed.\n");
     return 0;
   }
   return 1;
 }
 
-void FftFreeStateContents(struct FftState *state) {
+void FftFreeStateContents(struct FftState *state)
+{
   free(state->input);
   free(state->output);
   free(state->scratch);

src/filterbank_util.c

@@ -1,4 +1,5 @@
 /* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+Modifications copyright 2024 Kevin Ahrendt.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -18,10 +19,15 @@ limitations under the License.
 #include <math.h>
 #include <stdio.h>
 
+#ifdef USE_ESP32
+#include <esp_heap_caps.h>
+#endif
+
 #define kFilterbankIndexAlignment 4
 #define kFilterbankChannelBlockSize 4
 
-void FilterbankFillConfigWithDefaults(struct FilterbankConfig *config) {
+void FilterbankFillConfigWithDefaults(struct FilterbankConfig *config)
+{
   config->num_channels = 32;
   config->lower_band_limit = 125.0f;
   config->upper_band_limit = 7500.0f;
@@ -31,46 +37,101 @@ void FilterbankFillConfigWithDefaults(struct FilterbankConfig *config) {
 static float FreqToMel(float freq) { return 1127.0f * log1pf(freq / 700.0f); }
 
 static void CalculateCenterFrequencies(const int num_channels, const float lower_frequency_limit,
-                                       const float upper_frequency_limit, float *center_frequencies) {
+                                       const float upper_frequency_limit, float *center_frequencies)
+{
   assert(lower_frequency_limit >= 0.0f);
   assert(upper_frequency_limit > lower_frequency_limit);
 
   const float mel_low = FreqToMel(lower_frequency_limit);
   const float mel_hi = FreqToMel(upper_frequency_limit);
   const float mel_span = mel_hi - mel_low;
-  const float mel_spacing = mel_span / ((float) num_channels);
+  const float mel_spacing = mel_span / ((float)num_channels);
   int i;
-  for (i = 0; i < num_channels; ++i) {
+  for (i = 0; i < num_channels; ++i)
+  {
     center_frequencies[i] = mel_low + (mel_spacing * (i + 1));
   }
 }
 
-static void QuantizeFilterbankWeights(const float float_weight, int16_t *weight, int16_t *unweight) {
+static void QuantizeFilterbankWeights(const float float_weight, int16_t *weight, int16_t *unweight)
+{
   *weight = floorf(float_weight * (1 << kFilterbankBits) + 0.5f);
   *unweight = floorf((1.0f - float_weight) * (1 << kFilterbankBits) + 0.5f);
 }
 
 int FilterbankPopulateState(const struct FilterbankConfig *config, struct FilterbankState *state, int sample_rate,
-                            int spectrum_size) {
+                            int spectrum_size)
+{
   state->num_channels = config->num_channels;
   const int num_channels_plus_1 = config->num_channels + 1;
 
   // How should we align things to index counts given the byte alignment?
   const int index_alignment =
       (kFilterbankIndexAlignment < sizeof(int16_t) ? 1 : kFilterbankIndexAlignment / sizeof(int16_t));
 
-  state->channel_frequency_starts = (int16_t *) malloc(num_channels_plus_1 * sizeof(*state->channel_frequency_starts));
-  state->channel_weight_starts = (int16_t *) malloc(num_channels_plus_1 * sizeof(*state->channel_weight_starts));
-  state->channel_widths = (int16_t *) malloc(num_channels_plus_1 * sizeof(*state->channel_widths));
-  state->work = (uint64_t *) malloc(num_channels_plus_1 * sizeof(*state->work));
+#ifdef USE_ESP32
+  state->channel_frequency_starts = heap_caps_malloc(num_channels_plus_1 * sizeof(*state->channel_frequency_starts), MALLOC_CAP_SPIRAM | MALLOC_CAP_8BIT);
+  if (state->channel_frequency_starts == NULL)
+#endif
+  {
+    state->channel_frequency_starts = malloc(num_channels_plus_1 * sizeof(*state->channel_frequency_starts));
+  }
+
+#ifdef USE_ESP32
+  state->channel_weight_starts = heap_caps_malloc(num_channels_plus_1 * sizeof(*state->channel_weight_starts), MALLOC_CAP_SPIRAM | MALLOC_CAP_8BIT);
+  if (state->channel_weight_starts == NULL)
+#endif
+  {
+    state->channel_weight_starts = malloc(num_channels_plus_1 * sizeof(*state->channel_weight_starts));
+  }
+
+#ifdef USE_ESP32
+  state->channel_widths = heap_caps_malloc(num_channels_plus_1 * sizeof(*state->channel_widths), MALLOC_CAP_SPIRAM | MALLOC_CAP_8BIT);
+  if (state->channel_widths == NULL)
+#endif
+  {
+    state->channel_widths = malloc(num_channels_plus_1 * sizeof(*state->channel_widths));
+  }
+
+#ifdef USE_ESP32
+  state->work = heap_caps_malloc(num_channels_plus_1 * sizeof(*state->work), MALLOC_CAP_SPIRAM | MALLOC_CAP_8BIT);
+  if (state->work == NULL)
+#endif
+  {
+    state->work = malloc(num_channels_plus_1 * sizeof(*state->work));
+  }
+
+  float *center_mel_freqs = NULL;
+#ifdef USE_ESP32
+  center_mel_freqs = heap_caps_malloc(num_channels_plus_1 * sizeof(*center_mel_freqs), MALLOC_CAP_SPIRAM | MALLOC_CAP_8BIT);
+  if (center_mel_freqs == NULL)
+#endif
+  {
+    center_mel_freqs = malloc(num_channels_plus_1 * sizeof(*center_mel_freqs));
+  }
 
-  float *center_mel_freqs = (float *) malloc(num_channels_plus_1 * sizeof(*center_mel_freqs));
-  int16_t *actual_channel_starts = (int16_t *) malloc(num_channels_plus_1 * sizeof(*actual_channel_starts));
-  int16_t *actual_channel_widths = (int16_t *) malloc(num_channels_plus_1 * sizeof(*actual_channel_widths));
+  int16_t *actual_channel_starts = NULL;
+#ifdef USE_ESP32
+  actual_channel_starts = heap_caps_malloc(num_channels_plus_1 * sizeof(*actual_channel_starts), MALLOC_CAP_SPIRAM | MALLOC_CAP_8BIT);
+  if (actual_channel_starts == NULL)
+#endif
+  {
+    actual_channel_starts = malloc(num_channels_plus_1 * sizeof(*actual_channel_starts));
+  }
+
+  int16_t *actual_channel_widths = NULL;
+#ifdef USE_ESP32
+  actual_channel_widths = heap_caps_malloc(num_channels_plus_1 * sizeof(*actual_channel_widths), MALLOC_CAP_SPIRAM | MALLOC_CAP_8BIT);
+  if (actual_channel_widths == NULL)
+#endif
+  {
+    actual_channel_widths = malloc(num_channels_plus_1 * sizeof(*actual_channel_widths));
+  }
 
   if (state->channel_frequency_starts == NULL || state->channel_weight_starts == NULL ||
       state->channel_widths == NULL || center_mel_freqs == NULL || actual_channel_starts == NULL ||
-      actual_channel_widths == NULL) {
+      actual_channel_widths == NULL)
+  {
     free(center_mel_freqs);
     free(actual_channel_starts);
     free(actual_channel_widths);
@@ -81,9 +142,9 @@ int FilterbankPopulateState(const struct FilterbankConfig *config, struct Filter
   CalculateCenterFrequencies(num_channels_plus_1, config->lower_band_limit, config->upper_band_limit, center_mel_freqs);
 
   // Always exclude DC.
-  const float hz_per_sbin = 0.5f * sample_rate / ((float) spectrum_size - 1);
+  const float hz_per_sbin = 0.5f * sample_rate / ((float)spectrum_size - 1);
   state->start_index = 1.5f + config->lower_band_limit / hz_per_sbin;
-  state->end_index = 0;  // Initialized to zero here, but actually set below.
+  state->end_index = 0; // Initialized to zero here, but actually set below.
 
   // For each channel, we need to figure out what frequencies belong to it, and
   // how much padding we need to add so that we can efficiently multiply the
@@ -96,18 +157,21 @@ int FilterbankPopulateState(const struct FilterbankConfig *config, struct Filter
   int needs_zeros = 0;
 
   int chan;
-  for (chan = 0; chan < num_channels_plus_1; ++chan) {
+  for (chan = 0; chan < num_channels_plus_1; ++chan)
+  {
     // Keep jumping frequencies until we overshoot the bound on this channel.
     int freq_index = chan_freq_index_start;
-    while (FreqToMel((freq_index) *hz_per_sbin) <= center_mel_freqs[chan]) {
+    while (FreqToMel((freq_index)*hz_per_sbin) <= center_mel_freqs[chan])
+    {
       ++freq_index;
     }
 
     const int width = freq_index - chan_freq_index_start;
     actual_channel_starts[chan] = chan_freq_index_start;
     actual_channel_widths[chan] = width;
 
-    if (width == 0) {
+    if (width == 0)
+    {
       // This channel doesn't actually get anything from the frequencies, it's
       // always zero. We need then to insert some 'zero' weights into the
       // output, and just redirect this channel to do a single multiplication at
@@ -117,15 +181,19 @@ int FilterbankPopulateState(const struct FilterbankConfig *config, struct Filter
       state->channel_frequency_starts[chan] = 0;
       state->channel_weight_starts[chan] = 0;
       state->channel_widths[chan] = kFilterbankChannelBlockSize;
-      if (!needs_zeros) {
+      if (!needs_zeros)
+      {
         needs_zeros = 1;
         int j;
-        for (j = 0; j < chan; ++j) {
+        for (j = 0; j < chan; ++j)
+        {
           state->channel_weight_starts[j] += kFilterbankChannelBlockSize;
         }
         weight_index_start += kFilterbankChannelBlockSize;
       }
-    } else {
+    }
+    else
+    {
       // How far back do we need to go to ensure that we have the proper
       // alignment?
       const int aligned_start = (chan_freq_index_start / index_alignment) * index_alignment;
@@ -143,11 +211,12 @@ int FilterbankPopulateState(const struct FilterbankConfig *config, struct Filter
   // Allocate the two arrays to store the weights - weight_index_start contains
   // the index of what would be the next set of weights that we would need to
   // add, so that's how many weights we need to allocate.
-  state->weights = (int16_t *) calloc(weight_index_start, sizeof(*state->weights));
-  state->unweights = (int16_t *) calloc(weight_index_start, sizeof(*state->unweights));
+  state->weights = (int16_t *)calloc(weight_index_start, sizeof(*state->weights));
+  state->unweights = (int16_t *)calloc(weight_index_start, sizeof(*state->unweights));
 
   // If the alloc failed, we also need to nuke the arrays.
-  if (state->weights == NULL || state->unweights == NULL) {
+  if (state->weights == NULL || state->unweights == NULL)
+  {
     free(center_mel_freqs);
     free(actual_channel_starts);
     free(actual_channel_widths);
@@ -159,38 +228,43 @@ int FilterbankPopulateState(const struct FilterbankConfig *config, struct Filter
   // zero, we only need to fill in the weights that correspond to some frequency
   // for a channel.
   const float mel_low = FreqToMel(config->lower_band_limit);
-  for (chan = 0; chan < num_channels_plus_1; ++chan) {
+  for (chan = 0; chan < num_channels_plus_1; ++chan)
+  {
     int frequency = actual_channel_starts[chan];
     const int num_frequencies = actual_channel_widths[chan];
     const int frequency_offset = frequency - state->channel_frequency_starts[chan];
     const int weight_start = state->channel_weight_starts[chan];
     const float denom_val = (chan == 0) ? mel_low : center_mel_freqs[chan - 1];
 
     int j;
-    for (j = 0; j < num_frequencies; ++j, ++frequency) {
+    for (j = 0; j < num_frequencies; ++j, ++frequency)
+    {
       const float weight =
           (center_mel_freqs[chan] - FreqToMel(frequency * hz_per_sbin)) / (center_mel_freqs[chan] - denom_val);
 
       // Make the float into an integer for the weights (and unweights).
       const int weight_index = weight_start + frequency_offset + j;
       QuantizeFilterbankWeights(weight, state->weights + weight_index, state->unweights + weight_index);
     }
-    if (frequency > state->end_index) {
+    if (frequency > state->end_index)
+    {
       state->end_index = frequency;
     }
   }
 
   free(center_mel_freqs);
   free(actual_channel_starts);
   free(actual_channel_widths);
-  if (state->end_index >= spectrum_size) {
+  if (state->end_index >= spectrum_size)
+  {
     fprintf(stderr, "Filterbank end_index is above spectrum size.\n");
     return 0;
   }
   return 1;
 }
 
-void FilterbankFreeStateContents(struct FilterbankState *state) {
+void FilterbankFreeStateContents(struct FilterbankState *state)
+{
   free(state->channel_frequency_starts);
   free(state->channel_weight_starts);
   free(state->channel_widths);