Complete image_manager.

The ML Coordinator will validate the cpio model and collect the fsize
and msizes for each section. It will then pass that information along
with a Boxed version of the bundle_image, which implements the
kata_io::Read trait.

This CL adds the lower level (vec-core) function that writes to the TCM and
calls these functions from the image_manager.

Complete integration of the image_manager is pending splitting the ELF
image into all 6 sections.

Change-Id: I7a5706c588867b4aee04109e2a9edeca071d2ca8
GitOrigin-RevId: 89df1c81bade3ec4508f643a8ba83cae6a3e1f60

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

@@ -4,4 +4,5 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
+kata-io = { path = "../../DebugConsole/kata-io" }
 kata-ml-shared = { path = "../kata-ml-shared" }

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

@@ -1,7 +1,10 @@
 #![no_std]
 
 // fake-vec-core is a stubbed out version of kata-vec-core.
+extern crate alloc;
 
+use alloc::boxed::Box;
+use kata_io::Read;
 use kata_ml_shared::ModelSections;
 
 pub fn enable_interrupts(_enable: bool) {}
@@ -10,7 +13,14 @@ pub fn set_wmmu(_sections: &ModelSections) {}
 
 pub fn run() {}
 
-pub fn tcm_write(_offset: usize, _buf: &[u32]) {}
+pub fn write_image_part(
+    _image: &mut Box<dyn Read>,
+    _start_address: usize,
+    _on_flash_size: usize,
+    _unpacked_size: usize,
+) -> Result<(), &'static str> {
+    Ok(())
+}
 
 pub fn tcm_move(_src_offset: usize, _dest_offset: usize, _byte_length: usize) {}
 

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

@@ -10,7 +10,7 @@ use alloc::string::String;
 
 /// An image is uniquely identified by the bundle that owns it and the
 /// particular model id in that bundle.
-#[derive(Clone, Eq, PartialEq)]
+#[derive(Clone, Debug, Eq, PartialEq)]
 pub struct ImageId {
     pub bundle_id: String,
     pub model_id: String,
@@ -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, Default)]
+#[derive(Clone, Debug, Default)]
 pub struct ImageSizes {
     pub text: usize,
     pub model_input: usize,

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

@@ -5,6 +5,7 @@ edition = "2021"
 
 [dependencies]
 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" }

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

@@ -28,11 +28,14 @@
 
 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 log::trace;
+use log::{info, trace};
+
+use kata_io::Read;
 
 // XXX: Enable configuration when kata_vec_core does not depend on
 // kata-os-common.
@@ -61,6 +64,7 @@ use fake_vec_core as MlCore;
 //               +---------------+
 // Each segment is page aligned.
 
+#[derive(Debug)]
 struct Image {
     id: ImageId,
     data_top_addr: usize,
@@ -203,6 +207,7 @@ impl ImageManager {
     // 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) {
+        assert!(top_tcm_needed + temp_tcm_needed <= TCM_SIZE);
         let mut available_tcm = self.tcm_free_space();
         let mut space_needed_for_temp =
             space_needed(self.required_temporary_data(), temp_tcm_needed);
@@ -265,17 +270,57 @@ impl ImageManager {
         self.get_image_index(id).is_some()
     }
 
-    // XXX: Implement load_image.
-    pub fn load_image(&mut self, id: ImageId, sizes: ImageSizes) {
-        self.make_space(sizes.data_top_size(), sizes.temporary_data);
+    // 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;
 
-        // XXX: Do the write thing.
+        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,
+            sizes: unpacked_sizes,
             data_top_addr: self.tcm_top,
         });
+
+        Ok(())
     }
 
     // Unloads image |id| if loaded. Returns true if an image was unloaded.
@@ -290,7 +335,29 @@ impl ImageManager {
         false
     }
 
-    // XXX: Add debug_state fn, similar to MLCoordinator.
+    fn ids_at(&self, idx: ImageIdx) -> (&str, &str) {
+        match self.images[idx].as_ref() {
+            Some(image) => (&image.id.bundle_id, &image.id.model_id),
+            None => ("None", "None"),
+        }
+    }
+
+    pub fn debug_state(&self) {
+        info!("Loaded Images:");
+        for image in self.images.as_ref().iter().flatten() {
+            info!("  {:?}", image);
+        }
+
+        info!("Image Queue:");
+        for idx in &self.image_queue {
+            let (bundle, model) = self.ids_at(*idx);
+            info!("  {}:{}", bundle, model);
+        }
+
+        info!("Sensor Top: {}", self.sensor_top);
+        info!("TCM Top: {}", self.tcm_top);
+        info!("TCM Bottom: {}", self.tcm_bottom);
+    }
 }
 
 #[cfg(test)]
@@ -309,6 +376,25 @@ mod test {
         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,
@@ -331,11 +417,25 @@ mod test {
         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,
+        );
+    }
+
     // Load a model and see that is_loaded returns true. Unload and see false.
     #[test]
     fn load_unload() {
         let mut image_manager = ImageManager::default();
-        image_manager.load_image(default_id(), constant_image_size(0x1000));
+        load_image(
+            &mut image_manager,
+            default_id(),
+            constant_image_size(0x1000),
+        );
 
         let id = default_id();
 
@@ -393,19 +493,19 @@ mod test {
         let id2 = make_id(2);
         let id3 = make_id(3);
 
-        image_manager.load_image(id1.clone(), half_image());
-        image_manager.load_image(id2.clone(), half_image());
+        load_image(&mut image_manager, id1.clone(), half_image());
+        load_image(&mut image_manager, id2.clone(), half_image());
 
         assert!(image_manager.is_loaded(&id1));
         assert!(image_manager.is_loaded(&id2));
 
-        image_manager.load_image(id3.clone(), half_image());
+        load_image(&mut image_manager, id3.clone(), half_image());
 
         assert!(image_manager.is_loaded(&id1));
         assert!(image_manager.is_loaded(&id3));
 
         let id4 = make_id(4);
-        image_manager.load_image(id4.clone(), full_image());
+        load_image(&mut image_manager, id4.clone(), full_image());
         assert!(image_manager.is_loaded(&id4));
     }
 
@@ -447,9 +547,9 @@ mod test {
             temporary_data: 0x3000, // This will be the largest post unload
         };
 
-        image_manager.load_image(id1.clone(), sizes1.clone());
-        image_manager.load_image(id2.clone(), sizes2.clone());
-        image_manager.load_image(id3.clone(), sizes3.clone());
+        load_image(&mut image_manager, id1.clone(), sizes1.clone());
+        load_image(&mut image_manager, id2.clone(), sizes2.clone());
+        load_image(&mut image_manager, id3.clone(), sizes3.clone());
 
         // The third image will be available at images[2]. Before unloading we
         // validate that it's past the first two models.

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

@@ -3,13 +3,16 @@
 // kata-vec-core is the vector core driver. It is responsible for providing
 // convenient methods for interacting with the hardware.
 
+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_SIZE, TCM_PADDR};
+use kata_ml_shared::{TCM_PADDR, TCM_SIZE};
 use kata_proc_interface::BundleImage;
 
 use io::Read;
@@ -54,6 +57,29 @@ 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>,
+    start_address: usize,
+    on_flash_size: usize,
+    unpacked_size: usize,
+) -> Result<(), &'static str> {
+    let start = start_address - TCM_PADDR;
+
+    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);
@@ -90,9 +116,7 @@ pub fn load_image(frames: &ObjDescBundle) -> Result<ModelSections, &'static str>
     if !tcm_found {
         return Err("Incomplete");
     }
-    Ok(ModelSections {
-        tcm: window,
-    })
+    Ok(ModelSections { tcm: window })
 }
 
 // Interrupts are write 1 to clear.
@@ -123,8 +147,7 @@ pub fn clear_data_fault() {
 // TODO(jesionowski): Remove dead_code 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);
+    let init_start = vc_top::InitStart::new().with_address(start);
     vc_top::set_init_start(init_start);
 
     let init_end = vc_top::InitEnd::new().with_address(end).with_valid(true);