const fs = require('fs');
const path = require('path');
const os = require('os');
const ref = require('ref-napi');
const StructDi = require('ref-struct-di');
const ffi = require('ffi-napi');
const Struct = StructDi(ref);
const UInt8 = ref.types.uint8;
const UInt32 = ref.types.uint32;
const UInt64 = ref.types.uint64;
const Float = ref.types.float;
const Int32 = ref.types.int32;
const SizeT = ref.types.size_t;
const RnnFfiBenchmarkRecord = Struct({
model_name_ptr: ref.refType(UInt8),
model_name_len: SizeT,
precision_ptr: ref.refType(UInt8),
precision_len: SizeT,
elapsed_ms: UInt64,
iterations: UInt64,
avg_loss: Float,
last_loss: Float,
output_bytes: UInt64,
train_samples: UInt64,
total_params: UInt64,
layer_count: UInt32,
input_dim: UInt32,
output_dim: UInt32,
benchmark_flags: UInt64,
weights_bytes: UInt64,
biases_bytes: UInt64,
min_loss: Float,
max_loss: Float,
loss_stddev: Float,
iterations_per_sec: Float,
samples_per_sec: Float,
});
const RnnFfiBenchmarkView = Struct({
model_name_ptr: ref.refType(UInt8),
model_name_len: SizeT,
precision_ptr: ref.refType(UInt8),
precision_len: SizeT,
elapsed_ms: UInt64,
iterations: UInt64,
avg_loss: Float,
last_loss: Float,
output_bytes: UInt64,
train_samples: UInt64,
total_params: UInt64,
layer_count: UInt32,
input_dim: UInt32,
output_dim: UInt32,
benchmark_flags: UInt64,
weights_bytes: UInt64,
biases_bytes: UInt64,
min_loss: Float,
max_loss: Float,
loss_stddev: Float,
iterations_per_sec: Float,
samples_per_sec: Float,
});
const Codes = { OK: 0 };
function candidateLibNames() {
const p = os.platform();
if (p === 'win32') return ['rnn.dll'];
if (p === 'darwin') return ['librnn.dylib', 'rnn.dylib'];
return ['librnn.so', 'rnn.so'];
}
function discoverLibPath(explicitPath) {
if (explicitPath) return explicitPath;
if (process.env.RNN_FFI_LIB) return process.env.RNN_FFI_LIB;
const root = path.resolve(__dirname, '../../..');
const names = candidateLibNames();
for (const folder of [path.join(root, 'target', 'debug'), path.join(root, 'target', 'release')]) {
for (const name of names) {
const candidate = path.join(folder, name);
if (fs.existsSync(candidate)) return candidate;
}
}
return names[0];
}
class RnnFfiError extends Error {
constructor(code, message) {
super(`RNN FFI error ${code}: ${message}`);
this.code = code;
}
}
function readUtf8(ptr, len) {
if (!ptr || ref.isNull(ptr) || Number(len) === 0) return '';
return ref.reinterpret(ptr, Number(len), 0).toString('utf8');
}
class RnnFfi {
constructor(options = {}) {
this._lib = ffi.Library(discoverLibPath(options.libraryPath), {
rnn_ffi_api_version: [UInt32, []],
rnn_ffi_benchmark_encoded_size: [Int32, [ref.refType(RnnFfiBenchmarkRecord), ref.refType(SizeT)]],
rnn_ffi_encode_benchmark_blob: [Int32, [ref.refType(RnnFfiBenchmarkRecord), ref.refType(UInt8), SizeT, ref.refType(SizeT)]],
rnn_ffi_decode_benchmark_blob: [Int32, [ref.refType(UInt8), SizeT, ref.refType(RnnFfiBenchmarkView)]],
rnn_ffi_error_message: ['string', [Int32]],
});
}
_throwIfError(code) {
if (code !== Codes.OK) {
throw new RnnFfiError(code, this._lib.rnn_ffi_error_message(code));
}
}
apiVersion() {
return this._lib.rnn_ffi_api_version();
}
benchmarkEncodedSize(record) {
const blob = this.encodeBenchmarkBlob(record);
return blob.length;
}
encodeBenchmarkBlob(record) {
const modelName = Buffer.from(record.model_name || '', 'utf8');
const precision = Buffer.from(record.precision || '', 'utf8');
const ffiRecord = new RnnFfiBenchmarkRecord();
ffiRecord.model_name_ptr = modelName;
ffiRecord.model_name_len = modelName.length;
ffiRecord.precision_ptr = precision;
ffiRecord.precision_len = precision.length;
ffiRecord.elapsed_ms = BigInt(record.elapsed_ms || 0);
ffiRecord.iterations = BigInt(record.iterations || 0);
ffiRecord.avg_loss = Number(record.avg_loss || 0);
ffiRecord.last_loss = Number(record.last_loss || 0);
ffiRecord.output_bytes = BigInt(record.output_bytes || 0);
ffiRecord.train_samples = BigInt(record.train_samples || 0);
ffiRecord.total_params = BigInt(record.total_params || 0);
ffiRecord.layer_count = Number(record.layer_count || 0);
ffiRecord.input_dim = Number(record.input_dim || 0);
ffiRecord.output_dim = Number(record.output_dim || 0);
ffiRecord.benchmark_flags = BigInt(record.benchmark_flags || 0);
ffiRecord.weights_bytes = BigInt(record.weights_bytes || 0);
ffiRecord.biases_bytes = BigInt(record.biases_bytes || 0);
ffiRecord.min_loss = Number(record.min_loss || 0);
ffiRecord.max_loss = Number(record.max_loss || 0);
ffiRecord.loss_stddev = Number(record.loss_stddev || 0);
ffiRecord.iterations_per_sec = Number(record.iterations_per_sec || 0);
ffiRecord.samples_per_sec = Number(record.samples_per_sec || 0);
const outLen = ref.alloc(SizeT);
let code = this._lib.rnn_ffi_benchmark_encoded_size(ffiRecord.ref(), outLen);
this._throwIfError(code);
const size = Number(outLen.deref());
const out = Buffer.alloc(size);
const outUsed = ref.alloc(SizeT);
code = this._lib.rnn_ffi_encode_benchmark_blob(ffiRecord.ref(), out, out.length, outUsed);
this._throwIfError(code);
return out.subarray(0, Number(outUsed.deref()));
}
decodeBenchmarkBlob(blob) {
const bytes = Buffer.from(blob);
const out = new RnnFfiBenchmarkView();
const code = this._lib.rnn_ffi_decode_benchmark_blob(bytes, bytes.length, out.ref());
this._throwIfError(code);
return {
model_name: readUtf8(out.model_name_ptr, out.model_name_len),
precision: readUtf8(out.precision_ptr, out.precision_len),
elapsed_ms: Number(out.elapsed_ms),
iterations: Number(out.iterations),
avg_loss: Number(out.avg_loss),
last_loss: Number(out.last_loss),
output_bytes: Number(out.output_bytes),
train_samples: Number(out.train_samples),
total_params: Number(out.total_params),
layer_count: Number(out.layer_count),
input_dim: Number(out.input_dim),
output_dim: Number(out.output_dim),
benchmark_flags: Number(out.benchmark_flags),
weights_bytes: Number(out.weights_bytes),
biases_bytes: Number(out.biases_bytes),
min_loss: Number(out.min_loss),
max_loss: Number(out.max_loss),
loss_stddev: Number(out.loss_stddev),
iterations_per_sec: Number(out.iterations_per_sec),
samples_per_sec: Number(out.samples_per_sec),
};
}
}
module.exports = { RnnFfi, RnnFfiError };