kata: Integrate Image Manager and WMMU changes

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

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

@@ -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 }
 

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

@@ -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"

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

@@ -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) {
+    let id = match validate_ids(c_bundle_id, c_model_id) {
-        Ok(ids) => ids,
+        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) {
+    let id = match validate_ids(c_bundle_id, c_model_id) {
-        Ok(ids) => ids,
+        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) {
+    let id = match validate_ids(c_bundle_id, c_model_id) {
-        Ok(ids) => ids,
+        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;
     }
 

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

@@ -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" }

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

@@ -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,
+    id: ImageId,
-    model_id: String,
+    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,21 +71,83 @@ 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
+    // Validates the image by ensuring it has all the required loadable
-    /// already been loaded.
+    // sections and that it fits into the TCM. Returns a tuple of
-    fn load_model(&mut self, model_idx: ModelIdx) -> Result<(), MlCoordError> {
+    // |(on_flash_sizes, in_memory_sizes)|.
-        if self.loaded_model == Some(model_idx) {
+    fn validate_image(&self, id: &ImageId) -> Option<(ImageSizes, ImageSizes)> {
-            trace!("Model already loaded, skipping load");
+        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 next_idx = self.execution_queue.remove(0);
-        let model = self.models[model_idx]
+        let model = self.models[next_idx].as_ref().expect("Model get fail");
-            .as_ref()
-            .ok_or(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|.
@@ -91,50 +155,86 @@ impl MLCoordinator {
             // 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) {
+            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 ret_status = match MlCore::load_image(&model_frames) {
+                    let mut image = BundleImage::new(&model_frames);
-                    Err(e) => {
+
-                        error!("Load of {}:{} failed: {:?}", &model.bundle_id, &model.model_id, e);
+                    // Ask the image manager to make enough room and get
-                        // May have corrupted TCM.
+                    // the address to write to.
-                        self.loaded_model = None;
+                    let mut temp_top = self.image_manager.make_space(
-                        self.statistics.load_failures += 1;
+                        model.in_memory_sizes.data_top_size(),
-                        Err(MlCoordError::LoadModelFailed)
+                        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);
+                    // Inform the image manager the image has been written.
-                        Ok(())
+                    self.image_manager
-                    }
+                        .commit_image(model.id.clone(), model.in_memory_sizes);
-                };
+
+                    drop(image);
                     let _ = kata_object_free_in_cnode(&model_frames);
-                ret_status
                 }
                 Err(e) => {
                     error!(
                         "LoadModel of bundle {}:{} failed: {:?}",
-                    &model.bundle_id, &model.model_id, e
+                        &model.id.bundle_id, &model.id.model_id, e
                     );
                     self.statistics.load_failures += 1;
-                Err(MlCoordError::LoadModelFailed)
+                    return Err(MlCoordError::LoadModelFailed);
                 }
             }
         }
 
-    /// If there is a next model in the queue, load it onto the core and start
+        // TODO(jesionowski): Investigate if we need to clear the entire
-    /// running. If there's already a running model, don't do anything.
+        // temporary data section or just certain parts.
-    fn schedule_next_model(&mut self) -> Result<(), MlCoordError> {
+        // TODO(jesionowski): When hardware clear is enabled, we should
-        if !self.running_model.is_some() && !self.execution_queue.is_empty() {
+        // kick it off after the run instead.
-            let next_idx = self.execution_queue.remove(0);
+        self.image_manager.clear_temp_data();
-            // 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
+        self.image_manager.set_wmmu(&model.id);
-            // 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);
-        }
+        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,
+        id: ImageId,
-        model_id: String,
         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,
+            id,
-            model_id,
+            on_flash_sizes,
+            in_memory_sizes,
             rate_in_ms,
         });
 
         Ok(index)
     }
 
-    /// Start a one-time model execution, to be executed immediately.
+    // Returns the index for model |id| if it exists.
-    pub fn oneshot(&mut self, bundle_id: String, model_id: String) -> Result<(), MlCoordError> {
+    fn get_model_index(&self, id: &ImageId) -> Option<ModelIdx> {
-        let model_idx = self.ready_model(bundle_id, model_id, None)?;
+        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(
+    pub fn periodic(&mut self, id: ImageId, rate_in_ms: u32) -> Result<(), MlCoordError> {
-        &mut self,
+        // Check if we've loaded this model already.
-        bundle_id: String,
+        let idx = match self.get_model_index(&id) {
-        model_id: String,
+            Some(idx) => idx,
-        rate_in_ms: u32,
+            None => self.ready_model(id, Some(rate_in_ms))?,
-    ) -> Result<(), MlCoordError> {
+        };
-        let model_idx = self.ready_model(bundle_id, 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
+        let model_idx = self.get_model_index(id).ok_or(MlCoordError::NoSuchModel)?;
-            .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)?;
 
         // 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() {
+        for model in self.models.as_ref().iter().flatten() {
-            if let Some(m) = model {
+            info!("  {:x?}", model);
-                info!("  {:?}", m);
-            }
         }
 
         info!("Execution Queue:");
@@ -353,5 +452,7 @@ impl MLCoordinator {
         }
 
         info!("Statistics: {:?}", self.statistics);
+
+        self.image_manager.debug_state();
     }
 }

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

@@ -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,

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

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

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

@@ -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
+    pub fn total_size(&self) -> usize {
-// are enabled.
+        self.data_top_size() + self.temporary_data + self.model_input
-/// 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,
-}
 
-// XXX: Out-dated. Refactor when multiple sections are enabled.
+    // A set of sizes is considered valid if everything but model_input is
-/// When a model is loaded onto the Vector Core, the ML Coordinator needs to
+    // non-zero.
-/// track where each window is.
+    pub fn is_valid(&self) -> bool {
-pub struct ModelSections {
+        // TODO(jesionowski): Add `&& self.model_output != 0` when model output
-    pub tcm: Window,
+        // 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)
+    }
+}

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

@@ -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]

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

@@ -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::*;
-use kata_ml_shared::{MAX_MODELS, TCM_PADDR, TCM_SIZE, WMMU_PAGE_SIZE};
 use log::{info, trace};
 
-use kata_io::Read;
+#[cfg(test)]
-
-// 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)]
 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 {
+impl ImageManager {
-    fn default() -> Self {
+    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
+    /// Removes images in FILO order until the top TCM and temp TCM
-    // constraints are satisfied.
+    /// constraints are satisfied. Returns the address of the freed space.
-    fn make_space(&mut self, top_tcm_needed: usize, temp_tcm_needed: usize) {
+    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;
-    }
+        MlCore::set_wmmu_window(
-
+            WindowId::TempData,
-    // Updates the pointers after an image is written to TCM to ensure the
+            self.tcm_bottom,
-    // image is kept around.
+            temp_data_size,
-    fn update_image_bookkeeping(&mut self, image: Image) {
+            Permission::READ_WRITE,
-        // 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);
     }
 
     // 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
+    /// Appends an (already written) image to internal book-keeping. This class
-    // necessary. |on_flash_sizes| represents the on-disk sizes (ie fsize)
+    /// does not handle the write as it requires seL4 references. The
-    // for each section, while |unpacked_sizes| represents the aligned memory
+    /// MlCoordinator must call this function after the write.
-    // needed for execution (ie msize).
+    pub fn commit_image(&mut self, id: ImageId, sizes: ImageSizes) {
-    pub fn load_image(
+        let image = 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 {
             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!("Sensor Top: 0x{:x}", self.sensor_top);
-        info!("TCM Top: {}", self.tcm_top);
+        info!("TCM Top: 0x{:x}", self.tcm_top);
-        info!("TCM Bottom: {}", self.tcm_bottom);
+        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) {
+    fn load_image(image_manager: &mut ImageManager, id: ImageId, in_memory_sizes: ImageSizes) {
-        // fake_vec_core can't fail to load.
+        image_manager.make_space(in_memory_sizes.data_top_size(), in_memory_sizes.temporary_data);
-        let _ = image_manager.load_image(
+
-            &mut fake_image(),
+        image_manager.commit_image(id, in_memory_sizes);
-            id,
-            ignore_on_flash_sizes(),
-            unpacked_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);

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

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

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

@@ -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"

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

@@ -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_io::Read;
-use kata_ml_shared::{ModelSections, Window, WMMU_PAGE_SIZE};
+use kata_ml_shared::{Permission, WindowId, TCM_PADDR, TCM_SIZE};
-use kata_ml_shared::{TCM_PADDR, TCM_SIZE};
+use log::{error, trace};
-use kata_proc_interface::BundleImage;
-
-use io::Read;
-use kata_io as io;
 
 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) {
+pub fn set_wmmu_window(
-    // XXX: Support multiple sections.
+    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(window_id as usize, length - 1);
-    vc_top::set_mmu_window_length(0, sections.tcm.size - 1);
+    vc_top::set_mmu_window_permission(window_id as usize, permission);
-    vc_top::set_mmu_window_permission(0, vc_top::Permission::ReadWriteExecute);
 }
 
+/// 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
+/// Writes the section of the image from |start_address| to
-// |start_address + on_flash_size| into the TCM. Zeroes the section from
+/// |start_address + on_flash_size| into the TCM. Zeroes the section from
-// |on_flash_size| to |unpacked_size|.
+/// |on_flash_size| to |unpacked_size|. Returns None if the write failed.
-#[allow(dead_code)] // XXX: Remove when integrated.
+pub fn write_image_part<R: Read>(
-pub fn write_image_part(
+    image: &mut R,
-    image: &mut Box<dyn Read>,
     start_address: usize,
     on_flash_size: usize,
     unpacked_size: usize,
-) -> Result<(), &'static str> {
+) -> Option<()> {
     let start = start_address - TCM_PADDR;
 
-    let tcm_slice = unsafe { slice::from_raw_parts_mut(TCM as *mut u8, TCM_SIZE) };
+    trace!(
-    image
+        "Writing {:x} bytes to 0x{:x}, {:x} unpacked size",
-        .read_exact(&mut tcm_slice[start..on_flash_size])
+        on_flash_size,
-        .map_err(|_| "section read error")?;
+        start_address,
-    // TODO(jesionowski): Use hardware clear when TCM_SIZE fits into INIT_END.
+        unpacked_size
-    tcm_slice[on_flash_size..unpacked_size].fill(0x00);
+    );
 
-    Ok(())
+    let tcm_slice = unsafe { slice::from_raw_parts_mut(TCM as *mut u8, TCM_SIZE) };
+
+    if let Err(e) = image.read_exact(&mut tcm_slice[start..start + on_flash_size]) {
+        error!("Section read error {:?}", e);
+        return None;
+    };
+
+    // TODO(jesionowski): Use hardware clear when TCM_SIZE fits into INIT_END.
+    tcm_slice[start + on_flash_size..start + unpacked_size].fill(0x00);
+
+    Some(())
 }
 
-// XXX: Remove when write_image is integrated.
+/// Move |src_index..src_index + byte_length| to
-// Loads the model into the TCM.
+/// |dest_index..dest_index + byte_length|.
-pub fn load_image(frames: &ObjDescBundle) -> Result<ModelSections, &'static str> {
+pub fn tcm_move(src: usize, dest: usize, byte_length: usize) {
-    let mut image = BundleImage::new(frames);
+    trace!(
-    let mut tcm_found = false;
+        "Moving 0x{:x} bytes to 0x{:x} from 0x{:x}",
-    // Size of window is filled in below.
+        byte_length,
-    let mut window = Window {
+        dest as usize,
-        addr: TCM_PADDR,
+        src as usize,
-        size: 0,
+    );
-    };
 
-    clear_tcm();
+    let tcm_slice = get_tcm_slice();
-    // NB: we require a TCM section and that only one is present
+    let src_index = (src - TCM_PADDR) / size_of::<u32>();
-    while let Some(section) = image.next_section() {
+    let dest_index = (dest - TCM_PADDR) / size_of::<u32>();
-        let slice = if section.vaddr == TCM_PADDR {
+    let count: usize = byte_length / size_of::<u32>();
-            if tcm_found {
-                return Err("dup TCM section");
-            }
-            tcm_found = true;
 
-            if section.fsize > TCM_SIZE {
+    tcm_slice.copy_within(src_index..src_index + count, dest_index);
-                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 })
 }
 
 // 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() {
+/// Zeroes out |byte_length| bytes starting at |addr|.
-    // TODO(jesionowski): Enable when TCM_SIZE fits into INIT_END.
+pub fn clear_tcm(addr: usize, byte_length: usize) {
-    // clear_section(0, TCM_SIZE as u32, false);
+    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.

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

@@ -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);
     }
 }

apps/system/interfaces/MlCoordBindings.h

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