1 files changed, 54 insertions, 12 deletions

diff --git a/embassy-rp/src/pio/mod.rs b/embassy-rp/src/pio/mod.rs
index 5f554dfe3..1c370fdfc 100644
--- a/embassy-rp/src/pio/mod.rs
+++ b/embassy-rp/src/pio/mod.rs
@@ -2,21 +2,20 @@
 use core::future::Future;
 use core::marker::PhantomData;
 use core::pin::Pin as FuturePin;
-use core::sync::atomic::{compiler_fence, Ordering};
+use core::sync::atomic::{AtomicU8, AtomicU32, Ordering, compiler_fence};
 use core::task::{Context, Poll};
 
-use atomic_polyfill::{AtomicU64, AtomicU8};
 use embassy_hal_internal::{Peri, PeripheralType};
 use embassy_sync::waitqueue::AtomicWaker;
-use fixed::types::extra::U8;
 use fixed::FixedU32;
+use fixed::types::extra::U8;
 use pio::{Program, SideSet, Wrap};
 
 use crate::dma::{self, Channel, Transfer, Word};
 use crate::gpio::{self, AnyPin, Drive, Level, Pull, SealedPin, SlewRate};
 use crate::interrupt::typelevel::{Binding, Handler, Interrupt};
 use crate::relocate::RelocatedProgram;
-use crate::{pac, peripherals, RegExt};
+use crate::{RegExt, pac, peripherals};
 
 mod instr;
 
@@ -62,7 +61,7 @@ pub enum FifoJoin {
     #[cfg(feature = "_rp235x")]
     RxAsControl,
     /// FJOIN_RX_PUT | FJOIN_RX_GET: RX can be used as a scratch register,
-    /// not accesible from the CPU
+    /// not accessible from the CPU
     #[cfg(feature = "_rp235x")]
     PioScratch,
 }
@@ -984,11 +983,7 @@ impl<'d, PIO: Instance + 'd, const SM: usize> StateMachine<'d, PIO, SM> {
 
     #[cfg(feature = "_rp235x")]
     fn pin_base() -> u8 {
-        if PIO::PIO.gpiobase().read().gpiobase() {
-            16
-        } else {
-            0
-        }
+        if PIO::PIO.gpiobase().read().gpiobase() { 16 } else { 0 }
     }
 
     /// Sets pin directions. This pauses the current state machine to run `SET` commands
@@ -1236,7 +1231,13 @@ impl<'d, PIO: Instance> Common<'d, PIO> {
             w.set_pde(false);
         });
         // we can be relaxed about this because we're &mut here and nothing is cached
-        PIO::state().used_pins.fetch_or(1 << pin.pin_bank(), Ordering::Relaxed);
+        critical_section::with(|_| {
+            let val = PIO::state().used_pins.load(Ordering::Relaxed);
+            PIO::state()
+                .used_pins
+                .store(val | 1 << pin.pin_bank(), Ordering::Relaxed);
+        });
+
         Pin {
             pin: pin.into(),
             pio: PhantomData::default(),
@@ -1408,6 +1409,42 @@ impl<'d, PIO: Instance> Pio<'d, PIO> {
     }
 }
 
+struct AtomicU64 {
+    upper_32: AtomicU32,
+    lower_32: AtomicU32,
+}
+
+impl AtomicU64 {
+    const fn new(val: u64) -> Self {
+        let upper_32 = (val >> 32) as u32;
+        let lower_32 = val as u32;
+
+        Self {
+            upper_32: AtomicU32::new(upper_32),
+            lower_32: AtomicU32::new(lower_32),
+        }
+    }
+
+    fn load(&self, order: Ordering) -> u64 {
+        let (upper, lower) = critical_section::with(|_| (self.upper_32.load(order), self.lower_32.load(order)));
+
+        let upper = (upper as u64) << 32;
+        let lower = lower as u64;
+
+        upper | lower
+    }
+
+    fn store(&self, val: u64, order: Ordering) {
+        let upper_32 = (val >> 32) as u32;
+        let lower_32 = val as u32;
+
+        critical_section::with(|_| {
+            self.upper_32.store(upper_32, order);
+            self.lower_32.store(lower_32, order);
+        });
+    }
+}
+
 /// Representation of the PIO state keeping a record of which pins are assigned to
 /// each PIO.
 // make_pio_pin notionally takes ownership of the pin it is given, but the wrapped pin
@@ -1422,7 +1459,12 @@ pub struct State {
 
 fn on_pio_drop<PIO: Instance>() {
     let state = PIO::state();
-    if state.users.fetch_sub(1, Ordering::AcqRel) == 1 {
+    let users_state = critical_section::with(|_| {
+        let val = state.users.load(Ordering::Acquire) - 1;
+        state.users.store(val, Ordering::Release);
+        val
+    });
+    if users_state == 1 {
         let used_pins = state.used_pins.load(Ordering::Relaxed);
         let null = pac::io::vals::Gpio0ctrlFuncsel::NULL as _;
         for i in 0..crate::gpio::BANK0_PIN_COUNT {