2 files changed, 70 insertions, 12 deletions

diff --git a/embassy-rp/src/pio/instr.rs b/embassy-rp/src/pio/instr.rs
index b15d507de..304ddb20a 100644
--- a/embassy-rp/src/pio/instr.rs
+++ b/embassy-rp/src/pio/instr.rs
@@ -5,6 +5,10 @@ use crate::pio::{Instance, StateMachine};
 
 impl<'d, PIO: Instance, const SM: usize> StateMachine<'d, PIO, SM> {
     /// Set value of scratch register X.
+    ///
+    /// SAFETY: autopull enabled else it will write undefined data.
+    /// Make sure to have setup the according configuration see
+    /// [shift_out](crate::pio::Config::shift_out) and [auto_fill](crate::pio::ShiftConfig::auto_fill).
     pub unsafe fn set_x(&mut self, value: u32) {
         const OUT: u16 = InstructionOperands::OUT {
             destination: OutDestination::X,
@@ -16,6 +20,10 @@ impl<'d, PIO: Instance, const SM: usize> StateMachine<'d, PIO, SM> {
     }
 
     /// Get value of scratch register X.
+    ///
+    /// SAFETY: autopush enabled else it will read undefined data.
+    /// Make sure to have setup the according configurations see
+    /// [shift_in](crate::pio::Config::shift_in) and [auto_fill](crate::pio::ShiftConfig::auto_fill).
     pub unsafe fn get_x(&mut self) -> u32 {
         const IN: u16 = InstructionOperands::IN {
             source: InSource::X,
@@ -27,6 +35,10 @@ impl<'d, PIO: Instance, const SM: usize> StateMachine<'d, PIO, SM> {
     }
 
     /// Set value of scratch register Y.
+    ///
+    /// SAFETY: autopull enabled else it will write undefined data.
+    /// Make sure to have setup the according configuration see
+    /// [shift_out](crate::pio::Config::shift_out) and [auto_fill](crate::pio::ShiftConfig::auto_fill).
     pub unsafe fn set_y(&mut self, value: u32) {
         const OUT: u16 = InstructionOperands::OUT {
             destination: OutDestination::Y,
@@ -38,6 +50,10 @@ impl<'d, PIO: Instance, const SM: usize> StateMachine<'d, PIO, SM> {
     }
 
     /// Get value of scratch register Y.
+    ///
+    /// SAFETY: autopush enabled else it will read undefined data.
+    /// Make sure to have setup the according configurations see
+    /// [shift_in](crate::pio::Config::shift_in) and [auto_fill](crate::pio::ShiftConfig::auto_fill).
     pub unsafe fn get_y(&mut self) -> u32 {
         const IN: u16 = InstructionOperands::IN {
             source: InSource::Y,
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 {