Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix structs test on ARM #505

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open

Fix structs test on ARM #505

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
a330126
libia2: align allocated stacks on AArch64
fw-immunant Jan 13, 2025
4f0c0f8
rewriter: clarify emit_memcpy src_offset parameter
fw-immunant Feb 4, 2025
83b0333
rewriter: format
fw-immunant Feb 4, 2025
a33486b
rewriter: align space used for indirect arguments
fw-immunant Feb 4, 2025
4cb5171
rewriter: comment about copying of indirect arguments
fw-immunant Feb 4, 2025
dbe9929
tests: reenable structs test on AArch64
fw-immunant Feb 25, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions runtime/libia2/init.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,8 +32,12 @@ char *allocate_stack(int i) {
/* Tag the allocated stack pointer so it is accessed with the right pkey */
stack = (char *)((uint64_t)stack | (uint64_t)i << 56);
#endif
#ifdef __aarch64__
return stack + STACK_SIZE - 16;
#elif defined(__x86_64__)
/* Each stack frame start + 8 is initially 16-byte aligned. */
return stack + STACK_SIZE - 8;
#endif
}

void allocate_stack_0() {
Expand Down
3 changes: 1 addition & 2 deletions tests/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@ add_subdirectory(rewrite_fn_ptr_eq)
add_subdirectory(rewrite_macros)
add_subdirectory(sighandler)
add_subdirectory(static_addr_taken)
add_subdirectory(structs)

# The following tests are not supported on ARM64 yet
if (NOT LIBIA2_AARCH64)
Expand All @@ -74,8 +75,6 @@ if (NOT LIBIA2_AARCH64)
add_subdirectory(two_keys_minimal)
add_subdirectory(two_shared_ranges)
add_subdirectory(trusted_indirect)
# AArch64 call gates break this test sometimes
add_subdirectory(structs)
# LLVM patch seems to break this test
add_subdirectory(simple1)
# ARM does not support shared heap allocation yet
Expand Down
72 changes: 43 additions & 29 deletions tools/rewriter/GenCallAsm.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -237,50 +237,50 @@ static void emit_reg_pop(AsmWriter &aw, const ArgLocation &loc) {
}
}

// Emit code to copy `byte_count` bytes from `src` to `dst` using `scratch` as a
// temporary register.
// Emit code to copy `byte_count` bytes from `src_offset` bytes past the address stored in the `src`
// register to the region starting at the `dst` register using `scratch` as a temporary register.
static void emit_memcpy(AsmWriter &aw, unsigned byte_count, const std::string &dst,
const std::string &src, const std::string &scratch, size_t offset, Arch arch) {
const std::string &src, const std::string &scratch, size_t src_offset, Arch arch) {
add_comment_line(aw, "Copy " + std::to_string(byte_count) + " bytes from " + src + " to " + dst);
if (arch == Arch::X86) {
int i = 0;
for (; i + 8 <= byte_count; i += 8) {
add_asm_line(aw, "movq "s + std::to_string(i + offset) + "(%" + src + "), %" + scratch);
add_asm_line(aw, "movq "s + std::to_string(i + src_offset) + "(%" + src + "), %" + scratch);
add_asm_line(aw, "movq %" + scratch + ", "s + std::to_string(i) + "(%" + dst + ")");
}
if (i + 4 <= byte_count) {
add_asm_line(aw, "movl "s + std::to_string(i + offset) + "(%" + src + "), %" + scratch + "d");
add_asm_line(aw, "movl "s + std::to_string(i + src_offset) + "(%" + src + "), %" + scratch + "d");
add_asm_line(aw, "movl %" + scratch + "d, "s + std::to_string(i) + "(%" + dst + ")");
i += 4;
}
if (i + 2 <= byte_count) {
add_asm_line(aw, "movw "s + std::to_string(i + offset) + "(%" + src + "), %" + scratch + "w");
add_asm_line(aw, "movw "s + std::to_string(i + src_offset) + "(%" + src + "), %" + scratch + "w");
add_asm_line(aw, "movw %" + scratch + "w, "s + std::to_string(i) + "(%" + dst + ")");
i += 2;
}
if (i < byte_count) {
add_asm_line(aw, "movb "s + std::to_string(i + offset) + "(%" + src + "), %" + scratch + "b");
add_asm_line(aw, "movb "s + std::to_string(i + src_offset) + "(%" + src + "), %" + scratch + "b");
add_asm_line(aw, "movb %" + scratch + "b, "s + std::to_string(i) + "(%" + dst + ")");
}
} else if (arch == Arch::Aarch64) {
auto halfreg = "W"s + scratch.substr(1);
int i = 0;
for (; i + 8 <= byte_count; i += 8) {
add_asm_line(aw, "ldr "s + scratch + ", [" + src + ", #" + std::to_string(i + offset) + "]");
add_asm_line(aw, "ldr "s + scratch + ", [" + src + ", #" + std::to_string(i + src_offset) + "]");
add_asm_line(aw, "str "s + scratch + ", [" + dst + ", #" + std::to_string(i) + "]");
}
if (i + 4 <= byte_count) {
add_asm_line(aw, "ldr "s + halfreg + ", [" + src + ", #" + std::to_string(i + offset) + "]");
add_asm_line(aw, "ldr "s + halfreg + ", [" + src + ", #" + std::to_string(i + src_offset) + "]");
add_asm_line(aw, "str "s + halfreg + ", [" + dst + ", #" + std::to_string(i) + "]");
i += 4;
}
if (i + 2 <= byte_count) {
add_asm_line(aw, "ldrh "s + halfreg + ", [" + src + ", #" + std::to_string(i + offset) + "]");
add_asm_line(aw, "ldrh "s + halfreg + ", [" + src + ", #" + std::to_string(i + src_offset) + "]");
add_asm_line(aw, "strh "s + halfreg + ", [" + dst + ", #" + std::to_string(i) + "]");
i += 2;
}
if (i < byte_count) {
add_asm_line(aw, "ldrb "s + halfreg + ", [" + src + ", #" + std::to_string(i + offset) + "]");
add_asm_line(aw, "ldrb "s + halfreg + ", [" + src + ", #" + std::to_string(i + src_offset) + "]");
add_asm_line(aw, "strb "s + halfreg + ", [" + dst + ", #" + std::to_string(i) + "]");
}
}
Expand Down Expand Up @@ -525,8 +525,10 @@ static void x86_emit_intermediate_pkru(AsmWriter &aw, uint32_t caller_pkey, uint

static void emit_copy_args(AsmWriter &aw, const std::vector<ArgLocation> &args,
const std::optional<std::vector<ArgLocation>> &wrapper_args,
size_t stack_return_size, size_t stack_return_padding, int stack_alignment,
size_t stack_arg_size, size_t stack_arg_padding, size_t wrapper_stack_arg_size,
size_t stack_return_size, size_t stack_return_padding,
size_t indirect_arg_size, size_t indirect_arg_padding,
int stack_alignment,
size_t stack_arg_size, size_t stack_arg_padding, size_t wrapper_stack_arg_size,
uint32_t caller_pkey, Arch arch) {
if (arch == Arch::X86) {
// When returning via memory, the address of the return value is passed in
Expand Down Expand Up @@ -588,7 +590,7 @@ static void emit_copy_args(AsmWriter &aw, const std::vector<ArgLocation> &args,
src_arg++;
dest_arg++;
}

add_asm_line(aw, "movq %"s + src_arg->as_str() + ", %r12");
src_arg++;
for (; dest_arg != args.end(); src_arg++, dest_arg++) {
Expand All @@ -607,17 +609,7 @@ static void emit_copy_args(AsmWriter &aw, const std::vector<ArgLocation> &args,
}
}
} else if (arch == Arch::Aarch64) {
size_t indirect_arg_size = 0;
for (auto &arg : args) {
if (arg.is_indirect()) {
size_t align = std::max(arg.align(), (size_t)8);
if (stack_arg_size % align != 0) {
indirect_arg_size += align - (stack_arg_size % align);
}
indirect_arg_size += arg.size();
}
}
size_t total_stack_size = stack_return_size + stack_return_padding + indirect_arg_size + stack_arg_size + stack_arg_padding;
size_t total_stack_size = stack_return_size + stack_return_padding + indirect_arg_size + indirect_arg_padding + stack_arg_size + stack_arg_padding;
if (stack_return_size > 0) {
// Reserve space to save the original return value pointer. We only want
// this slot if we are returning a value on the stack.
Expand Down Expand Up @@ -653,6 +645,7 @@ static void emit_copy_args(AsmWriter &aw, const std::vector<ArgLocation> &args,
emit_memcpy(aw, stack_arg_size, "sp", "x12", "x9", src_offset, arch);
}

// Copy memory region for indirect arguments to target stack, updating indirect args themselves
size_t indirect_arg_current_offset = indirect_args_start;
for (auto &arg : args) {
if (arg.is_indirect()) {
Expand All @@ -661,6 +654,9 @@ static void emit_copy_args(AsmWriter &aw, const std::vector<ArgLocation> &args,
if (indirect_arg_current_offset % arg.align() != 0) {
indirect_arg_current_offset += arg.align() - (indirect_arg_current_offset % arg.align());
}
add_comment_line(
aw, llvm::formatv("Copy {} bytes for indirect argument in {} and update it by {}",
arg.size(), arg.as_str(), indirect_arg_current_offset));
if (arg.is_stack()) {
add_asm_line(aw, "ldr x10, [sp, #" + std::to_string(arg.stack_offset()) + "]");
add_asm_line(aw, "add x9, sp, #" + std::to_string(indirect_arg_current_offset));
Expand Down Expand Up @@ -974,6 +970,21 @@ std::string emit_asm_wrapper(AbiSignature sig,
llvm::errs() << "Generating wrapper for " << sig_string(sig, target_name) << "\n";
auto rets = allocate_return_locations(sig, arch);

size_t indirect_arg_size = 0;
for (auto &arg : args) {
if (arg.is_indirect()) { // Only ever true on AArch64 for now.
size_t align = std::max(arg.align(), (size_t)8);
if (stack_arg_size % align != 0) {
indirect_arg_size += align - (stack_arg_size % align);
}
indirect_arg_size += arg.size();
}
}
size_t indirect_arg_padding = 0;
if (indirect_arg_size % 16 != 0) {
indirect_arg_padding = 16 - indirect_arg_size % 16;
}

// std::stringstream ss;
// append_arg_kinds(ss, rets);
// llvm::errs() << "Return kinds: " << ss.str() << "\n";
Expand Down Expand Up @@ -1026,8 +1037,11 @@ std::string emit_asm_wrapper(AbiSignature sig,
| top |Top of the stack (stack grows down on AArch64). This address is
| |aligned to 16 bytes.
+-----+
|ind |Space for arguments passed indirectly (i.e. in memory). These arguments
|args |will point to this region.
|ind |Space for arguments passed indirectly (i.e. in memory). Indirect arguments
|args |are passed as pointers into this region; for access in another compartment
| |we must copy its contents and translate the pointers.
+-----+
|align|Space to align the stack back to 16 bytes after indirect arguments.
+-----+
| |Space for the compartment's return value if it has class MEMORY. This
|ret |space is only allocated if the pointer to the caller's return value
Expand Down Expand Up @@ -1077,7 +1091,7 @@ std::string emit_asm_wrapper(AbiSignature sig,

// Count room for for the ret align padding, return value, and our ret ptr.
size_t compartment_stack_space = start_of_ret_space + stack_return_size +
stack_return_padding + stack_arg_size + stack_arg_padding;
stack_return_padding + indirect_arg_size + indirect_arg_padding + stack_arg_size + stack_arg_padding;
size_t stack_alignment = 0;
if (arch == Arch::X86) {
// Compute what the stack alignment would be before calling the wrapped
Expand Down Expand Up @@ -1120,7 +1134,7 @@ std::string emit_asm_wrapper(AbiSignature sig,

emit_switch_stacks(aw, caller_pkey, target_pkey, arch);

emit_copy_args(aw, args, wrapper_args, stack_return_size, stack_return_padding, stack_alignment, stack_arg_size, stack_arg_padding, wrapper_stack_arg_size, caller_pkey, arch);
emit_copy_args(aw, args, wrapper_args, stack_return_size, stack_return_padding, indirect_arg_size, indirect_arg_padding, stack_alignment, stack_arg_size, stack_arg_padding, wrapper_stack_arg_size, caller_pkey, arch);

emit_scrub_regs(aw, caller_pkey, args, kind == WrapperKind::IndirectCallsite, arch);

Expand Down