kata: MLCoordinator and VectorCoreDriver

Add a VectorCoreDriver component that handles setting vector core CSRs.

Rewrite MLCoordinator to conform to other Kata components. The old code
wasn't useful.

Add `test_mlexecute` command for running ML. Add plumbing from shell to
coordinator.

Change-Id: I3d563f1a343361c95d3ad5b78231fbe9df32b851
GitOrigin-RevId: f3c38839f708743de596339d1b8173315283b772

apps/system/CMakeLists.txt

@@ -42,6 +42,18 @@ DeclareCAmkESComponent(ProcessManager
   INCLUDES interfaces
 )
 
+RustAddLibrary(
+  kata_ml_coordinator
+  SOURCE_DIR ${CMAKE_CURRENT_LIST_DIR}/components/MlCoordinator
+  TARGET "riscv32imc-unknown-none-elf"
+  LIB_FILENAME libkata_ml_coordinator.a
+)
+
+DeclareCAmkESComponent(MlCoordinator
+  LIBS kata_ml_coordinator
+  INCLUDES interfaces
+)
+
 DeclareCAmkESComponent(UartDriver
   SOURCES components/UartDriver/src/driver.c
   INCLUDES opentitan-gen/include
@@ -51,4 +63,8 @@ DeclareCAmkESComponent(SeL4Debug
   SOURCES components/SeL4Debug/src/wrappers.c
 )
 
+DeclareCAmkESComponent(VectorCoreDriver
+  SOURCES components/VectorCoreDriver/src/driver.c
+)
+
 DeclareCAmkESRootserver(system.camkes)

apps/system/components/DebugConsole/DebugConsole.camkes

@@ -1,6 +1,7 @@
 import <LoggerInterface.camkes>;
 import <ProcessControlInterface.camkes>;
 import <PackageManagementInterface.camkes>;
+import <MlCoordinatorInterface.camkes>;
 import <SeL4DebugInterface.camkes>;
 
 component DebugConsole {
@@ -13,4 +14,5 @@ component DebugConsole {
   uses ProcessControlInterface proc_ctrl;
   uses PackageManagementInterface pkg_mgmt;
   uses SeL4DebugInterface sel4debug;
+  uses MlCoordinatorInterface mlcoord;
 }

apps/system/components/DebugConsole/kata-shell/src/lib.rs

@@ -88,6 +88,7 @@ fn dispatch_command(cmdline: &str, output: &mut dyn io::Write) {
                 "test_alloc" => test_alloc_command(output),
                 "test_alloc_error" => test_alloc_error_command(output),
                 "test_panic" => test_panic_command(),
+                "test_mlexecute" => test_mlexecute_command(),
 
                 _ => Err(CommandError::UnknownCommand),
             };
@@ -327,3 +328,14 @@ fn test_alloc_error_command(output: &mut dyn io::Write) -> Result<(), CommandErr
 fn test_panic_command() -> Result<(), CommandError> {
     panic!("testing");
 }
+
+/// Implements a command that runs an ML execution.
+fn test_mlexecute_command() -> Result<(), CommandError> {
+    extern "C" {
+        fn mlcoord_execute();
+    }
+    unsafe {
+        mlcoord_execute();
+    }
+    Ok(())
+}

apps/system/components/MlCoordinator/Cargo.toml

@@ -1,8 +1,7 @@
 [workspace]
 
 members = [
-    "ml-common",
-    "ml-coordinator",
+    "kata-ml-coordinator",
 ]
 
 [profile.dev]

apps/system/components/MlCoordinator/MlCoordinator.camkes

@@ -0,0 +1,10 @@
+import <LoggerInterface.camkes>;
+import <MlCoordinatorInterface.camkes>;
+
+component MlCoordinator {
+  control;
+  provides MlCoordinatorInterface mlcoord;
+
+  uses LoggerInterface logger;
+  uses VectorCoreInterface vctop;
+}

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

@@ -0,0 +1,15 @@
+[package]
+name = "kata-ml-coordinator"
+version = "0.1.0"
+authors = ["Adam Jesionowski <jesionowski@google.com>"]
+edition = "2018"
+
+[dependencies]
+kata-logger = { path = "../../DebugConsole/kata-logger" }
+kata-panic = { path = "../../DebugConsole/kata-panic" }
+log = "0.4"
+
+[lib]
+name = "kata_ml_coordinator"
+path = "src/run.rs"
+crate-type = ["staticlib"]

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

@@ -0,0 +1,41 @@
+#![no_std]
+
+// ML Coordinator Design Doc: go/sparrow-ml-doc
+
+extern crate kata_panic;
+
+use kata_logger::KataLogger;
+use log::debug;
+
+static KATA_LOGGER: KataLogger = KataLogger;
+
+#[no_mangle]
+pub extern "C" fn pre_init() {
+    log::set_logger(&KATA_LOGGER).unwrap();
+    log::set_max_level(log::LevelFilter::Debug);
+}
+
+#[no_mangle]
+pub extern "C" fn run() {
+    debug!("run");
+}
+
+// TODO: Move out of this file into separate (auto-generated?) file.
+// TODO: Consider the modular_bitfield crate to represent bitfields.
+fn vctop_ctrl(freeze: u32, vc_reset: u32, pc_start: u32) -> u32  {
+    ((pc_start & 1) << 2) + ((vc_reset & 1) << 1) + freeze
+}
+
+#[no_mangle]
+pub extern "C" fn mlcoord_execute() {
+    // TODO: Call into MLCoordinator when available.
+    // TODO: Once multiple model support is in start by name.
+    // TODO: Read the fault registers after execution and report any errors found.
+    extern "C" {
+        fn vctop_set_ctrl(ctrl: u32);
+    }
+    unsafe {
+        // Unhalt, start at default PC.
+        vctop_set_ctrl(vctop_ctrl(0, 0, 0));
+    }
+}

apps/system/components/MlCoordinator/ml-common/Cargo.toml

@@ -1,5 +0,0 @@
-[package]
-name = "ml-common"
-version = "0.1.0"
-authors = ["Adam Jesionowski <jesionowski@google.com>"]
-edition = "2018"

apps/system/components/MlCoordinator/ml-common/src/lib.rs

@@ -1,27 +0,0 @@
-// TODO(jesionowski): What are the actual errors we may encounter?
-#[derive(Copy, Clone, Debug)]
-pub enum ExecutionError {
-    ModelNotLoaded,
-    InvalidInstruction,
-    InvalidFetch,
-    CoreReset,
-}
-
-// The abstraction layer over the "hardware" of running an execution.
-// Returns a slice of bytes, which is the output data, or an ExecutionError.
-pub trait ExecutiveInterface {
-    fn run_model(&self, model: &Model) -> Result<&[u8], ExecutionError>;
-}
-
-// Options for execution that may be set by the application.
-pub struct ModelOptions {
-    pub rate: u32,
-}
-
-// Immutable model attributes.
-#[derive(PartialEq, Debug)]
-pub struct Model {
-    pub output_activations_len: usize,
-}
-
-pub const MAX_MODELS: u32 = 10;

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

@@ -1,9 +0,0 @@
-[package]
-name = "ml-coordinator"
-version = "0.1.0"
-authors = ["Adam Jesionowski <jesionowski@google.com>"]
-edition = "2018"
-
-[dependencies]
-ml-common = { path = "../ml-common" }
-arrayvec = "0.5.2"

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

@@ -1,161 +0,0 @@
-#![no_std]
-use arrayvec::ArrayVec;
-use ml_common as ml;
-use ml_common::ExecutiveInterface;
-
-pub struct MlCoordinator<'a> {
-    executive: &'a dyn ExecutiveInterface,
-    // Application owns the model struct.
-    // ML Coordinator owns a reference to the model.
-    models: ArrayVec<[&'a ml::Model; 10]>,
-}
-
-impl<'a> MlCoordinator<'a> {
-    // Create a new ML Coordinator instance.
-    fn new(executive: &'a dyn ExecutiveInterface) -> MlCoordinator<'a> {
-        let mut models = ArrayVec::<[&ml::Model; 10]>::new();
-
-        MlCoordinator {
-            executive: executive,
-            models: models,
-        }
-    }
-
-    // Returns the index of the model if it has already been loaded.
-    fn has_model(&self, model: &'a ml::Model) -> Option<usize> {
-        self.models.iter().position(|&x| model == x)
-    }
-
-    // Load the passed model. Returns true if the model was loaded
-    // successfully, and false if there is no room or the model was already loaded.
-    // This function will eventually perform validation of the model, create transfer buffers, etc.
-    #[must_use]
-    fn load(&mut self, model: &'a ml::Model) -> bool {
-        if self.models.is_full() || self.has_model(model).is_some() {
-            return false;
-        }
-        self.models.push(model);
-        true
-    }
-
-    // Unload the passed model. Returns true if the model was unloaded successfully,
-    // false if the model was not loaded in the first place.
-    fn unload(&mut self, model: &'a ml::Model) -> bool {
-        let res = self.has_model(model);
-        if let Some(idx) = res {
-            self.models.remove(idx);
-        }
-        res.is_some()
-    }
-
-    // Execute the passed model.
-    fn execute(&self, model: &'a ml::Model) -> Result<&'a [u8], ml::ExecutionError> {
-        self.has_model(model).ok_or(ml::ExecutionError::ModelNotLoaded)?;
-        self.executive.run_model(model)
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    struct FakeExecutive {
-        output_memory: [u8; 512],
-        fake_error: Option<ml::ExecutionError>,
-    }
-
-    impl ExecutiveInterface for FakeExecutive {
-        fn run_model(&self, _model: &ml::Model) -> Result<&[u8], ml::ExecutionError> {
-            match self.fake_error {
-                Some(err) => Err(err),
-                None => Ok(&self.output_memory),
-            }
-        }
-    }
-
-    #[test]
-    fn test_load_unload() {
-        let exec = tests::FakeExecutive {
-            output_memory: [0x00; 512],
-            fake_error: None,
-        };
-        let model = ml::Model {
-            output_activations_len: 512,
-        };
-        let mut ml_coord = MlCoordinator::new(&exec);
-        assert!(ml_coord.load(&model));
-        assert!(ml_coord.unload(&model));
-    }
-
-    #[test]
-    fn test_unload_noload() {
-        let exec = tests::FakeExecutive {
-            output_memory: [0x00; 512],
-            fake_error: None,
-        };
-        let model = ml::Model {
-            output_activations_len: 512,
-        };
-        let mut ml_coord = MlCoordinator::new(&exec);
-        assert_eq!(ml_coord.unload(&model), false);
-    }
-
-
-    #[test]
-    fn test_unload_twice() {
-        let exec = tests::FakeExecutive {
-            output_memory: [0x00; 512],
-            fake_error: None,
-        };
-        let model = ml::Model {
-            output_activations_len: 512,
-        };
-        let mut ml_coord = MlCoordinator::new(&exec);
-        assert!(ml_coord.load(&model));
-        assert!(ml_coord.unload(&model));
-        assert_eq!(ml_coord.unload(&model), false);
-    }
-
-    #[test]
-    fn test_load_twice() {
-        let exec = tests::FakeExecutive {
-            output_memory: [0x00; 512],
-            fake_error: None,
-        };
-        let model = ml::Model {
-            output_activations_len: 512,
-        };
-        let mut ml_coord = MlCoordinator::new(&exec);
-        assert!(ml_coord.load(&model));
-        assert_eq!(ml_coord.load(&model), false);
-    }
-
-    #[test]
-    fn test_execute_noload() {
-        let exec = tests::FakeExecutive {
-            output_memory: [0x00; 512],
-            fake_error: None,
-        };
-        let model = ml::Model {
-            output_activations_len: 512,
-        };
-        let mut ml_coord = MlCoordinator::new(&exec);
-        assert!(ml_coord.execute(&model).is_err());
-    }
-
-    #[test]
-    fn test_execute() {
-        let exec = tests::FakeExecutive {
-            output_memory: [0xAD; 512],
-            fake_error: None,
-        };
-        let model = ml::Model {
-            output_activations_len: 512,
-        };
-        let mut ml_coord = MlCoordinator::new(&exec);
-        assert!(ml_coord.load(&model));
-
-        let res = ml_coord.execute(&model);
-        assert_eq!(res.unwrap()[0], 0xAD)
-    }
-}

apps/system/components/VectorCoreDriver/VectorCoreDriver.camkes

@@ -0,0 +1,5 @@
+component VectorCoreDriver {
+	dataport Buf csr;
+
+	provides VectorCoreInterface vctop;
+}

apps/system/components/VectorCoreDriver/src/driver.c

@@ -0,0 +1,9 @@
+#include <camkes.h>
+#include <stdint.h>
+
+// TODO: Set offsets into memory based on `csr`.
+#define CTRL (csr + 0x0)
+
+void vctop_set_ctrl(uint32_t ctrl) {
+    *((volatile uint32_t*)CTRL) = ctrl;
+}

apps/system/interfaces/MlCoordinatorInterface.camkes

@@ -0,0 +1,3 @@
+procedure MlCoordinatorInterface {
+  void execute();
+};

apps/system/interfaces/VectorCoreInterface.camkes

@@ -0,0 +1,3 @@
+procedure VectorCoreInterface {
+  void set_ctrl(in uint32_t ctrl);
+};

apps/system/system.camkes

@@ -13,10 +13,13 @@
 import <std_connector.camkes>;
 
 import "interfaces/uart.idl4";
+import "interfaces/VectorCoreInterface.camkes";
 import "components/UartDriver/UartDriver.camkes";
 import "components/DebugConsole/DebugConsole.camkes";
 import "components/ProcessManager/ProcessManager.camkes";
+import "components/MlCoordinator/MlCoordinator.camkes";
 import "components/SeL4Debug/SeL4Debug.camkes";
+import "components/VectorCoreDriver/VectorCoreDriver.camkes";
 
 component UART {
     hardware;
@@ -25,6 +28,11 @@ component UART {
     // emits Interrupt interrupt;
 }
 
+component VectorCoreHw {
+    hardware;
+    dataport Buf csr;
+}
+
 assembly {
     composition {
         component UART uart;
@@ -32,8 +40,13 @@ assembly {
         component DebugConsole debug_console;
         component ProcessManager process_manager;
         component SeL4Debug sel4debug;
+        component VectorCoreHw vctop;
+        component VectorCoreDriver vc_drv;
+        component MlCoordinator ml_coordinator;
 
         connection seL4HardwareMMIO uart_mem(from drv.mem, to uart.mem);
+        connection seL4HardwareMMIO vc_csr(from vc_drv.csr, to vctop.csr);
+
         // TODO(mattharvey): Make receives wait on interrupt.
         // connection seL4HardwareInterrupt uart_interrupt(
         //     from uart.interrupt, to drv.interrupt);
@@ -48,6 +61,7 @@ assembly {
         // Connect the LoggerInterface to each component that needs to log
         // to the console.
         connection seL4RPCCall LoggerInterface(from process_manager.logger,
+                                               from ml_coordinator.logger,
                                                to debug_console.logger);
 
         // Connect the SeL4Debug interface of each component that needs access.
@@ -55,6 +69,11 @@ assembly {
                                                   from process_manager.sel4debug,
                                                   to sel4debug.sel4debug);
 
+        connection seL4RPCCall MlCoordinatorInterface(from debug_console.mlcoord,
+                                                      to ml_coordinator.mlcoord);
+
+        connection seL4RPCCall VectorCoreInterface(from ml_coordinator.vctop, to vc_drv.vctop);
+
         connection seL4SharedData tx_channel(
             from debug_console.tx_dataport, to drv.tx_dataport);
         connection seL4SharedData rx_channel(
@@ -67,6 +86,9 @@ assembly {
         // seL4 claims 10 is bigger than irqMax if this is uncommented.
         // uart.interrupt_irq_number = 10;
 
+        vctop.csr_paddr = 0x48000000;
+        vctop.csr_size = 0x1000;
+
         random.ID = 1;
 
         uart.integrity_label = "drv";