github/sparrow-kata-full
1
0
Fork 0
mirror of https://github.com/AmbiML/sparrow-kata-full.git synced 2025-09-22 20:09:25 +00:00
Code Issues Projects Releases Wiki Activity

kata: Integrate Image Manager and WMMU changes

Browse Source 
This CL handles the integration of the Image Manager, which comes
concurrent with the necessary WMMU changes needed by the Image Manager.

The ML Coordinator now calls into Image Manager to make space, commit
images, and set the WMMU.

The MlCoordinator now first verifies that an image is valid by first making a pass through the section headers. It stores the two sizes that we're interested in per image: how big it is packed on flash, and how big it is unpacked in memory.

Known issues:
b/241799340: Refactor BundleImage to support unit testing
The writes to DMEM via the kata-vec-core crate was meant to be in image_manager.rs, but this interfered with the ability to run unit tests. We can refactor BundleImage to make this work.

b/241799866: Improve heap management
Right now I clear all of the "temporary data section" (bss, stack,
heap), but I suspect only the heap needs to be cleared. This needs more effort to check that that's correct, and track those locations.

Minor changes:
ImageId is used instead of (String, String) in the component.

Change-Id: I1505c6474fc60205323ce3bb13610fdac3702b89
GitOrigin-RevId: 5df9938a6cbd7ca5510ce8fcb500ce471f42b2cb
This commit is contained in:
Adam Jesionowski
2022-07-27 14:12:56 -07:00
committed by Sam Leffler
parent 50cd809320
commit fd7f31bcb2
15 changed files with 494 additions and 365 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
apps/system/components/MlCoordinator/fake-vec-core/src/lib.rs
View File

@@ -5,11 +5,17 @@ extern crate alloc;
use alloc::boxed::Box;
use kata_io::Read;
use kata_ml_shared::ModelSections;
use kata_ml_shared::Permission;
pub fn enable_interrupts(_enable: bool) {}
pub fn set_wmmu(_sections: &ModelSections) {}
pub fn set_wmmu_window(
_window_id: usize,
_start_address: usize,
_length: usize,
_permission: Permission,
) {
}
pub fn run() {}
@@ -32,7 +38,9 @@ pub fn clear_instruction_fault() {}
pub fn clear_data_fault() {}
pub fn clear_tcm() {}
pub fn clear_tcm(_addr: usize, _len: usize) {}
pub fn wait_for_clear_to_finish() {}
pub fn get_return_code() -> u32 { 0 }

1
apps/system/components/MlCoordinator/kata-ml-component/Cargo.toml
View File

@@ -10,6 +10,7 @@ kata-os-common = { path = "../../kata-os-common" }
kata-memory-interface = { path = "../../MemoryManager/kata-memory-interface" }
kata-ml-coordinator = { path = "../kata-ml-coordinator" }
kata-ml-interface = { path = "../kata-ml-interface" }
kata-ml-shared = { path = "../kata-ml-shared" }
kata-timer-interface = { path = "../../TimerService/kata-timer-interface" }
log = "0.4"
spin = "0.9"

26
apps/system/components/MlCoordinator/kata-ml-component/src/run.rs
View File

@@ -8,6 +8,7 @@ use cstr_core::CStr;
use kata_ml_coordinator::MLCoordinator;
use kata_ml_coordinator::ModelIdx;
use kata_ml_interface::MlCoordError;
use kata_ml_shared::ImageId;
use kata_os_common::camkes::Camkes;
use kata_timer_interface::*;
use log::error;
@@ -45,14 +46,17 @@ pub unsafe extern "C" fn run() {
unsafe fn validate_ids(
c_bundle_id: *const cstr_core::c_char,
c_model_id: *const cstr_core::c_char,
) -> Result<(String, String), MlCoordError> {
) -> Result<ImageId, MlCoordError> {
let bundle_id = CStr::from_ptr(c_bundle_id)
.to_str()
.map_err(|_| MlCoordError::InvalidBundleId)?;
let model_id = CStr::from_ptr(c_model_id)
.to_str()
.map_err(|_| MlCoordError::InvalidModelId)?;
Ok((String::from(bundle_id), String::from(model_id)))
Ok(ImageId {
bundle_id: String::from(bundle_id),
model_id: String::from(model_id),
})
}
#[no_mangle]
@@ -60,12 +64,12 @@ pub unsafe extern "C" fn mlcoord_oneshot(
c_bundle_id: *const cstr_core::c_char,
c_model_id: *const cstr_core::c_char,
) -> MlCoordError {
let (bundle_id, model_id) = match validate_ids(c_bundle_id, c_model_id) {
Ok(ids) => ids,
let id = match validate_ids(c_bundle_id, c_model_id) {
Ok(id) => id,
Err(e) => return e,
};
if let Err(e) = ML_COORD.lock().oneshot(bundle_id, model_id) {
if let Err(e) = ML_COORD.lock().oneshot(id) {
return e;
}
@@ -78,11 +82,11 @@ pub unsafe extern "C" fn mlcoord_periodic(
c_model_id: *const cstr_core::c_char,
rate_in_ms: u32,
) -> MlCoordError {
let (bundle_id, model_id) = match validate_ids(c_bundle_id, c_model_id) {
Ok(ids) => ids,
let id = match validate_ids(c_bundle_id, c_model_id) {
Ok(id) => id,
Err(e) => return e,
};
if let Err(e) = ML_COORD.lock().periodic(bundle_id, model_id, rate_in_ms) {
if let Err(e) = ML_COORD.lock().periodic(id, rate_in_ms) {
return e;
}
@@ -94,12 +98,12 @@ pub unsafe extern "C" fn mlcoord_cancel(
c_bundle_id: *const cstr_core::c_char,
c_model_id: *const cstr_core::c_char,
) -> MlCoordError {
let (bundle_id, model_id) = match validate_ids(c_bundle_id, c_model_id) {
Ok(ids) => ids,
let id = match validate_ids(c_bundle_id, c_model_id) {
Ok(id) => id,
Err(e) => return e,
};
if let Err(e) = ML_COORD.lock().cancel(bundle_id, model_id) {
if let Err(e) = ML_COORD.lock().cancel(&id) {
return e;
}

3
apps/system/components/MlCoordinator/kata-ml-coordinator/Cargo.toml
View File

@@ -6,10 +6,13 @@ edition = "2021"
[dependencies]
cstr_core = { version = "0.2.3", default-features = false }
kata-io = { path = "../../DebugConsole/kata-io" }
kata-os-common = { path = "../../kata-os-common" }
kata-memory-interface = { path = "../../MemoryManager/kata-memory-interface" }
kata-ml-interface = { path = "../kata-ml-interface" }
kata-ml-shared = { path = "../kata-ml-shared" }
kata-ml-support= { path = "../kata-ml-support" }
kata-proc-interface = { path = "../../ProcessManager/kata-proc-interface" }
kata-security-interface = { path = "../../SecurityCoordinator/kata-security-interface" }
kata-timer-interface = { path = "../../TimerService/kata-timer-interface" }
kata-vec-core = { path = "../kata-vec-core" }

317
apps/system/components/MlCoordinator/kata-ml-coordinator/src/lib.rs
View File

@@ -4,22 +4,24 @@
extern crate alloc;
use alloc::string::String;
use alloc::vec::Vec;
use kata_memory_interface::kata_object_free_in_cnode;
use kata_ml_interface::MlCoordError;
use kata_ml_shared::MAX_MODELS;
use kata_ml_shared::*;
use kata_ml_support::image_manager::ImageManager;
use kata_os_common::cspace_slot::CSpaceSlot;
use kata_proc_interface::BundleImage;
use kata_security_interface::*;
use kata_timer_interface::*;
use kata_vec_core as MlCore;
use log::{error, info, trace, warn};
use log::{error, info, warn};
/// Represents a single loadable model.
#[derive(Debug)]
struct LoadableModel {
bundle_id: String,
model_id: String,
id: ImageId,
on_flash_sizes: ImageSizes,
in_memory_sizes: ImageSizes,
rate_in_ms: Option<u32>,
}
@@ -33,15 +35,15 @@ struct Statistics {
pub struct MLCoordinator {
/// The currently running model index, if any.
running_model: Option<ModelIdx>,
/// The currently loaded model.
// NB: This will be removed once the WMMU allows for multiple models loaded
loaded_model: Option<ModelIdx>,
/// A list of all models that have been requested for oneshot or periodic
/// execution.
models: [Option<LoadableModel>; MAX_MODELS],
/// A queue of models that are ready for immediate execution on the vector
/// core, once the currently running model has finished.
execution_queue: Vec<ModelIdx>,
/// The image manager is responsible for tracking, loading, and unloading
/// images.
image_manager: ImageManager,
statistics: Statistics,
}
@@ -55,9 +57,9 @@ impl MLCoordinator {
pub const fn new() -> Self {
MLCoordinator {
running_model: None,
loaded_model: None,
models: [INIT_NONE; MAX_MODELS],
execution_queue: Vec::new(),
image_manager: ImageManager::new(),
statistics: Statistics {
load_failures: 0,
already_queued: 0,
@@ -69,72 +71,170 @@ impl MLCoordinator {
pub fn init(&mut self) {
MlCore::enable_interrupts(true);
self.execution_queue.reserve(MAX_MODELS);
self.image_manager.init();
}
/// Load a model by copying it into the Vector Core's TCM, if it's not
/// already been loaded.
fn load_model(&mut self, model_idx: ModelIdx) -> Result<(), MlCoordError> {
if self.loaded_model == Some(model_idx) {
trace!("Model already loaded, skipping load");
// Validates the image by ensuring it has all the required loadable
// sections and that it fits into the TCM. Returns a tuple of
// |(on_flash_sizes, in_memory_sizes)|.
fn validate_image(&self, id: &ImageId) -> Option<(ImageSizes, ImageSizes)> {
let mut container_slot = CSpaceSlot::new();
match kata_security_load_model(&id.bundle_id, &id.model_id, &container_slot) {
Ok(model_frames) => {
container_slot.release(); // NB: take ownership
let mut image = BundleImage::new(&model_frames);
let mut on_flash_sizes = ImageSizes::default();
let mut in_memory_sizes = ImageSizes::default();
while let Some(section) = image.next_section() {
match section.vaddr {
TEXT_VADDR => {
on_flash_sizes.text = section.fsize;
in_memory_sizes.text = round_up(section.msize, WMMU_PAGE_SIZE);
}
CONST_DATA_VADDR => {
on_flash_sizes.constant_data = section.fsize;
in_memory_sizes.constant_data = round_up(section.msize, WMMU_PAGE_SIZE);
}
MODEL_OUTPUT_VADDR => {
on_flash_sizes.model_output = section.fsize;
in_memory_sizes.model_output = round_up(section.msize, WMMU_PAGE_SIZE);
}
STATIC_DATA_VADDR => {
on_flash_sizes.static_data = section.fsize;
in_memory_sizes.static_data = round_up(section.msize, WMMU_PAGE_SIZE);
}
TEMP_DATA_VADDR => {
on_flash_sizes.temporary_data = section.fsize;
in_memory_sizes.temporary_data =
round_up(section.msize, WMMU_PAGE_SIZE);
}
vaddr => {
warn!("While validating found unexpected section at {}", vaddr);
}
}
}
if !in_memory_sizes.is_valid() {
error!("Image invalid, section missing: {:?}", in_memory_sizes);
return None;
}
if in_memory_sizes.total_size() > TCM_SIZE {
error!("Image too big to fit in TCM: {:?}", in_memory_sizes);
return None;
}
drop(image);
let _ = kata_object_free_in_cnode(&model_frames);
Some((on_flash_sizes, in_memory_sizes))
}
Err(status) => {
error!("Security Core error {:?}", status);
None
}
}
}
// If there is a next model in the queue, load it onto the vector core and
// start running. If there's already a running model, don't do anything.
fn schedule_next_model(&mut self) -> Result<(), MlCoordError> {
if self.running_model.is_some() || self.execution_queue.is_empty() {
return Ok(());
}
// Ensure we have a model at the passed index. This shouldn't error.
let model = self.models[model_idx]
.as_ref()
.ok_or(MlCoordError::LoadModelFailed)?;
let next_idx = self.execution_queue.remove(0);
let model = self.models[next_idx].as_ref().expect("Model get fail");
// Loads |model_id| associated with |bundle_id| from the
// SecurityCoordinator. The data are returned as unmapped
// page frames in a CNode container left in |container_slot|.
// To load the model into the vector core the pages must be
// mapped into the MlCoordinator's VSpace before being copied
// to their destination.
let mut container_slot = CSpaceSlot::new();
match kata_security_load_model(&model.bundle_id, &model.model_id, &container_slot) {
Ok(model_frames) => {
container_slot.release(); // NB: take ownership
let ret_status = match MlCore::load_image(&model_frames) {
Err(e) => {
error!("Load of {}:{} failed: {:?}", &model.bundle_id, &model.model_id, e);
// May have corrupted TCM.
self.loaded_model = None;
self.statistics.load_failures += 1;
Err(MlCoordError::LoadModelFailed)
if !self.image_manager.is_loaded(&model.id) {
// Loads |model_id| associated with |bundle_id| from the
// SecurityCoordinator. The data are returned as unmapped
// page frames in a CNode container left in |container_slot|.
// To load the model into the vector core the pages must be
// mapped into the MlCoordinator's VSpace before being copied
// to their destination.
let mut container_slot = CSpaceSlot::new();
match kata_security_load_model(&model.id.bundle_id, &model.id.model_id, &container_slot)
{
Ok(model_frames) => {
container_slot.release(); // NB: take ownership
let mut image = BundleImage::new(&model_frames);
// Ask the image manager to make enough room and get
// the address to write to.
let mut temp_top = self.image_manager.make_space(
model.in_memory_sizes.data_top_size(),
model.in_memory_sizes.temporary_data,
);
while let Some(section) = image.next_section() {
// TODO(jesionowski): Ensure these are in order.
if section.vaddr == TEXT_VADDR {
MlCore::write_image_part(
&mut image,
temp_top,
model.on_flash_sizes.text,
model.in_memory_sizes.text,
)
.ok_or(MlCoordError::LoadModelFailed)?;
temp_top += model.in_memory_sizes.text;
} else if section.vaddr == CONST_DATA_VADDR {
MlCore::write_image_part(
&mut image,
temp_top,
model.on_flash_sizes.constant_data,
model.in_memory_sizes.constant_data,
)
.ok_or(MlCoordError::LoadModelFailed)?;
temp_top += model.in_memory_sizes.constant_data;
} else if section.vaddr == MODEL_OUTPUT_VADDR {
// Don't load, but do skip.
temp_top += model.in_memory_sizes.model_output;
} else if section.vaddr == STATIC_DATA_VADDR {
MlCore::write_image_part(
&mut image,
temp_top,
model.on_flash_sizes.static_data,
model.in_memory_sizes.static_data,
)
.ok_or(MlCoordError::LoadModelFailed)?;
temp_top += model.in_memory_sizes.static_data;
}
}
Ok(sections) => {
info!("Load successful.");
MlCore::set_wmmu(&sections);
self.loaded_model = Some(model_idx);
Ok(())
}
};
let _ = kata_object_free_in_cnode(&model_frames);
ret_status
}
Err(e) => {
error!(
"LoadModel of bundle {}:{} failed: {:?}",
&model.bundle_id, &model.model_id, e
);
self.statistics.load_failures += 1;
Err(MlCoordError::LoadModelFailed)
info!("Load successful.");
// Inform the image manager the image has been written.
self.image_manager
.commit_image(model.id.clone(), model.in_memory_sizes);
drop(image);
let _ = kata_object_free_in_cnode(&model_frames);
}
Err(e) => {
error!(
"LoadModel of bundle {}:{} failed: {:?}",
&model.id.bundle_id, &model.id.model_id, e
);
self.statistics.load_failures += 1;
return Err(MlCoordError::LoadModelFailed);
}
}
}
}
/// If there is a next model in the queue, load it onto the core and start
/// running. If there's already a running model, don't do anything.
fn schedule_next_model(&mut self) -> Result<(), MlCoordError> {
if !self.running_model.is_some() && !self.execution_queue.is_empty() {
let next_idx = self.execution_queue.remove(0);
// If load model fails we won't try and re-queue this model.
// It's very unlikely for load errors to be transient, it should
// only happen in the case of a mal-formed model.
self.load_model(next_idx)?;
MlCore::run(); // Unhalt, start at default PC.
self.running_model = Some(next_idx);
}
// TODO(jesionowski): Investigate if we need to clear the entire
// temporary data section or just certain parts.
// TODO(jesionowski): When hardware clear is enabled, we should
// kick it off after the run instead.
self.image_manager.clear_temp_data();
self.image_manager.set_wmmu(&model.id);
self.running_model = Some(next_idx);
MlCore::run(); // Start core at default PC.
Ok(())
}
@@ -173,8 +273,7 @@ impl MLCoordinator {
// of that slot.
fn ready_model(
&mut self,
bundle_id: String,
model_id: String,
id: ImageId,
rate_in_ms: Option<u32>,
) -> Result<ModelIdx, MlCoordError> {
// Return None if all slots are full.
@@ -184,20 +283,39 @@ impl MLCoordinator {
.position(|m| m.is_none())
.ok_or(MlCoordError::NoModelSlotsLeft)?;
let (on_flash_sizes, in_memory_sizes) =
self.validate_image(&id).ok_or(MlCoordError::InvalidImage)?;
self.models[index] = Some(LoadableModel {
bundle_id,
model_id,
id,
on_flash_sizes,
in_memory_sizes,
rate_in_ms,
});
Ok(index)
}
/// Start a one-time model execution, to be executed immediately.
pub fn oneshot(&mut self, bundle_id: String, model_id: String) -> Result<(), MlCoordError> {
let model_idx = self.ready_model(bundle_id, model_id, None)?;
// Returns the index for model |id| if it exists.
fn get_model_index(&self, id: &ImageId) -> Option<ModelIdx> {
self.models.iter().position(|opti| {
if let Some(i) = opti {
i.id == *id
} else {
false
}
})
}
self.execution_queue.push(model_idx);
/// Start a one-time model execution, to be executed immediately.
pub fn oneshot(&mut self, id: ImageId) -> Result<(), MlCoordError> {
// Check if we've loaded this model already.
let idx = match self.get_model_index(&id) {
Some(idx) => idx,
None => self.ready_model(id, None)?,
};
self.execution_queue.push(idx);
self.schedule_next_model()?;
Ok(())
@@ -205,36 +323,25 @@ impl MLCoordinator {
/// Start a periodic model execution, to be executed immediately and
/// then every rate_in_ms.
pub fn periodic(
&mut self,
bundle_id: String,
model_id: String,
rate_in_ms: u32,
) -> Result<(), MlCoordError> {
let model_idx = self.ready_model(bundle_id, model_id, Some(rate_in_ms))?;
pub fn periodic(&mut self, id: ImageId, rate_in_ms: u32) -> Result<(), MlCoordError> {
// Check if we've loaded this model already.
let idx = match self.get_model_index(&id) {
Some(idx) => idx,
None => self.ready_model(id, Some(rate_in_ms))?,
};
self.execution_queue.push(model_idx);
self.execution_queue.push(idx);
self.schedule_next_model()?;
timer_service_periodic(model_idx as u32, rate_in_ms);
timer_service_periodic(idx as u32, rate_in_ms);
Ok(())
}
/// Cancels an outstanding execution.
pub fn cancel(&mut self, bundle_id: String, model_id: String) -> Result<(), MlCoordError> {
pub fn cancel(&mut self, id: &ImageId) -> Result<(), MlCoordError> {
// Find the model index matching the bundle/model id.
let model_idx = self
.models
.iter()
.position(|optm| {
if let Some(m) = optm {
m.bundle_id == bundle_id && m.model_id == model_id
} else {
false
}
})
.ok_or(MlCoordError::NoSuchModel)?;
let model_idx = self.get_model_index(id).ok_or(MlCoordError::NoSuchModel)?;
// If the model is periodic, cancel the timer.
if self.models[model_idx]
@@ -255,6 +362,8 @@ impl MLCoordinator {
self.execution_queue.remove(idx);
}
self.image_manager.unload_image(id);
self.models[model_idx] = None;
Ok(())
}
@@ -270,7 +379,7 @@ impl MLCoordinator {
let model = self.models[model_idx].as_ref().unwrap();
warn!(
"Dropping {}:{} periodic execution as it has an execution outstanding already.",
&model.bundle_id, &model.model_id
&model.id.bundle_id, &model.id.model_id
);
self.statistics.already_queued += 1;
return Ok(());
@@ -317,7 +426,7 @@ impl MLCoordinator {
fn ids_at(&self, idx: ModelIdx) -> (&str, &str) {
match self.models[idx].as_ref() {
Some(model) => (&model.bundle_id, &model.model_id),
Some(model) => (&model.id.bundle_id, &model.id.model_id),
None => ("None", "None"),
}
}
@@ -331,19 +440,9 @@ impl MLCoordinator {
None => info!("No running model."),
}
match self.loaded_model {
Some(idx) => {
let (bundle, model) = self.ids_at(idx);
info!("Loaded model: {}:{}", bundle, model);
}
None => info!("No loaded model."),
}
info!("Loadable Models:");
for model in self.models.as_ref() {
if let Some(m) = model {
info!(" {:?}", m);
}
for model in self.models.as_ref().iter().flatten() {
info!(" {:x?}", model);
}
info!("Execution Queue:");
@@ -353,5 +452,7 @@ impl MLCoordinator {
}
info!("Statistics: {:?}", self.statistics);
self.image_manager.debug_state();
}
}

2
apps/system/components/MlCoordinator/kata-ml-interface/src/lib.rs
View File

@@ -1,5 +1,4 @@
#![no_std]
#![allow(dead_code)]
use cstr_core::CString;
/// Errors that can occur when interacting with the MlCoordinator.
@@ -9,6 +8,7 @@ pub enum MlCoordError {
MlCoordOk,
InvalidModelId,
InvalidBundleId,
InvalidImage,
LoadModelFailed,
NoModelSlotsLeft,
NoSuchModel,

1
apps/system/components/MlCoordinator/kata-ml-shared/Cargo.toml
View File

@@ -4,3 +4,4 @@ version = "0.1.0"
edition = "2021"
[dependencies]
bitflags = "1.3.2"

69
apps/system/components/MlCoordinator/kata-ml-shared/src/lib.rs
View File

@@ -1,5 +1,4 @@
#![no_std]
#![allow(dead_code)]
// Data structures used throughout the Kata ML implementation that do not
// depend on kata-os-common.
@@ -7,6 +6,7 @@
extern crate alloc;
use alloc::string::String;
use bitflags::bitflags;
/// An image is uniquely identified by the bundle that owns it and the
/// particular model id in that bundle.
@@ -18,7 +18,7 @@ pub struct ImageId {
/// An image consists of five sections. See go/sparrow-vc-memory for a
/// description of each section. Sizes are in bytes.
#[derive(Clone, Debug, Default)]
#[derive(Clone, Copy, Debug, Default)]
pub struct ImageSizes {
pub text: usize,
pub model_input: usize,
@@ -33,24 +33,21 @@ impl ImageSizes {
pub fn data_top_size(&self) -> usize {
self.text + self.model_output + self.constant_data + self.static_data
}
}
// XXX: Out-dated and should use ImageSizes. Refactor when multiple sections
// are enabled.
/// The Vector Core uses a Windowed MMU (go/sparrow-wmmu) in order to prevent
/// models from interferring with each other. Before executing a model,
/// windows to only that model's code and data are opened.
/// A window is represented by an address and size of that window.
pub struct Window {
pub addr: usize,
pub size: usize,
}
pub fn total_size(&self) -> usize {
self.data_top_size() + self.temporary_data + self.model_input
}
// XXX: Out-dated. Refactor when multiple sections are enabled.
/// When a model is loaded onto the Vector Core, the ML Coordinator needs to
/// track where each window is.
pub struct ModelSections {
pub tcm: Window,
// A set of sizes is considered valid if everything but model_input is
// non-zero.
pub fn is_valid(&self) -> bool {
// TODO(jesionowski): Add `&& self.model_output != 0` when model output
// is integrated with model code.
self.text != 0
&& self.constant_data != 0
&& self.static_data != 0
&& self.temporary_data != 0
}
}
/// The page size of the WMMU.
@@ -66,3 +63,39 @@ pub const TCM_SIZE: usize = 0x1000000;
/// The address of the Vector Core's TCM, viewed from the SMC.
pub const TCM_PADDR: usize = 0x34000000;
// The virtualized address of each WMMU section (see: go/sparrow-vc-memory).
pub const TEXT_VADDR: usize = 0x80000000;
pub const CONST_DATA_VADDR: usize = 0x81000000;
pub const MODEL_OUTPUT_VADDR: usize = 0x82000000;
pub const STATIC_DATA_VADDR: usize = 0x83000000;
pub const MODEL_INPUT_VADDR: usize = 0x84000000;
pub const TEMP_DATA_VADDR: usize = 0x85000000;
#[derive(Clone, Copy, Debug)]
pub enum WindowId {
Text = 0,
ConstData = 1,
ModelOutput = 2,
StaticData = 3,
ModelInput = 4,
TempData = 5,
}
bitflags! {
pub struct Permission: u32 {
const READ = 0b00000001;
const WRITE = 0b00000010;
const EXECUTE = 0b00000100;
const READ_WRITE = Self::READ.bits | Self::WRITE.bits;
const READ_EXECUTE = Self::READ.bits | Self::EXECUTE.bits;
}
}
pub fn round_up(a: usize, b: usize) -> usize {
if (a % b) == 0 {
a
} else {
usize::checked_add(a, b).unwrap() - (a % b)
}
}

3
apps/system/components/MlCoordinator/kata-ml-support/Cargo.toml
View File

@@ -7,8 +7,7 @@ edition = "2021"
log = "0.4"
kata-io = { path = "../../DebugConsole/kata-io" }
kata-ml-shared = { path = "../kata-ml-shared" }
# XXX: Re-integrate when it does not depend on kata-os-common.
# kata-vec-core = { path = "../kata-vec-core" }
kata-vec-core = { path = "../kata-vec-core" }
fake-vec-core = { path = "../fake-vec-core" }
[dev-dependencies]

253
apps/system/components/MlCoordinator/kata-ml-support/src/image_manager.rs
View File

@@ -1,5 +1,3 @@
#![allow(dead_code)] // XXX: Supress warnings, remove once integrated.
// The Image Manager is responsible for loading and unloading multiple images
// into the Vector Core's tightly coupled memory. It tracks which image section
// is where and evicts images on the core when necessary.
@@ -28,21 +26,15 @@
extern crate alloc;
use alloc::boxed::Box;
use alloc::vec::Vec;
use core::cmp;
use kata_ml_shared::{ImageId, ImageSizes};
use kata_ml_shared::{MAX_MODELS, TCM_PADDR, TCM_SIZE, WMMU_PAGE_SIZE};
use kata_ml_shared::*;
use log::{info, trace};
use kata_io::Read;
// XXX: Enable configuration when kata_vec_core does not depend on
// kata-os-common.
// #[cfg(not(test))]
// use kata_vec_core as MlCore;
// #[cfg(test)]
#[cfg(test)]
use fake_vec_core as MlCore;
#[cfg(not(test))]
use kata_vec_core as MlCore;
// For each loaded image we need to track where the image's first segment is:
// data_top ---> +---------------+
@@ -98,7 +90,7 @@ const INIT_NONE: Option<Image> = None;
// | shared temp |
// | |
// +---------------+
struct ImageManager {
pub struct ImageManager {
images: [Option<Image>; MAX_MODELS],
image_queue: Vec<ImageIdx>,
@@ -107,33 +99,25 @@ struct ImageManager {
tcm_bottom: usize,
}
// TODO(jesionowski): Create kata-os-utils, move this to it.
fn round_up(a: usize, b: usize) -> usize {
if (a % b) == 0 {
a
} else {
usize::checked_add(a, b).unwrap() - (a % b)
}
}
// Returns the bytes needed above current_size to fit requested_size.
fn space_needed(current_size: usize, requested_size: usize) -> usize {
cmp::max(requested_size as isize - current_size as isize, 0) as usize
}
impl Default for ImageManager {
fn default() -> Self {
impl ImageManager {
pub const fn new() -> Self {
ImageManager {
images: [INIT_NONE; MAX_MODELS],
image_queue: Vec::with_capacity(MAX_MODELS),
image_queue: Vec::new(),
sensor_top: TCM_PADDR,
tcm_top: TCM_PADDR,
tcm_bottom: TCM_PADDR + TCM_SIZE,
}
}
}
impl ImageManager {
// Optional initilization step to reserve space for the queue.
pub fn init(&mut self) { self.image_queue.reserve(MAX_MODELS); }
// Allocate a block of memory for the SensorManager to use. Returns the
// address of the block. This function should only be called once during
// SensorManager initialization, before any images are loaded.
@@ -173,12 +157,6 @@ impl ImageManager {
// Only move data if the addresses are different.
if tcm_addr != image.data_top_addr {
trace!(
"Moving {:X} bytes from {:X} to {:X}",
size,
image.data_top_addr,
tcm_addr
);
MlCore::tcm_move(image.data_top_addr, tcm_addr, size);
image.data_top_addr = tcm_addr;
}
@@ -189,18 +167,21 @@ impl ImageManager {
self.tcm_top = tcm_addr;
}
// Remove the latest image loaded and return the size of the freed space.
// Removes the latest image loaded and return the size of the freed space.
fn unload_latest(&mut self) -> ImageSizes {
// We can assume there's an image in the queue and unwrap safely, as
// otherwise we wouldn't need to unload images to fit new ones.
let idx = self.image_queue.pop().unwrap();
let (bundle, model) = self.ids_at(idx);
info!("Unloading image {}:{}", bundle, model);
self.images[idx].take().unwrap().sizes
}
// Removes images in FILO order until the top TCM and temp TCM
// constraints are satisfied.
fn make_space(&mut self, top_tcm_needed: usize, temp_tcm_needed: usize) {
/// Removes images in FILO order until the top TCM and temp TCM
/// constraints are satisfied. Returns the address of the freed space.
pub fn make_space(&mut self, top_tcm_needed: usize, temp_tcm_needed: usize) -> usize {
assert!(top_tcm_needed + temp_tcm_needed <= TCM_SIZE);
let mut available_tcm = self.tcm_free_space();
let mut space_needed_for_temp =
@@ -211,6 +192,8 @@ impl ImageManager {
available_tcm += freed_sizes.data_top_size();
self.tcm_top -= freed_sizes.data_top_size();
// If we removed an image that had a temporary data size above the
// current temp data size, we add that new memory to the pool.
let remaining_temp = self.required_temporary_data();
@@ -221,31 +204,20 @@ impl ImageManager {
// Re-calculate space needed for temporary data given the new size.
space_needed_for_temp = space_needed(remaining_temp, temp_tcm_needed);
}
self.tcm_top
}
// Sets the size of the temporary section based on the remaining images.
fn set_tcm_bottom(&mut self) {
let temp_data_size = self.required_temporary_data();
self.tcm_bottom = TCM_PADDR + TCM_SIZE - temp_data_size;
}
// Updates the pointers after an image is written to TCM to ensure the
// image is kept around.
fn update_image_bookkeeping(&mut self, image: Image) {
// We expect to always have <32 models due to memory constraints,
// making this unwrap safe.
let index = self.images.iter().position(|i| i.is_none()).unwrap();
self.image_queue.push(index);
self.tcm_top += image.sizes.data_top_size();
self.set_tcm_bottom();
// If these pointers cross the memory is in an inconsistent state.
// (We shouldn't hit this unless our space calculations are wrong.)
assert!(self.tcm_bottom >= self.tcm_top);
self.images[index] = Some(image);
MlCore::set_wmmu_window(
WindowId::TempData,
self.tcm_bottom,
temp_data_size,
Permission::READ_WRITE,
);
}
// Returns the index for image |id| if it exists.
@@ -259,63 +231,43 @@ impl ImageManager {
})
}
// Returns true if the image is currently loaded in the TCM.
/// Returns true if the image |id| is currently loaded in the TCM.
pub fn is_loaded(&mut self, id: &ImageId) -> bool { self.get_image_index(id).is_some() }
// Loads an |image| onto the Vector Core's TCM, evicting models as
// necessary. |on_flash_sizes| represents the on-disk sizes (ie fsize)
// for each section, while |unpacked_sizes| represents the aligned memory
// needed for execution (ie msize).
pub fn load_image(
&mut self,
image: &mut Box<dyn Read>,
id: ImageId,
on_flash_sizes: ImageSizes,
unpacked_sizes: ImageSizes,
) -> Result<(), &'static str> {
self.make_space(unpacked_sizes.data_top_size(), unpacked_sizes.temporary_data);
let mut temp_top = self.tcm_top;
MlCore::write_image_part(image, temp_top, on_flash_sizes.text, unpacked_sizes.text)?;
temp_top += unpacked_sizes.text;
MlCore::write_image_part(
image,
temp_top,
on_flash_sizes.constant_data,
unpacked_sizes.constant_data,
)?;
temp_top += unpacked_sizes.constant_data;
MlCore::write_image_part(
image,
temp_top,
on_flash_sizes.model_output,
unpacked_sizes.model_output,
)?;
temp_top += unpacked_sizes.model_output;
MlCore::write_image_part(
image,
temp_top,
on_flash_sizes.static_data,
unpacked_sizes.static_data,
)?;
// Commit the image and update pointers.
self.update_image_bookkeeping(Image {
/// Appends an (already written) image to internal book-keeping. This class
/// does not handle the write as it requires seL4 references. The
/// MlCoordinator must call this function after the write.
pub fn commit_image(&mut self, id: ImageId, sizes: ImageSizes) {
let image = Image {
id,
sizes: unpacked_sizes,
sizes,
data_top_addr: self.tcm_top,
});
};
Ok(())
// We expect to always have <32 models due to memory constraints,
// making this unwrap safe.
let index = self.images.iter().position(|i| i.is_none()).unwrap();
trace!("Adding image: {:x?}", image);
self.image_queue.push(index);
self.tcm_top += image.sizes.data_top_size();
self.images[index] = Some(image);
self.set_tcm_bottom();
// If these pointers cross the memory is in an inconsistent state.
// (We shouldn't hit this unless our space calculations are wrong.)
assert!(self.tcm_bottom >= self.tcm_top);
}
// Unloads image |id| if loaded. Returns true if an image was unloaded.
pub fn unload_image(&mut self, id: &ImageId) -> bool {
if let Some(idx) = self.get_image_index(id) {
self.image_queue.remove(idx);
self.images[idx] = None;
self.compact_tcm_top();
self.set_tcm_bottom();
return true;
@@ -324,6 +276,59 @@ impl ImageManager {
false
}
/// Sets the WMMU to match the loaded image |id|. Returns true if that
/// image exists and the WMMU was set.
pub fn set_wmmu(&self, id: &ImageId) -> bool {
if let Some(idx) = self.get_image_index(id) {
let image = &self.images[idx].as_ref().unwrap();
let mut top = image.data_top_addr;
MlCore::set_wmmu_window(
WindowId::Text,
top,
image.sizes.text,
Permission::READ_EXECUTE,
);
top += image.sizes.text;
MlCore::set_wmmu_window(
WindowId::ConstData,
top,
image.sizes.constant_data,
Permission::READ,
);
top += image.sizes.constant_data;
MlCore::set_wmmu_window(
WindowId::ModelOutput,
top,
image.sizes.model_output,
Permission::READ_WRITE,
);
top += image.sizes.model_output;
MlCore::set_wmmu_window(
WindowId::StaticData,
top,
image.sizes.static_data,
Permission::READ_WRITE,
);
// TODO(jesionowski): Set model_input window when sensor manager
// is integrated.
// NB: TEMP_DATA_WINDOW is set in set_tcm_bottom.
return true;
}
false
}
/// Zeroes out the temporary data section.
pub fn clear_temp_data(&self) { MlCore::clear_tcm(self.tcm_bottom, self.tcm_bottom_size()); }
fn ids_at(&self, idx: ImageIdx) -> (&str, &str) {
match self.images[idx].as_ref() {
Some(image) => (&image.id.bundle_id, &image.id.model_id),
@@ -331,10 +336,11 @@ impl ImageManager {
}
}
/// Prints local state for debugging.
pub fn debug_state(&self) {
info!("Loaded Images:");
for image in self.images.as_ref().iter().flatten() {
info!(" {:?}", image);
info!(" {:x?}", image);
}
info!("Image Queue:");
@@ -343,9 +349,9 @@ impl ImageManager {
info!(" {}:{}", bundle, model);
}
info!("Sensor Top: {}", self.sensor_top);
info!("TCM Top: {}", self.tcm_top);
info!("TCM Bottom: {}", self.tcm_bottom);
info!("Sensor Top: 0x{:x}", self.sensor_top);
info!("TCM Top: 0x{:x}", self.tcm_top);
info!("TCM Bottom: 0x{:x}", self.tcm_bottom);
}
}
@@ -358,26 +364,13 @@ mod test {
#[test]
fn allocate_sensor() {
let mut image_manager = ImageManager::default();
let mut image_manager = ImageManager::new();
assert_eq_hex!(image_manager.allocate_sensor_input(0x1000), TCM_PADDR);
assert_eq_hex!(image_manager.tcm_top_size(), 0x1000);
}
// Stub out the Read trait to enable fake images.
struct FakeImage;
impl Read for FakeImage {
fn read(&mut self, _buf: &mut [u8]) -> Result<usize, kata_io::Error> { Ok(0) }
}
fn fake_image() -> Box<dyn Read> { Box::new(FakeImage {}) }
// The on_flash_sizes are only used when writing the image to memory. For
// these tests we want to ignore this, so just use zeroed sizes.
fn ignore_on_flash_sizes() -> ImageSizes { ImageSizes::default() }
fn constant_image_size(size: usize) -> ImageSizes {
ImageSizes {
text: size,
@@ -398,20 +391,16 @@ mod test {
fn default_id() -> ImageId { make_id(1) }
fn load_image(image_manager: &mut ImageManager, id: ImageId, unpacked_sizes: ImageSizes) {
// fake_vec_core can't fail to load.
let _ = image_manager.load_image(
&mut fake_image(),
id,
ignore_on_flash_sizes(),
unpacked_sizes,
);
fn load_image(image_manager: &mut ImageManager, id: ImageId, in_memory_sizes: ImageSizes) {
image_manager.make_space(in_memory_sizes.data_top_size(), in_memory_sizes.temporary_data);
image_manager.commit_image(id, in_memory_sizes);
}
// Load a model and see that is_loaded returns true. Unload and see false.
#[test]
fn load_unload() {
let mut image_manager = ImageManager::default();
let mut image_manager = ImageManager::new();
load_image(&mut image_manager, default_id(), constant_image_size(0x1000));
let id = default_id();
@@ -428,7 +417,7 @@ mod test {
// isn't loaded.
#[test]
fn is_loaded_no() {
let mut image_manager = ImageManager::default();
let mut image_manager = ImageManager::new();
assert!(!image_manager.is_loaded(&default_id()));
assert!(!image_manager.unload_image(&default_id()));
@@ -464,7 +453,7 @@ mod test {
// of the second model. Then, load a 4th that unloads the others.
#[test]
fn loads_force_unloads() {
let mut image_manager = ImageManager::default();
let mut image_manager = ImageManager::new();
let id1 = make_id(1);
let id2 = make_id(2);
@@ -490,7 +479,7 @@ mod test {
// others have been compacted.
#[test]
fn unloads_compact_tcm() {
let mut image_manager = ImageManager::default();
let mut image_manager = ImageManager::new();
let id1 = make_id(1);
let id2 = make_id(2);

2
apps/system/components/MlCoordinator/kata-ml-support/src/lib.rs
View File

@@ -1,3 +1,3 @@
#![no_std]
mod image_manager;
pub mod image_manager;

3
apps/system/components/MlCoordinator/kata-vec-core/Cargo.toml
View File

@@ -5,9 +5,6 @@ edition = "2021"
[dependencies]
kata-io = { path = "../../DebugConsole/kata-io" }
kata-memory-interface = { path = "../../MemoryManager/kata-memory-interface" }
kata-ml-shared = { path = "../kata-ml-shared" }
kata-proc-interface = { path = "../../ProcessManager/kata-proc-interface" }
kata-os-common = { path = "../../kata-os-common" }
modular-bitfield = "0.11.2"
log = "0.4"

153
apps/system/components/MlCoordinator/kata-vec-core/src/lib.rs
View File

@@ -7,29 +7,16 @@ extern crate alloc;
mod vc_top;
use alloc::boxed::Box;
use core::mem::size_of;
use core::slice;
use kata_memory_interface::ObjDescBundle;
use kata_ml_shared::{ModelSections, Window, WMMU_PAGE_SIZE};
use kata_ml_shared::{TCM_PADDR, TCM_SIZE};
use kata_proc_interface::BundleImage;
use io::Read;
use kata_io as io;
use kata_io::Read;
use kata_ml_shared::{Permission, WindowId, TCM_PADDR, TCM_SIZE};
use log::{error, trace};
extern "C" {
static TCM: *mut u32;
}
fn round_up(a: usize, b: usize) -> usize {
if (a % b) == 0 {
a
} else {
usize::checked_add(a, b).unwrap() - (a % b)
}
}
pub fn enable_interrupts(enable: bool) {
let intr_enable = vc_top::IntrEnable::new()
.with_host_req(enable)
@@ -39,16 +26,27 @@ pub fn enable_interrupts(enable: bool) {
vc_top::set_intr_enable(intr_enable);
}
pub fn set_wmmu(sections: &ModelSections) {
// XXX: Support multiple sections.
pub fn set_wmmu_window(
window_id: WindowId,
start_address: usize,
length: usize,
permission: Permission,
) {
trace!(
"Set window {:?} to addr {:x} len {:x}",
window_id,
start_address,
length
);
vc_top::set_mmu_window_offset(window_id as usize, start_address);
// The length of the window is not the size of the window, but rather
// the last address of the window. This saves us a bit in hardware:
// 0x400000 is 23 bits vs. 0x3FFFFF 22 bits.
vc_top::set_mmu_window_offset(0, sections.tcm.addr);
vc_top::set_mmu_window_length(0, sections.tcm.size - 1);
vc_top::set_mmu_window_permission(0, vc_top::Permission::ReadWriteExecute);
vc_top::set_mmu_window_length(window_id as usize, length - 1);
vc_top::set_mmu_window_permission(window_id as usize, permission);
}
/// Start the core at the default PC.
pub fn run() {
let ctrl = vc_top::Ctrl::new()
.with_freeze(false)
@@ -57,66 +55,53 @@ pub fn run() {
vc_top::set_ctrl(ctrl);
}
// Writes the section of the image from |start_address| to
// |start_address + on_flash_size| into the TCM. Zeroes the section from
// |on_flash_size| to |unpacked_size|.
#[allow(dead_code)] // XXX: Remove when integrated.
pub fn write_image_part(
image: &mut Box<dyn Read>,
/// Writes the section of the image from |start_address| to
/// |start_address + on_flash_size| into the TCM. Zeroes the section from
/// |on_flash_size| to |unpacked_size|. Returns None if the write failed.
pub fn write_image_part<R: Read>(
image: &mut R,
start_address: usize,
on_flash_size: usize,
unpacked_size: usize,
) -> Result<(), &'static str> {
) -> Option<()> {
let start = start_address - TCM_PADDR;
trace!(
"Writing {:x} bytes to 0x{:x}, {:x} unpacked size",
on_flash_size,
start_address,
unpacked_size
);
let tcm_slice = unsafe { slice::from_raw_parts_mut(TCM as *mut u8, TCM_SIZE) };
image
.read_exact(&mut tcm_slice[start..on_flash_size])
.map_err(|_| "section read error")?;
// TODO(jesionowski): Use hardware clear when TCM_SIZE fits into INIT_END.
tcm_slice[on_flash_size..unpacked_size].fill(0x00);
Ok(())
}
// XXX: Remove when write_image is integrated.
// Loads the model into the TCM.
pub fn load_image(frames: &ObjDescBundle) -> Result<ModelSections, &'static str> {
let mut image = BundleImage::new(frames);
let mut tcm_found = false;
// Size of window is filled in below.
let mut window = Window {
addr: TCM_PADDR,
size: 0,
if let Err(e) = image.read_exact(&mut tcm_slice[start..start + on_flash_size]) {
error!("Section read error {:?}", e);
return None;
};
clear_tcm();
// NB: we require a TCM section and that only one is present
while let Some(section) = image.next_section() {
let slice = if section.vaddr == TCM_PADDR {
if tcm_found {
return Err("dup TCM section");
}
tcm_found = true;
// TODO(jesionowski): Use hardware clear when TCM_SIZE fits into INIT_END.
tcm_slice[start + on_flash_size..start + unpacked_size].fill(0x00);
if section.fsize > TCM_SIZE {
return Err("TCM section too big");
}
window.size = round_up(section.msize, WMMU_PAGE_SIZE);
unsafe { slice::from_raw_parts_mut(TCM as *mut u8, TCM_SIZE) }
} else {
return Err("Unexpected section");
};
image
.read_exact(&mut slice[section.data_range()])
.map_err(|_| "section read error")?;
// TODO(jesionowski): Remove when clear_tcm is fully implemented.
slice[section.zero_range()].fill(0x00);
}
if !tcm_found {
return Err("Incomplete");
}
Ok(ModelSections { tcm: window })
Some(())
}
/// Move |src_index..src_index + byte_length| to
/// |dest_index..dest_index + byte_length|.
pub fn tcm_move(src: usize, dest: usize, byte_length: usize) {
trace!(
"Moving 0x{:x} bytes to 0x{:x} from 0x{:x}",
byte_length,
dest as usize,
src as usize,
);
let tcm_slice = get_tcm_slice();
let src_index = (src - TCM_PADDR) / size_of::<u32>();
let dest_index = (dest - TCM_PADDR) / size_of::<u32>();
let count: usize = byte_length / size_of::<u32>();
tcm_slice.copy_within(src_index..src_index + count, dest_index);
}
// Interrupts are write 1 to clear.
@@ -144,7 +129,7 @@ pub fn clear_data_fault() {
vc_top::set_intr_state(intr_state);
}
// TODO(jesionowski): Remove dead_code when TCM_SIZE fits into INIT_END.
// TODO(jesionowski): Use when TCM_SIZE fits into INIT_END.
#[allow(dead_code)]
fn clear_section(start: u32, end: u32) {
let init_start = vc_top::InitStart::new().with_address(start);
@@ -152,15 +137,29 @@ fn clear_section(start: u32, end: u32) {
let init_end = vc_top::InitEnd::new().with_address(end).with_valid(true);
vc_top::set_init_end(init_end);
while !vc_top::get_init_status().init_done() {}
}
pub fn clear_tcm() {
// TODO(jesionowski): Enable when TCM_SIZE fits into INIT_END.
// clear_section(0, TCM_SIZE as u32, false);
/// Zeroes out |byte_length| bytes starting at |addr|.
pub fn clear_tcm(addr: usize, byte_length: usize) {
assert!(addr >= TCM_PADDR);
assert!(addr + byte_length <= TCM_PADDR + TCM_SIZE);
trace!("Clearing 0x{:x} bytes at 0x{:x}", byte_length, addr);
let start = (addr - TCM_PADDR) / size_of::<u32>();
let count: usize = byte_length / size_of::<u32>();
// TODO(jesionowski): Use clear_section method when able.
let tcm_slice = get_tcm_slice();
tcm_slice[start..start + count].fill(0x00);
}
// TODO(jesionowski): Use when TCM_SIZE fits into INIT_END.
// We'll want to kick off the hardware clear after the execution is complete,
// holding off the busy-wait until we're ready to start another execution.
#[allow(dead_code)]
pub fn wait_for_clear_to_finish() { while !vc_top::get_init_status().init_done() {} }
// TODO(jesionowski): Remove these when error handling is refactored.
// The status will be faulty iff the interrupt line is raised, and
// we won't have the fault registers on Springbok.

11
apps/system/components/MlCoordinator/kata-vec-core/src/vc_top.rs
View File

@@ -3,6 +3,7 @@
// Setters and getters for the Vector Core CSRs.
use core::ptr;
use kata_ml_shared::Permission;
use modular_bitfield::prelude::*;
extern "C" {
@@ -219,17 +220,9 @@ pub fn set_mmu_window_length(window: usize, length: usize) {
}
}
pub enum Permission {
Read = 1,
Write = 2,
ReadWrite = 3,
Execute = 4,
ReadWriteExecute = 7,
}
pub fn set_mmu_window_permission(window: usize, permission: Permission) {
let addr = window_addr(window) + PERMISSIONS_ADDR;
unsafe {
core::ptr::write_volatile(addr as *mut usize, permission as usize);
core::ptr::write_volatile(addr as *mut usize, permission.bits() as usize);
}
}

1
apps/system/interfaces/MlCoordBindings.h
View File

@@ -13,6 +13,7 @@ typedef enum MlCoordError {
MlCoordOk,
InvalidModelId,
InvalidBundleId,
InvalidImage,
LoadModelFailed,
NoModelSlotsLeft,
NoSuchModel,