Update `align_up` function to return an u64 instead of usize. We also need to check all the usages of `align_up` and validate that they are lossless.