perf(virtio): use weaker memory barriers from the mem-barrier crate

Cargo.toml

@@ -63,7 +63,7 @@ tcp = ["net", "smoltcp", "smoltcp/socket-tcp"]
 trace = ["smoltcp?/log", "smoltcp?/verbose"]
 udp = ["net", "smoltcp", "smoltcp/socket-udp"]
 vga = []
-virtio = ["dep:virtio"]
+virtio = ["dep:virtio", "dep:mem-barrier"]
 virtio-net = ["net", "virtio"]
 vsock = ["virtio", "pci"]
 warn-prebuilt = []
@@ -123,6 +123,7 @@ hermit-entry = { version = "0.10.6", features = ["kernel"] }
 hermit-sync = "0.1"
 lock_api = "0.4"
 log = { version = "0.4", default-features = false }
+mem-barrier = { version = "0.1.0", features = ["nightly"], optional = true }
 num_enum = { version = "0.7", default-features = false }
 pci-ids = { version = "0.2", optional = true }
 pci_types = { version = "0.10" }

src/arch/aarch64/mod.rs

@@ -1,11 +1,2 @@
-use aarch64_cpu::asm::barrier::{ISH, dmb};
-
 pub mod kernel;
 pub mod mm;
-
-/// Force strict CPU ordering, serializes load and store operations.
-#[allow(dead_code)]
-#[inline(always)]
-pub(crate) fn memory_barrier() {
-	dmb(ISH);
-}

src/arch/riscv64/mod.rs

@@ -1,8 +1,2 @@
 pub mod kernel;
 pub mod mm;
-
-#[allow(dead_code)]
-#[inline(always)]
-pub(crate) fn memory_barrier() {
-	riscv::asm::sfence_vma_all();
-}

src/arch/x86_64/mod.rs

@@ -22,13 +22,3 @@ pub(crate) fn swapgs(stack_frame: &ExceptionStackFrame) {
 #[cfg(not(feature = "common-os"))]
 #[inline(always)]
 pub(crate) fn swapgs(_stack_frame: &ExceptionStackFrame) {}
-
-/// Force strict CPU ordering, serializes load and store operations.
-#[allow(dead_code)]
-#[inline(always)]
-pub(crate) fn memory_barrier() {
-	use core::arch::asm;
-	unsafe {
-		asm!("mfence", options(nostack, nomem, preserves_flags));
-	}
-}

src/drivers/virtio/transport/pci.rs

@@ -13,6 +13,7 @@ use alloc::vec::Vec;
 use core::ptr::NonNull;
 use core::{mem, ptr};
 
+use mem_barrier::{BarrierKind, BarrierType, mem_barrier};
 use memory_addresses::PhysAddr;
 use pci_types::capability::PciCapability;
 use virtio::pci::{
@@ -23,7 +24,6 @@ use virtio::{DeviceStatus, le16, le32};
 use volatile::access::ReadOnly;
 use volatile::{VolatilePtr, VolatileRef};
 
-use crate::arch::memory_barrier;
 use crate::arch::pci::PciConfigRegion;
 #[cfg(feature = "console")]
 use crate::drivers::console::VirtioConsoleDriver;
@@ -299,7 +299,7 @@ impl ComCfg {
 
 	/// Resets the device status field to zero.
 	pub fn reset_dev(&mut self) {
-		memory_barrier();
+		mem_barrier(BarrierKind::Mmio, BarrierType::Write);
 		self.com_cfg
 			.as_mut_ptr()
 			.device_status()
@@ -310,7 +310,7 @@ impl ComCfg {
 	/// A driver MUST NOT initialize and use the device any further after this.
 	/// A driver MAY use the device again after a proper reset of the device.
 	pub fn set_failed(&mut self) {
-		memory_barrier();
+		mem_barrier(BarrierKind::Mmio, BarrierType::Write);
 		self.com_cfg
 			.as_mut_ptr()
 			.device_status()
@@ -320,32 +320,29 @@ impl ComCfg {
 	/// Sets the ACKNOWLEDGE bit in the device status field. This indicates, the
 	/// OS has notived the device
 	pub fn ack_dev(&mut self) {
-		memory_barrier();
-		self.com_cfg
-			.as_mut_ptr()
-			.device_status()
-			.update(|s| s | DeviceStatus::ACKNOWLEDGE);
+		self.com_cfg.as_mut_ptr().device_status().update(|s| {
+			mem_barrier(BarrierKind::Mmio, BarrierType::General);
+			s | DeviceStatus::ACKNOWLEDGE
+		});
 	}
 
 	/// Sets the DRIVER bit in the device status field. This indicates, the OS
 	/// know how to run this device.
 	pub fn set_drv(&mut self) {
-		memory_barrier();
-		self.com_cfg
-			.as_mut_ptr()
-			.device_status()
-			.update(|s| s | DeviceStatus::DRIVER);
+		self.com_cfg.as_mut_ptr().device_status().update(|s| {
+			mem_barrier(BarrierKind::Mmio, BarrierType::General);
+			s | DeviceStatus::DRIVER
+		});
 	}
 
 	/// Sets the FEATURES_OK bit in the device status field.
 	///
 	/// Drivers MUST NOT accept new features after this step.
 	pub fn features_ok(&mut self) {
-		memory_barrier();
-		self.com_cfg
-			.as_mut_ptr()
-			.device_status()
-			.update(|s| s | DeviceStatus::FEATURES_OK);
+		self.com_cfg.as_mut_ptr().device_status().update(|s| {
+			mem_barrier(BarrierKind::Mmio, BarrierType::General);
+			s | DeviceStatus::FEATURES_OK
+		});
 	}
 
 	/// In order to correctly check feature negotiaten, this function
@@ -355,23 +352,19 @@ impl ComCfg {
 	/// Re-reads device status to ensure the FEATURES_OK bit is still set:
 	/// otherwise, the device does not support our subset of features and the device is unusable.
 	pub fn check_features(&self) -> bool {
-		memory_barrier();
-		self.com_cfg
-			.as_ptr()
-			.device_status()
-			.read()
-			.contains(DeviceStatus::FEATURES_OK)
+		let status = self.com_cfg.as_ptr().device_status().read();
+		mem_barrier(BarrierKind::Dma, BarrierType::Read);
+		status.contains(DeviceStatus::FEATURES_OK)
 	}
 
 	/// Sets the DRIVER_OK bit in the device status field.
 	///
 	/// After this call, the device is "live"!
 	pub fn drv_ok(&mut self) {
-		memory_barrier();
-		self.com_cfg
-			.as_mut_ptr()
-			.device_status()
-			.update(|s| s | DeviceStatus::DRIVER_OK);
+		self.com_cfg.as_mut_ptr().device_status().update(|s| {
+			mem_barrier(BarrierKind::Mmio, BarrierType::General);
+			s | DeviceStatus::DRIVER_OK
+		});
 	}
 
 	/// Returns the features offered by the device.
@@ -382,20 +375,22 @@ impl ComCfg {
 
 		// Indicate device to show high 32 bits in device_feature field.
 		// See Virtio specification v1.1. - 4.1.4.3
-		memory_barrier();
+		mem_barrier(BarrierKind::Mmio, BarrierType::Write);
 		device_feature_select.write(1.into());
-		memory_barrier();
 
 		// read high 32 bits of device features
+		mem_barrier(BarrierKind::Mmio, BarrierType::General);
 		let mut device_features = u64::from(device_feature.read().to_ne()) << 32;
 
 		// Indicate device to show low 32 bits in device_feature field.
 		// See Virtio specification v1.1. - 4.1.4.3
+		mem_barrier(BarrierKind::Mmio, BarrierType::General);
 		device_feature_select.write(0.into());
-		memory_barrier();
 
 		// read low 32 bits of device features
+		mem_barrier(BarrierKind::Mmio, BarrierType::General);
 		device_features |= u64::from(device_feature.read().to_ne());
+		mem_barrier(BarrierKind::Mmio, BarrierType::Read);
 
 		virtio::F::from_bits_retain(u128::from(device_features).into())
 	}
@@ -412,19 +407,20 @@ impl ComCfg {
 
 		// Indicate to device that driver_features field shows low 32 bits.
 		// See Virtio specification v1.1. - 4.1.4.3
-		memory_barrier();
+		mem_barrier(BarrierKind::Mmio, BarrierType::Write);
 		driver_feature_select.write(0.into());
-		memory_barrier();
 
 		// write low 32 bits of device features
+		mem_barrier(BarrierKind::Mmio, BarrierType::Write);
 		driver_feature.write(low.into());
 
 		// Indicate to device that driver_features field shows high 32 bits.
 		// See Virtio specification v1.1. - 4.1.4.3
+		mem_barrier(BarrierKind::Mmio, BarrierType::Write);
 		driver_feature_select.write(1.into());
-		memory_barrier();
 
 		// write high 32 bits of device features
+		mem_barrier(BarrierKind::Mmio, BarrierType::Write);
 		driver_feature.write(high.into());
 	}
 }

src/drivers/virtio/virtqueue/mod.rs

@@ -19,6 +19,7 @@ use core::mem::MaybeUninit;
 use core::{mem, ptr};
 
 use enum_dispatch::enum_dispatch;
+use mem_barrier::{BarrierKind, BarrierType};
 use smallvec::SmallVec;
 use virtio::{le32, le64, pvirtq, virtq};
 
@@ -27,6 +28,17 @@ use crate::drivers::virtio::virtqueue::packed::PackedVq;
 use crate::drivers::virtio::virtqueue::split::SplitVq;
 use crate::mm::device_alloc::DeviceAlloc;
 
+#[inline]
+pub fn virtio_mem_barrier(ty: BarrierType, order_platform: bool) {
+	let class = if order_platform {
+		BarrierKind::Smp
+	} else {
+		BarrierKind::Dma
+	};
+
+	mem_barrier::mem_barrier(class, ty);
+}
+
 // Public interface of Virtq
 
 /// The Virtq trait unifies access to the two different Virtqueue types

src/drivers/virtio/virtqueue/split.rs

@@ -12,6 +12,7 @@ use alloc::vec::Vec;
 use core::cell::UnsafeCell;
 use core::mem::{self, MaybeUninit};
 
+use mem_barrier::BarrierType;
 #[cfg(not(feature = "pci"))]
 use virtio::mmio::NotificationData;
 #[cfg(feature = "pci")]
@@ -25,7 +26,6 @@ use super::super::transport::pci::{ComCfg, NotifCfg, NotifCtrl};
 use super::error::VirtqError;
 use super::index_alloc::IndexAlloc;
 use super::{AvailBufferToken, BufferType, TransferToken, UsedBufferToken, Virtq, VirtqPrivate};
-use crate::arch::memory_barrier;
 use crate::mm::device_alloc::DeviceAlloc;
 
 struct DescrRing {
@@ -36,6 +36,7 @@ struct DescrRing {
 	descr_table_cell: Box<UnsafeCell<[MaybeUninit<virtq::Desc>]>, DeviceAlloc>,
 	avail_ring_cell: Box<UnsafeCell<virtq::Avail>, DeviceAlloc>,
 	used_ring_cell: Box<UnsafeCell<virtq::Used>, DeviceAlloc>,
+	order_platform: bool,
 }
 
 impl DescrRing {
@@ -88,14 +89,15 @@ impl DescrRing {
 		self.avail_ring_mut().ring_mut(true)[idx as usize % len] =
 			le16::from_ne(index.try_into().unwrap());
 
-		memory_barrier();
+		super::virtio_mem_barrier(BarrierType::Write, self.order_platform);
 		let next_idx = idx.wrapping_add(1);
 		self.avail_ring_mut().idx = next_idx.into();
 
 		Ok(next_idx)
 	}
 
 	fn try_recv(&mut self) -> Result<UsedBufferToken, VirtqError> {
+		super::virtio_mem_barrier(BarrierType::Read, self.order_platform);
 		if self.read_idx == self.used_ring().idx.to_ne() {
 			return Err(VirtqError::NoNewUsed);
 		}
@@ -123,7 +125,6 @@ impl DescrRing {
 			}
 		}
 
-		memory_barrier();
 		self.read_idx = self.read_idx.wrapping_add(1);
 		Ok(UsedBufferToken::from_avail_buffer_token(
 			tkn.buff_tkn,
@@ -291,6 +292,8 @@ impl SplitVq {
 		vq_handler.set_drv_ctrl_addr(DeviceAlloc.phys_addr_from(avail_ring_cell.as_mut()));
 		vq_handler.set_dev_ctrl_addr(DeviceAlloc.phys_addr_from(used_ring_cell.as_mut()));
 
+		let order_platform = features.contains(virtio::F::ORDER_PLATFORM);
+
 		let descr_ring = DescrRing {
 			read_idx: 0,
 			token_ring: core::iter::repeat_with(|| None)
@@ -302,6 +305,7 @@ impl SplitVq {
 			descr_table_cell,
 			avail_ring_cell,
 			used_ring_cell,
+			order_platform,
 		};
 
 		let mut notif_ctrl = NotifCtrl::new(notif_cfg.notification_location(&mut vq_handler));