Fully implemented hemem classifier

src/classifiers/hemem.rs

@@ -1,20 +1,235 @@
 //! Hemem classifier
 
+// Modules
+pub mod memories;
+pub mod page_table;
+
+// Exports
+pub use self::{
+	memories::{Memories, Memory},
+	page_table::{Page, PagePtr, PageTable},
+};
+
 // Imports
-use crate::sim;
+use {
+	self::memories::MemIdx,
+	crate::{pin_trace, sim},
+	anyhow::Context,
+};
 
 /// Hemem classifier
-pub struct HeMem {}
+#[derive(Debug)]
+pub struct HeMem {
+	/// Config
+	config: Config,
+
+	/// Memories
+	memories: Memories,
+
+	/// Page table
+	page_table: PageTable,
+}
 
 impl HeMem {
 	/// Creates a hemem classifier
-	pub fn new() -> Self {
-		Self {}
+	pub fn new(config: Config, memories: Vec<Memory>) -> Self {
+		Self {
+			config,
+			memories: Memories::new(memories),
+			page_table: PageTable::new(),
+		}
+	}
+
+	/// Maps a page to the first available memory and returns it.
+	///
+	/// # Errors
+	/// Returns an error if unable to insert
+	///
+	/// # Panics
+	/// Panics if the page is already mapped.
+	pub fn map_page(&mut self, page_ptr: PagePtr) -> Result<(), anyhow::Error> {
+		if self.page_table.contains(page_ptr) {
+			panic!("Page is already mapped: {page_ptr:?}");
+		}
+
+		for (mem_idx, mem) in self.memories.iter_mut() {
+			// Try to reserve a page on this memory
+			match mem.reserve_page() {
+				// If we got it, add the page to the page table
+				Ok(()) => {
+					let page = Page::new(page_ptr, mem_idx);
+					self.page_table.insert(page).expect("Unable to insert unmapped page");
+					return Ok(());
+				},
+
+				// If we didn't manage to, go to the next page
+				Err(err) => {
+					tracing::trace!(?page_ptr, ?mem_idx, ?err, "Unable to reserve page on memory");
+					continue;
+				},
+			}
+		}
+
+		// If we got here, all memories were full
+		anyhow::bail!("All memories were full");
+	}
+
+	/// Cools a memory by (at most) `count` pages.
+	///
+	/// Returns the number of pages cooled.
+	///
+	/// # Panics
+	/// Panics if `mem_idx` is an invalid memory index
+	pub fn cool_memory(&mut self, mem_idx: MemIdx, count: usize) -> usize {
+		let mut cooled_pages = 0;
+		for page_ptr in self.page_table.coldest_pages(mem_idx, count) {
+			if self.cool_page(page_ptr).is_ok() {
+				cooled_pages += 1;
+			}
+		}
+
+		cooled_pages
+	}
+
+	/// Migrates a page, possibly cooling the destination if full.
+	///
+	/// # Errors
+	/// Returns an error if unable to migrate the page to `dst_mem_idx`.
+	///
+	/// # Panics
+	/// Panics if `page_ptr` isn't a mapped page.
+	/// Panics if `dst_mem_idx` is an invalid memory index.
+	pub fn migrate_page(&mut self, page_ptr: PagePtr, dst_mem_idx: MemIdx) -> Result<(), anyhow::Error> {
+		let page = self.page_table.get_mut(page_ptr).expect("Page wasn't in page table");
+		let src_mem_idx = page.mem_idx();
+
+		match self.memories.migrate_page(src_mem_idx, dst_mem_idx) {
+			// If we managed to, move the page's memory
+			Ok(()) => page.move_mem(dst_mem_idx),
+
+			// Else try to cool the destination memory first, then try again
+			Err(err) => {
+				tracing::trace!(
+					?src_mem_idx,
+					?dst_mem_idx,
+					?err,
+					"Unable to migrate page, cooling destination"
+				);
+
+				// TODO: Cool for more than just 1 page at a time?
+				let pages_cooled = self.cool_memory(dst_mem_idx, 1);
+				match pages_cooled > 0 {
+					true => self
+						.memories
+						.migrate_page(src_mem_idx, dst_mem_idx)
+						.expect("Just freed some pages when cooling"),
+					false => anyhow::bail!("Cooler memory is full, even after cooling it"),
+				}
+			},
+		}
+
+		Ok(())
+	}
+
+	/// Cools a page.
+	///
+	/// # Errors
+	/// Returns an error if unable to cool the page.
+	///
+	/// # Panics
+	/// Panics if `page_ptr` isn't a mapped page.
+	pub fn cool_page(&mut self, page_ptr: PagePtr) -> Result<(), anyhow::Error> {
+		let page = self.page_table.get_mut(page_ptr).expect("Page wasn't in page table");
+
+		// Get the new memory index in the slower memory
+		let dst_mem_idx = match self.memories.slower_memory(page.mem_idx()) {
+			Some(mem_idx) => mem_idx,
+			None => anyhow::bail!("Page is already in the slowest memory"),
+		};
+
+		// Then try to migrate it
+		self.migrate_page(page_ptr, dst_mem_idx)
+			.context("Unable to migrate page to slower memory")
+	}
+
+	/// Warms a page.
+	///
+	/// # Errors
+	/// Returns an error if unable to warm the page.
+	///
+	/// # Panics
+	/// Panics if `page_ptr` isn't a mapped page.
+	pub fn warm_page(&mut self, page_ptr: PagePtr) -> Result<(), anyhow::Error> {
+		let page = self.page_table.get_mut(page_ptr).expect("Page wasn't in page table");
+
+		// Get the new memory index in the faster memory
+		let src_mem_idx = page.mem_idx();
+		let dst_mem_idx = match self.memories.faster_memory(src_mem_idx) {
+			Some(mem_idx) => mem_idx,
+			None => anyhow::bail!("Page is already in the hottest memory"),
+		};
+
+		// Then try to migrate it
+		self.migrate_page(page_ptr, dst_mem_idx)
+			.context("Unable to migrate page to faster memory")
 	}
 }
 
 impl sim::Classifier for HeMem {
-	fn handle_trace(&mut self, trace: sim::Trace) {
-		tracing::trace!(?trace, "Received trace")
+	fn handle_trace(&mut self, trace: sim::Trace) -> Result<(), anyhow::Error> {
+		tracing::trace!(?trace, "Received trace");
+		let page_ptr = PagePtr::new(trace.record.addr);
+
+		// Map the page if it doesn't exist
+		if !self.page_table.contains(page_ptr) {
+			tracing::trace!(?page_ptr, "Mapping page");
+			self.map_page(page_ptr).context("Unable to map page")?;
+		};
+		let page = self.page_table.get_mut(page_ptr).expect("Page wasn't in page table");
+		let page_was_hot = page.is_hot(self.config.read_hot_threshold, self.config.write_hot_threshold);
+
+		// Register the access on the page
+		match trace.record.kind {
+			pin_trace::RecordAccessKind::Read => page.register_read_access(),
+			pin_trace::RecordAccessKind::Write => page.register_write_access(),
+		};
+
+		// If the page is over the threshold, cool all pages
+		if page.over_threshold(self.config.global_cooling_threshold) {
+			self.page_table.cool_all_pages();
+		}
+
+		// Finally check if it's still hot and adjust if necessary
+		let page = self.page_table.get_mut(page_ptr).expect("Page wasn't in page table");
+		let page_is_hot = page.is_hot(self.config.read_hot_threshold, self.config.write_hot_threshold);
+
+		// If the page isn't hot and it was hot, cool it
+		if !page_is_hot && page_was_hot {
+			tracing::trace!(?page_ptr, "Page is no longer hot, cooling it");
+			if let Err(err) = self.cool_page(page_ptr) {
+				tracing::trace!(?page_ptr, ?err, "Unable to cool page");
+			}
+		}
+
+		// If the page was cold and is now hot, head it
+		if page_is_hot && !page_was_hot {
+			tracing::trace!(?page_ptr, "Page is now hot, warming it");
+			if let Err(err) = self.warm_page(page_ptr) {
+				tracing::trace!(?page_ptr, ?err, "Unable to warm page");
+			}
+		}
+
+		Ok(())
 	}
 }
+
+/// Configuration
+#[derive(Clone, Debug)]
+pub struct Config {
+	// R/W hotness threshold
+	pub read_hot_threshold:  usize,
+	pub write_hot_threshold: usize,
+
+	/// Max threshold for global cooling
+	pub global_cooling_threshold: usize,
+}

src/classifiers/hemem/memories.rs

@@ -0,0 +1,150 @@
+//! Memories
+
+// Imports
+use crate::util::FemtoDuration;
+
+/// Memories.
+///
+/// Maintains an array of memories, ordered from fastest to slowest.
+#[derive(Clone, Debug)]
+pub struct Memories {
+	/// All memories
+	memories: Vec<Memory>,
+}
+
+impl Memories {
+	/// Creates all the memories from an iterator of memories.
+	///
+	/// Memories are expected to be ordered from fastest to slowest.
+	pub fn new(memories: impl IntoIterator<Item = Memory>) -> Self {
+		Self {
+			memories: memories.into_iter().collect(),
+		}
+	}
+
+	/// Returns an iterator over all memories from fastest to slowest
+	pub fn iter_mut(&mut self) -> impl Iterator<Item = (MemIdx, &mut Memory)> {
+		self.memories
+			.iter_mut()
+			.enumerate()
+			.map(|(idx, mem)| (MemIdx(idx), mem))
+	}
+
+	/// Migrates a page from `src` to `dst`
+	///
+	/// Returns `Err` if the source memory is empty or the destination memory is full.
+	///
+	/// # Panics
+	/// Panics if either `src` or `dst` are invalid memory indexes
+	pub fn migrate_page(&mut self, src: MemIdx, dst: MemIdx) -> Result<(), anyhow::Error> {
+		// Get the memories
+		let [src, dst] = match self.memories.get_many_mut([src.0, dst.0]) {
+			Ok(mems) => mems,
+			Err(_) => match src == dst {
+				true => return Ok(()),
+				_ => panic!("Source or destination memory indexes were invalid"),
+			},
+		};
+
+		// Ensure they're not empty/full
+		anyhow::ensure!(!src.is_empty(), "Source memory was empty");
+		anyhow::ensure!(!dst.is_full(), "Destination memory was full");
+
+		// Then move them
+		dst.reserve_page().expect("Unable to reserve after checking non-full");
+		src.release_page().expect("Unable to release after checking non-empty");
+
+		Ok(())
+	}
+
+	/// Returns the faster memory after `mem_idx`
+	pub fn faster_memory(&self, mem_idx: MemIdx) -> Option<MemIdx> {
+		match mem_idx.0 {
+			0 => None,
+			_ => Some(MemIdx(mem_idx.0 - 1)),
+		}
+	}
+
+	/// Returns the slower memory after `mem_idx`
+	pub fn slower_memory(&self, mem_idx: MemIdx) -> Option<MemIdx> {
+		match mem_idx.0 + 1 >= self.memories.len() {
+			true => None,
+			false => Some(MemIdx(mem_idx.0 + 1)),
+		}
+	}
+}
+
+/// Memory index
+#[derive(PartialEq, Eq, Clone, Copy, Debug)]
+pub struct MemIdx(usize);
+
+/// Memory
+#[derive(Clone, Debug)]
+pub struct Memory {
+	/// Name
+	_name: String,
+
+	// Page size/capacity
+	page_len:      usize,
+	page_capacity: usize,
+
+	// Latencies
+	_latencies: AccessLatencies,
+}
+
+impl Memory {
+	/// Creates a new memory
+	pub fn new(name: impl Into<String>, page_capacity: usize, latencies: AccessLatencies) -> Self {
+		Self {
+			_name: name.into(),
+			page_len: 0,
+			page_capacity,
+			_latencies: latencies,
+		}
+	}
+
+	/// Attempts to release a page on this memory
+	pub fn release_page(&mut self) -> Result<(), anyhow::Error> {
+		// Ensure we're not empty
+		anyhow::ensure!(!self.is_empty(), "Memory is empty");
+
+		// Then release the page
+		self.page_len -= 1;
+
+		Ok(())
+	}
+
+	/// Attempts to reserve a page on this memory
+	pub fn reserve_page(&mut self) -> Result<(), anyhow::Error> {
+		// Ensure we're not full
+		anyhow::ensure!(!self.is_full(), "Memory is full");
+
+		// Then reserve the page
+		self.page_len += 1;
+
+		Ok(())
+	}
+
+	/// Returns if this memory is empty
+	pub fn is_empty(&self) -> bool {
+		self.page_len == 0
+	}
+
+	/// Returns if this memory is full
+	pub fn is_full(&self) -> bool {
+		self.page_len >= self.page_capacity
+	}
+}
+
+/// Access latencies
+#[derive(Clone, Copy, Debug)]
+pub struct AccessLatencies {
+	/// Read latency
+	pub read: FemtoDuration,
+
+	/// Write latency
+	pub write: FemtoDuration,
+
+	/// Fault latency
+	pub fault: FemtoDuration,
+}

src/classifiers/hemem/page_table.rs

@@ -0,0 +1,189 @@
+//! Page table
+
+// Imports
+use {
+	super::memories::MemIdx,
+	itertools::Itertools,
+	std::collections::{btree_map, BTreeMap},
+};
+
+/// Page table
+#[derive(Debug)]
+pub struct PageTable {
+	/// All pages, by their address
+	// TODO: `HashMap` with custom hash? We don't use the order
+	pages: BTreeMap<PagePtr, Page>,
+
+	/// Current cooling clock tick
+	cooling_clock_tick: usize,
+}
+
+impl PageTable {
+	/// Creates an empty page table
+	pub fn new() -> Self {
+		Self {
+			pages:              BTreeMap::new(),
+			cooling_clock_tick: 0,
+		}
+	}
+
+	/// Returns if a page exists in this page table
+	pub fn contains(&self, page_ptr: PagePtr) -> bool {
+		self.pages.contains_key(&page_ptr)
+	}
+
+	/// Returns a page from this page table.
+	pub fn get_mut(&mut self, page_ptr: PagePtr) -> Option<&mut Page> {
+		// Try to get the page
+		let page = self.pages.get_mut(&page_ptr)?;
+
+		// Then cool it before returning
+		page.cool_accesses(self.cooling_clock_tick);
+		Some(page)
+	}
+
+	/// Inserts a new page into this page table.
+	///
+	/// # Errors
+	/// Returns an error if the page already exists
+	pub fn insert(&mut self, mut page: Page) -> Result<(), anyhow::Error> {
+		match self.pages.entry(page.ptr) {
+			btree_map::Entry::Vacant(entry) => {
+				// Note: We cool it before inserting to ensure that the page is up to date.
+				page.cool_accesses(self.cooling_clock_tick);
+				entry.insert(page);
+
+				Ok(())
+			},
+			btree_map::Entry::Occupied(_) => anyhow::bail!("Page already existed: {page:?}"),
+		}
+	}
+
+	/// Cools all pages
+	pub fn cool_all_pages(&mut self) {
+		// Note: Instead of increasing all pages at once, we simply increase
+		//       our cooling clock and then, when accessing a page, we update
+		//       the pages's clock tick to match ours.
+		self.cooling_clock_tick += 1;
+	}
+
+	/// Returns the `count` coldest pages in memory `mem_idx`
+	// TODO: Optimize this function? Runs in `O(N)` with all memories
+	pub fn coldest_pages(&mut self, mem_idx: MemIdx, count: usize) -> Vec<PagePtr> {
+		self.pages
+			.iter_mut()
+			.update(|(_, page)| page.cool_accesses(self.cooling_clock_tick))
+			.filter(|(_, page)| page.mem_idx == mem_idx)
+			.sorted_by_key(|(_, page)| page.temperature())
+			.map(|(&page_ptr, _)| page_ptr)
+			.take(count)
+			.collect()
+	}
+}
+
+/// Page
+#[derive(Clone, Copy, Debug)]
+pub struct Page {
+	/// Pointer
+	ptr: PagePtr,
+
+	/// Memory index
+	mem_idx: MemIdx,
+
+	// Read/Write accesses (adjusted)
+	adjusted_read_accesses:  usize,
+	adjusted_write_accesses: usize,
+
+	// Current cooling clock tick
+	cur_cooling_clock_tick: usize,
+}
+
+impl Page {
+	/// Creates a new page
+	pub fn new(ptr: PagePtr, mem_idx: MemIdx) -> Self {
+		Self {
+			ptr,
+			mem_idx,
+			adjusted_read_accesses: 0,
+			adjusted_write_accesses: 0,
+			cur_cooling_clock_tick: 0,
+		}
+	}
+
+	/// Returns the memory index of this page
+	pub fn mem_idx(&self) -> MemIdx {
+		self.mem_idx
+	}
+
+	/// Moves this page to `mem_idx`
+	pub fn move_mem(&mut self, mem_idx: MemIdx) {
+		self.mem_idx = mem_idx;
+	}
+
+	/// Registers a read access
+	pub fn register_read_access(&mut self) {
+		self.adjusted_read_accesses += 1;
+	}
+
+	/// Registers a write access
+	pub fn register_write_access(&mut self) {
+		self.adjusted_write_accesses += 1;
+	}
+
+	/// Returns if this page is hot
+	pub fn is_hot(&self, read_hot_threshold: usize, write_hot_threshold: usize) -> bool {
+		self.adjusted_read_accesses >= read_hot_threshold || self.adjusted_write_accesses >= write_hot_threshold
+	}
+
+	/// Returns this page's temperature
+	pub fn temperature(&self) -> usize {
+		// TODO: Tune this definition?
+		self.adjusted_read_accesses * 1 + self.adjusted_write_accesses * 2
+	}
+
+	/// Returns if either read or write accesses are over a threshold
+	pub fn over_threshold(&self, threshold: usize) -> bool {
+		self.adjusted_read_accesses >= threshold || self.adjusted_write_accesses >= threshold
+	}
+
+	/// Cools this page's accesses to match the global cooling clock
+	fn cool_accesses(&mut self, global_access_cooling_clock_tick: usize) {
+		assert!(self.cur_cooling_clock_tick <= global_access_cooling_clock_tick);
+
+		let offset = (global_access_cooling_clock_tick - self.cur_cooling_clock_tick).min(usize::BITS as usize - 1);
+		self.adjusted_read_accesses >>= offset;
+		self.adjusted_write_accesses >>= offset;
+		self.cur_cooling_clock_tick = global_access_cooling_clock_tick;
+	}
+}
+
+/// Page pointer.
+///
+/// Guaranteed to be page-aligned
+#[derive(PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Hash)]
+pub struct PagePtr(u64);
+
+impl std::fmt::Debug for PagePtr {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		f.debug_tuple("PagePtr")
+			.field(&format_args!("{:#010x}", self.0))
+			.finish()
+	}
+}
+
+impl PagePtr {
+	/// Page mask
+	pub const PAGE_MASK: u64 = (1 << 12 - 1);
+
+	/// Creates a page pointer from a `u64`.
+	///
+	/// Will truncate any bits below the page mask.
+	pub fn new(page: u64) -> Self {
+		Self(page & !Self::PAGE_MASK)
+	}
+
+	/// Returns the page pointer as a u64
+	pub fn _to_u64(self) -> u64 {
+		self.0
+	}
+}

src/main.rs

@@ -14,7 +14,7 @@ mod util;
 // Imports
 use {
 	self::{args::Args, pin_trace::PinTrace},
-	crate::sim::Simulator,
+	crate::{classifiers::hemem, sim::Simulator, util::FemtoDuration},
 	anyhow::Context,
 	clap::Parser,
 	std::fs,
@@ -37,9 +37,28 @@ fn main() -> Result<(), anyhow::Error> {
 
 	// Run the simulator
 	let mut sim = Simulator::new(0);
-	let mut hemem_classifier = classifiers::hemem::HeMem::new();
+	let mut hemem = hemem::HeMem::new(
+		hemem::Config {
+			read_hot_threshold:       8,
+			write_hot_threshold:      4,
+			global_cooling_threshold: 18,
+		},
+		vec![
+			hemem::Memory::new("ram", 100, hemem::memories::AccessLatencies {
+				read:  FemtoDuration::from_nanos_f64(1.5),
+				write: FemtoDuration::from_nanos_f64(1.0),
+				fault: FemtoDuration::from_nanos_f64(10.0),
+			}),
+			hemem::Memory::new("optane", 800, hemem::memories::AccessLatencies {
+				read:  FemtoDuration::from_nanos_f64(5.0),
+				write: FemtoDuration::from_nanos_f64(4.0),
+				fault: FemtoDuration::from_nanos_f64(50.0),
+			}),
+		],
+	);
 
-	sim.run(pin_trace.records.iter().copied(), &mut hemem_classifier);
+	sim.run(pin_trace.records.iter().copied(), &mut hemem)
+		.context("Unable to run simulator")?;
 
 	Ok(())
 }

src/sim.rs

@@ -1,9 +1,10 @@
 //! Simulator
 
 // Imports
-use crate::pin_trace;
+use {crate::pin_trace, anyhow::Context};
 
 /// Simulator
+#[derive(Debug)]
 pub struct Simulator {
 	/// Trace skip
 	///
@@ -20,11 +21,19 @@ impl Simulator {
 	}
 
 	/// Runs the simulator on records `records` with classifier `classifier`
-	pub fn run<C: Classifier>(&mut self, records: impl IntoIterator<Item = pin_trace::Record>, classifier: &mut C) {
+	pub fn run<C: Classifier>(
+		&mut self,
+		records: impl IntoIterator<Item = pin_trace::Record>,
+		classifier: &mut C,
+	) -> Result<(), anyhow::Error> {
 		for record in records.into_iter().step_by(self.trace_skip + 1) {
 			let trace = Trace { record };
-			classifier.handle_trace(trace);
+			classifier
+				.handle_trace(trace)
+				.context("Unable to handle trace with classifier")?;
 		}
+
+		Ok(())
 	}
 }
 
@@ -32,7 +41,7 @@ impl Simulator {
 /// Classifier
 pub trait Classifier {
 	/// Handles a trace
-	fn handle_trace(&mut self, trace: Trace);
+	fn handle_trace(&mut self, trace: Trace) -> Result<(), anyhow::Error>;
 }
 
 /// Trace