executorch.js 36 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989
  1. // Experimental
  2. const executorch = {};
  3. const coreml = {};
  4. const vulkan = {};
  5. const xnnpack = {};
  6. import * as base from './base.js';
  7. import * as python from './python.js';
  8. import * as pytorch from './pytorch.js';
  9. executorch.ModelFactory = class {
  10. match(context) {
  11. const reader = executorch.Reader.open(context);
  12. if (reader) {
  13. context.type = 'executorch';
  14. context.target = reader;
  15. }
  16. }
  17. async open(context) {
  18. executorch.schema = await context.require('./executorch-schema');
  19. const target = context.target;
  20. await target.read();
  21. return new executorch.Model(target);
  22. }
  23. };
  24. executorch.Model = class {
  25. constructor(target) {
  26. this.format = `ExecuTorch v${target.program.version}`;
  27. this.graphs = [];
  28. for (const plan of target.program.execution_plan) {
  29. for (const chain of plan.chains) {
  30. const graph = new executorch.Graph(target, plan, chain);
  31. this.graphs.push(graph);
  32. }
  33. }
  34. }
  35. };
  36. executorch.Graph = class {
  37. constructor(target, plan, chain) {
  38. this.inputs = [];
  39. this.outputs = [];
  40. this.nodes = [];
  41. const values = new Map();
  42. values.map = (index, output) => {
  43. if (!values.has(index)) {
  44. const executorch_flatbuffer = executorch.schema.executorch_flatbuffer;
  45. const val = plan.values[index].val;
  46. const tensor = val instanceof executorch_flatbuffer.Tensor || val instanceof executorch_flatbuffer.TensorList || val instanceof executorch_flatbuffer.OptionalTensorList;
  47. if (output && !tensor) {
  48. const value = [new executorch.Value(index.toString(), null, null)];
  49. values.set(index, { type: null, value });
  50. } else if (tensor) {
  51. const tensors = val instanceof executorch_flatbuffer.Tensor ? [val] : Array.from(val.items).map((arg) => plan.values[arg].val);
  52. const list = [];
  53. for (let i = 0; i < tensors.length; i++) {
  54. const tensor = tensors[i];
  55. const type = new executorch.TensorType(tensor);
  56. let initializer = null;
  57. if (val.data_buffer_idx > 0) {
  58. initializer = new executorch.Tensor(tensor, target);
  59. }
  60. const identifier = tensors.length > 1 ? `${index}.${i}` : index.toString();
  61. const value = new executorch.Value(identifier, type, initializer);
  62. list.push(value);
  63. }
  64. values.set(index, { type: null, value: list });
  65. } else if (val instanceof executorch_flatbuffer.Bool) {
  66. values.set(index, { type: 'int64', value: val.bool_val });
  67. } else if (val instanceof executorch_flatbuffer.Int) {
  68. values.set(index, { type: 'int64', value: val.int_val });
  69. } else if (val instanceof executorch_flatbuffer.IntList) {
  70. const list = val.items.map((index) => plan.values[index].val.int_val);
  71. values.set(index, { type: 'int64[]', value: list });
  72. } else if (val instanceof executorch_flatbuffer.Double) {
  73. values.set(index, { type: 'float64', value: val.double_val });
  74. } else if (val instanceof executorch_flatbuffer.String) {
  75. values.set(index, { type: 'string', value: val.string_val });
  76. } else if (val instanceof executorch_flatbuffer.Null) {
  77. values.set(index, { type: 'attribute', value: null });
  78. } else {
  79. throw new Error(`Value type '${val.constructor.name}' not implemented.`);
  80. }
  81. }
  82. return values.get(index);
  83. };
  84. for (let i = 0; i < plan.inputs.length; i++) {
  85. const input = plan.inputs[i];
  86. const value = values.map(input);
  87. const name = plan.inputs.length === 1 ? 'input' : `input.${i}`;
  88. const argument = new executorch.Argument(name, value.value, value.type);
  89. this.inputs.push(argument);
  90. }
  91. for (let i = 0; i < plan.outputs.length; i++) {
  92. const output = plan.outputs[i];
  93. const value = values.map(output);
  94. const name = plan.outputs.length === 1 ? 'output' : `output.${i}`;
  95. const argument = new executorch.Argument(name, value.value, value.type);
  96. this.outputs.push(argument);
  97. }
  98. for (const instruction of chain.instructions) {
  99. const node = new executorch.Node(target, plan, chain, instruction, values);
  100. this.nodes.push(node);
  101. }
  102. }
  103. };
  104. executorch.Argument = class {
  105. constructor(name, value, type, visible) {
  106. this.name = name;
  107. this.value = value;
  108. this.type = type || null;
  109. this.visible = visible !== false;
  110. }
  111. };
  112. executorch.Value = class Value {
  113. constructor(name, type, initializer) {
  114. if (typeof name !== 'string') {
  115. throw new executorch.Error(`Invalid value identifier '${JSON.stringify(name)}'.`);
  116. }
  117. this.name = name;
  118. this.type = initializer && initializer.type ? initializer.type : type || null;
  119. this.initializer = initializer || null;
  120. }
  121. };
  122. executorch.Node = class {
  123. constructor(target, plan, chain, instruction, values) {
  124. this.name = '';
  125. this.inputs = [];
  126. this.outputs = [];
  127. this.attributes = [];
  128. const instr_args = instruction.instr_args;
  129. const executorch_flatbuffer = executorch.schema.executorch_flatbuffer;
  130. if (instr_args instanceof executorch_flatbuffer.KernelCall) {
  131. const op = plan.operators[instr_args.op_index];
  132. const name = op.name.split('::').pop();
  133. const identifier = op.overload ? `${op.name}.${op.overload}` : op.name;
  134. const schemas = target.execution.invoke('torch._C._jit_get_schemas_for_operator', [op.name]);
  135. const schema = schemas.find((schema) => schema.name === op.name && schema.overload_name === op.overload);
  136. if (!schema) {
  137. throw new executorch.Error(`Operator schema for '${identifier}' not found.`);
  138. }
  139. const category = schema && schema.category ? schema.category : '';
  140. const alias = (arg) => arg && arg.alias_info && arg.alias_info.before_set.length === 1 ? arg.alias_info.before_set[0] : null;
  141. const outputs = new Set(schema && Array.isArray(schema.returns) ? schema.returns.map((arg) => alias(arg)).filter((alias) => alias !== null) : []);
  142. const inputs = new Map();
  143. this.type = { name, identifier, category };
  144. let i = 0;
  145. const args = instr_args.args;
  146. for (; i < schema.arguments.length; i++) {
  147. const index = args[i];
  148. const arg = schema && i < schema.arguments.length ? schema.arguments[i] : null;
  149. const output = arg ? alias(schema.arguments[i]) : null;
  150. if (output && outputs.has(output)) {
  151. inputs.set(output, index);
  152. continue;
  153. }
  154. const name = arg ? arg.name : i.toString();
  155. const value = values.map(index);
  156. const argument = new executorch.Argument(name, value.value, value.type);
  157. this.inputs.push(argument);
  158. }
  159. for (let j = 0; j < schema.returns.length; j++) {
  160. const ret = schema.returns[j];
  161. const output = alias(ret);
  162. let index = args[i++];
  163. index = output && inputs.has(output) ? inputs.get(output) : index;
  164. const name = ret.name;
  165. const value = values.map(index, true);
  166. const argument = new executorch.Argument(name || '', value.value, value.type);
  167. this.outputs.push(argument);
  168. }
  169. } else if (instr_args instanceof executorch_flatbuffer.DelegateCall) {
  170. const delegate = plan.delegates[instr_args.delegate_index];
  171. const args = instr_args.args;
  172. if (!delegate.backend || !delegate.backend.type) {
  173. throw new executorch.Error(`ExecuTorch delegate '${delegate.id}' not implemented.`);
  174. }
  175. this.type = delegate.backend.type;
  176. const inputs = args.slice(0, this.type.inputs.length);
  177. for (let i = 0; i < inputs.length; i++) {
  178. const input = inputs[i];
  179. const value = values.map(input);
  180. const name = inputs.length === 1 ? 'input' : `input.${i}`;
  181. const argument = new executorch.Argument(name, value.value, value.type);
  182. this.inputs.push(argument);
  183. }
  184. const outputs = args.slice(this.type.inputs.length, this.type.inputs.length + this.type.outputs.length);
  185. for (let i = 0; i < outputs.length; i++) {
  186. const output = outputs[i];
  187. const value = values.map(output);
  188. const name = inputs.length === 1 ? 'output' : `output.${i}`;
  189. const argument = new executorch.Argument(name, value.value, value.type);
  190. this.outputs.push(argument);
  191. }
  192. for (const spec of delegate.compile_specs) {
  193. const value = spec.value instanceof Uint8Array ? new TextDecoder('utf-8').decode(spec.value) : spec.value;
  194. const attribute = new executorch.Argument(spec.key, value, 'attribute');
  195. this.attributes.push(attribute);
  196. }
  197. } else {
  198. throw new Error(`Instruction type '${instr_args.constructor.name}' not implemented.`);
  199. }
  200. }
  201. };
  202. executorch.TensorType = class {
  203. constructor(tensor) {
  204. executorch.TensorType._types = executorch.TensorType._types || [
  205. 'uint8',
  206. 'int8', 'int16', 'int32', 'int64',
  207. 'float16', 'float32', 'float64',
  208. 'complex16', 'complex32', 'complex64',
  209. 'boolean',
  210. 'qint8', 'quint8', 'qint32',
  211. 'bfloat16',
  212. 'quint4x2', 'quint2x4', 'bits1x8', 'bits2x4', 'bits4x2', 'bits8', 'bits16',
  213. 'float8e5m2', 'float8e4m3fn', 'float8e5m2fnuz', 'float8e4m3fnuz',
  214. 'uint16', 'uint32', 'uint64'
  215. ];
  216. if (tensor.scalar_type >= executorch.TensorType._types.length) {
  217. throw new executorch.Error(`Unknown tensor data type '${tensor.scalar_type}'.`);
  218. }
  219. this.dataType = executorch.TensorType._types[tensor.scalar_type];
  220. this.shape = new executorch.TensorShape(Array.from(tensor.sizes));
  221. }
  222. toString() {
  223. return this.dataType + this.shape.toString();
  224. }
  225. };
  226. executorch.TensorShape = class {
  227. constructor(dimensions) {
  228. this.dimensions = dimensions || [];
  229. }
  230. toString() {
  231. if (this.dimensions && this.dimensions.length > 0) {
  232. return `[${this.dimensions.map((dimension) => dimension.toString()).join(',')}]`;
  233. }
  234. return '';
  235. }
  236. };
  237. executorch.Tensor = class {
  238. constructor(tensor, target) {
  239. this.type = new executorch.TensorType(tensor);
  240. const data_buffer_idx = tensor.data_buffer_idx;
  241. const program = target.program;
  242. if (tensor.extra_tensor_info) {
  243. throw new executorch.Error('Extra tensor info not implemented.');
  244. } else if (program.constant_buffers) {
  245. throw new executorch.Error('Constant buffers not implemented.');
  246. } else if (tensor.allocation_info === null) {
  247. const constant_segment = program.constant_segment;
  248. const data_segment = program.segments[constant_segment.segment_index];
  249. const offset = constant_segment.offsets[data_buffer_idx].toNumber();
  250. const next = data_buffer_idx + 1 < constant_segment.offsets.length ? constant_segment.offsets[data_buffer_idx + 1].toNumber() : data_segment.size.toNumber();
  251. const size = next - offset;
  252. this.values = target.blob(data_segment.offset.toNumber() + offset, size);
  253. this.encoding = '<';
  254. } else {
  255. throw new executorch.Error('Tensor allocation info not implemented.');
  256. }
  257. }
  258. };
  259. executorch.Reader = class {
  260. static open(context) {
  261. const reader = context.peek('flatbuffers.binary');
  262. if (reader && reader.identifier === 'ET12') {
  263. return new executorch.Reader(context, reader);
  264. }
  265. return null;
  266. }
  267. constructor(context, reader) {
  268. this.context = context;
  269. this.reader = reader;
  270. }
  271. async read() {
  272. this.metadata = await pytorch.Metadata.open(this.context);
  273. this.execution = new python.Execution();
  274. this.metadata.register(this.execution);
  275. const executorch_flatbuffer = executorch.schema.executorch_flatbuffer;
  276. this.program = executorch_flatbuffer.Program.create(this.reader);
  277. this.reader = this.context.read('binary');
  278. if (this.reader.length >= 32) {
  279. this.reader.seek(8);
  280. const magic = String.fromCharCode(...this.reader.read(4));
  281. if (magic === 'eh00') {
  282. this.extended_file_header = {
  283. length: this.reader.uint32(),
  284. program_size: this.reader.uint64().toNumber(),
  285. segment_base_offset: this.reader.uint64().toNumber(),
  286. };
  287. }
  288. this.reader.seek(0);
  289. }
  290. for (const plan of this.program.execution_plan) {
  291. for (const chain of plan.chains) {
  292. for (const instruction of chain.instructions) {
  293. const instr_args = instruction.instr_args;
  294. if (instr_args instanceof executorch_flatbuffer.DelegateCall) {
  295. const delegate = plan.delegates[instr_args.delegate_index];
  296. if (delegate.backend) {
  297. continue;
  298. }
  299. let data = null;
  300. switch (delegate.processed.location) {
  301. case executorch_flatbuffer.DataLocation.INLINE: {
  302. data = this.program.backend_delegate_data[delegate.processed.index].data;
  303. break;
  304. }
  305. case executorch_flatbuffer.DataLocation.SEGMENT: {
  306. const segment = this.program.segments[delegate.processed.index];
  307. data = this.blob(segment.offset.toNumber(), segment.size.toNumber());
  308. break;
  309. }
  310. default: {
  311. throw new executorch.Error(`Delegate data location '${delegate.processed.location}' not implemented.`);
  312. }
  313. }
  314. switch (delegate.id) {
  315. case 'XnnpackBackend': {
  316. delegate.backend = xnnpack.Reader.open(data, this);
  317. break;
  318. }
  319. case 'CoreMLBackend': {
  320. delegate.backend = coreml.Reader.open(data, this);
  321. break;
  322. }
  323. case 'VulkanBackend': {
  324. delegate.backend = vulkan.Reader.open(data, this);
  325. break;
  326. }
  327. default: {
  328. throw new executorch.Error(`ExecuTorch delegate '${delegate.id}' not implemented.`);
  329. }
  330. }
  331. /* eslint-disable no-await-in-loop */
  332. await delegate.backend.read();
  333. /* eslint-enable no-await-in-loop */
  334. }
  335. }
  336. }
  337. }
  338. }
  339. blob(offset, size) {
  340. if (this.extended_file_header) {
  341. this.reader.seek(this.extended_file_header.segment_base_offset + offset);
  342. const data = this.reader.read(size);
  343. this.reader.seek(0);
  344. return data;
  345. }
  346. return null;
  347. }
  348. };
  349. executorch.Error = class extends Error {
  350. constructor(message) {
  351. super(message);
  352. this.name = 'Error loading ExecuTorch model.';
  353. }
  354. };
  355. xnnpack.Reader = class {
  356. static open(data, target) {
  357. if (data.length >= 30) {
  358. const reader = base.BinaryReader.open(data);
  359. reader.skip(4);
  360. const magic = String.fromCharCode(...reader.read(4));
  361. if (magic === 'XH00') {
  362. return new xnnpack.Reader(reader, target);
  363. }
  364. }
  365. return null;
  366. }
  367. constructor(reader, target) {
  368. this.reader = reader;
  369. this.target = target;
  370. reader.skip(2);
  371. this.flatbuffer = {
  372. offset: reader.uint32(),
  373. size: reader.uint32(),
  374. };
  375. this.constants = {
  376. offset: reader.uint32(),
  377. size: reader.uint32(),
  378. };
  379. }
  380. async read() {
  381. this.reader.seek(this.flatbuffer.offset);
  382. const flatbuffers = await import('./flatbuffers.js');
  383. const data = this.reader.read(this.flatbuffer.size);
  384. const reader = flatbuffers.BinaryReader.open(data);
  385. if (!executorch.schema.fb_xnnpack.XNNGraph.identifier(reader)) {
  386. throw new xnnpack.Error('Invalid XNNPACK data.');
  387. }
  388. this.graph = executorch.schema.fb_xnnpack.XNNGraph.create(reader);
  389. this.reader.seek(0);
  390. const metadata = new xnnpack.Metadata();
  391. this.type = new xnnpack.Graph(metadata, this.graph, this);
  392. }
  393. constant(idx) {
  394. const constant_data = this.graph.constant_data[idx];
  395. this.reader.seek(this.constants.offset + constant_data.offset.toNumber());
  396. const data = this.reader.read(constant_data.size.toNumber());
  397. this.reader.seek(0);
  398. return data;
  399. }
  400. };
  401. xnnpack.Graph = class {
  402. constructor(metadata, graph, reader) {
  403. this.name = 'XnnpackBackend';
  404. this.type = 'graph';
  405. this.inputs = [];
  406. this.outputs = [];
  407. this.nodes = [];
  408. const values = new Map();
  409. values.map = (id) => {
  410. if (!values.has(id)) {
  411. const fb_xnnpack = executorch.schema.fb_xnnpack;
  412. const name = id.toString();
  413. const xvalue = graph.xvalues[id].xvalue_union;
  414. if (xvalue instanceof fb_xnnpack.XNNTensorValue) {
  415. const type = new xnnpack.TensorType(xvalue);
  416. const initializer = xvalue.constant_buffer_idx === 0 ? null : new xnnpack.Tensor(xvalue, reader);
  417. const value = new xnnpack.Value(name, type, initializer);
  418. values.set(id, value);
  419. } else if (xvalue instanceof fb_xnnpack.XNNQuantizedTensorValue) {
  420. const value = new xnnpack.Value(name, null, null);
  421. values.set(id, value);
  422. } else {
  423. throw new xnnpack.Error(`Value type '${xvalue.constructor.name}' not implemented.`);
  424. }
  425. }
  426. return values.get(id);
  427. };
  428. for (let i = 0; i < graph.input_ids.length; i++) {
  429. const id = graph.input_ids[i];
  430. const value = values.map(id);
  431. const name = graph.input_ids.length === 1 ? 'input' : `input.${i}`;
  432. const argument = new xnnpack.Argument(name, [value]);
  433. this.inputs.push(argument);
  434. }
  435. for (let i = 0; i < graph.output_ids.length; i++) {
  436. const id = graph.output_ids[i];
  437. const value = values.map(id);
  438. const name = graph.output_ids.length === 1 ? 'output' : `output.${i}`;
  439. const argument = new xnnpack.Argument(name, [value]);
  440. this.outputs.push(argument);
  441. }
  442. for (const xnode of graph.xnodes) {
  443. const node = new xnnpack.Node(metadata, xnode, values);
  444. this.nodes.push(node);
  445. }
  446. }
  447. };
  448. xnnpack.Node = class {
  449. constructor(metadata, xnode, values) {
  450. const node = xnode.xnode_union;
  451. this.type = metadata.type(node.constructor.name) || { name: node.constructor.name };
  452. this.name = '';
  453. this.inputs = [];
  454. this.outputs = [];
  455. for (const [name, obj] of Object.entries(node)) {
  456. let value = ArrayBuffer.isView(obj) ? Array.from(obj) : obj;
  457. let type = 'attribute';
  458. if (name.endsWith('_id')) {
  459. value = obj === -1 ? [] : [values.map(obj)];
  460. type = null;
  461. }
  462. const argument = new xnnpack.Argument(name, value, type);
  463. if (name === 'output_id') {
  464. this.outputs.push(argument);
  465. } else {
  466. this.inputs.push(argument);
  467. }
  468. }
  469. }
  470. };
  471. xnnpack.Argument = class {
  472. constructor(name, value, type, visible) {
  473. this.name = name;
  474. this.value = value;
  475. this.type = type || null;
  476. this.visible = visible !== false;
  477. }
  478. };
  479. xnnpack.Value = class Value {
  480. constructor(name, type, initializer) {
  481. if (typeof name !== 'string') {
  482. throw new executorch.Error(`Invalid value identifier '${JSON.stringify(name)}'.`);
  483. }
  484. this.name = name;
  485. this.type = initializer && initializer.type ? initializer.type : type || null;
  486. this.initializer = initializer || null;
  487. }
  488. };
  489. xnnpack.Metadata = class {
  490. constructor() {
  491. this._types = new Map();
  492. this.register('XNNStaticTranspose', 'Transform');
  493. this.register('_XNNNodeConv', 'Layer');
  494. this.register('XNNFullyConnected', 'Layer');
  495. this.register('_XNNCat', 'Tensor');
  496. }
  497. register(name, category) {
  498. this._types.set(name, { name, category });
  499. }
  500. type(name) {
  501. return this._types.get(name);
  502. }
  503. };
  504. xnnpack.TensorType = class {
  505. constructor(tensor) {
  506. xnnpack.TensorType._types = executorch.TensorType._types || [
  507. 'invalid', 'float32', 'float16',
  508. 'qint8', 'quint8', 'qint32',
  509. 'qcint8', 'qcint32', 'qcint4',
  510. 'qdint8', 'qbint4'
  511. ];
  512. if (tensor.datatype >= xnnpack.TensorType._types.length) {
  513. throw new xnnpack.Error(`Unknown tensor data type '${tensor.datatype}'.`);
  514. }
  515. this.dataType = xnnpack.TensorType._types[tensor.datatype];
  516. this.shape = new xnnpack.TensorShape(Array.from(tensor.dims));
  517. }
  518. toString() {
  519. return this.dataType + this.shape.toString();
  520. }
  521. };
  522. xnnpack.TensorShape = class {
  523. constructor(dimensions) {
  524. this.dimensions = dimensions || [];
  525. }
  526. toString() {
  527. if (this.dimensions && this.dimensions.length > 0) {
  528. return `[${this.dimensions.map((dimension) => dimension.toString()).join(',')}]`;
  529. }
  530. return '';
  531. }
  532. };
  533. xnnpack.Tensor = class {
  534. constructor(tensor, reader) {
  535. this.type = new xnnpack.TensorType(tensor);
  536. this.values = reader.constant(tensor.constant_buffer_idx);
  537. this.encoding = '<';
  538. }
  539. };
  540. xnnpack.Error = class extends Error {
  541. constructor(message) {
  542. super(message);
  543. this.name = 'Error loading XNNPACK model.';
  544. }
  545. };
  546. vulkan.Reader = class {
  547. static open(data, target) {
  548. if (data.length >= 30) {
  549. const reader = base.BinaryReader.open(data);
  550. reader.skip(4);
  551. const magic = String.fromCharCode(...reader.read(4));
  552. if (magic === 'VH00') {
  553. return new vulkan.Reader(reader, target);
  554. }
  555. }
  556. return null;
  557. }
  558. constructor(reader, target) {
  559. this.reader = reader;
  560. this.target = target;
  561. reader.skip(2);
  562. this.flatbuffer = {
  563. offset: reader.uint32(),
  564. size: reader.uint32(),
  565. };
  566. this.constants = {
  567. offset: reader.uint32(),
  568. size: reader.uint32(),
  569. };
  570. }
  571. async read() {
  572. this.reader.seek(this.flatbuffer.offset);
  573. const metadata = new vulkan.Metadata(this.target.execution);
  574. metadata.register('conv_with_clamp(Tensor input, Tensor weight, Tensor? bias, SymInt[] stride, SymInt[] padding, SymInt[] dilation, bool transposed, SymInt[] output_padding, SymInt groups, Scalar? output_min, Scalar? output_max) -> Tensor)');
  575. const flatbuffers = await import('./flatbuffers.js');
  576. const data = this.reader.read(this.flatbuffer.size);
  577. const reader = flatbuffers.BinaryReader.open(data);
  578. if (!executorch.schema.vkgraph.VkGraph.identifier(reader)) {
  579. throw new xnnpack.Error('Invalid Vuklan data.');
  580. }
  581. this.graph = executorch.schema.vkgraph.VkGraph.create(reader);
  582. this.reader.seek(0);
  583. this.type = new vulkan.Graph(metadata, this.graph, this);
  584. }
  585. constant(id) {
  586. const constant = this.graph.constants[id];
  587. this.reader.seek(this.constants.offset + constant.offset.toNumber());
  588. const data = this.reader.read(constant.length.toNumber());
  589. this.reader.seek(0);
  590. return data;
  591. }
  592. };
  593. vulkan.Graph = class {
  594. constructor(metadata, graph, reader) {
  595. this.name = 'VulkanBackend';
  596. this.inputs = [];
  597. this.outputs = [];
  598. this.nodes = [];
  599. const values = new Map();
  600. values.map = (id) => {
  601. if (!values.has(id)) {
  602. const vkgraph = executorch.schema.vkgraph;
  603. const arg = graph.values[id].value;
  604. if (arg instanceof vkgraph.VkTensor) {
  605. const type = new vulkan.TensorType(arg);
  606. const initializer = arg.constant_id === -1 ? null : new vulkan.Tensor(arg, reader);
  607. const value = new vulkan.Value(id.toString(), type, initializer);
  608. values.set(id, { type: null, value: [value] });
  609. } else if (arg instanceof vkgraph.Int) {
  610. values.set(id, { type: 'int64', value: arg.int_val });
  611. } else if (arg instanceof vkgraph.IntList) {
  612. values.set(id, { type: 'int64[]', value: Array.from(arg.items) });
  613. } else if (arg instanceof vkgraph.Double) {
  614. values.set(id, { type: 'float64', value: arg.double_val });
  615. } else if (arg instanceof vkgraph.Bool) {
  616. values.set(id, { type: 'boolean', value: arg.bool_val });
  617. } else if (arg instanceof vkgraph.Null) {
  618. values.set(id, { type: 'attribute', value: null });
  619. } else {
  620. throw new Error(`Value type '${arg.constructor.name}' not implemented.`);
  621. }
  622. }
  623. return values.get(id);
  624. };
  625. for (let i = 0; i < graph.input_ids.length; i++) {
  626. const id = graph.input_ids[i];
  627. const value = values.map(id);
  628. const name = graph.input_ids.length === 1 ? 'input' : `input.${i}`;
  629. const argument = new vulkan.Argument(name, value.value, value.type);
  630. this.inputs.push(argument);
  631. }
  632. for (let i = 0; i < graph.output_ids.length; i++) {
  633. const id = graph.output_ids[i];
  634. const value = values.map(id);
  635. const name = graph.output_ids.length === 1 ? 'output' : `output.${i}`;
  636. const argument = new vulkan.Argument(name, value.value, value.type);
  637. this.outputs.push(argument);
  638. }
  639. for (const op of graph.chain) {
  640. const node = new vulkan.Node(metadata, op, values);
  641. this.nodes.push(node);
  642. }
  643. }
  644. };
  645. vulkan.Node = class {
  646. constructor(metadata, op, values) {
  647. const schema = metadata.type(op.name);
  648. if (!schema) {
  649. throw new vulkan.Error(`Operator schema for '${op.name}' not found.`);
  650. }
  651. this.type = {
  652. name: op.name.split(/\.([^.]*)$/)[0],
  653. identifier: op.name,
  654. category: schema.category || ''
  655. };
  656. this.name = op.node_id.toString();
  657. this.inputs = [];
  658. this.outputs = [];
  659. this.attributes = [];
  660. for (let i = 0; i < op.args.length; i++) {
  661. const arg = op.args[i];
  662. const input = schema && i < schema.arguments.length;
  663. const def = input ? schema.arguments[i] : schema.returns[i - schema.arguments.length];
  664. const value = values.map(arg);
  665. const argument = new vulkan.Argument(def.name || '', value.value, value.type);
  666. if (input) {
  667. this.inputs.push(argument);
  668. } else {
  669. this.outputs.push(argument);
  670. }
  671. }
  672. }
  673. };
  674. vulkan.Argument = class {
  675. constructor(name, value, type, visible) {
  676. this.name = name;
  677. this.value = value;
  678. this.type = type || null;
  679. this.visible = visible !== false;
  680. }
  681. };
  682. vulkan.Value = class Value {
  683. constructor(name, type, initializer) {
  684. if (typeof name !== 'string') {
  685. throw new executorch.Error(`Invalid value identifier '${JSON.stringify(name)}'.`);
  686. }
  687. this.name = name;
  688. this.type = initializer && initializer.type ? initializer.type : type || null;
  689. this.initializer = initializer || null;
  690. }
  691. };
  692. vulkan.TensorType = class {
  693. constructor(tensor) {
  694. const types = ['bool', 'uint8', 'int8', 'int32', 'float16', 'float32'];
  695. if (tensor.datatype >= types.length) {
  696. throw new vulkan.Error(`Unknown tensor data type '${tensor.datatype}'.`);
  697. }
  698. this.dataType = types[tensor.datatype];
  699. this.shape = new vulkan.TensorShape(Array.from(tensor.dims));
  700. }
  701. toString() {
  702. return this.dataType + this.shape.toString();
  703. }
  704. };
  705. vulkan.TensorShape = class {
  706. constructor(dimensions) {
  707. this.dimensions = dimensions || [];
  708. }
  709. toString() {
  710. if (this.dimensions && this.dimensions.length > 0) {
  711. return `[${this.dimensions.map((dimension) => dimension.toString()).join(',')}]`;
  712. }
  713. return '';
  714. }
  715. };
  716. vulkan.Tensor = class {
  717. constructor(tensor, reader) {
  718. this.type = new vulkan.TensorType(tensor);
  719. this.values = reader.constant(tensor.constant_id);
  720. this.encoding = '<';
  721. }
  722. };
  723. vulkan.Metadata = class {
  724. constructor(execution) {
  725. this.execution = execution;
  726. }
  727. register(signature) {
  728. const torch = this.execution.register('torch');
  729. const registry = torch._C.getRegistry();
  730. const schema = torch.FunctionSchema.parse(signature);
  731. const op = new torch._C.Operator(schema);
  732. registry.registerOperator(op);
  733. }
  734. type(identifier) {
  735. identifier = identifier.split(/\.([^.]*)$/);
  736. const name = identifier[0].replace('.', '::');
  737. const overload = identifier[1] === 'default' ? '' : identifier[1];
  738. const schemas = this.execution.invoke('torch._C._jit_get_schemas_for_operator', [name]);
  739. const schema = schemas.find((schema) => schema.name === name && schema.overload_name === overload);
  740. return schema;
  741. }
  742. };
  743. vulkan.Error = class extends Error {
  744. constructor(message) {
  745. super(message);
  746. this.name = 'Error loading Vulkan model.';
  747. }
  748. };
  749. coreml.Reader = class {
  750. static open(data, target) {
  751. const reader = base.BinaryReader.open(data);
  752. return new coreml.Reader(reader, target);
  753. }
  754. constructor(reader, target) {
  755. this.reader = reader;
  756. this.target = target;
  757. }
  758. async factory() {
  759. const coreml = await import('./coreml.js');
  760. return new coreml.ModelFactory();
  761. }
  762. async read() {
  763. const entries = this.entries(this.reader);
  764. const factory = await this.factory();
  765. const protobuf = await import('./protobuf.js');
  766. for (const [key, value] of entries) {
  767. const path = key.split('/');
  768. const identifier = path.pop();
  769. const folder = path.length === 0 ? '' : `${path.join('/')}/`;
  770. const locals = new Map(Array.from(entries).filter(([key]) => key.startsWith(folder)).map(([key, value]) => [key.substring(folder.length), value]));
  771. const context = new coreml.Context(this, identifier, value, locals, protobuf);
  772. factory.match(context);
  773. if (context.type === 'coreml.manifest') {
  774. /* eslint-disable no-await-in-loop */
  775. const model = await factory.open(context);
  776. /* eslint-enable no-await-in-loop */
  777. [this.type] = model.graphs;
  778. this.type.name = 'CoreMLBackend';
  779. return;
  780. }
  781. }
  782. }
  783. stream(offset, size) {
  784. this.reader.seek(offset);
  785. const stream = this.reader.stream(size);
  786. this.reader.seek(0);
  787. return stream;
  788. }
  789. entries(reader) {
  790. const files = new Map();
  791. reader.seek(reader.length - 1);
  792. const str = [];
  793. let depth = 0;
  794. do {
  795. const c = String.fromCharCode(reader.byte());
  796. reader.skip(-2);
  797. if (c === '{') {
  798. depth++;
  799. } else if (c === '}') {
  800. depth--;
  801. }
  802. str.push(c);
  803. } while (depth > 0);
  804. const metadata = JSON.parse(str.join(''));
  805. const nodes = metadata.nodes;
  806. const roots = Array.from(nodes);
  807. for (const root of roots) {
  808. if (root !== null) {
  809. for (const index of Object.values(root.children)) {
  810. roots[index] = null;
  811. }
  812. }
  813. }
  814. const process = (path, node) => {
  815. path = path ? `${path}/${node.name}` : node.name;
  816. if (node.kind === 0) {
  817. files.set(path, node.dataRegion);
  818. } else if (node.kind === 1) {
  819. for (const index of Object.values(node.children)) {
  820. process(path, nodes[index]);
  821. }
  822. } else {
  823. throw new Error(`Node kind '${node.kind}' not implemented.`);
  824. }
  825. };
  826. for (const root of roots.filter((node) => node !== null)) {
  827. process('', root);
  828. }
  829. return files;
  830. }
  831. };
  832. coreml.Context = class {
  833. constructor(reader, identifier, location, entries, protobuf) {
  834. this._reader = reader;
  835. this._location = location;
  836. this._identifier = identifier;
  837. this._entries = entries;
  838. this._protobuf = protobuf;
  839. }
  840. get identifier() {
  841. return this._identifier;
  842. }
  843. get stream() {
  844. if (!this._stream) {
  845. this._stream = this._reader.stream(this._location.offset, this._location.size);
  846. }
  847. return this._stream;
  848. }
  849. tags(type) {
  850. if (type === 'pb' && this.identifier.endsWith('.mlmodel')) {
  851. return new Map([[1,0],[2,2]]);
  852. }
  853. return new Map();
  854. }
  855. peek(type) {
  856. if (type === 'json') {
  857. const data = this.stream.peek();
  858. const decoder = new TextDecoder('utf-8');
  859. const text = decoder.decode(data);
  860. return JSON.parse(text);
  861. }
  862. return null;
  863. }
  864. read(type) {
  865. if (type === 'protobuf.binary') {
  866. return this._protobuf.BinaryReader.open(this.stream);
  867. }
  868. return null;
  869. }
  870. async fetch(file) {
  871. if (this._entries.has(file)) {
  872. const location = this._entries.get(file);
  873. const identifier = file.split('/').pop();
  874. return new coreml.Context(this._reader, identifier, location, this._entries, this._protobuf);
  875. }
  876. return null;
  877. }
  878. async require(id) {
  879. return this._reader.target.context.require(id);
  880. }
  881. async metadata(name) {
  882. return this._reader.target.context.metadata(name);
  883. }
  884. };
  885. export const ModelFactory = executorch.ModelFactory;