From 92dc1019bedf54f155bedf827e84b47636510524 Mon Sep 17 00:00:00 2001 From: Adam Jesionowski Date: Thu, 7 Jul 2022 09:38:05 -0700 Subject: [PATCH] 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 --- .../MlCoordinator/fake-vec-core/Cargo.toml | 1 + .../MlCoordinator/fake-vec-core/src/lib.rs | 12 +- .../MlCoordinator/kata-ml-shared/src/lib.rs | 4 +- .../MlCoordinator/kata-ml-support/Cargo.toml | 1 + .../kata-ml-support/src/image_manager.rs | 130 ++++++++++++++++-- .../MlCoordinator/kata-vec-core/src/lib.rs | 35 ++++- 6 files changed, 159 insertions(+), 24 deletions(-) diff --git a/apps/system/components/MlCoordinator/fake-vec-core/Cargo.toml b/apps/system/components/MlCoordinator/fake-vec-core/Cargo.toml index 0c98985..b25f988 100644 --- a/apps/system/components/MlCoordinator/fake-vec-core/Cargo.toml +++ b/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" } diff --git a/apps/system/components/MlCoordinator/fake-vec-core/src/lib.rs b/apps/system/components/MlCoordinator/fake-vec-core/src/lib.rs index 00a983a..1dc03e6 100644 --- a/apps/system/components/MlCoordinator/fake-vec-core/src/lib.rs +++ b/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, + _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) {} diff --git a/apps/system/components/MlCoordinator/kata-ml-shared/src/lib.rs b/apps/system/components/MlCoordinator/kata-ml-shared/src/lib.rs index dee3eed..11ef9f7 100644 --- a/apps/system/components/MlCoordinator/kata-ml-shared/src/lib.rs +++ b/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, diff --git a/apps/system/components/MlCoordinator/kata-ml-support/Cargo.toml b/apps/system/components/MlCoordinator/kata-ml-support/Cargo.toml index 458a889..7ea14a3 100644 --- a/apps/system/components/MlCoordinator/kata-ml-support/Cargo.toml +++ b/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" } diff --git a/apps/system/components/MlCoordinator/kata-ml-support/src/image_manager.rs b/apps/system/components/MlCoordinator/kata-ml-support/src/image_manager.rs index b30f757..fd4bcb7 100644 --- a/apps/system/components/MlCoordinator/kata-ml-support/src/image_manager.rs +++ b/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, + 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 { + Ok(0) + } + } + + fn fake_image() -> Box { + 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. diff --git a/apps/system/components/MlCoordinator/kata-vec-core/src/lib.rs b/apps/system/components/MlCoordinator/kata-vec-core/src/lib.rs index 15da2e1..038e3e3 100644 --- a/apps/system/components/MlCoordinator/kata-vec-core/src/lib.rs +++ b/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, + 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 { let mut image = BundleImage::new(frames); @@ -90,9 +116,7 @@ pub fn load_image(frames: &ObjDescBundle) -> Result 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);