Skip to content

Commit

Permalink
Move all NIF functions out of main.zig (#305)
Browse files Browse the repository at this point in the history
  • Loading branch information
@jackalcooper
jackalcooper committed Apr 25, 2024
1 parent 170dece commit 1effcb9
Show file tree
Hide file tree
Showing 8 changed files with 413 additions and 394 deletions.
78 changes: 39 additions & 39 deletions native/mlir-zig-proj/src/enif_support.zig
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,13 @@ const beam = @import("beam");
const mlir_capi = @import("mlir_capi.zig");
pub const c = @import("prelude.zig");
const e = @import("runtime.zig");
const result = @import("result.zig");

const Invocation = struct {
arg_terms: []beam.term = undefined,
res_term: beam.term = undefined,
packed_args: []?*anyopaque = undefined, // [arg0, arg1... result]
pub fn init(self: *@This(), environment: beam.env, list: beam.term) !void {
fn init(self: *@This(), environment: beam.env, list: beam.term) !void {
const size = try beam.get_list_length(environment, list);
var head: beam.term = undefined;
self.arg_terms = try beam.allocator.alloc(beam.term, size);
Expand All @@ -21,26 +23,25 @@ const Invocation = struct {
errdefer beam.allocator.free(self.arg_terms);
errdefer beam.allocator.free(self.packed_args);
}
pub fn deinit(self: *@This()) void {
fn deinit(self: *@This()) void {
beam.allocator.free(self.arg_terms);
beam.allocator.free(self.packed_args);
}
pub fn invoke(self: *@This(), environment: beam.env, jit: mlir_capi.ExecutionEngine.T, name: beam.binary) callconv(.C) mlir_capi.LogicalResult.T {
fn invoke(self: *@This(), environment: beam.env, jit: mlir_capi.ExecutionEngine.T, name: beam.binary) callconv(.C) mlir_capi.LogicalResult.T {
self.packed_args[0] = @ptrCast(@constCast(&environment));
return c.mlirExecutionEngineInvokePacked(jit, c.mlirStringRefCreate(name.data, name.size), &self.packed_args[0]);
}
};

pub fn beaver_raw_jit_invoke_with_terms(env: beam.env, _: c_int, args: [*c]const beam.term) callconv(.C) beam.term {
var jit: mlir_capi.ExecutionEngine.T = mlir_capi.ExecutionEngine.resource.fetch(env, args[0]) catch
return beam.make_error_binary(env, "fail to fetch resource for ExecutionEngine, expected: " ++ @typeName(mlir_capi.ExecutionEngine.T));
var name: beam.binary = beam.get_binary(env, args[1]) catch
return beam.make_error_binary(env, "fail to get binary for jit func name");
fn beaver_raw_jit_invoke_with_terms(env: beam.env, _: c_int, args: [*c]const beam.term) !beam.term {
const Error = error{JITFunctionCallFailure};
var jit: mlir_capi.ExecutionEngine.T = try mlir_capi.ExecutionEngine.resource.fetch(env, args[0]);
var name: beam.binary = try beam.get_binary(env, args[1]);
var invocation = Invocation{};
invocation.init(env, args[2]) catch return beam.make_error_binary(env, "fail to init jit invocation");
try invocation.init(env, args[2]);
defer invocation.deinit();
if (c.beaverLogicalResultIsFailure(invocation.invoke(env, jit, name))) {
return beam.make_error_binary(env, "fail to call jit function");
return Error.JITFunctionCallFailure;
}
return invocation.res_term;
}
Expand All @@ -56,9 +57,8 @@ fn register_jit_symbol(jit: mlir_capi.ExecutionEngine.T, comptime name: []const
c.mlirExecutionEngineRegisterSymbol(jit, name_str_ref, @ptrCast(@constCast(&f)));
}

pub fn beaver_raw_jit_register_enif(env: beam.env, _: c_int, args: [*c]const beam.term) callconv(.C) beam.term {
var jit: mlir_capi.ExecutionEngine.T = mlir_capi.ExecutionEngine.resource.fetch(env, args[0]) catch
return beam.make_error_binary(env, "fail to fetch resource for ExecutionEngine, expected: " ++ @typeName(mlir_capi.ExecutionEngine.T));
fn beaver_raw_jit_register_enif(env: beam.env, _: c_int, args: [*c]const beam.term) !beam.term {
var jit = try mlir_capi.ExecutionEngine.resource.fetch(env, args[0]);
inline for (enif_function_names) |name| {
register_jit_symbol(jit, name, @field(e, name));
}
Expand Down Expand Up @@ -122,26 +122,21 @@ fn dump_type_info(env: beam.env, ctx: mlir_capi.Context.T, comptime t: type) !be
return beam.make_tuple(env, type_info_slice);
}

pub fn beaver_raw_enif_signatures(env: beam.env, _: c_int, args: [*c]const beam.term) callconv(.C) beam.term {
const ctx = mlir_capi.Context.resource.fetch(env, args[0]) catch
return beam.make_error_binary(env, "fail to fetch resource for argument #0, expected: " ++ @typeName(mlir_capi.Context.T));
var signatures: []beam.term = beam.allocator.alloc(beam.term, enif_function_names.len) catch
return beam.make_error_binary(env, "fail to allocate");
fn beaver_raw_enif_signatures(env: beam.env, _: c_int, args: [*c]const beam.term) !beam.term {
const ctx = try mlir_capi.Context.resource.fetch(env, args[0]);
var signatures: []beam.term = try beam.allocator.alloc(beam.term, enif_function_names.len);
inline for (enif_function_names, 0..) |name, i| {
const f = @field(e, name);
const FTI = @typeInfo(@TypeOf(f)).Fn;
var signature_slice: []beam.term = beam.allocator.alloc(beam.term, 3) catch
return beam.make_error_binary(env, "fail to allocate");
var signature_slice: []beam.term = try beam.allocator.alloc(beam.term, 3);
defer beam.allocator.free(signature_slice);
var arg_type_slice: []beam.term = beam.allocator.alloc(beam.term, FTI.params.len) catch
return beam.make_error_binary(env, "fail to allocate");
var arg_type_slice: []beam.term = try beam.allocator.alloc(beam.term, FTI.params.len);
defer beam.allocator.free(arg_type_slice);
inline for (FTI.params, 0..) |p, arg_i| {
if (p.type) |t| {
arg_type_slice[arg_i] = dump_type_info(env, ctx, t) catch
return beam.make_error_binary(env, "fail to allocate");
arg_type_slice[arg_i] = try dump_type_info(env, ctx, t);
} else if (@TypeOf(f) == @TypeOf(e.enif_compare_pids)) {
arg_type_slice[arg_i] = dump_type_info(env, ctx, [*c]u8) catch return beam.make_error_binary(env, "fail to dump type");
arg_type_slice[arg_i] = try dump_type_info(env, ctx, [*c]u8);
} else {
@compileError("param type not found, function: " ++ name);
}
Expand All @@ -155,15 +150,14 @@ pub fn beaver_raw_enif_signatures(env: beam.env, _: c_int, args: [*c]const beam.
ret_size = 0;
}
}
var ret_slice: []beam.term = beam.allocator.alloc(beam.term, ret_size) catch
return beam.make_error_binary(env, "fail to allocate");
var ret_slice: []beam.term = try beam.allocator.alloc(beam.term, ret_size);
defer beam.allocator.free(ret_slice);
if (FTI.return_type) |t| {
if (t != void) {
ret_slice[0] = dump_type_info(env, ctx, t) catch return beam.make_error_binary(env, "fail to dump type");
ret_slice[0] = try dump_type_info(env, ctx, t);
}
} else if (f == e.enif_compare_pids) {
ret_slice[0] = dump_type_info(env, ctx, c_int) catch return beam.make_error_binary(env, "fail to dump type");
ret_slice[0] = try dump_type_info(env, ctx, c_int);
} else {
@compileError("return type not found, function: " ++ name);
}
Expand All @@ -173,25 +167,31 @@ pub fn beaver_raw_enif_signatures(env: beam.env, _: c_int, args: [*c]const beam.
return beam.make_term_list(env, signatures);
}

pub fn beaver_raw_enif_functions(env: beam.env, _: c_int, _: [*c]const beam.term) callconv(.C) beam.term {
var names: []beam.term = beam.allocator.alloc(beam.term, enif_function_names.len) catch
return beam.make_error_binary(env, "fail to allocate");
fn beaver_raw_enif_functions(env: beam.env, _: c_int, _: [*c]const beam.term) !beam.term {
var names: []beam.term = try beam.allocator.alloc(beam.term, enif_function_names.len);
inline for (enif_function_names, 0..) |name, i| {
names[i] = beam.make_atom(env, name);
}
return beam.make_term_list(env, names);
}

pub fn beaver_raw_mlir_type_of_enif_obj(env: beam.env, _: c_int, args: [*c]const beam.term) callconv(.C) beam.term {
const ctx = mlir_capi.Context.resource.fetch(env, args[0]) catch
return beam.make_error_binary(env, "fail to fetch resource for argument #0, expected: " ++ @typeName(mlir_capi.Context.T));
var name = beam.get_atom_slice(env, args[1]) catch
return beam.make_error_binary(env, "fail to get name");
fn beaver_raw_mlir_type_of_enif_obj(env: beam.env, _: c_int, args: [*c]const beam.term) !beam.term {
const Error = error{MLIRTypeForEnifObjNotFound};
const ctx = try mlir_capi.Context.resource.fetch(env, args[0]);
var name = try beam.get_atom_slice(env, args[1]);
inline for (.{ "term", "env" }) |obj| {
if (@import("std").mem.eql(u8, name, obj)) {
const t = @field(beam, obj);
return enif_mlir_type(env, ctx, t) catch return beam.make_error_binary(env, "fail to get mlir type");
return try enif_mlir_type(env, ctx, t);
}
}
return beam.make_error_binary(env, "mlir type not found for enif obj");
return Error.MLIRTypeForEnifObjNotFound;
}

pub const nifs = .{
result.nif("beaver_raw_jit_invoke_with_terms", 3, beaver_raw_jit_invoke_with_terms).entry,
result.nif("beaver_raw_jit_register_enif", 1, beaver_raw_jit_register_enif).entry,
result.nif("beaver_raw_enif_signatures", 1, beaver_raw_enif_signatures).entry,
result.nif("beaver_raw_enif_functions", 0, beaver_raw_enif_functions).entry,
result.nif("beaver_raw_mlir_type_of_enif_obj", 2, beaver_raw_mlir_type_of_enif_obj).entry,
};
Loading

0 comments on commit 1effcb9

Please sign in to comment.