Skip to content

Commit ef35b8b

Browse files
lutzroederlutzroeder
committed
Update mlir.js (#1044)
1 parent aa41bf2 commit ef35b8b

File tree

2 files changed

+220
-95
lines changed

2 files changed

+220
-95
lines changed

‎source/mlir.js

Lines changed: 219 additions & 94 deletions
Original file line numberDiff line numberDiff line change
@@ -1193,108 +1193,221 @@ mlir.BytecodeReader = class {
11931193

11941194
constructor(context) {
11951195
this._reader = new mlir.BinaryReader(context);
1196-
this._decoder = new TextDecoder('utf-8');
11971196
}
11981197

11991198
read() {
12001199
const reader = this._reader;
12011200
reader.read(4); // signature 'ML\xEFR'
12021201
this.version = reader.varint().toNumber();
12031202
this.producer = reader.string();
1204-
this.sections = [];
1203+
this.sections = new Map();
12051204
while (reader.position < reader.length) {
12061205
// https://mlir.llvm.org/docs/BytecodeFormat/
1207-
const code = reader.byte();
1208-
const identifier = code & 0x7F;
1206+
// https://github.com/llvm/llvm-project/blob/main/mlir/lib/Bytecode/Reader/BytecodeReader.cpp
1207+
const sectionIDAndHasAlignment = reader.byte();
1208+
const sectionID = sectionIDAndHasAlignment & 0x7F;
12091209
const length = reader.varint().toNumber();
1210-
if (code >> 7) {
1210+
const hasAlignment = sectionIDAndHasAlignment & 0x80;
1211+
if (sectionID >= 9) {
1212+
throw new mlir.Error(`Unsupported section identifier '${sectionID}'.`);
1213+
}
1214+
if (hasAlignment) {
12111215
const alignment = reader.varint();
12121216
reader.skip(alignment);
12131217
}
1214-
const next = reader.position + length;
1215-
switch (identifier) {
1216-
case 0: { // string
1217-
const lengths = new Array(reader.varint().toNumber());
1218-
for (let i = 0; i < lengths.length; i++) {
1219-
lengths[i] = reader.varint().toNumber();
1220-
}
1221-
this.strings = new Array(lengths.length);
1222-
for (let i = 0; i < this.strings.length; i++) {
1223-
const size = lengths[lengths.length - i - 1];
1224-
const buffer = reader.read(size);
1225-
this.strings[i] = this._decoder.decode(buffer);
1226-
}
1227-
break;
1228-
}
1229-
case 1: { // dialect
1230-
const numDialects = reader.varint().toNumber();
1231-
this.dialectNames = new Array(numDialects);
1232-
this.opNames = new Array(numDialects);
1233-
for (let i = 0; i < this.dialectNames.length; i++) {
1234-
const group = {};
1235-
const nameAndIsVersioned = reader.varint();
1236-
group.name = (nameAndIsVersioned >> 1n).toNumber();
1237-
if (nameAndIsVersioned & 1n) {
1238-
const size = reader.varint().toNumber();
1239-
group.version = reader.read(size);
1240-
}
1241-
this.dialectNames[i] = group;
1242-
}
1243-
for (let i = 0; i < this.opNames.length; i++) {
1244-
const dialect_ops_group = {};
1245-
dialect_ops_group.dialect = reader.varint();
1246-
dialect_ops_group.opNames = new Array(reader.varint().toNumber());
1247-
for (let j = 0; j < dialect_ops_group.opNames.length; j++) {
1248-
const op_name_group = {};
1249-
const nameAndIsRegistered = reader.varint();
1250-
op_name_group.isRegistered = (nameAndIsRegistered & 1n) === 1n;
1251-
op_name_group.name = (nameAndIsRegistered >> 1n).toNumber();
1252-
dialect_ops_group.opNames[j] = op_name_group;
1253-
}
1254-
this.opNames[i] = dialect_ops_group;
1255-
}
1256-
break;
1257-
}
1258-
case 2: // attrType
1259-
case 3: { // attrTypeOffset
1260-
/*
1261-
const numAttrs = reader.varint().toNumber();
1262-
const numTypes = reader.varint().toNumber();
1263-
for (let i = 0; i < (numAttrs + numTypes); i++) {
1218+
const offset = reader.position;
1219+
reader.skip(length);
1220+
this.sections.set(sectionID, { start: offset, end: reader.position });
1221+
}
1222+
if (!this.sections.has(0) || !this.sections.has(1) ||
1223+
!this.sections.has(2) || !this.sections.has(3) ||
1224+
!this.sections.has(4) || (this.version >= 5 && !this.sections.has(8))) {
1225+
throw new mlir.Error('Missing required section.');
1226+
}
1227+
this._parseStringSection();
1228+
if (this.sections.has(8)) {
1229+
this._parsePropertiesSection();
1230+
}
1231+
this._parseDialectSection();
1232+
this._parseResourceSection();
1233+
this._parseAttrTypeSection();
1234+
}
12641235

1265-
}
1266-
break;
1267-
*/
1268-
this.attrType = reader.stream(length);
1269-
break;
1270-
}
1271-
case 4: { // IR
1272-
reader.skip(length);
1273-
break;
1274-
}
1275-
case 5: { // resource
1276-
this.resource = reader.stream(length);
1277-
break;
1278-
}
1279-
case 6: { // resourceOffset
1280-
reader.skip(length);
1281-
break;
1282-
}
1283-
case 7: { // dialectVersions
1284-
reader.skip(length);
1285-
break;
1236+
_parseStringSection() {
1237+
const section = this.sections.get(0);
1238+
const reader = this._reader;
1239+
reader.seek(section.start);
1240+
const lengths = new Array(reader.varint().toNumber());
1241+
for (let i = 0; i < lengths.length; i++) {
1242+
lengths[i] = reader.varint().toNumber();
1243+
}
1244+
const decoder = new TextDecoder('utf-8');
1245+
this.strings = new Array(lengths.length);
1246+
for (let i = 0; i < this.strings.length; i++) {
1247+
const size = lengths[lengths.length - 1 - i];
1248+
const buffer = reader.read(size);
1249+
this.strings[i] = decoder.decode(buffer);
1250+
}
1251+
if (reader.position !== section.end) {
1252+
throw new mlir.Error(`Invalid string section size.`);
1253+
}
1254+
}
1255+
1256+
_parseDialectSection() {
1257+
const section = this.sections.get(1);
1258+
const reader = this._reader;
1259+
reader.seek(section.start);
1260+
const numDialects = reader.varint().toNumber();
1261+
this.dialects = new Array(numDialects);
1262+
for (let i = 0; i < this.dialects.length; i++) {
1263+
this.dialects[i] = {};
1264+
if (this.version < 1) { // kDialectVersioning
1265+
const entryIdx = reader.varint().toNumber();
1266+
this.dialects[i].name = this.strings[entryIdx];
1267+
continue;
1268+
}
1269+
const nameAndIsVersioned = reader.varint();
1270+
const dialectNameIdx = (nameAndIsVersioned >> 1n).toNumber();
1271+
this.dialects[i].name = this.strings[dialectNameIdx];
1272+
if (nameAndIsVersioned & 1n) {
1273+
const size = reader.varint().toNumber();
1274+
this.dialects[i].version = reader.read(size);
1275+
}
1276+
}
1277+
let numOps = -1;
1278+
this.opNames = [];
1279+
if (this.version > 4) { // kElideUnknownBlockArgLocation
1280+
numOps = reader.varint().toNumber();
1281+
this.opNames = new Array(numOps);
1282+
}
1283+
let i = 0;
1284+
while (reader.position < section.end) {
1285+
const dialect = this.dialects[reader.varint().toNumber()];
1286+
const numEntries = reader.varint().toNumber();
1287+
for (let j = 0; j < numEntries; j++) {
1288+
const opName = {};
1289+
if (this.version < 5) { // kNativePropertiesEncoding
1290+
opName.name = this.strings[reader.varint().toNumber()];
1291+
opName.dialect = dialect;
1292+
} else {
1293+
const nameAndIsRegistered = reader.varint();
1294+
opName.name = this.strings[(nameAndIsRegistered >> 1n).toNumber()];
1295+
opName.dialect = dialect;
1296+
opName.isRegistered = (nameAndIsRegistered & 1n) === 1n;
12861297
}
1287-
case 8: { // properties
1288-
reader.skip(length);
1289-
break;
1298+
if (numOps < 0) {
1299+
this.opNames.push(opName);
1300+
} else {
1301+
this.opNames[i++] = opName;
12901302
}
1291-
default: {
1292-
throw new mlir.Error(`Unsupported section identifier '${identifier}'.`);
1303+
}
1304+
}
1305+
if (reader.position !== section.end) {
1306+
throw new mlir.Error(`Invalid dialect section size.`);
1307+
}
1308+
}
1309+
1310+
_parseResourceSection() {
1311+
const section = this.sections.get(6);
1312+
const reader = this._reader;
1313+
reader.seek(section.start);
1314+
const numExternalResourceGroups = reader.varint().toNumber();
1315+
if (numExternalResourceGroups > 0) {
1316+
throw new mlir.Error(`Unsupported resource section.`);
1317+
}
1318+
/*
1319+
for (let i = 0; i < numExternalResourceGroups; i++) {
1320+
const numResources = reader.varint().toNumber();
1321+
for (let j = 0; j < numResources; j++) {
1322+
const resource = {};
1323+
resource.key = this.strings[reader.varint().toNumber()];
1324+
resource.offset = reader.varint().toNumber();
1325+
resource.kind = reader.byte();
1326+
}
1327+
}
1328+
*/
1329+
if (reader.position !== section.end) {
1330+
throw new mlir.Error(`Invalid dialect section size.`);
1331+
}
1332+
}
1333+
1334+
_parseAttrTypeSection() {
1335+
const section = this.sections.get(3);
1336+
const reader = this._reader;
1337+
reader.seek(section.start);
1338+
this.attributes = new Array(reader.varint().toNumber());
1339+
this.types = new Array(reader.varint().toNumber());
1340+
let offset = 0;
1341+
const parseEntries = (range) => {
1342+
for (let i = 0; i < range.length;) {
1343+
const dialect = this.dialects[reader.varint().toNumber()];
1344+
const numEntries = reader.varint().toNumber();
1345+
for (let j = 0; j < numEntries; j++) {
1346+
const entry = {};
1347+
const entrySizeWithFlag = reader.varint();
1348+
entry.hasCustomEncoding = (entrySizeWithFlag & 1n) === 1n;
1349+
entry.size = (entrySizeWithFlag >> 1n).toNumber();
1350+
entry.offset = offset;
1351+
entry.dialect = dialect;
1352+
offset += entry.size;
1353+
range[i++] = entry;
12931354
}
12941355
}
1295-
if (reader.position !== next) {
1296-
throw new mlir.Error('Invalid section length.');
1356+
};
1357+
parseEntries(this.attributes);
1358+
parseEntries(this.types);
1359+
if (reader.position !== section.end) {
1360+
throw new mlir.Error(`Invalid dialect section size.`);
1361+
}
1362+
offset = this.sections.get(2).start;
1363+
const parseCustomEntry = (entry, reader, entryType) => {
1364+
// throw new mlir.Error(`Unsupported custom encoding.`);
1365+
if (entryType === 'type') {
1366+
// debugger;
1367+
} else {
1368+
// debugger;
12971369
}
1370+
};
1371+
const parseAsmEntry = (entry, reader, entryType) => {
1372+
if (entryType === 'type') {
1373+
// debugger;
1374+
} else {
1375+
// debugger;
1376+
}
1377+
};
1378+
const resolveEntries = (range, entryType) => {
1379+
for (const entry of this.attributes) {
1380+
reader.seek(offset + entry.offset);
1381+
if (entry.hasCustomEncoding) {
1382+
parseCustomEntry(entry, reader);
1383+
} else {
1384+
parseAsmEntry(entry, reader, entryType);
1385+
}
1386+
// if (reader.position !== (offset + entry.offset + entry.size)) {
1387+
// throw new mlir.Error(`Invalid '${entryType}' section size.`);
1388+
// }
1389+
// delete entry.offset;
1390+
// delete entry.size;
1391+
}
1392+
};
1393+
resolveEntries(this.attributes, 'attribute');
1394+
resolveEntries(this.types, 'type');
1395+
}
1396+
1397+
_parsePropertiesSection() {
1398+
const section = this.sections.get(8);
1399+
const reader = this._reader;
1400+
reader.seek(section.start);
1401+
const count = reader.varint().toNumber();
1402+
const offsetTable = new Array(count);
1403+
for (let i = 0; i < offsetTable.length; i++) {
1404+
const offset = reader.position;
1405+
const size = reader.varint().toNumber();
1406+
const data = reader.read(size);
1407+
offsetTable[i] = { offset, data };
1408+
}
1409+
if (reader.position !== section.end) {
1410+
throw new mlir.Error(`Invalid properties section size.`);
12981411
}
12991412
}
13001413
};
@@ -1317,6 +1430,10 @@ mlir.BinaryReader = class {
13171430
this._reader.skip(length);
13181431
}
13191432

1433+
seek(offset) {
1434+
this._reader.seek(offset);
1435+
}
1436+
13201437
read(length) {
13211438
return this._reader.read(length);
13221439
}
@@ -1330,17 +1447,25 @@ mlir.BinaryReader = class {
13301447
}
13311448

13321449
varint() {
1333-
let value = 0n;
1334-
let shift = 0n;
1335-
for (let i = 0; i < 10 && this._reader.position < this._reader.length; i++) {
1336-
const byte = this._reader.byte();
1337-
value |= BigInt(byte >> 1) << shift;
1338-
if ((byte & 1) === 1) {
1339-
return value;
1340-
}
1341-
shift += 7n;
1342-
}
1343-
throw new mlir.Error('Invalid varint value.');
1450+
let result = this._reader.byte();
1451+
if (result & 1) {
1452+
return BigInt(result >> 1);
1453+
}
1454+
if (result === 0) {
1455+
return this._reader.uint64();
1456+
}
1457+
result = BigInt(result);
1458+
let mask = 1n;
1459+
let numBytes = 0n;
1460+
let shift = 8n;
1461+
while (result > 0n && (result & mask) === 0n) {
1462+
result |= (BigInt(this._reader.byte()) << shift);
1463+
mask <<= 1n;
1464+
shift += 8n;
1465+
numBytes++;
1466+
}
1467+
result >>= BigInt(numBytes + 1n);
1468+
return result;
13441469
}
13451470

13461471
string() {

‎test/models.json

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3113,7 +3113,7 @@
31133113
"target": "model.mlirbc",
31143114
"source": "https://github.com/user-attachments/files/17179955/model.mlirbc.zip[model.mlirbc]",
31153115
"format": "MLIR",
3116-
"error": "Invalid section length.",
3116+
"error": "Invalid content. File contains MLIR bytecode data.",
31173117
"link": "https://github.com/lutzroeder/netron/issues/1044"
31183118
},
31193119
{

0 commit comments

Comments
 (0)