Writing Safe Abstractions for Unsafe Rust Code

26 Jan 2025

Introduction

Rust is celebrated for its emphasis on safety and performance, largely thanks to its robust compile-time checks. However, there are situations where you need to bypass these checks to perform low-level operations—this is where Rust’s unsafe keyword comes in. While unsafe opens the door to powerful features, it also comes with significant risks.

The solution?

Encapsulating unsafe code in safe abstractions.

This post explores what that means, why it’s important, and how to do it effectively.

Understanding unsafe in Rust

Rust enforces strict memory safety guarantees by default. However, some operations are inherently unsafe and require explicit acknowledgment from the programmer. These include:

Raw pointer manipulation: Directly accessing memory without bounds or validity checks.
Foreign Function Interface (FFI): Interacting with non-Rust code (e.g., calling C functions).
Manual memory management: Allocating and freeing memory without Rust’s usual safeguards.
Concurrency primitives: Implementing data structures that require custom synchronization logic.

When you write unsafe code, you’re essentially telling the compiler, “I know what I’m doing; trust me.”

While this is sometimes necessary, it’s critical to minimize the potential for misuse by others.

Why Wrap Unsafe Code in Safe Abstractions?

Using unsafe is a trade-off. It gives you access to low-level features and optimizations but requires you to manually uphold the invariants that Rust would otherwise enforce. Safe abstractions address this challenge by:

Avoiding Undefined Behavior: Preventing common pitfalls like null pointer dereferences, data races, or buffer overflows.
Improving Maintainability: Reducing the scattering of unsafe blocks across the codebase makes it easier to audit and debug.
Providing Ease of Use: Enabling most developers to rely on Rust’s safety guarantees without needing to understand the intricacies of the underlying unsafe implementation.

What is a Safe Abstraction?

A safe abstraction is an API or module where the internal implementation may use unsafe code, but the external interface ensures that incorrect usage is either impossible or extremely difficult.

Let’s look at how to create one.

Example: Safe Wrapping of Unsafe Memory Allocation

Here’s a simplified example of wrapping unsafe memory management into a safe abstraction:

pub struct SafeAllocator {
    // Internal raw pointer or other unsafe constructs
    ptr: *mut u8,
    size: usize,
}

impl SafeAllocator {
    pub fn new(size: usize) -> Self {
        let ptr = unsafe { libc::malloc(size) as *mut u8 };
        if ptr.is_null() {
            panic!("Failed to allocate memory");
        }
        Self { ptr, size }
    }

    pub fn allocate(&self, offset: usize, len: usize) -> &[u8] {
        if offset + len > self.size {
            panic!("Out of bounds access");
        }
        unsafe {
            std::slice::from_raw_parts(self.ptr.add(offset), len)
        }
    }

    pub fn deallocate(self) {
        unsafe {
            libc::free(self.ptr as *mut libc::c_void);
        }
    }
}

impl Drop for SafeAllocator {
    fn drop(&mut self) {
        unsafe {
            libc::free(self.ptr as *mut libc::c_void);
        }
    }
}

In this example:

unsafe is confined to specific, well-defined sections of the code.
The API ensures that users cannot misuse the allocator (e.g., by accessing out-of-bounds memory).
Drop ensures memory is automatically freed when the allocator goes out of scope.

Example Usage of `SafeAllocator`

Here’s how you might use the SafeAllocator in practice:

fn main() {
    // Create a new SafeAllocator with 1024 bytes of memory
    let allocator = SafeAllocator::new(1024);

    // Allocate a slice of 128 bytes starting from offset 0
    let slice = allocator.allocate(0, 128);
    println!("Allocated slice of length: {}", slice.len());

    // The allocator will automatically deallocate memory when it goes out of scope
}

This usage demonstrates:

How to create and interact with the SafeAllocator API.
That memory is automatically managed via Rust’s Drop trait, preventing leaks.

Leveraging Rust’s Type System

Rust’s type system is another powerful tool for enforcing invariants. For example, you can use:

Lifetimes: To ensure references don’t outlive the data they point to.
PhantomData: To associate types or lifetimes with otherwise untyped data.
Ownership and Borrowing Rules: To enforce safe access patterns at compile time.

Documentation of Safety Contracts

Any unsafe code should include clear documentation of the invariants it relies on. For example:

// Safety:
// - `ptr` must be non-null and point to a valid memory region.
// - `len` must not exceed the bounds of the allocated memory.
unsafe {
    std::slice::from_raw_parts(ptr, len)
}

This makes it easier for future maintainers to understand and verify the correctness of the code.

Real-World Examples of Safe Abstractions

Many Rust libraries provide excellent examples of safe abstractions over unsafe code:

std::sync::Mutex: Internally uses unsafe for thread synchronization but exposes a safe API for locking and unlocking.
Vec: The Rust standard library’s Vec type uses unsafe for raw memory allocation and resizing but ensures bounds checks and proper memory management externally.
crossbeam: Provides safe concurrency primitives built on low-level atomic operations.

Costs and Benefits

While writing safe abstractions requires extra effort and careful thought, the benefits outweigh the costs:

Benefits:

Reduced Risk of Bugs: Encapsulating unsafe code minimizes the chance of introducing undefined behavior.
Improved Developer Experience: Safe APIs make it easier for others to use your code without worrying about low-level details.
Easier Auditing: With unsafe code isolated, it’s easier to review and verify its correctness.

Costs:

Initial Effort: Designing a robust safe abstraction takes time and expertise.
Performance Overhead: In rare cases, adding safety layers may incur slight overhead (though usually negligible in well-designed abstractions).

Conclusion

Writing safe abstractions for unsafe Rust code is both an art and a science. It involves understanding the invariants of your unsafe code, leveraging Rust’s type system to enforce safety, and documenting your assumptions clearly. By doing so, you can harness the power of unsafe while maintaining Rust’s guarantees of memory safety and concurrency correctness—the best of both worlds.

Writing a Key Value Server in Rust

26 Jan 2025

Introduction

In today’s post, we’ll build a simple key value server; but we’ll do it in an iterative way. We’ll build it up simple and then add safety, concurrency, and networking as we go.

Implementation

Now we’ll get started with our iterations. The finished code will be available at the end of this post.

Baseline

All of our implementations will deal with a KeyValueStore struct. This struct will hold all of the variables that we want to keep track of in our server.

use std::collections::HashMap;

struct KeyValueStore {
    data: HashMap<String, String>,
}

We define data as the in-memory representation of our database. We use String keys and store String values.

Our implementation is very basic. All we’re really doing is shadowing the functionality that HashMap provides.

impl KeyValueStore {
    fn new() -> Self {
        Self {
            data: HashMap::new(),
        }
    }

    fn insert(&mut self, key: String, value: String) {
        self.data.insert(key, value);
    }

    fn get(&self, key: &str) -> Option<&String> {
        self.data.get(key)
    }

    fn delete(&mut self, key: &str) {
        self.data.remove(key);
    }
}

This is a pretty decent starting point. We can use our KeyValueStore in some basic tests:

fn main() {
    let mut store = KeyValueStore::new();
    store.insert("key1".to_string(), "value1".to_string());
    println!("{:?}", store.get("key1"));
    store.delete("key1");
    println!("{:?}", store.get("key1"));
}

Variants

String is pretty limiting to store as far as the value side is concerned. We can upgrade this to specifically use data types that we will find useful via an enum:

#[derive(Debug, Clone)]
enum Value {
    String(String),
    Integer(i64),
    Float(f64),
    Boolean(bool),
    Binary(Vec<u8>),
    // Add more variants as needed
}

We can swap out the value side of our data member now, too.

struct KeyValueStore {
    data: HashMap<String, Value>,
}

The implementation simply swaps the String for Value:

impl KeyValueStore {
    fn new() -> Self {
        Self {
            data: HashMap::new(),
        }
    }

    fn insert(&mut self, key: String, value: Value) {
        self.data.insert(key, value);
    }

    fn get(&self, key: &str) -> Option<&Value> {
        self.data.get(key)
    }

    fn delete(&mut self, key: &str) {
        self.data.remove(key);
    }
}

We’re now able to not only store strings. We can store integers, floats, binary, and booleans. This makes our key value store a lot more versatile.

Thread Safety

We will have multiple threads of execution trying to perform actions on this structure at the same time, so we will add some thread safety to the process now. Wrapping data in Arc will give us a thread safe, reference counting pointer. We’re also going to need to lock this data structure for reading and for writing. We can use RwLock to take care of that for us.

We update our data structure to include these new types:

struct KeyValueStore {
    data: Arc<RwLock<HashMap<String, Value>>>,
}

Now our implementation functions need to change to work with these new structures. We can keep the structure of functions the same though.

impl KeyValueStore {
    fn new() -> Self {
        Self {
            data: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    fn insert(&self, key: String, value: Value) {
        let mut locked = self.data.write().unwrap();
        locked.insert(key, value);
    }

    fn get(&self, key: &str) -> Option<Value> {
        let mut locked = self.data.read().unwrap();
        locked.get(key).cloned()
    }

    fn delete(&self, key: &str) {
        let mut locked = self.data.write().unwrap();
        locked.remove(key);
    }
}

These functions are now safe, which means calling code can be multithreaded and we can guaranteed that our data structure will be treated consistently.

fn main() {
    let store = Arc::new(KeyValueStore::new());

    // Create a vector to hold thread handles
    let mut handles = vec![];

    // Spawn threads to perform inserts
    for i in 0..5 {
        let store = Arc::clone(&store);
        let handle = thread::spawn(move || {
            let key = format!("key{}", i);
            let value = Value::Integer(i * 10);
            store.insert(key.clone(), value);
            println!("Thread {} inserted: {}", i, key);
        });
        handles.push(handle);
    }

    // Spawn threads to read values
    for i in 0..5 {
        let store = Arc::clone(&store);
        let handle = thread::spawn(move || {
            let key = format!("key{}", i);
            if let Some(value) = store.get(&key) {
                println!("Thread {} read: {} -> {:?}", i, key, value);
            } else {
                println!("Thread {} could not find: {}", i, key);
            }
        });
        handles.push(handle);
    }

    // Spawn threads to delete keys
    for i in 0..5 {
        let store = Arc::clone(&store);
        let handle = thread::spawn(move || {
            let key = format!("key{}", i);
            store.delete(&key);
            println!("Thread {} deleted: {}", i, key);
        });
        handles.push(handle);
    }

    // Wait for all threads to complete
    for handle in handles {
        handle.join().unwrap();
    }

    println!("Final state of the store: {:?}", store.data.read().unwrap());
}

Error handling

You can see that we’re using unwrap in the implementation functions, which might be ok for tests or short scripts. If we’re going to expect to run this code in production, we’d be best replacing these with actual error handling counterparts.

In order to do that, we need to define our error domain first. We create an enum called StoreError. As we fill out our implementation, we’ll run into a number of different error cases. We’ll use StoreError to centralise all of these errors so we can express them clearly.

#[derive(Debug)]
enum StoreError {
    LockError(String),
    KeyNotFound(String),
}

impl<T> From<PoisonError<T>> for StoreError {
    fn from(err: PoisonError<T>) -> Self {
        StoreError::LockError(format!("Lock poisoned: {}", err))
    }
}

We’ve implemented PoisonError for our StoreError because the PoisonError type is an error which can be returned whenever a lock is acquired. If something goes wrong and we’ve acquired a lock, it’s a PoisonError that’s used.

Our insert, get, and delete methods now need an upgrade. We’ll be returning Result<T, E> values from our functions now to accomodate potential failures.

fn insert(&self, key: String, value: Value) -> Result<(), StoreError> {
    let mut locked = self.data.write()?;
    locked.insert(key, value);
    Ok(())
}

fn get(&self, key: &str) -> Result<Option<Value>, StoreError> {
    let locked = self.data.read()?;
    Ok(locked.get(key).cloned()) // Clone the value to return an owned copy
}

fn delete(&self, key: &str) -> Result<(), StoreError> {
    let mut locked = self.data.write()?;
    if locked.remove(key).is_none() {
        return Err(StoreError::KeyNotFound(key.to_string()));
    }
    Ok(())
}

We’ve removed the use of unwrap now, swapping out to using the ? operator. This will allow us to actually handle any failure that is bubbled out of calling code.

Using the File System

We need to be able to persist the state of our key value store out to disk for durability. In order to do this, we need to keep track of where we’ll write the file. We add a file_path member to our structure:

struct KeyValueStore {
    data: Arc<RwLock<HashMap<String, Value>>>,
    file_path: Option<String>,
}

impl KeyValueStore {
    fn new(file_path: Option<String>) -> Self {
        Self {
            data: Arc::new(RwLock::new(HashMap::new())),
            file_path,
        }
    }
}

Starting out this implementation simply, we just write a load and save function that we can call at any time. Before we do this we need some extra dependencies added for serialisation:

[dependencies]
serde = { version = "1.0.217", features = ["derive"] }
serde_json = "1.0.137"

This will allow us to reduce our internal state to JSON.

Loading the database off disk

/// Load the state from a file
fn load(&self) -> Result<(), StoreError> {
    if let Some(ref path) = self.file_path {
        match fs::read_to_string(path) {
            Ok(contents) => {
                let deserialized: HashMap<String, Value> = serde_json::from_str(&contents)?;
                let mut locked = self.data.write()?;
                *locked = deserialized; // Replace the current state with the loaded one
                Ok(())
            }
            Err(e) if e.kind() == ErrorKind::NotFound => {
                // File doesn't exist, just return Ok (no data to load)
                Ok(())
            }
            Err(e) => Err(e.into()),
        }
    } else {
        Err(StoreError::IoError("File path not set".to_string()))
    }
}

We need to make sure that a file_path was specified. We read everything off from the file into contents as a big string. Using serde_json::from_str we can turn that contents into the deserialised representation. From there, we simply swap out the underlying content.

We’ve got some new errors to deal with here in IoError.

#[derive(Debug)]
enum StoreError {
    LockError(String),
    KeyNotFound(String),
    IoError(String),
    SerdeError(String),
}

This will be used for our write implementation which looks like this:

/// Save the current state to a file
fn save(&self) -> Result<(), StoreError> {
    if let Some(ref path) = self.file_path {
        let locked = self.data.read()?;
        let serialized = serde_json::to_string(&*locked)?;
        fs::write(path, serialized)?;
        Ok(())
    } else {
        Err(StoreError::IoError("File path not set".to_string()))
    }
}

The magic here really is the serde_json::to_string taking our internal state and writing it as json.

An example of how this looks is like this:

{
    "key2":{"Integer":20},
    "key4":{"Integer":40},
    "key1":{"Integer":10},
    "key3":{"Integer":30},
    "key0":{"Integer":0}
}

Networking

Finally, we’ll add some networking to the solution. A really basic network interface will allow remote clients to perform the get, set, and delete operations for us.

The handle_client function is the heart of the server process, performing the needed processing on incoming requests and routing them to the database instance:

fn handle_client(mut stream: TcpStream, store: Arc<KeyValueStore>) {
    let mut buffer = [0; 512];

    // Read the incoming request
    match stream.read(&mut buffer) {
        Ok(_) => {
            let request = String::from_utf8_lossy(&buffer);
            let mut parts = request.trim().split_whitespace();
            let command = parts.next();

            let response = match command {
                Some("SET") => {
                    let key = parts.next().unwrap_or_default().to_string();
                    let value = parts.next().unwrap_or_default().to_string();
                    store.insert(key, Value::String(value));
                    "OK\n".to_string()
                }
                Some("GET") => {
                    let key = parts.next().unwrap_or_default();
                    if let Ok(Some(value)) = store.get(key) {
                        format!("{:?}\n", value)
                    } else {
                        "Key not found\n".to_string()
                    }
                }
                Some("DEL") => {
                    let key = parts.next().unwrap_or_default();
                    store.delete(key);
                    "OK\n".to_string()
                }
                _ => "Unknown command\n".to_string(),
            };

            // Send the response back to the client
            stream.write_all(response.as_bytes()).unwrap();
        }
        Err(e) => eprintln!("Failed to read from socket: {}", e),
    }
}

Out networking “protocol” looks like this:

-- set the key "key1" to the value "hello"
SET key1 hello

-- get the value of the key "key1"
GET key1

-- remove the value and key "key1"
DEL key1

This is all made possible by the following:

let request = String::from_utf8_lossy(&buffer);
let mut parts = request.trim().split_whitespace();
let command = parts.next();

We read in the request data from the client into request. This gets split up on white spaces into parts with command given the first of these parts. The code is expecting command to be either SET, GET, or DEL that is then handled in the following pattern match.

This function gets mounted onto the server in the main function which now looks like this:

fn main() {
    let store = Arc::new(
        KeyValueStore::new(None)
    );
    let listener = TcpListener::bind("127.0.0.1:7878").unwrap();

    println!("Server running on 127.0.0.1:7878");

    for stream in listener.incoming() {
        match stream {
            Ok(stream) => {
                let store = Arc::clone(&store);
                std::thread::spawn(move || handle_client(stream, store));
            }
            Err(e) => eprintln!("Connection failed: {}", e),
        }
    }
}

We’re starting our server on port 7878 and handling each connection with our handle_client function.

Running this and giving it a test with telnet gives us the following:

➜  telnet 127.0.0.1 7878
Trying 127.0.0.1...
Connected to 127.0.0.1.
Escape character is '^]'.
SET key1 hello
OK
Connection closed by foreign host.

➜  telnet 127.0.0.1 7878
Trying 127.0.0.1...
Connected to 127.0.0.1.
Escape character is '^]'.
GET key1
String("hello")
Connection closed by foreign host.

➜  telnet 127.0.0.1 7878
Trying 127.0.0.1...
Connected to 127.0.0.1.
Escape character is '^]'.
DEL key1
OK
Connection closed by foreign host.

➜  telnet 127.0.0.1 7878
Trying 127.0.0.1...
Connected to 127.0.0.1.
Escape character is '^]'.
GET key1
Key not found
Connection closed by foreign host.

So, it works. It’s crude and needs to be patched to be a little more production ready than this - but this is a start.

Conclusion

In this article, we walked through building a thread-safe, persistent key-value store in Rust. We started with a simple in-memory implementation and iteratively improved it by:

Adding support for multiple data types using an enum.
Ensuring thread safety with RwLock and Arc.
Replacing unwrap with proper error handling.
Adding file persistence using JSON serialization and deserialization.
Added some basic network access

This provides a solid foundation for a more robust and scalable key-value server. Next steps could include:

Implementing advanced features like snapshots or replication.
Optimizing for performance with tools like async I/O or a custom storage engine.
Partial reads and memory mapping
Clustering

The full implementation can be found here.

Implementing an ML Model in Rust

26 Jan 2025

Introduction

Rust, known for its performance, memory safety, and low-level control, is gaining traction in domains traditionally dominated by Python, such as machine learning (ML). While Python is the go-to for prototyping ML models due to its mature ecosystem, Rust shines in scenarios demanding high performance, safety, and seamless system-level integration.

In this post, we’ll explore how to implement logistic regression in Rust and discuss the implications of the model’s output.

Why use Rust?

Before diving into code, it’s worth asking: why choose Rust for ML when Python’s libraries like TensorFlow and PyTorch exist?

Benefits of Rust:

Performance: Rust offers near-C speeds, making it ideal for performance-critical tasks.
Memory Safety: Its ownership model ensures memory safety, preventing bugs like segmentation faults and data races.
Integration: Rust can easily integrate with low-level systems, making it a great choice for embedding ML models into IoT, edge devices, or game engines.
Control: Rust provides fine-grained control over execution, allowing developers to optimize their models at a deeper level.

While Rust’s ML ecosystem is still evolving, libraries like ndarray, linfa, and smartcore provide foundational tools for implementing machine learning models.

Logistic Regression

Logistic regression is a simple yet powerful algorithm for binary classification. It predicts whether a data point belongs to class 0 or 1 based on a weighted sum of features passed through a sigmoid function.

Below is a Rust implementation of logistic regression using the ndarray crate for numerical operations.

use ndarray::{Array2, Array1};
use ndarray_rand::RandomExt;
use ndarray_rand::rand_distr::Uniform;

fn sigmoid(x: f64) -> f64 {
    1.0 / (1.0 + (-x).exp())
}

fn logistic_regression(X: &Array2<f64>, y: &Array1<f64>, learning_rate: f64, epochs: usize) -> Array1<f64> {
    let (n_samples, n_features) = X.dim();
    let mut weights = Array1::<f64>::random(n_features, Uniform::new(-0.01, 0.01));
    let mut bias = 0.0;

    for _ in 0..epochs {
        let linear_model = X.dot(&weights) + bias;
        let predictions = linear_model.mapv(sigmoid);

        // Compute the error
        let error = &predictions - y;

        // Compute gradients
        let gradient_weights = X.t().dot(&error) / n_samples as f64;
        let gradient_bias = error.sum() / n_samples as f64;

        // Update weights and bias
        weights -= &(learning_rate * gradient_weights);
        bias -= learning_rate * gradient_bias;
    }

    weights
}

fn main() {
    let X = Array2::random((100, 2), Uniform::new(-1.0, 1.0)); // Random features
    let y = Array1::random(100, Uniform::new(0.0, 1.0)).mapv(|v| if v > 0.5 { 1.0 } else { 0.0 }); // Random labels

    let weights = logistic_regression(&X, &y, 0.01, 1000);
    println!("Trained Weights: {:?}", weights);
}

Key Concepts:

Sigmoid Function: Converts the linear combination of inputs into a value between 0 and 1.
Gradient Descent: Updates weights and bias iteratively to minimize the error between predictions and actual labels.
Random Initialization: Weights start with small random values and are fine-tuned during training.

Output

When you run the code, you’ll see output similar to this:

Trained Weights: [0.034283492207871635, 0.3083430316223569], shape=[2], strides=[1], layout=CFcf (0xf), const ndim=1

What Does This Mean?

Weights: Each weight corresponds to a feature in your dataset. For example, with 2 input features, the model learns two weights.
- A positive weight means the feature increases the likelihood of predicting 1.
- A negative weight means the feature decreases the likelihood of predicting 1.
Bias (Optional): The bias adjusts the decision boundary and accounts for data not centered at zero. To view the bias, modify the print statement:

println!("Trained Weights: {:?}, Bias: {}", weights, bias);

Predictions: To test the model, use new data points and calculate their predictions:

let new_data = array![0.5, -0.2];
let linear_combination = new_data.dot(&weights) + bias;
let prediction = sigmoid(linear_combination);
println!("Prediction Probability: {}", prediction);

Predictions close to 1 indicate class 1, while predictions close to 0 indicate class 0.

Why Does This Matter?

This simple implementation demonstrates the flexibility and control Rust provides for machine learning tasks. While Python excels in rapid prototyping, Rust’s performance and safety make it ideal for deploying models in production, especially in resource-constrained or latency-critical environments.

When Should You Use Rust for ML?

Rust is a great choice if:

Performance is critical: For example, in real-time systems or embedded devices.
Memory safety is a priority: Rust eliminates common bugs like memory leaks.
Integration with system-level components is needed: Rust can seamlessly work in environments where Python may not be ideal.
Custom ML Implementations: You want more control over how the algorithms are built and optimized.

For research or quick prototyping, Python remains the best choice due to its rich ecosystem and community. However, for production-grade systems, Rust’s strengths make it a compelling alternative.

Conclusion

While Rust’s machine learning ecosystem is still maturing, it’s already capable of handling fundamental ML tasks like logistic regression. By combining performance, safety, and control, Rust offers a unique proposition for ML developers building high-performance or production-critical applications.

Writing an ARM Bootloader

24 Jan 2025

Introduction

In this blog post, we’ll dive into the fascinating world of ARM assembly programming by writing and running a basic bootloader. ARM’s dominance in mobile and embedded systems makes it an essential architecture to understand for developers working with low-level programming or optimization.

We’re going to use QEMU, an open-source emulator, we can develop and test our code right on your PC. So we won’t need any hardware (just yet).

What is ARM?

ARM, short for Advanced RISC Machine, is a family of Reduced Instruction Set Computing (RISC) architectures. ARM processors power billions of devices, from smartphones and tablets to embedded systems and IoT devices. Its popularity stems from its energy efficiency and simplicity compared to complex instruction set computing (CISC) architectures like x86.

Why Emulation?

While ARM assembly is usually executed on physical devices, emulation tools like QEMU allow you to:

Test code without requiring hardware.
Experiment with different ARM-based architectures and peripherals.
Debug programs more effectively using tools like GDB.

Supported ARM Hardware

Before we begin coding, let’s take a brief look at some popular ARM-based platforms:

Raspberry Pi: A widely used single-board computer.
BeagleBone Black: A powerful option for embedded projects.
STM32 Microcontrollers: Common in IoT and robotics applications.

Setup

Before we begin, we need to setup our development and build environment. I’m using Manjaro so package names might be slightly different for your distro of choice.

QEMU emulates a variety of hardware architectures, including ARM.

sudo pacman -Ss qemu-system-arm

Now we need to install the ARM toolchain which will include the assembler (as), linker (ld), and other essential tools.

sudo pacman -S arm-none-eabi-gcc binutils

Now you should have everything you need to get going.

Bootloader

Our goal is to write a minimal ARM assembly program that outputs “Hello, World!” via the UART interface.

The Code

Here is the source code for our bootloader, saved as boot.s:

.section .text
.global _start

_start:
    LDR R0, =0x101f1000     @ UART0 base address
    LDR R1, =message        @ Address of the message
    LDR R2, =message_end    @ Address of the end of the message

loop:
    LDRB R3, [R1], #1       @ Load a byte from the message and increment the pointer
    CMP R1, R2              @ Check if we’ve reached the end of the message
    BEQ done                @ If yes, branch to done
    STRB R3, [R0]           @ Output the character to UART0
    B loop                  @ Repeat for the next character

done:
    B done                  @ Infinite loop to prevent execution from going beyond

message:
    .asciz "Hello, World!\n"  @ Null-terminated string
message_end:

Breaking this down line by line, we get:

LDR R0, =0x101f1000: Load the memory address of UART0 (used for serial output) into register R0.
LDR R1, =message: Load the starting address of the message into R1.
LDR R2, =message_end: Load the end address of the message into R2.

After this setup, we move into a loop.

Load a byte from the message (R3).
Compare R1 (current pointer) with R2 (end of the message).
Write the character to UART0 and repeat.

Finally, we finish up with an infinite loop to prevent the program from running into uninitialized memory.

Building

First we need to assemble the code into an object file:

arm-none-eabi-as -o boot.o boot.s

Next, we link the object file to produce an executable:

arm-none-eabi-ld -Ttext=0x10000 -o boot.elf boot.o

The -Ttext=0x10000 flag specifies the memory address where the program will start executing.

Running

We can give our bootloader a go now using the versatilepb machine in QEMU:

qemu-system-arm -M versatilepb -nographic -kernel boot.elf

-nographic here redirects UART ouput to the terminal, so we should see:

Hello, World!

Debugging

If you run into problems with your program, you do have an option to attach gdb for debugging:

qemu-system-arm -M versatilepb -kernel boot.elf -S -gdb tcp::1234

You can then connect to gdb with the following:

arm-none-eabi-gdb boot.elf
target remote :1234

Deployment

For deployment, we’ll use a Raspberry Pi as an example. This process is similar for other ARM-based boards.

Flashing

First, we need to convert the ELF file to a raw binary format suitable for booting:

arm-none-eabi-objcopy -O binary boot.elf boot.bin

Use a tool like dd to write the binary to an SD card:

dd if=boot.bin of=/dev/sdX bs=512 seek=2048

Running

Insert the SD card into the board.
Power up the device and connect to its UART output (e.g., using a USB-to-serial adapter).
You should see “Hello, World!” printed on the serial console.

Conclusion

Congratulations! You’ve successfully written, emulated, and deployed a simple ARM bootloader. Along the way, you learned:

How to write and debug ARM assembly.
How to use QEMU for emulation.
How to deploy code to real hardware.

From here, you can explore more advanced topics like interrupts, timers, or even writing a simple operating system kernel. The journey into ARM assembly has just begun!

Writing a Custom Loss Function for a Neural Network

24 Jan 2025

Introdution

Loss functions are the unsung heroes of machine learning. They guide the learning process by quantifying the difference between the predicted and actual outputs. While frameworks like PyTorch and TensorFlow offer a plethora of standard loss functions such as Cross-Entropy and Mean Squared Error, there are times when a custom loss function is necessary.

In this post, we’ll explore the why and how of custom loss functions by:

Setting up a simple neural network.
Using standard loss functions to train the model.
Introducing and implementing custom loss functions tailored to specific needs.

Pre-reqs

Before we begin, you’ll need to setup a python project and install some dependencies. We’ll be using PyTorch and torchvision. To install these dependencies, use the following command:

pip install torch torchvision

Once installed, verify the installation by running:

python -c "import torch; print(torch.__version__)"

Network Setup

Let’s start by creating a simple neural network to classify data. For simplicity, we’ll use a toy dataset like the MNIST digits dataset.

Dataet preparation

Use the MNIST dataset (handwritten digits) as an example.
Normalize the dataset for faster convergence during training.

import torch
import torch.optim as optim
from torchvision import datasets, transforms

# Data preparation
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
train_data = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
train_loader = torch.utils.data.DataLoader(train_data, batch_size=64, shuffle=True)

Model Architecture

Input layer flattens the 28x28 pixel images into a single vector.
Two hidden layers with 128 and 64 neurons, each followed by a ReLU activation.
An output layer with 10 neurons (one for each digit) and no activation (handled by the loss function).

# Simple Neural Network
import torch.nn as nn

class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(28 * 28, 128)
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, 10)

    def forward(self, x):
        x = x.view(x.size(0), -1)  # Flatten the input
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

Training Setup:

Use an optimizer (e.g., Adam) and CrossEntropyLoss for training.
Loop over the dataset for a fixed number of epochs, computing loss and updating weights.

# Initialize model, optimizer, and device
model = SimpleNN()
optimizer = optim.Adam(model.parameters(), lr=0.001)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

Standard Loss

Let’s train the model using the standard Cross-Entropy Loss, which is suitable for classification tasks.

Combines log_softmax and negative log likelihood into one step.
Suitable for classification tasks as it penalizes incorrect predictions heavily.

# Standard loss function
criterion = nn.CrossEntropyLoss()

# Training loop
def train_model(model, train_loader, criterion, optimizer, epochs=5):
    model.train()
    for epoch in range(epochs):
        total_loss = 0
        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)

            # Forward pass
            outputs = model(images)
            loss = criterion(outputs, labels)

            # Backward pass and optimization
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total_loss += loss.item()

        print(f'Epoch {epoch+1}/{epochs}, Loss: {total_loss/len(train_loader):.4f}')

train_model(model, train_loader, criterion, optimizer)

The output of this training session should look something like this:

Epoch 1/5, Loss: 0.3932
Epoch 2/5, Loss: 0.1834
Epoch 3/5, Loss: 0.1352
Epoch 4/5, Loss: 0.1054
Epoch 5/5, Loss: 0.0914

Custom Loss

Why Custom Loss Functions?

Standard loss functions may not work well in cases like:

Imbalanced Datasets: Classes have significantly different frequencies.
Multi-Task Learning: Different tasks require different weights.
Task-Specific Goals: Optimizing for metrics like precision or recall rather than accuracy.

Example: Weighted Loss

Suppose we want to penalize misclassifying certain classes more heavily. We can achieve this by implementing a weighted Cross-Entropy Loss.

# Custom weighted loss function
class WeightedCrossEntropyLoss(nn.Module):
    def __init__(self, class_weights):
        super(WeightedCrossEntropyLoss, self).__init__()
        self.class_weights = torch.tensor(class_weights).to(device)

    def forward(self, outputs, targets):
        log_probs = torch.log_softmax(outputs, dim=1)
        loss = -torch.sum(self.class_weights[targets] * log_probs[range(len(targets)), targets]) / len(targets)
        return loss

# Example: Higher weight for class 0
class_weights = [2.0 if i == 0 else 1.0 for i in range(10)]
custom_criterion = WeightedCrossEntropyLoss(class_weights)

# Training with custom loss function
train_model(model, train_loader, custom_criterion, optimizer)

After running this, you should see output like the following:

Epoch 1/5, Loss: 0.4222
Epoch 2/5, Loss: 0.1970
Epoch 3/5, Loss: 0.1390
Epoch 4/5, Loss: 0.1124
Epoch 5/5, Loss: 0.0976

Example: Combining Losses

Sometimes, you might want to combine multiple objectives into a single loss function.

# Custom loss combining Cross-Entropy and L1 regularization
class CombinedLoss(nn.Module):
    def __init__(self, alpha=0.1):
        super(CombinedLoss, self).__init__()
        self.ce_loss = nn.CrossEntropyLoss()
        self.alpha = alpha

    def forward(self, outputs, targets, model):
        ce_loss = self.ce_loss(outputs, targets)
        l1_loss = sum(torch.sum(torch.abs(param)) for param in model.parameters())
        return ce_loss + self.alpha * l1_loss

custom_criterion = CombinedLoss(alpha=0.01)

# Training with combined loss
train_model(model, train_loader, lambda outputs, targets: custom_criterion(outputs, targets, model), optimizer)

Comparing Results

To compare the results of standard and custom loss functions, you need to evaluate the following:

Training Loss:
- Plot the loss per epoch for both standard and custom loss functions.
Accuracy:
- Measure training and validation accuracy after each epoch.
- Compare how well the model performs in predicting each class.
Precision and Recall:
- Useful for imbalanced datasets to measure performance on minority classes.
Visualization:
- Confusion matrix: Visualize how often each class is misclassified.
- Loss curve: Show convergence speed and stability for different loss functions.

We can use graphs to visualise how these metrics perform:

from sklearn.metrics import classification_report, confusion_matrix
import matplotlib.pyplot as plt
import numpy as np

# After training
model.eval()
all_preds, all_labels = [], []
with torch.no_grad():
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        preds = torch.argmax(outputs, dim=1)
        all_preds.extend(preds.cpu().numpy())
        all_labels.extend(labels.cpu().numpy())

# Confusion Matrix
cm = confusion_matrix(all_labels, all_preds)
plt.imshow(cm, cmap='Blues')
plt.title('Confusion Matrix')
plt.colorbar()
plt.show()

# Classification Report
print(classification_report(all_labels, all_preds))

We can also produce visualisations of our loss curves:

# Assuming loss values are stored during training
plt.plot(range(len(train_losses)), train_losses, label="Standard Loss")
plt.plot(range(len(custom_losses)), custom_losses, label="Custom Loss")
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.title('Loss Curve')
plt.show()

Conclusion

Custom loss functions empower you to fine-tune your neural networks for unique problems. By carefully designing and experimenting with loss functions, you can align your model’s learning process with the specific goals of your application.

Some closing tips for custom loss functions:

Always start with a simple baseline (e.g., Cross-Entropy Loss) to understand your model’s behavior.
Visualize performance across metrics, especially when using weighted or multi-objective losses.
Experiment with different weights and loss combinations to find the optimal setup for your task.

The key is to balance complexity and interpretability—sometimes, even simple tweaks can significantly impact performance.

Older Newer

Cogs and Levers A blog full of technical stuff

Writing Safe Abstractions for Unsafe Rust Code

Introduction

Understanding unsafe in Rust

Why Wrap Unsafe Code in Safe Abstractions?

What is a Safe Abstraction?

Example: Safe Wrapping of Unsafe Memory Allocation

Example Usage of SafeAllocator

Leveraging Rust’s Type System

Documentation of Safety Contracts

Real-World Examples of Safe Abstractions

Costs and Benefits

Benefits:

Costs:

Conclusion

Writing a Key Value Server in Rust

Introduction

Implementation

Baseline

Variants

Thread Safety

Error handling

Using the File System

Networking

Conclusion

Implementing an ML Model in Rust

Introduction

Why use Rust?

Logistic Regression

Output

What Does This Mean?

Why Does This Matter?

When Should You Use Rust for ML?

Conclusion

Writing an ARM Bootloader

Introduction

What is ARM?

Why Emulation?

Supported ARM Hardware

Setup

Bootloader

The Code

Building

Running

Debugging

Deployment

Flashing

Running

Conclusion

Writing a Custom Loss Function for a Neural Network

Introdution

Pre-reqs

Network Setup

Dataet preparation

Model Architecture

Training Setup:

Standard Loss

Custom Loss

Why Custom Loss Functions?

Example: Weighted Loss

Example: Combining Losses

Comparing Results

Conclusion

Example Usage of `SafeAllocator`