ESP32-S3: Support execute in place from PSRAM

[EliteTK]

Submission Checklist 📝

• (NA?) I have updated existing examples or added new ones (if applicable).
• I have used cargo xtask fmt-packages command to ensure that all changed code is formatted correctly.
• My changes were added to the CHANGELOG.md in the proper section.
• (NA?) I have added necessary changes to user code to the Migration Guide.
• (I believe so) My changes are in accordance to the esp-rs API guidelines
• I have read the CONTRIBUTING.md guide and followed its instructions.

Pull Request Details 📖

Description

This implementation mirrors how the ESP-IDF implementation of this feature (which is based on the Cache_Flash_To_SPIRAM_Copy rom function) works except it differs in a few key ways:

The ESP-IDF seems to map .text and .rodata into the first and second 128 cache pages respectively (although looking at the linker scripts, I'm not sure how, but a runtime check confirmed this seemed to be the case). This is reflected in how the Cache_Count_Flash_Pages, Cache_Flash_To_SPIRAM_Copy rom functions and the ESP-IDF code executing them works. The count function can only be made to count flash pages within the first 256 pages (of which there are 512 on the ESP32-S3). Likewise, the copy function will only copy flash pages which are mapped within the first 256 entries (across two calls). As the esp-hal handles mapping .text and .rodata differently, these ROM functions are technically not appropriate if more than 256 pages of flash (.text and .rodata combined) are in use by the application.

Additionally, the functions both contain bugs, one of which the IDF attempts to work around incorrectly, and the other which the IDF does not appear to be aware of. Details of these bugs can be found on the IDF issue/PR tracker (espressif/esp-idf#15262, espressif/esp-idf#15263).

As a result, this PR contains a heavily modified/adjusted rust re-write of the reverse engineered ROM code combined with a vague port of the ESP-IDF code.

There are three additional noteworthy differences from the ESP-IDF version of the code:

1. The ESP-IDF allows the .text and .rodata segments to be mapped independently and separately allowing only one to be mapped. But the current version of the code does not allow this flexibility. This can be implemented by checking the address of each page entry against the segment locations to determine which segment each address belongs to.
2. The ESP-IDF calls cache_ll_l1_enable_bus(..., cache_ll_l1_get_bus(..., SOC_EXTRAM_DATA_HIGH, 0)); (functions from the ESP-IDF) in order to "Enable the most high bus, which is used for copying FLASH .text to PSRAM" but on the ESP32-S3 after careful inspection these calls result in a no-op as the address passed to cache_ll_l1_get_bus will result in an empty cache bus mask. It's currently unclear to me if this is a bug in the ESP-IDF code, or if this code (which from cursory investigation is probably not a no-op on the -S2) is solely targetting the ESP32-S3.
3. The ESP-IDF calls Cache_Flash_To_SPIRAM_Copy with an icache address when copying .text and a dcache address when copying .rodata. This affects which cache the reads will occur through. But the writes always go through a "spare page" (name I came up with during reverse engineering) via the dcache. This code performs all reads through the dcache. I don't know if there's a proper reason to read through the correct cache when doing the copy and this doesn't appear to have any negative impact.

Please let me know what this PR needs as I am not particulary experienced both with ESP32 development, embedded rust, unsafe rust, or this project.

Let me know if this change should have more documentation and where it's appropriate. As the psram stuff is unstable, I have assumed that it doesn't need to go in the migration guide, but I could be wrong about this.

Testing

I've done some basic testing of this version of the code on an ESP32-S3. I plan on doing some more testing (especially of the special handling of the spare page) but I felt like it was a good idea to get this on github for early comments.

[Dominaezzz]

I'm thinking the bulk of the code should go into a mmu.rs file.

[Dominaezzz]

Mind adding the size of both flash and psram to the panic string?

[bjoernQ]

minor, but maybe we can find a better name and doc-string

[EliteTK]

Yeah I am not sure about the name either, I just tried to come up with something slightly clearer than the ESP-IDF versions.

Options I've come up with as replacements:

• execute_in_place - The issue I have here is that, even though it seems like a common thing to say XiP from PSRAM, since this process involves copying pages to PSRAM, it's no longer really "in place" any more.
• map_flash
• mirror_flash
• replace_flash

Feel free to suggest something else. I'm also open to adding a mention of psram although I've so far avoided it given this field is in a "PsramConfig" struct.

Regarding the doc string. Would you prefer for it to be short, and for me to add a subheading to the module doc "Overview" section explaining the feature in more detail? Or would you prefer a multi line doc string with this information in line? For a short version I thought maybe: Copy and remap flash-mapped pages to PSRAM

Let me know what you think.

[bjoernQ]

esp-idf calls it

CONFIG_SPIRAM_XIP_FROM_PSRAM
Enable Executable in place from (XiP) from PSRAM feature

maybe s.th. like execute_from_psram and the description could mention that code and data is copied from flash to PSRAM and then memory gets remapped and code is executing from PSRAM - maybe link to https://docs.espressif.com/projects/esp-idf/en/stable/esp32s3/api-guides/external-ram.html#execute-in-place-xip-from-psram to avoid explaining why that might be useful to do

[bjoernQ]

This already looks very good from my point of view. Thanks for taking care of this

[EliteTK]

I'm thinking the bulk of the code should go into a mmu.rs file.

Do you think it's a good idea to move all the SOC_MMU_ constants into mmu.rs as well (further dropping the MMU_ prefix) and change the existing code references to refer to them?

I would just like to avoid duplicating the constants.

[Dominaezzz]

Yeah I'd say move the mmu constants as well

[EliteTK]

I've made all the requested changes.

Let me know if the explicitly limited visibility (pub(super)) is excessive and you want it removed.

Let me know if I've gone too far in the new MMU module refactoring.

I've also tried to make indices and sizes all usize instead of the previous mishmash of u32 and usize.

I have a few ideas for how the MMU module could be better but I didn't want to get too carried away (or more accurately, I already got too carried away on another branch and then scrapped the idea).

I noticed a "bug" in the previous PSRAM code. To be fair I think actually it wouldn't be a bug but more just a confusing error message.

Don't hold back any nitpicks.

[EliteTK]

I thought I could get away with running lint on only esp-hal. Apparently not. Whoops.

[Dominaezzz]

This is looking really good, I'll try and give it a spin myself soon.

[Dominaezzz]

Thinking about this some more, given there's more than one way to cut this XiP cake, especially once we start looking at static linking in PSRAM, I think this boolean should be separate from the PsramConfig.
I'm currently thinking it should go in esp_hal::Config, then init_psram can take a bool to do the right thing.

[Dominaezzz]

Would be awesome to have some kind of Entry struct here like this.

#[repr(transparent)]
struct Entry(u32);

enum Access {
    Flash,
    SpiRam
}

impl Entry {
    const ACCESS_BIT: u32 = 1 << 15;
    const VALIDITY_BIT: u32 = 1 << 14;
    const DATA_MASK: u32 = 0x3FFF;

    pub const fn invalid() -> Self {
        Self(Self::VALIDITY_BIT)
    }

    pub const fn new(access: Access, page_idx: usize) -> Self {
        let mut result = 0;
        if let Access::SpiRam = access {
            result |= Self::ACCESS_BIT;
        }
        result |= (page_idx as u32) & Self::DATA_MASK;

        Self(result)
    }

    #[procmacros::ram]
    pub fn access(&self) -> Access {
        if self.0 & Self::ACCESS_BIT != 0 {
            Access::SpiRam
        } else {
            Access::Flash
        }
    }

    #[procmacros::ram]
    pub fn is_valid(&self) -> bool {
        self.0 & Self::VALIDITY_BIT != 0
    }

    #[procmacros::ram]
    pub fn page_idx(&self) -> usize {
        (self.0 & Self::DATA_MASK) as usize
    }

    #[procmacros::ram]
    pub fn set_valid(&mut self, valid: bool) {
        todo!()
    }
}

Then you can do stuff like this let mmu_table_ptr = DR_REG_MMU_TABLE as *const Entry;

[EliteTK]

I thought about repr(transparent) kinds of things but how do you ensure the reads/writes are volatile?

My approach to this was something more like this:

mod entry {
    const INVALID: u32 = 1 << 14;
    const VALID: u32 = 0;
    const TYPE: u32 = 1 << 15;
    const ACCESS_FLASH: u32 = 0;
    const ACCESS_SPIRAM: u32 = 1 << 15;
    const VALID_VAL_MASK: u32 = 0x3fff;

    #[derive(Copy, Clone, Debug)]
    pub enum Entry {
        Flash(PageNum),
        Psram(PageNum),
        Invalid,
    }

    impl From<u32> for Entry {
        fn from(value: u32) -> Self {
            match value & INVALID {
                VALID => match value & TYPE {
                    ACCESS_FLASH => Entry::Flash(PageNum(value & VALID_VAL_MASK)),
                    ACCESS_SPIRAM => Entry::Psram(PageNum(value & VALID_VAL_MASK)),
                },
                INVALID => Entry::Invalid,
            }
        }
    }

    impl From<Entry> for u32 {
        fn from(value: Entry) -> Self {
            match value {
                Flash(page_num) => VALID | ACCESS_FLASH | page_num.into(),
                Psram(page_num) => VALID | ACCESS_SPIRAM | page_num.into(),
                Invalid => INVALID,
            }
        }
    }
}

pub use entry::Entry;

#[derive(Copy, Clone, Debug)]
pub struct BadPageNumError;

pub struct Index(usize);

impl TryFrom<usize> for Index {
    type 
}

pub struct Mmu;

impl Mmu {
    // I actually have no idea if this is a "steal" thing or not. I don't know what I'm doing, on my to-do list was to read how you're supposed to do these things properly as I forgot
    unsafe pub fn steal() -> Self {
        Mmu
    }

    pub fn iter(&self) -> Iter {
        Iter::new()
    }

    // TODO
    //pub fn iter_mut(&mut self) -> IterMut {
    //    IterMut { entry: 0 }
    //}

   pub  fn get(&self, index: Index) -> Entry {
        unsafe {
            (DR_REG_MMU_TABLE as *const _)
            .add(index.0)
            .read_volatile()
            .into()
        }
    }

    // Unsafe as it's never really "safe" to move memory pages around
    // Maybe this should also do cache invalidation
    unsafe pub fn set(&mut self, index: Index, entry: Entry) {
        unsafe {
            (DR_REG_MMU_TABLE as *const _)
            .add(index.0)
            .write_volatile(entry.into())
        }
    }

    // unsafe pub fn set_range(&mut self, range: Range<Index>, ... TODO
}

struct Iter {
    entry: u32,
    end: u32,
}

impl Iter {
    fn new() -> Self {
        Self {
            entry: 0,
            end: ENTRY_NUM,
        }
    }
}

impl Iterator for Iter { /* ... I've actually written all of these, just trimmed them for brevity ... */ }
impl DoubleEndedIterator for Iter { /* ... */ }
impl ExactSizeIterator for Iter { /* ... */ }
impl FusedIterator for Iter {}

[Dominaezzz]

but how do you ensure the reads/writes are volatile?

read_volatile work with any type, like *const Entry for example.

I like your Entry enum though, makes it more obvious what's going on.
Not sure about the Mmu struct yet.

[EliteTK]

I had not considered directly using the *const Entry for some reason. My mind was focused on slice::from_raw_parts for some reason.

I also think the Mmu struct is too much (aside from the fact that I'm pretty sure I'm doing it wrong). That's why I scrapped it, but I hadn't considered using the Entry enum without the Mmu struct.

If I understood correctly, you're suggesting something like this?

//! Thin MMU bindings
//!
//! More general information about the MMU can be found here:
//! <https://docs.espressif.com/projects/esp-idf/en/stable/esp32s3/api-reference/system/mm.html#introduction>

const DBUS_VADDR_BASE: u32 = 0x3C000000;
const DR_REG_MMU_TABLE: u32 = 0x600C5000;
const ICACHE_MMU_SIZE: usize = 0x800;
const MAX_PADDR_PAGE_NUM: usize = 16384;

//pub(super) const ENTRY_ACCESS_FLASH: u32 = 0;
pub(super) const ENTRY_ACCESS_SPIRAM: u32 = 1 << 15;
//pub(super) const ENTRY_INVALID: u32 = 1 << 14;
pub(super) const PAGE_SIZE: usize = 0x10000;
pub(super) const TABLE_SIZE: usize = ICACHE_MMU_SIZE / core::mem::size_of::<u32>();

extern "C" {
    /// Set DCache mmu mapping.
    ///
    /// [`ext_ram`]: u32 ENTRY_ACCESS_FLASH for flash, ENTRY_ACCESS_SPIRAM for spiram, ENTRY_INVALID for invalid.
    /// [`vaddr`]: u32 Virtual address in CPU address space.
    /// [`paddr`]: u32 Physical address in external memory. Should be aligned by psize.
    /// [`psize`]: u32 Page size of DCache, in kilobytes. Should be 64 here.
    /// [`num`]: u32 Pages to be set.
    /// [`fixes`]: u32 0 for physical pages grow with virtual pages, other for virtual pages map to same physical page.
    pub(super) fn cache_dbus_mmu_set(
        ext_ram: u32,
        vaddr: u32,
        paddr: u32,
        psize: u32,
        num: u32,
        fixed: u32,
    ) -> i32;

    fn Cache_Invalidate_Addr(addr: u32, size: u32);
    fn Cache_WriteBack_All();
    fn rom_Cache_WriteBack_Addr(addr: u32, size: u32);
}

#[derive(Copy, Clone, Debug)]
pub struct BadPageNumError;

#[derive(Copy, Clone, Debug)]
pub struct PageNum(usize);

impl TryFrom<usize> for PageNum {
    type Error = BadPageNumError;
    fn try_from(value: usize) -> Result<Self, BadPageNumError> {
        if value < MAX_PADDR_PAGE_NUM {
            Ok(Self(value))
        } else {
            Err(BadPageNumError)
        }
    }
}

impl From<PageNum> for usize {
    fn from(value: PageNum) -> Self {
        value.0
    }
}

impl From<PageNum> for u32 {
    fn from(value: PageNum) -> Self {
        value.0 as u32
    }
}

mod entry {
    use super::*;

    const INVALID: u32 = 1 << 14;
    const VALID: u32 = 0;
    const TYPE: u32 = 1 << 15;
    const ACCESS_FLASH: u32 = 0;
    const ACCESS_SPIRAM: u32 = 1 << 15;
    const VALID_VAL_MASK: u32 = 0x3fff;

    #[derive(Copy, Clone, Debug)]
    pub enum Entry {
        Flash(PageNum),
        Psram(PageNum),
        Invalid,
    }

    impl From<u32> for Entry {
        fn from(value: u32) -> Self {
            match value & INVALID {
                VALID => match value & TYPE {
                    ACCESS_FLASH => Entry::Flash(PageNum((value & VALID_VAL_MASK) as usize)),
                    ACCESS_SPIRAM => Entry::Psram(PageNum((value & VALID_VAL_MASK) as usize)),
                    _ => unreachable!(), // Maybe this would be better as an if
                                         // statement?
                },
                INVALID => Entry::Invalid,
                _ => unreachable!(), // Ditto
            }
        }
    }

    impl From<Entry> for u32 {
        fn from(value: Entry) -> Self {
            match value {
                Entry::Flash(page_num) => VALID | ACCESS_FLASH | u32::from(page_num),
                Entry::Psram(page_num) => VALID | ACCESS_SPIRAM | u32::from(page_num),
                Entry::Invalid => INVALID,
            }
        }
    }
}

use entry::Entry;

#[procmacros::ram]
pub(super) fn last_mapped_index() -> Option<usize> {
    let mmu_table_ptr = DR_REG_MMU_TABLE as *const u32;
    (0..TABLE_SIZE).rev().find(|&i| {
        !matches!(
            unsafe { mmu_table_ptr.add(i).read_volatile() }.into(),
            Entry::Invalid
        )
    })
}

#[procmacros::ram]
pub(super) fn index_to_data_address(index: usize) -> u32 {
    DBUS_VADDR_BASE + (PAGE_SIZE * index) as u32
}

/// Count flash-mapped pages, de-duplicating mappings which refer to flash page
/// 0
#[procmacros::ram]
pub(super) fn count_effective_flash_pages() -> usize {
    let mmu_table_ptr = DR_REG_MMU_TABLE as *const u32;
    let mut page0_seen = false;
    let mut flash_pages = 0;
    for i in 0..(TABLE_SIZE - 1) {
        if let Entry::Flash(page_num) = unsafe { mmu_table_ptr.add(i).read_volatile() }.into() {
            if u32::from(page_num) == 0 {
                if page0_seen {
                    continue;
                }
                page0_seen = true;
            }
            flash_pages += 1;
        }
    }
    flash_pages
}

#[procmacros::ram]
unsafe fn move_flash_to_psram_with_spare(
    target_entry: usize,
    psram_page: PageNum,
    spare_entry: usize,
) {
    let mmu_table_ptr = DR_REG_MMU_TABLE as *mut u32;
    let target_entry_addr = DBUS_VADDR_BASE + (target_entry * PAGE_SIZE) as u32;
    let spare_entry_addr = DBUS_VADDR_BASE + (spare_entry * PAGE_SIZE) as u32;
    unsafe {
        mmu_table_ptr
            .add(spare_entry)
            .write_volatile(Entry::Psram(psram_page).into());
        Cache_Invalidate_Addr(spare_entry_addr, PAGE_SIZE as u32);
        core::ptr::copy_nonoverlapping(
            target_entry_addr as *const u8,
            spare_entry_addr as *mut u8,
            PAGE_SIZE,
        );
        rom_Cache_WriteBack_Addr(spare_entry_addr, PAGE_SIZE as u32);
        mmu_table_ptr
            .add(target_entry)
            .write_volatile(Entry::Psram(psram_page).into());
    }
}

/// Copy flash-mapped pages to PSRAM, copying flash-page 0 only once, and re-map
/// those pages to the PSRAM copies
#[procmacros::ram]
pub(super) unsafe fn copy_flash_to_psram_and_remap(free_page: PageNum) -> usize {
    let mmu_table_ptr = DR_REG_MMU_TABLE as *mut u32;

    const SPARE_PAGE: usize = TABLE_SIZE - 1;

    let spare_page_mapping = unsafe { mmu_table_ptr.add(SPARE_PAGE).read_volatile() };
    let mut page0_page = None;
    let mut psram_page = free_page;

    unsafe { Cache_WriteBack_All() };
    for i in 0..(TABLE_SIZE - 1) {
        let mapping = unsafe { mmu_table_ptr.add(i).read_volatile() }.into();
        let Entry::Flash(page_num) = mapping else {
            continue;
        };
        if u32::from(page_num) == 0 {
            match page0_page {
                Some(page) => {
                    unsafe {
                        mmu_table_ptr
                            .add(i)
                            .write_volatile(page)
                    };
                    continue;
                }
                None => page0_page = Some(Entry::Psram(psram_page).into()),
            }
        }
        unsafe { move_flash_to_psram_with_spare(i, psram_page, SPARE_PAGE) };
        // Makes me wonder if there's a better way (aside from implementing Add
        // etc)
        psram_page = PageNum::try_from(usize::from(psram_page) + 1).unwrap();
    }

    // Restore spare page mapping
    unsafe {
        mmu_table_ptr
            .add(SPARE_PAGE)
            .write_volatile(spare_page_mapping);
        Cache_Invalidate_Addr(index_to_data_address(SPARE_PAGE), PAGE_SIZE as u32);
    }

    // Special handling if the spare page was mapped to flash
    if let Entry::Flash(_) = spare_page_mapping.into() {
        unsafe {
            // We're running from ram so using the first page should not cause issues
            const SECOND_SPARE: usize = 0;
            let second_spare_mapping = mmu_table_ptr.add(SECOND_SPARE).read_volatile();

            move_flash_to_psram_with_spare(SPARE_PAGE, psram_page, SECOND_SPARE);

            // Restore spare page mapping
            mmu_table_ptr.add(SECOND_SPARE).write_volatile(second_spare_mapping);
            Cache_Invalidate_Addr(
                DBUS_VADDR_BASE + (SECOND_SPARE * PAGE_SIZE) as u32,
                PAGE_SIZE as u32,
            );
        }
        psram_page = PageNum::try_from(usize::from(psram_page) + 1).unwrap();
    }
    usize::from(psram_page) - usize::from(free_page)
}

(probably with more strategically placed procmacros::ram)

I like the Entry enum but it seems to come at a cost (either you're constantly dealing with the PageNum or if you don't use a PageNum then the From<Entry> for u32 has to become TryFrom<Entry> which has its own annoyances).

I can push this version if you think the trade-offs are worth it. Or maybe you have more ideas for how to make this cleaner. It's also possible I misunderstood. Let me know.