lisyarus/pslang
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Implement array-types function arguments & return values in Aarch64 compiler

Browse source
This commit is contained in:
Nikita Lisitsa 2026-03-29 22:52:48 +03:00
parent 1404437324
commit e771542773
2 changed files with 72 additions and 52 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
examples/ir_test.psl
View file

@ -63,9 +63,21 @@ struct ray:
func intersect_plane(ray: ray, normal: vec3, value: f32) -> f32: func intersect_plane(ray: ray, normal: vec3, value: f32) -> f32:
return (value - dot(ray.origin, normal)) / dot(ray.direction, normal) return (value - dot(ray.origin, normal)) / dot(ray.direction, normal)
func test() -> i32[3]:
return [70, 60, 50]
func print_i32_3(a: i32[3]):
print('[')
print_i32(a[0])
print(',')
print_i32(a[1])
print(',')
print_i32(a[2])
print(']')
mut a = [1, 2, 3] mut a = [1, 2, 3]
a = test()
a[0] = 10 a[0] = 10
let p = a as i32* print_i32_3(test())
print_i32(p[1])
print('\n') print('\n')

94
libs/jit/source/arch/aarch64/compiler_v2.cpp
View file

@ -61,7 +61,7 @@ namespace pslang::jit::aarch64
return 1 << mode; return 1 << mode;
} }
std::optional<hfa_data> get_hfa_data(local_context & lcontext, ast::struct_definition const * node); std::optional<hfa_data> get_hfa_data(local_context & lcontext, types::type_ptr const & type);
std::optional<hfa_data> compute_hfa_data(local_context & lcontext, ast::struct_definition const * node) std::optional<hfa_data> compute_hfa_data(local_context & lcontext, ast::struct_definition const * node)
{ {
@ -71,21 +71,9 @@ namespace pslang::jit::aarch64
for (std::size_t i = 0; i < node->fields.size(); ++i) for (std::size_t i = 0; i < node->fields.size(); ++i)
{ {
auto const & field = node->fields[i]; auto const & field = node->fields[i];
if (types::is_builtin_type(*field.inferred_type))
{
if (!types::is_floating_point_type(*field.inferred_type))
return std::nullopt;
if (type && !types::equal(*type, *field.inferred_type))
return std::nullopt;
type = field.inferred_type;
++count;
}
else if (auto struct_type = std::get_if<types::struct_type>(field.inferred_type.get()))
{
// NB: recursion must be impossible due to prior checks in type checker // NB: recursion must be impossible due to prior checks in type checker
if (auto subdata = get_hfa_data(lcontext, struct_type->node)) if (auto subdata = get_hfa_data(lcontext, field.inferred_type))
{ {
if (type && !types::equal(*type, *subdata->element_type)) if (type && !types::equal(*type, *subdata->element_type))
return std::nullopt; return std::nullopt;
@ -96,9 +84,6 @@ namespace pslang::jit::aarch64
else else
return std::nullopt; return std::nullopt;
} }
else
return std::nullopt;
}
if (count == 0) if (count == 0)
return std::nullopt; return std::nullopt;
@ -106,15 +91,29 @@ namespace pslang::jit::aarch64
return hfa_data{type, count}; return hfa_data{type, count};
} }
std::optional<hfa_data> get_hfa_data(local_context & lcontext, ast::struct_definition const * node) std::optional<hfa_data> get_hfa_data(local_context & lcontext, types::type_ptr const & type)
{ {
if (auto it = lcontext.struct_hfa.find(node); it != lcontext.struct_hfa.end()) if (auto struct_type = std::get_if<types::struct_type>(type.get()))
{
if (auto it = lcontext.struct_hfa.find(struct_type->node); it != lcontext.struct_hfa.end())
return it->second; return it->second;
auto result = compute_hfa_data(lcontext, node); auto result = compute_hfa_data(lcontext, struct_type->node);
lcontext.struct_hfa[node] = result; lcontext.struct_hfa[struct_type->node] = result;
return result; return result;
} }
else if (auto array_type = std::get_if<types::array_type>(type.get()))
{
if (auto subdata = get_hfa_data(lcontext, array_type->element_type))
return hfa_data{subdata->element_type, subdata->count * array_type->size};
else
return std::nullopt;
}
else if (types::is_floating_point_type(*type))
return hfa_data{type, 1};
else
return std::nullopt;
}
struct populate_const_data_visitor struct populate_const_data_visitor
{ {
@ -737,16 +736,20 @@ namespace pslang::jit::aarch64
template <typename Node, typename DoCall> template <typename Node, typename DoCall>
void apply_call(ir::node_ref it, Node const & node, types::type_ptr const & type, DoCall && do_call) void apply_call(ir::node_ref it, Node const & node, types::type_ptr const & type, DoCall && do_call)
{ {
// TODO: array arguments?
// TODO: handle the case when there weren't enough registers // TODO: handle the case when there weren't enough registers
std::uint8_t reg = 0; std::uint8_t reg = 0;
std::uint8_t fp_reg = 0; std::uint8_t fp_reg = 0;
for (auto const & argument : node.arguments) for (auto const & argument : node.arguments)
{ {
if (auto struct_type = std::get_if<types::struct_type>(argument->inferred_type.get())) auto struct_type = std::get_if<types::struct_type>(argument->inferred_type.get());
auto array_type = std::get_if<types::array_type>(argument->inferred_type.get());
if (struct_type || array_type)
{ {
auto node = struct_type->node; // NB: fixed-size arrays are handled in the same way
if (auto hfa = get_hfa_data(lcontext, node); hfa && hfa->count <= 4) // as structs of N identical fields
auto size = ast::type_size(*argument->inferred_type);
if (auto hfa = get_hfa_data(lcontext, argument->inferred_type); hfa && hfa->count <= 4)
{ {
// HFA - passed in consecutive FP registers // HFA - passed in consecutive FP registers
std::int32_t base_offset = stack_size - stack_position.at(argument); std::int32_t base_offset = stack_size - stack_position.at(argument);
@ -755,16 +758,14 @@ namespace pslang::jit::aarch64
for (std::size_t i = 0; i < hfa->count; ++i) for (std::size_t i = 0; i < hfa->count; ++i)
builder.ldr_fp(fp_reg++, fp_mode, 31, (base_offset + i * size) / size); builder.ldr_fp(fp_reg++, fp_mode, 31, (base_offset + i * size) / size);
} }
else if (node->layout.size <= 16) else if (size <= 16)
{ {
// Small struct - passed in up to 2 GP registers // Small struct - passed in up to 2 GP registers
std::int32_t base_offset = stack_size - stack_position.at(argument); std::int32_t base_offset = stack_size - stack_position.at(argument);
std::int32_t size = node->layout.size;
std::int32_t offset = 0; std::int32_t offset = 0;
while (size > 0) while (offset < size)
{ {
builder.ldr(reg++, 31, (base_offset + offset) / 8); builder.ldr(reg++, 31, (base_offset + offset) / 8);
size -= 8;
offset += 8; offset += 8;
} }
} }
@ -806,12 +807,16 @@ namespace pslang::jit::aarch64
// TODO: array return value? // TODO: array return value?
auto size = ast::type_size(*type); auto size = ast::type_size(*type);
auto struct_type = std::get_if<types::struct_type>(type.get());
auto array_type = std::get_if<types::array_type>(type.get());
if (size == 0) if (size == 0)
{} {}
else if (auto struct_type = std::get_if<types::struct_type>(type.get())) else if (struct_type || array_type)
{ {
// NB: fixed-size arrays are handled in the same way
// as structs of N identical fields
auto base_offset = stack_size - stack_position.at(it); auto base_offset = stack_size - stack_position.at(it);
if (auto hfa = get_hfa_data(lcontext, struct_type->node); hfa && hfa->count <= 4) if (auto hfa = get_hfa_data(lcontext, type); hfa && hfa->count <= 4)
{ {
auto fp_mode = fp_mode_for(*hfa->element_type); auto fp_mode = fp_mode_for(*hfa->element_type);
auto size = fp_size(fp_mode); auto size = fp_size(fp_mode);
@ -863,13 +868,14 @@ namespace pslang::jit::aarch64
{ {
auto type = (*node.value)->inferred_type; auto type = (*node.value)->inferred_type;
auto size = ast::type_size(*type); auto size = ast::type_size(*type);
auto struct_type = std::get_if<types::struct_type>(type.get());
auto array_type = std::get_if<types::array_type>(type.get());
if (size == 0) if (size == 0)
{} {}
else if (auto struct_type = std::get_if<types::struct_type>(type.get())) else if (struct_type || array_type)
{ {
auto base_offset = stack_size - stack_position.at(*node.value); auto base_offset = stack_size - stack_position.at(*node.value);
auto struct_node = struct_type->node; if (auto hfa = get_hfa_data(lcontext, type); hfa && hfa->count <= 4)
if (auto hfa = get_hfa_data(lcontext, struct_node); hfa && hfa->count <= 4)
{ {
auto fp_mode = fp_mode_for(*hfa->element_type); auto fp_mode = fp_mode_for(*hfa->element_type);
auto size = fp_size(fp_mode); auto size = fp_size(fp_mode);
@ -909,8 +915,11 @@ namespace pslang::jit::aarch64
void compile(ast::function_definition const * function_definition, ir::node_ref begin, ir::node_ref end) void compile(ast::function_definition const * function_definition, ir::node_ref begin, ir::node_ref end)
{ {
if (auto struct_type = std::get_if<types::struct_type>(function_definition->inferred_result_type.get())) auto result_type = function_definition->inferred_result_type;
if (!get_hfa_data(lcontext, struct_type->node) && struct_type->node->layout.size > 16) auto struct_type = std::get_if<types::struct_type>(result_type.get());
auto array_type = std::get_if<types::array_type>(result_type.get());
if (struct_type || array_type)
if (!get_hfa_data(lcontext, result_type) && ast::type_size(*result_type) > 16)
return_value_is_large_struct = true; return_value_is_large_struct = true;
stack_size = 0; stack_size = 0;
@ -967,11 +976,12 @@ namespace pslang::jit::aarch64
{ {
auto const & argument = function_definition->arguments[i]; auto const & argument = function_definition->arguments[i];
auto size = ast::type_size(*argument.inferred_type); auto size = ast::type_size(*argument.inferred_type);
auto struct_type = std::get_if<types::struct_type>(argument.inferred_type.get());
auto array_type = std::get_if<types::array_type>(argument.inferred_type.get());
if (size == 0) continue; if (size == 0) continue;
if (auto struct_type = std::get_if<types::struct_type>(argument.inferred_type.get())) if (struct_type || array_type)
{ {
auto node = struct_type->node; if (auto hfa = get_hfa_data(lcontext, argument.inferred_type); hfa && hfa->count <= 4)
if (auto hfa = get_hfa_data(lcontext, node); hfa && hfa->count <= 4)
{ {
// HFA - passed in consecutive FP registers // HFA - passed in consecutive FP registers
std::int32_t base_offset = stack_size - argument_position[i]; std::int32_t base_offset = stack_size - argument_position[i];
@ -980,17 +990,15 @@ namespace pslang::jit::aarch64
for (std::size_t i = 0; i < hfa->count; ++i) for (std::size_t i = 0; i < hfa->count; ++i)
builder.str_fp(fp_reg++, fp_mode, 31, (base_offset + i * size) / size); builder.str_fp(fp_reg++, fp_mode, 31, (base_offset + i * size) / size);
} }
else if (node->layout.size <= 16) else if (size <= 16)
{ {
// Small struct - passed in up to 2 GP registers // Small struct - passed in up to 2 GP registers
std::int32_t base_offset = stack_size - argument_position[i]; std::int32_t base_offset = stack_size - argument_position[i];
std::int32_t size = node->layout.size;
std::int32_t offset = 0; std::int32_t offset = 0;
while (size > 0) while (offset < size)
{ {
builder.str(reg++, 31, (base_offset + offset) / 8); builder.str(reg++, 31, (base_offset + offset) / 8);
offset += 8; offset += 8;
size -= 8;
} }
} }
else else