FFT HW Acceleration

[ramtej]

I am currently benchmarking some DSP routines on the ESP32 S3 platform in Rust. Several issues have already arisen, see #180.

Upon reading the 'ESP32-S3 Technical Reference Manual', it became apparent that the S3 platform implements some SIMD as well as DSP operations in hardware, such as EE.FFT.R2BF.S16 or EE.CMUL.S16. I would be willing to invest some time and and implement a hardware-accelerated FFT. I am more of a mathematician and do not know the ESP32 architecture well enough, so I need some support.

Would it therefore be possible for someone to guide me and show me where I need to start? I looked at the SHA HW acceleration and can understand most things, but not all.

Thanks, Jiri

[ramtej]

To me, it seems that the EPS DSP "C" examples are merely ASM optimized and do not use any DSP functions!? Is that the case?

See https://github.com/espressif/esp-dsp/blob/71514173b58b960173b40c4ade9d15d372770a74/modules/fft/float/dsps_fft2r_fc32_ae32_.S#L60

[igrr]

Since these extra instructions are all integer instructions, only fixed-point FFT uses them: https://github.com/espressif/esp-dsp/blob/71514173b58b960173b40c4ade9d15d372770a74/modules/fft/fixed/dsps_fft2r_sc16_aes3.S

[ramtej]

Yes, that makes sense. This means that the DSP primitives have already been spotted in the wild and not only described in the above paper.

My approach now would be to fork or extend e.g. https://gitlab.com/teskje/microfft-rs that similar to the SHA HW acceleration use the DSP primitives. Specifically, I would want to replace the radix-2 butterfly computation with the DSP HW functions.

[ramtej]

Another (maybe stupid) idea could be to use the inline assembler asm!{} macro. Especially because the S3 target supports beside the DSP also some nice SIMD operations.

The follow compiles and even runs on the xtensa platform:

    std::hint::black_box(unsafe {
        asm!("nop");
    });

Now there is a lot between the above code and the universe that finally runs my code, especially the LLVM. I have dug into the esp Rust branch a bit and it looks like at least something was done in that direction:

pub enum InlineAsmArch {
    X86,
    X86_64,
    Arm,
  ..
    Xtensa,
  ..
}

rust/compiler/rustc_target/src/asm/mod.rs

Line 201 in ed3726b

pub enum InlineAsmArch {

Is that enough? I had naively just used a random DSP mnemonic - the error message is "mnemonic unknown".

Does anyone have an idea on the topic?

Thanks,
Jiri

[zRedShift]

@ramtej AFAIK, LLVM plain old doesn't support/recognize the hundreds of DSP instructions, so to use inline assembly it would require either adding them all, or using the escaped binary opcode (but that also means specifying registers by number)

I wish support could be added, but it's a non-trivial amount of search and replace in the llvm-project codebase.

https://github.com/esp-rs/xtensa-lx-rt/blob/39256baa9ff78950f502262fdbd4bce77bb31e76/src/exception/assembly_lx6.rs#L424

@MabezDev showed me this "trick" a good while ago

[ramtej]

@zRedShift Is the ".byte 0x00, 0x30, 0x00" sequence an inline assembly instruction specified as a sequence of bytes, serving as a workaround when the desired assembly instruction isn't supported by the LLVM? Indeed, this is a neat trick; I'll give it a try.

Thanks!

[ramtej]

There is a nice crate that does the '.byte' encoding for the RISC-V V extension instructions - https://github.com/cryptape/rvv-encoder. It would be exciting to have something similar for the Xtensa extension instructions.

unsafe {
    xtensa_asm::asm!(
         ..
        "ee.cmul.s16	q3,q2,q1,3",
        ..
    );
}

[zRedShift]

@ramtej awesome find. This will also help with ESP32P4's custom RISC V DSP extensions, when it comes out.

I haven't investigated it yet (I was actually planning on just writing .S files and compiling/linking them), but I wonder, is there support for rur.*/wsr.* etc instructions in llvm (chapter 1.6.10 in the ESP32S3 technical reference, Processor Control Instructions). They are neccessary to manipulate special registers that control stuff like FFT width in fixed point mode.

[ramtej]

Yes, probably the easiest thing for now will be to just link in the .S files. I think with the rur.*/wsr.* it will be similar to the other instructions. I need (i)FFT for my application and therefore I try to get the maximum out of the S3. The DSP benchmarks are promising, but I need the functions on Rust level. Maybe it makes sense to develop an esp-dsp-rs crate for the current Xtensa and the future RISC V DSP extensions.

[zRedShift]

I've been warming up with some Xtensa LLVM backend contributions over the last few days. My end goal is fast (scalar or vector) DSP, for ESP32S3 in Rust, since I'm working on real time audio processing/encoding and I need to maximize the performance to cram in as much processing into the pipeline. This is obviously DCT/FFT/FIR etc. heavy, among other things.

So a few more scalar instruction PRs, and I will move on to adding the 128-bit registers and instructions, so that the inline assembly can be supported directly. It's been a dream of mine for about a year to do it, at the time the ESP32S3 techinical manual still didn't include the extensions, and reverse engineering bfd/xtensa-modules.c was a major pain, and by the time they released the instructions in the reference, the priorities shifted. Now is the time to do the work.

Maybe even add auto-vectorization support in the (far) future, but it's a daunting task since the instructions are pipelined, the cost tables need to be populated, and there's a huge amount of user registers that control the runtime behavior of the instructions.

[igrr]

Just in case, cc @sstefan1 who is planning to merge the initial support for ESP32-S3 DSP instructions into Espressif's LLVM fork soon.

[zRedShift]

Don't want to step on anyone's toes, if @sstefan1 has already started work on this, I won't pursue, unless I can somehow assist?

[sstefan1]

Currently we have all ESP32-S3 DSP instructions implemented in LLVM. All instructions are available in clang through clang's builtins, which translate to llvm intrinsics and then to appropriate instructions.

For example:
__builtin_xtensa_ee_vld_128_ip(1, data, 0); --> ee.vld.128.ip q1, a9, 0

This work should be merged soon.

I need to investigate how that should be done in rust, though. If anybody already knows, please let me know.

[zRedShift]

@sstefan1 Well, we don't currently have any xtensa intrinsics support, but it would go to stdarch/core_arch here. It lives out of tree, so will need to be forked by the esp-rs org and the .gitmodules at esp-rs/rust should point to it. Then they can be added just like in clang. This can be done separately, whenever, and not a blocker for initial support.

For inline assembly support, it's much simpler since most of the base work has already been done by @MabezDev here
All that needs to be done is add the qregs support, and the user regs (FFT_WIDTH, QACC_H_0, etc.).

If you can get a branch/PR running on espressif/llvm-project, I can start work on initial support/testing on Rust for this. I'm already working on those files, adding rust support for the clamps/minmax features based on this PR.

[ramtej]

Ok, it looks to me like there is enough incentive and brainpower to tackle the ESP+Rust+DSP challenge. How are we going to coordinate this? Who does what?